diff --git a/api/streaming.py b/api/streaming.py index d0144bf4..38ed5ad6 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -4560,7 +4560,7 @@ def _generate_llm_session_title_for_agent(agent, user_text: str, assistant_text: return None, 'llm_invalid', str(raw)[:120] -def _generate_llm_session_title_via_aux(user_text: str, assistant_text: str, agent=None, *, use_agent_model: bool = False) -> tuple[Optional[str], str, str]: +def _generate_llm_session_title_via_aux(user_text: str, assistant_text: str, agent=None, *, use_agent_model: bool = False, conversation_id: str = '') -> tuple[Optional[str], str, str]: """Generate a title via dedicated auxiliary LLM route, then sanitize/validate result. When use_agent_model is False (default), the auxiliary client resolves @@ -4568,6 +4568,11 @@ def _generate_llm_session_title_via_aux(user_text: str, assistant_text: str, age prevents the session's chat model (e.g. a Chinese model) from overriding the dedicated title model. When True, the agent's attrs are passed through (legacy fallback behaviour). + + conversation_id republishes the webui session id as the Agent's ambient + conversation context for the duration of the aux call, so OpenCode relay + targets receive the same ``x-opencode-session`` sticky key as the session's + main turns (#7470). The context is reset in all exit paths. """ if use_agent_model and agent: provider = getattr(agent, 'provider', '') @@ -4577,13 +4582,30 @@ def _generate_llm_session_title_via_aux(user_text: str, assistant_text: str, age provider = '' model = '' base_url = '' - raw, status = generate_title_raw_via_aux( - user_text, - assistant_text, - provider=provider, - model=model, - base_url=base_url, - ) + ctx_token = None + if conversation_id: + try: + from agent.portal_tags import set_conversation_context + ctx_token = set_conversation_context(str(conversation_id)) + except Exception: + # Older/absent agent runtime: proceed without conversation context + # (previous behaviour) rather than failing title generation. + ctx_token = None + try: + raw, status = generate_title_raw_via_aux( + user_text, + assistant_text, + provider=provider, + model=model, + base_url=base_url, + ) + finally: + if ctx_token is not None: + try: + from agent.portal_tags import reset_conversation_context + reset_conversation_context(ctx_token) + except Exception: + pass if not raw: return None, status, '' title = _sanitize_generated_title(raw) @@ -4733,9 +4755,9 @@ def _run_background_title_update(session_id: str, user_text: str, assistant_text if agent and not aux_title_configured: next_title, llm_status, raw_preview = _generate_llm_session_title_for_agent(agent, user_text, assistant_text) if not next_title and llm_status in ('llm_error', 'llm_invalid'): - next_title, llm_status, raw_preview = _generate_llm_session_title_via_aux(user_text, assistant_text, agent=agent, use_agent_model=True) + next_title, llm_status, raw_preview = _generate_llm_session_title_via_aux(user_text, assistant_text, agent=agent, use_agent_model=True, conversation_id=session_id) else: - next_title, llm_status, raw_preview = _generate_llm_session_title_via_aux(user_text, assistant_text) + next_title, llm_status, raw_preview = _generate_llm_session_title_via_aux(user_text, assistant_text, conversation_id=session_id) if not next_title and agent and llm_status in ('llm_error_aux', 'llm_invalid_aux'): next_title, llm_status, raw_preview = _generate_llm_session_title_for_agent(agent, user_text, assistant_text) source = llm_status @@ -4831,9 +4853,9 @@ def _run_background_title_refresh(session_id: str, user_text: str, assistant_tex if agent and not aux_title_configured: next_title, llm_status, raw_preview = _generate_llm_session_title_for_agent(agent, user_text, assistant_text) if not next_title and llm_status in ('llm_error', 'llm_invalid'): - next_title, llm_status, raw_preview = _generate_llm_session_title_via_aux(user_text, assistant_text, agent=agent, use_agent_model=True) + next_title, llm_status, raw_preview = _generate_llm_session_title_via_aux(user_text, assistant_text, agent=agent, use_agent_model=True, conversation_id=session_id) else: - next_title, llm_status, raw_preview = _generate_llm_session_title_via_aux(user_text, assistant_text) + next_title, llm_status, raw_preview = _generate_llm_session_title_via_aux(user_text, assistant_text, conversation_id=session_id) if not next_title and agent and llm_status in ('llm_error_aux', 'llm_invalid_aux'): next_title, llm_status, raw_preview = _generate_llm_session_title_for_agent(agent, user_text, assistant_text) if not next_title: @@ -4896,7 +4918,7 @@ def generate_session_title_for_session(session, *, prefer_latest: bool = False, with profiles_api.profile_env_for_background_worker(session, "manual title regeneration", logger_override=logger): if not _aux_title_generation_enabled(): return None, 'title_generation_disabled', '' - next_title, llm_status, raw_preview = _generate_llm_session_title_via_aux(user_text, assistant_text, agent=agent) + next_title, llm_status, raw_preview = _generate_llm_session_title_via_aux(user_text, assistant_text, agent=agent, conversation_id=getattr(session, 'session_id', '') or '') if next_title: return next_title, llm_status, raw_preview fallback_title = _fallback_title_from_exchange(user_text, assistant_text) diff --git a/tests/_aux_client_helpers.py b/tests/_aux_client_helpers.py index 68cd493d..8dfd7f01 100644 --- a/tests/_aux_client_helpers.py +++ b/tests/_aux_client_helpers.py @@ -33,13 +33,45 @@ def auxiliary_client_modules(): The stubs are fresh objects owned by this helper, so the ``auxiliary_client`` attribute is set on our own module and no pre-existing ``agent`` module is mutated. Restoration is therefore entirely ``patch.dict``'s job. + + ``agent.portal_tags`` is stubbed with a real ``ContextVar`` behind the same + set/reset/get surface as the Agent runtime (tokens are truthy objects, and + reset restores the previous value), so title-generation tests can assert + conversation-context publication/restoration without a real Agent runtime + (#7470). """ + import contextvars as _contextvars + agent_stub = types.ModuleType("agent") auxiliary_client_stub = types.ModuleType("agent.auxiliary_client") + portal_tags_stub = types.ModuleType("agent.portal_tags") agent_stub.auxiliary_client = auxiliary_client_stub + portal_state = {"set_calls": [], "reset_calls": []} + _conversation_cv = _contextvars.ContextVar("webui_test_conversation_id", default=None) + + def _set_conversation_context(conversation_id): + portal_state["set_calls"].append(conversation_id) + return _conversation_cv.set(conversation_id or None) + + def _reset_conversation_context(token): + portal_state["reset_calls"].append(token) + _conversation_cv.reset(token) + + def _get_conversation_context(): + return _conversation_cv.get() + + portal_tags_stub.set_conversation_context = _set_conversation_context + portal_tags_stub.reset_conversation_context = _reset_conversation_context + portal_tags_stub.get_conversation_context = _get_conversation_context + portal_tags_stub._portal_state = portal_state + agent_stub.portal_tags = portal_tags_stub return mock.patch.dict( sys.modules, - {"agent": agent_stub, "agent.auxiliary_client": auxiliary_client_stub}, + { + "agent": agent_stub, + "agent.auxiliary_client": auxiliary_client_stub, + "agent.portal_tags": portal_tags_stub, + }, ) diff --git a/tests/test_title_aux_routing.py b/tests/test_title_aux_routing.py index 5b643b79..e5654da9 100644 --- a/tests/test_title_aux_routing.py +++ b/tests/test_title_aux_routing.py @@ -1082,5 +1082,196 @@ class TestAuxInvalidAuxTriggersAgentFallback(unittest.TestCase): mock_agent_title.assert_not_called() +class TestAuxTitleConversationContext(unittest.TestCase): + """#7470: aux title-generation calls republish the session id as the Agent's + ambient conversation context so OpenCode relay targets get the same + x-opencode-session sticky key as the session's main turns.""" + + MOCK_RESP = types.SimpleNamespace( + choices=[ + types.SimpleNamespace( + message=types.SimpleNamespace(content='Weather Title'), + finish_reason='stop', + ) + ] + ) + + def _portal_state(self): + from agent.portal_tags import _portal_state + return _portal_state + + def _current_context(self): + from agent.portal_tags import get_conversation_context + return get_conversation_context() + + def test_direct_call_publishes_and_restores_context_after_success(self): + from api.streaming import _generate_llm_session_title_via_aux + + seen = {} + + def fake_call_llm(**kwargs): + from agent.portal_tags import get_conversation_context + seen['context_during_call'] = get_conversation_context() + return self.MOCK_RESP + + with _patch_tg_config({'provider': '', 'model': 'gpt-4o', 'base_url': ''}): + with patch('agent.auxiliary_client.call_llm', side_effect=fake_call_llm, create=True): + title, status, _raw = _generate_llm_session_title_via_aux( + 'What is the weather?', 'It is sunny.', conversation_id='sid-123', + ) + + self.assertEqual(title, 'Weather Title') + self.assertEqual(status, 'llm_aux') + # Context published before the aux call and visible inside call_llm. + self.assertEqual(seen['context_during_call'], 'sid-123') + state = self._portal_state() + self.assertEqual(state['set_calls'], ['sid-123']) + # Reset restored the pre-call value (None) in the finally block. + self.assertIsNone(self._current_context()) + self.assertEqual(len(state['reset_calls']), 1) + + def test_context_restored_after_exception_escapes_aux_call(self): + from api.streaming import _generate_llm_session_title_via_aux + + with _patch_tg_config({'provider': '', 'model': 'gpt-4o', 'base_url': ''}): + with patch('api.streaming.generate_title_raw_via_aux', side_effect=RuntimeError('boom')): + with self.assertRaises(RuntimeError): + _generate_llm_session_title_via_aux( + 'What is the weather?', 'It is sunny.', conversation_id='sid-456', + ) + state = self._portal_state() + self.assertEqual(state['set_calls'], ['sid-456']) + self.assertIsNone(self._current_context()) + self.assertEqual(len(state['reset_calls']), 1) + + def test_no_conversation_id_skips_context_publication(self): + from api.streaming import _generate_llm_session_title_via_aux + + def fake_call_llm(**kwargs): + return self.MOCK_RESP + + with _patch_tg_config({'provider': '', 'model': 'gpt-4o', 'base_url': ''}): + with patch('agent.auxiliary_client.call_llm', side_effect=fake_call_llm, create=True): + title, status, _raw = _generate_llm_session_title_via_aux( + 'What is the weather?', 'It is sunny.', + ) + + self.assertEqual(title, 'Weather Title') + state = self._portal_state() + self.assertEqual(state['set_calls'], []) + self.assertEqual(state['reset_calls'], []) + + def test_opencode_target_agent_model_publishes_non_empty_session_context(self): + """use_agent_model=True through an OpenCode-style target still carries the key.""" + from api.streaming import _generate_llm_session_title_via_aux + + seen = {} + agent = types.SimpleNamespace( + provider='opencode', model='opencode-go', base_url='http://localhost:3456', + ) + + def fake_call_llm(**kwargs): + from agent.portal_tags import get_conversation_context + seen['context_during_call'] = get_conversation_context() + seen['provider'] = kwargs.get('provider') + seen['base_url'] = kwargs.get('base_url') + return self.MOCK_RESP + + with patch('agent.auxiliary_client.call_llm', side_effect=fake_call_llm, create=True): + title, status, _raw = _generate_llm_session_title_via_aux( + 'What is the weather?', 'It is sunny.', + agent=agent, use_agent_model=True, conversation_id='oc-sid', + ) + + self.assertEqual(title, 'Weather Title') + self.assertEqual(seen['context_during_call'], 'oc-sid') + self.assertEqual(seen['provider'], 'opencode') + self.assertEqual(seen['base_url'], 'http://localhost:3456') + self.assertIsNone(self._current_context()) + + @patch('api.streaming._aux_title_configured', return_value=True) + @patch('api.streaming.get_session') + @patch('api.streaming._generate_llm_session_title_for_agent') + @patch('api.streaming._generate_llm_session_title_via_aux') + def test_background_update_entry_passes_session_id( + self, mock_aux_title, mock_agent_title, mock_get_session, mock_configured, + ): + """_run_background_title_update forwards the session id as conversation id.""" + from api.streaming import _run_background_title_update + + mock_session = MagicMock() + mock_session.title = 'Untitled' + mock_session.llm_title_generated = False + mock_session.messages = [ + {'role': 'user', 'content': 'Hello'}, + {'role': 'assistant', 'content': 'Hi there'}, + ] + mock_get_session.return_value = mock_session + mock_aux_title.return_value = ('Greeting', 'llm_aux', '') + + _run_background_title_update( + session_id='entry-update-sid', + user_text='Hello', + assistant_text='Hi there', + placeholder_title='Untitled', + put_event=lambda _t, _d: None, + agent=MagicMock(), + ) + + self.assertEqual(mock_aux_title.call_args.kwargs.get('conversation_id'), 'entry-update-sid') + mock_agent_title.assert_not_called() + + @patch('api.streaming._aux_title_configured', return_value=True) + @patch('api.streaming.get_session') + @patch('api.streaming._generate_llm_session_title_for_agent') + @patch('api.streaming._generate_llm_session_title_via_aux') + def test_background_refresh_entry_passes_session_id( + self, mock_aux_title, mock_agent_title, mock_get_session, mock_configured, + ): + """_run_background_title_refresh forwards the session id as conversation id.""" + from api.streaming import _run_background_title_refresh + + mock_session = MagicMock() + mock_session.title = 'Old Title' + mock_session.messages = [ + {'role': 'user', 'content': 'Hello again'}, + {'role': 'assistant', 'content': 'More context'}, + ] + mock_get_session.return_value = mock_session + mock_aux_title.return_value = ('Newer Title', 'llm_aux', '') + + _run_background_title_refresh( + session_id='entry-refresh-sid', + user_text='Hello again', + assistant_text='More context', + current_title='Old Title', + put_event=lambda _t, _d: None, + agent=MagicMock(), + ) + + self.assertEqual(mock_aux_title.call_args.kwargs.get('conversation_id'), 'entry-refresh-sid') + mock_agent_title.assert_not_called() + + @patch('api.streaming._generate_llm_session_title_via_aux') + def test_sync_regenerate_entry_passes_session_id(self, mock_aux_title): + """generate_session_title_for_session forwards the session id (sync path).""" + from api.streaming import generate_session_title_for_session + + mock_session = MagicMock() + mock_session.session_id = 'entry-sync-sid' + mock_session.title = 'Untitled' + mock_session.messages = [ + {'role': 'user', 'content': 'Hello sync'}, + {'role': 'assistant', 'content': 'Sync reply'}, + ] + mock_aux_title.return_value = ('Sync Title', 'llm_aux', '') + + title, status, _raw = generate_session_title_for_session(mock_session) + + self.assertEqual(title, 'Sync Title') + self.assertEqual(status, 'llm_aux') + self.assertEqual(mock_aux_title.call_args.kwargs.get('conversation_id'), 'entry-sync-sid') + + if __name__ == '__main__': unittest.main()