fix(agent): preserve request provenance across checkpoints (toby)

src/xagent/core/agent/context/enrichment.py

@@ -2,7 +2,8 @@
 
 import asyncio
 import logging
-from typing import Any, cast
+from dataclasses import dataclass
+from typing import Any, Literal, cast
 
 from ...agent.trace import (
     trace_memory_retrieve_end,
@@ -20,6 +21,63 @@
 # with no image tools at all leaves it False.
 IMAGE_EDIT_UNAVAILABLE_METADATA_KEY = "image_edit_unavailable"
 
+DisplayMessageState = Literal["missing", "empty", "text"]
+TOP_LEVEL_USER_REQUEST_METADATA_KEY = "_xagent_top_level_user_request"
+
+
+@dataclass(frozen=True)
+class TopLevelUserRequest:
+    """Canonical execution and user-authored text for an independent request."""
+
+    execution_text: str
+    language_text: str
+    display_state: DisplayMessageState
+    has_pending_response: bool = False
+
+
+def _stored_top_level_user_request(context: Any) -> TopLevelUserRequest | None:
+    metadata = getattr(context, "metadata", None)
+    if not isinstance(metadata, dict):
+        return None
+    payload = metadata.get(TOP_LEVEL_USER_REQUEST_METADATA_KEY)
+    if not isinstance(payload, dict):
+        return None
+    execution_text = payload.get("execution_text")
+    language_text = payload.get("language_text")
+    display_state = payload.get("display_state")
+    if (
+        not isinstance(execution_text, str)
+        or not isinstance(language_text, str)
+        or display_state not in {"missing", "empty", "text"}
+    ):
+        return None
+    return TopLevelUserRequest(
+        execution_text=execution_text,
+        language_text=language_text,
+        display_state=display_state,
+    )
+
+
+def _persist_top_level_user_request(context: Any, request: TopLevelUserRequest) -> None:
+    metadata = getattr(context, "metadata", None)
+    if not isinstance(metadata, dict):
+        return
+    metadata[TOP_LEVEL_USER_REQUEST_METADATA_KEY] = {
+        "execution_text": request.execution_text,
+        "language_text": request.language_text,
+        "display_state": request.display_state,
+    }
+
+
+def hydrate_top_level_user_request(context: Any, root_context: Any) -> None:
+    """Backfill a legacy child snapshot from its canonical root request."""
+    if _stored_top_level_user_request(context) is not None:
+        return
+    root_request = _stored_top_level_user_request(root_context)
+    if root_request is None:
+        root_request = top_level_user_request(root_context)
+    _persist_top_level_user_request(context, root_request)
+
 
 async def enrich_context_with_memory(
     *,
@@ -110,6 +168,80 @@ def build_skill_context(skill: dict[str, Any]) -> str:
     return f"## Available Skill: {name}\n\n{content}".strip()
 
 
+def display_message_override(metadata: Any) -> str | None:
+    """Return a supported display-message override, including an empty one.
+
+    Missing keys and non-string values in directly constructed or restored
+    contexts keep the execution-content fallback. The runner normalizes a
+    present non-string value to an authoritative empty string at ingress.
+    """
+    if not isinstance(metadata, dict) or "display_message" not in metadata:
+        return None
+    display = metadata["display_message"]
+    if not isinstance(display, str):
+        return None
+    return display.strip()
+
+
+def top_level_user_request(context: Any) -> TopLevelUserRequest:
+    """Return and persist the latest independent top-level user request."""
+    has_pending_response = False
+    for message in reversed(getattr(context, "messages", []) or []):
+        if getattr(message, "role", None) != "user" or getattr(
+            message, "hidden", False
+        ):
+            continue
+        metadata = getattr(message, "metadata", None)
+        metadata = metadata if isinstance(metadata, dict) else {}
+        if metadata.get("response_to_waiting_for_user"):
+            has_pending_response = True
+            continue
+        if metadata.get("dag_step_id"):
+            continue
+
+        execution_text = str(getattr(message, "content", "") or "").strip()
+        display_text = display_message_override(metadata)
+        if display_text is None:
+            if not execution_text:
+                continue
+            request = TopLevelUserRequest(
+                execution_text=execution_text,
+                language_text=execution_text,
+                display_state="missing",
+                has_pending_response=has_pending_response,
+            )
+        else:
+            request = TopLevelUserRequest(
+                execution_text=execution_text,
+                language_text=display_text,
+                display_state="text" if display_text else "empty",
+                has_pending_response=has_pending_response,
+            )
+        _persist_top_level_user_request(context, request)
+        return request
+
+    stored = _stored_top_level_user_request(context)
+    if stored is not None:
+        return TopLevelUserRequest(
+            execution_text=stored.execution_text,
+            language_text=stored.language_text,
+            display_state=stored.display_state,
+            has_pending_response=has_pending_response,
+        )
+
+    metadata = getattr(context, "metadata", None)
+    task = metadata.get("task") if isinstance(metadata, dict) else None
+    task_text = str(task or "").strip()
+    request = TopLevelUserRequest(
+        execution_text=task_text,
+        language_text=task_text,
+        display_state="missing",
+        has_pending_response=has_pending_response,
+    )
+    _persist_top_level_user_request(context, request)
+    return request
+
+
 def latest_user_text(context: Any, *, prefer_display: bool = False) -> str:
     """Return the latest user turn's text.
 

src/xagent/core/agent/context/execution.py

@@ -42,6 +42,7 @@
     IMAGE_EDIT_UNAVAILABLE_METADATA_KEY,
     MEMORY_CONTEXT_METADATA_KEY,
     SKILL_CONTEXT_METADATA_KEY,
+    top_level_user_request,
 )
 from .memory_tool import MEMORY_TOOLS_METADATA_KEY
 from .message import LLMCallRecord, Message
@@ -932,6 +933,9 @@ def create_child_context(
         include_system_prompt: bool = True,
         metadata: dict[str, Any] | None = None,
     ) -> "ExecutionContext":
+        # Child compaction may discard the copied root message, so snapshot its
+        # request provenance before metadata is cloned.
+        top_level_user_request(self)
         child_metadata = dict(self.metadata)
         if metadata:
             child_metadata.update(metadata)
@@ -1113,6 +1117,7 @@ def compact_if_needed(self) -> CompactResult:
                 strategy="none",
             )
 
+        top_level_user_request(self)
         total_tokens = self._get_total_tokens()
         if total_tokens > self.compact_config.threshold:
             result = self._drop_oldest_messages()
@@ -1129,6 +1134,7 @@ def build_llm_compact_request_if_needed(self) -> dict[str, Any] | None:
         if not self.compact_config.enabled:
             return None
 
+        top_level_user_request(self)
         total_tokens = self._get_total_tokens()
         if total_tokens <= self.compact_config.threshold:
             return None
@@ -1205,6 +1211,7 @@ def compact_with_llm_response(
         llm: Any = None,
         original_tokens: int | None = None,
     ) -> CompactResult:
+        top_level_user_request(self)
         original_count = len(self.messages)
         summary = (
             ""

src/xagent/core/agent/pattern/dag/dag.py

@@ -14,6 +14,7 @@
 )
 from ...context.enrichment import (
     enrich_context_with_memory,
+    hydrate_top_level_user_request,
     latest_user_text,
 )
 from ...frame import ExecutionFrame, ExecutionSnapshot, ExecutionStatus
@@ -1929,6 +1930,8 @@ def _refresh_restored_step_runtime_metadata(
     ) -> None:
         """Refresh volatile routing metadata on a checkpoint-restored DAG step."""
 
+        hydrate_top_level_user_request(child_context, root_context)
+
         root_metadata = getattr(root_context, "metadata", {})
         preferred_modalities = normalize_input_modalities(
             root_metadata.get(PREFERRED_INPUT_MODALITIES_METADATA_KEY)