Fix Claude Code / Opus 4.8 compatibility and web_search reliability

.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)
 # ===========================================
@@ -194,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)
 # ===========================================

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:

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

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)
@@ -360,6 +440,15 @@ def _get_raw_env_value(var_name: str, env_file: str = ".env") -> Optional[str]:
 # 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