refactor: correlate lineage by request ID

src/rlm/client.py

@@ -16,16 +16,8 @@
 
 from rlm.config import ProviderConfig
 from rlm.lineage import (
-    COMPACTION_ID_HEADER,
-    CONTEXT_ID_HEADER,
-    DEPTH_HEADER,
     LINEAGE_HEADER_NAMES,
-    PARENT_SESSION_ID_HEADER,
-    PREVIOUS_CONTEXT_ID_HEADER,
     REQUEST_ID_HEADER,
-    SESSION_ID_HEADER,
-    TRANSITION_HEADER,
-    RequestProvenance,
 )
 from rlm.types import TokenUsage
 
@@ -92,25 +84,12 @@ def make_client(provider: ProviderConfig | None = None) -> AsyncOpenAI:
     )
 
 
-def model_call_headers(provenance: RequestProvenance | str) -> dict[str, str]:
+def model_call_headers(request_id: str) -> dict[str, str]:
     """Build transport headers for one idempotent, attributable model call."""
-    if isinstance(provenance, str):
-        return {IDEMPOTENCY_KEY_HEADER: provenance}
-    headers = {
-        IDEMPOTENCY_KEY_HEADER: provenance.request_id,
-        REQUEST_ID_HEADER: provenance.request_id,
-        SESSION_ID_HEADER: provenance.session_id,
-        CONTEXT_ID_HEADER: provenance.context_id,
-        TRANSITION_HEADER: provenance.transition,
-        DEPTH_HEADER: str(provenance.depth),
+    return {
+        IDEMPOTENCY_KEY_HEADER: request_id,
+        REQUEST_ID_HEADER: request_id,
     }
-    if provenance.parent_session_id is not None:
-        headers[PARENT_SESSION_ID_HEADER] = provenance.parent_session_id
-    if provenance.previous_context_id is not None:
-        headers[PREVIOUS_CONTEXT_ID_HEADER] = provenance.previous_context_id
-    if provenance.compaction_id is not None:
-        headers[COMPACTION_ID_HEADER] = provenance.compaction_id
-    return headers
 
 
 async def call_with_retries(

src/rlm/engine.py

@@ -399,15 +399,15 @@ async def _run_loop(self) -> RLMResult:
         for turn in itertools.count(self._turn):
             self._turn = turn + 1
             # Call LLM
-            provenance = self._lineage.start_request(
+            request_id = self._lineage.start_request(
                 self._invocation_id,
                 kind="turn",
             )
-            call_id = provenance.request_id
+            call_id = request_id
             request_kwargs = {
                 "model": self.model,
                 "messages": messages,
-                "extra_headers": model_call_headers(provenance),
+                "extra_headers": model_call_headers(request_id),
             }
             if self._active_tool_schemas:
                 request_kwargs["tools"] = self._active_tool_schemas
@@ -747,15 +747,15 @@ async def _compact_branch(
             checkpoint_prompt += REPL_RESTART_NOTE
         messages.append({"role": "user", "content": checkpoint_prompt})
         compaction = self._lineage.begin_compaction(self._invocation_id)
-        provenance = self._lineage.start_request(
+        request_id = self._lineage.start_request(
             self._invocation_id,
             kind="compaction",
             compaction_id=compaction.compaction_id,
         )
         request_kwargs: dict = {
             "model": self.model,
             "messages": messages,
-            "extra_headers": model_call_headers(provenance),
+            "extra_headers": model_call_headers(request_id),
         }
         if active_tools:
             request_kwargs["tools"] = active_tools

src/rlm/lineage.py

@@ -13,35 +13,7 @@
 CompactionStatus = Literal["in_progress", "completed", "failed", "cancelled"]
 
 REQUEST_ID_HEADER = "X-ACP-Lineage-Request-ID"
-SESSION_ID_HEADER = "X-ACP-Lineage-Session-ID"
-PARENT_SESSION_ID_HEADER = "X-ACP-Lineage-Parent-Session-ID"
-CONTEXT_ID_HEADER = "X-ACP-Lineage-Context-ID"
-PREVIOUS_CONTEXT_ID_HEADER = "X-ACP-Lineage-Previous-Context-ID"
-TRANSITION_HEADER = "X-ACP-Lineage-Transition"
-COMPACTION_ID_HEADER = "X-ACP-Lineage-Compaction-ID"
-DEPTH_HEADER = "X-ACP-Lineage-Depth"
-LINEAGE_HEADER_NAMES = (
-    REQUEST_ID_HEADER,
-    SESSION_ID_HEADER,
-    PARENT_SESSION_ID_HEADER,
-    CONTEXT_ID_HEADER,
-    PREVIOUS_CONTEXT_ID_HEADER,
-    TRANSITION_HEADER,
-    COMPACTION_ID_HEADER,
-    DEPTH_HEADER,
-)
-
-
-@dataclass(frozen=True)
-class RequestProvenance:
-    request_id: str
-    session_id: str
-    parent_session_id: str | None
-    context_id: str
-    previous_context_id: str | None
-    transition: ContextTransition
-    compaction_id: str | None
-    depth: int
+LINEAGE_HEADER_NAMES = (REQUEST_ID_HEADER,)
 
 
 @dataclass(frozen=True)
@@ -121,10 +93,9 @@ def start_request(
         *,
         kind: RequestKind,
         compaction_id: str | None = None,
-    ) -> RequestProvenance:
+    ) -> str:
         if kind == "compaction" and compaction_id is None:
             raise ValueError("compaction requests require a compaction ID")
-        session = self._sessions[session_id]
         context_id = self._active_contexts[session_id]
         context = self._contexts[context_id]
         effective_compaction_id = compaction_id or context.get("compaction_id")
@@ -142,16 +113,7 @@ def start_request(
         if kind == "compaction":
             self._compactions[compaction_id]["summary_request_id"] = request_id
 
-        return RequestProvenance(
-            request_id=request_id,
-            session_id=session_id,
-            parent_session_id=session.get("parent_session_id"),
-            context_id=context_id,
-            previous_context_id=context.get("previous_context_id"),
-            transition=context["transition"],
-            compaction_id=effective_compaction_id,
-            depth=session["depth"],
-        )
+        return request_id
 
     def begin_compaction(self, session_id: str) -> Compaction:
         compaction_id = uuid.uuid4().hex

tests/test_acp.py

@@ -357,19 +357,21 @@ async def flaky_first_call(**kwargs):
         header["Idempotency-Key"] == header["X-ACP-Lineage-Request-ID"]
         for header in (turn, compaction, resumed)
     )
-    assert turn["X-ACP-Lineage-Session-ID"] == engine._invocation_id
-    assert turn["X-ACP-Lineage-Context-ID"] == compaction["X-ACP-Lineage-Context-ID"]
-    assert turn["X-ACP-Lineage-Transition"] == "root"
-    assert "X-ACP-Lineage-Compaction-ID" not in turn
-    assert (
-        compaction["X-ACP-Lineage-Compaction-ID"]
-        == resumed["X-ACP-Lineage-Compaction-ID"]
-    )
-    assert (
-        resumed["X-ACP-Lineage-Previous-Context-ID"] == turn["X-ACP-Lineage-Context-ID"]
-    )
-    assert resumed["X-ACP-Lineage-Context-ID"] != turn["X-ACP-Lineage-Context-ID"]
-    assert resumed["X-ACP-Lineage-Transition"] == "compact"
+    lineage = engine.execution_snapshot()["lineage"]
+    requests = {request["request_id"]: request for request in lineage["requests"]}
+    contexts = {context["context_id"]: context for context in lineage["contexts"]}
+    turn_request, compaction_request, resumed_request = [
+        requests[headers["X-ACP-Lineage-Request-ID"]]
+        for headers in (turn, compaction, resumed)
+    ]
+    assert turn_request["session_id"] == engine._invocation_id
+    assert turn_request["context_id"] == compaction_request["context_id"]
+    assert compaction_request["kind"] == "compaction"
+    assert resumed_request["context_id"] != turn_request["context_id"]
+    resumed_context = contexts[resumed_request["context_id"]]
+    assert resumed_context["previous_context_id"] == turn_request["context_id"]
+    assert resumed_context["transition"] == "compact"
+    assert resumed_request["compaction_id"] == compaction_request["compaction_id"]
 
 
 async def test_latest_cancelled_prompt_does_not_finalize_prior_result(session):
@@ -496,14 +498,18 @@ async def test_failed_prompt_restores_pre_compaction_context(session):
     finally:
         engine.close()
 
-    initial, summary, compacted, retried = [
-        call["extra_headers"] for call in client.calls
+    request_ids = [
+        call["extra_headers"]["X-ACP-Lineage-Request-ID"] for call in client.calls
     ]
+    lineage = engine.execution_snapshot()["lineage"]
+    requests = {request["request_id"]: request for request in lineage["requests"]}
+    contexts = {context["context_id"]: context for context in lineage["contexts"]}
+    initial, summary, compacted, retried = [requests[item] for item in request_ids]
     assert result.answer == "continued"
-    assert summary["X-ACP-Lineage-Context-ID"] == initial["X-ACP-Lineage-Context-ID"]
-    assert compacted["X-ACP-Lineage-Context-ID"] != initial["X-ACP-Lineage-Context-ID"]
-    assert retried["X-ACP-Lineage-Context-ID"] == initial["X-ACP-Lineage-Context-ID"]
-    assert retried["X-ACP-Lineage-Transition"] == "root"
+    assert summary["context_id"] == initial["context_id"]
+    assert compacted["context_id"] != initial["context_id"]
+    assert retried["context_id"] == initial["context_id"]
+    assert contexts[retried["context_id"]]["transition"] == "root"
 
 
 async def test_engine_cancel_masks_tool_cleanup_error(monkeypatch, session):

tests/test_lineage.py

@@ -18,7 +18,7 @@ def test_root_and_child_requests_have_exact_parentage():
         "child-1",
         parent_session_id="trace-1",
         depth=1,
-        spawned_by_request_id=root_request.request_id,
+        spawned_by_request_id=root_request,
     )
     child_request = lineage.start_request("child-1", kind="turn")
 
@@ -27,18 +27,28 @@ def test_root_and_child_requests_have_exact_parentage():
     snapshot = lineage.snapshot()
 
     assert root_headers == {
-        "Idempotency-Key": root_request.request_id,
-        "X-ACP-Lineage-Request-ID": root_request.request_id,
-        "X-ACP-Lineage-Session-ID": "trace-1",
-        "X-ACP-Lineage-Context-ID": root_context_id,
-        "X-ACP-Lineage-Transition": "root",
-        "X-ACP-Lineage-Depth": "0",
+        "Idempotency-Key": root_request,
+        "X-ACP-Lineage-Request-ID": root_request,
     }
-    assert child_headers["X-ACP-Lineage-Parent-Session-ID"] == "trace-1"
-    assert child_headers["X-ACP-Lineage-Context-ID"] == child_context_id
-    assert child_headers["X-ACP-Lineage-Transition"] == "spawn"
-    assert child_headers["X-ACP-Lineage-Depth"] == "1"
-    assert snapshot["sessions"][1]["spawned_by_request_id"] == (root_request.request_id)
+    assert child_headers == {
+        "Idempotency-Key": child_request,
+        "X-ACP-Lineage-Request-ID": child_request,
+    }
+    assert snapshot["sessions"][1]["spawned_by_request_id"] == root_request
+    assert snapshot["requests"] == [
+        {
+            "request_id": root_request,
+            "session_id": "trace-1",
+            "context_id": root_context_id,
+            "kind": "turn",
+        },
+        {
+            "request_id": child_request,
+            "session_id": "child-1",
+            "context_id": child_context_id,
+            "kind": "turn",
+        },
+    ]
 
 
 def test_terminal_session_status_cannot_be_overwritten():
@@ -64,26 +74,31 @@ def test_completed_compaction_starts_linked_context_epoch():
     lineage.finish_compaction(compaction.compaction_id, "completed")
     resumed_request = lineage.start_request("trace-1", kind="turn")
 
-    summary_headers = model_call_headers(summary_request)
-    resumed_headers = model_call_headers(resumed_request)
     snapshot = lineage.snapshot()
-
-    assert summary_headers["X-ACP-Lineage-Context-ID"] == source_context_id
-    assert summary_headers["X-ACP-Lineage-Compaction-ID"] == compaction.compaction_id
-    assert "X-ACP-Lineage-Previous-Context-ID" not in summary_headers
-    assert resumed_headers["X-ACP-Lineage-Context-ID"] == compaction.target_context_id
-    assert resumed_headers["X-ACP-Lineage-Previous-Context-ID"] == source_context_id
-    assert resumed_headers["X-ACP-Lineage-Transition"] == "compact"
-    assert resumed_headers["X-ACP-Lineage-Compaction-ID"] == compaction.compaction_id
-    assert snapshot["requests"][-1]["compaction_id"] == compaction.compaction_id
+    requests = {request["request_id"]: request for request in snapshot["requests"]}
+    contexts = {context["context_id"]: context for context in snapshot["contexts"]}
+
+    assert requests[summary_request] == {
+        "request_id": summary_request,
+        "session_id": "trace-1",
+        "context_id": source_context_id,
+        "kind": "compaction",
+        "compaction_id": compaction.compaction_id,
+    }
+    assert requests[resumed_request]["context_id"] == compaction.target_context_id
+    assert requests[resumed_request]["compaction_id"] == compaction.compaction_id
+    assert contexts[compaction.target_context_id]["previous_context_id"] == (
+        source_context_id
+    )
+    assert contexts[compaction.target_context_id]["transition"] == "compact"
     assert snapshot["compactions"] == [
         {
             "compaction_id": compaction.compaction_id,
             "session_id": "trace-1",
             "source_context_id": source_context_id,
             "target_context_id": compaction.target_context_id,
             "status": "completed",
-            "summary_request_id": summary_request.request_id,
+            "summary_request_id": summary_request,
         }
     ]
 
@@ -109,7 +124,7 @@ async def start_request():
     )
 
 
-def test_new_compaction_id_overrides_source_context_origin():
+def test_new_compaction_request_keeps_context_origin_in_manifest():
     lineage = LineageTracker()
     first_context_id = lineage.register_session(
         "trace-1", parent_session_id=None, depth=0
@@ -130,20 +145,16 @@ def test_new_compaction_id_overrides_source_context_origin():
         compaction_id=second_compaction.compaction_id,
     )
 
-    ordinary_headers = model_call_headers(ordinary_request)
-    summary_headers = model_call_headers(second_summary)
-    assert ordinary_headers["X-ACP-Lineage-Compaction-ID"] == (
-        first_compaction.compaction_id
-    )
-    assert (
-        summary_headers["X-ACP-Lineage-Context-ID"]
-        == first_compaction.target_context_id
-    )
-    assert summary_headers["X-ACP-Lineage-Previous-Context-ID"] == first_context_id
-    assert summary_headers["X-ACP-Lineage-Transition"] == "compact"
-    assert summary_headers["X-ACP-Lineage-Compaction-ID"] == (
-        second_compaction.compaction_id
-    )
+    snapshot = lineage.snapshot()
+    requests = {request["request_id"]: request for request in snapshot["requests"]}
+    contexts = {context["context_id"]: context for context in snapshot["contexts"]}
+
+    assert requests[ordinary_request]["compaction_id"] == first_compaction.compaction_id
+    assert requests[second_summary]["context_id"] == first_compaction.target_context_id
+    assert requests[second_summary]["compaction_id"] == second_compaction.compaction_id
+    compacted = contexts[first_compaction.target_context_id]
+    assert compacted["previous_context_id"] == first_context_id
+    assert compacted["transition"] == "compact"
 
 
 def test_failed_compaction_does_not_activate_target_context():
@@ -159,7 +170,8 @@ def test_failed_compaction_does_not_activate_target_context():
     next_request = lineage.start_request("trace-1", kind="turn")
     snapshot = lineage.snapshot()
 
-    assert next_request.context_id == source_context_id
+    assert snapshot["requests"][-1]["request_id"] == next_request
+    assert snapshot["requests"][-1]["context_id"] == source_context_id
     assert snapshot["compactions"] == [
         {
             "compaction_id": compaction.compaction_id,