From bc9bea382d4080adc90af46ca02267f8b68246fe Mon Sep 17 00:00:00 2001 From: rohanbuilds-ai Date: Thu, 15 Oct 2026 08:18:24 +0530 Subject: [PATCH] fix(agents): handle unresolvable task_data IDs (#589) - Gracefully handle missing/unresolvable vendor_id and invoice_id in specialized prompt builders - Move BaseAgent pre-loop setup into protected execution block - Return structured task_status='failed' instead of raising unhandled exceptions - Ensure MCP disconnection and workflow context cleanup in finally block even on setup failure - Add comprehensive regression tests for agent error handling --- finbot/agents/base.py | 63 +-- finbot/agents/specialized/invoice.py | 16 +- finbot/agents/specialized/onboarding.py | 16 +- .../unit/agents/test_agent_error_handling.py | 389 ++++++++++++++++++ 4 files changed, 452 insertions(+), 32 deletions(-) create mode 100644 tests/unit/agents/test_agent_error_handling.py diff --git a/finbot/agents/base.py b/finbot/agents/base.py index 5e4c0480..ec7af822 100644 --- a/finbot/agents/base.py +++ b/finbot/agents/base.py @@ -80,29 +80,38 @@ async def _run_agent_loop( """ await self.log_task_start(task_data=task_data) - # Connect to MCP servers if the agent has any configured - await self._connect_mcp_servers() - - system_prompt = self._get_final_system_prompt() - user_prompt = await self._get_user_prompt(task_data=task_data) - - # Store the user prompt on the workflow so every event - # (agent + business) in this workflow carries it. - event_bus.set_workflow_context(self.workflow_id, user_prompt=user_prompt) - - messages = [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_prompt}, - ] - - tools = self._get_final_tool_definitions() - - max_iterations = self._get_max_iterations() - max_stall_iterations = self._get_max_stall_iterations() - callables = self._get_final_callables() - stall_count = 0 - try: + try: + # Connect to MCP servers if the agent has any configured + await self._connect_mcp_servers() + + system_prompt = self._get_final_system_prompt() + user_prompt = await self._get_user_prompt(task_data=task_data) + + # Store the user prompt on the workflow so every event + # (agent + business) in this workflow carries it. + event_bus.set_workflow_context(self.workflow_id, user_prompt=user_prompt) + + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ] + + tools = self._get_final_tool_definitions() + + max_iterations = self._get_max_iterations() + max_stall_iterations = self._get_max_stall_iterations() + callables = self._get_final_callables() + stall_count = 0 + except Exception as e: # pylint: disable=broad-exception-caught + logger.error("Agent setup failed: %s", e) + task_result = await self._complete_task( + task_status="failed", + task_summary=f"Agent setup failed: {e}", + ) + await self.log_task_completion(task_result=task_result) + return task_result + for iteration in range(max_iterations): # Emit iteration start event await event_bus.emit_agent_event( @@ -305,8 +314,14 @@ async def _run_agent_loop( await self.log_task_completion(task_result=task_result) return task_result finally: - await self._disconnect_mcp_servers() - event_bus.clear_workflow_context(self.workflow_id) + try: + await self._disconnect_mcp_servers() + except Exception as e: # pylint: disable=broad-exception-caught + logger.error("Failed to disconnect MCP servers: %s", e) + try: + event_bus.clear_workflow_context(self.workflow_id) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error("Failed to clear workflow context: %s", e) def _get_system_prompt(self) -> str: """ diff --git a/finbot/agents/specialized/invoice.py b/finbot/agents/specialized/invoice.py index 18319f67..94a5f85d 100644 --- a/finbot/agents/specialized/invoice.py +++ b/finbot/agents/specialized/invoice.py @@ -227,10 +227,18 @@ async def _get_user_prompt(self, task_data: dict[str, Any] | None = None) -> str invoice_id = task_data.get("invoice_id", None) description = "" agent_notes = "" - if invoice_id: - invoice = await get_invoice_details(invoice_id, self.session_context) - description = invoice.get("description", "") - agent_notes = invoice.get("agent_notes", "") + if invoice_id is not None: + try: + numeric_invoice_id = int(invoice_id) + invoice = await get_invoice_details(numeric_invoice_id, self.session_context) + description = invoice.get("description", "") + agent_notes = invoice.get("agent_notes", "") + except (ValueError, TypeError) as e: + logger.warning( + "Could not retrieve invoice details for invoice_id %s in prompt: %s", + invoice_id, + e, + ) user_prompt += f""" Here is the description of the invoice. Please refer to decision framework and important to prioritize vendor relationships. diff --git a/finbot/agents/specialized/onboarding.py b/finbot/agents/specialized/onboarding.py index 75c696e6..d758c3b3 100644 --- a/finbot/agents/specialized/onboarding.py +++ b/finbot/agents/specialized/onboarding.py @@ -192,10 +192,18 @@ async def _get_user_prompt(self, task_data: dict[str, Any] | None = None) -> str vendor_id = task_data.get("vendor_id", None) services = "" agent_notes = "" - if vendor_id: - vendor = await get_vendor_details(vendor_id, self.session_context) - services = vendor.get("services", "") - agent_notes = vendor.get("agent_notes", "") + if vendor_id is not None: + try: + numeric_vendor_id = int(vendor_id) + vendor = await get_vendor_details(numeric_vendor_id, self.session_context) + services = vendor.get("services", "") + agent_notes = vendor.get("agent_notes", "") + except (ValueError, TypeError) as e: + logger.warning( + "Could not retrieve vendor details for vendor_id %s in prompt: %s", + vendor_id, + e, + ) user_prompt += f""" Here are the services provided by the vendor, please refer to decision framework and important to prioritize vendor relationships. diff --git a/tests/unit/agents/test_agent_error_handling.py b/tests/unit/agents/test_agent_error_handling.py new file mode 100644 index 00000000..a996d98b --- /dev/null +++ b/tests/unit/agents/test_agent_error_handling.py @@ -0,0 +1,389 @@ +""" +Unit tests for agent error handling -- Issue #589. + +Verifies that unresolvable vendor_id or invoice_id in task_data does not +escape the agent loop's error handling, that prompt construction degrades +gracefully when database entities are absent, that setup exceptions +in BaseAgent._run_agent_loop follow the established task completion path, +and that resource cleanup in finally executes reliably even if MCP disconnect fails. +""" + +from datetime import UTC, datetime, timedelta +from typing import Any, Callable, Generator +from unittest.mock import AsyncMock, patch + +import pytest + +from finbot.agents.base import BaseAgent +from finbot.agents.specialized.invoice import InvoiceAgent +from finbot.agents.specialized.onboarding import VendorOnboardingAgent +from finbot.core.auth.session import SessionContext, session_manager +from finbot.core.data.database import db_session +from finbot.core.data.models import Invoice, LLMResponse, Vendor +from finbot.core.data.repositories import VendorRepository +from finbot.core.messaging import event_bus + + +@pytest.fixture(autouse=True) +def mock_event_bus_emission() -> Generator[None, None, None]: + """Mock async Redis event emissions so unit tests are isolated and offline.""" + with ( + patch.object(event_bus, "emit_agent_event", new_callable=AsyncMock), + patch.object(event_bus, "emit_business_event", new_callable=AsyncMock), + ): + yield + + +def _create_session_context(email: str) -> SessionContext: + """Helper to create SessionContext for testing.""" + session = session_manager.create_session( + email=email, + user_agent="AgentErrorHandlingTest/1.0", + ) + created_at = datetime.now(UTC) + expires_at = created_at + timedelta(hours=24) + + return SessionContext( + session_id=session.session_id, + user_id=f"user_{email.split('@')[0]}", + email=email, + namespace=f"user_{email.split('@')[0]}", + is_temporary=False, + created_at=created_at, + expires_at=expires_at, + ) + + +# ============================================================================ +# VendorOnboardingAgent Prompt & Lookup Error Handling Tests +# ============================================================================ + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_vendor_prompt_with_valid_vendor_id() -> None: + """Valid seeded vendor ID populates services and agent_notes in prompt.""" + session_context = _create_session_context("valid_vendor@example.com") + with db_session() as db: + repo = VendorRepository(db, session_context) + vendor = repo.create_vendor( + company_name="Alpha Tech", + vendor_category="Technology", + industry="Software", + services="Cloud Hosting and Security", + contact_name="Alice Alpha", + email="alice@alphatech.com", + tin="11-1111111", + bank_account_number="123456789012", + bank_name="Test Bank", + bank_routing_number="021000021", + bank_account_holder_name="Alice Alpha", + ) + vendor.agent_notes = "Verified tier-1 strategic partner" # type: ignore[assignment] + db.commit() + vendor_id = vendor.id + + agent = VendorOnboardingAgent(session_context=session_context) + prompt = await agent._get_user_prompt({"vendor_id": vendor_id}) + + assert "Cloud Hosting and Security" in prompt + assert "Verified tier-1 strategic partner" in prompt + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_vendor_prompt_with_missing_integer_vendor_id() -> None: + """Non-existent integer vendor ID does not raise ValueError; prompt is generated.""" + session_context = _create_session_context("missing_vendor_int@example.com") + agent = VendorOnboardingAgent(session_context=session_context) + + # Vendor 999999 does not exist in this namespace + prompt = await agent._get_user_prompt({"vendor_id": 999999}) + + assert isinstance(prompt, str) + assert len(prompt) > 0 + assert "vendor" in prompt.lower() + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_vendor_prompt_with_string_vendor_id() -> None: + """Non-integer string vendor ID does not raise ValueError; prompt is generated.""" + session_context = _create_session_context("missing_vendor_str@example.com") + agent = VendorOnboardingAgent(session_context=session_context) + + prompt = await agent._get_user_prompt({"vendor_id": "VEND-001"}) + + assert isinstance(prompt, str) + assert len(prompt) > 0 + assert "VEND-001" in prompt + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_vendor_process_with_unresolvable_vendor_id() -> None: + """VendorOnboardingAgent.process() with unresolvable vendor ID passes prompt to LLM and completes.""" + session_context = _create_session_context("vendor_process_missing@example.com") + agent = VendorOnboardingAgent(session_context=session_context) + + mock_llm_response = LLMResponse( + content=None, + tool_calls=[ + { + "name": "complete_task", + "call_id": "call_vendor_001", + "arguments": { + "task_status": "failed", + "task_summary": "Vendor identifier VEND-001 could not be resolved in the database.", + }, + } + ], + ) + + with patch.object( + agent.llm_client, "chat", new_callable=AsyncMock, return_value=mock_llm_response + ) as mock_chat: + result = await agent.process({"action": "get_progress", "vendor_id": "VEND-001"}) + + # Verify LLM was invoked with the constructed prompt containing the identifier + mock_chat.assert_awaited_once() + assert mock_chat.await_args is not None + called_request = mock_chat.await_args.kwargs["request"] + user_message = next(m for m in called_request.messages if m["role"] == "user") + assert "VEND-001" in user_message["content"] + + # Verify task completion was returned from the tool call rather than a crash or stall + assert isinstance(result, dict) + assert result["task_status"] == "failed" + assert ( + result["task_summary"] + == "Vendor identifier VEND-001 could not be resolved in the database." + ) + + +# ============================================================================ +# InvoiceAgent Prompt & Lookup Error Handling Tests +# ============================================================================ + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_invoice_prompt_with_valid_invoice_id() -> None: + """Valid seeded invoice ID populates description and agent_notes in prompt.""" + session_context = _create_session_context("valid_invoice@example.com") + with db_session() as db: + vendor_repo = VendorRepository(db, session_context) + vendor = vendor_repo.create_vendor( + company_name="Beta Services", + vendor_category="Consulting", + industry="Finance", + services="Advisory", + contact_name="Bob Beta", + email="bob@betaservices.com", + tin="22-2222222", + bank_account_number="987654321098", + bank_name="Beta Bank", + bank_routing_number="021000021", + bank_account_holder_name="Bob Beta", + ) + invoice = Invoice( + vendor_id=vendor.id, + invoice_number="INV-2026-001", + amount=5000.0, + description="Q3 Security Architecture Review", + agent_notes="Pre-approved under MSA 2026", + status="pending_approval", + namespace=session_context.namespace, + invoice_date=datetime.now(UTC) - timedelta(days=5), + due_date=datetime.now(UTC) + timedelta(days=25), + ) + db.add(invoice) + db.commit() + db.refresh(invoice) + invoice_id = invoice.id + + agent = InvoiceAgent(session_context=session_context) + prompt = await agent._get_user_prompt({"invoice_id": invoice_id}) + + assert "Q3 Security Architecture Review" in prompt + assert "Pre-approved under MSA 2026" in prompt + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_invoice_prompt_with_missing_integer_invoice_id() -> None: + """Non-existent integer invoice ID does not raise ValueError; prompt is generated.""" + session_context = _create_session_context("missing_invoice_int@example.com") + agent = InvoiceAgent(session_context=session_context) + + # Invoice 999999 does not exist in this namespace + prompt = await agent._get_user_prompt({"invoice_id": 999999}) + + assert isinstance(prompt, str) + assert len(prompt) > 0 + assert "invoice" in prompt.lower() + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_invoice_prompt_with_string_invoice_id() -> None: + """Non-integer string invoice ID does not raise ValueError; prompt is generated.""" + session_context = _create_session_context("missing_invoice_str@example.com") + agent = InvoiceAgent(session_context=session_context) + + prompt = await agent._get_user_prompt({"invoice_id": "INV-001"}) + + assert isinstance(prompt, str) + assert len(prompt) > 0 + assert "INV-001" in prompt + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_invoice_process_with_unresolvable_invoice_id() -> None: + """InvoiceAgent.process() with unresolvable invoice ID passes prompt to LLM and completes.""" + session_context = _create_session_context("invoice_process_missing@example.com") + agent = InvoiceAgent(session_context=session_context) + + mock_llm_response = LLMResponse( + content=None, + tool_calls=[ + { + "name": "complete_task", + "call_id": "call_invoice_001", + "arguments": { + "task_status": "failed", + "task_summary": "Invoice identifier INV-001 could not be resolved.", + }, + } + ], + ) + + with patch.object( + agent.llm_client, "chat", new_callable=AsyncMock, return_value=mock_llm_response + ) as mock_chat: + result = await agent.process({"action": "review_invoice", "invoice_id": "INV-001"}) + + # Verify LLM was invoked with the constructed prompt containing the identifier + mock_chat.assert_awaited_once() + assert mock_chat.await_args is not None + called_request = mock_chat.await_args.kwargs["request"] + user_message = next(m for m in called_request.messages if m["role"] == "user") + assert "INV-001" in user_message["content"] + + # Verify task completion was returned from the tool call rather than a crash or stall + assert isinstance(result, dict) + assert result["task_status"] == "failed" + assert result["task_summary"] == "Invoice identifier INV-001 could not be resolved." + + +# ============================================================================ +# BaseAgent Setup Exception Handling & MCP Cleanup Tests +# ============================================================================ + + +class FailingSetupAgent(BaseAgent): + """Test agent whose prompt construction fails unexpectedly.""" + + def _load_config(self) -> dict[str, Any]: + return {} + + def _get_system_prompt(self) -> str: + return "System prompt for failing setup agent" + + async def _get_user_prompt(self, task_data: dict[str, Any] | None = None) -> str: + raise RuntimeError("Simulated unexpected prompt construction failure") + + def _get_tool_definitions(self) -> list[dict[str, Any]]: + return [] + + def _get_callables(self) -> dict[str, Callable[..., Any]]: + return {} + + async def process(self, task_data: dict[str, Any], **kwargs: Any) -> dict[str, Any]: + return await self._run_agent_loop(task_data=task_data) + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_base_agent_catches_setup_exception_and_returns_failed_status() -> None: + """Unexpected exception in _get_user_prompt is caught by setup handler and returns failed task.""" + session_context = _create_session_context("setup_failure@example.com") + agent = FailingSetupAgent( + session_context=session_context, + workflow_id="wf_setup_fail_test", + ) + + result = await agent.process({"action": "do_something"}) + + assert isinstance(result, dict) + assert result["task_status"] == "failed" + assert "Agent setup failed" in result["task_summary"] + assert "Simulated unexpected prompt construction failure" in result["task_summary"] + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_base_agent_cleanup_on_setup_exception() -> None: + """MCP servers and workflow context are cleaned up even when setup fails.""" + session_context = _create_session_context("cleanup_on_failure@example.com") + agent = FailingSetupAgent( + session_context=session_context, + workflow_id="wf_cleanup_test", + ) + + with patch.object(agent, "_disconnect_mcp_servers", new_callable=AsyncMock) as mock_disconnect: + result = await agent.process({"action": "do_something"}) + + assert result["task_status"] == "failed" + mock_disconnect.assert_awaited_once() + + # Workflow context should be cleared + assert "wf_cleanup_test" not in event_bus._workflow_ctx + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_base_agent_cleanup_when_mcp_disconnect_fails() -> None: + """Workflow context is cleared even if MCP disconnection raises an exception.""" + session_context = _create_session_context("mcp_disconnect_fail@example.com") + agent = FailingSetupAgent( + session_context=session_context, + workflow_id="wf_mcp_fail_test", + ) + + with patch.object( + agent, + "_disconnect_mcp_servers", + new_callable=AsyncMock, + side_effect=RuntimeError("Simulated MCP disconnect crash"), + ): + result = await agent.process({"action": "do_something"}) + + assert result["task_status"] == "failed" + assert "Simulated unexpected prompt construction failure" in result["task_summary"] + + # Workflow context MUST be cleared despite MCP disconnect failure + assert "wf_mcp_fail_test" not in event_bus._workflow_ctx + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_base_agent_stall_termination() -> None: + """Agent terminates with failed status after max stall iterations when no tools called.""" + session_context = _create_session_context("stall_test@example.com") + agent = VendorOnboardingAgent(session_context=session_context) + + # Return text-only response with no tool calls to deterministically trigger stall detection + text_only_response = LLMResponse(content="Just thinking out loud...", tool_calls=[]) + + with patch.object( + agent.llm_client, "chat", new_callable=AsyncMock, return_value=text_only_response + ) as mock_chat: + result = await agent.process({}) + + assert isinstance(result, dict) + assert result["task_status"] == "failed" + assert "stalled" in result["task_summary"].lower() + # Verify it stalled after exactly 2 iterations (max_stall_iterations) + assert mock_chat.await_count == 2