diff --git a/src/xagent/core/agent/context/enrichment.py b/src/xagent/core/agent/context/enrichment.py index 441a43311f..4edc3dee1e 100644 --- a/src/xagent/core/agent/context/enrichment.py +++ b/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. diff --git a/src/xagent/core/agent/context/execution.py b/src/xagent/core/agent/context/execution.py index 994307f3dc..4a0b9e60ab 100644 --- a/src/xagent/core/agent/context/execution.py +++ b/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 = ( "" diff --git a/src/xagent/core/agent/pattern/dag/dag.py b/src/xagent/core/agent/pattern/dag/dag.py index 3b0315238c..8adf709d6a 100644 --- a/src/xagent/core/agent/pattern/dag/dag.py +++ b/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) diff --git a/tests/core/agent/test_request_provenance.py b/tests/core/agent/test_request_provenance.py new file mode 100644 index 0000000000..3d8da4b694 --- /dev/null +++ b/tests/core/agent/test_request_provenance.py @@ -0,0 +1,304 @@ +from __future__ import annotations + +import json +from typing import Any + +import pytest + +from xagent.core.agent.context import CompactConfig, ExecutionContext +from xagent.core.agent.context.enrichment import ( + TOP_LEVEL_USER_REQUEST_METADATA_KEY, + TopLevelUserRequest, + display_message_override, + top_level_user_request, +) +from xagent.core.agent.pattern.dag.dag import DAGPattern + +EXECUTION_REQUEST = "Summarize the email.\n[Connector context: contacto@example.es]" +CLEAN_REQUEST = "Summarize the email." + + +@pytest.mark.parametrize( + ("metadata", "language_text", "display_state"), + [ + pytest.param({}, EXECUTION_REQUEST, "missing", id="missing"), + pytest.param( + {"display_message": CLEAN_REQUEST}, CLEAN_REQUEST, "text", id="text" + ), + pytest.param({"display_message": ""}, "", "empty", id="blank"), + pytest.param({"display_message": " \n\t"}, "", "empty", id="whitespace"), + ], +) +def test_top_level_request_preserves_display_tri_state( + metadata: dict[str, Any], language_text: str, display_state: str +) -> None: + context = ExecutionContext(metadata={"task": EXECUTION_REQUEST}) + context.add_user_message(EXECUTION_REQUEST, metadata=metadata) + + request = top_level_user_request(context) + + assert request == TopLevelUserRequest( + execution_text=EXECUTION_REQUEST, + language_text=language_text, + display_state=display_state, + ) + assert context.metadata[TOP_LEVEL_USER_REQUEST_METADATA_KEY] == { + "execution_text": EXECUTION_REQUEST, + "language_text": language_text, + "display_state": display_state, + } + + +@pytest.mark.parametrize("value", [None, 17, [], {}]) +def test_direct_non_string_display_falls_back_to_execution_text(value: Any) -> None: + context = ExecutionContext() + context.add_user_message( + EXECUTION_REQUEST, + metadata={"display_message": value}, + ) + + request = top_level_user_request(context) + + assert display_message_override({"display_message": value}) is None + assert request.language_text == EXECUTION_REQUEST + assert request.display_state == "missing" + + +def test_new_independent_request_replaces_stored_snapshot() -> None: + context = ExecutionContext() + context.add_user_message( + EXECUTION_REQUEST, + metadata={"display_message": CLEAN_REQUEST}, + ) + assert top_level_user_request(context).language_text == CLEAN_REQUEST + + follow_up = "Switch to Spanish now." + context.add_user_message( + f"{follow_up}\n[Connector context in English]", + metadata={"display_message": follow_up}, + ) + + request = top_level_user_request(context) + assert request.execution_text == f"{follow_up}\n[Connector context in English]" + assert request.language_text == follow_up + assert request.display_state == "text" + + +@pytest.mark.parametrize( + ("display_message", "language_text", "display_state"), + [ + pytest.param(CLEAN_REQUEST, CLEAN_REQUEST, "text", id="text"), + pytest.param("", "", "empty", id="blank"), + pytest.param(" \n\t", "", "empty", id="whitespace"), + ], +) +@pytest.mark.parametrize("compaction", ["summary", "truncate"]) +@pytest.mark.parametrize("cold_restore", [False, True], ids=["live", "restored"]) +def test_request_provenance_survives_compaction_and_restore( + display_message: str, + language_text: str, + display_state: str, + compaction: str, + cold_restore: bool, +) -> None: + context = ExecutionContext(metadata={"task": EXECUTION_REQUEST}) + context.add_user_message( + EXECUTION_REQUEST, + metadata={"display_message": display_message}, + ) + context.add_user_message( + "DAG step instruction", + metadata={"dag_step_id": "draft", "kind": "dag_step_instruction"}, + ) + + if compaction == "summary": + assert context.compact_with_llm_response({"content": "Work summary"}).compacted + else: + context.compact_config = CompactConfig( + enabled=True, + threshold=1, + max_messages=1, + ) + assert context.compact_if_needed().compacted + if cold_restore: + context = ExecutionContext.from_dict(context.to_dict()) + + request = top_level_user_request(context) + assert request.execution_text == EXECUTION_REQUEST + assert request.language_text == language_text + assert request.display_state == display_state + + +def test_llm_compaction_request_persists_provenance_before_history_changes() -> None: + context = ExecutionContext( + metadata={"task": EXECUTION_REQUEST}, + compact_config=CompactConfig(enabled=True, threshold=1, max_messages=1), + ) + context.add_user_message( + EXECUTION_REQUEST, + metadata={"display_message": CLEAN_REQUEST}, + ) + + assert context.build_llm_compact_request_if_needed() is not None + context.messages = [] + + assert top_level_user_request(context) == TopLevelUserRequest( + execution_text=EXECUTION_REQUEST, + language_text=CLEAN_REQUEST, + display_state="text", + ) + + +def test_child_creation_snapshots_request_before_metadata_clone() -> None: + root = ExecutionContext(metadata={"task": EXECUTION_REQUEST}) + root.add_user_message( + EXECUTION_REQUEST, + metadata={"display_message": CLEAN_REQUEST}, + ) + + child = root.create_child_context(metadata={"dag_step_id": "draft"}) + child.messages = [] + + assert top_level_user_request(child).language_text == CLEAN_REQUEST + + +def test_provenance_roundtrip_is_checkpoint_compatible() -> None: + context = ExecutionContext(metadata={"task": EXECUTION_REQUEST}) + context.add_user_message( + EXECUTION_REQUEST, + metadata={"display_message": CLEAN_REQUEST}, + ) + expected = top_level_user_request(context) + + restored = ExecutionContext.from_dict(json.loads(json.dumps(context.to_dict()))) + + assert top_level_user_request(restored) == expected + + +@pytest.mark.parametrize( + ("display_message", "language_text", "display_state"), + [ + pytest.param(CLEAN_REQUEST, CLEAN_REQUEST, "text", id="text"), + pytest.param("", "", "empty", id="blank"), + pytest.param(" \n\t", "", "empty", id="whitespace"), + ], +) +def test_legacy_restored_child_hydrates_provenance_from_root( + display_message: str, + language_text: str, + display_state: str, +) -> None: + root = ExecutionContext(metadata={"task": EXECUTION_REQUEST}) + root.add_user_message( + EXECUTION_REQUEST, + metadata={"display_message": display_message}, + ) + root = ExecutionContext.from_dict(root.to_dict()) + child = ExecutionContext( + metadata={"task": EXECUTION_REQUEST, "dag_step_id": "draft"} + ) + child.add_user_message( + "DAG step instruction", + metadata={"dag_step_id": "draft", "kind": "dag_step_instruction"}, + ) + + DAGPattern._refresh_restored_step_runtime_metadata(child, root) + + request = top_level_user_request(child) + assert request.execution_text == EXECUTION_REQUEST + assert request.language_text == language_text + assert request.display_state == display_state + + +@pytest.mark.parametrize( + "snapshot", + [ + None, + {}, + {"execution_text": 1, "language_text": "x", "display_state": "text"}, + { + "execution_text": "x", + "language_text": "x", + "display_state": "unknown", + }, + ], +) +def test_invalid_legacy_child_snapshot_is_hydrated(snapshot: Any) -> None: + root = ExecutionContext() + root.add_user_message( + EXECUTION_REQUEST, metadata={"display_message": CLEAN_REQUEST} + ) + child = ExecutionContext( + metadata={ + "dag_step_id": "draft", + TOP_LEVEL_USER_REQUEST_METADATA_KEY: snapshot, + } + ) + + DAGPattern._refresh_restored_step_runtime_metadata(child, root) + + assert top_level_user_request(child).language_text == CLEAN_REQUEST + + +def test_valid_child_snapshot_wins_over_root_hydration() -> None: + root = ExecutionContext() + root.add_user_message( + EXECUTION_REQUEST, metadata={"display_message": CLEAN_REQUEST} + ) + child = ExecutionContext( + metadata={ + "dag_step_id": "draft", + TOP_LEVEL_USER_REQUEST_METADATA_KEY: { + "execution_text": "Child execution request", + "language_text": "Child language request", + "display_state": "text", + }, + } + ) + + DAGPattern._refresh_restored_step_runtime_metadata(child, root) + + assert top_level_user_request(child).language_text == "Child language request" + + +def test_legacy_waiting_child_uses_shared_restore_hydration_seam() -> None: + root = ExecutionContext(execution_id="waiting-root") + root.add_user_message( + EXECUTION_REQUEST, metadata={"display_message": CLEAN_REQUEST} + ) + child = root.create_child_context( + execution_id="waiting-child", metadata={"dag_step_id": "confirm"} + ) + child.messages = [] + child.metadata.pop(TOP_LEVEL_USER_REQUEST_METADATA_KEY) + pattern = DAGPattern(lambda **_: None) + pattern.status = "waiting_for_user" + pattern.active_step_id = "confirm" + pattern.active_step_ids = ["confirm"] + pattern.active_step_contexts = {"confirm": child.to_dict()} + pattern.active_step_pattern_states = { + "confirm": { + "status": "waiting_for_user", + "waiting_for_user_request": {"message": "Which date?"}, + } + } + pattern.planned_user_message_count = 1 + root.add_user_message("Friday") + + assert pattern._forward_user_response_to_waiting_step(root) + + restored_child = ExecutionContext.from_dict(pattern.active_step_contexts["confirm"]) + request = top_level_user_request(restored_child) + assert request.execution_text == EXECUTION_REQUEST + assert request.language_text == CLEAN_REQUEST + assert request.display_state == "text" + + +def test_persisting_provenance_does_not_change_rendered_prompt() -> None: + context = ExecutionContext() + context.add_user_message(EXECUTION_REQUEST, metadata={"display_message": ""}) + before = context.get_messages_for_llm() + + top_level_user_request(context) + + assert context.get_messages_for_llm() == before diff --git a/tests/core/agent/test_runner.py b/tests/core/agent/test_runner.py index c02b4fa84b..2a9301e293 100644 --- a/tests/core/agent/test_runner.py +++ b/tests/core/agent/test_runner.py @@ -1645,6 +1645,31 @@ async def test_runner_initial_user_message_preserves_display_metadata( assert user_event["data"]["turn_id"] == turn_id +@pytest.mark.parametrize( + ("request_context", "expected"), + [ + pytest.param({}, None, id="missing"), + pytest.param({"display_message": None}, "", id="null"), + pytest.param({"display_message": 17}, "", id="non-string"), + pytest.param({"display_message": ""}, "", id="blank"), + pytest.param({"display_message": " \n\t"}, " \n\t", id="whitespace"), + pytest.param({"display_message": "Read file"}, "Read file", id="text"), + ], +) +def test_runner_normalizes_initial_display_message_state( + request_context: dict[str, Any], expected: str | None +) -> None: + runner = AgentRunner(agent=Agent(name="writer", patterns=[])) + context = ExecutionContext(metadata={"request_context": request_context}) + + metadata = runner._initial_user_message_metadata(context) + + if expected is None: + assert "display_message" not in metadata + else: + assert metadata["display_message"] == expected + + @pytest.mark.asyncio async def test_runner_attaches_uploaded_image_refs_to_initial_user_message( tmp_path: Path,