diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 2a817d32..1979f746 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -3882,6 +3882,8 @@ class AgentManager: f"\n{user_prompt[:2000]}\n" ), }], + # Binds this aux call to its query's free-trial run; ignored off the free lane. + extra_headers={"X-Openswarm-Task-Id": session_id}, ) as stream: async for text in stream.text_stream: chunks.append(text) @@ -3936,6 +3938,8 @@ class AgentManager: max_tokens=1, system="You are a helpful assistant. Reply with one character.", messages=[{"role": "user", "content": "ping"}], + # Binds the cache-warm ping to its query's free-trial run; ignored off the free lane. + extra_headers={"X-Openswarm-Task-Id": session_id}, ) logger.debug(f"Cache pre-warm fired for session {session_id}") except Exception as e: @@ -4008,6 +4012,8 @@ class AgentManager: max_tokens=aux_max_tokens_for(aux_model, base=300), system=system, messages=[{"role": "user", "content": user_content}], + # Binds this aux call to its query's free-trial run; ignored off the free lane. + extra_headers={"X-Openswarm-Task-Id": session_id}, ) as stream: async for text in stream.text_stream: chunks.append(text) diff --git a/backend/apps/agents/agents.py b/backend/apps/agents/agents.py index 1cab53f8..cec4a8c6 100644 --- a/backend/apps/agents/agents.py +++ b/backend/apps/agents/agents.py @@ -80,7 +80,7 @@ async def send_message(session_id: str, body: dict): async def _emit_preflight(): try: - result = await run_preflight(prompt) + result = await run_preflight(prompt, task_id=session_id) if result.get("suggestions") or result.get("is_vague"): await _ws.send_to_session(session_id, "agent:mcp_suggestions", { "session_id": session_id, diff --git a/backend/apps/agents/core/mcp_preflight.py b/backend/apps/agents/core/mcp_preflight.py index 9de5ca70..fad2c739 100644 --- a/backend/apps/agents/core/mcp_preflight.py +++ b/backend/apps/agents/core/mcp_preflight.py @@ -85,7 +85,7 @@ def _is_obviously_local(prompt: str) -> bool: return False -async def run_preflight(prompt: str, timeout_s: float = 2.0) -> dict: +async def run_preflight(prompt: str, timeout_s: float = 2.0, task_id: str | None = None) -> dict: """Classify the prompt and return {is_vague, suggestions}; never raises.""" default: dict[str, Any] = {"is_vague": False, "suggestions": []} @@ -100,7 +100,7 @@ async def run_preflight(prompt: str, timeout_s: float = 2.0) -> dict: available = _build_available_shortlist(settings) result = await asyncio.wait_for( - _call_classifier(settings, prompt, available), + _call_classifier(settings, prompt, available, task_id), timeout=timeout_s, ) # Re-validate ids against the curated shortlist so hallucinations can't reach the frontend. @@ -151,7 +151,7 @@ def _decorate(llm_suggestion: dict, available: list[CuratedEntry]) -> dict | Non } -async def _call_classifier(settings, prompt: str, available: list[CuratedEntry]) -> dict: +async def _call_classifier(settings, prompt: str, available: list[CuratedEntry], task_id: str | None = None) -> dict: """One aux-model call, returns validated JSON {is_vague, suggestions}.""" aux_model, _base = await resolve_aux_model(settings, preferred_tier="haiku") client = get_anthropic_client_for_model(settings, aux_model) @@ -191,6 +191,8 @@ async def _call_classifier(settings, prompt: str, available: list[CuratedEntry]) max_tokens=300, system=system, messages=[{"role": "user", "content": user_turn}], + # Rides on its query's free-trial run instead of opening its own; ignored off the free lane. + extra_headers={"X-Openswarm-Task-Id": task_id} if task_id else {}, ) text = ""