From 84737acea59cfc635cc10a921e6a01368f5424b4 Mon Sep 17 00:00:00 2001 From: Jean Luca Bez Date: Mon, 6 Jul 2026 14:23:50 -0700 Subject: [PATCH] fixes to track input and output and responses in mlflow --- src/dsagt/mcp/server.py | 8 +++- src/dsagt/memory.py | 15 ++++-- src/dsagt/provenance.py | 55 ++++++++++++++-------- tests/test_dsagt_server.py | 33 +++++++++++++ tests/test_memory_extractor.py | 18 +++++++- tests/test_tool_executions.py | 84 ++++++++++++++++++++++++++-------- 6 files changed, 169 insertions(+), 44 deletions(-) diff --git a/src/dsagt/mcp/server.py b/src/dsagt/mcp/server.py index e010140..9fee520 100644 --- a/src/dsagt/mcp/server.py +++ b/src/dsagt/mcp/server.py @@ -92,7 +92,7 @@ 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: @@ -100,6 +100,12 @@ async def call_tool(tool_name: str, arguments: dict) -> list[types.TextContent]: 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) diff --git a/src/dsagt/memory.py b/src/dsagt/memory.py index e2d249f..87e2c28 100644 --- a/src/dsagt/memory.py +++ b/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) diff --git a/src/dsagt/provenance.py b/src/dsagt/provenance.py index 9abe20c..892b520 100644 --- a/src/dsagt/provenance.py +++ b/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=`` 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") # --------------------------------------------------------------------------- diff --git a/tests/test_dsagt_server.py b/tests/test_dsagt_server.py index 0174744..87c7ffe 100644 --- a/tests/test_dsagt_server.py +++ b/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) diff --git a/tests/test_memory_extractor.py b/tests/test_memory_extractor.py index 7015035..a7b2f2a 100644 --- a/tests/test_memory_extractor.py +++ b/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" diff --git a/tests/test_tool_executions.py b/tests/test_tool_executions.py index 2e60c67..9dabc6a 100644 --- a/tests/test_tool_executions.py +++ b/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() # ---------------------------------------------------------------------------