fixes to track input and output and responses in mlflow

src/dsagt/mcp/server.py

@@ -92,14 +92,20 @@ async def list_tools() -> list[types.Tool]:
     @server.call_tool()
     async def call_tool(tool_name: str, arguments: dict) -> list[types.TextContent]:
         handler = handlers[tool_name]  # KeyError = bug in list_tools schema
-        with open_span(tool_name, source=tool_category.get(tool_name)):
+        with open_span(tool_name, source=tool_category.get(tool_name)) as span:
             try:
                 result = await handler(arguments)
             except ValueError as e:
                 result = {"status": "error", "error": str(e)}
             except Exception as e:
                 logger.exception("Unexpected error in tool '%s'", tool_name)
                 result = {"status": "error", "error": f"Unexpected error: {e}"}
+            if span is not None:
+                # The trace-level Request/Inputs/Outputs are read from this
+                # categorization root; record the call's arguments and result so
+                # the MLflow UI shows them instead of a null request.
+                span.set_inputs(arguments)
+                span.set_outputs(result)
         text = (
             result
             if isinstance(result, str)

src/dsagt/memory.py

@@ -381,10 +381,16 @@ def write(self, trace: "Trace") -> None:
         # (This tags the *observability* span, not the memory chunks.)
         from dsagt.observability import open_span
 
-        with open_span("memory.extract", source="episodic"):
-            self._index_turns(exchanges)
-
-    def _index_turns(self, exchanges: list[dict]) -> None:
+        with open_span("memory.extract", source="episodic") as span:
+            n = self._index_turns(exchanges)
+            if span is not None:
+                # Trace-level Request/Inputs/Outputs are read from this root, so
+                # record the turn count in and chunk count out to keep the trace
+                # from showing a null request in the MLflow UI.
+                span.set_inputs({"n_turns": len(exchanges)})
+                span.set_outputs({"chunks_indexed": n})
+
+    def _index_turns(self, exchanges: list[dict]) -> int:
         """Chunk each turn per-block and embed into session_memory."""
         tool_names = _resolve_tool_names(exchanges)
         texts, metas = [], []
@@ -407,3 +413,4 @@ def _index_turns(self, exchanges: list[dict]) -> None:
             self._kb.add_entries(
                 texts=texts, collection=SESSION_MEMORY_COLLECTION, metadatas=metas
             )
+        return len(texts)

src/dsagt/provenance.py

@@ -284,6 +284,8 @@ def index_trace_archive(
     trace_dir: Path,
     kb: KnowledgeBase,
     indexed_ids: set[str] | None = None,
+    *,
+    source: str | None = None,
 ) -> dict:
     """Batch-index all code execution records in a trace archive directory."""
     if indexed_ids is None:
@@ -323,11 +325,25 @@ def index_trace_archive(
 
     # No ``route=`` — see ``index_execution_record`` above.
     if texts:
-        kb.add_entries(
-            texts=texts,
-            collection=CODE_USE_COLLECTION,
-            metadatas=metadatas,
-        )
+        from contextlib import nullcontext
+
+        from dsagt.observability import open_span
+
+        # Open a categorization root only when there is work to index — a quiet
+        # heartbeat produces no child spans, so wrapping it would just emit an
+        # empty, null-request trace.  Only the tagged background triggers pass a
+        # ``source``; the reconstruct-pipeline caller passes none and lets its
+        # kb.* writes inherit the tool's own trace.
+        cm = open_span("code_use.index", source=source) if source else nullcontext(None)
+        with cm as span:
+            kb.add_entries(
+                texts=texts,
+                collection=CODE_USE_COLLECTION,
+                metadatas=metadatas,
+            )
+            if span is not None:
+                span.set_inputs({"trace_dir": str(trace_dir), "n_records": len(texts)})
+                span.set_outputs({"indexed": len(texts)})
 
     return {
         "indexed": len(texts),
@@ -381,21 +397,25 @@ def _lock(self):
             finally:
                 fcntl.flock(lf, fcntl.LOCK_UN)
 
-    def tick(self) -> int:
+    def tick(self, *, source: str | None = None) -> int:
         """Index newly-arrived records; return how many were indexed this tick.
 
-        Emits no categorization root of its own: at the ``reconstruct_pipeline``
-        call site this runs *inside* the registry tool's trace, so its ``kb.*``
+        Passing ``source`` opens a ``dsagt.source=<source>`` categorization root
+        around the actual indexing (see :func:`index_trace_archive`) — but only
+        when records are indexed.  At the ``reconstruct_pipeline`` call site this
+        runs *inside* the registry tool's trace with no source, so its ``kb.*``
         writes correctly inherit ``dsagt.source=registry``.  The background
-        callers (heartbeat / startup catch-up) run outside any trace and use
-        :meth:`tick_traced` so their writes don't orphan as untagged roots.
+        callers use :meth:`tick_traced` so their writes don't orphan as untagged
+        roots.
         """
         with self._lock():
             acks = self._load_acks()
             before = len(acks)
             # index_trace_archive skips record_ids already in ``acks`` and adds
             # the newly-indexed ones to it (mutates the set we pass).
-            result = index_trace_archive(self._trace_dir, self._kb, indexed_ids=acks)
+            result = index_trace_archive(
+                self._trace_dir, self._kb, indexed_ids=acks, source=source
+            )
             if len(acks) != before:
                 self._save_acks(acks)
             return result.get("indexed", 0)
@@ -406,15 +426,12 @@ def tick_traced(self) -> int:
         For the background triggers (heartbeat, startup catch-up) that run off
         any tool-call trace — otherwise the indexer's ``kb.add_entries`` /
         ``kb.embed`` spans start their own untagged top-level traces, landing in
-        the ``unknown`` bucket and detached from the executions they index.
-        Must run on the same thread as the embedding (open the span inside the
-        worker), so callers dispatch *this* to the thread, not a wrapped
-        :meth:`tick`.
+        the ``unknown`` bucket and detached from the executions they index.  The
+        root is opened only when a tick actually indexes records, so a quiet
+        heartbeat emits no empty trace.  Runs on the caller's thread (callers
+        dispatch *this* to the embedding worker), so the span opens there.
         """
-        from dsagt.observability import open_span
-
-        with open_span("code_use.index", source="code_use"):
-            return self.tick()
+        return self.tick(source="code_use")
 
 
 # ---------------------------------------------------------------------------

tests/test_dsagt_server.py

@@ -84,6 +84,39 @@ def test_merged_server_exposes_all_tools(tmp_path):
     assert len(set(names)) == len(names)  # no name collision
 
 
+def test_dispatch_root_span_records_tool_inputs_and_outputs(tmp_path, monkeypatch):
+    """The categorization-root span must carry the tool arguments as its inputs
+    and the handler result as its outputs.
+
+    Without this, every MCP tool trace shows a null Request and empty
+    Inputs/Outputs in the MLflow UI, because the trace-level fields are read
+    from the root span and the dispatch wrapper is the root.
+    """
+    import mlflow
+
+    import dsagt.observability as obs_module
+    from dsagt.mcp.server import build_dispatch_server
+
+    mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/mlflow.db")
+    mlflow.set_experiment("test")
+    monkeypatch.setattr(obs_module, "_initialized", True)
+    monkeypatch.setattr(obs_module, "_default_session_id", None)
+
+    async def echo(args):
+        return {"echoed": args["q"]}
+
+    tools = [types.Tool(name="demo", description="d", inputSchema={"type": "object"})]
+    server = build_dispatch_server("test", tools, {"demo": echo}, {"demo": "knowledge"})
+
+    out = _call(server, "demo", {"q": "hello"})
+    assert json.loads(out) == {"echoed": "hello"}
+
+    trace = mlflow.MlflowClient().get_trace(mlflow.get_last_active_trace_id())
+    root = next(s for s in trace.data.spans if s.name == "demo")
+    assert root.inputs == {"q": "hello"}
+    assert root.outputs == {"echoed": "hello"}
+
+
 def test_registry_tool_returns_plain_string(tmp_path):
     """Registry handlers return a bare string — passed through unchanged."""
     server = _make_merged_server(tmp_path)

tests/test_memory_extractor.py

@@ -109,10 +109,11 @@ def test_empty_trace_writes_nothing(tmp_path):
     assert kb.calls == []
 
 
-def test_extraction_span_tagged_episodic_not_memory(tmp_path):
+def test_extraction_span_tagged_episodic_with_inputs_outputs(tmp_path):
     """The per-turn embedding runs off the heartbeat, so it must carry
     dsagt.source=episodic — filtering apart from the user-facing memory tools
-    (kb_remember / kb_get_memories) that carry dsagt.source=memory."""
+    (kb_remember / kb_get_memories) that carry dsagt.source=memory — and record
+    its turn/chunk counts as inputs/outputs so the trace isn't null-request."""
     from unittest.mock import patch
 
     opened = {}
@@ -124,7 +125,14 @@ def __enter__(self):
         def __exit__(self, *a):
             return False
 
+        def set_inputs(self, v):
+            opened["inputs"] = v
+
+        def set_outputs(self, v):
+            opened["outputs"] = v
+
     def _fake_open_span(name, span_type=None, source=None):
+        opened["name"] = name
         opened["source"] = source
         return _Span()
 
@@ -133,4 +141,10 @@ def _fake_open_span(name, span_type=None, source=None):
     with patch("dsagt.observability.open_span", _fake_open_span):
         ext.write(_one_turn_trace())
 
+    assert opened["name"] == "memory.extract"
+    assert opened["source"] == "episodic"
+    # one exchange in; two chunks (user question + assistant answer) out.
+    assert opened["inputs"] == {"n_turns": 1}
+    assert opened["outputs"] == {"chunks_indexed": 2}
+
     assert opened["source"] == "episodic"

tests/test_tool_executions.py

@@ -10,6 +10,7 @@
 """
 
 import json
+from contextlib import nullcontext
 from pathlib import Path
 from unittest.mock import MagicMock, patch
 
@@ -201,35 +202,82 @@ def test_incremental_and_idempotent(self, tmp_path):
             assert set(acks) == {"r1", "r2"}
             kb.close()
 
-    def test_tick_traced_wraps_in_code_use_source_span(self, tmp_path):
-        """The background triggers use tick_traced so the indexer's kb.* writes
-        nest under a dsagt.source=code_use root instead of orphaning."""
+    def test_tick_traced_opens_no_span_when_nothing_to_index(self, tmp_path):
+        """A quiet heartbeat (no new records) must open NO categorization root —
+        otherwise the MLflow trace list fills with empty, null-request traces."""
         pdir = tmp_path / "proj"
         (pdir / ".dsagt").mkdir(parents=True)
         kb = MagicMock()
         indexer = CodeUseIndexer(kb, pdir)
 
-        opened = {}
-
-        class _Span:
-            def __enter__(self):
-                return self
-
-            def __exit__(self, *a):
-                return False
+        opened = []
 
         def _fake_open_span(name, span_type=None, source=None):
-            opened["name"] = name
-            opened["source"] = source
-            return _Span()
+            opened.append((name, source))
+            return nullcontext(None)
 
         with patch("dsagt.observability.open_span", _fake_open_span):
-            # No records → tick returns 0, but the span must still be opened
-            # with the right source.
             assert indexer.tick_traced() == 0
 
-        assert opened["source"] == "code_use"
-        assert opened["name"] == "code_use.index"
+        assert opened == []  # nothing indexed → no span, no trace
+
+    def test_tick_traced_wraps_real_work_in_code_use_span(self, tmp_path):
+        """When a tick indexes records, it opens one code_use.index root tagged
+        dsagt.source=code_use with populated inputs/outputs (not a null trace)."""
+        with patch("dsagt.knowledge.Embedder.create") as mock_make:
+            mock_embedder = MagicMock()
+            mock_embedder.embed = fake_embed
+            mock_make.return_value = mock_embedder
+
+            pdir = tmp_path / "proj"
+            (pdir / ".dsagt").mkdir(parents=True)
+            kb = KnowledgeBase(index_dir=pdir / "kb")
+            indexer = CodeUseIndexer(kb, pdir)
+            r1 = make_wrapper_record()
+            r1["record_id"] = "r1"
+            self._write(pdir / "trace_archive", r1)
+
+            spans = []
+
+            class _RecSpan:
+                def __init__(self, name, source):
+                    self.name, self.source = name, source
+                    self.inputs = self.outputs = None
+
+                def __enter__(self):
+                    return self
+
+                def __exit__(self, *a):
+                    return False
+
+                def set_inputs(self, v):
+                    self.inputs = v
+
+                def set_outputs(self, v):
+                    self.outputs = v
+
+            def _fake_open_span(name, span_type=None, source=None):
+                # Only intercept the categorization root; nested kb.* @traced
+                # spans no-op as they would with tracing off.
+                if name != "code_use.index":
+                    return nullcontext(None)
+                s = _RecSpan(name, source)
+                spans.append(s)
+                return s
+
+            with patch("dsagt.observability.open_span", _fake_open_span):
+                assert indexer.tick_traced() == 1
+
+            assert len(spans) == 1
+            span = spans[0]
+            assert span.name == "code_use.index"
+            assert span.source == "code_use"
+            assert span.inputs == {
+                "trace_dir": str(pdir / "trace_archive"),
+                "n_records": 1,
+            }
+            assert span.outputs == {"indexed": 1}
+            kb.close()
 
 
 # ---------------------------------------------------------------------------