From ba889b0ae0726191b43a76a7cd9e85ad31eff859 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=A2=81?= Date: Wed, 3 Jun 2026 01:22:23 +0800 Subject: [PATCH] fix: add tool count validation with env-configurable MAX_TOOLS Prevents cryptic 'There was an issue with your request' errors from OpenAI's API when too many functions are selected (e.g., enabling all Gmail functions in the playground). Changes: - Add MAX_TOOLS=64 default in config.py, overridable via SERVER_MAX_TOOLS - Add ToolCountExceeded exception with actionable 400 response - Add early validation in /chat route before any API call is made - Add defensive warning log in openai_chat_stream for other call-sites - Add comprehensive tests for exception, config, and /chat endpoint Closes #337 Co-Authored-By: Claude Opus 4.8 --- backend/aci/common/exceptions.py | 18 ++ backend/aci/server/agent/prompt.py | 11 + backend/aci/server/config.py | 9 + backend/aci/server/routes/agent.py | 5 + .../routes/agent/test_chat_tool_validation.py | 283 ++++++++++++++++++ 5 files changed, 326 insertions(+) create mode 100644 backend/aci/server/tests/routes/agent/test_chat_tool_validation.py diff --git a/backend/aci/common/exceptions.py b/backend/aci/common/exceptions.py index 75056b1f2..983b0a556 100644 --- a/backend/aci/common/exceptions.py +++ b/backend/aci/common/exceptions.py @@ -524,3 +524,21 @@ def __init__(self, message: str | None = None): message=message, error_code=status.HTTP_400_BAD_REQUEST, ) + + +class ToolCountExceeded(ACIException): + """ + Exception raised when the number of tools in a chat request exceeds the configured limit. + """ + + def __init__(self, tool_count: int, max_tools: int): + message = ( + f"Too many functions selected ({tool_count}). " + f"The maximum allowed is {max_tools}. " + f"Please reduce the number of selected functions." + ) + super().__init__( + title="Tool count exceeded", + message=message, + error_code=status.HTTP_400_BAD_REQUEST, + ) diff --git a/backend/aci/server/agent/prompt.py b/backend/aci/server/agent/prompt.py index 6d9e68c64..e19188322 100644 --- a/backend/aci/server/agent/prompt.py +++ b/backend/aci/server/agent/prompt.py @@ -72,6 +72,17 @@ async def openai_chat_stream( """ client = OpenAI(api_key=config.OPENAI_API_KEY) + # Defensive check: log a warning if tool count exceeds the configured limit. + # The primary guard is in the /chat route, but this catch-all ensures we never + # silently send an oversized tool payload to the OpenAI API from other call-sites. + if len(tools) > config.MAX_TOOLS: + logger.warning( + "openai_chat_stream called with %d tools, exceeding MAX_TOOLS=%d. " + "This may cause OpenAI API errors.", + len(tools), + config.MAX_TOOLS, + ) + # TODO: support different meta function mode ACI_META_FUNCTIONS_SCHEMA_LIST stream = client.responses.create(model="gpt-4o", input=messages, stream=True, tools=tools) diff --git a/backend/aci/server/config.py b/backend/aci/server/config.py index 0bff53817..bcde05518 100644 --- a/backend/aci/server/config.py +++ b/backend/aci/server/config.py @@ -1,3 +1,5 @@ +import os + from aci.common.utils import check_and_get_env_variable, construct_db_url ENVIRONMENT = check_and_get_env_variable("SERVER_ENVIRONMENT") @@ -92,3 +94,10 @@ VECTOR_DB_FULL_URL = check_and_get_env_variable("SERVER_VECTOR_DB_FULL_URL") SENTRY_DSN = check_and_get_env_variable("SERVER_SENTRY_DSN") + +# Maximum number of tools allowed per chat request. +# OpenAI's API has undocumented limits on tool count; exceeding them produces a +# cryptic "There was an issue with your request" error. Default 64 is conservative +# (the actual limit is closer to 128 or ~200k tokens of tool definitions). +# Override via SERVER_MAX_TOOLS. +MAX_TOOLS = int(os.getenv("SERVER_MAX_TOOLS", "64")) diff --git a/backend/aci/server/routes/agent.py b/backend/aci/server/routes/agent.py index c3adf9541..79d24a5a3 100644 --- a/backend/aci/server/routes/agent.py +++ b/backend/aci/server/routes/agent.py @@ -6,6 +6,7 @@ from pydantic import BaseModel from aci.common.enums import FunctionDefinitionFormat +from aci.common.exceptions import ToolCountExceeded from aci.common.logging_setup import get_logger from aci.common.schemas.function import OpenAIResponsesFunctionDefinition from aci.server import config @@ -62,6 +63,10 @@ async def handle_chat( func for func in selected_functions if isinstance(func, OpenAIResponsesFunctionDefinition) ] + # Validate tool count before making the API call to avoid cryptic OpenAI errors. + if len(tools) > config.MAX_TOOLS: + raise ToolCountExceeded(tool_count=len(tools), max_tools=config.MAX_TOOLS) + response = StreamingResponse(openai_chat_stream(openai_messages, tools=tools)) response.headers["x-vercel-ai-data-stream"] = "v1" diff --git a/backend/aci/server/tests/routes/agent/test_chat_tool_validation.py b/backend/aci/server/tests/routes/agent/test_chat_tool_validation.py new file mode 100644 index 000000000..88ed5bcff --- /dev/null +++ b/backend/aci/server/tests/routes/agent/test_chat_tool_validation.py @@ -0,0 +1,283 @@ +"""Tests for tool-count validation on the /chat endpoint.""" + +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi import status +from fastapi.testclient import TestClient + +from aci.common.exceptions import ToolCountExceeded +from aci.common.schemas.function import OpenAIResponsesFunctionDefinition +from aci.server import config + + +# --------------------------------------------------------------------------- +# Helper: build a list of dummy tool definitions of a given size +# --------------------------------------------------------------------------- +def _make_tools(count: int) -> list[OpenAIResponsesFunctionDefinition]: + return [ + OpenAIResponsesFunctionDefinition( + name=f"test_tool_{i}", + description=f"Test tool number {i}", + parameters={ + "type": "object", + "properties": { + "arg": {"type": "string", "description": "A dummy argument"}, + }, + "required": ["arg"], + }, + ) + for i in range(count) + ] + + +# --------------------------------------------------------------------------- +# Unit: ToolCountExceeded exception +# --------------------------------------------------------------------------- +class TestToolCountExceededException: + """Verify the exception carries the right shape and content.""" + + def test_error_code_is_400(self) -> None: + exc = ToolCountExceeded(tool_count=70, max_tools=64) + assert exc.error_code == status.HTTP_400_BAD_REQUEST + + def test_title_is_descriptive(self) -> None: + exc = ToolCountExceeded(tool_count=70, max_tools=64) + assert exc.title == "Tool count exceeded" + + def test_message_includes_actual_tool_count(self) -> None: + exc = ToolCountExceeded(tool_count=70, max_tools=64) + assert "70" in exc.message + + def test_message_includes_max_tools_limit(self) -> None: + exc = ToolCountExceeded(tool_count=70, max_tools=64) + assert "64" in exc.message + + def test_message_is_actionable(self) -> None: + exc = ToolCountExceeded(tool_count=100, max_tools=50) + assert "Too many functions" in exc.message + assert "reduce" in exc.message.lower() + + def test_expands_to_full_string(self) -> None: + exc = ToolCountExceeded(tool_count=3, max_tools=2) + as_str = str(exc) + assert as_str == ( + "Tool count exceeded: " + "Too many functions selected (3). The maximum allowed is 2. " + "Please reduce the number of selected functions." + ) + + +# --------------------------------------------------------------------------- +# Config: MAX_TOOLS +# --------------------------------------------------------------------------- +class TestMaxToolsConfig: + """Verify MAX_TOOLS has a sensible default.""" + + def test_default_is_64(self) -> None: + assert config.MAX_TOOLS == 64 + + def test_is_positive_integer(self) -> None: + assert isinstance(config.MAX_TOOLS, int) + assert config.MAX_TOOLS > 0 + + +# --------------------------------------------------------------------------- +# Helpers for integration tests +# --------------------------------------------------------------------------- +@pytest.fixture(scope="function") +def test_client_with_chat_mocks(test_client: TestClient) -> TestClient: + """Override the request-context dependency so /chat doesn't need + a full DB / API-key setup. We also mock `openai_chat_stream` so no + real OpenAI call is made.""" + + class _Stub: + pass + + stub_project = _Stub() + stub_project.id = "00000000-0000-0000-0000-000000000001" + + stub_agent = _Stub() + stub_agent.id = "00000000-0000-0000-0000-000000000002" + + from aci.server.dependencies import RequestContext + + class _FakeDBSession: + def close(self) -> None: + pass + + def _fake_context() -> RequestContext: + return RequestContext( + db_session=_FakeDBSession(), # type: ignore[arg-type] + api_key_id="00000000-0000-0000-0000-000000000003", # type: ignore[arg-type] + project=stub_project, # type: ignore[arg-type] + agent=stub_agent, # type: ignore[arg-type] + ) + + import aci.server.dependencies as deps + + test_client.app.dependency_overrides[deps.get_request_context] = _fake_context + + yield test_client + + test_client.app.dependency_overrides.clear() + + +# --------------------------------------------------------------------------- +# Integration: /chat endpoint tool-count validation +# --------------------------------------------------------------------------- +class TestChatToolValidation: + """End-to-end tests for tool-count validation in the /chat route.""" + + def test_returns_400_when_tools_exceed_limit( + self, test_client_with_chat_mocks: TestClient + ) -> None: + """When the resolved tool list is larger than MAX_TOOLS, the route + must return 400 with a clear error message.""" + too_many = _make_tools(config.MAX_TOOLS + 1) + + with patch( + "aci.server.routes.agent.get_functions_definitions", + new=AsyncMock(return_value=too_many), + ): + payload = { + "id": "test-chat-id", + "linked_account_owner_id": "test_user", + "selected_apps": [], + "selected_functions": [f"test_tool_{i}" for i in range(len(too_many))], + "messages": [], + } + response = test_client_with_chat_mocks.post( + f"{config.ROUTER_PREFIX_AGENT}/chat", json=payload + ) + + assert response.status_code == status.HTTP_400_BAD_REQUEST + body = response.json() + assert "error" in body + assert "Tool count exceeded" in body["error"] + assert str(config.MAX_TOOLS) in body["error"] + + def test_returns_400_when_tools_exceed_custom_limit( + self, test_client_with_chat_mocks: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + """When MAX_TOOLS is overridden via env var, the custom limit is + used in the error message.""" + monkeypatch.setattr(config, "MAX_TOOLS", 10) + too_many = _make_tools(11) + + with patch( + "aci.server.routes.agent.get_functions_definitions", + new=AsyncMock(return_value=too_many), + ): + payload = { + "id": "test-chat-id", + "linked_account_owner_id": "test_user", + "selected_apps": [], + "selected_functions": [f"test_tool_{i}" for i in range(len(too_many))], + "messages": [], + } + response = test_client_with_chat_mocks.post( + f"{config.ROUTER_PREFIX_AGENT}/chat", json=payload + ) + + assert response.status_code == status.HTTP_400_BAD_REQUEST + body = response.json() + assert "11" in body["error"] + assert "10" in body["error"] + + def test_succeeds_when_tools_exactly_at_limit( + self, test_client_with_chat_mocks: TestClient + ) -> None: + """When the tool count exactly equals MAX_TOOLS, validation passes.""" + exact = _make_tools(config.MAX_TOOLS) + + async def _fake_stream(*args, **kwargs): # noqa: ARG001 + yield 'd:{"finishReason":"stop","usage":{"promptTokens":10,"completionTokens":5}}\n' + + with ( + patch( + "aci.server.routes.agent.get_functions_definitions", + new=AsyncMock(return_value=exact), + ), + patch("aci.server.routes.agent.openai_chat_stream", new=_fake_stream), + ): + payload = { + "id": "test-chat-id", + "linked_account_owner_id": "test_user", + "selected_apps": [], + "selected_functions": [f"test_tool_{i}" for i in range(len(exact))], + "messages": [], + } + response = test_client_with_chat_mocks.post( + f"{config.ROUTER_PREFIX_AGENT}/chat", json=payload + ) + + assert response.status_code == status.HTTP_200_OK + + def test_succeeds_when_tools_under_limit( + self, test_client_with_chat_mocks: TestClient + ) -> None: + """A normal request with a handful of tools streams normally.""" + few = _make_tools(3) + + async def _fake_stream(*args, **kwargs): # noqa: ARG001 + yield 'd:{"finishReason":"stop","usage":{"promptTokens":10,"completionTokens":5}}\n' + + with ( + patch( + "aci.server.routes.agent.get_functions_definitions", + new=AsyncMock(return_value=few), + ), + patch("aci.server.routes.agent.openai_chat_stream", new=_fake_stream), + ): + payload = { + "id": "test-chat-id", + "linked_account_owner_id": "test_user", + "selected_apps": [], + "selected_functions": ["test_tool_0", "test_tool_1", "test_tool_2"], + "messages": [], + } + response = test_client_with_chat_mocks.post( + f"{config.ROUTER_PREFIX_AGENT}/chat", json=payload + ) + + assert response.status_code == status.HTTP_200_OK + + def test_succeeds_when_no_openai_response_tools( + self, test_client_with_chat_mocks: TestClient + ) -> None: + """When selected functions produce zero OpenAIResponsesFunctionDefinition + tools (e.g. all are non-OpenAI formats), the request passes through.""" + from aci.common.schemas.function import OpenAIFunctionDefinition + + non_openai_tools = [ + OpenAIFunctionDefinition( + name=f"openai_tool_{i}", + description=f"OpenAI (non-responses) tool {i}", + parameters={"type": "object", "properties": {}}, + ) + for i in range(70) + ] + + async def _fake_stream(*args, **kwargs): # noqa: ARG001 + yield 'd:{"finishReason":"stop","usage":{"promptTokens":10,"completionTokens":5}}\n' + + with ( + patch( + "aci.server.routes.agent.get_functions_definitions", + new=AsyncMock(return_value=non_openai_tools), + ), + patch("aci.server.routes.agent.openai_chat_stream", new=_fake_stream), + ): + payload = { + "id": "test-chat-id", + "linked_account_owner_id": "test_user", + "selected_apps": [], + "selected_functions": ["non_openai_responses_tool"] * 70, + "messages": [], + } + response = test_client_with_chat_mocks.post( + f"{config.ROUTER_PREFIX_AGENT}/chat", json=payload + ) + + assert response.status_code == status.HTTP_200_OK