+
+ 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 .
+
+Also add information on how to contact you by electronic and paper mail.
+
+ If your software can interact with users remotely through a computer
+network, you should also make sure that it provides a way for users to
+get its source. For example, if your program is a web application, its
+interface could display a "Source" link that leads users to an archive
+of the code. There are many ways you could offer source, and different
+solutions will be better for different programs; see section 13 for the
+specific requirements.
+
+ You should also get your employer (if you work as a programmer) or school,
+if any, to sign a "copyright disclaimer" for the program, if necessary.
+For more information on this, and how to apply and follow the GNU AGPL, see
+.
diff --git a/kiro-gateway/README.md b/kiro-gateway/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..ba290f03b0ec7ad2f5a5b58e036af84d8d0beae8
--- /dev/null
+++ b/kiro-gateway/README.md
@@ -0,0 +1,626 @@
+
+
+# 👻 Kiro Gateway
+
+**Proxy gateway for Kiro API (Amazon Q Developer / AWS CodeWhisperer)**
+
+[🇷🇺 Русский](docs/ru/README.md) • [🇨🇳 中文](docs/zh/README.md) • [🇪🇸 Español](docs/es/README.md) • [🇮🇩 Indonesia](docs/id/README.md) • [🇧🇷 Português](docs/pt/README.md) • [🇯🇵 日本語](docs/ja/README.md) • [🇰🇷 한국어](docs/ko/README.md)
+
+Made with ❤️ by [@Jwadow](https://github.com/jwadow)
+
+[](https://www.gnu.org/licenses/agpl-3.0)
+[](https://www.python.org/downloads/)
+[](https://fastapi.tiangolo.com/)
+[](#-support-the-project)
+
+*Use Claude models from Kiro with Claude Code, OpenCode, Cursor, Cline, Roo Code, Kilo Code, Obsidian, OpenAI SDK, LangChain, Continue and other OpenAI or Anthropic compatible tools*
+
+[Models](#-supported-models) • [Features](#-features) • [Quick Start](#-quick-start) • [Configuration](#%EF%B8%8F-configuration) • [💖 Sponsor](#-support-the-project)
+
+
+
+---
+
+## 🤖 Available Models
+
+> ⚠️ **Important:** Model availability depends on your Kiro tier (free/paid). The gateway provides access to whatever models are available in your IDE or CLI based on your subscription. The list below shows models commonly available on the **free tier**.
+
+> 🔒 **Claude Opus 4.5** was removed from the free tier on January 17, 2026. It may be available on paid tiers — check your IDE/CLI model list.
+
+🚀 **Claude Sonnet 4.5** — Balanced performance. Great for coding, writing, and general-purpose tasks.
+
+⚡ **Claude Haiku 4.5** — Lightning fast. Perfect for quick responses, simple tasks, and chat.
+
+📦 **Claude Sonnet 4** — Previous generation. Still powerful and reliable for most use cases.
+
+📦 **Claude 3.7 Sonnet** — Legacy model. Available for backward compatibility.
+
+> 💡 **Smart Model Resolution:** Use any model name format — `claude-sonnet-4-5`, `claude-sonnet-4.5`, or even versioned names like `claude-sonnet-4-5-20250929`. The gateway normalizes them automatically.
+
+---
+
+## ✨ Features
+
+| Feature | Description |
+|---------|-------------|
+| 🔌 **OpenAI-compatible API** | Works with any OpenAI-compatible tool |
+| 🔌 **Anthropic-compatible API** | Native `/v1/messages` endpoint |
+| 🌐 **VPN/Proxy Support** | HTTP/SOCKS5 proxy for restricted networks |
+| 🧠 **Extended Thinking** | Reasoning is exclusive to our project |
+| 👁️ **Vision Support** | Send images to model |
+| 🛠️ **Tool Calling** | Supports function calling |
+| 💬 **Full message history** | Passes complete conversation context |
+| 📡 **Streaming** | Full SSE streaming support |
+| 🔄 **Retry Logic** | Automatic retries on errors (403, 429, 5xx) |
+| 📋 **Extended model list** | Including versioned models |
+| 🔐 **Smart token management** | Automatic refresh before expiration |
+
+---
+
+## 🚀 Quick Start
+
+### Prerequisites
+
+- Python 3.10+
+- One of the following:
+ - [Kiro IDE](https://kiro.dev/) with logged in account, OR
+ - [Kiro CLI](https://kiro.dev/cli/) with AWS SSO (AWS IAM Identity Center, OIDC) - free Builder ID or corporate account
+
+### Installation
+
+```bash
+# Clone the repository (requires Git)
+git clone https://github.com/Jwadow/kiro-gateway.git
+cd kiro-gateway
+
+# Or download ZIP: Code → Download ZIP → extract → open kiro-gateway folder
+
+# Install dependencies
+pip install -r requirements.txt
+
+# Configure (see Configuration section)
+cp .env.example .env
+# Copy and edit .env with your credentials
+
+# Start the server
+python main.py
+
+# Or with custom port (if 8000 is busy)
+python main.py --port 9000
+```
+
+The server will be available at `http://localhost:8000`
+
+---
+
+## ⚙️ Configuration
+
+### Option 1: JSON Credentials File (Kiro IDE / Enterprise)
+
+Specify the path to the credentials file:
+
+Works with:
+- **Kiro IDE** (standard) - for personal accounts
+- **Enterprise** - for corporate accounts with SSO
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/kiro-auth-token.json"
+
+# Password to protect YOUR proxy server (make up any secure string)
+# You'll use this as api_key when connecting to your gateway
+PROXY_API_KEY="my-super-secret-password-123"
+```
+
+
+📄 JSON file format
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:...",
+ "region": "us-east-1",
+ "clientIdHash": "abc123..." // Optional: for corporate SSO setups
+}
+```
+
+> **Note:** If you have two JSON files in `~/.aws/sso/cache/` (e.g., `kiro-auth-token.json` and a file with a hash name), use `kiro-auth-token.json` in `KIRO_CREDS_FILE`. The gateway will automatically load the other file.
+
+
+
+### Option 2: Environment Variables (.env file)
+
+Create a `.env` file in the project root:
+
+```env
+# Required
+REFRESH_TOKEN="your_kiro_refresh_token"
+
+# Password to protect YOUR proxy server (make up any secure string)
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Optional
+PROFILE_ARN="arn:aws:codewhisperer:us-east-1:..."
+KIRO_REGION="us-east-1"
+```
+
+### Option 3: AWS SSO Credentials (kiro-cli / Enterprise)
+
+If you use `kiro-cli` or Kiro IDE with AWS SSO (AWS IAM Identity Center), the gateway will automatically detect and use the appropriate authentication.
+
+Works with both free Builder ID accounts and corporate accounts.
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/your-sso-cache-file.json"
+
+# Password to protect YOUR proxy server
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Note: PROFILE_ARN is NOT needed for AWS SSO (Builder ID and corporate accounts)
+# The gateway will work without it
+```
+
+
+📄 AWS SSO JSON file format
+
+AWS SSO credentials files (from `~/.aws/sso/cache/`) contain:
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "region": "us-east-1",
+ "clientId": "...",
+ "clientSecret": "..."
+}
+```
+
+**Note:** AWS SSO (Builder ID and corporate accounts) users do NOT need `profileArn`. The gateway will work without it (if specified, it will be ignored).
+
+
+
+
+🔍 How it works
+
+The gateway automatically detects the authentication type based on the credentials file:
+
+- **Kiro Desktop Auth** (default): Used when `clientId` and `clientSecret` are NOT present
+ - Endpoint: `https://prod.{region}.auth.desktop.kiro.dev/refreshToken`
+
+- **AWS SSO (OIDC)**: Used when `clientId` and `clientSecret` ARE present
+ - Endpoint: `https://oidc.{region}.amazonaws.com/token`
+
+No additional configuration is needed — just point to your credentials file!
+
+
+
+### Option 4: kiro-cli SQLite Database
+
+If you use `kiro-cli` and prefer to use its SQLite database directly:
+
+```env
+KIRO_CLI_DB_FILE="~/.local/share/kiro-cli/data.sqlite3"
+
+# Password to protect YOUR proxy server
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Note: PROFILE_ARN is NOT needed for AWS SSO (Builder ID and corporate accounts)
+# The gateway will work without it
+```
+
+
+📄 Database locations
+
+| CLI Tool | Database Path |
+|----------|---------------|
+| kiro-cli | `~/.local/share/kiro-cli/data.sqlite3` |
+| amazon-q-developer-cli | `~/.local/share/amazon-q/data.sqlite3` |
+
+The gateway reads credentials from the `auth_kv` table which stores:
+- `kirocli:odic:token` or `codewhisperer:odic:token` — access token, refresh token, expiration
+- `kirocli:odic:device-registration` or `codewhisperer:odic:device-registration` — client ID and secret
+
+Both key formats are supported for compatibility with different kiro-cli versions.
+
+
+
+### Getting Credentials
+
+**For Kiro IDE users:**
+- Log in to Kiro IDE and use Option 1 above (JSON credentials file)
+- The credentials file is created automatically after login
+
+**For Kiro CLI users:**
+- Log in with `kiro-cli login` and use Option 3 or Option 4 above
+- No manual token extraction needed!
+
+
+🔧 Advanced: Manual token extraction
+
+If you need to manually extract the refresh token (e.g., for debugging), you can intercept Kiro IDE traffic:
+- Look for requests to: `prod.us-east-1.auth.desktop.kiro.dev/refreshToken`
+
+
+
+---
+
+## 🌐 VPN/Proxy Support
+
+**For users in China, corporate networks, or regions with connectivity issues to AWS services.**
+
+The gateway supports routing all Kiro API requests through a VPN or proxy server. This is essential if you experience connection problems to AWS endpoints or need to use a corporate proxy.
+
+### Configuration
+
+Add to your `.env` file:
+
+```env
+# HTTP proxy
+VPN_PROXY_URL=http://127.0.0.1:7890
+
+# SOCKS5 proxy
+VPN_PROXY_URL=socks5://127.0.0.1:1080
+
+# With authentication (corporate proxies)
+VPN_PROXY_URL=http://username:password@proxy.company.com:8080
+
+# Without protocol (defaults to http://)
+VPN_PROXY_URL=192.168.1.100:8080
+```
+
+### Supported Protocols
+
+- ✅ **HTTP** — Standard proxy protocol
+- ✅ **HTTPS** — Secure proxy connections
+- ✅ **SOCKS5** — Advanced proxy protocol (common in VPN software)
+- ✅ **Authentication** — Username/password embedded in URL
+
+### When You Need This
+
+| Situation | Solution |
+|-----------|----------|
+| Connection timeouts to AWS | Use VPN/proxy to route traffic |
+| Corporate network restrictions | Configure your company's proxy |
+| Regional connectivity issues | Use a VPN service with proxy support |
+| Privacy requirements | Route through your own proxy server |
+
+### Popular VPN Software with Proxy Support
+
+Most VPN clients provide a local proxy server you can use:
+- **Sing-box** — Modern VPN client with HTTP/SOCKS5 proxy
+- **Clash** — Usually runs on `http://127.0.0.1:7890`
+- **V2Ray** — Configurable SOCKS5/HTTP proxy
+- **Shadowsocks** — SOCKS5 proxy support
+- **Corporate VPN** — Check your IT department for proxy settings
+
+Leave `VPN_PROXY_URL` empty (default) if you don't need proxy support.
+
+---
+
+## 📡 API Reference
+
+### Endpoints
+
+| Endpoint | Method | Description |
+|----------|--------|-------------|
+| `/` | GET | Health check |
+| `/health` | GET | Detailed health check |
+| `/v1/models` | GET | List available models |
+| `/v1/chat/completions` | POST | OpenAI Chat Completions API |
+| `/v1/messages` | POST | Anthropic Messages API |
+
+---
+
+## 💡 Usage Examples
+
+### OpenAI API
+
+
+🔹 Simple cURL Request
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello!"}],
+ "stream": true
+ }'
+```
+
+> **Note:** Replace `my-super-secret-password-123` with the `PROXY_API_KEY` you set in your `.env` file.
+
+
+
+
+🔹 Streaming Request
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "What is 2+2?"}
+ ],
+ "stream": true
+ }'
+```
+
+
+
+
+🛠️ With Tool Calling
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "What is the weather in London?"}],
+ "tools": [{
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get weather for a location",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "City name"}
+ },
+ "required": ["location"]
+ }
+ }
+ }]
+ }'
+```
+
+
+
+
+🐍 Python OpenAI SDK
+
+```python
+from openai import OpenAI
+
+client = OpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123" # Your PROXY_API_KEY from .env
+)
+
+response = client.chat.completions.create(
+ model="claude-sonnet-4-5",
+ messages=[
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "Hello!"}
+ ],
+ stream=True
+)
+
+for chunk in response:
+ if chunk.choices[0].delta.content:
+ print(chunk.choices[0].delta.content, end="")
+```
+
+
+
+
+🦜 LangChain
+
+```python
+from langchain_openai import ChatOpenAI
+
+llm = ChatOpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123", # Your PROXY_API_KEY from .env
+ model="claude-sonnet-4-5"
+)
+
+response = llm.invoke("Hello, how are you?")
+print(response.content)
+```
+
+
+
+### Anthropic API
+
+
+🔹 Simple cURL Request
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello!"}]
+ }'
+```
+
+> **Note:** Anthropic API uses `x-api-key` header instead of `Authorization: Bearer`. Both are supported.
+
+
+
+
+🔹 With System Prompt
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "system": "You are a helpful assistant.",
+ "messages": [{"role": "user", "content": "Hello!"}]
+ }'
+```
+
+> **Note:** In Anthropic API, `system` is a separate field, not a message.
+
+
+
+
+📡 Streaming
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "stream": true,
+ "messages": [{"role": "user", "content": "Hello!"}]
+ }'
+```
+
+
+
+
+🐍 Python Anthropic SDK
+
+```python
+import anthropic
+
+client = anthropic.Anthropic(
+ api_key="my-super-secret-password-123", # Your PROXY_API_KEY from .env
+ base_url="http://localhost:8000"
+)
+
+# Non-streaming
+response = client.messages.create(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "Hello!"}]
+)
+print(response.content[0].text)
+
+# Streaming
+with client.messages.stream(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "Hello!"}]
+) as stream:
+ for text in stream.text_stream:
+ print(text, end="", flush=True)
+```
+
+
+
+---
+
+## 🔧 Debugging
+
+Debug logging is **disabled by default**. To enable, add to your `.env`:
+
+```env
+# Debug logging mode:
+# - off: disabled (default)
+# - errors: save logs only for failed requests (4xx, 5xx) - recommended for troubleshooting
+# - all: save logs for every request (overwrites on each request)
+DEBUG_MODE=errors
+```
+
+### Debug Modes
+
+| Mode | Description | Use Case |
+|------|-------------|----------|
+| `off` | Disabled (default) | Production |
+| `errors` | Save logs only for failed requests (4xx, 5xx) | **Recommended for troubleshooting** |
+| `all` | Save logs for every request | Development/debugging |
+
+### Debug Files
+
+When enabled, requests are logged to the `debug_logs/` folder:
+
+| File | Description |
+|------|-------------|
+| `request_body.json` | Incoming request from client (OpenAI format) |
+| `kiro_request_body.json` | Request sent to Kiro API |
+| `response_stream_raw.txt` | Raw stream from Kiro |
+| `response_stream_modified.txt` | Transformed stream (OpenAI format) |
+| `app_logs.txt` | Application logs for the request |
+| `error_info.json` | Error details (only on errors) |
+
+---
+
+## 📜 License
+
+This project is licensed under the **GNU Affero General Public License v3.0 (AGPL-3.0)**.
+
+This means:
+- ✅ You can use, modify, and distribute this software
+- ✅ You can use it for commercial purposes
+- ⚠️ **You must disclose source code** when you distribute the software
+- ⚠️ **Network use is distribution** — if you run a modified version on a server and let others interact with it, you must make the source code available to them
+- ⚠️ Modifications must be released under the same license
+
+See the [LICENSE](LICENSE) file for the full license text.
+
+### Why AGPL-3.0?
+
+AGPL-3.0 ensures that improvements to this software benefit the entire community. If you modify this gateway and deploy it as a service, you must share your improvements with your users.
+
+### Contributor License Agreement (CLA)
+
+By submitting a contribution to this project, you agree to the terms of our [Contributor License Agreement (CLA)](CLA.md). This ensures that:
+- You have the right to submit the contribution
+- You grant the maintainer rights to use and relicense your contribution
+- The project remains legally protected
+
+---
+
+## 💖 Support the Project
+
+
+
+
+
+**If this project saved you time or money, consider supporting it!**
+
+Every contribution helps keep this project alive and growing
+
+
+
+### 🤑 Donate
+
+[**☕ One-time Donation**](https://app.lava.top/jwadow?tabId=donate) • [**💎 Monthly Support**](https://app.lava.top/jwadow?tabId=subscriptions)
+
+
+
+### 🪙 Or send crypto
+
+| Currency | Network | Address |
+|:--------:|:-------:|:--------|
+| **USDT** | TRC20 | `TSVtgRc9pkC1UgcbVeijBHjFmpkYHDRu26` |
+| **BTC** | Bitcoin | `12GZqxqpcBsqJ4Vf1YreLqwoMGvzBPgJq6` |
+| **ETH** | Ethereum | `0xc86eab3bba3bbaf4eb5b5fff8586f1460f1fd395` |
+| **SOL** | Solana | `9amykF7KibZmdaw66a1oqYJyi75fRqgdsqnG66AK3jvh` |
+| **TON** | TON | `UQBVh8T1H3GI7gd7b-_PPNnxHYYxptrcCVf3qQk5v41h3QTM` |
+
+
+
+**[⬆ Back to Top](#-kiro-gateway)**
+
+
diff --git a/kiro-gateway/docs/en/ARCHITECTURE.md b/kiro-gateway/docs/en/ARCHITECTURE.md
new file mode 100644
index 0000000000000000000000000000000000000000..1c5d4b1ee3d75a7a0f7d5304a887d3d20508726e
--- /dev/null
+++ b/kiro-gateway/docs/en/ARCHITECTURE.md
@@ -0,0 +1,821 @@
+# Architectural Overview: Kiro Gateway
+
+## 1. System Purpose and Goals
+
+The project is a high-level proxy gateway implementing the **"Adapter"** structural design pattern.
+
+The main goal of the system is to provide transparent compatibility between multiple heterogeneous interfaces:
+
+### Supported API Formats
+
+| API | Endpoints | Status |
+|-----|-----------|--------|
+| **OpenAI** | `/v1/models`, `/v1/chat/completions` | ✅ Supported |
+| **Anthropic** | `/v1/messages` | ✅ Supported |
+
+### Architectural Model
+
+```
+┌─────────────────────────────────────────────────────────────────┐
+│ Clients │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │ OpenAI SDK/Tools │ │ Anthropic SDK/Tools │ │
+│ │ (Cursor, Cline, │ │ (Claude Code, │ │
+│ │ Continue, etc.) │ │ Anthropic SDK) │ │
+│ └──────────┬──────────┘ └──────────┬──────────┘ │
+└─────────────┼──────────────────────────────┼───────────────────┘
+ │ │
+ ▼ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ Kiro Gateway │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │ OpenAI Adapter │ │ Anthropic Adapter │ │
+│ │ /v1/chat/... │ │ /v1/messages │ │
+│ └──────────┬──────────┘ └──────────┬──────────┘ │
+│ └──────────────┬───────────────┘ │
+│ ▼ │
+│ ┌─────────────────────────────┐ │
+│ │ Core Layer │ │
+│ │ (Shared conversion logic) │ │
+│ └──────────────┬──────────────┘ │
+└────────────────────────────┼────────────────────────────────────┘
+ │
+ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ Kiro API │
+│ (AWS CodeWhisperer Backend) │
+└─────────────────────────────────────────────────────────────────┘
+```
+
+The system acts as a "translator", allowing the use of any tools, libraries, and IDE plugins developed for OpenAI and Anthropic ecosystems with Claude models through the Kiro API.
+
+**Both APIs work simultaneously** on the same server without any configuration switching.
+
+## 2. Project Structure
+
+The project is organized as a modular Python package `kiro/`:
+
+```
+kiro-gateway/
+├── main.py # Entry point, FastAPI application creation
+├── requirements.txt # Python dependencies
+├── .env.example # Environment configuration example
+│
+├── kiro/ # Main package
+│ ├── __init__.py # Package exports, version
+│ │
+│ │ # ═══════════════════════════════════════════════════════
+│ │ # SHARED LAYER - Reused by all APIs
+│ │ # ═══════════════════════════════════════════════════════
+│ ├── config.py # Configuration and constants
+│ ├── auth.py # KiroAuthManager - token management
+│ ├── cache.py # ModelInfoCache - model cache
+│ ├── http_client.py # HTTP client with retry logic
+│ ├── parsers.py # AWS SSE stream parsers
+│ ├── utils.py # Helper utilities
+│ ├── tokenizer.py # Token counting (tiktoken)
+│ ├── debug_logger.py # Debug request logging
+│ ├── exceptions.py # Exception handlers
+│ ├── thinking_parser.py # Thinking blocks parser
+│ │
+│ │ # ═══════════════════════════════════════════════════════
+│ │ # CORE LAYER - Shared core for all APIs
+│ │ # ═══════════════════════════════════════════════════════
+│ ├── converters_core.py # Shared Kiro payload building logic
+│ ├── streaming_core.py # Shared Kiro stream parsing logic
+│ │
+│ │ # ═══════════════════════════════════════════════════════
+│ │ # OPENAI API LAYER
+│ │ # ═══════════════════════════════════════════════════════
+│ ├── models_openai.py # Pydantic models for OpenAI API
+│ ├── converters_openai.py # OpenAI → Kiro adapter
+│ ├── routes_openai.py # FastAPI routes for OpenAI
+│ ├── streaming_openai.py # Kiro → OpenAI SSE formatter
+│ │
+│ │ # ═══════════════════════════════════════════════════════
+│ │ # ANTHROPIC API LAYER
+│ │ # ═══════════════════════════════════════════════════════
+│ ├── models_anthropic.py # Pydantic models for Anthropic API
+│ ├── converters_anthropic.py # Anthropic → Kiro adapter
+│ ├── routes_anthropic.py # FastAPI routes for Anthropic
+│ └── streaming_anthropic.py # Kiro → Anthropic SSE formatter
+│
+├── tests/ # Tests
+│ ├── conftest.py # Pytest fixtures
+│ ├── unit/ # Unit tests
+│ └── integration/ # Integration tests
+│
+├── docs/ # Documentation
+│ ├── ru/ # Russian version
+│ └── en/ # English version
+│
+└── debug_logs/ # Debug logs (generated when DEBUG_LAST_REQUEST=true)
+```
+
+### Organization Principle: Shared Core + Thin Adapters
+
+The architecture is built on the principle of **maximum code reuse**:
+
+| Layer | Purpose | Files |
+|-------|---------|-------|
+| **Shared Layer** | Infrastructure independent of API format | `auth.py`, `http_client.py`, `cache.py`, `parsers.py`, `tokenizer.py` |
+| **Core Layer** | Shared business logic for conversion | `converters_core.py`, `streaming_core.py` |
+| **API Layer** | Thin adapters for specific formats | `*_openai.py`, `*_anthropic.py` |
+
+## 3. Architectural Topology and Components
+
+The system is built on the asynchronous `FastAPI` framework and uses an event-driven lifecycle management model (`Lifespan Events`).
+
+### 3.1. Entry Point (`main.py`)
+
+The `main.py` file is responsible for:
+
+1. **Logging configuration** — Loguru setup with colored output
+2. **Configuration validation** — `validate_configuration()` function checks:
+ - Presence of `.env` file
+ - Presence of credentials (REFRESH_TOKEN or KIRO_CREDS_FILE)
+3. **Lifespan Manager** — creation and initialization of:
+ - `KiroAuthManager` for token management
+ - `ModelInfoCache` for model caching
+4. **Error handler registration** — `validation_exception_handler` for 422 errors
+5. **Route connection** — `app.include_router(router)`
+
+### 3.2. Configuration Module (`kiro/config.py`)
+
+Centralized storage of all settings:
+
+| Parameter | Description | Default Value |
+|-----------|-------------|---------------|
+| `PROXY_API_KEY` | API key for proxy access | `changeme_proxy_secret` |
+| `REFRESH_TOKEN` | Kiro refresh token | from `.env` |
+| `PROFILE_ARN` | AWS CodeWhisperer profile ARN | from `.env` |
+| `REGION` | AWS region | `us-east-1` |
+| `KIRO_CREDS_FILE` | Path to JSON credentials file | from `.env` |
+| `TOKEN_REFRESH_THRESHOLD` | Time before token refresh | 600 sec (10 min) |
+| `MAX_RETRIES` | Max retry attempts | 3 |
+| `BASE_RETRY_DELAY` | Base retry delay | 1.0 sec |
+| `MODEL_CACHE_TTL` | Model cache TTL | 3600 sec (1 hour) |
+| `DEFAULT_MAX_INPUT_TOKENS` | Default max input tokens | 200000 |
+| `TOOL_DESCRIPTION_MAX_LENGTH` | Max tool description length | 10000 characters |
+| `DEBUG_LAST_REQUEST` | Enable debug logging | `false` |
+| `DEBUG_DIR` | Debug logs directory | `debug_logs` |
+| `APP_VERSION` | Application version | `0.0.0` |
+
+**Helper functions:**
+- `get_kiro_refresh_url(region)` — URL for token refresh
+- `get_kiro_api_host(region)` — main API host
+- `get_kiro_q_host(region)` — Q API host
+- `get_internal_model_id(external_model)` — model name conversion
+
+### 3.3. Pydantic Models (`kiro/models_openai.py`)
+
+#### Models for `/v1/models`
+
+| Model | Description |
+|-------|-------------|
+| `OpenAIModel` | AI model description (id, object, created, owned_by) |
+| `ModelList` | Model list for endpoint response |
+
+#### Models for `/v1/chat/completions`
+
+| Model | Description |
+|-------|-------------|
+| `ChatMessage` | Chat message (role, content, tool_calls, tool_call_id) |
+| `ToolFunction` | Tool function description (name, description, parameters) |
+| `Tool` | OpenAI format tool (type, function) |
+| `ChatCompletionRequest` | Generation request (model, messages, stream, tools, ...) |
+
+#### Response Models
+
+| Model | Description |
+|-------|-------------|
+| `ChatCompletionChoice` | Single response variant |
+| `ChatCompletionUsage` | Token information (prompt_tokens, completion_tokens, credits_used) |
+| `ChatCompletionResponse` | Full response (non-streaming) |
+| `ChatCompletionChunk` | Streaming chunk |
+| `ChatCompletionChunkDelta` | Delta changes in chunk |
+| `ChatCompletionChunkChoice` | Variant in streaming chunk |
+
+### 3.4. State Management Layer
+
+#### KiroAuthManager (`kiro/auth.py`)
+
+**Role:** Stateful singleton encapsulating Kiro token management logic.
+
+**Capabilities:**
+- Loading credentials from `.env` or JSON file
+- Support for `expiresAt` to check token expiration time
+- Automatic token refresh 10 minutes before expiration
+- Saving updated tokens back to JSON file
+- Support for different AWS regions
+- Unique fingerprint generation for User-Agent
+
+**Concurrency Control:** Uses `asyncio.Lock` to protect against race conditions.
+
+**Main methods:**
+- `get_access_token()` — returns valid token, refreshing if necessary
+- `force_refresh()` — forced token refresh (on 403)
+- `is_token_expiring_soon()` — expiration time check
+
+**Properties:**
+- `profile_arn` — profile ARN
+- `region` — AWS region
+- `api_host` — API host for region
+- `q_host` — Q API host for region
+- `fingerprint` — unique machine fingerprint
+
+```python
+# Usage example
+auth_manager = KiroAuthManager(
+ refresh_token="your_token",
+ region="us-east-1",
+ creds_file="~/.aws/sso/cache/kiro-auth-token.json"
+)
+token = await auth_manager.get_access_token()
+```
+
+#### ModelInfoCache (`kiro/cache.py`)
+
+**Role:** Thread-safe storage for model configurations.
+
+**Population Strategy:**
+- Lazy Loading via `/ListAvailableModels`
+- Cache TTL: 1 hour
+- Fallback to static model list
+
+**Main methods:**
+- `update(models_data)` — cache update
+- `get(model_id)` — get model information
+- `get_max_input_tokens(model_id)` — get token limit
+- `is_empty()` / `is_stale()` — cache state check
+- `get_all_model_ids()` — list of all model IDs
+
+### 3.5. Helper Utilities (`kiro/utils.py`)
+
+| Function | Description |
+|----------|-------------|
+| `get_machine_fingerprint()` | SHA256 hash of `{hostname}-{username}-kiro-gateway` |
+| `get_kiro_headers(auth_manager, token)` | Form headers for Kiro API |
+| `generate_completion_id()` | ID in format `chatcmpl-{uuid_hex}` |
+| `generate_conversation_id()` | UUID for conversation |
+| `generate_tool_call_id()` | ID in format `call_{uuid_hex[:8]}` |
+
+### 3.6. Conversion Layer (`kiro/converters_openai.py`)
+
+#### Message Conversion
+
+OpenAI messages are transformed into Kiro conversationState:
+
+1. **System prompt** — added to the first user message
+2. **Message history** — fully passed in `history` array
+3. **Adjacent message merging** — messages with the same role are merged
+4. **Tool calls** — OpenAI tools format support
+5. **Tool results** — correct transmission of tool call results
+
+#### Long Tool Description Handling
+
+**Problem:** Kiro API returns error 400 for too long descriptions in `toolSpecification.description`.
+
+**Solution:** Tool Documentation Reference Pattern
+- If `description ≤ TOOL_DESCRIPTION_MAX_LENGTH` → leave as is
+- If `description > TOOL_DESCRIPTION_MAX_LENGTH`:
+ * In `toolSpecification.description` → reference: `"[Full documentation in system prompt under '## Tool: {name}']"`
+ * In system prompt, section `"## Tool: {name}"` with full description is added
+
+**Function:** `process_tools_with_long_descriptions(tools)` → `(processed_tools, tool_documentation)`
+
+#### Main Functions
+
+| Function | Description |
+|----------|-------------|
+| `extract_text_content(content)` | Extract text from various formats |
+| `merge_adjacent_messages(messages)` | Merge adjacent messages with same role |
+| `build_kiro_history(messages, model_id)` | Build history array for Kiro |
+| `build_kiro_payload(request_data, conversation_id, profile_arn)` | Full payload for request |
+
+#### Model Mapping
+
+External model names are converted to internal Kiro IDs:
+
+| External Name | Internal Kiro ID |
+|---------------|------------------|
+| `claude-opus-4-5` | `claude-opus-4.5` |
+| `claude-opus-4-5-20251101` | `claude-opus-4.5` |
+| `claude-haiku-4-5` | `claude-haiku-4.5` |
+| `claude-haiku-4.5` | `claude-haiku-4.5` (direct passthrough) |
+| `claude-sonnet-4-5` | `CLAUDE_SONNET_4_5_20250929_V1_0` |
+| `claude-sonnet-4-5-20250929` | `CLAUDE_SONNET_4_5_20250929_V1_0` |
+| `claude-sonnet-4` | `CLAUDE_SONNET_4_20250514_V1_0` |
+| `claude-sonnet-4-20250514` | `CLAUDE_SONNET_4_20250514_V1_0` |
+| `claude-3-7-sonnet-20250219` | `CLAUDE_3_7_SONNET_20250219_V1_0` |
+| `auto` | `claude-sonnet-4.5` (alias) |
+
+### 3.7. Parsing Layer (`kiro/parsers.py`)
+
+#### AwsEventStreamParser
+
+Advanced AWS SSE format parser with support for:
+
+- **Bracket counting** — correct parsing of nested JSON objects
+- **Content deduplication** — filtering of duplicate events
+- **Tool calls** — parsing of structured and bracket-style tool calls
+- **Escape sequences** — decoding of `\n` and others
+
+#### Event Types
+
+| Event | Description |
+|-------|-------------|
+| `content` | Text content of the response |
+| `tool_start` | Start of tool call (name, toolUseId) |
+| `tool_input` | Continuation of input for tool call |
+| `tool_stop` | End of tool call |
+| `usage` | Credit consumption information |
+| `context_usage` | Context usage percentage |
+
+#### Helper Functions
+
+| Function | Description |
+|----------|-------------|
+| `find_matching_brace(text, start_pos)` | Find closing brace with nesting support |
+| `parse_bracket_tool_calls(response_text)` | Parse `[Called func with args: {...}]` |
+| `deduplicate_tool_calls(tool_calls)` | Remove duplicate tool calls |
+
+### 3.8. Streaming (`kiro/streaming_openai.py`)
+
+#### stream_kiro_to_openai
+
+Async generator for transforming Kiro stream to OpenAI format.
+
+**Functionality:**
+- Parse AWS SSE stream via `AwsEventStreamParser`
+- Form OpenAI `chat.completion.chunk`
+- Handle tool calls (structured and bracket-style)
+- Calculate usage based on `contextUsagePercentage`
+- Debug logging via `debug_logger`
+
+#### collect_stream_response
+
+Collects full response from streaming for non-streaming mode.
+
+### 3.9. HTTP Client (`kiro/http_client.py`)
+
+#### KiroHttpClient
+
+Automatic error handling with exponential backoff:
+
+| Error Code | Action |
+|------------|--------|
+| `403` | Token refresh via `force_refresh()` + retry |
+| `429` | Exponential backoff: `BASE_RETRY_DELAY * (2 ** attempt)` |
+| `5xx` | Exponential backoff (up to MAX_RETRIES attempts) |
+| Timeout | Exponential backoff |
+
+**Delay formula:** `1s, 2s, 4s` (with `BASE_RETRY_DELAY=1.0`)
+
+**Methods:**
+- `request_with_retry(method, url, json_data, stream)` — request with retry
+- `close()` — close client
+
+Supports async context manager (`async with`).
+
+### 3.10. Routes (`kiro/routes_openai.py`)
+
+| Endpoint | Method | Description |
+|----------|--------|-------------|
+| `/` | GET | Health check (status, message, version) |
+| `/health` | GET | Detailed health check (status, timestamp, version) |
+| `/v1/models` | GET | List of available models (requires API key) |
+| `/v1/chat/completions` | POST | Chat completions (requires API key) |
+
+**Authentication:** Bearer token in `Authorization` header
+
+### 3.11. Exception Handling (`kiro/exceptions.py`)
+
+| Function | Description |
+|----------|-------------|
+| `sanitize_validation_errors(errors)` | Convert bytes to strings for JSON serialization |
+| `validation_exception_handler(request, exc)` | Pydantic validation error handler (422) |
+
+### 3.12. Debug Logging (`kiro/debug_logger.py`)
+
+**Class:** `DebugLogger` (singleton)
+
+**Activation:** `DEBUG_LAST_REQUEST=true` in `.env`
+
+**Methods:**
+| Method | Description |
+|--------|-------------|
+| `prepare_new_request()` | Clear directory for new request |
+| `log_request_body(body)` | Save incoming request |
+| `log_kiro_request_body(body)` | Save request to Kiro API |
+| `log_raw_chunk(chunk)` | Append raw chunk from Kiro |
+| `log_modified_chunk(chunk)` | Append transformed chunk |
+
+**Files in `debug_logs/`:**
+- `request_body.json` — incoming request (OpenAI format)
+- `kiro_request_body.json` — request to Kiro API
+- `response_stream_raw.txt` — raw stream from Kiro
+- `response_stream_modified.txt` — transformed stream (OpenAI format)
+
+### 3.13. Tokenizer (`kiro/tokenizer.py`)
+
+**Problem:** Kiro API does not return token counts directly. Instead, the API only provides `context_usage_percentage` — the percentage of model context usage.
+
+**Solution:** Tokenizer module based on `tiktoken` (OpenAI's Rust library) for fast token counting.
+
+**Features:**
+- Uses `cl100k_base` encoding (GPT-4), close to Claude tokenization
+- Correction factor `CLAUDE_CORRECTION_FACTOR = 1.15` for improved accuracy
+- Lazy initialization for faster imports
+- Fallback to rough estimation if tiktoken is unavailable
+
+**Token calculation formula in response:**
+```
+total_tokens = context_usage_percentage × max_input_tokens (from Kiro API)
+completion_tokens = tiktoken(response) (our calculation)
+prompt_tokens = total_tokens - completion_tokens (subtraction)
+```
+
+**Main functions:**
+
+| Function | Description |
+|----------|-------------|
+| `count_tokens(text)` | Count tokens in text |
+| `count_message_tokens(messages)` | Count tokens in message list |
+| `count_tools_tokens(tools)` | Count tokens in tool definitions |
+| `estimate_request_tokens(messages, tools)` | Full request token estimation |
+
+**Debug log:**
+```
+[Usage] claude-opus-4-5: prompt_tokens=142211 (subtraction), completion_tokens=769 (tiktoken), total_tokens=142980 (API Kiro)
+```
+
+**Accuracy:** ~97-99.7% compared to API data.
+
+### 3.14. Kiro API Endpoints
+
+All URLs are dynamically formed based on the region:
+
+* **Token Refresh:** `POST https://prod.{region}.auth.desktop.kiro.dev/refreshToken`
+* **List Models:** `GET https://q.{region}.amazonaws.com/ListAvailableModels`
+* **Generate Response:** `POST https://codewhisperer.{region}.amazonaws.com/generateAssistantResponse`
+
+## 4. Detailed Data Flow
+
+### 4.1 Multi-API Overview
+
+```
+┌─────────────────────────────────────────────────────────────────┐
+│ CLIENTS │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │ OpenAI Client │ │ Anthropic Client │ │
+│ └──────────┬──────────┘ └──────────┬──────────┘ │
+└─────────────┼──────────────────────────────┼───────────────────┘
+ │ │
+ │ POST /v1/chat/completions │ POST /v1/messages
+ ▼ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ API LAYER │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │ routes_openai.py │ │ routes_anthropic.py │ │
+│ │ Security Gate │ │ Security Gate │ │
+│ └──────────┬──────────┘ └──────────┬──────────┘ │
+│ │ │ │
+│ ▼ ▼ │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │converters_openai.py │ │converters_anthropic │ │
+│ │ Extract system │ │ System already │ │
+│ │ from messages │ │ separate in request │ │
+│ └──────────┬──────────┘ └──────────┬──────────┘ │
+└─────────────┼──────────────────────────────┼───────────────────┘
+ │ │
+ └──────────────┬───────────────┘
+ │
+ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ CORE LAYER │
+│ ┌─────────────────────────────┐ │
+│ │ converters_core.py │ │
+│ │ build_kiro_payload() │ │
+│ │ build_kiro_history() │ │
+│ │ process_tools() │ │
+│ └──────────────┬──────────────┘ │
+└────────────────────────────┼────────────────────────────────────┘
+ │
+ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ SHARED LAYER │
+│ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │
+│ │ KiroAuthManager │ │ KiroHttpClient │ │ ModelInfoCache │ │
+│ │ (auth.py) │ │(http_client.py) │ │ (cache.py) │ │
+│ └────────┬────────┘ └────────┬────────┘ └─────────────────┘ │
+└───────────┼────────────────────┼────────────────────────────────┘
+ │ │
+ │ │ POST /generateAssistantResponse
+ │ ▼
+ │ ┌─────────────────────────────────────┐
+ │ │ Kiro API │
+ │ └──────────────────┬──────────────────────┘
+ │ │
+ │ │ AWS SSE Stream
+ │ ▼
+┌───────────┼────────────────────────────────────────────────────┐
+│ │ CORE LAYER │
+│ │ ┌─────────────────────────────┐ │
+│ │ │ streaming_core.py │ │
+│ │ │ parse_kiro_stream() │ │
+│ │ │ → KiroEvent objects │ │
+│ │ └──────────────┬──────────────┘ │
+└────────────────────────────┼───────────────────────────────────┘
+ │
+ ┌──────────────┴───────────────┐
+ │ │
+ ▼ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ OUTPUT LAYER │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │streaming_openai.py │ │streaming_anthropic │ │
+│ │ format_openai_sse() │ │format_anthropic_sse │ │
+│ │ │ │ │ │
+│ │ data: {...} │ │ event: type │ │
+│ │ data: [DONE] │ │ data: {...} │ │
+│ └──────────┬──────────┘ └──────────┬──────────┘ │
+└─────────────┼──────────────────────────────┼───────────────────┘
+ │ │
+ ▼ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ CLIENTS │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │ OpenAI Client │ │ Anthropic Client │ │
+│ └─────────────────────┘ └─────────────────────┘ │
+└─────────────────────────────────┘
+```
+
+### 4.2 OpenAI API Flow
+
+```
+OpenAI Client
+ │ POST /v1/chat/completions
+ ▼
+routes_openai.py ──► converters_openai.py ──► converters_core.py
+ │ │
+ │ ▼
+ │ Kiro Payload
+ │ │
+ ▼ ▼
+KiroAuthManager ──────────────────────────► KiroHttpClient
+ │
+ ▼
+ Kiro API
+ │
+ ▼
+streaming_core.py ◄─────────────────────── AWS SSE Stream
+ │
+ ▼
+streaming_openai.py
+ │
+ ▼
+OpenAI SSE Format ──────────────────────► OpenAI Client
+```
+
+### 4.3 Anthropic API Flow
+
+```
+Anthropic Client
+ │ POST /v1/messages
+ ▼
+routes_anthropic.py ──► converters_anthropic.py ──► converters_core.py
+ │ │
+ │ ▼
+ │ Kiro Payload
+ │ │
+ ▼ ▼
+KiroAuthManager ──────────────────────────────────► KiroHttpClient
+ │
+ ▼
+ Kiro API
+ │
+ ▼
+streaming_core.py ◄─────────────────────────────── AWS SSE Stream
+ │
+ ▼
+streaming_anthropic.py
+ │
+ ▼
+Anthropic SSE Format ──────────────────────────► Anthropic Client
+```
+
+## 5. Available Models
+
+| Model | Description | Credits |
+|-------|-------------|---------|
+| `claude-opus-4-5` | Top-tier model | ~2.2 |
+| `claude-opus-4-5-20251101` | Top-tier model (version) | ~2.2 |
+| `claude-sonnet-4-5` | Enhanced model | ~1.3 |
+| `claude-sonnet-4-5-20250929` | Enhanced model (version) | ~1.3 |
+| `claude-sonnet-4` | Balanced model | ~1.3 |
+| `claude-sonnet-4-20250514` | Balanced (version) | ~1.3 |
+| `claude-haiku-4-5` | Fast model | ~0.4 |
+| `claude-3-7-sonnet-20250219` | Legacy model | ~1.0 |
+
+## 6. Configuration
+
+### Environment Variables (.env)
+
+```env
+# Required
+REFRESH_TOKEN="your_kiro_refresh_token"
+PROXY_API_KEY="your_proxy_secret"
+
+# Optional
+PROFILE_ARN="arn:aws:codewhisperer:..."
+KIRO_REGION="us-east-1"
+KIRO_CREDS_FILE="~/.aws/sso/cache/kiro-auth-token.json"
+
+# Debug
+DEBUG_LAST_REQUEST="false"
+DEBUG_DIR="debug_logs"
+
+# Limits
+TOOL_DESCRIPTION_MAX_LENGTH="10000"
+```
+
+### JSON Credentials File (optional)
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:...",
+ "region": "us-east-1"
+}
+```
+
+## 7. API Endpoints
+
+### 7.1 Common Endpoints
+
+| Endpoint | Method | Description |
+|----------|--------|-------------|
+| `/` | GET | Health check |
+| `/health` | GET | Detailed health check |
+
+### 7.2 OpenAI-compatible Endpoints
+
+| Endpoint | Method | Description |
+|----------|--------|-------------|
+| `/v1/models` | GET | List of available models |
+| `/v1/chat/completions` | POST | Chat completions (streaming/non-streaming) |
+
+**Authentication:** `Authorization: Bearer {PROXY_API_KEY}`
+
+### 7.3 Anthropic-compatible Endpoints
+
+| Endpoint | Method | Description |
+|----------|--------|-------------|
+| `/v1/messages` | POST | Messages API (streaming/non-streaming) |
+
+**Authentication:** `x-api-key: {PROXY_API_KEY}` + `anthropic-version: 2023-06-01`
+
+### 7.4 Format Comparison
+
+| Aspect | OpenAI | Anthropic |
+|--------|--------|-----------|
+| System prompt | In `messages` with `role: "system"` | Separate `system` field |
+| Content | String or array | Always array of content blocks |
+| Stop reason | `finish_reason: "stop"` | `stop_reason: "end_turn"` |
+| Usage | `prompt_tokens`, `completion_tokens` | `input_tokens`, `output_tokens` |
+| Streaming | `data: {...}\n\n` + `data: [DONE]` | `event: type\ndata: {...}\n\n` |
+| Tool format | `{type: "function", function: {...}}` | `{name: "...", input_schema: {...}}` |
+
+## 8. Implementation Features
+
+### Tool Calling
+
+Support for OpenAI-compatible tools format:
+
+```json
+{
+ "tools": [{
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get weather for a location",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string"}
+ }
+ }
+ }
+ }]
+}
+```
+
+### Streaming
+
+Full SSE streaming support with correct OpenAI format:
+
+```
+data: {"id":"chatcmpl-...","object":"chat.completion.chunk",...}
+
+data: [DONE]
+```
+
+### Debugging
+
+When `DEBUG_LAST_REQUEST=true`, all requests and responses are logged in `debug_logs/`:
+- `request_body.json` — incoming request
+- `kiro_request_body.json` — request to Kiro API
+- `response_stream_raw.txt` — raw stream from Kiro
+- `response_stream_modified.txt` — transformed stream
+
+## 9. Extensibility
+
+### Adding a New API Format
+
+The modular architecture allows easy addition of support for other API formats. Thanks to the Core Layer, most of the logic is already implemented.
+
+#### Steps to Add a New Format (e.g., Gemini)
+
+1. **Create models** — `models_gemini.py`
+ ```python
+ class GeminiRequest(BaseModel):
+ """Pydantic model for Gemini request."""
+ contents: List[GeminiContent]
+ ...
+ ```
+
+2. **Create conversion adapter** — `converters_gemini.py`
+ ```python
+ from kiro.converters_core import build_kiro_payload
+
+ def gemini_to_kiro(request: GeminiRequest, ...) -> dict:
+ """Converts Gemini request to Kiro payload."""
+ # Extract data from Gemini format
+ system_prompt = extract_system_instruction(request)
+ messages = convert_gemini_contents(request.contents)
+ tools = convert_gemini_tools(request.tools)
+
+ # Use shared core
+ return build_kiro_payload(
+ messages=messages,
+ system_prompt=system_prompt,
+ tools=tools,
+ ...
+ )
+ ```
+
+3. **Create streaming formatter** — `streaming_gemini.py`
+ ```python
+ from kiro.streaming_core import parse_kiro_stream
+
+ async def stream_to_gemini(response, ...) -> AsyncGenerator[str, None]:
+ """Formats Kiro events to Gemini SSE."""
+ async for event in parse_kiro_stream(response):
+ yield format_gemini_chunk(event)
+ ```
+
+4. **Create routes** — `routes_gemini.py`
+ ```python
+ router = APIRouter()
+
+ @router.post("/v1beta/models/{model}:generateContent")
+ async def generate_content(request: GeminiRequest):
+ ...
+ ```
+
+5. **Connect in main.py**
+ ```python
+ from kiro.routes_gemini import router as gemini_router
+ app.include_router(gemini_router)
+ ```
+
+### What Gets Reused Automatically
+
+When adding a new format, the following components work out of the box:
+
+| Component | Functionality |
+|-----------|---------------|
+| `auth.py` | Kiro token management |
+| `http_client.py` | HTTP with retry logic |
+| `cache.py` | Model cache |
+| `parsers.py` | AWS SSE parsing |
+| `tokenizer.py` | Token counting |
+| `converters_core.py` | Kiro payload building |
+| `streaming_core.py` | Kiro stream parsing |
+
+## 10. Dependencies
+
+Main project dependencies (from `requirements.txt`):
+
+| Package | Purpose |
+|---------|---------|
+| `fastapi` | Asynchronous web framework |
+| `uvicorn` | ASGI server |
+| `httpx` | Asynchronous HTTP client |
+| `pydantic` | Data validation and models |
+| `python-dotenv` | Environment variable loading |
+| `loguru` | Advanced logging |
+| `tiktoken` | Fast token counting |
diff --git a/kiro-gateway/docs/es/README.md b/kiro-gateway/docs/es/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..1bab382986a16d5e1171c9b15f5d35f35c41c2b4
--- /dev/null
+++ b/kiro-gateway/docs/es/README.md
@@ -0,0 +1,626 @@
+
+
+# 👻 Kiro Gateway
+
+**Gateway proxy para Kiro API (Amazon Q Developer / AWS CodeWhisperer)**
+
+[🇬🇧 English](../../README.md) • [🇷🇺 Русский](../ru/README.md) • [🇨🇳 中文](../zh/README.md) • [🇮🇩 Indonesia](../id/README.md) • [🇧🇷 Português](../pt/README.md) • [🇯🇵 日本語](../ja/README.md) • [🇰🇷 한국어](../ko/README.md)
+
+Hecho con ❤️ por [@Jwadow](https://github.com/jwadow)
+
+[](https://www.gnu.org/licenses/agpl-3.0)
+[](https://www.python.org/downloads/)
+[](https://fastapi.tiangolo.com/)
+[](#-apoya-el-proyecto)
+
+*Usa modelos Claude de Kiro con Claude Code, OpenCode, Cursor, Cline, Roo Code, Kilo Code, Obsidian, OpenAI SDK, LangChain, Continue y otras herramientas compatibles con OpenAI o Anthropic*
+
+[Modelos](#-modelos-soportados) • [Características](#-características) • [Inicio Rápido](#-inicio-rápido) • [Configuración](#%EF%B8%8F-configuración) • [💖 Apoyar](#-apoya-el-proyecto)
+
+
+
+---
+
+## 🤖 Modelos Disponibles
+
+> ⚠️ **Importante:** La disponibilidad de modelos depende de tu plan de Kiro (gratuito/pago). El gateway proporciona acceso a los modelos disponibles en tu IDE o CLI según tu suscripción. La lista a continuación muestra los modelos comúnmente disponibles en el **plan gratuito**.
+
+> 🔒 **Claude Opus 4.5** fue eliminado del plan gratuito el 17 de enero de 2026. Puede estar disponible en planes de pago — verifica la lista de modelos en tu IDE/CLI.
+
+🚀 **Claude Sonnet 4.5** — Rendimiento equilibrado. Excelente para programación, escritura y tareas de propósito general.
+
+⚡ **Claude Haiku 4.5** — Velocidad relámpago. Perfecto para respuestas rápidas, tareas simples y chat.
+
+📦 **Claude Sonnet 4** — Generación anterior. Todavía potente y confiable para la mayoría de casos de uso.
+
+📦 **Claude 3.7 Sonnet** — Modelo heredado. Disponible para compatibilidad retroactiva.
+
+> 💡 **Resolución Inteligente de Modelos:** Usa cualquier formato de nombre de modelo — `claude-sonnet-4-5`, `claude-sonnet-4.5`, o incluso nombres versionados como `claude-sonnet-4-5-20250929`. El gateway los normaliza automáticamente.
+
+---
+
+## ✨ Características
+
+| Característica | Descripción |
+|----------------|-------------|
+| 🔌 **API compatible con OpenAI** | Funciona con cualquier herramienta compatible con OpenAI |
+| 🔌 **API compatible con Anthropic** | Endpoint nativo `/v1/messages` |
+| 🌐 **Soporte de VPN/Proxy** | Proxy HTTP/SOCKS5 para redes restringidas |
+| 🧠 **Pensamiento Extendido** | El razonamiento es exclusivo de nuestro proyecto |
+| 👁️ **Soporte de Visión** | Envía imágenes al modelo |
+| 🛠️ **Llamada de Herramientas** | Soporta llamada de funciones |
+| 💬 **Historial completo de mensajes** | Pasa el contexto completo de la conversación |
+| 📡 **Streaming** | Soporte completo de streaming SSE |
+| 🔄 **Lógica de Reintentos** | Reintentos automáticos en errores (403, 429, 5xx) |
+| 📋 **Lista extendida de modelos** | Incluyendo modelos versionados |
+| 🔐 **Gestión inteligente de tokens** | Actualización automática antes de la expiración |
+
+---
+
+## 🚀 Inicio Rápido
+
+### Prerrequisitos
+
+- Python 3.10+
+- Uno de los siguientes:
+ - [Kiro IDE](https://kiro.dev/) con cuenta iniciada, O
+ - [Kiro CLI](https://kiro.dev/cli/) con AWS SSO (AWS IAM Identity Center, OIDC) - Builder ID gratuito o cuenta empresarial
+
+### Instalación
+
+```bash
+# Clona el repositorio (requiere Git)
+git clone https://github.com/Jwadow/kiro-gateway.git
+cd kiro-gateway
+
+# O descarga el ZIP: Code → Download ZIP → extrae → abre la carpeta kiro-gateway
+
+# Instala las dependencias
+pip install -r requirements.txt
+
+# Configura (ver sección Configuración)
+cp .env.example .env
+# Copia y edita .env con tus credenciales
+
+# Inicia el servidor
+python main.py
+
+# O con puerto personalizado (si 8000 está ocupado)
+python main.py --port 9000
+```
+
+El servidor estará disponible en `http://localhost:8000`
+
+---
+
+## ⚙️ Configuración
+
+### Opción 1: Archivo JSON de Credenciales (Kiro IDE / Enterprise)
+
+Especifica la ruta al archivo de credenciales:
+
+Funciona con:
+- **Kiro IDE** (estándar) - para cuentas personales
+- **Enterprise** - para cuentas empresariales con SSO
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/kiro-auth-token.json"
+
+# Contraseña para proteger TU servidor proxy (crea cualquier cadena segura)
+# Usarás esto como api_key al conectarte a tu gateway
+PROXY_API_KEY="my-super-secret-password-123"
+```
+
+
+📄 Formato del archivo JSON
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:...",
+ "region": "us-east-1",
+ "clientIdHash": "abc123..." // Optional: for corporate SSO setups
+}
+```
+
+> **Nota:** Si tienes dos archivos JSON en `~/.aws/sso/cache/` (por ejemplo, `kiro-auth-token.json` y un archivo con nombre hash), usa `kiro-auth-token.json` en `KIRO_CREDS_FILE`. El gateway cargará automáticamente el otro archivo.
+
+
+
+### Opción 2: Variables de Entorno (archivo .env)
+
+Crea un archivo `.env` en la raíz del proyecto:
+
+```env
+# Requerido
+REFRESH_TOKEN="tu_kiro_refresh_token"
+
+# Contraseña para proteger TU servidor proxy (crea cualquier cadena segura)
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Opcional
+PROFILE_ARN="arn:aws:codewhisperer:us-east-1:..."
+KIRO_REGION="us-east-1"
+```
+
+### Opción 3: Credenciales AWS SSO (kiro-cli / Enterprise)
+
+Si usas `kiro-cli` o Kiro IDE con AWS SSO (AWS IAM Identity Center), el gateway detectará y usará automáticamente la autenticación apropiada.
+
+Funciona tanto con cuentas Builder ID gratuitas como con cuentas empresariales.
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/your-sso-cache-file.json"
+
+# Contraseña para proteger TU servidor proxy
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Nota: PROFILE_ARN NO es necesario para AWS SSO (Builder ID y cuentas empresariales)
+# El gateway funcionará sin él
+```
+
+
+📄 Formato del archivo JSON de AWS SSO
+
+Los archivos de credenciales de AWS SSO (de `~/.aws/sso/cache/`) contienen:
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "region": "us-east-1",
+ "clientId": "...",
+ "clientSecret": "..."
+}
+```
+
+**Nota:** Los usuarios de AWS SSO (Builder ID y cuentas empresariales) NO necesitan `profileArn`. El gateway funcionará sin él (si se especifica, será ignorado).
+
+
+
+
+🔍 Cómo funciona
+
+El gateway detecta automáticamente el tipo de autenticación basándose en el archivo de credenciales:
+
+- **Kiro Desktop Auth** (predeterminado): Usado cuando `clientId` y `clientSecret` NO están presentes
+ - Endpoint: `https://prod.{region}.auth.desktop.kiro.dev/refreshToken`
+
+- **AWS SSO (OIDC)**: Usado cuando `clientId` y `clientSecret` están presentes
+ - Endpoint: `https://oidc.{region}.amazonaws.com/token`
+
+¡No se necesita configuración adicional — solo apunta a tu archivo de credenciales!
+
+
+
+### Opción 4: Base de datos SQLite de kiro-cli
+
+Si usas `kiro-cli` y prefieres usar su base de datos SQLite directamente:
+
+```env
+KIRO_CLI_DB_FILE="~/.local/share/kiro-cli/data.sqlite3"
+
+# Contraseña para proteger TU servidor proxy
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Nota: PROFILE_ARN NO es necesario para AWS SSO (Builder ID y cuentas empresariales)
+# El gateway funcionará sin él
+```
+
+
+📄 Ubicaciones de la base de datos
+
+| Herramienta CLI | Ruta de la Base de Datos |
+|-----------------|--------------------------|
+| kiro-cli | `~/.local/share/kiro-cli/data.sqlite3` |
+| amazon-q-developer-cli | `~/.local/share/amazon-q/data.sqlite3` |
+
+El gateway lee las credenciales de la tabla `auth_kv` que almacena:
+- `kirocli:odic:token` o `codewhisperer:odic:token` — token de acceso, token de actualización, expiración
+- `kirocli:odic:device-registration` o `codewhisperer:odic:device-registration` — ID de cliente y secreto
+
+Ambos formatos de clave son soportados para compatibilidad con diferentes versiones de kiro-cli.
+
+
+
+### Obtener Credenciales
+
+**Para usuarios de Kiro IDE:**
+- Inicia sesión en Kiro IDE y usa la Opción 1 arriba (archivo JSON de credenciales)
+- El archivo de credenciales se crea automáticamente después de iniciar sesión
+
+**Para usuarios de Kiro CLI:**
+- Inicia sesión con `kiro-cli login` y usa la Opción 3 u Opción 4 arriba
+- ¡No se necesita extracción manual de tokens!
+
+
+🔧 Avanzado: Extracción manual de token
+
+Si necesitas extraer manualmente el refresh token (por ejemplo, para depuración), puedes interceptar el tráfico de Kiro IDE:
+- Busca solicitudes a: `prod.us-east-1.auth.desktop.kiro.dev/refreshToken`
+
+
+
+---
+
+## 🌐 Soporte de VPN/Proxy
+
+**Para usuarios en China, redes corporativas o regiones con problemas de conectividad a servicios de AWS.**
+
+El gateway admite enrutar todas las solicitudes de Kiro API a través de un servidor VPN o proxy. Esto es esencial si experimenta problemas de conexión a puntos finales de AWS o necesita usar un proxy corporativo.
+
+### Configuración
+
+Añade a tu archivo `.env`:
+
+```env
+# Proxy HTTP
+VPN_PROXY_URL=http://127.0.0.1:7890
+
+# Proxy SOCKS5
+VPN_PROXY_URL=socks5://127.0.0.1:1080
+
+# Con autenticación (proxies corporativos)
+VPN_PROXY_URL=http://username:password@proxy.company.com:8080
+
+# Sin protocolo (por defecto http://)
+VPN_PROXY_URL=192.168.1.100:8080
+```
+
+### Protocolos Soportados
+
+- ✅ **HTTP** — Protocolo proxy estándar
+- ✅ **HTTPS** — Conexiones proxy seguras
+- ✅ **SOCKS5** — Protocolo proxy avanzado (común en software VPN)
+- ✅ **Autenticación** — Usuario/contraseña incrustados en URL
+
+### Cuándo lo Necesitas
+
+| Situación | Solución |
+|-----------|----------|
+| Tiempos de espera de conexión a AWS | Usa VPN/proxy para enrutar tráfico |
+| Restricciones de red corporativa | Configura el proxy de tu empresa |
+| Problemas de conectividad regional | Usa un servicio VPN con soporte proxy |
+| Requisitos de privacidad | Enruta a través de tu propio servidor proxy |
+
+### Software VPN Popular con Soporte Proxy
+
+La mayoría de clientes VPN proporcionan un servidor proxy local:
+- **Sing-box** — Cliente VPN moderno con soporte HTTP/SOCKS5 proxy
+- **Clash** — Generalmente se ejecuta en `http://127.0.0.1:7890`
+- **V2Ray** — Proxy SOCKS5/HTTP configurable
+- **Shadowsocks** — Soporte proxy SOCKS5
+- **VPN Corporativo** — Consulta a tu departamento de TI para configuración de proxy
+
+Deja `VPN_PROXY_URL` vacío (por defecto) si no necesitas soporte proxy.
+
+---
+
+## 📡 Referencia de API
+
+### Endpoints
+
+| Endpoint | Método | Descripción |
+|----------|--------|-------------|
+| `/` | GET | Verificación de salud |
+| `/health` | GET | Verificación de salud detallada |
+| `/v1/models` | GET | Lista modelos disponibles |
+| `/v1/chat/completions` | POST | OpenAI Chat Completions API |
+| `/v1/messages` | POST | Anthropic Messages API |
+
+---
+
+## 💡 Ejemplos de Uso
+
+### OpenAI API
+
+
+🔹 Solicitud cURL Simple
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "¡Hola!"}],
+ "stream": true
+ }'
+```
+
+> **Nota:** Reemplaza `my-super-secret-password-123` con el `PROXY_API_KEY` que configuraste en tu archivo `.env`.
+
+
+
+
+🔹 Solicitud con Streaming
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {"role": "system", "content": "Eres un asistente útil."},
+ {"role": "user", "content": "¿Cuánto es 2+2?"}
+ ],
+ "stream": true
+ }'
+```
+
+
+
+
+🛠️ Con Llamada de Herramientas
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "¿Cómo está el clima en Londres?"}],
+ "tools": [{
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Obtener el clima para una ubicación",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "Nombre de la ciudad"}
+ },
+ "required": ["location"]
+ }
+ }
+ }]
+ }'
+```
+
+
+
+
+🐍 Python OpenAI SDK
+
+```python
+from openai import OpenAI
+
+client = OpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123" # Tu PROXY_API_KEY del .env
+)
+
+response = client.chat.completions.create(
+ model="claude-sonnet-4-5",
+ messages=[
+ {"role": "system", "content": "Eres un asistente útil."},
+ {"role": "user", "content": "¡Hola!"}
+ ],
+ stream=True
+)
+
+for chunk in response:
+ if chunk.choices[0].delta.content:
+ print(chunk.choices[0].delta.content, end="")
+```
+
+
+
+
+🦜 LangChain
+
+```python
+from langchain_openai import ChatOpenAI
+
+llm = ChatOpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123", # Tu PROXY_API_KEY del .env
+ model="claude-sonnet-4-5"
+)
+
+response = llm.invoke("Hola, ¿cómo estás?")
+print(response.content)
+```
+
+
+
+### Anthropic API
+
+
+🔹 Solicitud cURL Simple
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "¡Hola!"}]
+ }'
+```
+
+> **Nota:** La API de Anthropic usa el header `x-api-key` en lugar de `Authorization: Bearer`. Ambos son soportados.
+
+
+
+
+🔹 Con Prompt de Sistema
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "system": "Eres un asistente útil.",
+ "messages": [{"role": "user", "content": "¡Hola!"}]
+ }'
+```
+
+> **Nota:** En la API de Anthropic, `system` es un campo separado, no un mensaje.
+
+
+
+
+📡 Streaming
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "stream": true,
+ "messages": [{"role": "user", "content": "¡Hola!"}]
+ }'
+```
+
+
+
+
+🐍 Python Anthropic SDK
+
+```python
+import anthropic
+
+client = anthropic.Anthropic(
+ api_key="my-super-secret-password-123", # Tu PROXY_API_KEY del .env
+ base_url="http://localhost:8000"
+)
+
+# Sin streaming
+response = client.messages.create(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "¡Hola!"}]
+)
+print(response.content[0].text)
+
+# Con streaming
+with client.messages.stream(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "¡Hola!"}]
+) as stream:
+ for text in stream.text_stream:
+ print(text, end="", flush=True)
+```
+
+
+
+---
+
+## 🔧 Depuración
+
+El registro de depuración está **deshabilitado por defecto**. Para habilitar, añade a tu `.env`:
+
+```env
+# Modo de registro de depuración:
+# - off: deshabilitado (predeterminado)
+# - errors: guardar logs solo para solicitudes fallidas (4xx, 5xx) - recomendado para solución de problemas
+# - all: guardar logs para cada solicitud (sobrescribe en cada solicitud)
+DEBUG_MODE=errors
+```
+
+### Modos de Depuración
+
+| Modo | Descripción | Caso de Uso |
+|------|-------------|-------------|
+| `off` | Deshabilitado (predeterminado) | Producción |
+| `errors` | Guardar logs solo para solicitudes fallidas (4xx, 5xx) | **Recomendado para solución de problemas** |
+| `all` | Guardar logs para cada solicitud | Desarrollo/depuración |
+
+### Archivos de Depuración
+
+Cuando está habilitado, las solicitudes se registran en la carpeta `debug_logs/`:
+
+| Archivo | Descripción |
+|---------|-------------|
+| `request_body.json` | Solicitud entrante del cliente (formato OpenAI) |
+| `kiro_request_body.json` | Solicitud enviada a la API de Kiro |
+| `response_stream_raw.txt` | Stream crudo de Kiro |
+| `response_stream_modified.txt` | Stream transformado (formato OpenAI) |
+| `app_logs.txt` | Logs de la aplicación para la solicitud |
+| `error_info.json` | Detalles del error (solo en errores) |
+
+---
+
+## 📜 Licencia
+
+Este proyecto está licenciado bajo la **GNU Affero General Public License v3.0 (AGPL-3.0)**.
+
+Esto significa:
+- ✅ Puedes usar, modificar y distribuir este software
+- ✅ Puedes usarlo con fines comerciales
+- ⚠️ **Debes revelar el código fuente** cuando distribuyas el software
+- ⚠️ **El uso en red es distribución** — si ejecutas una versión modificada en un servidor y permites que otros interactúen con ella, debes hacer el código fuente disponible para ellos
+- ⚠️ Las modificaciones deben ser liberadas bajo la misma licencia
+
+Consulta el archivo [LICENSE](../../LICENSE) para el texto completo de la licencia.
+
+### ¿Por qué AGPL-3.0?
+
+AGPL-3.0 asegura que las mejoras a este software beneficien a toda la comunidad. Si modificas este gateway y lo despliegas como un servicio, debes compartir tus mejoras con tus usuarios.
+
+### Acuerdo de Licencia de Contribuidor (CLA)
+
+Al enviar una contribución a este proyecto, aceptas los términos de nuestro [Acuerdo de Licencia de Contribuidor (CLA)](../../CLA.md). Esto asegura que:
+- Tienes el derecho de enviar la contribución
+- Otorgas al mantenedor derechos para usar y relicenciar tu contribución
+- El proyecto permanece legalmente protegido
+
+---
+
+## 💖 Apoya el Proyecto
+
+
+
+
+
+**¡Si este proyecto te ahorró tiempo o dinero, considera apoyarlo!**
+
+Cada contribución ayuda a mantener este proyecto vivo y creciendo
+
+
+
+### 🤑 Donar
+
+[**☕ Donación Única**](https://app.lava.top/jwadow?tabId=donate) • [**💎 Apoyo Mensual**](https://app.lava.top/jwadow?tabId=subscriptions)
+
+
+
+### 🪙 O envía criptomonedas
+
+| Moneda | Red | Dirección |
+|:------:|:---:|:----------|
+| **USDT** | TRC20 | `TSVtgRc9pkC1UgcbVeijBHjFmpkYHDRu26` |
+| **BTC** | Bitcoin | `12GZqxqpcBsqJ4Vf1YreLqwoMGvzBPgJq6` |
+| **ETH** | Ethereum | `0xc86eab3bba3bbaf4eb5b5fff8586f1460f1fd395` |
+| **SOL** | Solana | `9amykF7KibZmdaw66a1oqYJyi75fRqgdsqnG66AK3jvh` |
+| **TON** | TON | `UQBVh8T1H3GI7gd7b-_PPNnxHYYxptrcCVf3qQk5v41h3QTM` |
+
+
+
+**[⬆ Volver Arriba](#-kiro-gateway)**
+
+
diff --git a/kiro-gateway/docs/id/README.md b/kiro-gateway/docs/id/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..9f16cbc1b667caf169f127579bf6a1318fc5c0e0
--- /dev/null
+++ b/kiro-gateway/docs/id/README.md
@@ -0,0 +1,626 @@
+
+
+# 👻 Kiro Gateway
+
+**Gateway proxy untuk Kiro API (Amazon Q Developer / AWS CodeWhisperer)**
+
+[🇬🇧 English](../../README.md) • [🇷🇺 Русский](../ru/README.md) • [🇨🇳 中文](../zh/README.md) • [🇪🇸 Español](../es/README.md) • [🇧🇷 Português](../pt/README.md) • [🇯🇵 日本語](../ja/README.md) • [🇰🇷 한국어](../ko/README.md)
+
+Dibuat dengan ❤️ oleh [@Jwadow](https://github.com/jwadow)
+
+[](https://www.gnu.org/licenses/agpl-3.0)
+[](https://www.python.org/downloads/)
+[](https://fastapi.tiangolo.com/)
+[](#-dukung-proyek)
+
+*Gunakan model Claude dari Kiro dengan Claude Code, OpenCode, Cursor, Cline, Roo Code, Kilo Code, Obsidian, OpenAI SDK, LangChain, Continue dan alat lain yang kompatibel dengan OpenAI atau Anthropic*
+
+[Model](#-model-yang-didukung) • [Fitur](#-fitur) • [Mulai Cepat](#-mulai-cepat) • [Konfigurasi](#%EF%B8%8F-konfigurasi) • [💖 Dukung](#-dukung-proyek)
+
+
+
+---
+
+## 🤖 Model yang Tersedia
+
+> ⚠️ **Penting:** Ketersediaan model bergantung pada paket Kiro Anda (gratis/berbayar). Gateway menyediakan akses ke model yang tersedia di IDE atau CLI Anda berdasarkan langganan Anda. Daftar di bawah menunjukkan model yang umumnya tersedia di **paket gratis**.
+
+> 🔒 **Claude Opus 4.5** telah dihapus dari paket gratis pada 17 Januari 2026. Mungkin tersedia di paket berbayar — periksa daftar model di IDE/CLI Anda.
+
+🚀 **Claude Sonnet 4.5** — Performa seimbang. Bagus untuk coding, menulis, dan tugas umum.
+
+⚡ **Claude Haiku 4.5** — Secepat kilat. Sempurna untuk respons cepat, tugas sederhana, dan chat.
+
+📦 **Claude Sonnet 4** — Generasi sebelumnya. Masih kuat dan andal untuk sebagian besar kasus penggunaan.
+
+📦 **Claude 3.7 Sonnet** — Model lama. Tersedia untuk kompatibilitas mundur.
+
+> 💡 **Resolusi Model Cerdas:** Gunakan format nama model apa pun — `claude-sonnet-4-5`, `claude-sonnet-4.5`, atau bahkan nama berversi seperti `claude-sonnet-4-5-20250929`. Gateway akan menormalisasi secara otomatis.
+
+---
+
+## ✨ Fitur
+
+| Fitur | Deskripsi |
+|-------|-----------|
+| 🔌 **API kompatibel OpenAI** | Bekerja dengan alat apa pun yang kompatibel dengan OpenAI |
+| 🔌 **API kompatibel Anthropic** | Endpoint native `/v1/messages` |
+| 🌐 **Dukungan VPN/Proxy** | Proxy HTTP/SOCKS5 untuk jaringan terbatas |
+| 🧠 **Pemikiran Diperluas** | Penalaran adalah eksklusif proyek kami |
+| 👁️ **Dukungan Visi** | Kirim gambar ke model |
+| 🛠️ **Pemanggilan Alat** | Mendukung pemanggilan fungsi |
+| 💬 **Riwayat pesan lengkap** | Meneruskan konteks percakapan lengkap |
+| 📡 **Streaming** | Dukungan streaming SSE penuh |
+| 🔄 **Logika Retry** | Retry otomatis saat error (403, 429, 5xx) |
+| 📋 **Daftar model diperluas** | Termasuk model berversi |
+| 🔐 **Manajemen token cerdas** | Refresh otomatis sebelum kedaluwarsa |
+
+---
+
+## 🚀 Mulai Cepat
+
+### Prasyarat
+
+- Python 3.10+
+- Salah satu dari berikut:
+ - [Kiro IDE](https://kiro.dev/) dengan akun yang sudah login, ATAU
+ - [Kiro CLI](https://kiro.dev/cli/) dengan AWS SSO (AWS IAM Identity Center, OIDC) - Builder ID gratis atau akun perusahaan
+
+### Instalasi
+
+```bash
+# Clone repositori (memerlukan Git)
+git clone https://github.com/Jwadow/kiro-gateway.git
+cd kiro-gateway
+
+# Atau unduh ZIP: Code → Download ZIP → ekstrak → buka folder kiro-gateway
+
+# Instal dependensi
+pip install -r requirements.txt
+
+# Konfigurasi (lihat bagian Konfigurasi)
+cp .env.example .env
+# Salin dan edit .env dengan kredensial Anda
+
+# Jalankan server
+python main.py
+
+# Atau dengan port kustom (jika 8000 sedang digunakan)
+python main.py --port 9000
+```
+
+Server akan tersedia di `http://localhost:8000`
+
+---
+
+## ⚙️ Konfigurasi
+
+### Opsi 1: File JSON Kredensial (Kiro IDE / Enterprise)
+
+Tentukan path ke file kredensial:
+
+Bekerja dengan:
+- **Kiro IDE** (standar) - untuk akun pribadi
+- **Enterprise** - untuk akun perusahaan dengan SSO
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/kiro-auth-token.json"
+
+# Password untuk melindungi server proxy ANDA (buat string aman apa pun)
+# Anda akan menggunakan ini sebagai api_key saat menghubungkan ke gateway Anda
+PROXY_API_KEY="my-super-secret-password-123"
+```
+
+
+📄 Format file JSON
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:...",
+ "region": "us-east-1",
+ "clientIdHash": "abc123..." // Optional: for corporate SSO setups
+}
+```
+
+> **Catatan:** Jika Anda memiliki dua file JSON di `~/.aws/sso/cache/` (misalnya, `kiro-auth-token.json` dan file dengan nama hash), gunakan `kiro-auth-token.json` di `KIRO_CREDS_FILE`. Gateway akan secara otomatis memuat file lainnya.
+
+
+
+### Opsi 2: Variabel Lingkungan (file .env)
+
+Buat file `.env` di root proyek:
+
+```env
+# Wajib
+REFRESH_TOKEN="kiro_refresh_token_anda"
+
+# Password untuk melindungi server proxy ANDA (buat string aman apa pun)
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Opsional
+PROFILE_ARN="arn:aws:codewhisperer:us-east-1:..."
+KIRO_REGION="us-east-1"
+```
+
+### Opsi 3: Kredensial AWS SSO (kiro-cli / Enterprise)
+
+Jika Anda menggunakan `kiro-cli` atau Kiro IDE dengan AWS SSO (AWS IAM Identity Center), gateway akan secara otomatis mendeteksi dan menggunakan autentikasi yang sesuai.
+
+Bekerja dengan akun Builder ID gratis dan akun perusahaan.
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/your-sso-cache-file.json"
+
+# Password untuk melindungi server proxy ANDA
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Catatan: PROFILE_ARN TIDAK diperlukan untuk AWS SSO (Builder ID dan akun perusahaan)
+# Gateway akan bekerja tanpanya
+```
+
+
+📄 Format file JSON AWS SSO
+
+File kredensial AWS SSO (dari `~/.aws/sso/cache/`) berisi:
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "region": "us-east-1",
+ "clientId": "...",
+ "clientSecret": "..."
+}
+```
+
+**Catatan:** Pengguna AWS SSO (Builder ID dan akun perusahaan) TIDAK memerlukan `profileArn`. Gateway akan bekerja tanpanya (jika ditentukan, akan diabaikan).
+
+
+
+
+🔍 Cara kerjanya
+
+Gateway secara otomatis mendeteksi tipe autentikasi berdasarkan file kredensial:
+
+- **Kiro Desktop Auth** (default): Digunakan ketika `clientId` dan `clientSecret` TIDAK ada
+ - Endpoint: `https://prod.{region}.auth.desktop.kiro.dev/refreshToken`
+
+- **AWS SSO (OIDC)**: Digunakan ketika `clientId` dan `clientSecret` ada
+ - Endpoint: `https://oidc.{region}.amazonaws.com/token`
+
+Tidak perlu konfigurasi tambahan — cukup arahkan ke file kredensial Anda!
+
+
+
+### Opsi 4: Database SQLite kiro-cli
+
+Jika Anda menggunakan `kiro-cli` dan lebih suka menggunakan database SQLite-nya secara langsung:
+
+```env
+KIRO_CLI_DB_FILE="~/.local/share/kiro-cli/data.sqlite3"
+
+# Password untuk melindungi server proxy ANDA
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Catatan: PROFILE_ARN TIDAK diperlukan untuk AWS SSO (Builder ID dan akun perusahaan)
+# Gateway akan bekerja tanpanya
+```
+
+
+📄 Lokasi database
+
+| Alat CLI | Path Database |
+|----------|---------------|
+| kiro-cli | `~/.local/share/kiro-cli/data.sqlite3` |
+| amazon-q-developer-cli | `~/.local/share/amazon-q/data.sqlite3` |
+
+Gateway membaca kredensial dari tabel `auth_kv` yang menyimpan:
+- `kirocli:odic:token` atau `codewhisperer:odic:token` — access token, refresh token, kedaluwarsa
+- `kirocli:odic:device-registration` atau `codewhisperer:odic:device-registration` — client ID dan secret
+
+Kedua format kunci didukung untuk kompatibilitas dengan versi kiro-cli yang berbeda.
+
+
+
+### Mendapatkan Kredensial
+
+**Untuk pengguna Kiro IDE:**
+- Login ke Kiro IDE dan gunakan Opsi 1 di atas (file JSON kredensial)
+- File kredensial dibuat secara otomatis setelah login
+
+**Untuk pengguna Kiro CLI:**
+- Login dengan `kiro-cli login` dan gunakan Opsi 3 atau Opsi 4 di atas
+- Tidak perlu ekstraksi token manual!
+
+
+🔧 Lanjutan: Ekstraksi token manual
+
+Jika Anda perlu mengekstrak refresh token secara manual (misalnya, untuk debugging), Anda dapat mencegat traffic Kiro IDE:
+- Cari request ke: `prod.us-east-1.auth.desktop.kiro.dev/refreshToken`
+
+
+
+---
+
+## 🌐 Dukungan VPN/Proxy
+
+**Untuk pengguna di China, jaringan korporat, atau wilayah dengan masalah konektivitas ke layanan AWS.**
+
+Gateway mendukung perutean semua permintaan Kiro API melalui server VPN atau proxy. Ini penting jika Anda mengalami masalah koneksi ke endpoint AWS atau perlu menggunakan proxy korporat.
+
+### Konfigurasi
+
+Tambahkan ke file `.env` Anda:
+
+```env
+# Proxy HTTP
+VPN_PROXY_URL=http://127.0.0.1:7890
+
+# Proxy SOCKS5
+VPN_PROXY_URL=socks5://127.0.0.1:1080
+
+# Dengan autentikasi (proxy korporat)
+VPN_PROXY_URL=http://username:password@proxy.company.com:8080
+
+# Tanpa protokol (default ke http://)
+VPN_PROXY_URL=192.168.1.100:8080
+```
+
+### Protokol yang Didukung
+
+- ✅ **HTTP** — Protokol proxy standar
+- ✅ **HTTPS** — Koneksi proxy aman
+- ✅ **SOCKS5** — Protokol proxy lanjutan (umum di software VPN)
+- ✅ **Autentikasi** — Username/password tertanam di URL
+
+### Kapan Anda Membutuhkannya
+
+| Situasi | Solusi |
+|---------|--------|
+| Timeout koneksi ke AWS | Gunakan VPN/proxy untuk merutekan lalu lintas |
+| Pembatasan jaringan korporat | Konfigurasi proxy perusahaan Anda |
+| Masalah konektivitas regional | Gunakan layanan VPN dengan dukungan proxy |
+| Persyaratan privasi | Rutekan melalui server proxy Anda sendiri |
+
+### Software VPN Populer dengan Dukungan Proxy
+
+Sebagian besar klien VPN menyediakan server proxy lokal:
+- **Sing-box** — Klien VPN modern dengan dukungan proxy HTTP/SOCKS5
+- **Clash** — Biasanya berjalan di `http://127.0.0.1:7890`
+- **V2Ray** — Proxy SOCKS5/HTTP yang dapat dikonfigurasi
+- **Shadowsocks** — Dukungan proxy SOCKS5
+- **VPN Korporat** — Tanyakan departemen IT Anda untuk pengaturan proxy
+
+Biarkan `VPN_PROXY_URL` kosong (default) jika Anda tidak memerlukan dukungan proxy.
+
+---
+
+## 📡 Referensi API
+
+### Endpoint
+
+| Endpoint | Metode | Deskripsi |
+|----------|--------|-----------|
+| `/` | GET | Health check |
+| `/health` | GET | Health check detail |
+| `/v1/models` | GET | Daftar model yang tersedia |
+| `/v1/chat/completions` | POST | OpenAI Chat Completions API |
+| `/v1/messages` | POST | Anthropic Messages API |
+
+---
+
+## 💡 Contoh Penggunaan
+
+### OpenAI API
+
+
+🔹 Request cURL Sederhana
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Halo!"}],
+ "stream": true
+ }'
+```
+
+> **Catatan:** Ganti `my-super-secret-password-123` dengan `PROXY_API_KEY` yang Anda atur di file `.env`.
+
+
+
+
+🔹 Request dengan Streaming
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {"role": "system", "content": "Kamu adalah asisten yang membantu."},
+ {"role": "user", "content": "Berapa 2+2?"}
+ ],
+ "stream": true
+ }'
+```
+
+
+
+
+🛠️ Dengan Pemanggilan Alat
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Bagaimana cuaca di London?"}],
+ "tools": [{
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Dapatkan cuaca untuk suatu lokasi",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "Nama kota"}
+ },
+ "required": ["location"]
+ }
+ }
+ }]
+ }'
+```
+
+
+
+
+🐍 Python OpenAI SDK
+
+```python
+from openai import OpenAI
+
+client = OpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123" # PROXY_API_KEY Anda dari .env
+)
+
+response = client.chat.completions.create(
+ model="claude-sonnet-4-5",
+ messages=[
+ {"role": "system", "content": "Kamu adalah asisten yang membantu."},
+ {"role": "user", "content": "Halo!"}
+ ],
+ stream=True
+)
+
+for chunk in response:
+ if chunk.choices[0].delta.content:
+ print(chunk.choices[0].delta.content, end="")
+```
+
+
+
+
+🦜 LangChain
+
+```python
+from langchain_openai import ChatOpenAI
+
+llm = ChatOpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123", # PROXY_API_KEY Anda dari .env
+ model="claude-sonnet-4-5"
+)
+
+response = llm.invoke("Halo, apa kabar?")
+print(response.content)
+```
+
+
+
+### Anthropic API
+
+
+🔹 Request cURL Sederhana
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Halo!"}]
+ }'
+```
+
+> **Catatan:** Anthropic API menggunakan header `x-api-key` bukan `Authorization: Bearer`. Keduanya didukung.
+
+
+
+
+🔹 Dengan System Prompt
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "system": "Kamu adalah asisten yang membantu.",
+ "messages": [{"role": "user", "content": "Halo!"}]
+ }'
+```
+
+> **Catatan:** Di Anthropic API, `system` adalah field terpisah, bukan pesan.
+
+
+
+
+📡 Streaming
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "stream": true,
+ "messages": [{"role": "user", "content": "Halo!"}]
+ }'
+```
+
+
+
+
+🐍 Python Anthropic SDK
+
+```python
+import anthropic
+
+client = anthropic.Anthropic(
+ api_key="my-super-secret-password-123", # PROXY_API_KEY Anda dari .env
+ base_url="http://localhost:8000"
+)
+
+# Tanpa streaming
+response = client.messages.create(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "Halo!"}]
+)
+print(response.content[0].text)
+
+# Dengan streaming
+with client.messages.stream(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "Halo!"}]
+) as stream:
+ for text in stream.text_stream:
+ print(text, end="", flush=True)
+```
+
+
+
+---
+
+## 🔧 Debugging
+
+Logging debug **dinonaktifkan secara default**. Untuk mengaktifkan, tambahkan ke `.env` Anda:
+
+```env
+# Mode logging debug:
+# - off: dinonaktifkan (default)
+# - errors: simpan log hanya untuk request yang gagal (4xx, 5xx) - direkomendasikan untuk troubleshooting
+# - all: simpan log untuk setiap request (ditimpa setiap request)
+DEBUG_MODE=errors
+```
+
+### Mode Debug
+
+| Mode | Deskripsi | Kasus Penggunaan |
+|------|-----------|------------------|
+| `off` | Dinonaktifkan (default) | Produksi |
+| `errors` | Simpan log hanya untuk request yang gagal (4xx, 5xx) | **Direkomendasikan untuk troubleshooting** |
+| `all` | Simpan log untuk setiap request | Pengembangan/debugging |
+
+### File Debug
+
+Ketika diaktifkan, request dicatat ke folder `debug_logs/`:
+
+| File | Deskripsi |
+|------|-----------|
+| `request_body.json` | Request masuk dari klien (format OpenAI) |
+| `kiro_request_body.json` | Request yang dikirim ke Kiro API |
+| `response_stream_raw.txt` | Stream mentah dari Kiro |
+| `response_stream_modified.txt` | Stream yang ditransformasi (format OpenAI) |
+| `app_logs.txt` | Log aplikasi untuk request |
+| `error_info.json` | Detail error (hanya saat error) |
+
+---
+
+## 📜 Lisensi
+
+Proyek ini dilisensikan di bawah **GNU Affero General Public License v3.0 (AGPL-3.0)**.
+
+Ini berarti:
+- ✅ Anda dapat menggunakan, memodifikasi, dan mendistribusikan software ini
+- ✅ Anda dapat menggunakannya untuk tujuan komersial
+- ⚠️ **Anda harus mengungkapkan kode sumber** ketika Anda mendistribusikan software
+- ⚠️ **Penggunaan jaringan adalah distribusi** — jika Anda menjalankan versi yang dimodifikasi di server dan membiarkan orang lain berinteraksi dengannya, Anda harus membuat kode sumber tersedia untuk mereka
+- ⚠️ Modifikasi harus dirilis di bawah lisensi yang sama
+
+Lihat file [LICENSE](../../LICENSE) untuk teks lisensi lengkap.
+
+### Mengapa AGPL-3.0?
+
+AGPL-3.0 memastikan bahwa perbaikan pada software ini menguntungkan seluruh komunitas. Jika Anda memodifikasi gateway ini dan menerapkannya sebagai layanan, Anda harus membagikan perbaikan Anda dengan pengguna Anda.
+
+### Perjanjian Lisensi Kontributor (CLA)
+
+Dengan mengirimkan kontribusi ke proyek ini, Anda menyetujui ketentuan [Perjanjian Lisensi Kontributor (CLA)](../../CLA.md) kami. Ini memastikan bahwa:
+- Anda memiliki hak untuk mengirimkan kontribusi
+- Anda memberikan hak kepada pengelola untuk menggunakan dan melisensi ulang kontribusi Anda
+- Proyek tetap dilindungi secara hukum
+
+---
+
+## 💖 Dukung Proyek
+
+
+
+
+
+**Jika proyek ini menghemat waktu atau uang Anda, pertimbangkan untuk mendukungnya!**
+
+Setiap kontribusi membantu menjaga proyek ini tetap hidup dan berkembang
+
+
+
+### 🤑 Donasi
+
+[**☕ Donasi Sekali**](https://app.lava.top/jwadow?tabId=donate) • [**💎 Dukungan Bulanan**](https://app.lava.top/jwadow?tabId=subscriptions)
+
+
+
+### 🪙 Atau kirim crypto
+
+| Mata Uang | Jaringan | Alamat |
+|:---------:|:--------:|:-------|
+| **USDT** | TRC20 | `TSVtgRc9pkC1UgcbVeijBHjFmpkYHDRu26` |
+| **BTC** | Bitcoin | `12GZqxqpcBsqJ4Vf1YreLqwoMGvzBPgJq6` |
+| **ETH** | Ethereum | `0xc86eab3bba3bbaf4eb5b5fff8586f1460f1fd395` |
+| **SOL** | Solana | `9amykF7KibZmdaw66a1oqYJyi75fRqgdsqnG66AK3jvh` |
+| **TON** | TON | `UQBVh8T1H3GI7gd7b-_PPNnxHYYxptrcCVf3qQk5v41h3QTM` |
+
+
+
+**[⬆ Kembali ke Atas](#-kiro-gateway)**
+
+
diff --git a/kiro-gateway/docs/ja/README.md b/kiro-gateway/docs/ja/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..f8eece16ec19c749eb1f32d4a0df380f36284bd8
--- /dev/null
+++ b/kiro-gateway/docs/ja/README.md
@@ -0,0 +1,626 @@
+
+
+# 👻 Kiro Gateway
+
+**Kiro API (Amazon Q Developer / AWS CodeWhisperer) 用プロキシゲートウェイ**
+
+[🇬🇧 English](../../README.md) • [🇷🇺 Русский](../ru/README.md) • [🇨🇳 中文](../zh/README.md) • [🇪🇸 Español](../es/README.md) • [🇮🇩 Indonesia](../id/README.md) • [🇧🇷 Português](../pt/README.md) • [🇰🇷 한국어](../ko/README.md)
+
+[@Jwadow](https://github.com/jwadow) が ❤️ を込めて作成
+
+[](https://www.gnu.org/licenses/agpl-3.0)
+[](https://www.python.org/downloads/)
+[](https://fastapi.tiangolo.com/)
+[](#-プロジェクトを支援)
+
+*Kiro の Claude モデルを Claude Code、OpenCode、Cursor、Cline、Roo Code、Kilo Code、Obsidian、OpenAI SDK、LangChain、Continue などの OpenAI または Anthropic 互換ツールで使用*
+
+[モデル](#-対応モデル) • [機能](#-機能) • [クイックスタート](#-クイックスタート) • [設定](#%EF%B8%8F-設定) • [💖 サポート](#-プロジェクトを支援)
+
+
+
+---
+
+## 🤖 利用可能なモデル
+
+> ⚠️ **重要：** モデルの利用可能性は Kiro プラン（無料/有料）によって異なります。ゲートウェイは、サブスクリプションに基づいて IDE または CLI で利用可能なモデルへのアクセスを提供します。以下のリストは**無料プラン**で一般的に利用可能なモデルを示しています。
+
+> 🔒 **Claude Opus 4.5** は 2026年1月17日に無料プランから削除されました。有料プランで利用可能な場合があります — IDE/CLI のモデルリストを確認してください。
+
+🚀 **Claude Sonnet 4.5** — バランスの取れたパフォーマンス。コーディング、ライティング、汎用タスクに最適。
+
+⚡ **Claude Haiku 4.5** — 超高速。クイックレスポンス、シンプルなタスク、チャットに最適。
+
+📦 **Claude Sonnet 4** — 前世代モデル。ほとんどのユースケースで依然として強力で信頼性が高い。
+
+📦 **Claude 3.7 Sonnet** — レガシーモデル。後方互換性のために利用可能。
+
+> 💡 **スマートモデル解決:** どんなモデル名形式でも使用可能 — `claude-sonnet-4-5`、`claude-sonnet-4.5`、または `claude-sonnet-4-5-20250929` のようなバージョン付き名前も。ゲートウェイが自動的に正規化します。
+
+---
+
+## ✨ 機能
+
+| 機能 | 説明 |
+|------|------|
+| 🔌 **OpenAI 互換 API** | OpenAI 互換のあらゆるツールで動作 |
+| 🔌 **Anthropic 互換 API** | ネイティブ `/v1/messages` エンドポイント |
+| 🌐 **VPN/プロキシサポート** | 制限されたネットワーク向けの HTTP/SOCKS5 プロキシ |
+| 🧠 **拡張思考** | 推論機能は本プロジェクト独自の機能 |
+| 👁️ **ビジョンサポート** | モデルに画像を送信 |
+| 🛠️ **ツール呼び出し** | 関数呼び出しをサポート |
+| 💬 **完全なメッセージ履歴** | 完全な会話コンテキストを渡す |
+| 📡 **ストリーミング** | 完全な SSE ストリーミングサポート |
+| 🔄 **リトライロジック** | エラー時の自動リトライ（403、429、5xx） |
+| 📋 **拡張モデルリスト** | バージョン付きモデルを含む |
+| 🔐 **スマートトークン管理** | 有効期限前に自動更新 |
+
+---
+
+## 🚀 クイックスタート
+
+### 前提条件
+
+- Python 3.10+
+- 以下のいずれか：
+ - ログイン済みアカウントの [Kiro IDE](https://kiro.dev/)、または
+ - AWS SSO (AWS IAM Identity Center, OIDC) を使用した [Kiro CLI](https://kiro.dev/cli/) - 無料の Builder ID または企業アカウント
+
+### インストール
+
+```bash
+# リポジトリをクローン（Git が必要）
+git clone https://github.com/Jwadow/kiro-gateway.git
+cd kiro-gateway
+
+# または ZIP をダウンロード：Code → Download ZIP → 解凍 → kiro-gateway フォルダを開く
+
+# 依存関係をインストール
+pip install -r requirements.txt
+
+# 設定（設定セクションを参照）
+cp .env.example .env
+# .env をコピーして認証情報を編集
+
+# サーバーを起動
+python main.py
+
+# またはカスタムポートで（8000 が使用中の場合）
+python main.py --port 9000
+```
+
+サーバーは `http://localhost:8000` で利用可能になります
+
+---
+
+## ⚙️ 設定
+
+### オプション 1：JSON 認証情報ファイル (Kiro IDE / Enterprise)
+
+認証情報ファイルへのパスを指定：
+
+対応環境：
+- **Kiro IDE**（標準）- 個人アカウント用
+- **Enterprise** - SSO を使用した企業アカウント用
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/kiro-auth-token.json"
+
+# プロキシサーバーを保護するパスワード（任意の安全な文字列を設定）
+# ゲートウェイに接続する際に api_key として使用します
+PROXY_API_KEY="my-super-secret-password-123"
+```
+
+
+📄 JSON ファイル形式
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:...",
+ "region": "us-east-1",
+ "clientIdHash": "abc123..." // Optional: for corporate SSO setups
+}
+```
+
+> **注意：** `~/.aws/sso/cache/` に 2 つの JSON ファイルがある場合（例：`kiro-auth-token.json` とハッシュ名のファイル）、`KIRO_CREDS_FILE` で `kiro-auth-token.json` を使用してください。ゲートウェイが他のファイルを自動的に読み込みます。
+
+
+
+### オプション 2：環境変数（.env ファイル）
+
+プロジェクトルートに `.env` ファイルを作成：
+
+```env
+# 必須
+REFRESH_TOKEN="your_kiro_refresh_token"
+
+# プロキシサーバーを保護するパスワード（任意の安全な文字列を設定）
+PROXY_API_KEY="my-super-secret-password-123"
+
+# オプション
+PROFILE_ARN="arn:aws:codewhisperer:us-east-1:..."
+KIRO_REGION="us-east-1"
+```
+
+### オプション 3：AWS SSO 認証情報 (kiro-cli / Enterprise)
+
+AWS SSO (AWS IAM Identity Center) で `kiro-cli` または Kiro IDE を使用している場合、ゲートウェイは自動的に適切な認証を検出して使用します。
+
+無料の Builder ID アカウントと企業アカウントの両方で動作します。
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/your-sso-cache-file.json"
+
+# プロキシサーバーを保護するパスワード
+PROXY_API_KEY="my-super-secret-password-123"
+
+# 注意：AWS SSO (Builder ID および企業アカウント) ユーザーは PROFILE_ARN 不要
+# ゲートウェイはそれなしで動作します
+```
+
+
+📄 AWS SSO JSON ファイル形式
+
+AWS SSO 認証情報ファイル（`~/.aws/sso/cache/` から）には以下が含まれます：
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "region": "us-east-1",
+ "clientId": "...",
+ "clientSecret": "..."
+}
+```
+
+**注意：** AWS SSO (Builder ID および企業アカウント) ユーザーは `profileArn` 不要。ゲートウェイはそれなしで動作します（指定された場合は無視されます）。
+
+
+
+
+🔍 仕組み
+
+ゲートウェイは認証情報ファイルに基づいて認証タイプを自動検出します：
+
+- **Kiro Desktop Auth**（デフォルト）：`clientId` と `clientSecret` が存在しない場合に使用
+ - エンドポイント：`https://prod.{region}.auth.desktop.kiro.dev/refreshToken`
+
+- **AWS SSO (OIDC)**：`clientId` と `clientSecret` が存在する場合に使用
+ - エンドポイント：`https://oidc.{region}.amazonaws.com/token`
+
+追加設定は不要 — 認証情報ファイルを指定するだけ！
+
+
+
+### オプション 4：kiro-cli SQLite データベース
+
+`kiro-cli` を使用していて、その SQLite データベースを直接使用したい場合：
+
+```env
+KIRO_CLI_DB_FILE="~/.local/share/kiro-cli/data.sqlite3"
+
+# プロキシサーバーを保護するパスワード
+PROXY_API_KEY="my-super-secret-password-123"
+
+# 注意：AWS SSO (Builder ID および企業アカウント) ユーザーは PROFILE_ARN 不要
+# ゲートウェイはそれなしで動作します
+```
+
+
+📄 データベースの場所
+
+| CLI ツール | データベースパス |
+|-----------|-----------------|
+| kiro-cli | `~/.local/share/kiro-cli/data.sqlite3` |
+| amazon-q-developer-cli | `~/.local/share/amazon-q/data.sqlite3` |
+
+ゲートウェイは `auth_kv` テーブルから認証情報を読み取ります：
+- `kirocli:odic:token` または `codewhisperer:odic:token` — アクセストークン、リフレッシュトークン、有効期限
+- `kirocli:odic:device-registration` または `codewhisperer:odic:device-registration` — クライアント ID とシークレット
+
+異なる kiro-cli バージョンとの互換性のため、両方のキー形式がサポートされています。
+
+
+
+### 認証情報の取得
+
+**Kiro IDE ユーザー向け：**
+- Kiro IDE にログインして上記のオプション 1（JSON 認証情報ファイル）を使用
+- 認証情報ファイルはログイン後に自動作成されます
+
+**Kiro CLI ユーザー向け：**
+- `kiro-cli login` でログインして上記のオプション 3 または 4 を使用
+- 手動でのトークン抽出は不要！
+
+
+🔧 上級者向け：手動トークン抽出
+
+リフレッシュトークンを手動で抽出する必要がある場合（例：デバッグ用）、Kiro IDE のトラフィックをインターセプトできます：
+- 以下へのリクエストを探す：`prod.us-east-1.auth.desktop.kiro.dev/refreshToken`
+
+
+
+---
+
+## 🌐 VPN/プロキシサポート
+
+**中国、企業ネットワーク、または AWS サービスへの接続に問題がある地域のユーザー向け。**
+
+ゲートウェイは、すべての Kiro API リクエストを VPN またはプロキシサーバーでルーティングすることをサポートしています。AWS エンドポイントへの接続に問題が発生した場合、または企業プロキシを使用する必要がある場合に必須です。
+
+### 設定
+
+`.env` ファイルに追加：
+
+```env
+# HTTP プロキシ
+VPN_PROXY_URL=http://127.0.0.1:7890
+
+# SOCKS5 プロキシ
+VPN_PROXY_URL=socks5://127.0.0.1:1080
+
+# 認証付き（企業プロキシ）
+VPN_PROXY_URL=http://username:password@proxy.company.com:8080
+
+# プロトコルなし（デフォルトは http://）
+VPN_PROXY_URL=192.168.1.100:8080
+```
+
+### サポートされるプロトコル
+
+- ✅ **HTTP** — 標準プロキシプロトコル
+- ✅ **HTTPS** — セキュアプロキシ接続
+- ✅ **SOCKS5** — 高度なプロキシプロトコル（VPN ソフトウェアで一般的）
+- ✅ **認証** — URL に埋め込まれたユーザー名/パスワード
+
+### 必要な場合
+
+| 状況 | 解決策 |
+|------|--------|
+| AWS への接続タイムアウト | VPN/プロキシを使用してトラフィックをルーティング |
+| 企業ネットワーク制限 | 企業のプロキシを設定 |
+| 地域的な接続問題 | プロキシサポート付き VPN サービスを使用 |
+| プライバシー要件 | 独自のプロキシサーバーでルーティング |
+
+### プロキシサポート付きの人気 VPN ソフトウェア
+
+ほとんどの VPN クライアントはローカルプロキシサーバーを提供します：
+- **Sing-box** — HTTP/SOCKS5 プロキシサポート付きの最新 VPN クライアント
+- **Clash** — 通常 `http://127.0.0.1:7890` で実行
+- **V2Ray** — 設定可能な SOCKS5/HTTP プロキシ
+- **Shadowsocks** — SOCKS5 プロキシサポート
+- **企業 VPN** — プロキシ設定について IT 部門に確認
+
+プロキシサポートが不要な場合は、`VPN_PROXY_URL` を空のままにしてください（デフォルト）。
+
+---
+
+## 📡 API リファレンス
+
+### エンドポイント
+
+| エンドポイント | メソッド | 説明 |
+|---------------|---------|------|
+| `/` | GET | ヘルスチェック |
+| `/health` | GET | 詳細ヘルスチェック |
+| `/v1/models` | GET | 利用可能なモデル一覧 |
+| `/v1/chat/completions` | POST | OpenAI Chat Completions API |
+| `/v1/messages` | POST | Anthropic Messages API |
+
+---
+
+## 💡 使用例
+
+### OpenAI API
+
+
+🔹 シンプルな cURL リクエスト
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "こんにちは！"}],
+ "stream": true
+ }'
+```
+
+> **注意：** `my-super-secret-password-123` を `.env` ファイルで設定した `PROXY_API_KEY` に置き換えてください。
+
+
+
+
+🔹 ストリーミングリクエスト
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {"role": "system", "content": "あなたは親切なアシスタントです。"},
+ {"role": "user", "content": "2+2 は何ですか？"}
+ ],
+ "stream": true
+ }'
+```
+
+
+
+
+🛠️ ツール呼び出し付き
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "ロンドンの天気は？"}],
+ "tools": [{
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "場所の天気を取得",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "都市名"}
+ },
+ "required": ["location"]
+ }
+ }
+ }]
+ }'
+```
+
+
+
+
+🐍 Python OpenAI SDK
+
+```python
+from openai import OpenAI
+
+client = OpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123" # .env の PROXY_API_KEY
+)
+
+response = client.chat.completions.create(
+ model="claude-sonnet-4-5",
+ messages=[
+ {"role": "system", "content": "あなたは親切なアシスタントです。"},
+ {"role": "user", "content": "こんにちは！"}
+ ],
+ stream=True
+)
+
+for chunk in response:
+ if chunk.choices[0].delta.content:
+ print(chunk.choices[0].delta.content, end="")
+```
+
+
+
+
+🦜 LangChain
+
+```python
+from langchain_openai import ChatOpenAI
+
+llm = ChatOpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123", # .env の PROXY_API_KEY
+ model="claude-sonnet-4-5"
+)
+
+response = llm.invoke("こんにちは、お元気ですか？")
+print(response.content)
+```
+
+
+
+### Anthropic API
+
+
+🔹 シンプルな cURL リクエスト
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "こんにちは！"}]
+ }'
+```
+
+> **注意：** Anthropic API は `Authorization: Bearer` の代わりに `x-api-key` ヘッダーを使用します。両方サポートされています。
+
+
+
+
+🔹 システムプロンプト付き
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "system": "あなたは親切なアシスタントです。",
+ "messages": [{"role": "user", "content": "こんにちは！"}]
+ }'
+```
+
+> **注意：** Anthropic API では `system` はメッセージではなく別のフィールドです。
+
+
+
+
+📡 ストリーミング
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "stream": true,
+ "messages": [{"role": "user", "content": "こんにちは！"}]
+ }'
+```
+
+
+
+
+🐍 Python Anthropic SDK
+
+```python
+import anthropic
+
+client = anthropic.Anthropic(
+ api_key="my-super-secret-password-123", # .env の PROXY_API_KEY
+ base_url="http://localhost:8000"
+)
+
+# 非ストリーミング
+response = client.messages.create(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "こんにちは！"}]
+)
+print(response.content[0].text)
+
+# ストリーミング
+with client.messages.stream(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "こんにちは！"}]
+) as stream:
+ for text in stream.text_stream:
+ print(text, end="", flush=True)
+```
+
+
+
+---
+
+## 🔧 デバッグ
+
+デバッグログは**デフォルトで無効**です。有効にするには `.env` に追加：
+
+```env
+# デバッグログモード：
+# - off：無効（デフォルト）
+# - errors：失敗したリクエストのみログを保存（4xx、5xx）- トラブルシューティングに推奨
+# - all：すべてのリクエストのログを保存（リクエストごとに上書き）
+DEBUG_MODE=errors
+```
+
+### デバッグモード
+
+| モード | 説明 | 用途 |
+|--------|------|------|
+| `off` | 無効（デフォルト） | 本番環境 |
+| `errors` | 失敗したリクエストのみログを保存（4xx、5xx） | **トラブルシューティングに推奨** |
+| `all` | すべてのリクエストのログを保存 | 開発/デバッグ |
+
+### デバッグファイル
+
+有効にすると、リクエストは `debug_logs/` フォルダにログされます：
+
+| ファイル | 説明 |
+|---------|------|
+| `request_body.json` | クライアントからの受信リクエスト（OpenAI 形式） |
+| `kiro_request_body.json` | Kiro API に送信されたリクエスト |
+| `response_stream_raw.txt` | Kiro からの生ストリーム |
+| `response_stream_modified.txt` | 変換されたストリーム（OpenAI 形式） |
+| `app_logs.txt` | リクエストのアプリケーションログ |
+| `error_info.json` | エラー詳細（エラー時のみ） |
+
+---
+
+## 📜 ライセンス
+
+このプロジェクトは **GNU Affero General Public License v3.0 (AGPL-3.0)** でライセンスされています。
+
+これは以下を意味します：
+- ✅ このソフトウェアを使用、変更、配布できます
+- ✅ 商用目的で使用できます
+- ⚠️ ソフトウェアを配布する際は**ソースコードを公開する必要があります**
+- ⚠️ **ネットワーク使用は配布です** — 変更したバージョンをサーバーで実行し、他者がそれと対話できるようにする場合、ソースコードを彼らに提供する必要があります
+- ⚠️ 変更は同じライセンスでリリースする必要があります
+
+完全なライセンステキストは [LICENSE](../../LICENSE) ファイルを参照してください。
+
+### なぜ AGPL-3.0？
+
+AGPL-3.0 は、このソフトウェアへの改善がコミュニティ全体に利益をもたらすことを保証します。このゲートウェイを変更してサービスとしてデプロイする場合、改善をユーザーと共有する必要があります。
+
+### コントリビューターライセンス契約 (CLA)
+
+このプロジェクトへの貢献を提出することで、[コントリビューターライセンス契約 (CLA)](../../CLA.md) の条件に同意したことになります。これにより以下が保証されます：
+- 貢献を提出する権利があること
+- メンテナーに貢献を使用および再ライセンスする権利を付与すること
+- プロジェクトが法的に保護されること
+
+---
+
+## 💖 プロジェクトを支援
+
+
+
+
+
+**このプロジェクトが時間やお金を節約したなら、支援をご検討ください！**
+
+すべての貢献がこのプロジェクトの維持と成長に役立ちます
+
+
+
+### 🤑 寄付
+
+[**☕ 一回限りの寄付**](https://app.lava.top/jwadow?tabId=donate) • [**💎 月額サポート**](https://app.lava.top/jwadow?tabId=subscriptions)
+
+
+
+### 🪙 または暗号通貨を送信
+
+| 通貨 | ネットワーク | アドレス |
+|:----:|:----------:|:--------|
+| **USDT** | TRC20 | `TSVtgRc9pkC1UgcbVeijBHjFmpkYHDRu26` |
+| **BTC** | Bitcoin | `12GZqxqpcBsqJ4Vf1YreLqwoMGvzBPgJq6` |
+| **ETH** | Ethereum | `0xc86eab3bba3bbaf4eb5b5fff8586f1460f1fd395` |
+| **SOL** | Solana | `9amykF7KibZmdaw66a1oqYJyi75fRqgdsqnG66AK3jvh` |
+| **TON** | TON | `UQBVh8T1H3GI7gd7b-_PPNnxHYYxptrcCVf3qQk5v41h3QTM` |
+
+
+
+**[⬆ トップに戻る](#-kiro-gateway)**
+
+
diff --git a/kiro-gateway/docs/ko/README.md b/kiro-gateway/docs/ko/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..ddb48a91433276c2bdeb402e79e7c75479dd5d46
--- /dev/null
+++ b/kiro-gateway/docs/ko/README.md
@@ -0,0 +1,626 @@
+
+
+# 👻 Kiro Gateway
+
+**Kiro API (Amazon Q Developer / AWS CodeWhisperer) 프록시 게이트웨이**
+
+[🇬🇧 English](../../README.md) • [🇷🇺 Русский](../ru/README.md) • [🇨🇳 中文](../zh/README.md) • [🇪🇸 Español](../es/README.md) • [🇮🇩 Indonesia](../id/README.md) • [🇧🇷 Português](../pt/README.md) • [🇯🇵 日本語](../ja/README.md)
+
+[@Jwadow](https://github.com/jwadow)가 ❤️를 담아 제작
+
+[](https://www.gnu.org/licenses/agpl-3.0)
+[](https://www.python.org/downloads/)
+[](https://fastapi.tiangolo.com/)
+[](#-프로젝트-후원)
+
+*Kiro의 Claude 모델을 Claude Code, OpenCode, Cursor, Cline, Roo Code, Kilo Code, Obsidian, OpenAI SDK, LangChain, Continue 및 기타 OpenAI 또는 Anthropic 호환 도구와 함께 사용*
+
+[모델](#-지원-모델) • [기능](#-기능) • [빠른-시작](#-빠른-시작) • [설정](#%EF%B8%8F-설정) • [💖 후원](#-프로젝트-후원)
+
+
+
+---
+
+## 🤖 사용 가능한 모델
+
+> ⚠️ **중요:** 모델 가용성은 Kiro 플랜(무료/유료)에 따라 다릅니다. 게이트웨이는 구독에 따라 IDE 또는 CLI에서 사용 가능한 모델에 대한 액세스를 제공합니다. 아래 목록은 **무료 플랜**에서 일반적으로 사용 가능한 모델을 보여줍니다.
+
+> 🔒 **Claude Opus 4.5**는 2026년 1월 17일에 무료 플랜에서 제거되었습니다. 유료 플랜에서 사용 가능할 수 있습니다 — IDE/CLI의 모델 목록을 확인하세요.
+
+🚀 **Claude Sonnet 4.5** — 균형 잡힌 성능. 코딩, 글쓰기, 범용 작업에 적합.
+
+⚡ **Claude Haiku 4.5** — 번개처럼 빠름. 빠른 응답, 간단한 작업, 채팅에 완벽.
+
+📦 **Claude Sonnet 4** — 이전 세대. 대부분의 사용 사례에서 여전히 강력하고 신뢰할 수 있음.
+
+📦 **Claude 3.7 Sonnet** — 레거시 모델. 하위 호환성을 위해 제공.
+
+> 💡 **스마트 모델 해석:** 어떤 모델 이름 형식이든 사용 가능 — `claude-sonnet-4-5`, `claude-sonnet-4.5`, 또는 `claude-sonnet-4-5-20250929`와 같은 버전 이름도. 게이트웨이가 자동으로 정규화합니다.
+
+---
+
+## ✨ 기능
+
+| 기능 | 설명 |
+|------|------|
+| 🔌 **OpenAI 호환 API** | OpenAI 호환 도구와 함께 작동 |
+| 🔌 **Anthropic 호환 API** | 네이티브 `/v1/messages` 엔드포인트 |
+| 🌐 **VPN/프록시 지원** | 제한된 네트워크용 HTTP/SOCKS5 프록시 |
+| 🧠 **확장 사고** | 추론 기능은 우리 프로젝트만의 독점 기능 |
+| 👁️ **비전 지원** | 모델에 이미지 전송 |
+| 🛠️ **도구 호출** | 함수 호출 지원 |
+| 💬 **전체 메시지 기록** | 완전한 대화 컨텍스트 전달 |
+| 📡 **스트리밍** | 완전한 SSE 스트리밍 지원 |
+| 🔄 **재시도 로직** | 오류 시 자동 재시도 (403, 429, 5xx) |
+| 📋 **확장 모델 목록** | 버전 모델 포함 |
+| 🔐 **스마트 토큰 관리** | 만료 전 자동 갱신 |
+
+---
+
+## 🚀 빠른 시작
+
+### 사전 요구 사항
+
+- Python 3.10+
+- 다음 중 하나:
+ - 로그인된 계정이 있는 [Kiro IDE](https://kiro.dev/), 또는
+ - AWS SSO (AWS IAM Identity Center, OIDC)가 있는 [Kiro CLI](https://kiro.dev/cli/) - 무료 Builder ID 또는 기업 계정
+
+### 설치
+
+```bash
+# 저장소 클론 (Git 필요)
+git clone https://github.com/Jwadow/kiro-gateway.git
+cd kiro-gateway
+
+# 또는 ZIP 다운로드: Code → Download ZIP → 압축 해제 → kiro-gateway 폴더 열기
+
+# 의존성 설치
+pip install -r requirements.txt
+
+# 설정 (설정 섹션 참조)
+cp .env.example .env
+# .env를 복사하고 자격 증명으로 편집
+
+# 서버 시작
+python main.py
+
+# 또는 사용자 정의 포트로 (8000이 사용 중인 경우)
+python main.py --port 9000
+```
+
+서버는 `http://localhost:8000`에서 사용 가능합니다
+
+---
+
+## ⚙️ 설정
+
+### 옵션 1: JSON 자격 증명 파일 (Kiro IDE / Enterprise)
+
+자격 증명 파일 경로 지정:
+
+다음과 함께 작동:
+- **Kiro IDE** (표준) - 개인 계정용
+- **Enterprise** - SSO가 있는 기업 계정용
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/kiro-auth-token.json"
+
+# 프록시 서버를 보호하는 비밀번호 (안전한 문자열 설정)
+# 게이트웨이에 연결할 때 api_key로 사용합니다
+PROXY_API_KEY="my-super-secret-password-123"
+```
+
+
+📄 JSON 파일 형식
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:...",
+ "region": "us-east-1",
+ "clientIdHash": "abc123..." // Optional: for corporate SSO setups
+}
+```
+
+> **참고:** `~/.aws/sso/cache/`에 두 개의 JSON 파일이 있는 경우 (예: `kiro-auth-token.json` 및 해시 이름의 파일), `KIRO_CREDS_FILE`에서 `kiro-auth-token.json`을 사용하세요. 게이트웨이가 다른 파일을 자동으로 로드합니다.
+
+
+
+### 옵션 2: 환경 변수 (.env 파일)
+
+프로젝트 루트에 `.env` 파일 생성:
+
+```env
+# 필수
+REFRESH_TOKEN="your_kiro_refresh_token"
+
+# 프록시 서버를 보호하는 비밀번호 (안전한 문자열 설정)
+PROXY_API_KEY="my-super-secret-password-123"
+
+# 선택 사항
+PROFILE_ARN="arn:aws:codewhisperer:us-east-1:..."
+KIRO_REGION="us-east-1"
+```
+
+### 옵션 3: AWS SSO 자격 증명 (kiro-cli / Enterprise)
+
+AWS SSO (AWS IAM Identity Center)와 함께 `kiro-cli` 또는 Kiro IDE를 사용하는 경우, 게이트웨이가 자동으로 적절한 인증을 감지하고 사용합니다.
+
+무료 Builder ID 계정과 기업 계정 모두에서 작동합니다.
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/your-sso-cache-file.json"
+
+# 프록시 서버를 보호하는 비밀번호
+PROXY_API_KEY="my-super-secret-password-123"
+
+# 참고: AWS SSO (Builder ID 및 기업 계정) 사용자는 PROFILE_ARN 불필요
+# 게이트웨이는 그것 없이도 작동합니다
+```
+
+
+📄 AWS SSO JSON 파일 형식
+
+AWS SSO 자격 증명 파일 (`~/.aws/sso/cache/`에서)에는 다음이 포함됩니다:
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "region": "us-east-1",
+ "clientId": "...",
+ "clientSecret": "..."
+}
+```
+
+**참고:** AWS SSO (Builder ID 및 기업 계정) 사용자는 `profileArn`이 필요 없습니다. 게이트웨이는 그것 없이도 작동합니다 (지정된 경우 무시됨).
+
+
+
+
+🔍 작동 방식
+
+게이트웨이는 자격 증명 파일을 기반으로 인증 유형을 자동 감지합니다:
+
+- **Kiro Desktop Auth** (기본값): `clientId`와 `clientSecret`이 없을 때 사용
+ - 엔드포인트: `https://prod.{region}.auth.desktop.kiro.dev/refreshToken`
+
+- **AWS SSO (OIDC)**: `clientId`와 `clientSecret`이 있을 때 사용
+ - 엔드포인트: `https://oidc.{region}.amazonaws.com/token`
+
+추가 설정 불필요 — 자격 증명 파일만 지정하면 됩니다!
+
+
+
+### 옵션 4: kiro-cli SQLite 데이터베이스
+
+`kiro-cli`를 사용하고 SQLite 데이터베이스를 직접 사용하려는 경우:
+
+```env
+KIRO_CLI_DB_FILE="~/.local/share/kiro-cli/data.sqlite3"
+
+# 프록시 서버를 보호하는 비밀번호
+PROXY_API_KEY="my-super-secret-password-123"
+
+# 참고: AWS SSO (Builder ID 및 기업 계정) 사용자는 PROFILE_ARN 불필요
+# 게이트웨이는 그것 없이도 작동합니다
+```
+
+
+📄 데이터베이스 위치
+
+| CLI 도구 | 데이터베이스 경로 |
+|----------|------------------|
+| kiro-cli | `~/.local/share/kiro-cli/data.sqlite3` |
+| amazon-q-developer-cli | `~/.local/share/amazon-q/data.sqlite3` |
+
+게이트웨이는 `auth_kv` 테이블에서 자격 증명을 읽습니다:
+- `kirocli:odic:token` 또는 `codewhisperer:odic:token` — 액세스 토큰, 리프레시 토큰, 만료 시간
+- `kirocli:odic:device-registration` 또는 `codewhisperer:odic:device-registration` — 클라이언트 ID와 시크릿
+
+다양한 kiro-cli 버전과의 호환성을 위해 두 키 형식 모두 지원됩니다.
+
+
+
+### 자격 증명 얻기
+
+**Kiro IDE 사용자:**
+- Kiro IDE에 로그인하고 위의 옵션 1 (JSON 자격 증명 파일) 사용
+- 자격 증명 파일은 로그인 후 자동 생성됩니다
+
+**Kiro CLI 사용자:**
+- `kiro-cli login`으로 로그인하고 위의 옵션 3 또는 4 사용
+- 수동 토큰 추출 불필요!
+
+
+🔧 고급: 수동 토큰 추출
+
+리프레시 토큰을 수동으로 추출해야 하는 경우 (예: 디버깅용), Kiro IDE 트래픽을 가로챌 수 있습니다:
+- 다음으로의 요청 찾기: `prod.us-east-1.auth.desktop.kiro.dev/refreshToken`
+
+
+
+---
+
+## 🌐 VPN/프록시 지원
+
+**중국, 기업 네트워크 또는 AWS 서비스 연결에 문제가 있는 지역의 사용자를 위한 것입니다.**
+
+게이트웨이는 모든 Kiro API 요청을 VPN 또는 프록시 서버를 통해 라우팅하는 것을 지원합니다. AWS 엔드포인트에 대한 연결 문제가 발생하거나 기업 프록시를 사용해야 하는 경우 필수입니다.
+
+### 설정
+
+`.env` 파일에 추가:
+
+```env
+# HTTP 프록시
+VPN_PROXY_URL=http://127.0.0.1:7890
+
+# SOCKS5 프록시
+VPN_PROXY_URL=socks5://127.0.0.1:1080
+
+# 인증 포함 (기업 프록시)
+VPN_PROXY_URL=http://username:password@proxy.company.com:8080
+
+# 프로토콜 없음 (기본값 http://)
+VPN_PROXY_URL=192.168.1.100:8080
+```
+
+### 지원되는 프로토콜
+
+- ✅ **HTTP** — 표준 프록시 프로토콜
+- ✅ **HTTPS** — 보안 프록시 연결
+- ✅ **SOCKS5** — 고급 프록시 프로토콜 (VPN 소프트웨어에서 일반적)
+- ✅ **인증** — URL에 포함된 사용자명/비밀번호
+
+### 필요한 경우
+
+| 상황 | 해결책 |
+|------|--------|
+| AWS 연결 타임아웃 | VPN/프록시를 사용하여 트래픽 라우팅 |
+| 기업 네트워크 제한 | 회사 프록시 구성 |
+| 지역 연결 문제 | 프록시 지원이 있는 VPN 서비스 사용 |
+| 개인정보 보호 요구사항 | 자신의 프록시 서버를 통해 라우팅 |
+
+### 프록시 지원이 있는 인기 VPN 소프트웨어
+
+대부분의 VPN 클라이언트는 로컬 프록시 서버를 제공합니다:
+- **Sing-box** — HTTP/SOCKS5 프록시 지원이 있는 최신 VPN 클라이언트
+- **Clash** — 일반적으로 `http://127.0.0.1:7890`에서 실행
+- **V2Ray** — 구성 가능한 SOCKS5/HTTP 프록시
+- **Shadowsocks** — SOCKS5 프록시 지원
+- **기업 VPN** — 프록시 설정에 대해 IT 부서에 문의
+
+프록시 지원이 필요하지 않으면 `VPN_PROXY_URL`을 비워두세요 (기본값).
+
+---
+
+## 📡 API 레퍼런스
+
+### 엔드포인트
+
+| 엔드포인트 | 메서드 | 설명 |
+|-----------|--------|------|
+| `/` | GET | 헬스 체크 |
+| `/health` | GET | 상세 헬스 체크 |
+| `/v1/models` | GET | 사용 가능한 모델 목록 |
+| `/v1/chat/completions` | POST | OpenAI Chat Completions API |
+| `/v1/messages` | POST | Anthropic Messages API |
+
+---
+
+## 💡 사용 예시
+
+### OpenAI API
+
+
+🔹 간단한 cURL 요청
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "안녕하세요!"}],
+ "stream": true
+ }'
+```
+
+> **참고:** `my-super-secret-password-123`을 `.env` 파일에 설정한 `PROXY_API_KEY`로 교체하세요.
+
+
+
+
+🔹 스트리밍 요청
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {"role": "system", "content": "당신은 도움이 되는 어시스턴트입니다."},
+ {"role": "user", "content": "2+2는 얼마인가요?"}
+ ],
+ "stream": true
+ }'
+```
+
+
+
+
+🛠️ 도구 호출 포함
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "런던 날씨는 어때요?"}],
+ "tools": [{
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "위치의 날씨 가져오기",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "도시 이름"}
+ },
+ "required": ["location"]
+ }
+ }
+ }]
+ }'
+```
+
+
+
+
+🐍 Python OpenAI SDK
+
+```python
+from openai import OpenAI
+
+client = OpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123" # .env의 PROXY_API_KEY
+)
+
+response = client.chat.completions.create(
+ model="claude-sonnet-4-5",
+ messages=[
+ {"role": "system", "content": "당신은 도움이 되는 어시스턴트입니다."},
+ {"role": "user", "content": "안녕하세요!"}
+ ],
+ stream=True
+)
+
+for chunk in response:
+ if chunk.choices[0].delta.content:
+ print(chunk.choices[0].delta.content, end="")
+```
+
+
+
+
+🦜 LangChain
+
+```python
+from langchain_openai import ChatOpenAI
+
+llm = ChatOpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123", # .env의 PROXY_API_KEY
+ model="claude-sonnet-4-5"
+)
+
+response = llm.invoke("안녕하세요, 어떻게 지내세요?")
+print(response.content)
+```
+
+
+
+### Anthropic API
+
+
+🔹 간단한 cURL 요청
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "안녕하세요!"}]
+ }'
+```
+
+> **참고:** Anthropic API는 `Authorization: Bearer` 대신 `x-api-key` 헤더를 사용합니다. 둘 다 지원됩니다.
+
+
+
+
+🔹 시스템 프롬프트 포함
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "system": "당신은 도움이 되는 어시스턴트입니다.",
+ "messages": [{"role": "user", "content": "안녕하세요!"}]
+ }'
+```
+
+> **참고:** Anthropic API에서 `system`은 메시지가 아닌 별도의 필드입니다.
+
+
+
+
+📡 스트리밍
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "stream": true,
+ "messages": [{"role": "user", "content": "안녕하세요!"}]
+ }'
+```
+
+
+
+
+🐍 Python Anthropic SDK
+
+```python
+import anthropic
+
+client = anthropic.Anthropic(
+ api_key="my-super-secret-password-123", # .env의 PROXY_API_KEY
+ base_url="http://localhost:8000"
+)
+
+# 비스트리밍
+response = client.messages.create(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "안녕하세요!"}]
+)
+print(response.content[0].text)
+
+# 스트리밍
+with client.messages.stream(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "안녕하세요!"}]
+) as stream:
+ for text in stream.text_stream:
+ print(text, end="", flush=True)
+```
+
+
+
+---
+
+## 🔧 디버깅
+
+디버그 로깅은 **기본적으로 비활성화**되어 있습니다. 활성화하려면 `.env`에 추가:
+
+```env
+# 디버그 로깅 모드:
+# - off: 비활성화 (기본값)
+# - errors: 실패한 요청만 로그 저장 (4xx, 5xx) - 문제 해결에 권장
+# - all: 모든 요청 로그 저장 (요청마다 덮어쓰기)
+DEBUG_MODE=errors
+```
+
+### 디버그 모드
+
+| 모드 | 설명 | 용도 |
+|------|------|------|
+| `off` | 비활성화 (기본값) | 프로덕션 |
+| `errors` | 실패한 요청만 로그 저장 (4xx, 5xx) | **문제 해결에 권장** |
+| `all` | 모든 요청 로그 저장 | 개발/디버깅 |
+
+### 디버그 파일
+
+활성화되면 요청이 `debug_logs/` 폴더에 기록됩니다:
+
+| 파일 | 설명 |
+|------|------|
+| `request_body.json` | 클라이언트로부터의 수신 요청 (OpenAI 형식) |
+| `kiro_request_body.json` | Kiro API로 전송된 요청 |
+| `response_stream_raw.txt` | Kiro로부터의 원시 스트림 |
+| `response_stream_modified.txt` | 변환된 스트림 (OpenAI 형식) |
+| `app_logs.txt` | 요청에 대한 애플리케이션 로그 |
+| `error_info.json` | 오류 세부 정보 (오류 시에만) |
+
+---
+
+## 📜 라이선스
+
+이 프로젝트는 **GNU Affero General Public License v3.0 (AGPL-3.0)**으로 라이선스됩니다.
+
+이것은 다음을 의미합니다:
+- ✅ 이 소프트웨어를 사용, 수정, 배포할 수 있습니다
+- ✅ 상업적 목적으로 사용할 수 있습니다
+- ⚠️ 소프트웨어를 배포할 때 **소스 코드를 공개해야 합니다**
+- ⚠️ **네트워크 사용은 배포입니다** — 수정된 버전을 서버에서 실행하고 다른 사람이 상호 작용할 수 있게 하면 소스 코드를 그들에게 제공해야 합니다
+- ⚠️ 수정 사항은 동일한 라이선스로 릴리스해야 합니다
+
+전체 라이선스 텍스트는 [LICENSE](../../LICENSE) 파일을 참조하세요.
+
+### 왜 AGPL-3.0인가?
+
+AGPL-3.0은 이 소프트웨어에 대한 개선이 전체 커뮤니티에 이익이 되도록 보장합니다. 이 게이트웨이를 수정하고 서비스로 배포하는 경우 사용자와 개선 사항을 공유해야 합니다.
+
+### 기여자 라이선스 계약 (CLA)
+
+이 프로젝트에 기여를 제출함으로써 [기여자 라이선스 계약 (CLA)](../../CLA.md)의 조건에 동의하게 됩니다. 이것은 다음을 보장합니다:
+- 기여를 제출할 권리가 있음
+- 메인테이너에게 기여를 사용하고 재라이선스할 권리를 부여함
+- 프로젝트가 법적으로 보호됨
+
+---
+
+## 💖 프로젝트 후원
+
+
+
+
+
+**이 프로젝트가 시간이나 돈을 절약해 주었다면 후원을 고려해 주세요!**
+
+모든 기여가 이 프로젝트를 유지하고 성장시키는 데 도움이 됩니다
+
+
+
+### 🤑 기부
+
+[**☕ 일회성 기부**](https://app.lava.top/jwadow?tabId=donate) • [**💎 월간 후원**](https://app.lava.top/jwadow?tabId=subscriptions)
+
+
+
+### 🪙 또는 암호화폐 전송
+
+| 통화 | 네트워크 | 주소 |
+|:----:|:-------:|:-----|
+| **USDT** | TRC20 | `TSVtgRc9pkC1UgcbVeijBHjFmpkYHDRu26` |
+| **BTC** | Bitcoin | `12GZqxqpcBsqJ4Vf1YreLqwoMGvzBPgJq6` |
+| **ETH** | Ethereum | `0xc86eab3bba3bbaf4eb5b5fff8586f1460f1fd395` |
+| **SOL** | Solana | `9amykF7KibZmdaw66a1oqYJyi75fRqgdsqnG66AK3jvh` |
+| **TON** | TON | `UQBVh8T1H3GI7gd7b-_PPNnxHYYxptrcCVf3qQk5v41h3QTM` |
+
+
+
+**[⬆ 맨 위로](#-kiro-gateway)**
+
+
diff --git a/kiro-gateway/docs/pt/README.md b/kiro-gateway/docs/pt/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..ef3828feaa0f54726e79332f94bae5faec5d52cc
--- /dev/null
+++ b/kiro-gateway/docs/pt/README.md
@@ -0,0 +1,626 @@
+
+
+# 👻 Kiro Gateway
+
+**Gateway proxy para Kiro API (Amazon Q Developer / AWS CodeWhisperer)**
+
+[🇬🇧 English](../../README.md) • [🇷🇺 Русский](../ru/README.md) • [🇨🇳 中文](../zh/README.md) • [🇪🇸 Español](../es/README.md) • [🇮🇩 Indonesia](../id/README.md) • [🇯🇵 日本語](../ja/README.md) • [🇰🇷 한국어](../ko/README.md)
+
+Feito com ❤️ por [@Jwadow](https://github.com/jwadow)
+
+[](https://www.gnu.org/licenses/agpl-3.0)
+[](https://www.python.org/downloads/)
+[](https://fastapi.tiangolo.com/)
+[](#-apoie-o-projeto)
+
+*Use modelos Claude do Kiro com Claude Code, OpenCode, Cursor, Cline, Roo Code, Kilo Code, Obsidian, OpenAI SDK, LangChain, Continue e outras ferramentas compatíveis com OpenAI ou Anthropic*
+
+[Modelos](#-modelos-suportados) • [Recursos](#-recursos) • [Início Rápido](#-início-rápido) • [Configuração](#%EF%B8%8F-configuração) • [💖 Apoiar](#-apoie-o-projeto)
+
+
+
+---
+
+## 🤖 Modelos Disponíveis
+
+> ⚠️ **Importante:** A disponibilidade de modelos depende do seu plano Kiro (gratuito/pago). O gateway fornece acesso aos modelos disponíveis no seu IDE ou CLI com base na sua assinatura. A lista abaixo mostra os modelos comumente disponíveis no **plano gratuito**.
+
+> 🔒 **Claude Opus 4.5** foi removido do plano gratuito em 17 de janeiro de 2026. Pode estar disponível em planos pagos — verifique a lista de modelos no seu IDE/CLI.
+
+🚀 **Claude Sonnet 4.5** — Desempenho equilibrado. Ótimo para programação, escrita e tarefas de uso geral.
+
+⚡ **Claude Haiku 4.5** — Velocidade relâmpago. Perfeito para respostas rápidas, tarefas simples e chat.
+
+📦 **Claude Sonnet 4** — Geração anterior. Ainda poderoso e confiável para a maioria dos casos de uso.
+
+📦 **Claude 3.7 Sonnet** — Modelo legado. Disponível para compatibilidade retroativa.
+
+> 💡 **Resolução Inteligente de Modelos:** Use qualquer formato de nome de modelo — `claude-sonnet-4-5`, `claude-sonnet-4.5`, ou até nomes versionados como `claude-sonnet-4-5-20250929`. O gateway normaliza automaticamente.
+
+---
+
+## ✨ Recursos
+
+| Recurso | Descrição |
+|---------|-----------|
+| 🔌 **API compatível com OpenAI** | Funciona com qualquer ferramenta compatível com OpenAI |
+| 🔌 **API compatível com Anthropic** | Endpoint nativo `/v1/messages` |
+| 🌐 **Suporte a VPN/Proxy** | Proxy HTTP/SOCKS5 para redes restritas |
+| 🧠 **Pensamento Estendido** | Raciocínio é exclusivo do nosso projeto |
+| 👁️ **Suporte a Visão** | Envie imagens para o modelo |
+| 🛠️ **Chamada de Ferramentas** | Suporta chamada de funções |
+| 💬 **Histórico completo de mensagens** | Passa o contexto completo da conversa |
+| 📡 **Streaming** | Suporte completo a streaming SSE |
+| 🔄 **Lógica de Retry** | Retentativas automáticas em erros (403, 429, 5xx) |
+| 📋 **Lista estendida de modelos** | Incluindo modelos versionados |
+| 🔐 **Gerenciamento inteligente de tokens** | Atualização automática antes da expiração |
+
+---
+
+## 🚀 Início Rápido
+
+### Pré-requisitos
+
+- Python 3.10+
+- Um dos seguintes:
+ - [Kiro IDE](https://kiro.dev/) com conta logada, OU
+ - [Kiro CLI](https://kiro.dev/cli/) com AWS SSO (AWS IAM Identity Center, OIDC) - Builder ID gratuito ou conta corporativa
+
+### Instalação
+
+```bash
+# Clone o repositório (requer Git)
+git clone https://github.com/Jwadow/kiro-gateway.git
+cd kiro-gateway
+
+# Ou baixe o ZIP: Code → Download ZIP → extraia → abra a pasta kiro-gateway
+
+# Instale as dependências
+pip install -r requirements.txt
+
+# Configure (veja a seção Configuração)
+cp .env.example .env
+# Copie e edite o .env com suas credenciais
+
+# Inicie o servidor
+python main.py
+
+# Ou com porta personalizada (se 8000 estiver ocupada)
+python main.py --port 9000
+```
+
+O servidor estará disponível em `http://localhost:8000`
+
+---
+
+## ⚙️ Configuração
+
+### Opção 1: Arquivo JSON de Credenciais (Kiro IDE / Enterprise)
+
+Especifique o caminho para o arquivo de credenciais:
+
+Funciona com:
+- **Kiro IDE** (padrão) - para contas pessoais
+- **Enterprise** - para contas corporativas com SSO
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/kiro-auth-token.json"
+
+# Senha para proteger SEU servidor proxy (crie qualquer string segura)
+# Você usará isso como api_key ao conectar ao seu gateway
+PROXY_API_KEY="my-super-secret-password-123"
+```
+
+
+📄 Formato do arquivo JSON
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:...",
+ "region": "us-east-1",
+ "clientIdHash": "abc123..." // Optional: for corporate SSO setups
+}
+```
+
+> **Nota:** Se você tiver dois arquivos JSON em `~/.aws/sso/cache/` (por exemplo, `kiro-auth-token.json` e um arquivo com nome hash), use `kiro-auth-token.json` em `KIRO_CREDS_FILE`. O gateway carregará automaticamente o outro arquivo.
+
+
+
+### Opção 2: Variáveis de Ambiente (arquivo .env)
+
+Crie um arquivo `.env` na raiz do projeto:
+
+```env
+# Obrigatório
+REFRESH_TOKEN="seu_kiro_refresh_token"
+
+# Senha para proteger SEU servidor proxy (crie qualquer string segura)
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Opcional
+PROFILE_ARN="arn:aws:codewhisperer:us-east-1:..."
+KIRO_REGION="us-east-1"
+```
+
+### Opção 3: Credenciais AWS SSO (kiro-cli / Enterprise)
+
+Se você usa `kiro-cli` ou Kiro IDE com AWS SSO (AWS IAM Identity Center), o gateway detectará e usará automaticamente a autenticação apropriada.
+
+Funciona tanto com contas Builder ID gratuitas quanto com contas corporativas.
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/your-sso-cache-file.json"
+
+# Senha para proteger SEU servidor proxy
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Nota: PROFILE_ARN NÃO é necessário para AWS SSO (Builder ID e contas corporativas)
+# O gateway funcionará sem ele
+```
+
+
+📄 Formato do arquivo JSON AWS SSO
+
+Arquivos de credenciais AWS SSO (de `~/.aws/sso/cache/`) contêm:
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "region": "us-east-1",
+ "clientId": "...",
+ "clientSecret": "..."
+}
+```
+
+**Nota:** Usuários AWS SSO (Builder ID e contas corporativas) NÃO precisam de `profileArn`. O gateway funcionará sem ele (se especificado, será ignorado).
+
+
+
+
+🔍 Como funciona
+
+O gateway detecta automaticamente o tipo de autenticação com base no arquivo de credenciais:
+
+- **Kiro Desktop Auth** (padrão): Usado quando `clientId` e `clientSecret` NÃO estão presentes
+ - Endpoint: `https://prod.{region}.auth.desktop.kiro.dev/refreshToken`
+
+- **AWS SSO (OIDC)**: Usado quando `clientId` e `clientSecret` estão presentes
+ - Endpoint: `https://oidc.{region}.amazonaws.com/token`
+
+Nenhuma configuração adicional necessária — apenas aponte para seu arquivo de credenciais!
+
+
+
+### Opção 4: Banco de dados SQLite do kiro-cli
+
+Se você usa `kiro-cli` e prefere usar seu banco de dados SQLite diretamente:
+
+```env
+KIRO_CLI_DB_FILE="~/.local/share/kiro-cli/data.sqlite3"
+
+# Senha para proteger SEU servidor proxy
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Nota: PROFILE_ARN NÃO é necessário para AWS SSO (Builder ID e contas corporativas)
+# O gateway funcionará sem ele
+```
+
+
+📄 Localizações do banco de dados
+
+| Ferramenta CLI | Caminho do Banco de Dados |
+|----------------|---------------------------|
+| kiro-cli | `~/.local/share/kiro-cli/data.sqlite3` |
+| amazon-q-developer-cli | `~/.local/share/amazon-q/data.sqlite3` |
+
+O gateway lê credenciais da tabela `auth_kv` que armazena:
+- `kirocli:odic:token` ou `codewhisperer:odic:token` — token de acesso, token de atualização, expiração
+- `kirocli:odic:device-registration` ou `codewhisperer:odic:device-registration` — ID do cliente e segredo
+
+Ambos os formatos de chave são suportados para compatibilidade com diferentes versões do kiro-cli.
+
+
+
+### Obtendo Credenciais
+
+**Para usuários do Kiro IDE:**
+- Faça login no Kiro IDE e use a Opção 1 acima (arquivo JSON de credenciais)
+- O arquivo de credenciais é criado automaticamente após o login
+
+**Para usuários do Kiro CLI:**
+- Faça login com `kiro-cli login` e use a Opção 3 ou Opção 4 acima
+- Não é necessário extrair tokens manualmente!
+
+
+🔧 Avançado: Extração manual de token
+
+Se você precisar extrair manualmente o refresh token (por exemplo, para depuração), você pode interceptar o tráfego do Kiro IDE:
+- Procure por requisições para: `prod.us-east-1.auth.desktop.kiro.dev/refreshToken`
+
+
+
+---
+
+## 🌐 Suporte a VPN/Proxy
+
+**Para usuários na China, redes corporativas ou regiões com problemas de conectividade com serviços AWS.**
+
+O gateway suporta rotear todas as solicitações da Kiro API através de um servidor VPN ou proxy. Isso é essencial se você enfrentar problemas de conexão com endpoints AWS ou precisar usar um proxy corporativo.
+
+### Configuração
+
+Adicione ao seu arquivo `.env`:
+
+```env
+# Proxy HTTP
+VPN_PROXY_URL=http://127.0.0.1:7890
+
+# Proxy SOCKS5
+VPN_PROXY_URL=socks5://127.0.0.1:1080
+
+# Com autenticação (proxies corporativos)
+VPN_PROXY_URL=http://username:password@proxy.company.com:8080
+
+# Sem protocolo (padrão para http://)
+VPN_PROXY_URL=192.168.1.100:8080
+```
+
+### Protocolos Suportados
+
+- ✅ **HTTP** — Protocolo proxy padrão
+- ✅ **HTTPS** — Conexões proxy seguras
+- ✅ **SOCKS5** — Protocolo proxy avançado (comum em software VPN)
+- ✅ **Autenticação** — Nome de usuário/senha incorporados na URL
+
+### Quando Você Precisa Disso
+
+| Situação | Solução |
+|----------|---------|
+| Timeouts de conexão com AWS | Use VPN/proxy para rotear tráfego |
+| Restrições de rede corporativa | Configure o proxy da sua empresa |
+| Problemas de conectividade regional | Use um serviço VPN com suporte a proxy |
+| Requisitos de privacidade | Roteie através do seu próprio servidor proxy |
+
+### Software VPN Popular com Suporte a Proxy
+
+A maioria dos clientes VPN fornece um servidor proxy local:
+- **Sing-box** — Cliente VPN moderno com suporte a proxy HTTP/SOCKS5
+- **Clash** — Geralmente executado em `http://127.0.0.1:7890`
+- **V2Ray** — Proxy SOCKS5/HTTP configurável
+- **Shadowsocks** — Suporte a proxy SOCKS5
+- **VPN Corporativo** — Consulte seu departamento de TI para configurações de proxy
+
+Deixe `VPN_PROXY_URL` vazio (padrão) se você não precisar de suporte a proxy.
+
+---
+
+## 📡 Referência da API
+
+### Endpoints
+
+| Endpoint | Método | Descrição |
+|----------|--------|-----------|
+| `/` | GET | Verificação de saúde |
+| `/health` | GET | Verificação de saúde detalhada |
+| `/v1/models` | GET | Lista modelos disponíveis |
+| `/v1/chat/completions` | POST | OpenAI Chat Completions API |
+| `/v1/messages` | POST | Anthropic Messages API |
+
+---
+
+## 💡 Exemplos de Uso
+
+### OpenAI API
+
+
+🔹 Requisição cURL Simples
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Olá!"}],
+ "stream": true
+ }'
+```
+
+> **Nota:** Substitua `my-super-secret-password-123` pelo `PROXY_API_KEY` que você definiu no arquivo `.env`.
+
+
+
+
+🔹 Requisição com Streaming
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {"role": "system", "content": "Você é um assistente útil."},
+ {"role": "user", "content": "Quanto é 2+2?"}
+ ],
+ "stream": true
+ }'
+```
+
+
+
+
+🛠️ Com Chamada de Ferramentas
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Como está o tempo em Londres?"}],
+ "tools": [{
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Obter o tempo para uma localização",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "Nome da cidade"}
+ },
+ "required": ["location"]
+ }
+ }
+ }]
+ }'
+```
+
+
+
+
+🐍 Python OpenAI SDK
+
+```python
+from openai import OpenAI
+
+client = OpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123" # Seu PROXY_API_KEY do .env
+)
+
+response = client.chat.completions.create(
+ model="claude-sonnet-4-5",
+ messages=[
+ {"role": "system", "content": "Você é um assistente útil."},
+ {"role": "user", "content": "Olá!"}
+ ],
+ stream=True
+)
+
+for chunk in response:
+ if chunk.choices[0].delta.content:
+ print(chunk.choices[0].delta.content, end="")
+```
+
+
+
+
+🦜 LangChain
+
+```python
+from langchain_openai import ChatOpenAI
+
+llm = ChatOpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123", # Seu PROXY_API_KEY do .env
+ model="claude-sonnet-4-5"
+)
+
+response = llm.invoke("Olá, como você está?")
+print(response.content)
+```
+
+
+
+### Anthropic API
+
+
+🔹 Requisição cURL Simples
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Olá!"}]
+ }'
+```
+
+> **Nota:** A API Anthropic usa o header `x-api-key` em vez de `Authorization: Bearer`. Ambos são suportados.
+
+
+
+
+🔹 Com Prompt de Sistema
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "system": "Você é um assistente útil.",
+ "messages": [{"role": "user", "content": "Olá!"}]
+ }'
+```
+
+> **Nota:** Na API Anthropic, `system` é um campo separado, não uma mensagem.
+
+
+
+
+📡 Streaming
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "stream": true,
+ "messages": [{"role": "user", "content": "Olá!"}]
+ }'
+```
+
+
+
+
+🐍 Python Anthropic SDK
+
+```python
+import anthropic
+
+client = anthropic.Anthropic(
+ api_key="my-super-secret-password-123", # Seu PROXY_API_KEY do .env
+ base_url="http://localhost:8000"
+)
+
+# Sem streaming
+response = client.messages.create(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "Olá!"}]
+)
+print(response.content[0].text)
+
+# Com streaming
+with client.messages.stream(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "Olá!"}]
+) as stream:
+ for text in stream.text_stream:
+ print(text, end="", flush=True)
+```
+
+
+
+---
+
+## 🔧 Depuração
+
+O log de depuração está **desabilitado por padrão**. Para habilitar, adicione ao seu `.env`:
+
+```env
+# Modo de log de depuração:
+# - off: desabilitado (padrão)
+# - errors: salvar logs apenas para requisições com falha (4xx, 5xx) - recomendado para solução de problemas
+# - all: salvar logs para cada requisição (sobrescreve a cada requisição)
+DEBUG_MODE=errors
+```
+
+### Modos de Depuração
+
+| Modo | Descrição | Caso de Uso |
+|------|-----------|-------------|
+| `off` | Desabilitado (padrão) | Produção |
+| `errors` | Salvar logs apenas para requisições com falha (4xx, 5xx) | **Recomendado para solução de problemas** |
+| `all` | Salvar logs para cada requisição | Desenvolvimento/depuração |
+
+### Arquivos de Depuração
+
+Quando habilitado, as requisições são registradas na pasta `debug_logs/`:
+
+| Arquivo | Descrição |
+|---------|-----------|
+| `request_body.json` | Requisição recebida do cliente (formato OpenAI) |
+| `kiro_request_body.json` | Requisição enviada para a API Kiro |
+| `response_stream_raw.txt` | Stream bruto do Kiro |
+| `response_stream_modified.txt` | Stream transformado (formato OpenAI) |
+| `app_logs.txt` | Logs da aplicação para a requisição |
+| `error_info.json` | Detalhes do erro (apenas em erros) |
+
+---
+
+## 📜 Licença
+
+Este projeto está licenciado sob a **GNU Affero General Public License v3.0 (AGPL-3.0)**.
+
+Isso significa:
+- ✅ Você pode usar, modificar e distribuir este software
+- ✅ Você pode usá-lo para fins comerciais
+- ⚠️ **Você deve divulgar o código-fonte** quando distribuir o software
+- ⚠️ **Uso em rede é distribuição** — se você executar uma versão modificada em um servidor e permitir que outros interajam com ela, você deve disponibilizar o código-fonte para eles
+- ⚠️ Modificações devem ser lançadas sob a mesma licença
+
+Veja o arquivo [LICENSE](../../LICENSE) para o texto completo da licença.
+
+### Por que AGPL-3.0?
+
+AGPL-3.0 garante que melhorias neste software beneficiem toda a comunidade. Se você modificar este gateway e implantá-lo como um serviço, você deve compartilhar suas melhorias com seus usuários.
+
+### Acordo de Licença de Contribuidor (CLA)
+
+Ao enviar uma contribuição para este projeto, você concorda com os termos do nosso [Acordo de Licença de Contribuidor (CLA)](../../CLA.md). Isso garante que:
+- Você tem o direito de enviar a contribuição
+- Você concede ao mantenedor direitos de usar e relicenciar sua contribuição
+- O projeto permanece legalmente protegido
+
+---
+
+## 💖 Apoie o Projeto
+
+
+
+
+
+**Se este projeto economizou seu tempo ou dinheiro, considere apoiá-lo!**
+
+Cada contribuição ajuda a manter este projeto vivo e crescendo
+
+
+
+### 🤑 Doar
+
+[**☕ Doação Única**](https://app.lava.top/jwadow?tabId=donate) • [**💎 Apoio Mensal**](https://app.lava.top/jwadow?tabId=subscriptions)
+
+
+
+### 🪙 Ou envie criptomoedas
+
+| Moeda | Rede | Endereço |
+|:-----:|:----:|:---------|
+| **USDT** | TRC20 | `TSVtgRc9pkC1UgcbVeijBHjFmpkYHDRu26` |
+| **BTC** | Bitcoin | `12GZqxqpcBsqJ4Vf1YreLqwoMGvzBPgJq6` |
+| **ETH** | Ethereum | `0xc86eab3bba3bbaf4eb5b5fff8586f1460f1fd395` |
+| **SOL** | Solana | `9amykF7KibZmdaw66a1oqYJyi75fRqgdsqnG66AK3jvh` |
+| **TON** | TON | `UQBVh8T1H3GI7gd7b-_PPNnxHYYxptrcCVf3qQk5v41h3QTM` |
+
+
+
+**[⬆ Voltar ao Topo](#-kiro-gateway)**
+
+
diff --git a/kiro-gateway/docs/ru/ARCHITECTURE.md b/kiro-gateway/docs/ru/ARCHITECTURE.md
new file mode 100644
index 0000000000000000000000000000000000000000..b67265d77b6c9d5daa5dc4b3a515f08fd4a60ec6
--- /dev/null
+++ b/kiro-gateway/docs/ru/ARCHITECTURE.md
@@ -0,0 +1,821 @@
+# Архитектурный Обзор: Kiro Gateway
+
+## 1. Назначение и Цели Системы
+
+Проект представляет собой высокоуровневый прокси-шлюз, реализующий структурный паттерн проектирования **"Адаптер" (Adapter)**.
+
+Основная цель системы — обеспечить прозрачную совместимость между несколькими гетерогенными интерфейсами:
+
+### Поддерживаемые API форматы
+
+| API | Эндпоинты | Статус |
+|-----|-----------|--------|
+| **OpenAI** | `/v1/models`, `/v1/chat/completions` | ✅ Поддерживается |
+| **Anthropic** | `/v1/messages` | ✅ Поддерживается |
+
+### Архитектурная модель
+
+```
+┌─────────────────────────────────────────────────────────────────┐
+│ Клиенты │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │ OpenAI SDK/Tools │ │ Anthropic SDK/Tools │ │
+│ │ (Cursor, Cline, │ │ (Claude Code, │ │
+│ │ Continue, etc.) │ │ Anthropic SDK) │ │
+│ └──────────┬──────────┘ └──────────┬──────────┘ │
+└─────────────┼──────────────────────────────┼───────────────────┘
+ │ │
+ ▼ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ Kiro Gateway │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │ OpenAI Adapter │ │ Anthropic Adapter │ │
+│ │ /v1/chat/... │ │ /v1/messages │ │
+│ └──────────┬──────────┘ └──────────┬──────────┘ │
+│ └──────────────┬───────────────┘ │
+│ ▼ │
+│ ┌─────────────────────────────┐ │
+│ │ Core Layer │ │
+│ │ (Общая логика конвертации) │ │
+│ └──────────────┬──────────────┘ │
+└────────────────────────────┼────────────────────────────────────┘
+ │
+ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ Kiro API │
+│ (AWS CodeWhisperer Backend) │
+└─────────────────────────────────────────────────────────────────┘
+```
+
+Система выступает в роли "переводчика", позволяя использовать любые инструменты, библиотеки и IDE-плагины, разработанные для экосистем OpenAI и Anthropic, с моделями Claude через Kiro API.
+
+**Оба API работают одновременно** на одном сервере без необходимости переключения в настройках.
+
+## 2. Структура Проекта
+
+Проект организован в виде модульного Python-пакета `kiro/`:
+
+```
+kiro-gateway/
+├── main.py # Точка входа, создание FastAPI приложения
+├── requirements.txt # Зависимости Python
+├── .env.example # Пример конфигурации окружения
+│
+├── kiro/ # Основной пакет
+│ ├── __init__.py # Экспорты пакета, версия
+│ │
+│ │ # ═══════════════════════════════════════════════════════
+│ │ # SHARED LAYER - Переиспользуется всеми API
+│ │ # ═══════════════════════════════════════════════════════
+│ ├── config.py # Конфигурация и константы
+│ ├── auth.py # KiroAuthManager - управление токенами
+│ ├── cache.py # ModelInfoCache - кэш моделей
+│ ├── http_client.py # HTTP клиент с retry логикой
+│ ├── parsers.py # Парсеры AWS SSE потоков
+│ ├── utils.py # Вспомогательные утилиты
+│ ├── tokenizer.py # Подсчёт токенов (tiktoken)
+│ ├── debug_logger.py # Отладочное логирование запросов
+│ ├── exceptions.py # Обработчики исключений
+│ ├── thinking_parser.py # Парсер thinking блоков
+│ │
+│ │ # ═══════════════════════════════════════════════════════
+│ │ # CORE LAYER - Общее ядро для всех API
+│ │ # ═══════════════════════════════════════════════════════
+│ ├── converters_core.py # Общая логика построения Kiro payload
+│ ├── streaming_core.py # Общая логика парсинга Kiro stream
+│ │
+│ │ # ═══════════════════════════════════════════════════════
+│ │ # OPENAI API LAYER
+│ │ # ═══════════════════════════════════════════════════════
+│ ├── models_openai.py # Pydantic модели OpenAI API
+│ ├── converters_openai.py # OpenAI → Kiro адаптер
+│ ├── routes_openai.py # FastAPI роуты OpenAI
+│ ├── streaming_openai.py # Kiro → OpenAI SSE форматтер
+│ │
+│ │ # ═══════════════════════════════════════════════════════
+│ │ # ANTHROPIC API LAYER
+│ │ # ═══════════════════════════════════════════════════════
+│ ├── models_anthropic.py # Pydantic модели Anthropic API
+│ ├── converters_anthropic.py # Anthropic → Kiro адаптер
+│ ├── routes_anthropic.py # FastAPI роуты Anthropic
+│ └── streaming_anthropic.py # Kiro → Anthropic SSE форматтер
+│
+├── tests/ # Тесты
+│ ├── conftest.py # Pytest fixtures
+│ ├── unit/ # Юнит-тесты
+│ └── integration/ # Интеграционные тесты
+│
+├── docs/ # Документация
+│ ├── ru/ # Русская версия
+│ └── en/ # Английская версия
+│
+└── debug_logs/ # Отладочные логи (генерируются при DEBUG_LAST_REQUEST=true)
+```
+
+### Принцип организации: Общее ядро + тонкие адаптеры
+
+Архитектура построена на принципе **максимального переиспользования кода**:
+
+| Слой | Назначение | Файлы |
+|------|------------|-------|
+| **Shared Layer** | Инфраструктура, не зависящая от формата API | `auth.py`, `http_client.py`, `cache.py`, `parsers.py`, `tokenizer.py` |
+| **Core Layer** | Общая бизнес-логика конвертации | `converters_core.py`, `streaming_core.py` |
+| **API Layer** | Тонкие адаптеры для конкретных форматов | `*_openai.py`, `*_anthropic.py` |
+
+## 3. Архитектурная Топология и Компоненты
+
+Система построена на базе асинхронного фреймворка `FastAPI` и использует событийную модель управления жизненным циклом (`Lifespan Events`).
+
+### 3.1. Точка входа (`main.py`)
+
+Файл `main.py` отвечает за:
+
+1. **Конфигурацию логирования** — настройка Loguru с цветным выводом
+2. **Валидацию конфигурации** — функция `validate_configuration()` проверяет:
+ - Наличие файла `.env`
+ - Наличие credentials (REFRESH_TOKEN или KIRO_CREDS_FILE)
+3. **Lifespan Manager** — создание и инициализация:
+ - `KiroAuthManager` для управления токенами
+ - `ModelInfoCache` для кэширования моделей
+4. **Регистрация обработчиков ошибок** — `validation_exception_handler` для ошибок 422
+5. **Подключение роутов** — `app.include_router(router)`
+
+### 3.2. Модуль конфигурации (`kiro/config.py`)
+
+Централизованное хранение всех настроек:
+
+| Параметр | Описание | Значение по умолчанию |
+|----------|----------|----------------------|
+| `PROXY_API_KEY` | API ключ для доступа к прокси | `changeme_proxy_secret` |
+| `REFRESH_TOKEN` | Refresh token Kiro | из `.env` |
+| `PROFILE_ARN` | ARN профиля AWS CodeWhisperer | из `.env` |
+| `REGION` | Регион AWS | `us-east-1` |
+| `KIRO_CREDS_FILE` | Путь к JSON файлу credentials | из `.env` |
+| `TOKEN_REFRESH_THRESHOLD` | Время до обновления токена | 600 сек (10 мин) |
+| `MAX_RETRIES` | Макс. количество повторов | 3 |
+| `BASE_RETRY_DELAY` | Базовая задержка retry | 1.0 сек |
+| `MODEL_CACHE_TTL` | TTL кэша моделей | 3600 сек (1 час) |
+| `DEFAULT_MAX_INPUT_TOKENS` | Макс. input токенов по умолчанию | 200000 |
+| `TOOL_DESCRIPTION_MAX_LENGTH` | Макс. длина описания tool | 10000 символов |
+| `DEBUG_LAST_REQUEST` | Включить отладочное логирование | `false` |
+| `DEBUG_DIR` | Директория для debug логов | `debug_logs` |
+| `APP_VERSION` | Версия приложения | `0.0.0` |
+
+**Вспомогательные функции:**
+- `get_kiro_refresh_url(region)` — URL для обновления токена
+- `get_kiro_api_host(region)` — хост основного API
+- `get_kiro_q_host(region)` — хост Q API
+- `get_internal_model_id(external_model)` — конвертация имени модели
+
+### 3.3. Pydantic Модели (`kiro/models_openai.py`)
+
+#### Модели для `/v1/models`
+
+| Модель | Описание |
+|--------|----------|
+| `OpenAIModel` | Описание AI модели (id, object, created, owned_by) |
+| `ModelList` | Список моделей для ответа endpoint |
+
+#### Модели для `/v1/chat/completions`
+
+| Модель | Описание |
+|--------|----------|
+| `ChatMessage` | Сообщение чата (role, content, tool_calls, tool_call_id) |
+| `ToolFunction` | Описание функции инструмента (name, description, parameters) |
+| `Tool` | Инструмент OpenAI формата (type, function) |
+| `ChatCompletionRequest` | Запрос на генерацию (model, messages, stream, tools, ...) |
+
+#### Модели ответов
+
+| Модель | Описание |
+|--------|----------|
+| `ChatCompletionChoice` | Один вариант ответа |
+| `ChatCompletionUsage` | Информация о токенах (prompt_tokens, completion_tokens, credits_used) |
+| `ChatCompletionResponse` | Полный ответ (non-streaming) |
+| `ChatCompletionChunk` | Streaming chunk |
+| `ChatCompletionChunkDelta` | Дельта изменений в chunk |
+| `ChatCompletionChunkChoice` | Вариант в streaming chunk |
+
+### 3.4. Управление Состоянием (State Management Layer)
+
+#### KiroAuthManager (`kiro/auth.py`)
+
+**Роль:** Stateful-синглтон, инкапсулирующий логику управления токенами Kiro.
+
+**Возможности:**
+- Загрузка credentials из `.env` или JSON файла
+- Поддержка `expiresAt` для проверки времени истечения токена
+- Автоматическое обновление токена за 10 минут до истечения
+- Сохранение обновлённых токенов обратно в JSON файл
+- Поддержка разных регионов AWS
+- Генерация уникального fingerprint для User-Agent
+
+**Concurrency Control:** Использует `asyncio.Lock` для защиты от состояния гонки.
+
+**Основные методы:**
+- `get_access_token()` — возвращает действительный токен, обновляя при необходимости
+- `force_refresh()` — принудительное обновление токена (при 403)
+- `is_token_expiring_soon()` — проверка времени истечения
+
+**Properties:**
+- `profile_arn` — ARN профиля
+- `region` — регион AWS
+- `api_host` — хост API для региона
+- `q_host` — хост Q API для региона
+- `fingerprint` — уникальный fingerprint машины
+
+```python
+# Пример использования
+auth_manager = KiroAuthManager(
+ refresh_token="your_token",
+ region="us-east-1",
+ creds_file="~/.aws/sso/cache/kiro-auth-token.json"
+)
+token = await auth_manager.get_access_token()
+```
+
+#### ModelInfoCache (`kiro/cache.py`)
+
+**Роль:** Потокобезопасное хранилище конфигураций моделей.
+
+**Стратегия Заполнения:**
+- Lazy Loading через `/ListAvailableModels`
+- TTL кэша: 1 час
+- Fallback на статический список моделей
+
+**Основные методы:**
+- `update(models_data)` — обновление кэша
+- `get(model_id)` — получение информации о модели
+- `get_max_input_tokens(model_id)` — получение лимита токенов
+- `is_empty()` / `is_stale()` — проверка состояния кэша
+- `get_all_model_ids()` — список всех ID моделей
+
+### 3.5. Вспомогательные Утилиты (`kiro/utils.py`)
+
+| Функция | Описание |
+|---------|----------|
+| `get_machine_fingerprint()` | SHA256 хеш `{hostname}-{username}-kiro-gateway` |
+| `get_kiro_headers(auth_manager, token)` | Формирование заголовков для Kiro API |
+| `generate_completion_id()` | ID в формате `chatcmpl-{uuid_hex}` |
+| `generate_conversation_id()` | UUID для разговора |
+| `generate_tool_call_id()` | ID в формате `call_{uuid_hex[:8]}` |
+
+### 3.6. Слой Конвертации (`kiro/converters_openai.py`)
+
+#### Конвертация сообщений
+
+OpenAI messages преобразуются в Kiro conversationState:
+
+1. **System prompt** — добавляется к первому user сообщению
+2. **История сообщений** — полностью передаётся в `history` array
+3. **Объединение соседних сообщений** — сообщения с одинаковой ролью мерджатся
+4. **Tool calls** — поддержка OpenAI tools формата
+5. **Tool results** — корректная передача результатов вызова инструментов
+
+#### Обработка длинных описаний Tools
+
+**Проблема:** Kiro API возвращает ошибку 400 при слишком длинных описаниях в `toolSpecification.description`.
+
+**Решение:** Tool Documentation Reference Pattern
+- Если `description ≤ TOOL_DESCRIPTION_MAX_LENGTH` → оставляем как есть
+- Если `description > TOOL_DESCRIPTION_MAX_LENGTH`:
+ * В `toolSpecification.description` → ссылка: `"[Full documentation in system prompt under '## Tool: {name}']"`
+ * В system prompt добавляется секция `"## Tool: {name}"` с полным описанием
+
+**Функция:** `process_tools_with_long_descriptions(tools)` → `(processed_tools, tool_documentation)`
+
+#### Основные функции
+
+| Функция | Описание |
+|---------|----------|
+| `extract_text_content(content)` | Извлечение текста из различных форматов |
+| `merge_adjacent_messages(messages)` | Объединение соседних сообщений с одной ролью |
+| `build_kiro_history(messages, model_id)` | Построение массива history для Kiro |
+| `build_kiro_payload(request_data, conversation_id, profile_arn)` | Полный payload для запроса |
+
+#### Маппинг моделей
+
+Внешние имена моделей преобразуются во внутренние ID Kiro:
+
+| Внешнее имя | Внутренний ID Kiro |
+|-------------|-------------------|
+| `claude-opus-4-5` | `claude-opus-4.5` |
+| `claude-opus-4-5-20251101` | `claude-opus-4.5` |
+| `claude-haiku-4-5` | `claude-haiku-4.5` |
+| `claude-haiku-4.5` | `claude-haiku-4.5` (прямой проброс) |
+| `claude-sonnet-4-5` | `CLAUDE_SONNET_4_5_20250929_V1_0` |
+| `claude-sonnet-4-5-20250929` | `CLAUDE_SONNET_4_5_20250929_V1_0` |
+| `claude-sonnet-4` | `CLAUDE_SONNET_4_20250514_V1_0` |
+| `claude-sonnet-4-20250514` | `CLAUDE_SONNET_4_20250514_V1_0` |
+| `claude-3-7-sonnet-20250219` | `CLAUDE_3_7_SONNET_20250219_V1_0` |
+| `auto` | `claude-sonnet-4.5` (алиас) |
+
+### 3.7. Слой Парсинга (`kiro/parsers.py`)
+
+#### AwsEventStreamParser
+
+Продвинутый парсер AWS SSE формата с поддержкой:
+
+- **Bracket counting** — корректный парсинг вложенных JSON объектов
+- **Дедупликация контента** — фильтрация повторяющихся событий
+- **Tool calls** — парсинг структурированных и bracket-style tool calls
+- **Escape-последовательности** — декодирование `\n` и других
+
+#### Типы событий
+
+| Событие | Описание |
+|---------|----------|
+| `content` | Текстовый контент ответа |
+| `tool_start` | Начало tool call (name, toolUseId) |
+| `tool_input` | Продолжение input для tool call |
+| `tool_stop` | Завершение tool call |
+| `usage` | Информация о потреблении кредитов |
+| `context_usage` | Процент использования контекста |
+
+#### Вспомогательные функции
+
+| Функция | Описание |
+|---------|----------|
+| `find_matching_brace(text, start_pos)` | Поиск закрывающей скобки с учётом вложенности |
+| `parse_bracket_tool_calls(response_text)` | Парсинг `[Called func with args: {...}]` |
+| `deduplicate_tool_calls(tool_calls)` | Удаление дубликатов tool calls |
+
+### 3.8. Streaming (`kiro/streaming_openai.py`)
+
+#### stream_kiro_to_openai
+
+Асинхронный генератор для преобразования потока Kiro в OpenAI формат.
+
+**Функциональность:**
+- Парсинг AWS SSE stream через `AwsEventStreamParser`
+- Формирование OpenAI `chat.completion.chunk`
+- Обработка tool calls (структурированных и bracket-style)
+- Вычисление usage на основе `contextUsagePercentage`
+- Отладочное логирование через `debug_logger`
+
+#### collect_stream_response
+
+Собирает полный ответ из streaming потока для non-streaming режима.
+
+### 3.9. HTTP Клиент (`kiro/http_client.py`)
+
+#### KiroHttpClient
+
+Автоматическая обработка ошибок с exponential backoff:
+
+| Код ошибки | Действие |
+|------------|----------|
+| `403` | Refresh токена через `force_refresh()` + повтор |
+| `429` | Exponential backoff: `BASE_RETRY_DELAY * (2 ** attempt)` |
+| `5xx` | Exponential backoff (до MAX_RETRIES попыток) |
+| Timeout | Exponential backoff |
+
+**Формула задержки:** `1s, 2s, 4s` (при `BASE_RETRY_DELAY=1.0`)
+
+**Методы:**
+- `request_with_retry(method, url, json_data, stream)` — запрос с retry
+- `close()` — закрытие клиента
+
+Поддерживает async context manager (`async with`).
+
+### 3.10. Роуты (`kiro/routes_openai.py`)
+
+| Endpoint | Метод | Описание |
+|----------|-------|----------|
+| `/` | GET | Health check (status, message, version) |
+| `/health` | GET | Детальный health check (status, timestamp, version) |
+| `/v1/models` | GET | Список доступных моделей (требует API key) |
+| `/v1/chat/completions` | POST | Chat completions (требует API key) |
+
+**Аутентификация:** Bearer token в заголовке `Authorization`
+
+### 3.11. Обработка Исключений (`kiro/exceptions.py`)
+
+| Функция | Описание |
+|---------|----------|
+| `sanitize_validation_errors(errors)` | Конвертация bytes в строки для JSON-сериализации |
+| `validation_exception_handler(request, exc)` | Обработчик ошибок валидации Pydantic (422) |
+
+### 3.12. Отладочное Логирование (`kiro/debug_logger.py`)
+
+**Класс:** `DebugLogger` (синглтон)
+
+**Активация:** `DEBUG_LAST_REQUEST=true` в `.env`
+
+**Методы:**
+| Метод | Описание |
+|-------|----------|
+| `prepare_new_request()` | Очистка директории для нового запроса |
+| `log_request_body(body)` | Сохранение входящего запроса |
+| `log_kiro_request_body(body)` | Сохранение запроса к Kiro API |
+| `log_raw_chunk(chunk)` | Дописывание сырого chunk от Kiro |
+| `log_modified_chunk(chunk)` | Дописывание преобразованного chunk |
+
+**Файлы в `debug_logs/`:**
+- `request_body.json` — входящий запрос (OpenAI формат)
+- `kiro_request_body.json` — запрос к Kiro API
+- `response_stream_raw.txt` — сырой поток от Kiro
+- `response_stream_modified.txt` — преобразованный поток (OpenAI формат)
+
+### 3.13. Токенизатор (`kiro/tokenizer.py`)
+
+**Проблема:** Kiro API не возвращает напрямую количество токенов. Вместо этого API предоставляет только `context_usage_percentage` — процент использования контекста модели.
+
+**Решение:** Модуль токенизатора на базе `tiktoken` (библиотека OpenAI на Rust) для быстрого подсчёта токенов.
+
+**Особенности:**
+- Использует кодировку `cl100k_base` (GPT-4), близкую к токенизации Claude
+- Коэффициент коррекции `CLAUDE_CORRECTION_FACTOR = 1.15` для повышения точности
+- Ленивая инициализация для ускорения импорта
+- Fallback на грубую оценку если tiktoken недоступен
+
+**Формула расчёта токенов в ответе:**
+```
+total_tokens = context_usage_percentage × max_input_tokens (от API Kiro)
+completion_tokens = tiktoken(ответ) (наш подсчёт)
+prompt_tokens = total_tokens - completion_tokens (вычитание)
+```
+
+**Основные функции:**
+
+| Функция | Описание |
+|---------|----------|
+| `count_tokens(text)` | Подсчёт токенов в тексте |
+| `count_message_tokens(messages)` | Подсчёт токенов в списке сообщений |
+| `count_tools_tokens(tools)` | Подсчёт токенов в определениях инструментов |
+| `estimate_request_tokens(messages, tools)` | Полная оценка токенов запроса |
+
+**Дебаг-лог:**
+```
+[Usage] claude-opus-4-5: prompt_tokens=142211 (subtraction), completion_tokens=769 (tiktoken), total_tokens=142980 (API Kiro)
+```
+
+**Точность:** ~97-99.7% по сравнению с данными от API.
+
+### 3.14. Kiro API Endpoints
+
+Все URL динамически формируются на основе региона:
+
+* **Token Refresh:** `POST https://prod.{region}.auth.desktop.kiro.dev/refreshToken`
+* **List Models:** `GET https://q.{region}.amazonaws.com/ListAvailableModels`
+* **Generate Response:** `POST https://codewhisperer.{region}.amazonaws.com/generateAssistantResponse`
+
+## 4. Детальный Поток Данных
+
+### 4.1 Общая схема (мульти-API)
+
+```
+┌─────────────────────────────────────────────────────────────────┐
+│ КЛИЕНТЫ │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │ OpenAI Client │ │ Anthropic Client │ │
+│ └──────────┬──────────┘ └──────────┬──────────┘ │
+└─────────────┼──────────────────────────────┼───────────────────┘
+ │ │
+ │ POST /v1/chat/completions │ POST /v1/messages
+ ▼ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ API LAYER │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │ routes_openai.py │ │ routes_anthropic.py │ │
+│ │ Security Gate │ │ Security Gate │ │
+│ └──────────┬──────────┘ └──────────┬──────────┘ │
+│ │ │ │
+│ ▼ ▼ │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │converters_openai.py │ │converters_anthropic │ │
+│ │ Извлечение system │ │ System уже отдельно │ │
+│ │ из messages │ │ в запросе │ │
+│ └──────────┬──────────┘ └──────────┬──────────┘ │
+└─────────────┼──────────────────────────────┼───────────────────┘
+ │ │
+ └──────────────┬───────────────┘
+ │
+ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ CORE LAYER │
+│ ┌─────────────────────────────┐ │
+│ │ converters_core.py │ │
+│ │ build_kiro_payload() │ │
+│ │ build_kiro_history() │ │
+│ │ process_tools() │ │
+│ └──────────────┬──────────────┘ │
+└────────────────────────────┼────────────────────────────────────┘
+ │
+ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ SHARED LAYER │
+│ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │
+│ │ KiroAuthManager │ │ KiroHttpClient │ │ ModelInfoCache │ │
+│ │ (auth.py) │ │(http_client.py) │ │ (cache.py) │ │
+│ └────────┬────────┘ └────────┬────────┘ └─────────────────┘ │
+└───────────┼────────────────────┼────────────────────────────────┘
+ │ │
+ │ │ POST /generateAssistantResponse
+ │ ▼
+ │ ┌─────────────────────────────────────────┐
+ │ │ Kiro API │
+ │ └──────────────────┬──────────────────────┘
+ │ │
+ │ │ AWS SSE Stream
+ │ ▼
+┌───────────┼────────────────────────────────────────────────────┐
+│ │ CORE LAYER │
+│ │ ┌─────────────────────────────┐ │
+│ │ │ streaming_core.py │ │
+│ │ │ parse_kiro_stream() │ │
+│ │ │ → KiroEvent objects │ │
+│ │ └──────────────┬──────────────┘ │
+└────────────────────────────┼───────────────────────────────────┘
+ │
+ ┌──────────────┴───────────────┐
+ │ │
+ ▼ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ OUTPUT LAYER │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │streaming_openai.py │ │streaming_anthropic │ │
+│ │ format_openai_sse() │ │format_anthropic_sse │ │
+│ │ │ │ │ │
+│ │ data: {...} │ │ event: type │ │
+│ │ data: [DONE] │ │ data: {...} │ │
+│ └──────────┬──────────┘ └──────────┬──────────┘ │
+└─────────────┼──────────────────────────────┼───────────────────┘
+ │ │
+ ▼ ▼
+┌─────────────────────────────────────────────────────────────────┐
+│ КЛИЕНТЫ │
+│ ┌─────────────────────┐ ┌─────────────────────┐ │
+│ │ OpenAI Client │ │ Anthropic Client │ │
+│ └─────────────────────┘ └─────────────────────┘ │
+└─────────────────────────────────┘
+```
+
+### 4.2 Поток OpenAI API
+
+```
+OpenAI Client
+ │ POST /v1/chat/completions
+ ▼
+routes_openai.py ──► converters_openai.py ──► converters_core.py
+ │ │
+ │ ▼
+ │ Kiro Payload
+ │ │
+ ▼ ▼
+KiroAuthManager ──────────────────────────► KiroHttpClient
+ │
+ ▼
+ Kiro API
+ │
+ ▼
+streaming_core.py ◄─────────────────────── AWS SSE Stream
+ │
+ ▼
+streaming_openai.py
+ │
+ ▼
+OpenAI SSE Format ──────────────────────► OpenAI Client
+```
+
+### 4.3 Поток Anthropic API
+
+```
+Anthropic Client
+ │ POST /v1/messages
+ ▼
+routes_anthropic.py ──► converters_anthropic.py ──► converters_core.py
+ │ │
+ │ ▼
+ │ Kiro Payload
+ │ │
+ ▼ ▼
+KiroAuthManager ──────────────────────────────────► KiroHttpClient
+ │
+ ▼
+ Kiro API
+ │
+ ▼
+streaming_core.py ◄─────────────────────────────── AWS SSE Stream
+ │
+ ▼
+streaming_anthropic.py
+ │
+ ▼
+Anthropic SSE Format ──────────────────────────► Anthropic Client
+```
+
+## 5. Доступные Модели
+
+| Модель | Описание | Credits |
+|--------|----------|---------|
+| `claude-opus-4-5` | Топовая модель | ~2.2 |
+| `claude-opus-4-5-20251101` | Топовая модель (версия) | ~2.2 |
+| `claude-sonnet-4-5` | Улучшенная модель | ~1.3 |
+| `claude-sonnet-4-5-20250929` | Улучшенная модель (версия) | ~1.3 |
+| `claude-sonnet-4` | Сбалансированная модель | ~1.3 |
+| `claude-sonnet-4-20250514` | Сбалансированная (версия) | ~1.3 |
+| `claude-haiku-4-5` | Быстрая модель | ~0.4 |
+| `claude-3-7-sonnet-20250219` | Legacy модель | ~1.0 |
+
+## 6. Конфигурация
+
+### Переменные окружения (.env)
+
+```env
+# Обязательные
+REFRESH_TOKEN="your_kiro_refresh_token"
+PROXY_API_KEY="your_proxy_secret"
+
+# Опциональные
+PROFILE_ARN="arn:aws:codewhisperer:..."
+KIRO_REGION="us-east-1"
+KIRO_CREDS_FILE="~/.aws/sso/cache/kiro-auth-token.json"
+
+# Отладка
+DEBUG_LAST_REQUEST="false"
+DEBUG_DIR="debug_logs"
+
+# Лимиты
+TOOL_DESCRIPTION_MAX_LENGTH="10000"
+```
+
+### JSON файл credentials (опционально)
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:...",
+ "region": "us-east-1"
+}
+```
+
+## 7. API Endpoints
+
+### 7.1 Общие эндпоинты
+
+| Endpoint | Метод | Описание |
+|----------|-------|----------|
+| `/` | GET | Health check |
+| `/health` | GET | Детальный health check |
+
+### 7.2 OpenAI-совместимые эндпоинты
+
+| Endpoint | Метод | Описание |
+|----------|-------|----------|
+| `/v1/models` | GET | Список доступных моделей |
+| `/v1/chat/completions` | POST | Chat completions (streaming/non-streaming) |
+
+**Аутентификация:** `Authorization: Bearer {PROXY_API_KEY}`
+
+### 7.3 Anthropic-совместимые эндпоинты
+
+| Endpoint | Метод | Описание |
+|----------|-------|----------|
+| `/v1/messages` | POST | Messages API (streaming/non-streaming) |
+
+**Аутентификация:** `x-api-key: {PROXY_API_KEY}` + `anthropic-version: 2023-06-01`
+
+### 7.4 Сравнение форматов
+
+| Аспект | OpenAI | Anthropic |
+|--------|--------|-----------|
+| System prompt | В `messages` с `role: "system"` | Отдельное поле `system` |
+| Content | Строка или массив | Всегда массив content blocks |
+| Stop reason | `finish_reason: "stop"` | `stop_reason: "end_turn"` |
+| Usage | `prompt_tokens`, `completion_tokens` | `input_tokens`, `output_tokens` |
+| Streaming | `data: {...}\n\n` + `data: [DONE]` | `event: type\ndata: {...}\n\n` |
+| Tool format | `{type: "function", function: {...}}` | `{name: "...", input_schema: {...}}` |
+
+## 8. Особенности Реализации
+
+### Tool Calling
+
+Поддержка OpenAI-совместимого формата tools:
+
+```json
+{
+ "tools": [{
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get weather for a location",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string"}
+ }
+ }
+ }
+ }]
+}
+```
+
+### Streaming
+
+Полная поддержка SSE streaming с корректным форматом OpenAI:
+
+```
+data: {"id":"chatcmpl-...","object":"chat.completion.chunk",...}
+
+data: [DONE]
+```
+
+### Отладка
+
+При `DEBUG_LAST_REQUEST=true` все запросы и ответы логируются в `debug_logs/`:
+- `request_body.json` — входящий запрос
+- `kiro_request_body.json` — запрос к Kiro API
+- `response_stream_raw.txt` — сырой поток от Kiro
+- `response_stream_modified.txt` — преобразованный поток
+
+## 9. Расширяемость
+
+### Добавление нового API формата
+
+Модульная архитектура позволяет легко добавить поддержку других API форматов. Благодаря Core Layer, большая часть логики уже реализована.
+
+#### Шаги для добавления нового формата (например, Gemini)
+
+1. **Создать модели** — `models_gemini.py`
+ ```python
+ class GeminiRequest(BaseModel):
+ """Pydantic модель запроса Gemini."""
+ contents: List[GeminiContent]
+ ...
+ ```
+
+2. **Создать адаптер конвертации** — `converters_gemini.py`
+ ```python
+ from kiro.converters_core import build_kiro_payload
+
+ def gemini_to_kiro(request: GeminiRequest, ...) -> dict:
+ """Конвертирует Gemini запрос в Kiro payload."""
+ # Извлекаем данные из Gemini формата
+ system_prompt = extract_system_instruction(request)
+ messages = convert_gemini_contents(request.contents)
+ tools = convert_gemini_tools(request.tools)
+
+ # Используем общее ядро
+ return build_kiro_payload(
+ messages=messages,
+ system_prompt=system_prompt,
+ tools=tools,
+ ...
+ )
+ ```
+
+3. **Создать форматтер streaming** — `streaming_gemini.py`
+ ```python
+ from kiro.streaming_core import parse_kiro_stream
+
+ async def stream_to_gemini(response, ...) -> AsyncGenerator[str, None]:
+ """Форматирует Kiro события в Gemini SSE."""
+ async for event in parse_kiro_stream(response):
+ yield format_gemini_chunk(event)
+ ```
+
+4. **Создать роуты** — `routes_gemini.py`
+ ```python
+ router = APIRouter()
+
+ @router.post("/v1beta/models/{model}:generateContent")
+ async def generate_content(request: GeminiRequest):
+ ...
+ ```
+
+5. **Подключить в main.py**
+ ```python
+ from kiro.routes_gemini import router as gemini_router
+ app.include_router(gemini_router)
+ ```
+
+### Что переиспользуется автоматически
+
+При добавлении нового формата следующие компоненты работают "из коробки":
+
+| Компонент | Функциональность |
+|-----------|------------------|
+| `auth.py` | Управление токенами Kiro |
+| `http_client.py` | HTTP с retry логикой |
+| `cache.py` | Кэш моделей |
+| `parsers.py` | Парсинг AWS SSE |
+| `tokenizer.py` | Подсчёт токенов |
+| `converters_core.py` | Построение Kiro payload |
+| `streaming_core.py` | Парсинг Kiro stream |
+
+## 10. Зависимости
+
+Основные зависимости проекта (из `requirements.txt`):
+
+| Пакет | Назначение |
+|-------|------------|
+| `fastapi` | Асинхронный веб-фреймворк |
+| `uvicorn` | ASGI сервер |
+| `httpx` | Асинхронный HTTP клиент |
+| `pydantic` | Валидация данных и модели |
+| `python-dotenv` | Загрузка переменных окружения |
+| `loguru` | Продвинутое логирование |
+| `tiktoken` | Быстрый подсчёт токенов |
diff --git a/kiro-gateway/docs/ru/README.md b/kiro-gateway/docs/ru/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..e8c774ea4a107d297a9afd3ecbb263a1f59aed97
--- /dev/null
+++ b/kiro-gateway/docs/ru/README.md
@@ -0,0 +1,626 @@
+
+
+# 👻 Kiro Gateway
+
+**Прокси-шлюз для Kiro API (Amazon Q Developer / AWS CodeWhisperer)**
+
+[🇬🇧 English](../../README.md) • [🇨🇳 中文](../zh/README.md) • [🇪🇸 Español](../es/README.md) • [🇮🇩 Indonesia](../id/README.md) • [🇧🇷 Português](../pt/README.md) • [🇯🇵 日本語](../ja/README.md) • [🇰🇷 한국어](../ko/README.md)
+
+Сделано с ❤️ от [@Jwadow](https://github.com/jwadow)
+
+[](https://www.gnu.org/licenses/agpl-3.0)
+[](https://www.python.org/downloads/)
+[](https://fastapi.tiangolo.com/)
+[](#-поддержать-проект)
+
+*Используйте модели Claude из Kiro с Claude Code, OpenCode, Cursor, Cline, Roo Code, Kilo Code, Obsidian, OpenAI SDK, LangChain, Continue и другими инструментами, совместимыми с OpenAI или Anthropic*
+
+[Модели](#-поддерживаемые-модели) • [Возможности](#-возможности) • [Быстрый старт](#-быстрый-старт) • [Конфигурация](#%EF%B8%8F-конфигурация) • [💖 Поддержать](#-поддержать-проект)
+
+
+
+---
+
+## 🤖 Доступные модели
+
+> ⚠️ **Важно:** Доступность моделей зависит от вашего тарифа Kiro (бесплатный/платный). Шлюз предоставляет доступ к тем моделям, которые доступны в вашей IDE или CLI в зависимости от вашей подписки. Список ниже показывает модели, обычно доступные на **бесплатном тарифе**.
+
+> 🔒 **Claude Opus 4.5** был удалён из бесплатного тарифа 17 января 2026 года. Он может быть доступен на платных тарифах — проверьте список моделей в вашей IDE/CLI.
+
+🚀 **Claude Sonnet 4.5** — Сбалансированная производительность. Отлично подходит для программирования, написания текстов и задач общего назначения.
+
+⚡ **Claude Haiku 4.5** — Молниеносная скорость. Идеальна для быстрых ответов, простых задач и чата.
+
+📦 **Claude Sonnet 4** — Предыдущее поколение. По-прежнему мощная и надёжная для большинства задач.
+
+📦 **Claude 3.7 Sonnet** — Устаревшая модель. Доступна для обратной совместимости.
+
+> 💡 **Умное разрешение моделей:** Используйте любой формат названия модели — `claude-sonnet-4-5`, `claude-sonnet-4.5` или даже версионные названия вроде `claude-sonnet-4-5-20250929`. Шлюз автоматически нормализует их.
+
+---
+
+## ✨ Возможности
+
+| Возможность | Описание |
+|-------------|----------|
+| 🔌 **API, совместимый с OpenAI** | Работает с любым инструментом, совместимым с OpenAI |
+| 🔌 **API, совместимый с Anthropic** | Нативный эндпоинт `/v1/messages` |
+| 🌐 **Поддержка VPN/Proxy** | HTTP/SOCKS5 прокси для ограниченных сетей |
+| 🧠 **Расширенное мышление** | Режим рассуждений — эксклюзив нашего проекта |
+| 👁️ **Поддержка изображений** | Отправляйте изображения модели |
+| 🛠️ **Вызов инструментов** | Поддержка вызова функций |
+| 💬 **Полная история сообщений** | Передаёт полный контекст разговора |
+| 📡 **Стриминг** | Полная поддержка SSE-стриминга |
+| 🔄 **Логика повторных попыток** | Автоматические повторы при ошибках (403, 429, 5xx) |
+| 📋 **Расширенный список моделей** | Включая версионные модели |
+| 🔐 **Умное управление токенами** | Автоматическое обновление до истечения срока |
+
+---
+
+## 🚀 Быстрый старт
+
+### Предварительные требования
+
+- Python 3.10+
+- Одно из следующего:
+ - [Kiro IDE](https://kiro.dev/) с авторизованным аккаунтом, ИЛИ
+ - [Kiro CLI](https://kiro.dev/cli/) с AWS SSO (AWS IAM Identity Center, OIDC) - бесплатный Builder ID или корпоративный аккаунт
+
+### Установка
+
+```bash
+# Клонируйте репозиторий (требуется Git)
+git clone https://github.com/Jwadow/kiro-gateway.git
+cd kiro-gateway
+
+# Или скачайте ZIP: Code → Download ZIP → распакуйте → откройте папку kiro-gateway
+
+# Установите зависимости
+pip install -r requirements.txt
+
+# Настройте (см. раздел Конфигурация)
+cp .env.example .env
+# Скопируйте и отредактируйте .env с вашими учётными данными
+
+# Запустите сервер
+python main.py
+
+# Или с другим портом (если 8000 занят)
+python main.py --port 9000
+```
+
+Сервер будет доступен по адресу `http://localhost:8000`
+
+---
+
+## ⚙️ Конфигурация
+
+### Вариант 1: JSON-файл с учётными данными (Kiro IDE / Enterprise)
+
+Укажите путь к файлу с учётными данными:
+
+Работает с:
+- **Kiro IDE** (стандартный) - для личных аккаунтов
+- **Enterprise** - для корпоративных аккаунтов с SSO
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/kiro-auth-token.json"
+
+# Пароль для защиты ВАШЕГО прокси-сервера (придумайте любую надёжную строку)
+# Вы будете использовать его как api_key при подключении к вашему шлюзу
+PROXY_API_KEY="my-super-secret-password-123"
+```
+
+
+📄 Формат JSON-файла
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:...",
+ "region": "us-east-1",
+ "clientIdHash": "abc123..." // Optional: for corporate SSO setups
+}
+```
+
+> **Примечание:** Если у вас есть два JSON файла в `~/.aws/sso/cache/` (например, `kiro-auth-token.json` и файл с хешированным названием), используйте `kiro-auth-token.json` в `KIRO_CREDS_FILE`. Шлюз автоматически загрузит другой файл.
+
+
+
+### Вариант 2: Переменные окружения (файл .env)
+
+Создайте файл `.env` в корне проекта:
+
+```env
+# Обязательно
+REFRESH_TOKEN="ваш_kiro_refresh_token"
+
+# Пароль для защиты ВАШЕГО прокси-сервера (придумайте любую надёжную строку)
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Опционально
+PROFILE_ARN="arn:aws:codewhisperer:us-east-1:..."
+KIRO_REGION="us-east-1"
+```
+
+### Вариант 3: Учётные данные AWS SSO (kiro-cli / Enterprise)
+
+Если вы используете `kiro-cli` или Kiro IDE с AWS SSO (AWS IAM Identity Center), шлюз автоматически обнаружит и использует соответствующую аутентификацию.
+
+Работает как с бесплатными аккаунтами Builder ID, так и с корпоративными аккаунтами.
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/your-sso-cache-file.json"
+
+# Пароль для защиты ВАШЕГО прокси-сервера
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Примечание: PROFILE_ARN НЕ нужен для AWS SSO (Builder ID и корпоративные аккаунты)
+# Шлюз будет работать без него
+```
+
+
+📄 Формат JSON-файла AWS SSO
+
+Файлы учётных данных AWS SSO (из `~/.aws/sso/cache/`) содержат:
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "region": "us-east-1",
+ "clientId": "...",
+ "clientSecret": "..."
+}
+```
+
+**Примечание:** Пользователям AWS SSO (Builder ID и корпоративные аккаунты) НЕ нужен `profileArn`. Шлюз будет работать без него (если указан, он будет проигнорирован).
+
+
+
+
+🔍 Как это работает
+
+Шлюз автоматически определяет тип аутентификации на основе файла учётных данных:
+
+- **Kiro Desktop Auth** (по умолчанию): Используется, когда `clientId` и `clientSecret` НЕ присутствуют
+ - Эндпоинт: `https://prod.{region}.auth.desktop.kiro.dev/refreshToken`
+
+- **AWS SSO (OIDC)**: Используется, когда `clientId` и `clientSecret` присутствуют
+ - Эндпоинт: `https://oidc.{region}.amazonaws.com/token`
+
+Дополнительная настройка не требуется — просто укажите путь к вашему файлу учётных данных!
+
+
+
+### Вариант 4: SQLite-база данных kiro-cli
+
+Если вы используете `kiro-cli` и предпочитаете использовать его SQLite-базу данных напрямую:
+
+```env
+KIRO_CLI_DB_FILE="~/.local/share/kiro-cli/data.sqlite3"
+
+# Пароль для защиты ВАШЕГО прокси-сервер
+PROXY_API_KEY="my-super-secret-password-123"
+
+# Примечание: PROFILE_ARN НЕ нужен для AWS SSO (Builder ID и корпоративные аккаунты)
+# Шлюз будет работать без него
+```
+
+
+📄 Расположение баз данных
+
+| CLI-инструмент | Путь к базе данных |
+|----------------|-------------------|
+| kiro-cli | `~/.local/share/kiro-cli/data.sqlite3` |
+| amazon-q-developer-cli | `~/.local/share/amazon-q/data.sqlite3` |
+
+Шлюз читает учётные данные из таблицы `auth_kv`, которая хранит:
+- `kirocli:odic:token` или `codewhisperer:odic:token` — токен доступа, токен обновления, срок действия
+- `kirocli:odic:device-registration` или `codewhisperer:odic:device-registration` — ID клиента и секрет
+
+Оба формата ключей поддерживаются для совместимости с разными версиями kiro-cli.
+
+
+
+### Получение учётных данных
+
+**Для пользователей Kiro IDE:**
+- Войдите в Kiro IDE и используйте Вариант 1 выше (JSON-файл с учётными данными)
+- Файл учётных данных создаётся автоматически после входа
+
+**Для пользователей Kiro CLI:**
+- Войдите с помощью `kiro-cli login` и используйте Вариант 3 или Вариант 4 выше
+- Ручное извлечение токена не требуется!
+
+
+🔧 Продвинутое: Ручное извлечение токена
+
+Если вам нужно вручную извлечь refresh token (например, для отладки), вы можете перехватить трафик Kiro IDE:
+- Ищите запросы к: `prod.us-east-1.auth.desktop.kiro.dev/refreshToken`
+
+
+
+---
+
+## 🌐 Поддержка VPN/Proxy
+
+**Для пользователей в Китае, корпоративных сетях или регионах с проблемами подключения к сервисам AWS.**
+
+Шлюз поддерживает маршрутизацию всех запросов Kiro API через VPN или прокси-сервер. Это необходимо, если у вас возникают проблемы с подключением к конечным точкам AWS или вам нужно использовать корпоративный прокси.
+
+### Конфигурация
+
+Добавьте в ваш файл `.env`:
+
+```env
+# HTTP прокси
+VPN_PROXY_URL=http://127.0.0.1:7890
+
+# SOCKS5 прокси
+VPN_PROXY_URL=socks5://127.0.0.1:1080
+
+# С аутентификацией (корпоративные прокси)
+VPN_PROXY_URL=http://username:password@proxy.company.com:8080
+
+# Без протокола (по умолчанию http://)
+VPN_PROXY_URL=192.168.1.100:8080
+```
+
+### Поддерживаемые протоколы
+
+- ✅ **HTTP** — Стандартный протокол прокси
+- ✅ **HTTPS** — Безопасные соединения прокси
+- ✅ **SOCKS5** — Продвинутый протокол прокси (распространён в ПО VPN)
+- ✅ **Аутентификация** — Имя пользователя/пароль встроены в URL
+
+### Когда это нужно
+
+| Ситуация | Решение |
+|----------|---------|
+| Таймауты подключения к AWS | Используйте VPN/прокси для маршрутизации трафика |
+| Ограничения корпоративной сети | Настройте прокси вашей компании |
+| Проблемы с региональным подключением | Используйте VPN-сервис с поддержкой прокси |
+| Требования конфиденциальности | Маршрутизируйте через собственный прокси-сервер |
+
+### Популярное ПО VPN с поддержкой прокси
+
+Большинство VPN-клиентов предоставляют локальный прокси-сервер:
+- **Sing-box** — Современный VPN-клиент с поддержкой HTTP/SOCKS5 прокси
+- **Clash** — Обычно работает на `http://127.0.0.1:7890`
+- **V2Ray** — Настраиваемый SOCKS5/HTTP прокси
+- **Shadowsocks** — Поддержка SOCKS5 прокси
+- **Корпоративный VPN** — Уточните параметры прокси у вашего IT-отдела
+
+Оставьте `VPN_PROXY_URL` пустым (по умолчанию), если вам не нужна поддержка прокси.
+
+---
+
+## 📡 Справочник API
+
+### Эндпоинты
+
+| Эндпоинт | Метод | Описание |
+|----------|-------|----------|
+| `/` | GET | Проверка работоспособности |
+| `/health` | GET | Детальная проверка работоспособности |
+| `/v1/models` | GET | Список доступных моделей |
+| `/v1/chat/completions` | POST | OpenAI Chat Completions API |
+| `/v1/messages` | POST | Anthropic Messages API |
+
+---
+
+## 💡 Примеры использования
+
+### OpenAI API
+
+
+🔹 Простой cURL-запрос
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Привет!"}],
+ "stream": true
+ }'
+```
+
+> **Примечание:** Замените `my-super-secret-password-123` на `PROXY_API_KEY`, который вы указали в файле `.env`.
+
+
+
+
+🔹 Запрос со стримингом
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {"role": "system", "content": "Ты полезный ассистент."},
+ {"role": "user", "content": "Сколько будет 2+2?"}
+ ],
+ "stream": true
+ }'
+```
+
+
+
+
+🛠️ С вызовом инструментов
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Какая погода в Лондоне?"}],
+ "tools": [{
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Получить погоду для местоположения",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "Название города"}
+ },
+ "required": ["location"]
+ }
+ }
+ }]
+ }'
+```
+
+
+
+
+🐍 Python OpenAI SDK
+
+```python
+from openai import OpenAI
+
+client = OpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123" # Ваш PROXY_API_KEY из .env
+)
+
+response = client.chat.completions.create(
+ model="claude-sonnet-4-5",
+ messages=[
+ {"role": "system", "content": "Ты полезный ассистент."},
+ {"role": "user", "content": "Привет!"}
+ ],
+ stream=True
+)
+
+for chunk in response:
+ if chunk.choices[0].delta.content:
+ print(chunk.choices[0].delta.content, end="")
+```
+
+
+
+
+🦜 LangChain
+
+```python
+from langchain_openai import ChatOpenAI
+
+llm = ChatOpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123", # Ваш PROXY_API_KEY из .env
+ model="claude-sonnet-4-5"
+)
+
+response = llm.invoke("Привет, как дела?")
+print(response.content)
+```
+
+
+
+### Anthropic API
+
+
+🔹 Простой cURL-запрос
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Привет!"}]
+ }'
+```
+
+> **Примечание:** Anthropic API использует заголовок `x-api-key` вместо `Authorization: Bearer`. Оба варианта поддерживаются.
+
+
+
+
+🔹 С системным промптом
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "system": "Ты полезный ассистент.",
+ "messages": [{"role": "user", "content": "Привет!"}]
+ }'
+```
+
+> **Примечание:** В Anthropic API `system` — это отдельное поле, а не сообщение.
+
+
+
+
+📡 Стриминг
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "stream": true,
+ "messages": [{"role": "user", "content": "Привет!"}]
+ }'
+```
+
+
+
+
+🐍 Python Anthropic SDK
+
+```python
+import anthropic
+
+client = anthropic.Anthropic(
+ api_key="my-super-secret-password-123", # Ваш PROXY_API_KEY из .env
+ base_url="http://localhost:8000"
+)
+
+# Без стриминга
+response = client.messages.create(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "Привет!"}]
+)
+print(response.content[0].text)
+
+# Со стримингом
+with client.messages.stream(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "Привет!"}]
+) as stream:
+ for text in stream.text_stream:
+ print(text, end="", flush=True)
+```
+
+
+
+---
+
+## 🔧 Отладка
+
+Логирование отладки **отключено по умолчанию**. Чтобы включить, добавьте в ваш `.env`:
+
+```env
+# Режим логирования отладки:
+# - off: отключено (по умолчанию)
+# - errors: сохранять логи только для неудачных запросов (4xx, 5xx) - рекомендуется для устранения неполадок
+# - all: сохранять логи для каждого запроса (перезаписывается при каждом запросе)
+DEBUG_MODE=errors
+```
+
+### Режимы отладки
+
+| Режим | Описание | Случай использования |
+|-------|----------|---------------------|
+| `off` | Отключено (по умолчанию) | Продакшен |
+| `errors` | Сохранять логи только для неудачных запросов (4xx, 5xx) | **Рекомендуется для устранения неполадок** |
+| `all` | Сохранять логи для каждого запроса | Разработка/отладка |
+
+### Файлы отладки
+
+При включении запросы логируются в папку `debug_logs/`:
+
+| Файл | Описание |
+|------|----------|
+| `request_body.json` | Входящий запрос от клиента (формат OpenAI) |
+| `kiro_request_body.json` | Запрос, отправленный в Kiro API |
+| `response_stream_raw.txt` | Сырой поток от Kiro |
+| `response_stream_modified.txt` | Преобразованный поток (формат OpenAI) |
+| `app_logs.txt` | Логи приложения для запроса |
+| `error_info.json` | Детали ошибки (только при ошибках) |
+
+---
+
+## 📜 Лицензия
+
+Этот проект лицензирован под **GNU Affero General Public License v3.0 (AGPL-3.0)**.
+
+Это означает:
+- ✅ Вы можете использовать, модифицировать и распространять это программное обеспечение
+- ✅ Вы можете использовать его в коммерческих целях
+- ⚠️ **Вы должны раскрыть исходный код** при распространении программного обеспечения
+- ⚠️ **Сетевое использование является распространением** — если вы запускаете модифицированную версию на сервере и позволяете другим взаимодействовать с ней, вы должны сделать исходный код доступным для них
+- ⚠️ Модификации должны быть выпущены под той же лицензией
+
+Полный текст лицензии см. в файле [LICENSE](../../LICENSE).
+
+### Почему AGPL-3.0?
+
+AGPL-3.0 гарантирует, что улучшения этого программного обеспечения принесут пользу всему сообществу. Если вы модифицируете этот шлюз и развёртываете его как сервис, вы должны поделиться своими улучшениями с вашими пользователями.
+
+### Лицензионное соглашение участника (CLA)
+
+Отправляя вклад в этот проект, вы соглашаетесь с условиями нашего [Лицензионного соглашения участника (CLA)](../../CLA.md). Это гарантирует, что:
+- Вы имеете право отправить вклад
+- Вы предоставляете мейнтейнеру права на использование и перелицензирование вашего вклада
+- Проект остаётся юридически защищённым
+
+---
+
+## 💖 Поддержать проект
+
+
+
+
+
+**Если этот проект сэкономил вам время или деньги, рассмотрите возможность его поддержки!**
+
+Каждый вклад помогает поддерживать жизнь и развитие этого проекта
+
+
+
+### 🤑 Пожертвовать
+
+[**☕ Разовое пожертвование**](https://app.lava.top/jwadow?tabId=donate) • [**💎 Ежемесячная поддержка**](https://app.lava.top/jwadow?tabId=subscriptions)
+
+
+
+### 🪙 Или отправьте криптовалюту
+
+| Валюта | Сеть | Адрес |
+|:------:|:----:|:------|
+| **USDT** | TRC20 | `TSVtgRc9pkC1UgcbVeijBHjFmpkYHDRu26` |
+| **BTC** | Bitcoin | `12GZqxqpcBsqJ4Vf1YreLqwoMGvzBPgJq6` |
+| **ETH** | Ethereum | `0xc86eab3bba3bbaf4eb5b5fff8586f1460f1fd395` |
+| **SOL** | Solana | `9amykF7KibZmdaw66a1oqYJyi75fRqgdsqnG66AK3jvh` |
+| **TON** | TON | `UQBVh8T1H3GI7gd7b-_PPNnxHYYxptrcCVf3qQk5v41h3QTM` |
+
+
+
+**[⬆ Вернуться наверх](#-kiro-gateway)**
+
+
diff --git a/kiro-gateway/docs/zh/README.md b/kiro-gateway/docs/zh/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..8936dd394b213c89597f31a29a7603a27af49af7
--- /dev/null
+++ b/kiro-gateway/docs/zh/README.md
@@ -0,0 +1,626 @@
+
+
+# 👻 Kiro Gateway
+
+**Kiro API (Amazon Q Developer / AWS CodeWhisperer) 代理网关**
+
+[🇬🇧 English](../../README.md) • [🇷🇺 Русский](../ru/README.md) • [🇪🇸 Español](../es/README.md) • [🇮🇩 Indonesia](../id/README.md) • [🇧🇷 Português](../pt/README.md) • [🇯🇵 日本語](../ja/README.md) • [🇰🇷 한국어](../ko/README.md)
+
+由 [@Jwadow](https://github.com/jwadow) 用 ❤️ 制作
+
+[](https://www.gnu.org/licenses/agpl-3.0)
+[](https://www.python.org/downloads/)
+[](https://fastapi.tiangolo.com/)
+[](#-支持项目)
+
+*通过 Claude Code、OpenCode、Cursor、Cline、Roo Code、Kilo Code、Obsidian、OpenAI SDK、LangChain、Continue 和其他兼容 OpenAI 或 Anthropic 的工具使用 Kiro 的 Claude 模型*
+
+[模型](#-支持的模型) • [功能](#-功能特性) • [快速开始](#-快速开始) • [配置](#%EF%B8%8F-配置) • [💖 支持](#-支持项目)
+
+
+
+---
+
+## 🤖 可用模型
+
+> ⚠️ **重要：** 模型可用性取决于您的 Kiro 套餐（免费/付费）。网关提供对您的 IDE 或 CLI 中基于订阅可用的模型的访问。下面的列表显示**免费套餐**上通常可用的模型。
+
+> 🔒 **Claude Opus 4.5** 已于 2026 年 1 月 17 日从免费套餐中移除。它可能在付费套餐上可用 — 请检查您的 IDE/CLI 模型列表。
+
+🚀 **Claude Sonnet 4.5** — 性能均衡。非常适合编程、写作和通用任务。
+
+⚡ **Claude Haiku 4.5** — 闪电般快速。非常适合快速响应、简单任务和聊天。
+
+📦 **Claude Sonnet 4** — 上一代模型。对于大多数用例仍然强大可靠。
+
+📦 **Claude 3.7 Sonnet** — 旧版模型。为向后兼容而保留。
+
+> 💡 **智能模型解析：** 使用任何模型名称格式 — `claude-sonnet-4-5`、`claude-sonnet-4.5`，甚至版本化名称如 `claude-sonnet-4-5-20250929`。网关会自动标准化它们。
+
+---
+
+## ✨ 功能特性
+
+| 功能 | 描述 |
+|------|------|
+| 🔌 **兼容 OpenAI 的 API** | 与任何兼容 OpenAI 的工具配合使用 |
+| 🔌 **兼容 Anthropic 的 API** | 原生 `/v1/messages` 端点 |
+| 🌐 **VPN/代理支持** | 用于受限网络的 HTTP/SOCKS5 代理 |
+| 🧠 **扩展思维** | 推理功能是我们项目的独家特性 |
+| 👁️ **视觉支持** | 向模型发送图像 |
+| 🛠️ **工具调用** | 支持函数调用 |
+| 💬 **完整消息历史** | 传递完整的对话上下文 |
+| 📡 **流式传输** | 完整的 SSE 流式传输支持 |
+| 🔄 **重试逻辑** | 错误时自动重试（403、429、5xx） |
+| 📋 **扩展模型列表** | 包括版本化模型 |
+| 🔐 **智能令牌管理** | 到期前自动刷新 |
+
+---
+
+## 🚀 快速开始
+
+### 前置要求
+
+- Python 3.10+
+- 以下之一：
+ - 已登录账户的 [Kiro IDE](https://kiro.dev/)，或
+ - 带有 AWS SSO (AWS IAM Identity Center, OIDC) 的 [Kiro CLI](https://kiro.dev/cli/) - 免费 Builder ID 或企业账户
+
+### 安装
+
+```bash
+# 克隆仓库（需要 Git）
+git clone https://github.com/Jwadow/kiro-gateway.git
+cd kiro-gateway
+
+# 或下载 ZIP：Code → Download ZIP → 解压 → 打开 kiro-gateway 文件夹
+
+# 安装依赖
+pip install -r requirements.txt
+
+# 配置（参见配置部分）
+cp .env.example .env
+# 复制并编辑 .env 文件，填入您的凭据
+
+# 启动服务器
+python main.py
+
+# 或使用自定义端口（如果 8000 被占用）
+python main.py --port 9000
+```
+
+服务器将在 `http://localhost:8000` 上可用
+
+---
+
+## ⚙️ 配置
+
+### 选项 1：JSON 凭据文件 (Kiro IDE / Enterprise)
+
+指定凭据文件的路径：
+
+适用于：
+- **Kiro IDE**（标准）- 用于个人账户
+- **Enterprise** - 用于带有 SSO 的企业账户
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/kiro-auth-token.json"
+
+# 保护您的代理服务器的密码（设置任何安全字符串）
+# 连接到您的网关时，您将使用它作为 api_key
+PROXY_API_KEY="my-super-secret-password-123"
+```
+
+
+📄 JSON 文件格式
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:...",
+ "region": "us-east-1",
+ "clientIdHash": "abc123..." // Optional: for corporate SSO setups
+}
+```
+
+> **注意：** 如果您在 `~/.aws/sso/cache/` 中有两个 JSON 文件（例如 `kiro-auth-token.json` 和一个带有哈希名称的文件），请在 `KIRO_CREDS_FILE` 中使用 `kiro-auth-token.json`。网关将自动加载另一个文件。
+
+
+
+### 选项 2：环境变量（.env 文件）
+
+在项目根目录创建 `.env` 文件：
+
+```env
+# 必需
+REFRESH_TOKEN="您的_kiro_refresh_token"
+
+# 保护您的代理服务器的密码（设置任何安全字符串）
+PROXY_API_KEY="my-super-secret-password-123"
+
+# 可选
+PROFILE_ARN="arn:aws:codewhisperer:us-east-1:..."
+KIRO_REGION="us-east-1"
+```
+
+### 选项 3：AWS SSO 凭据 (kiro-cli / Enterprise)
+
+如果您使用带有 AWS SSO (AWS IAM Identity Center) 的 `kiro-cli` 或 Kiro IDE，网关将自动检测并使用相应的认证。
+
+适用于免费 Builder ID 账户和企业账户。
+
+```env
+KIRO_CREDS_FILE="~/.aws/sso/cache/your-sso-cache-file.json"
+
+# 保护您的代理服务器的密码
+PROXY_API_KEY="my-super-secret-password-123"
+
+# 注意：AWS SSO (Builder ID 和企业账户) 用户不需要 PROFILE_ARN
+# 网关无需它即可工作
+```
+
+
+📄 AWS SSO JSON 文件格式
+
+AWS SSO 凭据文件（来自 `~/.aws/sso/cache/`）包含：
+
+```json
+{
+ "accessToken": "eyJ...",
+ "refreshToken": "eyJ...",
+ "expiresAt": "2025-01-12T23:00:00.000Z",
+ "region": "us-east-1",
+ "clientId": "...",
+ "clientSecret": "..."
+}
+```
+
+**注意：** AWS SSO (Builder ID 和企业账户) 用户不需要 `profileArn`。网关无需它即可工作（如果指定，将被忽略）。
+
+
+
+
+🔍 工作原理
+
+网关根据凭据文件自动检测认证类型：
+
+- **Kiro Desktop Auth**（默认）：当 `clientId` 和 `clientSecret` 不存在时使用
+ - 端点：`https://prod.{region}.auth.desktop.kiro.dev/refreshToken`
+
+- **AWS SSO (OIDC)**：当 `clientId` 和 `clientSecret` 存在时使用
+ - 端点：`https://oidc.{region}.amazonaws.com/token`
+
+无需额外配置 — 只需指向您的凭据文件！
+
+
+
+### 选项 4：kiro-cli SQLite 数据库
+
+如果您使用 `kiro-cli` 并希望直接使用其 SQLite 数据库：
+
+```env
+KIRO_CLI_DB_FILE="~/.local/share/kiro-cli/data.sqlite3"
+
+# 保护您的代理服务器的密码
+PROXY_API_KEY="my-super-secret-password-123"
+
+# 注意：AWS SSO (Builder ID 和企业账户) 用户不需要 PROFILE_ARN
+# 网关无需它即可工作
+```
+
+
+📄 数据库位置
+
+| CLI 工具 | 数据库路径 |
+|----------|-----------|
+| kiro-cli | `~/.local/share/kiro-cli/data.sqlite3` |
+| amazon-q-developer-cli | `~/.local/share/amazon-q/data.sqlite3` |
+
+网关从 `auth_kv` 表读取凭据，该表存储：
+- `kirocli:odic:token` 或 `codewhisperer:odic:token` — 访问令牌、刷新令牌、过期时间
+- `kirocli:odic:device-registration` 或 `codewhisperer:odic:device-registration` — 客户端 ID 和密钥
+
+两种键格式都支持，以兼容不同版本的 kiro-cli。
+
+
+
+### 获取凭据
+
+**Kiro IDE 用户：**
+- 登录 Kiro IDE 并使用上面的选项 1（JSON 凭据文件）
+- 凭据文件在登录后自动创建
+
+**Kiro CLI 用户：**
+- 使用 `kiro-cli login` 登录并使用上面的选项 3 或选项 4
+- 无需手动提取令牌！
+
+
+🔧 高级：手动提取令牌
+
+如果您需要手动提取 refresh token（例如用于调试），您可以拦截 Kiro IDE 流量：
+- 查找发往以下地址的请求：`prod.us-east-1.auth.desktop.kiro.dev/refreshToken`
+
+
+
+---
+
+## 🌐 VPN/代理支持
+
+**适用于中国、企业网络或与 AWS 服务连接存在问题的地区的用户。**
+
+网关支持通过 VPN 或代理服务器路由所有 Kiro API 请求。如果您遇到与 AWS 端点的连接问题或需要使用企业代理，这是必需的。
+
+### 配置
+
+添加到您的 `.env` 文件：
+
+```env
+# HTTP 代理
+VPN_PROXY_URL=http://127.0.0.1:7890
+
+# SOCKS5 代理
+VPN_PROXY_URL=socks5://127.0.0.1:1080
+
+# 带身份验证（企业代理）
+VPN_PROXY_URL=http://username:password@proxy.company.com:8080
+
+# 无协议（默认为 http://）
+VPN_PROXY_URL=192.168.1.100:8080
+```
+
+### 支持的协议
+
+- ✅ **HTTP** — 标准代理协议
+- ✅ **HTTPS** — 安全代理连接
+- ✅ **SOCKS5** — 高级代理协议（VPN 软件中常见）
+- ✅ **身份验证** — URL 中嵌入的用户名/密码
+
+### 何时需要
+
+| 情况 | 解决方案 |
+|------|---------|
+| 与 AWS 连接超时 | 使用 VPN/代理路由流量 |
+| 企业网络限制 | 配置您公司的代理 |
+| 区域连接问题 | 使用支持代理的 VPN 服务 |
+| 隐私要求 | 通过您自己的代理服务器路由 |
+
+### 支持代理的流行 VPN 软件
+
+大多数 VPN 客户端提供本地代理服务器：
+- **Sing-box** — 支持 HTTP/SOCKS5 代理的现代 VPN 客户端
+- **Clash** — 通常在 `http://127.0.0.1:7890` 上运行
+- **V2Ray** — 可配置的 SOCKS5/HTTP 代理
+- **Shadowsocks** — SOCKS5 代理支持
+- **企业 VPN** — 向您的 IT 部门咨询代理设置
+
+如果您不需要代理支持，请将 `VPN_PROXY_URL` 留空（默认）。
+
+---
+
+## 📡 API 参考
+
+### 端点
+
+| 端点 | 方法 | 描述 |
+|------|------|------|
+| `/` | GET | 健康检查 |
+| `/health` | GET | 详细健康检查 |
+| `/v1/models` | GET | 列出可用模型 |
+| `/v1/chat/completions` | POST | OpenAI Chat Completions API |
+| `/v1/messages` | POST | Anthropic Messages API |
+
+---
+
+## 💡 使用示例
+
+### OpenAI API
+
+
+🔹 简单 cURL 请求
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "你好！"}],
+ "stream": true
+ }'
+```
+
+> **注意：** 将 `my-super-secret-password-123` 替换为您在 `.env` 文件中设置的 `PROXY_API_KEY`。
+
+
+
+
+🔹 流式请求
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {"role": "system", "content": "你是一个有帮助的助手。"},
+ {"role": "user", "content": "2+2 等于多少？"}
+ ],
+ "stream": true
+ }'
+```
+
+
+
+
+🛠️ 带工具调用
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Authorization: Bearer my-super-secret-password-123" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "伦敦的天气怎么样？"}],
+ "tools": [{
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "获取某个位置的天气",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "城市名称"}
+ },
+ "required": ["location"]
+ }
+ }
+ }]
+ }'
+```
+
+
+
+
+🐍 Python OpenAI SDK
+
+```python
+from openai import OpenAI
+
+client = OpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123" # 您在 .env 中的 PROXY_API_KEY
+)
+
+response = client.chat.completions.create(
+ model="claude-sonnet-4-5",
+ messages=[
+ {"role": "system", "content": "你是一个有帮助的助手。"},
+ {"role": "user", "content": "你好！"}
+ ],
+ stream=True
+)
+
+for chunk in response:
+ if chunk.choices[0].delta.content:
+ print(chunk.choices[0].delta.content, end="")
+```
+
+
+
+
+🦜 LangChain
+
+```python
+from langchain_openai import ChatOpenAI
+
+llm = ChatOpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="my-super-secret-password-123", # 您在 .env 中的 PROXY_API_KEY
+ model="claude-sonnet-4-5"
+)
+
+response = llm.invoke("你好，你好吗？")
+print(response.content)
+```
+
+
+
+### Anthropic API
+
+
+🔹 简单 cURL 请求
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "你好！"}]
+ }'
+```
+
+> **注意：** Anthropic API 使用 `x-api-key` 头而不是 `Authorization: Bearer`。两者都支持。
+
+
+
+
+🔹 带系统提示
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "system": "你是一个有帮助的助手。",
+ "messages": [{"role": "user", "content": "你好！"}]
+ }'
+```
+
+> **注意：** 在 Anthropic API 中，`system` 是一个单独的字段，而不是消息。
+
+
+
+
+📡 流式传输
+
+```bash
+curl http://localhost:8000/v1/messages \
+ -H "x-api-key: my-super-secret-password-123" \
+ -H "anthropic-version: 2023-06-01" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "stream": true,
+ "messages": [{"role": "user", "content": "你好！"}]
+ }'
+```
+
+
+
+
+🐍 Python Anthropic SDK
+
+```python
+import anthropic
+
+client = anthropic.Anthropic(
+ api_key="my-super-secret-password-123", # 您在 .env 中的 PROXY_API_KEY
+ base_url="http://localhost:8000"
+)
+
+# 非流式
+response = client.messages.create(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "你好！"}]
+)
+print(response.content[0].text)
+
+# 流式
+with client.messages.stream(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[{"role": "user", "content": "你好！"}]
+) as stream:
+ for text in stream.text_stream:
+ print(text, end="", flush=True)
+```
+
+
+
+---
+
+## 🔧 调试
+
+调试日志**默认禁用**。要启用，请在您的 `.env` 中添加：
+
+```env
+# 调试日志模式：
+# - off：禁用（默认）
+# - errors：仅保存失败请求的日志（4xx、5xx）- 推荐用于故障排除
+# - all：保存每个请求的日志（每次请求时覆盖）
+DEBUG_MODE=errors
+```
+
+### 调试模式
+
+| 模式 | 描述 | 使用场景 |
+|------|------|----------|
+| `off` | 禁用（默认） | 生产环境 |
+| `errors` | 仅保存失败请求的日志（4xx、5xx） | **推荐用于故障排除** |
+| `all` | 保存每个请求的日志 | 开发/调试 |
+
+### 调试文件
+
+启用后，请求将记录到 `debug_logs/` 文件夹：
+
+| 文件 | 描述 |
+|------|------|
+| `request_body.json` | 来自客户端的传入请求（OpenAI 格式） |
+| `kiro_request_body.json` | 发送到 Kiro API 的请求 |
+| `response_stream_raw.txt` | 来自 Kiro 的原始流 |
+| `response_stream_modified.txt` | 转换后的流（OpenAI 格式） |
+| `app_logs.txt` | 请求的应用程序日志 |
+| `error_info.json` | 错误详情（仅在出错时） |
+
+---
+
+## 📜 许可证
+
+本项目采用 **GNU Affero 通用公共许可证 v3.0 (AGPL-3.0)** 许可。
+
+这意味着：
+- ✅ 您可以使用、修改和分发此软件
+- ✅ 您可以将其用于商业目的
+- ⚠️ **您必须公开源代码** 当您分发软件时
+- ⚠️ **网络使用即为分发** — 如果您在服务器上运行修改版本并让他人与之交互，您必须向他们提供源代码
+- ⚠️ 修改必须在相同许可证下发布
+
+完整许可证文本请参见 [LICENSE](../../LICENSE) 文件。
+
+### 为什么选择 AGPL-3.0？
+
+AGPL-3.0 确保对此软件的改进惠及整个社区。如果您修改此网关并将其部署为服务，您必须与您的用户分享您的改进。
+
+### 贡献者许可协议 (CLA)
+
+通过向本项目提交贡献，您同意我们的[贡献者许可协议 (CLA)](../../CLA.md) 的条款。这确保：
+- 您有权提交贡献
+- 您授予维护者使用和重新许可您的贡献的权利
+- 项目保持法律保护
+
+---
+
+## 💖 支持项目
+
+
+
+
+
+**如果这个项目为您节省了时间或金钱，请考虑支持它！**
+
+每一份贡献都有助于保持这个项目的活力和发展
+
+
+
+### 🤑 捐赠
+
+[**☕ 一次性捐赠**](https://app.lava.top/jwadow?tabId=donate) • [**💎 每月支持**](https://app.lava.top/jwadow?tabId=subscriptions)
+
+
+
+### 🪙 或发送加密货币
+
+| 货币 | 网络 | 地址 |
+|:----:|:----:|:-----|
+| **USDT** | TRC20 | `TSVtgRc9pkC1UgcbVeijBHjFmpkYHDRu26` |
+| **BTC** | Bitcoin | `12GZqxqpcBsqJ4Vf1YreLqwoMGvzBPgJq6` |
+| **ETH** | Ethereum | `0xc86eab3bba3bbaf4eb5b5fff8586f1460f1fd395` |
+| **SOL** | Solana | `9amykF7KibZmdaw66a1oqYJyi75fRqgdsqnG66AK3jvh` |
+| **TON** | TON | `UQBVh8T1H3GI7gd7b-_PPNnxHYYxptrcCVf3qQk5v41h3QTM` |
+
+
+
+**[⬆ 返回顶部](#-kiro-gateway)**
+
+
diff --git a/kiro-gateway/kiro/__init__.py b/kiro-gateway/kiro/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..42c86b1f58eebe3263772cd59b2d426f8056d141
--- /dev/null
+++ b/kiro-gateway/kiro/__init__.py
@@ -0,0 +1,137 @@
+# -*- 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 .
+
+"""
+Kiro Gateway - Proxy for Kiro API.
+
+This package provides a modular architecture for proxying
+OpenAI API requests to Kiro (AWS CodeWhisperer).
+
+Modules:
+ - config: Configuration and constants
+ - models: Pydantic models for OpenAI API
+ - auth: Kiro authentication manager
+ - cache: Model metadata cache
+ - utils: Helper utilities
+ - converters: OpenAI <-> Kiro format conversion
+ - parsers: AWS SSE stream parsers
+ - streaming: Response streaming logic
+ - http_client: HTTP client with retry logic
+ - routes: FastAPI routes
+ - exceptions: Exception handlers
+"""
+
+# Version is imported from config.py — the single source of truth
+# This allows changing the version in only one place
+from kiro.config import APP_VERSION as __version__
+
+__author__ = "Jwadow"
+
+# Main components for convenient import
+from kiro.auth import KiroAuthManager
+from kiro.cache import ModelInfoCache
+from kiro.http_client import KiroHttpClient
+from kiro.routes_openai import router
+from kiro.model_resolver import ModelResolver, normalize_model_name, get_model_id_for_kiro
+
+# Configuration
+from kiro.config import (
+ PROXY_API_KEY,
+ REGION,
+ HIDDEN_MODELS,
+ APP_VERSION,
+)
+
+# Models
+from kiro.models_openai import (
+ ChatCompletionRequest,
+ ChatMessage,
+ OpenAIModel,
+ ModelList,
+)
+
+# Converters
+from kiro.converters_openai import build_kiro_payload
+from kiro.converters_core import (
+ extract_text_content,
+ merge_adjacent_messages,
+)
+
+# Parsers
+from kiro.parsers import (
+ AwsEventStreamParser,
+ parse_bracket_tool_calls,
+)
+
+# Streaming
+from kiro.streaming_openai import (
+ stream_kiro_to_openai,
+ collect_stream_response,
+)
+
+# Exceptions
+from kiro.exceptions import (
+ validation_exception_handler,
+ sanitize_validation_errors,
+)
+
+__all__ = [
+ # Version
+ "__version__",
+
+ # Main classes
+ "KiroAuthManager",
+ "ModelInfoCache",
+ "KiroHttpClient",
+ "ModelResolver",
+ "router",
+
+ # Configuration
+ "PROXY_API_KEY",
+ "REGION",
+ "HIDDEN_MODELS",
+ "APP_VERSION",
+
+ # Model resolution
+ "normalize_model_name",
+ "get_model_id_for_kiro",
+
+ # Models
+ "ChatCompletionRequest",
+ "ChatMessage",
+ "OpenAIModel",
+ "ModelList",
+
+ # Converters
+ "build_kiro_payload",
+ "extract_text_content",
+ "merge_adjacent_messages",
+
+ # Parsers
+ "AwsEventStreamParser",
+ "parse_bracket_tool_calls",
+
+ # Streaming
+ "stream_kiro_to_openai",
+ "collect_stream_response",
+
+ # Exceptions
+ "validation_exception_handler",
+ "sanitize_validation_errors",
+]
\ No newline at end of file
diff --git a/kiro-gateway/kiro/auth.py b/kiro-gateway/kiro/auth.py
new file mode 100644
index 0000000000000000000000000000000000000000..37fbc237447edeaa509f077ec8b04efb7aa4172a
--- /dev/null
+++ b/kiro-gateway/kiro/auth.py
@@ -0,0 +1,863 @@
+# -*- 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 .
+
+"""
+Authentication manager for Kiro API.
+
+Manages the lifecycle of access tokens:
+- Loading credentials from .env or JSON file
+- Automatic token refresh on expiration
+- Thread-safe refresh using asyncio.Lock
+- Support for both Kiro Desktop Auth and AWS SSO OIDC (kiro-cli)
+"""
+
+import asyncio
+import json
+import sqlite3
+from datetime import datetime, timezone, timedelta
+from enum import Enum
+from pathlib import Path
+from typing import Optional
+
+import httpx
+from loguru import logger
+
+from kiro.config import (
+ TOKEN_REFRESH_THRESHOLD,
+ get_kiro_refresh_url,
+ get_kiro_api_host,
+ get_kiro_q_host,
+ get_aws_sso_oidc_url,
+)
+from kiro.utils import get_machine_fingerprint
+
+
+# Supported SQLite token keys (searched in priority order)
+SQLITE_TOKEN_KEYS = [
+ "kirocli:social:token", # Social login (Google, GitHub, Microsoft, etc.)
+ "kirocli:odic:token", # AWS SSO OIDC (kiro-cli corporate)
+ "codewhisperer:odic:token", # Legacy AWS SSO OIDC
+]
+
+# Device registration keys (for AWS SSO OIDC only)
+SQLITE_REGISTRATION_KEYS = [
+ "kirocli:odic:device-registration",
+ "codewhisperer:odic:device-registration",
+]
+
+
+class AuthType(Enum):
+ """
+ Type of authentication mechanism.
+
+ KIRO_DESKTOP: Kiro IDE credentials (default)
+ - Uses https://prod.{region}.auth.desktop.kiro.dev/refreshToken
+ - JSON body: {"refreshToken": "..."}
+
+ AWS_SSO_OIDC: AWS SSO credentials from kiro-cli
+ - Uses https://oidc.{region}.amazonaws.com/token
+ - Form body: grant_type=refresh_token&client_id=...&client_secret=...&refresh_token=...
+ - Requires clientId and clientSecret from credentials file
+ """
+ KIRO_DESKTOP = "kiro_desktop"
+ AWS_SSO_OIDC = "aws_sso_oidc"
+
+
+class KiroAuthManager:
+ """
+ Manages the token lifecycle for accessing Kiro API.
+
+ Supports:
+ - Loading credentials from .env or JSON file
+ - Automatic token refresh on expiration
+ - Expiration time validation (expiresAt)
+ - Saving updated tokens to file
+ - Both Kiro Desktop Auth and AWS SSO OIDC (kiro-cli) authentication
+
+ Attributes:
+ profile_arn: AWS CodeWhisperer profile ARN
+ region: AWS region
+ api_host: API host for the current region
+ q_host: Q API host for the current region
+ fingerprint: Unique machine fingerprint
+ auth_type: Type of authentication (KIRO_DESKTOP or AWS_SSO_OIDC)
+
+ Example:
+ >>> # Kiro Desktop Auth (default)
+ >>> auth_manager = KiroAuthManager(
+ ... refresh_token="your_refresh_token",
+ ... region="us-east-1"
+ ... )
+ >>> token = await auth_manager.get_access_token()
+
+ >>> # AWS SSO OIDC (kiro-cli) - auto-detected from credentials file
+ >>> auth_manager = KiroAuthManager(
+ ... creds_file="~/.aws/sso/cache/your-cache.json"
+ ... )
+ >>> token = await auth_manager.get_access_token()
+ """
+
+ def __init__(
+ self,
+ refresh_token: Optional[str] = None,
+ profile_arn: Optional[str] = None,
+ region: str = "us-east-1",
+ creds_file: Optional[str] = None,
+ client_id: Optional[str] = None,
+ client_secret: Optional[str] = None,
+ sqlite_db: Optional[str] = None,
+ ):
+ """
+ Initializes the authentication manager.
+
+ Args:
+ refresh_token: Refresh token for obtaining access token
+ profile_arn: AWS CodeWhisperer profile ARN
+ region: AWS region (default: us-east-1)
+ creds_file: Path to JSON file with credentials (optional)
+ client_id: OAuth client ID (for AWS SSO OIDC, optional)
+ client_secret: OAuth client secret (for AWS SSO OIDC, optional)
+ sqlite_db: Path to kiro-cli SQLite database (optional)
+ Default location: ~/.local/share/kiro-cli/data.sqlite3
+ """
+ self._refresh_token = refresh_token
+ self._profile_arn = profile_arn
+ self._region = region
+ self._creds_file = creds_file
+ self._sqlite_db = sqlite_db
+
+ # AWS SSO OIDC specific fields
+ self._client_id: Optional[str] = client_id
+ self._client_secret: Optional[str] = client_secret
+ self._scopes: Optional[list] = None # OAuth scopes for AWS SSO OIDC
+ self._sso_region: Optional[str] = None # SSO region for OIDC token refresh (may differ from API region)
+
+ # Enterprise Kiro IDE specific fields
+ self._client_id_hash: Optional[str] = None # clientIdHash from Enterprise Kiro IDE
+
+ # Track which SQLite key we loaded credentials from (for saving back to correct location)
+ self._sqlite_token_key: Optional[str] = None
+
+ self._access_token: Optional[str] = None
+ self._expires_at: Optional[datetime] = None
+ self._lock = asyncio.Lock()
+
+ # Auth type will be determined after loading credentials
+ self._auth_type: AuthType = AuthType.KIRO_DESKTOP
+
+ # Dynamic URLs based on region
+ self._refresh_url = get_kiro_refresh_url(region)
+ self._api_host = get_kiro_api_host(region)
+ self._q_host = get_kiro_q_host(region)
+
+ # Fingerprint for User-Agent
+ self._fingerprint = get_machine_fingerprint()
+
+ # Load credentials from SQLite if specified (takes priority over JSON)
+ if sqlite_db:
+ self._load_credentials_from_sqlite(sqlite_db)
+ # Load credentials from JSON file if specified
+ elif creds_file:
+ self._load_credentials_from_file(creds_file)
+
+ # Determine auth type based on available credentials
+ self._detect_auth_type()
+
+ def _detect_auth_type(self) -> None:
+ """
+ Detects authentication type based on available credentials.
+
+ AWS SSO OIDC credentials contain clientId and clientSecret.
+ Kiro Desktop credentials do not contain these fields.
+ """
+ if self._client_id and self._client_secret:
+ self._auth_type = AuthType.AWS_SSO_OIDC
+ logger.info("Detected auth type: AWS SSO OIDC (kiro-cli)")
+ else:
+ self._auth_type = AuthType.KIRO_DESKTOP
+ logger.info("Detected auth type: Kiro Desktop")
+
+ def _load_credentials_from_sqlite(self, db_path: str) -> None:
+ """
+ Loads credentials from kiro-cli SQLite database.
+
+ The database contains an auth_kv table with key-value pairs.
+ Supports multiple authentication types:
+
+ Token keys (searched in priority order):
+ - 'kirocli:social:token': Social login (Google, GitHub, etc.)
+ - 'kirocli:odic:token': AWS SSO OIDC (kiro-cli corporate)
+ - 'codewhisperer:odic:token': Legacy AWS SSO OIDC
+
+ Device registration keys (for AWS SSO OIDC only):
+ - 'kirocli:odic:device-registration': Client ID and secret
+ - 'codewhisperer:odic:device-registration': Legacy format
+
+ The method remembers which key was used for loading, so credentials
+ can be saved back to the correct location after refresh.
+
+ Args:
+ db_path: Path to SQLite database file
+ """
+ try:
+ path = Path(db_path).expanduser()
+ if not path.exists():
+ logger.warning(f"SQLite database not found: {db_path}")
+ return
+
+ conn = sqlite3.connect(str(path))
+ cursor = conn.cursor()
+
+ # Try all possible token keys in priority order
+ token_row = None
+ for key in SQLITE_TOKEN_KEYS:
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", (key,))
+ token_row = cursor.fetchone()
+ if token_row:
+ self._sqlite_token_key = key # Remember which key we loaded from
+ logger.debug(f"Loaded credentials from SQLite key: {key}")
+ break
+
+ if token_row:
+ token_data = json.loads(token_row[0])
+ if token_data:
+ # Load token fields (using snake_case as in Rust struct)
+ if 'access_token' in token_data:
+ self._access_token = token_data['access_token']
+ if 'refresh_token' in token_data:
+ self._refresh_token = token_data['refresh_token']
+ if 'profile_arn' in token_data:
+ self._profile_arn = token_data['profile_arn']
+ if 'region' in token_data:
+ # Store SSO region for OIDC token refresh only
+ # IMPORTANT: CodeWhisperer API is only available in us-east-1,
+ # so we don't update _api_host and _q_host here.
+ # The SSO region (e.g., ap-southeast-1) is only used for OIDC token refresh.
+ self._sso_region = token_data['region']
+ logger.debug(f"SSO region from SQLite: {self._sso_region} (API stays at {self._region})")
+
+ # Load scopes if available
+ if 'scopes' in token_data:
+ self._scopes = token_data['scopes']
+
+ # Parse expires_at (RFC3339 format)
+ if 'expires_at' in token_data:
+ try:
+ expires_str = token_data['expires_at']
+ # Handle various ISO 8601 formats
+ if expires_str.endswith('Z'):
+ self._expires_at = datetime.fromisoformat(expires_str.replace('Z', '+00:00'))
+ else:
+ self._expires_at = datetime.fromisoformat(expires_str)
+ except Exception as e:
+ logger.warning(f"Failed to parse expires_at from SQLite: {e}")
+
+ # Load device registration (client_id, client_secret) - try all possible keys
+ registration_row = None
+ for key in SQLITE_REGISTRATION_KEYS:
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", (key,))
+ registration_row = cursor.fetchone()
+ if registration_row:
+ logger.debug(f"Loaded device registration from SQLite key: {key}")
+ break
+
+ if registration_row:
+ registration_data = json.loads(registration_row[0])
+ if registration_data:
+ if 'client_id' in registration_data:
+ self._client_id = registration_data['client_id']
+ if 'client_secret' in registration_data:
+ self._client_secret = registration_data['client_secret']
+ # SSO region from registration (fallback if not in token data)
+ if 'region' in registration_data and not self._sso_region:
+ self._sso_region = registration_data['region']
+ logger.debug(f"SSO region from device-registration: {self._sso_region}")
+
+ conn.close()
+ logger.info(f"Credentials loaded from SQLite database: {db_path}")
+
+ except sqlite3.Error as e:
+ logger.error(f"SQLite error loading credentials: {e}")
+ except json.JSONDecodeError as e:
+ logger.error(f"JSON decode error in SQLite data: {e}")
+ except Exception as e:
+ logger.error(f"Error loading credentials from SQLite: {e}")
+
+ def _load_credentials_from_file(self, file_path: str) -> None:
+ """
+ Loads credentials from a JSON file.
+
+ Supported JSON fields (Kiro Desktop):
+ - refreshToken: Refresh token
+ - accessToken: Access token (if already available)
+ - profileArn: Profile ARN
+ - region: AWS region
+ - expiresAt: Token expiration time (ISO 8601)
+
+ Additional fields for AWS SSO OIDC (kiro-cli):
+ - clientId: OAuth client ID
+ - clientSecret: OAuth client secret
+
+ For Enterprise Kiro IDE:
+ - clientIdHash: Hash of client ID (Enterprise Kiro IDE)
+ - When clientIdHash is present, automatically loads clientId and clientSecret
+ from ~/.aws/sso/cache/{clientIdHash}.json (device registration file)
+
+ Args:
+ file_path: Path to JSON file
+ """
+ try:
+ path = Path(file_path).expanduser()
+ if not path.exists():
+ logger.warning(f"Credentials file not found: {file_path}")
+ return
+
+ with open(path, 'r', encoding='utf-8') as f:
+ data = json.load(f)
+
+ # Load common data from file
+ if 'refreshToken' in data:
+ self._refresh_token = data['refreshToken']
+ if 'accessToken' in data:
+ self._access_token = data['accessToken']
+ if 'profileArn' in data:
+ self._profile_arn = data['profileArn']
+ if 'region' in data:
+ self._region = data['region']
+ # Update URLs for new region
+ self._refresh_url = get_kiro_refresh_url(self._region)
+ self._api_host = get_kiro_api_host(self._region)
+ self._q_host = get_kiro_q_host(self._region)
+
+ # Load clientIdHash and device registration for Enterprise Kiro IDE
+ if 'clientIdHash' in data:
+ self._client_id_hash = data['clientIdHash']
+ self._load_enterprise_device_registration(self._client_id_hash)
+
+ # Load AWS SSO OIDC specific fields (if directly in credentials file)
+ if 'clientId' in data:
+ self._client_id = data['clientId']
+ if 'clientSecret' in data:
+ self._client_secret = data['clientSecret']
+
+ # Parse expiresAt
+ if 'expiresAt' in data:
+ try:
+ expires_str = data['expiresAt']
+ # Support for different date formats
+ if expires_str.endswith('Z'):
+ self._expires_at = datetime.fromisoformat(expires_str.replace('Z', '+00:00'))
+ else:
+ self._expires_at = datetime.fromisoformat(expires_str)
+ except Exception as e:
+ logger.warning(f"Failed to parse expiresAt: {e}")
+
+ logger.info(f"Credentials loaded from {file_path}")
+
+ except Exception as e:
+ logger.error(f"Error loading credentials from file: {e}")
+
+ def _load_enterprise_device_registration(self, client_id_hash: str) -> None:
+ """
+ Loads clientId and clientSecret from Enterprise Kiro IDE device registration file.
+
+ Enterprise Kiro IDE uses AWS SSO OIDC authentication. Device registration is stored at:
+ ~/.aws/sso/cache/{clientIdHash}.json
+
+ Args:
+ client_id_hash: Client ID hash used to locate the device registration file
+ """
+ try:
+ device_reg_path = Path.home() / ".aws" / "sso" / "cache" / f"{client_id_hash}.json"
+
+ if not device_reg_path.exists():
+ logger.warning(f"Enterprise device registration file not found: {device_reg_path}")
+ return
+
+ with open(device_reg_path, 'r', encoding='utf-8') as f:
+ device_data = json.load(f)
+
+ if 'clientId' in device_data:
+ self._client_id = device_data['clientId']
+
+ if 'clientSecret' in device_data:
+ self._client_secret = device_data['clientSecret']
+
+ logger.info(f"Enterprise device registration loaded from {device_reg_path}")
+
+ except Exception as e:
+ logger.error(f"Error loading enterprise device registration: {e}")
+
+ def _save_credentials_to_file(self) -> None:
+ """
+ Saves updated credentials to a JSON file.
+
+ Updates the existing file while preserving other fields.
+ """
+ if not self._creds_file:
+ return
+
+ try:
+ path = Path(self._creds_file).expanduser()
+
+ # Read existing data
+ existing_data = {}
+ if path.exists():
+ with open(path, 'r', encoding='utf-8') as f:
+ existing_data = json.load(f)
+
+ # Update data
+ existing_data['accessToken'] = self._access_token
+ existing_data['refreshToken'] = self._refresh_token
+ if self._expires_at:
+ existing_data['expiresAt'] = self._expires_at.isoformat()
+ if self._profile_arn:
+ existing_data['profileArn'] = self._profile_arn
+
+ # Save
+ with open(path, 'w', encoding='utf-8') as f:
+ json.dump(existing_data, f, indent=2, ensure_ascii=False)
+
+ logger.debug(f"Credentials saved to {self._creds_file}")
+
+ except Exception as e:
+ logger.error(f"Error saving credentials: {e}")
+
+ def _save_credentials_to_sqlite(self) -> None:
+ """
+ Saves updated credentials back to SQLite database.
+
+ This ensures that tokens refreshed by the gateway are persisted
+ and available after gateway restart or for other processes reading
+ the same SQLite database.
+
+ Strategy:
+ 1. If we know which key we loaded from (_sqlite_token_key), save to that key
+ 2. If that fails or key is unknown, try all supported keys as fallback
+
+ This approach ensures credentials are saved to the correct location
+ regardless of authentication type (social login, AWS SSO OIDC, legacy).
+
+ Updates the auth_kv table with fresh access_token, refresh_token,
+ and expires_at values after successful token refresh.
+ """
+ if not self._sqlite_db:
+ return
+
+ try:
+ path = Path(self._sqlite_db).expanduser()
+ if not path.exists():
+ logger.warning(f"SQLite database not found for writing: {self._sqlite_db}")
+ return
+
+ # Use timeout to avoid blocking if database is locked
+ conn = sqlite3.connect(str(path), timeout=5.0)
+ cursor = conn.cursor()
+
+ # Prepare token data matching the structure from _load_credentials_from_sqlite
+ token_data = {
+ "access_token": self._access_token,
+ "refresh_token": self._refresh_token,
+ "expires_at": self._expires_at.isoformat() if self._expires_at else None,
+ "region": self._sso_region or self._region,
+ }
+ if self._scopes:
+ token_data["scopes"] = self._scopes
+
+ token_json = json.dumps(token_data)
+
+ # Save back to the same key we loaded from (if known)
+ if self._sqlite_token_key:
+ cursor.execute(
+ "UPDATE auth_kv SET value = ? WHERE key = ?",
+ (token_json, self._sqlite_token_key)
+ )
+ if cursor.rowcount > 0:
+ conn.commit()
+ conn.close()
+ logger.debug(f"Credentials saved to SQLite key: {self._sqlite_token_key}")
+ return
+ else:
+ logger.warning(f"Failed to update SQLite key: {self._sqlite_token_key}, trying fallback")
+
+ # Fallback: try all keys (for edge cases where source key is unknown)
+ for key in SQLITE_TOKEN_KEYS:
+ cursor.execute(
+ "UPDATE auth_kv SET value = ? WHERE key = ?",
+ (token_json, key)
+ )
+ if cursor.rowcount > 0:
+ conn.commit()
+ conn.close()
+ logger.debug(f"Credentials saved to SQLite key: {key} (fallback)")
+ 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}")
+
+ def is_token_expiring_soon(self) -> bool:
+ """
+ Checks if the token is expiring soon.
+
+ Returns:
+ True if the token expires within TOKEN_REFRESH_THRESHOLD seconds
+ or if expiration time information is not available
+ """
+ if not self._expires_at:
+ return True # If no expiration info available, assume refresh is needed
+
+ now = datetime.now(timezone.utc)
+ threshold = now.timestamp() + TOKEN_REFRESH_THRESHOLD
+
+ return self._expires_at.timestamp() <= threshold
+
+ def is_token_expired(self) -> bool:
+ """
+ Checks if the token is actually expired (not just expiring soon).
+
+ This is used for graceful degradation when refresh fails but
+ the access token might still be valid for a short time.
+
+ Returns:
+ True if the token has already expired or if expiration time
+ information is not available
+ """
+ if not self._expires_at:
+ return True # If no expiration info available, assume expired
+
+ now = datetime.now(timezone.utc)
+ return now >= self._expires_at
+
+ async def _refresh_token_request(self) -> None:
+ """
+ Performs a token refresh request.
+
+ Routes to appropriate refresh method based on auth type:
+ - KIRO_DESKTOP: Uses Kiro Desktop Auth endpoint
+ - AWS_SSO_OIDC: Uses AWS SSO OIDC endpoint
+
+ Raises:
+ ValueError: If refresh token is not set or response doesn't contain accessToken
+ httpx.HTTPError: On HTTP request error
+ """
+ if self._auth_type == AuthType.AWS_SSO_OIDC:
+ await self._refresh_token_aws_sso_oidc()
+ else:
+ await self._refresh_token_kiro_desktop()
+
+ async def _refresh_token_kiro_desktop(self) -> None:
+ """
+ Refreshes token using Kiro Desktop Auth endpoint.
+
+ Endpoint: https://prod.{region}.auth.desktop.kiro.dev/refreshToken
+ Method: POST
+ Content-Type: application/json
+ Body: {"refreshToken": "..."}
+
+ Raises:
+ ValueError: If refresh token is not set or response doesn't contain accessToken
+ httpx.HTTPError: On HTTP request error
+ """
+ if not self._refresh_token:
+ raise ValueError("Refresh token is not set")
+
+ logger.info("Refreshing Kiro token via Kiro Desktop Auth...")
+
+ payload = {'refreshToken': self._refresh_token}
+ headers = {
+ "Content-Type": "application/json",
+ "User-Agent": f"KiroIDE-0.7.45-{self._fingerprint}",
+ }
+
+ async with httpx.AsyncClient(timeout=30) as client:
+ response = await client.post(self._refresh_url, json=payload, headers=headers)
+ response.raise_for_status()
+ data = response.json()
+
+ new_access_token = data.get("accessToken")
+ new_refresh_token = data.get("refreshToken")
+ expires_in = data.get("expiresIn", 3600)
+ new_profile_arn = data.get("profileArn")
+
+ if not new_access_token:
+ raise ValueError(f"Response does not contain accessToken: {data}")
+
+ # Update data
+ self._access_token = new_access_token
+ if new_refresh_token:
+ self._refresh_token = new_refresh_token
+ if new_profile_arn:
+ self._profile_arn = new_profile_arn
+
+ # Calculate expiration time with buffer (minus 60 seconds)
+ self._expires_at = datetime.now(timezone.utc).replace(microsecond=0)
+ self._expires_at = datetime.fromtimestamp(
+ self._expires_at.timestamp() + expires_in - 60,
+ tz=timezone.utc
+ )
+
+ logger.info(f"Token refreshed via Kiro Desktop Auth, expires: {self._expires_at.isoformat()}")
+
+ # Save to file or SQLite depending on configuration
+ if self._sqlite_db:
+ self._save_credentials_to_sqlite()
+ else:
+ self._save_credentials_to_file()
+
+ async def _refresh_token_aws_sso_oidc(self) -> None:
+ """
+ Refreshes token using AWS SSO OIDC endpoint.
+
+ Used by kiro-cli which authenticates via AWS IAM Identity Center.
+
+ Strategy: Try with current in-memory token first. If it fails with 400
+ (invalid_request - token was invalidated by kiro-cli re-login), reload
+ credentials from SQLite and retry once.
+
+ This approach handles both scenarios:
+ 1. Container successfully refreshed token (uses in-memory token)
+ 2. kiro-cli re-login invalidated token (reloads from SQLite on failure)
+
+ Endpoint: https://oidc.{region}.amazonaws.com/token
+ Method: POST
+ Content-Type: application/x-www-form-urlencoded
+ Body: grant_type=refresh_token&client_id=...&client_secret=...&refresh_token=...
+
+ Raises:
+ ValueError: If required credentials are not set
+ httpx.HTTPError: On HTTP request error
+ """
+ try:
+ await self._do_aws_sso_oidc_refresh()
+ except httpx.HTTPStatusError as e:
+ # 400 = invalid_request, likely stale token after kiro-cli re-login
+ if e.response.status_code == 400 and self._sqlite_db:
+ logger.warning("Token refresh failed with 400, reloading credentials from SQLite and retrying...")
+ self._load_credentials_from_sqlite(self._sqlite_db)
+ await self._do_aws_sso_oidc_refresh()
+ else:
+ raise
+
+ async def _do_aws_sso_oidc_refresh(self) -> None:
+ """
+ Performs the actual AWS SSO OIDC token refresh.
+
+ This is the internal implementation called by _refresh_token_aws_sso_oidc().
+ It performs a single refresh attempt with current in-memory credentials.
+
+ Uses AWS SSO OIDC CreateToken API format:
+ - Content-Type: application/json (not form-urlencoded)
+ - Parameter names: camelCase (clientId, not client_id)
+ - Payload: JSON object
+
+ Raises:
+ ValueError: If required credentials are not set
+ httpx.HTTPStatusError: On HTTP error (including 400 for invalid token)
+ """
+ if not self._refresh_token:
+ raise ValueError("Refresh token is not set")
+ if not self._client_id:
+ raise ValueError("Client ID is not set (required for AWS SSO OIDC)")
+ if not self._client_secret:
+ raise ValueError("Client secret is not set (required for AWS SSO OIDC)")
+
+ logger.info("Refreshing Kiro token via AWS SSO OIDC...")
+
+ # AWS SSO OIDC CreateToken API uses JSON with camelCase parameters
+ # Use SSO region for OIDC endpoint (may differ from API region)
+ sso_region = self._sso_region or self._region
+ url = get_aws_sso_oidc_url(sso_region)
+
+ # IMPORTANT: AWS SSO OIDC CreateToken API requires:
+ # 1. JSON payload (not form-urlencoded)
+ # 2. camelCase parameter names (clientId, not client_id)
+ payload = {
+ "grantType": "refresh_token",
+ "clientId": self._client_id,
+ "clientSecret": self._client_secret,
+ "refreshToken": self._refresh_token,
+ }
+
+ headers = {
+ "Content-Type": "application/json",
+ }
+
+ # Log request details (without secrets) for debugging
+ logger.debug(f"AWS SSO OIDC refresh request: url={url}, sso_region={sso_region}, "
+ f"api_region={self._region}, client_id={self._client_id[:8]}...")
+
+ async with httpx.AsyncClient(timeout=30) as client:
+ response = await client.post(url, json=payload, headers=headers)
+
+ # Log response details for debugging (especially on errors)
+ if response.status_code != 200:
+ error_body = response.text
+ logger.error(f"AWS SSO OIDC refresh failed: status={response.status_code}, "
+ f"body={error_body}")
+ # Try to parse AWS error for more details
+ try:
+ error_json = response.json()
+ error_code = error_json.get("error", "unknown")
+ error_desc = error_json.get("error_description", "no description")
+ logger.error(f"AWS SSO OIDC error details: error={error_code}, "
+ f"description={error_desc}")
+ except Exception:
+ pass # Body wasn't JSON, already logged as text
+ response.raise_for_status()
+
+ result = response.json()
+
+ # AWS SSO OIDC CreateToken API returns camelCase fields
+ new_access_token = result.get("accessToken")
+ new_refresh_token = result.get("refreshToken")
+ expires_in = result.get("expiresIn", 3600)
+
+ if not new_access_token:
+ raise ValueError(f"AWS SSO OIDC response does not contain accessToken: {result}")
+
+ # Update data
+ self._access_token = new_access_token
+ if new_refresh_token:
+ self._refresh_token = new_refresh_token
+
+ # Calculate expiration time with buffer (minus 60 seconds)
+ self._expires_at = datetime.now(timezone.utc) + timedelta(seconds=expires_in - 60)
+
+ logger.info(f"Token refreshed via AWS SSO OIDC, expires: {self._expires_at.isoformat()}")
+
+ # Save to file or SQLite depending on configuration
+ if self._sqlite_db:
+ self._save_credentials_to_sqlite()
+ else:
+ self._save_credentials_to_file()
+
+ async def get_access_token(self) -> str:
+ """
+ Returns a valid access_token, refreshing it if necessary.
+
+ Thread-safe method using asyncio.Lock.
+ Automatically refreshes the token if it has expired or is about to expire.
+
+ For SQLite mode (kiro-cli): implements graceful degradation when refresh fails.
+ If kiro-cli has been running and refreshing tokens in memory (without persisting
+ to SQLite), the refresh_token in SQLite becomes stale. In this case, we fall back
+ to using the access_token directly until it actually expires.
+
+ Returns:
+ Valid access token
+
+ Raises:
+ ValueError: If unable to obtain access token
+ """
+ async with self._lock:
+ # Token is valid and not expiring soon - just return it
+ if self._access_token and not self.is_token_expiring_soon():
+ return self._access_token
+
+ # SQLite mode: reload credentials first, kiro-cli might have updated them
+ if self._sqlite_db and self.is_token_expiring_soon():
+ logger.debug("SQLite mode: reloading credentials before refresh attempt")
+ self._load_credentials_from_sqlite(self._sqlite_db)
+ # Check if reloaded token is now valid
+ if self._access_token and not self.is_token_expiring_soon():
+ logger.debug("SQLite reload provided fresh token, no refresh needed")
+ return self._access_token
+
+ # Try to refresh the token
+ try:
+ await self._refresh_token_request()
+ except httpx.HTTPStatusError as e:
+ # Graceful degradation for SQLite mode when refresh fails twice
+ # This happens when kiro-cli refreshed tokens in memory without persisting
+ if e.response.status_code == 400 and self._sqlite_db:
+ logger.warning(
+ "Token refresh failed with 400 after SQLite reload. "
+ "This may happen if kiro-cli refreshed tokens in memory without persisting."
+ )
+ # Check if access_token is still usable
+ if self._access_token and not self.is_token_expired():
+ logger.warning(
+ "Using existing access_token until it expires. "
+ "Run 'kiro-cli login' when convenient to refresh credentials."
+ )
+ return self._access_token
+ else:
+ raise ValueError(
+ "Token expired and refresh failed. "
+ "Please run 'kiro-cli login' to refresh your credentials."
+ )
+ # Non-SQLite mode or non-400 error - propagate the exception
+ raise
+ except Exception:
+ # For any other exception, propagate it
+ raise
+
+ if not self._access_token:
+ raise ValueError("Failed to obtain access token")
+
+ return self._access_token
+
+ async def force_refresh(self) -> str:
+ """
+ Forces a token refresh.
+
+ Used when receiving a 403 error from the API.
+
+ Returns:
+ New access token
+ """
+ async with self._lock:
+ await self._refresh_token_request()
+ return self._access_token
+
+ @property
+ def profile_arn(self) -> Optional[str]:
+ """AWS CodeWhisperer profile ARN."""
+ return self._profile_arn
+
+ @property
+ def region(self) -> str:
+ """AWS region."""
+ return self._region
+
+ @property
+ def api_host(self) -> str:
+ """API host for the current region."""
+ return self._api_host
+
+ @property
+ def q_host(self) -> str:
+ """Q API host for the current region."""
+ return self._q_host
+
+ @property
+ def fingerprint(self) -> str:
+ """Unique machine fingerprint."""
+ return self._fingerprint
+
+ @property
+ def auth_type(self) -> AuthType:
+ """Authentication type (KIRO_DESKTOP or AWS_SSO_OIDC)."""
+ return self._auth_type
\ No newline at end of file
diff --git a/kiro-gateway/kiro/cache.py b/kiro-gateway/kiro/cache.py
new file mode 100644
index 0000000000000000000000000000000000000000..f0be72af73d928faaf8edba37c45ee0da00d85bd
--- /dev/null
+++ b/kiro-gateway/kiro/cache.py
@@ -0,0 +1,182 @@
+# -*- 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 .
+
+"""
+Model metadata cache for Kiro Gateway.
+
+Thread-safe storage for available model information
+with TTL and lazy loading support.
+"""
+
+import asyncio
+import time
+from typing import Any, Dict, List, Optional
+
+from loguru import logger
+
+from kiro.config import MODEL_CACHE_TTL, DEFAULT_MAX_INPUT_TOKENS
+
+
+class ModelInfoCache:
+ """
+ Thread-safe cache for storing model metadata.
+
+ Uses Lazy Loading for population - data is loaded
+ only on first access or when cache is stale.
+
+ Attributes:
+ cache_ttl: Cache time-to-live in seconds
+
+ Example:
+ >>> cache = ModelInfoCache()
+ >>> await cache.update([{"modelId": "claude-sonnet-4", "tokenLimits": {...}}])
+ >>> info = cache.get("claude-sonnet-4")
+ >>> max_tokens = cache.get_max_input_tokens("claude-sonnet-4")
+ """
+
+ def __init__(self, cache_ttl: int = MODEL_CACHE_TTL):
+ """
+ Initializes the model cache.
+
+ Args:
+ cache_ttl: Cache time-to-live in seconds (default from config)
+ """
+ self._cache: Dict[str, Dict[str, Any]] = {}
+ self._lock = asyncio.Lock()
+ self._last_update: Optional[float] = None
+ self._cache_ttl = cache_ttl
+
+ async def update(self, models_data: List[Dict[str, Any]]) -> None:
+ """
+ Updates the model cache.
+
+ Thread-safely replaces cache contents with new data.
+
+ Args:
+ models_data: List of dictionaries with model information.
+ Each dictionary must contain the "modelId" key.
+ """
+ async with self._lock:
+ logger.info(f"Updating model cache. Found {len(models_data)} models.")
+ self._cache = {model["modelId"]: model for model in models_data}
+ self._last_update = time.time()
+
+ def get(self, model_id: str) -> Optional[Dict[str, Any]]:
+ """
+ Returns model information.
+
+ Args:
+ model_id: Model ID
+
+ Returns:
+ Dictionary with model information or None if model not found
+ """
+ return self._cache.get(model_id)
+
+ def is_valid_model(self, model_id: str) -> bool:
+ """
+ Check if model exists in dynamic cache.
+
+ Used by ModelResolver to verify if a model is available.
+
+ Args:
+ model_id: Model ID to check
+
+ Returns:
+ True if model exists in cache, False otherwise
+ """
+ return model_id in self._cache
+
+ def add_hidden_model(self, display_name: str, internal_id: str) -> None:
+ """
+ Add a hidden model to the cache.
+
+ Hidden models are not returned by Kiro /ListAvailableModels API
+ but are still functional. They are added to the cache so they
+ appear in our /v1/models endpoint.
+
+ Args:
+ display_name: Model name to display (e.g., "claude-3.7-sonnet")
+ internal_id: Internal Kiro ID (e.g., "CLAUDE_3_7_SONNET_20250219_V1_0")
+ """
+ if display_name not in self._cache:
+ self._cache[display_name] = {
+ "modelId": display_name,
+ "modelName": display_name,
+ "description": f"Hidden model (internal: {internal_id})",
+ "tokenLimits": {"maxInputTokens": DEFAULT_MAX_INPUT_TOKENS},
+ "_internal_id": internal_id, # Store internal ID for reference
+ "_is_hidden": True, # Mark as hidden model
+ }
+ logger.debug(f"Added hidden model: {display_name} → {internal_id}")
+
+ def get_max_input_tokens(self, model_id: str) -> int:
+ """
+ Returns maxInputTokens for the model.
+
+ Args:
+ model_id: Model ID
+
+ Returns:
+ Maximum number of input tokens or DEFAULT_MAX_INPUT_TOKENS
+ """
+ model = self._cache.get(model_id)
+ if model and model.get("tokenLimits"):
+ return model["tokenLimits"].get("maxInputTokens") or DEFAULT_MAX_INPUT_TOKENS
+ return DEFAULT_MAX_INPUT_TOKENS
+
+ def is_empty(self) -> bool:
+ """
+ Checks if the cache is empty.
+
+ Returns:
+ True if cache is empty
+ """
+ return not self._cache
+
+ def is_stale(self) -> bool:
+ """
+ Checks if the cache is stale.
+
+ Returns:
+ True if cache is stale (more than cache_ttl seconds have passed)
+ or if cache was never updated
+ """
+ if not self._last_update:
+ return True
+ return time.time() - self._last_update > self._cache_ttl
+
+ def get_all_model_ids(self) -> List[str]:
+ """
+ Returns a list of all model IDs in the cache.
+
+ Returns:
+ List of model IDs
+ """
+ return list(self._cache.keys())
+
+ @property
+ def size(self) -> int:
+ """Number of models in the cache."""
+ return len(self._cache)
+
+ @property
+ def last_update_time(self) -> Optional[float]:
+ """Last update time (timestamp) or None."""
+ return self._last_update
\ No newline at end of file
diff --git a/kiro-gateway/kiro/config.py b/kiro-gateway/kiro/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..247d0fdbeca134d034da1b6ca37a6eae4364cc8e
--- /dev/null
+++ b/kiro-gateway/kiro/config.py
@@ -0,0 +1,450 @@
+# -*- 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 .
+
+"""
+Kiro Gateway Configuration.
+
+Centralized storage for all settings, constants, and mappings.
+Loads environment variables and provides typed access to them.
+"""
+
+import os
+import re
+from pathlib import Path
+from typing import Dict, List, Optional
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+
+def _get_raw_env_value(var_name: str, env_file: str = ".env") -> Optional[str]:
+ """
+ Read variable value from .env file without processing escape sequences.
+
+ This is necessary for correct handling of Windows paths where backslashes
+ (e.g., D:\\Projects\\file.json) may be incorrectly interpreted
+ as escape sequences (\\a -> bell, \\n -> newline, etc.).
+
+ Args:
+ var_name: Environment variable name
+ env_file: Path to .env file (default ".env")
+
+ Returns:
+ Raw variable value or None if not found
+ """
+ env_path = Path(env_file)
+ if not env_path.exists():
+ return None
+
+ try:
+ # Read file as-is, without interpretation
+ content = env_path.read_text(encoding="utf-8")
+
+ # Search for variable considering different formats:
+ # VAR="value" or VAR='value' or VAR=value
+ # Pattern captures value with or without quotes
+ pattern = rf'^{re.escape(var_name)}=(["\']?)(.+?)\1\s*$'
+
+ for line in content.splitlines():
+ line = line.strip()
+ if line.startswith("#") or not line:
+ continue
+
+ match = re.match(pattern, line)
+ if match:
+ # Return value as-is, without processing escape sequences
+ return match.group(2)
+ except Exception:
+ pass
+
+ return None
+
+# ==================================================================================================
+# Server Settings
+# ==================================================================================================
+
+# Server host (default: 0.0.0.0 - listen on all interfaces)
+# Use "127.0.0.1" to only allow local connections
+DEFAULT_SERVER_HOST: str = "0.0.0.0"
+SERVER_HOST: str = os.getenv("SERVER_HOST", DEFAULT_SERVER_HOST)
+
+# Server port (default: 8000)
+# Can be overridden by CLI: python main.py --port 9000
+# Or by uvicorn directly: uvicorn main:app --port 9000
+DEFAULT_SERVER_PORT: int = 8000
+SERVER_PORT: int = int(os.getenv("SERVER_PORT", str(DEFAULT_SERVER_PORT)))
+
+# ==================================================================================================
+# Proxy Server Settings
+# ==================================================================================================
+
+# API key for proxy access (clients must pass it in Authorization header)
+PROXY_API_KEY: str = os.getenv("PROXY_API_KEY", "my-super-secret-password-123")
+
+# ==================================================================================================
+# VPN/Proxy Settings for Kiro API Access
+# ==================================================================================================
+
+# VPN/Proxy URL for accessing Kiro API through a proxy server.
+# Leave empty to connect directly (default).
+#
+# Use cases:
+# - China: GFW (Great Firewall) blocks AWS endpoints
+# - Corporate networks: Often require mandatory proxy
+# - Privacy: Hide your IP address from AWS
+#
+# Supports HTTP and SOCKS5 protocols.
+# Authentication can be embedded in the URL.
+#
+# Examples:
+# VPN_PROXY_URL=http://127.0.0.1:7890
+# VPN_PROXY_URL=socks5://127.0.0.1:1080
+# VPN_PROXY_URL=http://user:password@proxy.company.com:8080
+# VPN_PROXY_URL=192.168.1.100:8080 (defaults to http://)
+VPN_PROXY_URL: str = os.getenv("VPN_PROXY_URL", "")
+
+# ==================================================================================================
+# Kiro API Credentials
+# ==================================================================================================
+
+# Refresh token for updating access token
+REFRESH_TOKEN: str = os.getenv("REFRESH_TOKEN", "")
+
+# Profile ARN for AWS CodeWhisperer
+PROFILE_ARN: str = os.getenv("PROFILE_ARN", "")
+
+# AWS region (default us-east-1)
+REGION: str = os.getenv("KIRO_REGION", "us-east-1")
+
+# Path to credentials file (optional, alternative to .env)
+# Read directly from .env to avoid escape sequence issues on Windows
+# (e.g., \a in path D:\Projects\adolf is interpreted as bell character)
+_raw_creds_file = _get_raw_env_value("KIRO_CREDS_FILE") or os.getenv("KIRO_CREDS_FILE", "")
+# Normalize path for cross-platform compatibility
+KIRO_CREDS_FILE: str = str(Path(_raw_creds_file)) if _raw_creds_file else ""
+
+# Path to kiro-cli SQLite database (optional, for AWS SSO OIDC authentication)
+# Default location: ~/.local/share/kiro-cli/data.sqlite3 (Linux/macOS)
+# or ~/.local/share/amazon-q/data.sqlite3 (amazon-q-developer-cli)
+_raw_cli_db_file = _get_raw_env_value("KIRO_CLI_DB_FILE") or os.getenv("KIRO_CLI_DB_FILE", "")
+KIRO_CLI_DB_FILE: str = str(Path(_raw_cli_db_file)) if _raw_cli_db_file else ""
+
+# ==================================================================================================
+# Kiro API URL Templates
+# ==================================================================================================
+
+# URL for token refresh (Kiro Desktop Auth)
+KIRO_REFRESH_URL_TEMPLATE: str = "https://prod.{region}.auth.desktop.kiro.dev/refreshToken"
+
+# URL for token refresh (AWS SSO OIDC - used by kiro-cli)
+AWS_SSO_OIDC_URL_TEMPLATE: str = "https://oidc.{region}.amazonaws.com/token"
+
+# Host for main API (generateAssistantResponse)
+KIRO_API_HOST_TEMPLATE: str = "https://codewhisperer.{region}.amazonaws.com"
+
+# Host for Q API (ListAvailableModels)
+KIRO_Q_HOST_TEMPLATE: str = "https://q.{region}.amazonaws.com"
+
+# ==================================================================================================
+# Token Settings
+# ==================================================================================================
+
+# Time before token expiration when refresh is needed (in seconds)
+# Default 10 minutes - refresh token in advance to avoid errors
+TOKEN_REFRESH_THRESHOLD: int = 600
+
+# ==================================================================================================
+# Retry Configuration
+# ==================================================================================================
+
+# Maximum number of retry attempts on errors
+MAX_RETRIES: int = 3
+
+# Base delay between attempts (seconds)
+# Uses exponential backoff: delay * (2 ** attempt)
+BASE_RETRY_DELAY: float = 1.0
+
+# ==================================================================================================
+# Hidden Models Configuration
+# ==================================================================================================
+
+# Hidden models - not returned by Kiro /ListAvailableModels API but still functional.
+# These ARE shown in our /v1/models endpoint!
+# Use dot format for consistency with API models.
+#
+# Format: "display_name" → "internal_kiro_id"
+# Display names use dots (e.g., "claude-3.7-sonnet") for consistency with Kiro API.
+#
+# Why "hidden"? These models work but are not advertised by Kiro's /ListAvailableModels.
+# We expose them to our users because they're useful.
+HIDDEN_MODELS: Dict[str, str] = {
+ # Claude 3.7 Sonnet - legacy flagship model, still works!
+ # Hidden in Kiro API but functional. Great for users who prefer it.
+ "claude-3.7-sonnet": "CLAUDE_3_7_SONNET_20250219_V1_0",
+
+ # Add other hidden/experimental models here as discovered.
+ # Example: "claude-secret-model": "INTERNAL_SECRET_MODEL_ID",
+}
+
+# ==================================================================================================
+# Fallback Models Configuration (DNS Failure Recovery)
+# ==================================================================================================
+
+# Fallback model list - used when /ListAvailableModels API is unreachable.
+# This ensures basic functionality even with DNS/network issues.
+#
+# IMPORTANT: This list represents known models at the time of this gateway version.
+# - Some models may not be available on your Kiro plan (e.g., Opus on free tier)
+# - New models released after this version won't appear here
+# - Update gateway regularly to get the latest model list
+FALLBACK_MODELS: List[Dict[str, str]] = [
+ {"modelId": "auto"},
+ {"modelId": "claude-sonnet-4"},
+ {"modelId": "claude-haiku-4.5"},
+ {"modelId": "claude-sonnet-4.5"},
+ {"modelId": "claude-opus-4.5"},
+]
+
+# ==================================================================================================
+# Model Cache Settings
+# ==================================================================================================
+
+# Model cache TTL in seconds (1 hour)
+MODEL_CACHE_TTL: int = 3600
+
+# Default maximum number of input tokens
+DEFAULT_MAX_INPUT_TOKENS: int = 200000
+
+# ==================================================================================================
+# Tool Description Handling (Kiro API Limitations)
+# ==================================================================================================
+
+# Kiro API returns 400 "Improperly formed request" error when tool descriptions
+# in toolSpecification.description are too long.
+#
+# Solution: Tool Documentation Reference Pattern
+# - If description ≤ limit → keep as is
+# - If description > limit:
+# * In toolSpecification.description → reference to system prompt:
+# "[Full documentation in system prompt under '## Tool: {name}']"
+# * In system prompt, a section "## Tool: {name}" with full description is added
+#
+# The model sees an explicit reference and knows exactly where to find full documentation.
+
+# Maximum length of tool description in characters.
+# Descriptions longer than this limit will be moved to system prompt.
+# Set to 0 to disable (not recommended - will cause Kiro API errors).
+TOOL_DESCRIPTION_MAX_LENGTH: int = int(os.getenv("TOOL_DESCRIPTION_MAX_LENGTH", "10000"))
+
+# ==================================================================================================
+# Logging Settings
+# ==================================================================================================
+
+# Log level for the application
+# Available levels: TRACE, DEBUG, INFO, WARNING, ERROR, CRITICAL
+# Default: INFO (recommended for production)
+# Set to DEBUG for detailed troubleshooting
+LOG_LEVEL: str = os.getenv("LOG_LEVEL", "INFO").upper()
+
+# ==================================================================================================
+# First Token Timeout Settings (Streaming Retry)
+# ==================================================================================================
+
+# 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 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"))
+
+# Read timeout for streaming responses (in seconds).
+# This is the maximum time to wait for data between chunks during streaming.
+# Should be longer than FIRST_TOKEN_TIMEOUT since the model may pause between chunks
+# while "thinking" (especially for tool calls or complex reasoning).
+# Default: 300 seconds (5 minutes) - generous timeout to avoid premature disconnects.
+STREAMING_READ_TIMEOUT: float = float(os.getenv("STREAMING_READ_TIMEOUT", "300"))
+
+# Maximum number of attempts on first token timeout.
+# After exhausting all attempts, an error will be returned.
+# Default: 3 attempts
+FIRST_TOKEN_MAX_RETRIES: int = int(os.getenv("FIRST_TOKEN_MAX_RETRIES", "3"))
+
+# ==================================================================================================
+# Debug Settings
+# ==================================================================================================
+
+# Legacy option (deprecated, will be removed in future releases)
+# Use DEBUG_MODE instead
+_DEBUG_LAST_REQUEST_RAW: str = os.getenv("DEBUG_LAST_REQUEST", "").lower()
+DEBUG_LAST_REQUEST: bool = _DEBUG_LAST_REQUEST_RAW in ("true", "1", "yes")
+
+# Debug logging mode:
+# - off: disabled (default)
+# - errors: save logs only for failed requests (4xx, 5xx)
+# - all: save logs for every request (overwrites on each request)
+_DEBUG_MODE_RAW: str = os.getenv("DEBUG_MODE", "").lower()
+
+# Priority logic:
+# 1. If DEBUG_MODE is explicitly set → use it
+# 2. If DEBUG_MODE is not set but DEBUG_LAST_REQUEST=true → mode "all" (backward compatibility)
+# 3. Otherwise → mode "off"
+if _DEBUG_MODE_RAW in ("off", "errors", "all"):
+ DEBUG_MODE: str = _DEBUG_MODE_RAW
+elif DEBUG_LAST_REQUEST:
+ DEBUG_MODE: str = "all"
+else:
+ DEBUG_MODE: str = "off"
+
+# Directory for debug log files
+DEBUG_DIR: str = os.getenv("DEBUG_DIR", "debug_logs")
+
+
+def _warn_deprecated_debug_setting():
+ """
+ Print warning if deprecated DEBUG_LAST_REQUEST is used.
+ Called at application startup.
+ """
+ if _DEBUG_LAST_REQUEST_RAW and not _DEBUG_MODE_RAW:
+ import sys
+ # ANSI escape codes: yellow text
+ YELLOW = "\033[93m"
+ RESET = "\033[0m"
+
+ warning_text = f"""
+{YELLOW}⚠️ DEPRECATED: DEBUG_LAST_REQUEST will be removed in future releases.
+ Please use DEBUG_MODE instead:
+ - DEBUG_MODE=off (disabled, default)
+ - DEBUG_MODE=errors (save logs only for failed requests)
+ - DEBUG_MODE=all (save logs for every request)
+
+ DEBUG_LAST_REQUEST=true is equivalent to DEBUG_MODE=all
+ See .env.example for more details.{RESET}
+"""
+ print(warning_text, file=sys.stderr)
+
+
+def _warn_timeout_configuration():
+ """
+ Print warning if timeout configuration is suboptimal.
+ Called at application startup.
+
+ FIRST_TOKEN_TIMEOUT should be less than STREAMING_READ_TIMEOUT:
+ - FIRST_TOKEN_TIMEOUT: time to wait for model to START responding
+ - STREAMING_READ_TIMEOUT: time to wait BETWEEN chunks during streaming
+ """
+ if FIRST_TOKEN_TIMEOUT >= STREAMING_READ_TIMEOUT:
+ import sys
+ YELLOW = "\033[93m"
+ RESET = "\033[0m"
+
+ warning_text = f"""
+{YELLOW}⚠️ WARNING: Suboptimal timeout configuration detected.
+
+ FIRST_TOKEN_TIMEOUT ({FIRST_TOKEN_TIMEOUT}s) >= STREAMING_READ_TIMEOUT ({STREAMING_READ_TIMEOUT}s)
+
+ These timeouts serve different purposes:
+ - FIRST_TOKEN_TIMEOUT: time to wait for model to START responding (default: 15s)
+ - STREAMING_READ_TIMEOUT: time to wait BETWEEN chunks during streaming (default: 300s)
+
+ Recommendation: FIRST_TOKEN_TIMEOUT should be LESS than STREAMING_READ_TIMEOUT.
+
+ Example configuration:
+ FIRST_TOKEN_TIMEOUT=15
+ STREAMING_READ_TIMEOUT=300{RESET}
+"""
+ print(warning_text, file=sys.stderr)
+
+# ==================================================================================================
+# Fake Reasoning Settings (Extended Thinking via Tag Injection)
+# ==================================================================================================
+
+# Enable fake reasoning - injects special tags into requests to enable model reasoning.
+# When enabled, the model will include its reasoning process in the response wrapped in tags.
+# The response is then parsed and converted to OpenAI-compatible reasoning_content format.
+#
+# WHY "FAKE"? This is NOT native extended thinking API support. Instead, we inject
+# enabled tags into the prompt, and the model responds
+# with ... blocks that we parse and convert to reasoning_content.
+# It works great, but it's a hack - hence "fake" reasoning.
+#
+# Default: true (enabled) - provides premium experience out of the box
+_FAKE_REASONING_RAW: str = os.getenv("FAKE_REASONING", "").lower()
+# Default is True - if env var is not set or empty, enable fake reasoning
+FAKE_REASONING_ENABLED: bool = _FAKE_REASONING_RAW not in ("false", "0", "no", "disabled", "off")
+
+# Maximum thinking length in tokens.
+# This value is injected into the request as {value}
+# Higher values allow for more detailed reasoning but increase response time and token usage.
+# Default: 4000 tokens
+FAKE_REASONING_MAX_TOKENS: int = int(os.getenv("FAKE_REASONING_MAX_TOKENS", "4000"))
+
+# How to handle the thinking block in responses:
+# - "as_reasoning_content": Extract to reasoning_content field (OpenAI-compatible, recommended)
+# - "remove": Remove thinking block completely, return only final answer
+# - "pass": Pass through as-is with original tags in content
+# - "strip_tags": Remove tags but keep thinking content in regular content
+#
+# Default: "as_reasoning_content"
+_FAKE_REASONING_HANDLING_RAW: str = os.getenv("FAKE_REASONING_HANDLING", "as_reasoning_content").lower()
+if _FAKE_REASONING_HANDLING_RAW in ("as_reasoning_content", "remove", "pass", "strip_tags"):
+ FAKE_REASONING_HANDLING: str = _FAKE_REASONING_HANDLING_RAW
+else:
+ FAKE_REASONING_HANDLING: str = "as_reasoning_content"
+
+# List of opening tags to detect thinking blocks.
+# The parser will look for any of these tags at the start of the response.
+# Order matters - first match wins.
+FAKE_REASONING_OPEN_TAGS: List[str] = ["", "", "", ""]
+
+# Maximum size of initial buffer for tag detection (characters).
+# If no thinking tag is found within this limit, content is treated as regular response.
+# Lower values = faster first token, but may miss tags with leading whitespace.
+# Default: 30 characters (enough for longest tag + some whitespace)
+FAKE_REASONING_INITIAL_BUFFER_SIZE: int = int(os.getenv("FAKE_REASONING_INITIAL_BUFFER_SIZE", "20"))
+
+
+# ==================================================================================================
+# Application Version
+# ==================================================================================================
+
+APP_VERSION: str = "2.1"
+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"
+
+
+def get_kiro_refresh_url(region: str) -> str:
+ """Return Kiro Desktop Auth token refresh URL for the specified region."""
+ return KIRO_REFRESH_URL_TEMPLATE.format(region=region)
+
+
+def get_aws_sso_oidc_url(region: str) -> str:
+ """Return AWS SSO OIDC token URL for the specified region."""
+ return AWS_SSO_OIDC_URL_TEMPLATE.format(region=region)
+
+
+def get_kiro_api_host(region: str) -> str:
+ """Return API host for the specified region."""
+ return KIRO_API_HOST_TEMPLATE.format(region=region)
+
+
+def get_kiro_q_host(region: str) -> str:
+ """Return Q API host for the specified region."""
+ return KIRO_Q_HOST_TEMPLATE.format(region=region)
+
diff --git a/kiro-gateway/kiro/converters_anthropic.py b/kiro-gateway/kiro/converters_anthropic.py
new file mode 100644
index 0000000000000000000000000000000000000000..ab69dfcffd7f525c4278deaf6ad5b52f00d794bf
--- /dev/null
+++ b/kiro-gateway/kiro/converters_anthropic.py
@@ -0,0 +1,369 @@
+# -*- 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 Anthropic Messages API format to Kiro format.
+
+This module is an adapter layer that converts Anthropic-specific formats
+to the unified format used by converters_core.py.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from loguru import logger
+
+from kiro.config import HIDDEN_MODELS
+from kiro.model_resolver import get_model_id_for_kiro
+from kiro.models_anthropic import (
+ AnthropicMessagesRequest,
+ AnthropicMessage,
+ AnthropicTool,
+)
+from kiro.converters_core import (
+ UnifiedMessage,
+ UnifiedTool,
+ build_kiro_payload,
+ extract_text_content,
+ extract_images_from_content,
+)
+
+
+def convert_anthropic_content_to_text(content: Any) -> str:
+ """
+ Extracts text content from Anthropic message content.
+
+ Anthropic content can be:
+ - String: "Hello, world!"
+ - List of content blocks: [{"type": "text", "text": "Hello"}]
+
+ Args:
+ content: Anthropic message content
+
+ Returns:
+ Extracted text content
+ """
+ if isinstance(content, str):
+ return content
+
+ if isinstance(content, list):
+ text_parts = []
+ for block in content:
+ if isinstance(block, dict):
+ if block.get("type") == "text":
+ text_parts.append(block.get("text", ""))
+ elif hasattr(block, "type") and block.type == "text":
+ text_parts.append(block.text)
+ return "".join(text_parts)
+
+ return str(content) if content else ""
+
+
+def extract_system_prompt(system: Any) -> str:
+ """
+ Extracts system prompt text from Anthropic system field.
+
+ Anthropic API supports system in two formats:
+ 1. String: "You are helpful"
+ 2. List of content blocks: [{"type": "text", "text": "...", "cache_control": {...}}]
+
+ The second format is used for prompt caching with cache_control.
+ We extract only the text, ignoring cache_control (not supported by Kiro).
+
+ Args:
+ system: System prompt in string or list format
+
+ Returns:
+ Extracted system prompt as string
+ """
+ if system is None:
+ return ""
+
+ if isinstance(system, str):
+ return system
+
+ if isinstance(system, list):
+ text_parts = []
+ for block in system:
+ if isinstance(block, dict):
+ # Handle {"type": "text", "text": "...", "cache_control": {...}}
+ if block.get("type") == "text":
+ text_parts.append(block.get("text", ""))
+ elif hasattr(block, "type") and block.type == "text":
+ # Handle Pydantic model
+ text_parts.append(getattr(block, "text", ""))
+ return "\n".join(text_parts)
+
+ return str(system)
+
+
+def extract_tool_results_from_anthropic_content(content: Any) -> List[Dict[str, Any]]:
+ """
+ Extracts tool results from Anthropic message content.
+
+ Looks for content blocks with type="tool_result".
+
+ Args:
+ content: Anthropic message content (list of content blocks)
+
+ Returns:
+ List of tool results in unified format
+ """
+ tool_results = []
+
+ if not isinstance(content, list):
+ return tool_results
+
+ for block in content:
+ block_type = None
+ tool_use_id = None
+ result_content = ""
+
+ if isinstance(block, dict):
+ block_type = block.get("type")
+ tool_use_id = block.get("tool_use_id")
+ result_content = block.get("content", "")
+ elif hasattr(block, "type"):
+ block_type = block.type
+ tool_use_id = getattr(block, "tool_use_id", None)
+ result_content = getattr(block, "content", "")
+
+ if block_type == "tool_result" and tool_use_id:
+ # Convert content to text if it's a list
+ if isinstance(result_content, list):
+ result_content = extract_text_content(result_content)
+ elif not isinstance(result_content, str):
+ result_content = str(result_content) if result_content else ""
+
+ tool_results.append({
+ "type": "tool_result",
+ "tool_use_id": tool_use_id,
+ "content": result_content or "(empty result)"
+ })
+
+ return tool_results
+
+
+def extract_tool_uses_from_anthropic_content(content: Any) -> List[Dict[str, Any]]:
+ """
+ Extracts tool uses from Anthropic assistant message content.
+
+ Looks for content blocks with type="tool_use".
+
+ Args:
+ content: Anthropic message content (list of content blocks)
+
+ Returns:
+ List of tool calls in unified format
+ """
+ tool_calls = []
+
+ if not isinstance(content, list):
+ return tool_calls
+
+ for block in content:
+ block_type = None
+ tool_id = None
+ tool_name = None
+ tool_input = {}
+
+ if isinstance(block, dict):
+ block_type = block.get("type")
+ tool_id = block.get("id")
+ tool_name = block.get("name")
+ tool_input = block.get("input", {})
+ elif hasattr(block, "type"):
+ block_type = block.type
+ tool_id = getattr(block, "id", None)
+ tool_name = getattr(block, "name", None)
+ tool_input = getattr(block, "input", {})
+
+ if block_type == "tool_use" and tool_id and tool_name:
+ tool_calls.append({
+ "id": tool_id,
+ "type": "function",
+ "function": {
+ "name": tool_name,
+ "arguments": tool_input if isinstance(tool_input, str) else tool_input
+ }
+ })
+
+ return tool_calls
+
+
+def convert_anthropic_messages(messages: List[AnthropicMessage]) -> List[UnifiedMessage]:
+ """
+ Converts Anthropic messages to unified format.
+
+ Handles:
+ - Text content (string or list of text blocks)
+ - Tool use blocks (assistant messages)
+ - Tool result blocks (user messages)
+
+ Args:
+ messages: List of Anthropic messages
+
+ Returns:
+ List of messages in unified format
+ """
+
+ unified_messages = []
+ total_tool_calls = 0
+ total_tool_results = 0
+ total_images = 0
+
+ for msg in messages:
+ role = msg.role
+ content = msg.content
+
+ # Extract text content
+ text_content = convert_anthropic_content_to_text(content)
+
+ # Extract tool-related data and images based on role
+ tool_calls = None
+ tool_results = None
+ images = None
+
+ if role == "assistant":
+ # Assistant messages may contain tool_use blocks
+ tool_calls = extract_tool_uses_from_anthropic_content(content)
+ if tool_calls:
+ total_tool_calls += len(tool_calls)
+
+ elif role == "user":
+ # User messages may contain tool_result blocks and images
+ tool_results = extract_tool_results_from_anthropic_content(content)
+ if tool_results:
+ total_tool_results += len(tool_results)
+
+ # Extract images from user messages
+ images = extract_images_from_content(content)
+ if images:
+ total_images += len(images)
+
+ unified_msg = UnifiedMessage(
+ role=role,
+ content=text_content,
+ tool_calls=tool_calls if tool_calls else None,
+ tool_results=tool_results if tool_results else None,
+ images=images if images else None
+ )
+ unified_messages.append(unified_msg)
+
+ # Log summary if any tool content or images were found
+ if total_tool_calls > 0 or total_tool_results > 0 or total_images > 0:
+ logger.debug(
+ f"Converted {len(messages)} Anthropic messages: "
+ f"{total_tool_calls} tool_calls, {total_tool_results} tool_results, {total_images} images"
+ )
+
+ return unified_messages
+
+
+def convert_anthropic_tools(tools: Optional[List[AnthropicTool]]) -> Optional[List[UnifiedTool]]:
+ """
+ Converts Anthropic tools to unified format.
+
+ Args:
+ tools: List of Anthropic tools
+
+ Returns:
+ List of tools in unified format, or None if no tools
+ """
+ if not tools:
+ return None
+
+ unified_tools = []
+ for tool in tools:
+ # Handle both dict and Pydantic model
+ if isinstance(tool, dict):
+ name = tool.get("name", "")
+ description = tool.get("description")
+ input_schema = tool.get("input_schema", {})
+ else:
+ name = tool.name
+ description = tool.description
+ input_schema = tool.input_schema
+
+ unified_tools.append(UnifiedTool(
+ name=name,
+ description=description,
+ input_schema=input_schema
+ ))
+
+ return unified_tools if unified_tools else None
+
+
+def anthropic_to_kiro(
+ request: AnthropicMessagesRequest,
+ conversation_id: str,
+ profile_arn: str
+) -> dict:
+ """
+ Converts Anthropic Messages API request to Kiro API payload.
+
+ This is the main entry point for Anthropic → Kiro conversion.
+
+ Key differences from OpenAI:
+ - System prompt is a separate field (not in messages)
+ - Content can be string or list of content blocks
+ - Tool format uses input_schema instead of parameters
+
+ Args:
+ request: Anthropic MessagesRequest
+ conversation_id: Unique conversation ID
+ profile_arn: AWS CodeWhisperer profile ARN
+
+ Returns:
+ Payload dictionary for POST request to Kiro API
+
+ Raises:
+ ValueError: If there are no messages to send
+ """
+ # Convert messages to unified format
+ unified_messages = convert_anthropic_messages(request.messages)
+
+ # Convert tools to unified format
+ unified_tools = convert_anthropic_tools(request.tools)
+
+ # System prompt is already separate in Anthropic format!
+ # It can be a string or list of content blocks (for prompt caching)
+ system_prompt = extract_system_prompt(request.system)
+
+ # 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)
+
+ 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)}"
+ )
+
+ # Use core function to build payload
+ result = 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,
+ inject_thinking=True
+ )
+
+ return result.payload
\ No newline at end of file
diff --git a/kiro-gateway/kiro/converters_core.py b/kiro-gateway/kiro/converters_core.py
new file mode 100644
index 0000000000000000000000000000000000000000..34daa8197940e49ed9696a99eedfcf857943924f
--- /dev/null
+++ b/kiro-gateway/kiro/converters_core.py
@@ -0,0 +1,1310 @@
+# -*- 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 .
+
+"""
+Core converters for transforming API formats to Kiro format.
+
+This module contains shared logic used by both OpenAI and Anthropic converters:
+- Text content extraction from various formats
+- Message merging and processing
+- Kiro payload building
+- Tool processing and sanitization
+
+The core layer provides a unified interface that API-specific adapters use
+to convert their formats to Kiro API format.
+"""
+
+import json
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+from loguru import logger
+
+from kiro.config import (
+ TOOL_DESCRIPTION_MAX_LENGTH,
+ FAKE_REASONING_ENABLED,
+ FAKE_REASONING_MAX_TOKENS,
+)
+
+
+# ==================================================================================================
+# Data Classes for Unified Message Format
+# ==================================================================================================
+
+@dataclass
+class UnifiedMessage:
+ """
+ Unified message format used internally by converters.
+
+ This format is API-agnostic and can be created from both OpenAI and Anthropic formats.
+ Serves as the canonical representation for all message data before conversion to Kiro API.
+
+ Attributes:
+ role: Message role (user, assistant, system)
+ content: Text content or list of content blocks
+ tool_calls: List of tool calls (for assistant messages)
+ tool_results: List of tool results (for user messages with tool responses)
+ images: List of images in unified format (for multimodal user messages)
+ Format: [{"media_type": "image/jpeg", "data": "base64..."}]
+ """
+ role: str
+ content: Any = ""
+ tool_calls: Optional[List[Dict[str, Any]]] = None
+ tool_results: Optional[List[Dict[str, Any]]] = None
+ images: Optional[List[Dict[str, Any]]] = None
+
+
+@dataclass
+class UnifiedTool:
+ """
+ Unified tool format used internally by converters.
+
+ Attributes:
+ name: Tool name
+ description: Tool description
+ input_schema: JSON Schema for tool parameters
+ """
+ name: str
+ description: Optional[str] = None
+ input_schema: Optional[Dict[str, Any]] = None
+
+
+@dataclass
+class KiroPayloadResult:
+ """
+ Result of building Kiro payload.
+
+ Attributes:
+ payload: The complete Kiro API payload
+ tool_documentation: Documentation for tools with long descriptions (to add to system prompt)
+ """
+ payload: Dict[str, Any]
+ tool_documentation: str = ""
+
+
+# ==================================================================================================
+# Text Content Extraction
+# ==================================================================================================
+
+def extract_text_content(content: Any) -> str:
+ """
+ Extracts text content from various formats.
+
+ Supports multiple content formats used by different APIs:
+ - String: "Hello, world!"
+ - List of content blocks: [{"type": "text", "text": "Hello"}]
+ - None: empty message
+
+ Args:
+ content: Content in any supported format
+
+ Returns:
+ Extracted text or empty string
+
+ Example:
+ >>> extract_text_content("Hello")
+ 'Hello'
+ >>> extract_text_content([{"type": "text", "text": "World"}])
+ 'World'
+ >>> extract_text_content(None)
+ ''
+ """
+ if content is None:
+ return ""
+ if isinstance(content, str):
+ return content
+ if isinstance(content, list):
+ text_parts = []
+ for item in content:
+ if isinstance(item, dict):
+ # Skip image blocks - they're handled separately
+ if item.get("type") in ("image", "image_url"):
+ continue
+ if item.get("type") == "text":
+ text_parts.append(item.get("text", ""))
+ elif "text" in item:
+ text_parts.append(item["text"])
+ elif hasattr(item, "text"):
+ # Handle Pydantic models like TextContentBlock
+ text_parts.append(getattr(item, "text", ""))
+ elif isinstance(item, str):
+ text_parts.append(item)
+ return "".join(text_parts)
+ return str(content)
+
+
+def extract_images_from_content(content: Any) -> List[Dict[str, Any]]:
+ """
+ Extracts images from message content in unified format.
+
+ Supports multiple image formats used by different APIs:
+
+ OpenAI format (image_url with data URL):
+ {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,/9j/..."}}
+
+ Anthropic format (image with source):
+ {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": "/9j/..."}}
+
+ Args:
+ content: Content in any supported format (usually a list of content blocks)
+
+ Returns:
+ List of images in unified format: [{"media_type": "image/jpeg", "data": "base64..."}]
+ Empty list if no images found or content is not a list.
+
+ Example:
+ >>> extract_images_from_content([{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "abc123"}}])
+ [{'media_type': 'image/png', 'data': 'abc123'}]
+ """
+ images: List[Dict[str, Any]] = []
+
+ if not isinstance(content, list):
+ return images
+
+ for item in content:
+ # Handle both dict and Pydantic model objects
+ if isinstance(item, dict):
+ item_type = item.get("type")
+ elif hasattr(item, "type"):
+ item_type = item.type
+ else:
+ continue
+
+ # OpenAI format: {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}}
+ if item_type == "image_url":
+ if isinstance(item, dict):
+ image_url_obj = item.get("image_url", {})
+ else:
+ image_url_obj = getattr(item, "image_url", {})
+
+ if isinstance(image_url_obj, dict):
+ url = image_url_obj.get("url", "")
+ elif hasattr(image_url_obj, "url"):
+ url = image_url_obj.url
+ else:
+ url = ""
+
+ if url.startswith("data:"):
+ # Parse data URL: data:image/jpeg;base64,/9j/4AAQ...
+ try:
+ header, data = url.split(",", 1)
+ # Extract media type from "data:image/jpeg;base64"
+ media_part = header.split(";")[0] # "data:image/jpeg"
+ media_type = media_part.replace("data:", "") # "image/jpeg"
+
+ if data:
+ images.append({
+ "media_type": media_type,
+ "data": data
+ })
+ except (ValueError, IndexError) as e:
+ logger.warning(f"Failed to parse image data URL: {e}")
+ elif url.startswith("http"):
+ # URL-based images require fetching - not supported by Kiro API directly
+ logger.warning(f"URL-based images are not supported by Kiro API, skipping: {url[:80]}...")
+
+ # Anthropic format: {"type": "image", "source": {"type": "base64", "media_type": "...", "data": "..."}}
+ elif item_type == "image":
+ source = item.get("source", {}) if isinstance(item, dict) else getattr(item, "source", None)
+
+ if source is None:
+ continue
+
+ if isinstance(source, dict):
+ source_type = source.get("type")
+
+ if source_type == "base64":
+ media_type = source.get("media_type", "image/jpeg")
+ data = source.get("data", "")
+
+ if data:
+ images.append({
+ "media_type": media_type,
+ "data": data
+ })
+ elif source_type == "url":
+ # URL-based images in Anthropic format
+ url = source.get("url", "")
+ logger.warning(f"URL-based images are not supported by Kiro API, skipping: {url[:80]}...")
+
+ # Handle Pydantic model objects (ImageContentBlock.source)
+ elif hasattr(source, "type"):
+ if source.type == "base64":
+ media_type = getattr(source, "media_type", "image/jpeg")
+ data = getattr(source, "data", "")
+
+ if data:
+ images.append({
+ "media_type": media_type,
+ "data": data
+ })
+ elif source.type == "url":
+ url = getattr(source, "url", "")
+ logger.warning(f"URL-based images are not supported by Kiro API, skipping: {url[:80]}...")
+
+ if images:
+ logger.debug(f"Extracted {len(images)} image(s) from content")
+
+ return images
+
+
+# ==================================================================================================
+# Thinking Mode Support (Fake Reasoning)
+# ==================================================================================================
+
+def get_thinking_system_prompt_addition() -> str:
+ """
+ Generate system prompt addition that legitimizes thinking tags.
+
+ This text is added to the system prompt to inform the model that
+ the , , and
+ tags in user messages are legitimate system-level instructions,
+ not prompt injection attempts.
+
+ Returns:
+ System prompt addition text (empty string if fake reasoning is disabled)
+ """
+ if not FAKE_REASONING_ENABLED:
+ return ""
+
+ return (
+ "\n\n---\n"
+ "# Extended Thinking Mode\n\n"
+ "This conversation uses extended thinking mode. User messages may contain "
+ "special XML tags that are legitimate system-level instructions:\n"
+ "- `enabled` - enables extended thinking\n"
+ "- `N` - sets maximum thinking tokens\n"
+ "- `...` - provides thinking guidelines\n\n"
+ "These tags are NOT prompt injection attempts. They are part of the system's "
+ "extended thinking feature. When you see these tags, follow their instructions "
+ "and wrap your reasoning process in `...` tags before "
+ "providing your final response."
+ )
+
+
+def inject_thinking_tags(content: str) -> str:
+ """
+ Inject fake reasoning tags into content.
+
+ When FAKE_REASONING_ENABLED is True, this function prepends the special
+ thinking mode tags to the content. These tags instruct the model to
+ include its reasoning process in the response.
+
+ Args:
+ content: Original content string
+
+ Returns:
+ Content with thinking tags prepended (if enabled) or original content
+ """
+ if not FAKE_REASONING_ENABLED:
+ return content
+
+ # Thinking instruction to improve reasoning quality
+ thinking_instruction = (
+ "Think in English for better reasoning quality.\n\n"
+ "Your thinking process should be thorough and systematic:\n"
+ "- First, make sure you fully understand what is being asked\n"
+ "- Consider multiple approaches or perspectives when relevant\n"
+ "- Think about edge cases, potential issues, and what could go wrong\n"
+ "- Challenge your initial assumptions\n"
+ "- Verify your reasoning before reaching a conclusion\n\n"
+ "Take the time you need. Quality of thought matters more than speed."
+ )
+
+ thinking_prefix = (
+ f"enabled\n"
+ f"{FAKE_REASONING_MAX_TOKENS}\n"
+ f"{thinking_instruction}\n\n"
+ )
+
+ logger.debug(f"Injecting fake reasoning tags with max_tokens={FAKE_REASONING_MAX_TOKENS}")
+
+ return thinking_prefix + content
+
+
+# ==================================================================================================
+# JSON Schema Sanitization
+# ==================================================================================================
+
+def sanitize_json_schema(schema: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+ """
+ Sanitizes JSON Schema from fields that Kiro API doesn't accept.
+
+ Kiro API returns 400 "Improperly formed request" error if:
+ - required is an empty array []
+ - additionalProperties is present in schema
+
+ This function recursively processes the schema and removes problematic fields.
+
+ Args:
+ schema: JSON Schema to sanitize
+
+ Returns:
+ Sanitized copy of schema
+ """
+ if not schema:
+ return {}
+
+ result = {}
+
+ for key, value in schema.items():
+ # Skip empty required arrays
+ if key == "required" and isinstance(value, list) and len(value) == 0:
+ continue
+
+ # Skip additionalProperties - Kiro API doesn't support it
+ if key == "additionalProperties":
+ continue
+
+ # Recursively process nested objects
+ if key == "properties" and isinstance(value, dict):
+ result[key] = {
+ prop_name: sanitize_json_schema(prop_value) if isinstance(prop_value, dict) else prop_value
+ for prop_name, prop_value in value.items()
+ }
+ elif isinstance(value, dict):
+ result[key] = sanitize_json_schema(value)
+ elif isinstance(value, list):
+ # Process lists (e.g., anyOf, oneOf)
+ result[key] = [
+ sanitize_json_schema(item) if isinstance(item, dict) else item
+ for item in value
+ ]
+ else:
+ result[key] = value
+
+ return result
+
+
+# ==================================================================================================
+# Tool Processing
+# ==================================================================================================
+
+def process_tools_with_long_descriptions(
+ tools: Optional[List[UnifiedTool]]
+) -> Tuple[Optional[List[UnifiedTool]], str]:
+ """
+ Processes tools with long descriptions.
+
+ Kiro API has a limit on description length in toolSpecification.
+ If description exceeds the limit, full description is moved to system prompt,
+ and a reference to documentation remains in the tool.
+
+ Args:
+ tools: List of tools in unified format
+
+ Returns:
+ Tuple of:
+ - List of tools with processed descriptions (or None if tools is empty)
+ - String with documentation to add to system prompt (empty if all descriptions are short)
+ """
+ if not tools:
+ return None, ""
+
+ # If limit is disabled (0), return tools unchanged
+ if TOOL_DESCRIPTION_MAX_LENGTH <= 0:
+ return tools, ""
+
+ tool_documentation_parts = []
+ processed_tools = []
+
+ for tool in tools:
+ description = tool.description or ""
+
+ if len(description) <= TOOL_DESCRIPTION_MAX_LENGTH:
+ # Description is short - leave as is
+ processed_tools.append(tool)
+ else:
+ # Description is too long - move to system prompt
+ logger.debug(
+ f"Tool '{tool.name}' has long description ({len(description)} chars > {TOOL_DESCRIPTION_MAX_LENGTH}), "
+ f"moving to system prompt"
+ )
+
+ # Create documentation for system prompt
+ tool_documentation_parts.append(f"## Tool: {tool.name}\n\n{description}")
+
+ # Create copy of tool with reference description
+ reference_description = f"[Full documentation in system prompt under '## Tool: {tool.name}']"
+
+ processed_tool = UnifiedTool(
+ name=tool.name,
+ description=reference_description,
+ input_schema=tool.input_schema
+ )
+ processed_tools.append(processed_tool)
+
+ # Form final documentation
+ tool_documentation = ""
+ if tool_documentation_parts:
+ tool_documentation = (
+ "\n\n---\n"
+ "# Tool Documentation\n"
+ "The following tools have detailed documentation that couldn't fit in the tool definition.\n\n"
+ + "\n\n---\n\n".join(tool_documentation_parts)
+ )
+
+ return processed_tools if processed_tools else None, tool_documentation
+
+
+def validate_tool_names(tools: Optional[List[UnifiedTool]]) -> None:
+ """
+ Validates tool names against Kiro API 64-character limit.
+
+ Logs WARNING for each problematic tool and raises ValueError
+ with complete list of violations.
+
+ Args:
+ tools: List of tools to validate
+
+ Raises:
+ ValueError: If any tool name exceeds 64 characters
+
+ Example:
+ >>> validate_tool_names([UnifiedTool(name="short_name", description="test")])
+ # No error
+ >>> validate_tool_names([UnifiedTool(name="a" * 70, description="test")])
+ # Raises ValueError with detailed message
+ """
+ if not tools:
+ return
+
+ problematic_tools = []
+ for tool in tools:
+ if len(tool.name) > 64:
+ problematic_tools.append((tool.name, len(tool.name)))
+
+ if problematic_tools:
+ # Build detailed error message for client (no logging here - routes will log)
+ tool_list = "\n".join([
+ f" - '{name}' ({length} characters)"
+ for name, length in problematic_tools
+ ])
+
+ raise ValueError(
+ f"Tool name(s) exceed Kiro API limit of 64 characters:\n"
+ f"{tool_list}\n\n"
+ f"Solution: Use shorter tool names (max 64 characters).\n"
+ f"Example: 'get_user_data' instead of 'get_authenticated_user_profile_data_with_extended_information_about_it'"
+ )
+
+
+def convert_tools_to_kiro_format(tools: Optional[List[UnifiedTool]]) -> List[Dict[str, Any]]:
+ """
+ Converts unified tools to Kiro API format.
+
+ Args:
+ tools: List of tools in unified format
+
+ Returns:
+ List of tools in Kiro toolSpecification format
+ """
+ if not tools:
+ return []
+
+ kiro_tools = []
+ for tool in tools:
+ # Sanitize parameters from fields that Kiro API doesn't accept
+ sanitized_params = sanitize_json_schema(tool.input_schema)
+
+ # Kiro API requires non-empty description
+ description = tool.description
+ if not description or not description.strip():
+ description = f"Tool: {tool.name}"
+ logger.debug(f"Tool '{tool.name}' has empty description, using placeholder")
+
+ kiro_tools.append({
+ "toolSpecification": {
+ "name": tool.name,
+ "description": description,
+ "inputSchema": {"json": sanitized_params}
+ }
+ })
+
+ return kiro_tools
+
+
+# ==================================================================================================
+# Image Conversion to Kiro Format
+# ==================================================================================================
+
+def convert_images_to_kiro_format(images: Optional[List[Dict[str, Any]]]) -> List[Dict[str, Any]]:
+ """
+ Converts unified images to Kiro API format.
+
+ Unified format: [{"media_type": "image/jpeg", "data": "base64..."}]
+ Kiro format: [{"format": "jpeg", "source": {"bytes": "base64..."}}]
+
+ IMPORTANT: Images must be placed directly in userInputMessage.images,
+ NOT in userInputMessageContext.images. This matches the native Kiro IDE format.
+
+ Also handles the case where data contains a full data URL (data:image/jpeg;base64,...)
+ by stripping the prefix and extracting pure base64.
+
+ Args:
+ images: List of images in unified format
+
+ Returns:
+ List of images in Kiro format, ready for userInputMessage.images
+
+ Example:
+ >>> convert_images_to_kiro_format([{"media_type": "image/png", "data": "abc123"}])
+ [{'format': 'png', 'source': {'bytes': 'abc123'}}]
+ """
+ if not images:
+ return []
+
+ kiro_images = []
+ for img in images:
+ media_type = img.get("media_type", "image/jpeg")
+ data = img.get("data", "")
+
+ if not data:
+ logger.warning("Skipping image with empty data")
+ continue
+
+ # Strip data URL prefix if present (some clients send "data:image/jpeg;base64,..." in data field)
+ # Kiro API expects pure base64 without the prefix
+ if data.startswith("data:"):
+ try:
+ header, actual_data = data.split(",", 1)
+ # Extract media type from header if present
+ media_part = header.split(";")[0] # "data:image/jpeg"
+ extracted_media_type = media_part.replace("data:", "")
+ if extracted_media_type:
+ media_type = extracted_media_type
+ data = actual_data
+ logger.debug(f"Stripped data URL prefix, extracted media_type: {media_type}")
+ except (ValueError, IndexError) as e:
+ logger.warning(f"Failed to parse data URL prefix: {e}")
+
+ # Extract format from media_type: "image/jpeg" -> "jpeg"
+ format_str = media_type.split("/")[-1] if "/" in media_type else media_type
+
+ kiro_images.append({
+ "format": format_str,
+ "source": {
+ "bytes": data
+ }
+ })
+
+ if kiro_images:
+ logger.debug(f"Converted {len(kiro_images)} image(s) to Kiro format")
+
+ return kiro_images
+
+
+# ==================================================================================================
+# Tool Results and Tool Uses Extraction
+# ==================================================================================================
+
+def convert_tool_results_to_kiro_format(tool_results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Converts unified tool results to Kiro API format.
+
+ Unified format: {"type": "tool_result", "tool_use_id": "...", "content": "..."}
+ Kiro format: {"content": [{"text": "..."}], "status": "success", "toolUseId": "..."}
+
+ Args:
+ tool_results: List of tool results in unified format
+
+ Returns:
+ List of tool results in Kiro format
+ """
+ kiro_results = []
+ for tr in tool_results:
+ content = tr.get("content", "")
+ if isinstance(content, str):
+ content_text = content
+ else:
+ content_text = extract_text_content(content)
+
+ # Ensure content is not empty - Kiro API requires non-empty content
+ if not content_text:
+ content_text = "(empty result)"
+
+ kiro_results.append({
+ "content": [{"text": content_text}],
+ "status": "success",
+ "toolUseId": tr.get("tool_use_id", "")
+ })
+
+ return kiro_results
+
+
+def extract_tool_results_from_content(content: Any) -> List[Dict[str, Any]]:
+ """
+ Extracts tool results from message content.
+
+ Looks for content blocks with type="tool_result" and converts them
+ to Kiro API format.
+
+ Args:
+ content: Message content (can be a list of content blocks)
+
+ Returns:
+ List of tool results in Kiro format
+ """
+ tool_results = []
+
+ if isinstance(content, list):
+ for item in content:
+ if isinstance(item, dict) and item.get("type") == "tool_result":
+ tool_results.append({
+ "content": [{"text": extract_text_content(item.get("content", "")) or "(empty result)"}],
+ "status": "success",
+ "toolUseId": item.get("tool_use_id", "")
+ })
+
+ return tool_results
+
+
+def extract_tool_uses_from_message(
+ content: Any,
+ tool_calls: Optional[List[Dict[str, Any]]] = None
+) -> List[Dict[str, Any]]:
+ """
+ Extracts tool uses from assistant message.
+
+ Looks for tool calls in both:
+ - tool_calls field (OpenAI format)
+ - content blocks with type="tool_use" (Anthropic format)
+
+ Args:
+ content: Message content
+ tool_calls: List of tool calls (OpenAI format)
+
+ Returns:
+ List of tool uses in Kiro format
+ """
+ tool_uses = []
+
+ # From tool_calls field (OpenAI format or unified format from Anthropic)
+ if tool_calls:
+ for tc in tool_calls:
+ if isinstance(tc, dict):
+ func = tc.get("function", {})
+ arguments = func.get("arguments", "{}")
+ # Handle both string (OpenAI) and dict (Anthropic unified) formats
+ if isinstance(arguments, str):
+ input_data = json.loads(arguments) if arguments else {}
+ else:
+ input_data = arguments if arguments else {}
+ tool_uses.append({
+ "name": func.get("name", ""),
+ "input": input_data,
+ "toolUseId": tc.get("id", "")
+ })
+
+ # From content blocks (Anthropic format)
+ if isinstance(content, list):
+ for item in content:
+ if isinstance(item, dict) and item.get("type") == "tool_use":
+ tool_uses.append({
+ "name": item.get("name", ""),
+ "input": item.get("input", {}),
+ "toolUseId": item.get("id", "")
+ })
+
+ return tool_uses
+
+
+# ==================================================================================================
+# Tool Content to Text Conversion (for stripping when no tools defined)
+# ==================================================================================================
+
+def tool_calls_to_text(tool_calls: List[Dict[str, Any]]) -> str:
+ """
+ Converts tool_calls to human-readable text representation.
+
+ This is used when stripping tool content from messages (when no tools are defined).
+ Instead of losing the context, we convert tool calls to text so the model
+ can still understand what happened in the conversation.
+
+ Args:
+ tool_calls: List of tool calls in unified format
+
+ Returns:
+ Text representation of tool calls
+
+ Example:
+ >>> tool_calls_to_text([{"id": "call_123", "function": {"name": "bash", "arguments": '{"command": "ls"}'}}])
+ '[Tool: bash] (call_123)\\n{"command": "ls"}'
+ """
+ if not tool_calls:
+ return ""
+
+ parts = []
+ for tc in tool_calls:
+ func = tc.get("function", {})
+ name = func.get("name", "unknown")
+ arguments = func.get("arguments", "{}")
+ tool_id = tc.get("id", "")
+
+ # Format: [Tool: name] (id)\narguments
+ if tool_id:
+ parts.append(f"[Tool: {name} ({tool_id})]\n{arguments}")
+ else:
+ parts.append(f"[Tool: {name}]\n{arguments}")
+
+ return "\n\n".join(parts)
+
+
+def tool_results_to_text(tool_results: List[Dict[str, Any]]) -> str:
+ """
+ Converts tool_results to human-readable text representation.
+
+ This is used when stripping tool content from messages (when no tools are defined).
+ Instead of losing the context, we convert tool results to text so the model
+ can still understand what happened in the conversation.
+
+ Args:
+ tool_results: List of tool results in unified format
+
+ Returns:
+ Text representation of tool results
+
+ Example:
+ >>> tool_results_to_text([{"tool_use_id": "call_123", "content": "file1.txt\\nfile2.txt"}])
+ '[Tool Result] (call_123)\\nfile1.txt\\nfile2.txt'
+ """
+ if not tool_results:
+ return ""
+
+ parts = []
+ for tr in tool_results:
+ content = tr.get("content", "")
+ tool_use_id = tr.get("tool_use_id", "")
+
+ if isinstance(content, str):
+ content_text = content
+ else:
+ content_text = extract_text_content(content)
+
+ # Use placeholder if content is empty
+ if not content_text:
+ content_text = "(empty result)"
+
+ # Format: [Tool Result] (id)\ncontent
+ if tool_use_id:
+ parts.append(f"[Tool Result ({tool_use_id})]\n{content_text}")
+ else:
+ parts.append(f"[Tool Result]\n{content_text}")
+
+ return "\n\n".join(parts)
+
+
+# ==================================================================================================
+# Message Merging
+# ==================================================================================================
+
+def strip_all_tool_content(messages: List[UnifiedMessage]) -> Tuple[List[UnifiedMessage], bool]:
+ """
+ Strips ALL tool-related content from messages, converting it to text representation.
+
+ This is used when no tools are defined in the request. Kiro API rejects
+ requests that have toolResults but no tools defined.
+
+ Instead of simply removing tool content, this function converts tool_calls
+ and tool_results to human-readable text, preserving the context for
+ summarization and other use cases.
+
+ Args:
+ messages: List of messages in unified format
+
+ Returns:
+ Tuple of:
+ - List of messages with tool content converted to text
+ - Boolean indicating whether any tool content was converted
+ """
+ if not messages:
+ return [], False
+
+ result = []
+ total_tool_calls_stripped = 0
+ total_tool_results_stripped = 0
+
+ for msg in messages:
+ # Check if this message has any tool content
+ has_tool_calls = bool(msg.tool_calls)
+ has_tool_results = bool(msg.tool_results)
+
+ if has_tool_calls or has_tool_results:
+ if has_tool_calls:
+ total_tool_calls_stripped += len(msg.tool_calls)
+ if has_tool_results:
+ total_tool_results_stripped += len(msg.tool_results)
+
+ # Start with existing text content
+ existing_content = extract_text_content(msg.content)
+ content_parts = []
+
+ if existing_content:
+ content_parts.append(existing_content)
+
+ # Convert tool_calls to text (for assistant messages)
+ if has_tool_calls:
+ tool_text = tool_calls_to_text(msg.tool_calls)
+ if tool_text:
+ content_parts.append(tool_text)
+
+ # Convert tool_results to text (for user messages)
+ if has_tool_results:
+ result_text = tool_results_to_text(msg.tool_results)
+ if result_text:
+ content_parts.append(result_text)
+
+ # Join all parts with double newline
+ content = "\n\n".join(content_parts) if content_parts else "(empty)"
+
+ # Create a copy of the message without tool content but with text representation
+ cleaned_msg = UnifiedMessage(
+ role=msg.role,
+ content=content,
+ tool_calls=None,
+ tool_results=None
+ )
+ result.append(cleaned_msg)
+ else:
+ result.append(msg)
+
+ had_tool_content = total_tool_calls_stripped > 0 or total_tool_results_stripped > 0
+
+ # Log summary once (DEBUG level - this is normal for clients like Cline/Roo/Cursor)
+ if had_tool_content:
+ logger.debug(
+ f"Converted tool content to text (no tools defined): "
+ f"{total_tool_calls_stripped} tool_calls, {total_tool_results_stripped} tool_results"
+ )
+
+ return result, had_tool_content
+
+
+def ensure_assistant_before_tool_results(messages: List[UnifiedMessage]) -> Tuple[List[UnifiedMessage], bool]:
+ """
+ Ensures that messages with tool_results have a preceding assistant message with tool_calls.
+
+ Kiro API requires that when toolResults are present, there must be a preceding
+ assistantResponseMessage with toolUses. Some clients (like Cline/Roo/Cursor) may send
+ truncated conversations where the assistant message is missing.
+
+ Since we don't know the original tool name and arguments when the assistant message
+ is missing, we cannot create a valid synthetic assistant message. Instead, we convert
+ the tool_results to text representation and append to the message content, preserving
+ the context for the model while avoiding Kiro API rejection.
+
+ Args:
+ messages: List of messages in unified format
+
+ Returns:
+ Tuple of:
+ - List of messages with orphaned tool_results converted to text
+ - Boolean indicating whether any tool_results were converted (used to skip thinking tag injection)
+ """
+ if not messages:
+ return [], False
+
+ result = []
+ converted_any_tool_results = False
+
+ for msg in messages:
+ # Check if this message has tool_results
+ if msg.tool_results:
+ # Check if the previous message is an assistant with tool_calls
+ has_preceding_assistant = (
+ result and
+ result[-1].role == "assistant" and
+ result[-1].tool_calls
+ )
+
+ if not has_preceding_assistant:
+ # We cannot create a valid synthetic assistant message because we don't know
+ # the original tool name and arguments. Kiro API validates tool names.
+ # Convert tool_results to text to preserve context for the model.
+ logger.debug(
+ f"Converting {len(msg.tool_results)} orphaned tool_results to text "
+ f"(no preceding assistant message with tool_calls). "
+ f"Tool IDs: {[tr.get('tool_use_id', 'unknown') for tr in msg.tool_results]}"
+ )
+
+ # Convert tool_results to text representation
+ tool_results_text = tool_results_to_text(msg.tool_results)
+
+ # Append to existing content
+ original_content = extract_text_content(msg.content) or ""
+ if original_content and tool_results_text:
+ new_content = f"{original_content}\n\n{tool_results_text}"
+ elif tool_results_text:
+ new_content = tool_results_text
+ else:
+ new_content = original_content
+
+ # Create a copy of the message with tool_results converted to text
+ cleaned_msg = UnifiedMessage(
+ role=msg.role,
+ content=new_content,
+ tool_calls=msg.tool_calls,
+ tool_results=None, # Remove orphaned tool_results (now in text)
+ images=msg.images
+ )
+ result.append(cleaned_msg)
+ converted_any_tool_results = True
+ continue
+
+ result.append(msg)
+
+ return result, converted_any_tool_results
+
+
+def merge_adjacent_messages(messages: List[UnifiedMessage]) -> List[UnifiedMessage]:
+ """
+ Merges adjacent messages with the same role.
+
+ Kiro API does not accept multiple consecutive messages from the same role.
+ This function merges such messages into one.
+
+ Args:
+ messages: List of messages in unified format
+
+ Returns:
+ List of messages with merged adjacent messages
+ """
+ if not messages:
+ return []
+
+ merged = []
+ # Statistics for summary logging
+ merge_counts = {"user": 0, "assistant": 0}
+ total_tool_calls_merged = 0
+ total_tool_results_merged = 0
+
+ for msg in messages:
+ if not merged:
+ merged.append(msg)
+ continue
+
+ last = merged[-1]
+ if msg.role == last.role:
+ # Merge content
+ if isinstance(last.content, list) and isinstance(msg.content, list):
+ last.content = last.content + msg.content
+ elif isinstance(last.content, list):
+ last.content = last.content + [{"type": "text", "text": extract_text_content(msg.content)}]
+ elif isinstance(msg.content, list):
+ last.content = [{"type": "text", "text": extract_text_content(last.content)}] + msg.content
+ else:
+ last_text = extract_text_content(last.content)
+ current_text = extract_text_content(msg.content)
+ last.content = f"{last_text}\n{current_text}"
+
+ # Merge tool_calls for assistant messages
+ if msg.role == "assistant" and msg.tool_calls:
+ if last.tool_calls is None:
+ last.tool_calls = []
+ last.tool_calls = list(last.tool_calls) + list(msg.tool_calls)
+ total_tool_calls_merged += len(msg.tool_calls)
+
+ # Merge tool_results for user messages
+ if msg.role == "user" and msg.tool_results:
+ if last.tool_results is None:
+ last.tool_results = []
+ last.tool_results = list(last.tool_results) + list(msg.tool_results)
+ total_tool_results_merged += len(msg.tool_results)
+
+ # Count merges by role
+ if msg.role in merge_counts:
+ merge_counts[msg.role] += 1
+ else:
+ merged.append(msg)
+
+ # Log summary if any merges occurred
+ total_merges = sum(merge_counts.values())
+ if total_merges > 0:
+ parts = []
+ for role, count in merge_counts.items():
+ if count > 0:
+ parts.append(f"{count} {role}")
+ merge_summary = ", ".join(parts)
+
+ extras = []
+ if total_tool_calls_merged > 0:
+ extras.append(f"{total_tool_calls_merged} tool_calls")
+ if total_tool_results_merged > 0:
+ extras.append(f"{total_tool_results_merged} tool_results")
+
+ if extras:
+ logger.debug(f"Merged {total_merges} adjacent messages ({merge_summary}), including {', '.join(extras)}")
+ else:
+ logger.debug(f"Merged {total_merges} adjacent messages ({merge_summary})")
+
+ return merged
+
+
+# ==================================================================================================
+# Kiro History Building
+# ==================================================================================================
+
+def build_kiro_history(messages: List[UnifiedMessage], model_id: str) -> List[Dict[str, Any]]:
+ """
+ Builds history array for Kiro API from unified messages.
+
+ Kiro API expects alternating userInputMessage and assistantResponseMessage.
+ This function converts unified format to Kiro format.
+
+ Args:
+ messages: List of messages in unified format
+ model_id: Internal Kiro model ID
+
+ Returns:
+ List of dictionaries for history field in Kiro API
+ """
+ history = []
+
+ for msg in messages:
+ if msg.role == "user":
+ content = extract_text_content(msg.content)
+
+ # Fallback for empty content - Kiro API requires non-empty content
+ if not content:
+ content = "(empty)"
+
+ user_input = {
+ "content": content,
+ "modelId": model_id,
+ "origin": "AI_EDITOR",
+ }
+
+ # Process images - extract from message or content
+ # IMPORTANT: images go directly into userInputMessage, NOT into userInputMessageContext
+ # This matches the native Kiro IDE format
+ images = msg.images or extract_images_from_content(msg.content)
+ if images:
+ kiro_images = convert_images_to_kiro_format(images)
+ if kiro_images:
+ user_input["images"] = kiro_images
+
+ # Build userInputMessageContext for tools and toolResults only
+ user_input_context: Dict[str, Any] = {}
+
+ # Process tool_results - convert to Kiro format if present
+ if msg.tool_results:
+ kiro_tool_results = convert_tool_results_to_kiro_format(msg.tool_results)
+ if kiro_tool_results:
+ user_input_context["toolResults"] = kiro_tool_results
+ else:
+ # Try to extract from content (already in Kiro format)
+ tool_results = extract_tool_results_from_content(msg.content)
+ if tool_results:
+ user_input_context["toolResults"] = tool_results
+
+ # Add context if not empty (contains toolResults only, not images)
+ if user_input_context:
+ user_input["userInputMessageContext"] = user_input_context
+
+ history.append({"userInputMessage": user_input})
+
+ elif msg.role == "assistant":
+ content = extract_text_content(msg.content)
+
+ # Fallback for empty content - Kiro API requires non-empty content
+ if not content:
+ content = "(empty)"
+
+ assistant_response = {"content": content}
+
+ # Process tool_calls
+ tool_uses = extract_tool_uses_from_message(msg.content, msg.tool_calls)
+ if tool_uses:
+ assistant_response["toolUses"] = tool_uses
+
+ history.append({"assistantResponseMessage": assistant_response})
+
+ return history
+
+
+# ==================================================================================================
+# Main Payload Building
+# ==================================================================================================
+
+def build_kiro_payload(
+ messages: List[UnifiedMessage],
+ system_prompt: str,
+ model_id: str,
+ tools: Optional[List[UnifiedTool]],
+ conversation_id: str,
+ profile_arn: str,
+ inject_thinking: bool = True
+) -> KiroPayloadResult:
+ """
+ Builds complete payload for Kiro API from unified data.
+
+ This is the main function that assembles the Kiro API payload from
+ API-agnostic unified message and tool formats.
+
+ Args:
+ messages: List of messages in unified format (without system messages)
+ system_prompt: Already extracted system prompt
+ model_id: Internal Kiro model ID
+ tools: List of tools in unified format (or None)
+ conversation_id: Unique conversation ID
+ profile_arn: AWS CodeWhisperer profile ARN
+ inject_thinking: Whether to inject thinking tags (default True)
+
+ Returns:
+ KiroPayloadResult with payload and tool documentation
+
+ Raises:
+ ValueError: If there are no messages to send
+ """
+ # Process tools with long descriptions
+ processed_tools, tool_documentation = process_tools_with_long_descriptions(tools)
+
+ # Validate tool names against Kiro API 64-character limit
+ validate_tool_names(processed_tools)
+
+ # Add tool documentation to system prompt if present
+ full_system_prompt = system_prompt
+ if tool_documentation:
+ full_system_prompt = full_system_prompt + tool_documentation if full_system_prompt else tool_documentation.strip()
+
+ # Add thinking mode legitimization to system prompt if enabled
+ thinking_system_addition = get_thinking_system_prompt_addition()
+ if thinking_system_addition:
+ full_system_prompt = full_system_prompt + thinking_system_addition if full_system_prompt else thinking_system_addition.strip()
+
+ # If no tools are defined, strip ALL tool-related content from messages
+ # Kiro API rejects requests with toolResults but no tools
+ if not tools:
+ messages_without_tools, had_tool_content = strip_all_tool_content(messages)
+ messages_with_assistants = messages_without_tools
+ converted_tool_results = had_tool_content
+ else:
+ # Ensure assistant messages exist before tool_results (Kiro API requirement)
+ # Also returns flag if any tool_results were converted (to skip thinking tag injection)
+ messages_with_assistants, converted_tool_results = ensure_assistant_before_tool_results(messages)
+
+ # Merge adjacent messages with the same role
+ merged_messages = merge_adjacent_messages(messages_with_assistants)
+
+ if not merged_messages:
+ raise ValueError("No messages to send")
+
+ # 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}"
+
+ history = build_kiro_history(history_messages, model_id)
+
+ # 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 "Continue"
+ if current_message.role == "assistant":
+ history.append({
+ "assistantResponseMessage": {
+ "content": current_content
+ }
+ })
+ current_content = "Continue"
+
+ # If content is empty - use "Continue"
+ if not current_content:
+ current_content = "Continue"
+
+ # Process images in current message - extract from message or content
+ # IMPORTANT: images go directly into userInputMessage, NOT into userInputMessageContext
+ # This matches the native Kiro IDE format
+ images = current_message.images or extract_images_from_content(current_message.content)
+ kiro_images = None
+ if images:
+ kiro_images = convert_images_to_kiro_format(images)
+ if kiro_images:
+ logger.debug(f"Added {len(kiro_images)} image(s) to current message")
+
+ # Build user_input_context for tools and toolResults only (NOT images)
+ user_input_context: Dict[str, Any] = {}
+
+ # Add tools if present
+ kiro_tools = convert_tools_to_kiro_format(processed_tools)
+ if kiro_tools:
+ user_input_context["tools"] = kiro_tools
+
+ # Process tool_results in current message - convert to Kiro format if present
+ if current_message.tool_results:
+ # Convert unified format to Kiro format
+ kiro_tool_results = convert_tool_results_to_kiro_format(current_message.tool_results)
+ if kiro_tool_results:
+ user_input_context["toolResults"] = kiro_tool_results
+ else:
+ # Try to extract from content (already in Kiro format)
+ tool_results = extract_tool_results_from_content(current_message.content)
+ if tool_results:
+ user_input_context["toolResults"] = tool_results
+
+ # Inject thinking tags if enabled (only for the current/last user message)
+ if inject_thinking and current_message.role == "user":
+ current_content = inject_thinking_tags(current_content)
+
+ # Build userInputMessage
+ user_input_message = {
+ "content": current_content,
+ "modelId": model_id,
+ "origin": "AI_EDITOR",
+ }
+
+ # Add images directly to userInputMessage (NOT to userInputMessageContext)
+ if kiro_images:
+ user_input_message["images"] = kiro_images
+
+ # Add user_input_context if present (contains tools and toolResults only)
+ if user_input_context:
+ user_input_message["userInputMessageContext"] = user_input_context
+
+ # Assemble final payload
+ payload = {
+ "conversationState": {
+ "chatTriggerType": "MANUAL",
+ "conversationId": conversation_id,
+ "currentMessage": {
+ "userInputMessage": user_input_message
+ }
+ }
+ }
+
+ # Add history only if not empty
+ if history:
+ payload["conversationState"]["history"] = history
+
+ # Add profileArn
+ if profile_arn:
+ payload["profileArn"] = profile_arn
+
+ return KiroPayloadResult(payload=payload, tool_documentation=tool_documentation)
\ No newline at end of file
diff --git a/kiro-gateway/kiro/converters_openai.py b/kiro-gateway/kiro/converters_openai.py
new file mode 100644
index 0000000000000000000000000000000000000000..37011a7a634f899b8abf97457fef4e101c3942f1
--- /dev/null
+++ b/kiro-gateway/kiro/converters_openai.py
@@ -0,0 +1,303 @@
+# -*- 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 format to Kiro format.
+
+This module is an adapter layer that converts OpenAI-specific formats
+to the unified format used by converters_core.py.
+
+Contains functions for:
+- Converting OpenAI messages to unified format
+- Converting OpenAI tools to unified format
+- Building Kiro payload from OpenAI requests
+"""
+
+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_openai import ChatMessage, ChatCompletionRequest, Tool
+
+# Import from core - reuse shared logic
+from kiro.converters_core import (
+ extract_text_content,
+ extract_images_from_content,
+ UnifiedMessage,
+ UnifiedTool,
+ build_kiro_payload as core_build_kiro_payload,
+)
+
+
+# ==================================================================================================
+# OpenAI-specific Message Processing
+# ==================================================================================================
+
+def _extract_tool_results_from_openai(content: Any) -> List[Dict[str, Any]]:
+ """
+ Extracts tool results from OpenAI message content.
+
+ Args:
+ content: Message content (can be a list with tool_result blocks)
+
+ Returns:
+ List of tool results in unified format for UnifiedMessage
+ """
+ tool_results = []
+
+ if isinstance(content, list):
+ for item in content:
+ if isinstance(item, dict) and item.get("type") == "tool_result":
+ tool_results.append({
+ "type": "tool_result",
+ "tool_use_id": item.get("tool_use_id", ""),
+ "content": extract_text_content(item.get("content", "")) or "(empty result)"
+ })
+
+ return tool_results
+
+
+def _extract_tool_calls_from_openai(msg: ChatMessage) -> List[Dict[str, Any]]:
+ """
+ Extracts tool calls from OpenAI assistant message.
+
+ Args:
+ msg: OpenAI ChatMessage
+
+ Returns:
+ List of tool calls in unified format
+ """
+ tool_calls = []
+
+ if msg.tool_calls:
+ for tc in msg.tool_calls:
+ if isinstance(tc, dict):
+ tool_calls.append({
+ "id": tc.get("id", ""),
+ "type": "function",
+ "function": {
+ "name": tc.get("function", {}).get("name", ""),
+ "arguments": tc.get("function", {}).get("arguments", "{}")
+ }
+ })
+
+ return tool_calls
+
+
+def convert_openai_messages_to_unified(messages: List[ChatMessage]) -> Tuple[str, List[UnifiedMessage]]:
+ """
+ Converts OpenAI messages to unified format.
+
+ Handles:
+ - System messages (extracted as system prompt)
+ - Tool messages (converted to user messages with tool_results)
+ - Tool calls in assistant messages
+
+ Args:
+ messages: List of OpenAI ChatMessage objects
+
+ Returns:
+ Tuple of (system_prompt, unified_messages)
+ """
+ # Extract system prompt
+ system_prompt = ""
+ non_system_messages = []
+
+ for msg in messages:
+ if msg.role == "system":
+ system_prompt += extract_text_content(msg.content) + "\n"
+ else:
+ non_system_messages.append(msg)
+
+ system_prompt = system_prompt.strip()
+
+ # Process tool messages - convert to user messages with tool_results
+ processed = []
+ pending_tool_results = []
+ total_tool_calls = 0
+ total_tool_results = 0
+ total_images = 0
+
+ for msg in non_system_messages:
+ if msg.role == "tool":
+ # Collect tool results
+ tool_result = {
+ "type": "tool_result",
+ "tool_use_id": msg.tool_call_id or "",
+ "content": extract_text_content(msg.content) or "(empty result)"
+ }
+ pending_tool_results.append(tool_result)
+ total_tool_results += 1
+ else:
+ # If there are accumulated tool results, create user message with them
+ if pending_tool_results:
+ unified_msg = UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=pending_tool_results.copy()
+ )
+ processed.append(unified_msg)
+ pending_tool_results.clear()
+
+ # Convert regular message
+ tool_calls = None
+ tool_results = None
+ images = None
+
+ if msg.role == "assistant":
+ tool_calls = _extract_tool_calls_from_openai(msg) or None
+ if tool_calls:
+ total_tool_calls += len(tool_calls)
+ elif msg.role == "user":
+ tool_results = _extract_tool_results_from_openai(msg.content) or None
+ if tool_results:
+ total_tool_results += len(tool_results)
+ # Extract images from user messages
+ images = extract_images_from_content(msg.content) or None
+ if images:
+ total_images += len(images)
+
+ unified_msg = UnifiedMessage(
+ role=msg.role,
+ content=extract_text_content(msg.content),
+ tool_calls=tool_calls,
+ tool_results=tool_results,
+ images=images
+ )
+ processed.append(unified_msg)
+
+ # If tool results remain at the end
+ if pending_tool_results:
+ unified_msg = UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=pending_tool_results.copy()
+ )
+ processed.append(unified_msg)
+
+ # Log summary if any tool content or images were found
+ if total_tool_calls > 0 or total_tool_results > 0 or total_images > 0:
+ logger.debug(
+ f"Converted {len(messages)} OpenAI messages: "
+ f"{total_tool_calls} tool_calls, {total_tool_results} tool_results, {total_images} images"
+ )
+
+ return system_prompt, processed
+
+
+def convert_openai_tools_to_unified(tools: Optional[List[Tool]]) -> Optional[List[UnifiedTool]]:
+ """
+ Converts OpenAI tools to unified format.
+
+ Supports two formats:
+ 1. Standard OpenAI format: {"type": "function", "function": {"name": "...", ...}}
+ 2. Flat format (Cursor-style): {"name": "...", "description": "...", "input_schema": {...}}
+
+ Args:
+ tools: List of OpenAI Tool objects
+
+ Returns:
+ List of UnifiedTool objects, or None if no tools
+ """
+ if not tools:
+ return None
+
+ unified_tools = []
+ for tool in tools:
+ if tool.type != "function":
+ continue
+
+ # Standard OpenAI format (function field) takes priority
+ if tool.function is not None:
+ unified_tools.append(UnifiedTool(
+ name=tool.function.name,
+ description=tool.function.description,
+ input_schema=tool.function.parameters
+ ))
+ # Flat format compatibility (Cursor-style)
+ elif tool.name is not None:
+ unified_tools.append(UnifiedTool(
+ name=tool.name,
+ description=tool.description,
+ input_schema=tool.input_schema
+ ))
+ # Skip invalid tools
+ else:
+ logger.warning(f"Skipping invalid tool: no function or name field found")
+ continue
+
+ return unified_tools if unified_tools else None
+
+
+# ==================================================================================================
+# Main Entry Point
+# ==================================================================================================
+
+def build_kiro_payload(
+ request_data: ChatCompletionRequest,
+ conversation_id: str,
+ profile_arn: str
+) -> dict:
+ """
+ Builds complete payload for Kiro API from OpenAI request.
+
+ This is the main entry point for OpenAI → Kiro conversion.
+ Uses the core build_kiro_payload function with OpenAI-specific adapters.
+
+ Args:
+ request_data: Request in OpenAI format
+ conversation_id: Unique conversation ID
+ profile_arn: AWS CodeWhisperer profile ARN
+
+ Returns:
+ Payload dictionary for POST request to Kiro API
+
+ Raises:
+ ValueError: If there are no messages to send
+ """
+ # Convert messages to unified format
+ system_prompt, unified_messages = convert_openai_messages_to_unified(request_data.messages)
+
+ # Convert tools to unified format
+ unified_tools = convert_openai_tools_to_unified(request_data.tools)
+
+ # 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_data.model, HIDDEN_MODELS)
+
+ 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)}"
+ )
+
+ # Use core function to build payload
+ 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,
+ inject_thinking=True
+ )
+
+ return result.payload
\ No newline at end of file
diff --git a/kiro-gateway/kiro/debug_logger.py b/kiro-gateway/kiro/debug_logger.py
new file mode 100644
index 0000000000000000000000000000000000000000..688020760708d7d6cae9167d39f302ab69edf335
--- /dev/null
+++ b/kiro-gateway/kiro/debug_logger.py
@@ -0,0 +1,403 @@
+# -*- 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 .
+
+"""
+Debug logging module for requests.
+
+Supports three modes (DEBUG_MODE):
+- off: logging disabled
+- errors: logs are saved only on errors (4xx, 5xx)
+- all: logs are overwritten on each request
+
+In "errors" mode, data is buffered in memory and flushed to files
+only when flush_on_error() is called.
+
+Also captures application logs (loguru) for each request and saves
+them to app_logs.txt file for debugging convenience.
+"""
+
+import io
+import json
+import shutil
+from pathlib import Path
+from typing import Optional
+from loguru import logger
+
+from kiro.config import DEBUG_MODE, DEBUG_DIR
+
+
+class DebugLogger:
+ """
+ Singleton for managing debug request logs.
+
+ Operating modes:
+ - off: does nothing
+ - errors: buffers data, flushes to files only on errors
+ - all: writes data immediately to files (as before)
+ """
+ _instance = None
+
+ def __new__(cls):
+ if cls._instance is None:
+ cls._instance = super(DebugLogger, cls).__new__(cls)
+ cls._instance._initialized = False
+ return cls._instance
+
+ def __init__(self):
+ if self._initialized:
+ return
+ self.debug_dir = Path(DEBUG_DIR)
+ self._initialized = True
+
+ # Buffers for "errors" mode
+ self._request_body_buffer: Optional[bytes] = None
+ self._kiro_request_body_buffer: Optional[bytes] = None
+ self._raw_chunks_buffer: bytearray = bytearray()
+ self._modified_chunks_buffer: bytearray = bytearray()
+
+ # Buffer for application logs (loguru)
+ self._app_logs_buffer: io.StringIO = io.StringIO()
+ self._loguru_sink_id: Optional[int] = None
+
+ def _is_enabled(self) -> bool:
+ """Checks if logging is enabled."""
+ return DEBUG_MODE in ("errors", "all")
+
+ def _is_immediate_write(self) -> bool:
+ """Checks if immediate file writing is needed (all mode)."""
+ return DEBUG_MODE == "all"
+
+ def _clear_buffers(self):
+ """Clears all buffers."""
+ self._request_body_buffer = None
+ self._kiro_request_body_buffer = None
+ self._raw_chunks_buffer.clear()
+ self._modified_chunks_buffer.clear()
+ self._clear_app_logs_buffer()
+
+ def _clear_app_logs_buffer(self):
+ """Clears the application logs buffer and removes sink."""
+ # Remove sink from loguru
+ if self._loguru_sink_id is not None:
+ try:
+ logger.remove(self._loguru_sink_id)
+ except ValueError:
+ # Sink already removed
+ pass
+ self._loguru_sink_id = None
+
+ # Clear buffer
+ self._app_logs_buffer = io.StringIO()
+
+ def _setup_app_logs_capture(self):
+ """
+ Sets up application log capture to buffer.
+
+ Adds a temporary sink to loguru that writes to StringIO buffer.
+ Captures ALL logs without filtering, as sink is active only
+ during processing of a specific request.
+ """
+ # Remove previous sink if exists
+ self._clear_app_logs_buffer()
+
+ # Add new sink to capture ALL logs
+ # Format: time | level | module:function:line | message
+ self._loguru_sink_id = logger.add(
+ self._app_logs_buffer,
+ format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} | {message}",
+ level="DEBUG", # Capture all levels from DEBUG and above
+ colorize=False, # No ANSI colors in file
+ # No filter - capture ALL logs during request processing
+ )
+
+ def prepare_new_request(self):
+ """
+ Prepares the logger for a new request.
+
+ In "all" mode: clears the logs folder.
+ In "errors" mode: clears buffers.
+ In both modes: sets up application log capture.
+ """
+ if not self._is_enabled():
+ return
+
+ # Clear buffers in any case
+ self._clear_buffers()
+
+ # Set up application log capture
+ self._setup_app_logs_capture()
+
+ if self._is_immediate_write():
+ # "all" mode - clear folder and recreate
+ 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.")
+ except Exception as e:
+ logger.error(f"[DebugLogger] Error preparing directory: {e}")
+
+ def log_request_body(self, body: bytes):
+ """
+ Saves the request body (from client, OpenAI format).
+
+ In "all" mode: writes immediately to file.
+ In "errors" mode: buffers.
+ """
+ if not self._is_enabled():
+ return
+
+ if self._is_immediate_write():
+ self._write_request_body_to_file(body)
+ else:
+ # "errors" mode - buffer
+ self._request_body_buffer = body
+
+ def log_kiro_request_body(self, body: bytes):
+ """
+ Saves the modified request body (to Kiro API).
+
+ In "all" mode: writes immediately to file.
+ In "errors" mode: buffers.
+ """
+ if not self._is_enabled():
+ return
+
+ if self._is_immediate_write():
+ self._write_kiro_request_body_to_file(body)
+ else:
+ # "errors" mode - buffer
+ self._kiro_request_body_buffer = body
+
+ def log_raw_chunk(self, chunk: bytes):
+ """
+ Appends raw response chunk (from provider).
+
+ In "all" mode: writes immediately to file.
+ In "errors" mode: buffers.
+ """
+ if not self._is_enabled():
+ return
+
+ if self._is_immediate_write():
+ self._append_raw_chunk_to_file(chunk)
+ else:
+ # "errors" mode - buffer
+ self._raw_chunks_buffer.extend(chunk)
+
+ def log_modified_chunk(self, chunk: bytes):
+ """
+ Appends modified chunk (to client).
+
+ In "all" mode: writes immediately to file.
+ In "errors" mode: buffers.
+ """
+ if not self._is_enabled():
+ return
+
+ if self._is_immediate_write():
+ self._append_modified_chunk_to_file(chunk)
+ else:
+ # "errors" mode - buffer
+ self._modified_chunks_buffer.extend(chunk)
+
+ def log_error_info(self, status_code: int, error_message: str = ""):
+ """
+ Writes error information to file.
+
+ Works in both modes (errors and all).
+ In "all" mode writes immediately to file.
+ In "errors" mode called from flush_on_error().
+
+ Args:
+ status_code: HTTP error status code
+ error_message: Error message (optional)
+ """
+ if not self._is_enabled():
+ return
+
+ try:
+ # Ensure directory exists
+ self.debug_dir.mkdir(parents=True, exist_ok=True)
+
+ error_info = {
+ "status_code": status_code,
+ "error_message": error_message
+ }
+ error_file = self.debug_dir / "error_info.json"
+ with open(error_file, "w", encoding="utf-8") as f:
+ json.dump(error_info, f, indent=2, ensure_ascii=False)
+
+ logger.debug(f"[DebugLogger] Error info saved (status={status_code})")
+ except Exception as e:
+ logger.error(f"[DebugLogger] Error writing error_info: {e}")
+
+ def flush_on_error(self, status_code: int, error_message: str = ""):
+ """
+ Flushes buffers to files on error.
+
+ In "errors" mode: flushes buffers and saves error_info.
+ In "all" mode: only saves error_info (data already written).
+
+ Args:
+ status_code: HTTP error status code
+ error_message: Error message (optional)
+ """
+ if not self._is_enabled():
+ return
+
+ # In "all" mode data is already written, add error_info and app logs
+ if self._is_immediate_write():
+ self.log_error_info(status_code, error_message)
+ self._write_app_logs_to_file()
+ self._clear_app_logs_buffer()
+ return
+
+ # Check if there's anything to flush
+ if not any([
+ self._request_body_buffer,
+ self._kiro_request_body_buffer,
+ self._raw_chunks_buffer,
+ self._modified_chunks_buffer
+ ]):
+ return
+
+ try:
+ # Create directory if not exists
+ if self.debug_dir.exists():
+ shutil.rmtree(self.debug_dir)
+ self.debug_dir.mkdir(parents=True, exist_ok=True)
+
+ # Flush buffers to files
+ if self._request_body_buffer:
+ self._write_request_body_to_file(self._request_body_buffer)
+
+ if self._kiro_request_body_buffer:
+ self._write_kiro_request_body_to_file(self._kiro_request_body_buffer)
+
+ if self._raw_chunks_buffer:
+ file_path = self.debug_dir / "response_stream_raw.txt"
+ with open(file_path, "wb") as f:
+ f.write(self._raw_chunks_buffer)
+
+ if self._modified_chunks_buffer:
+ file_path = self.debug_dir / "response_stream_modified.txt"
+ with open(file_path, "wb") as f:
+ f.write(self._modified_chunks_buffer)
+
+ # Save error information
+ self.log_error_info(status_code, error_message)
+
+ # Save application logs
+ self._write_app_logs_to_file()
+
+ logger.info(f"[DebugLogger] Error logs flushed to {self.debug_dir} (status={status_code})")
+
+ except Exception as e:
+ logger.error(f"[DebugLogger] Error flushing buffers: {e}")
+ finally:
+ # Clear buffers after flush
+ self._clear_buffers()
+
+ def discard_buffers(self):
+ """
+ Clears buffers without writing to files.
+
+ Called when request completed successfully in "errors" mode.
+ Also called in "all" mode to save logs of successful request.
+ """
+ if DEBUG_MODE == "errors":
+ self._clear_buffers()
+ elif DEBUG_MODE == "all":
+ # In "all" mode save logs even for successful requests
+ self._write_app_logs_to_file()
+ self._clear_app_logs_buffer()
+
+ # ==================== Private file writing methods ====================
+
+ 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)
+ 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)
+ 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)
+ 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)
+ except Exception:
+ pass
+
+ def _write_app_logs_to_file(self):
+ """Writes captured application logs to file."""
+ try:
+ # Get buffer contents
+ logs_content = self._app_logs_buffer.getvalue()
+
+ if not logs_content.strip():
+ return
+
+ # 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)
+
+ logger.debug(f"[DebugLogger] App logs saved to {file_path}")
+ except Exception as e:
+ # Don't log error via logger to avoid recursion
+ pass
+
+
+# Global instance
+debug_logger = DebugLogger()
\ No newline at end of file
diff --git a/kiro-gateway/kiro/debug_middleware.py b/kiro-gateway/kiro/debug_middleware.py
new file mode 100644
index 0000000000000000000000000000000000000000..f02778a5c369ac859ac86d5e32db36371b3ceb77
--- /dev/null
+++ b/kiro-gateway/kiro/debug_middleware.py
@@ -0,0 +1,116 @@
+# -*- 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 .
+
+"""
+Debug logging middleware for Kiro Gateway.
+
+This middleware initializes debug logging BEFORE Pydantic validation,
+which allows capturing validation errors (422) in debug logs.
+
+The middleware:
+1. Intercepts requests to API endpoints (/v1/chat/completions, /v1/messages)
+2. Calls prepare_new_request() to initialize buffers and loguru sink
+3. Reads and logs the raw request body
+4. Passes the request to the next handler
+
+Flush/discard operations are handled by:
+- Route handlers (for successful requests and Kiro API errors)
+- Exception handlers (for validation errors and other exceptions)
+"""
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+from loguru import logger
+
+from kiro.config import DEBUG_MODE
+
+
+# API endpoints that should have debug logging enabled
+# These are the main API endpoints that process user requests
+LOGGED_ENDPOINTS = frozenset({
+ "/v1/chat/completions", # OpenAI-compatible endpoint
+ "/v1/messages", # Anthropic-compatible endpoint
+})
+
+
+class DebugLoggerMiddleware(BaseHTTPMiddleware):
+ """
+ Middleware for initializing debug logging on API requests.
+
+ This middleware runs BEFORE Pydantic validation, which means it can
+ capture the raw request body even for requests that fail validation.
+
+ The middleware only activates for API endpoints defined in LOGGED_ENDPOINTS.
+ Health checks, documentation, and other endpoints are not logged.
+
+ Lifecycle:
+ - prepare_new_request(): Called here (before validation)
+ - log_request_body(): Called here (raw body from client)
+ - log_kiro_request_body(): Called in route handlers (transformed payload)
+ - flush_on_error() / discard_buffers(): Called in routes or exception handlers
+ """
+
+ async def dispatch(self, request: Request, call_next) -> Response:
+ """
+ Process the request and initialize debug logging if needed.
+
+ Args:
+ request: The incoming HTTP request
+ call_next: The next middleware or route handler
+
+ Returns:
+ The response from the next handler
+ """
+ # Skip logging for non-API endpoints (health, docs, etc.)
+ if request.url.path not in LOGGED_ENDPOINTS:
+ return await call_next(request)
+
+ # Skip if debug mode is disabled
+ if DEBUG_MODE == "off":
+ return await call_next(request)
+
+ # Import here to avoid circular imports and allow graceful degradation
+ try:
+ from kiro.debug_logger import debug_logger
+ except ImportError:
+ logger.warning("debug_logger not available, skipping debug logging")
+ return await call_next(request)
+
+ # Initialize debug logging for this request
+ # This sets up buffers and creates a loguru sink to capture app logs
+ debug_logger.prepare_new_request()
+
+ # Read and log the raw request body
+ # FastAPI caches the body after first read, so this is safe
+ try:
+ body = await request.body()
+ if body:
+ debug_logger.log_request_body(body)
+ except Exception as e:
+ logger.warning(f"Failed to read request body for debug logging: {e}")
+
+ # Continue to validation and route handler
+ # flush_on_error() or discard_buffers() will be called by:
+ # - Route handlers (for successful requests and Kiro API errors)
+ # - validation_exception_handler (for 422 validation errors)
+ # - Generic exception handlers (for other errors)
+ response = await call_next(request)
+
+ return response
diff --git a/kiro-gateway/kiro/exceptions.py b/kiro-gateway/kiro/exceptions.py
new file mode 100644
index 0000000000000000000000000000000000000000..96a88e8fce6ed7c624d1e84768e743912479d863
--- /dev/null
+++ b/kiro-gateway/kiro/exceptions.py
@@ -0,0 +1,106 @@
+# -*- 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 .
+
+"""
+Exception handlers for Kiro Gateway.
+
+Contains functions for handling validation errors and other exceptions
+in a JSON-serialization compatible format.
+"""
+
+from typing import Any, List, Dict
+
+from fastapi import Request
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+from loguru import logger
+
+
+def sanitize_validation_errors(errors: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Converts validation errors to JSON-serializable format.
+
+ Pydantic may include bytes objects in the 'input' field, which
+ are not JSON-serializable. This function converts them to strings.
+
+ Args:
+ errors: List of validation errors from Pydantic
+
+ Returns:
+ List of errors with bytes converted to strings
+ """
+ sanitized = []
+ for error in errors:
+ sanitized_error = {}
+ for key, value in error.items():
+ if isinstance(value, bytes):
+ # Convert bytes to string
+ sanitized_error[key] = value.decode("utf-8", errors="replace")
+ elif isinstance(value, (list, tuple)):
+ # Recursively process lists
+ sanitized_error[key] = [
+ v.decode("utf-8", errors="replace") if isinstance(v, bytes) else v
+ for v in value
+ ]
+ else:
+ sanitized_error[key] = value
+ sanitized.append(sanitized_error)
+ return sanitized
+
+
+async def validation_exception_handler(request: Request, exc: RequestValidationError) -> JSONResponse:
+ """
+ Pydantic validation error handler.
+
+ Logs error details and returns an informative response.
+ Correctly handles bytes objects in errors by converting them to strings.
+ Also flushes debug logs for validation errors when DEBUG_MODE is enabled.
+
+ Args:
+ request: FastAPI Request object
+ exc: Validation exception from Pydantic
+
+ Returns:
+ JSONResponse with error details and status 422
+ """
+ body = await request.body()
+ body_str = body.decode("utf-8", errors="replace")
+
+ # Sanitize errors for JSON serialization
+ sanitized_errors = sanitize_validation_errors(exc.errors())
+
+ logger.error(f"Validation error (422): {sanitized_errors}")
+ # Log body at DEBUG level to avoid cluttering console with potentially large payloads
+ # logger.debug(f"Request body: {body_str[:500]}...")
+
+ # Flush debug logs for validation errors
+ # This is called AFTER middleware has initialized debug logging,
+ # so all app logs during request processing will be captured
+ try:
+ from kiro.debug_logger import debug_logger
+ if debug_logger:
+ error_message = f"Validation error: {sanitized_errors}"
+ debug_logger.flush_on_error(422, error_message)
+ except ImportError:
+ pass # debug_logger not available
+
+ return JSONResponse(
+ status_code=422,
+ content={"detail": sanitized_errors, "body": body_str[:500]},
+ )
\ No newline at end of file
diff --git a/kiro-gateway/kiro/http_client.py b/kiro-gateway/kiro/http_client.py
new file mode 100644
index 0000000000000000000000000000000000000000..943fa4991d639c641afaeb3e93f8f79ee2ed7a77
--- /dev/null
+++ b/kiro-gateway/kiro/http_client.py
@@ -0,0 +1,326 @@
+# -*- 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 .
+
+"""
+HTTP client for Kiro API with retry logic support.
+
+Handles:
+- 403: automatic token refresh and retry
+- 429: exponential backoff
+- 5xx: exponential backoff
+- Timeouts: exponential backoff
+
+Supports both per-request clients and shared application-level client
+with connection pooling for better resource management.
+"""
+
+import asyncio
+from typing import Optional
+
+import httpx
+from fastapi import HTTPException
+from loguru import logger
+
+from kiro.config import MAX_RETRIES, BASE_RETRY_DELAY, FIRST_TOKEN_MAX_RETRIES, STREAMING_READ_TIMEOUT
+from kiro.auth import KiroAuthManager
+from kiro.utils import get_kiro_headers
+from kiro.network_errors import classify_network_error, get_short_error_message, NetworkErrorInfo
+
+
+class KiroHttpClient:
+ """
+ HTTP client for Kiro API with retry logic support.
+
+ Automatically handles errors and retries requests:
+ - 403: refreshes token and retries
+ - 429: waits with exponential backoff
+ - 5xx: waits with exponential backoff
+ - Timeouts: waits with exponential backoff
+
+ Supports two modes of operation:
+ 1. Per-request client: Creates and owns its own httpx.AsyncClient
+ 2. Shared client: Uses an application-level shared client (recommended)
+
+ Using a shared client reduces memory usage and enables connection pooling,
+ which is especially important for handling concurrent requests.
+
+ Attributes:
+ auth_manager: Authentication manager for obtaining tokens
+ client: httpx HTTP client (owned or shared)
+
+ Example:
+ >>> # Per-request client (legacy mode)
+ >>> client = KiroHttpClient(auth_manager)
+ >>> response = await client.request_with_retry(...)
+
+ >>> # Shared client (recommended)
+ >>> shared = httpx.AsyncClient(limits=httpx.Limits(...))
+ >>> client = KiroHttpClient(auth_manager, shared_client=shared)
+ >>> response = await client.request_with_retry(...)
+ """
+
+ def __init__(
+ self,
+ auth_manager: KiroAuthManager,
+ shared_client: Optional[httpx.AsyncClient] = None
+ ):
+ """
+ Initializes the HTTP client.
+
+ Args:
+ auth_manager: Authentication manager
+ shared_client: Optional shared httpx.AsyncClient for connection pooling.
+ If provided, this client will be used instead of creating
+ a new one. The shared client will NOT be closed by close().
+ """
+ self.auth_manager = auth_manager
+ self._shared_client = shared_client
+ self._owns_client = shared_client is None
+ self.client: Optional[httpx.AsyncClient] = shared_client
+
+ async def _get_client(self, stream: bool = False) -> httpx.AsyncClient:
+ """
+ Returns or creates an HTTP client with proper timeouts.
+
+ If a shared client was provided at initialization, it is returned as-is.
+ Otherwise, creates a new client with appropriate timeout configuration.
+
+ httpx timeouts:
+ - connect: TCP handshake (DNS + TCP SYN/ACK)
+ - read: waiting for data from server between chunks
+ - write: sending data to server
+ - pool: waiting for free connection from pool
+
+ IMPORTANT: FIRST_TOKEN_TIMEOUT is NOT used here!
+ It is applied in streaming_openai.py via asyncio.wait_for() to control
+ the wait time for the first token from the model (retry business logic).
+
+ Args:
+ stream: If True, uses STREAMING_READ_TIMEOUT for read (only for new clients)
+
+ Returns:
+ Active HTTP client
+ """
+ # If using shared client, return it directly
+ # Shared client should be pre-configured with appropriate timeouts
+ if self._shared_client is not None:
+ return self._shared_client
+
+ # Create new client if needed (per-request mode)
+ if self.client is None or self.client.is_closed:
+ if stream:
+ # For streaming:
+ # - connect: 30 sec (TCP connection, usually < 1 sec)
+ # - read: STREAMING_READ_TIMEOUT (300 sec) - model may "think" between chunks
+ # - write/pool: standard values
+ timeout_config = httpx.Timeout(
+ connect=30.0,
+ read=STREAMING_READ_TIMEOUT,
+ write=30.0,
+ pool=30.0
+ )
+ logger.debug(f"Creating streaming HTTP client (read_timeout={STREAMING_READ_TIMEOUT}s)")
+ else:
+ # For regular requests: single timeout of 300 sec
+ timeout_config = httpx.Timeout(timeout=300.0)
+ logger.debug("Creating non-streaming HTTP client (timeout=300s)")
+
+ self.client = httpx.AsyncClient(timeout=timeout_config, follow_redirects=True)
+ return self.client
+
+ async def close(self) -> None:
+ """
+ Closes the HTTP client if this instance owns it.
+
+ If using a shared client, this method does nothing - the shared client
+ should be closed by the application lifecycle manager.
+
+ Uses graceful exception handling to prevent errors during cleanup
+ from masking the original exception in finally blocks.
+ """
+ # Don't close shared clients - they're managed by the application
+ if not self._owns_client:
+ return
+
+ if self.client and not self.client.is_closed:
+ try:
+ await self.client.aclose()
+ except Exception as e:
+ # Log but don't propagate - we're in cleanup code
+ # Propagating here could mask the original exception
+ logger.warning(f"Error closing HTTP client: {e}")
+
+ async def request_with_retry(
+ self,
+ method: str,
+ url: str,
+ json_data: dict,
+ stream: bool = False
+ ) -> httpx.Response:
+ """
+ Executes an HTTP request with retry logic.
+
+ Automatically handles various error types:
+ - 403: refreshes token via auth_manager.force_refresh() and retries
+ - 429: waits with exponential backoff (1s, 2s, 4s)
+ - 5xx: waits with exponential backoff
+ - Timeouts: waits with exponential backoff
+
+ For streaming, STREAMING_READ_TIMEOUT is used for waiting between chunks.
+ First token timeout is controlled separately in streaming_openai.py via asyncio.wait_for().
+
+ Args:
+ method: HTTP method (GET, POST, etc.)
+ url: Request URL
+ json_data: Request body (JSON)
+ stream: Use streaming (default False)
+
+ Returns:
+ httpx.Response with successful response
+
+ Raises:
+ HTTPException: On failure after all attempts (502/504)
+ """
+ # Determine the number of retry attempts
+ # FIRST_TOKEN_TIMEOUT is used in streaming_openai.py, not here
+ max_retries = FIRST_TOKEN_MAX_RETRIES if stream else MAX_RETRIES
+
+ client = await self._get_client(stream=stream)
+ last_error = None
+ last_error_info: Optional[NetworkErrorInfo] = None
+
+ for attempt in range(max_retries):
+ try:
+ # Get current token
+ token = await self.auth_manager.get_access_token()
+ headers = get_kiro_headers(self.auth_manager, token)
+
+ if stream:
+ # Prevent CLOSE_WAIT connection leak (issue #38)
+ headers["Connection"] = "close"
+ req = client.build_request(method, url, json=json_data, headers=headers)
+ logger.debug("Sending request to Kiro API...")
+ response = await client.send(req, stream=True)
+ else:
+ logger.debug("Sending request to Kiro API...")
+ response = await client.request(method, url, json=json_data, headers=headers)
+
+ # Check status
+ if response.status_code == 200:
+ return response
+
+ # 403 - token expired, refresh and retry
+ if response.status_code == 403:
+ logger.warning(f"Received 403, refreshing token (attempt {attempt + 1}/{MAX_RETRIES})")
+ await self.auth_manager.force_refresh()
+ continue
+
+ # 429 - rate limit, wait and retry
+ if response.status_code == 429:
+ delay = BASE_RETRY_DELAY * (2 ** attempt)
+ logger.warning(f"Received 429, waiting {delay}s (attempt {attempt + 1}/{MAX_RETRIES})")
+ await asyncio.sleep(delay)
+ continue
+
+ # 5xx - server error, wait and retry
+ if 500 <= response.status_code < 600:
+ delay = BASE_RETRY_DELAY * (2 ** attempt)
+ logger.warning(f"Received {response.status_code}, waiting {delay}s (attempt {attempt + 1}/{MAX_RETRIES})")
+ await asyncio.sleep(delay)
+ continue
+
+ # Other errors - return as is
+ return response
+
+ except httpx.TimeoutException as e:
+ last_error = e
+
+ # Classify timeout error for user-friendly messaging
+ error_info = classify_network_error(e)
+ last_error_info = error_info
+
+ # Log with user-friendly message
+ short_msg = get_short_error_message(error_info)
+
+ if error_info.is_retryable and attempt < max_retries - 1:
+ delay = BASE_RETRY_DELAY * (2 ** attempt)
+ logger.warning(f"{short_msg} - waiting {delay}s (attempt {attempt + 1}/{max_retries})")
+ await asyncio.sleep(delay)
+ else:
+ logger.error(f"{short_msg} - no more retries (attempt {attempt + 1}/{max_retries})")
+ if not error_info.is_retryable:
+ break # Don't retry non-retryable errors
+
+ except httpx.RequestError as e:
+ last_error = e
+
+ # Classify the error for user-friendly messaging
+ error_info = classify_network_error(e)
+ last_error_info = error_info
+
+ # Log with user-friendly message
+ short_msg = get_short_error_message(error_info)
+
+ if error_info.is_retryable and attempt < max_retries - 1:
+ delay = BASE_RETRY_DELAY * (2 ** attempt)
+ logger.warning(f"{short_msg} - waiting {delay}s (attempt {attempt + 1}/{max_retries})")
+ await asyncio.sleep(delay)
+ else:
+ logger.error(f"{short_msg} - no more retries (attempt {attempt + 1}/{max_retries})")
+ if not error_info.is_retryable:
+ break # Don't retry non-retryable errors
+
+ # All attempts exhausted - provide detailed, user-friendly error message
+ if last_error_info:
+ # Use classified error information
+ error_message = last_error_info.user_message
+
+ # Add troubleshooting steps
+ if last_error_info.troubleshooting_steps:
+ error_message += "\n\nTroubleshooting:\n"
+ for i, step in enumerate(last_error_info.troubleshooting_steps, 1):
+ error_message += f"{i}. {step}\n"
+
+ # Add technical details for debugging
+ error_message += f"\nTechnical details: {last_error_info.technical_details}"
+
+ raise HTTPException(
+ status_code=last_error_info.suggested_http_code,
+ detail=error_message.strip()
+ )
+ else:
+ # Fallback if no error was captured (shouldn't happen)
+ if stream:
+ raise HTTPException(
+ status_code=504,
+ detail=f"Streaming failed after {max_retries} attempts. Unknown error."
+ )
+ else:
+ raise HTTPException(
+ status_code=502,
+ detail=f"Request failed after {max_retries} attempts. Unknown error."
+ )
+
+ async def __aenter__(self) -> "KiroHttpClient":
+ """Async context manager support."""
+ return self
+
+ async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+ """Closes the client when exiting context."""
+ await self.close()
\ No newline at end of file
diff --git a/kiro-gateway/kiro/model_resolver.py b/kiro-gateway/kiro/model_resolver.py
new file mode 100644
index 0000000000000000000000000000000000000000..828bed133e512042102609948e186975a117bd99
--- /dev/null
+++ b/kiro-gateway/kiro/model_resolver.py
@@ -0,0 +1,376 @@
+# -*- 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 .
+
+"""
+Dynamic Model Resolution System for Kiro Gateway.
+
+Implements a 4-layer resolution pipeline:
+1. Normalize Name - Convert client formats to Kiro format (dashes→dots, strip dates)
+2. Check Dynamic Cache - Models from /ListAvailableModels API
+3. Check Hidden Models - Manual config for undocumented models
+4. Pass-through - Unknown models sent to Kiro (let Kiro decide)
+
+Key Principle: We are a gateway, not a gatekeeper. Kiro API is the final arbiter.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+from loguru import logger
+
+if TYPE_CHECKING:
+ from kiro.cache import ModelInfoCache
+
+
+@dataclass(frozen=True)
+class ModelResolution:
+ """
+ Result of model resolution.
+
+ Attributes:
+ internal_id: ID to send to Kiro API
+ source: Resolution source - "cache", "hidden", or "passthrough"
+ original_request: What client originally sent
+ normalized: Model name after normalization
+ is_verified: True if found in cache/hidden, False if passthrough
+ """
+ internal_id: str
+ source: str
+ original_request: str
+ normalized: str
+ is_verified: bool
+
+
+def normalize_model_name(name: str) -> str:
+ """
+ Normalize client model name to Kiro format.
+
+ Transformations applied:
+ 1. claude-haiku-4-5 → claude-haiku-4.5 (dash to dot for minor version)
+ 2. claude-haiku-4-5-20251001 → claude-haiku-4.5 (strip date suffix)
+ 3. claude-haiku-4-5-latest → claude-haiku-4.5 (strip 'latest' suffix)
+ 4. claude-sonnet-4-20250514 → claude-sonnet-4 (strip date, no minor)
+ 5. claude-3-7-sonnet → claude-3.7-sonnet (legacy format normalization)
+ 6. claude-3-7-sonnet-20250219 → claude-3.7-sonnet (legacy + strip date)
+ 7. claude-4.5-opus-high → claude-opus-4.5 (inverted format with suffix)
+
+ Args:
+ name: External model name from client
+
+ Returns:
+ Normalized model name in Kiro format
+
+ Examples:
+ >>> normalize_model_name("claude-haiku-4-5-20251001")
+ 'claude-haiku-4.5'
+ >>> normalize_model_name("claude-sonnet-4-5")
+ 'claude-sonnet-4.5'
+ >>> normalize_model_name("claude-opus-4-5")
+ 'claude-opus-4.5'
+ >>> normalize_model_name("claude-sonnet-4")
+ 'claude-sonnet-4'
+ >>> normalize_model_name("claude-sonnet-4-20250514")
+ 'claude-sonnet-4'
+ >>> normalize_model_name("claude-3-7-sonnet")
+ 'claude-3.7-sonnet'
+ >>> normalize_model_name("claude-3-7-sonnet-20250219")
+ 'claude-3.7-sonnet'
+ >>> normalize_model_name("claude-4.5-opus-high")
+ 'claude-opus-4.5'
+ >>> normalize_model_name("claude-4.5-sonnet-low")
+ 'claude-sonnet-4.5'
+ >>> normalize_model_name("auto")
+ 'auto'
+ """
+ if not name:
+ return name
+
+ # Lowercase for consistent matching
+ name_lower = name.lower()
+
+ # Pattern 1: Standard format - claude-{family}-{major}-{minor}(-{suffix})?
+ # Matches: claude-haiku-4-5, claude-haiku-4-5-20251001, claude-haiku-4-5-latest
+ # Groups: (claude-haiku-4), (5), optional suffix
+ # IMPORTANT: Minor version is 1-2 digits only! 8-digit dates should NOT match here.
+ standard_pattern = r'^(claude-(?:haiku|sonnet|opus)-\d+)-(\d{1,2})(?:-(?:\d{8}|latest|\d+))?$'
+ match = re.match(standard_pattern, name_lower)
+ if match:
+ base = match.group(1) # claude-haiku-4
+ minor = match.group(2) # 5
+ return f"{base}.{minor}" # claude-haiku-4.5
+
+ # Pattern 2: Standard format without minor - claude-{family}-{major}(-{date})?
+ # Matches: claude-sonnet-4, claude-sonnet-4-20250514
+ # Groups: (claude-sonnet-4), optional date
+ no_minor_pattern = r'^(claude-(?:haiku|sonnet|opus)-\d+)(?:-\d{8})?$'
+ match = re.match(no_minor_pattern, name_lower)
+ if match:
+ return match.group(1) # claude-sonnet-4
+
+ # Pattern 3: Legacy format - claude-{major}-{minor}-{family}(-{suffix})?
+ # Matches: claude-3-7-sonnet, claude-3-7-sonnet-20250219
+ # Groups: (claude), (3), (7), (sonnet), optional suffix
+ legacy_pattern = r'^(claude)-(\d+)-(\d+)-(haiku|sonnet|opus)(?:-(?:\d{8}|latest|\d+))?$'
+ match = re.match(legacy_pattern, name_lower)
+ if match:
+ prefix = match.group(1) # claude
+ major = match.group(2) # 3
+ minor = match.group(3) # 7
+ family = match.group(4) # sonnet
+ return f"{prefix}-{major}.{minor}-{family}" # claude-3.7-sonnet
+
+ # Pattern 4: Already normalized with dot but has date suffix
+ # Matches: claude-haiku-4.5-20251001, claude-3.7-sonnet-20250219
+ dot_with_date_pattern = r'^(claude-(?:\d+\.\d+-)?(?:haiku|sonnet|opus)(?:-\d+\.\d+)?)-\d{8}$'
+ match = re.match(dot_with_date_pattern, name_lower)
+ if match:
+ return match.group(1)
+
+ # Pattern 5: Inverted format with suffix - claude-{major}.{minor}-{family}-{suffix}
+ # Matches: claude-4.5-opus-high, claude-4.5-sonnet-low, claude-4.5-opus-high-thinking
+ # Convert to: claude-{family}-{major}.{minor}
+ # Groups: (4), (5), (opus), any suffix
+ # NOTE: This pattern REQUIRES a suffix to avoid matching already-normalized formats like claude-3.7-sonnet
+ inverted_with_suffix_pattern = r'^claude-(\d+)\.(\d+)-(haiku|sonnet|opus)-(.+)$'
+ match = re.match(inverted_with_suffix_pattern, name_lower)
+ if match:
+ major = match.group(1) # 4
+ minor = match.group(2) # 5
+ family = match.group(3) # opus
+ return f"claude-{family}-{major}.{minor}" # claude-opus-4.5
+
+ # No transformation needed - return as-is (preserving original case for passthrough)
+ return name
+
+
+def get_model_id_for_kiro(model_name: str, hidden_models: Dict[str, str]) -> 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.
+
+ For hidden models (like claude-3.7-sonnet), returns the internal Kiro ID.
+ For regular models, returns the normalized name.
+
+ Args:
+ model_name: External model name from client
+ hidden_models: Dict mapping display names to internal Kiro IDs
+
+ Returns:
+ Model ID to send to Kiro API
+
+ Examples:
+ >>> get_model_id_for_kiro("claude-haiku-4-5-20251001", {})
+ 'claude-haiku-4.5'
+ >>> 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("claude-3-7-sonnet", {"claude-3.7-sonnet": "CLAUDE_3_7_SONNET_20250219_V1_0"})
+ 'CLAUDE_3_7_SONNET_20250219_V1_0'
+ """
+ normalized = normalize_model_name(model_name)
+ return hidden_models.get(normalized, normalized)
+
+
+def extract_model_family(model_name: str) -> Optional[str]:
+ """
+ Extract model family from model name.
+
+ Args:
+ model_name: Model name (normalized or not)
+
+ Returns:
+ Family name ('haiku', 'sonnet', 'opus') or None if not a Claude model
+
+ Examples:
+ >>> extract_model_family("claude-haiku-4.5")
+ 'haiku'
+ >>> extract_model_family("claude-sonnet-4-5")
+ 'sonnet'
+ >>> extract_model_family("claude-3.7-sonnet")
+ 'sonnet'
+ >>> extract_model_family("gpt-4")
+ None
+ """
+ family_match = re.search(r'(haiku|sonnet|opus)', model_name, re.IGNORECASE)
+ if family_match:
+ return family_match.group(1).lower()
+ return None
+
+
+class ModelResolver:
+ """
+ Dynamic model resolver with normalization and optimistic pass-through.
+
+ Key principle: We are a gateway, not a gatekeeper.
+ Kiro API is the final arbiter of what models exist.
+
+ Resolution layers:
+ 1. Normalize name (dashes→dots, strip dates)
+ 2. Check dynamic cache (from /ListAvailableModels)
+ 3. Check hidden models (manual config)
+ 4. Pass-through (let Kiro decide)
+
+ Attributes:
+ cache: ModelInfoCache instance for dynamic model lookup
+ hidden_models: Dict mapping display names to internal Kiro IDs
+
+ Example:
+ >>> resolver = ModelResolver(cache, hidden_models)
+ >>> resolution = resolver.resolve("claude-haiku-4-5-20251001")
+ >>> resolution.internal_id
+ 'claude-haiku-4.5'
+ >>> resolution.source
+ 'cache'
+ """
+
+ def __init__(
+ self,
+ cache: ModelInfoCache,
+ hidden_models: Optional[Dict[str, str]] = None
+ ):
+ """
+ Initialize the model resolver.
+
+ Args:
+ cache: ModelInfoCache instance for dynamic model lookup
+ hidden_models: Dict mapping display names to internal Kiro IDs.
+ Display names should use dot format (e.g., "claude-3.7-sonnet")
+ """
+ self.cache = cache
+ self.hidden_models = hidden_models or {}
+
+ def resolve(self, external_model: str) -> ModelResolution:
+ """
+ Resolve external model name to internal Kiro ID.
+
+ NEVER raises - always returns a resolution.
+ If model is not in cache/hidden, we pass it through to Kiro.
+ Kiro will be the final judge.
+
+ Args:
+ external_model: Model name from client request
+
+ Returns:
+ ModelResolution with internal ID and metadata
+ """
+ # Layer 1: Normalize name (dashes→dots, strip date)
+ normalized = normalize_model_name(external_model)
+
+ logger.debug(
+ f"Model resolution: '{external_model}' → normalized: '{normalized}'"
+ )
+
+ # Layer 2: Check dynamic cache (from /ListAvailableModels)
+ if self.cache.is_valid_model(normalized):
+ logger.debug(f"Model '{normalized}' found in dynamic cache")
+ return ModelResolution(
+ internal_id=normalized,
+ source="cache",
+ original_request=external_model,
+ normalized=normalized,
+ is_verified=True
+ )
+
+ # Layer 3: Check hidden models
+ if normalized in self.hidden_models:
+ internal_id = self.hidden_models[normalized]
+ logger.debug(
+ f"Model '{normalized}' found in hidden models → '{internal_id}'"
+ )
+ return ModelResolution(
+ internal_id=internal_id,
+ source="hidden",
+ original_request=external_model,
+ normalized=normalized,
+ is_verified=True
+ )
+
+ # Layer 4: Pass-through - let Kiro decide!
+ # We don't know all models, Kiro might have hidden ones
+ logger.info(
+ f"Model '{external_model}' (normalized: '{normalized}') not in cache, "
+ f"passing through to Kiro API"
+ )
+ return ModelResolution(
+ internal_id=normalized, # Send normalized name to Kiro
+ source="passthrough",
+ original_request=external_model,
+ normalized=normalized,
+ is_verified=False # Not verified locally, Kiro will judge
+ )
+
+ def get_available_models(self) -> List[str]:
+ """
+ Get list of all available model IDs for /v1/models endpoint.
+
+ Combines:
+ - Models from dynamic cache (Kiro API)
+ - Hidden models (manual config)
+
+ Returns:
+ List of model IDs in consistent format (with dots)
+ """
+ # Start with cache models
+ models = set(self.cache.get_all_model_ids())
+
+ # Add hidden model display names (they use dot format)
+ models.update(self.hidden_models.keys())
+
+ return sorted(models)
+
+ def get_models_by_family(self, family: str) -> List[str]:
+ """
+ Get available models filtered by family.
+
+ Used for error messages to suggest alternatives from the same family.
+
+ Args:
+ family: Model family ('haiku', 'sonnet', 'opus')
+
+ Returns:
+ List of model IDs from the specified family
+ """
+ all_models = self.get_available_models()
+ return [m for m in all_models if family.lower() in m.lower()]
+
+ def get_suggestions_for_model(self, model_name: str) -> List[str]:
+ """
+ Get available models from the SAME family for error message.
+
+ IMPORTANT: Never suggests models from different family!
+ Opus request → only Opus suggestions
+ Sonnet request → only Sonnet suggestions
+
+ Args:
+ model_name: The model that was requested but not found
+
+ Returns:
+ List of available models from the same family, or all models
+ if family cannot be determined
+ """
+ family = extract_model_family(model_name)
+ if family:
+ return self.get_models_by_family(family)
+
+ # If we can't determine family, return all models
+ return self.get_available_models()
diff --git a/kiro-gateway/kiro/models_anthropic.py b/kiro-gateway/kiro/models_anthropic.py
new file mode 100644
index 0000000000000000000000000000000000000000..193746c584fa64521bf2c1b44b704b98fc5ffd3a
--- /dev/null
+++ b/kiro-gateway/kiro/models_anthropic.py
@@ -0,0 +1,442 @@
+# -*- 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 Anthropic Messages API.
+
+Defines data schemas for requests and responses compatible with
+Anthropic's Messages API specification.
+
+Reference: https://docs.anthropic.com/en/api/messages
+"""
+
+import time
+from typing import Any, Dict, List, Literal, Optional, Union
+from pydantic import BaseModel, Field
+
+
+# ==================================================================================================
+# Content Block Models
+# ==================================================================================================
+
+class TextContentBlock(BaseModel):
+ """
+ Text content block in Anthropic format.
+
+ Used in both requests and responses for text content.
+ """
+ type: Literal["text"] = "text"
+ text: str
+
+
+class ThinkingContentBlock(BaseModel):
+ """
+ Thinking content block in Anthropic format.
+
+ Represents the model's reasoning/thinking process.
+ Used when extended thinking is enabled.
+
+ Attributes:
+ type: Always "thinking"
+ thinking: The thinking/reasoning content
+ signature: Cryptographic signature for verification (placeholder in our case)
+ """
+ type: Literal["thinking"] = "thinking"
+ thinking: str
+ signature: str = ""
+
+
+class ToolUseContentBlock(BaseModel):
+ """
+ Tool use content block in Anthropic format.
+
+ Represents a tool call made by the assistant.
+ """
+ type: Literal["tool_use"] = "tool_use"
+ id: str
+ name: str
+ input: Dict[str, Any]
+
+
+class ToolResultContentBlock(BaseModel):
+ """
+ Tool result content block in Anthropic format.
+
+ Represents the result of a tool call, sent by the user.
+ """
+ type: Literal["tool_result"] = "tool_result"
+ tool_use_id: str
+ content: Optional[Union[str, List["TextContentBlock"]]] = None
+ is_error: Optional[bool] = None
+
+
+# ==================================================================================================
+# Image Content Block Models
+# ==================================================================================================
+
+class Base64ImageSource(BaseModel):
+ """
+ Base64-encoded image source in Anthropic format.
+
+ Attributes:
+ type: Always "base64"
+ media_type: MIME type (e.g., "image/jpeg", "image/png", "image/gif", "image/webp")
+ data: Base64-encoded image data
+ """
+ type: Literal["base64"] = "base64"
+ media_type: str
+ data: str
+
+
+class URLImageSource(BaseModel):
+ """
+ URL-based image source in Anthropic format.
+
+ Note: URL images require fetching and converting to base64 for Kiro API.
+ Currently logged as warning and skipped.
+
+ Attributes:
+ type: Always "url"
+ url: HTTP(S) URL to the image
+ """
+ type: Literal["url"] = "url"
+ url: str
+
+
+class ImageContentBlock(BaseModel):
+ """
+ Image content block in Anthropic format.
+
+ Represents an image in a message. Supports both base64-encoded
+ images and URL references.
+
+ Attributes:
+ type: Always "image"
+ source: Image source (base64 or URL)
+ """
+ type: Literal["image"] = "image"
+ source: Union[Base64ImageSource, URLImageSource]
+
+
+# Union type for all content blocks (including images and thinking)
+ContentBlock = Union[
+ TextContentBlock,
+ ThinkingContentBlock,
+ ImageContentBlock,
+ ToolUseContentBlock,
+ ToolResultContentBlock,
+]
+
+
+# ==================================================================================================
+# Message Models
+# ==================================================================================================
+
+class AnthropicMessage(BaseModel):
+ """
+ Message in Anthropic format.
+
+ Attributes:
+ role: Message role (user or assistant)
+ content: Message content (string or list of content blocks)
+ """
+ role: Literal["user", "assistant"]
+ content: Union[str, List[ContentBlock]]
+
+ model_config = {"extra": "allow"}
+
+
+# ==================================================================================================
+# Tool Models
+# ==================================================================================================
+
+class AnthropicTool(BaseModel):
+ """
+ Tool definition in Anthropic format.
+
+ Attributes:
+ name: Tool name (must match pattern ^[a-zA-Z0-9_-]{1,64}$)
+ description: Tool description (optional but recommended)
+ input_schema: JSON Schema for tool parameters
+ """
+ name: str
+ description: Optional[str] = None
+ input_schema: Dict[str, Any]
+
+
+class ToolChoiceAuto(BaseModel):
+ """Auto tool choice - model decides whether to use tools."""
+ type: Literal["auto"] = "auto"
+
+
+class ToolChoiceAny(BaseModel):
+ """Any tool choice - model must use at least one tool."""
+ type: Literal["any"] = "any"
+
+
+class ToolChoiceTool(BaseModel):
+ """Specific tool choice - model must use the specified tool."""
+ type: Literal["tool"] = "tool"
+ name: str
+
+
+ToolChoice = Union[ToolChoiceAuto, ToolChoiceAny, ToolChoiceTool]
+
+
+# ==================================================================================================
+# Request Models
+# ==================================================================================================
+
+class SystemContentBlock(BaseModel):
+ """
+ System content block for prompt caching.
+
+ Anthropic API supports system as a list of content blocks
+ with optional cache_control for prompt caching.
+ """
+ type: Literal["text"] = "text"
+ text: str
+ cache_control: Optional[Dict[str, Any]] = None
+
+ model_config = {"extra": "allow"}
+
+
+# System can be a string or list of content blocks (for prompt caching)
+SystemPrompt = Union[str, List[SystemContentBlock], List[Dict[str, Any]]]
+
+
+class AnthropicMessagesRequest(BaseModel):
+ """
+ Request to Anthropic Messages API (/v1/messages).
+
+ Attributes:
+ model: Model ID (e.g., "claude-sonnet-4-5")
+ messages: List of conversation messages
+ max_tokens: Maximum tokens in response (required)
+ system: System prompt (optional, string or list of content blocks for caching)
+ stream: Whether to stream the response
+ tools: List of available tools
+ tool_choice: Tool selection strategy
+ temperature: Sampling temperature (0-1)
+ top_p: Top-p sampling
+ top_k: Top-k sampling
+ stop_sequences: Custom stop sequences
+ metadata: Request metadata
+ """
+ model: str
+ messages: List[AnthropicMessage] = Field(min_length=1)
+ max_tokens: int
+
+ # Optional parameters - system can be string or list of content blocks
+ system: Optional[SystemPrompt] = None
+ stream: bool = False
+
+ # Tools
+ tools: Optional[List[AnthropicTool]] = None
+ tool_choice: Optional[Union[ToolChoice, Dict[str, Any]]] = None
+
+ # Sampling parameters
+ temperature: Optional[float] = Field(default=None, ge=0, le=1)
+ top_p: Optional[float] = Field(default=None, ge=0, le=1)
+ top_k: Optional[int] = Field(default=None, ge=0)
+
+ # Other parameters
+ stop_sequences: Optional[List[str]] = None
+ metadata: Optional[Dict[str, Any]] = None
+
+ model_config = {"extra": "allow"}
+
+
+# ==================================================================================================
+# Response Models
+# ==================================================================================================
+
+class AnthropicUsage(BaseModel):
+ """
+ Token usage information in Anthropic format.
+
+ Attributes:
+ input_tokens: Number of input tokens
+ output_tokens: Number of output tokens
+ """
+ input_tokens: int
+ output_tokens: int
+
+
+class AnthropicMessagesResponse(BaseModel):
+ """
+ Response from Anthropic Messages API (non-streaming).
+
+ Attributes:
+ id: Unique message ID
+ type: Always "message"
+ role: Always "assistant"
+ content: List of content blocks (may include thinking, text, tool_use)
+ model: Model used
+ stop_reason: Why generation stopped
+ stop_sequence: Stop sequence that triggered stop (if any)
+ usage: Token usage information
+ """
+ id: str
+ type: Literal["message"] = "message"
+ role: Literal["assistant"] = "assistant"
+ content: List[Union[ThinkingContentBlock, TextContentBlock, ToolUseContentBlock]]
+ model: str
+ stop_reason: Optional[Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"]] = None
+ stop_sequence: Optional[str] = None
+ usage: AnthropicUsage
+
+
+# ==================================================================================================
+# Streaming Event Models
+# ==================================================================================================
+
+class MessageStartEvent(BaseModel):
+ """
+ Event sent at the start of a message stream.
+
+ Contains the initial message object with empty content.
+ """
+ type: Literal["message_start"] = "message_start"
+ message: Dict[str, Any]
+
+
+class ContentBlockStartEvent(BaseModel):
+ """
+ Event sent at the start of a content block.
+
+ Attributes:
+ index: Index of the content block
+ content_block: Initial content block (with empty text for text blocks)
+ """
+ type: Literal["content_block_start"] = "content_block_start"
+ index: int
+ content_block: Dict[str, Any]
+
+
+class TextDelta(BaseModel):
+ """Delta for text content."""
+ type: Literal["text_delta"] = "text_delta"
+ text: str
+
+
+class ThinkingDelta(BaseModel):
+ """Delta for thinking content."""
+ type: Literal["thinking_delta"] = "thinking_delta"
+ thinking: str
+
+
+class InputJsonDelta(BaseModel):
+ """Delta for tool input JSON."""
+ type: Literal["input_json_delta"] = "input_json_delta"
+ partial_json: str
+
+
+class ContentBlockDeltaEvent(BaseModel):
+ """
+ Event sent when content block is updated.
+
+ Attributes:
+ index: Index of the content block being updated
+ delta: The delta update (text_delta, thinking_delta, or input_json_delta)
+ """
+ type: Literal["content_block_delta"] = "content_block_delta"
+ index: int
+ delta: Union[TextDelta, ThinkingDelta, InputJsonDelta, Dict[str, Any]]
+
+
+class ContentBlockStopEvent(BaseModel):
+ """
+ Event sent when a content block is complete.
+ """
+ type: Literal["content_block_stop"] = "content_block_stop"
+ index: int
+
+
+class MessageDeltaUsage(BaseModel):
+ """Usage information in message_delta event."""
+ output_tokens: int
+
+
+class MessageDeltaEvent(BaseModel):
+ """
+ Event sent near the end of the stream with final message data.
+
+ Attributes:
+ delta: Contains stop_reason and stop_sequence
+ usage: Output token count
+ """
+ type: Literal["message_delta"] = "message_delta"
+ delta: Dict[str, Any]
+ usage: MessageDeltaUsage
+
+
+class MessageStopEvent(BaseModel):
+ """
+ Event sent at the end of the message stream.
+ """
+ type: Literal["message_stop"] = "message_stop"
+
+
+class PingEvent(BaseModel):
+ """
+ Ping event sent periodically to keep connection alive.
+ """
+ type: Literal["ping"] = "ping"
+
+
+class ErrorEvent(BaseModel):
+ """
+ Error event sent when an error occurs during streaming.
+ """
+ type: Literal["error"] = "error"
+ error: Dict[str, Any]
+
+
+# Union of all streaming events
+StreamingEvent = Union[
+ MessageStartEvent,
+ ContentBlockStartEvent,
+ ContentBlockDeltaEvent,
+ ContentBlockStopEvent,
+ MessageDeltaEvent,
+ MessageStopEvent,
+ PingEvent,
+ ErrorEvent,
+]
+
+
+# ==================================================================================================
+# Error Models
+# ==================================================================================================
+
+class AnthropicErrorDetail(BaseModel):
+ """
+ Error detail in Anthropic format.
+ """
+ type: str
+ message: str
+
+
+class AnthropicErrorResponse(BaseModel):
+ """
+ Error response in Anthropic format.
+ """
+ type: Literal["error"] = "error"
+ error: AnthropicErrorDetail
\ No newline at end of file
diff --git a/kiro-gateway/kiro/models_openai.py b/kiro-gateway/kiro/models_openai.py
new file mode 100644
index 0000000000000000000000000000000000000000..46167c1b6e0bf8cdb5f89afb1f6abbe5c409c8d2
--- /dev/null
+++ b/kiro-gateway/kiro/models_openai.py
@@ -0,0 +1,282 @@
+# -*- 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-compatible API.
+
+Defines data schemas for requests and responses,
+providing validation and serialization.
+"""
+
+import time
+from typing import Any, Dict, List, Optional, Union
+from typing_extensions import Annotated
+from pydantic import BaseModel, Field
+
+
+# ==================================================================================================
+# Models for /v1/models endpoint
+# ==================================================================================================
+
+class OpenAIModel(BaseModel):
+ """
+ Data model for describing an AI model in OpenAI format.
+
+ Used in the /v1/models endpoint response.
+ """
+ id: str
+ object: str = "model"
+ created: int = Field(default_factory=lambda: int(time.time()))
+ owned_by: str = "anthropic"
+ description: Optional[str] = None
+
+
+class ModelList(BaseModel):
+ """
+ List of models in OpenAI format.
+
+ Response of GET /v1/models endpoint.
+ """
+ object: str = "list"
+ data: List[OpenAIModel]
+
+
+# ==================================================================================================
+# Models for /v1/chat/completions endpoint
+# ==================================================================================================
+
+class ChatMessage(BaseModel):
+ """
+ Chat message in OpenAI format.
+
+ Supports various roles (user, assistant, system, tool)
+ and various content formats (string, list, object).
+
+ Attributes:
+ role: Sender role (user, assistant, system, tool)
+ content: Message content (can be string, list, or None)
+ name: Optional sender name
+ tool_calls: List of tool calls (for assistant)
+ tool_call_id: Tool call ID (for tool)
+ """
+ role: str
+ content: Optional[Union[str, List[Any], Any]] = None
+ name: Optional[str] = None
+ tool_calls: Optional[List[Any]] = None
+ tool_call_id: Optional[str] = None
+
+ model_config = {"extra": "allow"}
+
+
+class ToolFunction(BaseModel):
+ """
+ Tool function description.
+
+ Attributes:
+ name: Function name
+ description: Function description
+ parameters: JSON Schema of function parameters
+ """
+ name: str
+ description: Optional[str] = None
+ parameters: Optional[Dict[str, Any]] = None
+
+
+class Tool(BaseModel):
+ """
+ Tool in OpenAI format.
+
+ Supports two formats:
+ 1. Standard OpenAI format: {"type": "function", "function": {...}}
+ 2. Flat format (Cursor-style): {"name": "...", "description": "...", "input_schema": {...}}
+
+ Attributes:
+ type: Tool type (usually "function")
+ function: Function description (standard format)
+ name: Function name (flat format)
+ description: Function description (flat format)
+ input_schema: Function parameters (flat format)
+ """
+ # Standard OpenAI format fields
+ type: str = "function"
+ function: Optional[ToolFunction] = None
+
+ # Flat format fields (Cursor-style)
+ name: Optional[str] = None
+ description: Optional[str] = None
+ input_schema: Optional[Dict[str, Any]] = None
+
+ model_config = {"extra": "allow"}
+
+
+class ChatCompletionRequest(BaseModel):
+ """
+ Request for response generation in OpenAI Chat Completions API format.
+
+ Supports all standard OpenAI API fields, including:
+ - Basic parameters (model, messages, stream)
+ - Generation parameters (temperature, top_p, max_tokens)
+ - Tools (function calling)
+ - Additional parameters (ignored but accepted for compatibility)
+
+ Attributes:
+ model: Model ID for generation
+ messages: List of chat messages
+ stream: Use streaming (default False)
+ temperature: Generation temperature (0-2)
+ top_p: Top-p sampling
+ n: Number of response variants
+ max_tokens: Maximum number of tokens in response
+ max_completion_tokens: Alternative field for max_tokens
+ stop: Stop sequences
+ presence_penalty: Penalty for topic repetition
+ frequency_penalty: Penalty for word repetition
+ tools: List of available tools
+ tool_choice: Tool selection strategy
+ """
+ model: str
+ messages: Annotated[List[ChatMessage], Field(min_length=1)]
+ stream: bool = False
+
+ # Generation parameters
+ temperature: Optional[float] = None
+ top_p: Optional[float] = None
+ n: Optional[int] = 1
+ max_tokens: Optional[int] = None
+ max_completion_tokens: Optional[int] = None
+ stop: Optional[Union[str, List[str]]] = None
+ presence_penalty: Optional[float] = None
+ frequency_penalty: Optional[float] = None
+
+ # Tools (function calling)
+ tools: Optional[List[Tool]] = None
+ tool_choice: Optional[Union[str, Dict]] = None
+
+ # Compatibility fields (ignored)
+ stream_options: Optional[Dict[str, Any]] = None
+ logit_bias: Optional[Dict[str, float]] = None
+ logprobs: Optional[bool] = None
+ top_logprobs: Optional[int] = None
+ user: Optional[str] = None
+ seed: Optional[int] = None
+ parallel_tool_calls: Optional[bool] = None
+
+ model_config = {"extra": "allow"}
+
+
+# ==================================================================================================
+# Models for responses
+# ==================================================================================================
+
+class ChatCompletionChoice(BaseModel):
+ """
+ Single response variant in Chat Completion.
+
+ Attributes:
+ index: Variant index
+ message: Response message
+ finish_reason: Completion reason (stop, tool_calls, length)
+ """
+ index: int = 0
+ message: Dict[str, Any]
+ finish_reason: Optional[str] = None
+
+
+class ChatCompletionUsage(BaseModel):
+ """
+ Token usage information.
+
+ Attributes:
+ prompt_tokens: Number of tokens in request
+ completion_tokens: Number of tokens in response
+ total_tokens: Total number of tokens
+ credits_used: Credits used (Kiro-specific)
+ """
+ prompt_tokens: int = 0
+ completion_tokens: int = 0
+ total_tokens: int = 0
+ credits_used: Optional[float] = None
+
+
+class ChatCompletionResponse(BaseModel):
+ """
+ Full Chat Completion response (non-streaming).
+
+ Attributes:
+ id: Unique response ID
+ object: Object type ("chat.completion")
+ created: Creation timestamp
+ model: Model used
+ choices: List of response variants
+ usage: Token usage information
+ """
+ id: str
+ object: str = "chat.completion"
+ created: int = Field(default_factory=lambda: int(time.time()))
+ model: str
+ choices: List[ChatCompletionChoice]
+ usage: ChatCompletionUsage
+
+
+class ChatCompletionChunkDelta(BaseModel):
+ """
+ Delta of changes in streaming chunk.
+
+ Attributes:
+ role: Role (only in first chunk)
+ content: New content
+ tool_calls: New tool calls
+ """
+ role: Optional[str] = None
+ content: Optional[str] = None
+ tool_calls: Optional[List[Dict[str, Any]]] = None
+
+
+class ChatCompletionChunkChoice(BaseModel):
+ """
+ Single variant in streaming chunk.
+
+ Attributes:
+ index: Variant index
+ delta: Delta of changes
+ finish_reason: Completion reason (only in last chunk)
+ """
+ index: int = 0
+ delta: ChatCompletionChunkDelta
+ finish_reason: Optional[str] = None
+
+
+class ChatCompletionChunk(BaseModel):
+ """
+ Streaming chunk in OpenAI format.
+
+ Attributes:
+ id: Unique response ID
+ object: Object type ("chat.completion.chunk")
+ created: Creation timestamp
+ model: Model used
+ choices: List of variants
+ usage: Usage information (only in last chunk)
+ """
+ id: str
+ object: str = "chat.completion.chunk"
+ created: int = Field(default_factory=lambda: int(time.time()))
+ model: str
+ choices: List[ChatCompletionChunkChoice]
+ usage: Optional[ChatCompletionUsage] = None
\ No newline at end of file
diff --git a/kiro-gateway/kiro/network_errors.py b/kiro-gateway/kiro/network_errors.py
new file mode 100644
index 0000000000000000000000000000000000000000..a32ae98cd83f2d44c1a5602a9ac383c4a1f86eff
--- /dev/null
+++ b/kiro-gateway/kiro/network_errors.py
@@ -0,0 +1,436 @@
+# -*- 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 .
+
+"""
+Network error classification and user-friendly message formatting.
+
+This module provides a centralized system for classifying network errors
+and converting them into actionable, user-friendly messages with troubleshooting steps.
+
+Architecture:
+- ErrorCategory: Enum of all possible network error types
+- NetworkErrorInfo: Structured information about an error
+- classify_network_error(): Analyzes exceptions and returns NetworkErrorInfo
+- format_error_for_user(): Formats errors for API responses (OpenAI/Anthropic)
+"""
+
+import socket
+from dataclasses import dataclass
+from enum import Enum
+from typing import List, Dict, Any, Optional
+
+import httpx
+from loguru import logger
+
+
+class ErrorCategory(str, Enum):
+ """
+ Categories of network errors.
+
+ Each category represents a distinct type of network failure
+ with specific troubleshooting steps.
+ """
+ DNS_RESOLUTION = "dns_resolution"
+ CONNECTION_REFUSED = "connection_refused"
+ CONNECTION_RESET = "connection_reset"
+ NETWORK_UNREACHABLE = "network_unreachable"
+ TIMEOUT_CONNECT = "timeout_connect"
+ TIMEOUT_READ = "timeout_read"
+ SSL_ERROR = "ssl_error"
+ PROXY_ERROR = "proxy_error"
+ TOO_MANY_REDIRECTS = "too_many_redirects"
+ UNKNOWN = "unknown"
+
+
+@dataclass
+class NetworkErrorInfo:
+ """
+ Structured information about a network error.
+
+ Attributes:
+ category: Error category for classification
+ user_message: Clear, non-technical message for end users
+ troubleshooting_steps: List of actionable steps to resolve the issue
+ technical_details: Technical error details for logging and debugging
+ is_retryable: Whether retrying the request might succeed
+ suggested_http_code: Appropriate HTTP status code (502, 504, etc.)
+ """
+ category: ErrorCategory
+ user_message: str
+ troubleshooting_steps: List[str]
+ technical_details: str
+ is_retryable: bool
+ suggested_http_code: int
+
+
+def classify_network_error(error: Exception) -> NetworkErrorInfo:
+ """
+ Classifies a network error and returns structured information.
+
+ Analyzes the exception type, error message, and underlying cause
+ to determine the specific type of network failure and provide
+ appropriate user-facing messages and troubleshooting steps.
+
+ Args:
+ error: The exception that occurred (typically httpx.RequestError)
+
+ Returns:
+ NetworkErrorInfo with classification and user-friendly details
+
+ Example:
+ >>> try:
+ ... response = await client.get("https://example.com")
+ ... except httpx.RequestError as e:
+ ... error_info = classify_network_error(e)
+ ... logger.error(f"[{error_info.category}] {error_info.user_message}")
+ """
+ error_type = type(error).__name__
+ error_str = str(error)
+
+ # Extract technical details for logging
+ technical_details = f"{error_type}: {error_str}"
+
+ # Analyze httpx.ConnectError (connection establishment failures)
+ if isinstance(error, httpx.ConnectError):
+ return _classify_connect_error(error, technical_details)
+
+ # Analyze httpx.TimeoutException (various timeout types)
+ if isinstance(error, httpx.TimeoutException):
+ return _classify_timeout_error(error, technical_details)
+
+ # Analyze httpx.TooManyRedirects
+ if isinstance(error, httpx.TooManyRedirects):
+ return NetworkErrorInfo(
+ category=ErrorCategory.TOO_MANY_REDIRECTS,
+ user_message="Too many redirects - the server is redirecting in a loop.",
+ troubleshooting_steps=[
+ "This is likely a server-side configuration issue",
+ "Try accessing the service directly without the gateway",
+ "Contact the service provider if the issue persists"
+ ],
+ technical_details=technical_details,
+ is_retryable=False,
+ suggested_http_code=502
+ )
+
+ # Analyze httpx.ProxyError
+ if isinstance(error, httpx.ProxyError):
+ return NetworkErrorInfo(
+ category=ErrorCategory.PROXY_ERROR,
+ user_message="Proxy connection failed - cannot connect through the configured proxy.",
+ troubleshooting_steps=[
+ "Check proxy configuration (HTTP_PROXY, HTTPS_PROXY environment variables)",
+ "Verify proxy server is accessible",
+ "Try disabling proxy temporarily",
+ "Check proxy authentication credentials if required"
+ ],
+ technical_details=technical_details,
+ is_retryable=True,
+ suggested_http_code=502
+ )
+
+ # Generic httpx.RequestError (catch-all)
+ if isinstance(error, httpx.RequestError):
+ return NetworkErrorInfo(
+ category=ErrorCategory.UNKNOWN,
+ user_message="Network request failed due to an unexpected error.",
+ troubleshooting_steps=[
+ "Check your internet connection",
+ "Verify firewall/antivirus settings",
+ "Try again in a few moments",
+ "Check the debug logs for more details"
+ ],
+ technical_details=technical_details,
+ is_retryable=True,
+ suggested_http_code=502
+ )
+
+ # Non-httpx errors (shouldn't happen, but handle gracefully)
+ return NetworkErrorInfo(
+ category=ErrorCategory.UNKNOWN,
+ user_message="An unexpected error occurred.",
+ troubleshooting_steps=[
+ "Check the debug logs for details",
+ "Try again in a few moments",
+ "Report this issue if it persists"
+ ],
+ technical_details=technical_details,
+ is_retryable=True,
+ suggested_http_code=500
+ )
+
+
+def _classify_connect_error(error: httpx.ConnectError, technical_details: str) -> NetworkErrorInfo:
+ """
+ Classifies httpx.ConnectError into specific subcategories.
+
+ Args:
+ error: The ConnectError exception
+ technical_details: Technical error string for logging
+
+ Returns:
+ NetworkErrorInfo with specific classification
+ """
+ error_str = str(error)
+
+ # Check underlying cause chain for more specific errors
+ cause = error.__cause__
+
+ # Check for DNS errors (socket.gaierror)
+ if cause and isinstance(cause, socket.gaierror):
+ # DNS resolution failed
+ # Common errno values:
+ # - 11001 (Windows): WSAHOST_NOT_FOUND
+ # - -2, -3, -5 (Unix): EAI_NONAME, EAI_AGAIN, EAI_NODATA
+ errno = getattr(cause, 'errno', None)
+
+ return NetworkErrorInfo(
+ category=ErrorCategory.DNS_RESOLUTION,
+ user_message="DNS resolution failed - cannot resolve the provider's domain name.",
+ troubleshooting_steps=[
+ "Check your internet connection",
+ "Try changing DNS servers to Google DNS (8.8.8.8, 8.8.4.4) or Cloudflare (1.1.1.1, 1.0.0.1)",
+ "Temporarily disable VPN if you're using one",
+ "Check if firewall/antivirus is blocking DNS requests",
+ "Verify the domain name is correct and the service is operational"
+ ],
+ technical_details=f"{technical_details} (errno: {errno})",
+ is_retryable=True,
+ suggested_http_code=502
+ )
+
+ # Check for connection refused
+ if "Connection refused" in error_str or "ECONNREFUSED" in error_str:
+ return NetworkErrorInfo(
+ category=ErrorCategory.CONNECTION_REFUSED,
+ user_message="Connection refused - the server is not accepting connections.",
+ troubleshooting_steps=[
+ "The service may be temporarily down",
+ "Check if the service is running and accessible",
+ "Verify firewall is not blocking the connection",
+ "Try again in a few moments"
+ ],
+ technical_details=technical_details,
+ is_retryable=True,
+ suggested_http_code=502
+ )
+
+ # Check for connection reset
+ if "Connection reset" in error_str or "ECONNRESET" in error_str:
+ return NetworkErrorInfo(
+ category=ErrorCategory.CONNECTION_RESET,
+ user_message="Connection reset - the server closed the connection unexpectedly.",
+ troubleshooting_steps=[
+ "This is usually a temporary server issue",
+ "Try again in a few moments",
+ "Check if VPN/proxy is interfering with the connection",
+ "Verify network stability"
+ ],
+ technical_details=technical_details,
+ is_retryable=True,
+ suggested_http_code=502
+ )
+
+ # Check for network unreachable
+ if "Network is unreachable" in error_str or "No route to host" in error_str or "ENETUNREACH" in error_str:
+ return NetworkErrorInfo(
+ category=ErrorCategory.NETWORK_UNREACHABLE,
+ user_message="Network unreachable - cannot reach the server's network.",
+ troubleshooting_steps=[
+ "Check your internet connection",
+ "Verify network adapter is enabled and working",
+ "Check routing table if using VPN",
+ "Try disabling VPN temporarily",
+ "Restart network adapter or router"
+ ],
+ technical_details=technical_details,
+ is_retryable=True,
+ suggested_http_code=502
+ )
+
+ # Check for SSL/TLS errors
+ if "SSL" in error_str or "TLS" in error_str or "certificate" in error_str.lower():
+ return NetworkErrorInfo(
+ category=ErrorCategory.SSL_ERROR,
+ user_message="SSL/TLS error - secure connection could not be established.",
+ troubleshooting_steps=[
+ "Check system date and time (incorrect time causes SSL errors)",
+ "Update SSL certificates on your system",
+ "Check if antivirus/firewall is intercepting HTTPS traffic",
+ "Verify the server's SSL certificate is valid"
+ ],
+ technical_details=technical_details,
+ is_retryable=False,
+ suggested_http_code=502
+ )
+
+ # Generic connection error
+ return NetworkErrorInfo(
+ category=ErrorCategory.UNKNOWN,
+ user_message="Connection failed - unable to establish connection to the server.",
+ troubleshooting_steps=[
+ "Check your internet connection",
+ "Verify firewall/antivirus settings",
+ "Try disabling VPN temporarily",
+ "Check if the service is accessible from other devices"
+ ],
+ technical_details=technical_details,
+ is_retryable=True,
+ suggested_http_code=502
+ )
+
+
+def _classify_timeout_error(error: httpx.TimeoutException, technical_details: str) -> NetworkErrorInfo:
+ """
+ Classifies httpx.TimeoutException into specific subcategories.
+
+ Args:
+ error: The TimeoutException
+ technical_details: Technical error string for logging
+
+ Returns:
+ NetworkErrorInfo with specific classification
+ """
+ # ConnectTimeout: TCP handshake timeout
+ if isinstance(error, httpx.ConnectTimeout):
+ return NetworkErrorInfo(
+ category=ErrorCategory.TIMEOUT_CONNECT,
+ user_message="Connection timeout - server did not respond to connection attempt.",
+ troubleshooting_steps=[
+ "Check your internet connection speed",
+ "The server may be overloaded or slow to respond",
+ "Try again in a few moments",
+ "Check if firewall is delaying connections"
+ ],
+ technical_details=technical_details,
+ is_retryable=True,
+ suggested_http_code=504
+ )
+
+ # ReadTimeout: Server stopped sending data
+ if isinstance(error, httpx.ReadTimeout):
+ return NetworkErrorInfo(
+ category=ErrorCategory.TIMEOUT_READ,
+ user_message="Read timeout - server stopped responding during data transfer.",
+ troubleshooting_steps=[
+ "The server may be processing a complex request",
+ "Check your internet connection stability",
+ "Try again with a simpler request",
+ "The service may be experiencing high load"
+ ],
+ technical_details=technical_details,
+ is_retryable=True,
+ suggested_http_code=504
+ )
+
+ # Generic timeout
+ return NetworkErrorInfo(
+ category=ErrorCategory.TIMEOUT_READ,
+ user_message="Request timeout - operation took too long to complete.",
+ troubleshooting_steps=[
+ "Check your internet connection",
+ "The server may be slow or overloaded",
+ "Try again in a few moments"
+ ],
+ technical_details=technical_details,
+ is_retryable=True,
+ suggested_http_code=504
+ )
+
+
+def format_error_for_user(
+ error_info: NetworkErrorInfo,
+ format_type: str = "openai",
+ include_troubleshooting: bool = True
+) -> Dict[str, Any]:
+ """
+ Formats NetworkErrorInfo for API response.
+
+ Converts structured error information into the appropriate format
+ for OpenAI or Anthropic API responses.
+
+ Args:
+ error_info: The classified error information
+ format_type: "openai" or "anthropic" format
+ include_troubleshooting: Whether to include troubleshooting steps
+
+ Returns:
+ Dictionary formatted for API response
+
+ Example:
+ >>> error_info = classify_network_error(exception)
+ >>> response = format_error_for_user(error_info, format_type="openai")
+ >>> return JSONResponse(status_code=502, content=response)
+ """
+ # Build the message
+ message = error_info.user_message
+
+ if include_troubleshooting and error_info.troubleshooting_steps:
+ message += "\n\nTroubleshooting steps:\n"
+ for i, step in enumerate(error_info.troubleshooting_steps, 1):
+ message += f"{i}. {step}\n"
+
+ # Format for OpenAI API
+ if format_type == "openai":
+ return {
+ "error": {
+ "message": message.strip(),
+ "type": "connectivity_error",
+ "code": error_info.category.value,
+ "param": None
+ }
+ }
+
+ # Format for Anthropic API
+ elif format_type == "anthropic":
+ return {
+ "type": "error",
+ "error": {
+ "type": "connectivity_error",
+ "message": message.strip()
+ }
+ }
+
+ # Generic format (fallback)
+ else:
+ return {
+ "error": {
+ "type": "connectivity_error",
+ "category": error_info.category.value,
+ "message": message.strip(),
+ "technical_details": error_info.technical_details
+ }
+ }
+
+
+def get_short_error_message(error_info: NetworkErrorInfo) -> str:
+ """
+ Returns a short, single-line error message for logging.
+
+ Args:
+ error_info: The classified error information
+
+ Returns:
+ Short error message suitable for log files
+
+ Example:
+ >>> error_info = classify_network_error(exception)
+ >>> logger.warning(get_short_error_message(error_info))
+ """
+ return error_info.user_message
diff --git a/kiro-gateway/kiro/parsers.py b/kiro-gateway/kiro/parsers.py
new file mode 100644
index 0000000000000000000000000000000000000000..661d248a1830ae9e2d894991e2f1b54b2d3496ec
--- /dev/null
+++ b/kiro-gateway/kiro/parsers.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 .
+
+"""
+Parsers for AWS Event Stream format.
+
+Contains classes and functions for:
+- Parsing binary AWS SSE stream
+- Extracting JSON events
+- Processing tool calls
+- Content deduplication
+"""
+
+import json
+import re
+from typing import Any, Dict, List, Optional
+
+from loguru import logger
+
+from kiro.utils import generate_tool_call_id
+
+
+def find_matching_brace(text: str, start_pos: int) -> int:
+ """
+ Finds the position of the closing brace considering nesting and strings.
+
+ Uses bracket counting for correct parsing of nested JSON.
+ Accounts for quoted strings and escape sequences.
+
+ Args:
+ text: Text to search
+ start_pos: Position of opening brace '{'
+
+ Returns:
+ Position of closing brace or -1 if not found
+
+ Example:
+ >>> find_matching_brace('{"a": {"b": 1}}', 0)
+ 14
+ >>> find_matching_brace('{"a": "{}"}', 0)
+ 10
+ """
+ if start_pos >= len(text) or text[start_pos] != '{':
+ return -1
+
+ brace_count = 0
+ in_string = False
+ escape_next = False
+
+ for i in range(start_pos, len(text)):
+ char = text[i]
+
+ if escape_next:
+ escape_next = False
+ continue
+
+ if char == '\\' and in_string:
+ escape_next = True
+ continue
+
+ if char == '"' and not escape_next:
+ in_string = not in_string
+ continue
+
+ if not in_string:
+ if char == '{':
+ brace_count += 1
+ elif char == '}':
+ brace_count -= 1
+ if brace_count == 0:
+ return i
+
+ return -1
+
+
+def parse_bracket_tool_calls(response_text: str) -> List[Dict[str, Any]]:
+ """
+ Parses tool calls in [Called func_name with args: {...}] format.
+
+ Some models return tool calls in text format instead of
+ structured JSON. This function extracts them.
+
+ Args:
+ response_text: Model response text
+
+ Returns:
+ List of tool calls in OpenAI format
+
+ Example:
+ >>> text = "[Called get_weather with args: {\"city\": \"London\"}]"
+ >>> calls = parse_bracket_tool_calls(text)
+ >>> calls[0]["function"]["name"]
+ 'get_weather'
+ """
+ if not response_text or "[Called" not in response_text:
+ return []
+
+ tool_calls = []
+ pattern = r'\[Called\s+(\w+)\s+with\s+args:\s*'
+
+ for match in re.finditer(pattern, response_text, re.IGNORECASE):
+ func_name = match.group(1)
+ args_start = match.end()
+
+ # Find JSON start
+ json_start = response_text.find('{', args_start)
+ if json_start == -1:
+ continue
+
+ # Find JSON end considering nesting
+ json_end = find_matching_brace(response_text, json_start)
+ if json_end == -1:
+ continue
+
+ json_str = response_text[json_start:json_end + 1]
+
+ try:
+ args = json.loads(json_str)
+ tool_call_id = generate_tool_call_id()
+ # index will be added later when forming the final response
+ tool_calls.append({
+ "id": tool_call_id,
+ "type": "function",
+ "function": {
+ "name": func_name,
+ "arguments": json.dumps(args)
+ }
+ })
+ except json.JSONDecodeError:
+ logger.warning(f"Failed to parse tool call arguments: {json_str[:100]}")
+
+ return tool_calls
+
+
+def deduplicate_tool_calls(tool_calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Removes duplicate tool calls.
+
+ Deduplication occurs by two criteria:
+ 1. By id - if there are multiple tool calls with the same id, keep the one with
+ more arguments (not empty "{}")
+ 2. By name+arguments - remove complete duplicates
+
+ Args:
+ tool_calls: List of tool calls
+
+ Returns:
+ List of unique tool calls
+ """
+ # First deduplicate by id - keep tool call with non-empty arguments
+ by_id: Dict[str, Dict[str, Any]] = {}
+ for tc in tool_calls:
+ tc_id = tc.get("id", "")
+ if not tc_id:
+ # Without id - add as is (will be deduplicated by name+args)
+ continue
+
+ existing = by_id.get(tc_id)
+ if existing is None:
+ by_id[tc_id] = tc
+ else:
+ # Duplicate by id exists - keep the one with more arguments
+ existing_args = existing.get("function", {}).get("arguments", "{}")
+ current_args = tc.get("function", {}).get("arguments", "{}")
+
+ # Prefer non-empty arguments
+ if current_args != "{}" and (existing_args == "{}" or len(current_args) > len(existing_args)):
+ logger.debug(f"Replacing tool call {tc_id} with better arguments: {len(existing_args)} -> {len(current_args)}")
+ by_id[tc_id] = tc
+
+ # Collect tool calls: first those with id, then without id
+ result_with_id = list(by_id.values())
+ result_without_id = [tc for tc in tool_calls if not tc.get("id")]
+
+ # Now deduplicate by name+arguments for all
+ seen = set()
+ unique = []
+
+ for tc in result_with_id + result_without_id:
+ # Protection against None in function
+ func = tc.get("function") or {}
+ func_name = func.get("name") or ""
+ func_args = func.get("arguments") or "{}"
+ key = f"{func_name}-{func_args}"
+ if key not in seen:
+ seen.add(key)
+ unique.append(tc)
+
+ if len(tool_calls) != len(unique):
+ logger.debug(f"Deduplicated tool calls: {len(tool_calls)} -> {len(unique)}")
+
+ return unique
+
+
+class AwsEventStreamParser:
+ """
+ Parser for AWS Event Stream format.
+
+ AWS returns events in binary format with :message-type...event delimiters.
+ This class extracts JSON events from the stream and converts them to a convenient format.
+
+ Supported event types:
+ - content: Text content of response
+ - tool_start: Start of tool call (name, toolUseId)
+ - tool_input: Continuation of input for tool call
+ - tool_stop: End of tool call
+ - usage: Credit consumption information
+ - context_usage: Context usage percentage
+
+ Attributes:
+ buffer: Buffer for accumulating data
+ last_content: Last processed content (for deduplication)
+ current_tool_call: Current incomplete tool call
+ tool_calls: List of completed tool calls
+
+ Example:
+ >>> parser = AwsEventStreamParser()
+ >>> events = parser.feed(chunk)
+ >>> for event in events:
+ ... if event["type"] == "content":
+ ... print(event["data"])
+ """
+
+ # Patterns for finding JSON events
+ EVENT_PATTERNS = [
+ ('{"content":', 'content'),
+ ('{"name":', 'tool_start'),
+ ('{"input":', 'tool_input'),
+ ('{"stop":', 'tool_stop'),
+ ('{"followupPrompt":', 'followup'),
+ ('{"usage":', 'usage'),
+ ('{"contextUsagePercentage":', 'context_usage'),
+ ]
+
+ def __init__(self):
+ """Initializes the parser."""
+ self.buffer = ""
+ 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]] = []
+
+ def feed(self, chunk: bytes) -> List[Dict[str, Any]]:
+ """
+ Adds chunk to buffer and returns parsed events.
+
+ Args:
+ chunk: Bytes of data from stream
+
+ Returns:
+ List of events in {"type": str, "data": Any} format
+ """
+ try:
+ self.buffer += chunk.decode('utf-8', errors='ignore')
+ except Exception:
+ return []
+
+ events = []
+
+ while True:
+ # Find nearest pattern
+ earliest_pos = -1
+ earliest_type = None
+
+ for pattern, event_type in self.EVENT_PATTERNS:
+ pos = self.buffer.find(pattern)
+ if pos != -1 and (earliest_pos == -1 or pos < earliest_pos):
+ earliest_pos = pos
+ earliest_type = event_type
+
+ if earliest_pos == -1:
+ break
+
+ # Find JSON end
+ json_end = find_matching_brace(self.buffer, earliest_pos)
+ if json_end == -1:
+ # JSON not complete, wait for more data
+ break
+
+ json_str = self.buffer[earliest_pos:json_end + 1]
+ self.buffer = self.buffer[json_end + 1:]
+
+ try:
+ data = json.loads(json_str)
+ event = self._process_event(data, earliest_type)
+ if event:
+ events.append(event)
+ except json.JSONDecodeError:
+ logger.warning(f"Failed to parse JSON: {json_str[:100]}")
+
+ return events
+
+ def _process_event(self, data: dict, event_type: str) -> Optional[Dict[str, Any]]:
+ """
+ Processes a parsed event.
+
+ Args:
+ data: Parsed JSON
+ event_type: Event type
+
+ Returns:
+ Processed event or None
+ """
+ if event_type == 'content':
+ return self._process_content_event(data)
+ elif event_type == 'tool_start':
+ return self._process_tool_start_event(data)
+ elif event_type == 'tool_input':
+ return self._process_tool_input_event(data)
+ elif event_type == 'tool_stop':
+ return self._process_tool_stop_event(data)
+ elif event_type == 'usage':
+ return {"type": "usage", "data": data.get('usage', 0)}
+ elif event_type == 'context_usage':
+ return {"type": "context_usage", "data": data.get('contextUsagePercentage', 0)}
+
+ return None
+
+ def _process_content_event(self, data: dict) -> Optional[Dict[str, Any]]:
+ """Processes content event."""
+ content = data.get('content', '')
+
+ # Skip followupPrompt
+ if data.get('followupPrompt'):
+ return None
+
+ # Deduplicate repeating content
+ if content == self.last_content:
+ return None
+
+ self.last_content = content
+
+ return {"type": "content", "data": content}
+
+ def _process_tool_start_event(self, data: dict) -> Optional[Dict[str, Any]]:
+ """Processes tool call start."""
+ # Finalize previous tool call if exists
+ if self.current_tool_call:
+ self._finalize_tool_call()
+
+ # input can be string or object
+ input_data = data.get('input', '')
+ if isinstance(input_data, dict):
+ input_str = json.dumps(input_data)
+ else:
+ input_str = str(input_data) if input_data else ''
+
+ self.current_tool_call = {
+ "id": data.get('toolUseId', generate_tool_call_id()),
+ "type": "function",
+ "function": {
+ "name": data.get('name', ''),
+ "arguments": input_str
+ }
+ }
+
+ if data.get('stop'):
+ self._finalize_tool_call()
+
+ return None
+
+ def _process_tool_input_event(self, data: dict) -> Optional[Dict[str, Any]]:
+ """Processes input continuation for tool call."""
+ if self.current_tool_call:
+ # input can be string or object
+ input_data = data.get('input', '')
+ if isinstance(input_data, dict):
+ input_str = json.dumps(input_data)
+ else:
+ input_str = str(input_data) if input_data else ''
+ self.current_tool_call['function']['arguments'] += input_str
+ return None
+
+ def _process_tool_stop_event(self, data: dict) -> Optional[Dict[str, Any]]:
+ """Processes tool call end."""
+ if self.current_tool_call and data.get('stop'):
+ self._finalize_tool_call()
+ return None
+
+ def _finalize_tool_call(self) -> None:
+ """Finalizes current tool call and adds to list."""
+ if not self.current_tool_call:
+ return
+
+ # Try to parse and normalize arguments as JSON
+ args = self.current_tool_call['function']['arguments']
+ tool_name = self.current_tool_call['function'].get('name', 'unknown')
+
+ logger.debug(f"Finalizing tool call '{tool_name}' with raw arguments: {repr(args)[:200]}")
+
+ if isinstance(args, str):
+ if args.strip():
+ try:
+ parsed = json.loads(args)
+ # Ensure result is a JSON string
+ self.current_tool_call['function']['arguments'] = json.dumps(parsed)
+ logger.debug(f"Tool '{tool_name}' arguments parsed successfully: {list(parsed.keys()) if isinstance(parsed, dict) else type(parsed)}")
+ except json.JSONDecodeError as e:
+ # Analyze the failure to provide better diagnostics
+ truncation_info = self._diagnose_json_truncation(args)
+
+ if truncation_info["is_truncated"]:
+ # This is likely an upstream issue - Kiro API truncated the stream
+ logger.warning(
+ f"Tool '{tool_name}' arguments appear truncated "
+ f"({truncation_info['size_bytes']} bytes received, {truncation_info['reason']}). "
+ f"This is NOT a Kiro Gateway bug — the stream was cut off before complete data arrived. "
+ f"Large tool call arguments (like writing big files) may trigger this limitation in Kiro API. "
+ f"Preview: {args[:100]}..."
+ )
+ else:
+ # Regular JSON parse error
+ logger.warning(f"Failed to parse tool '{tool_name}' arguments: {e}. Raw: {args[:200]}")
+
+ self.current_tool_call['function']['arguments'] = "{}"
+ else:
+ # Empty string - use empty object
+ # This is normal behavior for duplicate tool calls from Kiro
+ logger.debug(f"Tool '{tool_name}' has empty arguments string (will be deduplicated)")
+ self.current_tool_call['function']['arguments'] = "{}"
+ elif isinstance(args, dict):
+ # If already an object - serialize to string
+ self.current_tool_call['function']['arguments'] = json.dumps(args)
+ logger.debug(f"Tool '{tool_name}' arguments already dict with keys: {list(args.keys())}")
+ else:
+ # Unknown type - empty object
+ logger.warning(f"Tool '{tool_name}' has unexpected arguments type: {type(args)}")
+ self.current_tool_call['function']['arguments'] = "{}"
+
+ self.tool_calls.append(self.current_tool_call)
+ self.current_tool_call = None
+
+ def _diagnose_json_truncation(self, json_str: str) -> Dict[str, Any]:
+ """
+ Analyzes a malformed JSON string to determine if it was truncated.
+
+ This helps distinguish between upstream issues (Kiro API cutting off
+ large tool call arguments) and actual malformed JSON from the model.
+
+ Args:
+ json_str: The raw JSON string that failed to parse
+
+ Returns:
+ Dictionary with diagnostic information:
+ - is_truncated: True if the JSON appears to be cut off
+ - reason: Human-readable explanation of why it's truncated
+ - size_bytes: Size of the received data
+ """
+ size_bytes = len(json_str.encode('utf-8'))
+ stripped = json_str.strip()
+
+ # Check for obvious truncation signs
+ if not stripped:
+ return {"is_truncated": False, "reason": "empty string", "size_bytes": size_bytes}
+
+ # Count braces and brackets (simplified, doesn't account for strings perfectly)
+ open_braces = stripped.count('{')
+ close_braces = stripped.count('}')
+ open_brackets = stripped.count('[')
+ close_brackets = stripped.count(']')
+
+ # Check if JSON starts with { but doesn't end with }
+ if stripped.startswith('{') and not stripped.endswith('}'):
+ missing = open_braces - close_braces
+ return {
+ "is_truncated": True,
+ "reason": f"missing {missing} closing brace(s)",
+ "size_bytes": size_bytes
+ }
+
+ # Check if JSON starts with [ but doesn't end with ]
+ if stripped.startswith('[') and not stripped.endswith(']'):
+ missing = open_brackets - close_brackets
+ return {
+ "is_truncated": True,
+ "reason": f"missing {missing} closing bracket(s)",
+ "size_bytes": size_bytes
+ }
+
+ # Check for unbalanced braces/brackets
+ if open_braces != close_braces:
+ diff = open_braces - close_braces
+ return {
+ "is_truncated": True,
+ "reason": f"unbalanced braces ({open_braces} open, {close_braces} close)",
+ "size_bytes": size_bytes
+ }
+
+ if open_brackets != close_brackets:
+ diff = open_brackets - close_brackets
+ return {
+ "is_truncated": True,
+ "reason": f"unbalanced brackets ({open_brackets} open, {close_brackets} close)",
+ "size_bytes": size_bytes
+ }
+
+ # Check for unclosed string (ends with backslash or inside quotes)
+ # This is a heuristic - count unescaped quotes
+ quote_count = 0
+ i = 0
+ while i < len(stripped):
+ if stripped[i] == '\\' and i + 1 < len(stripped):
+ i += 2 # Skip escaped character
+ continue
+ if stripped[i] == '"':
+ quote_count += 1
+ i += 1
+
+ if quote_count % 2 != 0:
+ return {
+ "is_truncated": True,
+ "reason": "unclosed string literal",
+ "size_bytes": size_bytes
+ }
+
+ # Doesn't look truncated, probably just malformed
+ return {"is_truncated": False, "reason": "malformed JSON", "size_bytes": size_bytes}
+
+ def get_tool_calls(self) -> List[Dict[str, Any]]:
+ """
+ Returns all collected tool calls.
+
+ Finalizes current tool call if not finished.
+ Removes duplicates.
+
+ Returns:
+ List of unique tool calls
+ """
+ if self.current_tool_call:
+ self._finalize_tool_call()
+ return deduplicate_tool_calls(self.tool_calls)
+
+ def reset(self) -> None:
+ """Resets parser state."""
+ self.buffer = ""
+ self.last_content = None
+ self.current_tool_call = None
+ self.tool_calls = []
\ No newline at end of file
diff --git a/kiro-gateway/kiro/routes_anthropic.py b/kiro-gateway/kiro/routes_anthropic.py
new file mode 100644
index 0000000000000000000000000000000000000000..f6c235deeeec4b814ac8d624de207083d23d96b4
--- /dev/null
+++ b/kiro-gateway/kiro/routes_anthropic.py
@@ -0,0 +1,355 @@
+# -*- 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 routes for Anthropic Messages API.
+
+Contains the /v1/messages endpoint compatible with Anthropic's Messages API.
+
+Reference: https://docs.anthropic.com/en/api/messages
+"""
+
+import json
+from typing import Optional
+
+import httpx
+from fastapi import APIRouter, Depends, HTTPException, Request, Security, Header
+from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.security import APIKeyHeader
+from loguru import logger
+
+from kiro.config import PROXY_API_KEY
+from kiro.models_anthropic import (
+ AnthropicMessagesRequest,
+ AnthropicMessagesResponse,
+ AnthropicErrorResponse,
+ AnthropicErrorDetail,
+)
+from kiro.auth import KiroAuthManager, AuthType
+from kiro.cache import ModelInfoCache
+from kiro.converters_anthropic import anthropic_to_kiro
+from kiro.streaming_anthropic import (
+ stream_kiro_to_anthropic,
+ collect_anthropic_response,
+)
+from kiro.http_client import KiroHttpClient
+from kiro.utils import generate_conversation_id
+from kiro.tokenizer import count_tools_tokens
+
+# Import debug_logger
+try:
+ from kiro.debug_logger import debug_logger
+except ImportError:
+ debug_logger = None
+
+
+# --- Security scheme ---
+# Anthropic uses x-api-key header instead of Authorization: Bearer
+anthropic_api_key_header = APIKeyHeader(name="x-api-key", auto_error=False)
+# Also support Authorization: Bearer for compatibility
+auth_header = APIKeyHeader(name="Authorization", auto_error=False)
+
+
+async def verify_anthropic_api_key(
+ x_api_key: Optional[str] = Security(anthropic_api_key_header),
+ authorization: Optional[str] = Security(auth_header)
+) -> bool:
+ """
+ Verify API key for Anthropic API.
+
+ Supports two authentication methods:
+ 1. x-api-key header (Anthropic native)
+ 2. Authorization: Bearer header (for compatibility)
+
+ Args:
+ x_api_key: Value from x-api-key header
+ authorization: Value from Authorization header
+
+ Returns:
+ True if key is valid
+
+ Raises:
+ HTTPException: 401 if key is invalid or missing
+ """
+ # Check x-api-key first (Anthropic native)
+ if x_api_key and x_api_key == PROXY_API_KEY:
+ return True
+
+ # Fall back to Authorization: Bearer
+ if authorization and authorization == f"Bearer {PROXY_API_KEY}":
+ return True
+
+ logger.warning("Access attempt with invalid API key (Anthropic endpoint)")
+ raise HTTPException(
+ status_code=401,
+ detail={
+ "type": "error",
+ "error": {
+ "type": "authentication_error",
+ "message": "Invalid or missing API key. Use x-api-key header or Authorization: Bearer."
+ }
+ }
+ )
+
+
+# --- Router ---
+router = APIRouter(tags=["Anthropic API"])
+
+
+@router.post("/v1/messages", dependencies=[Depends(verify_anthropic_api_key)])
+async def messages(
+ request: Request,
+ request_data: AnthropicMessagesRequest,
+ anthropic_version: Optional[str] = Header(None, alias="anthropic-version")
+):
+ """
+ Anthropic Messages API endpoint.
+
+ Compatible with Anthropic's /v1/messages endpoint.
+ Accepts requests in Anthropic format and translates them to Kiro API.
+
+ Required headers:
+ - x-api-key: Your API key (or Authorization: Bearer)
+ - anthropic-version: API version (optional, for compatibility)
+ - Content-Type: application/json
+
+ Args:
+ request: FastAPI Request for accessing app.state
+ request_data: Request in Anthropic MessagesRequest format
+ anthropic_version: Anthropic API version header (optional)
+
+ Returns:
+ StreamingResponse for streaming mode (SSE)
+ JSONResponse for non-streaming mode
+
+ Raises:
+ HTTPException: On validation or API errors
+ """
+ logger.info(f"Request to /v1/messages (model={request_data.model}, stream={request_data.stream})")
+
+ if anthropic_version:
+ logger.debug(f"Anthropic-Version header: {anthropic_version}")
+
+ auth_manager: KiroAuthManager = request.app.state.auth_manager
+ model_cache: ModelInfoCache = request.app.state.model_cache
+
+ # Note: prepare_new_request() and log_request_body() are now called by DebugLoggerMiddleware
+ # This ensures debug logging works even for requests that fail Pydantic validation (422 errors)
+
+ # Generate conversation ID
+ conversation_id = generate_conversation_id()
+
+ # Build payload for Kiro
+ # profileArn is only needed for Kiro Desktop auth
+ profile_arn_for_payload = ""
+ if auth_manager.auth_type == AuthType.KIRO_DESKTOP and auth_manager.profile_arn:
+ profile_arn_for_payload = auth_manager.profile_arn
+
+ try:
+ kiro_payload = anthropic_to_kiro(
+ request_data,
+ conversation_id,
+ profile_arn_for_payload
+ )
+ except ValueError as e:
+ logger.error(f"Conversion error: {e}")
+ return JSONResponse(
+ status_code=400,
+ content={
+ "type": "error",
+ "error": {
+ "type": "invalid_request_error",
+ "message": str(e)
+ }
+ }
+ )
+
+ # Log Kiro payload
+ 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}")
+
+ # Create HTTP client with retry logic
+ # For streaming: use per-request client to avoid CLOSE_WAIT leak on VPN disconnect (issue #54)
+ # For non-streaming: use shared client for connection pooling
+ url = f"{auth_manager.api_host}/generateAssistantResponse"
+
+ if request_data.stream:
+ # Streaming mode: per-request client prevents orphaned connections
+ # 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)
+
+ # Prepare data for token counting
+ # Convert Pydantic models to dicts for tokenizer
+ messages_for_tokenizer = [msg.model_dump() for msg in request_data.messages]
+ tools_for_tokenizer = [tool.model_dump() for tool in request_data.tools] if request_data.tools else None
+
+ try:
+ # Make request to Kiro API (for both streaming and non-streaming modes)
+ # Important: we wait for Kiro response BEFORE returning StreamingResponse,
+ # so that we can return proper HTTP error codes if Kiro fails
+ 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')
+ logger.error(f"Error from Kiro API: {response.status_code} - {error_text}")
+
+ # Try to parse JSON response from Kiro to extract error message
+ error_message = error_text
+ try:
+ error_json = json.loads(error_text)
+ if "message" in error_json:
+ error_message = error_json["message"]
+ if "reason" in error_json:
+ error_message = f"{error_message} (reason: {error_json['reason']})"
+ except (json.JSONDecodeError, KeyError):
+ pass
+
+ # Log access log for error (before flush, so it gets into app_logs)
+ logger.warning(
+ f"HTTP {response.status_code} - POST /v1/messages - {error_message[:100]}"
+ )
+
+ # Flush debug logs on error
+ if debug_logger:
+ debug_logger.flush_on_error(response.status_code, error_message)
+
+ # Return error in Anthropic format
+ return JSONResponse(
+ status_code=response.status_code,
+ content={
+ "type": "error",
+ "error": {
+ "type": "api_error",
+ "message": error_message
+ }
+ }
+ )
+
+ if request_data.stream:
+ # Streaming mode - Kiro already returned 200, now stream the response
+ async def stream_wrapper():
+ streaming_error = None
+ client_disconnected = False
+ try:
+ async for chunk in stream_kiro_to_anthropic(
+ response,
+ request_data.model,
+ model_cache,
+ auth_manager,
+ request_messages=messages_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
+ # Send error event to client, then gracefully end the stream
+ try:
+ error_event = f'event: error\ndata: {json.dumps({"type": "error", "error": {"type": "api_error", "message": str(e)}})}\n\n'
+ yield error_event
+ except Exception:
+ pass
+ 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/messages (streaming) - [{error_type}] {error_msg[:100]}")
+ elif client_disconnected:
+ logger.info(f"HTTP 200 - POST /v1/messages (streaming) - client disconnected")
+ else:
+ logger.info(f"HTTP 200 - POST /v1/messages (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",
+ headers={
+ "Cache-Control": "no-cache",
+ "Connection": "keep-alive",
+ }
+ )
+
+ else:
+ # Non-streaming mode - collect entire response
+ anthropic_response = await collect_anthropic_response(
+ response,
+ request_data.model,
+ model_cache,
+ auth_manager,
+ request_messages=messages_for_tokenizer
+ )
+
+ await http_client.close()
+
+ logger.info(f"HTTP 200 - POST /v1/messages (non-streaming) - completed")
+
+ if debug_logger:
+ debug_logger.discard_buffers()
+
+ return JSONResponse(content=anthropic_response)
+
+ except HTTPException as e:
+ await http_client.close()
+ logger.error(f"HTTP {e.status_code} - POST /v1/messages - {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/messages - {str(e)[:100]}")
+ if debug_logger:
+ debug_logger.flush_on_error(500, str(e))
+
+ return JSONResponse(
+ status_code=500,
+ content={
+ "type": "error",
+ "error": {
+ "type": "api_error",
+ "message": f"Internal Server Error: {str(e)}"
+ }
+ }
+ )
\ No newline at end of file
diff --git a/kiro-gateway/kiro/routes_openai.py b/kiro-gateway/kiro/routes_openai.py
new file mode 100644
index 0000000000000000000000000000000000000000..80c3826fe0419cc38cba40294bfc10b43ccc886f
--- /dev/null
+++ b/kiro-gateway/kiro/routes_openai.py
@@ -0,0 +1,367 @@
+# -*- 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 routes for Kiro Gateway.
+
+Contains all API endpoints:
+- / and /health: Health check
+- /v1/models: Models list
+- /v1/chat/completions: Chat completions
+"""
+
+import json
+from datetime import datetime, timezone
+
+from fastapi import APIRouter, Depends, HTTPException, Request, Response, Security
+from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.security import APIKeyHeader
+from loguru import logger
+
+from kiro.config import (
+ PROXY_API_KEY,
+ APP_VERSION,
+)
+from kiro.models_openai import (
+ OpenAIModel,
+ ModelList,
+ ChatCompletionRequest,
+)
+from kiro.auth import KiroAuthManager, AuthType
+from kiro.cache import ModelInfoCache
+from kiro.model_resolver import ModelResolver
+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
+
+# Import debug_logger
+try:
+ from kiro.debug_logger import debug_logger
+except ImportError:
+ debug_logger = None
+
+
+# --- Security scheme ---
+api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
+
+
+async def verify_api_key(auth_header: str = Security(api_key_header)) -> bool:
+ """
+ Verify API key in Authorization header.
+
+ Expects format: "Bearer {PROXY_API_KEY}"
+
+ Args:
+ auth_header: Authorization header value
+
+ Returns:
+ True if key is valid
+
+ 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
+
+
+# --- Router ---
+router = APIRouter()
+
+
+@router.get("/")
+async def root():
+ """
+ Health check endpoint.
+
+ Returns:
+ Status and application version
+ """
+ return {
+ "status": "ok",
+ "message": "Kiro Gateway is running",
+ "version": APP_VERSION
+ }
+
+
+@router.get("/health")
+async def health():
+ """
+ Detailed health check.
+
+ Returns:
+ Status, timestamp and version
+ """
+ return {
+ "status": "healthy",
+ "timestamp": datetime.now(timezone.utc).isoformat(),
+ "version": APP_VERSION
+ }
+
+@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.
+
+ Args:
+ request: FastAPI Request for accessing app.state
+
+ Returns:
+ ModelList with available models in consistent format (with dots)
+ """
+ logger.info("Request to /v1/models")
+
+ model_resolver: ModelResolver = request.app.state.model_resolver
+
+ # Get all available models from resolver (cache + hidden models)
+ available_model_ids = model_resolver.get_available_models()
+
+ # Build OpenAI-compatible model list
+ openai_models = [
+ OpenAIModel(
+ id=model_id,
+ owned_by="anthropic",
+ description="Claude model via Kiro API"
+ )
+ for model_id in available_model_ids
+ ]
+
+ return ModelList(data=openai_models)
+
+
+@router.post("/v1/chat/completions", dependencies=[Depends(verify_api_key)])
+async def chat_completions(request: Request, request_data: ChatCompletionRequest):
+ """
+ Chat completions endpoint - compatible with OpenAI API.
+
+ Accepts requests in OpenAI format and translates them to Kiro API.
+ Supports streaming and non-streaming modes.
+
+ Args:
+ request: FastAPI Request for accessing app.state
+ request_data: Request in OpenAI ChatCompletionRequest format
+
+ Returns:
+ StreamingResponse for streaming mode
+ JSONResponse for non-streaming mode
+
+ Raises:
+ HTTPException: On validation or API errors
+ """
+ logger.info(f"Request to /v1/chat/completions (model={request_data.model}, stream={request_data.stream})")
+
+ auth_manager: KiroAuthManager = request.app.state.auth_manager
+ model_cache: ModelInfoCache = request.app.state.model_cache
+
+ # Note: prepare_new_request() and log_request_body() are now called by DebugLoggerMiddleware
+ # This ensures debug logging works even for requests that fail Pydantic validation (422 errors)
+
+ # Generate conversation ID
+ conversation_id = generate_conversation_id()
+
+ # Build payload for Kiro
+ # profileArn is only needed for Kiro Desktop auth
+ # AWS SSO OIDC (Builder ID) users don't need profileArn and it causes 403 if sent
+ profile_arn_for_payload = ""
+ if auth_manager.auth_type == AuthType.KIRO_DESKTOP and auth_manager.profile_arn:
+ profile_arn_for_payload = auth_manager.profile_arn
+
+ try:
+ kiro_payload = build_kiro_payload(
+ request_data,
+ conversation_id,
+ profile_arn_for_payload
+ )
+ except ValueError as e:
+ raise HTTPException(status_code=400, detail=str(e))
+
+ # Log Kiro payload
+ 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}")
+
+ # Create HTTP client with retry logic
+ # For streaming: use per-request client to avoid CLOSE_WAIT leak on VPN disconnect (issue #54)
+ # For non-streaming: use shared client for connection pooling
+ url = f"{auth_manager.api_host}/generateAssistantResponse"
+
+ if request_data.stream:
+ # Streaming mode: per-request client prevents orphaned connections
+ # 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)
+ try:
+ # Make request to Kiro API (for both streaming and non-streaming modes)
+ # Important: we wait for Kiro response BEFORE returning StreamingResponse,
+ # so that 200 OK means Kiro accepted the request and started responding
+ 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')
+ logger.error(f"Error from Kiro API: {response.status_code} - {error_text}")
+
+ # Try to parse JSON response from Kiro to extract error message
+ error_message = error_text
+ try:
+ error_json = json.loads(error_text)
+ if "message" in error_json:
+ error_message = error_json["message"]
+ if "reason" in error_json:
+ error_message = f"{error_message} (reason: {error_json['reason']})"
+ except (json.JSONDecodeError, KeyError):
+ pass
+
+ # Log access log for error (before flush, so it gets into app_logs)
+ logger.warning(
+ f"HTTP {response.status_code} - POST /v1/chat/completions - {error_message[:100]}"
+ )
+
+ # Flush debug logs on error ("errors" mode)
+ if debug_logger:
+ debug_logger.flush_on_error(response.status_code, error_message)
+
+ # Return error in OpenAI API format
+ return JSONResponse(
+ status_code=response.status_code,
+ content={
+ "error": {
+ "message": error_message,
+ "type": "kiro_api_error",
+ "code": response.status_code
+ }
+ }
+ )
+
+ # Prepare data for fallback token counting
+ # Convert Pydantic models to dicts for tokenizer
+ messages_for_tokenizer = [msg.model_dump() for msg in request_data.messages]
+ tools_for_tokenizer = [tool.model_dump() for tool in request_data.tools] if request_data.tools else None
+
+ if request_data.stream:
+ # Streaming mode
+ async def stream_wrapper():
+ streaming_error = None
+ client_disconnected = False
+ try:
+ async for chunk in stream_kiro_to_openai(
+ http_client.client,
+ response,
+ request_data.model,
+ model_cache,
+ auth_manager,
+ request_messages=messages_for_tokenizer,
+ request_tools=tools_for_tokenizer
+ ):
+ yield chunk
+ except GeneratorExit:
+ # Client disconnected - this is normal
+ client_disconnected = True
+ logger.debug("Client disconnected during streaming (GeneratorExit in routes)")
+ except Exception as e:
+ streaming_error = e
+ # Try to send [DONE] to client before finishing
+ # so client doesn't "hang" waiting for data
+ try:
+ yield "data: [DONE]\n\n"
+ except Exception:
+ pass # Client already disconnected
+ raise
+ finally:
+ await http_client.close()
+ # Log access log for streaming (success or error)
+ 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/chat/completions (streaming) - [{error_type}] {error_msg[:100]}")
+ elif client_disconnected:
+ logger.info(f"HTTP 200 - POST /v1/chat/completions (streaming) - client disconnected")
+ else:
+ logger.info(f"HTTP 200 - POST /v1/chat/completions (streaming) - completed")
+ # Write debug logs AFTER streaming completes
+ 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")
+
+ else:
+
+ # Non-streaming mode - collect entire response
+ openai_response = await collect_stream_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()
+
+ # Log access log for non-streaming success
+ logger.info(f"HTTP 200 - POST /v1/chat/completions (non-streaming) - completed")
+
+ # Write debug logs after non-streaming request completes
+ if debug_logger:
+ debug_logger.discard_buffers()
+
+ return JSONResponse(content=openai_response)
+
+ except HTTPException as e:
+ await http_client.close()
+ # Log access log for HTTP error
+ logger.error(f"HTTP {e.status_code} - POST /v1/chat/completions - {e.detail}")
+ # Flush debug logs on HTTP error ("errors" mode)
+ 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)
+ # Log access log for internal error
+ logger.error(f"HTTP 500 - POST /v1/chat/completions - {str(e)[:100]}")
+ # Flush debug logs on internal error ("errors" mode)
+ if debug_logger:
+ debug_logger.flush_on_error(500, str(e))
+ raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}")
\ No newline at end of file
diff --git a/kiro-gateway/kiro/streaming_anthropic.py b/kiro-gateway/kiro/streaming_anthropic.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a1cfbf3d576fd7de9c6d68615efad14313a1907
--- /dev/null
+++ b/kiro-gateway/kiro/streaming_anthropic.py
@@ -0,0 +1,671 @@
+# -*- 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 Anthropic Messages API format.
+
+This module formats Kiro events into Anthropic SSE format:
+- event: message_start
+- event: content_block_start
+- event: content_block_delta
+- event: content_block_stop
+- event: message_delta
+- event: message_stop
+
+Reference: https://docs.anthropic.com/en/api/messages-streaming
+"""
+
+import json
+import time
+import uuid
+from typing import TYPE_CHECKING, AsyncGenerator, Dict, List, Optional, Any
+
+import httpx
+from loguru import logger
+
+from kiro.streaming_core import (
+ parse_kiro_stream,
+ collect_stream_to_result,
+ FirstTokenTimeoutError,
+ KiroEvent,
+ calculate_tokens_from_context_usage,
+ stream_with_first_token_retry,
+)
+from kiro.tokenizer import count_tokens, count_message_tokens, count_tools_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
+
+if TYPE_CHECKING:
+ from kiro.auth import KiroAuthManager
+ from kiro.cache import ModelInfoCache
+
+# Import debug_logger for logging
+try:
+ from kiro.debug_logger import debug_logger
+except ImportError:
+ debug_logger = None
+
+
+def generate_message_id() -> str:
+ """Generate unique message ID in Anthropic format."""
+ return f"msg_{uuid.uuid4().hex[:24]}"
+
+
+def format_sse_event(event_type: str, data: Dict[str, Any]) -> str:
+ """
+ Format data as Anthropic SSE event.
+
+ Anthropic SSE format:
+ event: {event_type}
+ data: {json_data}
+
+ Args:
+ event_type: Event type (message_start, content_block_delta, etc.)
+ data: Event data dictionary
+
+ Returns:
+ Formatted SSE string
+ """
+ return f"event: {event_type}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
+
+
+def generate_thinking_signature() -> str:
+ """
+ Generate a placeholder signature for thinking content blocks.
+
+ In real Anthropic API, this is a cryptographic signature for verification.
+ Since we're using fake reasoning via tag injection, we generate a placeholder.
+
+ Returns:
+ Placeholder signature string
+ """
+ return f"sig_{uuid.uuid4().hex[:32]}"
+
+
+async def stream_kiro_to_anthropic(
+ response: httpx.Response,
+ model: str,
+ model_cache: "ModelInfoCache",
+ auth_manager: "KiroAuthManager",
+ first_token_timeout: float = FIRST_TOKEN_TIMEOUT,
+ request_messages: Optional[list] = None
+) -> AsyncGenerator[str, None]:
+ """
+ Generator for converting Kiro stream to Anthropic SSE format.
+
+ Parses Kiro AWS SSE stream and converts events to Anthropic format.
+ Supports thinking content blocks when FAKE_REASONING_HANDLING=as_reasoning_content.
+
+ Args:
+ response: HTTP response with data stream
+ model: Model name to include in response
+ model_cache: Model cache for getting token limits
+ auth_manager: Authentication manager
+ first_token_timeout: First token wait timeout (seconds)
+ request_messages: Original request messages (for token counting)
+
+ Yields:
+ Strings in Anthropic SSE format
+
+ Raises:
+ FirstTokenTimeoutError: If first token not received within timeout
+ """
+ message_id = generate_message_id()
+ input_tokens = 0
+ output_tokens = 0
+ full_content = ""
+ full_thinking_content = ""
+
+ # Count input tokens from request messages
+ if request_messages:
+ input_tokens = count_message_tokens(request_messages, apply_claude_correction=False)
+
+ # Track content blocks - thinking block is index 0, text block is index 1 (when thinking enabled)
+ current_block_index = 0
+ thinking_block_started = False
+ thinking_block_index: Optional[int] = None
+ text_block_started = False
+ text_block_index: Optional[int] = None
+ tool_blocks: List[Dict[str, Any]] = []
+ tool_input_buffers: Dict[int, str] = {} # index -> accumulated JSON
+
+ # Generate signature for thinking block (used if thinking is present)
+ thinking_signature = generate_thinking_signature()
+
+ # Track context usage for token calculation
+ context_usage_percentage: Optional[float] = None
+
+ try:
+ # Send message_start event
+ yield format_sse_event("message_start", {
+ "type": "message_start",
+ "message": {
+ "id": message_id,
+ "type": "message",
+ "role": "assistant",
+ "content": [],
+ "model": model,
+ "stop_reason": None,
+ "stop_sequence": None,
+ "usage": {
+ "input_tokens": input_tokens,
+ "output_tokens": 0
+ }
+ }
+ })
+
+ async for event in parse_kiro_stream(response, first_token_timeout):
+ if event.type == "content":
+ content = event.content or ""
+ full_content += content
+
+ # Close thinking block if it was open and we're now getting regular content
+ if thinking_block_started and thinking_block_index is not None:
+ yield format_sse_event("content_block_stop", {
+ "type": "content_block_stop",
+ "index": thinking_block_index
+ })
+ thinking_block_started = False
+ current_block_index += 1
+
+ # Start text block if not started
+ if not text_block_started:
+ text_block_index = current_block_index
+ yield format_sse_event("content_block_start", {
+ "type": "content_block_start",
+ "index": text_block_index,
+ "content_block": {
+ "type": "text",
+ "text": ""
+ }
+ })
+ text_block_started = True
+
+ # Send content delta
+ if content:
+ yield format_sse_event("content_block_delta", {
+ "type": "content_block_delta",
+ "index": text_block_index,
+ "delta": {
+ "type": "text_delta",
+ "text": content
+ }
+ })
+
+ elif event.type == "thinking":
+ thinking_content = event.thinking_content or ""
+ full_thinking_content += thinking_content
+
+ # Handle thinking content based on mode
+ if FAKE_REASONING_HANDLING == "as_reasoning_content":
+ # Use native Anthropic thinking content blocks
+ if not thinking_block_started:
+ thinking_block_index = current_block_index
+ yield format_sse_event("content_block_start", {
+ "type": "content_block_start",
+ "index": thinking_block_index,
+ "content_block": {
+ "type": "thinking",
+ "thinking": "",
+ "signature": thinking_signature
+ }
+ })
+ thinking_block_started = True
+
+ if thinking_content:
+ yield format_sse_event("content_block_delta", {
+ "type": "content_block_delta",
+ "index": thinking_block_index,
+ "delta": {
+ "type": "thinking_delta",
+ "thinking": thinking_content
+ }
+ })
+
+ elif FAKE_REASONING_HANDLING == "include_as_text":
+ # Include thinking as regular text content
+ # Close thinking block if it was open (shouldn't happen in this mode)
+ if thinking_block_started and thinking_block_index is not None:
+ yield format_sse_event("content_block_stop", {
+ "type": "content_block_stop",
+ "index": thinking_block_index
+ })
+ thinking_block_started = False
+ current_block_index += 1
+
+ # Start text block if not started
+ if not text_block_started:
+ text_block_index = current_block_index
+ yield format_sse_event("content_block_start", {
+ "type": "content_block_start",
+ "index": text_block_index,
+ "content_block": {
+ "type": "text",
+ "text": ""
+ }
+ })
+ text_block_started = True
+
+ if thinking_content:
+ yield format_sse_event("content_block_delta", {
+ "type": "content_block_delta",
+ "index": text_block_index,
+ "delta": {
+ "type": "text_delta",
+ "text": thinking_content
+ }
+ })
+ # For "strip" mode, we just skip the thinking content
+
+ elif event.type == "tool_use" and event.tool_use:
+ # Close thinking block if open
+ if thinking_block_started and thinking_block_index is not None:
+ yield format_sse_event("content_block_stop", {
+ "type": "content_block_stop",
+ "index": thinking_block_index
+ })
+ thinking_block_started = False
+ current_block_index += 1
+
+ # Close text block if open
+ if text_block_started and text_block_index is not None:
+ yield format_sse_event("content_block_stop", {
+ "type": "content_block_stop",
+ "index": text_block_index
+ })
+ text_block_started = False
+ current_block_index += 1
+
+ tool = event.tool_use
+ tool_id = tool.get("id") or f"toolu_{uuid.uuid4().hex[:24]}"
+ tool_name = tool.get("function", {}).get("name", "") or tool.get("name", "")
+ tool_input = tool.get("function", {}).get("arguments", {}) or tool.get("input", {})
+
+ # Parse arguments if string
+ if isinstance(tool_input, str):
+ try:
+ tool_input = json.loads(tool_input)
+ except json.JSONDecodeError:
+ tool_input = {}
+
+ # Send tool_use block start
+ yield format_sse_event("content_block_start", {
+ "type": "content_block_start",
+ "index": current_block_index,
+ "content_block": {
+ "type": "tool_use",
+ "id": tool_id,
+ "name": tool_name,
+ "input": {}
+ }
+ })
+
+ # Send tool input as delta
+ input_json = json.dumps(tool_input, ensure_ascii=False)
+ yield format_sse_event("content_block_delta", {
+ "type": "content_block_delta",
+ "index": current_block_index,
+ "delta": {
+ "type": "input_json_delta",
+ "partial_json": input_json
+ }
+ })
+
+ # Close tool block
+ yield format_sse_event("content_block_stop", {
+ "type": "content_block_stop",
+ "index": current_block_index
+ })
+
+ tool_blocks.append({
+ "id": tool_id,
+ "name": tool_name,
+ "input": tool_input
+ })
+ current_block_index += 1
+
+ elif event.type == "context_usage" and event.context_usage_percentage is not None:
+ context_usage_percentage = event.context_usage_percentage
+
+ # Check for bracket-style tool calls in full content
+ bracket_tool_calls = parse_bracket_tool_calls(full_content)
+ if bracket_tool_calls:
+ # Close thinking block if open
+ if thinking_block_started and thinking_block_index is not None:
+ yield format_sse_event("content_block_stop", {
+ "type": "content_block_stop",
+ "index": thinking_block_index
+ })
+ thinking_block_started = False
+ current_block_index += 1
+
+ # Close text block if open
+ if text_block_started and text_block_index is not None:
+ yield format_sse_event("content_block_stop", {
+ "type": "content_block_stop",
+ "index": text_block_index
+ })
+ text_block_started = False
+ current_block_index += 1
+
+ for tc in bracket_tool_calls:
+ tool_id = tc.get("id") or f"toolu_{uuid.uuid4().hex[:24]}"
+ tool_name = tc.get("function", {}).get("name", "")
+ tool_input = tc.get("function", {}).get("arguments", {})
+
+ if isinstance(tool_input, str):
+ try:
+ tool_input = json.loads(tool_input)
+ except json.JSONDecodeError:
+ tool_input = {}
+
+ yield format_sse_event("content_block_start", {
+ "type": "content_block_start",
+ "index": current_block_index,
+ "content_block": {
+ "type": "tool_use",
+ "id": tool_id,
+ "name": tool_name,
+ "input": {}
+ }
+ })
+
+ input_json = json.dumps(tool_input, ensure_ascii=False)
+ yield format_sse_event("content_block_delta", {
+ "type": "content_block_delta",
+ "index": current_block_index,
+ "delta": {
+ "type": "input_json_delta",
+ "partial_json": input_json
+ }
+ })
+
+ yield format_sse_event("content_block_stop", {
+ "type": "content_block_stop",
+ "index": current_block_index
+ })
+
+ tool_blocks.append({
+ "id": tool_id,
+ "name": tool_name,
+ "input": tool_input
+ })
+ current_block_index += 1
+
+ # Close thinking block if still open
+ if thinking_block_started and thinking_block_index is not None:
+ yield format_sse_event("content_block_stop", {
+ "type": "content_block_stop",
+ "index": thinking_block_index
+ })
+ current_block_index += 1
+
+ # Close text block if still open
+ if text_block_started and text_block_index is not None:
+ yield format_sse_event("content_block_stop", {
+ "type": "content_block_stop",
+ "index": text_block_index
+ })
+
+ # Calculate output tokens
+ output_tokens = count_tokens(full_content + full_thinking_content)
+
+ # Calculate total tokens from context usage if available
+ if context_usage_percentage is not None:
+ prompt_tokens, total_tokens, _, _ = calculate_tokens_from_context_usage(
+ context_usage_percentage, output_tokens, model_cache, model
+ )
+ input_tokens = prompt_tokens
+
+ # Determine stop reason
+ stop_reason = "tool_use" if tool_blocks else "end_turn"
+
+ # Send message_delta with stop_reason and usage
+ yield format_sse_event("message_delta", {
+ "type": "message_delta",
+ "delta": {
+ "stop_reason": stop_reason,
+ "stop_sequence": None
+ },
+ "usage": {
+ "output_tokens": output_tokens
+ }
+ })
+
+ # Send message_stop
+ yield format_sse_event("message_stop", {
+ "type": "message_stop"
+ })
+
+ logger.debug(
+ f"[Anthropic Streaming] Completed: "
+ f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
+ f"tool_blocks={len(tool_blocks)}, stop_reason={stop_reason}"
+ )
+
+ except FirstTokenTimeoutError:
+ raise
+ except GeneratorExit:
+ logger.debug("Client disconnected (GeneratorExit)")
+ raise
+ except Exception as e:
+ error_type = type(e).__name__
+ error_msg = str(e) if str(e) else "(empty message)"
+ logger.error(f"Error during Anthropic streaming: [{error_type}] {error_msg}", exc_info=True)
+
+ # Send error event
+ yield format_sse_event("error", {
+ "type": "error",
+ "error": {
+ "type": "api_error",
+ "message": f"Internal error: {error_msg}"
+ }
+ })
+ raise
+ finally:
+ try:
+ await response.aclose()
+ except Exception as close_error:
+ logger.debug(f"Error closing response: {close_error}")
+
+
+async def collect_anthropic_response(
+ response: httpx.Response,
+ model: str,
+ model_cache: "ModelInfoCache",
+ auth_manager: "KiroAuthManager",
+ request_messages: Optional[list] = None
+) -> dict:
+ """
+ Collect full response from Kiro stream in Anthropic format.
+
+ Used for non-streaming mode.
+
+ Args:
+ response: HTTP response with stream
+ model: Model name
+ model_cache: Model cache
+ auth_manager: Authentication manager
+ request_messages: Original request messages (for token counting)
+
+ Returns:
+ Dictionary with full response in Anthropic Messages format
+ """
+ message_id = generate_message_id()
+
+ # Count input tokens
+ input_tokens = 0
+ if request_messages:
+ input_tokens = count_message_tokens(request_messages, apply_claude_correction=False)
+
+ # Collect stream result
+ result = await collect_stream_to_result(response)
+
+ # Build content blocks
+ content_blocks = []
+
+ # Add thinking block FIRST if there's thinking content and mode is as_reasoning_content
+ if result.thinking_content and FAKE_REASONING_HANDLING == "as_reasoning_content":
+ content_blocks.append({
+ "type": "thinking",
+ "thinking": result.thinking_content,
+ "signature": generate_thinking_signature()
+ })
+
+ # Add text block if there's content
+ # For include_as_text mode, prepend thinking content to regular content
+ text_content = result.content
+ if result.thinking_content and FAKE_REASONING_HANDLING == "include_as_text":
+ text_content = result.thinking_content + text_content
+
+ if text_content:
+ content_blocks.append({
+ "type": "text",
+ "text": text_content
+ })
+
+ # Add tool use blocks
+ for tc in result.tool_calls:
+ tool_id = tc.get("id") or f"toolu_{uuid.uuid4().hex[:24]}"
+ tool_name = tc.get("function", {}).get("name", "") or tc.get("name", "")
+ tool_input = tc.get("function", {}).get("arguments", {}) or tc.get("input", {})
+
+ if isinstance(tool_input, str):
+ try:
+ tool_input = json.loads(tool_input)
+ except json.JSONDecodeError:
+ tool_input = {}
+
+ content_blocks.append({
+ "type": "tool_use",
+ "id": tool_id,
+ "name": tool_name,
+ "input": tool_input
+ })
+
+ # Calculate output tokens
+ output_tokens = count_tokens(result.content + result.thinking_content)
+
+ # Calculate from context usage if available
+ if result.context_usage_percentage is not None:
+ prompt_tokens, _, _, _ = calculate_tokens_from_context_usage(
+ result.context_usage_percentage, output_tokens, model_cache, model
+ )
+ input_tokens = prompt_tokens
+
+ # Determine stop reason
+ stop_reason = "tool_use" if result.tool_calls else "end_turn"
+
+ logger.debug(
+ f"[Anthropic Non-Streaming] Completed: "
+ f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
+ f"tool_calls={len(result.tool_calls)}, stop_reason={stop_reason}"
+ )
+
+ return {
+ "id": message_id,
+ "type": "message",
+ "role": "assistant",
+ "content": content_blocks,
+ "model": model,
+ "stop_reason": stop_reason,
+ "stop_sequence": None,
+ "usage": {
+ "input_tokens": input_tokens,
+ "output_tokens": output_tokens
+ }
+ }
+
+
+async def stream_with_first_token_retry_anthropic(
+ make_request,
+ model: str,
+ model_cache: "ModelInfoCache",
+ auth_manager: "KiroAuthManager",
+ 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 for Anthropic API.
+
+ If model doesn't respond within first_token_timeout seconds,
+ request is cancelled and a new one is made. Maximum max_retries attempts.
+
+ This is seamless for user - they just see a delay,
+ but eventually get a response (or error after all attempts).
+
+ Args:
+ make_request: Function to create new HTTP request
+ model: Model name
+ model_cache: Model cache
+ auth_manager: Authentication manager
+ max_retries: Maximum number of attempts
+ first_token_timeout: First token wait timeout (seconds)
+ request_messages: Original request messages (for fallback token counting)
+ request_tools: Original request tools (for fallback token counting)
+
+ Yields:
+ Strings in Anthropic SSE format
+
+ Raises:
+ Exception with Anthropic error format after exhausting all attempts
+ """
+ def create_http_error(status_code: int, error_text: str) -> Exception:
+ """Create exception for HTTP errors in Anthropic format."""
+ return Exception(json.dumps({
+ "type": "error",
+ "error": {
+ "type": "api_error",
+ "message": f"Upstream API error: {error_text}"
+ }
+ }))
+
+ def create_timeout_error(retries: int, timeout: float) -> Exception:
+ """Create exception for timeout errors in Anthropic format."""
+ return Exception(json.dumps({
+ "type": "error",
+ "error": {
+ "type": "timeout_error",
+ "message": f"Model did not respond within {timeout}s after {retries} attempts. Please try again."
+ }
+ }))
+
+ async def stream_processor(response: httpx.Response) -> AsyncGenerator[str, None]:
+ """Process response and yield Anthropic SSE chunks."""
+ async for chunk in stream_kiro_to_anthropic(
+ response,
+ model,
+ model_cache,
+ auth_manager,
+ first_token_timeout=first_token_timeout,
+ request_messages=request_messages
+ ):
+ yield chunk
+
+ async for chunk in stream_with_first_token_retry(
+ make_request=make_request,
+ stream_processor=stream_processor,
+ 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
\ No newline at end of file
diff --git a/kiro-gateway/kiro/streaming_core.py b/kiro-gateway/kiro/streaming_core.py
new file mode 100644
index 0000000000000000000000000000000000000000..5957a3eadb3d5f5e4d771163c9c5db45bc0c3f86
--- /dev/null
+++ b/kiro-gateway/kiro/streaming_core.py
@@ -0,0 +1,494 @@
+# -*- 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 .
+
+"""
+Core streaming logic for parsing Kiro API responses.
+
+This module contains shared logic used by both OpenAI and Anthropic streaming:
+- KiroEvent dataclass for unified events
+- Kiro SSE stream parsing
+- Full response collection
+- First token timeout handling
+
+The core layer provides a unified interface that API-specific formatters use
+to convert Kiro events to their respective SSE formats.
+"""
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Awaitable, Dict, List, Optional, Tuple
+
+import httpx
+from loguru import logger
+
+from kiro.parsers import AwsEventStreamParser, parse_bracket_tool_calls, deduplicate_tool_calls
+from kiro.config import (
+ FIRST_TOKEN_TIMEOUT,
+ FIRST_TOKEN_MAX_RETRIES,
+ FAKE_REASONING_ENABLED,
+ FAKE_REASONING_HANDLING,
+)
+from kiro.thinking_parser import ThinkingParser
+
+if TYPE_CHECKING:
+ from kiro.cache import ModelInfoCache
+
+# Import debug_logger for logging
+try:
+ from kiro.debug_logger import debug_logger
+except ImportError:
+ debug_logger = None
+
+
+# ==================================================================================================
+# Data Classes
+# ==================================================================================================
+
+@dataclass
+class KiroEvent:
+ """
+ Unified event from Kiro API stream.
+
+ This format is API-agnostic and can be converted to both OpenAI and Anthropic formats.
+
+ Attributes:
+ type: Event type (content, thinking, tool_use, usage, context_usage, error)
+ content: Text content (for content events)
+ thinking_content: Thinking/reasoning content (for thinking events)
+ tool_use: Tool use data (for tool_use events)
+ usage: Usage/metering data (for usage events)
+ context_usage_percentage: Context usage percentage (for context_usage events)
+ is_first_thinking_chunk: Whether this is the first thinking chunk
+ is_last_thinking_chunk: Whether this is the last thinking chunk
+ """
+ type: str
+ content: Optional[str] = None
+ thinking_content: Optional[str] = None
+ tool_use: Optional[Dict[str, Any]] = None
+ usage: Optional[Dict[str, Any]] = None
+ context_usage_percentage: Optional[float] = None
+ is_first_thinking_chunk: bool = False
+ is_last_thinking_chunk: bool = False
+
+
+@dataclass
+class StreamResult:
+ """
+ Result of collecting a complete stream response.
+
+ Attributes:
+ content: Full text content
+ thinking_content: Full thinking/reasoning content
+ tool_calls: List of tool calls
+ usage: Usage information
+ context_usage_percentage: Context usage percentage from Kiro API
+ """
+ content: str = ""
+ thinking_content: str = ""
+ tool_calls: List[Dict[str, Any]] = field(default_factory=list)
+ usage: Optional[Dict[str, Any]] = None
+ context_usage_percentage: Optional[float] = None
+
+
+class FirstTokenTimeoutError(Exception):
+ """Exception raised when first token timeout occurs."""
+ pass
+
+
+# ==================================================================================================
+# Kiro Stream Parsing
+# ==================================================================================================
+
+async def parse_kiro_stream(
+ response: httpx.Response,
+ first_token_timeout: float = FIRST_TOKEN_TIMEOUT,
+ enable_thinking_parser: bool = True
+) -> AsyncGenerator[KiroEvent, None]:
+ """
+ Parses Kiro SSE stream and yields unified events.
+
+ This is the core parsing function that converts Kiro's AWS SSE format
+ into unified KiroEvent objects that can be formatted for any API.
+
+ Args:
+ response: HTTP response with data stream
+ first_token_timeout: First token wait timeout (seconds)
+ enable_thinking_parser: Whether to enable thinking block parsing
+
+ Yields:
+ KiroEvent objects representing stream events
+
+ Raises:
+ FirstTokenTimeoutError: If first token not received within timeout
+ """
+ parser = AwsEventStreamParser()
+ first_token_received = False
+
+ # Initialize thinking parser if fake reasoning is enabled
+ thinking_parser: Optional[ThinkingParser] = None
+ if FAKE_REASONING_ENABLED and enable_thinking_parser:
+ thinking_parser = ThinkingParser(handling_mode=FAKE_REASONING_HANDLING)
+ logger.debug(f"Thinking parser initialized with mode: {FAKE_REASONING_HANDLING}")
+
+ try:
+ # Create iterator for reading bytes
+ byte_iterator = response.aiter_bytes()
+
+ # Wait for first chunk with timeout
+ try:
+ logger.debug(f"Waiting for first token (timeout={first_token_timeout}s)...")
+ first_byte_chunk = await asyncio.wait_for(
+ byte_iterator.__anext__(),
+ timeout=first_token_timeout
+ )
+ logger.debug("First token received")
+ except asyncio.TimeoutError:
+ logger.warning(f"[FirstTokenTimeout] Model did not respond within {first_token_timeout}s")
+ raise FirstTokenTimeoutError(f"No response within {first_token_timeout} seconds")
+ except StopAsyncIteration:
+ # Empty response - this is normal, just finish
+ logger.debug("Empty response from Kiro API")
+ return
+
+ # Process first chunk
+ if debug_logger:
+ debug_logger.log_raw_chunk(first_byte_chunk)
+
+ async for event in _process_chunk(parser, first_byte_chunk, thinking_parser):
+ if event.type == "content" or event.type == "thinking":
+ first_token_received = True
+ 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
+
+ # Finalize thinking parser and yield any remaining content
+ if thinking_parser:
+ final_result = thinking_parser.finalize()
+
+ if final_result.thinking_content:
+ processed_thinking = thinking_parser.process_for_output(
+ final_result.thinking_content,
+ final_result.is_first_thinking_chunk,
+ final_result.is_last_thinking_chunk,
+ )
+ if processed_thinking:
+ yield KiroEvent(
+ type="thinking",
+ thinking_content=processed_thinking,
+ is_first_thinking_chunk=final_result.is_first_thinking_chunk,
+ is_last_thinking_chunk=final_result.is_last_thinking_chunk,
+ )
+
+ if final_result.regular_content:
+ yield KiroEvent(type="content", content=final_result.regular_content)
+
+ if thinking_parser.found_thinking_block:
+ logger.debug("Thinking block processing completed")
+
+ # Check bracket-style tool calls in accumulated content
+ all_tool_calls = parser.get_tool_calls()
+ # Note: bracket tool calls are checked by the caller using full content
+
+ # Yield tool calls if any
+ for tc in all_tool_calls:
+ yield KiroEvent(type="tool_use", tool_use=tc)
+
+ except FirstTokenTimeoutError:
+ raise
+ except GeneratorExit:
+ logger.debug("Client disconnected (GeneratorExit)")
+ raise
+ except Exception as e:
+ error_type = type(e).__name__
+ 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
+
+
+async def _process_chunk(
+ parser: AwsEventStreamParser,
+ chunk: bytes,
+ thinking_parser: Optional[ThinkingParser]
+) -> AsyncGenerator[KiroEvent, None]:
+ """
+ Process a single chunk from Kiro stream.
+
+ Args:
+ parser: AWS event stream parser
+ chunk: Raw bytes chunk
+ thinking_parser: Optional thinking parser for fake reasoning
+
+ Yields:
+ KiroEvent objects
+ """
+ events = parser.feed(chunk)
+
+ for event in events:
+ if event["type"] == "content":
+ content = event["data"]
+
+ # Process through thinking parser if enabled
+ if thinking_parser:
+ parse_result = thinking_parser.feed(content)
+
+ # Yield thinking content if any
+ if parse_result.thinking_content:
+ processed_thinking = thinking_parser.process_for_output(
+ parse_result.thinking_content,
+ parse_result.is_first_thinking_chunk,
+ parse_result.is_last_thinking_chunk,
+ )
+ if processed_thinking:
+ yield KiroEvent(
+ type="thinking",
+ thinking_content=processed_thinking,
+ is_first_thinking_chunk=parse_result.is_first_thinking_chunk,
+ is_last_thinking_chunk=parse_result.is_last_thinking_chunk,
+ )
+
+ # Yield regular content if any
+ if parse_result.regular_content:
+ yield KiroEvent(type="content", content=parse_result.regular_content)
+ else:
+ # No thinking parser - pass through as-is
+ yield KiroEvent(type="content", content=content)
+
+ elif event["type"] == "usage":
+ yield KiroEvent(type="usage", usage=event["data"])
+
+ elif event["type"] == "context_usage":
+ yield KiroEvent(type="context_usage", context_usage_percentage=event["data"])
+
+
+# ==================================================================================================
+# Full Response Collection
+# ==================================================================================================
+
+async def collect_stream_to_result(
+ response: httpx.Response,
+ first_token_timeout: float = FIRST_TOKEN_TIMEOUT,
+ enable_thinking_parser: bool = True
+) -> StreamResult:
+ """
+ Collects full response from Kiro stream.
+
+ This function consumes the entire stream and returns a StreamResult
+ with all accumulated data.
+
+ Args:
+ response: HTTP response with stream
+ first_token_timeout: First token wait timeout
+ enable_thinking_parser: Whether to enable thinking block parsing
+
+ Returns:
+ StreamResult with full content, thinking, tool calls, and usage
+ """
+ result = StreamResult()
+ full_content_for_bracket_tools = ""
+
+ async for event in parse_kiro_stream(response, first_token_timeout, enable_thinking_parser):
+ if event.type == "content" and event.content:
+ result.content += event.content
+ full_content_for_bracket_tools += event.content
+ elif event.type == "thinking" and event.thinking_content:
+ result.thinking_content += event.thinking_content
+ full_content_for_bracket_tools += event.thinking_content
+ elif event.type == "tool_use" and event.tool_use:
+ result.tool_calls.append(event.tool_use)
+ elif event.type == "usage" and event.usage:
+ result.usage = event.usage
+ elif event.type == "context_usage" and event.context_usage_percentage is not None:
+ result.context_usage_percentage = event.context_usage_percentage
+
+ # Check for bracket-style tool calls in full content
+ 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)
+
+ return result
+
+
+# ==================================================================================================
+# Token Counting Utilities
+# ==================================================================================================
+
+def calculate_tokens_from_context_usage(
+ context_usage_percentage: Optional[float],
+ completion_tokens: int,
+ model_cache: "ModelInfoCache",
+ model: str
+) -> Tuple[int, int, str, str]:
+ """
+ Calculate token counts from Kiro's context usage percentage.
+
+ Args:
+ context_usage_percentage: Context usage percentage from Kiro API
+ completion_tokens: Number of completion tokens (counted via tiktoken)
+ model_cache: Model cache for getting max input tokens
+ model: Model name
+
+ Returns:
+ Tuple of (prompt_tokens, total_tokens, prompt_source, total_source)
+ """
+ if context_usage_percentage is not None and context_usage_percentage > 0:
+ max_input_tokens = model_cache.get_max_input_tokens(model)
+ total_tokens = int((context_usage_percentage / 100) * max_input_tokens)
+ prompt_tokens = max(0, total_tokens - completion_tokens)
+ return prompt_tokens, total_tokens, "subtraction", "API Kiro"
+
+ # Fallback: no context usage data
+ return 0, completion_tokens, "unknown", "tiktoken"
+
+
+# ==================================================================================================
+# First Token Retry Logic
+# ==================================================================================================
+
+async def stream_with_first_token_retry(
+ make_request: Callable[[], Awaitable[httpx.Response]],
+ stream_processor: Callable[[httpx.Response], AsyncGenerator[str, None]],
+ max_retries: int = FIRST_TOKEN_MAX_RETRIES,
+ first_token_timeout: float = FIRST_TOKEN_TIMEOUT,
+ on_http_error: Optional[Callable[[int, str], Exception]] = None,
+ on_all_retries_failed: Optional[Callable[[int, float], Exception]] = None,
+) -> AsyncGenerator[str, None]:
+ """
+ Generic streaming with automatic retry on first token timeout.
+
+ If model doesn't respond within first_token_timeout seconds,
+ request is cancelled and a new one is made. Maximum max_retries attempts.
+
+ This is seamless for user - they just see a delay,
+ but eventually get a response (or error after all attempts).
+
+ Args:
+ make_request: Function to create new HTTP request (returns httpx.Response)
+ stream_processor: Function that processes response and yields SSE strings.
+ Must use parse_kiro_stream internally for timeout handling.
+ max_retries: Maximum number of attempts
+ first_token_timeout: First token wait timeout (seconds)
+ on_http_error: Optional callback to create exception for HTTP errors.
+ Receives (status_code, error_text), returns Exception.
+ If None, raises generic Exception.
+ on_all_retries_failed: Optional callback to create exception when all retries fail.
+ Receives (max_retries, timeout), returns Exception.
+ If None, raises generic Exception.
+
+ Yields:
+ Strings in SSE format (format depends on stream_processor)
+
+ Raises:
+ Exception from on_http_error or on_all_retries_failed callbacks
+
+ Example:
+ >>> async def make_req():
+ ... return await http_client.request_with_retry("POST", url, payload, stream=True)
+ >>> async def process(response):
+ ... async for chunk in stream_kiro_to_openai(response, ...):
+ ... yield chunk
+ >>> async for chunk in stream_with_first_token_retry(make_req, process):
+ ... print(chunk)
+ """
+ last_error: Optional[Exception] = None
+
+ for attempt in range(max_retries):
+ response: Optional[httpx.Response] = None
+ try:
+ # Make request
+ if attempt > 0:
+ logger.warning(f"Retry attempt {attempt + 1}/{max_retries} after first token timeout")
+
+ response = await make_request()
+
+ if response.status_code != 200:
+ # Error from API - close response and raise exception
+ try:
+ error_content = await response.aread()
+ error_text = error_content.decode('utf-8', errors='replace')
+ except Exception:
+ error_text = "Unknown error"
+
+ try:
+ await response.aclose()
+ except Exception:
+ pass
+
+ logger.error(f"Error from Kiro API: {response.status_code} - {error_text}")
+
+ if on_http_error:
+ raise on_http_error(response.status_code, error_text)
+ else:
+ raise Exception(f"Upstream API error ({response.status_code}): {error_text}")
+
+ # Try to stream with first token timeout
+ async for chunk in stream_processor(response):
+ yield chunk
+
+ # Successfully completed - exit
+ return
+
+ except FirstTokenTimeoutError as e:
+ last_error = e
+ logger.warning(
+ f"[FirstTokenTimeout] Attempt {attempt + 1}/{max_retries} failed - "
+ f"model did not respond within {first_token_timeout}s"
+ )
+
+ # Close current response if open
+ if response:
+ try:
+ await response.aclose()
+ except Exception:
+ pass
+
+ # Continue to next attempt
+ continue
+
+ except Exception as e:
+ # Other errors - no retry, propagate
+ # Use positional argument to avoid loguru interpreting curly braces in error message as format placeholders
+ # f-string with repr() doesn't work because loguru still sees {type} inside the string
+ error_msg = str(e) if str(e) else "(empty message)"
+ logger.error("Unexpected error during streaming: {}", error_msg, exc_info=True)
+ if response:
+ try:
+ await response.aclose()
+ except Exception:
+ pass
+ raise
+
+ # All attempts exhausted - raise error
+ logger.error(
+ f"[FirstTokenTimeout] All {max_retries} attempts exhausted - "
+ f"model never responded within {first_token_timeout}s per attempt"
+ )
+
+ if on_all_retries_failed:
+ raise on_all_retries_failed(max_retries, first_token_timeout)
+ else:
+ raise Exception(
+ f"Model did not respond within {first_token_timeout}s after {max_retries} attempts. "
+ "Please try again."
+ )
\ No newline at end of file
diff --git a/kiro-gateway/kiro/streaming_openai.py b/kiro-gateway/kiro/streaming_openai.py
new file mode 100644
index 0000000000000000000000000000000000000000..981bfe65a3babdc6596f98d029bf329be19d4b0e
--- /dev/null
+++ b/kiro-gateway/kiro/streaming_openai.py
@@ -0,0 +1,549 @@
+# -*- 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 format.
+
+Contains generators for:
+- Converting AWS SSE to OpenAI SSE
+- Forming streaming chunks
+- Processing tool calls in stream
+
+Uses streaming_core.py for parsing Kiro stream into unified KiroEvent objects.
+"""
+
+import json
+import time
+from typing import TYPE_CHECKING, AsyncGenerator, Callable, Awaitable, Optional
+
+import httpx
+from fastapi import HTTPException
+from loguru import logger
+
+from kiro.parsers import parse_bracket_tool_calls, deduplicate_tool_calls
+from kiro.utils import generate_completion_id
+from kiro.config import (
+ FIRST_TOKEN_TIMEOUT,
+ FIRST_TOKEN_MAX_RETRIES,
+ FAKE_REASONING_HANDLING,
+)
+from kiro.tokenizer import count_tokens, count_message_tokens, count_tools_tokens
+
+# Import from streaming_core - reuse shared parsing logic
+from kiro.streaming_core import (
+ parse_kiro_stream,
+ FirstTokenTimeoutError,
+ KiroEvent,
+ 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
+
+# Import debug_logger for logging
+try:
+ from kiro.debug_logger import debug_logger
+except ImportError:
+ debug_logger = None
+
+
+# Re-export FirstTokenTimeoutError for backward compatibility
+__all__ = ['FirstTokenTimeoutError', 'stream_kiro_to_openai', 'stream_with_first_token_retry', 'collect_stream_response']
+
+
+async def stream_kiro_to_openai_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 for converting Kiro stream to OpenAI format.
+
+ Parses AWS SSE stream and converts events to OpenAI chat.completion.chunk.
+ Supports tool calls and usage calculation.
+
+ IMPORTANT: This function raises FirstTokenTimeoutError if first token
+ is not received within first_token_timeout seconds.
+
+ Args:
+ client: HTTP client (for connection management)
+ response: HTTP response with data stream
+ model: Model name to include in response
+ model_cache: Model cache for getting token limits
+ auth_manager: Authentication manager
+ first_token_timeout: First token wait timeout (seconds)
+ request_messages: Original request messages (for fallback token counting)
+ request_tools: Original request tools (for fallback token counting)
+
+ Yields:
+ Strings in SSE format: "data: {...}\\n\\n" or "data: [DONE]\\n\\n"
+
+ Raises:
+ FirstTokenTimeoutError: If first token not received within timeout
+
+ Example:
+ >>> async for chunk in stream_kiro_to_openai_internal(client, response, "claude-sonnet-4", cache, auth):
+ ... print(chunk)
+ data: {"id":"chatcmpl-...","object":"chat.completion.chunk",...}
+
+ data: [DONE]
+ """
+ completion_id = generate_completion_id()
+ created_time = int(time.time())
+ first_chunk = True
+
+ metering_data = None
+ context_usage_percentage = None
+ full_content = ""
+ full_thinking_content = "" # Accumulated thinking content for non-streaming
+
+ streaming_error_occurred = False
+ tool_calls_from_stream = []
+
+ try:
+ # Use streaming_core.parse_kiro_stream for unified event parsing
+ # This handles AWS SSE parsing, first token timeout, and thinking parser
+ async for event in parse_kiro_stream(response, first_token_timeout):
+ if event.type == "content" and event.content:
+ # Accumulate content for bracket tool call detection
+ full_content += event.content
+
+ # Format as OpenAI chunk
+ delta = {"content": event.content}
+ if first_chunk:
+ delta["role"] = "assistant"
+ first_chunk = False
+
+ openai_chunk = {
+ "id": completion_id,
+ "object": "chat.completion.chunk",
+ "created": created_time,
+ "model": model,
+ "choices": [{"index": 0, "delta": delta, "finish_reason": None}]
+ }
+
+ chunk_text = f"data: {json.dumps(openai_chunk, ensure_ascii=False)}\n\n"
+
+ if debug_logger:
+ debug_logger.log_modified_chunk(chunk_text.encode('utf-8'))
+
+ yield chunk_text
+
+ elif event.type == "thinking" and event.thinking_content:
+ # Accumulate thinking content
+ full_thinking_content += event.thinking_content
+
+ # Send as reasoning_content or content based on mode
+ if FAKE_REASONING_HANDLING == "as_reasoning_content":
+ delta = {"reasoning_content": event.thinking_content}
+ else:
+ delta = {"content": event.thinking_content}
+
+ if first_chunk:
+ delta["role"] = "assistant"
+ first_chunk = False
+
+ openai_chunk = {
+ "id": completion_id,
+ "object": "chat.completion.chunk",
+ "created": created_time,
+ "model": model,
+ "choices": [{"index": 0, "delta": delta, "finish_reason": None}]
+ }
+
+ chunk_text = f"data: {json.dumps(openai_chunk, ensure_ascii=False)}\n\n"
+
+ if debug_logger:
+ debug_logger.log_modified_chunk(chunk_text.encode('utf-8'))
+
+ yield chunk_text
+
+ elif event.type == "tool_use" and event.tool_use:
+ # Collect tool calls from stream
+ 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
+
+ # 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 = deduplicate_tool_calls(all_tool_calls)
+
+ # Determine finish_reason
+ finish_reason = "tool_calls" if all_tool_calls else "stop"
+
+ # Count completion_tokens (output) using tiktoken
+ completion_tokens = count_tokens(full_content + full_thinking_content)
+
+ # Calculate total_tokens based on context_usage_percentage from Kiro API
+ # context_usage shows TOTAL percentage of context usage (input + output)
+ prompt_tokens, total_tokens, prompt_source, total_source = calculate_tokens_from_context_usage(
+ context_usage_percentage, completion_tokens, model_cache, model
+ )
+
+ # Fallback: Kiro API didn't return context_usage, use tiktoken
+ # Count prompt_tokens from original messages
+ # IMPORTANT: Don't apply correction coefficient for prompt_tokens,
+ # as it was calibrated for completion_tokens
+ 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
+ prompt_source = "tiktoken"
+ total_source = "tiktoken"
+
+ # Send tool calls if present
+ if all_tool_calls:
+ logger.debug(f"Processing {len(all_tool_calls)} tool calls for streaming response")
+
+ # Add required index field to each tool_call
+ # according to OpenAI API specification for streaming
+ indexed_tool_calls = []
+ for idx, tc in enumerate(all_tool_calls):
+ # Extract function with None protection
+ func = tc.get("function") or {}
+ # Use "or" for protection against explicit None in values
+ tool_name = func.get("name") or ""
+ tool_args = func.get("arguments") or "{}"
+
+ logger.debug(f"Tool call [{idx}] '{tool_name}': id={tc.get('id')}, args_length={len(tool_args)}")
+
+ indexed_tc = {
+ "index": idx,
+ "id": tc.get("id"),
+ "type": tc.get("type", "function"),
+ "function": {
+ "name": tool_name,
+ "arguments": tool_args
+ }
+ }
+ indexed_tool_calls.append(indexed_tc)
+
+ tool_calls_chunk = {
+ "id": completion_id,
+ "object": "chat.completion.chunk",
+ "created": created_time,
+ "model": model,
+ "choices": [{
+ "index": 0,
+ "delta": {"tool_calls": indexed_tool_calls},
+ "finish_reason": None
+ }]
+ }
+ yield f"data: {json.dumps(tool_calls_chunk, ensure_ascii=False)}\n\n"
+
+ # Final chunk with usage
+ final_chunk = {
+ "id": completion_id,
+ "object": "chat.completion.chunk",
+ "created": created_time,
+ "model": model,
+ "choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}],
+ "usage": {
+ "prompt_tokens": prompt_tokens,
+ "completion_tokens": completion_tokens,
+ "total_tokens": total_tokens,
+ }
+ }
+
+ if metering_data:
+ final_chunk["usage"]["credits_used"] = metering_data
+
+ # Log final token values being sent to client
+ logger.debug(
+ f"[Usage] {model}: "
+ f"prompt_tokens={prompt_tokens} ({prompt_source}), "
+ f"completion_tokens={completion_tokens} (tiktoken), "
+ f"total_tokens={total_tokens} ({total_source})"
+ )
+
+ yield f"data: {json.dumps(final_chunk, ensure_ascii=False)}\n\n"
+ yield "data: [DONE]\n\n"
+
+ except FirstTokenTimeoutError:
+ # Propagate timeout up for retry
+ raise
+ except GeneratorExit:
+ # Client disconnected - this is normal, don't log as error
+ logger.debug("Client disconnected (GeneratorExit)")
+ streaming_error_occurred = True
+ except Exception as e:
+ streaming_error_occurred = True
+ # Log exception type and message for better diagnostics
+ error_type = type(e).__name__
+ error_msg = str(e) if str(e) else "(empty message)"
+ logger.error(
+ f"Error during streaming: [{error_type}] {error_msg}",
+ exc_info=True
+ )
+ # Propagate error up for proper handling in routes_openai.py
+ raise
+ finally:
+ # Always close response
+ try:
+ await response.aclose()
+ except Exception as close_error:
+ logger.debug(f"Error closing response: {close_error}")
+
+ if streaming_error_occurred:
+ logger.debug("Streaming completed with error")
+ else:
+ logger.debug("Streaming completed successfully")
+
+
+async def stream_kiro_to_openai(
+ 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]:
+ """
+ Generator for converting Kiro stream to OpenAI format.
+
+ This is a wrapper over stream_kiro_to_openai_internal that does NOT retry.
+ Retry logic is implemented in stream_with_first_token_retry.
+
+ Args:
+ client: HTTP client (for connection management)
+ response: HTTP response with data stream
+ model: Model name to include in response
+ model_cache: Model cache for getting token limits
+ auth_manager: Authentication manager
+ request_messages: Original request messages (for fallback token counting)
+ request_tools: Original request tools (for fallback token counting)
+
+ Yields:
+ Strings in SSE format: "data: {...}\\n\\n" or "data: [DONE]\\n\\n"
+ """
+ async for chunk in stream_kiro_to_openai_internal(
+ client, response, model, model_cache, auth_manager,
+ request_messages=request_messages,
+ request_tools=request_tools
+ ):
+ yield chunk
+
+
+async def stream_with_first_token_retry(
+ make_request: Callable[[], Awaitable[httpx.Response]],
+ client: httpx.AsyncClient,
+ model: str,
+ model_cache: "ModelInfoCache",
+ auth_manager: "KiroAuthManager",
+ 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.
+
+ If model doesn't respond within first_token_timeout seconds,
+ request is cancelled and a new one is made. Maximum max_retries attempts.
+
+ This is seamless for user - they just see a delay,
+ but eventually get a response (or error after all attempts).
+
+ Uses generic stream_with_first_token_retry from streaming_core.py.
+
+ Args:
+ make_request: Function to create new HTTP request
+ client: HTTP client
+ model: Model name
+ model_cache: Model cache
+ auth_manager: Authentication manager
+ max_retries: Maximum number of attempts
+ first_token_timeout: First token wait timeout (seconds)
+ request_messages: Original request messages (for fallback token counting)
+ request_tools: Original request tools (for fallback token counting)
+
+ Yields:
+ Strings in SSE format
+
+ Raises:
+ HTTPException: After exhausting all attempts
+
+ Example:
+ >>> async def make_req():
+ ... return await http_client.request_with_retry("POST", url, payload, stream=True)
+ >>> async for chunk in stream_with_first_token_retry(make_req, client, model, cache, auth):
+ ... print(chunk)
+ """
+ def create_http_error(status_code: int, error_text: str) -> HTTPException:
+ """Create HTTPException for HTTP errors."""
+ return HTTPException(
+ status_code=status_code,
+ detail=f"Upstream API error: {error_text}"
+ )
+
+ def create_timeout_error(retries: int, timeout: float) -> HTTPException:
+ """Create HTTPException for timeout errors."""
+ return HTTPException(
+ status_code=504,
+ detail=f"Model did not respond within {timeout}s after {retries} attempts. Please try again."
+ )
+
+ async def stream_processor(response: httpx.Response) -> AsyncGenerator[str, None]:
+ """Process response and yield OpenAI SSE chunks."""
+ async for chunk in stream_kiro_to_openai_internal(
+ client,
+ response,
+ 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,
+ 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_stream_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 full response from streaming stream.
+
+ Used for non-streaming mode - collects all chunks
+ and forms a single response.
+
+ Args:
+ client: HTTP client
+ response: HTTP response with stream
+ model: Model name
+ model_cache: Model cache
+ auth_manager: Authentication manager
+ request_messages: Original request messages (for fallback token counting)
+ request_tools: Original request tools (for fallback token counting)
+
+ Returns:
+ Dictionary with full response in OpenAI chat.completion format
+ """
+ full_content = ""
+ full_reasoning_content = ""
+ final_usage = None
+ tool_calls = []
+ completion_id = generate_completion_id()
+
+ async for chunk_str in stream_kiro_to_openai(
+ client,
+ response,
+ model,
+ model_cache,
+ auth_manager,
+ request_messages=request_messages,
+ request_tools=request_tools
+ ):
+ if not chunk_str.startswith("data:"):
+ continue
+
+ data_str = chunk_str[len("data:"):].strip()
+ if not data_str or data_str == "[DONE]":
+ continue
+
+ try:
+ chunk_data = json.loads(data_str)
+
+ # Extract data from chunk
+ delta = chunk_data.get("choices", [{}])[0].get("delta", {})
+ if "content" in delta:
+ full_content += delta["content"]
+ if "reasoning_content" in delta:
+ full_reasoning_content += delta["reasoning_content"]
+ if "tool_calls" in delta:
+ tool_calls.extend(delta["tool_calls"])
+
+ # Save usage from last chunk
+ if "usage" in chunk_data:
+ final_usage = chunk_data["usage"]
+
+ except (json.JSONDecodeError, IndexError):
+ continue
+
+ # Form final response
+ message = {"role": "assistant", "content": full_content}
+ if full_reasoning_content:
+ message["reasoning_content"] = full_reasoning_content
+ if tool_calls:
+ # For non-streaming response remove index field from tool_calls,
+ # as it's only required for streaming chunks
+ cleaned_tool_calls = []
+ for tc in tool_calls:
+ # Extract function with None protection
+ func = tc.get("function") or {}
+ cleaned_tc = {
+ "id": tc.get("id"),
+ "type": tc.get("type", "function"),
+ "function": {
+ "name": func.get("name", ""),
+ "arguments": func.get("arguments", "{}")
+ }
+ }
+ cleaned_tool_calls.append(cleaned_tc)
+ message["tool_calls"] = cleaned_tool_calls
+
+ finish_reason = "tool_calls" if tool_calls else "stop"
+
+ # Form usage for response
+ usage = final_usage or {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+
+ # Log token info for debugging (non-streaming uses same logs from streaming)
+
+ return {
+ "id": completion_id,
+ "object": "chat.completion",
+ "created": int(time.time()),
+ "model": model,
+ "choices": [{
+ "index": 0,
+ "message": message,
+ "finish_reason": finish_reason
+ }],
+ "usage": usage
+ }
\ No newline at end of file
diff --git a/kiro-gateway/kiro/thinking_parser.py b/kiro-gateway/kiro/thinking_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..3610244cd8fd093ac65d4e7a84b93593bd6d4cc5
--- /dev/null
+++ b/kiro-gateway/kiro/thinking_parser.py
@@ -0,0 +1,385 @@
+# -*- 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 .
+
+"""
+Thinking block parser for streaming responses.
+
+Implements a finite state machine (FSM) for reliable parsing of thinking blocks
+(, , , etc.) that may be split across multiple
+network chunks.
+
+Key features:
+- Tag detection ONLY at the start of response
+- "Cautious" sending - buffers potential tag fragments to avoid splitting tags
+- After closing tag - all content is treated as regular content
+- Support for multiple tag formats
+"""
+
+from enum import IntEnum
+from typing import Optional, List
+from dataclasses import dataclass, field
+
+from loguru import logger
+
+from kiro.config import (
+ FAKE_REASONING_HANDLING,
+ FAKE_REASONING_OPEN_TAGS,
+ FAKE_REASONING_INITIAL_BUFFER_SIZE,
+)
+
+
+class ParserState(IntEnum):
+ """
+ States of the thinking block parser FSM.
+
+ PRE_CONTENT: Initial state, buffering to detect opening tag
+ IN_THINKING: Inside thinking block, buffering until closing tag
+ STREAMING: Regular streaming, no more thinking block detection
+ """
+ PRE_CONTENT = 0
+ IN_THINKING = 1
+ STREAMING = 2
+
+
+@dataclass
+class ThinkingParseResult:
+ """
+ Result of processing a content chunk through the parser.
+
+ Attributes:
+ thinking_content: Content to be sent as reasoning_content (or processed per mode)
+ regular_content: Regular content to be sent as delta.content
+ is_first_thinking_chunk: True if this is the first chunk of thinking content
+ is_last_thinking_chunk: True if thinking block just closed
+ state_changed: True if parser state changed during this feed
+ """
+ thinking_content: Optional[str] = None
+ regular_content: Optional[str] = None
+ is_first_thinking_chunk: bool = False
+ is_last_thinking_chunk: bool = False
+ state_changed: bool = False
+
+
+class ThinkingParser:
+ """
+ Finite state machine parser for thinking blocks in streaming responses.
+
+ The parser detects thinking tags ONLY at the start of the response.
+ Once a thinking block is found and closed, all subsequent content
+ is treated as regular content (even if it contains thinking tags).
+
+ This implements "cautious" buffering to handle tags split across chunks:
+ - In PRE_CONTENT: buffer until tag found or buffer exceeds limit
+ - In IN_THINKING: buffer last MAX_TAG_LENGTH chars to avoid splitting closing tag
+
+ Example:
+ >>> parser = ThinkingParser()
+ >>> result = parser.feed(">> result.thinking_content # None - still buffering
+ >>> result = parser.feed("ing>Hello")
+ >>> result.thinking_content # "Hello" (or None if buffering)
+ >>> result = parser.feed("World")
+ >>> result.thinking_content # remaining thinking content
+ >>> result.regular_content # "World"
+ """
+
+ def __init__(
+ self,
+ handling_mode: Optional[str] = None,
+ open_tags: Optional[List[str]] = None,
+ initial_buffer_size: int = FAKE_REASONING_INITIAL_BUFFER_SIZE,
+ ):
+ """
+ Initialize the thinking parser.
+
+ Args:
+ handling_mode: How to handle thinking blocks. One of:
+ - "as_reasoning_content": Extract to reasoning_content field
+ - "remove": Remove thinking block completely
+ - "pass": Pass through with original tags
+ - "strip_tags": Remove tags but keep content
+ If None, uses FAKE_REASONING_HANDLING from config.
+ open_tags: List of opening tags to detect. If None, uses config.
+ initial_buffer_size: Max chars to buffer while looking for opening tag.
+ """
+ self.handling_mode = handling_mode or FAKE_REASONING_HANDLING
+ self.open_tags = open_tags or FAKE_REASONING_OPEN_TAGS
+ self.initial_buffer_size = initial_buffer_size
+
+ # Calculate max tag length for cautious buffering
+ # We need to buffer enough to not split a closing tag
+ self.max_tag_length = max(len(tag) for tag in self.open_tags) * 2
+
+ # State
+ self.state = ParserState.PRE_CONTENT
+ self.initial_buffer = ""
+ self.thinking_buffer = ""
+ self.open_tag: Optional[str] = None
+ self.close_tag: Optional[str] = None
+ self.is_first_thinking_chunk = True
+ self._thinking_block_found = False
+
+ def feed(self, content: str) -> ThinkingParseResult:
+ """
+ Process a chunk of content through the parser.
+
+ Args:
+ content: New content from delta.content
+
+ Returns:
+ ThinkingParseResult with processed content
+ """
+ result = ThinkingParseResult()
+
+ if not content:
+ return result
+
+ # Handle based on current state
+ if self.state == ParserState.PRE_CONTENT:
+ result = self._handle_pre_content(content)
+
+ # If state changed to IN_THINKING, process remaining content
+ if self.state == ParserState.IN_THINKING and result.state_changed:
+ # Content after tag is already in thinking_buffer from _handle_pre_content
+ pass
+ elif self.state == ParserState.IN_THINKING and not result.state_changed:
+ result = self._handle_in_thinking(content)
+
+ # If state changed to STREAMING, regular_content is already set
+ if self.state == ParserState.STREAMING and not result.state_changed:
+ result.regular_content = content
+
+ return result
+
+ def _handle_pre_content(self, content: str) -> ThinkingParseResult:
+ """
+ Handle content in PRE_CONTENT state.
+
+ Buffers content and looks for opening tag at the start.
+ """
+ result = ThinkingParseResult()
+ self.initial_buffer += content
+
+ # Strip leading whitespace for tag detection
+ stripped = self.initial_buffer.lstrip()
+
+ # Check if buffer starts with any of the opening tags
+ for tag in self.open_tags:
+ if stripped.startswith(tag):
+ # Tag found! Transition to IN_THINKING
+ self.state = ParserState.IN_THINKING
+ self.open_tag = tag
+ self.close_tag = f" ->
+ self._thinking_block_found = True
+ result.state_changed = True
+
+ logger.debug(f"Thinking tag '{tag}' detected. Transitioning to IN_THINKING.")
+
+ # Content after the tag goes to thinking buffer
+ content_after_tag = stripped[len(tag):]
+ self.thinking_buffer = content_after_tag
+ self.initial_buffer = ""
+
+ # Now process the thinking buffer for potential closing tag
+ thinking_result = self._process_thinking_buffer()
+ if thinking_result.thinking_content:
+ result.thinking_content = thinking_result.thinking_content
+ result.is_first_thinking_chunk = thinking_result.is_first_thinking_chunk
+ if thinking_result.is_last_thinking_chunk:
+ result.is_last_thinking_chunk = True
+ if thinking_result.regular_content:
+ result.regular_content = thinking_result.regular_content
+
+ return result
+
+ # Check if we might still be receiving the tag
+ # (buffer is shorter than longest tag and could be a prefix)
+ for tag in self.open_tags:
+ if tag.startswith(stripped) and len(stripped) < len(tag):
+ # Could still be receiving the tag, keep buffering
+ return result
+
+ # No tag found and buffer is either:
+ # 1. Too long (exceeds initial_buffer_size)
+ # 2. Doesn't match any tag prefix
+ if len(self.initial_buffer) > self.initial_buffer_size or not self._could_be_tag_prefix(stripped):
+ # No thinking block, transition to STREAMING
+ self.state = ParserState.STREAMING
+ result.state_changed = True
+ result.regular_content = self.initial_buffer
+ self.initial_buffer = ""
+
+ logger.debug("No thinking tag detected. Transitioning to STREAMING.")
+
+ return result
+
+ def _could_be_tag_prefix(self, text: str) -> bool:
+ """Check if text could be the start of any opening tag."""
+ if not text:
+ return True # Empty could be anything
+
+ for tag in self.open_tags:
+ if tag.startswith(text):
+ return True
+ return False
+
+ def _handle_in_thinking(self, content: str) -> ThinkingParseResult:
+ """
+ Handle content in IN_THINKING state.
+
+ Buffers content and looks for closing tag.
+ Uses "cautious" sending to avoid splitting the closing tag.
+ """
+ self.thinking_buffer += content
+ return self._process_thinking_buffer()
+
+ def _process_thinking_buffer(self) -> ThinkingParseResult:
+ """
+ Process the thinking buffer, looking for closing tag.
+
+ Implements "cautious" sending - keeps last max_tag_length chars
+ in buffer to avoid splitting the closing tag across chunks.
+ """
+ result = ThinkingParseResult()
+
+ if not self.close_tag:
+ return result
+
+ # Check for closing tag
+ if self.close_tag in self.thinking_buffer:
+ # Found closing tag!
+ idx = self.thinking_buffer.find(self.close_tag)
+ thinking_content = self.thinking_buffer[:idx]
+ after_tag = self.thinking_buffer[idx + len(self.close_tag):]
+
+ # Send all thinking content
+ if thinking_content:
+ result.thinking_content = thinking_content
+ result.is_first_thinking_chunk = self.is_first_thinking_chunk
+ self.is_first_thinking_chunk = False
+
+ result.is_last_thinking_chunk = True
+
+ # Transition to STREAMING
+ self.state = ParserState.STREAMING
+ result.state_changed = True
+ self.thinking_buffer = ""
+
+ logger.debug(f"Closing tag '{self.close_tag}' found. Transitioning to STREAMING.")
+
+ # Content after closing tag is regular content
+ # Strip leading whitespace/newlines that often follow the closing tag
+ if after_tag:
+ stripped_after = after_tag.lstrip()
+ if stripped_after:
+ result.regular_content = stripped_after
+
+ return result
+
+ # No closing tag yet - use "cautious" sending
+ # Keep last max_tag_length chars in buffer to avoid splitting tag
+ if len(self.thinking_buffer) > self.max_tag_length:
+ send_part = self.thinking_buffer[:-self.max_tag_length]
+ self.thinking_buffer = self.thinking_buffer[-self.max_tag_length:]
+
+ result.thinking_content = send_part
+ result.is_first_thinking_chunk = self.is_first_thinking_chunk
+ self.is_first_thinking_chunk = False
+
+ return result
+
+ def finalize(self) -> ThinkingParseResult:
+ """
+ Finalize parsing when stream ends.
+
+ Flushes any remaining buffered content.
+
+ Returns:
+ ThinkingParseResult with any remaining content
+ """
+ result = ThinkingParseResult()
+
+ # Flush thinking buffer if we're still in thinking state
+ if self.thinking_buffer:
+ if self.state == ParserState.IN_THINKING:
+ result.thinking_content = self.thinking_buffer
+ result.is_first_thinking_chunk = self.is_first_thinking_chunk
+ result.is_last_thinking_chunk = True
+ logger.warning("Stream ended while still in thinking block. Flushing remaining content.")
+ else:
+ result.regular_content = self.thinking_buffer
+ self.thinking_buffer = ""
+
+ # Flush initial buffer if we never found a tag
+ if self.initial_buffer:
+ result.regular_content = (result.regular_content or "") + self.initial_buffer
+ self.initial_buffer = ""
+
+ return result
+
+ def reset(self) -> None:
+ """Reset parser to initial state."""
+ self.state = ParserState.PRE_CONTENT
+ self.initial_buffer = ""
+ self.thinking_buffer = ""
+ self.open_tag = None
+ self.close_tag = None
+ self.is_first_thinking_chunk = True
+ self._thinking_block_found = False
+
+ @property
+ def found_thinking_block(self) -> bool:
+ """Returns True if a thinking block was detected in this response."""
+ return self._thinking_block_found
+
+ def process_for_output(
+ self,
+ thinking_content: Optional[str],
+ is_first: bool,
+ is_last: bool,
+ ) -> Optional[str]:
+ """
+ Process thinking content according to handling mode.
+
+ Args:
+ thinking_content: Raw thinking content
+ is_first: True if this is the first thinking chunk
+ is_last: True if this is the last thinking chunk
+
+ Returns:
+ Processed content string or None (for "remove" mode)
+ """
+ if not thinking_content:
+ return None
+
+ if self.handling_mode == "remove":
+ return None
+
+ if self.handling_mode == "pass":
+ # Add tags back
+ prefix = self.open_tag if is_first and self.open_tag else ""
+ suffix = self.close_tag if is_last and self.close_tag else ""
+ return f"{prefix}{thinking_content}{suffix}"
+
+ if self.handling_mode == "strip_tags":
+ # Return content without tags
+ return thinking_content
+
+ # "as_reasoning_content" - return as-is, caller will put in reasoning_content field
+ return thinking_content
\ No newline at end of file
diff --git a/kiro-gateway/kiro/tokenizer.py b/kiro-gateway/kiro/tokenizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..7450f6f293a9e3de6b3e772877255908ee145a03
--- /dev/null
+++ b/kiro-gateway/kiro/tokenizer.py
@@ -0,0 +1,245 @@
+# -*- 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 .
+
+"""
+Module for fast token counting.
+
+Uses tiktoken (OpenAI's Rust library) for approximate
+token counting. The cl100k_base encoding is close to Claude tokenization.
+
+Note: This is an approximate count, as the exact Claude tokenizer
+is not public. Anthropic does not publish their tokenizer,
+so tiktoken with a correction coefficient is used.
+
+The correction coefficient CLAUDE_CORRECTION_FACTOR = 1.15 is based on
+empirical observations: Claude tokenizes text approximately 15%
+more than GPT-4 (cl100k_base). This is due to differences in BPE vocabularies.
+"""
+
+from typing import List, Dict, Any, Optional
+from loguru import logger
+
+# Lazy loading of tiktoken to speed up import
+_encoding = None
+
+# Correction coefficient for Claude models
+# Claude tokenizes text approximately 15% more than GPT-4 (cl100k_base)
+# This is an empirical value based on comparison with context_usage from API
+CLAUDE_CORRECTION_FACTOR = 1.15
+
+
+def _get_encoding():
+ """
+ Lazy initialization of tokenizer.
+
+ Uses cl100k_base - encoding for GPT-4/ChatGPT,
+ which is close enough to Claude tokenization.
+
+ Returns:
+ tiktoken.Encoding or None if tiktoken is unavailable
+ """
+ global _encoding
+ if _encoding is None:
+ try:
+ import tiktoken
+ _encoding = tiktoken.get_encoding("cl100k_base")
+ logger.debug("[Tokenizer] Initialized tiktoken with cl100k_base encoding")
+ except ImportError:
+ logger.warning(
+ "[Tokenizer] tiktoken not installed. "
+ "Token counting will use fallback estimation. "
+ "Install with: pip install tiktoken"
+ )
+ _encoding = False # Marker that import failed
+ except Exception as e:
+ logger.error(f"[Tokenizer] Failed to initialize tiktoken: {e}")
+ _encoding = False
+ return _encoding if _encoding else None
+
+
+def count_tokens(text: str, apply_claude_correction: bool = True) -> int:
+ """
+ Counts the number of tokens in text.
+
+ Args:
+ text: Text to count tokens for
+ apply_claude_correction: Apply correction coefficient for Claude (default True)
+
+ Returns:
+ Number of tokens (approximate, with Claude correction)
+ """
+ if not text:
+ return 0
+
+ encoding = _get_encoding()
+ if encoding:
+ try:
+ base_tokens = len(encoding.encode(text))
+ if apply_claude_correction:
+ return int(base_tokens * CLAUDE_CORRECTION_FACTOR)
+ return base_tokens
+ except Exception as e:
+ logger.warning(f"[Tokenizer] Error encoding text: {e}")
+
+ # Fallback: rough estimate ~4 characters per token for English,
+ # ~2-3 characters for other languages (taking average ~3.5)
+ # For Claude we add correction
+ base_estimate = len(text) // 4 + 1
+ if apply_claude_correction:
+ return int(base_estimate * CLAUDE_CORRECTION_FACTOR)
+ return base_estimate
+
+
+def count_message_tokens(messages: List[Dict[str, Any]], apply_claude_correction: bool = True) -> int:
+ """
+ Counts tokens in a list of chat messages.
+
+ Accounts for OpenAI/Claude message structure:
+ - role: ~1 token
+ - content: text tokens
+ - Service tokens between messages: ~3-4 tokens
+
+ Args:
+ messages: List of messages in OpenAI format
+ apply_claude_correction: Apply correction coefficient for Claude
+
+ Returns:
+ Approximate number of tokens (with Claude correction)
+ """
+ if not messages:
+ return 0
+
+ total_tokens = 0
+
+ for message in messages:
+ # Base tokens per message (role, delimiters)
+ total_tokens += 4 # ~4 tokens for service information
+
+ # Role tokens (without correction, these are short strings)
+ role = message.get("role", "")
+ total_tokens += count_tokens(role, apply_claude_correction=False)
+
+ # Content tokens
+ content = message.get("content")
+ if content:
+ if isinstance(content, str):
+ total_tokens += count_tokens(content, apply_claude_correction=False)
+ elif isinstance(content, list):
+ # Multimodal content (text + images)
+ for item in content:
+ if isinstance(item, dict):
+ if item.get("type") == "text":
+ total_tokens += count_tokens(item.get("text", ""), apply_claude_correction=False)
+ elif item.get("type") == "image_url":
+ # Images take ~85-170 tokens depending on size
+ total_tokens += 100 # Average estimate
+
+ # tool_calls tokens (if present)
+ tool_calls = message.get("tool_calls")
+ if tool_calls:
+ for tc in tool_calls:
+ total_tokens += 4 # Service tokens
+ func = tc.get("function", {})
+ total_tokens += count_tokens(func.get("name", ""), apply_claude_correction=False)
+ total_tokens += count_tokens(func.get("arguments", ""), apply_claude_correction=False)
+
+ # tool_call_id tokens (for tool responses)
+ if message.get("tool_call_id"):
+ total_tokens += count_tokens(message["tool_call_id"], apply_claude_correction=False)
+
+ # Final service tokens
+ total_tokens += 3
+
+ # Apply correction to total count
+ if apply_claude_correction:
+ return int(total_tokens * CLAUDE_CORRECTION_FACTOR)
+ return total_tokens
+
+
+def count_tools_tokens(tools: Optional[List[Dict[str, Any]]], apply_claude_correction: bool = True) -> int:
+ """
+ Counts tokens in tool definitions.
+
+ Args:
+ tools: List of tools in OpenAI format
+ apply_claude_correction: Apply correction coefficient for Claude
+
+ Returns:
+ Approximate number of tokens (with Claude correction)
+ """
+ if not tools:
+ return 0
+
+ total_tokens = 0
+
+ for tool in tools:
+ total_tokens += 4 # Service tokens
+
+ if tool.get("type") == "function":
+ func = tool.get("function", {})
+
+ # Function name
+ total_tokens += count_tokens(func.get("name", ""), apply_claude_correction=False)
+
+ # Function description
+ total_tokens += count_tokens(func.get("description", ""), apply_claude_correction=False)
+
+ # Parameters (JSON schema)
+ params = func.get("parameters")
+ if params:
+ import json
+ params_str = json.dumps(params, ensure_ascii=False)
+ total_tokens += count_tokens(params_str, apply_claude_correction=False)
+
+ # Apply correction to total count
+ if apply_claude_correction:
+ return int(total_tokens * CLAUDE_CORRECTION_FACTOR)
+ return total_tokens
+
+
+def estimate_request_tokens(
+ messages: List[Dict[str, Any]],
+ tools: Optional[List[Dict[str, Any]]] = None,
+ system_prompt: Optional[str] = None
+) -> Dict[str, int]:
+ """
+ Estimates total number of tokens in request.
+
+ Args:
+ messages: List of messages
+ tools: List of tools (optional)
+ system_prompt: System prompt (optional, if not in messages)
+
+ Returns:
+ Dictionary with token breakdown:
+ - messages_tokens: message tokens
+ - tools_tokens: tool tokens
+ - system_tokens: system prompt tokens
+ - total_tokens: total count
+ """
+ messages_tokens = count_message_tokens(messages)
+ tools_tokens = count_tools_tokens(tools)
+ system_tokens = count_tokens(system_prompt) if system_prompt else 0
+
+ return {
+ "messages_tokens": messages_tokens,
+ "tools_tokens": tools_tokens,
+ "system_tokens": system_tokens,
+ "total_tokens": messages_tokens + tools_tokens + system_tokens
+ }
\ No newline at end of file
diff --git a/kiro-gateway/kiro/utils.py b/kiro-gateway/kiro/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..fb2c18f220a42482bbe608a873edde71861dac3b
--- /dev/null
+++ b/kiro-gateway/kiro/utils.py
@@ -0,0 +1,117 @@
+# -*- 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 .
+
+"""
+Utility functions for Kiro Gateway.
+
+Contains functions for fingerprint generation, header formatting,
+and other common utilities.
+"""
+
+import hashlib
+import uuid
+from typing import TYPE_CHECKING
+
+from loguru import logger
+
+if TYPE_CHECKING:
+ from kiro.auth import KiroAuthManager
+
+
+def get_machine_fingerprint() -> str:
+ """
+ Generates a unique machine fingerprint based on hostname and username.
+
+ Used for User-Agent formation to identify a specific gateway installation.
+
+ Returns:
+ SHA256 hash of the string "{hostname}-{username}-kiro-gateway"
+ """
+ try:
+ import socket
+ import getpass
+
+ hostname = socket.gethostname()
+ username = getpass.getuser()
+ unique_string = f"{hostname}-{username}-kiro-gateway"
+
+ return hashlib.sha256(unique_string.encode()).hexdigest()
+ except Exception as e:
+ logger.warning(f"Failed to get machine fingerprint: {e}")
+ return hashlib.sha256(b"default-kiro-gateway").hexdigest()
+
+
+def get_kiro_headers(auth_manager: "KiroAuthManager", token: str) -> dict:
+ """
+ Builds headers for Kiro API requests.
+
+ Includes all necessary headers for authentication and identification:
+ - Authorization with Bearer token
+ - User-Agent with fingerprint
+ - AWS CodeWhisperer specific headers
+
+ Args:
+ auth_manager: Authentication manager for obtaining fingerprint
+ token: Access token for authorization
+
+ Returns:
+ Dictionary with headers for HTTP request
+ """
+ fingerprint = auth_manager.fingerprint
+
+ return {
+ "Authorization": f"Bearer {token}",
+ "Content-Type": "application/json",
+ "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}",
+ "x-amzn-codewhisperer-optout": "true",
+ "x-amzn-kiro-agent-mode": "vibe",
+ "amz-sdk-invocation-id": str(uuid.uuid4()),
+ "amz-sdk-request": "attempt=1; max=3",
+ }
+
+
+def generate_completion_id() -> str:
+ """
+ Generates a unique ID for chat completion.
+
+ Returns:
+ ID in format "chatcmpl-{uuid_hex}"
+ """
+ return f"chatcmpl-{uuid.uuid4().hex}"
+
+
+def generate_conversation_id() -> str:
+ """
+ Generates a unique ID for conversation.
+
+ Returns:
+ UUID in string format
+ """
+ return str(uuid.uuid4())
+
+
+def generate_tool_call_id() -> str:
+ """
+ Generates a unique ID for tool call.
+
+ Returns:
+ ID in format "call_{uuid_hex[:8]}"
+ """
+ return f"call_{uuid.uuid4().hex[:8]}"
\ No newline at end of file
diff --git a/kiro-gateway/main.py b/kiro-gateway/main.py
new file mode 100644
index 0000000000000000000000000000000000000000..549b9023ae461198d46467f97b997efa97df179d
--- /dev/null
+++ b/kiro-gateway/main.py
@@ -0,0 +1,637 @@
+# -*- 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 .
+
+"""
+Kiro Gateway - OpenAI-compatible interface for Kiro API.
+
+Application entry point. Creates FastAPI app and connects routes.
+
+Usage:
+ # Using default settings (host: 0.0.0.0, port: 8000)
+ python main.py
+
+ # With CLI arguments (highest priority)
+ python main.py --port 9000
+ python main.py --host 127.0.0.1 --port 9000
+
+ # With environment variables (medium priority)
+ SERVER_PORT=9000 python main.py
+
+ # Using uvicorn directly (uvicorn handles its own CLI args)
+ uvicorn main:app --host 0.0.0.0 --port 8000
+
+Priority: CLI args > Environment variables > Default values
+"""
+
+import argparse
+import logging
+import sys
+import os
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+import httpx
+from fastapi import FastAPI
+from fastapi.exceptions import RequestValidationError
+from fastapi.middleware.cors import CORSMiddleware
+from loguru import logger
+
+from kiro.config import (
+ APP_TITLE,
+ APP_DESCRIPTION,
+ APP_VERSION,
+ REFRESH_TOKEN,
+ PROFILE_ARN,
+ REGION,
+ KIRO_CREDS_FILE,
+ KIRO_CLI_DB_FILE,
+ PROXY_API_KEY,
+ LOG_LEVEL,
+ SERVER_HOST,
+ SERVER_PORT,
+ DEFAULT_SERVER_HOST,
+ DEFAULT_SERVER_PORT,
+ STREAMING_READ_TIMEOUT,
+ HIDDEN_MODELS,
+ FALLBACK_MODELS,
+ VPN_PROXY_URL,
+ _warn_deprecated_debug_setting,
+ _warn_timeout_configuration,
+)
+from kiro.auth import KiroAuthManager
+from kiro.cache import ModelInfoCache
+from kiro.model_resolver import ModelResolver
+from kiro.routes_openai import router as openai_router
+from kiro.routes_anthropic import router as anthropic_router
+from kiro.exceptions import validation_exception_handler
+from kiro.debug_middleware import DebugLoggerMiddleware
+
+
+# --- Loguru Configuration ---
+logger.remove()
+logger.add(
+ sys.stderr,
+ level=LOG_LEVEL,
+ colorize=True,
+ format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}"
+)
+
+
+class InterceptHandler(logging.Handler):
+ """
+ Intercepts logs from standard logging and redirects them to loguru.
+
+ This allows capturing logs from uvicorn, FastAPI and other libraries
+ that use standard logging instead of loguru.
+
+ Also filters out noisy shutdown-related exceptions (CancelledError, KeyboardInterrupt)
+ that are normal during Ctrl+C but uvicorn logs as ERROR.
+ """
+
+ # Exceptions that are normal during shutdown and should not be logged as errors
+ SHUTDOWN_EXCEPTIONS = (
+ "CancelledError",
+ "KeyboardInterrupt",
+ "asyncio.exceptions.CancelledError",
+ )
+
+ def emit(self, record: logging.LogRecord) -> None:
+ # Filter out shutdown-related exceptions that uvicorn logs as ERROR
+ # These are normal during Ctrl+C and don't need to spam the console
+ if record.exc_info:
+ exc_type = record.exc_info[0]
+ if exc_type is not None:
+ exc_name = exc_type.__name__
+ if exc_name in self.SHUTDOWN_EXCEPTIONS:
+ # Suppress the full traceback, just log a simple message
+ logger.info("Server shutdown in progress...")
+ return
+
+ # Also filter by message content for cases where exc_info is not set
+ msg = record.getMessage()
+ if any(exc in msg for exc in self.SHUTDOWN_EXCEPTIONS):
+ return
+
+ # Get the corresponding loguru level
+ try:
+ level = logger.level(record.levelname).name
+ except ValueError:
+ level = record.levelno
+
+ # Find the caller frame for correct source display
+ frame, depth = logging.currentframe(), 2
+ while frame.f_code.co_filename == logging.__file__:
+ frame = frame.f_back
+ depth += 1
+
+ logger.opt(depth=depth, exception=record.exc_info).log(level, record.getMessage())
+
+
+def setup_logging_intercept():
+ """
+ Configures log interception from standard logging to loguru.
+
+ Intercepts logs from:
+ - uvicorn (access logs, error logs)
+ - uvicorn.error
+ - uvicorn.access
+ - fastapi
+ """
+ # List of loggers to intercept
+ loggers_to_intercept = [
+ "uvicorn",
+ "uvicorn.error",
+ "uvicorn.access",
+ "fastapi",
+ ]
+
+ for logger_name in loggers_to_intercept:
+ logging_logger = logging.getLogger(logger_name)
+ logging_logger.handlers = [InterceptHandler()]
+ logging_logger.propagate = False
+
+
+# Configure uvicorn/fastapi log interception
+setup_logging_intercept()
+
+
+# ==================================================================================================
+# VPN/Proxy Configuration
+# ==================================================================================================
+# Must be set BEFORE creating any httpx clients (including in lifespan)
+# httpx automatically picks up HTTP_PROXY, HTTPS_PROXY, ALL_PROXY from environment
+
+if VPN_PROXY_URL:
+ # Normalize URL - add http:// if no scheme specified
+ proxy_url_with_scheme = VPN_PROXY_URL if "://" in VPN_PROXY_URL else f"http://{VPN_PROXY_URL}"
+
+ # Set environment variables for httpx to pick up automatically
+ os.environ['HTTP_PROXY'] = proxy_url_with_scheme
+ os.environ['HTTPS_PROXY'] = proxy_url_with_scheme
+ os.environ['ALL_PROXY'] = proxy_url_with_scheme
+
+ # Exclude localhost from proxy to avoid routing local requests through it
+ no_proxy_hosts = os.environ.get("NO_PROXY", "")
+ local_hosts = "127.0.0.1,localhost"
+ if no_proxy_hosts:
+ os.environ["NO_PROXY"] = f"{no_proxy_hosts},{local_hosts}"
+ else:
+ os.environ["NO_PROXY"] = local_hosts
+
+ logger.info(f"Proxy configured: {proxy_url_with_scheme}")
+ logger.debug(f"NO_PROXY: {os.environ['NO_PROXY']}")
+
+
+# --- Configuration Validation ---
+def validate_configuration() -> None:
+ """
+ Validates that required configuration is present.
+
+ Checks:
+ - .env file exists
+ - Either REFRESH_TOKEN or KIRO_CREDS_FILE is configured
+
+ Raises:
+ SystemExit: If critical configuration is missing
+ """
+ errors = []
+
+ # Check if .env file exists
+ env_file = Path(".env")
+ env_example = Path(".env.example")
+
+ if not env_file.exists():
+ errors.append(
+ ".env file not found!\n"
+ "\n"
+ "To get started:\n"
+ "1. Create .env or rename from .env.example:\n"
+ " cp .env.example .env\n"
+ "\n"
+ "2. Edit .env and configure your credentials:\n"
+ " 2.1. Set you super-secret password as PROXY_API_KEY\n"
+ " 2.2. Set your Kiro credentials:\n"
+ " - 1 way: KIRO_CREDS_FILE to your Kiro credentials JSON file\n"
+ " - 2 way: REFRESH_TOKEN from Kiro IDE traffic\n"
+ "\n"
+ "See README.md for detailed instructions."
+ )
+ else:
+ # .env exists, check for credentials
+ has_refresh_token = bool(REFRESH_TOKEN)
+ has_creds_file = bool(KIRO_CREDS_FILE)
+ has_cli_db = bool(KIRO_CLI_DB_FILE)
+
+ # Check if creds file actually exists
+ if KIRO_CREDS_FILE:
+ creds_path = Path(KIRO_CREDS_FILE).expanduser()
+ if not creds_path.exists():
+ has_creds_file = False
+ logger.warning(f"KIRO_CREDS_FILE not found: {KIRO_CREDS_FILE}")
+
+ # Check if CLI database file actually exists
+ if KIRO_CLI_DB_FILE:
+ cli_db_path = Path(KIRO_CLI_DB_FILE).expanduser()
+ if not cli_db_path.exists():
+ has_cli_db = False
+ logger.warning(f"KIRO_CLI_DB_FILE not found: {KIRO_CLI_DB_FILE}")
+
+ if not has_refresh_token and not has_creds_file and not has_cli_db:
+ errors.append(
+ "No Kiro credentials configured!\n"
+ "\n"
+ " Configure one of the following in your .env file:\n"
+ "\n"
+ "Set you super-secret password as PROXY_API_KEY\n"
+ " PROXY_API_KEY=\"my-super-secret-password-123\"\n"
+ "\n"
+ " Option 1 (Recommended): JSON credentials file\n"
+ " KIRO_CREDS_FILE=\"path/to/your/kiro-credentials.json\"\n"
+ "\n"
+ " Option 2: Refresh token\n"
+ " REFRESH_TOKEN=\"your_refresh_token_here\"\n"
+ "\n"
+ " Option 3: kiro-cli SQLite database (AWS SSO)\n"
+ " KIRO_CLI_DB_FILE=\"~/.local/share/kiro-cli/data.sqlite3\"\n"
+ "\n"
+ " See README.md for how to obtain credentials."
+ )
+
+ # Print errors and exit if any
+ if errors:
+ logger.error("")
+ logger.error("=" * 60)
+ logger.error(" CONFIGURATION ERROR")
+ logger.error("=" * 60)
+ for error in errors:
+ for line in error.split('\n'):
+ logger.error(f" {line}")
+ logger.error("=" * 60)
+ logger.error("")
+ sys.exit(1)
+
+ # Note: Credential loading details are logged by KiroAuthManager
+
+
+# Run configuration validation on import
+validate_configuration()
+
+# Warn about deprecated DEBUG_LAST_REQUEST if used
+_warn_deprecated_debug_setting()
+
+# Warn about suboptimal timeout configuration
+_warn_timeout_configuration()
+
+
+# --- Lifespan Manager ---
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+ """
+ Manages the application lifecycle.
+
+ Creates and initializes:
+ - Shared HTTP client with connection pooling
+ - KiroAuthManager for token management
+ - ModelInfoCache for model caching
+
+ The shared HTTP client is used by all requests to reduce memory usage
+ and enable connection reuse. This is especially important for handling
+ concurrent requests efficiently (fixes issue #24).
+ """
+ logger.info("Starting application... Creating state managers.")
+
+ # 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
+ limits = httpx.Limits(
+ max_connections=100,
+ max_keepalive_connections=20,
+ keepalive_expiry=30.0 # Close idle connections after 30 seconds
+ )
+ # Timeout configuration for streaming (long read timeout for model "thinking")
+ timeout = httpx.Timeout(
+ connect=30.0,
+ read=STREAMING_READ_TIMEOUT, # 300 seconds for streaming
+ write=30.0,
+ pool=30.0
+ )
+ app.state.http_client = httpx.AsyncClient(
+ limits=limits,
+ timeout=timeout,
+ follow_redirects=True
+ )
+ logger.info("Shared HTTP client created with connection pooling")
+
+ # Create AuthManager
+ # Priority: SQLite DB > JSON file > environment variables
+ app.state.auth_manager = KiroAuthManager(
+ refresh_token=REFRESH_TOKEN,
+ profile_arn=PROFILE_ARN,
+ region=REGION,
+ creds_file=KIRO_CREDS_FILE if KIRO_CREDS_FILE else None,
+ sqlite_db=KIRO_CLI_DB_FILE if KIRO_CLI_DB_FILE else None,
+ )
+
+ # Create model cache
+ app.state.model_cache = ModelInfoCache()
+
+ # BLOCKING: Load models from Kiro API at startup
+ # This ensures the cache is populated BEFORE accepting any requests.
+ # No race conditions - requests only start after yield.
+ logger.info("Loading models from Kiro API...")
+ try:
+ token = await app.state.auth_manager.get_access_token()
+ from kiro.utils import get_kiro_headers
+ from kiro.auth import AuthType
+ headers = get_kiro_headers(app.state.auth_manager, token)
+
+ # Build params - profileArn is only needed for Kiro Desktop auth
+ params = {"origin": "AI_EDITOR"}
+ if app.state.auth_manager.auth_type == AuthType.KIRO_DESKTOP and app.state.auth_manager.profile_arn:
+ params["profileArn"] = app.state.auth_manager.profile_arn
+
+ async with httpx.AsyncClient(timeout=30) as client:
+ response = await client.get(
+ f"{app.state.auth_manager.q_host}/ListAvailableModels",
+ headers=headers,
+ params=params
+ )
+
+ if response.status_code == 200:
+ data = response.json()
+ models_list = data.get("models", [])
+ await app.state.model_cache.update(models_list)
+ logger.debug(f"Successfully loaded {len(models_list)} models from Kiro API")
+ else:
+ raise Exception(f"HTTP {response.status_code}")
+ except Exception as e:
+ # FALLBACK: Use built-in model list
+ logger.error(f"Failed to fetch models from Kiro API: {e}")
+ logger.error("Using pre-configured fallback models. Not all models may be available on your plan, or the list may be outdated.")
+
+ # Populate cache with fallback models
+ await app.state.model_cache.update(FALLBACK_MODELS)
+ logger.debug(f"Loaded {len(FALLBACK_MODELS)} fallback models")
+
+ # Add hidden models to cache (they appear in /v1/models but not in Kiro API)
+ # Hidden models are added ALWAYS, regardless of API success/failure
+ for display_name, internal_id in HIDDEN_MODELS.items():
+ app.state.model_cache.add_hidden_model(display_name, internal_id)
+
+ if HIDDEN_MODELS:
+ logger.debug(f"Added {len(HIDDEN_MODELS)} hidden models to cache")
+
+ # Log final cache state
+ all_models = app.state.model_cache.get_all_model_ids()
+ logger.info(f"Model cache ready: {len(all_models)} models total")
+
+ # Create model resolver (uses cache + hidden models for resolution)
+ app.state.model_resolver = ModelResolver(
+ cache=app.state.model_cache,
+ hidden_models=HIDDEN_MODELS
+ )
+ logger.info("Model resolver initialized")
+
+ yield
+
+ # Graceful shutdown
+ logger.info("Shutting down application...")
+ try:
+ await app.state.http_client.aclose()
+ logger.info("Shared HTTP client closed")
+ except Exception as e:
+ logger.warning(f"Error closing shared HTTP client: {e}")
+
+
+# --- FastAPI Application ---
+app = FastAPI(
+ title=APP_TITLE,
+ description=APP_DESCRIPTION,
+ version=APP_VERSION,
+ lifespan=lifespan
+)
+
+
+# --- CORS Middleware ---
+# Allow CORS for all origins to support browser clients
+# and tools that send preflight OPTIONS requests
+app.add_middleware(
+ CORSMiddleware,
+ allow_origins=["*"], # Allow all origins
+ allow_credentials=True,
+ allow_methods=["*"], # Allow all methods (GET, POST, OPTIONS, etc.)
+ allow_headers=["*"], # Allow all headers
+)
+
+
+# --- Debug Logger Middleware ---
+# Initializes debug logging BEFORE Pydantic validation
+# This allows capturing validation errors (422) in debug logs
+app.add_middleware(DebugLoggerMiddleware)
+
+
+# --- Validation Error Handler Registration ---
+app.add_exception_handler(RequestValidationError, validation_exception_handler)
+
+
+# --- Route Registration ---
+# OpenAI-compatible API: /v1/models, /v1/chat/completions
+app.include_router(openai_router)
+
+# Anthropic-compatible API: /v1/messages
+app.include_router(anthropic_router)
+
+
+# --- Uvicorn log config ---
+# Minimal configuration for redirecting uvicorn logs to loguru.
+# Uses InterceptHandler which intercepts logs and passes them to loguru.
+UVICORN_LOG_CONFIG = {
+ "version": 1,
+ "disable_existing_loggers": False,
+ "handlers": {
+ "default": {
+ "class": "main.InterceptHandler",
+ },
+ },
+ "loggers": {
+ "uvicorn": {"handlers": ["default"], "level": "INFO", "propagate": False},
+ "uvicorn.error": {"handlers": ["default"], "level": "INFO", "propagate": False},
+ "uvicorn.access": {"handlers": ["default"], "level": "INFO", "propagate": False},
+ },
+}
+
+
+def parse_cli_args() -> argparse.Namespace:
+ """
+ Parse command-line arguments for server configuration.
+
+ CLI arguments have the highest priority, overriding both
+ environment variables and default values.
+
+ Returns:
+ Parsed arguments namespace with host and port values
+ """
+ parser = argparse.ArgumentParser(
+ description=f"{APP_TITLE} - {APP_DESCRIPTION}",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog="""
+Configuration Priority (highest to lowest):
+ 1. CLI arguments (--host, --port)
+ 2. Environment variables (SERVER_HOST, SERVER_PORT)
+ 3. Default values (0.0.0.0:8000)
+
+Examples:
+ python main.py # Use defaults or env vars
+ python main.py --port 9000 # Override port only
+ python main.py --host 127.0.0.1 # Local connections only
+ python main.py -H 0.0.0.0 -p 8080 # Short form
+
+ SERVER_PORT=9000 python main.py # Via environment
+ uvicorn main:app --port 9000 # Via uvicorn directly
+ """
+ )
+
+ parser.add_argument(
+ "-H", "--host",
+ type=str,
+ default=None, # None means "use env or default"
+ metavar="HOST",
+ help=f"Server host address (default: {DEFAULT_SERVER_HOST}, env: SERVER_HOST)"
+ )
+
+ parser.add_argument(
+ "-p", "--port",
+ type=int,
+ default=None, # None means "use env or default"
+ metavar="PORT",
+ help=f"Server port (default: {DEFAULT_SERVER_PORT}, env: SERVER_PORT)"
+ )
+
+ parser.add_argument(
+ "-v", "--version",
+ action="version",
+ version=f"%(prog)s {APP_VERSION}"
+ )
+
+ return parser.parse_args()
+
+
+def resolve_server_config(args: argparse.Namespace) -> tuple[str, int]:
+ """
+ Resolve final server configuration using priority hierarchy.
+
+ Priority (highest to lowest):
+ 1. CLI arguments (--host, --port)
+ 2. Environment variables (SERVER_HOST, SERVER_PORT)
+ 3. Default values (0.0.0.0:8000)
+
+ Args:
+ args: Parsed CLI arguments
+
+ Returns:
+ Tuple of (host, port) with resolved values
+ """
+ # Host resolution: CLI > ENV > Default
+ if args.host is not None:
+ final_host = args.host
+ host_source = "CLI argument"
+ elif SERVER_HOST != DEFAULT_SERVER_HOST:
+ final_host = SERVER_HOST
+ host_source = "environment variable"
+ else:
+ final_host = DEFAULT_SERVER_HOST
+ host_source = "default"
+
+ # Port resolution: CLI > ENV > Default
+ if args.port is not None:
+ final_port = args.port
+ port_source = "CLI argument"
+ elif SERVER_PORT != DEFAULT_SERVER_PORT:
+ final_port = SERVER_PORT
+ port_source = "environment variable"
+ else:
+ final_port = DEFAULT_SERVER_PORT
+ port_source = "default"
+
+ # Log configuration sources for transparency
+ logger.debug(f"Host: {final_host} (from {host_source})")
+ logger.debug(f"Port: {final_port} (from {port_source})")
+
+ return final_host, final_port
+
+
+def print_startup_banner(host: str, port: int) -> None:
+ """
+ Print a startup banner with server information.
+
+ Args:
+ host: Server host address
+ port: Server port
+ """
+ # ANSI color codes
+ GREEN = "\033[92m"
+ CYAN = "\033[96m"
+ YELLOW = "\033[93m"
+ WHITE = "\033[97m"
+ BOLD = "\033[1m"
+ DIM = "\033[2m"
+ RESET = "\033[0m"
+
+ # Determine display URL
+ display_host = "localhost" if host == "0.0.0.0" else host
+ url = f"http://{display_host}:{port}"
+
+ print()
+ print(f" {WHITE}{BOLD}👻 {APP_TITLE} v{APP_VERSION}{RESET}")
+ print()
+ print(f" {WHITE}Server running at:{RESET}")
+ print(f" {GREEN}{BOLD}➜ {url}{RESET}")
+ print()
+ print(f" {DIM}API Docs: {url}/docs{RESET}")
+ print(f" {DIM}Health Check: {url}/health{RESET}")
+ print()
+ print(f" {DIM}{'─' * 48}{RESET}")
+ print(f" {WHITE}💬 Found a bug? Need help? Have questions?{RESET}")
+ print(f" {YELLOW}➜ https://github.com/jwadow/kiro-gateway/issues{RESET}")
+ print(f" {DIM}{'─' * 48}{RESET}")
+ print()
+
+
+# --- Entry Point ---
+if __name__ == "__main__":
+ import uvicorn
+
+ # Parse CLI arguments
+ args = parse_cli_args()
+
+ # Resolve final configuration with priority hierarchy
+ final_host, final_port = resolve_server_config(args)
+
+ # Print startup banner
+ print_startup_banner(final_host, final_port)
+
+ logger.info(f"Starting Uvicorn server on {final_host}:{final_port}...")
+
+ # Use string reference to avoid double module import
+ uvicorn.run(
+ "main:app",
+ host=final_host,
+ port=final_port,
+ log_config=UVICORN_LOG_CONFIG,
+ )
diff --git a/kiro-gateway/manual_api_test.py b/kiro-gateway/manual_api_test.py
new file mode 100644
index 0000000000000000000000000000000000000000..80d663f61d2ba9ac835b1e46664211af51d9a2b7
--- /dev/null
+++ b/kiro-gateway/manual_api_test.py
@@ -0,0 +1,471 @@
+# -*- 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 .
+
+import json
+import os
+import sqlite3
+import sys
+import uuid
+from pathlib import Path
+from enum import Enum
+
+import requests
+from dotenv import load_dotenv
+from loguru import logger
+
+# --- Load environment variables ---
+load_dotenv()
+
+
+class AuthType(Enum):
+ """Type of authentication mechanism."""
+ KIRO_DESKTOP = "kiro_desktop"
+ AWS_SSO_OIDC = "aws_sso_oidc"
+
+
+# --- Configuration ---
+# API region - CodeWhisperer API is only available in us-east-1
+API_REGION = "us-east-1"
+KIRO_API_HOST = f"https://q.{API_REGION}.amazonaws.com"
+KIRO_DESKTOP_TOKEN_URL = f"https://prod.{API_REGION}.auth.desktop.kiro.dev/refreshToken"
+
+# SSO region - may differ from API region (e.g., ap-southeast-1 for Singapore users)
+# This is used only for AWS SSO OIDC token refresh
+SSO_REGION = None
+AWS_SSO_OIDC_TOKEN_URL = None # Will be set when SSO_REGION is known
+
+REFRESH_TOKEN = os.getenv("REFRESH_TOKEN")
+PROFILE_ARN = os.getenv("PROFILE_ARN", "arn:aws:codewhisperer:us-east-1:699475941385:profile/EHGA3GRVQMUK")
+KIRO_CREDS_FILE = os.getenv("KIRO_CREDS_FILE", "")
+KIRO_CLI_DB_FILE = os.getenv("KIRO_CLI_DB_FILE", "")
+
+# AWS SSO OIDC specific credentials
+CLIENT_ID = None
+CLIENT_SECRET = None
+SCOPES = None
+AUTH_TOKEN = None
+AUTH_TYPE = AuthType.KIRO_DESKTOP
+
+
+def load_credentials_from_json(file_path: str) -> bool:
+ """Load credentials from JSON file."""
+ global REFRESH_TOKEN, PROFILE_ARN, CLIENT_ID, CLIENT_SECRET, AUTH_TYPE
+ global SSO_REGION, AWS_SSO_OIDC_TOKEN_URL
+
+ try:
+ creds_path = Path(file_path).expanduser()
+ if not creds_path.exists():
+ logger.warning(f"Credentials file not found: {file_path}")
+ return False
+
+ with open(creds_path, 'r', encoding='utf-8') as f:
+ creds_data = json.load(f)
+
+ # Load common fields
+ if 'refreshToken' in creds_data:
+ REFRESH_TOKEN = creds_data['refreshToken']
+ if 'profileArn' in creds_data:
+ PROFILE_ARN = creds_data['profileArn']
+ if 'region' in creds_data:
+ # Store as SSO region for OIDC token refresh only
+ # IMPORTANT: CodeWhisperer API is only available in us-east-1,
+ # so we don't update KIRO_API_HOST here
+ SSO_REGION = creds_data['region']
+ AWS_SSO_OIDC_TOKEN_URL = f"https://oidc.{SSO_REGION}.amazonaws.com/token"
+ logger.debug(f"SSO region from JSON: {SSO_REGION} (API stays at {API_REGION})")
+
+ # Load AWS SSO OIDC specific fields
+ if 'clientId' in creds_data:
+ CLIENT_ID = creds_data['clientId']
+ if 'clientSecret' in creds_data:
+ CLIENT_SECRET = creds_data['clientSecret']
+
+ # Detect auth type
+ if CLIENT_ID and CLIENT_SECRET:
+ AUTH_TYPE = AuthType.AWS_SSO_OIDC
+ logger.info(f"Detected auth type: AWS SSO OIDC")
+ else:
+ AUTH_TYPE = AuthType.KIRO_DESKTOP
+ logger.info(f"Detected auth type: Kiro Desktop")
+
+ logger.info(f"Credentials loaded from {file_path}")
+ return True
+
+ except Exception as e:
+ logger.error(f"Error loading credentials from file: {e}")
+ return False
+
+
+def load_credentials_from_sqlite(db_path: str) -> bool:
+ """Load credentials from kiro-cli SQLite database."""
+ global REFRESH_TOKEN, CLIENT_ID, CLIENT_SECRET, AUTH_TYPE, SCOPES, AUTH_TOKEN
+ global SSO_REGION, AWS_SSO_OIDC_TOKEN_URL
+
+ try:
+ path = Path(db_path).expanduser()
+ if not path.exists():
+ logger.warning(f"SQLite database not found: {db_path}")
+ return False
+
+ conn = sqlite3.connect(str(path))
+ cursor = conn.cursor()
+
+ # Load token data (try both kiro-cli and codewhisperer key formats)
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", ("kirocli:odic:token",))
+ token_row = cursor.fetchone()
+ if not token_row:
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", ("codewhisperer:odic:token",))
+ token_row = cursor.fetchone()
+
+ if token_row:
+ token_data = json.loads(token_row[0])
+ if token_data:
+ # Check if we have a valid access token
+ if 'access_token' in token_data and 'expires_at' in token_data:
+ from datetime import datetime
+ expires_at = datetime.fromisoformat(token_data['expires_at'].replace('Z', '+00:00'))
+ if expires_at > datetime.now().astimezone():
+ AUTH_TOKEN = token_data['access_token']
+ logger.info("Found valid access token in database (will use after HEADERS init)")
+ if 'refresh_token' in token_data:
+ REFRESH_TOKEN = token_data['refresh_token']
+ if 'scopes' in token_data:
+ SCOPES = token_data['scopes']
+ if 'region' in token_data:
+ # Store as SSO region for OIDC token refresh only
+ # IMPORTANT: CodeWhisperer API is only available in us-east-1,
+ # so we don't update KIRO_API_HOST here
+ SSO_REGION = token_data['region']
+ AWS_SSO_OIDC_TOKEN_URL = f"https://oidc.{SSO_REGION}.amazonaws.com/token"
+ logger.debug(f"SSO region from SQLite: {SSO_REGION} (API stays at {API_REGION})")
+
+ # Load device registration (client_id, client_secret) - try both key formats
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", ("kirocli:odic:device-registration",))
+ registration_row = cursor.fetchone()
+ if not registration_row:
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", ("codewhisperer:odic:device-registration",))
+ registration_row = cursor.fetchone()
+
+ if registration_row:
+ registration_data = json.loads(registration_row[0])
+ if registration_data:
+ if 'client_id' in registration_data:
+ CLIENT_ID = registration_data['client_id']
+ if 'client_secret' in registration_data:
+ CLIENT_SECRET = registration_data['client_secret']
+
+ conn.close()
+
+ # Detect auth type
+ if CLIENT_ID and CLIENT_SECRET:
+ AUTH_TYPE = AuthType.AWS_SSO_OIDC
+ logger.info(f"Detected auth type: AWS SSO OIDC (from SQLite)")
+ else:
+ AUTH_TYPE = AuthType.KIRO_DESKTOP
+ logger.info(f"Detected auth type: Kiro Desktop (from SQLite)")
+
+ logger.info(f"Credentials loaded from SQLite: {db_path}")
+ return True
+
+ except sqlite3.Error as e:
+ logger.error(f"SQLite error: {e}")
+ return False
+ except Exception as e:
+ logger.error(f"Error loading credentials from SQLite: {e}")
+ return False
+
+
+# --- Load credentials (priority: SQLite > JSON > env) ---
+cred_source = "REFRESH_TOKEN"
+
+if KIRO_CLI_DB_FILE:
+ if load_credentials_from_sqlite(KIRO_CLI_DB_FILE):
+ cred_source = "KIRO_CLI_DB_FILE (SQLite)"
+elif KIRO_CREDS_FILE:
+ if load_credentials_from_json(KIRO_CREDS_FILE):
+ cred_source = "KIRO_CREDS_FILE (JSON)"
+
+# --- Validate required credentials ---
+if not REFRESH_TOKEN:
+ logger.error("No credentials configured. Set REFRESH_TOKEN, KIRO_CREDS_FILE, or KIRO_CLI_DB_FILE. Exiting.")
+ sys.exit(1)
+
+# Additional validation for AWS SSO OIDC
+if AUTH_TYPE == AuthType.AWS_SSO_OIDC and (not CLIENT_ID or not CLIENT_SECRET):
+ logger.error("AWS SSO OIDC requires clientId and clientSecret. Exiting.")
+ sys.exit(1)
+
+# Global variables
+AUTH_TOKEN = None
+HEADERS = {
+ "Authorization": None,
+ "Content-Type": "application/json",
+ "User-Agent": "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-31c325a0ff0a9c8dec5d13048f4257462d751fe5b8af4cb1088f1fca45856c64",
+ "x-amz-user-agent": "aws-sdk-js/1.0.27 KiroIDE-0.7.45-31c325a0ff0a9c8dec5d13048f4257462d751fe5b8af4cb1088f1fca45856c64",
+ "x-amzn-codewhisperer-optout": "true",
+ "x-amzn-kiro-agent-mode": "vibe",
+}
+
+
+def refresh_auth_token():
+ """Refreshes AUTH_TOKEN via appropriate endpoint based on auth type."""
+ global AUTH_TOKEN, HEADERS
+
+ if AUTH_TYPE == AuthType.AWS_SSO_OIDC:
+ return refresh_auth_token_aws_sso_oidc()
+ else:
+ return refresh_auth_token_kiro_desktop()
+
+
+def refresh_auth_token_kiro_desktop():
+ """Refreshes AUTH_TOKEN via Kiro Desktop Auth endpoint."""
+ global AUTH_TOKEN, HEADERS
+ logger.info("Refreshing Kiro token via Kiro Desktop Auth...")
+
+ payload = {"refreshToken": REFRESH_TOKEN}
+ headers = {
+ "Content-Type": "application/json",
+ "User-Agent": "KiroIDE-0.7.45-31c325a0ff0a9c8dec5d13048f4257462d751fe5b8af4cb1088f1fca45856c64",
+ }
+
+ try:
+ response = requests.post(KIRO_DESKTOP_TOKEN_URL, json=payload, headers=headers)
+ response.raise_for_status()
+ data = response.json()
+
+ new_token = data.get("accessToken")
+ expires_in = data.get("expiresIn")
+
+ if not new_token:
+ logger.error("Failed to get accessToken from response")
+ return False
+
+ logger.success(f"Token refreshed via Kiro Desktop Auth. Expires in: {expires_in}s")
+ AUTH_TOKEN = new_token
+ HEADERS['Authorization'] = f"Bearer {AUTH_TOKEN}"
+ return True
+
+ except requests.exceptions.RequestException as e:
+ logger.error(f"Error refreshing token via Kiro Desktop Auth: {e}")
+ if hasattr(e, 'response') and e.response:
+ logger.error(f"Server response: {e.response.status_code} {e.response.text}")
+ return False
+
+
+def refresh_auth_token_aws_sso_oidc():
+ """Refreshes AUTH_TOKEN via AWS SSO OIDC endpoint."""
+ global AUTH_TOKEN, HEADERS
+ logger.info("Refreshing Kiro token via AWS SSO OIDC...")
+
+ # Determine SSO OIDC URL (use SSO_REGION if set, otherwise fall back to API_REGION)
+ sso_region = SSO_REGION or API_REGION
+ oidc_url = AWS_SSO_OIDC_TOKEN_URL or f"https://oidc.{sso_region}.amazonaws.com/token"
+
+ # AWS SSO OIDC uses form-urlencoded data
+ data = {
+ "grant_type": "refresh_token",
+ "client_id": CLIENT_ID,
+ "client_secret": CLIENT_SECRET,
+ "refresh_token": REFRESH_TOKEN,
+ }
+
+ # Note: scope parameter is NOT sent during refresh per OAuth 2.0 RFC 6749 Section 6
+ # AWS SSO OIDC uses the originally granted scopes automatically
+ headers = {
+ "Content-Type": "application/x-www-form-urlencoded",
+ }
+
+ # Log request details (without secrets) for debugging
+ logger.debug(f"AWS SSO OIDC refresh request: url={oidc_url}, "
+ f"sso_region={sso_region}, api_region={API_REGION}, "
+ f"client_id={CLIENT_ID[:8] if CLIENT_ID else 'None'}...")
+
+ try:
+ response = requests.post(oidc_url, data=data, headers=headers)
+
+ # Log response details for debugging (especially on errors)
+ if response.status_code != 200:
+ logger.error(f"AWS SSO OIDC refresh failed: status={response.status_code}")
+ logger.error(f"AWS SSO OIDC response body: {response.text}")
+ # Try to parse AWS error for more details
+ try:
+ error_json = response.json()
+ error_code = error_json.get("error", "unknown")
+ error_desc = error_json.get("error_description", "no description")
+ logger.error(f"AWS SSO OIDC error details: error={error_code}, "
+ f"description={error_desc}")
+ except Exception:
+ pass # Body wasn't JSON, already logged as text
+ response.raise_for_status()
+
+ result = response.json()
+
+ new_token = result.get("accessToken")
+ expires_in = result.get("expiresIn", 3600)
+
+ if not new_token:
+ logger.error(f"Failed to get accessToken from AWS SSO OIDC response: {result}")
+ return False
+
+ logger.success(f"Token refreshed via AWS SSO OIDC. Expires in: {expires_in}s")
+ AUTH_TOKEN = new_token
+ HEADERS['Authorization'] = f"Bearer {AUTH_TOKEN}"
+ return True
+
+ except requests.exceptions.RequestException as e:
+ logger.error(f"Error refreshing token via AWS SSO OIDC: {e}")
+ if hasattr(e, 'response') and e.response is not None:
+ logger.error(f"Server response: {e.response.status_code} {e.response.text}")
+ return False
+
+
+def get_profile_arn():
+ """Gets the profile ARN from ListAvailableProfiles endpoint."""
+ global PROFILE_ARN
+ logger.info("Getting profile ARN from /ListAvailableProfiles...")
+ url = f"{KIRO_API_HOST}/ListAvailableProfiles"
+
+ try:
+ response = requests.get(url, headers=HEADERS)
+ response.raise_for_status()
+ data = response.json()
+
+ profiles = data.get("profiles", [])
+ if profiles:
+ # Use the first available profile
+ PROFILE_ARN = profiles[0].get("arn")
+ logger.info(f"Found profile ARN: {PROFILE_ARN}")
+ return True
+ else:
+ logger.warning("No profiles found")
+ return False
+ except requests.exceptions.RequestException as e:
+ logger.error(f"ListAvailableProfiles failed: {e}")
+ if hasattr(e, 'response') and e.response is not None:
+ logger.error(f"Server response: {e.response.status_code} {e.response.text}")
+ return False
+
+
+def test_get_models():
+ """Tests the ListAvailableModels endpoint."""
+ logger.info("Testing /ListAvailableModels...")
+ url = f"{KIRO_API_HOST}/ListAvailableModels"
+ params = {
+ "origin": "AI_EDITOR",
+ "profileArn": PROFILE_ARN
+ }
+
+ try:
+ response = requests.get(url, headers=HEADERS, params=params)
+ response.raise_for_status()
+
+ logger.info(f"Response status: {response.status_code}")
+ logger.debug(f"Response (JSON):\n{json.dumps(response.json(), indent=2, ensure_ascii=False)}")
+ logger.success("ListAvailableModels test COMPLETED SUCCESSFULLY")
+ return True
+ except requests.exceptions.RequestException as e:
+ logger.error(f"ListAvailableModels test failed: {e}")
+ return False
+
+
+def test_generate_content():
+ """Tests the generateAssistantResponse endpoint."""
+ logger.info("Testing /generateAssistantResponse...")
+ url = f"{KIRO_API_HOST}/generateAssistantResponse"
+
+ payload = {
+ "conversationState": {
+ "agentContinuationId": str(uuid.uuid4()),
+ "agentTaskType": "vibe",
+ "chatTriggerType": "MANUAL",
+ "conversationId": str(uuid.uuid4()),
+ "currentMessage": {
+ "userInputMessage": {
+ "content": "Hello! Say something short.",
+ "modelId": "claude-haiku-4.5",
+ "origin": "AI_EDITOR",
+ "userInputMessageContext": {
+ "tools": []
+ }
+ }
+ },
+ "history": []
+ }
+ }
+
+ # Only add profileArn if it's set and not AWS SSO OIDC
+ # AWS SSO OIDC (Builder ID) users don't need profileArn and it causes 403 if sent
+ if PROFILE_ARN and AUTH_TYPE != AuthType.AWS_SSO_OIDC:
+ payload["profileArn"] = PROFILE_ARN
+
+ try:
+ with requests.post(url, headers=HEADERS, json=payload, stream=True) as response:
+ response.raise_for_status()
+ logger.info(f"Response status: {response.status_code}")
+ logger.info("Streaming response:")
+
+ for chunk in response.iter_content(chunk_size=1024):
+ if chunk:
+ # Try to decode and find JSON
+ chunk_str = chunk.decode('utf-8', errors='ignore')
+ logger.debug(f"Chunk: {chunk_str[:200]}...")
+
+ logger.success("generateAssistantResponse test COMPLETED")
+ return True
+ except requests.exceptions.RequestException as e:
+ logger.error(f"generateAssistantResponse test failed: {e}")
+ return False
+
+
+if __name__ == "__main__":
+ logger.info(f"Starting Kiro API tests...")
+ logger.info(f" Credentials source: {cred_source}")
+ logger.info(f" Auth type: {AUTH_TYPE.value}")
+ logger.info(f" API Region: {API_REGION}")
+ logger.info(f" SSO Region: {SSO_REGION or 'not set (using API region)'}")
+ logger.info(f" API Host: {KIRO_API_HOST}")
+
+ # Check if we already have a valid token from the database
+ if AUTH_TOKEN:
+ HEADERS['Authorization'] = f"Bearer {AUTH_TOKEN}"
+ logger.info("Using existing valid access token from database")
+ token_ok = True
+ else:
+ token_ok = refresh_auth_token()
+
+ if token_ok:
+ # Get profile ARN dynamically for AWS SSO OIDC users
+ if AUTH_TYPE == AuthType.AWS_SSO_OIDC:
+ get_profile_arn()
+
+ models_ok = test_get_models()
+ generate_ok = test_generate_content()
+
+ if models_ok and generate_ok:
+ logger.success(f"All tests passed successfully!")
+ logger.success(f" Auth type: {AUTH_TYPE.value}")
+ logger.success(f" Credentials: {cred_source}")
+ else:
+ logger.warning(f"One or more tests failed.")
+ else:
+ logger.error("Failed to refresh token. Tests not started.")
+ logger.error(f" Auth type: {AUTH_TYPE.value}")
+ sso_region = SSO_REGION or API_REGION
+ oidc_url = AWS_SSO_OIDC_TOKEN_URL or f"https://oidc.{sso_region}.amazonaws.com/token"
+ logger.error(f" Token URL: {oidc_url if AUTH_TYPE == AuthType.AWS_SSO_OIDC else KIRO_DESKTOP_TOKEN_URL}")
diff --git a/kiro-gateway/pytest.ini b/kiro-gateway/pytest.ini
new file mode 100644
index 0000000000000000000000000000000000000000..8925504660f1364f5a4446fefddc841cfc67f6a9
--- /dev/null
+++ b/kiro-gateway/pytest.ini
@@ -0,0 +1,14 @@
+[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
diff --git a/kiro-gateway/requirements.txt b/kiro-gateway/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..a429fd0aebfa8e933b2a24448d0caca0b375117d
--- /dev/null
+++ b/kiro-gateway/requirements.txt
@@ -0,0 +1,13 @@
+# Prod dependencies
+fastapi
+uvicorn[standard]
+httpx
+loguru
+requests
+python-dotenv
+tiktoken
+
+# Testing dependencies
+pytest
+pytest-asyncio
+hypothesis
\ No newline at end of file
diff --git a/kiro-gateway/tests/README.md b/kiro-gateway/tests/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..9e37eb31e02179e7a11bdd98074a6a89aa4e6125
--- /dev/null
+++ b/kiro-gateway/tests/README.md
@@ -0,0 +1,180 @@
+# Tests for Kiro Gateway
+
+A comprehensive set of unit and integration tests for Kiro Gateway, providing full coverage of all system components.
+
+## Testing Philosophy: Complete Network Isolation
+
+**The key principle of this test suite is 100% isolation from real network requests.**
+
+This is achieved through a global, automatically applied fixture `block_all_network_calls` in `tests/conftest.py`. It intercepts and blocks any attempts by `httpx.AsyncClient` to establish connections at the application level.
+
+**Benefits:**
+1. **Reliability**: Tests don't depend on external API availability or network state.
+2. **Speed**: Absence of real network delays makes test execution instant.
+3. **Security**: Guarantees that test runs never use real credentials.
+
+Any attempt to make an unauthorized network call will result in immediate test failure with an error, ensuring strict isolation control.
+
+## Running Tests
+
+### Installing Dependencies
+
+```bash
+# Main project dependencies
+pip install -r requirements.txt
+
+# Additional testing dependencies
+pip install pytest pytest-asyncio hypothesis
+```
+
+### Running All Tests
+
+```bash
+# Run the entire test suite
+pytest
+
+# Run with verbose output
+pytest -v
+
+# Run with verbose output and coverage
+pytest -v -s --tb=short
+
+# Run only unit tests
+pytest tests/unit/ -v
+
+# Run only integration tests
+pytest tests/integration/ -v
+
+# Run a specific file
+pytest tests/unit/test_auth_manager.py -v
+
+# Run a specific test
+pytest tests/unit/test_auth_manager.py::TestKiroAuthManagerInitialization::test_initialization_stores_credentials -v
+```
+
+### pytest Options
+
+```bash
+# Stop on first failure
+pytest -x
+
+# Show local variables on errors
+pytest -l
+
+# Run in parallel mode (requires pytest-xdist)
+pip install pytest-xdist
+pytest -n auto
+```
+
+## Test Structure
+
+```
+tests/
+├── conftest.py # Shared fixtures and utilities
+├── unit/ # Unit tests for individual components
+│ ├── test_auth_manager.py # KiroAuthManager tests
+│ ├── test_cache.py # ModelInfoCache tests (is_valid_model, add_hidden_model)
+│ ├── test_config.py # Configuration tests (SERVER_HOST, SERVER_PORT, LOG_LEVEL, etc.)
+│ ├── test_converters_anthropic.py # Anthropic Messages API → Kiro converter tests
+│ ├── test_converters_core.py # Shared conversion logic tests (UnifiedMessage, merging, etc.)
+│ ├── test_converters_openai.py # OpenAI Chat API → Kiro converter tests
+│ ├── test_debug_logger.py # DebugLogger tests (off/errors/all modes)
+│ ├── test_debug_middleware.py # DebugLoggerMiddleware tests (endpoint filtering, mode handling)
+│ ├── test_exceptions.py # Exception handlers tests (validation_exception_handler, sanitize_validation_errors)
+│ ├── test_http_client.py # KiroHttpClient tests
+│ ├── test_main_cli.py # CLI argument parsing tests (--host, --port)
+│ ├── test_model_resolver.py # Dynamic Model Resolution System tests
+│ ├── test_models_anthropic.py # Anthropic Pydantic models tests (all content blocks, tools, streaming)
+│ ├── test_models_openai.py # OpenAI Pydantic models tests (messages, tools, responses, streaming)
+│ ├── test_parsers.py # AwsEventStreamParser tests (including JSON truncation diagnostics)
+│ ├── test_routes_anthropic.py # Anthropic API endpoint tests (/v1/messages)
+│ ├── test_routes_openai.py # OpenAI API endpoint tests (/v1/chat/completions)
+│ ├── test_streaming_anthropic.py # Anthropic streaming response tests
+│ ├── test_streaming_core.py # Shared streaming logic tests
+│ ├── test_streaming_openai.py # OpenAI streaming response tests
+│ ├── test_thinking_parser.py # ThinkingParser tests (FSM for thinking blocks)
+│ ├── test_tokenizer.py # Tokenizer tests (tiktoken)
+│ └── test_vpn_proxy.py # VPN/Proxy configuration tests (environment variables, URL normalization, NO_PROXY)
+├── integration/ # Integration tests for full flow
+│ └── test_full_flow.py # End-to-end tests
+└── README.md # This file
+```
+
+## Testing Philosophy
+
+### Principles
+
+1. **Isolation**: Each test is completely isolated from external services through mocks
+2. **Detail**: Abundant print() for understanding test flow during debugging
+3. **Coverage**: Tests cover not only happy path, but also edge cases and errors
+4. **Security**: All tests use mock credentials, never real ones
+
+### Test Structure (Arrange-Act-Assert)
+
+Each test follows the pattern:
+1. **Arrange** (Setup): Prepare mocks and data
+2. **Act** (Action): Execute the tested action
+3. **Assert** (Verify): Verify result with explicit comparison
+
+### Test Types
+
+- **Unit tests**: Test individual functions/classes in isolation
+- **Integration tests**: Verify component interactions
+- **Security tests**: Verify security system
+- **Edge case tests**: Paranoid edge case checks
+
+## Adding New Tests
+
+When adding new tests:
+
+1. Follow existing class structure (`Test*Success`, `Test*Errors`, `Test*EdgeCases`)
+2. Use descriptive names: `test__`
+3. Add docstring with "What it does" and "Purpose"
+4. Use print() for logging test steps
+
+## Troubleshooting
+
+### Tests fail with ImportError
+
+```bash
+# Make sure you're in project root
+cd /path/to/kiro-gateway
+
+# pytest.ini already contains pythonpath = .
+# Just run pytest
+pytest
+```
+
+### Tests pass locally but fail in CI
+
+- Check dependency versions in requirements.txt
+- Ensure all mocks correctly isolate external calls
+
+### Async tests don't work
+
+```bash
+# Make sure pytest-asyncio is installed
+pip install pytest-asyncio
+
+# Check for @pytest.mark.asyncio decorator
+```
+
+## Coverage Metrics
+
+To check code coverage:
+
+```bash
+# Install coverage
+pip install pytest-cov
+
+# Run with coverage report
+pytest --cov=kiro --cov-report=html
+
+# View report
+open htmlcov/index.html # macOS/Linux
+start htmlcov/index.html # Windows
+```
+
+## Contacts and Support
+
+If you find bugs or have suggestions for test improvements, create an issue in the project repository.
diff --git a/kiro-gateway/tests/conftest.py b/kiro-gateway/tests/conftest.py
new file mode 100644
index 0000000000000000000000000000000000000000..76796acde6e4f6e939dfd23c6521be7753504836
--- /dev/null
+++ b/kiro-gateway/tests/conftest.py
@@ -0,0 +1,986 @@
+# -*- coding: utf-8 -*-
+
+"""
+Common fixtures and utilities for testing Kiro Gateway.
+
+Provides test isolation from external services and global state.
+All tests MUST be completely isolated from the network.
+"""
+
+import asyncio
+import json
+import pytest
+import time
+from typing import AsyncGenerator, Dict, Any, List
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
+from datetime import datetime, timezone
+
+import httpx
+from fastapi.testclient import TestClient
+
+
+# =============================================================================
+# Event Loop Fixtures
+# =============================================================================
+
+@pytest.fixture(scope="session")
+def event_loop():
+ """
+ Creates an event loop for the entire test session.
+ Required for proper async fixture operation.
+ """
+ print("Creating event loop for test session...")
+ loop = asyncio.get_event_loop_policy().new_event_loop()
+ yield loop
+ print("Closing event loop...")
+ loop.close()
+
+
+# =============================================================================
+# Environment Fixtures
+# =============================================================================
+
+@pytest.fixture
+def mock_env_vars(monkeypatch):
+ """
+ Mocks environment variables for isolation from real credentials.
+ """
+ print("Setting up mocked environment variables...")
+ monkeypatch.setenv("REFRESH_TOKEN", "test_refresh_token_abcdef")
+ monkeypatch.setenv("PROXY_API_KEY", "test_proxy_key_12345")
+ monkeypatch.setenv("PROFILE_ARN", "arn:aws:codewhisperer:us-east-1:123456789:profile/test")
+ monkeypatch.setenv("KIRO_REGION", "us-east-1")
+ return {
+ "REFRESH_TOKEN": "test_refresh_token_abcdef",
+ "PROXY_API_KEY": "test_proxy_key_12345",
+ "PROFILE_ARN": "arn:aws:codewhisperer:us-east-1:123456789:profile/test",
+ "KIRO_REGION": "us-east-1"
+ }
+
+
+# =============================================================================
+# Token and Authentication Fixtures
+# =============================================================================
+
+@pytest.fixture
+def valid_kiro_token():
+ """Returns a valid mock Kiro access token."""
+ return "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.test_kiro_access_token"
+
+
+@pytest.fixture
+def mock_kiro_token_response(valid_kiro_token):
+ """
+ Factory for creating mock Kiro token refresh endpoint responses.
+ """
+ def _create_response(expires_in: int = 3600, token: str = None):
+ return {
+ "accessToken": token or valid_kiro_token,
+ "refreshToken": "new_refresh_token_xyz",
+ "expiresIn": expires_in,
+ "profileArn": "arn:aws:codewhisperer:us-east-1:123456789:profile/test"
+ }
+ return _create_response
+
+
+@pytest.fixture
+def valid_proxy_api_key():
+ """Returns a valid proxy API key (from config)."""
+ return "changeme_proxy_secret"
+
+
+@pytest.fixture
+def invalid_proxy_api_key():
+ """Returns an invalid API key for negative tests."""
+ return "invalid_wrong_secret_key"
+
+
+@pytest.fixture
+def auth_headers(valid_proxy_api_key):
+ """
+ Factory for creating valid and invalid Authorization headers.
+ """
+ def _create_headers(api_key: str = None, invalid: bool = False):
+ if invalid:
+ return {"Authorization": "Bearer wrong_key_123"}
+ key = api_key or valid_proxy_api_key
+ return {"Authorization": f"Bearer {key}"}
+
+ return _create_headers
+
+
+# =============================================================================
+# Kiro Models Fixtures
+# =============================================================================
+
+@pytest.fixture
+def mock_kiro_models_response():
+ """
+ Mock successful response from Kiro API for ListAvailableModels.
+ """
+ return {
+ "models": [
+ {
+ "modelId": "claude-sonnet-4.5",
+ "displayName": "Claude Sonnet 4.5",
+ "tokenLimits": {
+ "maxInputTokens": 200000,
+ "maxOutputTokens": 8192
+ }
+ },
+ {
+ "modelId": "claude-opus-4.5",
+ "displayName": "Claude Opus 4.5",
+ "tokenLimits": {
+ "maxInputTokens": 200000,
+ "maxOutputTokens": 8192
+ }
+ },
+ {
+ "modelId": "claude-haiku-4.5",
+ "displayName": "Claude Haiku 4.5",
+ "tokenLimits": {
+ "maxInputTokens": 200000,
+ "maxOutputTokens": 8192
+ }
+ }
+ ]
+ }
+
+
+# =============================================================================
+# Kiro Streaming Response Fixtures
+# =============================================================================
+
+@pytest.fixture
+def mock_kiro_streaming_chunks():
+ """
+ Returns a list of mock SSE chunks from Kiro API for streaming response.
+ Covers: regular text, tool calls, usage.
+ """
+ return [
+ # Chunk 1: Text start
+ b'{"content":"Hello"}',
+ # Chunk 2: Text continuation
+ b'{"content":" World!"}',
+ # Chunk 3: Tool call start
+ b'{"name":"get_weather","toolUseId":"call_abc123"}',
+ # Chunk 4: Tool call input
+ b'{"input":"{\\"location\\": \\"Moscow\\"}"}',
+ # Chunk 5: Tool call stop
+ b'{"stop":true}',
+ # Chunk 6: Usage
+ b'{"usage":1.5}',
+ # Chunk 7: Context usage
+ b'{"contextUsagePercentage":25.5}',
+ ]
+
+@pytest.fixture
+def mock_kiro_simple_text_chunks():
+ """
+ Mock simple text response from Kiro (without tool calls).
+ """
+ return [
+ b'{"content":"This is a complete response."}',
+ b'{"usage":0.5}',
+ b'{"contextUsagePercentage":10.0}',
+ ]
+
+
+@pytest.fixture
+def mock_kiro_stream_with_usage():
+ """
+ Mock Kiro SSE response with usage information.
+ """
+ return [
+ b'{"content":"Final text."}',
+ b'{"usage":1.3}',
+ b'{"contextUsagePercentage":50.0}',
+ ]
+
+
+# =============================================================================
+# OpenAI Request Fixtures
+# =============================================================================
+
+@pytest.fixture
+def sample_openai_chat_request():
+ """
+ Factory for creating valid OpenAI chat completion requests.
+ """
+ def _create_request(
+ model: str = "claude-sonnet-4-5",
+ messages: list = None,
+ stream: bool = False,
+ temperature: float = None,
+ max_tokens: int = None,
+ tools: list = None,
+ **kwargs
+ ):
+ if messages is None:
+ messages = [{"role": "user", "content": "Hello, AI!"}]
+
+ request = {
+ "model": model,
+ "messages": messages,
+ "stream": stream
+ }
+
+ if temperature is not None:
+ request["temperature"] = temperature
+ if max_tokens is not None:
+ request["max_tokens"] = max_tokens
+ if tools is not None:
+ request["tools"] = tools
+
+ request.update(kwargs)
+ return request
+
+ return _create_request
+
+
+@pytest.fixture
+def sample_tool_definition():
+ """
+ Sample tool definition for testing tool calling.
+ """
+ return {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get weather for a location",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "City name"}
+ },
+ "required": ["location"]
+ }
+ }
+ }
+
+
+# =============================================================================
+# HTTP Client Fixtures
+# =============================================================================
+
+@pytest.fixture
+async def mock_httpx_client():
+ """
+ Creates a mocked httpx.AsyncClient for isolation from network requests.
+ """
+ print("Creating mocked httpx.AsyncClient...")
+ mock_client = AsyncMock(spec=httpx.AsyncClient)
+
+ # Mock methods
+ mock_client.post = AsyncMock()
+ mock_client.get = AsyncMock()
+ mock_client.aclose = AsyncMock()
+ mock_client.build_request = Mock()
+ mock_client.send = AsyncMock()
+ mock_client.is_closed = False
+
+ return mock_client
+
+
+@pytest.fixture
+def mock_httpx_response():
+ """
+ Factory for creating mocked httpx.Response objects.
+ """
+ def _create_response(
+ status_code: int = 200,
+ json_data: Dict[str, Any] = None,
+ text: str = None,
+ stream_chunks: list = None
+ ):
+ print(f"Creating mock httpx.Response (status={status_code})...")
+ mock_response = AsyncMock(spec=httpx.Response)
+ mock_response.status_code = status_code
+
+ if json_data is not None:
+ mock_response.json = Mock(return_value=json_data)
+
+ if text is not None:
+ mock_response.text = text
+ mock_response.content = text.encode()
+
+ if stream_chunks is not None:
+ # For streaming responses
+ async def mock_aiter_bytes():
+ for chunk in stream_chunks:
+ yield chunk
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ mock_response.raise_for_status = Mock()
+ mock_response.aclose = AsyncMock()
+ mock_response.aread = AsyncMock(return_value=b'{"error": "mocked error"}')
+
+ return mock_response
+
+ return _create_response
+
+
+# =============================================================================
+# Global Network Blocking
+# =============================================================================
+
+@pytest.fixture(scope="session", autouse=True)
+def block_all_network_calls():
+ """
+ CRITICAL FIXTURE: Globally blocks ALL network calls.
+ Ensures that NO test can make a real network request.
+ """
+
+ # Create a mock that will be used for all AsyncClient instances
+ mock_async_client = AsyncMock(spec=httpx.AsyncClient)
+
+ async def network_call_error(*args, **kwargs):
+ raise RuntimeError(
+ "🚨 CRITICAL ERROR: Real network request attempt detected! "
+ "Test did not provide a mock for httpx.AsyncClient. "
+ "All HTTP calls must be explicitly mocked."
+ )
+
+ mock_async_client.post.side_effect = network_call_error
+ mock_async_client.get.side_effect = network_call_error
+ mock_async_client.send.side_effect = network_call_error
+
+ # Mock context manager
+ mock_async_client.__aenter__ = AsyncMock(return_value=mock_async_client)
+ mock_async_client.__aexit__ = AsyncMock()
+ mock_async_client.aclose = AsyncMock()
+ mock_async_client.is_closed = False
+
+ # Patch AsyncClient in modules where it's used
+ patchers = [
+ patch('kiro.auth.httpx.AsyncClient', return_value=mock_async_client),
+ patch('kiro.http_client.httpx.AsyncClient', return_value=mock_async_client),
+ patch('kiro.streaming_openai.httpx.AsyncClient', return_value=mock_async_client),
+ ]
+
+ # Start patchers
+ for patcher in patchers:
+ patcher.start()
+
+ print("🛡️ GLOBAL NETWORK BLOCKING ACTIVATED")
+
+ yield
+
+ # Stop patchers
+ for patcher in patchers:
+ patcher.stop()
+
+ print("🛡️ GLOBAL NETWORK BLOCKING DEACTIVATED")
+
+
+# =============================================================================
+# Application Fixtures
+# =============================================================================
+
+@pytest.fixture
+def clean_app():
+ """
+ Returns a "clean" application instance for each test.
+ """
+ print("Importing application for test...")
+ from main import app
+ # Reset all dependency overrides before test
+ app.dependency_overrides = {}
+ return app
+
+
+@pytest.fixture
+def test_client(clean_app):
+ """
+ Creates a FastAPI TestClient for synchronous endpoint tests,
+ properly handling lifespan events.
+ """
+ print("Creating TestClient with lifespan support...")
+ with TestClient(clean_app) as client:
+ yield client
+ print("Closing TestClient...")
+
+
+@pytest.fixture
+async def async_test_client(clean_app):
+ """
+ Creates an asynchronous test client for async endpoints.
+ """
+ print("Creating async test client...")
+ from httpx import AsyncClient, ASGITransport
+
+ transport = ASGITransport(app=clean_app)
+ async with AsyncClient(transport=transport, base_url="http://test") as client:
+ yield client
+
+ print("Closing async test client...")
+
+
+# =============================================================================
+# KiroAuthManager Fixtures
+# =============================================================================
+
+@pytest.fixture
+def mock_auth_manager():
+ """
+ Creates a mocked KiroAuthManager for tests.
+ """
+ from kiro.auth import KiroAuthManager
+
+ manager = KiroAuthManager(
+ refresh_token="test_refresh_token",
+ profile_arn="arn:aws:codewhisperer:us-east-1:123456789:profile/test",
+ region="us-east-1"
+ )
+
+ # Set valid token
+ manager._access_token = "test_access_token"
+ manager._expires_at = datetime.now(timezone.utc).replace(
+ year=2099 # Far in the future
+ )
+
+ return manager
+
+
+@pytest.fixture
+def expired_auth_manager():
+ """
+ Creates a KiroAuthManager with an expired token.
+ """
+ from kiro.auth import KiroAuthManager
+
+ manager = KiroAuthManager(
+ refresh_token="test_refresh_token",
+ profile_arn="arn:aws:codewhisperer:us-east-1:123456789:profile/test",
+ region="us-east-1"
+ )
+
+ # Set expired token
+ manager._access_token = "expired_token"
+ manager._expires_at = datetime.now(timezone.utc).replace(
+ year=2020 # In the past
+ )
+
+ return manager
+
+
+# =============================================================================
+# ModelInfoCache Fixtures
+# =============================================================================
+
+@pytest.fixture
+def sample_models_data():
+ """
+ Returns a list of models for testing ModelInfoCache.
+ """
+ return [
+ {
+ "modelId": "claude-sonnet-4",
+ "displayName": "Claude Sonnet 4",
+ "tokenLimits": {
+ "maxInputTokens": 200000,
+ "maxOutputTokens": 8192
+ }
+ },
+ {
+ "modelId": "claude-opus-4.5",
+ "displayName": "Claude Opus 4.5",
+ "tokenLimits": {
+ "maxInputTokens": 200000,
+ "maxOutputTokens": 8192
+ }
+ },
+ {
+ "modelId": "claude-haiku-4.5",
+ "displayName": "Claude Haiku 4.5",
+ "tokenLimits": {
+ "maxInputTokens": 100000,
+ "maxOutputTokens": 4096
+ }
+ }
+ ]
+
+
+@pytest.fixture
+def empty_model_cache():
+ """
+ Creates an empty ModelInfoCache.
+ """
+ from kiro.cache import ModelInfoCache
+ return ModelInfoCache()
+
+
+@pytest.fixture
+async def populated_model_cache(mock_kiro_models_response):
+ """
+ Creates a ModelInfoCache with pre-populated data.
+ """
+ from kiro.cache import ModelInfoCache
+
+ cache = ModelInfoCache()
+ await cache.update(mock_kiro_models_response["models"])
+ return cache
+
+
+# =============================================================================
+# Time Fixtures
+# =============================================================================
+
+@pytest.fixture
+def mock_time():
+ """
+ Mocks time.time() for predictable behavior in tests.
+ """
+ with patch('time.time') as mock:
+ # Fixed point in time: 2024-01-01 12:00:00
+ mock.return_value = 1704110400.0
+ yield mock
+
+
+@pytest.fixture
+def mock_datetime():
+ """
+ Mocks datetime.now() for predictable behavior.
+ """
+ fixed_time = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
+
+ with patch('kiro.auth.datetime') as mock_dt:
+ mock_dt.now.return_value = fixed_time
+ mock_dt.fromisoformat = datetime.fromisoformat
+ mock_dt.fromtimestamp = datetime.fromtimestamp
+ yield mock_dt
+
+
+# =============================================================================
+# Temporary File Fixtures
+# =============================================================================
+
+@pytest.fixture
+def temp_creds_file(tmp_path):
+ """
+ Creates a temporary credentials file for tests (Kiro Desktop format).
+ """
+ creds_file = tmp_path / "kiro-auth-token.json"
+ creds_data = {
+ "accessToken": "file_access_token",
+ "refreshToken": "file_refresh_token",
+ "expiresAt": "2099-01-01T00:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:123456789:profile/test",
+ "region": "us-east-1"
+ }
+ creds_file.write_text(json.dumps(creds_data))
+ return str(creds_file)
+
+
+@pytest.fixture
+def temp_aws_sso_creds_file(tmp_path):
+ """
+ Creates a temporary credentials file for tests (AWS SSO OIDC format).
+ Contains clientId and clientSecret, indicating AWS SSO OIDC authentication.
+ """
+ creds_file = tmp_path / "aws-sso-cache.json"
+ creds_data = {
+ "accessToken": "aws_sso_access_token",
+ "refreshToken": "aws_sso_refresh_token",
+ "expiresAt": "2099-01-01T00:00:00.000Z",
+ "region": "us-east-1",
+ "clientId": "test_client_id_12345",
+ "clientSecret": "test_client_secret_67890"
+ }
+ creds_file.write_text(json.dumps(creds_data))
+ return str(creds_file)
+
+
+@pytest.fixture
+def temp_sqlite_db(tmp_path):
+ """
+ Creates a temporary SQLite database for tests (kiro-cli format).
+
+ Contains auth_kv table with keys:
+ - 'codewhisperer:odic:token': JSON with access_token, refresh_token, expires_at, region
+ - 'codewhisperer:odic:device-registration': JSON with client_id, client_secret
+ """
+ import sqlite3
+
+ db_file = tmp_path / "data.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ # Create auth_kv table
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ # Insert token data
+ token_data = {
+ "access_token": "sqlite_access_token",
+ "refresh_token": "sqlite_refresh_token",
+ "expires_at": "2099-01-01T00:00:00Z",
+ "region": "eu-west-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:token", json.dumps(token_data))
+ )
+
+ # Insert device registration data
+ registration_data = {
+ "client_id": "sqlite_client_id",
+ "client_secret": "sqlite_client_secret",
+ "region": "eu-west-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:device-registration", json.dumps(registration_data))
+ )
+
+ conn.commit()
+ conn.close()
+
+ return str(db_file)
+
+
+@pytest.fixture
+def temp_sqlite_db_token_only(tmp_path):
+ """
+ Creates a SQLite database with token only (without device-registration).
+ Used for testing partial loading.
+ """
+ import sqlite3
+
+ db_file = tmp_path / "data_token_only.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ token_data = {
+ "access_token": "partial_access_token",
+ "refresh_token": "partial_refresh_token",
+ "region": "ap-southeast-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:token", json.dumps(token_data))
+ )
+
+ conn.commit()
+ conn.close()
+
+ return str(db_file)
+
+
+@pytest.fixture
+def temp_sqlite_db_invalid_json(tmp_path):
+ """
+ Creates a SQLite database with invalid JSON in value.
+ Used for testing error handling.
+ """
+ import sqlite3
+
+ db_file = tmp_path / "data_invalid.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ # Insert invalid JSON
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:token", "not a valid json {{{")
+ )
+
+ conn.commit()
+ conn.close()
+
+ return str(db_file)
+
+
+@pytest.fixture
+def mock_aws_sso_oidc_token_response():
+ """
+ Factory for creating mock AWS SSO OIDC token endpoint responses.
+ """
+ def _create_response(
+ access_token: str = "new_aws_sso_access_token",
+ refresh_token: str = "new_aws_sso_refresh_token",
+ expires_in: int = 3600
+ ):
+ return {
+ "accessToken": access_token,
+ "refreshToken": refresh_token,
+ "expiresIn": expires_in,
+ "tokenType": "Bearer"
+ }
+ return _create_response
+
+
+@pytest.fixture
+def temp_debug_dir(tmp_path):
+ """
+ Creates a temporary directory for debug files.
+ """
+ debug_dir = tmp_path / "debug_logs"
+ debug_dir.mkdir()
+ return debug_dir
+
+
+# =============================================================================
+# Parser Fixtures
+# =============================================================================
+
+@pytest.fixture
+def aws_event_parser():
+ """
+ Creates an AwsEventStreamParser instance for tests.
+ """
+ from kiro.parsers import AwsEventStreamParser
+ return AwsEventStreamParser()
+
+
+# =============================================================================
+# Test Utilities
+# =============================================================================
+
+def create_kiro_content_chunk(content: str) -> bytes:
+ """Utility for creating a Kiro SSE chunk with content."""
+ return f'{{"content":"{content}"}}'.encode()
+
+
+def create_kiro_tool_start_chunk(name: str, tool_id: str) -> bytes:
+ """Utility for creating a Kiro SSE chunk with tool call start."""
+ return f'{{"name":"{name}","toolUseId":"{tool_id}"}}'.encode()
+
+
+def create_kiro_tool_input_chunk(input_json: str) -> bytes:
+ """Utility for creating a Kiro SSE chunk with tool call input."""
+ escaped = input_json.replace('"', '\\"')
+ return f'{{"input":"{escaped}"}}'.encode()
+
+
+def create_kiro_tool_stop_chunk() -> bytes:
+ """Utility for creating a Kiro SSE chunk with tool call stop."""
+ return b'{"stop":true}'
+
+
+def create_kiro_usage_chunk(usage: float) -> bytes:
+ """Utility for creating a Kiro SSE chunk with usage."""
+ return f'{{"usage":{usage}}}'.encode()
+
+
+def create_kiro_context_usage_chunk(percentage: float) -> bytes:
+ """Utility for creating a Kiro SSE chunk with context usage."""
+ return f'{{"contextUsagePercentage":{percentage}}}'.encode()
+
+
+# =============================================================================
+# Social Login Fixtures (for new functionality)
+# =============================================================================
+
+@pytest.fixture
+def temp_sqlite_db_social(tmp_path):
+ """
+ Creates a temporary SQLite database with social login credentials.
+
+ Contains auth_kv table with key:
+ - 'kirocli:social:token': JSON with access_token, refresh_token, expires_at, provider
+
+ This simulates kiro-cli with Google/GitHub social login (no client_id/client_secret).
+ """
+ import sqlite3
+
+ db_file = tmp_path / "data_social.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ # Create auth_kv table
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ # Insert social login token data
+ token_data = {
+ "access_token": "social_access_token",
+ "refresh_token": "social_refresh_token",
+ "expires_at": "2099-01-01T00:00:00Z",
+ "provider": "google",
+ "profile_arn": "arn:aws:codewhisperer:us-east-1:123456789:profile/social",
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("kirocli:social:token", json.dumps(token_data))
+ )
+
+ conn.commit()
+ conn.close()
+
+ return str(db_file)
+
+
+@pytest.fixture
+def temp_sqlite_db_all_keys(tmp_path):
+ """
+ Creates a SQLite database with ALL three token keys.
+
+ Used for testing key priority:
+ 1. kirocli:social:token (highest priority)
+ 2. kirocli:odic:token
+ 3. codewhisperer:odic:token (lowest priority)
+ """
+ import sqlite3
+
+ db_file = tmp_path / "data_all_keys.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ # Insert all three keys with different tokens
+ social_data = {
+ "access_token": "social_token",
+ "refresh_token": "social_refresh",
+ "expires_at": "2099-01-01T00:00:00Z",
+ "provider": "google"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("kirocli:social:token", json.dumps(social_data))
+ )
+
+ odic_data = {
+ "access_token": "odic_token",
+ "refresh_token": "odic_refresh",
+ "expires_at": "2099-01-01T00:00:00Z"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("kirocli:odic:token", json.dumps(odic_data))
+ )
+
+ legacy_data = {
+ "access_token": "legacy_token",
+ "refresh_token": "legacy_refresh",
+ "expires_at": "2099-01-01T00:00:00Z"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:token", json.dumps(legacy_data))
+ )
+
+ conn.commit()
+ conn.close()
+
+ return str(db_file)
+
+
+# =============================================================================
+# Enterprise Kiro IDE Fixtures (Issue #45)
+# =============================================================================
+
+@pytest.fixture
+def temp_enterprise_ide_creds_file(tmp_path):
+ """
+ Creates a temporary credentials file for Enterprise Kiro IDE.
+
+ Contains:
+ - clientIdHash: Hash used to locate device registration file
+ - refreshToken, accessToken, expiresAt, region
+
+ This simulates Enterprise Kiro IDE with IdC (AWS IAM Identity Center) login.
+ """
+ creds_file = tmp_path / "kiro-auth-token.json"
+ creds_data = {
+ "accessToken": "enterprise_access_token",
+ "refreshToken": "enterprise_refresh_token",
+ "expiresAt": "2099-01-01T00:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:123456789:profile/enterprise",
+ "region": "us-east-1",
+ "clientIdHash": "abc123def456"
+ }
+ creds_file.write_text(json.dumps(creds_data))
+ return str(creds_file)
+
+
+@pytest.fixture
+def temp_enterprise_device_registration(tmp_path):
+ """
+ Creates a temporary device registration file for Enterprise Kiro IDE.
+
+ Located at: ~/.aws/sso/cache/{clientIdHash}.json
+ Contains: clientId, clientSecret
+ """
+ # Create .aws/sso/cache directory structure
+ aws_dir = tmp_path / ".aws" / "sso" / "cache"
+ aws_dir.mkdir(parents=True, exist_ok=True)
+
+ # Create device registration file
+ device_reg_file = aws_dir / "abc123def456.json"
+ device_reg_data = {
+ "clientId": "enterprise_client_id_12345",
+ "clientSecret": "enterprise_client_secret_67890",
+ "region": "us-east-1"
+ }
+ device_reg_file.write_text(json.dumps(device_reg_data))
+
+ return str(device_reg_file)
+
+
+@pytest.fixture
+def temp_enterprise_ide_complete(tmp_path, monkeypatch):
+ """
+ Creates a complete Enterprise IDE setup with both credentials and device registration.
+
+ Returns tuple: (creds_file_path, device_reg_file_path)
+ """
+ # Mock Path.home() to return tmp_path
+ monkeypatch.setattr('pathlib.Path.home', lambda: tmp_path)
+
+ # Create credentials file
+ creds_file = tmp_path / "kiro-auth-token.json"
+ creds_data = {
+ "accessToken": "enterprise_access_token",
+ "refreshToken": "enterprise_refresh_token",
+ "expiresAt": "2099-01-01T00:00:00.000Z",
+ "profileArn": "arn:aws:codewhisperer:us-east-1:123456789:profile/enterprise",
+ "region": "us-east-1",
+ "clientIdHash": "abc123def456"
+ }
+ creds_file.write_text(json.dumps(creds_data))
+
+ # Create device registration file
+ aws_dir = tmp_path / ".aws" / "sso" / "cache"
+ aws_dir.mkdir(parents=True, exist_ok=True)
+
+ device_reg_file = aws_dir / "abc123def456.json"
+ device_reg_data = {
+ "clientId": "enterprise_client_id_12345",
+ "clientSecret": "enterprise_client_secret_67890",
+ "region": "us-east-1"
+ }
+ device_reg_file.write_text(json.dumps(device_reg_data))
+
+ return (str(creds_file), str(device_reg_file))
diff --git a/kiro-gateway/tests/integration/test_full_flow.py b/kiro-gateway/tests/integration/test_full_flow.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f0f7b5c784b86dc760c5aefcce336a37222b80d
--- /dev/null
+++ b/kiro-gateway/tests/integration/test_full_flow.py
@@ -0,0 +1,397 @@
+# -*- coding: utf-8 -*-
+
+"""
+Integration tests for complete end-to-end flow.
+Checks interaction of all system components.
+"""
+
+import pytest
+import json
+from unittest.mock import AsyncMock, Mock, patch, MagicMock
+from datetime import datetime, timezone, timedelta
+
+from fastapi.testclient import TestClient
+import httpx
+
+from kiro.config import PROXY_API_KEY
+
+
+class TestFullChatCompletionFlow:
+ """Integration tests for complete chat completions flow."""
+
+ def test_full_flow_health_to_models_to_chat(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Checks complete flow from health check to chat completions.
+ Goal: Ensure all endpoints work together.
+ """
+ print("Step 1: Health check...")
+ health_response = test_client.get("/health")
+ assert health_response.status_code == 200
+ assert health_response.json()["status"] == "healthy"
+ print(f"Health: {health_response.json()}")
+
+ print("Step 2: Getting models list...")
+ models_response = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"}
+ )
+ assert models_response.status_code == 200
+ assert len(models_response.json()["data"]) > 0
+ print(f"Models: {[m['id'] for m in models_response.json()['data']]}")
+
+ print("Step 3: Validating chat completions request...")
+ # This request will pass validation but fail on HTTP due to network blocking
+ chat_response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+ # Request should pass validation (not 422)
+ assert chat_response.status_code != 422
+ print(f"Chat response status: {chat_response.status_code}")
+
+ def test_authentication_flow(self, test_client, valid_proxy_api_key, invalid_proxy_api_key):
+ """
+ What it does: Checks authentication flow.
+ Goal: Ensure protected endpoints require authorization.
+ """
+ print("Step 1: Request without authorization...")
+ no_auth_response = test_client.get("/v1/models")
+ assert no_auth_response.status_code == 401
+ print(f"Without authorization: {no_auth_response.status_code}")
+
+ print("Step 2: Request with invalid key...")
+ wrong_auth_response = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {invalid_proxy_api_key}"}
+ )
+ assert wrong_auth_response.status_code == 401
+ print(f"Invalid key: {wrong_auth_response.status_code}")
+
+ print("Step 3: Request with valid key...")
+ valid_auth_response = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"}
+ )
+ assert valid_auth_response.status_code == 200
+ print(f"Valid key: {valid_auth_response.status_code}")
+
+ def test_openai_compatibility_format(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Checks response format compatibility with OpenAI API.
+ Goal: Ensure responses conform to OpenAI specification.
+ """
+ print("Checking /v1/models format...")
+ models_response = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"}
+ )
+
+ assert models_response.status_code == 200
+ data = models_response.json()
+
+ # Check OpenAI response structure
+ assert "object" in data
+ assert data["object"] == "list"
+ assert "data" in data
+ assert isinstance(data["data"], list)
+
+ # Check structure of each model
+ for model in data["data"]:
+ assert "id" in model
+ assert "object" in model
+ assert model["object"] == "model"
+ assert "owned_by" in model
+ assert "created" in model
+
+ print(f"Format conforms to OpenAI API: {len(data['data'])} models")
+
+
+class TestRequestValidationFlow:
+ """Integration tests for request validation."""
+
+ def test_chat_completions_request_validation(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Checks validation of various request formats.
+ Goal: Ensure validation works correctly.
+ """
+ print("Test 1: Empty messages...")
+ empty_messages = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={"model": "claude-sonnet-4-5", "messages": []}
+ )
+ assert empty_messages.status_code == 422
+ print(f"Empty messages: {empty_messages.status_code}")
+
+ print("Test 2: Missing model...")
+ no_model = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={"messages": [{"role": "user", "content": "Hello"}]}
+ )
+ assert no_model.status_code == 422
+ print(f"Without model: {no_model.status_code}")
+
+ print("Test 3: Missing messages...")
+ no_messages = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={"model": "claude-sonnet-4-5"}
+ )
+ assert no_messages.status_code == 422
+ print(f"Without messages: {no_messages.status_code}")
+
+ print("Test 4: Valid request...")
+ valid_request = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+ # Validation should pass (not 422)
+ assert valid_request.status_code != 422
+ print(f"Valid request: {valid_request.status_code}")
+
+ def test_complex_message_formats(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Checks handling of complex message formats.
+ Goal: Ensure multimodal and tool formats are accepted.
+ """
+ print("Test 1: System + User messages...")
+ system_user = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {"role": "system", "content": "You are helpful"},
+ {"role": "user", "content": "Hello"}
+ ]
+ }
+ )
+ assert system_user.status_code != 422
+ print(f"System + User: {system_user.status_code}")
+
+ print("Test 2: Multi-turn conversation...")
+ multi_turn = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {"role": "user", "content": "Hello"},
+ {"role": "assistant", "content": "Hi there!"},
+ {"role": "user", "content": "How are you?"}
+ ]
+ }
+ )
+ assert multi_turn.status_code != 422
+ print(f"Multi-turn: {multi_turn.status_code}")
+
+ print("Test 3: With tools...")
+ with_tools = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "What's the weather?"}],
+ "tools": [{
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get weather",
+ "parameters": {"type": "object", "properties": {}}
+ }
+ }]
+ }
+ )
+ assert with_tools.status_code != 422
+ print(f"With tools: {with_tools.status_code}")
+
+
+class TestErrorHandlingFlow:
+ """Integration tests for error handling."""
+
+ def test_invalid_json_handling(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Checks handling of invalid JSON.
+ Goal: Ensure invalid JSON returns clear error.
+ """
+ print("Sending invalid JSON...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={
+ "Authorization": f"Bearer {valid_proxy_api_key}",
+ "Content-Type": "application/json"
+ },
+ content=b"not valid json"
+ )
+
+ assert response.status_code == 422
+ print(f"Invalid JSON: {response.status_code}")
+
+ def test_wrong_content_type_handling(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Checks handling of wrong Content-Type.
+ Goal: Ensure wrong Content-Type is handled.
+ """
+ print("Sending with wrong Content-Type...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={
+ "Authorization": f"Bearer {valid_proxy_api_key}",
+ "Content-Type": "text/plain"
+ },
+ content=b"Hello"
+ )
+
+ # Should be validation error
+ assert response.status_code == 422
+ print(f"Wrong Content-Type: {response.status_code}")
+
+
+class TestModelsEndpointIntegration:
+ """Integration tests for /v1/models endpoint."""
+
+ def test_models_returns_all_available_models(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Checks that all models from config are returned.
+ Goal: Ensure completeness of models list.
+ """
+ print("Getting models list...")
+ response = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"}
+ )
+
+ assert response.status_code == 200
+
+ returned_ids = {m["id"] for m in response.json()["data"]}
+
+ print(f"Returned models: {returned_ids}")
+
+ # At minimum, hidden models should be available
+ assert len(returned_ids) >= 1, "Expected at least one model (hidden models)"
+
+ def test_models_caching_behavior(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Checks models caching behavior.
+ Goal: Ensure repeated requests work correctly.
+ """
+ print("First models request...")
+ response1 = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"}
+ )
+ assert response1.status_code == 200
+
+ print("Second models request...")
+ response2 = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"}
+ )
+ assert response2.status_code == 200
+
+ # Responses should be identical
+ assert response1.json()["data"] == response2.json()["data"]
+ print("Caching works correctly")
+
+
+class TestStreamingFlagHandling:
+ """Integration tests for stream flag handling."""
+
+ def test_stream_true_accepted(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Checks that stream=true is accepted.
+ Goal: Ensure streaming mode is available.
+
+ Note: Streaming mode requires HTTP client mock,
+ as request is executed inside generator.
+ """
+ print("Request with stream=true...")
+
+ # Create mock response for streaming
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+
+ async def mock_aiter_bytes():
+ yield b'{"content":"Hello"}'
+ yield b'{"usage":0.5}'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+ mock_response.aclose = AsyncMock()
+
+ # Mock request_with_retry to return our mock response
+ with patch('kiro.routes_openai.KiroHttpClient') as MockHttpClient:
+ mock_client_instance = AsyncMock()
+ mock_client_instance.request_with_retry = AsyncMock(return_value=mock_response)
+ mock_client_instance.client = AsyncMock()
+ mock_client_instance.close = AsyncMock()
+ MockHttpClient.return_value = mock_client_instance
+
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "stream": True
+ }
+ )
+
+ # Validation should pass and streaming should work
+ assert response.status_code == 200
+ print(f"stream=true: {response.status_code}")
+
+ def test_stream_false_accepted(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Checks that stream=false is accepted.
+ Goal: Ensure non-streaming mode is available.
+ """
+ print("Request with stream=false...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "stream": False
+ }
+ )
+
+ # Validation should pass
+ assert response.status_code != 422
+ print(f"stream=false: {response.status_code}")
+
+
+class TestHealthEndpointIntegration:
+ """Integration tests for health endpoints."""
+
+ def test_root_and_health_consistency(self, test_client):
+ """
+ What it does: Checks consistency of / and /health.
+ Goal: Ensure both endpoints return correct status.
+ """
+ print("Request to /...")
+ root_response = test_client.get("/")
+
+ print("Request to /health...")
+ health_response = test_client.get("/health")
+
+ assert root_response.status_code == 200
+ assert health_response.status_code == 200
+
+ # Both should show "ok" status
+ assert root_response.json()["status"] == "ok"
+ assert health_response.json()["status"] == "healthy"
+
+ # Versions should match
+ assert root_response.json()["version"] == health_response.json()["version"]
+
+ print("Health endpoints are consistent")
diff --git a/kiro-gateway/tests/unit/test_auth_manager.py b/kiro-gateway/tests/unit/test_auth_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..99a61ad33309f07a10812baea50f741049903b1f
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_auth_manager.py
@@ -0,0 +1,2969 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for KiroAuthManager.
+Tests token management logic for Kiro without real network requests.
+"""
+
+import asyncio
+import json
+import pytest
+from datetime import datetime, timezone, timedelta
+from unittest.mock import AsyncMock, Mock, patch
+import httpx
+
+from kiro.auth import KiroAuthManager, AuthType
+from kiro.config import TOKEN_REFRESH_THRESHOLD, get_aws_sso_oidc_url
+
+
+class TestKiroAuthManagerInitialization:
+ """Tests for KiroAuthManager initialization."""
+
+ def test_initialization_stores_credentials(self):
+ """
+ What it does: Verifies correct storage of credentials during initialization.
+ Purpose: Ensure all constructor parameters are stored in private fields.
+ """
+ print("Setup: Creating KiroAuthManager with test credentials...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh_123",
+ profile_arn="arn:aws:codewhisperer:us-east-1:123456789:profile/test",
+ region="us-east-1"
+ )
+
+ print("Verification: All credentials stored correctly...")
+ print(f"Comparing refresh_token: Expected 'test_refresh_123', Got '{manager._refresh_token}'")
+ assert manager._refresh_token == "test_refresh_123"
+
+ print(f"Comparing profile_arn: Expected 'arn:aws:...', Got '{manager._profile_arn}'")
+ assert manager._profile_arn == "arn:aws:codewhisperer:us-east-1:123456789:profile/test"
+
+ print(f"Comparing region: Expected 'us-east-1', Got '{manager._region}'")
+ assert manager._region == "us-east-1"
+
+ print("Verification: Token is initially empty...")
+ assert manager._access_token is None
+ assert manager._expires_at is None
+
+ def test_initialization_sets_correct_urls_for_region(self):
+ """
+ What it does: Verifies URL formation based on region.
+ Purpose: Ensure URLs are dynamically formed with the correct region.
+ """
+ print("Setup: Creating KiroAuthManager with region eu-west-1...")
+ manager = KiroAuthManager(
+ refresh_token="test_token",
+ region="eu-west-1"
+ )
+
+ print("Verification: URLs contain correct region...")
+ print(f"Comparing refresh_url: Expected 'eu-west-1' in URL, Got '{manager._refresh_url}'")
+ assert "eu-west-1" in manager._refresh_url
+
+ print(f"Comparing api_host: Expected 'eu-west-1' in URL, Got '{manager._api_host}'")
+ assert "eu-west-1" in manager._api_host
+
+ print(f"Comparing q_host: Expected 'eu-west-1' in URL, Got '{manager._q_host}'")
+ assert "eu-west-1" in manager._q_host
+
+ def test_initialization_generates_fingerprint(self):
+ """
+ What it does: Verifies unique fingerprint generation.
+ Purpose: Ensure fingerprint is generated and has correct format.
+ """
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(refresh_token="test_token")
+
+ print("Verification: Fingerprint generated...")
+ print(f"Fingerprint: {manager._fingerprint}")
+ assert manager._fingerprint is not None
+ assert len(manager._fingerprint) == 64 # SHA256 hex digest
+
+
+class TestKiroAuthManagerCredentialsFile:
+ """Tests for loading credentials from file."""
+
+ def test_load_credentials_from_file(self, temp_creds_file):
+ """
+ What it does: Verifies loading credentials from JSON file.
+ Purpose: Ensure data is correctly read from file.
+ """
+ print(f"Setup: Creating KiroAuthManager with credentials file: {temp_creds_file}")
+ manager = KiroAuthManager(creds_file=temp_creds_file)
+
+ print("Verification: Data loaded from file...")
+ print(f"Comparing access_token: Expected 'file_access_token', Got '{manager._access_token}'")
+ assert manager._access_token == "file_access_token"
+
+ print(f"Comparing refresh_token: Expected 'file_refresh_token', Got '{manager._refresh_token}'")
+ assert manager._refresh_token == "file_refresh_token"
+
+ print(f"Comparing region: Expected 'us-east-1', Got '{manager._region}'")
+ assert manager._region == "us-east-1"
+
+ print("Verification: expiresAt parsed correctly...")
+ assert manager._expires_at is not None
+ assert manager._expires_at.year == 2099
+
+ def test_load_credentials_file_not_found(self, tmp_path):
+ """
+ What it does: Verifies handling of missing credentials file.
+ Purpose: Ensure application doesn't crash when file is missing.
+ """
+ print("Setup: Creating KiroAuthManager with non-existent file...")
+ non_existent_file = str(tmp_path / "non_existent.json")
+
+ manager = KiroAuthManager(
+ refresh_token="fallback_token",
+ creds_file=non_existent_file
+ )
+
+ print("Verification: Fallback refresh_token is used...")
+ print(f"Comparing refresh_token: Expected 'fallback_token', Got '{manager._refresh_token}'")
+ assert manager._refresh_token == "fallback_token"
+
+
+class TestKiroAuthManagerTokenExpiration:
+ """Tests for token expiration checking."""
+
+ def test_is_token_expiring_soon_returns_true_when_no_expires_at(self):
+ """
+ What it does: Verifies that without expires_at token is considered expiring.
+ Purpose: Ensure safe behavior when time information is missing.
+ """
+ print("Setup: Creating KiroAuthManager without expires_at...")
+ manager = KiroAuthManager(refresh_token="test_token")
+ manager._expires_at = None
+
+ print("Verification: is_token_expiring_soon returns True...")
+ result = manager.is_token_expiring_soon()
+ print(f"Comparing result: Expected True, Got {result}")
+ assert result is True
+
+ def test_is_token_expiring_soon_returns_true_when_expired(self):
+ """
+ What it does: Verifies that expired token is correctly identified.
+ Purpose: Ensure token in the past is considered expiring.
+ """
+ print("Setup: Creating KiroAuthManager with expired token...")
+ manager = KiroAuthManager(refresh_token="test_token")
+ manager._expires_at = datetime.now(timezone.utc) - timedelta(hours=1)
+
+ print("Verification: is_token_expiring_soon returns True for expired token...")
+ result = manager.is_token_expiring_soon()
+ print(f"Comparing result: Expected True, Got {result}")
+ assert result is True
+
+ def test_is_token_expiring_soon_returns_true_within_threshold(self):
+ """
+ What it does: Verifies that token within threshold is considered expiring.
+ Purpose: Ensure token is refreshed in advance (10 minutes before expiration).
+ """
+ print("Setup: Creating KiroAuthManager with token expiring in 5 minutes...")
+ manager = KiroAuthManager(refresh_token="test_token")
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(minutes=5)
+
+ print(f"TOKEN_REFRESH_THRESHOLD = {TOKEN_REFRESH_THRESHOLD} seconds")
+ print("Verification: is_token_expiring_soon returns True (5 min < 10 min threshold)...")
+ result = manager.is_token_expiring_soon()
+ print(f"Comparing result: Expected True, Got {result}")
+ assert result is True
+
+ def test_is_token_expiring_soon_returns_false_when_valid(self):
+ """
+ What it does: Verifies that valid token is not considered expiring.
+ Purpose: Ensure token far in the future doesn't require refresh.
+ """
+ print("Setup: Creating KiroAuthManager with token expiring in 1 hour...")
+ manager = KiroAuthManager(refresh_token="test_token")
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
+
+ print("Verification: is_token_expiring_soon returns False...")
+ result = manager.is_token_expiring_soon()
+ print(f"Comparing result: Expected False, Got {result}")
+ assert result is False
+
+
+class TestKiroAuthManagerTokenRefresh:
+ """Tests for token refresh mechanism."""
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_successful(self, valid_kiro_token, mock_kiro_token_response):
+ """
+ What it does: Tests successful token refresh via Kiro API.
+ Purpose: Verify that on successful response token and expiration time are set.
+ """
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ region="us-east-1"
+ )
+
+ print("Setup: Mocking successful response from Kiro...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_kiro_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling _refresh_token_request()...")
+ await manager._refresh_token_request()
+
+ print("Verification: Token set correctly...")
+ print(f"Comparing access_token: Expected '{valid_kiro_token}', Got '{manager._access_token}'")
+ assert manager._access_token == valid_kiro_token
+
+ print("Verification: Expiration time set...")
+ assert manager._expires_at is not None
+
+ print("Verification: POST request was made...")
+ mock_client.post.assert_called_once()
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_updates_refresh_token(self, mock_kiro_token_response):
+ """
+ What it does: Verifies refresh_token update from response.
+ Purpose: Ensure new refresh_token is saved.
+ """
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(refresh_token="old_refresh_token")
+
+ print("Setup: Mocking response with new refresh_token...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_kiro_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Refreshing token...")
+ await manager._refresh_token_request()
+
+ print("Verification: refresh_token updated...")
+ print(f"Comparing refresh_token: Expected 'new_refresh_token_xyz', Got '{manager._refresh_token}'")
+ assert manager._refresh_token == "new_refresh_token_xyz"
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_missing_access_token_raises(self):
+ """
+ What it does: Verifies handling of response without accessToken.
+ Purpose: Ensure exception is raised on invalid response.
+ """
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(refresh_token="test_refresh")
+
+ print("Setup: Mocking response without accessToken...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value={"expiresIn": 3600}) # No accessToken!
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Attempting token refresh...")
+ with pytest.raises(ValueError) as exc_info:
+ await manager._refresh_token_request()
+
+ print(f"Verification: ValueError raised with message: {exc_info.value}")
+ assert "accessToken" in str(exc_info.value)
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_no_refresh_token_raises(self):
+ """
+ What it does: Verifies handling of missing refresh_token.
+ Purpose: Ensure exception is raised without refresh_token.
+ """
+ print("Setup: Creating KiroAuthManager without refresh_token...")
+ manager = KiroAuthManager()
+ manager._refresh_token = None
+
+ print("Action: Attempting token refresh without refresh_token...")
+ with pytest.raises(ValueError) as exc_info:
+ await manager._refresh_token_request()
+
+ print(f"Verification: ValueError raised: {exc_info.value}")
+ assert "Refresh token" in str(exc_info.value)
+
+
+class TestKiroAuthManagerGetAccessToken:
+ """Tests for public get_access_token method."""
+
+ @pytest.mark.asyncio
+ async def test_get_access_token_refreshes_when_expired(self, valid_kiro_token, mock_kiro_token_response):
+ """
+ What it does: Verifies automatic refresh of expired token.
+ Purpose: Ensure stale token is refreshed before returning.
+ """
+ print("Setup: Creating KiroAuthManager with expired token...")
+ manager = KiroAuthManager(refresh_token="test_refresh")
+ manager._access_token = "old_expired_token"
+ manager._expires_at = datetime.now(timezone.utc) - timedelta(hours=1)
+
+ print("Setup: Mocking successful refresh...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_kiro_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Requesting token via get_access_token()...")
+ token = await manager.get_access_token()
+
+ print("Verification: Got new token, not expired one...")
+ print(f"Comparing token: Expected '{valid_kiro_token}', Got '{token}'")
+ assert token == valid_kiro_token
+ assert token != "old_expired_token"
+
+ print("Verification: _refresh_token_request was called...")
+ mock_client.post.assert_called_once()
+
+ @pytest.mark.asyncio
+ async def test_get_access_token_returns_valid_without_refresh(self, valid_kiro_token):
+ """
+ What it does: Verifies valid token is returned without refresh.
+ Purpose: Ensure no unnecessary requests are made if token is valid.
+ """
+ print("Setup: Creating KiroAuthManager with valid token...")
+ manager = KiroAuthManager(refresh_token="test_refresh")
+ manager._access_token = valid_kiro_token
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
+
+ print("Setup: Mocking httpx to track calls...")
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock()
+ mock_client_class.return_value = mock_client
+
+ print("Action: Requesting valid token...")
+ token = await manager.get_access_token()
+
+ print("Verification: Existing token returned...")
+ print(f"Comparing token: Expected '{valid_kiro_token}', Got '{token}'")
+ assert token == valid_kiro_token
+
+ print("Verification: _refresh_token was NOT called (no network requests)...")
+ mock_client.post.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_get_access_token_thread_safety(self, valid_kiro_token, mock_kiro_token_response):
+ """
+ What it does: Verifies thread safety via asyncio.Lock.
+ Purpose: Ensure parallel calls don't cause race conditions.
+ """
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(refresh_token="test_refresh")
+ manager._access_token = None
+ manager._expires_at = None
+
+ refresh_call_count = 0
+
+ async def mock_refresh():
+ nonlocal refresh_call_count
+ refresh_call_count += 1
+ await asyncio.sleep(0.1) # Simulate delay
+ manager._access_token = valid_kiro_token
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
+
+ print("Setup: Patching _refresh_token_request to track calls...")
+ with patch.object(manager, '_refresh_token_request', side_effect=mock_refresh):
+ print("Action: 5 parallel get_access_token() calls...")
+ tokens = await asyncio.gather(*[
+ manager.get_access_token() for _ in range(5)
+ ])
+
+ print("Verification: All calls got the same token...")
+ assert all(token == valid_kiro_token for token in tokens)
+
+ print(f"Verification: _refresh_token called ONLY ONCE (thanks to lock)...")
+ print(f"Comparing call count: Expected 1, Got {refresh_call_count}")
+ assert refresh_call_count == 1
+
+
+class TestKiroAuthManagerForceRefresh:
+ """Tests for forced token refresh."""
+
+ @pytest.mark.asyncio
+ async def test_force_refresh_updates_token(self, valid_kiro_token, mock_kiro_token_response):
+ """
+ What it does: Verifies forced token refresh.
+ Purpose: Ensure force_refresh always refreshes the token.
+ """
+ print("Setup: Creating KiroAuthManager with valid token...")
+ manager = KiroAuthManager(refresh_token="test_refresh")
+ manager._access_token = "old_but_valid_token"
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
+
+ print("Setup: Mocking refresh...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_kiro_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Force refreshing token...")
+ token = await manager.force_refresh()
+
+ print("Verification: Token refreshed despite old one being valid...")
+ print(f"Comparing token: Expected '{valid_kiro_token}', Got '{token}'")
+ assert token == valid_kiro_token
+
+ print("Verification: POST request was made...")
+ mock_client.post.assert_called_once()
+
+
+class TestKiroAuthManagerProperties:
+ """Tests for KiroAuthManager properties."""
+
+ def test_profile_arn_property(self):
+ """
+ What it does: Verifies profile_arn property.
+ Purpose: Ensure profile_arn is accessible via property.
+ """
+ print("Setup: Creating KiroAuthManager with profile_arn...")
+ manager = KiroAuthManager(
+ refresh_token="test",
+ profile_arn="arn:aws:test:profile"
+ )
+
+ print("Verification: profile_arn accessible...")
+ print(f"Comparing profile_arn: Expected 'arn:aws:test:profile', Got '{manager.profile_arn}'")
+ assert manager.profile_arn == "arn:aws:test:profile"
+
+ def test_region_property(self):
+ """
+ What it does: Verifies region property.
+ Purpose: Ensure region is accessible via property.
+ """
+ print("Setup: Creating KiroAuthManager with region...")
+ manager = KiroAuthManager(
+ refresh_token="test",
+ region="eu-west-1"
+ )
+
+ print("Verification: region accessible...")
+ print(f"Comparing region: Expected 'eu-west-1', Got '{manager.region}'")
+ assert manager.region == "eu-west-1"
+
+ def test_api_host_property(self):
+ """
+ What it does: Verifies api_host property.
+ Purpose: Ensure api_host is formed correctly.
+ """
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(
+ refresh_token="test",
+ region="us-east-1"
+ )
+
+ print("Verification: api_host contains codewhisperer and region...")
+ print(f"api_host: {manager.api_host}")
+ assert "codewhisperer" in manager.api_host
+ assert "us-east-1" in manager.api_host
+
+ def test_fingerprint_property(self):
+ """
+ What it does: Verifies fingerprint property.
+ Purpose: Ensure fingerprint is accessible via property.
+ """
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(refresh_token="test")
+
+ print("Verification: fingerprint accessible and has correct length...")
+ print(f"fingerprint: {manager.fingerprint}")
+ assert len(manager.fingerprint) == 64
+
+
+# =============================================================================
+# Tests for AuthType enum
+# =============================================================================
+
+class TestAuthTypeEnum:
+ """Tests for AuthType enum."""
+
+ def test_auth_type_enum_values(self):
+ """
+ What it does: Verifies AuthType enum values.
+ Purpose: Ensure enum contains KIRO_DESKTOP and AWS_SSO_OIDC.
+ """
+ print("Verification: AuthType contains KIRO_DESKTOP...")
+ assert AuthType.KIRO_DESKTOP.value == "kiro_desktop"
+
+ print("Verification: AuthType contains AWS_SSO_OIDC...")
+ assert AuthType.AWS_SSO_OIDC.value == "aws_sso_oidc"
+
+ print(f"Comparing value count: Expected 2, Got {len(AuthType)}")
+ assert len(AuthType) == 2
+
+
+# =============================================================================
+# Tests for _detect_auth_type()
+# =============================================================================
+
+class TestKiroAuthManagerDetectAuthType:
+ """Tests for _detect_auth_type() method."""
+
+ def test_detect_auth_type_kiro_desktop_when_no_client_credentials(self):
+ """
+ What it does: Verifies KIRO_DESKTOP type detection without client credentials.
+ Purpose: Ensure KIRO_DESKTOP is used without clientId/clientSecret.
+ """
+ print("Setup: Creating KiroAuthManager without client credentials...")
+ manager = KiroAuthManager(refresh_token="test_token")
+
+ print("Verification: auth_type = KIRO_DESKTOP...")
+ print(f"Comparing auth_type: Expected KIRO_DESKTOP, Got {manager.auth_type}")
+ assert manager.auth_type == AuthType.KIRO_DESKTOP
+
+ def test_detect_auth_type_aws_sso_oidc_when_client_credentials_present(self):
+ """
+ What it does: Verifies AWS_SSO_OIDC type detection with client credentials.
+ Purpose: Ensure AWS_SSO_OIDC is used with clientId and clientSecret.
+ """
+ print("Setup: Creating KiroAuthManager with client credentials...")
+ manager = KiroAuthManager(
+ refresh_token="test_token",
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+
+ print("Verification: auth_type = AWS_SSO_OIDC...")
+ print(f"Comparing auth_type: Expected AWS_SSO_OIDC, Got {manager.auth_type}")
+ assert manager.auth_type == AuthType.AWS_SSO_OIDC
+
+ def test_detect_auth_type_kiro_desktop_when_only_client_id(self):
+ """
+ What it does: Verifies type detection with only clientId (no secret).
+ Purpose: Ensure KIRO_DESKTOP is used without clientSecret.
+ """
+ print("Setup: Creating KiroAuthManager with only client_id...")
+ manager = KiroAuthManager(
+ refresh_token="test_token",
+ client_id="test_client_id"
+ )
+
+ print("Verification: auth_type = KIRO_DESKTOP (both id and secret required)...")
+ print(f"Comparing auth_type: Expected KIRO_DESKTOP, Got {manager.auth_type}")
+ assert manager.auth_type == AuthType.KIRO_DESKTOP
+
+
+# =============================================================================
+# Tests for loading AWS SSO credentials from JSON file
+# =============================================================================
+
+class TestKiroAuthManagerAwsSsoCredentialsFile:
+ """Tests for loading AWS SSO OIDC credentials from JSON file."""
+
+ def test_load_credentials_from_file_with_client_id_and_secret(self, temp_aws_sso_creds_file):
+ """
+ What it does: Verifies loading clientId and clientSecret from JSON file.
+ Purpose: Ensure AWS SSO fields are correctly read from file.
+ """
+ print(f"Setup: Creating KiroAuthManager with AWS SSO file: {temp_aws_sso_creds_file}")
+ manager = KiroAuthManager(creds_file=temp_aws_sso_creds_file)
+
+ print("Verification: clientId loaded...")
+ print(f"Comparing client_id: Expected 'test_client_id_12345', Got '{manager._client_id}'")
+ assert manager._client_id == "test_client_id_12345"
+
+ print("Verification: clientSecret loaded...")
+ print(f"Comparing client_secret: Expected 'test_client_secret_67890', Got '{manager._client_secret}'")
+ assert manager._client_secret == "test_client_secret_67890"
+
+ def test_load_credentials_from_file_auto_detects_aws_sso_oidc(self, temp_aws_sso_creds_file):
+ """
+ What it does: Verifies auto-detection of auth type after loading from file.
+ Purpose: Ensure auth_type automatically becomes AWS_SSO_OIDC.
+ """
+ print(f"Setup: Creating KiroAuthManager with AWS SSO file: {temp_aws_sso_creds_file}")
+ manager = KiroAuthManager(creds_file=temp_aws_sso_creds_file)
+
+ print("Verification: auth_type automatically detected as AWS_SSO_OIDC...")
+ print(f"Comparing auth_type: Expected AWS_SSO_OIDC, Got {manager.auth_type}")
+ assert manager.auth_type == AuthType.AWS_SSO_OIDC
+
+ def test_load_kiro_desktop_file_stays_kiro_desktop(self, temp_creds_file):
+ """
+ What it does: Verifies that Kiro Desktop file doesn't change type to AWS SSO.
+ Purpose: Ensure file without clientId/clientSecret stays KIRO_DESKTOP.
+ """
+ print(f"Setup: Creating KiroAuthManager with Kiro Desktop file: {temp_creds_file}")
+ manager = KiroAuthManager(creds_file=temp_creds_file)
+
+ print("Verification: auth_type stays KIRO_DESKTOP...")
+ print(f"Comparing auth_type: Expected KIRO_DESKTOP, Got {manager.auth_type}")
+ assert manager.auth_type == AuthType.KIRO_DESKTOP
+
+
+# =============================================================================
+# Tests for loading credentials from SQLite
+# =============================================================================
+
+class TestKiroAuthManagerSqliteCredentials:
+ """Tests for loading credentials from SQLite database (kiro-cli format)."""
+
+ def test_load_credentials_from_sqlite_success(self, temp_sqlite_db):
+ """
+ What it does: Verifies successful loading of credentials from SQLite.
+ Purpose: Ensure all data is correctly read from database.
+ """
+ print(f"Setup: Creating KiroAuthManager with SQLite: {temp_sqlite_db}")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db)
+
+ print("Verification: access_token loaded...")
+ print(f"Comparing access_token: Expected 'sqlite_access_token', Got '{manager._access_token}'")
+ assert manager._access_token == "sqlite_access_token"
+
+ print("Verification: refresh_token loaded...")
+ print(f"Comparing refresh_token: Expected 'sqlite_refresh_token', Got '{manager._refresh_token}'")
+ assert manager._refresh_token == "sqlite_refresh_token"
+
+ def test_load_credentials_from_sqlite_file_not_found(self, tmp_path):
+ """
+ What it does: Verifies handling of missing SQLite file.
+ Purpose: Ensure application doesn't crash when file is missing.
+ """
+ print("Setup: Creating KiroAuthManager with non-existent SQLite file...")
+ non_existent_db = str(tmp_path / "non_existent.sqlite3")
+
+ manager = KiroAuthManager(
+ refresh_token="fallback_token",
+ sqlite_db=non_existent_db
+ )
+
+ print("Verification: Fallback refresh_token is used...")
+ print(f"Comparing refresh_token: Expected 'fallback_token', Got '{manager._refresh_token}'")
+ assert manager._refresh_token == "fallback_token"
+
+ def test_load_credentials_from_sqlite_loads_token_data(self, temp_sqlite_db):
+ """
+ What it does: Verifies loading token data from SQLite.
+ Purpose: Ensure access_token, refresh_token, sso_region are loaded.
+ Note: API region stays at us-east-1 (CodeWhisperer API only exists there),
+ SSO region is stored separately for OIDC token refresh.
+ """
+ print(f"Setup: Creating KiroAuthManager with SQLite: {temp_sqlite_db}")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db)
+
+ print("Verification: SSO region loaded from SQLite...")
+ print(f"Comparing sso_region: Expected 'eu-west-1', Got '{manager._sso_region}'")
+ assert manager._sso_region == "eu-west-1"
+
+ print("Verification: API region stays at us-east-1...")
+ print(f"Comparing region: Expected 'us-east-1', Got '{manager._region}'")
+ assert manager._region == "us-east-1"
+
+ print("Verification: expires_at parsed...")
+ assert manager._expires_at is not None
+ assert manager._expires_at.year == 2099
+
+ def test_load_credentials_from_sqlite_loads_device_registration(self, temp_sqlite_db):
+ """
+ What it does: Verifies loading device registration from SQLite.
+ Purpose: Ensure client_id and client_secret are loaded.
+ """
+ print(f"Setup: Creating KiroAuthManager with SQLite: {temp_sqlite_db}")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db)
+
+ print("Verification: client_id loaded...")
+ print(f"Comparing client_id: Expected 'sqlite_client_id', Got '{manager._client_id}'")
+ assert manager._client_id == "sqlite_client_id"
+
+ print("Verification: client_secret loaded...")
+ print(f"Comparing client_secret: Expected 'sqlite_client_secret', Got '{manager._client_secret}'")
+ assert manager._client_secret == "sqlite_client_secret"
+
+ def test_load_credentials_from_sqlite_auto_detects_aws_sso_oidc(self, temp_sqlite_db):
+ """
+ What it does: Verifies auto-detection of auth type after loading from SQLite.
+ Purpose: Ensure auth_type automatically becomes AWS_SSO_OIDC.
+ """
+ print(f"Setup: Creating KiroAuthManager with SQLite: {temp_sqlite_db}")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db)
+
+ print("Verification: auth_type automatically detected as AWS_SSO_OIDC...")
+ print(f"Comparing auth_type: Expected AWS_SSO_OIDC, Got {manager.auth_type}")
+ assert manager.auth_type == AuthType.AWS_SSO_OIDC
+
+ def test_load_credentials_from_sqlite_handles_missing_registration_key(self, temp_sqlite_db_token_only):
+ """
+ What it does: Verifies handling of missing device-registration key.
+ Purpose: Ensure application doesn't crash without device-registration.
+ """
+ print(f"Setup: Creating KiroAuthManager with SQLite without device-registration...")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db_token_only)
+
+ print("Verification: refresh_token loaded...")
+ assert manager._refresh_token == "partial_refresh_token"
+
+ print("Verification: client_id stayed None...")
+ assert manager._client_id is None
+
+ print("Verification: auth_type = KIRO_DESKTOP (no client credentials)...")
+ assert manager.auth_type == AuthType.KIRO_DESKTOP
+
+ def test_load_credentials_from_sqlite_handles_invalid_json(self, temp_sqlite_db_invalid_json):
+ """
+ What it does: Verifies handling of invalid JSON in SQLite.
+ Purpose: Ensure application doesn't crash on invalid JSON.
+ """
+ print("Setup: Creating KiroAuthManager with SQLite with invalid JSON...")
+ manager = KiroAuthManager(
+ refresh_token="fallback_token",
+ sqlite_db=temp_sqlite_db_invalid_json
+ )
+
+ print("Verification: Fallback refresh_token is used...")
+ print(f"Comparing refresh_token: Expected 'fallback_token', Got '{manager._refresh_token}'")
+ assert manager._refresh_token == "fallback_token"
+
+ def test_sqlite_takes_priority_over_json_file(self, temp_sqlite_db, temp_creds_file):
+ """
+ What it does: Verifies SQLite priority over JSON file.
+ Purpose: Ensure SQLite is loaded instead of JSON when both specified.
+ """
+ print("Setup: Creating KiroAuthManager with SQLite and JSON file...")
+ manager = KiroAuthManager(
+ sqlite_db=temp_sqlite_db,
+ creds_file=temp_creds_file
+ )
+
+ print("Verification: Data from SQLite (not from JSON)...")
+ print(f"Comparing access_token: Expected 'sqlite_access_token', Got '{manager._access_token}'")
+ assert manager._access_token == "sqlite_access_token"
+
+ print("Verification: SSO region from SQLite...")
+ print(f"Comparing sso_region: Expected 'eu-west-1', Got '{manager._sso_region}'")
+ assert manager._sso_region == "eu-west-1"
+
+ print("Verification: API region stays at us-east-1...")
+ print(f"Comparing region: Expected 'us-east-1', Got '{manager._region}'")
+ assert manager._region == "us-east-1"
+
+
+# =============================================================================
+# Tests for _refresh_token_request() routing
+# =============================================================================
+
+class TestKiroAuthManagerRefreshTokenRouting:
+ """Tests for _refresh_token_request() routing based on auth_type."""
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_request_routes_to_kiro_desktop(self):
+ """
+ What it does: Verifies that KIRO_DESKTOP calls _refresh_token_kiro_desktop.
+ Purpose: Ensure correct routing for Kiro Desktop auth.
+ """
+ print("Setup: Creating KiroAuthManager with KIRO_DESKTOP...")
+ manager = KiroAuthManager(refresh_token="test_refresh")
+ assert manager.auth_type == AuthType.KIRO_DESKTOP
+
+ print("Setup: Mocking _refresh_token_kiro_desktop...")
+ with patch.object(manager, '_refresh_token_kiro_desktop', new_callable=AsyncMock) as mock_desktop:
+ with patch.object(manager, '_refresh_token_aws_sso_oidc', new_callable=AsyncMock) as mock_sso:
+ await manager._refresh_token_request()
+
+ print("Verification: _refresh_token_kiro_desktop was called...")
+ mock_desktop.assert_called_once()
+
+ print("Verification: _refresh_token_aws_sso_oidc was NOT called...")
+ mock_sso.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_request_routes_to_aws_sso_oidc(self):
+ """
+ What it does: Verifies that AWS_SSO_OIDC calls _refresh_token_aws_sso_oidc.
+ Purpose: Ensure correct routing for AWS SSO OIDC auth.
+ """
+ print("Setup: Creating KiroAuthManager with AWS_SSO_OIDC...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+ assert manager.auth_type == AuthType.AWS_SSO_OIDC
+
+ print("Setup: Mocking _refresh_token_aws_sso_oidc...")
+ with patch.object(manager, '_refresh_token_kiro_desktop', new_callable=AsyncMock) as mock_desktop:
+ with patch.object(manager, '_refresh_token_aws_sso_oidc', new_callable=AsyncMock) as mock_sso:
+ await manager._refresh_token_request()
+
+ print("Verification: _refresh_token_aws_sso_oidc was called...")
+ mock_sso.assert_called_once()
+
+ print("Verification: _refresh_token_kiro_desktop was NOT called...")
+ mock_desktop.assert_not_called()
+
+
+# =============================================================================
+# Tests for _refresh_token_aws_sso_oidc()
+# =============================================================================
+
+class TestKiroAuthManagerAwsSsoOidcRefresh:
+ """Tests for _refresh_token_aws_sso_oidc() method."""
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_success(self, mock_aws_sso_oidc_token_response):
+ """
+ What it does: Tests successful token refresh via AWS SSO OIDC.
+ Purpose: Verify that on successful response token and expiration time are set.
+ """
+ print("Setup: Creating KiroAuthManager with AWS SSO OIDC...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret",
+ region="us-east-1"
+ )
+
+ print("Setup: Mocking successful response from AWS SSO OIDC...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling _refresh_token_aws_sso_oidc()...")
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: Token set correctly...")
+ print(f"Comparing access_token: Expected 'new_aws_sso_access_token', Got '{manager._access_token}'")
+ assert manager._access_token == "new_aws_sso_access_token"
+
+ print("Verification: Expiration time set...")
+ assert manager._expires_at is not None
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_raises_without_refresh_token(self):
+ """
+ What it does: Verifies handling of missing refresh_token.
+ Purpose: Ensure ValueError is raised without refresh_token.
+ """
+ print("Setup: Creating KiroAuthManager without refresh_token...")
+ manager = KiroAuthManager(
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+ manager._refresh_token = None
+
+ print("Action: Attempting token refresh without refresh_token...")
+ with pytest.raises(ValueError) as exc_info:
+ await manager._refresh_token_aws_sso_oidc()
+
+ print(f"Verification: ValueError raised: {exc_info.value}")
+ assert "Refresh token" in str(exc_info.value)
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_raises_without_client_id(self):
+ """
+ What it does: Verifies handling of missing client_id.
+ Purpose: Ensure ValueError is raised without client_id.
+ """
+ print("Setup: Creating KiroAuthManager without client_id...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_secret="test_client_secret"
+ )
+ manager._client_id = None
+ manager._auth_type = AuthType.AWS_SSO_OIDC
+
+ print("Action: Attempting token refresh without client_id...")
+ with pytest.raises(ValueError) as exc_info:
+ await manager._refresh_token_aws_sso_oidc()
+
+ print(f"Verification: ValueError raised: {exc_info.value}")
+ assert "Client ID" in str(exc_info.value)
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_raises_without_client_secret(self):
+ """
+ What it does: Verifies handling of missing client_secret.
+ Purpose: Ensure ValueError is raised without client_secret.
+ """
+ print("Setup: Creating KiroAuthManager without client_secret...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id"
+ )
+ manager._client_secret = None
+ manager._auth_type = AuthType.AWS_SSO_OIDC
+
+ print("Action: Attempting token refresh without client_secret...")
+ with pytest.raises(ValueError) as exc_info:
+ await manager._refresh_token_aws_sso_oidc()
+
+ print(f"Verification: ValueError raised: {exc_info.value}")
+ assert "Client secret" in str(exc_info.value)
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_uses_correct_endpoint(self, mock_aws_sso_oidc_token_response):
+ """
+ What it does: Verifies correct endpoint usage.
+ Purpose: Ensure request goes to https://oidc.{region}.amazonaws.com/token.
+ """
+ print("Setup: Creating KiroAuthManager with region=eu-west-1...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret",
+ region="eu-west-1"
+ )
+
+ print("Setup: Mocking HTTP client...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: POST request to correct URL...")
+ call_args = mock_client.post.call_args
+ url = call_args[0][0]
+ expected_url = "https://oidc.eu-west-1.amazonaws.com/token"
+ print(f"Comparing URL: Expected '{expected_url}', Got '{url}'")
+ assert url == expected_url
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_uses_json_format(self, mock_aws_sso_oidc_token_response):
+ """
+ What it does: Verifies JSON format usage (AWS SSO OIDC CreateToken API).
+ Purpose: Ensure Content-Type = application/json (not form-urlencoded).
+ """
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+
+ print("Setup: Mocking HTTP client...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: Content-Type = application/json...")
+ call_args = mock_client.post.call_args
+ headers = call_args[1].get('headers', {})
+ print(f"Comparing Content-Type: Expected 'application/json', Got '{headers.get('Content-Type')}'")
+ assert headers.get('Content-Type') == 'application/json'
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_sends_correct_grant_type(self, mock_aws_sso_oidc_token_response):
+ """
+ What it does: Verifies correct grantType is sent (camelCase).
+ Purpose: Ensure grantType=refresh_token in JSON payload.
+ """
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+
+ print("Setup: Mocking HTTP client...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: grantType = refresh_token (camelCase in JSON)...")
+ call_args = mock_client.post.call_args
+ json_payload = call_args[1].get('json', {})
+ print(f"Comparing grantType: Expected 'refresh_token', Got '{json_payload.get('grantType')}'")
+ assert json_payload.get('grantType') == 'refresh_token'
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_updates_tokens(self, mock_aws_sso_oidc_token_response):
+ """
+ What it does: Verifies access_token and refresh_token update.
+ Purpose: Ensure both tokens are updated from response.
+ """
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(
+ refresh_token="old_refresh_token",
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+
+ print("Setup: Mocking HTTP client...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: access_token updated...")
+ assert manager._access_token == "new_aws_sso_access_token"
+
+ print("Verification: refresh_token updated...")
+ assert manager._refresh_token == "new_aws_sso_refresh_token"
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_calculates_expiration(self, mock_aws_sso_oidc_token_response):
+ """
+ What it does: Verifies correct expiration time calculation.
+ Purpose: Ensure expires_at is calculated based on expiresIn.
+ """
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+
+ print("Setup: Mocking HTTP client with expiresIn=7200...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response(expires_in=7200))
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: expires_at set...")
+ assert manager._expires_at is not None
+
+ print("Verification: expires_at in the future...")
+ from datetime import datetime, timezone
+ now = datetime.now(timezone.utc)
+ assert manager._expires_at > now
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_does_not_send_scopes(self, mock_aws_sso_oidc_token_response):
+ """
+ What it does: Verifies that scopes are NOT sent in refresh request.
+ Purpose: Per OAuth 2.0 RFC 6749 Section 6, scope is optional in refresh and
+ AWS SSO OIDC returns invalid_request if scope is sent.
+ """
+ print("Setup: Creating KiroAuthManager with scopes...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+ # Simulate scopes loaded from SQLite (this is what caused the bug)
+ manager._scopes = ["codewhisperer:completions", "codewhisperer:analysis"]
+
+ print("Setup: Mocking HTTP client...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: scope NOT in JSON payload...")
+ call_args = mock_client.post.call_args
+ json_payload = call_args[1].get('json', {})
+ print(f"Request JSON keys: {list(json_payload.keys())}")
+ assert 'scope' not in json_payload, "scope should NOT be sent in refresh request"
+
+ print("Verification: only required fields sent (camelCase)...")
+ expected_keys = {'grantType', 'clientId', 'clientSecret', 'refreshToken'}
+ print(f"Comparing keys: Expected {expected_keys}, Got {set(json_payload.keys())}")
+ assert set(json_payload.keys()) == expected_keys
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_works_without_scopes(self, mock_aws_sso_oidc_token_response):
+ """
+ What it does: Verifies refresh works when scopes are None.
+ Purpose: Ensure backward compatibility with credentials that don't have scopes.
+ """
+ print("Setup: Creating KiroAuthManager without scopes...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+ # Explicitly set scopes to None (default state)
+ manager._scopes = None
+
+ print("Setup: Mocking HTTP client...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: Token refreshed successfully...")
+ assert manager._access_token == "new_aws_sso_access_token"
+
+ print("Verification: scope NOT in request JSON payload...")
+ call_args = mock_client.post.call_args
+ json_payload = call_args[1].get('json', {})
+ assert 'scope' not in json_payload
+
+
+# =============================================================================
+# Tests for auth_type property and constructor with new parameters
+# =============================================================================
+
+class TestKiroAuthManagerAuthTypeProperty:
+ """Tests for auth_type property and constructor."""
+
+ def test_auth_type_property_returns_correct_value(self):
+ """
+ What it does: Verifies that auth_type property returns correct value.
+ Purpose: Ensure property works correctly.
+ """
+ print("Setup: Creating KiroAuthManager with KIRO_DESKTOP...")
+ manager_desktop = KiroAuthManager(refresh_token="test")
+
+ print("Verification: auth_type = KIRO_DESKTOP...")
+ assert manager_desktop.auth_type == AuthType.KIRO_DESKTOP
+
+ print("Setup: Creating KiroAuthManager with AWS_SSO_OIDC...")
+ manager_sso = KiroAuthManager(
+ refresh_token="test",
+ client_id="id",
+ client_secret="secret"
+ )
+
+ print("Verification: auth_type = AWS_SSO_OIDC...")
+ assert manager_sso.auth_type == AuthType.AWS_SSO_OIDC
+
+ def test_init_with_client_id_and_secret(self):
+ """
+ What it does: Verifies initialization with client_id and client_secret.
+ Purpose: Ensure parameters are stored in private fields.
+ """
+ print("Setup: Creating KiroAuthManager with client credentials...")
+ manager = KiroAuthManager(
+ refresh_token="test",
+ client_id="my_client_id",
+ client_secret="my_client_secret"
+ )
+
+ print("Verification: client_id stored...")
+ assert manager._client_id == "my_client_id"
+
+ print("Verification: client_secret stored...")
+ assert manager._client_secret == "my_client_secret"
+
+ def test_init_with_sqlite_db_parameter(self, temp_sqlite_db):
+ """
+ What it does: Verifies initialization with sqlite_db parameter.
+ Purpose: Ensure data is loaded from SQLite.
+ """
+ print(f"Setup: Creating KiroAuthManager with sqlite_db: {temp_sqlite_db}")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db)
+
+ print("Verification: Data loaded from SQLite...")
+ assert manager._access_token == "sqlite_access_token"
+ assert manager._refresh_token == "sqlite_refresh_token"
+
+ def test_detect_auth_type_kiro_desktop_when_only_client_secret(self):
+ """
+ What it does: Verifies type detection with only clientSecret (no id).
+ Purpose: Ensure KIRO_DESKTOP is used without clientId.
+ """
+ print("Setup: Creating KiroAuthManager with only client_secret...")
+ manager = KiroAuthManager(
+ refresh_token="test_token",
+ client_secret="test_client_secret"
+ )
+
+ print("Verification: auth_type = KIRO_DESKTOP (both id and secret required)...")
+ print(f"Comparing auth_type: Expected KIRO_DESKTOP, Got {manager.auth_type}")
+ assert manager.auth_type == AuthType.KIRO_DESKTOP
+
+
+# =============================================================================
+# Tests for SSO region separation (Issue #16)
+# =============================================================================
+
+class TestKiroAuthManagerSsoRegionSeparation:
+ """Tests for SSO region separation from API region (Issue #16 fix).
+
+ Background: CodeWhisperer API only exists in us-east-1, but users may have
+ SSO credentials from other regions (e.g., ap-southeast-1 for Singapore).
+ The fix separates SSO region (for OIDC token refresh) from API region.
+ """
+
+ def test_api_region_stays_us_east_1_when_loading_from_sqlite(self, temp_sqlite_db):
+ """
+ What it does: Verifies API region doesn't change when loading from SQLite.
+ Purpose: Ensure CodeWhisperer API calls go to us-east-1 regardless of SSO region.
+ """
+ print(f"Setup: Creating KiroAuthManager with SQLite (region=eu-west-1)...")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db)
+
+ print("Verification: API region stays at us-east-1...")
+ print(f"Comparing _region: Expected 'us-east-1', Got '{manager._region}'")
+ assert manager._region == "us-east-1"
+
+ print("Verification: api_host contains us-east-1...")
+ print(f"api_host: {manager._api_host}")
+ assert "us-east-1" in manager._api_host
+
+ print("Verification: q_host contains us-east-1...")
+ print(f"q_host: {manager._q_host}")
+ assert "us-east-1" in manager._q_host
+
+ def test_sso_region_stored_separately_from_api_region(self, temp_sqlite_db):
+ """
+ What it does: Verifies SSO region is stored in _sso_region field.
+ Purpose: Ensure SSO region is available for OIDC token refresh.
+ """
+ print(f"Setup: Creating KiroAuthManager with SQLite (region=eu-west-1)...")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db)
+
+ print("Verification: SSO region stored in _sso_region...")
+ print(f"Comparing _sso_region: Expected 'eu-west-1', Got '{manager._sso_region}'")
+ assert manager._sso_region == "eu-west-1"
+
+ print("Verification: API region is different from SSO region...")
+ assert manager._region != manager._sso_region
+
+ def test_sso_region_none_when_not_loaded_from_sqlite(self):
+ """
+ What it does: Verifies _sso_region is None when not loading from SQLite.
+ Purpose: Ensure backward compatibility with direct credential initialization.
+ """
+ print("Setup: Creating KiroAuthManager with direct credentials...")
+ manager = KiroAuthManager(
+ refresh_token="test_token",
+ region="us-east-1"
+ )
+
+ print("Verification: _sso_region is None...")
+ print(f"Comparing _sso_region: Expected None, Got '{manager._sso_region}'")
+ assert manager._sso_region is None
+
+ @pytest.mark.asyncio
+ async def test_oidc_refresh_uses_sso_region(self, mock_aws_sso_oidc_token_response):
+ """
+ What it does: Verifies OIDC token refresh uses SSO region, not API region.
+ Purpose: Ensure token refresh goes to correct regional OIDC endpoint.
+ """
+ print("Setup: Creating KiroAuthManager with SSO region=ap-southeast-1...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret",
+ region="us-east-1" # API region
+ )
+ # Simulate SSO region loaded from SQLite
+ manager._sso_region = "ap-southeast-1"
+
+ print("Setup: Mocking HTTP client...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: OIDC request went to SSO region (ap-southeast-1)...")
+ call_args = mock_client.post.call_args
+ url = call_args[0][0]
+ expected_url = "https://oidc.ap-southeast-1.amazonaws.com/token"
+ print(f"Comparing URL: Expected '{expected_url}', Got '{url}'")
+ assert url == expected_url
+ assert "ap-southeast-1" in url
+ assert "us-east-1" not in url
+
+ @pytest.mark.asyncio
+ async def test_oidc_refresh_falls_back_to_api_region_when_no_sso_region(self, mock_aws_sso_oidc_token_response):
+ """
+ What it does: Verifies OIDC refresh uses API region when SSO region not set.
+ Purpose: Ensure backward compatibility when _sso_region is None.
+ """
+ print("Setup: Creating KiroAuthManager without SSO region...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret",
+ region="eu-west-1" # API region (also used for OIDC when no SSO region)
+ )
+ # Ensure _sso_region is None
+ manager._sso_region = None
+
+ print("Setup: Mocking HTTP client...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: OIDC request fell back to API region (eu-west-1)...")
+ call_args = mock_client.post.call_args
+ url = call_args[0][0]
+ expected_url = "https://oidc.eu-west-1.amazonaws.com/token"
+ print(f"Comparing URL: Expected '{expected_url}', Got '{url}'")
+ assert url == expected_url
+
+ def test_api_hosts_not_updated_when_loading_from_sqlite(self, temp_sqlite_db):
+ """
+ What it does: Verifies API hosts don't change when loading from SQLite.
+ Purpose: Ensure all API calls go to us-east-1 where CodeWhisperer exists.
+ """
+ print(f"Setup: Creating KiroAuthManager with SQLite (region=eu-west-1)...")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db)
+
+ print("Verification: _api_host points to us-east-1...")
+ assert "us-east-1" in manager._api_host
+ assert "eu-west-1" not in manager._api_host
+
+ print("Verification: _q_host points to us-east-1...")
+ assert "us-east-1" in manager._q_host
+ assert "eu-west-1" not in manager._q_host
+
+ print("Verification: _refresh_url points to us-east-1...")
+ assert "us-east-1" in manager._refresh_url
+ assert "eu-west-1" not in manager._refresh_url
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_uses_memory_token_first(
+ self, mock_aws_sso_oidc_token_response
+ ):
+ """
+ What it does: Verifies that in-memory token is used first, not SQLite.
+ Purpose: Ensure container's successfully refreshed token is used (not overwritten by SQLite).
+ """
+ print("Setup: Creating KiroAuthManager with in-memory credentials...")
+ manager = KiroAuthManager(
+ refresh_token="memory_refresh_token",
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+ # Simulate SQLite path being set (but we won't actually use it)
+ manager._sqlite_db = "/fake/path/data.sqlite3"
+
+ print("Setup: Mocking HTTP client for successful refresh...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ # Patch _load_credentials_from_sqlite to track if it's called
+ with patch.object(manager, '_load_credentials_from_sqlite') as mock_load:
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: SQLite was NOT reloaded (success on first try)...")
+ mock_load.assert_not_called()
+
+ print("Verification: Request used in-memory token...")
+ call_args = mock_client.post.call_args
+ json_payload = call_args[1].get('json', {})
+ print(f"Refresh token sent: {json_payload.get('refreshToken')}")
+ assert json_payload.get('refreshToken') == "memory_refresh_token"
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_reloads_sqlite_on_400_error(
+ self, tmp_path, mock_aws_sso_oidc_token_response
+ ):
+ """
+ What it does: Verifies SQLite is reloaded and retry happens on 400 error.
+ Purpose: Pick up fresh tokens after kiro-cli re-login when in-memory token is stale.
+ """
+ import sqlite3
+ import json
+
+ # Setup: Create initial SQLite database
+ db_file = tmp_path / "data.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ # Initial token data (will become stale)
+ initial_token_data = {
+ "access_token": "old_access_token",
+ "refresh_token": "old_refresh_token",
+ "expires_at": "2099-01-01T00:00:00Z",
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:token", json.dumps(initial_token_data))
+ )
+
+ registration_data = {
+ "client_id": "test_client_id",
+ "client_secret": "test_client_secret",
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:device-registration", json.dumps(registration_data))
+ )
+
+ conn.commit()
+ conn.close()
+
+ print("Setup: Creating KiroAuthManager with SQLite...")
+ manager = KiroAuthManager(sqlite_db=str(db_file))
+
+ print("Verification: Initial refresh_token loaded...")
+ assert manager._refresh_token == "old_refresh_token"
+
+ # Simulate kiro-cli updating the SQLite with fresh tokens
+ print("Action: Simulating kiro-cli token refresh (updating SQLite)...")
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ new_token_data = {
+ "access_token": "new_access_token",
+ "refresh_token": "new_refresh_token_from_kiro_cli",
+ "expires_at": "2099-01-01T00:00:00Z",
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "UPDATE auth_kv SET value = ? WHERE key = ?",
+ (json.dumps(new_token_data), "codewhisperer:odic:token")
+ )
+ conn.commit()
+ conn.close()
+
+ # Manager still has old token in memory
+ print("Verification: Manager still has old refresh_token in memory...")
+ assert manager._refresh_token == "old_refresh_token"
+
+ # Mock HTTP client: first call fails with 400, second succeeds
+ print("Setup: Mocking HTTP client (first=400, second=200)...")
+
+ # First response: 400 error (stale token)
+ mock_error_response = AsyncMock()
+ mock_error_response.status_code = 400
+ mock_error_response.text = '{"error":"invalid_request","error_description":"Invalid request"}'
+ mock_error_response.json = Mock(return_value={"error": "invalid_request"})
+ mock_error_response.raise_for_status = Mock(
+ side_effect=httpx.HTTPStatusError(
+ "400 Bad Request",
+ request=Mock(),
+ response=mock_error_response
+ )
+ )
+
+ # Second response: success
+ mock_success_response = AsyncMock()
+ mock_success_response.status_code = 200
+ mock_success_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_success_response.raise_for_status = Mock()
+
+ call_count = 0
+ sent_tokens = []
+
+ async def mock_post(*args, **kwargs):
+ nonlocal call_count
+ call_count += 1
+ sent_tokens.append(kwargs.get('json', {}).get('refreshToken'))
+ if call_count == 1:
+ return mock_error_response
+ return mock_success_response
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = mock_post
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling _refresh_token_aws_sso_oidc...")
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: Two requests were made (retry on 400)...")
+ print(f"Call count: {call_count}")
+ assert call_count == 2, "Should retry after 400 error"
+
+ print("Verification: First request used OLD token from memory...")
+ print(f"First token sent: {sent_tokens[0]}")
+ assert sent_tokens[0] == "old_refresh_token"
+
+ print("Verification: Second request used NEW token from SQLite...")
+ print(f"Second token sent: {sent_tokens[1]}")
+ assert sent_tokens[1] == "new_refresh_token_from_kiro_cli"
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_no_retry_on_non_400_error(
+ self, mock_aws_sso_oidc_token_response
+ ):
+ """
+ What it does: Verifies that non-400 errors are not retried.
+ Purpose: Ensure only 400 (invalid_request) triggers SQLite reload.
+ """
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+ manager._sqlite_db = "/fake/path/data.sqlite3"
+
+ print("Setup: Mocking HTTP client with 500 error...")
+ mock_error_response = AsyncMock()
+ mock_error_response.status_code = 500
+ mock_error_response.text = "Internal Server Error"
+ mock_error_response.json = Mock(side_effect=Exception("Not JSON"))
+ mock_error_response.raise_for_status = Mock(
+ side_effect=httpx.HTTPStatusError(
+ "500 Internal Server Error",
+ request=Mock(),
+ response=mock_error_response
+ )
+ )
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_error_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ with patch.object(manager, '_load_credentials_from_sqlite') as mock_load:
+ print("Action: Calling _refresh_token_aws_sso_oidc (expecting 500 error)...")
+ with pytest.raises(httpx.HTTPStatusError) as exc_info:
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: 500 error was raised (not retried)...")
+ assert exc_info.value.response.status_code == 500
+
+ print("Verification: SQLite was NOT reloaded (500 != 400)...")
+ mock_load.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_no_retry_without_sqlite_db(
+ self, mock_aws_sso_oidc_token_response
+ ):
+ """
+ What it does: Verifies that 400 error is not retried when sqlite_db is not set.
+ Purpose: Ensure retry only happens when SQLite source is available.
+ """
+ print("Setup: Creating KiroAuthManager WITHOUT sqlite_db...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+ # Explicitly ensure no sqlite_db
+ manager._sqlite_db = None
+
+ print("Setup: Mocking HTTP client with 400 error...")
+ mock_error_response = AsyncMock()
+ mock_error_response.status_code = 400
+ mock_error_response.text = '{"error":"invalid_request"}'
+ mock_error_response.json = Mock(return_value={"error": "invalid_request"})
+ mock_error_response.raise_for_status = Mock(
+ side_effect=httpx.HTTPStatusError(
+ "400 Bad Request",
+ request=Mock(),
+ response=mock_error_response
+ )
+ )
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_error_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling _refresh_token_aws_sso_oidc (expecting 400 error)...")
+ with pytest.raises(httpx.HTTPStatusError) as exc_info:
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: 400 error was raised (no retry without sqlite_db)...")
+ assert exc_info.value.response.status_code == 400
+
+ print("Verification: Only one request was made...")
+ assert mock_client.post.call_count == 1
+
+
+# =============================================================================
+# Tests for is_token_expired() method
+# =============================================================================
+
+class TestKiroAuthManagerIsTokenExpired:
+ """Tests for is_token_expired() method.
+
+ This method checks if the token has actually expired (not just expiring soon).
+ Used for graceful degradation when refresh fails.
+ """
+
+ def test_is_token_expired_returns_true_when_no_expires_at(self):
+ """
+ What it does: Verifies that without expires_at token is considered expired.
+ Purpose: Ensure safe behavior when time information is missing.
+ """
+ print("Setup: Creating KiroAuthManager without expires_at...")
+ manager = KiroAuthManager(refresh_token="test_token")
+ manager._expires_at = None
+
+ print("Verification: is_token_expired returns True...")
+ result = manager.is_token_expired()
+ print(f"Comparing result: Expected True, Got {result}")
+ assert result is True
+
+ def test_is_token_expired_returns_true_when_expired(self):
+ """
+ What it does: Verifies that expired token is correctly identified.
+ Purpose: Ensure token in the past is considered expired.
+ """
+ print("Setup: Creating KiroAuthManager with expired token...")
+ manager = KiroAuthManager(refresh_token="test_token")
+ manager._expires_at = datetime.now(timezone.utc) - timedelta(hours=1)
+
+ print("Verification: is_token_expired returns True for expired token...")
+ result = manager.is_token_expired()
+ print(f"Comparing result: Expected True, Got {result}")
+ assert result is True
+
+ def test_is_token_expired_returns_false_when_valid(self):
+ """
+ What it does: Verifies that valid token is not considered expired.
+ Purpose: Ensure token in the future is not considered expired.
+ """
+ print("Setup: Creating KiroAuthManager with valid token...")
+ manager = KiroAuthManager(refresh_token="test_token")
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
+
+ print("Verification: is_token_expired returns False...")
+ result = manager.is_token_expired()
+ print(f"Comparing result: Expected False, Got {result}")
+ assert result is False
+
+ def test_is_token_expired_returns_false_when_expiring_soon_but_not_expired(self):
+ """
+ What it does: Verifies difference between expiring soon and actually expired.
+ Purpose: Ensure token expiring in 5 minutes is NOT considered expired yet.
+ """
+ print("Setup: Creating KiroAuthManager with token expiring in 5 minutes...")
+ manager = KiroAuthManager(refresh_token="test_token")
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(minutes=5)
+
+ print("Verification: is_token_expiring_soon returns True (within threshold)...")
+ assert manager.is_token_expiring_soon() is True
+
+ print("Verification: is_token_expired returns False (not actually expired)...")
+ result = manager.is_token_expired()
+ print(f"Comparing result: Expected False, Got {result}")
+ assert result is False
+
+
+# =============================================================================
+# Tests for graceful degradation in get_access_token() (SQLite mode)
+# =============================================================================
+
+class TestKiroAuthManagerGracefulDegradation:
+ """Tests for graceful degradation when refresh fails in SQLite mode.
+
+ Background: When kiro-cli refreshes tokens in memory without persisting to SQLite,
+ the refresh_token in SQLite becomes stale. The gateway should gracefully fall back
+ to using the access_token directly until it actually expires.
+ """
+
+ @pytest.mark.asyncio
+ async def test_get_access_token_reloads_sqlite_when_expiring_soon(self, tmp_path):
+ """
+ What it does: Verifies SQLite is reloaded when token is expiring soon.
+ Purpose: Pick up fresh tokens from kiro-cli before attempting refresh.
+ """
+ import sqlite3
+ import json
+
+ print("Setup: Creating SQLite database with fresh token...")
+ db_file = tmp_path / "data.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ # Token that expires in 1 hour (fresh)
+ fresh_token_data = {
+ "access_token": "fresh_access_token",
+ "refresh_token": "fresh_refresh_token",
+ "expires_at": (datetime.now(timezone.utc) + timedelta(hours=1)).isoformat(),
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:token", json.dumps(fresh_token_data))
+ )
+
+ registration_data = {
+ "client_id": "test_client_id",
+ "client_secret": "test_client_secret",
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:device-registration", json.dumps(registration_data))
+ )
+ conn.commit()
+ conn.close()
+
+ print("Setup: Creating KiroAuthManager with expiring token...")
+ manager = KiroAuthManager(sqlite_db=str(db_file))
+ # Simulate token expiring soon (within threshold)
+ manager._access_token = "old_expiring_token"
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(minutes=5)
+
+ print("Verification: Token is expiring soon...")
+ assert manager.is_token_expiring_soon() is True
+
+ print("Action: Calling get_access_token()...")
+ token = await manager.get_access_token()
+
+ print("Verification: Got fresh token from SQLite reload...")
+ print(f"Comparing token: Expected 'fresh_access_token', Got '{token}'")
+ assert token == "fresh_access_token"
+
+ @pytest.mark.asyncio
+ async def test_get_access_token_graceful_fallback_when_refresh_fails_but_token_valid(
+ self, tmp_path
+ ):
+ """
+ What it does: Verifies graceful fallback when refresh fails with 400 but access_token still valid.
+ Purpose: Use existing access_token until it actually expires when kiro-cli owns refresh.
+ """
+ import sqlite3
+ import json
+
+ print("Setup: Creating SQLite database...")
+ db_file = tmp_path / "data.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ # Token that is expiring soon but NOT expired yet
+ token_data = {
+ "access_token": "still_valid_access_token",
+ "refresh_token": "stale_refresh_token",
+ "expires_at": (datetime.now(timezone.utc) + timedelta(minutes=5)).isoformat(),
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:token", json.dumps(token_data))
+ )
+
+ registration_data = {
+ "client_id": "test_client_id",
+ "client_secret": "test_client_secret",
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:device-registration", json.dumps(registration_data))
+ )
+ conn.commit()
+ conn.close()
+
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(sqlite_db=str(db_file))
+
+ print("Verification: Token is expiring soon but NOT expired...")
+ assert manager.is_token_expiring_soon() is True
+ assert manager.is_token_expired() is False
+
+ print("Setup: Mocking HTTP client to return 400 twice (stale refresh token)...")
+ mock_error_response = AsyncMock()
+ mock_error_response.status_code = 400
+ mock_error_response.text = '{"error":"invalid_request"}'
+ mock_error_response.json = Mock(return_value={"error": "invalid_request"})
+ mock_error_response.raise_for_status = Mock(
+ side_effect=httpx.HTTPStatusError(
+ "400 Bad Request",
+ request=Mock(),
+ response=mock_error_response
+ )
+ )
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_error_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling get_access_token() (expecting graceful fallback)...")
+ token = await manager.get_access_token()
+
+ print("Verification: Got existing access_token (graceful fallback)...")
+ print(f"Comparing token: Expected 'still_valid_access_token', Got '{token}'")
+ assert token == "still_valid_access_token"
+
+ @pytest.mark.asyncio
+ async def test_get_access_token_raises_when_refresh_fails_and_token_expired(
+ self, tmp_path
+ ):
+ """
+ What it does: Verifies error is raised when refresh fails and access_token is expired.
+ Purpose: Clear error message when user needs to run 'kiro-cli login'.
+ """
+ import sqlite3
+ import json
+
+ print("Setup: Creating SQLite database with expired token...")
+ db_file = tmp_path / "data.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ # Token that is already expired
+ token_data = {
+ "access_token": "expired_access_token",
+ "refresh_token": "stale_refresh_token",
+ "expires_at": (datetime.now(timezone.utc) - timedelta(hours=1)).isoformat(),
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:token", json.dumps(token_data))
+ )
+
+ registration_data = {
+ "client_id": "test_client_id",
+ "client_secret": "test_client_secret",
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:device-registration", json.dumps(registration_data))
+ )
+ conn.commit()
+ conn.close()
+
+ print("Setup: Creating KiroAuthManager...")
+ manager = KiroAuthManager(sqlite_db=str(db_file))
+
+ print("Verification: Token is expired...")
+ assert manager.is_token_expired() is True
+
+ print("Setup: Mocking HTTP client to return 400 (stale refresh token)...")
+ mock_error_response = AsyncMock()
+ mock_error_response.status_code = 400
+ mock_error_response.text = '{"error":"invalid_request"}'
+ mock_error_response.json = Mock(return_value={"error": "invalid_request"})
+ mock_error_response.raise_for_status = Mock(
+ side_effect=httpx.HTTPStatusError(
+ "400 Bad Request",
+ request=Mock(),
+ response=mock_error_response
+ )
+ )
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_error_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling get_access_token() (expecting ValueError)...")
+ with pytest.raises(ValueError) as exc_info:
+ await manager.get_access_token()
+
+ print(f"Verification: ValueError raised with helpful message: {exc_info.value}")
+ assert "kiro-cli login" in str(exc_info.value).lower()
+
+ @pytest.mark.asyncio
+ async def test_get_access_token_non_sqlite_mode_propagates_400_error(self):
+ """
+ What it does: Verifies 400 error is propagated in non-SQLite mode.
+ Purpose: Ensure graceful degradation only applies to SQLite mode.
+ """
+ print("Setup: Creating KiroAuthManager WITHOUT sqlite_db...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ client_id="test_client_id",
+ client_secret="test_client_secret"
+ )
+ manager._access_token = "expiring_token"
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(minutes=5)
+
+ print("Verification: No sqlite_db set...")
+ assert manager._sqlite_db is None
+
+ print("Setup: Mocking HTTP client to return 400...")
+ mock_error_response = AsyncMock()
+ mock_error_response.status_code = 400
+ mock_error_response.text = '{"error":"invalid_request"}'
+ mock_error_response.json = Mock(return_value={"error": "invalid_request"})
+ mock_error_response.raise_for_status = Mock(
+ side_effect=httpx.HTTPStatusError(
+ "400 Bad Request",
+ request=Mock(),
+ response=mock_error_response
+ )
+ )
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_error_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling get_access_token() (expecting HTTPStatusError)...")
+ with pytest.raises(httpx.HTTPStatusError) as exc_info:
+ await manager.get_access_token()
+
+ print("Verification: 400 error was propagated (no graceful degradation)...")
+ assert exc_info.value.response.status_code == 400
+
+
+# =============================================================================
+# Tests for _save_credentials_to_sqlite() - NEW FUNCTIONALITY
+# =============================================================================
+
+class TestKiroAuthManagerSaveCredentialsToSqlite:
+ """Tests for _save_credentials_to_sqlite() method (Issue #43 fix).
+
+ Background: Gateway was not persisting refreshed tokens back to SQLite,
+ causing stale tokens to be reloaded after 1-2 hours.
+ """
+
+ def test_save_credentials_to_sqlite_writes_token_data(self, tmp_path):
+ """
+ What it does: Verifies that _save_credentials_to_sqlite writes token data.
+ Purpose: Ensure tokens are persisted to SQLite after refresh.
+ """
+ import sqlite3
+ import json
+
+ print("Setup: Creating SQLite database...")
+ db_file = tmp_path / "data.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ # Initial token data
+ initial_token_data = {
+ "access_token": "old_access_token",
+ "refresh_token": "old_refresh_token",
+ "expires_at": "2099-01-01T00:00:00Z",
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:token", json.dumps(initial_token_data))
+ )
+ conn.commit()
+ conn.close()
+
+ print("Setup: Creating KiroAuthManager with SQLite...")
+ manager = KiroAuthManager(sqlite_db=str(db_file))
+
+ print("Action: Updating tokens in memory...")
+ manager._access_token = "new_access_token"
+ manager._refresh_token = "new_refresh_token"
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
+
+ print("Action: Calling _save_credentials_to_sqlite()...")
+ manager._save_credentials_to_sqlite()
+
+ print("Verification: Reading SQLite to check saved data...")
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", ("codewhisperer:odic:token",))
+ row = cursor.fetchone()
+ conn.close()
+
+ assert row is not None
+ saved_data = json.loads(row[0])
+
+ print(f"Comparing access_token: Expected 'new_access_token', Got '{saved_data['access_token']}'")
+ assert saved_data['access_token'] == "new_access_token"
+
+ print(f"Comparing refresh_token: Expected 'new_refresh_token', Got '{saved_data['refresh_token']}'")
+ assert saved_data['refresh_token'] == "new_refresh_token"
+
+ def test_save_credentials_to_sqlite_handles_missing_database(self, tmp_path):
+ """
+ What it does: Verifies handling of missing SQLite file.
+ Purpose: Ensure application doesn't crash when database is missing.
+ """
+ print("Setup: Creating KiroAuthManager with non-existent SQLite...")
+ non_existent_db = str(tmp_path / "non_existent.sqlite3")
+
+ manager = KiroAuthManager(
+ refresh_token="test_token",
+ sqlite_db=non_existent_db
+ )
+ manager._access_token = "new_token"
+
+ print("Action: Calling _save_credentials_to_sqlite() with missing database...")
+ # Should not raise exception
+ manager._save_credentials_to_sqlite()
+
+ print("Verification: No exception raised...")
+ assert True
+
+ def test_save_credentials_to_sqlite_returns_early_when_no_sqlite_db(self):
+ """
+ What it does: Verifies early return when sqlite_db is None.
+ Purpose: Ensure method is no-op when SQLite is not configured.
+ """
+ print("Setup: Creating KiroAuthManager without sqlite_db...")
+ manager = KiroAuthManager(refresh_token="test_token")
+ manager._sqlite_db = None
+ manager._access_token = "new_token"
+
+ print("Action: Calling _save_credentials_to_sqlite()...")
+ # Should return early without doing anything
+ manager._save_credentials_to_sqlite()
+
+ print("Verification: No exception raised...")
+ assert True
+
+
+# =============================================================================
+# Tests for token persistence after refresh (Issue #43 fix)
+# =============================================================================
+
+class TestKiroAuthManagerTokenPersistence:
+ """Tests for token persistence after refresh.
+
+ Background: After refresh, tokens must be saved to SQLite so they're
+ available after gateway restart or when reloaded.
+ """
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_aws_sso_oidc_saves_to_sqlite(self, tmp_path, mock_aws_sso_oidc_token_response):
+ """
+ What it does: Verifies tokens are saved to SQLite after AWS SSO OIDC refresh.
+ Purpose: Ensure refreshed tokens are persisted (Issue #43 fix).
+ """
+ import sqlite3
+ import json
+
+ print("Setup: Creating SQLite database...")
+ db_file = tmp_path / "data.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ initial_token_data = {
+ "access_token": "old_access_token",
+ "refresh_token": "old_refresh_token",
+ "expires_at": "2099-01-01T00:00:00Z",
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:token", json.dumps(initial_token_data))
+ )
+
+ registration_data = {
+ "client_id": "test_client_id",
+ "client_secret": "test_client_secret",
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:device-registration", json.dumps(registration_data))
+ )
+ conn.commit()
+ conn.close()
+
+ print("Setup: Creating KiroAuthManager with SQLite...")
+ manager = KiroAuthManager(sqlite_db=str(db_file))
+
+ print("Setup: Mocking HTTP client for successful refresh...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling _do_aws_sso_oidc_refresh()...")
+ await manager._do_aws_sso_oidc_refresh()
+
+ print("Verification: Tokens updated in memory...")
+ assert manager._access_token == "new_aws_sso_access_token"
+ assert manager._refresh_token == "new_aws_sso_refresh_token"
+
+ print("Verification: Reading SQLite to check persistence...")
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", ("codewhisperer:odic:token",))
+ row = cursor.fetchone()
+ conn.close()
+
+ assert row is not None
+ saved_data = json.loads(row[0])
+
+ print(f"Comparing saved access_token: Expected 'new_aws_sso_access_token', Got '{saved_data['access_token']}'")
+ assert saved_data['access_token'] == "new_aws_sso_access_token"
+
+ print(f"Comparing saved refresh_token: Expected 'new_aws_sso_refresh_token', Got '{saved_data['refresh_token']}'")
+ assert saved_data['refresh_token'] == "new_aws_sso_refresh_token"
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_kiro_desktop_saves_to_sqlite(self, tmp_path, mock_kiro_token_response):
+ """
+ What it does: Verifies tokens are saved to SQLite after Kiro Desktop refresh.
+ Purpose: Ensure consistency between both refresh methods.
+ """
+ import sqlite3
+ import json
+
+ print("Setup: Creating SQLite database...")
+ db_file = tmp_path / "data.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ initial_token_data = {
+ "access_token": "old_access_token",
+ "refresh_token": "old_refresh_token",
+ "expires_at": "2099-01-01T00:00:00Z",
+ "region": "us-east-1"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("codewhisperer:odic:token", json.dumps(initial_token_data))
+ )
+ conn.commit()
+ conn.close()
+
+ print("Setup: Creating KiroAuthManager with SQLite and Kiro Desktop auth...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ sqlite_db=str(db_file)
+ )
+
+ print("Setup: Mocking HTTP client for successful refresh...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_kiro_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling _refresh_token_kiro_desktop()...")
+ await manager._refresh_token_kiro_desktop()
+
+ print("Verification: Reading SQLite to check persistence...")
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", ("codewhisperer:odic:token",))
+ row = cursor.fetchone()
+ conn.close()
+
+ assert row is not None
+ saved_data = json.loads(row[0])
+
+ print(f"Comparing saved refresh_token: Expected 'new_refresh_token_xyz', Got '{saved_data['refresh_token']}'")
+ assert saved_data['refresh_token'] == "new_refresh_token_xyz"
+
+
+# =============================================================================
+# Tests for Social Login Support (kirocli:social:token)
+# =============================================================================
+
+class TestKiroAuthManagerSocialLogin:
+ """Tests for social login support (Google, GitHub, etc.).
+
+ Background: kiro-cli supports social login (Google, GitHub) for free-tier users.
+ These credentials are stored in SQLite with key 'kirocli:social:token' instead of
+ 'kirocli:odic:token'. Social login uses the same Kiro Desktop Auth endpoint
+ (no client_id/client_secret required).
+ """
+
+ def test_load_credentials_from_sqlite_social_token(self, temp_sqlite_db_social):
+ """
+ What it does: Verifies loading credentials from kirocli:social:token key.
+ Purpose: Ensure social login credentials are loaded correctly.
+ """
+ print(f"Setup: Creating KiroAuthManager with social login SQLite: {temp_sqlite_db_social}")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db_social)
+
+ print("Verification: access_token loaded from social key...")
+ print(f"Comparing access_token: Expected 'social_access_token', Got '{manager._access_token}'")
+ assert manager._access_token == "social_access_token"
+
+ print("Verification: refresh_token loaded from social key...")
+ print(f"Comparing refresh_token: Expected 'social_refresh_token', Got '{manager._refresh_token}'")
+ assert manager._refresh_token == "social_refresh_token"
+
+ print("Verification: profile_arn loaded...")
+ assert manager._profile_arn == "arn:aws:codewhisperer:us-east-1:123456789:profile/social"
+
+ def test_social_login_detected_as_kiro_desktop(self, temp_sqlite_db_social):
+ """
+ What it does: Verifies social login is detected as KIRO_DESKTOP auth type.
+ Purpose: Ensure social login uses Kiro Desktop Auth endpoint (no AWS SSO OIDC).
+ """
+ print(f"Setup: Creating KiroAuthManager with social login SQLite...")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db_social)
+
+ print("Verification: No client_id loaded (social login doesn't have it)...")
+ assert manager._client_id is None
+
+ print("Verification: No client_secret loaded...")
+ assert manager._client_secret is None
+
+ print("Verification: auth_type = KIRO_DESKTOP...")
+ print(f"Comparing auth_type: Expected KIRO_DESKTOP, Got {manager.auth_type}")
+ assert manager.auth_type == AuthType.KIRO_DESKTOP
+
+ def test_social_token_key_has_highest_priority(self, temp_sqlite_db_all_keys):
+ """
+ What it does: Verifies kirocli:social:token has highest priority.
+ Purpose: Ensure correct key is loaded when multiple keys exist.
+ """
+ print("Setup: Creating KiroAuthManager with database containing all three keys...")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db_all_keys)
+
+ print("Verification: Loaded from kirocli:social:token (highest priority)...")
+ print(f"Comparing access_token: Expected 'social_token', Got '{manager._access_token}'")
+ assert manager._access_token == "social_token"
+
+ print(f"Comparing refresh_token: Expected 'social_refresh', Got '{manager._refresh_token}'")
+ assert manager._refresh_token == "social_refresh"
+
+ print("Verification: _sqlite_token_key tracks source...")
+ print(f"Comparing _sqlite_token_key: Expected 'kirocli:social:token', Got '{manager._sqlite_token_key}'")
+ assert manager._sqlite_token_key == "kirocli:social:token"
+
+ def test_sqlite_token_key_tracked_for_social_login(self, temp_sqlite_db_social):
+ """
+ What it does: Verifies _sqlite_token_key is set when loading from social key.
+ Purpose: Ensure tokens are saved back to correct key after refresh.
+ """
+ print("Setup: Creating KiroAuthManager with social login SQLite...")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db_social)
+
+ print("Verification: _sqlite_token_key set to kirocli:social:token...")
+ print(f"Comparing _sqlite_token_key: Expected 'kirocli:social:token', Got '{manager._sqlite_token_key}'")
+ assert manager._sqlite_token_key == "kirocli:social:token"
+
+ def test_sqlite_token_key_tracked_for_odic(self, temp_sqlite_db):
+ """
+ What it does: Verifies _sqlite_token_key is set when loading from OIDC key.
+ Purpose: Ensure backward compatibility with existing OIDC credentials.
+ """
+ print("Setup: Creating KiroAuthManager with OIDC SQLite...")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db)
+
+ print("Verification: _sqlite_token_key set to codewhisperer:odic:token...")
+ print(f"Comparing _sqlite_token_key: Expected 'codewhisperer:odic:token', Got '{manager._sqlite_token_key}'")
+ assert manager._sqlite_token_key == "codewhisperer:odic:token"
+
+ def test_save_credentials_to_sqlite_uses_source_key(self, temp_sqlite_db_social):
+ """
+ What it does: Verifies tokens are saved back to the same key they were loaded from.
+ Purpose: Ensure social login tokens go to kirocli:social:token, not OIDC keys.
+ """
+ import sqlite3
+ import json
+
+ print("Setup: Creating KiroAuthManager with social login SQLite...")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db_social)
+
+ print("Verification: Loaded from kirocli:social:token...")
+ assert manager._sqlite_token_key == "kirocli:social:token"
+
+ print("Action: Updating tokens in memory...")
+ manager._access_token = "updated_social_access"
+ manager._refresh_token = "updated_social_refresh"
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
+
+ print("Action: Calling _save_credentials_to_sqlite()...")
+ manager._save_credentials_to_sqlite()
+
+ print("Verification: Reading SQLite to check saved data...")
+ conn = sqlite3.connect(temp_sqlite_db_social)
+ cursor = conn.cursor()
+
+ # Check that kirocli:social:token was updated
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", ("kirocli:social:token",))
+ row = cursor.fetchone()
+ conn.close()
+
+ assert row is not None
+ saved_data = json.loads(row[0])
+
+ print(f"Comparing saved access_token: Expected 'updated_social_access', Got '{saved_data['access_token']}'")
+ assert saved_data['access_token'] == "updated_social_access"
+
+ print(f"Comparing saved refresh_token: Expected 'updated_social_refresh', Got '{saved_data['refresh_token']}'")
+ assert saved_data['refresh_token'] == "updated_social_refresh"
+
+ @pytest.mark.asyncio
+ async def test_refresh_token_kiro_desktop_saves_to_social_key(
+ self, temp_sqlite_db_social, mock_kiro_token_response
+ ):
+ """
+ What it does: Verifies tokens are saved to kirocli:social:token after Kiro Desktop refresh.
+ Purpose: Ensure social login tokens persist correctly after refresh.
+ """
+ import sqlite3
+ import json
+
+ print("Setup: Creating KiroAuthManager with social login SQLite...")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db_social)
+
+ print("Verification: Loaded from kirocli:social:token...")
+ assert manager._sqlite_token_key == "kirocli:social:token"
+ assert manager.auth_type == AuthType.KIRO_DESKTOP
+
+ print("Setup: Mocking HTTP client for successful refresh...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_kiro_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling _refresh_token_kiro_desktop()...")
+ await manager._refresh_token_kiro_desktop()
+
+ print("Verification: Reading SQLite to check persistence...")
+ conn = sqlite3.connect(temp_sqlite_db_social)
+ cursor = conn.cursor()
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", ("kirocli:social:token",))
+ row = cursor.fetchone()
+ conn.close()
+
+ assert row is not None
+ saved_data = json.loads(row[0])
+
+ print(f"Comparing saved refresh_token: Expected 'new_refresh_token_xyz', Got '{saved_data['refresh_token']}'")
+ assert saved_data['refresh_token'] == "new_refresh_token_xyz"
+
+ def test_save_credentials_fallback_when_source_key_unknown(self, tmp_path):
+ """
+ What it does: Verifies fallback behavior when _sqlite_token_key is None.
+ Purpose: Ensure robustness when source key is not tracked.
+ """
+ import sqlite3
+ import json
+
+ print("Setup: Creating SQLite database with kirocli:social:token...")
+ db_file = tmp_path / "data_fallback.sqlite3"
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+
+ cursor.execute("""
+ CREATE TABLE auth_kv (
+ key TEXT PRIMARY KEY,
+ value TEXT
+ )
+ """)
+
+ token_data = {
+ "access_token": "old_token",
+ "refresh_token": "old_refresh",
+ "expires_at": "2099-01-01T00:00:00Z"
+ }
+ cursor.execute(
+ "INSERT INTO auth_kv (key, value) VALUES (?, ?)",
+ ("kirocli:social:token", json.dumps(token_data))
+ )
+ conn.commit()
+ conn.close()
+
+ print("Setup: Creating KiroAuthManager with direct credentials (not from SQLite)...")
+ manager = KiroAuthManager(
+ refresh_token="test_refresh",
+ sqlite_db=str(db_file)
+ )
+
+ # Simulate scenario where _sqlite_token_key is None (edge case)
+ manager._sqlite_token_key = None
+ manager._access_token = "new_fallback_token"
+ manager._refresh_token = "new_fallback_refresh"
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
+
+ print("Action: Calling _save_credentials_to_sqlite() with unknown source key...")
+ manager._save_credentials_to_sqlite()
+
+ print("Verification: Fallback should try all keys and update first match...")
+ conn = sqlite3.connect(str(db_file))
+ cursor = conn.cursor()
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", ("kirocli:social:token",))
+ row = cursor.fetchone()
+ conn.close()
+
+ assert row is not None
+ saved_data = json.loads(row[0])
+
+ print(f"Comparing saved access_token: Expected 'new_fallback_token', Got '{saved_data['access_token']}'")
+ assert saved_data['access_token'] == "new_fallback_token"
+
+ def test_social_login_no_device_registration_key(self, temp_sqlite_db_social):
+ """
+ What it does: Verifies social login works without device-registration key.
+ Purpose: Ensure social login doesn't require AWS SSO OIDC device registration.
+ """
+ import sqlite3
+
+ print("Setup: Verifying database has no device-registration key...")
+ conn = sqlite3.connect(temp_sqlite_db_social)
+ cursor = conn.cursor()
+ cursor.execute("SELECT COUNT(*) FROM auth_kv WHERE key LIKE '%device-registration%'")
+ count = cursor.fetchone()[0]
+ conn.close()
+
+ print(f"Verification: No device-registration keys found (count={count})...")
+ assert count == 0
+
+ print("Setup: Creating KiroAuthManager with social login SQLite...")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db_social)
+
+ print("Verification: Manager initialized successfully without device-registration...")
+ assert manager._access_token == "social_access_token"
+ assert manager._client_id is None
+ assert manager._client_secret is None
+
+ def test_provider_field_preserved_in_social_token(self, temp_sqlite_db_social):
+ """
+ What it does: Verifies provider field is preserved when saving social tokens.
+ Purpose: Ensure metadata like 'provider: google' is not lost.
+ """
+ import sqlite3
+ import json
+
+ print("Setup: Creating KiroAuthManager with social login SQLite...")
+ manager = KiroAuthManager(sqlite_db=temp_sqlite_db_social)
+
+ print("Action: Updating tokens and saving...")
+ manager._access_token = "new_social_token"
+ manager._refresh_token = "new_social_refresh"
+ manager._expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
+ manager._save_credentials_to_sqlite()
+
+ print("Verification: Reading SQLite to check provider field...")
+ conn = sqlite3.connect(temp_sqlite_db_social)
+ cursor = conn.cursor()
+ cursor.execute("SELECT value FROM auth_kv WHERE key = ?", ("kirocli:social:token",))
+ row = cursor.fetchone()
+ conn.close()
+
+ saved_data = json.loads(row[0])
+
+ # Note: provider field is NOT explicitly saved by gateway (it's metadata from kiro-cli)
+ # Gateway only saves: access_token, refresh_token, expires_at, region, scopes
+ # This is acceptable because provider is not needed for token refresh
+ print("Verification: Core token fields saved correctly...")
+ assert saved_data['access_token'] == "new_social_token"
+ assert saved_data['refresh_token'] == "new_social_refresh"
+
+
+# =============================================================================
+# Tests for Enterprise Kiro IDE Support (Issue #45)
+# =============================================================================
+
+class TestKiroAuthManagerEnterpriseIDE:
+ """Tests for Enterprise Kiro IDE support (IdC login with clientIdHash).
+
+ Background: Enterprise Kiro IDE uses AWS IAM Identity Center (IdC) for authentication.
+ Credentials are stored in JSON file with clientIdHash field that points to a separate
+ device registration file containing clientId and clientSecret.
+
+ This is different from:
+ - Personal Kiro IDE (social login): Uses Kiro Desktop Auth, no clientId/clientSecret
+ - kiro-cli (SQLite): Uses AWS SSO OIDC, credentials in SQLite database
+ """
+
+ def test_load_credentials_from_file_with_client_id_hash(self, temp_enterprise_ide_complete):
+ """
+ What it does: Verifies loading credentials from JSON file with clientIdHash.
+ Purpose: Ensure clientIdHash is detected and stored.
+ """
+ creds_file, device_reg_file = temp_enterprise_ide_complete
+
+ print(f"Setup: Creating KiroAuthManager with Enterprise IDE credentials: {creds_file}")
+ manager = KiroAuthManager(creds_file=creds_file)
+
+ print("Verification: clientIdHash loaded...")
+ print(f"Comparing _client_id_hash: Expected 'abc123def456', Got '{manager._client_id_hash}'")
+ assert manager._client_id_hash == "abc123def456"
+
+ print("Verification: Basic credentials loaded...")
+ assert manager._access_token == "enterprise_access_token"
+ assert manager._refresh_token == "enterprise_refresh_token"
+
+ def test_load_enterprise_device_registration_success(self, temp_enterprise_ide_complete):
+ """
+ What it does: Verifies successful loading of device registration.
+ Purpose: Ensure clientId and clientSecret are loaded from device registration file.
+ """
+ creds_file, device_reg_file = temp_enterprise_ide_complete
+
+ print("Setup: Creating KiroAuthManager with Enterprise IDE credentials...")
+ manager = KiroAuthManager(creds_file=creds_file)
+
+ print("Verification: clientId loaded from device registration...")
+ print(f"Comparing _client_id: Expected 'enterprise_client_id_12345', Got '{manager._client_id}'")
+ assert manager._client_id == "enterprise_client_id_12345"
+
+ print("Verification: clientSecret loaded from device registration...")
+ print(f"Comparing _client_secret: Expected 'enterprise_client_secret_67890', Got '{manager._client_secret}'")
+ assert manager._client_secret == "enterprise_client_secret_67890"
+
+ def test_enterprise_ide_detected_as_aws_sso_oidc(self, temp_enterprise_ide_complete):
+ """
+ What it does: Verifies Enterprise IDE is detected as AWS_SSO_OIDC auth type.
+ Purpose: Ensure correct authentication method is used (not Kiro Desktop Auth).
+ """
+ creds_file, device_reg_file = temp_enterprise_ide_complete
+
+ print("Setup: Creating KiroAuthManager with Enterprise IDE credentials...")
+ manager = KiroAuthManager(creds_file=creds_file)
+
+ print("Verification: auth_type = AWS_SSO_OIDC...")
+ print(f"Comparing auth_type: Expected AWS_SSO_OIDC, Got {manager.auth_type}")
+ assert manager.auth_type == AuthType.AWS_SSO_OIDC
+
+ def test_load_enterprise_device_registration_file_not_found(self, tmp_path, monkeypatch):
+ """
+ What it does: Verifies handling of missing device registration file.
+ Purpose: Ensure application doesn't crash when device registration is missing.
+ """
+ monkeypatch.setattr('pathlib.Path.home', lambda: tmp_path)
+
+ print("Setup: Creating credentials file with clientIdHash but no device registration...")
+ creds_file = tmp_path / "kiro-auth-token.json"
+ creds_data = {
+ "accessToken": "enterprise_access_token",
+ "refreshToken": "enterprise_refresh_token",
+ "expiresAt": "2099-01-01T00:00:00.000Z",
+ "region": "us-east-1",
+ "clientIdHash": "nonexistent_hash"
+ }
+ creds_file.write_text(json.dumps(creds_data))
+
+ print("Action: Creating KiroAuthManager...")
+ manager = KiroAuthManager(creds_file=str(creds_file))
+
+ print("Verification: clientIdHash stored...")
+ assert manager._client_id_hash == "nonexistent_hash"
+
+ print("Verification: clientId and clientSecret are None (file not found)...")
+ assert manager._client_id is None
+ assert manager._client_secret is None
+
+ print("Verification: auth_type = KIRO_DESKTOP (no client credentials)...")
+ assert manager.auth_type == AuthType.KIRO_DESKTOP
+
+ def test_load_enterprise_device_registration_invalid_json(self, tmp_path, monkeypatch):
+ """
+ What it does: Verifies handling of invalid JSON in device registration file.
+ Purpose: Ensure application doesn't crash on corrupted device registration.
+ """
+ monkeypatch.setattr('pathlib.Path.home', lambda: tmp_path)
+
+ print("Setup: Creating device registration file with invalid JSON...")
+ aws_dir = tmp_path / ".aws" / "sso" / "cache"
+ aws_dir.mkdir(parents=True, exist_ok=True)
+
+ device_reg_file = aws_dir / "invalid_hash.json"
+ device_reg_file.write_text("not a valid json {{{")
+
+ print("Setup: Creating credentials file...")
+ creds_file = tmp_path / "kiro-auth-token.json"
+ creds_data = {
+ "accessToken": "enterprise_access_token",
+ "refreshToken": "enterprise_refresh_token",
+ "expiresAt": "2099-01-01T00:00:00.000Z",
+ "region": "us-east-1",
+ "clientIdHash": "invalid_hash"
+ }
+ creds_file.write_text(json.dumps(creds_data))
+
+ print("Action: Creating KiroAuthManager (should handle error gracefully)...")
+ manager = KiroAuthManager(creds_file=str(creds_file))
+
+ print("Verification: clientId and clientSecret are None (JSON parse error)...")
+ assert manager._client_id is None
+ assert manager._client_secret is None
+
+ def test_load_enterprise_device_registration_missing_fields(self, tmp_path, monkeypatch):
+ """
+ What it does: Verifies handling of device registration without clientId/clientSecret.
+ Purpose: Ensure partial data doesn't cause crashes.
+ """
+ monkeypatch.setattr('pathlib.Path.home', lambda: tmp_path)
+
+ print("Setup: Creating device registration file without clientId/clientSecret...")
+ aws_dir = tmp_path / ".aws" / "sso" / "cache"
+ aws_dir.mkdir(parents=True, exist_ok=True)
+
+ device_reg_file = aws_dir / "partial_hash.json"
+ device_reg_data = {
+ "region": "us-east-1",
+ "someOtherField": "value"
+ }
+ device_reg_file.write_text(json.dumps(device_reg_data))
+
+ print("Setup: Creating credentials file...")
+ creds_file = tmp_path / "kiro-auth-token.json"
+ creds_data = {
+ "accessToken": "enterprise_access_token",
+ "refreshToken": "enterprise_refresh_token",
+ "expiresAt": "2099-01-01T00:00:00.000Z",
+ "region": "us-east-1",
+ "clientIdHash": "partial_hash"
+ }
+ creds_file.write_text(json.dumps(creds_data))
+
+ print("Action: Creating KiroAuthManager...")
+ manager = KiroAuthManager(creds_file=str(creds_file))
+
+ print("Verification: clientId and clientSecret are None (missing in file)...")
+ assert manager._client_id is None
+ assert manager._client_secret is None
+
+ @pytest.mark.asyncio
+ async def test_enterprise_ide_refresh_uses_json_format(
+ self, temp_enterprise_ide_complete, mock_aws_sso_oidc_token_response
+ ):
+ """
+ What it does: Verifies Enterprise IDE uses JSON format for token refresh.
+ Purpose: Ensure correct request format (not form-urlencoded).
+ """
+ creds_file, device_reg_file = temp_enterprise_ide_complete
+
+ print("Setup: Creating KiroAuthManager with Enterprise IDE credentials...")
+ manager = KiroAuthManager(creds_file=creds_file)
+
+ print("Verification: auth_type = AWS_SSO_OIDC...")
+ assert manager.auth_type == AuthType.AWS_SSO_OIDC
+
+ print("Setup: Mocking HTTP client...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling _refresh_token_aws_sso_oidc()...")
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: POST request made...")
+ mock_client.post.assert_called_once()
+
+ print("Verification: Request uses JSON format (not form-urlencoded)...")
+ call_args = mock_client.post.call_args
+ assert 'json' in call_args[1], "Request should use json= parameter"
+ assert 'data' not in call_args[1], "Request should NOT use data= parameter"
+
+ print("Verification: Content-Type = application/json...")
+ headers = call_args[1].get('headers', {})
+ assert headers.get('Content-Type') == 'application/json'
+
+ @pytest.mark.asyncio
+ async def test_enterprise_ide_refresh_uses_camel_case(
+ self, temp_enterprise_ide_complete, mock_aws_sso_oidc_token_response
+ ):
+ """
+ What it does: Verifies Enterprise IDE uses camelCase parameters.
+ Purpose: Ensure correct parameter naming (not snake_case).
+ """
+ creds_file, device_reg_file = temp_enterprise_ide_complete
+
+ print("Setup: Creating KiroAuthManager with Enterprise IDE credentials...")
+ manager = KiroAuthManager(creds_file=creds_file)
+
+ print("Setup: Mocking HTTP client...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling _refresh_token_aws_sso_oidc()...")
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: Request uses camelCase parameters...")
+ call_args = mock_client.post.call_args
+ json_payload = call_args[1].get('json', {})
+
+ print(f"JSON payload keys: {list(json_payload.keys())}")
+ assert 'grantType' in json_payload, "Should use grantType (camelCase)"
+ assert 'clientId' in json_payload, "Should use clientId (camelCase)"
+ assert 'clientSecret' in json_payload, "Should use clientSecret (camelCase)"
+ assert 'refreshToken' in json_payload, "Should use refreshToken (camelCase)"
+
+ print("Verification: NOT using snake_case...")
+ assert 'grant_type' not in json_payload, "Should NOT use grant_type (snake_case)"
+ assert 'client_id' not in json_payload, "Should NOT use client_id (snake_case)"
+ assert 'client_secret' not in json_payload, "Should NOT use client_secret (snake_case)"
+ assert 'refresh_token' not in json_payload, "Should NOT use refresh_token (snake_case)"
+
+ @pytest.mark.asyncio
+ async def test_enterprise_ide_refresh_uses_correct_endpoint(
+ self, temp_enterprise_ide_complete, mock_aws_sso_oidc_token_response
+ ):
+ """
+ What it does: Verifies Enterprise IDE uses AWS SSO OIDC endpoint.
+ Purpose: Ensure correct endpoint (not Kiro Desktop Auth).
+ """
+ creds_file, device_reg_file = temp_enterprise_ide_complete
+
+ print("Setup: Creating KiroAuthManager with Enterprise IDE credentials...")
+ manager = KiroAuthManager(creds_file=creds_file)
+
+ print("Setup: Mocking HTTP client...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Calling _refresh_token_aws_sso_oidc()...")
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: Request went to AWS SSO OIDC endpoint...")
+ call_args = mock_client.post.call_args
+ url = call_args[0][0]
+
+ print(f"Comparing URL: Expected AWS SSO OIDC endpoint, Got '{url}'")
+ assert "oidc" in url, "Should use AWS SSO OIDC endpoint"
+ assert "amazonaws.com" in url, "Should use AWS endpoint"
+ assert "/token" in url, "Should use /token endpoint"
+
+ print("Verification: NOT using Kiro Desktop Auth endpoint...")
+ assert "auth.desktop.kiro.dev" not in url, "Should NOT use Kiro Desktop Auth"
+
+ @pytest.mark.asyncio
+ async def test_enterprise_ide_full_refresh_flow(
+ self, temp_enterprise_ide_complete, mock_aws_sso_oidc_token_response
+ ):
+ """
+ What it does: Tests complete refresh flow for Enterprise IDE.
+ Purpose: Integration test covering load → refresh → verify.
+ """
+ creds_file, device_reg_file = temp_enterprise_ide_complete
+
+ print("Setup: Creating KiroAuthManager with Enterprise IDE credentials...")
+ manager = KiroAuthManager(creds_file=creds_file)
+
+ print("Verification: Initial state correct...")
+ assert manager._client_id_hash == "abc123def456"
+ assert manager._client_id == "enterprise_client_id_12345"
+ assert manager._client_secret == "enterprise_client_secret_67890"
+ assert manager.auth_type == AuthType.AWS_SSO_OIDC
+
+ print("Setup: Mocking HTTP client for successful refresh...")
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_aws_sso_oidc_token_response())
+ mock_response.raise_for_status = Mock()
+
+ with patch('kiro.auth.httpx.AsyncClient') as mock_client_class:
+ mock_client = AsyncMock()
+ mock_client.post = AsyncMock(return_value=mock_response)
+ mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+ mock_client.__aexit__ = AsyncMock(return_value=None)
+ mock_client_class.return_value = mock_client
+
+ print("Action: Refreshing token...")
+ await manager._refresh_token_aws_sso_oidc()
+
+ print("Verification: Tokens updated...")
+ assert manager._access_token == "new_aws_sso_access_token"
+ assert manager._refresh_token == "new_aws_sso_refresh_token"
+
+ print("Verification: Expiration time set...")
+ assert manager._expires_at is not None
+ assert manager._expires_at > datetime.now(timezone.utc)
+
+ def test_enterprise_ide_and_kiro_cli_use_same_format(self):
+ """
+ What it does: Verifies Enterprise IDE and kiro-cli use identical request format.
+ Purpose: Ensure architectural consistency (both use JSON with camelCase).
+ """
+ print("This test documents the architectural decision:")
+ print("Both Enterprise IDE (JSON file) and kiro-cli (SQLite) use:")
+ print(" - AWS SSO OIDC endpoint")
+ print(" - JSON format (Content-Type: application/json)")
+ print(" - camelCase parameters (grantType, clientId, etc.)")
+ print("")
+ print("The ONLY difference is where credentials are stored:")
+ print(" - Enterprise IDE: JSON file + device registration file")
+ print(" - kiro-cli: SQLite database")
+ print("")
+ print("This is verified by other tests in this class and")
+ print("TestKiroAuthManagerSsoRegionSeparation class.")
+ assert True # Documentation test
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_cache.py b/kiro-gateway/tests/unit/test_cache.py
new file mode 100644
index 0000000000000000000000000000000000000000..89396ec55b3feb6ab0b9149432a14c08eba5d857
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_cache.py
@@ -0,0 +1,437 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit-тесты для ModelInfoCache.
+Проверяет логику кэширования метаданных моделей.
+"""
+
+import asyncio
+import time
+import pytest
+
+from kiro.cache import ModelInfoCache
+from kiro.config import DEFAULT_MAX_INPUT_TOKENS
+
+
+class TestModelInfoCacheInitialization:
+ """Тесты инициализации ModelInfoCache."""
+
+ def test_initialization_creates_empty_cache(self):
+ """
+ Что он делает: Проверяет, что кэш создаётся пустым.
+ Цель: Убедиться в корректной инициализации.
+ """
+ print("Настройка: Создание ModelInfoCache...")
+ cache = ModelInfoCache()
+
+ print("Проверка: Кэш пуст при создании...")
+ print(f"Сравниваем is_empty(): Ожидалось True, Получено {cache.is_empty()}")
+ assert cache.is_empty() is True
+
+ print(f"Сравниваем size: Ожидалось 0, Получено {cache.size}")
+ assert cache.size == 0
+
+ def test_initialization_with_custom_ttl(self):
+ """
+ Что он делает: Проверяет создание кэша с кастомным TTL.
+ Цель: Убедиться, что TTL можно настроить.
+ """
+ print("Настройка: Создание ModelInfoCache с TTL=7200...")
+ cache = ModelInfoCache(cache_ttl=7200)
+
+ print("Проверка: TTL установлен корректно...")
+ print(f"Сравниваем _cache_ttl: Ожидалось 7200, Получено {cache._cache_ttl}")
+ assert cache._cache_ttl == 7200
+
+ def test_initialization_last_update_is_none(self):
+ """
+ Что он делает: Проверяет, что last_update_time изначально None.
+ Цель: Убедиться, что время обновления не установлено до первого update.
+ """
+ print("Настройка: Создание ModelInfoCache...")
+ cache = ModelInfoCache()
+
+ print("Проверка: last_update_time изначально None...")
+ print(f"Сравниваем last_update_time: Ожидалось None, Получено {cache.last_update_time}")
+ assert cache.last_update_time is None
+
+
+class TestModelInfoCacheUpdate:
+ """Тесты обновления кэша."""
+
+ @pytest.mark.asyncio
+ async def test_update_populates_cache(self, sample_models_data):
+ """
+ Что он делает: Проверяет заполнение кэша данными.
+ Цель: Убедиться, что update() корректно сохраняет модели.
+ """
+ print("Настройка: Создание ModelInfoCache...")
+ cache = ModelInfoCache()
+
+ print(f"Действие: Обновление кэша с {len(sample_models_data)} моделями...")
+ await cache.update(sample_models_data)
+
+ print("Проверка: Кэш заполнен...")
+ print(f"Сравниваем is_empty(): Ожидалось False, Получено {cache.is_empty()}")
+ assert cache.is_empty() is False
+
+ print(f"Сравниваем size: Ожидалось {len(sample_models_data)}, Получено {cache.size}")
+ assert cache.size == len(sample_models_data)
+
+ @pytest.mark.asyncio
+ async def test_update_sets_last_update_time(self, sample_models_data):
+ """
+ Что он делает: Проверяет установку времени последнего обновления.
+ Цель: Убедиться, что last_update_time устанавливается после update.
+ """
+ print("Настройка: Создание ModelInfoCache...")
+ cache = ModelInfoCache()
+
+ before_update = time.time()
+ print(f"Действие: Обновление кэша (время до: {before_update})...")
+ await cache.update(sample_models_data)
+ after_update = time.time()
+
+ print("Проверка: last_update_time установлен в разумных пределах...")
+ print(f"last_update_time: {cache.last_update_time}")
+ assert cache.last_update_time is not None
+ assert before_update <= cache.last_update_time <= after_update
+
+ @pytest.mark.asyncio
+ async def test_update_replaces_existing_data(self, sample_models_data):
+ """
+ Что он делает: Проверяет замену данных при повторном update.
+ Цель: Убедиться, что старые данные полностью заменяются.
+ """
+ print("Настройка: Создание ModelInfoCache и первое обновление...")
+ cache = ModelInfoCache()
+ await cache.update(sample_models_data)
+
+ print("Действие: Обновление с новыми данными...")
+ new_data = [{"modelId": "new-model", "tokenLimits": {"maxInputTokens": 50000}}]
+ await cache.update(new_data)
+
+ print("Проверка: Старые данные заменены...")
+ print(f"Сравниваем size: Ожидалось 1, Получено {cache.size}")
+ assert cache.size == 1
+
+ print("Проверка: Старая модель недоступна...")
+ assert cache.get("claude-sonnet-4") is None
+
+ print("Проверка: Новая модель доступна...")
+ assert cache.get("new-model") is not None
+
+ @pytest.mark.asyncio
+ async def test_update_with_empty_list(self):
+ """
+ Что он делает: Проверяет обновление пустым списком.
+ Цель: Убедиться, что кэш очищается при пустом update.
+ """
+ print("Настройка: Создание ModelInfoCache с данными...")
+ cache = ModelInfoCache()
+ await cache.update([{"modelId": "test-model"}])
+
+ print("Действие: Обновление пустым списком...")
+ await cache.update([])
+
+ print("Проверка: Кэш пуст...")
+ print(f"Сравниваем is_empty(): Ожидалось True, Получено {cache.is_empty()}")
+ assert cache.is_empty() is True
+
+
+class TestModelInfoCacheGet:
+ """Тесты получения данных из кэша."""
+
+ @pytest.mark.asyncio
+ async def test_get_returns_model_info(self, sample_models_data):
+ """
+ Что он делает: Проверяет получение информации о модели.
+ Цель: Убедиться, что get() возвращает корректные данные.
+ """
+ print("Настройка: Создание и заполнение кэша...")
+ cache = ModelInfoCache()
+ await cache.update(sample_models_data)
+
+ print("Действие: Получение информации о claude-sonnet-4...")
+ model_info = cache.get("claude-sonnet-4")
+
+ print("Проверка: Информация получена...")
+ print(f"model_info: {model_info}")
+ assert model_info is not None
+ assert model_info["modelId"] == "claude-sonnet-4"
+
+ @pytest.mark.asyncio
+ async def test_get_returns_none_for_unknown_model(self, sample_models_data):
+ """
+ Что он делает: Проверяет возврат None для неизвестной модели.
+ Цель: Убедиться, что get() не падает при отсутствии модели.
+ """
+ print("Настройка: Создание и заполнение кэша...")
+ cache = ModelInfoCache()
+ await cache.update(sample_models_data)
+
+ print("Действие: Получение информации о несуществующей модели...")
+ model_info = cache.get("non-existent-model")
+
+ print("Проверка: Возвращён None...")
+ print(f"Сравниваем model_info: Ожидалось None, Получено {model_info}")
+ assert model_info is None
+
+ def test_get_from_empty_cache(self):
+ """
+ Что он делает: Проверяет get() из пустого кэша.
+ Цель: Убедиться, что пустой кэш не вызывает ошибок.
+ """
+ print("Настройка: Создание пустого кэша...")
+ cache = ModelInfoCache()
+
+ print("Действие: Получение из пустого кэша...")
+ model_info = cache.get("any-model")
+
+ print("Проверка: Возвращён None...")
+ print(f"Сравниваем model_info: Ожидалось None, Получено {model_info}")
+ assert model_info is None
+
+
+class TestModelInfoCacheGetMaxInputTokens:
+ """Тесты получения maxInputTokens."""
+
+ @pytest.mark.asyncio
+ async def test_get_max_input_tokens_returns_value(self, sample_models_data):
+ """
+ Что он делает: Проверяет получение maxInputTokens для модели.
+ Цель: Убедиться, что значение извлекается из tokenLimits.
+ """
+ print("Настройка: Создание и заполнение кэша...")
+ cache = ModelInfoCache()
+ await cache.update(sample_models_data)
+
+ print("Действие: Получение maxInputTokens для claude-sonnet-4...")
+ max_tokens = cache.get_max_input_tokens("claude-sonnet-4")
+
+ print("Проверка: Значение корректно...")
+ print(f"Сравниваем max_tokens: Ожидалось 200000, Получено {max_tokens}")
+ assert max_tokens == 200000
+
+ @pytest.mark.asyncio
+ async def test_get_max_input_tokens_returns_default_for_unknown(self, sample_models_data):
+ """
+ Что он делает: Проверяет возврат дефолта для неизвестной модели.
+ Цель: Убедиться, что возвращается DEFAULT_MAX_INPUT_TOKENS.
+ """
+ print("Настройка: Создание и заполнение кэша...")
+ cache = ModelInfoCache()
+ await cache.update(sample_models_data)
+
+ print("Действие: Получение maxInputTokens для неизвестной модели...")
+ max_tokens = cache.get_max_input_tokens("unknown-model")
+
+ print("Проверка: Возвращён дефолт...")
+ print(f"Сравниваем max_tokens: Ожидалось {DEFAULT_MAX_INPUT_TOKENS}, Получено {max_tokens}")
+ assert max_tokens == DEFAULT_MAX_INPUT_TOKENS
+
+ @pytest.mark.asyncio
+ async def test_get_max_input_tokens_returns_default_when_no_token_limits(self):
+ """
+ Что он делает: Проверяет возврат дефолта при отсутствии tokenLimits.
+ Цель: Убедиться, что модель без tokenLimits не ломает логику.
+ """
+ print("Настройка: Создание кэша с моделью без tokenLimits...")
+ cache = ModelInfoCache()
+ await cache.update([{"modelId": "model-without-limits"}])
+
+ print("Действие: Получение maxInputTokens...")
+ max_tokens = cache.get_max_input_tokens("model-without-limits")
+
+ print("Проверка: Возвращён дефолт...")
+ print(f"Сравниваем max_tokens: Ожидалось {DEFAULT_MAX_INPUT_TOKENS}, Получено {max_tokens}")
+ assert max_tokens == DEFAULT_MAX_INPUT_TOKENS
+
+ @pytest.mark.asyncio
+ async def test_get_max_input_tokens_returns_default_when_max_input_is_none(self):
+ """
+ Что он делает: Проверяет возврат дефолта при maxInputTokens=None.
+ Цель: Убедиться, что None в tokenLimits обрабатывается корректно.
+ """
+ print("Настройка: Создание кэша с моделью с maxInputTokens=None...")
+ cache = ModelInfoCache()
+ await cache.update([{
+ "modelId": "model-with-null",
+ "tokenLimits": {"maxInputTokens": None}
+ }])
+
+ print("Действие: Получение maxInputTokens...")
+ max_tokens = cache.get_max_input_tokens("model-with-null")
+
+ print("Проверка: Возвращён дефолт...")
+ print(f"Сравниваем max_tokens: Ожидалось {DEFAULT_MAX_INPUT_TOKENS}, Получено {max_tokens}")
+ assert max_tokens == DEFAULT_MAX_INPUT_TOKENS
+
+
+class TestModelInfoCacheIsEmpty:
+ """Тесты проверки пустоты кэша."""
+
+ def test_is_empty_returns_true_for_new_cache(self):
+ """
+ Что он делает: Проверяет is_empty() для нового кэша.
+ Цель: Убедиться, что новый кэш считается пустым.
+ """
+ print("Настройка: Создание нового кэша...")
+ cache = ModelInfoCache()
+
+ print("Проверка: is_empty() возвращает True...")
+ print(f"Сравниваем is_empty(): Ожидалось True, Получено {cache.is_empty()}")
+ assert cache.is_empty() is True
+
+ @pytest.mark.asyncio
+ async def test_is_empty_returns_false_after_update(self, sample_models_data):
+ """
+ Что он делает: Проверяет is_empty() после заполнения.
+ Цель: Убедиться, что заполненный кэш не считается пустым.
+ """
+ print("Настройка: Создание и заполнение кэша...")
+ cache = ModelInfoCache()
+ await cache.update(sample_models_data)
+
+ print("Проверка: is_empty() возвращает False...")
+ print(f"Сравниваем is_empty(): Ожидалось False, Получено {cache.is_empty()}")
+ assert cache.is_empty() is False
+
+
+class TestModelInfoCacheIsStale:
+ """Тесты проверки устаревания кэша."""
+
+ def test_is_stale_returns_true_for_new_cache(self):
+ """
+ Что он делает: Проверяет is_stale() для нового кэша.
+ Цель: Убедиться, что кэш без обновлений считается устаревшим.
+ """
+ print("Настройка: Создание нового кэша...")
+ cache = ModelInfoCache()
+
+ print("Проверка: is_stale() возвращает True...")
+ print(f"Сравниваем is_stale(): Ожидалось True, Получено {cache.is_stale()}")
+ assert cache.is_stale() is True
+
+ @pytest.mark.asyncio
+ async def test_is_stale_returns_false_after_recent_update(self, sample_models_data):
+ """
+ Что он делает: Проверяет is_stale() сразу после обновления.
+ Цель: Убедиться, что свежий кэш не считается устаревшим.
+ """
+ print("Настройка: Создание и заполнение кэша...")
+ cache = ModelInfoCache()
+ await cache.update(sample_models_data)
+
+ print("Проверка: is_stale() возвращает False...")
+ print(f"Сравниваем is_stale(): Ожидалось False, Получено {cache.is_stale()}")
+ assert cache.is_stale() is False
+
+ @pytest.mark.asyncio
+ async def test_is_stale_returns_true_after_ttl_expires(self, sample_models_data):
+ """
+ Что он делает: Проверяет is_stale() после истечения TTL.
+ Цель: Убедиться, что кэш считается устаревшим после TTL.
+ """
+ print("Настройка: Создание кэша с TTL=0.1 секунды...")
+ cache = ModelInfoCache(cache_ttl=0.1)
+ await cache.update(sample_models_data)
+
+ print("Действие: Ожидание истечения TTL...")
+ await asyncio.sleep(0.2)
+
+ print("Проверка: is_stale() возвращает True...")
+ print(f"Сравниваем is_stale(): Ожидалось True, Получено {cache.is_stale()}")
+ assert cache.is_stale() is True
+
+
+class TestModelInfoCacheGetAllModelIds:
+ """Тесты получения списка ID моделей."""
+
+ def test_get_all_model_ids_returns_empty_for_new_cache(self):
+ """
+ Что он делает: Проверяет get_all_model_ids() для пустого кэша.
+ Цель: Убедиться, что возвращается пустой список.
+ """
+ print("Настройка: Создание пустого кэша...")
+ cache = ModelInfoCache()
+
+ print("Действие: Получение списка ID моделей...")
+ model_ids = cache.get_all_model_ids()
+
+ print("Проверка: Список пуст...")
+ print(f"Сравниваем model_ids: Ожидалось [], Получено {model_ids}")
+ assert model_ids == []
+
+ @pytest.mark.asyncio
+ async def test_get_all_model_ids_returns_all_ids(self, sample_models_data):
+ """
+ Что он делает: Проверяет get_all_model_ids() для заполненного кэша.
+ Цель: Убедиться, что возвращаются все ID моделей.
+ """
+ print("Настройка: Создание и заполнение кэша...")
+ cache = ModelInfoCache()
+ await cache.update(sample_models_data)
+
+ print("Действие: Получение списка ID моделей...")
+ model_ids = cache.get_all_model_ids()
+
+ print("Проверка: Все ID присутствуют...")
+ expected_ids = [m["modelId"] for m in sample_models_data]
+ print(f"Сравниваем model_ids: Ожидалось {expected_ids}, Получено {model_ids}")
+ assert set(model_ids) == set(expected_ids)
+
+
+class TestModelInfoCacheThreadSafety:
+ """Тесты потокобезопасности кэша."""
+
+ @pytest.mark.asyncio
+ async def test_concurrent_updates_dont_corrupt_cache(self, sample_models_data):
+ """
+ Что он делает: Проверяет потокобезопасность при параллельных update.
+ Цель: Убедиться, что asyncio.Lock защищает от race conditions.
+ """
+ print("Настройка: Создание кэша...")
+ cache = ModelInfoCache()
+
+ async def update_with_data(data):
+ await cache.update(data)
+
+ print("Действие: 10 параллельных обновлений...")
+ tasks = []
+ for i in range(10):
+ data = [{"modelId": f"model-{i}", "tokenLimits": {"maxInputTokens": 100000 + i}}]
+ tasks.append(update_with_data(data))
+
+ await asyncio.gather(*tasks)
+
+ print("Проверка: Кэш содержит данные последнего обновления...")
+ # Из-за race condition, мы не знаем какое обновление было последним,
+ # но кэш должен содержать ровно одну модель
+ print(f"Сравниваем size: Ожидалось 1, Получено {cache.size}")
+ assert cache.size == 1
+
+ print("Проверка: Кэш не повреждён...")
+ model_ids = cache.get_all_model_ids()
+ assert len(model_ids) == 1
+ assert model_ids[0].startswith("model-")
+
+ @pytest.mark.asyncio
+ async def test_concurrent_reads_are_safe(self, sample_models_data):
+ """
+ Что он делает: Проверяет безопасность параллельных чтений.
+ Цель: Убедиться, что множественные get() не вызывают проблем.
+ """
+ print("Настройка: Создание и заполнение кэша...")
+ cache = ModelInfoCache()
+ await cache.update(sample_models_data)
+
+ print("Действие: 100 параллельных чтений...")
+ async def read_model():
+ return cache.get("claude-sonnet-4")
+
+ results = await asyncio.gather(*[read_model() for _ in range(100)])
+
+ print("Проверка: Все чтения вернули одинаковый результат...")
+ assert all(r is not None for r in results)
+ assert all(r["modelId"] == "claude-sonnet-4" for r in results)
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_config.py b/kiro-gateway/tests/unit/test_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..248bca56ea9ac06948326187be3211d3538ec0cd
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_config.py
@@ -0,0 +1,690 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for the configuration module.
+Verifies loading settings from environment variables.
+"""
+
+import pytest
+import os
+from unittest.mock import patch
+
+
+class TestLogLevelConfig:
+ """Tests for LOG_LEVEL configuration."""
+
+ def test_default_log_level_is_info(self):
+ """
+ What it does: Verifies that LOG_LEVEL defaults to INFO.
+ Purpose: Ensure that INFO is used when no environment variable is set.
+
+ Note: This test verifies the config.py code logic, not the actual
+ value from the .env file. We mock os.getenv to simulate
+ the absence of the environment variable.
+ """
+ print("Setup: Mocking os.getenv for LOG_LEVEL...")
+
+ # Create a mock that returns None for LOG_LEVEL (simulating missing variable)
+ original_getenv = os.getenv
+
+ def mock_getenv(key, default=None):
+ if key == "LOG_LEVEL":
+ print(f"os.getenv('{key}') -> None (mocked)")
+ return default # Return default, simulating missing variable
+ return original_getenv(key, default)
+
+ with patch.object(os, 'getenv', side_effect=mock_getenv):
+ # Reload config module with mocked getenv
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"LOG_LEVEL: {config_module.LOG_LEVEL}")
+ print(f"Comparing: Expected 'INFO', Got '{config_module.LOG_LEVEL}'")
+ assert config_module.LOG_LEVEL == "INFO"
+
+ # Restore module with real values
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ def test_log_level_from_environment(self):
+ """
+ What it does: Verifies loading LOG_LEVEL from environment variable.
+ Purpose: Ensure that the value from environment is used.
+ """
+ print("Setup: Setting LOG_LEVEL=DEBUG...")
+
+ with patch.dict(os.environ, {"LOG_LEVEL": "DEBUG"}):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"LOG_LEVEL: {config_module.LOG_LEVEL}")
+ print(f"Comparing: Expected 'DEBUG', Got '{config_module.LOG_LEVEL}'")
+ assert config_module.LOG_LEVEL == "DEBUG"
+
+ def test_log_level_uppercase_conversion(self):
+ """
+ What it does: Verifies LOG_LEVEL conversion to uppercase.
+ Purpose: Ensure that lowercase value is converted to uppercase.
+ """
+ print("Setup: Setting LOG_LEVEL=warning (lowercase)...")
+
+ with patch.dict(os.environ, {"LOG_LEVEL": "warning"}):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"LOG_LEVEL: {config_module.LOG_LEVEL}")
+ print(f"Comparing: Expected 'WARNING', Got '{config_module.LOG_LEVEL}'")
+ assert config_module.LOG_LEVEL == "WARNING"
+
+ def test_log_level_trace(self):
+ """
+ What it does: Verifies setting LOG_LEVEL=TRACE.
+ Purpose: Ensure that TRACE level is supported.
+ """
+ print("Setup: Setting LOG_LEVEL=TRACE...")
+
+ with patch.dict(os.environ, {"LOG_LEVEL": "TRACE"}):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"LOG_LEVEL: {config_module.LOG_LEVEL}")
+ assert config_module.LOG_LEVEL == "TRACE"
+
+ def test_log_level_error(self):
+ """
+ What it does: Verifies setting LOG_LEVEL=ERROR.
+ Purpose: Ensure that ERROR level is supported.
+ """
+ print("Setup: Setting LOG_LEVEL=ERROR...")
+
+ with patch.dict(os.environ, {"LOG_LEVEL": "ERROR"}):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"LOG_LEVEL: {config_module.LOG_LEVEL}")
+ assert config_module.LOG_LEVEL == "ERROR"
+
+ def test_log_level_critical(self):
+ """
+ What it does: Verifies setting LOG_LEVEL=CRITICAL.
+ Purpose: Ensure that CRITICAL level is supported.
+ """
+ print("Setup: Setting LOG_LEVEL=CRITICAL...")
+
+ with patch.dict(os.environ, {"LOG_LEVEL": "CRITICAL"}):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"LOG_LEVEL: {config_module.LOG_LEVEL}")
+ assert config_module.LOG_LEVEL == "CRITICAL"
+
+
+class TestToolDescriptionMaxLengthConfig:
+ """Tests for TOOL_DESCRIPTION_MAX_LENGTH configuration."""
+
+ def test_default_tool_description_max_length(self):
+ """
+ What it does: Verifies the default value for TOOL_DESCRIPTION_MAX_LENGTH.
+ Purpose: Ensure that 10000 is used by default.
+ """
+ print("Setup: Removing TOOL_DESCRIPTION_MAX_LENGTH from environment...")
+
+ with patch.dict(os.environ, {}, clear=False):
+ if "TOOL_DESCRIPTION_MAX_LENGTH" in os.environ:
+ del os.environ["TOOL_DESCRIPTION_MAX_LENGTH"]
+
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"TOOL_DESCRIPTION_MAX_LENGTH: {config_module.TOOL_DESCRIPTION_MAX_LENGTH}")
+ assert config_module.TOOL_DESCRIPTION_MAX_LENGTH == 10000
+
+ def test_tool_description_max_length_from_environment(self):
+ """
+ What it does: Verifies loading TOOL_DESCRIPTION_MAX_LENGTH from environment.
+ Purpose: Ensure that the value from environment is used.
+ """
+ print("Setup: Setting TOOL_DESCRIPTION_MAX_LENGTH=5000...")
+
+ with patch.dict(os.environ, {"TOOL_DESCRIPTION_MAX_LENGTH": "5000"}):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"TOOL_DESCRIPTION_MAX_LENGTH: {config_module.TOOL_DESCRIPTION_MAX_LENGTH}")
+ assert config_module.TOOL_DESCRIPTION_MAX_LENGTH == 5000
+
+ def test_tool_description_max_length_zero_disables(self):
+ """
+ What it does: Verifies that 0 disables the feature.
+ Purpose: Ensure that TOOL_DESCRIPTION_MAX_LENGTH=0 works.
+ """
+ print("Setup: Setting TOOL_DESCRIPTION_MAX_LENGTH=0...")
+
+ with patch.dict(os.environ, {"TOOL_DESCRIPTION_MAX_LENGTH": "0"}):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"TOOL_DESCRIPTION_MAX_LENGTH: {config_module.TOOL_DESCRIPTION_MAX_LENGTH}")
+ assert config_module.TOOL_DESCRIPTION_MAX_LENGTH == 0
+
+
+class TestTimeoutConfigurationWarning:
+ """Tests for _warn_timeout_configuration() function."""
+
+ def test_no_warning_when_first_token_less_than_streaming(self, capsys):
+ """
+ What it does: Verifies that warning is NOT shown with correct configuration.
+ Purpose: Ensure that no warning when FIRST_TOKEN_TIMEOUT < STREAMING_READ_TIMEOUT.
+ """
+ print("Setup: FIRST_TOKEN_TIMEOUT=15, STREAMING_READ_TIMEOUT=300...")
+
+ with patch.dict(os.environ, {
+ "FIRST_TOKEN_TIMEOUT": "15",
+ "STREAMING_READ_TIMEOUT": "300"
+ }):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ # Call the warning function
+ config_module._warn_timeout_configuration()
+
+ captured = capsys.readouterr()
+ print(f"Captured stderr: {captured.err}")
+
+ # Warning should NOT be shown
+ assert "WARNING" not in captured.err
+ assert "Suboptimal timeout configuration" not in captured.err
+
+ def test_warning_when_first_token_equals_streaming(self, capsys):
+ """
+ What it does: Verifies that warning is shown when timeouts are equal.
+ Purpose: Ensure that warning when FIRST_TOKEN_TIMEOUT == STREAMING_READ_TIMEOUT.
+ """
+ print("Setup: FIRST_TOKEN_TIMEOUT=300, STREAMING_READ_TIMEOUT=300...")
+
+ with patch.dict(os.environ, {
+ "FIRST_TOKEN_TIMEOUT": "300",
+ "STREAMING_READ_TIMEOUT": "300"
+ }):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ # Call the warning function
+ config_module._warn_timeout_configuration()
+
+ captured = capsys.readouterr()
+ print(f"Captured stderr: {captured.err}")
+
+ # Warning SHOULD be shown
+ assert "WARNING" in captured.err or "Suboptimal timeout configuration" in captured.err
+
+ def test_warning_when_first_token_greater_than_streaming(self, capsys):
+ """
+ What it does: Verifies that warning is shown when FIRST_TOKEN > STREAMING.
+ Purpose: Ensure that warning when FIRST_TOKEN_TIMEOUT > STREAMING_READ_TIMEOUT.
+ """
+ print("Setup: FIRST_TOKEN_TIMEOUT=500, STREAMING_READ_TIMEOUT=300...")
+
+ with patch.dict(os.environ, {
+ "FIRST_TOKEN_TIMEOUT": "500",
+ "STREAMING_READ_TIMEOUT": "300"
+ }):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ # Call the warning function
+ config_module._warn_timeout_configuration()
+
+ captured = capsys.readouterr()
+ print(f"Captured stderr: {captured.err}")
+
+ # Warning SHOULD be shown
+ assert "WARNING" in captured.err or "Suboptimal timeout configuration" in captured.err
+ # Verify that timeout values are mentioned in warning
+ assert "500" in captured.err
+ assert "300" in captured.err
+
+ def test_warning_contains_recommendation(self, capsys):
+ """
+ What it does: Verifies that warning contains a recommendation.
+ Purpose: Ensure that user receives useful information.
+ """
+ print("Setup: FIRST_TOKEN_TIMEOUT=400, STREAMING_READ_TIMEOUT=300...")
+
+ with patch.dict(os.environ, {
+ "FIRST_TOKEN_TIMEOUT": "400",
+ "STREAMING_READ_TIMEOUT": "300"
+ }):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ # Call the warning function
+ config_module._warn_timeout_configuration()
+
+ captured = capsys.readouterr()
+ print(f"Captured stderr: {captured.err}")
+
+ # Warning should contain recommendation
+ assert "Recommendation" in captured.err or "LESS than" in captured.err
+
+
+class TestAwsSsoOidcUrlConfig:
+ """Tests for AWS SSO OIDC URL configuration."""
+
+ def test_aws_sso_oidc_url_template_exists(self):
+ """
+ What it does: Verifies that AWS_SSO_OIDC_URL_TEMPLATE constant exists.
+ Purpose: Ensure the template is defined in config.
+ """
+ print("Setup: Importing config module...")
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print("Verification: AWS_SSO_OIDC_URL_TEMPLATE exists...")
+ assert hasattr(config_module, 'AWS_SSO_OIDC_URL_TEMPLATE')
+
+ print(f"AWS_SSO_OIDC_URL_TEMPLATE: {config_module.AWS_SSO_OIDC_URL_TEMPLATE}")
+ assert "oidc" in config_module.AWS_SSO_OIDC_URL_TEMPLATE
+ assert "amazonaws.com" in config_module.AWS_SSO_OIDC_URL_TEMPLATE
+ assert "{region}" in config_module.AWS_SSO_OIDC_URL_TEMPLATE
+
+ def test_get_aws_sso_oidc_url_returns_correct_url(self):
+ """
+ What it does: Verifies that get_aws_sso_oidc_url returns correct URL.
+ Purpose: Ensure the function formats URL correctly.
+ """
+ print("Setup: Importing get_aws_sso_oidc_url...")
+ from kiro.config import get_aws_sso_oidc_url
+
+ print("Action: Calling get_aws_sso_oidc_url('us-east-1')...")
+ url = get_aws_sso_oidc_url("us-east-1")
+
+ print(f"Verification: URL is correct...")
+ expected = "https://oidc.us-east-1.amazonaws.com/token"
+ print(f"Comparing: Expected '{expected}', Got '{url}'")
+ assert url == expected
+
+ def test_get_aws_sso_oidc_url_with_different_regions(self):
+ """
+ What it does: Verifies URL generation for different regions.
+ Purpose: Ensure the function works with various AWS regions.
+ """
+ print("Setup: Importing get_aws_sso_oidc_url...")
+ from kiro.config import get_aws_sso_oidc_url
+
+ test_cases = [
+ ("us-east-1", "https://oidc.us-east-1.amazonaws.com/token"),
+ ("eu-west-1", "https://oidc.eu-west-1.amazonaws.com/token"),
+ ("ap-southeast-1", "https://oidc.ap-southeast-1.amazonaws.com/token"),
+ ("us-west-2", "https://oidc.us-west-2.amazonaws.com/token"),
+ ]
+
+ for region, expected in test_cases:
+ print(f"Action: Calling get_aws_sso_oidc_url('{region}')...")
+ url = get_aws_sso_oidc_url(region)
+ print(f"Comparing: Expected '{expected}', Got '{url}'")
+ assert url == expected
+
+
+class TestServerHostConfig:
+ """Tests for SERVER_HOST configuration."""
+
+ def test_default_server_host_is_0_0_0_0(self):
+ """
+ What it does: Verifies that SERVER_HOST defaults to 0.0.0.0.
+ Purpose: Ensure that 0.0.0.0 (all interfaces) is used when no environment variable is set.
+ """
+ print("Setup: Removing SERVER_HOST from environment...")
+
+ with patch.dict(os.environ, {}, clear=False):
+ if "SERVER_HOST" in os.environ:
+ del os.environ["SERVER_HOST"]
+
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"SERVER_HOST: {config_module.SERVER_HOST}")
+ print(f"DEFAULT_SERVER_HOST: {config_module.DEFAULT_SERVER_HOST}")
+ print(f"Comparing: Expected '0.0.0.0', Got '{config_module.SERVER_HOST}'")
+ assert config_module.SERVER_HOST == "0.0.0.0"
+ assert config_module.DEFAULT_SERVER_HOST == "0.0.0.0"
+
+ def test_server_host_from_environment(self):
+ """
+ What it does: Verifies loading SERVER_HOST from environment variable.
+ Purpose: Ensure that the value from environment is used.
+ """
+ print("Setup: Setting SERVER_HOST=127.0.0.1...")
+
+ with patch.dict(os.environ, {"SERVER_HOST": "127.0.0.1"}):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"SERVER_HOST: {config_module.SERVER_HOST}")
+ print(f"Comparing: Expected '127.0.0.1', Got '{config_module.SERVER_HOST}'")
+ assert config_module.SERVER_HOST == "127.0.0.1"
+
+ def test_server_host_custom_value(self):
+ """
+ What it does: Verifies setting SERVER_HOST to a custom IP address.
+ Purpose: Ensure that any valid IP address can be used.
+ """
+ print("Setup: Setting SERVER_HOST=192.168.1.100...")
+
+ with patch.dict(os.environ, {"SERVER_HOST": "192.168.1.100"}):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"SERVER_HOST: {config_module.SERVER_HOST}")
+ assert config_module.SERVER_HOST == "192.168.1.100"
+
+
+class TestServerPortConfig:
+ """Tests for SERVER_PORT configuration."""
+
+ def test_default_server_port_is_8000(self):
+ """
+ What it does: Verifies that SERVER_PORT defaults to 8000.
+ Purpose: Ensure that 8000 is used when no environment variable is set.
+ """
+ print("Setup: Removing SERVER_PORT from environment...")
+
+ with patch.dict(os.environ, {}, clear=False):
+ if "SERVER_PORT" in os.environ:
+ del os.environ["SERVER_PORT"]
+
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"SERVER_PORT: {config_module.SERVER_PORT}")
+ print(f"DEFAULT_SERVER_PORT: {config_module.DEFAULT_SERVER_PORT}")
+ print(f"Comparing: Expected 8000, Got {config_module.SERVER_PORT}")
+ assert config_module.SERVER_PORT == 8000
+ assert config_module.DEFAULT_SERVER_PORT == 8000
+
+ def test_server_port_from_environment(self):
+ """
+ What it does: Verifies loading SERVER_PORT from environment variable.
+ Purpose: Ensure that the value from environment is used.
+ """
+ print("Setup: Setting SERVER_PORT=9000...")
+
+ with patch.dict(os.environ, {"SERVER_PORT": "9000"}):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"SERVER_PORT: {config_module.SERVER_PORT}")
+ print(f"Comparing: Expected 9000, Got {config_module.SERVER_PORT}")
+ assert config_module.SERVER_PORT == 9000
+
+ def test_server_port_custom_value(self):
+ """
+ What it does: Verifies setting SERVER_PORT to a custom port number.
+ Purpose: Ensure that any valid port number can be used.
+ """
+ print("Setup: Setting SERVER_PORT=3000...")
+
+ with patch.dict(os.environ, {"SERVER_PORT": "3000"}):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"SERVER_PORT: {config_module.SERVER_PORT}")
+ assert config_module.SERVER_PORT == 3000
+
+ def test_server_port_is_integer(self):
+ """
+ What it does: Verifies that SERVER_PORT is converted to integer.
+ Purpose: Ensure that string from environment is converted to int.
+ """
+ print("Setup: Setting SERVER_PORT=8080 (as string)...")
+
+ with patch.dict(os.environ, {"SERVER_PORT": "8080"}):
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print(f"SERVER_PORT: {config_module.SERVER_PORT}")
+ print(f"Type: {type(config_module.SERVER_PORT)}")
+ assert isinstance(config_module.SERVER_PORT, int)
+ assert config_module.SERVER_PORT == 8080
+
+
+class TestKiroCliDbFileConfig:
+ """Tests for KIRO_CLI_DB_FILE configuration."""
+
+ def test_kiro_cli_db_file_config_exists(self):
+ """
+ What it does: Verifies that KIRO_CLI_DB_FILE constant exists.
+ Purpose: Ensure the config parameter is defined.
+ """
+ print("Setup: Importing config module...")
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print("Verification: KIRO_CLI_DB_FILE exists...")
+ assert hasattr(config_module, 'KIRO_CLI_DB_FILE')
+
+ print(f"KIRO_CLI_DB_FILE: '{config_module.KIRO_CLI_DB_FILE}'")
+ # Default should be empty string
+ assert isinstance(config_module.KIRO_CLI_DB_FILE, str)
+
+ def test_kiro_cli_db_file_from_environment(self):
+ """
+ What it does: Verifies loading KIRO_CLI_DB_FILE from environment variable.
+ Purpose: Ensure the value from environment is used and normalized.
+ """
+ print("Setup: Importing config module...")
+ import importlib
+ import kiro.config as config_module
+
+ # Test that KIRO_CLI_DB_FILE is loaded and is a string
+ print(f"KIRO_CLI_DB_FILE: {config_module.KIRO_CLI_DB_FILE}")
+ assert isinstance(config_module.KIRO_CLI_DB_FILE, str)
+
+ # If value is set (not empty), verify it's a normalized path
+ if config_module.KIRO_CLI_DB_FILE:
+ # Path should be normalized (no raw ~ or forward slashes on Windows)
+ assert not config_module.KIRO_CLI_DB_FILE.startswith("~")
+ # Should be a valid path string (contains path separators or is absolute)
+ from pathlib import Path
+ path = Path(config_module.KIRO_CLI_DB_FILE)
+ # Path should be constructable (doesn't raise exception)
+ assert str(path) == config_module.KIRO_CLI_DB_FILE
+
+
+class TestFallbackModelsConfig:
+ """Tests for FALLBACK_MODELS configuration."""
+
+ def test_fallback_models_exists(self):
+ """
+ What it does: Verifies that FALLBACK_MODELS constant exists.
+ Purpose: Ensure the fallback model list is defined in config.
+ """
+ print("Setup: Importing config module...")
+ import importlib
+ import kiro.config as config_module
+ importlib.reload(config_module)
+
+ print("Verification: FALLBACK_MODELS exists...")
+ assert hasattr(config_module, 'FALLBACK_MODELS')
+
+ print(f"FALLBACK_MODELS type: {type(config_module.FALLBACK_MODELS)}")
+ assert isinstance(config_module.FALLBACK_MODELS, list)
+
+ def test_fallback_models_not_empty(self):
+ """
+ What it does: Verifies that FALLBACK_MODELS contains at least one model.
+ Purpose: Ensure fallback list is populated for DNS failure recovery.
+ """
+ print("Setup: Importing FALLBACK_MODELS...")
+ from kiro.config import FALLBACK_MODELS
+
+ print(f"FALLBACK_MODELS length: {len(FALLBACK_MODELS)}")
+ print(f"Comparing: Expected > 0, Got {len(FALLBACK_MODELS)}")
+ assert len(FALLBACK_MODELS) > 0
+
+ def test_fallback_models_structure(self):
+ """
+ What it does: Verifies that each fallback model has required modelId field.
+ Purpose: Ensure fallback models have correct structure for cache.update().
+ """
+ print("Setup: Importing FALLBACK_MODELS...")
+ from kiro.config import FALLBACK_MODELS
+
+ print(f"Action: Checking structure of {len(FALLBACK_MODELS)} models...")
+ for i, model in enumerate(FALLBACK_MODELS):
+ print(f"Checking model {i}: {model}")
+
+ print(f" Verification: model is dict...")
+ assert isinstance(model, dict), f"Model {i} is not a dict"
+
+ print(f" Verification: model has 'modelId'...")
+ assert "modelId" in model, f"Model {i} missing 'modelId'"
+
+ print(f" Verification: modelId is string...")
+ assert isinstance(model["modelId"], str), f"Model {i} modelId is not string"
+
+ print(f" Verification: modelId is not empty...")
+ assert len(model["modelId"]) > 0, f"Model {i} modelId is empty"
+
+ def test_fallback_models_contain_claude_models(self):
+ """
+ What it does: Verifies that fallback models include Claude models.
+ Purpose: Ensure fallback list contains expected Claude 4/4.5 models.
+ """
+ print("Setup: Importing FALLBACK_MODELS...")
+ from kiro.config import FALLBACK_MODELS
+
+ model_ids = [m["modelId"] for m in FALLBACK_MODELS]
+ print(f"Model IDs in fallback list: {model_ids}")
+
+ print("Verification: Contains at least one Claude model...")
+ has_claude = any("claude" in mid.lower() for mid in model_ids)
+ assert has_claude, "No Claude models in fallback list"
+
+ def test_fallback_models_use_dot_format(self):
+ """
+ What it does: Verifies that model IDs use dot format (e.g., claude-4.5).
+ Purpose: Ensure consistency with Kiro API format.
+ """
+ print("Setup: Importing FALLBACK_MODELS...")
+ from kiro.config import FALLBACK_MODELS
+
+ print("Action: Checking model ID format...")
+ for model in FALLBACK_MODELS:
+ model_id = model["modelId"]
+ print(f"Checking: {model_id}")
+
+ # If model has version number, it should use dot format
+ if any(char.isdigit() for char in model_id):
+ # Check for patterns like "4.5" or "4-5"
+ if "-4-5" in model_id or "-4-0" in model_id:
+ print(f" WARNING: {model_id} uses dash format instead of dot")
+ # This is acceptable but not ideal
+ pass
+
+
+class TestFallbackModelsIntegration:
+ """Integration tests for FALLBACK_MODELS with ModelResolver."""
+
+ @pytest.mark.asyncio
+ async def test_fallback_models_work_with_model_resolver(self):
+ """
+ What it does: Verifies that fallback models work with ModelResolver normalization.
+ Purpose: Ensure that model name normalization (claude-opus-4-5 → claude-opus-4.5)
+ works correctly with fallback models, just like with API models.
+ """
+ print("Setup: Importing FALLBACK_MODELS and creating cache...")
+ from kiro.config import FALLBACK_MODELS
+ from kiro.cache import ModelInfoCache
+ from kiro.model_resolver import ModelResolver
+
+ # Simulate DNS failure scenario - populate cache with fallback models
+ cache = ModelInfoCache()
+ await cache.update(FALLBACK_MODELS)
+
+ print(f"Cache populated with {cache.size} fallback models")
+ print(f"Model IDs in cache: {cache.get_all_model_ids()}")
+
+ # Create resolver
+ resolver = ModelResolver(cache=cache, hidden_models={})
+
+ print("\nAction: Testing normalization with dash format...")
+ # Test that dash format (claude-opus-4-5) is normalized and found
+ test_cases = [
+ ("claude-opus-4-5", "claude-opus-4.5"), # Dash → Dot
+ ("claude-sonnet-4-5", "claude-sonnet-4.5"), # Dash → Dot
+ ("claude-haiku-4-5", "claude-haiku-4.5"), # Dash → Dot
+ ]
+
+ for input_name, expected_normalized in test_cases:
+ print(f"\n Testing: {input_name} → {expected_normalized}")
+ resolution = resolver.resolve(input_name)
+
+ print(f" Resolution source: {resolution.source}")
+ print(f" Normalized: {resolution.normalized}")
+ print(f" Internal ID: {resolution.internal_id}")
+ print(f" Is verified: {resolution.is_verified}")
+
+ # Verify normalization happened
+ print(f" Comparing normalized: Expected '{expected_normalized}', Got '{resolution.normalized}'")
+ assert resolution.normalized == expected_normalized
+
+ # Verify model was found in cache (not passthrough)
+ print(f" Comparing source: Expected 'cache', Got '{resolution.source}'")
+ assert resolution.source == "cache", f"Model {input_name} should be found in fallback cache"
+
+ print(f" Comparing is_verified: Expected True, Got {resolution.is_verified}")
+ assert resolution.is_verified is True
+
+ @pytest.mark.asyncio
+ async def test_fallback_models_appear_in_available_models(self):
+ """
+ What it does: Verifies that fallback models appear in get_available_models().
+ Purpose: Ensure that /v1/models endpoint will show fallback models.
+ """
+ print("Setup: Importing FALLBACK_MODELS and creating cache...")
+ from kiro.config import FALLBACK_MODELS
+ from kiro.cache import ModelInfoCache
+ from kiro.model_resolver import ModelResolver
+
+ cache = ModelInfoCache()
+ await cache.update(FALLBACK_MODELS)
+
+ resolver = ModelResolver(cache=cache, hidden_models={})
+
+ print("Action: Getting available models...")
+ available = resolver.get_available_models()
+
+ print(f"Available models: {available}")
+ print(f"Comparing length: Expected {len(FALLBACK_MODELS)}, Got {len(available)}")
+ assert len(available) == len(FALLBACK_MODELS)
+
+ # Verify all fallback models are present
+ fallback_ids = {m["modelId"] for m in FALLBACK_MODELS}
+ available_set = set(available)
+
+ print(f"Comparing sets: Expected {fallback_ids}, Got {available_set}")
+ assert fallback_ids == available_set
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_converters_anthropic.py b/kiro-gateway/tests/unit/test_converters_anthropic.py
new file mode 100644
index 0000000000000000000000000000000000000000..578baa1aabbe2dcc85da6329203a7280715cfa4e
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_converters_anthropic.py
@@ -0,0 +1,1339 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for converters_anthropic module.
+
+Tests for Anthropic Messages API to Kiro format conversion:
+- Content extraction from Anthropic format
+- Tool results extraction
+- Tool uses extraction
+- Message conversion to unified format
+- Tool conversion to unified format
+- Full Anthropic → Kiro payload conversion
+"""
+
+import pytest
+from unittest.mock import patch, MagicMock
+
+from kiro.converters_anthropic import (
+ convert_anthropic_content_to_text,
+ extract_system_prompt,
+ extract_tool_results_from_anthropic_content,
+ extract_tool_uses_from_anthropic_content,
+ convert_anthropic_messages,
+ convert_anthropic_tools,
+ anthropic_to_kiro,
+)
+from kiro.converters_core import UnifiedMessage, UnifiedTool
+from kiro.models_anthropic import (
+ AnthropicMessagesRequest,
+ AnthropicMessage,
+ AnthropicTool,
+ TextContentBlock,
+ ToolUseContentBlock,
+ ToolResultContentBlock,
+ SystemContentBlock,
+)
+
+
+# ==================================================================================================
+# Tests for convert_anthropic_content_to_text
+# ==================================================================================================
+
+class TestConvertAnthropicContentToText:
+ """Tests for convert_anthropic_content_to_text function."""
+
+ def test_extracts_from_string(self):
+ """
+ What it does: Verifies text extraction from a string.
+ Purpose: Ensure string is returned as-is.
+ """
+ print("Setup: Simple string content...")
+ content = "Hello, World!"
+
+ print("Action: Extracting text...")
+ result = convert_anthropic_content_to_text(content)
+
+ print(f"Comparing result: Expected 'Hello, World!', Got '{result}'")
+ assert result == "Hello, World!"
+
+ def test_extracts_from_list_with_text_blocks(self):
+ """
+ What it does: Verifies extraction from list of text content blocks.
+ Purpose: Ensure Anthropic multimodal format is handled.
+ """
+ print("Setup: List with text content blocks...")
+ content = [
+ {"type": "text", "text": "Hello"},
+ {"type": "text", "text": " World"}
+ ]
+
+ print("Action: Extracting text...")
+ result = convert_anthropic_content_to_text(content)
+
+ print(f"Comparing result: Expected 'Hello World', Got '{result}'")
+ assert result == "Hello World"
+
+ def test_extracts_from_pydantic_text_blocks(self):
+ """
+ What it does: Verifies extraction from Pydantic TextContentBlock objects.
+ Purpose: Ensure Pydantic models are handled correctly.
+ """
+ print("Setup: List with Pydantic TextContentBlock objects...")
+ content = [
+ TextContentBlock(type="text", text="Part 1"),
+ TextContentBlock(type="text", text=" Part 2")
+ ]
+
+ print("Action: Extracting text...")
+ result = convert_anthropic_content_to_text(content)
+
+ print(f"Comparing result: Expected 'Part 1 Part 2', Got '{result}'")
+ assert result == "Part 1 Part 2"
+
+ def test_ignores_non_text_blocks(self):
+ """
+ What it does: Verifies that non-text blocks are ignored.
+ Purpose: Ensure tool_use and tool_result blocks don't contribute to text.
+ """
+ print("Setup: List with mixed content blocks...")
+ content = [
+ {"type": "text", "text": "Hello"},
+ {"type": "tool_use", "id": "call_123", "name": "test", "input": {}},
+ {"type": "text", "text": " World"}
+ ]
+
+ print("Action: Extracting text...")
+ result = convert_anthropic_content_to_text(content)
+
+ print(f"Comparing result: Expected 'Hello World', Got '{result}'")
+ assert result == "Hello World"
+
+ def test_handles_none(self):
+ """
+ What it does: Verifies None handling.
+ Purpose: Ensure None returns empty string.
+ """
+ print("Setup: None content...")
+
+ print("Action: Extracting text...")
+ result = convert_anthropic_content_to_text(None)
+
+ print(f"Comparing result: Expected '', Got '{result}'")
+ assert result == ""
+
+ def test_handles_empty_list(self):
+ """
+ What it does: Verifies empty list handling.
+ Purpose: Ensure empty list returns empty string.
+ """
+ print("Setup: Empty list...")
+ content = []
+
+ print("Action: Extracting text...")
+ result = convert_anthropic_content_to_text(content)
+
+ print(f"Comparing result: Expected '', Got '{result}'")
+ assert result == ""
+
+ def test_converts_other_types_to_string(self):
+ """
+ What it does: Verifies conversion of other types to string.
+ Purpose: Ensure numbers and other types are converted.
+ """
+ print("Setup: Number content...")
+ content = 42
+
+ print("Action: Extracting text...")
+ result = convert_anthropic_content_to_text(content)
+
+ print(f"Comparing result: Expected '42', Got '{result}'")
+ assert result == "42"
+
+
+# ==================================================================================================
+# Tests for extract_system_prompt
+# ==================================================================================================
+
+class TestExtractSystemPrompt:
+ """Tests for extract_system_prompt function (Support System commit)."""
+
+ def test_extracts_from_string(self):
+ """
+ What it does: Verifies extraction from simple string.
+ Purpose: Ensure string system prompt is returned as-is.
+ """
+ print("Setup: Simple string system prompt...")
+ system = "You are a helpful assistant."
+
+ print("Action: Extracting system prompt...")
+ result = extract_system_prompt(system)
+
+ print(f"Comparing result: Expected 'You are a helpful assistant.', Got '{result}'")
+ assert result == "You are a helpful assistant."
+
+ def test_extracts_from_list_with_text_blocks(self):
+ """
+ What it does: Verifies extraction from list of content blocks.
+ Purpose: Ensure Anthropic prompt caching format is handled.
+ """
+ print("Setup: List with text content blocks (prompt caching format)...")
+ system = [
+ {"type": "text", "text": "You are helpful."},
+ {"type": "text", "text": "Be concise."}
+ ]
+
+ print("Action: Extracting system prompt...")
+ result = extract_system_prompt(system)
+
+ print(f"Comparing result: Expected 'You are helpful.\\nBe concise.', Got '{result}'")
+ assert result == "You are helpful.\nBe concise."
+
+ def test_extracts_from_list_with_cache_control(self):
+ """
+ What it does: Verifies extraction ignores cache_control field.
+ Purpose: Ensure cache_control is stripped (not supported by Kiro).
+ """
+ print("Setup: List with cache_control (prompt caching format)...")
+ system = [
+ {
+ "type": "text",
+ "text": "You are a helpful assistant.",
+ "cache_control": {"type": "ephemeral"}
+ }
+ ]
+
+ print("Action: Extracting system prompt...")
+ result = extract_system_prompt(system)
+
+ print(f"Comparing result: Expected 'You are a helpful assistant.', Got '{result}'")
+ assert result == "You are a helpful assistant."
+
+ def test_extracts_from_pydantic_system_content_blocks(self):
+ """
+ What it does: Verifies extraction from Pydantic SystemContentBlock objects.
+ Purpose: Ensure Pydantic models are handled correctly.
+ """
+ print("Setup: List with Pydantic SystemContentBlock objects...")
+ system = [
+ SystemContentBlock(type="text", text="Part 1"),
+ SystemContentBlock(type="text", text="Part 2")
+ ]
+
+ print("Action: Extracting system prompt...")
+ result = extract_system_prompt(system)
+
+ print(f"Comparing result: Expected 'Part 1\\nPart 2', Got '{result}'")
+ assert result == "Part 1\nPart 2"
+
+ def test_handles_none(self):
+ """
+ What it does: Verifies None handling.
+ Purpose: Ensure None returns empty string.
+ """
+ print("Setup: None system prompt...")
+
+ print("Action: Extracting system prompt...")
+ result = extract_system_prompt(None)
+
+ print(f"Comparing result: Expected '', Got '{result}'")
+ assert result == ""
+
+ def test_handles_empty_list(self):
+ """
+ What it does: Verifies empty list handling.
+ Purpose: Ensure empty list returns empty string.
+ """
+ print("Setup: Empty list...")
+ system = []
+
+ print("Action: Extracting system prompt...")
+ result = extract_system_prompt(system)
+
+ print(f"Comparing result: Expected '', Got '{result}'")
+ assert result == ""
+
+ def test_handles_mixed_content_blocks(self):
+ """
+ What it does: Verifies handling of list with non-text blocks.
+ Purpose: Ensure only text blocks are extracted.
+ """
+ print("Setup: List with mixed content blocks...")
+ system = [
+ {"type": "text", "text": "Hello"},
+ {"type": "image", "source": {"type": "base64", "data": "..."}},
+ {"type": "text", "text": "World"}
+ ]
+
+ print("Action: Extracting system prompt...")
+ result = extract_system_prompt(system)
+
+ print(f"Comparing result: Expected 'Hello\\nWorld', Got '{result}'")
+ assert result == "Hello\nWorld"
+
+ def test_converts_other_types_to_string(self):
+ """
+ What it does: Verifies conversion of other types to string.
+ Purpose: Ensure numbers and other types are converted.
+ """
+ print("Setup: Number as system prompt...")
+ system = 42
+
+ print("Action: Extracting system prompt...")
+ result = extract_system_prompt(system)
+
+ print(f"Comparing result: Expected '42', Got '{result}'")
+ assert result == "42"
+
+ def test_handles_single_text_block(self):
+ """
+ What it does: Verifies extraction from single text block in list.
+ Purpose: Ensure single block list works correctly.
+ """
+ print("Setup: Single text block in list...")
+ system = [{"type": "text", "text": "Single block"}]
+
+ print("Action: Extracting system prompt...")
+ result = extract_system_prompt(system)
+
+ print(f"Comparing result: Expected 'Single block', Got '{result}'")
+ assert result == "Single block"
+
+ def test_handles_empty_text_in_block(self):
+ """
+ What it does: Verifies handling of empty text in content block.
+ Purpose: Ensure empty text doesn't cause errors.
+ """
+ print("Setup: Content block with empty text...")
+ system = [{"type": "text", "text": ""}]
+
+ print("Action: Extracting system prompt...")
+ result = extract_system_prompt(system)
+
+ print(f"Comparing result: Expected '', Got '{result}'")
+ assert result == ""
+
+ def test_handles_missing_text_key(self):
+ """
+ What it does: Verifies handling of content block without text key.
+ Purpose: Ensure missing text key doesn't cause errors.
+ """
+ print("Setup: Content block without text key...")
+ system = [{"type": "text"}]
+
+ print("Action: Extracting system prompt...")
+ result = extract_system_prompt(system)
+
+ print(f"Comparing result: Expected '', Got '{result}'")
+ assert result == ""
+
+
+# ==================================================================================================
+# Tests for extract_tool_results_from_anthropic_content
+# ==================================================================================================
+
+class TestExtractToolResultsFromAnthropicContent:
+ """Tests for extract_tool_results_from_anthropic_content function."""
+
+ def test_extracts_tool_result_from_dict(self):
+ """
+ What it does: Verifies extraction of tool result from dict content block.
+ Purpose: Ensure tool_result blocks are extracted correctly.
+ """
+ print("Setup: Content with tool_result block...")
+ content = [
+ {"type": "tool_result", "tool_use_id": "call_123", "content": "Result text"}
+ ]
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_anthropic_content(content)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["type"] == "tool_result"
+ assert result[0]["tool_use_id"] == "call_123"
+ assert result[0]["content"] == "Result text"
+
+ def test_extracts_tool_result_from_pydantic_model(self):
+ """
+ What it does: Verifies extraction from Pydantic ToolResultContentBlock.
+ Purpose: Ensure Pydantic models are handled correctly.
+ """
+ print("Setup: Content with Pydantic ToolResultContentBlock...")
+ content = [
+ ToolResultContentBlock(
+ type="tool_result",
+ tool_use_id="call_456",
+ content="Pydantic result"
+ )
+ ]
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_anthropic_content(content)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["tool_use_id"] == "call_456"
+ assert result[0]["content"] == "Pydantic result"
+
+ def test_extracts_multiple_tool_results(self):
+ """
+ What it does: Verifies extraction of multiple tool results.
+ Purpose: Ensure all tool_result blocks are extracted.
+ """
+ print("Setup: Content with multiple tool_results...")
+ content = [
+ {"type": "tool_result", "tool_use_id": "call_1", "content": "Result 1"},
+ {"type": "text", "text": "Some text"},
+ {"type": "tool_result", "tool_use_id": "call_2", "content": "Result 2"}
+ ]
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_anthropic_content(content)
+
+ print(f"Result: {result}")
+ assert len(result) == 2
+ assert result[0]["tool_use_id"] == "call_1"
+ assert result[1]["tool_use_id"] == "call_2"
+
+ def test_returns_empty_for_string_content(self):
+ """
+ What it does: Verifies empty list return for string content.
+ Purpose: Ensure string doesn't contain tool results.
+ """
+ print("Setup: String content...")
+ content = "Just a string"
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_anthropic_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_returns_empty_for_list_without_tool_results(self):
+ """
+ What it does: Verifies empty list return without tool_result blocks.
+ Purpose: Ensure regular elements are not extracted.
+ """
+ print("Setup: List without tool_result...")
+ content = [{"type": "text", "text": "Hello"}]
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_anthropic_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_handles_empty_content_in_tool_result(self):
+ """
+ What it does: Verifies handling of empty content in tool_result.
+ Purpose: Ensure empty content is replaced with "(empty result)".
+ """
+ print("Setup: Tool result with empty content...")
+ content = [
+ {"type": "tool_result", "tool_use_id": "call_123", "content": ""}
+ ]
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_anthropic_content(content)
+
+ print(f"Result: {result}")
+ assert result[0]["content"] == "(empty result)"
+
+ def test_handles_none_content_in_tool_result(self):
+ """
+ What it does: Verifies handling of None content in tool_result.
+ Purpose: Ensure None content is replaced with "(empty result)".
+ """
+ print("Setup: Tool result with None content...")
+ content = [
+ {"type": "tool_result", "tool_use_id": "call_123", "content": None}
+ ]
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_anthropic_content(content)
+
+ print(f"Result: {result}")
+ assert result[0]["content"] == "(empty result)"
+
+ def test_handles_list_content_in_tool_result(self):
+ """
+ What it does: Verifies handling of list content in tool_result.
+ Purpose: Ensure list content is converted to text.
+ """
+ print("Setup: Tool result with list content...")
+ content = [
+ {
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": [{"type": "text", "text": "List result"}]
+ }
+ ]
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_anthropic_content(content)
+
+ print(f"Result: {result}")
+ assert result[0]["content"] == "List result"
+
+ def test_skips_tool_result_without_tool_use_id(self):
+ """
+ What it does: Verifies that tool_result without tool_use_id is skipped.
+ Purpose: Ensure invalid tool_result blocks are ignored.
+ """
+ print("Setup: Tool result without tool_use_id...")
+ content = [
+ {"type": "tool_result", "content": "Result without ID"}
+ ]
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_anthropic_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+
+# ==================================================================================================
+# Tests for extract_tool_uses_from_anthropic_content
+# ==================================================================================================
+
+class TestExtractToolUsesFromAnthropicContent:
+ """Tests for extract_tool_uses_from_anthropic_content function."""
+
+ def test_extracts_tool_use_from_dict(self):
+ """
+ What it does: Verifies extraction of tool use from dict content block.
+ Purpose: Ensure tool_use blocks are extracted correctly.
+ """
+ print("Setup: Content with tool_use block...")
+ content = [
+ {"type": "tool_use", "id": "call_123", "name": "get_weather", "input": {"location": "Moscow"}}
+ ]
+
+ print("Action: Extracting tool uses...")
+ result = extract_tool_uses_from_anthropic_content(content)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["id"] == "call_123"
+ assert result[0]["type"] == "function"
+ assert result[0]["function"]["name"] == "get_weather"
+ assert result[0]["function"]["arguments"] == {"location": "Moscow"}
+
+ def test_extracts_tool_use_from_pydantic_model(self):
+ """
+ What it does: Verifies extraction from Pydantic ToolUseContentBlock.
+ Purpose: Ensure Pydantic models are handled correctly.
+ """
+ print("Setup: Content with Pydantic ToolUseContentBlock...")
+ content = [
+ ToolUseContentBlock(
+ type="tool_use",
+ id="call_456",
+ name="search",
+ input={"query": "test"}
+ )
+ ]
+
+ print("Action: Extracting tool uses...")
+ result = extract_tool_uses_from_anthropic_content(content)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["id"] == "call_456"
+ assert result[0]["function"]["name"] == "search"
+
+ def test_extracts_multiple_tool_uses(self):
+ """
+ What it does: Verifies extraction of multiple tool uses.
+ Purpose: Ensure all tool_use blocks are extracted.
+ """
+ print("Setup: Content with multiple tool_uses...")
+ content = [
+ {"type": "tool_use", "id": "call_1", "name": "tool1", "input": {}},
+ {"type": "text", "text": "Some text"},
+ {"type": "tool_use", "id": "call_2", "name": "tool2", "input": {}}
+ ]
+
+ print("Action: Extracting tool uses...")
+ result = extract_tool_uses_from_anthropic_content(content)
+
+ print(f"Result: {result}")
+ assert len(result) == 2
+ assert result[0]["id"] == "call_1"
+ assert result[1]["id"] == "call_2"
+
+ def test_returns_empty_for_string_content(self):
+ """
+ What it does: Verifies empty list return for string content.
+ Purpose: Ensure string doesn't contain tool uses.
+ """
+ print("Setup: String content...")
+ content = "Just a string"
+
+ print("Action: Extracting tool uses...")
+ result = extract_tool_uses_from_anthropic_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_returns_empty_for_list_without_tool_uses(self):
+ """
+ What it does: Verifies empty list return without tool_use blocks.
+ Purpose: Ensure regular elements are not extracted.
+ """
+ print("Setup: List without tool_use...")
+ content = [{"type": "text", "text": "Hello"}]
+
+ print("Action: Extracting tool uses...")
+ result = extract_tool_uses_from_anthropic_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_skips_tool_use_without_id(self):
+ """
+ What it does: Verifies that tool_use without id is skipped.
+ Purpose: Ensure invalid tool_use blocks are ignored.
+ """
+ print("Setup: Tool use without id...")
+ content = [
+ {"type": "tool_use", "name": "test", "input": {}}
+ ]
+
+ print("Action: Extracting tool uses...")
+ result = extract_tool_uses_from_anthropic_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_skips_tool_use_without_name(self):
+ """
+ What it does: Verifies that tool_use without name is skipped.
+ Purpose: Ensure invalid tool_use blocks are ignored.
+ """
+ print("Setup: Tool use without name...")
+ content = [
+ {"type": "tool_use", "id": "call_123", "input": {}}
+ ]
+
+ print("Action: Extracting tool uses...")
+ result = extract_tool_uses_from_anthropic_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+
+# ==================================================================================================
+# Tests for convert_anthropic_messages
+# ==================================================================================================
+
+class TestConvertAnthropicMessages:
+ """Tests for convert_anthropic_messages function."""
+
+ def test_converts_simple_user_message(self):
+ """
+ What it does: Verifies conversion of simple user message.
+ Purpose: Ensure basic user message is converted to UnifiedMessage.
+ """
+ print("Setup: Simple user message...")
+ messages = [
+ AnthropicMessage(role="user", content="Hello!")
+ ]
+
+ print("Action: Converting messages...")
+ result = convert_anthropic_messages(messages)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0].role == "user"
+ assert result[0].content == "Hello!"
+ assert result[0].tool_calls is None
+ assert result[0].tool_results is None
+
+ def test_converts_simple_assistant_message(self):
+ """
+ What it does: Verifies conversion of simple assistant message.
+ Purpose: Ensure basic assistant message is converted to UnifiedMessage.
+ """
+ print("Setup: Simple assistant message...")
+ messages = [
+ AnthropicMessage(role="assistant", content="Hi there!")
+ ]
+
+ print("Action: Converting messages...")
+ result = convert_anthropic_messages(messages)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0].role == "assistant"
+ assert result[0].content == "Hi there!"
+
+ def test_converts_user_message_with_content_blocks(self):
+ """
+ What it does: Verifies conversion of user message with content blocks.
+ Purpose: Ensure multimodal content is handled.
+ """
+ print("Setup: User message with content blocks...")
+ messages = [
+ AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "Part 1"},
+ {"type": "text", "text": " Part 2"}
+ ]
+ )
+ ]
+
+ print("Action: Converting messages...")
+ result = convert_anthropic_messages(messages)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0].content == "Part 1 Part 2"
+
+ def test_converts_assistant_message_with_tool_use(self):
+ """
+ What it does: Verifies conversion of assistant message with tool_use.
+ Purpose: Ensure tool_use blocks are extracted as tool_calls.
+ """
+ print("Setup: Assistant message with tool_use...")
+ messages = [
+ AnthropicMessage(
+ role="assistant",
+ content=[
+ {"type": "text", "text": "I'll check the weather"},
+ {"type": "tool_use", "id": "call_123", "name": "get_weather", "input": {"location": "Moscow"}}
+ ]
+ )
+ ]
+
+ print("Action: Converting messages...")
+ result = convert_anthropic_messages(messages)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0].role == "assistant"
+ assert result[0].content == "I'll check the weather"
+ assert result[0].tool_calls is not None
+ assert len(result[0].tool_calls) == 1
+ assert result[0].tool_calls[0]["function"]["name"] == "get_weather"
+
+ def test_converts_user_message_with_tool_result(self):
+ """
+ What it does: Verifies conversion of user message with tool_result.
+ Purpose: Ensure tool_result blocks are extracted as tool_results.
+ """
+ print("Setup: User message with tool_result...")
+ messages = [
+ AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "tool_result", "tool_use_id": "call_123", "content": "Weather: Sunny, 25°C"}
+ ]
+ )
+ ]
+
+ print("Action: Converting messages...")
+ result = convert_anthropic_messages(messages)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0].role == "user"
+ assert result[0].tool_results is not None
+ assert len(result[0].tool_results) == 1
+ assert result[0].tool_results[0]["tool_use_id"] == "call_123"
+
+ def test_converts_full_conversation(self):
+ """
+ What it does: Verifies conversion of full conversation.
+ Purpose: Ensure multi-turn conversation is converted correctly.
+ """
+ print("Setup: Full conversation...")
+ messages = [
+ AnthropicMessage(role="user", content="Hello"),
+ AnthropicMessage(role="assistant", content="Hi! How can I help?"),
+ AnthropicMessage(role="user", content="What's the weather?")
+ ]
+
+ print("Action: Converting messages...")
+ result = convert_anthropic_messages(messages)
+
+ print(f"Result: {result}")
+ assert len(result) == 3
+ assert result[0].role == "user"
+ assert result[1].role == "assistant"
+ assert result[2].role == "user"
+
+ def test_handles_empty_messages_list(self):
+ """
+ What it does: Verifies handling of empty messages list.
+ Purpose: Ensure empty list returns empty list.
+ """
+ print("Setup: Empty messages list...")
+ messages = []
+
+ print("Action: Converting messages...")
+ result = convert_anthropic_messages(messages)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ # ==================================================================================
+ # Image extraction tests (Issue #30 fix)
+ # ==================================================================================
+
+ def test_extracts_images_from_user_message(self):
+ """
+ What it does: Verifies that images are extracted from user messages.
+ Purpose: Ensure Anthropic image content blocks are converted to unified format.
+
+ This test verifies the fix for Issue #30 - 422 Validation Error for image content.
+ """
+ print("Setup: User message with image content block...")
+ # Base64 1x1 pixel JPEG
+ test_image_base64 = "/9j/4AAQSkZJRgABAQEASABIAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/wAALCAABAAEBAREA/8QAFAABAAAAAAAAAAAAAAAAAAAACf/EABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEAAD8AVN//2Q=="
+
+ messages = [
+ AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "What's in this image?"},
+ {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "media_type": "image/jpeg",
+ "data": test_image_base64
+ }
+ }
+ ]
+ )
+ ]
+
+ print("Action: Converting messages...")
+ result = convert_anthropic_messages(messages)
+
+ print(f"Result: {result}")
+ print(f"Images: {result[0].images}")
+
+ assert len(result) == 1
+ assert result[0].role == "user"
+ assert result[0].content == "What's in this image?"
+
+ print("Checking images field...")
+ assert result[0].images is not None, "images field should not be None"
+ assert len(result[0].images) == 1, f"Expected 1 image, got {len(result[0].images)}"
+
+ image = result[0].images[0]
+ print(f"Comparing image: Expected media_type='image/jpeg', Got '{image.get('media_type')}'")
+ assert image["media_type"] == "image/jpeg"
+
+ print(f"Comparing image data: Expected {test_image_base64[:20]}..., Got {image.get('data', '')[:20]}...")
+ assert image["data"] == test_image_base64
+
+ def test_images_only_extracted_from_user_role(self):
+ """
+ What it does: Verifies that images are only extracted from user messages.
+ Purpose: Ensure assistant messages don't have images extracted (they shouldn't contain images).
+ """
+ print("Setup: Conversation with image in user message only...")
+ test_image_base64 = "/9j/4AAQSkZJRgABAQEASABIAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/wAALCAABAAEBAREA/8QAFAABAAAAAAAAAAAAAAAAAAAACf/EABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEAAD8AVN//2Q=="
+
+ messages = [
+ AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "Describe this image"},
+ {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "media_type": "image/png",
+ "data": test_image_base64
+ }
+ }
+ ]
+ ),
+ AnthropicMessage(
+ role="assistant",
+ content="I can see a small image."
+ )
+ ]
+
+ print("Action: Converting messages...")
+ result = convert_anthropic_messages(messages)
+
+ print(f"Result: {result}")
+
+ print("Checking user message has images...")
+ assert result[0].images is not None
+ assert len(result[0].images) == 1
+
+ print("Checking assistant message has no images...")
+ assert result[1].images is None, "Assistant messages should not have images extracted"
+
+ def test_extracts_multiple_images_from_user_message(self):
+ """
+ What it does: Verifies extraction of multiple images from a single user message.
+ Purpose: Ensure all images in a message are extracted.
+ """
+ print("Setup: User message with multiple images...")
+ test_image_base64 = "/9j/4AAQSkZJRgABAQEASABIAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/wAALCAABAAEBAREA/8QAFAABAAAAAAAAAAAAAAAAAAAACf/EABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEAAD8AVN//2Q=="
+
+ messages = [
+ AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "Compare these images"},
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/jpeg", "data": test_image_base64}
+ },
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/png", "data": test_image_base64}
+ },
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/webp", "data": test_image_base64}
+ }
+ ]
+ )
+ ]
+
+ print("Action: Converting messages...")
+ result = convert_anthropic_messages(messages)
+
+ print(f"Result images count: {len(result[0].images) if result[0].images else 0}")
+
+ assert result[0].images is not None
+ assert len(result[0].images) == 3, f"Expected 3 images, got {len(result[0].images)}"
+
+ print("Checking image media types...")
+ media_types = [img["media_type"] for img in result[0].images]
+ print(f"Media types: {media_types}")
+ assert "image/jpeg" in media_types
+ assert "image/png" in media_types
+ assert "image/webp" in media_types
+
+ def test_counts_images_in_debug_log(self, caplog):
+ """
+ What it does: Verifies that image count is logged in debug message.
+ Purpose: Ensure logging includes image statistics for debugging.
+ """
+ import logging
+
+ print("Setup: User message with images for logging test...")
+ test_image_base64 = "/9j/4AAQSkZJRgABAQEASABIAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/wAALCAABAAEBAREA/8QAFAABAAAAAAAAAAAAAAAAAAAACf/EABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEAAD8AVN//2Q=="
+
+ messages = [
+ AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "Analyze this"},
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/jpeg", "data": test_image_base64}
+ },
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/png", "data": test_image_base64}
+ }
+ ]
+ )
+ ]
+
+ print("Action: Converting messages with logging enabled...")
+ with caplog.at_level(logging.DEBUG):
+ result = convert_anthropic_messages(messages)
+
+ print(f"Log records: {[r.message for r in caplog.records]}")
+
+ # Check that images were extracted
+ assert result[0].images is not None
+ assert len(result[0].images) == 2
+
+ # Note: loguru doesn't integrate with caplog by default
+ # The function logs "Converted X Anthropic messages: Y tool_calls, Z tool_results, W images"
+ # We verify the images are extracted correctly, which proves the counting works
+ print("Images extracted successfully - logging verification complete")
+
+
+# ==================================================================================================
+# Tests for convert_anthropic_tools
+# ==================================================================================================
+
+class TestConvertAnthropicTools:
+ """Tests for convert_anthropic_tools function."""
+
+ def test_returns_none_for_none(self):
+ """
+ What it does: Verifies handling of None.
+ Purpose: Ensure None returns None.
+ """
+ print("Setup: None tools...")
+
+ print("Action: Converting tools...")
+ result = convert_anthropic_tools(None)
+
+ print(f"Comparing result: Expected None, Got {result}")
+ assert result is None
+
+ def test_returns_none_for_empty_list(self):
+ """
+ What it does: Verifies handling of empty list.
+ Purpose: Ensure empty list returns None.
+ """
+ print("Setup: Empty tools list...")
+
+ print("Action: Converting tools...")
+ result = convert_anthropic_tools([])
+
+ print(f"Comparing result: Expected None, Got {result}")
+ assert result is None
+
+ def test_converts_tool_from_pydantic_model(self):
+ """
+ What it does: Verifies conversion of Pydantic AnthropicTool.
+ Purpose: Ensure Pydantic models are converted to UnifiedTool.
+ """
+ print("Setup: Pydantic AnthropicTool...")
+ tools = [
+ AnthropicTool(
+ name="get_weather",
+ description="Get weather for a location",
+ input_schema={"type": "object", "properties": {"location": {"type": "string"}}}
+ )
+ ]
+
+ print("Action: Converting tools...")
+ result = convert_anthropic_tools(tools)
+
+ print(f"Result: {result}")
+ assert result is not None
+ assert len(result) == 1
+ assert isinstance(result[0], UnifiedTool)
+ assert result[0].name == "get_weather"
+ assert result[0].description == "Get weather for a location"
+ assert result[0].input_schema == {"type": "object", "properties": {"location": {"type": "string"}}}
+
+ def test_converts_tool_from_dict(self):
+ """
+ What it does: Verifies conversion of dict tool.
+ Purpose: Ensure dict tools are converted to UnifiedTool.
+ """
+ print("Setup: Dict tool...")
+ tools = [
+ {
+ "name": "search",
+ "description": "Search the web",
+ "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}
+ }
+ ]
+
+ print("Action: Converting tools...")
+ result = convert_anthropic_tools(tools)
+
+ print(f"Result: {result}")
+ assert result is not None
+ assert len(result) == 1
+ assert result[0].name == "search"
+ assert result[0].description == "Search the web"
+
+ def test_converts_multiple_tools(self):
+ """
+ What it does: Verifies conversion of multiple tools.
+ Purpose: Ensure all tools are converted.
+ """
+ print("Setup: Multiple tools...")
+ tools = [
+ AnthropicTool(name="tool1", description="Tool 1", input_schema={}),
+ AnthropicTool(name="tool2", description="Tool 2", input_schema={})
+ ]
+
+ print("Action: Converting tools...")
+ result = convert_anthropic_tools(tools)
+
+ print(f"Result: {result}")
+ assert result is not None
+ assert len(result) == 2
+ assert result[0].name == "tool1"
+ assert result[1].name == "tool2"
+
+ def test_handles_tool_without_description(self):
+ """
+ What it does: Verifies handling of tool without description.
+ Purpose: Ensure None description is preserved.
+ """
+ print("Setup: Tool without description...")
+ tools = [
+ AnthropicTool(name="test_tool", input_schema={})
+ ]
+
+ print("Action: Converting tools...")
+ result = convert_anthropic_tools(tools)
+
+ print(f"Result: {result}")
+ assert result is not None
+ assert result[0].description is None
+
+
+# ==================================================================================================
+# Tests for anthropic_to_kiro
+# ==================================================================================================
+
+class TestAnthropicToKiro:
+ """Tests for anthropic_to_kiro function - main entry point."""
+
+ def test_builds_simple_payload(self):
+ """
+ What it does: Verifies building of simple Kiro payload.
+ Purpose: Ensure basic request is converted correctly.
+ """
+ print("Setup: Simple Anthropic request...")
+ request = AnthropicMessagesRequest(
+ model="claude-sonnet-4-5",
+ messages=[AnthropicMessage(role="user", content="Hello!")],
+ max_tokens=1024
+ )
+
+ print("Action: Converting to Kiro payload...")
+ with patch('kiro.converters_anthropic.get_model_id_for_kiro', return_value='claude-sonnet-4.5'):
+ with patch('kiro.converters_core.FAKE_REASONING_ENABLED', False):
+ result = anthropic_to_kiro(request, "conv-123", "arn:aws:test")
+
+ print(f"Result: {result}")
+ assert "conversationState" in result
+ assert result["conversationState"]["conversationId"] == "conv-123"
+ assert "currentMessage" in result["conversationState"]
+ assert "userInputMessage" in result["conversationState"]["currentMessage"]
+ assert result["profileArn"] == "arn:aws:test"
+
+ def test_includes_system_prompt(self):
+ """
+ What it does: Verifies that system prompt is included.
+ Purpose: Ensure Anthropic's separate system field is handled.
+ """
+ print("Setup: Request with system prompt...")
+ request = AnthropicMessagesRequest(
+ model="claude-sonnet-4-5",
+ messages=[AnthropicMessage(role="user", content="Hello!")],
+ max_tokens=1024,
+ system="You are a helpful assistant."
+ )
+
+ print("Action: Converting to Kiro payload...")
+ with patch('kiro.converters_anthropic.get_model_id_for_kiro', return_value='claude-sonnet-4.5'):
+ with patch('kiro.converters_core.FAKE_REASONING_ENABLED', False):
+ 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
+
+ def test_includes_tools(self):
+ """
+ What it does: Verifies that tools are included in payload.
+ Purpose: Ensure Anthropic tools are converted to Kiro format.
+ """
+ print("Setup: Request with tools...")
+ request = AnthropicMessagesRequest(
+ model="claude-sonnet-4-5",
+ messages=[AnthropicMessage(role="user", content="What's the weather?")],
+ max_tokens=1024,
+ tools=[
+ AnthropicTool(
+ name="get_weather",
+ description="Get weather for a location",
+ input_schema={"type": "object", "properties": {"location": {"type": "string"}}}
+ )
+ ]
+ )
+
+ print("Action: Converting to Kiro payload...")
+ with patch('kiro.converters_anthropic.get_model_id_for_kiro', return_value='claude-sonnet-4.5'):
+ with patch('kiro.converters_core.FAKE_REASONING_ENABLED', False):
+ result = anthropic_to_kiro(request, "conv-123", "arn:aws:test")
+
+ print(f"Result: {result}")
+ context = result["conversationState"]["currentMessage"]["userInputMessage"].get("userInputMessageContext", {})
+ tools = context.get("tools", [])
+ print(f"Tools in payload: {tools}")
+ assert len(tools) == 1
+ assert tools[0]["toolSpecification"]["name"] == "get_weather"
+
+ def test_builds_history_for_multi_turn(self):
+ """
+ What it does: Verifies building of history for multi-turn conversation.
+ Purpose: Ensure conversation history is included in payload.
+ """
+ print("Setup: Multi-turn conversation...")
+ request = AnthropicMessagesRequest(
+ model="claude-sonnet-4-5",
+ messages=[
+ AnthropicMessage(role="user", content="Hello"),
+ AnthropicMessage(role="assistant", content="Hi! How can I help?"),
+ AnthropicMessage(role="user", content="What's the weather?")
+ ],
+ max_tokens=1024
+ )
+
+ print("Action: Converting to Kiro payload...")
+ with patch('kiro.converters_anthropic.get_model_id_for_kiro', return_value='claude-sonnet-4.5'):
+ with patch('kiro.converters_core.FAKE_REASONING_ENABLED', False):
+ result = anthropic_to_kiro(request, "conv-123", "arn:aws:test")
+
+ 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]
+
+ def test_handles_tool_use_and_result_flow(self):
+ """
+ What it does: Verifies handling of tool use and result flow.
+ Purpose: Ensure full tool flow is converted correctly.
+ """
+ print("Setup: Tool use and result flow...")
+ request = AnthropicMessagesRequest(
+ model="claude-sonnet-4-5",
+ messages=[
+ AnthropicMessage(role="user", content="What's the weather in Moscow?"),
+ AnthropicMessage(
+ role="assistant",
+ content=[
+ {"type": "text", "text": "I'll check the weather"},
+ {"type": "tool_use", "id": "call_123", "name": "get_weather", "input": {"location": "Moscow"}}
+ ]
+ ),
+ AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "tool_result", "tool_use_id": "call_123", "content": "Weather: Sunny, 25°C"}
+ ]
+ )
+ ],
+ max_tokens=1024,
+ # Tools must be defined for tool_results to be preserved
+ tools=[
+ AnthropicTool(
+ name="get_weather",
+ description="Get weather for a location",
+ input_schema={"type": "object", "properties": {"location": {"type": "string"}}}
+ )
+ ]
+ )
+
+ print("Action: Converting to Kiro payload...")
+ with patch('kiro.converters_anthropic.get_model_id_for_kiro', return_value='claude-sonnet-4.5'):
+ with patch('kiro.converters_core.FAKE_REASONING_ENABLED', False):
+ result = anthropic_to_kiro(request, "conv-123", "arn:aws:test")
+
+ print(f"Result: {result}")
+
+ # Check history contains tool use
+ history = result["conversationState"].get("history", [])
+ print(f"History: {history}")
+
+ # Check current message contains tool result
+ current_msg = result["conversationState"]["currentMessage"]["userInputMessage"]
+ context = current_msg.get("userInputMessageContext", {})
+ tool_results = context.get("toolResults", [])
+ print(f"Tool results: {tool_results}")
+ assert len(tool_results) == 1
+
+ def test_raises_for_empty_messages(self):
+ """
+ What it does: Verifies that empty messages raise Pydantic ValidationError.
+ Purpose: Ensure Pydantic validation works correctly (min_length=1).
+
+ Note: AnthropicMessagesRequest has min_length=1 validation on messages field,
+ so empty messages are rejected at the Pydantic level, not at anthropic_to_kiro.
+ """
+ from pydantic import ValidationError
+
+ print("Setup: Attempting to create request with empty messages...")
+
+ print("Action: Creating AnthropicMessagesRequest (should raise ValidationError)...")
+ with pytest.raises(ValidationError):
+ AnthropicMessagesRequest(
+ model="claude-sonnet-4-5",
+ messages=[],
+ max_tokens=1024
+ )
+
+ print("ValidationError raised as expected - Pydantic rejects empty messages")
+
+ def test_injects_thinking_tags_when_enabled(self):
+ """
+ What it does: Verifies that thinking tags are injected when enabled.
+ Purpose: Ensure fake reasoning feature works with Anthropic API.
+ """
+ print("Setup: Request with fake reasoning enabled...")
+ request = AnthropicMessagesRequest(
+ model="claude-sonnet-4-5",
+ messages=[AnthropicMessage(role="user", content="What is 2+2?")],
+ max_tokens=1024
+ )
+
+ print("Action: Converting to Kiro payload with fake reasoning...")
+ 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_MAX_TOKENS', 4000):
+ 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 (first 200 chars): {current_content[:200]}...")
+
+ print("Checking that thinking tags are present...")
+ assert "enabled" in current_content
+ assert "What is 2+2?" in current_content
+
+ def test_injects_thinking_tags_even_when_tool_results_present(self):
+ """
+ What it does: Verifies that thinking tags ARE injected even when tool results are present.
+ Purpose: Extended thinking should work in all scenarios including tool use flows.
+ """
+ print("Setup: Request with tool results and fake reasoning enabled...")
+ request = AnthropicMessagesRequest(
+ model="claude-sonnet-4-5",
+ messages=[
+ AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "tool_result", "tool_use_id": "call_123", "content": "Result"}
+ ]
+ )
+ ],
+ max_tokens=1024,
+ # Tools must be defined for tool_results to be preserved
+ tools=[
+ AnthropicTool(
+ name="test_tool",
+ description="A test tool",
+ input_schema={"type": "object", "properties": {}}
+ )
+ ]
+ )
+
+ print("Action: Converting to Kiro payload...")
+ 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_MAX_TOKENS', 4000):
+ 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 (first 100 chars): {current_content[:100]}...")
+
+ print("Checking that thinking tags ARE present...")
+ assert "enabled" in current_content, \
+ "thinking tags SHOULD be injected even with tool results"
+
+ print("Checking that tag IS present...")
+ assert "4000" in current_content, \
+ "max_thinking_length tag SHOULD be present even with tool results"
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_converters_core.py b/kiro-gateway/tests/unit/test_converters_core.py
new file mode 100644
index 0000000000000000000000000000000000000000..92981ca6c1fcc937e06afc3455e20087e5f1345c
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_converters_core.py
@@ -0,0 +1,5248 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for converters_core module.
+
+Tests for shared conversion logic used by both OpenAI and Anthropic adapters:
+- Text content extraction
+- Message merging
+- JSON Schema sanitization
+- Tool processing
+- Thinking tag injection
+"""
+
+import pytest
+from unittest.mock import patch
+
+from kiro.converters_core import (
+ extract_text_content,
+ extract_images_from_content,
+ convert_images_to_kiro_format,
+ merge_adjacent_messages,
+ ensure_assistant_before_tool_results,
+ strip_all_tool_content,
+ build_kiro_history,
+ build_kiro_payload,
+ process_tools_with_long_descriptions,
+ inject_thinking_tags,
+ extract_tool_results_from_content,
+ extract_tool_uses_from_message,
+ sanitize_json_schema,
+ convert_tools_to_kiro_format,
+ convert_tool_results_to_kiro_format,
+ tool_calls_to_text,
+ tool_results_to_text,
+ UnifiedMessage,
+ UnifiedTool,
+)
+
+# Test data for images - 1x1 pixel JPEG
+TEST_IMAGE_BASE64 = "/9j/4AAQSkZJRgABAQEASABIAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/wAALCAABAAEBAREA/8QAFAABAAAAAAAAAAAAAAAAAAAACf/EABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEAAD8AVN//2Q=="
+
+
+# ==================================================================================================
+# Tests for extract_text_content
+# ==================================================================================================
+
+class TestExtractTextContent:
+ """Tests for extract_text_content function."""
+
+ def test_extracts_from_string(self):
+ """
+ What it does: Verifies text extraction from a string.
+ Purpose: Ensure string is returned as-is.
+ """
+ print("Setup: Simple string...")
+ content = "Hello, World!"
+
+ print("Action: Extracting text...")
+ result = extract_text_content(content)
+
+ print(f"Comparing result: Expected 'Hello, World!', Got '{result}'")
+ assert result == "Hello, World!"
+
+ def test_extracts_from_none(self):
+ """
+ What it does: Verifies None handling.
+ Purpose: Ensure None returns empty string.
+ """
+ print("Setup: None...")
+
+ print("Action: Extracting text...")
+ result = extract_text_content(None)
+
+ print(f"Comparing result: Expected '', Got '{result}'")
+ assert result == ""
+
+ def test_extracts_from_list_with_text_type(self):
+ """
+ What it does: Verifies extraction from list with type=text.
+ Purpose: Ensure multimodal format is handled.
+ """
+ print("Setup: List with type=text...")
+ content = [
+ {"type": "text", "text": "Hello"},
+ {"type": "text", "text": " World"}
+ ]
+
+ print("Action: Extracting text...")
+ result = extract_text_content(content)
+
+ print(f"Comparing result: Expected 'Hello World', Got '{result}'")
+ assert result == "Hello World"
+
+ def test_extracts_from_list_with_text_key(self):
+ """
+ What it does: Verifies extraction from list with text key.
+ Purpose: Ensure alternative format is handled.
+ """
+ print("Setup: List with text key...")
+ content = [{"text": "Hello"}, {"text": " World"}]
+
+ print("Action: Extracting text...")
+ result = extract_text_content(content)
+
+ print(f"Comparing result: Expected 'Hello World', Got '{result}'")
+ assert result == "Hello World"
+
+ def test_extracts_from_list_with_strings(self):
+ """
+ What it does: Verifies extraction from list of strings.
+ Purpose: Ensure string list is concatenated.
+ """
+ print("Setup: List of strings...")
+ content = ["Hello", " ", "World"]
+
+ print("Action: Extracting text...")
+ result = extract_text_content(content)
+
+ print(f"Comparing result: Expected 'Hello World', Got '{result}'")
+ assert result == "Hello World"
+
+ def test_extracts_from_mixed_list(self):
+ """
+ What it does: Verifies extraction from mixed list.
+ Purpose: Ensure different formats in one list are handled.
+ """
+ print("Setup: Mixed list...")
+ content = [
+ {"type": "text", "text": "Part1"},
+ "Part2",
+ {"text": "Part3"}
+ ]
+
+ print("Action: Extracting text...")
+ result = extract_text_content(content)
+
+ print(f"Comparing result: Expected 'Part1Part2Part3', Got '{result}'")
+ assert result == "Part1Part2Part3"
+
+ def test_converts_other_types_to_string(self):
+ """
+ What it does: Verifies conversion of other types to string.
+ Purpose: Ensure numbers and other types are converted.
+ """
+ print("Setup: Number...")
+ content = 42
+
+ print("Action: Extracting text...")
+ result = extract_text_content(content)
+
+ print(f"Comparing result: Expected '42', Got '{result}'")
+ assert result == "42"
+
+ def test_handles_empty_list(self):
+ """
+ What it does: Verifies empty list handling.
+ Purpose: Ensure empty list returns empty string.
+ """
+ print("Setup: Empty list...")
+ content = []
+
+ print("Action: Extracting text...")
+ result = extract_text_content(content)
+
+ print(f"Comparing result: Expected '', Got '{result}'")
+ assert result == ""
+
+ def test_extracts_from_pydantic_text_content_block(self):
+ """
+ What it does: Verifies extraction from Pydantic TextContentBlock objects.
+ Purpose: Ensure Pydantic models are handled correctly (Issue #46/#50 fix).
+
+ This is the critical test for Issue #46/#50 - the original bug was that
+ Pydantic TextContentBlock objects weren't being handled, causing MCP tool
+ results to return "(empty result)" instead of actual data.
+ """
+ from kiro.models_anthropic import TextContentBlock
+
+ print("Setup: Pydantic TextContentBlock...")
+ content = [
+ TextContentBlock(type="text", text="Hello from MCP")
+ ]
+
+ print("Action: Extracting text...")
+ result = extract_text_content(content)
+
+ print(f"Result: '{result}'")
+ print(f"Comparing result: Expected 'Hello from MCP', Got '{result}'")
+ assert result == "Hello from MCP"
+
+ def test_extracts_from_multiple_pydantic_text_blocks(self):
+ """
+ What it does: Verifies extraction from multiple Pydantic TextContentBlock objects.
+ Purpose: Ensure multiple Pydantic models are concatenated correctly.
+ """
+ from kiro.models_anthropic import TextContentBlock
+
+ print("Setup: Multiple Pydantic TextContentBlocks...")
+ content = [
+ TextContentBlock(type="text", text="Part 1"),
+ TextContentBlock(type="text", text=" Part 2"),
+ TextContentBlock(type="text", text=" Part 3")
+ ]
+
+ print("Action: Extracting text...")
+ result = extract_text_content(content)
+
+ print(f"Result: '{result}'")
+ print(f"Comparing result: Expected 'Part 1 Part 2 Part 3', Got '{result}'")
+ assert result == "Part 1 Part 2 Part 3"
+
+ def test_extracts_from_mixed_dict_and_pydantic(self):
+ """
+ What it does: Verifies extraction from mixed dict and Pydantic content.
+ Purpose: Ensure dict and Pydantic models can coexist in the same list.
+
+ This simulates real-world scenarios where some content is parsed as dict
+ and some as Pydantic models.
+ """
+ from kiro.models_anthropic import TextContentBlock
+
+ print("Setup: Mixed dict and Pydantic content...")
+ content = [
+ {"type": "text", "text": "Dict text"},
+ TextContentBlock(type="text", text=" Pydantic text"),
+ " String text"
+ ]
+
+ print("Action: Extracting text...")
+ result = extract_text_content(content)
+
+ print(f"Result: '{result}'")
+ print(f"Comparing result: Expected 'Dict text Pydantic text String text', Got '{result}'")
+ assert result == "Dict text Pydantic text String text"
+
+ def test_handles_pydantic_with_empty_text(self):
+ """
+ What it does: Verifies handling of Pydantic TextContentBlock with empty text.
+ Purpose: Ensure empty text in Pydantic models doesn't cause errors.
+ """
+ from kiro.models_anthropic import TextContentBlock
+
+ print("Setup: Pydantic TextContentBlock with empty text...")
+ content = [
+ TextContentBlock(type="text", text="")
+ ]
+
+ print("Action: Extracting text...")
+ result = extract_text_content(content)
+
+ print(f"Result: '{result}'")
+ print(f"Comparing result: Expected '', Got '{result}'")
+ assert result == ""
+
+ def test_extracts_text_ignoring_other_pydantic_types(self):
+ """
+ What it does: Verifies that only text-containing Pydantic models are extracted.
+ Purpose: Ensure non-text Pydantic models (like ToolUseContentBlock) are ignored.
+
+ This simulates MCP tool results that contain both text and tool_use blocks.
+ """
+ from kiro.models_anthropic import TextContentBlock, ToolUseContentBlock
+
+ print("Setup: Mixed Pydantic content with text and tool_use...")
+ content = [
+ TextContentBlock(type="text", text="Before tool"),
+ ToolUseContentBlock(type="tool_use", id="call_123", name="test_tool", input={}),
+ TextContentBlock(type="text", text="After tool")
+ ]
+
+ print("Action: Extracting text...")
+ result = extract_text_content(content)
+
+ print(f"Result: '{result}'")
+ print(f"Comparing result: Expected 'Before toolAfter tool', Got '{result}'")
+ assert result == "Before toolAfter tool"
+
+
+# ==================================================================================================
+# Tests for extract_images_from_content (Issue #30 fix)
+# ==================================================================================================
+
+class TestExtractImagesFromContent:
+ """
+ Tests for extract_images_from_content function.
+
+ This function extracts images from message content in unified format.
+ Supports both OpenAI (image_url with data URL) and Anthropic (image with source) formats.
+
+ This is a critical function for Issue #30 fix - 422 Validation Error for image content blocks.
+ """
+
+ def test_extracts_from_openai_format_data_url(self):
+ """
+ What it does: Verifies extraction from OpenAI image_url format with data URL.
+ Purpose: Ensure OpenAI Vision API format is handled correctly.
+
+ OpenAI format: {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}}
+ """
+ print("Setup: OpenAI format image content...")
+ content = [
+ {"type": "text", "text": "What's in this image?"},
+ {
+ "type": "image_url",
+ "image_url": {"url": f"data:image/jpeg;base64,{TEST_IMAGE_BASE64}"}
+ }
+ ]
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Result: {result}")
+ print(f"Comparing count: Expected 1, Got {len(result)}")
+ assert len(result) == 1
+
+ print("Checking media_type...")
+ assert result[0]["media_type"] == "image/jpeg"
+
+ print("Checking data...")
+ assert result[0]["data"] == TEST_IMAGE_BASE64
+
+ def test_extracts_from_anthropic_format_base64(self):
+ """
+ What it does: Verifies extraction from Anthropic image format with base64 source.
+ Purpose: Ensure Anthropic Messages API format is handled correctly.
+
+ Anthropic format: {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": "..."}}
+ """
+ print("Setup: Anthropic format image content...")
+ content = [
+ {"type": "text", "text": "Describe this image"},
+ {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "media_type": "image/png",
+ "data": TEST_IMAGE_BASE64
+ }
+ }
+ ]
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Result: {result}")
+ print(f"Comparing count: Expected 1, Got {len(result)}")
+ assert len(result) == 1
+
+ print("Checking media_type...")
+ assert result[0]["media_type"] == "image/png"
+
+ print("Checking data...")
+ assert result[0]["data"] == TEST_IMAGE_BASE64
+
+ def test_extracts_from_mixed_content(self):
+ """
+ What it does: Verifies extraction from mixed content (text + multiple images).
+ Purpose: Ensure all images are extracted from multimodal content.
+ """
+ print("Setup: Mixed content with multiple images...")
+ content = [
+ {"type": "text", "text": "Compare these images:"},
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/jpeg", "data": "image1_data"}
+ },
+ {"type": "text", "text": "and"},
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/png", "data": "image2_data"}
+ }
+ ]
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Result: {result}")
+ print(f"Comparing count: Expected 2, Got {len(result)}")
+ assert len(result) == 2
+
+ print("Checking first image...")
+ assert result[0]["media_type"] == "image/jpeg"
+ assert result[0]["data"] == "image1_data"
+
+ print("Checking second image...")
+ assert result[1]["media_type"] == "image/png"
+ assert result[1]["data"] == "image2_data"
+
+ def test_returns_empty_for_string_content(self):
+ """
+ What it does: Verifies empty list return for string content.
+ Purpose: Ensure string content doesn't contain images.
+ """
+ print("Setup: String content...")
+ content = "Just a text message"
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_returns_empty_for_empty_content(self):
+ """
+ What it does: Verifies empty list return for empty content.
+ Purpose: Ensure empty list returns empty list.
+ """
+ print("Setup: Empty list...")
+ content = []
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_returns_empty_for_none_content(self):
+ """
+ What it does: Verifies empty list return for None content.
+ Purpose: Ensure None doesn't cause errors.
+ """
+ print("Setup: None content...")
+ content = None
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_returns_empty_for_text_only_content(self):
+ """
+ What it does: Verifies empty list return for text-only content.
+ Purpose: Ensure text blocks don't produce images.
+ """
+ print("Setup: Text-only content...")
+ content = [
+ {"type": "text", "text": "Hello"},
+ {"type": "text", "text": "World"}
+ ]
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_handles_url_images_with_warning(self):
+ """
+ What it does: Verifies URL-based images are skipped with warning.
+ Purpose: Ensure URL images don't crash but are logged as unsupported.
+
+ URL-based images require fetching and are not supported by Kiro API directly.
+ """
+ print("Setup: URL-based image content...")
+ content = [
+ {
+ "type": "image_url",
+ "image_url": {"url": "https://example.com/image.jpg"}
+ }
+ ]
+
+ print("Action: Extracting images (should skip URL images)...")
+ result = extract_images_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == [] # URL images are skipped
+
+ def test_handles_anthropic_url_source_with_warning(self):
+ """
+ What it does: Verifies Anthropic URL source images are skipped with warning.
+ Purpose: Ensure Anthropic URL format doesn't crash but is logged as unsupported.
+ """
+ print("Setup: Anthropic URL source image...")
+ content = [
+ {
+ "type": "image",
+ "source": {
+ "type": "url",
+ "url": "https://example.com/image.png"
+ }
+ }
+ ]
+
+ print("Action: Extracting images (should skip URL images)...")
+ result = extract_images_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == [] # URL images are skipped
+
+ def test_handles_invalid_data_url(self):
+ """
+ What it does: Verifies handling of invalid data URL format.
+ Purpose: Ensure malformed data URLs don't crash the function.
+ """
+ print("Setup: Invalid data URL...")
+ content = [
+ {
+ "type": "image_url",
+ "image_url": {"url": "data:invalid_format_without_comma"}
+ }
+ ]
+
+ print("Action: Extracting images (should handle gracefully)...")
+ result = extract_images_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == [] # Invalid data URL is skipped
+
+ def test_handles_empty_data_in_image(self):
+ """
+ What it does: Verifies handling of image with empty data.
+ Purpose: Ensure images with empty data are skipped.
+ """
+ print("Setup: Image with empty data...")
+ content = [
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/jpeg", "data": ""}
+ }
+ ]
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == [] # Empty data is skipped
+
+ def test_extracts_from_pydantic_image_content_block(self):
+ """
+ What it does: Verifies extraction from Pydantic ImageContentBlock objects.
+ Purpose: Ensure Pydantic models are handled correctly (Issue #30 fix).
+
+ This is the critical test for Issue #30 - the original bug was that
+ Pydantic ImageContentBlock objects weren't being handled.
+ """
+ from kiro.models_anthropic import ImageContentBlock, Base64ImageSource
+
+ print("Setup: Pydantic ImageContentBlock...")
+ content = [
+ ImageContentBlock(
+ type="image",
+ source=Base64ImageSource(
+ type="base64",
+ media_type="image/webp",
+ data=TEST_IMAGE_BASE64
+ )
+ )
+ ]
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Result: {result}")
+ print(f"Comparing count: Expected 1, Got {len(result)}")
+ assert len(result) == 1
+
+ print("Checking media_type...")
+ assert result[0]["media_type"] == "image/webp"
+
+ print("Checking data...")
+ assert result[0]["data"] == TEST_IMAGE_BASE64
+
+ def test_extracts_from_pydantic_url_image_source(self):
+ """
+ What it does: Verifies handling of Pydantic URLImageSource objects.
+ Purpose: Ensure Pydantic URL sources are skipped with warning.
+ """
+ from kiro.models_anthropic import ImageContentBlock, URLImageSource
+
+ print("Setup: Pydantic ImageContentBlock with URL source...")
+ content = [
+ ImageContentBlock(
+ type="image",
+ source=URLImageSource(
+ type="url",
+ url="https://example.com/image.gif"
+ )
+ )
+ ]
+
+ print("Action: Extracting images (should skip URL images)...")
+ result = extract_images_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == [] # URL images are skipped
+
+ def test_extracts_multiple_formats_mixed(self):
+ """
+ What it does: Verifies extraction from mixed OpenAI and Anthropic formats.
+ Purpose: Ensure both formats can coexist in the same content list.
+ """
+ print("Setup: Mixed OpenAI and Anthropic formats...")
+ content = [
+ # OpenAI format
+ {
+ "type": "image_url",
+ "image_url": {"url": f"data:image/jpeg;base64,openai_image_data"}
+ },
+ # Anthropic format
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/png", "data": "anthropic_image_data"}
+ }
+ ]
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Result: {result}")
+ print(f"Comparing count: Expected 2, Got {len(result)}")
+ assert len(result) == 2
+
+ print("Checking OpenAI image...")
+ assert result[0]["media_type"] == "image/jpeg"
+ assert result[0]["data"] == "openai_image_data"
+
+ print("Checking Anthropic image...")
+ assert result[1]["media_type"] == "image/png"
+ assert result[1]["data"] == "anthropic_image_data"
+
+ def test_handles_missing_source_in_anthropic_format(self):
+ """
+ What it does: Verifies handling of Anthropic image without source.
+ Purpose: Ensure malformed Anthropic images don't crash.
+ """
+ print("Setup: Anthropic image without source...")
+ content = [
+ {"type": "image"} # Missing source
+ ]
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_handles_missing_image_url_in_openai_format(self):
+ """
+ What it does: Verifies handling of OpenAI image_url without image_url field.
+ Purpose: Ensure malformed OpenAI images don't crash.
+ """
+ print("Setup: OpenAI image_url without image_url field...")
+ content = [
+ {"type": "image_url"} # Missing image_url
+ ]
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_extracts_gif_format(self):
+ """
+ What it does: Verifies extraction of GIF images.
+ Purpose: Ensure GIF format is supported.
+ """
+ print("Setup: GIF image...")
+ content = [
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/gif", "data": "gif_data"}
+ }
+ ]
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["media_type"] == "image/gif"
+
+ def test_extracts_webp_format(self):
+ """
+ What it does: Verifies extraction of WebP images.
+ Purpose: Ensure WebP format is supported.
+ """
+ print("Setup: WebP image...")
+ content = [
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/webp", "data": "webp_data"}
+ }
+ ]
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["media_type"] == "image/webp"
+
+ def test_uses_default_media_type_when_missing(self):
+ """
+ What it does: Verifies default media_type is used when not specified.
+ Purpose: Ensure missing media_type defaults to image/jpeg.
+ """
+ print("Setup: Image without media_type...")
+ content = [
+ {
+ "type": "image",
+ "source": {"type": "base64", "data": "some_data"} # No media_type
+ }
+ ]
+
+ print("Action: Extracting images...")
+ result = extract_images_from_content(content)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["media_type"] == "image/jpeg" # Default
+
+
+# ==================================================================================================
+# Tests for convert_images_to_kiro_format
+# ==================================================================================================
+
+class TestConvertImagesToKiroFormat:
+ """
+ Tests for convert_images_to_kiro_format function.
+
+ This function converts unified images to Kiro API format.
+
+ Unified format: [{"media_type": "image/jpeg", "data": "base64..."}]
+ Kiro format: [{"format": "jpeg", "source": {"bytes": "base64..."}}]
+ """
+
+ def test_converts_single_image(self):
+ """
+ What it does: Verifies conversion of a single image.
+ Purpose: Ensure basic conversion from unified to Kiro format works.
+ """
+ print("Setup: Single image in unified format...")
+ images = [{"media_type": "image/jpeg", "data": TEST_IMAGE_BASE64}]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result: {result}")
+ print(f"Comparing count: Expected 1, Got {len(result)}")
+ assert len(result) == 1
+
+ print("Checking format...")
+ assert result[0]["format"] == "jpeg"
+
+ print("Checking source.bytes...")
+ assert result[0]["source"]["bytes"] == TEST_IMAGE_BASE64
+
+ def test_converts_multiple_images(self):
+ """
+ What it does: Verifies conversion of multiple images.
+ Purpose: Ensure all images are converted correctly.
+ """
+ print("Setup: Multiple images...")
+ images = [
+ {"media_type": "image/jpeg", "data": "jpeg_data"},
+ {"media_type": "image/png", "data": "png_data"},
+ {"media_type": "image/gif", "data": "gif_data"}
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result: {result}")
+ print(f"Comparing count: Expected 3, Got {len(result)}")
+ assert len(result) == 3
+
+ print("Checking formats...")
+ assert result[0]["format"] == "jpeg"
+ assert result[1]["format"] == "png"
+ assert result[2]["format"] == "gif"
+
+ def test_returns_empty_for_none(self):
+ """
+ What it does: Verifies handling of None.
+ Purpose: Ensure None returns empty list.
+ """
+ print("Setup: None images...")
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(None)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_returns_empty_for_empty_list(self):
+ """
+ What it does: Verifies handling of empty list.
+ Purpose: Ensure empty list returns empty list.
+ """
+ print("Setup: Empty images list...")
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format([])
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_skips_images_with_empty_data(self):
+ """
+ What it does: Verifies skipping of images with empty data.
+ Purpose: Ensure images without data are not included.
+ """
+ print("Setup: Image with empty data...")
+ images = [
+ {"media_type": "image/jpeg", "data": ""},
+ {"media_type": "image/png", "data": "valid_data"}
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result: {result}")
+ print(f"Comparing count: Expected 1, Got {len(result)}")
+ assert len(result) == 1
+ assert result[0]["format"] == "png"
+
+ def test_extracts_format_from_media_type(self):
+ """
+ What it does: Verifies extraction of format from media_type.
+ Purpose: Ensure "image/jpeg" becomes "jpeg".
+ """
+ print("Setup: Various media types...")
+ images = [
+ {"media_type": "image/jpeg", "data": "data1"},
+ {"media_type": "image/png", "data": "data2"},
+ {"media_type": "image/gif", "data": "data3"},
+ {"media_type": "image/webp", "data": "data4"}
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result formats: {[r['format'] for r in result]}")
+ assert result[0]["format"] == "jpeg"
+ assert result[1]["format"] == "png"
+ assert result[2]["format"] == "gif"
+ assert result[3]["format"] == "webp"
+
+ def test_handles_media_type_without_slash(self):
+ """
+ What it does: Verifies handling of media_type without slash.
+ Purpose: Ensure edge case media_type is handled.
+ """
+ print("Setup: Media type without slash...")
+ images = [{"media_type": "jpeg", "data": "data"}]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["format"] == "jpeg"
+
+ def test_uses_default_media_type_when_missing(self):
+ """
+ What it does: Verifies default media_type is used when not specified.
+ Purpose: Ensure missing media_type defaults to image/jpeg.
+ """
+ print("Setup: Image without media_type...")
+ images = [{"data": "some_data"}] # No media_type
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["format"] == "jpeg" # Default from "image/jpeg"
+
+ def test_preserves_large_image_data(self):
+ """
+ What it does: Verifies large image data is preserved.
+ Purpose: Ensure large images are not truncated.
+ """
+ print("Setup: Large image data...")
+ large_data = "A" * 100000 # 100KB of data
+ images = [{"media_type": "image/png", "data": large_data}]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result data length: {len(result[0]['source']['bytes'])}")
+ assert len(result[0]["source"]["bytes"]) == 100000
+
+ # ==================================================================================
+ # Data URL Prefix Stripping Tests (Issue #32 fix)
+ # ==================================================================================
+
+ def test_strips_data_url_prefix_jpeg(self):
+ """
+ What it does: Verifies that data URL prefix is stripped from JPEG image data.
+ Purpose: Ensure Kiro API receives pure base64 without the data URL prefix (Issue #32 fix).
+
+ Some clients send the full data URL in the data field instead of pure base64.
+ Kiro API expects pure base64 without the "data:image/jpeg;base64," prefix.
+ """
+ print("Setup: Image with data URL prefix (JPEG)...")
+ pure_base64 = "/9j/4AAQSkZJRg==" # Sample JPEG base64
+ images = [{"media_type": "image/jpeg", "data": f"data:image/jpeg;base64,{pure_base64}"}]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result: {result}")
+ print(f"Comparing bytes: Expected '{pure_base64}', Got '{result[0]['source']['bytes']}'")
+ assert result[0]["source"]["bytes"] == pure_base64
+ assert result[0]["format"] == "jpeg"
+
+ def test_strips_data_url_prefix_png(self):
+ """
+ What it does: Verifies that data URL prefix is stripped from PNG image data.
+ Purpose: Ensure PNG images with data URL prefix are handled correctly (Issue #32 fix).
+ """
+ print("Setup: Image with data URL prefix (PNG)...")
+ pure_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
+ images = [{"media_type": "image/png", "data": f"data:image/png;base64,{pure_base64}"}]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result: {result}")
+ print(f"Comparing bytes: Expected pure base64, Got '{result[0]['source']['bytes'][:50]}...'")
+ assert result[0]["source"]["bytes"] == pure_base64
+ assert result[0]["format"] == "png"
+
+ def test_extracts_media_type_from_data_url(self):
+ """
+ What it does: Verifies that media_type is extracted from data URL header.
+ Purpose: Ensure media_type from data URL overrides the original media_type (Issue #32 fix).
+
+ When data URL contains media type info, it should be used instead of the
+ original media_type field (which might be incorrect or generic).
+ """
+ print("Setup: Image with mismatched media_type and data URL...")
+ pure_base64 = "R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7" # GIF
+ # Original media_type says jpeg, but data URL says gif
+ images = [{"media_type": "image/jpeg", "data": f"data:image/gif;base64,{pure_base64}"}]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result: {result}")
+ print("Checking that media_type from data URL is used...")
+ assert result[0]["format"] == "gif" # Should use gif from data URL, not jpeg
+ assert result[0]["source"]["bytes"] == pure_base64
+
+ def test_handles_malformed_data_url_no_comma(self):
+ """
+ What it does: Verifies graceful handling of malformed data URL without comma.
+ Purpose: Ensure function doesn't crash on malformed data URLs (Issue #32 fix).
+
+ If data URL is malformed (no comma separator), the function should
+ log a warning and use the original data as-is.
+ """
+ print("Setup: Malformed data URL without comma...")
+ malformed_data = "data:image/jpeg;base64_without_comma"
+ images = [{"media_type": "image/jpeg", "data": malformed_data}]
+
+ print("Action: Converting to Kiro format (should handle gracefully)...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result: {result}")
+ # The function should still produce output, using the malformed data as-is
+ # (since split(",", 1) will fail and the except block will catch it)
+ assert len(result) == 1
+ # After the fix, malformed data URL should be preserved as-is
+ assert result[0]["source"]["bytes"] == malformed_data
+
+ def test_preserves_pure_base64_data(self):
+ """
+ What it does: Verifies that pure base64 data (without prefix) is preserved.
+ Purpose: Ensure normal base64 data is not modified (Issue #32 fix).
+
+ When data is already pure base64 (doesn't start with "data:"),
+ it should be passed through unchanged.
+ """
+ print("Setup: Pure base64 data without prefix...")
+ pure_base64 = TEST_IMAGE_BASE64
+ images = [{"media_type": "image/jpeg", "data": pure_base64}]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result: {result}")
+ print("Checking that pure base64 is preserved unchanged...")
+ assert result[0]["source"]["bytes"] == pure_base64
+ assert result[0]["format"] == "jpeg"
+
+ def test_strips_data_url_prefix_webp(self):
+ """
+ What it does: Verifies that data URL prefix is stripped from WebP image data.
+ Purpose: Ensure WebP images with data URL prefix are handled correctly (Issue #32 fix).
+ """
+ print("Setup: Image with data URL prefix (WebP)...")
+ pure_base64 = "UklGRh4AAABXRUJQVlA4TBEAAAAvAAAAAAfQ//73v/+BiOh/AAA="
+ images = [{"media_type": "image/webp", "data": f"data:image/webp;base64,{pure_base64}"}]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result: {result}")
+ assert result[0]["source"]["bytes"] == pure_base64
+ assert result[0]["format"] == "webp"
+
+ def test_handles_data_url_with_empty_base64(self):
+ """
+ What it does: Verifies handling of data URL with empty base64 part.
+ Purpose: Ensure empty data after prefix is handled correctly (Issue #32 fix).
+
+ Note: The function strips the prefix but doesn't re-check for empty data after stripping.
+ This means an image with "data:image/jpeg;base64," will result in empty bytes.
+ This is acceptable behavior as Kiro API will handle the validation.
+ """
+ print("Setup: Data URL with empty base64 part...")
+ images = [{"media_type": "image/jpeg", "data": "data:image/jpeg;base64,"}]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_images_to_kiro_format(images)
+
+ print(f"Result: {result}")
+ print("Checking that image is converted (with empty bytes)...")
+ # The function strips the prefix but doesn't re-check for empty data
+ # This results in an image with empty bytes
+ assert len(result) == 1
+ assert result[0]["source"]["bytes"] == ""
+ assert result[0]["format"] == "jpeg"
+
+
+# ==================================================================================================
+# Tests for merge_adjacent_messages
+# ==================================================================================================
+
+class TestMergeAdjacentMessages:
+ """Tests for merge_adjacent_messages function using UnifiedMessage."""
+
+ def test_merges_adjacent_user_messages(self):
+ """
+ What it does: Verifies merging of adjacent user messages.
+ Purpose: Ensure messages with the same role are merged.
+ """
+ print("Setup: Two consecutive user messages...")
+ messages = [
+ UnifiedMessage(role="user", content="Hello"),
+ UnifiedMessage(role="user", content="World")
+ ]
+
+ print("Action: Merging messages...")
+ result = merge_adjacent_messages(messages)
+
+ print(f"Comparing length: Expected 1, Got {len(result)}")
+ assert len(result) == 1
+ assert "Hello" in result[0].content
+ assert "World" in result[0].content
+
+ def test_preserves_alternating_messages(self):
+ """
+ What it does: Verifies preservation of alternating messages.
+ Purpose: Ensure different roles are not merged.
+ """
+ print("Setup: Alternating messages...")
+ messages = [
+ UnifiedMessage(role="user", content="Hello"),
+ UnifiedMessage(role="assistant", content="Hi"),
+ UnifiedMessage(role="user", content="How are you?")
+ ]
+
+ print("Action: Merging messages...")
+ result = merge_adjacent_messages(messages)
+
+ print(f"Comparing length: Expected 3, Got {len(result)}")
+ assert len(result) == 3
+
+ def test_handles_empty_list(self):
+ """
+ What it does: Verifies empty list handling.
+ Purpose: Ensure empty list doesn't cause errors.
+ """
+ print("Setup: Empty list...")
+
+ print("Action: Merging messages...")
+ result = merge_adjacent_messages([])
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_handles_single_message(self):
+ """
+ What it does: Verifies single message handling.
+ Purpose: Ensure single message is returned as-is.
+ """
+ print("Setup: Single message...")
+ messages = [UnifiedMessage(role="user", content="Hello")]
+
+ print("Action: Merging messages...")
+ result = merge_adjacent_messages(messages)
+
+ print(f"Comparing length: Expected 1, Got {len(result)}")
+ assert len(result) == 1
+ assert result[0].content == "Hello"
+
+ def test_merges_multiple_adjacent_groups(self):
+ """
+ What it does: Verifies merging of multiple groups.
+ Purpose: Ensure multiple groups of adjacent messages are merged.
+ """
+ print("Setup: Multiple groups of adjacent messages...")
+ messages = [
+ UnifiedMessage(role="user", content="A"),
+ UnifiedMessage(role="user", content="B"),
+ UnifiedMessage(role="assistant", content="C"),
+ UnifiedMessage(role="assistant", content="D"),
+ UnifiedMessage(role="user", content="E")
+ ]
+
+ print("Action: Merging messages...")
+ result = merge_adjacent_messages(messages)
+
+ print(f"Comparing length: Expected 3, Got {len(result)}")
+ assert len(result) == 3
+ assert result[0].role == "user"
+ assert result[1].role == "assistant"
+ assert result[2].role == "user"
+
+ def test_merges_list_contents_correctly(self):
+ """
+ What it does: Verifies merging of list contents.
+ Purpose: Ensure lists are merged correctly.
+ """
+ print("Setup: Two user messages with list content...")
+ messages = [
+ UnifiedMessage(role="user", content=[{"type": "text", "text": "Part 1"}]),
+ UnifiedMessage(role="user", content=[{"type": "text", "text": "Part 2"}])
+ ]
+
+ print("Action: Merging messages...")
+ result = merge_adjacent_messages(messages)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert isinstance(result[0].content, list)
+ assert len(result[0].content) == 2
+
+ def test_merges_adjacent_assistant_tool_calls(self):
+ """
+ What it does: Verifies merging of tool_calls when merging adjacent assistant messages.
+ Purpose: Ensure tool_calls from all assistant messages are preserved when merging.
+
+ This is a critical test for a bug where multiple assistant messages with tool_calls
+ were sent in a row, and the second tool_call was lost.
+ """
+ print("Setup: Two assistant messages with different tool_calls...")
+ messages = [
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{
+ "id": "tooluse_first",
+ "type": "function",
+ "function": {"name": "shell", "arguments": '{"command": ["ls"]}'}
+ }]
+ ),
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{
+ "id": "tooluse_second",
+ "type": "function",
+ "function": {"name": "shell", "arguments": '{"command": ["pwd"]}'}
+ }]
+ )
+ ]
+
+ print("Action: Merging messages...")
+ result = merge_adjacent_messages(messages)
+
+ print(f"Result: {result}")
+ print(f"Comparing length: Expected 1, Got {len(result)}")
+ assert len(result) == 1
+ assert result[0].role == "assistant"
+
+ print("Checking that both tool_calls are preserved...")
+ assert result[0].tool_calls is not None
+ print(f"Comparing tool_calls count: Expected 2, Got {len(result[0].tool_calls)}")
+ assert len(result[0].tool_calls) == 2
+
+ tool_ids = [tc["id"] for tc in result[0].tool_calls]
+ print(f"Tool IDs: {tool_ids}")
+ assert "tooluse_first" in tool_ids
+ assert "tooluse_second" in tool_ids
+
+ def test_merges_three_adjacent_assistant_tool_calls(self):
+ """
+ What it does: Verifies merging of tool_calls from three assistant messages.
+ Purpose: Ensure all tool_calls are preserved when merging more than two messages.
+ """
+ print("Setup: Three assistant messages with tool_calls...")
+ messages = [
+ UnifiedMessage(role="assistant", content="", tool_calls=[
+ {"id": "call_1", "type": "function", "function": {"name": "tool1", "arguments": "{}"}}
+ ]),
+ UnifiedMessage(role="assistant", content="", tool_calls=[
+ {"id": "call_2", "type": "function", "function": {"name": "tool2", "arguments": "{}"}}
+ ]),
+ UnifiedMessage(role="assistant", content="", tool_calls=[
+ {"id": "call_3", "type": "function", "function": {"name": "tool3", "arguments": "{}"}}
+ ])
+ ]
+
+ print("Action: Merging messages...")
+ result = merge_adjacent_messages(messages)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert len(result[0].tool_calls) == 3
+
+ tool_ids = [tc["id"] for tc in result[0].tool_calls]
+ print(f"Comparing tool IDs: Expected ['call_1', 'call_2', 'call_3'], Got {tool_ids}")
+ assert tool_ids == ["call_1", "call_2", "call_3"]
+
+ def test_merges_assistant_with_and_without_tool_calls(self):
+ """
+ What it does: Verifies merging of assistant with and without tool_calls.
+ Purpose: Ensure tool_calls are correctly initialized when merging.
+ """
+ print("Setup: Assistant without tool_calls + assistant with tool_calls...")
+ messages = [
+ UnifiedMessage(role="assistant", content="Thinking...", tool_calls=None),
+ UnifiedMessage(role="assistant", content="", tool_calls=[
+ {"id": "call_1", "type": "function", "function": {"name": "tool1", "arguments": "{}"}}
+ ])
+ ]
+
+ print("Action: Merging messages...")
+ result = merge_adjacent_messages(messages)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0].tool_calls is not None
+ print(f"Comparing tool_calls count: Expected 1, Got {len(result[0].tool_calls)}")
+ assert len(result[0].tool_calls) == 1
+ assert result[0].tool_calls[0]["id"] == "call_1"
+
+ def test_merges_user_messages_with_tool_results(self):
+ """
+ What it does: Verifies merging of user messages with tool_results.
+ Purpose: Ensure tool_results are preserved when merging user messages.
+ """
+ print("Setup: Two user messages with tool_results...")
+ messages = [
+ UnifiedMessage(role="user", content="", tool_results=[
+ {"type": "tool_result", "tool_use_id": "call_1", "content": "Result 1"}
+ ]),
+ UnifiedMessage(role="user", content="", tool_results=[
+ {"type": "tool_result", "tool_use_id": "call_2", "content": "Result 2"}
+ ])
+ ]
+
+ print("Action: Merging messages...")
+ result = merge_adjacent_messages(messages)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0].tool_results is not None
+ assert len(result[0].tool_results) == 2
+
+
+# ==================================================================================================
+# Tests for ensure_assistant_before_tool_results
+# ==================================================================================================
+
+class TestEnsureAssistantBeforeToolResults:
+ """
+ Tests for ensure_assistant_before_tool_results function.
+
+ This function handles the case when clients (like Cline/Roo/Cursor) send truncated
+ conversations with tool_results but without the preceding assistant message
+ that contains the tool_calls. Since we don't know the original tool name,
+ we strip the orphaned tool_results to avoid Kiro API rejection.
+ """
+
+ def test_returns_empty_list_for_empty_input(self):
+ """
+ What it does: Verifies empty list handling.
+ Purpose: Ensure empty input returns empty output.
+ """
+ print("Setup: Empty list...")
+
+ print("Action: Processing messages...")
+ result, stripped = ensure_assistant_before_tool_results([])
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+ assert stripped is False
+
+ def test_preserves_messages_without_tool_results(self):
+ """
+ What it does: Verifies messages without tool_results are unchanged.
+ Purpose: Ensure regular messages pass through unmodified.
+ """
+ print("Setup: Messages without tool_results...")
+ messages = [
+ UnifiedMessage(role="user", content="Hello"),
+ UnifiedMessage(role="assistant", content="Hi there"),
+ UnifiedMessage(role="user", content="How are you?")
+ ]
+
+ print("Action: Processing messages...")
+ result, stripped = ensure_assistant_before_tool_results(messages)
+
+ print(f"Comparing length: Expected 3, Got {len(result)}")
+ assert len(result) == 3
+ assert result[0].content == "Hello"
+ assert result[1].content == "Hi there"
+ assert result[2].content == "How are you?"
+ assert stripped is False
+
+ def test_preserves_tool_results_with_preceding_assistant(self):
+ """
+ What it does: Verifies tool_results are preserved when assistant with tool_calls precedes.
+ Purpose: Ensure valid tool_results are not stripped.
+ """
+ print("Setup: Valid conversation with assistant tool_calls followed by user tool_results...")
+ messages = [
+ UnifiedMessage(role="user", content="Call a tool"),
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "get_weather", "arguments": '{"location": "Moscow"}'}
+ }]
+ ),
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Weather is sunny"
+ }]
+ )
+ ]
+
+ print("Action: Processing messages...")
+ result, stripped = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result: {result}")
+ print(f"Comparing length: Expected 3, Got {len(result)}")
+ assert len(result) == 3
+
+ print("Checking that tool_results are preserved...")
+ assert result[2].tool_results is not None
+ assert len(result[2].tool_results) == 1
+ assert result[2].tool_results[0]["tool_use_id"] == "call_123"
+ assert stripped is False
+
+ def test_strips_orphaned_tool_results_at_start(self):
+ """
+ What it does: Verifies orphaned tool_results at the start are converted to text.
+ Purpose: Ensure tool_results without preceding assistant are converted to text representation.
+
+ This is the critical bug fix test - when a client sends a truncated
+ conversation starting with tool_results, they should be converted to text.
+ """
+ print("Setup: Conversation starting with orphaned tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_orphan",
+ "content": "Orphaned result"
+ }]
+ ),
+ UnifiedMessage(role="user", content="Continue the conversation")
+ ]
+
+ print("Action: Processing messages...")
+ result, converted = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result: {result}")
+ print(f"Comparing length: Expected 2, Got {len(result)}")
+ assert len(result) == 2
+
+ print("Checking that orphaned tool_results are converted to text...")
+ assert result[0].tool_results is None
+
+ print("Checking that content now contains the tool result as text...")
+ print(f"Content: '{result[0].content}'")
+ assert "[Tool Result (call_orphan)]" in result[0].content
+ assert "Orphaned result" in result[0].content
+
+ assert result[1].content == "Continue the conversation"
+ assert converted is True
+
+ def test_converts_tool_results_after_assistant_without_tool_calls(self):
+ """
+ What it does: Verifies tool_results are converted when preceding assistant has no tool_calls.
+ Purpose: Ensure tool_results require assistant with tool_calls, not just any assistant.
+ """
+ print("Setup: Assistant without tool_calls followed by user with tool_results...")
+ messages = [
+ UnifiedMessage(role="user", content="Hello"),
+ UnifiedMessage(role="assistant", content="Let me think...", tool_calls=None),
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Result"
+ }]
+ )
+ ]
+
+ print("Action: Processing messages...")
+ result, converted = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result: {result}")
+ print("Checking that tool_results are converted to text...")
+ assert result[2].tool_results is None
+
+ print(f"Content after conversion: '{result[2].content}'")
+ assert "[Tool Result (call_123)]" in result[2].content
+ assert "Result" in result[2].content
+
+ assert converted is True
+
+ def test_converts_tool_results_after_user_message(self):
+ """
+ What it does: Verifies tool_results are converted when preceded by user message.
+ Purpose: Ensure tool_results require assistant, not user.
+ """
+ print("Setup: User message followed by user with tool_results...")
+ messages = [
+ UnifiedMessage(role="user", content="First message"),
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Result"
+ }]
+ )
+ ]
+
+ print("Action: Processing messages...")
+ result, converted = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result: {result}")
+ print("Checking that tool_results are converted to text...")
+ assert result[1].tool_results is None
+
+ print(f"Content after conversion: '{result[1].content}'")
+ assert "[Tool Result (call_123)]" in result[1].content
+ assert "Result" in result[1].content
+
+ assert converted is True
+
+ def test_preserves_content_when_converting_tool_results(self):
+ """
+ What it does: Verifies message content is preserved and tool_results are appended as text.
+ Purpose: Ensure original content is kept and tool_results are converted to text representation.
+ """
+ print("Setup: Message with both content and orphaned tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="Here is some context",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Result"
+ }]
+ )
+ ]
+
+ print("Action: Processing messages...")
+ result, converted = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result: {result}")
+ print(f"Content after conversion: '{result[0].content}'")
+
+ print("Checking that original content is preserved...")
+ assert "Here is some context" in result[0].content
+
+ print("Checking that tool_results are converted to text and appended...")
+ assert "[Tool Result (call_123)]" in result[0].content
+ assert "Result" in result[0].content
+
+ print("Checking that tool_results field is removed...")
+ assert result[0].tool_results is None
+
+ assert converted is True
+
+ def test_preserves_tool_calls_when_converting_tool_results(self):
+ """
+ What it does: Verifies tool_calls are preserved when tool_results are converted.
+ Purpose: Ensure only tool_results are converted, tool_calls stay.
+ """
+ print("Setup: Message with tool_calls and orphaned tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{
+ "id": "call_new",
+ "type": "function",
+ "function": {"name": "new_tool", "arguments": "{}"}
+ }],
+ tool_results=[{ # This shouldn't happen but let's test it
+ "type": "tool_result",
+ "tool_use_id": "call_old",
+ "content": "Old result"
+ }]
+ )
+ ]
+
+ print("Action: Processing messages...")
+ result, converted = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result: {result}")
+ print("Checking that tool_calls are preserved...")
+ assert result[0].tool_calls is not None
+ assert len(result[0].tool_calls) == 1
+
+ print("Checking that tool_results are converted to text...")
+ assert result[0].tool_results is None
+ assert "[Tool Result (call_old)]" in result[0].content
+ assert "Old result" in result[0].content
+
+ assert converted is True
+
+ def test_handles_multiple_orphaned_tool_results(self):
+ """
+ What it does: Verifies multiple orphaned tool_results are all converted.
+ Purpose: Ensure all tool_results in the list are converted to text.
+ """
+ print("Setup: Message with multiple orphaned tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[
+ {"type": "tool_result", "tool_use_id": "call_1", "content": "Result 1"},
+ {"type": "tool_result", "tool_use_id": "call_2", "content": "Result 2"},
+ {"type": "tool_result", "tool_use_id": "call_3", "content": "Result 3"}
+ ]
+ )
+ ]
+
+ print("Action: Processing messages...")
+ result, converted = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result: {result}")
+ print(f"Content after conversion: '{result[0].content}'")
+
+ print("Checking that all tool_results are converted to text...")
+ assert result[0].tool_results is None
+ assert "[Tool Result (call_1)]" in result[0].content
+ assert "Result 1" in result[0].content
+ assert "[Tool Result (call_2)]" in result[0].content
+ assert "Result 2" in result[0].content
+ assert "[Tool Result (call_3)]" in result[0].content
+ assert "Result 3" in result[0].content
+
+ assert converted is True
+
+ # ==================================================================================
+ # New tests for tool_results conversion (PR #49)
+ # ==================================================================================
+
+ def test_conversion_preserves_images(self):
+ """
+ What it does: Verifies that images field is preserved when converting tool_results.
+ Purpose: Ensure images=msg.images is set correctly in converted message.
+ """
+ print("Setup: Message with images and orphaned tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="Here's an image and tool result",
+ images=[{"media_type": "image/jpeg", "data": "image_data"}],
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Tool output"
+ }]
+ )
+ ]
+
+ print("Action: Processing messages...")
+ result, converted = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result: {result}")
+ print("Checking that images are preserved...")
+ assert result[0].images is not None
+ assert len(result[0].images) == 1
+ assert result[0].images[0]["media_type"] == "image/jpeg"
+
+ print("Checking that tool_results are converted...")
+ assert result[0].tool_results is None
+ assert "[Tool Result" in result[0].content
+
+ assert converted is True
+
+ def test_conversion_appends_to_existing_content(self):
+ """
+ What it does: Verifies tool_results are appended with double newline.
+ Purpose: Ensure formatting: "original\\n\\n[Tool Result]\\ndata".
+ """
+ print("Setup: Message with content and orphaned tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="Original content here",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_abc",
+ "content": "Tool data"
+ }]
+ )
+ ]
+
+ print("Action: Processing messages...")
+ result, converted = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result content: '{result[0].content}'")
+
+ print("Checking formatting...")
+ assert "Original content here" in result[0].content
+ assert "[Tool Result (call_abc)]" in result[0].content
+ assert "Tool data" in result[0].content
+
+ # Check double newline separator
+ assert "\n\n" in result[0].content
+
+ assert converted is True
+
+ def test_conversion_handles_empty_original_content(self):
+ """
+ What it does: Verifies conversion works when original content is empty.
+ Purpose: Ensure that only tool_results text is used when content is empty.
+ """
+ print("Setup: Message with empty content and orphaned tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_xyz",
+ "content": "Only tool result"
+ }]
+ )
+ ]
+
+ print("Action: Processing messages...")
+ result, converted = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result content: '{result[0].content}'")
+
+ print("Checking that only tool result text is present...")
+ assert "[Tool Result (call_xyz)]" in result[0].content
+ assert "Only tool result" in result[0].content
+
+ # Should not have leading/trailing whitespace from empty original content
+ assert result[0].content.strip() == result[0].content
+
+ assert converted is True
+
+ def test_conversion_returns_correct_flag(self):
+ """
+ What it does: Verifies that converted_any_tool_results flag is returned correctly.
+ Purpose: Ensure return value accurately reflects whether conversion happened.
+ """
+ print("Setup: Two scenarios - with and without orphaned tool_results...")
+
+ # Scenario 1: With orphaned tool_results (should return True)
+ messages_with_orphaned = [
+ UnifiedMessage(
+ role="user",
+ content="Test",
+ tool_results=[{"type": "tool_result", "tool_use_id": "call_1", "content": "Result"}]
+ )
+ ]
+
+ print("Action: Processing messages with orphaned tool_results...")
+ result1, converted1 = ensure_assistant_before_tool_results(messages_with_orphaned)
+
+ print(f"Comparing converted flag: Expected True, Got {converted1}")
+ assert converted1 is True
+
+ # Scenario 2: Without orphaned tool_results (should return False)
+ messages_without_orphaned = [
+ UnifiedMessage(role="user", content="Hello"),
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{"id": "call_1", "type": "function", "function": {"name": "tool", "arguments": "{}"}}]
+ ),
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{"type": "tool_result", "tool_use_id": "call_1", "content": "Result"}]
+ )
+ ]
+
+ print("Action: Processing messages without orphaned tool_results...")
+ result2, converted2 = ensure_assistant_before_tool_results(messages_without_orphaned)
+
+ print(f"Comparing converted flag: Expected False, Got {converted2}")
+ assert converted2 is False
+
+ def test_normal_tool_results_unchanged(self):
+ """
+ What it does: Verifies that normal (non-orphaned) tool_results are NOT converted.
+ Purpose: CRITICAL - ensure 99% of cases (normal tool use) have zero change.
+
+ This is the most important backward compatibility test. Normal tool_results
+ (with preceding assistant message with tool_calls) should pass through unchanged.
+ """
+ print("Setup: Normal conversation with valid tool_results...")
+ messages = [
+ UnifiedMessage(role="user", content="Call a tool"),
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{
+ "id": "call_valid",
+ "type": "function",
+ "function": {"name": "test_tool", "arguments": "{}"}
+ }]
+ ),
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_valid",
+ "content": "Tool executed successfully"
+ }]
+ )
+ ]
+
+ print("Action: Processing messages...")
+ result, converted = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result: {result}")
+ print(f"Comparing converted flag: Expected False, Got {converted}")
+ assert converted is False # No conversion happened
+
+ print("Checking that tool_results are preserved (NOT converted)...")
+ assert result[2].tool_results is not None # Still has tool_results
+ assert len(result[2].tool_results) == 1
+ assert result[2].tool_results[0]["tool_use_id"] == "call_valid"
+ assert result[2].tool_results[0]["content"] == "Tool executed successfully"
+
+ print("Checking that content is NOT modified...")
+ assert result[2].content == "" # Original empty content preserved
+ assert "[Tool Result" not in result[2].content # NOT converted to text
+
+ def test_mixed_valid_and_orphaned_tool_results(self):
+ """
+ What it does: Verifies correct handling of mixed valid and orphaned tool_results.
+ Purpose: Ensure valid tool_results are preserved while orphaned are stripped.
+ """
+ print("Setup: Mixed conversation with valid and orphaned tool_results...")
+ messages = [
+ # Orphaned tool_results at start
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_orphan",
+ "content": "Orphaned"
+ }]
+ ),
+ # Valid assistant with tool_calls
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{
+ "id": "call_valid",
+ "type": "function",
+ "function": {"name": "valid_tool", "arguments": "{}"}
+ }]
+ ),
+ # Valid tool_results
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_valid",
+ "content": "Valid result"
+ }]
+ )
+ ]
+
+ print("Action: Processing messages...")
+ result, stripped = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result: {result}")
+ print("Checking orphaned tool_results are stripped...")
+ assert result[0].tool_results is None
+
+ print("Checking valid tool_results are preserved...")
+ assert result[2].tool_results is not None
+ assert result[2].tool_results[0]["tool_use_id"] == "call_valid"
+ assert stripped is True # Because orphaned ones were stripped
+
+ def test_single_message_with_tool_results(self):
+ """
+ What it does: Verifies handling of single message with tool_results.
+ Purpose: Ensure single orphaned message is handled correctly.
+ """
+ print("Setup: Single message with tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Result"
+ }]
+ )
+ ]
+
+ print("Action: Processing messages...")
+ result, stripped = ensure_assistant_before_tool_results(messages)
+
+ print(f"Result: {result}")
+ print("Checking that tool_results are stripped...")
+ assert len(result) == 1
+ assert result[0].tool_results is None
+ assert stripped is True
+
+
+# ==================================================================================================
+# Tests for sanitize_json_schema
+# ==================================================================================================
+
+class TestSanitizeJsonSchema:
+ """
+ Tests for sanitize_json_schema function.
+
+ This function cleans JSON Schema from fields that Kiro API doesn't accept:
+ - Empty required arrays []
+ - additionalProperties
+ """
+
+ def test_returns_empty_dict_for_none(self):
+ """
+ What it does: Verifies handling of None.
+ Purpose: Ensure None returns empty dict.
+ """
+ print("Setup: None schema...")
+
+ print("Action: Sanitizing schema...")
+ result = sanitize_json_schema(None)
+
+ print(f"Comparing result: Expected {{}}, Got {result}")
+ assert result == {}
+
+ def test_returns_empty_dict_for_empty_dict(self):
+ """
+ What it does: Verifies handling of empty dict.
+ Purpose: Ensure empty dict is returned as-is.
+ """
+ print("Setup: Empty dict...")
+
+ print("Action: Sanitizing schema...")
+ result = sanitize_json_schema({})
+
+ print(f"Comparing result: Expected {{}}, Got {result}")
+ assert result == {}
+
+ def test_removes_empty_required_array(self):
+ """
+ What it does: Verifies removal of empty required array.
+ Purpose: Ensure required: [] is removed from schema.
+
+ This is a critical test for a bug where tools with required: []
+ caused a 400 "Improperly formed request" error from Kiro API.
+ """
+ print("Setup: Schema with empty required...")
+ schema = {
+ "type": "object",
+ "properties": {},
+ "required": []
+ }
+
+ print("Action: Sanitizing schema...")
+ result = sanitize_json_schema(schema)
+
+ print(f"Result: {result}")
+ print("Checking that required is removed...")
+ assert "required" not in result
+ assert result["type"] == "object"
+ assert result["properties"] == {}
+
+ def test_preserves_non_empty_required_array(self):
+ """
+ What it does: Verifies preservation of non-empty required array.
+ Purpose: Ensure required with elements is preserved.
+ """
+ print("Setup: Schema with non-empty required...")
+ schema = {
+ "type": "object",
+ "properties": {"location": {"type": "string"}},
+ "required": ["location"]
+ }
+
+ print("Action: Sanitizing schema...")
+ result = sanitize_json_schema(schema)
+
+ print(f"Result: {result}")
+ print("Checking that required is preserved...")
+ assert "required" in result
+ assert result["required"] == ["location"]
+
+ def test_removes_additional_properties(self):
+ """
+ What it does: Verifies removal of additionalProperties.
+ Purpose: Ensure additionalProperties is removed from schema.
+
+ Kiro API doesn't support additionalProperties in JSON Schema.
+ """
+ print("Setup: Schema with additionalProperties...")
+ schema = {
+ "type": "object",
+ "properties": {},
+ "additionalProperties": False
+ }
+
+ print("Action: Sanitizing schema...")
+ result = sanitize_json_schema(schema)
+
+ print(f"Result: {result}")
+ print("Checking that additionalProperties is removed...")
+ assert "additionalProperties" not in result
+ assert result["type"] == "object"
+
+ def test_removes_both_empty_required_and_additional_properties(self):
+ """
+ What it does: Verifies removal of both problematic fields.
+ Purpose: Ensure both fields are removed simultaneously.
+ """
+ print("Setup: Schema with both problematic fields...")
+ schema = {
+ "type": "object",
+ "properties": {},
+ "required": [],
+ "additionalProperties": False
+ }
+
+ print("Action: Sanitizing schema...")
+ result = sanitize_json_schema(schema)
+
+ print(f"Result: {result}")
+ print("Checking that both fields are removed...")
+ assert "required" not in result
+ assert "additionalProperties" not in result
+ assert result == {"type": "object", "properties": {}}
+
+ def test_recursively_sanitizes_nested_properties(self):
+ """
+ What it does: Verifies recursive sanitization of nested properties.
+ Purpose: Ensure nested schemas are also sanitized.
+ """
+ print("Setup: Schema with nested properties...")
+ schema = {
+ "type": "object",
+ "properties": {
+ "nested": {
+ "type": "object",
+ "properties": {},
+ "required": [],
+ "additionalProperties": False
+ }
+ }
+ }
+
+ print("Action: Sanitizing schema...")
+ result = sanitize_json_schema(schema)
+
+ print(f"Result: {result}")
+ print("Checking nested object...")
+ nested = result["properties"]["nested"]
+ assert "required" not in nested
+ assert "additionalProperties" not in nested
+
+ def test_sanitizes_items_in_lists(self):
+ """
+ What it does: Verifies sanitization of items in lists (anyOf, oneOf).
+ Purpose: Ensure list elements are also sanitized.
+ """
+ print("Setup: Schema with anyOf...")
+ schema = {
+ "anyOf": [
+ {"type": "string", "additionalProperties": False},
+ {"type": "number", "required": []}
+ ]
+ }
+
+ print("Action: Sanitizing schema...")
+ result = sanitize_json_schema(schema)
+
+ print(f"Result: {result}")
+ print("Checking anyOf elements...")
+ assert "additionalProperties" not in result["anyOf"][0]
+ assert "required" not in result["anyOf"][1]
+
+ def test_preserves_non_dict_list_items(self):
+ """
+ What it does: Verifies preservation of non-dict list items.
+ Purpose: Ensure strings and other types in lists are preserved.
+ """
+ print("Setup: Schema with enum...")
+ schema = {
+ "type": "string",
+ "enum": ["value1", "value2", "value3"]
+ }
+
+ print("Action: Sanitizing schema...")
+ result = sanitize_json_schema(schema)
+
+ print(f"Result: {result}")
+ print("Checking enum is preserved...")
+ assert result["enum"] == ["value1", "value2", "value3"]
+
+ def test_complex_real_world_schema(self):
+ """
+ What it does: Verifies sanitization of real complex schema.
+ Purpose: Ensure real schemas are handled correctly.
+ """
+ print("Setup: Real schema...")
+ schema = {
+ "type": "object",
+ "properties": {
+ "question": {"type": "string", "description": "The question to ask"},
+ "options": {"type": "string", "description": "Array of options"}
+ },
+ "required": ["question", "options"],
+ "additionalProperties": False
+ }
+
+ print("Action: Sanitizing schema...")
+ result = sanitize_json_schema(schema)
+
+ print(f"Result: {result}")
+ print("Checking result...")
+ assert "additionalProperties" not in result
+ assert result["required"] == ["question", "options"] # Non-empty required is preserved
+ assert result["properties"]["question"]["type"] == "string"
+
+
+# ==================================================================================================
+# Tests for extract_tool_results_from_content
+# ==================================================================================================
+
+class TestExtractToolResults:
+ """Tests for extract_tool_results_from_content function."""
+
+ def test_extracts_tool_results_from_list(self):
+ """
+ What it does: Verifies extraction of tool results from list.
+ Purpose: Ensure tool_result elements are extracted.
+ """
+ print("Setup: List with tool_result...")
+ content = [
+ {"type": "tool_result", "tool_use_id": "call_123", "content": "Result text"}
+ ]
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_content(content)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["toolUseId"] == "call_123"
+ assert result[0]["status"] == "success"
+
+ def test_returns_empty_for_string_content(self):
+ """
+ What it does: Verifies empty list return for string.
+ Purpose: Ensure string doesn't contain tool results.
+ """
+ print("Setup: String...")
+ content = "Just a string"
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_returns_empty_for_list_without_tool_results(self):
+ """
+ What it does: Verifies empty list return without tool_result.
+ Purpose: Ensure regular elements are not extracted.
+ """
+ print("Setup: List without tool_result...")
+ content = [{"type": "text", "text": "Hello"}]
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_content(content)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_extracts_multiple_tool_results(self):
+ """
+ What it does: Verifies extraction of multiple tool results.
+ Purpose: Ensure all tool_result elements are extracted.
+ """
+ print("Setup: List with multiple tool_results...")
+ content = [
+ {"type": "tool_result", "tool_use_id": "call_1", "content": "Result 1"},
+ {"type": "text", "text": "Some text"},
+ {"type": "tool_result", "tool_use_id": "call_2", "content": "Result 2"}
+ ]
+
+ print("Action: Extracting tool results...")
+ result = extract_tool_results_from_content(content)
+
+ print(f"Result: {result}")
+ assert len(result) == 2
+ assert result[0]["toolUseId"] == "call_1"
+ assert result[1]["toolUseId"] == "call_2"
+
+
+# ==================================================================================================
+# Tests for convert_tool_results_to_kiro_format
+# ==================================================================================================
+
+class TestConvertToolResultsToKiroFormat:
+ """
+ Tests for convert_tool_results_to_kiro_format function.
+
+ This function converts unified tool results format (snake_case) to Kiro API format (camelCase).
+
+ Unified format: {"type": "tool_result", "tool_use_id": "...", "content": "..."}
+ Kiro format: {"content": [{"text": "..."}], "status": "success", "toolUseId": "..."}
+
+ This is a critical function for fixing the 400 "Improperly formed request" bug
+ where tool_results were sent in unified format instead of Kiro format.
+ """
+
+ def test_converts_single_tool_result(self):
+ """
+ What it does: Verifies conversion of a single tool result.
+ Purpose: Ensure basic conversion from unified to Kiro format works.
+ """
+ print("Setup: Single tool result in unified format...")
+ tool_results = [
+ {"type": "tool_result", "tool_use_id": "call_123", "content": "Result text"}
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_tool_results_to_kiro_format(tool_results)
+
+ print(f"Result: {result}")
+ print("Checking structure...")
+ assert len(result) == 1
+
+ print("Checking toolUseId (camelCase)...")
+ assert result[0]["toolUseId"] == "call_123"
+
+ print("Checking status...")
+ assert result[0]["status"] == "success"
+
+ print("Checking content structure...")
+ assert "content" in result[0]
+ assert isinstance(result[0]["content"], list)
+ assert len(result[0]["content"]) == 1
+ assert result[0]["content"][0]["text"] == "Result text"
+
+ def test_converts_multiple_tool_results(self):
+ """
+ What it does: Verifies conversion of multiple tool results.
+ Purpose: Ensure all tool results are converted correctly.
+ """
+ print("Setup: Multiple tool results...")
+ tool_results = [
+ {"type": "tool_result", "tool_use_id": "call_1", "content": "Result 1"},
+ {"type": "tool_result", "tool_use_id": "call_2", "content": "Result 2"},
+ {"type": "tool_result", "tool_use_id": "call_3", "content": "Result 3"}
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_tool_results_to_kiro_format(tool_results)
+
+ print(f"Result: {result}")
+ print(f"Comparing count: Expected 3, Got {len(result)}")
+ assert len(result) == 3
+
+ print("Checking all toolUseIds...")
+ assert result[0]["toolUseId"] == "call_1"
+ assert result[1]["toolUseId"] == "call_2"
+ assert result[2]["toolUseId"] == "call_3"
+
+ print("Checking all contents...")
+ assert result[0]["content"][0]["text"] == "Result 1"
+ assert result[1]["content"][0]["text"] == "Result 2"
+ assert result[2]["content"][0]["text"] == "Result 3"
+
+ def test_returns_empty_list_for_empty_input(self):
+ """
+ What it does: Verifies empty list handling.
+ Purpose: Ensure empty input returns empty output.
+ """
+ print("Setup: Empty list...")
+
+ print("Action: Converting to Kiro format...")
+ result = convert_tool_results_to_kiro_format([])
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_replaces_empty_content_with_placeholder(self):
+ """
+ What it does: Verifies empty content is replaced with placeholder.
+ Purpose: Ensure Kiro API receives non-empty content (required by API).
+ """
+ print("Setup: Tool result with empty content...")
+ tool_results = [
+ {"type": "tool_result", "tool_use_id": "call_123", "content": ""}
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_tool_results_to_kiro_format(tool_results)
+
+ print(f"Result: {result}")
+ print("Checking that empty content is replaced with placeholder...")
+ assert result[0]["content"][0]["text"] == "(empty result)"
+
+ def test_replaces_none_content_with_placeholder(self):
+ """
+ What it does: Verifies None content is replaced with placeholder.
+ Purpose: Ensure Kiro API receives non-empty content when content is None.
+ """
+ print("Setup: Tool result with None content...")
+ tool_results = [
+ {"type": "tool_result", "tool_use_id": "call_123", "content": None}
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_tool_results_to_kiro_format(tool_results)
+
+ print(f"Result: {result}")
+ print("Checking that None content is replaced with placeholder...")
+ assert result[0]["content"][0]["text"] == "(empty result)"
+
+ def test_handles_missing_content_key(self):
+ """
+ What it does: Verifies handling of missing content key.
+ Purpose: Ensure function doesn't crash when content key is missing.
+ """
+ print("Setup: Tool result without content key...")
+ tool_results = [
+ {"type": "tool_result", "tool_use_id": "call_123"}
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_tool_results_to_kiro_format(tool_results)
+
+ print(f"Result: {result}")
+ print("Checking that missing content is replaced with placeholder...")
+ assert result[0]["content"][0]["text"] == "(empty result)"
+
+ def test_handles_missing_tool_use_id(self):
+ """
+ What it does: Verifies handling of missing tool_use_id.
+ Purpose: Ensure function returns empty string for missing tool_use_id.
+ """
+ print("Setup: Tool result without tool_use_id...")
+ tool_results = [
+ {"type": "tool_result", "content": "Result text"}
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_tool_results_to_kiro_format(tool_results)
+
+ print(f"Result: {result}")
+ print("Checking that missing tool_use_id becomes empty string...")
+ assert result[0]["toolUseId"] == ""
+ assert result[0]["content"][0]["text"] == "Result text"
+
+ def test_extracts_text_from_list_content(self):
+ """
+ What it does: Verifies extraction of text from list content.
+ Purpose: Ensure multimodal content format is handled correctly.
+ """
+ print("Setup: Tool result with list content...")
+ tool_results = [
+ {
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": [
+ {"type": "text", "text": "Part 1"},
+ {"type": "text", "text": " Part 2"}
+ ]
+ }
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_tool_results_to_kiro_format(tool_results)
+
+ print(f"Result: {result}")
+ print("Checking that list content is extracted correctly...")
+ assert result[0]["content"][0]["text"] == "Part 1 Part 2"
+
+ def test_preserves_long_content(self):
+ """
+ What it does: Verifies long content is preserved.
+ Purpose: Ensure large tool results are not truncated.
+ """
+ print("Setup: Tool result with long content...")
+ long_content = "A" * 10000
+ tool_results = [
+ {"type": "tool_result", "tool_use_id": "call_123", "content": long_content}
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_tool_results_to_kiro_format(tool_results)
+
+ print(f"Result content length: {len(result[0]['content'][0]['text'])}")
+ print("Checking that long content is preserved...")
+ assert result[0]["content"][0]["text"] == long_content
+ assert len(result[0]["content"][0]["text"]) == 10000
+
+ def test_all_results_have_success_status(self):
+ """
+ What it does: Verifies all results have status="success".
+ Purpose: Ensure Kiro API receives correct status field.
+ """
+ print("Setup: Multiple tool results...")
+ tool_results = [
+ {"type": "tool_result", "tool_use_id": "call_1", "content": "Result 1"},
+ {"type": "tool_result", "tool_use_id": "call_2", "content": "Result 2"}
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_tool_results_to_kiro_format(tool_results)
+
+ print("Checking all statuses...")
+ for i, r in enumerate(result):
+ print(f"Result {i}: status = {r['status']}")
+ assert r["status"] == "success"
+
+ def test_handles_unicode_content(self):
+ """
+ What it does: Verifies Unicode content is preserved.
+ Purpose: Ensure non-ASCII characters are handled correctly.
+ """
+ print("Setup: Tool result with Unicode content...")
+ tool_results = [
+ {"type": "tool_result", "tool_use_id": "call_123", "content": "Привет мир! 你好世界! 🎉"}
+ ]
+
+ print("Action: Converting to Kiro format...")
+ result = convert_tool_results_to_kiro_format(tool_results)
+
+ print(f"Result: {result}")
+ print("Checking that Unicode content is preserved...")
+ assert result[0]["content"][0]["text"] == "Привет мир! 你好世界! 🎉"
+
+
+# ==================================================================================================
+# Tests for extract_tool_uses_from_message
+# ==================================================================================================
+
+class TestExtractToolUses:
+ """Tests for extract_tool_uses_from_message function."""
+
+ def test_extracts_from_tool_calls_field(self):
+ """
+ What it does: Verifies extraction from tool_calls field.
+ Purpose: Ensure OpenAI tool_calls format is handled.
+ """
+ print("Setup: tool_calls list...")
+ tool_calls = [{
+ "id": "call_123",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"location": "Moscow"}'
+ }
+ }]
+
+ print("Action: Extracting tool uses...")
+ result = extract_tool_uses_from_message(content="", tool_calls=tool_calls)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["name"] == "get_weather"
+ assert result[0]["toolUseId"] == "call_123"
+
+ def test_extracts_from_content_list(self):
+ """
+ What it does: Verifies extraction from content list.
+ Purpose: Ensure tool_use in content is handled (Anthropic format).
+ """
+ print("Setup: Content with tool_use...")
+ content = [{
+ "type": "tool_use",
+ "id": "call_456",
+ "name": "search",
+ "input": {"query": "test"}
+ }]
+
+ print("Action: Extracting tool uses...")
+ result = extract_tool_uses_from_message(content=content, tool_calls=None)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["name"] == "search"
+ assert result[0]["toolUseId"] == "call_456"
+
+ def test_returns_empty_for_no_tool_uses(self):
+ """
+ What it does: Verifies empty list return without tool uses.
+ Purpose: Ensure regular message doesn't contain tool uses.
+ """
+ print("Setup: Regular content...")
+
+ print("Action: Extracting tool uses...")
+ result = extract_tool_uses_from_message(content="Hello", tool_calls=None)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_extracts_from_both_sources(self):
+ """
+ What it does: Verifies extraction from both tool_calls and content.
+ Purpose: Ensure both sources are combined.
+ """
+ print("Setup: Both tool_calls and content with tool_use...")
+ tool_calls = [{
+ "id": "call_1",
+ "function": {"name": "tool1", "arguments": "{}"}
+ }]
+ content = [{
+ "type": "tool_use",
+ "id": "call_2",
+ "name": "tool2",
+ "input": {}
+ }]
+
+ print("Action: Extracting tool uses...")
+ result = extract_tool_uses_from_message(content=content, tool_calls=tool_calls)
+
+ print(f"Result: {result}")
+ assert len(result) == 2
+
+
+# ==================================================================================================
+# Tests for process_tools_with_long_descriptions
+# ==================================================================================================
+
+class TestProcessToolsWithLongDescriptions:
+ """Tests for process_tools_with_long_descriptions function using UnifiedTool."""
+
+ def test_returns_none_and_empty_string_for_none_tools(self):
+ """
+ What it does: Verifies handling of None instead of tools list.
+ Purpose: Ensure None returns (None, "").
+ """
+ print("Setup: None instead of tools...")
+
+ print("Action: Processing tools...")
+ processed, doc = process_tools_with_long_descriptions(None)
+
+ print(f"Comparing result: Expected (None, ''), Got ({processed}, '{doc}')")
+ assert processed is None
+ assert doc == ""
+
+ def test_returns_none_and_empty_string_for_empty_list(self):
+ """
+ What it does: Verifies handling of empty tools list.
+ Purpose: Ensure empty list returns (None, "").
+ """
+ print("Setup: Empty tools list...")
+
+ print("Action: Processing tools...")
+ processed, doc = process_tools_with_long_descriptions([])
+
+ print(f"Comparing result: Expected (None, ''), Got ({processed}, '{doc}')")
+ assert processed is None
+ assert doc == ""
+
+ def test_short_description_unchanged(self):
+ """
+ What it does: Verifies short descriptions are unchanged.
+ Purpose: Ensure tools with short descriptions remain as-is.
+ """
+ print("Setup: Tool with short description...")
+ tools = [UnifiedTool(
+ name="get_weather",
+ description="Get weather for a location",
+ input_schema={"type": "object", "properties": {}}
+ )]
+
+ print("Action: Processing tools...")
+ with patch('kiro.converters_core.TOOL_DESCRIPTION_MAX_LENGTH', 10000):
+ processed, doc = process_tools_with_long_descriptions(tools)
+
+ print(f"Comparing description: Expected 'Get weather for a location', Got '{processed[0].description}'")
+ assert len(processed) == 1
+ assert processed[0].description == "Get weather for a location"
+ assert doc == ""
+
+ def test_long_description_moved_to_system_prompt(self):
+ """
+ What it does: Verifies moving long description to system prompt.
+ Purpose: Ensure long descriptions are moved correctly.
+ """
+ print("Setup: Tool with very long description...")
+ long_description = "A" * 15000 # 15000 chars - exceeds limit
+ tools = [UnifiedTool(
+ name="bash",
+ description=long_description,
+ input_schema={"type": "object", "properties": {"command": {"type": "string"}}}
+ )]
+
+ print("Action: Processing tools with limit 10000...")
+ with patch('kiro.converters_core.TOOL_DESCRIPTION_MAX_LENGTH', 10000):
+ processed, doc = process_tools_with_long_descriptions(tools)
+
+ print("Checking reference in description...")
+ assert len(processed) == 1
+ assert "[Full documentation in system prompt under '## Tool: bash']" in processed[0].description
+
+ print("Checking documentation in system prompt...")
+ assert "## Tool: bash" in doc
+ assert long_description in doc
+ assert "# Tool Documentation" in doc
+
+ def test_mixed_short_and_long_descriptions(self):
+ """
+ What it does: Verifies handling of mixed tools list.
+ Purpose: Ensure short ones stay, long ones are moved.
+ """
+ print("Setup: Two tools - short and long...")
+ short_desc = "Short description"
+ long_desc = "B" * 15000
+ tools = [
+ UnifiedTool(name="short_tool", description=short_desc, input_schema={}),
+ UnifiedTool(name="long_tool", description=long_desc, input_schema={})
+ ]
+
+ print("Action: Processing tools...")
+ with patch('kiro.converters_core.TOOL_DESCRIPTION_MAX_LENGTH', 10000):
+ processed, doc = process_tools_with_long_descriptions(tools)
+
+ print(f"Checking tools count: Expected 2, Got {len(processed)}")
+ assert len(processed) == 2
+
+ print("Checking short tool...")
+ assert processed[0].description == short_desc
+
+ print("Checking long tool...")
+ assert "[Full documentation in system prompt" in processed[1].description
+ assert "## Tool: long_tool" in doc
+ assert long_desc in doc
+
+ def test_disabled_when_limit_is_zero(self):
+ """
+ What it does: Verifies function is disabled when limit is 0.
+ Purpose: Ensure tools are unchanged when TOOL_DESCRIPTION_MAX_LENGTH=0.
+ """
+ print("Setup: Tool with long description and limit 0...")
+ long_desc = "D" * 15000
+ tools = [UnifiedTool(name="test_tool", description=long_desc, input_schema={})]
+
+ print("Action: Processing tools with limit 0...")
+ with patch('kiro.converters_core.TOOL_DESCRIPTION_MAX_LENGTH', 0):
+ processed, doc = process_tools_with_long_descriptions(tools)
+
+ print("Checking that description is unchanged...")
+ assert processed[0].description == long_desc
+ assert doc == ""
+
+ def test_multiple_long_descriptions_all_moved(self):
+ """
+ What it does: Verifies moving of multiple long descriptions.
+ Purpose: Ensure all long descriptions are moved.
+ """
+ print("Setup: Three tools with long descriptions...")
+ tools = [
+ UnifiedTool(name="tool1", description="F" * 15000, input_schema={}),
+ UnifiedTool(name="tool2", description="G" * 15000, input_schema={}),
+ UnifiedTool(name="tool3", description="H" * 15000, input_schema={})
+ ]
+
+ print("Action: Processing tools...")
+ with patch('kiro.converters_core.TOOL_DESCRIPTION_MAX_LENGTH', 10000):
+ processed, doc = process_tools_with_long_descriptions(tools)
+
+ print("Checking all three tools...")
+ assert len(processed) == 3
+ for tool in processed:
+ assert "[Full documentation in system prompt" in tool.description
+
+ print("Checking documentation contains all three sections...")
+ assert "## Tool: tool1" in doc
+ assert "## Tool: tool2" in doc
+ assert "## Tool: tool3" in doc
+
+ def test_empty_description_unchanged(self):
+ """
+ What it does: Verifies handling of empty description.
+ Purpose: Ensure empty description doesn't cause errors.
+ """
+ print("Setup: Tool with empty description...")
+ tools = [UnifiedTool(name="empty_desc_tool", description="", input_schema={})]
+
+ print("Action: Processing tools...")
+ with patch('kiro.converters_core.TOOL_DESCRIPTION_MAX_LENGTH', 10000):
+ processed, doc = process_tools_with_long_descriptions(tools)
+
+ print("Checking that empty description remains empty...")
+ assert processed[0].description == ""
+ assert doc == ""
+
+ def test_none_description_unchanged(self):
+ """
+ What it does: Verifies handling of None description.
+ Purpose: Ensure None description doesn't cause errors.
+ """
+ print("Setup: Tool with None description...")
+ tools = [UnifiedTool(name="none_desc_tool", description=None, input_schema={})]
+
+ print("Action: Processing tools...")
+ with patch('kiro.converters_core.TOOL_DESCRIPTION_MAX_LENGTH', 10000):
+ processed, doc = process_tools_with_long_descriptions(tools)
+
+ print("Checking that None description is handled correctly...")
+ # None should remain None or become empty string
+ assert processed[0].description is None or processed[0].description == ""
+ assert doc == ""
+
+ def test_preserves_tool_input_schema(self):
+ """
+ What it does: Verifies input_schema preservation when moving description.
+ Purpose: Ensure input_schema is not lost.
+ """
+ print("Setup: Tool with input_schema and long description...")
+ input_schema = {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "City name"},
+ "units": {"type": "string", "enum": ["celsius", "fahrenheit"]}
+ },
+ "required": ["location"]
+ }
+ tools = [UnifiedTool(
+ name="weather",
+ description="C" * 15000,
+ input_schema=input_schema
+ )]
+
+ print("Action: Processing tools...")
+ with patch('kiro.converters_core.TOOL_DESCRIPTION_MAX_LENGTH', 10000):
+ processed, doc = process_tools_with_long_descriptions(tools)
+
+ print("Checking input_schema preservation...")
+ assert processed[0].input_schema == input_schema
+
+
+# ==================================================================================================
+# Tests for convert_tools_to_kiro_format
+# ==================================================================================================
+
+class TestConvertToolsToKiroFormat:
+ """Tests for convert_tools_to_kiro_format function."""
+
+ def test_returns_empty_list_for_none(self):
+ """
+ What it does: Verifies handling of None.
+ Purpose: Ensure None returns empty list.
+ """
+ print("Setup: None tools...")
+
+ print("Action: Converting tools...")
+ result = convert_tools_to_kiro_format(None)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_returns_empty_list_for_empty_list(self):
+ """
+ What it does: Verifies handling of empty list.
+ Purpose: Ensure empty list returns empty list.
+ """
+ print("Setup: Empty tools list...")
+
+ print("Action: Converting tools...")
+ result = convert_tools_to_kiro_format([])
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_converts_tool_to_kiro_format(self):
+ """
+ What it does: Verifies conversion of tool to Kiro format.
+ Purpose: Ensure toolSpecification structure is correct.
+ """
+ print("Setup: Tool...")
+ tools = [UnifiedTool(
+ name="get_weather",
+ description="Get weather for a location",
+ input_schema={"type": "object", "properties": {"location": {"type": "string"}}}
+ )]
+
+ print("Action: Converting tools...")
+ result = convert_tools_to_kiro_format(tools)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert "toolSpecification" in result[0]
+ spec = result[0]["toolSpecification"]
+ assert spec["name"] == "get_weather"
+ assert spec["description"] == "Get weather for a location"
+ assert "inputSchema" in spec
+ assert "json" in spec["inputSchema"]
+
+ def test_replaces_empty_description_with_placeholder(self):
+ """
+ What it does: Verifies replacement of empty description.
+ Purpose: Ensure empty description is replaced with "Tool: {name}".
+ """
+ print("Setup: Tool with empty description...")
+ tools = [UnifiedTool(name="focus_chain", description="", input_schema={})]
+
+ print("Action: Converting tools...")
+ result = convert_tools_to_kiro_format(tools)
+
+ print(f"Result: {result}")
+ spec = result[0]["toolSpecification"]
+ assert spec["description"] == "Tool: focus_chain"
+
+ def test_replaces_none_description_with_placeholder(self):
+ """
+ What it does: Verifies replacement of None description.
+ Purpose: Ensure None description is replaced with "Tool: {name}".
+ """
+ print("Setup: Tool with None description...")
+ tools = [UnifiedTool(name="test_tool", description=None, input_schema={})]
+
+ print("Action: Converting tools...")
+ result = convert_tools_to_kiro_format(tools)
+
+ print(f"Result: {result}")
+ spec = result[0]["toolSpecification"]
+ assert spec["description"] == "Tool: test_tool"
+
+ def test_sanitizes_input_schema(self):
+ """
+ What it does: Verifies sanitization of input schema.
+ Purpose: Ensure problematic fields are removed from schema.
+ """
+ print("Setup: Tool with problematic schema...")
+ tools = [UnifiedTool(
+ name="test_tool",
+ description="Test",
+ input_schema={
+ "type": "object",
+ "properties": {},
+ "required": [],
+ "additionalProperties": False
+ }
+ )]
+
+ print("Action: Converting tools...")
+ result = convert_tools_to_kiro_format(tools)
+
+ print(f"Result: {result}")
+ schema = result[0]["toolSpecification"]["inputSchema"]["json"]
+ assert "required" not in schema
+ assert "additionalProperties" not in schema
+
+
+# ==================================================================================================
+# Tests for inject_thinking_tags
+# ==================================================================================================
+
+class TestInjectThinkingTags:
+ """
+ Tests for inject_thinking_tags function.
+
+ This function injects thinking mode tags into content when FAKE_REASONING_ENABLED is True.
+ """
+
+ def test_returns_original_content_when_disabled(self):
+ """
+ What it does: Verifies that content is returned unchanged when fake reasoning is disabled.
+ Purpose: Ensure no modification occurs when FAKE_REASONING_ENABLED=False.
+ """
+ print("Setup: Content with fake reasoning disabled...")
+ content = "Hello, world!"
+
+ print("Action: Inject thinking tags with FAKE_REASONING_ENABLED=False...")
+ with patch('kiro.converters_core.FAKE_REASONING_ENABLED', False):
+ result = inject_thinking_tags(content)
+
+ print(f"Comparing result: Expected 'Hello, world!', Got '{result}'")
+ assert result == "Hello, world!"
+
+ def test_injects_tags_when_enabled(self):
+ """
+ What it does: Verifies that thinking tags are injected when enabled.
+ Purpose: Ensure tags are prepended to content when FAKE_REASONING_ENABLED=True.
+ """
+ print("Setup: Content with fake reasoning enabled...")
+ 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)
+
+ print(f"Result: {result[:200]}...")
+ print("Checking that thinking_mode tag is present...")
+ assert "enabled" in result
+
+ print("Checking that max_thinking_length tag is present...")
+ assert "4000" in result
+
+ print("Checking that original content is preserved at the end...")
+ assert result.endswith("What is 2+2?")
+
+ def test_injects_thinking_instruction_tag(self):
+ """
+ What it does: Verifies that thinking_instruction tag is injected.
+ Purpose: Ensure the quality improvement prompt is included.
+ """
+ print("Setup: Content with fake reasoning enabled...")
+ 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)
+
+ print(f"Result length: {len(result)} chars")
+ print("Checking that thinking_instruction tag is present...")
+ assert "" in result
+ assert "" in result
+
+ def test_thinking_instruction_contains_english_directive(self):
+ """
+ What it does: Verifies that thinking instruction includes English language directive.
+ Purpose: Ensure model is instructed to think in English for better reasoning quality.
+ """
+ print("Setup: Content with fake reasoning enabled...")
+ 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)
+
+ print("Checking for English directive...")
+ assert "Think in English" in result
+
+ def test_uses_configured_max_tokens(self):
+ """
+ What it does: Verifies that FAKE_REASONING_MAX_TOKENS config value is used.
+ Purpose: Ensure the configured max tokens value is injected into the tag.
+ """
+ print("Setup: Content with custom max tokens...")
+ 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):
+ result = inject_thinking_tags(content)
+
+ print(f"Result: {result[:300]}...")
+ print("Checking that max_thinking_length uses configured value...")
+ assert "16000" in result
+
+ def test_preserves_empty_content(self):
+ """
+ What it does: Verifies that empty content is handled correctly.
+ Purpose: Ensure empty string doesn't cause issues.
+ """
+ print("Setup: Empty content with fake reasoning enabled...")
+ 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)
+
+ print(f"Result length: {len(result)} chars")
+ print("Checking that tags are present even with empty content...")
+ assert "enabled" in result
+ assert "" in result
+
+ def test_preserves_multiline_content(self):
+ """
+ What it does: Verifies that multiline content is preserved correctly.
+ Purpose: Ensure newlines in original content are not corrupted.
+ """
+ print("Setup: Multiline content...")
+ 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)
+
+ print("Checking that multiline content is preserved...")
+ assert "Line 1\nLine 2\nLine 3" in result
+
+ def test_preserves_special_characters(self):
+ """
+ What it does: Verifies that special characters in content are preserved.
+ Purpose: Ensure XML-like content in user message doesn't break injection.
+ """
+ print("Setup: Content with special characters...")
+ 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)
+
+ print("Checking that special characters are preserved...")
+ assert "example" in result
+ assert "{json: 'value'}" in result
+
+ def test_thinking_instruction_contains_systematic_approach(self):
+ """
+ What it does: Verifies that thinking instruction includes systematic approach guidance.
+ Purpose: Ensure model is instructed to think systematically.
+ """
+ print("Setup: Content with fake reasoning enabled...")
+ 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)
+
+ print("Checking for systematic approach keywords...")
+ assert "thorough" in result.lower() or "systematic" in result.lower()
+
+ def test_thinking_instruction_contains_understanding_step(self):
+ """
+ What it does: Verifies that thinking instruction includes understanding step.
+ Purpose: Ensure model is instructed to understand the problem first.
+ """
+ print("Setup: Content with fake reasoning enabled...")
+ 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)
+
+ print("Checking for understanding step...")
+ assert "understand" in result.lower()
+
+ def test_thinking_instruction_contains_verification_step(self):
+ """
+ What it does: Verifies that thinking instruction includes verification step.
+ Purpose: Ensure model is instructed to verify reasoning before concluding.
+ """
+ print("Setup: Content with fake reasoning enabled...")
+ 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)
+
+ print("Checking for verification step...")
+ assert "verify" in result.lower()
+
+ def test_thinking_instruction_contains_quality_emphasis(self):
+ """
+ What it does: Verifies that thinking instruction emphasizes quality over speed.
+ Purpose: Ensure model is instructed to prioritize quality of thought.
+ """
+ print("Setup: Content with fake reasoning enabled...")
+ 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)
+
+ print("Checking for quality emphasis...")
+ assert "quality" in result.lower()
+
+ def test_tag_order_is_correct(self):
+ """
+ What it does: Verifies that tags are in the correct order.
+ Purpose: Ensure thinking_mode comes first, then max_thinking_length, then instruction, then content.
+ """
+ print("Setup: Content...")
+ 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)
+
+ print("Checking tag order...")
+ thinking_mode_pos = result.find("")
+ max_length_pos = result.find("")
+ instruction_pos = result.find("")
+ content_pos = result.find("USER_CONTENT_HERE")
+
+ print(f"Positions: thinking_mode={thinking_mode_pos}, max_length={max_length_pos}, instruction={instruction_pos}, content={content_pos}")
+
+ assert thinking_mode_pos < max_length_pos, "thinking_mode should come before max_thinking_length"
+ assert max_length_pos < instruction_pos, "max_thinking_length should come before thinking_instruction"
+ assert instruction_pos < content_pos, "thinking_instruction should come before user content"
+
+
+# ==================================================================================================
+# Tests for build_kiro_history
+# ==================================================================================================
+
+class TestBuildKiroHistory:
+ """Tests for build_kiro_history function using UnifiedMessage."""
+
+ def test_builds_user_message(self):
+ """
+ What it does: Verifies building of user message.
+ Purpose: Ensure user message is converted to userInputMessage.
+ """
+ print("Setup: User message...")
+ messages = [UnifiedMessage(role="user", content="Hello")]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert "userInputMessage" in result[0]
+ assert result[0]["userInputMessage"]["content"] == "Hello"
+ assert result[0]["userInputMessage"]["modelId"] == "claude-sonnet-4"
+
+ def test_builds_assistant_message(self):
+ """
+ What it does: Verifies building of assistant message.
+ Purpose: Ensure assistant message is converted to assistantResponseMessage.
+ """
+ print("Setup: Assistant message...")
+ messages = [UnifiedMessage(role="assistant", content="Hi there")]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert "assistantResponseMessage" in result[0]
+ assert result[0]["assistantResponseMessage"]["content"] == "Hi there"
+
+ def test_ignores_system_messages(self):
+ """
+ What it does: Verifies ignoring of system messages.
+ Purpose: Ensure system messages are not added to history.
+ """
+ print("Setup: System message...")
+ messages = [UnifiedMessage(role="system", content="You are helpful")]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Comparing length: Expected 0, Got {len(result)}")
+ assert len(result) == 0
+
+ def test_builds_conversation_history(self):
+ """
+ What it does: Verifies building of full conversation history.
+ Purpose: Ensure user/assistant alternation is preserved.
+ """
+ print("Setup: Full conversation history...")
+ messages = [
+ UnifiedMessage(role="user", content="Hello"),
+ UnifiedMessage(role="assistant", content="Hi"),
+ UnifiedMessage(role="user", content="How are you?")
+ ]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ assert len(result) == 3
+ assert "userInputMessage" in result[0]
+ assert "assistantResponseMessage" in result[1]
+ assert "userInputMessage" in result[2]
+
+ def test_handles_empty_list(self):
+ """
+ What it does: Verifies empty list handling.
+ Purpose: Ensure empty list returns empty history.
+ """
+ print("Setup: Empty list...")
+
+ print("Action: Building history...")
+ result = build_kiro_history([], "claude-sonnet-4")
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_builds_user_message_with_tool_results(self):
+ """
+ What it does: Verifies building of user message with tool_results.
+ Purpose: Ensure tool_results are included in userInputMessageContext.
+ """
+ print("Setup: User message with tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="Here are the results",
+ tool_results=[
+ {"type": "tool_result", "tool_use_id": "call_123", "content": "Result text"}
+ ]
+ )
+ ]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert "userInputMessage" in result[0]
+ user_msg = result[0]["userInputMessage"]
+ assert "userInputMessageContext" in user_msg
+ assert "toolResults" in user_msg["userInputMessageContext"]
+
+ def test_builds_assistant_message_with_tool_calls(self):
+ """
+ What it does: Verifies building of assistant message with tool_calls.
+ Purpose: Ensure tool_calls are converted to toolUses.
+ """
+ print("Setup: Assistant message with tool_calls...")
+ messages = [
+ UnifiedMessage(
+ role="assistant",
+ content="I'll call a tool",
+ tool_calls=[{
+ "id": "call_123",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"location": "Moscow"}'
+ }
+ }]
+ )
+ ]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert "assistantResponseMessage" in result[0]
+ assistant_msg = result[0]["assistantResponseMessage"]
+ assert "toolUses" in assistant_msg
+
+ def test_adds_empty_placeholder_for_empty_user_content(self):
+ """
+ What it does: Verifies that "(empty)" placeholder is added for user messages with empty content.
+ Purpose: Ensure Kiro API receives non-empty content in history.
+
+ This is a fallback test for issue #20 - ensures any edge case with empty content
+ is handled even if strip_all_tool_content didn't add a placeholder.
+ """
+ print("Setup: User message with empty content...")
+ messages = [UnifiedMessage(role="user", content="")]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ print(f"Content: '{result[0]['userInputMessage']['content']}'")
+ print("Checking that '(empty)' placeholder is added...")
+ assert result[0]["userInputMessage"]["content"] == "(empty)"
+
+ def test_adds_empty_placeholder_for_empty_assistant_content(self):
+ """
+ What it does: Verifies that "(empty)" placeholder is added for assistant messages with empty content.
+ Purpose: Ensure Kiro API receives non-empty content in history.
+
+ This is a fallback test for issue #20 - ensures any edge case with empty content
+ is handled even if strip_all_tool_content didn't add a placeholder.
+ """
+ print("Setup: Assistant message with empty content...")
+ messages = [UnifiedMessage(role="assistant", content="")]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ print(f"Content: '{result[0]['assistantResponseMessage']['content']}'")
+ print("Checking that '(empty)' placeholder is added...")
+ assert result[0]["assistantResponseMessage"]["content"] == "(empty)"
+
+ def test_adds_empty_placeholder_for_none_user_content(self):
+ """
+ What it does: Verifies that "(empty)" placeholder is added for user messages with None content.
+ Purpose: Ensure Kiro API receives non-empty content when content is None.
+ """
+ print("Setup: User message with None content...")
+ messages = [UnifiedMessage(role="user", content=None)]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ print(f"Content: '{result[0]['userInputMessage']['content']}'")
+ print("Checking that '(empty)' placeholder is added...")
+ assert result[0]["userInputMessage"]["content"] == "(empty)"
+
+ def test_adds_empty_placeholder_for_none_assistant_content(self):
+ """
+ What it does: Verifies that "(empty)" placeholder is added for assistant messages with None content.
+ Purpose: Ensure Kiro API receives non-empty content when content is None.
+ """
+ print("Setup: Assistant message with None content...")
+ messages = [UnifiedMessage(role="assistant", content=None)]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ print(f"Content: '{result[0]['assistantResponseMessage']['content']}'")
+ print("Checking that '(empty)' placeholder is added...")
+ assert result[0]["assistantResponseMessage"]["content"] == "(empty)"
+
+ def test_preserves_non_empty_content_in_history(self):
+ """
+ What it does: Verifies that non-empty content is preserved (not replaced with placeholder).
+ Purpose: Ensure placeholder is only added when content is actually empty.
+ """
+ print("Setup: Messages with actual content...")
+ messages = [
+ UnifiedMessage(role="user", content="Hello"),
+ UnifiedMessage(role="assistant", content="Hi there")
+ ]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ print("Checking that original content is preserved...")
+ assert result[0]["userInputMessage"]["content"] == "Hello"
+ assert result[1]["assistantResponseMessage"]["content"] == "Hi there"
+
+ def test_mixed_empty_and_non_empty_content_in_history(self):
+ """
+ What it does: Verifies correct handling of mixed empty and non-empty content.
+ Purpose: Ensure only empty messages get placeholders.
+
+ This simulates a conversation where some messages have content and some don't.
+ """
+ print("Setup: Mixed conversation with empty and non-empty content...")
+ messages = [
+ UnifiedMessage(role="user", content="Start"),
+ UnifiedMessage(role="assistant", content=""), # Empty - should get placeholder
+ UnifiedMessage(role="user", content=""), # Empty - should get placeholder
+ UnifiedMessage(role="assistant", content="Response")
+ ]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ print("Checking each message...")
+
+ print(f"Message 0 content: '{result[0]['userInputMessage']['content']}'")
+ assert result[0]["userInputMessage"]["content"] == "Start"
+
+ print(f"Message 1 content: '{result[1]['assistantResponseMessage']['content']}'")
+ assert result[1]["assistantResponseMessage"]["content"] == "(empty)"
+
+ print(f"Message 2 content: '{result[2]['userInputMessage']['content']}'")
+ assert result[2]["userInputMessage"]["content"] == "(empty)"
+
+ print(f"Message 3 content: '{result[3]['assistantResponseMessage']['content']}'")
+ assert result[3]["assistantResponseMessage"]["content"] == "Response"
+
+ def test_builds_user_message_with_images(self):
+ """
+ What it does: Verifies building of user message with images.
+ Purpose: Ensure images are included directly in userInputMessage.images (Issue #32 fix).
+
+ This is a critical test for Issue #30/#32 fix - images should be in Kiro format
+ and placed directly in userInputMessage, NOT in userInputMessageContext.
+ """
+ print("Setup: User message with images...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="What's in this image?",
+ images=[{"media_type": "image/jpeg", "data": TEST_IMAGE_BASE64}]
+ )
+ ]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert "userInputMessage" in result[0]
+
+ user_msg = result[0]["userInputMessage"]
+ print(f"User message: {user_msg}")
+
+ print("Checking that images are directly in userInputMessage (Issue #32 fix)...")
+ assert "images" in user_msg
+
+ print("Checking image format (Kiro format)...")
+ images = user_msg["images"]
+ assert len(images) == 1
+ assert images[0]["format"] == "jpeg"
+ assert images[0]["source"]["bytes"] == TEST_IMAGE_BASE64
+
+ def test_builds_user_message_with_multiple_images(self):
+ """
+ What it does: Verifies building of user message with multiple images.
+ Purpose: Ensure all images are included directly in userInputMessage (Issue #32 fix).
+ """
+ print("Setup: User message with multiple images...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="Compare these images",
+ images=[
+ {"media_type": "image/jpeg", "data": "image1_data"},
+ {"media_type": "image/png", "data": "image2_data"}
+ ]
+ )
+ ]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ user_msg = result[0]["userInputMessage"]
+ images = user_msg["images"]
+
+ print(f"Comparing image count: Expected 2, Got {len(images)}")
+ assert len(images) == 2
+
+ print("Checking first image...")
+ assert images[0]["format"] == "jpeg"
+ assert images[0]["source"]["bytes"] == "image1_data"
+
+ print("Checking second image...")
+ assert images[1]["format"] == "png"
+ assert images[1]["source"]["bytes"] == "image2_data"
+
+ def test_builds_user_message_with_images_and_tool_results(self):
+ """
+ What it does: Verifies building of user message with both images and tool_results.
+ Purpose: Ensure images are in userInputMessage and toolResults are in userInputMessageContext (Issue #32 fix).
+ """
+ print("Setup: User message with images and tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="Here's the image and tool result",
+ images=[{"media_type": "image/png", "data": "image_data"}],
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Tool output"
+ }]
+ )
+ ]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ user_msg = result[0]["userInputMessage"]
+ context = user_msg.get("userInputMessageContext", {})
+
+ print("Checking that images are directly in userInputMessage (Issue #32 fix)...")
+ assert "images" in user_msg
+
+ print("Checking that toolResults are in userInputMessageContext...")
+ assert "toolResults" in context
+
+ print("Checking images...")
+ assert len(user_msg["images"]) == 1
+ assert user_msg["images"][0]["format"] == "png"
+
+ print("Checking toolResults...")
+ assert len(context["toolResults"]) == 1
+ assert context["toolResults"][0]["toolUseId"] == "call_123"
+
+ def test_no_images_context_when_no_images(self):
+ """
+ What it does: Verifies that images key is not added when there are no images.
+ Purpose: Ensure clean payload without empty images array.
+ """
+ print("Setup: User message without images...")
+ messages = [
+ UnifiedMessage(role="user", content="Hello, no images here")
+ ]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ user_msg = result[0]["userInputMessage"]
+
+ print("Checking that images key is not present...")
+ # Either no context at all, or context without images
+ if "userInputMessageContext" in user_msg:
+ context = user_msg["userInputMessageContext"]
+ assert "images" not in context or context.get("images") == []
+ else:
+ print("No userInputMessageContext - OK")
+
+ def test_builds_user_message_with_webp_image(self):
+ """
+ What it does: Verifies building of user message with WebP image.
+ Purpose: Ensure WebP format is correctly converted to Kiro format in userInputMessage (Issue #32 fix).
+ """
+ print("Setup: User message with WebP image...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="Analyze this WebP image",
+ images=[{"media_type": "image/webp", "data": "webp_image_data"}]
+ )
+ ]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ user_msg = result[0]["userInputMessage"]
+ images = user_msg["images"]
+
+ print("Checking WebP format...")
+ assert len(images) == 1
+ assert images[0]["format"] == "webp"
+ assert images[0]["source"]["bytes"] == "webp_image_data"
+
+ def test_builds_user_message_with_gif_image(self):
+ """
+ What it does: Verifies building of user message with GIF image.
+ Purpose: Ensure GIF format is correctly converted to Kiro format in userInputMessage (Issue #32 fix).
+ """
+ print("Setup: User message with GIF image...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="What's happening in this GIF?",
+ images=[{"media_type": "image/gif", "data": "gif_image_data"}]
+ )
+ ]
+
+ print("Action: Building history...")
+ result = build_kiro_history(messages, "claude-sonnet-4")
+
+ print(f"Result: {result}")
+ user_msg = result[0]["userInputMessage"]
+ images = user_msg["images"]
+
+ print("Checking GIF format...")
+ assert len(images) == 1
+ assert images[0]["format"] == "gif"
+ assert images[0]["source"]["bytes"] == "gif_image_data"
+
+
+# ==================================================================================================
+# Tests for strip_all_tool_content
+# ==================================================================================================
+
+class TestStripAllToolContent:
+ """
+ Tests for strip_all_tool_content function.
+
+ This function strips ALL tool-related content (tool_calls and tool_results)
+ from messages. It is used when no tools are defined in the request, because
+ Kiro API rejects requests that have toolResults but no tools defined.
+
+ This is a critical function for handling clients like Cline/Roo/Cursor that may
+ send tool-related content even when tools are not available.
+ """
+
+ def test_returns_empty_list_for_empty_input(self):
+ """
+ What it does: Verifies empty list handling.
+ Purpose: Ensure empty input returns empty output.
+ """
+ print("Setup: Empty list...")
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content([])
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+ assert had_content is False
+
+ def test_preserves_messages_without_tool_content(self):
+ """
+ What it does: Verifies messages without tool content are unchanged.
+ Purpose: Ensure regular messages pass through unmodified.
+ """
+ print("Setup: Messages without tool content...")
+ messages = [
+ UnifiedMessage(role="user", content="Hello"),
+ UnifiedMessage(role="assistant", content="Hi there"),
+ UnifiedMessage(role="user", content="How are you?")
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Comparing length: Expected 3, Got {len(result)}")
+ assert len(result) == 3
+ assert result[0].content == "Hello"
+ assert result[1].content == "Hi there"
+ assert result[2].content == "How are you?"
+ assert had_content is False
+
+ def test_strips_tool_calls_from_assistant(self):
+ """
+ What it does: Verifies tool_calls are stripped and converted to text.
+ Purpose: Ensure tool_calls are converted to text representation when no tools are defined.
+ """
+ print("Setup: Assistant message with tool_calls...")
+ messages = [
+ UnifiedMessage(
+ role="assistant",
+ content="I'll call a tool",
+ tool_calls=[{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "get_weather", "arguments": '{"location": "Moscow"}'}
+ }]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print("Checking that tool_calls are stripped and converted to text...")
+ assert len(result) == 1
+ assert result[0].tool_calls is None
+ # Original content is preserved AND tool text is appended
+ assert "I'll call a tool" in result[0].content
+ assert "[Tool: get_weather" in result[0].content
+ assert had_content is True
+
+ def test_strips_tool_results_from_user(self):
+ """
+ What it does: Verifies tool_results are stripped and converted to text.
+ Purpose: Ensure tool_results are converted to text representation when no tools are defined.
+ """
+ print("Setup: User message with tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="Here are the results",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Weather is sunny"
+ }]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print("Checking that tool_results are stripped and converted to text...")
+ assert len(result) == 1
+ assert result[0].tool_results is None
+ # Original content is preserved AND tool result text is appended
+ assert "Here are the results" in result[0].content
+ assert "[Tool Result" in result[0].content
+ assert "Weather is sunny" in result[0].content
+ assert had_content is True
+
+ def test_strips_both_tool_calls_and_tool_results(self):
+ """
+ What it does: Verifies both tool_calls and tool_results are stripped.
+ Purpose: Ensure all tool content is removed in a conversation.
+ """
+ print("Setup: Conversation with tool_calls and tool_results...")
+ messages = [
+ UnifiedMessage(role="user", content="Call a tool"),
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "get_weather", "arguments": "{}"}
+ }]
+ ),
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Result"
+ }]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print("Checking that all tool content is stripped...")
+ assert len(result) == 3
+ assert result[0].tool_calls is None
+ assert result[0].tool_results is None
+ assert result[1].tool_calls is None
+ assert result[1].tool_results is None
+ assert result[2].tool_calls is None
+ assert result[2].tool_results is None
+ assert had_content is True
+
+ def test_strips_multiple_tool_calls(self):
+ """
+ What it does: Verifies multiple tool_calls are all stripped.
+ Purpose: Ensure all tool_calls in a message are removed.
+ """
+ print("Setup: Assistant message with multiple tool_calls...")
+ messages = [
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[
+ {"id": "call_1", "type": "function", "function": {"name": "tool1", "arguments": "{}"}},
+ {"id": "call_2", "type": "function", "function": {"name": "tool2", "arguments": "{}"}},
+ {"id": "call_3", "type": "function", "function": {"name": "tool3", "arguments": "{}"}}
+ ]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print("Checking that all tool_calls are stripped...")
+ assert result[0].tool_calls is None
+ assert had_content is True
+
+ def test_strips_multiple_tool_results(self):
+ """
+ What it does: Verifies multiple tool_results are all stripped.
+ Purpose: Ensure all tool_results in a message are removed.
+ """
+ print("Setup: User message with multiple tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[
+ {"type": "tool_result", "tool_use_id": "call_1", "content": "Result 1"},
+ {"type": "tool_result", "tool_use_id": "call_2", "content": "Result 2"},
+ {"type": "tool_result", "tool_use_id": "call_3", "content": "Result 3"}
+ ]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print("Checking that all tool_results are stripped...")
+ assert result[0].tool_results is None
+ assert had_content is True
+
+ def test_preserves_message_content_when_stripping(self):
+ """
+ What it does: Verifies message content is preserved and tool content is appended as text.
+ Purpose: Ensure original content is kept and tool content is converted to text.
+ """
+ print("Setup: Messages with both content and tool content...")
+ messages = [
+ UnifiedMessage(
+ role="assistant",
+ content="Let me help you with that",
+ tool_calls=[{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "helper", "arguments": "{}"}
+ }]
+ ),
+ UnifiedMessage(
+ role="user",
+ content="Thanks for the result",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Done"
+ }]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print("Checking that original content is preserved and tool text is appended...")
+ assert "Let me help you with that" in result[0].content
+ assert "[Tool: helper" in result[0].content
+ assert "Thanks for the result" in result[1].content
+ assert "[Tool Result" in result[1].content
+ assert had_content is True
+
+ def test_preserves_message_role_when_stripping(self):
+ """
+ What it does: Verifies message role is preserved when tool content is stripped.
+ Purpose: Ensure role is not modified during stripping.
+ """
+ print("Setup: Messages with tool content...")
+ messages = [
+ UnifiedMessage(role="assistant", content="", tool_calls=[
+ {"id": "call_1", "type": "function", "function": {"name": "tool", "arguments": "{}"}}
+ ]),
+ UnifiedMessage(role="user", content="", tool_results=[
+ {"type": "tool_result", "tool_use_id": "call_1", "content": "Result"}
+ ])
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print("Checking that roles are preserved...")
+ assert result[0].role == "assistant"
+ assert result[1].role == "user"
+ assert had_content is True
+
+ def test_mixed_messages_with_and_without_tool_content(self):
+ """
+ What it does: Verifies correct handling of mixed messages.
+ Purpose: Ensure only messages with tool content are modified.
+ """
+ print("Setup: Mixed messages...")
+ messages = [
+ UnifiedMessage(role="user", content="Hello"), # No tool content
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{"id": "call_1", "type": "function", "function": {"name": "tool", "arguments": "{}"}}]
+ ), # Has tool content
+ UnifiedMessage(role="user", content="Continue"), # No tool content
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print("Checking mixed handling...")
+ assert result[0].content == "Hello"
+ assert result[0].tool_calls is None
+ assert result[1].tool_calls is None # Stripped
+ assert result[2].content == "Continue"
+ assert result[2].tool_calls is None
+ assert had_content is True
+
+ def test_returns_false_when_no_tool_content_stripped(self):
+ """
+ What it does: Verifies had_content flag is False when no tool content exists.
+ Purpose: Ensure correct flag value for messages without tool content.
+ """
+ print("Setup: Messages without any tool content...")
+ messages = [
+ UnifiedMessage(role="user", content="Hello"),
+ UnifiedMessage(role="assistant", content="Hi"),
+ UnifiedMessage(role="user", content="Bye")
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"had_content: {had_content}")
+ assert had_content is False
+
+ def test_returns_true_when_tool_content_stripped(self):
+ """
+ What it does: Verifies had_content flag is True when tool content is stripped.
+ Purpose: Ensure correct flag value for messages with tool content.
+ """
+ print("Setup: Message with tool content...")
+ messages = [
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{"id": "call_1", "type": "function", "function": {"name": "tool", "arguments": "{}"}}]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"had_content: {had_content}")
+ assert had_content is True
+
+ def test_handles_empty_tool_calls_list(self):
+ """
+ What it does: Verifies handling of empty tool_calls list.
+ Purpose: Ensure empty list is treated as no tool content.
+ """
+ print("Setup: Message with empty tool_calls list...")
+ messages = [
+ UnifiedMessage(role="assistant", content="Hello", tool_calls=[])
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print(f"had_content: {had_content}")
+ # Empty list is falsy, so should not be considered as having tool content
+ assert had_content is False
+
+ def test_handles_empty_tool_results_list(self):
+ """
+ What it does: Verifies handling of empty tool_results list.
+ Purpose: Ensure empty list is treated as no tool content.
+ """
+ print("Setup: Message with empty tool_results list...")
+ messages = [
+ UnifiedMessage(role="user", content="Hello", tool_results=[])
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print(f"had_content: {had_content}")
+ # Empty list is falsy, so should not be considered as having tool content
+ assert had_content is False
+
+ def test_adds_tool_text_for_empty_content_with_tool_calls(self):
+ """
+ What it does: Verifies that tool_calls are converted to text when content is empty.
+ Purpose: Ensure Kiro API receives non-empty content for messages that only had tool_calls.
+
+ This is a critical test for issue #20 - OpenCode compaction returns 400 error
+ because messages with only tool_calls become empty after stripping.
+ Now we convert tool_calls to text representation instead of simple placeholder.
+ """
+ print("Setup: Assistant message with only tool_calls (empty content)...")
+ messages = [
+ UnifiedMessage(
+ role="assistant",
+ content="", # Empty content - only tool_calls
+ tool_calls=[{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "read_file", "arguments": '{"path": "test.py"}'}
+ }]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print(f"Content after stripping: '{result[0].content}'")
+ print("Checking that tool_calls are converted to text representation...")
+ assert "[Tool: read_file" in result[0].content
+ assert "call_123" in result[0].content
+ assert '{"path": "test.py"}' in result[0].content
+ assert result[0].tool_calls is None
+ assert had_content is True
+
+ def test_adds_tool_text_for_empty_content_with_tool_results(self):
+ """
+ What it does: Verifies that tool_results are converted to text when content is empty.
+ Purpose: Ensure Kiro API receives non-empty content for messages that only had tool_results.
+
+ This is a critical test for issue #20 - OpenCode compaction returns 400 error
+ because messages with only tool_results become empty after stripping.
+ Now we convert tool_results to text representation instead of simple placeholder.
+ """
+ print("Setup: User message with only tool_results (empty content)...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="", # Empty content - only tool_results
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "File contents here"
+ }]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print(f"Content after stripping: '{result[0].content}'")
+ print("Checking that tool_results are converted to text representation...")
+ assert "[Tool Result" in result[0].content
+ assert "call_123" in result[0].content
+ assert "File contents here" in result[0].content
+ assert result[0].tool_results is None
+ assert had_content is True
+
+ def test_preserves_existing_content_when_stripping_tool_calls(self):
+ """
+ What it does: Verifies that existing content is preserved and tool text is appended.
+ Purpose: Ensure original content is kept and tool_calls are converted to text.
+ """
+ print("Setup: Assistant message with both content and tool_calls...")
+ messages = [
+ UnifiedMessage(
+ role="assistant",
+ content="I'll read the file for you",
+ tool_calls=[{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "read_file", "arguments": "{}"}
+ }]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print(f"Content after stripping: '{result[0].content}'")
+ print("Checking that original content is preserved and tool text is appended...")
+ assert "I'll read the file for you" in result[0].content
+ assert "[Tool: read_file" in result[0].content
+ assert result[0].tool_calls is None
+ assert had_content is True
+
+ def test_preserves_existing_content_when_stripping_tool_results(self):
+ """
+ What it does: Verifies that existing content is preserved and tool result text is appended.
+ Purpose: Ensure original content is kept and tool_results are converted to text.
+ """
+ print("Setup: User message with both content and tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="Here are the results you requested",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Result data"
+ }]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print(f"Content after stripping: '{result[0].content}'")
+ print("Checking that original content is preserved and tool result text is appended...")
+ assert "Here are the results you requested" in result[0].content
+ assert "[Tool Result" in result[0].content
+ assert "Result data" in result[0].content
+ assert result[0].tool_results is None
+ assert had_content is True
+
+ def test_both_tool_calls_and_results_converted_to_text(self):
+ """
+ What it does: Verifies that both tool_calls and tool_results are converted to text.
+ Purpose: Ensure all tool content is preserved when message has both types.
+
+ Note: This is an edge case - normally assistant messages have tool_calls and user messages have tool_results.
+ """
+ print("Setup: Message with both tool_calls and tool_results (edge case)...")
+ messages = [
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{"id": "call_1", "type": "function", "function": {"name": "my_tool", "arguments": '{"x": 1}'}}],
+ tool_results=[{"type": "tool_result", "tool_use_id": "call_0", "content": "Previous result"}]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print(f"Content after stripping: '{result[0].content}'")
+ print("Checking that both tool_calls and tool_results are converted to text...")
+ assert "[Tool: my_tool" in result[0].content
+ assert "[Tool Result" in result[0].content
+ assert "Previous result" in result[0].content
+ assert had_content is True
+
+ def test_multiple_messages_with_empty_content_get_text_representation(self):
+ """
+ What it does: Verifies correct text representation for multiple messages in a conversation.
+ Purpose: Ensure each message gets the appropriate text representation based on its tool content type.
+
+ This simulates the OpenCode compaction scenario from issue #20 where multiple
+ tool-only messages are sent without text content.
+ """
+ print("Setup: Conversation with multiple tool-only messages...")
+ messages = [
+ UnifiedMessage(role="user", content="Read these files"),
+ UnifiedMessage(
+ role="assistant",
+ content="", # Only tool_calls
+ tool_calls=[{"id": "call_1", "type": "function", "function": {"name": "read_file", "arguments": '{"path": "a.txt"}'}}]
+ ),
+ UnifiedMessage(
+ role="user",
+ content="", # Only tool_results
+ tool_results=[{"type": "tool_result", "tool_use_id": "call_1", "content": "File content ABC"}]
+ ),
+ UnifiedMessage(
+ role="assistant",
+ content="", # Only tool_calls
+ tool_calls=[{"id": "call_2", "type": "function", "function": {"name": "write_file", "arguments": '{"path": "b.txt"}'}}]
+ ),
+ UnifiedMessage(
+ role="user",
+ content="", # Only tool_results
+ tool_results=[{"type": "tool_result", "tool_use_id": "call_2", "content": "Write completed"}]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result: {result}")
+ print("Checking text representation for each message...")
+
+ print(f"Message 0 content: '{result[0].content}'")
+ assert result[0].content == "Read these files" # Original content preserved
+
+ print(f"Message 1 content: '{result[1].content}'")
+ assert "[Tool: read_file" in result[1].content # Text representation for tool_calls
+ assert "call_1" in result[1].content
+
+ print(f"Message 2 content: '{result[2].content}'")
+ assert "[Tool Result" in result[2].content # Text representation for tool_results
+ assert "File content ABC" in result[2].content
+
+ print(f"Message 3 content: '{result[3].content}'")
+ assert "[Tool: write_file" in result[3].content # Text representation for tool_calls
+ assert "call_2" in result[3].content
+
+ print(f"Message 4 content: '{result[4].content}'")
+ assert "[Tool Result" in result[4].content # Text representation for tool_results
+ assert "Write completed" in result[4].content
+
+ assert had_content is True
+
+ def test_converts_tool_calls_to_text_representation(self):
+ """
+ What it does: Verifies that tool_calls are converted to text representation.
+ Purpose: Ensure tool context is preserved as readable text when stripping.
+
+ This is a critical test for issue #20 - instead of losing tool context,
+ we convert it to human-readable text.
+ """
+ print("Setup: Assistant message with tool_calls...")
+ messages = [
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{
+ "id": "call_abc123",
+ "type": "function",
+ "function": {"name": "read_file", "arguments": '{"path": "test.py"}'}
+ }]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result content: '{result[0].content}'")
+ print("Checking that tool name is in text representation...")
+ assert "[Tool: read_file" in result[0].content
+ print("Checking that tool_id is in text representation...")
+ assert "call_abc123" in result[0].content
+ print("Checking that arguments are in text representation...")
+ assert '{"path": "test.py"}' in result[0].content
+ assert had_content is True
+
+ def test_converts_tool_results_to_text_representation(self):
+ """
+ What it does: Verifies that tool_results are converted to text representation.
+ Purpose: Ensure tool result context is preserved as readable text when stripping.
+
+ This is a critical test for issue #20 - instead of losing tool context,
+ we convert it to human-readable text.
+ """
+ print("Setup: User message with tool_results...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_xyz789",
+ "content": "File contents:\ndef hello():\n print('world')"
+ }]
+ )
+ ]
+
+ print("Action: Stripping tool content...")
+ result, had_content = strip_all_tool_content(messages)
+
+ print(f"Result content: '{result[0].content}'")
+ print("Checking that [Tool Result] marker is present...")
+ assert "[Tool Result" in result[0].content
+ print("Checking that tool_use_id is in text representation...")
+ assert "call_xyz789" in result[0].content
+ print("Checking that result content is preserved...")
+ assert "def hello():" in result[0].content
+ assert had_content is True
+
+
+# ==================================================================================================
+# Tests for tool_calls_to_text
+# ==================================================================================================
+
+class TestToolCallsToText:
+ """
+ Tests for tool_calls_to_text function.
+
+ This function converts tool_calls to human-readable text representation.
+ Used when stripping tool content from messages (when no tools are defined).
+ """
+
+ def test_converts_single_tool_call_to_text(self):
+ """
+ What it does: Verifies conversion of a single tool call to text.
+ Purpose: Ensure basic conversion works correctly.
+ """
+ print("Setup: Single tool call...")
+ tool_calls = [{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "bash", "arguments": '{"command": "ls -la"}'}
+ }]
+
+ print("Action: Converting to text...")
+ result = tool_calls_to_text(tool_calls)
+
+ print(f"Result: '{result}'")
+ print("Checking that tool name is present...")
+ assert "[Tool: bash" in result
+ print("Checking that arguments are present...")
+ assert '{"command": "ls -la"}' in result
+
+ def test_converts_multiple_tool_calls_to_text(self):
+ """
+ What it does: Verifies conversion of multiple tool calls to text.
+ Purpose: Ensure all tool calls are converted and separated.
+ """
+ print("Setup: Multiple tool calls...")
+ tool_calls = [
+ {"id": "call_1", "type": "function", "function": {"name": "read_file", "arguments": '{"path": "a.txt"}'}},
+ {"id": "call_2", "type": "function", "function": {"name": "write_file", "arguments": '{"path": "b.txt"}'}}
+ ]
+
+ print("Action: Converting to text...")
+ result = tool_calls_to_text(tool_calls)
+
+ print(f"Result: '{result}'")
+ print("Checking that both tools are present...")
+ assert "[Tool: read_file" in result
+ assert "[Tool: write_file" in result
+ assert '{"path": "a.txt"}' in result
+ assert '{"path": "b.txt"}' in result
+
+ def test_includes_tool_id_in_output(self):
+ """
+ What it does: Verifies that tool_id is included in output.
+ Purpose: Ensure traceability between tool calls and results.
+ """
+ print("Setup: Tool call with id...")
+ tool_calls = [{
+ "id": "tooluse_abc123xyz",
+ "type": "function",
+ "function": {"name": "search", "arguments": "{}"}
+ }]
+
+ print("Action: Converting to text...")
+ result = tool_calls_to_text(tool_calls)
+
+ print(f"Result: '{result}'")
+ print("Checking that tool_id is present...")
+ assert "tooluse_abc123xyz" in result
+
+ def test_handles_missing_tool_id(self):
+ """
+ What it does: Verifies handling of tool call without id.
+ Purpose: Ensure function doesn't crash when id is missing.
+ """
+ print("Setup: Tool call without id...")
+ tool_calls = [{
+ "type": "function",
+ "function": {"name": "test_tool", "arguments": "{}"}
+ }]
+
+ print("Action: Converting to text...")
+ result = tool_calls_to_text(tool_calls)
+
+ print(f"Result: '{result}'")
+ print("Checking that tool name is still present...")
+ assert "[Tool: test_tool]" in result
+
+ def test_returns_empty_string_for_empty_list(self):
+ """
+ What it does: Verifies empty list handling.
+ Purpose: Ensure empty input returns empty output.
+ """
+ print("Setup: Empty list...")
+
+ print("Action: Converting to text...")
+ result = tool_calls_to_text([])
+
+ print(f"Comparing result: Expected '', Got '{result}'")
+ assert result == ""
+
+ def test_handles_missing_function_key(self):
+ """
+ What it does: Verifies handling of malformed tool call without function key.
+ Purpose: Ensure function doesn't crash on malformed input.
+ """
+ print("Setup: Tool call without function key...")
+ tool_calls = [{"id": "call_123", "type": "function"}]
+
+ print("Action: Converting to text...")
+ result = tool_calls_to_text(tool_calls)
+
+ print(f"Result: '{result}'")
+ print("Checking that 'unknown' is used as fallback...")
+ assert "[Tool: unknown" in result
+
+ def test_handles_complex_json_arguments(self):
+ """
+ What it does: Verifies handling of complex JSON arguments.
+ Purpose: Ensure nested JSON is preserved correctly.
+ """
+ print("Setup: Tool call with complex arguments...")
+ complex_args = '{"files": ["a.py", "b.py"], "options": {"recursive": true}}'
+ tool_calls = [{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "process", "arguments": complex_args}
+ }]
+
+ print("Action: Converting to text...")
+ result = tool_calls_to_text(tool_calls)
+
+ print(f"Result: '{result}'")
+ print("Checking that complex arguments are preserved...")
+ assert complex_args in result
+
+
+# ==================================================================================================
+# Tests for tool_results_to_text
+# ==================================================================================================
+
+class TestToolResultsToText:
+ """
+ Tests for tool_results_to_text function.
+
+ This function converts tool_results to human-readable text representation.
+ Used when stripping tool content from messages (when no tools are defined).
+ """
+
+ def test_converts_single_tool_result_to_text(self):
+ """
+ What it does: Verifies conversion of a single tool result to text.
+ Purpose: Ensure basic conversion works correctly.
+ """
+ print("Setup: Single tool result...")
+ tool_results = [{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Operation completed successfully"
+ }]
+
+ print("Action: Converting to text...")
+ result = tool_results_to_text(tool_results)
+
+ print(f"Result: '{result}'")
+ print("Checking that [Tool Result] marker is present...")
+ assert "[Tool Result" in result
+ print("Checking that content is present...")
+ assert "Operation completed successfully" in result
+
+ def test_converts_multiple_tool_results_to_text(self):
+ """
+ What it does: Verifies conversion of multiple tool results to text.
+ Purpose: Ensure all tool results are converted and separated.
+ """
+ print("Setup: Multiple tool results...")
+ tool_results = [
+ {"type": "tool_result", "tool_use_id": "call_1", "content": "Result 1"},
+ {"type": "tool_result", "tool_use_id": "call_2", "content": "Result 2"}
+ ]
+
+ print("Action: Converting to text...")
+ result = tool_results_to_text(tool_results)
+
+ print(f"Result: '{result}'")
+ print("Checking that both results are present...")
+ assert "Result 1" in result
+ assert "Result 2" in result
+ assert "call_1" in result
+ assert "call_2" in result
+
+ def test_includes_tool_use_id_in_output(self):
+ """
+ What it does: Verifies that tool_use_id is included in output.
+ Purpose: Ensure traceability between tool calls and results.
+ """
+ print("Setup: Tool result with tool_use_id...")
+ tool_results = [{
+ "type": "tool_result",
+ "tool_use_id": "tooluse_xyz789abc",
+ "content": "Done"
+ }]
+
+ print("Action: Converting to text...")
+ result = tool_results_to_text(tool_results)
+
+ print(f"Result: '{result}'")
+ print("Checking that tool_use_id is present...")
+ assert "tooluse_xyz789abc" in result
+
+ def test_handles_missing_tool_use_id(self):
+ """
+ What it does: Verifies handling of tool result without tool_use_id.
+ Purpose: Ensure function doesn't crash when tool_use_id is missing.
+ """
+ print("Setup: Tool result without tool_use_id...")
+ tool_results = [{
+ "type": "tool_result",
+ "content": "Some result"
+ }]
+
+ print("Action: Converting to text...")
+ result = tool_results_to_text(tool_results)
+
+ print(f"Result: '{result}'")
+ print("Checking that content is still present...")
+ assert "Some result" in result
+ assert "[Tool Result]" in result
+
+ def test_handles_empty_content(self):
+ """
+ What it does: Verifies handling of empty content.
+ Purpose: Ensure empty content is replaced with placeholder.
+ """
+ print("Setup: Tool result with empty content...")
+ tool_results = [{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": ""
+ }]
+
+ print("Action: Converting to text...")
+ result = tool_results_to_text(tool_results)
+
+ print(f"Result: '{result}'")
+ print("Checking that placeholder is used...")
+ assert "(empty result)" in result
+
+ def test_returns_empty_string_for_empty_list(self):
+ """
+ What it does: Verifies empty list handling.
+ Purpose: Ensure empty input returns empty output.
+ """
+ print("Setup: Empty list...")
+
+ print("Action: Converting to text...")
+ result = tool_results_to_text([])
+
+ print(f"Comparing result: Expected '', Got '{result}'")
+ assert result == ""
+
+ def test_handles_multiline_content(self):
+ """
+ What it does: Verifies handling of multiline content.
+ Purpose: Ensure newlines in content are preserved.
+ """
+ print("Setup: Tool result with multiline content...")
+ multiline_content = "Line 1\nLine 2\nLine 3"
+ tool_results = [{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": multiline_content
+ }]
+
+ print("Action: Converting to text...")
+ result = tool_results_to_text(tool_results)
+
+ print(f"Result: '{result}'")
+ print("Checking that multiline content is preserved...")
+ assert "Line 1\nLine 2\nLine 3" in result
+
+ def test_handles_list_content(self):
+ """
+ What it does: Verifies handling of list content (multimodal format).
+ Purpose: Ensure list content is extracted correctly.
+ """
+ print("Setup: Tool result with list content...")
+ tool_results = [{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": [{"type": "text", "text": "Extracted text"}]
+ }]
+
+ print("Action: Converting to text...")
+ result = tool_results_to_text(tool_results)
+
+ print(f"Result: '{result}'")
+ print("Checking that text is extracted from list...")
+ assert "Extracted text" in result
+
+
+# ==================================================================================================
+# Tests for build_kiro_payload with Issue #20 Scenario
+# ==================================================================================================
+
+class TestBuildKiroPayloadIssue20:
+ """
+ Tests for build_kiro_payload function specifically for Issue #20 scenario.
+
+ Issue #20: OpenCode compaction returns 400 "Improperly formed request"
+ because it sends tool_calls/tool_results in history but WITHOUT tools definitions.
+
+ Kiro API requires tools definitions if toolUses/toolResults are present.
+ The fix converts tool content to text representation when no tools are defined.
+ """
+
+ def test_compaction_without_tools_converts_tool_content_to_text(self):
+ """
+ What it does: Simulates OpenCode compaction scenario - messages with tool content but no tools.
+ Purpose: Ensure build_kiro_payload doesn't crash and converts tool content to text.
+
+ This is THE critical test for issue #20. If this test passes but the fix is removed,
+ the actual API call would fail with 400 error.
+ """
+ print("Setup: Simulating OpenCode compaction scenario...")
+ messages = [
+ UnifiedMessage(role="user", content="Read the file test.py"),
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{
+ "id": "tooluse_abc123",
+ "type": "function",
+ "function": {"name": "read_file", "arguments": '{"path": "test.py"}'}
+ }]
+ ),
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "tooluse_abc123",
+ "content": "def hello():\n print('world')"
+ }]
+ ),
+ UnifiedMessage(role="assistant", content="I see the file contains a hello function."),
+ UnifiedMessage(role="user", content="Summarize what we did")
+ ]
+
+ print("Action: Building Kiro payload WITHOUT tools (compaction scenario)...")
+ result = build_kiro_payload(
+ messages=messages,
+ system_prompt="You are a helpful assistant.",
+ model_id="claude-sonnet-4",
+ tools=None, # NO TOOLS - this is the compaction scenario
+ conversation_id="test-conv-123",
+ profile_arn="arn:aws:codewhisperer:us-east-1:123456789:profile/test",
+ inject_thinking=False
+ )
+
+ print(f"Result payload keys: {result.payload.keys()}")
+ print("Checking that payload was built successfully...")
+ assert "conversationState" in result.payload
+ assert "currentMessage" in result.payload["conversationState"]
+
+ print("Checking that history exists...")
+ history = result.payload["conversationState"].get("history", [])
+ print(f"History length: {len(history)}")
+ assert len(history) > 0
+
+ print("Checking that NO toolUses in history (they should be converted to text)...")
+ for i, msg in enumerate(history):
+ if "assistantResponseMessage" in msg:
+ assistant_msg = msg["assistantResponseMessage"]
+ print(f"History[{i}] assistant content: '{assistant_msg.get('content', '')[:100]}...'")
+ assert "toolUses" not in assistant_msg, f"toolUses should not be in history[{i}]"
+
+ print("Checking that NO toolResults in history (they should be converted to text)...")
+ for i, msg in enumerate(history):
+ if "userInputMessage" in msg:
+ user_msg = msg["userInputMessage"]
+ context = user_msg.get("userInputMessageContext", {})
+ print(f"History[{i}] user content: '{user_msg.get('content', '')[:100]}...'")
+ assert "toolResults" not in context, f"toolResults should not be in history[{i}]"
+
+ print("Checking that tool content was converted to text (preserved context)...")
+ # Find the assistant message that had tool_calls
+ found_tool_text = False
+ for msg in history:
+ if "assistantResponseMessage" in msg:
+ content = msg["assistantResponseMessage"].get("content", "")
+ if "[Tool: read_file" in content:
+ found_tool_text = True
+ print(f"Found tool text representation: '{content[:200]}...'")
+ break
+ assert found_tool_text, "Tool calls should be converted to text representation"
+
+ def test_compaction_preserves_tool_result_content_as_text(self):
+ """
+ What it does: Verifies that tool result content is preserved as text.
+ Purpose: Ensure the actual tool output is not lost during compaction.
+ """
+ print("Setup: Message with tool result containing important data...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "IMPORTANT_DATA_12345"
+ }]
+ ),
+ UnifiedMessage(role="user", content="What was in that result?")
+ ]
+
+ print("Action: Building Kiro payload without tools...")
+ result = build_kiro_payload(
+ messages=messages,
+ system_prompt="",
+ model_id="claude-sonnet-4",
+ tools=None,
+ conversation_id="test-conv",
+ profile_arn="arn:test",
+ inject_thinking=False
+ )
+
+ print("Checking that important data is preserved...")
+ # The data could be in history OR in current message (after merging adjacent user messages)
+ payload = result.payload
+
+ found_data = False
+
+ # Check history
+ history = payload["conversationState"].get("history", [])
+ for msg in history:
+ if "userInputMessage" in msg:
+ content = msg["userInputMessage"].get("content", "")
+ if "IMPORTANT_DATA_12345" in content:
+ found_data = True
+ print(f"Found preserved data in history: '{content[:100]}...'")
+ break
+
+ # Check current message (adjacent user messages are merged)
+ if not found_data:
+ current_content = payload["conversationState"]["currentMessage"]["userInputMessage"].get("content", "")
+ if "IMPORTANT_DATA_12345" in current_content:
+ found_data = True
+ print(f"Found preserved data in current message: '{current_content[:100]}...'")
+
+ assert found_data, "Tool result content should be preserved as text"
+
+ def test_with_tools_defined_keeps_tool_structure(self):
+ """
+ What it does: Verifies that when tools ARE defined, tool structure is preserved.
+ Purpose: Ensure the fix doesn't break normal tool usage.
+ """
+ print("Setup: Messages with tool content AND tools defined...")
+ messages = [
+ UnifiedMessage(role="user", content="Call a tool"),
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "test_tool", "arguments": "{}"}
+ }]
+ ),
+ UnifiedMessage(
+ role="user",
+ content="",
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Tool executed"
+ }]
+ ),
+ UnifiedMessage(role="user", content="Continue")
+ ]
+
+ tools = [UnifiedTool(
+ name="test_tool",
+ description="A test tool",
+ input_schema={"type": "object", "properties": {}}
+ )]
+
+ print("Action: Building Kiro payload WITH tools...")
+ result = build_kiro_payload(
+ messages=messages,
+ system_prompt="",
+ model_id="claude-sonnet-4",
+ tools=tools, # TOOLS DEFINED
+ conversation_id="test-conv",
+ profile_arn="arn:test",
+ inject_thinking=False
+ )
+
+ print("Checking that tools are in payload...")
+ current_msg = result.payload["conversationState"]["currentMessage"]["userInputMessage"]
+ context = current_msg.get("userInputMessageContext", {})
+ assert "tools" in context, "Tools should be in payload when defined"
+
+ print("Checking that toolUses are preserved in history...")
+ history = result.payload["conversationState"].get("history", [])
+ found_tool_uses = False
+ for msg in history:
+ if "assistantResponseMessage" in msg:
+ if "toolUses" in msg["assistantResponseMessage"]:
+ found_tool_uses = True
+ break
+ assert found_tool_uses, "toolUses should be preserved when tools are defined"
+
+ def test_empty_tools_list_triggers_stripping(self):
+ """
+ What it does: Verifies that empty tools list (tools=[]) triggers tool content stripping.
+ Purpose: Ensure edge case of empty tools list is handled correctly.
+ """
+ print("Setup: Messages with tool content and empty tools list...")
+ messages = [
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "some_tool", "arguments": "{}"}
+ }]
+ ),
+ UnifiedMessage(role="user", content="Continue")
+ ]
+
+ print("Action: Building Kiro payload with empty tools list...")
+ result = build_kiro_payload(
+ messages=messages,
+ system_prompt="",
+ model_id="claude-sonnet-4",
+ tools=[], # EMPTY TOOLS LIST
+ conversation_id="test-conv",
+ profile_arn="arn:test",
+ inject_thinking=False
+ )
+
+ print("Checking that NO tools in payload...")
+ current_msg = result.payload["conversationState"]["currentMessage"]["userInputMessage"]
+ context = current_msg.get("userInputMessageContext", {})
+ assert "tools" not in context, "Empty tools list should result in no tools in payload"
+
+ print("Checking that tool content was converted to text...")
+ history = result.payload["conversationState"].get("history", [])
+ for msg in history:
+ if "assistantResponseMessage" in msg:
+ assert "toolUses" not in msg["assistantResponseMessage"]
+
+
+# ==================================================================================================
+# Tests for build_kiro_payload with Images (Issue #30)
+# ==================================================================================================
+
+class TestBuildKiroPayloadImages:
+ """
+ Tests for build_kiro_payload function with image content.
+
+ Issue #30: 422 Validation Error when sending image content blocks.
+ The fix adds support for image content blocks in messages.
+
+ These tests verify that images are correctly included in the Kiro payload.
+ """
+
+ def test_includes_images_in_current_message(self):
+ """
+ What it does: Verifies that images are included in the current message.
+ Purpose: Ensure images from the last user message are directly in userInputMessage (Issue #32 fix).
+
+ This is a critical test for Issue #30/#32 fix - images should be in userInputMessage, NOT in userInputMessageContext.
+ """
+ print("Setup: User message with image as current message...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="What's in this image?",
+ images=[{"media_type": "image/jpeg", "data": TEST_IMAGE_BASE64}]
+ )
+ ]
+
+ print("Action: Building Kiro payload...")
+ result = build_kiro_payload(
+ messages=messages,
+ system_prompt="You are a helpful assistant.",
+ model_id="claude-sonnet-4",
+ tools=None,
+ conversation_id="test-conv-123",
+ profile_arn="arn:aws:codewhisperer:us-east-1:123456789:profile/test",
+ inject_thinking=False
+ )
+
+ print(f"Result payload keys: {result.payload.keys()}")
+ print("Checking that payload was built successfully...")
+ assert "conversationState" in result.payload
+
+ current_msg = result.payload["conversationState"]["currentMessage"]["userInputMessage"]
+ print(f"Current message: {current_msg}")
+
+ print("Checking that images are directly in userInputMessage (Issue #32 fix)...")
+ assert "images" in current_msg
+
+ images = current_msg["images"]
+ print(f"Images: {images}")
+ assert len(images) == 1
+
+ print("Checking image format (Kiro format)...")
+ assert images[0]["format"] == "jpeg"
+ assert images[0]["source"]["bytes"] == TEST_IMAGE_BASE64
+
+ def test_includes_multiple_images_in_current_message(self):
+ """
+ What it does: Verifies that multiple images are included in the current message.
+ Purpose: Ensure all images from the last user message are directly in userInputMessage (Issue #32 fix).
+ """
+ print("Setup: User message with multiple images...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="Compare these images",
+ images=[
+ {"media_type": "image/jpeg", "data": "image1_data"},
+ {"media_type": "image/png", "data": "image2_data"},
+ {"media_type": "image/gif", "data": "image3_data"}
+ ]
+ )
+ ]
+
+ print("Action: Building Kiro payload...")
+ result = build_kiro_payload(
+ messages=messages,
+ system_prompt="",
+ model_id="claude-sonnet-4",
+ tools=None,
+ conversation_id="test-conv",
+ profile_arn="arn:test",
+ inject_thinking=False
+ )
+
+ current_msg = result.payload["conversationState"]["currentMessage"]["userInputMessage"]
+ images = current_msg["images"]
+
+ print(f"Comparing image count: Expected 3, Got {len(images)}")
+ assert len(images) == 3
+
+ print("Checking image formats...")
+ assert images[0]["format"] == "jpeg"
+ assert images[1]["format"] == "png"
+ assert images[2]["format"] == "gif"
+
+ def test_includes_images_in_history(self):
+ """
+ What it does: Verifies that images are included in history messages.
+ Purpose: Ensure images from previous user messages are directly in userInputMessage (Issue #32 fix).
+ """
+ print("Setup: Conversation with images in history...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="What's in this image?",
+ images=[{"media_type": "image/jpeg", "data": "history_image_data"}]
+ ),
+ UnifiedMessage(role="assistant", content="I see a cat in the image."),
+ UnifiedMessage(role="user", content="What color is the cat?")
+ ]
+
+ print("Action: Building Kiro payload...")
+ result = build_kiro_payload(
+ messages=messages,
+ system_prompt="",
+ model_id="claude-sonnet-4",
+ tools=None,
+ conversation_id="test-conv",
+ profile_arn="arn:test",
+ inject_thinking=False
+ )
+
+ print("Checking history...")
+ history = result.payload["conversationState"]["history"]
+ 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
+
+ images = first_msg["images"]
+ print(f"History images: {images}")
+ assert len(images) == 1
+ assert images[0]["format"] == "jpeg"
+ assert images[0]["source"]["bytes"] == "history_image_data"
+
+ def test_images_with_tools(self):
+ """
+ What it does: Verifies that images work correctly with tools.
+ Purpose: Ensure images are in userInputMessage and tools are in userInputMessageContext (Issue #32 fix).
+ """
+ print("Setup: User message with image and tools defined...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="Analyze this image and use tools if needed",
+ images=[{"media_type": "image/png", "data": "image_with_tools_data"}]
+ )
+ ]
+
+ tools = [UnifiedTool(
+ name="analyze_image",
+ description="Analyze an image",
+ input_schema={"type": "object", "properties": {}}
+ )]
+
+ print("Action: Building Kiro payload with tools...")
+ result = build_kiro_payload(
+ messages=messages,
+ system_prompt="",
+ model_id="claude-sonnet-4",
+ tools=tools,
+ conversation_id="test-conv",
+ profile_arn="arn:test",
+ inject_thinking=False
+ )
+
+ current_msg = result.payload["conversationState"]["currentMessage"]["userInputMessage"]
+ context = current_msg.get("userInputMessageContext", {})
+
+ print("Checking that images are directly in userInputMessage (Issue #32 fix)...")
+ assert "images" in current_msg
+
+ print("Checking that tools are in userInputMessageContext...")
+ assert "tools" in context
+
+ print("Checking images...")
+ assert len(current_msg["images"]) == 1
+ assert current_msg["images"][0]["format"] == "png"
+
+ print("Checking tools...")
+ assert len(context["tools"]) == 1
+ assert context["tools"][0]["toolSpecification"]["name"] == "analyze_image"
+
+ def test_images_with_tool_results(self):
+ """
+ What it does: Verifies that images work correctly with tool results.
+ Purpose: Ensure images are in userInputMessage and tool_results are in userInputMessageContext (Issue #32 fix).
+ """
+ print("Setup: User message with image and tool_results...")
+ messages = [
+ UnifiedMessage(role="user", content="Call a tool"),
+ UnifiedMessage(
+ role="assistant",
+ content="",
+ tool_calls=[{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "get_data", "arguments": "{}"}
+ }]
+ ),
+ UnifiedMessage(
+ role="user",
+ content="Here's the result and an image",
+ images=[{"media_type": "image/jpeg", "data": "image_with_result_data"}],
+ tool_results=[{
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Tool output"
+ }]
+ )
+ ]
+
+ tools = [UnifiedTool(
+ name="get_data",
+ description="Get data",
+ input_schema={"type": "object", "properties": {}}
+ )]
+
+ print("Action: Building Kiro payload...")
+ result = build_kiro_payload(
+ messages=messages,
+ system_prompt="",
+ model_id="claude-sonnet-4",
+ tools=tools,
+ conversation_id="test-conv",
+ profile_arn="arn:test",
+ inject_thinking=False
+ )
+
+ # The last user message becomes current message
+ current_msg = result.payload["conversationState"]["currentMessage"]["userInputMessage"]
+ context = current_msg.get("userInputMessageContext", {})
+
+ print("Checking that images are directly in userInputMessage (Issue #32 fix)...")
+ assert "images" in current_msg
+
+ print("Checking that toolResults are in userInputMessageContext...")
+ assert "toolResults" in context
+
+ print("Checking images...")
+ assert len(current_msg["images"]) == 1
+ assert current_msg["images"][0]["format"] == "jpeg"
+
+ print("Checking toolResults...")
+ assert len(context["toolResults"]) == 1
+
+ def test_no_images_when_none_provided(self):
+ """
+ What it does: Verifies that images key is not added when no images are provided.
+ Purpose: Ensure clean payload without unnecessary empty arrays.
+ """
+ print("Setup: User message without images...")
+ messages = [
+ UnifiedMessage(role="user", content="Hello, no images here")
+ ]
+
+ print("Action: Building Kiro payload...")
+ result = build_kiro_payload(
+ messages=messages,
+ system_prompt="",
+ model_id="claude-sonnet-4",
+ tools=None,
+ conversation_id="test-conv",
+ profile_arn="arn:test",
+ inject_thinking=False
+ )
+
+ context = result.payload["conversationState"]["currentMessage"]["userInputMessage"].get("userInputMessageContext", {})
+
+ print("Checking that images key is not present or empty...")
+ # Either no images key, or empty images array
+ if "images" in context:
+ assert context["images"] == [], "Images should be empty when none provided"
+ else:
+ print("No images key - OK")
+
+ def test_large_image_data_preserved(self):
+ """
+ What it does: Verifies that large image data is preserved without truncation.
+ Purpose: Ensure large images are not corrupted during conversion (Issue #32 fix).
+ """
+ print("Setup: User message with large image data...")
+ large_image_data = "A" * 500000 # 500KB of data
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="Analyze this large image",
+ images=[{"media_type": "image/png", "data": large_image_data}]
+ )
+ ]
+
+ print("Action: Building Kiro payload...")
+ result = build_kiro_payload(
+ messages=messages,
+ system_prompt="",
+ model_id="claude-sonnet-4",
+ tools=None,
+ conversation_id="test-conv",
+ profile_arn="arn:test",
+ inject_thinking=False
+ )
+
+ current_msg = result.payload["conversationState"]["currentMessage"]["userInputMessage"]
+ images = current_msg["images"]
+
+ print(f"Checking image data length: Expected 500000, Got {len(images[0]['source']['bytes'])}")
+ assert len(images[0]["source"]["bytes"]) == 500000
+ assert images[0]["source"]["bytes"] == large_image_data
+
+ def test_images_with_thinking_injection(self):
+ """
+ What it does: Verifies that images work correctly with thinking injection.
+ Purpose: Ensure images are preserved in userInputMessage when fake reasoning is enabled (Issue #32 fix).
+ """
+ print("Setup: User message with image and thinking injection...")
+ messages = [
+ UnifiedMessage(
+ role="user",
+ content="What's in this image?",
+ images=[{"media_type": "image/jpeg", "data": "thinking_test_image"}]
+ )
+ ]
+
+ 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",
+ inject_thinking=True
+ )
+
+ current_msg = result.payload["conversationState"]["currentMessage"]["userInputMessage"]
+
+ print("Checking that images are directly in userInputMessage (Issue #32 fix)...")
+ assert "images" in current_msg
+ assert len(current_msg["images"]) == 1
+ assert current_msg["images"][0]["source"]["bytes"] == "thinking_test_image"
+
+ print("Checking that thinking tags were injected in content...")
+ content = current_msg["content"]
+ assert "" in content
+
+
+# ==================================================================================================
+# Tests for validate_tool_names (Issue #41 fix)
+# ==================================================================================================
+
+class TestValidateToolNames:
+ """
+ Tests for validate_tool_names function.
+
+ This function validates tool names against Kiro API 64-character limit.
+ Issue #41: 400 Improperly formed request with long tool names from MCP servers.
+ """
+
+ def test_accepts_short_tool_names(self):
+ """
+ What it does: Verifies that short tool names are accepted.
+ Purpose: Ensure normal tool names pass validation.
+ """
+ print("Setup: Tool with short name...")
+ tools = [UnifiedTool(name="get_weather", description="Get weather")]
+
+ print("Action: Validating tool names...")
+ try:
+ from kiro.converters_core import validate_tool_names
+ validate_tool_names(tools)
+ print("Validation passed - OK")
+ except ValueError as e:
+ print(f"ERROR: Validation failed: {e}")
+ raise AssertionError("Short tool names should be accepted")
+
+ def test_accepts_exactly_64_character_name(self):
+ """
+ What it does: Verifies that exactly 64-character names are accepted (boundary).
+ Purpose: Ensure boundary case is handled correctly.
+ """
+ print("Setup: Tool with exactly 64-character name...")
+ name_64 = "a" * 64
+ tools = [UnifiedTool(name=name_64, description="Test")]
+
+ print(f"Tool name length: {len(name_64)}")
+ print("Action: Validating tool names...")
+ try:
+ from kiro.converters_core import validate_tool_names
+ validate_tool_names(tools)
+ print("Validation passed - OK")
+ except ValueError as e:
+ print(f"ERROR: Validation failed: {e}")
+ raise AssertionError("64-character names should be accepted")
+
+ def test_rejects_65_character_name(self):
+ """
+ What it does: Verifies that 65-character names are rejected.
+ Purpose: Ensure names exceeding limit are caught.
+ """
+ print("Setup: Tool with 65-character name...")
+ name_65 = "a" * 65
+ tools = [UnifiedTool(name=name_65, description="Test")]
+
+ print(f"Tool name length: {len(name_65)}")
+ print("Action: Validating tool names (should raise ValueError)...")
+ try:
+ from kiro.converters_core import validate_tool_names
+ validate_tool_names(tools)
+ print("ERROR: Validation passed but should have failed")
+ raise AssertionError("65-character names should be rejected")
+ except ValueError as e:
+ print(f"Validation correctly rejected: {str(e)[:100]}...")
+ assert "exceed Kiro API limit" in str(e)
+ assert name_65 in str(e)
+
+ def test_rejects_very_long_tool_names(self):
+ """
+ What it does: Verifies that very long tool names are rejected.
+ Purpose: Ensure the validation works for extreme cases.
+ """
+ print("Setup: Tool with 100-character name...")
+ name_100 = "mcp__GitHub__" + "a" * 87
+ tools = [UnifiedTool(name=name_100, description="Test")]
+
+ print(f"Tool name length: {len(name_100)}")
+ print("Action: Validating tool names (should raise ValueError)...")
+ try:
+ from kiro.converters_core import validate_tool_names
+ validate_tool_names(tools)
+ raise AssertionError("Very long names should be rejected")
+ except ValueError as e:
+ print(f"Validation correctly rejected: {str(e)[:100]}...")
+ assert "exceed Kiro API limit" in str(e)
+ assert "100 characters" in str(e)
+
+ def test_rejects_multiple_long_names(self):
+ """
+ What it does: Verifies that all long names are listed in error message.
+ Purpose: Ensure user sees all problematic tools at once.
+ """
+ print("Setup: Multiple tools with long names...")
+ tools = [
+ UnifiedTool(name="a" * 65, description="Test 1"),
+ UnifiedTool(name="short", description="Test 2"),
+ UnifiedTool(name="b" * 70, description="Test 3")
+ ]
+
+ print("Action: Validating tool names (should raise ValueError)...")
+ try:
+ from kiro.converters_core import validate_tool_names
+ validate_tool_names(tools)
+ raise AssertionError("Should reject multiple long names")
+ except ValueError as e:
+ error_msg = str(e)
+ print(f"Error message: {error_msg[:200]}...")
+
+ print("Checking that both long names are listed...")
+ assert "65 characters" in error_msg
+ assert "70 characters" in error_msg
+
+ def test_handles_none_tools(self):
+ """
+ What it does: Verifies that None tools list is handled gracefully.
+ Purpose: Ensure function doesn't crash on None input.
+ """
+ print("Setup: None tools...")
+
+ print("Action: Validating None...")
+ try:
+ from kiro.converters_core import validate_tool_names
+ validate_tool_names(None)
+ print("Validation passed - OK")
+ except Exception as e:
+ print(f"ERROR: Unexpected exception: {e}")
+ raise AssertionError("None should be handled gracefully")
+
+ def test_handles_empty_tools_list(self):
+ """
+ What it does: Verifies that empty tools list is handled gracefully.
+ Purpose: Ensure function doesn't crash on empty list.
+ """
+ print("Setup: Empty tools list...")
+
+ print("Action: Validating empty list...")
+ try:
+ from kiro.converters_core import validate_tool_names
+ validate_tool_names([])
+ print("Validation passed - OK")
+ except Exception as e:
+ print(f"ERROR: Unexpected exception: {e}")
+ raise AssertionError("Empty list should be handled gracefully")
+
+ def test_error_message_includes_solution(self):
+ """
+ What it does: Verifies that error message includes solution guidance.
+ Purpose: Ensure user knows how to fix the problem.
+ """
+ print("Setup: Tool with long name...")
+ tools = [UnifiedTool(name="mcp__GitHub__" + "a" * 60, description="Test")]
+
+ print("Action: Validating tool names (should raise ValueError)...")
+ try:
+ from kiro.converters_core import validate_tool_names
+ validate_tool_names(tools)
+ raise AssertionError("Should reject long name")
+ except ValueError as e:
+ error_msg = str(e)
+ print(f"Error message: {error_msg[:300]}...")
+
+ print("Checking that error message includes solution...")
+ assert "Solution:" in error_msg
+ assert "64 characters" in error_msg
+ assert "Example:" in error_msg
+
+ def test_real_world_mcp_tool_names(self):
+ """
+ What it does: Verifies rejection of real MCP tool names from Issue #41.
+ Purpose: Ensure the fix works for actual problematic tool names.
+ """
+ print("Setup: Real MCP tool names from Issue #41...")
+ problematic_names = [
+ "mcp__GitHub__check_if_a_person_is_followed_by_the_authenticated_user",
+ "mcp__GitHub__check_if_a_repository_is_starred_by_the_authenticated_user",
+ "mcp__GitHub__remove_interaction_restrictions_from_your_public_repositories",
+ ]
+
+ tools = [UnifiedTool(name=name, description="Test") for name in problematic_names]
+
+ print("Action: Validating real MCP tool names (should raise ValueError)...")
+ try:
+ from kiro.converters_core import validate_tool_names
+ validate_tool_names(tools)
+ raise AssertionError("Should reject real MCP tool names")
+ except ValueError as e:
+ error_msg = str(e)
+ print(f"Error message length: {len(error_msg)} chars")
+ print(f"Error message: {error_msg[:400]}...")
+
+ print("Checking that all problematic names are listed...")
+ for name in problematic_names:
+ assert name in error_msg, f"Tool name '{name}' should be in error message"
+
+ print("Checking that character counts are shown...")
+ assert "68 characters" in error_msg
+ assert "71 characters" in error_msg
+ assert "74 characters" in error_msg
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_converters_openai.py b/kiro-gateway/tests/unit/test_converters_openai.py
new file mode 100644
index 0000000000000000000000000000000000000000..2a3050651b19458d38dafdfc707bdbc944388f6b
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_converters_openai.py
@@ -0,0 +1,1362 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for converters_openai module.
+
+Tests for OpenAI-specific conversion logic:
+- Converting OpenAI messages to unified format
+- Converting OpenAI tools to unified format
+- Building Kiro payload from OpenAI requests
+"""
+
+import pytest
+from unittest.mock import patch
+
+from kiro.converters_openai import (
+ build_kiro_payload,
+ convert_openai_messages_to_unified,
+ convert_openai_tools_to_unified,
+)
+from kiro.models_openai import ChatMessage, ChatCompletionRequest, Tool, ToolFunction
+
+
+# ==================================================================================================
+# Tests for convert_openai_messages_to_unified
+# ==================================================================================================
+
+class TestConvertOpenAIMessagesToUnified:
+ """Tests for convert_openai_messages_to_unified function."""
+
+ def test_extracts_system_prompt(self):
+ """
+ What it does: Verifies extraction of system prompt from messages.
+ Purpose: Ensure system messages are extracted separately.
+ """
+ print("Setup: Messages with system prompt...")
+ messages = [
+ ChatMessage(role="system", content="You are helpful"),
+ ChatMessage(role="user", content="Hello")
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"System prompt: '{system_prompt}'")
+ print(f"Unified messages: {len(unified)}")
+ assert system_prompt == "You are helpful"
+ assert len(unified) == 1
+ assert unified[0].role == "user"
+
+ def test_combines_multiple_system_messages(self):
+ """
+ What it does: Verifies combining of multiple system messages.
+ Purpose: Ensure all system messages are concatenated.
+ """
+ print("Setup: Multiple system messages...")
+ messages = [
+ ChatMessage(role="system", content="You are helpful."),
+ ChatMessage(role="system", content="Be concise."),
+ ChatMessage(role="user", content="Hello")
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"System prompt: '{system_prompt}'")
+ assert "You are helpful." in system_prompt
+ assert "Be concise." in system_prompt
+ assert len(unified) == 1
+
+ def test_converts_tool_message_to_user_with_tool_results(self):
+ """
+ What it does: Verifies conversion of tool message to user message with tool_results.
+ Purpose: Ensure role="tool" is converted correctly.
+ """
+ print("Setup: Tool message...")
+ messages = [
+ ChatMessage(role="tool", content="Tool result text", tool_call_id="call_123")
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Unified messages: {unified}")
+ assert len(unified) == 1
+ assert unified[0].role == "user"
+ assert unified[0].tool_results is not None
+ assert len(unified[0].tool_results) == 1
+ assert unified[0].tool_results[0]["tool_use_id"] == "call_123"
+
+ def test_converts_multiple_tool_messages(self):
+ """
+ What it does: Verifies conversion of multiple consecutive tool messages.
+ Purpose: Ensure all tool results are collected into one user message.
+ """
+ print("Setup: Multiple tool messages...")
+ messages = [
+ ChatMessage(role="tool", content="Result 1", tool_call_id="call_1"),
+ ChatMessage(role="tool", content="Result 2", tool_call_id="call_2"),
+ ChatMessage(role="tool", content="Result 3", tool_call_id="call_3")
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Unified messages: {unified}")
+ assert len(unified) == 1
+ assert unified[0].role == "user"
+ assert len(unified[0].tool_results) == 3
+
+ def test_extracts_tool_calls_from_assistant(self):
+ """
+ What it does: Verifies extraction of tool_calls from assistant message.
+ Purpose: Ensure tool_calls are preserved in unified format.
+ """
+ print("Setup: Assistant message with tool_calls...")
+ messages = [
+ ChatMessage(
+ role="assistant",
+ content="I'll call a tool",
+ tool_calls=[{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "get_weather", "arguments": '{"location": "Moscow"}'}
+ }]
+ )
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Unified messages: {unified}")
+ assert len(unified) == 1
+ assert unified[0].role == "assistant"
+ assert unified[0].tool_calls is not None
+ assert len(unified[0].tool_calls) == 1
+ assert unified[0].tool_calls[0]["id"] == "call_123"
+
+ def test_handles_empty_tool_call_id(self):
+ """
+ What it does: Verifies handling of None tool_call_id.
+ Purpose: Ensure None is replaced with empty string.
+ """
+ print("Setup: Tool message with None tool_call_id...")
+ messages = [
+ ChatMessage(role="tool", content="Result", tool_call_id=None)
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Unified messages: {unified}")
+ assert unified[0].tool_results[0]["tool_use_id"] == ""
+
+ def test_handles_empty_tool_content(self):
+ """
+ What it does: Verifies handling of empty tool content.
+ Purpose: Ensure empty content is replaced with "(empty result)".
+ """
+ print("Setup: Tool message with empty content...")
+ messages = [
+ ChatMessage(role="tool", content="", tool_call_id="call_1")
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Unified messages: {unified}")
+ assert unified[0].tool_results[0]["content"] == "(empty result)"
+
+ def test_tool_messages_followed_by_user_message(self):
+ """
+ What it does: Verifies tool messages followed by user message.
+ Purpose: Ensure tool results are in separate message from user content.
+ """
+ print("Setup: Tool messages + user message...")
+ messages = [
+ ChatMessage(role="tool", content="Result 1", tool_call_id="call_1"),
+ ChatMessage(role="user", content="Continue please")
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Unified messages: {unified}")
+ # Tool results should be in first message, user content in second
+ assert len(unified) == 2
+ assert unified[0].role == "user"
+ assert unified[0].tool_results is not None
+ assert unified[1].role == "user"
+ assert unified[1].content == "Continue please"
+
+ # ==================================================================================
+ # Image extraction tests (Issue #30 fix)
+ # ==================================================================================
+
+ def test_extracts_images_from_user_message(self):
+ """
+ What it does: Verifies that images are extracted from user messages.
+ Purpose: Ensure OpenAI image_url content blocks are converted to unified format.
+
+ This test verifies the fix for Issue #30 - 422 Validation Error for image content.
+ """
+ print("Setup: User message with image_url content block...")
+ # Base64 1x1 pixel JPEG
+ test_image_base64 = "/9j/4AAQSkZJRgABAQEASABIAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/wAALCAABAAEBAREA/8QAFAABAAAAAAAAAAAAAAAAAAAACf/EABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEAAD8AVN//2Q=="
+
+ messages = [
+ ChatMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "What's in this image?"},
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": f"data:image/jpeg;base64,{test_image_base64}"
+ }
+ }
+ ]
+ )
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Result: {unified}")
+ print(f"Images: {unified[0].images}")
+
+ assert len(unified) == 1
+ assert unified[0].role == "user"
+ assert unified[0].content == "What's in this image?"
+
+ print("Checking images field...")
+ assert unified[0].images is not None, "images field should not be None"
+ assert len(unified[0].images) == 1, f"Expected 1 image, got {len(unified[0].images)}"
+
+ image = unified[0].images[0]
+ print(f"Comparing image: Expected media_type='image/jpeg', Got '{image.get('media_type')}'")
+ assert image["media_type"] == "image/jpeg"
+
+ print(f"Comparing image data: Expected {test_image_base64[:20]}..., Got {image.get('data', '')[:20]}...")
+ assert image["data"] == test_image_base64
+
+ def test_images_only_extracted_from_user_role(self):
+ """
+ What it does: Verifies that images are only extracted from user messages.
+ Purpose: Ensure assistant messages don't have images extracted.
+ """
+ print("Setup: Conversation with image in user message only...")
+ test_image_base64 = "/9j/4AAQSkZJRgABAQEASABIAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/wAALCAABAAEBAREA/8QAFAABAAAAAAAAAAAAAAAAAAAACf/EABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEAAD8AVN//2Q=="
+
+ messages = [
+ ChatMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "Describe this image"},
+ {
+ "type": "image_url",
+ "image_url": {"url": f"data:image/png;base64,{test_image_base64}"}
+ }
+ ]
+ ),
+ ChatMessage(
+ role="assistant",
+ content="I can see a small image."
+ )
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Result: {unified}")
+
+ print("Checking user message has images...")
+ assert unified[0].images is not None
+ assert len(unified[0].images) == 1
+
+ print("Checking assistant message has no images...")
+ assert unified[1].images is None, "Assistant messages should not have images extracted"
+
+ def test_extracts_multiple_images_from_user_message(self):
+ """
+ What it does: Verifies extraction of multiple images from a single user message.
+ Purpose: Ensure all images in a message are extracted.
+ """
+ print("Setup: User message with multiple images...")
+ test_image_base64 = "/9j/4AAQSkZJRgABAQEASABIAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/wAALCAABAAEBAREA/8QAFAABAAAAAAAAAAAAAAAAAAAACf/EABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEAAD8AVN//2Q=="
+
+ messages = [
+ ChatMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "Compare these images"},
+ {
+ "type": "image_url",
+ "image_url": {"url": f"data:image/jpeg;base64,{test_image_base64}"}
+ },
+ {
+ "type": "image_url",
+ "image_url": {"url": f"data:image/png;base64,{test_image_base64}"}
+ },
+ {
+ "type": "image_url",
+ "image_url": {"url": f"data:image/webp;base64,{test_image_base64}"}
+ }
+ ]
+ )
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Result images count: {len(unified[0].images) if unified[0].images else 0}")
+
+ assert unified[0].images is not None
+ assert len(unified[0].images) == 3, f"Expected 3 images, got {len(unified[0].images)}"
+
+ print("Checking image media types...")
+ media_types = [img["media_type"] for img in unified[0].images]
+ print(f"Media types: {media_types}")
+ assert "image/jpeg" in media_types
+ assert "image/png" in media_types
+ assert "image/webp" in media_types
+
+ def test_counts_images_in_debug_log(self, caplog):
+ """
+ What it does: Verifies that image count is logged in debug message.
+ Purpose: Ensure logging includes image statistics for debugging.
+ """
+ import logging
+
+ print("Setup: User message with images for logging test...")
+ test_image_base64 = "/9j/4AAQSkZJRgABAQEASABIAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/wAALCAABAAEBAREA/8QAFAABAAAAAAAAAAAAAAAAAAAACf/EABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEAAD8AVN//2Q=="
+
+ messages = [
+ ChatMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "Analyze this"},
+ {
+ "type": "image_url",
+ "image_url": {"url": f"data:image/jpeg;base64,{test_image_base64}"}
+ },
+ {
+ "type": "image_url",
+ "image_url": {"url": f"data:image/png;base64,{test_image_base64}"}
+ }
+ ]
+ )
+ ]
+
+ print("Action: Converting messages with logging enabled...")
+ with caplog.at_level(logging.DEBUG):
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Log records: {[r.message for r in caplog.records]}")
+
+ # Check that images were extracted
+ assert unified[0].images is not None
+ assert len(unified[0].images) == 2
+
+ # Note: loguru doesn't integrate with caplog by default
+ # The function logs "Converted X OpenAI messages: Y tool_calls, Z tool_results, W images"
+ # We verify the images are extracted correctly, which proves the counting works
+ print("Images extracted successfully - logging verification complete")
+
+
+# ==================================================================================================
+# Tests for convert_openai_tools_to_unified
+# ==================================================================================================
+
+class TestConvertOpenAIToolsToUnified:
+ """Tests for convert_openai_tools_to_unified function."""
+
+ def test_returns_none_for_none(self):
+ """
+ What it does: Verifies handling of None.
+ Purpose: Ensure None returns None.
+ """
+ print("Setup: None tools...")
+
+ print("Action: Converting tools...")
+ result = convert_openai_tools_to_unified(None)
+
+ print(f"Result: {result}")
+ assert result is None
+
+ def test_returns_none_for_empty_list(self):
+ """
+ What it does: Verifies handling of empty list.
+ Purpose: Ensure empty list returns None.
+ """
+ print("Setup: Empty tools list...")
+
+ print("Action: Converting tools...")
+ result = convert_openai_tools_to_unified([])
+
+ print(f"Result: {result}")
+ assert result is None
+
+ def test_converts_function_tool(self):
+ """
+ What it does: Verifies conversion of function tool.
+ Purpose: Ensure Tool is converted to UnifiedTool.
+ """
+ print("Setup: Function tool...")
+ tools = [Tool(
+ type="function",
+ function=ToolFunction(
+ name="get_weather",
+ description="Get weather for a location",
+ parameters={"type": "object", "properties": {"location": {"type": "string"}}}
+ )
+ )]
+
+ print("Action: Converting tools...")
+ result = convert_openai_tools_to_unified(tools)
+
+ print(f"Result: {result}")
+ assert result is not None
+ assert len(result) == 1
+ assert result[0].name == "get_weather"
+ assert result[0].description == "Get weather for a location"
+ assert result[0].input_schema == {"type": "object", "properties": {"location": {"type": "string"}}}
+
+ def test_skips_non_function_tools(self):
+ """
+ What it does: Verifies skipping of non-function tools.
+ Purpose: Ensure only function tools are converted.
+ """
+ print("Setup: Non-function tool...")
+ tools = [Tool(
+ type="other_type",
+ function=ToolFunction(name="test", description="Test", parameters={})
+ )]
+
+ print("Action: Converting tools...")
+ result = convert_openai_tools_to_unified(tools)
+
+ print(f"Result: {result}")
+ assert result is None # No function tools, so None
+
+ def test_converts_multiple_tools(self):
+ """
+ What it does: Verifies conversion of multiple tools.
+ Purpose: Ensure all function tools are converted.
+ """
+ print("Setup: Multiple tools...")
+ tools = [
+ Tool(type="function", function=ToolFunction(name="tool1", description="Tool 1", parameters={})),
+ Tool(type="function", function=ToolFunction(name="tool2", description="Tool 2", parameters={})),
+ Tool(type="function", function=ToolFunction(name="tool3", description="Tool 3", parameters={}))
+ ]
+
+ print("Action: Converting tools...")
+ result = convert_openai_tools_to_unified(tools)
+
+ print(f"Result: {result}")
+ assert len(result) == 3
+ assert result[0].name == "tool1"
+ assert result[1].name == "tool2"
+ assert result[2].name == "tool3"
+
+ # ==================================================================================
+ # Cursor IDE Flat Tool Format Tests (PR #49)
+ # ==================================================================================
+
+ def test_converts_flat_format_tool(self):
+ """
+ What it does: Verifies conversion of flat format tool (Cursor-style).
+ Purpose: Ensure Cursor IDE flat format is supported.
+
+ Cursor IDE sends tools in flat format:
+ {"type": "function", "name": "...", "description": "...", "input_schema": {...}}
+ instead of standard OpenAI nested format.
+ """
+ print("Setup: Flat format tool (Cursor-style)...")
+ tools = [Tool(
+ type="function",
+ name="cursor_tool",
+ description="A tool from Cursor IDE",
+ input_schema={"type": "object", "properties": {"param": {"type": "string"}}}
+ )]
+
+ print("Action: Converting tools...")
+ result = convert_openai_tools_to_unified(tools)
+
+ print(f"Result: {result}")
+ print(f"Comparing count: Expected 1, Got {len(result) if result else 0}")
+ assert result is not None
+ assert len(result) == 1
+
+ print(f"Comparing name: Expected 'cursor_tool', Got '{result[0].name}'")
+ assert result[0].name == "cursor_tool"
+
+ print(f"Comparing description: Expected 'A tool from Cursor IDE', Got '{result[0].description}'")
+ assert result[0].description == "A tool from Cursor IDE"
+
+ print(f"Comparing input_schema: Got {result[0].input_schema}")
+ assert result[0].input_schema == {"type": "object", "properties": {"param": {"type": "string"}}}
+
+ def test_converts_mixed_format_tools(self):
+ """
+ What it does: Verifies conversion of mixed format tools.
+ Purpose: Ensure both standard and flat format can coexist in same request.
+
+ This simulates a scenario where some tools are in standard OpenAI format
+ and some are in Cursor flat format (though unlikely in practice).
+ """
+ print("Setup: Mixed format tools...")
+ tools = [
+ # Standard OpenAI format
+ Tool(
+ type="function",
+ function=ToolFunction(
+ name="standard_tool",
+ description="Standard format",
+ parameters={"type": "object"}
+ )
+ ),
+ # Cursor flat format
+ Tool(
+ type="function",
+ name="flat_tool",
+ description="Flat format",
+ input_schema={"type": "object"}
+ )
+ ]
+
+ print("Action: Converting tools...")
+ result = convert_openai_tools_to_unified(tools)
+
+ print(f"Result: {result}")
+ print(f"Comparing count: Expected 2, Got {len(result)}")
+ assert len(result) == 2
+
+ print("Checking standard format tool...")
+ assert result[0].name == "standard_tool"
+ assert result[0].description == "Standard format"
+
+ print("Checking flat format tool...")
+ assert result[1].name == "flat_tool"
+ assert result[1].description == "Flat format"
+
+ def test_standard_format_takes_priority(self):
+ """
+ What it does: Verifies that standard format takes priority over flat format.
+ Purpose: Ensure function field is used when both formats are present (edge case).
+
+ This is an edge case where a tool has BOTH function and name fields.
+ The standard format (function) should take priority.
+ """
+ print("Setup: Tool with BOTH formats (edge case)...")
+ tools = [Tool(
+ type="function",
+ # Standard format
+ function=ToolFunction(
+ name="standard_name",
+ description="Standard description",
+ parameters={"type": "object", "properties": {"a": {"type": "string"}}}
+ ),
+ # Flat format (should be ignored)
+ name="flat_name",
+ description="Flat description",
+ input_schema={"type": "object", "properties": {"b": {"type": "string"}}}
+ )]
+
+ print("Action: Converting tools...")
+ result = convert_openai_tools_to_unified(tools)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+
+ print("Checking that standard format was used (not flat)...")
+ print(f"Comparing name: Expected 'standard_name', Got '{result[0].name}'")
+ assert result[0].name == "standard_name"
+
+ print(f"Comparing description: Expected 'Standard description', Got '{result[0].description}'")
+ assert result[0].description == "Standard description"
+
+ print(f"Comparing input_schema: Got {result[0].input_schema}")
+ assert result[0].input_schema == {"type": "object", "properties": {"a": {"type": "string"}}}
+
+ def test_skips_invalid_tools(self):
+ """
+ What it does: Verifies that tools without function OR name are skipped.
+ Purpose: Ensure invalid tools don't crash the conversion.
+
+ This tests the error handling when a tool has neither function nor name field.
+ """
+ print("Setup: Invalid tool (no function, no name)...")
+ tools = [
+ # Valid tool
+ Tool(
+ type="function",
+ function=ToolFunction(name="valid_tool", description="Valid")
+ ),
+ # Invalid tool (neither function nor name)
+ Tool(type="function"),
+ # Another valid tool
+ Tool(
+ type="function",
+ name="another_valid",
+ description="Also valid",
+ input_schema={}
+ )
+ ]
+
+ print("Action: Converting tools...")
+ result = convert_openai_tools_to_unified(tools)
+
+ print(f"Result: {result}")
+ print(f"Comparing count: Expected 2 (invalid skipped), Got {len(result)}")
+ assert len(result) == 2
+
+ print("Checking that only valid tools were converted...")
+ assert result[0].name == "valid_tool"
+ assert result[1].name == "another_valid"
+
+ def test_backward_compat_standard_openai_tools(self):
+ """
+ What it does: Verifies that standard OpenAI format is not broken.
+ Purpose: Regression test for existing clients (non-Cursor).
+
+ This is a critical backward compatibility test. After adding support for
+ Cursor's flat format, we must ensure standard OpenAI format still works.
+ """
+ print("Setup: Standard OpenAI tools (regression test)...")
+ tools = [
+ Tool(
+ type="function",
+ function=ToolFunction(
+ name="get_weather",
+ description="Get weather for a location",
+ parameters={
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "City name"}
+ },
+ "required": ["location"]
+ }
+ )
+ )
+ ]
+
+ print("Action: Converting tools...")
+ result = convert_openai_tools_to_unified(tools)
+
+ print(f"Result: {result}")
+ assert result is not None
+ assert len(result) == 1
+
+ print(f"Comparing name: Expected 'get_weather', Got '{result[0].name}'")
+ assert result[0].name == "get_weather"
+
+ print(f"Comparing description: Expected 'Get weather for a location', Got '{result[0].description}'")
+ assert result[0].description == "Get weather for a location"
+
+ print(f"Comparing input_schema: Got {result[0].input_schema}")
+ assert result[0].input_schema["required"] == ["location"]
+ assert result[0].input_schema["properties"]["location"]["type"] == "string"
+
+
+# ==================================================================================================
+# Tests for build_kiro_payload
+# ==================================================================================================
+
+class TestBuildKiroPayload:
+ """Tests for build_kiro_payload function."""
+
+ def test_builds_simple_payload(self):
+ """
+ What it does: Verifies building of simple payload.
+ Purpose: Ensure basic request is converted correctly.
+ """
+ print("Setup: Simple request...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="user", content="Hello")]
+ )
+
+ print("Action: Building payload...")
+ result = build_kiro_payload(request, "conv-123", "arn:aws:test")
+
+ print(f"Result: {result}")
+ assert "conversationState" in result
+ assert result["conversationState"]["conversationId"] == "conv-123"
+ assert "currentMessage" in result["conversationState"]
+ assert result["profileArn"] == "arn:aws:test"
+
+ 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.
+ """
+ print("Setup: Request with system prompt...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[
+ ChatMessage(role="system", content="You are helpful"),
+ ChatMessage(role="user", content="Hello")
+ ]
+ )
+
+ print("Action: Building payload...")
+ 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
+
+ def test_builds_history_for_multi_turn(self):
+ """
+ What it does: Verifies building history for multi-turn.
+ Purpose: Ensure previous messages go into history.
+ """
+ print("Setup: Multi-turn request...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[
+ ChatMessage(role="user", content="Hello"),
+ ChatMessage(role="assistant", content="Hi"),
+ ChatMessage(role="user", content="How are you?")
+ ]
+ )
+
+ print("Action: Building payload...")
+ result = build_kiro_payload(request, "conv-123", "")
+
+ print(f"Result: {result}")
+ assert "history" in result["conversationState"]
+ assert len(result["conversationState"]["history"]) == 2
+
+ def test_handles_assistant_as_last_message(self):
+ """
+ What it does: Verifies handling of assistant as last message.
+ Purpose: Ensure "Continue" message is created.
+ """
+ print("Setup: Request with assistant at the end...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[
+ ChatMessage(role="user", content="Hello"),
+ ChatMessage(role="assistant", content="Hi there")
+ ]
+ )
+
+ print("Action: Building payload...")
+ result = build_kiro_payload(request, "conv-123", "")
+
+ print(f"Result: {result}")
+ current_content = result["conversationState"]["currentMessage"]["userInputMessage"]["content"]
+ assert current_content == "Continue"
+
+ def test_raises_for_empty_messages(self):
+ """
+ What it does: Verifies exception raising for empty messages.
+ Purpose: Ensure empty request raises ValueError.
+ """
+ print("Setup: Request with only system message...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="system", content="You are helpful")]
+ )
+
+ print("Action: Attempting to build payload...")
+ with pytest.raises(ValueError) as exc_info:
+ build_kiro_payload(request, "conv-123", "")
+
+ print(f"Exception: {exc_info.value}")
+ assert "No messages to send" in str(exc_info.value)
+
+ def test_uses_continue_for_empty_content(self):
+ """
+ What it does: Verifies using "Continue" for empty content.
+ Purpose: Ensure empty message is replaced with "Continue".
+ """
+ print("Setup: Request with empty content...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="user", content="")]
+ )
+
+ print("Action: Building payload (with fake reasoning disabled)...")
+ with patch('kiro.converters_core.FAKE_REASONING_ENABLED', False):
+ result = build_kiro_payload(request, "conv-123", "")
+
+ print(f"Result: {result}")
+ current_content = result["conversationState"]["currentMessage"]["userInputMessage"]["content"]
+ assert current_content == "Continue"
+
+ def test_normalizes_model_id_correctly(self):
+ """
+ What it does: Verifies normalization of external model ID to Kiro format.
+ Purpose: Ensure model name normalization is applied (dashes→dots, strip dates).
+
+ Note: The new Dynamic Model Resolution System normalizes model names
+ (e.g., claude-sonnet-4-5 → claude-sonnet-4.5) instead of mapping to
+ internal IDs. Kiro API accepts the normalized format directly.
+ """
+ print("Setup: Request with external model ID...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="user", content="Hello")]
+ )
+
+ print("Action: Building payload...")
+ result = build_kiro_payload(request, "conv-123", "")
+
+ print(f"Result: {result}")
+ model_id = result["conversationState"]["currentMessage"]["userInputMessage"]["modelId"]
+ # claude-sonnet-4-5 should normalize to claude-sonnet-4.5 (dashes→dots)
+ print(f"Comparing model_id: Expected 'claude-sonnet-4.5', Got '{model_id}'")
+ assert model_id == "claude-sonnet-4.5"
+
+ def test_includes_tools_in_context(self):
+ """
+ What it does: Verifies including tools in userInputMessageContext.
+ Purpose: Ensure tools are converted and included.
+ """
+ print("Setup: Request with tools...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="user", content="Hello")],
+ tools=[Tool(
+ type="function",
+ function=ToolFunction(
+ name="get_weather",
+ description="Get weather",
+ parameters={"type": "object", "properties": {}}
+ )
+ )]
+ )
+
+ print("Action: Building payload...")
+ result = build_kiro_payload(request, "conv-123", "")
+
+ print(f"Result: {result}")
+ context = result["conversationState"]["currentMessage"]["userInputMessage"]["userInputMessageContext"]
+ assert "tools" in context
+ assert len(context["tools"]) == 1
+ assert context["tools"][0]["toolSpecification"]["name"] == "get_weather"
+
+ def test_injects_thinking_tags_even_when_tool_results_present(self):
+ """
+ What it does: Verifies thinking tags ARE injected even when toolResults are present.
+ Purpose: Extended thinking should work in all scenarios including tool use flows.
+ """
+ print("Setup: Request where last message is a tool result...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[
+ ChatMessage(role="user", content="Run a command"),
+ ChatMessage(
+ role="assistant",
+ content="I'll run the command",
+ tool_calls=[{
+ "id": "tool_1",
+ "type": "function",
+ "function": {"name": "bash", "arguments": "{}"}
+ }]
+ ),
+ ChatMessage(role="tool", content="Command output here", tool_call_id="tool_1"),
+ ],
+ # Tools must be defined for tool_results to be preserved
+ tools=[
+ Tool(
+ type="function",
+ function=ToolFunction(
+ name="bash",
+ description="Run a bash command",
+ parameters={"type": "object", "properties": {}}
+ )
+ )
+ ]
+ )
+
+ 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", "")
+
+ current_msg = result["conversationState"]["currentMessage"]["userInputMessage"]
+ content = current_msg["content"]
+ context = current_msg.get("userInputMessageContext", {})
+
+ print(f"Content: {repr(content[:100] if len(content) > 100 else content)}")
+ print(f"Has toolResults: {'toolResults' in context}")
+
+ assert "toolResults" in context, "toolResults should be present"
+ assert "enabled" in content, "thinking tags SHOULD be injected even with toolResults"
+ assert "4000" in content, "max_thinking_length should be present"
+
+ def test_injects_thinking_tags_when_no_tool_results(self):
+ """
+ What it does: Verifies thinking tags ARE injected for normal user messages.
+ Purpose: Ensure fix for issue #20 doesn't break normal thinking tag injection.
+ """
+ print("Setup: Normal user message without tool results...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="user", content="Hello")]
+ )
+
+ 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", "")
+
+ current_msg = result["conversationState"]["currentMessage"]["userInputMessage"]
+ content = current_msg["content"]
+ context = current_msg.get("userInputMessageContext", {})
+
+ print(f"Content starts with thinking tags: {'' in content}")
+ print(f"Has toolResults: {'toolResults' in context}")
+
+ assert "toolResults" not in context, "toolResults should NOT be present"
+ assert "" in content, "thinking tags SHOULD be injected for normal messages"
+ assert "Hello" in content, "Original content should be preserved"
+
+
+# ==================================================================================================
+# Tests for tool message handling
+# ==================================================================================================
+
+class TestToolMessageHandling:
+ """Tests for OpenAI tool message (role="tool") handling."""
+
+ def test_converts_multiple_tool_messages_to_single_user_message(self):
+ """
+ What it does: Verifies merging of multiple tool messages into single user message.
+ Purpose: Ensure multiple tool results are merged into one user message.
+ """
+ print("Setup: Multiple consecutive tool messages...")
+ messages = [
+ ChatMessage(role="tool", content="Result 1", tool_call_id="call_1"),
+ ChatMessage(role="tool", content="Result 2", tool_call_id="call_2"),
+ ChatMessage(role="tool", content="Result 3", tool_call_id="call_3")
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Result: {unified}")
+ print(f"Comparing length: Expected 1, Got {len(unified)}")
+ assert len(unified) == 1
+ assert unified[0].role == "user"
+
+ print("Checking content contains all tool_results...")
+ assert unified[0].tool_results is not None
+ assert len(unified[0].tool_results) == 3
+
+ tool_use_ids = [item["tool_use_id"] for item in unified[0].tool_results]
+ assert "call_1" in tool_use_ids
+ assert "call_2" in tool_use_ids
+ assert "call_3" in tool_use_ids
+
+ def test_assistant_tool_user_sequence(self):
+ """
+ What it does: Verifies assistant -> tool -> user sequence.
+ Purpose: Ensure tool message is correctly inserted between assistant and user.
+ """
+ print("Setup: assistant -> tool -> user...")
+ messages = [
+ ChatMessage(role="assistant", content="I'll call a tool"),
+ ChatMessage(role="tool", content="Tool output", tool_call_id="call_abc"),
+ ChatMessage(role="user", content="Thanks!")
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Result: {unified}")
+ # assistant stays, tool becomes user with tool_results, then user
+ assert len(unified) == 3
+ assert unified[0].role == "assistant"
+ assert unified[1].role == "user"
+ assert unified[1].tool_results is not None
+ assert unified[2].role == "user"
+
+ def test_tool_message_with_empty_content(self):
+ """
+ What it does: Verifies tool message with empty content.
+ Purpose: Ensure empty result is replaced with "(empty result)".
+ """
+ print("Setup: Tool message with empty content...")
+ messages = [
+ ChatMessage(role="tool", content="", tool_call_id="call_empty")
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Result: {unified}")
+ assert len(unified) == 1
+ assert unified[0].tool_results[0]["content"] == "(empty result)"
+
+ def test_tool_message_with_none_tool_call_id(self):
+ """
+ What it does: Verifies tool message without tool_call_id.
+ Purpose: Ensure missing tool_call_id is replaced with empty string.
+ """
+ print("Setup: Tool message without tool_call_id...")
+ messages = [
+ ChatMessage(role="tool", content="Result", tool_call_id=None)
+ ]
+
+ print("Action: Converting messages...")
+ system_prompt, unified = convert_openai_messages_to_unified(messages)
+
+ print(f"Result: {unified}")
+ assert len(unified) == 1
+ assert unified[0].tool_results[0]["tool_use_id"] == ""
+
+
+# ==================================================================================================
+# Tests for tool description handling
+# ==================================================================================================
+
+class TestToolDescriptionHandling:
+ """Tests for handling empty/whitespace tool descriptions."""
+
+ def test_empty_description_replaced_with_placeholder(self):
+ """
+ What it does: Verifies replacement of empty description with placeholder.
+ Purpose: Ensure empty description is replaced with "Tool: {name}".
+
+ This is a critical test for a Cline bug where tool focus_chain had
+ empty description "", which caused a 400 error from Kiro API.
+ """
+ print("Setup: Tool with empty description...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="user", content="Hello")],
+ tools=[Tool(
+ type="function",
+ function=ToolFunction(
+ name="focus_chain",
+ description="",
+ parameters={"type": "object", "properties": {}}
+ )
+ )]
+ )
+
+ print("Action: Building payload...")
+ result = build_kiro_payload(request, "conv-123", "")
+
+ print(f"Result: {result}")
+ print("Checking that description is replaced with placeholder...")
+ context = result["conversationState"]["currentMessage"]["userInputMessage"]["userInputMessageContext"]
+ tool_spec = context["tools"][0]["toolSpecification"]
+ assert tool_spec["description"] == "Tool: focus_chain"
+
+ def test_whitespace_only_description_replaced_with_placeholder(self):
+ """
+ What it does: Verifies replacement of whitespace-only description with placeholder.
+ Purpose: Ensure description with only whitespace is replaced.
+ """
+ print("Setup: Tool with whitespace-only description...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="user", content="Hello")],
+ tools=[Tool(
+ type="function",
+ function=ToolFunction(
+ name="whitespace_tool",
+ description=" ",
+ parameters={}
+ )
+ )]
+ )
+
+ print("Action: Building payload...")
+ result = build_kiro_payload(request, "conv-123", "")
+
+ print(f"Result: {result}")
+ print("Checking that description is replaced with placeholder...")
+ context = result["conversationState"]["currentMessage"]["userInputMessage"]["userInputMessageContext"]
+ tool_spec = context["tools"][0]["toolSpecification"]
+ assert tool_spec["description"] == "Tool: whitespace_tool"
+
+ def test_none_description_replaced_with_placeholder(self):
+ """
+ What it does: Verifies replacement of None description with placeholder.
+ Purpose: Ensure None description is replaced with "Tool: {name}".
+ """
+ print("Setup: Tool with None description...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="user", content="Hello")],
+ tools=[Tool(
+ type="function",
+ function=ToolFunction(
+ name="none_desc_tool",
+ description=None,
+ parameters={}
+ )
+ )]
+ )
+
+ print("Action: Building payload...")
+ result = build_kiro_payload(request, "conv-123", "")
+
+ print(f"Result: {result}")
+ print("Checking that description is replaced with placeholder...")
+ context = result["conversationState"]["currentMessage"]["userInputMessage"]["userInputMessageContext"]
+ tool_spec = context["tools"][0]["toolSpecification"]
+ assert tool_spec["description"] == "Tool: none_desc_tool"
+
+ def test_non_empty_description_preserved(self):
+ """
+ What it does: Verifies preservation of non-empty description.
+ Purpose: Ensure normal description is not changed.
+ """
+ print("Setup: Tool with normal description...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="user", content="Hello")],
+ tools=[Tool(
+ type="function",
+ function=ToolFunction(
+ name="get_weather",
+ description="Get weather for a location",
+ parameters={}
+ )
+ )]
+ )
+
+ print("Action: Building payload...")
+ result = build_kiro_payload(request, "conv-123", "")
+
+ print(f"Result: {result}")
+ print("Checking that description is preserved...")
+ context = result["conversationState"]["currentMessage"]["userInputMessage"]["userInputMessageContext"]
+ tool_spec = context["tools"][0]["toolSpecification"]
+ assert tool_spec["description"] == "Get weather for a location"
+
+ def test_sanitizes_tool_parameters(self):
+ """
+ What it does: Verifies sanitization of parameters from problematic fields.
+ Purpose: Ensure sanitize_json_schema is applied to parameters.
+ """
+ print("Setup: Tool with problematic parameters...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="user", content="Hello")],
+ tools=[Tool(
+ type="function",
+ function=ToolFunction(
+ name="test_tool",
+ description="Test tool",
+ parameters={
+ "type": "object",
+ "properties": {},
+ "required": [],
+ "additionalProperties": False
+ }
+ )
+ )]
+ )
+
+ print("Action: Building payload...")
+ result = build_kiro_payload(request, "conv-123", "")
+
+ print(f"Result: {result}")
+ print("Checking that parameters are sanitized...")
+ context = result["conversationState"]["currentMessage"]["userInputMessage"]["userInputMessageContext"]
+ input_schema = context["tools"][0]["toolSpecification"]["inputSchema"]["json"]
+ assert "required" not in input_schema
+ assert "additionalProperties" not in input_schema
+
+ def test_mixed_tools_with_empty_and_normal_descriptions(self):
+ """
+ What it does: Verifies handling of mixed tools list.
+ Purpose: Ensure empty descriptions are replaced while normal ones are preserved.
+
+ This is a real scenario from Cline where most tools have
+ normal descriptions, but focus_chain has an empty one.
+ """
+ print("Setup: Mixed list of tools...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="user", content="Hello")],
+ tools=[
+ Tool(
+ type="function",
+ function=ToolFunction(
+ name="read_file",
+ description="Read contents of a file",
+ parameters={}
+ )
+ ),
+ Tool(
+ type="function",
+ function=ToolFunction(
+ name="focus_chain",
+ description="",
+ parameters={}
+ )
+ ),
+ Tool(
+ type="function",
+ function=ToolFunction(
+ name="write_file",
+ description="Write content to a file",
+ parameters={}
+ )
+ )
+ ]
+ )
+
+ print("Action: Building payload...")
+ result = build_kiro_payload(request, "conv-123", "")
+
+ print(f"Result: {result}")
+ print("Checking descriptions...")
+ context = result["conversationState"]["currentMessage"]["userInputMessage"]["userInputMessageContext"]
+ tools = context["tools"]
+ assert tools[0]["toolSpecification"]["description"] == "Read contents of a file"
+ assert tools[1]["toolSpecification"]["description"] == "Tool: focus_chain"
+ assert tools[2]["toolSpecification"]["description"] == "Write content to a file"
+
+
+# ==================================================================================================
+# Integration tests for full flow
+# ==================================================================================================
+
+class TestBuildKiroPayloadToolCallsIntegration:
+ """
+ Integration tests for build_kiro_payload with tool_calls.
+ Tests full flow from OpenAI format to Kiro format.
+ """
+
+ def test_multiple_assistant_tool_calls_with_results(self):
+ """
+ What it does: Verifies full scenario with multiple assistant tool_calls and their results.
+ Purpose: Ensure all toolUses and toolResults are correctly linked in Kiro payload.
+
+ This is an integration test for a Codex CLI bug where multiple assistant
+ messages with tool_calls were sent in a row, followed by tool results.
+ """
+ print("Setup: Full scenario with two tool_calls and their results...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[
+ ChatMessage(role="user", content="Run two commands"),
+ # First assistant with tool_call
+ ChatMessage(
+ role="assistant",
+ content=None,
+ tool_calls=[{
+ "id": "tooluse_first",
+ "type": "function",
+ "function": {"name": "shell", "arguments": '{"command": ["ls"]}'}
+ }]
+ ),
+ # Second assistant with tool_call (consecutive!)
+ ChatMessage(
+ role="assistant",
+ content=None,
+ tool_calls=[{
+ "id": "tooluse_second",
+ "type": "function",
+ "function": {"name": "shell", "arguments": '{"command": ["pwd"]}'}
+ }]
+ ),
+ # Results of both tool_calls
+ ChatMessage(role="tool", content="file1.txt\nfile2.txt", tool_call_id="tooluse_first"),
+ ChatMessage(role="tool", content="/home/user", tool_call_id="tooluse_second")
+ ],
+ # Tools must be defined for tool_results to be preserved
+ tools=[
+ Tool(
+ type="function",
+ function=ToolFunction(
+ name="shell",
+ description="Run a shell command",
+ parameters={"type": "object", "properties": {"command": {"type": "array"}}}
+ )
+ )
+ ]
+ )
+
+ print("Action: Building Kiro payload...")
+ result = build_kiro_payload(request, "conv-123", "arn:aws:test")
+
+ print(f"Result: {result}")
+
+ # Check history
+ history = result["conversationState"].get("history", [])
+ print(f"History: {history}")
+
+ # Should have userInputMessage and assistantResponseMessage in history
+ assert len(history) >= 2, f"Expected at least 2 elements in history, got {len(history)}"
+
+ # Find assistantResponseMessage
+ assistant_msgs = [h for h in history if "assistantResponseMessage" in h]
+ print(f"Assistant messages in history: {assistant_msgs}")
+ assert len(assistant_msgs) >= 1, "Should have at least one assistantResponseMessage"
+
+ # 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)}")
+ assert len(tool_uses) == 2, f"Should have 2 toolUses, got {len(tool_uses)}"
+
+ tool_use_ids = [tu["toolUseId"] for tu in tool_uses]
+ print(f"ToolUse IDs: {tool_use_ids}")
+ assert "tooluse_first" in tool_use_ids
+ assert "tooluse_second" in tool_use_ids
+
+ # Check currentMessage contains toolResults
+ current_msg = result["conversationState"]["currentMessage"]["userInputMessage"]
+ context = current_msg.get("userInputMessageContext", {})
+ tool_results = context.get("toolResults", [])
+ print(f"ToolResults in currentMessage: {tool_results}")
+ print(f"Comparing toolResults count: Expected 2, Got {len(tool_results)}")
+ assert len(tool_results) == 2, f"Should have 2 toolResults, got {len(tool_results)}"
+
+ # Note: tool_results in Kiro payload use camelCase (toolUseId)
+ tool_result_ids = [tr["toolUseId"] for tr in tool_results]
+ print(f"ToolResult IDs: {tool_result_ids}")
+ assert "tooluse_first" in tool_result_ids
+ assert "tooluse_second" in tool_result_ids
+
+ def test_long_tool_description_added_to_system_prompt(self):
+ """
+ What it does: Verifies integration of long tool descriptions into payload.
+ Purpose: Ensure long descriptions are added to system prompt in payload.
+ """
+ print("Setup: Request with tool with long description...")
+ long_desc = "X" * 15000
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[
+ ChatMessage(role="system", content="You are helpful"),
+ ChatMessage(role="user", content="Hello")
+ ],
+ tools=[Tool(
+ type="function",
+ function=ToolFunction(
+ name="long_tool",
+ description=long_desc,
+ parameters={}
+ )
+ )]
+ )
+
+ print("Action: Building payload...")
+ with patch('kiro.converters_core.TOOL_DESCRIPTION_MAX_LENGTH', 10000):
+ 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
+
+ print("Checking that tool in context has reference description...")
+ tools_context = result["conversationState"]["currentMessage"]["userInputMessage"]["userInputMessageContext"]["tools"]
+ assert "[Full documentation in system prompt" in tools_context[0]["toolSpecification"]["description"]
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_debug_logger.py b/kiro-gateway/tests/unit/test_debug_logger.py
new file mode 100644
index 0000000000000000000000000000000000000000..a1138a6183115db0f40afe4dfee7cd575015b783
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_debug_logger.py
@@ -0,0 +1,690 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit-тесты для DebugLogger.
+Проверяет логику буферизации и записи debug логов в разных режимах.
+"""
+
+import json
+import pytest
+from pathlib import Path
+from unittest.mock import patch, MagicMock
+
+
+class TestDebugLoggerModeOff:
+ """Тесты для режима DEBUG_MODE=off."""
+
+ def test_prepare_new_request_does_nothing(self, tmp_path):
+ """
+ Что он делает: Проверяет, что prepare_new_request ничего не делает в режиме off.
+ Цель: Убедиться, что в режиме off директория не создаётся.
+ """
+ print("Настройка: Режим off...")
+ with patch('kiro.debug_logger.DEBUG_MODE', 'off'):
+ with patch('kiro.debug_logger.DEBUG_DIR', str(tmp_path / "debug_logs")):
+ # Пересоздаём экземпляр с новыми настройками
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = tmp_path / "debug_logs"
+
+ print("Действие: Вызов prepare_new_request...")
+ logger.prepare_new_request()
+
+ print(f"Проверяем, что директория не создана...")
+ assert not (tmp_path / "debug_logs").exists()
+
+ def test_log_request_body_does_nothing(self, tmp_path):
+ """
+ Что он делает: Проверяет, что log_request_body ничего не делает в режиме off.
+ Цель: Убедиться, что данные не записываются.
+ """
+ print("Настройка: Режим off...")
+ with patch('kiro.debug_logger.DEBUG_MODE', 'off'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = tmp_path / "debug_logs"
+
+ print("Действие: Вызов log_request_body...")
+ logger.log_request_body(b'{"test": "data"}')
+
+ print(f"Проверяем, что файл не создан...")
+ assert not (tmp_path / "debug_logs" / "request_body.json").exists()
+
+
+class TestDebugLoggerModeAll:
+ """Тесты для режима DEBUG_MODE=all."""
+
+ def test_prepare_new_request_clears_directory(self, tmp_path):
+ """
+ Что он делает: Проверяет, что prepare_new_request очищает директорию в режиме all.
+ Цель: Убедиться, что старые логи удаляются.
+ """
+ print("Настройка: Режим all, создаём старый файл...")
+ debug_dir = tmp_path / "debug_logs"
+ debug_dir.mkdir()
+ old_file = debug_dir / "old_file.txt"
+ old_file.write_text("old content")
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ print("Действие: Вызов prepare_new_request...")
+ logger.prepare_new_request()
+
+ print(f"Проверяем, что старый файл удалён...")
+ assert not old_file.exists()
+ print(f"Проверяем, что директория существует...")
+ assert debug_dir.exists()
+
+ def test_log_request_body_writes_immediately(self, tmp_path):
+ """
+ Что он делает: Проверяет, что log_request_body пишет сразу в файл в режиме all.
+ Цель: Убедиться, что данные записываются немедленно.
+ """
+ print("Настройка: Режим all...")
+ debug_dir = tmp_path / "debug_logs"
+ debug_dir.mkdir()
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ print("Действие: Вызов log_request_body...")
+ test_data = b'{"model": "test", "messages": []}'
+ logger.log_request_body(test_data)
+
+ print(f"Проверяем, что файл создан...")
+ file_path = debug_dir / "request_body.json"
+ assert file_path.exists()
+
+ print(f"Проверяем содержимое файла...")
+ content = json.loads(file_path.read_text())
+ assert content["model"] == "test"
+
+ def test_log_kiro_request_body_writes_immediately(self, tmp_path):
+ """
+ Что он делает: Проверяет, что log_kiro_request_body пишет сразу в файл в режиме all.
+ Цель: Убедиться, что Kiro payload записывается немедленно.
+ """
+ print("Настройка: Режим all...")
+ debug_dir = tmp_path / "debug_logs"
+ debug_dir.mkdir()
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ print("Действие: Вызов log_kiro_request_body...")
+ test_data = b'{"conversationState": {}}'
+ logger.log_kiro_request_body(test_data)
+
+ print(f"Проверяем, что файл создан...")
+ file_path = debug_dir / "kiro_request_body.json"
+ assert file_path.exists()
+
+ def test_log_raw_chunk_appends_to_file(self, tmp_path):
+ """
+ Что он делает: Проверяет, что log_raw_chunk дописывает в файл в режиме all.
+ Цель: Убедиться, что чанки накапливаются.
+ """
+ print("Настройка: Режим all...")
+ debug_dir = tmp_path / "debug_logs"
+ debug_dir.mkdir()
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ print("Действие: Вызов log_raw_chunk дважды...")
+ logger.log_raw_chunk(b'chunk1')
+ logger.log_raw_chunk(b'chunk2')
+
+ print(f"Проверяем содержимое файла...")
+ file_path = debug_dir / "response_stream_raw.txt"
+ content = file_path.read_bytes()
+ assert content == b'chunk1chunk2'
+
+
+class TestDebugLoggerModeErrors:
+ """Тесты для режима DEBUG_MODE=errors."""
+
+ def test_log_request_body_buffers_data(self, tmp_path):
+ """
+ Что он делает: Проверяет, что log_request_body буферизует данные в режиме errors.
+ Цель: Убедиться, что данные не записываются сразу.
+ """
+ print("Настройка: Режим errors...")
+ debug_dir = tmp_path / "debug_logs"
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'errors'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ print("Действие: Вызов log_request_body...")
+ test_data = b'{"test": "buffered"}'
+ logger.log_request_body(test_data)
+
+ print(f"Проверяем, что файл НЕ создан...")
+ assert not debug_dir.exists()
+
+ print(f"Проверяем, что данные в буфере...")
+ assert logger._request_body_buffer == test_data
+
+ def test_flush_on_error_writes_buffers(self, tmp_path):
+ """
+ Что он делает: Проверяет, что flush_on_error записывает буферы в файлы.
+ Цель: Убедиться, что при ошибке данные сохраняются.
+ """
+ print("Настройка: Режим errors, заполняем буферы...")
+ debug_dir = tmp_path / "debug_logs"
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'errors'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ # Заполняем буферы
+ logger.log_request_body(b'{"request": "body"}')
+ logger.log_kiro_request_body(b'{"kiro": "request"}')
+ logger.log_raw_chunk(b'raw_chunk')
+ logger.log_modified_chunk(b'modified_chunk')
+
+ print("Действие: Вызов flush_on_error...")
+ logger.flush_on_error(400, "Bad Request")
+
+ print(f"Проверяем, что все файлы созданы...")
+ assert (debug_dir / "request_body.json").exists()
+ assert (debug_dir / "kiro_request_body.json").exists()
+ assert (debug_dir / "response_stream_raw.txt").exists()
+ assert (debug_dir / "response_stream_modified.txt").exists()
+ assert (debug_dir / "error_info.json").exists()
+
+ print(f"Проверяем error_info.json...")
+ error_info = json.loads((debug_dir / "error_info.json").read_text())
+ assert error_info["status_code"] == 400
+ assert error_info["error_message"] == "Bad Request"
+
+ def test_flush_on_error_clears_buffers(self, tmp_path):
+ """
+ Что он делает: Проверяет, что flush_on_error очищает буферы после записи.
+ Цель: Убедиться, что буферы не накапливаются между запросами.
+ """
+ print("Настройка: Режим errors...")
+ debug_dir = tmp_path / "debug_logs"
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'errors'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ logger.log_request_body(b'{"test": "data"}')
+
+ print("Действие: Вызов flush_on_error...")
+ logger.flush_on_error(500, "Error")
+
+ print(f"Проверяем, что буферы очищены...")
+ assert logger._request_body_buffer is None
+ assert logger._kiro_request_body_buffer is None
+ assert len(logger._raw_chunks_buffer) == 0
+ assert len(logger._modified_chunks_buffer) == 0
+
+ def test_discard_buffers_clears_without_writing(self, tmp_path):
+ """
+ Что он делает: Проверяет, что discard_buffers очищает буферы без записи.
+ Цель: Убедиться, что успешные запросы не оставляют логов.
+ """
+ print("Настройка: Режим errors, заполняем буферы...")
+ debug_dir = tmp_path / "debug_logs"
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'errors'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ logger.log_request_body(b'{"test": "data"}')
+ logger.log_raw_chunk(b'chunk')
+
+ print("Действие: Вызов discard_buffers...")
+ logger.discard_buffers()
+
+ print(f"Проверяем, что директория НЕ создана...")
+ assert not debug_dir.exists()
+
+ print(f"Проверяем, что буферы очищены...")
+ assert logger._request_body_buffer is None
+ assert len(logger._raw_chunks_buffer) == 0
+
+ def test_flush_on_error_writes_error_info_in_mode_all(self, tmp_path):
+ """
+ Что он делает: Проверяет, что flush_on_error записывает error_info.json в режиме all.
+ Цель: Убедиться, что информация об ошибке сохраняется в обоих режимах.
+ """
+ print("Настройка: Режим all...")
+ debug_dir = tmp_path / "debug_logs"
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ print("Действие: Вызов flush_on_error...")
+ logger.flush_on_error(400, "Bad Request")
+
+ print(f"Проверяем, что error_info.json создан...")
+ assert (debug_dir / "error_info.json").exists()
+
+ print(f"Проверяем содержимое error_info.json...")
+ error_info = json.loads((debug_dir / "error_info.json").read_text())
+ assert error_info["status_code"] == 400
+ assert error_info["error_message"] == "Bad Request"
+
+
+class TestDebugLoggerLogErrorInfo:
+ """Тесты для метода log_error_info()."""
+
+ def test_log_error_info_writes_in_mode_all(self, tmp_path):
+ """
+ Что он делает: Проверяет, что log_error_info записывает файл в режиме all.
+ Цель: Убедиться, что error_info.json создаётся при ошибках.
+ """
+ print("Настройка: Режим all...")
+ debug_dir = tmp_path / "debug_logs"
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ print("Действие: Вызов log_error_info...")
+ logger.log_error_info(500, "Internal Server Error")
+
+ print(f"Проверяем, что error_info.json создан...")
+ error_file = debug_dir / "error_info.json"
+ assert error_file.exists()
+
+ print(f"Проверяем содержимое...")
+ error_info = json.loads(error_file.read_text())
+ assert error_info["status_code"] == 500
+ assert error_info["error_message"] == "Internal Server Error"
+
+ def test_log_error_info_writes_in_mode_errors(self, tmp_path):
+ """
+ Что он делает: Проверяет, что log_error_info записывает файл в режиме errors.
+ Цель: Убедиться, что метод работает в обоих режимах.
+ """
+ print("Настройка: Режим errors...")
+ debug_dir = tmp_path / "debug_logs"
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'errors'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ print("Действие: Вызов log_error_info...")
+ logger.log_error_info(404, "Not Found")
+
+ print(f"Проверяем, что error_info.json создан...")
+ error_file = debug_dir / "error_info.json"
+ assert error_file.exists()
+
+ def test_log_error_info_does_nothing_in_mode_off(self, tmp_path):
+ """
+ Что он делает: Проверяет, что log_error_info ничего не делает в режиме off.
+ Цель: Убедиться, что в режиме off файлы не создаются.
+ """
+ print("Настройка: Режим off...")
+ debug_dir = tmp_path / "debug_logs"
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'off'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ print("Действие: Вызов log_error_info...")
+ logger.log_error_info(500, "Error")
+
+ print(f"Проверяем, что директория НЕ создана...")
+ assert not debug_dir.exists()
+
+
+class TestDebugLoggerHelperMethods:
+ """Тесты для вспомогательных методов DebugLogger."""
+
+ def test_is_enabled_returns_true_for_errors(self):
+ """
+ Что он делает: Проверяет _is_enabled() для режима errors.
+ Цель: Убедиться, что режим errors считается включённым.
+ """
+ print("Настройка: Режим errors...")
+ with patch('kiro.debug_logger.DEBUG_MODE', 'errors'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+
+ print(f"Проверяем _is_enabled()...")
+ assert logger._is_enabled() is True
+
+ def test_is_enabled_returns_true_for_all(self):
+ """
+ Что он делает: Проверяет _is_enabled() для режима all.
+ Цель: Убедиться, что режим all считается включённым.
+ """
+ print("Настройка: Режим all...")
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+
+ print(f"Проверяем _is_enabled()...")
+ assert logger._is_enabled() is True
+
+ def test_is_enabled_returns_false_for_off(self):
+ """
+ Что он делает: Проверяет _is_enabled() для режима off.
+ Цель: Убедиться, что режим off считается выключенным.
+ """
+ print("Настройка: Режим off...")
+ with patch('kiro.debug_logger.DEBUG_MODE', 'off'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+
+ print(f"Проверяем _is_enabled()...")
+ assert logger._is_enabled() is False
+
+ def test_is_immediate_write_returns_true_for_all(self):
+ """
+ Что он делает: Проверяет _is_immediate_write() для режима all.
+ Цель: Убедиться, что режим all пишет сразу.
+ """
+ print("Настройка: Режим all...")
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+
+ print(f"Проверяем _is_immediate_write()...")
+ assert logger._is_immediate_write() is True
+
+ def test_is_immediate_write_returns_false_for_errors(self):
+ """
+ Что он делает: Проверяет _is_immediate_write() для режима errors.
+ Цель: Убедиться, что режим errors буферизует.
+ """
+ print("Настройка: Режим errors...")
+ with patch('kiro.debug_logger.DEBUG_MODE', 'errors'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+
+ print(f"Проверяем _is_immediate_write()...")
+ assert logger._is_immediate_write() is False
+
+
+class TestDebugLoggerJsonHandling:
+ """Тесты для обработки JSON в DebugLogger."""
+
+ def test_log_request_body_formats_json_pretty(self, tmp_path):
+ """
+ Что он делает: Проверяет, что JSON форматируется красиво.
+ Цель: Убедиться, что JSON читаем в файле.
+ """
+ print("Настройка: Режим all...")
+ debug_dir = tmp_path / "debug_logs"
+ debug_dir.mkdir()
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ print("Действие: Вызов log_request_body с JSON...")
+ logger.log_request_body(b'{"key":"value"}')
+
+ print(f"Проверяем форматирование...")
+ content = (debug_dir / "request_body.json").read_text()
+ # Должен быть отформатирован с отступами
+ assert " " in content or "\n" in content
+
+ def test_log_request_body_handles_invalid_json(self, tmp_path):
+ """
+ Что он делает: Проверяет обработку невалидного JSON.
+ Цель: Убедиться, что невалидный JSON записывается как есть.
+ """
+ print("Настройка: Режим all...")
+ debug_dir = tmp_path / "debug_logs"
+ debug_dir.mkdir()
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+ logger = DebugLogger.__new__(DebugLogger)
+ logger._initialized = False
+ logger.__init__()
+ logger.debug_dir = debug_dir
+
+ print("Действие: Вызов log_request_body с невалидным JSON...")
+ invalid_data = b'not a json {{'
+ logger.log_request_body(invalid_data)
+
+ print(f"Проверяем, что данные записаны как есть...")
+ content = (debug_dir / "request_body.json").read_bytes()
+ assert content == invalid_data
+
+
+class TestDebugLoggerAppLogsCapture:
+ """Тесты для захвата логов приложения (app_logs.txt)."""
+
+ def test_prepare_new_request_sets_up_log_capture(self, tmp_path):
+ """
+ Что он делает: Проверяет, что prepare_new_request настраивает захват логов.
+ Цель: Убедиться, что sink для логов создаётся.
+ """
+ print("Настройка: Режим all...")
+ debug_dir = tmp_path / "debug_logs"
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+ dbg_logger = DebugLogger.__new__(DebugLogger)
+ dbg_logger._initialized = False
+ dbg_logger.__init__()
+ dbg_logger.debug_dir = debug_dir
+
+ print("Действие: Вызов prepare_new_request...")
+ dbg_logger.prepare_new_request()
+
+ print(f"Проверяем, что sink создан...")
+ assert dbg_logger._loguru_sink_id is not None
+
+ # Очистка
+ dbg_logger._clear_app_logs_buffer()
+
+ def test_flush_on_error_writes_app_logs_in_mode_errors(self, tmp_path):
+ """
+ Что он делает: Проверяет, что flush_on_error записывает app_logs.txt в режиме errors.
+ Цель: Убедиться, что логи приложения сохраняются при ошибках.
+ """
+ print("Настройка: Режим errors...")
+ debug_dir = tmp_path / "debug_logs"
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'errors'):
+ from kiro.debug_logger import DebugLogger
+ from loguru import logger as loguru_logger
+
+ dbg_logger = DebugLogger.__new__(DebugLogger)
+ dbg_logger._initialized = False
+ dbg_logger.__init__()
+ dbg_logger.debug_dir = debug_dir
+
+ # Настраиваем захват логов
+ dbg_logger.prepare_new_request()
+
+ # Добавляем данные в буфер чтобы flush сработал
+ dbg_logger.log_request_body(b'{"test": "data"}')
+
+ # Пишем тестовый лог напрямую в буфер (имитация)
+ dbg_logger._app_logs_buffer.write("Test log message\n")
+
+ print("Действие: Вызов flush_on_error...")
+ dbg_logger.flush_on_error(500, "Test Error")
+
+ print(f"Проверяем, что app_logs.txt создан...")
+ app_logs_file = debug_dir / "app_logs.txt"
+ assert app_logs_file.exists()
+
+ print(f"Проверяем содержимое...")
+ content = app_logs_file.read_text()
+ assert "Test log message" in content
+
+ def test_discard_buffers_saves_logs_in_mode_all(self, tmp_path):
+ """
+ Что он делает: Проверяет, что discard_buffers сохраняет логи в режиме all.
+ Цель: Убедиться, что даже успешные запросы сохраняют логи в режиме all.
+ """
+ print("Настройка: Режим all...")
+ debug_dir = tmp_path / "debug_logs"
+ debug_dir.mkdir()
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+
+ dbg_logger = DebugLogger.__new__(DebugLogger)
+ dbg_logger._initialized = False
+ dbg_logger.__init__()
+ dbg_logger.debug_dir = debug_dir
+
+ # Настраиваем захват логов
+ dbg_logger.prepare_new_request()
+
+ # Пишем тестовый лог напрямую в буфер
+ dbg_logger._app_logs_buffer.write("Success log message\n")
+
+ print("Действие: Вызов discard_buffers...")
+ dbg_logger.discard_buffers()
+
+ print(f"Проверяем, что app_logs.txt создан...")
+ app_logs_file = debug_dir / "app_logs.txt"
+ assert app_logs_file.exists()
+
+ print(f"Проверяем содержимое...")
+ content = app_logs_file.read_text()
+ assert "Success log message" in content
+
+ def test_discard_buffers_does_not_save_logs_in_mode_errors(self, tmp_path):
+ """
+ Что он делает: Проверяет, что discard_buffers НЕ сохраняет логи в режиме errors.
+ Цель: Убедиться, что успешные запросы не оставляют логов в режиме errors.
+ """
+ print("Настройка: Режим errors...")
+ debug_dir = tmp_path / "debug_logs"
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'errors'):
+ from kiro.debug_logger import DebugLogger
+
+ dbg_logger = DebugLogger.__new__(DebugLogger)
+ dbg_logger._initialized = False
+ dbg_logger.__init__()
+ dbg_logger.debug_dir = debug_dir
+
+ # Настраиваем захват логов
+ dbg_logger.prepare_new_request()
+
+ # Пишем тестовый лог напрямую в буфер
+ dbg_logger._app_logs_buffer.write("Should not be saved\n")
+
+ print("Действие: Вызов discard_buffers...")
+ dbg_logger.discard_buffers()
+
+ print(f"Проверяем, что директория НЕ создана...")
+ assert not debug_dir.exists()
+
+ def test_clear_app_logs_buffer_removes_sink(self, tmp_path):
+ """
+ Что он делает: Проверяет, что _clear_app_logs_buffer удаляет sink.
+ Цель: Убедиться, что sink корректно удаляется.
+ """
+ print("Настройка: Режим all...")
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+
+ dbg_logger = DebugLogger.__new__(DebugLogger)
+ dbg_logger._initialized = False
+ dbg_logger.__init__()
+ dbg_logger.debug_dir = tmp_path / "debug_logs"
+
+ # Настраиваем захват логов
+ dbg_logger.prepare_new_request()
+ sink_id = dbg_logger._loguru_sink_id
+ assert sink_id is not None
+
+ print("Действие: Вызов _clear_app_logs_buffer...")
+ dbg_logger._clear_app_logs_buffer()
+
+ print(f"Проверяем, что sink_id сброшен...")
+ assert dbg_logger._loguru_sink_id is None
+
+ def test_app_logs_not_saved_when_empty(self, tmp_path):
+ """
+ Что он делает: Проверяет, что пустые логи не создают файл.
+ Цель: Убедиться, что app_logs.txt не создаётся если логов нет.
+ """
+ print("Настройка: Режим all...")
+ debug_dir = tmp_path / "debug_logs"
+ debug_dir.mkdir()
+
+ with patch('kiro.debug_logger.DEBUG_MODE', 'all'):
+ from kiro.debug_logger import DebugLogger
+
+ dbg_logger = DebugLogger.__new__(DebugLogger)
+ dbg_logger._initialized = False
+ dbg_logger.__init__()
+ dbg_logger.debug_dir = debug_dir
+
+ # НЕ пишем ничего в буфер
+
+ print("Действие: Вызов _write_app_logs_to_file...")
+ dbg_logger._write_app_logs_to_file()
+
+ print(f"Проверяем, что app_logs.txt НЕ создан...")
+ app_logs_file = debug_dir / "app_logs.txt"
+ assert not app_logs_file.exists()
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_debug_middleware.py b/kiro-gateway/tests/unit/test_debug_middleware.py
new file mode 100644
index 0000000000000000000000000000000000000000..6f80491cd334cc075fd92bb3b0d7d2a758618dae
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_debug_middleware.py
@@ -0,0 +1,383 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for DebugLoggerMiddleware.
+Tests debug logging initialization at the middleware level.
+"""
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+from starlette.requests import Request
+from starlette.responses import Response
+
+
+class TestDebugLoggerMiddlewareEndpointFiltering:
+ """Tests for endpoint filtering in middleware."""
+
+ @pytest.mark.asyncio
+ async def test_skips_health_endpoint(self):
+ """
+ What it does: Verifies that middleware skips /health endpoint.
+ Purpose: Ensure health checks are not logged.
+ """
+ print("Setup: Creating mock request for /health...")
+
+ with patch('kiro.debug_middleware.DEBUG_MODE', 'all'):
+ from kiro.debug_middleware import DebugLoggerMiddleware
+
+ middleware = DebugLoggerMiddleware(app=MagicMock())
+
+ # Mock request
+ mock_request = MagicMock(spec=Request)
+ mock_request.url.path = "/health"
+
+ # Mock call_next
+ mock_response = MagicMock(spec=Response)
+ mock_call_next = AsyncMock(return_value=mock_response)
+
+ # Mock debug_logger at the source module
+ with patch('kiro.debug_logger.debug_logger') as mock_logger:
+ print("Action: Calling dispatch for /health...")
+ response = await middleware.dispatch(mock_request, mock_call_next)
+
+ print("Verifying prepare_new_request was NOT called...")
+ mock_logger.prepare_new_request.assert_not_called()
+
+ print("Verifying call_next was called...")
+ mock_call_next.assert_called_once_with(mock_request)
+
+ print(f"Comparing response: Expected {mock_response}, Got {response}")
+ assert response == mock_response
+
+ @pytest.mark.asyncio
+ async def test_skips_docs_endpoint(self):
+ """
+ What it does: Verifies that middleware skips /docs endpoint.
+ Purpose: Ensure documentation is not logged.
+ """
+ print("Setup: Creating mock request for /docs...")
+
+ with patch('kiro.debug_middleware.DEBUG_MODE', 'all'):
+ from kiro.debug_middleware import DebugLoggerMiddleware
+
+ middleware = DebugLoggerMiddleware(app=MagicMock())
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.url.path = "/docs"
+
+ mock_response = MagicMock(spec=Response)
+ mock_call_next = AsyncMock(return_value=mock_response)
+
+ with patch('kiro.debug_logger.debug_logger') as mock_logger:
+ print("Action: Calling dispatch for /docs...")
+ response = await middleware.dispatch(mock_request, mock_call_next)
+
+ print("Verifying prepare_new_request was NOT called...")
+ mock_logger.prepare_new_request.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_skips_root_endpoint(self):
+ """
+ What it does: Verifies that middleware skips / endpoint.
+ Purpose: Ensure root endpoint is not logged.
+ """
+ print("Setup: Creating mock request for /...")
+
+ with patch('kiro.debug_middleware.DEBUG_MODE', 'all'):
+ from kiro.debug_middleware import DebugLoggerMiddleware
+
+ middleware = DebugLoggerMiddleware(app=MagicMock())
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.url.path = "/"
+
+ mock_response = MagicMock(spec=Response)
+ mock_call_next = AsyncMock(return_value=mock_response)
+
+ with patch('kiro.debug_logger.debug_logger') as mock_logger:
+ print("Action: Calling dispatch for /...")
+ await middleware.dispatch(mock_request, mock_call_next)
+
+ print("Verifying prepare_new_request was NOT called...")
+ mock_logger.prepare_new_request.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_processes_chat_completions_endpoint(self):
+ """
+ What it does: Verifies that middleware processes /v1/chat/completions.
+ Purpose: Ensure OpenAI endpoint is logged.
+ """
+ print("Setup: Creating mock request for /v1/chat/completions...")
+
+ with patch('kiro.debug_middleware.DEBUG_MODE', 'all'):
+ from kiro.debug_middleware import DebugLoggerMiddleware
+
+ middleware = DebugLoggerMiddleware(app=MagicMock())
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.url.path = "/v1/chat/completions"
+ mock_request.body = AsyncMock(return_value=b'{"model": "test"}')
+
+ mock_response = MagicMock(spec=Response)
+ mock_call_next = AsyncMock(return_value=mock_response)
+
+ with patch('kiro.debug_logger.debug_logger') as mock_logger:
+ print("Action: Calling dispatch for /v1/chat/completions...")
+ await middleware.dispatch(mock_request, mock_call_next)
+
+ print("Verifying prepare_new_request was called...")
+ mock_logger.prepare_new_request.assert_called_once()
+
+ print("Verifying log_request_body was called...")
+ mock_logger.log_request_body.assert_called_once_with(b'{"model": "test"}')
+
+ @pytest.mark.asyncio
+ async def test_processes_messages_endpoint(self):
+ """
+ What it does: Verifies that middleware processes /v1/messages.
+ Purpose: Ensure Anthropic endpoint is logged.
+ """
+ print("Setup: Creating mock request for /v1/messages...")
+
+ with patch('kiro.debug_middleware.DEBUG_MODE', 'all'):
+ from kiro.debug_middleware import DebugLoggerMiddleware
+
+ middleware = DebugLoggerMiddleware(app=MagicMock())
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.url.path = "/v1/messages"
+ mock_request.body = AsyncMock(return_value=b'{"model": "claude"}')
+
+ mock_response = MagicMock(spec=Response)
+ mock_call_next = AsyncMock(return_value=mock_response)
+
+ with patch('kiro.debug_logger.debug_logger') as mock_logger:
+ print("Action: Calling dispatch for /v1/messages...")
+ await middleware.dispatch(mock_request, mock_call_next)
+
+ print("Verifying prepare_new_request was called...")
+ mock_logger.prepare_new_request.assert_called_once()
+
+
+class TestDebugLoggerMiddlewareModeHandling:
+ """Tests for DEBUG_MODE handling in middleware."""
+
+ @pytest.mark.asyncio
+ async def test_skips_when_debug_mode_off(self):
+ """
+ What it does: Verifies that middleware skips requests when DEBUG_MODE=off.
+ Purpose: Ensure logging is disabled in off mode.
+ """
+ print("Setup: DEBUG_MODE=off...")
+
+ with patch('kiro.debug_middleware.DEBUG_MODE', 'off'):
+ from kiro.debug_middleware import DebugLoggerMiddleware
+
+ middleware = DebugLoggerMiddleware(app=MagicMock())
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.url.path = "/v1/chat/completions"
+
+ mock_response = MagicMock(spec=Response)
+ mock_call_next = AsyncMock(return_value=mock_response)
+
+ with patch('kiro.debug_logger.debug_logger') as mock_logger:
+ print("Action: Calling dispatch with DEBUG_MODE=off...")
+ response = await middleware.dispatch(mock_request, mock_call_next)
+
+ print("Verifying prepare_new_request was NOT called...")
+ mock_logger.prepare_new_request.assert_not_called()
+
+ print("Verifying call_next was called...")
+ mock_call_next.assert_called_once()
+
+ @pytest.mark.asyncio
+ async def test_processes_when_debug_mode_errors(self):
+ """
+ What it does: Verifies that middleware works when DEBUG_MODE=errors.
+ Purpose: Ensure errors mode activates logging.
+ """
+ print("Setup: DEBUG_MODE=errors...")
+
+ with patch('kiro.debug_middleware.DEBUG_MODE', 'errors'):
+ from kiro.debug_middleware import DebugLoggerMiddleware
+
+ middleware = DebugLoggerMiddleware(app=MagicMock())
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.url.path = "/v1/chat/completions"
+ mock_request.body = AsyncMock(return_value=b'{"test": "data"}')
+
+ mock_response = MagicMock(spec=Response)
+ mock_call_next = AsyncMock(return_value=mock_response)
+
+ with patch('kiro.debug_logger.debug_logger') as mock_logger:
+ print("Action: Calling dispatch with DEBUG_MODE=errors...")
+ await middleware.dispatch(mock_request, mock_call_next)
+
+ print("Verifying prepare_new_request was called...")
+ mock_logger.prepare_new_request.assert_called_once()
+
+ @pytest.mark.asyncio
+ async def test_processes_when_debug_mode_all(self):
+ """
+ What it does: Verifies that middleware works when DEBUG_MODE=all.
+ Purpose: Ensure all mode activates logging.
+ """
+ print("Setup: DEBUG_MODE=all...")
+
+ with patch('kiro.debug_middleware.DEBUG_MODE', 'all'):
+ from kiro.debug_middleware import DebugLoggerMiddleware
+
+ middleware = DebugLoggerMiddleware(app=MagicMock())
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.url.path = "/v1/messages"
+ mock_request.body = AsyncMock(return_value=b'{"test": "data"}')
+
+ mock_response = MagicMock(spec=Response)
+ mock_call_next = AsyncMock(return_value=mock_response)
+
+ with patch('kiro.debug_logger.debug_logger') as mock_logger:
+ print("Action: Calling dispatch with DEBUG_MODE=all...")
+ await middleware.dispatch(mock_request, mock_call_next)
+
+ print("Verifying prepare_new_request was called...")
+ mock_logger.prepare_new_request.assert_called_once()
+
+
+class TestDebugLoggerMiddlewareErrorHandling:
+ """Tests for error handling in middleware."""
+
+ @pytest.mark.asyncio
+ async def test_handles_body_read_error_gracefully(self):
+ """
+ What it does: Verifies that middleware handles body read errors gracefully.
+ Purpose: Ensure body read errors don't break the request.
+ """
+ print("Setup: Simulating body read error...")
+
+ with patch('kiro.debug_middleware.DEBUG_MODE', 'all'):
+ from kiro.debug_middleware import DebugLoggerMiddleware
+
+ middleware = DebugLoggerMiddleware(app=MagicMock())
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.url.path = "/v1/chat/completions"
+ mock_request.body = AsyncMock(side_effect=Exception("Body read error"))
+
+ mock_response = MagicMock(spec=Response)
+ mock_call_next = AsyncMock(return_value=mock_response)
+
+ with patch('kiro.debug_logger.debug_logger') as mock_logger:
+ print("Action: Calling dispatch with body read error...")
+ response = await middleware.dispatch(mock_request, mock_call_next)
+
+ print("Verifying prepare_new_request was called...")
+ mock_logger.prepare_new_request.assert_called_once()
+
+ print("Verifying log_request_body was NOT called (due to error)...")
+ mock_logger.log_request_body.assert_not_called()
+
+ print("Verifying call_next was called (request continued)...")
+ mock_call_next.assert_called_once()
+
+ @pytest.mark.asyncio
+ async def test_skips_empty_body(self):
+ """
+ What it does: Verifies that middleware doesn't log empty body.
+ Purpose: Ensure empty requests don't create unnecessary logs.
+ """
+ print("Setup: Creating request with empty body...")
+
+ with patch('kiro.debug_middleware.DEBUG_MODE', 'all'):
+ from kiro.debug_middleware import DebugLoggerMiddleware
+
+ middleware = DebugLoggerMiddleware(app=MagicMock())
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.url.path = "/v1/chat/completions"
+ mock_request.body = AsyncMock(return_value=b'') # Empty body
+
+ mock_response = MagicMock(spec=Response)
+ mock_call_next = AsyncMock(return_value=mock_response)
+
+ with patch('kiro.debug_logger.debug_logger') as mock_logger:
+ print("Action: Calling dispatch with empty body...")
+ await middleware.dispatch(mock_request, mock_call_next)
+
+ print("Verifying prepare_new_request was called...")
+ mock_logger.prepare_new_request.assert_called_once()
+
+ print("Verifying log_request_body was NOT called (body is empty)...")
+ mock_logger.log_request_body.assert_not_called()
+
+
+class TestDebugLoggerMiddlewareResponsePassthrough:
+ """Tests for transparent response passthrough."""
+
+ @pytest.mark.asyncio
+ async def test_returns_response_from_call_next(self):
+ """
+ What it does: Verifies that middleware returns response from call_next.
+ Purpose: Ensure middleware doesn't modify the response.
+ """
+ print("Setup: Creating mock response...")
+
+ with patch('kiro.debug_middleware.DEBUG_MODE', 'all'):
+ from kiro.debug_middleware import DebugLoggerMiddleware
+
+ middleware = DebugLoggerMiddleware(app=MagicMock())
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.url.path = "/v1/chat/completions"
+ mock_request.body = AsyncMock(return_value=b'{"test": "data"}')
+
+ expected_response = MagicMock(spec=Response)
+ expected_response.status_code = 200
+ mock_call_next = AsyncMock(return_value=expected_response)
+
+ with patch('kiro.debug_logger.debug_logger'):
+ print("Action: Calling dispatch...")
+ actual_response = await middleware.dispatch(mock_request, mock_call_next)
+
+ print(f"Comparing response: Expected {expected_response}, Got {actual_response}")
+ assert actual_response == expected_response
+ assert actual_response.status_code == 200
+
+
+class TestLoggedEndpointsConstant:
+ """Tests for LOGGED_ENDPOINTS constant."""
+
+ def test_logged_endpoints_contains_chat_completions(self):
+ """
+ What it does: Verifies that LOGGED_ENDPOINTS contains /v1/chat/completions.
+ Purpose: Ensure OpenAI endpoint is included in logging.
+ """
+ print("Checking LOGGED_ENDPOINTS...")
+ from kiro.debug_middleware import LOGGED_ENDPOINTS
+
+ print(f"LOGGED_ENDPOINTS contents: {LOGGED_ENDPOINTS}")
+ assert "/v1/chat/completions" in LOGGED_ENDPOINTS
+
+ def test_logged_endpoints_contains_messages(self):
+ """
+ What it does: Verifies that LOGGED_ENDPOINTS contains /v1/messages.
+ Purpose: Ensure Anthropic endpoint is included in logging.
+ """
+ print("Checking LOGGED_ENDPOINTS...")
+ from kiro.debug_middleware import LOGGED_ENDPOINTS
+
+ print(f"LOGGED_ENDPOINTS contents: {LOGGED_ENDPOINTS}")
+ assert "/v1/messages" in LOGGED_ENDPOINTS
+
+ def test_logged_endpoints_is_frozenset(self):
+ """
+ What it does: Verifies that LOGGED_ENDPOINTS is a frozenset.
+ Purpose: Ensure the constant is immutable.
+ """
+ print("Checking LOGGED_ENDPOINTS type...")
+ from kiro.debug_middleware import LOGGED_ENDPOINTS
+
+ print(f"LOGGED_ENDPOINTS type: {type(LOGGED_ENDPOINTS)}")
+ assert isinstance(LOGGED_ENDPOINTS, frozenset)
diff --git a/kiro-gateway/tests/unit/test_exceptions.py b/kiro-gateway/tests/unit/test_exceptions.py
new file mode 100644
index 0000000000000000000000000000000000000000..0c65033070538cae165d85acb04c6c86d02efca2
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_exceptions.py
@@ -0,0 +1,291 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for exception handlers.
+Tests validation error handling and debug logging integration.
+"""
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+from fastapi import Request
+from fastapi.exceptions import RequestValidationError
+
+
+class TestSanitizeValidationErrors:
+ """Tests for sanitize_validation_errors function."""
+
+ def test_sanitizes_bytes_in_input_field(self):
+ """
+ What it does: Verifies that bytes in 'input' field are converted to strings.
+ Purpose: Ensure JSON serialization works for bytes objects.
+ """
+ print("Setup: Creating error with bytes in input field...")
+ from kiro.exceptions import sanitize_validation_errors
+
+ errors = [
+ {
+ "type": "json_invalid",
+ "loc": ["body", 0],
+ "msg": "Invalid JSON",
+ "input": b'{"invalid": json}'
+ }
+ ]
+
+ print("Action: Calling sanitize_validation_errors...")
+ result = sanitize_validation_errors(errors)
+
+ print(f"Comparing input type: Expected str, Got {type(result[0]['input'])}")
+ assert isinstance(result[0]["input"], str)
+ assert result[0]["input"] == '{"invalid": json}'
+
+ def test_sanitizes_bytes_in_list_values(self):
+ """
+ What it does: Verifies that bytes in list values are converted to strings.
+ Purpose: Ensure nested bytes are handled.
+ """
+ print("Setup: Creating error with bytes in list...")
+ from kiro.exceptions import sanitize_validation_errors
+
+ errors = [
+ {
+ "type": "value_error",
+ "loc": ["body", "messages"],
+ "msg": "Invalid value",
+ "input": [b'bytes1', "string", b'bytes2']
+ }
+ ]
+
+ print("Action: Calling sanitize_validation_errors...")
+ result = sanitize_validation_errors(errors)
+
+ print(f"Checking list values are converted...")
+ assert result[0]["input"] == ["bytes1", "string", "bytes2"]
+
+ def test_preserves_non_bytes_values(self):
+ """
+ What it does: Verifies that non-bytes values are preserved.
+ Purpose: Ensure normal values are not modified.
+ """
+ print("Setup: Creating error with normal values...")
+ from kiro.exceptions import sanitize_validation_errors
+
+ errors = [
+ {
+ "type": "missing",
+ "loc": ["body", "model"],
+ "msg": "Field required",
+ "input": {"messages": []}
+ }
+ ]
+
+ print("Action: Calling sanitize_validation_errors...")
+ result = sanitize_validation_errors(errors)
+
+ print(f"Checking values are preserved...")
+ assert result[0]["input"] == {"messages": []}
+ assert result[0]["type"] == "missing"
+
+
+class TestValidationExceptionHandler:
+ """Tests for validation_exception_handler function."""
+
+ @pytest.mark.asyncio
+ async def test_returns_422_status_code(self):
+ """
+ What it does: Verifies that handler returns 422 status code.
+ Purpose: Ensure proper HTTP status for validation errors.
+ """
+ print("Setup: Creating mock request and exception...")
+ from kiro.exceptions import validation_exception_handler
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.body = AsyncMock(return_value=b'{"invalid": json}')
+
+ mock_exc = MagicMock(spec=RequestValidationError)
+ mock_exc.errors.return_value = [
+ {"type": "json_invalid", "loc": ["body"], "msg": "Invalid JSON", "input": {}}
+ ]
+
+ # Patch debug_logger at the source module
+ with patch('kiro.debug_logger.debug_logger') as mock_logger:
+ print("Action: Calling validation_exception_handler...")
+ response = await validation_exception_handler(mock_request, mock_exc)
+
+ print(f"Comparing status_code: Expected 422, Got {response.status_code}")
+ assert response.status_code == 422
+
+ @pytest.mark.asyncio
+ async def test_calls_flush_on_error_with_422(self):
+ """
+ What it does: Verifies that handler calls flush_on_error(422).
+ Purpose: Ensure debug logs are flushed for validation errors.
+ """
+ print("Setup: Creating mock request and exception...")
+ from kiro.exceptions import validation_exception_handler
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.body = AsyncMock(return_value=b'{"test": "data"}')
+
+ mock_exc = MagicMock(spec=RequestValidationError)
+ mock_exc.errors.return_value = [
+ {"type": "missing", "loc": ["body", "model"], "msg": "Field required", "input": {}}
+ ]
+
+ # Patch debug_logger at the source module
+ with patch('kiro.debug_logger.debug_logger') as mock_logger:
+ print("Action: Calling validation_exception_handler...")
+ await validation_exception_handler(mock_request, mock_exc)
+
+ print("Verifying flush_on_error was called with 422...")
+ mock_logger.flush_on_error.assert_called_once()
+ call_args = mock_logger.flush_on_error.call_args
+ assert call_args[0][0] == 422 # First positional argument is status_code
+
+ @pytest.mark.asyncio
+ async def test_includes_sanitized_errors_in_response(self):
+ """
+ What it does: Verifies that response includes sanitized errors.
+ Purpose: Ensure error details are returned to client.
+ """
+ print("Setup: Creating mock request and exception...")
+ from kiro.exceptions import validation_exception_handler
+ import json
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.body = AsyncMock(return_value=b'{"test": "data"}')
+
+ mock_exc = MagicMock(spec=RequestValidationError)
+ mock_exc.errors.return_value = [
+ {"type": "missing", "loc": ["body", "model"], "msg": "Field required", "input": {}}
+ ]
+
+ with patch('kiro.debug_logger.debug_logger'):
+ print("Action: Calling validation_exception_handler...")
+ response = await validation_exception_handler(mock_request, mock_exc)
+
+ print("Parsing response body...")
+ body = json.loads(response.body.decode())
+
+ print(f"Verifying 'detail' is in response...")
+ assert "detail" in body
+ assert len(body["detail"]) == 1
+ assert body["detail"][0]["type"] == "missing"
+
+ @pytest.mark.asyncio
+ async def test_truncates_body_in_response(self):
+ """
+ What it does: Verifies that body is truncated to 500 chars in response.
+ Purpose: Ensure large bodies don't bloat error responses.
+ """
+ print("Setup: Creating mock request with large body...")
+ from kiro.exceptions import validation_exception_handler
+ import json
+
+ large_body = b'{"data": "' + b'x' * 1000 + b'"}'
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.body = AsyncMock(return_value=large_body)
+
+ mock_exc = MagicMock(spec=RequestValidationError)
+ mock_exc.errors.return_value = [
+ {"type": "json_invalid", "loc": ["body"], "msg": "Invalid", "input": {}}
+ ]
+
+ with patch('kiro.debug_logger.debug_logger'):
+ print("Action: Calling validation_exception_handler...")
+ response = await validation_exception_handler(mock_request, mock_exc)
+
+ print("Parsing response body...")
+ body = json.loads(response.body.decode())
+
+ print(f"Verifying body is truncated to 500 chars...")
+ assert len(body["body"]) <= 500
+
+
+class TestValidationExceptionHandlerLogging:
+ """Tests for logging behavior in validation_exception_handler."""
+
+ @pytest.mark.asyncio
+ async def test_logs_error_at_error_level(self):
+ """
+ What it does: Verifies that validation error is logged at ERROR level.
+ Purpose: Ensure errors are visible in logs.
+ """
+ print("Setup: Creating mock request and exception...")
+ from kiro.exceptions import validation_exception_handler
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.body = AsyncMock(return_value=b'{"test": "data"}')
+
+ mock_exc = MagicMock(spec=RequestValidationError)
+ mock_exc.errors.return_value = [
+ {"type": "missing", "loc": ["body", "model"], "msg": "Field required", "input": {}}
+ ]
+
+ with patch('kiro.debug_logger.debug_logger'):
+ with patch('kiro.exceptions.logger') as mock_logger:
+ print("Action: Calling validation_exception_handler...")
+ await validation_exception_handler(mock_request, mock_exc)
+
+ print("Verifying logger.error was called...")
+ mock_logger.error.assert_called()
+
+
+class TestValidationExceptionHandlerEdgeCases:
+ """Tests for edge cases in validation_exception_handler."""
+
+ @pytest.mark.asyncio
+ async def test_handles_empty_errors_list(self):
+ """
+ What it does: Verifies that handler works with empty errors list.
+ Purpose: Ensure edge case doesn't cause crash.
+ """
+ print("Setup: Creating mock request with empty errors...")
+ from kiro.exceptions import validation_exception_handler
+ import json
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.body = AsyncMock(return_value=b'{}')
+
+ mock_exc = MagicMock(spec=RequestValidationError)
+ mock_exc.errors.return_value = []
+
+ with patch('kiro.debug_logger.debug_logger'):
+ print("Action: Calling validation_exception_handler...")
+ response = await validation_exception_handler(mock_request, mock_exc)
+
+ print(f"Verifying response is valid...")
+ assert response.status_code == 422
+
+ body = json.loads(response.body.decode())
+ assert body["detail"] == []
+
+ @pytest.mark.asyncio
+ async def test_handles_unicode_in_body(self):
+ """
+ What it does: Verifies that handler works with unicode in body.
+ Purpose: Ensure international characters are handled.
+ """
+ print("Setup: Creating mock request with unicode body...")
+ from kiro.exceptions import validation_exception_handler
+ import json
+
+ unicode_body = '{"message": "Привет мир 🌍"}'.encode('utf-8')
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.body = AsyncMock(return_value=unicode_body)
+
+ mock_exc = MagicMock(spec=RequestValidationError)
+ mock_exc.errors.return_value = [
+ {"type": "missing", "loc": ["body", "model"], "msg": "Field required", "input": {}}
+ ]
+
+ with patch('kiro.debug_logger.debug_logger'):
+ print("Action: Calling validation_exception_handler...")
+ response = await validation_exception_handler(mock_request, mock_exc)
+
+ print(f"Verifying response is valid...")
+ assert response.status_code == 422
+
+ body = json.loads(response.body.decode())
+ assert "Привет мир" in body["body"]
diff --git a/kiro-gateway/tests/unit/test_http_client.py b/kiro-gateway/tests/unit/test_http_client.py
new file mode 100644
index 0000000000000000000000000000000000000000..a559e6467252ef981eff1b1e83ec307d07c64c5e
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_http_client.py
@@ -0,0 +1,1158 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for KiroHttpClient.
+Tests retry logic, error handling, and HTTP client management.
+"""
+
+import asyncio
+import pytest
+from unittest.mock import AsyncMock, Mock, patch, MagicMock
+from datetime import datetime, timezone, timedelta
+
+import httpx
+from fastapi import HTTPException
+
+from kiro.http_client import KiroHttpClient
+from kiro.auth import KiroAuthManager
+from kiro.config import MAX_RETRIES, BASE_RETRY_DELAY, FIRST_TOKEN_MAX_RETRIES, STREAMING_READ_TIMEOUT
+
+
+@pytest.fixture
+def mock_auth_manager_for_http():
+ """Creates a mocked KiroAuthManager for HTTP client tests."""
+ manager = Mock(spec=KiroAuthManager)
+ manager.get_access_token = AsyncMock(return_value="test_access_token")
+ manager.force_refresh = AsyncMock(return_value="new_access_token")
+ manager.fingerprint = "test_fingerprint_12345678"
+ manager._fingerprint = "test_fingerprint_12345678"
+ return manager
+
+
+class TestKiroHttpClientInitialization:
+ """Tests for KiroHttpClient initialization."""
+
+ def test_initialization_stores_auth_manager(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies auth_manager is stored during initialization.
+ Purpose: Ensure auth_manager is available for obtaining tokens.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ client = KiroHttpClient(mock_auth_manager_for_http)
+
+ print("Verification: auth_manager is stored...")
+ assert client.auth_manager is mock_auth_manager_for_http
+
+ def test_initialization_client_is_none(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that HTTP client is initially None.
+ Purpose: Ensure lazy initialization.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ client = KiroHttpClient(mock_auth_manager_for_http)
+
+ print("Verification: client is initially None...")
+ assert client.client is None
+
+
+class TestKiroHttpClientGetClient:
+ """Tests for _get_client method."""
+
+ @pytest.mark.asyncio
+ async def test_get_client_creates_new_client(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies creation of a new HTTP client.
+ Purpose: Ensure client is created on first call.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ print("Action: Getting client...")
+ with patch('kiro.http_client.httpx.AsyncClient') as mock_async_client:
+ mock_instance = AsyncMock()
+ mock_instance.is_closed = False
+ mock_async_client.return_value = mock_instance
+
+ client = await http_client._get_client()
+
+ print("Verification: Client created...")
+ mock_async_client.assert_called_once()
+ assert client is mock_instance
+
+ @pytest.mark.asyncio
+ async def test_get_client_reuses_existing_client(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies reuse of existing client.
+ Purpose: Ensure client is not recreated unnecessarily.
+ """
+ print("Setup: Creating KiroHttpClient with existing client...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_existing = AsyncMock()
+ mock_existing.is_closed = False
+ http_client.client = mock_existing
+
+ print("Action: Getting client...")
+ client = await http_client._get_client()
+
+ print("Verification: Existing client returned...")
+ assert client is mock_existing
+
+ @pytest.mark.asyncio
+ async def test_get_client_recreates_closed_client(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies recreation of closed client.
+ Purpose: Ensure closed client is replaced with a new one.
+ """
+ print("Setup: Creating KiroHttpClient with closed client...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_closed = AsyncMock()
+ mock_closed.is_closed = True
+ http_client.client = mock_closed
+
+ print("Action: Getting client...")
+ with patch('kiro.http_client.httpx.AsyncClient') as mock_async_client:
+ mock_new = AsyncMock()
+ mock_new.is_closed = False
+ mock_async_client.return_value = mock_new
+
+ client = await http_client._get_client()
+
+ print("Verification: New client created...")
+ mock_async_client.assert_called_once()
+ assert client is mock_new
+
+
+class TestKiroHttpClientClose:
+ """Tests for close method."""
+
+ @pytest.mark.asyncio
+ async def test_close_closes_client(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies HTTP client closure.
+ Purpose: Ensure aclose() is called.
+ """
+ print("Setup: Creating KiroHttpClient with client...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.aclose = AsyncMock()
+ http_client.client = mock_client
+
+ print("Action: Closing client...")
+ await http_client.close()
+
+ print("Verification: aclose() called...")
+ mock_client.aclose.assert_called_once()
+
+ @pytest.mark.asyncio
+ async def test_close_does_nothing_for_none_client(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that close() doesn't fail for None client.
+ Purpose: Ensure safe close() call without client.
+ """
+ print("Setup: Creating KiroHttpClient without client...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ print("Action: Closing client...")
+ await http_client.close() # Should not raise an error
+
+ print("Verification: No errors...")
+
+ @pytest.mark.asyncio
+ async def test_close_does_nothing_for_closed_client(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that close() doesn't fail for closed client.
+ Purpose: Ensure safe repeated close() call.
+ """
+ print("Setup: Creating KiroHttpClient with closed client...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = True
+ http_client.client = mock_client
+
+ print("Action: Closing client...")
+ await http_client.close()
+
+ print("Verification: aclose() NOT called...")
+ mock_client.aclose.assert_not_called()
+
+
+class TestKiroHttpClientRequestWithRetry:
+ """Tests for request_with_retry method."""
+
+ @pytest.mark.asyncio
+ async def test_successful_request_returns_response(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies successful request.
+ Purpose: Ensure 200 response is returned immediately.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.request = AsyncMock(return_value=mock_response)
+
+ print("Action: Executing request...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"}
+ )
+
+ print("Verification: Response received...")
+ assert response.status_code == 200
+ mock_client.request.assert_called_once()
+
+ @pytest.mark.asyncio
+ async def test_403_triggers_token_refresh(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies token refresh on 403.
+ Purpose: Ensure force_refresh() is called on 403.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response_403 = AsyncMock()
+ mock_response_403.status_code = 403
+
+ mock_response_200 = AsyncMock()
+ mock_response_200.status_code = 200
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.request = AsyncMock(side_effect=[mock_response_403, mock_response_200])
+
+ print("Action: Executing request...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"}
+ )
+
+ print("Verification: force_refresh() called...")
+ mock_auth_manager_for_http.force_refresh.assert_called_once()
+ assert response.status_code == 200
+
+ @pytest.mark.asyncio
+ async def test_429_triggers_backoff(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies exponential backoff on 429.
+ Purpose: Ensure request is retried after delay.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response_429 = AsyncMock()
+ mock_response_429.status_code = 429
+
+ mock_response_200 = AsyncMock()
+ mock_response_200.status_code = 200
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.request = AsyncMock(side_effect=[mock_response_429, mock_response_200])
+
+ print("Action: Executing request...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ with patch('kiro.http_client.asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"}
+ )
+
+ print("Verification: sleep() called for backoff...")
+ mock_sleep.assert_called_once()
+ assert response.status_code == 200
+
+ @pytest.mark.asyncio
+ async def test_5xx_triggers_backoff(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies exponential backoff on 5xx.
+ Purpose: Ensure server errors are handled with retry.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response_500 = AsyncMock()
+ mock_response_500.status_code = 500
+
+ mock_response_200 = AsyncMock()
+ mock_response_200.status_code = 200
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.request = AsyncMock(side_effect=[mock_response_500, mock_response_200])
+
+ print("Action: Executing request...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ with patch('kiro.http_client.asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"}
+ )
+
+ print("Verification: sleep() called for backoff...")
+ mock_sleep.assert_called_once()
+ assert response.status_code == 200
+
+ @pytest.mark.asyncio
+ async def test_timeout_triggers_backoff(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies exponential backoff on timeout.
+ Purpose: Ensure timeouts are handled with retry.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response_200 = AsyncMock()
+ mock_response_200.status_code = 200
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.request = AsyncMock(side_effect=[
+ httpx.TimeoutException("Timeout"),
+ mock_response_200
+ ])
+
+ print("Action: Executing request...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ with patch('kiro.http_client.asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"}
+ )
+
+ print("Verification: sleep() called for backoff...")
+ mock_sleep.assert_called_once()
+ assert response.status_code == 200
+
+ @pytest.mark.asyncio
+ async def test_request_error_triggers_backoff(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies exponential backoff on request error.
+ Purpose: Ensure network errors are handled with retry.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response_200 = AsyncMock()
+ mock_response_200.status_code = 200
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.request = AsyncMock(side_effect=[
+ httpx.RequestError("Connection error"),
+ mock_response_200
+ ])
+
+ print("Action: Executing request...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ with patch('kiro.http_client.asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"}
+ )
+
+ print("Verification: sleep() called for backoff...")
+ mock_sleep.assert_called_once()
+ assert response.status_code == 200
+
+ @pytest.mark.asyncio
+ async def test_max_retries_exceeded_raises_502(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies HTTPException is raised after exhausting retries.
+ Purpose: Ensure 504 is raised after MAX_RETRIES for timeout errors.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.request = AsyncMock(side_effect=httpx.TimeoutException("Timeout"))
+
+ print("Action: Executing request...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ with patch('kiro.http_client.asyncio.sleep', new_callable=AsyncMock):
+ with pytest.raises(HTTPException) as exc_info:
+ await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"}
+ )
+
+ print(f"Verification: HTTPException with code 504 (timeout errors now return 504)...")
+ assert exc_info.value.status_code == 504
+ print(f"Verification: Error detail contains user-friendly message...")
+ assert "timeout" in exc_info.value.detail.lower()
+
+ @pytest.mark.asyncio
+ async def test_other_status_codes_returned_as_is(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies other status codes are returned without retry.
+ Purpose: Ensure 400, 404, etc. are returned immediately.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 400
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.request = AsyncMock(return_value=mock_response)
+
+ print("Action: Executing request...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"}
+ )
+
+ print("Verification: 400 response returned without retry...")
+ assert response.status_code == 400
+ mock_client.request.assert_called_once()
+
+ @pytest.mark.asyncio
+ async def test_streaming_request_uses_send(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies send() is used for streaming.
+ Purpose: Ensure stream=True uses build_request + send.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+
+ mock_request = Mock()
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.build_request = Mock(return_value=mock_request)
+ mock_client.send = AsyncMock(return_value=mock_response)
+
+ print("Action: Executing streaming request...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"},
+ stream=True
+ )
+
+ print("Verification: build_request and send called...")
+ mock_client.build_request.assert_called_once()
+ mock_client.send.assert_called_once_with(mock_request, stream=True)
+ assert response.status_code == 200
+
+
+class TestKiroHttpClientContextManager:
+ """Tests for async context manager."""
+
+ @pytest.mark.asyncio
+ async def test_context_manager_returns_self(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that __aenter__ returns self.
+ Purpose: Ensure correct async with behavior.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ print("Action: Entering context...")
+ result = await http_client.__aenter__()
+
+ print("Verification: self returned...")
+ assert result is http_client
+
+ @pytest.mark.asyncio
+ async def test_context_manager_closes_on_exit(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies client closure on context exit.
+ Purpose: Ensure close() is called in __aexit__.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.aclose = AsyncMock()
+ http_client.client = mock_client
+
+ print("Action: Exiting context...")
+ await http_client.__aexit__(None, None, None)
+
+ print("Verification: aclose() called...")
+ mock_client.aclose.assert_called_once()
+
+
+class TestKiroHttpClientExponentialBackoff:
+ """Tests for exponential backoff logic."""
+
+ @pytest.mark.asyncio
+ async def test_backoff_delay_increases_exponentially(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies exponential delay increase.
+ Purpose: Ensure delay = BASE_RETRY_DELAY * (2 ** attempt).
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response_429 = AsyncMock()
+ mock_response_429.status_code = 429
+
+ mock_response_200 = AsyncMock()
+ mock_response_200.status_code = 200
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ # 2 errors 429, then success (to verify 2 backoff delays)
+ mock_client.request = AsyncMock(side_effect=[
+ mock_response_429,
+ mock_response_429,
+ mock_response_200
+ ])
+
+ sleep_delays = []
+
+ async def capture_sleep(delay):
+ sleep_delays.append(delay)
+
+ print("Action: Executing request with multiple retries...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ with patch('kiro.http_client.asyncio.sleep', side_effect=capture_sleep):
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"}
+ )
+
+ print(f"Verification: Delays increase exponentially...")
+ print(f"Delays: {sleep_delays}")
+ assert len(sleep_delays) == 2
+ assert sleep_delays[0] == BASE_RETRY_DELAY * (2 ** 0) # 1.0
+ assert sleep_delays[1] == BASE_RETRY_DELAY * (2 ** 1) # 2.0
+
+
+class TestKiroHttpClientStreamingTimeout:
+ """Tests for streaming request timeout logic."""
+
+ @pytest.mark.asyncio
+ async def test_streaming_uses_streaming_read_timeout(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that streaming requests use STREAMING_READ_TIMEOUT.
+ Purpose: Ensure stream=True uses httpx.Timeout with correct values.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+
+ mock_request = Mock()
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.build_request = Mock(return_value=mock_request)
+ mock_client.send = AsyncMock(return_value=mock_response)
+
+ print("Action: Executing streaming request...")
+ with patch('kiro.http_client.httpx.AsyncClient') as mock_async_client:
+ mock_async_client.return_value = mock_client
+
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"},
+ stream=True
+ )
+
+ print("Verification: AsyncClient created with httpx.Timeout for streaming...")
+ call_args = mock_async_client.call_args
+ timeout_arg = call_args.kwargs.get('timeout')
+ assert timeout_arg is not None, f"timeout not found in call_args: {call_args}"
+ print(f"Comparing connect: Expected 30.0, Got {timeout_arg.connect}")
+ assert timeout_arg.connect == 30.0, f"Expected connect=30.0, got {timeout_arg.connect}"
+ print(f"Comparing read: Expected {STREAMING_READ_TIMEOUT}, Got {timeout_arg.read}")
+ assert timeout_arg.read == STREAMING_READ_TIMEOUT, f"Expected read={STREAMING_READ_TIMEOUT}, got {timeout_arg.read}"
+ assert call_args.kwargs.get('follow_redirects') == True
+ assert response.status_code == 200
+
+ @pytest.mark.asyncio
+ async def test_streaming_uses_first_token_max_retries(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that streaming requests use FIRST_TOKEN_MAX_RETRIES.
+ Purpose: Ensure stream=True uses separate retry counter.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_request = Mock()
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.build_request = Mock(return_value=mock_request)
+ mock_client.send = AsyncMock(side_effect=httpx.TimeoutException("Timeout"))
+
+ print("Action: Executing streaming request with timeouts...")
+ with patch('kiro.http_client.httpx.AsyncClient', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ with patch('kiro.http_client.asyncio.sleep', new_callable=AsyncMock):
+ with pytest.raises(HTTPException) as exc_info:
+ await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"},
+ stream=True
+ )
+
+ print(f"Verification: HTTPException with code 504...")
+ assert exc_info.value.status_code == 504
+ assert str(FIRST_TOKEN_MAX_RETRIES) in exc_info.value.detail
+
+ print(f"Verification: Attempt count = FIRST_TOKEN_MAX_RETRIES ({FIRST_TOKEN_MAX_RETRIES})...")
+ assert mock_client.send.call_count == FIRST_TOKEN_MAX_RETRIES
+
+ @pytest.mark.asyncio
+ async def test_streaming_timeout_retry_without_delay(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that streaming timeout retry happens with exponential backoff.
+ Purpose: Ensure timeouts are retried with proper delay (new behavior with classifier).
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+
+ mock_request = Mock()
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.build_request = Mock(return_value=mock_request)
+ # First timeout, then success
+ mock_client.send = AsyncMock(side_effect=[
+ httpx.TimeoutException("Timeout"),
+ mock_response
+ ])
+
+ sleep_called = False
+
+ async def capture_sleep(delay):
+ nonlocal sleep_called
+ sleep_called = True
+
+ print("Action: Executing streaming request with one timeout...")
+ with patch('kiro.http_client.httpx.AsyncClient', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ with patch('kiro.http_client.asyncio.sleep', side_effect=capture_sleep):
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"},
+ stream=True
+ )
+
+ print("Verification: sleep() IS called for timeout retry (new behavior)...")
+ assert sleep_called
+ assert response.status_code == 200
+
+ @pytest.mark.asyncio
+ async def test_non_streaming_uses_default_timeout(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that non-streaming requests use 300 seconds.
+ Purpose: Ensure stream=False uses unified httpx.Timeout.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.request = AsyncMock(return_value=mock_response)
+
+ print("Action: Executing non-streaming request...")
+ with patch('kiro.http_client.httpx.AsyncClient') as mock_async_client:
+ mock_async_client.return_value = mock_client
+
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"},
+ stream=False
+ )
+
+ print("Verification: AsyncClient created with httpx.Timeout(timeout=300)...")
+ call_args = mock_async_client.call_args
+ timeout_arg = call_args.kwargs.get('timeout')
+ assert timeout_arg is not None, f"timeout not found in call_args: {call_args}"
+ # httpx.Timeout(timeout=300) sets all timeouts to 300
+ print(f"Comparing timeout: Expected 300.0 for all, Got connect={timeout_arg.connect}")
+ assert timeout_arg.connect == 300.0
+ assert timeout_arg.read == 300.0
+ assert response.status_code == 200
+
+ @pytest.mark.asyncio
+ async def test_connect_timeout_logged_correctly(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies ConnectTimeout logging.
+ Purpose: Ensure ConnectTimeout is logged with user-friendly message.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+
+ mock_request = Mock()
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.build_request = Mock(return_value=mock_request)
+ # First ConnectTimeout, then success
+ mock_client.send = AsyncMock(side_effect=[
+ httpx.ConnectTimeout("Connection timeout"),
+ mock_response
+ ])
+
+ print("Action: Executing streaming request with ConnectTimeout...")
+ with patch('kiro.http_client.httpx.AsyncClient', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ with patch('kiro.http_client.asyncio.sleep', new_callable=AsyncMock):
+ with patch('kiro.http_client.logger') as mock_logger:
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"},
+ stream=True
+ )
+
+ print("Verification: logger.warning called with user-friendly timeout message...")
+ warning_calls = [str(call) for call in mock_logger.warning.call_args_list]
+ assert any("timeout" in call.lower() for call in warning_calls), f"Timeout message not found in: {warning_calls}"
+ assert response.status_code == 200
+
+ @pytest.mark.asyncio
+ async def test_read_timeout_logged_correctly(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies ReadTimeout logging.
+ Purpose: Ensure ReadTimeout is logged with user-friendly message.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+
+ mock_request = Mock()
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.build_request = Mock(return_value=mock_request)
+ # First ReadTimeout, then success
+ mock_client.send = AsyncMock(side_effect=[
+ httpx.ReadTimeout("Read timeout"),
+ mock_response
+ ])
+
+ print("Action: Executing streaming request with ReadTimeout...")
+ with patch('kiro.http_client.httpx.AsyncClient', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ with patch('kiro.http_client.asyncio.sleep', new_callable=AsyncMock):
+ with patch('kiro.http_client.logger') as mock_logger:
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"},
+ stream=True
+ )
+
+ print("Verification: logger.warning called with user-friendly timeout message...")
+ warning_calls = [str(call) for call in mock_logger.warning.call_args_list]
+ assert any("timeout" in call.lower() for call in warning_calls), f"Timeout message not found in: {warning_calls}"
+ assert response.status_code == 200
+
+ @pytest.mark.asyncio
+ async def test_streaming_timeout_returns_504_with_error_type(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that streaming timeout returns 504 with error type.
+ Purpose: Ensure 504 is returned with error info after exhausting retries.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_request = Mock()
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.build_request = Mock(return_value=mock_request)
+ mock_client.send = AsyncMock(side_effect=httpx.ReadTimeout("Timeout"))
+
+ print("Action: Executing streaming request with persistent timeouts...")
+ with patch('kiro.http_client.httpx.AsyncClient', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ with patch('kiro.http_client.asyncio.sleep', new_callable=AsyncMock):
+ with pytest.raises(HTTPException) as exc_info:
+ await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"},
+ stream=True
+ )
+
+ print("Verification: HTTPException with code 504 and user-friendly message...")
+ print(f"Comparing status_code: Expected 504, Got {exc_info.value.status_code}")
+ assert exc_info.value.status_code == 504
+ print(f"Comparing detail: Expected timeout message with troubleshooting in '{exc_info.value.detail}'")
+ assert "timeout" in exc_info.value.detail.lower()
+ assert "Troubleshooting" in exc_info.value.detail or "Technical details" in exc_info.value.detail
+
+ @pytest.mark.asyncio
+ async def test_non_streaming_timeout_returns_502(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that non-streaming timeout returns 504.
+ Purpose: Ensure timeouts consistently return 504 (new behavior with classifier).
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.request = AsyncMock(side_effect=httpx.TimeoutException("Timeout"))
+
+ print("Action: Executing non-streaming request with persistent timeouts...")
+ with patch('kiro.http_client.httpx.AsyncClient', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={}):
+ with patch('kiro.http_client.asyncio.sleep', new_callable=AsyncMock):
+ with pytest.raises(HTTPException) as exc_info:
+ await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"},
+ stream=False
+ )
+
+ print("Verification: HTTPException with code 504 (timeouts now consistently return 504)...")
+ assert exc_info.value.status_code == 504
+
+
+class TestKiroHttpClientSharedClient:
+ """Tests for shared client functionality (connection pooling support)."""
+
+ def test_initialization_with_shared_client(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies shared_client is stored during initialization.
+ Purpose: Ensure shared client is available for connection pooling.
+ """
+ print("Setup: Creating mock shared client...")
+ mock_shared = AsyncMock()
+ mock_shared.is_closed = False
+
+ print("Action: Creating KiroHttpClient with shared client...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http, shared_client=mock_shared)
+
+ print("Verification: shared_client is stored...")
+ print(f"Comparing _shared_client: Expected mock_shared, Got {http_client._shared_client}")
+ assert http_client._shared_client is mock_shared
+ print(f"Comparing client: Expected mock_shared, Got {http_client.client}")
+ assert http_client.client is mock_shared
+
+ def test_initialization_without_shared_client_owns_client(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies _owns_client is True when no shared client provided.
+ Purpose: Ensure client ownership is tracked correctly for cleanup.
+ """
+ print("Setup: Creating KiroHttpClient without shared client...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ print("Verification: _owns_client is True...")
+ print(f"Comparing _owns_client: Expected True, Got {http_client._owns_client}")
+ assert http_client._owns_client is True
+ print(f"Comparing _shared_client: Expected None, Got {http_client._shared_client}")
+ assert http_client._shared_client is None
+
+ def test_initialization_with_shared_client_does_not_own(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies _owns_client is False when shared client provided.
+ Purpose: Ensure shared client is not closed by this instance.
+ """
+ print("Setup: Creating mock shared client...")
+ mock_shared = AsyncMock()
+ mock_shared.is_closed = False
+
+ print("Action: Creating KiroHttpClient with shared client...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http, shared_client=mock_shared)
+
+ print("Verification: _owns_client is False...")
+ print(f"Comparing _owns_client: Expected False, Got {http_client._owns_client}")
+ assert http_client._owns_client is False
+
+ @pytest.mark.asyncio
+ async def test_get_client_returns_shared_client(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies _get_client returns shared client directly.
+ Purpose: Ensure shared client is used without creating new one.
+ """
+ print("Setup: Creating mock shared client...")
+ mock_shared = AsyncMock()
+ mock_shared.is_closed = False
+
+ print("Action: Creating KiroHttpClient with shared client...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http, shared_client=mock_shared)
+
+ print("Action: Getting client...")
+ with patch('kiro.http_client.httpx.AsyncClient') as mock_async_client:
+ client = await http_client._get_client(stream=True)
+
+ print("Verification: Shared client returned, no new client created...")
+ print(f"Comparing client: Expected mock_shared, Got {client}")
+ assert client is mock_shared
+ mock_async_client.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_close_does_not_close_shared_client(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies close() does NOT close shared client.
+ Purpose: Ensure shared client lifecycle is managed by application.
+ """
+ print("Setup: Creating mock shared client...")
+ mock_shared = AsyncMock()
+ mock_shared.is_closed = False
+ mock_shared.aclose = AsyncMock()
+
+ print("Action: Creating KiroHttpClient with shared client...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http, shared_client=mock_shared)
+
+ print("Action: Closing client...")
+ await http_client.close()
+
+ print("Verification: aclose() NOT called on shared client...")
+ mock_shared.aclose.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_close_closes_owned_client(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies close() DOES close owned client.
+ Purpose: Ensure owned client is properly cleaned up.
+ """
+ print("Setup: Creating KiroHttpClient without shared client...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_owned = AsyncMock()
+ mock_owned.is_closed = False
+ mock_owned.aclose = AsyncMock()
+ http_client.client = mock_owned
+
+ print("Action: Closing client...")
+ await http_client.close()
+
+ print("Verification: aclose() called on owned client...")
+ mock_owned.aclose.assert_called_once()
+
+
+class TestKiroHttpClientGracefulClose:
+ """Tests for graceful exception handling in close() method."""
+
+ @pytest.mark.asyncio
+ async def test_close_handles_aclose_exception_gracefully(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies exception in aclose() is caught and doesn't propagate.
+ Purpose: Ensure cleanup errors don't mask original exceptions.
+ """
+ print("Setup: Creating KiroHttpClient with client that raises on close...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.aclose = AsyncMock(side_effect=Exception("Connection reset"))
+ http_client.client = mock_client
+
+ print("Action: Closing client (should not raise)...")
+ # Should not raise - exception should be caught
+ await http_client.close()
+
+ print("Verification: No exception propagated...")
+ # If we get here, the test passed
+ assert True
+
+ @pytest.mark.asyncio
+ async def test_close_logs_warning_on_exception(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies warning is logged when aclose() fails.
+ Purpose: Ensure errors are visible in logs for debugging.
+ """
+ print("Setup: Creating KiroHttpClient with client that raises on close...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.aclose = AsyncMock(side_effect=Exception("Connection reset"))
+ http_client.client = mock_client
+
+ print("Action: Closing client with logger mock...")
+ with patch('kiro.http_client.logger') as mock_logger:
+ await http_client.close()
+
+ print("Verification: logger.warning called...")
+ mock_logger.warning.assert_called_once()
+ warning_message = str(mock_logger.warning.call_args)
+ print(f"Warning message: {warning_message}")
+ assert "Connection reset" in warning_message or "Error closing" in warning_message
+
+
+class TestKiroHttpClientConnectionCloseHeader:
+ """Tests for Connection: close header on streaming requests (issue #38)."""
+
+ @pytest.mark.asyncio
+ async def test_streaming_request_includes_connection_close_header(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that streaming requests include Connection: close header.
+ Purpose: Prevent CLOSE_WAIT connection leak by disabling connection reuse for streaming.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+
+ mock_request = Mock()
+ captured_headers = {}
+
+ def capture_build_request(method, url, json, headers):
+ captured_headers.update(headers)
+ return mock_request
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.build_request = Mock(side_effect=capture_build_request)
+ mock_client.send = AsyncMock(return_value=mock_response)
+
+ print("Action: Executing streaming request...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={"Authorization": "Bearer test"}):
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"},
+ stream=True
+ )
+
+ print("Verification: Connection: close header is present...")
+ print(f"Captured headers: {captured_headers}")
+ assert "Connection" in captured_headers, f"Connection header not found in: {captured_headers}"
+ print(f"Comparing Connection: Expected 'close', Got '{captured_headers['Connection']}'")
+ assert captured_headers["Connection"] == "close"
+ assert response.status_code == 200
+
+ @pytest.mark.asyncio
+ async def test_non_streaming_request_does_not_include_connection_close_header(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that non-streaming requests do NOT include Connection: close header.
+ Purpose: Ensure connection pooling is preserved for non-streaming requests.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+
+ captured_headers = {}
+
+ async def capture_request(method, url, json, headers):
+ captured_headers.update(headers)
+ return mock_response
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.request = AsyncMock(side_effect=capture_request)
+
+ print("Action: Executing non-streaming request...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value={"Authorization": "Bearer test"}):
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"},
+ stream=False
+ )
+
+ print("Verification: Connection: close header is NOT present...")
+ print(f"Captured headers: {captured_headers}")
+ assert "Connection" not in captured_headers, f"Connection header should not be present for non-streaming: {captured_headers}"
+ assert response.status_code == 200
+
+ @pytest.mark.asyncio
+ async def test_streaming_connection_close_preserves_other_headers(self, mock_auth_manager_for_http):
+ """
+ What it does: Verifies that adding Connection: close doesn't remove other headers.
+ Purpose: Ensure Authorization and other headers are preserved.
+ """
+ print("Setup: Creating KiroHttpClient...")
+ http_client = KiroHttpClient(mock_auth_manager_for_http)
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+
+ mock_request = Mock()
+ captured_headers = {}
+
+ def capture_build_request(method, url, json, headers):
+ captured_headers.update(headers)
+ return mock_request
+
+ mock_client = AsyncMock()
+ mock_client.is_closed = False
+ mock_client.build_request = Mock(side_effect=capture_build_request)
+ mock_client.send = AsyncMock(return_value=mock_response)
+
+ original_headers = {
+ "Authorization": "Bearer test_token",
+ "Content-Type": "application/json",
+ "X-Custom-Header": "custom_value"
+ }
+
+ print("Action: Executing streaming request with multiple headers...")
+ with patch.object(http_client, '_get_client', return_value=mock_client):
+ with patch('kiro.http_client.get_kiro_headers', return_value=original_headers.copy()):
+ response = await http_client.request_with_retry(
+ "POST",
+ "https://api.example.com/test",
+ {"data": "value"},
+ stream=True
+ )
+
+ print("Verification: All original headers preserved plus Connection: close...")
+ print(f"Captured headers: {captured_headers}")
+ assert captured_headers["Authorization"] == "Bearer test_token"
+ assert captured_headers["Content-Type"] == "application/json"
+ assert captured_headers["X-Custom-Header"] == "custom_value"
+ assert captured_headers["Connection"] == "close"
+ assert response.status_code == 200
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_main_cli.py b/kiro-gateway/tests/unit/test_main_cli.py
new file mode 100644
index 0000000000000000000000000000000000000000..994ac986588255706cdb4ee01d1fdf3ae951c206
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_main_cli.py
@@ -0,0 +1,409 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for main.py CLI functions.
+Tests for parse_cli_args(), resolve_server_config(), and print_startup_banner().
+"""
+
+import pytest
+import argparse
+import sys
+from unittest.mock import patch, MagicMock
+from io import StringIO
+
+
+class TestParseCliArgs:
+ """Tests for parse_cli_args() function."""
+
+ def test_default_values_are_none(self):
+ """
+ What it does: Verifies that default values for host and port are None.
+ Purpose: Ensure that None indicates "use env or default" in priority resolution.
+ """
+ print("Setup: Importing parse_cli_args...")
+ from main import parse_cli_args
+
+ print("Action: Calling parse_cli_args with no arguments...")
+ with patch.object(sys, 'argv', ['main.py']):
+ args = parse_cli_args()
+
+ print(f"args.host: {args.host}")
+ print(f"args.port: {args.port}")
+ print(f"Comparing: Expected host=None, port=None")
+ assert args.host is None
+ assert args.port is None
+
+ def test_port_argument_long_form(self):
+ """
+ What it does: Verifies that --port argument is parsed correctly.
+ Purpose: Ensure long form --port works.
+ """
+ print("Setup: Importing parse_cli_args...")
+ from main import parse_cli_args
+
+ print("Action: Calling parse_cli_args with --port 9000...")
+ with patch.object(sys, 'argv', ['main.py', '--port', '9000']):
+ args = parse_cli_args()
+
+ print(f"args.port: {args.port}")
+ print(f"Comparing: Expected 9000, Got {args.port}")
+ assert args.port == 9000
+
+ def test_port_argument_short_form(self):
+ """
+ What it does: Verifies that -p argument is parsed correctly.
+ Purpose: Ensure short form -p works.
+ """
+ print("Setup: Importing parse_cli_args...")
+ from main import parse_cli_args
+
+ print("Action: Calling parse_cli_args with -p 8080...")
+ with patch.object(sys, 'argv', ['main.py', '-p', '8080']):
+ args = parse_cli_args()
+
+ print(f"args.port: {args.port}")
+ print(f"Comparing: Expected 8080, Got {args.port}")
+ assert args.port == 8080
+
+ def test_host_argument_long_form(self):
+ """
+ What it does: Verifies that --host argument is parsed correctly.
+ Purpose: Ensure long form --host works.
+ """
+ print("Setup: Importing parse_cli_args...")
+ from main import parse_cli_args
+
+ print("Action: Calling parse_cli_args with --host 127.0.0.1...")
+ with patch.object(sys, 'argv', ['main.py', '--host', '127.0.0.1']):
+ args = parse_cli_args()
+
+ print(f"args.host: {args.host}")
+ print(f"Comparing: Expected '127.0.0.1', Got '{args.host}'")
+ assert args.host == "127.0.0.1"
+
+ def test_host_argument_short_form(self):
+ """
+ What it does: Verifies that -H argument is parsed correctly.
+ Purpose: Ensure short form -H works.
+ """
+ print("Setup: Importing parse_cli_args...")
+ from main import parse_cli_args
+
+ print("Action: Calling parse_cli_args with -H 192.168.1.1...")
+ with patch.object(sys, 'argv', ['main.py', '-H', '192.168.1.1']):
+ args = parse_cli_args()
+
+ print(f"args.host: {args.host}")
+ print(f"Comparing: Expected '192.168.1.1', Got '{args.host}'")
+ assert args.host == "192.168.1.1"
+
+ def test_both_arguments_together(self):
+ """
+ What it does: Verifies that both --host and --port can be used together.
+ Purpose: Ensure both arguments work simultaneously.
+ """
+ print("Setup: Importing parse_cli_args...")
+ from main import parse_cli_args
+
+ print("Action: Calling parse_cli_args with --host 0.0.0.0 --port 3000...")
+ with patch.object(sys, 'argv', ['main.py', '--host', '0.0.0.0', '--port', '3000']):
+ args = parse_cli_args()
+
+ print(f"args.host: {args.host}")
+ print(f"args.port: {args.port}")
+ assert args.host == "0.0.0.0"
+ assert args.port == 3000
+
+ def test_short_forms_together(self):
+ """
+ What it does: Verifies that both -H and -p can be used together.
+ Purpose: Ensure short forms work simultaneously.
+ """
+ print("Setup: Importing parse_cli_args...")
+ from main import parse_cli_args
+
+ print("Action: Calling parse_cli_args with -H 127.0.0.1 -p 5000...")
+ with patch.object(sys, 'argv', ['main.py', '-H', '127.0.0.1', '-p', '5000']):
+ args = parse_cli_args()
+
+ print(f"args.host: {args.host}")
+ print(f"args.port: {args.port}")
+ assert args.host == "127.0.0.1"
+ assert args.port == 5000
+
+
+class TestResolveServerConfig:
+ """Tests for resolve_server_config() function - priority hierarchy."""
+
+ def test_cli_args_take_priority_over_env(self):
+ """
+ What it does: Verifies that CLI arguments have highest priority.
+ Purpose: Ensure CLI args override environment variables.
+ """
+ print("Setup: Importing resolve_server_config...")
+ from main import resolve_server_config
+
+ print("Setup: Creating args with host=127.0.0.1, port=9000...")
+ args = argparse.Namespace(host="127.0.0.1", port=9000)
+
+ print("Action: Calling resolve_server_config with CLI args...")
+ # Even if env vars are set, CLI should win
+ with patch('main.SERVER_HOST', '0.0.0.0'), \
+ patch('main.SERVER_PORT', 8000), \
+ patch('main.DEFAULT_SERVER_HOST', '0.0.0.0'), \
+ patch('main.DEFAULT_SERVER_PORT', 8000):
+ host, port = resolve_server_config(args)
+
+ print(f"Resolved host: {host}")
+ print(f"Resolved port: {port}")
+ print(f"Comparing: Expected ('127.0.0.1', 9000)")
+ assert host == "127.0.0.1"
+ assert port == 9000
+
+ def test_env_vars_take_priority_over_defaults(self):
+ """
+ What it does: Verifies that env vars have priority over defaults.
+ Purpose: Ensure env vars are used when CLI args are not provided.
+ """
+ print("Setup: Importing resolve_server_config...")
+ from main import resolve_server_config
+
+ print("Setup: Creating args with host=None, port=None (no CLI args)...")
+ args = argparse.Namespace(host=None, port=None)
+
+ print("Action: Calling resolve_server_config with env vars set...")
+ # SERVER_HOST and SERVER_PORT are different from defaults
+ with patch('main.SERVER_HOST', '192.168.1.100'), \
+ patch('main.SERVER_PORT', 3000), \
+ patch('main.DEFAULT_SERVER_HOST', '0.0.0.0'), \
+ patch('main.DEFAULT_SERVER_PORT', 8000):
+ host, port = resolve_server_config(args)
+
+ print(f"Resolved host: {host}")
+ print(f"Resolved port: {port}")
+ print(f"Comparing: Expected ('192.168.1.100', 3000)")
+ assert host == "192.168.1.100"
+ assert port == 3000
+
+ def test_defaults_used_when_nothing_set(self):
+ """
+ What it does: Verifies that defaults are used when nothing else is set.
+ Purpose: Ensure default values work correctly.
+ """
+ print("Setup: Importing resolve_server_config...")
+ from main import resolve_server_config
+
+ print("Setup: Creating args with host=None, port=None...")
+ args = argparse.Namespace(host=None, port=None)
+
+ print("Action: Calling resolve_server_config with defaults...")
+ # SERVER_HOST and SERVER_PORT equal to defaults (no env override)
+ with patch('main.SERVER_HOST', '0.0.0.0'), \
+ patch('main.SERVER_PORT', 8000), \
+ patch('main.DEFAULT_SERVER_HOST', '0.0.0.0'), \
+ patch('main.DEFAULT_SERVER_PORT', 8000):
+ host, port = resolve_server_config(args)
+
+ print(f"Resolved host: {host}")
+ print(f"Resolved port: {port}")
+ print(f"Comparing: Expected ('0.0.0.0', 8000)")
+ assert host == "0.0.0.0"
+ assert port == 8000
+
+ def test_cli_host_only_env_port(self):
+ """
+ What it does: Verifies mixed priority - CLI host with env port.
+ Purpose: Ensure each argument is resolved independently.
+ """
+ print("Setup: Importing resolve_server_config...")
+ from main import resolve_server_config
+
+ print("Setup: Creating args with host='127.0.0.1', port=None...")
+ args = argparse.Namespace(host="127.0.0.1", port=None)
+
+ print("Action: Calling resolve_server_config...")
+ with patch('main.SERVER_HOST', '0.0.0.0'), \
+ patch('main.SERVER_PORT', 9000), \
+ patch('main.DEFAULT_SERVER_HOST', '0.0.0.0'), \
+ patch('main.DEFAULT_SERVER_PORT', 8000):
+ host, port = resolve_server_config(args)
+
+ print(f"Resolved host: {host}")
+ print(f"Resolved port: {port}")
+ print(f"Comparing: Expected ('127.0.0.1', 9000)")
+ assert host == "127.0.0.1" # From CLI
+ assert port == 9000 # From env (different from default)
+
+ def test_cli_port_only_env_host(self):
+ """
+ What it does: Verifies mixed priority - CLI port with env host.
+ Purpose: Ensure each argument is resolved independently.
+ """
+ print("Setup: Importing resolve_server_config...")
+ from main import resolve_server_config
+
+ print("Setup: Creating args with host=None, port=5000...")
+ args = argparse.Namespace(host=None, port=5000)
+
+ print("Action: Calling resolve_server_config...")
+ with patch('main.SERVER_HOST', '192.168.1.1'), \
+ patch('main.SERVER_PORT', 8000), \
+ patch('main.DEFAULT_SERVER_HOST', '0.0.0.0'), \
+ patch('main.DEFAULT_SERVER_PORT', 8000):
+ host, port = resolve_server_config(args)
+
+ print(f"Resolved host: {host}")
+ print(f"Resolved port: {port}")
+ print(f"Comparing: Expected ('192.168.1.1', 5000)")
+ assert host == "192.168.1.1" # From env (different from default)
+ assert port == 5000 # From CLI
+
+
+class TestPrintStartupBanner:
+ """Tests for print_startup_banner() function."""
+
+ def test_banner_contains_url(self, capsys):
+ """
+ What it does: Verifies that banner contains the server URL.
+ Purpose: Ensure URL is displayed to user.
+ """
+ print("Setup: Importing print_startup_banner...")
+ from main import print_startup_banner
+
+ print("Action: Calling print_startup_banner('0.0.0.0', 8000)...")
+ print_startup_banner("0.0.0.0", 8000)
+
+ captured = capsys.readouterr()
+ print(f"Captured output length: {len(captured.out)}")
+
+ # When host is 0.0.0.0, display should show localhost
+ assert "localhost:8000" in captured.out or "8000" in captured.out
+
+ def test_banner_contains_custom_port(self, capsys):
+ """
+ What it does: Verifies that banner shows custom port.
+ Purpose: Ensure custom port is displayed correctly.
+ """
+ print("Setup: Importing print_startup_banner...")
+ from main import print_startup_banner
+
+ print("Action: Calling print_startup_banner('127.0.0.1', 9000)...")
+ print_startup_banner("127.0.0.1", 9000)
+
+ captured = capsys.readouterr()
+ print(f"Captured output contains '9000': {'9000' in captured.out}")
+
+ assert "9000" in captured.out
+
+ def test_banner_contains_docs_url(self, capsys):
+ """
+ What it does: Verifies that banner contains API docs URL.
+ Purpose: Ensure /docs endpoint is mentioned.
+ """
+ print("Setup: Importing print_startup_banner...")
+ from main import print_startup_banner
+
+ print("Action: Calling print_startup_banner('0.0.0.0', 8000)...")
+ print_startup_banner("0.0.0.0", 8000)
+
+ captured = capsys.readouterr()
+ print(f"Captured output contains '/docs': {'/docs' in captured.out}")
+
+ assert "/docs" in captured.out
+
+ def test_banner_contains_health_url(self, capsys):
+ """
+ What it does: Verifies that banner contains health check URL.
+ Purpose: Ensure /health endpoint is mentioned.
+ """
+ print("Setup: Importing print_startup_banner...")
+ from main import print_startup_banner
+
+ print("Action: Calling print_startup_banner('0.0.0.0', 8000)...")
+ print_startup_banner("0.0.0.0", 8000)
+
+ captured = capsys.readouterr()
+ print(f"Captured output contains '/health': {'/health' in captured.out}")
+
+ assert "/health" in captured.out
+
+
+class TestCliHelp:
+ """Tests for CLI help output."""
+
+ def test_help_shows_port_option(self):
+ """
+ What it does: Verifies that --help shows port option.
+ Purpose: Ensure help is informative.
+ """
+ print("Setup: Importing parse_cli_args...")
+ from main import parse_cli_args
+
+ print("Action: Calling parse_cli_args with --help...")
+ with patch.object(sys, 'argv', ['main.py', '--help']):
+ with pytest.raises(SystemExit) as exc_info:
+ parse_cli_args()
+
+ print(f"Exit code: {exc_info.value.code}")
+ # --help exits with code 0
+ assert exc_info.value.code == 0
+
+ def test_help_shows_host_option(self, capsys):
+ """
+ What it does: Verifies that --help output contains host option.
+ Purpose: Ensure host option is documented.
+ """
+ print("Setup: Importing parse_cli_args...")
+ from main import parse_cli_args
+
+ print("Action: Calling parse_cli_args with --help...")
+ with patch.object(sys, 'argv', ['main.py', '--help']):
+ with pytest.raises(SystemExit):
+ parse_cli_args()
+
+ captured = capsys.readouterr()
+ print(f"Help output contains '--host': {'--host' in captured.out}")
+ print(f"Help output contains '-H': {'-H' in captured.out}")
+
+ assert "--host" in captured.out
+ assert "-H" in captured.out
+
+
+class TestCliVersion:
+ """Tests for CLI version output."""
+
+ def test_version_flag_exits_with_zero(self):
+ """
+ What it does: Verifies that --version exits with code 0.
+ Purpose: Ensure version flag works correctly.
+ """
+ print("Setup: Importing parse_cli_args...")
+ from main import parse_cli_args
+
+ print("Action: Calling parse_cli_args with --version...")
+ with patch.object(sys, 'argv', ['main.py', '--version']):
+ with pytest.raises(SystemExit) as exc_info:
+ parse_cli_args()
+
+ print(f"Exit code: {exc_info.value.code}")
+ assert exc_info.value.code == 0
+
+ def test_version_shows_app_version(self, capsys):
+ """
+ What it does: Verifies that --version shows application version.
+ Purpose: Ensure version is displayed.
+ """
+ print("Setup: Importing parse_cli_args and APP_VERSION...")
+ from main import parse_cli_args
+ from kiro.config import APP_VERSION
+
+ print("Action: Calling parse_cli_args with --version...")
+ with patch.object(sys, 'argv', ['main.py', '--version']):
+ with pytest.raises(SystemExit):
+ parse_cli_args()
+
+ captured = capsys.readouterr()
+ print(f"Version output: {captured.out}")
+ print(f"APP_VERSION: {APP_VERSION}")
+
+ assert APP_VERSION in captured.out
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_model_resolver.py b/kiro-gateway/tests/unit/test_model_resolver.py
new file mode 100644
index 0000000000000000000000000000000000000000..e52575f3bd03960a09f9dd275a1f73fdcc455298
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_model_resolver.py
@@ -0,0 +1,1288 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for Dynamic Model Resolution System.
+
+Tests 4-layer model resolution architecture:
+1. Normalize Name - convert client formats to Kiro format
+2. Check Dynamic Cache - models from /ListAvailableModels API
+3. Check Hidden Models - manual config for undocumented models
+4. Pass-through - unknown models are sent to Kiro
+"""
+
+import pytest
+from dataclasses import FrozenInstanceError
+
+from kiro.model_resolver import (
+ normalize_model_name,
+ get_model_id_for_kiro,
+ extract_model_family,
+ ModelResolver,
+ ModelResolution,
+)
+from kiro.cache import ModelInfoCache
+
+
+# =============================================================================
+# Fixtures
+# =============================================================================
+
+@pytest.fixture
+def mock_model_cache():
+ """
+ Creates ModelInfoCache with pre-populated models.
+ Simulates data from Kiro /ListAvailableModels API.
+ """
+ print("Setup: Creating ModelInfoCache with test models...")
+ cache = ModelInfoCache()
+ # Directly populate cache (without async update)
+ cache._cache = {
+ "auto": {"modelId": "auto", "modelName": "Auto"},
+ "claude-sonnet-4.5": {"modelId": "claude-sonnet-4.5", "modelName": "Claude Sonnet 4.5"},
+ "claude-sonnet-4": {"modelId": "claude-sonnet-4", "modelName": "Claude Sonnet 4"},
+ "claude-haiku-4.5": {"modelId": "claude-haiku-4.5", "modelName": "Claude Haiku 4.5"},
+ "claude-opus-4.5": {"modelId": "claude-opus-4.5", "modelName": "Claude Opus 4.5"},
+ }
+ return cache
+
+
+@pytest.fixture
+def empty_model_cache():
+ """Creates empty ModelInfoCache."""
+ print("Setup: Creating empty ModelInfoCache...")
+ return ModelInfoCache()
+
+
+@pytest.fixture
+def hidden_models():
+ """Hidden models for tests."""
+ return {
+ "claude-3.7-sonnet": "CLAUDE_3_7_SONNET_20250219_V1_0",
+ }
+
+
+@pytest.fixture
+def model_resolver(mock_model_cache, hidden_models):
+ """Ready-to-use ModelResolver for tests."""
+ print("Setup: Creating ModelResolver with cache and hidden models...")
+ return ModelResolver(cache=mock_model_cache, hidden_models=hidden_models)
+
+
+@pytest.fixture
+def resolver_without_hidden(mock_model_cache):
+ """ModelResolver without hidden models."""
+ print("Setup: Creating ModelResolver without hidden models...")
+ return ModelResolver(cache=mock_model_cache, hidden_models={})
+
+
+# =============================================================================
+# TestNormalizeModelName - Tests for model name normalization
+# =============================================================================
+
+class TestNormalizeModelName:
+ """
+ Tests for normalize_model_name() function.
+
+ Checks conversion of client formats to Kiro format:
+ - Dashes → dots for minor versions
+ - Removal of date suffix (20251001)
+ - Removal of 'latest' suffix
+ - Legacy format (claude-3-7-sonnet)
+ """
+
+ # === Standard format with minor version ===
+
+ def test_normalizes_haiku_dash_to_dot(self):
+ """
+ What it does: claude-haiku-4-5 → claude-haiku-4.5
+ Goal: Check dash-to-dot conversion for Haiku.
+ """
+ print("Action: Normalizing 'claude-haiku-4-5'...")
+ result = normalize_model_name("claude-haiku-4-5")
+
+ print(f"Comparing result: Expected 'claude-haiku-4.5', Got '{result}'")
+ assert result == "claude-haiku-4.5"
+
+ def test_normalizes_sonnet_dash_to_dot(self):
+ """
+ What it does: claude-sonnet-4-5 → claude-sonnet-4.5
+ Goal: Check dash-to-dot conversion for Sonnet.
+ """
+ print("Action: Normalizing 'claude-sonnet-4-5'...")
+ result = normalize_model_name("claude-sonnet-4-5")
+
+ print(f"Comparing result: Expected 'claude-sonnet-4.5', Got '{result}'")
+ assert result == "claude-sonnet-4.5"
+
+ def test_normalizes_opus_dash_to_dot(self):
+ """
+ What it does: claude-opus-4-5 → claude-opus-4.5
+ Goal: Check dash-to-dot conversion for Opus.
+ """
+ print("Action: Normalizing 'claude-opus-4-5'...")
+ result = normalize_model_name("claude-opus-4-5")
+
+ print(f"Comparing result: Expected 'claude-opus-4.5', Got '{result}'")
+ assert result == "claude-opus-4.5"
+
+ # === Removal of date suffix ===
+
+ def test_strips_date_suffix_haiku(self):
+ """
+ What it does: claude-haiku-4-5-20251001 → claude-haiku-4.5
+ Goal: Check date suffix removal for Haiku (Claude Code format).
+ """
+ print("Action: Normalizing 'claude-haiku-4-5-20251001'...")
+ result = normalize_model_name("claude-haiku-4-5-20251001")
+
+ print(f"Comparing result: Expected 'claude-haiku-4.5', Got '{result}'")
+ assert result == "claude-haiku-4.5"
+
+ def test_strips_date_suffix_sonnet(self):
+ """
+ What it does: claude-sonnet-4-5-20250929 → claude-sonnet-4.5
+ Goal: Check date suffix removal for Sonnet.
+ """
+ print("Action: Normalizing 'claude-sonnet-4-5-20250929'...")
+ result = normalize_model_name("claude-sonnet-4-5-20250929")
+
+ print(f"Comparing result: Expected 'claude-sonnet-4.5', Got '{result}'")
+ assert result == "claude-sonnet-4.5"
+
+ def test_strips_date_suffix_opus(self):
+ """
+ What it does: claude-opus-4-5-20251101 → claude-opus-4.5
+ Goal: Check date suffix removal for Opus.
+ """
+ print("Action: Normalizing 'claude-opus-4-5-20251101'...")
+ result = normalize_model_name("claude-opus-4-5-20251101")
+
+ print(f"Comparing result: Expected 'claude-opus-4.5', Got '{result}'")
+ assert result == "claude-opus-4.5"
+
+ # === Removal of 'latest' suffix ===
+
+ def test_strips_latest_suffix(self):
+ """
+ What it does: claude-haiku-4-5-latest → claude-haiku-4.5
+ Goal: Check 'latest' suffix removal.
+ """
+ print("Action: Normalizing 'claude-haiku-4-5-latest'...")
+ result = normalize_model_name("claude-haiku-4-5-latest")
+
+ print(f"Comparing result: Expected 'claude-haiku-4.5', Got '{result}'")
+ assert result == "claude-haiku-4.5"
+
+ # === Standard format without minor version ===
+
+ def test_keeps_model_without_minor(self):
+ """
+ What it does: claude-sonnet-4 → claude-sonnet-4
+ Goal: Check that models without minor version are unchanged.
+ """
+ print("Action: Normalizing 'claude-sonnet-4'...")
+ result = normalize_model_name("claude-sonnet-4")
+
+ print(f"Comparing result: Expected 'claude-sonnet-4', Got '{result}'")
+ assert result == "claude-sonnet-4"
+
+ def test_strips_date_from_model_without_minor(self):
+ """
+ What it does: claude-sonnet-4-20250514 → claude-sonnet-4
+ Goal: Check date suffix removal for model without minor version.
+ """
+ print("Action: Normalizing 'claude-sonnet-4-20250514'...")
+ result = normalize_model_name("claude-sonnet-4-20250514")
+
+ print(f"Comparing result: Expected 'claude-sonnet-4', Got '{result}'")
+ assert result == "claude-sonnet-4"
+
+ # === Legacy format (claude-X-Y-family) ===
+
+ def test_normalizes_legacy_format(self):
+ """
+ What it does: claude-3-7-sonnet → claude-3.7-sonnet
+ Goal: Check legacy format normalization.
+ """
+ print("Action: Normalizing 'claude-3-7-sonnet'...")
+ result = normalize_model_name("claude-3-7-sonnet")
+
+ print(f"Comparing result: Expected 'claude-3.7-sonnet', Got '{result}'")
+ assert result == "claude-3.7-sonnet"
+
+ def test_normalizes_legacy_format_with_date(self):
+ """
+ What it does: claude-3-7-sonnet-20250219 → claude-3.7-sonnet
+ Goal: Check legacy format normalization with date suffix.
+ """
+ print("Action: Normalizing 'claude-3-7-sonnet-20250219'...")
+ result = normalize_model_name("claude-3-7-sonnet-20250219")
+
+ print(f"Comparing result: Expected 'claude-3.7-sonnet', Got '{result}'")
+ assert result == "claude-3.7-sonnet"
+
+ def test_normalizes_legacy_haiku(self):
+ """
+ What it does: claude-3-5-haiku → claude-3.5-haiku
+ Goal: Check legacy format normalization for Haiku.
+ """
+ print("Action: Normalizing 'claude-3-5-haiku'...")
+ result = normalize_model_name("claude-3-5-haiku")
+
+ print(f"Comparing result: Expected 'claude-3.5-haiku', Got '{result}'")
+ assert result == "claude-3.5-haiku"
+
+ def test_normalizes_legacy_opus(self):
+ """
+ What it does: claude-3-0-opus → claude-3.0-opus
+ Goal: Check legacy format normalization for Opus.
+ """
+ print("Action: Normalizing 'claude-3-0-opus'...")
+ result = normalize_model_name("claude-3-0-opus")
+
+ print(f"Comparing result: Expected 'claude-3.0-opus', Got '{result}'")
+ assert result == "claude-3.0-opus"
+
+ # === Inverted format with suffix (Pattern 5 - Cursor IDE) ===
+
+ def test_inverted_format_with_high_suffix(self):
+ """
+ What it does: claude-4.5-opus-high → claude-opus-4.5
+ Goal: Check inverted format normalization with 'high' suffix (Cursor IDE).
+
+ Cursor IDE sends model names in inverted format with priority suffix.
+ This is Pattern 5 from PR #49.
+ """
+ print("Action: Normalizing 'claude-4.5-opus-high'...")
+ result = normalize_model_name("claude-4.5-opus-high")
+
+ print(f"Comparing result: Expected 'claude-opus-4.5', Got '{result}'")
+ assert result == "claude-opus-4.5"
+
+ def test_inverted_format_with_low_suffix(self):
+ """
+ What it does: claude-4.5-sonnet-low → claude-sonnet-4.5
+ Goal: Check inverted format normalization with 'low' suffix (Cursor IDE).
+ """
+ print("Action: Normalizing 'claude-4.5-sonnet-low'...")
+ result = normalize_model_name("claude-4.5-sonnet-low")
+
+ print(f"Comparing result: Expected 'claude-sonnet-4.5', Got '{result}'")
+ assert result == "claude-sonnet-4.5"
+
+ def test_inverted_format_with_thinking_suffix(self):
+ """
+ What it does: claude-4.5-opus-high-thinking → claude-opus-4.5
+ Goal: Check inverted format with compound suffix (high-thinking).
+
+ The pattern strips ALL suffixes after the family name.
+ """
+ print("Action: Normalizing 'claude-4.5-opus-high-thinking'...")
+ result = normalize_model_name("claude-4.5-opus-high-thinking")
+
+ print(f"Comparing result: Expected 'claude-opus-4.5', Got '{result}'")
+ assert result == "claude-opus-4.5"
+
+ def test_inverted_format_all_families(self):
+ """
+ What it does: Verifies inverted format works for all families.
+ Goal: Check haiku, sonnet, opus all work with inverted format.
+ """
+ print("Action: Normalizing inverted format for all families...")
+
+ print(" Testing haiku...")
+ result_haiku = normalize_model_name("claude-4.5-haiku-high")
+ print(f" Comparing: Expected 'claude-haiku-4.5', Got '{result_haiku}'")
+ assert result_haiku == "claude-haiku-4.5"
+
+ print(" Testing sonnet...")
+ result_sonnet = normalize_model_name("claude-4.5-sonnet-low")
+ print(f" Comparing: Expected 'claude-sonnet-4.5', Got '{result_sonnet}'")
+ assert result_sonnet == "claude-sonnet-4.5"
+
+ print(" Testing opus...")
+ result_opus = normalize_model_name("claude-4.5-opus-high")
+ print(f" Comparing: Expected 'claude-opus-4.5', Got '{result_opus}'")
+ assert result_opus == "claude-opus-4.5"
+
+ def test_inverted_format_requires_suffix(self):
+ """
+ What it does: Verifies that suffix is required (doesn't match claude-3.7-sonnet).
+ Goal: CRITICAL - ensure Pattern 5 doesn't break already-normalized formats.
+
+ This is the most important test for Pattern 5. The regex MUST require a suffix
+ to avoid matching already-normalized formats like claude-3.7-sonnet.
+ """
+ print("Action: Normalizing 'claude-3.7-sonnet' (should NOT match Pattern 5)...")
+ result = normalize_model_name("claude-3.7-sonnet")
+
+ print(f"Comparing result: Expected 'claude-3.7-sonnet' (unchanged), Got '{result}'")
+ assert result == "claude-3.7-sonnet"
+
+ print("Action: Normalizing 'claude-4.5-sonnet' (should NOT match Pattern 5)...")
+ result2 = normalize_model_name("claude-4.5-sonnet")
+
+ print(f"Comparing result: Expected 'claude-4.5-sonnet' (unchanged), Got '{result2}'")
+ assert result2 == "claude-4.5-sonnet"
+
+ def test_inverted_format_case_insensitive(self):
+ """
+ What it does: CLAUDE-4.5-OPUS-HIGH → claude-opus-4.5
+ Goal: Check case insensitivity for inverted format.
+ """
+ print("Action: Normalizing 'CLAUDE-4.5-OPUS-HIGH'...")
+ result = normalize_model_name("CLAUDE-4.5-OPUS-HIGH")
+
+ print(f"Comparing result: Expected 'claude-opus-4.5', Got '{result}'")
+ assert result == "claude-opus-4.5"
+
+ # === Already normalized (passthrough) ===
+
+ def test_passthrough_already_normalized_haiku(self):
+ """
+ What it does: claude-haiku-4.5 → claude-haiku-4.5
+ Goal: Check that already normalized models are unchanged.
+ """
+ print("Action: Normalizing 'claude-haiku-4.5'...")
+ result = normalize_model_name("claude-haiku-4.5")
+
+ print(f"Comparing result: Expected 'claude-haiku-4.5', Got '{result}'")
+ assert result == "claude-haiku-4.5"
+
+ def test_passthrough_already_normalized_sonnet(self):
+ """
+ What it does: claude-sonnet-4.5 → claude-sonnet-4.5
+ Goal: Check passthrough for Sonnet.
+ """
+ print("Action: Normalizing 'claude-sonnet-4.5'...")
+ result = normalize_model_name("claude-sonnet-4.5")
+
+ print(f"Comparing result: Expected 'claude-sonnet-4.5', Got '{result}'")
+ assert result == "claude-sonnet-4.5"
+
+ def test_passthrough_auto(self):
+ """
+ What it does: auto → auto
+ Goal: Check passthrough for 'auto'.
+ """
+ print("Action: Normalizing 'auto'...")
+ result = normalize_model_name("auto")
+
+ print(f"Comparing result: Expected 'auto', Got '{result}'")
+ assert result == "auto"
+
+ # === Edge cases ===
+
+ def test_handles_empty_string(self):
+ """
+ What it does: "" → ""
+ Goal: Check empty string handling.
+ """
+ print("Action: Normalizing empty string...")
+ result = normalize_model_name("")
+
+ print(f"Comparing result: Expected '', Got '{result}'")
+ assert result == ""
+
+ def test_handles_unknown_format(self):
+ """
+ What it does: gpt-4 → gpt-4 (passthrough)
+ Goal: Check passthrough for unknown formats.
+ """
+ print("Action: Normalizing 'gpt-4'...")
+ result = normalize_model_name("gpt-4")
+
+ print(f"Comparing result: Expected 'gpt-4', Got '{result}'")
+ assert result == "gpt-4"
+
+ def test_handles_random_model_name(self):
+ """
+ What it does: some-random-model → some-random-model
+ Goal: Check passthrough for arbitrary names.
+ """
+ print("Action: Normalizing 'some-random-model'...")
+ result = normalize_model_name("some-random-model")
+
+ print(f"Comparing result: Expected 'some-random-model', Got '{result}'")
+ assert result == "some-random-model"
+
+
+# =============================================================================
+# TestNormalizeModelNameParametrized - Parametrized tests
+# =============================================================================
+
+class TestNormalizeModelNameParametrized:
+ """Parametrized tests for complete coverage of scenarios."""
+
+ @pytest.mark.parametrize("input_model,expected", [
+ # Standard format with minor version
+ ("claude-haiku-4-5", "claude-haiku-4.5"),
+ ("claude-haiku-4-5-20251001", "claude-haiku-4.5"),
+ ("claude-haiku-4-5-latest", "claude-haiku-4.5"),
+ ("claude-sonnet-4-5", "claude-sonnet-4.5"),
+ ("claude-sonnet-4-5-20250929", "claude-sonnet-4.5"),
+ ("claude-opus-4-5", "claude-opus-4.5"),
+ ("claude-opus-4-5-20251101", "claude-opus-4.5"),
+ # Without minor version
+ ("claude-sonnet-4", "claude-sonnet-4"),
+ ("claude-sonnet-4-20250514", "claude-sonnet-4"),
+ ("claude-haiku-4", "claude-haiku-4"),
+ ("claude-opus-4", "claude-opus-4"),
+ # Legacy format
+ ("claude-3-7-sonnet", "claude-3.7-sonnet"),
+ ("claude-3-7-sonnet-20250219", "claude-3.7-sonnet"),
+ ("claude-3-5-haiku", "claude-3.5-haiku"),
+ ("claude-3-0-opus", "claude-3.0-opus"),
+ # Already normalized
+ ("claude-haiku-4.5", "claude-haiku-4.5"),
+ ("claude-sonnet-4.5", "claude-sonnet-4.5"),
+ ("claude-opus-4.5", "claude-opus-4.5"),
+ ("claude-3.7-sonnet", "claude-3.7-sonnet"),
+ ("auto", "auto"),
+ # Passthrough for unknown
+ ("gpt-4", "gpt-4"),
+ ("gpt-4-turbo", "gpt-4-turbo"),
+ ("unknown-model", "unknown-model"),
+ ])
+ def test_normalize_model_name_all_scenarios(self, input_model, expected):
+ """
+ What it does: Checks all normalization scenarios.
+ Goal: Complete coverage of scenario table.
+ """
+ print(f"Action: Normalizing '{input_model}'...")
+ result = normalize_model_name(input_model)
+
+ print(f"Comparing result: Expected '{expected}', Got '{result}'")
+ assert result == expected
+
+
+# =============================================================================
+# TestExtractModelFamily - Tests for model family extraction
+# =============================================================================
+
+class TestExtractModelFamily:
+ """
+ Tests for extract_model_family() function.
+
+ Checks extraction of model family (haiku, sonnet, opus) from name.
+ """
+
+ def test_extracts_haiku_from_standard_format(self):
+ """
+ What it does: claude-haiku-4.5 → haiku
+ Goal: Check Haiku family extraction.
+ """
+ print("Action: Extracting family from 'claude-haiku-4.5'...")
+ result = extract_model_family("claude-haiku-4.5")
+
+ print(f"Comparing result: Expected 'haiku', Got '{result}'")
+ assert result == "haiku"
+
+ def test_extracts_sonnet_from_standard_format(self):
+ """
+ What it does: claude-sonnet-4.5 → sonnet
+ Goal: Check Sonnet family extraction.
+ """
+ print("Action: Extracting family from 'claude-sonnet-4.5'...")
+ result = extract_model_family("claude-sonnet-4.5")
+
+ print(f"Comparing result: Expected 'sonnet', Got '{result}'")
+ assert result == "sonnet"
+
+ def test_extracts_opus_from_standard_format(self):
+ """
+ What it does: claude-opus-4.5 → opus
+ Goal: Check Opus family extraction.
+ """
+ print("Action: Extracting family from 'claude-opus-4.5'...")
+ result = extract_model_family("claude-opus-4.5")
+
+ print(f"Comparing result: Expected 'opus', Got '{result}'")
+ assert result == "opus"
+
+ def test_extracts_sonnet_from_legacy_format(self):
+ """
+ What it does: claude-3.7-sonnet → sonnet
+ Goal: Check family extraction from legacy format.
+ """
+ print("Action: Extracting family from 'claude-3.7-sonnet'...")
+ result = extract_model_family("claude-3.7-sonnet")
+
+ print(f"Comparing result: Expected 'sonnet', Got '{result}'")
+ assert result == "sonnet"
+
+ def test_extracts_haiku_from_unnormalized(self):
+ """
+ What it does: claude-haiku-4-5-20251001 → haiku
+ Goal: Check family extraction from unnormalized name.
+ """
+ print("Action: Extracting family from 'claude-haiku-4-5-20251001'...")
+ result = extract_model_family("claude-haiku-4-5-20251001")
+
+ print(f"Comparing result: Expected 'haiku', Got '{result}'")
+ assert result == "haiku"
+
+ def test_returns_none_for_non_claude(self):
+ """
+ What it does: gpt-4 → None
+ Goal: Check None return for non-Claude models.
+ """
+ print("Action: Extracting family from 'gpt-4'...")
+ result = extract_model_family("gpt-4")
+
+ print(f"Comparing result: Expected None, Got {result}")
+ assert result is None
+
+ def test_returns_none_for_auto(self):
+ """
+ What it does: auto → None
+ Goal: Check None return for 'auto'.
+ """
+ print("Action: Extracting family from 'auto'...")
+ result = extract_model_family("auto")
+
+ print(f"Comparing result: Expected None, Got {result}")
+ assert result is None
+
+ def test_case_insensitive(self):
+ """
+ What it does: CLAUDE-HAIKU-4.5 → haiku
+ Goal: Check case insensitivity.
+ """
+ print("Action: Extracting family from 'CLAUDE-HAIKU-4.5'...")
+ result = extract_model_family("CLAUDE-HAIKU-4.5")
+
+ print(f"Comparing result: Expected 'haiku', Got '{result}'")
+ assert result == "haiku"
+
+
+# =============================================================================
+# TestGetModelIdForKiro - Tests for converter helper
+# =============================================================================
+
+class TestGetModelIdForKiro:
+ """
+ Tests for get_model_id_for_kiro() function.
+
+ Checks getting model ID for sending to Kiro API.
+ """
+
+ def test_normalizes_without_hidden_models(self):
+ """
+ What it does: Normalizes model without hidden models.
+ Goal: Check basic normalization.
+ """
+ print("Action: get_model_id_for_kiro('claude-haiku-4-5-20251001', {})...")
+ result = get_model_id_for_kiro("claude-haiku-4-5-20251001", {})
+
+ print(f"Comparing result: Expected 'claude-haiku-4.5', Got '{result}'")
+ assert result == "claude-haiku-4.5"
+
+ def test_returns_internal_id_for_hidden_model(self):
+ """
+ What it does: Returns internal ID for hidden model.
+ Goal: Check hidden model resolution.
+ """
+ hidden = {"claude-3.7-sonnet": "CLAUDE_3_7_SONNET_20250219_V1_0"}
+
+ print("Action: get_model_id_for_kiro('claude-3.7-sonnet', hidden)...")
+ result = get_model_id_for_kiro("claude-3.7-sonnet", hidden)
+
+ print(f"Comparing result: Expected 'CLAUDE_3_7_SONNET_20250219_V1_0', Got '{result}'")
+ assert result == "CLAUDE_3_7_SONNET_20250219_V1_0"
+
+ def test_normalizes_then_checks_hidden(self):
+ """
+ What it does: Normalizes first, then checks hidden.
+ Goal: Check operation order.
+ """
+ hidden = {"claude-3.7-sonnet": "CLAUDE_3_7_SONNET_20250219_V1_0"}
+
+ print("Action: get_model_id_for_kiro('claude-3-7-sonnet', hidden)...")
+ result = get_model_id_for_kiro("claude-3-7-sonnet", hidden)
+
+ print(f"Comparing result: Expected 'CLAUDE_3_7_SONNET_20250219_V1_0', Got '{result}'")
+ assert result == "CLAUDE_3_7_SONNET_20250219_V1_0"
+
+ def test_normalizes_with_date_then_checks_hidden(self):
+ """
+ What it does: Normalizes with date suffix, then checks hidden.
+ Goal: Check full normalization chain.
+ """
+ hidden = {"claude-3.7-sonnet": "CLAUDE_3_7_SONNET_20250219_V1_0"}
+
+ print("Action: get_model_id_for_kiro('claude-3-7-sonnet-20250219', hidden)...")
+ result = get_model_id_for_kiro("claude-3-7-sonnet-20250219", hidden)
+
+ print(f"Comparing result: Expected 'CLAUDE_3_7_SONNET_20250219_V1_0', Got '{result}'")
+ assert result == "CLAUDE_3_7_SONNET_20250219_V1_0"
+
+ def test_passthrough_unknown_model(self):
+ """
+ What it does: Passthrough for unknown models.
+ Goal: Check that unknown models pass through normalized.
+ """
+ print("Action: get_model_id_for_kiro('claude-unknown-model', {})...")
+ result = get_model_id_for_kiro("claude-unknown-model", {})
+
+ print(f"Comparing result: Expected 'claude-unknown-model', Got '{result}'")
+ assert result == "claude-unknown-model"
+
+
+# =============================================================================
+# TestModelResolver - Tests for ModelResolver class
+# =============================================================================
+
+class TestModelResolverInitialization:
+ """Tests for ModelResolver initialization."""
+
+ def test_init_with_cache_and_hidden_models(self, mock_model_cache, hidden_models):
+ """
+ What it does: Creates ModelResolver with cache and hidden models.
+ Goal: Check correct initialization.
+ """
+ print("Action: Creating ModelResolver...")
+ resolver = ModelResolver(cache=mock_model_cache, hidden_models=hidden_models)
+
+ print("Check: Attributes set correctly...")
+ assert resolver.cache is mock_model_cache
+ assert resolver.hidden_models == hidden_models
+
+ def test_init_with_empty_hidden_models(self, mock_model_cache):
+ """
+ What it does: Creates ModelResolver without hidden models.
+ Goal: Check work with empty dict.
+ """
+ print("Action: Creating ModelResolver without hidden models...")
+ resolver = ModelResolver(cache=mock_model_cache, hidden_models={})
+
+ print("Check: hidden_models is empty...")
+ assert resolver.hidden_models == {}
+
+ def test_init_with_none_hidden_models(self, mock_model_cache):
+ """
+ What it does: Creates ModelResolver with hidden_models=None.
+ Goal: Check default value.
+ """
+ print("Action: Creating ModelResolver with hidden_models=None...")
+ resolver = ModelResolver(cache=mock_model_cache, hidden_models=None)
+
+ print("Check: hidden_models initialized as empty dict...")
+ assert resolver.hidden_models == {}
+
+
+class TestModelResolverResolve:
+ """Tests for resolve() method of ModelResolver class."""
+
+ def test_resolve_finds_model_in_cache(self, model_resolver):
+ """
+ What it does: Finds model in cache.
+ Goal: Check Layer 2 (Dynamic Cache).
+ """
+ print("Action: Resolving 'claude-haiku-4-5'...")
+ result = model_resolver.resolve("claude-haiku-4-5")
+
+ print(f"Check result: {result}")
+ print(f"Comparing internal_id: Expected 'claude-haiku-4.5', Got '{result.internal_id}'")
+ assert result.internal_id == "claude-haiku-4.5"
+
+ print(f"Comparing source: Expected 'cache', Got '{result.source}'")
+ assert result.source == "cache"
+
+ print(f"Comparing is_verified: Expected True, Got {result.is_verified}")
+ assert result.is_verified is True
+
+ print(f"Comparing normalized: Expected 'claude-haiku-4.5', Got '{result.normalized}'")
+ assert result.normalized == "claude-haiku-4.5"
+
+ print(f"Comparing original_request: Expected 'claude-haiku-4-5', Got '{result.original_request}'")
+ assert result.original_request == "claude-haiku-4-5"
+
+ def test_resolve_finds_model_in_hidden(self, model_resolver):
+ """
+ What it does: Finds model in hidden models.
+ Goal: Check Layer 3 (Hidden Models).
+ """
+ print("Action: Resolving 'claude-3-7-sonnet'...")
+ result = model_resolver.resolve("claude-3-7-sonnet")
+
+ print(f"Check result: {result}")
+ print(f"Comparing internal_id: Expected 'CLAUDE_3_7_SONNET_20250219_V1_0', Got '{result.internal_id}'")
+ assert result.internal_id == "CLAUDE_3_7_SONNET_20250219_V1_0"
+
+ print(f"Comparing source: Expected 'hidden', Got '{result.source}'")
+ assert result.source == "hidden"
+
+ print(f"Comparing is_verified: Expected True, Got {result.is_verified}")
+ assert result.is_verified is True
+
+ def test_resolve_passthrough_for_unknown(self, model_resolver):
+ """
+ What it does: Passthrough for unknown model.
+ Goal: Check Layer 4 (Pass-through).
+ """
+ print("Action: Resolving 'claude-haiku-4-6' (does not exist)...")
+ result = model_resolver.resolve("claude-haiku-4-6")
+
+ print(f"Check result: {result}")
+ print(f"Comparing internal_id: Expected 'claude-haiku-4.6', Got '{result.internal_id}'")
+ assert result.internal_id == "claude-haiku-4.6"
+
+ print(f"Comparing source: Expected 'passthrough', Got '{result.source}'")
+ assert result.source == "passthrough"
+
+ print(f"Comparing is_verified: Expected False, Got {result.is_verified}")
+ assert result.is_verified is False
+
+ def test_resolve_normalizes_before_lookup(self, model_resolver):
+ """
+ What it does: Normalizes name before cache lookup.
+ Goal: Check Layer 1 (Normalize Name).
+ """
+ print("Action: Resolving 'claude-haiku-4-5-20251001'...")
+ result = model_resolver.resolve("claude-haiku-4-5-20251001")
+
+ print(f"Comparing normalized: Expected 'claude-haiku-4.5', Got '{result.normalized}'")
+ assert result.normalized == "claude-haiku-4.5"
+
+ print(f"Comparing source: Expected 'cache', Got '{result.source}'")
+ assert result.source == "cache"
+
+ def test_resolve_never_raises(self, model_resolver):
+ """
+ What it does: Never raises exception.
+ Goal: Check that resolve() always returns ModelResolution.
+ """
+ print("Action: Resolving strange input data...")
+
+ # Empty string
+ result1 = model_resolver.resolve("")
+ print(f"Empty string: {result1}")
+ assert isinstance(result1, ModelResolution)
+
+ # Special characters
+ result2 = model_resolver.resolve("!@#$%^&*()")
+ print(f"Special characters: {result2}")
+ assert isinstance(result2, ModelResolution)
+
+ # Very long name
+ result3 = model_resolver.resolve("a" * 1000)
+ print(f"Long name: source={result3.source}")
+ assert isinstance(result3, ModelResolution)
+
+ def test_resolve_auto_model(self, model_resolver):
+ """
+ What it does: Resolves 'auto' model.
+ Goal: Check that 'auto' is in cache.
+ """
+ print("Action: Resolving 'auto'...")
+ result = model_resolver.resolve("auto")
+
+ print(f"Comparing internal_id: Expected 'auto', Got '{result.internal_id}'")
+ assert result.internal_id == "auto"
+
+ print(f"Comparing source: Expected 'cache', Got '{result.source}'")
+ assert result.source == "cache"
+
+ def test_resolve_with_empty_cache(self, empty_model_cache, hidden_models):
+ """
+ What it does: Resolves model with empty cache.
+ Goal: Check work with only hidden models.
+ """
+ print("Setup: Creating resolver with empty cache...")
+ resolver = ModelResolver(cache=empty_model_cache, hidden_models=hidden_models)
+
+ print("Action: Resolving 'claude-3.7-sonnet'...")
+ result = resolver.resolve("claude-3.7-sonnet")
+
+ print(f"Comparing source: Expected 'hidden', Got '{result.source}'")
+ assert result.source == "hidden"
+
+ print("Action: Resolving 'claude-haiku-4.5' (not in cache)...")
+ result2 = resolver.resolve("claude-haiku-4.5")
+
+ print(f"Comparing source: Expected 'passthrough', Got '{result2.source}'")
+ assert result2.source == "passthrough"
+
+
+class TestModelResolverGetAvailableModels:
+ """Tests for get_available_models() method."""
+
+ def test_get_available_models_combines_cache_and_hidden(self, model_resolver):
+ """
+ What it does: Returns models from cache and hidden.
+ Goal: Check combining sources.
+ """
+ print("Action: Getting list of available models...")
+ models = model_resolver.get_available_models()
+
+ print(f"Received models: {models}")
+
+ # Check cache models
+ print("Check: Cache models present...")
+ assert "claude-haiku-4.5" in models
+ assert "claude-sonnet-4.5" in models
+ assert "claude-opus-4.5" in models
+ assert "auto" in models
+
+ # Check hidden models
+ print("Check: Hidden models present...")
+ assert "claude-3.7-sonnet" in models
+
+ def test_get_available_models_returns_sorted_list(self, model_resolver):
+ """
+ What it does: Returns sorted list.
+ Goal: Check sorting.
+ """
+ print("Action: Getting list of available models...")
+ models = model_resolver.get_available_models()
+
+ print(f"Received models: {models}")
+ print(f"Sorted: {sorted(models)}")
+
+ assert models == sorted(models)
+
+ def test_get_available_models_no_duplicates(self, mock_model_cache):
+ """
+ What it does: Does not return duplicates.
+ Goal: Check uniqueness.
+ """
+ # Add hidden model that already exists in cache
+ hidden = {"claude-haiku-4.5": "SOME_INTERNAL_ID"}
+ resolver = ModelResolver(cache=mock_model_cache, hidden_models=hidden)
+
+ print("Action: Getting list with potential duplicate...")
+ models = resolver.get_available_models()
+
+ print(f"Received models: {models}")
+
+ # Check uniqueness
+ assert len(models) == len(set(models))
+
+
+class TestModelResolverGetModelsByFamily:
+ """Tests for get_models_by_family() method."""
+
+ def test_get_models_by_family_haiku(self, model_resolver):
+ """
+ What it does: Returns only Haiku models.
+ Goal: Check filtering by family.
+ """
+ print("Action: Getting Haiku models...")
+ models = model_resolver.get_models_by_family("haiku")
+
+ print(f"Received models: {models}")
+
+ assert "claude-haiku-4.5" in models
+ assert "claude-sonnet-4.5" not in models
+ assert "claude-opus-4.5" not in models
+
+ def test_get_models_by_family_sonnet(self, model_resolver):
+ """
+ What it does: Returns only Sonnet models.
+ Goal: Check Sonnet filtering.
+ """
+ print("Action: Getting Sonnet models...")
+ models = model_resolver.get_models_by_family("sonnet")
+
+ print(f"Received models: {models}")
+
+ assert "claude-sonnet-4.5" in models
+ assert "claude-sonnet-4" in models
+ assert "claude-3.7-sonnet" in models # Hidden model
+ assert "claude-haiku-4.5" not in models
+
+ def test_get_models_by_family_opus(self, model_resolver):
+ """
+ What it does: Returns only Opus models.
+ Goal: Check Opus filtering.
+ """
+ print("Action: Getting Opus models...")
+ models = model_resolver.get_models_by_family("opus")
+
+ print(f"Received models: {models}")
+
+ assert "claude-opus-4.5" in models
+ assert "claude-sonnet-4.5" not in models
+
+ def test_get_models_by_family_case_insensitive(self, model_resolver):
+ """
+ What it does: Filtering is case insensitive.
+ Goal: Check case-insensitivity.
+ """
+ print("Action: Getting HAIKU models (uppercase)...")
+ models = model_resolver.get_models_by_family("HAIKU")
+
+ print(f"Received models: {models}")
+
+ assert "claude-haiku-4.5" in models
+
+
+class TestModelResolverGetSuggestionsForModel:
+ """Tests for get_suggestions_for_model() method."""
+
+ def test_get_suggestions_returns_same_family(self, model_resolver):
+ """
+ What it does: Returns models of same family.
+ Goal: Check that suggestions are from same family.
+ """
+ print("Action: Getting suggestions for 'claude-haiku-4-6'...")
+ suggestions = model_resolver.get_suggestions_for_model("claude-haiku-4-6")
+
+ print(f"Received suggestions: {suggestions}")
+
+ # All suggestions should be Haiku
+ for s in suggestions:
+ print(f"Check: '{s}' contains 'haiku'...")
+ assert "haiku" in s.lower()
+
+ def test_get_suggestions_no_cross_family(self, model_resolver):
+ """
+ What it does: NEVER suggests models from other family.
+ Goal: Critical check - Opus never becomes Sonnet!
+ """
+ print("Action: Getting suggestions for 'claude-opus-5'...")
+ suggestions = model_resolver.get_suggestions_for_model("claude-opus-5")
+
+ print(f"Received suggestions: {suggestions}")
+
+ # Should NOT be Sonnet or Haiku
+ for s in suggestions:
+ print(f"Check: '{s}' does NOT contain 'sonnet' or 'haiku'...")
+ assert "sonnet" not in s.lower()
+ assert "haiku" not in s.lower()
+
+ def test_get_suggestions_returns_all_for_unknown_family(self, model_resolver):
+ """
+ What it does: Returns all models for unknown family.
+ Goal: Check fallback for non-Claude models.
+ """
+ print("Action: Getting suggestions for 'gpt-4'...")
+ suggestions = model_resolver.get_suggestions_for_model("gpt-4")
+
+ print(f"Received suggestions: {suggestions}")
+
+ # Should be all models
+ all_models = model_resolver.get_available_models()
+ assert set(suggestions) == set(all_models)
+
+
+# =============================================================================
+# TestModelResolution - Tests for ModelResolution dataclass
+# =============================================================================
+
+class TestModelResolution:
+ """Tests for ModelResolution dataclass."""
+
+ def test_model_resolution_fields(self):
+ """
+ What it does: Checks all ModelResolution fields.
+ Goal: Ensure correct structure.
+ """
+ print("Action: Creating ModelResolution...")
+ resolution = ModelResolution(
+ internal_id="claude-haiku-4.5",
+ source="cache",
+ original_request="claude-haiku-4-5",
+ normalized="claude-haiku-4.5",
+ is_verified=True
+ )
+
+ print(f"Check fields: {resolution}")
+ assert resolution.internal_id == "claude-haiku-4.5"
+ assert resolution.source == "cache"
+ assert resolution.original_request == "claude-haiku-4-5"
+ assert resolution.normalized == "claude-haiku-4.5"
+ assert resolution.is_verified is True
+
+ def test_model_resolution_is_frozen(self):
+ """
+ What it does: Checks that ModelResolution is immutable.
+ Goal: Ensure immutability (frozen=True).
+ """
+ print("Action: Creating ModelResolution...")
+ resolution = ModelResolution(
+ internal_id="test",
+ source="cache",
+ original_request="test",
+ normalized="test",
+ is_verified=True
+ )
+
+ print("Check: Attempt to modify field should raise error...")
+ with pytest.raises(FrozenInstanceError):
+ resolution.internal_id = "changed"
+
+ def test_model_resolution_equality(self):
+ """
+ What it does: Checks comparison of two ModelResolution objects.
+ Goal: Ensure correct __eq__ implementation.
+ """
+ print("Action: Creating two identical ModelResolution objects...")
+ resolution1 = ModelResolution(
+ internal_id="test",
+ source="cache",
+ original_request="test",
+ normalized="test",
+ is_verified=True
+ )
+ resolution2 = ModelResolution(
+ internal_id="test",
+ source="cache",
+ original_request="test",
+ normalized="test",
+ is_verified=True
+ )
+
+ print(f"Comparing: {resolution1} == {resolution2}")
+ assert resolution1 == resolution2
+
+ def test_model_resolution_inequality(self):
+ """
+ What it does: Checks inequality of different ModelResolution objects.
+ Goal: Ensure correct __eq__ implementation.
+ """
+ print("Action: Creating two different ModelResolution objects...")
+ resolution1 = ModelResolution(
+ internal_id="test1",
+ source="cache",
+ original_request="test",
+ normalized="test",
+ is_verified=True
+ )
+ resolution2 = ModelResolution(
+ internal_id="test2",
+ source="hidden",
+ original_request="test",
+ normalized="test",
+ is_verified=True
+ )
+
+ print(f"Comparing: {resolution1} != {resolution2}")
+ assert resolution1 != resolution2
+
+
+# =============================================================================
+# TestModelInfoCacheNewMethods - Tests for new cache methods
+# =============================================================================
+
+class TestModelInfoCacheIsValidModel:
+ """Tests for is_valid_model() method in ModelInfoCache."""
+
+ @pytest.mark.asyncio
+ async def test_is_valid_model_returns_true_for_cached(self):
+ """
+ What it does: Returns True for model in cache.
+ Goal: Check basic functionality.
+ """
+ print("Setup: Creating and populating cache...")
+ cache = ModelInfoCache()
+ await cache.update([{"modelId": "claude-sonnet-4.5"}])
+
+ print("Action: Checking is_valid_model('claude-sonnet-4.5')...")
+ result = cache.is_valid_model("claude-sonnet-4.5")
+
+ print(f"Comparing result: Expected True, Got {result}")
+ assert result is True
+
+ @pytest.mark.asyncio
+ async def test_is_valid_model_returns_false_for_unknown(self):
+ """
+ What it does: Returns False for unknown model.
+ Goal: Check negative case.
+ """
+ print("Setup: Creating and populating cache...")
+ cache = ModelInfoCache()
+ await cache.update([{"modelId": "claude-sonnet-4.5"}])
+
+ print("Action: Checking is_valid_model('unknown-model')...")
+ result = cache.is_valid_model("unknown-model")
+
+ print(f"Comparing result: Expected False, Got {result}")
+ assert result is False
+
+ def test_is_valid_model_on_empty_cache(self):
+ """
+ What it does: Returns False for empty cache.
+ Goal: Check edge case.
+ """
+ print("Setup: Creating empty cache...")
+ cache = ModelInfoCache()
+
+ print("Action: Checking is_valid_model('any-model')...")
+ result = cache.is_valid_model("any-model")
+
+ print(f"Comparing result: Expected False, Got {result}")
+ assert result is False
+
+
+class TestModelInfoCacheAddHiddenModel:
+ """Tests for add_hidden_model() method in ModelInfoCache."""
+
+ def test_add_hidden_model_adds_to_cache(self):
+ """
+ What it does: Adds hidden model to cache.
+ Goal: Check basic functionality.
+ """
+ print("Setup: Creating empty cache...")
+ cache = ModelInfoCache()
+
+ print("Action: Adding hidden model...")
+ cache.add_hidden_model("claude-3.7-sonnet", "CLAUDE_3_7_SONNET_20250219_V1_0")
+
+ print("Check: Model added to cache...")
+ assert cache.is_valid_model("claude-3.7-sonnet") is True
+
+ def test_add_hidden_model_stores_internal_id(self):
+ """
+ What it does: Stores internal ID in _internal_id field.
+ Goal: Check data structure.
+ """
+ print("Setup: Creating empty cache...")
+ cache = ModelInfoCache()
+
+ print("Action: Adding hidden model...")
+ cache.add_hidden_model("claude-3.7-sonnet", "CLAUDE_3_7_SONNET_20250219_V1_0")
+
+ print("Check: _internal_id saved...")
+ model_info = cache.get("claude-3.7-sonnet")
+ print(f"model_info: {model_info}")
+
+ assert model_info["_internal_id"] == "CLAUDE_3_7_SONNET_20250219_V1_0"
+ assert model_info["_is_hidden"] is True
+
+ def test_add_hidden_model_sets_model_id(self):
+ """
+ What it does: Sets modelId equal to display_name.
+ Goal: Check data consistency.
+ """
+ print("Setup: Creating empty cache...")
+ cache = ModelInfoCache()
+
+ print("Action: Adding hidden model...")
+ cache.add_hidden_model("claude-3.7-sonnet", "INTERNAL_ID")
+
+ print("Check: modelId set...")
+ model_info = cache.get("claude-3.7-sonnet")
+
+ assert model_info["modelId"] == "claude-3.7-sonnet"
+ assert model_info["modelName"] == "claude-3.7-sonnet"
+
+ @pytest.mark.asyncio
+ async def test_add_hidden_model_does_not_overwrite_existing(self):
+ """
+ What it does: Does not overwrite existing model.
+ Goal: Check protection from overwriting.
+ """
+ print("Setup: Creating cache with model...")
+ cache = ModelInfoCache()
+ await cache.update([{
+ "modelId": "claude-3.7-sonnet",
+ "modelName": "Original Name",
+ "tokenLimits": {"maxInputTokens": 200000}
+ }])
+
+ print("Action: Attempting to add hidden model with same ID...")
+ cache.add_hidden_model("claude-3.7-sonnet", "NEW_INTERNAL_ID")
+
+ print("Check: Original data preserved...")
+ model_info = cache.get("claude-3.7-sonnet")
+
+ assert model_info["modelName"] == "Original Name"
+ assert "_internal_id" not in model_info # Should not be added
+
+ def test_add_hidden_model_appears_in_get_all_model_ids(self):
+ """
+ What it does: Hidden model appears in list of all models.
+ Goal: Check integration with get_all_model_ids().
+ """
+ print("Setup: Creating empty cache...")
+ cache = ModelInfoCache()
+
+ print("Action: Adding hidden model...")
+ cache.add_hidden_model("claude-3.7-sonnet", "INTERNAL_ID")
+
+ print("Check: Model in list...")
+ model_ids = cache.get_all_model_ids()
+
+ assert "claude-3.7-sonnet" in model_ids
+
+
+# =============================================================================
+# TestCriticalSafetyPrinciple - Critical security tests
+# =============================================================================
+
+class TestCriticalSafetyPrinciple:
+ """
+ Critical security tests: Family Isolation.
+
+ IMPORTANT: Resolver MUST NEVER cross model family boundaries!
+ """
+
+ def test_opus_never_becomes_sonnet(self, model_resolver):
+ """
+ What it does: Opus request NEVER becomes Sonnet.
+ Goal: Critical check for Family Isolation.
+ """
+ print("Action: Resolving non-existent Opus model...")
+ result = model_resolver.resolve("claude-opus-5")
+
+ print(f"Result: {result}")
+
+ # Should be passthrough, NOT fallback to Sonnet
+ print("Check: Does NOT contain 'sonnet'...")
+ assert "sonnet" not in result.internal_id.lower()
+
+ print("Check: Does NOT contain 'haiku'...")
+ assert "haiku" not in result.internal_id.lower()
+
+ def test_haiku_never_becomes_opus(self, model_resolver):
+ """
+ What it does: Haiku request NEVER becomes Opus.
+ Goal: Critical check for Family Isolation.
+ """
+ print("Action: Resolving non-existent Haiku model...")
+ result = model_resolver.resolve("claude-haiku-5")
+
+ print(f"Result: {result}")
+
+ # Should be passthrough, NOT fallback to Opus
+ print("Check: Does NOT contain 'opus'...")
+ assert "opus" not in result.internal_id.lower()
+
+ print("Check: Does NOT contain 'sonnet'...")
+ assert "sonnet" not in result.internal_id.lower()
+
+ def test_sonnet_never_becomes_haiku(self, model_resolver):
+ """
+ What it does: Sonnet request NEVER becomes Haiku.
+ Goal: Critical check for Family Isolation.
+ """
+ print("Action: Resolving non-existent Sonnet model...")
+ result = model_resolver.resolve("claude-sonnet-5")
+
+ print(f"Result: {result}")
+
+ # Should be passthrough, NOT fallback to Haiku
+ print("Check: Does NOT contain 'haiku'...")
+ assert "haiku" not in result.internal_id.lower()
+
+ print("Check: Does NOT contain 'opus'...")
+ assert "opus" not in result.internal_id.lower()
+
+ def test_suggestions_respect_family_boundaries(self, model_resolver):
+ """
+ What it does: Suggestions only from same family.
+ Goal: Check that get_suggestions_for_model() respects boundaries.
+ """
+ families = ["haiku", "sonnet", "opus"]
+
+ for family in families:
+ print(f"Check family: {family}...")
+ suggestions = model_resolver.get_suggestions_for_model(f"claude-{family}-99")
+
+ for suggestion in suggestions:
+ print(f" Suggestion: {suggestion}")
+ # Each suggestion must contain same family
+ assert family in suggestion.lower(), \
+ f"Suggestion '{suggestion}' not from family '{family}'!"
diff --git a/kiro-gateway/tests/unit/test_models_anthropic.py b/kiro-gateway/tests/unit/test_models_anthropic.py
new file mode 100644
index 0000000000000000000000000000000000000000..9e7ef2aadeeca91272d5015b3b9a8c7e5be286f6
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_models_anthropic.py
@@ -0,0 +1,1504 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for Anthropic Pydantic models.
+
+Comprehensive tests for all Anthropic API models:
+- Content blocks (text, image, tool_use, tool_result, thinking)
+- Image sources (base64, URL)
+- Messages and requests
+- Tools and tool choice
+- Responses and streaming events
+- Error models
+"""
+
+import pytest
+from pydantic import ValidationError
+
+from kiro.models_anthropic import (
+ # Content blocks
+ TextContentBlock,
+ ThinkingContentBlock,
+ ToolUseContentBlock,
+ ToolResultContentBlock,
+ # Image models
+ Base64ImageSource,
+ URLImageSource,
+ ImageContentBlock,
+ ContentBlock,
+ # Message models
+ AnthropicMessage,
+ # Tool models
+ AnthropicTool,
+ ToolChoiceAuto,
+ ToolChoiceAny,
+ ToolChoiceTool,
+ ToolChoice,
+ # Request models
+ SystemContentBlock,
+ AnthropicMessagesRequest,
+ # Response models
+ AnthropicUsage,
+ AnthropicMessagesResponse,
+ # Streaming models
+ MessageStartEvent,
+ ContentBlockStartEvent,
+ TextDelta,
+ ThinkingDelta,
+ InputJsonDelta,
+ ContentBlockDeltaEvent,
+ ContentBlockStopEvent,
+ MessageDeltaUsage,
+ MessageDeltaEvent,
+ MessageStopEvent,
+ PingEvent,
+ ErrorEvent,
+ # Error models
+ AnthropicErrorDetail,
+ AnthropicErrorResponse,
+)
+
+
+# Base64 1x1 pixel JPEG for testing
+TEST_IMAGE_BASE64 = "/9j/4AAQSkZJRgABAQEASABIAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/wAALCAABAAEBAREA/8QAFAABAAAAAAAAAAAAAAAAAAAACf/EABQQAQAAAAAAAAAAAAAAAAAAAAD/2gAIAQEAAD8AVN//2Q=="
+
+
+# ==================================================================================================
+# Tests for Base64ImageSource
+# ==================================================================================================
+
+class TestBase64ImageSource:
+ """Tests for Base64ImageSource Pydantic model."""
+
+ def test_valid_base64_source(self):
+ """
+ What it does: Verifies creation of valid Base64ImageSource.
+ Purpose: Ensure model accepts valid base64 image data.
+ """
+ print("Setup: Creating Base64ImageSource with valid data...")
+ source = Base64ImageSource(
+ type="base64",
+ media_type="image/jpeg",
+ data=TEST_IMAGE_BASE64
+ )
+
+ print(f"Result: {source}")
+ print(f"Comparing type: Expected 'base64', Got '{source.type}'")
+ assert source.type == "base64"
+
+ print(f"Comparing media_type: Expected 'image/jpeg', Got '{source.media_type}'")
+ assert source.media_type == "image/jpeg"
+
+ print(f"Comparing data: Expected {TEST_IMAGE_BASE64[:20]}..., Got {source.data[:20]}...")
+ assert source.data == TEST_IMAGE_BASE64
+
+ def test_type_defaults_to_base64(self):
+ """
+ What it does: Verifies that type defaults to "base64".
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating Base64ImageSource without explicit type...")
+ source = Base64ImageSource(
+ media_type="image/png",
+ data=TEST_IMAGE_BASE64
+ )
+
+ print(f"Comparing type: Expected 'base64', Got '{source.type}'")
+ assert source.type == "base64"
+
+ def test_requires_media_type(self):
+ """
+ What it does: Verifies that media_type is required.
+ Purpose: Ensure validation fails without media_type.
+ """
+ print("Setup: Attempting to create Base64ImageSource without media_type...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ Base64ImageSource(data=TEST_IMAGE_BASE64)
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "media_type" in str(exc_info.value)
+
+ def test_requires_data(self):
+ """
+ What it does: Verifies that data is required.
+ Purpose: Ensure validation fails without data.
+ """
+ print("Setup: Attempting to create Base64ImageSource without data...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ Base64ImageSource(media_type="image/jpeg")
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "data" in str(exc_info.value)
+
+ def test_accepts_various_media_types(self):
+ """
+ What it does: Verifies acceptance of various image media types.
+ Purpose: Ensure all common image formats are supported.
+ """
+ print("Setup: Testing various media types...")
+ media_types = ["image/jpeg", "image/png", "image/gif", "image/webp"]
+
+ for media_type in media_types:
+ print(f"Testing media_type: {media_type}")
+ source = Base64ImageSource(media_type=media_type, data=TEST_IMAGE_BASE64)
+ assert source.media_type == media_type
+
+ print("All media types accepted successfully")
+
+
+# ==================================================================================================
+# Tests for URLImageSource
+# ==================================================================================================
+
+class TestURLImageSource:
+ """Tests for URLImageSource Pydantic model."""
+
+ def test_valid_url_source(self):
+ """
+ What it does: Verifies creation of valid URLImageSource.
+ Purpose: Ensure model accepts valid URL.
+ """
+ print("Setup: Creating URLImageSource with valid URL...")
+ source = URLImageSource(
+ type="url",
+ url="https://example.com/image.jpg"
+ )
+
+ print(f"Result: {source}")
+ print(f"Comparing type: Expected 'url', Got '{source.type}'")
+ assert source.type == "url"
+
+ print(f"Comparing url: Expected 'https://example.com/image.jpg', Got '{source.url}'")
+ assert source.url == "https://example.com/image.jpg"
+
+ def test_type_defaults_to_url(self):
+ """
+ What it does: Verifies that type defaults to "url".
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating URLImageSource without explicit type...")
+ source = URLImageSource(url="https://example.com/image.png")
+
+ print(f"Comparing type: Expected 'url', Got '{source.type}'")
+ assert source.type == "url"
+
+ def test_requires_url(self):
+ """
+ What it does: Verifies that url is required.
+ Purpose: Ensure validation fails without url.
+ """
+ print("Setup: Attempting to create URLImageSource without url...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ URLImageSource()
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "url" in str(exc_info.value)
+
+
+# ==================================================================================================
+# Tests for ImageContentBlock
+# ==================================================================================================
+
+class TestImageContentBlock:
+ """Tests for ImageContentBlock Pydantic model."""
+
+ def test_with_base64_source(self):
+ """
+ What it does: Verifies creation of ImageContentBlock with base64 source.
+ Purpose: Ensure model accepts Base64ImageSource.
+ """
+ print("Setup: Creating ImageContentBlock with base64 source...")
+ block = ImageContentBlock(
+ type="image",
+ source=Base64ImageSource(
+ media_type="image/jpeg",
+ data=TEST_IMAGE_BASE64
+ )
+ )
+
+ print(f"Result: {block}")
+ print(f"Comparing type: Expected 'image', Got '{block.type}'")
+ assert block.type == "image"
+
+ print(f"Comparing source.type: Expected 'base64', Got '{block.source.type}'")
+ assert block.source.type == "base64"
+ assert block.source.media_type == "image/jpeg"
+
+ def test_with_url_source(self):
+ """
+ What it does: Verifies creation of ImageContentBlock with URL source.
+ Purpose: Ensure model accepts URLImageSource.
+ """
+ print("Setup: Creating ImageContentBlock with URL source...")
+ block = ImageContentBlock(
+ type="image",
+ source=URLImageSource(url="https://example.com/image.jpg")
+ )
+
+ print(f"Result: {block}")
+ print(f"Comparing type: Expected 'image', Got '{block.type}'")
+ assert block.type == "image"
+
+ print(f"Comparing source.type: Expected 'url', Got '{block.source.type}'")
+ assert block.source.type == "url"
+ assert block.source.url == "https://example.com/image.jpg"
+
+ def test_with_dict_base64_source(self):
+ """
+ What it does: Verifies creation of ImageContentBlock with dict source.
+ Purpose: Ensure model accepts dict that matches Base64ImageSource schema.
+ """
+ print("Setup: Creating ImageContentBlock with dict source...")
+ block = ImageContentBlock(
+ type="image",
+ source={
+ "type": "base64",
+ "media_type": "image/png",
+ "data": TEST_IMAGE_BASE64
+ }
+ )
+
+ print(f"Result: {block}")
+ print(f"Comparing source.type: Expected 'base64', Got '{block.source.type}'")
+ assert block.source.type == "base64"
+ assert block.source.media_type == "image/png"
+
+ def test_with_dict_url_source(self):
+ """
+ What it does: Verifies creation of ImageContentBlock with dict URL source.
+ Purpose: Ensure model accepts dict that matches URLImageSource schema.
+ """
+ print("Setup: Creating ImageContentBlock with dict URL source...")
+ block = ImageContentBlock(
+ type="image",
+ source={
+ "type": "url",
+ "url": "https://example.com/test.gif"
+ }
+ )
+
+ print(f"Result: {block}")
+ print(f"Comparing source.type: Expected 'url', Got '{block.source.type}'")
+ assert block.source.type == "url"
+ assert block.source.url == "https://example.com/test.gif"
+
+ def test_type_literal_is_image(self):
+ """
+ What it does: Verifies that type must be "image".
+ Purpose: Ensure type literal validation works.
+ """
+ print("Setup: Creating ImageContentBlock with correct type...")
+ block = ImageContentBlock(
+ source=Base64ImageSource(media_type="image/jpeg", data=TEST_IMAGE_BASE64)
+ )
+
+ print(f"Comparing type: Expected 'image', Got '{block.type}'")
+ assert block.type == "image"
+
+ def test_requires_source(self):
+ """
+ What it does: Verifies that source is required.
+ Purpose: Ensure validation fails without source.
+ """
+ print("Setup: Attempting to create ImageContentBlock without source...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ImageContentBlock(type="image")
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "source" in str(exc_info.value)
+
+
+# ==================================================================================================
+# Tests for ContentBlock Union
+# ==================================================================================================
+
+class TestContentBlockUnion:
+ """Tests for ContentBlock union type accepting ImageContentBlock."""
+
+ def test_accepts_text_content_block(self):
+ """
+ What it does: Verifies ContentBlock accepts TextContentBlock.
+ Purpose: Ensure union includes text blocks.
+ """
+ print("Setup: Creating TextContentBlock...")
+ block: ContentBlock = TextContentBlock(text="Hello, world!")
+
+ print(f"Result: {block}")
+ print(f"Comparing type: Expected 'text', Got '{block.type}'")
+ assert block.type == "text"
+ assert block.text == "Hello, world!"
+
+ def test_accepts_image_content_block(self):
+ """
+ What it does: Verifies ContentBlock accepts ImageContentBlock.
+ Purpose: Ensure union includes image blocks (Issue #30 fix).
+
+ This is the key test that verifies the fix for Issue #30.
+ Before the fix, ContentBlock union did not include ImageContentBlock,
+ causing 422 Validation Error when image content was sent.
+ """
+ print("Setup: Creating ImageContentBlock...")
+ block: ContentBlock = ImageContentBlock(
+ source=Base64ImageSource(media_type="image/jpeg", data=TEST_IMAGE_BASE64)
+ )
+
+ print(f"Result: {block}")
+ print(f"Comparing type: Expected 'image', Got '{block.type}'")
+ assert block.type == "image"
+ assert block.source.type == "base64"
+
+ def test_accepts_tool_use_content_block(self):
+ """
+ What it does: Verifies ContentBlock accepts ToolUseContentBlock.
+ Purpose: Ensure union includes tool_use blocks.
+ """
+ print("Setup: Creating ToolUseContentBlock...")
+ block: ContentBlock = ToolUseContentBlock(
+ id="call_123",
+ name="get_weather",
+ input={"location": "Moscow"}
+ )
+
+ print(f"Result: {block}")
+ print(f"Comparing type: Expected 'tool_use', Got '{block.type}'")
+ assert block.type == "tool_use"
+
+ def test_accepts_tool_result_content_block(self):
+ """
+ What it does: Verifies ContentBlock accepts ToolResultContentBlock.
+ Purpose: Ensure union includes tool_result blocks.
+ """
+ print("Setup: Creating ToolResultContentBlock...")
+ block: ContentBlock = ToolResultContentBlock(
+ tool_use_id="call_123",
+ content="Weather: Sunny, 25°C"
+ )
+
+ print(f"Result: {block}")
+ print(f"Comparing type: Expected 'tool_result', Got '{block.type}'")
+ assert block.type == "tool_result"
+
+
+# ==================================================================================================
+# Tests for AnthropicMessage with Image Content (Issue #30 fix verification)
+# ==================================================================================================
+
+class TestAnthropicMessageWithImages:
+ """
+ Tests for AnthropicMessage with image content.
+
+ These tests verify the fix for Issue #30 - 422 Validation Error
+ when sending image content blocks in messages.
+ """
+
+ def test_message_with_image_content_validates(self):
+ """
+ What it does: Verifies AnthropicMessage accepts image content blocks.
+ Purpose: This is the PRIMARY test for Issue #30 fix.
+
+ Before the fix, this would raise a ValidationError because
+ ContentBlock union did not include ImageContentBlock.
+ """
+ print("Setup: Creating AnthropicMessage with image content...")
+ message = AnthropicMessage(
+ role="user",
+ content=[
+ TextContentBlock(text="What's in this image?"),
+ ImageContentBlock(
+ source=Base64ImageSource(
+ media_type="image/jpeg",
+ data=TEST_IMAGE_BASE64
+ )
+ )
+ ]
+ )
+
+ print(f"Result: {message}")
+ print(f"Comparing role: Expected 'user', Got '{message.role}'")
+ assert message.role == "user"
+
+ print(f"Comparing content length: Expected 2, Got {len(message.content)}")
+ assert len(message.content) == 2
+
+ print(f"Comparing content[0].type: Expected 'text', Got '{message.content[0].type}'")
+ assert message.content[0].type == "text"
+
+ print(f"Comparing content[1].type: Expected 'image', Got '{message.content[1].type}'")
+ assert message.content[1].type == "image"
+
+ def test_message_with_dict_image_content_validates(self):
+ """
+ What it does: Verifies AnthropicMessage accepts dict image content.
+ Purpose: Ensure raw dict format (as received from API) validates correctly.
+
+ This is how the actual API request comes in - as raw dicts, not Pydantic models.
+ """
+ print("Setup: Creating AnthropicMessage with dict image content...")
+ message = AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "Describe this image"},
+ {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "media_type": "image/png",
+ "data": TEST_IMAGE_BASE64
+ }
+ }
+ ]
+ )
+
+ print(f"Result: {message}")
+ print(f"Comparing content length: Expected 2, Got {len(message.content)}")
+ assert len(message.content) == 2
+
+ print(f"Comparing content[1].type: Expected 'image', Got '{message.content[1].type}'")
+ assert message.content[1].type == "image"
+ assert message.content[1].source.type == "base64"
+
+ def test_message_with_multiple_images_validates(self):
+ """
+ What it does: Verifies AnthropicMessage accepts multiple images.
+ Purpose: Ensure multiple image blocks in one message work correctly.
+ """
+ print("Setup: Creating AnthropicMessage with multiple images...")
+ message = AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "Compare these images"},
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/jpeg", "data": TEST_IMAGE_BASE64}
+ },
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/png", "data": TEST_IMAGE_BASE64}
+ },
+ {
+ "type": "image",
+ "source": {"type": "base64", "media_type": "image/webp", "data": TEST_IMAGE_BASE64}
+ }
+ ]
+ )
+
+ print(f"Result content length: {len(message.content)}")
+ assert len(message.content) == 4
+
+ image_blocks = [b for b in message.content if b.type == "image"]
+ print(f"Image blocks count: {len(image_blocks)}")
+ assert len(image_blocks) == 3
+
+ def test_message_with_url_image_validates(self):
+ """
+ What it does: Verifies AnthropicMessage accepts URL image source.
+ Purpose: Ensure URL-based images are accepted (even if not fully supported).
+ """
+ print("Setup: Creating AnthropicMessage with URL image...")
+ message = AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "What's in this image?"},
+ {
+ "type": "image",
+ "source": {
+ "type": "url",
+ "url": "https://example.com/image.jpg"
+ }
+ }
+ ]
+ )
+
+ print(f"Result: {message}")
+ print(f"Comparing content[1].source.type: Expected 'url', Got '{message.content[1].source.type}'")
+ assert message.content[1].source.type == "url"
+ assert message.content[1].source.url == "https://example.com/image.jpg"
+
+
+# ==================================================================================================
+# Tests for AnthropicMessagesRequest with Image Content
+# ==================================================================================================
+
+class TestAnthropicMessagesRequestWithImages:
+ """Tests for full AnthropicMessagesRequest with image content."""
+
+ def test_request_with_image_message_validates(self):
+ """
+ What it does: Verifies full request with image content validates.
+ Purpose: End-to-end validation test for Issue #30 fix.
+
+ This simulates the actual request that was failing with 422 error.
+ """
+ print("Setup: Creating full AnthropicMessagesRequest with image...")
+ request = AnthropicMessagesRequest(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[
+ AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "What's in this image?"},
+ {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "media_type": "image/jpeg",
+ "data": TEST_IMAGE_BASE64
+ }
+ }
+ ]
+ )
+ ]
+ )
+
+ print(f"Result: {request}")
+ print(f"Comparing model: Expected 'claude-sonnet-4-5', Got '{request.model}'")
+ assert request.model == "claude-sonnet-4-5"
+
+ print(f"Comparing messages count: Expected 1, Got {len(request.messages)}")
+ assert len(request.messages) == 1
+
+ print(f"Comparing content count: Expected 2, Got {len(request.messages[0].content)}")
+ assert len(request.messages[0].content) == 2
+
+ print("Request with image content validated successfully!")
+
+ def test_request_with_conversation_including_images(self):
+ """
+ What it does: Verifies multi-turn conversation with images validates.
+ Purpose: Ensure images work in conversation context.
+ """
+ print("Setup: Creating multi-turn conversation with images...")
+ request = AnthropicMessagesRequest(
+ model="claude-sonnet-4-5",
+ max_tokens=1024,
+ messages=[
+ AnthropicMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "What's in this image?"},
+ {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "media_type": "image/jpeg",
+ "data": TEST_IMAGE_BASE64
+ }
+ }
+ ]
+ ),
+ AnthropicMessage(
+ role="assistant",
+ content="I can see a small test image."
+ ),
+ AnthropicMessage(
+ role="user",
+ content="Can you describe it in more detail?"
+ )
+ ]
+ )
+
+ print(f"Result messages count: {len(request.messages)}")
+ assert len(request.messages) == 3
+
+ # First message has image
+ assert request.messages[0].content[1].type == "image"
+
+ # Second message is string (assistant)
+ assert request.messages[1].content == "I can see a small test image."
+
+ # Third message is string (user follow-up)
+ assert request.messages[2].content == "Can you describe it in more detail?"
+
+ print("Multi-turn conversation with images validated successfully!")
+
+
+# ==================================================================================================
+# Tests for TextContentBlock
+# ==================================================================================================
+
+class TestTextContentBlock:
+ """Tests for TextContentBlock Pydantic model."""
+
+ def test_valid_text_block(self):
+ """
+ What it does: Verifies creation of valid TextContentBlock.
+ Purpose: Ensure model accepts valid text content.
+ """
+ print("Setup: Creating TextContentBlock with valid text...")
+ block = TextContentBlock(text="Hello, world!")
+
+ print(f"Result: {block}")
+ print(f"Comparing type: Expected 'text', Got '{block.type}'")
+ assert block.type == "text"
+
+ print(f"Comparing text: Expected 'Hello, world!', Got '{block.text}'")
+ assert block.text == "Hello, world!"
+
+ def test_type_defaults_to_text(self):
+ """
+ What it does: Verifies that type defaults to "text".
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating TextContentBlock without explicit type...")
+ block = TextContentBlock(text="Test")
+
+ print(f"Comparing type: Expected 'text', Got '{block.type}'")
+ assert block.type == "text"
+
+ def test_requires_text(self):
+ """
+ What it does: Verifies that text is required.
+ Purpose: Ensure validation fails without text.
+ """
+ print("Setup: Attempting to create TextContentBlock without text...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ TextContentBlock()
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "text" in str(exc_info.value)
+
+ def test_accepts_empty_string(self):
+ """
+ What it does: Verifies that empty string is accepted.
+ Purpose: Ensure empty text is valid.
+ """
+ print("Setup: Creating TextContentBlock with empty string...")
+ block = TextContentBlock(text="")
+
+ print(f"Comparing text: Expected '', Got '{block.text}'")
+ assert block.text == ""
+
+ def test_accepts_multiline_text(self):
+ """
+ What it does: Verifies that multiline text is accepted.
+ Purpose: Ensure newlines are preserved.
+ """
+ print("Setup: Creating TextContentBlock with multiline text...")
+ multiline = "Line 1\nLine 2\nLine 3"
+ block = TextContentBlock(text=multiline)
+
+ print(f"Comparing text: Expected multiline, Got '{block.text}'")
+ assert block.text == multiline
+ assert "\n" in block.text
+
+
+# ==================================================================================================
+# Tests for ThinkingContentBlock
+# ==================================================================================================
+
+class TestThinkingContentBlock:
+ """Tests for ThinkingContentBlock Pydantic model."""
+
+ def test_valid_thinking_block(self):
+ """
+ What it does: Verifies creation of valid ThinkingContentBlock.
+ Purpose: Ensure model accepts valid thinking content.
+ """
+ print("Setup: Creating ThinkingContentBlock with valid thinking...")
+ block = ThinkingContentBlock(
+ thinking="Let me analyze this step by step...",
+ signature="abc123"
+ )
+
+ print(f"Result: {block}")
+ print(f"Comparing type: Expected 'thinking', Got '{block.type}'")
+ assert block.type == "thinking"
+
+ print(f"Comparing thinking: Got '{block.thinking[:30]}...'")
+ assert block.thinking == "Let me analyze this step by step..."
+
+ print(f"Comparing signature: Expected 'abc123', Got '{block.signature}'")
+ assert block.signature == "abc123"
+
+ def test_type_defaults_to_thinking(self):
+ """
+ What it does: Verifies that type defaults to "thinking".
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating ThinkingContentBlock without explicit type...")
+ block = ThinkingContentBlock(thinking="Test thinking")
+
+ print(f"Comparing type: Expected 'thinking', Got '{block.type}'")
+ assert block.type == "thinking"
+
+ def test_signature_defaults_to_empty(self):
+ """
+ What it does: Verifies that signature defaults to empty string.
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating ThinkingContentBlock without signature...")
+ block = ThinkingContentBlock(thinking="Test")
+
+ print(f"Comparing signature: Expected '', Got '{block.signature}'")
+ assert block.signature == ""
+
+ def test_requires_thinking(self):
+ """
+ What it does: Verifies that thinking is required.
+ Purpose: Ensure validation fails without thinking.
+ """
+ print("Setup: Attempting to create ThinkingContentBlock without thinking...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ThinkingContentBlock()
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "thinking" in str(exc_info.value)
+
+
+# ==================================================================================================
+# Tests for ToolUseContentBlock
+# ==================================================================================================
+
+class TestToolUseContentBlock:
+ """Tests for ToolUseContentBlock Pydantic model."""
+
+ def test_valid_tool_use_block(self):
+ """
+ What it does: Verifies creation of valid ToolUseContentBlock.
+ Purpose: Ensure model accepts valid tool use data.
+ """
+ print("Setup: Creating ToolUseContentBlock with valid data...")
+ block = ToolUseContentBlock(
+ id="call_123",
+ name="get_weather",
+ input={"location": "Moscow", "units": "celsius"}
+ )
+
+ print(f"Result: {block}")
+ print(f"Comparing type: Expected 'tool_use', Got '{block.type}'")
+ assert block.type == "tool_use"
+
+ print(f"Comparing id: Expected 'call_123', Got '{block.id}'")
+ assert block.id == "call_123"
+
+ print(f"Comparing name: Expected 'get_weather', Got '{block.name}'")
+ assert block.name == "get_weather"
+
+ print(f"Comparing input: Got {block.input}")
+ assert block.input == {"location": "Moscow", "units": "celsius"}
+
+ def test_type_defaults_to_tool_use(self):
+ """
+ What it does: Verifies that type defaults to "tool_use".
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating ToolUseContentBlock without explicit type...")
+ block = ToolUseContentBlock(id="call_1", name="test", input={})
+
+ print(f"Comparing type: Expected 'tool_use', Got '{block.type}'")
+ assert block.type == "tool_use"
+
+ def test_requires_id(self):
+ """
+ What it does: Verifies that id is required.
+ Purpose: Ensure validation fails without id.
+ """
+ print("Setup: Attempting to create ToolUseContentBlock without id...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ToolUseContentBlock(name="test", input={})
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "id" in str(exc_info.value)
+
+ def test_requires_name(self):
+ """
+ What it does: Verifies that name is required.
+ Purpose: Ensure validation fails without name.
+ """
+ print("Setup: Attempting to create ToolUseContentBlock without name...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ToolUseContentBlock(id="call_1", input={})
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "name" in str(exc_info.value)
+
+ def test_requires_input(self):
+ """
+ What it does: Verifies that input is required.
+ Purpose: Ensure validation fails without input.
+ """
+ print("Setup: Attempting to create ToolUseContentBlock without input...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ToolUseContentBlock(id="call_1", name="test")
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "input" in str(exc_info.value)
+
+ def test_accepts_empty_input(self):
+ """
+ What it does: Verifies that empty input dict is accepted.
+ Purpose: Ensure tools without parameters work.
+ """
+ print("Setup: Creating ToolUseContentBlock with empty input...")
+ block = ToolUseContentBlock(id="call_1", name="no_params_tool", input={})
+
+ print(f"Comparing input: Expected {{}}, Got {block.input}")
+ assert block.input == {}
+
+ def test_accepts_complex_input(self):
+ """
+ What it does: Verifies that complex nested input is accepted.
+ Purpose: Ensure nested structures work.
+ """
+ print("Setup: Creating ToolUseContentBlock with complex input...")
+ complex_input = {
+ "query": "test",
+ "options": {"limit": 10, "offset": 0},
+ "filters": ["active", "recent"]
+ }
+ block = ToolUseContentBlock(id="call_1", name="search", input=complex_input)
+
+ print(f"Comparing input: Got {block.input}")
+ assert block.input == complex_input
+
+
+# ==================================================================================================
+# Tests for ToolResultContentBlock
+# ==================================================================================================
+
+class TestToolResultContentBlock:
+ """Tests for ToolResultContentBlock Pydantic model."""
+
+ def test_valid_tool_result_block(self):
+ """
+ What it does: Verifies creation of valid ToolResultContentBlock.
+ Purpose: Ensure model accepts valid tool result data.
+ """
+ print("Setup: Creating ToolResultContentBlock with valid data...")
+ block = ToolResultContentBlock(
+ tool_use_id="call_123",
+ content="Weather in Moscow: Sunny, 25°C"
+ )
+
+ print(f"Result: {block}")
+ print(f"Comparing type: Expected 'tool_result', Got '{block.type}'")
+ assert block.type == "tool_result"
+
+ print(f"Comparing tool_use_id: Expected 'call_123', Got '{block.tool_use_id}'")
+ assert block.tool_use_id == "call_123"
+
+ print(f"Comparing content: Got '{block.content}'")
+ assert block.content == "Weather in Moscow: Sunny, 25°C"
+
+ def test_type_defaults_to_tool_result(self):
+ """
+ What it does: Verifies that type defaults to "tool_result".
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating ToolResultContentBlock without explicit type...")
+ block = ToolResultContentBlock(tool_use_id="call_1")
+
+ print(f"Comparing type: Expected 'tool_result', Got '{block.type}'")
+ assert block.type == "tool_result"
+
+ def test_requires_tool_use_id(self):
+ """
+ What it does: Verifies that tool_use_id is required.
+ Purpose: Ensure validation fails without tool_use_id.
+ """
+ print("Setup: Attempting to create ToolResultContentBlock without tool_use_id...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ToolResultContentBlock(content="Result")
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "tool_use_id" in str(exc_info.value)
+
+ def test_content_is_optional(self):
+ """
+ What it does: Verifies that content is optional.
+ Purpose: Ensure tool results without content work.
+ """
+ print("Setup: Creating ToolResultContentBlock without content...")
+ block = ToolResultContentBlock(tool_use_id="call_1")
+
+ print(f"Comparing content: Expected None, Got {block.content}")
+ assert block.content is None
+
+ def test_accepts_list_content(self):
+ """
+ What it does: Verifies that list content is accepted.
+ Purpose: Ensure content can be list of TextContentBlock.
+ """
+ print("Setup: Creating ToolResultContentBlock with list content...")
+ block = ToolResultContentBlock(
+ tool_use_id="call_1",
+ content=[TextContentBlock(text="Part 1"), TextContentBlock(text="Part 2")]
+ )
+
+ print(f"Comparing content type: Expected list, Got {type(block.content)}")
+ assert isinstance(block.content, list)
+ assert len(block.content) == 2
+
+ def test_is_error_field(self):
+ """
+ What it does: Verifies that is_error field works.
+ Purpose: Ensure error results can be marked.
+ """
+ print("Setup: Creating ToolResultContentBlock with is_error=True...")
+ block = ToolResultContentBlock(
+ tool_use_id="call_1",
+ content="Error: File not found",
+ is_error=True
+ )
+
+ print(f"Comparing is_error: Expected True, Got {block.is_error}")
+ assert block.is_error is True
+
+ def test_is_error_defaults_to_none(self):
+ """
+ What it does: Verifies that is_error defaults to None.
+ Purpose: Ensure default value is correct.
+ """
+ print("Setup: Creating ToolResultContentBlock without is_error...")
+ block = ToolResultContentBlock(tool_use_id="call_1", content="Success")
+
+ print(f"Comparing is_error: Expected None, Got {block.is_error}")
+ assert block.is_error is None
+
+
+# ==================================================================================================
+# Tests for AnthropicTool
+# ==================================================================================================
+
+class TestAnthropicTool:
+ """Tests for AnthropicTool Pydantic model."""
+
+ def test_valid_tool(self):
+ """
+ What it does: Verifies creation of valid AnthropicTool.
+ Purpose: Ensure model accepts valid tool definition.
+ """
+ print("Setup: Creating AnthropicTool with valid data...")
+ tool = AnthropicTool(
+ name="get_weather",
+ description="Get weather for a location",
+ input_schema={
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "City name"}
+ },
+ "required": ["location"]
+ }
+ )
+
+ print(f"Result: {tool}")
+ print(f"Comparing name: Expected 'get_weather', Got '{tool.name}'")
+ assert tool.name == "get_weather"
+
+ print(f"Comparing description: Got '{tool.description}'")
+ assert tool.description == "Get weather for a location"
+
+ print(f"Comparing input_schema: Got {tool.input_schema}")
+ assert "properties" in tool.input_schema
+
+ def test_requires_name(self):
+ """
+ What it does: Verifies that name is required.
+ Purpose: Ensure validation fails without name.
+ """
+ print("Setup: Attempting to create AnthropicTool without name...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ AnthropicTool(input_schema={})
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "name" in str(exc_info.value)
+
+ def test_requires_input_schema(self):
+ """
+ What it does: Verifies that input_schema is required.
+ Purpose: Ensure validation fails without input_schema.
+ """
+ print("Setup: Attempting to create AnthropicTool without input_schema...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ AnthropicTool(name="test")
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "input_schema" in str(exc_info.value)
+
+ def test_description_is_optional(self):
+ """
+ What it does: Verifies that description is optional.
+ Purpose: Ensure tools without description work.
+ """
+ print("Setup: Creating AnthropicTool without description...")
+ tool = AnthropicTool(name="simple_tool", input_schema={})
+
+ print(f"Comparing description: Expected None, Got {tool.description}")
+ assert tool.description is None
+
+
+# ==================================================================================================
+# Tests for ToolChoice models
+# ==================================================================================================
+
+class TestToolChoiceModels:
+ """Tests for ToolChoice Pydantic models."""
+
+ def test_tool_choice_auto(self):
+ """
+ What it does: Verifies creation of ToolChoiceAuto.
+ Purpose: Ensure auto tool choice works.
+ """
+ print("Setup: Creating ToolChoiceAuto...")
+ choice = ToolChoiceAuto()
+
+ print(f"Result: {choice}")
+ print(f"Comparing type: Expected 'auto', Got '{choice.type}'")
+ assert choice.type == "auto"
+
+ def test_tool_choice_any(self):
+ """
+ What it does: Verifies creation of ToolChoiceAny.
+ Purpose: Ensure any tool choice works.
+ """
+ print("Setup: Creating ToolChoiceAny...")
+ choice = ToolChoiceAny()
+
+ print(f"Result: {choice}")
+ print(f"Comparing type: Expected 'any', Got '{choice.type}'")
+ assert choice.type == "any"
+
+ def test_tool_choice_tool(self):
+ """
+ What it does: Verifies creation of ToolChoiceTool.
+ Purpose: Ensure specific tool choice works.
+ """
+ print("Setup: Creating ToolChoiceTool...")
+ choice = ToolChoiceTool(name="get_weather")
+
+ print(f"Result: {choice}")
+ print(f"Comparing type: Expected 'tool', Got '{choice.type}'")
+ assert choice.type == "tool"
+
+ print(f"Comparing name: Expected 'get_weather', Got '{choice.name}'")
+ assert choice.name == "get_weather"
+
+ def test_tool_choice_tool_requires_name(self):
+ """
+ What it does: Verifies that ToolChoiceTool requires name.
+ Purpose: Ensure validation fails without name.
+ """
+ print("Setup: Attempting to create ToolChoiceTool without name...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ToolChoiceTool()
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "name" in str(exc_info.value)
+
+
+# ==================================================================================================
+# Tests for SystemContentBlock
+# ==================================================================================================
+
+class TestSystemContentBlock:
+ """Tests for SystemContentBlock Pydantic model."""
+
+ def test_valid_system_block(self):
+ """
+ What it does: Verifies creation of valid SystemContentBlock.
+ Purpose: Ensure model accepts valid system content.
+ """
+ print("Setup: Creating SystemContentBlock with valid data...")
+ block = SystemContentBlock(text="You are a helpful assistant.")
+
+ print(f"Result: {block}")
+ print(f"Comparing type: Expected 'text', Got '{block.type}'")
+ assert block.type == "text"
+
+ print(f"Comparing text: Got '{block.text}'")
+ assert block.text == "You are a helpful assistant."
+
+ def test_with_cache_control(self):
+ """
+ What it does: Verifies SystemContentBlock with cache_control.
+ Purpose: Ensure prompt caching format works.
+ """
+ print("Setup: Creating SystemContentBlock with cache_control...")
+ block = SystemContentBlock(
+ text="You are helpful.",
+ cache_control={"type": "ephemeral"}
+ )
+
+ print(f"Result: {block}")
+ print(f"Comparing cache_control: Got {block.cache_control}")
+ assert block.cache_control == {"type": "ephemeral"}
+
+ def test_cache_control_is_optional(self):
+ """
+ What it does: Verifies that cache_control is optional.
+ Purpose: Ensure blocks without cache_control work.
+ """
+ print("Setup: Creating SystemContentBlock without cache_control...")
+ block = SystemContentBlock(text="Test")
+
+ print(f"Comparing cache_control: Expected None, Got {block.cache_control}")
+ assert block.cache_control is None
+
+ def test_requires_text(self):
+ """
+ What it does: Verifies that text is required.
+ Purpose: Ensure validation fails without text.
+ """
+ print("Setup: Attempting to create SystemContentBlock without text...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ SystemContentBlock()
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "text" in str(exc_info.value)
+
+
+# ==================================================================================================
+# Tests for AnthropicUsage
+# ==================================================================================================
+
+class TestAnthropicUsage:
+ """Tests for AnthropicUsage Pydantic model."""
+
+ def test_valid_usage(self):
+ """
+ What it does: Verifies creation of valid AnthropicUsage.
+ Purpose: Ensure model accepts valid usage data.
+ """
+ print("Setup: Creating AnthropicUsage with valid data...")
+ usage = AnthropicUsage(input_tokens=100, output_tokens=50)
+
+ print(f"Result: {usage}")
+ print(f"Comparing input_tokens: Expected 100, Got {usage.input_tokens}")
+ assert usage.input_tokens == 100
+
+ print(f"Comparing output_tokens: Expected 50, Got {usage.output_tokens}")
+ assert usage.output_tokens == 50
+
+ def test_requires_input_tokens(self):
+ """
+ What it does: Verifies that input_tokens is required.
+ Purpose: Ensure validation fails without input_tokens.
+ """
+ print("Setup: Attempting to create AnthropicUsage without input_tokens...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ AnthropicUsage(output_tokens=50)
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "input_tokens" in str(exc_info.value)
+
+ def test_requires_output_tokens(self):
+ """
+ What it does: Verifies that output_tokens is required.
+ Purpose: Ensure validation fails without output_tokens.
+ """
+ print("Setup: Attempting to create AnthropicUsage without output_tokens...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ AnthropicUsage(input_tokens=100)
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "output_tokens" in str(exc_info.value)
+
+
+# ==================================================================================================
+# Tests for AnthropicMessagesResponse
+# ==================================================================================================
+
+class TestAnthropicMessagesResponse:
+ """Tests for AnthropicMessagesResponse Pydantic model."""
+
+ def test_valid_response(self):
+ """
+ What it does: Verifies creation of valid AnthropicMessagesResponse.
+ Purpose: Ensure model accepts valid response data.
+ """
+ print("Setup: Creating AnthropicMessagesResponse with valid data...")
+ response = AnthropicMessagesResponse(
+ id="msg_123",
+ model="claude-sonnet-4-5",
+ content=[TextContentBlock(text="Hello!")],
+ usage=AnthropicUsage(input_tokens=10, output_tokens=5)
+ )
+
+ print(f"Result: {response}")
+ print(f"Comparing id: Expected 'msg_123', Got '{response.id}'")
+ assert response.id == "msg_123"
+
+ print(f"Comparing type: Expected 'message', Got '{response.type}'")
+ assert response.type == "message"
+
+ print(f"Comparing role: Expected 'assistant', Got '{response.role}'")
+ assert response.role == "assistant"
+
+ print(f"Comparing model: Expected 'claude-sonnet-4-5', Got '{response.model}'")
+ assert response.model == "claude-sonnet-4-5"
+
+ def test_stop_reason_values(self):
+ """
+ What it does: Verifies that stop_reason accepts valid values.
+ Purpose: Ensure all stop reasons work.
+ """
+ print("Setup: Testing various stop_reason values...")
+ stop_reasons = ["end_turn", "max_tokens", "stop_sequence", "tool_use"]
+
+ for reason in stop_reasons:
+ print(f"Testing stop_reason: {reason}")
+ response = AnthropicMessagesResponse(
+ id="msg_1",
+ model="claude-sonnet-4-5",
+ content=[TextContentBlock(text="Test")],
+ usage=AnthropicUsage(input_tokens=1, output_tokens=1),
+ stop_reason=reason
+ )
+ assert response.stop_reason == reason
+
+ print("All stop_reason values accepted successfully")
+
+ def test_stop_reason_is_optional(self):
+ """
+ What it does: Verifies that stop_reason is optional.
+ Purpose: Ensure responses without stop_reason work.
+ """
+ print("Setup: Creating response without stop_reason...")
+ response = AnthropicMessagesResponse(
+ id="msg_1",
+ model="claude-sonnet-4-5",
+ content=[TextContentBlock(text="Test")],
+ usage=AnthropicUsage(input_tokens=1, output_tokens=1)
+ )
+
+ print(f"Comparing stop_reason: Expected None, Got {response.stop_reason}")
+ assert response.stop_reason is None
+
+
+# ==================================================================================================
+# Tests for Streaming Event Models
+# ==================================================================================================
+
+class TestStreamingEvents:
+ """Tests for streaming event Pydantic models."""
+
+ def test_message_start_event(self):
+ """
+ What it does: Verifies creation of MessageStartEvent.
+ Purpose: Ensure message_start event works.
+ """
+ print("Setup: Creating MessageStartEvent...")
+ event = MessageStartEvent(
+ message={"id": "msg_1", "type": "message", "role": "assistant"}
+ )
+
+ print(f"Result: {event}")
+ print(f"Comparing type: Expected 'message_start', Got '{event.type}'")
+ assert event.type == "message_start"
+ assert event.message["id"] == "msg_1"
+
+ def test_content_block_start_event(self):
+ """
+ What it does: Verifies creation of ContentBlockStartEvent.
+ Purpose: Ensure content_block_start event works.
+ """
+ print("Setup: Creating ContentBlockStartEvent...")
+ event = ContentBlockStartEvent(
+ index=0,
+ content_block={"type": "text", "text": ""}
+ )
+
+ print(f"Result: {event}")
+ print(f"Comparing type: Expected 'content_block_start', Got '{event.type}'")
+ assert event.type == "content_block_start"
+ assert event.index == 0
+
+ def test_text_delta(self):
+ """
+ What it does: Verifies creation of TextDelta.
+ Purpose: Ensure text_delta works.
+ """
+ print("Setup: Creating TextDelta...")
+ delta = TextDelta(text="Hello")
+
+ print(f"Result: {delta}")
+ print(f"Comparing type: Expected 'text_delta', Got '{delta.type}'")
+ assert delta.type == "text_delta"
+ assert delta.text == "Hello"
+
+ def test_thinking_delta(self):
+ """
+ What it does: Verifies creation of ThinkingDelta.
+ Purpose: Ensure thinking_delta works.
+ """
+ print("Setup: Creating ThinkingDelta...")
+ delta = ThinkingDelta(thinking="Let me think...")
+
+ print(f"Result: {delta}")
+ print(f"Comparing type: Expected 'thinking_delta', Got '{delta.type}'")
+ assert delta.type == "thinking_delta"
+ assert delta.thinking == "Let me think..."
+
+ def test_input_json_delta(self):
+ """
+ What it does: Verifies creation of InputJsonDelta.
+ Purpose: Ensure input_json_delta works.
+ """
+ print("Setup: Creating InputJsonDelta...")
+ delta = InputJsonDelta(partial_json='{"loc')
+
+ print(f"Result: {delta}")
+ print(f"Comparing type: Expected 'input_json_delta', Got '{delta.type}'")
+ assert delta.type == "input_json_delta"
+ assert delta.partial_json == '{"loc'
+
+ def test_content_block_delta_event(self):
+ """
+ What it does: Verifies creation of ContentBlockDeltaEvent.
+ Purpose: Ensure content_block_delta event works.
+ """
+ print("Setup: Creating ContentBlockDeltaEvent...")
+ event = ContentBlockDeltaEvent(
+ index=0,
+ delta=TextDelta(text="Hello")
+ )
+
+ print(f"Result: {event}")
+ print(f"Comparing type: Expected 'content_block_delta', Got '{event.type}'")
+ assert event.type == "content_block_delta"
+ assert event.index == 0
+
+ def test_content_block_stop_event(self):
+ """
+ What it does: Verifies creation of ContentBlockStopEvent.
+ Purpose: Ensure content_block_stop event works.
+ """
+ print("Setup: Creating ContentBlockStopEvent...")
+ event = ContentBlockStopEvent(index=0)
+
+ print(f"Result: {event}")
+ print(f"Comparing type: Expected 'content_block_stop', Got '{event.type}'")
+ assert event.type == "content_block_stop"
+ assert event.index == 0
+
+ def test_message_delta_event(self):
+ """
+ What it does: Verifies creation of MessageDeltaEvent.
+ Purpose: Ensure message_delta event works.
+ """
+ print("Setup: Creating MessageDeltaEvent...")
+ event = MessageDeltaEvent(
+ delta={"stop_reason": "end_turn"},
+ usage=MessageDeltaUsage(output_tokens=10)
+ )
+
+ print(f"Result: {event}")
+ print(f"Comparing type: Expected 'message_delta', Got '{event.type}'")
+ assert event.type == "message_delta"
+ assert event.delta["stop_reason"] == "end_turn"
+
+ def test_message_stop_event(self):
+ """
+ What it does: Verifies creation of MessageStopEvent.
+ Purpose: Ensure message_stop event works.
+ """
+ print("Setup: Creating MessageStopEvent...")
+ event = MessageStopEvent()
+
+ print(f"Result: {event}")
+ print(f"Comparing type: Expected 'message_stop', Got '{event.type}'")
+ assert event.type == "message_stop"
+
+ def test_ping_event(self):
+ """
+ What it does: Verifies creation of PingEvent.
+ Purpose: Ensure ping event works.
+ """
+ print("Setup: Creating PingEvent...")
+ event = PingEvent()
+
+ print(f"Result: {event}")
+ print(f"Comparing type: Expected 'ping', Got '{event.type}'")
+ assert event.type == "ping"
+
+ def test_error_event(self):
+ """
+ What it does: Verifies creation of ErrorEvent.
+ Purpose: Ensure error event works.
+ """
+ print("Setup: Creating ErrorEvent...")
+ event = ErrorEvent(error={"type": "invalid_request", "message": "Bad request"})
+
+ print(f"Result: {event}")
+ print(f"Comparing type: Expected 'error', Got '{event.type}'")
+ assert event.type == "error"
+ assert event.error["type"] == "invalid_request"
+
+
+# ==================================================================================================
+# Tests for Error Models
+# ==================================================================================================
+
+class TestErrorModels:
+ """Tests for error Pydantic models."""
+
+ def test_anthropic_error_detail(self):
+ """
+ What it does: Verifies creation of AnthropicErrorDetail.
+ Purpose: Ensure error detail model works.
+ """
+ print("Setup: Creating AnthropicErrorDetail...")
+ detail = AnthropicErrorDetail(
+ type="invalid_request_error",
+ message="Invalid API key"
+ )
+
+ print(f"Result: {detail}")
+ print(f"Comparing type: Expected 'invalid_request_error', Got '{detail.type}'")
+ assert detail.type == "invalid_request_error"
+
+ print(f"Comparing message: Got '{detail.message}'")
+ assert detail.message == "Invalid API key"
+
+ def test_anthropic_error_response(self):
+ """
+ What it does: Verifies creation of AnthropicErrorResponse.
+ Purpose: Ensure error response model works.
+ """
+ print("Setup: Creating AnthropicErrorResponse...")
+ response = AnthropicErrorResponse(
+ error=AnthropicErrorDetail(
+ type="authentication_error",
+ message="Invalid API key provided"
+ )
+ )
+
+ print(f"Result: {response}")
+ print(f"Comparing type: Expected 'error', Got '{response.type}'")
+ assert response.type == "error"
+
+ print(f"Comparing error.type: Got '{response.error.type}'")
+ assert response.error.type == "authentication_error"
diff --git a/kiro-gateway/tests/unit/test_models_openai.py b/kiro-gateway/tests/unit/test_models_openai.py
new file mode 100644
index 0000000000000000000000000000000000000000..9f769badc0bf1d64283b0f6c5146df8d363e3ee9
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_models_openai.py
@@ -0,0 +1,1056 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for OpenAI Pydantic models.
+
+Comprehensive tests for all OpenAI-compatible API models:
+- Model listing (OpenAIModel, ModelList)
+- Chat messages (ChatMessage)
+- Tools (ToolFunction, Tool)
+- Requests (ChatCompletionRequest)
+- Responses (ChatCompletionChoice, ChatCompletionUsage, ChatCompletionResponse)
+- Streaming (ChatCompletionChunk, ChatCompletionChunkChoice, ChatCompletionChunkDelta)
+"""
+
+import pytest
+from pydantic import ValidationError
+
+from kiro.models_openai import (
+ # Model listing
+ OpenAIModel,
+ ModelList,
+ # Chat messages
+ ChatMessage,
+ # Tools
+ ToolFunction,
+ Tool,
+ # Requests
+ ChatCompletionRequest,
+ # Responses
+ ChatCompletionChoice,
+ ChatCompletionUsage,
+ ChatCompletionResponse,
+ # Streaming
+ ChatCompletionChunkDelta,
+ ChatCompletionChunkChoice,
+ ChatCompletionChunk,
+)
+
+
+# ==================================================================================================
+# Tests for OpenAIModel
+# ==================================================================================================
+
+class TestOpenAIModel:
+ """Tests for OpenAIModel Pydantic model."""
+
+ def test_valid_model(self):
+ """
+ What it does: Verifies creation of valid OpenAIModel.
+ Purpose: Ensure model accepts valid data.
+ """
+ print("Setup: Creating OpenAIModel with valid data...")
+ model = OpenAIModel(
+ id="claude-sonnet-4-5",
+ description="Claude Sonnet 4.5 model"
+ )
+
+ print(f"Result: {model}")
+ print(f"Comparing id: Expected 'claude-sonnet-4-5', Got '{model.id}'")
+ assert model.id == "claude-sonnet-4-5"
+
+ print(f"Comparing object: Expected 'model', Got '{model.object}'")
+ assert model.object == "model"
+
+ print(f"Comparing owned_by: Expected 'anthropic', Got '{model.owned_by}'")
+ assert model.owned_by == "anthropic"
+
+ print(f"Comparing description: Got '{model.description}'")
+ assert model.description == "Claude Sonnet 4.5 model"
+
+ def test_requires_id(self):
+ """
+ What it does: Verifies that id is required.
+ Purpose: Ensure validation fails without id.
+ """
+ print("Setup: Attempting to create OpenAIModel without id...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ OpenAIModel()
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "id" in str(exc_info.value)
+
+ def test_object_defaults_to_model(self):
+ """
+ What it does: Verifies that object defaults to "model".
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating OpenAIModel without explicit object...")
+ model = OpenAIModel(id="test-model")
+
+ print(f"Comparing object: Expected 'model', Got '{model.object}'")
+ assert model.object == "model"
+
+ def test_owned_by_defaults_to_anthropic(self):
+ """
+ What it does: Verifies that owned_by defaults to "anthropic".
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating OpenAIModel without explicit owned_by...")
+ model = OpenAIModel(id="test-model")
+
+ print(f"Comparing owned_by: Expected 'anthropic', Got '{model.owned_by}'")
+ assert model.owned_by == "anthropic"
+
+ def test_created_is_auto_generated(self):
+ """
+ What it does: Verifies that created timestamp is auto-generated.
+ Purpose: Ensure timestamp is set automatically.
+ """
+ print("Setup: Creating OpenAIModel without explicit created...")
+ model = OpenAIModel(id="test-model")
+
+ print(f"Comparing created: Got {model.created}")
+ assert model.created > 0
+ assert isinstance(model.created, int)
+
+ def test_description_is_optional(self):
+ """
+ What it does: Verifies that description is optional.
+ Purpose: Ensure models without description work.
+ """
+ print("Setup: Creating OpenAIModel without description...")
+ model = OpenAIModel(id="test-model")
+
+ print(f"Comparing description: Expected None, Got {model.description}")
+ assert model.description is None
+
+
+# ==================================================================================================
+# Tests for ModelList
+# ==================================================================================================
+
+class TestModelList:
+ """Tests for ModelList Pydantic model."""
+
+ def test_valid_model_list(self):
+ """
+ What it does: Verifies creation of valid ModelList.
+ Purpose: Ensure model list accepts valid data.
+ """
+ print("Setup: Creating ModelList with valid data...")
+ model_list = ModelList(
+ data=[
+ OpenAIModel(id="claude-sonnet-4-5"),
+ OpenAIModel(id="claude-opus-4")
+ ]
+ )
+
+ print(f"Result: {model_list}")
+ print(f"Comparing object: Expected 'list', Got '{model_list.object}'")
+ assert model_list.object == "list"
+
+ print(f"Comparing data length: Expected 2, Got {len(model_list.data)}")
+ assert len(model_list.data) == 2
+
+ def test_object_defaults_to_list(self):
+ """
+ What it does: Verifies that object defaults to "list".
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating ModelList without explicit object...")
+ model_list = ModelList(data=[])
+
+ print(f"Comparing object: Expected 'list', Got '{model_list.object}'")
+ assert model_list.object == "list"
+
+ def test_requires_data(self):
+ """
+ What it does: Verifies that data is required.
+ Purpose: Ensure validation fails without data.
+ """
+ print("Setup: Attempting to create ModelList without data...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ModelList()
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "data" in str(exc_info.value)
+
+ def test_accepts_empty_list(self):
+ """
+ What it does: Verifies that empty list is accepted.
+ Purpose: Ensure empty model list works.
+ """
+ print("Setup: Creating ModelList with empty data...")
+ model_list = ModelList(data=[])
+
+ print(f"Comparing data: Expected [], Got {model_list.data}")
+ assert model_list.data == []
+
+
+# ==================================================================================================
+# Tests for ChatMessage
+# ==================================================================================================
+
+class TestChatMessage:
+ """Tests for ChatMessage Pydantic model."""
+
+ def test_valid_user_message(self):
+ """
+ What it does: Verifies creation of valid user message.
+ Purpose: Ensure model accepts valid user message.
+ """
+ print("Setup: Creating ChatMessage with user role...")
+ message = ChatMessage(role="user", content="Hello!")
+
+ print(f"Result: {message}")
+ print(f"Comparing role: Expected 'user', Got '{message.role}'")
+ assert message.role == "user"
+
+ print(f"Comparing content: Expected 'Hello!', Got '{message.content}'")
+ assert message.content == "Hello!"
+
+ def test_valid_assistant_message(self):
+ """
+ What it does: Verifies creation of valid assistant message.
+ Purpose: Ensure model accepts valid assistant message.
+ """
+ print("Setup: Creating ChatMessage with assistant role...")
+ message = ChatMessage(role="assistant", content="Hi there!")
+
+ print(f"Result: {message}")
+ print(f"Comparing role: Expected 'assistant', Got '{message.role}'")
+ assert message.role == "assistant"
+
+ def test_valid_system_message(self):
+ """
+ What it does: Verifies creation of valid system message.
+ Purpose: Ensure model accepts valid system message.
+ """
+ print("Setup: Creating ChatMessage with system role...")
+ message = ChatMessage(role="system", content="You are helpful.")
+
+ print(f"Result: {message}")
+ print(f"Comparing role: Expected 'system', Got '{message.role}'")
+ assert message.role == "system"
+
+ def test_valid_tool_message(self):
+ """
+ What it does: Verifies creation of valid tool message.
+ Purpose: Ensure model accepts valid tool message.
+ """
+ print("Setup: Creating ChatMessage with tool role...")
+ message = ChatMessage(
+ role="tool",
+ content="Tool result",
+ tool_call_id="call_123"
+ )
+
+ print(f"Result: {message}")
+ print(f"Comparing role: Expected 'tool', Got '{message.role}'")
+ assert message.role == "tool"
+
+ print(f"Comparing tool_call_id: Expected 'call_123', Got '{message.tool_call_id}'")
+ assert message.tool_call_id == "call_123"
+
+ def test_requires_role(self):
+ """
+ What it does: Verifies that role is required.
+ Purpose: Ensure validation fails without role.
+ """
+ print("Setup: Attempting to create ChatMessage without role...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ChatMessage(content="Hello")
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "role" in str(exc_info.value)
+
+ def test_content_is_optional(self):
+ """
+ What it does: Verifies that content is optional.
+ Purpose: Ensure messages without content work (e.g., tool calls).
+ """
+ print("Setup: Creating ChatMessage without content...")
+ message = ChatMessage(role="assistant")
+
+ print(f"Comparing content: Expected None, Got {message.content}")
+ assert message.content is None
+
+ def test_accepts_list_content(self):
+ """
+ What it does: Verifies that list content is accepted.
+ Purpose: Ensure multimodal content works.
+ """
+ print("Setup: Creating ChatMessage with list content...")
+ message = ChatMessage(
+ role="user",
+ content=[
+ {"type": "text", "text": "What's in this image?"},
+ {"type": "image_url", "image_url": {"url": "https://example.com/img.jpg"}}
+ ]
+ )
+
+ print(f"Result: {message}")
+ print(f"Comparing content type: Expected list, Got {type(message.content)}")
+ assert isinstance(message.content, list)
+ assert len(message.content) == 2
+
+ def test_accepts_tool_calls(self):
+ """
+ What it does: Verifies that tool_calls is accepted.
+ Purpose: Ensure assistant messages with tool calls work.
+ """
+ print("Setup: Creating ChatMessage with tool_calls...")
+ message = ChatMessage(
+ role="assistant",
+ content="I'll call a tool",
+ tool_calls=[{
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "get_weather", "arguments": '{"location": "Moscow"}'}
+ }]
+ )
+
+ print(f"Result: {message}")
+ print(f"Comparing tool_calls: Got {message.tool_calls}")
+ assert message.tool_calls is not None
+ assert len(message.tool_calls) == 1
+
+ def test_name_is_optional(self):
+ """
+ What it does: Verifies that name is optional.
+ Purpose: Ensure messages without name work.
+ """
+ print("Setup: Creating ChatMessage without name...")
+ message = ChatMessage(role="user", content="Hello")
+
+ print(f"Comparing name: Expected None, Got {message.name}")
+ assert message.name is None
+
+ def test_accepts_name(self):
+ """
+ What it does: Verifies that name is accepted.
+ Purpose: Ensure named messages work.
+ """
+ print("Setup: Creating ChatMessage with name...")
+ message = ChatMessage(role="user", content="Hello", name="John")
+
+ print(f"Comparing name: Expected 'John', Got '{message.name}'")
+ assert message.name == "John"
+
+ def test_extra_fields_allowed(self):
+ """
+ What it does: Verifies that extra fields are allowed.
+ Purpose: Ensure model_config extra="allow" works.
+ """
+ print("Setup: Creating ChatMessage with extra field...")
+ message = ChatMessage(role="user", content="Hello", custom_field="value")
+
+ print(f"Comparing custom_field: Got '{message.custom_field}'")
+ assert message.custom_field == "value"
+
+
+# ==================================================================================================
+# Tests for ToolFunction
+# ==================================================================================================
+
+class TestToolFunction:
+ """Tests for ToolFunction Pydantic model."""
+
+ def test_valid_tool_function(self):
+ """
+ What it does: Verifies creation of valid ToolFunction.
+ Purpose: Ensure model accepts valid tool function.
+ """
+ print("Setup: Creating ToolFunction with valid data...")
+ func = ToolFunction(
+ name="get_weather",
+ description="Get weather for a location",
+ parameters={
+ "type": "object",
+ "properties": {"location": {"type": "string"}}
+ }
+ )
+
+ print(f"Result: {func}")
+ print(f"Comparing name: Expected 'get_weather', Got '{func.name}'")
+ assert func.name == "get_weather"
+
+ print(f"Comparing description: Got '{func.description}'")
+ assert func.description == "Get weather for a location"
+
+ print(f"Comparing parameters: Got {func.parameters}")
+ assert "properties" in func.parameters
+
+ def test_requires_name(self):
+ """
+ What it does: Verifies that name is required.
+ Purpose: Ensure validation fails without name.
+ """
+ print("Setup: Attempting to create ToolFunction without name...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ToolFunction(description="Test")
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "name" in str(exc_info.value)
+
+ def test_description_is_optional(self):
+ """
+ What it does: Verifies that description is optional.
+ Purpose: Ensure functions without description work.
+ """
+ print("Setup: Creating ToolFunction without description...")
+ func = ToolFunction(name="test_func")
+
+ print(f"Comparing description: Expected None, Got {func.description}")
+ assert func.description is None
+
+ def test_parameters_is_optional(self):
+ """
+ What it does: Verifies that parameters is optional.
+ Purpose: Ensure functions without parameters work.
+ """
+ print("Setup: Creating ToolFunction without parameters...")
+ func = ToolFunction(name="no_params_func")
+
+ print(f"Comparing parameters: Expected None, Got {func.parameters}")
+ assert func.parameters is None
+
+
+# ==================================================================================================
+# Tests for Tool
+# ==================================================================================================
+
+class TestTool:
+ """Tests for Tool Pydantic model."""
+
+ def test_valid_tool(self):
+ """
+ What it does: Verifies creation of valid Tool.
+ Purpose: Ensure model accepts valid tool.
+ """
+ print("Setup: Creating Tool with valid data...")
+ tool = Tool(
+ type="function",
+ function=ToolFunction(
+ name="get_weather",
+ description="Get weather",
+ parameters={}
+ )
+ )
+
+ print(f"Result: {tool}")
+ print(f"Comparing type: Expected 'function', Got '{tool.type}'")
+ assert tool.type == "function"
+
+ print(f"Comparing function.name: Expected 'get_weather', Got '{tool.function.name}'")
+ assert tool.function.name == "get_weather"
+
+ def test_type_defaults_to_function(self):
+ """
+ What it does: Verifies that type defaults to "function".
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating Tool without explicit type...")
+ tool = Tool(function=ToolFunction(name="test"))
+
+ print(f"Comparing type: Expected 'function', Got '{tool.type}'")
+ assert tool.type == "function"
+
+ def test_function_is_optional_for_flat_format(self):
+ """
+ What it does: Verifies that function is optional (for flat format compatibility).
+ Purpose: Ensure flat format (Cursor-style) is supported without function field.
+ """
+ print("Setup: Creating Tool with flat format (name, description, input_schema)...")
+ tool = Tool(
+ type="function",
+ name="test_tool",
+ description="A test tool",
+ input_schema={"type": "object", "properties": {}}
+ )
+
+ print(f"Result: {tool}")
+ print(f"Comparing name: Expected 'test_tool', Got '{tool.name}'")
+ assert tool.name == "test_tool"
+
+ print(f"Comparing function: Expected None, Got {tool.function}")
+ assert tool.function is None
+
+ print(f"Comparing description: Expected 'A test tool', Got '{tool.description}'")
+ assert tool.description == "A test tool"
+
+ def test_standard_format_still_works(self):
+ """
+ What it does: Verifies that standard OpenAI format still works.
+ Purpose: Ensure backward compatibility with standard format.
+ """
+ print("Setup: Creating Tool with standard OpenAI format (function field)...")
+ tool = Tool(
+ type="function",
+ function=ToolFunction(name="standard_tool", description="Standard")
+ )
+
+ print(f"Result: {tool}")
+ print(f"Comparing function.name: Expected 'standard_tool', Got '{tool.function.name}'")
+ assert tool.function.name == "standard_tool"
+
+ print(f"Comparing name: Expected None, Got {tool.name}")
+ assert tool.name is None
+
+
+# ==================================================================================================
+# Tests for ChatCompletionRequest
+# ==================================================================================================
+
+class TestChatCompletionRequest:
+ """Tests for ChatCompletionRequest Pydantic model."""
+
+ def test_valid_request(self):
+ """
+ What it does: Verifies creation of valid ChatCompletionRequest.
+ Purpose: Ensure model accepts valid request.
+ """
+ print("Setup: Creating ChatCompletionRequest with valid data...")
+ request = ChatCompletionRequest(
+ model="claude-sonnet-4-5",
+ messages=[ChatMessage(role="user", content="Hello")]
+ )
+
+ print(f"Result: {request}")
+ print(f"Comparing model: Expected 'claude-sonnet-4-5', Got '{request.model}'")
+ assert request.model == "claude-sonnet-4-5"
+
+ print(f"Comparing messages length: Expected 1, Got {len(request.messages)}")
+ assert len(request.messages) == 1
+
+ print(f"Comparing stream: Expected False, Got {request.stream}")
+ assert request.stream is False
+
+ def test_requires_model(self):
+ """
+ What it does: Verifies that model is required.
+ Purpose: Ensure validation fails without model.
+ """
+ print("Setup: Attempting to create ChatCompletionRequest without model...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ChatCompletionRequest(messages=[ChatMessage(role="user", content="Hi")])
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "model" in str(exc_info.value)
+
+ def test_requires_messages(self):
+ """
+ What it does: Verifies that messages is required.
+ Purpose: Ensure validation fails without messages.
+ """
+ print("Setup: Attempting to create ChatCompletionRequest without messages...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ChatCompletionRequest(model="claude-sonnet-4-5")
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "messages" in str(exc_info.value)
+
+ def test_requires_at_least_one_message(self):
+ """
+ What it does: Verifies that at least one message is required.
+ Purpose: Ensure validation fails with empty messages.
+ """
+ print("Setup: Attempting to create ChatCompletionRequest with empty messages...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ChatCompletionRequest(model="claude-sonnet-4-5", messages=[])
+
+ print(f"ValidationError raised: {exc_info.value}")
+
+ def test_stream_defaults_to_false(self):
+ """
+ What it does: Verifies that stream defaults to False.
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating ChatCompletionRequest without explicit stream...")
+ request = ChatCompletionRequest(
+ model="test",
+ messages=[ChatMessage(role="user", content="Hi")]
+ )
+
+ print(f"Comparing stream: Expected False, Got {request.stream}")
+ assert request.stream is False
+
+ def test_accepts_stream_true(self):
+ """
+ What it does: Verifies that stream=True is accepted.
+ Purpose: Ensure streaming requests work.
+ """
+ print("Setup: Creating ChatCompletionRequest with stream=True...")
+ request = ChatCompletionRequest(
+ model="test",
+ messages=[ChatMessage(role="user", content="Hi")],
+ stream=True
+ )
+
+ print(f"Comparing stream: Expected True, Got {request.stream}")
+ assert request.stream is True
+
+ def test_accepts_tools(self):
+ """
+ What it does: Verifies that tools are accepted.
+ Purpose: Ensure function calling works.
+ """
+ print("Setup: Creating ChatCompletionRequest with tools...")
+ request = ChatCompletionRequest(
+ model="test",
+ messages=[ChatMessage(role="user", content="Hi")],
+ tools=[Tool(function=ToolFunction(name="test_tool"))]
+ )
+
+ print(f"Comparing tools: Got {request.tools}")
+ assert request.tools is not None
+ assert len(request.tools) == 1
+
+ def test_accepts_generation_parameters(self):
+ """
+ What it does: Verifies that generation parameters are accepted.
+ Purpose: Ensure temperature, top_p, max_tokens work.
+ """
+ print("Setup: Creating ChatCompletionRequest with generation params...")
+ request = ChatCompletionRequest(
+ model="test",
+ messages=[ChatMessage(role="user", content="Hi")],
+ temperature=0.7,
+ top_p=0.9,
+ max_tokens=1000
+ )
+
+ print(f"Comparing temperature: Expected 0.7, Got {request.temperature}")
+ assert request.temperature == 0.7
+
+ print(f"Comparing top_p: Expected 0.9, Got {request.top_p}")
+ assert request.top_p == 0.9
+
+ print(f"Comparing max_tokens: Expected 1000, Got {request.max_tokens}")
+ assert request.max_tokens == 1000
+
+
+# ==================================================================================================
+# Tests for ChatCompletionUsage
+# ==================================================================================================
+
+class TestChatCompletionUsage:
+ """Tests for ChatCompletionUsage Pydantic model."""
+
+ def test_valid_usage(self):
+ """
+ What it does: Verifies creation of valid ChatCompletionUsage.
+ Purpose: Ensure model accepts valid usage data.
+ """
+ print("Setup: Creating ChatCompletionUsage with valid data...")
+ usage = ChatCompletionUsage(
+ prompt_tokens=100,
+ completion_tokens=50,
+ total_tokens=150
+ )
+
+ print(f"Result: {usage}")
+ print(f"Comparing prompt_tokens: Expected 100, Got {usage.prompt_tokens}")
+ assert usage.prompt_tokens == 100
+
+ print(f"Comparing completion_tokens: Expected 50, Got {usage.completion_tokens}")
+ assert usage.completion_tokens == 50
+
+ print(f"Comparing total_tokens: Expected 150, Got {usage.total_tokens}")
+ assert usage.total_tokens == 150
+
+ def test_defaults_to_zero(self):
+ """
+ What it does: Verifies that all fields default to 0.
+ Purpose: Ensure default values are set correctly.
+ """
+ print("Setup: Creating ChatCompletionUsage without explicit values...")
+ usage = ChatCompletionUsage()
+
+ print(f"Comparing prompt_tokens: Expected 0, Got {usage.prompt_tokens}")
+ assert usage.prompt_tokens == 0
+
+ print(f"Comparing completion_tokens: Expected 0, Got {usage.completion_tokens}")
+ assert usage.completion_tokens == 0
+
+ print(f"Comparing total_tokens: Expected 0, Got {usage.total_tokens}")
+ assert usage.total_tokens == 0
+
+ def test_credits_used_is_optional(self):
+ """
+ What it does: Verifies that credits_used is optional.
+ Purpose: Ensure Kiro-specific field is optional.
+ """
+ print("Setup: Creating ChatCompletionUsage without credits_used...")
+ usage = ChatCompletionUsage()
+
+ print(f"Comparing credits_used: Expected None, Got {usage.credits_used}")
+ assert usage.credits_used is None
+
+
+# ==================================================================================================
+# Tests for ChatCompletionChoice
+# ==================================================================================================
+
+class TestChatCompletionChoice:
+ """Tests for ChatCompletionChoice Pydantic model."""
+
+ def test_valid_choice(self):
+ """
+ What it does: Verifies creation of valid ChatCompletionChoice.
+ Purpose: Ensure model accepts valid choice data.
+ """
+ print("Setup: Creating ChatCompletionChoice with valid data...")
+ choice = ChatCompletionChoice(
+ index=0,
+ message={"role": "assistant", "content": "Hello!"},
+ finish_reason="stop"
+ )
+
+ print(f"Result: {choice}")
+ print(f"Comparing index: Expected 0, Got {choice.index}")
+ assert choice.index == 0
+
+ print(f"Comparing message: Got {choice.message}")
+ assert choice.message["role"] == "assistant"
+
+ print(f"Comparing finish_reason: Expected 'stop', Got '{choice.finish_reason}'")
+ assert choice.finish_reason == "stop"
+
+ def test_index_defaults_to_zero(self):
+ """
+ What it does: Verifies that index defaults to 0.
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating ChatCompletionChoice without explicit index...")
+ choice = ChatCompletionChoice(message={"role": "assistant", "content": "Hi"})
+
+ print(f"Comparing index: Expected 0, Got {choice.index}")
+ assert choice.index == 0
+
+ def test_requires_message(self):
+ """
+ What it does: Verifies that message is required.
+ Purpose: Ensure validation fails without message.
+ """
+ print("Setup: Attempting to create ChatCompletionChoice without message...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ChatCompletionChoice(index=0, finish_reason="stop")
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "message" in str(exc_info.value)
+
+ def test_finish_reason_is_optional(self):
+ """
+ What it does: Verifies that finish_reason is optional.
+ Purpose: Ensure choices without finish_reason work.
+ """
+ print("Setup: Creating ChatCompletionChoice without finish_reason...")
+ choice = ChatCompletionChoice(message={"role": "assistant", "content": "Hi"})
+
+ print(f"Comparing finish_reason: Expected None, Got {choice.finish_reason}")
+ assert choice.finish_reason is None
+
+
+# ==================================================================================================
+# Tests for ChatCompletionResponse
+# ==================================================================================================
+
+class TestChatCompletionResponse:
+ """Tests for ChatCompletionResponse Pydantic model."""
+
+ def test_valid_response(self):
+ """
+ What it does: Verifies creation of valid ChatCompletionResponse.
+ Purpose: Ensure model accepts valid response data.
+ """
+ print("Setup: Creating ChatCompletionResponse with valid data...")
+ response = ChatCompletionResponse(
+ id="chatcmpl-123",
+ model="claude-sonnet-4-5",
+ choices=[ChatCompletionChoice(
+ message={"role": "assistant", "content": "Hello!"},
+ finish_reason="stop"
+ )],
+ usage=ChatCompletionUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
+ )
+
+ print(f"Result: {response}")
+ print(f"Comparing id: Expected 'chatcmpl-123', Got '{response.id}'")
+ assert response.id == "chatcmpl-123"
+
+ print(f"Comparing object: Expected 'chat.completion', Got '{response.object}'")
+ assert response.object == "chat.completion"
+
+ print(f"Comparing model: Expected 'claude-sonnet-4-5', Got '{response.model}'")
+ assert response.model == "claude-sonnet-4-5"
+
+ print(f"Comparing choices length: Expected 1, Got {len(response.choices)}")
+ assert len(response.choices) == 1
+
+ def test_object_defaults_to_chat_completion(self):
+ """
+ What it does: Verifies that object defaults to "chat.completion".
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating ChatCompletionResponse without explicit object...")
+ response = ChatCompletionResponse(
+ id="test",
+ model="test",
+ choices=[ChatCompletionChoice(message={"role": "assistant", "content": "Hi"})],
+ usage=ChatCompletionUsage()
+ )
+
+ print(f"Comparing object: Expected 'chat.completion', Got '{response.object}'")
+ assert response.object == "chat.completion"
+
+ def test_created_is_auto_generated(self):
+ """
+ What it does: Verifies that created timestamp is auto-generated.
+ Purpose: Ensure timestamp is set automatically.
+ """
+ print("Setup: Creating ChatCompletionResponse without explicit created...")
+ response = ChatCompletionResponse(
+ id="test",
+ model="test",
+ choices=[ChatCompletionChoice(message={"role": "assistant", "content": "Hi"})],
+ usage=ChatCompletionUsage()
+ )
+
+ print(f"Comparing created: Got {response.created}")
+ assert response.created > 0
+ assert isinstance(response.created, int)
+
+ def test_requires_id(self):
+ """
+ What it does: Verifies that id is required.
+ Purpose: Ensure validation fails without id.
+ """
+ print("Setup: Attempting to create ChatCompletionResponse without id...")
+
+ print("Action: Creating model (should raise ValidationError)...")
+ with pytest.raises(ValidationError) as exc_info:
+ ChatCompletionResponse(
+ model="test",
+ choices=[ChatCompletionChoice(message={"role": "assistant", "content": "Hi"})],
+ usage=ChatCompletionUsage()
+ )
+
+ print(f"ValidationError raised: {exc_info.value}")
+ assert "id" in str(exc_info.value)
+
+
+# ==================================================================================================
+# Tests for Streaming Models
+# ==================================================================================================
+
+class TestChatCompletionChunkDelta:
+ """Tests for ChatCompletionChunkDelta Pydantic model."""
+
+ def test_valid_delta_with_content(self):
+ """
+ What it does: Verifies creation of valid delta with content.
+ Purpose: Ensure model accepts content delta.
+ """
+ print("Setup: Creating ChatCompletionChunkDelta with content...")
+ delta = ChatCompletionChunkDelta(content="Hello")
+
+ print(f"Result: {delta}")
+ print(f"Comparing content: Expected 'Hello', Got '{delta.content}'")
+ assert delta.content == "Hello"
+
+ def test_valid_delta_with_role(self):
+ """
+ What it does: Verifies creation of valid delta with role.
+ Purpose: Ensure model accepts role delta (first chunk).
+ """
+ print("Setup: Creating ChatCompletionChunkDelta with role...")
+ delta = ChatCompletionChunkDelta(role="assistant")
+
+ print(f"Result: {delta}")
+ print(f"Comparing role: Expected 'assistant', Got '{delta.role}'")
+ assert delta.role == "assistant"
+
+ def test_all_fields_optional(self):
+ """
+ What it does: Verifies that all fields are optional.
+ Purpose: Ensure empty delta works.
+ """
+ print("Setup: Creating empty ChatCompletionChunkDelta...")
+ delta = ChatCompletionChunkDelta()
+
+ print(f"Comparing role: Expected None, Got {delta.role}")
+ assert delta.role is None
+
+ print(f"Comparing content: Expected None, Got {delta.content}")
+ assert delta.content is None
+
+ print(f"Comparing tool_calls: Expected None, Got {delta.tool_calls}")
+ assert delta.tool_calls is None
+
+ def test_accepts_tool_calls(self):
+ """
+ What it does: Verifies that tool_calls is accepted.
+ Purpose: Ensure streaming tool calls work.
+ """
+ print("Setup: Creating ChatCompletionChunkDelta with tool_calls...")
+ delta = ChatCompletionChunkDelta(
+ tool_calls=[{"index": 0, "id": "call_1", "function": {"name": "test"}}]
+ )
+
+ print(f"Comparing tool_calls: Got {delta.tool_calls}")
+ assert delta.tool_calls is not None
+ assert len(delta.tool_calls) == 1
+
+
+class TestChatCompletionChunkChoice:
+ """Tests for ChatCompletionChunkChoice Pydantic model."""
+
+ def test_valid_chunk_choice(self):
+ """
+ What it does: Verifies creation of valid chunk choice.
+ Purpose: Ensure model accepts valid chunk choice.
+ """
+ print("Setup: Creating ChatCompletionChunkChoice with valid data...")
+ choice = ChatCompletionChunkChoice(
+ index=0,
+ delta=ChatCompletionChunkDelta(content="Hello")
+ )
+
+ print(f"Result: {choice}")
+ print(f"Comparing index: Expected 0, Got {choice.index}")
+ assert choice.index == 0
+
+ print(f"Comparing delta.content: Expected 'Hello', Got '{choice.delta.content}'")
+ assert choice.delta.content == "Hello"
+
+ def test_index_defaults_to_zero(self):
+ """
+ What it does: Verifies that index defaults to 0.
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating ChatCompletionChunkChoice without explicit index...")
+ choice = ChatCompletionChunkChoice(delta=ChatCompletionChunkDelta())
+
+ print(f"Comparing index: Expected 0, Got {choice.index}")
+ assert choice.index == 0
+
+ def test_finish_reason_is_optional(self):
+ """
+ What it does: Verifies that finish_reason is optional.
+ Purpose: Ensure intermediate chunks work.
+ """
+ print("Setup: Creating ChatCompletionChunkChoice without finish_reason...")
+ choice = ChatCompletionChunkChoice(delta=ChatCompletionChunkDelta(content="Hi"))
+
+ print(f"Comparing finish_reason: Expected None, Got {choice.finish_reason}")
+ assert choice.finish_reason is None
+
+ def test_accepts_finish_reason(self):
+ """
+ What it does: Verifies that finish_reason is accepted.
+ Purpose: Ensure final chunk works.
+ """
+ print("Setup: Creating ChatCompletionChunkChoice with finish_reason...")
+ choice = ChatCompletionChunkChoice(
+ delta=ChatCompletionChunkDelta(),
+ finish_reason="stop"
+ )
+
+ print(f"Comparing finish_reason: Expected 'stop', Got '{choice.finish_reason}'")
+ assert choice.finish_reason == "stop"
+
+
+class TestChatCompletionChunk:
+ """Tests for ChatCompletionChunk Pydantic model."""
+
+ def test_valid_chunk(self):
+ """
+ What it does: Verifies creation of valid chunk.
+ Purpose: Ensure model accepts valid chunk data.
+ """
+ print("Setup: Creating ChatCompletionChunk with valid data...")
+ chunk = ChatCompletionChunk(
+ id="chatcmpl-123",
+ model="claude-sonnet-4-5",
+ choices=[ChatCompletionChunkChoice(
+ delta=ChatCompletionChunkDelta(content="Hello")
+ )]
+ )
+
+ print(f"Result: {chunk}")
+ print(f"Comparing id: Expected 'chatcmpl-123', Got '{chunk.id}'")
+ assert chunk.id == "chatcmpl-123"
+
+ print(f"Comparing object: Expected 'chat.completion.chunk', Got '{chunk.object}'")
+ assert chunk.object == "chat.completion.chunk"
+
+ print(f"Comparing model: Expected 'claude-sonnet-4-5', Got '{chunk.model}'")
+ assert chunk.model == "claude-sonnet-4-5"
+
+ def test_object_defaults_to_chunk(self):
+ """
+ What it does: Verifies that object defaults to "chat.completion.chunk".
+ Purpose: Ensure default value is set correctly.
+ """
+ print("Setup: Creating ChatCompletionChunk without explicit object...")
+ chunk = ChatCompletionChunk(
+ id="test",
+ model="test",
+ choices=[ChatCompletionChunkChoice(delta=ChatCompletionChunkDelta())]
+ )
+
+ print(f"Comparing object: Expected 'chat.completion.chunk', Got '{chunk.object}'")
+ assert chunk.object == "chat.completion.chunk"
+
+ def test_usage_is_optional(self):
+ """
+ What it does: Verifies that usage is optional.
+ Purpose: Ensure intermediate chunks work without usage.
+ """
+ print("Setup: Creating ChatCompletionChunk without usage...")
+ chunk = ChatCompletionChunk(
+ id="test",
+ model="test",
+ choices=[ChatCompletionChunkChoice(delta=ChatCompletionChunkDelta())]
+ )
+
+ print(f"Comparing usage: Expected None, Got {chunk.usage}")
+ assert chunk.usage is None
+
+ def test_accepts_usage(self):
+ """
+ What it does: Verifies that usage is accepted.
+ Purpose: Ensure final chunk with usage works.
+ """
+ print("Setup: Creating ChatCompletionChunk with usage...")
+ chunk = ChatCompletionChunk(
+ id="test",
+ model="test",
+ choices=[ChatCompletionChunkChoice(
+ delta=ChatCompletionChunkDelta(),
+ finish_reason="stop"
+ )],
+ usage=ChatCompletionUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
+ )
+
+ print(f"Comparing usage: Got {chunk.usage}")
+ assert chunk.usage is not None
+ assert chunk.usage.total_tokens == 15
diff --git a/kiro-gateway/tests/unit/test_network_errors.py b/kiro-gateway/tests/unit/test_network_errors.py
new file mode 100644
index 0000000000000000000000000000000000000000..84203662b457dcf3c97332eafffbe7516c4e6588
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_network_errors.py
@@ -0,0 +1,671 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for network error classification system.
+Tests classify_network_error(), format_error_for_user(), and get_short_error_message().
+"""
+
+import socket
+import pytest
+
+import httpx
+
+from kiro.network_errors import (
+ ErrorCategory,
+ NetworkErrorInfo,
+ classify_network_error,
+ format_error_for_user,
+ get_short_error_message
+)
+
+
+class TestClassifyNetworkErrorDNS:
+ """Tests for DNS resolution error classification."""
+
+ def test_dns_error_with_socket_gaierror_windows(self):
+ """
+ What it does: Verifies DNS errors are classified correctly on Windows.
+ Purpose: Ensure socket.gaierror with errno 11001 is detected as DNS_RESOLUTION (issue #53).
+ """
+ print("Setup: Creating ConnectError with socket.gaierror (Windows errno 11001)...")
+ dns_error = socket.gaierror(11001, "getaddrinfo failed")
+ connect_error = httpx.ConnectError("All connection attempts failed")
+ connect_error.__cause__ = dns_error
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is DNS_RESOLUTION...")
+ print(f"Comparing category: Expected {ErrorCategory.DNS_RESOLUTION}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.DNS_RESOLUTION
+ assert "DNS resolution failed" in error_info.user_message
+ assert "cannot resolve" in error_info.user_message.lower()
+ assert error_info.is_retryable is True
+ assert error_info.suggested_http_code == 502
+ assert "11001" in error_info.technical_details
+
+ def test_dns_error_with_socket_gaierror_unix(self):
+ """
+ What it does: Verifies DNS errors are classified correctly on Unix.
+ Purpose: Ensure socket.gaierror with Unix errno is detected as DNS_RESOLUTION.
+ """
+ print("Setup: Creating ConnectError with socket.gaierror (Unix errno -2)...")
+ dns_error = socket.gaierror(-2, "Name or service not known")
+ connect_error = httpx.ConnectError("Connection failed")
+ connect_error.__cause__ = dns_error
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is DNS_RESOLUTION...")
+ assert error_info.category == ErrorCategory.DNS_RESOLUTION
+ assert "DNS" in error_info.user_message
+ assert "-2" in error_info.technical_details
+
+ def test_dns_error_includes_troubleshooting_steps(self):
+ """
+ What it does: Verifies DNS errors include actionable troubleshooting steps.
+ Purpose: Ensure users get clear guidance on fixing DNS issues.
+ """
+ print("Setup: Creating DNS error...")
+ dns_error = socket.gaierror(11001, "getaddrinfo failed")
+ connect_error = httpx.ConnectError("Connection failed")
+ connect_error.__cause__ = dns_error
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Troubleshooting steps present...")
+ steps = error_info.troubleshooting_steps
+ assert len(steps) >= 3
+ assert any("DNS" in step for step in steps)
+ assert any("8.8.8.8" in step or "1.1.1.1" in step for step in steps)
+ assert any("VPN" in step for step in steps)
+ assert any("firewall" in step.lower() or "antivirus" in step.lower() for step in steps)
+
+ def test_dns_error_technical_details_include_errno(self):
+ """
+ What it does: Verifies technical details include errno for debugging.
+ Purpose: Ensure developers can identify specific DNS error codes.
+ """
+ print("Setup: Creating DNS error with specific errno...")
+ dns_error = socket.gaierror(11001, "getaddrinfo failed")
+ connect_error = httpx.ConnectError("Failed")
+ connect_error.__cause__ = dns_error
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Technical details include errno...")
+ assert "errno" in error_info.technical_details.lower()
+ assert "11001" in error_info.technical_details
+
+
+class TestClassifyNetworkErrorConnection:
+ """Tests for connection error classification."""
+
+ def test_connection_refused_error(self):
+ """
+ What it does: Verifies connection refused errors are classified correctly.
+ Purpose: Ensure "Connection refused" is detected as CONNECTION_REFUSED.
+ """
+ print("Setup: Creating ConnectError with 'Connection refused'...")
+ connect_error = httpx.ConnectError("Connection refused")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is CONNECTION_REFUSED...")
+ print(f"Comparing category: Expected {ErrorCategory.CONNECTION_REFUSED}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.CONNECTION_REFUSED
+ assert "Connection refused" in error_info.user_message
+ assert "not accepting connections" in error_info.user_message
+ assert error_info.is_retryable is True
+ assert error_info.suggested_http_code == 502
+
+ def test_connection_refused_with_econnrefused(self):
+ """
+ What it does: Verifies ECONNREFUSED is detected as CONNECTION_REFUSED.
+ Purpose: Ensure Unix-style error codes are recognized.
+ """
+ print("Setup: Creating ConnectError with ECONNREFUSED...")
+ connect_error = httpx.ConnectError("[Errno 111] ECONNREFUSED")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is CONNECTION_REFUSED...")
+ assert error_info.category == ErrorCategory.CONNECTION_REFUSED
+
+ def test_connection_reset_error(self):
+ """
+ What it does: Verifies connection reset errors are classified correctly.
+ Purpose: Ensure "Connection reset" is detected as CONNECTION_RESET.
+ """
+ print("Setup: Creating ConnectError with 'Connection reset'...")
+ connect_error = httpx.ConnectError("Connection reset by peer")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is CONNECTION_RESET...")
+ print(f"Comparing category: Expected {ErrorCategory.CONNECTION_RESET}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.CONNECTION_RESET
+ assert "Connection reset" in error_info.user_message
+ assert "closed the connection" in error_info.user_message
+ assert error_info.is_retryable is True
+
+ def test_connection_reset_with_econnreset(self):
+ """
+ What it does: Verifies ECONNRESET is detected as CONNECTION_RESET.
+ Purpose: Ensure Unix-style error codes are recognized.
+ """
+ print("Setup: Creating ConnectError with ECONNRESET...")
+ connect_error = httpx.ConnectError("[Errno 104] ECONNRESET")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is CONNECTION_RESET...")
+ assert error_info.category == ErrorCategory.CONNECTION_RESET
+
+ def test_network_unreachable_error(self):
+ """
+ What it does: Verifies network unreachable errors are classified correctly.
+ Purpose: Ensure "Network is unreachable" is detected as NETWORK_UNREACHABLE.
+ """
+ print("Setup: Creating ConnectError with 'Network is unreachable'...")
+ connect_error = httpx.ConnectError("Network is unreachable")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is NETWORK_UNREACHABLE...")
+ print(f"Comparing category: Expected {ErrorCategory.NETWORK_UNREACHABLE}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.NETWORK_UNREACHABLE
+ assert "Network unreachable" in error_info.user_message
+ assert error_info.is_retryable is True
+
+ def test_network_unreachable_with_no_route_to_host(self):
+ """
+ What it does: Verifies "No route to host" is detected as NETWORK_UNREACHABLE.
+ Purpose: Ensure alternative error messages are recognized.
+ """
+ print("Setup: Creating ConnectError with 'No route to host'...")
+ connect_error = httpx.ConnectError("No route to host")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is NETWORK_UNREACHABLE...")
+ assert error_info.category == ErrorCategory.NETWORK_UNREACHABLE
+
+ def test_network_unreachable_with_enetunreach(self):
+ """
+ What it does: Verifies ENETUNREACH is detected as NETWORK_UNREACHABLE.
+ Purpose: Ensure Unix-style error codes are recognized.
+ """
+ print("Setup: Creating ConnectError with ENETUNREACH...")
+ connect_error = httpx.ConnectError("[Errno 101] ENETUNREACH")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is NETWORK_UNREACHABLE...")
+ assert error_info.category == ErrorCategory.NETWORK_UNREACHABLE
+
+ def test_generic_connect_error_classified_as_unknown(self):
+ """
+ What it does: Verifies generic connection errors fall back to UNKNOWN.
+ Purpose: Ensure unrecognized errors have a fallback category.
+ """
+ print("Setup: Creating generic ConnectError...")
+ connect_error = httpx.ConnectError("All connection attempts failed")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is UNKNOWN...")
+ print(f"Comparing category: Expected {ErrorCategory.UNKNOWN}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.UNKNOWN
+ assert "Connection failed" in error_info.user_message
+ assert error_info.is_retryable is True
+ assert error_info.suggested_http_code == 502
+
+
+class TestClassifyNetworkErrorTimeout:
+ """Tests for timeout error classification."""
+
+ def test_connect_timeout_error(self):
+ """
+ What it does: Verifies ConnectTimeout is classified correctly.
+ Purpose: Ensure TCP handshake timeouts are detected as TIMEOUT_CONNECT.
+ """
+ print("Setup: Creating ConnectTimeout...")
+ timeout_error = httpx.ConnectTimeout("Connection timeout")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(timeout_error)
+
+ print("Verification: Category is TIMEOUT_CONNECT...")
+ print(f"Comparing category: Expected {ErrorCategory.TIMEOUT_CONNECT}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.TIMEOUT_CONNECT
+ assert "Connection timeout" in error_info.user_message
+ assert "did not respond" in error_info.user_message
+ assert error_info.is_retryable is True
+ assert error_info.suggested_http_code == 504
+
+ def test_connect_timeout_includes_troubleshooting(self):
+ """
+ What it does: Verifies connect timeout includes troubleshooting steps.
+ Purpose: Ensure users get guidance on fixing timeout issues.
+ """
+ print("Setup: Creating ConnectTimeout...")
+ timeout_error = httpx.ConnectTimeout("Timeout")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(timeout_error)
+
+ print("Verification: Troubleshooting steps present...")
+ steps = error_info.troubleshooting_steps
+ assert len(steps) >= 2
+ assert any("internet connection" in step.lower() for step in steps)
+ assert any("server" in step.lower() or "overloaded" in step.lower() for step in steps)
+
+ def test_read_timeout_error(self):
+ """
+ What it does: Verifies ReadTimeout is classified correctly.
+ Purpose: Ensure read timeouts are detected as TIMEOUT_READ.
+ """
+ print("Setup: Creating ReadTimeout...")
+ timeout_error = httpx.ReadTimeout("Read timeout")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(timeout_error)
+
+ print("Verification: Category is TIMEOUT_READ...")
+ print(f"Comparing category: Expected {ErrorCategory.TIMEOUT_READ}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.TIMEOUT_READ
+ assert "Read timeout" in error_info.user_message
+ assert "stopped responding" in error_info.user_message
+ assert error_info.is_retryable is True
+ assert error_info.suggested_http_code == 504
+
+ def test_read_timeout_includes_troubleshooting(self):
+ """
+ What it does: Verifies read timeout includes troubleshooting steps.
+ Purpose: Ensure users get guidance on fixing read timeout issues.
+ """
+ print("Setup: Creating ReadTimeout...")
+ timeout_error = httpx.ReadTimeout("Timeout")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(timeout_error)
+
+ print("Verification: Troubleshooting steps present...")
+ steps = error_info.troubleshooting_steps
+ assert len(steps) >= 2
+ assert any("server" in step.lower() or "processing" in step.lower() for step in steps)
+
+ def test_generic_timeout_error(self):
+ """
+ What it does: Verifies generic TimeoutException is classified as TIMEOUT_READ.
+ Purpose: Ensure unspecified timeouts have a fallback.
+ """
+ print("Setup: Creating generic TimeoutException...")
+ timeout_error = httpx.TimeoutException("Timeout")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(timeout_error)
+
+ print("Verification: Category is TIMEOUT_READ...")
+ print(f"Comparing category: Expected {ErrorCategory.TIMEOUT_READ}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.TIMEOUT_READ
+ assert "timeout" in error_info.user_message.lower()
+ assert error_info.is_retryable is True
+
+
+class TestClassifyNetworkErrorSSL:
+ """Tests for SSL/TLS error classification."""
+
+ def test_ssl_error_detection(self):
+ """
+ What it does: Verifies SSL errors are classified correctly.
+ Purpose: Ensure SSL/TLS errors are detected as SSL_ERROR.
+ """
+ print("Setup: Creating ConnectError with SSL in message...")
+ connect_error = httpx.ConnectError("SSL handshake failed")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is SSL_ERROR...")
+ print(f"Comparing category: Expected {ErrorCategory.SSL_ERROR}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.SSL_ERROR
+ assert "SSL/TLS error" in error_info.user_message
+ assert "secure connection" in error_info.user_message
+ assert error_info.is_retryable is False
+ assert error_info.suggested_http_code == 502
+
+ def test_tls_error_detection(self):
+ """
+ What it does: Verifies TLS errors are detected as SSL_ERROR.
+ Purpose: Ensure TLS keyword is recognized.
+ """
+ print("Setup: Creating ConnectError with TLS in message...")
+ connect_error = httpx.ConnectError("TLS connection failed")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is SSL_ERROR...")
+ assert error_info.category == ErrorCategory.SSL_ERROR
+
+ def test_certificate_error_detection(self):
+ """
+ What it does: Verifies certificate errors are detected as SSL_ERROR.
+ Purpose: Ensure certificate keyword is recognized.
+ """
+ print("Setup: Creating ConnectError with certificate in message...")
+ connect_error = httpx.ConnectError("Certificate verification failed")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Category is SSL_ERROR...")
+ assert error_info.category == ErrorCategory.SSL_ERROR
+
+ def test_ssl_error_includes_troubleshooting(self):
+ """
+ What it does: Verifies SSL errors include troubleshooting steps.
+ Purpose: Ensure users get guidance on fixing SSL issues.
+ """
+ print("Setup: Creating SSL error...")
+ connect_error = httpx.ConnectError("SSL error")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(connect_error)
+
+ print("Verification: Troubleshooting steps present...")
+ steps = error_info.troubleshooting_steps
+ assert len(steps) >= 2
+ assert any("certificate" in step.lower() for step in steps)
+ assert any("date" in step.lower() or "time" in step.lower() for step in steps)
+
+
+class TestClassifyNetworkErrorProxy:
+ """Tests for proxy error classification."""
+
+ def test_proxy_error_detection(self):
+ """
+ What it does: Verifies proxy errors are classified correctly.
+ Purpose: Ensure ProxyError is detected as PROXY_ERROR.
+ """
+ print("Setup: Creating ProxyError...")
+ proxy_error = httpx.ProxyError("Proxy connection failed")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(proxy_error)
+
+ print("Verification: Category is PROXY_ERROR...")
+ print(f"Comparing category: Expected {ErrorCategory.PROXY_ERROR}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.PROXY_ERROR
+ assert "Proxy" in error_info.user_message
+ assert "cannot connect through" in error_info.user_message
+ assert error_info.is_retryable is True
+ assert error_info.suggested_http_code == 502
+
+ def test_proxy_error_includes_troubleshooting(self):
+ """
+ What it does: Verifies proxy errors include troubleshooting steps.
+ Purpose: Ensure users get guidance on fixing proxy issues.
+ """
+ print("Setup: Creating ProxyError...")
+ proxy_error = httpx.ProxyError("Proxy failed")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(proxy_error)
+
+ print("Verification: Troubleshooting steps present...")
+ steps = error_info.troubleshooting_steps
+ assert len(steps) >= 2
+ assert any("HTTP_PROXY" in step or "HTTPS_PROXY" in step for step in steps)
+ assert any("proxy" in step.lower() for step in steps)
+
+
+class TestClassifyNetworkErrorRedirects:
+ """Tests for redirect error classification."""
+
+ def test_too_many_redirects_error(self):
+ """
+ What it does: Verifies TooManyRedirects is classified correctly.
+ Purpose: Ensure redirect loops are detected as TOO_MANY_REDIRECTS.
+ """
+ print("Setup: Creating TooManyRedirects...")
+ redirect_error = httpx.TooManyRedirects("Too many redirects")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(redirect_error)
+
+ print("Verification: Category is TOO_MANY_REDIRECTS...")
+ print(f"Comparing category: Expected {ErrorCategory.TOO_MANY_REDIRECTS}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.TOO_MANY_REDIRECTS
+ assert "redirect" in error_info.user_message.lower()
+ assert "loop" in error_info.user_message.lower()
+ assert error_info.is_retryable is False
+ assert error_info.suggested_http_code == 502
+
+
+class TestClassifyNetworkErrorGeneric:
+ """Tests for generic error classification."""
+
+ def test_generic_request_error_classified_as_unknown(self):
+ """
+ What it does: Verifies generic RequestError falls back to UNKNOWN.
+ Purpose: Ensure unrecognized httpx errors have a fallback.
+ """
+ print("Setup: Creating generic RequestError...")
+ request_error = httpx.RequestError("Unknown network error")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(request_error)
+
+ print("Verification: Category is UNKNOWN...")
+ print(f"Comparing category: Expected {ErrorCategory.UNKNOWN}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.UNKNOWN
+ assert "unexpected error" in error_info.user_message.lower()
+ assert error_info.is_retryable is True
+ assert error_info.suggested_http_code == 502
+
+ def test_non_httpx_error_classified_as_unknown(self):
+ """
+ What it does: Verifies non-httpx errors fall back to UNKNOWN.
+ Purpose: Ensure graceful handling of unexpected exception types.
+ """
+ print("Setup: Creating generic Exception...")
+ generic_error = Exception("Something went wrong")
+
+ print("Action: Classifying error...")
+ error_info = classify_network_error(generic_error)
+
+ print("Verification: Category is UNKNOWN...")
+ print(f"Comparing category: Expected {ErrorCategory.UNKNOWN}, Got {error_info.category}")
+ assert error_info.category == ErrorCategory.UNKNOWN
+ assert error_info.suggested_http_code == 500
+
+
+class TestFormatErrorForUser:
+ """Tests for format_error_for_user() function."""
+
+ def test_format_openai_includes_troubleshooting(self):
+ """
+ What it does: Verifies OpenAI format includes troubleshooting steps.
+ Purpose: Ensure users get actionable guidance in API responses.
+ """
+ print("Setup: Creating NetworkErrorInfo...")
+ error_info = NetworkErrorInfo(
+ category=ErrorCategory.DNS_RESOLUTION,
+ user_message="DNS failed",
+ troubleshooting_steps=["Step 1", "Step 2"],
+ technical_details="Technical info",
+ is_retryable=True,
+ suggested_http_code=502
+ )
+
+ print("Action: Formatting for OpenAI...")
+ formatted = format_error_for_user(error_info, format_type="openai", include_troubleshooting=True)
+
+ print("Verification: OpenAI format structure...")
+ assert "error" in formatted
+ assert "message" in formatted["error"]
+ assert "type" in formatted["error"]
+ assert "code" in formatted["error"]
+ assert formatted["error"]["type"] == "connectivity_error"
+ assert formatted["error"]["code"] == "dns_resolution"
+ assert "Step 1" in formatted["error"]["message"]
+ assert "Step 2" in formatted["error"]["message"]
+
+ def test_format_openai_without_troubleshooting(self):
+ """
+ What it does: Verifies OpenAI format can exclude troubleshooting.
+ Purpose: Ensure flexibility in error message verbosity.
+ """
+ print("Setup: Creating NetworkErrorInfo...")
+ error_info = NetworkErrorInfo(
+ category=ErrorCategory.TIMEOUT_CONNECT,
+ user_message="Connection timeout",
+ troubleshooting_steps=["Step 1"],
+ technical_details="Technical info",
+ is_retryable=True,
+ suggested_http_code=504
+ )
+
+ print("Action: Formatting for OpenAI without troubleshooting...")
+ formatted = format_error_for_user(error_info, format_type="openai", include_troubleshooting=False)
+
+ print("Verification: No troubleshooting steps in message...")
+ assert "Step 1" not in formatted["error"]["message"]
+ assert formatted["error"]["message"] == "Connection timeout"
+
+ def test_format_anthropic_structure(self):
+ """
+ What it does: Verifies Anthropic format structure.
+ Purpose: Ensure compatibility with Anthropic API error format.
+ """
+ print("Setup: Creating NetworkErrorInfo...")
+ error_info = NetworkErrorInfo(
+ category=ErrorCategory.CONNECTION_REFUSED,
+ user_message="Connection refused",
+ troubleshooting_steps=["Step 1"],
+ technical_details="Technical info",
+ is_retryable=True,
+ suggested_http_code=502
+ )
+
+ print("Action: Formatting for Anthropic...")
+ formatted = format_error_for_user(error_info, format_type="anthropic")
+
+ print("Verification: Anthropic format structure...")
+ assert "type" in formatted
+ assert formatted["type"] == "error"
+ assert "error" in formatted
+ assert "type" in formatted["error"]
+ assert "message" in formatted["error"]
+ assert formatted["error"]["type"] == "connectivity_error"
+
+ def test_format_generic_includes_technical_details(self):
+ """
+ What it does: Verifies generic format includes technical details.
+ Purpose: Ensure debugging information is available in generic format.
+ """
+ print("Setup: Creating NetworkErrorInfo...")
+ error_info = NetworkErrorInfo(
+ category=ErrorCategory.SSL_ERROR,
+ user_message="SSL error",
+ troubleshooting_steps=[],
+ technical_details="ConnectError: SSL handshake failed",
+ is_retryable=False,
+ suggested_http_code=502
+ )
+
+ print("Action: Formatting with generic format...")
+ formatted = format_error_for_user(error_info, format_type="generic")
+
+ print("Verification: Technical details present...")
+ assert "technical_details" in formatted["error"]
+ assert formatted["error"]["technical_details"] == "ConnectError: SSL handshake failed"
+
+
+class TestGetShortErrorMessage:
+ """Tests for get_short_error_message() function."""
+
+ def test_short_message_no_brackets(self):
+ """
+ What it does: Verifies short message doesn't include brackets.
+ Purpose: Ensure clean log output without category brackets.
+ """
+ print("Setup: Creating NetworkErrorInfo...")
+ error_info = NetworkErrorInfo(
+ category=ErrorCategory.DNS_RESOLUTION,
+ user_message="DNS resolution failed",
+ troubleshooting_steps=[],
+ technical_details="Technical info",
+ is_retryable=True,
+ suggested_http_code=502
+ )
+
+ print("Action: Getting short message...")
+ short_msg = get_short_error_message(error_info)
+
+ print("Verification: No brackets in message...")
+ print(f"Short message: {short_msg}")
+ assert short_msg == "DNS resolution failed"
+ assert "[" not in short_msg
+ assert "]" not in short_msg
+
+ def test_short_message_different_categories(self):
+ """
+ What it does: Verifies short messages for different error categories.
+ Purpose: Ensure consistent format across all error types.
+ """
+ print("Setup: Creating multiple NetworkErrorInfo instances...")
+ errors = [
+ NetworkErrorInfo(ErrorCategory.TIMEOUT_CONNECT, "Timeout", [], "Tech", True, 504),
+ NetworkErrorInfo(ErrorCategory.CONNECTION_REFUSED, "Refused", [], "Tech", True, 502),
+ NetworkErrorInfo(ErrorCategory.UNKNOWN, "Unknown", [], "Tech", True, 502),
+ ]
+
+ print("Action: Getting short messages...")
+ for error_info in errors:
+ short_msg = get_short_error_message(error_info)
+
+ print(f"Verification: {error_info.category} -> {short_msg}")
+ assert short_msg == error_info.user_message
+ assert "[" not in short_msg
+
+
+class TestNetworkErrorInfoDataclass:
+ """Tests for NetworkErrorInfo dataclass."""
+
+ def test_network_error_info_creation(self):
+ """
+ What it does: Verifies NetworkErrorInfo can be created with all fields.
+ Purpose: Ensure dataclass structure is correct.
+ """
+ print("Setup: Creating NetworkErrorInfo...")
+ error_info = NetworkErrorInfo(
+ category=ErrorCategory.DNS_RESOLUTION,
+ user_message="Test message",
+ troubleshooting_steps=["Step 1", "Step 2"],
+ technical_details="Technical details",
+ is_retryable=True,
+ suggested_http_code=502
+ )
+
+ print("Verification: All fields accessible...")
+ assert error_info.category == ErrorCategory.DNS_RESOLUTION
+ assert error_info.user_message == "Test message"
+ assert len(error_info.troubleshooting_steps) == 2
+ assert error_info.technical_details == "Technical details"
+ assert error_info.is_retryable is True
+ assert error_info.suggested_http_code == 502
diff --git a/kiro-gateway/tests/unit/test_parsers.py b/kiro-gateway/tests/unit/test_parsers.py
new file mode 100644
index 0000000000000000000000000000000000000000..34409cd8ad55ce8dce718dcf9a865a917944f04f
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_parsers.py
@@ -0,0 +1,1225 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for AwsEventStreamParser and auxiliary parsing functions.
+Tests the parsing logic for AWS SSE stream from Kiro API.
+"""
+
+import pytest
+
+from kiro.parsers import (
+ AwsEventStreamParser,
+ find_matching_brace,
+ parse_bracket_tool_calls,
+ deduplicate_tool_calls
+)
+
+
+class TestFindMatchingBrace:
+ """Tests for find_matching_brace function."""
+
+ def test_simple_json_object(self):
+ """
+ What it does: Tests finding closing brace for simple JSON.
+ Goal: Ensure the basic case works.
+ """
+ print("Setup: Simple JSON object...")
+ text = '{"key": "value"}'
+
+ print("Action: Finding closing brace...")
+ result = find_matching_brace(text, 0)
+
+ print(f"Comparing result: Expected 15, Got {result}")
+ assert result == 15
+
+ def test_nested_json_object(self):
+ """
+ What it does: Tests finding brace for nested JSON.
+ Goal: Ensure nesting is handled correctly.
+ """
+ print("Setup: Nested JSON object...")
+ text = '{"outer": {"inner": "value"}}'
+
+ print("Action: Finding closing brace...")
+ result = find_matching_brace(text, 0)
+
+ # String length 29, last character index 28
+ print(f"Comparing result: Expected 28, Got {result}")
+ assert result == 28
+
+ def test_json_with_braces_in_string(self):
+ """
+ What it does: Tests ignoring braces inside strings.
+ Goal: Ensure braces in strings don't affect counting.
+ """
+ print("Setup: JSON with braces in string...")
+ text = '{"text": "Hello {world}"}'
+
+ print("Action: Finding closing brace...")
+ result = find_matching_brace(text, 0)
+
+ print(f"Comparing result: Expected 24, Got {result}")
+ assert result == 24
+
+ def test_json_with_escaped_quotes(self):
+ """
+ What it does: Tests handling of escaped quotes.
+ Goal: Ensure escape sequences don't break parsing.
+ """
+ print("Setup: JSON with escaped quotes...")
+ text = '{"text": "Say \\"hello\\""}'
+
+ print("Action: Finding closing brace...")
+ result = find_matching_brace(text, 0)
+
+ # String length 25, last character index 24
+ print(f"Comparing result: Expected 24, Got {result}")
+ assert result == 24
+
+ def test_incomplete_json(self):
+ """
+ What it does: Tests handling of incomplete JSON.
+ Goal: Ensure -1 is returned for incomplete JSON.
+ """
+ print("Setup: Incomplete JSON...")
+ text = '{"key": "value"'
+
+ print("Action: Finding closing brace...")
+ result = find_matching_brace(text, 0)
+
+ print(f"Comparing result: Expected -1, Got {result}")
+ assert result == -1
+
+ def test_invalid_start_position(self):
+ """
+ What it does: Tests handling of invalid start position.
+ Goal: Ensure -1 is returned if start_pos is not on '{'.
+ """
+ print("Setup: Text without brace at start_pos...")
+ text = 'hello {"key": "value"}'
+
+ print("Action: Finding from position 0 (not a brace)...")
+ result = find_matching_brace(text, 0)
+
+ print(f"Comparing result: Expected -1, Got {result}")
+ assert result == -1
+
+ def test_start_position_out_of_bounds(self):
+ """
+ What it does: Tests handling of position beyond text bounds.
+ Goal: Ensure -1 is returned for invalid position.
+ """
+ print("Setup: Short text...")
+ text = '{"a":1}'
+
+ print("Action: Finding from position 100...")
+ result = find_matching_brace(text, 100)
+
+ print(f"Comparing result: Expected -1, Got {result}")
+ assert result == -1
+
+
+class TestParseBracketToolCalls:
+ """Tests for parse_bracket_tool_calls function."""
+
+ def test_parses_single_tool_call(self):
+ """
+ What it does: Tests parsing of a single tool call.
+ Goal: Ensure bracket-style tool call is extracted correctly.
+ """
+ print("Setup: Text with one tool call...")
+ text = '[Called get_weather with args: {"location": "Moscow"}]'
+
+ print("Action: Parsing tool calls...")
+ result = parse_bracket_tool_calls(text)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["function"]["name"] == "get_weather"
+ assert '"location"' in result[0]["function"]["arguments"]
+
+ def test_parses_multiple_tool_calls(self):
+ """
+ What it does: Tests parsing of multiple tool calls.
+ Goal: Ensure all tool calls are extracted.
+ """
+ print("Setup: Text with multiple tool calls...")
+ text = '''
+ [Called get_weather with args: {"location": "Moscow"}]
+ Some text in between
+ [Called get_time with args: {"timezone": "UTC"}]
+ '''
+
+ print("Action: Parsing tool calls...")
+ result = parse_bracket_tool_calls(text)
+
+ print(f"Result: {result}")
+ assert len(result) == 2
+ assert result[0]["function"]["name"] == "get_weather"
+ assert result[1]["function"]["name"] == "get_time"
+
+ def test_returns_empty_for_no_tool_calls(self):
+ """
+ What it does: Tests returning empty list without tool calls.
+ Goal: Ensure regular text is not parsed as tool call.
+ """
+ print("Setup: Regular text without tool calls...")
+ text = "This is just regular text without any tool calls."
+
+ print("Action: Parsing tool calls...")
+ result = parse_bracket_tool_calls(text)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_returns_empty_for_empty_string(self):
+ """
+ What it does: Tests handling of empty string.
+ Goal: Ensure empty string doesn't cause errors.
+ """
+ print("Setup: Empty string...")
+
+ print("Action: Parsing tool calls...")
+ result = parse_bracket_tool_calls("")
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_returns_empty_for_none(self):
+ """
+ What it does: Tests handling of None.
+ Goal: Ensure None doesn't cause errors.
+ """
+ print("Setup: None...")
+
+ print("Action: Parsing tool calls...")
+ result = parse_bracket_tool_calls(None)
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_handles_nested_json_in_args(self):
+ """
+ What it does: Tests parsing of nested JSON in arguments.
+ Goal: Ensure complex arguments are parsed correctly.
+ """
+ print("Setup: Tool call with nested JSON...")
+ text = '[Called complex_func with args: {"data": {"nested": {"deep": "value"}}}]'
+
+ print("Action: Parsing tool calls...")
+ result = parse_bracket_tool_calls(text)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["function"]["name"] == "complex_func"
+ assert "nested" in result[0]["function"]["arguments"]
+
+ def test_generates_unique_ids(self):
+ """
+ What it does: Tests generation of unique IDs for tool calls.
+ Goal: Ensure each tool call has a unique ID.
+ """
+ print("Setup: Two identical tool calls...")
+ text = '''
+ [Called func with args: {"a": 1}]
+ [Called func with args: {"a": 1}]
+ '''
+
+ print("Action: Parsing tool calls...")
+ result = parse_bracket_tool_calls(text)
+
+ print(f"IDs: {[r['id'] for r in result]}")
+ assert len(result) == 2
+ assert result[0]["id"] != result[1]["id"]
+
+
+class TestDeduplicateToolCalls:
+ """Tests for deduplicate_tool_calls function."""
+
+ def test_removes_duplicates(self):
+ """
+ What it does: Tests removal of duplicates.
+ Goal: Ensure identical tool calls are removed.
+ """
+ print("Setup: List with duplicates...")
+ tool_calls = [
+ {"id": "1", "function": {"name": "func", "arguments": '{"a": 1}'}},
+ {"id": "2", "function": {"name": "func", "arguments": '{"a": 1}'}},
+ {"id": "3", "function": {"name": "other", "arguments": '{"b": 2}'}},
+ ]
+
+ print("Action: Deduplication...")
+ result = deduplicate_tool_calls(tool_calls)
+
+ print(f"Comparing length: Expected 2, Got {len(result)}")
+ assert len(result) == 2
+
+ def test_preserves_first_occurrence(self):
+ """
+ What it does: Tests preservation of first occurrence.
+ Goal: Ensure the first tool call from duplicates is preserved.
+ """
+ print("Setup: List with duplicates...")
+ tool_calls = [
+ {"id": "first", "function": {"name": "func", "arguments": '{"a": 1}'}},
+ {"id": "second", "function": {"name": "func", "arguments": '{"a": 1}'}},
+ ]
+
+ print("Action: Deduplication...")
+ result = deduplicate_tool_calls(tool_calls)
+
+ print(f"Comparing ID: Expected 'first', Got '{result[0]['id']}'")
+ assert result[0]["id"] == "first"
+
+ def test_handles_empty_list(self):
+ """
+ What it does: Tests handling of empty list.
+ Goal: Ensure empty list doesn't cause errors.
+ """
+ print("Setup: Empty list...")
+
+ print("Action: Deduplication...")
+ result = deduplicate_tool_calls([])
+
+ print(f"Comparing result: Expected [], Got {result}")
+ assert result == []
+
+ def test_deduplicates_by_id_keeps_one_with_arguments(self):
+ """
+ What it does: Tests deduplication by id keeping tool call with arguments.
+ Goal: Ensure that when duplicates by id exist, the one with arguments is kept.
+ """
+ print("Setup: Two tool calls with same id, one with arguments, one empty...")
+ tool_calls = [
+ {"id": "call_123", "function": {"name": "func", "arguments": "{}"}},
+ {"id": "call_123", "function": {"name": "func", "arguments": '{"location": "Moscow"}'}},
+ ]
+
+ print("Action: Deduplication...")
+ result = deduplicate_tool_calls(tool_calls)
+
+ print(f"Result: {result}")
+ print(f"Comparing length: Expected 1, Got {len(result)}")
+ assert len(result) == 1
+
+ print("Verifying that tool call with arguments was kept...")
+ assert "Moscow" in result[0]["function"]["arguments"]
+
+ def test_deduplicates_by_id_prefers_longer_arguments(self):
+ """
+ What it does: Tests that duplicates by id prefer longer arguments.
+ Goal: Ensure tool call with more complete arguments is kept.
+ """
+ print("Setup: Two tool calls with same id, different argument lengths...")
+ tool_calls = [
+ {"id": "call_abc", "function": {"name": "search", "arguments": '{"q": "test"}'}},
+ {"id": "call_abc", "function": {"name": "search", "arguments": '{"q": "test", "limit": 10, "offset": 0}'}},
+ ]
+
+ print("Action: Deduplication...")
+ result = deduplicate_tool_calls(tool_calls)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+
+ print("Verifying that tool call with longer arguments was kept...")
+ assert "limit" in result[0]["function"]["arguments"]
+
+ def test_deduplicates_empty_arguments_replaced_by_non_empty(self):
+ """
+ What it does: Tests replacement of empty arguments with non-empty.
+ Goal: Ensure "{}" is replaced with actual arguments.
+ """
+ print("Setup: First tool call with empty arguments, second with real ones...")
+ tool_calls = [
+ {"id": "call_xyz", "function": {"name": "get_weather", "arguments": "{}"}},
+ {"id": "call_xyz", "function": {"name": "get_weather", "arguments": '{"city": "London"}'}},
+ ]
+
+ print("Action: Deduplication...")
+ result = deduplicate_tool_calls(tool_calls)
+
+ print(f"Result: {result}")
+ assert len(result) == 1
+ assert result[0]["function"]["arguments"] == '{"city": "London"}'
+
+ def test_handles_tool_calls_without_id(self):
+ """
+ What it does: Tests handling of tool calls without id.
+ Goal: Ensure tool calls without id are deduplicated by name+arguments.
+ """
+ print("Setup: Tool calls without id...")
+ tool_calls = [
+ {"id": "", "function": {"name": "func", "arguments": '{"a": 1}'}},
+ {"id": "", "function": {"name": "func", "arguments": '{"a": 1}'}},
+ {"id": "", "function": {"name": "func", "arguments": '{"b": 2}'}},
+ ]
+
+ print("Action: Deduplication...")
+ result = deduplicate_tool_calls(tool_calls)
+
+ print(f"Result: {result}")
+ # Two unique by name+arguments
+ assert len(result) == 2
+
+ def test_mixed_with_and_without_id(self):
+ """
+ What it does: Tests mixed list with and without id.
+ Goal: Ensure both types are handled correctly.
+ """
+ print("Setup: Mixed list...")
+ tool_calls = [
+ {"id": "call_1", "function": {"name": "func1", "arguments": '{"x": 1}'}},
+ {"id": "call_1", "function": {"name": "func1", "arguments": "{}"}}, # Duplicate by id
+ {"id": "", "function": {"name": "func2", "arguments": '{"y": 2}'}},
+ {"id": "", "function": {"name": "func2", "arguments": '{"y": 2}'}}, # Duplicate by name+args
+ ]
+
+ print("Action: Deduplication...")
+ result = deduplicate_tool_calls(tool_calls)
+
+ print(f"Result: {result}")
+ # call_1 with arguments + func2 once
+ assert len(result) == 2
+
+ # Verify that call_1 kept its arguments
+ call_1 = next(tc for tc in result if tc["id"] == "call_1")
+ assert call_1["function"]["arguments"] == '{"x": 1}'
+
+
+class TestAwsEventStreamParserInitialization:
+ """Tests for AwsEventStreamParser initialization."""
+
+ def test_initialization_creates_empty_state(self):
+ """
+ What it does: Tests initial parser state.
+ Goal: Ensure parser is created with empty state.
+ """
+ print("Setup: Creating parser...")
+ parser = AwsEventStreamParser()
+
+ print("Check: Buffer is empty...")
+ assert parser.buffer == ""
+
+ print("Check: last_content is None...")
+ assert parser.last_content is None
+
+ print("Check: current_tool_call is None...")
+ assert parser.current_tool_call is None
+
+ print("Check: tool_calls is empty...")
+ assert parser.tool_calls == []
+
+
+class TestAwsEventStreamParserFeed:
+ """Tests for parser feed method."""
+
+ def test_parses_content_event(self, aws_event_parser):
+ """
+ What it does: Tests parsing of content event.
+ Goal: Ensure text content is extracted.
+ """
+ print("Setup: Chunk with content...")
+ chunk = b'{"content":"Hello World"}'
+
+ print("Action: Parsing chunk...")
+ events = aws_event_parser.feed(chunk)
+
+ print(f"Result: {events}")
+ assert len(events) == 1
+ assert events[0]["type"] == "content"
+ assert events[0]["data"] == "Hello World"
+
+ def test_parses_multiple_content_events(self, aws_event_parser):
+ """
+ What it does: Tests parsing of multiple content events.
+ Goal: Ensure all events are extracted.
+ """
+ print("Setup: Chunk with multiple events...")
+ chunk = b'{"content":"First"}{"content":"Second"}'
+
+ print("Action: Parsing chunk...")
+ events = aws_event_parser.feed(chunk)
+
+ print(f"Result: {events}")
+ assert len(events) == 2
+ assert events[0]["data"] == "First"
+ assert events[1]["data"] == "Second"
+
+ def test_deduplicates_repeated_content(self, aws_event_parser):
+ """
+ What it does: Tests deduplication of repeated content.
+ Goal: Ensure identical content is not duplicated.
+ """
+ print("Setup: Chunks with repeated content...")
+
+ print("Action: Parsing first chunk...")
+ events1 = aws_event_parser.feed(b'{"content":"Same"}')
+
+ print("Action: Parsing second chunk with same content...")
+ events2 = aws_event_parser.feed(b'{"content":"Same"}')
+
+ print(f"First result: {events1}")
+ print(f"Second result: {events2}")
+ assert len(events1) == 1
+ assert len(events2) == 0 # Duplicate filtered out
+
+ def test_parses_usage_event(self, aws_event_parser):
+ """
+ What it does: Tests parsing of usage event.
+ Goal: Ensure credits information is extracted.
+ """
+ print("Setup: Chunk with usage...")
+ chunk = b'{"usage":1.5}'
+
+ print("Action: Parsing chunk...")
+ events = aws_event_parser.feed(chunk)
+
+ print(f"Result: {events}")
+ assert len(events) == 1
+ assert events[0]["type"] == "usage"
+ assert events[0]["data"] == 1.5
+
+ def test_parses_context_usage_event(self, aws_event_parser):
+ """
+ What it does: Tests parsing of context_usage event.
+ Goal: Ensure context usage percentage is extracted.
+ """
+ print("Setup: Chunk with context usage...")
+ chunk = b'{"contextUsagePercentage":25.5}'
+
+ print("Action: Parsing chunk...")
+ events = aws_event_parser.feed(chunk)
+
+ print(f"Result: {events}")
+ assert len(events) == 1
+ assert events[0]["type"] == "context_usage"
+ assert events[0]["data"] == 25.5
+
+ def test_handles_incomplete_json(self, aws_event_parser):
+ """
+ What it does: Tests handling of incomplete JSON.
+ Goal: Ensure incomplete JSON is buffered.
+ """
+ print("Setup: Incomplete chunk...")
+ chunk = b'{"content":"Hel'
+
+ print("Action: Parsing incomplete chunk...")
+ events = aws_event_parser.feed(chunk)
+
+ print(f"Result: {events}")
+ assert len(events) == 0 # Nothing parsed
+
+ print("Check: Data in buffer...")
+ assert 'content' in aws_event_parser.buffer
+
+ def test_completes_json_across_chunks(self, aws_event_parser):
+ """
+ What it does: Tests assembling JSON from multiple chunks.
+ Goal: Ensure JSON is assembled from parts.
+ """
+ print("Setup: First part of JSON...")
+ events1 = aws_event_parser.feed(b'{"content":"Hel')
+
+ print("Action: Second part of JSON...")
+ events2 = aws_event_parser.feed(b'lo World"}')
+
+ print(f"First result: {events1}")
+ print(f"Second result: {events2}")
+ assert len(events1) == 0
+ assert len(events2) == 1
+ assert events2[0]["data"] == "Hello World"
+
+ def test_decodes_escape_sequences(self, aws_event_parser):
+ """
+ What it does: Tests decoding of escape sequences.
+ Goal: Ensure \\n is converted to actual newline.
+ """
+ print("Setup: Chunk with escape sequence...")
+ # Using correct escape sequence format
+ chunk = b'{"content":"Line1\\nLine2"}'
+
+ print("Action: Parsing chunk...")
+ events = aws_event_parser.feed(chunk)
+
+ print(f"Result: {events}")
+ assert len(events) == 1
+ assert "\n" in events[0]["data"]
+ def test_handles_invalid_bytes(self, aws_event_parser):
+ """
+ What it does: Tests handling of invalid bytes.
+ Goal: Ensure invalid data doesn't break the parser.
+ """
+ print("Setup: Invalid bytes...")
+ chunk = b'\xff\xfe{"content":"test"}'
+
+ print("Action: Parsing chunk...")
+ events = aws_event_parser.feed(chunk)
+
+ print(f"Result: {events}")
+ # Parser should continue working
+ assert len(events) == 1
+
+
+class TestAwsEventStreamParserToolCalls:
+ """Tests for tool calls parsing."""
+
+ def test_parses_tool_start_event(self, aws_event_parser):
+ """
+ What it does: Tests parsing of tool call start.
+ Goal: Ensure tool_start creates current_tool_call.
+ """
+ print("Setup: Chunk with tool call start...")
+ chunk = b'{"name":"get_weather","toolUseId":"call_123"}'
+
+ print("Action: Parsing chunk...")
+ events = aws_event_parser.feed(chunk)
+
+ print(f"Result: {events}")
+ print(f"current_tool_call: {aws_event_parser.current_tool_call}")
+
+ # tool_start doesn't return event, but creates current_tool_call
+ assert aws_event_parser.current_tool_call is not None
+ assert aws_event_parser.current_tool_call["function"]["name"] == "get_weather"
+
+ def test_parses_tool_input_event(self, aws_event_parser):
+ """
+ What it does: Tests parsing of input for tool call.
+ Goal: Ensure input is added to current_tool_call.
+ """
+ print("Setup: Tool call start...")
+ aws_event_parser.feed(b'{"name":"func","toolUseId":"call_1"}')
+
+ print("Action: Parsing input...")
+ aws_event_parser.feed(b'{"input":"{\\"key\\": \\"value\\"}"}')
+
+ print(f"current_tool_call: {aws_event_parser.current_tool_call}")
+ assert '{"key": "value"}' in aws_event_parser.current_tool_call["function"]["arguments"]
+
+ def test_parses_tool_stop_event(self, aws_event_parser):
+ """
+ What it does: Tests tool call completion.
+ Goal: Ensure tool call is added to the list.
+ """
+ print("Setup: Complete tool call...")
+ aws_event_parser.feed(b'{"name":"func","toolUseId":"call_1"}')
+ aws_event_parser.feed(b'{"input":"{}"}')
+
+ print("Action: Parsing stop...")
+ aws_event_parser.feed(b'{"stop":true}')
+
+ print(f"tool_calls: {aws_event_parser.tool_calls}")
+ assert len(aws_event_parser.tool_calls) == 1
+ assert aws_event_parser.current_tool_call is None
+
+ def test_get_tool_calls_returns_all(self, aws_event_parser):
+ """
+ What it does: Tests getting all tool calls.
+ Goal: Ensure get_tool_calls returns completed calls.
+ """
+ print("Setup: Multiple tool calls...")
+ aws_event_parser.feed(b'{"name":"func1","toolUseId":"call_1"}')
+ aws_event_parser.feed(b'{"stop":true}')
+ aws_event_parser.feed(b'{"name":"func2","toolUseId":"call_2"}')
+ aws_event_parser.feed(b'{"stop":true}')
+
+ print("Action: Getting tool calls...")
+ tool_calls = aws_event_parser.get_tool_calls()
+
+ print(f"Result: {tool_calls}")
+ assert len(tool_calls) == 2
+
+ def test_get_tool_calls_finalizes_current(self, aws_event_parser):
+ """
+ What it does: Tests finalization of incomplete tool call.
+ Goal: Ensure get_tool_calls finalizes current_tool_call.
+ """
+ print("Setup: Incomplete tool call...")
+ aws_event_parser.feed(b'{"name":"func","toolUseId":"call_1"}')
+
+ print("Action: Getting tool calls...")
+ tool_calls = aws_event_parser.get_tool_calls()
+
+ print(f"Result: {tool_calls}")
+ assert len(tool_calls) == 1
+ assert aws_event_parser.current_tool_call is None
+
+
+class TestAwsEventStreamParserReset:
+ """Tests for reset method."""
+
+ def test_reset_clears_state(self, aws_event_parser):
+ """
+ What it does: Tests parser state reset.
+ Goal: Ensure reset clears all data.
+ """
+ print("Setup: Filling parser with data...")
+ aws_event_parser.feed(b'{"content":"test"}')
+ aws_event_parser.feed(b'{"name":"func","toolUseId":"call_1"}')
+
+ print("Action: Resetting parser...")
+ aws_event_parser.reset()
+
+ print("Check: All data cleared...")
+ assert aws_event_parser.buffer == ""
+ assert aws_event_parser.last_content is None
+ assert aws_event_parser.current_tool_call is None
+ assert aws_event_parser.tool_calls == []
+
+
+class TestAwsEventStreamParserFinalizeToolCall:
+ """Tests for _finalize_tool_call method handling different input types."""
+
+ def test_finalize_with_string_arguments(self, aws_event_parser):
+ """
+ What it does: Tests finalization of tool call with string arguments.
+ Goal: Ensure JSON string is parsed and serialized back.
+ """
+ print("Setup: Tool call with string arguments...")
+ aws_event_parser.current_tool_call = {
+ "id": "call_1",
+ "type": "function",
+ "function": {
+ "name": "test_func",
+ "arguments": '{"key": "value"}'
+ }
+ }
+
+ print("Action: Finalizing tool call...")
+ aws_event_parser._finalize_tool_call()
+
+ print(f"Result: {aws_event_parser.tool_calls}")
+ assert len(aws_event_parser.tool_calls) == 1
+ assert aws_event_parser.tool_calls[0]["function"]["arguments"] == '{"key": "value"}'
+
+ def test_finalize_with_dict_arguments(self, aws_event_parser):
+ """
+ What it does: Tests finalization of tool call with dict arguments.
+ Goal: Ensure dict is serialized to JSON string.
+ """
+ print("Setup: Tool call with dict arguments...")
+ aws_event_parser.current_tool_call = {
+ "id": "call_2",
+ "type": "function",
+ "function": {
+ "name": "test_func",
+ "arguments": {"location": "Moscow", "units": "celsius"}
+ }
+ }
+
+ print("Action: Finalizing tool call...")
+ aws_event_parser._finalize_tool_call()
+
+ print(f"Result: {aws_event_parser.tool_calls}")
+ assert len(aws_event_parser.tool_calls) == 1
+
+ args = aws_event_parser.tool_calls[0]["function"]["arguments"]
+ print(f"Arguments: {args}")
+ assert isinstance(args, str)
+ assert "Moscow" in args
+ assert "celsius" in args
+
+ def test_finalize_with_empty_string_arguments(self, aws_event_parser):
+ """
+ What it does: Tests finalization of tool call with empty string arguments.
+ Goal: Ensure empty string is replaced with "{}".
+ """
+ print("Setup: Tool call with empty string arguments...")
+ aws_event_parser.current_tool_call = {
+ "id": "call_3",
+ "type": "function",
+ "function": {
+ "name": "test_func",
+ "arguments": ""
+ }
+ }
+
+ print("Action: Finalizing tool call...")
+ aws_event_parser._finalize_tool_call()
+
+ print(f"Result: {aws_event_parser.tool_calls}")
+ assert len(aws_event_parser.tool_calls) == 1
+ assert aws_event_parser.tool_calls[0]["function"]["arguments"] == "{}"
+
+ def test_finalize_with_whitespace_only_arguments(self, aws_event_parser):
+ """
+ What it does: Tests finalization of tool call with whitespace arguments.
+ Goal: Ensure whitespace string is replaced with "{}".
+ """
+ print("Setup: Tool call with whitespace arguments...")
+ aws_event_parser.current_tool_call = {
+ "id": "call_4",
+ "type": "function",
+ "function": {
+ "name": "test_func",
+ "arguments": " "
+ }
+ }
+
+ print("Action: Finalizing tool call...")
+ aws_event_parser._finalize_tool_call()
+
+ print(f"Result: {aws_event_parser.tool_calls}")
+ assert len(aws_event_parser.tool_calls) == 1
+ assert aws_event_parser.tool_calls[0]["function"]["arguments"] == "{}"
+
+ def test_finalize_with_invalid_json_arguments(self, aws_event_parser):
+ """
+ What it does: Tests finalization of tool call with invalid JSON.
+ Goal: Ensure invalid JSON is replaced with "{}".
+ """
+ print("Setup: Tool call with invalid JSON...")
+ aws_event_parser.current_tool_call = {
+ "id": "call_5",
+ "type": "function",
+ "function": {
+ "name": "test_func",
+ "arguments": "not valid json {"
+ }
+ }
+
+ print("Action: Finalizing tool call...")
+ aws_event_parser._finalize_tool_call()
+
+ print(f"Result: {aws_event_parser.tool_calls}")
+ assert len(aws_event_parser.tool_calls) == 1
+ assert aws_event_parser.tool_calls[0]["function"]["arguments"] == "{}"
+
+ def test_finalize_with_none_current_tool_call(self, aws_event_parser):
+ """
+ What it does: Tests finalization when current_tool_call is None.
+ Goal: Ensure nothing happens with None.
+ """
+ print("Setup: current_tool_call = None...")
+ aws_event_parser.current_tool_call = None
+
+ print("Action: Finalizing tool call...")
+ aws_event_parser._finalize_tool_call()
+
+ print(f"Result: {aws_event_parser.tool_calls}")
+ assert len(aws_event_parser.tool_calls) == 0
+
+ def test_finalize_clears_current_tool_call(self, aws_event_parser):
+ """
+ What it does: Tests that finalization clears current_tool_call.
+ Goal: Ensure current_tool_call = None after finalization.
+ """
+ print("Setup: Tool call...")
+ aws_event_parser.current_tool_call = {
+ "id": "call_6",
+ "type": "function",
+ "function": {
+ "name": "test_func",
+ "arguments": "{}"
+ }
+ }
+
+ print("Action: Finalizing tool call...")
+ aws_event_parser._finalize_tool_call()
+
+ print(f"current_tool_call after finalization: {aws_event_parser.current_tool_call}")
+ assert aws_event_parser.current_tool_call is None
+
+
+class TestAwsEventStreamParserEdgeCases:
+ """Tests for edge cases."""
+
+ def test_handles_followup_prompt(self, aws_event_parser):
+ """
+ What it does: Tests ignoring followupPrompt.
+ Goal: Ensure followupPrompt doesn't create an event.
+ """
+ print("Setup: Chunk with followupPrompt...")
+ chunk = b'{"content":"text","followupPrompt":"suggestion"}'
+
+ print("Action: Parsing chunk...")
+ events = aws_event_parser.feed(chunk)
+
+ print(f"Result: {events}")
+ assert len(events) == 0 # followupPrompt is ignored
+
+ def test_handles_mixed_events(self, aws_event_parser):
+ """
+ What it does: Tests parsing of mixed events.
+ Goal: Ensure different event types are handled together.
+ """
+ print("Setup: Chunk with mixed events...")
+ chunk = b'{"content":"Hello"}{"usage":1.0}{"contextUsagePercentage":50}'
+
+ print("Action: Parsing chunk...")
+ events = aws_event_parser.feed(chunk)
+
+ print(f"Result: {events}")
+ assert len(events) == 3
+ assert events[0]["type"] == "content"
+ assert events[1]["type"] == "usage"
+ assert events[2]["type"] == "context_usage"
+
+ def test_handles_garbage_between_events(self, aws_event_parser):
+ """
+ What it does: Tests handling of garbage between events.
+ Goal: Ensure parser finds JSON among garbage.
+ """
+ print("Setup: Chunk with garbage between JSON...")
+ chunk = b'garbage{"content":"valid"}more garbage{"usage":1}'
+
+ print("Action: Parsing chunk...")
+ events = aws_event_parser.feed(chunk)
+
+ print(f"Result: {events}")
+ assert len(events) == 2
+
+ def test_handles_empty_chunk(self, aws_event_parser):
+ """
+ What it does: Tests handling of empty chunk.
+ Goal: Ensure empty chunk doesn't cause errors.
+ """
+ print("Setup: Empty chunk...")
+
+ print("Action: Parsing empty chunk...")
+ events = aws_event_parser.feed(b'')
+
+ print(f"Comparing result: Expected [], Got {events}")
+ assert events == []
+
+
+class TestDiagnoseJsonTruncation:
+ """
+ Tests for _diagnose_json_truncation method for diagnosing truncated JSON.
+
+ This method helps distinguish upstream issues (Kiro API truncates large
+ tool call arguments) from actually invalid JSON from the model.
+ """
+
+ def test_empty_string_not_truncated(self, aws_event_parser):
+ """
+ What it does: Tests handling of empty string.
+ Goal: Ensure empty string is not considered truncated.
+ """
+ print("Setup: Empty string...")
+ json_str = ""
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ print(f"Comparing is_truncated: Expected False, Got {result['is_truncated']}")
+ assert result["is_truncated"] is False
+ assert result["reason"] == "empty string"
+ assert result["size_bytes"] == 0
+
+ def test_whitespace_only_not_truncated(self, aws_event_parser):
+ """
+ What it does: Tests handling of whitespace-only string.
+ Goal: Ensure whitespace string is not considered truncated.
+ """
+ print("Setup: Whitespace string...")
+ json_str = " \t\n "
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ print(f"Comparing is_truncated: Expected False, Got {result['is_truncated']}")
+ assert result["is_truncated"] is False
+ assert result["reason"] == "empty string"
+
+ def test_valid_json_not_truncated(self, aws_event_parser):
+ """
+ What it does: Tests handling of valid JSON.
+ Goal: Ensure valid JSON is not considered truncated.
+ """
+ print("Setup: Valid JSON...")
+ json_str = '{"key": "value", "number": 42}'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ print(f"Comparing is_truncated: Expected False, Got {result['is_truncated']}")
+ assert result["is_truncated"] is False
+ assert result["reason"] == "malformed JSON" # Function doesn't check validity, only structure
+
+ def test_valid_nested_json_not_truncated(self, aws_event_parser):
+ """
+ What it does: Tests handling of nested valid JSON.
+ Goal: Ensure complex JSON is not considered truncated.
+ """
+ print("Setup: Nested valid JSON...")
+ json_str = '{"outer": {"inner": {"deep": [1, 2, 3]}}}'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ assert result["is_truncated"] is False
+
+ def test_missing_closing_brace_truncated(self, aws_event_parser):
+ """
+ What it does: Tests detection of missing closing brace.
+ Goal: Ensure JSON without closing } is considered truncated.
+ """
+ print("Setup: JSON without closing brace...")
+ json_str = '{"filePath": "/path/to/file.md"'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ print(f"Comparing is_truncated: Expected True, Got {result['is_truncated']}")
+ assert result["is_truncated"] is True
+ assert "missing" in result["reason"] and "brace" in result["reason"]
+
+ def test_real_world_truncation_from_issue_34(self, aws_event_parser):
+ """
+ What it does: Tests real example from Issue #34.
+ Goal: Ensure real truncated JSON from bug is detected.
+ """
+ print("Setup: Real example from Issue #34...")
+ # This is exact example from log: JSON truncated after filePath
+ json_str = '{"filePath": "/Users/cc/Documents/Code/mock-all/docs/plans/2026-01-12-mock-all-impl.md"'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ print(f"Comparing is_truncated: Expected True, Got {result['is_truncated']}")
+ assert result["is_truncated"] is True
+ assert "brace" in result["reason"]
+ assert result["size_bytes"] == 87 # Exact size from log (char 87 = error position)
+
+ def test_multiple_missing_braces_truncated(self, aws_event_parser):
+ """
+ What it does: Tests detection of multiple missing braces.
+ Goal: Ensure nested JSON without closing braces is detected.
+ """
+ print("Setup: Nested JSON without closing braces...")
+ json_str = '{"outer": {"inner": {"deep": "value"'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ assert result["is_truncated"] is True
+ assert "3" in result["reason"] or "brace" in result["reason"]
+
+ def test_missing_closing_bracket_truncated(self, aws_event_parser):
+ """
+ What it does: Tests detection of missing closing square bracket.
+ Goal: Ensure array without ] is considered truncated.
+ """
+ print("Setup: Array without closing bracket...")
+ json_str = '[1, 2, 3, {"key": "value"}'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ print(f"Comparing is_truncated: Expected True, Got {result['is_truncated']}")
+ assert result["is_truncated"] is True
+ assert "bracket" in result["reason"]
+
+ def test_array_start_truncated(self, aws_event_parser):
+ """
+ What it does: Tests detection of truncated array at start.
+ Goal: Ensure [ without ] is detected.
+ """
+ print("Setup: Array start without end...")
+ json_str = '["item1", "item2"'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ assert result["is_truncated"] is True
+ assert "bracket" in result["reason"]
+
+ def test_unbalanced_braces_truncated(self, aws_event_parser):
+ """
+ What it does: Tests detection of unbalanced curly braces.
+ Goal: Ensure different count of { and } is detected.
+ """
+ print("Setup: JSON with unbalanced braces...")
+ # Ends with }, but has extra opening inside
+ json_str = '{"a": {"b": 1}}'[:-1] # Remove last }
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ assert result["is_truncated"] is True
+
+ def test_unbalanced_brackets_truncated(self, aws_event_parser):
+ """
+ What it does: Tests detection of unbalanced square brackets.
+ Goal: Ensure different count of [ and ] is detected.
+ """
+ print("Setup: JSON with unbalanced square brackets...")
+ json_str = '{"items": [[1, 2], [3, 4]}' # Missing one ]
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ assert result["is_truncated"] is True
+ assert "bracket" in result["reason"]
+
+ def test_unclosed_string_truncated(self, aws_event_parser):
+ """
+ What it does: Tests detection of unclosed string.
+ Goal: Ensure odd number of quotes is detected.
+ """
+ print("Setup: JSON with unclosed string...")
+ json_str = '{"content": "This is a very long string that was cut off'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ print(f"Comparing is_truncated: Expected True, Got {result['is_truncated']}")
+ assert result["is_truncated"] is True
+ assert "string" in result["reason"] or "brace" in result["reason"]
+
+ def test_escaped_quotes_handled_correctly(self, aws_event_parser):
+ """
+ What it does: Tests correct handling of escaped quotes.
+ Goal: Ensure \\" doesn't break quote counting.
+ """
+ print("Setup: JSON with escaped quotes...")
+ json_str = '{"text": "Say \\"hello\\" to everyone"}'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ print(f"Comparing is_truncated: Expected False, Got {result['is_truncated']}")
+ assert result["is_truncated"] is False
+
+ def test_truncated_in_middle_of_escaped_sequence(self, aws_event_parser):
+ """
+ What it does: Tests truncation in middle of escape sequence.
+ Goal: Ensure truncation after \\ is detected.
+ """
+ print("Setup: JSON truncated after backslash...")
+ json_str = '{"text": "Line1\\nLine2\\'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ assert result["is_truncated"] is True
+
+ def test_size_bytes_calculated_correctly(self, aws_event_parser):
+ """
+ What it does: Tests correct byte size calculation.
+ Goal: Ensure UTF-8 characters are counted correctly.
+ """
+ print("Setup: JSON with Unicode characters...")
+ json_str = '{"city": "Москва"' # Cyrillic = 2 bytes per character
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ expected_size = len(json_str.encode('utf-8'))
+ print(f"Comparing size_bytes: Expected {expected_size}, Got {result['size_bytes']}")
+ assert result["size_bytes"] == expected_size
+ assert result["is_truncated"] is True # No closing }
+
+ def test_large_truncated_json(self, aws_event_parser):
+ """
+ What it does: Tests handling of large truncated JSON.
+ Goal: Ensure large data is handled correctly.
+ """
+ print("Setup: Large truncated JSON...")
+ # Simulate large file that was truncated
+ content = "x" * 10000
+ json_str = f'{{"filePath": "/path/to/file.md", "content": "{content}'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: is_truncated={result['is_truncated']}, size_bytes={result['size_bytes']}")
+ assert result["is_truncated"] is True
+ assert result["size_bytes"] > 10000
+
+ def test_malformed_but_not_truncated(self, aws_event_parser):
+ """
+ What it does: Tests invalid but not truncated JSON.
+ Goal: Ensure syntax errors are not confused with truncation.
+ """
+ print("Setup: Invalid JSON (trailing comma)...")
+ json_str = '{"key": "value",}' # Trailing comma - invalid, but not truncated
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ print(f"Comparing is_truncated: Expected False, Got {result['is_truncated']}")
+ assert result["is_truncated"] is False
+ assert result["reason"] == "malformed JSON"
+
+ def test_json_with_only_opening_brace(self, aws_event_parser):
+ """
+ What it does: Tests JSON with only opening brace.
+ Goal: Ensure minimal truncated JSON is detected.
+ """
+ print("Setup: Only opening brace...")
+ json_str = '{'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ assert result["is_truncated"] is True
+ assert "brace" in result["reason"]
+
+ def test_json_with_only_opening_bracket(self, aws_event_parser):
+ """
+ What it does: Tests JSON with only opening square bracket.
+ Goal: Ensure minimal truncated array is detected.
+ """
+ print("Setup: Only opening square bracket...")
+ json_str = '['
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ assert result["is_truncated"] is True
+ assert "bracket" in result["reason"]
+
+ def test_braces_inside_string_not_counted(self, aws_event_parser):
+ """
+ What it does: Tests that braces inside strings don't affect counting.
+ Goal: Ensure "{}" inside string doesn't break diagnosis.
+
+ Note: Current implementation uses simplified counting,
+ which doesn't account for string context. This is a known limitation.
+ """
+ print("Setup: JSON with braces inside string...")
+ json_str = '{"text": "Hello {world}"}'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ # Function uses simplified counting, so this may be False
+ # Main thing - it doesn't crash and returns correct structure
+ assert "is_truncated" in result
+ assert "reason" in result
+ assert "size_bytes" in result
+
+ def test_complex_nested_truncation(self, aws_event_parser):
+ """
+ What it does: Tests complex nested truncated JSON.
+ Goal: Ensure deep nesting is handled.
+ """
+ print("Setup: Complex nested truncated JSON...")
+ json_str = '{"level1": {"level2": {"level3": [{"item": "value'
+
+ print("Action: Diagnosis...")
+ result = aws_event_parser._diagnose_json_truncation(json_str)
+
+ print(f"Result: {result}")
+ assert result["is_truncated"] is True
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_routes_anthropic.py b/kiro-gateway/tests/unit/test_routes_anthropic.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f0a71b3285f2310eb1d67691b93d322b6f978c0
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_routes_anthropic.py
@@ -0,0 +1,1084 @@
+
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for Anthropic API endpoints (routes_anthropic.py).
+
+Tests the following endpoint:
+- POST /v1/messages - Anthropic Messages API
+
+For OpenAI API tests, see test_routes_openai.py.
+"""
+
+import pytest
+from unittest.mock import AsyncMock, Mock, patch, MagicMock
+from datetime import datetime, timezone
+import json
+
+from fastapi import HTTPException
+from fastapi.testclient import TestClient
+
+from kiro.routes_anthropic import verify_anthropic_api_key, router
+from kiro.config import PROXY_API_KEY
+
+
+# =============================================================================
+# Tests for verify_anthropic_api_key function
+# =============================================================================
+
+class TestVerifyAnthropicApiKey:
+ """Tests for the verify_anthropic_api_key authentication function."""
+
+ @pytest.mark.asyncio
+ async def test_valid_x_api_key_returns_true(self):
+ """
+ What it does: Verifies that a valid x-api-key header passes authentication.
+ Purpose: Ensure Anthropic native authentication works.
+ """
+ print("Setup: Creating valid x-api-key...")
+
+ print("Action: Calling verify_anthropic_api_key...")
+ result = await verify_anthropic_api_key(x_api_key=PROXY_API_KEY, authorization=None)
+
+ print(f"Comparing result: Expected True, Got {result}")
+ assert result is True
+
+ @pytest.mark.asyncio
+ async def test_valid_bearer_token_returns_true(self):
+ """
+ What it does: Verifies that a valid Bearer token passes authentication.
+ Purpose: Ensure OpenAI-style authentication also works.
+ """
+ print("Setup: Creating valid Bearer token...")
+ valid_auth = f"Bearer {PROXY_API_KEY}"
+
+ print("Action: Calling verify_anthropic_api_key...")
+ result = await verify_anthropic_api_key(x_api_key=None, authorization=valid_auth)
+
+ print(f"Comparing result: Expected True, Got {result}")
+ assert result is True
+
+ @pytest.mark.asyncio
+ async def test_x_api_key_takes_precedence(self):
+ """
+ What it does: Verifies x-api-key is checked before Authorization header.
+ Purpose: Ensure Anthropic native auth has priority.
+ """
+ print("Setup: Both headers provided...")
+
+ print("Action: Calling verify_anthropic_api_key with both headers...")
+ result = await verify_anthropic_api_key(
+ x_api_key=PROXY_API_KEY,
+ authorization="Bearer wrong_key"
+ )
+
+ print(f"Comparing result: Expected True, Got {result}")
+ assert result is True
+
+ @pytest.mark.asyncio
+ async def test_invalid_x_api_key_raises_401(self):
+ """
+ What it does: Verifies that an invalid x-api-key is rejected.
+ Purpose: Ensure unauthorized access is blocked.
+ """
+ print("Setup: Creating invalid x-api-key...")
+
+ print("Action: Calling verify_anthropic_api_key with invalid key...")
+ with pytest.raises(HTTPException) as exc_info:
+ await verify_anthropic_api_key(x_api_key="wrong_key", authorization=None)
+
+ print(f"Checking: HTTPException with status 401...")
+ assert exc_info.value.status_code == 401
+
+ @pytest.mark.asyncio
+ async def test_invalid_bearer_token_raises_401(self):
+ """
+ What it does: Verifies that an invalid Bearer token is rejected.
+ Purpose: Ensure unauthorized access is blocked.
+ """
+ print("Setup: Creating invalid Bearer token...")
+
+ print("Action: Calling verify_anthropic_api_key with invalid token...")
+ with pytest.raises(HTTPException) as exc_info:
+ await verify_anthropic_api_key(x_api_key=None, authorization="Bearer wrong_key")
+
+ print(f"Checking: HTTPException with status 401...")
+ assert exc_info.value.status_code == 401
+
+ @pytest.mark.asyncio
+ async def test_missing_both_headers_raises_401(self):
+ """
+ What it does: Verifies that missing both headers is rejected.
+ Purpose: Ensure authentication is required.
+ """
+ print("Setup: No authentication headers...")
+
+ print("Action: Calling verify_anthropic_api_key with no headers...")
+ with pytest.raises(HTTPException) as exc_info:
+ await verify_anthropic_api_key(x_api_key=None, authorization=None)
+
+ print(f"Checking: HTTPException with status 401...")
+ assert exc_info.value.status_code == 401
+
+ @pytest.mark.asyncio
+ async def test_empty_x_api_key_raises_401(self):
+ """
+ What it does: Verifies that empty x-api-key is rejected.
+ Purpose: Ensure empty credentials are blocked.
+ """
+ print("Setup: Empty x-api-key...")
+
+ print("Action: Calling verify_anthropic_api_key with empty key...")
+ with pytest.raises(HTTPException) as exc_info:
+ await verify_anthropic_api_key(x_api_key="", authorization=None)
+
+ print(f"Checking: HTTPException with status 401...")
+ assert exc_info.value.status_code == 401
+
+ @pytest.mark.asyncio
+ async def test_error_response_format_is_anthropic_style(self):
+ """
+ What it does: Verifies error response follows Anthropic format.
+ Purpose: Ensure error format matches Anthropic API.
+ """
+ print("Setup: Invalid credentials...")
+
+ print("Action: Calling verify_anthropic_api_key...")
+ with pytest.raises(HTTPException) as exc_info:
+ await verify_anthropic_api_key(x_api_key="wrong", authorization=None)
+
+ print(f"Checking: Error format...")
+ detail = exc_info.value.detail
+ assert "type" in detail
+ assert "error" in detail
+ assert detail["error"]["type"] == "authentication_error"
+
+
+# =============================================================================
+# Tests for /v1/messages endpoint authentication
+# =============================================================================
+
+class TestMessagesAuthentication:
+ """Tests for authentication on /v1/messages endpoint."""
+
+ def test_messages_requires_authentication(self, test_client):
+ """
+ What it does: Verifies messages endpoint requires authentication.
+ Purpose: Ensure protected endpoint is secured.
+ """
+ print("Action: POST /v1/messages without auth...")
+ response = test_client.post(
+ "/v1/messages",
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 401
+
+ def test_messages_accepts_x_api_key(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies messages endpoint accepts x-api-key header.
+ Purpose: Ensure Anthropic native authentication works.
+ """
+ print("Action: POST /v1/messages with x-api-key...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Should pass auth (not 401)
+ assert response.status_code != 401
+
+ def test_messages_accepts_bearer_token(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies messages endpoint accepts Bearer token.
+ Purpose: Ensure OpenAI-style authentication also works.
+ """
+ print("Action: POST /v1/messages with Bearer token...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Should pass auth (not 401)
+ assert response.status_code != 401
+
+ def test_messages_rejects_invalid_x_api_key(self, test_client, invalid_proxy_api_key):
+ """
+ What it does: Verifies messages endpoint rejects invalid x-api-key.
+ Purpose: Ensure authentication is enforced.
+ """
+ print("Action: POST /v1/messages with invalid x-api-key...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": invalid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 401
+
+
+# =============================================================================
+# Tests for /v1/messages endpoint validation
+# =============================================================================
+
+class TestMessagesValidation:
+ """Tests for request validation on /v1/messages endpoint."""
+
+ def test_validates_missing_model(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies missing model field is rejected.
+ Purpose: Ensure model is required.
+ """
+ print("Action: POST /v1/messages without model...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 422
+
+ def test_validates_missing_max_tokens(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies missing max_tokens field is rejected.
+ Purpose: Ensure max_tokens is required (Anthropic API requirement).
+ """
+ print("Action: POST /v1/messages without max_tokens...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 422
+
+ def test_validates_missing_messages(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies missing messages field is rejected.
+ Purpose: Ensure messages are required.
+ """
+ print("Action: POST /v1/messages without messages...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 422
+
+ def test_validates_empty_messages_array(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies empty messages array is rejected.
+ Purpose: Ensure at least one message is required.
+ """
+ print("Action: POST /v1/messages with empty messages...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": []
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 422
+
+ def test_validates_invalid_json(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies invalid JSON is rejected.
+ Purpose: Ensure proper JSON parsing.
+ """
+ print("Action: POST /v1/messages with invalid JSON...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={
+ "x-api-key": valid_proxy_api_key,
+ "Content-Type": "application/json"
+ },
+ content=b"not valid json {{{}"
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 422
+
+ def test_validates_invalid_role(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies invalid message role is rejected.
+ Purpose: Anthropic model strictly validates role (only 'user' or 'assistant').
+ """
+ print("Action: POST /v1/messages with invalid role...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "invalid_role", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Anthropic model strictly validates role - only 'user' or 'assistant' allowed
+ assert response.status_code == 422
+
+ def test_accepts_valid_request_format(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies valid request format passes validation.
+ Purpose: Ensure Pydantic validation works correctly.
+ """
+ print("Action: POST /v1/messages with valid format...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Should pass validation (not 422)
+ assert response.status_code != 422
+
+
+# =============================================================================
+# Tests for /v1/messages system prompt
+# =============================================================================
+
+class TestMessagesSystemPrompt:
+ """Tests for system prompt handling on /v1/messages endpoint."""
+
+ def test_accepts_system_as_separate_field(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies system prompt as separate field is accepted.
+ Purpose: Ensure Anthropic-style system prompt works.
+ """
+ print("Action: POST /v1/messages with system field...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "system": "You are a helpful assistant.",
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Should pass validation
+ assert response.status_code != 422
+
+ def test_accepts_empty_system_prompt(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies empty system prompt is accepted.
+ Purpose: Ensure system prompt is optional.
+ """
+ print("Action: POST /v1/messages with empty system...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "system": "",
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Should pass validation
+ assert response.status_code != 422
+
+ def test_accepts_no_system_prompt(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies request without system prompt is accepted.
+ Purpose: Ensure system prompt is optional.
+ """
+ print("Action: POST /v1/messages without system field...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Should pass validation
+ assert response.status_code != 422
+
+
+# =============================================================================
+# Tests for /v1/messages content blocks
+# =============================================================================
+
+class TestMessagesContentBlocks:
+ """Tests for content block handling on /v1/messages endpoint."""
+
+ def test_accepts_string_content(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies string content is accepted.
+ Purpose: Ensure simple string content works.
+ """
+ print("Action: POST /v1/messages with string content...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_content_block_array(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies content block array is accepted.
+ Purpose: Ensure Anthropic content block format works.
+ """
+ print("Action: POST /v1/messages with content blocks...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "Hello"}
+ ]
+ }
+ ]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_multiple_content_blocks(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies multiple content blocks are accepted.
+ Purpose: Ensure complex content works.
+ """
+ print("Action: POST /v1/messages with multiple content blocks...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "First part"},
+ {"type": "text", "text": "Second part"}
+ ]
+ }
+ ]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+
+# =============================================================================
+# Tests for /v1/messages tool use
+# =============================================================================
+
+class TestMessagesToolUse:
+ """Tests for tool use on /v1/messages endpoint."""
+
+ def test_accepts_tool_definition(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies tool definition is accepted.
+ Purpose: Ensure Anthropic tool format works.
+ """
+ print("Action: POST /v1/messages with tools...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "What's the weather?"}],
+ "tools": [
+ {
+ "name": "get_weather",
+ "description": "Get weather for a location",
+ "input_schema": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string"}
+ },
+ "required": ["location"]
+ }
+ }
+ ]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_multiple_tools(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies multiple tools are accepted.
+ Purpose: Ensure multiple tool definitions work.
+ """
+ print("Action: POST /v1/messages with multiple tools...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}],
+ "tools": [
+ {
+ "name": "get_weather",
+ "description": "Get weather",
+ "input_schema": {"type": "object", "properties": {}}
+ },
+ {
+ "name": "get_time",
+ "description": "Get time",
+ "input_schema": {"type": "object", "properties": {}}
+ }
+ ]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_tool_result_message(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies tool result message is accepted.
+ Purpose: Ensure tool result handling works.
+ """
+ print("Action: POST /v1/messages with tool result...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [
+ {"role": "user", "content": "What's the weather?"},
+ {
+ "role": "assistant",
+ "content": [
+ {
+ "type": "tool_use",
+ "id": "call_123",
+ "name": "get_weather",
+ "input": {"location": "Moscow"}
+ }
+ ]
+ },
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": "call_123",
+ "content": "Sunny, 25°C"
+ }
+ ]
+ }
+ ]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+
+# =============================================================================
+# Tests for /v1/messages optional parameters
+# =============================================================================
+
+class TestMessagesOptionalParams:
+ """Tests for optional parameters on /v1/messages endpoint."""
+
+ def test_accepts_temperature_parameter(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies temperature parameter is accepted.
+ Purpose: Ensure temperature control works.
+ """
+ print("Action: POST /v1/messages with temperature...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}],
+ "temperature": 0.7
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_top_p_parameter(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies top_p parameter is accepted.
+ Purpose: Ensure nucleus sampling control works.
+ """
+ print("Action: POST /v1/messages with top_p...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}],
+ "top_p": 0.9
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_top_k_parameter(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies top_k parameter is accepted.
+ Purpose: Ensure top-k sampling control works.
+ """
+ print("Action: POST /v1/messages with top_k...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}],
+ "top_k": 40
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_stream_true(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies stream=true is accepted.
+ Purpose: Ensure streaming mode is supported.
+ """
+ print("Action: POST /v1/messages with stream=true...")
+
+ # Mock the streaming function to avoid real HTTP requests
+ async def mock_stream(*args, **kwargs):
+ yield 'event: message_start\ndata: {"type":"message_start"}\n\n'
+ yield 'event: message_stop\ndata: {"type":"message_stop"}\n\n'
+
+ # Create mock response for HTTP client
+ mock_response = MagicMock()
+ mock_response.status_code = 200
+
+ with patch('kiro.routes_anthropic.stream_kiro_to_anthropic', mock_stream), \
+ patch('kiro.http_client.KiroHttpClient.request_with_retry', return_value=mock_response):
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}],
+ "stream": True
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_stop_sequences(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies stop_sequences parameter is accepted.
+ Purpose: Ensure stop sequence control works.
+ """
+ print("Action: POST /v1/messages with stop_sequences...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}],
+ "stop_sequences": ["END", "STOP"]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_metadata(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies metadata parameter is accepted.
+ Purpose: Ensure metadata passing works.
+ """
+ print("Action: POST /v1/messages with metadata...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}],
+ "metadata": {"user_id": "test_user"}
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+
+# =============================================================================
+# Tests for /v1/messages anthropic-version header
+# =============================================================================
+
+class TestMessagesAnthropicVersion:
+ """Tests for anthropic-version header handling."""
+
+ def test_accepts_anthropic_version_header(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies anthropic-version header is accepted.
+ Purpose: Ensure Anthropic SDK compatibility.
+ """
+ print("Action: POST /v1/messages with anthropic-version header...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={
+ "x-api-key": valid_proxy_api_key,
+ "anthropic-version": "2023-06-01"
+ },
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Should pass validation
+ assert response.status_code != 422
+
+ def test_works_without_anthropic_version_header(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies request works without anthropic-version header.
+ Purpose: Ensure header is optional.
+ """
+ print("Action: POST /v1/messages without anthropic-version header...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Should pass validation
+ assert response.status_code != 422
+
+
+# =============================================================================
+# Tests for router integration
+# =============================================================================
+
+class TestAnthropicRouterIntegration:
+ """Tests for Anthropic router configuration and integration."""
+
+ def test_router_has_messages_endpoint(self):
+ """
+ What it does: Verifies messages endpoint is registered.
+ Purpose: Ensure endpoint is available.
+ """
+ print("Checking: Router endpoints...")
+ routes = [route.path for route in router.routes]
+
+ print(f"Found routes: {routes}")
+ assert "/v1/messages" in routes
+
+ def test_messages_endpoint_uses_post_method(self):
+ """
+ What it does: Verifies messages endpoint uses POST method.
+ Purpose: Ensure correct HTTP method.
+ """
+ print("Checking: HTTP methods...")
+ for route in router.routes:
+ if route.path == "/v1/messages":
+ print(f"Route /v1/messages methods: {route.methods}")
+ assert "POST" in route.methods
+ return
+ pytest.fail("Messages endpoint not found")
+
+ def test_router_has_anthropic_tag(self):
+ """
+ What it does: Verifies router has Anthropic API tag.
+ Purpose: Ensure proper API documentation grouping.
+ """
+ print("Checking: Router tags...")
+ print(f"Router tags: {router.tags}")
+ assert "Anthropic API" in router.tags
+
+
+# =============================================================================
+# Tests for conversation history
+# =============================================================================
+
+class TestMessagesConversationHistory:
+ """Tests for conversation history handling on /v1/messages endpoint."""
+
+ def test_accepts_multi_turn_conversation(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies multi-turn conversation is accepted.
+ Purpose: Ensure conversation history works.
+ """
+ print("Action: POST /v1/messages with conversation history...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [
+ {"role": "user", "content": "Hello"},
+ {"role": "assistant", "content": "Hi there!"},
+ {"role": "user", "content": "How are you?"}
+ ]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_long_conversation(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies long conversation is accepted.
+ Purpose: Ensure many messages work.
+ """
+ print("Action: POST /v1/messages with long conversation...")
+ messages = []
+ for i in range(10):
+ messages.append({"role": "user", "content": f"Message {i}"})
+ messages.append({"role": "assistant", "content": f"Response {i}"})
+ messages.append({"role": "user", "content": "Final question"})
+
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": messages
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+
+# =============================================================================
+# Tests for error response format
+# =============================================================================
+
+class TestMessagesErrorFormat:
+ """Tests for error response format on /v1/messages endpoint."""
+
+ def test_validation_error_format(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies validation error response format.
+ Purpose: Ensure errors follow expected format.
+ """
+ print("Action: POST /v1/messages with invalid request...")
+ response = test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5"
+ # Missing required fields
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ print(f"Response: {response.json()}")
+ assert response.status_code == 422
+
+ def test_auth_error_format_is_anthropic_style(self, test_client):
+ """
+ What it does: Verifies auth error follows Anthropic format.
+ Purpose: Ensure error format matches Anthropic API.
+ """
+ print("Action: POST /v1/messages without auth...")
+ response = test_client.post(
+ "/v1/messages",
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 1024,
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ print(f"Response: {response.json()}")
+ assert response.status_code == 401
+
+ # Check Anthropic error format
+ data = response.json()
+ assert "detail" in data
+ detail = data["detail"]
+ assert "type" in detail
+ assert "error" in detail
+
+
+# =============================================================================
+# Tests for HTTP client selection (issue #54)
+# =============================================================================
+
+class TestAnthropicHTTPClientSelection:
+ """
+ Tests for HTTP client selection in Anthropic routes (issue #54).
+
+ Verifies that streaming requests use per-request clients to avoid CLOSE_WAIT leak
+ when network interface changes (VPN disconnect/reconnect), while non-streaming
+ requests use shared client for connection pooling.
+ """
+
+ @patch('kiro.routes_anthropic.KiroHttpClient')
+ def test_streaming_uses_per_request_client(
+ self,
+ mock_kiro_http_client_class,
+ test_client,
+ valid_proxy_api_key
+ ):
+ """
+ What it does: Verifies streaming requests create per-request HTTP client.
+ Purpose: Prevent CLOSE_WAIT leak on VPN disconnect (issue #54).
+ """
+ print("\n--- Test: Anthropic streaming uses per-request client ---")
+
+ # Setup mock
+ mock_client_instance = AsyncMock()
+ mock_client_instance.request_with_retry = AsyncMock(
+ side_effect=Exception("Network blocked")
+ )
+ mock_client_instance.close = AsyncMock()
+ mock_kiro_http_client_class.return_value = mock_client_instance
+
+ print("Action: POST /v1/messages with stream=true...")
+ try:
+ test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 100,
+ "messages": [{"role": "user", "content": "Hello"}],
+ "stream": True
+ }
+ )
+ except Exception:
+ pass
+
+ 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 None, \
+ "Streaming should use per-request client"
+ print("✅ Anthropic streaming correctly uses per-request client")
+
+ @patch('kiro.routes_anthropic.KiroHttpClient')
+ def test_non_streaming_uses_shared_client(
+ self,
+ mock_kiro_http_client_class,
+ test_client,
+ valid_proxy_api_key
+ ):
+ """
+ What it does: Verifies non-streaming requests use shared HTTP client.
+ Purpose: Ensure connection pooling for non-streaming requests.
+ """
+ print("\n--- Test: Anthropic non-streaming uses shared client ---")
+
+ # Setup mock
+ mock_client_instance = AsyncMock()
+ mock_client_instance.request_with_retry = AsyncMock(
+ side_effect=Exception("Network blocked")
+ )
+ mock_client_instance.close = AsyncMock()
+ mock_kiro_http_client_class.return_value = mock_client_instance
+
+ print("Action: POST /v1/messages with stream=false...")
+ try:
+ test_client.post(
+ "/v1/messages",
+ headers={"x-api-key": valid_proxy_api_key},
+ json={
+ "model": "claude-sonnet-4-5",
+ "max_tokens": 100,
+ "messages": [{"role": "user", "content": "Hello"}],
+ "stream": False
+ }
+ )
+ except Exception:
+ pass
+
+ print("Checking: KiroHttpClient(shared_client=app.state.http_client)...")
+ 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")
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_routes_openai.py b/kiro-gateway/tests/unit/test_routes_openai.py
new file mode 100644
index 0000000000000000000000000000000000000000..bf9c4f4abf36e8075697e700b0b7ca7dcb2b7173
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_routes_openai.py
@@ -0,0 +1,978 @@
+
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for OpenAI API endpoints (routes_openai.py).
+
+Tests the following endpoints:
+- GET / - Root endpoint
+- GET /health - Health check
+- GET /v1/models - List available models
+- POST /v1/chat/completions - Chat completions
+
+For Anthropic API tests, see test_routes_anthropic.py.
+"""
+
+import pytest
+from unittest.mock import AsyncMock, Mock, patch, MagicMock
+from datetime import datetime, timezone
+import json
+
+from fastapi import HTTPException
+from fastapi.testclient import TestClient
+
+from kiro.routes_openai import verify_api_key, router
+from kiro.config import PROXY_API_KEY, APP_VERSION
+
+
+# =============================================================================
+# Tests for verify_api_key function
+# =============================================================================
+
+class TestVerifyApiKey:
+ """Tests for the verify_api_key authentication function."""
+
+ @pytest.mark.asyncio
+ async def test_valid_bearer_token_returns_true(self):
+ """
+ What it does: Verifies that a valid Bearer token passes authentication.
+ Purpose: Ensure correct API keys are accepted.
+ """
+ print("Setup: Creating valid Bearer token...")
+ valid_header = f"Bearer {PROXY_API_KEY}"
+
+ print("Action: Calling verify_api_key...")
+ result = await verify_api_key(valid_header)
+
+ print(f"Comparing result: Expected True, Got {result}")
+ assert result is True
+
+ @pytest.mark.asyncio
+ async def test_invalid_api_key_raises_401(self):
+ """
+ What it does: Verifies that an invalid API key is rejected.
+ Purpose: Ensure unauthorized access is blocked.
+ """
+ print("Setup: Creating invalid Bearer token...")
+ invalid_header = "Bearer wrong_key_12345"
+
+ print("Action: Calling verify_api_key with invalid key...")
+ with pytest.raises(HTTPException) as exc_info:
+ await verify_api_key(invalid_header)
+
+ print(f"Checking: HTTPException with status 401...")
+ assert exc_info.value.status_code == 401
+ assert "Invalid or missing API Key" in exc_info.value.detail
+
+ @pytest.mark.asyncio
+ async def test_missing_api_key_raises_401(self):
+ """
+ What it does: Verifies that missing API key is rejected.
+ Purpose: Ensure requests without authentication are blocked.
+ """
+ print("Setup: No API key provided...")
+
+ print("Action: Calling verify_api_key with None...")
+ with pytest.raises(HTTPException) as exc_info:
+ await verify_api_key(None)
+
+ print(f"Checking: HTTPException with status 401...")
+ assert exc_info.value.status_code == 401
+
+ @pytest.mark.asyncio
+ async def test_empty_api_key_raises_401(self):
+ """
+ What it does: Verifies that empty string API key is rejected.
+ Purpose: Ensure empty credentials are blocked.
+ """
+ print("Setup: Empty API key...")
+
+ print("Action: Calling verify_api_key with empty string...")
+ with pytest.raises(HTTPException) as exc_info:
+ await verify_api_key("")
+
+ print(f"Checking: HTTPException with status 401...")
+ assert exc_info.value.status_code == 401
+
+ @pytest.mark.asyncio
+ async def test_key_without_bearer_prefix_raises_401(self):
+ """
+ What it does: Verifies that API key without Bearer prefix is rejected.
+ Purpose: Ensure proper Authorization header format is required.
+ """
+ print("Setup: API key without Bearer prefix...")
+ wrong_format = PROXY_API_KEY # Without "Bearer "
+
+ print("Action: Calling verify_api_key...")
+ with pytest.raises(HTTPException) as exc_info:
+ await verify_api_key(wrong_format)
+
+ print(f"Checking: HTTPException with status 401...")
+ assert exc_info.value.status_code == 401
+
+ @pytest.mark.asyncio
+ async def test_bearer_with_extra_spaces_raises_401(self):
+ """
+ What it does: Verifies that Bearer token with extra spaces is rejected.
+ Purpose: Ensure strict format validation.
+ """
+ print("Setup: Bearer token with extra spaces...")
+ malformed = f"Bearer {PROXY_API_KEY}" # Double space
+
+ print("Action: Calling verify_api_key...")
+ with pytest.raises(HTTPException) as exc_info:
+ await verify_api_key(malformed)
+
+ print(f"Checking: HTTPException with status 401...")
+ assert exc_info.value.status_code == 401
+
+ @pytest.mark.asyncio
+ async def test_lowercase_bearer_raises_401(self):
+ """
+ What it does: Verifies that lowercase 'bearer' is rejected.
+ Purpose: Ensure case-sensitive Bearer prefix.
+ """
+ print("Setup: Lowercase bearer prefix...")
+ lowercase = f"bearer {PROXY_API_KEY}"
+
+ print("Action: Calling verify_api_key...")
+ with pytest.raises(HTTPException) as exc_info:
+ await verify_api_key(lowercase)
+
+ print(f"Checking: HTTPException with status 401...")
+ assert exc_info.value.status_code == 401
+
+
+# =============================================================================
+# Tests for root endpoint (/)
+# =============================================================================
+
+class TestRootEndpoint:
+ """Tests for the GET / endpoint."""
+
+ def test_root_returns_status_ok(self, test_client):
+ """
+ What it does: Verifies root endpoint returns ok status.
+ Purpose: Ensure basic health check works.
+ """
+ print("Action: GET /...")
+ response = test_client.get("/")
+
+ print(f"Result: {response.json()}")
+ assert response.status_code == 200
+ assert response.json()["status"] == "ok"
+
+ def test_root_returns_gateway_message(self, test_client):
+ """
+ What it does: Verifies root endpoint returns gateway message.
+ Purpose: Ensure service identification is present.
+ """
+ print("Action: GET /...")
+ response = test_client.get("/")
+
+ print(f"Result: {response.json()}")
+ assert response.status_code == 200
+ assert "Kiro Gateway" in response.json()["message"]
+
+ def test_root_returns_version(self, test_client):
+ """
+ What it does: Verifies root endpoint returns application version.
+ Purpose: Ensure version information is available.
+ """
+ print("Action: GET /...")
+ response = test_client.get("/")
+
+ print(f"Result: {response.json()}")
+ assert response.status_code == 200
+ assert "version" in response.json()
+ assert response.json()["version"] == APP_VERSION
+
+ def test_root_does_not_require_auth(self, test_client):
+ """
+ What it does: Verifies root endpoint is accessible without authentication.
+ Purpose: Ensure public health check availability.
+ """
+ print("Action: GET / without auth headers...")
+ response = test_client.get("/")
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 200
+
+
+# =============================================================================
+# Tests for health endpoint (/health)
+# =============================================================================
+
+class TestHealthEndpoint:
+ """Tests for the GET /health endpoint."""
+
+ def test_health_returns_healthy_status(self, test_client):
+ """
+ What it does: Verifies health endpoint returns healthy status.
+ Purpose: Ensure health check indicates service is running.
+ """
+ print("Action: GET /health...")
+ response = test_client.get("/health")
+
+ print(f"Result: {response.json()}")
+ assert response.status_code == 200
+ assert response.json()["status"] == "healthy"
+
+ def test_health_returns_timestamp(self, test_client):
+ """
+ What it does: Verifies health endpoint returns timestamp.
+ Purpose: Ensure timestamp is present for monitoring.
+ """
+ print("Action: GET /health...")
+ response = test_client.get("/health")
+
+ print(f"Result: {response.json()}")
+ assert response.status_code == 200
+ assert "timestamp" in response.json()
+ # Verify timestamp is ISO format
+ timestamp = response.json()["timestamp"]
+ assert "T" in timestamp # ISO format contains T
+
+ def test_health_returns_version(self, test_client):
+ """
+ What it does: Verifies health endpoint returns version.
+ Purpose: Ensure version is available for monitoring.
+ """
+ print("Action: GET /health...")
+ response = test_client.get("/health")
+
+ print(f"Result: {response.json()}")
+ assert response.status_code == 200
+ assert response.json()["version"] == APP_VERSION
+
+ def test_health_does_not_require_auth(self, test_client):
+ """
+ What it does: Verifies health endpoint is accessible without authentication.
+ Purpose: Ensure health checks work for load balancers.
+ """
+ print("Action: GET /health without auth headers...")
+ response = test_client.get("/health")
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 200
+
+
+# =============================================================================
+# Tests for models endpoint (/v1/models)
+# =============================================================================
+
+class TestModelsEndpoint:
+ """Tests for the GET /v1/models endpoint."""
+
+ def test_models_requires_authentication(self, test_client):
+ """
+ What it does: Verifies models endpoint requires authentication.
+ Purpose: Ensure protected endpoints are secured.
+ """
+ print("Action: GET /v1/models without auth...")
+ response = test_client.get("/v1/models")
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 401
+
+ def test_models_rejects_invalid_key(self, test_client, invalid_proxy_api_key):
+ """
+ What it does: Verifies models endpoint rejects invalid API key.
+ Purpose: Ensure authentication is enforced.
+ """
+ print("Action: GET /v1/models with invalid key...")
+ response = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {invalid_proxy_api_key}"}
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 401
+
+ def test_models_returns_list_object(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies models endpoint returns list object type.
+ Purpose: Ensure OpenAI API compatibility.
+ """
+ print("Action: GET /v1/models with valid auth...")
+ response = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"}
+ )
+
+ print(f"Result: {response.json()}")
+ assert response.status_code == 200
+ assert response.json()["object"] == "list"
+
+ def test_models_returns_data_array(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies models endpoint returns data array.
+ Purpose: Ensure response structure matches OpenAI format.
+ """
+ print("Action: GET /v1/models with valid auth...")
+ response = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"}
+ )
+
+ print(f"Result: {response.json()}")
+ assert response.status_code == 200
+ assert "data" in response.json()
+ assert isinstance(response.json()["data"], list)
+
+ def test_models_contains_available_models(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies all configured models are returned.
+ Purpose: Ensure model list is complete.
+ """
+ print("Action: GET /v1/models with valid auth...")
+ response = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"}
+ )
+
+ print(f"Result: {response.json()}")
+ assert response.status_code == 200
+
+ model_ids = [m["id"] for m in response.json()["data"]]
+ print(f"Model IDs: {model_ids}")
+
+ # At minimum, hidden models should be present
+ # (even if Kiro API cache is empty)
+ assert len(model_ids) >= 1, "Expected at least one model (hidden models)"
+
+ def test_models_format_is_openai_compatible(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies model objects have OpenAI-compatible format.
+ Purpose: Ensure compatibility with OpenAI clients.
+ """
+ print("Action: GET /v1/models with valid auth...")
+ response = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"}
+ )
+
+ print(f"Result: {response.json()}")
+ assert response.status_code == 200
+
+ for model in response.json()["data"]:
+ print(f"Checking model format: {model}")
+ assert "id" in model, "Model missing 'id' field"
+ assert "object" in model, "Model missing 'object' field"
+ assert model["object"] == "model", "Model object type should be 'model'"
+ assert "owned_by" in model, "Model missing 'owned_by' field"
+
+ def test_models_owned_by_anthropic(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies models are owned by Anthropic.
+ Purpose: Ensure correct model attribution.
+ """
+ print("Action: GET /v1/models with valid auth...")
+ response = test_client.get(
+ "/v1/models",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"}
+ )
+
+ print(f"Result: {response.json()}")
+ assert response.status_code == 200
+
+ for model in response.json()["data"]:
+ assert model["owned_by"] == "anthropic"
+
+
+# =============================================================================
+# Tests for chat completions endpoint (/v1/chat/completions)
+# =============================================================================
+
+class TestChatCompletionsAuthentication:
+ """Tests for authentication on /v1/chat/completions endpoint."""
+
+ def test_chat_completions_requires_authentication(self, test_client):
+ """
+ What it does: Verifies chat completions requires authentication.
+ Purpose: Ensure protected endpoint is secured.
+ """
+ print("Action: POST /v1/chat/completions without auth...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 401
+
+ def test_chat_completions_rejects_invalid_key(self, test_client, invalid_proxy_api_key):
+ """
+ What it does: Verifies chat completions rejects invalid API key.
+ Purpose: Ensure authentication is enforced.
+ """
+ print("Action: POST /v1/chat/completions with invalid key...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {invalid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 401
+
+
+class TestChatCompletionsValidation:
+ """Tests for request validation on /v1/chat/completions endpoint."""
+
+ def test_validates_empty_messages_array(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies empty messages array is rejected.
+ Purpose: Ensure at least one message is required.
+ """
+ print("Action: POST /v1/chat/completions with empty messages...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": []
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 422
+
+ def test_validates_missing_model(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies missing model field is rejected.
+ Purpose: Ensure model is required.
+ """
+ print("Action: POST /v1/chat/completions without model...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "messages": [{"role": "user", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 422
+
+ def test_validates_missing_messages(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies missing messages field is rejected.
+ Purpose: Ensure messages are required.
+ """
+ print("Action: POST /v1/chat/completions without messages...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5"
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 422
+
+ def test_validates_invalid_json(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies invalid JSON is rejected.
+ Purpose: Ensure proper JSON parsing.
+ """
+ print("Action: POST /v1/chat/completions with invalid JSON...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={
+ "Authorization": f"Bearer {valid_proxy_api_key}",
+ "Content-Type": "application/json"
+ },
+ content=b"not valid json {{{}"
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code == 422
+
+ def test_validates_invalid_role(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies invalid message role passes Pydantic validation.
+ Purpose: Pydantic model accepts any string as role (validation happens later).
+ Note: The role validation is not strict at Pydantic level, so invalid roles
+ pass validation but may fail during processing.
+ """
+ print("Action: POST /v1/chat/completions with invalid role...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "invalid_role", "content": "Hello"}]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Pydantic model accepts any string as role, so validation passes (not 422)
+ # The request may fail later during processing (500) due to network blocking
+ assert response.status_code != 422
+
+ def test_accepts_valid_request_format(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies valid request format passes validation.
+ Purpose: Ensure Pydantic validation works correctly.
+ """
+ print("Action: POST /v1/chat/completions with valid format...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "stream": False
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Should pass validation (not 422)
+ # May fail on HTTP call due to network blocking, but that's expected
+ assert response.status_code != 422
+
+ def test_accepts_message_without_content(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies message without content is accepted.
+ Purpose: Ensure content is optional (for tool results).
+ """
+ print("Action: POST /v1/chat/completions with message without content...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user"}] # No content
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Should pass validation (content is optional)
+ assert response.status_code != 422 or "content" not in str(response.json())
+
+
+class TestChatCompletionsWithTools:
+ """Tests for tool calling on /v1/chat/completions endpoint."""
+
+ def test_accepts_valid_tool_definition(self, test_client, valid_proxy_api_key, sample_tool_definition):
+ """
+ What it does: Verifies valid tool definition is accepted.
+ Purpose: Ensure tool calling format is supported.
+ """
+ print("Action: POST /v1/chat/completions with tools...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "What's the weather?"}],
+ "tools": [sample_tool_definition]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ # Should pass validation
+ assert response.status_code != 422
+
+ def test_accepts_multiple_tools(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies multiple tools are accepted.
+ Purpose: Ensure multiple tool definitions work.
+ """
+ print("Action: POST /v1/chat/completions with multiple tools...")
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get weather",
+ "parameters": {"type": "object", "properties": {}}
+ }
+ },
+ {
+ "type": "function",
+ "function": {
+ "name": "get_time",
+ "description": "Get time",
+ "parameters": {"type": "object", "properties": {}}
+ }
+ }
+ ]
+
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "tools": tools
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+
+class TestChatCompletionsOptionalParams:
+ """Tests for optional parameters on /v1/chat/completions endpoint."""
+
+ def test_accepts_temperature_parameter(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies temperature parameter is accepted.
+ Purpose: Ensure temperature control works.
+ """
+ print("Action: POST /v1/chat/completions with temperature...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "temperature": 0.7
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_max_tokens_parameter(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies max_tokens parameter is accepted.
+ Purpose: Ensure output length control works.
+ """
+ print("Action: POST /v1/chat/completions with max_tokens...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "max_tokens": 100
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_stream_true(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies stream=true is accepted.
+ Purpose: Ensure streaming mode is supported.
+ """
+ print("Action: POST /v1/chat/completions with stream=true...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "stream": True
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_top_p_parameter(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies top_p parameter is accepted.
+ Purpose: Ensure nucleus sampling control works.
+ """
+ print("Action: POST /v1/chat/completions with top_p...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "top_p": 0.9
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+
+class TestChatCompletionsMessageTypes:
+ """Tests for different message types on /v1/chat/completions endpoint."""
+
+ def test_accepts_system_message(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies system message is accepted.
+ Purpose: Ensure system prompts work.
+ """
+ print("Action: POST /v1/chat/completions with system message...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {"role": "system", "content": "You are helpful."},
+ {"role": "user", "content": "Hello"}
+ ]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_assistant_message(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies assistant message is accepted.
+ Purpose: Ensure conversation history works.
+ """
+ print("Action: POST /v1/chat/completions with assistant message...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {"role": "user", "content": "Hello"},
+ {"role": "assistant", "content": "Hi there!"},
+ {"role": "user", "content": "How are you?"}
+ ]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+ def test_accepts_multipart_content(self, test_client, valid_proxy_api_key):
+ """
+ What it does: Verifies multipart content array is accepted.
+ Purpose: Ensure complex content format works.
+ """
+ print("Action: POST /v1/chat/completions with multipart content...")
+ response = test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "Hello"},
+ {"type": "text", "text": "World"}
+ ]
+ }
+ ]
+ }
+ )
+
+ print(f"Status: {response.status_code}")
+ assert response.status_code != 422
+
+
+# =============================================================================
+# Tests for router integration
+# =============================================================================
+
+class TestRouterIntegration:
+ """Tests for router configuration and integration."""
+
+ def test_router_has_root_endpoint(self):
+ """
+ What it does: Verifies root endpoint is registered.
+ Purpose: Ensure endpoint is available.
+ """
+ print("Checking: Router endpoints...")
+ routes = [route.path for route in router.routes]
+
+ print(f"Found routes: {routes}")
+ assert "/" in routes
+
+ def test_router_has_health_endpoint(self):
+ """
+ What it does: Verifies health endpoint is registered.
+ Purpose: Ensure endpoint is available.
+ """
+ print("Checking: Router endpoints...")
+ routes = [route.path for route in router.routes]
+
+ print(f"Found routes: {routes}")
+ assert "/health" in routes
+
+ def test_router_has_models_endpoint(self):
+ """
+ What it does: Verifies models endpoint is registered.
+ Purpose: Ensure endpoint is available.
+ """
+ print("Checking: Router endpoints...")
+ routes = [route.path for route in router.routes]
+
+ print(f"Found routes: {routes}")
+ assert "/v1/models" in routes
+
+ def test_router_has_chat_completions_endpoint(self):
+ """
+ What it does: Verifies chat completions endpoint is registered.
+ Purpose: Ensure endpoint is available.
+ """
+ print("Checking: Router endpoints...")
+ routes = [route.path for route in router.routes]
+
+ print(f"Found routes: {routes}")
+ assert "/v1/chat/completions" in routes
+
+ def test_root_endpoint_uses_get_method(self):
+ """
+ What it does: Verifies root endpoint uses GET method.
+ Purpose: Ensure correct HTTP method.
+ """
+ print("Checking: HTTP methods...")
+ for route in router.routes:
+ if route.path == "/":
+ print(f"Route / methods: {route.methods}")
+ assert "GET" in route.methods
+ return
+ pytest.fail("Root endpoint not found")
+
+ def test_health_endpoint_uses_get_method(self):
+ """
+ What it does: Verifies health endpoint uses GET method.
+ Purpose: Ensure correct HTTP method.
+ """
+ print("Checking: HTTP methods...")
+ for route in router.routes:
+ if route.path == "/health":
+ print(f"Route /health methods: {route.methods}")
+ assert "GET" in route.methods
+ return
+ pytest.fail("Health endpoint not found")
+
+ def test_models_endpoint_uses_get_method(self):
+ """
+ What it does: Verifies models endpoint uses GET method.
+ Purpose: Ensure correct HTTP method.
+ """
+ print("Checking: HTTP methods...")
+ for route in router.routes:
+ if route.path == "/v1/models":
+ print(f"Route /v1/models methods: {route.methods}")
+ assert "GET" in route.methods
+ return
+ pytest.fail("Models endpoint not found")
+
+ def test_chat_completions_endpoint_uses_post_method(self):
+ """
+ What it does: Verifies chat completions endpoint uses POST method.
+ Purpose: Ensure correct HTTP method.
+ """
+ print("Checking: HTTP methods...")
+ for route in router.routes:
+ if route.path == "/v1/chat/completions":
+ print(f"Route /v1/chat/completions methods: {route.methods}")
+ assert "POST" in route.methods
+ return
+ pytest.fail("Chat completions endpoint not found")
+
+
+# =============================================================================
+# Tests for HTTP client selection (issue #54)
+# =============================================================================
+
+class TestHTTPClientSelection:
+ """
+ Tests for HTTP client selection in routes (issue #54).
+
+ Verifies that streaming requests use per-request clients to avoid CLOSE_WAIT leak
+ when network interface changes (VPN disconnect/reconnect), while non-streaming
+ requests use shared client for connection pooling.
+ """
+
+ @patch('kiro.routes_openai.KiroHttpClient')
+ def test_streaming_uses_per_request_client(
+ self,
+ mock_kiro_http_client_class,
+ test_client,
+ valid_proxy_api_key
+ ):
+ """
+ What it does: Verifies streaming requests create per-request HTTP client.
+ Purpose: Prevent CLOSE_WAIT leak on VPN disconnect (issue #54).
+ """
+ print("\n--- Test: Streaming uses per-request client ---")
+
+ # Setup mock
+ mock_client_instance = AsyncMock()
+ mock_client_instance.request_with_retry = AsyncMock(
+ side_effect=Exception("Network blocked")
+ )
+ mock_client_instance.close = AsyncMock()
+ mock_kiro_http_client_class.return_value = mock_client_instance
+
+ print("Action: POST with stream=true...")
+ try:
+ test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "stream": True
+ }
+ )
+ except Exception:
+ pass
+
+ 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 None, \
+ "Streaming should use per-request client"
+ print("✅ Streaming correctly uses per-request client")
+
+ @patch('kiro.routes_openai.KiroHttpClient')
+ def test_non_streaming_uses_shared_client(
+ self,
+ mock_kiro_http_client_class,
+ test_client,
+ valid_proxy_api_key
+ ):
+ """
+ What it does: Verifies non-streaming requests use shared HTTP client.
+ Purpose: Ensure connection pooling for non-streaming requests.
+ """
+ print("\n--- Test: Non-streaming uses shared client ---")
+
+ # Setup mock
+ mock_client_instance = AsyncMock()
+ mock_client_instance.request_with_retry = AsyncMock(
+ side_effect=Exception("Network blocked")
+ )
+ mock_client_instance.close = AsyncMock()
+ mock_kiro_http_client_class.return_value = mock_client_instance
+
+ print("Action: POST with stream=false...")
+ try:
+ test_client.post(
+ "/v1/chat/completions",
+ headers={"Authorization": f"Bearer {valid_proxy_api_key}"},
+ json={
+ "model": "claude-sonnet-4-5",
+ "messages": [{"role": "user", "content": "Hello"}],
+ "stream": False
+ }
+ )
+ except Exception:
+ pass
+
+ print("Checking: KiroHttpClient(shared_client=app.state.http_client)...")
+ 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")
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_streaming_anthropic.py b/kiro-gateway/tests/unit/test_streaming_anthropic.py
new file mode 100644
index 0000000000000000000000000000000000000000..6932fad1b2115a0b7a0de4a3335a3c9a3590daeb
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_streaming_anthropic.py
@@ -0,0 +1,1446 @@
+
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for streaming_anthropic module.
+
+Tests for:
+- generate_message_id() function
+- format_sse_event() function
+- stream_kiro_to_anthropic() generator
+- collect_anthropic_response() function
+"""
+
+import pytest
+import json
+import uuid
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from kiro.streaming_anthropic import (
+ generate_message_id,
+ generate_thinking_signature,
+ format_sse_event,
+ stream_kiro_to_anthropic,
+ collect_anthropic_response,
+ stream_with_first_token_retry_anthropic,
+)
+from kiro.streaming_core import KiroEvent, StreamResult
+
+
+# ==================================================================================================
+# Fixtures
+# ==================================================================================================
+
+@pytest.fixture
+def mock_model_cache():
+ """Mock for ModelInfoCache."""
+ cache = MagicMock()
+ cache.get_max_input_tokens.return_value = 200000
+ return cache
+
+
+@pytest.fixture
+def mock_auth_manager():
+ """Mock for KiroAuthManager."""
+ manager = MagicMock()
+ return manager
+
+
+@pytest.fixture
+def mock_response():
+ """Mock for httpx.Response."""
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ return response
+
+
+# ==================================================================================================
+# Tests for generate_message_id()
+# ==================================================================================================
+
+class TestGenerateMessageId:
+ """Tests for generate_message_id() function."""
+
+ def test_generates_message_id_with_prefix(self):
+ """
+ What it does: Generates message ID with 'msg_' prefix.
+ Goal: Verify Anthropic message ID format.
+ """
+ print("Action: Generating message ID...")
+ message_id = generate_message_id()
+
+ print(f"Generated ID: {message_id}")
+ assert message_id.startswith("msg_")
+ print("✓ Message ID has correct prefix")
+
+ def test_generates_unique_ids(self):
+ """
+ What it does: Generates unique message IDs.
+ Goal: Verify IDs are unique.
+ """
+ print("Action: Generating multiple message IDs...")
+ ids = [generate_message_id() for _ in range(100)]
+
+ print(f"Generated {len(ids)} IDs")
+ unique_ids = set(ids)
+ print(f"Unique IDs: {len(unique_ids)}")
+
+ assert len(unique_ids) == 100
+ print("✓ All message IDs are unique")
+
+ def test_message_id_has_correct_length(self):
+ """
+ What it does: Verifies message ID length.
+ Goal: Ensure ID format matches Anthropic spec.
+ """
+ print("Action: Generating message ID...")
+ message_id = generate_message_id()
+
+ # Format: msg_ + 24 hex chars
+ print(f"Generated ID: {message_id}, length: {len(message_id)}")
+ assert len(message_id) == 4 + 24 # "msg_" + 24 chars
+ print("✓ Message ID has correct length")
+
+
+# ==================================================================================================
+# Tests for format_sse_event()
+# ==================================================================================================
+
+class TestFormatSseEvent:
+ """Tests for format_sse_event() function."""
+
+ def test_formats_message_start_event(self):
+ """
+ What it does: Formats message_start event.
+ Goal: Verify Anthropic SSE format.
+ """
+ print("Action: Formatting message_start event...")
+ data = {
+ "type": "message_start",
+ "message": {
+ "id": "msg_123",
+ "type": "message",
+ "role": "assistant"
+ }
+ }
+
+ result = format_sse_event("message_start", data)
+
+ print(f"Formatted event:\n{result}")
+ assert result.startswith("event: message_start\n")
+ assert "data: " in result
+ assert result.endswith("\n\n")
+ print("✓ Event formatted correctly")
+
+ def test_formats_content_block_delta_event(self):
+ """
+ What it does: Formats content_block_delta event.
+ Goal: Verify delta event format.
+ """
+ print("Action: Formatting content_block_delta event...")
+ data = {
+ "type": "content_block_delta",
+ "index": 0,
+ "delta": {
+ "type": "text_delta",
+ "text": "Hello"
+ }
+ }
+
+ result = format_sse_event("content_block_delta", data)
+
+ print(f"Formatted event:\n{result}")
+ assert "event: content_block_delta\n" in result
+ assert '"text": "Hello"' in result
+ print("✓ Delta event formatted correctly")
+
+ def test_formats_message_stop_event(self):
+ """
+ What it does: Formats message_stop event.
+ Goal: Verify stop event format.
+ """
+ print("Action: Formatting message_stop event...")
+ data = {"type": "message_stop"}
+
+ result = format_sse_event("message_stop", data)
+
+ print(f"Formatted event:\n{result}")
+ assert "event: message_stop\n" in result
+ print("✓ Stop event formatted correctly")
+
+ def test_handles_unicode_content(self):
+ """
+ What it does: Handles Unicode content in events.
+ Goal: Verify non-ASCII characters are preserved.
+ """
+ print("Action: Formatting event with Unicode...")
+ data = {
+ "type": "content_block_delta",
+ "delta": {"text": "Привет мир! 🌍"}
+ }
+
+ result = format_sse_event("content_block_delta", data)
+
+ print(f"Formatted event:\n{result}")
+ assert "Привет мир!" in result
+ assert "🌍" in result
+ print("✓ Unicode content preserved")
+
+ def test_json_data_is_valid(self):
+ """
+ What it does: Verifies JSON data is valid.
+ Goal: Ensure data can be parsed back.
+ """
+ print("Action: Formatting and parsing event...")
+ data = {
+ "type": "message_delta",
+ "delta": {"stop_reason": "end_turn"},
+ "usage": {"output_tokens": 100}
+ }
+
+ result = format_sse_event("message_delta", data)
+
+ # Extract JSON from result
+ lines = result.strip().split("\n")
+ data_line = [l for l in lines if l.startswith("data: ")][0]
+ json_str = data_line[6:] # Remove "data: " prefix
+
+ print(f"JSON string: {json_str}")
+ parsed = json.loads(json_str)
+
+ assert parsed["type"] == "message_delta"
+ assert parsed["delta"]["stop_reason"] == "end_turn"
+ print("✓ JSON data is valid and parseable")
+
+
+# ==================================================================================================
+# Tests for stream_kiro_to_anthropic()
+# ==================================================================================================
+
+class TestStreamKiroToAnthropic:
+ """Tests for stream_kiro_to_anthropic() generator."""
+
+ @pytest.mark.asyncio
+ async def test_yields_message_start_event(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields message_start event at beginning.
+ Goal: Verify Anthropic streaming protocol.
+ """
+ print("Setup: Mock empty stream...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ return
+ yield # Make it a generator
+
+ print("Action: Streaming to Anthropic format...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+
+ # First event should be message_start
+ assert len(events) > 0
+ assert "event: message_start" in events[0]
+ print("✓ message_start event yielded first")
+
+ @pytest.mark.asyncio
+ async def test_yields_content_block_start_on_first_content(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields content_block_start before first content.
+ Goal: Verify content block lifecycle.
+ """
+ print("Setup: Mock stream with content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Streaming to Anthropic format...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+
+ # Should have content_block_start
+ content_block_start_found = any("content_block_start" in e for e in events)
+ assert content_block_start_found
+ print("✓ content_block_start event yielded")
+
+ @pytest.mark.asyncio
+ async def test_yields_content_block_delta_for_content(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields content_block_delta for content events.
+ Goal: Verify content streaming.
+ """
+ print("Setup: Mock stream with content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ yield KiroEvent(type="content", content=" World")
+
+ print("Action: Streaming to Anthropic format...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+
+ # Should have content_block_delta events
+ delta_events = [e for e in events if "content_block_delta" in e]
+ print(f"Delta events: {len(delta_events)}")
+
+ assert len(delta_events) >= 2
+ assert "Hello" in delta_events[0]
+ assert "World" in delta_events[1]
+ print("✓ content_block_delta events yielded for content")
+
+ @pytest.mark.asyncio
+ async def test_yields_tool_use_block_for_tool_calls(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields tool_use block for tool calls.
+ Goal: Verify tool use streaming.
+ """
+ print("Setup: Mock stream with tool call...")
+
+ tool_use_data = {
+ "id": "toolu_123",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"city": "Moscow"}'
+ }
+ }
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Let me check")
+ yield KiroEvent(type="tool_use", tool_use=tool_use_data)
+
+ print("Action: Streaming to Anthropic format...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+
+ # Should have tool_use content block
+ tool_use_events = [e for e in events if "tool_use" in e and "content_block_start" in e]
+ print(f"Tool use events: {len(tool_use_events)}")
+
+ assert len(tool_use_events) >= 1
+ assert "get_weather" in tool_use_events[0]
+ print("✓ tool_use block yielded for tool calls")
+
+ @pytest.mark.asyncio
+ async def test_yields_message_delta_with_stop_reason(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields message_delta with stop_reason.
+ Goal: Verify message completion.
+ """
+ print("Setup: Mock stream with content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Streaming to Anthropic format...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+
+ # Should have message_delta with stop_reason
+ message_delta_events = [e for e in events if "message_delta" in e]
+ assert len(message_delta_events) >= 1
+ assert "end_turn" in message_delta_events[0]
+ print("✓ message_delta with stop_reason yielded")
+
+ @pytest.mark.asyncio
+ async def test_yields_message_stop_at_end(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields message_stop at end.
+ Goal: Verify stream termination.
+ """
+ print("Setup: Mock stream with content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Streaming to Anthropic format...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+
+ # Last event should be message_stop
+ assert "message_stop" in events[-1]
+ print("✓ message_stop yielded at end")
+
+ @pytest.mark.asyncio
+ async def test_stop_reason_is_tool_use_when_tools_present(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Sets stop_reason to tool_use when tools are present.
+ Goal: Verify correct stop reason for tool calls.
+ """
+ print("Setup: Mock stream with tool call...")
+
+ tool_use_data = {
+ "id": "toolu_123",
+ "function": {"name": "func1", "arguments": "{}"}
+ }
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="tool_use", tool_use=tool_use_data)
+
+ print("Action: Streaming to Anthropic format...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+
+ # message_delta should have stop_reason: tool_use
+ message_delta_events = [e for e in events if "message_delta" in e]
+ assert len(message_delta_events) >= 1
+ assert "tool_use" in message_delta_events[0]
+ print("✓ stop_reason is tool_use when tools present")
+
+ @pytest.mark.asyncio
+ async def test_handles_bracket_tool_calls(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Handles bracket-style tool calls in content.
+ Goal: Verify bracket tool call detection.
+ """
+ print("Setup: Mock stream with bracket tool calls...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="[tool_call: func1]")
+
+ bracket_tool_calls = [
+ {"id": "call_1", "function": {"name": "func1", "arguments": "{}"}}
+ ]
+
+ print("Action: Streaming to Anthropic format...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=bracket_tool_calls):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+
+ # Should have tool_use block from bracket tool calls
+ tool_use_events = [e for e in events if "tool_use" in e and "content_block_start" in e]
+ assert len(tool_use_events) >= 1
+ print("✓ Bracket tool calls handled correctly")
+
+ @pytest.mark.asyncio
+ async def test_closes_response_on_completion(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Closes response on completion.
+ Goal: Verify resource cleanup.
+ """
+ print("Setup: Mock stream...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Streaming to Anthropic format...")
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ pass
+
+ print("Check: response.aclose() should be called...")
+ mock_response.aclose.assert_called()
+ print("✓ Response closed on completion")
+
+ @pytest.mark.asyncio
+ async def test_closes_response_on_error(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Closes response on error.
+ Goal: Verify resource cleanup on error.
+ """
+ print("Setup: Mock stream that raises error...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ raise RuntimeError("Test error")
+
+ print("Action: Streaming to Anthropic format with error...")
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ try:
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ pass
+ except RuntimeError:
+ pass
+
+ print("Check: response.aclose() should be called...")
+ mock_response.aclose.assert_called()
+ print("✓ Response closed on error")
+
+
+# ==================================================================================================
+# Tests for collect_anthropic_response()
+# ==================================================================================================
+
+class TestCollectAnthropicResponse:
+ """Tests for collect_anthropic_response() function."""
+
+ @pytest.mark.asyncio
+ async def test_collects_text_content(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Collects text content into response.
+ Goal: Verify content collection.
+ """
+ print("Setup: Mock stream result with content...")
+
+ mock_result = StreamResult(
+ content="Hello, world!",
+ thinking_content="",
+ tool_calls=[],
+ usage=None,
+ context_usage_percentage=None
+ )
+
+ print("Action: Collecting Anthropic response...")
+
+ with patch('kiro.streaming_anthropic.collect_stream_to_result', return_value=mock_result):
+ result = await collect_anthropic_response(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Result: {result}")
+
+ assert result["type"] == "message"
+ assert result["role"] == "assistant"
+ assert len(result["content"]) == 1
+ assert result["content"][0]["type"] == "text"
+ assert result["content"][0]["text"] == "Hello, world!"
+ print("✓ Text content collected correctly")
+
+ @pytest.mark.asyncio
+ async def test_collects_tool_use_content(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Collects tool use into response.
+ Goal: Verify tool use collection.
+ """
+ print("Setup: Mock stream result with tool calls...")
+
+ mock_result = StreamResult(
+ content="Let me check",
+ thinking_content="",
+ tool_calls=[
+ {
+ "id": "toolu_123",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"city": "Moscow"}'
+ }
+ }
+ ],
+ usage=None,
+ context_usage_percentage=None
+ )
+
+ print("Action: Collecting Anthropic response...")
+
+ with patch('kiro.streaming_anthropic.collect_stream_to_result', return_value=mock_result):
+ result = await collect_anthropic_response(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Result: {result}")
+
+ # Should have text and tool_use blocks
+ assert len(result["content"]) == 2
+
+ text_block = result["content"][0]
+ assert text_block["type"] == "text"
+
+ tool_block = result["content"][1]
+ assert tool_block["type"] == "tool_use"
+ assert tool_block["name"] == "get_weather"
+ assert tool_block["input"] == {"city": "Moscow"}
+ print("✓ Tool use content collected correctly")
+
+ @pytest.mark.asyncio
+ async def test_sets_stop_reason_end_turn(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Sets stop_reason to end_turn for normal completion.
+ Goal: Verify stop reason.
+ """
+ print("Setup: Mock stream result without tool calls...")
+
+ mock_result = StreamResult(
+ content="Hello",
+ thinking_content="",
+ tool_calls=[],
+ usage=None,
+ context_usage_percentage=None
+ )
+
+ print("Action: Collecting Anthropic response...")
+
+ with patch('kiro.streaming_anthropic.collect_stream_to_result', return_value=mock_result):
+ result = await collect_anthropic_response(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ )
+
+ print(f"stop_reason: {result['stop_reason']}")
+ assert result["stop_reason"] == "end_turn"
+ print("✓ stop_reason is end_turn")
+
+ @pytest.mark.asyncio
+ async def test_sets_stop_reason_tool_use(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Sets stop_reason to tool_use when tools present.
+ Goal: Verify stop reason for tool calls.
+ """
+ print("Setup: Mock stream result with tool calls...")
+
+ mock_result = StreamResult(
+ content="",
+ thinking_content="",
+ tool_calls=[{"id": "call_1", "function": {"name": "func1", "arguments": "{}"}}],
+ usage=None,
+ context_usage_percentage=None
+ )
+
+ print("Action: Collecting Anthropic response...")
+
+ with patch('kiro.streaming_anthropic.collect_stream_to_result', return_value=mock_result):
+ result = await collect_anthropic_response(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ )
+
+ print(f"stop_reason: {result['stop_reason']}")
+ assert result["stop_reason"] == "tool_use"
+ print("✓ stop_reason is tool_use")
+
+ @pytest.mark.asyncio
+ async def test_includes_usage_info(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Includes usage information in response.
+ Goal: Verify usage is included.
+ """
+ print("Setup: Mock stream result...")
+
+ mock_result = StreamResult(
+ content="Hello, world!",
+ thinking_content="",
+ tool_calls=[],
+ usage=None,
+ context_usage_percentage=None
+ )
+
+ print("Action: Collecting Anthropic response...")
+
+ with patch('kiro.streaming_anthropic.collect_stream_to_result', return_value=mock_result):
+ with patch('kiro.streaming_anthropic.count_message_tokens', return_value=10):
+ with patch('kiro.streaming_anthropic.count_tokens', return_value=5):
+ result = await collect_anthropic_response(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager,
+ request_messages=[{"role": "user", "content": "Hi"}]
+ )
+
+ print(f"Usage: {result['usage']}")
+ assert "input_tokens" in result["usage"]
+ assert "output_tokens" in result["usage"]
+ print("✓ Usage info included")
+
+ @pytest.mark.asyncio
+ async def test_generates_message_id(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Generates message ID for response.
+ Goal: Verify message ID is present.
+ """
+ print("Setup: Mock stream result...")
+
+ mock_result = StreamResult(
+ content="Hello",
+ thinking_content="",
+ tool_calls=[],
+ usage=None,
+ context_usage_percentage=None
+ )
+
+ print("Action: Collecting Anthropic response...")
+
+ with patch('kiro.streaming_anthropic.collect_stream_to_result', return_value=mock_result):
+ result = await collect_anthropic_response(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Message ID: {result['id']}")
+ assert result["id"].startswith("msg_")
+ print("✓ Message ID generated")
+
+ @pytest.mark.asyncio
+ async def test_includes_model_name(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Includes model name in response.
+ Goal: Verify model is included.
+ """
+ print("Setup: Mock stream result...")
+
+ mock_result = StreamResult(
+ content="Hello",
+ thinking_content="",
+ tool_calls=[],
+ usage=None,
+ context_usage_percentage=None
+ )
+
+ print("Action: Collecting Anthropic response...")
+
+ with patch('kiro.streaming_anthropic.collect_stream_to_result', return_value=mock_result):
+ result = await collect_anthropic_response(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Model: {result['model']}")
+ assert result["model"] == "claude-sonnet-4"
+ print("✓ Model name included")
+
+ @pytest.mark.asyncio
+ async def test_parses_tool_arguments_from_string(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Parses tool arguments from JSON string.
+ Goal: Verify arguments are parsed to dict.
+ """
+ print("Setup: Mock stream result with string arguments...")
+
+ mock_result = StreamResult(
+ content="",
+ thinking_content="",
+ tool_calls=[
+ {
+ "id": "call_1",
+ "function": {
+ "name": "func1",
+ "arguments": '{"key": "value"}' # String, not dict
+ }
+ }
+ ],
+ usage=None,
+ context_usage_percentage=None
+ )
+
+ print("Action: Collecting Anthropic response...")
+
+ with patch('kiro.streaming_anthropic.collect_stream_to_result', return_value=mock_result):
+ result = await collect_anthropic_response(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Result: {result}")
+
+ # Tool input should be parsed to dict
+ tool_block = result["content"][0] # Only tool_use since content is empty
+ assert tool_block["type"] == "tool_use"
+ assert tool_block["input"] == {"key": "value"}
+ assert isinstance(tool_block["input"], dict)
+ print("✓ Tool arguments parsed from string to dict")
+
+ @pytest.mark.asyncio
+ async def test_handles_invalid_json_arguments(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Handles invalid JSON in tool arguments.
+ Goal: Verify graceful handling of invalid JSON.
+ """
+ print("Setup: Mock stream result with invalid JSON arguments...")
+
+ mock_result = StreamResult(
+ content="",
+ thinking_content="",
+ tool_calls=[
+ {
+ "id": "call_1",
+ "function": {
+ "name": "func1",
+ "arguments": "not valid json" # Invalid JSON
+ }
+ }
+ ],
+ usage=None,
+ context_usage_percentage=None
+ )
+
+ print("Action: Collecting Anthropic response...")
+
+ with patch('kiro.streaming_anthropic.collect_stream_to_result', return_value=mock_result):
+ result = await collect_anthropic_response(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Result: {result}")
+
+ # Should handle gracefully with empty dict
+ tool_block = result["content"][0]
+ assert tool_block["type"] == "tool_use"
+ assert tool_block["input"] == {}
+ print("✓ Invalid JSON arguments handled gracefully")
+
+ @pytest.mark.asyncio
+ async def test_handles_empty_content(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Handles empty content in response.
+ Goal: Verify empty content is handled.
+ """
+ print("Setup: Mock stream result with empty content...")
+
+ mock_result = StreamResult(
+ content="",
+ thinking_content="",
+ tool_calls=[],
+ usage=None,
+ context_usage_percentage=None
+ )
+
+ print("Action: Collecting Anthropic response...")
+
+ with patch('kiro.streaming_anthropic.collect_stream_to_result', return_value=mock_result):
+ result = await collect_anthropic_response(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Result: {result}")
+
+ # Content should be empty list
+ assert result["content"] == []
+ print("✓ Empty content handled correctly")
+
+
+# ==================================================================================================
+# Tests for error handling
+# ==================================================================================================
+
+class TestStreamingAnthropicErrorHandling:
+ """Tests for error handling in streaming_anthropic."""
+
+ @pytest.mark.asyncio
+ async def test_propagates_first_token_timeout_error(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Propagates FirstTokenTimeoutError.
+ Goal: Verify timeout error is not caught internally.
+ """
+ from kiro.streaming_core import FirstTokenTimeoutError
+
+ print("Setup: Mock stream that raises timeout...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ raise FirstTokenTimeoutError("Timeout!")
+ yield # Make it a generator
+
+ print("Action: Streaming to Anthropic format with timeout...")
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with pytest.raises(FirstTokenTimeoutError):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ pass
+
+ print("✓ FirstTokenTimeoutError propagated correctly")
+
+ @pytest.mark.asyncio
+ async def test_propagates_generator_exit(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Propagates GeneratorExit.
+ Goal: Verify client disconnect is handled.
+ """
+ print("Setup: Mock stream that raises GeneratorExit...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ raise GeneratorExit()
+
+ print("Action: Streaming to Anthropic format with GeneratorExit...")
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ with pytest.raises(GeneratorExit):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ pass
+
+ print("✓ GeneratorExit propagated correctly")
+
+ @pytest.mark.asyncio
+ async def test_yields_error_event_on_exception(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields error event on exception.
+ Goal: Verify error event is sent to client.
+ """
+ print("Setup: Mock stream that raises RuntimeError...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ raise RuntimeError("Test error")
+
+ print("Action: Streaming to Anthropic format with error...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ try:
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ events.append(event)
+ except RuntimeError:
+ pass
+
+ print(f"Received {len(events)} events")
+
+ # Should have error event
+ error_events = [e for e in events if "event: error" in e]
+ assert len(error_events) >= 1
+ assert "Test error" in error_events[0]
+ print("✓ Error event yielded on exception")
+
+ @pytest.mark.asyncio
+ async def test_closes_response_in_finally(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Closes response in finally block.
+ Goal: Verify resource cleanup always happens.
+ """
+ print("Setup: Mock stream that raises error...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ raise ValueError("Test error")
+ yield # Make it a generator
+
+ print("Action: Streaming to Anthropic format with error...")
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ try:
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ pass
+ except ValueError:
+ pass
+
+ print("Check: response.aclose() should be called...")
+ mock_response.aclose.assert_called()
+ print("✓ Response closed in finally block")
+
+
+# ==================================================================================================
+# Tests for thinking content handling
+# ==================================================================================================
+
+class TestStreamingAnthropicThinkingContent:
+ """Tests for thinking content handling in Anthropic streaming."""
+
+ @pytest.mark.asyncio
+ async def test_includes_thinking_as_text_when_configured(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Includes thinking content as text when configured.
+ Goal: Verify thinking content handling.
+ """
+ print("Setup: Mock stream with thinking content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="thinking", thinking_content="Let me think...")
+ yield KiroEvent(type="content", content="Here is my answer")
+
+ print("Action: Streaming to Anthropic format with thinking...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ with patch('kiro.streaming_anthropic.FAKE_REASONING_HANDLING', 'include_as_text'):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+
+ # Should have thinking content as text delta
+ delta_events = [e for e in events if "content_block_delta" in e]
+ thinking_found = any("Let me think" in e for e in delta_events)
+ assert thinking_found
+ print("✓ Thinking content included as text")
+
+ @pytest.mark.asyncio
+ async def test_strips_thinking_when_configured(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Strips thinking content when configured.
+ Goal: Verify thinking content is stripped.
+ """
+ print("Setup: Mock stream with thinking content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="thinking", thinking_content="Let me think...")
+ yield KiroEvent(type="content", content="Here is my answer")
+
+ print("Action: Streaming to Anthropic format with strip mode...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ with patch('kiro.streaming_anthropic.FAKE_REASONING_HANDLING', 'strip'):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+
+ # Should NOT have thinking content
+ delta_events = [e for e in events if "content_block_delta" in e]
+ thinking_found = any("Let me think" in e for e in delta_events)
+ assert not thinking_found
+ print("✓ Thinking content stripped")
+
+
+# ==================================================================================================
+# Tests for context usage calculation
+# ==================================================================================================
+
+class TestStreamingAnthropicContextUsage:
+ """Tests for context usage calculation in Anthropic streaming."""
+
+ @pytest.mark.asyncio
+ async def test_calculates_tokens_from_context_usage(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Calculates tokens from context usage percentage.
+ Goal: Verify token calculation.
+ """
+ print("Setup: Mock stream with context usage...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ yield KiroEvent(type="context_usage", context_usage_percentage=5.0)
+
+ print("Action: Streaming to Anthropic format...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager
+ ):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+
+ # message_delta should have usage with output_tokens
+ message_delta_events = [e for e in events if "message_delta" in e]
+ assert len(message_delta_events) >= 1
+ assert "output_tokens" in message_delta_events[0]
+ print("✓ Tokens calculated from context usage")
+
+ @pytest.mark.asyncio
+ async def test_uses_request_messages_for_input_tokens(self, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Uses request messages for input token count.
+ Goal: Verify input tokens are counted from request.
+ """
+ print("Setup: Mock stream...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ request_messages = [
+ {"role": "user", "content": "Hi there!"}
+ ]
+
+ print("Action: Streaming to Anthropic format with request messages...")
+ events = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ with patch('kiro.streaming_anthropic.count_message_tokens', return_value=10) as mock_count:
+ async for event in stream_kiro_to_anthropic(
+ mock_response, "claude-sonnet-4", mock_model_cache, mock_auth_manager,
+ request_messages=request_messages
+ ):
+ events.append(event)
+
+ # Verify count_message_tokens was called
+ mock_count.assert_called_once_with(request_messages, apply_claude_correction=False)
+
+ print("✓ Request messages used for input token count")
+
+
+# ==================================================================================================
+# Tests for generate_thinking_signature()
+# ==================================================================================================
+
+class TestGenerateThinkingSignature:
+ """
+ Tests for generate_thinking_signature() function.
+
+ This function generates placeholder signatures for thinking content blocks.
+ In real Anthropic API, this is a cryptographic signature for verification.
+ Since we use fake reasoning via tag injection, we generate a placeholder.
+ """
+
+ def test_generates_signature_with_prefix(self):
+ """
+ What it does: Generates signature with 'sig_' prefix.
+ Goal: Verify signature format matches expected pattern.
+ """
+ print("Action: Generating thinking signature...")
+ signature = generate_thinking_signature()
+
+ print(f"Generated signature: {signature}")
+ assert signature.startswith("sig_")
+ print("✓ Signature has correct prefix")
+
+ def test_generates_unique_signatures(self):
+ """
+ What it does: Generates unique signatures.
+ Goal: Verify signatures are unique across multiple calls.
+ """
+ print("Action: Generating multiple signatures...")
+ signatures = [generate_thinking_signature() for _ in range(100)]
+
+ print(f"Generated {len(signatures)} signatures")
+ unique_signatures = set(signatures)
+ print(f"Unique signatures: {len(unique_signatures)}")
+
+ assert len(unique_signatures) == 100
+ print("✓ All signatures are unique")
+
+ def test_signature_has_correct_length(self):
+ """
+ What it does: Verifies signature length.
+ Goal: Ensure signature format is consistent.
+ """
+ print("Action: Generating signature...")
+ signature = generate_thinking_signature()
+
+ # Format: sig_ + 32 hex chars
+ print(f"Generated signature: {signature}, length: {len(signature)}")
+ assert len(signature) == 4 + 32 # "sig_" + 32 chars
+ print("✓ Signature has correct length")
+
+ def test_signature_contains_only_valid_characters(self):
+ """
+ What it does: Verifies signature contains only valid hex characters.
+ Goal: Ensure signature is properly formatted.
+ """
+ print("Action: Generating signature...")
+ signature = generate_thinking_signature()
+
+ print(f"Generated signature: {signature}")
+ # Remove prefix and check remaining chars are hex
+ hex_part = signature[4:] # Remove "sig_"
+ assert all(c in '0123456789abcdef' for c in hex_part)
+ print("✓ Signature contains only valid hex characters")
+
+
+# ==================================================================================================
+# Tests for stream_with_first_token_retry_anthropic()
+# ==================================================================================================
+
+class TestStreamWithFirstTokenRetryAnthropic:
+ """
+ Tests for stream_with_first_token_retry_anthropic() function.
+
+ This function wraps stream_kiro_to_anthropic with automatic retry
+ on first token timeout. It uses the generic stream_with_first_token_retry
+ from streaming_core.py with Anthropic-specific error formatting.
+ """
+
+ @pytest.mark.asyncio
+ async def test_yields_chunks_on_success(self, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields chunks on successful streaming.
+ Goal: Verify normal operation without retries.
+ """
+ print("Setup: Mock successful request...")
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.aclose = AsyncMock()
+
+ async def mock_make_request():
+ return mock_response
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Streaming with retry wrapper...")
+ chunks = []
+
+ with patch('kiro.streaming_anthropic.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_anthropic.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_with_first_token_retry_anthropic(
+ make_request=mock_make_request,
+ model="claude-sonnet-4",
+ model_cache=mock_model_cache,
+ auth_manager=mock_auth_manager,
+ max_retries=3,
+ first_token_timeout=30
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+ assert len(chunks) > 0
+ assert any("message_start" in c for c in chunks)
+ print("✓ Chunks yielded on success")
+
+ @pytest.mark.asyncio
+ async def test_retries_on_first_token_timeout(self, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Retries on first token timeout.
+ Goal: Verify retry logic is triggered.
+ """
+ from kiro.streaming_core import FirstTokenTimeoutError
+
+ print("Setup: Mock request that times out then succeeds...")
+
+ call_count = 0
+
+ async def mock_make_request():
+ nonlocal call_count
+ call_count += 1
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ return response
+
+ async def mock_stream_kiro_to_anthropic(*args, **kwargs):
+ nonlocal call_count
+ if call_count == 1:
+ raise FirstTokenTimeoutError("Timeout on first attempt")
+ yield "event: message_start\ndata: {}\n\n"
+ yield "event: message_stop\ndata: {}\n\n"
+
+ print("Action: Streaming with retry on timeout...")
+ chunks = []
+
+ with patch('kiro.streaming_anthropic.stream_kiro_to_anthropic', mock_stream_kiro_to_anthropic):
+ async for chunk in stream_with_first_token_retry_anthropic(
+ make_request=mock_make_request,
+ model="claude-sonnet-4",
+ model_cache=mock_model_cache,
+ auth_manager=mock_auth_manager,
+ max_retries=3,
+ first_token_timeout=30
+ ):
+ chunks.append(chunk)
+
+ print(f"Call count: {call_count}")
+ print(f"Received {len(chunks)} chunks")
+
+ assert call_count == 2 # First timeout, second success
+ assert len(chunks) > 0
+ print("✓ Retry on timeout works correctly")
+
+ @pytest.mark.asyncio
+ async def test_raises_anthropic_error_after_all_retries(self, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Raises Anthropic-formatted error after all retries exhausted.
+ Goal: Verify error format matches Anthropic API.
+ """
+ from kiro.streaming_core import FirstTokenTimeoutError
+
+ print("Setup: Mock request that always times out...")
+
+ async def mock_make_request():
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ return response
+
+ async def mock_stream_kiro_to_anthropic(*args, **kwargs):
+ raise FirstTokenTimeoutError("Timeout!")
+ yield # Make it a generator
+
+ print("Action: Streaming with all retries failing...")
+
+ with patch('kiro.streaming_anthropic.stream_kiro_to_anthropic', mock_stream_kiro_to_anthropic):
+ with pytest.raises(Exception) as exc_info:
+ async for chunk in stream_with_first_token_retry_anthropic(
+ make_request=mock_make_request,
+ model="claude-sonnet-4",
+ model_cache=mock_model_cache,
+ auth_manager=mock_auth_manager,
+ max_retries=2,
+ first_token_timeout=30
+ ):
+ pass
+
+ print(f"Exception: {exc_info.value}")
+
+ # Error should be in Anthropic format (JSON)
+ error_json = json.loads(str(exc_info.value))
+ assert error_json["type"] == "error"
+ assert error_json["error"]["type"] == "timeout_error"
+ assert "30" in error_json["error"]["message"]
+ print("✓ Anthropic-formatted error raised after all retries")
+
+ @pytest.mark.asyncio
+ async def test_raises_anthropic_error_on_http_error(self, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Raises Anthropic-formatted error on HTTP error.
+ Goal: Verify HTTP errors are formatted correctly.
+ """
+ print("Setup: Mock request that returns HTTP error...")
+
+ async def mock_make_request():
+ response = AsyncMock()
+ response.status_code = 500
+ response.aread = AsyncMock(return_value=b"Internal Server Error")
+ response.aclose = AsyncMock()
+ return response
+
+ print("Action: Streaming with HTTP error...")
+
+ with pytest.raises(Exception) as exc_info:
+ async for chunk in stream_with_first_token_retry_anthropic(
+ make_request=mock_make_request,
+ model="claude-sonnet-4",
+ model_cache=mock_model_cache,
+ auth_manager=mock_auth_manager,
+ max_retries=2,
+ first_token_timeout=30
+ ):
+ pass
+
+ print(f"Exception: {exc_info.value}")
+
+ # Error should be in Anthropic format (JSON)
+ error_json = json.loads(str(exc_info.value))
+ assert error_json["type"] == "error"
+ assert error_json["error"]["type"] == "api_error"
+ assert "Upstream API error" in error_json["error"]["message"]
+ print("✓ Anthropic-formatted error raised on HTTP error")
+
+ @pytest.mark.asyncio
+ async def test_passes_request_messages_to_stream(self, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Passes request_messages to underlying stream function.
+ Goal: Verify token counting parameters are forwarded.
+ """
+ print("Setup: Mock request with messages...")
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.aclose = AsyncMock()
+
+ async def mock_make_request():
+ return mock_response
+
+ captured_kwargs = {}
+
+ async def mock_stream_kiro_to_anthropic(*args, **kwargs):
+ captured_kwargs.update(kwargs)
+ yield "event: message_start\ndata: {}\n\n"
+ yield "event: message_stop\ndata: {}\n\n"
+
+ request_messages = [{"role": "user", "content": "Hello"}]
+
+ print("Action: Streaming with request_messages...")
+
+ with patch('kiro.streaming_anthropic.stream_kiro_to_anthropic', mock_stream_kiro_to_anthropic):
+ async for chunk in stream_with_first_token_retry_anthropic(
+ make_request=mock_make_request,
+ model="claude-sonnet-4",
+ model_cache=mock_model_cache,
+ auth_manager=mock_auth_manager,
+ request_messages=request_messages
+ ):
+ pass
+
+ print(f"Captured kwargs: {captured_kwargs}")
+ assert captured_kwargs.get("request_messages") == request_messages
+ print("✓ request_messages passed to stream function")
+
+ @pytest.mark.asyncio
+ async def test_uses_configured_max_retries(self, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Uses configured max_retries value.
+ Goal: Verify max_retries parameter is respected.
+ """
+ from kiro.streaming_core import FirstTokenTimeoutError
+
+ print("Setup: Mock request that always times out...")
+
+ call_count = 0
+
+ async def mock_make_request():
+ nonlocal call_count
+ call_count += 1
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ return response
+
+ async def mock_stream_kiro_to_anthropic(*args, **kwargs):
+ raise FirstTokenTimeoutError("Timeout!")
+ yield # Make it a generator
+
+ print("Action: Streaming with max_retries=5...")
+
+ with patch('kiro.streaming_anthropic.stream_kiro_to_anthropic', mock_stream_kiro_to_anthropic):
+ try:
+ async for chunk in stream_with_first_token_retry_anthropic(
+ make_request=mock_make_request,
+ model="claude-sonnet-4",
+ model_cache=mock_model_cache,
+ auth_manager=mock_auth_manager,
+ max_retries=5,
+ first_token_timeout=30
+ ):
+ pass
+ except Exception:
+ pass
+
+ print(f"Call count: {call_count}")
+ assert call_count == 5 # Should try exactly 5 times
+ print("✓ max_retries parameter respected")
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_streaming_core.py b/kiro-gateway/tests/unit/test_streaming_core.py
new file mode 100644
index 0000000000000000000000000000000000000000..f1609950fdf5bd645f2d230e7ea47f7b9b8f2d5b
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_streaming_core.py
@@ -0,0 +1,1684 @@
+
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for streaming_core module.
+
+Tests for:
+- KiroEvent dataclass
+- StreamResult dataclass
+- FirstTokenTimeoutError exception
+- parse_kiro_stream() function
+- collect_stream_to_result() function
+- calculate_tokens_from_context_usage() function
+"""
+
+import pytest
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+from dataclasses import asdict
+
+from kiro.streaming_core import (
+ KiroEvent,
+ StreamResult,
+ FirstTokenTimeoutError,
+ parse_kiro_stream,
+ collect_stream_to_result,
+ calculate_tokens_from_context_usage,
+ stream_with_first_token_retry,
+ _process_chunk,
+)
+
+
+# ==================================================================================================
+# Fixtures
+# ==================================================================================================
+
+@pytest.fixture
+def mock_model_cache():
+ """Mock for ModelInfoCache."""
+ cache = MagicMock()
+ cache.get_max_input_tokens.return_value = 200000
+ return cache
+
+
+@pytest.fixture
+def mock_response():
+ """Mock for httpx.Response."""
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ return response
+
+
+@pytest.fixture
+def mock_parser():
+ """Mock for AwsEventStreamParser."""
+ parser = MagicMock()
+ parser.feed.return_value = []
+ parser.get_tool_calls.return_value = []
+ return parser
+
+
+# ==================================================================================================
+# Tests for KiroEvent dataclass
+# ==================================================================================================
+
+class TestKiroEvent:
+ """Tests for KiroEvent dataclass."""
+
+ def test_creates_content_event(self):
+ """
+ What it does: Creates a content event with text.
+ Goal: Verify KiroEvent can represent content events.
+ """
+ print("Action: Creating content event...")
+ event = KiroEvent(type="content", content="Hello, world!")
+
+ print(f"Comparing type: Expected 'content', Got '{event.type}'")
+ assert event.type == "content"
+ print(f"Comparing content: Expected 'Hello, world!', Got '{event.content}'")
+ assert event.content == "Hello, world!"
+ assert event.thinking_content is None
+ assert event.tool_use is None
+ print("✓ Content event created correctly")
+
+ def test_creates_thinking_event(self):
+ """
+ What it does: Creates a thinking event with reasoning content.
+ Goal: Verify KiroEvent can represent thinking events.
+ """
+ print("Action: Creating thinking event...")
+ event = KiroEvent(
+ type="thinking",
+ thinking_content="Let me think...",
+ is_first_thinking_chunk=True,
+ is_last_thinking_chunk=False
+ )
+
+ print(f"Comparing type: Expected 'thinking', Got '{event.type}'")
+ assert event.type == "thinking"
+ print(f"Comparing thinking_content: Expected 'Let me think...', Got '{event.thinking_content}'")
+ assert event.thinking_content == "Let me think..."
+ assert event.is_first_thinking_chunk is True
+ assert event.is_last_thinking_chunk is False
+ print("✓ Thinking event created correctly")
+
+ def test_creates_tool_use_event(self):
+ """
+ What it does: Creates a tool_use event with tool data.
+ Goal: Verify KiroEvent can represent tool use events.
+ """
+ print("Action: Creating tool_use event...")
+ tool_data = {
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "get_weather", "arguments": '{"city": "Moscow"}'}
+ }
+ event = KiroEvent(type="tool_use", tool_use=tool_data)
+
+ print(f"Comparing type: Expected 'tool_use', Got '{event.type}'")
+ assert event.type == "tool_use"
+ print(f"Comparing tool_use: Expected {tool_data}, Got {event.tool_use}")
+ assert event.tool_use == tool_data
+ print("✓ Tool use event created correctly")
+
+ def test_creates_usage_event(self):
+ """
+ What it does: Creates a usage event with metering data.
+ Goal: Verify KiroEvent can represent usage events.
+ """
+ print("Action: Creating usage event...")
+ usage_data = {"credits": 0.001}
+ event = KiroEvent(type="usage", usage=usage_data)
+
+ print(f"Comparing type: Expected 'usage', Got '{event.type}'")
+ assert event.type == "usage"
+ print(f"Comparing usage: Expected {usage_data}, Got {event.usage}")
+ assert event.usage == usage_data
+ print("✓ Usage event created correctly")
+
+ def test_creates_context_usage_event(self):
+ """
+ What it does: Creates a context_usage event with percentage.
+ Goal: Verify KiroEvent can represent context usage events.
+ """
+ print("Action: Creating context_usage event...")
+ event = KiroEvent(type="context_usage", context_usage_percentage=5.5)
+
+ print(f"Comparing type: Expected 'context_usage', Got '{event.type}'")
+ assert event.type == "context_usage"
+ print(f"Comparing context_usage_percentage: Expected 5.5, Got {event.context_usage_percentage}")
+ assert event.context_usage_percentage == 5.5
+ print("✓ Context usage event created correctly")
+
+ def test_default_values(self):
+ """
+ What it does: Verifies default values for optional fields.
+ Goal: Ensure all optional fields default to None/False.
+ """
+ print("Action: Creating minimal event...")
+ event = KiroEvent(type="content")
+
+ print("Checking default values...")
+ assert event.content is None
+ assert event.thinking_content is None
+ assert event.tool_use is None
+ assert event.usage is None
+ assert event.context_usage_percentage is None
+ assert event.is_first_thinking_chunk is False
+ assert event.is_last_thinking_chunk is False
+ print("✓ All default values are correct")
+
+
+# ==================================================================================================
+# Tests for StreamResult dataclass
+# ==================================================================================================
+
+class TestStreamResult:
+ """Tests for StreamResult dataclass."""
+
+ def test_creates_empty_result(self):
+ """
+ What it does: Creates an empty StreamResult.
+ Goal: Verify default values are correct.
+ """
+ print("Action: Creating empty StreamResult...")
+ result = StreamResult()
+
+ print("Checking default values...")
+ assert result.content == ""
+ assert result.thinking_content == ""
+ assert result.tool_calls == []
+ assert result.usage is None
+ assert result.context_usage_percentage is None
+ print("✓ Empty StreamResult created correctly")
+
+ def test_creates_result_with_content(self):
+ """
+ What it does: Creates StreamResult with content.
+ Goal: Verify content is stored correctly.
+ """
+ print("Action: Creating StreamResult with content...")
+ result = StreamResult(content="Hello, world!")
+
+ print(f"Comparing content: Expected 'Hello, world!', Got '{result.content}'")
+ assert result.content == "Hello, world!"
+ print("✓ StreamResult with content created correctly")
+
+ def test_creates_result_with_tool_calls(self):
+ """
+ What it does: Creates StreamResult with tool calls.
+ Goal: Verify tool calls are stored correctly.
+ """
+ print("Action: Creating StreamResult with tool calls...")
+ tool_calls = [
+ {"id": "call_1", "function": {"name": "func1"}},
+ {"id": "call_2", "function": {"name": "func2"}}
+ ]
+ result = StreamResult(tool_calls=tool_calls)
+
+ print(f"Comparing tool_calls count: Expected 2, Got {len(result.tool_calls)}")
+ assert len(result.tool_calls) == 2
+ assert result.tool_calls[0]["id"] == "call_1"
+ print("✓ StreamResult with tool calls created correctly")
+
+ def test_creates_result_with_usage(self):
+ """
+ What it does: Creates StreamResult with usage data.
+ Goal: Verify usage is stored correctly.
+ """
+ print("Action: Creating StreamResult with usage...")
+ usage = {"credits": 0.002}
+ result = StreamResult(usage=usage)
+
+ print(f"Comparing usage: Expected {usage}, Got {result.usage}")
+ assert result.usage == usage
+ print("✓ StreamResult with usage created correctly")
+
+ def test_creates_full_result(self):
+ """
+ What it does: Creates StreamResult with all fields.
+ Goal: Verify all fields work together.
+ """
+ print("Action: Creating full StreamResult...")
+ result = StreamResult(
+ content="Response text",
+ thinking_content="Thinking...",
+ tool_calls=[{"id": "call_1"}],
+ usage={"credits": 0.001},
+ context_usage_percentage=3.5
+ )
+
+ print("Checking all fields...")
+ assert result.content == "Response text"
+ assert result.thinking_content == "Thinking..."
+ assert len(result.tool_calls) == 1
+ assert result.usage == {"credits": 0.001}
+ assert result.context_usage_percentage == 3.5
+ print("✓ Full StreamResult created correctly")
+
+
+# ==================================================================================================
+# Tests for FirstTokenTimeoutError
+# ==================================================================================================
+
+class TestFirstTokenTimeoutError:
+ """Tests for FirstTokenTimeoutError exception."""
+
+ def test_creates_exception_with_message(self):
+ """
+ What it does: Creates exception with custom message.
+ Goal: Verify exception message is stored correctly.
+ """
+ print("Action: Creating FirstTokenTimeoutError...")
+ error = FirstTokenTimeoutError("No response within 30 seconds")
+
+ print(f"Comparing message: Expected 'No response within 30 seconds', Got '{str(error)}'")
+ assert str(error) == "No response within 30 seconds"
+ print("✓ Exception created correctly")
+
+ def test_exception_is_catchable(self):
+ """
+ What it does: Verifies exception can be caught.
+ Goal: Ensure exception inherits from Exception.
+ """
+ print("Action: Raising and catching FirstTokenTimeoutError...")
+
+ with pytest.raises(FirstTokenTimeoutError) as exc_info:
+ raise FirstTokenTimeoutError("Timeout!")
+
+ print(f"Caught exception: {exc_info.value}")
+ assert "Timeout!" in str(exc_info.value)
+ print("✓ Exception is catchable")
+
+ def test_exception_inherits_from_exception(self):
+ """
+ What it does: Verifies inheritance chain.
+ Goal: Ensure proper exception hierarchy.
+ """
+ print("Action: Checking inheritance...")
+ error = FirstTokenTimeoutError("Test")
+
+ assert isinstance(error, Exception)
+ print("✓ FirstTokenTimeoutError inherits from Exception")
+
+
+# ==================================================================================================
+# Tests for parse_kiro_stream()
+# ==================================================================================================
+
+class TestParseKiroStream:
+ """Tests for parse_kiro_stream() function."""
+
+ @pytest.mark.asyncio
+ async def test_parses_content_events(self, mock_response, mock_parser):
+ """
+ What it does: Parses content events from Kiro stream.
+ Goal: Verify content events are yielded correctly.
+ """
+ print("Setup: Mock parser to return content events...")
+ mock_parser.feed.return_value = [
+ {"type": "content", "data": "Hello"},
+ {"type": "content", "data": " World"}
+ ]
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ print("Action: Parsing stream...")
+ events = []
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ async for event in parse_kiro_stream(mock_response, first_token_timeout=30):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+ content_events = [e for e in events if e.type == "content"]
+ print(f"Content events: {len(content_events)}")
+
+ assert len(content_events) == 2
+ assert content_events[0].content == "Hello"
+ assert content_events[1].content == " World"
+ print("✓ Content events parsed correctly")
+
+ @pytest.mark.asyncio
+ async def test_parses_usage_events(self, mock_response, mock_parser):
+ """
+ What it does: Parses usage events from Kiro stream.
+ Goal: Verify usage events are yielded correctly.
+ """
+ print("Setup: Mock parser to return usage event...")
+ mock_parser.feed.return_value = [
+ {"type": "usage", "data": {"credits": 0.001}}
+ ]
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ print("Action: Parsing stream...")
+ events = []
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ async for event in parse_kiro_stream(mock_response, first_token_timeout=30):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+ usage_events = [e for e in events if e.type == "usage"]
+
+ assert len(usage_events) == 1
+ assert usage_events[0].usage == {"credits": 0.001}
+ print("✓ Usage events parsed correctly")
+
+ @pytest.mark.asyncio
+ async def test_parses_context_usage_events(self, mock_response, mock_parser):
+ """
+ What it does: Parses context_usage events from Kiro stream.
+ Goal: Verify context usage percentage is yielded correctly.
+ """
+ print("Setup: Mock parser to return context_usage event...")
+ mock_parser.feed.return_value = [
+ {"type": "context_usage", "data": 5.5}
+ ]
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ print("Action: Parsing stream...")
+ events = []
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ async for event in parse_kiro_stream(mock_response, first_token_timeout=30):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+ context_events = [e for e in events if e.type == "context_usage"]
+
+ assert len(context_events) == 1
+ assert context_events[0].context_usage_percentage == 5.5
+ print("✓ Context usage events parsed correctly")
+
+ @pytest.mark.asyncio
+ async def test_yields_tool_calls_at_end(self, mock_response, mock_parser):
+ """
+ What it does: Yields tool calls collected during parsing.
+ Goal: Verify tool calls are yielded as events.
+ """
+ print("Setup: Mock parser with tool calls...")
+ mock_parser.feed.return_value = [{"type": "content", "data": "text"}]
+ mock_parser.get_tool_calls.return_value = [
+ {"id": "call_1", "function": {"name": "func1", "arguments": "{}"}}
+ ]
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ print("Action: Parsing stream...")
+ events = []
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ async for event in parse_kiro_stream(mock_response, first_token_timeout=30):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+ tool_events = [e for e in events if e.type == "tool_use"]
+
+ assert len(tool_events) == 1
+ assert tool_events[0].tool_use["id"] == "call_1"
+ print("✓ Tool calls yielded correctly")
+
+ @pytest.mark.asyncio
+ async def test_raises_timeout_on_first_token(self, mock_response):
+ """
+ What it does: Raises FirstTokenTimeoutError on timeout.
+ Goal: Verify timeout handling for first token.
+ """
+ print("Setup: Mock response that times out...")
+
+ async def mock_aiter_bytes():
+ 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):
+ with pytest.raises(FirstTokenTimeoutError) as exc_info:
+ async for event in parse_kiro_stream(mock_response, first_token_timeout=30):
+ pass
+
+ print(f"Caught exception: {exc_info.value}")
+ assert "30" in str(exc_info.value)
+ print("✓ FirstTokenTimeoutError raised on timeout")
+
+ @pytest.mark.asyncio
+ async def test_handles_empty_response(self, mock_response):
+ """
+ What it does: Handles empty response gracefully.
+ Goal: Verify no events yielded for empty response.
+ """
+ print("Setup: Mock empty response...")
+
+ async def mock_aiter_bytes():
+ return
+ yield # Make it a generator
+
+ 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):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+ assert len(events) == 0
+ print("✓ Empty response handled correctly")
+
+ @pytest.mark.asyncio
+ async def test_handles_generator_exit(self, mock_response, mock_parser):
+ """
+ What it does: Handles GeneratorExit gracefully.
+ Goal: Verify client disconnect is handled.
+ """
+ print("Setup: Mock response that raises GeneratorExit...")
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+ raise GeneratorExit()
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+ mock_parser.feed.return_value = [{"type": "content", "data": "Hello"}]
+
+ print("Action: Parsing stream with GeneratorExit...")
+ events = []
+ generator_exit_raised = False
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ try:
+ async for event in parse_kiro_stream(mock_response, first_token_timeout=30):
+ events.append(event)
+ except GeneratorExit:
+ generator_exit_raised = True
+
+ print(f"GeneratorExit raised: {generator_exit_raised}")
+ assert generator_exit_raised
+ print("✓ GeneratorExit handled correctly")
+
+
+# ==================================================================================================
+# Tests for _process_chunk()
+# ==================================================================================================
+
+class TestProcessChunk:
+ """Tests for _process_chunk() helper function."""
+
+ @pytest.mark.asyncio
+ async def test_processes_content_event(self, mock_parser):
+ """
+ What it does: Processes content event from chunk.
+ Goal: Verify content is converted to KiroEvent.
+ """
+ print("Setup: Mock parser with content event...")
+ mock_parser.feed.return_value = [{"type": "content", "data": "Hello"}]
+
+ print("Action: Processing chunk...")
+ events = []
+ async for event in _process_chunk(mock_parser, b'chunk', None):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+ assert len(events) == 1
+ assert events[0].type == "content"
+ assert events[0].content == "Hello"
+ print("✓ Content event processed correctly")
+
+ @pytest.mark.asyncio
+ async def test_processes_usage_event(self, mock_parser):
+ """
+ What it does: Processes usage event from chunk.
+ Goal: Verify usage is converted to KiroEvent.
+ """
+ print("Setup: Mock parser with usage event...")
+ mock_parser.feed.return_value = [{"type": "usage", "data": {"credits": 0.001}}]
+
+ print("Action: Processing chunk...")
+ events = []
+ async for event in _process_chunk(mock_parser, b'chunk', None):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+ assert len(events) == 1
+ assert events[0].type == "usage"
+ assert events[0].usage == {"credits": 0.001}
+ print("✓ Usage event processed correctly")
+
+ @pytest.mark.asyncio
+ async def test_processes_context_usage_event(self, mock_parser):
+ """
+ What it does: Processes context_usage event from chunk.
+ Goal: Verify context usage is converted to KiroEvent.
+ """
+ print("Setup: Mock parser with context_usage event...")
+ mock_parser.feed.return_value = [{"type": "context_usage", "data": 7.5}]
+
+ print("Action: Processing chunk...")
+ events = []
+ async for event in _process_chunk(mock_parser, b'chunk', None):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+ assert len(events) == 1
+ assert events[0].type == "context_usage"
+ assert events[0].context_usage_percentage == 7.5
+ print("✓ Context usage event processed correctly")
+
+ @pytest.mark.asyncio
+ async def test_processes_multiple_events(self, mock_parser):
+ """
+ What it does: Processes multiple events from single chunk.
+ Goal: Verify all events are yielded.
+ """
+ print("Setup: Mock parser with multiple events...")
+ mock_parser.feed.return_value = [
+ {"type": "content", "data": "Hello"},
+ {"type": "content", "data": " World"},
+ {"type": "usage", "data": {"credits": 0.001}}
+ ]
+
+ print("Action: Processing chunk...")
+ events = []
+ async for event in _process_chunk(mock_parser, b'chunk', None):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+ assert len(events) == 3
+ assert events[0].type == "content"
+ assert events[1].type == "content"
+ assert events[2].type == "usage"
+ print("✓ Multiple events processed correctly")
+
+ @pytest.mark.asyncio
+ async def test_processes_with_thinking_parser(self, mock_parser):
+ """
+ What it does: Processes content through thinking parser.
+ Goal: Verify thinking parser integration.
+ """
+ print("Setup: Mock parser and thinking parser...")
+ mock_parser.feed.return_value = [{"type": "content", "data": "Hello"}]
+
+ 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
+ )
+
+ print("Action: Processing chunk with thinking parser...")
+ events = []
+ async for event in _process_chunk(mock_parser, b'chunk', mock_thinking_parser):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+ assert len(events) == 1
+ assert events[0].type == "content"
+ assert events[0].content == "Hello"
+ print("✓ Thinking parser integration works correctly")
+
+ @pytest.mark.asyncio
+ async def test_yields_thinking_content(self, mock_parser):
+ """
+ What it does: Yields thinking content from thinking parser.
+ Goal: Verify thinking events are created.
+ """
+ print("Setup: Mock parser and thinking parser with thinking content...")
+ mock_parser.feed.return_value = [{"type": "content", "data": "Let me think"}]
+
+ mock_thinking_parser = MagicMock()
+ mock_thinking_parser.feed.return_value = MagicMock(
+ thinking_content="Let me think",
+ regular_content=None,
+ is_first_thinking_chunk=True,
+ is_last_thinking_chunk=True
+ )
+ mock_thinking_parser.process_for_output.return_value = "Let me think"
+
+ print("Action: Processing chunk with thinking content...")
+ events = []
+ async for event in _process_chunk(mock_parser, b'chunk', mock_thinking_parser):
+ events.append(event)
+
+ print(f"Received {len(events)} events")
+ thinking_events = [e for e in events if e.type == "thinking"]
+ assert len(thinking_events) == 1
+ assert thinking_events[0].thinking_content == "Let me think"
+ print("✓ Thinking content yielded correctly")
+
+
+# ==================================================================================================
+# Tests for collect_stream_to_result()
+# ==================================================================================================
+
+class TestCollectStreamToResult:
+ """Tests for collect_stream_to_result() function."""
+
+ @pytest.mark.asyncio
+ async def test_collects_content(self, mock_response, mock_parser):
+ """
+ What it does: Collects content from stream.
+ Goal: Verify content is accumulated correctly.
+ """
+ print("Setup: Mock parser with content events...")
+ mock_parser.feed.return_value = [
+ {"type": "content", "data": "Hello"},
+ {"type": "content", "data": " World"}
+ ]
+ mock_parser.get_tool_calls.return_value = []
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ print("Action: Collecting stream...")
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ with patch('kiro.streaming_core.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_to_result(mock_response, first_token_timeout=30)
+
+ print(f"Collected content: '{result.content}'")
+ assert result.content == "Hello World"
+ print("✓ Content collected correctly")
+
+ @pytest.mark.asyncio
+ async def test_collects_tool_calls(self, mock_response, mock_parser):
+ """
+ What it does: Collects tool calls from stream.
+ Goal: Verify tool calls are accumulated correctly.
+ """
+ print("Setup: Mock parser with tool calls...")
+ mock_parser.feed.return_value = [{"type": "content", "data": "text"}]
+ mock_parser.get_tool_calls.return_value = [
+ {"id": "call_1", "function": {"name": "func1", "arguments": "{}"}}
+ ]
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ print("Action: Collecting stream...")
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ with patch('kiro.streaming_core.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_to_result(mock_response, first_token_timeout=30)
+
+ print(f"Collected tool calls: {len(result.tool_calls)}")
+ assert len(result.tool_calls) == 1
+ assert result.tool_calls[0]["id"] == "call_1"
+ print("✓ Tool calls collected correctly")
+
+ @pytest.mark.asyncio
+ async def test_collects_usage(self, mock_response, mock_parser):
+ """
+ What it does: Collects usage from stream.
+ Goal: Verify usage is stored correctly.
+ """
+ print("Setup: Mock parser with usage event...")
+ mock_parser.feed.return_value = [
+ {"type": "content", "data": "text"},
+ {"type": "usage", "data": {"credits": 0.002}}
+ ]
+ mock_parser.get_tool_calls.return_value = []
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ print("Action: Collecting stream...")
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ with patch('kiro.streaming_core.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_to_result(mock_response, first_token_timeout=30)
+
+ print(f"Collected usage: {result.usage}")
+ assert result.usage == {"credits": 0.002}
+ print("✓ Usage collected correctly")
+
+ @pytest.mark.asyncio
+ async def test_collects_context_usage_percentage(self, mock_response, mock_parser):
+ """
+ What it does: Collects context usage percentage from stream.
+ Goal: Verify context usage is stored correctly.
+ """
+ print("Setup: Mock parser with context_usage event...")
+ mock_parser.feed.return_value = [
+ {"type": "content", "data": "text"},
+ {"type": "context_usage", "data": 8.5}
+ ]
+ mock_parser.get_tool_calls.return_value = []
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ print("Action: Collecting stream...")
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ with patch('kiro.streaming_core.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_to_result(mock_response, first_token_timeout=30)
+
+ print(f"Collected context_usage_percentage: {result.context_usage_percentage}")
+ assert result.context_usage_percentage == 8.5
+ print("✓ Context usage percentage collected correctly")
+
+ @pytest.mark.asyncio
+ async def test_collects_thinking_content(self, mock_response, mock_parser):
+ """
+ What it does: Collects thinking content from stream.
+ Goal: Verify thinking content is accumulated correctly.
+ """
+ print("Setup: Mock parser with thinking content...")
+ # We need to mock the thinking parser behavior
+ mock_parser.feed.return_value = [{"type": "content", "data": "thinking text"}]
+ mock_parser.get_tool_calls.return_value = []
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ # Create mock events that include thinking
+ mock_events = [
+ KiroEvent(type="thinking", thinking_content="Let me think..."),
+ KiroEvent(type="content", content="Here is my answer")
+ ]
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ for event in mock_events:
+ yield event
+
+ print("Action: Collecting stream with thinking...")
+
+ with patch('kiro.streaming_core.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_core.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_to_result(mock_response, first_token_timeout=30)
+
+ print(f"Collected thinking_content: '{result.thinking_content}'")
+ print(f"Collected content: '{result.content}'")
+ assert result.thinking_content == "Let me think..."
+ assert result.content == "Here is my answer"
+ print("✓ Thinking content collected correctly")
+
+ @pytest.mark.asyncio
+ async def test_deduplicates_bracket_tool_calls(self, mock_response, mock_parser):
+ """
+ What it does: Deduplicates bracket-style tool calls.
+ Goal: Verify duplicate tool calls are removed.
+ """
+ print("Setup: Mock parser with tool calls and bracket tool calls...")
+ mock_parser.feed.return_value = [{"type": "content", "data": "text"}]
+ mock_parser.get_tool_calls.return_value = [
+ {"id": "call_1", "function": {"name": "func1", "arguments": "{}"}}
+ ]
+
+ bracket_tool_calls = [
+ {"id": "call_1", "function": {"name": "func1", "arguments": "{}"}}, # Duplicate
+ {"id": "call_2", "function": {"name": "func2", "arguments": "{}"}} # New
+ ]
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ print("Action: Collecting stream with duplicates...")
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ with patch('kiro.streaming_core.parse_bracket_tool_calls', return_value=bracket_tool_calls):
+ with patch('kiro.streaming_core.deduplicate_tool_calls') as mock_dedup:
+ mock_dedup.return_value = [
+ {"id": "call_1", "function": {"name": "func1", "arguments": "{}"}},
+ {"id": "call_2", "function": {"name": "func2", "arguments": "{}"}}
+ ]
+ result = await collect_stream_to_result(mock_response, first_token_timeout=30)
+
+ print(f"Collected tool calls: {len(result.tool_calls)}")
+ assert len(result.tool_calls) == 2
+ print("✓ Tool calls deduplicated correctly")
+
+
+# ==================================================================================================
+# Tests for calculate_tokens_from_context_usage()
+# ==================================================================================================
+
+class TestCalculateTokensFromContextUsage:
+ """Tests for calculate_tokens_from_context_usage() function."""
+
+ def test_calculates_tokens_from_percentage(self, mock_model_cache):
+ """
+ What it does: Calculates tokens from context usage percentage.
+ Goal: Verify token calculation is correct.
+ """
+ print("Setup: Context usage 10% with 200000 max tokens...")
+ context_usage_percentage = 10.0
+ completion_tokens = 100
+
+ print("Action: Calculating tokens...")
+ prompt_tokens, total_tokens, prompt_source, total_source = calculate_tokens_from_context_usage(
+ context_usage_percentage, completion_tokens, mock_model_cache, "claude-sonnet-4"
+ )
+
+ # 10% of 200000 = 20000 total tokens
+ # prompt_tokens = 20000 - 100 = 19900
+ print(f"Comparing total_tokens: Expected 20000, Got {total_tokens}")
+ assert total_tokens == 20000
+ print(f"Comparing prompt_tokens: Expected 19900, Got {prompt_tokens}")
+ assert prompt_tokens == 19900
+ assert prompt_source == "subtraction"
+ assert total_source == "API Kiro"
+ print("✓ Tokens calculated correctly")
+
+ def test_handles_zero_percentage(self, mock_model_cache):
+ """
+ What it does: Handles zero context usage percentage.
+ Goal: Verify fallback behavior for zero percentage.
+ """
+ print("Setup: Context usage 0%...")
+ context_usage_percentage = 0.0
+ completion_tokens = 100
+
+ print("Action: Calculating tokens...")
+ prompt_tokens, total_tokens, prompt_source, total_source = calculate_tokens_from_context_usage(
+ context_usage_percentage, completion_tokens, mock_model_cache, "claude-sonnet-4"
+ )
+
+ print(f"Comparing prompt_tokens: Expected 0, Got {prompt_tokens}")
+ assert prompt_tokens == 0
+ print(f"Comparing total_tokens: Expected 100, Got {total_tokens}")
+ assert total_tokens == 100
+ assert prompt_source == "unknown"
+ assert total_source == "tiktoken"
+ print("✓ Zero percentage handled correctly")
+
+ def test_handles_none_percentage(self, mock_model_cache):
+ """
+ What it does: Handles None context usage percentage.
+ Goal: Verify fallback behavior for None percentage.
+ """
+ print("Setup: Context usage None...")
+ context_usage_percentage = None
+ completion_tokens = 100
+
+ print("Action: Calculating tokens...")
+ prompt_tokens, total_tokens, prompt_source, total_source = calculate_tokens_from_context_usage(
+ context_usage_percentage, completion_tokens, mock_model_cache, "claude-sonnet-4"
+ )
+
+ print(f"Comparing prompt_tokens: Expected 0, Got {prompt_tokens}")
+ assert prompt_tokens == 0
+ print(f"Comparing total_tokens: Expected 100, Got {total_tokens}")
+ assert total_tokens == 100
+ assert prompt_source == "unknown"
+ assert total_source == "tiktoken"
+ print("✓ None percentage handled correctly")
+
+ def test_prevents_negative_prompt_tokens(self, mock_model_cache):
+ """
+ What it does: Prevents negative prompt tokens.
+ Goal: Verify prompt_tokens is never negative.
+ """
+ print("Setup: Very small context usage with large completion...")
+ context_usage_percentage = 0.01 # 0.01% of 200000 = 20 total tokens
+ completion_tokens = 100 # More than total!
+
+ print("Action: Calculating tokens...")
+ prompt_tokens, total_tokens, prompt_source, total_source = calculate_tokens_from_context_usage(
+ context_usage_percentage, completion_tokens, mock_model_cache, "claude-sonnet-4"
+ )
+
+ print(f"Comparing prompt_tokens: Expected >= 0, Got {prompt_tokens}")
+ assert prompt_tokens >= 0
+ print("✓ Negative prompt tokens prevented")
+
+ def test_uses_model_specific_max_tokens(self, mock_model_cache):
+ """
+ What it does: Uses model-specific max input tokens.
+ Goal: Verify model cache is queried correctly.
+ """
+ print("Setup: Different max tokens for model...")
+ mock_model_cache.get_max_input_tokens.return_value = 100000 # Different from default
+ context_usage_percentage = 10.0
+ completion_tokens = 100
+
+ print("Action: Calculating tokens...")
+ prompt_tokens, total_tokens, prompt_source, total_source = calculate_tokens_from_context_usage(
+ context_usage_percentage, completion_tokens, mock_model_cache, "claude-haiku-3"
+ )
+
+ # 10% of 100000 = 10000 total tokens
+ print(f"Comparing total_tokens: Expected 10000, Got {total_tokens}")
+ assert total_tokens == 10000
+
+ # Verify model cache was called with correct model
+ mock_model_cache.get_max_input_tokens.assert_called_with("claude-haiku-3")
+ print("✓ Model-specific max tokens used correctly")
+
+ def test_small_percentage_calculation(self, mock_model_cache):
+ """
+ What it does: Calculates tokens for small percentage.
+ Goal: Verify precision for small percentages.
+ """
+ print("Setup: Context usage 0.5%...")
+ context_usage_percentage = 0.5
+ completion_tokens = 50
+
+ print("Action: Calculating tokens...")
+ prompt_tokens, total_tokens, prompt_source, total_source = calculate_tokens_from_context_usage(
+ context_usage_percentage, completion_tokens, mock_model_cache, "claude-sonnet-4"
+ )
+
+ # 0.5% of 200000 = 1000 total tokens
+ # prompt_tokens = 1000 - 50 = 950
+ print(f"Comparing total_tokens: Expected 1000, Got {total_tokens}")
+ assert total_tokens == 1000
+ print(f"Comparing prompt_tokens: Expected 950, Got {prompt_tokens}")
+ assert prompt_tokens == 950
+ print("✓ Small percentage calculated correctly")
+
+ def test_large_percentage_calculation(self, mock_model_cache):
+ """
+ What it does: Calculates tokens for large percentage.
+ Goal: Verify calculation for high context usage.
+ """
+ print("Setup: Context usage 95%...")
+ context_usage_percentage = 95.0
+ completion_tokens = 1000
+
+ print("Action: Calculating tokens...")
+ prompt_tokens, total_tokens, prompt_source, total_source = calculate_tokens_from_context_usage(
+ context_usage_percentage, completion_tokens, mock_model_cache, "claude-sonnet-4"
+ )
+
+ # 95% of 200000 = 190000 total tokens
+ # prompt_tokens = 190000 - 1000 = 189000
+ print(f"Comparing total_tokens: Expected 190000, Got {total_tokens}")
+ assert total_tokens == 190000
+ print(f"Comparing prompt_tokens: Expected 189000, Got {prompt_tokens}")
+ assert prompt_tokens == 189000
+ print("✓ Large percentage calculated correctly")
+
+
+# ==================================================================================================
+# Tests for thinking parser integration
+# ==================================================================================================
+
+class TestThinkingParserIntegration:
+ """Tests for thinking parser integration in streaming."""
+
+ @pytest.mark.asyncio
+ async def test_thinking_parser_enabled_when_fake_reasoning_on(self, mock_response, mock_parser):
+ """
+ What it does: Enables thinking parser when FAKE_REASONING_ENABLED is True.
+ Goal: Verify thinking parser is created.
+ """
+ print("Setup: Enable fake reasoning...")
+ mock_parser.feed.return_value = [{"type": "content", "data": "Hello"}]
+ mock_parser.get_tool_calls.return_value = []
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ print("Action: Parsing stream with fake reasoning enabled...")
+ 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()
+
+ print("✓ Thinking parser enabled when fake reasoning is on")
+
+ @pytest.mark.asyncio
+ async def test_thinking_parser_disabled_when_fake_reasoning_off(self, mock_response, mock_parser):
+ """
+ What it does: Disables thinking parser when FAKE_REASONING_ENABLED is False.
+ Goal: Verify thinking parser is not created.
+ """
+ print("Setup: Disable fake reasoning...")
+ mock_parser.feed.return_value = [{"type": "content", "data": "Hello"}]
+ mock_parser.get_tool_calls.return_value = []
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ print("Action: Parsing stream with fake reasoning disabled...")
+ events = []
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ with patch('kiro.streaming_core.ThinkingParser') as mock_thinking_parser_class:
+ async for event in parse_kiro_stream(mock_response, first_token_timeout=30):
+ events.append(event)
+
+ # Verify ThinkingParser was NOT instantiated
+ mock_thinking_parser_class.assert_not_called()
+
+ print("✓ Thinking parser disabled when fake reasoning is off")
+
+ @pytest.mark.asyncio
+ async def test_thinking_parser_can_be_disabled_via_parameter(self, mock_response, mock_parser):
+ """
+ What it does: Disables thinking parser via enable_thinking_parser parameter.
+ Goal: Verify parameter overrides config.
+ """
+ print("Setup: Enable fake reasoning but disable via parameter...")
+ mock_parser.feed.return_value = [{"type": "content", "data": "Hello"}]
+ mock_parser.get_tool_calls.return_value = []
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+
+ print("Action: Parsing stream with thinking parser disabled via parameter...")
+ 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:
+ async for event in parse_kiro_stream(
+ mock_response,
+ first_token_timeout=30,
+ enable_thinking_parser=False
+ ):
+ events.append(event)
+
+ # Verify ThinkingParser was NOT instantiated
+ mock_thinking_parser_class.assert_not_called()
+
+ print("✓ Thinking parser disabled via parameter")
+
+
+# ==================================================================================================
+# Tests for error handling
+# ==================================================================================================
+
+class TestStreamingCoreErrorHandling:
+ """Tests for error handling in streaming_core."""
+
+ @pytest.mark.asyncio
+ async def test_propagates_first_token_timeout_error(self, mock_response):
+ """
+ What it does: Propagates FirstTokenTimeoutError.
+ Goal: Verify timeout error is not caught internally.
+ """
+ print("Setup: Mock response that times out...")
+
+ async def mock_aiter_bytes():
+ 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):
+ with pytest.raises(FirstTokenTimeoutError):
+ async for event in parse_kiro_stream(mock_response, first_token_timeout=30):
+ pass
+
+ print("✓ FirstTokenTimeoutError propagated correctly")
+
+ @pytest.mark.asyncio
+ async def test_propagates_generator_exit(self, mock_response, mock_parser):
+ """
+ What it does: Propagates GeneratorExit.
+ Goal: Verify client disconnect is handled.
+ """
+ print("Setup: Mock response that raises GeneratorExit...")
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+ raise GeneratorExit()
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+ mock_parser.feed.return_value = [{"type": "content", "data": "Hello"}]
+
+ print("Action: Parsing stream with GeneratorExit...")
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ with pytest.raises(GeneratorExit):
+ async for event in parse_kiro_stream(mock_response, first_token_timeout=30):
+ pass
+
+ print("✓ GeneratorExit propagated correctly")
+
+ @pytest.mark.asyncio
+ async def test_propagates_other_exceptions(self, mock_response, mock_parser):
+ """
+ What it does: Propagates other exceptions.
+ Goal: Verify errors are not swallowed.
+ """
+ print("Setup: Mock response that raises RuntimeError...")
+
+ async def mock_aiter_bytes():
+ yield b'chunk1'
+ raise RuntimeError("Test error")
+
+ mock_response.aiter_bytes = mock_aiter_bytes
+ mock_parser.feed.return_value = [{"type": "content", "data": "Hello"}]
+
+ print("Action: Parsing stream with RuntimeError...")
+
+ with patch('kiro.streaming_core.AwsEventStreamParser', return_value=mock_parser):
+ with patch('kiro.streaming_core.FAKE_REASONING_ENABLED', False):
+ with pytest.raises(RuntimeError) as exc_info:
+ async for event in parse_kiro_stream(mock_response, first_token_timeout=30):
+ pass
+
+ print(f"Caught exception: {exc_info.value}")
+ assert "Test error" in str(exc_info.value)
+ print("✓ RuntimeError propagated correctly")
+
+
+# ==================================================================================================
+# Tests for stream_with_first_token_retry()
+# ==================================================================================================
+
+class TestStreamWithFirstTokenRetryCore:
+ """
+ Tests for stream_with_first_token_retry() generic function.
+
+ This function provides automatic retry logic on first token timeout.
+ It is used by both OpenAI and Anthropic streaming implementations.
+ """
+
+ @pytest.mark.asyncio
+ async def test_yields_chunks_on_success(self):
+ """
+ What it does: Yields chunks on successful streaming.
+ Goal: Verify normal operation without retries.
+ """
+ print("Setup: Mock successful request...")
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.aclose = AsyncMock()
+
+ async def mock_make_request():
+ return mock_response
+
+ async def mock_stream_processor(response):
+ yield "chunk1"
+ yield "chunk2"
+ yield "chunk3"
+
+ print("Action: Streaming with retry wrapper...")
+ chunks = []
+
+ async for chunk in stream_with_first_token_retry(
+ make_request=mock_make_request,
+ stream_processor=mock_stream_processor,
+ max_retries=3,
+ first_token_timeout=30
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+ assert len(chunks) == 3
+ assert chunks == ["chunk1", "chunk2", "chunk3"]
+ print("✓ Chunks yielded on success")
+
+ @pytest.mark.asyncio
+ async def test_retries_on_first_token_timeout(self):
+ """
+ What it does: Retries on first token timeout.
+ Goal: Verify retry logic is triggered.
+ """
+ print("Setup: Mock request that times out then succeeds...")
+
+ call_count = 0
+
+ async def mock_make_request():
+ nonlocal call_count
+ call_count += 1
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ return response
+
+ async def mock_stream_processor(response):
+ nonlocal call_count
+ if call_count == 1:
+ raise FirstTokenTimeoutError("Timeout on first attempt")
+ yield "success_chunk"
+
+ print("Action: Streaming with retry on timeout...")
+ chunks = []
+
+ async for chunk in stream_with_first_token_retry(
+ make_request=mock_make_request,
+ stream_processor=mock_stream_processor,
+ max_retries=3,
+ first_token_timeout=30
+ ):
+ chunks.append(chunk)
+
+ print(f"Call count: {call_count}")
+ print(f"Received {len(chunks)} chunks")
+
+ assert call_count == 2 # First timeout, second success
+ assert len(chunks) == 1
+ assert chunks[0] == "success_chunk"
+ print("✓ Retry on timeout works correctly")
+
+ @pytest.mark.asyncio
+ async def test_raises_exception_after_all_retries(self):
+ """
+ What it does: Raises exception after all retries exhausted.
+ Goal: Verify error handling when all retries fail.
+ """
+ print("Setup: Mock request that always times out...")
+
+ call_count = 0
+
+ async def mock_make_request():
+ nonlocal call_count
+ call_count += 1
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ return response
+
+ async def mock_stream_processor(response):
+ raise FirstTokenTimeoutError("Timeout!")
+ yield # Make it a generator
+
+ print("Action: Streaming with all retries failing...")
+
+ with pytest.raises(Exception) as exc_info:
+ async for chunk in stream_with_first_token_retry(
+ make_request=mock_make_request,
+ stream_processor=mock_stream_processor,
+ max_retries=3,
+ first_token_timeout=30
+ ):
+ pass
+
+ print(f"Call count: {call_count}")
+ print(f"Exception: {exc_info.value}")
+
+ assert call_count == 3 # Should try exactly 3 times
+ assert "30" in str(exc_info.value) # Timeout value in message
+ assert "3" in str(exc_info.value) # Retry count in message
+ print("✓ Exception raised after all retries")
+
+ @pytest.mark.asyncio
+ async def test_uses_custom_error_callbacks(self):
+ """
+ What it does: Uses custom error callbacks.
+ Goal: Verify on_http_error and on_all_retries_failed callbacks.
+ """
+ print("Setup: Mock request that always times out with custom callbacks...")
+
+ async def mock_make_request():
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ return response
+
+ async def mock_stream_processor(response):
+ raise FirstTokenTimeoutError("Timeout!")
+ yield # Make it a generator
+
+ def custom_all_retries_failed(max_retries, timeout):
+ return ValueError(f"Custom error: {max_retries} retries, {timeout}s timeout")
+
+ print("Action: Streaming with custom callback...")
+
+ with pytest.raises(ValueError) as exc_info:
+ async for chunk in stream_with_first_token_retry(
+ make_request=mock_make_request,
+ stream_processor=mock_stream_processor,
+ max_retries=2,
+ first_token_timeout=15,
+ on_all_retries_failed=custom_all_retries_failed
+ ):
+ pass
+
+ print(f"Exception: {exc_info.value}")
+ assert "Custom error" in str(exc_info.value)
+ assert "2 retries" in str(exc_info.value)
+ assert "15" in str(exc_info.value)
+ print("✓ Custom callback used correctly")
+
+ @pytest.mark.asyncio
+ async def test_handles_http_error(self):
+ """
+ What it does: Handles HTTP error from API.
+ Goal: Verify HTTP errors are handled correctly.
+ """
+ print("Setup: Mock request that returns HTTP error...")
+
+ async def mock_make_request():
+ response = AsyncMock()
+ response.status_code = 500
+ response.aread = AsyncMock(return_value=b"Internal Server Error")
+ response.aclose = AsyncMock()
+ return response
+
+ async def mock_stream_processor(response):
+ yield "should not reach"
+
+ print("Action: Streaming with HTTP error...")
+
+ with pytest.raises(Exception) as exc_info:
+ async for chunk in stream_with_first_token_retry(
+ make_request=mock_make_request,
+ stream_processor=mock_stream_processor,
+ max_retries=3,
+ first_token_timeout=30
+ ):
+ pass
+
+ print(f"Exception: {exc_info.value}")
+ assert "500" in str(exc_info.value)
+ assert "Internal Server Error" in str(exc_info.value)
+ print("✓ HTTP error handled correctly")
+
+ @pytest.mark.asyncio
+ async def test_uses_custom_http_error_callback(self):
+ """
+ What it does: Uses custom HTTP error callback.
+ Goal: Verify on_http_error callback is used.
+ """
+ print("Setup: Mock request with custom HTTP error callback...")
+
+ async def mock_make_request():
+ response = AsyncMock()
+ response.status_code = 429
+ response.aread = AsyncMock(return_value=b"Rate limited")
+ response.aclose = AsyncMock()
+ return response
+
+ async def mock_stream_processor(response):
+ yield "should not reach"
+
+ def custom_http_error(status_code, error_text):
+ return RuntimeError(f"Custom HTTP error: {status_code} - {error_text}")
+
+ print("Action: Streaming with custom HTTP error callback...")
+
+ with pytest.raises(RuntimeError) as exc_info:
+ async for chunk in stream_with_first_token_retry(
+ make_request=mock_make_request,
+ stream_processor=mock_stream_processor,
+ max_retries=3,
+ first_token_timeout=30,
+ on_http_error=custom_http_error
+ ):
+ pass
+
+ print(f"Exception: {exc_info.value}")
+ assert "Custom HTTP error" in str(exc_info.value)
+ assert "429" in str(exc_info.value)
+ assert "Rate limited" in str(exc_info.value)
+ print("✓ Custom HTTP error callback used correctly")
+
+ @pytest.mark.asyncio
+ async def test_closes_response_on_timeout(self):
+ """
+ What it does: Closes response on timeout.
+ Goal: Verify response is properly closed after timeout.
+ """
+ print("Setup: Mock request that times out...")
+
+ responses = []
+
+ async def mock_make_request():
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ responses.append(response)
+ return response
+
+ async def mock_stream_processor(response):
+ raise FirstTokenTimeoutError("Timeout!")
+ yield # Make it a generator
+
+ print("Action: Streaming with timeout...")
+
+ try:
+ async for chunk in stream_with_first_token_retry(
+ make_request=mock_make_request,
+ stream_processor=mock_stream_processor,
+ max_retries=2,
+ first_token_timeout=30
+ ):
+ pass
+ except Exception:
+ pass
+
+ print(f"Created {len(responses)} responses")
+
+ # All responses should have been closed
+ for i, response in enumerate(responses):
+ print(f"Response {i} aclose called: {response.aclose.called}")
+ response.aclose.assert_called()
+
+ print("✓ Responses closed on timeout")
+
+ @pytest.mark.asyncio
+ async def test_propagates_non_timeout_exceptions(self):
+ """
+ What it does: Propagates non-timeout exceptions without retry.
+ Goal: Verify other exceptions are not retried.
+ """
+ print("Setup: Mock request that raises RuntimeError...")
+
+ call_count = 0
+
+ async def mock_make_request():
+ nonlocal call_count
+ call_count += 1
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ return response
+
+ async def mock_stream_processor(response):
+ raise RuntimeError("Not a timeout error")
+ yield # Make it a generator
+
+ print("Action: Streaming with non-timeout error...")
+
+ with pytest.raises(RuntimeError) as exc_info:
+ async for chunk in stream_with_first_token_retry(
+ make_request=mock_make_request,
+ stream_processor=mock_stream_processor,
+ max_retries=3,
+ first_token_timeout=30
+ ):
+ pass
+
+ print(f"Call count: {call_count}")
+ print(f"Exception: {exc_info.value}")
+
+ assert call_count == 1 # Should NOT retry
+ assert "Not a timeout error" in str(exc_info.value)
+ print("✓ Non-timeout exceptions propagated without retry")
+
+ @pytest.mark.asyncio
+ async def test_uses_configured_max_retries(self):
+ """
+ What it does: Uses configured max_retries value.
+ Goal: Verify max_retries parameter is respected.
+ """
+ print("Setup: Mock request that always times out...")
+
+ call_count = 0
+
+ async def mock_make_request():
+ nonlocal call_count
+ call_count += 1
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ return response
+
+ async def mock_stream_processor(response):
+ raise FirstTokenTimeoutError("Timeout!")
+ yield # Make it a generator
+
+ print("Action: Streaming with max_retries=5...")
+
+ try:
+ async for chunk in stream_with_first_token_retry(
+ make_request=mock_make_request,
+ stream_processor=mock_stream_processor,
+ max_retries=5,
+ first_token_timeout=30
+ ):
+ pass
+ except Exception:
+ pass
+
+ print(f"Call count: {call_count}")
+ assert call_count == 5 # Should try exactly 5 times
+ print("✓ max_retries parameter respected")
+
+ @pytest.mark.asyncio
+ async def test_multiple_retries_then_success(self):
+ """
+ What it does: Succeeds after multiple retries.
+ Goal: Verify recovery after multiple failures.
+ """
+ print("Setup: Mock request that fails twice then succeeds...")
+
+ call_count = 0
+
+ async def mock_make_request():
+ nonlocal call_count
+ call_count += 1
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ return response
+
+ async def mock_stream_processor(response):
+ nonlocal call_count
+ if call_count < 3:
+ raise FirstTokenTimeoutError(f"Timeout on attempt {call_count}")
+ yield "finally_success"
+
+ print("Action: Streaming with multiple retries...")
+ chunks = []
+
+ async for chunk in stream_with_first_token_retry(
+ make_request=mock_make_request,
+ stream_processor=mock_stream_processor,
+ max_retries=5,
+ first_token_timeout=30
+ ):
+ chunks.append(chunk)
+
+ print(f"Call count: {call_count}")
+ print(f"Received {len(chunks)} chunks")
+
+ assert call_count == 3 # Failed twice, succeeded on third
+ assert len(chunks) == 1
+ assert chunks[0] == "finally_success"
+ print("✓ Multiple retries then success works correctly")
+
+ @pytest.mark.asyncio
+ async def test_closes_response_on_http_error(self):
+ """
+ What it does: Closes response on HTTP error.
+ Goal: Verify response is properly closed after HTTP error.
+ """
+ print("Setup: Mock request that returns HTTP error...")
+
+ response = AsyncMock()
+ response.status_code = 503
+ response.aread = AsyncMock(return_value=b"Service Unavailable")
+ response.aclose = AsyncMock()
+
+ async def mock_make_request():
+ return response
+
+ async def mock_stream_processor(resp):
+ yield "should not reach"
+
+ print("Action: Streaming with HTTP error...")
+
+ try:
+ async for chunk in stream_with_first_token_retry(
+ make_request=mock_make_request,
+ stream_processor=mock_stream_processor,
+ max_retries=3,
+ first_token_timeout=30
+ ):
+ pass
+ except Exception:
+ pass
+
+ print(f"Response aclose called: {response.aclose.called}")
+ response.aclose.assert_called()
+ print("✓ Response closed on HTTP error")
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_streaming_openai.py b/kiro-gateway/tests/unit/test_streaming_openai.py
new file mode 100644
index 0000000000000000000000000000000000000000..6a97514019df40495c094b3adb3c6c9ad8ec111e
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_streaming_openai.py
@@ -0,0 +1,1353 @@
+
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for streaming_openai module.
+
+Tests for:
+- stream_kiro_to_openai() generator
+- stream_kiro_to_openai_internal() generator
+- stream_with_first_token_retry() function
+- collect_stream_response() function
+"""
+
+import pytest
+import json
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from kiro.streaming_openai import (
+ stream_kiro_to_openai,
+ stream_kiro_to_openai_internal,
+ stream_with_first_token_retry,
+ collect_stream_response,
+ FirstTokenTimeoutError,
+)
+from kiro.streaming_core import KiroEvent
+
+
+# ==================================================================================================
+# Fixtures
+# ==================================================================================================
+
+@pytest.fixture
+def mock_model_cache():
+ """Mock for ModelInfoCache."""
+ cache = MagicMock()
+ cache.get_max_input_tokens.return_value = 200000
+ return cache
+
+
+@pytest.fixture
+def mock_auth_manager():
+ """Mock for KiroAuthManager."""
+ manager = MagicMock()
+ return manager
+
+
+@pytest.fixture
+def mock_http_client():
+ """Mock for httpx.AsyncClient."""
+ client = AsyncMock()
+ return client
+
+
+@pytest.fixture
+def mock_response():
+ """Mock for httpx.Response."""
+ response = AsyncMock()
+ response.status_code = 200
+ response.aclose = AsyncMock()
+ return response
+
+
+# ==================================================================================================
+# Tests for stream_kiro_to_openai()
+# ==================================================================================================
+
+class TestStreamKiroToOpenai:
+ """Tests for stream_kiro_to_openai() generator."""
+
+ @pytest.mark.asyncio
+ async def test_yields_content_chunks(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields content chunks in OpenAI format.
+ Goal: Verify content streaming.
+ """
+ print("Setup: Mock stream with content events...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ yield KiroEvent(type="content", content=" World")
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Should have content chunks
+ content_chunks = [c for c in chunks if "content" in c and '"Hello"' in c or '" World"' in c]
+ assert len(content_chunks) >= 2
+ print("✓ Content chunks yielded correctly")
+
+ @pytest.mark.asyncio
+ async def test_first_chunk_has_role(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: First chunk includes role: assistant.
+ Goal: Verify OpenAI streaming protocol.
+ """
+ print("Setup: Mock stream with content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # First content chunk should have role
+ first_content_chunk = [c for c in chunks if '"content"' in c and '"Hello"' in c][0]
+ assert '"role": "assistant"' in first_content_chunk
+ print("✓ First chunk has role: assistant")
+
+ @pytest.mark.asyncio
+ async def test_yields_done_at_end(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields [DONE] at end of stream.
+ Goal: Verify stream termination.
+ """
+ print("Setup: Mock stream with content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Last chunk should be [DONE]
+ assert chunks[-1] == "data: [DONE]\n\n"
+ print("✓ [DONE] yielded at end")
+
+ @pytest.mark.asyncio
+ async def test_yields_final_chunk_with_usage(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields final chunk with usage info.
+ Goal: Verify usage is included.
+ """
+ print("Setup: Mock stream with content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ yield KiroEvent(type="context_usage", context_usage_percentage=5.0)
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Should have chunk with usage before [DONE]
+ usage_chunks = [c for c in chunks if '"usage"' in c]
+ assert len(usage_chunks) >= 1
+ print("✓ Final chunk with usage yielded")
+
+ @pytest.mark.asyncio
+ async def test_yields_tool_calls_chunk(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields tool_calls chunk when tools present.
+ Goal: Verify tool call streaming.
+ """
+ print("Setup: Mock stream with tool call...")
+
+ tool_use_data = {
+ "id": "call_123",
+ "type": "function",
+ "function": {"name": "get_weather", "arguments": '{"city": "Moscow"}'}
+ }
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Let me check")
+ yield KiroEvent(type="tool_use", tool_use=tool_use_data)
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Should have tool_calls chunk
+ tool_chunks = [c for c in chunks if '"tool_calls"' in c]
+ assert len(tool_chunks) >= 1
+ assert "get_weather" in tool_chunks[0]
+ print("✓ Tool calls chunk yielded")
+
+ @pytest.mark.asyncio
+ async def test_tool_calls_have_index(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Tool calls have index field.
+ Goal: Verify OpenAI streaming spec compliance.
+ """
+ print("Setup: Mock stream with multiple tool calls...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="tool_use", tool_use={
+ "id": "call_1", "type": "function",
+ "function": {"name": "func1", "arguments": "{}"}
+ })
+ yield KiroEvent(type="tool_use", tool_use={
+ "id": "call_2", "type": "function",
+ "function": {"name": "func2", "arguments": "{}"}
+ })
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Find tool_calls chunk and verify indices
+ tool_chunks = [c for c in chunks if '"tool_calls"' in c]
+ assert len(tool_chunks) >= 1
+
+ # Parse and check indices
+ for chunk in tool_chunks:
+ if chunk.startswith("data: "):
+ json_str = chunk[6:].strip()
+ if json_str != "[DONE]":
+ data = json.loads(json_str)
+ if "choices" in data and data["choices"]:
+ delta = data["choices"][0].get("delta", {})
+ if "tool_calls" in delta:
+ for tc in delta["tool_calls"]:
+ assert "index" in tc
+
+ print("✓ Tool calls have index field")
+
+ @pytest.mark.asyncio
+ async def test_finish_reason_is_tool_calls_when_tools_present(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Sets finish_reason to tool_calls when tools present.
+ Goal: Verify correct finish reason.
+ """
+ print("Setup: Mock stream with tool call...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="tool_use", tool_use={
+ "id": "call_1", "type": "function",
+ "function": {"name": "func1", "arguments": "{}"}
+ })
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Final chunk before [DONE] should have finish_reason: tool_calls
+ final_chunk = chunks[-2] # Before [DONE]
+ assert '"finish_reason": "tool_calls"' in final_chunk
+ print("✓ finish_reason is tool_calls")
+
+ @pytest.mark.asyncio
+ async def test_finish_reason_is_stop_without_tools(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Sets finish_reason to stop without tools.
+ Goal: Verify correct finish reason.
+ """
+ print("Setup: Mock stream without tool calls...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Final chunk before [DONE] should have finish_reason: stop
+ final_chunk = chunks[-2] # Before [DONE]
+ assert '"finish_reason": "stop"' in final_chunk
+ print("✓ finish_reason is stop")
+
+ @pytest.mark.asyncio
+ async def test_closes_response_on_completion(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Closes response on completion.
+ Goal: Verify resource cleanup.
+ """
+ print("Setup: Mock stream...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Streaming to OpenAI format...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ pass
+
+ print("Check: response.aclose() should be called...")
+ mock_response.aclose.assert_called()
+ print("✓ Response closed on completion")
+
+ @pytest.mark.asyncio
+ async def test_closes_response_on_error(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Closes response on error.
+ Goal: Verify resource cleanup on error.
+ """
+ print("Setup: Mock stream that raises error...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ raise RuntimeError("Test error")
+
+ print("Action: Streaming to OpenAI format with error...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ try:
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ pass
+ except RuntimeError:
+ pass
+
+ print("Check: response.aclose() should be called...")
+ mock_response.aclose.assert_called()
+ print("✓ Response closed on error")
+
+
+# ==================================================================================================
+# Tests for thinking content handling
+# ==================================================================================================
+
+class TestStreamingOpenaiThinkingContent:
+ """Tests for thinking content handling in OpenAI streaming."""
+
+ @pytest.mark.asyncio
+ async def test_yields_thinking_as_reasoning_content(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields thinking as reasoning_content when configured.
+ Goal: Verify thinking content handling.
+ """
+ print("Setup: Mock stream with thinking content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="thinking", thinking_content="Let me think...")
+ yield KiroEvent(type="content", content="Here is my answer")
+
+ print("Action: Streaming to OpenAI format with reasoning mode...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ with patch('kiro.streaming_openai.FAKE_REASONING_HANDLING', 'as_reasoning_content'):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Should have reasoning_content
+ reasoning_chunks = [c for c in chunks if '"reasoning_content"' in c]
+ assert len(reasoning_chunks) >= 1
+ assert "Let me think" in reasoning_chunks[0]
+ print("✓ Thinking yielded as reasoning_content")
+
+ @pytest.mark.asyncio
+ async def test_yields_thinking_as_content_when_configured(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Yields thinking as content when configured.
+ Goal: Verify thinking content handling.
+ """
+ print("Setup: Mock stream with thinking content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="thinking", thinking_content="Let me think...")
+ yield KiroEvent(type="content", content="Here is my answer")
+
+ print("Action: Streaming to OpenAI format with content mode...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ with patch('kiro.streaming_openai.FAKE_REASONING_HANDLING', 'include_as_text'):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Should have thinking as content
+ content_chunks = [c for c in chunks if '"content"' in c and "Let me think" in c]
+ assert len(content_chunks) >= 1
+ print("✓ Thinking yielded as content")
+
+
+# ==================================================================================================
+# Tests for None protection in tool calls
+# ==================================================================================================
+
+class TestStreamingOpenaiNoneProtection:
+ """Tests for None protection in tool calls."""
+
+ @pytest.mark.asyncio
+ async def test_handles_none_function_name(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Handles None in function.name.
+ Goal: Verify None is replaced with empty string.
+ """
+ print("Setup: Mock stream with None function name...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="tool_use", tool_use={
+ "id": "call_1", "type": "function",
+ "function": {"name": None, "arguments": "{}"}
+ })
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Should handle None gracefully
+ tool_chunks = [c for c in chunks if '"tool_calls"' in c]
+ assert len(tool_chunks) >= 1
+
+ # Parse and verify name is empty string
+ for chunk in tool_chunks:
+ if chunk.startswith("data: "):
+ json_str = chunk[6:].strip()
+ if json_str != "[DONE]":
+ data = json.loads(json_str)
+ if "choices" in data and data["choices"]:
+ delta = data["choices"][0].get("delta", {})
+ if "tool_calls" in delta:
+ for tc in delta["tool_calls"]:
+ assert tc["function"]["name"] == ""
+
+ print("✓ None function name handled")
+
+ @pytest.mark.asyncio
+ async def test_handles_none_function_arguments(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Handles None in function.arguments.
+ Goal: Verify None is replaced with "{}".
+ """
+ print("Setup: Mock stream with None arguments...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="tool_use", tool_use={
+ "id": "call_1", "type": "function",
+ "function": {"name": "func1", "arguments": None}
+ })
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Should handle None gracefully
+ tool_chunks = [c for c in chunks if '"tool_calls"' in c]
+ assert len(tool_chunks) >= 1
+
+ # Parse and verify arguments is "{}"
+ for chunk in tool_chunks:
+ if chunk.startswith("data: "):
+ json_str = chunk[6:].strip()
+ if json_str != "[DONE]":
+ data = json.loads(json_str)
+ if "choices" in data and data["choices"]:
+ delta = data["choices"][0].get("delta", {})
+ if "tool_calls" in delta:
+ for tc in delta["tool_calls"]:
+ assert tc["function"]["arguments"] == "{}"
+
+ print("✓ None function arguments handled")
+
+ @pytest.mark.asyncio
+ async def test_handles_none_function_object(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Handles None function object.
+ Goal: Verify None function is handled.
+ """
+ print("Setup: Mock stream with None function...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="tool_use", tool_use={
+ "id": "call_1", "type": "function",
+ "function": None
+ })
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Should handle None gracefully without error
+ assert len(chunks) > 0
+ print("✓ None function object handled")
+
+
+# ==================================================================================================
+# Tests for stream_with_first_token_retry()
+# ==================================================================================================
+
+class TestStreamWithFirstTokenRetry:
+ """Tests for stream_with_first_token_retry() function."""
+
+ @pytest.mark.asyncio
+ async def test_retries_on_first_token_timeout(self, mock_http_client, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Retries on first token timeout.
+ Goal: Verify retry logic.
+ """
+ print("Setup: Mock make_request that succeeds on second attempt...")
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.aclose = AsyncMock()
+
+ call_count = 0
+
+ async def mock_make_request():
+ nonlocal call_count
+ call_count += 1
+ print(f"make_request called (attempt {call_count})")
+ return mock_response
+
+ # First call raises timeout, second succeeds
+ timeout_raised = False
+
+ async def mock_parse_kiro_stream_with_retry(*args, **kwargs):
+ nonlocal timeout_raised
+ if not timeout_raised:
+ timeout_raised = True
+ raise FirstTokenTimeoutError("Timeout!")
+ yield KiroEvent(type="content", content="Success")
+
+ print("Action: Running stream_with_first_token_retry...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream_with_retry):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_with_first_token_retry(
+ mock_make_request,
+ mock_http_client,
+ "claude-sonnet-4",
+ mock_model_cache,
+ mock_auth_manager,
+ max_retries=3,
+ first_token_timeout=15
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+ print(f"make_request was called {call_count} times")
+
+ assert call_count == 2
+ assert len(chunks) > 0
+ print("✓ Retry logic worked correctly")
+
+ @pytest.mark.asyncio
+ async def test_raises_504_after_all_retries_exhausted(self, mock_http_client, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Raises 504 after all retries exhausted.
+ Goal: Verify error handling.
+ """
+ from fastapi import HTTPException
+
+ print("Setup: Mock make_request that always times out...")
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.aclose = AsyncMock()
+
+ call_count = 0
+
+ async def mock_make_request():
+ nonlocal call_count
+ call_count += 1
+ return mock_response
+
+ async def mock_parse_kiro_stream_always_timeout(*args, **kwargs):
+ raise FirstTokenTimeoutError("Timeout!")
+ yield # Make it a generator
+
+ max_retries = 3
+
+ print(f"Action: Running stream_with_first_token_retry with max_retries={max_retries}...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream_always_timeout):
+ with pytest.raises(HTTPException) as exc_info:
+ async for chunk in stream_with_first_token_retry(
+ mock_make_request,
+ mock_http_client,
+ "claude-sonnet-4",
+ mock_model_cache,
+ mock_auth_manager,
+ max_retries=max_retries,
+ first_token_timeout=15
+ ):
+ pass
+
+ print(f"Caught HTTPException: {exc_info.value.status_code}")
+ print(f"make_request was called {call_count} times")
+
+ assert exc_info.value.status_code == 504
+ assert call_count == max_retries
+ print("✓ 504 raised after all retries exhausted")
+
+ @pytest.mark.asyncio
+ async def test_handles_api_error_response(self, mock_http_client, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Handles API error response.
+ Goal: Verify error response handling.
+ """
+ from fastapi import HTTPException
+
+ print("Setup: Mock make_request that returns error...")
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 500
+ # Use simple error text without curly braces to avoid loguru format issues
+ mock_response.aread = AsyncMock(return_value=b'Internal server error')
+ mock_response.aclose = AsyncMock()
+
+ async def mock_make_request():
+ return mock_response
+
+ print("Action: Running stream_with_first_token_retry with error response...")
+
+ with pytest.raises(HTTPException) as exc_info:
+ async for chunk in stream_with_first_token_retry(
+ mock_make_request,
+ mock_http_client,
+ "claude-sonnet-4",
+ mock_model_cache,
+ mock_auth_manager,
+ max_retries=3,
+ first_token_timeout=15
+ ):
+ pass
+
+ print(f"Caught HTTPException: {exc_info.value.status_code}")
+ assert exc_info.value.status_code == 500
+ print("✓ API error response handled")
+
+ @pytest.mark.asyncio
+ async def test_propagates_non_timeout_errors(self, mock_http_client, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Propagates non-timeout errors without retry.
+ Goal: Verify only timeout errors trigger retry.
+ """
+ print("Setup: Mock make_request that raises RuntimeError...")
+
+ mock_response = AsyncMock()
+ mock_response.status_code = 200
+ mock_response.aclose = AsyncMock()
+
+ call_count = 0
+
+ async def mock_make_request():
+ nonlocal call_count
+ call_count += 1
+ return mock_response
+
+ async def mock_parse_kiro_stream_error(*args, **kwargs):
+ raise RuntimeError("Test error")
+ yield # Make it a generator
+
+ print("Action: Running stream_with_first_token_retry with RuntimeError...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream_error):
+ with pytest.raises(RuntimeError) as exc_info:
+ async for chunk in stream_with_first_token_retry(
+ mock_make_request,
+ mock_http_client,
+ "claude-sonnet-4",
+ mock_model_cache,
+ mock_auth_manager,
+ max_retries=3,
+ first_token_timeout=15
+ ):
+ pass
+
+ print(f"Caught RuntimeError: {exc_info.value}")
+ print(f"make_request was called {call_count} times")
+
+ # Should only be called once - no retry for non-timeout errors
+ assert call_count == 1
+ assert "Test error" in str(exc_info.value)
+ print("✓ Non-timeout errors propagated without retry")
+
+ @pytest.mark.asyncio
+ async def test_closes_response_on_retry(self, mock_http_client, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Closes response when retrying.
+ Goal: Verify resource cleanup on retry.
+ """
+ print("Setup: Mock responses for retry...")
+
+ mock_response1 = AsyncMock()
+ mock_response1.status_code = 200
+ mock_response1.aclose = AsyncMock()
+
+ mock_response2 = AsyncMock()
+ mock_response2.status_code = 200
+ mock_response2.aclose = AsyncMock()
+
+ responses = [mock_response1, mock_response2]
+ call_count = 0
+
+ async def mock_make_request():
+ nonlocal call_count
+ response = responses[call_count]
+ call_count += 1
+ return response
+
+ # First call raises timeout, second succeeds
+ timeout_raised = False
+
+ async def mock_parse_kiro_stream_with_retry(*args, **kwargs):
+ nonlocal timeout_raised
+ if not timeout_raised:
+ timeout_raised = True
+ raise FirstTokenTimeoutError("Timeout!")
+ yield KiroEvent(type="content", content="Success")
+
+ print("Action: Running stream_with_first_token_retry...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream_with_retry):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_with_first_token_retry(
+ mock_make_request,
+ mock_http_client,
+ "claude-sonnet-4",
+ mock_model_cache,
+ mock_auth_manager,
+ max_retries=3,
+ first_token_timeout=15
+ ):
+ pass
+
+ print("Check: First response should be closed...")
+ mock_response1.aclose.assert_called()
+ print("✓ Response closed on retry")
+
+
+# ==================================================================================================
+# Tests for collect_stream_response()
+# ==================================================================================================
+
+class TestCollectStreamResponse:
+ """Tests for collect_stream_response() function."""
+
+ @pytest.mark.asyncio
+ async def test_collects_content(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Collects content from stream.
+ Goal: Verify content accumulation.
+ """
+ print("Setup: Mock stream with content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ yield KiroEvent(type="content", content=" World")
+
+ print("Action: Collecting stream response...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_response(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Result: {result}")
+
+ assert result["choices"][0]["message"]["content"] == "Hello World"
+ print("✓ Content collected correctly")
+
+ @pytest.mark.asyncio
+ async def test_collects_reasoning_content(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Collects reasoning content from stream.
+ Goal: Verify reasoning content accumulation.
+ """
+ print("Setup: Mock stream with thinking content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="thinking", thinking_content="Let me think...")
+ yield KiroEvent(type="content", content="Answer")
+
+ print("Action: Collecting stream response with reasoning mode...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ with patch('kiro.streaming_openai.FAKE_REASONING_HANDLING', 'as_reasoning_content'):
+ result = await collect_stream_response(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Result: {result}")
+
+ message = result["choices"][0]["message"]
+ assert "reasoning_content" in message
+ assert message["reasoning_content"] == "Let me think..."
+ print("✓ Reasoning content collected correctly")
+
+ @pytest.mark.asyncio
+ async def test_collects_tool_calls(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Collects tool calls from stream.
+ Goal: Verify tool call accumulation.
+ """
+ print("Setup: Mock stream with tool calls...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="tool_use", tool_use={
+ "id": "call_1", "type": "function",
+ "function": {"name": "func1", "arguments": '{"a": 1}'}
+ })
+
+ print("Action: Collecting stream response...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_response(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Result: {result}")
+
+ message = result["choices"][0]["message"]
+ assert "tool_calls" in message
+ assert len(message["tool_calls"]) == 1
+ assert message["tool_calls"][0]["function"]["name"] == "func1"
+ print("✓ Tool calls collected correctly")
+
+ @pytest.mark.asyncio
+ async def test_tool_calls_have_no_index(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Collected tool calls don't have index field.
+ Goal: Verify index is removed for non-streaming.
+ """
+ print("Setup: Mock stream with tool calls...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="tool_use", tool_use={
+ "id": "call_1", "type": "function",
+ "function": {"name": "func1", "arguments": "{}"}
+ })
+
+ print("Action: Collecting stream response...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_response(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Result: {result}")
+
+ message = result["choices"][0]["message"]
+ for tc in message.get("tool_calls", []):
+ assert "index" not in tc
+
+ print("✓ Tool calls have no index field")
+
+ @pytest.mark.asyncio
+ async def test_includes_usage(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Includes usage in response.
+ Goal: Verify usage is included.
+ """
+ print("Setup: Mock stream with content...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ yield KiroEvent(type="context_usage", context_usage_percentage=5.0)
+
+ print("Action: Collecting stream response...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_response(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Result: {result}")
+
+ assert "usage" in result
+ assert "prompt_tokens" in result["usage"]
+ assert "completion_tokens" in result["usage"]
+ assert "total_tokens" in result["usage"]
+ print("✓ Usage included in response")
+
+ @pytest.mark.asyncio
+ async def test_sets_finish_reason_tool_calls(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Sets finish_reason to tool_calls when tools present.
+ Goal: Verify correct finish reason.
+ """
+ print("Setup: Mock stream with tool calls...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="tool_use", tool_use={
+ "id": "call_1", "type": "function",
+ "function": {"name": "func1", "arguments": "{}"}
+ })
+
+ print("Action: Collecting stream response...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_response(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Result: {result}")
+
+ assert result["choices"][0]["finish_reason"] == "tool_calls"
+ print("✓ finish_reason is tool_calls")
+
+ @pytest.mark.asyncio
+ async def test_sets_finish_reason_stop(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Sets finish_reason to stop without tools.
+ Goal: Verify correct finish reason.
+ """
+ print("Setup: Mock stream without tool calls...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Collecting stream response...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_response(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Result: {result}")
+
+ assert result["choices"][0]["finish_reason"] == "stop"
+ print("✓ finish_reason is stop")
+
+ @pytest.mark.asyncio
+ async def test_generates_completion_id(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Generates completion ID.
+ Goal: Verify ID is present.
+ """
+ print("Setup: Mock stream...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Collecting stream response...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_response(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ )
+
+ print(f"ID: {result['id']}")
+
+ assert result["id"].startswith("chatcmpl-")
+ print("✓ Completion ID generated")
+
+ @pytest.mark.asyncio
+ async def test_includes_model_name(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Includes model name in response.
+ Goal: Verify model is included.
+ """
+ print("Setup: Mock stream...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Collecting stream response...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_response(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Model: {result['model']}")
+
+ assert result["model"] == "claude-sonnet-4"
+ print("✓ Model name included")
+
+ @pytest.mark.asyncio
+ async def test_object_is_chat_completion(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Sets object to chat.completion.
+ Goal: Verify OpenAI format.
+ """
+ print("Setup: Mock stream...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+
+ print("Action: Collecting stream response...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ result = await collect_stream_response(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ )
+
+ print(f"Object: {result['object']}")
+
+ assert result["object"] == "chat.completion"
+ print("✓ Object is chat.completion")
+
+
+# ==================================================================================================
+# Tests for error handling
+# ==================================================================================================
+
+class TestStreamingOpenaiErrorHandling:
+ """Tests for error handling in streaming_openai."""
+
+ @pytest.mark.asyncio
+ async def test_propagates_first_token_timeout_error(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Propagates FirstTokenTimeoutError.
+ Goal: Verify timeout error is propagated for retry.
+ """
+ print("Setup: Mock stream that raises timeout...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ raise FirstTokenTimeoutError("Timeout!")
+ yield # Make it a generator
+
+ print("Action: Streaming to OpenAI format with timeout...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with pytest.raises(FirstTokenTimeoutError):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ pass
+
+ print("✓ FirstTokenTimeoutError propagated correctly")
+
+ @pytest.mark.asyncio
+ async def test_handles_generator_exit_gracefully(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Handles GeneratorExit gracefully without re-raising.
+ Goal: Verify client disconnect is handled without error.
+ """
+ print("Setup: Mock stream that raises GeneratorExit...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ raise GeneratorExit()
+
+ print("Action: Streaming to OpenAI format with GeneratorExit...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ # GeneratorExit is caught internally and not re-raised
+ # This is correct behavior - client disconnect should be handled gracefully
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks before disconnect")
+ # Response should be closed
+ mock_response.aclose.assert_called()
+ print("✓ GeneratorExit handled gracefully")
+
+ @pytest.mark.asyncio
+ async def test_propagates_other_exceptions(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Propagates other exceptions.
+ Goal: Verify errors are not swallowed.
+ """
+ print("Setup: Mock stream that raises RuntimeError...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ raise RuntimeError("Test error")
+
+ print("Action: Streaming to OpenAI format with RuntimeError...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ with pytest.raises(RuntimeError) as exc_info:
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ pass
+
+ print(f"Caught exception: {exc_info.value}")
+ assert "Test error" in str(exc_info.value)
+ print("✓ RuntimeError propagated correctly")
+
+ @pytest.mark.asyncio
+ async def test_aclose_error_does_not_mask_original(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: aclose() error doesn't mask original error.
+ Goal: Verify original exception is propagated.
+ """
+ print("Setup: Mock response with error in aclose()...")
+
+ mock_response.aclose = AsyncMock(side_effect=ConnectionError("Connection lost"))
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ raise RuntimeError("Original error")
+
+ print("Action: Streaming to OpenAI format with error and aclose error...")
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ with pytest.raises(RuntimeError) as exc_info:
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ pass
+
+ print(f"Caught exception: {exc_info.value}")
+ assert "Original error" in str(exc_info.value)
+ print("✓ Original error not masked by aclose error")
+
+
+# ==================================================================================================
+# Tests for bracket tool calls
+# ==================================================================================================
+
+class TestStreamingOpenaiBracketToolCalls:
+ """Tests for bracket-style tool call handling."""
+
+ @pytest.mark.asyncio
+ async def test_detects_bracket_tool_calls(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Detects bracket-style tool calls in content.
+ Goal: Verify bracket tool call detection.
+ """
+ print("Setup: Mock stream with bracket tool calls...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="[tool_call: func1]")
+
+ bracket_tool_calls = [
+ {"id": "call_1", "type": "function", "function": {"name": "func1", "arguments": "{}"}}
+ ]
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=bracket_tool_calls):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Should have tool_calls chunk
+ tool_chunks = [c for c in chunks if '"tool_calls"' in c]
+ assert len(tool_chunks) >= 1
+ print("✓ Bracket tool calls detected")
+
+ @pytest.mark.asyncio
+ async def test_deduplicates_tool_calls(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Deduplicates tool calls from stream and bracket.
+ Goal: Verify deduplication.
+ """
+ print("Setup: Mock stream with duplicate tool calls...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="text")
+ yield KiroEvent(type="tool_use", tool_use={
+ "id": "call_1", "type": "function",
+ "function": {"name": "func1", "arguments": "{}"}
+ })
+
+ # Same tool call from bracket detection
+ bracket_tool_calls = [
+ {"id": "call_1", "type": "function", "function": {"name": "func1", "arguments": "{}"}}
+ ]
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=bracket_tool_calls):
+ with patch('kiro.streaming_openai.deduplicate_tool_calls') as mock_dedup:
+ mock_dedup.return_value = [
+ {"id": "call_1", "type": "function", "function": {"name": "func1", "arguments": "{}"}}
+ ]
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ # Verify deduplicate was called
+ mock_dedup.assert_called()
+
+ print("✓ Tool calls deduplicated")
+
+
+# ==================================================================================================
+# Tests for metering data
+# ==================================================================================================
+
+class TestStreamingOpenaiMeteringData:
+ """Tests for metering data handling."""
+
+ @pytest.mark.asyncio
+ async def test_includes_credits_used_in_usage(self, mock_http_client, mock_response, mock_model_cache, mock_auth_manager):
+ """
+ What it does: Includes credits_used in usage when metering data present.
+ Goal: Verify metering data is included.
+ """
+ print("Setup: Mock stream with metering data...")
+
+ async def mock_parse_kiro_stream(*args, **kwargs):
+ yield KiroEvent(type="content", content="Hello")
+ yield KiroEvent(type="usage", usage={"credits": 0.001})
+
+ print("Action: Streaming to OpenAI format...")
+ chunks = []
+
+ with patch('kiro.streaming_openai.parse_kiro_stream', mock_parse_kiro_stream):
+ with patch('kiro.streaming_openai.parse_bracket_tool_calls', return_value=[]):
+ async for chunk in stream_kiro_to_openai(
+ mock_http_client, mock_response, "claude-sonnet-4",
+ mock_model_cache, mock_auth_manager
+ ):
+ chunks.append(chunk)
+
+ print(f"Received {len(chunks)} chunks")
+
+ # Final chunk should have credits_used
+ final_chunk = chunks[-2] # Before [DONE]
+ assert '"credits_used"' in final_chunk
+ print("✓ credits_used included in usage")
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_thinking_parser.py b/kiro-gateway/tests/unit/test_thinking_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..9eff0be1d0429e738f23993df2bf9bf2e5e772e2
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_thinking_parser.py
@@ -0,0 +1,992 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for ThinkingParser - FSM-based parser for thinking blocks in streaming responses.
+
+Tests cover:
+- Parser state transitions (PRE_CONTENT -> IN_THINKING -> STREAMING)
+- Tag detection at response start
+- "Cautious" buffering for split tags
+- Different handling modes (as_reasoning_content, remove, pass, strip_tags)
+- Edge cases and error handling
+"""
+
+import pytest
+from unittest.mock import patch
+
+from kiro.thinking_parser import (
+ ThinkingParser,
+ ThinkingParseResult,
+ ParserState,
+)
+
+
+class TestParserStateEnum:
+ """Tests for ParserState enum."""
+
+ def test_pre_content_value(self):
+ """
+ What it does: Verifies PRE_CONTENT enum value.
+ Purpose: Ensure PRE_CONTENT is 0 (initial state).
+ """
+ print("Checking PRE_CONTENT enum value...")
+ assert ParserState.PRE_CONTENT == 0
+
+ def test_in_thinking_value(self):
+ """
+ What it does: Verifies IN_THINKING enum value.
+ Purpose: Ensure IN_THINKING is 1.
+ """
+ print("Checking IN_THINKING enum value...")
+ assert ParserState.IN_THINKING == 1
+
+ def test_streaming_value(self):
+ """
+ What it does: Verifies STREAMING enum value.
+ Purpose: Ensure STREAMING is 2.
+ """
+ print("Checking STREAMING enum value...")
+ assert ParserState.STREAMING == 2
+
+
+class TestThinkingParseResult:
+ """Tests for ThinkingParseResult dataclass."""
+
+ def test_default_values(self):
+ """
+ What it does: Verifies default values of ThinkingParseResult.
+ Purpose: Ensure all fields have correct defaults.
+ """
+ print("Creating ThinkingParseResult with defaults...")
+ result = ThinkingParseResult()
+
+ print(f"Comparing: Expected None, Got {result.thinking_content}")
+ assert result.thinking_content is None
+ assert result.regular_content is None
+ assert result.is_first_thinking_chunk is False
+ assert result.is_last_thinking_chunk is False
+ assert result.state_changed is False
+
+ def test_custom_values(self):
+ """
+ What it does: Verifies custom values in ThinkingParseResult.
+ Purpose: Ensure all fields can be set.
+ """
+ print("Creating ThinkingParseResult with custom values...")
+ result = ThinkingParseResult(
+ thinking_content="thinking",
+ regular_content="regular",
+ is_first_thinking_chunk=True,
+ is_last_thinking_chunk=True,
+ state_changed=True
+ )
+
+ print(f"Comparing thinking_content: Expected 'thinking', Got '{result.thinking_content}'")
+ assert result.thinking_content == "thinking"
+ assert result.regular_content == "regular"
+ assert result.is_first_thinking_chunk is True
+ assert result.is_last_thinking_chunk is True
+ assert result.state_changed is True
+
+
+class TestThinkingParserInitialization:
+ """Tests for ThinkingParser initialization."""
+
+ def test_default_initialization(self):
+ """
+ What it does: Verifies default initialization of ThinkingParser.
+ Purpose: Ensure parser starts in PRE_CONTENT state with empty buffers.
+ """
+ print("Creating ThinkingParser with defaults...")
+ parser = ThinkingParser()
+
+ print(f"Comparing state: Expected PRE_CONTENT, Got {parser.state}")
+ assert parser.state == ParserState.PRE_CONTENT
+ assert parser.initial_buffer == ""
+ assert parser.thinking_buffer == ""
+ assert parser.open_tag is None
+ assert parser.close_tag is None
+ assert parser.is_first_thinking_chunk is True
+ assert parser._thinking_block_found is False
+
+ def test_custom_handling_mode(self):
+ """
+ What it does: Verifies custom handling_mode parameter.
+ Purpose: Ensure handling_mode can be overridden.
+ """
+ print("Creating ThinkingParser with custom handling_mode...")
+ parser = ThinkingParser(handling_mode="remove")
+
+ print(f"Comparing handling_mode: Expected 'remove', Got '{parser.handling_mode}'")
+ assert parser.handling_mode == "remove"
+
+ def test_custom_open_tags(self):
+ """
+ What it does: Verifies custom open_tags parameter.
+ Purpose: Ensure open_tags can be overridden.
+ """
+ print("Creating ThinkingParser with custom open_tags...")
+ custom_tags = ["", ""]
+ parser = ThinkingParser(open_tags=custom_tags)
+
+ print(f"Comparing open_tags: Expected {custom_tags}, Got {parser.open_tags}")
+ assert parser.open_tags == custom_tags
+
+ def test_custom_initial_buffer_size(self):
+ """
+ What it does: Verifies custom initial_buffer_size parameter.
+ Purpose: Ensure initial_buffer_size can be overridden.
+ """
+ print("Creating ThinkingParser with custom initial_buffer_size...")
+ parser = ThinkingParser(initial_buffer_size=50)
+
+ print(f"Comparing initial_buffer_size: Expected 50, Got {parser.initial_buffer_size}")
+ assert parser.initial_buffer_size == 50
+
+ def test_max_tag_length_calculated(self):
+ """
+ What it does: Verifies max_tag_length is calculated from open_tags.
+ Purpose: Ensure cautious buffering uses correct buffer size.
+ """
+ print("Creating ThinkingParser and checking max_tag_length...")
+ parser = ThinkingParser(open_tags=["", ""])
+
+ # max_tag_length = max(len(tag) for tag in open_tags) * 2
+ # len("") = 10, so max_tag_length = 20
+ expected = 20
+ print(f"Comparing max_tag_length: Expected {expected}, Got {parser.max_tag_length}")
+ assert parser.max_tag_length == expected
+
+
+class TestThinkingParserFeedPreContent:
+ """Tests for ThinkingParser.feed() in PRE_CONTENT state."""
+
+ def test_empty_content_returns_empty_result(self):
+ """
+ What it does: Verifies empty content returns empty result.
+ Purpose: Ensure empty string doesn't change state.
+ """
+ print("Feeding empty content...")
+ parser = ThinkingParser()
+ result = parser.feed("")
+
+ print(f"Comparing result: Expected empty result")
+ assert result.thinking_content is None
+ assert result.regular_content is None
+ assert result.state_changed is False
+ assert parser.state == ParserState.PRE_CONTENT
+
+ def test_detects_thinking_tag(self):
+ """
+ What it does: Verifies tag detection.
+ Purpose: Ensure parser transitions to IN_THINKING on tag detection.
+ """
+ print("Feeding content with tag...")
+ parser = ThinkingParser()
+ result = parser.feed("Hello")
+
+ print(f"Comparing state: Expected IN_THINKING, Got {parser.state}")
+ assert parser.state == ParserState.IN_THINKING
+ assert parser.open_tag == ""
+ assert parser.close_tag == ""
+ assert result.state_changed is True
+ assert parser._thinking_block_found is True
+
+ def test_detects_think_tag(self):
+ """
+ What it does: Verifies tag detection.
+ Purpose: Ensure parser detects alternative tag format.
+ """
+ print("Feeding content with tag...")
+ parser = ThinkingParser()
+ result = parser.feed("Hello")
+
+ print(f"Comparing open_tag: Expected '', Got '{parser.open_tag}'")
+ assert parser.state == ParserState.IN_THINKING
+ assert parser.open_tag == ""
+ assert parser.close_tag == ""
+
+ def test_detects_reasoning_tag(self):
+ """
+ What it does: Verifies tag detection.
+ Purpose: Ensure parser detects reasoning tag format.
+ """
+ print("Feeding content with tag...")
+ parser = ThinkingParser()
+ result = parser.feed("Hello")
+
+ print(f"Comparing open_tag: Expected '', Got '{parser.open_tag}'")
+ assert parser.state == ParserState.IN_THINKING
+ assert parser.open_tag == ""
+ assert parser.close_tag == ""
+
+ def test_detects_thought_tag(self):
+ """
+ What it does: Verifies tag detection.
+ Purpose: Ensure parser detects thought tag format.
+ """
+ print("Feeding content with tag...")
+ parser = ThinkingParser()
+ result = parser.feed("Hello")
+
+ print(f"Comparing open_tag: Expected '', Got '{parser.open_tag}'")
+ assert parser.state == ParserState.IN_THINKING
+ assert parser.open_tag == ""
+ assert parser.close_tag == ""
+
+ def test_strips_leading_whitespace_for_tag_detection(self):
+ """
+ What it does: Verifies leading whitespace is stripped for tag detection.
+ Purpose: Ensure tags with leading whitespace are detected.
+ """
+ print("Feeding content with leading whitespace...")
+ parser = ThinkingParser()
+ result = parser.feed(" \n\nHello")
+
+ print(f"Comparing state: Expected IN_THINKING, Got {parser.state}")
+ assert parser.state == ParserState.IN_THINKING
+ assert parser.open_tag == ""
+
+ def test_buffers_partial_tag(self):
+ """
+ What it does: Verifies partial tag is buffered.
+ Purpose: Ensure parser waits for complete tag.
+ """
+ print("Feeding partial tag...")
+ parser = ThinkingParser()
+ result = parser.feed("Hello")
+ print(f"After second chunk: state={parser.state}")
+ assert parser.state == ParserState.IN_THINKING
+ assert parser.open_tag == ""
+
+ def test_no_tag_transitions_to_streaming(self):
+ """
+ What it does: Verifies transition to STREAMING when no tag found.
+ Purpose: Ensure regular content is passed through.
+ """
+ print("Feeding content without thinking tag...")
+ parser = ThinkingParser()
+ result = parser.feed("Hello, this is regular content without any thinking tags.")
+
+ print(f"Comparing state: Expected STREAMING, Got {parser.state}")
+ assert parser.state == ParserState.STREAMING
+ assert result.state_changed is True
+ assert result.regular_content == "Hello, this is regular content without any thinking tags."
+
+ def test_buffer_exceeds_limit_transitions_to_streaming(self):
+ """
+ What it does: Verifies transition to STREAMING when buffer exceeds limit.
+ Purpose: Ensure parser doesn't buffer indefinitely.
+ """
+ print("Feeding content that exceeds buffer limit...")
+ parser = ThinkingParser(initial_buffer_size=10)
+ result = parser.feed("This is a long content that exceeds the buffer limit")
+
+ print(f"Comparing state: Expected STREAMING, Got {parser.state}")
+ assert parser.state == ParserState.STREAMING
+ assert result.state_changed is True
+
+
+class TestThinkingParserFeedInThinking:
+ """Tests for ThinkingParser.feed() in IN_THINKING state."""
+
+ def test_accumulates_thinking_content(self):
+ """
+ What it does: Verifies thinking content is accumulated.
+ Purpose: Ensure content inside thinking block is captured.
+ """
+ print("Feeding thinking content...")
+ parser = ThinkingParser()
+ parser.feed("")
+
+ # Feed more content
+ result = parser.feed("This is thinking content")
+
+ print(f"Comparing thinking_buffer: Got '{parser.thinking_buffer}'")
+ # Content is in buffer due to cautious sending
+ assert "This is thinking content" in parser.thinking_buffer or result.thinking_content
+
+ def test_detects_closing_tag(self):
+ """
+ What it does: Verifies closing tag detection.
+ Purpose: Ensure parser transitions to STREAMING on closing tag.
+ """
+ print("Feeding content with closing tag...")
+ parser = ThinkingParser()
+ parser.feed("Hello")
+ result = parser.feed("World")
+
+ print(f"Comparing state: Expected STREAMING, Got {parser.state}")
+ assert parser.state == ParserState.STREAMING
+ assert result.is_last_thinking_chunk is True
+ assert result.state_changed is True
+
+ def test_regular_content_after_closing_tag(self):
+ """
+ What it does: Verifies regular content after closing tag.
+ Purpose: Ensure content after closing tag is returned as regular_content.
+ """
+ print("Feeding content with closing tag and regular content...")
+ parser = ThinkingParser()
+ parser.feed("Thinking")
+ result = parser.feed("Regular content")
+
+ print(f"Comparing regular_content: Got '{result.regular_content}'")
+ assert result.regular_content == "Regular content"
+
+ def test_strips_whitespace_after_closing_tag(self):
+ """
+ What it does: Verifies whitespace is stripped after closing tag.
+ Purpose: Ensure leading newlines after closing tag are removed.
+ """
+ print("Feeding content with whitespace after closing tag...")
+ parser = ThinkingParser()
+ parser.feed("Thinking")
+ result = parser.feed("\n\nRegular content")
+
+ print(f"Comparing regular_content: Got '{result.regular_content}'")
+ assert result.regular_content == "Regular content"
+
+ def test_cautious_buffering(self):
+ """
+ What it does: Verifies cautious buffering keeps last max_tag_length chars.
+ Purpose: Ensure closing tag is not split across chunks.
+ """
+ print("Testing cautious buffering...")
+ parser = ThinkingParser(open_tags=[""]) # Short tag for easier testing
+ parser.feed("")
+
+ # Feed content longer than max_tag_length
+ long_content = "A" * 50
+ result = parser.feed(long_content)
+
+ print(f"Comparing thinking_buffer length: Got {len(parser.thinking_buffer)}")
+ # Buffer should keep last max_tag_length chars
+ assert len(parser.thinking_buffer) <= parser.max_tag_length
+
+ def test_split_closing_tag(self):
+ """
+ What it does: Verifies split closing tag is handled.
+ Purpose: Ensure closing tag split across chunks is detected.
+ """
+ print("Feeding split closing tag...")
+ parser = ThinkingParser()
+ parser.feed("Hello")
+ parser.feed("World")
+
+ print(f"Comparing state: Expected STREAMING, Got {parser.state}")
+ assert parser.state == ParserState.STREAMING
+
+
+class TestThinkingParserFeedStreaming:
+ """Tests for ThinkingParser.feed() in STREAMING state."""
+
+ def test_passes_content_through(self):
+ """
+ What it does: Verifies content is passed through in STREAMING state.
+ Purpose: Ensure regular content is returned as-is.
+ """
+ print("Feeding content in STREAMING state...")
+ parser = ThinkingParser()
+ # Transition to STREAMING by feeding non-tag content
+ parser.feed("Regular content")
+
+ result = parser.feed("More content")
+
+ print(f"Comparing regular_content: Expected 'More content', Got '{result.regular_content}'")
+ assert result.regular_content == "More content"
+ assert result.thinking_content is None
+
+ def test_ignores_thinking_tags_in_streaming(self):
+ """
+ What it does: Verifies thinking tags are ignored in STREAMING state.
+ Purpose: Ensure tags after initial detection are passed through.
+ """
+ print("Feeding thinking tag in STREAMING state...")
+ parser = ThinkingParser()
+ parser.feed("Regular content") # Transition to STREAMING
+
+ result = parser.feed("This should be regular")
+
+ print(f"Comparing regular_content: Got '{result.regular_content}'")
+ assert result.regular_content == "This should be regular"
+ assert result.thinking_content is None
+
+
+class TestThinkingParserFinalize:
+ """Tests for ThinkingParser.finalize()."""
+
+ def test_flushes_thinking_buffer(self):
+ """
+ What it does: Verifies thinking buffer is flushed on finalize.
+ Purpose: Ensure remaining thinking content is returned.
+ """
+ print("Finalizing parser with thinking buffer...")
+ parser = ThinkingParser()
+ parser.feed("Incomplete thinking")
+
+ result = parser.finalize()
+
+ print(f"Comparing thinking_content: Got '{result.thinking_content}'")
+ assert result.thinking_content is not None
+ assert result.is_last_thinking_chunk is True
+
+ def test_flushes_initial_buffer(self):
+ """
+ What it does: Verifies initial buffer is flushed on finalize.
+ Purpose: Ensure buffered content is returned when no tag found.
+ """
+ print("Finalizing parser with initial buffer...")
+ parser = ThinkingParser()
+ parser.feed("Content")
+ parser.finalize()
+
+ print(f"Comparing buffers: thinking_buffer='{parser.thinking_buffer}', initial_buffer='{parser.initial_buffer}'")
+ assert parser.thinking_buffer == ""
+ assert parser.initial_buffer == ""
+
+
+class TestThinkingParserReset:
+ """Tests for ThinkingParser.reset()."""
+
+ def test_resets_to_initial_state(self):
+ """
+ What it does: Verifies reset returns parser to initial state.
+ Purpose: Ensure parser can be reused.
+ """
+ print("Resetting parser after use...")
+ parser = ThinkingParser()
+ parser.feed("ContentRegular")
+
+ parser.reset()
+
+ print(f"Comparing state: Expected PRE_CONTENT, Got {parser.state}")
+ assert parser.state == ParserState.PRE_CONTENT
+ assert parser.initial_buffer == ""
+ assert parser.thinking_buffer == ""
+ assert parser.open_tag is None
+ assert parser.close_tag is None
+ assert parser.is_first_thinking_chunk is True
+ assert parser._thinking_block_found is False
+
+
+class TestThinkingParserFoundThinkingBlock:
+ """Tests for ThinkingParser.found_thinking_block property."""
+
+ def test_false_initially(self):
+ """
+ What it does: Verifies found_thinking_block is False initially.
+ Purpose: Ensure property starts as False.
+ """
+ print("Checking found_thinking_block initially...")
+ parser = ThinkingParser()
+
+ print(f"Comparing: Expected False, Got {parser.found_thinking_block}")
+ assert parser.found_thinking_block is False
+
+ def test_true_after_tag_detection(self):
+ """
+ What it does: Verifies found_thinking_block is True after tag detection.
+ Purpose: Ensure property is set when thinking block is found.
+ """
+ print("Checking found_thinking_block after tag detection...")
+ parser = ThinkingParser()
+ parser.feed("Content")
+
+ print(f"Comparing: Expected True, Got {parser.found_thinking_block}")
+ assert parser.found_thinking_block is True
+
+ def test_false_when_no_tag(self):
+ """
+ What it does: Verifies found_thinking_block is False when no tag found.
+ Purpose: Ensure property stays False for regular content.
+ """
+ print("Checking found_thinking_block with no tag...")
+ parser = ThinkingParser()
+ parser.feed("Regular content without thinking tags")
+
+ print(f"Comparing: Expected False, Got {parser.found_thinking_block}")
+ assert parser.found_thinking_block is False
+
+
+class TestThinkingParserProcessForOutput:
+ """Tests for ThinkingParser.process_for_output()."""
+
+ def test_as_reasoning_content_mode(self):
+ """
+ What it does: Verifies as_reasoning_content mode returns content as-is.
+ Purpose: Ensure content is returned unchanged for reasoning_content field.
+ """
+ print("Testing as_reasoning_content mode...")
+ parser = ThinkingParser(handling_mode="as_reasoning_content")
+ parser.open_tag = ""
+ parser.close_tag = ""
+
+ result = parser.process_for_output("Thinking content", is_first=True, is_last=True)
+
+ print(f"Comparing: Expected 'Thinking content', Got '{result}'")
+ assert result == "Thinking content"
+
+ def test_remove_mode(self):
+ """
+ What it does: Verifies remove mode returns None.
+ Purpose: Ensure thinking content is removed.
+ """
+ print("Testing remove mode...")
+ parser = ThinkingParser(handling_mode="remove")
+
+ result = parser.process_for_output("Thinking content", is_first=True, is_last=True)
+
+ print(f"Comparing: Expected None, Got {result}")
+ assert result is None
+
+ def test_pass_mode_first_chunk(self):
+ """
+ What it does: Verifies pass mode adds opening tag to first chunk.
+ Purpose: Ensure tags are preserved in pass mode.
+ """
+ print("Testing pass mode with first chunk...")
+ parser = ThinkingParser(handling_mode="pass")
+ parser.open_tag = ""
+ parser.close_tag = ""
+
+ result = parser.process_for_output("Content", is_first=True, is_last=False)
+
+ print(f"Comparing: Expected 'Content', Got '{result}'")
+ assert result == "Content"
+
+ def test_pass_mode_last_chunk(self):
+ """
+ What it does: Verifies pass mode adds closing tag to last chunk.
+ Purpose: Ensure closing tag is added in pass mode.
+ """
+ print("Testing pass mode with last chunk...")
+ parser = ThinkingParser(handling_mode="pass")
+ parser.open_tag = ""
+ parser.close_tag = ""
+
+ result = parser.process_for_output("Content", is_first=False, is_last=True)
+
+ print(f"Comparing: Expected 'Content', Got '{result}'")
+ assert result == "Content"
+
+ def test_pass_mode_first_and_last_chunk(self):
+ """
+ What it does: Verifies pass mode adds both tags when first and last.
+ Purpose: Ensure both tags are added for single chunk.
+ """
+ print("Testing pass mode with first and last chunk...")
+ parser = ThinkingParser(handling_mode="pass")
+ parser.open_tag = ""
+ parser.close_tag = ""
+
+ result = parser.process_for_output("Content", is_first=True, is_last=True)
+
+ print(f"Comparing: Expected 'Content', Got '{result}'")
+ assert result == "Content"
+
+ def test_pass_mode_middle_chunk(self):
+ """
+ What it does: Verifies pass mode returns content as-is for middle chunk.
+ Purpose: Ensure no tags are added for middle chunks.
+ """
+ print("Testing pass mode with middle chunk...")
+ parser = ThinkingParser(handling_mode="pass")
+ parser.open_tag = ""
+ parser.close_tag = ""
+
+ result = parser.process_for_output("Content", is_first=False, is_last=False)
+
+ print(f"Comparing: Expected 'Content', Got '{result}'")
+ assert result == "Content"
+
+ def test_strip_tags_mode(self):
+ """
+ What it does: Verifies strip_tags mode returns content without tags.
+ Purpose: Ensure content is returned without tags.
+ """
+ print("Testing strip_tags mode...")
+ parser = ThinkingParser(handling_mode="strip_tags")
+
+ result = parser.process_for_output("Thinking content", is_first=True, is_last=True)
+
+ print(f"Comparing: Expected 'Thinking content', Got '{result}'")
+ assert result == "Thinking content"
+
+ def test_none_content_returns_none(self):
+ """
+ What it does: Verifies None content returns None.
+ Purpose: Ensure None is handled correctly.
+ """
+ print("Testing None content...")
+ parser = ThinkingParser()
+
+ result = parser.process_for_output(None, is_first=True, is_last=True)
+
+ print(f"Comparing: Expected None, Got {result}")
+ assert result is None
+
+ def test_empty_content_returns_none(self):
+ """
+ What it does: Verifies empty content returns None.
+ Purpose: Ensure empty string is handled correctly.
+ """
+ print("Testing empty content...")
+ parser = ThinkingParser()
+
+ result = parser.process_for_output("", is_first=True, is_last=True)
+
+ print(f"Comparing: Expected None, Got {result}")
+ assert result is None
+
+
+class TestThinkingParserFullFlow:
+ """Integration tests for full parsing flow."""
+
+ def test_complete_thinking_block(self):
+ """
+ What it does: Verifies complete thinking block parsing.
+ Purpose: Ensure full flow works correctly.
+ """
+ print("Testing complete thinking block flow...")
+ parser = ThinkingParser()
+
+ # Feed complete thinking block
+ result1 = parser.feed("This is my reasoning process.Here is the answer.")
+
+ print(f"State: {parser.state}")
+ print(f"Thinking content: {result1.thinking_content}")
+ print(f"Regular content: {result1.regular_content}")
+
+ assert parser.state == ParserState.STREAMING
+ assert parser.found_thinking_block is True
+ assert result1.regular_content == "Here is the answer."
+
+ def test_multi_chunk_thinking_block(self):
+ """
+ What it does: Verifies thinking block split across multiple chunks.
+ Purpose: Ensure chunked content is handled correctly.
+ """
+ print("Testing multi-chunk thinking block...")
+ parser = ThinkingParser()
+
+ # Feed in multiple chunks
+ result1 = parser.feed("Let me think")
+ print(f"After chunk 2: state={parser.state}")
+ assert parser.state == ParserState.IN_THINKING
+
+ result3 = parser.feed(" about this...The answer is 42.")
+ print(f"After chunk 4: state={parser.state}")
+ assert parser.state == ParserState.STREAMING
+ assert result4.regular_content == "The answer is 42."
+
+ def test_no_thinking_block(self):
+ """
+ What it does: Verifies handling of content without thinking block.
+ Purpose: Ensure regular content passes through unchanged.
+ """
+ print("Testing content without thinking block...")
+ parser = ThinkingParser()
+
+ result = parser.feed("This is just regular content without any thinking tags.")
+
+ print(f"State: {parser.state}")
+ print(f"Regular content: {result.regular_content}")
+
+ assert parser.state == ParserState.STREAMING
+ assert parser.found_thinking_block is False
+ assert result.regular_content == "This is just regular content without any thinking tags."
+
+ def test_thinking_block_with_newlines(self):
+ """
+ What it does: Verifies thinking block with newlines after closing tag.
+ Purpose: Ensure newlines are stripped from regular content.
+ """
+ print("Testing thinking block with newlines...")
+ parser = ThinkingParser()
+
+ result = parser.feed("Reasoning\n\n\nAnswer here")
+
+ print(f"Regular content: '{result.regular_content}'")
+ assert result.regular_content == "Answer here"
+
+ def test_empty_thinking_block(self):
+ """
+ What it does: Verifies empty thinking block handling.
+ Purpose: Ensure empty thinking block doesn't break parser.
+ """
+ print("Testing empty thinking block...")
+ parser = ThinkingParser()
+
+ result = parser.feed("Answer")
+
+ print(f"State: {parser.state}")
+ print(f"Regular content: '{result.regular_content}'")
+ assert parser.state == ParserState.STREAMING
+ assert result.regular_content == "Answer"
+
+ def test_thinking_block_only_whitespace_after(self):
+ """
+ What it does: Verifies thinking block with only whitespace after closing tag.
+ Purpose: Ensure whitespace-only content after tag returns None.
+ """
+ print("Testing thinking block with only whitespace after...")
+ parser = ThinkingParser()
+
+ result = parser.feed("Reasoning \n\n ")
+
+ print(f"Regular content: {result.regular_content}")
+ # Whitespace-only content should be stripped to None
+ assert result.regular_content is None or result.regular_content == ""
+
+
+class TestThinkingParserEdgeCases:
+ """Edge case tests for ThinkingParser."""
+
+ def test_nested_tags_not_supported(self):
+ """
+ What it does: Verifies nested tags are not specially handled.
+ Purpose: Ensure nested tags are treated as content.
+ """
+ print("Testing nested tags...")
+ parser = ThinkingParser()
+
+ result = parser.feed("OuterInnerStill outerAnswer")
+
+ print(f"State: {parser.state}")
+ # First closes the block
+ assert parser.state == ParserState.STREAMING
+
+ def test_tag_in_middle_of_content(self):
+ """
+ What it does: Verifies tag in middle of content is not detected.
+ Purpose: Ensure tags are only detected at start.
+ """
+ print("Testing tag in middle of content...")
+ parser = ThinkingParser()
+
+ result = parser.feed("Some text This is not a thinking block")
+
+ print(f"State: {parser.state}")
+ print(f"Regular content: '{result.regular_content}'")
+ assert parser.state == ParserState.STREAMING
+ assert parser.found_thinking_block is False
+ assert "" in result.regular_content
+
+ def test_malformed_closing_tag(self):
+ """
+ What it does: Verifies malformed closing tag is not detected.
+ Purpose: Ensure only exact closing tag is matched.
+ """
+ print("Testing malformed closing tag...")
+ parser = ThinkingParser()
+
+ parser.feed("Content")
+ result = parser.feed("More content") # Wrong case
+
+ print(f"State: {parser.state}")
+ # Should still be in thinking state
+ assert parser.state == ParserState.IN_THINKING
+
+ def test_unicode_content(self):
+ """
+ What it does: Verifies Unicode content is handled correctly.
+ Purpose: Ensure non-ASCII characters work.
+ """
+ print("Testing Unicode content...")
+ parser = ThinkingParser()
+
+ result = parser.feed("Думаю о проблеме 🤔Ответ: 42")
+
+ print(f"Regular content: '{result.regular_content}'")
+ assert parser.state == ParserState.STREAMING
+ assert result.regular_content == "Ответ: 42"
+
+ def test_very_long_thinking_content(self):
+ """
+ What it does: Verifies very long thinking content is handled.
+ Purpose: Ensure large content doesn't break parser.
+ """
+ print("Testing very long thinking content...")
+ parser = ThinkingParser()
+
+ long_content = "A" * 10000
+ result = parser.feed(f"{long_content}Done")
+
+ print(f"State: {parser.state}")
+ assert parser.state == ParserState.STREAMING
+ assert result.regular_content == "Done"
+
+ def test_special_characters_in_content(self):
+ """
+ What it does: Verifies special characters are handled.
+ Purpose: Ensure HTML-like content doesn't break parser.
+ """
+ print("Testing special characters...")
+ parser = ThinkingParser()
+
+ result = parser.feed("Content with bold and & entitiesAnswer")
+
+ print(f"State: {parser.state}")
+ assert parser.state == ParserState.STREAMING
+ assert result.regular_content == "Answer"
+
+ def test_multiple_feeds_after_streaming(self):
+ """
+ What it does: Verifies multiple feeds in STREAMING state.
+ Purpose: Ensure parser continues to work after transition.
+ """
+ print("Testing multiple feeds in STREAMING state...")
+ parser = ThinkingParser()
+
+ parser.feed("ThinkingFirst")
+ result2 = parser.feed(" Second")
+ result3 = parser.feed(" Third")
+
+ print(f"Result 2: '{result2.regular_content}'")
+ print(f"Result 3: '{result3.regular_content}'")
+ assert result2.regular_content == " Second"
+ assert result3.regular_content == " Third"
+
+
+class TestThinkingParserConfigIntegration:
+ """Tests for ThinkingParser integration with config."""
+
+ def test_uses_config_handling_mode(self):
+ """
+ What it does: Verifies parser uses FAKE_REASONING_HANDLING from config.
+ Purpose: Ensure config integration works.
+ """
+ print("Testing config handling mode...")
+ with patch('kiro.thinking_parser.FAKE_REASONING_HANDLING', 'remove'):
+ parser = ThinkingParser()
+
+ print(f"Handling mode: {parser.handling_mode}")
+ assert parser.handling_mode == "remove"
+
+ def test_uses_config_open_tags(self):
+ """
+ What it does: Verifies parser uses FAKE_REASONING_OPEN_TAGS from config.
+ Purpose: Ensure config integration works.
+ """
+ print("Testing config open tags...")
+ custom_tags = [""]
+ with patch('kiro.thinking_parser.FAKE_REASONING_OPEN_TAGS', custom_tags):
+ parser = ThinkingParser()
+
+ print(f"Open tags: {parser.open_tags}")
+ assert parser.open_tags == custom_tags
+
+ def test_default_initial_buffer_size_from_config(self):
+ """
+ What it does: Verifies parser uses default initial_buffer_size from config.
+ Purpose: Ensure config value is used when not overridden.
+
+ Note: We can't easily patch the config value after import, so we just
+ verify the default is used. Custom values are tested in
+ TestThinkingParserInitialization.test_custom_initial_buffer_size.
+ """
+ print("Testing default initial buffer size from config...")
+ from kiro.config import FAKE_REASONING_INITIAL_BUFFER_SIZE
+
+ parser = ThinkingParser()
+
+ print(f"Initial buffer size: {parser.initial_buffer_size}")
+ print(f"Config value: {FAKE_REASONING_INITIAL_BUFFER_SIZE}")
+ assert parser.initial_buffer_size == FAKE_REASONING_INITIAL_BUFFER_SIZE
+
+
+class TestInjectThinkingTags:
+ """Tests for inject_thinking_tags function in converters."""
+
+ def test_injects_tags_when_enabled(self):
+ """
+ What it does: Verifies tags are injected when FAKE_REASONING_ENABLED is True.
+ Purpose: Ensure tags are added to content.
+ """
+ print("Testing tag injection when enabled...")
+ from kiro.converters_core import 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("Hello")
+
+ print(f"Result: '{result}'")
+ assert "enabled" in result
+ assert "4000" in result
+ assert "Hello" in result
+
+ def test_no_injection_when_disabled(self):
+ """
+ What it does: Verifies tags are not injected when FAKE_REASONING_ENABLED is False.
+ Purpose: Ensure tags are not added when disabled.
+ """
+ print("Testing no tag injection when disabled...")
+ from kiro.converters_core import inject_thinking_tags
+
+ with patch('kiro.converters_core.FAKE_REASONING_ENABLED', False):
+ result = inject_thinking_tags("Hello")
+
+ print(f"Result: '{result}'")
+ assert result == "Hello"
+ assert "" not in result
+
+ def test_injection_preserves_content(self):
+ """
+ What it does: Verifies original content is preserved after injection.
+ Purpose: Ensure content is not modified.
+ """
+ print("Testing content preservation...")
+ from kiro.converters_core import inject_thinking_tags
+
+ original = "This is my original content with special chars: <>&"
+
+ with patch('kiro.converters_core.FAKE_REASONING_ENABLED', True):
+ with patch('kiro.converters_core.FAKE_REASONING_MAX_TOKENS', 4000):
+ result = inject_thinking_tags(original)
+
+ print(f"Result ends with original: {result.endswith(original)}")
+ assert result.endswith(original)
+
diff --git a/kiro-gateway/tests/unit/test_tokenizer.py b/kiro-gateway/tests/unit/test_tokenizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..a5f7ebeb2bd8e7befe9d6d76a5f42bc92aa406d1
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_tokenizer.py
@@ -0,0 +1,859 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit-тесты для модуля токенизатора (kiro/tokenizer.py).
+
+Проверяет:
+- Подсчёт токенов в тексте (count_tokens)
+- Подсчёт токенов в сообщениях (count_message_tokens)
+- Подсчёт токенов в инструментах (count_tools_tokens)
+- Оценку токенов запроса (estimate_request_tokens)
+- Коэффициент коррекции для Claude (CLAUDE_CORRECTION_FACTOR)
+- Fallback при отсутствии tiktoken
+"""
+
+import pytest
+from unittest.mock import patch, MagicMock
+
+from kiro.tokenizer import (
+ count_tokens,
+ count_message_tokens,
+ count_tools_tokens,
+ estimate_request_tokens,
+ CLAUDE_CORRECTION_FACTOR,
+ _get_encoding
+)
+
+
+class TestCountTokens:
+ """Тесты для функции count_tokens."""
+
+ def test_empty_string_returns_zero(self):
+ """
+ Что он делает: Проверяет, что пустая строка возвращает 0 токенов.
+ Цель: Убедиться в корректной обработке граничного случая.
+ """
+ print("Тест: Пустая строка...")
+ result = count_tokens("")
+ print(f"Результат: {result}")
+ assert result == 0, "Пустая строка должна возвращать 0 токенов"
+
+ def test_none_returns_zero(self):
+ """
+ Что он делает: Проверяет, что None возвращает 0 токенов.
+ Цель: Убедиться в корректной обработке None.
+ """
+ print("Тест: None...")
+ result = count_tokens(None)
+ print(f"Результат: {result}")
+ assert result == 0, "None должен возвращать 0 токенов"
+
+ def test_simple_text_returns_positive(self):
+ """
+ Что он делает: Проверяет, что простой текст возвращает положительное число токенов.
+ Цель: Убедиться в базовой работоспособности подсчёта.
+ """
+ print("Тест: Простой текст...")
+ result = count_tokens("Hello, world!")
+ print(f"Результат: {result}")
+ assert result > 0, "Простой текст должен возвращать положительное число токенов"
+
+ def test_longer_text_returns_more_tokens(self):
+ """
+ Что он делает: Проверяет, что более длинный текст возвращает больше токенов.
+ Цель: Убедиться в корректной пропорциональности подсчёта.
+ """
+ print("Тест: Сравнение длинного и короткого текста...")
+ short_text = "Hello"
+ long_text = "Hello, this is a much longer text that should have more tokens"
+
+ short_tokens = count_tokens(short_text)
+ long_tokens = count_tokens(long_text)
+
+ print(f"Короткий текст: {short_tokens} токенов")
+ print(f"Длинный текст: {long_tokens} токенов")
+
+ assert long_tokens > short_tokens, "Длинный текст должен иметь больше токенов"
+
+ def test_claude_correction_applied_by_default(self):
+ """
+ Что он делает: Проверяет, что коэффициент коррекции Claude применяется по умолчанию.
+ Цель: Убедиться, что apply_claude_correction=True по умолчанию.
+ """
+ print("Тест: Коэффициент коррекции Claude...")
+ text = "This is a test text for token counting"
+
+ with_correction = count_tokens(text, apply_claude_correction=True)
+ without_correction = count_tokens(text, apply_claude_correction=False)
+
+ print(f"С коррекцией: {with_correction}")
+ print(f"Без коррекции: {without_correction}")
+
+ # С коррекцией должно быть больше (коэффициент 1.15)
+ assert with_correction > without_correction, "С коррекцией должно быть больше токенов"
+
+ # Проверяем примерное соотношение
+ ratio = with_correction / without_correction
+ print(f"Соотношение: {ratio}")
+ assert 1.1 <= ratio <= 1.2, f"Соотношение должно быть около {CLAUDE_CORRECTION_FACTOR}"
+
+ def test_without_claude_correction(self):
+ """
+ Что он делает: Проверяет подсчёт без коэффициента коррекции.
+ Цель: Убедиться, что apply_claude_correction=False работает.
+ """
+ print("Тест: Без коэффициента коррекции...")
+ text = "Test text"
+
+ result = count_tokens(text, apply_claude_correction=False)
+ print(f"Результат: {result}")
+
+ assert result > 0, "Должен вернуть положительное число токенов"
+
+ def test_unicode_text(self):
+ """
+ Что он делает: Проверяет подсчёт токенов для Unicode текста.
+ Цель: Убедиться в корректной обработке не-ASCII символов.
+ """
+ print("Тест: Unicode текст...")
+ text = "Привет, мир! 你好世界 🌍"
+
+ result = count_tokens(text)
+ print(f"Результат: {result}")
+
+ assert result > 0, "Unicode текст должен возвращать положительное число токенов"
+
+ def test_multiline_text(self):
+ """
+ Что он делает: Проверяет подсчёт токенов для многострочного текста.
+ Цель: Убедиться в корректной обработке переносов строк.
+ """
+ print("Тест: Многострочный текст...")
+ text = """Line 1
+ Line 2
+ Line 3"""
+
+ result = count_tokens(text)
+ print(f"Результат: {result}")
+
+ assert result > 0, "Многострочный текст должен возвращать положительное число токенов"
+
+ def test_json_text(self):
+ """
+ Что он делает: Проверяет подсчёт токенов для JSON строки.
+ Цель: Убедиться в корректной обработке JSON.
+ """
+ print("Тест: JSON текст...")
+ text = '{"name": "test", "value": 123, "nested": {"key": "value"}}'
+
+ result = count_tokens(text)
+ print(f"Результат: {result}")
+
+ assert result > 0, "JSON текст должен возвращать положительное число токенов"
+
+
+class TestCountTokensFallback:
+ """Тесты для fallback логики при отсутствии tiktoken."""
+
+ def test_fallback_when_tiktoken_unavailable(self):
+ """
+ Что он делает: Проверяет fallback подсчёт когда tiktoken недоступен.
+ Цель: Убедиться, что система работает без tiktoken.
+ """
+ print("Тест: Fallback без tiktoken...")
+
+ # Мокируем _get_encoding чтобы вернуть None
+ with patch('kiro.tokenizer._get_encoding', return_value=None):
+ result = count_tokens("Hello world test")
+ print(f"Результат: {result}")
+
+ # Fallback: len(text) // 4 + 1, затем * 1.15
+ # "Hello world test" = 16 символов
+ # 16 // 4 + 1 = 5
+ # 5 * 1.15 = 5.75 -> 5
+ assert result > 0, "Fallback должен вернуть положительное число"
+
+ def test_fallback_without_correction(self):
+ """
+ Что он делает: Проверяет fallback без коэффициента коррекции.
+ Цель: Убедиться, что fallback работает с apply_claude_correction=False.
+ """
+ print("Тест: Fallback без коррекции...")
+
+ with patch('kiro.tokenizer._get_encoding', return_value=None):
+ result = count_tokens("Test", apply_claude_correction=False)
+ print(f"Результат: {result}")
+
+ # "Test" = 4 символа
+ # 4 // 4 + 1 = 2
+ assert result > 0, "Fallback должен вернуть положительное число"
+
+
+class TestCountMessageTokens:
+ """Тесты для функции count_message_tokens."""
+
+ def test_empty_list_returns_zero(self):
+ """
+ Что он делает: Проверяет, что пустой список возвращает 0 токенов.
+ Цель: Убедиться в корректной обработке пустого списка.
+ """
+ print("Тест: Пустой список сообщений...")
+ result = count_message_tokens([])
+ print(f"Результат: {result}")
+ assert result == 0, "Пустой список должен возвращать 0 токенов"
+
+ def test_none_returns_zero(self):
+ """
+ Что он делает: Проверяет, что None возвращает 0 токенов.
+ Цель: Убедиться в корректной обработке None.
+ """
+ print("Тест: None...")
+ result = count_message_tokens(None)
+ print(f"Результат: {result}")
+ assert result == 0, "None должен возвращать 0 токенов"
+
+ def test_single_user_message(self):
+ """
+ Что он делает: Проверяет подсчёт токенов для одного user сообщения.
+ Цель: Убедиться в базовой работоспособности.
+ """
+ print("Тест: Одно user сообщение...")
+ messages = [{"role": "user", "content": "Hello, AI!"}]
+
+ result = count_message_tokens(messages)
+ print(f"Результат: {result}")
+
+ assert result > 0, "Должен вернуть положительное число токенов"
+
+ def test_multiple_messages(self):
+ """
+ Что он делает: Проверяет подсчёт токенов для нескольких сообщений.
+ Цель: Убедиться, что токены суммируются корректно.
+ """
+ print("Тест: Несколько сообщений...")
+ messages = [
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "Hello!"},
+ {"role": "assistant", "content": "Hi there! How can I help you?"},
+ {"role": "user", "content": "What is the weather?"}
+ ]
+
+ result = count_message_tokens(messages)
+ print(f"Результат: {result}")
+
+ # Больше сообщений = больше токенов
+ single_message = count_message_tokens([messages[0]])
+ assert result > single_message, "Несколько сообщений должны иметь больше токенов"
+
+ def test_message_with_tool_calls(self):
+ """
+ Что он делает: Проверяет подсчёт токенов для сообщения с tool_calls.
+ Цель: Убедиться, что tool_calls учитываются.
+ """
+ print("Тест: Сообщение с tool_calls...")
+ messages = [
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_123",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"location": "Moscow"}'
+ }
+ }
+ ]
+ }
+ ]
+
+ result = count_message_tokens(messages)
+ print(f"Результат: {result}")
+
+ assert result > 0, "Сообщение с tool_calls должно иметь токены"
+
+ def test_message_with_tool_call_id(self):
+ """
+ Что он делает: Проверяет подсчёт токенов для tool response сообщения.
+ Цель: Убедиться, что tool_call_id учитывается.
+ """
+ print("Тест: Tool response сообщение...")
+ messages = [
+ {
+ "role": "tool",
+ "content": "The weather in Moscow is sunny, 25°C",
+ "tool_call_id": "call_123"
+ }
+ ]
+
+ result = count_message_tokens(messages)
+ print(f"Результат: {result}")
+
+ assert result > 0, "Tool response должен иметь токены"
+
+ def test_message_with_list_content(self):
+ """
+ Что он делает: Проверяет подсчёт токенов для мультимодального контента.
+ Цель: Убедиться, что list content обрабатывается.
+ """
+ print("Тест: Мультимодальный контент...")
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "What is in this image?"},
+ {"type": "image_url", "image_url": {"url": "https://example.com/image.jpg"}}
+ ]
+ }
+ ]
+
+ result = count_message_tokens(messages)
+ print(f"Результат: {result}")
+
+ assert result > 0, "Мультимодальный контент должен иметь токены"
+
+ def test_without_claude_correction(self):
+ """
+ Что он делает: Проверяет подсчёт без коэффициента коррекции.
+ Цель: Убедиться, что apply_claude_correction=False работает.
+ """
+ print("Тест: Без коэффициента коррекции...")
+ messages = [{"role": "user", "content": "Test message"}]
+
+ with_correction = count_message_tokens(messages, apply_claude_correction=True)
+ without_correction = count_message_tokens(messages, apply_claude_correction=False)
+
+ print(f"С коррекцией: {with_correction}")
+ print(f"Без коррекции: {without_correction}")
+
+ assert with_correction > without_correction, "С коррекцией должно быть больше"
+
+ def test_message_with_empty_content(self):
+ """
+ Что он делает: Проверяет подсчёт для сообщения с пустым content.
+ Цель: Убедиться, что пустой content не ломает подсчёт.
+ """
+ print("Тест: Пустой content...")
+ messages = [{"role": "user", "content": ""}]
+
+ result = count_message_tokens(messages)
+ print(f"Результат: {result}")
+
+ # Должны быть служебные токены (role, разделители)
+ assert result > 0, "Даже пустое сообщение должно иметь служебные токены"
+
+ def test_message_with_none_content(self):
+ """
+ Что он делает: Проверяет подсчёт для сообщения с None content.
+ Цель: Убедиться, что None content не ломает подсчёт.
+ """
+ print("Тест: None content...")
+ messages = [{"role": "assistant", "content": None}]
+
+ result = count_message_tokens(messages)
+ print(f"Результат: {result}")
+
+ assert result > 0, "Сообщение с None content должно иметь служебные токены"
+
+
+class TestCountToolsTokens:
+ """Тесты для функции count_tools_tokens."""
+
+ def test_none_returns_zero(self):
+ """
+ Что он делает: Проверяет, что None возвращает 0 токенов.
+ Цель: Убедиться в корректной обработке None.
+ """
+ print("Тест: None...")
+ result = count_tools_tokens(None)
+ print(f"Результат: {result}")
+ assert result == 0, "None должен возвращать 0 токенов"
+
+ def test_empty_list_returns_zero(self):
+ """
+ Что он делает: Проверяет, что пустой список возвращает 0 токенов.
+ Цель: Убедиться в корректной обработке пустого списка.
+ """
+ print("Тест: Пустой список...")
+ result = count_tools_tokens([])
+ print(f"Результат: {result}")
+ assert result == 0, "Пустой список должен возвращать 0 токенов"
+
+ def test_single_tool(self):
+ """
+ Что он делает: Проверяет подсчёт токенов для одного инструмента.
+ Цель: Убедиться в базовой работоспособности.
+ """
+ print("Тест: Один инструмент...")
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get the current weather for a location",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string", "description": "City name"}
+ },
+ "required": ["location"]
+ }
+ }
+ }
+ ]
+
+ result = count_tools_tokens(tools)
+ print(f"Результат: {result}")
+
+ assert result > 0, "Инструмент должен иметь токены"
+
+ def test_multiple_tools(self):
+ """
+ Что он делает: Проверяет подсчёт токенов для нескольких инструментов.
+ Цель: Убедиться, что токены суммируются.
+ """
+ print("Тест: Несколько инструментов...")
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get weather",
+ "parameters": {"type": "object", "properties": {}}
+ }
+ },
+ {
+ "type": "function",
+ "function": {
+ "name": "search_web",
+ "description": "Search the web",
+ "parameters": {"type": "object", "properties": {}}
+ }
+ }
+ ]
+
+ result = count_tools_tokens(tools)
+ single_tool = count_tools_tokens([tools[0]])
+
+ print(f"Два инструмента: {result}")
+ print(f"Один инструмент: {single_tool}")
+
+ assert result > single_tool, "Больше инструментов = больше токенов"
+
+ def test_tool_with_complex_parameters(self):
+ """
+ Что он делает: Проверяет подсчёт для инструмента со сложными параметрами.
+ Цель: Убедиться, что JSON schema параметров учитывается.
+ """
+ print("Тест: Сложные параметры...")
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "complex_function",
+ "description": "A function with complex parameters",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "name": {"type": "string", "description": "Name"},
+ "age": {"type": "integer", "description": "Age"},
+ "address": {
+ "type": "object",
+ "properties": {
+ "street": {"type": "string"},
+ "city": {"type": "string"},
+ "country": {"type": "string"}
+ }
+ },
+ "tags": {
+ "type": "array",
+ "items": {"type": "string"}
+ }
+ },
+ "required": ["name", "age"]
+ }
+ }
+ }
+ ]
+
+ result = count_tools_tokens(tools)
+ print(f"Результат: {result}")
+
+ assert result > 0, "Сложный инструмент должен иметь токены"
+
+ def test_tool_without_parameters(self):
+ """
+ Что он делает: Проверяет подсчёт для инструмента без параметров.
+ Цель: Убедиться, что отсутствие parameters не ломает подсчёт.
+ """
+ print("Тест: Без параметров...")
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "no_params_func",
+ "description": "A function without parameters"
+ }
+ }
+ ]
+
+ result = count_tools_tokens(tools)
+ print(f"Результат: {result}")
+
+ assert result > 0, "Инструмент без параметров должен иметь токены"
+
+ def test_tool_with_empty_description(self):
+ """
+ Что он делает: Проверяет подсчёт для инструмента с пустым description.
+ Цель: Убедиться, что пустой description не ломает подсчёт.
+ """
+ print("Тест: Пустой description...")
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "func",
+ "description": "",
+ "parameters": {"type": "object", "properties": {}}
+ }
+ }
+ ]
+
+ result = count_tools_tokens(tools)
+ print(f"Результат: {result}")
+
+ assert result > 0, "Инструмент с пустым description должен иметь токены"
+
+ def test_non_function_tool_type(self):
+ """
+ Что он делает: Проверяет обработку инструмента с type != "function".
+ Цель: Убедиться, что non-function tools обрабатываются.
+ """
+ print("Тест: Non-function tool...")
+ tools = [
+ {
+ "type": "other_type",
+ "some_field": "value"
+ }
+ ]
+
+ result = count_tools_tokens(tools)
+ print(f"Результат: {result}")
+
+ # Должны быть хотя бы служебные токены
+ assert result >= 0, "Non-function tool не должен ломать подсчёт"
+
+ def test_without_claude_correction(self):
+ """
+ Что он делает: Проверяет подсчёт без коэффициента коррекции.
+ Цель: Убедиться, что apply_claude_correction=False работает.
+ """
+ print("Тест: Без коэффициента коррекции...")
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "test_func",
+ "description": "Test function",
+ "parameters": {"type": "object", "properties": {}}
+ }
+ }
+ ]
+
+ with_correction = count_tools_tokens(tools, apply_claude_correction=True)
+ without_correction = count_tools_tokens(tools, apply_claude_correction=False)
+
+ print(f"С коррекцией: {with_correction}")
+ print(f"Без коррекции: {without_correction}")
+
+ assert with_correction > without_correction, "С коррекцией должно быть больше"
+
+
+class TestEstimateRequestTokens:
+ """Тесты для функции estimate_request_tokens."""
+
+ def test_messages_only(self):
+ """
+ Что он делает: Проверяет оценку токенов только для сообщений.
+ Цель: Убедиться в базовой работоспособности.
+ """
+ print("Тест: Только сообщения...")
+ messages = [{"role": "user", "content": "Hello!"}]
+
+ result = estimate_request_tokens(messages)
+ print(f"Результат: {result}")
+
+ assert "messages_tokens" in result
+ assert "tools_tokens" in result
+ assert "system_tokens" in result
+ assert "total_tokens" in result
+
+ assert result["messages_tokens"] > 0
+ assert result["tools_tokens"] == 0
+ assert result["system_tokens"] == 0
+ assert result["total_tokens"] == result["messages_tokens"]
+
+ def test_messages_with_tools(self):
+ """
+ Что он делает: Проверяет оценку токенов для сообщений с инструментами.
+ Цель: Убедиться, что tools учитываются.
+ """
+ print("Тест: Сообщения с инструментами...")
+ messages = [{"role": "user", "content": "What is the weather?"}]
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get weather",
+ "parameters": {"type": "object", "properties": {}}
+ }
+ }
+ ]
+
+ result = estimate_request_tokens(messages, tools=tools)
+ print(f"Результат: {result}")
+
+ assert result["messages_tokens"] > 0
+ assert result["tools_tokens"] > 0
+ assert result["total_tokens"] == result["messages_tokens"] + result["tools_tokens"]
+
+ def test_messages_with_system_prompt(self):
+ """
+ Что он делает: Проверяет оценку токенов с отдельным system prompt.
+ Цель: Убедиться, что system_prompt учитывается.
+ """
+ print("Тест: С system prompt...")
+ messages = [{"role": "user", "content": "Hello!"}]
+ system_prompt = "You are a helpful assistant."
+
+ result = estimate_request_tokens(messages, system_prompt=system_prompt)
+ print(f"Результат: {result}")
+
+ assert result["messages_tokens"] > 0
+ assert result["system_tokens"] > 0
+ assert result["total_tokens"] == result["messages_tokens"] + result["system_tokens"]
+
+ def test_full_request(self):
+ """
+ Что он делает: Проверяет оценку токенов для полного запроса.
+ Цель: Убедиться, что все компоненты суммируются.
+ """
+ print("Тест: Полный запрос...")
+ messages = [
+ {"role": "user", "content": "What is the weather in Moscow?"}
+ ]
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get weather for a location",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string"}
+ }
+ }
+ }
+ }
+ ]
+ system_prompt = "You are a weather assistant."
+
+ result = estimate_request_tokens(messages, tools=tools, system_prompt=system_prompt)
+ print(f"Результат: {result}")
+
+ expected_total = result["messages_tokens"] + result["tools_tokens"] + result["system_tokens"]
+ assert result["total_tokens"] == expected_total, "Total должен быть суммой компонентов"
+
+ def test_empty_messages(self):
+ """
+ Что он делает: Проверяет оценку для пустого списка сообщений.
+ Цель: Убедиться в корректной обработке граничного случая.
+ """
+ print("Тест: Пустые сообщения...")
+ result = estimate_request_tokens([])
+ print(f"Результат: {result}")
+
+ assert result["messages_tokens"] == 0
+ assert result["total_tokens"] == 0
+
+
+class TestClaudeCorrectionFactor:
+ """Тесты для коэффициента коррекции Claude."""
+
+ def test_correction_factor_value(self):
+ """
+ Что он делает: Проверяет значение коэффициента коррекции.
+ Цель: Убедиться, что коэффициент равен 1.15.
+ """
+ print(f"Коэффициент коррекции: {CLAUDE_CORRECTION_FACTOR}")
+ assert CLAUDE_CORRECTION_FACTOR == 1.15, "Коэффициент должен быть 1.15"
+
+ def test_correction_increases_token_count(self):
+ """
+ Что он делает: Проверяет, что коррекция увеличивает количество токенов.
+ Цель: Убедиться, что коэффициент применяется корректно.
+ """
+ print("Тест: Коррекция увеличивает токены...")
+ text = "This is a test text for checking the correction factor"
+
+ with_correction = count_tokens(text, apply_claude_correction=True)
+ without_correction = count_tokens(text, apply_claude_correction=False)
+
+ print(f"С коррекцией: {with_correction}")
+ print(f"Без коррекции: {without_correction}")
+
+ assert with_correction > without_correction
+
+ # Проверяем, что разница примерно 15%
+ increase_percent = (with_correction - without_correction) / without_correction * 100
+ print(f"Увеличение: {increase_percent:.1f}%")
+
+ # Допускаем погрешность из-за округления
+ assert 10 <= increase_percent <= 20, "Увеличение должно быть около 15%"
+class TestGetEncoding:
+ """Тесты для функции _get_encoding."""
+
+ def test_returns_encoding_when_tiktoken_available(self):
+ """
+ Что он делает: Проверяет, что _get_encoding возвращает encoding когда tiktoken доступен.
+ Цель: Убедиться в корректной инициализации tiktoken.
+ """
+ print("Тест: tiktoken доступен...")
+
+ # Сбрасываем глобальную переменную для чистого теста
+ import kiro.tokenizer as tokenizer_module
+ original_encoding = tokenizer_module._encoding
+ tokenizer_module._encoding = None
+
+ try:
+ encoding = _get_encoding()
+ print(f"Encoding: {encoding}")
+
+ # Если tiktoken установлен, должен вернуть encoding
+ if encoding is not None:
+ assert hasattr(encoding, 'encode'), "Encoding должен иметь метод encode"
+ finally:
+ # Восстанавливаем
+ tokenizer_module._encoding = original_encoding
+
+ def test_caches_encoding(self):
+ """
+ Что он делает: Проверяет, что encoding кэшируется.
+ Цель: Убедиться в ленивой инициализации.
+ """
+ print("Тест: Кэширование encoding...")
+
+ encoding1 = _get_encoding()
+ encoding2 = _get_encoding()
+
+ print(f"Encoding 1: {encoding1}")
+ print(f"Encoding 2: {encoding2}")
+
+ # Должен вернуть тот же объект
+ assert encoding1 is encoding2, "Encoding должен кэшироваться"
+
+ def test_handles_import_error(self):
+ """
+ Что он делает: Проверяет обработку ImportError при отсутствии tiktoken.
+ Цель: Убедиться, что система работает без tiktoken.
+ """
+ print("Тест: ImportError...")
+
+ import kiro.tokenizer as tokenizer_module
+ original_encoding = tokenizer_module._encoding
+ tokenizer_module._encoding = None
+
+ try:
+ # Мокируем import tiktoken чтобы выбросить ImportError
+ with patch.dict('sys.modules', {'tiktoken': None}):
+ with patch('builtins.__import__', side_effect=ImportError("No module named 'tiktoken'")):
+ # Сбрасываем кэш
+ tokenizer_module._encoding = None
+
+ # Должен вернуть None и не упасть
+ # Примечание: из-за кэширования этот тест может не работать идеально
+ # но главное - проверить что код не падает
+ pass
+ finally:
+ tokenizer_module._encoding = original_encoding
+
+
+class TestTokenizerIntegration:
+ """Интеграционные тесты для токенизатора."""
+
+ def test_realistic_chat_request(self):
+ """
+ Что он делает: Проверяет подсчёт токенов для реалистичного chat запроса.
+ Цель: Убедиться в корректной работе на реальных данных.
+ """
+ print("Тест: Реалистичный chat запрос...")
+
+ messages = [
+ {"role": "system", "content": "You are a helpful AI assistant. Be concise and accurate."},
+ {"role": "user", "content": "What is the capital of France?"},
+ {"role": "assistant", "content": "The capital of France is Paris."},
+ {"role": "user", "content": "What is its population?"}
+ ]
+
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "search_web",
+ "description": "Search the web for information",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "query": {"type": "string", "description": "Search query"}
+ },
+ "required": ["query"]
+ }
+ }
+ }
+ ]
+
+ result = estimate_request_tokens(messages, tools=tools)
+ print(f"Результат: {result}")
+
+ # Проверяем разумность значений
+ assert result["messages_tokens"] > 50, "Сообщения должны иметь > 50 токенов"
+ assert result["tools_tokens"] > 20, "Tools должны иметь > 20 токенов"
+ assert result["total_tokens"] > 70, "Total должен быть > 70 токенов"
+
+ def test_large_context(self):
+ """
+ Что он делает: Проверяет подсчёт токенов для большого контекста.
+ Цель: Убедиться в производительности на больших данных.
+ """
+ print("Тест: Большой контекст...")
+
+ # Создаём большой текст
+ large_text = "This is a test sentence. " * 1000 # ~5000 слов
+
+ messages = [{"role": "user", "content": large_text}]
+
+ result = estimate_request_tokens(messages)
+ print(f"Токенов в большом тексте: {result['total_tokens']}")
+
+ # Должно быть много токенов
+ assert result["total_tokens"] > 1000, "Большой текст должен иметь > 1000 токенов"
+
+ def test_consistency_across_calls(self):
+ """
+ Что он делает: Проверяет консистентность подсчёта при повторных вызовах.
+ Цель: Убедиться, что результаты детерминированы.
+ """
+ print("Тест: Консистентность...")
+
+ text = "This is a test for consistency checking"
+
+ results = [count_tokens(text) for _ in range(5)]
+ print(f"Результаты: {results}")
+
+ # Все результаты должны быть одинаковыми
+ assert len(set(results)) == 1, "Результаты должны быть консистентными"
+
+
\ No newline at end of file
diff --git a/kiro-gateway/tests/unit/test_vpn_proxy.py b/kiro-gateway/tests/unit/test_vpn_proxy.py
new file mode 100644
index 0000000000000000000000000000000000000000..36e2d06bbbaef5bb11eabf72dcef16e68d077be6
--- /dev/null
+++ b/kiro-gateway/tests/unit/test_vpn_proxy.py
@@ -0,0 +1,310 @@
+# -*- coding: utf-8 -*-
+
+"""
+Unit tests for VPN/Proxy configuration logic.
+
+Tests verify that proxy environment variables are set correctly
+for different input formats and scenarios.
+"""
+
+import os
+import pytest
+
+
+@pytest.mark.parametrize(
+ "test_id, initial_no_proxy, vpn_url, expected_http_proxy, expected_https_proxy, expected_no_proxy",
+ [
+ (
+ "proxy_with_http_scheme",
+ None,
+ "http://192.168.1.103:2080",
+ "http://192.168.1.103:2080",
+ "http://192.168.1.103:2080",
+ "127.0.0.1,localhost"
+ ),
+ (
+ "proxy_with_socks5_scheme",
+ None,
+ "socks5://192.168.1.103:1080",
+ "socks5://192.168.1.103:1080",
+ "socks5://192.168.1.103:1080",
+ "127.0.0.1,localhost"
+ ),
+ (
+ "proxy_without_scheme",
+ None,
+ "192.168.1.103:2080",
+ "http://192.168.1.103:2080",
+ "http://192.168.1.103:2080",
+ "127.0.0.1,localhost"
+ ),
+ (
+ "proxy_with_auth",
+ None,
+ "http://user:pass@192.168.1.103:2080",
+ "http://user:pass@192.168.1.103:2080",
+ "http://user:pass@192.168.1.103:2080",
+ "127.0.0.1,localhost"
+ ),
+ (
+ "proxy_preserves_existing_no_proxy",
+ "internal.corp,*.example.com",
+ "http://192.168.1.103:2080",
+ "http://192.168.1.103:2080",
+ "http://192.168.1.103:2080",
+ "internal.corp,*.example.com,127.0.0.1,localhost"
+ ),
+ (
+ "proxy_empty_url",
+ None,
+ "",
+ None,
+ None,
+ None
+ ),
+ ]
+)
+def test_vpn_proxy_environment_setup(
+ test_id,
+ initial_no_proxy,
+ vpn_url,
+ expected_http_proxy,
+ expected_https_proxy,
+ expected_no_proxy,
+ monkeypatch
+):
+ """
+ Parametrized test for VPN/Proxy setup via environment variables.
+
+ Verifies that:
+ - HTTP_PROXY, HTTPS_PROXY, ALL_PROXY are set correctly
+ - URL normalization works (adds http:// if no scheme)
+ - NO_PROXY includes localhost and preserves existing values
+ - Empty URL doesn't set any proxy variables
+ """
+ print(f"\n--- Running VPN/Proxy test: ID = {test_id} ---")
+
+ # Clear proxy environment variables before test
+ for key in ["HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY"]:
+ monkeypatch.delenv(key, raising=False)
+
+ # Set initial NO_PROXY if specified
+ if initial_no_proxy:
+ monkeypatch.setenv("NO_PROXY", initial_no_proxy)
+ print(f"Initial NO_PROXY: '{initial_no_proxy}'")
+
+ # Simulate VPN_PROXY_URL configuration
+ print(f"VPN_PROXY_URL set to: '{vpn_url}'")
+
+ # Replicate logic from main.py (lines 175-197)
+ if vpn_url:
+ proxy_url_with_scheme = vpn_url if "://" in vpn_url else f"http://{vpn_url}"
+ os.environ['HTTP_PROXY'] = proxy_url_with_scheme
+ os.environ['HTTPS_PROXY'] = proxy_url_with_scheme
+ os.environ['ALL_PROXY'] = proxy_url_with_scheme
+
+ no_proxy_hosts = os.environ.get("NO_PROXY", "")
+ local_hosts = "127.0.0.1,localhost"
+ if no_proxy_hosts:
+ os.environ["NO_PROXY"] = f"{no_proxy_hosts},{local_hosts}"
+ else:
+ os.environ["NO_PROXY"] = local_hosts
+
+ # --- Assertions ---
+ print("\n[Verification]")
+
+ if expected_http_proxy:
+ actual_http_proxy = os.environ.get("HTTP_PROXY")
+ print(f"HTTP_PROXY: Expected '{expected_http_proxy}', Got '{actual_http_proxy}'")
+ assert actual_http_proxy == expected_http_proxy, "HTTP_PROXY mismatch!"
+
+ actual_https_proxy = os.environ.get("HTTPS_PROXY")
+ print(f"HTTPS_PROXY: Expected '{expected_https_proxy}', Got '{actual_https_proxy}'")
+ assert actual_https_proxy == expected_https_proxy, "HTTPS_PROXY mismatch!"
+
+ actual_all_proxy = os.environ.get("ALL_PROXY")
+ print(f"ALL_PROXY: Expected '{expected_http_proxy}', Got '{actual_all_proxy}'")
+ assert actual_all_proxy == expected_http_proxy, "ALL_PROXY mismatch!"
+ else:
+ # If proxy should not be set
+ assert os.environ.get("HTTP_PROXY") is None, "HTTP_PROXY should be None!"
+ assert os.environ.get("HTTPS_PROXY") is None, "HTTPS_PROXY should be None!"
+ print("Proxy not set (as expected)")
+
+ if expected_no_proxy:
+ actual_no_proxy = os.environ.get("NO_PROXY")
+ print(f"NO_PROXY: Expected '{expected_no_proxy}', Got '{actual_no_proxy}'")
+ assert actual_no_proxy == expected_no_proxy, "NO_PROXY mismatch!"
+
+ print(f"--- Test '{test_id}' passed successfully ---")
+
+
+def test_proxy_scheme_normalization():
+ """
+ Verifies that URLs without scheme are correctly normalized to http://.
+
+ Tests various input formats:
+ - Plain host:port → http://host:port
+ - http:// → unchanged
+ - https:// → unchanged
+ - socks5:// → unchanged
+ """
+ print("\n--- Test: Proxy scheme normalization ---")
+
+ test_cases = [
+ ("192.168.1.100:8080", "http://192.168.1.100:8080"),
+ ("http://192.168.1.100:8080", "http://192.168.1.100:8080"),
+ ("https://192.168.1.100:8080", "https://192.168.1.100:8080"),
+ ("socks5://192.168.1.100:8080", "socks5://192.168.1.100:8080"),
+ ("127.0.0.1:7890", "http://127.0.0.1:7890"),
+ ]
+
+ for input_url, expected_url in test_cases:
+ print(f"\nInput: '{input_url}'")
+
+ # Logic from main.py
+ proxy_url_with_scheme = input_url if "://" in input_url else f"http://{input_url}"
+
+ print(f"Result: '{proxy_url_with_scheme}'")
+ print(f"Expected: '{expected_url}'")
+ assert proxy_url_with_scheme == expected_url, f"Normalization failed for '{input_url}'"
+
+ print("\n--- Test passed: all schemes normalized correctly ---")
+
+
+def test_no_proxy_list_merging(monkeypatch):
+ """
+ Verifies correct merging of existing and new NO_PROXY values.
+
+ Tests:
+ - Empty existing → only localhost
+ - Existing values → preserved and localhost added
+ - Duplicate localhost → acceptable (not a problem)
+ """
+ print("\n--- Test: NO_PROXY list merging ---")
+
+ test_cases = [
+ # (existing, expected_result)
+ ("", "127.0.0.1,localhost"),
+ ("internal.local", "internal.local,127.0.0.1,localhost"),
+ ("192.168.0.0/16,10.0.0.0/8", "192.168.0.0/16,10.0.0.0/8,127.0.0.1,localhost"),
+ ("*.corp.com,localhost", "*.corp.com,localhost,127.0.0.1,localhost"), # Duplicate localhost - OK
+ ]
+
+ for existing_value, expected_result in test_cases:
+ print(f"\nExisting NO_PROXY: '{existing_value}'")
+
+ # Simulate logic
+ if existing_value:
+ monkeypatch.setenv("NO_PROXY", existing_value)
+ else:
+ monkeypatch.delenv("NO_PROXY", raising=False)
+
+ no_proxy_hosts = os.environ.get("NO_PROXY", "")
+ local_hosts = "127.0.0.1,localhost"
+ if no_proxy_hosts:
+ result = f"{no_proxy_hosts},{local_hosts}"
+ else:
+ result = local_hosts
+
+ print(f"Result: '{result}'")
+ print(f"Expected: '{expected_result}'")
+ assert result == expected_result, f"Merging failed for '{existing_value}'"
+
+ print("\n--- Test passed: lists merged correctly ---")
+
+
+def test_proxy_does_not_affect_local_connections(monkeypatch):
+ """
+ Verifies that local addresses (127.0.0.1, localhost) are always in NO_PROXY.
+
+ This ensures that local tests don't go through VPN/proxy,
+ which would be slow and incorrect.
+ """
+ print("\n--- Test: Local addresses excluded from proxy ---")
+
+ # Simulate proxy setup
+ vpn_url = "http://vpn.example.com:8080"
+ os.environ['HTTP_PROXY'] = vpn_url
+ os.environ['HTTPS_PROXY'] = vpn_url
+
+ no_proxy_hosts = os.environ.get("NO_PROXY", "")
+ local_hosts = "127.0.0.1,localhost"
+ if no_proxy_hosts:
+ os.environ["NO_PROXY"] = f"{no_proxy_hosts},{local_hosts}"
+ else:
+ os.environ["NO_PROXY"] = local_hosts
+
+ no_proxy_value = os.environ.get("NO_PROXY")
+ print(f"NO_PROXY set to: '{no_proxy_value}'")
+
+ # Assertions
+ assert "127.0.0.1" in no_proxy_value, "127.0.0.1 must be in NO_PROXY!"
+ assert "localhost" in no_proxy_value, "localhost must be in NO_PROXY!"
+
+ print("✅ Local addresses correctly excluded from proxy")
+ print("--- Test passed ---")
+
+
+def test_proxy_with_special_characters():
+ """
+ Verifies that proxy URLs with special characters in credentials work correctly.
+
+ Tests authentication with:
+ - Special characters in password
+ - URL encoding (if needed)
+ """
+ print("\n--- Test: Proxy with special characters in credentials ---")
+
+ test_cases = [
+ # (input_url, expected_normalized)
+ ("http://user:p@ss@proxy.com:8080", "http://user:p@ss@proxy.com:8080"),
+ ("http://admin:P@ssw0rd!@192.168.1.1:3128", "http://admin:P@ssw0rd!@192.168.1.1:3128"),
+ ("socks5://user123:pass456@localhost:1080", "socks5://user123:pass456@localhost:1080"),
+ ]
+
+ for input_url, expected_url in test_cases:
+ print(f"\nInput: '{input_url}'")
+
+ # Normalization logic (should preserve special chars)
+ proxy_url_with_scheme = input_url if "://" in input_url else f"http://{input_url}"
+
+ print(f"Result: '{proxy_url_with_scheme}'")
+ print(f"Expected: '{expected_url}'")
+ assert proxy_url_with_scheme == expected_url, f"Special chars handling failed for '{input_url}'"
+
+ print("\n--- Test passed: special characters preserved correctly ---")
+
+
+def test_empty_vpn_proxy_url_does_not_set_variables(monkeypatch):
+ """
+ Verifies that empty VPN_PROXY_URL doesn't set any proxy variables.
+
+ This is the default behavior - direct connection without proxy.
+ """
+ print("\n--- Test: Empty VPN_PROXY_URL (direct connection) ---")
+
+ # Clear all proxy variables
+ for key in ["HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY"]:
+ monkeypatch.delenv(key, raising=False)
+
+ # Simulate empty VPN_PROXY_URL
+ vpn_url = ""
+
+ # Logic from main.py - should NOT execute if vpn_url is empty
+ if vpn_url:
+ proxy_url_with_scheme = vpn_url if "://" in vpn_url else f"http://{vpn_url}"
+ os.environ['HTTP_PROXY'] = proxy_url_with_scheme
+ os.environ['HTTPS_PROXY'] = proxy_url_with_scheme
+ os.environ['ALL_PROXY'] = proxy_url_with_scheme
+
+ # Verify no proxy variables are set
+ assert os.environ.get("HTTP_PROXY") is None, "HTTP_PROXY should not be set!"
+ assert os.environ.get("HTTPS_PROXY") is None, "HTTPS_PROXY should not be set!"
+ assert os.environ.get("ALL_PROXY") is None, "ALL_PROXY should not be set!"
+
+ print("✅ No proxy variables set (direct connection)")
+ print("--- Test passed ---")
+
+
+print("VPN/Proxy tests loaded. Will verify proxy setup logic!")
diff --git a/plans/data-pipeline-roadmap.md b/plans/data-pipeline-roadmap.md
new file mode 100644
index 0000000000000000000000000000000000000000..5c61cec835b96ea588aee0669c0e6dc4f68e6503
--- /dev/null
+++ b/plans/data-pipeline-roadmap.md
@@ -0,0 +1,107 @@
+# Data Pipeline & Streaming Services Upgrade Roadmap
+
+## Executive Summary
+This document outlines a technical roadmap to upgrade the data pipelines and streaming services of the CLI Proxy API. The current implementation suffers from significant buffering bottlenecks, excessive I/O operations (double-writes), and scalability limits in log retrieval. The proposed upgrades focus on zero-copy streaming, asynchronous processing, and optimized data persistence.
+
+## Current Architecture Assessment
+
+### 1. Proxy & Streaming Service
+**Current State:**
+- **Buffering:** The proxy buffers the entire response body in memory for non-streaming responses (and gzip streams) to perform content rewriting.
+- **Gzip Handling:** `internal/api/modules/amp/proxy.go` reads the full body into memory to decompress it if `Content-Encoding` is present, defeating the purpose of streaming for compressed upstream responses.
+- **Response Rewriting:** `ResponseRewriter` buffers non-streaming bodies entirely. For SSE, it attempts to parse chunks individually, which is brittle if JSON tokens span across network chunk boundaries.
+
+**Bottlenecks:**
+- High memory pressure during large response payloads.
+- Increased Time-To-First-Byte (TTFB) due to buffering.
+- Potential data corruption in SSE if network chunks split `data:` lines or JSON fields.
+
+### 2. Data Pipeline (Logging)
+**Current State:**
+- **Write Path:** `FileStreamingLogWriter` writes chunks to a temporary file asynchronously. However, `Close()` triggers a synchronous "assembly" phase that reads the temp file back and writes it to the final log file. This results in 2x Disk I/O (Write Temp -> Read Temp -> Write Final).
+- **Read Path:** `LogRepository` scans *all* files in the log directory to build a list or find logs. Reading a specific log involves iterating through lines in memory (`logAccumulator`).
+
+**Bottlenecks:**
+- Double I/O penalty for every logged request.
+- Log retrieval performance degrades linearly (O(N)) with the number of log files.
+- Synchronous blocking on file system operations during request finalization.
+
+---
+
+## Technical Roadmap
+
+### Phase 1: Zero-Buffer Streaming Proxy
+**Goal:** Eliminate memory buffering in the proxy layer to minimize latency and memory footprint.
+
+#### 1.1 Streaming Decompression
+- **Task:** Refactor `proxy.go` to use a streaming `gzip.Reader` (or `brotli`/`zstd` wrappers) that wraps the `http.Response.Body`.
+- **Implementation:** Create a `DecompressingReadCloser` that transparently decompresses as `Read()` is called, rather than pre-reading the whole body.
+- **Benefit:** Constant memory usage regardless of response size.
+
+#### 1.2 Streaming Response Rewriter
+- **Task:** Rewrite `ResponseRewriter` to use a streaming JSON parser (e.g., `json.Decoder` or a token-based replacer) instead of `gjson`/`sjson` on full buffers.
+- **Implementation:**
+ - Create a `TokenReplacingReader` that scans the stream for specific keys (`model`, `modelVersion`) and replaces values on the fly.
+ - Ensure it maintains state across `Read()` calls to handle tokens split across buffer boundaries.
+- **Benefit:** Zero-latency overhead for model name rewriting; safe for large JSON bodies.
+
+### Phase 2: Robust SSE Handling
+**Goal:** Ensure 100% reliability for streaming AI responses (Server-Sent Events).
+
+#### 2.1 Stateful SSE Parser
+- **Task:** Replace the naive line-splitting logic in `response_rewriter.go`.
+- **Implementation:**
+ - Implement a state machine that buffers only incomplete lines.
+ - Process full `data: {...}` lines as they become available.
+ - Handle multi-line JSON data correctly.
+- **Benefit:** Prevents corruption when network packets fragment SSE messages.
+
+### Phase 3: High-Performance Logging Pipeline
+**Goal:** Decouple logging from request latency and reduce I/O.
+
+#### 3.1 Eliminate Double-Writes
+- **Task:** Redesign the log storage format to allow append-only writing without post-request assembly.
+- **Implementation:**
+ - Change log format to a structured line-based JSON (NDJSON) or a format that doesn't require a specific "header-first, body-second" physical layout if possible.
+ - Alternatively, keep the temp file approach but use `sendfile` (via `io.Copy` optimizations) to merge files efficiently, or just move/rename the temp file to the final location if the order can be adjusted.
+- **Recommendation:** Switch to a directory-per-request or a pure append-only log file where request metadata and body chunks are interleaved but tagged with a Request ID. This allows writing directly to the final destination.
+
+#### 3.2 Async Log Persister
+- **Task:** Move file I/O entirely out of the request context.
+- **Implementation:**
+ - A background worker pool receives `LogEntry` objects (metadata, body chunks) via a buffered channel.
+ - Workers handle file opening/writing/closing independently of the HTTP handler.
+- **Benefit:** Zero impact of disk latency on API response times.
+
+### Phase 4: Scalable Data Access
+**Goal:** Make log retrieval instant regardless of history size.
+
+#### 4.1 Indexing Strategy
+- **Task:** Stop scanning all files for listing/searching.
+- **Implementation:**
+ - Maintain a lightweight `index.json` or SQLite DB that tracks: `RequestID`, `Timestamp`, `Path`, `StatusCode`, `Filename`.
+ - Update the index asynchronously when logs are finalized.
+- **Benefit:** O(1) lookup by Request ID; O(log N) lookup by time range.
+
+#### 4.2 Optimized Reader
+- **Task:** Read logs efficiently.
+- **Implementation:**
+ - When tailing logs (`latest`), read the file backwards from the end (using `Seek`) rather than scanning from the start.
+ - Implement pagination for log listing based on the index.
+
+---
+
+## Execution Plan
+
+1. **Step 1 (Critical):** Fix the Proxy buffering. This is the biggest risk for production stability.
+ - Refactor `proxy.go` gzip handling.
+ - Refactor `ResponseRewriter` for streaming JSON.
+
+2. **Step 2 (Reliability):** Fix SSE parsing in `ResponseRewriter`.
+ - Implement stateful line buffering.
+
+3. **Step 3 (Performance):** Optimize Log Writing.
+ - Refactor `RequestLogger` to avoid double-write.
+
+4. **Step 4 (Scalability):** Implement Log Indexing.
+ - Add `LogIndexService` and update `LogRepository` to use it.
diff --git a/plans/phase2-prd-mcp.md b/plans/phase2-prd-mcp.md
new file mode 100644
index 0000000000000000000000000000000000000000..c20f855a22cc85de45bead4263d387b4249ab947
--- /dev/null
+++ b/plans/phase2-prd-mcp.md
@@ -0,0 +1,111 @@
+# Phase 2: Streaming Stabilization & Observability - PRD & MCP
+
+## 1. Product Requirements Document (PRD)
+
+### 1.1 Objective
+The goal of Phase 2 is to harden the "Zero-Buffer Streaming Proxy" architecture implemented in Phase 1. This phase focuses on **reliability**, **observability**, and **error resilience**. We aim to ensure the system can handle high concurrency, network instability, and malformed upstream responses without crashing or leaking resources, while providing deep visibility into the streaming pipeline.
+
+### 1.2 Functional Requirements
+
+#### 1.2.1 Advanced Request Logging
+* **Streaming Support:** The logger must support true streaming for *both* request and response bodies. The current `[]byte` buffer for request bodies must be replaced with a stream-aware interface (`io.Reader`).
+* **Sanitization:** Sensitive data (Thinking blocks, Tool arguments) must be redacted in real-time with zero-latency overhead.
+* **Format:** Logs should optionally support **NDJSON** (Newline Delimited JSON) to facilitate machine parsing and ingestion into observability platforms.
+
+#### 1.2.2 AMP Response Rewriter Resilience
+* **Edge Case Handling:** The rewriter must gracefully handle:
+ * Split JSON tokens across chunk boundaries (already implemented, needs verification).
+ * Invalid or malformed JSON from upstream.
+ * Mixed content types (e.g., error responses sent as plain text instead of SSE).
+* **Fallback Strategy:** If rewriting fails (e.g., parsing error), the proxy must fallback to passing the raw chunk through to avoid disrupting the client, logging the error asynchronously.
+
+#### 1.2.3 Proxy Resilience & Timeouts
+* **Context Management:** Request context must be propagated correctly. Client disconnection must immediately cancel the upstream request to save costs.
+* **Timeouts:** The proxy must enforce explicit timeouts:
+ * **Connect Timeout:** Max 10s.
+ * **Header Timeout:** Max 30s (TTFB).
+ * **Idle Timeout:** Max 60s (for SSE streams).
+
+### 1.3 Non-Functional Requirements
+
+* **Performance:**
+ * **Latency Overhead:** < 5ms added by the proxy layer (decompression + rewriting + logging).
+ * **Memory Usage:** Constant memory usage per request (O(1)), independent of response size. Target: < 64KB overhead per active stream.
+* **Concurrency:** Support 1000+ concurrent streaming connections on a standard instance without OOM.
+* **Observability:** Expose Prometheus metrics for:
+ * Active streams.
+ * Log queue depth.
+ * Sanitization hit rate.
+ * Upstream latency histograms.
+
+---
+
+## 2. Master Control Plan (MCP)
+
+### 2.1 Architecture Review
+The Phase 1 refactor successfully removed full-body buffering from the Response path. However, the **Request path** still buffers the full body in memory (`RequestLogger` interface takes `body []byte`). Additionally, the current implementation lacks comprehensive error handling for network interruptions during streaming and doesn't expose internal metrics.
+
+**Refinement Areas:**
+* **Logging Interface:** Refactor `RequestLogger` to accept `io.Reader` for the request body.
+* **Async Safety:** Ensure the `eventLoop` in the logger handles channel overflows gracefully (drop strategy vs. block strategy).
+* **Transport Configuration:** The `httputil.ReverseProxy` needs a custom `Transport` with tuned timeouts.
+
+### 2.2 Implementation Roadmap
+
+#### Step 1: Logging Interface Refactor (True Zero-Buffer)
+* **Task:** Modify `RequestLogger.LogStreamingRequest` signature.
+ * **From:** `LogStreamingRequest(..., body []byte, ...)`
+ * **To:** `LogStreamingRequest(..., body io.Reader, ...)`
+* **Action:** Update `internal/logging/request_logger.go`.
+* **Action:** Update `internal/api/middleware/request_logging.go` to pass the request body stream directly (using `io.TeeReader` if necessary to log *and* process, though usually we log what we read).
+
+#### Step 2: Proxy Timeout & Transport Hardening
+* **Task:** Configure `httputil.ReverseProxy` with a custom `http.Transport`.
+* **Action:** In `internal/api/modules/amp/proxy.go`, define:
+ ```go
+ Transport: &http.Transport{
+ DialContext: (&net.Dialer{
+ Timeout: 10 * time.Second,
+ KeepAlive: 30 * time.Second,
+ }).DialContext,
+ ResponseHeaderTimeout: 30 * time.Second,
+ IdleConnTimeout: 90 * time.Second,
+ }
+ ```
+* **Task:** Ensure `req.Context()` cancellation propagates upstream.
+
+#### Step 3: Observability Integration
+* **Task:** Instrument the `FileStreamingLogWriter` and `ResponseRewriter`.
+* **Action:** Add counters for `log_events_dropped`, `chunks_processed`, `thinking_blocks_redacted`.
+
+#### Step 4: Testing & Validation
+* **Task:** Add Unit/Integration Tests.
+ * Test `StreamingSanitizer` with random chunk splits (Fuzzing).
+ * Test `ResponseRewriter` with invalid JSON.
+ * Test Memory usage under load (using `pprof`).
+
+### 2.3 Technical Specifications
+
+#### New Logger Interface
+```go
+type RequestLogger interface {
+ // ... existing LogRequest ...
+
+ // LogStreamingRequest now accepts a reader for the request body
+ LogStreamingRequest(ctx context.Context, url, method string, headers map[string][]string, bodyStream io.Reader, requestID string) (StreamingLogWriter, error)
+}
+```
+
+#### Struct Modifications
+**FileStreamingLogWriter:**
+* Add `dropCount atomic.Uint64` to track buffer overflows.
+* Add `metrics MetricsCollector` interface dependency.
+
+**ResponseRewriter:**
+* Add `FallbackMode bool` flag. If `true`, parser errors disable rewriting for the rest of the stream to ensure delivery.
+
+### 2.4 Execution Strategy
+1. **Refactor Logger Interface:** High impact, touches middleware. Do this first.
+2. **Harden Proxy:** Low risk, high value for reliability.
+3. **Add Tests:** Critical for verifying the stability of the complex streaming logic.
+4. **Add Metrics:** Final polish for operations.
diff --git a/plans/refactoring-architecture.md b/plans/refactoring-architecture.md
new file mode 100644
index 0000000000000000000000000000000000000000..ae7446627bfca6fe008351fdcb90140697fb1972
--- /dev/null
+++ b/plans/refactoring-architecture.md
@@ -0,0 +1,292 @@
+# Clean Architecture Refactoring Plan
+
+## Current State Analysis
+
+### API Handlers (`internal/api/handlers/management/`)
+The current handlers have several issues:
+1. **Tight coupling**: Handlers directly manipulate config, auth, and logging
+2. **Mixed concerns**: Business logic mixed with HTTP transport (Gin)
+3. **Inconsistent error handling**: Uses `gin.H{"error": ...}` directly
+4. **No clear separation**: Domain logic embedded in HTTP handlers
+5. **Large files**: `auth_files.go` is ~2200 lines with multiple responsibilities
+
+### Auth Module (`internal/auth/`)
+1. **Minimal interface**: Only `TokenStorage` interface defined
+2. **Provider-specific logic**: Each provider has its own auth implementation
+3. **No unified error types**: Each provider handles errors differently
+
+### Logging Module (`internal/logging/`)
+1. **Good interfaces**: `RequestLogger` and `StreamingLogWriter` already defined
+2. **Mixed transport concerns**: Some Gin-specific code in logging
+3. **Correlation IDs**: Basic request ID support exists but could be enhanced
+
+## Proposed Clean Architecture
+
+### Layer Structure
+
+```
+┌─────────────────────────────────────────────────────────────┐
+│ Transport Layer │
+│ (Gin handlers, middleware, routing) │
+├─────────────────────────────────────────────────────────────┤
+│ Application Layer │
+│ (Use cases, DTOs, request/response mapping) │
+├─────────────────────────────────────────────────────────────┤
+│ Domain Layer │
+│ (Domain services, entities, business rules, interfaces) │
+├─────────────────────────────────────────────────────────────┤
+│ Infrastructure Layer │
+│ (Config persistence, auth storage, file system, HTTP) │
+└─────────────────────────────────────────────────────────────┘
+```
+
+### New Directory Structure
+
+```
+internal/
+├── api/
+│ ├── handlers/ # HTTP transport (thin layer)
+│ │ └── management/
+│ │ ├── handler.go
+│ │ ├── config_handler.go
+│ │ ├── auth_handler.go
+│ │ └── logs_handler.go
+│ └── middleware/ # HTTP middleware
+│ ├── auth.go
+│ ├── logging.go
+│ └── error_handler.go
+├── application/ # Application layer (NEW)
+│ ├── dto/ # Data transfer objects
+│ │ ├── config_dto.go
+│ │ ├── auth_dto.go
+│ │ └── response_dto.go
+│ ├── mapper/ # DTO <-> Domain mapping
+│ │ ├── config_mapper.go
+│ │ └── auth_mapper.go
+│ └── usecase/ # Use cases
+│ ├── config_usecase.go
+│ ├── auth_usecase.go
+│ └── logs_usecase.go
+├── domain/ # Domain layer (NEW)
+│ ├── entity/ # Domain entities
+│ │ ├── config.go
+│ │ ├── auth.go
+│ │ └── log_entry.go
+│ ├── service/ # Domain services (business logic)
+│ │ ├── config_service.go
+│ │ ├── auth_service.go
+│ │ └── log_service.go
+│ ├── repository/ # Repository interfaces
+│ │ ├── config_repository.go
+│ │ └── auth_repository.go
+│ └── error/ # Domain errors
+│ └── errors.go
+├── infrastructure/ # Infrastructure layer (NEW)
+│ ├── persistence/ # Repository implementations
+│ │ ├── config_repository.go
+│ │ └── auth_repository.go
+│ ├── auth/ # Auth provider implementations
+│ │ ├── provider.go
+│ │ └── factory.go
+│ └── logging/ # Logging infrastructure
+│ ├── file_logger.go
+│ └── structured_logger.go
+└── interfaces/ # Existing - shared interfaces
+ └── types.go
+```
+
+## Key Components
+
+### 1. Domain Errors (`internal/domain/error/errors.go`)
+
+Standardized error types for the entire application:
+
+```go
+package error
+
+type DomainError struct {
+ Code string
+ Message string
+ Cause error
+}
+
+func (e *DomainError) Error() string { ... }
+
+// Specific error types
+var (
+ ErrNotFound = &DomainError{Code: "NOT_FOUND", ...}
+ ErrUnauthorized = &DomainError{Code: "UNAUTHORIZED", ...}
+ ErrInvalidInput = &DomainError{Code: "INVALID_INPUT", ...}
+ ErrInternal = &DomainError{Code: "INTERNAL_ERROR", ...}
+)
+```
+
+### 2. Domain Services
+
+Domain services contain pure business logic, no HTTP concerns:
+
+```go
+// internal/domain/service/config_service.go
+type ConfigService interface {
+ GetConfig(ctx context.Context) (*entity.Config, error)
+ UpdateConfig(ctx context.Context, cfg *entity.Config) error
+ UpdateField(ctx context.Context, field string, value any) error
+ ValidateConfig(ctx context.Context, cfg *entity.Config) error
+}
+
+// internal/domain/service/auth_service.go
+type AuthService interface {
+ ListAuthFiles(ctx context.Context) ([]*entity.AuthFile, error)
+ UploadAuthFile(ctx context.Context, file *entity.AuthFile) error
+ DeleteAuthFile(ctx context.Context, id string) error
+ RefreshToken(ctx context.Context, id string) (*entity.AuthToken, error)
+}
+```
+
+### 3. Repository Interfaces
+
+```go
+// internal/domain/repository/config_repository.go
+type ConfigRepository interface {
+ Load(ctx context.Context) (*entity.Config, error)
+ Save(ctx context.Context, cfg *entity.Config) error
+ Validate(ctx context.Context, cfg *entity.Config) error
+}
+
+// internal/domain/repository/auth_repository.go
+type AuthRepository interface {
+ List(ctx context.Context) ([]*entity.AuthFile, error)
+ GetByID(ctx context.Context, id string) (*entity.AuthFile, error)
+ Save(ctx context.Context, file *entity.AuthFile) error
+ Delete(ctx context.Context, id string) error
+}
+```
+
+### 4. Application Use Cases
+
+Use cases orchestrate domain services for specific operations:
+
+```go
+// internal/application/usecase/config_usecase.go
+type ConfigUseCase struct {
+ configService domain.ConfigService
+ logger logging.Logger
+}
+
+func (uc *ConfigUseCase) GetConfig(ctx context.Context) (*dto.ConfigResponse, error) {
+ cfg, err := uc.configService.GetConfig(ctx)
+ if err != nil {
+ uc.logger.Error("failed to get config", err)
+ return nil, err
+ }
+ return mapper.ToConfigResponse(cfg), nil
+}
+```
+
+### 5. Structured Logging with Correlation IDs
+
+```go
+// internal/infrastructure/logging/structured_logger.go
+type StructuredLogger struct {
+ logger *logrus.Logger
+}
+
+type LogEntry struct {
+ Timestamp time.Time
+ Level string
+ Message string
+ CorrelationID string
+ Service string
+ Operation string
+ Fields map[string]interface{}
+}
+
+func (l *StructuredLogger) WithCorrelationID(ctx context.Context) *logrus.Entry {
+ correlationID := logging.GetRequestID(ctx)
+ return l.logger.WithField("correlation_id", correlationID)
+}
+```
+
+### 6. HTTP Handlers (Thin Layer)
+
+Handlers only deal with HTTP concerns:
+
+```go
+// internal/api/handlers/management/config_handler.go
+type ConfigHandler struct {
+ useCase *usecase.ConfigUseCase
+}
+
+func (h *ConfigHandler) GetConfig(c *gin.Context) {
+ ctx := c.Request.Context()
+ response, err := h.useCase.GetConfig(ctx)
+ if err != nil {
+ h.handleError(c, err)
+ return
+ }
+ c.JSON(http.StatusOK, response)
+}
+
+func (h *ConfigHandler) handleError(c *gin.Context, err error) {
+ // Map domain errors to HTTP responses
+ var domainErr *domainerror.DomainError
+ if errors.As(err, &domainErr) {
+ status := h.mapErrorCodeToStatus(domainErr.Code)
+ c.JSON(status, gin.H{
+ "error": domainErr.Code,
+ "message": domainErr.Message,
+ })
+ return
+ }
+ c.JSON(http.StatusInternalServerError, gin.H{
+ "error": "INTERNAL_ERROR",
+ "message": "An unexpected error occurred",
+ })
+}
+```
+
+## Migration Strategy
+
+### Phase 1: Foundation
+1. Create domain error types
+2. Define repository interfaces
+3. Create domain service interfaces
+4. Set up structured logging infrastructure
+
+### Phase 2: Extract Domain Logic
+1. Move business logic from handlers to domain services
+2. Implement repository interfaces
+3. Create use cases
+4. Add comprehensive unit tests
+
+### Phase 3: Refactor Handlers
+1. Make handlers thin - delegate to use cases
+2. Standardize error handling
+3. Add correlation ID middleware
+
+### Phase 4: Cleanup
+1. Remove old code
+2. Verify backward compatibility
+3. Update documentation
+
+## Backward Compatibility
+
+- All existing API endpoints remain unchanged
+- Response formats preserved
+- Configuration file format unchanged
+- Auth file format unchanged
+
+## Testing Strategy
+
+1. **Unit tests**: Domain services with mocked repositories
+2. **Integration tests**: Use cases with real repositories
+3. **E2E tests**: Full HTTP request/response cycle
+4. **Mock implementations**: For all external dependencies
+
+## Benefits
+
+1. **Testability**: Domain logic can be tested without HTTP layer
+2. **Maintainability**: Clear separation of concerns
+3. **Flexibility**: Easy to swap implementations (e.g., different storage)
+4. **Observability**: Structured logging with correlation IDs
+5. **Scalability**: Clear boundaries for future extensions
diff --git a/rest/.next/dev/types/cache-life.d.ts b/rest/.next/dev/types/cache-life.d.ts
new file mode 100644
index 0000000000000000000000000000000000000000..a8c6997e7228736d34a48b9380bef497972c55cc
--- /dev/null
+++ b/rest/.next/dev/types/cache-life.d.ts
@@ -0,0 +1,145 @@
+// Type definitions for Next.js cacheLife configs
+
+declare module 'next/cache' {
+ export { unstable_cache } from 'next/dist/server/web/spec-extension/unstable-cache'
+ export {
+ updateTag,
+ revalidateTag,
+ revalidatePath,
+ refresh,
+ } from 'next/dist/server/web/spec-extension/revalidate'
+ export { unstable_noStore } from 'next/dist/server/web/spec-extension/unstable-no-store'
+
+
+ /**
+ * Cache this `"use cache"` for a timespan defined by the `"default"` profile.
+ * ```
+ * stale: 300 seconds (5 minutes)
+ * revalidate: 900 seconds (15 minutes)
+ * expire: never
+ * ```
+ *
+ * This cache may be stale on clients for 5 minutes before checking with the server.
+ * If the server receives a new request after 15 minutes, start revalidating new values in the background.
+ * It lives for the maximum age of the server cache. If this entry has no traffic for a while, it may serve an old value the next request.
+ */
+ export function cacheLife(profile: "default"): void
+
+ /**
+ * Cache this `"use cache"` for a timespan defined by the `"seconds"` profile.
+ * ```
+ * stale: 30 seconds
+ * revalidate: 1 seconds
+ * expire: 60 seconds (1 minute)
+ * ```
+ *
+ * This cache may be stale on clients for 30 seconds before checking with the server.
+ * If the server receives a new request after 1 seconds, start revalidating new values in the background.
+ * If this entry has no traffic for 1 minute it will expire. The next request will recompute it.
+ */
+ export function cacheLife(profile: "seconds"): void
+
+ /**
+ * Cache this `"use cache"` for a timespan defined by the `"minutes"` profile.
+ * ```
+ * stale: 300 seconds (5 minutes)
+ * revalidate: 60 seconds (1 minute)
+ * expire: 3600 seconds (1 hour)
+ * ```
+ *
+ * This cache may be stale on clients for 5 minutes before checking with the server.
+ * If the server receives a new request after 1 minute, start revalidating new values in the background.
+ * If this entry has no traffic for 1 hour it will expire. The next request will recompute it.
+ */
+ export function cacheLife(profile: "minutes"): void
+
+ /**
+ * Cache this `"use cache"` for a timespan defined by the `"hours"` profile.
+ * ```
+ * stale: 300 seconds (5 minutes)
+ * revalidate: 3600 seconds (1 hour)
+ * expire: 86400 seconds (1 day)
+ * ```
+ *
+ * This cache may be stale on clients for 5 minutes before checking with the server.
+ * If the server receives a new request after 1 hour, start revalidating new values in the background.
+ * If this entry has no traffic for 1 day it will expire. The next request will recompute it.
+ */
+ export function cacheLife(profile: "hours"): void
+
+ /**
+ * Cache this `"use cache"` for a timespan defined by the `"days"` profile.
+ * ```
+ * stale: 300 seconds (5 minutes)
+ * revalidate: 86400 seconds (1 day)
+ * expire: 604800 seconds (1 week)
+ * ```
+ *
+ * This cache may be stale on clients for 5 minutes before checking with the server.
+ * If the server receives a new request after 1 day, start revalidating new values in the background.
+ * If this entry has no traffic for 1 week it will expire. The next request will recompute it.
+ */
+ export function cacheLife(profile: "days"): void
+
+ /**
+ * Cache this `"use cache"` for a timespan defined by the `"weeks"` profile.
+ * ```
+ * stale: 300 seconds (5 minutes)
+ * revalidate: 604800 seconds (1 week)
+ * expire: 2592000 seconds (1 month)
+ * ```
+ *
+ * This cache may be stale on clients for 5 minutes before checking with the server.
+ * If the server receives a new request after 1 week, start revalidating new values in the background.
+ * If this entry has no traffic for 1 month it will expire. The next request will recompute it.
+ */
+ export function cacheLife(profile: "weeks"): void
+
+ /**
+ * Cache this `"use cache"` for a timespan defined by the `"max"` profile.
+ * ```
+ * stale: 300 seconds (5 minutes)
+ * revalidate: 2592000 seconds (1 month)
+ * expire: 31536000 seconds (365 days)
+ * ```
+ *
+ * This cache may be stale on clients for 5 minutes before checking with the server.
+ * If the server receives a new request after 1 month, start revalidating new values in the background.
+ * If this entry has no traffic for 365 days it will expire. The next request will recompute it.
+ */
+ export function cacheLife(profile: "max"): void
+
+ /**
+ * Cache this `"use cache"` using a custom timespan.
+ * ```
+ * stale: ... // seconds
+ * revalidate: ... // seconds
+ * expire: ... // seconds
+ * ```
+ *
+ * This is similar to Cache-Control: max-age=`stale`,s-max-age=`revalidate`,stale-while-revalidate=`expire-revalidate`
+ *
+ * If a value is left out, the lowest of other cacheLife() calls or the default, is used instead.
+ */
+ export function cacheLife(profile: {
+ /**
+ * This cache may be stale on clients for ... seconds before checking with the server.
+ */
+ stale?: number,
+ /**
+ * If the server receives a new request after ... seconds, start revalidating new values in the background.
+ */
+ revalidate?: number,
+ /**
+ * If this entry has no traffic for ... seconds it will expire. The next request will recompute it.
+ */
+ expire?: number
+ }): void
+
+
+ import { cacheTag } from 'next/dist/server/use-cache/cache-tag'
+ export { cacheTag }
+
+ export const unstable_cacheTag: typeof cacheTag
+ export const unstable_cacheLife: typeof cacheLife
+}
diff --git a/rest/.next/dev/types/routes.d.ts b/rest/.next/dev/types/routes.d.ts
new file mode 100644
index 0000000000000000000000000000000000000000..15617e113dac7d709d06f7a26ada4cdef8f01e33
--- /dev/null
+++ b/rest/.next/dev/types/routes.d.ts
@@ -0,0 +1,55 @@
+// This file is generated automatically by Next.js
+// Do not edit this file manually
+
+type AppRoutes = never
+type PageRoutes = never
+type LayoutRoutes = never
+type RedirectRoutes = never
+type RewriteRoutes = never
+type Routes = AppRoutes | PageRoutes | LayoutRoutes | RedirectRoutes | RewriteRoutes
+
+
+interface ParamMap {
+}
+
+
+export type ParamsOf = ParamMap[Route]
+
+interface LayoutSlotMap {
+}
+
+
+export type { AppRoutes, PageRoutes, LayoutRoutes, RedirectRoutes, RewriteRoutes, ParamMap }
+
+declare global {
+ /**
+ * Props for Next.js App Router page components
+ * @example
+ * ```tsx
+ * export default function Page(props: PageProps<'/blog/[slug]'>) {
+ * const { slug } = await props.params
+ * return Blog post: {slug}
+ * }
+ * ```
+ */
+ interface PageProps {
+ params: Promise
+ searchParams: Promise>
+ }
+
+ /**
+ * Props for Next.js App Router layout components
+ * @example
+ * ```tsx
+ * export default function Layout(props: LayoutProps<'/dashboard'>) {
+ * return {props.children}
+ * }
+ * ```
+ */
+ type LayoutProps = {
+ params: Promise
+ children: React.ReactNode
+ } & {
+ [K in LayoutSlotMap[LayoutRoute]]: React.ReactNode
+ }
+}
diff --git a/rest/.next/dev/types/validator.ts b/rest/.next/dev/types/validator.ts
new file mode 100644
index 0000000000000000000000000000000000000000..000dc8e06e2ae3f3c8354565987a5fa39a86c0d4
--- /dev/null
+++ b/rest/.next/dev/types/validator.ts
@@ -0,0 +1,16 @@
+// This file is generated automatically by Next.js
+// Do not edit this file manually
+// This file validates that all pages and layouts export the correct types
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/sdk/access/errors.go b/sdk/access/errors.go
new file mode 100644
index 0000000000000000000000000000000000000000..6ea2cc1a2b224cf55cf85425b59d7bc0a98916fa
--- /dev/null
+++ b/sdk/access/errors.go
@@ -0,0 +1,12 @@
+package access
+
+import "errors"
+
+var (
+ // ErrNoCredentials indicates no recognizable credentials were supplied.
+ ErrNoCredentials = errors.New("access: no credentials provided")
+ // ErrInvalidCredential signals that supplied credentials were rejected by a provider.
+ ErrInvalidCredential = errors.New("access: invalid credential")
+ // ErrNotHandled tells the manager to continue trying other providers.
+ ErrNotHandled = errors.New("access: not handled")
+)
diff --git a/sdk/access/manager.go b/sdk/access/manager.go
new file mode 100644
index 0000000000000000000000000000000000000000..fb5f8ccab6b317cc3c4d9a7d44b5cd026c790169
--- /dev/null
+++ b/sdk/access/manager.go
@@ -0,0 +1,89 @@
+package access
+
+import (
+ "context"
+ "errors"
+ "net/http"
+ "sync"
+)
+
+// Manager coordinates authentication providers.
+type Manager struct {
+ mu sync.RWMutex
+ providers []Provider
+}
+
+// NewManager constructs an empty manager.
+func NewManager() *Manager {
+ return &Manager{}
+}
+
+// SetProviders replaces the active provider list.
+func (m *Manager) SetProviders(providers []Provider) {
+ if m == nil {
+ return
+ }
+ cloned := make([]Provider, len(providers))
+ copy(cloned, providers)
+ m.mu.Lock()
+ m.providers = cloned
+ m.mu.Unlock()
+}
+
+// Providers returns a snapshot of the active providers.
+func (m *Manager) Providers() []Provider {
+ if m == nil {
+ return nil
+ }
+ m.mu.RLock()
+ defer m.mu.RUnlock()
+ snapshot := make([]Provider, len(m.providers))
+ copy(snapshot, m.providers)
+ return snapshot
+}
+
+// Authenticate evaluates providers until one succeeds.
+func (m *Manager) Authenticate(ctx context.Context, r *http.Request) (*Result, error) {
+ if m == nil {
+ return nil, nil
+ }
+ providers := m.Providers()
+ if len(providers) == 0 {
+ return nil, nil
+ }
+
+ var (
+ missing bool
+ invalid bool
+ )
+
+ for _, provider := range providers {
+ if provider == nil {
+ continue
+ }
+ res, err := provider.Authenticate(ctx, r)
+ if err == nil {
+ return res, nil
+ }
+ if errors.Is(err, ErrNotHandled) {
+ continue
+ }
+ if errors.Is(err, ErrNoCredentials) {
+ missing = true
+ continue
+ }
+ if errors.Is(err, ErrInvalidCredential) {
+ invalid = true
+ continue
+ }
+ return nil, err
+ }
+
+ if invalid {
+ return nil, ErrInvalidCredential
+ }
+ if missing {
+ return nil, ErrNoCredentials
+ }
+ return nil, ErrNoCredentials
+}
diff --git a/sdk/access/registry.go b/sdk/access/registry.go
new file mode 100644
index 0000000000000000000000000000000000000000..a29cdd96b619dc9b5b270e66d4495ebe63d43e50
--- /dev/null
+++ b/sdk/access/registry.go
@@ -0,0 +1,87 @@
+package access
+
+import (
+ "context"
+ "fmt"
+ "net/http"
+ "sync"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
+)
+
+// Provider validates credentials for incoming requests.
+type Provider interface {
+ Identifier() string
+ Authenticate(ctx context.Context, r *http.Request) (*Result, error)
+}
+
+// Result conveys authentication outcome.
+type Result struct {
+ Provider string
+ Principal string
+ Metadata map[string]string
+}
+
+// ProviderFactory builds a provider from configuration data.
+type ProviderFactory func(cfg *config.AccessProvider, root *config.SDKConfig) (Provider, error)
+
+var (
+ registryMu sync.RWMutex
+ registry = make(map[string]ProviderFactory)
+)
+
+// RegisterProvider registers a provider factory for a given type identifier.
+func RegisterProvider(typ string, factory ProviderFactory) {
+ if typ == "" || factory == nil {
+ return
+ }
+ registryMu.Lock()
+ registry[typ] = factory
+ registryMu.Unlock()
+}
+
+func BuildProvider(cfg *config.AccessProvider, root *config.SDKConfig) (Provider, error) {
+ if cfg == nil {
+ return nil, fmt.Errorf("access: nil provider config")
+ }
+ registryMu.RLock()
+ factory, ok := registry[cfg.Type]
+ registryMu.RUnlock()
+ if !ok {
+ return nil, fmt.Errorf("access: provider type %q is not registered", cfg.Type)
+ }
+ provider, err := factory(cfg, root)
+ if err != nil {
+ return nil, fmt.Errorf("access: failed to build provider %q: %w", cfg.Name, err)
+ }
+ return provider, nil
+}
+
+// BuildProviders constructs providers declared in configuration.
+func BuildProviders(root *config.SDKConfig) ([]Provider, error) {
+ if root == nil {
+ return nil, nil
+ }
+ providers := make([]Provider, 0, len(root.Access.Providers))
+ for i := range root.Access.Providers {
+ providerCfg := &root.Access.Providers[i]
+ if providerCfg.Type == "" {
+ continue
+ }
+ provider, err := BuildProvider(providerCfg, root)
+ if err != nil {
+ return nil, err
+ }
+ providers = append(providers, provider)
+ }
+ if len(providers) == 0 {
+ if inline := config.MakeInlineAPIKeyProvider(root.APIKeys); inline != nil {
+ provider, err := BuildProvider(inline, root)
+ if err != nil {
+ return nil, err
+ }
+ providers = append(providers, provider)
+ }
+ }
+ return providers, nil
+}
diff --git a/sdk/api/handlers/claude/code_handlers.go b/sdk/api/handlers/claude/code_handlers.go
new file mode 100644
index 0000000000000000000000000000000000000000..22e10fa59822d78be50db705f8a3045d98013419
--- /dev/null
+++ b/sdk/api/handlers/claude/code_handlers.go
@@ -0,0 +1,323 @@
+// Package claude provides HTTP handlers for Claude API code-related functionality.
+// This package implements Claude-compatible streaming chat completions with sophisticated
+// client rotation and quota management systems to ensure high availability and optimal
+// resource utilization across multiple backend clients. It handles request translation
+// between Claude API format and the underlying Gemini backend, providing seamless
+// API compatibility while maintaining robust error handling and connection management.
+package claude
+
+import (
+ "bytes"
+ "compress/gzip"
+ "context"
+ "encoding/json"
+ "fmt"
+ "io"
+ "net/http"
+
+ "github.com/gin-gonic/gin"
+ . "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
+ log "github.com/sirupsen/logrus"
+ "github.com/tidwall/gjson"
+)
+
+// ClaudeCodeAPIHandler contains the handlers for Claude API endpoints.
+// It holds a pool of clients to interact with the backend service.
+type ClaudeCodeAPIHandler struct {
+ *handlers.BaseAPIHandler
+}
+
+// NewClaudeCodeAPIHandler creates a new Claude API handlers instance.
+// It takes an BaseAPIHandler instance as input and returns a ClaudeCodeAPIHandler.
+//
+// Parameters:
+// - apiHandlers: The base API handler instance.
+//
+// Returns:
+// - *ClaudeCodeAPIHandler: A new Claude code API handler instance.
+func NewClaudeCodeAPIHandler(apiHandlers *handlers.BaseAPIHandler) *ClaudeCodeAPIHandler {
+ return &ClaudeCodeAPIHandler{
+ BaseAPIHandler: apiHandlers,
+ }
+}
+
+// HandlerType returns the identifier for this handler implementation.
+func (h *ClaudeCodeAPIHandler) HandlerType() string {
+ return Claude
+}
+
+// Models returns a list of models supported by this handler.
+func (h *ClaudeCodeAPIHandler) Models() []map[string]any {
+ // Get dynamic models from the global registry
+ modelRegistry := registry.GetGlobalRegistry()
+ return modelRegistry.GetAvailableModels("claude")
+}
+
+// ClaudeMessages handles Claude-compatible streaming chat completions.
+// This function implements a sophisticated client rotation and quota management system
+// to ensure high availability and optimal resource utilization across multiple backend clients.
+//
+// Parameters:
+// - c: The Gin context for the request.
+func (h *ClaudeCodeAPIHandler) ClaudeMessages(c *gin.Context) {
+ // Extract raw JSON data from the incoming request
+ rawJSON, err := c.GetRawData()
+ // If data retrieval fails, return a 400 Bad Request error.
+ if err != nil {
+ c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: fmt.Sprintf("Invalid request: %v", err),
+ Type: "invalid_request_error",
+ },
+ })
+ return
+ }
+
+ // Check if the client requested a streaming response.
+ streamResult := gjson.GetBytes(rawJSON, "stream")
+ if !streamResult.Exists() || streamResult.Type == gjson.False {
+ h.handleNonStreamingResponse(c, rawJSON)
+ } else {
+ h.handleStreamingResponse(c, rawJSON)
+ }
+}
+
+// ClaudeMessages handles Claude-compatible streaming chat completions.
+// This function implements a sophisticated client rotation and quota management system
+// to ensure high availability and optimal resource utilization across multiple backend clients.
+//
+// Parameters:
+// - c: The Gin context for the request.
+func (h *ClaudeCodeAPIHandler) ClaudeCountTokens(c *gin.Context) {
+ // Extract raw JSON data from the incoming request
+ rawJSON, err := c.GetRawData()
+ // If data retrieval fails, return a 400 Bad Request error.
+ if err != nil {
+ c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: fmt.Sprintf("Invalid request: %v", err),
+ Type: "invalid_request_error",
+ },
+ })
+ return
+ }
+
+ c.Header("Content-Type", "application/json")
+
+ alt := h.GetAlt(c)
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+
+ modelName := gjson.GetBytes(rawJSON, "model").String()
+
+ resp, errMsg := h.ExecuteCountWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
+ if errMsg != nil {
+ h.WriteErrorResponse(c, errMsg)
+ cliCancel(errMsg.Error)
+ return
+ }
+ _, _ = c.Writer.Write(resp)
+ cliCancel()
+}
+
+// ClaudeModels handles the Claude models listing endpoint.
+// It returns a JSON response containing available Claude models and their specifications.
+//
+// Parameters:
+// - c: The Gin context for the request.
+func (h *ClaudeCodeAPIHandler) ClaudeModels(c *gin.Context) {
+ models := h.Models()
+ firstID := ""
+ lastID := ""
+ if len(models) > 0 {
+ if id, ok := models[0]["id"].(string); ok {
+ firstID = id
+ }
+ if id, ok := models[len(models)-1]["id"].(string); ok {
+ lastID = id
+ }
+ }
+
+ c.JSON(http.StatusOK, gin.H{
+ "data": models,
+ "has_more": false,
+ "first_id": firstID,
+ "last_id": lastID,
+ })
+}
+
+// handleNonStreamingResponse handles non-streaming content generation requests for Claude models.
+// This function processes the request synchronously and returns the complete generated
+// response in a single API call. It supports various generation parameters and
+// response formats.
+//
+// Parameters:
+// - c: The Gin context for the request
+// - modelName: The name of the Gemini model to use for content generation
+// - rawJSON: The raw JSON request body containing generation parameters and content
+func (h *ClaudeCodeAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSON []byte) {
+ c.Header("Content-Type", "application/json")
+ alt := h.GetAlt(c)
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+ stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
+
+ modelName := gjson.GetBytes(rawJSON, "model").String()
+
+ resp, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
+ stopKeepAlive()
+ if errMsg != nil {
+ h.WriteErrorResponse(c, errMsg)
+ cliCancel(errMsg.Error)
+ return
+ }
+
+ // Decompress gzipped responses - Claude API sometimes returns gzip without Content-Encoding header
+ // This fixes title generation and other non-streaming responses that arrive compressed
+ if len(resp) >= 2 && resp[0] == 0x1f && resp[1] == 0x8b {
+ gzReader, errGzip := gzip.NewReader(bytes.NewReader(resp))
+ if errGzip != nil {
+ log.Warnf("failed to decompress gzipped Claude response: %v", errGzip)
+ } else {
+ defer func() {
+ if errClose := gzReader.Close(); errClose != nil {
+ log.Warnf("failed to close Claude gzip reader: %v", errClose)
+ }
+ }()
+ decompressed, errRead := io.ReadAll(gzReader)
+ if errRead != nil {
+ log.Warnf("failed to read decompressed Claude response: %v", errRead)
+ } else {
+ resp = decompressed
+ }
+ }
+ }
+
+ _, _ = c.Writer.Write(resp)
+ cliCancel()
+}
+
+// handleStreamingResponse streams Claude-compatible responses backed by Gemini.
+// It sets up SSE, selects a backend client with rotation/quota logic,
+// forwards chunks, and translates them to Claude CLI format.
+//
+// Parameters:
+// - c: The Gin context for the request.
+// - rawJSON: The raw JSON request body.
+func (h *ClaudeCodeAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON []byte) {
+ // Get the http.Flusher interface to manually flush the response.
+ // This is crucial for streaming as it allows immediate sending of data chunks
+ flusher, ok := c.Writer.(http.Flusher)
+ if !ok {
+ c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: "Streaming not supported",
+ Type: "server_error",
+ },
+ })
+ return
+ }
+
+ modelName := gjson.GetBytes(rawJSON, "model").String()
+
+ // Create a cancellable context for the backend client request
+ // This allows proper cleanup and cancellation of ongoing requests
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+
+ dataChan, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "")
+ setSSEHeaders := func() {
+ c.Header("Content-Type", "text/event-stream")
+ c.Header("Cache-Control", "no-cache")
+ c.Header("Connection", "keep-alive")
+ c.Header("Access-Control-Allow-Origin", "*")
+ }
+
+ // Peek at the first chunk to determine success or failure before setting headers
+ for {
+ select {
+ case <-c.Request.Context().Done():
+ cliCancel(c.Request.Context().Err())
+ return
+ case errMsg, ok := <-errChan:
+ if !ok {
+ // Err channel closed cleanly; wait for data channel.
+ errChan = nil
+ continue
+ }
+ // Upstream failed immediately. Return proper error status and JSON.
+ h.WriteErrorResponse(c, errMsg)
+ if errMsg != nil {
+ cliCancel(errMsg.Error)
+ } else {
+ cliCancel(nil)
+ }
+ return
+ case chunk, ok := <-dataChan:
+ if !ok {
+ // Stream closed without data? Send DONE or just headers.
+ setSSEHeaders()
+ flusher.Flush()
+ cliCancel(nil)
+ return
+ }
+
+ // Success! Set headers now.
+ setSSEHeaders()
+
+ // Write the first chunk
+ if len(chunk) > 0 {
+ _, _ = c.Writer.Write(chunk)
+ flusher.Flush()
+ }
+
+ // Continue streaming the rest
+ h.forwardClaudeStream(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan)
+ return
+ }
+ }
+}
+
+func (h *ClaudeCodeAPIHandler) forwardClaudeStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
+ h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
+ WriteChunk: func(chunk []byte) {
+ if len(chunk) == 0 {
+ return
+ }
+ _, _ = c.Writer.Write(chunk)
+ },
+ WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
+ if errMsg == nil {
+ return
+ }
+ status := http.StatusInternalServerError
+ if errMsg.StatusCode > 0 {
+ status = errMsg.StatusCode
+ }
+ c.Status(status)
+
+ errorBytes, _ := json.Marshal(h.toClaudeError(errMsg))
+ _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", errorBytes)
+ },
+ })
+}
+
+type claudeErrorDetail struct {
+ Type string `json:"type"`
+ Message string `json:"message"`
+}
+
+type claudeErrorResponse struct {
+ Type string `json:"type"`
+ Error claudeErrorDetail `json:"error"`
+}
+
+func (h *ClaudeCodeAPIHandler) toClaudeError(msg *interfaces.ErrorMessage) claudeErrorResponse {
+ return claudeErrorResponse{
+ Type: "error",
+ Error: claudeErrorDetail{
+ Type: "api_error",
+ Message: msg.Error.Error(),
+ },
+ }
+}
diff --git a/sdk/api/handlers/gemini/gemini-cli_handlers.go b/sdk/api/handlers/gemini/gemini-cli_handlers.go
new file mode 100644
index 0000000000000000000000000000000000000000..ea78657d6218a384e3b428d7205f526a64ae1540
--- /dev/null
+++ b/sdk/api/handlers/gemini/gemini-cli_handlers.go
@@ -0,0 +1,229 @@
+// Package gemini provides HTTP handlers for Gemini CLI API functionality.
+// This package implements handlers that process CLI-specific requests for Gemini API operations,
+// including content generation and streaming content generation endpoints.
+// The handlers restrict access to localhost only and manage communication with the backend service.
+package gemini
+
+import (
+ "bytes"
+ "context"
+ "fmt"
+ "io"
+ "net/http"
+ "strings"
+ "time"
+
+ "github.com/gin-gonic/gin"
+ . "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
+ log "github.com/sirupsen/logrus"
+ "github.com/tidwall/gjson"
+)
+
+// GeminiCLIAPIHandler contains the handlers for Gemini CLI API endpoints.
+// It holds a pool of clients to interact with the backend service.
+type GeminiCLIAPIHandler struct {
+ *handlers.BaseAPIHandler
+}
+
+// NewGeminiCLIAPIHandler creates a new Gemini CLI API handlers instance.
+// It takes an BaseAPIHandler instance as input and returns a GeminiCLIAPIHandler.
+func NewGeminiCLIAPIHandler(apiHandlers *handlers.BaseAPIHandler) *GeminiCLIAPIHandler {
+ return &GeminiCLIAPIHandler{
+ BaseAPIHandler: apiHandlers,
+ }
+}
+
+// HandlerType returns the type of this handler.
+func (h *GeminiCLIAPIHandler) HandlerType() string {
+ return GeminiCLI
+}
+
+// Models returns a list of models supported by this handler.
+func (h *GeminiCLIAPIHandler) Models() []map[string]any {
+ return make([]map[string]any, 0)
+}
+
+// CLIHandler handles CLI-specific requests for Gemini API operations.
+// It restricts access to localhost only and routes requests to appropriate internal handlers.
+func (h *GeminiCLIAPIHandler) CLIHandler(c *gin.Context) {
+ if !strings.HasPrefix(c.Request.RemoteAddr, "127.0.0.1:") {
+ c.JSON(http.StatusForbidden, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: "CLI reply only allow local access",
+ Type: "forbidden",
+ },
+ })
+ return
+ }
+
+ rawJSON, _ := c.GetRawData()
+ requestRawURI := c.Request.URL.Path
+
+ if requestRawURI == "/v1internal:generateContent" {
+ h.handleInternalGenerateContent(c, rawJSON)
+ } else if requestRawURI == "/v1internal:streamGenerateContent" {
+ h.handleInternalStreamGenerateContent(c, rawJSON)
+ } else {
+ reqBody := bytes.NewBuffer(rawJSON)
+ req, err := http.NewRequest("POST", fmt.Sprintf("https://cloudcode-pa.googleapis.com%s", c.Request.URL.RequestURI()), reqBody)
+ if err != nil {
+ c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: fmt.Sprintf("Invalid request: %v", err),
+ Type: "invalid_request_error",
+ },
+ })
+ return
+ }
+ for key, value := range c.Request.Header {
+ req.Header[key] = value
+ }
+
+ httpClient := util.SetProxy(h.Cfg, &http.Client{})
+
+ resp, err := httpClient.Do(req)
+ if err != nil {
+ c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: fmt.Sprintf("Invalid request: %v", err),
+ Type: "invalid_request_error",
+ },
+ })
+ return
+ }
+
+ if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+ defer func() {
+ if err = resp.Body.Close(); err != nil {
+ log.Printf("warn: failed to close response body: %v", err)
+ }
+ }()
+ bodyBytes, _ := io.ReadAll(resp.Body)
+
+ c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: string(bodyBytes),
+ Type: "invalid_request_error",
+ },
+ })
+ return
+ }
+
+ defer func() {
+ _ = resp.Body.Close()
+ }()
+
+ for key, value := range resp.Header {
+ c.Header(key, value[0])
+ }
+ output, err := io.ReadAll(resp.Body)
+ if err != nil {
+ log.Errorf("Failed to read response body: %v", err)
+ return
+ }
+ _, _ = c.Writer.Write(output)
+ c.Set("API_RESPONSE", output)
+ }
+}
+
+// handleInternalStreamGenerateContent handles streaming content generation requests.
+// It sets up a server-sent event stream and forwards the request to the backend client.
+// The function continuously proxies response chunks from the backend to the client.
+func (h *GeminiCLIAPIHandler) handleInternalStreamGenerateContent(c *gin.Context, rawJSON []byte) {
+ alt := h.GetAlt(c)
+
+ if alt == "" {
+ c.Header("Content-Type", "text/event-stream")
+ c.Header("Cache-Control", "no-cache")
+ c.Header("Connection", "keep-alive")
+ c.Header("Access-Control-Allow-Origin", "*")
+ }
+
+ // Get the http.Flusher interface to manually flush the response.
+ flusher, ok := c.Writer.(http.Flusher)
+ if !ok {
+ c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: "Streaming not supported",
+ Type: "server_error",
+ },
+ })
+ return
+ }
+
+ modelResult := gjson.GetBytes(rawJSON, "model")
+ modelName := modelResult.String()
+
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+ dataChan, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "")
+ h.forwardCLIStream(c, flusher, "", func(err error) { cliCancel(err) }, dataChan, errChan)
+ return
+}
+
+// handleInternalGenerateContent handles non-streaming content generation requests.
+// It sends a request to the backend client and proxies the entire response back to the client at once.
+func (h *GeminiCLIAPIHandler) handleInternalGenerateContent(c *gin.Context, rawJSON []byte) {
+ c.Header("Content-Type", "application/json")
+ modelResult := gjson.GetBytes(rawJSON, "model")
+ modelName := modelResult.String()
+
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+ resp, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "")
+ if errMsg != nil {
+ h.WriteErrorResponse(c, errMsg)
+ cliCancel(errMsg.Error)
+ return
+ }
+ _, _ = c.Writer.Write(resp)
+ cliCancel()
+}
+
+func (h *GeminiCLIAPIHandler) forwardCLIStream(c *gin.Context, flusher http.Flusher, alt string, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
+ var keepAliveInterval *time.Duration
+ if alt != "" {
+ disabled := time.Duration(0)
+ keepAliveInterval = &disabled
+ }
+
+ h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
+ KeepAliveInterval: keepAliveInterval,
+ WriteChunk: func(chunk []byte) {
+ if alt == "" {
+ if bytes.Equal(chunk, []byte("data: [DONE]")) || bytes.Equal(chunk, []byte("[DONE]")) {
+ return
+ }
+
+ if !bytes.HasPrefix(chunk, []byte("data:")) {
+ _, _ = c.Writer.Write([]byte("data: "))
+ }
+
+ _, _ = c.Writer.Write(chunk)
+ _, _ = c.Writer.Write([]byte("\n\n"))
+ } else {
+ _, _ = c.Writer.Write(chunk)
+ }
+ },
+ WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
+ if errMsg == nil {
+ return
+ }
+ status := http.StatusInternalServerError
+ if errMsg.StatusCode > 0 {
+ status = errMsg.StatusCode
+ }
+ errText := http.StatusText(status)
+ if errMsg.Error != nil && errMsg.Error.Error() != "" {
+ errText = errMsg.Error.Error()
+ }
+ body := handlers.BuildErrorResponseBody(status, errText)
+ if alt == "" {
+ _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", string(body))
+ } else {
+ _, _ = c.Writer.Write(body)
+ }
+ },
+ })
+}
diff --git a/sdk/api/handlers/gemini/gemini_handlers.go b/sdk/api/handlers/gemini/gemini_handlers.go
new file mode 100644
index 0000000000000000000000000000000000000000..71c485ad01257a20c2ef9d620a6ad99c76242188
--- /dev/null
+++ b/sdk/api/handlers/gemini/gemini_handlers.go
@@ -0,0 +1,338 @@
+// Package gemini provides HTTP handlers for Gemini API endpoints.
+// This package implements handlers for managing Gemini model operations including
+// model listing, content generation, streaming content generation, and token counting.
+// It serves as a proxy layer between clients and the Gemini backend service,
+// handling request translation, client management, and response processing.
+package gemini
+
+import (
+ "context"
+ "fmt"
+ "net/http"
+ "strings"
+ "time"
+
+ "github.com/gin-gonic/gin"
+ . "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
+)
+
+// GeminiAPIHandler contains the handlers for Gemini API endpoints.
+// It holds a pool of clients to interact with the backend service.
+type GeminiAPIHandler struct {
+ *handlers.BaseAPIHandler
+}
+
+// NewGeminiAPIHandler creates a new Gemini API handlers instance.
+// It takes an BaseAPIHandler instance as input and returns a GeminiAPIHandler.
+func NewGeminiAPIHandler(apiHandlers *handlers.BaseAPIHandler) *GeminiAPIHandler {
+ return &GeminiAPIHandler{
+ BaseAPIHandler: apiHandlers,
+ }
+}
+
+// HandlerType returns the identifier for this handler implementation.
+func (h *GeminiAPIHandler) HandlerType() string {
+ return Gemini
+}
+
+// Models returns the Gemini-compatible model metadata supported by this handler.
+func (h *GeminiAPIHandler) Models() []map[string]any {
+ // Get dynamic models from the global registry
+ modelRegistry := registry.GetGlobalRegistry()
+ return modelRegistry.GetAvailableModels("gemini")
+}
+
+// GeminiModels handles the Gemini models listing endpoint.
+// It returns a JSON response containing available Gemini models and their specifications.
+func (h *GeminiAPIHandler) GeminiModels(c *gin.Context) {
+ rawModels := h.Models()
+ normalizedModels := make([]map[string]any, 0, len(rawModels))
+ defaultMethods := []string{"generateContent"}
+ for _, model := range rawModels {
+ normalizedModel := make(map[string]any, len(model))
+ for k, v := range model {
+ normalizedModel[k] = v
+ }
+ if name, ok := normalizedModel["name"].(string); ok && name != "" {
+ if !strings.HasPrefix(name, "models/") {
+ normalizedModel["name"] = "models/" + name
+ }
+ if displayName, _ := normalizedModel["displayName"].(string); displayName == "" {
+ normalizedModel["displayName"] = name
+ }
+ if description, _ := normalizedModel["description"].(string); description == "" {
+ normalizedModel["description"] = name
+ }
+ }
+ if _, ok := normalizedModel["supportedGenerationMethods"]; !ok {
+ normalizedModel["supportedGenerationMethods"] = defaultMethods
+ }
+ normalizedModels = append(normalizedModels, normalizedModel)
+ }
+ c.JSON(http.StatusOK, gin.H{
+ "models": normalizedModels,
+ })
+}
+
+// GeminiGetHandler handles GET requests for specific Gemini model information.
+// It returns detailed information about a specific Gemini model based on the action parameter.
+func (h *GeminiAPIHandler) GeminiGetHandler(c *gin.Context) {
+ var request struct {
+ Action string `uri:"action" binding:"required"`
+ }
+ if err := c.ShouldBindUri(&request); err != nil {
+ c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: fmt.Sprintf("Invalid request: %v", err),
+ Type: "invalid_request_error",
+ },
+ })
+ return
+ }
+ action := strings.TrimPrefix(request.Action, "/")
+
+ // Get dynamic models from the global registry and find the matching one
+ availableModels := h.Models()
+ var targetModel map[string]any
+
+ for _, model := range availableModels {
+ name, _ := model["name"].(string)
+ // Match name with or without 'models/' prefix
+ if name == action || name == "models/"+action {
+ targetModel = model
+ break
+ }
+ }
+
+ if targetModel != nil {
+ // Ensure the name has 'models/' prefix in the output if it's a Gemini model
+ if name, ok := targetModel["name"].(string); ok && name != "" && !strings.HasPrefix(name, "models/") {
+ targetModel["name"] = "models/" + name
+ }
+ c.JSON(http.StatusOK, targetModel)
+ return
+ }
+
+ c.JSON(http.StatusNotFound, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: "Not Found",
+ Type: "not_found",
+ },
+ })
+}
+
+// GeminiHandler handles POST requests for Gemini API operations.
+// It routes requests to appropriate handlers based on the action parameter (model:method format).
+func (h *GeminiAPIHandler) GeminiHandler(c *gin.Context) {
+ var request struct {
+ Action string `uri:"action" binding:"required"`
+ }
+ if err := c.ShouldBindUri(&request); err != nil {
+ c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: fmt.Sprintf("Invalid request: %v", err),
+ Type: "invalid_request_error",
+ },
+ })
+ return
+ }
+ action := strings.Split(strings.TrimPrefix(request.Action, "/"), ":")
+ if len(action) != 2 {
+ c.JSON(http.StatusNotFound, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: fmt.Sprintf("%s not found.", c.Request.URL.Path),
+ Type: "invalid_request_error",
+ },
+ })
+ return
+ }
+
+ method := action[1]
+ rawJSON, _ := c.GetRawData()
+
+ switch method {
+ case "generateContent":
+ h.handleGenerateContent(c, action[0], rawJSON)
+ case "streamGenerateContent":
+ h.handleStreamGenerateContent(c, action[0], rawJSON)
+ case "countTokens":
+ h.handleCountTokens(c, action[0], rawJSON)
+ }
+}
+
+// handleStreamGenerateContent handles streaming content generation requests for Gemini models.
+// This function establishes a Server-Sent Events connection and streams the generated content
+// back to the client in real-time. It supports both SSE format and direct streaming based
+// on the 'alt' query parameter.
+//
+// Parameters:
+// - c: The Gin context for the request
+// - modelName: The name of the Gemini model to use for content generation
+// - rawJSON: The raw JSON request body containing generation parameters
+func (h *GeminiAPIHandler) handleStreamGenerateContent(c *gin.Context, modelName string, rawJSON []byte) {
+ alt := h.GetAlt(c)
+
+ // Get the http.Flusher interface to manually flush the response.
+ flusher, ok := c.Writer.(http.Flusher)
+ if !ok {
+ c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: "Streaming not supported",
+ Type: "server_error",
+ },
+ })
+ return
+ }
+
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+ dataChan, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
+
+ setSSEHeaders := func() {
+ c.Header("Content-Type", "text/event-stream")
+ c.Header("Cache-Control", "no-cache")
+ c.Header("Connection", "keep-alive")
+ c.Header("Access-Control-Allow-Origin", "*")
+ }
+
+ // Peek at the first chunk
+ for {
+ select {
+ case <-c.Request.Context().Done():
+ cliCancel(c.Request.Context().Err())
+ return
+ case errMsg, ok := <-errChan:
+ if !ok {
+ // Err channel closed cleanly; wait for data channel.
+ errChan = nil
+ continue
+ }
+ // Upstream failed immediately. Return proper error status and JSON.
+ h.WriteErrorResponse(c, errMsg)
+ if errMsg != nil {
+ cliCancel(errMsg.Error)
+ } else {
+ cliCancel(nil)
+ }
+ return
+ case chunk, ok := <-dataChan:
+ if !ok {
+ // Closed without data
+ if alt == "" {
+ setSSEHeaders()
+ }
+ flusher.Flush()
+ cliCancel(nil)
+ return
+ }
+
+ // Success! Set headers.
+ if alt == "" {
+ setSSEHeaders()
+ }
+
+ // Write first chunk
+ if alt == "" {
+ _, _ = c.Writer.Write([]byte("data: "))
+ _, _ = c.Writer.Write(chunk)
+ _, _ = c.Writer.Write([]byte("\n\n"))
+ } else {
+ _, _ = c.Writer.Write(chunk)
+ }
+ flusher.Flush()
+
+ // Continue
+ h.forwardGeminiStream(c, flusher, alt, func(err error) { cliCancel(err) }, dataChan, errChan)
+ return
+ }
+ }
+}
+
+// handleCountTokens handles token counting requests for Gemini models.
+// This function counts the number of tokens in the provided content without
+// generating a response. It's useful for quota management and content validation.
+//
+// Parameters:
+// - c: The Gin context for the request
+// - modelName: The name of the Gemini model to use for token counting
+// - rawJSON: The raw JSON request body containing the content to count
+func (h *GeminiAPIHandler) handleCountTokens(c *gin.Context, modelName string, rawJSON []byte) {
+ c.Header("Content-Type", "application/json")
+ alt := h.GetAlt(c)
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+ resp, errMsg := h.ExecuteCountWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
+ if errMsg != nil {
+ h.WriteErrorResponse(c, errMsg)
+ cliCancel(errMsg.Error)
+ return
+ }
+ _, _ = c.Writer.Write(resp)
+ cliCancel()
+}
+
+// handleGenerateContent handles non-streaming content generation requests for Gemini models.
+// This function processes the request synchronously and returns the complete generated
+// response in a single API call. It supports various generation parameters and
+// response formats.
+//
+// Parameters:
+// - c: The Gin context for the request
+// - modelName: The name of the Gemini model to use for content generation
+// - rawJSON: The raw JSON request body containing generation parameters and content
+func (h *GeminiAPIHandler) handleGenerateContent(c *gin.Context, modelName string, rawJSON []byte) {
+ c.Header("Content-Type", "application/json")
+ alt := h.GetAlt(c)
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+ stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
+ resp, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
+ stopKeepAlive()
+ if errMsg != nil {
+ h.WriteErrorResponse(c, errMsg)
+ cliCancel(errMsg.Error)
+ return
+ }
+ _, _ = c.Writer.Write(resp)
+ cliCancel()
+}
+
+func (h *GeminiAPIHandler) forwardGeminiStream(c *gin.Context, flusher http.Flusher, alt string, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
+ var keepAliveInterval *time.Duration
+ if alt != "" {
+ disabled := time.Duration(0)
+ keepAliveInterval = &disabled
+ }
+
+ h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
+ KeepAliveInterval: keepAliveInterval,
+ WriteChunk: func(chunk []byte) {
+ if alt == "" {
+ _, _ = c.Writer.Write([]byte("data: "))
+ _, _ = c.Writer.Write(chunk)
+ _, _ = c.Writer.Write([]byte("\n\n"))
+ } else {
+ _, _ = c.Writer.Write(chunk)
+ }
+ },
+ WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
+ if errMsg == nil {
+ return
+ }
+ status := http.StatusInternalServerError
+ if errMsg.StatusCode > 0 {
+ status = errMsg.StatusCode
+ }
+ errText := http.StatusText(status)
+ if errMsg.Error != nil && errMsg.Error.Error() != "" {
+ errText = errMsg.Error.Error()
+ }
+ body := handlers.BuildErrorResponseBody(status, errText)
+ if alt == "" {
+ _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", string(body))
+ } else {
+ _, _ = c.Writer.Write(body)
+ }
+ },
+ })
+}
diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go
new file mode 100644
index 0000000000000000000000000000000000000000..b1da966422dc7c2274265f0ee936dc49bf5e6cf5
--- /dev/null
+++ b/sdk/api/handlers/handlers.go
@@ -0,0 +1,745 @@
+// Package handlers provides core API handler functionality for the CLI Proxy API server.
+// It includes common types, client management, load balancing, and error handling
+// shared across all API endpoint handlers (OpenAI, Claude, Gemini).
+package handlers
+
+import (
+ "bytes"
+ "encoding/json"
+ "fmt"
+ "net/http"
+ "strings"
+ "sync"
+ "time"
+
+ "github.com/gin-gonic/gin"
+ "github.com/google/uuid"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ coreexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
+ sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
+ "golang.org/x/net/context"
+)
+
+// ErrorResponse represents a standard error response format for the API.
+// It contains a single ErrorDetail field.
+type ErrorResponse struct {
+ // Error contains detailed information about the error that occurred.
+ Error ErrorDetail `json:"error"`
+}
+
+// ErrorDetail provides specific information about an error that occurred.
+// It includes a human-readable message, an error type, and an optional error code.
+type ErrorDetail struct {
+ // Message is a human-readable message providing more details about the error.
+ Message string `json:"message"`
+
+ // Type is the category of error that occurred (e.g., "invalid_request_error").
+ Type string `json:"type"`
+
+ // Code is a short code identifying the error, if applicable.
+ Code string `json:"code,omitempty"`
+}
+
+const idempotencyKeyMetadataKey = "idempotency_key"
+
+const (
+ defaultStreamingKeepAliveSeconds = 0
+ defaultStreamingBootstrapRetries = 0
+)
+
+// BuildErrorResponseBody builds an OpenAI-compatible JSON error response body.
+// If errText is already valid JSON, it is returned as-is to preserve upstream error payloads.
+func BuildErrorResponseBody(status int, errText string) []byte {
+ if status <= 0 {
+ status = http.StatusInternalServerError
+ }
+ if strings.TrimSpace(errText) == "" {
+ errText = http.StatusText(status)
+ }
+
+ trimmed := strings.TrimSpace(errText)
+ if trimmed != "" && json.Valid([]byte(trimmed)) {
+ return []byte(trimmed)
+ }
+
+ errType := "invalid_request_error"
+ var code string
+ switch status {
+ case http.StatusUnauthorized:
+ errType = "authentication_error"
+ code = "invalid_api_key"
+ case http.StatusForbidden:
+ errType = "permission_error"
+ code = "insufficient_quota"
+ case http.StatusTooManyRequests:
+ errType = "rate_limit_error"
+ code = "rate_limit_exceeded"
+ case http.StatusNotFound:
+ errType = "invalid_request_error"
+ code = "model_not_found"
+ default:
+ if status >= http.StatusInternalServerError {
+ errType = "server_error"
+ code = "internal_server_error"
+ }
+ }
+
+ payload, err := json.Marshal(ErrorResponse{
+ Error: ErrorDetail{
+ Message: errText,
+ Type: errType,
+ Code: code,
+ },
+ })
+ if err != nil {
+ return []byte(fmt.Sprintf(`{"error":{"message":%q,"type":"server_error","code":"internal_server_error"}}`, errText))
+ }
+ return payload
+}
+
+// StreamingKeepAliveInterval returns the SSE keep-alive interval for this server.
+// Returning 0 disables keep-alives (default when unset).
+func StreamingKeepAliveInterval(cfg *config.SDKConfig) time.Duration {
+ seconds := defaultStreamingKeepAliveSeconds
+ if cfg != nil {
+ seconds = cfg.Streaming.KeepAliveSeconds
+ }
+ if seconds <= 0 {
+ return 0
+ }
+ return time.Duration(seconds) * time.Second
+}
+
+// NonStreamingKeepAliveInterval returns the keep-alive interval for non-streaming responses.
+// Returning 0 disables keep-alives (default when unset).
+func NonStreamingKeepAliveInterval(cfg *config.SDKConfig) time.Duration {
+ seconds := 0
+ if cfg != nil {
+ seconds = cfg.NonStreamKeepAliveInterval
+ }
+ if seconds <= 0 {
+ return 0
+ }
+ return time.Duration(seconds) * time.Second
+}
+
+// StreamingBootstrapRetries returns how many times a streaming request may be retried before any bytes are sent.
+func StreamingBootstrapRetries(cfg *config.SDKConfig) int {
+ retries := defaultStreamingBootstrapRetries
+ if cfg != nil {
+ retries = cfg.Streaming.BootstrapRetries
+ }
+ if retries < 0 {
+ retries = 0
+ }
+ return retries
+}
+
+func requestExecutionMetadata(ctx context.Context) map[string]any {
+ // Idempotency-Key is an optional client-supplied header used to correlate retries.
+ // It is forwarded as execution metadata; when absent we generate a UUID.
+ key := ""
+ if ctx != nil {
+ if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil {
+ key = strings.TrimSpace(ginCtx.GetHeader("Idempotency-Key"))
+ }
+ }
+ if key == "" {
+ key = uuid.NewString()
+ }
+ return map[string]any{idempotencyKeyMetadataKey: key}
+}
+
+func mergeMetadata(base, overlay map[string]any) map[string]any {
+ if len(base) == 0 && len(overlay) == 0 {
+ return nil
+ }
+ out := make(map[string]any, len(base)+len(overlay))
+ for k, v := range base {
+ out[k] = v
+ }
+ for k, v := range overlay {
+ out[k] = v
+ }
+ return out
+}
+
+// BaseAPIHandler contains the handlers for API endpoints.
+// It holds a pool of clients to interact with the backend service and manages
+// load balancing, client selection, and configuration.
+type BaseAPIHandler struct {
+ // AuthManager manages auth lifecycle and execution in the new architecture.
+ AuthManager *coreauth.Manager
+
+ // Cfg holds the current application configuration.
+ Cfg *config.SDKConfig
+}
+
+// NewBaseAPIHandlers creates a new API handlers instance.
+// It takes a slice of clients and configuration as input.
+//
+// Parameters:
+// - cliClients: A slice of AI service clients
+// - cfg: The application configuration
+//
+// Returns:
+// - *BaseAPIHandler: A new API handlers instance
+func NewBaseAPIHandlers(cfg *config.SDKConfig, authManager *coreauth.Manager) *BaseAPIHandler {
+ return &BaseAPIHandler{
+ Cfg: cfg,
+ AuthManager: authManager,
+ }
+}
+
+// UpdateClients updates the handlers' client list and configuration.
+// This method is called when the configuration or authentication tokens change.
+//
+// Parameters:
+// - clients: The new slice of AI service clients
+// - cfg: The new application configuration
+func (h *BaseAPIHandler) UpdateClients(cfg *config.SDKConfig) { h.Cfg = cfg }
+
+// GetAlt extracts the 'alt' parameter from the request query string.
+// It checks both 'alt' and '$alt' parameters and returns the appropriate value.
+//
+// Parameters:
+// - c: The Gin context containing the HTTP request
+//
+// Returns:
+// - string: The alt parameter value, or empty string if it's "sse"
+func (h *BaseAPIHandler) GetAlt(c *gin.Context) string {
+ var alt string
+ var hasAlt bool
+ alt, hasAlt = c.GetQuery("alt")
+ if !hasAlt {
+ alt, _ = c.GetQuery("$alt")
+ }
+ if alt == "sse" {
+ return ""
+ }
+ return alt
+}
+
+// GetContextWithCancel creates a new context with cancellation capabilities.
+// It embeds the Gin context and the API handler into the new context for later use.
+// The returned cancel function also handles logging the API response if request logging is enabled.
+//
+// Parameters:
+// - handler: The API handler associated with the request.
+// - c: The Gin context of the current request.
+// - ctx: The parent context (caller values/deadlines are preserved; request context adds cancellation and request ID).
+//
+// Returns:
+// - context.Context: The new context with cancellation and embedded values.
+// - APIHandlerCancelFunc: A function to cancel the context and log the response.
+func (h *BaseAPIHandler) GetContextWithCancel(handler interfaces.APIHandler, c *gin.Context, ctx context.Context) (context.Context, APIHandlerCancelFunc) {
+ parentCtx := ctx
+ if parentCtx == nil {
+ parentCtx = context.Background()
+ }
+
+ var requestCtx context.Context
+ if c != nil && c.Request != nil {
+ requestCtx = c.Request.Context()
+ }
+
+ if requestCtx != nil && logging.GetRequestID(parentCtx) == "" {
+ if requestID := logging.GetRequestID(requestCtx); requestID != "" {
+ parentCtx = logging.WithRequestID(parentCtx, requestID)
+ } else if requestID := logging.GetGinRequestID(c); requestID != "" {
+ parentCtx = logging.WithRequestID(parentCtx, requestID)
+ }
+ }
+ newCtx, cancel := context.WithCancel(parentCtx)
+ if requestCtx != nil && requestCtx != parentCtx {
+ go func() {
+ select {
+ case <-requestCtx.Done():
+ cancel()
+ case <-newCtx.Done():
+ }
+ }()
+ }
+ newCtx = context.WithValue(newCtx, "gin", c)
+ newCtx = context.WithValue(newCtx, "handler", handler)
+ return newCtx, func(params ...interface{}) {
+ if h.Cfg.RequestLog && len(params) == 1 {
+ if existing, exists := c.Get("API_RESPONSE"); exists {
+ if existingBytes, ok := existing.([]byte); ok && len(bytes.TrimSpace(existingBytes)) > 0 {
+ switch params[0].(type) {
+ case error, string:
+ cancel()
+ return
+ }
+ }
+ }
+
+ var payload []byte
+ switch data := params[0].(type) {
+ case []byte:
+ payload = data
+ case error:
+ if data != nil {
+ payload = []byte(data.Error())
+ }
+ case string:
+ payload = []byte(data)
+ }
+ if len(payload) > 0 {
+ if existing, exists := c.Get("API_RESPONSE"); exists {
+ if existingBytes, ok := existing.([]byte); ok && len(existingBytes) > 0 {
+ trimmedPayload := bytes.TrimSpace(payload)
+ if len(trimmedPayload) > 0 && bytes.Contains(existingBytes, trimmedPayload) {
+ cancel()
+ return
+ }
+ }
+ }
+ appendAPIResponse(c, payload)
+ }
+ }
+
+ cancel()
+ }
+}
+
+// StartNonStreamingKeepAlive emits blank lines every 5 seconds while waiting for a non-streaming response.
+// It returns a stop function that must be called before writing the final response.
+func (h *BaseAPIHandler) StartNonStreamingKeepAlive(c *gin.Context, ctx context.Context) func() {
+ if h == nil || c == nil {
+ return func() {}
+ }
+ interval := NonStreamingKeepAliveInterval(h.Cfg)
+ if interval <= 0 {
+ return func() {}
+ }
+ flusher, ok := c.Writer.(http.Flusher)
+ if !ok {
+ return func() {}
+ }
+ if ctx == nil {
+ ctx = context.Background()
+ }
+
+ stopChan := make(chan struct{})
+ var stopOnce sync.Once
+ var wg sync.WaitGroup
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ ticker := time.NewTicker(interval)
+ defer ticker.Stop()
+ for {
+ select {
+ case <-stopChan:
+ return
+ case <-ctx.Done():
+ return
+ case <-ticker.C:
+ _, _ = c.Writer.Write([]byte("\n"))
+ flusher.Flush()
+ }
+ }
+ }()
+
+ return func() {
+ stopOnce.Do(func() {
+ close(stopChan)
+ })
+ wg.Wait()
+ }
+}
+
+// appendAPIResponse preserves any previously captured API response and appends new data.
+func appendAPIResponse(c *gin.Context, data []byte) {
+ if c == nil || len(data) == 0 {
+ return
+ }
+
+ if existing, exists := c.Get("API_RESPONSE"); exists {
+ if existingBytes, ok := existing.([]byte); ok && len(existingBytes) > 0 {
+ combined := make([]byte, 0, len(existingBytes)+len(data)+1)
+ combined = append(combined, existingBytes...)
+ if existingBytes[len(existingBytes)-1] != '\n' {
+ combined = append(combined, '\n')
+ }
+ combined = append(combined, data...)
+ c.Set("API_RESPONSE", combined)
+ return
+ }
+ }
+
+ c.Set("API_RESPONSE", bytes.Clone(data))
+}
+
+// ExecuteWithAuthManager executes a non-streaming request via the core auth manager.
+// This path is the only supported execution route.
+func (h *BaseAPIHandler) ExecuteWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) ([]byte, *interfaces.ErrorMessage) {
+ providers, normalizedModel, errMsg := h.getRequestDetails(modelName)
+ if errMsg != nil {
+ return nil, errMsg
+ }
+ reqMeta := requestExecutionMetadata(ctx)
+ reqMeta[coreexecutor.RequestedModelMetadataKey] = normalizedModel
+ req := coreexecutor.Request{
+ Model: normalizedModel,
+ Payload: cloneBytes(rawJSON),
+ }
+ opts := coreexecutor.Options{
+ Stream: false,
+ Alt: alt,
+ OriginalRequest: cloneBytes(rawJSON),
+ SourceFormat: sdktranslator.FromString(handlerType),
+ }
+ opts.Metadata = reqMeta
+ resp, err := h.AuthManager.Execute(ctx, providers, req, opts)
+ if err != nil {
+ status := http.StatusInternalServerError
+ if se, ok := err.(interface{ StatusCode() int }); ok && se != nil {
+ if code := se.StatusCode(); code > 0 {
+ status = code
+ }
+ }
+ var addon http.Header
+ if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil {
+ if hdr := he.Headers(); hdr != nil {
+ addon = hdr.Clone()
+ }
+ }
+ return nil, &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon}
+ }
+ return cloneBytes(resp.Payload), nil
+}
+
+// ExecuteCountWithAuthManager executes a non-streaming request via the core auth manager.
+// This path is the only supported execution route.
+func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) ([]byte, *interfaces.ErrorMessage) {
+ providers, normalizedModel, errMsg := h.getRequestDetails(modelName)
+ if errMsg != nil {
+ return nil, errMsg
+ }
+ reqMeta := requestExecutionMetadata(ctx)
+ reqMeta[coreexecutor.RequestedModelMetadataKey] = normalizedModel
+ req := coreexecutor.Request{
+ Model: normalizedModel,
+ Payload: cloneBytes(rawJSON),
+ }
+ opts := coreexecutor.Options{
+ Stream: false,
+ Alt: alt,
+ OriginalRequest: cloneBytes(rawJSON),
+ SourceFormat: sdktranslator.FromString(handlerType),
+ }
+ opts.Metadata = reqMeta
+ resp, err := h.AuthManager.ExecuteCount(ctx, providers, req, opts)
+ if err != nil {
+ status := http.StatusInternalServerError
+ if se, ok := err.(interface{ StatusCode() int }); ok && se != nil {
+ if code := se.StatusCode(); code > 0 {
+ status = code
+ }
+ }
+ var addon http.Header
+ if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil {
+ if hdr := he.Headers(); hdr != nil {
+ addon = hdr.Clone()
+ }
+ }
+ return nil, &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon}
+ }
+ return cloneBytes(resp.Payload), nil
+}
+
+// ExecuteStreamWithAuthManager executes a streaming request via the core auth manager.
+// This path is the only supported execution route.
+func (h *BaseAPIHandler) ExecuteStreamWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) (<-chan []byte, <-chan *interfaces.ErrorMessage) {
+ providers, normalizedModel, errMsg := h.getRequestDetails(modelName)
+ if errMsg != nil {
+ errChan := make(chan *interfaces.ErrorMessage, 1)
+ errChan <- errMsg
+ close(errChan)
+ return nil, errChan
+ }
+ reqMeta := requestExecutionMetadata(ctx)
+ reqMeta[coreexecutor.RequestedModelMetadataKey] = normalizedModel
+ req := coreexecutor.Request{
+ Model: normalizedModel,
+ Payload: cloneBytes(rawJSON),
+ }
+ opts := coreexecutor.Options{
+ Stream: true,
+ Alt: alt,
+ OriginalRequest: cloneBytes(rawJSON),
+ SourceFormat: sdktranslator.FromString(handlerType),
+ }
+ opts.Metadata = reqMeta
+ chunks, err := h.AuthManager.ExecuteStream(ctx, providers, req, opts)
+ if err != nil {
+ errChan := make(chan *interfaces.ErrorMessage, 1)
+ status := http.StatusInternalServerError
+ if se, ok := err.(interface{ StatusCode() int }); ok && se != nil {
+ if code := se.StatusCode(); code > 0 {
+ status = code
+ }
+ }
+ var addon http.Header
+ if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil {
+ if hdr := he.Headers(); hdr != nil {
+ addon = hdr.Clone()
+ }
+ }
+ errChan <- &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon}
+ close(errChan)
+ return nil, errChan
+ }
+ dataChan := make(chan []byte)
+ errChan := make(chan *interfaces.ErrorMessage, 1)
+ go func() {
+ defer close(dataChan)
+ defer close(errChan)
+ sentPayload := false
+ bootstrapRetries := 0
+ maxBootstrapRetries := StreamingBootstrapRetries(h.Cfg)
+
+ sendErr := func(msg *interfaces.ErrorMessage) bool {
+ if ctx == nil {
+ errChan <- msg
+ return true
+ }
+ select {
+ case <-ctx.Done():
+ return false
+ case errChan <- msg:
+ return true
+ }
+ }
+
+ sendData := func(chunk []byte) bool {
+ if ctx == nil {
+ dataChan <- chunk
+ return true
+ }
+ select {
+ case <-ctx.Done():
+ return false
+ case dataChan <- chunk:
+ return true
+ }
+ }
+
+ bootstrapEligible := func(err error) bool {
+ status := statusFromError(err)
+ if status == 0 {
+ return true
+ }
+ switch status {
+ case http.StatusUnauthorized, http.StatusForbidden, http.StatusPaymentRequired,
+ http.StatusRequestTimeout, http.StatusTooManyRequests:
+ return true
+ default:
+ return status >= http.StatusInternalServerError
+ }
+ }
+
+ outer:
+ for {
+ for {
+ var chunk coreexecutor.StreamChunk
+ var ok bool
+ if ctx != nil {
+ select {
+ case <-ctx.Done():
+ return
+ case chunk, ok = <-chunks:
+ }
+ } else {
+ chunk, ok = <-chunks
+ }
+ if !ok {
+ return
+ }
+ if chunk.Err != nil {
+ streamErr := chunk.Err
+ // Safe bootstrap recovery: if the upstream fails before any payload bytes are sent,
+ // retry a few times (to allow auth rotation / transient recovery) and then attempt model fallback.
+ if !sentPayload {
+ if bootstrapRetries < maxBootstrapRetries && bootstrapEligible(streamErr) {
+ bootstrapRetries++
+ retryChunks, retryErr := h.AuthManager.ExecuteStream(ctx, providers, req, opts)
+ if retryErr == nil {
+ chunks = retryChunks
+ continue outer
+ }
+ streamErr = retryErr
+ }
+ }
+
+ status := http.StatusInternalServerError
+ if se, ok := streamErr.(interface{ StatusCode() int }); ok && se != nil {
+ if code := se.StatusCode(); code > 0 {
+ status = code
+ }
+ }
+ var addon http.Header
+ if he, ok := streamErr.(interface{ Headers() http.Header }); ok && he != nil {
+ if hdr := he.Headers(); hdr != nil {
+ addon = hdr.Clone()
+ }
+ }
+ _ = sendErr(&interfaces.ErrorMessage{StatusCode: status, Error: streamErr, Addon: addon})
+ return
+ }
+ if len(chunk.Payload) > 0 {
+ sentPayload = true
+ if okSendData := sendData(cloneBytes(chunk.Payload)); !okSendData {
+ return
+ }
+ }
+ }
+ }
+ }()
+ return dataChan, errChan
+}
+
+func statusFromError(err error) int {
+ if err == nil {
+ return 0
+ }
+ if se, ok := err.(interface{ StatusCode() int }); ok && se != nil {
+ if code := se.StatusCode(); code > 0 {
+ return code
+ }
+ }
+ return 0
+}
+
+func (h *BaseAPIHandler) getRequestDetails(modelName string) (providers []string, normalizedModel string, err *interfaces.ErrorMessage) {
+ resolvedModelName := modelName
+ initialSuffix := thinking.ParseSuffix(modelName)
+ if initialSuffix.ModelName == "auto" {
+ resolvedBase := util.ResolveAutoModel(initialSuffix.ModelName)
+ if initialSuffix.HasSuffix {
+ resolvedModelName = fmt.Sprintf("%s(%s)", resolvedBase, initialSuffix.RawSuffix)
+ } else {
+ resolvedModelName = resolvedBase
+ }
+ } else {
+ resolvedModelName = util.ResolveAutoModel(modelName)
+ }
+
+ parsed := thinking.ParseSuffix(resolvedModelName)
+ baseModel := strings.TrimSpace(parsed.ModelName)
+
+ providers = util.GetProviderName(baseModel)
+ // Fallback: if baseModel has no provider but differs from resolvedModelName,
+ // try using the full model name. This handles edge cases where custom models
+ // may be registered with their full suffixed name (e.g., "my-model(8192)").
+ // Evaluated in Story 11.8: This fallback is intentionally preserved to support
+ // custom model registrations that include thinking suffixes.
+ if len(providers) == 0 && baseModel != resolvedModelName {
+ providers = util.GetProviderName(resolvedModelName)
+ }
+
+ if len(providers) == 0 {
+ return nil, "", &interfaces.ErrorMessage{StatusCode: http.StatusBadRequest, Error: fmt.Errorf("unknown provider for model %s", modelName)}
+ }
+
+ // The thinking suffix is preserved in the model name itself, so no
+ // metadata-based configuration passing is needed.
+ return providers, resolvedModelName, nil
+}
+
+func cloneBytes(src []byte) []byte {
+ if len(src) == 0 {
+ return nil
+ }
+ dst := make([]byte, len(src))
+ copy(dst, src)
+ return dst
+}
+
+func cloneMetadata(src map[string]any) map[string]any {
+ if len(src) == 0 {
+ return nil
+ }
+ dst := make(map[string]any, len(src))
+ for k, v := range src {
+ dst[k] = v
+ }
+ return dst
+}
+
+// WriteErrorResponse writes an error message to the response writer using the HTTP status embedded in the message.
+func (h *BaseAPIHandler) WriteErrorResponse(c *gin.Context, msg *interfaces.ErrorMessage) {
+ status := http.StatusInternalServerError
+ if msg != nil && msg.StatusCode > 0 {
+ status = msg.StatusCode
+ }
+ if msg != nil && msg.Addon != nil {
+ for key, values := range msg.Addon {
+ if len(values) == 0 {
+ continue
+ }
+ c.Writer.Header().Del(key)
+ for _, value := range values {
+ c.Writer.Header().Add(key, value)
+ }
+ }
+ }
+
+ errText := http.StatusText(status)
+ if msg != nil && msg.Error != nil {
+ if v := strings.TrimSpace(msg.Error.Error()); v != "" {
+ errText = v
+ }
+ }
+
+ body := BuildErrorResponseBody(status, errText)
+ // Append first to preserve upstream response logs, then drop duplicate payloads if already recorded.
+ var previous []byte
+ if existing, exists := c.Get("API_RESPONSE"); exists {
+ if existingBytes, ok := existing.([]byte); ok && len(existingBytes) > 0 {
+ previous = bytes.Clone(existingBytes)
+ }
+ }
+ appendAPIResponse(c, body)
+ trimmedErrText := strings.TrimSpace(errText)
+ trimmedBody := bytes.TrimSpace(body)
+ if len(previous) > 0 {
+ if (trimmedErrText != "" && bytes.Contains(previous, []byte(trimmedErrText))) ||
+ (len(trimmedBody) > 0 && bytes.Contains(previous, trimmedBody)) {
+ c.Set("API_RESPONSE", previous)
+ }
+ }
+
+ if !c.Writer.Written() {
+ c.Writer.Header().Set("Content-Type", "application/json")
+ }
+ c.Status(status)
+ _, _ = c.Writer.Write(body)
+}
+
+func (h *BaseAPIHandler) LoggingAPIResponseError(ctx context.Context, err *interfaces.ErrorMessage) {
+ if h.Cfg.RequestLog {
+ if ginContext, ok := ctx.Value("gin").(*gin.Context); ok {
+ if apiResponseErrors, isExist := ginContext.Get("API_RESPONSE_ERROR"); isExist {
+ if slicesAPIResponseError, isOk := apiResponseErrors.([]*interfaces.ErrorMessage); isOk {
+ slicesAPIResponseError = append(slicesAPIResponseError, err)
+ ginContext.Set("API_RESPONSE_ERROR", slicesAPIResponseError)
+ }
+ } else {
+ // Create new response data entry
+ ginContext.Set("API_RESPONSE_ERROR", []*interfaces.ErrorMessage{err})
+ }
+ }
+ }
+}
+
+// APIHandlerCancelFunc is a function type for canceling an API handler's context.
+// It can optionally accept parameters, which are used for logging the response.
+type APIHandlerCancelFunc func(params ...interface{})
diff --git a/sdk/api/handlers/handlers_request_details_test.go b/sdk/api/handlers/handlers_request_details_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..b0f6b1326203cb42fc0ef96c759d71817c675b2c
--- /dev/null
+++ b/sdk/api/handlers/handlers_request_details_test.go
@@ -0,0 +1,118 @@
+package handlers
+
+import (
+ "reflect"
+ "testing"
+ "time"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
+)
+
+func TestGetRequestDetails_PreservesSuffix(t *testing.T) {
+ modelRegistry := registry.GetGlobalRegistry()
+ now := time.Now().Unix()
+
+ modelRegistry.RegisterClient("test-request-details-gemini", "gemini", []*registry.ModelInfo{
+ {ID: "gemini-2.5-pro", Created: now + 30},
+ {ID: "gemini-2.5-flash", Created: now + 25},
+ })
+ modelRegistry.RegisterClient("test-request-details-openai", "openai", []*registry.ModelInfo{
+ {ID: "gpt-5.2", Created: now + 20},
+ })
+ modelRegistry.RegisterClient("test-request-details-claude", "claude", []*registry.ModelInfo{
+ {ID: "claude-sonnet-4-5", Created: now + 5},
+ })
+
+ // Ensure cleanup of all test registrations.
+ clientIDs := []string{
+ "test-request-details-gemini",
+ "test-request-details-openai",
+ "test-request-details-claude",
+ }
+ for _, clientID := range clientIDs {
+ id := clientID
+ t.Cleanup(func() {
+ modelRegistry.UnregisterClient(id)
+ })
+ }
+
+ handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, coreauth.NewManager(nil, nil, nil))
+
+ tests := []struct {
+ name string
+ inputModel string
+ wantProviders []string
+ wantModel string
+ wantErr bool
+ }{
+ {
+ name: "numeric suffix preserved",
+ inputModel: "gemini-2.5-pro(8192)",
+ wantProviders: []string{"gemini"},
+ wantModel: "gemini-2.5-pro(8192)",
+ wantErr: false,
+ },
+ {
+ name: "level suffix preserved",
+ inputModel: "gpt-5.2(high)",
+ wantProviders: []string{"openai"},
+ wantModel: "gpt-5.2(high)",
+ wantErr: false,
+ },
+ {
+ name: "no suffix unchanged",
+ inputModel: "claude-sonnet-4-5",
+ wantProviders: []string{"claude"},
+ wantModel: "claude-sonnet-4-5",
+ wantErr: false,
+ },
+ {
+ name: "unknown model with suffix",
+ inputModel: "unknown-model(8192)",
+ wantProviders: nil,
+ wantModel: "",
+ wantErr: true,
+ },
+ {
+ name: "auto suffix resolved",
+ inputModel: "auto(high)",
+ wantProviders: []string{"gemini"},
+ wantModel: "gemini-2.5-pro(high)",
+ wantErr: false,
+ },
+ {
+ name: "special suffix none preserved",
+ inputModel: "gemini-2.5-flash(none)",
+ wantProviders: []string{"gemini"},
+ wantModel: "gemini-2.5-flash(none)",
+ wantErr: false,
+ },
+ {
+ name: "special suffix auto preserved",
+ inputModel: "claude-sonnet-4-5(auto)",
+ wantProviders: []string{"claude"},
+ wantModel: "claude-sonnet-4-5(auto)",
+ wantErr: false,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ providers, model, errMsg := handler.getRequestDetails(tt.inputModel)
+ if (errMsg != nil) != tt.wantErr {
+ t.Fatalf("getRequestDetails() error = %v, wantErr %v", errMsg, tt.wantErr)
+ }
+ if errMsg != nil {
+ return
+ }
+ if !reflect.DeepEqual(providers, tt.wantProviders) {
+ t.Fatalf("getRequestDetails() providers = %v, want %v", providers, tt.wantProviders)
+ }
+ if model != tt.wantModel {
+ t.Fatalf("getRequestDetails() model = %v, want %v", model, tt.wantModel)
+ }
+ })
+ }
+}
diff --git a/sdk/api/handlers/handlers_stream_bootstrap_test.go b/sdk/api/handlers/handlers_stream_bootstrap_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..3851746d4f26a319487290eb2af597a5b48613b2
--- /dev/null
+++ b/sdk/api/handlers/handlers_stream_bootstrap_test.go
@@ -0,0 +1,132 @@
+package handlers
+
+import (
+ "context"
+ "net/http"
+ "sync"
+ "testing"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ coreexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
+ sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
+)
+
+type failOnceStreamExecutor struct {
+ mu sync.Mutex
+ calls int
+}
+
+func (e *failOnceStreamExecutor) Identifier() string { return "codex" }
+
+func (e *failOnceStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
+ return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "Execute not implemented"}
+}
+
+func (e *failOnceStreamExecutor) ExecuteStream(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (<-chan coreexecutor.StreamChunk, error) {
+ e.mu.Lock()
+ e.calls++
+ call := e.calls
+ e.mu.Unlock()
+
+ ch := make(chan coreexecutor.StreamChunk, 1)
+ if call == 1 {
+ ch <- coreexecutor.StreamChunk{
+ Err: &coreauth.Error{
+ Code: "unauthorized",
+ Message: "unauthorized",
+ Retryable: false,
+ HTTPStatus: http.StatusUnauthorized,
+ },
+ }
+ close(ch)
+ return ch, nil
+ }
+
+ ch <- coreexecutor.StreamChunk{Payload: []byte("ok")}
+ close(ch)
+ return ch, nil
+}
+
+func (e *failOnceStreamExecutor) Refresh(ctx context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) {
+ return auth, nil
+}
+
+func (e *failOnceStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
+ return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"}
+}
+
+func (e *failOnceStreamExecutor) HttpRequest(ctx context.Context, auth *coreauth.Auth, req *http.Request) (*http.Response, error) {
+ return nil, &coreauth.Error{
+ Code: "not_implemented",
+ Message: "HttpRequest not implemented",
+ HTTPStatus: http.StatusNotImplemented,
+ }
+}
+
+func (e *failOnceStreamExecutor) Calls() int {
+ e.mu.Lock()
+ defer e.mu.Unlock()
+ return e.calls
+}
+
+func TestExecuteStreamWithAuthManager_RetriesBeforeFirstByte(t *testing.T) {
+ executor := &failOnceStreamExecutor{}
+ manager := coreauth.NewManager(nil, nil, nil)
+ manager.RegisterExecutor(executor)
+
+ auth1 := &coreauth.Auth{
+ ID: "auth1",
+ Provider: "codex",
+ Status: coreauth.StatusActive,
+ Metadata: map[string]any{"email": "test1@example.com"},
+ }
+ if _, err := manager.Register(context.Background(), auth1); err != nil {
+ t.Fatalf("manager.Register(auth1): %v", err)
+ }
+
+ auth2 := &coreauth.Auth{
+ ID: "auth2",
+ Provider: "codex",
+ Status: coreauth.StatusActive,
+ Metadata: map[string]any{"email": "test2@example.com"},
+ }
+ if _, err := manager.Register(context.Background(), auth2); err != nil {
+ t.Fatalf("manager.Register(auth2): %v", err)
+ }
+
+ registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}})
+ registry.GetGlobalRegistry().RegisterClient(auth2.ID, auth2.Provider, []*registry.ModelInfo{{ID: "test-model"}})
+ t.Cleanup(func() {
+ registry.GetGlobalRegistry().UnregisterClient(auth1.ID)
+ registry.GetGlobalRegistry().UnregisterClient(auth2.ID)
+ })
+
+ handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{
+ Streaming: sdkconfig.StreamingConfig{
+ BootstrapRetries: 1,
+ },
+ }, manager)
+ dataChan, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "test-model", []byte(`{"model":"test-model"}`), "")
+ if dataChan == nil || errChan == nil {
+ t.Fatalf("expected non-nil channels")
+ }
+
+ var got []byte
+ for chunk := range dataChan {
+ got = append(got, chunk...)
+ }
+
+ for msg := range errChan {
+ if msg != nil {
+ t.Fatalf("unexpected error: %+v", msg)
+ }
+ }
+
+ if string(got) != "ok" {
+ t.Fatalf("expected payload ok, got %q", string(got))
+ }
+ if executor.Calls() != 2 {
+ t.Fatalf("expected 2 stream attempts, got %d", executor.Calls())
+ }
+}
diff --git a/sdk/api/handlers/openai/openai_handlers.go b/sdk/api/handlers/openai/openai_handlers.go
new file mode 100644
index 0000000000000000000000000000000000000000..09471ce1d695eb32f56af2d2cf168bc7606d5237
--- /dev/null
+++ b/sdk/api/handlers/openai/openai_handlers.go
@@ -0,0 +1,672 @@
+// Package openai provides HTTP handlers for OpenAI API endpoints.
+// This package implements the OpenAI-compatible API interface, including model listing
+// and chat completion functionality. It supports both streaming and non-streaming responses,
+// and manages a pool of clients to interact with backend services.
+// The handlers translate OpenAI API requests to the appropriate backend format and
+// convert responses back to OpenAI-compatible format.
+package openai
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+ "net/http"
+ "sync"
+
+ "github.com/gin-gonic/gin"
+ . "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
+ responsesconverter "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/openai/openai/responses"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
+ "github.com/tidwall/gjson"
+ "github.com/tidwall/sjson"
+)
+
+// OpenAIAPIHandler contains the handlers for OpenAI API endpoints.
+// It holds a pool of clients to interact with the backend service.
+type OpenAIAPIHandler struct {
+ *handlers.BaseAPIHandler
+}
+
+// NewOpenAIAPIHandler creates a new OpenAI API handlers instance.
+// It takes an BaseAPIHandler instance as input and returns an OpenAIAPIHandler.
+//
+// Parameters:
+// - apiHandlers: The base API handlers instance
+//
+// Returns:
+// - *OpenAIAPIHandler: A new OpenAI API handlers instance
+func NewOpenAIAPIHandler(apiHandlers *handlers.BaseAPIHandler) *OpenAIAPIHandler {
+ return &OpenAIAPIHandler{
+ BaseAPIHandler: apiHandlers,
+ }
+}
+
+// HandlerType returns the identifier for this handler implementation.
+func (h *OpenAIAPIHandler) HandlerType() string {
+ return OpenAI
+}
+
+// Models returns the OpenAI-compatible model metadata supported by this handler.
+func (h *OpenAIAPIHandler) Models() []map[string]any {
+ // Get dynamic models from the global registry
+ modelRegistry := registry.GetGlobalRegistry()
+ return modelRegistry.GetAvailableModels("openai")
+}
+
+// OpenAIModels handles the /v1/models endpoint.
+// It returns a list of available AI models with their capabilities
+// and specifications in OpenAI-compatible format.
+func (h *OpenAIAPIHandler) OpenAIModels(c *gin.Context) {
+ // Get all available models
+ allModels := h.Models()
+
+ // Filter to only include the 4 required fields: id, object, created, owned_by
+ filteredModels := make([]map[string]any, len(allModels))
+ for i, model := range allModels {
+ filteredModel := map[string]any{
+ "id": model["id"],
+ "object": model["object"],
+ }
+
+ // Add created field if it exists
+ if created, exists := model["created"]; exists {
+ filteredModel["created"] = created
+ }
+
+ // Add owned_by field if it exists
+ if ownedBy, exists := model["owned_by"]; exists {
+ filteredModel["owned_by"] = ownedBy
+ }
+
+ filteredModels[i] = filteredModel
+ }
+
+ c.JSON(http.StatusOK, gin.H{
+ "object": "list",
+ "data": filteredModels,
+ })
+}
+
+// ChatCompletions handles the /v1/chat/completions endpoint.
+// It determines whether the request is for a streaming or non-streaming response
+// and calls the appropriate handler based on the model provider.
+//
+// Parameters:
+// - c: The Gin context containing the HTTP request and response
+func (h *OpenAIAPIHandler) ChatCompletions(c *gin.Context) {
+ rawJSON, err := c.GetRawData()
+ // If data retrieval fails, return a 400 Bad Request error.
+ if err != nil {
+ c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: fmt.Sprintf("Invalid request: %v", err),
+ Type: "invalid_request_error",
+ },
+ })
+ return
+ }
+
+ // Check if the client requested a streaming response.
+ streamResult := gjson.GetBytes(rawJSON, "stream")
+ stream := streamResult.Type == gjson.True
+
+ // Some clients send OpenAI Responses-format payloads to /v1/chat/completions.
+ // Convert them to Chat Completions so downstream translators preserve tool metadata.
+ if shouldTreatAsResponsesFormat(rawJSON) {
+ modelName := gjson.GetBytes(rawJSON, "model").String()
+ rawJSON = responsesconverter.ConvertOpenAIResponsesRequestToOpenAIChatCompletions(modelName, rawJSON, stream)
+ stream = gjson.GetBytes(rawJSON, "stream").Bool()
+ }
+
+ if stream {
+ h.handleStreamingResponse(c, rawJSON)
+ } else {
+ h.handleNonStreamingResponse(c, rawJSON)
+ }
+
+}
+
+// shouldTreatAsResponsesFormat detects OpenAI Responses-style payloads that are
+// accidentally sent to the Chat Completions endpoint.
+func shouldTreatAsResponsesFormat(rawJSON []byte) bool {
+ if gjson.GetBytes(rawJSON, "messages").Exists() {
+ return false
+ }
+ if gjson.GetBytes(rawJSON, "input").Exists() {
+ return true
+ }
+ if gjson.GetBytes(rawJSON, "instructions").Exists() {
+ return true
+ }
+ return false
+}
+
+// Completions handles the /v1/completions endpoint.
+// It determines whether the request is for a streaming or non-streaming response
+// and calls the appropriate handler based on the model provider.
+// This endpoint follows the OpenAI completions API specification.
+//
+// Parameters:
+// - c: The Gin context containing the HTTP request and response
+func (h *OpenAIAPIHandler) Completions(c *gin.Context) {
+ rawJSON, err := c.GetRawData()
+ // If data retrieval fails, return a 400 Bad Request error.
+ if err != nil {
+ c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: fmt.Sprintf("Invalid request: %v", err),
+ Type: "invalid_request_error",
+ },
+ })
+ return
+ }
+
+ // Check if the client requested a streaming response.
+ streamResult := gjson.GetBytes(rawJSON, "stream")
+ if streamResult.Type == gjson.True {
+ h.handleCompletionsStreamingResponse(c, rawJSON)
+ } else {
+ h.handleCompletionsNonStreamingResponse(c, rawJSON)
+ }
+
+}
+
+// convertCompletionsRequestToChatCompletions converts OpenAI completions API request to chat completions format.
+// This allows the completions endpoint to use the existing chat completions infrastructure.
+//
+// Parameters:
+// - rawJSON: The raw JSON bytes of the completions request
+//
+// Returns:
+// - []byte: The converted chat completions request
+func convertCompletionsRequestToChatCompletions(rawJSON []byte) []byte {
+ root := gjson.ParseBytes(rawJSON)
+
+ // Extract prompt from completions request
+ prompt := root.Get("prompt").String()
+ if prompt == "" {
+ prompt = "Complete this:"
+ }
+
+ // Create chat completions structure
+ out := `{"model":"","messages":[{"role":"user","content":""}]}`
+
+ // Set model
+ if model := root.Get("model"); model.Exists() {
+ out, _ = sjson.Set(out, "model", model.String())
+ }
+
+ // Set the prompt as user message content
+ out, _ = sjson.Set(out, "messages.0.content", prompt)
+
+ // Copy other parameters from completions to chat completions
+ if maxTokens := root.Get("max_tokens"); maxTokens.Exists() {
+ out, _ = sjson.Set(out, "max_tokens", maxTokens.Int())
+ }
+
+ if temperature := root.Get("temperature"); temperature.Exists() {
+ out, _ = sjson.Set(out, "temperature", temperature.Float())
+ }
+
+ if topP := root.Get("top_p"); topP.Exists() {
+ out, _ = sjson.Set(out, "top_p", topP.Float())
+ }
+
+ if frequencyPenalty := root.Get("frequency_penalty"); frequencyPenalty.Exists() {
+ out, _ = sjson.Set(out, "frequency_penalty", frequencyPenalty.Float())
+ }
+
+ if presencePenalty := root.Get("presence_penalty"); presencePenalty.Exists() {
+ out, _ = sjson.Set(out, "presence_penalty", presencePenalty.Float())
+ }
+
+ if stop := root.Get("stop"); stop.Exists() {
+ out, _ = sjson.SetRaw(out, "stop", stop.Raw)
+ }
+
+ if stream := root.Get("stream"); stream.Exists() {
+ out, _ = sjson.Set(out, "stream", stream.Bool())
+ }
+
+ if logprobs := root.Get("logprobs"); logprobs.Exists() {
+ out, _ = sjson.Set(out, "logprobs", logprobs.Bool())
+ }
+
+ if topLogprobs := root.Get("top_logprobs"); topLogprobs.Exists() {
+ out, _ = sjson.Set(out, "top_logprobs", topLogprobs.Int())
+ }
+
+ if echo := root.Get("echo"); echo.Exists() {
+ out, _ = sjson.Set(out, "echo", echo.Bool())
+ }
+
+ return []byte(out)
+}
+
+// convertChatCompletionsResponseToCompletions converts chat completions API response back to completions format.
+// This ensures the completions endpoint returns data in the expected format.
+//
+// Parameters:
+// - rawJSON: The raw JSON bytes of the chat completions response
+//
+// Returns:
+// - []byte: The converted completions response
+func convertChatCompletionsResponseToCompletions(rawJSON []byte) []byte {
+ root := gjson.ParseBytes(rawJSON)
+
+ // Base completions response structure
+ out := `{"id":"","object":"text_completion","created":0,"model":"","choices":[]}`
+
+ // Copy basic fields
+ if id := root.Get("id"); id.Exists() {
+ out, _ = sjson.Set(out, "id", id.String())
+ }
+
+ if created := root.Get("created"); created.Exists() {
+ out, _ = sjson.Set(out, "created", created.Int())
+ }
+
+ if model := root.Get("model"); model.Exists() {
+ out, _ = sjson.Set(out, "model", model.String())
+ }
+
+ if usage := root.Get("usage"); usage.Exists() {
+ out, _ = sjson.SetRaw(out, "usage", usage.Raw)
+ }
+
+ // Convert choices from chat completions to completions format
+ var choices []interface{}
+ if chatChoices := root.Get("choices"); chatChoices.Exists() && chatChoices.IsArray() {
+ chatChoices.ForEach(func(_, choice gjson.Result) bool {
+ completionsChoice := map[string]interface{}{
+ "index": choice.Get("index").Int(),
+ }
+
+ // Extract text content from message.content
+ if message := choice.Get("message"); message.Exists() {
+ if content := message.Get("content"); content.Exists() {
+ completionsChoice["text"] = content.String()
+ }
+ } else if delta := choice.Get("delta"); delta.Exists() {
+ // For streaming responses, use delta.content
+ if content := delta.Get("content"); content.Exists() {
+ completionsChoice["text"] = content.String()
+ }
+ }
+
+ // Copy finish_reason
+ if finishReason := choice.Get("finish_reason"); finishReason.Exists() {
+ completionsChoice["finish_reason"] = finishReason.String()
+ }
+
+ // Copy logprobs if present
+ if logprobs := choice.Get("logprobs"); logprobs.Exists() {
+ completionsChoice["logprobs"] = logprobs.Value()
+ }
+
+ choices = append(choices, completionsChoice)
+ return true
+ })
+ }
+
+ if len(choices) > 0 {
+ choicesJSON, _ := json.Marshal(choices)
+ out, _ = sjson.SetRaw(out, "choices", string(choicesJSON))
+ }
+
+ return []byte(out)
+}
+
+// convertChatCompletionsStreamChunkToCompletions converts a streaming chat completions chunk to completions format.
+// This handles the real-time conversion of streaming response chunks and filters out empty text responses.
+//
+// Parameters:
+// - chunkData: The raw JSON bytes of a single chat completions stream chunk
+//
+// Returns:
+// - []byte: The converted completions stream chunk, or nil if should be filtered out
+func convertChatCompletionsStreamChunkToCompletions(chunkData []byte) []byte {
+ root := gjson.ParseBytes(chunkData)
+
+ // Check if this chunk has any meaningful content
+ hasContent := false
+ if chatChoices := root.Get("choices"); chatChoices.Exists() && chatChoices.IsArray() {
+ chatChoices.ForEach(func(_, choice gjson.Result) bool {
+ // Check if delta has content or finish_reason
+ if delta := choice.Get("delta"); delta.Exists() {
+ if content := delta.Get("content"); content.Exists() && content.String() != "" {
+ hasContent = true
+ return false // Break out of forEach
+ }
+ }
+ // Also check for finish_reason to ensure we don't skip final chunks
+ if finishReason := choice.Get("finish_reason"); finishReason.Exists() && finishReason.String() != "" && finishReason.String() != "null" {
+ hasContent = true
+ return false // Break out of forEach
+ }
+ return true
+ })
+ }
+
+ // If no meaningful content, return nil to indicate this chunk should be skipped
+ if !hasContent {
+ return nil
+ }
+
+ // Base completions stream response structure
+ out := `{"id":"","object":"text_completion","created":0,"model":"","choices":[]}`
+
+ // Copy basic fields
+ if id := root.Get("id"); id.Exists() {
+ out, _ = sjson.Set(out, "id", id.String())
+ }
+
+ if created := root.Get("created"); created.Exists() {
+ out, _ = sjson.Set(out, "created", created.Int())
+ }
+
+ if model := root.Get("model"); model.Exists() {
+ out, _ = sjson.Set(out, "model", model.String())
+ }
+
+ // Convert choices from chat completions delta to completions format
+ var choices []interface{}
+ if chatChoices := root.Get("choices"); chatChoices.Exists() && chatChoices.IsArray() {
+ chatChoices.ForEach(func(_, choice gjson.Result) bool {
+ completionsChoice := map[string]interface{}{
+ "index": choice.Get("index").Int(),
+ }
+
+ // Extract text content from delta.content
+ if delta := choice.Get("delta"); delta.Exists() {
+ if content := delta.Get("content"); content.Exists() && content.String() != "" {
+ completionsChoice["text"] = content.String()
+ } else {
+ completionsChoice["text"] = ""
+ }
+ } else {
+ completionsChoice["text"] = ""
+ }
+
+ // Copy finish_reason
+ if finishReason := choice.Get("finish_reason"); finishReason.Exists() && finishReason.String() != "null" {
+ completionsChoice["finish_reason"] = finishReason.String()
+ }
+
+ // Copy logprobs if present
+ if logprobs := choice.Get("logprobs"); logprobs.Exists() {
+ completionsChoice["logprobs"] = logprobs.Value()
+ }
+
+ choices = append(choices, completionsChoice)
+ return true
+ })
+ }
+
+ if len(choices) > 0 {
+ choicesJSON, _ := json.Marshal(choices)
+ out, _ = sjson.SetRaw(out, "choices", string(choicesJSON))
+ }
+
+ return []byte(out)
+}
+
+// handleNonStreamingResponse handles non-streaming chat completion responses
+// for Gemini models. It selects a client from the pool, sends the request, and
+// aggregates the response before sending it back to the client in OpenAI format.
+//
+// Parameters:
+// - c: The Gin context containing the HTTP request and response
+// - rawJSON: The raw JSON bytes of the OpenAI-compatible request
+func (h *OpenAIAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSON []byte) {
+ c.Header("Content-Type", "application/json")
+
+ modelName := gjson.GetBytes(rawJSON, "model").String()
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+ resp, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, h.GetAlt(c))
+ if errMsg != nil {
+ h.WriteErrorResponse(c, errMsg)
+ cliCancel(errMsg.Error)
+ return
+ }
+ _, _ = c.Writer.Write(resp)
+ cliCancel()
+}
+
+// handleStreamingResponse handles streaming responses for Gemini models.
+// It establishes a streaming connection with the backend service and forwards
+// the response chunks to the client in real-time using Server-Sent Events.
+//
+// Parameters:
+// - c: The Gin context containing the HTTP request and response
+// - rawJSON: The raw JSON bytes of the OpenAI-compatible request
+func (h *OpenAIAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON []byte) {
+ // Get the http.Flusher interface to manually flush the response.
+ flusher, ok := c.Writer.(http.Flusher)
+ if !ok {
+ c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: "Streaming not supported",
+ Type: "server_error",
+ },
+ })
+ return
+ }
+
+ modelName := gjson.GetBytes(rawJSON, "model").String()
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+ dataChan, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, h.GetAlt(c))
+
+ setSSEHeaders := func() {
+ c.Header("Content-Type", "text/event-stream")
+ c.Header("Cache-Control", "no-cache")
+ c.Header("Connection", "keep-alive")
+ c.Header("Access-Control-Allow-Origin", "*")
+ }
+
+ // Peek at the first chunk to determine success or failure before setting headers
+ for {
+ select {
+ case <-c.Request.Context().Done():
+ cliCancel(c.Request.Context().Err())
+ return
+ case errMsg, ok := <-errChan:
+ if !ok {
+ // Err channel closed cleanly; wait for data channel.
+ errChan = nil
+ continue
+ }
+ // Upstream failed immediately. Return proper error status and JSON.
+ h.WriteErrorResponse(c, errMsg)
+ if errMsg != nil {
+ cliCancel(errMsg.Error)
+ } else {
+ cliCancel(nil)
+ }
+ return
+ case chunk, ok := <-dataChan:
+ if !ok {
+ // Stream closed without data? Send DONE or just headers.
+ setSSEHeaders()
+ _, _ = fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
+ flusher.Flush()
+ cliCancel(nil)
+ return
+ }
+
+ // Success! Commit to streaming headers.
+ setSSEHeaders()
+
+ _, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(chunk))
+ flusher.Flush()
+
+ // Continue streaming the rest
+ h.handleStreamResult(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan)
+ return
+ }
+ }
+}
+
+// handleCompletionsNonStreamingResponse handles non-streaming completions responses.
+// It converts completions request to chat completions format, sends to backend,
+// then converts the response back to completions format before sending to client.
+//
+// Parameters:
+// - c: The Gin context containing the HTTP request and response
+// - rawJSON: The raw JSON bytes of the OpenAI-compatible completions request
+func (h *OpenAIAPIHandler) handleCompletionsNonStreamingResponse(c *gin.Context, rawJSON []byte) {
+ c.Header("Content-Type", "application/json")
+
+ // Convert completions request to chat completions format
+ chatCompletionsJSON := convertCompletionsRequestToChatCompletions(rawJSON)
+
+ modelName := gjson.GetBytes(chatCompletionsJSON, "model").String()
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+ stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
+ resp, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, chatCompletionsJSON, "")
+ stopKeepAlive()
+ if errMsg != nil {
+ h.WriteErrorResponse(c, errMsg)
+ cliCancel(errMsg.Error)
+ return
+ }
+ completionsResp := convertChatCompletionsResponseToCompletions(resp)
+ _, _ = c.Writer.Write(completionsResp)
+ cliCancel()
+}
+
+// handleCompletionsStreamingResponse handles streaming completions responses.
+// It converts completions request to chat completions format, streams from backend,
+// then converts each response chunk back to completions format before sending to client.
+//
+// Parameters:
+// - c: The Gin context containing the HTTP request and response
+// - rawJSON: The raw JSON bytes of the OpenAI-compatible completions request
+func (h *OpenAIAPIHandler) handleCompletionsStreamingResponse(c *gin.Context, rawJSON []byte) {
+ // Get the http.Flusher interface to manually flush the response.
+ flusher, ok := c.Writer.(http.Flusher)
+ if !ok {
+ c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: "Streaming not supported",
+ Type: "server_error",
+ },
+ })
+ return
+ }
+
+ // Convert completions request to chat completions format
+ chatCompletionsJSON := convertCompletionsRequestToChatCompletions(rawJSON)
+
+ modelName := gjson.GetBytes(chatCompletionsJSON, "model").String()
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+ dataChan, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, chatCompletionsJSON, "")
+
+ setSSEHeaders := func() {
+ c.Header("Content-Type", "text/event-stream")
+ c.Header("Cache-Control", "no-cache")
+ c.Header("Connection", "keep-alive")
+ c.Header("Access-Control-Allow-Origin", "*")
+ }
+
+ // Peek at the first chunk
+ for {
+ select {
+ case <-c.Request.Context().Done():
+ cliCancel(c.Request.Context().Err())
+ return
+ case errMsg, ok := <-errChan:
+ if !ok {
+ // Err channel closed cleanly; wait for data channel.
+ errChan = nil
+ continue
+ }
+ h.WriteErrorResponse(c, errMsg)
+ if errMsg != nil {
+ cliCancel(errMsg.Error)
+ } else {
+ cliCancel(nil)
+ }
+ return
+ case chunk, ok := <-dataChan:
+ if !ok {
+ setSSEHeaders()
+ _, _ = fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
+ flusher.Flush()
+ cliCancel(nil)
+ return
+ }
+
+ // Success! Set headers.
+ setSSEHeaders()
+
+ // Write the first chunk
+ converted := convertChatCompletionsStreamChunkToCompletions(chunk)
+ if converted != nil {
+ _, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(converted))
+ flusher.Flush()
+ }
+
+ done := make(chan struct{})
+ var doneOnce sync.Once
+ stop := func() { doneOnce.Do(func() { close(done) }) }
+
+ convertedChan := make(chan []byte)
+ go func() {
+ defer close(convertedChan)
+ for {
+ select {
+ case <-done:
+ return
+ case chunk, ok := <-dataChan:
+ if !ok {
+ return
+ }
+ converted := convertChatCompletionsStreamChunkToCompletions(chunk)
+ if converted == nil {
+ continue
+ }
+ select {
+ case <-done:
+ return
+ case convertedChan <- converted:
+ }
+ }
+ }
+ }()
+
+ h.handleStreamResult(c, flusher, func(err error) {
+ stop()
+ cliCancel(err)
+ }, convertedChan, errChan)
+ return
+ }
+ }
+}
+func (h *OpenAIAPIHandler) handleStreamResult(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
+ h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
+ WriteChunk: func(chunk []byte) {
+ _, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(chunk))
+ },
+ WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
+ if errMsg == nil {
+ return
+ }
+ status := http.StatusInternalServerError
+ if errMsg.StatusCode > 0 {
+ status = errMsg.StatusCode
+ }
+ errText := http.StatusText(status)
+ if errMsg.Error != nil && errMsg.Error.Error() != "" {
+ errText = errMsg.Error.Error()
+ }
+ body := handlers.BuildErrorResponseBody(status, errText)
+ _, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(body))
+ },
+ WriteDone: func() {
+ _, _ = fmt.Fprint(c.Writer, "data: [DONE]\n\n")
+ },
+ })
+}
diff --git a/sdk/api/handlers/openai/openai_responses_handlers.go b/sdk/api/handlers/openai/openai_responses_handlers.go
new file mode 100644
index 0000000000000000000000000000000000000000..31099f818a2bee5bdb71059b0bfe6353c32a8940
--- /dev/null
+++ b/sdk/api/handlers/openai/openai_responses_handlers.go
@@ -0,0 +1,227 @@
+// Package openai provides HTTP handlers for OpenAIResponses API endpoints.
+// This package implements the OpenAIResponses-compatible API interface, including model listing
+// and chat completion functionality. It supports both streaming and non-streaming responses,
+// and manages a pool of clients to interact with backend services.
+// The handlers translate OpenAIResponses API requests to the appropriate backend format and
+// convert responses back to OpenAIResponses-compatible format.
+package openai
+
+import (
+ "bytes"
+ "context"
+ "fmt"
+ "net/http"
+
+ "github.com/gin-gonic/gin"
+ . "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
+ "github.com/tidwall/gjson"
+)
+
+// OpenAIResponsesAPIHandler contains the handlers for OpenAIResponses API endpoints.
+// It holds a pool of clients to interact with the backend service.
+type OpenAIResponsesAPIHandler struct {
+ *handlers.BaseAPIHandler
+}
+
+// NewOpenAIResponsesAPIHandler creates a new OpenAIResponses API handlers instance.
+// It takes an BaseAPIHandler instance as input and returns an OpenAIResponsesAPIHandler.
+//
+// Parameters:
+// - apiHandlers: The base API handlers instance
+//
+// Returns:
+// - *OpenAIResponsesAPIHandler: A new OpenAIResponses API handlers instance
+func NewOpenAIResponsesAPIHandler(apiHandlers *handlers.BaseAPIHandler) *OpenAIResponsesAPIHandler {
+ return &OpenAIResponsesAPIHandler{
+ BaseAPIHandler: apiHandlers,
+ }
+}
+
+// HandlerType returns the identifier for this handler implementation.
+func (h *OpenAIResponsesAPIHandler) HandlerType() string {
+ return OpenaiResponse
+}
+
+// Models returns the OpenAIResponses-compatible model metadata supported by this handler.
+func (h *OpenAIResponsesAPIHandler) Models() []map[string]any {
+ // Get dynamic models from the global registry
+ modelRegistry := registry.GetGlobalRegistry()
+ return modelRegistry.GetAvailableModels("openai")
+}
+
+// OpenAIResponsesModels handles the /v1/models endpoint.
+// It returns a list of available AI models with their capabilities
+// and specifications in OpenAIResponses-compatible format.
+func (h *OpenAIResponsesAPIHandler) OpenAIResponsesModels(c *gin.Context) {
+ c.JSON(http.StatusOK, gin.H{
+ "object": "list",
+ "data": h.Models(),
+ })
+}
+
+// Responses handles the /v1/responses endpoint.
+// It determines whether the request is for a streaming or non-streaming response
+// and calls the appropriate handler based on the model provider.
+//
+// Parameters:
+// - c: The Gin context containing the HTTP request and response
+func (h *OpenAIResponsesAPIHandler) Responses(c *gin.Context) {
+ rawJSON, err := c.GetRawData()
+ // If data retrieval fails, return a 400 Bad Request error.
+ if err != nil {
+ c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: fmt.Sprintf("Invalid request: %v", err),
+ Type: "invalid_request_error",
+ },
+ })
+ return
+ }
+
+ // Check if the client requested a streaming response.
+ streamResult := gjson.GetBytes(rawJSON, "stream")
+ if streamResult.Type == gjson.True {
+ h.handleStreamingResponse(c, rawJSON)
+ } else {
+ h.handleNonStreamingResponse(c, rawJSON)
+ }
+
+}
+
+// handleNonStreamingResponse handles non-streaming chat completion responses
+// for Gemini models. It selects a client from the pool, sends the request, and
+// aggregates the response before sending it back to the client in OpenAIResponses format.
+//
+// Parameters:
+// - c: The Gin context containing the HTTP request and response
+// - rawJSON: The raw JSON bytes of the OpenAIResponses-compatible request
+func (h *OpenAIResponsesAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSON []byte) {
+ c.Header("Content-Type", "application/json")
+
+ modelName := gjson.GetBytes(rawJSON, "model").String()
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+ stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
+
+ resp, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "")
+ stopKeepAlive()
+ if errMsg != nil {
+ h.WriteErrorResponse(c, errMsg)
+ cliCancel(errMsg.Error)
+ return
+ }
+ _, _ = c.Writer.Write(resp)
+ cliCancel()
+}
+
+// handleStreamingResponse handles streaming responses for Gemini models.
+// It establishes a streaming connection with the backend service and forwards
+// the response chunks to the client in real-time using Server-Sent Events.
+//
+// Parameters:
+// - c: The Gin context containing the HTTP request and response
+// - rawJSON: The raw JSON bytes of the OpenAIResponses-compatible request
+func (h *OpenAIResponsesAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON []byte) {
+ // Get the http.Flusher interface to manually flush the response.
+ flusher, ok := c.Writer.(http.Flusher)
+ if !ok {
+ c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{
+ Error: handlers.ErrorDetail{
+ Message: "Streaming not supported",
+ Type: "server_error",
+ },
+ })
+ return
+ }
+
+ // New core execution path
+ modelName := gjson.GetBytes(rawJSON, "model").String()
+ cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+ dataChan, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "")
+
+ setSSEHeaders := func() {
+ c.Header("Content-Type", "text/event-stream")
+ c.Header("Cache-Control", "no-cache")
+ c.Header("Connection", "keep-alive")
+ c.Header("Access-Control-Allow-Origin", "*")
+ }
+
+ // Peek at the first chunk
+ for {
+ select {
+ case <-c.Request.Context().Done():
+ cliCancel(c.Request.Context().Err())
+ return
+ case errMsg, ok := <-errChan:
+ if !ok {
+ // Err channel closed cleanly; wait for data channel.
+ errChan = nil
+ continue
+ }
+ // Upstream failed immediately. Return proper error status and JSON.
+ h.WriteErrorResponse(c, errMsg)
+ if errMsg != nil {
+ cliCancel(errMsg.Error)
+ } else {
+ cliCancel(nil)
+ }
+ return
+ case chunk, ok := <-dataChan:
+ if !ok {
+ // Stream closed without data? Send headers and done.
+ setSSEHeaders()
+ _, _ = c.Writer.Write([]byte("\n"))
+ flusher.Flush()
+ cliCancel(nil)
+ return
+ }
+
+ // Success! Set headers.
+ setSSEHeaders()
+
+ // Write first chunk logic (matching forwardResponsesStream)
+ if bytes.HasPrefix(chunk, []byte("event:")) {
+ _, _ = c.Writer.Write([]byte("\n"))
+ }
+ _, _ = c.Writer.Write(chunk)
+ _, _ = c.Writer.Write([]byte("\n"))
+ flusher.Flush()
+
+ // Continue
+ h.forwardResponsesStream(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan)
+ return
+ }
+ }
+}
+
+func (h *OpenAIResponsesAPIHandler) forwardResponsesStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
+ h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
+ WriteChunk: func(chunk []byte) {
+ if bytes.HasPrefix(chunk, []byte("event:")) {
+ _, _ = c.Writer.Write([]byte("\n"))
+ }
+ _, _ = c.Writer.Write(chunk)
+ _, _ = c.Writer.Write([]byte("\n"))
+ },
+ WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
+ if errMsg == nil {
+ return
+ }
+ status := http.StatusInternalServerError
+ if errMsg.StatusCode > 0 {
+ status = errMsg.StatusCode
+ }
+ errText := http.StatusText(status)
+ if errMsg.Error != nil && errMsg.Error.Error() != "" {
+ errText = errMsg.Error.Error()
+ }
+ body := handlers.BuildErrorResponseBody(status, errText)
+ _, _ = fmt.Fprintf(c.Writer, "\nevent: error\ndata: %s\n\n", string(body))
+ },
+ WriteDone: func() {
+ _, _ = c.Writer.Write([]byte("\n"))
+ },
+ })
+}
diff --git a/sdk/api/handlers/stream_forwarder.go b/sdk/api/handlers/stream_forwarder.go
new file mode 100644
index 0000000000000000000000000000000000000000..401baca8fae38cde32d841e5b70f729ae3cca9dd
--- /dev/null
+++ b/sdk/api/handlers/stream_forwarder.go
@@ -0,0 +1,121 @@
+package handlers
+
+import (
+ "net/http"
+ "time"
+
+ "github.com/gin-gonic/gin"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
+)
+
+type StreamForwardOptions struct {
+ // KeepAliveInterval overrides the configured streaming keep-alive interval.
+ // If nil, the configured default is used. If set to <= 0, keep-alives are disabled.
+ KeepAliveInterval *time.Duration
+
+ // WriteChunk writes a single data chunk to the response body. It should not flush.
+ WriteChunk func(chunk []byte)
+
+ // WriteTerminalError writes an error payload to the response body when streaming fails
+ // after headers have already been committed. It should not flush.
+ WriteTerminalError func(errMsg *interfaces.ErrorMessage)
+
+ // WriteDone optionally writes a terminal marker when the upstream data channel closes
+ // without an error (e.g. OpenAI's `[DONE]`). It should not flush.
+ WriteDone func()
+
+ // WriteKeepAlive optionally writes a keep-alive heartbeat. It should not flush.
+ // When nil, a standard SSE comment heartbeat is used.
+ WriteKeepAlive func()
+}
+
+func (h *BaseAPIHandler) ForwardStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage, opts StreamForwardOptions) {
+ if c == nil {
+ return
+ }
+ if cancel == nil {
+ return
+ }
+
+ writeChunk := opts.WriteChunk
+ if writeChunk == nil {
+ writeChunk = func([]byte) {}
+ }
+
+ writeKeepAlive := opts.WriteKeepAlive
+ if writeKeepAlive == nil {
+ writeKeepAlive = func() {
+ _, _ = c.Writer.Write([]byte(": keep-alive\n\n"))
+ }
+ }
+
+ keepAliveInterval := StreamingKeepAliveInterval(h.Cfg)
+ if opts.KeepAliveInterval != nil {
+ keepAliveInterval = *opts.KeepAliveInterval
+ }
+ var keepAlive *time.Ticker
+ var keepAliveC <-chan time.Time
+ if keepAliveInterval > 0 {
+ keepAlive = time.NewTicker(keepAliveInterval)
+ defer keepAlive.Stop()
+ keepAliveC = keepAlive.C
+ }
+
+ var terminalErr *interfaces.ErrorMessage
+ for {
+ select {
+ case <-c.Request.Context().Done():
+ cancel(c.Request.Context().Err())
+ return
+ case chunk, ok := <-data:
+ if !ok {
+ // Prefer surfacing a terminal error if one is pending.
+ if terminalErr == nil {
+ select {
+ case errMsg, ok := <-errs:
+ if ok && errMsg != nil {
+ terminalErr = errMsg
+ }
+ default:
+ }
+ }
+ if terminalErr != nil {
+ if opts.WriteTerminalError != nil {
+ opts.WriteTerminalError(terminalErr)
+ }
+ flusher.Flush()
+ cancel(terminalErr.Error)
+ return
+ }
+ if opts.WriteDone != nil {
+ opts.WriteDone()
+ }
+ flusher.Flush()
+ cancel(nil)
+ return
+ }
+ writeChunk(chunk)
+ flusher.Flush()
+ case errMsg, ok := <-errs:
+ if !ok {
+ continue
+ }
+ if errMsg != nil {
+ terminalErr = errMsg
+ if opts.WriteTerminalError != nil {
+ opts.WriteTerminalError(errMsg)
+ flusher.Flush()
+ }
+ }
+ var execErr error
+ if errMsg != nil {
+ execErr = errMsg.Error
+ }
+ cancel(execErr)
+ return
+ case <-keepAliveC:
+ writeKeepAlive()
+ flusher.Flush()
+ }
+ }
+}
diff --git a/sdk/api/management.go b/sdk/api/management.go
new file mode 100644
index 0000000000000000000000000000000000000000..66af41ae91d105db1494b7ad27e8556c01ecc864
--- /dev/null
+++ b/sdk/api/management.go
@@ -0,0 +1,72 @@
+// Package api exposes helpers for embedding CLIProxyAPI.
+//
+// It wraps internal management handler types so external projects can integrate
+// management endpoints without importing internal packages.
+package api
+
+import (
+ "github.com/gin-gonic/gin"
+ internalmanagement "github.com/router-for-me/CLIProxyAPI/v6/internal/api/handlers/management"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
+)
+
+// ManagementTokenRequester exposes a limited subset of management endpoints for requesting tokens.
+type ManagementTokenRequester interface {
+ RequestAnthropicToken(*gin.Context)
+ RequestGeminiCLIToken(*gin.Context)
+ RequestCodexToken(*gin.Context)
+ RequestAntigravityToken(*gin.Context)
+ RequestQwenToken(*gin.Context)
+ RequestIFlowToken(*gin.Context)
+ RequestIFlowCookieToken(*gin.Context)
+ GetAuthStatus(c *gin.Context)
+ PostOAuthCallback(c *gin.Context)
+}
+
+type managementTokenRequester struct {
+ handler *internalmanagement.Handler
+}
+
+// NewManagementTokenRequester creates a limited management handler exposing only token request endpoints.
+func NewManagementTokenRequester(cfg *config.Config, manager *coreauth.Manager) ManagementTokenRequester {
+ return &managementTokenRequester{
+ handler: internalmanagement.NewHandlerWithoutConfigFilePath(cfg, manager),
+ }
+}
+
+func (m *managementTokenRequester) RequestAnthropicToken(c *gin.Context) {
+ m.handler.RequestAnthropicToken(c)
+}
+
+func (m *managementTokenRequester) RequestGeminiCLIToken(c *gin.Context) {
+ m.handler.RequestGeminiCLIToken(c)
+}
+
+func (m *managementTokenRequester) RequestCodexToken(c *gin.Context) {
+ m.handler.RequestCodexToken(c)
+}
+
+func (m *managementTokenRequester) RequestAntigravityToken(c *gin.Context) {
+ m.handler.RequestAntigravityToken(c)
+}
+
+func (m *managementTokenRequester) RequestQwenToken(c *gin.Context) {
+ m.handler.RequestQwenToken(c)
+}
+
+func (m *managementTokenRequester) RequestIFlowToken(c *gin.Context) {
+ m.handler.RequestIFlowToken(c)
+}
+
+func (m *managementTokenRequester) RequestIFlowCookieToken(c *gin.Context) {
+ m.handler.RequestIFlowCookieToken(c)
+}
+
+func (m *managementTokenRequester) GetAuthStatus(c *gin.Context) {
+ m.handler.GetAuthStatus(c)
+}
+
+func (m *managementTokenRequester) PostOAuthCallback(c *gin.Context) {
+ m.handler.PostOAuthCallback(c)
+}
diff --git a/sdk/api/options.go b/sdk/api/options.go
new file mode 100644
index 0000000000000000000000000000000000000000..8497884bf0bf85a2d6fabe5f37d15adeb8e4bbf5
--- /dev/null
+++ b/sdk/api/options.go
@@ -0,0 +1,46 @@
+// Package api exposes server option helpers for embedding CLIProxyAPI.
+//
+// It wraps internal server option types so external projects can configure the embedded
+// HTTP server without importing internal packages.
+package api
+
+import (
+ "time"
+
+ "github.com/gin-gonic/gin"
+ internalapi "github.com/router-for-me/CLIProxyAPI/v6/internal/api"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/logging"
+)
+
+// ServerOption customises HTTP server construction.
+type ServerOption = internalapi.ServerOption
+
+// WithMiddleware appends additional Gin middleware during server construction.
+func WithMiddleware(mw ...gin.HandlerFunc) ServerOption { return internalapi.WithMiddleware(mw...) }
+
+// WithEngineConfigurator allows callers to mutate the Gin engine prior to middleware setup.
+func WithEngineConfigurator(fn func(*gin.Engine)) ServerOption {
+ return internalapi.WithEngineConfigurator(fn)
+}
+
+// WithRouterConfigurator appends a callback after default routes are registered.
+func WithRouterConfigurator(fn func(*gin.Engine, *handlers.BaseAPIHandler, *config.Config)) ServerOption {
+ return internalapi.WithRouterConfigurator(fn)
+}
+
+// WithLocalManagementPassword stores a runtime-only management password accepted for localhost requests.
+func WithLocalManagementPassword(password string) ServerOption {
+ return internalapi.WithLocalManagementPassword(password)
+}
+
+// WithKeepAliveEndpoint enables a keep-alive endpoint with the provided timeout and callback.
+func WithKeepAliveEndpoint(timeout time.Duration, onTimeout func()) ServerOption {
+ return internalapi.WithKeepAliveEndpoint(timeout, onTimeout)
+}
+
+// WithRequestLoggerFactory customises request logger creation.
+func WithRequestLoggerFactory(factory func(*config.Config, string) logging.RequestLogger) ServerOption {
+ return internalapi.WithRequestLoggerFactory(factory)
+}
diff --git a/sdk/auth/antigravity.go b/sdk/auth/antigravity.go
new file mode 100644
index 0000000000000000000000000000000000000000..ecca0a0041295db06660a00bbd5b6020298ab975
--- /dev/null
+++ b/sdk/auth/antigravity.go
@@ -0,0 +1,266 @@
+package auth
+
+import (
+ "context"
+ "fmt"
+ "net"
+ "net/http"
+ "strings"
+ "time"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/antigravity"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ log "github.com/sirupsen/logrus"
+)
+
+// AntigravityAuthenticator implements OAuth login for the antigravity provider.
+type AntigravityAuthenticator struct{}
+
+// NewAntigravityAuthenticator constructs a new authenticator instance.
+func NewAntigravityAuthenticator() Authenticator { return &AntigravityAuthenticator{} }
+
+// Provider returns the provider key for antigravity.
+func (AntigravityAuthenticator) Provider() string { return "antigravity" }
+
+// RefreshLead instructs the manager to refresh five minutes before expiry.
+func (AntigravityAuthenticator) RefreshLead() *time.Duration {
+ lead := 5 * time.Minute
+ return &lead
+}
+
+// Login launches a local OAuth flow to obtain antigravity tokens and persists them.
+func (AntigravityAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
+ if cfg == nil {
+ return nil, fmt.Errorf("cliproxy auth: configuration is required")
+ }
+ if ctx == nil {
+ ctx = context.Background()
+ }
+ if opts == nil {
+ opts = &LoginOptions{}
+ }
+
+ callbackPort := antigravity.CallbackPort
+ if opts.CallbackPort > 0 {
+ callbackPort = opts.CallbackPort
+ }
+
+ authSvc := antigravity.NewAntigravityAuth(cfg, nil)
+
+ state, err := misc.GenerateRandomState()
+ if err != nil {
+ return nil, fmt.Errorf("antigravity: failed to generate state: %w", err)
+ }
+
+ srv, port, cbChan, errServer := startAntigravityCallbackServer(callbackPort)
+ if errServer != nil {
+ return nil, fmt.Errorf("antigravity: failed to start callback server: %w", errServer)
+ }
+ defer func() {
+ shutdownCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+ defer cancel()
+ _ = srv.Shutdown(shutdownCtx)
+ }()
+
+ redirectURI := fmt.Sprintf("http://localhost:%d/oauth-callback", port)
+ authURL := authSvc.BuildAuthURL(state, redirectURI)
+
+ if !opts.NoBrowser {
+ fmt.Println("Opening browser for antigravity authentication")
+ if !browser.IsAvailable() {
+ log.Warn("No browser available; please open the URL manually")
+ util.PrintSSHTunnelInstructions(port)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ } else if errOpen := browser.OpenURL(authURL); errOpen != nil {
+ log.Warnf("Failed to open browser automatically: %v", errOpen)
+ util.PrintSSHTunnelInstructions(port)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ }
+ } else {
+ util.PrintSSHTunnelInstructions(port)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ }
+
+ fmt.Println("Waiting for antigravity authentication callback...")
+
+ var cbRes callbackResult
+ timeoutTimer := time.NewTimer(5 * time.Minute)
+ defer timeoutTimer.Stop()
+
+ var manualPromptTimer *time.Timer
+ var manualPromptC <-chan time.Time
+ if opts.Prompt != nil {
+ manualPromptTimer = time.NewTimer(15 * time.Second)
+ manualPromptC = manualPromptTimer.C
+ defer manualPromptTimer.Stop()
+ }
+
+waitForCallback:
+ for {
+ select {
+ case res := <-cbChan:
+ cbRes = res
+ break waitForCallback
+ case <-manualPromptC:
+ manualPromptC = nil
+ if manualPromptTimer != nil {
+ manualPromptTimer.Stop()
+ }
+ select {
+ case res := <-cbChan:
+ cbRes = res
+ break waitForCallback
+ default:
+ }
+ input, errPrompt := opts.Prompt("Paste the antigravity callback URL (or press Enter to keep waiting): ")
+ if errPrompt != nil {
+ return nil, errPrompt
+ }
+ parsed, errParse := misc.ParseOAuthCallback(input)
+ if errParse != nil {
+ return nil, errParse
+ }
+ if parsed == nil {
+ continue
+ }
+ cbRes = callbackResult{
+ Code: parsed.Code,
+ State: parsed.State,
+ Error: parsed.Error,
+ }
+ break waitForCallback
+ case <-timeoutTimer.C:
+ return nil, fmt.Errorf("antigravity: authentication timed out")
+ }
+ }
+
+ if cbRes.Error != "" {
+ return nil, fmt.Errorf("antigravity: authentication failed: %s", cbRes.Error)
+ }
+ if cbRes.State != state {
+ return nil, fmt.Errorf("antigravity: invalid state")
+ }
+ if cbRes.Code == "" {
+ return nil, fmt.Errorf("antigravity: missing authorization code")
+ }
+
+ tokenResp, errToken := authSvc.ExchangeCodeForTokens(ctx, cbRes.Code, redirectURI)
+ if errToken != nil {
+ return nil, fmt.Errorf("antigravity: token exchange failed: %w", errToken)
+ }
+
+ accessToken := strings.TrimSpace(tokenResp.AccessToken)
+ if accessToken == "" {
+ return nil, fmt.Errorf("antigravity: token exchange returned empty access token")
+ }
+
+ email, errInfo := authSvc.FetchUserInfo(ctx, accessToken)
+ if errInfo != nil {
+ return nil, fmt.Errorf("antigravity: fetch user info failed: %w", errInfo)
+ }
+ email = strings.TrimSpace(email)
+ if email == "" {
+ return nil, fmt.Errorf("antigravity: empty email returned from user info")
+ }
+
+ // Fetch project ID via loadCodeAssist (same approach as Gemini CLI)
+ projectID := ""
+ if accessToken != "" {
+ fetchedProjectID, errProject := authSvc.FetchProjectID(ctx, accessToken)
+ if errProject != nil {
+ log.Warnf("antigravity: failed to fetch project ID: %v", errProject)
+ } else {
+ projectID = fetchedProjectID
+ log.Infof("antigravity: obtained project ID %s", projectID)
+ }
+ }
+
+ now := time.Now()
+ metadata := map[string]any{
+ "type": "antigravity",
+ "access_token": tokenResp.AccessToken,
+ "refresh_token": tokenResp.RefreshToken,
+ "expires_in": tokenResp.ExpiresIn,
+ "timestamp": now.UnixMilli(),
+ "expired": now.Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339),
+ }
+ if email != "" {
+ metadata["email"] = email
+ }
+ if projectID != "" {
+ metadata["project_id"] = projectID
+ }
+
+ fileName := antigravity.CredentialFileName(email)
+ label := email
+ if label == "" {
+ label = "antigravity"
+ }
+
+ fmt.Println("Antigravity authentication successful")
+ if projectID != "" {
+ fmt.Printf("Using GCP project: %s\n", projectID)
+ }
+ return &coreauth.Auth{
+ ID: fileName,
+ Provider: "antigravity",
+ FileName: fileName,
+ Label: label,
+ Metadata: metadata,
+ }, nil
+}
+
+type callbackResult struct {
+ Code string
+ Error string
+ State string
+}
+
+func startAntigravityCallbackServer(port int) (*http.Server, int, <-chan callbackResult, error) {
+ if port <= 0 {
+ port = antigravity.CallbackPort
+ }
+ addr := fmt.Sprintf(":%d", port)
+ listener, err := net.Listen("tcp", addr)
+ if err != nil {
+ return nil, 0, nil, err
+ }
+ port = listener.Addr().(*net.TCPAddr).Port
+ resultCh := make(chan callbackResult, 1)
+
+ mux := http.NewServeMux()
+ mux.HandleFunc("/oauth-callback", func(w http.ResponseWriter, r *http.Request) {
+ q := r.URL.Query()
+ res := callbackResult{
+ Code: strings.TrimSpace(q.Get("code")),
+ Error: strings.TrimSpace(q.Get("error")),
+ State: strings.TrimSpace(q.Get("state")),
+ }
+ resultCh <- res
+ if res.Code != "" && res.Error == "" {
+ _, _ = w.Write([]byte("Login successful
You can close this window.
"))
+ } else {
+ _, _ = w.Write([]byte("Login failed
Please check the CLI output.
"))
+ }
+ })
+
+ srv := &http.Server{Handler: mux}
+ go func() {
+ if errServe := srv.Serve(listener); errServe != nil && !strings.Contains(errServe.Error(), "Server closed") {
+ log.Warnf("antigravity callback server error: %v", errServe)
+ }
+ }()
+
+ return srv, port, resultCh, nil
+}
+
+// FetchAntigravityProjectID exposes project discovery for external callers.
+func FetchAntigravityProjectID(ctx context.Context, accessToken string, httpClient *http.Client) (string, error) {
+ cfg := &config.Config{}
+ authSvc := antigravity.NewAntigravityAuth(cfg, httpClient)
+ return authSvc.FetchProjectID(ctx, accessToken)
+}
diff --git a/sdk/auth/claude.go b/sdk/auth/claude.go
new file mode 100644
index 0000000000000000000000000000000000000000..2c7a89888a09ccc4b5830086c1b64b40a360fc27
--- /dev/null
+++ b/sdk/auth/claude.go
@@ -0,0 +1,212 @@
+package auth
+
+import (
+ "context"
+ "fmt"
+ "net/http"
+ "strings"
+ "time"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/claude"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
+ // legacy client removed
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ log "github.com/sirupsen/logrus"
+)
+
+// ClaudeAuthenticator implements the OAuth login flow for Anthropic Claude accounts.
+type ClaudeAuthenticator struct {
+ CallbackPort int
+}
+
+// NewClaudeAuthenticator constructs a Claude authenticator with default settings.
+func NewClaudeAuthenticator() *ClaudeAuthenticator {
+ return &ClaudeAuthenticator{CallbackPort: 54545}
+}
+
+func (a *ClaudeAuthenticator) Provider() string {
+ return "claude"
+}
+
+func (a *ClaudeAuthenticator) RefreshLead() *time.Duration {
+ d := 4 * time.Hour
+ return &d
+}
+
+func (a *ClaudeAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
+ if cfg == nil {
+ return nil, fmt.Errorf("cliproxy auth: configuration is required")
+ }
+ if ctx == nil {
+ ctx = context.Background()
+ }
+ if opts == nil {
+ opts = &LoginOptions{}
+ }
+
+ callbackPort := a.CallbackPort
+ if opts.CallbackPort > 0 {
+ callbackPort = opts.CallbackPort
+ }
+
+ pkceCodes, err := claude.GeneratePKCECodes()
+ if err != nil {
+ return nil, fmt.Errorf("claude pkce generation failed: %w", err)
+ }
+
+ state, err := misc.GenerateRandomState()
+ if err != nil {
+ return nil, fmt.Errorf("claude state generation failed: %w", err)
+ }
+
+ oauthServer := claude.NewOAuthServer(callbackPort)
+ if err = oauthServer.Start(); err != nil {
+ if strings.Contains(err.Error(), "already in use") {
+ return nil, claude.NewAuthenticationError(claude.ErrPortInUse, err)
+ }
+ return nil, claude.NewAuthenticationError(claude.ErrServerStartFailed, err)
+ }
+ defer func() {
+ stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+ defer cancel()
+ if stopErr := oauthServer.Stop(stopCtx); stopErr != nil {
+ log.Warnf("claude oauth server stop error: %v", stopErr)
+ }
+ }()
+
+ authSvc := claude.NewClaudeAuth(cfg)
+
+ authURL, returnedState, err := authSvc.GenerateAuthURL(state, pkceCodes)
+ if err != nil {
+ return nil, fmt.Errorf("claude authorization url generation failed: %w", err)
+ }
+ state = returnedState
+
+ if !opts.NoBrowser {
+ fmt.Println("Opening browser for Claude authentication")
+ if !browser.IsAvailable() {
+ log.Warn("No browser available; please open the URL manually")
+ util.PrintSSHTunnelInstructions(callbackPort)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ } else if err = browser.OpenURL(authURL); err != nil {
+ log.Warnf("Failed to open browser automatically: %v", err)
+ util.PrintSSHTunnelInstructions(callbackPort)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ }
+ } else {
+ util.PrintSSHTunnelInstructions(callbackPort)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ }
+
+ fmt.Println("Waiting for Claude authentication callback...")
+
+ callbackCh := make(chan *claude.OAuthResult, 1)
+ callbackErrCh := make(chan error, 1)
+ manualDescription := ""
+
+ go func() {
+ result, errWait := oauthServer.WaitForCallback(5 * time.Minute)
+ if errWait != nil {
+ callbackErrCh <- errWait
+ return
+ }
+ callbackCh <- result
+ }()
+
+ var result *claude.OAuthResult
+ var manualPromptTimer *time.Timer
+ var manualPromptC <-chan time.Time
+ if opts.Prompt != nil {
+ manualPromptTimer = time.NewTimer(15 * time.Second)
+ manualPromptC = manualPromptTimer.C
+ defer manualPromptTimer.Stop()
+ }
+
+waitForCallback:
+ for {
+ select {
+ case result = <-callbackCh:
+ break waitForCallback
+ case err = <-callbackErrCh:
+ if strings.Contains(err.Error(), "timeout") {
+ return nil, claude.NewAuthenticationError(claude.ErrCallbackTimeout, err)
+ }
+ return nil, err
+ case <-manualPromptC:
+ manualPromptC = nil
+ if manualPromptTimer != nil {
+ manualPromptTimer.Stop()
+ }
+ select {
+ case result = <-callbackCh:
+ break waitForCallback
+ case err = <-callbackErrCh:
+ if strings.Contains(err.Error(), "timeout") {
+ return nil, claude.NewAuthenticationError(claude.ErrCallbackTimeout, err)
+ }
+ return nil, err
+ default:
+ }
+ input, errPrompt := opts.Prompt("Paste the Claude callback URL (or press Enter to keep waiting): ")
+ if errPrompt != nil {
+ return nil, errPrompt
+ }
+ parsed, errParse := misc.ParseOAuthCallback(input)
+ if errParse != nil {
+ return nil, errParse
+ }
+ if parsed == nil {
+ continue
+ }
+ manualDescription = parsed.ErrorDescription
+ result = &claude.OAuthResult{
+ Code: parsed.Code,
+ State: parsed.State,
+ Error: parsed.Error,
+ }
+ break waitForCallback
+ }
+ }
+
+ if result.Error != "" {
+ return nil, claude.NewOAuthError(result.Error, manualDescription, http.StatusBadRequest)
+ }
+
+ if result.State != state {
+ return nil, claude.NewAuthenticationError(claude.ErrInvalidState, fmt.Errorf("state mismatch"))
+ }
+
+ log.Debug("Claude authorization code received; exchanging for tokens")
+
+ authBundle, err := authSvc.ExchangeCodeForTokens(ctx, result.Code, state, pkceCodes)
+ if err != nil {
+ return nil, claude.NewAuthenticationError(claude.ErrCodeExchangeFailed, err)
+ }
+
+ tokenStorage := authSvc.CreateTokenStorage(authBundle)
+
+ if tokenStorage == nil || tokenStorage.Email == "" {
+ return nil, fmt.Errorf("claude token storage missing account information")
+ }
+
+ fileName := fmt.Sprintf("claude-%s.json", tokenStorage.Email)
+ metadata := map[string]any{
+ "email": tokenStorage.Email,
+ }
+
+ fmt.Println("Claude authentication successful")
+ if authBundle.APIKey != "" {
+ fmt.Println("Claude API key obtained and stored")
+ }
+
+ return &coreauth.Auth{
+ ID: fileName,
+ Provider: a.Provider(),
+ FileName: fileName,
+ Storage: tokenStorage,
+ Metadata: metadata,
+ }, nil
+}
diff --git a/sdk/auth/codex.go b/sdk/auth/codex.go
new file mode 100644
index 0000000000000000000000000000000000000000..b655a23945e2a00b495400ef56dc2e5fd4753df1
--- /dev/null
+++ b/sdk/auth/codex.go
@@ -0,0 +1,225 @@
+package auth
+
+import (
+ "context"
+ "crypto/sha256"
+ "encoding/hex"
+ "fmt"
+ "net/http"
+ "strings"
+ "time"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/codex"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
+ // legacy client removed
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ log "github.com/sirupsen/logrus"
+)
+
+// CodexAuthenticator implements the OAuth login flow for Codex accounts.
+type CodexAuthenticator struct {
+ CallbackPort int
+}
+
+// NewCodexAuthenticator constructs a Codex authenticator with default settings.
+func NewCodexAuthenticator() *CodexAuthenticator {
+ return &CodexAuthenticator{CallbackPort: 1455}
+}
+
+func (a *CodexAuthenticator) Provider() string {
+ return "codex"
+}
+
+func (a *CodexAuthenticator) RefreshLead() *time.Duration {
+ d := 5 * 24 * time.Hour
+ return &d
+}
+
+func (a *CodexAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
+ if cfg == nil {
+ return nil, fmt.Errorf("cliproxy auth: configuration is required")
+ }
+ if ctx == nil {
+ ctx = context.Background()
+ }
+ if opts == nil {
+ opts = &LoginOptions{}
+ }
+
+ callbackPort := a.CallbackPort
+ if opts.CallbackPort > 0 {
+ callbackPort = opts.CallbackPort
+ }
+
+ pkceCodes, err := codex.GeneratePKCECodes()
+ if err != nil {
+ return nil, fmt.Errorf("codex pkce generation failed: %w", err)
+ }
+
+ state, err := misc.GenerateRandomState()
+ if err != nil {
+ return nil, fmt.Errorf("codex state generation failed: %w", err)
+ }
+
+ oauthServer := codex.NewOAuthServer(callbackPort)
+ if err = oauthServer.Start(); err != nil {
+ if strings.Contains(err.Error(), "already in use") {
+ return nil, codex.NewAuthenticationError(codex.ErrPortInUse, err)
+ }
+ return nil, codex.NewAuthenticationError(codex.ErrServerStartFailed, err)
+ }
+ defer func() {
+ stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+ defer cancel()
+ if stopErr := oauthServer.Stop(stopCtx); stopErr != nil {
+ log.Warnf("codex oauth server stop error: %v", stopErr)
+ }
+ }()
+
+ authSvc := codex.NewCodexAuth(cfg)
+
+ authURL, err := authSvc.GenerateAuthURL(state, pkceCodes)
+ if err != nil {
+ return nil, fmt.Errorf("codex authorization url generation failed: %w", err)
+ }
+
+ if !opts.NoBrowser {
+ fmt.Println("Opening browser for Codex authentication")
+ if !browser.IsAvailable() {
+ log.Warn("No browser available; please open the URL manually")
+ util.PrintSSHTunnelInstructions(callbackPort)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ } else if err = browser.OpenURL(authURL); err != nil {
+ log.Warnf("Failed to open browser automatically: %v", err)
+ util.PrintSSHTunnelInstructions(callbackPort)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ }
+ } else {
+ util.PrintSSHTunnelInstructions(callbackPort)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ }
+
+ fmt.Println("Waiting for Codex authentication callback...")
+
+ callbackCh := make(chan *codex.OAuthResult, 1)
+ callbackErrCh := make(chan error, 1)
+ manualDescription := ""
+
+ go func() {
+ result, errWait := oauthServer.WaitForCallback(5 * time.Minute)
+ if errWait != nil {
+ callbackErrCh <- errWait
+ return
+ }
+ callbackCh <- result
+ }()
+
+ var result *codex.OAuthResult
+ var manualPromptTimer *time.Timer
+ var manualPromptC <-chan time.Time
+ if opts.Prompt != nil {
+ manualPromptTimer = time.NewTimer(15 * time.Second)
+ manualPromptC = manualPromptTimer.C
+ defer manualPromptTimer.Stop()
+ }
+
+waitForCallback:
+ for {
+ select {
+ case result = <-callbackCh:
+ break waitForCallback
+ case err = <-callbackErrCh:
+ if strings.Contains(err.Error(), "timeout") {
+ return nil, codex.NewAuthenticationError(codex.ErrCallbackTimeout, err)
+ }
+ return nil, err
+ case <-manualPromptC:
+ manualPromptC = nil
+ if manualPromptTimer != nil {
+ manualPromptTimer.Stop()
+ }
+ select {
+ case result = <-callbackCh:
+ break waitForCallback
+ case err = <-callbackErrCh:
+ if strings.Contains(err.Error(), "timeout") {
+ return nil, codex.NewAuthenticationError(codex.ErrCallbackTimeout, err)
+ }
+ return nil, err
+ default:
+ }
+ input, errPrompt := opts.Prompt("Paste the Codex callback URL (or press Enter to keep waiting): ")
+ if errPrompt != nil {
+ return nil, errPrompt
+ }
+ parsed, errParse := misc.ParseOAuthCallback(input)
+ if errParse != nil {
+ return nil, errParse
+ }
+ if parsed == nil {
+ continue
+ }
+ manualDescription = parsed.ErrorDescription
+ result = &codex.OAuthResult{
+ Code: parsed.Code,
+ State: parsed.State,
+ Error: parsed.Error,
+ }
+ break waitForCallback
+ }
+ }
+
+ if result.Error != "" {
+ return nil, codex.NewOAuthError(result.Error, manualDescription, http.StatusBadRequest)
+ }
+
+ if result.State != state {
+ return nil, codex.NewAuthenticationError(codex.ErrInvalidState, fmt.Errorf("state mismatch"))
+ }
+
+ log.Debug("Codex authorization code received; exchanging for tokens")
+
+ authBundle, err := authSvc.ExchangeCodeForTokens(ctx, result.Code, pkceCodes)
+ if err != nil {
+ return nil, codex.NewAuthenticationError(codex.ErrCodeExchangeFailed, err)
+ }
+
+ tokenStorage := authSvc.CreateTokenStorage(authBundle)
+
+ if tokenStorage == nil || tokenStorage.Email == "" {
+ return nil, fmt.Errorf("codex token storage missing account information")
+ }
+
+ planType := ""
+ hashAccountID := ""
+ if tokenStorage.IDToken != "" {
+ if claims, errParse := codex.ParseJWTToken(tokenStorage.IDToken); errParse == nil && claims != nil {
+ planType = strings.TrimSpace(claims.CodexAuthInfo.ChatgptPlanType)
+ accountID := strings.TrimSpace(claims.CodexAuthInfo.ChatgptAccountID)
+ if accountID != "" {
+ digest := sha256.Sum256([]byte(accountID))
+ hashAccountID = hex.EncodeToString(digest[:])[:8]
+ }
+ }
+ }
+ fileName := codex.CredentialFileName(tokenStorage.Email, planType, hashAccountID, true)
+ metadata := map[string]any{
+ "email": tokenStorage.Email,
+ }
+
+ fmt.Println("Codex authentication successful")
+ if authBundle.APIKey != "" {
+ fmt.Println("Codex API key obtained and stored")
+ }
+
+ return &coreauth.Auth{
+ ID: fileName,
+ Provider: a.Provider(),
+ FileName: fileName,
+ Storage: tokenStorage,
+ Metadata: metadata,
+ }, nil
+}
diff --git a/sdk/auth/errors.go b/sdk/auth/errors.go
new file mode 100644
index 0000000000000000000000000000000000000000..78fe9a17bd25420d088aab471cba921032612b48
--- /dev/null
+++ b/sdk/auth/errors.go
@@ -0,0 +1,40 @@
+package auth
+
+import (
+ "fmt"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
+)
+
+// ProjectSelectionError indicates that the user must choose a specific project ID.
+type ProjectSelectionError struct {
+ Email string
+ Projects []interfaces.GCPProjectProjects
+}
+
+func (e *ProjectSelectionError) Error() string {
+ if e == nil {
+ return "cliproxy auth: project selection required"
+ }
+ return fmt.Sprintf("cliproxy auth: project selection required for %s", e.Email)
+}
+
+// ProjectsDisplay returns the projects list for caller presentation.
+func (e *ProjectSelectionError) ProjectsDisplay() []interfaces.GCPProjectProjects {
+ if e == nil {
+ return nil
+ }
+ return e.Projects
+}
+
+// EmailRequiredError indicates that the calling context must provide an email or alias.
+type EmailRequiredError struct {
+ Prompt string
+}
+
+func (e *EmailRequiredError) Error() string {
+ if e == nil || e.Prompt == "" {
+ return "cliproxy auth: email is required"
+ }
+ return e.Prompt
+}
diff --git a/sdk/auth/filestore.go b/sdk/auth/filestore.go
new file mode 100644
index 0000000000000000000000000000000000000000..0bb7ff7da3ac2ebc5dbb7fece50e1b13ebdcc8cb
--- /dev/null
+++ b/sdk/auth/filestore.go
@@ -0,0 +1,368 @@
+package auth
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+ "io/fs"
+ "net/http"
+ "os"
+ "path/filepath"
+ "strings"
+ "sync"
+ "time"
+
+ cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+)
+
+// FileTokenStore persists token records and auth metadata using the filesystem as backing storage.
+type FileTokenStore struct {
+ mu sync.Mutex
+ dirLock sync.RWMutex
+ baseDir string
+}
+
+// NewFileTokenStore creates a token store that saves credentials to disk through the
+// TokenStorage implementation embedded in the token record.
+func NewFileTokenStore() *FileTokenStore {
+ return &FileTokenStore{}
+}
+
+// SetBaseDir updates the default directory used for auth JSON persistence when no explicit path is provided.
+func (s *FileTokenStore) SetBaseDir(dir string) {
+ s.dirLock.Lock()
+ s.baseDir = strings.TrimSpace(dir)
+ s.dirLock.Unlock()
+}
+
+// Save persists token storage and metadata to the resolved auth file path.
+func (s *FileTokenStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (string, error) {
+ if auth == nil {
+ return "", fmt.Errorf("auth filestore: auth is nil")
+ }
+
+ path, err := s.resolveAuthPath(auth)
+ if err != nil {
+ return "", err
+ }
+ if path == "" {
+ return "", fmt.Errorf("auth filestore: missing file path attribute for %s", auth.ID)
+ }
+
+ if auth.Disabled {
+ if _, statErr := os.Stat(path); os.IsNotExist(statErr) {
+ return "", nil
+ }
+ }
+
+ s.mu.Lock()
+ defer s.mu.Unlock()
+
+ if err = os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
+ return "", fmt.Errorf("auth filestore: create dir failed: %w", err)
+ }
+
+ switch {
+ case auth.Storage != nil:
+ if err = auth.Storage.SaveTokenToFile(path); err != nil {
+ return "", err
+ }
+ case auth.Metadata != nil:
+ auth.Metadata["disabled"] = auth.Disabled
+ raw, errMarshal := json.Marshal(auth.Metadata)
+ if errMarshal != nil {
+ return "", fmt.Errorf("auth filestore: marshal metadata failed: %w", errMarshal)
+ }
+ if existing, errRead := os.ReadFile(path); errRead == nil {
+ if jsonEqual(existing, raw) {
+ return path, nil
+ }
+ file, errOpen := os.OpenFile(path, os.O_WRONLY|os.O_TRUNC, 0o600)
+ if errOpen != nil {
+ return "", fmt.Errorf("auth filestore: open existing failed: %w", errOpen)
+ }
+ if _, errWrite := file.Write(raw); errWrite != nil {
+ _ = file.Close()
+ return "", fmt.Errorf("auth filestore: write existing failed: %w", errWrite)
+ }
+ if errClose := file.Close(); errClose != nil {
+ return "", fmt.Errorf("auth filestore: close existing failed: %w", errClose)
+ }
+ return path, nil
+ } else if !os.IsNotExist(errRead) {
+ return "", fmt.Errorf("auth filestore: read existing failed: %w", errRead)
+ }
+ if errWrite := os.WriteFile(path, raw, 0o600); errWrite != nil {
+ return "", fmt.Errorf("auth filestore: write file failed: %w", errWrite)
+ }
+ default:
+ return "", fmt.Errorf("auth filestore: nothing to persist for %s", auth.ID)
+ }
+
+ if auth.Attributes == nil {
+ auth.Attributes = make(map[string]string)
+ }
+ auth.Attributes["path"] = path
+
+ if strings.TrimSpace(auth.FileName) == "" {
+ auth.FileName = auth.ID
+ }
+
+ return path, nil
+}
+
+// List enumerates all auth JSON files under the configured directory.
+func (s *FileTokenStore) List(ctx context.Context) ([]*cliproxyauth.Auth, error) {
+ dir := s.baseDirSnapshot()
+ if dir == "" {
+ return nil, fmt.Errorf("auth filestore: directory not configured")
+ }
+ entries := make([]*cliproxyauth.Auth, 0)
+ err := filepath.WalkDir(dir, func(path string, d fs.DirEntry, walkErr error) error {
+ if walkErr != nil {
+ return walkErr
+ }
+ if d.IsDir() {
+ return nil
+ }
+ if !strings.HasSuffix(strings.ToLower(d.Name()), ".json") {
+ return nil
+ }
+ auth, err := s.readAuthFile(path, dir)
+ if err != nil {
+ return nil
+ }
+ if auth != nil {
+ entries = append(entries, auth)
+ }
+ return nil
+ })
+ if err != nil {
+ return nil, err
+ }
+ return entries, nil
+}
+
+// Delete removes the auth file.
+func (s *FileTokenStore) Delete(ctx context.Context, id string) error {
+ id = strings.TrimSpace(id)
+ if id == "" {
+ return fmt.Errorf("auth filestore: id is empty")
+ }
+ path, err := s.resolveDeletePath(id)
+ if err != nil {
+ return err
+ }
+ if err = os.Remove(path); err != nil && !os.IsNotExist(err) {
+ return fmt.Errorf("auth filestore: delete failed: %w", err)
+ }
+ return nil
+}
+
+func (s *FileTokenStore) resolveDeletePath(id string) (string, error) {
+ if strings.ContainsRune(id, os.PathSeparator) || filepath.IsAbs(id) {
+ return id, nil
+ }
+ dir := s.baseDirSnapshot()
+ if dir == "" {
+ return "", fmt.Errorf("auth filestore: directory not configured")
+ }
+ return filepath.Join(dir, id), nil
+}
+
+func (s *FileTokenStore) readAuthFile(path, baseDir string) (*cliproxyauth.Auth, error) {
+ data, err := os.ReadFile(path)
+ if err != nil {
+ return nil, fmt.Errorf("read file: %w", err)
+ }
+ if len(data) == 0 {
+ return nil, nil
+ }
+ metadata := make(map[string]any)
+ if err = json.Unmarshal(data, &metadata); err != nil {
+ return nil, fmt.Errorf("unmarshal auth json: %w", err)
+ }
+ provider, _ := metadata["type"].(string)
+ if provider == "" {
+ provider = "unknown"
+ }
+ if provider == "antigravity" {
+ projectID := ""
+ if pid, ok := metadata["project_id"].(string); ok {
+ projectID = strings.TrimSpace(pid)
+ }
+ if projectID == "" {
+ accessToken := ""
+ if token, ok := metadata["access_token"].(string); ok {
+ accessToken = strings.TrimSpace(token)
+ }
+ if accessToken != "" {
+ fetchedProjectID, errFetch := FetchAntigravityProjectID(context.Background(), accessToken, http.DefaultClient)
+ if errFetch == nil && strings.TrimSpace(fetchedProjectID) != "" {
+ metadata["project_id"] = strings.TrimSpace(fetchedProjectID)
+ if raw, errMarshal := json.Marshal(metadata); errMarshal == nil {
+ if file, errOpen := os.OpenFile(path, os.O_WRONLY|os.O_TRUNC, 0o600); errOpen == nil {
+ _, _ = file.Write(raw)
+ _ = file.Close()
+ }
+ }
+ }
+ }
+ }
+ }
+ info, err := os.Stat(path)
+ if err != nil {
+ return nil, fmt.Errorf("stat file: %w", err)
+ }
+ id := s.idFor(path, baseDir)
+ disabled, _ := metadata["disabled"].(bool)
+ status := cliproxyauth.StatusActive
+ if disabled {
+ status = cliproxyauth.StatusDisabled
+ }
+ auth := &cliproxyauth.Auth{
+ ID: id,
+ Provider: provider,
+ FileName: id,
+ Label: s.labelFor(metadata),
+ Status: status,
+ Disabled: disabled,
+ Attributes: map[string]string{"path": path},
+ Metadata: metadata,
+ CreatedAt: info.ModTime(),
+ UpdatedAt: info.ModTime(),
+ LastRefreshedAt: time.Time{},
+ NextRefreshAfter: time.Time{},
+ }
+ if email, ok := metadata["email"].(string); ok && email != "" {
+ auth.Attributes["email"] = email
+ }
+ return auth, nil
+}
+
+func (s *FileTokenStore) idFor(path, baseDir string) string {
+ if baseDir == "" {
+ return path
+ }
+ rel, err := filepath.Rel(baseDir, path)
+ if err != nil {
+ return path
+ }
+ return rel
+}
+
+func (s *FileTokenStore) resolveAuthPath(auth *cliproxyauth.Auth) (string, error) {
+ if auth == nil {
+ return "", fmt.Errorf("auth filestore: auth is nil")
+ }
+ if auth.Attributes != nil {
+ if p := strings.TrimSpace(auth.Attributes["path"]); p != "" {
+ return p, nil
+ }
+ }
+ if fileName := strings.TrimSpace(auth.FileName); fileName != "" {
+ if filepath.IsAbs(fileName) {
+ return fileName, nil
+ }
+ if dir := s.baseDirSnapshot(); dir != "" {
+ return filepath.Join(dir, fileName), nil
+ }
+ return fileName, nil
+ }
+ if auth.ID == "" {
+ return "", fmt.Errorf("auth filestore: missing id")
+ }
+ if filepath.IsAbs(auth.ID) {
+ return auth.ID, nil
+ }
+ dir := s.baseDirSnapshot()
+ if dir == "" {
+ return "", fmt.Errorf("auth filestore: directory not configured")
+ }
+ return filepath.Join(dir, auth.ID), nil
+}
+
+func (s *FileTokenStore) labelFor(metadata map[string]any) string {
+ if metadata == nil {
+ return ""
+ }
+ if v, ok := metadata["label"].(string); ok && v != "" {
+ return v
+ }
+ if v, ok := metadata["email"].(string); ok && v != "" {
+ return v
+ }
+ if project, ok := metadata["project_id"].(string); ok && project != "" {
+ return project
+ }
+ return ""
+}
+
+func (s *FileTokenStore) baseDirSnapshot() string {
+ s.dirLock.RLock()
+ defer s.dirLock.RUnlock()
+ return s.baseDir
+}
+
+// jsonEqual compares two JSON blobs by parsing them into Go objects and deep comparing.
+func jsonEqual(a, b []byte) bool {
+ var objA any
+ var objB any
+ if err := json.Unmarshal(a, &objA); err != nil {
+ return false
+ }
+ if err := json.Unmarshal(b, &objB); err != nil {
+ return false
+ }
+ return deepEqualJSON(objA, objB)
+}
+
+func deepEqualJSON(a, b any) bool {
+ switch valA := a.(type) {
+ case map[string]any:
+ valB, ok := b.(map[string]any)
+ if !ok || len(valA) != len(valB) {
+ return false
+ }
+ for key, subA := range valA {
+ subB, ok1 := valB[key]
+ if !ok1 || !deepEqualJSON(subA, subB) {
+ return false
+ }
+ }
+ return true
+ case []any:
+ sliceB, ok := b.([]any)
+ if !ok || len(valA) != len(sliceB) {
+ return false
+ }
+ for i := range valA {
+ if !deepEqualJSON(valA[i], sliceB[i]) {
+ return false
+ }
+ }
+ return true
+ case float64:
+ valB, ok := b.(float64)
+ if !ok {
+ return false
+ }
+ return valA == valB
+ case string:
+ valB, ok := b.(string)
+ if !ok {
+ return false
+ }
+ return valA == valB
+ case bool:
+ valB, ok := b.(bool)
+ if !ok {
+ return false
+ }
+ return valA == valB
+ case nil:
+ return b == nil
+ default:
+ return false
+ }
+}
diff --git a/sdk/auth/gemini.go b/sdk/auth/gemini.go
new file mode 100644
index 0000000000000000000000000000000000000000..2b8f9c2b88b854d54fda2e0f54c702eaac6b6173
--- /dev/null
+++ b/sdk/auth/gemini.go
@@ -0,0 +1,73 @@
+package auth
+
+import (
+ "context"
+ "fmt"
+ "time"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/gemini"
+ // legacy client removed
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+)
+
+// GeminiAuthenticator implements the login flow for Google Gemini CLI accounts.
+type GeminiAuthenticator struct{}
+
+// NewGeminiAuthenticator constructs a Gemini authenticator.
+func NewGeminiAuthenticator() *GeminiAuthenticator {
+ return &GeminiAuthenticator{}
+}
+
+func (a *GeminiAuthenticator) Provider() string {
+ return "gemini"
+}
+
+func (a *GeminiAuthenticator) RefreshLead() *time.Duration {
+ return nil
+}
+
+func (a *GeminiAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
+ if cfg == nil {
+ return nil, fmt.Errorf("cliproxy auth: configuration is required")
+ }
+ if ctx == nil {
+ ctx = context.Background()
+ }
+ if opts == nil {
+ opts = &LoginOptions{}
+ }
+
+ var ts gemini.GeminiTokenStorage
+ if opts.ProjectID != "" {
+ ts.ProjectID = opts.ProjectID
+ }
+
+ geminiAuth := gemini.NewGeminiAuth()
+ _, err := geminiAuth.GetAuthenticatedClient(ctx, &ts, cfg, &gemini.WebLoginOptions{
+ NoBrowser: opts.NoBrowser,
+ CallbackPort: opts.CallbackPort,
+ Prompt: opts.Prompt,
+ })
+ if err != nil {
+ return nil, fmt.Errorf("gemini authentication failed: %w", err)
+ }
+
+ // Skip onboarding here; rely on upstream configuration
+
+ fileName := fmt.Sprintf("%s-%s.json", ts.Email, ts.ProjectID)
+ metadata := map[string]any{
+ "email": ts.Email,
+ "project_id": ts.ProjectID,
+ }
+
+ fmt.Println("Gemini authentication successful")
+
+ return &coreauth.Auth{
+ ID: fileName,
+ Provider: a.Provider(),
+ FileName: fileName,
+ Storage: &ts,
+ Metadata: metadata,
+ }, nil
+}
diff --git a/sdk/auth/iflow.go b/sdk/auth/iflow.go
new file mode 100644
index 0000000000000000000000000000000000000000..6d4ff9466b019078b1ef20db8d05306b8cefaf88
--- /dev/null
+++ b/sdk/auth/iflow.go
@@ -0,0 +1,191 @@
+package auth
+
+import (
+ "context"
+ "fmt"
+ "strings"
+ "time"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/iflow"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ log "github.com/sirupsen/logrus"
+)
+
+// IFlowAuthenticator implements the OAuth login flow for iFlow accounts.
+type IFlowAuthenticator struct{}
+
+// NewIFlowAuthenticator constructs a new authenticator instance.
+func NewIFlowAuthenticator() *IFlowAuthenticator { return &IFlowAuthenticator{} }
+
+// Provider returns the provider key for the authenticator.
+func (a *IFlowAuthenticator) Provider() string { return "iflow" }
+
+// RefreshLead indicates how soon before expiry a refresh should be attempted.
+func (a *IFlowAuthenticator) RefreshLead() *time.Duration {
+ d := 24 * time.Hour
+ return &d
+}
+
+// Login performs the OAuth code flow using a local callback server.
+func (a *IFlowAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
+ if cfg == nil {
+ return nil, fmt.Errorf("cliproxy auth: configuration is required")
+ }
+ if ctx == nil {
+ ctx = context.Background()
+ }
+ if opts == nil {
+ opts = &LoginOptions{}
+ }
+
+ callbackPort := iflow.CallbackPort
+ if opts.CallbackPort > 0 {
+ callbackPort = opts.CallbackPort
+ }
+
+ authSvc := iflow.NewIFlowAuth(cfg)
+
+ oauthServer := iflow.NewOAuthServer(callbackPort)
+ if err := oauthServer.Start(); err != nil {
+ if strings.Contains(err.Error(), "already in use") {
+ return nil, fmt.Errorf("iflow authentication server port in use: %w", err)
+ }
+ return nil, fmt.Errorf("iflow authentication server failed: %w", err)
+ }
+ defer func() {
+ stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+ defer cancel()
+ if stopErr := oauthServer.Stop(stopCtx); stopErr != nil {
+ log.Warnf("iflow oauth server stop error: %v", stopErr)
+ }
+ }()
+
+ state, err := misc.GenerateRandomState()
+ if err != nil {
+ return nil, fmt.Errorf("iflow auth: failed to generate state: %w", err)
+ }
+
+ authURL, redirectURI := authSvc.AuthorizationURL(state, callbackPort)
+
+ if !opts.NoBrowser {
+ fmt.Println("Opening browser for iFlow authentication")
+ if !browser.IsAvailable() {
+ log.Warn("No browser available; please open the URL manually")
+ util.PrintSSHTunnelInstructions(callbackPort)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ } else if err = browser.OpenURL(authURL); err != nil {
+ log.Warnf("Failed to open browser automatically: %v", err)
+ util.PrintSSHTunnelInstructions(callbackPort)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ }
+ } else {
+ util.PrintSSHTunnelInstructions(callbackPort)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ }
+
+ fmt.Println("Waiting for iFlow authentication callback...")
+
+ callbackCh := make(chan *iflow.OAuthResult, 1)
+ callbackErrCh := make(chan error, 1)
+
+ go func() {
+ result, errWait := oauthServer.WaitForCallback(5 * time.Minute)
+ if errWait != nil {
+ callbackErrCh <- errWait
+ return
+ }
+ callbackCh <- result
+ }()
+
+ var result *iflow.OAuthResult
+ var manualPromptTimer *time.Timer
+ var manualPromptC <-chan time.Time
+ if opts.Prompt != nil {
+ manualPromptTimer = time.NewTimer(15 * time.Second)
+ manualPromptC = manualPromptTimer.C
+ defer manualPromptTimer.Stop()
+ }
+
+waitForCallback:
+ for {
+ select {
+ case result = <-callbackCh:
+ break waitForCallback
+ case err = <-callbackErrCh:
+ return nil, fmt.Errorf("iflow auth: callback wait failed: %w", err)
+ case <-manualPromptC:
+ manualPromptC = nil
+ if manualPromptTimer != nil {
+ manualPromptTimer.Stop()
+ }
+ select {
+ case result = <-callbackCh:
+ break waitForCallback
+ case err = <-callbackErrCh:
+ return nil, fmt.Errorf("iflow auth: callback wait failed: %w", err)
+ default:
+ }
+ input, errPrompt := opts.Prompt("Paste the iFlow callback URL (or press Enter to keep waiting): ")
+ if errPrompt != nil {
+ return nil, errPrompt
+ }
+ parsed, errParse := misc.ParseOAuthCallback(input)
+ if errParse != nil {
+ return nil, errParse
+ }
+ if parsed == nil {
+ continue
+ }
+ result = &iflow.OAuthResult{
+ Code: parsed.Code,
+ State: parsed.State,
+ Error: parsed.Error,
+ }
+ break waitForCallback
+ }
+ }
+ if result.Error != "" {
+ return nil, fmt.Errorf("iflow auth: provider returned error %s", result.Error)
+ }
+ if result.State != state {
+ return nil, fmt.Errorf("iflow auth: state mismatch")
+ }
+
+ tokenData, err := authSvc.ExchangeCodeForTokens(ctx, result.Code, redirectURI)
+ if err != nil {
+ return nil, fmt.Errorf("iflow authentication failed: %w", err)
+ }
+
+ tokenStorage := authSvc.CreateTokenStorage(tokenData)
+
+ email := strings.TrimSpace(tokenStorage.Email)
+ if email == "" {
+ return nil, fmt.Errorf("iflow authentication failed: missing account identifier")
+ }
+
+ fileName := fmt.Sprintf("iflow-%s-%d.json", email, time.Now().Unix())
+ metadata := map[string]any{
+ "email": email,
+ "api_key": tokenStorage.APIKey,
+ "access_token": tokenStorage.AccessToken,
+ "refresh_token": tokenStorage.RefreshToken,
+ "expired": tokenStorage.Expire,
+ }
+
+ fmt.Println("iFlow authentication successful")
+
+ return &coreauth.Auth{
+ ID: fileName,
+ Provider: a.Provider(),
+ FileName: fileName,
+ Storage: tokenStorage,
+ Metadata: metadata,
+ Attributes: map[string]string{
+ "api_key": tokenStorage.APIKey,
+ },
+ }, nil
+}
diff --git a/sdk/auth/interfaces.go b/sdk/auth/interfaces.go
new file mode 100644
index 0000000000000000000000000000000000000000..64cf8ed035a325f021a78e0a364a60cb72c25784
--- /dev/null
+++ b/sdk/auth/interfaces.go
@@ -0,0 +1,29 @@
+package auth
+
+import (
+ "context"
+ "errors"
+ "time"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+)
+
+var ErrRefreshNotSupported = errors.New("cliproxy auth: refresh not supported")
+
+// LoginOptions captures generic knobs shared across authenticators.
+// Provider-specific logic can inspect Metadata for extra parameters.
+type LoginOptions struct {
+ NoBrowser bool
+ ProjectID string
+ CallbackPort int
+ Metadata map[string]string
+ Prompt func(prompt string) (string, error)
+}
+
+// Authenticator manages login and optional refresh flows for a provider.
+type Authenticator interface {
+ Provider() string
+ Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error)
+ RefreshLead() *time.Duration
+}
diff --git a/sdk/auth/manager.go b/sdk/auth/manager.go
new file mode 100644
index 0000000000000000000000000000000000000000..c6469a7d1991ad9a1c9f02f746cd1d7b9a9f6032
--- /dev/null
+++ b/sdk/auth/manager.go
@@ -0,0 +1,76 @@
+package auth
+
+import (
+ "context"
+ "fmt"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+)
+
+// Manager aggregates authenticators and coordinates persistence via a token store.
+type Manager struct {
+ authenticators map[string]Authenticator
+ store coreauth.Store
+}
+
+// NewManager constructs a manager with the provided token store and authenticators.
+// If store is nil, the caller must set it later using SetStore.
+func NewManager(store coreauth.Store, authenticators ...Authenticator) *Manager {
+ mgr := &Manager{
+ authenticators: make(map[string]Authenticator),
+ store: store,
+ }
+ for i := range authenticators {
+ mgr.Register(authenticators[i])
+ }
+ return mgr
+}
+
+// Register adds or replaces an authenticator keyed by its provider identifier.
+func (m *Manager) Register(a Authenticator) {
+ if a == nil {
+ return
+ }
+ if m.authenticators == nil {
+ m.authenticators = make(map[string]Authenticator)
+ }
+ m.authenticators[a.Provider()] = a
+}
+
+// SetStore updates the token store used for persistence.
+func (m *Manager) SetStore(store coreauth.Store) {
+ m.store = store
+}
+
+// Login executes the provider login flow and persists the resulting auth record.
+func (m *Manager) Login(ctx context.Context, provider string, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, string, error) {
+ auth, ok := m.authenticators[provider]
+ if !ok {
+ return nil, "", fmt.Errorf("cliproxy auth: authenticator %s not registered", provider)
+ }
+
+ record, err := auth.Login(ctx, cfg, opts)
+ if err != nil {
+ return nil, "", err
+ }
+ if record == nil {
+ return nil, "", fmt.Errorf("cliproxy auth: authenticator %s returned nil record", provider)
+ }
+
+ if m.store == nil {
+ return record, "", nil
+ }
+
+ if cfg != nil {
+ if dirSetter, ok := m.store.(interface{ SetBaseDir(string) }); ok {
+ dirSetter.SetBaseDir(cfg.AuthDir)
+ }
+ }
+
+ savedPath, err := m.store.Save(ctx, record)
+ if err != nil {
+ return record, "", err
+ }
+ return record, savedPath, nil
+}
diff --git a/sdk/auth/qwen.go b/sdk/auth/qwen.go
new file mode 100644
index 0000000000000000000000000000000000000000..151fba6816e279ae04d4f8645c0a837dcce53414
--- /dev/null
+++ b/sdk/auth/qwen.go
@@ -0,0 +1,114 @@
+package auth
+
+import (
+ "context"
+ "fmt"
+ "strings"
+ "time"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/qwen"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
+ // legacy client removed
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ log "github.com/sirupsen/logrus"
+)
+
+// QwenAuthenticator implements the device flow login for Qwen accounts.
+type QwenAuthenticator struct{}
+
+// NewQwenAuthenticator constructs a Qwen authenticator.
+func NewQwenAuthenticator() *QwenAuthenticator {
+ return &QwenAuthenticator{}
+}
+
+func (a *QwenAuthenticator) Provider() string {
+ return "qwen"
+}
+
+func (a *QwenAuthenticator) RefreshLead() *time.Duration {
+ d := 3 * time.Hour
+ return &d
+}
+
+func (a *QwenAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
+ if cfg == nil {
+ return nil, fmt.Errorf("cliproxy auth: configuration is required")
+ }
+ if ctx == nil {
+ ctx = context.Background()
+ }
+ if opts == nil {
+ opts = &LoginOptions{}
+ }
+
+ authSvc := qwen.NewQwenAuth(cfg)
+
+ deviceFlow, err := authSvc.InitiateDeviceFlow(ctx)
+ if err != nil {
+ return nil, fmt.Errorf("qwen device flow initiation failed: %w", err)
+ }
+
+ authURL := deviceFlow.VerificationURIComplete
+
+ if !opts.NoBrowser {
+ fmt.Println("Opening browser for Qwen authentication")
+ if !browser.IsAvailable() {
+ log.Warn("No browser available; please open the URL manually")
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ } else if err = browser.OpenURL(authURL); err != nil {
+ log.Warnf("Failed to open browser automatically: %v", err)
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ }
+ } else {
+ fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL)
+ }
+
+ fmt.Println("Waiting for Qwen authentication...")
+
+ tokenData, err := authSvc.PollForToken(deviceFlow.DeviceCode, deviceFlow.CodeVerifier)
+ if err != nil {
+ return nil, fmt.Errorf("qwen authentication failed: %w", err)
+ }
+
+ tokenStorage := authSvc.CreateTokenStorage(tokenData)
+
+ email := ""
+ if opts.Metadata != nil {
+ email = opts.Metadata["email"]
+ if email == "" {
+ email = opts.Metadata["alias"]
+ }
+ }
+
+ if email == "" && opts.Prompt != nil {
+ email, err = opts.Prompt("Please input your email address or alias for Qwen:")
+ if err != nil {
+ return nil, err
+ }
+ }
+
+ email = strings.TrimSpace(email)
+ if email == "" {
+ return nil, &EmailRequiredError{Prompt: "Please provide an email address or alias for Qwen."}
+ }
+
+ tokenStorage.Email = email
+
+ // no legacy client construction
+
+ fileName := fmt.Sprintf("qwen-%s.json", tokenStorage.Email)
+ metadata := map[string]any{
+ "email": tokenStorage.Email,
+ }
+
+ fmt.Println("Qwen authentication successful")
+
+ return &coreauth.Auth{
+ ID: fileName,
+ Provider: a.Provider(),
+ FileName: fileName,
+ Storage: tokenStorage,
+ Metadata: metadata,
+ }, nil
+}
diff --git a/sdk/auth/refresh_registry.go b/sdk/auth/refresh_registry.go
new file mode 100644
index 0000000000000000000000000000000000000000..e82ac68487d02c9f584c6b94df67badbf9a7acce
--- /dev/null
+++ b/sdk/auth/refresh_registry.go
@@ -0,0 +1,30 @@
+package auth
+
+import (
+ "time"
+
+ cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+)
+
+func init() {
+ registerRefreshLead("codex", func() Authenticator { return NewCodexAuthenticator() })
+ registerRefreshLead("claude", func() Authenticator { return NewClaudeAuthenticator() })
+ registerRefreshLead("qwen", func() Authenticator { return NewQwenAuthenticator() })
+ registerRefreshLead("iflow", func() Authenticator { return NewIFlowAuthenticator() })
+ registerRefreshLead("gemini", func() Authenticator { return NewGeminiAuthenticator() })
+ registerRefreshLead("gemini-cli", func() Authenticator { return NewGeminiAuthenticator() })
+ registerRefreshLead("antigravity", func() Authenticator { return NewAntigravityAuthenticator() })
+}
+
+func registerRefreshLead(provider string, factory func() Authenticator) {
+ cliproxyauth.RegisterRefreshLeadProvider(provider, func() *time.Duration {
+ if factory == nil {
+ return nil
+ }
+ auth := factory()
+ if auth == nil {
+ return nil
+ }
+ return auth.RefreshLead()
+ })
+}
diff --git a/sdk/auth/store_registry.go b/sdk/auth/store_registry.go
new file mode 100644
index 0000000000000000000000000000000000000000..760449f8cf6fa8004964f796ec317f4cf00ab87b
--- /dev/null
+++ b/sdk/auth/store_registry.go
@@ -0,0 +1,35 @@
+package auth
+
+import (
+ "sync"
+
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+)
+
+var (
+ storeMu sync.RWMutex
+ registeredStore coreauth.Store
+)
+
+// RegisterTokenStore sets the global token store used by the authentication helpers.
+func RegisterTokenStore(store coreauth.Store) {
+ storeMu.Lock()
+ registeredStore = store
+ storeMu.Unlock()
+}
+
+// GetTokenStore returns the globally registered token store.
+func GetTokenStore() coreauth.Store {
+ storeMu.RLock()
+ s := registeredStore
+ storeMu.RUnlock()
+ if s != nil {
+ return s
+ }
+ storeMu.Lock()
+ defer storeMu.Unlock()
+ if registeredStore == nil {
+ registeredStore = NewFileTokenStore()
+ }
+ return registeredStore
+}
diff --git a/sdk/cliproxy/auth/api_key_model_alias_test.go b/sdk/cliproxy/auth/api_key_model_alias_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..70915d9e373a0dbfbd91ebfa3cb543670a7ec1ff
--- /dev/null
+++ b/sdk/cliproxy/auth/api_key_model_alias_test.go
@@ -0,0 +1,180 @@
+package auth
+
+import (
+ "context"
+ "testing"
+
+ internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+)
+
+func TestLookupAPIKeyUpstreamModel(t *testing.T) {
+ cfg := &internalconfig.Config{
+ GeminiKey: []internalconfig.GeminiKey{
+ {
+ APIKey: "k",
+ BaseURL: "https://example.com",
+ Models: []internalconfig.GeminiModel{
+ {Name: "gemini-2.5-pro-exp-03-25", Alias: "g25p"},
+ {Name: "gemini-2.5-flash(low)", Alias: "g25f"},
+ },
+ },
+ },
+ }
+
+ mgr := NewManager(nil, nil, nil)
+ mgr.SetConfig(cfg)
+
+ ctx := context.Background()
+ _, _ = mgr.Register(ctx, &Auth{ID: "a1", Provider: "gemini", Attributes: map[string]string{"api_key": "k", "base_url": "https://example.com"}})
+
+ tests := []struct {
+ name string
+ authID string
+ input string
+ want string
+ }{
+ // Fast path + suffix preservation
+ {"alias with suffix", "a1", "g25p(8192)", "gemini-2.5-pro-exp-03-25(8192)"},
+ {"alias without suffix", "a1", "g25p", "gemini-2.5-pro-exp-03-25"},
+
+ // Config suffix takes priority
+ {"config suffix priority", "a1", "g25f(high)", "gemini-2.5-flash(low)"},
+ {"config suffix no user suffix", "a1", "g25f", "gemini-2.5-flash(low)"},
+
+ // Case insensitive
+ {"uppercase alias", "a1", "G25P", "gemini-2.5-pro-exp-03-25"},
+ {"mixed case with suffix", "a1", "G25p(4096)", "gemini-2.5-pro-exp-03-25(4096)"},
+
+ // Direct name lookup
+ {"upstream name direct", "a1", "gemini-2.5-pro-exp-03-25", "gemini-2.5-pro-exp-03-25"},
+ {"upstream name with suffix", "a1", "gemini-2.5-pro-exp-03-25(8192)", "gemini-2.5-pro-exp-03-25(8192)"},
+
+ // Cache miss scenarios
+ {"non-existent auth", "non-existent", "g25p", ""},
+ {"unknown alias", "a1", "unknown-alias", ""},
+ {"empty auth ID", "", "g25p", ""},
+ {"empty model", "a1", "", ""},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ resolved := mgr.lookupAPIKeyUpstreamModel(tt.authID, tt.input)
+ if resolved != tt.want {
+ t.Errorf("lookupAPIKeyUpstreamModel(%q, %q) = %q, want %q", tt.authID, tt.input, resolved, tt.want)
+ }
+ })
+ }
+}
+
+func TestAPIKeyModelAlias_ConfigHotReload(t *testing.T) {
+ cfg := &internalconfig.Config{
+ GeminiKey: []internalconfig.GeminiKey{
+ {
+ APIKey: "k",
+ Models: []internalconfig.GeminiModel{{Name: "gemini-2.5-pro-exp-03-25", Alias: "g25p"}},
+ },
+ },
+ }
+
+ mgr := NewManager(nil, nil, nil)
+ mgr.SetConfig(cfg)
+
+ ctx := context.Background()
+ _, _ = mgr.Register(ctx, &Auth{ID: "a1", Provider: "gemini", Attributes: map[string]string{"api_key": "k"}})
+
+ // Initial alias
+ if resolved := mgr.lookupAPIKeyUpstreamModel("a1", "g25p"); resolved != "gemini-2.5-pro-exp-03-25" {
+ t.Fatalf("before reload: got %q, want %q", resolved, "gemini-2.5-pro-exp-03-25")
+ }
+
+ // Hot reload with new alias
+ mgr.SetConfig(&internalconfig.Config{
+ GeminiKey: []internalconfig.GeminiKey{
+ {
+ APIKey: "k",
+ Models: []internalconfig.GeminiModel{{Name: "gemini-2.5-flash", Alias: "g25p"}},
+ },
+ },
+ })
+
+ // New alias should take effect
+ if resolved := mgr.lookupAPIKeyUpstreamModel("a1", "g25p"); resolved != "gemini-2.5-flash" {
+ t.Fatalf("after reload: got %q, want %q", resolved, "gemini-2.5-flash")
+ }
+}
+
+func TestAPIKeyModelAlias_MultipleProviders(t *testing.T) {
+ cfg := &internalconfig.Config{
+ GeminiKey: []internalconfig.GeminiKey{{APIKey: "gemini-key", Models: []internalconfig.GeminiModel{{Name: "gemini-2.5-pro", Alias: "gp"}}}},
+ ClaudeKey: []internalconfig.ClaudeKey{{APIKey: "claude-key", Models: []internalconfig.ClaudeModel{{Name: "claude-sonnet-4", Alias: "cs4"}}}},
+ CodexKey: []internalconfig.CodexKey{{APIKey: "codex-key", Models: []internalconfig.CodexModel{{Name: "o3", Alias: "o"}}}},
+ }
+
+ mgr := NewManager(nil, nil, nil)
+ mgr.SetConfig(cfg)
+
+ ctx := context.Background()
+ _, _ = mgr.Register(ctx, &Auth{ID: "gemini-auth", Provider: "gemini", Attributes: map[string]string{"api_key": "gemini-key"}})
+ _, _ = mgr.Register(ctx, &Auth{ID: "claude-auth", Provider: "claude", Attributes: map[string]string{"api_key": "claude-key"}})
+ _, _ = mgr.Register(ctx, &Auth{ID: "codex-auth", Provider: "codex", Attributes: map[string]string{"api_key": "codex-key"}})
+
+ tests := []struct {
+ authID, input, want string
+ }{
+ {"gemini-auth", "gp", "gemini-2.5-pro"},
+ {"claude-auth", "cs4", "claude-sonnet-4"},
+ {"codex-auth", "o", "o3"},
+ }
+
+ for _, tt := range tests {
+ if resolved := mgr.lookupAPIKeyUpstreamModel(tt.authID, tt.input); resolved != tt.want {
+ t.Errorf("lookupAPIKeyUpstreamModel(%q, %q) = %q, want %q", tt.authID, tt.input, resolved, tt.want)
+ }
+ }
+}
+
+func TestApplyAPIKeyModelAlias(t *testing.T) {
+ cfg := &internalconfig.Config{
+ GeminiKey: []internalconfig.GeminiKey{
+ {APIKey: "k", Models: []internalconfig.GeminiModel{{Name: "gemini-2.5-pro-exp-03-25", Alias: "g25p"}}},
+ },
+ }
+
+ mgr := NewManager(nil, nil, nil)
+ mgr.SetConfig(cfg)
+
+ ctx := context.Background()
+ apiKeyAuth := &Auth{ID: "a1", Provider: "gemini", Attributes: map[string]string{"api_key": "k"}}
+ oauthAuth := &Auth{ID: "oauth-auth", Provider: "gemini", Attributes: map[string]string{"auth_kind": "oauth"}}
+ _, _ = mgr.Register(ctx, apiKeyAuth)
+
+ tests := []struct {
+ name string
+ auth *Auth
+ inputModel string
+ wantModel string
+ }{
+ {
+ name: "api_key auth with alias",
+ auth: apiKeyAuth,
+ inputModel: "g25p(8192)",
+ wantModel: "gemini-2.5-pro-exp-03-25(8192)",
+ },
+ {
+ name: "oauth auth passthrough",
+ auth: oauthAuth,
+ inputModel: "some-model",
+ wantModel: "some-model",
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ resolvedModel := mgr.applyAPIKeyModelAlias(tt.auth, tt.inputModel)
+
+ if resolvedModel != tt.wantModel {
+ t.Errorf("model = %q, want %q", resolvedModel, tt.wantModel)
+ }
+ })
+ }
+}
diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go
new file mode 100644
index 0000000000000000000000000000000000000000..3a64c8c3476c29db7f6f756380862f0ac8025f56
--- /dev/null
+++ b/sdk/cliproxy/auth/conductor.go
@@ -0,0 +1,2232 @@
+package auth
+
+import (
+ "bytes"
+ "context"
+ "encoding/json"
+ "errors"
+ "io"
+ "net/http"
+ "path/filepath"
+ "strconv"
+ "strings"
+ "sync"
+ "sync/atomic"
+ "time"
+
+ "github.com/google/uuid"
+ internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
+ cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
+ log "github.com/sirupsen/logrus"
+)
+
+// ProviderExecutor defines the contract required by Manager to execute provider calls.
+type ProviderExecutor interface {
+ // Identifier returns the provider key handled by this executor.
+ Identifier() string
+ // Execute handles non-streaming execution and returns the provider response payload.
+ Execute(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error)
+ // ExecuteStream handles streaming execution and returns a channel of provider chunks.
+ ExecuteStream(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (<-chan cliproxyexecutor.StreamChunk, error)
+ // Refresh attempts to refresh provider credentials and returns the updated auth state.
+ Refresh(ctx context.Context, auth *Auth) (*Auth, error)
+ // CountTokens returns the token count for the given request.
+ CountTokens(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error)
+ // HttpRequest injects provider credentials into the supplied HTTP request and executes it.
+ // Callers must close the response body when non-nil.
+ HttpRequest(ctx context.Context, auth *Auth, req *http.Request) (*http.Response, error)
+}
+
+// RefreshEvaluator allows runtime state to override refresh decisions.
+type RefreshEvaluator interface {
+ ShouldRefresh(now time.Time, auth *Auth) bool
+}
+
+const (
+ refreshCheckInterval = 5 * time.Second
+ refreshPendingBackoff = time.Minute
+ refreshFailureBackoff = 5 * time.Minute
+ quotaBackoffBase = time.Second
+ quotaBackoffMax = 30 * time.Minute
+)
+
+var quotaCooldownDisabled atomic.Bool
+
+// SetQuotaCooldownDisabled toggles quota cooldown scheduling globally.
+func SetQuotaCooldownDisabled(disable bool) {
+ quotaCooldownDisabled.Store(disable)
+}
+
+func quotaCooldownDisabledForAuth(auth *Auth) bool {
+ if auth != nil {
+ if override, ok := auth.DisableCoolingOverride(); ok {
+ return override
+ }
+ }
+ return quotaCooldownDisabled.Load()
+}
+
+// Result captures execution outcome used to adjust auth state.
+type Result struct {
+ // AuthID references the auth that produced this result.
+ AuthID string
+ // Provider is copied for convenience when emitting hooks.
+ Provider string
+ // Model is the upstream model identifier used for the request.
+ Model string
+ // Success marks whether the execution succeeded.
+ Success bool
+ // RetryAfter carries a provider supplied retry hint (e.g. 429 retryDelay).
+ RetryAfter *time.Duration
+ // Error describes the failure when Success is false.
+ Error *Error
+}
+
+// Selector chooses an auth candidate for execution.
+type Selector interface {
+ Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error)
+}
+
+// Hook captures lifecycle callbacks for observing auth changes.
+type Hook interface {
+ // OnAuthRegistered fires when a new auth is registered.
+ OnAuthRegistered(ctx context.Context, auth *Auth)
+ // OnAuthUpdated fires when an existing auth changes state.
+ OnAuthUpdated(ctx context.Context, auth *Auth)
+ // OnResult fires when execution result is recorded.
+ OnResult(ctx context.Context, result Result)
+}
+
+// NoopHook provides optional hook defaults.
+type NoopHook struct{}
+
+// OnAuthRegistered implements Hook.
+func (NoopHook) OnAuthRegistered(context.Context, *Auth) {}
+
+// OnAuthUpdated implements Hook.
+func (NoopHook) OnAuthUpdated(context.Context, *Auth) {}
+
+// OnResult implements Hook.
+func (NoopHook) OnResult(context.Context, Result) {}
+
+// Manager orchestrates auth lifecycle, selection, execution, and persistence.
+type Manager struct {
+ store Store
+ executors map[string]ProviderExecutor
+ selector Selector
+ hook Hook
+ mu sync.RWMutex
+ auths map[string]*Auth
+ // providerOffsets tracks per-model provider rotation state for multi-provider routing.
+ providerOffsets map[string]int
+
+ // Retry controls request retry behavior.
+ requestRetry atomic.Int32
+ maxRetryInterval atomic.Int64
+
+ // oauthModelAlias stores global OAuth model alias mappings (alias -> upstream name) keyed by channel.
+ oauthModelAlias atomic.Value
+
+ // apiKeyModelAlias caches resolved model alias mappings for API-key auths.
+ // Keyed by auth.ID, value is alias(lower) -> upstream model (including suffix).
+ apiKeyModelAlias atomic.Value
+
+ // runtimeConfig stores the latest application config for request-time decisions.
+ // It is initialized in NewManager; never Load() before first Store().
+ runtimeConfig atomic.Value
+
+ // Optional HTTP RoundTripper provider injected by host.
+ rtProvider RoundTripperProvider
+
+ // Auto refresh state
+ refreshCancel context.CancelFunc
+}
+
+// NewManager constructs a manager with optional custom selector and hook.
+func NewManager(store Store, selector Selector, hook Hook) *Manager {
+ if selector == nil {
+ selector = &RoundRobinSelector{}
+ }
+ if hook == nil {
+ hook = NoopHook{}
+ }
+ manager := &Manager{
+ store: store,
+ executors: make(map[string]ProviderExecutor),
+ selector: selector,
+ hook: hook,
+ auths: make(map[string]*Auth),
+ providerOffsets: make(map[string]int),
+ }
+ // atomic.Value requires non-nil initial value.
+ manager.runtimeConfig.Store(&internalconfig.Config{})
+ manager.apiKeyModelAlias.Store(apiKeyModelAliasTable(nil))
+ return manager
+}
+
+func (m *Manager) SetSelector(selector Selector) {
+ if m == nil {
+ return
+ }
+ if selector == nil {
+ selector = &RoundRobinSelector{}
+ }
+ m.mu.Lock()
+ m.selector = selector
+ m.mu.Unlock()
+}
+
+// SetStore swaps the underlying persistence store.
+func (m *Manager) SetStore(store Store) {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+ m.store = store
+}
+
+// SetRoundTripperProvider register a provider that returns a per-auth RoundTripper.
+func (m *Manager) SetRoundTripperProvider(p RoundTripperProvider) {
+ m.mu.Lock()
+ m.rtProvider = p
+ m.mu.Unlock()
+}
+
+// SetConfig updates the runtime config snapshot used by request-time helpers.
+// Callers should provide the latest config on reload so per-credential alias mapping stays in sync.
+func (m *Manager) SetConfig(cfg *internalconfig.Config) {
+ if m == nil {
+ return
+ }
+ if cfg == nil {
+ cfg = &internalconfig.Config{}
+ }
+ m.runtimeConfig.Store(cfg)
+ m.rebuildAPIKeyModelAliasFromRuntimeConfig()
+}
+
+func (m *Manager) lookupAPIKeyUpstreamModel(authID, requestedModel string) string {
+ if m == nil {
+ return ""
+ }
+ authID = strings.TrimSpace(authID)
+ if authID == "" {
+ return ""
+ }
+ requestedModel = strings.TrimSpace(requestedModel)
+ if requestedModel == "" {
+ return ""
+ }
+ table, _ := m.apiKeyModelAlias.Load().(apiKeyModelAliasTable)
+ if table == nil {
+ return ""
+ }
+ byAlias := table[authID]
+ if len(byAlias) == 0 {
+ return ""
+ }
+ key := strings.ToLower(thinking.ParseSuffix(requestedModel).ModelName)
+ if key == "" {
+ key = strings.ToLower(requestedModel)
+ }
+ resolved := strings.TrimSpace(byAlias[key])
+ if resolved == "" {
+ return ""
+ }
+ // Preserve thinking suffix from the client's requested model unless config already has one.
+ requestResult := thinking.ParseSuffix(requestedModel)
+ if thinking.ParseSuffix(resolved).HasSuffix {
+ return resolved
+ }
+ if requestResult.HasSuffix && requestResult.RawSuffix != "" {
+ return resolved + "(" + requestResult.RawSuffix + ")"
+ }
+ return resolved
+
+}
+
+func (m *Manager) rebuildAPIKeyModelAliasFromRuntimeConfig() {
+ if m == nil {
+ return
+ }
+ cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
+ if cfg == nil {
+ cfg = &internalconfig.Config{}
+ }
+ m.mu.Lock()
+ defer m.mu.Unlock()
+ m.rebuildAPIKeyModelAliasLocked(cfg)
+}
+
+func (m *Manager) rebuildAPIKeyModelAliasLocked(cfg *internalconfig.Config) {
+ if m == nil {
+ return
+ }
+ if cfg == nil {
+ cfg = &internalconfig.Config{}
+ }
+
+ out := make(apiKeyModelAliasTable)
+ for _, auth := range m.auths {
+ if auth == nil {
+ continue
+ }
+ if strings.TrimSpace(auth.ID) == "" {
+ continue
+ }
+ kind, _ := auth.AccountInfo()
+ if !strings.EqualFold(strings.TrimSpace(kind), "api_key") {
+ continue
+ }
+
+ byAlias := make(map[string]string)
+ provider := strings.ToLower(strings.TrimSpace(auth.Provider))
+ switch provider {
+ case "gemini":
+ if entry := resolveGeminiAPIKeyConfig(cfg, auth); entry != nil {
+ compileAPIKeyModelAliasForModels(byAlias, entry.Models)
+ }
+ case "claude":
+ if entry := resolveClaudeAPIKeyConfig(cfg, auth); entry != nil {
+ compileAPIKeyModelAliasForModels(byAlias, entry.Models)
+ }
+ case "codex":
+ if entry := resolveCodexAPIKeyConfig(cfg, auth); entry != nil {
+ compileAPIKeyModelAliasForModels(byAlias, entry.Models)
+ }
+ case "vertex":
+ if entry := resolveVertexAPIKeyConfig(cfg, auth); entry != nil {
+ compileAPIKeyModelAliasForModels(byAlias, entry.Models)
+ }
+ default:
+ // OpenAI-compat uses config selection from auth.Attributes.
+ providerKey := ""
+ compatName := ""
+ if auth.Attributes != nil {
+ providerKey = strings.TrimSpace(auth.Attributes["provider_key"])
+ compatName = strings.TrimSpace(auth.Attributes["compat_name"])
+ }
+ if compatName != "" || strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") {
+ if entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider); entry != nil {
+ compileAPIKeyModelAliasForModels(byAlias, entry.Models)
+ }
+ }
+ }
+
+ if len(byAlias) > 0 {
+ out[auth.ID] = byAlias
+ }
+ }
+
+ m.apiKeyModelAlias.Store(out)
+}
+
+func compileAPIKeyModelAliasForModels[T interface {
+ GetName() string
+ GetAlias() string
+}](out map[string]string, models []T) {
+ if out == nil {
+ return
+ }
+ for i := range models {
+ alias := strings.TrimSpace(models[i].GetAlias())
+ name := strings.TrimSpace(models[i].GetName())
+ if alias == "" || name == "" {
+ continue
+ }
+ aliasKey := strings.ToLower(thinking.ParseSuffix(alias).ModelName)
+ if aliasKey == "" {
+ aliasKey = strings.ToLower(alias)
+ }
+ // Config priority: first alias wins.
+ if _, exists := out[aliasKey]; exists {
+ continue
+ }
+ out[aliasKey] = name
+ // Also allow direct lookup by upstream name (case-insensitive), so lookups on already-upstream
+ // models remain a cheap no-op.
+ nameKey := strings.ToLower(thinking.ParseSuffix(name).ModelName)
+ if nameKey == "" {
+ nameKey = strings.ToLower(name)
+ }
+ if nameKey != "" {
+ if _, exists := out[nameKey]; !exists {
+ out[nameKey] = name
+ }
+ }
+ // Preserve config suffix priority by seeding a base-name lookup when name already has suffix.
+ nameResult := thinking.ParseSuffix(name)
+ if nameResult.HasSuffix {
+ baseKey := strings.ToLower(strings.TrimSpace(nameResult.ModelName))
+ if baseKey != "" {
+ if _, exists := out[baseKey]; !exists {
+ out[baseKey] = name
+ }
+ }
+ }
+ }
+}
+
+// SetRetryConfig updates retry attempts and cooldown wait interval.
+func (m *Manager) SetRetryConfig(retry int, maxRetryInterval time.Duration) {
+ if m == nil {
+ return
+ }
+ if retry < 0 {
+ retry = 0
+ }
+ if maxRetryInterval < 0 {
+ maxRetryInterval = 0
+ }
+ m.requestRetry.Store(int32(retry))
+ m.maxRetryInterval.Store(maxRetryInterval.Nanoseconds())
+}
+
+// RegisterExecutor registers a provider executor with the manager.
+func (m *Manager) RegisterExecutor(executor ProviderExecutor) {
+ if executor == nil {
+ return
+ }
+ m.mu.Lock()
+ defer m.mu.Unlock()
+ m.executors[executor.Identifier()] = executor
+}
+
+// UnregisterExecutor removes the executor associated with the provider key.
+func (m *Manager) UnregisterExecutor(provider string) {
+ provider = strings.ToLower(strings.TrimSpace(provider))
+ if provider == "" {
+ return
+ }
+ m.mu.Lock()
+ delete(m.executors, provider)
+ m.mu.Unlock()
+}
+
+// Register inserts a new auth entry into the manager.
+func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) {
+ if auth == nil {
+ return nil, nil
+ }
+ if auth.ID == "" {
+ auth.ID = uuid.NewString()
+ }
+ auth.EnsureIndex()
+ m.mu.Lock()
+ m.auths[auth.ID] = auth.Clone()
+ m.mu.Unlock()
+ m.rebuildAPIKeyModelAliasFromRuntimeConfig()
+ _ = m.persist(ctx, auth)
+ m.hook.OnAuthRegistered(ctx, auth.Clone())
+ return auth.Clone(), nil
+}
+
+// Update replaces an existing auth entry and notifies hooks.
+func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) {
+ if auth == nil || auth.ID == "" {
+ return nil, nil
+ }
+ m.mu.Lock()
+ if existing, ok := m.auths[auth.ID]; ok && existing != nil && !auth.indexAssigned && auth.Index == "" {
+ auth.Index = existing.Index
+ auth.indexAssigned = existing.indexAssigned
+ }
+ auth.EnsureIndex()
+ m.auths[auth.ID] = auth.Clone()
+ m.mu.Unlock()
+ m.rebuildAPIKeyModelAliasFromRuntimeConfig()
+ _ = m.persist(ctx, auth)
+ m.hook.OnAuthUpdated(ctx, auth.Clone())
+ return auth.Clone(), nil
+}
+
+// Load resets manager state from the backing store.
+func (m *Manager) Load(ctx context.Context) error {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+ if m.store == nil {
+ return nil
+ }
+ items, err := m.store.List(ctx)
+ if err != nil {
+ return err
+ }
+ m.auths = make(map[string]*Auth, len(items))
+ for _, auth := range items {
+ if auth == nil || auth.ID == "" {
+ continue
+ }
+ auth.EnsureIndex()
+ m.auths[auth.ID] = auth.Clone()
+ }
+ cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
+ if cfg == nil {
+ cfg = &internalconfig.Config{}
+ }
+ m.rebuildAPIKeyModelAliasLocked(cfg)
+ return nil
+}
+
+// Execute performs a non-streaming execution using the configured selector and executor.
+// It supports multiple providers for the same model and round-robins the starting provider per model.
+func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
+ normalized := m.normalizeProviders(providers)
+ if len(normalized) == 0 {
+ return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
+ }
+
+ _, maxWait := m.retrySettings()
+
+ var lastErr error
+ for attempt := 0; ; attempt++ {
+ resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts)
+ if errExec == nil {
+ return resp, nil
+ }
+ lastErr = errExec
+ wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, req.Model, maxWait)
+ if !shouldRetry {
+ break
+ }
+ if errWait := waitForCooldown(ctx, wait); errWait != nil {
+ return cliproxyexecutor.Response{}, errWait
+ }
+ }
+ if lastErr != nil {
+ return cliproxyexecutor.Response{}, lastErr
+ }
+ return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
+}
+
+// ExecuteCount performs a non-streaming execution using the configured selector and executor.
+// It supports multiple providers for the same model and round-robins the starting provider per model.
+func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
+ normalized := m.normalizeProviders(providers)
+ if len(normalized) == 0 {
+ return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
+ }
+
+ _, maxWait := m.retrySettings()
+
+ var lastErr error
+ for attempt := 0; ; attempt++ {
+ resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts)
+ if errExec == nil {
+ return resp, nil
+ }
+ lastErr = errExec
+ wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, req.Model, maxWait)
+ if !shouldRetry {
+ break
+ }
+ if errWait := waitForCooldown(ctx, wait); errWait != nil {
+ return cliproxyexecutor.Response{}, errWait
+ }
+ }
+ if lastErr != nil {
+ return cliproxyexecutor.Response{}, lastErr
+ }
+ return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
+}
+
+// ExecuteStream performs a streaming execution using the configured selector and executor.
+// It supports multiple providers for the same model and round-robins the starting provider per model.
+func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (<-chan cliproxyexecutor.StreamChunk, error) {
+ normalized := m.normalizeProviders(providers)
+ if len(normalized) == 0 {
+ return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"}
+ }
+
+ _, maxWait := m.retrySettings()
+
+ var lastErr error
+ for attempt := 0; ; attempt++ {
+ chunks, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts)
+ if errStream == nil {
+ return chunks, nil
+ }
+ lastErr = errStream
+ wait, shouldRetry := m.shouldRetryAfterError(errStream, attempt, normalized, req.Model, maxWait)
+ if !shouldRetry {
+ break
+ }
+ if errWait := waitForCooldown(ctx, wait); errWait != nil {
+ return nil, errWait
+ }
+ }
+ if lastErr != nil {
+ return nil, lastErr
+ }
+ return nil, &Error{Code: "auth_not_found", Message: "no auth available"}
+}
+
+func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
+ if len(providers) == 0 {
+ return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
+ }
+ routeModel := req.Model
+ opts = ensureRequestedModelMetadata(opts, routeModel)
+ tried := make(map[string]struct{})
+ var lastErr error
+ for {
+ auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried)
+ if errPick != nil {
+ if lastErr != nil {
+ return cliproxyexecutor.Response{}, lastErr
+ }
+ return cliproxyexecutor.Response{}, errPick
+ }
+
+ entry := logEntryWithRequestID(ctx)
+ debugLogAuthSelection(entry, auth, provider, req.Model)
+
+ tried[auth.ID] = struct{}{}
+ execCtx := ctx
+ if rt := m.roundTripperFor(auth); rt != nil {
+ execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt)
+ execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt)
+ }
+ execReq := req
+ execReq.Model = rewriteModelForAuth(routeModel, auth)
+ execReq.Model = m.applyOAuthModelAlias(auth, execReq.Model)
+ execReq.Model = m.applyAPIKeyModelAlias(auth, execReq.Model)
+ resp, errExec := executor.Execute(execCtx, auth, execReq, opts)
+ result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: errExec == nil}
+ if errExec != nil {
+ if errCtx := execCtx.Err(); errCtx != nil {
+ return cliproxyexecutor.Response{}, errCtx
+ }
+ result.Error = &Error{Message: errExec.Error()}
+ var se cliproxyexecutor.StatusError
+ if errors.As(errExec, &se) && se != nil {
+ result.Error.HTTPStatus = se.StatusCode()
+ }
+ if ra := retryAfterFromError(errExec); ra != nil {
+ result.RetryAfter = ra
+ }
+ m.MarkResult(execCtx, result)
+ lastErr = errExec
+ continue
+ }
+ m.MarkResult(execCtx, result)
+ return resp, nil
+ }
+}
+
+func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
+ if len(providers) == 0 {
+ return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
+ }
+ routeModel := req.Model
+ opts = ensureRequestedModelMetadata(opts, routeModel)
+ tried := make(map[string]struct{})
+ var lastErr error
+ for {
+ auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried)
+ if errPick != nil {
+ if lastErr != nil {
+ return cliproxyexecutor.Response{}, lastErr
+ }
+ return cliproxyexecutor.Response{}, errPick
+ }
+
+ entry := logEntryWithRequestID(ctx)
+ debugLogAuthSelection(entry, auth, provider, req.Model)
+
+ tried[auth.ID] = struct{}{}
+ execCtx := ctx
+ if rt := m.roundTripperFor(auth); rt != nil {
+ execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt)
+ execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt)
+ }
+ execReq := req
+ execReq.Model = rewriteModelForAuth(routeModel, auth)
+ execReq.Model = m.applyOAuthModelAlias(auth, execReq.Model)
+ execReq.Model = m.applyAPIKeyModelAlias(auth, execReq.Model)
+ resp, errExec := executor.CountTokens(execCtx, auth, execReq, opts)
+ result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: errExec == nil}
+ if errExec != nil {
+ if errCtx := execCtx.Err(); errCtx != nil {
+ return cliproxyexecutor.Response{}, errCtx
+ }
+ result.Error = &Error{Message: errExec.Error()}
+ var se cliproxyexecutor.StatusError
+ if errors.As(errExec, &se) && se != nil {
+ result.Error.HTTPStatus = se.StatusCode()
+ }
+ if ra := retryAfterFromError(errExec); ra != nil {
+ result.RetryAfter = ra
+ }
+ m.MarkResult(execCtx, result)
+ lastErr = errExec
+ continue
+ }
+ m.MarkResult(execCtx, result)
+ return resp, nil
+ }
+}
+
+func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (<-chan cliproxyexecutor.StreamChunk, error) {
+ if len(providers) == 0 {
+ return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"}
+ }
+ routeModel := req.Model
+ opts = ensureRequestedModelMetadata(opts, routeModel)
+ tried := make(map[string]struct{})
+ var lastErr error
+ for {
+ auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried)
+ if errPick != nil {
+ if lastErr != nil {
+ return nil, lastErr
+ }
+ return nil, errPick
+ }
+
+ entry := logEntryWithRequestID(ctx)
+ debugLogAuthSelection(entry, auth, provider, req.Model)
+
+ tried[auth.ID] = struct{}{}
+ execCtx := ctx
+ if rt := m.roundTripperFor(auth); rt != nil {
+ execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt)
+ execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt)
+ }
+ execReq := req
+ execReq.Model = rewriteModelForAuth(routeModel, auth)
+ execReq.Model = m.applyOAuthModelAlias(auth, execReq.Model)
+ execReq.Model = m.applyAPIKeyModelAlias(auth, execReq.Model)
+ chunks, errStream := executor.ExecuteStream(execCtx, auth, execReq, opts)
+ if errStream != nil {
+ if errCtx := execCtx.Err(); errCtx != nil {
+ return nil, errCtx
+ }
+ rerr := &Error{Message: errStream.Error()}
+ var se cliproxyexecutor.StatusError
+ if errors.As(errStream, &se) && se != nil {
+ rerr.HTTPStatus = se.StatusCode()
+ }
+ result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: rerr}
+ result.RetryAfter = retryAfterFromError(errStream)
+ m.MarkResult(execCtx, result)
+ lastErr = errStream
+ continue
+ }
+ out := make(chan cliproxyexecutor.StreamChunk)
+ go func(streamCtx context.Context, streamAuth *Auth, streamProvider string, streamChunks <-chan cliproxyexecutor.StreamChunk) {
+ defer close(out)
+ var failed bool
+ forward := true
+ for chunk := range streamChunks {
+ if chunk.Err != nil && !failed {
+ failed = true
+ rerr := &Error{Message: chunk.Err.Error()}
+ var se cliproxyexecutor.StatusError
+ if errors.As(chunk.Err, &se) && se != nil {
+ rerr.HTTPStatus = se.StatusCode()
+ }
+ m.MarkResult(streamCtx, Result{AuthID: streamAuth.ID, Provider: streamProvider, Model: routeModel, Success: false, Error: rerr})
+ }
+ if !forward {
+ continue
+ }
+ if streamCtx == nil {
+ out <- chunk
+ continue
+ }
+ select {
+ case <-streamCtx.Done():
+ forward = false
+ case out <- chunk:
+ }
+ }
+ if !failed {
+ m.MarkResult(streamCtx, Result{AuthID: streamAuth.ID, Provider: streamProvider, Model: routeModel, Success: true})
+ }
+ }(execCtx, auth.Clone(), provider, chunks)
+ return out, nil
+ }
+}
+
+func ensureRequestedModelMetadata(opts cliproxyexecutor.Options, requestedModel string) cliproxyexecutor.Options {
+ requestedModel = strings.TrimSpace(requestedModel)
+ if requestedModel == "" {
+ return opts
+ }
+ if hasRequestedModelMetadata(opts.Metadata) {
+ return opts
+ }
+ if len(opts.Metadata) == 0 {
+ opts.Metadata = map[string]any{cliproxyexecutor.RequestedModelMetadataKey: requestedModel}
+ return opts
+ }
+ meta := make(map[string]any, len(opts.Metadata)+1)
+ for k, v := range opts.Metadata {
+ meta[k] = v
+ }
+ meta[cliproxyexecutor.RequestedModelMetadataKey] = requestedModel
+ opts.Metadata = meta
+ return opts
+}
+
+func hasRequestedModelMetadata(meta map[string]any) bool {
+ if len(meta) == 0 {
+ return false
+ }
+ raw, ok := meta[cliproxyexecutor.RequestedModelMetadataKey]
+ if !ok || raw == nil {
+ return false
+ }
+ switch v := raw.(type) {
+ case string:
+ return strings.TrimSpace(v) != ""
+ case []byte:
+ return strings.TrimSpace(string(v)) != ""
+ default:
+ return false
+ }
+}
+
+func rewriteModelForAuth(model string, auth *Auth) string {
+ if auth == nil || model == "" {
+ return model
+ }
+ prefix := strings.TrimSpace(auth.Prefix)
+ if prefix == "" {
+ return model
+ }
+ needle := prefix + "/"
+ if !strings.HasPrefix(model, needle) {
+ return model
+ }
+ return strings.TrimPrefix(model, needle)
+}
+
+func (m *Manager) applyAPIKeyModelAlias(auth *Auth, requestedModel string) string {
+ if m == nil || auth == nil {
+ return requestedModel
+ }
+
+ kind, _ := auth.AccountInfo()
+ if !strings.EqualFold(strings.TrimSpace(kind), "api_key") {
+ return requestedModel
+ }
+
+ requestedModel = strings.TrimSpace(requestedModel)
+ if requestedModel == "" {
+ return requestedModel
+ }
+
+ // Fast path: lookup per-auth mapping table (keyed by auth.ID).
+ if resolved := m.lookupAPIKeyUpstreamModel(auth.ID, requestedModel); resolved != "" {
+ return resolved
+ }
+
+ // Slow path: scan config for the matching credential entry and resolve alias.
+ // This acts as a safety net if mappings are stale or auth.ID is missing.
+ cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
+ if cfg == nil {
+ cfg = &internalconfig.Config{}
+ }
+
+ provider := strings.ToLower(strings.TrimSpace(auth.Provider))
+ upstreamModel := ""
+ switch provider {
+ case "gemini":
+ upstreamModel = resolveUpstreamModelForGeminiAPIKey(cfg, auth, requestedModel)
+ case "claude":
+ upstreamModel = resolveUpstreamModelForClaudeAPIKey(cfg, auth, requestedModel)
+ case "codex":
+ upstreamModel = resolveUpstreamModelForCodexAPIKey(cfg, auth, requestedModel)
+ case "vertex":
+ upstreamModel = resolveUpstreamModelForVertexAPIKey(cfg, auth, requestedModel)
+ default:
+ upstreamModel = resolveUpstreamModelForOpenAICompatAPIKey(cfg, auth, requestedModel)
+ }
+
+ // Return upstream model if found, otherwise return requested model.
+ if upstreamModel != "" {
+ return upstreamModel
+ }
+ return requestedModel
+}
+
+// APIKeyConfigEntry is a generic interface for API key configurations.
+type APIKeyConfigEntry interface {
+ GetAPIKey() string
+ GetBaseURL() string
+}
+
+func resolveAPIKeyConfig[T APIKeyConfigEntry](entries []T, auth *Auth) *T {
+ if auth == nil || len(entries) == 0 {
+ return nil
+ }
+ attrKey, attrBase := "", ""
+ if auth.Attributes != nil {
+ attrKey = strings.TrimSpace(auth.Attributes["api_key"])
+ attrBase = strings.TrimSpace(auth.Attributes["base_url"])
+ }
+ for i := range entries {
+ entry := &entries[i]
+ cfgKey := strings.TrimSpace((*entry).GetAPIKey())
+ cfgBase := strings.TrimSpace((*entry).GetBaseURL())
+ if attrKey != "" && attrBase != "" {
+ if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) {
+ return entry
+ }
+ continue
+ }
+ if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
+ if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
+ return entry
+ }
+ }
+ if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
+ return entry
+ }
+ }
+ if attrKey != "" {
+ for i := range entries {
+ entry := &entries[i]
+ if strings.EqualFold(strings.TrimSpace((*entry).GetAPIKey()), attrKey) {
+ return entry
+ }
+ }
+ }
+ return nil
+}
+
+func resolveGeminiAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.GeminiKey {
+ if cfg == nil {
+ return nil
+ }
+ return resolveAPIKeyConfig(cfg.GeminiKey, auth)
+}
+
+func resolveClaudeAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.ClaudeKey {
+ if cfg == nil {
+ return nil
+ }
+ return resolveAPIKeyConfig(cfg.ClaudeKey, auth)
+}
+
+func resolveCodexAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.CodexKey {
+ if cfg == nil {
+ return nil
+ }
+ return resolveAPIKeyConfig(cfg.CodexKey, auth)
+}
+
+func resolveVertexAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.VertexCompatKey {
+ if cfg == nil {
+ return nil
+ }
+ return resolveAPIKeyConfig(cfg.VertexCompatAPIKey, auth)
+}
+
+func resolveUpstreamModelForGeminiAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
+ entry := resolveGeminiAPIKeyConfig(cfg, auth)
+ if entry == nil {
+ return ""
+ }
+ return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
+}
+
+func resolveUpstreamModelForClaudeAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
+ entry := resolveClaudeAPIKeyConfig(cfg, auth)
+ if entry == nil {
+ return ""
+ }
+ return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
+}
+
+func resolveUpstreamModelForCodexAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
+ entry := resolveCodexAPIKeyConfig(cfg, auth)
+ if entry == nil {
+ return ""
+ }
+ return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
+}
+
+func resolveUpstreamModelForVertexAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
+ entry := resolveVertexAPIKeyConfig(cfg, auth)
+ if entry == nil {
+ return ""
+ }
+ return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
+}
+
+func resolveUpstreamModelForOpenAICompatAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
+ providerKey := ""
+ compatName := ""
+ if auth != nil && len(auth.Attributes) > 0 {
+ providerKey = strings.TrimSpace(auth.Attributes["provider_key"])
+ compatName = strings.TrimSpace(auth.Attributes["compat_name"])
+ }
+ if compatName == "" && !strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") {
+ return ""
+ }
+ entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider)
+ if entry == nil {
+ return ""
+ }
+ return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
+}
+
+type apiKeyModelAliasTable map[string]map[string]string
+
+func resolveOpenAICompatConfig(cfg *internalconfig.Config, providerKey, compatName, authProvider string) *internalconfig.OpenAICompatibility {
+ if cfg == nil {
+ return nil
+ }
+ candidates := make([]string, 0, 3)
+ if v := strings.TrimSpace(compatName); v != "" {
+ candidates = append(candidates, v)
+ }
+ if v := strings.TrimSpace(providerKey); v != "" {
+ candidates = append(candidates, v)
+ }
+ if v := strings.TrimSpace(authProvider); v != "" {
+ candidates = append(candidates, v)
+ }
+ for i := range cfg.OpenAICompatibility {
+ compat := &cfg.OpenAICompatibility[i]
+ for _, candidate := range candidates {
+ if candidate != "" && strings.EqualFold(strings.TrimSpace(candidate), compat.Name) {
+ return compat
+ }
+ }
+ }
+ return nil
+}
+
+func asModelAliasEntries[T interface {
+ GetName() string
+ GetAlias() string
+}](models []T) []modelAliasEntry {
+ if len(models) == 0 {
+ return nil
+ }
+ out := make([]modelAliasEntry, 0, len(models))
+ for i := range models {
+ out = append(out, models[i])
+ }
+ return out
+}
+
+func (m *Manager) normalizeProviders(providers []string) []string {
+ if len(providers) == 0 {
+ return nil
+ }
+ result := make([]string, 0, len(providers))
+ seen := make(map[string]struct{}, len(providers))
+ for _, provider := range providers {
+ p := strings.TrimSpace(strings.ToLower(provider))
+ if p == "" {
+ continue
+ }
+ if _, ok := seen[p]; ok {
+ continue
+ }
+ seen[p] = struct{}{}
+ result = append(result, p)
+ }
+ return result
+}
+
+func (m *Manager) retrySettings() (int, time.Duration) {
+ if m == nil {
+ return 0, 0
+ }
+ return int(m.requestRetry.Load()), time.Duration(m.maxRetryInterval.Load())
+}
+
+func (m *Manager) closestCooldownWait(providers []string, model string, attempt int) (time.Duration, bool) {
+ if m == nil || len(providers) == 0 {
+ return 0, false
+ }
+ now := time.Now()
+ defaultRetry := int(m.requestRetry.Load())
+ if defaultRetry < 0 {
+ defaultRetry = 0
+ }
+ providerSet := make(map[string]struct{}, len(providers))
+ for i := range providers {
+ key := strings.TrimSpace(strings.ToLower(providers[i]))
+ if key == "" {
+ continue
+ }
+ providerSet[key] = struct{}{}
+ }
+ m.mu.RLock()
+ defer m.mu.RUnlock()
+ var (
+ found bool
+ minWait time.Duration
+ )
+ for _, auth := range m.auths {
+ if auth == nil {
+ continue
+ }
+ providerKey := strings.TrimSpace(strings.ToLower(auth.Provider))
+ if _, ok := providerSet[providerKey]; !ok {
+ continue
+ }
+ effectiveRetry := defaultRetry
+ if override, ok := auth.RequestRetryOverride(); ok {
+ effectiveRetry = override
+ }
+ if effectiveRetry < 0 {
+ effectiveRetry = 0
+ }
+ if attempt >= effectiveRetry {
+ continue
+ }
+ blocked, reason, next := isAuthBlockedForModel(auth, model, now)
+ if !blocked || next.IsZero() || reason == blockReasonDisabled {
+ continue
+ }
+ wait := next.Sub(now)
+ if wait < 0 {
+ continue
+ }
+ if !found || wait < minWait {
+ minWait = wait
+ found = true
+ }
+ }
+ return minWait, found
+}
+
+func (m *Manager) shouldRetryAfterError(err error, attempt int, providers []string, model string, maxWait time.Duration) (time.Duration, bool) {
+ if err == nil {
+ return 0, false
+ }
+ if maxWait <= 0 {
+ return 0, false
+ }
+ if status := statusCodeFromError(err); status == http.StatusOK {
+ return 0, false
+ }
+ wait, found := m.closestCooldownWait(providers, model, attempt)
+ if !found || wait > maxWait {
+ return 0, false
+ }
+ return wait, true
+}
+
+func waitForCooldown(ctx context.Context, wait time.Duration) error {
+ if wait <= 0 {
+ return nil
+ }
+ timer := time.NewTimer(wait)
+ defer timer.Stop()
+ select {
+ case <-ctx.Done():
+ return ctx.Err()
+ case <-timer.C:
+ return nil
+ }
+}
+
+// MarkResult records an execution result and notifies hooks.
+func (m *Manager) MarkResult(ctx context.Context, result Result) {
+ if result.AuthID == "" {
+ return
+ }
+
+ shouldResumeModel := false
+ shouldSuspendModel := false
+ suspendReason := ""
+ clearModelQuota := false
+ setModelQuota := false
+
+ m.mu.Lock()
+ if auth, ok := m.auths[result.AuthID]; ok && auth != nil {
+ now := time.Now()
+
+ if result.Success {
+ if result.Model != "" {
+ state := ensureModelState(auth, result.Model)
+ resetModelState(state, now)
+ updateAggregatedAvailability(auth, now)
+ if !hasModelError(auth, now) {
+ auth.LastError = nil
+ auth.StatusMessage = ""
+ auth.Status = StatusActive
+ }
+ auth.UpdatedAt = now
+ shouldResumeModel = true
+ clearModelQuota = true
+ } else {
+ clearAuthStateOnSuccess(auth, now)
+ }
+ } else {
+ if result.Model != "" {
+ state := ensureModelState(auth, result.Model)
+ state.Unavailable = true
+ state.Status = StatusError
+ state.UpdatedAt = now
+ if result.Error != nil {
+ state.LastError = cloneError(result.Error)
+ state.StatusMessage = result.Error.Message
+ auth.LastError = cloneError(result.Error)
+ auth.StatusMessage = result.Error.Message
+ }
+
+ statusCode := statusCodeFromResult(result.Error)
+ switch statusCode {
+ case 401:
+ next := now.Add(30 * time.Minute)
+ state.NextRetryAfter = next
+ suspendReason = "unauthorized"
+ shouldSuspendModel = true
+ case 402, 403:
+ next := now.Add(30 * time.Minute)
+ state.NextRetryAfter = next
+ suspendReason = "payment_required"
+ shouldSuspendModel = true
+ case 404:
+ next := now.Add(12 * time.Hour)
+ state.NextRetryAfter = next
+ suspendReason = "not_found"
+ shouldSuspendModel = true
+ case 429:
+ var next time.Time
+ backoffLevel := state.Quota.BackoffLevel
+ if result.RetryAfter != nil {
+ next = now.Add(*result.RetryAfter)
+ } else {
+ cooldown, nextLevel := nextQuotaCooldown(backoffLevel, quotaCooldownDisabledForAuth(auth))
+ if cooldown > 0 {
+ next = now.Add(cooldown)
+ }
+ backoffLevel = nextLevel
+ }
+ state.NextRetryAfter = next
+ state.Quota = QuotaState{
+ Exceeded: true,
+ Reason: "quota",
+ NextRecoverAt: next,
+ BackoffLevel: backoffLevel,
+ }
+ suspendReason = "quota"
+ shouldSuspendModel = true
+ setModelQuota = true
+ case 408, 500, 502, 503, 504:
+ if quotaCooldownDisabledForAuth(auth) {
+ state.NextRetryAfter = time.Time{}
+ } else {
+ next := now.Add(1 * time.Minute)
+ state.NextRetryAfter = next
+ }
+ default:
+ state.NextRetryAfter = time.Time{}
+ }
+
+ auth.Status = StatusError
+ auth.UpdatedAt = now
+ updateAggregatedAvailability(auth, now)
+ } else {
+ applyAuthFailureState(auth, result.Error, result.RetryAfter, now)
+ }
+ }
+
+ _ = m.persist(ctx, auth)
+ }
+ m.mu.Unlock()
+
+ if clearModelQuota && result.Model != "" {
+ registry.GetGlobalRegistry().ClearModelQuotaExceeded(result.AuthID, result.Model)
+ }
+ if setModelQuota && result.Model != "" {
+ registry.GetGlobalRegistry().SetModelQuotaExceeded(result.AuthID, result.Model)
+ }
+ if shouldResumeModel {
+ registry.GetGlobalRegistry().ResumeClientModel(result.AuthID, result.Model)
+ } else if shouldSuspendModel {
+ registry.GetGlobalRegistry().SuspendClientModel(result.AuthID, result.Model, suspendReason)
+ }
+
+ m.hook.OnResult(ctx, result)
+}
+
+func ensureModelState(auth *Auth, model string) *ModelState {
+ if auth == nil || model == "" {
+ return nil
+ }
+ if auth.ModelStates == nil {
+ auth.ModelStates = make(map[string]*ModelState)
+ }
+ if state, ok := auth.ModelStates[model]; ok && state != nil {
+ return state
+ }
+ state := &ModelState{Status: StatusActive}
+ auth.ModelStates[model] = state
+ return state
+}
+
+func resetModelState(state *ModelState, now time.Time) {
+ if state == nil {
+ return
+ }
+ state.Unavailable = false
+ state.Status = StatusActive
+ state.StatusMessage = ""
+ state.NextRetryAfter = time.Time{}
+ state.LastError = nil
+ state.Quota = QuotaState{}
+ state.UpdatedAt = now
+}
+
+func updateAggregatedAvailability(auth *Auth, now time.Time) {
+ if auth == nil || len(auth.ModelStates) == 0 {
+ return
+ }
+ allUnavailable := true
+ earliestRetry := time.Time{}
+ quotaExceeded := false
+ quotaRecover := time.Time{}
+ maxBackoffLevel := 0
+ for _, state := range auth.ModelStates {
+ if state == nil {
+ continue
+ }
+ stateUnavailable := false
+ if state.Status == StatusDisabled {
+ stateUnavailable = true
+ } else if state.Unavailable {
+ if state.NextRetryAfter.IsZero() {
+ stateUnavailable = true
+ } else if state.NextRetryAfter.After(now) {
+ stateUnavailable = true
+ if earliestRetry.IsZero() || state.NextRetryAfter.Before(earliestRetry) {
+ earliestRetry = state.NextRetryAfter
+ }
+ } else {
+ state.Unavailable = false
+ state.NextRetryAfter = time.Time{}
+ }
+ }
+ if !stateUnavailable {
+ allUnavailable = false
+ }
+ if state.Quota.Exceeded {
+ quotaExceeded = true
+ if quotaRecover.IsZero() || (!state.Quota.NextRecoverAt.IsZero() && state.Quota.NextRecoverAt.Before(quotaRecover)) {
+ quotaRecover = state.Quota.NextRecoverAt
+ }
+ if state.Quota.BackoffLevel > maxBackoffLevel {
+ maxBackoffLevel = state.Quota.BackoffLevel
+ }
+ }
+ }
+ auth.Unavailable = allUnavailable
+ if allUnavailable {
+ auth.NextRetryAfter = earliestRetry
+ } else {
+ auth.NextRetryAfter = time.Time{}
+ }
+ if quotaExceeded {
+ auth.Quota.Exceeded = true
+ auth.Quota.Reason = "quota"
+ auth.Quota.NextRecoverAt = quotaRecover
+ auth.Quota.BackoffLevel = maxBackoffLevel
+ } else {
+ auth.Quota.Exceeded = false
+ auth.Quota.Reason = ""
+ auth.Quota.NextRecoverAt = time.Time{}
+ auth.Quota.BackoffLevel = 0
+ }
+}
+
+func hasModelError(auth *Auth, now time.Time) bool {
+ if auth == nil || len(auth.ModelStates) == 0 {
+ return false
+ }
+ for _, state := range auth.ModelStates {
+ if state == nil {
+ continue
+ }
+ if state.LastError != nil {
+ return true
+ }
+ if state.Status == StatusError {
+ if state.Unavailable && (state.NextRetryAfter.IsZero() || state.NextRetryAfter.After(now)) {
+ return true
+ }
+ }
+ }
+ return false
+}
+
+func clearAuthStateOnSuccess(auth *Auth, now time.Time) {
+ if auth == nil {
+ return
+ }
+ auth.Unavailable = false
+ auth.Status = StatusActive
+ auth.StatusMessage = ""
+ auth.Quota.Exceeded = false
+ auth.Quota.Reason = ""
+ auth.Quota.NextRecoverAt = time.Time{}
+ auth.Quota.BackoffLevel = 0
+ auth.LastError = nil
+ auth.NextRetryAfter = time.Time{}
+ auth.UpdatedAt = now
+}
+
+func cloneError(err *Error) *Error {
+ if err == nil {
+ return nil
+ }
+ return &Error{
+ Code: err.Code,
+ Message: err.Message,
+ Retryable: err.Retryable,
+ HTTPStatus: err.HTTPStatus,
+ }
+}
+
+func statusCodeFromError(err error) int {
+ if err == nil {
+ return 0
+ }
+ type statusCoder interface {
+ StatusCode() int
+ }
+ var sc statusCoder
+ if errors.As(err, &sc) && sc != nil {
+ return sc.StatusCode()
+ }
+ return 0
+}
+
+func retryAfterFromError(err error) *time.Duration {
+ if err == nil {
+ return nil
+ }
+ type retryAfterProvider interface {
+ RetryAfter() *time.Duration
+ }
+ rap, ok := err.(retryAfterProvider)
+ if !ok || rap == nil {
+ return nil
+ }
+ retryAfter := rap.RetryAfter()
+ if retryAfter == nil {
+ return nil
+ }
+ val := *retryAfter
+ return &val
+}
+
+func statusCodeFromResult(err *Error) int {
+ if err == nil {
+ return 0
+ }
+ return err.StatusCode()
+}
+
+func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Duration, now time.Time) {
+ if auth == nil {
+ return
+ }
+ auth.Unavailable = true
+ auth.Status = StatusError
+ auth.UpdatedAt = now
+ if resultErr != nil {
+ auth.LastError = cloneError(resultErr)
+ if resultErr.Message != "" {
+ auth.StatusMessage = resultErr.Message
+ }
+ }
+ statusCode := statusCodeFromResult(resultErr)
+ switch statusCode {
+ case 401:
+ auth.StatusMessage = "unauthorized"
+ auth.NextRetryAfter = now.Add(30 * time.Minute)
+ case 402, 403:
+ auth.StatusMessage = "payment_required"
+ auth.NextRetryAfter = now.Add(30 * time.Minute)
+ case 404:
+ auth.StatusMessage = "not_found"
+ auth.NextRetryAfter = now.Add(12 * time.Hour)
+ case 429:
+ auth.StatusMessage = "quota exhausted"
+ auth.Quota.Exceeded = true
+ auth.Quota.Reason = "quota"
+ var next time.Time
+ if retryAfter != nil {
+ next = now.Add(*retryAfter)
+ } else {
+ cooldown, nextLevel := nextQuotaCooldown(auth.Quota.BackoffLevel, quotaCooldownDisabledForAuth(auth))
+ if cooldown > 0 {
+ next = now.Add(cooldown)
+ }
+ auth.Quota.BackoffLevel = nextLevel
+ }
+ auth.Quota.NextRecoverAt = next
+ auth.NextRetryAfter = next
+ case 408, 500, 502, 503, 504:
+ auth.StatusMessage = "transient upstream error"
+ if quotaCooldownDisabledForAuth(auth) {
+ auth.NextRetryAfter = time.Time{}
+ } else {
+ auth.NextRetryAfter = now.Add(1 * time.Minute)
+ }
+ default:
+ if auth.StatusMessage == "" {
+ auth.StatusMessage = "request failed"
+ }
+ }
+}
+
+// nextQuotaCooldown returns the next cooldown duration and updated backoff level for repeated quota errors.
+func nextQuotaCooldown(prevLevel int, disableCooling bool) (time.Duration, int) {
+ if prevLevel < 0 {
+ prevLevel = 0
+ }
+ if disableCooling {
+ return 0, prevLevel
+ }
+ cooldown := quotaBackoffBase * time.Duration(1<= quotaBackoffMax {
+ return quotaBackoffMax, prevLevel
+ }
+ return cooldown, prevLevel + 1
+}
+
+// List returns all auth entries currently known by the manager.
+func (m *Manager) List() []*Auth {
+ m.mu.RLock()
+ defer m.mu.RUnlock()
+ list := make([]*Auth, 0, len(m.auths))
+ for _, auth := range m.auths {
+ list = append(list, auth.Clone())
+ }
+ return list
+}
+
+// GetByID retrieves an auth entry by its ID.
+
+func (m *Manager) GetByID(id string) (*Auth, bool) {
+ if id == "" {
+ return nil, false
+ }
+ m.mu.RLock()
+ defer m.mu.RUnlock()
+ auth, ok := m.auths[id]
+ if !ok {
+ return nil, false
+ }
+ return auth.Clone(), true
+}
+
+func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) {
+ m.mu.RLock()
+ executor, okExecutor := m.executors[provider]
+ if !okExecutor {
+ m.mu.RUnlock()
+ return nil, nil, &Error{Code: "executor_not_found", Message: "executor not registered"}
+ }
+ candidates := make([]*Auth, 0, len(m.auths))
+ modelKey := strings.TrimSpace(model)
+ // Always use base model name (without thinking suffix) for auth matching.
+ if modelKey != "" {
+ parsed := thinking.ParseSuffix(modelKey)
+ if parsed.ModelName != "" {
+ modelKey = strings.TrimSpace(parsed.ModelName)
+ }
+ }
+ registryRef := registry.GetGlobalRegistry()
+ for _, candidate := range m.auths {
+ if candidate.Provider != provider || candidate.Disabled {
+ continue
+ }
+ if _, used := tried[candidate.ID]; used {
+ continue
+ }
+ if modelKey != "" && registryRef != nil && !registryRef.ClientSupportsModel(candidate.ID, modelKey) {
+ continue
+ }
+ candidates = append(candidates, candidate)
+ }
+ if len(candidates) == 0 {
+ m.mu.RUnlock()
+ return nil, nil, &Error{Code: "auth_not_found", Message: "no auth available"}
+ }
+ selected, errPick := m.selector.Pick(ctx, provider, model, opts, candidates)
+ if errPick != nil {
+ m.mu.RUnlock()
+ return nil, nil, errPick
+ }
+ if selected == nil {
+ m.mu.RUnlock()
+ return nil, nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"}
+ }
+ authCopy := selected.Clone()
+ m.mu.RUnlock()
+ if !selected.indexAssigned {
+ m.mu.Lock()
+ if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned {
+ current.EnsureIndex()
+ authCopy = current.Clone()
+ }
+ m.mu.Unlock()
+ }
+ return authCopy, executor, nil
+}
+
+func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) {
+ providerSet := make(map[string]struct{}, len(providers))
+ for _, provider := range providers {
+ p := strings.TrimSpace(strings.ToLower(provider))
+ if p == "" {
+ continue
+ }
+ providerSet[p] = struct{}{}
+ }
+ if len(providerSet) == 0 {
+ return nil, nil, "", &Error{Code: "provider_not_found", Message: "no provider supplied"}
+ }
+
+ m.mu.RLock()
+ candidates := make([]*Auth, 0, len(m.auths))
+ modelKey := strings.TrimSpace(model)
+ // Always use base model name (without thinking suffix) for auth matching.
+ if modelKey != "" {
+ parsed := thinking.ParseSuffix(modelKey)
+ if parsed.ModelName != "" {
+ modelKey = strings.TrimSpace(parsed.ModelName)
+ }
+ }
+ registryRef := registry.GetGlobalRegistry()
+ for _, candidate := range m.auths {
+ if candidate == nil || candidate.Disabled {
+ continue
+ }
+ providerKey := strings.TrimSpace(strings.ToLower(candidate.Provider))
+ if providerKey == "" {
+ continue
+ }
+ if _, ok := providerSet[providerKey]; !ok {
+ continue
+ }
+ if _, used := tried[candidate.ID]; used {
+ continue
+ }
+ if _, ok := m.executors[providerKey]; !ok {
+ continue
+ }
+ if modelKey != "" && registryRef != nil && !registryRef.ClientSupportsModel(candidate.ID, modelKey) {
+ continue
+ }
+ candidates = append(candidates, candidate)
+ }
+ if len(candidates) == 0 {
+ m.mu.RUnlock()
+ return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"}
+ }
+ selected, errPick := m.selector.Pick(ctx, "mixed", model, opts, candidates)
+ if errPick != nil {
+ m.mu.RUnlock()
+ return nil, nil, "", errPick
+ }
+ if selected == nil {
+ m.mu.RUnlock()
+ return nil, nil, "", &Error{Code: "auth_not_found", Message: "selector returned no auth"}
+ }
+ providerKey := strings.TrimSpace(strings.ToLower(selected.Provider))
+ executor, okExecutor := m.executors[providerKey]
+ if !okExecutor {
+ m.mu.RUnlock()
+ return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"}
+ }
+ authCopy := selected.Clone()
+ m.mu.RUnlock()
+ if !selected.indexAssigned {
+ m.mu.Lock()
+ if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned {
+ current.EnsureIndex()
+ authCopy = current.Clone()
+ }
+ m.mu.Unlock()
+ }
+ return authCopy, executor, providerKey, nil
+}
+
+func (m *Manager) persist(ctx context.Context, auth *Auth) error {
+ if m.store == nil || auth == nil {
+ return nil
+ }
+ if shouldSkipPersist(ctx) {
+ return nil
+ }
+ if auth.Attributes != nil {
+ if v := strings.ToLower(strings.TrimSpace(auth.Attributes["runtime_only"])); v == "true" {
+ return nil
+ }
+ }
+ // Skip persistence when metadata is absent (e.g., runtime-only auths).
+ if auth.Metadata == nil {
+ return nil
+ }
+ _, err := m.store.Save(ctx, auth)
+ return err
+}
+
+// StartAutoRefresh launches a background loop that evaluates auth freshness
+// every few seconds and triggers refresh operations when required.
+// Only one loop is kept alive; starting a new one cancels the previous run.
+func (m *Manager) StartAutoRefresh(parent context.Context, interval time.Duration) {
+ if interval <= 0 || interval > refreshCheckInterval {
+ interval = refreshCheckInterval
+ } else {
+ interval = refreshCheckInterval
+ }
+ if m.refreshCancel != nil {
+ m.refreshCancel()
+ m.refreshCancel = nil
+ }
+ ctx, cancel := context.WithCancel(parent)
+ m.refreshCancel = cancel
+ go func() {
+ ticker := time.NewTicker(interval)
+ defer ticker.Stop()
+ m.checkRefreshes(ctx)
+ for {
+ select {
+ case <-ctx.Done():
+ return
+ case <-ticker.C:
+ m.checkRefreshes(ctx)
+ }
+ }
+ }()
+}
+
+// StopAutoRefresh cancels the background refresh loop, if running.
+func (m *Manager) StopAutoRefresh() {
+ if m.refreshCancel != nil {
+ m.refreshCancel()
+ m.refreshCancel = nil
+ }
+}
+
+func (m *Manager) checkRefreshes(ctx context.Context) {
+ // log.Debugf("checking refreshes")
+ now := time.Now()
+ snapshot := m.snapshotAuths()
+ for _, a := range snapshot {
+ typ, _ := a.AccountInfo()
+ if typ != "api_key" {
+ if !m.shouldRefresh(a, now) {
+ continue
+ }
+ log.Debugf("checking refresh for %s, %s, %s", a.Provider, a.ID, typ)
+
+ if exec := m.executorFor(a.Provider); exec == nil {
+ continue
+ }
+ if !m.markRefreshPending(a.ID, now) {
+ continue
+ }
+ go m.refreshAuth(ctx, a.ID)
+ }
+ }
+}
+
+func (m *Manager) snapshotAuths() []*Auth {
+ m.mu.RLock()
+ defer m.mu.RUnlock()
+ out := make([]*Auth, 0, len(m.auths))
+ for _, a := range m.auths {
+ out = append(out, a.Clone())
+ }
+ return out
+}
+
+func (m *Manager) shouldRefresh(a *Auth, now time.Time) bool {
+ if a == nil || a.Disabled {
+ return false
+ }
+ if !a.NextRefreshAfter.IsZero() && now.Before(a.NextRefreshAfter) {
+ return false
+ }
+ if evaluator, ok := a.Runtime.(RefreshEvaluator); ok && evaluator != nil {
+ return evaluator.ShouldRefresh(now, a)
+ }
+
+ lastRefresh := a.LastRefreshedAt
+ if lastRefresh.IsZero() {
+ if ts, ok := authLastRefreshTimestamp(a); ok {
+ lastRefresh = ts
+ }
+ }
+
+ expiry, hasExpiry := a.ExpirationTime()
+
+ if interval := authPreferredInterval(a); interval > 0 {
+ if hasExpiry && !expiry.IsZero() {
+ if !expiry.After(now) {
+ return true
+ }
+ if expiry.Sub(now) <= interval {
+ return true
+ }
+ }
+ if lastRefresh.IsZero() {
+ return true
+ }
+ return now.Sub(lastRefresh) >= interval
+ }
+
+ provider := strings.ToLower(a.Provider)
+ lead := ProviderRefreshLead(provider, a.Runtime)
+ if lead == nil {
+ return false
+ }
+ if *lead <= 0 {
+ if hasExpiry && !expiry.IsZero() {
+ return now.After(expiry)
+ }
+ return false
+ }
+ if hasExpiry && !expiry.IsZero() {
+ return time.Until(expiry) <= *lead
+ }
+ if !lastRefresh.IsZero() {
+ return now.Sub(lastRefresh) >= *lead
+ }
+ return true
+}
+
+func authPreferredInterval(a *Auth) time.Duration {
+ if a == nil {
+ return 0
+ }
+ if d := durationFromMetadata(a.Metadata, "refresh_interval_seconds", "refreshIntervalSeconds", "refresh_interval", "refreshInterval"); d > 0 {
+ return d
+ }
+ if d := durationFromAttributes(a.Attributes, "refresh_interval_seconds", "refreshIntervalSeconds", "refresh_interval", "refreshInterval"); d > 0 {
+ return d
+ }
+ return 0
+}
+
+func durationFromMetadata(meta map[string]any, keys ...string) time.Duration {
+ if len(meta) == 0 {
+ return 0
+ }
+ for _, key := range keys {
+ if val, ok := meta[key]; ok {
+ if dur := parseDurationValue(val); dur > 0 {
+ return dur
+ }
+ }
+ }
+ return 0
+}
+
+func durationFromAttributes(attrs map[string]string, keys ...string) time.Duration {
+ if len(attrs) == 0 {
+ return 0
+ }
+ for _, key := range keys {
+ if val, ok := attrs[key]; ok {
+ if dur := parseDurationString(val); dur > 0 {
+ return dur
+ }
+ }
+ }
+ return 0
+}
+
+func parseDurationValue(val any) time.Duration {
+ switch v := val.(type) {
+ case time.Duration:
+ if v <= 0 {
+ return 0
+ }
+ return v
+ case int:
+ if v <= 0 {
+ return 0
+ }
+ return time.Duration(v) * time.Second
+ case int32:
+ if v <= 0 {
+ return 0
+ }
+ return time.Duration(v) * time.Second
+ case int64:
+ if v <= 0 {
+ return 0
+ }
+ return time.Duration(v) * time.Second
+ case uint:
+ if v == 0 {
+ return 0
+ }
+ return time.Duration(v) * time.Second
+ case uint32:
+ if v == 0 {
+ return 0
+ }
+ return time.Duration(v) * time.Second
+ case uint64:
+ if v == 0 {
+ return 0
+ }
+ return time.Duration(v) * time.Second
+ case float32:
+ if v <= 0 {
+ return 0
+ }
+ return time.Duration(float64(v) * float64(time.Second))
+ case float64:
+ if v <= 0 {
+ return 0
+ }
+ return time.Duration(v * float64(time.Second))
+ case json.Number:
+ if i, err := v.Int64(); err == nil {
+ if i <= 0 {
+ return 0
+ }
+ return time.Duration(i) * time.Second
+ }
+ if f, err := v.Float64(); err == nil && f > 0 {
+ return time.Duration(f * float64(time.Second))
+ }
+ case string:
+ return parseDurationString(v)
+ }
+ return 0
+}
+
+func parseDurationString(raw string) time.Duration {
+ s := strings.TrimSpace(raw)
+ if s == "" {
+ return 0
+ }
+ if dur, err := time.ParseDuration(s); err == nil && dur > 0 {
+ return dur
+ }
+ if secs, err := strconv.ParseFloat(s, 64); err == nil && secs > 0 {
+ return time.Duration(secs * float64(time.Second))
+ }
+ return 0
+}
+
+func authLastRefreshTimestamp(a *Auth) (time.Time, bool) {
+ if a == nil {
+ return time.Time{}, false
+ }
+ if a.Metadata != nil {
+ if ts, ok := lookupMetadataTime(a.Metadata, "last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"); ok {
+ return ts, true
+ }
+ }
+ if a.Attributes != nil {
+ for _, key := range []string{"last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"} {
+ if val := strings.TrimSpace(a.Attributes[key]); val != "" {
+ if ts, ok := parseTimeValue(val); ok {
+ return ts, true
+ }
+ }
+ }
+ }
+ return time.Time{}, false
+}
+
+func lookupMetadataTime(meta map[string]any, keys ...string) (time.Time, bool) {
+ for _, key := range keys {
+ if val, ok := meta[key]; ok {
+ if ts, ok1 := parseTimeValue(val); ok1 {
+ return ts, true
+ }
+ }
+ }
+ return time.Time{}, false
+}
+
+func (m *Manager) markRefreshPending(id string, now time.Time) bool {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+ auth, ok := m.auths[id]
+ if !ok || auth == nil || auth.Disabled {
+ return false
+ }
+ if !auth.NextRefreshAfter.IsZero() && now.Before(auth.NextRefreshAfter) {
+ return false
+ }
+ auth.NextRefreshAfter = now.Add(refreshPendingBackoff)
+ m.auths[id] = auth
+ return true
+}
+
+func (m *Manager) refreshAuth(ctx context.Context, id string) {
+ if ctx == nil {
+ ctx = context.Background()
+ }
+ m.mu.RLock()
+ auth := m.auths[id]
+ var exec ProviderExecutor
+ if auth != nil {
+ exec = m.executors[auth.Provider]
+ }
+ m.mu.RUnlock()
+ if auth == nil || exec == nil {
+ return
+ }
+ cloned := auth.Clone()
+ updated, err := exec.Refresh(ctx, cloned)
+ if err != nil && errors.Is(err, context.Canceled) {
+ log.Debugf("refresh canceled for %s, %s", auth.Provider, auth.ID)
+ return
+ }
+ log.Debugf("refreshed %s, %s, %v", auth.Provider, auth.ID, err)
+ now := time.Now()
+ if err != nil {
+ m.mu.Lock()
+ if current := m.auths[id]; current != nil {
+ current.NextRefreshAfter = now.Add(refreshFailureBackoff)
+ current.LastError = &Error{Message: err.Error()}
+ m.auths[id] = current
+ }
+ m.mu.Unlock()
+ return
+ }
+ if updated == nil {
+ updated = cloned
+ }
+ // Preserve runtime created by the executor during Refresh.
+ // If executor didn't set one, fall back to the previous runtime.
+ if updated.Runtime == nil {
+ updated.Runtime = auth.Runtime
+ }
+ updated.LastRefreshedAt = now
+ updated.NextRefreshAfter = time.Time{}
+ updated.LastError = nil
+ updated.UpdatedAt = now
+ _, _ = m.Update(ctx, updated)
+}
+
+func (m *Manager) executorFor(provider string) ProviderExecutor {
+ m.mu.RLock()
+ defer m.mu.RUnlock()
+ return m.executors[provider]
+}
+
+// roundTripperContextKey is an unexported context key type to avoid collisions.
+type roundTripperContextKey struct{}
+
+// roundTripperFor retrieves an HTTP RoundTripper for the given auth if a provider is registered.
+func (m *Manager) roundTripperFor(auth *Auth) http.RoundTripper {
+ m.mu.RLock()
+ p := m.rtProvider
+ m.mu.RUnlock()
+ if p == nil || auth == nil {
+ return nil
+ }
+ return p.RoundTripperFor(auth)
+}
+
+// RoundTripperProvider defines a minimal provider of per-auth HTTP transports.
+type RoundTripperProvider interface {
+ RoundTripperFor(auth *Auth) http.RoundTripper
+}
+
+// RequestPreparer is an optional interface that provider executors can implement
+// to mutate outbound HTTP requests with provider credentials.
+type RequestPreparer interface {
+ PrepareRequest(req *http.Request, auth *Auth) error
+}
+
+func executorKeyFromAuth(auth *Auth) string {
+ if auth == nil {
+ return ""
+ }
+ if auth.Attributes != nil {
+ providerKey := strings.TrimSpace(auth.Attributes["provider_key"])
+ compatName := strings.TrimSpace(auth.Attributes["compat_name"])
+ if compatName != "" {
+ if providerKey == "" {
+ providerKey = compatName
+ }
+ return strings.ToLower(providerKey)
+ }
+ }
+ return strings.ToLower(strings.TrimSpace(auth.Provider))
+}
+
+// logEntryWithRequestID returns a logrus entry with request_id field if available in context.
+func logEntryWithRequestID(ctx context.Context) *log.Entry {
+ if ctx == nil {
+ return log.NewEntry(log.StandardLogger())
+ }
+ if reqID := logging.GetRequestID(ctx); reqID != "" {
+ return log.WithField("request_id", reqID)
+ }
+ return log.NewEntry(log.StandardLogger())
+}
+
+func debugLogAuthSelection(entry *log.Entry, auth *Auth, provider string, model string) {
+ if !log.IsLevelEnabled(log.DebugLevel) {
+ return
+ }
+ if entry == nil || auth == nil {
+ return
+ }
+ accountType, accountInfo := auth.AccountInfo()
+ proxyInfo := auth.ProxyInfo()
+ suffix := ""
+ if proxyInfo != "" {
+ suffix = " " + proxyInfo
+ }
+ switch accountType {
+ case "api_key":
+ entry.Debugf("Use API key %s for model %s%s", util.HideAPIKey(accountInfo), model, suffix)
+ case "oauth":
+ ident := formatOauthIdentity(auth, provider, accountInfo)
+ entry.Debugf("Use OAuth %s for model %s%s", ident, model, suffix)
+ }
+}
+
+func formatOauthIdentity(auth *Auth, provider string, accountInfo string) string {
+ if auth == nil {
+ return ""
+ }
+ // Prefer the auth's provider when available.
+ providerName := strings.TrimSpace(auth.Provider)
+ if providerName == "" {
+ providerName = strings.TrimSpace(provider)
+ }
+ // Only log the basename to avoid leaking host paths.
+ // FileName may be unset for some auth backends; fall back to ID.
+ authFile := strings.TrimSpace(auth.FileName)
+ if authFile == "" {
+ authFile = strings.TrimSpace(auth.ID)
+ }
+ if authFile != "" {
+ authFile = filepath.Base(authFile)
+ }
+ parts := make([]string, 0, 3)
+ if providerName != "" {
+ parts = append(parts, "provider="+providerName)
+ }
+ if authFile != "" {
+ parts = append(parts, "auth_file="+authFile)
+ }
+ if len(parts) == 0 {
+ return accountInfo
+ }
+ return strings.Join(parts, " ")
+}
+
+// InjectCredentials delegates per-provider HTTP request preparation when supported.
+// If the registered executor for the auth provider implements RequestPreparer,
+// it will be invoked to modify the request (e.g., add headers).
+func (m *Manager) InjectCredentials(req *http.Request, authID string) error {
+ if req == nil || authID == "" {
+ return nil
+ }
+ m.mu.RLock()
+ a := m.auths[authID]
+ var exec ProviderExecutor
+ if a != nil {
+ exec = m.executors[executorKeyFromAuth(a)]
+ }
+ m.mu.RUnlock()
+ if a == nil || exec == nil {
+ return nil
+ }
+ if p, ok := exec.(RequestPreparer); ok && p != nil {
+ return p.PrepareRequest(req, a)
+ }
+ return nil
+}
+
+// PrepareHttpRequest injects provider credentials into the supplied HTTP request.
+func (m *Manager) PrepareHttpRequest(ctx context.Context, auth *Auth, req *http.Request) error {
+ if m == nil {
+ return &Error{Code: "provider_not_found", Message: "manager is nil"}
+ }
+ if auth == nil {
+ return &Error{Code: "auth_not_found", Message: "auth is nil"}
+ }
+ if req == nil {
+ return &Error{Code: "invalid_request", Message: "http request is nil"}
+ }
+ if ctx != nil {
+ *req = *req.WithContext(ctx)
+ }
+ providerKey := executorKeyFromAuth(auth)
+ if providerKey == "" {
+ return &Error{Code: "provider_not_found", Message: "auth provider is empty"}
+ }
+ exec := m.executorFor(providerKey)
+ if exec == nil {
+ return &Error{Code: "provider_not_found", Message: "executor not registered for provider: " + providerKey}
+ }
+ preparer, ok := exec.(RequestPreparer)
+ if !ok || preparer == nil {
+ return &Error{Code: "not_supported", Message: "executor does not support http request preparation"}
+ }
+ return preparer.PrepareRequest(req, auth)
+}
+
+// NewHttpRequest constructs a new HTTP request and injects provider credentials into it.
+func (m *Manager) NewHttpRequest(ctx context.Context, auth *Auth, method, targetURL string, body []byte, headers http.Header) (*http.Request, error) {
+ if ctx == nil {
+ ctx = context.Background()
+ }
+ method = strings.TrimSpace(method)
+ if method == "" {
+ method = http.MethodGet
+ }
+ var reader io.Reader
+ if body != nil {
+ reader = bytes.NewReader(body)
+ }
+ httpReq, err := http.NewRequestWithContext(ctx, method, targetURL, reader)
+ if err != nil {
+ return nil, err
+ }
+ if headers != nil {
+ httpReq.Header = headers.Clone()
+ }
+ if errPrepare := m.PrepareHttpRequest(ctx, auth, httpReq); errPrepare != nil {
+ return nil, errPrepare
+ }
+ return httpReq, nil
+}
+
+// HttpRequest injects provider credentials into the supplied HTTP request and executes it.
+func (m *Manager) HttpRequest(ctx context.Context, auth *Auth, req *http.Request) (*http.Response, error) {
+ if m == nil {
+ return nil, &Error{Code: "provider_not_found", Message: "manager is nil"}
+ }
+ if auth == nil {
+ return nil, &Error{Code: "auth_not_found", Message: "auth is nil"}
+ }
+ if req == nil {
+ return nil, &Error{Code: "invalid_request", Message: "http request is nil"}
+ }
+ providerKey := executorKeyFromAuth(auth)
+ if providerKey == "" {
+ return nil, &Error{Code: "provider_not_found", Message: "auth provider is empty"}
+ }
+ exec := m.executorFor(providerKey)
+ if exec == nil {
+ return nil, &Error{Code: "provider_not_found", Message: "executor not registered for provider: " + providerKey}
+ }
+ return exec.HttpRequest(ctx, auth, req)
+}
diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..ef39ed829c36bcf995586a2bf4fc08d3755423c1
--- /dev/null
+++ b/sdk/cliproxy/auth/conductor_overrides_test.go
@@ -0,0 +1,97 @@
+package auth
+
+import (
+ "context"
+ "testing"
+ "time"
+)
+
+func TestManager_ShouldRetryAfterError_RespectsAuthRequestRetryOverride(t *testing.T) {
+ m := NewManager(nil, nil, nil)
+ m.SetRetryConfig(3, 30*time.Second)
+
+ model := "test-model"
+ next := time.Now().Add(5 * time.Second)
+
+ auth := &Auth{
+ ID: "auth-1",
+ Provider: "claude",
+ Metadata: map[string]any{
+ "request_retry": float64(0),
+ },
+ ModelStates: map[string]*ModelState{
+ model: {
+ Unavailable: true,
+ Status: StatusError,
+ NextRetryAfter: next,
+ },
+ },
+ }
+ if _, errRegister := m.Register(context.Background(), auth); errRegister != nil {
+ t.Fatalf("register auth: %v", errRegister)
+ }
+
+ _, maxWait := m.retrySettings()
+ wait, shouldRetry := m.shouldRetryAfterError(&Error{HTTPStatus: 500, Message: "boom"}, 0, []string{"claude"}, model, maxWait)
+ if shouldRetry {
+ t.Fatalf("expected shouldRetry=false for request_retry=0, got true (wait=%v)", wait)
+ }
+
+ auth.Metadata["request_retry"] = float64(1)
+ if _, errUpdate := m.Update(context.Background(), auth); errUpdate != nil {
+ t.Fatalf("update auth: %v", errUpdate)
+ }
+
+ wait, shouldRetry = m.shouldRetryAfterError(&Error{HTTPStatus: 500, Message: "boom"}, 0, []string{"claude"}, model, maxWait)
+ if !shouldRetry {
+ t.Fatalf("expected shouldRetry=true for request_retry=1, got false")
+ }
+ if wait <= 0 {
+ t.Fatalf("expected wait > 0, got %v", wait)
+ }
+
+ _, shouldRetry = m.shouldRetryAfterError(&Error{HTTPStatus: 500, Message: "boom"}, 1, []string{"claude"}, model, maxWait)
+ if shouldRetry {
+ t.Fatalf("expected shouldRetry=false on attempt=1 for request_retry=1, got true")
+ }
+}
+
+func TestManager_MarkResult_RespectsAuthDisableCoolingOverride(t *testing.T) {
+ prev := quotaCooldownDisabled.Load()
+ quotaCooldownDisabled.Store(false)
+ t.Cleanup(func() { quotaCooldownDisabled.Store(prev) })
+
+ m := NewManager(nil, nil, nil)
+
+ auth := &Auth{
+ ID: "auth-1",
+ Provider: "claude",
+ Metadata: map[string]any{
+ "disable_cooling": true,
+ },
+ }
+ if _, errRegister := m.Register(context.Background(), auth); errRegister != nil {
+ t.Fatalf("register auth: %v", errRegister)
+ }
+
+ model := "test-model"
+ m.MarkResult(context.Background(), Result{
+ AuthID: "auth-1",
+ Provider: "claude",
+ Model: model,
+ Success: false,
+ Error: &Error{HTTPStatus: 500, Message: "boom"},
+ })
+
+ updated, ok := m.GetByID("auth-1")
+ if !ok || updated == nil {
+ t.Fatalf("expected auth to be present")
+ }
+ state := updated.ModelStates[model]
+ if state == nil {
+ t.Fatalf("expected model state to be present")
+ }
+ if !state.NextRetryAfter.IsZero() {
+ t.Fatalf("expected NextRetryAfter to be zero when disable_cooling=true, got %v", state.NextRetryAfter)
+ }
+}
diff --git a/sdk/cliproxy/auth/errors.go b/sdk/cliproxy/auth/errors.go
new file mode 100644
index 0000000000000000000000000000000000000000..72bca1fcf87181481d2ed5284f1539e6626b6f35
--- /dev/null
+++ b/sdk/cliproxy/auth/errors.go
@@ -0,0 +1,32 @@
+package auth
+
+// Error describes an authentication related failure in a provider agnostic format.
+type Error struct {
+ // Code is a short machine readable identifier.
+ Code string `json:"code,omitempty"`
+ // Message is a human readable description of the failure.
+ Message string `json:"message"`
+ // Retryable indicates whether a retry might fix the issue automatically.
+ Retryable bool `json:"retryable"`
+ // HTTPStatus optionally records an HTTP-like status code for the error.
+ HTTPStatus int `json:"http_status,omitempty"`
+}
+
+// Error implements the error interface.
+func (e *Error) Error() string {
+ if e == nil {
+ return ""
+ }
+ if e.Code == "" {
+ return e.Message
+ }
+ return e.Code + ": " + e.Message
+}
+
+// StatusCode implements optional status accessor for manager decision making.
+func (e *Error) StatusCode() int {
+ if e == nil {
+ return 0
+ }
+ return e.HTTPStatus
+}
diff --git a/sdk/cliproxy/auth/oauth_model_alias.go b/sdk/cliproxy/auth/oauth_model_alias.go
new file mode 100644
index 0000000000000000000000000000000000000000..4111663e9768ca8d6af73e9036be8983ab3c1fe1
--- /dev/null
+++ b/sdk/cliproxy/auth/oauth_model_alias.go
@@ -0,0 +1,253 @@
+package auth
+
+import (
+ "strings"
+
+ internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
+)
+
+type modelAliasEntry interface {
+ GetName() string
+ GetAlias() string
+}
+
+type oauthModelAliasTable struct {
+ // reverse maps channel -> alias (lower) -> original upstream model name.
+ reverse map[string]map[string]string
+}
+
+func compileOAuthModelAliasTable(aliases map[string][]internalconfig.OAuthModelAlias) *oauthModelAliasTable {
+ if len(aliases) == 0 {
+ return &oauthModelAliasTable{}
+ }
+ out := &oauthModelAliasTable{
+ reverse: make(map[string]map[string]string, len(aliases)),
+ }
+ for rawChannel, entries := range aliases {
+ channel := strings.ToLower(strings.TrimSpace(rawChannel))
+ if channel == "" || len(entries) == 0 {
+ continue
+ }
+ rev := make(map[string]string, len(entries))
+ for _, entry := range entries {
+ name := strings.TrimSpace(entry.Name)
+ alias := strings.TrimSpace(entry.Alias)
+ if name == "" || alias == "" {
+ continue
+ }
+ if strings.EqualFold(name, alias) {
+ continue
+ }
+ aliasKey := strings.ToLower(alias)
+ if _, exists := rev[aliasKey]; exists {
+ continue
+ }
+ rev[aliasKey] = name
+ }
+ if len(rev) > 0 {
+ out.reverse[channel] = rev
+ }
+ }
+ if len(out.reverse) == 0 {
+ out.reverse = nil
+ }
+ return out
+}
+
+// SetOAuthModelAlias updates the OAuth model name alias table used during execution.
+// The alias is applied per-auth channel to resolve the upstream model name while keeping the
+// client-visible model name unchanged for translation/response formatting.
+func (m *Manager) SetOAuthModelAlias(aliases map[string][]internalconfig.OAuthModelAlias) {
+ if m == nil {
+ return
+ }
+ table := compileOAuthModelAliasTable(aliases)
+ // atomic.Value requires non-nil store values.
+ if table == nil {
+ table = &oauthModelAliasTable{}
+ }
+ m.oauthModelAlias.Store(table)
+}
+
+// applyOAuthModelAlias resolves the upstream model from OAuth model alias.
+// If an alias exists, the returned model is the upstream model.
+func (m *Manager) applyOAuthModelAlias(auth *Auth, requestedModel string) string {
+ upstreamModel := m.resolveOAuthUpstreamModel(auth, requestedModel)
+ if upstreamModel == "" {
+ return requestedModel
+ }
+ return upstreamModel
+}
+
+func resolveModelAliasFromConfigModels(requestedModel string, models []modelAliasEntry) string {
+ requestedModel = strings.TrimSpace(requestedModel)
+ if requestedModel == "" {
+ return ""
+ }
+ if len(models) == 0 {
+ return ""
+ }
+
+ requestResult := thinking.ParseSuffix(requestedModel)
+ base := requestResult.ModelName
+ candidates := []string{base}
+ if base != requestedModel {
+ candidates = append(candidates, requestedModel)
+ }
+
+ preserveSuffix := func(resolved string) string {
+ resolved = strings.TrimSpace(resolved)
+ if resolved == "" {
+ return ""
+ }
+ if thinking.ParseSuffix(resolved).HasSuffix {
+ return resolved
+ }
+ if requestResult.HasSuffix && requestResult.RawSuffix != "" {
+ return resolved + "(" + requestResult.RawSuffix + ")"
+ }
+ return resolved
+ }
+
+ for i := range models {
+ name := strings.TrimSpace(models[i].GetName())
+ alias := strings.TrimSpace(models[i].GetAlias())
+ for _, candidate := range candidates {
+ if candidate == "" {
+ continue
+ }
+ if alias != "" && strings.EqualFold(alias, candidate) {
+ if name != "" {
+ return preserveSuffix(name)
+ }
+ return preserveSuffix(candidate)
+ }
+ if name != "" && strings.EqualFold(name, candidate) {
+ return preserveSuffix(name)
+ }
+ }
+ }
+ return ""
+}
+
+// resolveOAuthUpstreamModel resolves the upstream model name from OAuth model alias.
+// If an alias exists, returns the original (upstream) model name that corresponds
+// to the requested alias.
+//
+// If the requested model contains a thinking suffix (e.g., "gemini-2.5-pro(8192)"),
+// the suffix is preserved in the returned model name. However, if the alias's
+// original name already contains a suffix, the config suffix takes priority.
+func (m *Manager) resolveOAuthUpstreamModel(auth *Auth, requestedModel string) string {
+ return resolveUpstreamModelFromAliasTable(m, auth, requestedModel, modelAliasChannel(auth))
+}
+
+func resolveUpstreamModelFromAliasTable(m *Manager, auth *Auth, requestedModel, channel string) string {
+ if m == nil || auth == nil {
+ return ""
+ }
+ if channel == "" {
+ return ""
+ }
+
+ // Extract thinking suffix from requested model using ParseSuffix
+ requestResult := thinking.ParseSuffix(requestedModel)
+ baseModel := requestResult.ModelName
+
+ // Candidate keys to match: base model and raw input (handles suffix-parsing edge cases).
+ candidates := []string{baseModel}
+ if baseModel != requestedModel {
+ candidates = append(candidates, requestedModel)
+ }
+
+ raw := m.oauthModelAlias.Load()
+ table, _ := raw.(*oauthModelAliasTable)
+ if table == nil || table.reverse == nil {
+ return ""
+ }
+ rev := table.reverse[channel]
+ if rev == nil {
+ return ""
+ }
+
+ for _, candidate := range candidates {
+ key := strings.ToLower(strings.TrimSpace(candidate))
+ if key == "" {
+ continue
+ }
+ original := strings.TrimSpace(rev[key])
+ if original == "" {
+ continue
+ }
+ if strings.EqualFold(original, baseModel) {
+ return ""
+ }
+
+ // If config already has suffix, it takes priority.
+ if thinking.ParseSuffix(original).HasSuffix {
+ return original
+ }
+ // Preserve user's thinking suffix on the resolved model.
+ if requestResult.HasSuffix && requestResult.RawSuffix != "" {
+ return original + "(" + requestResult.RawSuffix + ")"
+ }
+ return original
+ }
+
+ return ""
+}
+
+// modelAliasChannel extracts the OAuth model alias channel from an Auth object.
+// It determines the provider and auth kind from the Auth's attributes and delegates
+// to OAuthModelAliasChannel for the actual channel resolution.
+func modelAliasChannel(auth *Auth) string {
+ if auth == nil {
+ return ""
+ }
+ provider := strings.ToLower(strings.TrimSpace(auth.Provider))
+ authKind := ""
+ if auth.Attributes != nil {
+ authKind = strings.ToLower(strings.TrimSpace(auth.Attributes["auth_kind"]))
+ }
+ if authKind == "" {
+ if kind, _ := auth.AccountInfo(); strings.EqualFold(kind, "api_key") {
+ authKind = "apikey"
+ }
+ }
+ return OAuthModelAliasChannel(provider, authKind)
+}
+
+// OAuthModelAliasChannel returns the OAuth model alias channel name for a given provider
+// and auth kind. Returns empty string if the provider/authKind combination doesn't support
+// OAuth model alias (e.g., API key authentication).
+//
+// Supported channels: gemini-cli, vertex, aistudio, antigravity, claude, codex, qwen, iflow.
+func OAuthModelAliasChannel(provider, authKind string) string {
+ provider = strings.ToLower(strings.TrimSpace(provider))
+ authKind = strings.ToLower(strings.TrimSpace(authKind))
+ switch provider {
+ case "gemini":
+ // gemini provider uses gemini-api-key config, not oauth-model-alias.
+ // OAuth-based gemini auth is converted to "gemini-cli" by the synthesizer.
+ return ""
+ case "vertex":
+ if authKind == "apikey" {
+ return ""
+ }
+ return "vertex"
+ case "claude":
+ if authKind == "apikey" {
+ return ""
+ }
+ return "claude"
+ case "codex":
+ if authKind == "apikey" {
+ return ""
+ }
+ return "codex"
+ case "gemini-cli", "aistudio", "antigravity", "qwen", "iflow":
+ return provider
+ default:
+ return ""
+ }
+}
diff --git a/sdk/cliproxy/auth/oauth_model_alias_test.go b/sdk/cliproxy/auth/oauth_model_alias_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..6956411c97a4a9ea9d2596e4a97390330cbb1dc8
--- /dev/null
+++ b/sdk/cliproxy/auth/oauth_model_alias_test.go
@@ -0,0 +1,177 @@
+package auth
+
+import (
+ "testing"
+
+ internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+)
+
+func TestResolveOAuthUpstreamModel_SuffixPreservation(t *testing.T) {
+ t.Parallel()
+
+ tests := []struct {
+ name string
+ aliases map[string][]internalconfig.OAuthModelAlias
+ channel string
+ input string
+ want string
+ }{
+ {
+ name: "numeric suffix preserved",
+ aliases: map[string][]internalconfig.OAuthModelAlias{
+ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}},
+ },
+ channel: "gemini-cli",
+ input: "gemini-2.5-pro(8192)",
+ want: "gemini-2.5-pro-exp-03-25(8192)",
+ },
+ {
+ name: "level suffix preserved",
+ aliases: map[string][]internalconfig.OAuthModelAlias{
+ "claude": {{Name: "claude-sonnet-4-5-20250514", Alias: "claude-sonnet-4-5"}},
+ },
+ channel: "claude",
+ input: "claude-sonnet-4-5(high)",
+ want: "claude-sonnet-4-5-20250514(high)",
+ },
+ {
+ name: "no suffix unchanged",
+ aliases: map[string][]internalconfig.OAuthModelAlias{
+ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}},
+ },
+ channel: "gemini-cli",
+ input: "gemini-2.5-pro",
+ want: "gemini-2.5-pro-exp-03-25",
+ },
+ {
+ name: "config suffix takes priority",
+ aliases: map[string][]internalconfig.OAuthModelAlias{
+ "claude": {{Name: "claude-sonnet-4-5-20250514(low)", Alias: "claude-sonnet-4-5"}},
+ },
+ channel: "claude",
+ input: "claude-sonnet-4-5(high)",
+ want: "claude-sonnet-4-5-20250514(low)",
+ },
+ {
+ name: "auto suffix preserved",
+ aliases: map[string][]internalconfig.OAuthModelAlias{
+ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}},
+ },
+ channel: "gemini-cli",
+ input: "gemini-2.5-pro(auto)",
+ want: "gemini-2.5-pro-exp-03-25(auto)",
+ },
+ {
+ name: "none suffix preserved",
+ aliases: map[string][]internalconfig.OAuthModelAlias{
+ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}},
+ },
+ channel: "gemini-cli",
+ input: "gemini-2.5-pro(none)",
+ want: "gemini-2.5-pro-exp-03-25(none)",
+ },
+ {
+ name: "case insensitive alias lookup with suffix",
+ aliases: map[string][]internalconfig.OAuthModelAlias{
+ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "Gemini-2.5-Pro"}},
+ },
+ channel: "gemini-cli",
+ input: "gemini-2.5-pro(high)",
+ want: "gemini-2.5-pro-exp-03-25(high)",
+ },
+ {
+ name: "no alias returns empty",
+ aliases: map[string][]internalconfig.OAuthModelAlias{
+ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}},
+ },
+ channel: "gemini-cli",
+ input: "unknown-model(high)",
+ want: "",
+ },
+ {
+ name: "wrong channel returns empty",
+ aliases: map[string][]internalconfig.OAuthModelAlias{
+ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}},
+ },
+ channel: "claude",
+ input: "gemini-2.5-pro(high)",
+ want: "",
+ },
+ {
+ name: "empty suffix filtered out",
+ aliases: map[string][]internalconfig.OAuthModelAlias{
+ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}},
+ },
+ channel: "gemini-cli",
+ input: "gemini-2.5-pro()",
+ want: "gemini-2.5-pro-exp-03-25",
+ },
+ {
+ name: "incomplete suffix treated as no suffix",
+ aliases: map[string][]internalconfig.OAuthModelAlias{
+ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro(high"}},
+ },
+ channel: "gemini-cli",
+ input: "gemini-2.5-pro(high",
+ want: "gemini-2.5-pro-exp-03-25",
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ t.Parallel()
+
+ mgr := NewManager(nil, nil, nil)
+ mgr.SetConfig(&internalconfig.Config{})
+ mgr.SetOAuthModelAlias(tt.aliases)
+
+ auth := createAuthForChannel(tt.channel)
+ got := mgr.resolveOAuthUpstreamModel(auth, tt.input)
+ if got != tt.want {
+ t.Errorf("resolveOAuthUpstreamModel(%q) = %q, want %q", tt.input, got, tt.want)
+ }
+ })
+ }
+}
+
+func createAuthForChannel(channel string) *Auth {
+ switch channel {
+ case "gemini-cli":
+ return &Auth{Provider: "gemini-cli"}
+ case "claude":
+ return &Auth{Provider: "claude", Attributes: map[string]string{"auth_kind": "oauth"}}
+ case "vertex":
+ return &Auth{Provider: "vertex", Attributes: map[string]string{"auth_kind": "oauth"}}
+ case "codex":
+ return &Auth{Provider: "codex", Attributes: map[string]string{"auth_kind": "oauth"}}
+ case "aistudio":
+ return &Auth{Provider: "aistudio"}
+ case "antigravity":
+ return &Auth{Provider: "antigravity"}
+ case "qwen":
+ return &Auth{Provider: "qwen"}
+ case "iflow":
+ return &Auth{Provider: "iflow"}
+ default:
+ return &Auth{Provider: channel}
+ }
+}
+
+func TestApplyOAuthModelAlias_SuffixPreservation(t *testing.T) {
+ t.Parallel()
+
+ aliases := map[string][]internalconfig.OAuthModelAlias{
+ "gemini-cli": {{Name: "gemini-2.5-pro-exp-03-25", Alias: "gemini-2.5-pro"}},
+ }
+
+ mgr := NewManager(nil, nil, nil)
+ mgr.SetConfig(&internalconfig.Config{})
+ mgr.SetOAuthModelAlias(aliases)
+
+ auth := &Auth{ID: "test-auth-id", Provider: "gemini-cli"}
+
+ resolvedModel := mgr.applyOAuthModelAlias(auth, "gemini-2.5-pro(8192)")
+ if resolvedModel != "gemini-2.5-pro-exp-03-25(8192)" {
+ t.Errorf("applyOAuthModelAlias() model = %q, want %q", resolvedModel, "gemini-2.5-pro-exp-03-25(8192)")
+ }
+}
diff --git a/sdk/cliproxy/auth/persist_policy.go b/sdk/cliproxy/auth/persist_policy.go
new file mode 100644
index 0000000000000000000000000000000000000000..35423c304c95c8a5ac62e83330c6862e7669d814
--- /dev/null
+++ b/sdk/cliproxy/auth/persist_policy.go
@@ -0,0 +1,24 @@
+package auth
+
+import "context"
+
+type skipPersistContextKey struct{}
+
+// WithSkipPersist returns a derived context that disables persistence for Manager Update/Register calls.
+// It is intended for code paths that are reacting to file watcher events, where the file on disk is
+// already the source of truth and persisting again would create a write-back loop.
+func WithSkipPersist(ctx context.Context) context.Context {
+ if ctx == nil {
+ ctx = context.Background()
+ }
+ return context.WithValue(ctx, skipPersistContextKey{}, true)
+}
+
+func shouldSkipPersist(ctx context.Context) bool {
+ if ctx == nil {
+ return false
+ }
+ v := ctx.Value(skipPersistContextKey{})
+ enabled, ok := v.(bool)
+ return ok && enabled
+}
diff --git a/sdk/cliproxy/auth/persist_policy_test.go b/sdk/cliproxy/auth/persist_policy_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..f408c872dcca5b3b2cb764f040cddb122dfe5895
--- /dev/null
+++ b/sdk/cliproxy/auth/persist_policy_test.go
@@ -0,0 +1,62 @@
+package auth
+
+import (
+ "context"
+ "sync/atomic"
+ "testing"
+)
+
+type countingStore struct {
+ saveCount atomic.Int32
+}
+
+func (s *countingStore) List(context.Context) ([]*Auth, error) { return nil, nil }
+
+func (s *countingStore) Save(context.Context, *Auth) (string, error) {
+ s.saveCount.Add(1)
+ return "", nil
+}
+
+func (s *countingStore) Delete(context.Context, string) error { return nil }
+
+func TestWithSkipPersist_DisablesUpdatePersistence(t *testing.T) {
+ store := &countingStore{}
+ mgr := NewManager(store, nil, nil)
+ auth := &Auth{
+ ID: "auth-1",
+ Provider: "antigravity",
+ Metadata: map[string]any{"type": "antigravity"},
+ }
+
+ if _, err := mgr.Update(context.Background(), auth); err != nil {
+ t.Fatalf("Update returned error: %v", err)
+ }
+ if got := store.saveCount.Load(); got != 1 {
+ t.Fatalf("expected 1 Save call, got %d", got)
+ }
+
+ ctxSkip := WithSkipPersist(context.Background())
+ if _, err := mgr.Update(ctxSkip, auth); err != nil {
+ t.Fatalf("Update(skipPersist) returned error: %v", err)
+ }
+ if got := store.saveCount.Load(); got != 1 {
+ t.Fatalf("expected Save call count to remain 1, got %d", got)
+ }
+}
+
+func TestWithSkipPersist_DisablesRegisterPersistence(t *testing.T) {
+ store := &countingStore{}
+ mgr := NewManager(store, nil, nil)
+ auth := &Auth{
+ ID: "auth-1",
+ Provider: "antigravity",
+ Metadata: map[string]any{"type": "antigravity"},
+ }
+
+ if _, err := mgr.Register(WithSkipPersist(context.Background()), auth); err != nil {
+ t.Fatalf("Register(skipPersist) returned error: %v", err)
+ }
+ if got := store.saveCount.Load(); got != 0 {
+ t.Fatalf("expected 0 Save calls, got %d", got)
+ }
+}
diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go
new file mode 100644
index 0000000000000000000000000000000000000000..7febf219da61f91d2e7cd93517ef07ed9bc0a213
--- /dev/null
+++ b/sdk/cliproxy/auth/selector.go
@@ -0,0 +1,267 @@
+package auth
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+ "math"
+ "net/http"
+ "sort"
+ "strconv"
+ "strings"
+ "sync"
+ "time"
+
+ cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
+)
+
+// RoundRobinSelector provides a simple provider scoped round-robin selection strategy.
+type RoundRobinSelector struct {
+ mu sync.Mutex
+ cursors map[string]int
+}
+
+// FillFirstSelector selects the first available credential (deterministic ordering).
+// This "burns" one account before moving to the next, which can help stagger
+// rolling-window subscription caps (e.g. chat message limits).
+type FillFirstSelector struct{}
+
+type blockReason int
+
+const (
+ blockReasonNone blockReason = iota
+ blockReasonCooldown
+ blockReasonDisabled
+ blockReasonOther
+)
+
+type modelCooldownError struct {
+ model string
+ resetIn time.Duration
+ provider string
+}
+
+func newModelCooldownError(model, provider string, resetIn time.Duration) *modelCooldownError {
+ if resetIn < 0 {
+ resetIn = 0
+ }
+ return &modelCooldownError{
+ model: model,
+ provider: provider,
+ resetIn: resetIn,
+ }
+}
+
+func (e *modelCooldownError) Error() string {
+ modelName := e.model
+ if modelName == "" {
+ modelName = "requested model"
+ }
+ message := fmt.Sprintf("All credentials for model %s are cooling down", modelName)
+ if e.provider != "" {
+ message = fmt.Sprintf("%s via provider %s", message, e.provider)
+ }
+ resetSeconds := int(math.Ceil(e.resetIn.Seconds()))
+ if resetSeconds < 0 {
+ resetSeconds = 0
+ }
+ displayDuration := e.resetIn
+ if displayDuration > 0 && displayDuration < time.Second {
+ displayDuration = time.Second
+ } else {
+ displayDuration = displayDuration.Round(time.Second)
+ }
+ errorBody := map[string]any{
+ "code": "model_cooldown",
+ "message": message,
+ "model": e.model,
+ "reset_time": displayDuration.String(),
+ "reset_seconds": resetSeconds,
+ }
+ if e.provider != "" {
+ errorBody["provider"] = e.provider
+ }
+ payload := map[string]any{"error": errorBody}
+ data, err := json.Marshal(payload)
+ if err != nil {
+ return fmt.Sprintf(`{"error":{"code":"model_cooldown","message":"%s"}}`, message)
+ }
+ return string(data)
+}
+
+func (e *modelCooldownError) StatusCode() int {
+ return http.StatusTooManyRequests
+}
+
+func (e *modelCooldownError) Headers() http.Header {
+ headers := make(http.Header)
+ headers.Set("Content-Type", "application/json")
+ resetSeconds := int(math.Ceil(e.resetIn.Seconds()))
+ if resetSeconds < 0 {
+ resetSeconds = 0
+ }
+ headers.Set("Retry-After", strconv.Itoa(resetSeconds))
+ return headers
+}
+
+func authPriority(auth *Auth) int {
+ if auth == nil || auth.Attributes == nil {
+ return 0
+ }
+ raw := strings.TrimSpace(auth.Attributes["priority"])
+ if raw == "" {
+ return 0
+ }
+ parsed, err := strconv.Atoi(raw)
+ if err != nil {
+ return 0
+ }
+ return parsed
+}
+
+func collectAvailableByPriority(auths []*Auth, model string, now time.Time) (available map[int][]*Auth, cooldownCount int, earliest time.Time) {
+ available = make(map[int][]*Auth)
+ for i := 0; i < len(auths); i++ {
+ candidate := auths[i]
+ blocked, reason, next := isAuthBlockedForModel(candidate, model, now)
+ if !blocked {
+ priority := authPriority(candidate)
+ available[priority] = append(available[priority], candidate)
+ continue
+ }
+ if reason == blockReasonCooldown {
+ cooldownCount++
+ if !next.IsZero() && (earliest.IsZero() || next.Before(earliest)) {
+ earliest = next
+ }
+ }
+ }
+ return available, cooldownCount, earliest
+}
+
+func getAvailableAuths(auths []*Auth, provider, model string, now time.Time) ([]*Auth, error) {
+ if len(auths) == 0 {
+ return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"}
+ }
+
+ availableByPriority, cooldownCount, earliest := collectAvailableByPriority(auths, model, now)
+ if len(availableByPriority) == 0 {
+ if cooldownCount == len(auths) && !earliest.IsZero() {
+ providerForError := provider
+ if providerForError == "mixed" {
+ providerForError = ""
+ }
+ resetIn := earliest.Sub(now)
+ if resetIn < 0 {
+ resetIn = 0
+ }
+ return nil, newModelCooldownError(model, providerForError, resetIn)
+ }
+ return nil, &Error{Code: "auth_unavailable", Message: "no auth available"}
+ }
+
+ bestPriority := 0
+ found := false
+ for priority := range availableByPriority {
+ if !found || priority > bestPriority {
+ bestPriority = priority
+ found = true
+ }
+ }
+
+ available := availableByPriority[bestPriority]
+ if len(available) > 1 {
+ sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID })
+ }
+ return available, nil
+}
+
+// Pick selects the next available auth for the provider in a round-robin manner.
+func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
+ _ = ctx
+ _ = opts
+ now := time.Now()
+ available, err := getAvailableAuths(auths, provider, model, now)
+ if err != nil {
+ return nil, err
+ }
+ key := provider + ":" + model
+ s.mu.Lock()
+ if s.cursors == nil {
+ s.cursors = make(map[string]int)
+ }
+ index := s.cursors[key]
+
+ if index >= 2_147_483_640 {
+ index = 0
+ }
+
+ s.cursors[key] = index + 1
+ s.mu.Unlock()
+ // log.Debugf("available: %d, index: %d, key: %d", len(available), index, index%len(available))
+ return available[index%len(available)], nil
+}
+
+// Pick selects the first available auth for the provider in a deterministic manner.
+func (s *FillFirstSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
+ _ = ctx
+ _ = opts
+ now := time.Now()
+ available, err := getAvailableAuths(auths, provider, model, now)
+ if err != nil {
+ return nil, err
+ }
+ return available[0], nil
+}
+
+func isAuthBlockedForModel(auth *Auth, model string, now time.Time) (bool, blockReason, time.Time) {
+ if auth == nil {
+ return true, blockReasonOther, time.Time{}
+ }
+ if auth.Disabled || auth.Status == StatusDisabled {
+ return true, blockReasonDisabled, time.Time{}
+ }
+ if model != "" {
+ if len(auth.ModelStates) > 0 {
+ if state, ok := auth.ModelStates[model]; ok && state != nil {
+ if state.Status == StatusDisabled {
+ return true, blockReasonDisabled, time.Time{}
+ }
+ if state.Unavailable {
+ if state.NextRetryAfter.IsZero() {
+ return false, blockReasonNone, time.Time{}
+ }
+ if state.NextRetryAfter.After(now) {
+ next := state.NextRetryAfter
+ if !state.Quota.NextRecoverAt.IsZero() && state.Quota.NextRecoverAt.After(now) {
+ next = state.Quota.NextRecoverAt
+ }
+ if next.Before(now) {
+ next = now
+ }
+ if state.Quota.Exceeded {
+ return true, blockReasonCooldown, next
+ }
+ return true, blockReasonOther, next
+ }
+ }
+ return false, blockReasonNone, time.Time{}
+ }
+ }
+ return false, blockReasonNone, time.Time{}
+ }
+ if auth.Unavailable && auth.NextRetryAfter.After(now) {
+ next := auth.NextRetryAfter
+ if !auth.Quota.NextRecoverAt.IsZero() && auth.Quota.NextRecoverAt.After(now) {
+ next = auth.Quota.NextRecoverAt
+ }
+ if next.Before(now) {
+ next = now
+ }
+ if auth.Quota.Exceeded {
+ return true, blockReasonCooldown, next
+ }
+ return true, blockReasonOther, next
+ }
+ return false, blockReasonNone, time.Time{}
+}
diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..91a7ed14f073188074b1a95b1b174461a998395c
--- /dev/null
+++ b/sdk/cliproxy/auth/selector_test.go
@@ -0,0 +1,177 @@
+package auth
+
+import (
+ "context"
+ "errors"
+ "sync"
+ "testing"
+ "time"
+
+ cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
+)
+
+func TestFillFirstSelectorPick_Deterministic(t *testing.T) {
+ t.Parallel()
+
+ selector := &FillFirstSelector{}
+ auths := []*Auth{
+ {ID: "b"},
+ {ID: "a"},
+ {ID: "c"},
+ }
+
+ got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths)
+ if err != nil {
+ t.Fatalf("Pick() error = %v", err)
+ }
+ if got == nil {
+ t.Fatalf("Pick() auth = nil")
+ }
+ if got.ID != "a" {
+ t.Fatalf("Pick() auth.ID = %q, want %q", got.ID, "a")
+ }
+}
+
+func TestRoundRobinSelectorPick_CyclesDeterministic(t *testing.T) {
+ t.Parallel()
+
+ selector := &RoundRobinSelector{}
+ auths := []*Auth{
+ {ID: "b"},
+ {ID: "a"},
+ {ID: "c"},
+ }
+
+ want := []string{"a", "b", "c", "a", "b"}
+ for i, id := range want {
+ got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths)
+ if err != nil {
+ t.Fatalf("Pick() #%d error = %v", i, err)
+ }
+ if got == nil {
+ t.Fatalf("Pick() #%d auth = nil", i)
+ }
+ if got.ID != id {
+ t.Fatalf("Pick() #%d auth.ID = %q, want %q", i, got.ID, id)
+ }
+ }
+}
+
+func TestRoundRobinSelectorPick_PriorityBuckets(t *testing.T) {
+ t.Parallel()
+
+ selector := &RoundRobinSelector{}
+ auths := []*Auth{
+ {ID: "c", Attributes: map[string]string{"priority": "0"}},
+ {ID: "a", Attributes: map[string]string{"priority": "10"}},
+ {ID: "b", Attributes: map[string]string{"priority": "10"}},
+ }
+
+ want := []string{"a", "b", "a", "b"}
+ for i, id := range want {
+ got, err := selector.Pick(context.Background(), "mixed", "", cliproxyexecutor.Options{}, auths)
+ if err != nil {
+ t.Fatalf("Pick() #%d error = %v", i, err)
+ }
+ if got == nil {
+ t.Fatalf("Pick() #%d auth = nil", i)
+ }
+ if got.ID != id {
+ t.Fatalf("Pick() #%d auth.ID = %q, want %q", i, got.ID, id)
+ }
+ if got.ID == "c" {
+ t.Fatalf("Pick() #%d unexpectedly selected lower priority auth", i)
+ }
+ }
+}
+
+func TestFillFirstSelectorPick_PriorityFallbackCooldown(t *testing.T) {
+ t.Parallel()
+
+ selector := &FillFirstSelector{}
+ now := time.Now()
+ model := "test-model"
+
+ high := &Auth{
+ ID: "high",
+ Attributes: map[string]string{"priority": "10"},
+ ModelStates: map[string]*ModelState{
+ model: {
+ Status: StatusActive,
+ Unavailable: true,
+ NextRetryAfter: now.Add(30 * time.Minute),
+ Quota: QuotaState{
+ Exceeded: true,
+ },
+ },
+ },
+ }
+ low := &Auth{ID: "low", Attributes: map[string]string{"priority": "0"}}
+
+ got, err := selector.Pick(context.Background(), "mixed", model, cliproxyexecutor.Options{}, []*Auth{high, low})
+ if err != nil {
+ t.Fatalf("Pick() error = %v", err)
+ }
+ if got == nil {
+ t.Fatalf("Pick() auth = nil")
+ }
+ if got.ID != "low" {
+ t.Fatalf("Pick() auth.ID = %q, want %q", got.ID, "low")
+ }
+}
+
+func TestRoundRobinSelectorPick_Concurrent(t *testing.T) {
+ selector := &RoundRobinSelector{}
+ auths := []*Auth{
+ {ID: "b"},
+ {ID: "a"},
+ {ID: "c"},
+ }
+
+ start := make(chan struct{})
+ var wg sync.WaitGroup
+ errCh := make(chan error, 1)
+
+ goroutines := 32
+ iterations := 100
+ for i := 0; i < goroutines; i++ {
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ <-start
+ for j := 0; j < iterations; j++ {
+ got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths)
+ if err != nil {
+ select {
+ case errCh <- err:
+ default:
+ }
+ return
+ }
+ if got == nil {
+ select {
+ case errCh <- errors.New("Pick() returned nil auth"):
+ default:
+ }
+ return
+ }
+ if got.ID == "" {
+ select {
+ case errCh <- errors.New("Pick() returned auth with empty ID"):
+ default:
+ }
+ return
+ }
+ }
+ }()
+ }
+
+ close(start)
+ wg.Wait()
+
+ select {
+ case err := <-errCh:
+ t.Fatalf("concurrent Pick() error = %v", err)
+ default:
+ }
+}
diff --git a/sdk/cliproxy/auth/status.go b/sdk/cliproxy/auth/status.go
new file mode 100644
index 0000000000000000000000000000000000000000..fa60ed82919034ca47f804e041faefcedd69f895
--- /dev/null
+++ b/sdk/cliproxy/auth/status.go
@@ -0,0 +1,19 @@
+package auth
+
+// Status represents the lifecycle state of an Auth entry.
+type Status string
+
+const (
+ // StatusUnknown means the auth state could not be determined.
+ StatusUnknown Status = "unknown"
+ // StatusActive indicates the auth is valid and ready for execution.
+ StatusActive Status = "active"
+ // StatusPending indicates the auth is waiting for an external action, such as MFA.
+ StatusPending Status = "pending"
+ // StatusRefreshing indicates the auth is undergoing a refresh flow.
+ StatusRefreshing Status = "refreshing"
+ // StatusError indicates the auth is temporarily unavailable due to errors.
+ StatusError Status = "error"
+ // StatusDisabled marks the auth as intentionally disabled.
+ StatusDisabled Status = "disabled"
+)
diff --git a/sdk/cliproxy/auth/store.go b/sdk/cliproxy/auth/store.go
new file mode 100644
index 0000000000000000000000000000000000000000..0594a77dd37f1405a2a5c9f6d3437c37b6b7f7de
--- /dev/null
+++ b/sdk/cliproxy/auth/store.go
@@ -0,0 +1,13 @@
+package auth
+
+import "context"
+
+// Store abstracts persistence of Auth state across restarts.
+type Store interface {
+ // List returns all auth records stored in the backend.
+ List(ctx context.Context) ([]*Auth, error)
+ // Save persists the provided auth record, replacing any existing one with same ID.
+ Save(ctx context.Context, auth *Auth) (string, error)
+ // Delete removes the auth record identified by id.
+ Delete(ctx context.Context, id string) error
+}
diff --git a/sdk/cliproxy/auth/types.go b/sdk/cliproxy/auth/types.go
new file mode 100644
index 0000000000000000000000000000000000000000..b2bbe0a2eafccaa4bfd14de62c054f4cc3e49e07
--- /dev/null
+++ b/sdk/cliproxy/auth/types.go
@@ -0,0 +1,479 @@
+package auth
+
+import (
+ "crypto/sha256"
+ "encoding/hex"
+ "encoding/json"
+ "strconv"
+ "strings"
+ "sync"
+ "time"
+
+ baseauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth"
+)
+
+// Auth encapsulates the runtime state and metadata associated with a single credential.
+type Auth struct {
+ // ID uniquely identifies the auth record across restarts.
+ ID string `json:"id"`
+ // Index is a stable runtime identifier derived from auth metadata (not persisted).
+ Index string `json:"-"`
+ // Provider is the upstream provider key (e.g. "gemini", "claude").
+ Provider string `json:"provider"`
+ // Prefix optionally namespaces models for routing (e.g., "teamA/gemini-3-pro-preview").
+ Prefix string `json:"prefix,omitempty"`
+ // FileName stores the relative or absolute path of the backing auth file.
+ FileName string `json:"-"`
+ // Storage holds the token persistence implementation used during login flows.
+ Storage baseauth.TokenStorage `json:"-"`
+ // Label is an optional human readable label for logging.
+ Label string `json:"label,omitempty"`
+ // Status is the lifecycle status managed by the AuthManager.
+ Status Status `json:"status"`
+ // StatusMessage holds a short description for the current status.
+ StatusMessage string `json:"status_message,omitempty"`
+ // Disabled indicates the auth is intentionally disabled by operator.
+ Disabled bool `json:"disabled"`
+ // Unavailable flags transient provider unavailability (e.g. quota exceeded).
+ Unavailable bool `json:"unavailable"`
+ // ProxyURL overrides the global proxy setting for this auth if provided.
+ ProxyURL string `json:"proxy_url,omitempty"`
+ // Attributes stores provider specific metadata needed by executors (immutable configuration).
+ Attributes map[string]string `json:"attributes,omitempty"`
+ // Metadata stores runtime mutable provider state (e.g. tokens, cookies).
+ Metadata map[string]any `json:"metadata,omitempty"`
+ // Quota captures recent quota information for load balancers.
+ Quota QuotaState `json:"quota"`
+ // LastError stores the last failure encountered while executing or refreshing.
+ LastError *Error `json:"last_error,omitempty"`
+ // CreatedAt is the creation timestamp in UTC.
+ CreatedAt time.Time `json:"created_at"`
+ // UpdatedAt is the last modification timestamp in UTC.
+ UpdatedAt time.Time `json:"updated_at"`
+ // LastRefreshedAt records the last successful refresh time in UTC.
+ LastRefreshedAt time.Time `json:"last_refreshed_at"`
+ // NextRefreshAfter is the earliest time a refresh should retrigger.
+ NextRefreshAfter time.Time `json:"next_refresh_after"`
+ // NextRetryAfter is the earliest time a retry should retrigger.
+ NextRetryAfter time.Time `json:"next_retry_after"`
+ // ModelStates tracks per-model runtime availability data.
+ ModelStates map[string]*ModelState `json:"model_states,omitempty"`
+
+ // Runtime carries non-serialisable data used during execution (in-memory only).
+ Runtime any `json:"-"`
+
+ indexAssigned bool `json:"-"`
+}
+
+// QuotaState contains limiter tracking data for a credential.
+type QuotaState struct {
+ // Exceeded indicates the credential recently hit a quota error.
+ Exceeded bool `json:"exceeded"`
+ // Reason provides an optional provider specific human readable description.
+ Reason string `json:"reason,omitempty"`
+ // NextRecoverAt is when the credential may become available again.
+ NextRecoverAt time.Time `json:"next_recover_at"`
+ // BackoffLevel stores the progressive cooldown exponent used for rate limits.
+ BackoffLevel int `json:"backoff_level,omitempty"`
+}
+
+// ModelState captures the execution state for a specific model under an auth entry.
+type ModelState struct {
+ // Status reflects the lifecycle status for this model.
+ Status Status `json:"status"`
+ // StatusMessage provides an optional short description of the status.
+ StatusMessage string `json:"status_message,omitempty"`
+ // Unavailable mirrors whether the model is temporarily blocked for retries.
+ Unavailable bool `json:"unavailable"`
+ // NextRetryAfter defines the per-model retry time.
+ NextRetryAfter time.Time `json:"next_retry_after"`
+ // LastError records the latest error observed for this model.
+ LastError *Error `json:"last_error,omitempty"`
+ // Quota retains quota information if this model hit rate limits.
+ Quota QuotaState `json:"quota"`
+ // UpdatedAt tracks the last update timestamp for this model state.
+ UpdatedAt time.Time `json:"updated_at"`
+}
+
+// Clone shallow copies the Auth structure, duplicating maps to avoid accidental mutation.
+func (a *Auth) Clone() *Auth {
+ if a == nil {
+ return nil
+ }
+ copyAuth := *a
+ if len(a.Attributes) > 0 {
+ copyAuth.Attributes = make(map[string]string, len(a.Attributes))
+ for key, value := range a.Attributes {
+ copyAuth.Attributes[key] = value
+ }
+ }
+ if len(a.Metadata) > 0 {
+ copyAuth.Metadata = make(map[string]any, len(a.Metadata))
+ for key, value := range a.Metadata {
+ copyAuth.Metadata[key] = value
+ }
+ }
+ if len(a.ModelStates) > 0 {
+ copyAuth.ModelStates = make(map[string]*ModelState, len(a.ModelStates))
+ for key, state := range a.ModelStates {
+ copyAuth.ModelStates[key] = state.Clone()
+ }
+ }
+ copyAuth.Runtime = a.Runtime
+ return ©Auth
+}
+
+func stableAuthIndex(seed string) string {
+ seed = strings.TrimSpace(seed)
+ if seed == "" {
+ return ""
+ }
+ sum := sha256.Sum256([]byte(seed))
+ return hex.EncodeToString(sum[:8])
+}
+
+// EnsureIndex returns a stable index derived from the auth file name or API key.
+func (a *Auth) EnsureIndex() string {
+ if a == nil {
+ return ""
+ }
+ if a.indexAssigned && a.Index != "" {
+ return a.Index
+ }
+
+ seed := strings.TrimSpace(a.FileName)
+ if seed != "" {
+ seed = "file:" + seed
+ } else if a.Attributes != nil {
+ if apiKey := strings.TrimSpace(a.Attributes["api_key"]); apiKey != "" {
+ seed = "api_key:" + apiKey
+ }
+ }
+ if seed == "" {
+ if id := strings.TrimSpace(a.ID); id != "" {
+ seed = "id:" + id
+ } else {
+ return ""
+ }
+ }
+
+ idx := stableAuthIndex(seed)
+ a.Index = idx
+ a.indexAssigned = true
+ return idx
+}
+
+// Clone duplicates a model state including nested error details.
+func (m *ModelState) Clone() *ModelState {
+ if m == nil {
+ return nil
+ }
+ copyState := *m
+ if m.LastError != nil {
+ copyState.LastError = &Error{
+ Code: m.LastError.Code,
+ Message: m.LastError.Message,
+ Retryable: m.LastError.Retryable,
+ HTTPStatus: m.LastError.HTTPStatus,
+ }
+ }
+ return ©State
+}
+
+func (a *Auth) ProxyInfo() string {
+ if a == nil {
+ return ""
+ }
+ proxyStr := strings.TrimSpace(a.ProxyURL)
+ if proxyStr == "" {
+ return ""
+ }
+ if idx := strings.Index(proxyStr, "://"); idx > 0 {
+ return "via " + proxyStr[:idx] + " proxy"
+ }
+ return "via proxy"
+}
+
+// DisableCoolingOverride returns the auth-file scoped disable_cooling override when present.
+// The value is read from metadata key "disable_cooling" (or legacy "disable-cooling").
+func (a *Auth) DisableCoolingOverride() (bool, bool) {
+ if a == nil || a.Metadata == nil {
+ return false, false
+ }
+ if val, ok := a.Metadata["disable_cooling"]; ok {
+ if parsed, okParse := parseBoolAny(val); okParse {
+ return parsed, true
+ }
+ }
+ if val, ok := a.Metadata["disable-cooling"]; ok {
+ if parsed, okParse := parseBoolAny(val); okParse {
+ return parsed, true
+ }
+ }
+ return false, false
+}
+
+// RequestRetryOverride returns the auth-file scoped request_retry override when present.
+// The value is read from metadata key "request_retry" (or legacy "request-retry").
+func (a *Auth) RequestRetryOverride() (int, bool) {
+ if a == nil || a.Metadata == nil {
+ return 0, false
+ }
+ if val, ok := a.Metadata["request_retry"]; ok {
+ if parsed, okParse := parseIntAny(val); okParse {
+ if parsed < 0 {
+ parsed = 0
+ }
+ return parsed, true
+ }
+ }
+ if val, ok := a.Metadata["request-retry"]; ok {
+ if parsed, okParse := parseIntAny(val); okParse {
+ if parsed < 0 {
+ parsed = 0
+ }
+ return parsed, true
+ }
+ }
+ return 0, false
+}
+
+func parseBoolAny(val any) (bool, bool) {
+ switch typed := val.(type) {
+ case bool:
+ return typed, true
+ case string:
+ trimmed := strings.TrimSpace(typed)
+ if trimmed == "" {
+ return false, false
+ }
+ parsed, err := strconv.ParseBool(trimmed)
+ if err != nil {
+ return false, false
+ }
+ return parsed, true
+ case float64:
+ return typed != 0, true
+ case json.Number:
+ parsed, err := typed.Int64()
+ if err != nil {
+ return false, false
+ }
+ return parsed != 0, true
+ default:
+ return false, false
+ }
+}
+
+func parseIntAny(val any) (int, bool) {
+ switch typed := val.(type) {
+ case int:
+ return typed, true
+ case int32:
+ return int(typed), true
+ case int64:
+ return int(typed), true
+ case float64:
+ return int(typed), true
+ case json.Number:
+ parsed, err := typed.Int64()
+ if err != nil {
+ return 0, false
+ }
+ return int(parsed), true
+ case string:
+ trimmed := strings.TrimSpace(typed)
+ if trimmed == "" {
+ return 0, false
+ }
+ parsed, err := strconv.Atoi(trimmed)
+ if err != nil {
+ return 0, false
+ }
+ return parsed, true
+ default:
+ return 0, false
+ }
+}
+
+func (a *Auth) AccountInfo() (string, string) {
+ if a == nil {
+ return "", ""
+ }
+ // For Gemini CLI, include project ID in the OAuth account info if present.
+ if strings.ToLower(a.Provider) == "gemini-cli" {
+ if a.Metadata != nil {
+ email, _ := a.Metadata["email"].(string)
+ email = strings.TrimSpace(email)
+ if email != "" {
+ if p, ok := a.Metadata["project_id"].(string); ok {
+ p = strings.TrimSpace(p)
+ if p != "" {
+ return "oauth", email + " (" + p + ")"
+ }
+ }
+ return "oauth", email
+ }
+ }
+ }
+
+ // For iFlow provider, prioritize OAuth type if email is present
+ if strings.ToLower(a.Provider) == "iflow" {
+ if a.Metadata != nil {
+ if email, ok := a.Metadata["email"].(string); ok {
+ email = strings.TrimSpace(email)
+ if email != "" {
+ return "oauth", email
+ }
+ }
+ }
+ }
+
+ // Check metadata for email first (OAuth-style auth)
+ if a.Metadata != nil {
+ if v, ok := a.Metadata["email"].(string); ok {
+ email := strings.TrimSpace(v)
+ if email != "" {
+ return "oauth", email
+ }
+ }
+ }
+ // Fall back to API key (API-key auth)
+ if a.Attributes != nil {
+ if v := a.Attributes["api_key"]; v != "" {
+ return "api_key", v
+ }
+ }
+ return "", ""
+}
+
+// ExpirationTime attempts to extract the credential expiration timestamp from metadata.
+// It inspects common keys such as "expired", "expire", "expires_at", and also
+// nested "token" objects to remain compatible with legacy auth file formats.
+func (a *Auth) ExpirationTime() (time.Time, bool) {
+ if a == nil {
+ return time.Time{}, false
+ }
+ if ts, ok := expirationFromMap(a.Metadata); ok {
+ return ts, true
+ }
+ return time.Time{}, false
+}
+
+var (
+ refreshLeadMu sync.RWMutex
+ refreshLeadFactories = make(map[string]func() *time.Duration)
+)
+
+func RegisterRefreshLeadProvider(provider string, factory func() *time.Duration) {
+ provider = strings.ToLower(strings.TrimSpace(provider))
+ if provider == "" || factory == nil {
+ return
+ }
+ refreshLeadMu.Lock()
+ refreshLeadFactories[provider] = factory
+ refreshLeadMu.Unlock()
+}
+
+var expireKeys = [...]string{"expired", "expire", "expires_at", "expiresAt", "expiry", "expires"}
+
+func expirationFromMap(meta map[string]any) (time.Time, bool) {
+ if meta == nil {
+ return time.Time{}, false
+ }
+ for _, key := range expireKeys {
+ if v, ok := meta[key]; ok {
+ if ts, ok1 := parseTimeValue(v); ok1 {
+ return ts, true
+ }
+ }
+ }
+ for _, nestedKey := range []string{"token", "Token"} {
+ if nested, ok := meta[nestedKey]; ok {
+ switch val := nested.(type) {
+ case map[string]any:
+ if ts, ok1 := expirationFromMap(val); ok1 {
+ return ts, true
+ }
+ case map[string]string:
+ temp := make(map[string]any, len(val))
+ for k, v := range val {
+ temp[k] = v
+ }
+ if ts, ok1 := expirationFromMap(temp); ok1 {
+ return ts, true
+ }
+ }
+ }
+ }
+ return time.Time{}, false
+}
+
+func ProviderRefreshLead(provider string, runtime any) *time.Duration {
+ provider = strings.ToLower(strings.TrimSpace(provider))
+ if runtime != nil {
+ if eval, ok := runtime.(interface{ RefreshLead() *time.Duration }); ok {
+ if lead := eval.RefreshLead(); lead != nil && *lead > 0 {
+ return lead
+ }
+ }
+ }
+ refreshLeadMu.RLock()
+ factory := refreshLeadFactories[provider]
+ refreshLeadMu.RUnlock()
+ if factory == nil {
+ return nil
+ }
+ if lead := factory(); lead != nil && *lead > 0 {
+ return lead
+ }
+ return nil
+}
+
+func parseTimeValue(v any) (time.Time, bool) {
+ switch value := v.(type) {
+ case string:
+ s := strings.TrimSpace(value)
+ if s == "" {
+ return time.Time{}, false
+ }
+ layouts := []string{
+ time.RFC3339,
+ time.RFC3339Nano,
+ "2006-01-02 15:04:05",
+ "2006-01-02 15:04",
+ "2006-01-02T15:04:05Z07:00",
+ }
+ for _, layout := range layouts {
+ if ts, err := time.Parse(layout, s); err == nil {
+ return ts, true
+ }
+ }
+ if unix, err := strconv.ParseInt(s, 10, 64); err == nil {
+ return normaliseUnix(unix), true
+ }
+ case float64:
+ return normaliseUnix(int64(value)), true
+ case int64:
+ return normaliseUnix(value), true
+ case json.Number:
+ if i, err := value.Int64(); err == nil {
+ return normaliseUnix(i), true
+ }
+ if f, err := value.Float64(); err == nil {
+ return normaliseUnix(int64(f)), true
+ }
+ }
+ return time.Time{}, false
+}
+
+func normaliseUnix(raw int64) time.Time {
+ if raw <= 0 {
+ return time.Time{}
+ }
+ // Heuristic: treat values with millisecond precision (>1e12) accordingly.
+ if raw > 1_000_000_000_000 {
+ return time.UnixMilli(raw)
+ }
+ return time.Unix(raw, 0)
+}
diff --git a/sdk/cliproxy/builder.go b/sdk/cliproxy/builder.go
new file mode 100644
index 0000000000000000000000000000000000000000..5eba18a01dfd4ff283d63cac2cd38547c7d4ec02
--- /dev/null
+++ b/sdk/cliproxy/builder.go
@@ -0,0 +1,234 @@
+// Package cliproxy provides the core service implementation for the CLI Proxy API.
+// It includes service lifecycle management, authentication handling, file watching,
+// and integration with various AI service providers through a unified interface.
+package cliproxy
+
+import (
+ "fmt"
+ "strings"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/api"
+ sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
+ sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
+)
+
+// Builder constructs a Service instance with customizable providers.
+// It provides a fluent interface for configuring all aspects of the service
+// including authentication, file watching, HTTP server options, and lifecycle hooks.
+type Builder struct {
+ // cfg holds the application configuration.
+ cfg *config.Config
+
+ // configPath is the path to the configuration file.
+ configPath string
+
+ // tokenProvider handles loading token-based clients.
+ tokenProvider TokenClientProvider
+
+ // apiKeyProvider handles loading API key-based clients.
+ apiKeyProvider APIKeyClientProvider
+
+ // watcherFactory creates file watcher instances.
+ watcherFactory WatcherFactory
+
+ // hooks provides lifecycle callbacks.
+ hooks Hooks
+
+ // authManager handles legacy authentication operations.
+ authManager *sdkAuth.Manager
+
+ // accessManager handles request authentication providers.
+ accessManager *sdkaccess.Manager
+
+ // coreManager handles core authentication and execution.
+ coreManager *coreauth.Manager
+
+ // serverOptions contains additional server configuration options.
+ serverOptions []api.ServerOption
+}
+
+// Hooks allows callers to plug into service lifecycle stages.
+// These callbacks provide opportunities to perform custom initialization
+// and cleanup operations during service startup and shutdown.
+type Hooks struct {
+ // OnBeforeStart is called before the service starts, allowing configuration
+ // modifications or additional setup.
+ OnBeforeStart func(*config.Config)
+
+ // OnAfterStart is called after the service has started successfully,
+ // providing access to the service instance for additional operations.
+ OnAfterStart func(*Service)
+}
+
+// NewBuilder creates a Builder with default dependencies left unset.
+// Use the fluent interface methods to configure the service before calling Build().
+//
+// Returns:
+// - *Builder: A new builder instance ready for configuration
+func NewBuilder() *Builder {
+ return &Builder{}
+}
+
+// WithConfig sets the configuration instance used by the service.
+//
+// Parameters:
+// - cfg: The application configuration
+//
+// Returns:
+// - *Builder: The builder instance for method chaining
+func (b *Builder) WithConfig(cfg *config.Config) *Builder {
+ b.cfg = cfg
+ return b
+}
+
+// WithConfigPath sets the absolute configuration file path used for reload watching.
+//
+// Parameters:
+// - path: The absolute path to the configuration file
+//
+// Returns:
+// - *Builder: The builder instance for method chaining
+func (b *Builder) WithConfigPath(path string) *Builder {
+ b.configPath = path
+ return b
+}
+
+// WithTokenClientProvider overrides the provider responsible for token-backed clients.
+func (b *Builder) WithTokenClientProvider(provider TokenClientProvider) *Builder {
+ b.tokenProvider = provider
+ return b
+}
+
+// WithAPIKeyClientProvider overrides the provider responsible for API key-backed clients.
+func (b *Builder) WithAPIKeyClientProvider(provider APIKeyClientProvider) *Builder {
+ b.apiKeyProvider = provider
+ return b
+}
+
+// WithWatcherFactory allows customizing the watcher factory that handles reloads.
+func (b *Builder) WithWatcherFactory(factory WatcherFactory) *Builder {
+ b.watcherFactory = factory
+ return b
+}
+
+// WithHooks registers lifecycle hooks executed around service startup.
+func (b *Builder) WithHooks(h Hooks) *Builder {
+ b.hooks = h
+ return b
+}
+
+// WithAuthManager overrides the authentication manager used for token lifecycle operations.
+func (b *Builder) WithAuthManager(mgr *sdkAuth.Manager) *Builder {
+ b.authManager = mgr
+ return b
+}
+
+// WithRequestAccessManager overrides the request authentication manager.
+func (b *Builder) WithRequestAccessManager(mgr *sdkaccess.Manager) *Builder {
+ b.accessManager = mgr
+ return b
+}
+
+// WithCoreAuthManager overrides the runtime auth manager responsible for request execution.
+func (b *Builder) WithCoreAuthManager(mgr *coreauth.Manager) *Builder {
+ b.coreManager = mgr
+ return b
+}
+
+// WithServerOptions appends server configuration options used during construction.
+func (b *Builder) WithServerOptions(opts ...api.ServerOption) *Builder {
+ b.serverOptions = append(b.serverOptions, opts...)
+ return b
+}
+
+// WithLocalManagementPassword configures a password that is only accepted from localhost management requests.
+func (b *Builder) WithLocalManagementPassword(password string) *Builder {
+ if password == "" {
+ return b
+ }
+ b.serverOptions = append(b.serverOptions, api.WithLocalManagementPassword(password))
+ return b
+}
+
+// Build validates inputs, applies defaults, and returns a ready-to-run service.
+func (b *Builder) Build() (*Service, error) {
+ if b.cfg == nil {
+ return nil, fmt.Errorf("cliproxy: configuration is required")
+ }
+ if b.configPath == "" {
+ return nil, fmt.Errorf("cliproxy: configuration path is required")
+ }
+
+ tokenProvider := b.tokenProvider
+ if tokenProvider == nil {
+ tokenProvider = NewFileTokenClientProvider()
+ }
+
+ apiKeyProvider := b.apiKeyProvider
+ if apiKeyProvider == nil {
+ apiKeyProvider = NewAPIKeyClientProvider()
+ }
+
+ watcherFactory := b.watcherFactory
+ if watcherFactory == nil {
+ watcherFactory = defaultWatcherFactory
+ }
+
+ authManager := b.authManager
+ if authManager == nil {
+ authManager = newDefaultAuthManager()
+ }
+
+ accessManager := b.accessManager
+ if accessManager == nil {
+ accessManager = sdkaccess.NewManager()
+ }
+
+ providers, err := sdkaccess.BuildProviders(&b.cfg.SDKConfig)
+ if err != nil {
+ return nil, err
+ }
+ accessManager.SetProviders(providers)
+
+ coreManager := b.coreManager
+ if coreManager == nil {
+ tokenStore := sdkAuth.GetTokenStore()
+ if dirSetter, ok := tokenStore.(interface{ SetBaseDir(string) }); ok && b.cfg != nil {
+ dirSetter.SetBaseDir(b.cfg.AuthDir)
+ }
+
+ strategy := ""
+ if b.cfg != nil {
+ strategy = strings.ToLower(strings.TrimSpace(b.cfg.Routing.Strategy))
+ }
+ var selector coreauth.Selector
+ switch strategy {
+ case "fill-first", "fillfirst", "ff":
+ selector = &coreauth.FillFirstSelector{}
+ default:
+ selector = &coreauth.RoundRobinSelector{}
+ }
+
+ coreManager = coreauth.NewManager(tokenStore, selector, nil)
+ }
+ // Attach a default RoundTripper provider so providers can opt-in per-auth transports.
+ coreManager.SetRoundTripperProvider(newDefaultRoundTripperProvider())
+ coreManager.SetConfig(b.cfg)
+ coreManager.SetOAuthModelAlias(b.cfg.OAuthModelAlias)
+
+ service := &Service{
+ cfg: b.cfg,
+ configPath: b.configPath,
+ tokenProvider: tokenProvider,
+ apiKeyProvider: apiKeyProvider,
+ watcherFactory: watcherFactory,
+ hooks: b.hooks,
+ authManager: authManager,
+ accessManager: accessManager,
+ coreManager: coreManager,
+ serverOptions: append([]api.ServerOption(nil), b.serverOptions...),
+ }
+ return service, nil
+}
diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go
new file mode 100644
index 0000000000000000000000000000000000000000..8c11bbc463067ddcc8218e7d62fb2615798d3d96
--- /dev/null
+++ b/sdk/cliproxy/executor/types.go
@@ -0,0 +1,65 @@
+package executor
+
+import (
+ "net/http"
+ "net/url"
+
+ sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
+)
+
+// RequestedModelMetadataKey stores the client-requested model name in Options.Metadata.
+const RequestedModelMetadataKey = "requested_model"
+
+// Request encapsulates the translated payload that will be sent to a provider executor.
+type Request struct {
+ // Model is the upstream model identifier after translation.
+ Model string
+ // Payload is the provider specific JSON payload.
+ Payload []byte
+ // Format represents the provider payload schema.
+ Format sdktranslator.Format
+ // Metadata carries optional provider specific execution hints.
+ Metadata map[string]any
+}
+
+// Options controls execution behavior for both streaming and non-streaming calls.
+type Options struct {
+ // Stream toggles streaming mode.
+ Stream bool
+ // Alt carries optional alternate format hint (e.g. SSE JSON key).
+ Alt string
+ // Headers are forwarded to the provider request builder.
+ Headers http.Header
+ // Query contains optional query string parameters.
+ Query url.Values
+ // OriginalRequest preserves the inbound request bytes prior to translation.
+ OriginalRequest []byte
+ // SourceFormat identifies the inbound schema.
+ SourceFormat sdktranslator.Format
+ // Metadata carries extra execution hints shared across selection and executors.
+ Metadata map[string]any
+}
+
+// Response wraps either a full provider response or metadata for streaming flows.
+type Response struct {
+ // Payload is the provider response in the executor format.
+ Payload []byte
+ // Metadata exposes optional structured data for translators.
+ Metadata map[string]any
+}
+
+// StreamChunk represents a single streaming payload unit emitted by provider executors.
+type StreamChunk struct {
+ // Payload is the raw provider chunk payload.
+ Payload []byte
+ // Err reports any terminal error encountered while producing chunks.
+ Err error
+}
+
+// StatusError represents an error that carries an HTTP-like status code.
+// Provider executors should implement this when possible to enable
+// better auth state updates on failures (e.g., 401/402/429).
+type StatusError interface {
+ error
+ StatusCode() int
+}
diff --git a/sdk/cliproxy/model_registry.go b/sdk/cliproxy/model_registry.go
new file mode 100644
index 0000000000000000000000000000000000000000..01cea5b71583dbcba819359db5d8a2a04db7ab44
--- /dev/null
+++ b/sdk/cliproxy/model_registry.go
@@ -0,0 +1,30 @@
+package cliproxy
+
+import "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
+
+// ModelInfo re-exports the registry model info structure.
+type ModelInfo = registry.ModelInfo
+
+// ModelRegistryHook re-exports the registry hook interface for external integrations.
+type ModelRegistryHook = registry.ModelRegistryHook
+
+// ModelRegistry describes registry operations consumed by external callers.
+type ModelRegistry interface {
+ RegisterClient(clientID, clientProvider string, models []*ModelInfo)
+ UnregisterClient(clientID string)
+ SetModelQuotaExceeded(clientID, modelID string)
+ ClearModelQuotaExceeded(clientID, modelID string)
+ ClientSupportsModel(clientID, modelID string) bool
+ GetAvailableModels(handlerType string) []map[string]any
+ GetAvailableModelsByProvider(provider string) []*ModelInfo
+}
+
+// GlobalModelRegistry returns the shared registry instance.
+func GlobalModelRegistry() ModelRegistry {
+ return registry.GetGlobalRegistry()
+}
+
+// SetGlobalModelRegistryHook registers an optional hook on the shared global registry instance.
+func SetGlobalModelRegistryHook(hook ModelRegistryHook) {
+ registry.GetGlobalRegistry().SetHook(hook)
+}
diff --git a/sdk/cliproxy/pipeline/context.go b/sdk/cliproxy/pipeline/context.go
new file mode 100644
index 0000000000000000000000000000000000000000..fc6754eb977541d72f4da3412b5952845bc24f14
--- /dev/null
+++ b/sdk/cliproxy/pipeline/context.go
@@ -0,0 +1,64 @@
+package pipeline
+
+import (
+ "context"
+ "net/http"
+
+ cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
+ sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
+)
+
+// Context encapsulates execution state shared across middleware, translators, and executors.
+type Context struct {
+ // Request encapsulates the provider facing request payload.
+ Request cliproxyexecutor.Request
+ // Options carries execution flags (streaming, headers, etc.).
+ Options cliproxyexecutor.Options
+ // Auth references the credential selected for execution.
+ Auth *cliproxyauth.Auth
+ // Translator represents the pipeline responsible for schema adaptation.
+ Translator *sdktranslator.Pipeline
+ // HTTPClient allows middleware to customise the outbound transport per request.
+ HTTPClient *http.Client
+}
+
+// Hook captures middleware callbacks around execution.
+type Hook interface {
+ BeforeExecute(ctx context.Context, execCtx *Context)
+ AfterExecute(ctx context.Context, execCtx *Context, resp cliproxyexecutor.Response, err error)
+ OnStreamChunk(ctx context.Context, execCtx *Context, chunk cliproxyexecutor.StreamChunk)
+}
+
+// HookFunc aggregates optional hook implementations.
+type HookFunc struct {
+ Before func(context.Context, *Context)
+ After func(context.Context, *Context, cliproxyexecutor.Response, error)
+ Stream func(context.Context, *Context, cliproxyexecutor.StreamChunk)
+}
+
+// BeforeExecute implements Hook.
+func (h HookFunc) BeforeExecute(ctx context.Context, execCtx *Context) {
+ if h.Before != nil {
+ h.Before(ctx, execCtx)
+ }
+}
+
+// AfterExecute implements Hook.
+func (h HookFunc) AfterExecute(ctx context.Context, execCtx *Context, resp cliproxyexecutor.Response, err error) {
+ if h.After != nil {
+ h.After(ctx, execCtx, resp, err)
+ }
+}
+
+// OnStreamChunk implements Hook.
+func (h HookFunc) OnStreamChunk(ctx context.Context, execCtx *Context, chunk cliproxyexecutor.StreamChunk) {
+ if h.Stream != nil {
+ h.Stream(ctx, execCtx, chunk)
+ }
+}
+
+// RoundTripperProvider allows injection of custom HTTP transports per auth entry.
+type RoundTripperProvider interface {
+ RoundTripperFor(auth *cliproxyauth.Auth) http.RoundTripper
+}
diff --git a/sdk/cliproxy/providers.go b/sdk/cliproxy/providers.go
new file mode 100644
index 0000000000000000000000000000000000000000..7ce89f76fe7744b7112cf39d165dec2eca87ef84
--- /dev/null
+++ b/sdk/cliproxy/providers.go
@@ -0,0 +1,47 @@
+package cliproxy
+
+import (
+ "context"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/watcher"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
+)
+
+// NewFileTokenClientProvider returns the default token-backed client loader.
+func NewFileTokenClientProvider() TokenClientProvider {
+ return &fileTokenClientProvider{}
+}
+
+type fileTokenClientProvider struct{}
+
+func (p *fileTokenClientProvider) Load(ctx context.Context, cfg *config.Config) (*TokenClientResult, error) {
+ // Stateless executors handle tokens
+ _ = ctx
+ _ = cfg
+ return &TokenClientResult{SuccessfulAuthed: 0}, nil
+}
+
+// NewAPIKeyClientProvider returns the default API key client loader that reuses existing logic.
+func NewAPIKeyClientProvider() APIKeyClientProvider {
+ return &apiKeyClientProvider{}
+}
+
+type apiKeyClientProvider struct{}
+
+func (p *apiKeyClientProvider) Load(ctx context.Context, cfg *config.Config) (*APIKeyClientResult, error) {
+ geminiCount, vertexCompatCount, claudeCount, codexCount, openAICompat := watcher.BuildAPIKeyClients(cfg)
+ if ctx != nil {
+ select {
+ case <-ctx.Done():
+ return nil, ctx.Err()
+ default:
+ }
+ }
+ return &APIKeyClientResult{
+ GeminiKeyCount: geminiCount,
+ VertexCompatKeyCount: vertexCompatCount,
+ ClaudeKeyCount: claudeCount,
+ CodexKeyCount: codexCount,
+ OpenAICompatCount: openAICompat,
+ }, nil
+}
diff --git a/sdk/cliproxy/rtprovider.go b/sdk/cliproxy/rtprovider.go
new file mode 100644
index 0000000000000000000000000000000000000000..dad4fc23870484677a2e8f7e5d29f16ca8d3b691
--- /dev/null
+++ b/sdk/cliproxy/rtprovider.go
@@ -0,0 +1,77 @@
+package cliproxy
+
+import (
+ "context"
+ "net"
+ "net/http"
+ "net/url"
+ "strings"
+ "sync"
+
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ log "github.com/sirupsen/logrus"
+ "golang.org/x/net/proxy"
+)
+
+// defaultRoundTripperProvider returns a per-auth HTTP RoundTripper based on
+// the Auth.ProxyURL value. It caches transports per proxy URL string.
+type defaultRoundTripperProvider struct {
+ mu sync.RWMutex
+ cache map[string]http.RoundTripper
+}
+
+func newDefaultRoundTripperProvider() *defaultRoundTripperProvider {
+ return &defaultRoundTripperProvider{cache: make(map[string]http.RoundTripper)}
+}
+
+// RoundTripperFor implements coreauth.RoundTripperProvider.
+func (p *defaultRoundTripperProvider) RoundTripperFor(auth *coreauth.Auth) http.RoundTripper {
+ if auth == nil {
+ return nil
+ }
+ proxyStr := strings.TrimSpace(auth.ProxyURL)
+ if proxyStr == "" {
+ return nil
+ }
+ p.mu.RLock()
+ rt := p.cache[proxyStr]
+ p.mu.RUnlock()
+ if rt != nil {
+ return rt
+ }
+ // Parse the proxy URL to determine the scheme.
+ proxyURL, errParse := url.Parse(proxyStr)
+ if errParse != nil {
+ log.Errorf("parse proxy URL failed: %v", errParse)
+ return nil
+ }
+ var transport *http.Transport
+ // Handle different proxy schemes.
+ if proxyURL.Scheme == "socks5" {
+ // Configure SOCKS5 proxy with optional authentication.
+ username := proxyURL.User.Username()
+ password, _ := proxyURL.User.Password()
+ proxyAuth := &proxy.Auth{User: username, Password: password}
+ dialer, errSOCKS5 := proxy.SOCKS5("tcp", proxyURL.Host, proxyAuth, proxy.Direct)
+ if errSOCKS5 != nil {
+ log.Errorf("create SOCKS5 dialer failed: %v", errSOCKS5)
+ return nil
+ }
+ // Set up a custom transport using the SOCKS5 dialer.
+ transport = &http.Transport{
+ DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+ return dialer.Dial(network, addr)
+ },
+ }
+ } else if proxyURL.Scheme == "http" || proxyURL.Scheme == "https" {
+ // Configure HTTP or HTTPS proxy.
+ transport = &http.Transport{Proxy: http.ProxyURL(proxyURL)}
+ } else {
+ log.Errorf("unsupported proxy scheme: %s", proxyURL.Scheme)
+ return nil
+ }
+ p.mu.Lock()
+ p.cache[proxyStr] = transport
+ p.mu.Unlock()
+ return transport
+}
diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go
new file mode 100644
index 0000000000000000000000000000000000000000..3cd00d5158410bb7fac8b0f699dd65f8203b99c0
--- /dev/null
+++ b/sdk/cliproxy/service.go
@@ -0,0 +1,1337 @@
+// Package cliproxy provides the core service implementation for the CLI Proxy API.
+// It includes service lifecycle management, authentication handling, file watching,
+// and integration with various AI service providers through a unified interface.
+package cliproxy
+
+import (
+ "context"
+ "errors"
+ "fmt"
+ "os"
+ "strings"
+ "sync"
+ "time"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/api"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor"
+ _ "github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/watcher"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/wsrelay"
+ sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
+ sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
+ log "github.com/sirupsen/logrus"
+)
+
+// Service wraps the proxy server lifecycle so external programs can embed the CLI proxy.
+// It manages the complete lifecycle including authentication, file watching, HTTP server,
+// and integration with various AI service providers.
+type Service struct {
+ // cfg holds the current application configuration.
+ cfg *config.Config
+
+ // cfgMu protects concurrent access to the configuration.
+ cfgMu sync.RWMutex
+
+ // configPath is the path to the configuration file.
+ configPath string
+
+ // tokenProvider handles loading token-based clients.
+ tokenProvider TokenClientProvider
+
+ // apiKeyProvider handles loading API key-based clients.
+ apiKeyProvider APIKeyClientProvider
+
+ // watcherFactory creates file watcher instances.
+ watcherFactory WatcherFactory
+
+ // hooks provides lifecycle callbacks.
+ hooks Hooks
+
+ // serverOptions contains additional server configuration options.
+ serverOptions []api.ServerOption
+
+ // server is the HTTP API server instance.
+ server *api.Server
+
+ // serverErr channel for server startup/shutdown errors.
+ serverErr chan error
+
+ // watcher handles file system monitoring.
+ watcher *WatcherWrapper
+
+ // watcherCancel cancels the watcher context.
+ watcherCancel context.CancelFunc
+
+ // authUpdates channel for authentication updates.
+ authUpdates chan watcher.AuthUpdate
+
+ // authQueueStop cancels the auth update queue processing.
+ authQueueStop context.CancelFunc
+
+ // authManager handles legacy authentication operations.
+ authManager *sdkAuth.Manager
+
+ // accessManager handles request authentication providers.
+ accessManager *sdkaccess.Manager
+
+ // coreManager handles core authentication and execution.
+ coreManager *coreauth.Manager
+
+ // shutdownOnce ensures shutdown is called only once.
+ shutdownOnce sync.Once
+
+ // wsGateway manages websocket Gemini providers.
+ wsGateway *wsrelay.Manager
+}
+
+// RegisterUsagePlugin registers a usage plugin on the global usage manager.
+// This allows external code to monitor API usage and token consumption.
+//
+// Parameters:
+// - plugin: The usage plugin to register
+func (s *Service) RegisterUsagePlugin(plugin usage.Plugin) {
+ usage.RegisterPlugin(plugin)
+}
+
+// newDefaultAuthManager creates a default authentication manager with all supported providers.
+func newDefaultAuthManager() *sdkAuth.Manager {
+ return sdkAuth.NewManager(
+ sdkAuth.GetTokenStore(),
+ sdkAuth.NewGeminiAuthenticator(),
+ sdkAuth.NewCodexAuthenticator(),
+ sdkAuth.NewClaudeAuthenticator(),
+ sdkAuth.NewQwenAuthenticator(),
+ )
+}
+
+func (s *Service) ensureAuthUpdateQueue(ctx context.Context) {
+ if s == nil {
+ return
+ }
+ if s.authUpdates == nil {
+ s.authUpdates = make(chan watcher.AuthUpdate, 256)
+ }
+ if s.authQueueStop != nil {
+ return
+ }
+ queueCtx, cancel := context.WithCancel(ctx)
+ s.authQueueStop = cancel
+ go s.consumeAuthUpdates(queueCtx)
+}
+
+func (s *Service) consumeAuthUpdates(ctx context.Context) {
+ ctx = coreauth.WithSkipPersist(ctx)
+ for {
+ select {
+ case <-ctx.Done():
+ return
+ case update, ok := <-s.authUpdates:
+ if !ok {
+ return
+ }
+ s.handleAuthUpdate(ctx, update)
+ labelDrain:
+ for {
+ select {
+ case nextUpdate := <-s.authUpdates:
+ s.handleAuthUpdate(ctx, nextUpdate)
+ default:
+ break labelDrain
+ }
+ }
+ }
+ }
+}
+
+func (s *Service) emitAuthUpdate(ctx context.Context, update watcher.AuthUpdate) {
+ if s == nil {
+ return
+ }
+ if ctx == nil {
+ ctx = context.Background()
+ }
+ if s.watcher != nil && s.watcher.DispatchRuntimeAuthUpdate(update) {
+ return
+ }
+ if s.authUpdates != nil {
+ select {
+ case s.authUpdates <- update:
+ return
+ default:
+ log.Debugf("auth update queue saturated, applying inline action=%v id=%s", update.Action, update.ID)
+ }
+ }
+ s.handleAuthUpdate(ctx, update)
+}
+
+func (s *Service) handleAuthUpdate(ctx context.Context, update watcher.AuthUpdate) {
+ if s == nil {
+ return
+ }
+ s.cfgMu.RLock()
+ cfg := s.cfg
+ s.cfgMu.RUnlock()
+ if cfg == nil || s.coreManager == nil {
+ return
+ }
+ switch update.Action {
+ case watcher.AuthUpdateActionAdd, watcher.AuthUpdateActionModify:
+ if update.Auth == nil || update.Auth.ID == "" {
+ return
+ }
+ s.applyCoreAuthAddOrUpdate(ctx, update.Auth)
+ case watcher.AuthUpdateActionDelete:
+ id := update.ID
+ if id == "" && update.Auth != nil {
+ id = update.Auth.ID
+ }
+ if id == "" {
+ return
+ }
+ s.applyCoreAuthRemoval(ctx, id)
+ default:
+ log.Debugf("received unknown auth update action: %v", update.Action)
+ }
+}
+
+func (s *Service) ensureWebsocketGateway() {
+ if s == nil {
+ return
+ }
+ if s.wsGateway != nil {
+ return
+ }
+ opts := wsrelay.Options{
+ Path: "/v1/ws",
+ OnConnected: s.wsOnConnected,
+ OnDisconnected: s.wsOnDisconnected,
+ LogDebugf: log.Debugf,
+ LogInfof: log.Infof,
+ LogWarnf: log.Warnf,
+ }
+ s.wsGateway = wsrelay.NewManager(opts)
+}
+
+func (s *Service) wsOnConnected(channelID string) {
+ if s == nil || channelID == "" {
+ return
+ }
+ if !strings.HasPrefix(strings.ToLower(channelID), "aistudio-") {
+ return
+ }
+ if s.coreManager != nil {
+ if existing, ok := s.coreManager.GetByID(channelID); ok && existing != nil {
+ if !existing.Disabled && existing.Status == coreauth.StatusActive {
+ return
+ }
+ }
+ }
+ now := time.Now().UTC()
+ auth := &coreauth.Auth{
+ ID: channelID, // keep channel identifier as ID
+ Provider: "aistudio", // logical provider for switch routing
+ Label: channelID, // display original channel id
+ Status: coreauth.StatusActive,
+ CreatedAt: now,
+ UpdatedAt: now,
+ Attributes: map[string]string{"runtime_only": "true"},
+ Metadata: map[string]any{"email": channelID}, // metadata drives logging and usage tracking
+ }
+ log.Infof("websocket provider connected: %s", channelID)
+ s.emitAuthUpdate(context.Background(), watcher.AuthUpdate{
+ Action: watcher.AuthUpdateActionAdd,
+ ID: auth.ID,
+ Auth: auth,
+ })
+}
+
+func (s *Service) wsOnDisconnected(channelID string, reason error) {
+ if s == nil || channelID == "" {
+ return
+ }
+ if reason != nil {
+ if strings.Contains(reason.Error(), "replaced by new connection") {
+ log.Infof("websocket provider replaced: %s", channelID)
+ return
+ }
+ log.Warnf("websocket provider disconnected: %s (%v)", channelID, reason)
+ } else {
+ log.Infof("websocket provider disconnected: %s", channelID)
+ }
+ ctx := context.Background()
+ s.emitAuthUpdate(ctx, watcher.AuthUpdate{
+ Action: watcher.AuthUpdateActionDelete,
+ ID: channelID,
+ })
+}
+
+func (s *Service) applyCoreAuthAddOrUpdate(ctx context.Context, auth *coreauth.Auth) {
+ if s == nil || auth == nil || auth.ID == "" {
+ return
+ }
+ if s.coreManager == nil {
+ return
+ }
+ auth = auth.Clone()
+ s.ensureExecutorsForAuth(auth)
+ s.registerModelsForAuth(auth)
+ if existing, ok := s.coreManager.GetByID(auth.ID); ok && existing != nil {
+ auth.CreatedAt = existing.CreatedAt
+ auth.LastRefreshedAt = existing.LastRefreshedAt
+ auth.NextRefreshAfter = existing.NextRefreshAfter
+ if _, err := s.coreManager.Update(ctx, auth); err != nil {
+ log.Errorf("failed to update auth %s: %v", auth.ID, err)
+ }
+ return
+ }
+ if _, err := s.coreManager.Register(ctx, auth); err != nil {
+ log.Errorf("failed to register auth %s: %v", auth.ID, err)
+ }
+}
+
+func (s *Service) applyCoreAuthRemoval(ctx context.Context, id string) {
+ if s == nil || id == "" {
+ return
+ }
+ if s.coreManager == nil {
+ return
+ }
+ GlobalModelRegistry().UnregisterClient(id)
+ if existing, ok := s.coreManager.GetByID(id); ok && existing != nil {
+ existing.Disabled = true
+ existing.Status = coreauth.StatusDisabled
+ if _, err := s.coreManager.Update(ctx, existing); err != nil {
+ log.Errorf("failed to disable auth %s: %v", id, err)
+ }
+ }
+}
+
+func (s *Service) applyRetryConfig(cfg *config.Config) {
+ if s == nil || s.coreManager == nil || cfg == nil {
+ return
+ }
+ maxInterval := time.Duration(cfg.MaxRetryInterval) * time.Second
+ s.coreManager.SetRetryConfig(cfg.RequestRetry, maxInterval)
+}
+
+func openAICompatInfoFromAuth(a *coreauth.Auth) (providerKey string, compatName string, ok bool) {
+ if a == nil {
+ return "", "", false
+ }
+ if len(a.Attributes) > 0 {
+ providerKey = strings.TrimSpace(a.Attributes["provider_key"])
+ compatName = strings.TrimSpace(a.Attributes["compat_name"])
+ if compatName != "" {
+ if providerKey == "" {
+ providerKey = compatName
+ }
+ return strings.ToLower(providerKey), compatName, true
+ }
+ }
+ if strings.EqualFold(strings.TrimSpace(a.Provider), "openai-compatibility") {
+ return "openai-compatibility", strings.TrimSpace(a.Label), true
+ }
+ return "", "", false
+}
+
+func (s *Service) ensureExecutorsForAuth(a *coreauth.Auth) {
+ if s == nil || a == nil {
+ return
+ }
+ // Skip disabled auth entries when (re)binding executors.
+ // Disabled auths can linger during config reloads (e.g., removed OpenAI-compat entries)
+ // and must not override active provider executors (such as iFlow OAuth accounts).
+ if a.Disabled {
+ return
+ }
+ if compatProviderKey, _, isCompat := openAICompatInfoFromAuth(a); isCompat {
+ if compatProviderKey == "" {
+ compatProviderKey = strings.ToLower(strings.TrimSpace(a.Provider))
+ }
+ if compatProviderKey == "" {
+ compatProviderKey = "openai-compatibility"
+ }
+ s.coreManager.RegisterExecutor(executor.NewOpenAICompatExecutor(compatProviderKey, s.cfg))
+ return
+ }
+ switch strings.ToLower(a.Provider) {
+ case "gemini":
+ s.coreManager.RegisterExecutor(executor.NewGeminiExecutor(s.cfg))
+ case "vertex":
+ s.coreManager.RegisterExecutor(executor.NewGeminiVertexExecutor(s.cfg))
+ case "gemini-cli":
+ s.coreManager.RegisterExecutor(executor.NewGeminiCLIExecutor(s.cfg))
+ case "aistudio":
+ if s.wsGateway != nil {
+ s.coreManager.RegisterExecutor(executor.NewAIStudioExecutor(s.cfg, a.ID, s.wsGateway))
+ }
+ return
+ case "antigravity":
+ s.coreManager.RegisterExecutor(executor.NewAntigravityExecutor(s.cfg))
+ case "claude":
+ s.coreManager.RegisterExecutor(executor.NewClaudeExecutor(s.cfg))
+ case "codex":
+ s.coreManager.RegisterExecutor(executor.NewCodexExecutor(s.cfg))
+ case "qwen":
+ s.coreManager.RegisterExecutor(executor.NewQwenExecutor(s.cfg))
+ case "iflow":
+ s.coreManager.RegisterExecutor(executor.NewIFlowExecutor(s.cfg))
+ case "kiro":
+ s.coreManager.RegisterExecutor(executor.NewKiroExecutor(s.cfg))
+ default:
+ providerKey := strings.ToLower(strings.TrimSpace(a.Provider))
+ if providerKey == "" {
+ providerKey = "openai-compatibility"
+ }
+ s.coreManager.RegisterExecutor(executor.NewOpenAICompatExecutor(providerKey, s.cfg))
+ }
+}
+
+// rebindExecutors refreshes provider executors so they observe the latest configuration.
+func (s *Service) rebindExecutors() {
+ if s == nil || s.coreManager == nil {
+ return
+ }
+ auths := s.coreManager.List()
+ for _, auth := range auths {
+ s.ensureExecutorsForAuth(auth)
+ }
+}
+
+// Run starts the service and blocks until the context is cancelled or the server stops.
+// It initializes all components including authentication, file watching, HTTP server,
+// and starts processing requests. The method blocks until the context is cancelled.
+//
+// Parameters:
+// - ctx: The context for controlling the service lifecycle
+//
+// Returns:
+// - error: An error if the service fails to start or run
+func (s *Service) Run(ctx context.Context) error {
+ if s == nil {
+ return fmt.Errorf("cliproxy: service is nil")
+ }
+ if ctx == nil {
+ ctx = context.Background()
+ }
+
+ usage.StartDefault(ctx)
+
+ shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
+ defer shutdownCancel()
+ defer func() {
+ if err := s.Shutdown(shutdownCtx); err != nil {
+ log.Errorf("service shutdown returned error: %v", err)
+ }
+ }()
+
+ if err := s.ensureAuthDir(); err != nil {
+ return err
+ }
+
+ s.applyRetryConfig(s.cfg)
+
+ if s.coreManager != nil {
+ if errLoad := s.coreManager.Load(ctx); errLoad != nil {
+ log.Warnf("failed to load auth store: %v", errLoad)
+ }
+ }
+
+ tokenResult, err := s.tokenProvider.Load(ctx, s.cfg)
+ if err != nil && !errors.Is(err, context.Canceled) {
+ return err
+ }
+ if tokenResult == nil {
+ tokenResult = &TokenClientResult{}
+ }
+
+ apiKeyResult, err := s.apiKeyProvider.Load(ctx, s.cfg)
+ if err != nil && !errors.Is(err, context.Canceled) {
+ return err
+ }
+ if apiKeyResult == nil {
+ apiKeyResult = &APIKeyClientResult{}
+ }
+
+ // legacy clients removed; no caches to refresh
+
+ // handlers no longer depend on legacy clients; pass nil slice initially
+ s.server = api.NewServer(s.cfg, s.coreManager, s.accessManager, s.configPath, s.serverOptions...)
+
+ if s.authManager == nil {
+ s.authManager = newDefaultAuthManager()
+ }
+
+ s.ensureWebsocketGateway()
+ if s.server != nil && s.wsGateway != nil {
+ s.server.AttachWebsocketRoute(s.wsGateway.Path(), s.wsGateway.Handler())
+ s.server.SetWebsocketAuthChangeHandler(func(oldEnabled, newEnabled bool) {
+ if oldEnabled == newEnabled {
+ return
+ }
+ if !oldEnabled && newEnabled {
+ ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+ defer cancel()
+ if errStop := s.wsGateway.Stop(ctx); errStop != nil {
+ log.Warnf("failed to reset websocket connections after ws-auth change %t -> %t: %v", oldEnabled, newEnabled, errStop)
+ return
+ }
+ log.Debugf("ws-auth enabled; existing websocket sessions terminated to enforce authentication")
+ return
+ }
+ log.Debugf("ws-auth disabled; existing websocket sessions remain connected")
+ })
+ }
+
+ if s.hooks.OnBeforeStart != nil {
+ s.hooks.OnBeforeStart(s.cfg)
+ }
+
+ s.serverErr = make(chan error, 1)
+ go func() {
+ if errStart := s.server.Start(); errStart != nil {
+ s.serverErr <- errStart
+ } else {
+ s.serverErr <- nil
+ }
+ }()
+
+ time.Sleep(100 * time.Millisecond)
+ fmt.Printf("API server started successfully on: %s:%d\n", s.cfg.Host, s.cfg.Port)
+
+ if s.hooks.OnAfterStart != nil {
+ s.hooks.OnAfterStart(s)
+ }
+
+ var watcherWrapper *WatcherWrapper
+ reloadCallback := func(newCfg *config.Config) {
+ previousStrategy := ""
+ s.cfgMu.RLock()
+ if s.cfg != nil {
+ previousStrategy = strings.ToLower(strings.TrimSpace(s.cfg.Routing.Strategy))
+ }
+ s.cfgMu.RUnlock()
+
+ if newCfg == nil {
+ s.cfgMu.RLock()
+ newCfg = s.cfg
+ s.cfgMu.RUnlock()
+ }
+ if newCfg == nil {
+ return
+ }
+
+ nextStrategy := strings.ToLower(strings.TrimSpace(newCfg.Routing.Strategy))
+ normalizeStrategy := func(strategy string) string {
+ switch strategy {
+ case "fill-first", "fillfirst", "ff":
+ return "fill-first"
+ default:
+ return "round-robin"
+ }
+ }
+ previousStrategy = normalizeStrategy(previousStrategy)
+ nextStrategy = normalizeStrategy(nextStrategy)
+ if s.coreManager != nil && previousStrategy != nextStrategy {
+ var selector coreauth.Selector
+ switch nextStrategy {
+ case "fill-first":
+ selector = &coreauth.FillFirstSelector{}
+ default:
+ selector = &coreauth.RoundRobinSelector{}
+ }
+ s.coreManager.SetSelector(selector)
+ log.Infof("routing strategy updated to %s", nextStrategy)
+ }
+
+ s.applyRetryConfig(newCfg)
+ if s.server != nil {
+ s.server.UpdateClients(newCfg)
+ }
+ s.cfgMu.Lock()
+ s.cfg = newCfg
+ s.cfgMu.Unlock()
+ if s.coreManager != nil {
+ s.coreManager.SetConfig(newCfg)
+ s.coreManager.SetOAuthModelAlias(newCfg.OAuthModelAlias)
+ }
+ s.rebindExecutors()
+ }
+
+ watcherWrapper, err = s.watcherFactory(s.configPath, s.cfg.AuthDir, reloadCallback)
+ if err != nil {
+ return fmt.Errorf("cliproxy: failed to create watcher: %w", err)
+ }
+ s.watcher = watcherWrapper
+ s.ensureAuthUpdateQueue(ctx)
+ if s.authUpdates != nil {
+ watcherWrapper.SetAuthUpdateQueue(s.authUpdates)
+ }
+ watcherWrapper.SetConfig(s.cfg)
+
+ watcherCtx, watcherCancel := context.WithCancel(context.Background())
+ s.watcherCancel = watcherCancel
+ if err = watcherWrapper.Start(watcherCtx); err != nil {
+ return fmt.Errorf("cliproxy: failed to start watcher: %w", err)
+ }
+ log.Info("file watcher started for config and auth directory changes")
+
+ // Prefer core auth manager auto refresh if available.
+ if s.coreManager != nil {
+ interval := 15 * time.Minute
+ s.coreManager.StartAutoRefresh(context.Background(), interval)
+ log.Infof("core auth auto-refresh started (interval=%s)", interval)
+ }
+
+ select {
+ case <-ctx.Done():
+ log.Debug("service context cancelled, shutting down...")
+ return ctx.Err()
+ case err = <-s.serverErr:
+ return err
+ }
+}
+
+// Shutdown gracefully stops background workers and the HTTP server.
+// It ensures all resources are properly cleaned up and connections are closed.
+// The shutdown is idempotent and can be called multiple times safely.
+//
+// Parameters:
+// - ctx: The context for controlling the shutdown timeout
+//
+// Returns:
+// - error: An error if shutdown fails
+func (s *Service) Shutdown(ctx context.Context) error {
+ if s == nil {
+ return nil
+ }
+ var shutdownErr error
+ s.shutdownOnce.Do(func() {
+ if ctx == nil {
+ ctx = context.Background()
+ }
+
+ // legacy refresh loop removed; only stopping core auth manager below
+
+ if s.watcherCancel != nil {
+ s.watcherCancel()
+ }
+ if s.coreManager != nil {
+ s.coreManager.StopAutoRefresh()
+ }
+ if s.watcher != nil {
+ if err := s.watcher.Stop(); err != nil {
+ log.Errorf("failed to stop file watcher: %v", err)
+ shutdownErr = err
+ }
+ }
+ if s.wsGateway != nil {
+ if err := s.wsGateway.Stop(ctx); err != nil {
+ log.Errorf("failed to stop websocket gateway: %v", err)
+ if shutdownErr == nil {
+ shutdownErr = err
+ }
+ }
+ }
+ if s.authQueueStop != nil {
+ s.authQueueStop()
+ s.authQueueStop = nil
+ }
+
+ // no legacy clients to persist
+
+ if s.server != nil {
+ shutdownCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+ defer cancel()
+ if err := s.server.Stop(shutdownCtx); err != nil {
+ log.Errorf("error stopping API server: %v", err)
+ if shutdownErr == nil {
+ shutdownErr = err
+ }
+ }
+ }
+
+ usage.StopDefault()
+ })
+ return shutdownErr
+}
+
+func (s *Service) ensureAuthDir() error {
+ info, err := os.Stat(s.cfg.AuthDir)
+ if err != nil {
+ if os.IsNotExist(err) {
+ if mkErr := os.MkdirAll(s.cfg.AuthDir, 0o755); mkErr != nil {
+ return fmt.Errorf("cliproxy: failed to create auth directory %s: %w", s.cfg.AuthDir, mkErr)
+ }
+ log.Infof("created missing auth directory: %s", s.cfg.AuthDir)
+ return nil
+ }
+ return fmt.Errorf("cliproxy: error checking auth directory %s: %w", s.cfg.AuthDir, err)
+ }
+ if !info.IsDir() {
+ return fmt.Errorf("cliproxy: auth path exists but is not a directory: %s", s.cfg.AuthDir)
+ }
+ return nil
+}
+
+// registerModelsForAuth (re)binds provider models in the global registry using the core auth ID as client identifier.
+func (s *Service) registerModelsForAuth(a *coreauth.Auth) {
+ if a == nil || a.ID == "" {
+ return
+ }
+ if a.Disabled {
+ GlobalModelRegistry().UnregisterClient(a.ID)
+ return
+ }
+ authKind := strings.ToLower(strings.TrimSpace(a.Attributes["auth_kind"]))
+ if authKind == "" {
+ if kind, _ := a.AccountInfo(); strings.EqualFold(kind, "api_key") {
+ authKind = "apikey"
+ }
+ }
+ if a.Attributes != nil {
+ if v := strings.TrimSpace(a.Attributes["gemini_virtual_primary"]); strings.EqualFold(v, "true") {
+ GlobalModelRegistry().UnregisterClient(a.ID)
+ return
+ }
+ }
+ // Unregister legacy client ID (if present) to avoid double counting
+ if a.Runtime != nil {
+ if idGetter, ok := a.Runtime.(interface{ GetClientID() string }); ok {
+ if rid := idGetter.GetClientID(); rid != "" && rid != a.ID {
+ GlobalModelRegistry().UnregisterClient(rid)
+ }
+ }
+ }
+ provider := strings.ToLower(strings.TrimSpace(a.Provider))
+ compatProviderKey, compatDisplayName, compatDetected := openAICompatInfoFromAuth(a)
+ if compatDetected {
+ provider = "openai-compatibility"
+ }
+ excluded := s.oauthExcludedModels(provider, authKind)
+ var models []*ModelInfo
+ switch provider {
+ case "gemini":
+ models = registry.GetGeminiModels()
+ if entry := s.resolveConfigGeminiKey(a); entry != nil {
+ if len(entry.Models) > 0 {
+ models = buildGeminiConfigModels(entry)
+ }
+ if authKind == "apikey" {
+ excluded = entry.ExcludedModels
+ }
+ }
+ models = applyExcludedModels(models, excluded)
+ case "vertex":
+ // Vertex AI Gemini supports the same model identifiers as Gemini.
+ models = registry.GetGeminiVertexModels()
+ if authKind == "apikey" {
+ if entry := s.resolveConfigVertexCompatKey(a); entry != nil && len(entry.Models) > 0 {
+ models = buildVertexCompatConfigModels(entry)
+ }
+ }
+ models = applyExcludedModels(models, excluded)
+ case "gemini-cli":
+ models = registry.GetGeminiCLIModels()
+ models = applyExcludedModels(models, excluded)
+ case "aistudio":
+ models = registry.GetAIStudioModels()
+ models = applyExcludedModels(models, excluded)
+ case "antigravity":
+ ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+ models = executor.FetchAntigravityModels(ctx, a, s.cfg)
+ cancel()
+ models = applyExcludedModels(models, excluded)
+ case "claude":
+ models = registry.GetClaudeModels()
+ if entry := s.resolveConfigClaudeKey(a); entry != nil {
+ if len(entry.Models) > 0 {
+ models = buildClaudeConfigModels(entry)
+ }
+ if authKind == "apikey" {
+ excluded = entry.ExcludedModels
+ }
+ }
+ models = applyExcludedModels(models, excluded)
+ case "codex":
+ models = registry.GetOpenAIModels()
+ if entry := s.resolveConfigCodexKey(a); entry != nil {
+ if len(entry.Models) > 0 {
+ models = buildCodexConfigModels(entry)
+ }
+ if authKind == "apikey" {
+ excluded = entry.ExcludedModels
+ }
+ }
+ models = applyExcludedModels(models, excluded)
+ case "qwen":
+ models = registry.GetQwenModels()
+ models = applyExcludedModels(models, excluded)
+ case "iflow":
+ models = registry.GetIFlowModels()
+ models = applyExcludedModels(models, excluded)
+ default:
+ // Handle OpenAI-compatibility providers by name using config
+ if s.cfg != nil {
+ providerKey := provider
+ compatName := strings.TrimSpace(a.Provider)
+ isCompatAuth := false
+ if compatDetected {
+ if compatProviderKey != "" {
+ providerKey = compatProviderKey
+ }
+ if compatDisplayName != "" {
+ compatName = compatDisplayName
+ }
+ isCompatAuth = true
+ }
+ if strings.EqualFold(providerKey, "openai-compatibility") {
+ isCompatAuth = true
+ if a.Attributes != nil {
+ if v := strings.TrimSpace(a.Attributes["compat_name"]); v != "" {
+ compatName = v
+ }
+ if v := strings.TrimSpace(a.Attributes["provider_key"]); v != "" {
+ providerKey = strings.ToLower(v)
+ isCompatAuth = true
+ }
+ }
+ if providerKey == "openai-compatibility" && compatName != "" {
+ providerKey = strings.ToLower(compatName)
+ }
+ } else if a.Attributes != nil {
+ if v := strings.TrimSpace(a.Attributes["compat_name"]); v != "" {
+ compatName = v
+ isCompatAuth = true
+ }
+ if v := strings.TrimSpace(a.Attributes["provider_key"]); v != "" {
+ providerKey = strings.ToLower(v)
+ isCompatAuth = true
+ }
+ }
+ for i := range s.cfg.OpenAICompatibility {
+ compat := &s.cfg.OpenAICompatibility[i]
+ if strings.EqualFold(compat.Name, compatName) {
+ isCompatAuth = true
+ // Convert compatibility models to registry models
+ ms := make([]*ModelInfo, 0, len(compat.Models))
+ for j := range compat.Models {
+ m := compat.Models[j]
+ // Use alias as model ID, fallback to name if alias is empty
+ modelID := m.Alias
+ if modelID == "" {
+ modelID = m.Name
+ }
+ ms = append(ms, &ModelInfo{
+ ID: modelID,
+ Object: "model",
+ Created: time.Now().Unix(),
+ OwnedBy: compat.Name,
+ Type: "openai-compatibility",
+ DisplayName: modelID,
+ UserDefined: true,
+ })
+ }
+ // Register and return
+ if len(ms) > 0 {
+ if providerKey == "" {
+ providerKey = "openai-compatibility"
+ }
+ GlobalModelRegistry().RegisterClient(a.ID, providerKey, applyModelPrefixes(ms, a.Prefix, s.cfg.ForceModelPrefix))
+ } else {
+ // Ensure stale registrations are cleared when model list becomes empty.
+ GlobalModelRegistry().UnregisterClient(a.ID)
+ }
+ return
+ }
+ }
+ if isCompatAuth {
+ // No matching provider found or models removed entirely; drop any prior registration.
+ GlobalModelRegistry().UnregisterClient(a.ID)
+ return
+ }
+ }
+ }
+ models = applyOAuthModelAlias(s.cfg, provider, authKind, models)
+ if len(models) > 0 {
+ key := provider
+ if key == "" {
+ key = strings.ToLower(strings.TrimSpace(a.Provider))
+ }
+ GlobalModelRegistry().RegisterClient(a.ID, key, applyModelPrefixes(models, a.Prefix, s.cfg != nil && s.cfg.ForceModelPrefix))
+ return
+ }
+
+ GlobalModelRegistry().UnregisterClient(a.ID)
+}
+
+func (s *Service) resolveConfigClaudeKey(auth *coreauth.Auth) *config.ClaudeKey {
+ if auth == nil || s.cfg == nil {
+ return nil
+ }
+ var attrKey, attrBase string
+ if auth.Attributes != nil {
+ attrKey = strings.TrimSpace(auth.Attributes["api_key"])
+ attrBase = strings.TrimSpace(auth.Attributes["base_url"])
+ }
+ for i := range s.cfg.ClaudeKey {
+ entry := &s.cfg.ClaudeKey[i]
+ cfgKey := strings.TrimSpace(entry.APIKey)
+ cfgBase := strings.TrimSpace(entry.BaseURL)
+ if attrKey != "" && attrBase != "" {
+ if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) {
+ return entry
+ }
+ continue
+ }
+ if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
+ if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
+ return entry
+ }
+ }
+ if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
+ return entry
+ }
+ }
+ if attrKey != "" {
+ for i := range s.cfg.ClaudeKey {
+ entry := &s.cfg.ClaudeKey[i]
+ if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) {
+ return entry
+ }
+ }
+ }
+ return nil
+}
+
+func (s *Service) resolveConfigGeminiKey(auth *coreauth.Auth) *config.GeminiKey {
+ if auth == nil || s.cfg == nil {
+ return nil
+ }
+ var attrKey, attrBase string
+ if auth.Attributes != nil {
+ attrKey = strings.TrimSpace(auth.Attributes["api_key"])
+ attrBase = strings.TrimSpace(auth.Attributes["base_url"])
+ }
+ for i := range s.cfg.GeminiKey {
+ entry := &s.cfg.GeminiKey[i]
+ cfgKey := strings.TrimSpace(entry.APIKey)
+ cfgBase := strings.TrimSpace(entry.BaseURL)
+ if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
+ if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
+ return entry
+ }
+ continue
+ }
+ if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
+ return entry
+ }
+ }
+ return nil
+}
+
+func (s *Service) resolveConfigVertexCompatKey(auth *coreauth.Auth) *config.VertexCompatKey {
+ if auth == nil || s.cfg == nil {
+ return nil
+ }
+ var attrKey, attrBase string
+ if auth.Attributes != nil {
+ attrKey = strings.TrimSpace(auth.Attributes["api_key"])
+ attrBase = strings.TrimSpace(auth.Attributes["base_url"])
+ }
+ for i := range s.cfg.VertexCompatAPIKey {
+ entry := &s.cfg.VertexCompatAPIKey[i]
+ cfgKey := strings.TrimSpace(entry.APIKey)
+ cfgBase := strings.TrimSpace(entry.BaseURL)
+ if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
+ if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
+ return entry
+ }
+ continue
+ }
+ if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
+ return entry
+ }
+ }
+ if attrKey != "" {
+ for i := range s.cfg.VertexCompatAPIKey {
+ entry := &s.cfg.VertexCompatAPIKey[i]
+ if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) {
+ return entry
+ }
+ }
+ }
+ return nil
+}
+
+func (s *Service) resolveConfigCodexKey(auth *coreauth.Auth) *config.CodexKey {
+ if auth == nil || s.cfg == nil {
+ return nil
+ }
+ var attrKey, attrBase string
+ if auth.Attributes != nil {
+ attrKey = strings.TrimSpace(auth.Attributes["api_key"])
+ attrBase = strings.TrimSpace(auth.Attributes["base_url"])
+ }
+ for i := range s.cfg.CodexKey {
+ entry := &s.cfg.CodexKey[i]
+ cfgKey := strings.TrimSpace(entry.APIKey)
+ cfgBase := strings.TrimSpace(entry.BaseURL)
+ if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
+ if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
+ return entry
+ }
+ continue
+ }
+ if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
+ return entry
+ }
+ }
+ return nil
+}
+
+func (s *Service) oauthExcludedModels(provider, authKind string) []string {
+ cfg := s.cfg
+ if cfg == nil {
+ return nil
+ }
+ authKindKey := strings.ToLower(strings.TrimSpace(authKind))
+ providerKey := strings.ToLower(strings.TrimSpace(provider))
+ if authKindKey == "apikey" {
+ return nil
+ }
+ return cfg.OAuthExcludedModels[providerKey]
+}
+
+func applyExcludedModels(models []*ModelInfo, excluded []string) []*ModelInfo {
+ if len(models) == 0 || len(excluded) == 0 {
+ return models
+ }
+
+ patterns := make([]string, 0, len(excluded))
+ for _, item := range excluded {
+ if trimmed := strings.TrimSpace(item); trimmed != "" {
+ patterns = append(patterns, strings.ToLower(trimmed))
+ }
+ }
+ if len(patterns) == 0 {
+ return models
+ }
+
+ filtered := make([]*ModelInfo, 0, len(models))
+ for _, model := range models {
+ if model == nil {
+ continue
+ }
+ modelID := strings.ToLower(strings.TrimSpace(model.ID))
+ blocked := false
+ for _, pattern := range patterns {
+ if matchWildcard(pattern, modelID) {
+ blocked = true
+ break
+ }
+ }
+ if !blocked {
+ filtered = append(filtered, model)
+ }
+ }
+ return filtered
+}
+
+func applyModelPrefixes(models []*ModelInfo, prefix string, forceModelPrefix bool) []*ModelInfo {
+ trimmedPrefix := strings.TrimSpace(prefix)
+ if trimmedPrefix == "" || len(models) == 0 {
+ return models
+ }
+
+ out := make([]*ModelInfo, 0, len(models)*2)
+ seen := make(map[string]struct{}, len(models)*2)
+
+ addModel := func(model *ModelInfo) {
+ if model == nil {
+ return
+ }
+ id := strings.TrimSpace(model.ID)
+ if id == "" {
+ return
+ }
+ if _, exists := seen[id]; exists {
+ return
+ }
+ seen[id] = struct{}{}
+ out = append(out, model)
+ }
+
+ for _, model := range models {
+ if model == nil {
+ continue
+ }
+ baseID := strings.TrimSpace(model.ID)
+ if baseID == "" {
+ continue
+ }
+ if !forceModelPrefix || trimmedPrefix == baseID {
+ addModel(model)
+ }
+ clone := *model
+ clone.ID = trimmedPrefix + "/" + baseID
+ addModel(&clone)
+ }
+ return out
+}
+
+// matchWildcard performs case-insensitive wildcard matching where '*' matches any substring.
+func matchWildcard(pattern, value string) bool {
+ if pattern == "" {
+ return false
+ }
+
+ // Fast path for exact match (no wildcard present).
+ if !strings.Contains(pattern, "*") {
+ return pattern == value
+ }
+
+ parts := strings.Split(pattern, "*")
+ // Handle prefix.
+ if prefix := parts[0]; prefix != "" {
+ if !strings.HasPrefix(value, prefix) {
+ return false
+ }
+ value = value[len(prefix):]
+ }
+
+ // Handle suffix.
+ if suffix := parts[len(parts)-1]; suffix != "" {
+ if !strings.HasSuffix(value, suffix) {
+ return false
+ }
+ value = value[:len(value)-len(suffix)]
+ }
+
+ // Handle middle segments in order.
+ for i := 1; i < len(parts)-1; i++ {
+ segment := parts[i]
+ if segment == "" {
+ continue
+ }
+ idx := strings.Index(value, segment)
+ if idx < 0 {
+ return false
+ }
+ value = value[idx+len(segment):]
+ }
+
+ return true
+}
+
+type modelEntry interface {
+ GetName() string
+ GetAlias() string
+}
+
+func buildConfigModels[T modelEntry](models []T, ownedBy, modelType string) []*ModelInfo {
+ if len(models) == 0 {
+ return nil
+ }
+ now := time.Now().Unix()
+ out := make([]*ModelInfo, 0, len(models))
+ seen := make(map[string]struct{}, len(models))
+ for i := range models {
+ model := models[i]
+ name := strings.TrimSpace(model.GetName())
+ alias := strings.TrimSpace(model.GetAlias())
+ if alias == "" {
+ alias = name
+ }
+ if alias == "" {
+ continue
+ }
+ key := strings.ToLower(alias)
+ if _, exists := seen[key]; exists {
+ continue
+ }
+ seen[key] = struct{}{}
+ display := name
+ if display == "" {
+ display = alias
+ }
+ info := &ModelInfo{
+ ID: alias,
+ Object: "model",
+ Created: now,
+ OwnedBy: ownedBy,
+ Type: modelType,
+ DisplayName: display,
+ UserDefined: true,
+ }
+ if name != "" {
+ if upstream := registry.LookupStaticModelInfo(name); upstream != nil && upstream.Thinking != nil {
+ info.Thinking = upstream.Thinking
+ }
+ }
+ out = append(out, info)
+ }
+ return out
+}
+
+func buildVertexCompatConfigModels(entry *config.VertexCompatKey) []*ModelInfo {
+ if entry == nil {
+ return nil
+ }
+ return buildConfigModels(entry.Models, "google", "vertex")
+}
+
+func buildGeminiConfigModels(entry *config.GeminiKey) []*ModelInfo {
+ if entry == nil {
+ return nil
+ }
+ return buildConfigModels(entry.Models, "google", "gemini")
+}
+
+func buildClaudeConfigModels(entry *config.ClaudeKey) []*ModelInfo {
+ if entry == nil {
+ return nil
+ }
+ return buildConfigModels(entry.Models, "anthropic", "claude")
+}
+
+func buildCodexConfigModels(entry *config.CodexKey) []*ModelInfo {
+ if entry == nil {
+ return nil
+ }
+ return buildConfigModels(entry.Models, "openai", "openai")
+}
+
+func rewriteModelInfoName(name, oldID, newID string) string {
+ trimmed := strings.TrimSpace(name)
+ if trimmed == "" {
+ return name
+ }
+ oldID = strings.TrimSpace(oldID)
+ newID = strings.TrimSpace(newID)
+ if oldID == "" || newID == "" {
+ return name
+ }
+ if strings.EqualFold(oldID, newID) {
+ return name
+ }
+ if strings.EqualFold(trimmed, oldID) {
+ return newID
+ }
+ if strings.HasSuffix(trimmed, "/"+oldID) {
+ prefix := strings.TrimSuffix(trimmed, oldID)
+ return prefix + newID
+ }
+ if trimmed == "models/"+oldID {
+ return "models/" + newID
+ }
+ return name
+}
+
+func applyOAuthModelAlias(cfg *config.Config, provider, authKind string, models []*ModelInfo) []*ModelInfo {
+ if cfg == nil || len(models) == 0 {
+ return models
+ }
+ channel := coreauth.OAuthModelAliasChannel(provider, authKind)
+ if channel == "" || len(cfg.OAuthModelAlias) == 0 {
+ return models
+ }
+ aliases := cfg.OAuthModelAlias[channel]
+ if len(aliases) == 0 {
+ return models
+ }
+
+ type aliasEntry struct {
+ alias string
+ fork bool
+ }
+
+ forward := make(map[string][]aliasEntry, len(aliases))
+ for i := range aliases {
+ name := strings.TrimSpace(aliases[i].Name)
+ alias := strings.TrimSpace(aliases[i].Alias)
+ if name == "" || alias == "" {
+ continue
+ }
+ if strings.EqualFold(name, alias) {
+ continue
+ }
+ key := strings.ToLower(name)
+ forward[key] = append(forward[key], aliasEntry{alias: alias, fork: aliases[i].Fork})
+ }
+ if len(forward) == 0 {
+ return models
+ }
+
+ out := make([]*ModelInfo, 0, len(models))
+ seen := make(map[string]struct{}, len(models))
+ for _, model := range models {
+ if model == nil {
+ continue
+ }
+ id := strings.TrimSpace(model.ID)
+ if id == "" {
+ continue
+ }
+ key := strings.ToLower(id)
+ entries := forward[key]
+ if len(entries) == 0 {
+ if _, exists := seen[key]; exists {
+ continue
+ }
+ seen[key] = struct{}{}
+ out = append(out, model)
+ continue
+ }
+
+ keepOriginal := false
+ for _, entry := range entries {
+ if entry.fork {
+ keepOriginal = true
+ break
+ }
+ }
+ if keepOriginal {
+ if _, exists := seen[key]; !exists {
+ seen[key] = struct{}{}
+ out = append(out, model)
+ }
+ }
+
+ addedAlias := false
+ for _, entry := range entries {
+ mappedID := strings.TrimSpace(entry.alias)
+ if mappedID == "" {
+ continue
+ }
+ if strings.EqualFold(mappedID, id) {
+ continue
+ }
+ aliasKey := strings.ToLower(mappedID)
+ if _, exists := seen[aliasKey]; exists {
+ continue
+ }
+ seen[aliasKey] = struct{}{}
+ clone := *model
+ clone.ID = mappedID
+ if clone.Name != "" {
+ clone.Name = rewriteModelInfoName(clone.Name, id, mappedID)
+ }
+ out = append(out, &clone)
+ addedAlias = true
+ }
+
+ if !keepOriginal && !addedAlias {
+ if _, exists := seen[key]; exists {
+ continue
+ }
+ seen[key] = struct{}{}
+ out = append(out, model)
+ }
+ }
+ return out
+}
diff --git a/sdk/cliproxy/service_oauth_model_alias_test.go b/sdk/cliproxy/service_oauth_model_alias_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..2caf7a178fbc660e111450f89b074b9e9f2e9864
--- /dev/null
+++ b/sdk/cliproxy/service_oauth_model_alias_test.go
@@ -0,0 +1,92 @@
+package cliproxy
+
+import (
+ "testing"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
+)
+
+func TestApplyOAuthModelAlias_Rename(t *testing.T) {
+ cfg := &config.Config{
+ OAuthModelAlias: map[string][]config.OAuthModelAlias{
+ "codex": {
+ {Name: "gpt-5", Alias: "g5"},
+ },
+ },
+ }
+ models := []*ModelInfo{
+ {ID: "gpt-5", Name: "models/gpt-5"},
+ }
+
+ out := applyOAuthModelAlias(cfg, "codex", "oauth", models)
+ if len(out) != 1 {
+ t.Fatalf("expected 1 model, got %d", len(out))
+ }
+ if out[0].ID != "g5" {
+ t.Fatalf("expected model id %q, got %q", "g5", out[0].ID)
+ }
+ if out[0].Name != "models/g5" {
+ t.Fatalf("expected model name %q, got %q", "models/g5", out[0].Name)
+ }
+}
+
+func TestApplyOAuthModelAlias_ForkAddsAlias(t *testing.T) {
+ cfg := &config.Config{
+ OAuthModelAlias: map[string][]config.OAuthModelAlias{
+ "codex": {
+ {Name: "gpt-5", Alias: "g5", Fork: true},
+ },
+ },
+ }
+ models := []*ModelInfo{
+ {ID: "gpt-5", Name: "models/gpt-5"},
+ }
+
+ out := applyOAuthModelAlias(cfg, "codex", "oauth", models)
+ if len(out) != 2 {
+ t.Fatalf("expected 2 models, got %d", len(out))
+ }
+ if out[0].ID != "gpt-5" {
+ t.Fatalf("expected first model id %q, got %q", "gpt-5", out[0].ID)
+ }
+ if out[1].ID != "g5" {
+ t.Fatalf("expected second model id %q, got %q", "g5", out[1].ID)
+ }
+ if out[1].Name != "models/g5" {
+ t.Fatalf("expected forked model name %q, got %q", "models/g5", out[1].Name)
+ }
+}
+
+func TestApplyOAuthModelAlias_ForkAddsMultipleAliases(t *testing.T) {
+ cfg := &config.Config{
+ OAuthModelAlias: map[string][]config.OAuthModelAlias{
+ "codex": {
+ {Name: "gpt-5", Alias: "g5", Fork: true},
+ {Name: "gpt-5", Alias: "g5-2", Fork: true},
+ },
+ },
+ }
+ models := []*ModelInfo{
+ {ID: "gpt-5", Name: "models/gpt-5"},
+ }
+
+ out := applyOAuthModelAlias(cfg, "codex", "oauth", models)
+ if len(out) != 3 {
+ t.Fatalf("expected 3 models, got %d", len(out))
+ }
+ if out[0].ID != "gpt-5" {
+ t.Fatalf("expected first model id %q, got %q", "gpt-5", out[0].ID)
+ }
+ if out[1].ID != "g5" {
+ t.Fatalf("expected second model id %q, got %q", "g5", out[1].ID)
+ }
+ if out[1].Name != "models/g5" {
+ t.Fatalf("expected forked model name %q, got %q", "models/g5", out[1].Name)
+ }
+ if out[2].ID != "g5-2" {
+ t.Fatalf("expected third model id %q, got %q", "g5-2", out[2].ID)
+ }
+ if out[2].Name != "models/g5-2" {
+ t.Fatalf("expected forked model name %q, got %q", "models/g5-2", out[2].Name)
+ }
+}
diff --git a/sdk/cliproxy/types.go b/sdk/cliproxy/types.go
new file mode 100644
index 0000000000000000000000000000000000000000..1521dffee442e4a8890ead23455ec602dccb8872
--- /dev/null
+++ b/sdk/cliproxy/types.go
@@ -0,0 +1,148 @@
+// Package cliproxy provides the core service implementation for the CLI Proxy API.
+// It includes service lifecycle management, authentication handling, file watching,
+// and integration with various AI service providers through a unified interface.
+package cliproxy
+
+import (
+ "context"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/watcher"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
+)
+
+// TokenClientProvider loads clients backed by stored authentication tokens.
+// It provides an interface for loading authentication tokens from various sources
+// and creating clients for AI service providers.
+type TokenClientProvider interface {
+ // Load loads token-based clients from the configured source.
+ //
+ // Parameters:
+ // - ctx: The context for the loading operation
+ // - cfg: The application configuration
+ //
+ // Returns:
+ // - *TokenClientResult: The result containing loaded clients
+ // - error: An error if loading fails
+ Load(ctx context.Context, cfg *config.Config) (*TokenClientResult, error)
+}
+
+// TokenClientResult represents clients generated from persisted tokens.
+// It contains metadata about the loading operation and the number of successful authentications.
+type TokenClientResult struct {
+ // SuccessfulAuthed is the number of successfully authenticated clients.
+ SuccessfulAuthed int
+}
+
+// APIKeyClientProvider loads clients backed directly by configured API keys.
+// It provides an interface for loading API key-based clients for various AI service providers.
+type APIKeyClientProvider interface {
+ // Load loads API key-based clients from the configuration.
+ //
+ // Parameters:
+ // - ctx: The context for the loading operation
+ // - cfg: The application configuration
+ //
+ // Returns:
+ // - *APIKeyClientResult: The result containing loaded clients
+ // - error: An error if loading fails
+ Load(ctx context.Context, cfg *config.Config) (*APIKeyClientResult, error)
+}
+
+// APIKeyClientResult is returned by APIKeyClientProvider.Load()
+type APIKeyClientResult struct {
+ // GeminiKeyCount is the number of Gemini API keys loaded
+ GeminiKeyCount int
+
+ // VertexCompatKeyCount is the number of Vertex-compatible API keys loaded
+ VertexCompatKeyCount int
+
+ // ClaudeKeyCount is the number of Claude API keys loaded
+ ClaudeKeyCount int
+
+ // CodexKeyCount is the number of Codex API keys loaded
+ CodexKeyCount int
+
+ // OpenAICompatCount is the number of OpenAI compatibility API keys loaded
+ OpenAICompatCount int
+}
+
+// WatcherFactory creates a watcher for configuration and token changes.
+// The reload callback receives the updated configuration when changes are detected.
+//
+// Parameters:
+// - configPath: The path to the configuration file to watch
+// - authDir: The directory containing authentication tokens to watch
+// - reload: The callback function to call when changes are detected
+//
+// Returns:
+// - *WatcherWrapper: A watcher wrapper instance
+// - error: An error if watcher creation fails
+type WatcherFactory func(configPath, authDir string, reload func(*config.Config)) (*WatcherWrapper, error)
+
+// WatcherWrapper exposes the subset of watcher methods required by the SDK.
+type WatcherWrapper struct {
+ start func(ctx context.Context) error
+ stop func() error
+
+ setConfig func(cfg *config.Config)
+ snapshotAuths func() []*coreauth.Auth
+ setUpdateQueue func(queue chan<- watcher.AuthUpdate)
+ dispatchRuntimeUpdate func(update watcher.AuthUpdate) bool
+}
+
+// Start proxies to the underlying watcher Start implementation.
+func (w *WatcherWrapper) Start(ctx context.Context) error {
+ if w == nil || w.start == nil {
+ return nil
+ }
+ return w.start(ctx)
+}
+
+// Stop proxies to the underlying watcher Stop implementation.
+func (w *WatcherWrapper) Stop() error {
+ if w == nil || w.stop == nil {
+ return nil
+ }
+ return w.stop()
+}
+
+// SetConfig updates the watcher configuration cache.
+func (w *WatcherWrapper) SetConfig(cfg *config.Config) {
+ if w == nil || w.setConfig == nil {
+ return
+ }
+ w.setConfig(cfg)
+}
+
+// DispatchRuntimeAuthUpdate forwards runtime auth updates (e.g., websocket providers)
+// into the watcher-managed auth update queue when available.
+// Returns true if the update was enqueued successfully.
+func (w *WatcherWrapper) DispatchRuntimeAuthUpdate(update watcher.AuthUpdate) bool {
+ if w == nil || w.dispatchRuntimeUpdate == nil {
+ return false
+ }
+ return w.dispatchRuntimeUpdate(update)
+}
+
+// SetClients updates the watcher file-backed clients registry.
+// SetClients and SetAPIKeyClients removed; watcher manages its own caches
+
+// SnapshotClients returns the current combined clients snapshot from the underlying watcher.
+// SnapshotClients removed; use SnapshotAuths
+
+// SnapshotAuths returns the current auth entries derived from legacy clients.
+func (w *WatcherWrapper) SnapshotAuths() []*coreauth.Auth {
+ if w == nil || w.snapshotAuths == nil {
+ return nil
+ }
+ return w.snapshotAuths()
+}
+
+// SetAuthUpdateQueue registers the channel used to propagate auth updates.
+func (w *WatcherWrapper) SetAuthUpdateQueue(queue chan<- watcher.AuthUpdate) {
+ if w == nil || w.setUpdateQueue == nil {
+ return
+ }
+ w.setUpdateQueue(queue)
+}
diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go
new file mode 100644
index 0000000000000000000000000000000000000000..58b036076142d05f19e6ce1ef046e51a8245d153
--- /dev/null
+++ b/sdk/cliproxy/usage/manager.go
@@ -0,0 +1,181 @@
+package usage
+
+import (
+ "context"
+ "sync"
+ "time"
+
+ log "github.com/sirupsen/logrus"
+)
+
+// Record contains the usage statistics captured for a single provider request.
+type Record struct {
+ Provider string
+ Model string
+ APIKey string
+ AuthID string
+ AuthIndex string
+ Source string
+ RequestedAt time.Time
+ Failed bool
+ Detail Detail
+}
+
+// Detail holds the token usage breakdown.
+type Detail struct {
+ InputTokens int64
+ OutputTokens int64
+ ReasoningTokens int64
+ CachedTokens int64
+ TotalTokens int64
+}
+
+// Plugin consumes usage records emitted by the proxy runtime.
+type Plugin interface {
+ HandleUsage(ctx context.Context, record Record)
+}
+
+type queueItem struct {
+ ctx context.Context
+ record Record
+}
+
+// Manager maintains a queue of usage records and delivers them to registered plugins.
+type Manager struct {
+ once sync.Once
+ stopOnce sync.Once
+ cancel context.CancelFunc
+
+ mu sync.Mutex
+ cond *sync.Cond
+ queue []queueItem
+ closed bool
+
+ pluginsMu sync.RWMutex
+ plugins []Plugin
+}
+
+// NewManager constructs a manager with a buffered queue.
+func NewManager(buffer int) *Manager {
+ m := &Manager{}
+ m.cond = sync.NewCond(&m.mu)
+ return m
+}
+
+// Start launches the background dispatcher. Calling Start multiple times is safe.
+func (m *Manager) Start(ctx context.Context) {
+ if m == nil {
+ return
+ }
+ m.once.Do(func() {
+ if ctx == nil {
+ ctx = context.Background()
+ }
+ var workerCtx context.Context
+ workerCtx, m.cancel = context.WithCancel(ctx)
+ go m.run(workerCtx)
+ })
+}
+
+// Stop stops the dispatcher and drains the queue.
+func (m *Manager) Stop() {
+ if m == nil {
+ return
+ }
+ m.stopOnce.Do(func() {
+ if m.cancel != nil {
+ m.cancel()
+ }
+ m.mu.Lock()
+ m.closed = true
+ m.mu.Unlock()
+ m.cond.Broadcast()
+ })
+}
+
+// Register appends a plugin to the delivery list.
+func (m *Manager) Register(plugin Plugin) {
+ if m == nil || plugin == nil {
+ return
+ }
+ m.pluginsMu.Lock()
+ m.plugins = append(m.plugins, plugin)
+ m.pluginsMu.Unlock()
+}
+
+// Publish enqueues a usage record for processing. If no plugin is registered
+// the record will be discarded downstream.
+func (m *Manager) Publish(ctx context.Context, record Record) {
+ if m == nil {
+ return
+ }
+ // ensure worker is running even if Start was not called explicitly
+ m.Start(context.Background())
+ m.mu.Lock()
+ if m.closed {
+ m.mu.Unlock()
+ return
+ }
+ m.queue = append(m.queue, queueItem{ctx: ctx, record: record})
+ m.mu.Unlock()
+ m.cond.Signal()
+}
+
+func (m *Manager) run(ctx context.Context) {
+ for {
+ m.mu.Lock()
+ for !m.closed && len(m.queue) == 0 {
+ m.cond.Wait()
+ }
+ if len(m.queue) == 0 && m.closed {
+ m.mu.Unlock()
+ return
+ }
+ item := m.queue[0]
+ m.queue = m.queue[1:]
+ m.mu.Unlock()
+ m.dispatch(item)
+ }
+}
+
+func (m *Manager) dispatch(item queueItem) {
+ m.pluginsMu.RLock()
+ plugins := make([]Plugin, len(m.plugins))
+ copy(plugins, m.plugins)
+ m.pluginsMu.RUnlock()
+ if len(plugins) == 0 {
+ return
+ }
+ for _, plugin := range plugins {
+ if plugin == nil {
+ continue
+ }
+ safeInvoke(plugin, item.ctx, item.record)
+ }
+}
+
+func safeInvoke(plugin Plugin, ctx context.Context, record Record) {
+ defer func() {
+ if r := recover(); r != nil {
+ log.Errorf("usage: plugin panic recovered: %v", r)
+ }
+ }()
+ plugin.HandleUsage(ctx, record)
+}
+
+var defaultManager = NewManager(512)
+
+// DefaultManager returns the global usage manager instance.
+func DefaultManager() *Manager { return defaultManager }
+
+// RegisterPlugin registers a plugin on the default manager.
+func RegisterPlugin(plugin Plugin) { DefaultManager().Register(plugin) }
+
+// PublishRecord publishes a record using the default manager.
+func PublishRecord(ctx context.Context, record Record) { DefaultManager().Publish(ctx, record) }
+
+// StartDefault starts the default manager's dispatcher.
+func StartDefault(ctx context.Context) { DefaultManager().Start(ctx) }
+
+// StopDefault stops the default manager's dispatcher.
+func StopDefault() { DefaultManager().Stop() }
diff --git a/sdk/cliproxy/watcher.go b/sdk/cliproxy/watcher.go
new file mode 100644
index 0000000000000000000000000000000000000000..caeadf19b910daa64250751fa5f9589b9135fab2
--- /dev/null
+++ b/sdk/cliproxy/watcher.go
@@ -0,0 +1,35 @@
+package cliproxy
+
+import (
+ "context"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/watcher"
+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
+)
+
+func defaultWatcherFactory(configPath, authDir string, reload func(*config.Config)) (*WatcherWrapper, error) {
+ w, err := watcher.NewWatcher(configPath, authDir, reload)
+ if err != nil {
+ return nil, err
+ }
+
+ return &WatcherWrapper{
+ start: func(ctx context.Context) error {
+ return w.Start(ctx)
+ },
+ stop: func() error {
+ return w.Stop()
+ },
+ setConfig: func(cfg *config.Config) {
+ w.SetConfig(cfg)
+ },
+ snapshotAuths: func() []*coreauth.Auth { return w.SnapshotCoreAuths() },
+ setUpdateQueue: func(queue chan<- watcher.AuthUpdate) {
+ w.SetAuthUpdateQueue(queue)
+ },
+ dispatchRuntimeUpdate: func(update watcher.AuthUpdate) bool {
+ return w.DispatchRuntimeAuthUpdate(update)
+ },
+ }, nil
+}
diff --git a/sdk/config/config.go b/sdk/config/config.go
new file mode 100644
index 0000000000000000000000000000000000000000..304ccdd8c34d02d6115bde1fe6e658c9e198b36e
--- /dev/null
+++ b/sdk/config/config.go
@@ -0,0 +1,61 @@
+// Package config provides the public SDK configuration API.
+//
+// It re-exports the server configuration types and helpers so external projects can
+// embed CLIProxyAPI without importing internal packages.
+package config
+
+import internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+
+type SDKConfig = internalconfig.SDKConfig
+type AccessConfig = internalconfig.AccessConfig
+type AccessProvider = internalconfig.AccessProvider
+
+type Config = internalconfig.Config
+
+type StreamingConfig = internalconfig.StreamingConfig
+type TLSConfig = internalconfig.TLSConfig
+type RemoteManagement = internalconfig.RemoteManagement
+type AmpCode = internalconfig.AmpCode
+type OAuthModelAlias = internalconfig.OAuthModelAlias
+type PayloadConfig = internalconfig.PayloadConfig
+type PayloadRule = internalconfig.PayloadRule
+type PayloadModelRule = internalconfig.PayloadModelRule
+
+type GeminiKey = internalconfig.GeminiKey
+type CodexKey = internalconfig.CodexKey
+type ClaudeKey = internalconfig.ClaudeKey
+type VertexCompatKey = internalconfig.VertexCompatKey
+type VertexCompatModel = internalconfig.VertexCompatModel
+type OpenAICompatibility = internalconfig.OpenAICompatibility
+type OpenAICompatibilityAPIKey = internalconfig.OpenAICompatibilityAPIKey
+type OpenAICompatibilityModel = internalconfig.OpenAICompatibilityModel
+
+type TLS = internalconfig.TLSConfig
+
+const (
+ AccessProviderTypeConfigAPIKey = internalconfig.AccessProviderTypeConfigAPIKey
+ DefaultAccessProviderName = internalconfig.DefaultAccessProviderName
+ DefaultPanelGitHubRepository = internalconfig.DefaultPanelGitHubRepository
+)
+
+func MakeInlineAPIKeyProvider(keys []string) *AccessProvider {
+ return internalconfig.MakeInlineAPIKeyProvider(keys)
+}
+
+func LoadConfig(configFile string) (*Config, error) { return internalconfig.LoadConfig(configFile) }
+
+func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
+ return internalconfig.LoadConfigOptional(configFile, optional)
+}
+
+func SaveConfigPreserveComments(configFile string, cfg *Config) error {
+ return internalconfig.SaveConfigPreserveComments(configFile, cfg)
+}
+
+func SaveConfigPreserveCommentsUpdateNestedScalar(configFile string, path []string, value string) error {
+ return internalconfig.SaveConfigPreserveCommentsUpdateNestedScalar(configFile, path, value)
+}
+
+func NormalizeCommentIndentation(data []byte) []byte {
+ return internalconfig.NormalizeCommentIndentation(data)
+}
diff --git a/sdk/logging/request_logger.go b/sdk/logging/request_logger.go
new file mode 100644
index 0000000000000000000000000000000000000000..39ff5ba8361f894d3cb7fc7cf0874e90e7cc05c9
--- /dev/null
+++ b/sdk/logging/request_logger.go
@@ -0,0 +1,18 @@
+// Package logging re-exports request logging primitives for SDK consumers.
+package logging
+
+import internallogging "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
+
+// RequestLogger defines the interface for logging HTTP requests and responses.
+type RequestLogger = internallogging.RequestLogger
+
+// StreamingLogWriter handles real-time logging of streaming response chunks.
+type StreamingLogWriter = internallogging.StreamingLogWriter
+
+// FileRequestLogger implements RequestLogger using file-based storage.
+type FileRequestLogger = internallogging.FileRequestLogger
+
+// NewFileRequestLogger creates a new file-based request logger.
+func NewFileRequestLogger(enabled bool, logsDir string, configDir string) *FileRequestLogger {
+ return internallogging.NewFileRequestLogger(enabled, logsDir, configDir)
+}
diff --git a/sdk/translator/builtin/builtin.go b/sdk/translator/builtin/builtin.go
new file mode 100644
index 0000000000000000000000000000000000000000..798e43f1a97160168e862fed3dc9f41a10156d80
--- /dev/null
+++ b/sdk/translator/builtin/builtin.go
@@ -0,0 +1,18 @@
+// Package builtin exposes the built-in translator registrations for SDK users.
+package builtin
+
+import (
+ sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
+
+ _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator"
+)
+
+// Registry exposes the default registry populated with all built-in translators.
+func Registry() *sdktranslator.Registry {
+ return sdktranslator.Default()
+}
+
+// Pipeline returns a pipeline that already contains the built-in translators.
+func Pipeline() *sdktranslator.Pipeline {
+ return sdktranslator.NewPipeline(sdktranslator.Default())
+}
diff --git a/sdk/translator/format.go b/sdk/translator/format.go
new file mode 100644
index 0000000000000000000000000000000000000000..ec0f37f65d3fbef46d7482a9ac45a83912fa6c96
--- /dev/null
+++ b/sdk/translator/format.go
@@ -0,0 +1,14 @@
+package translator
+
+// Format identifies a request/response schema used inside the proxy.
+type Format string
+
+// FromString converts an arbitrary identifier to a translator format.
+func FromString(v string) Format {
+ return Format(v)
+}
+
+// String returns the raw schema identifier.
+func (f Format) String() string {
+ return string(f)
+}
diff --git a/sdk/translator/formats.go b/sdk/translator/formats.go
new file mode 100644
index 0000000000000000000000000000000000000000..aafe9e056cc0619ccbad59decfebc90de2dc0757
--- /dev/null
+++ b/sdk/translator/formats.go
@@ -0,0 +1,12 @@
+package translator
+
+// Common format identifiers exposed for SDK users.
+const (
+ FormatOpenAI Format = "openai"
+ FormatOpenAIResponse Format = "openai-response"
+ FormatClaude Format = "claude"
+ FormatGemini Format = "gemini"
+ FormatGeminiCLI Format = "gemini-cli"
+ FormatCodex Format = "codex"
+ FormatAntigravity Format = "antigravity"
+)
diff --git a/sdk/translator/helpers.go b/sdk/translator/helpers.go
new file mode 100644
index 0000000000000000000000000000000000000000..bf8cfbf79d75e2be001dbe3656a21fbb366c15e3
--- /dev/null
+++ b/sdk/translator/helpers.go
@@ -0,0 +1,28 @@
+package translator
+
+import "context"
+
+// TranslateRequestByFormatName converts a request payload between schemas by their string identifiers.
+func TranslateRequestByFormatName(from, to Format, model string, rawJSON []byte, stream bool) []byte {
+ return TranslateRequest(from, to, model, rawJSON, stream)
+}
+
+// HasResponseTransformerByFormatName reports whether a response translator exists between two schemas.
+func HasResponseTransformerByFormatName(from, to Format) bool {
+ return HasResponseTransformer(from, to)
+}
+
+// TranslateStreamByFormatName converts streaming responses between schemas by their string identifiers.
+func TranslateStreamByFormatName(ctx context.Context, from, to Format, model string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []string {
+ return TranslateStream(ctx, from, to, model, originalRequestRawJSON, requestRawJSON, rawJSON, param)
+}
+
+// TranslateNonStreamByFormatName converts non-streaming responses between schemas by their string identifiers.
+func TranslateNonStreamByFormatName(ctx context.Context, from, to Format, model string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) string {
+ return TranslateNonStream(ctx, from, to, model, originalRequestRawJSON, requestRawJSON, rawJSON, param)
+}
+
+// TranslateTokenCountByFormatName converts token counts between schemas by their string identifiers.
+func TranslateTokenCountByFormatName(ctx context.Context, from, to Format, count int64, rawJSON []byte) string {
+ return TranslateTokenCount(ctx, from, to, count, rawJSON)
+}
diff --git a/sdk/translator/pipeline.go b/sdk/translator/pipeline.go
new file mode 100644
index 0000000000000000000000000000000000000000..5fa6c66a0abc019145acc7211d15e8589de91406
--- /dev/null
+++ b/sdk/translator/pipeline.go
@@ -0,0 +1,106 @@
+package translator
+
+import "context"
+
+// RequestEnvelope represents a request in the translation pipeline.
+type RequestEnvelope struct {
+ Format Format
+ Model string
+ Stream bool
+ Body []byte
+}
+
+// ResponseEnvelope represents a response in the translation pipeline.
+type ResponseEnvelope struct {
+ Format Format
+ Model string
+ Stream bool
+ Body []byte
+ Chunks []string
+}
+
+// RequestMiddleware decorates request translation.
+type RequestMiddleware func(ctx context.Context, req RequestEnvelope, next RequestHandler) (RequestEnvelope, error)
+
+// ResponseMiddleware decorates response translation.
+type ResponseMiddleware func(ctx context.Context, resp ResponseEnvelope, next ResponseHandler) (ResponseEnvelope, error)
+
+// RequestHandler performs request translation between formats.
+type RequestHandler func(ctx context.Context, req RequestEnvelope) (RequestEnvelope, error)
+
+// ResponseHandler performs response translation between formats.
+type ResponseHandler func(ctx context.Context, resp ResponseEnvelope) (ResponseEnvelope, error)
+
+// Pipeline orchestrates request/response transformation with middleware support.
+type Pipeline struct {
+ registry *Registry
+ requestMiddleware []RequestMiddleware
+ responseMiddleware []ResponseMiddleware
+}
+
+// NewPipeline constructs a pipeline bound to the provided registry.
+func NewPipeline(registry *Registry) *Pipeline {
+ if registry == nil {
+ registry = Default()
+ }
+ return &Pipeline{registry: registry}
+}
+
+// UseRequest adds request middleware executed in registration order.
+func (p *Pipeline) UseRequest(mw RequestMiddleware) {
+ if mw != nil {
+ p.requestMiddleware = append(p.requestMiddleware, mw)
+ }
+}
+
+// UseResponse adds response middleware executed in registration order.
+func (p *Pipeline) UseResponse(mw ResponseMiddleware) {
+ if mw != nil {
+ p.responseMiddleware = append(p.responseMiddleware, mw)
+ }
+}
+
+// TranslateRequest applies middleware and registry transformations.
+func (p *Pipeline) TranslateRequest(ctx context.Context, from, to Format, req RequestEnvelope) (RequestEnvelope, error) {
+ terminal := func(ctx context.Context, input RequestEnvelope) (RequestEnvelope, error) {
+ translated := p.registry.TranslateRequest(from, to, input.Model, input.Body, input.Stream)
+ input.Body = translated
+ input.Format = to
+ return input, nil
+ }
+
+ handler := terminal
+ for i := len(p.requestMiddleware) - 1; i >= 0; i-- {
+ mw := p.requestMiddleware[i]
+ next := handler
+ handler = func(ctx context.Context, r RequestEnvelope) (RequestEnvelope, error) {
+ return mw(ctx, r, next)
+ }
+ }
+
+ return handler(ctx, req)
+}
+
+// TranslateResponse applies middleware and registry transformations.
+func (p *Pipeline) TranslateResponse(ctx context.Context, from, to Format, resp ResponseEnvelope, originalReq, translatedReq []byte, param *any) (ResponseEnvelope, error) {
+ terminal := func(ctx context.Context, input ResponseEnvelope) (ResponseEnvelope, error) {
+ if input.Stream {
+ input.Chunks = p.registry.TranslateStream(ctx, from, to, input.Model, originalReq, translatedReq, input.Body, param)
+ } else {
+ input.Body = []byte(p.registry.TranslateNonStream(ctx, from, to, input.Model, originalReq, translatedReq, input.Body, param))
+ }
+ input.Format = to
+ return input, nil
+ }
+
+ handler := terminal
+ for i := len(p.responseMiddleware) - 1; i >= 0; i-- {
+ mw := p.responseMiddleware[i]
+ next := handler
+ handler = func(ctx context.Context, r ResponseEnvelope) (ResponseEnvelope, error) {
+ return mw(ctx, r, next)
+ }
+ }
+
+ return handler(ctx, resp)
+}
diff --git a/sdk/translator/registry.go b/sdk/translator/registry.go
new file mode 100644
index 0000000000000000000000000000000000000000..ace9713711b6989d229d95fd2a5b3d1c9a81c71a
--- /dev/null
+++ b/sdk/translator/registry.go
@@ -0,0 +1,142 @@
+package translator
+
+import (
+ "context"
+ "sync"
+)
+
+// Registry manages translation functions across schemas.
+type Registry struct {
+ mu sync.RWMutex
+ requests map[Format]map[Format]RequestTransform
+ responses map[Format]map[Format]ResponseTransform
+}
+
+// NewRegistry constructs an empty translator registry.
+func NewRegistry() *Registry {
+ return &Registry{
+ requests: make(map[Format]map[Format]RequestTransform),
+ responses: make(map[Format]map[Format]ResponseTransform),
+ }
+}
+
+// Register stores request/response transforms between two formats.
+func (r *Registry) Register(from, to Format, request RequestTransform, response ResponseTransform) {
+ r.mu.Lock()
+ defer r.mu.Unlock()
+
+ if _, ok := r.requests[from]; !ok {
+ r.requests[from] = make(map[Format]RequestTransform)
+ }
+ if request != nil {
+ r.requests[from][to] = request
+ }
+
+ if _, ok := r.responses[from]; !ok {
+ r.responses[from] = make(map[Format]ResponseTransform)
+ }
+ r.responses[from][to] = response
+}
+
+// TranslateRequest converts a payload between schemas, returning the original payload
+// if no translator is registered.
+func (r *Registry) TranslateRequest(from, to Format, model string, rawJSON []byte, stream bool) []byte {
+ r.mu.RLock()
+ defer r.mu.RUnlock()
+
+ if byTarget, ok := r.requests[from]; ok {
+ if fn, isOk := byTarget[to]; isOk && fn != nil {
+ return fn(model, rawJSON, stream)
+ }
+ }
+ return rawJSON
+}
+
+// HasResponseTransformer indicates whether a response translator exists.
+func (r *Registry) HasResponseTransformer(from, to Format) bool {
+ r.mu.RLock()
+ defer r.mu.RUnlock()
+
+ if byTarget, ok := r.responses[from]; ok {
+ if _, isOk := byTarget[to]; isOk {
+ return true
+ }
+ }
+ return false
+}
+
+// TranslateStream applies the registered streaming response translator.
+func (r *Registry) TranslateStream(ctx context.Context, from, to Format, model string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []string {
+ r.mu.RLock()
+ defer r.mu.RUnlock()
+
+ if byTarget, ok := r.responses[to]; ok {
+ if fn, isOk := byTarget[from]; isOk && fn.Stream != nil {
+ return fn.Stream(ctx, model, originalRequestRawJSON, requestRawJSON, rawJSON, param)
+ }
+ }
+ return []string{string(rawJSON)}
+}
+
+// TranslateNonStream applies the registered non-stream response translator.
+func (r *Registry) TranslateNonStream(ctx context.Context, from, to Format, model string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) string {
+ r.mu.RLock()
+ defer r.mu.RUnlock()
+
+ if byTarget, ok := r.responses[to]; ok {
+ if fn, isOk := byTarget[from]; isOk && fn.NonStream != nil {
+ return fn.NonStream(ctx, model, originalRequestRawJSON, requestRawJSON, rawJSON, param)
+ }
+ }
+ return string(rawJSON)
+}
+
+// TranslateNonStream applies the registered non-stream response translator.
+func (r *Registry) TranslateTokenCount(ctx context.Context, from, to Format, count int64, rawJSON []byte) string {
+ r.mu.RLock()
+ defer r.mu.RUnlock()
+
+ if byTarget, ok := r.responses[to]; ok {
+ if fn, isOk := byTarget[from]; isOk && fn.TokenCount != nil {
+ return fn.TokenCount(ctx, count)
+ }
+ }
+ return string(rawJSON)
+}
+
+var defaultRegistry = NewRegistry()
+
+// Default exposes the package-level registry for shared use.
+func Default() *Registry {
+ return defaultRegistry
+}
+
+// Register attaches transforms to the default registry.
+func Register(from, to Format, request RequestTransform, response ResponseTransform) {
+ defaultRegistry.Register(from, to, request, response)
+}
+
+// TranslateRequest is a helper on the default registry.
+func TranslateRequest(from, to Format, model string, rawJSON []byte, stream bool) []byte {
+ return defaultRegistry.TranslateRequest(from, to, model, rawJSON, stream)
+}
+
+// HasResponseTransformer inspects the default registry.
+func HasResponseTransformer(from, to Format) bool {
+ return defaultRegistry.HasResponseTransformer(from, to)
+}
+
+// TranslateStream is a helper on the default registry.
+func TranslateStream(ctx context.Context, from, to Format, model string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []string {
+ return defaultRegistry.TranslateStream(ctx, from, to, model, originalRequestRawJSON, requestRawJSON, rawJSON, param)
+}
+
+// TranslateNonStream is a helper on the default registry.
+func TranslateNonStream(ctx context.Context, from, to Format, model string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) string {
+ return defaultRegistry.TranslateNonStream(ctx, from, to, model, originalRequestRawJSON, requestRawJSON, rawJSON, param)
+}
+
+// TranslateTokenCount is a helper on the default registry.
+func TranslateTokenCount(ctx context.Context, from, to Format, count int64, rawJSON []byte) string {
+ return defaultRegistry.TranslateTokenCount(ctx, from, to, count, rawJSON)
+}
diff --git a/sdk/translator/types.go b/sdk/translator/types.go
new file mode 100644
index 0000000000000000000000000000000000000000..ff69340a5737b1eb7d06dc5d5b4291dab6c0ab62
--- /dev/null
+++ b/sdk/translator/types.go
@@ -0,0 +1,34 @@
+// Package translator provides types and functions for converting chat requests and responses between different schemas.
+package translator
+
+import "context"
+
+// RequestTransform is a function type that converts a request payload from a source schema to a target schema.
+// It takes the model name, the raw JSON payload of the request, and a boolean indicating if the request is for a streaming response.
+// It returns the converted request payload as a byte slice.
+type RequestTransform func(model string, rawJSON []byte, stream bool) []byte
+
+// ResponseStreamTransform is a function type that converts a streaming response from a source schema to a target schema.
+// It takes a context, the model name, the raw JSON of the original and converted requests, the raw JSON of the current response chunk, and an optional parameter.
+// It returns a slice of strings, where each string is a chunk of the converted streaming response.
+type ResponseStreamTransform func(ctx context.Context, model string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []string
+
+// ResponseNonStreamTransform is a function type that converts a non-streaming response from a source schema to a target schema.
+// It takes a context, the model name, the raw JSON of the original and converted requests, the raw JSON of the response, and an optional parameter.
+// It returns the converted response as a single string.
+type ResponseNonStreamTransform func(ctx context.Context, model string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) string
+
+// ResponseTokenCountTransform is a function type that transforms a token count from a source format to a target format.
+// It takes a context and the token count as an int64, and returns the transformed token count as a string.
+type ResponseTokenCountTransform func(ctx context.Context, count int64) string
+
+// ResponseTransform is a struct that groups together the functions for transforming streaming and non-streaming responses,
+// as well as token counts.
+type ResponseTransform struct {
+ // Stream is the function for transforming streaming responses.
+ Stream ResponseStreamTransform
+ // NonStream is the function for transforming non-streaming responses.
+ NonStream ResponseNonStreamTransform
+ // TokenCount is the function for transforming token counts.
+ TokenCount ResponseTokenCountTransform
+}
diff --git a/test/amp_management_test.go b/test/amp_management_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..e384ef0e8bf909bdb10b33ae3c8b417e1fce8eb3
--- /dev/null
+++ b/test/amp_management_test.go
@@ -0,0 +1,915 @@
+package test
+
+import (
+ "bytes"
+ "encoding/json"
+ "net/http"
+ "net/http/httptest"
+ "os"
+ "path/filepath"
+ "testing"
+
+ "github.com/gin-gonic/gin"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/api/handlers/management"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+)
+
+func init() {
+ gin.SetMode(gin.TestMode)
+}
+
+// newAmpTestHandler creates a test handler with default ampcode configuration.
+func newAmpTestHandler(t *testing.T) (*management.Handler, string) {
+ t.Helper()
+ tmpDir := t.TempDir()
+ configPath := filepath.Join(tmpDir, "config.yaml")
+
+ cfg := &config.Config{
+ AmpCode: config.AmpCode{
+ UpstreamURL: "https://example.com",
+ UpstreamAPIKey: "test-api-key-12345",
+ RestrictManagementToLocalhost: true,
+ ForceModelMappings: false,
+ ModelMappings: []config.AmpModelMapping{
+ {From: "gpt-4", To: "gemini-pro"},
+ },
+ },
+ }
+
+ if err := os.WriteFile(configPath, []byte("port: 8080\n"), 0644); err != nil {
+ t.Fatalf("failed to write config file: %v", err)
+ }
+
+ h := management.NewHandler(cfg, configPath, nil)
+ return h, configPath
+}
+
+// setupAmpRouter creates a test router with all ampcode management endpoints.
+func setupAmpRouter(h *management.Handler) *gin.Engine {
+ r := gin.New()
+ mgmt := r.Group("/v0/management")
+ {
+ mgmt.GET("/ampcode", h.GetAmpCode)
+ mgmt.GET("/ampcode/upstream-url", h.GetAmpUpstreamURL)
+ mgmt.PUT("/ampcode/upstream-url", h.PutAmpUpstreamURL)
+ mgmt.DELETE("/ampcode/upstream-url", h.DeleteAmpUpstreamURL)
+ mgmt.GET("/ampcode/upstream-api-key", h.GetAmpUpstreamAPIKey)
+ mgmt.PUT("/ampcode/upstream-api-key", h.PutAmpUpstreamAPIKey)
+ mgmt.DELETE("/ampcode/upstream-api-key", h.DeleteAmpUpstreamAPIKey)
+ mgmt.GET("/ampcode/upstream-api-keys", h.GetAmpUpstreamAPIKeys)
+ mgmt.PUT("/ampcode/upstream-api-keys", h.PutAmpUpstreamAPIKeys)
+ mgmt.PATCH("/ampcode/upstream-api-keys", h.PatchAmpUpstreamAPIKeys)
+ mgmt.DELETE("/ampcode/upstream-api-keys", h.DeleteAmpUpstreamAPIKeys)
+ mgmt.GET("/ampcode/restrict-management-to-localhost", h.GetAmpRestrictManagementToLocalhost)
+ mgmt.PUT("/ampcode/restrict-management-to-localhost", h.PutAmpRestrictManagementToLocalhost)
+ mgmt.GET("/ampcode/model-mappings", h.GetAmpModelMappings)
+ mgmt.PUT("/ampcode/model-mappings", h.PutAmpModelMappings)
+ mgmt.PATCH("/ampcode/model-mappings", h.PatchAmpModelMappings)
+ mgmt.DELETE("/ampcode/model-mappings", h.DeleteAmpModelMappings)
+ mgmt.GET("/ampcode/force-model-mappings", h.GetAmpForceModelMappings)
+ mgmt.PUT("/ampcode/force-model-mappings", h.PutAmpForceModelMappings)
+ }
+ return r
+}
+
+// TestGetAmpCode verifies GET /v0/management/ampcode returns full ampcode config.
+func TestGetAmpCode(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+
+ var resp map[string]config.AmpCode
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal response: %v", err)
+ }
+
+ ampcode := resp["ampcode"]
+ if ampcode.UpstreamURL != "https://example.com" {
+ t.Errorf("expected upstream-url %q, got %q", "https://example.com", ampcode.UpstreamURL)
+ }
+ if len(ampcode.ModelMappings) != 1 {
+ t.Errorf("expected 1 model mapping, got %d", len(ampcode.ModelMappings))
+ }
+}
+
+// TestGetAmpUpstreamURL verifies GET /v0/management/ampcode/upstream-url returns the upstream URL.
+func TestGetAmpUpstreamURL(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-url", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+
+ var resp map[string]string
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal response: %v", err)
+ }
+
+ if resp["upstream-url"] != "https://example.com" {
+ t.Errorf("expected %q, got %q", "https://example.com", resp["upstream-url"])
+ }
+}
+
+// TestPutAmpUpstreamURL verifies PUT /v0/management/ampcode/upstream-url updates the upstream URL.
+func TestPutAmpUpstreamURL(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": "https://new-upstream.com"}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/upstream-url", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d: %s", http.StatusOK, w.Code, w.Body.String())
+ }
+}
+
+// TestDeleteAmpUpstreamURL verifies DELETE /v0/management/ampcode/upstream-url clears the upstream URL.
+func TestDeleteAmpUpstreamURL(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/upstream-url", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+}
+
+// TestGetAmpUpstreamAPIKey verifies GET /v0/management/ampcode/upstream-api-key returns the API key.
+func TestGetAmpUpstreamAPIKey(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-api-key", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+
+ var resp map[string]any
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal response: %v", err)
+ }
+
+ key := resp["upstream-api-key"].(string)
+ if key != "test-api-key-12345" {
+ t.Errorf("expected key %q, got %q", "test-api-key-12345", key)
+ }
+}
+
+// TestPutAmpUpstreamAPIKey verifies PUT /v0/management/ampcode/upstream-api-key updates the API key.
+func TestPutAmpUpstreamAPIKey(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": "new-secret-key"}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/upstream-api-key", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+}
+
+func TestPutAmpUpstreamAPIKeys_PersistsAndReturns(t *testing.T) {
+ h, configPath := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value":[{"upstream-api-key":" u1 ","api-keys":[" k1 ","","k2"]}]}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/upstream-api-keys", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d: %s", http.StatusOK, w.Code, w.Body.String())
+ }
+
+ // Verify it was persisted to disk
+ loaded, err := config.LoadConfig(configPath)
+ if err != nil {
+ t.Fatalf("failed to load config from disk: %v", err)
+ }
+ if len(loaded.AmpCode.UpstreamAPIKeys) != 1 {
+ t.Fatalf("expected 1 upstream-api-keys entry, got %d", len(loaded.AmpCode.UpstreamAPIKeys))
+ }
+ entry := loaded.AmpCode.UpstreamAPIKeys[0]
+ if entry.UpstreamAPIKey != "u1" {
+ t.Fatalf("expected upstream-api-key u1, got %q", entry.UpstreamAPIKey)
+ }
+ if len(entry.APIKeys) != 2 || entry.APIKeys[0] != "k1" || entry.APIKeys[1] != "k2" {
+ t.Fatalf("expected api-keys [k1 k2], got %#v", entry.APIKeys)
+ }
+
+ // Verify it is returned by GET /ampcode
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+ var resp map[string]config.AmpCode
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal response: %v", err)
+ }
+ if got := resp["ampcode"].UpstreamAPIKeys; len(got) != 1 || got[0].UpstreamAPIKey != "u1" {
+ t.Fatalf("expected upstream-api-keys to be present after update, got %#v", got)
+ }
+}
+
+func TestDeleteAmpUpstreamAPIKeys_ClearsAll(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ // Seed with one entry
+ putBody := `{"value":[{"upstream-api-key":"u1","api-keys":["k1"]}]}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/upstream-api-keys", bytes.NewBufferString(putBody))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d: %s", http.StatusOK, w.Code, w.Body.String())
+ }
+
+ deleteBody := `{"value":[]}`
+ req = httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/upstream-api-keys", bytes.NewBufferString(deleteBody))
+ req.Header.Set("Content-Type", "application/json")
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-api-keys", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+ var resp map[string][]config.AmpUpstreamAPIKeyEntry
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal response: %v", err)
+ }
+ if resp["upstream-api-keys"] != nil && len(resp["upstream-api-keys"]) != 0 {
+ t.Fatalf("expected cleared list, got %#v", resp["upstream-api-keys"])
+ }
+}
+
+// TestDeleteAmpUpstreamAPIKey verifies DELETE /v0/management/ampcode/upstream-api-key clears the API key.
+func TestDeleteAmpUpstreamAPIKey(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/upstream-api-key", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+}
+
+// TestGetAmpRestrictManagementToLocalhost verifies GET returns the localhost restriction setting.
+func TestGetAmpRestrictManagementToLocalhost(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/restrict-management-to-localhost", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+
+ var resp map[string]bool
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal response: %v", err)
+ }
+
+ if resp["restrict-management-to-localhost"] != true {
+ t.Error("expected restrict-management-to-localhost to be true")
+ }
+}
+
+// TestPutAmpRestrictManagementToLocalhost verifies PUT updates the localhost restriction setting.
+func TestPutAmpRestrictManagementToLocalhost(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": false}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/restrict-management-to-localhost", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+}
+
+// TestGetAmpModelMappings verifies GET /v0/management/ampcode/model-mappings returns all mappings.
+func TestGetAmpModelMappings(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+
+ var resp map[string][]config.AmpModelMapping
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal response: %v", err)
+ }
+
+ mappings := resp["model-mappings"]
+ if len(mappings) != 1 {
+ t.Fatalf("expected 1 mapping, got %d", len(mappings))
+ }
+ if mappings[0].From != "gpt-4" || mappings[0].To != "gemini-pro" {
+ t.Errorf("unexpected mapping: %+v", mappings[0])
+ }
+}
+
+// TestPutAmpModelMappings verifies PUT /v0/management/ampcode/model-mappings replaces all mappings.
+func TestPutAmpModelMappings(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": [{"from": "claude-3", "to": "gpt-4o"}, {"from": "gemini", "to": "claude"}]}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d: %s", http.StatusOK, w.Code, w.Body.String())
+ }
+}
+
+// TestPatchAmpModelMappings verifies PATCH updates existing mappings and adds new ones.
+func TestPatchAmpModelMappings(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": [{"from": "gpt-4", "to": "updated-model"}, {"from": "new-model", "to": "target"}]}`
+ req := httptest.NewRequest(http.MethodPatch, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d: %s", http.StatusOK, w.Code, w.Body.String())
+ }
+}
+
+// TestDeleteAmpModelMappings_Specific verifies DELETE removes specified mappings by "from" field.
+func TestDeleteAmpModelMappings_Specific(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": ["gpt-4"]}`
+ req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+}
+
+// TestDeleteAmpModelMappings_All verifies DELETE with empty body removes all mappings.
+func TestDeleteAmpModelMappings_All(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/model-mappings", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+}
+
+// TestGetAmpForceModelMappings verifies GET returns the force-model-mappings setting.
+func TestGetAmpForceModelMappings(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/force-model-mappings", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+
+ var resp map[string]bool
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal response: %v", err)
+ }
+
+ if resp["force-model-mappings"] != false {
+ t.Error("expected force-model-mappings to be false")
+ }
+}
+
+// TestPutAmpForceModelMappings verifies PUT updates the force-model-mappings setting.
+func TestPutAmpForceModelMappings(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": true}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/force-model-mappings", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+}
+
+// TestPutAmpModelMappings_VerifyState verifies PUT replaces mappings and state is persisted.
+func TestPutAmpModelMappings_VerifyState(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": [{"from": "model-a", "to": "model-b"}, {"from": "model-c", "to": "model-d"}, {"from": "model-e", "to": "model-f"}]}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("PUT failed: status %d, body: %s", w.Code, w.Body.String())
+ }
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ var resp map[string][]config.AmpModelMapping
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ mappings := resp["model-mappings"]
+ if len(mappings) != 3 {
+ t.Fatalf("expected 3 mappings, got %d", len(mappings))
+ }
+
+ expected := map[string]string{"model-a": "model-b", "model-c": "model-d", "model-e": "model-f"}
+ for _, m := range mappings {
+ if expected[m.From] != m.To {
+ t.Errorf("mapping %q -> expected %q, got %q", m.From, expected[m.From], m.To)
+ }
+ }
+}
+
+// TestPatchAmpModelMappings_VerifyState verifies PATCH merges mappings correctly.
+func TestPatchAmpModelMappings_VerifyState(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": [{"from": "gpt-4", "to": "updated-target"}, {"from": "new-model", "to": "new-target"}]}`
+ req := httptest.NewRequest(http.MethodPatch, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("PATCH failed: status %d", w.Code)
+ }
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ var resp map[string][]config.AmpModelMapping
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ mappings := resp["model-mappings"]
+ if len(mappings) != 2 {
+ t.Fatalf("expected 2 mappings (1 updated + 1 new), got %d", len(mappings))
+ }
+
+ found := make(map[string]string)
+ for _, m := range mappings {
+ found[m.From] = m.To
+ }
+
+ if found["gpt-4"] != "updated-target" {
+ t.Errorf("gpt-4 should map to updated-target, got %q", found["gpt-4"])
+ }
+ if found["new-model"] != "new-target" {
+ t.Errorf("new-model should map to new-target, got %q", found["new-model"])
+ }
+}
+
+// TestDeleteAmpModelMappings_VerifyState verifies DELETE removes specific mappings and keeps others.
+func TestDeleteAmpModelMappings_VerifyState(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ putBody := `{"value": [{"from": "a", "to": "1"}, {"from": "b", "to": "2"}, {"from": "c", "to": "3"}]}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(putBody))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ delBody := `{"value": ["a", "c"]}`
+ req = httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(delBody))
+ req.Header.Set("Content-Type", "application/json")
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("DELETE failed: status %d", w.Code)
+ }
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ var resp map[string][]config.AmpModelMapping
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ mappings := resp["model-mappings"]
+ if len(mappings) != 1 {
+ t.Fatalf("expected 1 mapping remaining, got %d", len(mappings))
+ }
+ if mappings[0].From != "b" || mappings[0].To != "2" {
+ t.Errorf("expected b->2, got %s->%s", mappings[0].From, mappings[0].To)
+ }
+}
+
+// TestDeleteAmpModelMappings_NonExistent verifies DELETE with non-existent mapping doesn't affect existing ones.
+func TestDeleteAmpModelMappings_NonExistent(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ delBody := `{"value": ["non-existent-model"]}`
+ req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(delBody))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ var resp map[string][]config.AmpModelMapping
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ if len(resp["model-mappings"]) != 1 {
+ t.Errorf("original mapping should remain, got %d mappings", len(resp["model-mappings"]))
+ }
+}
+
+// TestPutAmpModelMappings_Empty verifies PUT with empty array clears all mappings.
+func TestPutAmpModelMappings_Empty(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": []}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ var resp map[string][]config.AmpModelMapping
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ if len(resp["model-mappings"]) != 0 {
+ t.Errorf("expected 0 mappings, got %d", len(resp["model-mappings"]))
+ }
+}
+
+// TestPutAmpUpstreamURL_VerifyState verifies PUT updates upstream URL and persists state.
+func TestPutAmpUpstreamURL_VerifyState(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": "https://new-api.example.com"}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/upstream-url", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("PUT failed: status %d", w.Code)
+ }
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-url", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ var resp map[string]string
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ if resp["upstream-url"] != "https://new-api.example.com" {
+ t.Errorf("expected %q, got %q", "https://new-api.example.com", resp["upstream-url"])
+ }
+}
+
+// TestDeleteAmpUpstreamURL_VerifyState verifies DELETE clears upstream URL.
+func TestDeleteAmpUpstreamURL_VerifyState(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/upstream-url", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("DELETE failed: status %d", w.Code)
+ }
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-url", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ var resp map[string]string
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ if resp["upstream-url"] != "" {
+ t.Errorf("expected empty string, got %q", resp["upstream-url"])
+ }
+}
+
+// TestPutAmpUpstreamAPIKey_VerifyState verifies PUT updates API key and persists state.
+func TestPutAmpUpstreamAPIKey_VerifyState(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": "new-secret-api-key-xyz"}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/upstream-api-key", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("PUT failed: status %d", w.Code)
+ }
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-api-key", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ var resp map[string]string
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ if resp["upstream-api-key"] != "new-secret-api-key-xyz" {
+ t.Errorf("expected %q, got %q", "new-secret-api-key-xyz", resp["upstream-api-key"])
+ }
+}
+
+// TestDeleteAmpUpstreamAPIKey_VerifyState verifies DELETE clears API key.
+func TestDeleteAmpUpstreamAPIKey_VerifyState(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/upstream-api-key", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("DELETE failed: status %d", w.Code)
+ }
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-api-key", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ var resp map[string]string
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ if resp["upstream-api-key"] != "" {
+ t.Errorf("expected empty string, got %q", resp["upstream-api-key"])
+ }
+}
+
+// TestPutAmpRestrictManagementToLocalhost_VerifyState verifies PUT updates localhost restriction.
+func TestPutAmpRestrictManagementToLocalhost_VerifyState(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": false}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/restrict-management-to-localhost", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("PUT failed: status %d", w.Code)
+ }
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/restrict-management-to-localhost", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ var resp map[string]bool
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ if resp["restrict-management-to-localhost"] != false {
+ t.Error("expected false after update")
+ }
+}
+
+// TestPutAmpForceModelMappings_VerifyState verifies PUT updates force-model-mappings setting.
+func TestPutAmpForceModelMappings_VerifyState(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{"value": true}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/force-model-mappings", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("PUT failed: status %d", w.Code)
+ }
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/force-model-mappings", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ var resp map[string]bool
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ if resp["force-model-mappings"] != true {
+ t.Error("expected true after update")
+ }
+}
+
+// TestPutBoolField_EmptyObject verifies PUT with empty object returns 400.
+func TestPutBoolField_EmptyObject(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ body := `{}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/force-model-mappings", bytes.NewBufferString(body))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusBadRequest {
+ t.Fatalf("expected status %d for empty object, got %d", http.StatusBadRequest, w.Code)
+ }
+}
+
+// TestComplexMappingsWorkflow tests a full workflow: PUT, PATCH, DELETE, and GET.
+func TestComplexMappingsWorkflow(t *testing.T) {
+ h, _ := newAmpTestHandler(t)
+ r := setupAmpRouter(h)
+
+ putBody := `{"value": [{"from": "m1", "to": "t1"}, {"from": "m2", "to": "t2"}, {"from": "m3", "to": "t3"}, {"from": "m4", "to": "t4"}]}`
+ req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(putBody))
+ req.Header.Set("Content-Type", "application/json")
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ patchBody := `{"value": [{"from": "m2", "to": "t2-updated"}, {"from": "m5", "to": "t5"}]}`
+ req = httptest.NewRequest(http.MethodPatch, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(patchBody))
+ req.Header.Set("Content-Type", "application/json")
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ delBody := `{"value": ["m1", "m3"]}`
+ req = httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(delBody))
+ req.Header.Set("Content-Type", "application/json")
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
+ w = httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ var resp map[string][]config.AmpModelMapping
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ mappings := resp["model-mappings"]
+ if len(mappings) != 3 {
+ t.Fatalf("expected 3 mappings (m2, m4, m5), got %d", len(mappings))
+ }
+
+ expected := map[string]string{"m2": "t2-updated", "m4": "t4", "m5": "t5"}
+ found := make(map[string]string)
+ for _, m := range mappings {
+ found[m.From] = m.To
+ }
+
+ for from, to := range expected {
+ if found[from] != to {
+ t.Errorf("mapping %s: expected %q, got %q", from, to, found[from])
+ }
+ }
+}
+
+// TestNilHandlerGetAmpCode verifies handler works with empty config.
+func TestNilHandlerGetAmpCode(t *testing.T) {
+ cfg := &config.Config{}
+ h := management.NewHandler(cfg, "", nil)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+}
+
+// TestEmptyConfigGetAmpModelMappings verifies GET returns empty array for fresh config.
+func TestEmptyConfigGetAmpModelMappings(t *testing.T) {
+ cfg := &config.Config{}
+ tmpDir := t.TempDir()
+ configPath := filepath.Join(tmpDir, "config.yaml")
+ if err := os.WriteFile(configPath, []byte("port: 8080\n"), 0644); err != nil {
+ t.Fatalf("failed to write config: %v", err)
+ }
+
+ h := management.NewHandler(cfg, configPath, nil)
+ r := setupAmpRouter(h)
+
+ req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
+ w := httptest.NewRecorder()
+ r.ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
+ }
+
+ var resp map[string][]config.AmpModelMapping
+ if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("failed to unmarshal: %v", err)
+ }
+
+ if len(resp["model-mappings"]) != 0 {
+ t.Errorf("expected 0 mappings, got %d", len(resp["model-mappings"]))
+ }
+}
diff --git a/test/builtin_tools_translation_test.go b/test/builtin_tools_translation_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..b4ca7b0da6cf367744fbc87ce19dbf15280afce4
--- /dev/null
+++ b/test/builtin_tools_translation_test.go
@@ -0,0 +1,54 @@
+package test
+
+import (
+ "testing"
+
+ _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator"
+
+ sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
+ "github.com/tidwall/gjson"
+)
+
+func TestOpenAIToCodex_PreservesBuiltinTools(t *testing.T) {
+ in := []byte(`{
+ "model":"gpt-5",
+ "messages":[{"role":"user","content":"hi"}],
+ "tools":[{"type":"web_search","search_context_size":"high"}],
+ "tool_choice":{"type":"web_search"}
+ }`)
+
+ out := sdktranslator.TranslateRequest(sdktranslator.FormatOpenAI, sdktranslator.FormatCodex, "gpt-5", in, false)
+
+ if got := gjson.GetBytes(out, "tools.#").Int(); got != 1 {
+ t.Fatalf("expected 1 tool, got %d: %s", got, string(out))
+ }
+ if got := gjson.GetBytes(out, "tools.0.type").String(); got != "web_search" {
+ t.Fatalf("expected tools[0].type=web_search, got %q: %s", got, string(out))
+ }
+ if got := gjson.GetBytes(out, "tools.0.search_context_size").String(); got != "high" {
+ t.Fatalf("expected tools[0].search_context_size=high, got %q: %s", got, string(out))
+ }
+ if got := gjson.GetBytes(out, "tool_choice.type").String(); got != "web_search" {
+ t.Fatalf("expected tool_choice.type=web_search, got %q: %s", got, string(out))
+ }
+}
+
+func TestOpenAIResponsesToOpenAI_PreservesBuiltinTools(t *testing.T) {
+ in := []byte(`{
+ "model":"gpt-5",
+ "input":[{"role":"user","content":[{"type":"input_text","text":"hi"}]}],
+ "tools":[{"type":"web_search","search_context_size":"low"}]
+ }`)
+
+ out := sdktranslator.TranslateRequest(sdktranslator.FormatOpenAIResponse, sdktranslator.FormatOpenAI, "gpt-5", in, false)
+
+ if got := gjson.GetBytes(out, "tools.#").Int(); got != 1 {
+ t.Fatalf("expected 1 tool, got %d: %s", got, string(out))
+ }
+ if got := gjson.GetBytes(out, "tools.0.type").String(); got != "web_search" {
+ t.Fatalf("expected tools[0].type=web_search, got %q: %s", got, string(out))
+ }
+ if got := gjson.GetBytes(out, "tools.0.search_context_size").String(); got != "low" {
+ t.Fatalf("expected tools[0].search_context_size=low, got %q: %s", got, string(out))
+ }
+}
diff --git a/test/config_migration_test.go b/test/config_migration_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..2ed878827769ade4dd95a2cf91b63865fd8dc97b
--- /dev/null
+++ b/test/config_migration_test.go
@@ -0,0 +1,195 @@
+package test
+
+import (
+ "os"
+ "path/filepath"
+ "strings"
+ "testing"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+)
+
+func TestLegacyConfigMigration(t *testing.T) {
+ t.Run("onlyLegacyFields", func(t *testing.T) {
+ path := writeConfig(t, `
+port: 8080
+generative-language-api-key:
+ - "legacy-gemini-1"
+openai-compatibility:
+ - name: "legacy-provider"
+ base-url: "https://example.com"
+ api-keys:
+ - "legacy-openai-1"
+amp-upstream-url: "https://amp.example.com"
+amp-upstream-api-key: "amp-legacy-key"
+amp-restrict-management-to-localhost: false
+amp-model-mappings:
+ - from: "old-model"
+ to: "new-model"
+`)
+ cfg, err := config.LoadConfig(path)
+ if err != nil {
+ t.Fatalf("load legacy config: %v", err)
+ }
+ if got := len(cfg.GeminiKey); got != 1 || cfg.GeminiKey[0].APIKey != "legacy-gemini-1" {
+ t.Fatalf("gemini migration mismatch: %+v", cfg.GeminiKey)
+ }
+ if got := len(cfg.OpenAICompatibility); got != 1 {
+ t.Fatalf("expected 1 openai-compat provider, got %d", got)
+ }
+ if entries := cfg.OpenAICompatibility[0].APIKeyEntries; len(entries) != 1 || entries[0].APIKey != "legacy-openai-1" {
+ t.Fatalf("openai-compat migration mismatch: %+v", entries)
+ }
+ if cfg.AmpCode.UpstreamURL != "https://amp.example.com" || cfg.AmpCode.UpstreamAPIKey != "amp-legacy-key" {
+ t.Fatalf("amp migration failed: %+v", cfg.AmpCode)
+ }
+ if cfg.AmpCode.RestrictManagementToLocalhost {
+ t.Fatalf("expected amp restriction to be false after migration")
+ }
+ if got := len(cfg.AmpCode.ModelMappings); got != 1 || cfg.AmpCode.ModelMappings[0].From != "old-model" {
+ t.Fatalf("amp mappings migration mismatch: %+v", cfg.AmpCode.ModelMappings)
+ }
+ updated := readFile(t, path)
+ if strings.Contains(updated, "generative-language-api-key") {
+ t.Fatalf("legacy gemini key still present:\n%s", updated)
+ }
+ if strings.Contains(updated, "amp-upstream-url") || strings.Contains(updated, "amp-restrict-management-to-localhost") {
+ t.Fatalf("legacy amp keys still present:\n%s", updated)
+ }
+ if strings.Contains(updated, "\n api-keys:") {
+ t.Fatalf("legacy openai compat keys still present:\n%s", updated)
+ }
+ })
+
+ t.Run("mixedLegacyAndNewFields", func(t *testing.T) {
+ path := writeConfig(t, `
+gemini-api-key:
+ - api-key: "new-gemini"
+generative-language-api-key:
+ - "new-gemini"
+ - "legacy-gemini-only"
+openai-compatibility:
+ - name: "mixed-provider"
+ base-url: "https://mixed.example.com"
+ api-key-entries:
+ - api-key: "new-entry"
+ api-keys:
+ - "legacy-entry"
+ - "new-entry"
+`)
+ cfg, err := config.LoadConfig(path)
+ if err != nil {
+ t.Fatalf("load mixed config: %v", err)
+ }
+ if got := len(cfg.GeminiKey); got != 2 {
+ t.Fatalf("expected 2 gemini entries, got %d: %+v", got, cfg.GeminiKey)
+ }
+ seen := make(map[string]struct{}, len(cfg.GeminiKey))
+ for _, entry := range cfg.GeminiKey {
+ if _, exists := seen[entry.APIKey]; exists {
+ t.Fatalf("duplicate gemini key %q after migration", entry.APIKey)
+ }
+ seen[entry.APIKey] = struct{}{}
+ }
+ provider := cfg.OpenAICompatibility[0]
+ if got := len(provider.APIKeyEntries); got != 2 {
+ t.Fatalf("expected 2 openai entries, got %d: %+v", got, provider.APIKeyEntries)
+ }
+ entrySeen := make(map[string]struct{}, len(provider.APIKeyEntries))
+ for _, entry := range provider.APIKeyEntries {
+ if _, ok := entrySeen[entry.APIKey]; ok {
+ t.Fatalf("duplicate openai key %q after migration", entry.APIKey)
+ }
+ entrySeen[entry.APIKey] = struct{}{}
+ }
+ })
+
+ t.Run("onlyNewFields", func(t *testing.T) {
+ path := writeConfig(t, `
+gemini-api-key:
+ - api-key: "new-only"
+openai-compatibility:
+ - name: "new-only-provider"
+ base-url: "https://new-only.example.com"
+ api-key-entries:
+ - api-key: "new-only-entry"
+ampcode:
+ upstream-url: "https://amp.new"
+ upstream-api-key: "new-amp-key"
+ restrict-management-to-localhost: true
+ model-mappings:
+ - from: "a"
+ to: "b"
+`)
+ cfg, err := config.LoadConfig(path)
+ if err != nil {
+ t.Fatalf("load new config: %v", err)
+ }
+ if len(cfg.GeminiKey) != 1 || cfg.GeminiKey[0].APIKey != "new-only" {
+ t.Fatalf("unexpected gemini entries: %+v", cfg.GeminiKey)
+ }
+ if len(cfg.OpenAICompatibility) != 1 || len(cfg.OpenAICompatibility[0].APIKeyEntries) != 1 {
+ t.Fatalf("unexpected openai compat entries: %+v", cfg.OpenAICompatibility)
+ }
+ if cfg.AmpCode.UpstreamURL != "https://amp.new" || cfg.AmpCode.UpstreamAPIKey != "new-amp-key" {
+ t.Fatalf("unexpected amp config: %+v", cfg.AmpCode)
+ }
+ })
+
+ t.Run("duplicateNamesDifferentBase", func(t *testing.T) {
+ path := writeConfig(t, `
+openai-compatibility:
+ - name: "dup-provider"
+ base-url: "https://provider-a"
+ api-keys:
+ - "key-a"
+ - name: "dup-provider"
+ base-url: "https://provider-b"
+ api-keys:
+ - "key-b"
+`)
+ cfg, err := config.LoadConfig(path)
+ if err != nil {
+ t.Fatalf("load duplicate config: %v", err)
+ }
+ if len(cfg.OpenAICompatibility) != 2 {
+ t.Fatalf("expected 2 providers, got %d", len(cfg.OpenAICompatibility))
+ }
+ for _, entry := range cfg.OpenAICompatibility {
+ if len(entry.APIKeyEntries) != 1 {
+ t.Fatalf("expected 1 key entry per provider: %+v", entry)
+ }
+ switch entry.BaseURL {
+ case "https://provider-a":
+ if entry.APIKeyEntries[0].APIKey != "key-a" {
+ t.Fatalf("provider-a key mismatch: %+v", entry.APIKeyEntries)
+ }
+ case "https://provider-b":
+ if entry.APIKeyEntries[0].APIKey != "key-b" {
+ t.Fatalf("provider-b key mismatch: %+v", entry.APIKeyEntries)
+ }
+ default:
+ t.Fatalf("unexpected provider base url: %s", entry.BaseURL)
+ }
+ }
+ })
+}
+
+func writeConfig(t *testing.T, content string) string {
+ t.Helper()
+ dir := t.TempDir()
+ path := filepath.Join(dir, "config.yaml")
+ if err := os.WriteFile(path, []byte(strings.TrimSpace(content)+"\n"), 0o644); err != nil {
+ t.Fatalf("write temp config: %v", err)
+ }
+ return path
+}
+
+func readFile(t *testing.T, path string) string {
+ t.Helper()
+ data, err := os.ReadFile(path)
+ if err != nil {
+ t.Fatalf("read temp config: %v", err)
+ }
+ return string(data)
+}
diff --git a/test/thinking_conversion_test.go b/test/thinking_conversion_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..fc20199ed43b742f9e19d38460793980de9b0947
--- /dev/null
+++ b/test/thinking_conversion_test.go
@@ -0,0 +1,2798 @@
+package test
+
+import (
+ "fmt"
+ "strings"
+ "testing"
+ "time"
+
+ _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator"
+
+ // Import provider packages to trigger init() registration of ProviderAppliers
+ _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/antigravity"
+ _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/claude"
+ _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/codex"
+ _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/gemini"
+ _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/geminicli"
+ _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/iflow"
+ _ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking/provider/openai"
+
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
+ "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
+ sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
+ "github.com/tidwall/gjson"
+ "github.com/tidwall/sjson"
+)
+
+// thinkingTestCase represents a common test case structure for both suffix and body tests.
+type thinkingTestCase struct {
+ name string
+ from string
+ to string
+ model string
+ inputJSON string
+ expectField string
+ expectValue string
+ includeThoughts string
+ expectErr bool
+}
+
+// TestThinkingE2EMatrix_Suffix tests the thinking configuration transformation using model name suffix.
+// Data flow: Input JSON → TranslateRequest → ApplyThinking → Validate Output
+// No helper functions are used; all test data is inline.
+func TestThinkingE2EMatrix_Suffix(t *testing.T) {
+ reg := registry.GetGlobalRegistry()
+ uid := fmt.Sprintf("thinking-e2e-suffix-%d", time.Now().UnixNano())
+
+ reg.RegisterClient(uid, "test", getTestModels())
+ defer reg.UnregisterClient(uid)
+
+ cases := []thinkingTestCase{
+ // level-model (Levels=minimal/low/medium/high, ZeroAllowed=false, DynamicAllowed=false)
+
+ // Case 1: No suffix → injected default → medium
+ {
+ name: "1",
+ from: "openai",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 2: Specified medium → medium
+ {
+ name: "2",
+ from: "openai",
+ to: "codex",
+ model: "level-model(medium)",
+ inputJSON: `{"model":"level-model(medium)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 3: Specified xhigh → out of range error
+ {
+ name: "3",
+ from: "openai",
+ to: "codex",
+ model: "level-model(xhigh)",
+ inputJSON: `{"model":"level-model(xhigh)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 4: Level none → clamped to minimal (ZeroAllowed=false)
+ {
+ name: "4",
+ from: "openai",
+ to: "codex",
+ model: "level-model(none)",
+ inputJSON: `{"model":"level-model(none)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "minimal",
+ expectErr: false,
+ },
+ // Case 5: Level auto → DynamicAllowed=false → medium (mid-range)
+ {
+ name: "5",
+ from: "openai",
+ to: "codex",
+ model: "level-model(auto)",
+ inputJSON: `{"model":"level-model(auto)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 6: No suffix from gemini → injected default → medium
+ {
+ name: "6",
+ from: "gemini",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 7: Budget 8192 → medium
+ {
+ name: "7",
+ from: "gemini",
+ to: "codex",
+ model: "level-model(8192)",
+ inputJSON: `{"model":"level-model(8192)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 8: Budget 64000 → clamped to high
+ {
+ name: "8",
+ from: "gemini",
+ to: "codex",
+ model: "level-model(64000)",
+ inputJSON: `{"model":"level-model(64000)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "high",
+ expectErr: false,
+ },
+ // Case 9: Budget 0 → clamped to minimal (ZeroAllowed=false)
+ {
+ name: "9",
+ from: "gemini",
+ to: "codex",
+ model: "level-model(0)",
+ inputJSON: `{"model":"level-model(0)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "minimal",
+ expectErr: false,
+ },
+ // Case 10: Budget -1 → auto → DynamicAllowed=false → medium (mid-range)
+ {
+ name: "10",
+ from: "gemini",
+ to: "codex",
+ model: "level-model(-1)",
+ inputJSON: `{"model":"level-model(-1)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 11: Claude source no suffix → passthrough (no thinking)
+ {
+ name: "11",
+ from: "claude",
+ to: "openai",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 12: Budget 8192 → medium
+ {
+ name: "12",
+ from: "claude",
+ to: "openai",
+ model: "level-model(8192)",
+ inputJSON: `{"model":"level-model(8192)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning_effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 13: Budget 64000 → clamped to high
+ {
+ name: "13",
+ from: "claude",
+ to: "openai",
+ model: "level-model(64000)",
+ inputJSON: `{"model":"level-model(64000)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning_effort",
+ expectValue: "high",
+ expectErr: false,
+ },
+ // Case 14: Budget 0 → clamped to minimal (ZeroAllowed=false)
+ {
+ name: "14",
+ from: "claude",
+ to: "openai",
+ model: "level-model(0)",
+ inputJSON: `{"model":"level-model(0)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning_effort",
+ expectValue: "minimal",
+ expectErr: false,
+ },
+ // Case 15: Budget -1 → auto → DynamicAllowed=false → medium (mid-range)
+ {
+ name: "15",
+ from: "claude",
+ to: "openai",
+ model: "level-model(-1)",
+ inputJSON: `{"model":"level-model(-1)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning_effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+
+ // level-subset-model (Levels=low/high, ZeroAllowed=false, DynamicAllowed=false)
+
+ // Case 16: Budget 8192 → medium → rounded down to low
+ {
+ name: "16",
+ from: "gemini",
+ to: "openai",
+ model: "level-subset-model(8192)",
+ inputJSON: `{"model":"level-subset-model(8192)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning_effort",
+ expectValue: "low",
+ expectErr: false,
+ },
+ // Case 17: Budget 1 → minimal → clamped to low (min supported)
+ {
+ name: "17",
+ from: "claude",
+ to: "gemini",
+ model: "level-subset-model(1)",
+ inputJSON: `{"model":"level-subset-model(1)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingLevel",
+ expectValue: "low",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+
+ // gemini-budget-model (Min=128, Max=20000, ZeroAllowed=false, DynamicAllowed=true)
+
+ // Case 18: No suffix → passthrough
+ {
+ name: "18",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 19: Effort medium → 8192
+ {
+ name: "19",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-budget-model(medium)",
+ inputJSON: `{"model":"gemini-budget-model(medium)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 20: Effort xhigh → clamped to 20000 (max)
+ {
+ name: "20",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-budget-model(xhigh)",
+ inputJSON: `{"model":"gemini-budget-model(xhigh)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 21: Effort none → clamped to 128 (min) → includeThoughts=false
+ {
+ name: "21",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-budget-model(none)",
+ inputJSON: `{"model":"gemini-budget-model(none)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "128",
+ includeThoughts: "false",
+ expectErr: false,
+ },
+ // Case 22: Effort auto → DynamicAllowed=true → -1
+ {
+ name: "22",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-budget-model(auto)",
+ inputJSON: `{"model":"gemini-budget-model(auto)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "-1",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 23: Claude source no suffix → passthrough
+ {
+ name: "23",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 24: Budget 8192 → 8192
+ {
+ name: "24",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-budget-model(8192)",
+ inputJSON: `{"model":"gemini-budget-model(8192)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 25: Budget 64000 → clamped to 20000 (max)
+ {
+ name: "25",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-budget-model(64000)",
+ inputJSON: `{"model":"gemini-budget-model(64000)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 26: Budget 0 → clamped to 128 (min) → includeThoughts=false
+ {
+ name: "26",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-budget-model(0)",
+ inputJSON: `{"model":"gemini-budget-model(0)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "128",
+ includeThoughts: "false",
+ expectErr: false,
+ },
+ // Case 27: Budget -1 → DynamicAllowed=true → -1
+ {
+ name: "27",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-budget-model(-1)",
+ inputJSON: `{"model":"gemini-budget-model(-1)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "-1",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+
+ // gemini-mixed-model (Min=128, Max=32768, Levels=low/high, ZeroAllowed=false, DynamicAllowed=true)
+
+ // Case 28: OpenAI source no suffix → passthrough
+ {
+ name: "28",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-mixed-model",
+ inputJSON: `{"model":"gemini-mixed-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 29: Effort high → low/high supported → high
+ {
+ name: "29",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-mixed-model(high)",
+ inputJSON: `{"model":"gemini-mixed-model(high)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingLevel",
+ expectValue: "high",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 30: Effort xhigh → not in low/high → error
+ {
+ name: "30",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-mixed-model(xhigh)",
+ inputJSON: `{"model":"gemini-mixed-model(xhigh)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 31: Effort none → clamped to low (min supported) → includeThoughts=false
+ {
+ name: "31",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-mixed-model(none)",
+ inputJSON: `{"model":"gemini-mixed-model(none)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingLevel",
+ expectValue: "low",
+ includeThoughts: "false",
+ expectErr: false,
+ },
+ // Case 32: Effort auto → DynamicAllowed=true → -1 (budget)
+ {
+ name: "32",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-mixed-model(auto)",
+ inputJSON: `{"model":"gemini-mixed-model(auto)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "-1",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 33: Claude source no suffix → passthrough
+ {
+ name: "33",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-mixed-model",
+ inputJSON: `{"model":"gemini-mixed-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 34: Budget 8192 → 8192 (keep budget)
+ {
+ name: "34",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-mixed-model(8192)",
+ inputJSON: `{"model":"gemini-mixed-model(8192)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 35: Budget 64000 → clamped to 32768 (max)
+ {
+ name: "35",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-mixed-model(64000)",
+ inputJSON: `{"model":"gemini-mixed-model(64000)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "32768",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 36: Budget 0 → minimal → clamped to low (min level) → includeThoughts=false
+ {
+ name: "36",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-mixed-model(0)",
+ inputJSON: `{"model":"gemini-mixed-model(0)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingLevel",
+ expectValue: "low",
+ includeThoughts: "false",
+ expectErr: false,
+ },
+ // Case 37: Budget -1 → DynamicAllowed=true → -1 (budget)
+ {
+ name: "37",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-mixed-model(-1)",
+ inputJSON: `{"model":"gemini-mixed-model(-1)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "-1",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+
+ // claude-budget-model (Min=1024, Max=128000, ZeroAllowed=true, DynamicAllowed=false)
+
+ // Case 38: OpenAI source no suffix → passthrough
+ {
+ name: "38",
+ from: "openai",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 39: Effort medium → 8192
+ {
+ name: "39",
+ from: "openai",
+ to: "claude",
+ model: "claude-budget-model(medium)",
+ inputJSON: `{"model":"claude-budget-model(medium)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "8192",
+ expectErr: false,
+ },
+ // Case 40: Effort xhigh → clamped to 32768 (matrix value)
+ {
+ name: "40",
+ from: "openai",
+ to: "claude",
+ model: "claude-budget-model(xhigh)",
+ inputJSON: `{"model":"claude-budget-model(xhigh)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "32768",
+ expectErr: false,
+ },
+ // Case 41: Effort none → ZeroAllowed=true → disabled
+ {
+ name: "41",
+ from: "openai",
+ to: "claude",
+ model: "claude-budget-model(none)",
+ inputJSON: `{"model":"claude-budget-model(none)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "thinking.type",
+ expectValue: "disabled",
+ expectErr: false,
+ },
+ // Case 42: Effort auto → DynamicAllowed=false → 64512 (mid-range)
+ {
+ name: "42",
+ from: "openai",
+ to: "claude",
+ model: "claude-budget-model(auto)",
+ inputJSON: `{"model":"claude-budget-model(auto)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "64512",
+ expectErr: false,
+ },
+ // Case 43: Gemini source no suffix → passthrough
+ {
+ name: "43",
+ from: "gemini",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 44: Budget 8192 → 8192
+ {
+ name: "44",
+ from: "gemini",
+ to: "claude",
+ model: "claude-budget-model(8192)",
+ inputJSON: `{"model":"claude-budget-model(8192)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "8192",
+ expectErr: false,
+ },
+ // Case 45: Budget 200000 → clamped to 128000 (max)
+ {
+ name: "45",
+ from: "gemini",
+ to: "claude",
+ model: "claude-budget-model(200000)",
+ inputJSON: `{"model":"claude-budget-model(200000)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "128000",
+ expectErr: false,
+ },
+ // Case 46: Budget 0 → ZeroAllowed=true → disabled
+ {
+ name: "46",
+ from: "gemini",
+ to: "claude",
+ model: "claude-budget-model(0)",
+ inputJSON: `{"model":"claude-budget-model(0)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "thinking.type",
+ expectValue: "disabled",
+ expectErr: false,
+ },
+ // Case 47: Budget -1 → auto → DynamicAllowed=false → 64512 (mid-range)
+ {
+ name: "47",
+ from: "gemini",
+ to: "claude",
+ model: "claude-budget-model(-1)",
+ inputJSON: `{"model":"claude-budget-model(-1)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "64512",
+ expectErr: false,
+ },
+
+ // antigravity-budget-model (Min=128, Max=20000, ZeroAllowed=true, DynamicAllowed=true)
+
+ // Case 48: Gemini to Antigravity no suffix → passthrough
+ {
+ name: "48",
+ from: "gemini",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 49: Effort medium → 8192
+ {
+ name: "49",
+ from: "gemini",
+ to: "antigravity",
+ model: "antigravity-budget-model(medium)",
+ inputJSON: `{"model":"antigravity-budget-model(medium)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 50: Effort xhigh → clamped to 20000 (max)
+ {
+ name: "50",
+ from: "gemini",
+ to: "antigravity",
+ model: "antigravity-budget-model(xhigh)",
+ inputJSON: `{"model":"antigravity-budget-model(xhigh)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 51: Effort none → ZeroAllowed=true → 0 → includeThoughts=false
+ {
+ name: "51",
+ from: "gemini",
+ to: "antigravity",
+ model: "antigravity-budget-model(none)",
+ inputJSON: `{"model":"antigravity-budget-model(none)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "0",
+ includeThoughts: "false",
+ expectErr: false,
+ },
+ // Case 52: Effort auto → DynamicAllowed=true → -1
+ {
+ name: "52",
+ from: "gemini",
+ to: "antigravity",
+ model: "antigravity-budget-model(auto)",
+ inputJSON: `{"model":"antigravity-budget-model(auto)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "-1",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 53: Claude to Antigravity no suffix → passthrough
+ {
+ name: "53",
+ from: "claude",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 54: Budget 8192 → 8192
+ {
+ name: "54",
+ from: "claude",
+ to: "antigravity",
+ model: "antigravity-budget-model(8192)",
+ inputJSON: `{"model":"antigravity-budget-model(8192)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 55: Budget 64000 → clamped to 20000 (max)
+ {
+ name: "55",
+ from: "claude",
+ to: "antigravity",
+ model: "antigravity-budget-model(64000)",
+ inputJSON: `{"model":"antigravity-budget-model(64000)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 56: Budget 0 → ZeroAllowed=true → 0 → includeThoughts=false
+ {
+ name: "56",
+ from: "claude",
+ to: "antigravity",
+ model: "antigravity-budget-model(0)",
+ inputJSON: `{"model":"antigravity-budget-model(0)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "0",
+ includeThoughts: "false",
+ expectErr: false,
+ },
+ // Case 57: Budget -1 → DynamicAllowed=true → -1
+ {
+ name: "57",
+ from: "claude",
+ to: "antigravity",
+ model: "antigravity-budget-model(-1)",
+ inputJSON: `{"model":"antigravity-budget-model(-1)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "-1",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+
+ // no-thinking-model (Thinking=nil)
+
+ // Case 58: No thinking support → no configuration
+ {
+ name: "58",
+ from: "gemini",
+ to: "openai",
+ model: "no-thinking-model",
+ inputJSON: `{"model":"no-thinking-model","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 59: Budget 8192 → no thinking support → suffix stripped → no configuration
+ {
+ name: "59",
+ from: "gemini",
+ to: "openai",
+ model: "no-thinking-model(8192)",
+ inputJSON: `{"model":"no-thinking-model(8192)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 60: Budget 0 → suffix stripped → no configuration
+ {
+ name: "60",
+ from: "gemini",
+ to: "openai",
+ model: "no-thinking-model(0)",
+ inputJSON: `{"model":"no-thinking-model(0)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 61: Budget -1 → suffix stripped → no configuration
+ {
+ name: "61",
+ from: "gemini",
+ to: "openai",
+ model: "no-thinking-model(-1)",
+ inputJSON: `{"model":"no-thinking-model(-1)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 62: Claude source no suffix → no configuration
+ {
+ name: "62",
+ from: "claude",
+ to: "openai",
+ model: "no-thinking-model",
+ inputJSON: `{"model":"no-thinking-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 63: Budget 8192 → suffix stripped → no configuration
+ {
+ name: "63",
+ from: "claude",
+ to: "openai",
+ model: "no-thinking-model(8192)",
+ inputJSON: `{"model":"no-thinking-model(8192)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 64: Budget 0 → suffix stripped → no configuration
+ {
+ name: "64",
+ from: "claude",
+ to: "openai",
+ model: "no-thinking-model(0)",
+ inputJSON: `{"model":"no-thinking-model(0)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 65: Budget -1 → suffix stripped → no configuration
+ {
+ name: "65",
+ from: "claude",
+ to: "openai",
+ model: "no-thinking-model(-1)",
+ inputJSON: `{"model":"no-thinking-model(-1)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+
+ // user-defined-model (UserDefined=true, Thinking=nil)
+
+ // Case 66: User defined model no suffix → passthrough
+ {
+ name: "66",
+ from: "gemini",
+ to: "openai",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 67: Budget 8192 → passthrough logic → medium
+ {
+ name: "67",
+ from: "gemini",
+ to: "openai",
+ model: "user-defined-model(8192)",
+ inputJSON: `{"model":"user-defined-model(8192)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning_effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 68: Budget 64000 → passthrough logic → xhigh
+ {
+ name: "68",
+ from: "gemini",
+ to: "openai",
+ model: "user-defined-model(64000)",
+ inputJSON: `{"model":"user-defined-model(64000)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning_effort",
+ expectValue: "xhigh",
+ expectErr: false,
+ },
+ // Case 69: Budget 0 → passthrough logic → none
+ {
+ name: "69",
+ from: "gemini",
+ to: "openai",
+ model: "user-defined-model(0)",
+ inputJSON: `{"model":"user-defined-model(0)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning_effort",
+ expectValue: "none",
+ expectErr: false,
+ },
+ // Case 70: Budget -1 → passthrough logic → auto
+ {
+ name: "70",
+ from: "gemini",
+ to: "openai",
+ model: "user-defined-model(-1)",
+ inputJSON: `{"model":"user-defined-model(-1)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning_effort",
+ expectValue: "auto",
+ expectErr: false,
+ },
+ // Case 71: Claude to Codex no suffix → injected default → medium
+ {
+ name: "71",
+ from: "claude",
+ to: "codex",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 72: Budget 8192 → passthrough logic → medium
+ {
+ name: "72",
+ from: "claude",
+ to: "codex",
+ model: "user-defined-model(8192)",
+ inputJSON: `{"model":"user-defined-model(8192)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 73: Budget 64000 → passthrough logic → xhigh
+ {
+ name: "73",
+ from: "claude",
+ to: "codex",
+ model: "user-defined-model(64000)",
+ inputJSON: `{"model":"user-defined-model(64000)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "xhigh",
+ expectErr: false,
+ },
+ // Case 74: Budget 0 → passthrough logic → none
+ {
+ name: "74",
+ from: "claude",
+ to: "codex",
+ model: "user-defined-model(0)",
+ inputJSON: `{"model":"user-defined-model(0)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "none",
+ expectErr: false,
+ },
+ // Case 75: Budget -1 → passthrough logic → auto
+ {
+ name: "75",
+ from: "claude",
+ to: "codex",
+ model: "user-defined-model(-1)",
+ inputJSON: `{"model":"user-defined-model(-1)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "auto",
+ expectErr: false,
+ },
+ // Case 76: OpenAI to Gemini budget 8192 → passthrough → 8192
+ {
+ name: "76",
+ from: "openai",
+ to: "gemini",
+ model: "user-defined-model(8192)",
+ inputJSON: `{"model":"user-defined-model(8192)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 77: OpenAI to Claude budget 8192 → passthrough → 8192
+ {
+ name: "77",
+ from: "openai",
+ to: "claude",
+ model: "user-defined-model(8192)",
+ inputJSON: `{"model":"user-defined-model(8192)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "8192",
+ expectErr: false,
+ },
+ // Case 78: OpenAI-Response to Gemini budget 8192 → passthrough → 8192
+ {
+ name: "78",
+ from: "openai-response",
+ to: "gemini",
+ model: "user-defined-model(8192)",
+ inputJSON: `{"model":"user-defined-model(8192)","input":[{"role":"user","content":"hi"}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 79: OpenAI-Response to Claude budget 8192 → passthrough → 8192
+ {
+ name: "79",
+ from: "openai-response",
+ to: "claude",
+ model: "user-defined-model(8192)",
+ inputJSON: `{"model":"user-defined-model(8192)","input":[{"role":"user","content":"hi"}]}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "8192",
+ expectErr: false,
+ },
+
+ // Same-protocol passthrough tests (80-89)
+
+ // Case 80: OpenAI to OpenAI, level high → passthrough reasoning_effort
+ {
+ name: "80",
+ from: "openai",
+ to: "openai",
+ model: "level-model(high)",
+ inputJSON: `{"model":"level-model(high)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning_effort",
+ expectValue: "high",
+ expectErr: false,
+ },
+ // Case 81: OpenAI to OpenAI, level xhigh → out of range error
+ {
+ name: "81",
+ from: "openai",
+ to: "openai",
+ model: "level-model(xhigh)",
+ inputJSON: `{"model":"level-model(xhigh)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 82: OpenAI-Response to Codex, level high → passthrough reasoning.effort
+ {
+ name: "82",
+ from: "openai-response",
+ to: "codex",
+ model: "level-model(high)",
+ inputJSON: `{"model":"level-model(high)","input":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "high",
+ expectErr: false,
+ },
+ // Case 83: OpenAI-Response to Codex, level xhigh → out of range error
+ {
+ name: "83",
+ from: "openai-response",
+ to: "codex",
+ model: "level-model(xhigh)",
+ inputJSON: `{"model":"level-model(xhigh)","input":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 84: Gemini to Gemini, budget 8192 → passthrough thinkingBudget
+ {
+ name: "84",
+ from: "gemini",
+ to: "gemini",
+ model: "gemini-budget-model(8192)",
+ inputJSON: `{"model":"gemini-budget-model(8192)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 85: Gemini to Gemini, budget 64000 → clamped to Max
+ {
+ name: "85",
+ from: "gemini",
+ to: "gemini",
+ model: "gemini-budget-model(64000)",
+ inputJSON: `{"model":"gemini-budget-model(64000)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 86: Claude to Claude, budget 8192 → passthrough thinking.budget_tokens
+ {
+ name: "86",
+ from: "claude",
+ to: "claude",
+ model: "claude-budget-model(8192)",
+ inputJSON: `{"model":"claude-budget-model(8192)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "8192",
+ expectErr: false,
+ },
+ // Case 87: Claude to Claude, budget 200000 → clamped to Max
+ {
+ name: "87",
+ from: "claude",
+ to: "claude",
+ model: "claude-budget-model(200000)",
+ inputJSON: `{"model":"claude-budget-model(200000)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "128000",
+ expectErr: false,
+ },
+ // Case 88: Gemini-CLI to Antigravity, budget 8192 → passthrough thinkingBudget
+ {
+ name: "88",
+ from: "gemini-cli",
+ to: "antigravity",
+ model: "antigravity-budget-model(8192)",
+ inputJSON: `{"model":"antigravity-budget-model(8192)","request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 89: Gemini-CLI to Antigravity, budget 64000 → clamped to Max
+ {
+ name: "89",
+ from: "gemini-cli",
+ to: "antigravity",
+ model: "antigravity-budget-model(64000)",
+ inputJSON: `{"model":"antigravity-budget-model(64000)","request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+
+ // iflow tests: glm-test and minimax-test (Cases 90-105)
+
+ // glm-test (from: openai, claude)
+ // Case 90: OpenAI to iflow, no suffix → passthrough
+ {
+ name: "90",
+ from: "openai",
+ to: "iflow",
+ model: "glm-test",
+ inputJSON: `{"model":"glm-test","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 91: OpenAI to iflow, (medium) → enable_thinking=true
+ {
+ name: "91",
+ from: "openai",
+ to: "iflow",
+ model: "glm-test(medium)",
+ inputJSON: `{"model":"glm-test(medium)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "chat_template_kwargs.enable_thinking",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 92: OpenAI to iflow, (auto) → enable_thinking=true
+ {
+ name: "92",
+ from: "openai",
+ to: "iflow",
+ model: "glm-test(auto)",
+ inputJSON: `{"model":"glm-test(auto)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "chat_template_kwargs.enable_thinking",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 93: OpenAI to iflow, (none) → enable_thinking=false
+ {
+ name: "93",
+ from: "openai",
+ to: "iflow",
+ model: "glm-test(none)",
+ inputJSON: `{"model":"glm-test(none)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "chat_template_kwargs.enable_thinking",
+ expectValue: "false",
+ expectErr: false,
+ },
+ // Case 94: Claude to iflow, no suffix → passthrough
+ {
+ name: "94",
+ from: "claude",
+ to: "iflow",
+ model: "glm-test",
+ inputJSON: `{"model":"glm-test","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 95: Claude to iflow, (8192) → enable_thinking=true
+ {
+ name: "95",
+ from: "claude",
+ to: "iflow",
+ model: "glm-test(8192)",
+ inputJSON: `{"model":"glm-test(8192)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "chat_template_kwargs.enable_thinking",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 96: Claude to iflow, (-1) → enable_thinking=true
+ {
+ name: "96",
+ from: "claude",
+ to: "iflow",
+ model: "glm-test(-1)",
+ inputJSON: `{"model":"glm-test(-1)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "chat_template_kwargs.enable_thinking",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 97: Claude to iflow, (0) → enable_thinking=false
+ {
+ name: "97",
+ from: "claude",
+ to: "iflow",
+ model: "glm-test(0)",
+ inputJSON: `{"model":"glm-test(0)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "chat_template_kwargs.enable_thinking",
+ expectValue: "false",
+ expectErr: false,
+ },
+
+ // minimax-test (from: openai, gemini)
+ // Case 98: OpenAI to iflow, no suffix → passthrough
+ {
+ name: "98",
+ from: "openai",
+ to: "iflow",
+ model: "minimax-test",
+ inputJSON: `{"model":"minimax-test","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 99: OpenAI to iflow, (medium) → reasoning_split=true
+ {
+ name: "99",
+ from: "openai",
+ to: "iflow",
+ model: "minimax-test(medium)",
+ inputJSON: `{"model":"minimax-test(medium)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning_split",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 100: OpenAI to iflow, (auto) → reasoning_split=true
+ {
+ name: "100",
+ from: "openai",
+ to: "iflow",
+ model: "minimax-test(auto)",
+ inputJSON: `{"model":"minimax-test(auto)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning_split",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 101: OpenAI to iflow, (none) → reasoning_split=false
+ {
+ name: "101",
+ from: "openai",
+ to: "iflow",
+ model: "minimax-test(none)",
+ inputJSON: `{"model":"minimax-test(none)","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning_split",
+ expectValue: "false",
+ expectErr: false,
+ },
+ // Case 102: Gemini to iflow, no suffix → passthrough
+ {
+ name: "102",
+ from: "gemini",
+ to: "iflow",
+ model: "minimax-test",
+ inputJSON: `{"model":"minimax-test","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 103: Gemini to iflow, (8192) → reasoning_split=true
+ {
+ name: "103",
+ from: "gemini",
+ to: "iflow",
+ model: "minimax-test(8192)",
+ inputJSON: `{"model":"minimax-test(8192)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning_split",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 104: Gemini to iflow, (-1) → reasoning_split=true
+ {
+ name: "104",
+ from: "gemini",
+ to: "iflow",
+ model: "minimax-test(-1)",
+ inputJSON: `{"model":"minimax-test(-1)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning_split",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 105: Gemini to iflow, (0) → reasoning_split=false
+ {
+ name: "105",
+ from: "gemini",
+ to: "iflow",
+ model: "minimax-test(0)",
+ inputJSON: `{"model":"minimax-test(0)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning_split",
+ expectValue: "false",
+ expectErr: false,
+ },
+
+ // Gemini Family Cross-Channel Consistency (Cases 106-114)
+ // Tests that gemini/gemini-cli/antigravity as same API family should have consistent validation behavior
+
+ // Case 106: Gemini to Antigravity, budget 64000 (suffix) → clamped to Max
+ {
+ name: "106",
+ from: "gemini",
+ to: "antigravity",
+ model: "gemini-budget-model(64000)",
+ inputJSON: `{"model":"gemini-budget-model(64000)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 107: Gemini to Gemini-CLI, budget 64000 (suffix) → clamped to Max
+ {
+ name: "107",
+ from: "gemini",
+ to: "gemini-cli",
+ model: "gemini-budget-model(64000)",
+ inputJSON: `{"model":"gemini-budget-model(64000)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 108: Gemini-CLI to Antigravity, budget 64000 (suffix) → clamped to Max
+ {
+ name: "108",
+ from: "gemini-cli",
+ to: "antigravity",
+ model: "gemini-budget-model(64000)",
+ inputJSON: `{"model":"gemini-budget-model(64000)","request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 109: Gemini-CLI to Gemini, budget 64000 (suffix) → clamped to Max
+ {
+ name: "109",
+ from: "gemini-cli",
+ to: "gemini",
+ model: "gemini-budget-model(64000)",
+ inputJSON: `{"model":"gemini-budget-model(64000)","request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 110: Gemini to Antigravity, budget 8192 → passthrough (normal value)
+ {
+ name: "110",
+ from: "gemini",
+ to: "antigravity",
+ model: "gemini-budget-model(8192)",
+ inputJSON: `{"model":"gemini-budget-model(8192)","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 111: Gemini-CLI to Antigravity, budget 8192 → passthrough (normal value)
+ {
+ name: "111",
+ from: "gemini-cli",
+ to: "antigravity",
+ model: "gemini-budget-model(8192)",
+ inputJSON: `{"model":"gemini-budget-model(8192)","request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ }
+
+ runThinkingTests(t, cases)
+}
+
+// TestThinkingE2EMatrix_Body tests the thinking configuration transformation using request body parameters.
+// Data flow: Input JSON with thinking params → TranslateRequest → ApplyThinking → Validate Output
+func TestThinkingE2EMatrix_Body(t *testing.T) {
+ reg := registry.GetGlobalRegistry()
+ uid := fmt.Sprintf("thinking-e2e-body-%d", time.Now().UnixNano())
+
+ reg.RegisterClient(uid, "test", getTestModels())
+ defer reg.UnregisterClient(uid)
+
+ cases := []thinkingTestCase{
+ // level-model (Levels=minimal/low/medium/high, ZeroAllowed=false, DynamicAllowed=false)
+
+ // Case 1: No param → injected default → medium
+ {
+ name: "1",
+ from: "openai",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 2: reasoning_effort=medium → medium
+ {
+ name: "2",
+ from: "openai",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"medium"}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 3: reasoning_effort=xhigh → out of range error
+ {
+ name: "3",
+ from: "openai",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"xhigh"}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 4: reasoning_effort=none → clamped to minimal
+ {
+ name: "4",
+ from: "openai",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"none"}`,
+ expectField: "reasoning.effort",
+ expectValue: "minimal",
+ expectErr: false,
+ },
+ // Case 5: reasoning_effort=auto → medium (DynamicAllowed=false)
+ {
+ name: "5",
+ from: "openai",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"auto"}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 6: No param from gemini → injected default → medium
+ {
+ name: "6",
+ from: "gemini",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 7: thinkingBudget=8192 → medium
+ {
+ name: "7",
+ from: "gemini",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 8: thinkingBudget=64000 → clamped to high
+ {
+ name: "8",
+ from: "gemini",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":64000}}}`,
+ expectField: "reasoning.effort",
+ expectValue: "high",
+ expectErr: false,
+ },
+ // Case 9: thinkingBudget=0 → clamped to minimal
+ {
+ name: "9",
+ from: "gemini",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":0}}}`,
+ expectField: "reasoning.effort",
+ expectValue: "minimal",
+ expectErr: false,
+ },
+ // Case 10: thinkingBudget=-1 → medium (DynamicAllowed=false)
+ {
+ name: "10",
+ from: "gemini",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":-1}}}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 11: Claude no param → passthrough (no thinking)
+ {
+ name: "11",
+ from: "claude",
+ to: "openai",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 12: thinking.budget_tokens=8192 → medium
+ {
+ name: "12",
+ from: "claude",
+ to: "openai",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":8192}}`,
+ expectField: "reasoning_effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 13: thinking.budget_tokens=64000 → clamped to high
+ {
+ name: "13",
+ from: "claude",
+ to: "openai",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":64000}}`,
+ expectField: "reasoning_effort",
+ expectValue: "high",
+ expectErr: false,
+ },
+ // Case 14: thinking.budget_tokens=0 → clamped to minimal
+ {
+ name: "14",
+ from: "claude",
+ to: "openai",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":0}}`,
+ expectField: "reasoning_effort",
+ expectValue: "minimal",
+ expectErr: false,
+ },
+ // Case 15: thinking.budget_tokens=-1 → medium (DynamicAllowed=false)
+ {
+ name: "15",
+ from: "claude",
+ to: "openai",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":-1}}`,
+ expectField: "reasoning_effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+
+ // level-subset-model (Levels=low/high, ZeroAllowed=false, DynamicAllowed=false)
+
+ // Case 16: thinkingBudget=8192 → medium → rounded down to low
+ {
+ name: "16",
+ from: "gemini",
+ to: "openai",
+ model: "level-subset-model",
+ inputJSON: `{"model":"level-subset-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}`,
+ expectField: "reasoning_effort",
+ expectValue: "low",
+ expectErr: false,
+ },
+ // Case 17: thinking.budget_tokens=1 → minimal → clamped to low
+ {
+ name: "17",
+ from: "claude",
+ to: "gemini",
+ model: "level-subset-model",
+ inputJSON: `{"model":"level-subset-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":1}}`,
+ expectField: "generationConfig.thinkingConfig.thinkingLevel",
+ expectValue: "low",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+
+ // gemini-budget-model (Min=128, Max=20000, ZeroAllowed=false, DynamicAllowed=true)
+
+ // Case 18: No param → passthrough
+ {
+ name: "18",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 19: reasoning_effort=medium → 8192
+ {
+ name: "19",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"medium"}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 20: reasoning_effort=xhigh → clamped to 20000
+ {
+ name: "20",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"xhigh"}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 21: reasoning_effort=none → clamped to 128 → includeThoughts=false
+ {
+ name: "21",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"none"}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "128",
+ includeThoughts: "false",
+ expectErr: false,
+ },
+ // Case 22: reasoning_effort=auto → -1 (DynamicAllowed=true)
+ {
+ name: "22",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"auto"}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "-1",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 23: Claude no param → passthrough
+ {
+ name: "23",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 24: thinking.budget_tokens=8192 → 8192
+ {
+ name: "24",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":8192}}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 25: thinking.budget_tokens=64000 → clamped to 20000
+ {
+ name: "25",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":64000}}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 26: thinking.budget_tokens=0 → clamped to 128 → includeThoughts=false
+ {
+ name: "26",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":0}}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "128",
+ includeThoughts: "false",
+ expectErr: false,
+ },
+ // Case 27: thinking.budget_tokens=-1 → -1 (DynamicAllowed=true)
+ {
+ name: "27",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":-1}}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "-1",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+
+ // gemini-mixed-model (Min=128, Max=32768, Levels=low/high, ZeroAllowed=false, DynamicAllowed=true)
+
+ // Case 28: No param → passthrough
+ {
+ name: "28",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-mixed-model",
+ inputJSON: `{"model":"gemini-mixed-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 29: reasoning_effort=high → high
+ {
+ name: "29",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-mixed-model",
+ inputJSON: `{"model":"gemini-mixed-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"high"}`,
+ expectField: "generationConfig.thinkingConfig.thinkingLevel",
+ expectValue: "high",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 30: reasoning_effort=xhigh → error (not in low/high)
+ {
+ name: "30",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-mixed-model",
+ inputJSON: `{"model":"gemini-mixed-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"xhigh"}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 31: reasoning_effort=none → clamped to low → includeThoughts=false
+ {
+ name: "31",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-mixed-model",
+ inputJSON: `{"model":"gemini-mixed-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"none"}`,
+ expectField: "generationConfig.thinkingConfig.thinkingLevel",
+ expectValue: "low",
+ includeThoughts: "false",
+ expectErr: false,
+ },
+ // Case 32: reasoning_effort=auto → -1 (DynamicAllowed=true)
+ {
+ name: "32",
+ from: "openai",
+ to: "gemini",
+ model: "gemini-mixed-model",
+ inputJSON: `{"model":"gemini-mixed-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"auto"}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "-1",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 33: Claude no param → passthrough
+ {
+ name: "33",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-mixed-model",
+ inputJSON: `{"model":"gemini-mixed-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 34: thinking.budget_tokens=8192 → 8192 (keeps budget)
+ {
+ name: "34",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-mixed-model",
+ inputJSON: `{"model":"gemini-mixed-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":8192}}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 35: thinking.budget_tokens=64000 → clamped to 32768 (keeps budget)
+ {
+ name: "35",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-mixed-model",
+ inputJSON: `{"model":"gemini-mixed-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":64000}}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "32768",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 36: thinking.budget_tokens=0 → clamped to low → includeThoughts=false
+ {
+ name: "36",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-mixed-model",
+ inputJSON: `{"model":"gemini-mixed-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":0}}`,
+ expectField: "generationConfig.thinkingConfig.thinkingLevel",
+ expectValue: "low",
+ includeThoughts: "false",
+ expectErr: false,
+ },
+ // Case 37: thinking.budget_tokens=-1 → -1 (DynamicAllowed=true)
+ {
+ name: "37",
+ from: "claude",
+ to: "gemini",
+ model: "gemini-mixed-model",
+ inputJSON: `{"model":"gemini-mixed-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":-1}}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "-1",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+
+ // claude-budget-model (Min=1024, Max=128000, ZeroAllowed=true, DynamicAllowed=false)
+
+ // Case 38: No param → passthrough
+ {
+ name: "38",
+ from: "openai",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 39: reasoning_effort=medium → 8192
+ {
+ name: "39",
+ from: "openai",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"medium"}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "8192",
+ expectErr: false,
+ },
+ // Case 40: reasoning_effort=xhigh → clamped to 32768
+ {
+ name: "40",
+ from: "openai",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"xhigh"}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "32768",
+ expectErr: false,
+ },
+ // Case 41: reasoning_effort=none → disabled
+ {
+ name: "41",
+ from: "openai",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"none"}`,
+ expectField: "thinking.type",
+ expectValue: "disabled",
+ expectErr: false,
+ },
+ // Case 42: reasoning_effort=auto → 64512 (mid-range)
+ {
+ name: "42",
+ from: "openai",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"auto"}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "64512",
+ expectErr: false,
+ },
+ // Case 43: Gemini no param → passthrough
+ {
+ name: "43",
+ from: "gemini",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 44: thinkingBudget=8192 → 8192
+ {
+ name: "44",
+ from: "gemini",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "8192",
+ expectErr: false,
+ },
+ // Case 45: thinkingBudget=200000 → clamped to 128000
+ {
+ name: "45",
+ from: "gemini",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":200000}}}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "128000",
+ expectErr: false,
+ },
+ // Case 46: thinkingBudget=0 → disabled
+ {
+ name: "46",
+ from: "gemini",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":0}}}`,
+ expectField: "thinking.type",
+ expectValue: "disabled",
+ expectErr: false,
+ },
+ // Case 47: thinkingBudget=-1 → 64512 (mid-range)
+ {
+ name: "47",
+ from: "gemini",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":-1}}}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "64512",
+ expectErr: false,
+ },
+
+ // antigravity-budget-model (Min=128, Max=20000, ZeroAllowed=true, DynamicAllowed=true)
+
+ // Case 48: Gemini no param → passthrough
+ {
+ name: "48",
+ from: "gemini",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 49: thinkingLevel=medium → 8192
+ {
+ name: "49",
+ from: "gemini",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingLevel":"medium"}}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 50: thinkingLevel=xhigh → clamped to 20000
+ {
+ name: "50",
+ from: "gemini",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingLevel":"xhigh"}}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 51: thinkingLevel=none → 0 (ZeroAllowed=true)
+ {
+ name: "51",
+ from: "gemini",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingLevel":"none"}}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "0",
+ includeThoughts: "false",
+ expectErr: false,
+ },
+ // Case 52: thinkingBudget=-1 → -1 (DynamicAllowed=true)
+ {
+ name: "52",
+ from: "gemini",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":-1}}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "-1",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 53: Claude no param → passthrough
+ {
+ name: "53",
+ from: "claude",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 54: thinking.budget_tokens=8192 → 8192
+ {
+ name: "54",
+ from: "claude",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":8192}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 55: thinking.budget_tokens=64000 → clamped to 20000
+ {
+ name: "55",
+ from: "claude",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":64000}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "20000",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 56: thinking.budget_tokens=0 → 0 (ZeroAllowed=true)
+ {
+ name: "56",
+ from: "claude",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":0}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "0",
+ includeThoughts: "false",
+ expectErr: false,
+ },
+ // Case 57: thinking.budget_tokens=-1 → -1 (DynamicAllowed=true)
+ {
+ name: "57",
+ from: "claude",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":-1}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "-1",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+
+ // no-thinking-model (Thinking=nil)
+
+ // Case 58: Gemini no param → passthrough
+ {
+ name: "58",
+ from: "gemini",
+ to: "openai",
+ model: "no-thinking-model",
+ inputJSON: `{"model":"no-thinking-model","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 59: thinkingBudget=8192 → stripped
+ {
+ name: "59",
+ from: "gemini",
+ to: "openai",
+ model: "no-thinking-model",
+ inputJSON: `{"model":"no-thinking-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 60: thinkingBudget=0 → stripped
+ {
+ name: "60",
+ from: "gemini",
+ to: "openai",
+ model: "no-thinking-model",
+ inputJSON: `{"model":"no-thinking-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":0}}}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 61: thinkingBudget=-1 → stripped
+ {
+ name: "61",
+ from: "gemini",
+ to: "openai",
+ model: "no-thinking-model",
+ inputJSON: `{"model":"no-thinking-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":-1}}}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 62: Claude no param → passthrough
+ {
+ name: "62",
+ from: "claude",
+ to: "openai",
+ model: "no-thinking-model",
+ inputJSON: `{"model":"no-thinking-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 63: thinking.budget_tokens=8192 → stripped
+ {
+ name: "63",
+ from: "claude",
+ to: "openai",
+ model: "no-thinking-model",
+ inputJSON: `{"model":"no-thinking-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":8192}}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 64: thinking.budget_tokens=0 → stripped
+ {
+ name: "64",
+ from: "claude",
+ to: "openai",
+ model: "no-thinking-model",
+ inputJSON: `{"model":"no-thinking-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":0}}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 65: thinking.budget_tokens=-1 → stripped
+ {
+ name: "65",
+ from: "claude",
+ to: "openai",
+ model: "no-thinking-model",
+ inputJSON: `{"model":"no-thinking-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":-1}}`,
+ expectField: "",
+ expectErr: false,
+ },
+
+ // user-defined-model (UserDefined=true, Thinking=nil)
+
+ // Case 66: Gemini no param → passthrough
+ {
+ name: "66",
+ from: "gemini",
+ to: "openai",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 67: thinkingBudget=8192 → medium
+ {
+ name: "67",
+ from: "gemini",
+ to: "openai",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}`,
+ expectField: "reasoning_effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 68: thinkingBudget=64000 → xhigh (passthrough)
+ {
+ name: "68",
+ from: "gemini",
+ to: "openai",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":64000}}}`,
+ expectField: "reasoning_effort",
+ expectValue: "xhigh",
+ expectErr: false,
+ },
+ // Case 69: thinkingBudget=0 → none
+ {
+ name: "69",
+ from: "gemini",
+ to: "openai",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":0}}}`,
+ expectField: "reasoning_effort",
+ expectValue: "none",
+ expectErr: false,
+ },
+ // Case 70: thinkingBudget=-1 → auto
+ {
+ name: "70",
+ from: "gemini",
+ to: "openai",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":-1}}}`,
+ expectField: "reasoning_effort",
+ expectValue: "auto",
+ expectErr: false,
+ },
+ // Case 71: Claude no param → injected default → medium
+ {
+ name: "71",
+ from: "claude",
+ to: "codex",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 72: thinking.budget_tokens=8192 → medium
+ {
+ name: "72",
+ from: "claude",
+ to: "codex",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":8192}}`,
+ expectField: "reasoning.effort",
+ expectValue: "medium",
+ expectErr: false,
+ },
+ // Case 73: thinking.budget_tokens=64000 → xhigh (passthrough)
+ {
+ name: "73",
+ from: "claude",
+ to: "codex",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":64000}}`,
+ expectField: "reasoning.effort",
+ expectValue: "xhigh",
+ expectErr: false,
+ },
+ // Case 74: thinking.budget_tokens=0 → none
+ {
+ name: "74",
+ from: "claude",
+ to: "codex",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":0}}`,
+ expectField: "reasoning.effort",
+ expectValue: "none",
+ expectErr: false,
+ },
+ // Case 75: thinking.budget_tokens=-1 → auto
+ {
+ name: "75",
+ from: "claude",
+ to: "codex",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":-1}}`,
+ expectField: "reasoning.effort",
+ expectValue: "auto",
+ expectErr: false,
+ },
+ // Case 76: OpenAI reasoning_effort=medium to Gemini → 8192
+ {
+ name: "76",
+ from: "openai",
+ to: "gemini",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"medium"}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 77: OpenAI reasoning_effort=medium to Claude → 8192
+ {
+ name: "77",
+ from: "openai",
+ to: "claude",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"medium"}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "8192",
+ expectErr: false,
+ },
+ // Case 78: OpenAI-Response reasoning.effort=medium to Gemini → 8192
+ {
+ name: "78",
+ from: "openai-response",
+ to: "gemini",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","input":[{"role":"user","content":"hi"}],"reasoning":{"effort":"medium"}}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 79: OpenAI-Response reasoning.effort=medium to Claude → 8192
+ {
+ name: "79",
+ from: "openai-response",
+ to: "claude",
+ model: "user-defined-model",
+ inputJSON: `{"model":"user-defined-model","input":[{"role":"user","content":"hi"}],"reasoning":{"effort":"medium"}}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "8192",
+ expectErr: false,
+ },
+
+ // Same-protocol passthrough tests (80-89)
+
+ // Case 80: OpenAI to OpenAI, reasoning_effort=high → passthrough
+ {
+ name: "80",
+ from: "openai",
+ to: "openai",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"high"}`,
+ expectField: "reasoning_effort",
+ expectValue: "high",
+ expectErr: false,
+ },
+ // Case 81: OpenAI to OpenAI, reasoning_effort=xhigh → out of range error
+ {
+ name: "81",
+ from: "openai",
+ to: "openai",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"xhigh"}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 82: OpenAI-Response to Codex, reasoning.effort=high → passthrough
+ {
+ name: "82",
+ from: "openai-response",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","input":[{"role":"user","content":"hi"}],"reasoning":{"effort":"high"}}`,
+ expectField: "reasoning.effort",
+ expectValue: "high",
+ expectErr: false,
+ },
+ // Case 83: OpenAI-Response to Codex, reasoning.effort=xhigh → out of range error
+ {
+ name: "83",
+ from: "openai-response",
+ to: "codex",
+ model: "level-model",
+ inputJSON: `{"model":"level-model","input":[{"role":"user","content":"hi"}],"reasoning":{"effort":"xhigh"}}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 84: Gemini to Gemini, thinkingBudget=8192 → passthrough
+ {
+ name: "84",
+ from: "gemini",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}`,
+ expectField: "generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 85: Gemini to Gemini, thinkingBudget=64000 → exceeds Max error
+ {
+ name: "85",
+ from: "gemini",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":64000}}}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 86: Claude to Claude, thinking.budget_tokens=8192 → passthrough
+ {
+ name: "86",
+ from: "claude",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":8192}}`,
+ expectField: "thinking.budget_tokens",
+ expectValue: "8192",
+ expectErr: false,
+ },
+ // Case 87: Claude to Claude, thinking.budget_tokens=200000 → exceeds Max error
+ {
+ name: "87",
+ from: "claude",
+ to: "claude",
+ model: "claude-budget-model",
+ inputJSON: `{"model":"claude-budget-model","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":200000}}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 88: Gemini-CLI to Antigravity, thinkingBudget=8192 → passthrough
+ {
+ name: "88",
+ from: "gemini-cli",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 89: Gemini-CLI to Antigravity, thinkingBudget=64000 → exceeds Max error
+ {
+ name: "89",
+ from: "gemini-cli",
+ to: "antigravity",
+ model: "antigravity-budget-model",
+ inputJSON: `{"model":"antigravity-budget-model","request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":64000}}}}`,
+ expectField: "",
+ expectErr: true,
+ },
+
+ // iflow tests: glm-test and minimax-test (Cases 90-105)
+
+ // glm-test (from: openai, claude)
+ // Case 90: OpenAI to iflow, no param → passthrough
+ {
+ name: "90",
+ from: "openai",
+ to: "iflow",
+ model: "glm-test",
+ inputJSON: `{"model":"glm-test","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 91: OpenAI to iflow, reasoning_effort=medium → enable_thinking=true
+ {
+ name: "91",
+ from: "openai",
+ to: "iflow",
+ model: "glm-test",
+ inputJSON: `{"model":"glm-test","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"medium"}`,
+ expectField: "chat_template_kwargs.enable_thinking",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 92: OpenAI to iflow, reasoning_effort=auto → enable_thinking=true
+ {
+ name: "92",
+ from: "openai",
+ to: "iflow",
+ model: "glm-test",
+ inputJSON: `{"model":"glm-test","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"auto"}`,
+ expectField: "chat_template_kwargs.enable_thinking",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 93: OpenAI to iflow, reasoning_effort=none → enable_thinking=false
+ {
+ name: "93",
+ from: "openai",
+ to: "iflow",
+ model: "glm-test",
+ inputJSON: `{"model":"glm-test","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"none"}`,
+ expectField: "chat_template_kwargs.enable_thinking",
+ expectValue: "false",
+ expectErr: false,
+ },
+ // Case 94: Claude to iflow, no param → passthrough
+ {
+ name: "94",
+ from: "claude",
+ to: "iflow",
+ model: "glm-test",
+ inputJSON: `{"model":"glm-test","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 95: Claude to iflow, thinking.budget_tokens=8192 → enable_thinking=true
+ {
+ name: "95",
+ from: "claude",
+ to: "iflow",
+ model: "glm-test",
+ inputJSON: `{"model":"glm-test","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":8192}}`,
+ expectField: "chat_template_kwargs.enable_thinking",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 96: Claude to iflow, thinking.budget_tokens=-1 → enable_thinking=true
+ {
+ name: "96",
+ from: "claude",
+ to: "iflow",
+ model: "glm-test",
+ inputJSON: `{"model":"glm-test","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":-1}}`,
+ expectField: "chat_template_kwargs.enable_thinking",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 97: Claude to iflow, thinking.budget_tokens=0 → enable_thinking=false
+ {
+ name: "97",
+ from: "claude",
+ to: "iflow",
+ model: "glm-test",
+ inputJSON: `{"model":"glm-test","messages":[{"role":"user","content":"hi"}],"thinking":{"type":"enabled","budget_tokens":0}}`,
+ expectField: "chat_template_kwargs.enable_thinking",
+ expectValue: "false",
+ expectErr: false,
+ },
+
+ // minimax-test (from: openai, gemini)
+ // Case 98: OpenAI to iflow, no param → passthrough
+ {
+ name: "98",
+ from: "openai",
+ to: "iflow",
+ model: "minimax-test",
+ inputJSON: `{"model":"minimax-test","messages":[{"role":"user","content":"hi"}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 99: OpenAI to iflow, reasoning_effort=medium → reasoning_split=true
+ {
+ name: "99",
+ from: "openai",
+ to: "iflow",
+ model: "minimax-test",
+ inputJSON: `{"model":"minimax-test","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"medium"}`,
+ expectField: "reasoning_split",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 100: OpenAI to iflow, reasoning_effort=auto → reasoning_split=true
+ {
+ name: "100",
+ from: "openai",
+ to: "iflow",
+ model: "minimax-test",
+ inputJSON: `{"model":"minimax-test","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"auto"}`,
+ expectField: "reasoning_split",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 101: OpenAI to iflow, reasoning_effort=none → reasoning_split=false
+ {
+ name: "101",
+ from: "openai",
+ to: "iflow",
+ model: "minimax-test",
+ inputJSON: `{"model":"minimax-test","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"none"}`,
+ expectField: "reasoning_split",
+ expectValue: "false",
+ expectErr: false,
+ },
+ // Case 102: Gemini to iflow, no param → passthrough
+ {
+ name: "102",
+ from: "gemini",
+ to: "iflow",
+ model: "minimax-test",
+ inputJSON: `{"model":"minimax-test","contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
+ expectField: "",
+ expectErr: false,
+ },
+ // Case 103: Gemini to iflow, thinkingBudget=8192 → reasoning_split=true
+ {
+ name: "103",
+ from: "gemini",
+ to: "iflow",
+ model: "minimax-test",
+ inputJSON: `{"model":"minimax-test","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}`,
+ expectField: "reasoning_split",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 104: Gemini to iflow, thinkingBudget=-1 → reasoning_split=true
+ {
+ name: "104",
+ from: "gemini",
+ to: "iflow",
+ model: "minimax-test",
+ inputJSON: `{"model":"minimax-test","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":-1}}}`,
+ expectField: "reasoning_split",
+ expectValue: "true",
+ expectErr: false,
+ },
+ // Case 105: Gemini to iflow, thinkingBudget=0 → reasoning_split=false
+ {
+ name: "105",
+ from: "gemini",
+ to: "iflow",
+ model: "minimax-test",
+ inputJSON: `{"model":"minimax-test","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":0}}}`,
+ expectField: "reasoning_split",
+ expectValue: "false",
+ expectErr: false,
+ },
+
+ // Gemini Family Cross-Channel Consistency (Cases 106-114)
+ // Tests that gemini/gemini-cli/antigravity as same API family should have consistent validation behavior
+
+ // Case 106: Gemini to Antigravity, thinkingBudget=64000 → exceeds Max error (same family strict validation)
+ {
+ name: "106",
+ from: "gemini",
+ to: "antigravity",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":64000}}}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 107: Gemini to Gemini-CLI, thinkingBudget=64000 → exceeds Max error (same family strict validation)
+ {
+ name: "107",
+ from: "gemini",
+ to: "gemini-cli",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":64000}}}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 108: Gemini-CLI to Antigravity, thinkingBudget=64000 → exceeds Max error (same family strict validation)
+ {
+ name: "108",
+ from: "gemini-cli",
+ to: "antigravity",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":64000}}}}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 109: Gemini-CLI to Gemini, thinkingBudget=64000 → exceeds Max error (same family strict validation)
+ {
+ name: "109",
+ from: "gemini-cli",
+ to: "gemini",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":64000}}}}`,
+ expectField: "",
+ expectErr: true,
+ },
+ // Case 110: Gemini to Antigravity, thinkingBudget=8192 → passthrough (normal value)
+ {
+ name: "110",
+ from: "gemini",
+ to: "antigravity",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ // Case 111: Gemini-CLI to Antigravity, thinkingBudget=8192 → passthrough (normal value)
+ {
+ name: "111",
+ from: "gemini-cli",
+ to: "antigravity",
+ model: "gemini-budget-model",
+ inputJSON: `{"model":"gemini-budget-model","request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}}`,
+ expectField: "request.generationConfig.thinkingConfig.thinkingBudget",
+ expectValue: "8192",
+ includeThoughts: "true",
+ expectErr: false,
+ },
+ }
+
+ runThinkingTests(t, cases)
+}
+
+// getTestModels returns the shared model definitions for E2E tests.
+func getTestModels() []*registry.ModelInfo {
+ return []*registry.ModelInfo{
+ {
+ ID: "level-model",
+ Object: "model",
+ Created: 1700000000,
+ OwnedBy: "test",
+ Type: "openai",
+ DisplayName: "Level Model",
+ Thinking: ®istry.ThinkingSupport{Levels: []string{"minimal", "low", "medium", "high"}, ZeroAllowed: false, DynamicAllowed: false},
+ },
+ {
+ ID: "level-subset-model",
+ Object: "model",
+ Created: 1700000000,
+ OwnedBy: "test",
+ Type: "gemini",
+ DisplayName: "Level Subset Model",
+ Thinking: ®istry.ThinkingSupport{Levels: []string{"low", "high"}, ZeroAllowed: false, DynamicAllowed: false},
+ },
+ {
+ ID: "gemini-budget-model",
+ Object: "model",
+ Created: 1700000000,
+ OwnedBy: "test",
+ Type: "gemini",
+ DisplayName: "Gemini Budget Model",
+ Thinking: ®istry.ThinkingSupport{Min: 128, Max: 20000, ZeroAllowed: false, DynamicAllowed: true},
+ },
+ {
+ ID: "gemini-mixed-model",
+ Object: "model",
+ Created: 1700000000,
+ OwnedBy: "test",
+ Type: "gemini",
+ DisplayName: "Gemini Mixed Model",
+ Thinking: ®istry.ThinkingSupport{Min: 128, Max: 32768, Levels: []string{"low", "high"}, ZeroAllowed: false, DynamicAllowed: true},
+ },
+ {
+ ID: "claude-budget-model",
+ Object: "model",
+ Created: 1700000000,
+ OwnedBy: "test",
+ Type: "claude",
+ DisplayName: "Claude Budget Model",
+ Thinking: ®istry.ThinkingSupport{Min: 1024, Max: 128000, ZeroAllowed: true, DynamicAllowed: false},
+ },
+ {
+ ID: "antigravity-budget-model",
+ Object: "model",
+ Created: 1700000000,
+ OwnedBy: "test",
+ Type: "gemini-cli",
+ DisplayName: "Antigravity Budget Model",
+ Thinking: ®istry.ThinkingSupport{Min: 128, Max: 20000, ZeroAllowed: true, DynamicAllowed: true},
+ },
+ {
+ ID: "no-thinking-model",
+ Object: "model",
+ Created: 1700000000,
+ OwnedBy: "test",
+ Type: "openai",
+ DisplayName: "No Thinking Model",
+ Thinking: nil,
+ },
+ {
+ ID: "user-defined-model",
+ Object: "model",
+ Created: 1700000000,
+ OwnedBy: "test",
+ Type: "openai",
+ DisplayName: "User Defined Model",
+ UserDefined: true,
+ Thinking: nil,
+ },
+ {
+ ID: "glm-test",
+ Object: "model",
+ Created: 1700000000,
+ OwnedBy: "test",
+ Type: "iflow",
+ DisplayName: "GLM Test Model",
+ Thinking: ®istry.ThinkingSupport{Levels: []string{"none", "auto", "minimal", "low", "medium", "high", "xhigh"}},
+ },
+ {
+ ID: "minimax-test",
+ Object: "model",
+ Created: 1700000000,
+ OwnedBy: "test",
+ Type: "iflow",
+ DisplayName: "MiniMax Test Model",
+ Thinking: ®istry.ThinkingSupport{Levels: []string{"none", "auto", "minimal", "low", "medium", "high", "xhigh"}},
+ },
+ }
+}
+
+// runThinkingTests runs thinking test cases using the real data flow path.
+func runThinkingTests(t *testing.T, cases []thinkingTestCase) {
+ for _, tc := range cases {
+ tc := tc
+ testName := fmt.Sprintf("Case%s_%s->%s_%s", tc.name, tc.from, tc.to, tc.model)
+ t.Run(testName, func(t *testing.T) {
+ suffixResult := thinking.ParseSuffix(tc.model)
+ baseModel := suffixResult.ModelName
+
+ translateTo := tc.to
+ applyTo := tc.to
+ if tc.to == "iflow" {
+ translateTo = "openai"
+ applyTo = "iflow"
+ }
+
+ body := sdktranslator.TranslateRequest(
+ sdktranslator.FromString(tc.from),
+ sdktranslator.FromString(translateTo),
+ baseModel,
+ []byte(tc.inputJSON),
+ true,
+ )
+ if applyTo == "claude" {
+ body, _ = sjson.SetBytes(body, "max_tokens", 200000)
+ }
+
+ body, err := thinking.ApplyThinking(body, tc.model, tc.from, applyTo, applyTo)
+
+ if tc.expectErr {
+ if err == nil {
+ t.Fatalf("expected error but got none, body=%s", string(body))
+ }
+ return
+ }
+ if err != nil {
+ t.Fatalf("unexpected error: %v, body=%s", err, string(body))
+ }
+
+ if tc.expectField == "" {
+ var hasThinking bool
+ switch tc.to {
+ case "gemini":
+ hasThinking = gjson.GetBytes(body, "generationConfig.thinkingConfig").Exists()
+ case "gemini-cli":
+ hasThinking = gjson.GetBytes(body, "request.generationConfig.thinkingConfig").Exists()
+ case "antigravity":
+ hasThinking = gjson.GetBytes(body, "request.generationConfig.thinkingConfig").Exists()
+ case "claude":
+ hasThinking = gjson.GetBytes(body, "thinking").Exists()
+ case "openai":
+ hasThinking = gjson.GetBytes(body, "reasoning_effort").Exists()
+ case "codex":
+ hasThinking = gjson.GetBytes(body, "reasoning.effort").Exists() || gjson.GetBytes(body, "reasoning").Exists()
+ case "iflow":
+ hasThinking = gjson.GetBytes(body, "chat_template_kwargs.enable_thinking").Exists() || gjson.GetBytes(body, "reasoning_split").Exists()
+ }
+ if hasThinking {
+ t.Fatalf("expected no thinking field but found one, body=%s", string(body))
+ }
+ return
+ }
+
+ val := gjson.GetBytes(body, tc.expectField)
+ if !val.Exists() {
+ t.Fatalf("expected field %s not found, body=%s", tc.expectField, string(body))
+ }
+
+ actualValue := val.String()
+ if val.Type == gjson.Number {
+ actualValue = fmt.Sprintf("%d", val.Int())
+ }
+ if actualValue != tc.expectValue {
+ t.Fatalf("field %s: expected %q, got %q, body=%s", tc.expectField, tc.expectValue, actualValue, string(body))
+ }
+
+ if tc.includeThoughts != "" && (tc.to == "gemini" || tc.to == "gemini-cli" || tc.to == "antigravity") {
+ path := "generationConfig.thinkingConfig.includeThoughts"
+ if tc.to == "gemini-cli" || tc.to == "antigravity" {
+ path = "request.generationConfig.thinkingConfig.includeThoughts"
+ }
+ itVal := gjson.GetBytes(body, path)
+ if !itVal.Exists() {
+ t.Fatalf("expected includeThoughts field not found, body=%s", string(body))
+ }
+ actual := fmt.Sprintf("%v", itVal.Bool())
+ if actual != tc.includeThoughts {
+ t.Fatalf("includeThoughts: expected %s, got %s, body=%s", tc.includeThoughts, actual, string(body))
+ }
+ }
+
+ // Verify clear_thinking for iFlow GLM models when enable_thinking=true
+ if tc.to == "iflow" && tc.expectField == "chat_template_kwargs.enable_thinking" && tc.expectValue == "true" {
+ baseModel := thinking.ParseSuffix(tc.model).ModelName
+ isGLM := strings.HasPrefix(strings.ToLower(baseModel), "glm")
+ ctVal := gjson.GetBytes(body, "chat_template_kwargs.clear_thinking")
+ if isGLM {
+ if !ctVal.Exists() {
+ t.Fatalf("expected clear_thinking field not found for GLM model, body=%s", string(body))
+ }
+ if ctVal.Bool() != false {
+ t.Fatalf("clear_thinking: expected false, got %v, body=%s", ctVal.Bool(), string(body))
+ }
+ } else if ctVal.Exists() {
+ t.Fatalf("expected no clear_thinking field for non-GLM enable_thinking model, body=%s", string(body))
+ }
+ }
+ })
+ }
+}
