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

@@ -19,7 +19,6 @@ import logging
import os
from typing import Any, Literal, Optional
from mcp.types import TextContent
from pydantic import Field
from config import TEMPERATURE_ANALYTICAL
@@ -27,7 +26,6 @@ from systemprompts import REFACTOR_PROMPT
from utils.file_utils import translate_file_paths
from .base import BaseTool, ToolRequest
from .models import ToolOutput
logger = logging.getLogger(__name__)
@@ -154,25 +152,6 @@ class RefactorTool(BaseTool):
def get_request_model(self):
return RefactorRequest
async def execute(self, arguments: dict[str, Any]) -> list[TextContent]:
"""Override execute to check prompt size before processing"""
logger.info(f"[REFACTOR] execute called with arguments: {list(arguments.keys())}")
# First validate request
request_model = self.get_request_model()
request = request_model(**arguments)
# Check prompt size if provided
if request.prompt:
size_check = self.check_prompt_size(request.prompt)
if size_check:
logger.info("[REFACTOR] Prompt size check triggered, returning early")
return [TextContent(type="text", text=ToolOutput(**size_check).model_dump_json())]
logger.info("[REFACTOR] Prompt size OK, calling super().execute()")
# Continue with normal execution
return await super().execute(arguments)
def detect_primary_language(self, file_paths: list[str]) -> str:
"""
Detect the primary programming language from file extensions.
@@ -417,6 +396,14 @@ class RefactorTool(BaseTool):
logger.debug(f"[REFACTOR] Updated files list after prompt.txt processing: {len(updated_files)} files")
request.files = updated_files
# Check user input size at MCP transport boundary (before adding internal content)
user_content = request.prompt
size_check = self.check_prompt_size(user_content)
if size_check:
from tools.models import ToolOutput
raise ValueError(f"MCP_SIZE_CHECK:{ToolOutput(**size_check).model_dump_json()}")
# Calculate available token budget for dynamic allocation
continuation_id = getattr(request, "continuation_id", None)