"""Tests for the continuation infrastructure. Covers ContinuationService, ToolExecutor.check_pause, handler pause signaling, BaseAgent.gen_continuation, and request validation. """ import uuid from unittest.mock import Mock, MagicMock import pytest from application.agents.tool_executor import ToolExecutor from application.llm.handlers.base import LLMHandler, LLMResponse, ToolCall # --------------------------------------------------------------------------- # In-memory MongoDB collection mock (no mongomock / bson dependency) # --------------------------------------------------------------------------- class _InMemoryCollection: """Minimal dict-backed collection supporting find_one, replace_one, delete_one.""" def __init__(self): self._docs = [] def _matches(self, doc, query): return all(doc.get(k) == v for k, v in query.items()) def find_one(self, query): for doc in self._docs: if self._matches(doc, query): import copy return copy.deepcopy(doc) return None def replace_one(self, query, replacement, upsert=False): result = MagicMock() for i, doc in enumerate(self._docs): if self._matches(doc, query): self._docs[i] = dict(replacement) if "_id" not in self._docs[i]: self._docs[i]["_id"] = str(uuid.uuid4()) result.upserted_id = None return result if upsert: new_doc = dict(replacement) new_doc["_id"] = str(uuid.uuid4()) self._docs.append(new_doc) result.upserted_id = new_doc["_id"] else: result.upserted_id = None return result def delete_one(self, query): result = MagicMock() for i, doc in enumerate(self._docs): if self._matches(doc, query): self._docs.pop(i) result.deleted_count = 1 return result result.deleted_count = 0 return result def create_index(self, *args, **kwargs): pass # no-op class _InMemoryDB: def __init__(self): self._collections = {} def __getitem__(self, name): if name not in self._collections: self._collections[name] = _InMemoryCollection() return self._collections[name] @pytest.fixture def mock_mongo_continuation(monkeypatch): """Provide an in-memory MongoDB for ContinuationService (no bson/mongomock).""" db = _InMemoryDB() mock_client = {_get_mongo_db_name(): db} def _get_client(): return mock_client monkeypatch.setattr( "application.api.answer.services.continuation_service.MongoDB.get_client", _get_client, ) monkeypatch.setattr( "application.storage.db.dual_write.dual_write", lambda repo_cls, fn: None, ) return db def _get_mongo_db_name(): from application.core.settings import settings return settings.MONGO_DB_NAME # --------------------------------------------------------------------------- # ContinuationService # --------------------------------------------------------------------------- @pytest.mark.unit @pytest.mark.skip(reason="needs PG fixture rewrite — tracked as part of post-cutover test cleanup") class TestContinuationService: def test_save_and_load(self, mock_mongo_continuation): from application.api.answer.services.continuation_service import ( ContinuationService, ) svc = ContinuationService() svc.save_state( conversation_id="conv-1", user="alice", messages=[{"role": "user", "content": "hi"}], pending_tool_calls=[{"call_id": "c1", "pause_type": "awaiting_approval"}], tools_dict={"0": {"name": "test_tool"}}, tool_schemas=[{"type": "function", "function": {"name": "act_0"}}], agent_config={"model_id": "gpt-4"}, ) state = svc.load_state("conv-1", "alice") assert state is not None assert state["conversation_id"] == "conv-1" assert state["user"] == "alice" assert len(state["messages"]) == 1 assert len(state["pending_tool_calls"]) == 1 assert state["agent_config"]["model_id"] == "gpt-4" def test_load_returns_none_when_missing(self, mock_mongo_continuation): from application.api.answer.services.continuation_service import ( ContinuationService, ) svc = ContinuationService() assert svc.load_state("nonexistent", "alice") is None def test_delete_state(self, mock_mongo_continuation): from application.api.answer.services.continuation_service import ( ContinuationService, ) svc = ContinuationService() svc.save_state( conversation_id="conv-2", user="bob", messages=[], pending_tool_calls=[], tools_dict={}, tool_schemas=[], agent_config={}, ) assert svc.delete_state("conv-2", "bob") is True assert svc.load_state("conv-2", "bob") is None def test_delete_nonexistent(self, mock_mongo_continuation): from application.api.answer.services.continuation_service import ( ContinuationService, ) svc = ContinuationService() assert svc.delete_state("nope", "nope") is False def test_upsert_replaces_existing(self, mock_mongo_continuation): from application.api.answer.services.continuation_service import ( ContinuationService, ) svc = ContinuationService() svc.save_state( conversation_id="conv-3", user="carol", messages=[{"role": "user", "content": "v1"}], pending_tool_calls=[], tools_dict={}, tool_schemas=[], agent_config={}, ) svc.save_state( conversation_id="conv-3", user="carol", messages=[{"role": "user", "content": "v2"}], pending_tool_calls=[{"call_id": "c2"}], tools_dict={}, tool_schemas=[], agent_config={}, ) state = svc.load_state("conv-3", "carol") assert state["messages"][0]["content"] == "v2" assert len(state["pending_tool_calls"]) == 1 # --------------------------------------------------------------------------- # ToolExecutor.check_pause # --------------------------------------------------------------------------- @pytest.mark.unit class TestCheckPause: def _make_call(self, name="action_0", call_id="c1", arguments="{}"): call = Mock() call.name = name call.id = call_id call.arguments = arguments call.thought_signature = None return call def test_returns_none_for_normal_tool(self): executor = ToolExecutor() tools_dict = { "0": { "name": "brave", "actions": [ {"name": "search", "active": True, "parameters": {}}, ], } } call = self._make_call(name="search_0") result = executor.check_pause(tools_dict, call, "OpenAILLM") assert result is None def test_returns_pause_for_client_side_tool(self): executor = ToolExecutor() tools_dict = { "0": { "name": "get_weather", "client_side": True, "actions": [ {"name": "get_weather", "active": True, "parameters": {}}, ], } } call = self._make_call(name="get_weather_0") result = executor.check_pause(tools_dict, call, "OpenAILLM") assert result is not None assert result["pause_type"] == "requires_client_execution" assert result["call_id"] == "c1" assert result["tool_id"] == "0" def test_returns_pause_for_approval_required(self): executor = ToolExecutor() tools_dict = { "0": { "name": "telegram", "actions": [ { "name": "send_msg", "active": True, "require_approval": True, "parameters": {}, }, ], } } call = self._make_call(name="send_msg_0") result = executor.check_pause(tools_dict, call, "OpenAILLM") assert result is not None assert result["pause_type"] == "awaiting_approval" def test_returns_none_when_parse_fails(self): executor = ToolExecutor() call = self._make_call(name="bad_name_no_id", arguments="not json") # Bad arguments will cause parse error -> None result = executor.check_pause({}, call, "OpenAILLM") assert result is None def test_returns_none_when_tool_not_in_dict(self): executor = ToolExecutor() call = self._make_call(name="action_99") result = executor.check_pause({"0": {"name": "t"}}, call, "OpenAILLM") assert result is None def test_api_tool_approval(self): executor = ToolExecutor() tools_dict = { "0": { "name": "api_tool", "config": { "actions": { "delete_user": { "name": "delete_user", "require_approval": True, "url": "http://example.com", "method": "DELETE", "active": True, } } }, } } call = self._make_call(name="delete_user_0") result = executor.check_pause(tools_dict, call, "OpenAILLM") assert result is not None assert result["pause_type"] == "awaiting_approval" # --------------------------------------------------------------------------- # Handler pause signaling (handle_tool_calls returns pending_actions) # --------------------------------------------------------------------------- class ConcreteHandler(LLMHandler): """Minimal concrete handler for testing.""" def parse_response(self, response): return LLMResponse( content=str(response), tool_calls=[], finish_reason="stop", raw_response=response, ) def create_tool_message(self, tool_call, result): return { "role": "tool", "content": [ { "function_response": { "name": tool_call.name, "response": {"result": result}, "call_id": tool_call.id, } } ], } def _iterate_stream(self, response): for chunk in response: yield chunk @pytest.mark.unit class TestHandlerPauseSignaling: def _make_agent(self): agent = Mock() agent._check_context_limit = Mock(return_value=False) agent.context_limit_reached = False agent.llm.__class__.__name__ = "MockLLM" agent.tool_executor.check_pause = Mock(return_value=None) def fake_execute(tools_dict, call): yield {"type": "tool_call", "data": {"status": "pending"}} return ("tool result", call.id) agent._execute_tool_action = Mock(side_effect=fake_execute) return agent def test_no_pause_returns_none_pending(self): handler = ConcreteHandler() agent = self._make_agent() call = ToolCall(id="c1", name="action_0", arguments="{}") gen = handler.handle_tool_calls(agent, [call], {"0": {"name": "t"}}, []) events = [] messages = None pending = "NOT_SET" try: while True: events.append(next(gen)) except StopIteration as e: messages, pending = e.value assert pending is None assert messages is not None def test_pause_returns_pending_actions(self): handler = ConcreteHandler() agent = self._make_agent() agent.tool_executor.check_pause = Mock(return_value={ "call_id": "c1", "name": "send_msg_0", "tool_name": "telegram", "tool_id": "0", "action_name": "send_msg", "arguments": {"text": "hello"}, "pause_type": "awaiting_approval", "thought_signature": None, }) call = ToolCall(id="c1", name="send_msg_0", arguments='{"text": "hello"}') gen = handler.handle_tool_calls( agent, [call], {"0": {"name": "telegram"}}, [] ) events = [] pending = None try: while True: events.append(next(gen)) except StopIteration as e: messages, pending = e.value assert pending is not None assert len(pending) == 1 assert pending[0]["pause_type"] == "awaiting_approval" # Should have yielded a tool_call event with awaiting_approval status pause_events = [ e for e in events if e.get("type") == "tool_call" and e.get("data", {}).get("status") == "awaiting_approval" ] assert len(pause_events) == 1 def test_pause_propagates_device_id_for_remote_device(self): """``pause_info['device_id']`` (set in tool_executor for the remote_device tool) must be copied into the emitted ``tool_call`` event so the approval UI can render the "don't ask again" button.""" handler = ConcreteHandler() agent = self._make_agent() agent.tool_executor.check_pause = Mock(return_value={ "call_id": "c1", "name": "run_command_0", "tool_name": "remote_device", "tool_id": "0", "action_name": "run_command", "arguments": {"command": "ls"}, "pause_type": "awaiting_approval", "device_id": "dev_abc", "thought_signature": None, }) call = ToolCall(id="c1", name="run_command_0", arguments='{"command": "ls"}') gen = handler.handle_tool_calls( agent, [call], {"0": {"name": "remote_device"}}, [] ) events = [] try: while True: events.append(next(gen)) except StopIteration: pass pause_events = [ e for e in events if e.get("type") == "tool_call" and e.get("data", {}).get("status") == "awaiting_approval" ] assert len(pause_events) == 1 assert pause_events[0]["data"].get("device_id") == "dev_abc" def test_pause_omits_device_id_for_non_remote_tools(self): """``device_id`` must NOT leak into pause events for tools that don't ship one in ``pause_info``.""" handler = ConcreteHandler() agent = self._make_agent() agent.tool_executor.check_pause = Mock(return_value={ "call_id": "c1", "name": "send_msg_0", "tool_name": "telegram", "tool_id": "0", "action_name": "send_msg", "arguments": {"text": "hello"}, "pause_type": "awaiting_approval", "thought_signature": None, }) call = ToolCall(id="c1", name="send_msg_0", arguments='{"text": "hello"}') gen = handler.handle_tool_calls( agent, [call], {"0": {"name": "telegram"}}, [] ) events = [] try: while True: events.append(next(gen)) except StopIteration: pass pause_events = [ e for e in events if e.get("type") == "tool_call" and e.get("data", {}).get("status") == "awaiting_approval" ] assert len(pause_events) == 1 assert "device_id" not in pause_events[0]["data"] def test_mixed_execute_and_pause(self): """One tool executes, another needs approval.""" handler = ConcreteHandler() agent = self._make_agent() call_count = {"n": 0} def selective_pause(tools_dict, call, llm_class): call_count["n"] += 1 if call_count["n"] == 2: return { "call_id": "c2", "name": "danger_0", "tool_name": "danger", "tool_id": "0", "action_name": "danger", "arguments": {}, "pause_type": "awaiting_approval", "thought_signature": None, } return None agent.tool_executor.check_pause = Mock(side_effect=selective_pause) calls = [ ToolCall(id="c1", name="safe_0", arguments="{}"), ToolCall(id="c2", name="danger_0", arguments="{}"), ] gen = handler.handle_tool_calls( agent, calls, {"0": {"name": "multi"}}, [] ) events = [] try: while True: events.append(next(gen)) except StopIteration as e: messages, pending = e.value # First tool was executed normally assert agent._execute_tool_action.call_count == 1 # Second tool is pending assert pending is not None assert len(pending) == 1 assert pending[0]["call_id"] == "c2" # --------------------------------------------------------------------------- # handle_streaming yields tool_calls_pending # --------------------------------------------------------------------------- @pytest.mark.unit class TestStreamingPause: def test_streaming_yields_tool_calls_pending(self): handler = ConcreteHandler() agent = Mock() agent.llm = Mock() agent.model_id = "test" agent.tools = [] agent._check_context_limit = Mock(return_value=False) agent.context_limit_reached = False agent.llm.__class__.__name__ = "MockLLM" pause_info = { "call_id": "c1", "name": "fn_0", "tool_name": "test", "tool_id": "0", "action_name": "fn", "arguments": {}, "pause_type": "awaiting_approval", "thought_signature": None, } agent.tool_executor.check_pause = Mock(return_value=pause_info) chunk = LLMResponse( content="", tool_calls=[ToolCall(id="c1", name="fn_0", arguments="{}", index=0)], finish_reason="tool_calls", raw_response={}, ) handler.parse_response = lambda c: c def fake_iterate(response): yield from response handler._iterate_stream = fake_iterate gen = handler.handle_streaming(agent, [chunk], {"0": {"name": "t"}}, []) events = list(gen) # Should contain a tool_calls_pending event pending_events = [ e for e in events if isinstance(e, dict) and e.get("type") == "tool_calls_pending" ] assert len(pending_events) == 1 assert len(pending_events[0]["data"]["pending_tool_calls"]) == 1 # Agent should have _pending_continuation set assert hasattr(agent, "_pending_continuation") # --------------------------------------------------------------------------- # BaseAgent.gen_continuation # --------------------------------------------------------------------------- @pytest.mark.unit class TestGenContinuation: def test_approved_tool_executes(self): """When a tool action is approved, the tool is executed.""" from application.agents.classic_agent import ClassicAgent mock_llm = Mock() mock_llm._supports_tools = True mock_llm.gen_stream = Mock(return_value=iter(["Final answer"])) mock_llm._supports_structured_output = Mock(return_value=False) mock_llm.__class__.__name__ = "MockLLM" mock_handler = Mock() mock_handler.process_message_flow = Mock(return_value=iter([])) mock_handler.create_tool_message = Mock( return_value={"role": "tool", "content": [{"function_response": { "name": "act_0", "response": {"result": "done"}, "call_id": "c1" }}]} ) mock_executor = Mock() mock_executor.tool_calls = [] mock_executor.prepare_tools_for_llm = Mock(return_value=[]) mock_executor.get_truncated_tool_calls = Mock(return_value=[]) def fake_execute(tools_dict, call, llm_class): yield {"type": "tool_call", "data": {"status": "pending"}} return ("result_data", "c1") mock_executor.execute = Mock(side_effect=fake_execute) agent = ClassicAgent( endpoint="stream", llm_name="openai", model_id="gpt-4", api_key="test", llm=mock_llm, llm_handler=mock_handler, tool_executor=mock_executor, ) messages = [{"role": "system", "content": "You are helpful."}] tools_dict = {"0": {"name": "test_tool"}} pending = [ { "call_id": "c1", "name": "act_0", "tool_name": "test_tool", "tool_id": "0", "action_name": "act", "arguments": {"q": "test"}, "pause_type": "awaiting_approval", "thought_signature": None, } ] tool_actions = [{"call_id": "c1", "decision": "approved"}] list(agent.gen_continuation(messages, tools_dict, pending, tool_actions)) # Tool should have been executed assert mock_executor.execute.called def test_denied_tool_sends_denial(self): """When a tool action is denied, a denial message is added.""" from application.agents.classic_agent import ClassicAgent mock_llm = Mock() mock_llm._supports_tools = True mock_llm.gen_stream = Mock(return_value=iter(["Answer"])) mock_llm._supports_structured_output = Mock(return_value=False) mock_llm.__class__.__name__ = "MockLLM" mock_handler = Mock() mock_handler.process_message_flow = Mock(return_value=iter([])) mock_handler.create_tool_message = Mock( return_value={"role": "tool", "content": "denied"} ) mock_executor = Mock() mock_executor.tool_calls = [] mock_executor.prepare_tools_for_llm = Mock(return_value=[]) mock_executor.get_truncated_tool_calls = Mock(return_value=[]) agent = ClassicAgent( endpoint="stream", llm_name="openai", model_id="gpt-4", api_key="test", llm=mock_llm, llm_handler=mock_handler, tool_executor=mock_executor, ) messages = [{"role": "system", "content": "test"}] pending = [ { "call_id": "c1", "name": "danger_0", "tool_name": "danger", "tool_id": "0", "action_name": "danger", "arguments": {}, "pause_type": "awaiting_approval", "thought_signature": None, } ] tool_actions = [ {"call_id": "c1", "decision": "denied", "comment": "too risky"} ] events = list( agent.gen_continuation(messages, {"0": {"name": "danger"}}, pending, tool_actions) ) # Should have a denied tool_call event denied = [ e for e in events if isinstance(e, dict) and e.get("type") == "tool_call" and e.get("data", {}).get("status") == "denied" ] assert len(denied) == 1 # create_tool_message should have been called with denial text denial_arg = mock_handler.create_tool_message.call_args[0][1] assert "denied" in denial_arg.lower() assert "too risky" in denial_arg def test_client_result_appended(self): """Client-provided tool result is added to messages.""" from application.agents.classic_agent import ClassicAgent mock_llm = Mock() mock_llm._supports_tools = True mock_llm.gen_stream = Mock(return_value=iter(["Done"])) mock_llm._supports_structured_output = Mock(return_value=False) mock_llm.__class__.__name__ = "MockLLM" mock_handler = Mock() mock_handler.process_message_flow = Mock(return_value=iter([])) mock_handler.create_tool_message = Mock( return_value={"role": "tool", "content": "client result"} ) mock_executor = Mock() mock_executor.tool_calls = [] mock_executor.prepare_tools_for_llm = Mock(return_value=[]) mock_executor.get_truncated_tool_calls = Mock(return_value=[]) agent = ClassicAgent( endpoint="stream", llm_name="openai", model_id="gpt-4", api_key="test", llm=mock_llm, llm_handler=mock_handler, tool_executor=mock_executor, ) messages = [{"role": "system", "content": "test"}] pending = [ { "call_id": "c1", "name": "weather_0", "tool_name": "weather", "tool_id": "0", "action_name": "weather", "arguments": {"city": "SF"}, "pause_type": "requires_client_execution", "thought_signature": None, } ] tool_actions = [{"call_id": "c1", "result": {"temp": "72F"}}] events = list( agent.gen_continuation(messages, {"0": {"name": "weather"}}, pending, tool_actions) ) # create_tool_message was called with the client result result_arg = mock_handler.create_tool_message.call_args[0][1] assert "72F" in result_arg # Should have a completed tool_call event completed = [ e for e in events if isinstance(e, dict) and e.get("type") == "tool_call" and e.get("data", {}).get("status") == "completed" ] assert len(completed) == 1 # --------------------------------------------------------------------------- # validate_request # --------------------------------------------------------------------------- @pytest.mark.unit class TestValidateRequest: @pytest.fixture(autouse=True) def _app_context(self): from flask import Flask app = Flask(__name__) with app.app_context(): yield def test_continuation_request_without_question(self): from application.api.answer.routes.base import BaseAnswerResource base = BaseAnswerResource() data = { "conversation_id": "conv-1", "tool_actions": [{"call_id": "c1", "decision": "approved"}], } result = base.validate_request(data) assert result is None # Valid def test_continuation_request_missing_conversation_id(self): from application.api.answer.routes.base import BaseAnswerResource base = BaseAnswerResource() data = { "tool_actions": [{"call_id": "c1", "decision": "approved"}], } result = base.validate_request(data) assert result is not None # Error — missing conversation_id def test_normal_request_still_requires_question(self): from application.api.answer.routes.base import BaseAnswerResource base = BaseAnswerResource() data = {"conversation_id": "conv-1"} result = base.validate_request(data) assert result is not None # Error — missing question # --------------------------------------------------------------------------- # Resume durability: mark_resuming on resume, delete only on success # --------------------------------------------------------------------------- @pytest.mark.unit class TestResumeMarkResuming: """Resumed runs must mark state ``resuming`` instead of deleting it eagerly; the row stays in PG so a crashed resume can be retried.""" def test_resume_calls_mark_resuming_not_delete(self, monkeypatch): """``resume_from_tool_actions`` flips the row to 'resuming' and does not delete it before the run finishes.""" from application.api.answer.services import ( continuation_service as cont_mod, ) from application.api.answer.services import stream_processor as sp_mod from application.llm import llm_creator as llm_creator_mod from application.llm.handlers import handler_creator as handler_mod cont_service = MagicMock() cont_service.load_state.return_value = { "messages": [], "pending_tool_calls": [], "tools_dict": {}, "tool_schemas": [], "agent_config": { "model_id": "m1", "model_user_id": None, "llm_name": "openai", "api_key": "k", "user_api_key": None, "agent_id": None, "agent_type": "ClassicAgent", "prompt": "", "json_schema": None, "retriever_config": None, }, "client_tools": None, } cont_service.mark_resuming.return_value = True monkeypatch.setattr( cont_mod, "ContinuationService", lambda: cont_service ) monkeypatch.setattr( llm_creator_mod.LLMCreator, "create_llm", lambda *a, **kw: MagicMock(), ) monkeypatch.setattr( handler_mod.LLMHandlerCreator, "create_handler", lambda *a, **kw: MagicMock(), ) from application.agents import agent_creator as ac_mod from application.agents import tool_executor as te_mod monkeypatch.setattr( te_mod, "ToolExecutor", lambda **kw: MagicMock(client_tools=None) ) monkeypatch.setattr( ac_mod.AgentCreator, "create_agent", lambda *a, **kw: MagicMock() ) sp = sp_mod.StreamProcessor.__new__(sp_mod.StreamProcessor) sp.data = {} sp.decoded_token = {"sub": "alice"} sp.initial_user_id = "alice" sp.conversation_id = "00000000-0000-0000-0000-000000000001" sp.agent_config = {} sp.resume_from_tool_actions( tool_actions=[], conversation_id="00000000-0000-0000-0000-000000000001", ) cont_service.mark_resuming.assert_called_once_with( "00000000-0000-0000-0000-000000000001", "alice" ) cont_service.delete_state.assert_not_called() def test_resume_extracts_reserved_message_id_from_agent_config( self, monkeypatch ): """The WAL placeholder id stashed in ``agent_config`` at pause time must be hoisted onto the processor so the resumed ``complete_stream`` finalises the same row instead of stranding it.""" from application.api.answer.services import ( continuation_service as cont_mod, ) from application.api.answer.services import stream_processor as sp_mod from application.llm import llm_creator as llm_creator_mod from application.llm.handlers import handler_creator as handler_mod reserved_id = "22222222-2222-2222-2222-222222222222" cont_service = MagicMock() cont_service.load_state.return_value = { "messages": [], "pending_tool_calls": [], "tools_dict": {}, "tool_schemas": [], "agent_config": { "model_id": "m1", "model_user_id": None, "llm_name": "openai", "api_key": "k", "user_api_key": None, "agent_id": None, "agent_type": "ClassicAgent", "prompt": "", "json_schema": None, "retriever_config": None, "reserved_message_id": reserved_id, }, "client_tools": None, } cont_service.mark_resuming.return_value = True monkeypatch.setattr(cont_mod, "ContinuationService", lambda: cont_service) monkeypatch.setattr( llm_creator_mod.LLMCreator, "create_llm", lambda *a, **kw: MagicMock(), ) monkeypatch.setattr( handler_mod.LLMHandlerCreator, "create_handler", lambda *a, **kw: MagicMock(), ) from application.agents import agent_creator as ac_mod from application.agents import tool_executor as te_mod monkeypatch.setattr( te_mod, "ToolExecutor", lambda **kw: MagicMock(client_tools=None) ) monkeypatch.setattr( ac_mod.AgentCreator, "create_agent", lambda *a, **kw: MagicMock() ) sp = sp_mod.StreamProcessor.__new__(sp_mod.StreamProcessor) sp.data = {} sp.decoded_token = {"sub": "alice"} sp.initial_user_id = "alice" sp.conversation_id = "00000000-0000-0000-0000-000000000001" sp.agent_config = {} sp.reserved_message_id = None sp.resume_from_tool_actions( tool_actions=[], conversation_id="00000000-0000-0000-0000-000000000001", ) assert sp.reserved_message_id == reserved_id def test_resume_resolves_owner_from_api_key_when_no_jwt(self, monkeypatch): """api_key-authenticated resumes (no JWT) must resolve the agent owner before loading state. On the native /stream and /api/answer routes the agent key lives in the request body, so ``request.decoded_token`` — and hence ``initial_user_id`` — is None. The pending state was saved under the owner's id during the first turn, so the resume has to resolve the owner here or the lookup misses and the run 400s with "No pending tool state found for this conversation". """ from contextlib import contextmanager from application.api.answer.services import ( continuation_service as cont_mod, ) from application.api.answer.services import stream_processor as sp_mod from application.llm import llm_creator as llm_creator_mod from application.llm.handlers import handler_creator as handler_mod cont_service = MagicMock() cont_service.load_state.return_value = { "messages": [], "pending_tool_calls": [], "tools_dict": {}, "tool_schemas": [], "agent_config": { "model_id": "m1", "model_user_id": None, "llm_name": "openai", "api_key": "k", "user_api_key": None, "agent_id": None, "agent_type": "ClassicAgent", "prompt": "", "json_schema": None, "retriever_config": None, }, "client_tools": None, } cont_service.mark_resuming.return_value = True monkeypatch.setattr(cont_mod, "ContinuationService", lambda: cont_service) monkeypatch.setattr( llm_creator_mod.LLMCreator, "create_llm", lambda *a, **kw: MagicMock(), ) monkeypatch.setattr( handler_mod.LLMHandlerCreator, "create_handler", lambda *a, **kw: MagicMock(), ) from application.agents import agent_creator as ac_mod from application.agents import tool_executor as te_mod monkeypatch.setattr( te_mod, "ToolExecutor", lambda **kw: MagicMock(client_tools=None) ) monkeypatch.setattr( ac_mod.AgentCreator, "create_agent", lambda *a, **kw: MagicMock() ) # The body api_key resolves to its owning user. fake_repo = MagicMock() fake_repo.find_by_key.return_value = {"user_id": "owner-1"} @contextmanager def _fake_db_readonly(): yield MagicMock() monkeypatch.setattr(sp_mod, "db_readonly", _fake_db_readonly) monkeypatch.setattr(sp_mod, "AgentsRepository", lambda conn: fake_repo) conv_id = "00000000-0000-0000-0000-000000000009" sp = sp_mod.StreamProcessor.__new__(sp_mod.StreamProcessor) sp.data = {"api_key": "agent-key-1"} sp.decoded_token = None sp.initial_user_id = None sp.conversation_id = conv_id sp.agent_config = {} sp.reserved_message_id = None sp.resume_from_tool_actions(tool_actions=[], conversation_id=conv_id) fake_repo.find_by_key.assert_called_once_with("agent-key-1") # The lookup + claim now key on the owner id, not None. cont_service.load_state.assert_called_once_with(conv_id, "owner-1") cont_service.mark_resuming.assert_called_once_with(conv_id, "owner-1") assert sp.initial_user_id == "owner-1" assert sp.decoded_token == {"sub": "owner-1"} @pytest.mark.unit class TestContinuationServiceMarkResuming: """``ContinuationService.mark_resuming`` is the thin wrapper used by the resume path; it should flip the repository row in place.""" def test_mark_resuming_flips_pending_row(self, pg_engine, monkeypatch): from contextlib import contextmanager from application.api.answer.services import ( continuation_service as cont_mod, ) from application.storage.db.repositories.conversations import ( ConversationsRepository, ) from application.storage.db.repositories.pending_tool_state import ( PendingToolStateRepository, ) with pg_engine.begin() as conn: conv = ConversationsRepository(conn).create("alice", "c") PendingToolStateRepository(conn).save_state( conv["id"], "alice", messages=[], pending_tool_calls=[], tools_dict={}, tool_schemas=[], agent_config={}, ) @contextmanager def _session(): with pg_engine.begin() as conn: yield conn @contextmanager def _readonly(): with pg_engine.connect() as conn: yield conn monkeypatch.setattr(cont_mod, "db_session", _session) monkeypatch.setattr(cont_mod, "db_readonly", _readonly) svc = cont_mod.ContinuationService() flipped = svc.mark_resuming(conv["id"], "alice") assert flipped is True with pg_engine.connect() as conn: row = PendingToolStateRepository(conn).load_state( conv["id"], "alice" ) assert row["status"] == "resuming" assert row["resumed_at"] is not None def test_mark_resuming_returns_false_for_unknown_conv( self, pg_engine, monkeypatch ): from contextlib import contextmanager from application.api.answer.services import ( continuation_service as cont_mod, ) @contextmanager def _session(): with pg_engine.begin() as conn: yield conn @contextmanager def _readonly(): with pg_engine.connect() as conn: yield conn monkeypatch.setattr(cont_mod, "db_session", _session) monkeypatch.setattr(cont_mod, "db_readonly", _readonly) svc = cont_mod.ContinuationService() # Not a UUID and no legacy row exists. assert svc.mark_resuming("not-a-uuid", "alice") is False # --------------------------------------------------------------------------- # Refresh during pause: per-call ``tool_call`` events must reconstruct # into ``tool_calls`` so the approval bar re-renders. # --------------------------------------------------------------------------- @pytest.mark.integration class TestReconstructPartialToolCallReplay: """The pause path in ``llm/handlers/base.py`` emits a per-call ``tool_call`` event (status ``awaiting_approval``) — not a bulk ``tool_calls`` snapshot. ``reconstruct_partial`` must overlay these so a conversation refresh while paused still shows the approval bar. """ @staticmethod def _seed_message(conn): from sqlalchemy import text user_id = f"user-{uuid.uuid4().hex[:8]}" conv_id = uuid.uuid4() msg_id = uuid.uuid4() conn.execute( text("INSERT INTO users (user_id) VALUES (:u)"), {"u": user_id}, ) conn.execute( text( "INSERT INTO conversations (id, user_id, name) " "VALUES (:id, :u, 'test')" ), {"id": conv_id, "u": user_id}, ) conn.execute( text( "INSERT INTO conversation_messages (id, conversation_id, user_id, position) " "VALUES (:id, :c, :u, 0)" ), {"id": msg_id, "c": conv_id, "u": user_id}, ) return str(msg_id) def test_paused_tool_call_event_lands_in_tool_calls(self, pg_conn): from application.storage.db.repositories.message_events import ( MessageEventsRepository, ) message_id = self._seed_message(pg_conn) repo = MessageEventsRepository(pg_conn) # Mirror the exact shape emitted at base.py:941-950 on pause. repo.record( message_id, 0, "tool_call", { "type": "tool_call", "data": { "tool_name": "remote_device", "call_id": "call_remote_test_1", "action_name": "run_command", "arguments": {"command": "ls -la /tmp"}, "status": "awaiting_approval", "device_id": "dev_abc", }, }, ) partial = repo.reconstruct_partial(message_id) assert len(partial["tool_calls"]) == 1 tc = partial["tool_calls"][0] assert tc["status"] == "awaiting_approval" assert tc["device_id"] == "dev_abc" assert tc["call_id"] == "call_remote_test_1" def test_completed_event_replaces_paused_event(self, pg_conn): from application.storage.db.repositories.message_events import ( MessageEventsRepository, ) message_id = self._seed_message(pg_conn) repo = MessageEventsRepository(pg_conn) repo.record( message_id, 0, "tool_call", { "type": "tool_call", "data": { "tool_name": "remote_device", "call_id": "c1", "status": "awaiting_approval", }, }, ) repo.record( message_id, 1, "tool_call", { "type": "tool_call", "data": { "tool_name": "remote_device", "call_id": "c1", "status": "completed", "result": "/tmp listing here", }, }, ) partial = repo.reconstruct_partial(message_id) assert len(partial["tool_calls"]) == 1 assert partial["tool_calls"][0]["status"] == "completed" assert partial["tool_calls"][0]["result"] == "/tmp listing here"