diff --git a/agent/memory_manager.py b/agent/memory_manager.py index 7eda64fba4dd4..a45a6f6a29481 100644 --- a/agent/memory_manager.py +++ b/agent/memory_manager.py @@ -378,7 +378,8 @@ def handle_tool_call( def on_turn_start(self, turn_number: int, message: str, **kwargs) -> None: """Notify all providers of a new turn. - kwargs may include: remaining_tokens, model, platform, tool_count. + kwargs may include: remaining_tokens, model, platform, tool_count, + user_id, user_name, and session_title. """ for provider in self._providers: try: diff --git a/agent/memory_provider.py b/agent/memory_provider.py index c9abc48c7a92e..771ba706cfb33 100644 --- a/agent/memory_provider.py +++ b/agent/memory_provider.py @@ -144,9 +144,10 @@ def shutdown(self) -> None: def on_turn_start(self, turn_number: int, message: str, **kwargs) -> None: """Called at the start of each turn with the user message. - Use for turn-counting, scope management, periodic maintenance. + Use for turn-counting, scope management, attribution, and periodic maintenance. - kwargs may include: remaining_tokens, model, platform, tool_count. + kwargs may include: remaining_tokens, model, platform, tool_count, + user_id, user_name, and session_title. Providers use what they need; extras are ignored. """ diff --git a/gateway/run.py b/gateway/run.py index 46c508e4bde00..785a76d051b54 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -10302,6 +10302,8 @@ def run_sync(): return agent.run_conversation( user_message=prompt, task_id=task_id, + turn_user_id=source.user_id, + turn_user_name=source.user_name, ) finally: self._cleanup_agent_resources(agent) @@ -15327,7 +15329,13 @@ def _approval_notify_sync(approval_data: dict) -> None: else: _run_message = message - result = agent.run_conversation(_run_message, conversation_history=agent_history, task_id=session_id) + result = agent.run_conversation( + _run_message, + conversation_history=agent_history, + task_id=session_id, + turn_user_id=source.user_id, + turn_user_name=source.user_name, + ) finally: unregister_gateway_notify(_approval_session_key) # Cancel any pending clarify entries so blocked agent diff --git a/run_agent.py b/run_agent.py index f2f3379e0d78e..8b3922bc828f8 100644 --- a/run_agent.py +++ b/run_agent.py @@ -5658,6 +5658,35 @@ def _sync_external_memory_for_turn( except Exception: pass + def _build_memory_turn_context( + self, + *, + turn_user_id: Optional[str] = None, + turn_user_name: Optional[str] = None, + ) -> Dict[str, str]: + """Collect per-turn attribution context for external memory providers.""" + context: Dict[str, str] = {} + + if self.platform: + context["platform"] = str(self.platform) + + user_id = turn_user_id or self._user_id + if user_id: + context["user_id"] = str(user_id) + + if turn_user_name: + context["user_name"] = str(turn_user_name) + + if self._session_db and self.session_id: + try: + session_title = self._session_db.get_session_title(self.session_id) + except Exception: + session_title = None + if session_title: + context["session_title"] = str(session_title) + + return context + def release_clients(self) -> None: """Release LLM client resources WITHOUT tearing down session tool state. @@ -11690,6 +11719,8 @@ def run_conversation( task_id: str = None, stream_callback: Optional[callable] = None, persist_user_message: Optional[str] = None, + turn_user_id: Optional[str] = None, + turn_user_name: Optional[str] = None, ) -> Dict[str, Any]: """ Run a complete conversation with tool calling until completion. @@ -11705,7 +11736,12 @@ def run_conversation( persist_user_message: Optional clean user message to store in transcripts/history when user_message contains API-only synthetic prefixes. - or queuing follow-up prefetch work. + or queuing follow-up prefetch work. + turn_user_id: Optional per-message user identifier supplied by the + host. This can differ from the cached agent user in shared + thread sessions. + turn_user_name: Optional per-message display name supplied by the + host for memory-provider attribution in shared sessions. Returns: Dict: Complete conversation result with final response and message history @@ -12098,9 +12134,13 @@ def run_conversation( if self._memory_manager: try: _turn_msg = original_user_message if isinstance(original_user_message, str) else "" - self._memory_manager.on_turn_start(self._user_turn_count, _turn_msg) - except Exception: - pass + _turn_context = self._build_memory_turn_context( + turn_user_id=turn_user_id, + turn_user_name=turn_user_name, + ) + self._memory_manager.on_turn_start(self._user_turn_count, _turn_msg, **_turn_context) + except Exception as e: + logger.debug("External memory provider turn-start hook failed: %s", e) # External memory provider: prefetch once before the tool loop. # Reuse the cached result on every iteration to avoid re-calling diff --git a/tests/agent/test_memory_provider.py b/tests/agent/test_memory_provider.py index ca39da70f0819..ea232c00efd2c 100644 --- a/tests/agent/test_memory_provider.py +++ b/tests/agent/test_memory_provider.py @@ -64,8 +64,8 @@ def handle_tool_call(self, tool_name, args, **kwargs): def shutdown(self): self.shutdown_called = True - def on_turn_start(self, turn_number, message): - self.turn_starts.append((turn_number, message)) + def on_turn_start(self, turn_number, message, **kwargs): + self.turn_starts.append((turn_number, message, kwargs)) def on_session_end(self, messages): self.session_end_called = True @@ -310,7 +310,30 @@ def test_on_turn_start(self): p = FakeMemoryProvider("p") mgr.add_provider(p) mgr.on_turn_start(3, "hello") - assert p.turn_starts == [(3, "hello")] + assert p.turn_starts == [(3, "hello", {})] + + def test_on_turn_start_passes_turn_context(self): + mgr = MemoryManager() + p = FakeMemoryProvider("p") + mgr.add_provider(p) + mgr.on_turn_start( + 4, + "hello", + session_title="Research Thread", + user_id="u-42", + user_name="Alice", + ) + assert p.turn_starts == [ + ( + 4, + "hello", + { + "session_title": "Research Thread", + "user_id": "u-42", + "user_name": "Alice", + }, + ) + ] def test_on_session_end(self): mgr = MemoryManager() diff --git a/tests/run_agent/test_run_agent.py b/tests/run_agent/test_run_agent.py index dadb7b31ccee8..2693d2ab98427 100644 --- a/tests/run_agent/test_run_agent.py +++ b/tests/run_agent/test_run_agent.py @@ -13,7 +13,7 @@ from logging.handlers import RotatingFileHandler from pathlib import Path from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest from agent.codex_responses_adapter import _chat_messages_to_responses_input, _normalize_codex_response, _preflight_codex_input_items @@ -2476,6 +2476,68 @@ def test_stop_finish_reason_returns_response(self, agent): assert result["final_response"] == "Final answer" assert result["completed"] is True + def test_turn_context_reaches_memory_manager_before_prefetch(self, agent): + self._setup_agent(agent) + resp = _mock_response(content="Final answer", finish_reason="stop") + agent.client.chat.completions.create.return_value = resp + agent.platform = "telegram" + agent.session_id = "sess-123" + agent._user_id = "stale-user" + agent._memory_manager = MagicMock() + agent._memory_manager.prefetch_all.return_value = "" + agent._session_db = MagicMock() + agent._session_db.get_session_title.return_value = "Thread Research" + + with ( + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + ): + result = agent.run_conversation( + "hello", + turn_user_id="u-42", + turn_user_name="Alice", + ) + + assert result["final_response"] == "Final answer" + relevant_calls = [ + c for c in agent._memory_manager.mock_calls + if c[0] in {"on_turn_start", "prefetch_all"} + ] + assert relevant_calls[:2] == [ + call.on_turn_start( + 1, + "hello", + platform="telegram", + user_id="u-42", + user_name="Alice", + session_title="Thread Research", + ), + call.prefetch_all("hello"), + ] + + def test_empty_turn_user_id_falls_back_to_cached_user(self, agent): + self._setup_agent(agent) + resp = _mock_response(content="Final answer", finish_reason="stop") + agent.client.chat.completions.create.return_value = resp + agent._user_id = "cached-user" + agent._memory_manager = MagicMock() + agent._memory_manager.prefetch_all.return_value = "" + + with ( + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + ): + result = agent.run_conversation("hello", turn_user_id="") + + assert result["final_response"] == "Final answer" + agent._memory_manager.on_turn_start.assert_any_call( + 1, + "hello", + user_id="cached-user", + ) + def test_tool_calls_then_stop(self, agent): self._setup_agent(agent) tc = _mock_tool_call(name="web_search", arguments="{}", call_id="c1")