From 7b1cbed2f8b3ae176545f24416b526b419b45edc Mon Sep 17 00:00:00 2001 From: coderhisham Date: Thu, 25 Jun 2026 12:31:13 +0530 Subject: [PATCH 1/6] feat(models): add Claude Opus 4.8 (1M context) and configurable context windows (#208) Add claude-opus-4.8 to the fallback model list with a 1M maxInputTokens so it is recognized on the runtime.kiro.dev endpoint (which has no /ListAvailableModels). Make context-window accounting configurable: DEFAULT_MAX_INPUT_TOKENS is now an env var, and MODEL_CONTEXT_WINDOWS lets users declare per-model windows (e.g. 1M) that take precedence over discovered/default limits. This only affects token-usage estimation; requests remain pass-through. Adds tests for the override resolution chain, the env parser, and the opus-4.8 fallback entry. --- .env.example | 19 ++++ kiro/cache.py | 19 +++- kiro/config.py | 88 +++++++++++++++++- tests/unit/test_cache.py | 109 ++++++++++++++++++++++ tests/unit/test_config.py | 149 +++++++++++++++++++++++++++++- tests/unit/test_model_resolver.py | 23 +++++ 6 files changed, 398 insertions(+), 9 deletions(-) diff --git a/.env.example b/.env.example index 1a57c794..da93321b 100644 --- a/.env.example +++ b/.env.example @@ -171,6 +171,25 @@ PROXY_API_KEY="my-super-secret-password-123" # Interval for periodic state.json saving in seconds # STATE_SAVE_INTERVAL_SECONDS=10 +# =========================================== +# CONTEXT WINDOW (Token Accounting) +# =========================================== + +# Default max input tokens used for token-usage estimation when the upstream +# model limit is unknown. This is NOT an enforced cap - requests are always +# passed through to Kiro, which decides the real limit. It only affects the +# input_tokens reported in usage (context_usage_percentage * this value). +# Default: 200000 +# DEFAULT_MAX_INPUT_TOKENS=200000 + +# Per-model context window overrides (JSON object: model id -> max input tokens). +# Highest precedence - wins over both the discovered /ListAvailableModels value +# and DEFAULT_MAX_INPUT_TOKENS. Use this to declare a model's real window (e.g. +# 1M on Kiro) when the gateway cannot discover it - notably on the +# runtime.kiro.dev endpoint, which does not provide /ListAvailableModels. +# Keys are Kiro model ids (after normalization), e.g. "claude-sonnet-4". +# MODEL_CONTEXT_WINDOWS='{"claude-sonnet-4": 1000000, "claude-opus-4.7": 1000000}' + # =========================================== # FIRST TOKEN TIMEOUT (Streaming Retry) # =========================================== diff --git a/kiro/cache.py b/kiro/cache.py index f0be72af..40307ebb 100644 --- a/kiro/cache.py +++ b/kiro/cache.py @@ -30,7 +30,7 @@ from loguru import logger -from kiro.config import MODEL_CACHE_TTL, DEFAULT_MAX_INPUT_TOKENS +from kiro.config import MODEL_CACHE_TTL, DEFAULT_MAX_INPUT_TOKENS, MODEL_CONTEXT_WINDOWS class ModelInfoCache: @@ -129,13 +129,24 @@ def add_hidden_model(self, display_name: str, internal_id: str) -> None: def get_max_input_tokens(self, model_id: str) -> int: """ Returns maxInputTokens for the model. - + + Resolution order (highest precedence first): + 1. Explicit per-model override from MODEL_CONTEXT_WINDOWS env config + (lets users declare a real window, e.g. 1M, when the gateway cannot + discover it - notably on the runtime.kiro.dev endpoint). + 2. Upstream value from the cached /ListAvailableModels tokenLimits. + 3. DEFAULT_MAX_INPUT_TOKENS fallback. + Args: model_id: Model ID - + Returns: - Maximum number of input tokens or DEFAULT_MAX_INPUT_TOKENS + Maximum number of input tokens """ + override = MODEL_CONTEXT_WINDOWS.get(model_id) + if override: + return override + model = self._cache.get(model_id) if model and model.get("tokenLimits"): return model["tokenLimits"].get("maxInputTokens") or DEFAULT_MAX_INPUT_TOKENS diff --git a/kiro/config.py b/kiro/config.py index e0f3a527..911179d9 100644 --- a/kiro/config.py +++ b/kiro/config.py @@ -26,9 +26,11 @@ import os import re +import json from pathlib import Path -from typing import Dict, List, Optional +from typing import Any, Dict, List, Optional from dotenv import load_dotenv +from loguru import logger # Load environment variables load_dotenv() @@ -273,7 +275,12 @@ def _get_raw_env_value(var_name: str, env_file: str = ".env") -> Optional[str]: # - Some models may not be available on your Kiro plan (e.g., Opus on free tier) # - New models released after this version won't appear here # - Update gateway regularly to get the latest model list -FALLBACK_MODELS: List[Dict[str, str]] = [ +# +# tokenLimits.maxInputTokens is included for models with a non-default context +# window so token accounting is accurate on the runtime.kiro.dev endpoint, which +# does not expose /ListAvailableModels. Models without it fall back to +# DEFAULT_MAX_INPUT_TOKENS (and can still be overridden via MODEL_CONTEXT_WINDOWS). +FALLBACK_MODELS: List[Dict[str, Any]] = [ {"modelId": "auto"}, {"modelId": "claude-sonnet-4"}, {"modelId": "claude-sonnet-4.5"}, @@ -282,6 +289,9 @@ def _get_raw_env_value(var_name: str, env_file: str = ".env") -> Optional[str]: {"modelId": "claude-opus-4.5"}, {"modelId": "claude-opus-4.6"}, {"modelId": "claude-opus-4.7"}, + # Claude Opus 4.8 (released 2026-05-28): 1M token context window by default + # on the Claude API / Amazon Bedrock (Kiro is Bedrock-backed), 128k max output. + {"modelId": "claude-opus-4.8", "tokenLimits": {"maxInputTokens": 1000000}}, {"modelId": "deepseek-3.2"}, {"modelId": "glm-5"}, {"modelId": "minimax-m2.1"}, @@ -296,8 +306,78 @@ def _get_raw_env_value(var_name: str, env_file: str = ".env") -> Optional[str]: # Model cache TTL in seconds (1 hour) MODEL_CACHE_TTL: int = 3600 -# Default maximum number of input tokens -DEFAULT_MAX_INPUT_TOKENS: int = 200000 +# Default maximum number of input tokens. +# Used for token-usage estimation (context_usage_percentage * max_input_tokens) +# when the upstream model limit is unknown. NOT an enforced cap - requests are +# passed through to Kiro regardless. +# +# Configurable via env so larger context windows (e.g. 1M on Kiro) are accounted +# for accurately on the runtime.kiro.dev endpoint, which does not expose +# /ListAvailableModels and therefore has no per-model tokenLimits. +DEFAULT_MAX_INPUT_TOKENS: int = int(os.getenv("DEFAULT_MAX_INPUT_TOKENS", "200000")) + + +def _parse_model_context_overrides() -> Dict[str, int]: + """ + Parse per-model context window overrides from the environment. + + These overrides take precedence over both the cached upstream value (from + /ListAvailableModels) and DEFAULT_MAX_INPUT_TOKENS. They let users declare a + model's real context window when the gateway cannot discover it - most + notably on the runtime.kiro.dev endpoint, where Kiro may grant a 1M window + but provides no token limits to read. + + Accepted format (env var MODEL_CONTEXT_WINDOWS), JSON object mapping a Kiro + model id to its max input tokens: + + MODEL_CONTEXT_WINDOWS='{"claude-sonnet-4": 1000000, "claude-opus-4.7": 1000000}' + + Returns: + Mapping of model id -> max input tokens. Empty dict if unset or invalid. + """ + raw = os.getenv("MODEL_CONTEXT_WINDOWS") + if not raw: + return {} + + try: + parsed = json.loads(raw) + except (json.JSONDecodeError, ValueError): + logger.warning( + "MODEL_CONTEXT_WINDOWS is not valid JSON; ignoring. " + 'Expected e.g. {"claude-sonnet-4": 1000000}' + ) + return {} + + if not isinstance(parsed, dict): + logger.warning( + "MODEL_CONTEXT_WINDOWS must be a JSON object mapping model id -> int; ignoring." + ) + return {} + + overrides: Dict[str, int] = {} + for model_id, value in parsed.items(): + try: + tokens = int(value) + except (TypeError, ValueError): + logger.warning( + f"MODEL_CONTEXT_WINDOWS: ignoring non-integer limit for '{model_id}': {value!r}" + ) + continue + if tokens <= 0: + logger.warning( + f"MODEL_CONTEXT_WINDOWS: ignoring non-positive limit for '{model_id}': {tokens}" + ) + continue + overrides[str(model_id)] = tokens + + if overrides: + logger.info(f"Loaded {len(overrides)} model context window override(s) from env") + + return overrides + + +# Per-model context window overrides (highest precedence). See helper above. +MODEL_CONTEXT_WINDOWS: Dict[str, int] = _parse_model_context_overrides() # ================================================================================================== # Tool Description Handling (Kiro API Limitations) diff --git a/tests/unit/test_cache.py b/tests/unit/test_cache.py index 89396ec5..cdb28de0 100644 --- a/tests/unit/test_cache.py +++ b/tests/unit/test_cache.py @@ -8,6 +8,7 @@ import asyncio import time import pytest +from unittest.mock import patch from kiro.cache import ModelInfoCache from kiro.config import DEFAULT_MAX_INPUT_TOKENS @@ -268,6 +269,114 @@ async def test_get_max_input_tokens_returns_default_when_max_input_is_none(self) assert max_tokens == DEFAULT_MAX_INPUT_TOKENS +class TestModelInfoCacheContextWindowOverrides: + """ + Tests for per-model context window overrides (MODEL_CONTEXT_WINDOWS). + + These verify that an explicit override takes precedence over both the cached + upstream tokenLimits and the default - the mechanism that lets users declare + a 1M window on the runtime.kiro.dev endpoint (no /ListAvailableModels). + """ + + @pytest.mark.asyncio + async def test_override_takes_precedence_over_cached_value(self): + """ + What it does: Verifies an override beats the cached upstream tokenLimits. + Purpose: Explicit user config must win over discovered limits. + """ + print("Setup: cache with a 200k model, override declaring 1M...") + cache = ModelInfoCache() + await cache.update([ + {"modelId": "claude-sonnet-4", "tokenLimits": {"maxInputTokens": 200000}} + ]) + + with patch.dict( + "kiro.cache.MODEL_CONTEXT_WINDOWS", + {"claude-sonnet-4": 1000000}, + clear=True, + ): + result = cache.get_max_input_tokens("claude-sonnet-4") + + print(f"Result: {result}") + assert result == 1000000 + + @pytest.mark.asyncio + async def test_override_applies_when_model_not_in_cache(self): + """ + What it does: Verifies override is used even if the model isn't cached. + Purpose: Runtime endpoint uses a static fallback list with no tokenLimits, + so the override must apply without a cache entry. + """ + print("Setup: empty cache, override declaring 1M for a model...") + cache = ModelInfoCache() + + with patch.dict( + "kiro.cache.MODEL_CONTEXT_WINDOWS", + {"claude-opus-4.7": 1000000}, + clear=True, + ): + result = cache.get_max_input_tokens("claude-opus-4.7") + + print(f"Result: {result}") + assert result == 1000000 + + @pytest.mark.asyncio + async def test_no_override_falls_back_to_cached_value(self): + """ + What it does: Verifies models without an override still use cached limits. + Purpose: Overrides must not affect unrelated models. + """ + print("Setup: cached model, override only for a different model...") + cache = ModelInfoCache() + await cache.update([ + {"modelId": "claude-haiku-4.5", "tokenLimits": {"maxInputTokens": 200000}} + ]) + + with patch.dict( + "kiro.cache.MODEL_CONTEXT_WINDOWS", + {"some-other-model": 1000000}, + clear=True, + ): + result = cache.get_max_input_tokens("claude-haiku-4.5") + + print(f"Result: {result}") + assert result == 200000 + + @pytest.mark.asyncio + async def test_no_override_falls_back_to_default(self): + """ + What it does: Verifies default is used when neither override nor cache match. + Purpose: Guard the bottom of the resolution chain. + """ + print("Setup: empty cache, empty overrides...") + cache = ModelInfoCache() + + with patch.dict("kiro.cache.MODEL_CONTEXT_WINDOWS", {}, clear=True): + result = cache.get_max_input_tokens("unknown-model") + + print(f"Result: {result}") + assert result == DEFAULT_MAX_INPUT_TOKENS + + @pytest.mark.asyncio + async def test_opus_4_8_fallback_reports_1m(self): + """ + What it does: Verifies Opus 4.8 from FALLBACK_MODELS reports a 1M window. + Purpose: End-to-end check that the static fallback list (runtime endpoint) + surfaces the correct 1M context for Opus 4.8 with no env override. + """ + from kiro.config import FALLBACK_MODELS + + print("Setup: cache populated from FALLBACK_MODELS, no overrides...") + cache = ModelInfoCache() + await cache.update(FALLBACK_MODELS) + + with patch.dict("kiro.cache.MODEL_CONTEXT_WINDOWS", {}, clear=True): + result = cache.get_max_input_tokens("claude-opus-4.8") + + print(f"Result: {result}") + assert result == 1000000 + + class TestModelInfoCacheIsEmpty: """Тесты проверки пустоты кэша.""" diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index dae701e5..9503299a 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -968,4 +968,151 @@ def test_state_save_interval_seconds_default(self, monkeypatch): reload(config_module) print(f"Comparing STATE_SAVE_INTERVAL_SECONDS: Expected 10, Got {config_module.STATE_SAVE_INTERVAL_SECONDS}") - assert config_module.STATE_SAVE_INTERVAL_SECONDS == 10 \ No newline at end of file + assert config_module.STATE_SAVE_INTERVAL_SECONDS == 10 + + +class TestModelContextWindowOverrides: + """ + Tests for _parse_model_context_overrides (MODEL_CONTEXT_WINDOWS env var). + + Verifies parsing of per-model context window overrides used to declare a + model's real window (e.g. 1M on Kiro) when the gateway cannot discover it. + """ + + def _parse_with_env(self, value): + """Helper: set/unset MODEL_CONTEXT_WINDOWS and call the parser.""" + from kiro.config import _parse_model_context_overrides + + original_getenv = os.getenv + + def mock_getenv(key, default=None): + if key == "MODEL_CONTEXT_WINDOWS": + return value + return original_getenv(key, default) + + with patch.object(os, "getenv", side_effect=mock_getenv): + return _parse_model_context_overrides() + + def test_unset_returns_empty(self): + """ + What it does: Verifies unset env var yields no overrides. + Purpose: Default behavior must be a no-op. + """ + print("Action: Parsing with MODEL_CONTEXT_WINDOWS unset...") + result = self._parse_with_env(None) + print(f"Result: {result}") + assert result == {} + + def test_valid_single_override(self): + """ + What it does: Verifies a valid single-model JSON map is parsed. + Purpose: PRIMARY use case - declare 1M for one model. + """ + print("Action: Parsing valid single override...") + result = self._parse_with_env('{"claude-sonnet-4": 1000000}') + print(f"Result: {result}") + assert result == {"claude-sonnet-4": 1000000} + + def test_valid_multiple_overrides(self): + """ + What it does: Verifies multiple models are parsed. + Purpose: Users may run several 1M models. + """ + print("Action: Parsing multiple overrides...") + result = self._parse_with_env( + '{"claude-sonnet-4": 1000000, "claude-opus-4.7": 1000000}' + ) + print(f"Result: {result}") + assert result == {"claude-sonnet-4": 1000000, "claude-opus-4.7": 1000000} + + def test_string_integer_values_are_coerced(self): + """ + What it does: Verifies numeric strings are coerced to int. + Purpose: Tolerate '1000000' as well as 1000000 in the JSON. + """ + print("Action: Parsing override with string integer...") + result = self._parse_with_env('{"claude-sonnet-4": "1000000"}') + print(f"Result: {result}") + assert result == {"claude-sonnet-4": 1000000} + + def test_invalid_json_returns_empty(self): + """ + What it does: Verifies malformed JSON is ignored (no crash). + Purpose: A bad env var must not take down the gateway. + """ + print("Action: Parsing malformed JSON...") + result = self._parse_with_env("{not valid json") + print(f"Result: {result}") + assert result == {} + + def test_non_object_json_returns_empty(self): + """ + What it does: Verifies a JSON array/scalar is rejected. + Purpose: Only an object mapping is meaningful. + """ + print("Action: Parsing JSON array...") + result = self._parse_with_env('[1000000]') + print(f"Result: {result}") + assert result == {} + + def test_non_integer_value_is_skipped(self): + """ + What it does: Verifies non-numeric limits are skipped, others kept. + Purpose: Partial robustness - one bad entry must not drop the rest. + """ + print("Action: Parsing override with one bad and one good entry...") + result = self._parse_with_env( + '{"bad": "lots", "claude-sonnet-4": 1000000}' + ) + print(f"Result: {result}") + assert result == {"claude-sonnet-4": 1000000} + + def test_non_positive_value_is_skipped(self): + """ + What it does: Verifies zero/negative limits are skipped. + Purpose: A non-positive window is nonsensical for estimation. + """ + print("Action: Parsing override with non-positive limits...") + result = self._parse_with_env('{"a": 0, "b": -5, "claude-sonnet-4": 1000000}') + print(f"Result: {result}") + assert result == {"claude-sonnet-4": 1000000} + + +class TestFallbackModels: + """ + Tests for the FALLBACK_MODELS static list (used on the runtime.kiro.dev + endpoint, which has no /ListAvailableModels). + """ + + def test_opus_4_8_present(self): + """ + What it does: Verifies claude-opus-4.8 is in the fallback list. + Purpose: Ensure Opus 4.8 is recognized out-of-the-box even without + dynamic model discovery. + """ + from kiro.config import FALLBACK_MODELS + + model_ids = [m["modelId"] for m in FALLBACK_MODELS] + print(f"Fallback model ids: {model_ids}") + assert "claude-opus-4.8" in model_ids + + def test_opus_4_8_has_1m_context_window(self): + """ + What it does: Verifies the Opus 4.8 fallback entry declares a 1M window. + Purpose: Opus 4.8 has a 1M context on Bedrock/Kiro; token accounting must + reflect that without requiring MODEL_CONTEXT_WINDOWS env config. + """ + from kiro.config import FALLBACK_MODELS + + opus_48 = next(m for m in FALLBACK_MODELS if m["modelId"] == "claude-opus-4.8") + print(f"Opus 4.8 entry: {opus_48}") + assert opus_48["tokenLimits"]["maxInputTokens"] == 1000000 + + def test_all_fallback_entries_have_model_id(self): + """ + What it does: Verifies every fallback entry has a modelId. + Purpose: cache.update keys on modelId; a missing key would crash refresh. + """ + from kiro.config import FALLBACK_MODELS + + assert all("modelId" in m and m["modelId"] for m in FALLBACK_MODELS) diff --git a/tests/unit/test_model_resolver.py b/tests/unit/test_model_resolver.py index 63036f07..fe0229c9 100644 --- a/tests/unit/test_model_resolver.py +++ b/tests/unit/test_model_resolver.py @@ -125,6 +125,29 @@ def test_normalizes_opus_dash_to_dot(self): print(f"Comparing result: Expected 'claude-opus-4.5', Got '{result}'") assert result == "claude-opus-4.5" + + def test_normalizes_opus_4_8(self): + """ + What it does: claude-opus-4-8 → claude-opus-4.8 + Goal: Check Opus 4.8 (the 1M-context flagship) dash-to-dot conversion, + which is the exact id Claude Code sends (claude-opus-4-8). + """ + print("Action: Normalizing 'claude-opus-4-8'...") + result = normalize_model_name("claude-opus-4-8") + + print(f"Comparing result: Expected 'claude-opus-4.8', Got '{result}'") + assert result == "claude-opus-4.8" + + def test_strips_date_suffix_opus_4_8(self): + """ + What it does: claude-opus-4-8-20260528 → claude-opus-4.8 + Goal: Check date-stamped Opus 4.8 id normalizes to the Kiro form. + """ + print("Action: Normalizing 'claude-opus-4-8-20260528'...") + result = normalize_model_name("claude-opus-4-8-20260528") + + print(f"Comparing result: Expected 'claude-opus-4.8', Got '{result}'") + assert result == "claude-opus-4.8" # === Removal of date suffix === From 66a3d34960ec58d80e2d408b3ebd79e9f75191c0 Mon Sep 17 00:00:00 2001 From: coderhisham Date: Thu, 25 Jun 2026 12:31:32 +0530 Subject: [PATCH 2/6] fix(auth): return actionable 401 on invalid/expired refresh token When the auth provider rejects the refresh token (e.g. AWS SSO OIDC 400 invalid_grant), the gateway previously surfaced it as a generic HTTP 500 'Internal error'. Add InvalidRefreshTokenError, raise it from get_access_token when the provider response indicates an invalid/expired token, and map it to a clear 401 (authentication_error) with re-auth guidance in both API routes. Transient/server errors (e.g. 500) still propagate unchanged. Adds tests for the detector and for get_access_token raising the typed error vs propagating others. --- kiro/auth.py | 50 ++++++++++++ kiro/routes_anthropic.py | 19 ++++- kiro/routes_openai.py | 10 ++- tests/unit/test_auth_manager.py | 139 +++++++++++++++++++++++++++++++- 4 files changed, 215 insertions(+), 3 deletions(-) diff --git a/kiro/auth.py b/kiro/auth.py index 61ffe02b..b6ff57f7 100644 --- a/kiro/auth.py +++ b/kiro/auth.py @@ -58,6 +58,45 @@ "codewhisperer:odic:token", # Legacy AWS SSO OIDC ] + +class InvalidRefreshTokenError(Exception): + """ + Raised when the auth provider rejects the stored refresh token. + + This indicates the refresh token has expired or been revoked (e.g. AWS SSO + OIDC returns 400 ``invalid_grant``). It is a credentials/configuration + problem - retrying will not help. The user must re-authenticate and update + the gateway's credentials. Routes map this to HTTP 401 (not 500) with an + actionable message. + """ + + +def _is_invalid_refresh_token_response(response: "httpx.Response") -> bool: + """ + Detect whether an auth-endpoint error response means the refresh token is + no longer valid (expired/revoked), as opposed to a transient failure. + + Args: + response: The httpx response from a failed token refresh. + + Returns: + True if the response indicates an invalid/expired refresh token. + """ + if response.status_code not in (400, 401): + return False + error_code = "" + try: + error_code = (response.json().get("error") or "").lower() + except (ValueError, AttributeError): + error_code = "" + if error_code in ("invalid_grant", "invalid_client"): + return True + # Fall back to body text for providers that don't return a JSON error code. + try: + return "invalid refresh token" in response.text.lower() + except (ValueError, AttributeError): + return False + # Device registration keys (for AWS SSO OIDC only) SQLITE_REGISTRATION_KEYS = [ "kirocli:odic:device-registration", @@ -922,6 +961,17 @@ async def get_access_token(self) -> str: "Token expired and refresh failed. " "Please run 'kiro-cli login' to refresh your credentials." ) + # If the provider rejected the refresh token itself (expired or + # revoked), this is a credentials problem that retrying cannot + # fix. Surface a clear, actionable error that routes turn into a + # 401 instead of a confusing 500. + if _is_invalid_refresh_token_response(e.response): + raise InvalidRefreshTokenError( + "Authentication failed: the stored refresh token was rejected by the " + "auth provider (invalid_grant). It has expired or been revoked. " + "Re-authenticate (e.g. sign in to Kiro / run 'kiro-cli login', or update " + "REFRESH_TOKEN / your credentials file), then restart the gateway." + ) from e # Non-SQLite mode or non-400 error - propagate the exception raise except Exception: diff --git a/kiro/routes_anthropic.py b/kiro/routes_anthropic.py index dec797a0..c5af4c41 100644 --- a/kiro/routes_anthropic.py +++ b/kiro/routes_anthropic.py @@ -42,7 +42,7 @@ AnthropicErrorResponse, AnthropicErrorDetail, ) -from kiro.auth import KiroAuthManager, AuthType +from kiro.auth import KiroAuthManager, AuthType, InvalidRefreshTokenError from kiro.cache import ModelInfoCache from kiro.converters_anthropic import anthropic_to_kiro from kiro.streaming_anthropic import ( @@ -887,6 +887,23 @@ async def make_retry_request(): if debug_logger: debug_logger.flush_on_error(e.status_code, str(e.detail)) raise + except InvalidRefreshTokenError as e: + # Credentials problem (refresh token expired/revoked) - not an internal + # error. Return an actionable 401 instead of a confusing 500. + await http_client.close() + logger.error(f"HTTP 401 - POST /v1/messages - {e}") + if debug_logger: + debug_logger.flush_on_error(401, str(e)) + return JSONResponse( + status_code=401, + content={ + "type": "error", + "error": { + "type": "authentication_error", + "message": str(e) + } + } + ) except Exception as e: await http_client.close() logger.error(f"Internal error: {e}", exc_info=True) diff --git a/kiro/routes_openai.py b/kiro/routes_openai.py index 262cb0e1..3666bc1b 100644 --- a/kiro/routes_openai.py +++ b/kiro/routes_openai.py @@ -44,7 +44,7 @@ ModelList, ChatCompletionRequest, ) -from kiro.auth import KiroAuthManager, AuthType +from kiro.auth import KiroAuthManager, AuthType, InvalidRefreshTokenError from kiro.cache import ModelInfoCache from kiro.model_resolver import ModelResolver from kiro.converters_openai import build_kiro_payload @@ -759,6 +759,14 @@ async def make_retry_request(): if debug_logger: debug_logger.flush_on_error(e.status_code, str(e.detail)) raise + except InvalidRefreshTokenError as e: + # Credentials problem (refresh token expired/revoked) - not an internal + # error. Return an actionable 401 instead of a confusing 500. + await http_client.close() + logger.error(f"HTTP 401 - POST /v1/chat/completions - {e}") + if debug_logger: + debug_logger.flush_on_error(401, str(e)) + raise HTTPException(status_code=401, detail=str(e)) except Exception as e: await http_client.close() logger.error(f"Internal error: {e}", exc_info=True) diff --git a/tests/unit/test_auth_manager.py b/tests/unit/test_auth_manager.py index 30226c6c..962b9124 100644 --- a/tests/unit/test_auth_manager.py +++ b/tests/unit/test_auth_manager.py @@ -12,7 +12,7 @@ from unittest.mock import AsyncMock, Mock, patch import httpx -from kiro.auth import KiroAuthManager, AuthType +from kiro.auth import KiroAuthManager, AuthType, InvalidRefreshTokenError, _is_invalid_refresh_token_response from kiro.config import TOKEN_REFRESH_THRESHOLD, get_aws_sso_oidc_url @@ -4251,3 +4251,140 @@ def test_auth_manager_api_region_priority_hierarchy(self, temp_sqlite_db_with_pr print(f"Result: api_host={manager5._api_host}") assert "ap-south-1" in manager5._api_host + + +# ============================================================================= +# Tests for invalid refresh token handling (clear 401 instead of 500) +# ============================================================================= + +class TestIsInvalidRefreshTokenResponse: + """Tests for the _is_invalid_refresh_token_response helper.""" + + def _resp(self, status_code, json_value=None, text=""): + r = Mock() + r.status_code = status_code + if json_value is None: + r.json = Mock(side_effect=ValueError("not json")) + else: + r.json = Mock(return_value=json_value) + r.text = text + return r + + def test_detects_invalid_grant(self): + """ + What it does: Detects AWS SSO OIDC invalid_grant as an invalid token. + Purpose: PRIMARY signal for an expired/revoked refresh token. + """ + resp = self._resp(400, {"error": "invalid_grant", "error_description": "Invalid refresh token provided"}) + assert _is_invalid_refresh_token_response(resp) is True + + def test_detects_invalid_client(self): + """ + What it does: Treats invalid_client as an invalid-credential failure. + Purpose: clientId/secret rejection is also a non-retryable auth error. + """ + resp = self._resp(401, {"error": "invalid_client"}) + assert _is_invalid_refresh_token_response(resp) is True + + def test_detects_via_body_text_fallback(self): + """ + What it does: Falls back to body text when no JSON error code is present. + Purpose: Tolerate providers that don't return a structured error code. + """ + resp = self._resp(400, json_value=None, text="Invalid refresh token provided") + assert _is_invalid_refresh_token_response(resp) is True + + def test_ignores_non_auth_status(self): + """ + What it does: A 500 server error is NOT treated as an invalid token. + Purpose: Transient/server errors must keep their normal handling. + """ + resp = self._resp(500, {"error": "internal"}, text="server error") + assert _is_invalid_refresh_token_response(resp) is False + + def test_ignores_other_400_errors(self): + """ + What it does: A 400 that is not grant/client related is not flagged. + Purpose: Avoid masking unrelated bad-request failures as auth errors. + """ + resp = self._resp(400, {"error": "invalid_request"}, text="bad request") + assert _is_invalid_refresh_token_response(resp) is False + + +class TestGetAccessTokenInvalidRefreshToken: + """Tests that get_access_token surfaces a clear InvalidRefreshTokenError.""" + + def _make_error_post(self, status_code=400, error_code="invalid_grant"): + """Build a mock httpx post that returns an auth error response.""" + error_response = AsyncMock() + error_response.status_code = status_code + error_response.text = json.dumps({"error": error_code, "error_description": "Invalid refresh token provided"}) + error_response.json = Mock(return_value={"error": error_code, "error_description": "Invalid refresh token provided"}) + error_response.raise_for_status = Mock( + side_effect=httpx.HTTPStatusError( + f"{status_code} error", request=Mock(), response=error_response + ) + ) + + async def mock_post(*args, **kwargs): + return error_response + + return mock_post + + @pytest.mark.asyncio + async def test_raises_invalid_refresh_token_on_invalid_grant(self): + """ + What it does: get_access_token raises InvalidRefreshTokenError when the + SSO OIDC provider returns 400 invalid_grant. + Purpose: PRIMARY fix - a rejected refresh token must not surface as a 500. + """ + print("Setup: AWS SSO OIDC manager with an expired/missing token...") + manager = KiroAuthManager( + refresh_token="bad_refresh", + client_id="cid", + client_secret="csecret", + region="ap-south-1", + ) + manager._access_token = None # Force a refresh + + with patch("kiro.auth.httpx.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client.post = self._make_error_post(400, "invalid_grant") + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client_class.return_value = mock_client + + print("Action: get_access_token (expecting InvalidRefreshTokenError)...") + with pytest.raises(InvalidRefreshTokenError) as exc_info: + await manager.get_access_token() + + print(f"Raised: {exc_info.value}") + message = str(exc_info.value).lower() + assert "re-authenticate" in message or "expired" in message + + @pytest.mark.asyncio + async def test_propagates_server_error_unchanged(self): + """ + What it does: A 500 from the auth provider is NOT converted to + InvalidRefreshTokenError; it propagates as before. + Purpose: Only credential rejection becomes a 401; transient errors stay. + """ + print("Setup: AWS SSO OIDC manager, provider returns 500...") + manager = KiroAuthManager( + refresh_token="some_refresh", + client_id="cid", + client_secret="csecret", + region="ap-south-1", + ) + manager._access_token = None + + with patch("kiro.auth.httpx.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client.post = self._make_error_post(500, "internal") + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client_class.return_value = mock_client + + print("Action: get_access_token (expecting httpx.HTTPStatusError, not auth error)...") + with pytest.raises(httpx.HTTPStatusError): + await manager.get_access_token() From aaac78cc626f03dde53e1ed21bec351b5599a307 Mon Sep 17 00:00:00 2001 From: coderhisham Date: Thu, 25 Jun 2026 12:31:56 +0530 Subject: [PATCH 3/6] fix(mcp): make web_search reliable and never leak the tool call (#231, #173) Two root causes made web_search fail and surface as 'No such tool available' in clients: 1. The /mcp call sent only a bare Authorization header. It now sends the full Kiro client-identity headers (User-Agent/x-amz-user-agent/x-amzn-kiro-agent-mode) adjusted for JSON-RPC, and includes top-level profileArn for Enterprise accounts. Error logging now includes the response body. 2. On MCP failure the gateway fell through and emitted the raw web_search tool_use/tool_call to the client. All paths now degrade gracefully: - Anthropic streaming + non-streaming emit a native web_search_tool_result (error) instead of leaking a tool_use. - OpenAI streaming (and non-streaming via reuse) emit an in-band unavailable note instead of a tool_call. Adds shared outcome helpers and tests across both APIs and both modes. --- kiro/mcp_tools.py | 148 +++++++++++++- kiro/streaming_anthropic.py | 213 ++++++++++++--------- kiro/streaming_openai.py | 89 +++++---- tests/unit/test_mcp_tools.py | 255 ++++++++++++++++++++++++- tests/unit/test_streaming_anthropic.py | 152 +++++++++++++++ tests/unit/test_streaming_openai.py | 113 ++++++++++- 6 files changed, 833 insertions(+), 137 deletions(-) diff --git a/kiro/mcp_tools.py b/kiro/mcp_tools.py index c908ec55..113ea3ef 100644 --- a/kiro/mcp_tools.py +++ b/kiro/mcp_tools.py @@ -34,13 +34,14 @@ import random import string from datetime import datetime -from typing import Dict, Any, Optional, Tuple +from typing import Dict, Any, List, Optional, Tuple import httpx from fastapi.responses import JSONResponse, StreamingResponse from loguru import logger from kiro.tokenizer import count_message_tokens, count_tokens +from kiro.utils import get_kiro_headers # Import debug_logger try: @@ -135,6 +136,12 @@ async def call_kiro_mcp_api( "arguments": {"query": query} } } + + # Enterprise / Builder accounts require profileArn, and it goes at the TOP + # level of the request body (not inside params). Without it the /mcp call is + # rejected and web_search silently degrades to "unavailable". + if getattr(auth_manager, "profile_arn", None): + mcp_request["profileArn"] = auth_manager.profile_arn # Log MCP request try: @@ -147,12 +154,16 @@ async def call_kiro_mcp_api( try: token = await auth_manager.get_access_token() - # EXACT headers from architecture - headers = { - "Authorization": f"Bearer {token}", - "x-amzn-codewhisperer-optout": "false", - "Content-Type": "application/json" - } + # The /mcp endpoint requires the same Kiro client-identity headers as the + # completion path (User-Agent with KiroIDE+fingerprint, x-amz-user-agent, + # x-amzn-kiro-agent-mode). Without them Kiro returns 403, which would make + # this call fail and leak the web_search tool_use back to the client. + # We reuse get_kiro_headers and adjust for the JSON-RPC /mcp endpoint: + # it expects plain JSON (not x-amz-json-1.0) and has no x-amz-target. + headers = get_kiro_headers(auth_manager, token) + headers["Content-Type"] = "application/json" + headers.pop("x-amz-target", None) + headers["x-amzn-codewhisperer-optout"] = "false" mcp_url = f"{auth_manager.q_host}/mcp" logger.debug(f"Calling MCP API: {mcp_url}") @@ -161,7 +172,7 @@ async def call_kiro_mcp_api( response = await client.post(mcp_url, json=mcp_request, headers=headers) if response.status_code != 200: - logger.error(f"MCP API error: {response.status_code}") + logger.error(f"MCP API error: {response.status_code} - {response.text[:500]}") return None, None mcp_response = response.json() @@ -274,6 +285,127 @@ def generate_search_summary(query: str, results: Dict) -> str: return summary +# ================================================================================================== +# WebSearch Outcome Helpers (shared by streaming + non-streaming, both APIs) +# ================================================================================================== + +# Error code reported in an Anthropic web_search_tool_result when the search +# could not be completed (e.g. the Kiro /mcp call failed). +WEB_SEARCH_ERROR_CODE: str = "unavailable" + + +def generate_search_unavailable_summary(query: str) -> str: + """ + Human-readable note shown when a web_search could not be completed. + + Used as a graceful fallback so the model/user gets a clear, in-band message + instead of the gateway leaking an unhandled ``web_search`` tool call back to + a client that never defined that tool. + + Args: + query: The search query that could not be completed (may be empty) + + Returns: + A short note wrapped in tags. + """ + target = f' for "{query}"' if query else "" + return ( + f"\n\nSearch{target} could not be completed: the web search " + f"service is currently unavailable. Continue using your existing knowledge " + f"and note that information may be out of date.\n\n" + ) + + +def build_web_search_result_items(results: Optional[Dict]) -> List[Dict[str, Any]]: + """ + Convert MCP results into Anthropic web_search_result content items. + + Args: + results: Parsed MCP results dict (or None) + + Returns: + List of web_search_result blocks (empty if no results). + """ + items: List[Dict[str, Any]] = [] + for r in (results or {}).get("results", []): + items.append({ + "type": "web_search_result", + "title": r.get("title", ""), + "url": r.get("url", ""), + "encrypted_content": r.get("snippet", ""), + "page_age": None, + }) + return items + + +def web_search_tool_result_content(results: Optional[Dict]) -> Any: + """ + Build the ``content`` for an Anthropic web_search_tool_result block. + + Args: + results: Parsed MCP results dict, or None if the search failed. + + Returns: + - On success: a list of web_search_result items. + - On failure (results is None): a web_search_tool_result_error dict, the + native Anthropic shape for a server-side tool failure. This keeps the + response valid without leaking an executable ``web_search`` tool_use. + """ + if results is None: + return { + "type": "web_search_tool_result_error", + "error_code": WEB_SEARCH_ERROR_CODE, + } + return build_web_search_result_items(results) + + +def build_anthropic_web_search_blocks( + query: str, + mcp_tool_use_id: str, + results: Optional[Dict], +) -> List[Dict[str, Any]]: + """ + Build the non-streaming Anthropic content blocks for a web_search outcome. + + Produces the native server-side tool sequence so a client (e.g. Claude Code) + never receives an unhandled ``web_search`` tool_use: + 1. server_tool_use (the search request) + 2. web_search_tool_result (results, or an error on failure) + 3. text (human-readable summary, or an unavailable note on failure) + + Args: + query: The search query. + mcp_tool_use_id: ID linking server_tool_use and its result. + results: Parsed MCP results, or None on failure. + + Returns: + List of Anthropic content-block dicts. + """ + summary = ( + generate_search_summary(query, results) + if results is not None + else generate_search_unavailable_summary(query) + ) + return [ + { + "type": "server_tool_use", + "id": mcp_tool_use_id, + "name": "web_search", + "input": {"query": query}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": mcp_tool_use_id, + "content": web_search_tool_result_content(results), + }, + { + "type": "text", + "text": summary, + }, + ] + + + # ================================================================================================== # SSE Emulation (Anthropic Format) # ================================================================================================== diff --git a/kiro/streaming_anthropic.py b/kiro/streaming_anthropic.py index 4fee71e4..27240d2f 100644 --- a/kiro/streaming_anthropic.py +++ b/kiro/streaming_anthropic.py @@ -353,7 +353,12 @@ async def stream_kiro_to_anthropic( # INTERCEPT web_search tool calls (Path B - MCP emulation) if tool_name == "web_search": - from kiro.mcp_tools import call_kiro_mcp_api, generate_search_summary + from kiro.mcp_tools import ( + call_kiro_mcp_api, + generate_search_summary, + generate_search_unavailable_summary, + web_search_tool_result_content, + ) logger.info("Intercepted web_search tool call (Path B - MCP emulation)") @@ -375,97 +380,100 @@ async def stream_kiro_to_anthropic( # Call MCP API mcp_tool_use_id, results = await call_kiro_mcp_api(query, auth_manager) + # Ensure we always have a tool_use id, even on failure, so the + # server_tool_use and web_search_tool_result blocks stay linked. + if not mcp_tool_use_id: + mcp_tool_use_id = f"srvtoolu_{uuid.uuid4().hex[:32]}" + if results is None: - logger.error("MCP API call failed for web_search") - # Continue with normal tool_use processing (will show error to user) - else: - # Emit server_tool_use + web_search_tool_result + text summary - # (full SSE sequence as in mcp_tools.py) - - # Event: content_block_start (server_tool_use) - yield format_sse_event("content_block_start", { - "type": "content_block_start", - "index": current_block_index, - "content_block": { - "id": mcp_tool_use_id, - "type": "server_tool_use", - "name": "web_search", - "input": {} - } - }) - - # Event: content_block_delta (input_json_delta) + logger.error( + "MCP API call failed for web_search; emitting error result " + "(no tool_use leaked to client)" + ) + + # Always emit a native server-side tool sequence (success OR + # error). This guarantees the gateway never leaks an unhandled + # `web_search` tool_use to a client that did not define it. + summary = ( + generate_search_summary(query, results) + if results is not None + else generate_search_unavailable_summary(query) + ) + + # Event: content_block_start (server_tool_use) + yield format_sse_event("content_block_start", { + "type": "content_block_start", + "index": current_block_index, + "content_block": { + "id": mcp_tool_use_id, + "type": "server_tool_use", + "name": "web_search", + "input": {} + } + }) + + # Event: content_block_delta (input_json_delta) + yield format_sse_event("content_block_delta", { + "type": "content_block_delta", + "index": current_block_index, + "delta": { + "type": "input_json_delta", + "partial_json": json.dumps({"query": query}) + } + }) + + # Event: content_block_stop (server_tool_use) + yield format_sse_event("content_block_stop", { + "type": "content_block_stop", + "index": current_block_index + }) + current_block_index += 1 + + # Event: content_block_start (web_search_tool_result) + # content is a result list on success, or an error dict on failure. + yield format_sse_event("content_block_start", { + "type": "content_block_start", + "index": current_block_index, + "content_block": { + "type": "web_search_tool_result", + "tool_use_id": mcp_tool_use_id, + "content": web_search_tool_result_content(results) + } + }) + + # Event: content_block_stop (web_search_tool_result) + yield format_sse_event("content_block_stop", { + "type": "content_block_stop", + "index": current_block_index + }) + current_block_index += 1 + + # Event: content_block_start (text) + yield format_sse_event("content_block_start", { + "type": "content_block_start", + "index": current_block_index, + "content_block": {"type": "text", "text": ""} + }) + + # Events: content_block_delta (text_delta) - stream summary + chunk_size = 100 + for i in range(0, len(summary), chunk_size): + chunk = summary[i:i + chunk_size] yield format_sse_event("content_block_delta", { "type": "content_block_delta", "index": current_block_index, - "delta": { - "type": "input_json_delta", - "partial_json": json.dumps({"query": query}) - } - }) - - # Event: content_block_stop (server_tool_use) - yield format_sse_event("content_block_stop", { - "type": "content_block_stop", - "index": current_block_index - }) - current_block_index += 1 - - # Event: content_block_start (web_search_tool_result) - search_content = [] - for r in results.get("results", []): - search_content.append({ - "type": "web_search_result", - "title": r.get("title", ""), - "url": r.get("url", ""), - "encrypted_content": r.get("snippet", ""), - "page_age": None - }) - - yield format_sse_event("content_block_start", { - "type": "content_block_start", - "index": current_block_index, - "content_block": { - "type": "web_search_tool_result", - "tool_use_id": mcp_tool_use_id, - "content": search_content - } - }) - - # Event: content_block_stop (web_search_tool_result) - yield format_sse_event("content_block_stop", { - "type": "content_block_stop", - "index": current_block_index - }) - current_block_index += 1 - - # Event: content_block_start (text) - yield format_sse_event("content_block_start", { - "type": "content_block_start", - "index": current_block_index, - "content_block": {"type": "text", "text": ""} - }) - - # Events: content_block_delta (text_delta) - stream summary - summary = generate_search_summary(query, results) - chunk_size = 100 - for i in range(0, len(summary), chunk_size): - chunk = summary[i:i + chunk_size] - yield format_sse_event("content_block_delta", { - "type": "content_block_delta", - "index": current_block_index, - "delta": {"type": "text_delta", "text": chunk} - }) - - # Event: content_block_stop (text) - yield format_sse_event("content_block_stop", { - "type": "content_block_stop", - "index": current_block_index + "delta": {"type": "text_delta", "text": chunk} }) - current_block_index += 1 - - # Skip normal tool_use processing - continue + + # Event: content_block_stop (text) + yield format_sse_event("content_block_stop", { + "type": "content_block_stop", + "index": current_block_index + }) + current_block_index += 1 + + # Skip normal tool_use processing (never leak web_search) + continue # Check if this tool was truncated if tool.get('_truncation_detected'): @@ -796,6 +804,37 @@ async def collect_anthropic_response( except json.JSONDecodeError: tool_input = {} + # ============================================================================== + # WebSearch Support - Path B: MCP Tool Emulation (Non-Streaming Interception) + # ============================================================================== + # Mirror the streaming path: intercept web_search, run the MCP call, and + # emit native server_tool_use + web_search_tool_result + text blocks. + # This guarantees an auto-injected web_search call is never leaked to the + # client as an unhandled tool_use (which Claude Code rejects with + # "No such tool available: web_search"). + if tool_name == "web_search": + from kiro.mcp_tools import call_kiro_mcp_api, build_anthropic_web_search_blocks + + logger.info("Intercepted web_search tool call (Path B - non-streaming)") + query = tool_input.get("query", "") if isinstance(tool_input, dict) else "" + if not query: + logger.warning("web_search called without query, skipping MCP call") + continue + + mcp_tool_use_id, results = await call_kiro_mcp_api(query, auth_manager) + if not mcp_tool_use_id: + mcp_tool_use_id = f"srvtoolu_{uuid.uuid4().hex[:32]}" + if results is None: + logger.error( + "MCP API call failed for web_search (non-streaming); emitting " + "error result (no tool_use leaked to client)" + ) + + content_blocks.extend( + build_anthropic_web_search_blocks(query, mcp_tool_use_id, results) + ) + continue + content_blocks.append({ "type": "tool_use", "id": tool_id, diff --git a/kiro/streaming_openai.py b/kiro/streaming_openai.py index b6179540..4842fb83 100644 --- a/kiro/streaming_openai.py +++ b/kiro/streaming_openai.py @@ -198,7 +198,11 @@ async def stream_kiro_to_openai_internal( # INTERCEPT web_search tool calls (Path B - MCP emulation) if tool_name == "web_search": - from kiro.mcp_tools import call_kiro_mcp_api, generate_search_summary + from kiro.mcp_tools import ( + call_kiro_mcp_api, + generate_search_summary, + generate_search_unavailable_summary, + ) logger.info("Intercepted web_search tool call (Path B - MCP emulation)") @@ -212,52 +216,57 @@ async def stream_kiro_to_openai_internal( # Extract query query = tool_input.get("query", "") + if not query: - logger.warning("web_search called without query, skipping MCP call") - # Continue with normal tool_use processing + logger.warning( + "web_search called without query; emitting unavailable note " + "(no tool_call leaked to client)" + ) + summary = generate_search_unavailable_summary("") else: logger.debug(f"WebSearch query (Path B): {query}") - - # Call MCP API mcp_tool_use_id, results = await call_kiro_mcp_api(query, auth_manager) - if results is None: - logger.error("MCP API call failed for web_search") - # Continue with normal tool_use processing (will show error to user) + logger.error( + "MCP API call failed for web_search; emitting unavailable " + "note (no tool_call leaked to client)" + ) + summary = generate_search_unavailable_summary(query) else: - # Emit summary as content chunks (OpenAI format) summary = generate_search_summary(query, results) - - # Send content chunks - chunk_size = 100 - for i in range(0, len(summary), chunk_size): - content_chunk = summary[i:i + chunk_size] - - delta = {"content": content_chunk} - if first_chunk: - delta["role"] = "assistant" - first_chunk = False - - openai_chunk = { - "id": completion_id, - "object": "chat.completion.chunk", - "created": created_time, - "model": model, - "choices": [{"index": 0, "delta": delta, "finish_reason": None}] - } - - chunk_text = f"data: {json.dumps(openai_chunk, ensure_ascii=False)}\n\n" - - if debug_logger: - debug_logger.log_modified_chunk(chunk_text.encode('utf-8')) - - yield chunk_text - - # Accumulate for token counting - full_content += summary - - # Skip normal tool_use processing - continue + + # Emit the outcome as assistant content chunks (OpenAI format). + # web_search is never emitted as a tool_call the client cannot + # execute - on failure we degrade to an in-band note. + chunk_size = 100 + for i in range(0, len(summary), chunk_size): + content_chunk = summary[i:i + chunk_size] + + delta = {"content": content_chunk} + if first_chunk: + delta["role"] = "assistant" + first_chunk = False + + openai_chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created_time, + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": None}] + } + + chunk_text = f"data: {json.dumps(openai_chunk, ensure_ascii=False)}\n\n" + + if debug_logger: + debug_logger.log_modified_chunk(chunk_text.encode('utf-8')) + + yield chunk_text + + # Accumulate for token counting + full_content += summary + + # Skip normal tool_use processing (never leak web_search) + continue # Collect tool calls from stream (normal tools, not web_search) tool_calls_from_stream.append(event.tool_use) diff --git a/tests/unit/test_mcp_tools.py b/tests/unit/test_mcp_tools.py index 34d22ffd..40e19286 100644 --- a/tests/unit/test_mcp_tools.py +++ b/tests/unit/test_mcp_tools.py @@ -21,6 +21,10 @@ generate_random_id, call_kiro_mcp_api, generate_search_summary, + generate_search_unavailable_summary, + build_web_search_result_items, + web_search_tool_result_content, + build_anthropic_web_search_blocks, extract_query_from_messages, handle_native_web_search, generate_anthropic_web_search_sse, @@ -146,7 +150,133 @@ async def test_mcp_api_success(self, mock_auth_manager): assert results["totalResults"] == 1 assert results["results"][0]["title"] == "Python Tutorial" assert results["results"][0]["url"] == "https://python.org" - + + @pytest.mark.asyncio + async def test_mcp_api_sends_kiro_identity_headers(self, mock_auth_manager): + """ + What it does: Verifies the /mcp request carries the Kiro client-identity + headers (User-Agent with KiroIDE+fingerprint, x-amz-user-agent, + x-amzn-kiro-agent-mode) plus JSON-RPC adjustments. + Purpose: Guard against the 403 regression. Without these headers Kiro + rejects /mcp, call_kiro_mcp_api returns (None, None), and the + web_search tool_use leaks back to the client ("No such tool + available: web_search"). Also verifies the JSON-RPC overrides: + plain JSON content type, no x-amz-target, optout=false. + """ + print("Setup: Mocking MCP API response to capture posted headers...") + query = "Python tutorials" + + mock_response_data = { + "id": "web_search_tooluse_abc123_1234567890_xyz", + "jsonrpc": "2.0", + "result": { + "content": [{ + "type": "text", + "text": json.dumps({"results": [], "totalResults": 0, "query": query}) + }], + "isError": False + } + } + + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json = Mock(return_value=mock_response_data) + + mock_post = AsyncMock(return_value=mock_response) + mock_client = AsyncMock() + mock_client.__aenter__.return_value.post = mock_post + + print("Action: Calling call_kiro_mcp_api...") + with patch("kiro.mcp_tools.httpx.AsyncClient", return_value=mock_client): + await call_kiro_mcp_api(query, mock_auth_manager) + + print("Inspecting posted request headers...") + assert mock_post.call_count == 1 + headers = mock_post.call_args.kwargs["headers"] + print(f"Headers: {headers}") + + # Kiro client-identity headers must be present + assert "KiroIDE" in headers["User-Agent"] + assert mock_auth_manager.fingerprint in headers["User-Agent"] + assert "KiroIDE" in headers["x-amz-user-agent"] + assert headers["x-amzn-kiro-agent-mode"] == "vibe" + assert headers["Authorization"].startswith("Bearer ") + + # JSON-RPC /mcp adjustments + assert headers["Content-Type"] == "application/json" + assert "x-amz-target" not in headers + assert headers["x-amzn-codewhisperer-optout"] == "false" + + @pytest.mark.asyncio + async def test_mcp_request_includes_profile_arn_when_set(self, mock_auth_manager): + """ + What it does: Verifies profileArn is added at the TOP level of the MCP + request body when the auth manager has one. + Purpose: Enterprise/Builder accounts require profileArn or /mcp rejects + the call, silently degrading web_search to "unavailable". + """ + print("Setup: auth_manager fixture has a profile_arn...") + query = "Python" + assert mock_auth_manager.profile_arn is not None + + mock_response_data = { + "id": "web_search_tooluse_x", + "jsonrpc": "2.0", + "result": { + "content": [{"type": "text", "text": json.dumps({"results": [], "totalResults": 0, "query": query})}], + "isError": False, + }, + } + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json = Mock(return_value=mock_response_data) + mock_post = AsyncMock(return_value=mock_response) + mock_client = AsyncMock() + mock_client.__aenter__.return_value.post = mock_post + + print("Action: Calling call_kiro_mcp_api...") + with patch("kiro.mcp_tools.httpx.AsyncClient", return_value=mock_client): + await call_kiro_mcp_api(query, mock_auth_manager) + + sent_body = mock_post.await_args.kwargs["json"] + print(f"Sent body keys: {list(sent_body.keys())}") + # Top-level profileArn, never inside params + assert sent_body.get("profileArn") == mock_auth_manager.profile_arn + assert "profileArn" not in sent_body.get("params", {}) + + @pytest.mark.asyncio + async def test_mcp_request_omits_profile_arn_when_unset(self, mock_auth_manager): + """ + What it does: Verifies profileArn is omitted when the account has none. + Purpose: Avoid sending a null/empty profileArn for non-Enterprise accounts. + """ + print("Setup: clearing profile_arn on auth_manager...") + mock_auth_manager._profile_arn = None + query = "Python" + + mock_response_data = { + "id": "web_search_tooluse_y", + "jsonrpc": "2.0", + "result": { + "content": [{"type": "text", "text": json.dumps({"results": [], "totalResults": 0, "query": query})}], + "isError": False, + }, + } + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json = Mock(return_value=mock_response_data) + mock_post = AsyncMock(return_value=mock_response) + mock_client = AsyncMock() + mock_client.__aenter__.return_value.post = mock_post + + print("Action: Calling call_kiro_mcp_api...") + with patch("kiro.mcp_tools.httpx.AsyncClient", return_value=mock_client): + await call_kiro_mcp_api(query, mock_auth_manager) + + sent_body = mock_post.await_args.kwargs["json"] + print(f"Sent body keys: {list(sent_body.keys())}") + assert "profileArn" not in sent_body + @pytest.mark.asyncio async def test_mcp_api_error_response(self, mock_auth_manager): """ @@ -600,3 +730,126 @@ async def test_generate_openai_sse_structure(self): print("Checking for usage information...") assert any('"usage"' in chunk for chunk in chunks) + + +# ================================================================================================== +# Tests for WebSearch Outcome Helpers (anti-leak hardening) +# ================================================================================================== + +class TestWebSearchOutcomeHelpers: + """ + Tests for the shared helpers that build graceful web_search outcomes so the + gateway never leaks an unhandled web_search tool_use/tool_call to the client. + """ + + SAMPLE_RESULTS = { + "results": [ + { + "title": "Python", + "url": "https://python.org", + "snippet": "Official site", + "publishedDate": 1700000000000, + } + ], + "totalResults": 1, + "query": "python", + } + + def test_unavailable_summary_includes_query(self): + """ + What it does: Verifies the failure note mentions the query and is tagged. + Purpose: Give the model/user a clear in-band message on failure. + """ + summary = generate_search_unavailable_summary("python tutorials") + print(f"Summary: {summary}") + assert "python tutorials" in summary + assert "" in summary and "" in summary + assert "unavailable" in summary.lower() + + def test_unavailable_summary_handles_empty_query(self): + """ + What it does: Verifies the failure note works with an empty query. + Purpose: web_search called without a query must not crash the helper. + """ + summary = generate_search_unavailable_summary("") + print(f"Summary: {summary}") + assert "" in summary + # No dangling 'for ""' fragment + assert 'for ""' not in summary + + def test_build_result_items_maps_fields(self): + """ + What it does: Verifies MCP results map to web_search_result items. + Purpose: Ensure title/url/snippet are carried into the result blocks. + """ + items = build_web_search_result_items(self.SAMPLE_RESULTS) + print(f"Items: {items}") + assert len(items) == 1 + assert items[0]["type"] == "web_search_result" + assert items[0]["title"] == "Python" + assert items[0]["url"] == "https://python.org" + assert items[0]["encrypted_content"] == "Official site" + + def test_build_result_items_handles_none(self): + """ + What it does: Verifies None/empty results yield an empty list. + Purpose: Defensive boundary check. + """ + assert build_web_search_result_items(None) == [] + assert build_web_search_result_items({}) == [] + + def test_tool_result_content_success_is_list(self): + """ + What it does: Verifies success content is a list of result items. + Purpose: Native Anthropic web_search_tool_result success shape. + """ + content = web_search_tool_result_content(self.SAMPLE_RESULTS) + print(f"Content: {content}") + assert isinstance(content, list) + assert content[0]["type"] == "web_search_result" + + def test_tool_result_content_failure_is_error_dict(self): + """ + What it does: Verifies failure content is the native error dict. + Purpose: PRIMARY anti-leak behavior - failure stays a server-side tool + result error, never an executable tool_use. + """ + content = web_search_tool_result_content(None) + print(f"Content: {content}") + assert content["type"] == "web_search_tool_result_error" + assert content["error_code"] == "unavailable" + + def test_build_blocks_success_sequence(self): + """ + What it does: Verifies success blocks are server_tool_use + result + text. + Purpose: Non-streaming Anthropic shape must match the streaming sequence. + """ + blocks = build_anthropic_web_search_blocks("python", "srvtoolu_x", self.SAMPLE_RESULTS) + print(f"Block types: {[b['type'] for b in blocks]}") + assert [b["type"] for b in blocks] == [ + "server_tool_use", + "web_search_tool_result", + "text", + ] + assert blocks[0]["id"] == "srvtoolu_x" + assert blocks[0]["input"] == {"query": "python"} + assert blocks[1]["tool_use_id"] == "srvtoolu_x" + assert isinstance(blocks[1]["content"], list) + # No leaked executable tool_use + assert all(b["type"] != "tool_use" for b in blocks) + + def test_build_blocks_failure_sequence(self): + """ + What it does: Verifies failure blocks carry the error result, no tool_use. + Purpose: The non-streaming failure path must not leak web_search. + """ + blocks = build_anthropic_web_search_blocks("python", "srvtoolu_x", None) + print(f"Block types: {[b['type'] for b in blocks]}") + assert [b["type"] for b in blocks] == [ + "server_tool_use", + "web_search_tool_result", + "text", + ] + assert blocks[1]["content"]["type"] == "web_search_tool_result_error" + assert "unavailable" in blocks[2]["text"].lower() + assert all(b["type"] != "tool_use" for b in blocks) diff --git a/tests/unit/test_streaming_anthropic.py b/tests/unit/test_streaming_anthropic.py index c27d599e..4f427f79 100644 --- a/tests/unit/test_streaming_anthropic.py +++ b/tests/unit/test_streaming_anthropic.py @@ -1667,3 +1667,155 @@ async def test_collect_detects_truncation_in_non_streaming(self, mock_response, # Should detect truncation and set max_tokens assert result["stop_reason"] == "max_tokens" print("✓ collect_anthropic_response detects truncation correctly") + + +# ================================================================================================== +# Tests for web_search anti-leak hardening (Anthropic streaming + non-streaming) +# ================================================================================================== + +class TestAnthropicWebSearchNoLeak: + """ + Verifies that an intercepted web_search call is never emitted to the client + as an executable tool_use - on success or failure, in streaming and + non-streaming - so Claude Code never errors with "No such tool available". + """ + + WS_TOOL_USE = { + "id": "toolu_ws", + "function": {"name": "web_search", "arguments": '{"query": "python"}'}, + } + + SAMPLE_RESULTS = { + "results": [{"title": "Python", "url": "https://python.org", "snippet": "Official"}], + "totalResults": 1, + "query": "python", + } + + @pytest.mark.asyncio + async def test_streaming_success_emits_server_tool_result_not_tool_use( + self, mock_response, mock_model_cache, mock_auth_manager + ): + """ + What it does: On a successful MCP call, streaming emits server_tool_use + + web_search_tool_result, and NO `web_search` tool_use block. + Purpose: Confirm the happy path stays a server-side tool result. + """ + async def mock_parse_kiro_stream(*args, **kwargs): + yield KiroEvent(type="tool_use", tool_use=self.WS_TOOL_USE) + + events = [] + with patch("kiro.streaming_anthropic.parse_kiro_stream", mock_parse_kiro_stream): + with patch("kiro.streaming_anthropic.parse_bracket_tool_calls", return_value=[]): + with patch( + "kiro.mcp_tools.call_kiro_mcp_api", + new=AsyncMock(return_value=("srvtoolu_ok", self.SAMPLE_RESULTS)), + ): + async for event in stream_kiro_to_anthropic( + mock_response, "claude-opus-4.8", mock_model_cache, mock_auth_manager + ): + events.append(event) + + joined = "".join(events) + print(joined) + assert "web_search_tool_result" in joined + assert "server_tool_use" in joined + # The model-facing executable tool_use must NOT be emitted for web_search + assert '"type": "tool_use"' not in joined + + @pytest.mark.asyncio + async def test_streaming_failure_emits_error_result_not_tool_use( + self, mock_response, mock_model_cache, mock_auth_manager + ): + """ + What it does: On MCP failure, streaming emits a web_search_tool_result + ERROR block, never the raw web_search tool_use. + Purpose: PRIMARY regression guard for "No such tool available: web_search". + """ + async def mock_parse_kiro_stream(*args, **kwargs): + yield KiroEvent(type="tool_use", tool_use=self.WS_TOOL_USE) + + events = [] + with patch("kiro.streaming_anthropic.parse_kiro_stream", mock_parse_kiro_stream): + with patch("kiro.streaming_anthropic.parse_bracket_tool_calls", return_value=[]): + with patch( + "kiro.mcp_tools.call_kiro_mcp_api", + new=AsyncMock(return_value=(None, None)), + ): + async for event in stream_kiro_to_anthropic( + mock_response, "claude-opus-4.8", mock_model_cache, mock_auth_manager + ): + events.append(event) + + joined = "".join(events) + print(joined) + assert "web_search_tool_result_error" in joined + # No executable tool_use leaked to the client + assert '"type": "tool_use"' not in joined + + @pytest.mark.asyncio + async def test_non_streaming_failure_emits_error_result_not_tool_use( + self, mock_response, mock_model_cache, mock_auth_manager + ): + """ + What it does: Non-streaming collector intercepts web_search and, on + failure, emits server-side error blocks (no tool_use). + Purpose: Close the non-streaming leak path (previously absent entirely). + """ + mock_result = StreamResult( + content="", + thinking_content="", + tool_calls=[{"id": "toolu_ws", "name": "web_search", "input": {"query": "python"}}], + usage=None, + context_usage_percentage=5.0, + ) + + with patch("kiro.streaming_anthropic.collect_stream_to_result", return_value=mock_result): + with patch( + "kiro.mcp_tools.call_kiro_mcp_api", + new=AsyncMock(return_value=(None, None)), + ): + result = await collect_anthropic_response( + mock_response, "claude-opus-4.8", mock_model_cache, mock_auth_manager + ) + + types = [b["type"] for b in result["content"]] + print(f"Content block types: {types}") + assert "web_search_tool_result" in types + assert "server_tool_use" in types + # No leaked executable tool_use + assert "tool_use" not in types + # The web_search_tool_result must carry the native error + ws_result = next(b for b in result["content"] if b["type"] == "web_search_tool_result") + assert ws_result["content"]["type"] == "web_search_tool_result_error" + + @pytest.mark.asyncio + async def test_non_streaming_success_emits_results_not_tool_use( + self, mock_response, mock_model_cache, mock_auth_manager + ): + """ + What it does: Non-streaming collector emits results on success, no tool_use. + Purpose: Confirm the happy path for non-streaming. + """ + mock_result = StreamResult( + content="", + thinking_content="", + tool_calls=[{"id": "toolu_ws", "name": "web_search", "input": {"query": "python"}}], + usage=None, + context_usage_percentage=5.0, + ) + + with patch("kiro.streaming_anthropic.collect_stream_to_result", return_value=mock_result): + with patch( + "kiro.mcp_tools.call_kiro_mcp_api", + new=AsyncMock(return_value=("srvtoolu_ok", self.SAMPLE_RESULTS)), + ): + result = await collect_anthropic_response( + mock_response, "claude-opus-4.8", mock_model_cache, mock_auth_manager + ) + + types = [b["type"] for b in result["content"]] + print(f"Content block types: {types}") + assert "tool_use" not in types + ws_result = next(b for b in result["content"] if b["type"] == "web_search_tool_result") + assert isinstance(ws_result["content"], list) + assert ws_result["content"][0]["type"] == "web_search_result" diff --git a/tests/unit/test_streaming_openai.py b/tests/unit/test_streaming_openai.py index b52c1b60..78d9bc22 100644 --- a/tests/unit/test_streaming_openai.py +++ b/tests/unit/test_streaming_openai.py @@ -1484,4 +1484,115 @@ async def mock_parse_kiro_stream(*args, **kwargs): # Should extract "length" from streaming chunks assert result["choices"][0]["finish_reason"] == "length" - print("✓ collect_stream_response extracts finish_reason correctly") \ No newline at end of file + print("✓ collect_stream_response extracts finish_reason correctly") + + +# ================================================================================================== +# Tests for web_search anti-leak hardening (OpenAI streaming, covers non-streaming via reuse) +# ================================================================================================== + +def _reconstruct_openai_content(chunks: list) -> str: + """Reassemble assistant content from OpenAI SSE chunk strings.""" + content = "" + for chunk in chunks: + if not chunk.startswith("data:"): + continue + data_str = chunk[len("data:"):].strip() + if not data_str or data_str == "[DONE]": + continue + try: + data = json.loads(data_str) + except json.JSONDecodeError: + continue + delta = data.get("choices", [{}])[0].get("delta", {}) + if "content" in delta and delta["content"]: + content += delta["content"] + return content + + +class TestOpenAIWebSearchNoLeak: + """ + Verifies an intercepted web_search call never leaks to the OpenAI client as a + tool_call. On failure it degrades to assistant content. Because the OpenAI + non-streaming collector reuses this generator, fixing it covers both modes. + """ + + WS_TOOL_USE = { + "id": "call_ws", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "python"}'}, + } + + SAMPLE_RESULTS = { + "results": [{"title": "Python", "url": "https://python.org", "snippet": "Official"}], + "totalResults": 1, + "query": "python", + } + + @pytest.mark.asyncio + async def test_streaming_failure_degrades_to_content_no_tool_call( + self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager + ): + """ + What it does: On MCP failure, the generator emits assistant content + (an unavailable note) and NO web_search tool_call. + Purpose: PRIMARY regression guard - OpenAI clients must not receive an + unknown web_search tool_call. + """ + async def mock_parse_kiro_stream(*args, **kwargs): + yield KiroEvent(type="tool_use", tool_use=self.WS_TOOL_USE) + + chunks = [] + with patch("kiro.streaming_openai.parse_kiro_stream", mock_parse_kiro_stream): + with patch("kiro.streaming_openai.parse_bracket_tool_calls", return_value=[]): + with patch( + "kiro.mcp_tools.call_kiro_mcp_api", + new=AsyncMock(return_value=(None, None)), + ): + async for chunk in stream_kiro_to_openai( + mock_http_client, mock_response, "claude-opus-4.8", + mock_model_cache, mock_auth_manager + ): + chunks.append(chunk) + + joined = "".join(chunks) + # No web_search tool_call leaked to the client + assert '"tool_calls"' not in joined + + # Reconstruct the assistant content from the SSE deltas (word may be split + # across chunk boundaries, so check the reassembled text, not raw SSE). + content = _reconstruct_openai_content(chunks) + print(f"Reconstructed content: {content}") + assert "unavailable" in content.lower() + + @pytest.mark.asyncio + async def test_streaming_success_emits_content_no_tool_call( + self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager + ): + """ + What it does: On success, the generator emits the search summary as + content, not as a tool_call. + Purpose: Confirm the happy path stays content-only for OpenAI. + """ + async def mock_parse_kiro_stream(*args, **kwargs): + yield KiroEvent(type="tool_use", tool_use=self.WS_TOOL_USE) + + chunks = [] + with patch("kiro.streaming_openai.parse_kiro_stream", mock_parse_kiro_stream): + with patch("kiro.streaming_openai.parse_bracket_tool_calls", return_value=[]): + with patch( + "kiro.mcp_tools.call_kiro_mcp_api", + new=AsyncMock(return_value=("srvtoolu_ok", self.SAMPLE_RESULTS)), + ): + async for chunk in stream_kiro_to_openai( + mock_http_client, mock_response, "claude-opus-4.8", + mock_model_cache, mock_auth_manager + ): + chunks.append(chunk) + + joined = "".join(chunks) + assert '"tool_calls"' not in joined + # Search summary content present (reconstruct - URL may span chunks) + content = _reconstruct_openai_content(chunks) + print(f"Reconstructed content: {content}") + assert "python.org" in content From dc8be3b4600ff3f8916240f514769bb2efcbcc62 Mon Sep 17 00:00:00 2001 From: coderhisham Date: Thu, 25 Jun 2026 12:32:17 +0530 Subject: [PATCH 4/6] fix(anthropic): accept inline system role and server-side web_search blocks (#190, #219) Claude Code sends two shapes the strict models rejected with 422: 1. Inline messages with role 'system' in the messages array. AnthropicMessage.role is now a free-form string (mirroring OpenAI ChatMessage); inline system messages are hoisted into the top-level system prompt (separate_inline_system_messages), matching the OpenAI adapter and avoiding spurious user turns. 2. Follow-up turns echoing back the server_tool_use / web_search_tool_result blocks the gateway emitted. These block types are now part of the ContentBlock union, and convert_anthropic_content_to_text folds an echoed web_search_tool_result into text so search grounding survives conversion. Adds model-validation, converter, and route tests for both shapes. --- kiro/converters_anthropic.py | 142 ++++++++- kiro/models_anthropic.py | 70 ++++- tests/unit/test_converters_anthropic.py | 391 ++++++++++++++++++++++++ tests/unit/test_models_anthropic.py | 191 ++++++++++++ tests/unit/test_routes_anthropic.py | 23 +- 5 files changed, 798 insertions(+), 19 deletions(-) diff --git a/kiro/converters_anthropic.py b/kiro/converters_anthropic.py index ea297dc1..593da1a0 100644 --- a/kiro/converters_anthropic.py +++ b/kiro/converters_anthropic.py @@ -24,7 +24,7 @@ to the unified format used by converters_core.py. """ -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple from loguru import logger @@ -45,6 +45,47 @@ ) +def _render_web_search_result_block(block: Any) -> str: + """ + Render a web_search_tool_result content block as readable text. + + When a client echoes a prior web_search turn back to the gateway, the raw + result blocks carry the grounding the assistant used. We fold them into text + so that context survives conversion to Kiro (which has no native + web_search_tool_result representation). + + Args: + block: A web_search_tool_result block (dict or Pydantic model). + + Returns: + A readable summary string, or "" if there is nothing to render. + """ + content = block.get("content") if isinstance(block, dict) else getattr(block, "content", None) + + # Failure case: the native error object. + if isinstance(content, dict): + if content.get("type") == "web_search_tool_result_error": + return "\n\n(Previous web search was unavailable.)\n\n" + return "" + + if not isinstance(content, list): + return "" + + lines = [] + for item in content: + item_dict = item if isinstance(item, dict) else getattr(item, "__dict__", {}) + title = item_dict.get("title", "") + url = item_dict.get("url", "") + snippet = item_dict.get("encrypted_content", "") or item_dict.get("snippet", "") + if title or url or snippet: + lines.append(f"- {title} ({url})\n {snippet}".rstrip()) + + if not lines: + return "" + + return "\n\nPrevious search results:\n" + "\n".join(lines) + "\n\n" + + def convert_anthropic_content_to_text(content: Any) -> str: """ Extracts text content from Anthropic message content. @@ -53,6 +94,10 @@ def convert_anthropic_content_to_text(content: Any) -> str: - String: "Hello, world!" - List of content blocks: [{"type": "text", "text": "Hello"}] + Server-side web_search blocks (server_tool_use / web_search_tool_result) that + a client echoes back from a prior turn are folded into text so the search + grounding is preserved through conversion to Kiro. + Args: content: Anthropic message content @@ -66,10 +111,16 @@ def convert_anthropic_content_to_text(content: Any) -> str: text_parts = [] for block in content: if isinstance(block, dict): - if block.get("type") == "text": + block_type = block.get("type") + if block_type == "text": text_parts.append(block.get("text", "")) - elif hasattr(block, "type") and block.type == "text": - text_parts.append(block.text) + elif block_type == "web_search_tool_result": + text_parts.append(_render_web_search_result_block(block)) + elif hasattr(block, "type"): + if block.type == "text": + text_parts.append(getattr(block, "text", "")) + elif block.type == "web_search_tool_result": + text_parts.append(_render_web_search_result_block(block)) return "".join(text_parts) return str(content) if content else "" @@ -426,6 +477,55 @@ def extract_thinking_config_from_anthropic(request: AnthropicMessagesRequest) -> return ThinkingConfig(enabled=True, budget_tokens=None) +def separate_inline_system_messages( + messages: List[AnthropicMessage], +) -> Tuple[List[str], List[AnthropicMessage]]: + """ + Separates inline ``system`` role messages from conversation turns. + + Some clients (notably Claude Code) inject messages with ``role == "system"`` + directly into the ``messages`` array (e.g. SessionStart hook context), even + though the Anthropic spec keeps ``system`` as a separate top-level field. + + To stay consistent with the OpenAI adapter (which extracts all ``system`` + messages into the system prompt, see ``convert_openai_messages_to_unified``), + we peel these inline system messages out here and return their text so the + caller can merge it into the system prompt. This keeps system instructions + out of the conversation history and avoids turning them into ``user`` turns + (which would trigger synthetic alternation placeholders downstream). + + Args: + messages: List of Anthropic messages (possibly containing inline system) + + Returns: + Tuple of: + - List of extracted system text fragments, in original order + - List of remaining conversation messages (user/assistant and any other + non-system roles, which are normalized later by the core layer) + + Example: + >>> msgs = [ + ... AnthropicMessage(role="user", content="Hi"), + ... AnthropicMessage(role="system", content="Hook context"), + ... ] + >>> parts, convo = separate_inline_system_messages(msgs) + >>> parts + ['Hook context'] + >>> [m.role for m in convo] + ['user'] + """ + inline_system_parts: List[str] = [] + conversation_messages: List[AnthropicMessage] = [] + + for msg in messages: + if msg.role == "system": + inline_system_parts.append(convert_anthropic_content_to_text(msg.content)) + else: + conversation_messages.append(msg) + + return inline_system_parts, conversation_messages + + def anthropic_to_kiro( request: AnthropicMessagesRequest, conversation_id: str, profile_arn: str ) -> dict: @@ -435,7 +535,9 @@ def anthropic_to_kiro( This is the main entry point for Anthropic → Kiro conversion. Key differences from OpenAI: - - System prompt is a separate field (not in messages) + - System prompt is a separate top-level field. Inline ``system`` messages in + the ``messages`` array (sent by some clients like Claude Code) are merged + into that system prompt, matching the OpenAI adapter's behavior. - Content can be string or list of content blocks - Tool format uses input_schema instead of parameters @@ -450,15 +552,37 @@ def anthropic_to_kiro( Raises: ValueError: If there are no messages to send """ - # Convert messages to unified format - unified_messages = convert_anthropic_messages(request.messages) + # Separate inline system messages (e.g. Claude Code hook context) from the + # actual conversation turns. This mirrors the OpenAI adapter, which extracts + # all system messages into the system prompt rather than the history. + inline_system_parts, conversation_messages = separate_inline_system_messages( + request.messages + ) + + if inline_system_parts: + logger.debug( + f"Merged {len(inline_system_parts)} inline system message(s) into system prompt" + ) + + # Convert conversation messages to unified format + unified_messages = convert_anthropic_messages(conversation_messages) # Convert tools to unified format unified_tools = convert_anthropic_tools(request.tools) - # System prompt is already separate in Anthropic format! - # It can be a string or list of content blocks (for prompt caching) + # System prompt comes from the top-level field first, then any inline system + # messages are appended (preserving Claude Code's intended ordering). system_prompt = extract_system_prompt(request.system) + if inline_system_parts: + inline_system_text = "\n\n".join( + part for part in inline_system_parts if part + ) + if inline_system_text: + system_prompt = ( + f"{system_prompt}\n\n{inline_system_text}" + if system_prompt + else inline_system_text + ) # Get model ID for Kiro API (normalizes + resolves hidden models) # Pass-through principle: we normalize and send to Kiro, Kiro decides if valid diff --git a/kiro/models_anthropic.py b/kiro/models_anthropic.py index c63d60ba..27cf8dc7 100644 --- a/kiro/models_anthropic.py +++ b/kiro/models_anthropic.py @@ -162,7 +162,53 @@ class ImageContentBlock(BaseModel): source: Union[Base64ImageSource, URLImageSource] -# Union type for all content blocks (including images and thinking) +class ServerToolUseContentBlock(BaseModel): + """ + Server-side tool use block (e.g. web_search), in Anthropic format. + + The gateway emits these for web_search (Path A and Path B). When a client + like Claude Code echoes the conversation back on the next turn, the assistant + message contains this block, so we must accept it on input as well. + + Attributes: + type: Always "server_tool_use" + id: Tool use ID (links to the matching web_search_tool_result) + name: Server-side tool name (e.g. "web_search") + input: Tool input (e.g. {"query": "..."}) + """ + + type: Literal["server_tool_use"] = "server_tool_use" + id: str + name: str + input: Dict[str, Any] = {} + + model_config = {"extra": "allow"} + + +class WebSearchToolResultContentBlock(BaseModel): + """ + web_search tool result block, in Anthropic format. + + Emitted by the gateway alongside server_tool_use and echoed back by clients + on subsequent turns. The content is either a list of web_search_result items + (success) or a web_search_tool_result_error object (failure), so it is kept + permissive here. + + Attributes: + type: Always "web_search_tool_result" + tool_use_id: ID linking back to the server_tool_use block + content: List of result items, an error object, or a string + """ + + type: Literal["web_search_tool_result"] = "web_search_tool_result" + tool_use_id: str + content: Union[str, List[Dict[str, Any]], Dict[str, Any], None] = None + + model_config = {"extra": "allow"} + + +# Union type for all content blocks (including images, thinking, and the +# server-side web_search blocks the gateway both emits and accepts back) ContentBlock = Union[ TextContentBlock, ThinkingContentBlock, @@ -170,6 +216,8 @@ class ImageContentBlock(BaseModel): ToolUseContentBlock, ToolResultContentBlock, ToolReferenceContentBlock, + ServerToolUseContentBlock, + WebSearchToolResultContentBlock, ] @@ -182,12 +230,28 @@ class AnthropicMessage(BaseModel): """ Message in Anthropic format. + Although the Anthropic spec only documents ``user`` and ``assistant`` roles + in the ``messages`` array (``system`` is a separate top-level field), some + clients (notably Claude Code) inject inline messages with non-standard roles + such as ``system`` mid-conversation. Rejecting these with a 422 would break + those clients, so the role is accepted as a free-form string here, mirroring + the OpenAI ``ChatMessage`` model. + + Downstream handling (in ``anthropic_to_kiro``): + - Inline ``system`` messages are hoisted into the top-level system prompt, + matching the OpenAI adapter's behavior (keeps system instructions out of + the conversation history). + - Any other non-standard role (e.g. ``developer``) is normalized to ``user`` + by ``normalize_message_roles`` in ``converters_core.py``. + Attributes: - role: Message role (user or assistant) + role: Message role. Standard values are ``user`` and ``assistant``; + inline ``system`` is hoisted to the system prompt, and other + non-standard roles are normalized to ``user`` downstream. content: Message content (string or list of content blocks) """ - role: Literal["user", "assistant"] + role: str content: Union[str, List[ContentBlock]] model_config = {"extra": "allow"} diff --git a/tests/unit/test_converters_anthropic.py b/tests/unit/test_converters_anthropic.py index b2b37402..af947d18 100644 --- a/tests/unit/test_converters_anthropic.py +++ b/tests/unit/test_converters_anthropic.py @@ -24,6 +24,7 @@ convert_anthropic_messages, convert_anthropic_tools, anthropic_to_kiro, + separate_inline_system_messages, extract_thinking_config_from_anthropic, ) from kiro.converters_core import UnifiedMessage, UnifiedTool @@ -980,6 +981,33 @@ def test_converts_simple_assistant_message(self): assert result[0].role == "assistant" assert result[0].content == "Hi there!" + def test_preserves_inline_system_role(self): + """ + What it does: Verifies an inline system message keeps its role through conversion. + Purpose: convert_anthropic_messages should pass the role through unchanged; + normalization to 'user' happens later in build_kiro_payload. This + guards the boundary between the adapter and the core normalization. + """ + print("Setup: Inline system message (Claude Code style)...") + messages = [ + AnthropicMessage(role="user", content="Hello"), + AnthropicMessage( + role="system", + content=[{"type": "text", "text": "be concise"}], + ), + ] + + print("Action: Converting messages...") + result = convert_anthropic_messages(messages) + + print(f"Result roles: {[m.role for m in result]}") + assert len(result) == 2 + assert result[1].role == "system" + assert result[1].content == "be concise" + # System messages carry no tool data + assert result[1].tool_calls is None + assert result[1].tool_results is None + def test_converts_user_message_with_content_blocks(self): """ What it does: Verifies conversion of user message with content blocks. @@ -1441,6 +1469,123 @@ def test_handles_tool_without_description(self): # ================================================================================================== +class TestSeparateInlineSystemMessages: + """Tests for separate_inline_system_messages function.""" + + def test_no_system_messages_returns_all_as_conversation(self): + """ + What it does: Verifies messages without system role are untouched. + Purpose: Ensure the common path adds no system parts and keeps all turns. + """ + print("Setup: user + assistant messages, no system...") + messages = [ + AnthropicMessage(role="user", content="Hi"), + AnthropicMessage(role="assistant", content="Hello"), + ] + + print("Action: Separating inline system messages...") + parts, convo = separate_inline_system_messages(messages) + + print(f"Result parts={parts}, convo_roles={[m.role for m in convo]}") + assert parts == [] + assert len(convo) == 2 + assert [m.role for m in convo] == ["user", "assistant"] + + def test_extracts_single_inline_system_message(self): + """ + What it does: Verifies a single inline system message is extracted. + Purpose: PRIMARY behavior - Claude Code SessionStart hook context. + """ + print("Setup: user + inline system message...") + messages = [ + AnthropicMessage(role="user", content="Hi"), + AnthropicMessage(role="system", content="Hook context"), + ] + + print("Action: Separating inline system messages...") + parts, convo = separate_inline_system_messages(messages) + + print(f"Result parts={parts}, convo_roles={[m.role for m in convo]}") + assert parts == ["Hook context"] + assert len(convo) == 1 + assert convo[0].role == "user" + + def test_extracts_multiple_inline_system_messages_in_order(self): + """ + What it does: Verifies multiple system messages are extracted in order. + Purpose: Ensure ordering is preserved for downstream merging. + """ + print("Setup: multiple inline system messages interleaved...") + messages = [ + AnthropicMessage(role="system", content="First"), + AnthropicMessage(role="user", content="Hi"), + AnthropicMessage(role="system", content="Second"), + ] + + print("Action: Separating inline system messages...") + parts, convo = separate_inline_system_messages(messages) + + print(f"Result parts={parts}, convo_roles={[m.role for m in convo]}") + assert parts == ["First", "Second"] + assert len(convo) == 1 + assert convo[0].role == "user" + + def test_extracts_system_message_with_content_blocks(self): + """ + What it does: Verifies system content as list-of-blocks is flattened to text. + Purpose: Claude Code sends system-reminder as text content blocks. + """ + print("Setup: system message with content blocks...") + messages = [ + AnthropicMessage(role="user", content="Hi"), + AnthropicMessage( + role="system", + content=[ + {"type": "text", "text": "Block A "}, + {"type": "text", "text": "Block B"}, + ], + ), + ] + + print("Action: Separating inline system messages...") + parts, convo = separate_inline_system_messages(messages) + + print(f"Result parts={parts}") + assert parts == ["Block A Block B"] + assert len(convo) == 1 + + def test_non_system_unknown_roles_remain_in_conversation(self): + """ + What it does: Verifies non-system unknown roles are NOT extracted here. + Purpose: Only 'system' is hoisted; other roles (e.g. 'developer') stay in + the conversation and are normalized to 'user' by the core layer. + """ + print("Setup: developer role message (not system)...") + messages = [ + AnthropicMessage(role="developer", content="Guidelines"), + AnthropicMessage(role="user", content="Hi"), + ] + + print("Action: Separating inline system messages...") + parts, convo = separate_inline_system_messages(messages) + + print(f"Result parts={parts}, convo_roles={[m.role for m in convo]}") + assert parts == [] + assert len(convo) == 2 + assert convo[0].role == "developer" + + def test_empty_messages_returns_empty(self): + """ + What it does: Verifies empty input is handled gracefully. + Purpose: Defensive boundary check. + """ + print("Action: Separating inline system messages on empty list...") + parts, convo = separate_inline_system_messages([]) + + assert parts == [] + assert convo == [] + + class TestAnthropicToKiro: """Tests for anthropic_to_kiro function - main entry point.""" @@ -1538,6 +1683,252 @@ def test_includes_tools(self): assert len(tools) == 1 assert tools[0]["toolSpecification"]["name"] == "get_weather" + def test_handles_inline_system_message_from_claude_code(self): + """ + What it does: Verifies a request with an inline 'system' role message + (as sent by Claude Code) merges that content into the + system prompt rather than a conversation turn. + Purpose: End-to-end regression test for the 422 literal_error + ("Input should be 'user' or 'assistant'", input='system'). + The inline system message must be hoisted into the system + prompt (consistent with the OpenAI adapter), not turned into a + user turn. + """ + print("Setup: Request mimicking Claude Code inline system message...") + request = AnthropicMessagesRequest( + model="claude-opus-4-8", + messages=[ + AnthropicMessage(role="user", content="First user turn"), + AnthropicMessage( + role="system", + content=[ + {"type": "text", "text": "be concise"} + ], + ), + ], + max_tokens=1024, + ) + + print("Action: Converting to Kiro payload...") + with patch( + "kiro.converters_anthropic.get_model_id_for_kiro", + return_value="claude-opus-4.8", + ): + with patch("kiro.converters_core.FAKE_REASONING_ENABLED", False): + result = anthropic_to_kiro(request, "conv-sys", "arn:aws:test") + + print(f"Result: {result}") + # Only the user turn remains as conversation, so no history is produced + history = result["conversationState"].get("history", []) + current_content = result["conversationState"]["currentMessage"][ + "userInputMessage" + ]["content"] + print(f"History: {history}") + print(f"Current content: {current_content}") + + # System content is merged into the (only) user message, history stays empty + assert history == [] + assert "be concise" in current_content + assert "First user turn" in current_content + + def test_inline_system_message_does_not_create_user_turn(self): + """ + What it does: Verifies an inline system message between two user turns is + hoisted to the system prompt and does NOT add an extra turn. + Purpose: Guards against the normalize-to-user behavior that would turn + hook context into a conversation turn and force synthetic + assistant placeholders for alternation. + """ + print("Setup: Request with user / inline-system / user...") + request = AnthropicMessagesRequest( + model="claude-opus-4-8", + messages=[ + AnthropicMessage(role="user", content="First user turn"), + AnthropicMessage(role="system", content="Inline hook context"), + AnthropicMessage(role="user", content="Second user turn"), + ], + max_tokens=1024, + ) + + print("Action: Converting to Kiro payload...") + with patch( + "kiro.converters_anthropic.get_model_id_for_kiro", + return_value="claude-opus-4.8", + ): + with patch("kiro.converters_core.FAKE_REASONING_ENABLED", False): + result = anthropic_to_kiro(request, "conv-sys2", "arn:aws:test") + + print(f"Result: {result}") + history = result["conversationState"].get("history", []) + current_content = result["conversationState"]["currentMessage"][ + "userInputMessage" + ]["content"] + print(f"History length: {len(history)}") + + # The inline system message was removed (hoisted to the system prompt), + # NOT converted to a user turn. The two surrounding user turns are then + # adjacent and merge into a single user message - so there is no history + # and no synthetic assistant placeholder. + assert history == [] + assert "Inline hook context" in current_content + assert "First user turn" in current_content + assert "Second user turn" in current_content + # No synthetic placeholder was needed for alternation + assert "(empty placeholder)" not in current_content + + def test_merges_top_level_and_inline_system_in_order(self): + """ + What it does: Verifies top-level system precedes inline system content. + Purpose: Preserve Claude Code's intended system prompt ordering when both + a top-level system field and inline system messages are present. + """ + print("Setup: Request with both top-level and inline system prompts...") + request = AnthropicMessagesRequest( + model="claude-opus-4-8", + messages=[ + AnthropicMessage(role="user", content="Hello"), + AnthropicMessage(role="system", content="Inline system context"), + ], + max_tokens=1024, + system="Top-level system prompt", + ) + + print("Action: Converting to Kiro payload...") + with patch( + "kiro.converters_anthropic.get_model_id_for_kiro", + return_value="claude-opus-4.8", + ): + with patch("kiro.converters_core.FAKE_REASONING_ENABLED", False): + result = anthropic_to_kiro(request, "conv-order", "arn:aws:test") + + current_content = result["conversationState"]["currentMessage"][ + "userInputMessage" + ]["content"] + print(f"Current content: {current_content}") + + top_level_index = current_content.index("Top-level system prompt") + inline_index = current_content.index("Inline system context") + assert top_level_index < inline_index + assert "Hello" in current_content + + def test_handles_echoed_web_search_blocks(self): + """ + What it does: Verifies a follow-up request whose assistant turn echoes + server_tool_use + web_search_tool_result blocks validates + and converts cleanly. + Purpose: Regression for the 422 that occurred after a successful + web_search - Claude Code sends the search blocks back in + history, and the gateway (which emitted them) must accept them. + Also verifies the search grounding survives as text. + """ + print("Setup: Follow-up request with echoed web_search blocks...") + request = AnthropicMessagesRequest( + model="claude-opus-4-8", + max_tokens=1024, + messages=[ + AnthropicMessage(role="user", content="Do a websearch for coderhisham"), + AnthropicMessage( + role="assistant", + content=[ + { + "id": "srvtoolu_abc", + "type": "server_tool_use", + "name": "web_search", + "input": {"query": "coderhisham"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_abc", + "content": [ + { + "type": "web_search_result", + "title": "coderhisham GitHub", + "url": "https://github.com/coderhisham", + "encrypted_content": "Developer & Tech Enthusiast", + "page_age": None, + } + ], + }, + {"type": "text", "text": "coderhisham is Muhammed Hisham."}, + ], + ), + AnthropicMessage(role="user", content="who is coderhisham"), + ], + ) + + # The assistant turn must validate with the server-side block types + block_types = [getattr(b, "type", None) for b in request.messages[1].content] + print(f"Assistant block types: {block_types}") + assert block_types == ["server_tool_use", "web_search_tool_result", "text"] + + print("Action: Converting to Kiro payload...") + with patch( + "kiro.converters_anthropic.get_model_id_for_kiro", + return_value="claude-opus-4.8", + ): + with patch("kiro.converters_core.FAKE_REASONING_ENABLED", False): + result = anthropic_to_kiro(request, "conv-ws", "arn:aws:test") + + print(f"Result: {result}") + assert "conversationState" in result + # Search grounding (the result url) is preserved in the converted history + history = result["conversationState"].get("history", []) + assistant_text = "".join( + entry.get("assistantResponseMessage", {}).get("content", "") for entry in history + ) + print(f"Assistant history text: {assistant_text}") + assert "github.com/coderhisham" in assistant_text + assert "coderhisham is Muhammed Hisham." in assistant_text + + def test_handles_echoed_web_search_error_block(self): + """ + What it does: Verifies an echoed web_search_tool_result ERROR block also + validates and converts (failure case round-trip). + Purpose: The gateway emits an error result on MCP failure; the client may + echo it back, so it must be accepted too. + """ + print("Setup: Follow-up request with echoed web_search error block...") + request = AnthropicMessagesRequest( + model="claude-opus-4-8", + max_tokens=1024, + messages=[ + AnthropicMessage(role="user", content="search please"), + AnthropicMessage( + role="assistant", + content=[ + { + "id": "srvtoolu_err", + "type": "server_tool_use", + "name": "web_search", + "input": {"query": "x"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_err", + "content": { + "type": "web_search_tool_result_error", + "error_code": "unavailable", + }, + }, + ], + ), + AnthropicMessage(role="user", content="continue"), + ], + ) + + block_types = [getattr(b, "type", None) for b in request.messages[1].content] + assert block_types == ["server_tool_use", "web_search_tool_result"] + + with patch( + "kiro.converters_anthropic.get_model_id_for_kiro", + return_value="claude-opus-4.8", + ): + with patch("kiro.converters_core.FAKE_REASONING_ENABLED", False): + result = anthropic_to_kiro(request, "conv-ws-err", "arn:aws:test") + + print(f"Result: {result}") + assert "conversationState" in result + def test_builds_history_for_multi_turn(self): """ What it does: Verifies building of history for multi-turn conversation. diff --git a/tests/unit/test_models_anthropic.py b/tests/unit/test_models_anthropic.py index 7c364c9c..d1285458 100644 --- a/tests/unit/test_models_anthropic.py +++ b/tests/unit/test_models_anthropic.py @@ -524,10 +524,201 @@ def test_message_with_url_image_validates(self): assert message.content[1].source.url == "https://example.com/image.jpg" +# ================================================================================================== +# Tests for AnthropicMessage with Non-Standard Roles (inline system / developer messages) +# ================================================================================================== + +class TestAnthropicMessageNonStandardRoles: + """ + Tests for AnthropicMessage accepting non-standard roles in the messages array. + + Some clients (notably Claude Code) inject inline messages with a 'system' + role mid-conversation, even though the Anthropic spec keeps 'system' as a + separate top-level field. Previously the strict Literal["user", "assistant"] + role rejected these with a 422 ValidationError before the request could reach + the conversion pipeline (which normalizes such roles to 'user'). + + These tests verify the role field is now permissive while standard roles + continue to work. + """ + + def test_inline_system_role_string_content_validates(self): + """ + What it does: Verifies AnthropicMessage accepts role='system' with string content. + Purpose: PRIMARY test for the Claude Code inline system message 422 fix. + + Before the fix this raised: + literal_error: Input should be 'user' or 'assistant', input='system' + """ + print("Setup: Creating AnthropicMessage with inline system role...") + message = AnthropicMessage(role="system", content="You are a helpful assistant.") + + print(f"Comparing role: Expected 'system', Got '{message.role}'") + assert message.role == "system" + assert message.content == "You are a helpful assistant." + + def test_inline_system_role_list_content_validates(self): + """ + What it does: Verifies role='system' works with list-of-blocks content. + Purpose: Claude Code sends system-reminder content as text content blocks. + """ + print("Setup: Creating system role message with content blocks...") + message = AnthropicMessage( + role="system", + content=[{"type": "text", "text": "context"}], + ) + + print(f"Comparing role: Expected 'system', Got '{message.role}'") + assert message.role == "system" + assert len(message.content) == 1 + assert message.content[0].type == "text" + + def test_developer_role_validates(self): + """ + What it does: Verifies role='developer' (OpenAI o1-style) is accepted. + Purpose: Ensure the whole class of non-standard roles is handled, not just 'system'. + """ + print("Setup: Creating AnthropicMessage with developer role...") + message = AnthropicMessage(role="developer", content="Follow these guidelines.") + + print(f"Comparing role: Expected 'developer', Got '{message.role}'") + assert message.role == "developer" + + def test_standard_user_and_assistant_roles_still_valid(self): + """ + What it does: Verifies the permissive role field does not break standard roles. + Purpose: Guard against regression of the common path. + """ + print("Setup: Creating standard user and assistant messages...") + user_msg = AnthropicMessage(role="user", content="Hello") + assistant_msg = AnthropicMessage(role="assistant", content="Hi there") + + print(f"Comparing user role: Expected 'user', Got '{user_msg.role}'") + assert user_msg.role == "user" + assert assistant_msg.role == "assistant" + + def test_missing_role_still_raises(self): + """ + What it does: Verifies role is still required. + Purpose: Permissiveness on value must not make the field optional. + """ + print("Setup: Attempting to create AnthropicMessage without role...") + print("Action: Creating model (should raise ValidationError)...") + with pytest.raises(ValidationError) as exc_info: + AnthropicMessage(content="No role here") + + print(f"ValidationError raised: {exc_info.value}") + assert "role" in str(exc_info.value) + + def test_request_with_inline_system_message_validates(self): + """ + What it does: Verifies a full request with an inline system message validates. + Purpose: Reproduces the exact Claude Code payload shape that caused the 422 + (system message at index 1 in the messages array). + """ + print("Setup: Creating request with user + inline system + user messages...") + request = AnthropicMessagesRequest( + model="claude-opus-4-8", + max_tokens=1024, + messages=[ + AnthropicMessage(role="user", content="First user turn"), + AnthropicMessage( + role="system", + content=[{"type": "text", "text": "be concise"}], + ), + AnthropicMessage(role="user", content="Second user turn"), + ], + ) + + print(f"Comparing message count: Expected 3, Got {len(request.messages)}") + assert len(request.messages) == 3 + print(f"Comparing messages[1].role: Expected 'system', Got '{request.messages[1].role}'") + assert request.messages[1].role == "system" + + # ================================================================================================== # Tests for AnthropicMessagesRequest with Image Content # ================================================================================================== +class TestAnthropicMessageServerToolBlocks: + """ + Tests that the assistant content union accepts the server-side web_search + blocks the gateway emits (server_tool_use, web_search_tool_result), so a + client echoing them back on the next turn does not trigger a 422. + """ + + def test_accepts_server_tool_use_and_result_blocks(self): + """ + What it does: Verifies an assistant message with server_tool_use + + web_search_tool_result (success) validates. + Purpose: PRIMARY regression for the post-web_search 422. + """ + print("Setup: assistant message echoing web_search blocks...") + message = AnthropicMessage( + role="assistant", + content=[ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "x"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + {"type": "web_search_result", "title": "T", "url": "https://e.com", "encrypted_content": "s"} + ], + }, + {"type": "text", "text": "answer"}, + ], + ) + types = [b.type for b in message.content] + print(f"Block types: {types}") + assert types == ["server_tool_use", "web_search_tool_result", "text"] + + def test_accepts_web_search_tool_result_error_content(self): + """ + What it does: Verifies a web_search_tool_result with an error object validates. + Purpose: The failure result the gateway emits must round-trip too. + """ + print("Setup: assistant message with web_search error result...") + message = AnthropicMessage( + role="assistant", + content=[ + {"type": "server_tool_use", "id": "srvtoolu_2", "name": "web_search", "input": {}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_2", + "content": {"type": "web_search_tool_result_error", "error_code": "unavailable"}, + }, + ], + ) + result_block = message.content[1] + assert result_block.type == "web_search_tool_result" + assert result_block.content["error_code"] == "unavailable" + + def test_full_request_with_echoed_web_search_validates(self): + """ + What it does: Verifies a full request with an echoed web_search turn validates. + Purpose: Reproduce the exact shape that produced the 422. + """ + print("Setup: full request with echoed web_search assistant turn...") + request = AnthropicMessagesRequest( + model="claude-opus-4-8", + max_tokens=1024, + messages=[ + AnthropicMessage(role="user", content="search"), + AnthropicMessage( + role="assistant", + content=[ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {"query": "q"}}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": []}, + {"type": "text", "text": "done"}, + ], + ), + AnthropicMessage(role="user", content="follow up"), + ], + ) + assert len(request.messages) == 3 + assert request.messages[1].content[0].type == "server_tool_use" + + class TestAnthropicMessagesRequestWithImages: """Tests for full AnthropicMessagesRequest with image content.""" diff --git a/tests/unit/test_routes_anthropic.py b/tests/unit/test_routes_anthropic.py index d156f66a..bdec516c 100644 --- a/tests/unit/test_routes_anthropic.py +++ b/tests/unit/test_routes_anthropic.py @@ -337,25 +337,34 @@ def test_validates_invalid_json(self, test_client, valid_proxy_api_key): print(f"Status: {response.status_code}") assert response.status_code == 422 - def test_validates_invalid_role(self, test_client, valid_proxy_api_key): + def test_accepts_non_standard_role(self, test_client, valid_proxy_api_key): """ - What it does: Verifies invalid message role is rejected. - Purpose: Anthropic model strictly validates role (only 'user' or 'assistant'). + What it does: Verifies a non-standard message role (e.g. inline 'system') + is NOT rejected at the validation layer. + Purpose: Some clients (notably Claude Code) inject inline 'system' role + messages into the messages array. These previously triggered a + 422 literal_error. The role field is now permissive and such + roles are normalized to 'user' downstream, so validation must + pass (any non-422 status). """ - print("Action: POST /v1/messages with invalid role...") + print("Action: POST /v1/messages with inline system role...") response = test_client.post( "/v1/messages", headers={"x-api-key": valid_proxy_api_key}, json={ "model": "claude-sonnet-4-5", "max_tokens": 1024, - "messages": [{"role": "invalid_role", "content": "Hello"}] + "messages": [ + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "Inline system reminder"}, + {"role": "user", "content": "Continue"}, + ] } ) print(f"Status: {response.status_code}") - # Anthropic model strictly validates role - only 'user' or 'assistant' allowed - assert response.status_code == 422 + # Role is now permissive - must pass validation (normalized to 'user' downstream) + assert response.status_code != 422 def test_accepts_valid_request_format(self, test_client, valid_proxy_api_key): """ From 98dfd8a291aa91ecd719657d4f0d02cb868b4218 Mon Sep 17 00:00:00 2001 From: coderhisham Date: Thu, 25 Jun 2026 12:38:11 +0530 Subject: [PATCH 5/6] feat(cli): add --stop and --force to free the server port Stopping a stuck/orphaned gateway previously meant manually running lsof + kill. Add cross-platform port utilities (kiro/port_utils.py) and two CLI options: - --stop: stop any process listening on the resolved port, then exit. - --force: if the port is in use at startup, stop the occupant and start anyway. Also adds a preflight check so an occupied port produces a clear, actionable message (with the exact --stop/--force commands) instead of a raw 'address already in use' traceback. Only processes LISTENING on the exact port are targeted, and the current process is always excluded. Adds test_port_utils.py covering detection, lsof/netstat parsing, self-exclusion, graceful/forceful termination, and free_port orchestration. --- kiro/port_utils.py | 180 +++++++++++++++++++++++++++++ main.py | 48 ++++++++ tests/README.md | 1 + tests/unit/test_port_utils.py | 210 ++++++++++++++++++++++++++++++++++ 4 files changed, 439 insertions(+) create mode 100644 kiro/port_utils.py create mode 100644 tests/unit/test_port_utils.py diff --git a/kiro/port_utils.py b/kiro/port_utils.py new file mode 100644 index 00000000..62f78d19 --- /dev/null +++ b/kiro/port_utils.py @@ -0,0 +1,180 @@ +# -*- coding: utf-8 -*- + +# Kiro Gateway +# https://github.com/jwadow/kiro-gateway +# Copyright (C) 2025 Jwadow +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +""" +Cross-platform helpers for inspecting and freeing the server port. + +These power the ``--stop`` and ``--force`` CLI options so a stuck/orphaned +gateway process can be stopped without manually running ``lsof``/``kill``. + +Only processes *listening* on the exact target port are considered, so we never +touch unrelated processes. +""" + +import os +import socket +import subprocess +import time +from typing import List + +from loguru import logger + + +def is_port_in_use(host: str, port: int) -> bool: + """ + Check whether a TCP port is already bound on the given host. + + Uses a bind test (the same operation the server performs at startup), so a + True result reliably predicts an "address already in use" failure. + + Args: + host: Host address the server would bind to (e.g. "0.0.0.0"). + port: TCP port number. + + Returns: + True if the port cannot be bound (already in use), False otherwise. + """ + # "0.0.0.0" binds all interfaces; test against that to mirror the server. + bind_host = host if host not in ("", "*") else "0.0.0.0" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind((bind_host, port)) + return False + except OSError: + return True + + +def find_listening_pids(port: int) -> List[int]: + """ + Find process IDs listening on the given TCP port. + + Cross-platform: uses ``lsof`` on macOS/Linux and ``netstat`` on Windows. + The current process is excluded so callers never target themselves. + + Args: + port: TCP port number. + + Returns: + Sorted list of PIDs listening on the port (empty if none, or if the + platform tool is unavailable). + """ + pids: set[int] = set() + + try: + if os.name == "nt": + result = subprocess.run( + ["netstat", "-ano", "-p", "TCP"], + capture_output=True, text=True, timeout=10, + ) + needle = f":{port}" + for line in result.stdout.splitlines(): + if needle in line and "LISTENING" in line.upper(): + parts = line.split() + if parts: + try: + pids.add(int(parts[-1])) + except ValueError: + continue + else: + result = subprocess.run( + ["lsof", "-nP", f"-iTCP:{port}", "-sTCP:LISTEN", "-t"], + capture_output=True, text=True, timeout=10, + ) + for token in result.stdout.split(): + try: + pids.add(int(token)) + except ValueError: + continue + except FileNotFoundError: + logger.warning( + "Could not inspect port {}: required tool ({}) not found.", + port, "netstat" if os.name == "nt" else "lsof", + ) + except subprocess.TimeoutExpired: + logger.warning("Timed out inspecting processes on port {}.", port) + + pids.discard(os.getpid()) + return sorted(pids) + + +def _terminate_pid(pid: int, timeout: float) -> bool: + """ + Terminate a single PID gracefully, escalating to a forceful kill. + + Args: + pid: Process ID to terminate. + timeout: Seconds to wait for graceful exit before forcing. + + Returns: + True if the process is gone after the attempt, False otherwise. + """ + if os.name == "nt": + subprocess.run( + ["taskkill", "/PID", str(pid), "/F"], + capture_output=True, text=True, + ) + return True + + # POSIX: SIGTERM, wait, then SIGKILL. + try: + os.kill(pid, 15) # SIGTERM + except ProcessLookupError: + return True + except PermissionError: + logger.error("No permission to stop PID {} (try running with sufficient privileges).", pid) + return False + + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + os.kill(pid, 0) # Probe: raises if process is gone. + except ProcessLookupError: + return True + time.sleep(0.1) + + # Still alive - force kill. + try: + os.kill(pid, 9) # SIGKILL + except ProcessLookupError: + return True + return True + + +def free_port(port: int, timeout: float = 5.0) -> List[int]: + """ + Stop every process listening on the given port. + + Args: + port: TCP port to free. + timeout: Seconds to wait for each process to exit gracefully. + + Returns: + List of PIDs that were stopped (empty if the port was already free). + """ + pids = find_listening_pids(port) + if not pids: + return [] + + stopped: List[int] = [] + for pid in pids: + logger.info("Stopping process {} listening on port {}...", pid, port) + if _terminate_pid(pid, timeout): + stopped.append(pid) + return stopped diff --git a/main.py b/main.py index 40cc2c30..04dd4fa9 100644 --- a/main.py +++ b/main.py @@ -36,6 +36,13 @@ # Using uvicorn directly (uvicorn handles its own CLI args) uvicorn main:app --host 0.0.0.0 --port 8000 +Stopping / freeing the port: + # Stop whatever is listening on the resolved port, then exit + python main.py --port 9000 --stop + + # If the port is busy at startup, stop the occupant and start anyway + python main.py --port 9000 --force + Priority: CLI args > Environment variables > Default values """ @@ -88,6 +95,7 @@ from kiro.routes_anthropic import router as anthropic_router from kiro.exceptions import validation_exception_handler from kiro.debug_middleware import DebugLoggerMiddleware +from kiro.port_utils import is_port_in_use, free_port # --- Loguru Configuration --- @@ -637,6 +645,20 @@ def parse_cli_args() -> argparse.Namespace: help=f"Server port (default: {DEFAULT_SERVER_PORT}, env: SERVER_PORT)" ) + parser.add_argument( + "--stop", + action="store_true", + help="Stop any process listening on the resolved port, then exit " + "(frees the port without starting the server)." + ) + + parser.add_argument( + "-f", "--force", + action="store_true", + help="If the port is already in use at startup, stop the occupying " + "process(es) and start anyway." + ) + parser.add_argument( "-v", "--version", action="version", @@ -734,6 +756,17 @@ def print_startup_banner(host: str, port: int) -> None: # Parse CLI arguments first (handles --version, --help without requiring config) args = parse_cli_args() + # --stop: free the resolved port and exit (no credentials needed). + if args.stop: + final_host, final_port = resolve_server_config(args) + logger.info(f"Stopping any process listening on port {final_port}...") + stopped = free_port(final_port) + if stopped: + logger.info(f"Stopped process(es) {stopped}; port {final_port} is now free.") + else: + logger.info(f"Nothing was listening on port {final_port}; already free.") + sys.exit(0) + # Run configuration validation before starting server validate_configuration() @@ -743,6 +776,21 @@ def print_startup_banner(host: str, port: int) -> None: # Resolve final configuration with priority hierarchy final_host, final_port = resolve_server_config(args) + # Preflight: make sure the port is available, with an actionable message. + if is_port_in_use(final_host, final_port): + if args.force: + logger.warning(f"Port {final_port} is in use; --force given, stopping the occupant(s)...") + stopped = free_port(final_port) + logger.info(f"Stopped process(es) {stopped}; continuing startup.") + else: + logger.error( + f"Port {final_port} is already in use. Stop the existing process with:\n" + f" python main.py --port {final_port} --stop\n" + f"or start anyway (stopping the occupant) with:\n" + f" python main.py --port {final_port} --force" + ) + sys.exit(1) + # Print startup banner print_startup_banner(final_host, final_port) diff --git a/tests/README.md b/tests/README.md index a66cee2c..cfa0f153 100644 --- a/tests/README.md +++ b/tests/README.md @@ -93,6 +93,7 @@ tests/ │ ├── test_models_openai.py # OpenAI Pydantic models tests (messages, tools, responses, streaming) │ ├── test_network_errors.py # Network error handling tests │ ├── test_parsers.py # AwsEventStreamParser tests (JSON truncation diagnostics, truncation recovery integration) +│ ├── test_port_utils.py # Port utilities tests (is_port_in_use, listening-PID discovery, free_port for --stop/--force) │ ├── test_routes_anthropic.py # Anthropic API endpoint tests (/v1/messages, truncation recovery, WebSearch, Account System failover) │ ├── test_routes_openai.py # OpenAI API endpoint tests (/v1/chat/completions, truncation recovery, WebSearch, Account System failover) │ ├── test_streaming_anthropic.py # Anthropic streaming response tests (truncation detection, stop_reason priority, initial_response reuse) diff --git a/tests/unit/test_port_utils.py b/tests/unit/test_port_utils.py new file mode 100644 index 00000000..cad8c079 --- /dev/null +++ b/tests/unit/test_port_utils.py @@ -0,0 +1,210 @@ +# -*- coding: utf-8 -*- + +""" +Unit tests for kiro.port_utils. + +Covers port-in-use detection, listening-PID discovery (lsof/netstat parsing, +cross-platform), self-exclusion, graceful termination/escalation, and free_port. +All subprocess and signal interactions are mocked - no real processes are killed. +""" + +import os +import socket +import subprocess +from unittest.mock import patch, MagicMock + +import pytest + +from kiro.port_utils import ( + is_port_in_use, + find_listening_pids, + free_port, + _terminate_pid, +) + + +class TestIsPortInUse: + """Tests for is_port_in_use (real local sockets, no network).""" + + def test_free_port_is_not_in_use(self): + """A port nobody is bound to reports not-in-use.""" + # Find a definitely-free port by binding then releasing it. + s = socket.socket() + s.bind(("127.0.0.1", 0)) + free = s.getsockname()[1] + s.close() + print(f"Testing freed port {free}") + assert is_port_in_use("127.0.0.1", free) is False + + def test_bound_port_is_in_use(self): + """A port held by a live listener reports in-use.""" + s = socket.socket() + s.bind(("127.0.0.1", 0)) + s.listen() + port = s.getsockname()[1] + try: + print(f"Testing bound port {port}") + assert is_port_in_use("127.0.0.1", port) is True + finally: + s.close() + + def test_wildcard_host_is_normalized(self): + """An empty/wildcard host is treated as 0.0.0.0 without error.""" + s = socket.socket() + s.bind(("127.0.0.1", 0)) + free = s.getsockname()[1] + s.close() + # Should not raise and should report free. + assert is_port_in_use("", free) is False + + +class TestFindListeningPids: + """Tests for find_listening_pids parsing and self-exclusion.""" + + def test_parses_lsof_output(self): + """lsof -t output (one PID per line) is parsed to ints.""" + fake = MagicMock(stdout="4242\n4243\n") + with patch("kiro.port_utils.os.name", "posix"): + with patch("kiro.port_utils.subprocess.run", return_value=fake) as run: + with patch("kiro.port_utils.os.getpid", return_value=1): + pids = find_listening_pids(9000) + print(f"PIDs: {pids}, cmd: {run.call_args.args[0]}") + assert pids == [4242, 4243] + assert run.call_args.args[0][0] == "lsof" + + def test_excludes_current_process(self): + """The current PID is never returned (no self-kill).""" + fake = MagicMock(stdout="111\n222\n") + with patch("kiro.port_utils.os.name", "posix"): + with patch("kiro.port_utils.subprocess.run", return_value=fake): + with patch("kiro.port_utils.os.getpid", return_value=111): + pids = find_listening_pids(9000) + print(f"PIDs: {pids}") + assert pids == [222] + + def test_parses_windows_netstat_output(self): + """Windows netstat LISTENING lines are parsed for the PID (last column).""" + netstat = ( + " Proto Local Address Foreign Address State PID\n" + " TCP 0.0.0.0:9000 0.0.0.0:0 LISTENING 7777\n" + " TCP 0.0.0.0:8000 0.0.0.0:0 LISTENING 8888\n" + ) + fake = MagicMock(stdout=netstat) + with patch("kiro.port_utils.os.name", "nt"): + with patch("kiro.port_utils.subprocess.run", return_value=fake) as run: + with patch("kiro.port_utils.os.getpid", return_value=1): + pids = find_listening_pids(9000) + print(f"PIDs: {pids}, cmd: {run.call_args.args[0]}") + assert pids == [7777] # only the :9000 listener + assert run.call_args.args[0][0] == "netstat" + + def test_missing_tool_returns_empty(self): + """If lsof/netstat isn't installed, return [] instead of crashing.""" + with patch("kiro.port_utils.subprocess.run", side_effect=FileNotFoundError()): + pids = find_listening_pids(9000) + assert pids == [] + + def test_timeout_returns_empty(self): + """A hung inspection command degrades to [].""" + with patch( + "kiro.port_utils.subprocess.run", + side_effect=subprocess.TimeoutExpired(cmd="lsof", timeout=10), + ): + pids = find_listening_pids(9000) + assert pids == [] + + def test_ignores_non_integer_tokens(self): + """Garbage tokens in tool output are skipped.""" + fake = MagicMock(stdout="123\nnot_a_pid\n456\n") + with patch("kiro.port_utils.os.name", "posix"): + with patch("kiro.port_utils.subprocess.run", return_value=fake): + with patch("kiro.port_utils.os.getpid", return_value=1): + pids = find_listening_pids(9000) + assert pids == [123, 456] + + +class TestTerminatePid: + """Tests for _terminate_pid graceful/forceful behavior (POSIX).""" + + def test_graceful_termination(self): + """SIGTERM that takes effect (process gone) returns True without SIGKILL.""" + calls = [] + + def fake_kill(pid, sig): + calls.append(sig) + # SIGTERM (15) then probe (0): probe raises -> process gone. + if sig == 0: + raise ProcessLookupError() + + with patch("kiro.port_utils.os.name", "posix"): + with patch("kiro.port_utils.os.kill", side_effect=fake_kill): + result = _terminate_pid(999, timeout=1.0) + print(f"signals: {calls}") + assert result is True + assert 15 in calls # SIGTERM sent + assert 9 not in calls # never escalated + + def test_force_kill_when_still_alive(self): + """A process that ignores SIGTERM gets SIGKILL.""" + sent = [] + + def fake_kill(pid, sig): + sent.append(sig) + # Probe (0) always succeeds -> still alive until SIGKILL. + + with patch("kiro.port_utils.os.name", "posix"): + with patch("kiro.port_utils.os.kill", side_effect=fake_kill): + with patch("kiro.port_utils.time.monotonic", side_effect=[0.0, 0.05, 1.0, 2.0]): + with patch("kiro.port_utils.time.sleep"): + result = _terminate_pid(999, timeout=0.1) + print(f"signals: {sent}") + assert result is True + assert 9 in sent # SIGKILL escalation happened + + def test_already_gone(self): + """If the process is already gone, SIGTERM raising is treated as success.""" + with patch("kiro.port_utils.os.name", "posix"): + with patch("kiro.port_utils.os.kill", side_effect=ProcessLookupError()): + result = _terminate_pid(999, timeout=1.0) + assert result is True + + def test_permission_error_returns_false(self): + """Lack of permission to signal the process is reported as failure.""" + with patch("kiro.port_utils.os.name", "posix"): + with patch("kiro.port_utils.os.kill", side_effect=PermissionError()): + result = _terminate_pid(999, timeout=1.0) + assert result is False + + def test_windows_uses_taskkill(self): + """On Windows, termination shells out to taskkill /F.""" + with patch("kiro.port_utils.os.name", "nt"): + with patch("kiro.port_utils.subprocess.run") as run: + result = _terminate_pid(999, timeout=1.0) + assert result is True + assert run.call_args.args[0][:1] == ["taskkill"] + + +class TestFreePort: + """Tests for free_port orchestration.""" + + def test_no_listeners_returns_empty(self): + """Freeing an already-free port stops nothing.""" + with patch("kiro.port_utils.find_listening_pids", return_value=[]): + assert free_port(9000) == [] + + def test_stops_all_listeners(self): + """Every listening PID is terminated and reported.""" + with patch("kiro.port_utils.find_listening_pids", return_value=[10, 20]): + with patch("kiro.port_utils._terminate_pid", return_value=True) as term: + stopped = free_port(9000) + print(f"stopped: {stopped}") + assert stopped == [10, 20] + assert term.call_count == 2 + + def test_excludes_failed_terminations(self): + """PIDs that could not be stopped are not reported as stopped.""" + with patch("kiro.port_utils.find_listening_pids", return_value=[10, 20]): + with patch("kiro.port_utils._terminate_pid", side_effect=[True, False]): + stopped = free_port(9000) + print(f"stopped: {stopped}") + assert stopped == [10] From 72feef66c774cf116d0d0b092e30d35e20ef2428 Mon Sep 17 00:00:00 2001 From: coderhisham Date: Thu, 25 Jun 2026 12:43:18 +0530 Subject: [PATCH 6/6] fix(server): bound graceful shutdown so Ctrl+C frees the port promptly On Ctrl+C/SIGTERM, uvicorn waits for in-flight connections to drain. With long-lived streaming (STREAMING_READ_TIMEOUT=300), an open SSE connection kept the process alive and the port bound for minutes - which is why a stopped gateway sometimes lingered on the port. Add SHUTDOWN_TIMEOUT (env, default 10s) and pass it to uvicorn as timeout_graceful_shutdown, so shutdown is bounded and the port is freed promptly. A second Ctrl+C still forces an immediate exit. Adds CLI-parsing tests for --stop/--force and config tests for SHUTDOWN_TIMEOUT. --- .env.example | 6 ++++++ kiro/config.py | 9 +++++++++ main.py | 5 +++++ tests/unit/test_config.py | 34 ++++++++++++++++++++++++++++++++++ tests/unit/test_main_cli.py | 35 +++++++++++++++++++++++++++++++++++ 5 files changed, 89 insertions(+) diff --git a/.env.example b/.env.example index da93321b..a8f0d24c 100644 --- a/.env.example +++ b/.env.example @@ -213,6 +213,12 @@ PROXY_API_KEY="my-super-secret-password-123" # Default: 300 seconds (5 minutes) - generous timeout to avoid premature disconnects. # STREAMING_READ_TIMEOUT="300" +# Graceful shutdown timeout (in seconds) for Ctrl+C / SIGTERM. +# Caps how long the server waits for in-flight (streaming) connections to drain +# before exiting, so Ctrl+C frees the port promptly instead of hanging on an open +# SSE stream. A second Ctrl+C forces an immediate exit. Default: 10 seconds. +# SHUTDOWN_TIMEOUT="10" + # =========================================== # FAKE REASONING (Extended Thinking via Tag Injection) # =========================================== diff --git a/kiro/config.py b/kiro/config.py index 911179d9..6aa72725 100644 --- a/kiro/config.py +++ b/kiro/config.py @@ -440,6 +440,15 @@ def _parse_model_context_overrides() -> Dict[str, int]: # Default: 300 seconds (5 minutes) - generous timeout to avoid premature disconnects. STREAMING_READ_TIMEOUT: float = float(os.getenv("STREAMING_READ_TIMEOUT", "300")) +# Graceful shutdown timeout (in seconds) for Ctrl+C / SIGTERM. +# When the server is stopped, uvicorn waits for in-flight connections to drain. +# With long-lived streaming (STREAMING_READ_TIMEOUT=300), an open SSE connection +# would otherwise keep the process alive (and the port bound) for minutes. This +# caps that wait so Ctrl+C frees the port promptly. A second Ctrl+C still forces +# an immediate exit. +# Default: 10 seconds. +SHUTDOWN_TIMEOUT: int = int(os.getenv("SHUTDOWN_TIMEOUT", "10")) + # Maximum number of attempts on first token timeout. # After exhausting all attempts, an error will be returned. # Default: 3 attempts diff --git a/main.py b/main.py index 04dd4fa9..dda63a37 100644 --- a/main.py +++ b/main.py @@ -77,6 +77,7 @@ DEFAULT_SERVER_HOST, DEFAULT_SERVER_PORT, STREAMING_READ_TIMEOUT, + SHUTDOWN_TIMEOUT, HIDDEN_MODELS, MODEL_ALIASES, HIDDEN_FROM_LIST, @@ -795,6 +796,7 @@ def print_startup_banner(host: str, port: int) -> None: print_startup_banner(final_host, final_port) logger.info(f"Starting Uvicorn server on {final_host}:{final_port}...") + logger.debug(f"Graceful shutdown timeout: {SHUTDOWN_TIMEOUT}s (Ctrl+C; press again to force quit)") # Use string reference to avoid double module import uvicorn.run( @@ -802,4 +804,7 @@ def print_startup_banner(host: str, port: int) -> None: host=final_host, port=final_port, log_config=UVICORN_LOG_CONFIG, + # Cap how long shutdown waits for in-flight (streaming) connections so + # Ctrl+C frees the port promptly instead of hanging on open SSE streams. + timeout_graceful_shutdown=SHUTDOWN_TIMEOUT, ) diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 9503299a..139b93ca 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -1116,3 +1116,37 @@ def test_all_fallback_entries_have_model_id(self): from kiro.config import FALLBACK_MODELS assert all("modelId" in m and m["modelId"] for m in FALLBACK_MODELS) + + +class TestShutdownTimeoutConfig: + """Tests for SHUTDOWN_TIMEOUT configuration.""" + + def test_default_shutdown_timeout_is_10(self, monkeypatch): + """ + What it does: Verifies SHUTDOWN_TIMEOUT defaults to 10 seconds. + Purpose: Ctrl+C should free the port within a bounded time by default. + """ + print("Setup: Removing SHUTDOWN_TIMEOUT from environment...") + monkeypatch.delenv("SHUTDOWN_TIMEOUT", raising=False) + + from importlib import reload + import kiro.config as config_module + reload(config_module) + + print(f"SHUTDOWN_TIMEOUT: {config_module.SHUTDOWN_TIMEOUT}") + assert config_module.SHUTDOWN_TIMEOUT == 10 + + def test_shutdown_timeout_env_override(self, monkeypatch): + """ + What it does: Verifies SHUTDOWN_TIMEOUT honors the env override. + Purpose: Operators can tune how long shutdown waits for streams to drain. + """ + print("Setup: Setting SHUTDOWN_TIMEOUT=3...") + monkeypatch.setenv("SHUTDOWN_TIMEOUT", "3") + + from importlib import reload + import kiro.config as config_module + reload(config_module) + + print(f"SHUTDOWN_TIMEOUT: {config_module.SHUTDOWN_TIMEOUT}") + assert config_module.SHUTDOWN_TIMEOUT == 3 diff --git a/tests/unit/test_main_cli.py b/tests/unit/test_main_cli.py index 994ac986..c98d452b 100644 --- a/tests/unit/test_main_cli.py +++ b/tests/unit/test_main_cli.py @@ -33,6 +33,41 @@ def test_default_values_are_none(self): assert args.host is None assert args.port is None + def test_stop_and_force_default_false(self): + """ + What it does: Verifies --stop and --force default to False. + Purpose: Normal startup must not be treated as a stop/force request. + """ + from main import parse_cli_args + with patch.object(sys, 'argv', ['main.py']): + args = parse_cli_args() + print(f"stop={args.stop}, force={args.force}") + assert args.stop is False + assert args.force is False + + def test_stop_flag_parsed(self): + """ + What it does: Verifies --stop sets args.stop True. + Purpose: Enable the 'free the port and exit' path. + """ + from main import parse_cli_args + with patch.object(sys, 'argv', ['main.py', '--port', '9000', '--stop']): + args = parse_cli_args() + print(f"stop={args.stop}, port={args.port}") + assert args.stop is True + assert args.port == 9000 + + def test_force_flag_long_and_short(self): + """ + What it does: Verifies --force and -f both set args.force True. + Purpose: Enable the 'stop occupant then start' path. + """ + from main import parse_cli_args + with patch.object(sys, 'argv', ['main.py', '--force']): + assert parse_cli_args().force is True + with patch.object(sys, 'argv', ['main.py', '-f']): + assert parse_cli_args().force is True + def test_port_argument_long_form(self): """ What it does: Verifies that --port argument is parsed correctly.