chore: handlers tests

tests/agents/test_base_agent.py

@@ -1,4 +1,4 @@
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import pytest
 from application.agents.classic_agent import ClassicAgent
@@ -56,6 +56,73 @@ class TestBaseAgentInitialization:
         agent = ClassicAgent(**agent_base_params)
         assert agent.user == "user123"
 
+    def test_dependency_injection_llm(self, agent_base_params, mock_llm_handler_creator):
+        """When llm is provided, LLMCreator.create_llm is NOT called."""
+        injected_llm = Mock()
+        agent_base_params["llm"] = injected_llm
+        agent = ClassicAgent(**agent_base_params)
+        assert agent.llm is injected_llm
+
+    def test_dependency_injection_llm_handler(self, agent_base_params, mock_llm_creator):
+        """When llm_handler is provided, LLMHandlerCreator is NOT called."""
+        injected_handler = Mock()
+        agent_base_params["llm_handler"] = injected_handler
+        agent = ClassicAgent(**agent_base_params)
+        assert agent.llm_handler is injected_handler
+
+    def test_dependency_injection_tool_executor(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        """When tool_executor is provided, a new one is NOT created."""
+        injected_executor = Mock()
+        injected_executor.tool_calls = []
+        agent_base_params["tool_executor"] = injected_executor
+        agent = ClassicAgent(**agent_base_params)
+        assert agent.tool_executor is injected_executor
+
+    def test_json_schema_normalized(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent_base_params["json_schema"] = {"type": "object"}
+        agent = ClassicAgent(**agent_base_params)
+        assert agent.json_schema == {"type": "object"}
+
+    def test_json_schema_wrapped(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent_base_params["json_schema"] = {"schema": {"type": "string"}}
+        agent = ClassicAgent(**agent_base_params)
+        assert agent.json_schema == {"type": "string"}
+
+    def test_json_schema_invalid_ignored(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent_base_params["json_schema"] = {"bad": "no type"}
+        agent = ClassicAgent(**agent_base_params)
+        assert agent.json_schema is None
+
+    def test_retrieved_docs_defaults_to_empty(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+        assert agent.retrieved_docs == []
+
+    def test_attachments_defaults_to_empty(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent_base_params["attachments"] = None
+        agent = ClassicAgent(**agent_base_params)
+        assert agent.attachments == []
+
+    def test_limited_token_mode_defaults(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+        assert agent.limited_token_mode is False
+        assert agent.limited_request_mode is False
+        assert agent.current_token_count == 0
+        assert agent.context_limit_reached is False
+
 
 @pytest.mark.unit
 class TestBaseAgentBuildMessages:
@@ -602,3 +669,656 @@ class TestBaseAgentHandleResponse:
         assert len(results) == 2
         assert results[0]["type"] == "tool_call"
         assert results[1]["answer"] == "Final answer"
+
+    def test_handle_response_dict_event_passthrough(
+        self,
+        agent_base_params,
+        mock_llm_handler,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+        log_context,
+    ):
+        """Dict events with 'type' key pass through without wrapping."""
+
+        def mock_process(*args):
+            yield {"type": "info", "data": {"message": "processing"}}
+
+        mock_llm_handler.process_message_flow = Mock(side_effect=mock_process)
+
+        agent = ClassicAgent(**agent_base_params)
+        response = Mock()
+        response.message = None
+
+        results = list(agent._handle_response(response, {}, [], log_context))
+        assert results == [{"type": "info", "data": {"message": "processing"}}]
+
+    def test_handle_response_message_object_from_handler(
+        self,
+        agent_base_params,
+        mock_llm_handler,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+        log_context,
+    ):
+        """Response objects with .message.content from handler are unwrapped."""
+        event = Mock()
+        event.message = Mock()
+        event.message.content = "from handler"
+
+        def mock_process(*args):
+            yield event
+
+        mock_llm_handler.process_message_flow = Mock(side_effect=mock_process)
+
+        agent = ClassicAgent(**agent_base_params)
+        response = Mock()
+        response.message = None
+
+        results = list(agent._handle_response(response, {}, [], log_context))
+        assert results[0]["answer"] == "from handler"
+
+
+# ---------------------------------------------------------------------------
+# gen() — the @log_activity decorated entry point
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestBaseAgentGen:
+
+    def test_gen_delegates_to_gen_inner(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+
+        # ClassicAgent._gen_inner is abstract — we patch it
+        with patch.object(agent, "_gen_inner") as mock_inner:
+            mock_inner.return_value = iter([{"answer": "ok"}])
+            results = list(agent.gen("hello"))
+
+        assert any(r.get("answer") == "ok" for r in results)
+
+
+# ---------------------------------------------------------------------------
+# tool_calls property
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestBaseAgentToolCallsProperty:
+
+    def test_getter(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+        agent.tool_executor.tool_calls = ["a", "b"]
+        assert agent.tool_calls == ["a", "b"]
+
+    def test_setter(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+        agent.tool_calls = ["x"]
+        assert agent.tool_executor.tool_calls == ["x"]
+
+
+# ---------------------------------------------------------------------------
+# _calculate_current_context_tokens
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestCalculateContextTokens:
+
+    def test_delegates_to_token_counter(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+        messages = [{"role": "user", "content": "hello"}]
+
+        with patch(
+            "application.api.answer.services.compression.token_counter.TokenCounter"
+        ) as MockTC:
+            MockTC.count_message_tokens.return_value = 42
+            result = agent._calculate_current_context_tokens(messages)
+            assert result == 42
+            MockTC.count_message_tokens.assert_called_once_with(messages)
+
+
+# ---------------------------------------------------------------------------
+# _check_context_limit
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestCheckContextLimit:
+
+    def _make_agent(self, agent_base_params, mock_llm_creator, mock_llm_handler_creator):
+        return ClassicAgent(**agent_base_params)
+
+    def test_below_threshold_returns_false(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = self._make_agent(
+            agent_base_params, mock_llm_creator, mock_llm_handler_creator
+        )
+        messages = [{"role": "user", "content": "hi"}]
+
+        with patch.object(agent, "_calculate_current_context_tokens", return_value=100):
+            with patch(
+                "application.core.model_utils.get_token_limit", return_value=10000
+            ):
+                result = agent._check_context_limit(messages)
+                assert result is False
+                assert agent.current_token_count == 100
+
+    def test_at_threshold_returns_true(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = self._make_agent(
+            agent_base_params, mock_llm_creator, mock_llm_handler_creator
+        )
+        messages = [{"role": "user", "content": "hi"}]
+
+        # threshold = 10000 * 0.8 = 8000; tokens = 8001 → True
+        with patch.object(agent, "_calculate_current_context_tokens", return_value=8001):
+            with patch(
+                "application.core.model_utils.get_token_limit", return_value=10000
+            ):
+                result = agent._check_context_limit(messages)
+                assert result is True
+
+    def test_error_returns_false(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = self._make_agent(
+            agent_base_params, mock_llm_creator, mock_llm_handler_creator
+        )
+        with patch.object(
+            agent,
+            "_calculate_current_context_tokens",
+            side_effect=RuntimeError("boom"),
+        ):
+            result = agent._check_context_limit([])
+            assert result is False
+
+
+# ---------------------------------------------------------------------------
+# _validate_context_size
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestValidateContextSize:
+
+    def test_at_limit_logs_warning(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+        with patch.object(agent, "_calculate_current_context_tokens", return_value=10000):
+            with patch(
+                "application.core.model_utils.get_token_limit", return_value=10000
+            ):
+                # Should not raise
+                agent._validate_context_size([{"role": "user", "content": "x"}])
+                assert agent.current_token_count == 10000
+
+    def test_below_threshold_no_warning(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+        with patch.object(agent, "_calculate_current_context_tokens", return_value=100):
+            with patch(
+                "application.core.model_utils.get_token_limit", return_value=10000
+            ):
+                agent._validate_context_size([])
+                assert agent.current_token_count == 100
+
+    def test_approaching_threshold(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+        # 8500 / 10000 = 85% → above 80% threshold but below 100%
+        with patch.object(agent, "_calculate_current_context_tokens", return_value=8500):
+            with patch(
+                "application.core.model_utils.get_token_limit", return_value=10000
+            ):
+                agent._validate_context_size([])
+                assert agent.current_token_count == 8500
+
+
+# ---------------------------------------------------------------------------
+# _truncate_text_middle
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestTruncateTextMiddle:
+
+    def test_short_text_unchanged(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+        with patch("application.utils.num_tokens_from_string", return_value=5):
+            result = agent._truncate_text_middle("short", max_tokens=100)
+            assert result == "short"
+
+    def test_long_text_truncated(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+        long_text = "A" * 1000
+
+        def fake_tokens(text):
+            return len(text) // 4
+
+        with patch("application.utils.num_tokens_from_string", side_effect=fake_tokens):
+            result = agent._truncate_text_middle(long_text, max_tokens=50)
+            assert "[... content truncated to fit context limit ...]" in result
+            assert len(result) < len(long_text)
+
+    def test_zero_target_returns_empty(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+        with patch("application.utils.num_tokens_from_string", return_value=100):
+            result = agent._truncate_text_middle("some text", max_tokens=0)
+            assert result == ""
+
+
+# ---------------------------------------------------------------------------
+# _truncate_history_to_fit
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestTruncateHistoryToFit:
+
+    def _make_agent(self, agent_base_params, mock_llm_creator, mock_llm_handler_creator):
+        return ClassicAgent(**agent_base_params)
+
+    def test_empty_history(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = self._make_agent(
+            agent_base_params, mock_llm_creator, mock_llm_handler_creator
+        )
+        assert agent._truncate_history_to_fit([], 100) == []
+
+    def test_zero_budget(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = self._make_agent(
+            agent_base_params, mock_llm_creator, mock_llm_handler_creator
+        )
+        history = [{"prompt": "a", "response": "b"}]
+        assert agent._truncate_history_to_fit(history, 0) == []
+
+    def test_fits_all(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = self._make_agent(
+            agent_base_params, mock_llm_creator, mock_llm_handler_creator
+        )
+        history = [
+            {"prompt": "q1", "response": "a1"},
+            {"prompt": "q2", "response": "a2"},
+        ]
+        with patch("application.utils.num_tokens_from_string", return_value=5):
+            result = agent._truncate_history_to_fit(history, 10000)
+            assert len(result) == 2
+
+    def test_partial_fit_keeps_most_recent(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = self._make_agent(
+            agent_base_params, mock_llm_creator, mock_llm_handler_creator
+        )
+        history = [
+            {"prompt": "old", "response": "old_ans"},
+            {"prompt": "new", "response": "new_ans"},
+        ]
+        # Each message = 10 tokens (prompt + response), budget = 15 → only 1 fits
+        with patch("application.utils.num_tokens_from_string", return_value=5):
+            result = agent._truncate_history_to_fit(history, 15)
+            assert len(result) == 1
+            assert result[0]["prompt"] == "new"
+
+    def test_history_with_tool_calls(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = self._make_agent(
+            agent_base_params, mock_llm_creator, mock_llm_handler_creator
+        )
+        history = [
+            {
+                "prompt": "q",
+                "response": "a",
+                "tool_calls": [
+                    {
+                        "tool_name": "t",
+                        "action_name": "act",
+                        "arguments": "{}",
+                        "result": "ok",
+                    }
+                ],
+            }
+        ]
+        with patch("application.utils.num_tokens_from_string", return_value=3):
+            result = agent._truncate_history_to_fit(history, 100)
+            assert len(result) == 1
+
+
+# ---------------------------------------------------------------------------
+# _build_messages — compressed_summary and query truncation
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestBuildMessagesAdvanced:
+
+    def test_compressed_summary_appended(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent_base_params["compressed_summary"] = "Previous conversation summary"
+        agent = ClassicAgent(**agent_base_params)
+
+        with patch(
+            "application.core.model_utils.get_token_limit", return_value=100000
+        ), patch("application.utils.num_tokens_from_string", return_value=10):
+            messages = agent._build_messages("System prompt", "query")
+
+        system_content = messages[0]["content"]
+        assert "Previous conversation summary" in system_content
+
+    def test_query_truncated_when_too_large(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        agent = ClassicAgent(**agent_base_params)
+
+        call_count = {"n": 0}
+
+        def fake_tokens(text):
+            call_count["n"] += 1
+            return len(text)
+
+        with patch(
+            "application.core.model_utils.get_token_limit", return_value=200
+        ), patch("application.utils.num_tokens_from_string", side_effect=fake_tokens):
+            with patch.object(agent, "_truncate_text_middle", return_value="truncated"):
+                with patch.object(agent, "_truncate_history_to_fit", return_value=[]):
+                    messages = agent._build_messages("sys", "A" * 500)
+
+        # The method should have been called for truncation
+        assert messages[-1]["role"] == "user"
+
+    def test_build_messages_with_tool_call_missing_call_id(
+        self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
+    ):
+        """Tool calls without call_id get a generated UUID."""
+        history = [
+            {
+                "tool_calls": [
+                    {
+                        "action_name": "search",
+                        "arguments": "{}",
+                        "result": "found",
+                    }
+                ]
+            }
+        ]
+        agent_base_params["chat_history"] = history
+        agent = ClassicAgent(**agent_base_params)
+
+        with patch(
+            "application.core.model_utils.get_token_limit", return_value=100000
+        ), patch("application.utils.num_tokens_from_string", return_value=5):
+            messages = agent._build_messages("sys", "q")
+
+        tool_msgs = [m for m in messages if m["role"] == "tool"]
+        assert len(tool_msgs) == 1
+
+
+# ---------------------------------------------------------------------------
+# _llm_gen — edge cases
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestLLMGenAdvanced:
+
+    def test_llm_gen_with_attachments(
+        self,
+        agent_base_params,
+        mock_llm,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+    ):
+        agent_base_params["attachments"] = [{"id": "att1", "mime_type": "image/png"}]
+        agent = ClassicAgent(**agent_base_params)
+
+        messages = [{"role": "user", "content": "test"}]
+        agent._llm_gen(messages)
+
+        call_kwargs = mock_llm.gen_stream.call_args[1]
+        assert "_usage_attachments" in call_kwargs
+
+    def test_llm_gen_without_log_context(
+        self,
+        agent_base_params,
+        mock_llm,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+    ):
+        agent = ClassicAgent(**agent_base_params)
+        messages = [{"role": "user", "content": "test"}]
+
+        # Should not raise even without log_context
+        agent._llm_gen(messages, log_context=None)
+        mock_llm.gen_stream.assert_called_once()
+
+    def test_llm_gen_google_structured_output(
+        self,
+        agent_base_params,
+        mock_llm,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+        log_context,
+    ):
+        mock_llm._supports_structured_output = Mock(return_value=True)
+        mock_llm.prepare_structured_output_format = Mock(
+            return_value={"schema": "test"}
+        )
+
+        agent_base_params["json_schema"] = {"type": "object"}
+        agent_base_params["llm_name"] = "google"
+        agent = ClassicAgent(**agent_base_params)
+
+        messages = [{"role": "user", "content": "test"}]
+        agent._llm_gen(messages, log_context)
+
+        call_kwargs = mock_llm.gen_stream.call_args[1]
+        assert "response_schema" in call_kwargs
+
+    def test_llm_gen_no_tools_when_unsupported(
+        self,
+        agent_base_params,
+        mock_llm,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+    ):
+        mock_llm._supports_tools = False
+        agent = ClassicAgent(**agent_base_params)
+        agent.tools = [{"type": "function", "function": {"name": "test"}}]
+
+        messages = [{"role": "user", "content": "test"}]
+        agent._llm_gen(messages)
+
+        call_kwargs = mock_llm.gen_stream.call_args[1]
+        assert "tools" not in call_kwargs
+
+    def test_llm_gen_no_structured_output_when_unsupported(
+        self,
+        agent_base_params,
+        mock_llm,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+    ):
+        mock_llm._supports_structured_output = Mock(return_value=False)
+        agent_base_params["json_schema"] = {"type": "object"}
+        agent = ClassicAgent(**agent_base_params)
+
+        messages = [{"role": "user", "content": "test"}]
+        agent._llm_gen(messages)
+
+        call_kwargs = mock_llm.gen_stream.call_args[1]
+        assert "response_format" not in call_kwargs
+        assert "response_schema" not in call_kwargs
+
+    def test_llm_gen_no_format_when_prepare_returns_none(
+        self,
+        agent_base_params,
+        mock_llm,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+    ):
+        mock_llm._supports_structured_output = Mock(return_value=True)
+        mock_llm.prepare_structured_output_format = Mock(return_value=None)
+
+        agent_base_params["json_schema"] = {"type": "object"}
+        agent_base_params["llm_name"] = "openai"
+        agent = ClassicAgent(**agent_base_params)
+
+        messages = [{"role": "user", "content": "test"}]
+        agent._llm_gen(messages)
+
+        call_kwargs = mock_llm.gen_stream.call_args[1]
+        assert "response_format" not in call_kwargs
+
+
+# ---------------------------------------------------------------------------
+# _llm_handler
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestLLMHandlerMethod:
+
+    def test_delegates_to_handler(
+        self,
+        agent_base_params,
+        mock_llm_handler,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+        log_context,
+    ):
+        mock_llm_handler.process_message_flow = Mock(return_value="result")
+
+        agent = ClassicAgent(**agent_base_params)
+        resp = Mock()
+        result = agent._llm_handler(resp, {}, [], log_context)
+
+        mock_llm_handler.process_message_flow.assert_called_once()
+        assert result == "result"
+        assert len(log_context.stacks) == 1
+        assert log_context.stacks[0]["component"] == "llm_handler"
+
+    def test_without_log_context(
+        self,
+        agent_base_params,
+        mock_llm_handler,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+    ):
+        mock_llm_handler.process_message_flow = Mock(return_value="r")
+        agent = ClassicAgent(**agent_base_params)
+        result = agent._llm_handler(Mock(), {}, [], log_context=None)
+        assert result == "r"
+
+
+# ---------------------------------------------------------------------------
+# _handle_response — structured output on all code paths
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestHandleResponseStructuredAllPaths:
+
+    def test_message_response_with_structured_output(
+        self,
+        agent_base_params,
+        mock_llm,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+        log_context,
+    ):
+        """Structured output on the message.content early-return path."""
+        mock_llm._supports_structured_output = Mock(return_value=True)
+        agent_base_params["json_schema"] = {"type": "object"}
+        agent = ClassicAgent(**agent_base_params)
+
+        response = Mock()
+        response.message = Mock()
+        response.message.content = "structured msg"
+
+        results = list(agent._handle_response(response, {}, [], log_context))
+        assert results[0]["structured"] is True
+        assert results[0]["schema"] == {"type": "object"}
+        assert results[0]["answer"] == "structured msg"
+
+    def test_handler_string_event_with_structured_output(
+        self,
+        agent_base_params,
+        mock_llm,
+        mock_llm_handler,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+        log_context,
+    ):
+        """Structured output on string events from the handler."""
+        mock_llm._supports_structured_output = Mock(return_value=True)
+        agent_base_params["json_schema"] = {"type": "array"}
+
+        def mock_process(*args):
+            yield "handler string"
+
+        mock_llm_handler.process_message_flow = Mock(side_effect=mock_process)
+
+        agent = ClassicAgent(**agent_base_params)
+        response = Mock()
+        response.message = None
+
+        results = list(agent._handle_response(response, {}, [], log_context))
+        assert results[0]["structured"] is True
+        assert results[0]["schema"] == {"type": "array"}
+
+    def test_handler_message_event_with_structured_output(
+        self,
+        agent_base_params,
+        mock_llm,
+        mock_llm_handler,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+        log_context,
+    ):
+        """Structured output on message-object events from the handler."""
+        mock_llm._supports_structured_output = Mock(return_value=True)
+        agent_base_params["json_schema"] = {"type": "number"}
+
+        event = Mock()
+        event.message = Mock()
+        event.message.content = "from handler msg"
+
+        def mock_process(*args):
+            yield event
+
+        mock_llm_handler.process_message_flow = Mock(side_effect=mock_process)
+
+        agent = ClassicAgent(**agent_base_params)
+        response = Mock()
+        response.message = None
+
+        results = list(agent._handle_response(response, {}, [], log_context))
+        assert results[0]["structured"] is True
+        assert results[0]["schema"] == {"type": "number"}
+        assert results[0]["answer"] == "from handler msg"