Perform prompt size checks only at the MCP boundary

New test to confirm history build-up and system prompt does not affect prompt size checks
Also check for large prompts in focus_on
Fixed .env.example incorrectly did not comment out CUSTOM_API causing the run-server script to think at least one key exists
This commit is contained in:
Fahad
2025-06-15 10:37:08 +04:00
parent 3b03783ea7
commit 4becd70a82
14 changed files with 404 additions and 198 deletions

View File

@@ -4,7 +4,6 @@ Chat tool - General development chat and collaborative thinking
from typing import TYPE_CHECKING, Any, Optional
from mcp.types import TextContent
from pydantic import Field
if TYPE_CHECKING:
@@ -14,7 +13,6 @@ from config import TEMPERATURE_BALANCED
from systemprompts import CHAT_PROMPT
from .base import BaseTool, ToolRequest
from .models import ToolOutput
class ChatRequest(ToolRequest):
@@ -102,20 +100,6 @@ class ChatTool(BaseTool):
def get_request_model(self):
return ChatRequest
async def execute(self, arguments: dict[str, Any]) -> list[TextContent]:
"""Override execute to check prompt size before processing"""
# First validate request
request_model = self.get_request_model()
request = request_model(**arguments)
# Check prompt size
size_check = self.check_prompt_size(request.prompt)
if size_check:
return [TextContent(type="text", text=ToolOutput(**size_check).model_dump_json())]
# Continue with normal execution
return await super().execute(arguments)
async def prepare_prompt(self, request: ChatRequest) -> str:
"""Prepare the chat prompt with optional context files"""
# Check for prompt.txt in files
@@ -124,6 +108,16 @@ class ChatTool(BaseTool):
# Use prompt.txt content if available, otherwise use the prompt field
user_content = prompt_content if prompt_content else request.prompt
# Check user input size at MCP transport boundary (before adding internal content)
size_check = self.check_prompt_size(user_content)
if size_check:
# Need to return error, but prepare_prompt returns str
# Use exception to handle this cleanly
from tools.models import ToolOutput
raise ValueError(f"MCP_SIZE_CHECK:{ToolOutput(**size_check).model_dump_json()}")
# Update request files list
if updated_files is not None:
request.files = updated_files