diff --git a/.env.example b/.env.example index 1a57c794..bc4d27c3 100644 --- a/.env.example +++ b/.env.example @@ -11,6 +11,12 @@ # Example: "my-super-secret-password-123" or any secure string PROXY_API_KEY="my-super-secret-password-123" +# Disable proxy authentication entirely (for trusted local networks only) +# When true, the proxy will skip API key validation, allowing unauthenticated access. +# Useful when local clients cannot send custom Authorization headers. +# Default: false +# PROXY_AUTH_DISABLED=false + # =========================================== # OPTION 1: Kiro IDE credentials (JSON file) # =========================================== @@ -56,6 +62,10 @@ PROXY_API_KEY="my-super-secret-password-123" # For kiro-cli (AWS SSO / Builder ID): not needed, will be ignored # PROFILE_ARN="arn:aws:codewhisperer:us-east-1:..." +# Custom path to Kiro IDE profile.json (auto-detected by default) +# Useful if Kiro stores profiles in a non-standard location +# KIRO_PROFILE_FILE="~/Library/Application Support/Kiro/User/globalStorage/kiro.kiroagent/profile.json" + # =========================================== # OPTIONAL # =========================================== @@ -178,9 +188,10 @@ PROXY_API_KEY="my-super-secret-password-123" # Timeout for waiting for the first token from the model (in seconds). # If the model doesn't respond within this time, the request will be cancelled and retried. # This helps handle "stuck" requests when the model takes too long to start responding. -# Default: 15 seconds (recommended for production) -# Set a lower value (e.g., 5-10) for more aggressive retry behavior. -# FIRST_TOKEN_TIMEOUT="15" +# NOTE: Extended thinking models (Opus 4.8/4.7) may need 30-120+ seconds first token. +# Default: 300 seconds (covers extended thinking scenarios) +# Set a lower value (e.g., 15-30) for faster retry on non-thinking models. +# FIRST_TOKEN_TIMEOUT="300" # Maximum number of retry attempts when first token timeout occurs. # After exhausting all attempts, a 504 Gateway Timeout error will be returned. @@ -241,11 +252,67 @@ PROXY_API_KEY="my-super-secret-password-123" # Default: 20 characters (enough for longest tag = 11 chars + some whitespace) # FAKE_REASONING_INITIAL_BUFFER_SIZE=20 +# =========================================== +# NATIVE REASONING (Native Kiro Extended Thinking) +# =========================================== + +# Enable native reasoning pass-through (default: true) +# When enabled, the Kiro backend's real reasoningContentEvent stream is captured +# and surfaced as reasoning_content (OpenAI) / native thinking blocks (Anthropic). +# The legacy fake-reasoning prompt injection is skipped for these requests. +# Set to false to fall back to fake-reasoning behavior. +# NATIVE_REASONING=true + +# Default effort level applied to all requests (when client doesn't specify). +# Valid values: low, medium, high, xhigh, max +# Empty/unset = no default (model decides). Only applies when model supports effort. +# DEFAULT_EFFORT_LEVEL="" + +# =========================================== +# NATIVE THINKING (Kiro/Claude Adaptive Thinking) +# =========================================== + +# Experimental native adaptive thinking pass-through (default: off) +# Modes: +# - "off": Disabled (preserves existing fake reasoning behavior) +# - "auto": Enable only when the client explicitly requests a reasoning effort +# - "force": Enable for supported models even without client request +# KIRO_NATIVE_THINKING_MODE="off" + +# Native thinking display mode (default: summarized) +# - "summarized": Return summarized reasoning text +# - "omitted": Omit thinking text entirely (only signature) +# KIRO_NATIVE_THINKING_DISPLAY="summarized" + +# =========================================== +# API HOST OVERRIDE +# =========================================== + +# Override the default Kiro API host template (default: https://runtime.{region}.kiro.dev) +# Allows falling back to https://q.{region}.amazonaws.com when runtime.kiro.dev is unavailable +# KIRO_API_HOST_TEMPLATE="https://runtime.{region}.kiro.dev" + +# Same as above but for API host template (default: https://runtime.{region}.kiro.dev) +# KIRO_Q_HOST_TEMPLATE="https://runtime.{region}.kiro.dev" + +# Complete API host override (takes precedence over templates) +# When set, overrides both API and Q hosts entirely. +# Useful when the default Kiro domain is blocked by corporate proxies. +# KIRO_API_HOST="https://custom.endpoint.com" + +# =========================================== +# BILLING HEADER STRIPPING +# =========================================== + +# Strip Claude Code's per-request billing attribution headers from system prompts +# (default: true). These random per-request headers invalidate prompt caches. +# STRIP_BILLING_HEADER=true + # =========================================== # WEBSEARCH (MCP Tool Emulation) # =========================================== -# Enable web_search tool auto-injection (default: true) +# Enable web_search tool auto-injection (default: false) # When enabled, web_search is automatically added as a tool for MCP emulation (Path B) # Model decides whether to use it or not # diff --git a/.gitignore b/.gitignore index 679a6d8a..11b179e4 100644 --- a/.gitignore +++ b/.gitignore @@ -30,4 +30,4 @@ requests/ # Testing .pytest_cache/ .coverage -htmlcov/ \ No newline at end of file +errors.log diff --git a/extensions/__init__.py b/extensions/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/extensions/tool_name_alias.py b/extensions/tool_name_alias.py new file mode 100644 index 00000000..320634e6 --- /dev/null +++ b/extensions/tool_name_alias.py @@ -0,0 +1,213 @@ +""" +Runtime tool-name aliasing for Kiro's 64-character tool name limit. + +This module lives outside the upstream kiro/ tree. main.py installs it at +startup so kiro-gateway can keep upstream files sync-friendly while still +handling long MCP tool names (e.g. names emitted by aggregating MCP gateways +that exceed Kiro's 64-character limit). + +The install hook patches a small surface in kiro.{converters_core,parsers, +streaming_*} to: + + - Replace strict >64-char validation with silent registration. + - Rewrite tool names client→Kiro to a stable hashed alias on the way in. + - Rewrite tool names Kiro→client back to the original on the way out + (both stream events and bracket-formatted tool calls). + +All replacements are idempotent and reversible via uninstall_tool_name_aliasing, +which is intended for unit tests. +""" + +import hashlib +import re +from typing import Any, AsyncGenerator, Dict, Iterable, Optional + +from loguru import logger + + +MAX_KIRO_TOOL_NAME_LENGTH = 64 +_KIRO_TOOL_NAME_RE = re.compile(r"^[A-Za-z0-9_-]{1,64}$") +_SAFE_SUFFIX_RE = re.compile(r"[^A-Za-z0-9_-]+") + +_alias_to_original: Dict[str, str] = {} +_original_to_alias: Dict[str, str] = {} +_reserved_kiro_names = set() +_installed = False +_originals: Dict[str, Any] = {} + + +def needs_alias(name: str) -> bool: + return not bool(_KIRO_TOOL_NAME_RE.match(name or "")) + + +def _build_alias(name: str, digest_len: int) -> str: + digest = hashlib.sha256(name.encode("utf-8")).hexdigest()[:digest_len] + suffix_source = name.rsplit("__", 1)[-1] or name + suffix = _SAFE_SUFFIX_RE.sub("_", suffix_source).strip("_-") or "tool" + prefix = f"t_{digest}_" + max_suffix_len = MAX_KIRO_TOOL_NAME_LENGTH - len(prefix) + return prefix + suffix[-max_suffix_len:] + + +def alias_for_tool_name(name: str) -> str: + """Return the Kiro-facing name for a client-facing tool name.""" + if not needs_alias(name): + return name + + if name in _original_to_alias: + return _original_to_alias[name] + + alias = _build_alias(name, 12) + digest_len = 16 + while ( + alias in _reserved_kiro_names + or (alias in _alias_to_original and _alias_to_original[alias] != name) + ): + alias = _build_alias(name, digest_len) + digest_len += 4 + + _original_to_alias[name] = alias + _alias_to_original[alias] = name + logger.debug(f"Aliased long Kiro tool name: '{name}' -> '{alias}'") + return alias + + +def original_for_tool_name(name: str) -> str: + """Return the client-facing name for a Kiro-facing alias.""" + return _alias_to_original.get(name, name) + + +def _register_tool_names(tools: Optional[Iterable[Any]]) -> None: + if not tools: + return + for tool in tools: + name = getattr(tool, "name", None) + if isinstance(name, str) and not needs_alias(name): + _reserved_kiro_names.add(name) + for tool in tools: + name = getattr(tool, "name", None) + if isinstance(name, str): + alias_for_tool_name(name) + + +def _alias_tool_use(tool_use: Dict[str, Any]) -> Dict[str, Any]: + name = tool_use.get("name") + if isinstance(name, str): + tool_use["name"] = alias_for_tool_name(name) + return tool_use + + +def _restore_tool_call(tool_call: Dict[str, Any]) -> Dict[str, Any]: + func = tool_call.get("function") + if isinstance(func, dict): + name = func.get("name") + if isinstance(name, str): + func["name"] = original_for_tool_name(name) + name = tool_call.get("name") + if isinstance(name, str): + tool_call["name"] = original_for_tool_name(name) + return tool_call + + +def _restore_event_tool_name(event: Any) -> Any: + if getattr(event, "type", None) == "tool_use" and getattr(event, "tool_use", None): + _restore_tool_call(event.tool_use) + return event + + +def install_tool_name_aliasing() -> None: + """Install runtime patches without modifying upstream kiro/ modules.""" + global _installed + if _installed: + return + + import kiro.converters_core as converters_core + import kiro.parsers as parsers + import kiro.streaming_core as streaming_core + import kiro.streaming_openai as streaming_openai + import kiro.streaming_anthropic as streaming_anthropic + + _originals.update( + validate_tool_names=converters_core.validate_tool_names, + convert_tools_to_kiro_format=converters_core.convert_tools_to_kiro_format, + extract_tool_uses_from_message=converters_core.extract_tool_uses_from_message, + parse_bracket_tool_calls=parsers.parse_bracket_tool_calls, + parse_kiro_stream=streaming_core.parse_kiro_stream, + ) + + def validate_tool_names_with_aliases(tools): + _register_tool_names(tools) + + def convert_tools_to_kiro_format_with_aliases(tools): + _register_tool_names(tools) + if not tools: + return _originals["convert_tools_to_kiro_format"](tools) + + aliased_tools = [] + for tool in tools: + aliased_tools.append( + converters_core.UnifiedTool( + name=alias_for_tool_name(tool.name), + description=tool.description, + input_schema=tool.input_schema, + ) + ) + return _originals["convert_tools_to_kiro_format"](aliased_tools) + + def extract_tool_uses_from_message_with_aliases(content, tool_calls=None): + tool_uses = _originals["extract_tool_uses_from_message"](content, tool_calls) + return [_alias_tool_use(tool_use) for tool_use in tool_uses] + + def parse_bracket_tool_calls_with_restore(response_text): + tool_calls = _originals["parse_bracket_tool_calls"](response_text) + return [_restore_tool_call(tool_call) for tool_call in tool_calls] + + async def parse_kiro_stream_with_restore(*args, **kwargs) -> AsyncGenerator[Any, None]: + async for event in _originals["parse_kiro_stream"](*args, **kwargs): + yield _restore_event_tool_name(event) + + converters_core.validate_tool_names = validate_tool_names_with_aliases + converters_core.convert_tools_to_kiro_format = convert_tools_to_kiro_format_with_aliases + converters_core.extract_tool_uses_from_message = extract_tool_uses_from_message_with_aliases + + parsers.parse_bracket_tool_calls = parse_bracket_tool_calls_with_restore + streaming_core.parse_bracket_tool_calls = parse_bracket_tool_calls_with_restore + streaming_openai.parse_bracket_tool_calls = parse_bracket_tool_calls_with_restore + streaming_anthropic.parse_bracket_tool_calls = parse_bracket_tool_calls_with_restore + + streaming_core.parse_kiro_stream = parse_kiro_stream_with_restore + streaming_openai.parse_kiro_stream = parse_kiro_stream_with_restore + streaming_anthropic.parse_kiro_stream = parse_kiro_stream_with_restore + + _installed = True + logger.info("Installed Kiro tool-name aliasing extension") + + +def uninstall_tool_name_aliasing() -> None: + """Restore patched functions. Intended for unit tests.""" + global _installed + if _originals: + import kiro.converters_core as converters_core + import kiro.parsers as parsers + import kiro.streaming_core as streaming_core + import kiro.streaming_openai as streaming_openai + import kiro.streaming_anthropic as streaming_anthropic + + converters_core.validate_tool_names = _originals["validate_tool_names"] + converters_core.convert_tools_to_kiro_format = _originals["convert_tools_to_kiro_format"] + converters_core.extract_tool_uses_from_message = _originals["extract_tool_uses_from_message"] + + parsers.parse_bracket_tool_calls = _originals["parse_bracket_tool_calls"] + streaming_core.parse_bracket_tool_calls = _originals["parse_bracket_tool_calls"] + streaming_openai.parse_bracket_tool_calls = _originals["parse_bracket_tool_calls"] + streaming_anthropic.parse_bracket_tool_calls = _originals["parse_bracket_tool_calls"] + + streaming_core.parse_kiro_stream = _originals["parse_kiro_stream"] + streaming_openai.parse_kiro_stream = _originals["parse_kiro_stream"] + streaming_anthropic.parse_kiro_stream = _originals["parse_kiro_stream"] + + _originals.clear() + _alias_to_original.clear() + _original_to_alias.clear() + _reserved_kiro_names.clear() + _installed = False diff --git a/install.sh b/install.sh new file mode 100755 index 00000000..367ec349 --- /dev/null +++ b/install.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Kiro Gateway - Install/Reinstall Script +# Installs as a user systemd service and (re)starts it. +# Requires: uv (https://docs.astral.sh/uv/) + +REPO_DIR="$(cd "$(dirname "$0")" && pwd)" +SERVICE_NAME="kiro-gateway" +SERVICE_FILE="$HOME/.config/systemd/user/${SERVICE_NAME}.service" +VENV_DIR="${REPO_DIR}/.venv" + +echo "==> Kiro Gateway installer" +echo " Repo: ${REPO_DIR}" + +# Ensure systemd user directory exists +mkdir -p "$HOME/.config/systemd/user" + +# Create/update virtualenv and install dependencies +echo "==> Installing dependencies with uv..." +uv venv "$VENV_DIR" --quiet +uv pip install --quiet -r "${REPO_DIR}/requirements.txt" -p "${VENV_DIR}/bin/python" + +# Write systemd service unit +echo "==> Writing systemd service to ${SERVICE_FILE}" +cat > "$SERVICE_FILE" < Reloading systemd and restarting ${SERVICE_NAME}..." +systemctl --user daemon-reload +systemctl --user enable "${SERVICE_NAME}.service" +systemctl --user restart "${SERVICE_NAME}.service" + +# Show status +echo "==> Done. Service status:" +systemctl --user status "${SERVICE_NAME}.service" --no-pager diff --git a/kiro/account_manager.py b/kiro/account_manager.py index ec133d92..77f377a6 100644 --- a/kiro/account_manager.py +++ b/kiro/account_manager.py @@ -62,35 +62,102 @@ from kiro.utils import get_kiro_headers from kiro.account_errors import ErrorType from kiro.http_client import KiroHttpClient +from kiro import native_reasoning -def _is_runtime_endpoint(auth_manager: KiroAuthManager) -> bool: +KIRO_IDE_PROFILE_PATHS = ( + Path.home() / "Library/Application Support/Kiro/User/globalStorage/kiro.kiroagent/profile.json", + Path.home() / ".config/Kiro/User/globalStorage/kiro.kiroagent/profile.json", +) +MODEL_CATALOG_MAX_RESULTS = 50 +MODEL_CATALOG_MAX_PAGES = 10 + + +def _resolve_model_profile_arn(auth_manager: KiroAuthManager) -> Optional[str]: """ - Check if auth manager uses runtime endpoint that doesn't provide /ListAvailableModels. - - Runtime endpoint pattern: https://runtime.{region}.kiro.dev - Old endpoint pattern: https://q.{region}.amazonaws.com - - Runtime endpoint does not provide /ListAvailableModels API (AWS limitation). - + Resolve the profile ARN required by Kiro's model control plane. + + Refreshed Kiro IDE credential files may omit profileArn even though the IDE + persists the selected profile separately. This mirrors the IDE lookup while + still preferring an ARN explicitly supplied with gateway credentials. + Args: - auth_manager: KiroAuthManager instance - + auth_manager: Authenticated Kiro account. + Returns: - True if using runtime endpoint, False otherwise - - Examples: - >>> auth_manager.api_host = "https://runtime.us-east-1.kiro.dev" - >>> _is_runtime_endpoint(auth_manager) - True - >>> auth_manager.api_host = "https://runtime.eu-central-1.kiro.dev" - >>> _is_runtime_endpoint(auth_manager) - True - >>> auth_manager.api_host = "https://q.us-east-1.amazonaws.com" - >>> _is_runtime_endpoint(auth_manager) - False + A Kiro profile ARN, or None when no valid local profile is available. """ - return "://runtime." in auth_manager.api_host + if auth_manager.profile_arn: + return auth_manager.profile_arn + + for profile_path in KIRO_IDE_PROFILE_PATHS: + try: + with profile_path.open("r", encoding="utf-8") as profile_file: + profile_data = json.load(profile_file) + except FileNotFoundError: + continue + except (OSError, json.JSONDecodeError) as error: + logger.warning( + f"Unable to read Kiro IDE profile cache at {profile_path}: {error}" + ) + continue + + profile_arn = profile_data.get("arn") + if isinstance(profile_arn, str) and profile_arn.startswith("arn:"): + logger.debug(f"Using Kiro IDE profile cache at {profile_path}") + return profile_arn + + logger.warning(f"Kiro IDE profile cache at {profile_path} has no valid ARN") + + return None + + +# Control-plane operation that returns the model catalog, including each model's +# additionalModelRequestFieldsSchema (native reasoning effort options). It is an +# AWS JSON 1.0 RPC (x-amz-target header), not a REST path — MITM-confirmed against +# management.{region}.kiro.dev. See kiro-analysis/api-reference.md. +_LIST_MODELS_TARGET = "AmazonCodeWhispererService.ListAvailableModels" + + +async def _fetch_model_catalog(auth_manager: KiroAuthManager) -> List[Dict]: + """ + Fetch the model catalog from the Kiro control plane. + + Unlike the legacy runtime GET, this hits management.{region}.kiro.dev with the + ListAvailableModels JSON-1.0 RPC, which returns per-model + ``additionalModelRequestFieldsSchema`` metadata (used to drive native reasoning). + + Args: + auth_manager: Authenticated Kiro account. + + Returns: + List of model dicts (``models`` field of the response). + + Raises: + Exception: On non-200 after retries, propagated from the HTTP client. + """ + body: Dict[str, str] = {"origin": "AI_EDITOR"} + if auth_manager.profile_arn: + body["profileArn"] = auth_manager.profile_arn + + http_client = KiroHttpClient(auth_manager, shared_client=None) + try: + response = await http_client.request_with_retry( + method="POST", + url=f"{auth_manager.management_host}/", + json_data=body, + params=dict(body), + stream=False, + extra_headers={ + "x-amz-target": _LIST_MODELS_TARGET, + "Content-Type": "application/x-amz-json-1.0", + }, + ) + if response.status_code != 200: + raise Exception(f"HTTP {response.status_code}") + return response.json().get("models", []) + finally: + await http_client.close() def _format_duration(seconds: float) -> str: @@ -497,49 +564,21 @@ async def _initialize_account(self, account_id: str) -> bool: # Get token to verify credentials token = await auth_manager.get_access_token() - # Determine if we should fetch models or use static list - if _is_runtime_endpoint(auth_manager): - # New runtime endpoint does not provide /ListAvailableModels (AWS limitation) - # Use static list without attempting request - logger.debug(f"Account {account_id}: Using static model list for runtime.kiro.dev endpoint") + # Fetch the live catalog from the control plane (management.{region}.kiro.dev), + # which serves ListAvailableModels even for runtime-endpoint accounts and carries + # per-model additionalModelRequestFieldsSchema. Fall back to static models on failure. + try: + models_list = await _fetch_model_catalog(auth_manager) + if not models_list: + raise Exception("empty model catalog") + except Exception as e: + logger.error(f"Failed to fetch models for {account_id} after retries: {e}") + logger.warning("Using pre-configured fallback models. Models will be refreshed on next TTL cycle when network recovers.") models_list = FALLBACK_MODELS - else: - # Old endpoint - attempt to fetch dynamic model list - # Fetch models list with retry + fallback - params = {"origin": "AI_EDITOR"} - if auth_manager.auth_type == AuthType.KIRO_DESKTOP and auth_manager.profile_arn: - params["profileArn"] = auth_manager.profile_arn - - list_models_url = f"{auth_manager.q_host}/ListAvailableModels" - - # Use KiroHttpClient for retry logic (3 attempts with exponential backoff) - http_client = KiroHttpClient(auth_manager, shared_client=None) - - try: - response = await http_client.request_with_retry( - method="GET", - url=list_models_url, - json_data=None, - params=params, - stream=False - ) - - if response.status_code == 200: - data = response.json() - models_list = data.get("models", []) - else: - # Shouldn't happen (retry handles non-200), but keep for safety - raise Exception(f"HTTP {response.status_code}") - - except Exception as e: - # All retries exhausted - use fallback - logger.error(f"Failed to fetch models for {account_id} after retries: {e}") - logger.warning("Using pre-configured fallback models. Models will be refreshed on next TTL cycle when network recovers.") - models_list = FALLBACK_MODELS - - finally: - await http_client.close() - + + # Register native reasoning effort schemas advertised by the live catalog. + native_reasoning.register_from_catalog(models_list) + # Create model cache and update model_cache = ModelInfoCache() await model_cache.update(models_list) @@ -588,59 +627,52 @@ async def _refresh_account_models(self, account_id: str) -> None: account = self._accounts.get(account_id) if not account or not account.auth_manager: return - - # Check if using runtime endpoint (no dynamic model list available) - if _is_runtime_endpoint(account.auth_manager): - # Runtime endpoint does not provide /ListAvailableModels - # Use static list and update cache timestamp - logger.debug(f"Account {account_id}: Skipping model refresh for runtime.kiro.dev endpoint (using static list)") - await account.model_cache.update(FALLBACK_MODELS) + + # Re-fetch the live catalog from the control plane. On failure, keep the stale cache. + try: + models_list = await _fetch_model_catalog(account.auth_manager) + if not models_list: + raise Exception("empty model catalog") + + await account.model_cache.update(models_list) account.models_cached_at = time.time() + native_reasoning.register_from_catalog(models_list) + + # Update model_to_accounts mapping (new models may have appeared) + available_models = account.model_resolver.get_available_models() + for model in available_models: + if model not in self._model_to_accounts: + self._model_to_accounts[model] = ModelAccountList() + if account_id not in self._model_to_accounts[model].accounts: + self._model_to_accounts[model].accounts.append(account_id) + + logger.debug(f"Refreshed models for {account_id}") self._dirty = True - return - - # Old endpoint - attempt to fetch dynamic model list - # Use KiroHttpClient for retry logic - http_client = KiroHttpClient(account.auth_manager, shared_client=None) - - try: - params = {"origin": "AI_EDITOR"} - if account.auth_manager.auth_type == AuthType.KIRO_DESKTOP and account.auth_manager.profile_arn: - params["profileArn"] = account.auth_manager.profile_arn - - list_models_url = f"{account.auth_manager.q_host}/ListAvailableModels" - - response = await http_client.request_with_retry( - method="GET", - url=list_models_url, - json_data=None, - params=params, - stream=False - ) - - if response.status_code == 200: - data = response.json() - models_list = data.get("models", []) - await account.model_cache.update(models_list) - account.models_cached_at = time.time() - - # Update model_to_accounts mapping (new models may have appeared) - available_models = account.model_resolver.get_available_models() - for model in available_models: - if model not in self._model_to_accounts: - self._model_to_accounts[model] = ModelAccountList() - if account_id not in self._model_to_accounts[model].accounts: - self._model_to_accounts[model].accounts.append(account_id) - - logger.debug(f"Refreshed models for {account_id}") - self._dirty = True - + except Exception as e: # All retries exhausted - keep using stale cache logger.warning(f"Failed to refresh models for {account_id} after retries: {e}") - - finally: - await http_client.close() + + async def refresh_initialized_account_models(self) -> None: + """ + Refresh model lists for every initialized account. + + This is used by the public model-list endpoint so that additions and + removals made by Kiro are visible on the next client request. A failed + refresh preserves the last successfully fetched cache for that account. + + Returns: + None. + """ + async with self._lock: + initialized_account_ids = [ + account_id + for account_id, account in self._accounts.items() + if account.auth_manager and account.model_cache and account.model_resolver + ] + + for account_id in initialized_account_ids: + await self._refresh_account_models(account_id) async def get_next_account(self, model: str, exclude_accounts: Optional[set] = None) -> Optional[Account]: """ diff --git a/kiro/auth.py b/kiro/auth.py index 61ffe02b..acc7f4eb 100644 --- a/kiro/auth.py +++ b/kiro/auth.py @@ -46,6 +46,7 @@ get_kiro_refresh_url, get_kiro_api_host, get_kiro_q_host, + get_kiro_management_host, get_aws_sso_oidc_url, ) from kiro.utils import get_machine_fingerprint @@ -221,6 +222,7 @@ def __init__( self._refresh_url = get_kiro_refresh_url(sso_region_for_oidc) self._api_host = get_kiro_api_host(final_api_region) self._q_host = get_kiro_q_host(final_api_region) + self._management_host = get_kiro_management_host(final_api_region) # Log initialized endpoints for diagnostics (helps with DNS issues like #58, #132, #133) logger.info( @@ -267,6 +269,7 @@ def _load_credentials_from_sqlite(self, db_path: str) -> None: Args: db_path: Path to SQLite database file """ + conn = None try: path = Path(db_path).expanduser() if not path.exists(): @@ -371,7 +374,6 @@ def _load_credentials_from_sqlite(self, db_path: str) -> None: except Exception as e: logger.debug(f"Failed to auto-detect API region from profile ARN: {e}") - conn.close() logger.info(f"Credentials loaded from SQLite database: {db_path}") except sqlite3.Error as e: @@ -380,6 +382,9 @@ def _load_credentials_from_sqlite(self, db_path: str) -> None: logger.error(f"JSON decode error in SQLite data: {e}") except Exception as e: logger.error(f"Error loading credentials from SQLite: {e}") + finally: + if conn is not None: + conn.close() def _load_credentials_from_file(self, file_path: str) -> None: """ @@ -544,6 +549,7 @@ def _save_credentials_to_sqlite(self) -> None: logger.debug("SQLite write-back disabled (SQLITE_READONLY=true)") return + conn = None try: path = Path(self._sqlite_db).expanduser() if not path.exists(): @@ -558,7 +564,6 @@ def _save_credentials_to_sqlite(self) -> None: if self._sqlite_token_key: if self._try_save_to_key(cursor, self._sqlite_token_key): conn.commit() - conn.close() logger.debug(f"Credentials saved to SQLite key: {self._sqlite_token_key} (merged)") return else: @@ -568,18 +573,19 @@ def _save_credentials_to_sqlite(self) -> None: for key in SQLITE_TOKEN_KEYS: if self._try_save_to_key(cursor, key): conn.commit() - conn.close() logger.debug(f"Credentials saved to SQLite key: {key} (fallback, merged)") return # If we get here, no keys were updated - conn.close() logger.warning(f"Failed to save credentials to SQLite: no matching keys found") except sqlite3.Error as e: logger.error(f"SQLite error saving credentials: {e}") except Exception as e: logger.error(f"Error saving credentials to SQLite: {e}") + finally: + if conn is not None: + conn.close() def _try_save_to_key(self, cursor: sqlite3.Cursor, key: str) -> bool: """ @@ -965,6 +971,11 @@ def api_host(self) -> str: def q_host(self) -> str: """Q API host for the current region.""" return self._q_host + + @property + def management_host(self) -> str: + """Kiro control-plane host (ListAvailableModels/GetProfile) for the current region.""" + return self._management_host @property def fingerprint(self) -> str: diff --git a/kiro/cache.py b/kiro/cache.py index f0be72af..d99a09fa 100644 --- a/kiro/cache.py +++ b/kiro/cache.py @@ -25,12 +25,15 @@ """ import asyncio +import hashlib +import json import time -from typing import Any, Dict, List, Optional +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple 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, PROMPT_CACHE_TTL_SECONDS, PROMPT_CACHE_ACCOUNTING_ENABLED class ModelInfoCache: @@ -179,4 +182,348 @@ def size(self) -> int: @property def last_update_time(self) -> Optional[float]: """Last update time (timestamp) or None.""" - return self._last_update \ No newline at end of file + return self._last_update + + +@dataclass +class PromptCacheUsage: + """Anthropic-compatible prompt cache usage fields.""" + + cache_read_input_tokens: int = 0 + cache_creation_input_tokens: int = 0 + + def to_anthropic_usage_fields(self) -> Dict[str, int]: + """Return non-zero Anthropic usage fields.""" + fields: Dict[str, int] = {} + if self.cache_read_input_tokens > 0: + fields["cache_read_input_tokens"] = self.cache_read_input_tokens + if self.cache_creation_input_tokens > 0: + fields["cache_creation_input_tokens"] = self.cache_creation_input_tokens + return fields + + +@dataclass +class _PromptCacheEntry: + tokens: int + expires_at: float + + +@dataclass +class _PromptCacheBlock: + prefix_fingerprint: str + cumulative_tokens: int + breakpoint: bool = False + implicit_breakpoint: bool = False + + +class PromptCacheTracker: + """ + In-memory prompt cache accounting using cacheable prompt prefixes. + + This only reports Anthropic cache usage fields. It does not cache model + responses and does not alter upstream requests. When Anthropic + ``cache_control`` markers are present, cache reads are matched on the + stable prompt prefix up to the most recent cacheable breakpoint, so + appended turns can reuse earlier cached context. Entry TTL is fixed when + the entry is created and is not refreshed by cache hits. + """ + + def __init__( + self, + cache_ttl: int = PROMPT_CACHE_TTL_SECONDS, + enabled: bool = PROMPT_CACHE_ACCOUNTING_ENABLED, + ): + self._cache_ttl = cache_ttl + self._enabled = enabled + self._entries: Dict[str, _PromptCacheEntry] = {} + self._lock = asyncio.Lock() + + @property + def enabled(self) -> bool: + return self._enabled + + @property + def size(self) -> int: + self._prune_expired(time.time()) + return len(self._entries) + + def clear(self) -> None: + self._entries.clear() + + async def record( + self, + *, + model: str, + messages: Optional[List[Dict[str, Any]]], + tools: Optional[List[Dict[str, Any]]], + system: Optional[Any], + input_tokens: int, + ) -> PromptCacheUsage: + """ + Record a successful request and return cache accounting usage. + + Args: + model: Requested model name + messages: Anthropic request messages as dictionaries + tools: Anthropic request tools as dictionaries + system: Anthropic system prompt + input_tokens: Estimated input tokens for this request + """ + if not self._enabled or input_tokens <= 0: + return PromptCacheUsage() + + blocks = self._build_prefix_blocks( + model=model, + messages=messages or [], + tools=tools, + system=system, + input_tokens=input_tokens, + ) + breakpoints = [block for block in blocks if block.breakpoint] + if not breakpoints: + return PromptCacheUsage() + + last_breakpoint_tokens = min(breakpoints[-1].cumulative_tokens, input_tokens) + now = time.time() + + async with self._lock: + self._prune_expired(now) + + matched_tokens = 0 + for block in reversed(breakpoints[-10:]): + entry = self._entries.get(block.prefix_fingerprint) + if entry and entry.expires_at > now: + matched_tokens = min(entry.tokens, block.cumulative_tokens, input_tokens) + break + + for block in breakpoints: + self._entries.setdefault( + block.prefix_fingerprint, + _PromptCacheEntry( + tokens=min(block.cumulative_tokens, input_tokens), + expires_at=now + self._cache_ttl, + ), + ) + + return PromptCacheUsage( + cache_read_input_tokens=matched_tokens, + cache_creation_input_tokens=max(last_breakpoint_tokens - matched_tokens, 0), + ) + + def _prune_expired(self, now: float) -> None: + expired_keys = [ + key for key, entry in self._entries.items() + if entry.expires_at <= now + ] + for key in expired_keys: + self._entries.pop(key, None) + + def _build_prefix_blocks( + self, + *, + model: str, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], + system: Optional[Any], + input_tokens: int, + ) -> List[_PromptCacheBlock]: + flattened = self._flatten_cache_blocks( + messages=messages, + tools=tools or [], + system=system, + ) + if not flattened: + return [] + + prelude = { + "model": model, + } + prefix_hash = self._hash_payload(prelude) + blocks: List[_PromptCacheBlock] = [] + cumulative_tokens = 0 + has_explicit_breakpoint = False + active_cache = False + + for value, tokens, has_cache_control, is_message_end in flattened: + cumulative_tokens += max(tokens, 0) + prefix_hash = self._hash_payload({ + "previous": prefix_hash, + "block": value, + }) + + if has_cache_control: + has_explicit_breakpoint = True + active_cache = True + blocks.append(_PromptCacheBlock( + prefix_fingerprint=prefix_hash, + cumulative_tokens=cumulative_tokens, + breakpoint=True, + )) + elif active_cache and is_message_end: + blocks.append(_PromptCacheBlock( + prefix_fingerprint=prefix_hash, + cumulative_tokens=cumulative_tokens, + breakpoint=True, + )) + else: + blocks.append(_PromptCacheBlock( + prefix_fingerprint=prefix_hash, + cumulative_tokens=cumulative_tokens, + breakpoint=False, + implicit_breakpoint=is_message_end, + )) + + if not has_explicit_breakpoint: + min_cacheable_tokens = self._minimum_cacheable_tokens(model) + implicit_blocks = [ + _PromptCacheBlock( + prefix_fingerprint=block.prefix_fingerprint, + cumulative_tokens=block.cumulative_tokens, + breakpoint=block.cumulative_tokens >= min_cacheable_tokens, + implicit_breakpoint=block.implicit_breakpoint, + ) + for block in blocks + ] + if any(block.breakpoint for block in implicit_blocks): + blocks = implicit_blocks + else: + # Preserve the old behavior for short callers without + # Anthropic cache_control markers: exact repeated requests + # still report cache reads, but appended short conversations do + # not accidentally match. + blocks[-1] = _PromptCacheBlock( + prefix_fingerprint=prefix_hash, + cumulative_tokens=input_tokens, + breakpoint=True, + ) + + return blocks + + def _flatten_cache_blocks( + self, + *, + messages: List[Dict[str, Any]], + tools: List[Dict[str, Any]], + system: Optional[Any], + ) -> List[Tuple[Any, int, bool, bool]]: + blocks: List[Tuple[Any, int, bool, bool]] = [] + + for index, tool in enumerate(tools): + normalized, has_cache_control = self._strip_cache_control(tool) + blocks.append(( + {"kind": "tool", "tool_index": index, "tool": normalized}, + self._estimate_tokens(normalized), + has_cache_control, + False, + )) + + system_blocks = system if isinstance(system, list) else ([system] if system else []) + for index, block in enumerate(system_blocks): + normalized, has_cache_control = self._strip_cache_control(block) + blocks.append(( + {"kind": "system", "system_index": index, "block": normalized}, + self._estimate_tokens(normalized), + has_cache_control, + False, + )) + + for message_index, message in enumerate(messages): + role = message.get("role", "") + content = message.get("content") + if isinstance(content, list): + last_block_index = len(content) - 1 + for block_index, content_block in enumerate(content): + normalized, has_cache_control = self._strip_cache_control(content_block) + blocks.append(( + { + "kind": "message", + "message_index": message_index, + "role": role, + "block_index": block_index, + "block": normalized, + }, + self._estimate_tokens(normalized), + has_cache_control, + block_index == last_block_index, + )) + else: + normalized, has_cache_control = self._strip_cache_control(content) + blocks.append(( + { + "kind": "message", + "message_index": message_index, + "role": role, + "block_index": 0, + "block": normalized, + }, + self._estimate_tokens(normalized), + has_cache_control, + True, + )) + + return blocks + + def _strip_cache_control(self, value: Any) -> Tuple[Any, bool]: + normalized = self._normalize(value) + found = False + + def strip(item: Any) -> Any: + nonlocal found + if isinstance(item, dict): + output = {} + for key, subvalue in item.items(): + if key == "cache_control": + if isinstance(subvalue, dict) and subvalue.get("type") == "ephemeral": + found = True + continue + output[key] = strip(subvalue) + return output + if isinstance(item, list): + return [strip(subvalue) for subvalue in item] + return item + + return strip(normalized), found + + def _estimate_tokens(self, value: Any) -> int: + if isinstance(value, str): + text = value + else: + text = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + try: + from kiro.tokenizer import count_tokens + return count_tokens(text) + except Exception: + return len(text) // 4 + 1 + + def _hash_payload(self, value: Any) -> str: + serialized = json.dumps( + self._normalize(value), + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ) + return hashlib.sha256(serialized.encode("utf-8")).hexdigest() + + def _minimum_cacheable_tokens(self, model: str) -> int: + model_lower = model.lower() + if "opus" in model_lower: + return 4096 + if "haiku-3" in model_lower or "haiku_3" in model_lower: + return 2048 + return 1024 + + def _normalize(self, value: Any) -> Any: + if hasattr(value, "model_dump"): + return self._normalize(value.model_dump(exclude_none=True)) + if isinstance(value, dict): + return { + str(key): self._normalize(item) + for key, item in sorted(value.items(), key=lambda kv: str(kv[0])) + if item is not None + } + if isinstance(value, list): + return [self._normalize(item) for item in value] + return value + + +prompt_cache_tracker = PromptCacheTracker() \ No newline at end of file diff --git a/kiro/config.py b/kiro/config.py index e0f3a527..6ebc28c5 100644 --- a/kiro/config.py +++ b/kiro/config.py @@ -120,6 +120,11 @@ def _get_raw_env_value(var_name: str, env_file: str = ".env") -> Optional[str]: # VPN_PROXY_URL=192.168.1.100:8080 (defaults to http://) VPN_PROXY_URL: str = os.getenv("VPN_PROXY_URL", "") +# Disable proxy API key authentication (default: false) +# When true, the gateway accepts requests without valid Authorization / x-api-key header. +# Only use in trusted local environments - do NOT enable on public or shared servers. +PROXY_AUTH_DISABLED: bool = os.getenv("PROXY_AUTH_DISABLED", "false").lower() in ("true", "1", "yes") + # ================================================================================================== # Kiro API Credentials # ================================================================================================== @@ -127,6 +132,20 @@ def _get_raw_env_value(var_name: str, env_file: str = ".env") -> Optional[str]: # Refresh token for updating access token REFRESH_TOKEN: str = os.getenv("REFRESH_TOKEN", "") +# Kiro IDE profile file used as a fallback source for profileArn. +DEFAULT_KIRO_PROFILE_FILE: Path = ( + Path.home() + / "Library" + / "Application Support" + / "Kiro" + / "User" + / "globalStorage" + / "kiro.kiroagent" + / "profile.json" +) +_raw_profile_file = _get_raw_env_value("KIRO_PROFILE_FILE") or os.getenv("KIRO_PROFILE_FILE", "") +KIRO_PROFILE_FILE: str = str(Path(_raw_profile_file)) if _raw_profile_file else str(DEFAULT_KIRO_PROFILE_FILE) + # Profile ARN for AWS CodeWhisperer PROFILE_ARN: str = os.getenv("PROFILE_ARN", "") @@ -179,10 +198,22 @@ def _get_raw_env_value(var_name: str, env_file: str = ".env") -> Optional[str]: # Universal endpoint for all regions (us-east-1, eu-central-1, etc.) # See: https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/security-data-perimeter.html # Fixed in issue #58 - codewhisperer.{region}.amazonaws.com doesn't exist for non-us-east-1 regions -KIRO_API_HOST_TEMPLATE: str = "https://runtime.{region}.kiro.dev" +KIRO_API_HOST_TEMPLATE: str = os.getenv("KIRO_API_HOST_TEMPLATE", "https://runtime.{region}.kiro.dev") # Host for Q API (ListAvailableModels) -KIRO_Q_HOST_TEMPLATE: str = "https://runtime.{region}.kiro.dev" +KIRO_Q_HOST_TEMPLATE: str = os.getenv("KIRO_Q_HOST_TEMPLATE", "https://runtime.{region}.kiro.dev") + +# Host for the Kiro control plane (ListAvailableModels, GetProfile). +# The runtime plane (runtime.{region}.kiro.dev) does not serve ListAvailableModels; +# the model catalog — including each model's additionalModelRequestFieldsSchema — is +# served here. See kiro-analysis: MITM-confirmed on management.us-east-1.kiro.dev. +KIRO_MANAGEMENT_HOST_TEMPLATE: str = os.getenv( + "KIRO_MANAGEMENT_HOST_TEMPLATE", "https://management.{region}.kiro.dev" +) + +# Override API host directly (bypasses template-based region resolution). +# When set, this takes precedence over KIRO_API_HOST_TEMPLATE. +KIRO_API_HOST_OVERRIDE: str = os.getenv("KIRO_API_HOST", "") # ================================================================================================== # Token Settings @@ -282,8 +313,10 @@ 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"}, + {"modelId": "claude-opus-4.8"}, {"modelId": "deepseek-3.2"}, {"modelId": "glm-5"}, + {"modelId": "gpt-5.6-sol"}, {"modelId": "minimax-m2.1"}, {"modelId": "minimax-m2.5"}, {"modelId": "qwen3-coder-next"}, @@ -296,6 +329,15 @@ 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 +# Prompt cache accounting TTL in seconds (5 minutes) +PROMPT_CACHE_TTL_SECONDS: int = int(os.getenv("PROMPT_CACHE_TTL_SECONDS", "300")) + +# Enable Anthropic prompt cache usage accounting. +# This only adds usage fields; responses are never cached. +PROMPT_CACHE_ACCOUNTING_ENABLED: bool = os.getenv( + "PROMPT_CACHE_ACCOUNTING_ENABLED", "true" +).lower() in ("true", "1", "yes") + # Default maximum number of input tokens DEFAULT_MAX_INPUT_TOKENS: int = 200000 @@ -351,7 +393,7 @@ def _get_raw_env_value(var_name: str, env_file: str = ".env") -> Optional[str]: # This helps handle "stuck" requests when the model takes too long to think. # Default: 30 seconds (recommended for production) # Set a lower value (e.g., 10-15) for more aggressive retry. -FIRST_TOKEN_TIMEOUT: float = float(os.getenv("FIRST_TOKEN_TIMEOUT", "15")) +FIRST_TOKEN_TIMEOUT: float = float(os.getenv("FIRST_TOKEN_TIMEOUT", "300")) # Read timeout for streaming responses (in seconds). # This is the maximum time to wait for data between chunks during streaming. @@ -360,6 +402,10 @@ 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")) +# Strip billing header from Anthropic API responses (default: true) +# When enabled, removes billing/usage metadata that may leak internal information. +STRIP_BILLING_HEADER: bool = os.getenv("STRIP_BILLING_HEADER", "true").lower() in ("true", "1", "yes") + # Maximum number of attempts on first token timeout. # After exhausting all attempts, an error will be returned. # Default: 3 attempts @@ -478,6 +524,69 @@ def _warn_timeout_configuration(): FAKE_REASONING_INITIAL_BUFFER_SIZE: int = int(os.getenv("FAKE_REASONING_INITIAL_BUFFER_SIZE", "20")) +# ================================================================================================== +# Native Reasoning Settings (Extended Thinking via Kiro API) +# ================================================================================================== + +# Enable native reasoning support. +# When enabled, requests can use the Kiro API's native extended thinking via output_config +# or reasoning schemas in the payload. +NATIVE_REASONING_ENABLED: bool = os.getenv("NATIVE_REASONING", "").lower() not in ("false", "0", "no", "disabled", "off") + +# Maps model IDs to their supported reasoning schema type. +# "output_config" models use outputConfig.effortLevel in the payload. +# "reasoning" models use a different reasoning block format. +NATIVE_EFFORT_SCHEMA_BY_MODEL: dict = { + "auto": "output_config", + "claude-opus-4.6": "output_config", + "claude-opus-4.7": "output_config", + "claude-opus-4.8": "output_config", + "claude-sonnet-4.6": "output_config", + "gpt-5.6-sol": "reasoning", +} + +# Valid effort levels accepted by the Kiro API. +VALID_EFFORT_LEVELS: list = ["low", "medium", "high", "xhigh", "max"] + +# Common aliases that map to canonical effort level strings. +EFFORT_LEVEL_ALIASES: dict = { + "minimal": "low", + "low": "low", + "medium": "medium", + "high": "high", + "xhigh": "xhigh", + "max": "max", +} + +# Default effort level when the client does not specify one. +_DEFAULT_EFFORT_RAW: str = os.getenv("DEFAULT_EFFORT_LEVEL", "").lower().strip() +DEFAULT_EFFORT_LEVEL: Optional[str] = ( + EFFORT_LEVEL_ALIASES.get(_DEFAULT_EFFORT_RAW) + if _DEFAULT_EFFORT_RAW in EFFORT_LEVEL_ALIASES + else None +) + + +# ================================================================================================== +# Native Thinking Settings +# ================================================================================================== + +# Native thinking mode - controls how the gateway handles thinking/reasoning. +# - "off": Native thinking is disabled (default) +# - "auto": Enable if client requests thinking via thinking budget +# - "force": Always inject thinking mode regardless of client request +KIRO_NATIVE_THINKING_MODE: str = os.getenv("KIRO_NATIVE_THINKING_MODE", "off").lower() +if KIRO_NATIVE_THINKING_MODE not in ("off", "auto", "force"): + KIRO_NATIVE_THINKING_MODE = "off" + +# How to display thinking content in responses: +# - "summarized": Show a short summary of the thinking process +# - "omitted": Do not include thinking content in the response +KIRO_NATIVE_THINKING_DISPLAY: str = os.getenv("KIRO_NATIVE_THINKING_DISPLAY", "summarized").lower() +if KIRO_NATIVE_THINKING_DISPLAY not in ("summarized", "omitted"): + KIRO_NATIVE_THINKING_DISPLAY = "summarized" + + # ================================================================================================== # Payload Size Guard Settings # ================================================================================================== @@ -555,7 +664,7 @@ def _warn_timeout_configuration(): # Application Version # ================================================================================================== -APP_VERSION: str = "2.4.dev.13" +APP_VERSION: str = "2.5" APP_TITLE: str = "Kiro Gateway" APP_DESCRIPTION: str = "Proxy gateway for Kiro API (Amazon Q Developer / AWS CodeWhisperer). OpenAI and Anthropic compatible. Made by @jwadow" @@ -579,3 +688,8 @@ def get_kiro_q_host(region: str) -> str: """Return Q API host for the specified region.""" return KIRO_Q_HOST_TEMPLATE.format(region=region) + +def get_kiro_management_host(region: str) -> str: + """Return management host for the specified region.""" + return KIRO_MANAGEMENT_HOST_TEMPLATE.format(region=region) + diff --git a/kiro/converters_anthropic.py b/kiro/converters_anthropic.py index ea297dc1..a0134887 100644 --- a/kiro/converters_anthropic.py +++ b/kiro/converters_anthropic.py @@ -24,11 +24,18 @@ to the unified format used by converters_core.py. """ -from typing import Any, Dict, List, Optional +import re +from typing import Any, Dict, List, Optional, Tuple from loguru import logger -from kiro.config import HIDDEN_MODELS +from kiro.config import ( + HIDDEN_MODELS, + MODEL_ALIASES, + STRIP_BILLING_HEADER, + KIRO_NATIVE_THINKING_MODE, + KIRO_NATIVE_THINKING_DISPLAY, +) from kiro.model_resolver import get_model_id_for_kiro from kiro.models_anthropic import ( AnthropicMessagesRequest, @@ -39,10 +46,283 @@ UnifiedMessage, UnifiedTool, ThinkingConfig, + NativeThinkingConfig, + build_native_thinking_config, build_kiro_payload, extract_text_content, extract_images_from_content, + extract_document_text_from_content_block, + KIRO_CACHE_POINT, +) + + +# ================================================================================================== +# Billing Attribution Stripping +# ================================================================================================== + +_BILLING_HEADER_LINE_PATTERN = re.compile( + r"^x-anthropic-billing-header:[^\n]*\n?", re.IGNORECASE ) +_BILLING_FOOTER_RE = re.compile(r"Response from [^.]+\.model\..+?\.(.+?) via", re.DOTALL) + + +def _strip_billing_attribution(text: str) -> str: + """ + Strip AWS Q Developer / Kiro billing attribution from response text. + + Handles two patterns: + 1. Claude Code's per-request billing header line: + ``x-anthropic-billing-header: cc_version=...; cc_entrypoint=...`` + 2. Kiro API's billing footer: + ``Response from prod.us-east-1.deflector.9x7g0y8h model.claude-3-5-sonnet-20241022-v2:0 via ...`` + + Both are billing/metering metadata, not assistant content. + + Args: + text: Response text that may contain billing attribution. + + Returns: + Text with billing attribution stripped. + """ + if not text or not STRIP_BILLING_HEADER: + return text + + # Strip Claude Code billing header line + stripped = _BILLING_HEADER_LINE_PATTERN.sub("", text, count=1) + if stripped != text: + text = stripped.lstrip("\n") + + # Strip Kiro billing footer + match = _BILLING_FOOTER_RE.search(text) + if match: + start = match.start() + text = text[:start].rstrip() + + return text + + +# ================================================================================================== +# Inline System Message Extraction +# ================================================================================================== + +def separate_inline_system_messages( + messages: List[UnifiedMessage], +) -> Tuple[Optional[str], List[UnifiedMessage]]: + """ + Separate inline system messages from conversation messages. + + Some clients send system instructions as user messages with specific prefixes. + This function extracts them and returns them as a system prompt string, + along with the remaining conversation messages. + + Args: + messages: List of unified messages. + + Returns: + Tuple of (system_prompt_or_None, remaining_messages). + """ + system_parts: List[str] = [] + remaining: List[UnifiedMessage] = [] + + for msg in messages: + if msg.role == "user": + text = extract_text_content(msg.content) if msg.content else "" + if text.startswith("[SYSTEM INSTRUCTION]"): + system_parts.append(text.replace("[SYSTEM INSTRUCTION]", "").strip()) + continue + remaining.append(msg) + + system_prompt = "\n\n".join(system_parts) if system_parts else None + return system_prompt, remaining + + +# ================================================================================================== +# Cache Control Detection +# ================================================================================================== + +def get_cache_control(value: Any) -> Optional[Dict[str, str]]: + """ + Extract cache_control from a value if present. + + Args: + value: Dict or object that may have a cache_control field. + + Returns: + Cache control dict if present, None otherwise. + """ + if isinstance(value, dict): + cc = value.get("cache_control") + if isinstance(cc, dict): + return cc + elif hasattr(value, "cache_control"): + cc = getattr(value, "cache_control", None) + if isinstance(cc, dict): + return cc + return None + + +def has_supported_cache_control(value: Any) -> bool: + """ + Check if a value has a supported cache_control field (type: "ephemeral"). + + Args: + value: Dict or object that may have a cache_control field. + + Returns: + True if cache_control is present and supported. + """ + cc = get_cache_control(value) + return cc is not None and cc.get("type") == "ephemeral" + + +# ================================================================================================== +# Budget to Effort Mapping +# =================================================================================================> + +def _budget_to_effort_level(budget_tokens: Optional[int], max_tokens: int) -> str: + """ + Map a fake thinking budget token count back to an effort level string + for native thinking support. + + Args: + budget_tokens: The fake thinking budget in tokens. + max_tokens: Maximum output tokens for the request. + + Returns: + Effort level string (low/medium/high/xhigh/max). + """ + if not budget_tokens or not max_tokens: + return "medium" + + ratio = budget_tokens / max_tokens + + if ratio >= 0.95: + return "max" + if ratio >= 0.80: + return "xhigh" + if ratio >= 0.50: + return "high" + if ratio >= 0.20: + return "medium" + return "low" + + +# ================================================================================================== +# Tool Choice Prompt Addition +# ================================================================================================== + +def build_tool_choice_prompt_addition(tool_choice: Optional[Dict[str, Any]]) -> str: + """ + Build system prompt addition for tool_choice instruction. + + When the client specifies a tool_choice, we need to add JSON instructions + to the system prompt so the model knows to output a tool call in the + specified format. + + Args: + tool_choice: Tool choice configuration from Anthropic request. + + Returns: + System prompt addition text (empty string if not applicable). + """ + if not tool_choice: + return "" + + choice_type = tool_choice.get("type", "") + + if choice_type == "auto": + return "" + + if choice_type == "any": + return ( + "\n\n---\n" + "# Tool Use Instruction\n\n" + "You MUST use a tool in your response. Select the most appropriate tool " + "for the task and provide the necessary input parameters." + ) + + if choice_type == "tool": + tool_name = tool_choice.get("name", "") + if tool_name: + return ( + f"\n\n---\n" + f"# Tool Use Instruction\n\n" + f"You MUST use the `{tool_name}` tool in your response. " + f"Provide the necessary input parameters for this tool." + ) + + return "" + + +# ================================================================================================== +# JSON Schema Output Support +# ================================================================================================== + +def extract_json_schema_output_config( + response_format: Optional[Dict[str, Any]], +) -> Optional[Dict[str, Any]]: + """ + Extract JSON schema configuration from response_format. + + Args: + response_format: Response format configuration from Anthropic request. + + Returns: + JSON schema dict if present and valid, None otherwise. + """ + if not response_format: + return None + + fmt_type = response_format.get("type", "") + if fmt_type != "json_schema": + return None + + json_schema = response_format.get("json_schema") + if not json_schema: + return None + + schema = json_schema.get("schema") + if not schema: + return None + + name = json_schema.get("name", "response") + return {"name": name, "schema": schema} + + +def build_json_schema_prompt_addition( + response_format: Optional[Dict[str, Any]], +) -> str: + """ + Build system prompt addition for JSON schema output instruction. + + Args: + response_format: Response format configuration from Anthropic request. + + Returns: + System prompt addition text (empty string if not applicable). + """ + json_schema_config = extract_json_schema_output_config(response_format) + if not json_schema_config: + return "" + + import json as _json + schema_str = _json.dumps(json_schema_config["schema"], indent=2) + name = json_schema_config["name"] + + return ( + f"\n\n---\n" + f"# Output Format Instruction\n\n" + f"You MUST respond with a valid JSON object that conforms to the following schema " + f"(named `{name}`):\n\n" + f"```json\n{schema_str}\n```\n\n" + f"Your response should be ONLY the JSON object, with no additional text, " + f"markdown formatting, or explanation outside the JSON." + ) + + +# ================================================================================================== +# Content Processing Helpers +# ================================================================================================== def convert_anthropic_content_to_text(content: Any) -> str: @@ -68,6 +348,10 @@ def convert_anthropic_content_to_text(content: Any) -> str: if isinstance(block, dict): if block.get("type") == "text": text_parts.append(block.get("text", "")) + elif block.get("type") == "document": + doc_text = extract_document_text_from_content_block(block) + if doc_text: + text_parts.append(doc_text) elif hasattr(block, "type") and block.type == "text": text_parts.append(block.text) return "".join(text_parts) @@ -96,7 +380,7 @@ def extract_system_prompt(system: Any) -> str: return "" if isinstance(system, str): - return system + return _strip_billing_attribution(system) if isinstance(system, list): text_parts = [] @@ -104,10 +388,12 @@ def extract_system_prompt(system: Any) -> str: if isinstance(block, dict): # Handle {"type": "text", "text": "...", "cache_control": {...}} if block.get("type") == "text": - text_parts.append(block.get("text", "")) + text = block.get("text", "") + text_parts.append(_strip_billing_attribution(text)) elif hasattr(block, "type") and block.type == "text": # Handle Pydantic model - text_parts.append(getattr(block, "text", "")) + text = getattr(block, "text", "") + text_parts.append(_strip_billing_attribution(text)) return "\n".join(text_parts) return str(system) @@ -201,8 +487,6 @@ def extract_images_from_tool_results(content: Any) -> List[Dict[str, Any]]: return images - return tool_results - def extract_tool_uses_from_anthropic_content(content: Any) -> List[Dict[str, Any]]: """ @@ -357,11 +641,17 @@ def convert_anthropic_tools( if isinstance(tool, dict): name = tool.get("name", "") description = tool.get("description") - input_schema = tool.get("input_schema", {}) + input_schema = tool.get("input_schema") else: name = tool.name description = tool.description - input_schema = tool.input_schema + input_schema = getattr(tool, "input_schema", None) + + # Skip server-managed tools without input_schema (e.g., web_search) + # Kiro API requires tools to have a valid input_schema + if input_schema is None: + logger.debug(f"Skipping tool '{name}' without input_schema (server-managed)") + continue unified_tools.append( UnifiedTool(name=name, description=description, input_schema=input_schema) @@ -376,6 +666,7 @@ def extract_thinking_config_from_anthropic(request: AnthropicMessagesRequest) -> Handles thinking parameter: - {"type": "enabled", "budget_tokens": N} → enabled with budget + - {"type": "adaptive", "effort": "max"} → enabled with effort-based budget - {"type": "disabled"} → disabled - None → enabled with default budget @@ -400,6 +691,11 @@ def extract_thinking_config_from_anthropic(request: AnthropicMessagesRequest) -> >>> request.thinking = {"type": "enabled", "budget_tokens": 8000} >>> extract_thinking_config_from_anthropic(request) ThinkingConfig(enabled=True, budget_tokens=8000) + + >>> # Adaptive effort translated to gateway fake thinking budget + >>> request.thinking = {"type": "adaptive", "effort": "max"} + >>> extract_thinking_config_from_anthropic(request) + ThinkingConfig(enabled=True, budget_tokens=4096) """ if not request.thinking: # No thinking specified → use defaults @@ -422,6 +718,31 @@ def extract_thinking_config_from_anthropic(request: AnthropicMessagesRequest) -> logger.debug(f"Extracted thinking config from Anthropic: type='enabled', budget={budget}") return ThinkingConfig(enabled=True, budget_tokens=budget) + if thinking_type == "adaptive": + effort = request.thinking.get("effort") + if not effort: + logger.debug("Extracted adaptive thinking config from Anthropic without effort") + return ThinkingConfig(enabled=True, budget_tokens=None) + + if effort == "none": + logger.debug("Extracted adaptive thinking config from Anthropic: effort='none'") + return ThinkingConfig(enabled=False, budget_tokens=None) + + try: + budget = reasoning_effort_to_budget(request.max_tokens, effort) + except ValueError: + logger.warning( + f"Unsupported Anthropic adaptive thinking effort '{effort}'. " + "Using default fake thinking budget." + ) + return ThinkingConfig(enabled=True, budget_tokens=None) + + logger.debug( + f"Extracted adaptive thinking config from Anthropic: effort='{effort}', " + f"max_tokens={request.max_tokens}, budget={budget}" + ) + return ThinkingConfig(enabled=True, budget_tokens=budget) + # Unknown type → use defaults return ThinkingConfig(enabled=True, budget_tokens=None) @@ -453,6 +774,9 @@ def anthropic_to_kiro( # Convert messages to unified format unified_messages = convert_anthropic_messages(request.messages) + # Separate inline system messages from conversation + inline_system_prompt, unified_messages = separate_inline_system_messages(unified_messages) + # Convert tools to unified format unified_tools = convert_anthropic_tools(request.tools) @@ -460,6 +784,27 @@ def anthropic_to_kiro( # It can be a string or list of content blocks (for prompt caching) system_prompt = extract_system_prompt(request.system) + # Merge inline system prompt with main system prompt + if inline_system_prompt: + if system_prompt: + system_prompt = f"{system_prompt}\n\n{inline_system_prompt}" + else: + system_prompt = inline_system_prompt + + # Add tool choice prompt addition if specified + tool_choice_addition = build_tool_choice_prompt_addition( + getattr(request, "tool_choice", None) + ) + if tool_choice_addition: + system_prompt = system_prompt + tool_choice_addition if system_prompt else tool_choice_addition + + # Add JSON schema output prompt addition if specified + json_schema_addition = build_json_schema_prompt_addition( + getattr(request, "response_format", None) + ) + if json_schema_addition: + system_prompt = system_prompt + json_schema_addition if system_prompt else json_schema_addition + # Get model ID for Kiro API (normalizes + resolves hidden models) # Pass-through principle: we normalize and send to Kiro, Kiro decides if valid model_id = get_model_id_for_kiro(request.model, HIDDEN_MODELS) @@ -467,11 +812,25 @@ def anthropic_to_kiro( # Extract thinking configuration from thinking parameter thinking_config = extract_thinking_config_from_anthropic(request) + # Build native thinking config from model and thinking effort + native_effort: Optional[str] = None + native_display: Optional[str] = None + if isinstance(request.thinking, dict) and request.thinking.get("type") == "adaptive": + native_effort = request.thinking.get("effort") or "high" + native_display = request.thinking.get("display") + native_thinking_config = build_native_thinking_config(model_id, native_effort) + if native_display in ("summarized", "omitted"): + native_thinking_config.display = native_display + if native_thinking_config.enabled: + # Native adaptive thinking supersedes fake tag injection for this request. + thinking_config = ThinkingConfig(enabled=False, budget_tokens=None) + logger.debug( f"Converting Anthropic request: model={request.model} -> {model_id}, " f"messages={len(unified_messages)}, tools={len(unified_tools) if unified_tools else 0}, " f"system_prompt_length={len(system_prompt)}, " - f"thinking_enabled={thinking_config.enabled}, thinking_budget={thinking_config.budget_tokens}" + f"thinking_enabled={thinking_config.enabled}, thinking_budget={thinking_config.budget_tokens}, " + f"native_thinking_enabled={native_thinking_config.enabled}, native_effort={native_thinking_config.effort}" ) # Use core function to build payload @@ -483,6 +842,7 @@ def anthropic_to_kiro( conversation_id=conversation_id, profile_arn=profile_arn, thinking_config=thinking_config, + native_thinking_config=native_thinking_config, ) return result.payload diff --git a/kiro/converters_core.py b/kiro/converters_core.py index bd890207..23f458b9 100644 --- a/kiro/converters_core.py +++ b/kiro/converters_core.py @@ -30,7 +30,9 @@ to convert their formats to Kiro API format. """ +import base64 import json +import uuid from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple @@ -43,10 +45,22 @@ FAKE_REASONING_BUDGET_CAP, KIRO_MAX_PAYLOAD_BYTES, AUTO_TRIM_PAYLOAD, + NATIVE_REASONING_ENABLED, + NATIVE_EFFORT_SCHEMA_BY_MODEL, + VALID_EFFORT_LEVELS, + EFFORT_LEVEL_ALIASES, + DEFAULT_EFFORT_LEVEL, + KIRO_NATIVE_THINKING_MODE, + KIRO_NATIVE_THINKING_DISPLAY, ) +from kiro import native_reasoning from kiro.payload_guards import check_payload_size, trim_payload_to_limit +SYSTEM_PROMPT_ACK = "I will follow these instructions." +KIRO_CACHE_POINT: Dict[str, str] = {"type": "default"} + + # ================================================================================================== # Data Classes for Unified Message Format # ================================================================================================== @@ -78,6 +92,22 @@ class ThinkingConfig: """ enabled: bool = True budget_tokens: Optional[int] = None + effort_level: Optional[str] = None # raw reasoning_effort (low/medium/high/xhigh/max) for native effort + + +@dataclass +class NativeThinkingConfig: + """ + Native Kiro/Claude adaptive thinking configuration. + + Attributes: + enabled: Whether to add native adaptive thinking fields. + effort: Adaptive thinking effort level. + display: Whether upstream should return summarized or omitted thinking text. + """ + enabled: bool = False + effort: Optional[str] = None + display: str = "summarized" @dataclass @@ -101,6 +131,7 @@ class UnifiedMessage: tool_calls: Optional[List[Dict[str, Any]]] = None tool_results: Optional[List[Dict[str, Any]]] = None images: Optional[List[Dict[str, Any]]] = None + cache_point: bool = False @dataclass @@ -131,6 +162,441 @@ class KiroPayloadResult: tool_documentation: str = "" +# ================================================================================================== +# Native Thinking Support +# ================================================================================================== + +REASONING_EFFORT_BUDGET_RATIOS: Dict[str, float] = { + "none": 0.0, + "minimal": 0.10, + "low": 0.20, + "medium": 0.50, + "high": 0.80, + "xhigh": 0.95, + "max": 1.0, +} + + +NATIVE_THINKING_SUPPORTED_MODELS = ( + "claude-opus-4.8", + "claude-opus-4.7", + "claude-opus-4.6", + "claude-sonnet-4.6", +) + + +def reasoning_effort_to_budget(max_tokens: int, effort: str) -> int: + """ + Convert a reasoning effort level to a fake thinking token budget. + + Args: + max_tokens: Maximum output tokens for the request. + effort: Reasoning effort level. + + Returns: + Thinking budget in tokens. + + Raises: + ValueError: If the effort level is not supported. + """ + try: + ratio = REASONING_EFFORT_BUDGET_RATIOS[effort] + except KeyError as exc: + supported_efforts = ", ".join(sorted(REASONING_EFFORT_BUDGET_RATIOS)) + raise ValueError( + f"Unsupported reasoning effort '{effort}'. Supported values: {supported_efforts}" + ) from exc + + return int(max_tokens * ratio) + + +def supports_native_adaptive_thinking(model_id: str) -> bool: + """ + Check whether a Kiro model is known to support native adaptive thinking. + + Args: + model_id: Internal Kiro model ID. + + Returns: + True when native adaptive thinking should be attempted. + """ + normalized_model = model_id.lower() + return any(model in normalized_model for model in NATIVE_THINKING_SUPPORTED_MODELS) + + +def normalize_native_thinking_effort(effort: Optional[str]) -> Optional[str]: + """ + Normalize client effort values to Claude adaptive thinking effort values. + + Args: + effort: Client-provided effort level. + + Returns: + Native effort level, or None if the value disables or cannot map to native thinking. + """ + if effort is None: + return None + + if effort == "none": + return None + + if effort == "minimal": + return "low" + + if effort in ("low", "medium", "high", "xhigh", "max"): + return effort + + logger.warning(f"Unsupported native thinking effort '{effort}'. Native thinking disabled.") + return None + + +def build_native_thinking_config(model_id: str, effort: Optional[str]) -> NativeThinkingConfig: + """ + Build native adaptive thinking configuration from model and client effort. + + Args: + model_id: Internal Kiro model ID. + effort: Client effort level. + + Returns: + NativeThinkingConfig for payload construction. + """ + if KIRO_NATIVE_THINKING_MODE == "off": + return NativeThinkingConfig(enabled=False) + + if not supports_native_adaptive_thinking(model_id): + return NativeThinkingConfig(enabled=False) + + native_effort = normalize_native_thinking_effort(effort) + if native_effort is None: + if KIRO_NATIVE_THINKING_MODE != "force": + return NativeThinkingConfig(enabled=False) + native_effort = "high" + + return NativeThinkingConfig( + enabled=True, + effort=native_effort, + display=KIRO_NATIVE_THINKING_DISPLAY, + ) + + +def apply_native_thinking_fields( + payload: Dict[str, Any], + native_thinking_config: Optional[NativeThinkingConfig], +) -> None: + """ + Add native Kiro/Claude adaptive thinking fields to the payload in-place. + + Args: + payload: Kiro API payload. + native_thinking_config: Native thinking configuration. + """ + if not native_thinking_config or not native_thinking_config.enabled: + return + + payload["thinking"] = { + "type": "adaptive", + "effort": native_thinking_config.effort, + "display": native_thinking_config.display, + } + + +def build_native_effort_fields( + model_id: str, + effort_level: Optional[str], +) -> Optional[Dict[str, Any]]: + """ + Build the Bedrock Converse additionalModelRequestFields object that carries + native reasoning effort to the Kiro backend. + + Args: + model_id: Resolved Kiro model ID (dot form, e.g. "claude-opus-4.8") + effort_level: Client-supplied reasoning_effort (low/medium/high/xhigh/max, + or aliases like "minimal"). None disables native effort. + + Returns: + The additionalModelRequestFields dict, or None if not applicable. + """ + if not effort_level: + return None + + # Normalize aliases (minimal->low, etc.) before matching against advertised efforts. + level = EFFORT_LEVEL_ALIASES.get(effort_level.lower(), effort_level.lower()) + + # Prefer the live catalog's advertised schema (per-model schema type + valid effort + # enum). Fall back to the static config table only when the catalog is unavailable. + descriptor = native_reasoning.get_descriptor(model_id) + if descriptor is not None: + schema = descriptor.schema_type + clamped = native_reasoning.clamp_effort(level, descriptor.valid_efforts) + if clamped is None: + logger.debug(f"Model '{model_id}' advertises no effort levels; skipping") + return None + if clamped != level: + logger.debug( + f"Clamped effort '{level}' -> '{clamped}' for '{model_id}' " + f"(advertised: {'/'.join(descriptor.valid_efforts)})" + ) + level = clamped + else: + schema = NATIVE_EFFORT_SCHEMA_BY_MODEL.get(model_id) + if not schema: + logger.debug(f"Model '{model_id}' does not support native effort; skipping") + return None + if level not in VALID_EFFORT_LEVELS: + logger.warning(f"Invalid effort level '{effort_level}' for native effort; skipping") + return None + + if schema == "output_config": + fields = { + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": level}, + } + elif schema == "reasoning": + fields = {"reasoning": {"effort": level}} + else: + logger.warning(f"Unknown effort schema '{schema}' for model '{model_id}'; skipping") + return None + + logger.debug(f"Native effort for '{model_id}': schema={schema}, effort={level}") + return fields + + +# ================================================================================================== +# Document / PDF Extraction +# ================================================================================================== + +import re as _re + +_PDF_TEXT_OPERAND_RE = _re.compile(rb"\((.+?)\)\s*Tj") +_PDF_STRING_RE = _re.compile(rb"^\((.+)\)$") + + +def _parse_data_url(data_url: str) -> Tuple[str, bytes]: + """ + Parse a data URL and return (mime_type, decoded_bytes). + + Raises: + ValueError: If the data URL is malformed or contains invalid base64 data. + """ + if not data_url.startswith("data:"): + raise ValueError("Not a data URL") + + header, data = data_url.split(",", 1) + mime = header.split(":", 1)[1].split(";", 1)[0] + try: + raw = base64.b64decode(data) + except Exception as exc: + raise ValueError(f"Invalid base64 data in data URL: {exc}") from exc + return mime, raw + + +def _decode_base64_document(base64_content: str) -> Tuple[str, bytes]: + """ + Decode a base64 document and detect its MIME type from magic bytes. + + Args: + base64_content: Base64-encoded document data. + + Returns: + Tuple of (mime_type, decoded_bytes). + """ + try: + raw = base64.b64decode(base64_content) + except Exception: + return "application/octet-stream", base64.b64decode(base64_content) + + if raw[:4] == b"%PDF": + return "application/pdf", raw + if raw[:4] == b"\x89PNG": + return "image/png", raw + if raw[:3] == b"\xff\xd8\xff": + return "image/jpeg", raw + if raw[:4] in (b"RIFF", b"\x00\x00\x00\x1c"): + return "image/webp", raw + if raw[:6] in (b"GIF87a", b"GIF89a"): + return "image/gif", raw + if raw[:4] == b"\x00\x00\x01\x00": + return "image/x-icon", raw + + return "application/octet-stream", raw + + +def _decode_pdf_string(raw: bytes) -> str: + """Decode a PDF literal string, handling common escape sequences.""" + try: + decoded = raw.decode("latin-1") + except Exception: + return raw.decode("utf-8", errors="replace") + + result: list[str] = [] + i = 0 + while i < len(decoded): + ch = decoded[i] + if ch == "\\" and i + 1 < len(decoded): + nxt = decoded[i + 1] + if nxt == "n": + result.append("\n") + elif nxt == "r": + result.append("\r") + elif nxt == "t": + result.append("\t") + elif nxt == "b": + result.append("\b") + elif nxt == "f": + result.append("\f") + elif nxt == "(": + result.append("(") + elif nxt == ")": + result.append(")") + elif nxt == "\\": + result.append("\\") + elif nxt.isdigit(): + octal = nxt + for _ in range(2): + if i + 2 < len(decoded) and decoded[i + 2].isdigit(): + octal += decoded[i + 2] + i += 1 + result.append(chr(int(octal, 8))) + i += 1 + else: + result.append(nxt) + i += 2 + else: + result.append(ch) + i += 1 + + return "".join(result) + + +def _extract_pdf_text(pdf_bytes: bytes) -> str: + """ + Extract text from a PDF using regex-based parsing of the raw content stream. + + Falls back to a simpler approach for PDFs that don't match the standard pattern. + """ + pages: list[str] = [] + pos = 0 + + while True: + page_start = pdf_bytes.find(b"BT", pos) + if page_start == -1: + break + + page_end = pdf_bytes.find(b"ET", page_start) + if page_end == -1: + break + + page_stream = pdf_bytes[page_start:page_end] + text_parts: list[str] = [] + + for match in _PDF_TEXT_OPERAND_RE.finditer(page_stream): + raw_group = match.group(1) + if not raw_group: + continue + + string_match = _PDF_STRING_RE.match(raw_group) + if string_match: + text_parts.append(_decode_pdf_string(string_match.group(1))) + else: + text_parts.append(_decode_pdf_string(raw_group)) + + pages.append(" ".join(text_parts)) + pos = page_end + 2 + + return "\n\n".join(page.strip() for page in pages if page.strip()) + + +async def _extract_document_source(source: Dict[str, Any], filename: str) -> Optional[Dict[str, Any]]: + """ + Normalize an Anthropic document source into a dict suitable for the + Kiro messages payload: ``{"data": "", "media_type": ""}``. + + Supports: + - ``{"type": "base64", "data": "...", "media_type": "..."}`` + - ``{"type": "url", "url": "data:..."}`` + + Returns ``None`` when the source is unsupported or empty. + """ + if not source: + return None + + source_type = source.get("type", "") + + if source_type == "base64": + data = source.get("data", "") + media_type = source.get("media_type", "application/octet-stream") + + if media_type == "application/pdf" and data: + try: + _, raw = _decode_base64_document(data) + text = _extract_pdf_text(raw) + if text: + return {"type": "text", "text": text} + except Exception as exc: + logger.warning(f"PDF extraction failed for '{filename}': {exc}") + + if data: + mime, _ = _decode_base64_document(data) + return {"data": data, "media_type": mime} + return None + + if source_type == "url": + url = source.get("url", "") + if url.startswith("data:"): + try: + mime, raw = _parse_data_url(url) + + if mime == "application/pdf": + text = _extract_pdf_text(raw) + if text: + return {"type": "text", "text": text} + + return {"data": base64.b64encode(raw).decode("ascii"), "media_type": mime} + except ValueError as exc: + logger.warning(f"Failed to parse data URL for '{filename}': {exc}") + return None + + return None + + +def extract_document_text_from_content_block(content_block: Dict[str, Any]) -> Optional[str]: + """ + Extract text from an Anthropic document content block. + + For PDFs, attempts to extract text content directly. + For other document types, returns None (caller should pass through as-is). + + Returns: + Extracted text if available, None otherwise. + """ + if content_block.get("type") != "document": + return None + + source = content_block.get("source", {}) + source_type = source.get("type", "") + media_type = source.get("media_type", "") + + if media_type != "application/pdf": + return None + + try: + if source_type == "base64": + data = source.get("data", "") + if data: + _, raw = _decode_base64_document(data) + return _extract_pdf_text(raw) + elif source_type == "url": + url = source.get("url", "") + if url.startswith("data:"): + _, raw = _parse_data_url(url) + return _extract_pdf_text(raw) + except Exception as exc: + logger.warning(f"PDF text extraction failed: {exc}") + + return None + + # ================================================================================================== # Text Content Extraction # ================================================================================================== @@ -166,9 +632,15 @@ def extract_text_content(content: Any) -> str: text_parts = [] for item in content: if isinstance(item, dict): - # Skip image and tool_reference blocks - they're handled separately + # Skip image, tool_reference blocks - they're handled separately if item.get("type") in ("image", "image_url", "tool_reference"): continue + # Extract text from document blocks (e.g., PDFs) + if item.get("type") == "document": + doc_text = extract_document_text_from_content_block(item) + if doc_text: + text_parts.append(doc_text) + continue if item.get("type") == "text": text_parts.append(item.get("text", "")) elif "text" in item: @@ -315,6 +787,8 @@ def get_thinking_system_prompt_addition() -> str: """ if not FAKE_REASONING_ENABLED: return "" + if NATIVE_REASONING_ENABLED: + return "" return ( "\n\n---\n" @@ -387,6 +861,8 @@ def inject_thinking_tags(content: str, thinking_config: ThinkingConfig) -> str: # Check if thinking is enabled globally if not FAKE_REASONING_ENABLED: return content + if NATIVE_REASONING_ENABLED: + return content # Check if thinking is enabled for this request if not thinking_config.enabled: @@ -1402,6 +1878,56 @@ def build_kiro_history(messages: List[UnifiedMessage], model_id: str) -> List[Di # Main Payload Building # ================================================================================================== +def build_kiro_system_history( + system_prompt: str, + model_id: str, + system_cache_point: bool = False, +) -> List[Dict[str, Any]]: + """ + Build a synthetic Kiro history entry wrapping the system prompt as a + user/assistant exchange. + + When Kiro API receives a ``systemPrompt``, it is passed directly to the + model but does not appear in conversation history. For clients like OpenAI + Codex that depend on the system prompt being present in history for + continuation, this function creates a minimal synthetic exchange. + + Args: + system_prompt: System prompt content to wrap. + model_id: Internal Kiro model ID. + system_cache_point: Whether to add a cache point to the system prompt. + + Returns: + List containing a single history entry, or empty list if the prompt is empty. + """ + if not system_prompt: + return [] + + system_content = system_prompt + + if system_cache_point: + system_content = [ + {"text": system_prompt}, + {"text": SYSTEM_PROMPT_ACK}, + KIRO_CACHE_POINT, + ] + + return [ + { + "userInputMessage": { + "content": system_content, + "modelId": model_id, + "origin": "AI_EDITOR", + } + }, + { + "assistantResponseMessage": { + "content": "", + } + }, + ] + + def build_kiro_payload( messages: List[UnifiedMessage], system_prompt: str, @@ -1409,7 +1935,9 @@ def build_kiro_payload( tools: Optional[List[UnifiedTool]], conversation_id: str, profile_arn: str, - thinking_config: ThinkingConfig + thinking_config: ThinkingConfig, + system_cache_point: bool = False, + native_thinking_config: Optional[NativeThinkingConfig] = None, ) -> KiroPayloadResult: """ Builds complete payload for Kiro API from unified data. @@ -1425,6 +1953,8 @@ def build_kiro_payload( conversation_id: Unique conversation ID profile_arn: AWS CodeWhisperer profile ARN thinking_config: Thinking configuration from API adapter + system_cache_point: Whether to add cache point to system prompt in history + native_thinking_config: Native Kiro/Claude adaptive thinking configuration Returns: KiroPayloadResult with payload and tool documentation @@ -1485,23 +2015,19 @@ def build_kiro_payload( # Build history (all messages except the last one) history_messages = merged_messages[:-1] if len(merged_messages) > 1 else [] - # If there's a system prompt, add it to the first user message in history - if full_system_prompt and history_messages: - first_msg = history_messages[0] - if first_msg.role == "user": - original_content = extract_text_content(first_msg.content) - first_msg.content = f"{full_system_prompt}\n\n{original_content}" + # Build system history from system prompt + system_history = build_kiro_system_history(full_system_prompt, model_id, system_cache_point) history = build_kiro_history(history_messages, model_id) + # Prepend system history to conversation history + if system_history: + history = system_history + history + # Current message (the last one) current_message = merged_messages[-1] current_content = extract_text_content(current_message.content) - # If system prompt exists but history is empty - add to current message - if full_system_prompt and not history: - current_content = f"{full_system_prompt}\n\n{current_content}" - # If current message is assistant, need to add it to history # and create user message placeholder if current_message.role == "assistant": @@ -1570,6 +2096,8 @@ def build_kiro_payload( "conversationState": { "chatTriggerType": "MANUAL", "conversationId": conversation_id, + "agentContinuationId": str(uuid.uuid4()), + "agentTaskType": "vibe", "currentMessage": { "userInputMessage": user_input_message } @@ -1584,6 +2112,14 @@ def build_kiro_payload( if profile_arn: payload["profileArn"] = profile_arn + # Build additionalModelRequestFields for native reasoning effort + native_effort_fields = build_native_effort_fields(model_id, thinking_config.effort_level) + if native_effort_fields: + payload["additionalModelRequestFields"] = native_effort_fields + + # Apply native thinking fields (thinking.type=adaptive, display, effort) + apply_native_thinking_fields(payload, native_thinking_config) + # Payload size guard — auto-trim if enabled if AUTO_TRIM_PAYLOAD: payload_size = check_payload_size(payload) diff --git a/kiro/converters_openai.py b/kiro/converters_openai.py index 137d041f..d2e20e6c 100644 --- a/kiro/converters_openai.py +++ b/kiro/converters_openai.py @@ -33,7 +33,7 @@ from loguru import logger -from kiro.config import HIDDEN_MODELS +from kiro.config import HIDDEN_MODELS, MODEL_ALIASES from kiro.model_resolver import get_model_id_for_kiro from kiro.models_openai import ChatMessage, ChatCompletionRequest, Tool @@ -44,6 +44,9 @@ UnifiedMessage, UnifiedTool, ThinkingConfig, + NativeThinkingConfig, + build_native_thinking_config, + reasoning_effort_to_budget, build_kiro_payload as core_build_kiro_payload, ) @@ -334,7 +337,7 @@ def extract_thinking_config_from_openai(request: ChatCompletionRequest) -> Think Handles reasoning_effort parameter: - "none" → disabled (no thinking tags injected) - - "minimal", "low", "medium", "high", "xhigh" → enabled with percentage-based budget + - "minimal", "low", "medium", "high", "xhigh", "max" → enabled with percentage-based budget - None (not specified) → enabled with default budget Args: @@ -358,7 +361,7 @@ def extract_thinking_config_from_openai(request: ChatCompletionRequest) -> Think >>> request.reasoning_effort = "high" >>> request.max_tokens = 4096 >>> extract_thinking_config_from_openai(request) - ThinkingConfig(enabled=True, budget_tokens=3276) # 80% of 4096 + ThinkingConfig(enabled=True, budget_tokens=3276, effort_level='high') # 80% of 4096 """ if not request.reasoning_effort: # No reasoning_effort specified → use defaults @@ -368,6 +371,10 @@ def extract_thinking_config_from_openai(request: ChatCompletionRequest) -> Think # Explicitly disabled return ThinkingConfig(enabled=False, budget_tokens=None) + # Resolve model aliases for downstream native effort support + raw_model = request.model + resolved_model = MODEL_ALIASES.get(raw_model, raw_model) + # Calculate budget from reasoning_effort # Get max_tokens from request (OUTPUT tokens limit) max_tokens = request.max_tokens or request.max_completion_tokens @@ -380,10 +387,11 @@ def extract_thinking_config_from_openai(request: ChatCompletionRequest) -> Think logger.debug( f"Extracted thinking config from OpenAI: reasoning_effort='{request.reasoning_effort}', " + f"model='{raw_model}' -> '{resolved_model}', " f"max_tokens={max_tokens}, budget={budget}" ) - return ThinkingConfig(enabled=True, budget_tokens=budget) + return ThinkingConfig(enabled=True, budget_tokens=budget, effort_level=request.reasoning_effort) # ================================================================================================== @@ -425,11 +433,15 @@ def build_kiro_payload( # Extract thinking configuration from reasoning_effort thinking_config = extract_thinking_config_from_openai(request_data) + # Build native thinking config from model and reasoning_effort + native_thinking_config = build_native_thinking_config(model_id, request_data.reasoning_effort) + logger.debug( f"Converting OpenAI request: model={request_data.model} -> {model_id}, " f"messages={len(unified_messages)}, tools={len(unified_tools) if unified_tools else 0}, " f"system_prompt_length={len(system_prompt)}, " - f"thinking_enabled={thinking_config.enabled}, thinking_budget={thinking_config.budget_tokens}" + f"thinking_enabled={thinking_config.enabled}, thinking_budget={thinking_config.budget_tokens}, " + f"native_thinking_enabled={native_thinking_config.enabled}, native_thinking_effort={native_thinking_config.effort}" ) # Use core function to build payload @@ -440,7 +452,8 @@ def build_kiro_payload( tools=unified_tools, conversation_id=conversation_id, profile_arn=profile_arn, - thinking_config=thinking_config + thinking_config=thinking_config, + native_thinking_config=native_thinking_config, ) return result.payload \ No newline at end of file diff --git a/kiro/converters_responses.py b/kiro/converters_responses.py new file mode 100644 index 00000000..a73dd2c3 --- /dev/null +++ b/kiro/converters_responses.py @@ -0,0 +1,363 @@ +# -*- 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 . + +""" +Converters for transforming OpenAI Responses API format to Kiro format. + +This is the adapter layer for the Responses API (used by Codex CLI). It converts +Responses-specific structures into the unified format used by converters_core.py. + +Key differences handled here: +- `input` may be a string or a list of typed items +- `instructions` (+ any system/developer messages) become the system prompt +- `function_call` / `function_call_output` are TOP-LEVEL items (not nested in a + message). We emit them as separate assistant/user UnifiedMessages carrying + tool_calls / tool_results; the core layer merges adjacent same-role messages. +- `reasoning` items (id=rs_..., encrypted_content) are IGNORED. Kiro produces + plaintext thinking only, so we cannot honor server-side reasoning references. + Passing them through would trigger "Item with id rs_... not found" style errors. +- Tools use a flat structure (no nested `function` object) +""" + +import json +from typing import Any, Dict, List, Optional, Tuple + +from loguru import logger + +from kiro.config import HIDDEN_MODELS +from kiro.model_resolver import get_model_id_for_kiro +from kiro.models_responses import ResponsesRequest, ResponsesTool + +from kiro.converters_core import ( + extract_text_content, + UnifiedMessage, + UnifiedTool, + ThinkingConfig, + build_kiro_payload as core_build_kiro_payload, +) +from kiro.converters_openai import reasoning_effort_to_budget + + +# ================================================================================================== +# Content part normalization +# ================================================================================================== + +def _normalize_content_parts(content: Any) -> Any: + """ + Normalize Responses content parts into blocks the core layer understands. + + Responses uses part types like `input_text`, `output_text`, `input_image`. + The core's extract_text_content()/extract_images_from_content() understand + `text` blocks and OpenAI-style `image_url` blocks, so we translate. + + Args: + content: Message content (string or list of content parts) + + Returns: + Normalized content (string passed through; list of blocks otherwise) + """ + if content is None: + return "" + if isinstance(content, str): + return content + if not isinstance(content, list): + return str(content) + + normalized: List[Dict[str, Any]] = [] + for part in content: + if not isinstance(part, dict): + if isinstance(part, str): + normalized.append({"type": "text", "text": part}) + continue + + part_type = part.get("type") + + if part_type in ("input_text", "output_text", "text", "summary_text"): + normalized.append({"type": "text", "text": part.get("text", "")}) + elif part_type in ("input_image", "image"): + # Responses puts the URL/data URL directly in `image_url` (a string) + url = part.get("image_url") + if isinstance(url, dict): + url = url.get("url", "") + if url: + normalized.append({"type": "image_url", "image_url": {"url": url}}) + elif "text" in part: + normalized.append({"type": "text", "text": part.get("text", "")}) + + return normalized + + +# ================================================================================================== +# Input items -> Unified messages +# ================================================================================================== + +def _function_call_to_unified(item: Dict[str, Any]) -> Dict[str, Any]: + """Convert a Responses function_call item to a unified tool_call.""" + arguments = item.get("arguments", "{}") + if not isinstance(arguments, str): + # Some clients may send arguments as an object + try: + arguments = json.dumps(arguments, ensure_ascii=False) + except (TypeError, ValueError): + arguments = "{}" + return { + "id": item.get("call_id") or item.get("id") or "", + "type": "function", + "function": { + "name": item.get("name", ""), + "arguments": arguments or "{}", + }, + } + + +def _function_call_output_to_unified(item: Dict[str, Any]) -> Dict[str, Any]: + """Convert a Responses function_call_output item to a unified tool_result.""" + output = item.get("output", "") + if isinstance(output, str): + content_text = output + else: + # output may be a list of content parts or structured data + content_text = extract_text_content(_normalize_content_parts(output)) + if not content_text and output is not None: + try: + content_text = json.dumps(output, ensure_ascii=False) + except (TypeError, ValueError): + content_text = str(output) + return { + "type": "tool_result", + "tool_use_id": item.get("call_id") or item.get("id") or "", + "content": content_text or "(empty result)", + } + + +def convert_responses_input_to_unified( + request: ResponsesRequest, +) -> Tuple[str, List[UnifiedMessage]]: + """ + Convert Responses `input` + `instructions` into (system_prompt, unified_messages). + + Args: + request: The parsed Responses request + + Returns: + Tuple of (system_prompt, unified_messages) + """ + system_parts: List[str] = [] + if request.instructions: + system_parts.append(request.instructions) + + unified: List[UnifiedMessage] = [] + + # Plain string input -> single user message + if isinstance(request.input, str): + unified.append(UnifiedMessage(role="user", content=request.input)) + return "\n".join(system_parts).strip(), unified + + ignored_reasoning = 0 + tool_calls_seen = 0 + tool_results_seen = 0 + + for item in request.input: + if not isinstance(item, dict): + # Bare string item -> user message + if isinstance(item, str): + unified.append(UnifiedMessage(role="user", content=item)) + continue + + item_type = item.get("type", "message") + + if item_type == "message": + role = item.get("role", "user") + normalized = _normalize_content_parts(item.get("content")) + + # system / developer messages fold into the system prompt + if role in ("system", "developer"): + text = extract_text_content(normalized) + if text: + system_parts.append(text) + continue + + unified.append(UnifiedMessage(role=role, content=normalized)) + + elif item_type == "function_call": + unified.append(UnifiedMessage( + role="assistant", + content="", + tool_calls=[_function_call_to_unified(item)], + )) + tool_calls_seen += 1 + + elif item_type == "function_call_output": + unified.append(UnifiedMessage( + role="user", + content="", + tool_results=[_function_call_output_to_unified(item)], + )) + tool_results_seen += 1 + + elif item_type == "reasoning": + # Ignore server-side reasoning items - Kiro has no encrypted reasoning + ignored_reasoning += 1 + + else: + # Unknown item type - try to salvage any text content + text = extract_text_content(_normalize_content_parts(item.get("content"))) + if text: + unified.append(UnifiedMessage(role=item.get("role", "user"), content=text)) + else: + logger.debug(f"Ignoring unsupported Responses input item type: {item_type}") + + if ignored_reasoning or tool_calls_seen or tool_results_seen: + logger.debug( + f"Converted Responses input: {len(unified)} messages, " + f"{tool_calls_seen} function_call(s), {tool_results_seen} function_call_output(s), " + f"{ignored_reasoning} reasoning item(s) ignored" + ) + + return "\n".join(system_parts).strip(), unified + + +# ================================================================================================== +# Tools +# ================================================================================================== + +def convert_responses_tools_to_unified( + tools: Optional[List[ResponsesTool]], +) -> Optional[List[UnifiedTool]]: + """ + Convert Responses tools (flat format) to unified tools. + + Only `function` type tools are supported; built-in tool types (web_search, + file_search, etc.) are skipped since Kiro only accepts custom functions. + + Args: + tools: List of ResponsesTool objects + + Returns: + List of UnifiedTool objects, or None if no usable tools + """ + if not tools: + return None + + unified_tools: List[UnifiedTool] = [] + for tool in tools: + if tool.type != "function": + logger.debug(f"Skipping non-function Responses tool: type={tool.type}") + continue + if not tool.name: + logger.warning("Skipping Responses function tool with no name") + continue + unified_tools.append(UnifiedTool( + name=tool.name, + description=tool.description, + input_schema=tool.parameters, + )) + + return unified_tools if unified_tools else None + + +# ================================================================================================== +# Thinking configuration +# ================================================================================================== + +def extract_thinking_config_from_responses(request: ResponsesRequest) -> ThinkingConfig: + """ + Extract thinking configuration from a Responses request's `reasoning.effort`. + + - No reasoning specified -> enabled with default budget + - effort "none" -> disabled + - effort minimal/low/medium/high/xhigh -> percentage-based budget + + Args: + request: The parsed Responses request + + Returns: + ThinkingConfig for the core layer + """ + effort = request.reasoning.effort if request.reasoning else None + + if not effort: + return ThinkingConfig(enabled=True, budget_tokens=None) + + if effort == "none": + return ThinkingConfig(enabled=False, budget_tokens=None) + + max_tokens = request.max_output_tokens or 4096 + + # Unknown effort levels fall back to the default budget + try: + budget = reasoning_effort_to_budget(max_tokens, effort) + except KeyError: + logger.debug(f"Unknown reasoning effort '{effort}', using default budget") + return ThinkingConfig(enabled=True, budget_tokens=None) + + logger.debug( + f"Extracted thinking config from Responses: effort='{effort}', " + f"max_output_tokens={max_tokens}, budget={budget}" + ) + return ThinkingConfig(enabled=True, budget_tokens=budget) + + +# ================================================================================================== +# Main entry point +# ================================================================================================== + +def build_kiro_payload( + request_data: ResponsesRequest, + conversation_id: str, + profile_arn: str, +) -> dict: + """ + Build a complete Kiro API payload from a Responses API request. + + Args: + request_data: Request in Responses API format + conversation_id: Unique conversation ID + profile_arn: AWS CodeWhisperer profile ARN + + Returns: + Payload dictionary for POST to Kiro API + + Raises: + ValueError: If there are no messages to send + """ + system_prompt, unified_messages = convert_responses_input_to_unified(request_data) + unified_tools = convert_responses_tools_to_unified(request_data.tools) + + model_id = get_model_id_for_kiro(request_data.model, HIDDEN_MODELS) + thinking_config = extract_thinking_config_from_responses(request_data) + + logger.debug( + f"Converting Responses request: model={request_data.model} -> {model_id}, " + f"messages={len(unified_messages)}, tools={len(unified_tools) if unified_tools else 0}, " + f"system_prompt_length={len(system_prompt)}, " + f"thinking_enabled={thinking_config.enabled}, thinking_budget={thinking_config.budget_tokens}" + ) + + result = core_build_kiro_payload( + messages=unified_messages, + system_prompt=system_prompt, + model_id=model_id, + tools=unified_tools, + conversation_id=conversation_id, + profile_arn=profile_arn, + thinking_config=thinking_config, + ) + + return result.payload diff --git a/kiro/debug_logger.py b/kiro/debug_logger.py index 68802076..ecbcf1d7 100644 --- a/kiro/debug_logger.py +++ b/kiro/debug_logger.py @@ -35,6 +35,9 @@ import io import json import shutil +from contextvars import ContextVar +from datetime import datetime +from itertools import count from pathlib import Path from typing import Optional from loguru import logger @@ -63,6 +66,11 @@ def __init__(self): if self._initialized: return self.debug_dir = Path(DEBUG_DIR) + self._request_counter = count(1) + self._active_request_dir: ContextVar[Optional[Path]] = ContextVar( + "debug_logger_active_request_dir", + default=None, + ) self._initialized = True # Buffers for "errors" mode @@ -90,6 +98,70 @@ def _clear_buffers(self): self._raw_chunks_buffer.clear() self._modified_chunks_buffer.clear() self._clear_app_logs_buffer() + self._active_request_dir.set(None) + + def _clear_latest_files(self) -> None: + """ + Remove top-level "latest request" files while preserving request archives. + + In DEBUG_MODE=all, every request is written to a dedicated directory under + debug_logs/requests/. The top-level files are still maintained as a quick + view of the latest request, so only those files are removed here. + """ + latest_files = ( + "request_body.json", + "kiro_request_body.json", + "response_body.json", + "response_stream_raw.txt", + "response_stream_modified.txt", + "error_info.json", + "app_logs.txt", + ) + for name in latest_files: + path = self.debug_dir / name + if path.exists() and path.is_file(): + path.unlink() + + def _create_request_dir(self) -> Path: + """ + Create a unique directory for the current DEBUG_MODE=all request. + + Returns: + Path to the per-request debug directory. + """ + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S-%f") + sequence = next(self._request_counter) + request_dir = self.debug_dir / "requests" / f"{timestamp}-{sequence:06d}" + request_dir.mkdir(parents=True, exist_ok=True) + return request_dir + + def _get_request_dir(self) -> Path: + """ + Return the current request's write directory. + + If a method is called without prepare_new_request(), fall back to the root + debug directory. This keeps direct unit tests and manual calls simple. + """ + if self._is_immediate_write(): + active_dir = self._active_request_dir.get() + if active_dir is not None: + return active_dir + return self.debug_dir + + def _get_output_paths(self, filename: str) -> list[Path]: + """ + Return output paths for a debug artifact. + + DEBUG_MODE=all writes both the archived per-request file and a top-level + latest copy. Other modes write only the top-level file. + """ + request_dir = self._get_request_dir() + paths = [request_dir / filename] + + if self._is_immediate_write() and request_dir != self.debug_dir: + paths.append(self.debug_dir / filename) + + return paths def _clear_app_logs_buffer(self): """Clears the application logs buffer and removes sink.""" @@ -130,7 +202,7 @@ def prepare_new_request(self): """ Prepares the logger for a new request. - In "all" mode: clears the logs folder. + In "all" mode: creates a new per-request log directory. In "errors" mode: clears buffers. In both modes: sets up application log capture. """ @@ -144,12 +216,13 @@ def prepare_new_request(self): self._setup_app_logs_capture() if self._is_immediate_write(): - # "all" mode - clear folder and recreate + # "all" mode - preserve history and rotate only the top-level latest files. try: - if self.debug_dir.exists(): - shutil.rmtree(self.debug_dir) self.debug_dir.mkdir(parents=True, exist_ok=True) - logger.debug(f"[DebugLogger] Directory {self.debug_dir} cleared for new request.") + self._clear_latest_files() + request_dir = self._create_request_dir() + self._active_request_dir.set(request_dir) + logger.debug(f"[DebugLogger] Request logs will be written to {request_dir}.") except Exception as e: logger.error(f"[DebugLogger] Error preparing directory: {e}") @@ -334,46 +407,46 @@ def discard_buffers(self): def _write_request_body_to_file(self, body: bytes): """Writes request body to file.""" try: - file_path = self.debug_dir / "request_body.json" - try: - json_obj = json.loads(body) - with open(file_path, "w", encoding="utf-8") as f: - json.dump(json_obj, f, indent=2, ensure_ascii=False) - except json.JSONDecodeError: - with open(file_path, "wb") as f: - f.write(body) + for file_path in self._get_output_paths("request_body.json"): + try: + json_obj = json.loads(body) + with open(file_path, "w", encoding="utf-8") as f: + json.dump(json_obj, f, indent=2, ensure_ascii=False) + except json.JSONDecodeError: + with open(file_path, "wb") as f: + f.write(body) except Exception as e: logger.error(f"[DebugLogger] Error writing request_body: {e}") def _write_kiro_request_body_to_file(self, body: bytes): """Writes Kiro request body to file.""" try: - file_path = self.debug_dir / "kiro_request_body.json" - try: - json_obj = json.loads(body) - with open(file_path, "w", encoding="utf-8") as f: - json.dump(json_obj, f, indent=2, ensure_ascii=False) - except json.JSONDecodeError: - with open(file_path, "wb") as f: - f.write(body) + for file_path in self._get_output_paths("kiro_request_body.json"): + try: + json_obj = json.loads(body) + with open(file_path, "w", encoding="utf-8") as f: + json.dump(json_obj, f, indent=2, ensure_ascii=False) + except json.JSONDecodeError: + with open(file_path, "wb") as f: + f.write(body) except Exception as e: logger.error(f"[DebugLogger] Error writing kiro_request_body: {e}") def _append_raw_chunk_to_file(self, chunk: bytes): """Appends raw chunk to file.""" try: - file_path = self.debug_dir / "response_stream_raw.txt" - with open(file_path, "ab") as f: - f.write(chunk) + for file_path in self._get_output_paths("response_stream_raw.txt"): + with open(file_path, "ab") as f: + f.write(chunk) except Exception: pass def _append_modified_chunk_to_file(self, chunk: bytes): """Appends modified chunk to file.""" try: - file_path = self.debug_dir / "response_stream_modified.txt" - with open(file_path, "ab") as f: - f.write(chunk) + for file_path in self._get_output_paths("response_stream_modified.txt"): + with open(file_path, "ab") as f: + f.write(chunk) except Exception: pass @@ -389,11 +462,11 @@ def _write_app_logs_to_file(self): # Ensure directory exists self.debug_dir.mkdir(parents=True, exist_ok=True) - file_path = self.debug_dir / "app_logs.txt" - with open(file_path, "w", encoding="utf-8") as f: - f.write(logs_content) + for file_path in self._get_output_paths("app_logs.txt"): + with open(file_path, "w", encoding="utf-8") as f: + f.write(logs_content) - logger.debug(f"[DebugLogger] App logs saved to {file_path}") + logger.debug(f"[DebugLogger] App logs saved") except Exception as e: # Don't log error via logger to avoid recursion pass diff --git a/kiro/http_client.py b/kiro/http_client.py index 798233f5..571d07df 100644 --- a/kiro/http_client.py +++ b/kiro/http_client.py @@ -173,7 +173,8 @@ async def request_with_retry( url: str, json_data: Optional[dict] = None, params: Optional[dict] = None, - stream: bool = False + stream: bool = False, + extra_headers: Optional[dict] = None ) -> httpx.Response: """ Executes an HTTP request with retry logic. @@ -214,7 +215,9 @@ async def request_with_retry( # Get current token token = await self.auth_manager.get_access_token() headers = get_kiro_headers(self.auth_manager, token) - + if extra_headers: + headers.update(extra_headers) + # Build request kwargs based on parameters request_kwargs = {"headers": headers} diff --git a/kiro/mcp_tools.py b/kiro/mcp_tools.py index c908ec55..4cd52d01 100644 --- a/kiro/mcp_tools.py +++ b/kiro/mcp_tools.py @@ -41,6 +41,8 @@ from loguru import logger from kiro.tokenizer import count_message_tokens, count_tokens +from kiro.profile_arn import profile_arn_for_payload +from kiro.utils import get_kiro_headers # Import debug_logger try: @@ -136,6 +138,11 @@ async def call_kiro_mcp_api( } } + # profileArn is required by runtime.kiro.dev for all auth types (mirror completion paths) + profile_arn = profile_arn_for_payload(auth_manager) + if profile_arn: + mcp_request["profileArn"] = profile_arn + # Log MCP request try: mcp_request_json = json.dumps(mcp_request, ensure_ascii=False, indent=2).encode('utf-8') @@ -147,12 +154,12 @@ 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" - } + # /mcp requires the same Kiro client-identity headers as the completion + # path (else 403); reuse get_kiro_headers and override what differs. + headers = get_kiro_headers(auth_manager, token) + headers["Content-Type"] = "application/json" # /mcp is JSON-RPC + 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}") diff --git a/kiro/model_resolver.py b/kiro/model_resolver.py index 41e367b2..414ec4c7 100644 --- a/kiro/model_resolver.py +++ b/kiro/model_resolver.py @@ -189,12 +189,16 @@ def normalize_model_name(name: str) -> str: return name -def get_model_id_for_kiro(model_name: str, hidden_models: Dict[str, str]) -> str: +def get_model_id_for_kiro( + model_name: str, + hidden_models: Dict[str, str], + aliases: Optional[Dict[str, str]] = None, +) -> str: """ Get the model ID to send to Kiro API. This is a simple helper for converters that don't have access to the full - ModelResolver. It normalizes the name and checks hidden models. + ModelResolver. It applies aliases, normalizes the name, and checks hidden models. For hidden models (like claude-3.7-sonnet), returns the internal Kiro ID. For regular models, returns the normalized name. @@ -202,6 +206,7 @@ def get_model_id_for_kiro(model_name: str, hidden_models: Dict[str, str]) -> str Args: model_name: External model name from client hidden_models: Dict mapping display names to internal Kiro IDs + aliases: Optional dict mapping alias names to real model IDs Returns: Model ID to send to Kiro API @@ -213,8 +218,14 @@ def get_model_id_for_kiro(model_name: str, hidden_models: Dict[str, str]) -> str 'CLAUDE_3_7_SONNET_20250219_V1_0' >>> get_model_id_for_kiro("claude-3-7-sonnet", {"claude-3.7-sonnet": "CLAUDE_3_7_SONNET_20250219_V1_0"}) 'CLAUDE_3_7_SONNET_20250219_V1_0' + >>> get_model_id_for_kiro("auto-kiro", {}, {"auto-kiro": "auto"}) + 'auto' """ - normalized = normalize_model_name(model_name) + resolved_model = (aliases or {}).get(model_name, model_name) + if resolved_model != model_name: + logger.debug(f"Alias resolved: '{model_name}' → '{resolved_model}'") + + normalized = normalize_model_name(resolved_model) internal = hidden_models.get(normalized, normalized) return to_runtime_model_id(internal) diff --git a/kiro/models_anthropic.py b/kiro/models_anthropic.py index c63d60ba..bded1910 100644 --- a/kiro/models_anthropic.py +++ b/kiro/models_anthropic.py @@ -65,6 +65,20 @@ class ThinkingContentBlock(BaseModel): signature: str = "" +class RedactedThinkingContentBlock(BaseModel): + """ + Redacted thinking content block in Anthropic format. + + Claude may return encrypted reasoning when extended thinking is enabled. + Clients can replay these blocks in later requests; the gateway accepts + them for schema compatibility without exposing the opaque data as prompt + text. + """ + + type: Literal["redacted_thinking"] = "redacted_thinking" + data: str + + class ToolUseContentBlock(BaseModel): """ Tool use content block in Anthropic format. @@ -162,14 +176,108 @@ class ImageContentBlock(BaseModel): source: Union[Base64ImageSource, URLImageSource] -# Union type for all content blocks (including images and thinking) +class Base64DocumentSource(BaseModel): + """ + Base64-encoded document source in Anthropic format. + + Anthropic supports PDF documents as content blocks. Kiro does not expose a + matching document input field, so converters extract text from supported + documents and append it to the prompt. + """ + + type: Literal["base64"] = "base64" + media_type: str + data: str + + +class DocumentContentBlock(BaseModel): + """ + Document content block in Anthropic format. + """ + + type: Literal["document"] = "document" + source: Base64DocumentSource + title: Optional[str] = None + + model_config = {"extra": "allow"} + + +# ================================================================================================== +# Server-Side Tool Content Block Models +# ================================================================================================== + + +class ServerToolUseContentBlock(BaseModel): + """ + Server-side tool use block (Anthropic server tools). + + Emitted by Anthropic server-side tools such as web_search, web_fetch and + code_execution, then replayed back into the conversation history by clients + like Claude Code. Accepted for validation only; the converters extract just + text/tool_use/tool_result blocks, so the raw server-tool call is ignored. + """ + + type: Literal["server_tool_use"] = "server_tool_use" + id: str + name: str + input: Dict[str, Any] = Field(default_factory=dict) + + model_config = {"extra": "allow"} + + +class WebSearchToolResultContentBlock(BaseModel): + """ + Web search tool result block (Anthropic server tool). + + Paired with a ServerToolUseContentBlock, this carries the search results + returned by Anthropic's server-side web_search tool. The ``content`` field + is a list of provider-specific result objects (or an error object), so it is + kept loosely typed. Accepted for validation; not forwarded to Kiro. + """ + + type: Literal["web_search_tool_result"] = "web_search_tool_result" + tool_use_id: str + content: Union[List[Dict[str, Any]], Dict[str, Any], str, None] = None + + model_config = {"extra": "allow"} + + +class UnknownContentBlock(BaseModel): + """ + Forward-compatible fallback for unrecognized content blocks. + + Anthropic periodically ships new server-side tool block types (e.g. + code_execution_tool_result, mcp_tool_use). Rather than 422 on every new + block, we accept any block that carries a ``type`` string and preserve its + extra fields. The converters ignore blocks they do not explicitly handle, + so unknown blocks are dropped safely instead of breaking the request. + + MUST be the last member of the ContentBlock union: as a catch-all with only + a ``type: str`` field it would otherwise shadow the specific typed blocks + under Pydantic union resolution. + """ + + type: str + + model_config = {"extra": "allow"} + + +# Union type for all content blocks (including images and thinking). +# Server-side tool blocks (server_tool_use, web_search_tool_result) are accepted +# so clients replaying Anthropic server-tool output into history don't 422; the +# catch-all UnknownContentBlock MUST stay last so it does not shadow typed blocks. ContentBlock = Union[ TextContentBlock, ThinkingContentBlock, + RedactedThinkingContentBlock, ImageContentBlock, + DocumentContentBlock, ToolUseContentBlock, ToolResultContentBlock, ToolReferenceContentBlock, + ServerToolUseContentBlock, + WebSearchToolResultContentBlock, + UnknownContentBlock, ] @@ -183,11 +291,14 @@ class AnthropicMessage(BaseModel): Message in Anthropic format. Attributes: - role: Message role (user or assistant) + role: Message role. Anthropic spec only allows user/assistant, but newer + Claude Code clients sometimes inline "system" messages in the array. + We accept it here and let normalize_message_roles() in converters_core + collapse it to "user" downstream. content: Message content (string or list of content blocks) """ - role: Literal["user", "assistant"] + role: Literal["user", "assistant", "system"] content: Union[str, List[ContentBlock]] model_config = {"extra": "allow"} diff --git a/kiro/models_openai.py b/kiro/models_openai.py index 9551a841..4498767d 100644 --- a/kiro/models_openai.py +++ b/kiro/models_openai.py @@ -125,6 +125,20 @@ class Tool(BaseModel): model_config = {"extra": "allow"} +class ResponseFormat(BaseModel): + """ + Response format request in OpenAI format. + + Supports json_object and json_schema modes. The gateway maps these to + system instructions because Kiro does not expose a native constrained + decoding parameter. + """ + type: str + json_schema: Optional[Dict[str, Any]] = None + + model_config = {"extra": "allow"} + + class ChatCompletionRequest(BaseModel): """ Request for response generation in OpenAI Chat Completions API format. @@ -166,12 +180,15 @@ class ChatCompletionRequest(BaseModel): # Reasoning (OpenAI reasoning models) # Supports all official reasoning_effort levels from OpenAI API - reasoning_effort: Optional[Literal["none", "minimal", "low", "medium", "high", "xhigh"]] = None + reasoning_effort: Optional[Literal["none", "minimal", "low", "medium", "high", "xhigh", "max"]] = None # Tools (function calling) tools: Optional[List[Tool]] = None tool_choice: Optional[Union[str, Dict]] = None + # Response format + response_format: Optional[ResponseFormat] = None + # Compatibility fields (ignored) stream_options: Optional[Dict[str, Any]] = None logit_bias: Optional[Dict[str, float]] = None diff --git a/kiro/models_responses.py b/kiro/models_responses.py new file mode 100644 index 00000000..862d9679 --- /dev/null +++ b/kiro/models_responses.py @@ -0,0 +1,139 @@ +# -*- 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 . + +""" +Pydantic models for OpenAI Responses API (/v1/responses). + +The Responses API is the protocol used by OpenAI Codex CLI. It differs from +Chat Completions in several ways: +- Uses `input` (string or list of typed items) instead of `messages` +- Uses `instructions` instead of a system message +- Tool calls / tool outputs are TOP-LEVEL items, not nested inside a message +- Tools use a flat structure (no nested `function` object) +- Reasoning is configured via a `reasoning` object + +Validation is intentionally lenient (extra="allow") because Codex CLI sends +many fields we do not need, and the `input` items are polymorphic. The +polymorphic items are processed as plain dicts in converters_responses.py. +""" + +from typing import Any, Dict, List, Literal, Optional, Union + +from pydantic import BaseModel, Field + + +# ================================================================================================== +# Reasoning configuration +# ================================================================================================== + +class ReasoningConfig(BaseModel): + """ + Reasoning configuration for the Responses API. + + Attributes: + effort: Reasoning effort level (minimal, low, medium, high, ...) + summary: Reasoning summary mode ("auto", "concise", "detailed", None) + """ + effort: Optional[str] = None + summary: Optional[str] = None + + model_config = {"extra": "allow"} + + +# ================================================================================================== +# Tools +# ================================================================================================== + +class ResponsesTool(BaseModel): + """ + Tool definition in Responses API format. + + Unlike Chat Completions, the function fields are FLAT (not nested under a + `function` key): + {"type": "function", "name": "shell", "description": "...", + "parameters": {...}, "strict": false} + + Non-function tool types (web_search, file_search, etc.) are accepted but + ignored by the converter since Kiro only supports custom function tools. + + Attributes: + type: Tool type (usually "function") + name: Function name (flat format) + description: Function description + parameters: JSON Schema for function parameters + """ + type: str = "function" + name: Optional[str] = None + description: Optional[str] = None + parameters: Optional[Dict[str, Any]] = None + + model_config = {"extra": "allow"} + + +# ================================================================================================== +# Request +# ================================================================================================== + +class ResponsesRequest(BaseModel): + """ + Request body for POST /v1/responses. + + Attributes: + model: Model ID for generation + input: User input - either a plain string or a list of typed items + (message / function_call / function_call_output / reasoning / ...) + instructions: System prompt equivalent + stream: Use streaming (default False) + tools: List of available tools (flat format) + tool_choice: Tool selection strategy + parallel_tool_calls: Whether parallel tool calls are allowed + reasoning: Reasoning configuration (effort / summary) + max_output_tokens: Maximum number of output tokens + temperature / top_p: Generation parameters + store: Whether the response is persisted server-side (Codex sends false) + previous_response_id: Prior response ID (stateful mode - we ignore it) + include: Extra fields to include (e.g. reasoning.encrypted_content) + """ + model: str + input: Union[str, List[Any]] + + instructions: Optional[str] = None + stream: bool = False + + # Tools + tools: Optional[List[ResponsesTool]] = None + tool_choice: Optional[Union[str, Dict[str, Any]]] = None + parallel_tool_calls: Optional[bool] = None + + # Reasoning + reasoning: Optional[ReasoningConfig] = None + + # Generation parameters + max_output_tokens: Optional[int] = None + temperature: Optional[float] = None + top_p: Optional[float] = None + + # Stateless/compat fields (accepted, mostly ignored) + store: Optional[bool] = None + previous_response_id: Optional[str] = None + include: Optional[List[str]] = None + metadata: Optional[Dict[str, Any]] = None + user: Optional[str] = None + + model_config = {"extra": "allow"} diff --git a/kiro/native_reasoning.py b/kiro/native_reasoning.py new file mode 100644 index 00000000..e76bf4e0 --- /dev/null +++ b/kiro/native_reasoning.py @@ -0,0 +1,180 @@ +""" +Catalog-driven native reasoning registry. + +The Kiro control plane (``management.{region}.kiro.dev``) advertises, per model, an +``additionalModelRequestFieldsSchema`` that declares whether the model supports native +adaptive thinking and which effort levels it accepts. For example, ``claude-opus-4.6`` +advertises:: + + { + "type": "object", + "properties": { + "thinking": {"type": "object", "properties": { + "type": {"enum": ["adaptive", "disabled"]}, + "display": {"enum": ["summarized", "omitted"]}}}, + "output_config": {"type": "object", "properties": { + "effort": {"enum": ["low", "medium", "high", "max"], "default": "high"}}}, + "max_tokens": {"type": "integer", "minimum": 1024, "maximum": 64000} + }, + "additionalProperties": false + } + +This module turns that advertised schema into ``NativeEffortDescriptor`` objects and keeps +a process-global registry keyed by model id. It is populated from the live catalog by the +account manager (``register_from_catalog``) and consulted by the converter pipeline when +building ``additionalModelRequestFields``. When the live catalog is unavailable (e.g. the +runtime endpoint returned only static fallback models), callers fall back to the static +config tables — this module only holds what the backend actually advertised. +""" +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +logger = logging.getLogger("kiro-gateway.native_reasoning") + +# Canonical ordering of effort levels from least to most effort. Used to clamp a client +# requested effort onto the set a model actually accepts (e.g. "xhigh" -> "max" when the +# model only advertises low/medium/high/max). +EFFORT_ORDER: List[str] = ["minimal", "low", "medium", "high", "xhigh", "max"] + + +@dataclass +class NativeEffortDescriptor: + """ + Describes how a single model accepts native reasoning effort, derived from the model's + advertised ``additionalModelRequestFieldsSchema``. + + Attributes: + schema_type: Where effort rides in ``additionalModelRequestFields`` — + ``"output_config"`` (``output_config.effort`` alongside a ``thinking`` block) or + ``"reasoning"`` (``reasoning.effort``). + valid_efforts: The effort strings the backend advertised in the enum, in advertised + order. The server strictly validates against this set. + default_effort: The advertised default effort, if any. + """ + + schema_type: str + valid_efforts: List[str] = field(default_factory=list) + default_effort: Optional[str] = None + + +# Process-global registry, keyed by model id (dot form, e.g. "claude-opus-4.6"). +_REGISTRY: Dict[str, NativeEffortDescriptor] = {} + + +def _descriptor_from_schema(schema: Dict[str, Any]) -> Optional[NativeEffortDescriptor]: + """ + Parse a single model's ``additionalModelRequestFieldsSchema`` into a descriptor. + + Returns None when the schema does not advertise a native effort field (i.e. the model + does not support native reasoning). + """ + if not isinstance(schema, dict): + return None + + properties = schema.get("properties") + if not isinstance(properties, dict): + return None + + # Two known shapes: output_config.effort (Claude adaptive) and reasoning.effort. + for schema_type in ("output_config", "reasoning"): + block = properties.get(schema_type) + if not isinstance(block, dict): + continue + effort = (block.get("properties") or {}).get("effort") + if not isinstance(effort, dict): + continue + enum = effort.get("enum") + valid = [str(e) for e in enum] if isinstance(enum, list) and enum else [] + return NativeEffortDescriptor( + schema_type=schema_type, + valid_efforts=valid, + default_effort=effort.get("default"), + ) + + return None + + +def register_from_catalog(models: List[Dict[str, Any]]) -> None: + """ + Rebuild the registry from a freshly fetched model catalog. + + Each entry is a model dict as returned by ``ListAvailableModels``; those that advertise + an ``additionalModelRequestFieldsSchema`` with a native effort field are registered. + Replaces any prior contents so stale models don't linger across a refresh. + + Args: + models: List of model dicts from the control-plane catalog. + """ + new_registry: Dict[str, NativeEffortDescriptor] = {} + for model in models: + if not isinstance(model, dict): + continue + model_id = model.get("modelId") + schema = model.get("additionalModelRequestFieldsSchema") + if not model_id or not schema: + continue + descriptor = _descriptor_from_schema(schema) + if descriptor is not None: + new_registry[model_id] = descriptor + + _REGISTRY.clear() + _REGISTRY.update(new_registry) + if new_registry: + logger.info( + "Native reasoning catalog: %d model(s) advertise effort — %s", + len(new_registry), + ", ".join( + f"{m}({d.schema_type}:{'/'.join(d.valid_efforts) or '?'})" + for m, d in new_registry.items() + ), + ) + else: + logger.debug("Native reasoning catalog: no models advertised effort schema") + + +def get_descriptor(model_id: str) -> Optional[NativeEffortDescriptor]: + """Return the advertised descriptor for a model, or None if the catalog had none.""" + return _REGISTRY.get(model_id) + + +def has_catalog() -> bool: + """True if any model advertised a native effort schema (live catalog was parsed).""" + return bool(_REGISTRY) + + +def clamp_effort(requested: str, valid_efforts: List[str]) -> Optional[str]: + """ + Map a client-requested effort onto the set a model actually accepts. + + Exact matches pass through. Otherwise the requested level is snapped to the nearest + advertised level by position in ``EFFORT_ORDER``, preferring the higher level on a tie + (so ``xhigh`` against ``[low, medium, high, max]`` resolves to ``max``, honoring the + intent of "more than high"). Returns None only when the model advertised no efforts. + + Args: + requested: The client's requested effort (already alias-normalized). + valid_efforts: The effort levels the model advertised. + + Returns: + A valid effort string, or None if ``valid_efforts`` is empty. + """ + if not valid_efforts: + return None + if requested in valid_efforts: + return requested + + if requested not in EFFORT_ORDER: + # Unknown token — fall back to the model's default or its highest advertised level. + return valid_efforts[-1] + + target = EFFORT_ORDER.index(requested) + + def distance(level: str) -> tuple[int, int]: + # Primary: absolute distance. Tie-break: prefer higher position (negative index). + pos = EFFORT_ORDER.index(level) if level in EFFORT_ORDER else 0 + return (abs(pos - target), -pos) + + return min(valid_efforts, key=distance) diff --git a/kiro/parsers.py b/kiro/parsers.py index 9bd7f433..89f3134b 100644 --- a/kiro/parsers.py +++ b/kiro/parsers.py @@ -148,6 +148,255 @@ def parse_bracket_tool_calls(response_text: str) -> List[Dict[str, Any]]: return tool_calls +# Anthropic/Claude models sometimes emit tool calls as XML inside the text +# channel instead of via the backend's structured tool events. The block looks +# like: +# +# +# +# value +# {"k": 1} +# +# +# +# When the backend forwards this verbatim (e.g. when text and a tool call share +# one turn), kg previously had no parser for it, so the raw XML leaked into the +# response content. The functions below recover those calls and strip the XML. +_XML_INVOKE_RE = re.compile( + r'(.*?)', + re.DOTALL, +) +_XML_PARAM_RE = re.compile( + r'(.*?)', + re.DOTALL, +) +# Matches a whole ... wrapper, or a bare run of +# blocks, so we can excise them from the visible content. +_XML_FUNCTION_CALLS_RE = re.compile( + r'.*?', + re.DOTALL, +) + + +def _coerce_xml_param_value(raw: str) -> Any: + """Best-effort typing of an XML body. + + Values arrive as text. If the body is valid JSON (object, array, number, + bool, null, or quoted string) we decode it so structured args round-trip; + otherwise we keep the trimmed string. This mirrors how Anthropic clients + interpret parameter bodies. + """ + stripped = raw.strip() + if stripped == "": + return "" + # Only attempt JSON decode for values that look like JSON containers or + # literals — avoids turning a plain word into a parse error path. + first = stripped[0] + if first in '{[' or stripped in ("true", "false", "null") or _looks_numeric(stripped): + try: + return json.loads(stripped) + except (json.JSONDecodeError, ValueError): + return stripped + return stripped + + +def _looks_numeric(s: str) -> bool: + """True if the string is a bare int/float literal.""" + if not s: + return False + try: + float(s) + return True + except ValueError: + return False + + +def parse_xml_tool_calls(response_text: str) -> List[Dict[str, Any]]: + """Parses tool calls in the Anthropic ```` XML format. + + Some models return tool calls as XML embedded in the text channel instead + of via the backend's structured tool events. This extracts them into the + OpenAI tool-call shape. + + Args: + response_text: Model response text (may contain other prose too). + + Returns: + List of tool calls in OpenAI format. + + Example: + >>> text = 'London' + >>> calls = parse_xml_tool_calls(text) + >>> calls[0]["function"]["name"] + 'get_weather' + """ + if not response_text or " str: + """Removes ````/```` XML blocks from content. + + Used after :func:`parse_xml_tool_calls` so the recovered XML does not leak + into the visible assistant text. Strips the ```` wrapper + when present, then any stray bare ``...`` blocks, and + tidies the surrounding whitespace. + """ + if not response_text or "", "").replace("", "") + return cleaned.strip() + + +# Any of these substrings, appearing at the tail of the streamed-so-far content, +# means an XML tool-call block MIGHT be starting and we must stop emitting until +# we know for sure. Kept short so we only ever hold back a tiny suffix. +_XML_TOOL_MARKERS = ("", " bool: + """True if ``buffer`` ends with a partial prefix of an XML tool marker. + + Lets the streaming gate hold back a short suffix like ``"``/ + ```` XML on the text channel. In a streaming response those chunks + would otherwise reach the client verbatim before the post-hoc parser runs. + + Feed each content delta through :meth:`feed`; it returns the text that is + safe to emit right now (everything up to a potential tool-call block), + buffering any suspected XML. Call :meth:`flush` at end-of-stream to recover + tool calls from the withheld buffer and get any remaining safe text. + + Usage: + gate = StreamingXmlToolGate() + safe = gate.feed(delta) # emit `safe` if non-empty + ... + leftover, tool_calls = gate.flush() # emit leftover, add tool_calls + """ + + def __init__(self) -> None: + self._buffer = "" # withheld text once a marker is seen + self._holding = False # currently buffering a suspected XML block + self._pending = "" # tiny tail that might be a partial marker + + def feed(self, chunk: str) -> str: + """Consume a content delta, return text safe to emit immediately.""" + if not chunk: + return "" + + if self._holding: + # Already inside a suspected XML block — keep swallowing until the + # block closes. If it closes and no more markers follow, whatever + # trails the close is safe to emit. + self._buffer += chunk + return self._maybe_release_after_close() + + # Not holding yet. Combine any pending partial-marker tail with the new + # chunk and look for a marker. + combined = self._pending + chunk + self._pending = "" + + idx = self._first_marker_index(combined) + if idx is None: + # No full marker. Hold back only a possible partial-marker suffix. + if _tail_could_start_marker(combined): + # Find how much to hold: the longest suffix that is a marker prefix. + hold = self._partial_marker_suffix_len(combined) + self._pending = combined[len(combined) - hold:] + return combined[: len(combined) - hold] + return combined + + # Found a marker — emit everything before it, start holding from there. + safe = combined[:idx] + self._buffer = combined[idx:] + self._holding = True + return safe + self._maybe_release_after_close() + + def _first_marker_index(self, text: str) -> Optional[int]: + positions = [text.find(m) for m in _XML_TOOL_MARKERS] + positions = [p for p in positions if p != -1] + return min(positions) if positions else None + + def _partial_marker_suffix_len(self, text: str) -> int: + window = text[-_XML_MAX_MARKER_LEN:] + best = 0 + for marker in _XML_TOOL_MARKERS: + for plen in range(min(len(marker), len(window)), 0, -1): + if window.endswith(marker[:plen]): + best = max(best, plen) + break + return best + + def _maybe_release_after_close(self) -> str: + """If the buffered block contains a closed tool call with trailing + non-XML text, we still keep holding (more calls may follow). We only + release trailing text at flush(). This keeps the logic simple and + correct: once holding, everything stays buffered until flush.""" + return "" + + def flush(self) -> tuple[str, List[Dict[str, Any]]]: + """End of stream: parse tool calls from the buffer, return leftover text. + + Returns: + (leftover_text, tool_calls) — leftover_text is any buffered content + with XML blocks stripped (normally empty); tool_calls are recovered + from the buffered XML. + """ + buffered = self._buffer + self._pending + self._buffer = "" + self._pending = "" + self._holding = False + if not buffered: + return "", [] + tool_calls = parse_xml_tool_calls(buffered) + leftover = strip_xml_tool_calls(buffered) if tool_calls else buffered + return leftover, tool_calls + + def deduplicate_tool_calls(tool_calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """ Removes duplicate tool calls. @@ -240,6 +489,8 @@ class AwsEventStreamParser: # Patterns for finding JSON events EVENT_PATTERNS = [ ('{"content":', 'content'), + ('{"text":', 'reasoning'), # reasoningContentEvent (native extended thinking) + ('{"signature":', 'reasoning_sig'), # reasoningContentEvent signature (metadata, ignored) ('{"name":', 'tool_start'), ('{"input":', 'tool_input'), ('{"stop":', 'tool_stop'), @@ -254,6 +505,7 @@ def __init__(self): self.last_content: Optional[str] = None # For deduplicating repeating content self.current_tool_call: Optional[Dict[str, Any]] = None self.tool_calls: List[Dict[str, Any]] = [] + self.native_thinking_started = False def feed(self, chunk: bytes) -> List[Dict[str, Any]]: """ @@ -318,6 +570,20 @@ def _process_event(self, data: dict, event_type: str) -> Optional[Dict[str, Any] """ if event_type == 'content': return self._process_content_event(data) + elif event_type == 'reasoning': + return self._process_reasoning_event(data) + elif event_type == 'reasoning_sig': + # Signature is cryptographic metadata for native thinking blocks. + # We don't surface it to clients; just swallow it. + self.native_thinking_started = False + return None + elif event_type == 'thinking': + is_first = not self.native_thinking_started + self.native_thinking_started = True + return {"type": "thinking", "data": data.get('text', ''), "is_first": is_first, "is_native": True} + elif event_type == 'thinking_signature': + self.native_thinking_started = False + return {"type": "thinking_signature", "data": data.get('signature', '')} elif event_type == 'tool_start': return self._process_tool_start_event(data) elif event_type == 'tool_input': @@ -347,6 +613,30 @@ def _process_content_event(self, data: dict) -> Optional[Dict[str, Any]]: return {"type": "content", "data": content} + def _process_reasoning_event(self, data: dict) -> Optional[Dict[str, Any]]: + """ + Processes a native reasoningContentEvent (extended thinking). + + The Kiro backend emits the model's native reasoning as a separate + event channel (reasoningContentEvent) with incremental {"text": ...} + deltas, followed by a {"signature": ...} metadata event. This is the + real model thinking, distinct from kg's legacy fake-reasoning prompt + injection. We surface it as a dedicated 'reasoning' event so the + streaming layer can map it to reasoning_content (OpenAI) or a native + thinking block (Anthropic). + """ + text = data.get('text', '') + if not text: + return None + is_first = not self.native_thinking_started + self.native_thinking_started = True + return { + "type": "reasoning", + "data": text, + "is_first": is_first, + "is_native": True, + } + def _process_tool_start_event(self, data: dict) -> Optional[Dict[str, Any]]: """Processes tool call start.""" # Finalize previous tool call if exists @@ -566,4 +856,5 @@ def reset(self) -> None: self.buffer = "" self.last_content = None self.current_tool_call = None - self.tool_calls = [] \ No newline at end of file + self.tool_calls = [] + self.native_thinking_started = False \ No newline at end of file diff --git a/kiro/profile_arn.py b/kiro/profile_arn.py new file mode 100644 index 00000000..6a2465ea --- /dev/null +++ b/kiro/profile_arn.py @@ -0,0 +1,47 @@ +# -*- coding: utf-8 -*- + +""" +Profile ARN selection for Kiro runtime payloads. + +Since the migration to runtime.kiro.dev (upstream commits 07d24fc, 90d0509), +profileArn is required for ALL auth types — Kiro Desktop, kiro-cli AWS SSO +OIDC, and Enterprise Kiro IDE. The previous gating that stripped profileArn +for plain SSO OIDC requests applied only to the legacy q.amazonaws.com +endpoint and now causes 400 "profileArn is required" errors against runtime. + +Falls back to the PROFILE_ARN environment variable if the auth manager has no +profile ARN of its own. +""" + +from typing import Optional, Protocol + +from kiro.auth import AuthType +from kiro.config import PROFILE_ARN + + +class ProfileArnCarrier(Protocol): + """Auth object fields needed to resolve the profileArn for a payload.""" + + @property + def auth_type(self) -> AuthType: + """Authentication type used by the account.""" + ... + + @property + def profile_arn(self) -> Optional[str]: + """AWS CodeWhisperer profile ARN if available.""" + ... + + +def profile_arn_for_payload(auth_manager: ProfileArnCarrier) -> str: + """ + Return the profileArn that should be sent to runtime.kiro.dev. + + Args: + auth_manager: Auth manager for the selected account. + + Returns: + The auth manager's profile ARN, falling back to the PROFILE_ARN env + variable, or an empty string if neither is available. + """ + return auth_manager.profile_arn or PROFILE_ARN or "" diff --git a/kiro/routes_anthropic.py b/kiro/routes_anthropic.py index dec797a0..e5ae8df6 100644 --- a/kiro/routes_anthropic.py +++ b/kiro/routes_anthropic.py @@ -34,7 +34,8 @@ from fastapi.security import APIKeyHeader from loguru import logger -from kiro.config import PROXY_API_KEY, PROFILE_ARN +from kiro.config import PROXY_API_KEY, PROXY_AUTH_DISABLED +from kiro.profile_arn import profile_arn_for_payload from kiro.models_anthropic import ( AnthropicMessagesRequest, AnthropicCountTokensRequest, @@ -51,7 +52,7 @@ stream_with_first_token_retry_anthropic, ) from kiro.http_client import KiroHttpClient -from kiro.utils import generate_conversation_id +from kiro.utils import generate_conversation_id, derive_conversation_id_from_metadata from kiro.tokenizer import estimate_request_tokens from kiro.config import WEB_SEARCH_ENABLED from kiro.mcp_tools import handle_native_web_search @@ -81,6 +82,9 @@ async def verify_anthropic_api_key( 1. x-api-key header (Anthropic native) 2. Authorization: Bearer header (for compatibility) + When PROXY_AUTH_DISABLED is True, authentication is bypassed entirely + for trusted local networks. + Args: x_api_key: Value from x-api-key header authorization: Value from Authorization header @@ -91,6 +95,9 @@ async def verify_anthropic_api_key( Raises: HTTPException: 401 if key is invalid or missing """ + if PROXY_AUTH_DISABLED: + return True + # Check x-api-key first (Anthropic native) if x_api_key and x_api_key == PROXY_API_KEY: return True @@ -373,17 +380,23 @@ async def messages( model_resolver = account.model_resolver # Generate conversation ID - conversation_id = generate_conversation_id() + # Claude Code Desktop provides metadata.user_id.session_id, which lets us + # keep a stable Kiro conversation ID across turns. Other clients keep the + # previous random UUID behavior. + conversation_id = ( + derive_conversation_id_from_metadata(request_data.metadata) + or generate_conversation_id() + ) # Build payload for Kiro # profileArn is required by runtime.kiro.dev for all auth types - profile_arn_for_payload = auth_manager.profile_arn or PROFILE_ARN or "" + profile_arn = profile_arn_for_payload(auth_manager) try: kiro_payload = anthropic_to_kiro( request_data, conversation_id, - profile_arn_for_payload + profile_arn ) except ValueError as e: logger.error(f"Conversion error: {e}") @@ -413,8 +426,7 @@ async def messages( if request_data.stream: http_client = KiroHttpClient(auth_manager, shared_client=None) else: - shared_client = request.app.state.http_client - http_client = KiroHttpClient(auth_manager, shared_client=shared_client) + http_client = KiroHttpClient(auth_manager, shared_client=None) # Prepare data for token counting messages_for_tokenizer = [msg.model_dump() for msg in request_data.messages] @@ -554,12 +566,20 @@ async def make_retry_request(): if debug_logger: debug_logger.flush_on_error(response.status_code, last_error_message) + # Map context overflow to the Anthropic error type Claude Code + # recognises, so /compact and auto-retry can recover natively + # instead of treating it as a generic API failure. + anthropic_error_type = ( + "invalid_request_error" + if error_reason == "CONTENT_LENGTH_EXCEEDS_THRESHOLD" + else "api_error" + ) return JSONResponse( status_code=response.status_code, content={ "type": "error", "error": { - "type": "api_error", + "type": anthropic_error_type, "message": last_error_message } } @@ -680,8 +700,14 @@ async def make_retry_request(): # Normal Flow (Path B will be intercepted in streaming, or no web_search) # ============================================================================== - # Generate conversation ID for Kiro API (random UUID, not used for tracking) - conversation_id = generate_conversation_id() + # Generate conversation ID for Kiro API. + # Claude Code Desktop provides metadata.user_id.session_id, which lets us + # keep a stable Kiro conversation ID across turns. Other clients keep the + # previous random UUID behavior. + conversation_id = ( + derive_conversation_id_from_metadata(request_data.metadata) + or generate_conversation_id() + ) # Build payload for Kiro # profileArn is required by runtime.kiro.dev for all auth types @@ -725,9 +751,7 @@ async def make_retry_request(): # when network interface changes (VPN disconnect/reconnect) http_client = KiroHttpClient(auth_manager, shared_client=None) else: - # Non-streaming mode: shared client for efficient connection reuse - shared_client = request.app.state.http_client - http_client = KiroHttpClient(auth_manager, shared_client=shared_client) + http_client = KiroHttpClient(auth_manager, shared_client=None) # Prepare data for token counting # Convert Pydantic models to dicts for tokenizer @@ -761,12 +785,14 @@ async def make_retry_request(): # Try to parse JSON response from Kiro to extract error message error_message = error_text + error_reason = None try: error_json = json.loads(error_text) # Enhance Kiro API errors with user-friendly messages from kiro.kiro_errors import enhance_kiro_error error_info = enhance_kiro_error(error_json) error_message = error_info.user_message + error_reason = error_info.reason # Log original error for debugging logger.debug(f"Original Kiro error: {error_info.original_message} (reason: {error_info.reason})") except (json.JSONDecodeError, KeyError): @@ -781,13 +807,20 @@ async def make_retry_request(): if debug_logger: debug_logger.flush_on_error(response.status_code, error_message) + # Map context overflow to the Anthropic error type Claude Code recognises, + # so /compact and auto-retry can recover natively. + anthropic_error_type = ( + "invalid_request_error" + if error_reason == "CONTENT_LENGTH_EXCEEDS_THRESHOLD" + else "api_error" + ) # Return error in Anthropic format return JSONResponse( status_code=response.status_code, content={ "type": "error", "error": { - "type": "api_error", + "type": anthropic_error_type, "message": error_message } } diff --git a/kiro/routes_openai.py b/kiro/routes_openai.py index 262cb0e1..a539308b 100644 --- a/kiro/routes_openai.py +++ b/kiro/routes_openai.py @@ -29,6 +29,8 @@ import json from datetime import datetime, timezone +from typing import Optional + from fastapi import APIRouter, Depends, HTTPException, Request, Response, Security from fastapi.responses import JSONResponse, StreamingResponse from fastapi.security import APIKeyHeader @@ -37,8 +39,9 @@ from kiro.config import ( PROXY_API_KEY, APP_VERSION, - PROFILE_ARN, + PROXY_AUTH_DISABLED, ) +from kiro.profile_arn import profile_arn_for_payload from kiro.models_openai import ( OpenAIModel, ModelList, @@ -50,7 +53,7 @@ from kiro.converters_openai import build_kiro_payload from kiro.streaming_openai import stream_kiro_to_openai, collect_stream_response, stream_with_first_token_retry from kiro.http_client import KiroHttpClient -from kiro.utils import generate_conversation_id +from kiro.utils import generate_conversation_id, get_kiro_headers from kiro.config import WEB_SEARCH_ENABLED from kiro.mcp_tools import handle_native_web_search @@ -63,16 +66,25 @@ # --- Security scheme --- api_key_header = APIKeyHeader(name="Authorization", auto_error=False) +x_api_key_header = APIKeyHeader(name="x-api-key", auto_error=False) -async def verify_api_key(auth_header: str = Security(api_key_header)) -> bool: +async def verify_api_key( + auth_header: Optional[str] = Security(api_key_header), + x_api_key: Optional[str] = Security(x_api_key_header), +) -> bool: """ - Verify API key in Authorization header. + Verify API key in Authorization or x-api-key header. + + Supports OpenAI-style "Authorization: Bearer {PROXY_API_KEY}" and + Anthropic-style "x-api-key: {PROXY_API_KEY}" for model discovery clients. - Expects format: "Bearer {PROXY_API_KEY}" + When PROXY_AUTH_DISABLED is True, authentication is bypassed entirely + for trusted local networks. Args: auth_header: Authorization header value + x_api_key: x-api-key header value Returns: True if key is valid @@ -80,10 +92,17 @@ async def verify_api_key(auth_header: str = Security(api_key_header)) -> bool: Raises: HTTPException: 401 if key is invalid or missing """ - if not auth_header or auth_header != f"Bearer {PROXY_API_KEY}": - logger.warning("Access attempt with invalid API key.") - raise HTTPException(status_code=401, detail="Invalid or missing API Key") - return True + if PROXY_AUTH_DISABLED: + return True + + if auth_header and auth_header == f"Bearer {PROXY_API_KEY}": + return True + + if x_api_key and x_api_key == PROXY_API_KEY: + return True + + logger.warning("Access attempt with invalid API key.") + raise HTTPException(status_code=401, detail="Invalid or missing API Key") # --- Router --- @@ -119,13 +138,50 @@ async def health(): "version": APP_VERSION } + +@router.get("/usage", dependencies=[Depends(verify_api_key)]) +async def get_usage(request: Request): + """ + Query Kiro credit usage via GetUsageLimits API. + + Returns subscription info, credit usage breakdown, overage config, + and reset date — the same data shown by kiro-cli's /usage command. + """ + logger.info("Request to /v1/usage") + + account = request.app.state.account_manager.get_first_account() + auth_manager: KiroAuthManager = account.auth_manager + shared_client = request.app.state.http_client + + token = await auth_manager.get_access_token() + headers = get_kiro_headers(auth_manager, token) + headers["x-amz-target"] = "com.amazon.aws.codewhisperer.runtime.AmazonCodeWhispererService.GetUsageLimits" + headers["Content-Type"] = "application/json" + + url = auth_manager.management_host + body = {"origin": "AI_EDITOR"} + if auth_manager.profile_arn: + body["profileArn"] = auth_manager.profile_arn + + try: + response = await shared_client.post(url, json=body, headers=headers) + if response.status_code != 200: + raise HTTPException(status_code=response.status_code, detail=response.text) + return JSONResponse(content=response.json()) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error fetching usage: {e}", exc_info=True) + raise HTTPException(status_code=502, detail=f"Failed to fetch usage: {str(e)}") + @router.get("/v1/models", response_model=ModelList, dependencies=[Depends(verify_api_key)]) async def get_models(request: Request): """ Return list of available models. - Models are loaded at startup (blocking) and cached. - This endpoint returns the cached list. + On each request, the model cache is refreshed for all initialized accounts + so that additions and removals made by Kiro are visible immediately. + A failed refresh preserves the last successfully fetched cache. Args: request: FastAPI Request for accessing app.state @@ -135,6 +191,12 @@ async def get_models(request: Request): """ logger.info("Request to /v1/models") + # Refresh model cache before returning + try: + await request.app.state.account_manager.refresh_initialized_account_models() + except Exception: + logger.warning("Model refresh failed, using cached models") + # Get available models based on mode if request.app.state.account_system: # Account system: collect models from all initialized accounts @@ -324,13 +386,13 @@ async def chat_completions(request: Request, request_data: ChatCompletionRequest # Build payload for Kiro # profileArn is required by runtime.kiro.dev for all auth types - profile_arn_for_payload = auth_manager.profile_arn or PROFILE_ARN or "" + profile_arn = profile_arn_for_payload(auth_manager) try: kiro_payload = build_kiro_payload( request_data, conversation_id, - profile_arn_for_payload + profile_arn ) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) @@ -350,8 +412,7 @@ async def chat_completions(request: Request, request_data: ChatCompletionRequest if request_data.stream: http_client = KiroHttpClient(auth_manager, shared_client=None) else: - shared_client = request.app.state.http_client - http_client = KiroHttpClient(auth_manager, shared_client=shared_client) + http_client = KiroHttpClient(auth_manager, shared_client=None) try: # Make request to Kiro API @@ -479,12 +540,20 @@ async def make_retry_request(): if debug_logger: debug_logger.flush_on_error(response.status_code, last_error_message) + # Map context overflow to the standard OpenAI error type so + # clients (Claude Code / claude-code-router) trigger context + # recovery instead of treating it as a generic API failure. + openai_error_type = ( + "context_length_exceeded" + if error_reason == "CONTENT_LENGTH_EXCEEDS_THRESHOLD" + else "kiro_api_error" + ) return JSONResponse( status_code=response.status_code, content={ "error": { "message": last_error_message, - "type": "kiro_api_error", + "type": openai_error_type, "code": response.status_code } } @@ -572,13 +641,13 @@ async def make_retry_request(): # Build payload for Kiro # profileArn is required by runtime.kiro.dev for all auth types - profile_arn_for_payload = auth_manager.profile_arn or PROFILE_ARN or "" + profile_arn = profile_arn_for_payload(auth_manager) try: kiro_payload = build_kiro_payload( request_data, conversation_id, - profile_arn_for_payload + profile_arn ) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) @@ -602,9 +671,7 @@ async def make_retry_request(): # when network interface changes (VPN disconnect/reconnect) http_client = KiroHttpClient(auth_manager, shared_client=None) else: - # Non-streaming mode: shared client for efficient connection reuse - shared_client = request.app.state.http_client - http_client = KiroHttpClient(auth_manager, shared_client=shared_client) + http_client = KiroHttpClient(auth_manager, shared_client=None) try: # Make request to Kiro API (for both streaming and non-streaming modes) # Important: we wait for Kiro response BEFORE returning StreamingResponse, @@ -627,12 +694,14 @@ async def make_retry_request(): # Try to parse JSON response from Kiro to extract error message error_message = error_text + error_reason = None try: error_json = json.loads(error_text) # Enhance Kiro API errors with user-friendly messages from kiro.kiro_errors import enhance_kiro_error error_info = enhance_kiro_error(error_json) error_message = error_info.user_message + error_reason = error_info.reason # Log original error for debugging logger.debug(f"Original Kiro error: {error_info.original_message} (reason: {error_info.reason})") except (json.JSONDecodeError, KeyError): @@ -647,13 +716,20 @@ async def make_retry_request(): if debug_logger: debug_logger.flush_on_error(response.status_code, error_message) + # Map context overflow to the standard OpenAI error type so clients + # trigger context recovery instead of treating it as a generic failure. + openai_error_type = ( + "context_length_exceeded" + if error_reason == "CONTENT_LENGTH_EXCEEDS_THRESHOLD" + else "kiro_api_error" + ) # Return error in OpenAI API format return JSONResponse( status_code=response.status_code, content={ "error": { "message": error_message, - "type": "kiro_api_error", + "type": openai_error_type, "code": response.status_code } } diff --git a/kiro/routes_responses.py b/kiro/routes_responses.py new file mode 100644 index 00000000..f9055ee8 --- /dev/null +++ b/kiro/routes_responses.py @@ -0,0 +1,408 @@ +# -*- 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 . + +""" +FastAPI route for the OpenAI Responses API (/v1/responses). + +This is the endpoint used by OpenAI Codex CLI. It mirrors the account-system +failover / legacy-mode logic of routes_openai.py but produces Responses API +output (a `response` object for non-streaming, or the Responses SSE event +sequence for streaming). +""" + +import json + +from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi.responses import JSONResponse, StreamingResponse +from loguru import logger + +from kiro.profile_arn import profile_arn_for_payload +from kiro.models_responses import ResponsesRequest +from kiro.converters_responses import ( + build_kiro_payload, + convert_responses_input_to_unified, + convert_responses_tools_to_unified, +) +from kiro.converters_core import extract_text_content +from kiro.streaming_responses import ( + stream_responses_with_first_token_retry, + collect_responses_response, +) +from kiro.http_client import KiroHttpClient +from kiro.utils import generate_conversation_id + +# Reuse the OpenAI endpoint's API-key verification (same Bearer scheme) +from kiro.routes_openai import verify_api_key + +try: + from kiro.debug_logger import debug_logger +except ImportError: + debug_logger = None + + +router = APIRouter() + + +def _build_tokenizer_inputs(request_data: ResponsesRequest): + """ + Build simple message/tool dicts for fallback (tiktoken) token counting. + + Kiro's context_usage percentage is the primary token source; this is only + used when Kiro returns no usage data. + """ + _, unified_messages = convert_responses_input_to_unified(request_data) + messages_for_tokenizer = [ + {"role": m.role, "content": extract_text_content(m.content)} + for m in unified_messages + ] + + tools_for_tokenizer = None + unified_tools = convert_responses_tools_to_unified(request_data.tools) + if unified_tools: + tools_for_tokenizer = [ + {"name": t.name, "description": t.description or "", "parameters": t.input_schema or {}} + for t in unified_tools + ] + + return messages_for_tokenizer, tools_for_tokenizer + + +def _error_response(status_code: int, message: str) -> JSONResponse: + """Build a Responses-style error JSON response.""" + return JSONResponse( + status_code=status_code, + content={"error": {"message": message, "type": "kiro_api_error", "code": status_code}}, + ) + + +@router.post("/v1/responses", dependencies=[Depends(verify_api_key)]) +async def responses(request: Request, request_data: ResponsesRequest): + """ + Responses API endpoint - compatible with OpenAI Codex CLI. + + Accepts requests in Responses API format, translates them to Kiro, and + returns either a `response` object (non-streaming) or the Responses SSE + event sequence (streaming). + """ + logger.info(f"Request to /v1/responses (model={request_data.model}, stream={request_data.stream})") + + messages_for_tokenizer, tools_for_tokenizer = _build_tokenizer_inputs(request_data) + + # ============================================================================== + # ACCOUNT SYSTEM ENABLED: Failover Loop + # ============================================================================== + if request.app.state.account_system: + from kiro.account_errors import classify_error, ErrorType + + account_manager = request.app.state.account_manager + all_accounts = list(account_manager._accounts.keys()) + MAX_ATTEMPTS = len(all_accounts) * 2 + + last_error_message = None + last_error_status = None + tried_accounts = set() + + for attempt in range(MAX_ATTEMPTS): + account = await account_manager.get_next_account( + request_data.model, exclude_accounts=tried_accounts + ) + + if account is None: + if len(all_accounts) == 1: + raise HTTPException( + status_code=last_error_status or 503, + detail=last_error_message or "Account unavailable", + ) + detail = "No available accounts for this model." + if last_error_message: + detail += f" Error from last account: {last_error_message}" + raise HTTPException(status_code=503, detail=detail) + + tried_accounts.add(account.id) + + auth_manager = account.auth_manager + model_cache = account.model_cache + + conversation_id = generate_conversation_id() + profile_arn = profile_arn_for_payload(auth_manager) + + try: + kiro_payload = build_kiro_payload( + request_data, conversation_id, profile_arn + ) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + + try: + kiro_request_body = json.dumps(kiro_payload, ensure_ascii=False, indent=2).encode("utf-8") + if debug_logger: + debug_logger.log_kiro_request_body(kiro_request_body) + except Exception as e: + logger.warning(f"Failed to log Kiro request: {e}") + + url = f"{auth_manager.api_host}/generateAssistantResponse" + logger.debug(f"Kiro API URL: {url} (account: {account.id})") + + if request_data.stream: + http_client = KiroHttpClient(auth_manager, shared_client=None) + else: + shared_client = request.app.state.http_client + http_client = KiroHttpClient(auth_manager, shared_client=shared_client) + + try: + response = await http_client.request_with_retry("POST", url, kiro_payload, stream=True) + + if response.status_code == 200: + await account_manager.report_success(account.id, request_data.model) + + if request_data.stream: + return _make_streaming_response( + http_client, url, kiro_payload, request_data, + model_cache, auth_manager, response, + messages_for_tokenizer, tools_for_tokenizer, + ) + + responses_obj = await collect_responses_response( + http_client.client, response, request_data.model, + model_cache, auth_manager, + request_messages=messages_for_tokenizer, + request_tools=tools_for_tokenizer, + ) + await http_client.close() + logger.info("HTTP 200 - POST /v1/responses (non-streaming) - completed") + if debug_logger: + debug_logger.discard_buffers() + return JSONResponse(content=responses_obj) + + # --- error path --- + try: + error_content = await response.aread() + except Exception: + error_content = b"Unknown error" + await http_client.close() + error_text = error_content.decode("utf-8", errors="replace") + + error_reason = None + try: + error_json = json.loads(error_text) + from kiro.kiro_errors import enhance_kiro_error + error_info = enhance_kiro_error(error_json) + error_reason = error_info.reason + last_error_message = error_info.user_message + last_error_status = response.status_code + except (json.JSONDecodeError, KeyError): + last_error_message = error_text + last_error_status = response.status_code + + error_type = classify_error(response.status_code, error_reason) + + if error_type == ErrorType.FATAL: + await account_manager.report_failure( + account.id, request_data.model, error_type, + response.status_code, error_reason, + ) + logger.warning(f"HTTP {response.status_code} - POST /v1/responses - {last_error_message[:100]}") + if debug_logger: + debug_logger.flush_on_error(response.status_code, last_error_message) + return _error_response(response.status_code, last_error_message) + + # RECOVERABLE - try next account + await account_manager.report_failure( + account.id, request_data.model, error_type, + response.status_code, error_reason, + ) + if len(all_accounts) == 1: + break + continue + + except HTTPException as e: + await http_client.close() + if e.status_code in (502, 504): + await account_manager.report_failure( + account.id, request_data.model, ErrorType.RECOVERABLE, e.status_code, None + ) + last_error_message = str(e.detail) + last_error_status = e.status_code + if len(all_accounts) == 1: + break + logger.warning(f"Network error on account {account.id}, trying next account") + continue + logger.error(f"HTTP {e.status_code} - POST /v1/responses - {e.detail}") + if debug_logger: + debug_logger.flush_on_error(e.status_code, str(e.detail)) + raise + except Exception as e: + await http_client.close() + logger.error(f"Internal error: {e}", exc_info=True) + logger.error(f"HTTP 500 - POST /v1/responses - {str(e)[:100]}") + if debug_logger: + debug_logger.flush_on_error(500, str(e)) + raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}") + + # All attempts exhausted + if len(all_accounts) == 1: + raise HTTPException(status_code=last_error_status, detail=last_error_message) + detail = "All accounts failed after full circle." + if last_error_message: + detail += f" Error from last account: {last_error_message}" + raise HTTPException(status_code=503, detail=detail) + + # ============================================================================== + # LEGACY MODE: Single Account (no failover) + # ============================================================================== + account = request.app.state.account_manager.get_first_account() + if not account.auth_manager: + logger.error("No initialized accounts available (legacy mode)") + raise HTTPException(503, "No initialized accounts available") + auth_manager = account.auth_manager + model_cache = account.model_cache + + conversation_id = generate_conversation_id() + profile_arn = profile_arn_for_payload(auth_manager) + + try: + kiro_payload = build_kiro_payload(request_data, conversation_id, profile_arn) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + + try: + kiro_request_body = json.dumps(kiro_payload, ensure_ascii=False, indent=2).encode("utf-8") + if debug_logger: + debug_logger.log_kiro_request_body(kiro_request_body) + except Exception as e: + logger.warning(f"Failed to log Kiro request: {e}") + + url = f"{auth_manager.api_host}/generateAssistantResponse" + logger.debug(f"Kiro API URL: {url}") + + if request_data.stream: + http_client = KiroHttpClient(auth_manager, shared_client=None) + else: + shared_client = request.app.state.http_client + http_client = KiroHttpClient(auth_manager, shared_client=shared_client) + + try: + response = await http_client.request_with_retry("POST", url, kiro_payload, stream=True) + + if response.status_code != 200: + try: + error_content = await response.aread() + except Exception: + error_content = b"Unknown error" + await http_client.close() + error_text = error_content.decode("utf-8", errors="replace") + + error_message = error_text + try: + error_json = json.loads(error_text) + from kiro.kiro_errors import enhance_kiro_error + error_info = enhance_kiro_error(error_json) + error_message = error_info.user_message + except (json.JSONDecodeError, KeyError): + pass + + logger.warning(f"HTTP {response.status_code} - POST /v1/responses - {error_message[:100]}") + if debug_logger: + debug_logger.flush_on_error(response.status_code, error_message) + return _error_response(response.status_code, error_message) + + if request_data.stream: + return _make_streaming_response( + http_client, url, kiro_payload, request_data, + model_cache, auth_manager, response, + messages_for_tokenizer, tools_for_tokenizer, + ) + + responses_obj = await collect_responses_response( + http_client.client, response, request_data.model, + model_cache, auth_manager, + request_messages=messages_for_tokenizer, + request_tools=tools_for_tokenizer, + ) + await http_client.close() + logger.info("HTTP 200 - POST /v1/responses (non-streaming) - completed") + if debug_logger: + debug_logger.discard_buffers() + return JSONResponse(content=responses_obj) + + except HTTPException as e: + await http_client.close() + if e.status_code in (502, 504): + logger.warning("Network error (legacy mode, no failover available)") + logger.error(f"HTTP {e.status_code} - POST /v1/responses - {e.detail}") + if debug_logger: + debug_logger.flush_on_error(e.status_code, str(e.detail)) + raise + except Exception as e: + await http_client.close() + logger.error(f"Internal error: {e}", exc_info=True) + logger.error(f"HTTP 500 - POST /v1/responses - {str(e)[:100]}") + if debug_logger: + debug_logger.flush_on_error(500, str(e)) + raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}") + + +def _make_streaming_response( + http_client, url, kiro_payload, request_data, + model_cache, auth_manager, initial_response, + messages_for_tokenizer, tools_for_tokenizer, +) -> StreamingResponse: + """Build the StreamingResponse wrapper for Responses SSE output.""" + + async def stream_wrapper(): + streaming_error = None + client_disconnected = False + try: + async def make_retry_request(): + return await http_client.request_with_retry("POST", url, kiro_payload, stream=True) + + async for chunk in stream_responses_with_first_token_retry( + make_request=make_retry_request, + client=http_client.client, + model=request_data.model, + model_cache=model_cache, + auth_manager=auth_manager, + initial_response=initial_response, + request_messages=messages_for_tokenizer, + request_tools=tools_for_tokenizer, + ): + yield chunk + except GeneratorExit: + client_disconnected = True + logger.debug("Client disconnected during streaming (GeneratorExit in routes)") + except Exception as e: + streaming_error = e + raise + finally: + await http_client.close() + if streaming_error: + error_type = type(streaming_error).__name__ + error_msg = str(streaming_error) if str(streaming_error) else "(empty message)" + logger.error(f"HTTP 500 - POST /v1/responses (streaming) - [{error_type}] {error_msg[:100]}") + elif client_disconnected: + logger.info("HTTP 200 - POST /v1/responses (streaming) - client disconnected") + else: + logger.info("HTTP 200 - POST /v1/responses (streaming) - completed") + if debug_logger: + if streaming_error: + debug_logger.flush_on_error(500, str(streaming_error)) + else: + debug_logger.discard_buffers() + + return StreamingResponse(stream_wrapper(), media_type="text/event-stream") diff --git a/kiro/streaming_anthropic.py b/kiro/streaming_anthropic.py index 4fee71e4..403eca1a 100644 --- a/kiro/streaming_anthropic.py +++ b/kiro/streaming_anthropic.py @@ -50,6 +50,7 @@ from kiro.tokenizer import count_tokens, estimate_request_tokens from kiro.parsers import parse_bracket_tool_calls, deduplicate_tool_calls from kiro.config import FIRST_TOKEN_TIMEOUT, FIRST_TOKEN_MAX_RETRIES, FAKE_REASONING_HANDLING +from kiro.cache import prompt_cache_tracker if TYPE_CHECKING: from kiro.auth import KiroAuthManager @@ -634,6 +635,16 @@ async def stream_kiro_to_anthropic( if prompt_source != "unknown": input_tokens = prompt_tokens + local_cache_usage = ( + await prompt_cache_tracker.record( + model=model, + messages=request_messages, + tools=request_tools, + system=request_system, + input_tokens=input_tokens, + ) + ).to_anthropic_usage_fields() + # Determine stop reason (truncation has highest priority) if content_was_truncated: stop_reason = "max_tokens" @@ -647,6 +658,8 @@ async def stream_kiro_to_anthropic( "output_tokens": output_tokens } usage_payload.update(upstream_cache_usage) + if not upstream_cache_usage: + usage_payload.update(local_cache_usage) yield format_sse_event("message_delta", { "type": "message_delta", @@ -815,6 +828,16 @@ async def collect_anthropic_response( if prompt_source != "unknown": input_tokens = prompt_tokens + local_cache_usage = ( + await prompt_cache_tracker.record( + model=model, + messages=request_messages, + tools=request_tools, + system=request_system, + input_tokens=input_tokens, + ) + ).to_anthropic_usage_fields() + # Detect content truncation (missing completion signals) stream_completed_normally = result.context_usage_percentage is not None content_was_truncated = ( @@ -850,6 +873,8 @@ async def collect_anthropic_response( "output_tokens": output_tokens } usage_payload.update(upstream_cache_usage) + if not upstream_cache_usage: + usage_payload.update(local_cache_usage) return { "id": message_id, diff --git a/kiro/streaming_core.py b/kiro/streaming_core.py index 4294ffa9..51960253 100644 --- a/kiro/streaming_core.py +++ b/kiro/streaming_core.py @@ -37,12 +37,13 @@ import httpx from loguru import logger -from kiro.parsers import AwsEventStreamParser, parse_bracket_tool_calls, deduplicate_tool_calls +from kiro.parsers import AwsEventStreamParser, parse_bracket_tool_calls, deduplicate_tool_calls, parse_xml_tool_calls, strip_xml_tool_calls from kiro.config import ( FIRST_TOKEN_TIMEOUT, FIRST_TOKEN_MAX_RETRIES, FAKE_REASONING_ENABLED, FAKE_REASONING_HANDLING, + NATIVE_REASONING_ENABLED, ) from kiro.thinking_parser import ThinkingParser @@ -85,6 +86,7 @@ class KiroEvent: context_usage_percentage: Optional[float] = None is_first_thinking_chunk: bool = False is_last_thinking_chunk: bool = False + is_native_thinking: bool = False @dataclass @@ -139,12 +141,19 @@ async def parse_kiro_stream( """ parser = AwsEventStreamParser() first_token_received = False + byte_iterator = None - # Initialize thinking parser if fake reasoning is enabled + # Initialize thinking parser only for legacy fake-reasoning mode. + # When native reasoning is active, the model emits a dedicated + # reasoningContentEvent stream (no thinking tags in content), so the + # tag-detecting ThinkingParser is unnecessary and would only add buffering + # latency to the content channel. thinking_parser: Optional[ThinkingParser] = None - if FAKE_REASONING_ENABLED and enable_thinking_parser: + if FAKE_REASONING_ENABLED and enable_thinking_parser and not NATIVE_REASONING_ENABLED: thinking_parser = ThinkingParser(handling_mode=FAKE_REASONING_HANDLING) logger.debug(f"Thinking parser initialized with mode: {FAKE_REASONING_HANDLING}") + elif NATIVE_REASONING_ENABLED: + logger.debug("Native reasoning enabled - thinking parser skipped (using reasoningContentEvent)") try: # Create iterator for reading bytes @@ -165,6 +174,16 @@ async def parse_kiro_stream( # Empty response - this is normal, just finish logger.debug("Empty response from Kiro API") return + except (httpx.RemoteProtocolError, httpx.ReadError) as e: + # Upstream dropped the connection before sending any token. Nothing has + # been emitted to the client yet, so it is safe to retry. Reuse the + # first-token retry path by raising FirstTokenTimeoutError. + logger.warning( + "Upstream closed connection before first token [{}: {}] - triggering retry", + type(e).__name__, + str(e) if str(e) else "(empty message)", + ) + raise FirstTokenTimeoutError("Upstream closed connection before first token") # Process first chunk if debug_logger: @@ -176,12 +195,24 @@ async def parse_kiro_stream( yield event # Continue reading remaining chunks - async for chunk in byte_iterator: - if debug_logger: - debug_logger.log_raw_chunk(chunk) - - async for event in _process_chunk(parser, chunk, thinking_parser): - yield event + try: + async for chunk in byte_iterator: + if debug_logger: + debug_logger.log_raw_chunk(chunk) + + async for event in _process_chunk(parser, chunk, thinking_parser): + yield event + except (httpx.RemoteProtocolError, httpx.ReadError) as e: + # Upstream Kiro closed the connection before completing the chunked + # body (incomplete chunked read). Content was already streamed, so we + # cannot retry without duplicating output to the client. End the stream + # gracefully and fall through to the finalization below so the client + # keeps the partial content + a proper finish marker instead of HTTP 500. + logger.warning( + "Upstream connection dropped mid-stream [{}: {}] - finalizing with partial content", + type(e).__name__, + str(e) if str(e) else "(empty message)", + ) # Finalize thinking parser and yield any remaining content if thinking_parser: @@ -225,6 +256,12 @@ async def parse_kiro_stream( error_msg = str(e) if str(e) else "(empty message)" logger.error(f"Error during stream parsing: [{error_type}] {error_msg}", exc_info=True) raise + finally: + if byte_iterator is not None: + try: + await byte_iterator.aclose() + except Exception: + pass async def _process_chunk( @@ -281,6 +318,27 @@ async def _process_chunk( elif event["type"] == "context_usage": yield KiroEvent(type="context_usage", context_usage_percentage=event["data"]) + elif event["type"] == "reasoning": + # Native extended thinking from the model (reasoningContentEvent). + # This is NOT wrapped in tags, so it bypasses the ThinkingParser + # entirely and is emitted directly as a thinking event. + if event["data"]: + yield KiroEvent( + type="thinking", + thinking_content=event["data"], + is_first_thinking_chunk=event.get("is_first", False), + is_native_thinking=event.get("is_native", False), + ) + + elif event["type"] == "thinking_signature": + logger.debug("Received native thinking signature from Kiro stream") + yield KiroEvent( + type="thinking", + thinking_content="", + is_last_thinking_chunk=True, + is_native_thinking=True, + ) + # ================================================================================================== # Full Response Collection @@ -326,7 +384,14 @@ async def collect_stream_to_result( bracket_tool_calls = parse_bracket_tool_calls(full_content_for_bracket_tools) if bracket_tool_calls: result.tool_calls = deduplicate_tool_calls(result.tool_calls + bracket_tool_calls) - + + # Recover XML tool calls that leaked into the text channel and strip + # the raw XML from the visible content so it doesn't render as text. + xml_tool_calls = parse_xml_tool_calls(result.content) + if xml_tool_calls: + result.tool_calls = deduplicate_tool_calls(result.tool_calls + xml_tool_calls) + result.content = strip_xml_tool_calls(result.content) + return result @@ -431,7 +496,35 @@ async def stream_with_first_token_retry( response = initial_response logger.debug("Reusing initial response for first attempt") else: - response = await make_request() + # Make request with keep-alive chunks during prefill (TTFT) to + # prevent client disconnect on large-context requests where the + # upstream withholds HTTP 200 headers until the first token is + # ready. Without this, the route stays byte-silent for the entire + # prefill window, and the client treats the silence as a dead stream. + async def _run_make_request() -> httpx.Response: + return await make_request() + + req_task = asyncio.create_task(_run_make_request()) + _ka_count = 0 + try: + while not req_task.done(): + try: + await asyncio.wait_for(asyncio.shield(req_task), timeout=3.0) + except asyncio.TimeoutError: + _ka_count += 1 + yield ": keepalive\n\n" + if _ka_count: + logger.debug(f"Sent {_ka_count} keep-alive chunk(s) during prefill ({_ka_count * 3}s+ TTFT)") + response = req_task.result() + finally: + if not req_task.done(): + req_task.cancel() + try: + await req_task + except BaseException: + pass + + assert response is not None if response.status_code != 200: # Error from API - close response and raise exception diff --git a/kiro/streaming_openai.py b/kiro/streaming_openai.py index b6179540..f4ae4709 100644 --- a/kiro/streaming_openai.py +++ b/kiro/streaming_openai.py @@ -36,7 +36,7 @@ from fastapi import HTTPException from loguru import logger -from kiro.parsers import parse_bracket_tool_calls, deduplicate_tool_calls +from kiro.parsers import parse_bracket_tool_calls, deduplicate_tool_calls, StreamingXmlToolGate from kiro.utils import generate_completion_id from kiro.config import ( FIRST_TOKEN_TIMEOUT, @@ -125,6 +125,7 @@ async def stream_kiro_to_openai_internal( streaming_error_occurred = False tool_calls_from_stream = [] + xml_gate = StreamingXmlToolGate() # withholds Anthropic XML tool-call blocks from streamed content try: # Use streaming_core.parse_kiro_stream for unified event parsing @@ -133,9 +134,17 @@ async def stream_kiro_to_openai_internal( if event.type == "content" and event.content: # Accumulate content for bracket tool call detection full_content += event.content - + + # Route through the XML gate: it withholds any Anthropic + # function_calls/invoke tool-call XML so it never streams to + # the client as text. Only the safe portion is emitted here; + # recovered tool calls are merged at end-of-stream via flush(). + safe_content = xml_gate.feed(event.content) + if not safe_content: + continue + # Format as OpenAI chunk - delta = {"content": event.content} + delta = {"content": safe_content} if first_chunk: delta["role"] = "assistant" first_chunk = False @@ -155,13 +164,20 @@ async def stream_kiro_to_openai_internal( yield chunk_text - elif event.type == "thinking" and event.thinking_content: + elif event.type == "thinking" and (event.thinking_content or event.is_last_thinking_chunk): # Accumulate thinking content - full_thinking_content += event.thinking_content + full_thinking_content += event.thinking_content or "" # Send as reasoning_content or content based on mode if FAKE_REASONING_HANDLING == "as_reasoning_content": delta = {"reasoning_content": event.thinking_content} + elif event.is_native_thinking and FAKE_REASONING_HANDLING == "pass": + thinking_text = event.thinking_content or "" + if event.is_first_thinking_chunk: + thinking_text = f"{thinking_text}" + if event.is_last_thinking_chunk: + thinking_text = f"{thinking_text}" + delta = {"content": thinking_text} else: delta = {"content": event.thinking_content} @@ -267,7 +283,24 @@ async def stream_kiro_to_openai_internal( elif event.type == "context_usage" and event.context_usage_percentage is not None: context_usage_percentage = event.context_usage_percentage - + + # Flush the XML gate: recover any withheld Anthropic XML tool calls and + # emit whatever safe (non-XML) text remained buffered. + xml_leftover, xml_tool_calls = xml_gate.flush() + if xml_leftover: + delta = {"content": xml_leftover} + if first_chunk: + delta["role"] = "assistant" + first_chunk = False + leftover_chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created_time, + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": None}] + } + yield f"data: {json.dumps(leftover_chunk, ensure_ascii=False)}\n\n" + # Track completion signals for truncation detection received_usage = metering_data is not None received_context_usage = context_usage_percentage is not None @@ -275,7 +308,7 @@ async def stream_kiro_to_openai_internal( # Check bracket-style tool calls in full content bracket_tool_calls = parse_bracket_tool_calls(full_content) - all_tool_calls = tool_calls_from_stream + bracket_tool_calls + all_tool_calls = tool_calls_from_stream + bracket_tool_calls + xml_tool_calls all_tool_calls = deduplicate_tool_calls(all_tool_calls) # Detect content truncation (missing completion signals) diff --git a/kiro/streaming_responses.py b/kiro/streaming_responses.py new file mode 100644 index 00000000..0c65756c --- /dev/null +++ b/kiro/streaming_responses.py @@ -0,0 +1,553 @@ +# -*- 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 . + +""" +Streaming logic for converting Kiro stream to OpenAI Responses API format. + +The Responses API SSE protocol is more verbose than Chat Completions: each +logical output (reasoning / message text / function call) is wrapped in an +"item" lifecycle (added -> content parts / deltas -> done) and the stream is +bookended by `response.created` and `response.completed`. + +This module implements a small state machine over the unified KiroEvent stream: +- thinking events -> a `reasoning` item with summary text deltas +- content events -> a `message` item with output_text deltas +- tool_use events -> one `function_call` item each (emitted at stream end) + +IMPORTANT (retry safety): `response.created` is emitted lazily, only after the +first KiroEvent arrives. parse_kiro_stream raises FirstTokenTimeoutError BEFORE +yielding anything, so on a first-token retry no events have been sent yet and +the retried attempt starts cleanly. +""" + +import json +import time +from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Awaitable, Dict, List, Optional + +import httpx +from fastapi import HTTPException +from loguru import logger + +from kiro.parsers import parse_bracket_tool_calls, deduplicate_tool_calls +from kiro.config import ( + FIRST_TOKEN_TIMEOUT, + FIRST_TOKEN_MAX_RETRIES, +) +from kiro.tokenizer import count_tokens, count_message_tokens, count_tools_tokens + +from kiro.streaming_core import ( + parse_kiro_stream, + FirstTokenTimeoutError, + calculate_tokens_from_context_usage, + stream_with_first_token_retry as stream_with_first_token_retry_core, +) + +if TYPE_CHECKING: + from kiro.auth import KiroAuthManager + from kiro.cache import ModelInfoCache + +try: + from kiro.debug_logger import debug_logger +except ImportError: + debug_logger = None + + +__all__ = [ + "FirstTokenTimeoutError", + "stream_kiro_to_responses", + "stream_responses_with_first_token_retry", + "collect_responses_response", +] + + +def _new_response_id() -> str: + import uuid + return f"resp_{uuid.uuid4().hex}" + + +def _new_item_id(prefix: str) -> str: + import uuid + return f"{prefix}_{uuid.uuid4().hex}" + + +class _ResponsesStreamState: + """ + Holds counters and emit helpers for building the Responses SSE event stream. + + The main generator drives this by calling ensure_reasoning_open/ + ensure_message_open as thinking/content arrives, add_function_call for each + tool, and finish() at the end. Every emit helper is a generator yielding SSE + strings so the caller can `yield from` them. + """ + + def __init__(self, model: str, created_time: int): + self.model = model + self.created_time = created_time + self.response_id = _new_response_id() + + self.sequence_number = 0 + self.output_index = 0 + + # "reasoning" | "message" | None + self.current_section: Optional[str] = None + self.current_item_id: Optional[str] = None + self.current_output_index: Optional[int] = None + + # Accumulators for the current open item + self._text_buffer = "" + self._reasoning_buffer = "" + + # Completed output items (for the final response object) + self.output_items: List[Dict[str, Any]] = [] + + self.created_emitted = False + + # -- low-level SSE formatting ------------------------------------------------ + + def _sse(self, event_type: str, payload: Dict[str, Any]) -> str: + payload = dict(payload) + payload["type"] = event_type + payload["sequence_number"] = self.sequence_number + self.sequence_number += 1 + chunk = f"event: {event_type}\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n" + if debug_logger: + debug_logger.log_modified_chunk(chunk.encode("utf-8")) + return chunk + + def _response_skeleton(self, status: str, usage: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + return { + "id": self.response_id, + "object": "response", + "created_at": self.created_time, + "status": status, + "model": self.model, + "output": list(self.output_items), + "usage": usage, + } + + # -- lifecycle --------------------------------------------------------------- + + def ensure_created(self) -> AsyncGenerator[str, None]: + async def gen(): + if self.created_emitted: + return + self.created_emitted = True + yield self._sse("response.created", {"response": self._response_skeleton("in_progress")}) + yield self._sse("response.in_progress", {"response": self._response_skeleton("in_progress")}) + return gen() + + def _close_current(self) -> AsyncGenerator[str, None]: + async def gen(): + if self.current_section == "message": + item_id = self.current_item_id + oidx = self.current_output_index + yield self._sse("response.output_text.done", { + "item_id": item_id, "output_index": oidx, + "content_index": 0, "text": self._text_buffer, + }) + yield self._sse("response.content_part.done", { + "item_id": item_id, "output_index": oidx, "content_index": 0, + "part": {"type": "output_text", "text": self._text_buffer, "annotations": []}, + }) + item = { + "id": item_id, "type": "message", "status": "completed", "role": "assistant", + "content": [{"type": "output_text", "text": self._text_buffer, "annotations": []}], + } + self.output_items.append(item) + yield self._sse("response.output_item.done", {"output_index": oidx, "item": item}) + + elif self.current_section == "reasoning": + item_id = self.current_item_id + oidx = self.current_output_index + yield self._sse("response.reasoning_summary_text.done", { + "item_id": item_id, "output_index": oidx, + "summary_index": 0, "text": self._reasoning_buffer, + }) + yield self._sse("response.reasoning_summary_part.done", { + "item_id": item_id, "output_index": oidx, "summary_index": 0, + "part": {"type": "summary_text", "text": self._reasoning_buffer}, + }) + item = { + "id": item_id, "type": "reasoning", + "summary": [{"type": "summary_text", "text": self._reasoning_buffer}], + } + self.output_items.append(item) + yield self._sse("response.output_item.done", {"output_index": oidx, "item": item}) + + self.current_section = None + self.current_item_id = None + self.current_output_index = None + self._text_buffer = "" + self._reasoning_buffer = "" + return gen() + + def ensure_reasoning_open(self) -> AsyncGenerator[str, None]: + async def gen(): + async for c in self.ensure_created(): + yield c + if self.current_section == "reasoning": + return + async for c in self._close_current(): + yield c + item_id = _new_item_id("rs") + oidx = self.output_index + self.output_index += 1 + self.current_section = "reasoning" + self.current_item_id = item_id + self.current_output_index = oidx + yield self._sse("response.output_item.added", { + "output_index": oidx, + "item": {"id": item_id, "type": "reasoning", "summary": []}, + }) + yield self._sse("response.reasoning_summary_part.added", { + "item_id": item_id, "output_index": oidx, "summary_index": 0, + "part": {"type": "summary_text", "text": ""}, + }) + return gen() + + def ensure_message_open(self) -> AsyncGenerator[str, None]: + async def gen(): + async for c in self.ensure_created(): + yield c + if self.current_section == "message": + return + async for c in self._close_current(): + yield c + item_id = _new_item_id("msg") + oidx = self.output_index + self.output_index += 1 + self.current_section = "message" + self.current_item_id = item_id + self.current_output_index = oidx + yield self._sse("response.output_item.added", { + "output_index": oidx, + "item": {"id": item_id, "type": "message", "status": "in_progress", + "role": "assistant", "content": []}, + }) + yield self._sse("response.content_part.added", { + "item_id": item_id, "output_index": oidx, "content_index": 0, + "part": {"type": "output_text", "text": "", "annotations": []}, + }) + return gen() + + def add_reasoning_delta(self, text: str) -> AsyncGenerator[str, None]: + async def gen(): + async for c in self.ensure_reasoning_open(): + yield c + self._reasoning_buffer += text + yield self._sse("response.reasoning_summary_text.delta", { + "item_id": self.current_item_id, "output_index": self.current_output_index, + "summary_index": 0, "delta": text, + }) + return gen() + + def add_text_delta(self, text: str) -> AsyncGenerator[str, None]: + async def gen(): + async for c in self.ensure_message_open(): + yield c + self._text_buffer += text + yield self._sse("response.output_text.delta", { + "item_id": self.current_item_id, "output_index": self.current_output_index, + "content_index": 0, "delta": text, + }) + return gen() + + def add_function_call(self, call_id: str, name: str, arguments: str) -> AsyncGenerator[str, None]: + async def gen(): + async for c in self._close_current(): + yield c + item_id = _new_item_id("fc") + oidx = self.output_index + self.output_index += 1 + call_id_final = call_id or _new_item_id("call") + yield self._sse("response.output_item.added", { + "output_index": oidx, + "item": {"id": item_id, "type": "function_call", "status": "in_progress", + "call_id": call_id_final, "name": name, "arguments": ""}, + }) + yield self._sse("response.function_call_arguments.delta", { + "item_id": item_id, "output_index": oidx, "delta": arguments, + }) + yield self._sse("response.function_call_arguments.done", { + "item_id": item_id, "output_index": oidx, "arguments": arguments, + }) + item = {"id": item_id, "type": "function_call", "status": "completed", + "call_id": call_id_final, "name": name, "arguments": arguments} + self.output_items.append(item) + yield self._sse("response.output_item.done", {"output_index": oidx, "item": item}) + return gen() + + def finish(self, usage: Dict[str, Any], status: str = "completed") -> AsyncGenerator[str, None]: + async def gen(): + async for c in self.ensure_created(): + yield c + async for c in self._close_current(): + yield c + yield self._sse("response.completed", {"response": self._response_skeleton(status, usage)}) + return gen() + + +async def stream_kiro_to_responses_internal( + client: httpx.AsyncClient, + response: httpx.Response, + model: str, + model_cache: "ModelInfoCache", + auth_manager: "KiroAuthManager", + first_token_timeout: float = FIRST_TOKEN_TIMEOUT, + request_messages: Optional[list] = None, + request_tools: Optional[list] = None, +) -> AsyncGenerator[str, None]: + """ + Internal generator converting a Kiro stream to Responses API SSE events. + + Raises: + FirstTokenTimeoutError: If first token not received within timeout + """ + state = _ResponsesStreamState(model, int(time.time())) + + full_content = "" + full_thinking = "" + tool_calls_from_stream: List[Dict[str, Any]] = [] + metering_data = None + context_usage_percentage = None + streaming_error_occurred = False + + try: + async for event in parse_kiro_stream(response, first_token_timeout): + if event.type == "content" and event.content: + full_content += event.content + async for c in state.add_text_delta(event.content): + yield c + + elif event.type == "thinking" and event.thinking_content: + full_thinking += event.thinking_content + async for c in state.add_reasoning_delta(event.thinking_content): + yield c + + elif event.type == "tool_use" and event.tool_use: + tool_calls_from_stream.append(event.tool_use) + + elif event.type == "usage" and event.usage: + metering_data = event.usage + + elif event.type == "context_usage" and event.context_usage_percentage is not None: + context_usage_percentage = event.context_usage_percentage + + # Merge bracket-style tool calls found in the accumulated content + bracket_tool_calls = parse_bracket_tool_calls(full_content) + all_tool_calls = deduplicate_tool_calls(tool_calls_from_stream + bracket_tool_calls) + + # Emit function_call items (each is a self-contained item) + for tc in all_tool_calls: + func = tc.get("function") or {} + name = func.get("name") or tc.get("name") or "" + arguments = func.get("arguments") + if arguments is None: + raw_input = tc.get("input", {}) + arguments = json.dumps(raw_input, ensure_ascii=False) if raw_input else "{}" + if not isinstance(arguments, str): + arguments = json.dumps(arguments, ensure_ascii=False) + async for c in state.add_function_call(tc.get("id", ""), name, arguments): + yield c + + # Token accounting (same approach as the OpenAI path) + completion_tokens = count_tokens(full_content + full_thinking) + prompt_tokens, total_tokens, prompt_source, total_source = calculate_tokens_from_context_usage( + context_usage_percentage, completion_tokens, model_cache, model + ) + if prompt_source == "unknown" and request_messages: + prompt_tokens = count_message_tokens(request_messages, apply_claude_correction=False) + if request_tools: + prompt_tokens += count_tools_tokens(request_tools, apply_claude_correction=False) + total_tokens = prompt_tokens + completion_tokens + + usage = { + "input_tokens": prompt_tokens, + "output_tokens": completion_tokens, + "total_tokens": total_tokens, + } + + logger.debug( + f"[Usage] {model}: input_tokens={prompt_tokens} ({prompt_source}), " + f"output_tokens={completion_tokens} (tiktoken), total_tokens={total_tokens} ({total_source})" + ) + + async for c in state.finish(usage): + yield c + + except FirstTokenTimeoutError: + raise + except GeneratorExit: + logger.debug("Client disconnected (GeneratorExit)") + streaming_error_occurred = True + except Exception as e: + streaming_error_occurred = True + error_type = type(e).__name__ + error_msg = str(e) if str(e) else "(empty message)" + logger.error(f"Error during Responses streaming: [{error_type}] {error_msg}", exc_info=True) + raise + finally: + try: + await response.aclose() + except Exception as close_error: + logger.debug(f"Error closing response: {close_error}") + if streaming_error_occurred: + logger.debug("Responses streaming completed with error") + else: + logger.debug("Responses streaming completed successfully") + + +async def stream_kiro_to_responses( + client: httpx.AsyncClient, + response: httpx.Response, + model: str, + model_cache: "ModelInfoCache", + auth_manager: "KiroAuthManager", + request_messages: Optional[list] = None, + request_tools: Optional[list] = None, +) -> AsyncGenerator[str, None]: + """Non-retrying wrapper over stream_kiro_to_responses_internal.""" + async for chunk in stream_kiro_to_responses_internal( + client, response, model, model_cache, auth_manager, + request_messages=request_messages, request_tools=request_tools, + ): + yield chunk + + +async def stream_responses_with_first_token_retry( + make_request: Callable[[], Awaitable[httpx.Response]], + client: httpx.AsyncClient, + model: str, + model_cache: "ModelInfoCache", + auth_manager: "KiroAuthManager", + initial_response: Optional[httpx.Response] = None, + max_retries: int = FIRST_TOKEN_MAX_RETRIES, + first_token_timeout: float = FIRST_TOKEN_TIMEOUT, + request_messages: Optional[list] = None, + request_tools: Optional[list] = None, +) -> AsyncGenerator[str, None]: + """Streaming with automatic retry on first token timeout (Responses format).""" + + def create_http_error(status_code: int, error_text: str) -> HTTPException: + return HTTPException(status_code=status_code, detail=f"Upstream API error: {error_text}") + + def create_timeout_error(retries: int, timeout: float) -> HTTPException: + return HTTPException( + status_code=504, + detail=f"Model did not respond within {timeout}s after {retries} attempts. Please try again.", + ) + + async def stream_processor(resp: httpx.Response) -> AsyncGenerator[str, None]: + async for chunk in stream_kiro_to_responses_internal( + client, resp, model, model_cache, auth_manager, + first_token_timeout=first_token_timeout, + request_messages=request_messages, request_tools=request_tools, + ): + yield chunk + + async for chunk in stream_with_first_token_retry_core( + make_request=make_request, + stream_processor=stream_processor, + initial_response=initial_response, + max_retries=max_retries, + first_token_timeout=first_token_timeout, + on_http_error=create_http_error, + on_all_retries_failed=create_timeout_error, + ): + yield chunk + + +async def collect_responses_response( + client: httpx.AsyncClient, + response: httpx.Response, + model: str, + model_cache: "ModelInfoCache", + auth_manager: "KiroAuthManager", + request_messages: Optional[list] = None, + request_tools: Optional[list] = None, +) -> dict: + """ + Collect a full non-streaming Responses API response object. + + Consumes the Kiro stream, accumulates content/reasoning/tool calls, and + builds the final `response` object with output items and usage. + """ + from kiro.streaming_core import collect_stream_to_result + + result = await collect_stream_to_result(response) + + output_items: List[Dict[str, Any]] = [] + + if result.thinking_content: + output_items.append({ + "id": _new_item_id("rs"), + "type": "reasoning", + "summary": [{"type": "summary_text", "text": result.thinking_content}], + }) + + if result.content: + output_items.append({ + "id": _new_item_id("msg"), + "type": "message", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": result.content, "annotations": []}], + }) + + for tc in result.tool_calls: + func = tc.get("function") or {} + name = func.get("name") or tc.get("name") or "" + arguments = func.get("arguments") + if arguments is None: + raw_input = tc.get("input", {}) + arguments = json.dumps(raw_input, ensure_ascii=False) if raw_input else "{}" + if not isinstance(arguments, str): + arguments = json.dumps(arguments, ensure_ascii=False) + output_items.append({ + "id": _new_item_id("fc"), + "type": "function_call", + "status": "completed", + "call_id": tc.get("id") or _new_item_id("call"), + "name": name, + "arguments": arguments or "{}", + }) + + completion_tokens = count_tokens(result.content + result.thinking_content) + prompt_tokens, total_tokens, prompt_source, _ = calculate_tokens_from_context_usage( + result.context_usage_percentage, completion_tokens, model_cache, model + ) + if prompt_source == "unknown" and request_messages: + prompt_tokens = count_message_tokens(request_messages, apply_claude_correction=False) + if request_tools: + prompt_tokens += count_tools_tokens(request_tools, apply_claude_correction=False) + total_tokens = prompt_tokens + completion_tokens + + return { + "id": _new_response_id(), + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "model": model, + "output": output_items, + "usage": { + "input_tokens": prompt_tokens, + "output_tokens": completion_tokens, + "total_tokens": total_tokens, + }, + } diff --git a/kiro/utils.py b/kiro/utils.py index 1ebe46b0..8ca1dc11 100644 --- a/kiro/utils.py +++ b/kiro/utils.py @@ -27,7 +27,7 @@ import hashlib import json import uuid -from typing import TYPE_CHECKING, List, Dict, Any +from typing import TYPE_CHECKING, List, Dict, Any, Optional from loguru import logger @@ -78,7 +78,7 @@ def get_kiro_headers(auth_manager: "KiroAuthManager", token: str) -> dict: return { "Authorization": f"Bearer {token}", - "Content-Type": "application/x-amz-json-1.0", + "Content-Type": "application/json", "x-amz-target": "AmazonCodeWhispererStreamingService.GenerateAssistantResponse", "User-Agent": f"aws-sdk-js/1.0.27 ua/2.1 os/win32#10.0.19044 lang/js md/nodejs#22.21.1 api/codewhispererstreaming#1.0.27 m/E KiroIDE-0.7.45-{fingerprint}", "x-amz-user-agent": f"aws-sdk-js/1.0.27 KiroIDE-0.7.45-{fingerprint}", @@ -163,6 +163,52 @@ def generate_conversation_id(messages: List[Dict[str, Any]] = None) -> str: return hash_digest[:16] +def derive_conversation_id_from_metadata(metadata: Optional[Dict[str, Any]]) -> Optional[str]: + """ + Derive a stable Kiro conversation ID from client request metadata. + + Claude Code Desktop sends ``metadata.user_id`` as a JSON string containing + a stable ``session_id`` for the current conversation. When present, this + value is used directly as Kiro's conversation ID to preserve the client's + native conversation identity. + + Args: + metadata: Anthropic request metadata, if provided by the client. + + Returns: + Raw session_id string when a non-empty session_id is available, otherwise + None so callers can keep the existing fallback behavior. + """ + if not metadata: + return None + + user_id = metadata.get("user_id") + if user_id is None: + return None + + if isinstance(user_id, str): + try: + user_id_data = json.loads(user_id) + except json.JSONDecodeError: + logger.debug("metadata.user_id is not JSON; using fallback conversation ID") + return None + elif isinstance(user_id, dict): + user_id_data = user_id + else: + logger.debug("metadata.user_id has unsupported type; using fallback conversation ID") + return None + + if not isinstance(user_id_data, dict): + logger.debug("metadata.user_id JSON is not an object; using fallback conversation ID") + return None + + session_id = user_id_data.get("session_id") + if not isinstance(session_id, str) or not session_id.strip(): + return None + + return session_id.strip() + + def generate_tool_call_id() -> str: """ Generates a unique ID for tool call. diff --git a/main.py b/main.py index 40cc2c30..7accd3e4 100644 --- a/main.py +++ b/main.py @@ -88,6 +88,8 @@ from kiro.routes_anthropic import router as anthropic_router from kiro.exceptions import validation_exception_handler from kiro.debug_middleware import DebugLoggerMiddleware +from kiro.routes_responses import router as responses_router +from extensions.tool_name_alias import install_tool_name_aliasing # --- Loguru Configuration --- @@ -99,6 +101,17 @@ format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}" ) +logger.add( + Path(__file__).parent / "errors.log", + level="ERROR", + rotation="10 MB", + retention=5, + encoding="utf-8", + backtrace=True, + diagnose=True, + format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}" +) + class InterceptHandler(logging.Handler): """ @@ -332,7 +345,10 @@ async def lifespan(app: FastAPI): concurrent requests efficiently (fixes issue #24). """ logger.info("Starting application... Creating state managers.") - + + # Install tool-name aliasing patch (handles MCP tool names exceeding Kiro's 64-char limit) + install_tool_name_aliasing() + # Create shared HTTP client with connection pooling # This reduces memory usage and enables connection reuse across requests # Limits: max 100 total connections, max 20 keep-alive connections @@ -571,6 +587,9 @@ def _add_env_overrides(entry: dict) -> None: # Anthropic-compatible API: /v1/messages app.include_router(anthropic_router) +# Responses API +app.include_router(responses_router) + # --- Uvicorn log config --- # Minimal configuration for redirecting uvicorn logs to loguru. diff --git a/pytest.ini b/pytest.ini index 89255046..9e5c4eaf 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,14 +1,12 @@ [pytest] -# Конфигурация pytest для проекта testpaths = tests python_files = test_*.py python_classes = Test* python_functions = test_* -# Добавляем корневую директорию в PYTHONPATH pythonpath = . -# Исключаем manual_api_test.py из автоматического запуска -# (это скрипт для ручного тестирования реального API, не unit-тест) -# Чтобы запустить его: python manual_api_test.py -norecursedirs = .git __pycache__ old requests _notes \ No newline at end of file +norecursedirs = .git __pycache__ old requests _notes + +markers = + slow: marks tests as slow (deselect with '-m "not slow"') \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index f2f74a28..ff5db917 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,10 +1,11 @@ # Prod dependencies fastapi uvicorn[standard] -httpx +httpx[socks] loguru python-dotenv tiktoken +pypdf # Testing dependencies pytest diff --git a/tests/conftest.py b/tests/conftest.py index 69ba58a8..dfbefe80 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,6 +8,7 @@ """ import asyncio +import warnings import json import pytest import time @@ -16,6 +17,11 @@ from datetime import datetime, timezone import httpx + +# Suppress StarletteDeprecationWarning: httpx2 package doesn't exist yet, +# and starlette.testclient always falls back to httpx with a deprecation warning. +warnings.filterwarnings("ignore", message=".*httpx.*starlette.testclient.*") + from fastapi.testclient import TestClient @@ -86,6 +92,13 @@ def setup_test_environment(tmp_path_factory): import kiro.config original_creds_file = kiro.config.ACCOUNTS_CONFIG_FILE original_state_file = kiro.config.ACCOUNTS_STATE_FILE + # Also set OS env vars so that config module reloads (e.g. in test_config.py) + # pick up the correct temp paths instead of defaults. + import os + original_env_creds = os.environ.get("ACCOUNTS_CONFIG_FILE") + original_env_state = os.environ.get("ACCOUNTS_STATE_FILE") + os.environ["ACCOUNTS_CONFIG_FILE"] = str(creds_file) + os.environ["ACCOUNTS_STATE_FILE"] = str(tmp_dir / "state.json") kiro.config.ACCOUNTS_CONFIG_FILE = str(creds_file) kiro.config.ACCOUNTS_STATE_FILE = str(tmp_dir / "state.json") @@ -98,6 +111,14 @@ def setup_test_environment(tmp_path_factory): # Restore original paths kiro.config.ACCOUNTS_CONFIG_FILE = original_creds_file kiro.config.ACCOUNTS_STATE_FILE = original_state_file + if original_env_creds is not None: + os.environ["ACCOUNTS_CONFIG_FILE"] = original_env_creds + else: + os.environ.pop("ACCOUNTS_CONFIG_FILE", None) + if original_env_state is not None: + os.environ["ACCOUNTS_STATE_FILE"] = original_env_state + else: + os.environ.pop("ACCOUNTS_STATE_FILE", None) print("🧹 Test environment cleaned up") @@ -525,7 +546,17 @@ async def __aexit__(self, *args): def clean_app(): """ Returns a "clean" application instance for each test. + + Ensures config module attributes are synchronized with OS env vars, + which is needed when test_config.py reloads the config module. """ + import os + import kiro.config + if "ACCOUNTS_CONFIG_FILE" in os.environ: + kiro.config.ACCOUNTS_CONFIG_FILE = os.environ["ACCOUNTS_CONFIG_FILE"] + if "ACCOUNTS_STATE_FILE" in os.environ: + kiro.config.ACCOUNTS_STATE_FILE = os.environ["ACCOUNTS_STATE_FILE"] + print("Importing application for test...") from main import app # Reset all dependency overrides before test @@ -543,6 +574,13 @@ def test_client(clean_app): with TestClient(clean_app) as client: yield client print("Closing TestClient...") + # Cleanup: uninstall tool name aliasing patches so subsequent tests + # in the same process see unpatched module state. + try: + from extensions.tool_name_alias import uninstall_tool_name_aliasing + uninstall_tool_name_aliasing() + except ImportError: + pass @pytest.fixture diff --git a/tests/unit/test_account_manager.py b/tests/unit/test_account_manager.py index 926d78b4..3b46e99b 100644 --- a/tests/unit/test_account_manager.py +++ b/tests/unit/test_account_manager.py @@ -23,7 +23,8 @@ AccountStats, ModelAccountList, AccountManager, - _format_duration + _format_duration, + _fetch_model_catalog, ) from kiro.account_errors import ErrorType from kiro.auth import KiroAuthManager, AuthType @@ -1300,3 +1301,76 @@ def test_format_duration_days(self): """Test formatting days.""" assert _format_duration(86400) == "1d" assert _format_duration(172800) == "2d" + + +class TestFetchModelCatalog: + """The control-plane ListAvailableModels call (management host, JSON-1.0 RPC).""" + + @pytest.mark.asyncio + async def test_hits_management_host_with_target_header(self): + auth = MagicMock() + auth.management_host = "https://management.us-east-1.kiro.dev" + auth.profile_arn = "arn:aws:codewhisperer:us-east-1:123:profile/ABC" + + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = {"models": [{"modelId": "claude-opus-4.6"}]} + + with patch('kiro.account_manager.KiroHttpClient') as mock_http_class: + mock_client = AsyncMock() + mock_client.request_with_retry = AsyncMock(return_value=mock_response) + mock_client.close = AsyncMock() + mock_http_class.return_value = mock_client + + models = await _fetch_model_catalog(auth) + + assert models == [{"modelId": "claude-opus-4.6"}] + _, kwargs = mock_client.request_with_retry.call_args + assert kwargs["method"] == "POST" + assert kwargs["url"] == "https://management.us-east-1.kiro.dev/" + assert kwargs["extra_headers"]["x-amz-target"] == \ + "AmazonCodeWhispererService.ListAvailableModels" + assert kwargs["extra_headers"]["Content-Type"] == "application/x-amz-json-1.0" + assert kwargs["json_data"]["profileArn"] == auth.profile_arn + mock_client.close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_omits_profile_arn_when_absent(self): + auth = MagicMock() + auth.management_host = "https://management.us-east-1.kiro.dev" + auth.profile_arn = None + + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = {"models": []} + + with patch('kiro.account_manager.KiroHttpClient') as mock_http_class: + mock_client = AsyncMock() + mock_client.request_with_retry = AsyncMock(return_value=mock_response) + mock_client.close = AsyncMock() + mock_http_class.return_value = mock_client + + await _fetch_model_catalog(auth) + + _, kwargs = mock_client.request_with_retry.call_args + assert "profileArn" not in kwargs["json_data"] + + @pytest.mark.asyncio + async def test_non_200_raises_and_closes(self): + auth = MagicMock() + auth.management_host = "https://management.us-east-1.kiro.dev" + auth.profile_arn = None + + mock_response = Mock() + mock_response.status_code = 500 + + with patch('kiro.account_manager.KiroHttpClient') as mock_http_class: + mock_client = AsyncMock() + mock_client.request_with_retry = AsyncMock(return_value=mock_response) + mock_client.close = AsyncMock() + mock_http_class.return_value = mock_client + + with pytest.raises(Exception): + await _fetch_model_catalog(auth) + + mock_client.close.assert_awaited_once() diff --git a/tests/unit/test_converters_anthropic.py b/tests/unit/test_converters_anthropic.py index b2b37402..19bed286 100644 --- a/tests/unit/test_converters_anthropic.py +++ b/tests/unit/test_converters_anthropic.py @@ -1493,11 +1493,21 @@ def test_includes_system_prompt(self): result = anthropic_to_kiro(request, "conv-123", "arn:aws:test") print(f"Result: {result}") - current_content = result["conversationState"]["currentMessage"][ - "userInputMessage" - ]["content"] - print(f"Current content: {current_content}") - assert "You are a helpful assistant." in current_content + # System prompt now goes into history via build_kiro_system_history + history = result["conversationState"].get("history", []) + found_system = False + for entry in history: + if "userInputMessage" in entry: + content = entry["userInputMessage"].get("content", "") + if isinstance(content, str) and "You are a helpful assistant." in content: + found_system = True + break + elif isinstance(content, list): + for block in content: + if isinstance(block, dict) and "You are a helpful assistant." in str(block.get("text", "")): + found_system = True + break + assert found_system, f"System prompt not found in history. History: {history}" def test_includes_tools(self): """ @@ -1565,9 +1575,13 @@ def test_builds_history_for_multi_turn(self): print(f"Result: {result}") history = result["conversationState"].get("history", []) print(f"History length: {len(history)}") - assert len(history) == 2 # First user + assistant - assert "userInputMessage" in history[0] - assert "assistantResponseMessage" in history[1] + # History includes system history entries + conversation history + assert len(history) >= 2 # At least conversation history + # Find conversation history entries (skip system history) + user_msgs = [h for h in history if "userInputMessage" in h and h["userInputMessage"].get("content") != ""] + assistant_msgs = [h for h in history if "assistantResponseMessage" in h and h["assistantResponseMessage"].get("content") != ""] + assert len(user_msgs) >= 1, "Should have at least one user message" + assert len(assistant_msgs) >= 1, "Should have at least one assistant message" def test_handles_tool_use_and_result_flow(self): """ @@ -1676,9 +1690,10 @@ def test_injects_thinking_tags_when_enabled(self): "kiro.converters_anthropic.get_model_id_for_kiro", return_value="claude-sonnet-4.5", ): - with patch("kiro.converters_core.FAKE_REASONING_ENABLED", True): - with patch("kiro.converters_core.FAKE_REASONING_MAX_TOKENS", 4000): - result = anthropic_to_kiro(request, "conv-123", "arn:aws:test") + with patch("kiro.converters_core.NATIVE_REASONING_ENABLED", False): + with patch("kiro.converters_core.FAKE_REASONING_ENABLED", True): + with patch("kiro.converters_core.FAKE_REASONING_MAX_TOKENS", 4000): + result = anthropic_to_kiro(request, "conv-123", "arn:aws:test") print(f"Result: {result}") current_content = result["conversationState"]["currentMessage"][ @@ -1726,9 +1741,10 @@ def test_injects_thinking_tags_even_when_tool_results_present(self): "kiro.converters_anthropic.get_model_id_for_kiro", return_value="claude-sonnet-4.5", ): - with patch("kiro.converters_core.FAKE_REASONING_ENABLED", True): - with patch("kiro.converters_core.FAKE_REASONING_MAX_TOKENS", 4000): - result = anthropic_to_kiro(request, "conv-123", "arn:aws:test") + with patch("kiro.converters_core.NATIVE_REASONING_ENABLED", False): + with patch("kiro.converters_core.FAKE_REASONING_ENABLED", True): + with patch("kiro.converters_core.FAKE_REASONING_MAX_TOKENS", 4000): + result = anthropic_to_kiro(request, "conv-123", "arn:aws:test") print(f"Result: {result}") current_content = result["conversationState"]["currentMessage"][ @@ -1873,9 +1889,10 @@ def test_extracts_and_passes_thinking_config(self): print("Calling anthropic_to_kiro...") with patch("kiro.converters_anthropic.get_model_id_for_kiro", return_value="claude-sonnet-4.5"): - with patch("kiro.converters_core.FAKE_REASONING_ENABLED", True): - with patch("kiro.converters_core.FAKE_REASONING_BUDGET_CAP", 10000): - payload = anthropic_to_kiro(request, "test-conv-123", "arn:aws:test") + with patch("kiro.converters_core.NATIVE_REASONING_ENABLED", False): + with patch("kiro.converters_core.FAKE_REASONING_ENABLED", True): + with patch("kiro.converters_core.FAKE_REASONING_BUDGET_CAP", 10000): + payload = anthropic_to_kiro(request, "test-conv-123", "arn:aws:test") print("Extracting userInputMessage content...") user_input = payload["conversationState"]["currentMessage"]["userInputMessage"] diff --git a/tests/unit/test_converters_core.py b/tests/unit/test_converters_core.py index 1b8ccd33..f4265b0c 100644 --- a/tests/unit/test_converters_core.py +++ b/tests/unit/test_converters_core.py @@ -3537,9 +3537,10 @@ def test_injects_tags_when_enabled(self): content = "What is 2+2?" print("Action: Inject thinking tags with FAKE_REASONING_ENABLED=True...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = inject_thinking_tags(content, ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = inject_thinking_tags(content, ThinkingConfig()) print(f"Result: {result[:200]}...") print("Checking that thinking_mode tag is present...") @@ -3560,9 +3561,10 @@ def test_injects_thinking_instruction_tag(self): content = "Analyze this code" print("Action: Inject thinking tags...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 8000): - result = inject_thinking_tags(content, ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 8000): + result = inject_thinking_tags(content, ThinkingConfig()) print(f"Result length: {len(result)} chars") print("Checking that thinking_instruction tag is present...") @@ -3578,9 +3580,10 @@ def test_thinking_instruction_contains_english_directive(self): content = "Test" print("Action: Inject thinking tags...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = inject_thinking_tags(content, ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = inject_thinking_tags(content, ThinkingConfig()) print("Checking for English directive...") assert "Think in English" in result @@ -3594,10 +3597,11 @@ def test_uses_configured_max_tokens(self): content = "Test" print("Action: Inject thinking tags with FAKE_REASONING_MAX_TOKENS=16000...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 16000): - with patch('kiro.converters_core.FAKE_REASONING_BUDGET_CAP', 0): # Disable cap - result = inject_thinking_tags(content, ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 16000): + with patch('kiro.converters_core.FAKE_REASONING_BUDGET_CAP', 0): # Disable cap + result = inject_thinking_tags(content, ThinkingConfig()) print(f"Result: {result[:300]}...") print("Checking that max_thinking_length uses configured value...") @@ -3612,9 +3616,10 @@ def test_preserves_empty_content(self): content = "" print("Action: Inject thinking tags...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = inject_thinking_tags(content, ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = inject_thinking_tags(content, ThinkingConfig()) print(f"Result length: {len(result)} chars") print("Checking that tags are present even with empty content...") @@ -3630,9 +3635,10 @@ def test_preserves_multiline_content(self): content = "Line 1\nLine 2\nLine 3" print("Action: Inject thinking tags...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = inject_thinking_tags(content, ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = inject_thinking_tags(content, ThinkingConfig()) print("Checking that multiline content is preserved...") assert "Line 1\nLine 2\nLine 3" in result @@ -3646,9 +3652,10 @@ def test_preserves_special_characters(self): content = "Check this example and {json: 'value'}" print("Action: Inject thinking tags...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = inject_thinking_tags(content, ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = inject_thinking_tags(content, ThinkingConfig()) print("Checking that special characters are preserved...") assert "example" in result @@ -3663,9 +3670,10 @@ def test_thinking_instruction_contains_systematic_approach(self): content = "Test" print("Action: Inject thinking tags...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = inject_thinking_tags(content, ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = inject_thinking_tags(content, ThinkingConfig()) print("Checking for systematic approach keywords...") assert "thorough" in result.lower() or "systematic" in result.lower() @@ -3679,9 +3687,10 @@ def test_thinking_instruction_contains_understanding_step(self): content = "Test" print("Action: Inject thinking tags...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = inject_thinking_tags(content, ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = inject_thinking_tags(content, ThinkingConfig()) print("Checking for understanding step...") assert "understand" in result.lower() @@ -3695,9 +3704,10 @@ def test_thinking_instruction_contains_verification_step(self): content = "Test" print("Action: Inject thinking tags...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = inject_thinking_tags(content, ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = inject_thinking_tags(content, ThinkingConfig()) print("Checking for verification step...") assert "verify" in result.lower() @@ -3711,9 +3721,10 @@ def test_thinking_instruction_contains_quality_emphasis(self): content = "Test" print("Action: Inject thinking tags...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = inject_thinking_tags(content, ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = inject_thinking_tags(content, ThinkingConfig()) print("Checking for quality emphasis...") assert "quality" in result.lower() @@ -3727,9 +3738,10 @@ def test_tag_order_is_correct(self): content = "USER_CONTENT_HERE" print("Action: Inject thinking tags...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = inject_thinking_tags(content, ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = inject_thinking_tags(content, ThinkingConfig()) print("Checking tag order...") thinking_mode_pos = result.find("") @@ -5715,11 +5727,18 @@ def test_includes_images_in_history(self): print(f"History length: {len(history)}") assert len(history) >= 1 - print("Checking that first history message has images directly in userInputMessage (Issue #32 fix)...") - first_msg = history[0]["userInputMessage"] - assert "images" in first_msg + print("Checking that a history message has images directly in userInputMessage (Issue #32 fix)...") + # Find the history entry with images (may not be first due to system history entries) + history_msg_with_images = None + for entry in history: + if "userInputMessage" in entry: + msg = entry["userInputMessage"] + if "images" in msg: + history_msg_with_images = msg + break + assert history_msg_with_images is not None, "No history entry with images found" - images = first_msg["images"] + images = history_msg_with_images["images"] print(f"History images: {images}") assert len(images) == 1 assert images[0]["format"] == "jpeg" @@ -5914,17 +5933,18 @@ def test_images_with_thinking_injection(self): ] print("Action: Building Kiro payload with thinking injection...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = build_kiro_payload( - messages=messages, - system_prompt="", - model_id="claude-sonnet-4", - tools=None, - conversation_id="test-conv", - profile_arn="arn:test", - thinking_config=ThinkingConfig(enabled=True) - ) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = build_kiro_payload( + messages=messages, + system_prompt="", + model_id="claude-sonnet-4", + tools=None, + conversation_id="test-conv", + profile_arn="arn:test", + thinking_config=ThinkingConfig(enabled=True) + ) current_msg = result.payload["conversationState"]["currentMessage"]["userInputMessage"] @@ -6301,6 +6321,7 @@ def test_disabled_by_client_request(self, monkeypatch): Purpose: Ensure client can disable thinking per-request """ print("Setting FAKE_REASONING_ENABLED=True...") + monkeypatch.setattr("kiro.converters_core.NATIVE_REASONING_ENABLED", False) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_ENABLED", True) config = ThinkingConfig(enabled=False, budget_tokens=None) @@ -6318,6 +6339,7 @@ def test_uses_default_budget(self, monkeypatch): Purpose: Ensure default budget fallback works """ print("Setting FAKE_REASONING_ENABLED=True, FAKE_REASONING_MAX_TOKENS=4000...") + monkeypatch.setattr("kiro.converters_core.NATIVE_REASONING_ENABLED", False) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_ENABLED", True) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_MAX_TOKENS", 4000) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_BUDGET_CAP", 10000) @@ -6339,6 +6361,7 @@ def test_uses_custom_budget(self, monkeypatch): Purpose: Ensure client-provided budget is respected """ print("Setting FAKE_REASONING_ENABLED=True...") + monkeypatch.setattr("kiro.converters_core.NATIVE_REASONING_ENABLED", False) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_ENABLED", True) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_BUDGET_CAP", 10000) @@ -6361,6 +6384,7 @@ def test_applies_cap_with_warning(self, monkeypatch): from unittest.mock import patch, call print("Setting FAKE_REASONING_ENABLED=True, cap=10000...") + monkeypatch.setattr("kiro.converters_core.NATIVE_REASONING_ENABLED", False) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_ENABLED", True) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_BUDGET_CAP", 10000) @@ -6391,6 +6415,7 @@ def test_uses_budget_when_below_cap(self, monkeypatch): Purpose: Ensure cap doesn't affect budgets below limit """ print("Setting FAKE_REASONING_ENABLED=True, cap=10000...") + monkeypatch.setattr("kiro.converters_core.NATIVE_REASONING_ENABLED", False) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_ENABLED", True) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_BUDGET_CAP", 10000) @@ -6409,6 +6434,7 @@ def test_cap_disabled_when_zero(self, monkeypatch): Purpose: Ensure users can disable capping """ print("Setting FAKE_REASONING_ENABLED=True, cap=0 (disabled)...") + monkeypatch.setattr("kiro.converters_core.NATIVE_REASONING_ENABLED", False) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_ENABLED", True) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_BUDGET_CAP", 0) @@ -6431,6 +6457,7 @@ def test_passes_thinking_config_to_inject(self, monkeypatch): Purpose: Ensure thinking configuration flows through the pipeline """ print("Setting up mocks...") + monkeypatch.setattr("kiro.converters_core.NATIVE_REASONING_ENABLED", False) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_ENABLED", True) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_BUDGET_CAP", 10000) diff --git a/tests/unit/test_converters_openai.py b/tests/unit/test_converters_openai.py index 1c3dc1e9..816503a8 100644 --- a/tests/unit/test_converters_openai.py +++ b/tests/unit/test_converters_openai.py @@ -692,8 +692,8 @@ def test_builds_simple_payload(self): def test_includes_system_prompt_in_first_message(self): """ - What it does: Verifies adding system prompt to first message. - Purpose: Ensure system prompt is merged with user message. + What it does: Verifies system prompt is included in history via build_kiro_system_history. + Purpose: Ensure system prompt appears in conversation history. """ print("Setup: Request with system prompt...") request = ChatCompletionRequest( @@ -708,9 +708,12 @@ def test_includes_system_prompt_in_first_message(self): result = build_kiro_payload(request, "conv-123", "") print(f"Result: {result}") - current_content = result["conversationState"]["currentMessage"]["userInputMessage"]["content"] - assert "You are helpful" in current_content - assert "Hello" in current_content + # System prompt now goes into history via build_kiro_system_history + history = result["conversationState"].get("history", []) + assert len(history) >= 2 + system_entry = history[0] + assert "userInputMessage" in system_entry + assert "You are helpful" in str(system_entry["userInputMessage"]["content"]) def test_builds_history_for_multi_turn(self): """ @@ -732,7 +735,8 @@ def test_builds_history_for_multi_turn(self): print(f"Result: {result}") assert "history" in result["conversationState"] - assert len(result["conversationState"]["history"]) == 2 + # History includes conversation messages (may also include system history entries) + assert len(result["conversationState"]["history"]) >= 2 def test_handles_assistant_as_last_message(self): """ @@ -880,9 +884,10 @@ def test_injects_thinking_tags_even_when_tool_results_present(self): ) print("Action: Building payload with FAKE_REASONING_ENABLED=True...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = build_kiro_payload(request, "conv-123", "") + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = build_kiro_payload(request, "conv-123", "") current_msg = result["conversationState"]["currentMessage"]["userInputMessage"] content = current_msg["content"] @@ -907,9 +912,10 @@ def test_injects_thinking_tags_when_no_tool_results(self): ) print("Action: Building payload with FAKE_REASONING_ENABLED=True...") - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = build_kiro_payload(request, "conv-123", "") + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = build_kiro_payload(request, "conv-123", "") current_msg = result["conversationState"]["currentMessage"]["userInputMessage"] content = current_msg["content"] @@ -1297,13 +1303,21 @@ def test_multiple_assistant_tool_calls_with_results(self): # Should have userInputMessage and assistantResponseMessage in history assert len(history) >= 2, f"Expected at least 2 elements in history, got {len(history)}" - # Find assistantResponseMessage + # Find assistantResponseMessage with toolUses (skip system history entries) assistant_msgs = [h for h in history if "assistantResponseMessage" in h] - print(f"Assistant messages in history: {assistant_msgs}") + print(f"Assistant messages in history: {len(assistant_msgs)}") assert len(assistant_msgs) >= 1, "Should have at least one assistantResponseMessage" + # Find the assistant message that has toolUses (not the system history one) + assistant_msg = None + for msg in assistant_msgs: + arm = msg["assistantResponseMessage"] + if arm.get("toolUses"): + assistant_msg = arm + break + assert assistant_msg is not None, "Should have an assistantResponseMessage with toolUses" + # Check that assistantResponseMessage has both toolUses - assistant_msg = assistant_msgs[0]["assistantResponseMessage"] tool_uses = assistant_msg.get("toolUses", []) print(f"ToolUses in assistant: {tool_uses}") print(f"Comparing toolUses count: Expected 2, Got {len(tool_uses)}") @@ -1356,10 +1370,18 @@ def test_long_tool_description_added_to_system_prompt(self): result = build_kiro_payload(request, "conv-123", "") print("Checking that system prompt contains tool documentation...") - current_content = result["conversationState"]["currentMessage"]["userInputMessage"]["content"] - assert "You are helpful" in current_content - assert "## Tool: long_tool" in current_content - assert long_desc in current_content + # System prompt + tool documentation now goes into history + history = result["conversationState"].get("history", []) + system_content = "" + for entry in history: + if "userInputMessage" in entry: + content = entry["userInputMessage"].get("content", "") + if isinstance(content, str) and "You are helpful" in content: + system_content = content + break + assert "You are helpful" in system_content, f"System prompt not found in history. History: {history}" + assert "## Tool: long_tool" in system_content + assert long_desc in system_content print("Checking that tool in context has reference description...") tools_context = result["conversationState"]["currentMessage"]["userInputMessage"]["userInputMessageContext"]["tools"] @@ -1846,6 +1868,7 @@ def test_extracts_and_passes_thinking_config(self, monkeypatch): Purpose: Ensure end-to-end thinking configuration flow works """ print("Setting up mocks...") + monkeypatch.setattr("kiro.converters_core.NATIVE_REASONING_ENABLED", False) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_ENABLED", True) monkeypatch.setattr("kiro.converters_core.FAKE_REASONING_BUDGET_CAP", 10000) diff --git a/tests/unit/test_debug_logger.py b/tests/unit/test_debug_logger.py index a1138a61..81c62650 100644 --- a/tests/unit/test_debug_logger.py +++ b/tests/unit/test_debug_logger.py @@ -58,16 +58,21 @@ def test_log_request_body_does_nothing(self, tmp_path): class TestDebugLoggerModeAll: """Тесты для режима DEBUG_MODE=all.""" - def test_prepare_new_request_clears_directory(self, tmp_path): + def test_prepare_new_request_clears_latest_files(self, tmp_path): """ - Что он делает: Проверяет, что prepare_new_request очищает директорию в режиме all. - Цель: Убедиться, что старые логи удаляются. + Checks that prepare_new_request clears known latest files in mode all. + In mode 'all', the logger archives old requests to timestamped directories + and clears only the known 'latest' file list at the top level. """ - print("Настройка: Режим all, создаём старый файл...") + print("Setup: mode all, create old files...") debug_dir = tmp_path / "debug_logs" debug_dir.mkdir() - old_file = debug_dir / "old_file.txt" - old_file.write_text("old content") + # Create a known latest file (should be cleared) + latest_file = debug_dir / "request_body.json" + latest_file.write_text("old content") + # Create an unknown file (should NOT be cleared in mode 'all') + unknown_file = debug_dir / "random.txt" + unknown_file.write_text("keep me") with patch('kiro.debug_logger.DEBUG_MODE', 'all'): from kiro.debug_logger import DebugLogger @@ -76,11 +81,21 @@ def test_prepare_new_request_clears_directory(self, tmp_path): logger.__init__() logger.debug_dir = debug_dir - print("Действие: Вызов prepare_new_request...") + print("Action: Call prepare_new_request...") logger.prepare_new_request() - print(f"Проверяем, что старый файл удалён...") - assert not old_file.exists() + print("Check: Known latest file should be cleared...") + assert not latest_file.exists(), \ + "known latest files should be cleared in mode 'all'" + + print("Check: Unknown files are preserved (archived in subdir)...") + # unknown_file might still exist or might have been moved + # The key behavior is that request_body.json is gone + + print("Check: request directory was created...") + requests_dir = debug_dir / "requests" + assert requests_dir.exists(), \ + "requests subdirectory should exist in mode 'all'" print(f"Проверяем, что директория существует...") assert debug_dir.exists() diff --git a/tests/unit/test_native_reasoning.py b/tests/unit/test_native_reasoning.py new file mode 100644 index 00000000..73ac66ec --- /dev/null +++ b/tests/unit/test_native_reasoning.py @@ -0,0 +1,150 @@ +# -*- coding: utf-8 -*- + +""" +Unit tests for catalog-driven native reasoning. + +Covers parsing per-model additionalModelRequestFieldsSchema into effort descriptors, +clamping requested effort onto the advertised set, and the end-to-end payload built by +build_native_effort_fields when a live catalog is (and isn't) registered. +""" + +import pytest + +from kiro import native_reasoning as nr +from kiro.converters_core import build_native_effort_fields + + +# Real advertised schema for claude-opus-4.6 (effort enum: low/medium/high/max). +OPUS_46 = { + "modelId": "claude-opus-4.6", + "additionalModelRequestFieldsSchema": { + "type": "object", + "properties": { + "thinking": { + "type": "object", + "properties": { + "type": {"enum": ["adaptive", "disabled"]}, + "display": {"enum": ["summarized", "omitted"]}, + }, + }, + "output_config": { + "type": "object", + "properties": { + "effort": { + "type": "string", + "enum": ["low", "medium", "high", "max"], + "default": "high", + } + }, + }, + "max_tokens": {"type": "integer", "minimum": 1024, "maximum": 64000}, + }, + "additionalProperties": False, + }, +} + +REASONING_MODEL = { + "modelId": "gpt-5.6-sol", + "additionalModelRequestFieldsSchema": { + "type": "object", + "properties": { + "reasoning": { + "type": "object", + "properties": {"effort": {"type": "string", "enum": ["low", "high"]}}, + } + }, + }, +} + +NO_SCHEMA_MODEL = {"modelId": "claude-sonnet-4", "description": "no schema"} + + +@pytest.fixture(autouse=True) +def _clear_registry(): + """Each test starts from an empty registry and leaves it empty.""" + nr.register_from_catalog([]) + yield + nr.register_from_catalog([]) + + +class TestRegisterFromCatalog: + def test_parses_output_config_schema(self): + nr.register_from_catalog([OPUS_46]) + d = nr.get_descriptor("claude-opus-4.6") + assert d is not None + assert d.schema_type == "output_config" + assert d.valid_efforts == ["low", "medium", "high", "max"] + assert d.default_effort == "high" + + def test_parses_reasoning_schema(self): + nr.register_from_catalog([REASONING_MODEL]) + d = nr.get_descriptor("gpt-5.6-sol") + assert d.schema_type == "reasoning" + assert d.valid_efforts == ["low", "high"] + + def test_models_without_schema_are_absent(self): + nr.register_from_catalog([OPUS_46, NO_SCHEMA_MODEL]) + assert nr.get_descriptor("claude-sonnet-4") is None + assert nr.has_catalog() is True + + def test_refresh_replaces_prior_contents(self): + nr.register_from_catalog([OPUS_46]) + nr.register_from_catalog([REASONING_MODEL]) + assert nr.get_descriptor("claude-opus-4.6") is None + assert nr.get_descriptor("gpt-5.6-sol") is not None + + def test_empty_catalog_clears(self): + nr.register_from_catalog([OPUS_46]) + nr.register_from_catalog([]) + assert nr.has_catalog() is False + + +class TestClampEffort: + EFFORTS = ["low", "medium", "high", "max"] + + def test_exact_match_passthrough(self): + assert nr.clamp_effort("high", self.EFFORTS) == "high" + + def test_xhigh_snaps_up_to_max(self): + # xhigh sits between high and max; tie-break prefers the higher level. + assert nr.clamp_effort("xhigh", self.EFFORTS) == "max" + + def test_minimal_snaps_to_low(self): + assert nr.clamp_effort("minimal", ["medium", "high"]) == "medium" + + def test_unknown_token_falls_back_to_highest(self): + assert nr.clamp_effort("bogus", self.EFFORTS) == "max" + + def test_empty_efforts_returns_none(self): + assert nr.clamp_effort("high", []) is None + + +class TestBuildNativeEffortFields: + def test_catalog_xhigh_clamped_not_rejected(self): + nr.register_from_catalog([OPUS_46]) + fields = build_native_effort_fields("claude-opus-4.6", "xhigh") + assert fields == { + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "max"}, + } + + def test_catalog_reasoning_schema_shape(self): + nr.register_from_catalog([REASONING_MODEL]) + fields = build_native_effort_fields("gpt-5.6-sol", "high") + assert fields == {"reasoning": {"effort": "high"}} + + def test_unsupported_model_returns_none(self): + nr.register_from_catalog([OPUS_46]) + assert build_native_effort_fields("claude-sonnet-4", "high") is None + + def test_no_effort_returns_none(self): + nr.register_from_catalog([OPUS_46]) + assert build_native_effort_fields("claude-opus-4.6", None) is None + + def test_static_fallback_when_no_catalog(self): + # Empty registry -> falls back to the static config table (auto is output_config). + fields = build_native_effort_fields("auto", "high") + assert fields == { + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "high"}, + } diff --git a/tests/unit/test_routes_anthropic.py b/tests/unit/test_routes_anthropic.py index d156f66a..0951d24c 100644 --- a/tests/unit/test_routes_anthropic.py +++ b/tests/unit/test_routes_anthropic.py @@ -1081,13 +1081,12 @@ def test_non_streaming_uses_shared_client( except Exception: pass - print("Checking: KiroHttpClient(shared_client=app.state.http_client)...") + print("Checking: KiroHttpClient(shared_client=None)...") assert mock_kiro_http_client_class.called call_args = mock_kiro_http_client_class.call_args print(f"Call args: {call_args}") - assert call_args[1]['shared_client'] is not None, \ - "Non-streaming should use shared client" - print("✅ Anthropic non-streaming correctly uses shared client") + assert call_args[1]['shared_client'] is None, \ + "All requests now use per-request clients (shared_client=None) for runtime stability" # ============================================================================= diff --git a/tests/unit/test_routes_openai.py b/tests/unit/test_routes_openai.py index 5b2bb4e7..c88d63f1 100644 --- a/tests/unit/test_routes_openai.py +++ b/tests/unit/test_routes_openai.py @@ -970,13 +970,13 @@ def test_non_streaming_uses_shared_client( except Exception: pass - print("Checking: KiroHttpClient(shared_client=app.state.http_client)...") + print("Checking: KiroHttpClient(shared_client=None)...") assert mock_kiro_http_client_class.called call_args = mock_kiro_http_client_class.call_args print(f"Call args: {call_args}") - assert call_args[1]['shared_client'] is not None, \ - "Non-streaming should use shared client" - print("✅ Non-streaming correctly uses shared client") + assert call_args[1]['shared_client'] is None, \ + "All requests now use per-request clients (shared_client=None) for runtime stability" + print("✅ Non-streaming correctly uses per-request client") # ============================================================================= diff --git a/tests/unit/test_streaming_core.py b/tests/unit/test_streaming_core.py index 4813b4ea..7fd2f2ee 100644 --- a/tests/unit/test_streaming_core.py +++ b/tests/unit/test_streaming_core.py @@ -448,22 +448,23 @@ async def test_raises_timeout_on_first_token(self, mock_response): print("Setup: Mock response that times out...") async def mock_aiter_bytes(): + await asyncio.sleep(float('inf')) yield b'chunk' mock_response.aiter_bytes = mock_aiter_bytes - async def mock_wait_for_timeout(*args, **kwargs): - raise asyncio.TimeoutError() - print("Action: Parsing stream with timeout...") - with patch('kiro.streaming_core.asyncio.wait_for', side_effect=mock_wait_for_timeout): + gen = parse_kiro_stream(mock_response, first_token_timeout=0.1) + try: with pytest.raises(FirstTokenTimeoutError) as exc_info: - async for event in parse_kiro_stream(mock_response, first_token_timeout=30): + async for event in gen: pass + finally: + await gen.aclose() print(f"Caught exception: {exc_info.value}") - assert "30" in str(exc_info.value) + assert "0.1" in str(exc_info.value) print("✓ FirstTokenTimeoutError raised on timeout") @pytest.mark.asyncio @@ -480,16 +481,15 @@ async def mock_aiter_bytes(): mock_response.aiter_bytes = mock_aiter_bytes - # Mock wait_for to raise StopAsyncIteration (empty response) - async def mock_wait_for_empty(*args, **kwargs): - raise StopAsyncIteration() - print("Action: Parsing empty stream...") events = [] - with patch('kiro.streaming_core.asyncio.wait_for', side_effect=mock_wait_for_empty): - async for event in parse_kiro_stream(mock_response, first_token_timeout=30): + gen = parse_kiro_stream(mock_response, first_token_timeout=30) + try: + async for event in gen: events.append(event) + finally: + await gen.aclose() print(f"Received {len(events)} events") assert len(events) == 0 @@ -1064,29 +1064,29 @@ async def mock_aiter_bytes(): events = [] with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser): - with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.streaming_core.ThinkingParser') as mock_thinking_parser_class: - mock_thinking_parser = MagicMock() - mock_thinking_parser.feed.return_value = MagicMock( - thinking_content=None, - regular_content="Hello", - is_first_thinking_chunk=False, - is_last_thinking_chunk=False - ) - mock_thinking_parser.finalize.return_value = MagicMock( - thinking_content=None, - regular_content=None, - is_first_thinking_chunk=False, - is_last_thinking_chunk=False - ) - mock_thinking_parser.found_thinking_block = False - mock_thinking_parser_class.return_value = mock_thinking_parser - - async for event in parse_kiro_stream(mock_response, first_token_timeout=30): - events.append(event) - - # Verify ThinkingParser was instantiated - mock_thinking_parser_class.assert_called_once() + with patch('kiro.streaming_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.streaming_core.ThinkingParser') as mock_thinking_parser_class: + mock_thinking_parser = MagicMock() + mock_thinking_parser.feed.return_value = MagicMock( + thinking_content=None, + regular_content="Hello", + is_first_thinking_chunk=False, + is_last_thinking_chunk=False + ) + mock_thinking_parser.finalize.return_value = MagicMock( + thinking_content=None, + regular_content=None, + is_first_thinking_chunk=False, + is_last_thinking_chunk=False + ) + mock_thinking_parser.found_thinking_block = False + mock_thinking_parser_class.return_value = mock_thinking_parser + + async for event in parse_kiro_stream(mock_response, first_token_timeout=30): + events.append(event) + # Verify ThinkingParser was instantiated + mock_thinking_parser_class.assert_called_once() print("✓ Thinking parser enabled when fake reasoning is on") @@ -1169,19 +1169,20 @@ async def test_propagates_first_token_timeout_error(self, mock_response): print("Setup: Mock response that times out...") async def mock_aiter_bytes(): + await asyncio.sleep(float('inf')) yield b'chunk' mock_response.aiter_bytes = mock_aiter_bytes - async def mock_wait_for_timeout(*args, **kwargs): - raise asyncio.TimeoutError() - print("Action: Parsing stream with timeout...") - with patch('kiro.streaming_core.asyncio.wait_for', side_effect=mock_wait_for_timeout): + gen = parse_kiro_stream(mock_response, first_token_timeout=0.1) + try: with pytest.raises(FirstTokenTimeoutError): - async for event in parse_kiro_stream(mock_response, first_token_timeout=30): + async for event in gen: pass + finally: + await gen.aclose() print("✓ FirstTokenTimeoutError propagated correctly") diff --git a/tests/unit/test_thinking_parser.py b/tests/unit/test_thinking_parser.py index 0b18ce35..9e856f24 100644 --- a/tests/unit/test_thinking_parser.py +++ b/tests/unit/test_thinking_parser.py @@ -949,9 +949,10 @@ def test_injects_tags_when_enabled(self): print("Testing tag injection when enabled...") from kiro.converters_core import inject_thinking_tags, ThinkingConfig - with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): - with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): - result = inject_thinking_tags("Hello", ThinkingConfig()) + with patch('kiro.converters_core.NATIVE_REASONING_ENABLED', False): + with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True): + with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000): + result = inject_thinking_tags("Hello", ThinkingConfig()) print(f"Result: '{result}'") assert "enabled" in result