fix(agent): preserve request provenance across checkpoints (toby) - #2068
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
rogercloud
left a comment
There was a problem hiding this comment.
Summary
This PR introduces a canonical "top-level user request" snapshot (TopLevelUserRequest) in src/xagent/core/agent/context/enrichment.py, persisted into context.metadata at four call sites in execution.py (before child-context creation and before each compaction/summarization/truncation boundary), and backfilled into legacy checkpoint-restored DAG children via a shared hydration seam in dag.py covering both the active-step and waiting-step restore paths. It is explicitly a behavior-neutral foundation layer for a future PR that will build a response-language policy on top of this data — nothing in the current codebase consumes the new snapshot yet, so this PR changes no user-visible behavior. The mechanism and hydration seam are correct and well-tested end-to-end (context.metadata including the new key survives the full production checkpoint path unfiltered: runtime.py → checkpoint.py → JSONB column), and it correctly closes the bulk of linked issue #2065's acceptance criteria, with the specific gaps noted below.
Blocking: no — recommended event: APPROVE
Round 0 design verdict: acceptable-with-reservations
The core design choice to store this snapshot in context.metadata rather than as a first-class dataclass field is correct given the project's "no schema change" constraint — to_dict/from_dict hand-enumerate fields with no versioning, so a first-class field would cost roughly five touch points versus zero for metadata. The main reservation is duplication: this PR adds a fourth-ish parallel "get the current/top-level user request" implementation into an area that already has two others doing similar work — ExecutionContext._current_user_request_text (pre-existing, untouched) and latest_user_text (pre-existing) — and the new implementation already produces different results than the one that actually feeds the LLM prompt today, for the same input (Finding 1). This is currently harmless because nothing reads the new snapshot yet, but it is exactly the kind of latent inconsistency that becomes a real bug the moment a consumer is wired in.
Findings
1. [MAJOR] Three overlapping "current/top-level user request" scanners already disagree on the same input
File: src/xagent/core/agent/context/enrichment.py:186-242 (new top_level_user_request), vs. src/xagent/core/agent/context/execution.py:542-565 (_current_user_request_text, pre-existing, untouched), vs. src/xagent/core/agent/context/enrichment.py:245-265 (latest_user_text, pre-existing)
Blocking: no
For a context with an earlier real user message followed by a later message whose content is blank/whitespace but carries a display_message metadata override, _current_user_request_text() — the only one of the three currently consumed for actual LLM prompt/system-context building (execution.py:571,652) — skips the blank message and returns the earlier real one. The new top_level_user_request() instead accepts the blank-content message immediately because a display override is present, yielding execution_text="". This is inert today (nothing reads the new function's output for prompt-building yet), but the entire purpose of persisting this snapshot is that something will consume it soon, and it will disagree with what is actually rendered to the LLM. Recommend consolidating to one canonical scanner — have top_level_user_request/latest_user_text delegate to _current_user_request_text's filtering logic (or vice versa) rather than maintaining three independent implementations.
2. [MINOR] Legacy-child hydration on the active-restore path can back-fill a stale root snapshot
File: src/xagent/core/agent/context/enrichment.py:72-79 (hydrate_top_level_user_request), called from src/xagent/core/agent/pattern/dag/dag.py:1927-1935 (_refresh_restored_step_runtime_metadata), active-restore branch src/xagent/core/agent/pattern/dag/dag.py:~951-967
Blocking: no
On the active-restore path, hydration prefers the root context's already-stored snapshot over live recomputation from the root's current messages. If the root was restored from a checkpoint carrying an old snapshot and then received a genuinely new independent user message before hydration runs, the legacy child is hydrated with the stale old request rather than the new one — nothing in the active-restore call chain refreshes the root's own snapshot first. This partially contradicts issue #2065's "a new independent request replaces the prior snapshot" criterion for this specific ordering; the direct-call case is correctly covered by test_new_independent_request_replaces_stored_snapshot, but none of the 34 existing tests build a root with both a stored snapshot and a newer message before hydration. Recommend preferring live recomputation (top_level_user_request(root_context)) over a possibly-stale stored snapshot in hydrate_top_level_user_request, as a fast-follow before any consumer lands.
Note: the equivalent concern was checked on the waiting-restore path and does not apply — that path intentionally preserves the original top-level task's provenance rather than overwriting it with a follow-up answer, and is correctly tested (test_request_provenance.py:264-294).
3. [MINOR] Stored-snapshot validation accepts a well-typed but empty payload as "valid," blocking real backfill
File: src/xagent/core/agent/context/enrichment.py:38-53 (_stored_top_level_user_request)
Blocking: no
Validation only checks that execution_text/language_text are strings and display_state is a valid enum value — it never checks for meaningful non-emptiness. A context with no qualifying user messages and no metadata["task"] fallback (reachable in production: an empty/whitespace task submitted with file attachments in DAG/"think" mode skips both the initial user-message add and the task fallback) causes top_level_user_request()'s fallback branch to persist {"execution_text": "", "language_text": "", "display_state": "missing"}, which passes validation as "valid." When this context later serves as a legacy-restore root, hydrate_top_level_user_request's early-return fires on this empty snapshot, so a real root request is never backfilled into a child that needs it. test_invalid_legacy_child_snapshot_is_hydrated (test_request_provenance.py:213-240) covers None/{}/wrong-type/bad-enum but not this valid-but-empty case. Recommend treating an all-empty stored payload as absent, or adding an emptiness check to _stored_top_level_user_request.
4. [MINOR] has_pending_response is dead and always resolves to a stale value after round-trip
File: src/xagent/core/agent/context/enrichment.py:35 (TopLevelUserRequest.has_pending_response)
Blocking: no
The field is never written by _persist_top_level_user_request (which persists only 3 of the 4 fields) and never read back by _stored_top_level_user_request, so it is always False after any checkpoint round-trip regardless of its value at persist time. It has zero consumers anywhere in src/ or tests/ outside its own definition and the constructor calls inside top_level_user_request itself, and in the "stored"/"fallback" reconstruction branches it is recomputed fresh from live messages while the other three fields come from the persisted snapshot — producing an object with mixed provenance. Safe to fix now since there are zero callers: either wire it into persist/restore properly, or drop the field until a consumer needs it.
5. [MINOR] Waiting-for-user answer stamping only touches the DAG child, and this PR can now durably lock in that gap
File: src/xagent/core/agent/pattern/react/react.py:1746 (response_to_waiting_for_user stamping) — not modified by this PR; noted here since this PR's persistence changes the impact of the gap
Blocking: no
This stamping only ever marks the context ReAct is directly running on (the DAG step's child context), never the DAG root's own copy of a waiting-answer message. This blind spot is pre-existing and shared with _current_user_request_text, but that function recomputes fresh on every call and self-corrects. This PR's new persistence path (_persist_top_level_user_request + _stored_top_level_user_request/hydrate_top_level_user_request) can instead durably lock a one-time misclassification (a waiting-for-user answer on the root looking like a brand-new independent top-level request) into checkpoint metadata, where it will keep being served on later restores instead of self-correcting. Recommend stamping the root's own copy when forwarding a waiting-step answer, or having top_level_user_request recognize the forwarded_from_root/dag_step_id forwarding metadata as a signal.
6. [MINOR] top_level_user_request reads as a pure accessor but mutates metadata, and the naming convention breaks
File: src/xagent/core/agent/context/enrichment.py:186 (top_level_user_request); call sites src/xagent/core/agent/context/execution.py:938,1120,1137,1214
Blocking: no
top_level_user_request(context) reads as a pure accessor by name and docstring ("Return and persist...") but mutates context.metadata on 2 of 3 return paths, and is called as a bare, return-value-discarded statement at four hot-path call sites — only one of those four has a comment explaining the mutation intent. Separately, hydrate_top_level_user_request (called from dag.py:1933 inside _refresh_restored_step_runtime_metadata, whose docstring says it refreshes the checkpoint-restored child step) can, via its fallback branch, mutate the root context instead — surprising given the enclosing method's stated scope. This is the sole exception to an otherwise consistent convention in this area, where every other metadata-writing function is verb-prefixed (_persist_top_level_user_request, enrich_context_with_memory, _attach_decision_metadata, _persist_injected_context) and every other get_*/noun-named accessor is pure. Recommend renaming to something like snapshot_top_level_user_request, or splitting into a pure read plus an explicit persist call.
7. [MINOR] Checkpoint-roundtrip test doesn't actually exercise the persisted snapshot
File: tests/core/agent/test_request_provenance.py:165-175 (test_provenance_roundtrip_is_checkpoint_compatible)
Blocking: no
This test round-trips a context through to_dict()/from_dict() but never clears restored.messages, unlike two nearby tests (lines 143 and 160) that correctly do so to isolate the persisted-snapshot path. Because top_level_user_request() checks live .messages for a qualifying message before ever consulting the persisted metadata snapshot, this test recomputes the identical answer from the still-present original message and would pass unchanged even if metadata persistence were deleted entirely — confirmed by corrupting/removing the persisted key from the round-tripped data and observing the assertion still passes. This undercuts the PR's stated verification claim that "JSON checkpoint round trips" are exercised by this test. Recommend a one-line fix: restored.messages = [] before the assertion, mirroring the correct pattern already used in the same file.
restored = ExecutionContext.from_dict(context.to_dict())
restored.messages = [] # isolate the persisted-metadata path, as done at line 143/1608. [MINOR] No allowlist on caller-supplied request_context keys copied into metadata
File: src/xagent/core/agent/runner.py:937-955 (_apply_request_context) — not modified by this PR; pre-existing behavior noted for context
Blocking: no
This copies arbitrary caller-supplied request_context keys verbatim into context.metadata (intentional per the comment at execution.py:107-108), with no allowlist and no origin/provenance marking — so a caller could set request_context["_xagent_top_level_user_request"] to forge an arbitrary snapshot. The _xagent_-prefix-stripping protection that exists for outbound chat message dicts (model/chat/basic/base.py:132-148) does not apply here, since ExecutionContext.metadata is a separate dict. However, request_context is always caller-owned data for that caller's own task execution — there is no cross-user/cross-tenant propagation path, and no security-sensitive consumer (auth, billing, tool-gating) currently reads this field. Flagging as a defense-in-depth note only, not an urgent fix.
9. [MINOR] "missing" display state not exercised through the full compaction/restore matrix
File: tests/core/agent/test_request_provenance.py:88-102
Blocking: no
The parametrized compaction/serialization/cold-restore matrix covers only text/blank/whitespace display states. The missing state is only exercised by a separate, narrower test (lines 19-40/47-64) that never drives it through compaction or cold-restore, even though issue #2065 explicitly lists "missing" as one of the four states required to survive exactly those paths. Recommend adding missing to the parametrized matrix for full coverage.
10. [MINOR] display_message: None and a missing key produce the same tri-state, but a different code path normalizes differently
File: src/xagent/core/agent/context/enrichment.py:171-183 (display_message_override); contrast src/xagent/core/agent/runner.py:902-908 (unmodified by this PR)
Blocking: no
display_message_override treats an explicit display_message: None the same as a missing key (both → display_state="missing"), but runner.py normalizes any present non-string value to "" before ingress (→ display_state="empty") — so the same conceptual "no override" input yields a different tri-state depending on which code path constructed the message. This split is even documented in the function's own docstring. Reachable for DAG-child-constructed messages and any pre-normalization checkpoint data bypassing the runner's ingress path. Worth a short note in the docstring or a follow-up to align the two normalizations.
11. [MINOR] Blank content with a display override skips the "look for earlier real content" scan
File: src/xagent/core/agent/context/enrichment.py:171-183 (display_message_override)
Blocking: no
When no display_message override is present, a blank/whitespace-only message is skipped and the scan continues backward for an earlier real message. But when a display_message override is present — even with blank execution_text — the message is accepted immediately without checking for an earlier real-content message, yielding an empty execution_text alongside a real language_text (confirmed: content=" " + display_message="Hola" → execution_text='', language_text='Hola', display_state='text'). This is plausibly intentional (an explicit override is a stronger signal than searching for prior content), but it's undocumented and untested. Recommend a docstring note, or a small test case, clarifying the intended behavior.
Simplification opportunities
L35: yagni — TopLevelUserRequest.has_pending_response has zero consumers outside its own definition and 4 constructor calls in the same function; drop it or wire it into persist/restore (overlaps Finding 4, don't fix twice).L51: the hardcoded enum-value set in _stored_top_level_user_request duplicates the DisplayMessageState Literal; consider typing.get_args(DisplayMessageState) instead of restating the values.- Net: dropping
has_pending_responseentirely would save roughly 8-12 lines; adataclasses.asdict-based pack/unpack simplification was investigated and rejected —has_pending_responseis deliberately excluded from persistence (derived from live state), soasdict()would need an explicit exclusion step anyway, netting no real improvement once per-field validation is accounted for.
Blocking status & recommended decision
Blocking: no. No confirmed finding produces an incorrect user-visible result today — this is an inert foundation layer with no production consumer of the new data yet; all findings above are pre-merge-quality or fast-follow recommendations, not merge blockers.
Recommended event: APPROVE
Summary
This pull request is the behavior-neutral foundation layer from #2062. It does not activate request-language guidance or change any existing language-policy prompt; policy representation and consumer activation will land in later ordered layers.
Verification
Closes #2065
Parent: #2062
Aggregate reference: #1990
#2041 remains separate and out of scope.