History

Merge branch 'main' into plus

Luis Pater committed 10 months ago

Commit e85c9d92 parents cb56cb2 5c9512975 files · +15212 −302Browse files

Unified · Split

75 files
  1. .github/workflows/docker-image.yml +1 −1
  2. .github/workflows/release.yaml +2 −1
  3. .gitignore +2 −0
  4. .goreleaser.yml +4 −4
  5. Dockerfile +3 −3
  6. README.md +10 −112
  7. README_CN.md +10 −119
  8. cmd/server/main.go +56 −0
  9. config.example.yaml +15 −0
  10. docker-compose.yml +2 −2
  11. go.mod +2 −1
  12. go.sum +3 −2
  13. internal/api/handlers/management/auth_files.go +318 −0
  14. internal/api/handlers/management/config_basic.go +2 −2
  15. internal/api/handlers/management/oauth_sessions.go +8 −1
  16. internal/api/modules/amp/proxy.go +34 −3
  17. internal/api/modules/amp/response_rewriter.go +58 −2
  18. internal/api/server.go +21 −0
  19. internal/auth/claude/oauth_server.go +11 −0
  20. internal/auth/codex/oauth_server.go +11 −0
  21. internal/auth/copilot/copilot_auth.go +225 −0
  22. internal/auth/copilot/errors.go +187 −0
  23. internal/auth/copilot/oauth.go +255 −0
  24. internal/auth/copilot/token.go +93 −0
  25. internal/auth/iflow/iflow_auth.go +17 −5
  26. internal/auth/kiro/aws.go +305 −0
  27. internal/auth/kiro/aws_auth.go +314 −0
  28. internal/auth/kiro/aws_test.go +161 −0
  29. internal/auth/kiro/codewhisperer_client.go +166 −0
  30. internal/auth/kiro/oauth.go +303 −0
  31. internal/auth/kiro/protocol_handler.go +725 −0
  32. internal/auth/kiro/social_auth.go +403 −0
  33. internal/auth/kiro/sso_oidc.go +1371 −0
  34. internal/auth/kiro/token.go +72 −0
  35. internal/browser/browser.go +416 −14
  36. internal/cmd/auth_manager.go +3 −1
  37. internal/cmd/github_copilot_login.go +44 −0
  38. internal/cmd/kiro_login.go +208 −0
  39. internal/config/config.go +62 −0
  40. internal/constant/constant.go +3 −0
  41. internal/logging/global_logger.go +1 −0
  42. internal/registry/model_definitions.go +361 −0
  43. internal/registry/model_registry.go +15 −1
  44. internal/runtime/executor/cache_helpers.go +30 −2
  45. internal/runtime/executor/codex_executor.go +2 −2
  46. internal/runtime/executor/github_copilot_executor.go +361 −0
  47. internal/runtime/executor/kiro_executor.go +3302 −0
  48. internal/runtime/executor/proxy_helpers.go +44 −5
  49. internal/runtime/executor/token_helpers.go +276 −15
  50. internal/translator/claude/openai/chat-completions/claude_openai_response.go +4 −0
  51. internal/translator/init.go +3 −0
  52. internal/translator/kiro/claude/init.go +20 −0
  53. internal/translator/kiro/claude/kiro_claude.go +21 −0
  54. internal/translator/kiro/claude/kiro_claude_request.go +809 −0
  55. internal/translator/kiro/claude/kiro_claude_response.go +204 −0
  56. internal/translator/kiro/claude/kiro_claude_stream.go +186 −0
  57. internal/translator/kiro/claude/kiro_claude_tools.go +522 −0
  58. internal/translator/kiro/common/constants.go +75 −0
  59. internal/translator/kiro/common/message_merge.go +132 −0
  60. internal/translator/kiro/common/utils.go +16 −0
  61. internal/translator/kiro/openai/init.go +20 −0
  62. internal/translator/kiro/openai/kiro_openai.go +371 −0
  63. internal/translator/kiro/openai/kiro_openai_request.go +863 −0
  64. internal/translator/kiro/openai/kiro_openai_request_test.go +386 −0
  65. internal/translator/kiro/openai/kiro_openai_response.go +277 −0
  66. internal/translator/kiro/openai/kiro_openai_stream.go +212 −0
  67. internal/watcher/events.go +67 −1
  68. internal/watcher/synthesizer/config.go +97 −0
  69. sdk/api/handlers/handlers.go +2 −1
  70. sdk/auth/github_copilot.go +129 −0
  71. sdk/auth/kiro.go +470 −0
  72. sdk/auth/manager.go +13 −0
  73. sdk/auth/refresh_registry.go +2 −0
  74. sdk/cliproxy/auth/conductor.go +4 −2
  75. sdk/cliproxy/service.go +9 −0
.github/workflows/docker-image.yml+1−1
@@ -7,7 +7,7 @@ on: · 6 unmodified lines
7 7
8env:8env:
9 APP_NAME: CLIProxyAPI9 APP_NAME: CLIProxyAPI
10 DOCKERHUB_REPO: eceasy/cli-proxy-api10 DOCKERHUB_REPO: eceasy/cli-proxy-api-plus
11 11
12jobs:12jobs:
13 docker:13 docker:
.github/workflows/release.yaml+2−1
@@ -23,7 +23,8 @@ jobs: · 22 unmodified lines
23 cache: true23 cache: true
24 - name: Generate Build Metadata24 - name: Generate Build Metadata
25 run: |25 run: |
26 echo VERSION=`git describe --tags --always --dirty` >> $GITHUB_ENV26 VERSION=$(git describe --tags --always --dirty)
27 echo "VERSION=${VERSION}" >> $GITHUB_ENV
27 echo COMMIT=`git rev-parse --short HEAD` >> $GITHUB_ENV28 echo COMMIT=`git rev-parse --short HEAD` >> $GITHUB_ENV
28 echo BUILD_DATE=`date -u +%Y-%m-%dT%H:%M:%SZ` >> $GITHUB_ENV29 echo BUILD_DATE=`date -u +%Y-%m-%dT%H:%M:%SZ` >> $GITHUB_ENV
29 - uses: goreleaser/goreleaser-action@v430 - uses: goreleaser/goreleaser-action@v4
.gitignore+2−0
@@ -1,5 +1,6 @@
1# Binaries1# Binaries
2cli-proxy-api2cli-proxy-api
3cliproxy
3*.exe4*.exe
4 5
5# Configuration6# Configuration
@@ -44,6 +45,7 @@ GEMINI.md · 38 unmodified lines
44.bmad/*45.bmad/*
45_bmad/*46_bmad/*
46_bmad-output/*47_bmad-output/*
48.mcp/cache/
47 49
48# macOS50# macOS
49.DS_Store51.DS_Store
.goreleaser.yml+4−4
@@ -1,5 +1,5 @@
1builds:1builds:
2 - id: "cli-proxy-api"2 - id: "cli-proxy-api-plus"
3 env:3 env:
4 - CGO_ENABLED=04 - CGO_ENABLED=0
5 goos:5 goos:
@@ -10,11 +10,11 @@ builds: · 4 unmodified lines
10 - amd6410 - amd64
11 - arm6411 - arm64
12 main: ./cmd/server/12 main: ./cmd/server/
13 binary: cli-proxy-api13 binary: cli-proxy-api-plus
14 ldflags:14 ldflags:
15 - -s -w -X 'main.Version={{.Version}}' -X 'main.Commit={{.ShortCommit}}' -X 'main.BuildDate={{.Date}}'15 - -s -w -X 'main.Version={{.Version}}-plus' -X 'main.Commit={{.ShortCommit}}' -X 'main.BuildDate={{.Date}}'
16archives:16archives:
17 - id: "cli-proxy-api"17 - id: "cli-proxy-api-plus"
18 format: tar.gz18 format: tar.gz
19 format_overrides:19 format_overrides:
20 - goos: windows20 - goos: windows
Dockerfile+3−3
@@ -12,7 +12,7 @@ ARG VERSION=dev · 11 unmodified lines
12ARG COMMIT=none12ARG COMMIT=none
13ARG BUILD_DATE=unknown13ARG BUILD_DATE=unknown
14 14
15RUN CGO_ENABLED=0 GOOS=linux go build -ldflags="-s -w -X 'main.Version=${VERSION}' -X 'main.Commit=${COMMIT}' -X 'main.BuildDate=${BUILD_DATE}'" -o ./CLIProxyAPI ./cmd/server/15RUN CGO_ENABLED=0 GOOS=linux go build -ldflags="-s -w -X 'main.Version=${VERSION}-plus' -X 'main.Commit=${COMMIT}' -X 'main.BuildDate=${BUILD_DATE}'" -o ./CLIProxyAPIPlus ./cmd/server/
16 16
17FROM alpine:3.22.017FROM alpine:3.22.0
18 18
@@ -20,7 +20,7 @@ RUN apk add --no-cache tzdata · 1 unmodified line
20 20
21RUN mkdir /CLIProxyAPI21RUN mkdir /CLIProxyAPI
22 22
23COPY --from=builder ./app/CLIProxyAPI /CLIProxyAPI/CLIProxyAPI23COPY --from=builder ./app/CLIProxyAPIPlus /CLIProxyAPI/CLIProxyAPIPlus
24 24
25COPY config.example.yaml /CLIProxyAPI/config.example.yaml25COPY config.example.yaml /CLIProxyAPI/config.example.yaml
26 26
@@ -32,4 +32,4 @@ ENV TZ=Asia/Shanghai · 5 unmodified lines
32 32
33RUN cp /usr/share/zoneinfo/${TZ} /etc/localtime && echo "${TZ}" > /etc/timezone33RUN cp /usr/share/zoneinfo/${TZ} /etc/localtime && echo "${TZ}" > /etc/timezone
34 34
35CMD ["./CLIProxyAPI"]35CMD ["./CLIProxyAPIPlus"]
README.md+10−112
@@ -1,125 +1,23 @@
1# CLI Proxy API1# CLIProxyAPI Plus
2 2
3English | [中文](README_CN.md)3English | [Chinese](README_CN.md)
4 4
5A proxy server that provides OpenAI/Gemini/Claude/Codex compatible API interfaces for CLI.5This is the Plus version of [CLIProxyAPI](https://github.com/router-for-me/CLIProxyAPI), adding support for third-party providers on top of the mainline project.
6 6
7It now also supports OpenAI Codex (GPT models) and Claude Code via OAuth.7All third-party provider support is maintained by community contributors; CLIProxyAPI does not provide technical support. Please contact the corresponding community maintainer if you need assistance.
8 8
9So you can use local or multi-account CLI access with OpenAI(include Responses)/Gemini/Claude-compatible clients and SDKs.9The Plus release stays in lockstep with the mainline features.
10 10
11## Sponsor11## Differences from the Mainline
12 12
13[![z.ai](https://assets.router-for.me/english-4.7.png)](https://z.ai/subscribe?ic=8JVLJQFSKB)13- Added GitHub Copilot support (OAuth login), provided by [em4go](https://github.com/em4go/CLIProxyAPI/tree/feature/github-copilot-auth)
14 14- Added Kiro (AWS CodeWhisperer) support (OAuth login), provided by [fuko2935](https://github.com/fuko2935/CLIProxyAPI/tree/feature/kiro-integration), [Ravens2121](https://github.com/Ravens2121/CLIProxyAPIPlus/)
15This project is sponsored by Z.ai, supporting us with their GLM CODING PLAN.
16
17GLM CODING PLAN is a subscription service designed for AI coding, starting at just $3/month. It provides access to their flagship GLM-4.7 model across 10+ popular AI coding tools (Claude Code, Cline, Roo Code, etc.), offering developers top-tier, fast, and stable coding experiences.
18
19Get 10% OFF GLM CODING PLAN：https://z.ai/subscribe?ic=8JVLJQFSKB
20
21---
22
23<table>
24<tbody>
25<tr>
26<td width="180"><a href="https://www.packyapi.com/register?aff=cliproxyapi"><img src="./assets/packycode.png" alt="PackyCode" width="150"></a></td>
27<td>Thanks to PackyCode for sponsoring this project! PackyCode is a reliable and efficient API relay service provider, offering relay services for Claude Code, Codex, Gemini, and more. PackyCode provides special discounts for our software users: register using <a href="https://www.packyapi.com/register?aff=cliproxyapi">this link</a> and enter the "cliproxyapi" promo code during recharge to get 10% off.</td>
28</tr>
29<tr>
30<td width="180"><a href="https://cubence.com/signup?code=CLIPROXYAPI&source=cpa"><img src="./assets/cubence.png" alt="Cubence" width="150"></a></td>
31<td>Thanks to Cubence for sponsoring this project! Cubence is a reliable and efficient API relay service provider, offering relay services for Claude Code, Codex, Gemini, and more. Cubence provides special discounts for our software users: register using <a href="https://cubence.com/signup?code=CLIPROXYAPI&source=cpa">this link</a> and enter the "CLIPROXYAPI" promo code during recharge to get 10% off.</td>
32</tr>
33</tbody>
34</table>
35
36## Overview
37
38- OpenAI/Gemini/Claude compatible API endpoints for CLI models
39- OpenAI Codex support (GPT models) via OAuth login
40- Claude Code support via OAuth login
41- Qwen Code support via OAuth login
42- iFlow support via OAuth login
43- Amp CLI and IDE extensions support with provider routing
44- Streaming and non-streaming responses
45- Function calling/tools support
46- Multimodal input support (text and images)
47- Multiple accounts with round-robin load balancing (Gemini, OpenAI, Claude, Qwen and iFlow)
48- Simple CLI authentication flows (Gemini, OpenAI, Claude, Qwen and iFlow)
49- Generative Language API Key support
50- AI Studio Build multi-account load balancing
51- Gemini CLI multi-account load balancing
52- Claude Code multi-account load balancing
53- Qwen Code multi-account load balancing
54- iFlow multi-account load balancing
55- OpenAI Codex multi-account load balancing
56- OpenAI-compatible upstream providers via config (e.g., OpenRouter)
57- Reusable Go SDK for embedding the proxy (see `docs/sdk-usage.md`)
58
59## Getting Started
60
61CLIProxyAPI Guides: [https://help.router-for.me/](https://help.router-for.me/)
62
63## Management API
64
65see [MANAGEMENT_API.md](https://help.router-for.me/management/api)
66
67## Amp CLI Support
68
69CLIProxyAPI includes integrated support for [Amp CLI](https://ampcode.com) and Amp IDE extensions, enabling you to use your Google/ChatGPT/Claude OAuth subscriptions with Amp's coding tools:
70
71- Provider route aliases for Amp's API patterns (`/api/provider/{provider}/v1...`)
72- Management proxy for OAuth authentication and account features
73- Smart model fallback with automatic routing
74- **Model mapping** to route unavailable models to alternatives (e.g., `claude-opus-4.5` → `claude-sonnet-4`)
75- Security-first design with localhost-only management endpoints
76
77**→ [Complete Amp CLI Integration Guide](https://help.router-for.me/agent-client/amp-cli.html)**
78
79## SDK Docs
80
81- Usage: [docs/sdk-usage.md](docs/sdk-usage.md)
82- Advanced (executors & translators): [docs/sdk-advanced.md](docs/sdk-advanced.md)
83- Access: [docs/sdk-access.md](docs/sdk-access.md)
84- Watcher: [docs/sdk-watcher.md](docs/sdk-watcher.md)
85- Custom Provider Example: `examples/custom-provider`
86 15
87## Contributing16## Contributing
88 17
89Contributions are welcome! Please feel free to submit a Pull Request.18This project only accepts pull requests that relate to third-party provider support. Any pull requests unrelated to third-party provider support will be rejected.
90
911. Fork the repository
922. Create your feature branch (`git checkout -b feature/amazing-feature`)
933. Commit your changes (`git commit -m 'Add some amazing feature'`)
944. Push to the branch (`git push origin feature/amazing-feature`)
955. Open a Pull Request
96
97## Who is with us?
98
99Those projects are based on CLIProxyAPI:
100
101### [vibeproxy](https://github.com/automazeio/vibeproxy)
102
103Native macOS menu bar app to use your Claude Code & ChatGPT subscriptions with AI coding tools - no API keys needed
104
105### [Subtitle Translator](https://github.com/VjayC/SRT-Subtitle-Translator-Validator)
106
107Browser-based tool to translate SRT subtitles using your Gemini subscription via CLIProxyAPI with automatic validation/error correction - no API keys needed
108
109### [CCS (Claude Code Switch)](https://github.com/kaitranntt/ccs)
110
111CLI wrapper for instant switching between multiple Claude accounts and alternative models (Gemini, Codex, Antigravity) via CLIProxyAPI OAuth - no API keys needed
112
113### [ProxyPal](https://github.com/heyhuynhgiabuu/proxypal)
114
115Native macOS GUI for managing CLIProxyAPI: configure providers, model mappings, and endpoints via OAuth - no API keys needed.
116
117### [Quotio](https://github.com/nguyenphutrong/quotio)
118
119Native macOS menu bar app that unifies Claude, Gemini, OpenAI, Qwen, and Antigravity subscriptions with real-time quota tracking and smart auto-failover for AI coding tools like Claude Code, OpenCode, and Droid - no API keys needed.
120 19
121> [!NOTE] 20If you need to submit any non-third-party provider changes, please open them against the mainline repository.
122> If you developed a project based on CLIProxyAPI, please open a PR to add it to this list.
123 21
124## License22## License
125 23
README_CN.md+10−119
@@ -1,133 +1,24 @@
1# CLI 代理 API1# CLIProxyAPI Plus
2 2
3[English](README.md) | 中文3[English](README.md) | 中文
4 4
5一个为 CLI 提供 OpenAI/Gemini/Claude/Codex 兼容 API 接口的代理服务器。5这是 [CLIProxyAPI](https://github.com/router-for-me/CLIProxyAPI) 的 Plus 版本，在原有基础上增加了第三方供应商的支持。
6 6
7现已支持通过 OAuth 登录接入 OpenAI Codex（GPT 系列）和 Claude Code。7所有的第三方供应商支持都由第三方社区维护者提供，CLIProxyAPI 不提供技术支持。如需取得支持，请与对应的社区维护者联系。
8 8
9您可以使用本地或多账户的CLI方式，通过任何与 OpenAI（包括Responses）/Gemini/Claude 兼容的客户端和SDK进行访问。9该 Plus 版本的主线功能与主线功能强制同步。
10 10
11## 赞助商11## 与主线版本版本差异
12 12
13[![bigmodel.cn](https://assets.router-for.me/chinese-4.7.png)](https://www.bigmodel.cn/claude-code?ic=RRVJPB5SII)13- 新增 GitHub Copilot 支持（OAuth 登录），由[em4go](https://github.com/em4go/CLIProxyAPI/tree/feature/github-copilot-auth)提供
14 14- 新增 Kiro (AWS CodeWhisperer) 支持 (OAuth 登录), 由[fuko2935](https://github.com/fuko2935/CLIProxyAPI/tree/feature/kiro-integration)、[Ravens2121](https://github.com/Ravens2121/CLIProxyAPIPlus/)提供
15本项目由 Z智谱 提供赞助, 他们通过 GLM CODING PLAN 对本项目提供技术支持。
16
17GLM CODING PLAN 是专为AI编码打造的订阅套餐，每月最低仅需20元，即可在十余款主流AI编码工具如 Claude Code、Cline、Roo Code 中畅享智谱旗舰模型GLM-4.7，为开发者提供顶尖的编码体验。
18
19智谱AI为本软件提供了特别优惠，使用以下链接购买可以享受九折优惠：https://www.bigmodel.cn/claude-code?ic=RRVJPB5SII
20
21---
22
23<table>
24<tbody>
25<tr>
26<td width="180"><a href="https://www.packyapi.com/register?aff=cliproxyapi"><img src="./assets/packycode.png" alt="PackyCode" width="150"></a></td>
27<td>感谢 PackyCode 对本项目的赞助！PackyCode 是一家可靠高效的 API 中转服务商，提供 Claude Code、Codex、Gemini 等多种服务的中转。PackyCode 为本软件用户提供了特别优惠：使用<a href="https://www.packyapi.com/register?aff=cliproxyapi">此链接</a>注册，并在充值时输入 "cliproxyapi" 优惠码即可享受九折优惠。</td>
28</tr>
29<tr>
30<td width="180"><a href="https://cubence.com/signup?code=CLIPROXYAPI&source=cpa"><img src="./assets/cubence.png" alt="Cubence" width="150"></a></td>
31<td>感谢 Cubence 对本项目的赞助！Cubence 是一家可靠高效的 API 中转服务商，提供 Claude Code、Codex、Gemini 等多种服务的中转。Cubence 为本软件用户提供了特别优惠：使用<a href="https://cubence.com/signup?code=CLIPROXYAPI&source=cpa">此链接</a>注册，并在充值时输入 "CLIPROXYAPI" 优惠码即可享受九折优惠。</td>
32</tr>
33</tbody>
34</table>
35
36
37## 功能特性
38
39- 为 CLI 模型提供 OpenAI/Gemini/Claude/Codex 兼容的 API 端点
40- 新增 OpenAI Codex（GPT 系列）支持（OAuth 登录）
41- 新增 Claude Code 支持（OAuth 登录）
42- 新增 Qwen Code 支持（OAuth 登录）
43- 新增 iFlow 支持（OAuth 登录）
44- 支持流式与非流式响应
45- 函数调用/工具支持
46- 多模态输入（文本、图片）
47- 多账户支持与轮询负载均衡（Gemini、OpenAI、Claude、Qwen 与 iFlow）
48- 简单的 CLI 身份验证流程（Gemini、OpenAI、Claude、Qwen 与 iFlow）
49- 支持 Gemini AIStudio API 密钥
50- 支持 AI Studio Build 多账户轮询
51- 支持 Gemini CLI 多账户轮询
52- 支持 Claude Code 多账户轮询
53- 支持 Qwen Code 多账户轮询
54- 支持 iFlow 多账户轮询
55- 支持 OpenAI Codex 多账户轮询
56- 通过配置接入上游 OpenAI 兼容提供商（例如 OpenRouter）
57- 可复用的 Go SDK（见 `docs/sdk-usage_CN.md`）
58
59## 新手入门
60
61CLIProxyAPI 用户手册： [https://help.router-for.me/](https://help.router-for.me/cn/)
62
63## 管理 API 文档
64
65请参见 [MANAGEMENT_API_CN.md](https://help.router-for.me/cn/management/api)
66
67## Amp CLI 支持
68
69CLIProxyAPI 已内置对 [Amp CLI](https://ampcode.com) 和 Amp IDE 扩展的支持，可让你使用自己的 Google/ChatGPT/Claude OAuth 订阅来配合 Amp 编码工具：
70
71- 提供商路由别名，兼容 Amp 的 API 路径模式（`/api/provider/{provider}/v1...`）
72- 管理代理，处理 OAuth 认证和账号功能
73- 智能模型回退与自动路由
74- 以安全为先的设计，管理端点仅限 localhost
75
76**→ [Amp CLI 完整集成指南](https://help.router-for.me/cn/agent-client/amp-cli.html)**
77
78## SDK 文档
79
80- 使用文档：[docs/sdk-usage_CN.md](docs/sdk-usage_CN.md)
81- 高级（执行器与翻译器）：[docs/sdk-advanced_CN.md](docs/sdk-advanced_CN.md)
82- 认证: [docs/sdk-access_CN.md](docs/sdk-access_CN.md)
83- 凭据加载/更新: [docs/sdk-watcher_CN.md](docs/sdk-watcher_CN.md)
84- 自定义 Provider 示例：`examples/custom-provider`
85 15
86## 贡献16## 贡献
87 17
88欢迎贡献！请随时提交 Pull Request。18该项目仅接受第三方供应商支持的 Pull Request。任何非第三方供应商支持的 Pull Request 都将被拒绝。
89 19
901. Fork 仓库20如果需要提交任何非第三方供应商支持的 Pull Request，请提交到主线版本。
912. 创建您的功能分支（`git checkout -b feature/amazing-feature`）
923. 提交您的更改（`git commit -m 'Add some amazing feature'`）
934. 推送到分支（`git push origin feature/amazing-feature`）
945. 打开 Pull Request
95
96## 谁与我们在一起？
97
98这些项目基于 CLIProxyAPI:
99
100### [vibeproxy](https://github.com/automazeio/vibeproxy)
101
102一个原生 macOS 菜单栏应用，让您可以使用 Claude Code & ChatGPT 订阅服务和 AI 编程工具，无需 API 密钥。
103
104### [Subtitle Translator](https://github.com/VjayC/SRT-Subtitle-Translator-Validator)
105
106一款基于浏览器的 SRT 字幕翻译工具，可通过 CLI 代理 API 使用您的 Gemini 订阅。内置自动验证与错误修正功能，无需 API 密钥。
107
108### [CCS (Claude Code Switch)](https://github.com/kaitranntt/ccs)
109
110CLI 封装器，用于通过 CLIProxyAPI OAuth 即时切换多个 Claude 账户和替代模型（Gemini, Codex, Antigravity），无需 API 密钥。
111
112### [ProxyPal](https://github.com/heyhuynhgiabuu/proxypal)
113
114基于 macOS 平台的原生 CLIProxyAPI GUI：配置供应商、模型映射以及OAuth端点，无需 API 密钥。
115
116### [Quotio](https://github.com/nguyenphutrong/quotio)
117
118原生 macOS 菜单栏应用，统一管理 Claude、Gemini、OpenAI、Qwen 和 Antigravity 订阅，提供实时配额追踪和智能自动故障转移，支持 Claude Code、OpenCode 和 Droid 等 AI 编程工具，无需 API 密钥。
119
120> [!NOTE]
121> 如果你开发了基于 CLIProxyAPI 的项目，请提交一个 PR（拉取请求）将其添加到此列表中。
122 21
123## 许可证22## 许可证
124 23
125此项目根据 MIT 许可证授权 - 有关详细信息，请参阅 [LICENSE](LICENSE) 文件。24此项目根据 MIT 许可证授权 - 有关详细信息，请参阅 [LICENSE](LICENSE) 文件。
126
127## 写给所有中国网友的
128
129QQ 群：188637136
130
131或
132
133Telegram 群：https://t.me/CLIProxyAPI
cmd/server/main.go+56−0
@@ -47,6 +47,19 @@ func init() { · 46 unmodified lines
47 buildinfo.BuildDate = BuildDate47 buildinfo.BuildDate = BuildDate
48}48}
49 49
50// setKiroIncognitoMode sets the incognito browser mode for Kiro authentication.
51// Kiro defaults to incognito mode for multi-account support.
52// Users can explicitly override with --incognito or --no-incognito flags.
53func setKiroIncognitoMode(cfg *config.Config, useIncognito, noIncognito bool) {
54 if useIncognito {
55 cfg.IncognitoBrowser = true
56 } else if noIncognito {
57 cfg.IncognitoBrowser = false
58 } else {
59 cfg.IncognitoBrowser = true // Kiro default
60 }
61}
62
50// main is the entry point of the application.63// main is the entry point of the application.
51// It parses command-line flags, loads configuration, and starts the appropriate64// It parses command-line flags, loads configuration, and starts the appropriate
52// service based on the provided flags (login, codex-login, or server mode).65// service based on the provided flags (login, codex-login, or server mode).
@@ -62,10 +75,18 @@ func main() { · 9 unmodified lines
62 var iflowCookie bool75 var iflowCookie bool
63 var noBrowser bool76 var noBrowser bool
64 var antigravityLogin bool77 var antigravityLogin bool
78 var kiroLogin bool
79 var kiroGoogleLogin bool
80 var kiroAWSLogin bool
81 var kiroAWSAuthCode bool
82 var kiroImport bool
83 var githubCopilotLogin bool
65 var projectID string84 var projectID string
66 var vertexImport string85 var vertexImport string
67 var configPath string86 var configPath string
68 var password string87 var password string
88 var noIncognito bool
89 var useIncognito bool
69 90
70 // Define command-line flags for different operation modes.91 // Define command-line flags for different operation modes.
71 flag.BoolVar(&login, "login", false, "Login Google Account")92 flag.BoolVar(&login, "login", false, "Login Google Account")
@@ -75,7 +96,15 @@ func main() { · 3 unmodified lines
75 flag.BoolVar(&iflowLogin, "iflow-login", false, "Login to iFlow using OAuth")96 flag.BoolVar(&iflowLogin, "iflow-login", false, "Login to iFlow using OAuth")
76 flag.BoolVar(&iflowCookie, "iflow-cookie", false, "Login to iFlow using Cookie")97 flag.BoolVar(&iflowCookie, "iflow-cookie", false, "Login to iFlow using Cookie")
77 flag.BoolVar(&noBrowser, "no-browser", false, "Don't open browser automatically for OAuth")98 flag.BoolVar(&noBrowser, "no-browser", false, "Don't open browser automatically for OAuth")
99 flag.BoolVar(&useIncognito, "incognito", false, "Open browser in incognito/private mode for OAuth (useful for multiple accounts)")
100 flag.BoolVar(&noIncognito, "no-incognito", false, "Force disable incognito mode (uses existing browser session)")
78 flag.BoolVar(&antigravityLogin, "antigravity-login", false, "Login to Antigravity using OAuth")101 flag.BoolVar(&antigravityLogin, "antigravity-login", false, "Login to Antigravity using OAuth")
102 flag.BoolVar(&kiroLogin, "kiro-login", false, "Login to Kiro using Google OAuth")
103 flag.BoolVar(&kiroGoogleLogin, "kiro-google-login", false, "Login to Kiro using Google OAuth (same as --kiro-login)")
104 flag.BoolVar(&kiroAWSLogin, "kiro-aws-login", false, "Login to Kiro using AWS Builder ID (device code flow)")
105 flag.BoolVar(&kiroAWSAuthCode, "kiro-aws-authcode", false, "Login to Kiro using AWS Builder ID (authorization code flow, better UX)")
106 flag.BoolVar(&kiroImport, "kiro-import", false, "Import Kiro token from Kiro IDE (~/.aws/sso/cache/kiro-auth-token.json)")
107 flag.BoolVar(&githubCopilotLogin, "github-copilot-login", false, "Login to GitHub Copilot using device flow")
79 flag.StringVar(&projectID, "project_id", "", "Project ID (Gemini only, not required)")108 flag.StringVar(&projectID, "project_id", "", "Project ID (Gemini only, not required)")
80 flag.StringVar(&configPath, "config", DefaultConfigPath, "Configure File Path")109 flag.StringVar(&configPath, "config", DefaultConfigPath, "Configure File Path")
81 flag.StringVar(&vertexImport, "vertex-import", "", "Import Vertex service account key JSON file")110 flag.StringVar(&vertexImport, "vertex-import", "", "Import Vertex service account key JSON file")
@@ -453,6 +482,9 @@ func main() { · 371 unmodified lines
453 } else if antigravityLogin {482 } else if antigravityLogin {
454 // Handle Antigravity login483 // Handle Antigravity login
455 cmd.DoAntigravityLogin(cfg, options)484 cmd.DoAntigravityLogin(cfg, options)
485 } else if githubCopilotLogin {
486 // Handle GitHub Copilot login
487 cmd.DoGitHubCopilotLogin(cfg, options)
456 } else if codexLogin {488 } else if codexLogin {
457 // Handle Codex login489 // Handle Codex login
458 cmd.DoCodexLogin(cfg, options)490 cmd.DoCodexLogin(cfg, options)
@@ -465,6 +497,30 @@ func main() { · 6 unmodified lines
465 cmd.DoIFlowLogin(cfg, options)497 cmd.DoIFlowLogin(cfg, options)
466 } else if iflowCookie {498 } else if iflowCookie {
467 cmd.DoIFlowCookieAuth(cfg, options)499 cmd.DoIFlowCookieAuth(cfg, options)
500 } else if kiroLogin {
501 // For Kiro auth, default to incognito mode for multi-account support
502 // Users can explicitly override with --no-incognito
503 // Note: This config mutation is safe - auth commands exit after completion
504 // and don't share config with StartService (which is in the else branch)
505 setKiroIncognitoMode(cfg, useIncognito, noIncognito)
506 cmd.DoKiroLogin(cfg, options)
507 } else if kiroGoogleLogin {
508 // For Kiro auth, default to incognito mode for multi-account support
509 // Users can explicitly override with --no-incognito
510 // Note: This config mutation is safe - auth commands exit after completion
511 setKiroIncognitoMode(cfg, useIncognito, noIncognito)
512 cmd.DoKiroGoogleLogin(cfg, options)
513 } else if kiroAWSLogin {
514 // For Kiro auth, default to incognito mode for multi-account support
515 // Users can explicitly override with --no-incognito
516 setKiroIncognitoMode(cfg, useIncognito, noIncognito)
517 cmd.DoKiroAWSLogin(cfg, options)
518 } else if kiroAWSAuthCode {
519 // For Kiro auth with authorization code flow (better UX)
520 setKiroIncognitoMode(cfg, useIncognito, noIncognito)
521 cmd.DoKiroAWSAuthCodeLogin(cfg, options)
522 } else if kiroImport {
523 cmd.DoKiroImport(cfg, options)
468 } else {524 } else {
469 // In cloud deploy mode without config file, just wait for shutdown signals525 // In cloud deploy mode without config file, just wait for shutdown signals
470 if isCloudDeploy && !configFileExists {526 if isCloudDeploy && !configFileExists {
config.example.yaml+15−0
@@ -43,6 +43,11 @@ debug: false · 42 unmodified lines
43# When true, disable high-overhead HTTP middleware features to reduce per-request memory usage under high concurrency.43# When true, disable high-overhead HTTP middleware features to reduce per-request memory usage under high concurrency.
44commercial-mode: false44commercial-mode: false
45 45
46# Open OAuth URLs in incognito/private browser mode.
47# Useful when you want to login with a different account without logging out from your current session.
48# Default: false (but Kiro auth defaults to true for multi-account support)
49incognito-browser: true
50
46# When true, write application logs to rotating files instead of stdout51# When true, write application logs to rotating files instead of stdout
47logging-to-file: false52logging-to-file: false
48 53
@@ -135,6 +140,16 @@ ws-auth: false · 86 unmodified lines
135# - "*-thinking" # wildcard matching suffix (e.g. claude-opus-4-5-thinking)140# - "*-thinking" # wildcard matching suffix (e.g. claude-opus-4-5-thinking)
136# - "*haiku*" # wildcard matching substring (e.g. claude-3-5-haiku-20241022)141# - "*haiku*" # wildcard matching substring (e.g. claude-3-5-haiku-20241022)
137 142
143# Kiro (AWS CodeWhisperer) configuration
144# Note: Kiro API currently only operates in us-east-1 region
145#kiro:
146# - token-file: "~/.aws/sso/cache/kiro-auth-token.json" # path to Kiro token file
147# agent-task-type: "" # optional: "vibe" or empty (API default)
148# - access-token: "aoaAAAAA..." # or provide tokens directly
149# refresh-token: "aorAAAAA..."
150# profile-arn: "arn:aws:codewhisperer:us-east-1:..."
151# proxy-url: "socks5://proxy.example.com:1080" # optional: proxy override
152
138# OpenAI compatibility providers153# OpenAI compatibility providers
139# openai-compatibility:154# openai-compatibility:
140# - name: "openrouter" # The name of the provider; it will be used in the user agent and other places.155# - name: "openrouter" # The name of the provider; it will be used in the user agent and other places.
docker-compose.yml+2−2
@@ -1,6 +1,6 @@
1services:1services:
2 cli-proxy-api:2 cli-proxy-api:
3 image: ${CLI_PROXY_IMAGE:-eceasy/cli-proxy-api:latest}3 image: ${CLI_PROXY_IMAGE:-eceasy/cli-proxy-api-plus:latest}
4 pull_policy: always4 pull_policy: always
5 build:5 build:
6 context: .6 context: .
@@ -9,7 +9,7 @@ services: · 2 unmodified lines
9 VERSION: ${VERSION:-dev}9 VERSION: ${VERSION:-dev}
10 COMMIT: ${COMMIT:-none}10 COMMIT: ${COMMIT:-none}
11 BUILD_DATE: ${BUILD_DATE:-unknown}11 BUILD_DATE: ${BUILD_DATE:-unknown}
12 container_name: cli-proxy-api12 container_name: cli-proxy-api-plus
13 # env_file:13 # env_file:
14 # - .env14 # - .env
15 environment:15 environment:
go.mod+2−1
@@ -13,14 +13,15 @@ require ( · 12 unmodified lines
13 github.com/joho/godotenv v1.5.113 github.com/joho/godotenv v1.5.1
14 github.com/klauspost/compress v1.17.414 github.com/klauspost/compress v1.17.4
15 github.com/minio/minio-go/v7 v7.0.6615 github.com/minio/minio-go/v7 v7.0.66
16 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
16 github.com/sirupsen/logrus v1.9.317 github.com/sirupsen/logrus v1.9.3
17 github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966
18 github.com/tidwall/gjson v1.18.018 github.com/tidwall/gjson v1.18.0
19 github.com/tidwall/sjson v1.2.519 github.com/tidwall/sjson v1.2.5
20 github.com/tiktoken-go/tokenizer v0.7.020 github.com/tiktoken-go/tokenizer v0.7.0
21 golang.org/x/crypto v0.45.021 golang.org/x/crypto v0.45.0
22 golang.org/x/net v0.47.022 golang.org/x/net v0.47.0
23 golang.org/x/oauth2 v0.30.023 golang.org/x/oauth2 v0.30.0
24 golang.org/x/term v0.37.0
24 gopkg.in/natefinch/lumberjack.v2 v2.2.125 gopkg.in/natefinch/lumberjack.v2 v2.2.1
25 gopkg.in/yaml.v3 v3.0.126 gopkg.in/yaml.v3 v3.0.1
26)27)
go.sum+3−2
@@ -116,6 +116,8 @@ github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6 · 115 unmodified lines
116github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=116github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
117github.com/pjbgf/sha1cd v0.5.0 h1:a+UkboSi1znleCDUNT3M5YxjOnN1fz2FhN48FlwCxs0=117github.com/pjbgf/sha1cd v0.5.0 h1:a+UkboSi1znleCDUNT3M5YxjOnN1fz2FhN48FlwCxs0=
118github.com/pjbgf/sha1cd v0.5.0/go.mod h1:lhpGlyHLpQZoxMv8HcgXvZEhcGs0PG/vsZnEJ7H0iCM=118github.com/pjbgf/sha1cd v0.5.0/go.mod h1:lhpGlyHLpQZoxMv8HcgXvZEhcGs0PG/vsZnEJ7H0iCM=
119github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
120github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
119github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=121github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
120github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=122github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
121github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=123github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
@@ -126,8 +128,6 @@ github.com/sergi/go-diff v1.4.0 h1:n/SP9D5ad1fORl+llWyN+D6qoUETXNZARKjyY2/KVCw= · 4 unmodified lines
126github.com/sergi/go-diff v1.4.0/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4=128github.com/sergi/go-diff v1.4.0/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4=
127github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=129github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
128github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=130github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
129github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 h1:JIAuq3EEf9cgbU6AtGPK4CTG3Zf6CKMNqf0MHTggAUA=
130github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966/go.mod h1:sUM3LWHvSMaG192sy56D9F7CNvL7jUJVXoqM1QKLnog=
131github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=131github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
132github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=132github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
133github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=133github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
@@ -169,6 +169,7 @@ golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKl · 35 unmodified lines
169golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I=169golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I=
170golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=170golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
171golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=171golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
172golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
172golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=173golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
173golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc=174golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc=
174golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=175golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
internal/api/handlers/management/auth_files.go+318−0
@@ -3,6 +3,9 @@ package management · 2 unmodified lines
3import (3import (
4 "bytes"4 "bytes"
5 "context"5 "context"
6 "crypto/rand"
7 "crypto/sha256"
8 "encoding/base64"
6 "encoding/json"9 "encoding/json"
7 "errors"10 "errors"
8 "fmt"11 "fmt"
@@ -23,6 +26,7 @@ import ( · 14 unmodified lines
23 "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/codex"26 "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/codex"
24 geminiAuth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/gemini"27 geminiAuth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/gemini"
25 iflowauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/iflow"28 iflowauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/iflow"
29 kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"
26 "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/qwen"30 "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/qwen"
27 "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"31 "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
28 "github.com/router-for-me/CLIProxyAPI/v6/internal/misc"32 "github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
@@ -2281,8 +2285,322 @@ func (h *Handler) GetAuthStatus(c *gin.Context) { · 2252 unmodified lines
2281 return2285 return
2282 }2286 }
2283 if status != "" {2287 if status != "" {
2288 if strings.HasPrefix(status, "device_code|") {
2289 parts := strings.SplitN(status, "|", 3)
2290 if len(parts) == 3 {
2291 c.JSON(http.StatusOK, gin.H{
2292 "status": "device_code",
2293 "verification_url": parts[1],
2294 "user_code": parts[2],
2295 })
2296 return
2297 }
2298 }
2299 if strings.HasPrefix(status, "auth_url|") {
2300 authURL := strings.TrimPrefix(status, "auth_url|")
2301 c.JSON(http.StatusOK, gin.H{
2302 "status": "auth_url",
2303 "url": authURL,
2304 })
2305 return
2306 }
2284 c.JSON(http.StatusOK, gin.H{"status": "error", "error": status})2307 c.JSON(http.StatusOK, gin.H{"status": "error", "error": status})
2285 return2308 return
2286 }2309 }
2287 c.JSON(http.StatusOK, gin.H{"status": "wait"})2310 c.JSON(http.StatusOK, gin.H{"status": "wait"})
2288}2311}
2312
2313const kiroCallbackPort = 9876
2314
2315func (h *Handler) RequestKiroToken(c *gin.Context) {
2316 ctx := context.Background()
2317
2318 // Get the login method from query parameter (default: aws for device code flow)
2319 method := strings.ToLower(strings.TrimSpace(c.Query("method")))
2320 if method == "" {
2321 method = "aws"
2322 }
2323
2324 fmt.Println("Initializing Kiro authentication...")
2325
2326 state := fmt.Sprintf("kiro-%d", time.Now().UnixNano())
2327
2328 switch method {
2329 case "aws", "builder-id":
2330 RegisterOAuthSession(state, "kiro")
2331
2332 // AWS Builder ID uses device code flow (no callback needed)
2333 go func() {
2334 ssoClient := kiroauth.NewSSOOIDCClient(h.cfg)
2335
2336 // Step 1: Register client
2337 fmt.Println("Registering client...")
2338 regResp, errRegister := ssoClient.RegisterClient(ctx)
2339 if errRegister != nil {
2340 log.Errorf("Failed to register client: %v", errRegister)
2341 SetOAuthSessionError(state, "Failed to register client")
2342 return
2343 }
2344
2345 // Step 2: Start device authorization
2346 fmt.Println("Starting device authorization...")
2347 authResp, errAuth := ssoClient.StartDeviceAuthorization(ctx, regResp.ClientID, regResp.ClientSecret)
2348 if errAuth != nil {
2349 log.Errorf("Failed to start device auth: %v", errAuth)
2350 SetOAuthSessionError(state, "Failed to start device authorization")
2351 return
2352 }
2353
2354 // Store the verification URL for the frontend to display.
2355 // Using "|" as separator because URLs contain ":".
2356 SetOAuthSessionError(state, "device_code|"+authResp.VerificationURIComplete+"|"+authResp.UserCode)
2357
2358 // Step 3: Poll for token
2359 fmt.Println("Waiting for authorization...")
2360 interval := 5 * time.Second
2361 if authResp.Interval > 0 {
2362 interval = time.Duration(authResp.Interval) * time.Second
2363 }
2364 deadline := time.Now().Add(time.Duration(authResp.ExpiresIn) * time.Second)
2365
2366 for time.Now().Before(deadline) {
2367 select {
2368 case <-ctx.Done():
2369 SetOAuthSessionError(state, "Authorization cancelled")
2370 return
2371 case <-time.After(interval):
2372 tokenResp, errToken := ssoClient.CreateToken(ctx, regResp.ClientID, regResp.ClientSecret, authResp.DeviceCode)
2373 if errToken != nil {
2374 errStr := errToken.Error()
2375 if strings.Contains(errStr, "authorization_pending") {
2376 continue
2377 }
2378 if strings.Contains(errStr, "slow_down") {
2379 interval += 5 * time.Second
2380 continue
2381 }
2382 log.Errorf("Token creation failed: %v", errToken)
2383 SetOAuthSessionError(state, "Token creation failed")
2384 return
2385 }
2386
2387 // Success! Save the token
2388 expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
2389 email := kiroauth.ExtractEmailFromJWT(tokenResp.AccessToken)
2390
2391 idPart := kiroauth.SanitizeEmailForFilename(email)
2392 if idPart == "" {
2393 idPart = fmt.Sprintf("%d", time.Now().UnixNano()%100000)
2394 }
2395
2396 now := time.Now()
2397 fileName := fmt.Sprintf("kiro-aws-%s.json", idPart)
2398
2399 record := &coreauth.Auth{
2400 ID: fileName,
2401 Provider: "kiro",
2402 FileName: fileName,
2403 Metadata: map[string]any{
2404 "type": "kiro",
2405 "access_token": tokenResp.AccessToken,
2406 "refresh_token": tokenResp.RefreshToken,
2407 "expires_at": expiresAt.Format(time.RFC3339),
2408 "auth_method": "builder-id",
2409 "provider": "AWS",
2410 "client_id": regResp.ClientID,
2411 "client_secret": regResp.ClientSecret,
2412 "email": email,
2413 "last_refresh": now.Format(time.RFC3339),
2414 },
2415 }
2416
2417 savedPath, errSave := h.saveTokenRecord(ctx, record)
2418 if errSave != nil {
2419 log.Errorf("Failed to save authentication tokens: %v", errSave)
2420 SetOAuthSessionError(state, "Failed to save authentication tokens")
2421 return
2422 }
2423
2424 fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
2425 if email != "" {
2426 fmt.Printf("Authenticated as: %s\n", email)
2427 }
2428 CompleteOAuthSession(state)
2429 return
2430 }
2431 }
2432
2433 SetOAuthSessionError(state, "Authorization timed out")
2434 }()
2435
2436 // Return immediately with the state for polling
2437 c.JSON(http.StatusOK, gin.H{"status": "ok", "state": state, "method": "device_code"})
2438
2439 case "google", "github":
2440 RegisterOAuthSession(state, "kiro")
2441
2442 // Social auth uses protocol handler - for WEB UI we use a callback forwarder
2443 provider := "Google"
2444 if method == "github" {
2445 provider = "Github"
2446 }
2447
2448 isWebUI := isWebUIRequest(c)
2449 if isWebUI {
2450 targetURL, errTarget := h.managementCallbackURL("/kiro/callback")
2451 if errTarget != nil {
2452 log.WithError(errTarget).Error("failed to compute kiro callback target")
2453 c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"})
2454 return
2455 }
2456 if _, errStart := startCallbackForwarder(kiroCallbackPort, "kiro", targetURL); errStart != nil {
2457 log.WithError(errStart).Error("failed to start kiro callback forwarder")
2458 c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"})
2459 return
2460 }
2461 }
2462
2463 go func() {
2464 if isWebUI {
2465 defer stopCallbackForwarder(kiroCallbackPort)
2466 }
2467
2468 socialClient := kiroauth.NewSocialAuthClient(h.cfg)
2469
2470 // Generate PKCE codes
2471 codeVerifier, codeChallenge, errPKCE := generateKiroPKCE()
2472 if errPKCE != nil {
2473 log.Errorf("Failed to generate PKCE: %v", errPKCE)
2474 SetOAuthSessionError(state, "Failed to generate PKCE")
2475 return
2476 }
2477
2478 // Build login URL
2479 authURL := fmt.Sprintf("%s/login?idp=%s&redirect_uri=%s&code_challenge=%s&code_challenge_method=S256&state=%s&prompt=select_account",
2480 "https://prod.us-east-1.auth.desktop.kiro.dev",
2481 provider,
2482 url.QueryEscape(kiroauth.KiroRedirectURI),
2483 codeChallenge,
2484 state,
2485 )
2486
2487 // Store auth URL for frontend.
2488 // Using "|" as separator because URLs contain ":".
2489 SetOAuthSessionError(state, "auth_url|"+authURL)
2490
2491 // Wait for callback file
2492 waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-kiro-%s.oauth", state))
2493 deadline := time.Now().Add(5 * time.Minute)
2494
2495 for {
2496 if time.Now().After(deadline) {
2497 log.Error("oauth flow timed out")
2498 SetOAuthSessionError(state, "OAuth flow timed out")
2499 return
2500 }
2501 if data, errRead := os.ReadFile(waitFile); errRead == nil {
2502 var m map[string]string
2503 _ = json.Unmarshal(data, &m)
2504 _ = os.Remove(waitFile)
2505 if errStr := m["error"]; errStr != "" {
2506 log.Errorf("Authentication failed: %s", errStr)
2507 SetOAuthSessionError(state, "Authentication failed")
2508 return
2509 }
2510 if m["state"] != state {
2511 log.Errorf("State mismatch")
2512 SetOAuthSessionError(state, "State mismatch")
2513 return
2514 }
2515 code := m["code"]
2516 if code == "" {
2517 log.Error("No authorization code received")
2518 SetOAuthSessionError(state, "No authorization code received")
2519 return
2520 }
2521
2522 // Exchange code for tokens
2523 tokenReq := &kiroauth.CreateTokenRequest{
2524 Code: code,
2525 CodeVerifier: codeVerifier,
2526 RedirectURI: kiroauth.KiroRedirectURI,
2527 }
2528
2529 tokenResp, errToken := socialClient.CreateToken(ctx, tokenReq)
2530 if errToken != nil {
2531 log.Errorf("Failed to exchange code for tokens: %v", errToken)
2532 SetOAuthSessionError(state, "Failed to exchange code for tokens")
2533 return
2534 }
2535
2536 // Save the token
2537 expiresIn := tokenResp.ExpiresIn
2538 if expiresIn <= 0 {
2539 expiresIn = 3600
2540 }
2541 expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
2542 email := kiroauth.ExtractEmailFromJWT(tokenResp.AccessToken)
2543
2544 idPart := kiroauth.SanitizeEmailForFilename(email)
2545 if idPart == "" {
2546 idPart = fmt.Sprintf("%d", time.Now().UnixNano()%100000)
2547 }
2548
2549 now := time.Now()
2550 fileName := fmt.Sprintf("kiro-%s-%s.json", strings.ToLower(provider), idPart)
2551
2552 record := &coreauth.Auth{
2553 ID: fileName,
2554 Provider: "kiro",
2555 FileName: fileName,
2556 Metadata: map[string]any{
2557 "type": "kiro",
2558 "access_token": tokenResp.AccessToken,
2559 "refresh_token": tokenResp.RefreshToken,
2560 "profile_arn": tokenResp.ProfileArn,
2561 "expires_at": expiresAt.Format(time.RFC3339),
2562 "auth_method": "social",
2563 "provider": provider,
2564 "email": email,
2565 "last_refresh": now.Format(time.RFC3339),
2566 },
2567 }
2568
2569 savedPath, errSave := h.saveTokenRecord(ctx, record)
2570 if errSave != nil {
2571 log.Errorf("Failed to save authentication tokens: %v", errSave)
2572 SetOAuthSessionError(state, "Failed to save authentication tokens")
2573 return
2574 }
2575
2576 fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
2577 if email != "" {
2578 fmt.Printf("Authenticated as: %s\n", email)
2579 }
2580 CompleteOAuthSession(state)
2581 return
2582 }
2583 time.Sleep(500 * time.Millisecond)
2584 }
2585 }()
2586
2587 c.JSON(http.StatusOK, gin.H{"status": "ok", "state": state, "method": "social"})
2588
2589 default:
2590 c.JSON(http.StatusBadRequest, gin.H{"error": "invalid method, use 'aws', 'google', or 'github'"})
2591 }
2592}
2593
2594// generateKiroPKCE generates PKCE code verifier and challenge for Kiro OAuth.
2595func generateKiroPKCE() (verifier, challenge string, err error) {
2596 b := make([]byte, 32)
2597 if _, errRead := io.ReadFull(rand.Reader, b); errRead != nil {
2598 return "", "", fmt.Errorf("failed to generate random bytes: %w", errRead)
2599 }
2600 verifier = base64.RawURLEncoding.EncodeToString(b)
2601
2602 h := sha256.Sum256([]byte(verifier))
2603 challenge = base64.RawURLEncoding.EncodeToString(h[:])
2604
2605 return verifier, challenge, nil
2606}
internal/api/handlers/management/config_basic.go+2−2
@@ -19,8 +19,8 @@ import ( · 18 unmodified lines
19)19)
20 20
21const (21const (
22 latestReleaseURL = "https://api.github.com/repos/router-for-me/CLIProxyAPI/releases/latest"22 latestReleaseURL = "https://api.github.com/repos/router-for-me/CLIProxyAPIPlus/releases/latest"
23 latestReleaseUserAgent = "CLIProxyAPI"23 latestReleaseUserAgent = "CLIProxyAPIPlus"
24)24)
25 25
26func (h *Handler) GetConfig(c *gin.Context) {26func (h *Handler) GetConfig(c *gin.Context) {
internal/api/handlers/management/oauth_sessions.go+8−1
@@ -158,7 +158,12 @@ func (s *oauthSessionStore) IsPending(state, provider string) bool { · 157 unmodified lines
158 return false158 return false
159 }159 }
160 if session.Status != "" {160 if session.Status != "" {
161 return false161 if !strings.EqualFold(session.Provider, "kiro") {
162 return false
163 }
164 if !strings.HasPrefix(session.Status, "device_code|") && !strings.HasPrefix(session.Status, "auth_url|") {
165 return false
166 }
162 }167 }
163 if provider == "" {168 if provider == "" {
164 return true169 return true
@@ -231,6 +236,8 @@ func NormalizeOAuthProvider(provider string) (string, error) { · 66 unmodified lines
231 return "antigravity", nil236 return "antigravity", nil
232 case "qwen":237 case "qwen":
233 return "qwen", nil238 return "qwen", nil
239 case "kiro":
240 return "kiro", nil
234 default:241 default:
235 return "", errUnsupportedOAuthFlow242 return "", errUnsupportedOAuthFlow
236 }243 }
internal/api/modules/amp/proxy.go+34−3
@@ -3,8 +3,11 @@ package amp · 2 unmodified lines
3import (3import (
4 "bytes"4 "bytes"
5 "compress/gzip"5 "compress/gzip"
6 "context"
7 "errors"
6 "fmt"8 "fmt"
7 "io"9 "io"
10 "net"
8 "net/http"11 "net/http"
9 "net/http/httputil"12 "net/http/httputil"
10 "net/url"13 "net/url"
@@ -102,7 +105,15 @@ func createReverseProxy(upstreamURL string, secretSource SecretSource) (*httputi · 91 unmodified lines
102 // Modify incoming responses to handle gzip without Content-Encoding105 // Modify incoming responses to handle gzip without Content-Encoding
103 // This addresses the same issue as inline handler gzip handling, but at the proxy level106 // This addresses the same issue as inline handler gzip handling, but at the proxy level
104 proxy.ModifyResponse = func(resp *http.Response) error {107 proxy.ModifyResponse = func(resp *http.Response) error {
105 // Only process successful responses108 // Log upstream error responses for diagnostics (502, 503, etc.)
109 // These are NOT proxy connection errors - the upstream responded with an error status
110 if resp.StatusCode >= 500 {
111 log.Errorf("amp upstream responded with error [%d] for %s %s", resp.StatusCode, resp.Request.Method, resp.Request.URL.Path)
112 } else if resp.StatusCode >= 400 {
113 log.Warnf("amp upstream responded with client error [%d] for %s %s", resp.StatusCode, resp.Request.Method, resp.Request.URL.Path)
114 }
115
116 // Only process successful responses for gzip decompression
106 if resp.StatusCode < 200 || resp.StatusCode >= 300 {117 if resp.StatusCode < 200 || resp.StatusCode >= 300 {
107 return nil118 return nil
108 }119 }
@@ -186,9 +197,29 @@ func createReverseProxy(upstreamURL string, secretSource SecretSource) (*httputi · 77 unmodified lines
186 return nil197 return nil
187 }198 }
188 199
189 // Error handler for proxy failures200 // Error handler for proxy failures with detailed error classification for diagnostics
190 proxy.ErrorHandler = func(rw http.ResponseWriter, req *http.Request, err error) {201 proxy.ErrorHandler = func(rw http.ResponseWriter, req *http.Request, err error) {
191 log.Errorf("amp upstream proxy error for %s %s: %v", req.Method, req.URL.Path, err)202 // Classify the error type for better diagnostics
203 var errType string
204 if errors.Is(err, context.DeadlineExceeded) {
205 errType = "timeout"
206 } else if errors.Is(err, context.Canceled) {
207 errType = "canceled"
208 } else if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
209 errType = "dial_timeout"
210 } else if _, ok := err.(net.Error); ok {
211 errType = "network_error"
212 } else {
213 errType = "connection_error"
214 }
215
216 // Don't log as error for context canceled - it's usually client closing connection
217 if errors.Is(err, context.Canceled) {
218 log.Debugf("amp upstream proxy [%s]: client canceled request for %s %s", errType, req.Method, req.URL.Path)
219 } else {
220 log.Errorf("amp upstream proxy error [%s] for %s %s: %v", errType, req.Method, req.URL.Path, err)
221 }
222
192 rw.Header().Set("Content-Type", "application/json")223 rw.Header().Set("Content-Type", "application/json")
193 rw.WriteHeader(http.StatusBadGateway)224 rw.WriteHeader(http.StatusBadGateway)
194 _, _ = rw.Write([]byte(`{"error":"amp_upstream_proxy_error","message":"Failed to reach Amp upstream"}`))225 _, _ = rw.Write([]byte(`{"error":"amp_upstream_proxy_error","message":"Failed to reach Amp upstream"}`))
internal/api/modules/amp/response_rewriter.go+58−2
@@ -29,15 +29,71 @@ func NewResponseRewriter(w gin.ResponseWriter, originalModel string) *ResponseRe · 28 unmodified lines
29 }29 }
30}30}
31 31
32const maxBufferedResponseBytes = 2 * 1024 * 1024 // 2MB safety cap
33
34func looksLikeSSEChunk(data []byte) bool {
35 // Fallback detection: some upstreams may omit/lie about Content-Type, causing SSE to be buffered.
36 // Heuristics are intentionally simple and cheap.
37 return bytes.Contains(data, []byte("data:")) ||
38 bytes.Contains(data, []byte("event:")) ||
39 bytes.Contains(data, []byte("message_start")) ||
40 bytes.Contains(data, []byte("message_delta")) ||
41 bytes.Contains(data, []byte("content_block_start")) ||
42 bytes.Contains(data, []byte("content_block_delta")) ||
43 bytes.Contains(data, []byte("content_block_stop")) ||
44 bytes.Contains(data, []byte("\n\n"))
45}
46
47func (rw *ResponseRewriter) enableStreaming(reason string) error {
48 if rw.isStreaming {
49 return nil
50 }
51 rw.isStreaming = true
52
53 // Flush any previously buffered data to avoid reordering or data loss.
54 if rw.body != nil && rw.body.Len() > 0 {
55 buf := rw.body.Bytes()
56 // Copy before Reset() to keep bytes stable.
57 toFlush := make([]byte, len(buf))
58 copy(toFlush, buf)
59 rw.body.Reset()
60
61 if _, err := rw.ResponseWriter.Write(rw.rewriteStreamChunk(toFlush)); err != nil {
62 return err
63 }
64 if flusher, ok := rw.ResponseWriter.(http.Flusher); ok {
65 flusher.Flush()
66 }
67 }
68
69 log.Debugf("amp response rewriter: switched to streaming (%s)", reason)
70 return nil
71}
72
32// Write intercepts response writes and buffers them for model name replacement73// Write intercepts response writes and buffers them for model name replacement
33func (rw *ResponseRewriter) Write(data []byte) (int, error) {74func (rw *ResponseRewriter) Write(data []byte) (int, error) {
34 // Detect streaming on first write75 // Detect streaming on first write (header-based)
35 if rw.body.Len() == 0 && !rw.isStreaming {76 if !rw.isStreaming && rw.body.Len() == 0 {
36 contentType := rw.Header().Get("Content-Type")77 contentType := rw.Header().Get("Content-Type")
37 rw.isStreaming = strings.Contains(contentType, "text/event-stream") ||78 rw.isStreaming = strings.Contains(contentType, "text/event-stream") ||
38 strings.Contains(contentType, "stream")79 strings.Contains(contentType, "stream")
39 }80 }
40 81
82 if !rw.isStreaming {
83 // Content-based fallback: detect SSE-like chunks even if Content-Type is missing/wrong.
84 if looksLikeSSEChunk(data) {
85 if err := rw.enableStreaming("sse heuristic"); err != nil {
86 return 0, err
87 }
88 } else if rw.body.Len()+len(data) > maxBufferedResponseBytes {
89 // Safety cap: avoid unbounded buffering on large responses.
90 log.Warnf("amp response rewriter: buffer exceeded %d bytes, switching to streaming", maxBufferedResponseBytes)
91 if err := rw.enableStreaming("buffer limit"); err != nil {
92 return 0, err
93 }
94 }
95 }
96
41 if rw.isStreaming {97 if rw.isStreaming {
42 n, err := rw.ResponseWriter.Write(rw.rewriteStreamChunk(data))98 n, err := rw.ResponseWriter.Write(rw.rewriteStreamChunk(data))
43 if err == nil {99 if err == nil {
internal/api/server.go+21−0
@@ -347,6 +347,12 @@ func (s *Server) setupRoutes() { · 346 unmodified lines
347 },347 },
348 })348 })
349 })349 })
350
351 // Event logging endpoint - handles Claude Code telemetry requests
352 // Returns 200 OK to prevent 404 errors in logs
353 s.engine.POST("/api/event_logging/batch", func(c *gin.Context) {
354 c.JSON(http.StatusOK, gin.H{"status": "ok"})
355 })
350 s.engine.POST("/v1internal:method", geminiCLIHandlers.CLIHandler)356 s.engine.POST("/v1internal:method", geminiCLIHandlers.CLIHandler)
351 357
352 // OAuth callback endpoints (reuse main server port)358 // OAuth callback endpoints (reuse main server port)
@@ -422,6 +428,20 @@ func (s *Server) setupRoutes() { · 69 unmodified lines
422 c.String(http.StatusOK, oauthCallbackSuccessHTML)428 c.String(http.StatusOK, oauthCallbackSuccessHTML)
423 })429 })
424 430
431 s.engine.GET("/kiro/callback", func(c *gin.Context) {
432 code := c.Query("code")
433 state := c.Query("state")
434 errStr := c.Query("error")
435 if errStr == "" {
436 errStr = c.Query("error_description")
437 }
438 if state != "" {
439 _, _ = managementHandlers.WriteOAuthCallbackFileForPendingSession(s.cfg.AuthDir, "kiro", state, code, errStr)
440 }
441 c.Header("Content-Type", "text/html; charset=utf-8")
442 c.String(http.StatusOK, oauthCallbackSuccessHTML)
443 })
444
425 // Management routes are registered lazily by registerManagementRoutes when a secret is configured.445 // Management routes are registered lazily by registerManagementRoutes when a secret is configured.
426}446}
427 447
@@ -597,6 +617,7 @@ func (s *Server) registerManagementRoutes() { · 169 unmodified lines
597 mgmt.GET("/qwen-auth-url", s.mgmt.RequestQwenToken)617 mgmt.GET("/qwen-auth-url", s.mgmt.RequestQwenToken)
598 mgmt.GET("/iflow-auth-url", s.mgmt.RequestIFlowToken)618 mgmt.GET("/iflow-auth-url", s.mgmt.RequestIFlowToken)
599 mgmt.POST("/iflow-auth-url", s.mgmt.RequestIFlowCookieToken)619 mgmt.POST("/iflow-auth-url", s.mgmt.RequestIFlowCookieToken)
620 mgmt.GET("/kiro-auth-url", s.mgmt.RequestKiroToken)
600 mgmt.POST("/oauth-callback", s.mgmt.PostOAuthCallback)621 mgmt.POST("/oauth-callback", s.mgmt.PostOAuthCallback)
601 mgmt.GET("/get-auth-status", s.mgmt.GetAuthStatus)622 mgmt.GET("/get-auth-status", s.mgmt.GetAuthStatus)
602 }623 }
internal/auth/claude/oauth_server.go+11−0
@@ -242,6 +242,11 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) { · 241 unmodified lines
242 platformURL = "https://console.anthropic.com/"242 platformURL = "https://console.anthropic.com/"
243 }243 }
244 244
245 // Validate platformURL to prevent XSS - only allow http/https URLs
246 if !isValidURL(platformURL) {
247 platformURL = "https://console.anthropic.com/"
248 }
249
245 // Generate success page HTML with dynamic content250 // Generate success page HTML with dynamic content
246 successHTML := s.generateSuccessHTML(setupRequired, platformURL)251 successHTML := s.generateSuccessHTML(setupRequired, platformURL)
247 252
@@ -251,6 +256,12 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) { · 3 unmodified lines
251 }256 }
252}257}
253 258
259// isValidURL checks if the URL is a valid http/https URL to prevent XSS
260func isValidURL(urlStr string) bool {
261 urlStr = strings.TrimSpace(urlStr)
262 return strings.HasPrefix(urlStr, "https://") || strings.HasPrefix(urlStr, "http://")
263}
264
254// generateSuccessHTML creates the HTML content for the success page.265// generateSuccessHTML creates the HTML content for the success page.
255// It customizes the page based on whether additional setup is required266// It customizes the page based on whether additional setup is required
256// and includes a link to the platform.267// and includes a link to the platform.
internal/auth/codex/oauth_server.go+11−0
@@ -239,6 +239,11 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) { · 238 unmodified lines
239 platformURL = "https://platform.openai.com"239 platformURL = "https://platform.openai.com"
240 }240 }
241 241
242 // Validate platformURL to prevent XSS - only allow http/https URLs
243 if !isValidURL(platformURL) {
244 platformURL = "https://platform.openai.com"
245 }
246
242 // Generate success page HTML with dynamic content247 // Generate success page HTML with dynamic content
243 successHTML := s.generateSuccessHTML(setupRequired, platformURL)248 successHTML := s.generateSuccessHTML(setupRequired, platformURL)
244 249
@@ -248,6 +253,12 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) { · 3 unmodified lines
248 }253 }
249}254}
250 255
256// isValidURL checks if the URL is a valid http/https URL to prevent XSS
257func isValidURL(urlStr string) bool {
258 urlStr = strings.TrimSpace(urlStr)
259 return strings.HasPrefix(urlStr, "https://") || strings.HasPrefix(urlStr, "http://")
260}
261
251// generateSuccessHTML creates the HTML content for the success page.262// generateSuccessHTML creates the HTML content for the success page.
252// It customizes the page based on whether additional setup is required263// It customizes the page based on whether additional setup is required
253// and includes a link to the platform.264// and includes a link to the platform.
internal/auth/copilot/copilot_auth.go+225−0
@@ -0,0 +1,225 @@
1// Package copilot provides authentication and token management for GitHub Copilot API.
2// It handles the OAuth2 device flow for secure authentication with the Copilot API.
3package copilot
4
5import (
6 "context"
7 "encoding/json"
8 "fmt"
9 "io"
10 "net/http"
11 "time"
12
13 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
14 "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
15 log "github.com/sirupsen/logrus"
16)
17
18const (
19 // copilotAPITokenURL is the endpoint for getting Copilot API tokens from GitHub token.
20 copilotAPITokenURL = "https://api.github.com/copilot_internal/v2/token"
21 // copilotAPIEndpoint is the base URL for making API requests.
22 copilotAPIEndpoint = "https://api.githubcopilot.com"
23
24 // Common HTTP header values for Copilot API requests.
25 copilotUserAgent = "GithubCopilot/1.0"
26 copilotEditorVersion = "vscode/1.100.0"
27 copilotPluginVersion = "copilot/1.300.0"
28 copilotIntegrationID = "vscode-chat"
29 copilotOpenAIIntent = "conversation-panel"
30)
31
32// CopilotAPIToken represents the Copilot API token response.
33type CopilotAPIToken struct {
34 // Token is the JWT token for authenticating with the Copilot API.
35 Token string `json:"token"`
36 // ExpiresAt is the Unix timestamp when the token expires.
37 ExpiresAt int64 `json:"expires_at"`
38 // Endpoints contains the available API endpoints.
39 Endpoints struct {
40 API string `json:"api"`
41 Proxy string `json:"proxy"`
42 OriginTracker string `json:"origin-tracker"`
43 Telemetry string `json:"telemetry"`
44 } `json:"endpoints,omitempty"`
45 // ErrorDetails contains error information if the request failed.
46 ErrorDetails *struct {
47 URL string `json:"url"`
48 Message string `json:"message"`
49 DocumentationURL string `json:"documentation_url"`
50 } `json:"error_details,omitempty"`
51}
52
53// CopilotAuth handles GitHub Copilot authentication flow.
54// It provides methods for device flow authentication and token management.
55type CopilotAuth struct {
56 httpClient *http.Client
57 deviceClient *DeviceFlowClient
58 cfg *config.Config
59}
60
61// NewCopilotAuth creates a new CopilotAuth service instance.
62// It initializes an HTTP client with proxy settings from the provided configuration.
63func NewCopilotAuth(cfg *config.Config) *CopilotAuth {
64 return &CopilotAuth{
65 httpClient: util.SetProxy(&cfg.SDKConfig, &http.Client{Timeout: 30 * time.Second}),
66 deviceClient: NewDeviceFlowClient(cfg),
67 cfg: cfg,
68 }
69}
70
71// StartDeviceFlow initiates the device flow authentication.
72// Returns the device code response containing the user code and verification URI.
73func (c *CopilotAuth) StartDeviceFlow(ctx context.Context) (*DeviceCodeResponse, error) {
74 return c.deviceClient.RequestDeviceCode(ctx)
75}
76
77// WaitForAuthorization polls for user authorization and returns the auth bundle.
78func (c *CopilotAuth) WaitForAuthorization(ctx context.Context, deviceCode *DeviceCodeResponse) (*CopilotAuthBundle, error) {
79 tokenData, err := c.deviceClient.PollForToken(ctx, deviceCode)
80 if err != nil {
81 return nil, err
82 }
83
84 // Fetch the GitHub username
85 username, err := c.deviceClient.FetchUserInfo(ctx, tokenData.AccessToken)
86 if err != nil {
87 log.Warnf("copilot: failed to fetch user info: %v", err)
88 username = "unknown"
89 }
90
91 return &CopilotAuthBundle{
92 TokenData: tokenData,
93 Username: username,
94 }, nil
95}
96
97// GetCopilotAPIToken exchanges a GitHub access token for a Copilot API token.
98// This token is used to make authenticated requests to the Copilot API.
99func (c *CopilotAuth) GetCopilotAPIToken(ctx context.Context, githubAccessToken string) (*CopilotAPIToken, error) {
100 if githubAccessToken == "" {
101 return nil, NewAuthenticationError(ErrTokenExchangeFailed, fmt.Errorf("github access token is empty"))
102 }
103
104 req, err := http.NewRequestWithContext(ctx, http.MethodGet, copilotAPITokenURL, nil)
105 if err != nil {
106 return nil, NewAuthenticationError(ErrTokenExchangeFailed, err)
107 }
108
109 req.Header.Set("Authorization", "token "+githubAccessToken)
110 req.Header.Set("Accept", "application/json")
111 req.Header.Set("User-Agent", copilotUserAgent)
112 req.Header.Set("Editor-Version", copilotEditorVersion)
113 req.Header.Set("Editor-Plugin-Version", copilotPluginVersion)
114
115 resp, err := c.httpClient.Do(req)
116 if err != nil {
117 return nil, NewAuthenticationError(ErrTokenExchangeFailed, err)
118 }
119 defer func() {
120 if errClose := resp.Body.Close(); errClose != nil {
121 log.Errorf("copilot api token: close body error: %v", errClose)
122 }
123 }()
124
125 bodyBytes, err := io.ReadAll(resp.Body)
126 if err != nil {
127 return nil, NewAuthenticationError(ErrTokenExchangeFailed, err)
128 }
129
130 if !isHTTPSuccess(resp.StatusCode) {
131 return nil, NewAuthenticationError(ErrTokenExchangeFailed,
132 fmt.Errorf("status %d: %s", resp.StatusCode, string(bodyBytes)))
133 }
134
135 var apiToken CopilotAPIToken
136 if err = json.Unmarshal(bodyBytes, &apiToken); err != nil {
137 return nil, NewAuthenticationError(ErrTokenExchangeFailed, err)
138 }
139
140 if apiToken.Token == "" {
141 return nil, NewAuthenticationError(ErrTokenExchangeFailed, fmt.Errorf("empty copilot api token"))
142 }
143
144 return &apiToken, nil
145}
146
147// ValidateToken checks if a GitHub access token is valid by attempting to fetch user info.
148func (c *CopilotAuth) ValidateToken(ctx context.Context, accessToken string) (bool, string, error) {
149 if accessToken == "" {
150 return false, "", nil
151 }
152
153 username, err := c.deviceClient.FetchUserInfo(ctx, accessToken)
154 if err != nil {
155 return false, "", err
156 }
157
158 return true, username, nil
159}
160
161// CreateTokenStorage creates a new CopilotTokenStorage from auth bundle.
162func (c *CopilotAuth) CreateTokenStorage(bundle *CopilotAuthBundle) *CopilotTokenStorage {
163 return &CopilotTokenStorage{
164 AccessToken: bundle.TokenData.AccessToken,
165 TokenType: bundle.TokenData.TokenType,
166 Scope: bundle.TokenData.Scope,
167 Username: bundle.Username,
168 Type: "github-copilot",
169 }
170}
171
172// LoadAndValidateToken loads a token from storage and validates it.
173// Returns the storage if valid, or an error if the token is invalid or expired.
174func (c *CopilotAuth) LoadAndValidateToken(ctx context.Context, storage *CopilotTokenStorage) (bool, error) {
175 if storage == nil || storage.AccessToken == "" {
176 return false, fmt.Errorf("no token available")
177 }
178
179 // Check if we can still use the GitHub token to get a Copilot API token
180 apiToken, err := c.GetCopilotAPIToken(ctx, storage.AccessToken)
181 if err != nil {
182 return false, err
183 }
184
185 // Check if the API token is expired
186 if apiToken.ExpiresAt > 0 && time.Now().Unix() >= apiToken.ExpiresAt {
187 return false, fmt.Errorf("copilot api token expired")
188 }
189
190 return true, nil
191}
192
193// GetAPIEndpoint returns the Copilot API endpoint URL.
194func (c *CopilotAuth) GetAPIEndpoint() string {
195 return copilotAPIEndpoint
196}
197
198// MakeAuthenticatedRequest creates an authenticated HTTP request to the Copilot API.
199func (c *CopilotAuth) MakeAuthenticatedRequest(ctx context.Context, method, url string, body io.Reader, apiToken *CopilotAPIToken) (*http.Request, error) {
200 req, err := http.NewRequestWithContext(ctx, method, url, body)
201 if err != nil {
202 return nil, fmt.Errorf("failed to create request: %w", err)
203 }
204
205 req.Header.Set("Authorization", "Bearer "+apiToken.Token)
206 req.Header.Set("Content-Type", "application/json")
207 req.Header.Set("Accept", "application/json")
208 req.Header.Set("User-Agent", copilotUserAgent)
209 req.Header.Set("Editor-Version", copilotEditorVersion)
210 req.Header.Set("Editor-Plugin-Version", copilotPluginVersion)
211 req.Header.Set("Openai-Intent", copilotOpenAIIntent)
212 req.Header.Set("Copilot-Integration-Id", copilotIntegrationID)
213
214 return req, nil
215}
216
217// buildChatCompletionURL builds the URL for chat completions API.
218func buildChatCompletionURL() string {
219 return copilotAPIEndpoint + "/chat/completions"
220}
221
222// isHTTPSuccess checks if the status code indicates success (2xx).
223func isHTTPSuccess(statusCode int) bool {
224 return statusCode >= 200 && statusCode < 300
225}
internal/auth/copilot/errors.go+187−0
@@ -0,0 +1,187 @@
1package copilot
2
3import (
4 "errors"
5 "fmt"
6 "net/http"
7)
8
9// OAuthError represents an OAuth-specific error.
10type OAuthError struct {
11 // Code is the OAuth error code.
12 Code string `json:"error"`
13 // Description is a human-readable description of the error.
14 Description string `json:"error_description,omitempty"`
15 // URI is a URI identifying a human-readable web page with information about the error.
16 URI string `json:"error_uri,omitempty"`
17 // StatusCode is the HTTP status code associated with the error.
18 StatusCode int `json:"-"`
19}
20
21// Error returns a string representation of the OAuth error.
22func (e *OAuthError) Error() string {
23 if e.Description != "" {
24 return fmt.Sprintf("OAuth error %s: %s", e.Code, e.Description)
25 }
26 return fmt.Sprintf("OAuth error: %s", e.Code)
27}
28
29// NewOAuthError creates a new OAuth error with the specified code, description, and status code.
30func NewOAuthError(code, description string, statusCode int) *OAuthError {
31 return &OAuthError{
32 Code: code,
33 Description: description,
34 StatusCode: statusCode,
35 }
36}
37
38// AuthenticationError represents authentication-related errors.
39type AuthenticationError struct {
40 // Type is the type of authentication error.
41 Type string `json:"type"`
42 // Message is a human-readable message describing the error.
43 Message string `json:"message"`
44 // Code is the HTTP status code associated with the error.
45 Code int `json:"code"`
46 // Cause is the underlying error that caused this authentication error.
47 Cause error `json:"-"`
48}
49
50// Error returns a string representation of the authentication error.
51func (e *AuthenticationError) Error() string {
52 if e.Cause != nil {
53 return fmt.Sprintf("%s: %s (caused by: %v)", e.Type, e.Message, e.Cause)
54 }
55 return fmt.Sprintf("%s: %s", e.Type, e.Message)
56}
57
58// Unwrap returns the underlying cause of the error.
59func (e *AuthenticationError) Unwrap() error {
60 return e.Cause
61}
62
63// Common authentication error types for GitHub Copilot device flow.
64var (
65 // ErrDeviceCodeFailed represents an error when requesting the device code fails.
66 ErrDeviceCodeFailed = &AuthenticationError{
67 Type: "device_code_failed",
68 Message: "Failed to request device code from GitHub",
69 Code: http.StatusBadRequest,
70 }
71
72 // ErrDeviceCodeExpired represents an error when the device code has expired.
73 ErrDeviceCodeExpired = &AuthenticationError{
74 Type: "device_code_expired",
75 Message: "Device code has expired. Please try again.",
76 Code: http.StatusGone,
77 }
78
79 // ErrAuthorizationPending represents a pending authorization state (not an error, used for polling).
80 ErrAuthorizationPending = &AuthenticationError{
81 Type: "authorization_pending",
82 Message: "Authorization is pending. Waiting for user to authorize.",
83 Code: http.StatusAccepted,
84 }
85
86 // ErrSlowDown represents a request to slow down polling.
87 ErrSlowDown = &AuthenticationError{
88 Type: "slow_down",
89 Message: "Polling too frequently. Slowing down.",
90 Code: http.StatusTooManyRequests,
91 }
92
93 // ErrAccessDenied represents an error when the user denies authorization.
94 ErrAccessDenied = &AuthenticationError{
95 Type: "access_denied",
96 Message: "User denied authorization",
97 Code: http.StatusForbidden,
98 }
99
100 // ErrTokenExchangeFailed represents an error when token exchange fails.
101 ErrTokenExchangeFailed = &AuthenticationError{
102 Type: "token_exchange_failed",
103 Message: "Failed to exchange device code for access token",
104 Code: http.StatusBadRequest,
105 }
106
107 // ErrPollingTimeout represents an error when polling times out.
108 ErrPollingTimeout = &AuthenticationError{
109 Type: "polling_timeout",
110 Message: "Timeout waiting for user authorization",
111 Code: http.StatusRequestTimeout,
112 }
113
114 // ErrUserInfoFailed represents an error when fetching user info fails.
115 ErrUserInfoFailed = &AuthenticationError{
116 Type: "user_info_failed",
117 Message: "Failed to fetch GitHub user information",
118 Code: http.StatusBadRequest,
119 }
120)
121
122// NewAuthenticationError creates a new authentication error with a cause based on a base error.
123func NewAuthenticationError(baseErr *AuthenticationError, cause error) *AuthenticationError {
124 return &AuthenticationError{
125 Type: baseErr.Type,
126 Message: baseErr.Message,
127 Code: baseErr.Code,
128 Cause: cause,
129 }
130}
131
132// IsAuthenticationError checks if an error is an authentication error.
133func IsAuthenticationError(err error) bool {
134 var authenticationError *AuthenticationError
135 ok := errors.As(err, &authenticationError)
136 return ok
137}
138
139// IsOAuthError checks if an error is an OAuth error.
140func IsOAuthError(err error) bool {
141 var oAuthError *OAuthError
142 ok := errors.As(err, &oAuthError)
143 return ok
144}
145
146// GetUserFriendlyMessage returns a user-friendly error message based on the error type.
147func GetUserFriendlyMessage(err error) string {
148 var authErr *AuthenticationError
149 if errors.As(err, &authErr) {
150 switch authErr.Type {
151 case "device_code_failed":
152 return "Failed to start GitHub authentication. Please check your network connection and try again."
153 case "device_code_expired":
154 return "The authentication code has expired. Please try again."
155 case "authorization_pending":
156 return "Waiting for you to authorize the application on GitHub."
157 case "slow_down":
158 return "Please wait a moment before trying again."
159 case "access_denied":
160 return "Authentication was cancelled or denied."
161 case "token_exchange_failed":
162 return "Failed to complete authentication. Please try again."
163 case "polling_timeout":
164 return "Authentication timed out. Please try again."
165 case "user_info_failed":
166 return "Failed to get your GitHub account information. Please try again."
167 default:
168 return "Authentication failed. Please try again."
169 }
170 }
171
172 var oauthErr *OAuthError
173 if errors.As(err, &oauthErr) {
174 switch oauthErr.Code {
175 case "access_denied":
176 return "Authentication was cancelled or denied."
177 case "invalid_request":
178 return "Invalid authentication request. Please try again."
179 case "server_error":
180 return "GitHub server error. Please try again later."
181 default:
182 return fmt.Sprintf("Authentication failed: %s", oauthErr.Description)
183 }
184 }
185
186 return "An unexpected error occurred. Please try again."
187}
internal/auth/copilot/oauth.go+255−0
@@ -0,0 +1,255 @@
1package copilot
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "io"
9 "net/http"
10 "net/url"
11 "strings"
12 "time"
13
14 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
15 "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
16 log "github.com/sirupsen/logrus"
17)
18
19const (
20 // copilotClientID is GitHub's Copilot CLI OAuth client ID.
21 copilotClientID = "Iv1.b507a08c87ecfe98"
22 // copilotDeviceCodeURL is the endpoint for requesting device codes.
23 copilotDeviceCodeURL = "https://github.com/login/device/code"
24 // copilotTokenURL is the endpoint for exchanging device codes for tokens.
25 copilotTokenURL = "https://github.com/login/oauth/access_token"
26 // copilotUserInfoURL is the endpoint for fetching GitHub user information.
27 copilotUserInfoURL = "https://api.github.com/user"
28 // defaultPollInterval is the default interval for polling token endpoint.
29 defaultPollInterval = 5 * time.Second
30 // maxPollDuration is the maximum time to wait for user authorization.
31 maxPollDuration = 15 * time.Minute
32)
33
34// DeviceFlowClient handles the OAuth2 device flow for GitHub Copilot.
35type DeviceFlowClient struct {
36 httpClient *http.Client
37 cfg *config.Config
38}
39
40// NewDeviceFlowClient creates a new device flow client.
41func NewDeviceFlowClient(cfg *config.Config) *DeviceFlowClient {
42 client := &http.Client{Timeout: 30 * time.Second}
43 if cfg != nil {
44 client = util.SetProxy(&cfg.SDKConfig, client)
45 }
46 return &DeviceFlowClient{
47 httpClient: client,
48 cfg: cfg,
49 }
50}
51
52// RequestDeviceCode initiates the device flow by requesting a device code from GitHub.
53func (c *DeviceFlowClient) RequestDeviceCode(ctx context.Context) (*DeviceCodeResponse, error) {
54 data := url.Values{}
55 data.Set("client_id", copilotClientID)
56 data.Set("scope", "user:email")
57
58 req, err := http.NewRequestWithContext(ctx, http.MethodPost, copilotDeviceCodeURL, strings.NewReader(data.Encode()))
59 if err != nil {
60 return nil, NewAuthenticationError(ErrDeviceCodeFailed, err)
61 }
62 req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
63 req.Header.Set("Accept", "application/json")
64
65 resp, err := c.httpClient.Do(req)
66 if err != nil {
67 return nil, NewAuthenticationError(ErrDeviceCodeFailed, err)
68 }
69 defer func() {
70 if errClose := resp.Body.Close(); errClose != nil {
71 log.Errorf("copilot device code: close body error: %v", errClose)
72 }
73 }()
74
75 if !isHTTPSuccess(resp.StatusCode) {
76 bodyBytes, _ := io.ReadAll(resp.Body)
77 return nil, NewAuthenticationError(ErrDeviceCodeFailed, fmt.Errorf("status %d: %s", resp.StatusCode, string(bodyBytes)))
78 }
79
80 var deviceCode DeviceCodeResponse
81 if err = json.NewDecoder(resp.Body).Decode(&deviceCode); err != nil {
82 return nil, NewAuthenticationError(ErrDeviceCodeFailed, err)
83 }
84
85 return &deviceCode, nil
86}
87
88// PollForToken polls the token endpoint until the user authorizes or the device code expires.
89func (c *DeviceFlowClient) PollForToken(ctx context.Context, deviceCode *DeviceCodeResponse) (*CopilotTokenData, error) {
90 if deviceCode == nil {
91 return nil, NewAuthenticationError(ErrTokenExchangeFailed, fmt.Errorf("device code is nil"))
92 }
93
94 interval := time.Duration(deviceCode.Interval) * time.Second
95 if interval < defaultPollInterval {
96 interval = defaultPollInterval
97 }
98
99 deadline := time.Now().Add(maxPollDuration)
100 if deviceCode.ExpiresIn > 0 {
101 codeDeadline := time.Now().Add(time.Duration(deviceCode.ExpiresIn) * time.Second)
102 if codeDeadline.Before(deadline) {
103 deadline = codeDeadline
104 }
105 }
106
107 ticker := time.NewTicker(interval)
108 defer ticker.Stop()
109
110 for {
111 select {
112 case <-ctx.Done():
113 return nil, NewAuthenticationError(ErrPollingTimeout, ctx.Err())
114 case <-ticker.C:
115 if time.Now().After(deadline) {
116 return nil, ErrPollingTimeout
117 }
118
119 token, err := c.exchangeDeviceCode(ctx, deviceCode.DeviceCode)
120 if err != nil {
121 var authErr *AuthenticationError
122 if errors.As(err, &authErr) {
123 switch authErr.Type {
124 case ErrAuthorizationPending.Type:
125 // Continue polling
126 continue
127 case ErrSlowDown.Type:
128 // Increase interval and continue
129 interval += 5 * time.Second
130 ticker.Reset(interval)
131 continue
132 case ErrDeviceCodeExpired.Type:
133 return nil, err
134 case ErrAccessDenied.Type:
135 return nil, err
136 }
137 }
138 return nil, err
139 }
140 return token, nil
141 }
142 }
143}
144
145// exchangeDeviceCode attempts to exchange the device code for an access token.
146func (c *DeviceFlowClient) exchangeDeviceCode(ctx context.Context, deviceCode string) (*CopilotTokenData, error) {
147 data := url.Values{}
148 data.Set("client_id", copilotClientID)
149 data.Set("device_code", deviceCode)
150 data.Set("grant_type", "urn:ietf:params:oauth:grant-type:device_code")
151
152 req, err := http.NewRequestWithContext(ctx, http.MethodPost, copilotTokenURL, strings.NewReader(data.Encode()))
153 if err != nil {
154 return nil, NewAuthenticationError(ErrTokenExchangeFailed, err)
155 }
156 req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
157 req.Header.Set("Accept", "application/json")
158
159 resp, err := c.httpClient.Do(req)
160 if err != nil {
161 return nil, NewAuthenticationError(ErrTokenExchangeFailed, err)
162 }
163 defer func() {
164 if errClose := resp.Body.Close(); errClose != nil {
165 log.Errorf("copilot token exchange: close body error: %v", errClose)
166 }
167 }()
168
169 bodyBytes, err := io.ReadAll(resp.Body)
170 if err != nil {
171 return nil, NewAuthenticationError(ErrTokenExchangeFailed, err)
172 }
173
174 // GitHub returns 200 for both success and error cases in device flow
175 // Check for OAuth error response first
176 var oauthResp struct {
177 Error string `json:"error"`
178 ErrorDescription string `json:"error_description"`
179 AccessToken string `json:"access_token"`
180 TokenType string `json:"token_type"`
181 Scope string `json:"scope"`
182 }
183
184 if err = json.Unmarshal(bodyBytes, &oauthResp); err != nil {
185 return nil, NewAuthenticationError(ErrTokenExchangeFailed, err)
186 }
187
188 if oauthResp.Error != "" {
189 switch oauthResp.Error {
190 case "authorization_pending":
191 return nil, ErrAuthorizationPending
192 case "slow_down":
193 return nil, ErrSlowDown
194 case "expired_token":
195 return nil, ErrDeviceCodeExpired
196 case "access_denied":
197 return nil, ErrAccessDenied
198 default:
199 return nil, NewOAuthError(oauthResp.Error, oauthResp.ErrorDescription, resp.StatusCode)
200 }
201 }
202
203 if oauthResp.AccessToken == "" {
204 return nil, NewAuthenticationError(ErrTokenExchangeFailed, fmt.Errorf("empty access token"))
205 }
206
207 return &CopilotTokenData{
208 AccessToken: oauthResp.AccessToken,
209 TokenType: oauthResp.TokenType,
210 Scope: oauthResp.Scope,
211 }, nil
212}
213
214// FetchUserInfo retrieves the GitHub username for the authenticated user.
215func (c *DeviceFlowClient) FetchUserInfo(ctx context.Context, accessToken string) (string, error) {
216 if accessToken == "" {
217 return "", NewAuthenticationError(ErrUserInfoFailed, fmt.Errorf("access token is empty"))
218 }
219
220 req, err := http.NewRequestWithContext(ctx, http.MethodGet, copilotUserInfoURL, nil)
221 if err != nil {
222 return "", NewAuthenticationError(ErrUserInfoFailed, err)
223 }
224 req.Header.Set("Authorization", "Bearer "+accessToken)
225 req.Header.Set("Accept", "application/json")
226 req.Header.Set("User-Agent", "CLIProxyAPI")
227
228 resp, err := c.httpClient.Do(req)
229 if err != nil {
230 return "", NewAuthenticationError(ErrUserInfoFailed, err)
231 }
232 defer func() {
233 if errClose := resp.Body.Close(); errClose != nil {
234 log.Errorf("copilot user info: close body error: %v", errClose)
235 }
236 }()
237
238 if !isHTTPSuccess(resp.StatusCode) {
239 bodyBytes, _ := io.ReadAll(resp.Body)
240 return "", NewAuthenticationError(ErrUserInfoFailed, fmt.Errorf("status %d: %s", resp.StatusCode, string(bodyBytes)))
241 }
242
243 var userInfo struct {
244 Login string `json:"login"`
245 }
246 if err = json.NewDecoder(resp.Body).Decode(&userInfo); err != nil {
247 return "", NewAuthenticationError(ErrUserInfoFailed, err)
248 }
249
250 if userInfo.Login == "" {
251 return "", NewAuthenticationError(ErrUserInfoFailed, fmt.Errorf("empty username"))
252 }
253
254 return userInfo.Login, nil
255}
internal/auth/copilot/token.go+93−0
@@ -0,0 +1,93 @@
1// Package copilot provides authentication and token management functionality
2// for GitHub Copilot AI services. It handles OAuth2 device flow token storage,
3// serialization, and retrieval for maintaining authenticated sessions with the Copilot API.
4package copilot
5
6import (
7 "encoding/json"
8 "fmt"
9 "os"
10 "path/filepath"
11
12 "github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
13)
14
15// CopilotTokenStorage stores OAuth2 token information for GitHub Copilot API authentication.
16// It maintains compatibility with the existing auth system while adding Copilot-specific fields
17// for managing access tokens and user account information.
18type CopilotTokenStorage struct {
19 // AccessToken is the OAuth2 access token used for authenticating API requests.
20 AccessToken string `json:"access_token"`
21 // TokenType is the type of token, typically "bearer".
22 TokenType string `json:"token_type"`
23 // Scope is the OAuth2 scope granted to the token.
24 Scope string `json:"scope"`
25 // ExpiresAt is the timestamp when the access token expires (if provided).
26 ExpiresAt string `json:"expires_at,omitempty"`
27 // Username is the GitHub username associated with this token.
28 Username string `json:"username"`
29 // Type indicates the authentication provider type, always "github-copilot" for this storage.
30 Type string `json:"type"`
31}
32
33// CopilotTokenData holds the raw OAuth token response from GitHub.
34type CopilotTokenData struct {
35 // AccessToken is the OAuth2 access token.
36 AccessToken string `json:"access_token"`
37 // TokenType is the type of token, typically "bearer".
38 TokenType string `json:"token_type"`
39 // Scope is the OAuth2 scope granted to the token.
40 Scope string `json:"scope"`
41}
42
43// CopilotAuthBundle bundles authentication data for storage.
44type CopilotAuthBundle struct {
45 // TokenData contains the OAuth token information.
46 TokenData *CopilotTokenData
47 // Username is the GitHub username.
48 Username string
49}
50
51// DeviceCodeResponse represents GitHub's device code response.
52type DeviceCodeResponse struct {
53 // DeviceCode is the device verification code.
54 DeviceCode string `json:"device_code"`
55 // UserCode is the code the user must enter at the verification URI.
56 UserCode string `json:"user_code"`
57 // VerificationURI is the URL where the user should enter the code.
58 VerificationURI string `json:"verification_uri"`
59 // ExpiresIn is the number of seconds until the device code expires.
60 ExpiresIn int `json:"expires_in"`
61 // Interval is the minimum number of seconds to wait between polling requests.
62 Interval int `json:"interval"`
63}
64
65// SaveTokenToFile serializes the Copilot token storage to a JSON file.
66// This method creates the necessary directory structure and writes the token
67// data in JSON format to the specified file path for persistent storage.
68//
69// Parameters:
70// - authFilePath: The full path where the token file should be saved
71//
72// Returns:
73// - error: An error if the operation fails, nil otherwise
74func (ts *CopilotTokenStorage) SaveTokenToFile(authFilePath string) error {
75 misc.LogSavingCredentials(authFilePath)
76 ts.Type = "github-copilot"
77 if err := os.MkdirAll(filepath.Dir(authFilePath), 0700); err != nil {
78 return fmt.Errorf("failed to create directory: %v", err)
79 }
80
81 f, err := os.Create(authFilePath)
82 if err != nil {
83 return fmt.Errorf("failed to create token file: %w", err)
84 }
85 defer func() {
86 _ = f.Close()
87 }()
88
89 if err = json.NewEncoder(f).Encode(ts); err != nil {
90 return fmt.Errorf("failed to write token to file: %w", err)
91 }
92 return nil
93}
internal/auth/iflow/iflow_auth.go+17−5
@@ -9,6 +9,7 @@ import ( · 8 unmodified lines
9 "io"9 "io"
10 "net/http"10 "net/http"
11 "net/url"11 "net/url"
12 "os"
12 "strings"13 "strings"
13 "time"14 "time"
14 15
@@ -28,10 +29,21 @@ const ( · 13 unmodified lines
28 iFlowAPIKeyEndpoint = "https://platform.iflow.cn/api/openapi/apikey"29 iFlowAPIKeyEndpoint = "https://platform.iflow.cn/api/openapi/apikey"
29 30
30 // Client credentials provided by iFlow for the Code Assist integration.31 // Client credentials provided by iFlow for the Code Assist integration.
31 iFlowOAuthClientID = "10009311001"32 iFlowOAuthClientID = "10009311001"
32 iFlowOAuthClientSecret = "4Z3YjXycVsQvyGF1etiNlIBB4RsqSDtW"33 // Default client secret (can be overridden via IFLOW_CLIENT_SECRET env var)
34 defaultIFlowClientSecret = "4Z3YjXycVsQvyGF1etiNlIBB4RsqSDtW"
33)35)
34 36
37// getIFlowClientSecret returns the iFlow OAuth client secret.
38// It first checks the IFLOW_CLIENT_SECRET environment variable,
39// falling back to the default value if not set.
40func getIFlowClientSecret() string {
41 if secret := os.Getenv("IFLOW_CLIENT_SECRET"); secret != "" {
42 return secret
43 }
44 return defaultIFlowClientSecret
45}
46
35// DefaultAPIBaseURL is the canonical chat completions endpoint.47// DefaultAPIBaseURL is the canonical chat completions endpoint.
36const DefaultAPIBaseURL = "https://apis.iflow.cn/v1"48const DefaultAPIBaseURL = "https://apis.iflow.cn/v1"
37 49
@@ -72,7 +84,7 @@ func (ia *IFlowAuth) ExchangeCodeForTokens(ctx context.Context, code, redirectUR · 34 unmodified lines
72 form.Set("code", code)84 form.Set("code", code)
73 form.Set("redirect_uri", redirectURI)85 form.Set("redirect_uri", redirectURI)
74 form.Set("client_id", iFlowOAuthClientID)86 form.Set("client_id", iFlowOAuthClientID)
75 form.Set("client_secret", iFlowOAuthClientSecret)87 form.Set("client_secret", getIFlowClientSecret())
76 88
77 req, err := ia.newTokenRequest(ctx, form)89 req, err := ia.newTokenRequest(ctx, form)
78 if err != nil {90 if err != nil {
@@ -88,7 +100,7 @@ func (ia *IFlowAuth) RefreshTokens(ctx context.Context, refreshToken string) (*I · 9 unmodified lines
88 form.Set("grant_type", "refresh_token")100 form.Set("grant_type", "refresh_token")
89 form.Set("refresh_token", refreshToken)101 form.Set("refresh_token", refreshToken)
90 form.Set("client_id", iFlowOAuthClientID)102 form.Set("client_id", iFlowOAuthClientID)
91 form.Set("client_secret", iFlowOAuthClientSecret)103 form.Set("client_secret", getIFlowClientSecret())
92 104
93 req, err := ia.newTokenRequest(ctx, form)105 req, err := ia.newTokenRequest(ctx, form)
94 if err != nil {106 if err != nil {
@@ -104,7 +116,7 @@ func (ia *IFlowAuth) newTokenRequest(ctx context.Context, form url.Values) (*htt · 9 unmodified lines
104 return nil, fmt.Errorf("iflow token: create request failed: %w", err)116 return nil, fmt.Errorf("iflow token: create request failed: %w", err)
105 }117 }
106 118
107 basic := base64.StdEncoding.EncodeToString([]byte(iFlowOAuthClientID + ":" + iFlowOAuthClientSecret))119 basic := base64.StdEncoding.EncodeToString([]byte(iFlowOAuthClientID + ":" + getIFlowClientSecret()))
108 req.Header.Set("Content-Type", "application/x-www-form-urlencoded")120 req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
109 req.Header.Set("Accept", "application/json")121 req.Header.Set("Accept", "application/json")
110 req.Header.Set("Authorization", "Basic "+basic)122 req.Header.Set("Authorization", "Basic "+basic)
internal/auth/kiro/aws.go+305−0
@@ -0,0 +1,305 @@
1// Package kiro provides authentication functionality for AWS CodeWhisperer (Kiro) API.
2// It includes interfaces and implementations for token storage and authentication methods.
3package kiro
4
5import (
6 "encoding/base64"
7 "encoding/json"
8 "fmt"
9 "os"
10 "path/filepath"
11 "strings"
12)
13
14// PKCECodes holds PKCE verification codes for OAuth2 PKCE flow
15type PKCECodes struct {
16 // CodeVerifier is the cryptographically random string used to correlate
17 // the authorization request to the token request
18 CodeVerifier string `json:"code_verifier"`
19 // CodeChallenge is the SHA256 hash of the code verifier, base64url-encoded
20 CodeChallenge string `json:"code_challenge"`
21}
22
23// KiroTokenData holds OAuth token information from AWS CodeWhisperer (Kiro)
24type KiroTokenData struct {
25 // AccessToken is the OAuth2 access token for API access
26 AccessToken string `json:"accessToken"`
27 // RefreshToken is used to obtain new access tokens
28 RefreshToken string `json:"refreshToken"`
29 // ProfileArn is the AWS CodeWhisperer profile ARN
30 ProfileArn string `json:"profileArn"`
31 // ExpiresAt is the timestamp when the token expires
32 ExpiresAt string `json:"expiresAt"`
33 // AuthMethod indicates the authentication method used (e.g., "builder-id", "social")
34 AuthMethod string `json:"authMethod"`
35 // Provider indicates the OAuth provider (e.g., "AWS", "Google")
36 Provider string `json:"provider"`
37 // ClientID is the OIDC client ID (needed for token refresh)
38 ClientID string `json:"clientId,omitempty"`
39 // ClientSecret is the OIDC client secret (needed for token refresh)
40 ClientSecret string `json:"clientSecret,omitempty"`
41 // Email is the user's email address (used for file naming)
42 Email string `json:"email,omitempty"`
43 // StartURL is the IDC/Identity Center start URL (only for IDC auth method)
44 StartURL string `json:"startUrl,omitempty"`
45 // Region is the AWS region for IDC authentication (only for IDC auth method)
46 Region string `json:"region,omitempty"`
47}
48
49// KiroAuthBundle aggregates authentication data after OAuth flow completion
50type KiroAuthBundle struct {
51 // TokenData contains the OAuth tokens from the authentication flow
52 TokenData KiroTokenData `json:"token_data"`
53 // LastRefresh is the timestamp of the last token refresh
54 LastRefresh string `json:"last_refresh"`
55}
56
57// KiroUsageInfo represents usage information from CodeWhisperer API
58type KiroUsageInfo struct {
59 // SubscriptionTitle is the subscription plan name (e.g., "KIRO FREE")
60 SubscriptionTitle string `json:"subscription_title"`
61 // CurrentUsage is the current credit usage
62 CurrentUsage float64 `json:"current_usage"`
63 // UsageLimit is the maximum credit limit
64 UsageLimit float64 `json:"usage_limit"`
65 // NextReset is the timestamp of the next usage reset
66 NextReset string `json:"next_reset"`
67}
68
69// KiroModel represents a model available through the CodeWhisperer API
70type KiroModel struct {
71 // ModelID is the unique identifier for the model
72 ModelID string `json:"modelId"`
73 // ModelName is the human-readable name
74 ModelName string `json:"modelName"`
75 // Description is the model description
76 Description string `json:"description"`
77 // RateMultiplier is the credit multiplier for this model
78 RateMultiplier float64 `json:"rateMultiplier"`
79 // RateUnit is the unit for rate calculation (e.g., "credit")
80 RateUnit string `json:"rateUnit"`
81 // MaxInputTokens is the maximum input token limit
82 MaxInputTokens int `json:"maxInputTokens,omitempty"`
83}
84
85// KiroIDETokenFile is the default path to Kiro IDE's token file
86const KiroIDETokenFile = ".aws/sso/cache/kiro-auth-token.json"
87
88// LoadKiroIDEToken loads token data from Kiro IDE's token file.
89func LoadKiroIDEToken() (*KiroTokenData, error) {
90 homeDir, err := os.UserHomeDir()
91 if err != nil {
92 return nil, fmt.Errorf("failed to get home directory: %w", err)
93 }
94
95 tokenPath := filepath.Join(homeDir, KiroIDETokenFile)
96 data, err := os.ReadFile(tokenPath)
97 if err != nil {
98 return nil, fmt.Errorf("failed to read Kiro IDE token file (%s): %w", tokenPath, err)
99 }
100
101 var token KiroTokenData
102 if err := json.Unmarshal(data, &token); err != nil {
103 return nil, fmt.Errorf("failed to parse Kiro IDE token: %w", err)
104 }
105
106 if token.AccessToken == "" {
107 return nil, fmt.Errorf("access token is empty in Kiro IDE token file")
108 }
109
110 return &token, nil
111}
112
113// LoadKiroTokenFromPath loads token data from a custom path.
114// This supports multiple accounts by allowing different token files.
115func LoadKiroTokenFromPath(tokenPath string) (*KiroTokenData, error) {
116 // Expand ~ to home directory
117 if len(tokenPath) > 0 && tokenPath[0] == '~' {
118 homeDir, err := os.UserHomeDir()
119 if err != nil {
120 return nil, fmt.Errorf("failed to get home directory: %w", err)
121 }
122 tokenPath = filepath.Join(homeDir, tokenPath[1:])
123 }
124
125 data, err := os.ReadFile(tokenPath)
126 if err != nil {
127 return nil, fmt.Errorf("failed to read token file (%s): %w", tokenPath, err)
128 }
129
130 var token KiroTokenData
131 if err := json.Unmarshal(data, &token); err != nil {
132 return nil, fmt.Errorf("failed to parse token file: %w", err)
133 }
134
135 if token.AccessToken == "" {
136 return nil, fmt.Errorf("access token is empty in token file")
137 }
138
139 return &token, nil
140}
141
142// ListKiroTokenFiles lists all Kiro token files in the cache directory.
143// This supports multiple accounts by finding all token files.
144func ListKiroTokenFiles() ([]string, error) {
145 homeDir, err := os.UserHomeDir()
146 if err != nil {
147 return nil, fmt.Errorf("failed to get home directory: %w", err)
148 }
149
150 cacheDir := filepath.Join(homeDir, ".aws", "sso", "cache")
151
152 // Check if directory exists
153 if _, err := os.Stat(cacheDir); os.IsNotExist(err) {
154 return nil, nil // No token files
155 }
156
157 entries, err := os.ReadDir(cacheDir)
158 if err != nil {
159 return nil, fmt.Errorf("failed to read cache directory: %w", err)
160 }
161
162 var tokenFiles []string
163 for _, entry := range entries {
164 if entry.IsDir() {
165 continue
166 }
167 name := entry.Name()
168 // Look for kiro token files only (avoid matching unrelated AWS SSO cache files)
169 if strings.HasSuffix(name, ".json") && strings.HasPrefix(name, "kiro") {
170 tokenFiles = append(tokenFiles, filepath.Join(cacheDir, name))
171 }
172 }
173
174 return tokenFiles, nil
175}
176
177// LoadAllKiroTokens loads all Kiro tokens from the cache directory.
178// This supports multiple accounts.
179func LoadAllKiroTokens() ([]*KiroTokenData, error) {
180 files, err := ListKiroTokenFiles()
181 if err != nil {
182 return nil, err
183 }
184
185 var tokens []*KiroTokenData
186 for _, file := range files {
187 token, err := LoadKiroTokenFromPath(file)
188 if err != nil {
189 // Skip invalid token files
190 continue
191 }
192 tokens = append(tokens, token)
193 }
194
195 return tokens, nil
196}
197
198// JWTClaims represents the claims we care about from a JWT token.
199// JWT tokens from Kiro/AWS contain user information in the payload.
200type JWTClaims struct {
201 Email string `json:"email,omitempty"`
202 Sub string `json:"sub,omitempty"`
203 PreferredUser string `json:"preferred_username,omitempty"`
204 Name string `json:"name,omitempty"`
205 Iss string `json:"iss,omitempty"`
206}
207
208// ExtractEmailFromJWT extracts the user's email from a JWT access token.
209// JWT tokens typically have format: header.payload.signature
210// The payload is base64url-encoded JSON containing user claims.
211func ExtractEmailFromJWT(accessToken string) string {
212 if accessToken == "" {
213 return ""
214 }
215
216 // JWT format: header.payload.signature
217 parts := strings.Split(accessToken, ".")
218 if len(parts) != 3 {
219 return ""
220 }
221
222 // Decode the payload (second part)
223 payload := parts[1]
224
225 // Add padding if needed (base64url requires padding)
226 switch len(payload) % 4 {
227 case 2:
228 payload += "=="
229 case 3:
230 payload += "="
231 }
232
233 decoded, err := base64.URLEncoding.DecodeString(payload)
234 if err != nil {
235 // Try RawURLEncoding (no padding)
236 decoded, err = base64.RawURLEncoding.DecodeString(parts[1])
237 if err != nil {
238 return ""
239 }
240 }
241
242 var claims JWTClaims
243 if err := json.Unmarshal(decoded, &claims); err != nil {
244 return ""
245 }
246
247 // Return email if available
248 if claims.Email != "" {
249 return claims.Email
250 }
251
252 // Fallback to preferred_username (some providers use this)
253 if claims.PreferredUser != "" && strings.Contains(claims.PreferredUser, "@") {
254 return claims.PreferredUser
255 }
256
257 // Fallback to sub if it looks like an email
258 if claims.Sub != "" && strings.Contains(claims.Sub, "@") {
259 return claims.Sub
260 }
261
262 return ""
263}
264
265// SanitizeEmailForFilename sanitizes an email address for use in a filename.
266// Replaces special characters with underscores and prevents path traversal attacks.
267// Also handles URL-encoded characters to prevent encoded path traversal attempts.
268func SanitizeEmailForFilename(email string) string {
269 if email == "" {
270 return ""
271 }
272
273 result := email
274
275 // First, handle URL-encoded path traversal attempts (%2F, %2E, %5C, etc.)
276 // This prevents encoded characters from bypassing the sanitization.
277 // Note: We replace % last to catch any remaining encodings including double-encoding (%252F)
278 result = strings.ReplaceAll(result, "%2F", "_") // /
279 result = strings.ReplaceAll(result, "%2f", "_")
280 result = strings.ReplaceAll(result, "%5C", "_") // \
281 result = strings.ReplaceAll(result, "%5c", "_")
282 result = strings.ReplaceAll(result, "%2E", "_") // .
283 result = strings.ReplaceAll(result, "%2e", "_")
284 result = strings.ReplaceAll(result, "%00", "_") // null byte
285 result = strings.ReplaceAll(result, "%", "_") // Catch remaining % to prevent double-encoding attacks
286
287 // Replace characters that are problematic in filenames
288 // Keep @ and . in middle but replace other special characters
289 for _, char := range []string{"/", "\\", ":", "*", "?", "\"", "<", ">", "|", " ", "\x00"} {
290 result = strings.ReplaceAll(result, char, "_")
291 }
292
293 // Prevent path traversal: replace leading dots in each path component
294 // This handles cases like "../../../etc/passwd" → "_.._.._.._etc_passwd"
295 parts := strings.Split(result, "_")
296 for i, part := range parts {
297 for strings.HasPrefix(part, ".") {
298 part = "_" + part[1:]
299 }
300 parts[i] = part
301 }
302 result = strings.Join(parts, "_")
303
304 return result
305}
internal/auth/kiro/aws_auth.go+314−0
@@ -0,0 +1,314 @@
1// Package kiro provides OAuth2 authentication functionality for AWS CodeWhisperer (Kiro) API.
2// This package implements token loading, refresh, and API communication with CodeWhisperer.
3package kiro
4
5import (
6 "context"
7 "encoding/json"
8 "fmt"
9 "io"
10 "net/http"
11 "os"
12 "path/filepath"
13 "strings"
14 "time"
15
16 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
17 "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
18 log "github.com/sirupsen/logrus"
19)
20
21const (
22 // awsKiroEndpoint is used for CodeWhisperer management APIs (GetUsageLimits, ListProfiles, etc.)
23 // Note: This is different from the Amazon Q streaming endpoint (q.us-east-1.amazonaws.com)
24 // used in kiro_executor.go for GenerateAssistantResponse. Both endpoints are correct
25 // for their respective API operations.
26 awsKiroEndpoint = "https://codewhisperer.us-east-1.amazonaws.com"
27 defaultTokenFile = "~/.aws/sso/cache/kiro-auth-token.json"
28 targetGetUsage = "AmazonCodeWhispererService.GetUsageLimits"
29 targetListModels = "AmazonCodeWhispererService.ListAvailableModels"
30 targetGenerateChat = "AmazonCodeWhispererStreamingService.GenerateAssistantResponse"
31)
32
33// KiroAuth handles AWS CodeWhisperer authentication and API communication.
34// It provides methods for loading tokens, refreshing expired tokens,
35// and communicating with the CodeWhisperer API.
36type KiroAuth struct {
37 httpClient *http.Client
38 endpoint string
39}
40
41// NewKiroAuth creates a new Kiro authentication service.
42// It initializes the HTTP client with proxy settings from the configuration.
43//
44// Parameters:
45// - cfg: The application configuration containing proxy settings
46//
47// Returns:
48// - *KiroAuth: A new Kiro authentication service instance
49func NewKiroAuth(cfg *config.Config) *KiroAuth {
50 return &KiroAuth{
51 httpClient: util.SetProxy(&cfg.SDKConfig, &http.Client{Timeout: 120 * time.Second}),
52 endpoint: awsKiroEndpoint,
53 }
54}
55
56// LoadTokenFromFile loads token data from a file path.
57// This method reads and parses the token file, expanding ~ to the home directory.
58//
59// Parameters:
60// - tokenFile: Path to the token file (supports ~ expansion)
61//
62// Returns:
63// - *KiroTokenData: The parsed token data
64// - error: An error if file reading or parsing fails
65func (k *KiroAuth) LoadTokenFromFile(tokenFile string) (*KiroTokenData, error) {
66 // Expand ~ to home directory
67 if strings.HasPrefix(tokenFile, "~") {
68 home, err := os.UserHomeDir()
69 if err != nil {
70 return nil, fmt.Errorf("failed to get home directory: %w", err)
71 }
72 tokenFile = filepath.Join(home, tokenFile[1:])
73 }
74
75 data, err := os.ReadFile(tokenFile)
76 if err != nil {
77 return nil, fmt.Errorf("failed to read token file: %w", err)
78 }
79
80 var tokenData KiroTokenData
81 if err := json.Unmarshal(data, &tokenData); err != nil {
82 return nil, fmt.Errorf("failed to parse token file: %w", err)
83 }
84
85 return &tokenData, nil
86}
87
88// IsTokenExpired checks if the token has expired.
89// This method parses the expiration timestamp and compares it with the current time.
90//
91// Parameters:
92// - tokenData: The token data to check
93//
94// Returns:
95// - bool: True if the token has expired, false otherwise
96func (k *KiroAuth) IsTokenExpired(tokenData *KiroTokenData) bool {
97 if tokenData.ExpiresAt == "" {
98 return true
99 }
100
101 expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
102 if err != nil {
103 // Try alternate format
104 expiresAt, err = time.Parse("2006-01-02T15:04:05.000Z", tokenData.ExpiresAt)
105 if err != nil {
106 return true
107 }
108 }
109
110 return time.Now().After(expiresAt)
111}
112
113// makeRequest sends a request to the CodeWhisperer API.
114// This is an internal method for making authenticated API calls.
115//
116// Parameters:
117// - ctx: The context for the request
118// - target: The API target (e.g., "AmazonCodeWhispererService.GetUsageLimits")
119// - accessToken: The OAuth access token
120// - payload: The request payload
121//
122// Returns:
123// - []byte: The response body
124// - error: An error if the request fails
125func (k *KiroAuth) makeRequest(ctx context.Context, target string, accessToken string, payload interface{}) ([]byte, error) {
126 jsonBody, err := json.Marshal(payload)
127 if err != nil {
128 return nil, fmt.Errorf("failed to marshal request: %w", err)
129 }
130
131 req, err := http.NewRequestWithContext(ctx, http.MethodPost, k.endpoint, strings.NewReader(string(jsonBody)))
132 if err != nil {
133 return nil, fmt.Errorf("failed to create request: %w", err)
134 }
135
136 req.Header.Set("Content-Type", "application/x-amz-json-1.0")
137 req.Header.Set("x-amz-target", target)
138 req.Header.Set("Authorization", "Bearer "+accessToken)
139 req.Header.Set("Accept", "application/json")
140
141 resp, err := k.httpClient.Do(req)
142 if err != nil {
143 return nil, fmt.Errorf("request failed: %w", err)
144 }
145 defer func() {
146 if errClose := resp.Body.Close(); errClose != nil {
147 log.Errorf("failed to close response body: %v", errClose)
148 }
149 }()
150
151 body, err := io.ReadAll(resp.Body)
152 if err != nil {
153 return nil, fmt.Errorf("failed to read response: %w", err)
154 }
155
156 if resp.StatusCode != http.StatusOK {
157 return nil, fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(body))
158 }
159
160 return body, nil
161}
162
163// GetUsageLimits retrieves usage information from the CodeWhisperer API.
164// This method fetches the current usage statistics and subscription information.
165//
166// Parameters:
167// - ctx: The context for the request
168// - tokenData: The token data containing access token and profile ARN
169//
170// Returns:
171// - *KiroUsageInfo: The usage information
172// - error: An error if the request fails
173func (k *KiroAuth) GetUsageLimits(ctx context.Context, tokenData *KiroTokenData) (*KiroUsageInfo, error) {
174 payload := map[string]interface{}{
175 "origin": "AI_EDITOR",
176 "profileArn": tokenData.ProfileArn,
177 "resourceType": "AGENTIC_REQUEST",
178 }
179
180 body, err := k.makeRequest(ctx, targetGetUsage, tokenData.AccessToken, payload)
181 if err != nil {
182 return nil, err
183 }
184
185 var result struct {
186 SubscriptionInfo struct {
187 SubscriptionTitle string `json:"subscriptionTitle"`
188 } `json:"subscriptionInfo"`
189 UsageBreakdownList []struct {
190 CurrentUsageWithPrecision float64 `json:"currentUsageWithPrecision"`
191 UsageLimitWithPrecision float64 `json:"usageLimitWithPrecision"`
192 } `json:"usageBreakdownList"`
193 NextDateReset float64 `json:"nextDateReset"`
194 }
195
196 if err := json.Unmarshal(body, &result); err != nil {
197 return nil, fmt.Errorf("failed to parse usage response: %w", err)
198 }
199
200 usage := &KiroUsageInfo{
201 SubscriptionTitle: result.SubscriptionInfo.SubscriptionTitle,
202 NextReset: fmt.Sprintf("%v", result.NextDateReset),
203 }
204
205 if len(result.UsageBreakdownList) > 0 {
206 usage.CurrentUsage = result.UsageBreakdownList[0].CurrentUsageWithPrecision
207 usage.UsageLimit = result.UsageBreakdownList[0].UsageLimitWithPrecision
208 }
209
210 return usage, nil
211}
212
213// ListAvailableModels retrieves available models from the CodeWhisperer API.
214// This method fetches the list of AI models available for the authenticated user.
215//
216// Parameters:
217// - ctx: The context for the request
218// - tokenData: The token data containing access token and profile ARN
219//
220// Returns:
221// - []*KiroModel: The list of available models
222// - error: An error if the request fails
223func (k *KiroAuth) ListAvailableModels(ctx context.Context, tokenData *KiroTokenData) ([]*KiroModel, error) {
224 payload := map[string]interface{}{
225 "origin": "AI_EDITOR",
226 "profileArn": tokenData.ProfileArn,
227 }
228
229 body, err := k.makeRequest(ctx, targetListModels, tokenData.AccessToken, payload)
230 if err != nil {
231 return nil, err
232 }
233
234 var result struct {
235 Models []struct {
236 ModelID string `json:"modelId"`
237 ModelName string `json:"modelName"`
238 Description string `json:"description"`
239 RateMultiplier float64 `json:"rateMultiplier"`
240 RateUnit string `json:"rateUnit"`
241 TokenLimits struct {
242 MaxInputTokens int `json:"maxInputTokens"`
243 } `json:"tokenLimits"`
244 } `json:"models"`
245 }
246
247 if err := json.Unmarshal(body, &result); err != nil {
248 return nil, fmt.Errorf("failed to parse models response: %w", err)
249 }
250
251 models := make([]*KiroModel, 0, len(result.Models))
252 for _, m := range result.Models {
253 models = append(models, &KiroModel{
254 ModelID: m.ModelID,
255 ModelName: m.ModelName,
256 Description: m.Description,
257 RateMultiplier: m.RateMultiplier,
258 RateUnit: m.RateUnit,
259 MaxInputTokens: m.TokenLimits.MaxInputTokens,
260 })
261 }
262
263 return models, nil
264}
265
266// CreateTokenStorage creates a new KiroTokenStorage from token data.
267// This method converts the token data into a storage structure suitable for persistence.
268//
269// Parameters:
270// - tokenData: The token data to convert
271//
272// Returns:
273// - *KiroTokenStorage: A new token storage instance
274func (k *KiroAuth) CreateTokenStorage(tokenData *KiroTokenData) *KiroTokenStorage {
275 return &KiroTokenStorage{
276 AccessToken: tokenData.AccessToken,
277 RefreshToken: tokenData.RefreshToken,
278 ProfileArn: tokenData.ProfileArn,
279 ExpiresAt: tokenData.ExpiresAt,
280 AuthMethod: tokenData.AuthMethod,
281 Provider: tokenData.Provider,
282 LastRefresh: time.Now().Format(time.RFC3339),
283 }
284}
285
286// ValidateToken checks if the token is valid by making a test API call.
287// This method verifies the token by attempting to fetch usage limits.
288//
289// Parameters:
290// - ctx: The context for the request
291// - tokenData: The token data to validate
292//
293// Returns:
294// - error: An error if the token is invalid
295func (k *KiroAuth) ValidateToken(ctx context.Context, tokenData *KiroTokenData) error {
296 _, err := k.GetUsageLimits(ctx, tokenData)
297 return err
298}
299
300// UpdateTokenStorage updates an existing token storage with new token data.
301// This method refreshes the token storage with newly obtained access and refresh tokens.
302//
303// Parameters:
304// - storage: The existing token storage to update
305// - tokenData: The new token data to apply
306func (k *KiroAuth) UpdateTokenStorage(storage *KiroTokenStorage, tokenData *KiroTokenData) {
307 storage.AccessToken = tokenData.AccessToken
308 storage.RefreshToken = tokenData.RefreshToken
309 storage.ProfileArn = tokenData.ProfileArn
310 storage.ExpiresAt = tokenData.ExpiresAt
311 storage.AuthMethod = tokenData.AuthMethod
312 storage.Provider = tokenData.Provider
313 storage.LastRefresh = time.Now().Format(time.RFC3339)
314}
internal/auth/kiro/aws_test.go+161−0
@@ -0,0 +1,161 @@
1package kiro
2
3import (
4 "encoding/base64"
5 "encoding/json"
6 "testing"
7)
8
9func TestExtractEmailFromJWT(t *testing.T) {
10 tests := []struct {
11 name string
12 token string
13 expected string
14 }{
15 {
16 name: "Empty token",
17 token: "",
18 expected: "",
19 },
20 {
21 name: "Invalid token format",
22 token: "not.a.valid.jwt",
23 expected: "",
24 },
25 {
26 name: "Invalid token - not base64",
27 token: "xxx.yyy.zzz",
28 expected: "",
29 },
30 {
31 name: "Valid JWT with email",
32 token: createTestJWT(map[string]any{"email": "test@example.com", "sub": "user123"}),
33 expected: "test@example.com",
34 },
35 {
36 name: "JWT without email but with preferred_username",
37 token: createTestJWT(map[string]any{"preferred_username": "user@domain.com", "sub": "user123"}),
38 expected: "user@domain.com",
39 },
40 {
41 name: "JWT with email-like sub",
42 token: createTestJWT(map[string]any{"sub": "another@test.com"}),
43 expected: "another@test.com",
44 },
45 {
46 name: "JWT without any email fields",
47 token: createTestJWT(map[string]any{"sub": "user123", "name": "Test User"}),
48 expected: "",
49 },
50 }
51
52 for _, tt := range tests {
53 t.Run(tt.name, func(t *testing.T) {
54 result := ExtractEmailFromJWT(tt.token)
55 if result != tt.expected {
56 t.Errorf("ExtractEmailFromJWT() = %q, want %q", result, tt.expected)
57 }
58 })
59 }
60}
61
62func TestSanitizeEmailForFilename(t *testing.T) {
63 tests := []struct {
64 name string
65 email string
66 expected string
67 }{
68 {
69 name: "Empty email",
70 email: "",
71 expected: "",
72 },
73 {
74 name: "Simple email",
75 email: "user@example.com",
76 expected: "user@example.com",
77 },
78 {
79 name: "Email with space",
80 email: "user name@example.com",
81 expected: "user_name@example.com",
82 },
83 {
84 name: "Email with special chars",
85 email: "user:name@example.com",
86 expected: "user_name@example.com",
87 },
88 {
89 name: "Email with multiple special chars",
90 email: "user/name:test@example.com",
91 expected: "user_name_test@example.com",
92 },
93 {
94 name: "Path traversal attempt",
95 email: "../../../etc/passwd",
96 expected: "_.__.__._etc_passwd",
97 },
98 {
99 name: "Path traversal with backslash",
100 email: `..\..\..\..\windows\system32`,
101 expected: "_.__.__.__._windows_system32",
102 },
103 {
104 name: "Null byte injection attempt",
105 email: "user\x00@evil.com",
106 expected: "user_@evil.com",
107 },
108 // URL-encoded path traversal tests
109 {
110 name: "URL-encoded slash",
111 email: "user%2Fpath@example.com",
112 expected: "user_path@example.com",
113 },
114 {
115 name: "URL-encoded backslash",
116 email: "user%5Cpath@example.com",
117 expected: "user_path@example.com",
118 },
119 {
120 name: "URL-encoded dot",
121 email: "%2E%2E%2Fetc%2Fpasswd",
122 expected: "___etc_passwd",
123 },
124 {
125 name: "URL-encoded null",
126 email: "user%00@evil.com",
127 expected: "user_@evil.com",
128 },
129 {
130 name: "Double URL-encoding attack",
131 email: "%252F%252E%252E",
132 expected: "_252F_252E_252E", // % replaced with _, remaining chars preserved (safe)
133 },
134 {
135 name: "Mixed case URL-encoding",
136 email: "%2f%2F%5c%5C",
137 expected: "____",
138 },
139 }
140
141 for _, tt := range tests {
142 t.Run(tt.name, func(t *testing.T) {
143 result := SanitizeEmailForFilename(tt.email)
144 if result != tt.expected {
145 t.Errorf("SanitizeEmailForFilename() = %q, want %q", result, tt.expected)
146 }
147 })
148 }
149}
150
151// createTestJWT creates a test JWT token with the given claims
152func createTestJWT(claims map[string]any) string {
153 header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"RS256","typ":"JWT"}`))
154
155 payloadBytes, _ := json.Marshal(claims)
156 payload := base64.RawURLEncoding.EncodeToString(payloadBytes)
157
158 signature := base64.RawURLEncoding.EncodeToString([]byte("fake-signature"))
159
160 return header + "." + payload + "." + signature
161}
internal/auth/kiro/codewhisperer_client.go+166−0
@@ -0,0 +1,166 @@
1// Package kiro provides CodeWhisperer API client for fetching user info.
2package kiro
3
4import (
5 "context"
6 "encoding/json"
7 "fmt"
8 "io"
9 "net/http"
10 "time"
11
12 "github.com/google/uuid"
13 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
14 "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
15 log "github.com/sirupsen/logrus"
16)
17
18const (
19 codeWhispererAPI = "https://codewhisperer.us-east-1.amazonaws.com"
20 kiroVersion = "0.6.18"
21)
22
23// CodeWhispererClient handles CodeWhisperer API calls.
24type CodeWhispererClient struct {
25 httpClient *http.Client
26 machineID string
27}
28
29// UsageLimitsResponse represents the getUsageLimits API response.
30type UsageLimitsResponse struct {
31 DaysUntilReset *int `json:"daysUntilReset,omitempty"`
32 NextDateReset *float64 `json:"nextDateReset,omitempty"`
33 UserInfo *UserInfo `json:"userInfo,omitempty"`
34 SubscriptionInfo *SubscriptionInfo `json:"subscriptionInfo,omitempty"`
35 UsageBreakdownList []UsageBreakdown `json:"usageBreakdownList,omitempty"`
36}
37
38// UserInfo contains user information from the API.
39type UserInfo struct {
40 Email string `json:"email,omitempty"`
41 UserID string `json:"userId,omitempty"`
42}
43
44// SubscriptionInfo contains subscription details.
45type SubscriptionInfo struct {
46 SubscriptionTitle string `json:"subscriptionTitle,omitempty"`
47 Type string `json:"type,omitempty"`
48}
49
50// UsageBreakdown contains usage details.
51type UsageBreakdown struct {
52 UsageLimit *int `json:"usageLimit,omitempty"`
53 CurrentUsage *int `json:"currentUsage,omitempty"`
54 UsageLimitWithPrecision *float64 `json:"usageLimitWithPrecision,omitempty"`
55 CurrentUsageWithPrecision *float64 `json:"currentUsageWithPrecision,omitempty"`
56 NextDateReset *float64 `json:"nextDateReset,omitempty"`
57 DisplayName string `json:"displayName,omitempty"`
58 ResourceType string `json:"resourceType,omitempty"`
59}
60
61// NewCodeWhispererClient creates a new CodeWhisperer client.
62func NewCodeWhispererClient(cfg *config.Config, machineID string) *CodeWhispererClient {
63 client := &http.Client{Timeout: 30 * time.Second}
64 if cfg != nil {
65 client = util.SetProxy(&cfg.SDKConfig, client)
66 }
67 if machineID == "" {
68 machineID = uuid.New().String()
69 }
70 return &CodeWhispererClient{
71 httpClient: client,
72 machineID: machineID,
73 }
74}
75
76// generateInvocationID generates a unique invocation ID.
77func generateInvocationID() string {
78 return uuid.New().String()
79}
80
81// GetUsageLimits fetches usage limits and user info from CodeWhisperer API.
82// This is the recommended way to get user email after login.
83func (c *CodeWhispererClient) GetUsageLimits(ctx context.Context, accessToken string) (*UsageLimitsResponse, error) {
84 url := fmt.Sprintf("%s/getUsageLimits?isEmailRequired=true&origin=AI_EDITOR&resourceType=AGENTIC_REQUEST", codeWhispererAPI)
85
86 req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
87 if err != nil {
88 return nil, fmt.Errorf("failed to create request: %w", err)
89 }
90
91 // Set headers to match Kiro IDE
92 xAmzUserAgent := fmt.Sprintf("aws-sdk-js/1.0.0 KiroIDE-%s-%s", kiroVersion, c.machineID)
93 userAgent := fmt.Sprintf("aws-sdk-js/1.0.0 ua/2.1 os/windows lang/js md/nodejs#20.16.0 api/codewhispererruntime#1.0.0 m/E KiroIDE-%s-%s", kiroVersion, c.machineID)
94
95 req.Header.Set("Authorization", "Bearer "+accessToken)
96 req.Header.Set("x-amz-user-agent", xAmzUserAgent)
97 req.Header.Set("User-Agent", userAgent)
98 req.Header.Set("amz-sdk-invocation-id", generateInvocationID())
99 req.Header.Set("amz-sdk-request", "attempt=1; max=1")
100 req.Header.Set("Connection", "close")
101
102 log.Debugf("codewhisperer: GET %s", url)
103
104 resp, err := c.httpClient.Do(req)
105 if err != nil {
106 return nil, fmt.Errorf("request failed: %w", err)
107 }
108 defer resp.Body.Close()
109
110 body, err := io.ReadAll(resp.Body)
111 if err != nil {
112 return nil, fmt.Errorf("failed to read response: %w", err)
113 }
114
115 log.Debugf("codewhisperer: status=%d, body=%s", resp.StatusCode, string(body))
116
117 if resp.StatusCode != http.StatusOK {
118 return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(body))
119 }
120
121 var result UsageLimitsResponse
122 if err := json.Unmarshal(body, &result); err != nil {
123 return nil, fmt.Errorf("failed to parse response: %w", err)
124 }
125
126 return &result, nil
127}
128
129// FetchUserEmailFromAPI fetches user email using CodeWhisperer getUsageLimits API.
130// This is more reliable than JWT parsing as it uses the official API.
131func (c *CodeWhispererClient) FetchUserEmailFromAPI(ctx context.Context, accessToken string) string {
132 resp, err := c.GetUsageLimits(ctx, accessToken)
133 if err != nil {
134 log.Debugf("codewhisperer: failed to get usage limits: %v", err)
135 return ""
136 }
137
138 if resp.UserInfo != nil && resp.UserInfo.Email != "" {
139 log.Debugf("codewhisperer: got email from API: %s", resp.UserInfo.Email)
140 return resp.UserInfo.Email
141 }
142
143 log.Debugf("codewhisperer: no email in response")
144 return ""
145}
146
147// FetchUserEmailWithFallback fetches user email with multiple fallback methods.
148// Priority: 1. CodeWhisperer API 2. userinfo endpoint 3. JWT parsing
149func FetchUserEmailWithFallback(ctx context.Context, cfg *config.Config, accessToken string) string {
150 // Method 1: Try CodeWhisperer API (most reliable)
151 cwClient := NewCodeWhispererClient(cfg, "")
152 email := cwClient.FetchUserEmailFromAPI(ctx, accessToken)
153 if email != "" {
154 return email
155 }
156
157 // Method 2: Try SSO OIDC userinfo endpoint
158 ssoClient := NewSSOOIDCClient(cfg)
159 email = ssoClient.FetchUserEmail(ctx, accessToken)
160 if email != "" {
161 return email
162 }
163
164 // Method 3: Fallback to JWT parsing
165 return ExtractEmailFromJWT(accessToken)
166}
internal/auth/kiro/oauth.go+303−0
@@ -0,0 +1,303 @@
1// Package kiro provides OAuth2 authentication for Kiro using native Google login.
2package kiro
3
4import (
5 "context"
6 "crypto/rand"
7 "crypto/sha256"
8 "encoding/base64"
9 "encoding/json"
10 "fmt"
11 "html"
12 "io"
13 "net"
14 "net/http"
15 "strings"
16 "time"
17
18 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
19 "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
20 log "github.com/sirupsen/logrus"
21)
22
23const (
24 // Kiro auth endpoint
25 kiroAuthEndpoint = "https://prod.us-east-1.auth.desktop.kiro.dev"
26
27 // Default callback port
28 defaultCallbackPort = 9876
29
30 // Auth timeout
31 authTimeout = 10 * time.Minute
32)
33
34// KiroTokenResponse represents the response from Kiro token endpoint.
35type KiroTokenResponse struct {
36 AccessToken string `json:"accessToken"`
37 RefreshToken string `json:"refreshToken"`
38 ProfileArn string `json:"profileArn"`
39 ExpiresIn int `json:"expiresIn"`
40}
41
42// KiroOAuth handles the OAuth flow for Kiro authentication.
43type KiroOAuth struct {
44 httpClient *http.Client
45 cfg *config.Config
46}
47
48// NewKiroOAuth creates a new Kiro OAuth handler.
49func NewKiroOAuth(cfg *config.Config) *KiroOAuth {
50 client := &http.Client{Timeout: 30 * time.Second}
51 if cfg != nil {
52 client = util.SetProxy(&cfg.SDKConfig, client)
53 }
54 return &KiroOAuth{
55 httpClient: client,
56 cfg: cfg,
57 }
58}
59
60// generateCodeVerifier generates a random code verifier for PKCE.
61func generateCodeVerifier() (string, error) {
62 b := make([]byte, 32)
63 if _, err := rand.Read(b); err != nil {
64 return "", err
65 }
66 return base64.RawURLEncoding.EncodeToString(b), nil
67}
68
69// generateCodeChallenge generates the code challenge from verifier.
70func generateCodeChallenge(verifier string) string {
71 h := sha256.Sum256([]byte(verifier))
72 return base64.RawURLEncoding.EncodeToString(h[:])
73}
74
75// generateState generates a random state parameter.
76func generateState() (string, error) {
77 b := make([]byte, 16)
78 if _, err := rand.Read(b); err != nil {
79 return "", err
80 }
81 return base64.RawURLEncoding.EncodeToString(b), nil
82}
83
84// AuthResult contains the authorization code and state from callback.
85type AuthResult struct {
86 Code string
87 State string
88 Error string
89}
90
91// startCallbackServer starts a local HTTP server to receive the OAuth callback.
92func (o *KiroOAuth) startCallbackServer(ctx context.Context, expectedState string) (string, <-chan AuthResult, error) {
93 // Try to find an available port - use localhost like Kiro does
94 listener, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", defaultCallbackPort))
95 if err != nil {
96 // Try with dynamic port (RFC 8252 allows dynamic ports for native apps)
97 log.Warnf("kiro oauth: default port %d is busy, falling back to dynamic port", defaultCallbackPort)
98 listener, err = net.Listen("tcp", "localhost:0")
99 if err != nil {
100 return "", nil, fmt.Errorf("failed to start callback server: %w", err)
101 }
102 }
103
104 port := listener.Addr().(*net.TCPAddr).Port
105 // Use http scheme for local callback server
106 redirectURI := fmt.Sprintf("http://localhost:%d/oauth/callback", port)
107 resultChan := make(chan AuthResult, 1)
108
109 server := &http.Server{
110 ReadHeaderTimeout: 10 * time.Second,
111 }
112
113 mux := http.NewServeMux()
114 mux.HandleFunc("/oauth/callback", func(w http.ResponseWriter, r *http.Request) {
115 code := r.URL.Query().Get("code")
116 state := r.URL.Query().Get("state")
117 errParam := r.URL.Query().Get("error")
118
119 if errParam != "" {
120 w.Header().Set("Content-Type", "text/html")
121 w.WriteHeader(http.StatusBadRequest)
122 fmt.Fprintf(w, `<html><body><h1>Login Failed</h1><p>%s</p><p>You can close this window.</p></body></html>`, html.EscapeString(errParam))
123 resultChan <- AuthResult{Error: errParam}
124 return
125 }
126
127 if state != expectedState {
128 w.Header().Set("Content-Type", "text/html")
129 w.WriteHeader(http.StatusBadRequest)
130 fmt.Fprint(w, `<html><body><h1>Login Failed</h1><p>Invalid state parameter</p><p>You can close this window.</p></body></html>`)
131 resultChan <- AuthResult{Error: "state mismatch"}
132 return
133 }
134
135 w.Header().Set("Content-Type", "text/html")
136 fmt.Fprint(w, `<html><body><h1>Login Successful!</h1><p>You can close this window and return to the terminal.</p></body></html>`)
137 resultChan <- AuthResult{Code: code, State: state}
138 })
139
140 server.Handler = mux
141
142 go func() {
143 if err := server.Serve(listener); err != nil && err != http.ErrServerClosed {
144 log.Debugf("callback server error: %v", err)
145 }
146 }()
147
148 go func() {
149 select {
150 case <-ctx.Done():
151 case <-time.After(authTimeout):
152 case <-resultChan:
153 }
154 _ = server.Shutdown(context.Background())
155 }()
156
157 return redirectURI, resultChan, nil
158}
159
160// LoginWithBuilderID performs OAuth login with AWS Builder ID using device code flow.
161func (o *KiroOAuth) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, error) {
162 ssoClient := NewSSOOIDCClient(o.cfg)
163 return ssoClient.LoginWithBuilderID(ctx)
164}
165
166// LoginWithBuilderIDAuthCode performs OAuth login with AWS Builder ID using authorization code flow.
167// This provides a better UX than device code flow as it uses automatic browser callback.
168func (o *KiroOAuth) LoginWithBuilderIDAuthCode(ctx context.Context) (*KiroTokenData, error) {
169 ssoClient := NewSSOOIDCClient(o.cfg)
170 return ssoClient.LoginWithBuilderIDAuthCode(ctx)
171}
172
173// exchangeCodeForToken exchanges the authorization code for tokens.
174func (o *KiroOAuth) exchangeCodeForToken(ctx context.Context, code, codeVerifier, redirectURI string) (*KiroTokenData, error) {
175 payload := map[string]string{
176 "code": code,
177 "code_verifier": codeVerifier,
178 "redirect_uri": redirectURI,
179 }
180
181 body, err := json.Marshal(payload)
182 if err != nil {
183 return nil, fmt.Errorf("failed to marshal request: %w", err)
184 }
185
186 tokenURL := kiroAuthEndpoint + "/oauth/token"
187 req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(string(body)))
188 if err != nil {
189 return nil, fmt.Errorf("failed to create request: %w", err)
190 }
191
192 req.Header.Set("Content-Type", "application/json")
193 req.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
194
195 resp, err := o.httpClient.Do(req)
196 if err != nil {
197 return nil, fmt.Errorf("token request failed: %w", err)
198 }
199 defer resp.Body.Close()
200
201 respBody, err := io.ReadAll(resp.Body)
202 if err != nil {
203 return nil, fmt.Errorf("failed to read response: %w", err)
204 }
205
206 if resp.StatusCode != http.StatusOK {
207 log.Debugf("token exchange failed (status %d): %s", resp.StatusCode, string(respBody))
208 return nil, fmt.Errorf("token exchange failed (status %d)", resp.StatusCode)
209 }
210
211 var tokenResp KiroTokenResponse
212 if err := json.Unmarshal(respBody, &tokenResp); err != nil {
213 return nil, fmt.Errorf("failed to parse token response: %w", err)
214 }
215
216 // Validate ExpiresIn - use default 1 hour if invalid
217 expiresIn := tokenResp.ExpiresIn
218 if expiresIn <= 0 {
219 expiresIn = 3600
220 }
221 expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
222
223 return &KiroTokenData{
224 AccessToken: tokenResp.AccessToken,
225 RefreshToken: tokenResp.RefreshToken,
226 ProfileArn: tokenResp.ProfileArn,
227 ExpiresAt: expiresAt.Format(time.RFC3339),
228 AuthMethod: "social",
229 Provider: "", // Caller should preserve original provider
230 }, nil
231}
232
233// RefreshToken refreshes an expired access token.
234func (o *KiroOAuth) RefreshToken(ctx context.Context, refreshToken string) (*KiroTokenData, error) {
235 payload := map[string]string{
236 "refreshToken": refreshToken,
237 }
238
239 body, err := json.Marshal(payload)
240 if err != nil {
241 return nil, fmt.Errorf("failed to marshal request: %w", err)
242 }
243
244 refreshURL := kiroAuthEndpoint + "/refreshToken"
245 req, err := http.NewRequestWithContext(ctx, http.MethodPost, refreshURL, strings.NewReader(string(body)))
246 if err != nil {
247 return nil, fmt.Errorf("failed to create request: %w", err)
248 }
249
250 req.Header.Set("Content-Type", "application/json")
251 req.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
252
253 resp, err := o.httpClient.Do(req)
254 if err != nil {
255 return nil, fmt.Errorf("refresh request failed: %w", err)
256 }
257 defer resp.Body.Close()
258
259 respBody, err := io.ReadAll(resp.Body)
260 if err != nil {
261 return nil, fmt.Errorf("failed to read response: %w", err)
262 }
263
264 if resp.StatusCode != http.StatusOK {
265 log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
266 return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
267 }
268
269 var tokenResp KiroTokenResponse
270 if err := json.Unmarshal(respBody, &tokenResp); err != nil {
271 return nil, fmt.Errorf("failed to parse token response: %w", err)
272 }
273
274 // Validate ExpiresIn - use default 1 hour if invalid
275 expiresIn := tokenResp.ExpiresIn
276 if expiresIn <= 0 {
277 expiresIn = 3600
278 }
279 expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
280
281 return &KiroTokenData{
282 AccessToken: tokenResp.AccessToken,
283 RefreshToken: tokenResp.RefreshToken,
284 ProfileArn: tokenResp.ProfileArn,
285 ExpiresAt: expiresAt.Format(time.RFC3339),
286 AuthMethod: "social",
287 Provider: "", // Caller should preserve original provider
288 }, nil
289}
290
291// LoginWithGoogle performs OAuth login with Google using Kiro's social auth.
292// This uses a custom protocol handler (kiro://) to receive the callback.
293func (o *KiroOAuth) LoginWithGoogle(ctx context.Context) (*KiroTokenData, error) {
294 socialClient := NewSocialAuthClient(o.cfg)
295 return socialClient.LoginWithGoogle(ctx)
296}
297
298// LoginWithGitHub performs OAuth login with GitHub using Kiro's social auth.
299// This uses a custom protocol handler (kiro://) to receive the callback.
300func (o *KiroOAuth) LoginWithGitHub(ctx context.Context) (*KiroTokenData, error) {
301 socialClient := NewSocialAuthClient(o.cfg)
302 return socialClient.LoginWithGitHub(ctx)
303}
internal/auth/kiro/protocol_handler.go+725−0
@@ -0,0 +1,725 @@
1// Package kiro provides custom protocol handler registration for Kiro OAuth.
2// This enables the CLI to intercept kiro:// URIs for social authentication (Google/GitHub).
3package kiro
4
5import (
6 "context"
7 "fmt"
8 "html"
9 "net"
10 "net/http"
11 "net/url"
12 "os"
13 "os/exec"
14 "path/filepath"
15 "runtime"
16 "strings"
17 "sync"
18 "time"
19
20 log "github.com/sirupsen/logrus"
21)
22
23const (
24 // KiroProtocol is the custom URI scheme used by Kiro
25 KiroProtocol = "kiro"
26
27 // KiroAuthority is the URI authority for authentication callbacks
28 KiroAuthority = "kiro.kiroAgent"
29
30 // KiroAuthPath is the path for successful authentication
31 KiroAuthPath = "/authenticate-success"
32
33 // KiroRedirectURI is the full redirect URI for social auth
34 KiroRedirectURI = "kiro://kiro.kiroAgent/authenticate-success"
35
36 // DefaultHandlerPort is the default port for the local callback server
37 DefaultHandlerPort = 19876
38
39 // HandlerTimeout is how long to wait for the OAuth callback
40 HandlerTimeout = 10 * time.Minute
41)
42
43// ProtocolHandler manages the custom kiro:// protocol handler for OAuth callbacks.
44type ProtocolHandler struct {
45 port int
46 server *http.Server
47 listener net.Listener
48 resultChan chan *AuthCallback
49 stopChan chan struct{}
50 mu sync.Mutex
51 running bool
52}
53
54// AuthCallback contains the OAuth callback parameters.
55type AuthCallback struct {
56 Code string
57 State string
58 Error string
59}
60
61// NewProtocolHandler creates a new protocol handler.
62func NewProtocolHandler() *ProtocolHandler {
63 return &ProtocolHandler{
64 port: DefaultHandlerPort,
65 resultChan: make(chan *AuthCallback, 1),
66 stopChan: make(chan struct{}),
67 }
68}
69
70// Start starts the local callback server that receives redirects from the protocol handler.
71func (h *ProtocolHandler) Start(ctx context.Context) (int, error) {
72 h.mu.Lock()
73 defer h.mu.Unlock()
74
75 if h.running {
76 return h.port, nil
77 }
78
79 // Drain any stale results from previous runs
80 select {
81 case <-h.resultChan:
82 default:
83 }
84
85 // Reset stopChan for reuse - close old channel first to unblock any waiting goroutines
86 if h.stopChan != nil {
87 select {
88 case <-h.stopChan:
89 // Already closed
90 default:
91 close(h.stopChan)
92 }
93 }
94 h.stopChan = make(chan struct{})
95
96 // Try ports in known range (must match handler script port range)
97 var listener net.Listener
98 var err error
99 portRange := []int{DefaultHandlerPort, DefaultHandlerPort + 1, DefaultHandlerPort + 2, DefaultHandlerPort + 3, DefaultHandlerPort + 4}
100
101 for _, port := range portRange {
102 listener, err = net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
103 if err == nil {
104 break
105 }
106 log.Debugf("kiro protocol handler: port %d busy, trying next", port)
107 }
108
109 if listener == nil {
110 return 0, fmt.Errorf("failed to start callback server: all ports %d-%d are busy", DefaultHandlerPort, DefaultHandlerPort+4)
111 }
112
113 h.listener = listener
114 h.port = listener.Addr().(*net.TCPAddr).Port
115
116 mux := http.NewServeMux()
117 mux.HandleFunc("/oauth/callback", h.handleCallback)
118
119 h.server = &http.Server{
120 Handler: mux,
121 ReadHeaderTimeout: 10 * time.Second,
122 }
123
124 go func() {
125 if err := h.server.Serve(listener); err != nil && err != http.ErrServerClosed {
126 log.Debugf("kiro protocol handler server error: %v", err)
127 }
128 }()
129
130 h.running = true
131 log.Debugf("kiro protocol handler started on port %d", h.port)
132
133 // Auto-shutdown after context done, timeout, or explicit stop
134 // Capture references to prevent race with new Start() calls
135 currentStopChan := h.stopChan
136 currentServer := h.server
137 currentListener := h.listener
138 go func() {
139 select {
140 case <-ctx.Done():
141 case <-time.After(HandlerTimeout):
142 case <-currentStopChan:
143 return // Already stopped, exit goroutine
144 }
145 // Only stop if this is still the current server/listener instance
146 h.mu.Lock()
147 if h.server == currentServer && h.listener == currentListener {
148 h.mu.Unlock()
149 h.Stop()
150 } else {
151 h.mu.Unlock()
152 }
153 }()
154
155 return h.port, nil
156}
157
158// Stop stops the callback server.
159func (h *ProtocolHandler) Stop() {
160 h.mu.Lock()
161 defer h.mu.Unlock()
162
163 if !h.running {
164 return
165 }
166
167 // Signal the auto-shutdown goroutine to exit.
168 // This select pattern is safe because stopChan is only modified while holding h.mu,
169 // and we hold the lock here. The select prevents panic from double-close.
170 select {
171 case <-h.stopChan:
172 // Already closed
173 default:
174 close(h.stopChan)
175 }
176
177 if h.server != nil {
178 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
179 defer cancel()
180 _ = h.server.Shutdown(ctx)
181 }
182
183 h.running = false
184 log.Debug("kiro protocol handler stopped")
185}
186
187// WaitForCallback waits for the OAuth callback and returns the result.
188func (h *ProtocolHandler) WaitForCallback(ctx context.Context) (*AuthCallback, error) {
189 select {
190 case <-ctx.Done():
191 return nil, ctx.Err()
192 case <-time.After(HandlerTimeout):
193 return nil, fmt.Errorf("timeout waiting for OAuth callback")
194 case result := <-h.resultChan:
195 return result, nil
196 }
197}
198
199// GetPort returns the port the handler is listening on.
200func (h *ProtocolHandler) GetPort() int {
201 return h.port
202}
203
204// handleCallback processes the OAuth callback from the protocol handler script.
205func (h *ProtocolHandler) handleCallback(w http.ResponseWriter, r *http.Request) {
206 code := r.URL.Query().Get("code")
207 state := r.URL.Query().Get("state")
208 errParam := r.URL.Query().Get("error")
209
210 result := &AuthCallback{
211 Code: code,
212 State: state,
213 Error: errParam,
214 }
215
216 // Send result
217 select {
218 case h.resultChan <- result:
219 default:
220 // Channel full, ignore duplicate callbacks
221 }
222
223 // Send success response
224 w.Header().Set("Content-Type", "text/html; charset=utf-8")
225 if errParam != "" {
226 w.WriteHeader(http.StatusBadRequest)
227 fmt.Fprintf(w, `<!DOCTYPE html>
228<html>
229<head><title>Login Failed</title></head>
230<body>
231<h1>Login Failed</h1>
232<p>Error: %s</p>
233<p>You can close this window.</p>
234</body>
235</html>`, html.EscapeString(errParam))
236 } else {
237 fmt.Fprint(w, `<!DOCTYPE html>
238<html>
239<head><title>Login Successful</title></head>
240<body>
241<h1>Login Successful!</h1>
242<p>You can close this window and return to the terminal.</p>
243<script>window.close();</script>
244</body>
245</html>`)
246 }
247}
248
249// IsProtocolHandlerInstalled checks if the kiro:// protocol handler is installed.
250func IsProtocolHandlerInstalled() bool {
251 switch runtime.GOOS {
252 case "linux":
253 return isLinuxHandlerInstalled()
254 case "windows":
255 return isWindowsHandlerInstalled()
256 case "darwin":
257 return isDarwinHandlerInstalled()
258 default:
259 return false
260 }
261}
262
263// InstallProtocolHandler installs the kiro:// protocol handler for the current platform.
264func InstallProtocolHandler(handlerPort int) error {
265 switch runtime.GOOS {
266 case "linux":
267 return installLinuxHandler(handlerPort)
268 case "windows":
269 return installWindowsHandler(handlerPort)
270 case "darwin":
271 return installDarwinHandler(handlerPort)
272 default:
273 return fmt.Errorf("unsupported platform: %s", runtime.GOOS)
274 }
275}
276
277// UninstallProtocolHandler removes the kiro:// protocol handler.
278func UninstallProtocolHandler() error {
279 switch runtime.GOOS {
280 case "linux":
281 return uninstallLinuxHandler()
282 case "windows":
283 return uninstallWindowsHandler()
284 case "darwin":
285 return uninstallDarwinHandler()
286 default:
287 return fmt.Errorf("unsupported platform: %s", runtime.GOOS)
288 }
289}
290
291// --- Linux Implementation ---
292
293func getLinuxDesktopPath() string {
294 homeDir, _ := os.UserHomeDir()
295 return filepath.Join(homeDir, ".local", "share", "applications", "kiro-oauth-handler.desktop")
296}
297
298func getLinuxHandlerScriptPath() string {
299 homeDir, _ := os.UserHomeDir()
300 return filepath.Join(homeDir, ".local", "bin", "kiro-oauth-handler")
301}
302
303func isLinuxHandlerInstalled() bool {
304 desktopPath := getLinuxDesktopPath()
305 _, err := os.Stat(desktopPath)
306 return err == nil
307}
308
309func installLinuxHandler(handlerPort int) error {
310 // Create directories
311 homeDir, err := os.UserHomeDir()
312 if err != nil {
313 return err
314 }
315
316 binDir := filepath.Join(homeDir, ".local", "bin")
317 appDir := filepath.Join(homeDir, ".local", "share", "applications")
318
319 if err := os.MkdirAll(binDir, 0755); err != nil {
320 return fmt.Errorf("failed to create bin directory: %w", err)
321 }
322 if err := os.MkdirAll(appDir, 0755); err != nil {
323 return fmt.Errorf("failed to create applications directory: %w", err)
324 }
325
326 // Create handler script - tries multiple ports to handle dynamic port allocation
327 scriptPath := getLinuxHandlerScriptPath()
328 scriptContent := fmt.Sprintf(`#!/bin/bash
329# Kiro OAuth Protocol Handler
330# Handles kiro:// URIs - tries CLI first, then forwards to Kiro IDE
331
332URL="$1"
333
334# Check curl availability
335if ! command -v curl &> /dev/null; then
336 echo "Error: curl is required for Kiro OAuth handler" >&2
337 exit 1
338fi
339
340# Extract code and state from URL
341[[ "$URL" =~ code=([^&]+) ]] && CODE="${BASH_REMATCH[1]}"
342[[ "$URL" =~ state=([^&]+) ]] && STATE="${BASH_REMATCH[1]}"
343[[ "$URL" =~ error=([^&]+) ]] && ERROR="${BASH_REMATCH[1]}"
344
345# Try CLI proxy on multiple possible ports (default + dynamic range)
346CLI_OK=0
347for PORT in %d %d %d %d %d; do
348 if [ -n "$ERROR" ]; then
349 curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?error=$ERROR" && CLI_OK=1 && break
350 elif [ -n "$CODE" ] && [ -n "$STATE" ]; then
351 curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?code=$CODE&state=$STATE" && CLI_OK=1 && break
352 fi
353done
354
355# If CLI not available, forward to Kiro IDE
356if [ $CLI_OK -eq 0 ] && [ -x "/usr/share/kiro/kiro" ]; then
357 /usr/share/kiro/kiro --open-url "$URL" &
358fi
359`, handlerPort, handlerPort+1, handlerPort+2, handlerPort+3, handlerPort+4)
360
361 if err := os.WriteFile(scriptPath, []byte(scriptContent), 0755); err != nil {
362 return fmt.Errorf("failed to write handler script: %w", err)
363 }
364
365 // Create .desktop file
366 desktopPath := getLinuxDesktopPath()
367 desktopContent := fmt.Sprintf(`[Desktop Entry]
368Name=Kiro OAuth Handler
369Comment=Handle kiro:// protocol for CLI Proxy API authentication
370Exec=%s %%u
371Type=Application
372Terminal=false
373NoDisplay=true
374MimeType=x-scheme-handler/kiro;
375Categories=Utility;
376`, scriptPath)
377
378 if err := os.WriteFile(desktopPath, []byte(desktopContent), 0644); err != nil {
379 return fmt.Errorf("failed to write desktop file: %w", err)
380 }
381
382 // Register handler with xdg-mime
383 cmd := exec.Command("xdg-mime", "default", "kiro-oauth-handler.desktop", "x-scheme-handler/kiro")
384 if err := cmd.Run(); err != nil {
385 log.Warnf("xdg-mime registration failed (may need manual setup): %v", err)
386 }
387
388 // Update desktop database
389 cmd = exec.Command("update-desktop-database", appDir)
390 _ = cmd.Run() // Ignore errors, not critical
391
392 log.Info("Kiro protocol handler installed for Linux")
393 return nil
394}
395
396func uninstallLinuxHandler() error {
397 desktopPath := getLinuxDesktopPath()
398 scriptPath := getLinuxHandlerScriptPath()
399
400 if err := os.Remove(desktopPath); err != nil && !os.IsNotExist(err) {
401 return fmt.Errorf("failed to remove desktop file: %w", err)
402 }
403 if err := os.Remove(scriptPath); err != nil && !os.IsNotExist(err) {
404 return fmt.Errorf("failed to remove handler script: %w", err)
405 }
406
407 log.Info("Kiro protocol handler uninstalled")
408 return nil
409}
410
411// --- Windows Implementation ---
412
413func isWindowsHandlerInstalled() bool {
414 // Check registry key existence
415 cmd := exec.Command("reg", "query", `HKCU\Software\Classes\kiro`, "/ve")
416 return cmd.Run() == nil
417}
418
419func installWindowsHandler(handlerPort int) error {
420 homeDir, err := os.UserHomeDir()
421 if err != nil {
422 return err
423 }
424
425 // Create handler script (PowerShell)
426 scriptDir := filepath.Join(homeDir, ".cliproxyapi")
427 if err := os.MkdirAll(scriptDir, 0755); err != nil {
428 return fmt.Errorf("failed to create script directory: %w", err)
429 }
430
431 scriptPath := filepath.Join(scriptDir, "kiro-oauth-handler.ps1")
432 scriptContent := fmt.Sprintf(`# Kiro OAuth Protocol Handler for Windows
433param([string]$url)
434
435# Load required assembly for HttpUtility
436Add-Type -AssemblyName System.Web
437
438# Parse URL parameters
439$uri = [System.Uri]$url
440$query = [System.Web.HttpUtility]::ParseQueryString($uri.Query)
441$code = $query["code"]
442$state = $query["state"]
443$errorParam = $query["error"]
444
445# Try multiple ports (default + dynamic range)
446$ports = @(%d, %d, %d, %d, %d)
447$success = $false
448
449foreach ($port in $ports) {
450 if ($success) { break }
451 $callbackUrl = "http://127.0.0.1:$port/oauth/callback"
452 try {
453 if ($errorParam) {
454 $fullUrl = $callbackUrl + "?error=" + $errorParam
455 Invoke-WebRequest -Uri $fullUrl -UseBasicParsing -TimeoutSec 1 -ErrorAction Stop | Out-Null
456 $success = $true
457 } elseif ($code -and $state) {
458 $fullUrl = $callbackUrl + "?code=" + $code + "&state=" + $state
459 Invoke-WebRequest -Uri $fullUrl -UseBasicParsing -TimeoutSec 1 -ErrorAction Stop | Out-Null
460 $success = $true
461 }
462 } catch {
463 # Try next port
464 }
465}
466`, handlerPort, handlerPort+1, handlerPort+2, handlerPort+3, handlerPort+4)
467
468 if err := os.WriteFile(scriptPath, []byte(scriptContent), 0644); err != nil {
469 return fmt.Errorf("failed to write handler script: %w", err)
470 }
471
472 // Create batch wrapper
473 batchPath := filepath.Join(scriptDir, "kiro-oauth-handler.bat")
474 batchContent := fmt.Sprintf("@echo off\npowershell -ExecutionPolicy Bypass -File \"%s\" %%1\n", scriptPath)
475
476 if err := os.WriteFile(batchPath, []byte(batchContent), 0644); err != nil {
477 return fmt.Errorf("failed to write batch wrapper: %w", err)
478 }
479
480 // Register in Windows registry
481 commands := [][]string{
482 {"reg", "add", `HKCU\Software\Classes\kiro`, "/ve", "/d", "URL:Kiro Protocol", "/f"},
483 {"reg", "add", `HKCU\Software\Classes\kiro`, "/v", "URL Protocol", "/d", "", "/f"},
484 {"reg", "add", `HKCU\Software\Classes\kiro\shell`, "/f"},
485 {"reg", "add", `HKCU\Software\Classes\kiro\shell\open`, "/f"},
486 {"reg", "add", `HKCU\Software\Classes\kiro\shell\open\command`, "/ve", "/d", fmt.Sprintf("\"%s\" \"%%1\"", batchPath), "/f"},
487 }
488
489 for _, args := range commands {
490 cmd := exec.Command(args[0], args[1:]...)
491 if err := cmd.Run(); err != nil {
492 return fmt.Errorf("failed to run registry command: %w", err)
493 }
494 }
495
496 log.Info("Kiro protocol handler installed for Windows")
497 return nil
498}
499
500func uninstallWindowsHandler() error {
501 // Remove registry keys
502 cmd := exec.Command("reg", "delete", `HKCU\Software\Classes\kiro`, "/f")
503 if err := cmd.Run(); err != nil {
504 log.Warnf("failed to remove registry key: %v", err)
505 }
506
507 // Remove scripts
508 homeDir, _ := os.UserHomeDir()
509 scriptDir := filepath.Join(homeDir, ".cliproxyapi")
510 _ = os.Remove(filepath.Join(scriptDir, "kiro-oauth-handler.ps1"))
511 _ = os.Remove(filepath.Join(scriptDir, "kiro-oauth-handler.bat"))
512
513 log.Info("Kiro protocol handler uninstalled")
514 return nil
515}
516
517// --- macOS Implementation ---
518
519func getDarwinAppPath() string {
520 homeDir, _ := os.UserHomeDir()
521 return filepath.Join(homeDir, "Applications", "KiroOAuthHandler.app")
522}
523
524func isDarwinHandlerInstalled() bool {
525 appPath := getDarwinAppPath()
526 _, err := os.Stat(appPath)
527 return err == nil
528}
529
530func installDarwinHandler(handlerPort int) error {
531 // Create app bundle structure
532 appPath := getDarwinAppPath()
533 contentsPath := filepath.Join(appPath, "Contents")
534 macOSPath := filepath.Join(contentsPath, "MacOS")
535
536 if err := os.MkdirAll(macOSPath, 0755); err != nil {
537 return fmt.Errorf("failed to create app bundle: %w", err)
538 }
539
540 // Create Info.plist
541 plistPath := filepath.Join(contentsPath, "Info.plist")
542 plistContent := `<?xml version="1.0" encoding="UTF-8"?>
543<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
544<plist version="1.0">
545<dict>
546 <key>CFBundleIdentifier</key>
547 <string>com.cliproxyapi.kiro-oauth-handler</string>
548 <key>CFBundleName</key>
549 <string>KiroOAuthHandler</string>
550 <key>CFBundleExecutable</key>
551 <string>kiro-oauth-handler</string>
552 <key>CFBundleVersion</key>
553 <string>1.0</string>
554 <key>CFBundleURLTypes</key>
555 <array>
556 <dict>
557 <key>CFBundleURLName</key>
558 <string>Kiro Protocol</string>
559 <key>CFBundleURLSchemes</key>
560 <array>
561 <string>kiro</string>
562 </array>
563 </dict>
564 </array>
565 <key>LSBackgroundOnly</key>
566 <true/>
567</dict>
568</plist>`
569
570 if err := os.WriteFile(plistPath, []byte(plistContent), 0644); err != nil {
571 return fmt.Errorf("failed to write Info.plist: %w", err)
572 }
573
574 // Create executable script - tries multiple ports to handle dynamic port allocation
575 execPath := filepath.Join(macOSPath, "kiro-oauth-handler")
576 execContent := fmt.Sprintf(`#!/bin/bash
577# Kiro OAuth Protocol Handler for macOS
578
579URL="$1"
580
581# Check curl availability (should always exist on macOS)
582if [ ! -x /usr/bin/curl ]; then
583 echo "Error: curl is required for Kiro OAuth handler" >&2
584 exit 1
585fi
586
587# Extract code and state from URL
588[[ "$URL" =~ code=([^&]+) ]] && CODE="${BASH_REMATCH[1]}"
589[[ "$URL" =~ state=([^&]+) ]] && STATE="${BASH_REMATCH[1]}"
590[[ "$URL" =~ error=([^&]+) ]] && ERROR="${BASH_REMATCH[1]}"
591
592# Try multiple ports (default + dynamic range)
593for PORT in %d %d %d %d %d; do
594 if [ -n "$ERROR" ]; then
595 /usr/bin/curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?error=$ERROR" && exit 0
596 elif [ -n "$CODE" ] && [ -n "$STATE" ]; then
597 /usr/bin/curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?code=$CODE&state=$STATE" && exit 0
598 fi
599done
600`, handlerPort, handlerPort+1, handlerPort+2, handlerPort+3, handlerPort+4)
601
602 if err := os.WriteFile(execPath, []byte(execContent), 0755); err != nil {
603 return fmt.Errorf("failed to write executable: %w", err)
604 }
605
606 // Register the app with Launch Services
607 cmd := exec.Command("/System/Library/Frameworks/CoreServices.framework/Frameworks/LaunchServices.framework/Support/lsregister",
608 "-f", appPath)
609 if err := cmd.Run(); err != nil {
610 log.Warnf("lsregister failed (handler may still work): %v", err)
611 }
612
613 log.Info("Kiro protocol handler installed for macOS")
614 return nil
615}
616
617func uninstallDarwinHandler() error {
618 appPath := getDarwinAppPath()
619
620 // Unregister from Launch Services
621 cmd := exec.Command("/System/Library/Frameworks/CoreServices.framework/Frameworks/LaunchServices.framework/Support/lsregister",
622 "-u", appPath)
623 _ = cmd.Run()
624
625 // Remove app bundle
626 if err := os.RemoveAll(appPath); err != nil && !os.IsNotExist(err) {
627 return fmt.Errorf("failed to remove app bundle: %w", err)
628 }
629
630 log.Info("Kiro protocol handler uninstalled")
631 return nil
632}
633
634// ParseKiroURI parses a kiro:// URI and extracts the callback parameters.
635func ParseKiroURI(rawURI string) (*AuthCallback, error) {
636 u, err := url.Parse(rawURI)
637 if err != nil {
638 return nil, fmt.Errorf("invalid URI: %w", err)
639 }
640
641 if u.Scheme != KiroProtocol {
642 return nil, fmt.Errorf("invalid scheme: expected %s, got %s", KiroProtocol, u.Scheme)
643 }
644
645 if u.Host != KiroAuthority {
646 return nil, fmt.Errorf("invalid authority: expected %s, got %s", KiroAuthority, u.Host)
647 }
648
649 query := u.Query()
650 return &AuthCallback{
651 Code: query.Get("code"),
652 State: query.Get("state"),
653 Error: query.Get("error"),
654 }, nil
655}
656
657// GetHandlerInstructions returns platform-specific instructions for manual handler setup.
658func GetHandlerInstructions() string {
659 switch runtime.GOOS {
660 case "linux":
661 return `To manually set up the Kiro protocol handler on Linux:
662
6631. Create ~/.local/share/applications/kiro-oauth-handler.desktop:
664 [Desktop Entry]
665 Name=Kiro OAuth Handler
666 Exec=~/.local/bin/kiro-oauth-handler %u
667 Type=Application
668 Terminal=false
669 MimeType=x-scheme-handler/kiro;
670
6712. Create ~/.local/bin/kiro-oauth-handler (make it executable):
672 #!/bin/bash
673 URL="$1"
674 # ... (see generated script for full content)
675
6763. Run: xdg-mime default kiro-oauth-handler.desktop x-scheme-handler/kiro`
677
678 case "windows":
679 return `To manually set up the Kiro protocol handler on Windows:
680
6811. Open Registry Editor (regedit.exe)
6822. Create key: HKEY_CURRENT_USER\Software\Classes\kiro
6833. Set default value to: URL:Kiro Protocol
6844. Create string value "URL Protocol" with empty data
6855. Create subkey: shell\open\command
6866. Set default value to: "C:\path\to\handler.bat" "%1"`
687
688 case "darwin":
689 return `To manually set up the Kiro protocol handler on macOS:
690
6911. Create ~/Applications/KiroOAuthHandler.app bundle
6922. Add Info.plist with CFBundleURLTypes containing "kiro" scheme
6933. Create executable in Contents/MacOS/
6944. Run: /System/Library/.../lsregister -f ~/Applications/KiroOAuthHandler.app`
695
696 default:
697 return "Protocol handler setup is not supported on this platform."
698 }
699}
700
701// SetupProtocolHandlerIfNeeded checks and installs the protocol handler if needed.
702func SetupProtocolHandlerIfNeeded(handlerPort int) error {
703 if IsProtocolHandlerInstalled() {
704 log.Debug("Kiro protocol handler already installed")
705 return nil
706 }
707
708 fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
709 fmt.Println("║ Kiro Protocol Handler Setup Required ║")
710 fmt.Println("╚══════════════════════════════════════════════════════════╝")
711 fmt.Println("\nTo enable Google/GitHub login, we need to install a protocol handler.")
712 fmt.Println("This allows your browser to redirect back to the CLI after authentication.")
713 fmt.Println("\nInstalling protocol handler...")
714
715 if err := InstallProtocolHandler(handlerPort); err != nil {
716 fmt.Printf("\n⚠ Automatic installation failed: %v\n", err)
717 fmt.Println("\nManual setup instructions:")
718 fmt.Println(strings.Repeat("-", 60))
719 fmt.Println(GetHandlerInstructions())
720 return err
721 }
722
723 fmt.Println("\n✓ Protocol handler installed successfully!")
724 return nil
725}
internal/auth/kiro/social_auth.go+403−0
@@ -0,0 +1,403 @@
1// Package kiro provides social authentication (Google/GitHub) for Kiro via AuthServiceClient.
2package kiro
3
4import (
5 "bufio"
6 "context"
7 "crypto/rand"
8 "crypto/sha256"
9 "encoding/base64"
10 "encoding/json"
11 "fmt"
12 "io"
13 "net/http"
14 "net/url"
15 "os"
16 "os/exec"
17 "runtime"
18 "strings"
19 "time"
20
21 "github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
22 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
23 "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
24 log "github.com/sirupsen/logrus"
25 "golang.org/x/term"
26)
27
28const (
29 // Kiro AuthService endpoint
30 kiroAuthServiceEndpoint = "https://prod.us-east-1.auth.desktop.kiro.dev"
31
32 // OAuth timeout
33 socialAuthTimeout = 10 * time.Minute
34)
35
36// SocialProvider represents the social login provider.
37type SocialProvider string
38
39const (
40 // ProviderGoogle is Google OAuth provider
41 ProviderGoogle SocialProvider = "Google"
42 // ProviderGitHub is GitHub OAuth provider
43 ProviderGitHub SocialProvider = "Github"
44 // Note: AWS Builder ID is NOT supported by Kiro's auth service.
45 // It only supports: Google, Github, Cognito
46 // AWS Builder ID must use device code flow via SSO OIDC.
47)
48
49// CreateTokenRequest is sent to Kiro's /oauth/token endpoint.
50type CreateTokenRequest struct {
51 Code string `json:"code"`
52 CodeVerifier string `json:"code_verifier"`
53 RedirectURI string `json:"redirect_uri"`
54 InvitationCode string `json:"invitation_code,omitempty"`
55}
56
57// SocialTokenResponse from Kiro's /oauth/token endpoint for social auth.
58type SocialTokenResponse struct {
59 AccessToken string `json:"accessToken"`
60 RefreshToken string `json:"refreshToken"`
61 ProfileArn string `json:"profileArn"`
62 ExpiresIn int `json:"expiresIn"`
63}
64
65// RefreshTokenRequest is sent to Kiro's /refreshToken endpoint.
66type RefreshTokenRequest struct {
67 RefreshToken string `json:"refreshToken"`
68}
69
70// SocialAuthClient handles social authentication with Kiro.
71type SocialAuthClient struct {
72 httpClient *http.Client
73 cfg *config.Config
74 protocolHandler *ProtocolHandler
75}
76
77// NewSocialAuthClient creates a new social auth client.
78func NewSocialAuthClient(cfg *config.Config) *SocialAuthClient {
79 client := &http.Client{Timeout: 30 * time.Second}
80 if cfg != nil {
81 client = util.SetProxy(&cfg.SDKConfig, client)
82 }
83 return &SocialAuthClient{
84 httpClient: client,
85 cfg: cfg,
86 protocolHandler: NewProtocolHandler(),
87 }
88}
89
90// generatePKCE generates PKCE code verifier and challenge.
91func generatePKCE() (verifier, challenge string, err error) {
92 // Generate 32 bytes of random data for verifier
93 b := make([]byte, 32)
94 if _, err := rand.Read(b); err != nil {
95 return "", "", fmt.Errorf("failed to generate random bytes: %w", err)
96 }
97 verifier = base64.RawURLEncoding.EncodeToString(b)
98
99 // Generate SHA256 hash of verifier for challenge
100 h := sha256.Sum256([]byte(verifier))
101 challenge = base64.RawURLEncoding.EncodeToString(h[:])
102
103 return verifier, challenge, nil
104}
105
106// generateState generates a random state parameter.
107func generateStateParam() (string, error) {
108 b := make([]byte, 16)
109 if _, err := rand.Read(b); err != nil {
110 return "", err
111 }
112 return base64.RawURLEncoding.EncodeToString(b), nil
113}
114
115// buildLoginURL constructs the Kiro OAuth login URL.
116// The login endpoint expects a GET request with query parameters.
117// Format: /login?idp=Google&redirect_uri=...&code_challenge=...&code_challenge_method=S256&state=...&prompt=select_account
118// The prompt=select_account parameter forces the account selection screen even if already logged in.
119func (c *SocialAuthClient) buildLoginURL(provider, redirectURI, codeChallenge, state string) string {
120 return fmt.Sprintf("%s/login?idp=%s&redirect_uri=%s&code_challenge=%s&code_challenge_method=S256&state=%s&prompt=select_account",
121 kiroAuthServiceEndpoint,
122 provider,
123 url.QueryEscape(redirectURI),
124 codeChallenge,
125 state,
126 )
127}
128
129// CreateToken exchanges the authorization code for tokens.
130func (c *SocialAuthClient) CreateToken(ctx context.Context, req *CreateTokenRequest) (*SocialTokenResponse, error) {
131 body, err := json.Marshal(req)
132 if err != nil {
133 return nil, fmt.Errorf("failed to marshal token request: %w", err)
134 }
135
136 tokenURL := kiroAuthServiceEndpoint + "/oauth/token"
137 httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(string(body)))
138 if err != nil {
139 return nil, fmt.Errorf("failed to create token request: %w", err)
140 }
141
142 httpReq.Header.Set("Content-Type", "application/json")
143 httpReq.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
144
145 resp, err := c.httpClient.Do(httpReq)
146 if err != nil {
147 return nil, fmt.Errorf("token request failed: %w", err)
148 }
149 defer resp.Body.Close()
150
151 respBody, err := io.ReadAll(resp.Body)
152 if err != nil {
153 return nil, fmt.Errorf("failed to read token response: %w", err)
154 }
155
156 if resp.StatusCode != http.StatusOK {
157 log.Debugf("token exchange failed (status %d): %s", resp.StatusCode, string(respBody))
158 return nil, fmt.Errorf("token exchange failed (status %d)", resp.StatusCode)
159 }
160
161 var tokenResp SocialTokenResponse
162 if err := json.Unmarshal(respBody, &tokenResp); err != nil {
163 return nil, fmt.Errorf("failed to parse token response: %w", err)
164 }
165
166 return &tokenResp, nil
167}
168
169// RefreshSocialToken refreshes an expired social auth token.
170func (c *SocialAuthClient) RefreshSocialToken(ctx context.Context, refreshToken string) (*KiroTokenData, error) {
171 body, err := json.Marshal(&RefreshTokenRequest{RefreshToken: refreshToken})
172 if err != nil {
173 return nil, fmt.Errorf("failed to marshal refresh request: %w", err)
174 }
175
176 refreshURL := kiroAuthServiceEndpoint + "/refreshToken"
177 httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, refreshURL, strings.NewReader(string(body)))
178 if err != nil {
179 return nil, fmt.Errorf("failed to create refresh request: %w", err)
180 }
181
182 httpReq.Header.Set("Content-Type", "application/json")
183 httpReq.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
184
185 resp, err := c.httpClient.Do(httpReq)
186 if err != nil {
187 return nil, fmt.Errorf("refresh request failed: %w", err)
188 }
189 defer resp.Body.Close()
190
191 respBody, err := io.ReadAll(resp.Body)
192 if err != nil {
193 return nil, fmt.Errorf("failed to read refresh response: %w", err)
194 }
195
196 if resp.StatusCode != http.StatusOK {
197 log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
198 return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
199 }
200
201 var tokenResp SocialTokenResponse
202 if err := json.Unmarshal(respBody, &tokenResp); err != nil {
203 return nil, fmt.Errorf("failed to parse refresh response: %w", err)
204 }
205
206 // Validate ExpiresIn - use default 1 hour if invalid
207 expiresIn := tokenResp.ExpiresIn
208 if expiresIn <= 0 {
209 expiresIn = 3600 // Default 1 hour
210 }
211 expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
212
213 return &KiroTokenData{
214 AccessToken: tokenResp.AccessToken,
215 RefreshToken: tokenResp.RefreshToken,
216 ProfileArn: tokenResp.ProfileArn,
217 ExpiresAt: expiresAt.Format(time.RFC3339),
218 AuthMethod: "social",
219 Provider: "", // Caller should preserve original provider
220 }, nil
221}
222
223// LoginWithSocial performs OAuth login with Google.
224func (c *SocialAuthClient) LoginWithSocial(ctx context.Context, provider SocialProvider) (*KiroTokenData, error) {
225 providerName := string(provider)
226
227 fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
228 fmt.Printf("║ Kiro Authentication (%s) ║\n", providerName)
229 fmt.Println("╚══════════════════════════════════════════════════════════╝")
230
231 // Step 1: Setup protocol handler
232 fmt.Println("\nSetting up authentication...")
233
234 // Start the local callback server
235 handlerPort, err := c.protocolHandler.Start(ctx)
236 if err != nil {
237 return nil, fmt.Errorf("failed to start callback server: %w", err)
238 }
239 defer c.protocolHandler.Stop()
240
241 // Ensure protocol handler is installed and set as default
242 if err := SetupProtocolHandlerIfNeeded(handlerPort); err != nil {
243 fmt.Println("\n⚠ Protocol handler setup failed. Trying alternative method...")
244 fmt.Println(" If you see a browser 'Open with' dialog, select your default browser.")
245 fmt.Println(" For manual setup instructions, run: cliproxy kiro --help-protocol")
246 log.Debugf("kiro: protocol handler setup error: %v", err)
247 // Continue anyway - user might have set it up manually or select browser manually
248 } else {
249 // Force set our handler as default (prevents "Open with" dialog)
250 forceDefaultProtocolHandler()
251 }
252
253 // Step 2: Generate PKCE codes
254 codeVerifier, codeChallenge, err := generatePKCE()
255 if err != nil {
256 return nil, fmt.Errorf("failed to generate PKCE: %w", err)
257 }
258
259 // Step 3: Generate state
260 state, err := generateStateParam()
261 if err != nil {
262 return nil, fmt.Errorf("failed to generate state: %w", err)
263 }
264
265 // Step 4: Build the login URL (Kiro uses GET request with query params)
266 authURL := c.buildLoginURL(providerName, KiroRedirectURI, codeChallenge, state)
267
268 // Set incognito mode based on config (defaults to true for Kiro, can be overridden with --no-incognito)
269 // Incognito mode enables multi-account support by bypassing cached sessions
270 if c.cfg != nil {
271 browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
272 if !c.cfg.IncognitoBrowser {
273 log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.")
274 } else {
275 log.Debug("kiro: using incognito mode for multi-account support")
276 }
277 } else {
278 browser.SetIncognitoMode(true) // Default to incognito if no config
279 log.Debug("kiro: using incognito mode for multi-account support (default)")
280 }
281
282 // Step 5: Open browser for user authentication
283 fmt.Println("\n════════════════════════════════════════════════════════════")
284 fmt.Printf(" Opening browser for %s authentication...\n", providerName)
285 fmt.Println("════════════════════════════════════════════════════════════")
286 fmt.Printf("\n URL: %s\n\n", authURL)
287
288 if err := browser.OpenURL(authURL); err != nil {
289 log.Warnf("Could not open browser automatically: %v", err)
290 fmt.Println(" ⚠ Could not open browser automatically.")
291 fmt.Println(" Please open the URL above in your browser manually.")
292 } else {
293 fmt.Println(" (Browser opened automatically)")
294 }
295
296 fmt.Println("\n Waiting for authentication callback...")
297
298 // Step 6: Wait for callback
299 callback, err := c.protocolHandler.WaitForCallback(ctx)
300 if err != nil {
301 return nil, fmt.Errorf("failed to receive callback: %w", err)
302 }
303
304 if callback.Error != "" {
305 return nil, fmt.Errorf("authentication error: %s", callback.Error)
306 }
307
308 if callback.State != state {
309 // Log state values for debugging, but don't expose in user-facing error
310 log.Debugf("kiro: OAuth state mismatch - expected %s, got %s", state, callback.State)
311 return nil, fmt.Errorf("OAuth state validation failed - please try again")
312 }
313
314 if callback.Code == "" {
315 return nil, fmt.Errorf("no authorization code received")
316 }
317
318 fmt.Println("\n✓ Authorization received!")
319
320 // Step 7: Exchange code for tokens
321 fmt.Println("Exchanging code for tokens...")
322
323 tokenReq := &CreateTokenRequest{
324 Code: callback.Code,
325 CodeVerifier: codeVerifier,
326 RedirectURI: KiroRedirectURI,
327 }
328
329 tokenResp, err := c.CreateToken(ctx, tokenReq)
330 if err != nil {
331 return nil, fmt.Errorf("failed to exchange code for tokens: %w", err)
332 }
333
334 fmt.Println("\n✓ Authentication successful!")
335
336 // Close the browser window
337 if err := browser.CloseBrowser(); err != nil {
338 log.Debugf("Failed to close browser: %v", err)
339 }
340
341 // Validate ExpiresIn - use default 1 hour if invalid
342 expiresIn := tokenResp.ExpiresIn
343 if expiresIn <= 0 {
344 expiresIn = 3600
345 }
346 expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
347
348 // Try to extract email from JWT access token first
349 email := ExtractEmailFromJWT(tokenResp.AccessToken)
350
351 // If no email in JWT, ask user for account label (only in interactive mode)
352 if email == "" && isInteractiveTerminal() {
353 fmt.Print("\n Enter account label for file naming (optional, press Enter to skip): ")
354 reader := bufio.NewReader(os.Stdin)
355 var err error
356 email, err = reader.ReadString('\n')
357 if err != nil {
358 log.Debugf("Failed to read account label: %v", err)
359 }
360 email = strings.TrimSpace(email)
361 }
362
363 return &KiroTokenData{
364 AccessToken: tokenResp.AccessToken,
365 RefreshToken: tokenResp.RefreshToken,
366 ProfileArn: tokenResp.ProfileArn,
367 ExpiresAt: expiresAt.Format(time.RFC3339),
368 AuthMethod: "social",
369 Provider: providerName,
370 Email: email, // JWT email or user-provided label
371 }, nil
372}
373
374// LoginWithGoogle performs OAuth login with Google.
375func (c *SocialAuthClient) LoginWithGoogle(ctx context.Context) (*KiroTokenData, error) {
376 return c.LoginWithSocial(ctx, ProviderGoogle)
377}
378
379// LoginWithGitHub performs OAuth login with GitHub.
380func (c *SocialAuthClient) LoginWithGitHub(ctx context.Context) (*KiroTokenData, error) {
381 return c.LoginWithSocial(ctx, ProviderGitHub)
382}
383
384// forceDefaultProtocolHandler sets our protocol handler as the default for kiro:// URLs.
385// This prevents the "Open with" dialog from appearing on Linux.
386// On non-Linux platforms, this is a no-op as they use different mechanisms.
387func forceDefaultProtocolHandler() {
388 if runtime.GOOS != "linux" {
389 return // Non-Linux platforms use different handler mechanisms
390 }
391
392 // Set our handler as default using xdg-mime
393 cmd := exec.Command("xdg-mime", "default", "kiro-oauth-handler.desktop", "x-scheme-handler/kiro")
394 if err := cmd.Run(); err != nil {
395 log.Warnf("Failed to set default protocol handler: %v. You may see a handler selection dialog.", err)
396 }
397}
398
399// isInteractiveTerminal checks if stdin is connected to an interactive terminal.
400// Returns false in CI/automated environments or when stdin is piped.
401func isInteractiveTerminal() bool {
402 return term.IsTerminal(int(os.Stdin.Fd()))
403}
internal/auth/kiro/sso_oidc.go+1371−0
@@ -0,0 +1,1371 @@
1// Package kiro provides AWS SSO OIDC authentication for Kiro.
2package kiro
3
4import (
5 "bufio"
6 "context"
7 "crypto/rand"
8 "crypto/sha256"
9 "encoding/base64"
10 "encoding/json"
11 "errors"
12 "fmt"
13 "html"
14 "io"
15 "net"
16 "net/http"
17 "os"
18 "strings"
19 "time"
20
21 "github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
22 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
23 "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
24 log "github.com/sirupsen/logrus"
25)
26
27const (
28 // AWS SSO OIDC endpoints
29 ssoOIDCEndpoint = "https://oidc.us-east-1.amazonaws.com"
30
31 // Kiro's start URL for Builder ID
32 builderIDStartURL = "https://view.awsapps.com/start"
33
34 // Default region for IDC
35 defaultIDCRegion = "us-east-1"
36
37 // Polling interval
38 pollInterval = 5 * time.Second
39
40 // Authorization code flow callback
41 authCodeCallbackPath = "/oauth/callback"
42 authCodeCallbackPort = 19877
43
44 // User-Agent to match official Kiro IDE
45 kiroUserAgent = "KiroIDE"
46
47 // IDC token refresh headers (matching Kiro IDE behavior)
48 idcAmzUserAgent = "aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE"
49)
50
51// Sentinel errors for OIDC token polling
52var (
53 ErrAuthorizationPending = errors.New("authorization_pending")
54 ErrSlowDown = errors.New("slow_down")
55)
56
57// SSOOIDCClient handles AWS SSO OIDC authentication.
58type SSOOIDCClient struct {
59 httpClient *http.Client
60 cfg *config.Config
61}
62
63// NewSSOOIDCClient creates a new SSO OIDC client.
64func NewSSOOIDCClient(cfg *config.Config) *SSOOIDCClient {
65 client := &http.Client{Timeout: 30 * time.Second}
66 if cfg != nil {
67 client = util.SetProxy(&cfg.SDKConfig, client)
68 }
69 return &SSOOIDCClient{
70 httpClient: client,
71 cfg: cfg,
72 }
73}
74
75// RegisterClientResponse from AWS SSO OIDC.
76type RegisterClientResponse struct {
77 ClientID string `json:"clientId"`
78 ClientSecret string `json:"clientSecret"`
79 ClientIDIssuedAt int64 `json:"clientIdIssuedAt"`
80 ClientSecretExpiresAt int64 `json:"clientSecretExpiresAt"`
81}
82
83// StartDeviceAuthResponse from AWS SSO OIDC.
84type StartDeviceAuthResponse struct {
85 DeviceCode string `json:"deviceCode"`
86 UserCode string `json:"userCode"`
87 VerificationURI string `json:"verificationUri"`
88 VerificationURIComplete string `json:"verificationUriComplete"`
89 ExpiresIn int `json:"expiresIn"`
90 Interval int `json:"interval"`
91}
92
93// CreateTokenResponse from AWS SSO OIDC.
94type CreateTokenResponse struct {
95 AccessToken string `json:"accessToken"`
96 TokenType string `json:"tokenType"`
97 ExpiresIn int `json:"expiresIn"`
98 RefreshToken string `json:"refreshToken"`
99}
100
101// getOIDCEndpoint returns the OIDC endpoint for the given region.
102func getOIDCEndpoint(region string) string {
103 if region == "" {
104 region = defaultIDCRegion
105 }
106 return fmt.Sprintf("https://oidc.%s.amazonaws.com", region)
107}
108
109// promptInput prompts the user for input with an optional default value.
110func promptInput(prompt, defaultValue string) string {
111 reader := bufio.NewReader(os.Stdin)
112 if defaultValue != "" {
113 fmt.Printf("%s [%s]: ", prompt, defaultValue)
114 } else {
115 fmt.Printf("%s: ", prompt)
116 }
117 input, err := reader.ReadString('\n')
118 if err != nil {
119 log.Warnf("Error reading input: %v", err)
120 return defaultValue
121 }
122 input = strings.TrimSpace(input)
123 if input == "" {
124 return defaultValue
125 }
126 return input
127}
128
129// promptSelect prompts the user to select from options using number input.
130func promptSelect(prompt string, options []string) int {
131 reader := bufio.NewReader(os.Stdin)
132
133 for {
134 fmt.Println(prompt)
135 for i, opt := range options {
136 fmt.Printf(" %d) %s\n", i+1, opt)
137 }
138 fmt.Printf("Enter selection (1-%d): ", len(options))
139
140 input, err := reader.ReadString('\n')
141 if err != nil {
142 log.Warnf("Error reading input: %v", err)
143 return 0 // Default to first option on error
144 }
145 input = strings.TrimSpace(input)
146
147 // Parse the selection
148 var selection int
149 if _, err := fmt.Sscanf(input, "%d", &selection); err != nil || selection < 1 || selection > len(options) {
150 fmt.Printf("Invalid selection '%s'. Please enter a number between 1 and %d.\n\n", input, len(options))
151 continue
152 }
153 return selection - 1
154 }
155}
156
157// RegisterClientWithRegion registers a new OIDC client with AWS using a specific region.
158func (c *SSOOIDCClient) RegisterClientWithRegion(ctx context.Context, region string) (*RegisterClientResponse, error) {
159 endpoint := getOIDCEndpoint(region)
160
161 payload := map[string]interface{}{
162 "clientName": "Kiro IDE",
163 "clientType": "public",
164 "scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"},
165 "grantTypes": []string{"urn:ietf:params:oauth:grant-type:device_code", "refresh_token"},
166 }
167
168 body, err := json.Marshal(payload)
169 if err != nil {
170 return nil, err
171 }
172
173 req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/client/register", strings.NewReader(string(body)))
174 if err != nil {
175 return nil, err
176 }
177 req.Header.Set("Content-Type", "application/json")
178 req.Header.Set("User-Agent", kiroUserAgent)
179
180 resp, err := c.httpClient.Do(req)
181 if err != nil {
182 return nil, err
183 }
184 defer resp.Body.Close()
185
186 respBody, err := io.ReadAll(resp.Body)
187 if err != nil {
188 return nil, err
189 }
190
191 if resp.StatusCode != http.StatusOK {
192 log.Debugf("register client failed (status %d): %s", resp.StatusCode, string(respBody))
193 return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode)
194 }
195
196 var result RegisterClientResponse
197 if err := json.Unmarshal(respBody, &result); err != nil {
198 return nil, err
199 }
200
201 return &result, nil
202}
203
204// StartDeviceAuthorizationWithIDC starts the device authorization flow for IDC.
205func (c *SSOOIDCClient) StartDeviceAuthorizationWithIDC(ctx context.Context, clientID, clientSecret, startURL, region string) (*StartDeviceAuthResponse, error) {
206 endpoint := getOIDCEndpoint(region)
207
208 payload := map[string]string{
209 "clientId": clientID,
210 "clientSecret": clientSecret,
211 "startUrl": startURL,
212 }
213
214 body, err := json.Marshal(payload)
215 if err != nil {
216 return nil, err
217 }
218
219 req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/device_authorization", strings.NewReader(string(body)))
220 if err != nil {
221 return nil, err
222 }
223 req.Header.Set("Content-Type", "application/json")
224 req.Header.Set("User-Agent", kiroUserAgent)
225
226 resp, err := c.httpClient.Do(req)
227 if err != nil {
228 return nil, err
229 }
230 defer resp.Body.Close()
231
232 respBody, err := io.ReadAll(resp.Body)
233 if err != nil {
234 return nil, err
235 }
236
237 if resp.StatusCode != http.StatusOK {
238 log.Debugf("start device auth failed (status %d): %s", resp.StatusCode, string(respBody))
239 return nil, fmt.Errorf("start device auth failed (status %d)", resp.StatusCode)
240 }
241
242 var result StartDeviceAuthResponse
243 if err := json.Unmarshal(respBody, &result); err != nil {
244 return nil, err
245 }
246
247 return &result, nil
248}
249
250// CreateTokenWithRegion polls for the access token after user authorization using a specific region.
251func (c *SSOOIDCClient) CreateTokenWithRegion(ctx context.Context, clientID, clientSecret, deviceCode, region string) (*CreateTokenResponse, error) {
252 endpoint := getOIDCEndpoint(region)
253
254 payload := map[string]string{
255 "clientId": clientID,
256 "clientSecret": clientSecret,
257 "deviceCode": deviceCode,
258 "grantType": "urn:ietf:params:oauth:grant-type:device_code",
259 }
260
261 body, err := json.Marshal(payload)
262 if err != nil {
263 return nil, err
264 }
265
266 req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/token", strings.NewReader(string(body)))
267 if err != nil {
268 return nil, err
269 }
270 req.Header.Set("Content-Type", "application/json")
271 req.Header.Set("User-Agent", kiroUserAgent)
272
273 resp, err := c.httpClient.Do(req)
274 if err != nil {
275 return nil, err
276 }
277 defer resp.Body.Close()
278
279 respBody, err := io.ReadAll(resp.Body)
280 if err != nil {
281 return nil, err
282 }
283
284 // Check for pending authorization
285 if resp.StatusCode == http.StatusBadRequest {
286 var errResp struct {
287 Error string `json:"error"`
288 }
289 if json.Unmarshal(respBody, &errResp) == nil {
290 if errResp.Error == "authorization_pending" {
291 return nil, ErrAuthorizationPending
292 }
293 if errResp.Error == "slow_down" {
294 return nil, ErrSlowDown
295 }
296 }
297 log.Debugf("create token failed: %s", string(respBody))
298 return nil, fmt.Errorf("create token failed")
299 }
300
301 if resp.StatusCode != http.StatusOK {
302 log.Debugf("create token failed (status %d): %s", resp.StatusCode, string(respBody))
303 return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode)
304 }
305
306 var result CreateTokenResponse
307 if err := json.Unmarshal(respBody, &result); err != nil {
308 return nil, err
309 }
310
311 return &result, nil
312}
313
314// RefreshTokenWithRegion refreshes an access token using the refresh token with a specific region.
315func (c *SSOOIDCClient) RefreshTokenWithRegion(ctx context.Context, clientID, clientSecret, refreshToken, region, startURL string) (*KiroTokenData, error) {
316 endpoint := getOIDCEndpoint(region)
317
318 payload := map[string]string{
319 "clientId": clientID,
320 "clientSecret": clientSecret,
321 "refreshToken": refreshToken,
322 "grantType": "refresh_token",
323 }
324
325 body, err := json.Marshal(payload)
326 if err != nil {
327 return nil, err
328 }
329
330 req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/token", strings.NewReader(string(body)))
331 if err != nil {
332 return nil, err
333 }
334
335 // Set headers matching kiro2api's IDC token refresh
336 // These headers are required for successful IDC token refresh
337 req.Header.Set("Content-Type", "application/json")
338 req.Header.Set("Host", fmt.Sprintf("oidc.%s.amazonaws.com", region))
339 req.Header.Set("Connection", "keep-alive")
340 req.Header.Set("x-amz-user-agent", idcAmzUserAgent)
341 req.Header.Set("Accept", "*/*")
342 req.Header.Set("Accept-Language", "*")
343 req.Header.Set("sec-fetch-mode", "cors")
344 req.Header.Set("User-Agent", "node")
345 req.Header.Set("Accept-Encoding", "br, gzip, deflate")
346
347 resp, err := c.httpClient.Do(req)
348 if err != nil {
349 return nil, err
350 }
351 defer resp.Body.Close()
352
353 respBody, err := io.ReadAll(resp.Body)
354 if err != nil {
355 return nil, err
356 }
357
358 if resp.StatusCode != http.StatusOK {
359 log.Warnf("IDC token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
360 return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
361 }
362
363 var result CreateTokenResponse
364 if err := json.Unmarshal(respBody, &result); err != nil {
365 return nil, err
366 }
367
368 expiresAt := time.Now().Add(time.Duration(result.ExpiresIn) * time.Second)
369
370 return &KiroTokenData{
371 AccessToken: result.AccessToken,
372 RefreshToken: result.RefreshToken,
373 ExpiresAt: expiresAt.Format(time.RFC3339),
374 AuthMethod: "idc",
375 Provider: "AWS",
376 ClientID: clientID,
377 ClientSecret: clientSecret,
378 StartURL: startURL,
379 Region: region,
380 }, nil
381}
382
383// LoginWithIDC performs the full device code flow for AWS Identity Center (IDC).
384func (c *SSOOIDCClient) LoginWithIDC(ctx context.Context, startURL, region string) (*KiroTokenData, error) {
385 fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
386 fmt.Println("║ Kiro Authentication (AWS Identity Center) ║")
387 fmt.Println("╚══════════════════════════════════════════════════════════╝")
388
389 // Step 1: Register client with the specified region
390 fmt.Println("\nRegistering client...")
391 regResp, err := c.RegisterClientWithRegion(ctx, region)
392 if err != nil {
393 return nil, fmt.Errorf("failed to register client: %w", err)
394 }
395 log.Debugf("Client registered: %s", regResp.ClientID)
396
397 // Step 2: Start device authorization with IDC start URL
398 fmt.Println("Starting device authorization...")
399 authResp, err := c.StartDeviceAuthorizationWithIDC(ctx, regResp.ClientID, regResp.ClientSecret, startURL, region)
400 if err != nil {
401 return nil, fmt.Errorf("failed to start device auth: %w", err)
402 }
403
404 // Step 3: Show user the verification URL
405 fmt.Printf("\n")
406 fmt.Println("════════════════════════════════════════════════════════════")
407 fmt.Printf(" Confirm the following code in the browser:\n")
408 fmt.Printf(" Code: %s\n", authResp.UserCode)
409 fmt.Println("════════════════════════════════════════════════════════════")
410 fmt.Printf("\n Open this URL: %s\n\n", authResp.VerificationURIComplete)
411
412 // Set incognito mode based on config
413 if c.cfg != nil {
414 browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
415 if !c.cfg.IncognitoBrowser {
416 log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.")
417 } else {
418 log.Debug("kiro: using incognito mode for multi-account support")
419 }
420 } else {
421 browser.SetIncognitoMode(true)
422 log.Debug("kiro: using incognito mode for multi-account support (default)")
423 }
424
425 // Open browser
426 if err := browser.OpenURL(authResp.VerificationURIComplete); err != nil {
427 log.Warnf("Could not open browser automatically: %v", err)
428 fmt.Println(" Please open the URL manually in your browser.")
429 } else {
430 fmt.Println(" (Browser opened automatically)")
431 }
432
433 // Step 4: Poll for token
434 fmt.Println("Waiting for authorization...")
435
436 interval := pollInterval
437 if authResp.Interval > 0 {
438 interval = time.Duration(authResp.Interval) * time.Second
439 }
440
441 deadline := time.Now().Add(time.Duration(authResp.ExpiresIn) * time.Second)
442
443 for time.Now().Before(deadline) {
444 select {
445 case <-ctx.Done():
446 browser.CloseBrowser()
447 return nil, ctx.Err()
448 case <-time.After(interval):
449 tokenResp, err := c.CreateTokenWithRegion(ctx, regResp.ClientID, regResp.ClientSecret, authResp.DeviceCode, region)
450 if err != nil {
451 if errors.Is(err, ErrAuthorizationPending) {
452 fmt.Print(".")
453 continue
454 }
455 if errors.Is(err, ErrSlowDown) {
456 interval += 5 * time.Second
457 continue
458 }
459 browser.CloseBrowser()
460 return nil, fmt.Errorf("token creation failed: %w", err)
461 }
462
463 fmt.Println("\n\n✓ Authorization successful!")
464
465 // Close the browser window
466 if err := browser.CloseBrowser(); err != nil {
467 log.Debugf("Failed to close browser: %v", err)
468 }
469
470 // Step 5: Get profile ARN from CodeWhisperer API
471 fmt.Println("Fetching profile information...")
472 profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken)
473
474 // Fetch user email
475 email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken)
476 if email != "" {
477 fmt.Printf(" Logged in as: %s\n", email)
478 }
479
480 expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
481
482 return &KiroTokenData{
483 AccessToken: tokenResp.AccessToken,
484 RefreshToken: tokenResp.RefreshToken,
485 ProfileArn: profileArn,
486 ExpiresAt: expiresAt.Format(time.RFC3339),
487 AuthMethod: "idc",
488 Provider: "AWS",
489 ClientID: regResp.ClientID,
490 ClientSecret: regResp.ClientSecret,
491 Email: email,
492 StartURL: startURL,
493 Region: region,
494 }, nil
495 }
496 }
497
498 // Close browser on timeout
499 if err := browser.CloseBrowser(); err != nil {
500 log.Debugf("Failed to close browser on timeout: %v", err)
501 }
502 return nil, fmt.Errorf("authorization timed out")
503}
504
505// LoginWithMethodSelection prompts the user to select between Builder ID and IDC, then performs the login.
506func (c *SSOOIDCClient) LoginWithMethodSelection(ctx context.Context) (*KiroTokenData, error) {
507 fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
508 fmt.Println("║ Kiro Authentication (AWS) ║")
509 fmt.Println("╚══════════════════════════════════════════════════════════╝")
510
511 // Prompt for login method
512 options := []string{
513 "Use with Builder ID (personal AWS account)",
514 "Use with IDC Account (organization SSO)",
515 }
516 selection := promptSelect("\n? Select login method:", options)
517
518 if selection == 0 {
519 // Builder ID flow - use existing implementation
520 return c.LoginWithBuilderID(ctx)
521 }
522
523 // IDC flow - prompt for start URL and region
524 fmt.Println()
525 startURL := promptInput("? Enter Start URL", "")
526 if startURL == "" {
527 return nil, fmt.Errorf("start URL is required for IDC login")
528 }
529
530 region := promptInput("? Enter Region", defaultIDCRegion)
531
532 return c.LoginWithIDC(ctx, startURL, region)
533}
534
535// RegisterClient registers a new OIDC client with AWS.
536func (c *SSOOIDCClient) RegisterClient(ctx context.Context) (*RegisterClientResponse, error) {
537 payload := map[string]interface{}{
538 "clientName": "Kiro IDE",
539 "clientType": "public",
540 "scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"},
541 "grantTypes": []string{"urn:ietf:params:oauth:grant-type:device_code", "refresh_token"},
542 }
543
544 body, err := json.Marshal(payload)
545 if err != nil {
546 return nil, err
547 }
548
549 req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/client/register", strings.NewReader(string(body)))
550 if err != nil {
551 return nil, err
552 }
553 req.Header.Set("Content-Type", "application/json")
554 req.Header.Set("User-Agent", kiroUserAgent)
555
556 resp, err := c.httpClient.Do(req)
557 if err != nil {
558 return nil, err
559 }
560 defer resp.Body.Close()
561
562 respBody, err := io.ReadAll(resp.Body)
563 if err != nil {
564 return nil, err
565 }
566
567 if resp.StatusCode != http.StatusOK {
568 log.Debugf("register client failed (status %d): %s", resp.StatusCode, string(respBody))
569 return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode)
570 }
571
572 var result RegisterClientResponse
573 if err := json.Unmarshal(respBody, &result); err != nil {
574 return nil, err
575 }
576
577 return &result, nil
578}
579
580// StartDeviceAuthorization starts the device authorization flow.
581func (c *SSOOIDCClient) StartDeviceAuthorization(ctx context.Context, clientID, clientSecret string) (*StartDeviceAuthResponse, error) {
582 payload := map[string]string{
583 "clientId": clientID,
584 "clientSecret": clientSecret,
585 "startUrl": builderIDStartURL,
586 }
587
588 body, err := json.Marshal(payload)
589 if err != nil {
590 return nil, err
591 }
592
593 req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/device_authorization", strings.NewReader(string(body)))
594 if err != nil {
595 return nil, err
596 }
597 req.Header.Set("Content-Type", "application/json")
598 req.Header.Set("User-Agent", kiroUserAgent)
599
600 resp, err := c.httpClient.Do(req)
601 if err != nil {
602 return nil, err
603 }
604 defer resp.Body.Close()
605
606 respBody, err := io.ReadAll(resp.Body)
607 if err != nil {
608 return nil, err
609 }
610
611 if resp.StatusCode != http.StatusOK {
612 log.Debugf("start device auth failed (status %d): %s", resp.StatusCode, string(respBody))
613 return nil, fmt.Errorf("start device auth failed (status %d)", resp.StatusCode)
614 }
615
616 var result StartDeviceAuthResponse
617 if err := json.Unmarshal(respBody, &result); err != nil {
618 return nil, err
619 }
620
621 return &result, nil
622}
623
624// CreateToken polls for the access token after user authorization.
625func (c *SSOOIDCClient) CreateToken(ctx context.Context, clientID, clientSecret, deviceCode string) (*CreateTokenResponse, error) {
626 payload := map[string]string{
627 "clientId": clientID,
628 "clientSecret": clientSecret,
629 "deviceCode": deviceCode,
630 "grantType": "urn:ietf:params:oauth:grant-type:device_code",
631 }
632
633 body, err := json.Marshal(payload)
634 if err != nil {
635 return nil, err
636 }
637
638 req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body)))
639 if err != nil {
640 return nil, err
641 }
642 req.Header.Set("Content-Type", "application/json")
643 req.Header.Set("User-Agent", kiroUserAgent)
644
645 resp, err := c.httpClient.Do(req)
646 if err != nil {
647 return nil, err
648 }
649 defer resp.Body.Close()
650
651 respBody, err := io.ReadAll(resp.Body)
652 if err != nil {
653 return nil, err
654 }
655
656 // Check for pending authorization
657 if resp.StatusCode == http.StatusBadRequest {
658 var errResp struct {
659 Error string `json:"error"`
660 }
661 if json.Unmarshal(respBody, &errResp) == nil {
662 if errResp.Error == "authorization_pending" {
663 return nil, ErrAuthorizationPending
664 }
665 if errResp.Error == "slow_down" {
666 return nil, ErrSlowDown
667 }
668 }
669 log.Debugf("create token failed: %s", string(respBody))
670 return nil, fmt.Errorf("create token failed")
671 }
672
673 if resp.StatusCode != http.StatusOK {
674 log.Debugf("create token failed (status %d): %s", resp.StatusCode, string(respBody))
675 return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode)
676 }
677
678 var result CreateTokenResponse
679 if err := json.Unmarshal(respBody, &result); err != nil {
680 return nil, err
681 }
682
683 return &result, nil
684}
685
686// RefreshToken refreshes an access token using the refresh token.
687func (c *SSOOIDCClient) RefreshToken(ctx context.Context, clientID, clientSecret, refreshToken string) (*KiroTokenData, error) {
688 payload := map[string]string{
689 "clientId": clientID,
690 "clientSecret": clientSecret,
691 "refreshToken": refreshToken,
692 "grantType": "refresh_token",
693 }
694
695 body, err := json.Marshal(payload)
696 if err != nil {
697 return nil, err
698 }
699
700 req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body)))
701 if err != nil {
702 return nil, err
703 }
704 req.Header.Set("Content-Type", "application/json")
705 req.Header.Set("User-Agent", kiroUserAgent)
706
707 resp, err := c.httpClient.Do(req)
708 if err != nil {
709 return nil, err
710 }
711 defer resp.Body.Close()
712
713 respBody, err := io.ReadAll(resp.Body)
714 if err != nil {
715 return nil, err
716 }
717
718 if resp.StatusCode != http.StatusOK {
719 log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
720 return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
721 }
722
723 var result CreateTokenResponse
724 if err := json.Unmarshal(respBody, &result); err != nil {
725 return nil, err
726 }
727
728 expiresAt := time.Now().Add(time.Duration(result.ExpiresIn) * time.Second)
729
730 return &KiroTokenData{
731 AccessToken: result.AccessToken,
732 RefreshToken: result.RefreshToken,
733 ExpiresAt: expiresAt.Format(time.RFC3339),
734 AuthMethod: "builder-id",
735 Provider: "AWS",
736 ClientID: clientID,
737 ClientSecret: clientSecret,
738 }, nil
739}
740
741// LoginWithBuilderID performs the full device code flow for AWS Builder ID.
742func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, error) {
743 fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
744 fmt.Println("║ Kiro Authentication (AWS Builder ID) ║")
745 fmt.Println("╚══════════════════════════════════════════════════════════╝")
746
747 // Step 1: Register client
748 fmt.Println("\nRegistering client...")
749 regResp, err := c.RegisterClient(ctx)
750 if err != nil {
751 return nil, fmt.Errorf("failed to register client: %w", err)
752 }
753 log.Debugf("Client registered: %s", regResp.ClientID)
754
755 // Step 2: Start device authorization
756 fmt.Println("Starting device authorization...")
757 authResp, err := c.StartDeviceAuthorization(ctx, regResp.ClientID, regResp.ClientSecret)
758 if err != nil {
759 return nil, fmt.Errorf("failed to start device auth: %w", err)
760 }
761
762 // Step 3: Show user the verification URL
763 fmt.Printf("\n")
764 fmt.Println("════════════════════════════════════════════════════════════")
765 fmt.Printf(" Open this URL in your browser:\n")
766 fmt.Printf(" %s\n", authResp.VerificationURIComplete)
767 fmt.Println("════════════════════════════════════════════════════════════")
768 fmt.Printf("\n Or go to: %s\n", authResp.VerificationURI)
769 fmt.Printf(" And enter code: %s\n\n", authResp.UserCode)
770
771 // Set incognito mode based on config (defaults to true for Kiro, can be overridden with --no-incognito)
772 // Incognito mode enables multi-account support by bypassing cached sessions
773 if c.cfg != nil {
774 browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
775 if !c.cfg.IncognitoBrowser {
776 log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.")
777 } else {
778 log.Debug("kiro: using incognito mode for multi-account support")
779 }
780 } else {
781 browser.SetIncognitoMode(true) // Default to incognito if no config
782 log.Debug("kiro: using incognito mode for multi-account support (default)")
783 }
784
785 // Open browser using cross-platform browser package
786 if err := browser.OpenURL(authResp.VerificationURIComplete); err != nil {
787 log.Warnf("Could not open browser automatically: %v", err)
788 fmt.Println(" Please open the URL manually in your browser.")
789 } else {
790 fmt.Println(" (Browser opened automatically)")
791 }
792
793 // Step 4: Poll for token
794 fmt.Println("Waiting for authorization...")
795
796 interval := pollInterval
797 if authResp.Interval > 0 {
798 interval = time.Duration(authResp.Interval) * time.Second
799 }
800
801 deadline := time.Now().Add(time.Duration(authResp.ExpiresIn) * time.Second)
802
803 for time.Now().Before(deadline) {
804 select {
805 case <-ctx.Done():
806 browser.CloseBrowser() // Cleanup on cancel
807 return nil, ctx.Err()
808 case <-time.After(interval):
809 tokenResp, err := c.CreateToken(ctx, regResp.ClientID, regResp.ClientSecret, authResp.DeviceCode)
810 if err != nil {
811 if errors.Is(err, ErrAuthorizationPending) {
812 fmt.Print(".")
813 continue
814 }
815 if errors.Is(err, ErrSlowDown) {
816 interval += 5 * time.Second
817 continue
818 }
819 // Close browser on error before returning
820 browser.CloseBrowser()
821 return nil, fmt.Errorf("token creation failed: %w", err)
822 }
823
824 fmt.Println("\n\n✓ Authorization successful!")
825
826 // Close the browser window
827 if err := browser.CloseBrowser(); err != nil {
828 log.Debugf("Failed to close browser: %v", err)
829 }
830
831 // Step 5: Get profile ARN from CodeWhisperer API
832 fmt.Println("Fetching profile information...")
833 profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken)
834
835 // Fetch user email (tries CodeWhisperer API first, then userinfo endpoint, then JWT parsing)
836 email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken)
837 if email != "" {
838 fmt.Printf(" Logged in as: %s\n", email)
839 }
840
841 expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
842
843 return &KiroTokenData{
844 AccessToken: tokenResp.AccessToken,
845 RefreshToken: tokenResp.RefreshToken,
846 ProfileArn: profileArn,
847 ExpiresAt: expiresAt.Format(time.RFC3339),
848 AuthMethod: "builder-id",
849 Provider: "AWS",
850 ClientID: regResp.ClientID,
851 ClientSecret: regResp.ClientSecret,
852 Email: email,
853 }, nil
854 }
855 }
856
857 // Close browser on timeout for better UX
858 if err := browser.CloseBrowser(); err != nil {
859 log.Debugf("Failed to close browser on timeout: %v", err)
860 }
861 return nil, fmt.Errorf("authorization timed out")
862}
863
864// FetchUserEmail retrieves the user's email from AWS SSO OIDC userinfo endpoint.
865// Falls back to JWT parsing if userinfo fails.
866func (c *SSOOIDCClient) FetchUserEmail(ctx context.Context, accessToken string) string {
867 // Method 1: Try userinfo endpoint (standard OIDC)
868 email := c.tryUserInfoEndpoint(ctx, accessToken)
869 if email != "" {
870 return email
871 }
872
873 // Method 2: Fallback to JWT parsing
874 return ExtractEmailFromJWT(accessToken)
875}
876
877// tryUserInfoEndpoint attempts to get user info from AWS SSO OIDC userinfo endpoint.
878func (c *SSOOIDCClient) tryUserInfoEndpoint(ctx context.Context, accessToken string) string {
879 req, err := http.NewRequestWithContext(ctx, http.MethodGet, ssoOIDCEndpoint+"/userinfo", nil)
880 if err != nil {
881 return ""
882 }
883 req.Header.Set("Authorization", "Bearer "+accessToken)
884 req.Header.Set("Accept", "application/json")
885
886 resp, err := c.httpClient.Do(req)
887 if err != nil {
888 log.Debugf("userinfo request failed: %v", err)
889 return ""
890 }
891 defer resp.Body.Close()
892
893 if resp.StatusCode != http.StatusOK {
894 respBody, _ := io.ReadAll(resp.Body)
895 log.Debugf("userinfo endpoint returned status %d: %s", resp.StatusCode, string(respBody))
896 return ""
897 }
898
899 respBody, err := io.ReadAll(resp.Body)
900 if err != nil {
901 return ""
902 }
903
904 log.Debugf("userinfo response: %s", string(respBody))
905
906 var userInfo struct {
907 Email string `json:"email"`
908 Sub string `json:"sub"`
909 PreferredUsername string `json:"preferred_username"`
910 Name string `json:"name"`
911 }
912
913 if err := json.Unmarshal(respBody, &userInfo); err != nil {
914 return ""
915 }
916
917 if userInfo.Email != "" {
918 return userInfo.Email
919 }
920 if userInfo.PreferredUsername != "" && strings.Contains(userInfo.PreferredUsername, "@") {
921 return userInfo.PreferredUsername
922 }
923 return ""
924}
925
926// fetchProfileArn retrieves the profile ARN from CodeWhisperer API.
927// This is needed for file naming since AWS SSO OIDC doesn't return profile info.
928func (c *SSOOIDCClient) fetchProfileArn(ctx context.Context, accessToken string) string {
929 // Try ListProfiles API first
930 profileArn := c.tryListProfiles(ctx, accessToken)
931 if profileArn != "" {
932 return profileArn
933 }
934
935 // Fallback: Try ListAvailableCustomizations
936 return c.tryListCustomizations(ctx, accessToken)
937}
938
939func (c *SSOOIDCClient) tryListProfiles(ctx context.Context, accessToken string) string {
940 payload := map[string]interface{}{
941 "origin": "AI_EDITOR",
942 }
943
944 body, err := json.Marshal(payload)
945 if err != nil {
946 return ""
947 }
948
949 req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body)))
950 if err != nil {
951 return ""
952 }
953
954 req.Header.Set("Content-Type", "application/x-amz-json-1.0")
955 req.Header.Set("x-amz-target", "AmazonCodeWhispererService.ListProfiles")
956 req.Header.Set("Authorization", "Bearer "+accessToken)
957 req.Header.Set("Accept", "application/json")
958
959 resp, err := c.httpClient.Do(req)
960 if err != nil {
961 return ""
962 }
963 defer resp.Body.Close()
964
965 respBody, _ := io.ReadAll(resp.Body)
966
967 if resp.StatusCode != http.StatusOK {
968 log.Debugf("ListProfiles failed (status %d): %s", resp.StatusCode, string(respBody))
969 return ""
970 }
971
972 log.Debugf("ListProfiles response: %s", string(respBody))
973
974 var result struct {
975 Profiles []struct {
976 Arn string `json:"arn"`
977 } `json:"profiles"`
978 ProfileArn string `json:"profileArn"`
979 }
980
981 if err := json.Unmarshal(respBody, &result); err != nil {
982 return ""
983 }
984
985 if result.ProfileArn != "" {
986 return result.ProfileArn
987 }
988
989 if len(result.Profiles) > 0 {
990 return result.Profiles[0].Arn
991 }
992
993 return ""
994}
995
996func (c *SSOOIDCClient) tryListCustomizations(ctx context.Context, accessToken string) string {
997 payload := map[string]interface{}{
998 "origin": "AI_EDITOR",
999 }
1000
1001 body, err := json.Marshal(payload)
1002 if err != nil {
1003 return ""
1004 }
1005
1006 req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body)))
1007 if err != nil {
1008 return ""
1009 }
1010
1011 req.Header.Set("Content-Type", "application/x-amz-json-1.0")
1012 req.Header.Set("x-amz-target", "AmazonCodeWhispererService.ListAvailableCustomizations")
1013 req.Header.Set("Authorization", "Bearer "+accessToken)
1014 req.Header.Set("Accept", "application/json")
1015
1016 resp, err := c.httpClient.Do(req)
1017 if err != nil {
1018 return ""
1019 }
1020 defer resp.Body.Close()
1021
1022 respBody, _ := io.ReadAll(resp.Body)
1023
1024 if resp.StatusCode != http.StatusOK {
1025 log.Debugf("ListAvailableCustomizations failed (status %d): %s", resp.StatusCode, string(respBody))
1026 return ""
1027 }
1028
1029 log.Debugf("ListAvailableCustomizations response: %s", string(respBody))
1030
1031 var result struct {
1032 Customizations []struct {
1033 Arn string `json:"arn"`
1034 } `json:"customizations"`
1035 ProfileArn string `json:"profileArn"`
1036 }
1037
1038 if err := json.Unmarshal(respBody, &result); err != nil {
1039 return ""
1040 }
1041
1042 if result.ProfileArn != "" {
1043 return result.ProfileArn
1044 }
1045
1046 if len(result.Customizations) > 0 {
1047 return result.Customizations[0].Arn
1048 }
1049
1050 return ""
1051}
1052
1053// RegisterClientForAuthCode registers a new OIDC client for authorization code flow.
1054func (c *SSOOIDCClient) RegisterClientForAuthCode(ctx context.Context, redirectURI string) (*RegisterClientResponse, error) {
1055 payload := map[string]interface{}{
1056 "clientName": "Kiro IDE",
1057 "clientType": "public",
1058 "scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"},
1059 "grantTypes": []string{"authorization_code", "refresh_token"},
1060 "redirectUris": []string{redirectURI},
1061 "issuerUrl": builderIDStartURL,
1062 }
1063
1064 body, err := json.Marshal(payload)
1065 if err != nil {
1066 return nil, err
1067 }
1068
1069 req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/client/register", strings.NewReader(string(body)))
1070 if err != nil {
1071 return nil, err
1072 }
1073 req.Header.Set("Content-Type", "application/json")
1074 req.Header.Set("User-Agent", kiroUserAgent)
1075
1076 resp, err := c.httpClient.Do(req)
1077 if err != nil {
1078 return nil, err
1079 }
1080 defer resp.Body.Close()
1081
1082 respBody, err := io.ReadAll(resp.Body)
1083 if err != nil {
1084 return nil, err
1085 }
1086
1087 if resp.StatusCode != http.StatusOK {
1088 log.Debugf("register client for auth code failed (status %d): %s", resp.StatusCode, string(respBody))
1089 return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode)
1090 }
1091
1092 var result RegisterClientResponse
1093 if err := json.Unmarshal(respBody, &result); err != nil {
1094 return nil, err
1095 }
1096
1097 return &result, nil
1098}
1099
1100// AuthCodeCallbackResult contains the result from authorization code callback.
1101type AuthCodeCallbackResult struct {
1102 Code string
1103 State string
1104 Error string
1105}
1106
1107// startAuthCodeCallbackServer starts a local HTTP server to receive the authorization code callback.
1108func (c *SSOOIDCClient) startAuthCodeCallbackServer(ctx context.Context, expectedState string) (string, <-chan AuthCodeCallbackResult, error) {
1109 // Try to find an available port
1110 listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", authCodeCallbackPort))
1111 if err != nil {
1112 // Try with dynamic port
1113 log.Warnf("sso oidc: default port %d is busy, falling back to dynamic port", authCodeCallbackPort)
1114 listener, err = net.Listen("tcp", "127.0.0.1:0")
1115 if err != nil {
1116 return "", nil, fmt.Errorf("failed to start callback server: %w", err)
1117 }
1118 }
1119
1120 port := listener.Addr().(*net.TCPAddr).Port
1121 redirectURI := fmt.Sprintf("http://127.0.0.1:%d%s", port, authCodeCallbackPath)
1122 resultChan := make(chan AuthCodeCallbackResult, 1)
1123
1124 server := &http.Server{
1125 ReadHeaderTimeout: 10 * time.Second,
1126 }
1127
1128 mux := http.NewServeMux()
1129 mux.HandleFunc(authCodeCallbackPath, func(w http.ResponseWriter, r *http.Request) {
1130 code := r.URL.Query().Get("code")
1131 state := r.URL.Query().Get("state")
1132 errParam := r.URL.Query().Get("error")
1133
1134 // Send response to browser
1135 w.Header().Set("Content-Type", "text/html; charset=utf-8")
1136 if errParam != "" {
1137 w.WriteHeader(http.StatusBadRequest)
1138 fmt.Fprintf(w, `<!DOCTYPE html>
1139<html><head><title>Login Failed</title></head>
1140<body><h1>Login Failed</h1><p>Error: %s</p><p>You can close this window.</p></body></html>`, html.EscapeString(errParam))
1141 resultChan <- AuthCodeCallbackResult{Error: errParam}
1142 return
1143 }
1144
1145 if state != expectedState {
1146 w.WriteHeader(http.StatusBadRequest)
1147 fmt.Fprint(w, `<!DOCTYPE html>
1148<html><head><title>Login Failed</title></head>
1149<body><h1>Login Failed</h1><p>Invalid state parameter</p><p>You can close this window.</p></body></html>`)
1150 resultChan <- AuthCodeCallbackResult{Error: "state mismatch"}
1151 return
1152 }
1153
1154 fmt.Fprint(w, `<!DOCTYPE html>
1155<html><head><title>Login Successful</title></head>
1156<body><h1>Login Successful!</h1><p>You can close this window and return to the terminal.</p>
1157<script>window.close();</script></body></html>`)
1158 resultChan <- AuthCodeCallbackResult{Code: code, State: state}
1159 })
1160
1161 server.Handler = mux
1162
1163 go func() {
1164 if err := server.Serve(listener); err != nil && err != http.ErrServerClosed {
1165 log.Debugf("auth code callback server error: %v", err)
1166 }
1167 }()
1168
1169 go func() {
1170 select {
1171 case <-ctx.Done():
1172 case <-time.After(10 * time.Minute):
1173 case <-resultChan:
1174 }
1175 _ = server.Shutdown(context.Background())
1176 }()
1177
1178 return redirectURI, resultChan, nil
1179}
1180
1181// generatePKCEForAuthCode generates PKCE code verifier and challenge for authorization code flow.
1182func generatePKCEForAuthCode() (verifier, challenge string, err error) {
1183 b := make([]byte, 32)
1184 if _, err := rand.Read(b); err != nil {
1185 return "", "", fmt.Errorf("failed to generate random bytes: %w", err)
1186 }
1187 verifier = base64.RawURLEncoding.EncodeToString(b)
1188 h := sha256.Sum256([]byte(verifier))
1189 challenge = base64.RawURLEncoding.EncodeToString(h[:])
1190 return verifier, challenge, nil
1191}
1192
1193// generateStateForAuthCode generates a random state parameter.
1194func generateStateForAuthCode() (string, error) {
1195 b := make([]byte, 16)
1196 if _, err := rand.Read(b); err != nil {
1197 return "", err
1198 }
1199 return base64.RawURLEncoding.EncodeToString(b), nil
1200}
1201
1202// CreateTokenWithAuthCode exchanges authorization code for tokens.
1203func (c *SSOOIDCClient) CreateTokenWithAuthCode(ctx context.Context, clientID, clientSecret, code, codeVerifier, redirectURI string) (*CreateTokenResponse, error) {
1204 payload := map[string]string{
1205 "clientId": clientID,
1206 "clientSecret": clientSecret,
1207 "code": code,
1208 "codeVerifier": codeVerifier,
1209 "redirectUri": redirectURI,
1210 "grantType": "authorization_code",
1211 }
1212
1213 body, err := json.Marshal(payload)
1214 if err != nil {
1215 return nil, err
1216 }
1217
1218 req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body)))
1219 if err != nil {
1220 return nil, err
1221 }
1222 req.Header.Set("Content-Type", "application/json")
1223 req.Header.Set("User-Agent", kiroUserAgent)
1224
1225 resp, err := c.httpClient.Do(req)
1226 if err != nil {
1227 return nil, err
1228 }
1229 defer resp.Body.Close()
1230
1231 respBody, err := io.ReadAll(resp.Body)
1232 if err != nil {
1233 return nil, err
1234 }
1235
1236 if resp.StatusCode != http.StatusOK {
1237 log.Debugf("create token with auth code failed (status %d): %s", resp.StatusCode, string(respBody))
1238 return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode)
1239 }
1240
1241 var result CreateTokenResponse
1242 if err := json.Unmarshal(respBody, &result); err != nil {
1243 return nil, err
1244 }
1245
1246 return &result, nil
1247}
1248
1249// LoginWithBuilderIDAuthCode performs the authorization code flow for AWS Builder ID.
1250// This provides a better UX than device code flow as it uses automatic browser callback.
1251func (c *SSOOIDCClient) LoginWithBuilderIDAuthCode(ctx context.Context) (*KiroTokenData, error) {
1252 fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
1253 fmt.Println("║ Kiro Authentication (AWS Builder ID - Auth Code) ║")
1254 fmt.Println("╚══════════════════════════════════════════════════════════╝")
1255
1256 // Step 1: Generate PKCE and state
1257 codeVerifier, codeChallenge, err := generatePKCEForAuthCode()
1258 if err != nil {
1259 return nil, fmt.Errorf("failed to generate PKCE: %w", err)
1260 }
1261
1262 state, err := generateStateForAuthCode()
1263 if err != nil {
1264 return nil, fmt.Errorf("failed to generate state: %w", err)
1265 }
1266
1267 // Step 2: Start callback server
1268 fmt.Println("\nStarting callback server...")
1269 redirectURI, resultChan, err := c.startAuthCodeCallbackServer(ctx, state)
1270 if err != nil {
1271 return nil, fmt.Errorf("failed to start callback server: %w", err)
1272 }
1273 log.Debugf("Callback server started, redirect URI: %s", redirectURI)
1274
1275 // Step 3: Register client with auth code grant type
1276 fmt.Println("Registering client...")
1277 regResp, err := c.RegisterClientForAuthCode(ctx, redirectURI)
1278 if err != nil {
1279 return nil, fmt.Errorf("failed to register client: %w", err)
1280 }
1281 log.Debugf("Client registered: %s", regResp.ClientID)
1282
1283 // Step 4: Build authorization URL
1284 scopes := "codewhisperer:completions,codewhisperer:analysis,codewhisperer:conversations"
1285 authURL := fmt.Sprintf("%s/authorize?response_type=code&client_id=%s&redirect_uri=%s&scopes=%s&state=%s&code_challenge=%s&code_challenge_method=S256",
1286 ssoOIDCEndpoint,
1287 regResp.ClientID,
1288 redirectURI,
1289 scopes,
1290 state,
1291 codeChallenge,
1292 )
1293
1294 // Step 5: Open browser
1295 fmt.Println("\n════════════════════════════════════════════════════════════")
1296 fmt.Println(" Opening browser for authentication...")
1297 fmt.Println("════════════════════════════════════════════════════════════")
1298 fmt.Printf("\n URL: %s\n\n", authURL)
1299
1300 // Set incognito mode
1301 if c.cfg != nil {
1302 browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
1303 } else {
1304 browser.SetIncognitoMode(true)
1305 }
1306
1307 if err := browser.OpenURL(authURL); err != nil {
1308 log.Warnf("Could not open browser automatically: %v", err)
1309 fmt.Println(" ⚠ Could not open browser automatically.")
1310 fmt.Println(" Please open the URL above in your browser manually.")
1311 } else {
1312 fmt.Println(" (Browser opened automatically)")
1313 }
1314
1315 fmt.Println("\n Waiting for authorization callback...")
1316
1317 // Step 6: Wait for callback
1318 select {
1319 case <-ctx.Done():
1320 browser.CloseBrowser()
1321 return nil, ctx.Err()
1322 case <-time.After(10 * time.Minute):
1323 browser.CloseBrowser()
1324 return nil, fmt.Errorf("authorization timed out")
1325 case result := <-resultChan:
1326 if result.Error != "" {
1327 browser.CloseBrowser()
1328 return nil, fmt.Errorf("authorization failed: %s", result.Error)
1329 }
1330
1331 fmt.Println("\n✓ Authorization received!")
1332
1333 // Close browser
1334 if err := browser.CloseBrowser(); err != nil {
1335 log.Debugf("Failed to close browser: %v", err)
1336 }
1337
1338 // Step 7: Exchange code for tokens
1339 fmt.Println("Exchanging code for tokens...")
1340 tokenResp, err := c.CreateTokenWithAuthCode(ctx, regResp.ClientID, regResp.ClientSecret, result.Code, codeVerifier, redirectURI)
1341 if err != nil {
1342 return nil, fmt.Errorf("failed to exchange code for tokens: %w", err)
1343 }
1344
1345 fmt.Println("\n✓ Authentication successful!")
1346
1347 // Step 8: Get profile ARN
1348 fmt.Println("Fetching profile information...")
1349 profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken)
1350
1351 // Fetch user email (tries CodeWhisperer API first, then userinfo endpoint, then JWT parsing)
1352 email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken)
1353 if email != "" {
1354 fmt.Printf(" Logged in as: %s\n", email)
1355 }
1356
1357 expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
1358
1359 return &KiroTokenData{
1360 AccessToken: tokenResp.AccessToken,
1361 RefreshToken: tokenResp.RefreshToken,
1362 ProfileArn: profileArn,
1363 ExpiresAt: expiresAt.Format(time.RFC3339),
1364 AuthMethod: "builder-id",
1365 Provider: "AWS",
1366 ClientID: regResp.ClientID,
1367 ClientSecret: regResp.ClientSecret,
1368 Email: email,
1369 }, nil
1370 }
1371}
internal/auth/kiro/token.go+72−0
@@ -0,0 +1,72 @@
1package kiro
2
3import (
4 "encoding/json"
5 "fmt"
6 "os"
7 "path/filepath"
8)
9
10// KiroTokenStorage holds the persistent token data for Kiro authentication.
11type KiroTokenStorage struct {
12 // AccessToken is the OAuth2 access token for API access
13 AccessToken string `json:"access_token"`
14 // RefreshToken is used to obtain new access tokens
15 RefreshToken string `json:"refresh_token"`
16 // ProfileArn is the AWS CodeWhisperer profile ARN
17 ProfileArn string `json:"profile_arn"`
18 // ExpiresAt is the timestamp when the token expires
19 ExpiresAt string `json:"expires_at"`
20 // AuthMethod indicates the authentication method used
21 AuthMethod string `json:"auth_method"`
22 // Provider indicates the OAuth provider
23 Provider string `json:"provider"`
24 // LastRefresh is the timestamp of the last token refresh
25 LastRefresh string `json:"last_refresh"`
26}
27
28// SaveTokenToFile persists the token storage to the specified file path.
29func (s *KiroTokenStorage) SaveTokenToFile(authFilePath string) error {
30 dir := filepath.Dir(authFilePath)
31 if err := os.MkdirAll(dir, 0700); err != nil {
32 return fmt.Errorf("failed to create directory: %w", err)
33 }
34
35 data, err := json.MarshalIndent(s, "", " ")
36 if err != nil {
37 return fmt.Errorf("failed to marshal token storage: %w", err)
38 }
39
40 if err := os.WriteFile(authFilePath, data, 0600); err != nil {
41 return fmt.Errorf("failed to write token file: %w", err)
42 }
43
44 return nil
45}
46
47// LoadFromFile loads token storage from the specified file path.
48func LoadFromFile(authFilePath string) (*KiroTokenStorage, error) {
49 data, err := os.ReadFile(authFilePath)
50 if err != nil {
51 return nil, fmt.Errorf("failed to read token file: %w", err)
52 }
53
54 var storage KiroTokenStorage
55 if err := json.Unmarshal(data, &storage); err != nil {
56 return nil, fmt.Errorf("failed to parse token file: %w", err)
57 }
58
59 return &storage, nil
60}
61
62// ToTokenData converts storage to KiroTokenData for API use.
63func (s *KiroTokenStorage) ToTokenData() *KiroTokenData {
64 return &KiroTokenData{
65 AccessToken: s.AccessToken,
66 RefreshToken: s.RefreshToken,
67 ProfileArn: s.ProfileArn,
68 ExpiresAt: s.ExpiresAt,
69 AuthMethod: s.AuthMethod,
70 Provider: s.Provider,
71 }
72}
internal/browser/browser.go+416−14
@@ -6,14 +6,49 @@ import ( · 5 unmodified lines
6 "fmt"6 "fmt"
7 "os/exec"7 "os/exec"
8 "runtime"8 "runtime"
9 "strings"
10 "sync"
9 11
12 pkgbrowser "github.com/pkg/browser"
10 log "github.com/sirupsen/logrus"13 log "github.com/sirupsen/logrus"
11 "github.com/skratchdot/open-golang/open"
12)14)
13 15
16// incognitoMode controls whether to open URLs in incognito/private mode.
17// This is useful for OAuth flows where you want to use a different account.
18var incognitoMode bool
19
20// lastBrowserProcess stores the last opened browser process for cleanup
21var lastBrowserProcess *exec.Cmd
22var browserMutex sync.Mutex
23
24// SetIncognitoMode enables or disables incognito/private browsing mode.
25func SetIncognitoMode(enabled bool) {
26 incognitoMode = enabled
27}
28
29// IsIncognitoMode returns whether incognito mode is enabled.
30func IsIncognitoMode() bool {
31 return incognitoMode
32}
33
34// CloseBrowser closes the last opened browser process.
35func CloseBrowser() error {
36 browserMutex.Lock()
37 defer browserMutex.Unlock()
38
39 if lastBrowserProcess == nil || lastBrowserProcess.Process == nil {
40 return nil
41 }
42
43 err := lastBrowserProcess.Process.Kill()
44 lastBrowserProcess = nil
45 return err
46}
47
14// OpenURL opens the specified URL in the default web browser.48// OpenURL opens the specified URL in the default web browser.
15// It first attempts to use a platform-agnostic library and falls back to49// It uses the pkg/browser library which provides robust cross-platform support
16// platform-specific commands if that fails.50// for Windows, macOS, and Linux.
51// If incognito mode is enabled, it will open in a private/incognito window.
17//52//
18// Parameters:53// Parameters:
19// - url: The URL to open.54// - url: The URL to open.
@@ -21,16 +56,22 @@ import ( · 1 unmodified line
21// Returns:56// Returns:
22// - An error if the URL cannot be opened, otherwise nil.57// - An error if the URL cannot be opened, otherwise nil.
23func OpenURL(url string) error {58func OpenURL(url string) error {
24 fmt.Printf("Attempting to open URL in browser: %s\n", url)59 log.Debugf("Opening URL in browser: %s (incognito=%v)", url, incognitoMode)
25 60
26 // Try using the open-golang library first61 // If incognito mode is enabled, use platform-specific incognito commands
27 err := open.Run(url)62 if incognitoMode {
63 log.Debug("Using incognito mode")
64 return openURLIncognito(url)
65 }
66
67 // Use pkg/browser for cross-platform support
68 err := pkgbrowser.OpenURL(url)
28 if err == nil {69 if err == nil {
29 log.Debug("Successfully opened URL using open-golang library")70 log.Debug("Successfully opened URL using pkg/browser library")
30 return nil71 return nil
31 }72 }
32 73
33 log.Debugf("open-golang failed: %v, trying platform-specific commands", err)74 log.Debugf("pkg/browser failed: %v, trying platform-specific commands", err)
34 75
35 // Fallback to platform-specific commands76 // Fallback to platform-specific commands
36 return openURLPlatformSpecific(url)77 return openURLPlatformSpecific(url)
@@ -78,18 +119,379 @@ func openURLPlatformSpecific(url string) error { · 41 unmodified lines
78 return nil119 return nil
79}120}
80 121
122// openURLIncognito opens a URL in incognito/private browsing mode.
123// It first tries to detect the default browser and use its incognito flag.
124// Falls back to a chain of known browsers if detection fails.
125//
126// Parameters:
127// - url: The URL to open.
128//
129// Returns:
130// - An error if the URL cannot be opened, otherwise nil.
131func openURLIncognito(url string) error {
132 // First, try to detect and use the default browser
133 if cmd := tryDefaultBrowserIncognito(url); cmd != nil {
134 log.Debugf("Using detected default browser: %s %v", cmd.Path, cmd.Args[1:])
135 if err := cmd.Start(); err == nil {
136 storeBrowserProcess(cmd)
137 log.Debug("Successfully opened URL in default browser's incognito mode")
138 return nil
139 }
140 log.Debugf("Failed to start default browser, trying fallback chain")
141 }
142
143 // Fallback to known browser chain
144 cmd := tryFallbackBrowsersIncognito(url)
145 if cmd == nil {
146 log.Warn("No browser with incognito support found, falling back to normal mode")
147 return openURLPlatformSpecific(url)
148 }
149
150 log.Debugf("Running incognito command: %s %v", cmd.Path, cmd.Args[1:])
151 err := cmd.Start()
152 if err != nil {
153 log.Warnf("Failed to open incognito browser: %v, falling back to normal mode", err)
154 return openURLPlatformSpecific(url)
155 }
156
157 storeBrowserProcess(cmd)
158 log.Debug("Successfully opened URL in incognito/private mode")
159 return nil
160}
161
162// storeBrowserProcess safely stores the browser process for later cleanup.
163func storeBrowserProcess(cmd *exec.Cmd) {
164 browserMutex.Lock()
165 lastBrowserProcess = cmd
166 browserMutex.Unlock()
167}
168
169// tryDefaultBrowserIncognito attempts to detect the default browser and return
170// an exec.Cmd configured with the appropriate incognito flag.
171func tryDefaultBrowserIncognito(url string) *exec.Cmd {
172 switch runtime.GOOS {
173 case "darwin":
174 return tryDefaultBrowserMacOS(url)
175 case "windows":
176 return tryDefaultBrowserWindows(url)
177 case "linux":
178 return tryDefaultBrowserLinux(url)
179 }
180 return nil
181}
182
183// tryDefaultBrowserMacOS detects the default browser on macOS.
184func tryDefaultBrowserMacOS(url string) *exec.Cmd {
185 // Try to get default browser from Launch Services
186 out, err := exec.Command("defaults", "read", "com.apple.LaunchServices/com.apple.launchservices.secure", "LSHandlers").Output()
187 if err != nil {
188 return nil
189 }
190
191 output := string(out)
192 var browserName string
193
194 // Parse the output to find the http/https handler
195 if containsBrowserID(output, "com.google.chrome") {
196 browserName = "chrome"
197 } else if containsBrowserID(output, "org.mozilla.firefox") {
198 browserName = "firefox"
199 } else if containsBrowserID(output, "com.apple.safari") {
200 browserName = "safari"
201 } else if containsBrowserID(output, "com.brave.browser") {
202 browserName = "brave"
203 } else if containsBrowserID(output, "com.microsoft.edgemac") {
204 browserName = "edge"
205 }
206
207 return createMacOSIncognitoCmd(browserName, url)
208}
209
210// containsBrowserID checks if the LaunchServices output contains a browser ID.
211func containsBrowserID(output, bundleID string) bool {
212 return strings.Contains(output, bundleID)
213}
214
215// createMacOSIncognitoCmd creates the appropriate incognito command for macOS browsers.
216func createMacOSIncognitoCmd(browserName, url string) *exec.Cmd {
217 switch browserName {
218 case "chrome":
219 // Try direct path first
220 chromePath := "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"
221 if _, err := exec.LookPath(chromePath); err == nil {
222 return exec.Command(chromePath, "--incognito", url)
223 }
224 return exec.Command("open", "-na", "Google Chrome", "--args", "--incognito", url)
225 case "firefox":
226 return exec.Command("open", "-na", "Firefox", "--args", "--private-window", url)
227 case "safari":
228 // Safari doesn't have CLI incognito, try AppleScript
229 return tryAppleScriptSafariPrivate(url)
230 case "brave":
231 return exec.Command("open", "-na", "Brave Browser", "--args", "--incognito", url)
232 case "edge":
233 return exec.Command("open", "-na", "Microsoft Edge", "--args", "--inprivate", url)
234 }
235 return nil
236}
237
238// tryAppleScriptSafariPrivate attempts to open Safari in private browsing mode using AppleScript.
239func tryAppleScriptSafariPrivate(url string) *exec.Cmd {
240 // AppleScript to open a new private window in Safari
241 script := fmt.Sprintf(`
242 tell application "Safari"
243 activate
244 tell application "System Events"
245 keystroke "n" using {command down, shift down}
246 delay 0.5
247 end tell
248 set URL of document 1 to "%s"
249 end tell
250 `, url)
251
252 cmd := exec.Command("osascript", "-e", script)
253 // Test if this approach works by checking if Safari is available
254 if _, err := exec.LookPath("/Applications/Safari.app/Contents/MacOS/Safari"); err != nil {
255 log.Debug("Safari not found, AppleScript private window not available")
256 return nil
257 }
258 log.Debug("Attempting Safari private window via AppleScript")
259 return cmd
260}
261
262// tryDefaultBrowserWindows detects the default browser on Windows via registry.
263func tryDefaultBrowserWindows(url string) *exec.Cmd {
264 // Query registry for default browser
265 out, err := exec.Command("reg", "query",
266 `HKEY_CURRENT_USER\Software\Microsoft\Windows\Shell\Associations\UrlAssociations\http\UserChoice`,
267 "/v", "ProgId").Output()
268 if err != nil {
269 return nil
270 }
271
272 output := string(out)
273 var browserName string
274
275 // Map ProgId to browser name
276 if strings.Contains(output, "ChromeHTML") {
277 browserName = "chrome"
278 } else if strings.Contains(output, "FirefoxURL") {
279 browserName = "firefox"
280 } else if strings.Contains(output, "MSEdgeHTM") {
281 browserName = "edge"
282 } else if strings.Contains(output, "BraveHTML") {
283 browserName = "brave"
284 }
285
286 return createWindowsIncognitoCmd(browserName, url)
287}
288
289// createWindowsIncognitoCmd creates the appropriate incognito command for Windows browsers.
290func createWindowsIncognitoCmd(browserName, url string) *exec.Cmd {
291 switch browserName {
292 case "chrome":
293 paths := []string{
294 "chrome",
295 `C:\Program Files\Google\Chrome\Application\chrome.exe`,
296 `C:\Program Files (x86)\Google\Chrome\Application\chrome.exe`,
297 }
298 for _, p := range paths {
299 if _, err := exec.LookPath(p); err == nil {
300 return exec.Command(p, "--incognito", url)
301 }
302 }
303 case "firefox":
304 if path, err := exec.LookPath("firefox"); err == nil {
305 return exec.Command(path, "--private-window", url)
306 }
307 case "edge":
308 paths := []string{
309 "msedge",
310 `C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe`,
311 `C:\Program Files\Microsoft\Edge\Application\msedge.exe`,
312 }
313 for _, p := range paths {
314 if _, err := exec.LookPath(p); err == nil {
315 return exec.Command(p, "--inprivate", url)
316 }
317 }
318 case "brave":
319 paths := []string{
320 `C:\Program Files\BraveSoftware\Brave-Browser\Application\brave.exe`,
321 `C:\Program Files (x86)\BraveSoftware\Brave-Browser\Application\brave.exe`,
322 }
323 for _, p := range paths {
324 if _, err := exec.LookPath(p); err == nil {
325 return exec.Command(p, "--incognito", url)
326 }
327 }
328 }
329 return nil
330}
331
332// tryDefaultBrowserLinux detects the default browser on Linux using xdg-settings.
333func tryDefaultBrowserLinux(url string) *exec.Cmd {
334 out, err := exec.Command("xdg-settings", "get", "default-web-browser").Output()
335 if err != nil {
336 return nil
337 }
338
339 desktop := string(out)
340 var browserName string
341
342 // Map .desktop file to browser name
343 if strings.Contains(desktop, "google-chrome") || strings.Contains(desktop, "chrome") {
344 browserName = "chrome"
345 } else if strings.Contains(desktop, "firefox") {
346 browserName = "firefox"
347 } else if strings.Contains(desktop, "chromium") {
348 browserName = "chromium"
349 } else if strings.Contains(desktop, "brave") {
350 browserName = "brave"
351 } else if strings.Contains(desktop, "microsoft-edge") || strings.Contains(desktop, "msedge") {
352 browserName = "edge"
353 }
354
355 return createLinuxIncognitoCmd(browserName, url)
356}
357
358// createLinuxIncognitoCmd creates the appropriate incognito command for Linux browsers.
359func createLinuxIncognitoCmd(browserName, url string) *exec.Cmd {
360 switch browserName {
361 case "chrome":
362 paths := []string{"google-chrome", "google-chrome-stable"}
363 for _, p := range paths {
364 if path, err := exec.LookPath(p); err == nil {
365 return exec.Command(path, "--incognito", url)
366 }
367 }
368 case "firefox":
369 paths := []string{"firefox", "firefox-esr"}
370 for _, p := range paths {
371 if path, err := exec.LookPath(p); err == nil {
372 return exec.Command(path, "--private-window", url)
373 }
374 }
375 case "chromium":
376 paths := []string{"chromium", "chromium-browser"}
377 for _, p := range paths {
378 if path, err := exec.LookPath(p); err == nil {
379 return exec.Command(path, "--incognito", url)
380 }
381 }
382 case "brave":
383 if path, err := exec.LookPath("brave-browser"); err == nil {
384 return exec.Command(path, "--incognito", url)
385 }
386 case "edge":
387 if path, err := exec.LookPath("microsoft-edge"); err == nil {
388 return exec.Command(path, "--inprivate", url)
389 }
390 }
391 return nil
392}
393
394// tryFallbackBrowsersIncognito tries a chain of known browsers as fallback.
395func tryFallbackBrowsersIncognito(url string) *exec.Cmd {
396 switch runtime.GOOS {
397 case "darwin":
398 return tryFallbackBrowsersMacOS(url)
399 case "windows":
400 return tryFallbackBrowsersWindows(url)
401 case "linux":
402 return tryFallbackBrowsersLinuxChain(url)
403 }
404 return nil
405}
406
407// tryFallbackBrowsersMacOS tries known browsers on macOS.
408func tryFallbackBrowsersMacOS(url string) *exec.Cmd {
409 // Try Chrome
410 chromePath := "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"
411 if _, err := exec.LookPath(chromePath); err == nil {
412 return exec.Command(chromePath, "--incognito", url)
413 }
414 // Try Firefox
415 if _, err := exec.LookPath("/Applications/Firefox.app/Contents/MacOS/firefox"); err == nil {
416 return exec.Command("open", "-na", "Firefox", "--args", "--private-window", url)
417 }
418 // Try Brave
419 if _, err := exec.LookPath("/Applications/Brave Browser.app/Contents/MacOS/Brave Browser"); err == nil {
420 return exec.Command("open", "-na", "Brave Browser", "--args", "--incognito", url)
421 }
422 // Try Edge
423 if _, err := exec.LookPath("/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge"); err == nil {
424 return exec.Command("open", "-na", "Microsoft Edge", "--args", "--inprivate", url)
425 }
426 // Last resort: try Safari with AppleScript
427 if cmd := tryAppleScriptSafariPrivate(url); cmd != nil {
428 log.Info("Using Safari with AppleScript for private browsing (may require accessibility permissions)")
429 return cmd
430 }
431 return nil
432}
433
434// tryFallbackBrowsersWindows tries known browsers on Windows.
435func tryFallbackBrowsersWindows(url string) *exec.Cmd {
436 // Chrome
437 chromePaths := []string{
438 "chrome",
439 `C:\Program Files\Google\Chrome\Application\chrome.exe`,
440 `C:\Program Files (x86)\Google\Chrome\Application\chrome.exe`,
441 }
442 for _, p := range chromePaths {
443 if _, err := exec.LookPath(p); err == nil {
444 return exec.Command(p, "--incognito", url)
445 }
446 }
447 // Firefox
448 if path, err := exec.LookPath("firefox"); err == nil {
449 return exec.Command(path, "--private-window", url)
450 }
451 // Edge (usually available on Windows 10+)
452 edgePaths := []string{
453 "msedge",
454 `C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe`,
455 `C:\Program Files\Microsoft\Edge\Application\msedge.exe`,
456 }
457 for _, p := range edgePaths {
458 if _, err := exec.LookPath(p); err == nil {
459 return exec.Command(p, "--inprivate", url)
460 }
461 }
462 return nil
463}
464
465// tryFallbackBrowsersLinuxChain tries known browsers on Linux.
466func tryFallbackBrowsersLinuxChain(url string) *exec.Cmd {
467 type browserConfig struct {
468 name string
469 flag string
470 }
471 browsers := []browserConfig{
472 {"google-chrome", "--incognito"},
473 {"google-chrome-stable", "--incognito"},
474 {"chromium", "--incognito"},
475 {"chromium-browser", "--incognito"},
476 {"firefox", "--private-window"},
477 {"firefox-esr", "--private-window"},
478 {"brave-browser", "--incognito"},
479 {"microsoft-edge", "--inprivate"},
480 }
481 for _, b := range browsers {
482 if path, err := exec.LookPath(b.name); err == nil {
483 return exec.Command(path, b.flag, url)
484 }
485 }
486 return nil
487}
488
81// IsAvailable checks if the system has a command available to open a web browser.489// IsAvailable checks if the system has a command available to open a web browser.
82// It verifies the presence of necessary commands for the current operating system.490// It verifies the presence of necessary commands for the current operating system.
83//491//
84// Returns:492// Returns:
85// - true if a browser can be opened, false otherwise.493// - true if a browser can be opened, false otherwise.
86func IsAvailable() bool {494func IsAvailable() bool {
87 // First check if open-golang can work
88 testErr := open.Run("about:blank")
89 if testErr == nil {
90 return true
91 }
92
93 // Check platform-specific commands495 // Check platform-specific commands
94 switch runtime.GOOS {496 switch runtime.GOOS {
95 case "darwin":497 case "darwin":
internal/cmd/auth_manager.go+3−1
@@ -6,7 +6,7 @@ import ( · 5 unmodified lines
6 6
7// newAuthManager creates a new authentication manager instance with all supported7// newAuthManager creates a new authentication manager instance with all supported
8// authenticators and a file-based token store. It initializes authenticators for8// authenticators and a file-based token store. It initializes authenticators for
9// Gemini, Codex, Claude, and Qwen providers.9// Gemini, Codex, Claude, Qwen, IFlow, Antigravity, and GitHub Copilot providers.
10//10//
11// Returns:11// Returns:
12// - *sdkAuth.Manager: A configured authentication manager instance12// - *sdkAuth.Manager: A configured authentication manager instance
@@ -19,6 +19,8 @@ func newAuthManager() *sdkAuth.Manager { · 6 unmodified lines
19 sdkAuth.NewQwenAuthenticator(),19 sdkAuth.NewQwenAuthenticator(),
20 sdkAuth.NewIFlowAuthenticator(),20 sdkAuth.NewIFlowAuthenticator(),
21 sdkAuth.NewAntigravityAuthenticator(),21 sdkAuth.NewAntigravityAuthenticator(),
22 sdkAuth.NewKiroAuthenticator(),
23 sdkAuth.NewGitHubCopilotAuthenticator(),
22 )24 )
23 return manager25 return manager
24}26}
internal/cmd/github_copilot_login.go+44−0
@@ -0,0 +1,44 @@
1package cmd
2
3import (
4 "context"
5 "fmt"
6
7 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
8 sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
9 log "github.com/sirupsen/logrus"
10)
11
12// DoGitHubCopilotLogin triggers the OAuth device flow for GitHub Copilot and saves tokens.
13// It initiates the device flow authentication, displays the user code for the user to enter
14// at GitHub's verification URL, and waits for authorization before saving the tokens.
15//
16// Parameters:
17// - cfg: The application configuration containing proxy and auth directory settings
18// - options: Login options including browser behavior settings
19func DoGitHubCopilotLogin(cfg *config.Config, options *LoginOptions) {
20 if options == nil {
21 options = &LoginOptions{}
22 }
23
24 manager := newAuthManager()
25 authOpts := &sdkAuth.LoginOptions{
26 NoBrowser: options.NoBrowser,
27 Metadata: map[string]string{},
28 Prompt: options.Prompt,
29 }
30
31 record, savedPath, err := manager.Login(context.Background(), "github-copilot", cfg, authOpts)
32 if err != nil {
33 log.Errorf("GitHub Copilot authentication failed: %v", err)
34 return
35 }
36
37 if savedPath != "" {
38 fmt.Printf("Authentication saved to %s\n", savedPath)
39 }
40 if record != nil && record.Label != "" {
41 fmt.Printf("Authenticated as %s\n", record.Label)
42 }
43 fmt.Println("GitHub Copilot authentication successful!")
44}
internal/cmd/kiro_login.go+208−0
@@ -0,0 +1,208 @@
1package cmd
2
3import (
4 "context"
5 "fmt"
6
7 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
8 sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
9 log "github.com/sirupsen/logrus"
10)
11
12// DoKiroLogin triggers the Kiro authentication flow with Google OAuth.
13// This is the default login method (same as --kiro-google-login).
14//
15// Parameters:
16// - cfg: The application configuration
17// - options: Login options including Prompt field
18func DoKiroLogin(cfg *config.Config, options *LoginOptions) {
19 // Use Google login as default
20 DoKiroGoogleLogin(cfg, options)
21}
22
23// DoKiroGoogleLogin triggers Kiro authentication with Google OAuth.
24// This uses a custom protocol handler (kiro://) to receive the callback.
25//
26// Parameters:
27// - cfg: The application configuration
28// - options: Login options including prompts
29func DoKiroGoogleLogin(cfg *config.Config, options *LoginOptions) {
30 if options == nil {
31 options = &LoginOptions{}
32 }
33
34 // Note: Kiro defaults to incognito mode for multi-account support.
35 // Users can override with --no-incognito if they want to use existing browser sessions.
36
37 manager := newAuthManager()
38
39 // Use KiroAuthenticator with Google login
40 authenticator := sdkAuth.NewKiroAuthenticator()
41 record, err := authenticator.LoginWithGoogle(context.Background(), cfg, &sdkAuth.LoginOptions{
42 NoBrowser: options.NoBrowser,
43 Metadata: map[string]string{},
44 Prompt: options.Prompt,
45 })
46 if err != nil {
47 log.Errorf("Kiro Google authentication failed: %v", err)
48 fmt.Println("\nTroubleshooting:")
49 fmt.Println("1. Make sure the protocol handler is installed")
50 fmt.Println("2. Complete the Google login in the browser")
51 fmt.Println("3. If callback fails, try: --kiro-import (after logging in via Kiro IDE)")
52 return
53 }
54
55 // Save the auth record
56 savedPath, err := manager.SaveAuth(record, cfg)
57 if err != nil {
58 log.Errorf("Failed to save auth: %v", err)
59 return
60 }
61
62 if savedPath != "" {
63 fmt.Printf("Authentication saved to %s\n", savedPath)
64 }
65 if record != nil && record.Label != "" {
66 fmt.Printf("Authenticated as %s\n", record.Label)
67 }
68 fmt.Println("Kiro Google authentication successful!")
69}
70
71// DoKiroAWSLogin triggers Kiro authentication with AWS Builder ID.
72// This uses the device code flow for AWS SSO OIDC authentication.
73//
74// Parameters:
75// - cfg: The application configuration
76// - options: Login options including prompts
77func DoKiroAWSLogin(cfg *config.Config, options *LoginOptions) {
78 if options == nil {
79 options = &LoginOptions{}
80 }
81
82 // Note: Kiro defaults to incognito mode for multi-account support.
83 // Users can override with --no-incognito if they want to use existing browser sessions.
84
85 manager := newAuthManager()
86
87 // Use KiroAuthenticator with AWS Builder ID login (device code flow)
88 authenticator := sdkAuth.NewKiroAuthenticator()
89 record, err := authenticator.Login(context.Background(), cfg, &sdkAuth.LoginOptions{
90 NoBrowser: options.NoBrowser,
91 Metadata: map[string]string{},
92 Prompt: options.Prompt,
93 })
94 if err != nil {
95 log.Errorf("Kiro AWS authentication failed: %v", err)
96 fmt.Println("\nTroubleshooting:")
97 fmt.Println("1. Make sure you have an AWS Builder ID")
98 fmt.Println("2. Complete the authorization in the browser")
99 fmt.Println("3. If callback fails, try: --kiro-import (after logging in via Kiro IDE)")
100 return
101 }
102
103 // Save the auth record
104 savedPath, err := manager.SaveAuth(record, cfg)
105 if err != nil {
106 log.Errorf("Failed to save auth: %v", err)
107 return
108 }
109
110 if savedPath != "" {
111 fmt.Printf("Authentication saved to %s\n", savedPath)
112 }
113 if record != nil && record.Label != "" {
114 fmt.Printf("Authenticated as %s\n", record.Label)
115 }
116 fmt.Println("Kiro AWS authentication successful!")
117}
118
119// DoKiroAWSAuthCodeLogin triggers Kiro authentication with AWS Builder ID using authorization code flow.
120// This provides a better UX than device code flow as it uses automatic browser callback.
121//
122// Parameters:
123// - cfg: The application configuration
124// - options: Login options including prompts
125func DoKiroAWSAuthCodeLogin(cfg *config.Config, options *LoginOptions) {
126 if options == nil {
127 options = &LoginOptions{}
128 }
129
130 // Note: Kiro defaults to incognito mode for multi-account support.
131 // Users can override with --no-incognito if they want to use existing browser sessions.
132
133 manager := newAuthManager()
134
135 // Use KiroAuthenticator with AWS Builder ID login (authorization code flow)
136 authenticator := sdkAuth.NewKiroAuthenticator()
137 record, err := authenticator.LoginWithAuthCode(context.Background(), cfg, &sdkAuth.LoginOptions{
138 NoBrowser: options.NoBrowser,
139 Metadata: map[string]string{},
140 Prompt: options.Prompt,
141 })
142 if err != nil {
143 log.Errorf("Kiro AWS authentication (auth code) failed: %v", err)
144 fmt.Println("\nTroubleshooting:")
145 fmt.Println("1. Make sure you have an AWS Builder ID")
146 fmt.Println("2. Complete the authorization in the browser")
147 fmt.Println("3. If callback fails, try: --kiro-aws-login (device code flow)")
148 return
149 }
150
151 // Save the auth record
152 savedPath, err := manager.SaveAuth(record, cfg)
153 if err != nil {
154 log.Errorf("Failed to save auth: %v", err)
155 return
156 }
157
158 if savedPath != "" {
159 fmt.Printf("Authentication saved to %s\n", savedPath)
160 }
161 if record != nil && record.Label != "" {
162 fmt.Printf("Authenticated as %s\n", record.Label)
163 }
164 fmt.Println("Kiro AWS authentication successful!")
165}
166
167// DoKiroImport imports Kiro token from Kiro IDE's token file.
168// This is useful for users who have already logged in via Kiro IDE
169// and want to use the same credentials in CLI Proxy API.
170//
171// Parameters:
172// - cfg: The application configuration
173// - options: Login options (currently unused for import)
174func DoKiroImport(cfg *config.Config, options *LoginOptions) {
175 if options == nil {
176 options = &LoginOptions{}
177 }
178
179 manager := newAuthManager()
180
181 // Use ImportFromKiroIDE instead of Login
182 authenticator := sdkAuth.NewKiroAuthenticator()
183 record, err := authenticator.ImportFromKiroIDE(context.Background(), cfg)
184 if err != nil {
185 log.Errorf("Kiro token import failed: %v", err)
186 fmt.Println("\nMake sure you have logged in to Kiro IDE first:")
187 fmt.Println("1. Open Kiro IDE")
188 fmt.Println("2. Click 'Sign in with Google' (or GitHub)")
189 fmt.Println("3. Complete the login process")
190 fmt.Println("4. Run this command again")
191 return
192 }
193
194 // Save the imported auth record
195 savedPath, err := manager.SaveAuth(record, cfg)
196 if err != nil {
197 log.Errorf("Failed to save auth: %v", err)
198 return
199 }
200
201 if savedPath != "" {
202 fmt.Printf("Authentication saved to %s\n", savedPath)
203 }
204 if record != nil && record.Label != "" {
205 fmt.Printf("Imported as %s\n", record.Label)
206 }
207 fmt.Println("Kiro token import successful!")
208}
internal/config/config.go+62−0
@@ -72,6 +72,13 @@ type Config struct { · 71 unmodified lines
72 // GeminiKey defines Gemini API key configurations with optional routing overrides.72 // GeminiKey defines Gemini API key configurations with optional routing overrides.
73 GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`73 GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`
74 74
75 // KiroKey defines a list of Kiro (AWS CodeWhisperer) configurations.
76 KiroKey []KiroKey `yaml:"kiro" json:"kiro"`
77
78 // KiroPreferredEndpoint sets the global default preferred endpoint for all Kiro providers.
79 // Values: "ide" (default, CodeWhisperer) or "cli" (Amazon Q).
80 KiroPreferredEndpoint string `yaml:"kiro-preferred-endpoint" json:"kiro-preferred-endpoint"`
81
75 // Codex defines a list of Codex API key configurations as specified in the YAML configuration file.82 // Codex defines a list of Codex API key configurations as specified in the YAML configuration file.
76 CodexKey []CodexKey `yaml:"codex-api-key" json:"codex-api-key"`83 CodexKey []CodexKey `yaml:"codex-api-key" json:"codex-api-key"`
77 84
@@ -102,6 +109,11 @@ type Config struct { · 24 unmodified lines
102 // Payload defines default and override rules for provider payload parameters.109 // Payload defines default and override rules for provider payload parameters.
103 Payload PayloadConfig `yaml:"payload" json:"payload"`110 Payload PayloadConfig `yaml:"payload" json:"payload"`
104 111
112 // IncognitoBrowser enables opening OAuth URLs in incognito/private browsing mode.
113 // This is useful when you want to login with a different account without logging out
114 // from your current session. Default: false.
115 IncognitoBrowser bool `yaml:"incognito-browser" json:"incognito-browser"`
116
105 legacyMigrationPending bool `yaml:"-" json:"-"`117 legacyMigrationPending bool `yaml:"-" json:"-"`
106}118}
107 119
@@ -346,6 +358,35 @@ type GeminiModel struct { · 238 unmodified lines
346func (m GeminiModel) GetName() string { return m.Name }358func (m GeminiModel) GetName() string { return m.Name }
347func (m GeminiModel) GetAlias() string { return m.Alias }359func (m GeminiModel) GetAlias() string { return m.Alias }
348 360
361// KiroKey represents the configuration for Kiro (AWS CodeWhisperer) authentication.
362type KiroKey struct {
363 // TokenFile is the path to the Kiro token file (default: ~/.aws/sso/cache/kiro-auth-token.json)
364 TokenFile string `yaml:"token-file,omitempty" json:"token-file,omitempty"`
365
366 // AccessToken is the OAuth access token for direct configuration.
367 AccessToken string `yaml:"access-token,omitempty" json:"access-token,omitempty"`
368
369 // RefreshToken is the OAuth refresh token for token renewal.
370 RefreshToken string `yaml:"refresh-token,omitempty" json:"refresh-token,omitempty"`
371
372 // ProfileArn is the AWS CodeWhisperer profile ARN.
373 ProfileArn string `yaml:"profile-arn,omitempty" json:"profile-arn,omitempty"`
374
375 // Region is the AWS region (default: us-east-1).
376 Region string `yaml:"region,omitempty" json:"region,omitempty"`
377
378 // ProxyURL optionally overrides the global proxy for this configuration.
379 ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"`
380
381 // AgentTaskType sets the Kiro API task type. Known values: "vibe", "dev", "chat".
382 // Leave empty to let API use defaults. Different values may inject different system prompts.
383 AgentTaskType string `yaml:"agent-task-type,omitempty" json:"agent-task-type,omitempty"`
384
385 // PreferredEndpoint sets the preferred Kiro API endpoint/quota.
386 // Values: "codewhisperer" (default, IDE quota) or "amazonq" (CLI quota).
387 PreferredEndpoint string `yaml:"preferred-endpoint,omitempty" json:"preferred-endpoint,omitempty"`
388}
389
349// OpenAICompatibility represents the configuration for OpenAI API compatibility390// OpenAICompatibility represents the configuration for OpenAI API compatibility
350// with external providers, allowing model aliases to be routed through OpenAI API format.391// with external providers, allowing model aliases to be routed through OpenAI API format.
351type OpenAICompatibility struct {392type OpenAICompatibility struct {
@@ -432,6 +473,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) { · 80 unmodified lines
432 cfg.DisableCooling = false473 cfg.DisableCooling = false
433 cfg.AmpCode.RestrictManagementToLocalhost = false // Default to false: API key auth is sufficient474 cfg.AmpCode.RestrictManagementToLocalhost = false // Default to false: API key auth is sufficient
434 cfg.RemoteManagement.PanelGitHubRepository = DefaultPanelGitHubRepository475 cfg.RemoteManagement.PanelGitHubRepository = DefaultPanelGitHubRepository
476 cfg.IncognitoBrowser = false // Default to normal browser (AWS uses incognito by force)
435 if err = yaml.Unmarshal(data, &cfg); err != nil {477 if err = yaml.Unmarshal(data, &cfg); err != nil {
436 if optional {478 if optional {
437 // In cloud deploy mode, if YAML parsing fails, return empty config instead of error.479 // In cloud deploy mode, if YAML parsing fails, return empty config instead of error.
@@ -491,6 +533,9 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) { · 53 unmodified lines
491 // Sanitize Claude key headers533 // Sanitize Claude key headers
492 cfg.SanitizeClaudeKeys()534 cfg.SanitizeClaudeKeys()
493 535
536 // Sanitize Kiro keys: trim whitespace from credential fields
537 cfg.SanitizeKiroKeys()
538
494 // Sanitize OpenAI compatibility providers: drop entries without base-url539 // Sanitize OpenAI compatibility providers: drop entries without base-url
495 cfg.SanitizeOpenAICompatibility()540 cfg.SanitizeOpenAICompatibility()
496 541
@@ -617,6 +662,23 @@ func (cfg *Config) SanitizeClaudeKeys() { · 120 unmodified lines
617 }662 }
618}663}
619 664
665// SanitizeKiroKeys trims whitespace from Kiro credential fields.
666func (cfg *Config) SanitizeKiroKeys() {
667 if cfg == nil || len(cfg.KiroKey) == 0 {
668 return
669 }
670 for i := range cfg.KiroKey {
671 entry := &cfg.KiroKey[i]
672 entry.TokenFile = strings.TrimSpace(entry.TokenFile)
673 entry.AccessToken = strings.TrimSpace(entry.AccessToken)
674 entry.RefreshToken = strings.TrimSpace(entry.RefreshToken)
675 entry.ProfileArn = strings.TrimSpace(entry.ProfileArn)
676 entry.Region = strings.TrimSpace(entry.Region)
677 entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
678 entry.PreferredEndpoint = strings.TrimSpace(entry.PreferredEndpoint)
679 }
680}
681
620// SanitizeGeminiKeys deduplicates and normalizes Gemini credentials.682// SanitizeGeminiKeys deduplicates and normalizes Gemini credentials.
621func (cfg *Config) SanitizeGeminiKeys() {683func (cfg *Config) SanitizeGeminiKeys() {
622 if cfg == nil {684 if cfg == nil {
internal/constant/constant.go+3−0
@@ -24,4 +24,7 @@ const ( · 23 unmodified lines
24 24
25 // Antigravity represents the Antigravity response format identifier.25 // Antigravity represents the Antigravity response format identifier.
26 Antigravity = "antigravity"26 Antigravity = "antigravity"
27
28 // Kiro represents the AWS CodeWhisperer (Kiro) provider identifier.
29 Kiro = "kiro"
27)30)
internal/logging/global_logger.go+1−0
@@ -68,6 +68,7 @@ func (m *LogFormatter) Format(entry *log.Entry) ([]byte, error) { · 67 unmodified lines
68func SetupBaseLogger() {68func SetupBaseLogger() {
69 setupOnce.Do(func() {69 setupOnce.Do(func() {
70 log.SetOutput(os.Stdout)70 log.SetOutput(os.Stdout)
71 log.SetLevel(log.InfoLevel)
71 log.SetReportCaller(true)72 log.SetReportCaller(true)
72 log.SetFormatter(&LogFormatter{})73 log.SetFormatter(&LogFormatter{})
73 74
internal/registry/model_definitions.go+361−0
@@ -807,3 +807,364 @@ func LookupStaticModelInfo(modelID string) *ModelInfo { · 806 unmodified lines
807 }807 }
808 return nil808 return nil
809}809}
810
811// GetGitHubCopilotModels returns the available models for GitHub Copilot.
812// These models are available through the GitHub Copilot API at api.githubcopilot.com.
813func GetGitHubCopilotModels() []*ModelInfo {
814 now := int64(1732752000) // 2024-11-27
815 return []*ModelInfo{
816 {
817 ID: "gpt-4.1",
818 Object: "model",
819 Created: now,
820 OwnedBy: "github-copilot",
821 Type: "github-copilot",
822 DisplayName: "GPT-4.1",
823 Description: "OpenAI GPT-4.1 via GitHub Copilot",
824 ContextLength: 128000,
825 MaxCompletionTokens: 16384,
826 },
827 {
828 ID: "gpt-5",
829 Object: "model",
830 Created: now,
831 OwnedBy: "github-copilot",
832 Type: "github-copilot",
833 DisplayName: "GPT-5",
834 Description: "OpenAI GPT-5 via GitHub Copilot",
835 ContextLength: 200000,
836 MaxCompletionTokens: 32768,
837 },
838 {
839 ID: "gpt-5-mini",
840 Object: "model",
841 Created: now,
842 OwnedBy: "github-copilot",
843 Type: "github-copilot",
844 DisplayName: "GPT-5 Mini",
845 Description: "OpenAI GPT-5 Mini via GitHub Copilot",
846 ContextLength: 128000,
847 MaxCompletionTokens: 16384,
848 },
849 {
850 ID: "gpt-5-codex",
851 Object: "model",
852 Created: now,
853 OwnedBy: "github-copilot",
854 Type: "github-copilot",
855 DisplayName: "GPT-5 Codex",
856 Description: "OpenAI GPT-5 Codex via GitHub Copilot",
857 ContextLength: 200000,
858 MaxCompletionTokens: 32768,
859 },
860 {
861 ID: "gpt-5.1",
862 Object: "model",
863 Created: now,
864 OwnedBy: "github-copilot",
865 Type: "github-copilot",
866 DisplayName: "GPT-5.1",
867 Description: "OpenAI GPT-5.1 via GitHub Copilot",
868 ContextLength: 200000,
869 MaxCompletionTokens: 32768,
870 },
871 {
872 ID: "gpt-5.1-codex",
873 Object: "model",
874 Created: now,
875 OwnedBy: "github-copilot",
876 Type: "github-copilot",
877 DisplayName: "GPT-5.1 Codex",
878 Description: "OpenAI GPT-5.1 Codex via GitHub Copilot",
879 ContextLength: 200000,
880 MaxCompletionTokens: 32768,
881 },
882 {
883 ID: "gpt-5.1-codex-mini",
884 Object: "model",
885 Created: now,
886 OwnedBy: "github-copilot",
887 Type: "github-copilot",
888 DisplayName: "GPT-5.1 Codex Mini",
889 Description: "OpenAI GPT-5.1 Codex Mini via GitHub Copilot",
890 ContextLength: 128000,
891 MaxCompletionTokens: 16384,
892 },
893 {
894 ID: "gpt-5.2",
895 Object: "model",
896 Created: now,
897 OwnedBy: "github-copilot",
898 Type: "github-copilot",
899 DisplayName: "GPT-5.2",
900 Description: "OpenAI GPT-5.2 via GitHub Copilot",
901 ContextLength: 200000,
902 MaxCompletionTokens: 32768,
903 },
904 {
905 ID: "claude-haiku-4.5",
906 Object: "model",
907 Created: now,
908 OwnedBy: "github-copilot",
909 Type: "github-copilot",
910 DisplayName: "Claude Haiku 4.5",
911 Description: "Anthropic Claude Haiku 4.5 via GitHub Copilot",
912 ContextLength: 200000,
913 MaxCompletionTokens: 64000,
914 },
915 {
916 ID: "claude-opus-4.1",
917 Object: "model",
918 Created: now,
919 OwnedBy: "github-copilot",
920 Type: "github-copilot",
921 DisplayName: "Claude Opus 4.1",
922 Description: "Anthropic Claude Opus 4.1 via GitHub Copilot",
923 ContextLength: 200000,
924 MaxCompletionTokens: 32000,
925 },
926 {
927 ID: "claude-opus-4.5",
928 Object: "model",
929 Created: now,
930 OwnedBy: "github-copilot",
931 Type: "github-copilot",
932 DisplayName: "Claude Opus 4.5",
933 Description: "Anthropic Claude Opus 4.5 via GitHub Copilot",
934 ContextLength: 200000,
935 MaxCompletionTokens: 64000,
936 },
937 {
938 ID: "claude-sonnet-4",
939 Object: "model",
940 Created: now,
941 OwnedBy: "github-copilot",
942 Type: "github-copilot",
943 DisplayName: "Claude Sonnet 4",
944 Description: "Anthropic Claude Sonnet 4 via GitHub Copilot",
945 ContextLength: 200000,
946 MaxCompletionTokens: 64000,
947 },
948 {
949 ID: "claude-sonnet-4.5",
950 Object: "model",
951 Created: now,
952 OwnedBy: "github-copilot",
953 Type: "github-copilot",
954 DisplayName: "Claude Sonnet 4.5",
955 Description: "Anthropic Claude Sonnet 4.5 via GitHub Copilot",
956 ContextLength: 200000,
957 MaxCompletionTokens: 64000,
958 },
959 {
960 ID: "gemini-2.5-pro",
961 Object: "model",
962 Created: now,
963 OwnedBy: "github-copilot",
964 Type: "github-copilot",
965 DisplayName: "Gemini 2.5 Pro",
966 Description: "Google Gemini 2.5 Pro via GitHub Copilot",
967 ContextLength: 1048576,
968 MaxCompletionTokens: 65536,
969 },
970 {
971 ID: "gemini-3-pro",
972 Object: "model",
973 Created: now,
974 OwnedBy: "github-copilot",
975 Type: "github-copilot",
976 DisplayName: "Gemini 3 Pro",
977 Description: "Google Gemini 3 Pro via GitHub Copilot",
978 ContextLength: 1048576,
979 MaxCompletionTokens: 65536,
980 },
981 {
982 ID: "grok-code-fast-1",
983 Object: "model",
984 Created: now,
985 OwnedBy: "github-copilot",
986 Type: "github-copilot",
987 DisplayName: "Grok Code Fast 1",
988 Description: "xAI Grok Code Fast 1 via GitHub Copilot",
989 ContextLength: 128000,
990 MaxCompletionTokens: 16384,
991 },
992 {
993 ID: "raptor-mini",
994 Object: "model",
995 Created: now,
996 OwnedBy: "github-copilot",
997 Type: "github-copilot",
998 DisplayName: "Raptor Mini",
999 Description: "Raptor Mini via GitHub Copilot",
1000 ContextLength: 128000,
1001 MaxCompletionTokens: 16384,
1002 },
1003 }
1004}
1005
1006// GetKiroModels returns the Kiro (AWS CodeWhisperer) model definitions
1007func GetKiroModels() []*ModelInfo {
1008 return []*ModelInfo{
1009 // --- Base Models ---
1010 {
1011 ID: "kiro-claude-opus-4-5",
1012 Object: "model",
1013 Created: 1732752000,
1014 OwnedBy: "aws",
1015 Type: "kiro",
1016 DisplayName: "Kiro Claude Opus 4.5",
1017 Description: "Claude Opus 4.5 via Kiro (2.2x credit)",
1018 ContextLength: 200000,
1019 MaxCompletionTokens: 64000,
1020 Thinking: &ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true},
1021 },
1022 {
1023 ID: "kiro-claude-sonnet-4-5",
1024 Object: "model",
1025 Created: 1732752000,
1026 OwnedBy: "aws",
1027 Type: "kiro",
1028 DisplayName: "Kiro Claude Sonnet 4.5",
1029 Description: "Claude Sonnet 4.5 via Kiro (1.3x credit)",
1030 ContextLength: 200000,
1031 MaxCompletionTokens: 64000,
1032 Thinking: &ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true},
1033 },
1034 {
1035 ID: "kiro-claude-sonnet-4",
1036 Object: "model",
1037 Created: 1732752000,
1038 OwnedBy: "aws",
1039 Type: "kiro",
1040 DisplayName: "Kiro Claude Sonnet 4",
1041 Description: "Claude Sonnet 4 via Kiro (1.3x credit)",
1042 ContextLength: 200000,
1043 MaxCompletionTokens: 64000,
1044 Thinking: &ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true},
1045 },
1046 {
1047 ID: "kiro-claude-haiku-4-5",
1048 Object: "model",
1049 Created: 1732752000,
1050 OwnedBy: "aws",
1051 Type: "kiro",
1052 DisplayName: "Kiro Claude Haiku 4.5",
1053 Description: "Claude Haiku 4.5 via Kiro (0.4x credit)",
1054 ContextLength: 200000,
1055 MaxCompletionTokens: 64000,
1056 Thinking: &ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true},
1057 },
1058 // --- Agentic Variants (Optimized for coding agents with chunked writes) ---
1059 {
1060 ID: "kiro-claude-opus-4-5-agentic",
1061 Object: "model",
1062 Created: 1732752000,
1063 OwnedBy: "aws",
1064 Type: "kiro",
1065 DisplayName: "Kiro Claude Opus 4.5 (Agentic)",
1066 Description: "Claude Opus 4.5 optimized for coding agents (chunked writes)",
1067 ContextLength: 200000,
1068 MaxCompletionTokens: 64000,
1069 Thinking: &ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true},
1070 },
1071 {
1072 ID: "kiro-claude-sonnet-4-5-agentic",
1073 Object: "model",
1074 Created: 1732752000,
1075 OwnedBy: "aws",
1076 Type: "kiro",
1077 DisplayName: "Kiro Claude Sonnet 4.5 (Agentic)",
1078 Description: "Claude Sonnet 4.5 optimized for coding agents (chunked writes)",
1079 ContextLength: 200000,
1080 MaxCompletionTokens: 64000,
1081 Thinking: &ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true},
1082 },
1083 {
1084 ID: "kiro-claude-sonnet-4-agentic",
1085 Object: "model",
1086 Created: 1732752000,
1087 OwnedBy: "aws",
1088 Type: "kiro",
1089 DisplayName: "Kiro Claude Sonnet 4 (Agentic)",
1090 Description: "Claude Sonnet 4 optimized for coding agents (chunked writes)",
1091 ContextLength: 200000,
1092 MaxCompletionTokens: 64000,
1093 Thinking: &ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true},
1094 },
1095 {
1096 ID: "kiro-claude-haiku-4-5-agentic",
1097 Object: "model",
1098 Created: 1732752000,
1099 OwnedBy: "aws",
1100 Type: "kiro",
1101 DisplayName: "Kiro Claude Haiku 4.5 (Agentic)",
1102 Description: "Claude Haiku 4.5 optimized for coding agents (chunked writes)",
1103 ContextLength: 200000,
1104 MaxCompletionTokens: 64000,
1105 Thinking: &ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true},
1106 },
1107 }
1108}
1109
1110// GetAmazonQModels returns the Amazon Q (AWS CodeWhisperer) model definitions.
1111// These models use the same API as Kiro and share the same executor.
1112func GetAmazonQModels() []*ModelInfo {
1113 return []*ModelInfo{
1114 {
1115 ID: "amazonq-auto",
1116 Object: "model",
1117 Created: 1732752000,
1118 OwnedBy: "aws",
1119 Type: "kiro", // Uses Kiro executor - same API
1120 DisplayName: "Amazon Q Auto",
1121 Description: "Automatic model selection by Amazon Q",
1122 ContextLength: 200000,
1123 MaxCompletionTokens: 64000,
1124 },
1125 {
1126 ID: "amazonq-claude-opus-4.5",
1127 Object: "model",
1128 Created: 1732752000,
1129 OwnedBy: "aws",
1130 Type: "kiro",
1131 DisplayName: "Amazon Q Claude Opus 4.5",
1132 Description: "Claude Opus 4.5 via Amazon Q (2.2x credit)",
1133 ContextLength: 200000,
1134 MaxCompletionTokens: 64000,
1135 },
1136 {
1137 ID: "amazonq-claude-sonnet-4.5",
1138 Object: "model",
1139 Created: 1732752000,
1140 OwnedBy: "aws",
1141 Type: "kiro",
1142 DisplayName: "Amazon Q Claude Sonnet 4.5",
1143 Description: "Claude Sonnet 4.5 via Amazon Q (1.3x credit)",
1144 ContextLength: 200000,
1145 MaxCompletionTokens: 64000,
1146 },
1147 {
1148 ID: "amazonq-claude-sonnet-4",
1149 Object: "model",
1150 Created: 1732752000,
1151 OwnedBy: "aws",
1152 Type: "kiro",
1153 DisplayName: "Amazon Q Claude Sonnet 4",
1154 Description: "Claude Sonnet 4 via Amazon Q (1.3x credit)",
1155 ContextLength: 200000,
1156 MaxCompletionTokens: 64000,
1157 },
1158 {
1159 ID: "amazonq-claude-haiku-4.5",
1160 Object: "model",
1161 Created: 1732752000,
1162 OwnedBy: "aws",
1163 Type: "kiro",
1164 DisplayName: "Amazon Q Claude Haiku 4.5",
1165 Description: "Claude Haiku 4.5 via Amazon Q (0.4x credit)",
1166 ContextLength: 200000,
1167 MaxCompletionTokens: 64000,
1168 },
1169 }
1170}
internal/registry/model_registry.go+15−1
@@ -766,7 +766,8 @@ func (r *ModelRegistry) convertModelToMap(model *ModelInfo, handlerType string) · 765 unmodified lines
766 }766 }
767 return result767 return result
768 768
769 case "claude":769 case "claude", "kiro", "antigravity":
770 // Claude, Kiro, and Antigravity all use Claude-compatible format for Claude Code client
770 result := map[string]any{771 result := map[string]any{
771 "id": model.ID,772 "id": model.ID,
772 "object": "model",773 "object": "model",
@@ -781,6 +782,19 @@ func (r *ModelRegistry) convertModelToMap(model *ModelInfo, handlerType string) · 8 unmodified lines
781 if model.DisplayName != "" {782 if model.DisplayName != "" {
782 result["display_name"] = model.DisplayName783 result["display_name"] = model.DisplayName
783 }784 }
785 // Add thinking support for Claude Code client
786 // Claude Code checks for "thinking" field (simple boolean) to enable tab toggle
787 // Also add "extended_thinking" for detailed budget info
788 if model.Thinking != nil {
789 result["thinking"] = true
790 result["extended_thinking"] = map[string]any{
791 "supported": true,
792 "min": model.Thinking.Min,
793 "max": model.Thinking.Max,
794 "zero_allowed": model.Thinking.ZeroAllowed,
795 "dynamic_allowed": model.Thinking.DynamicAllowed,
796 }
797 }
784 return result798 return result
785 799
786 case "gemini":800 case "gemini":
internal/runtime/executor/cache_helpers.go+30−2
@@ -1,10 +1,38 @@
1package executor1package executor
2 2
3import "time"3import (
4 "sync"
5 "time"
6)
4 7
5type codexCache struct {8type codexCache struct {
6 ID string9 ID string
7 Expire time.Time10 Expire time.Time
8}11}
9 12
10var codexCacheMap = map[string]codexCache{}13var (
14 codexCacheMap = map[string]codexCache{}
15 codexCacheMutex sync.RWMutex
16)
17
18// getCodexCache safely retrieves a cache entry
19func getCodexCache(key string) (codexCache, bool) {
20 codexCacheMutex.RLock()
21 defer codexCacheMutex.RUnlock()
22 cache, ok := codexCacheMap[key]
23 return cache, ok
24}
25
26// setCodexCache safely sets a cache entry
27func setCodexCache(key string, cache codexCache) {
28 codexCacheMutex.Lock()
29 defer codexCacheMutex.Unlock()
30 codexCacheMap[key] = cache
31}
32
33// deleteCodexCache safely deletes a cache entry
34func deleteCodexCache(key string) {
35 codexCacheMutex.Lock()
36 defer codexCacheMutex.Unlock()
37 delete(codexCacheMap, key)
38}
internal/runtime/executor/codex_executor.go+2−2
@@ -449,12 +449,12 @@ func (e *CodexExecutor) cacheHelper(ctx context.Context, from sdktranslator.Form · 448 unmodified lines
449 if userIDResult.Exists() {449 if userIDResult.Exists() {
450 var hasKey bool450 var hasKey bool
451 key := fmt.Sprintf("%s-%s", req.Model, userIDResult.String())451 key := fmt.Sprintf("%s-%s", req.Model, userIDResult.String())
452 if cache, hasKey = codexCacheMap[key]; !hasKey || cache.Expire.Before(time.Now()) {452 if cache, hasKey = getCodexCache(key); !hasKey || cache.Expire.Before(time.Now()) {
453 cache = codexCache{453 cache = codexCache{
454 ID: uuid.New().String(),454 ID: uuid.New().String(),
455 Expire: time.Now().Add(1 * time.Hour),455 Expire: time.Now().Add(1 * time.Hour),
456 }456 }
457 codexCacheMap[key] = cache457 setCodexCache(key, cache)
458 }458 }
459 }459 }
460 } else if from == "openai-response" {460 } else if from == "openai-response" {
internal/runtime/executor/github_copilot_executor.go+361−0
@@ -0,0 +1,361 @@
1package executor
2
3import (
4 "bufio"
5 "bytes"
6 "context"
7 "fmt"
8 "io"
9 "net/http"
10 "sync"
11 "time"
12
13 "github.com/google/uuid"
14 copilotauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/copilot"
15 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
16 cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
17 cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
18 sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
19 log "github.com/sirupsen/logrus"
20 "github.com/tidwall/sjson"
21)
22
23const (
24 githubCopilotBaseURL = "https://api.githubcopilot.com"
25 githubCopilotChatPath = "/chat/completions"
26 githubCopilotAuthType = "github-copilot"
27 githubCopilotTokenCacheTTL = 25 * time.Minute
28 // tokenExpiryBuffer is the time before expiry when we should refresh the token.
29 tokenExpiryBuffer = 5 * time.Minute
30 // maxScannerBufferSize is the maximum buffer size for SSE scanning (20MB).
31 maxScannerBufferSize = 20_971_520
32
33 // Copilot API header values.
34 copilotUserAgent = "GithubCopilot/1.0"
35 copilotEditorVersion = "vscode/1.100.0"
36 copilotPluginVersion = "copilot/1.300.0"
37 copilotIntegrationID = "vscode-chat"
38 copilotOpenAIIntent = "conversation-panel"
39)
40
41// GitHubCopilotExecutor handles requests to the GitHub Copilot API.
42type GitHubCopilotExecutor struct {
43 cfg *config.Config
44 mu sync.RWMutex
45 cache map[string]*cachedAPIToken
46}
47
48// cachedAPIToken stores a cached Copilot API token with its expiry.
49type cachedAPIToken struct {
50 token string
51 expiresAt time.Time
52}
53
54// NewGitHubCopilotExecutor constructs a new executor instance.
55func NewGitHubCopilotExecutor(cfg *config.Config) *GitHubCopilotExecutor {
56 return &GitHubCopilotExecutor{
57 cfg: cfg,
58 cache: make(map[string]*cachedAPIToken),
59 }
60}
61
62// Identifier implements ProviderExecutor.
63func (e *GitHubCopilotExecutor) Identifier() string { return githubCopilotAuthType }
64
65// PrepareRequest implements ProviderExecutor.
66func (e *GitHubCopilotExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error {
67 return nil
68}
69
70// Execute handles non-streaming requests to GitHub Copilot.
71func (e *GitHubCopilotExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
72 apiToken, errToken := e.ensureAPIToken(ctx, auth)
73 if errToken != nil {
74 return resp, errToken
75 }
76
77 reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
78 defer reporter.trackFailure(ctx, &err)
79
80 from := opts.SourceFormat
81 to := sdktranslator.FromString("openai")
82 body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
83 body = e.normalizeModel(req.Model, body)
84 body = applyPayloadConfig(e.cfg, req.Model, body)
85 body, _ = sjson.SetBytes(body, "stream", false)
86
87 url := githubCopilotBaseURL + githubCopilotChatPath
88 httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
89 if err != nil {
90 return resp, err
91 }
92 e.applyHeaders(httpReq, apiToken)
93
94 var authID, authLabel, authType, authValue string
95 if auth != nil {
96 authID = auth.ID
97 authLabel = auth.Label
98 authType, authValue = auth.AccountInfo()
99 }
100 recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
101 URL: url,
102 Method: http.MethodPost,
103 Headers: httpReq.Header.Clone(),
104 Body: body,
105 Provider: e.Identifier(),
106 AuthID: authID,
107 AuthLabel: authLabel,
108 AuthType: authType,
109 AuthValue: authValue,
110 })
111
112 httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
113 httpResp, err := httpClient.Do(httpReq)
114 if err != nil {
115 recordAPIResponseError(ctx, e.cfg, err)
116 return resp, err
117 }
118 defer func() {
119 if errClose := httpResp.Body.Close(); errClose != nil {
120 log.Errorf("github-copilot executor: close response body error: %v", errClose)
121 }
122 }()
123
124 recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
125
126 if !isHTTPSuccess(httpResp.StatusCode) {
127 data, _ := io.ReadAll(httpResp.Body)
128 appendAPIResponseChunk(ctx, e.cfg, data)
129 log.Debugf("github-copilot executor: upstream error status: %d, body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
130 err = statusErr{code: httpResp.StatusCode, msg: string(data)}
131 return resp, err
132 }
133
134 data, err := io.ReadAll(httpResp.Body)
135 if err != nil {
136 recordAPIResponseError(ctx, e.cfg, err)
137 return resp, err
138 }
139 appendAPIResponseChunk(ctx, e.cfg, data)
140
141 detail := parseOpenAIUsage(data)
142 if detail.TotalTokens > 0 {
143 reporter.publish(ctx, detail)
144 }
145
146 var param any
147 converted := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, data, &param)
148 resp = cliproxyexecutor.Response{Payload: []byte(converted)}
149 reporter.ensurePublished(ctx)
150 return resp, nil
151}
152
153// ExecuteStream handles streaming requests to GitHub Copilot.
154func (e *GitHubCopilotExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
155 apiToken, errToken := e.ensureAPIToken(ctx, auth)
156 if errToken != nil {
157 return nil, errToken
158 }
159
160 reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
161 defer reporter.trackFailure(ctx, &err)
162
163 from := opts.SourceFormat
164 to := sdktranslator.FromString("openai")
165 body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
166 body = e.normalizeModel(req.Model, body)
167 body = applyPayloadConfig(e.cfg, req.Model, body)
168 body, _ = sjson.SetBytes(body, "stream", true)
169 // Enable stream options for usage stats in stream
170 body, _ = sjson.SetBytes(body, "stream_options.include_usage", true)
171
172 url := githubCopilotBaseURL + githubCopilotChatPath
173 httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
174 if err != nil {
175 return nil, err
176 }
177 e.applyHeaders(httpReq, apiToken)
178
179 var authID, authLabel, authType, authValue string
180 if auth != nil {
181 authID = auth.ID
182 authLabel = auth.Label
183 authType, authValue = auth.AccountInfo()
184 }
185 recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
186 URL: url,
187 Method: http.MethodPost,
188 Headers: httpReq.Header.Clone(),
189 Body: body,
190 Provider: e.Identifier(),
191 AuthID: authID,
192 AuthLabel: authLabel,
193 AuthType: authType,
194 AuthValue: authValue,
195 })
196
197 httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
198 httpResp, err := httpClient.Do(httpReq)
199 if err != nil {
200 recordAPIResponseError(ctx, e.cfg, err)
201 return nil, err
202 }
203
204 recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
205
206 if !isHTTPSuccess(httpResp.StatusCode) {
207 data, readErr := io.ReadAll(httpResp.Body)
208 if errClose := httpResp.Body.Close(); errClose != nil {
209 log.Errorf("github-copilot executor: close response body error: %v", errClose)
210 }
211 if readErr != nil {
212 recordAPIResponseError(ctx, e.cfg, readErr)
213 return nil, readErr
214 }
215 appendAPIResponseChunk(ctx, e.cfg, data)
216 log.Debugf("github-copilot executor: upstream error status: %d, body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
217 err = statusErr{code: httpResp.StatusCode, msg: string(data)}
218 return nil, err
219 }
220
221 out := make(chan cliproxyexecutor.StreamChunk)
222 stream = out
223
224 go func() {
225 defer close(out)
226 defer func() {
227 if errClose := httpResp.Body.Close(); errClose != nil {
228 log.Errorf("github-copilot executor: close response body error: %v", errClose)
229 }
230 }()
231
232 scanner := bufio.NewScanner(httpResp.Body)
233 scanner.Buffer(nil, maxScannerBufferSize)
234 var param any
235
236 for scanner.Scan() {
237 line := scanner.Bytes()
238 appendAPIResponseChunk(ctx, e.cfg, line)
239
240 // Parse SSE data
241 if bytes.HasPrefix(line, dataTag) {
242 data := bytes.TrimSpace(line[5:])
243 if bytes.Equal(data, []byte("[DONE]")) {
244 continue
245 }
246 if detail, ok := parseOpenAIStreamUsage(line); ok {
247 reporter.publish(ctx, detail)
248 }
249 }
250
251 chunks := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, bytes.Clone(line), &param)
252 for i := range chunks {
253 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunks[i])}
254 }
255 }
256
257 if errScan := scanner.Err(); errScan != nil {
258 recordAPIResponseError(ctx, e.cfg, errScan)
259 reporter.publishFailure(ctx)
260 out <- cliproxyexecutor.StreamChunk{Err: errScan}
261 } else {
262 reporter.ensurePublished(ctx)
263 }
264 }()
265
266 return stream, nil
267}
268
269// CountTokens is not supported for GitHub Copilot.
270func (e *GitHubCopilotExecutor) CountTokens(_ context.Context, _ *cliproxyauth.Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
271 return cliproxyexecutor.Response{}, statusErr{code: http.StatusNotImplemented, msg: "count tokens not supported for github-copilot"}
272}
273
274// Refresh validates the GitHub token is still working.
275// GitHub OAuth tokens don't expire traditionally, so we just validate.
276func (e *GitHubCopilotExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
277 if auth == nil {
278 return nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"}
279 }
280
281 // Get the GitHub access token
282 accessToken := metaStringValue(auth.Metadata, "access_token")
283 if accessToken == "" {
284 return auth, nil
285 }
286
287 // Validate the token can still get a Copilot API token
288 copilotAuth := copilotauth.NewCopilotAuth(e.cfg)
289 _, err := copilotAuth.GetCopilotAPIToken(ctx, accessToken)
290 if err != nil {
291 return nil, statusErr{code: http.StatusUnauthorized, msg: fmt.Sprintf("github-copilot token validation failed: %v", err)}
292 }
293
294 return auth, nil
295}
296
297// ensureAPIToken gets or refreshes the Copilot API token.
298func (e *GitHubCopilotExecutor) ensureAPIToken(ctx context.Context, auth *cliproxyauth.Auth) (string, error) {
299 if auth == nil {
300 return "", statusErr{code: http.StatusUnauthorized, msg: "missing auth"}
301 }
302
303 // Get the GitHub access token
304 accessToken := metaStringValue(auth.Metadata, "access_token")
305 if accessToken == "" {
306 return "", statusErr{code: http.StatusUnauthorized, msg: "missing github access token"}
307 }
308
309 // Check for cached API token using thread-safe access
310 e.mu.RLock()
311 if cached, ok := e.cache[accessToken]; ok && cached.expiresAt.After(time.Now().Add(tokenExpiryBuffer)) {
312 e.mu.RUnlock()
313 return cached.token, nil
314 }
315 e.mu.RUnlock()
316
317 // Get a new Copilot API token
318 copilotAuth := copilotauth.NewCopilotAuth(e.cfg)
319 apiToken, err := copilotAuth.GetCopilotAPIToken(ctx, accessToken)
320 if err != nil {
321 return "", statusErr{code: http.StatusUnauthorized, msg: fmt.Sprintf("failed to get copilot api token: %v", err)}
322 }
323
324 // Cache the token with thread-safe access
325 expiresAt := time.Now().Add(githubCopilotTokenCacheTTL)
326 if apiToken.ExpiresAt > 0 {
327 expiresAt = time.Unix(apiToken.ExpiresAt, 0)
328 }
329 e.mu.Lock()
330 e.cache[accessToken] = &cachedAPIToken{
331 token: apiToken.Token,
332 expiresAt: expiresAt,
333 }
334 e.mu.Unlock()
335
336 return apiToken.Token, nil
337}
338
339// applyHeaders sets the required headers for GitHub Copilot API requests.
340func (e *GitHubCopilotExecutor) applyHeaders(r *http.Request, apiToken string) {
341 r.Header.Set("Content-Type", "application/json")
342 r.Header.Set("Authorization", "Bearer "+apiToken)
343 r.Header.Set("Accept", "application/json")
344 r.Header.Set("User-Agent", copilotUserAgent)
345 r.Header.Set("Editor-Version", copilotEditorVersion)
346 r.Header.Set("Editor-Plugin-Version", copilotPluginVersion)
347 r.Header.Set("Openai-Intent", copilotOpenAIIntent)
348 r.Header.Set("Copilot-Integration-Id", copilotIntegrationID)
349 r.Header.Set("X-Request-Id", uuid.NewString())
350}
351
352// normalizeModel is a no-op as GitHub Copilot accepts model names directly.
353// Model mapping should be done at the registry level if needed.
354func (e *GitHubCopilotExecutor) normalizeModel(_ string, body []byte) []byte {
355 return body
356}
357
358// isHTTPSuccess checks if the status code indicates success (2xx).
359func isHTTPSuccess(statusCode int) bool {
360 return statusCode >= 200 && statusCode < 300
361}
internal/runtime/executor/kiro_executor.go+3302−0
@@ -0,0 +1,3302 @@
1package executor
2
3import (
4 "bufio"
5 "bytes"
6 "context"
7 "encoding/base64"
8 "encoding/binary"
9 "encoding/json"
10 "fmt"
11 "io"
12 "net/http"
13 "os"
14 "path/filepath"
15 "strings"
16 "sync"
17 "time"
18
19 "github.com/google/uuid"
20 kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"
21 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
22 kiroclaude "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/claude"
23 kirocommon "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/common"
24 kiroopenai "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/openai"
25 "github.com/router-for-me/CLIProxyAPI/v6/internal/util"
26 cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
27 cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
28 "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
29 sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
30 log "github.com/sirupsen/logrus"
31
32)
33
34const (
35 // Kiro API common constants
36 kiroContentType = "application/x-amz-json-1.0"
37 kiroAcceptStream = "*/*"
38
39 // Event Stream frame size constants for boundary protection
40 // AWS Event Stream binary format: prelude (12 bytes) + headers + payload + message_crc (4 bytes)
41 // Prelude consists of: total_length (4) + headers_length (4) + prelude_crc (4)
42 minEventStreamFrameSize = 16 // Minimum: 4(total_len) + 4(headers_len) + 4(prelude_crc) + 4(message_crc)
43 maxEventStreamMsgSize = 10 << 20 // Maximum message length: 10MB
44
45 // Event Stream error type constants
46 ErrStreamFatal = "fatal" // Connection/authentication errors, not recoverable
47 ErrStreamMalformed = "malformed" // Format errors, data cannot be parsed
48 // kiroUserAgent matches amq2api format for User-Agent header (Amazon Q CLI style)
49 kiroUserAgent = "aws-sdk-rust/1.3.9 os/macos lang/rust/1.87.0"
50 // kiroFullUserAgent is the complete x-amz-user-agent header matching amq2api (Amazon Q CLI style)
51 kiroFullUserAgent = "aws-sdk-rust/1.3.9 ua/2.1 api/ssooidc/1.88.0 os/macos lang/rust/1.87.0 m/E app/AmazonQ-For-CLI"
52
53 // Kiro IDE style headers (from kiro2api - for IDC auth)
54 kiroIDEUserAgent = "aws-sdk-js/1.0.18 ua/2.1 os/darwin#25.0.0 lang/js md/nodejs#20.16.0 api/codewhispererstreaming#1.0.18 m/E KiroIDE-0.2.13-66c23a8c5d15afabec89ef9954ef52a119f10d369df04d548fc6c1eac694b0d1"
55 kiroIDEAmzUserAgent = "aws-sdk-js/1.0.18 KiroIDE-0.2.13-66c23a8c5d15afabec89ef9954ef52a119f10d369df04d548fc6c1eac694b0d1"
56 kiroIDEAgentModeSpec = "spec"
57)
58
59// Real-time usage estimation configuration
60// These control how often usage updates are sent during streaming
61var (
62 usageUpdateCharThreshold = 5000 // Send usage update every 5000 characters
63 usageUpdateTimeInterval = 15 * time.Second // Or every 15 seconds, whichever comes first
64)
65
66// kiroEndpointConfig bundles endpoint URL with its compatible Origin and AmzTarget values.
67// This solves the "triple mismatch" problem where different endpoints require matching
68// Origin and X-Amz-Target header values.
69//
70// Based on reference implementations:
71// - amq2api-main: Uses Amazon Q endpoint with CLI origin and AmazonQDeveloperStreamingService target
72// - AIClient-2-API: Uses CodeWhisperer endpoint with AI_EDITOR origin and AmazonCodeWhispererStreamingService target
73type kiroEndpointConfig struct {
74 URL string // Endpoint URL
75 Origin string // Request Origin: "CLI" for Amazon Q quota, "AI_EDITOR" for Kiro IDE quota
76 AmzTarget string // X-Amz-Target header value
77 Name string // Endpoint name for logging
78}
79
80// kiroEndpointConfigs defines the available Kiro API endpoints with their compatible configurations.
81// The order determines fallback priority: primary endpoint first, then fallbacks.
82//
83// CRITICAL: Each endpoint MUST use its compatible Origin and AmzTarget values:
84// - CodeWhisperer endpoint (codewhisperer.us-east-1.amazonaws.com): Uses AI_EDITOR origin and AmazonCodeWhispererStreamingService target
85// - Amazon Q endpoint (q.us-east-1.amazonaws.com): Uses CLI origin and AmazonQDeveloperStreamingService target
86//
87// Mismatched combinations will result in 403 Forbidden errors.
88//
89// NOTE: CodeWhisperer is set as the default endpoint because:
90// 1. Most tokens come from Kiro IDE / VSCode extensions (AWS Builder ID auth)
91// 2. These tokens use AI_EDITOR origin which is only compatible with CodeWhisperer endpoint
92// 3. Amazon Q endpoint requires CLI origin which is for Amazon Q CLI tokens
93// This matches the AIClient-2-API-main project's configuration.
94var kiroEndpointConfigs = []kiroEndpointConfig{
95 {
96 URL: "https://codewhisperer.us-east-1.amazonaws.com/generateAssistantResponse",
97 Origin: "AI_EDITOR",
98 AmzTarget: "AmazonCodeWhispererStreamingService.GenerateAssistantResponse",
99 Name: "CodeWhisperer",
100 },
101 {
102 URL: "https://q.us-east-1.amazonaws.com/",
103 Origin: "CLI",
104 AmzTarget: "AmazonQDeveloperStreamingService.SendMessage",
105 Name: "AmazonQ",
106 },
107}
108
109// getKiroEndpointConfigs returns the list of Kiro API endpoint configurations to try in order.
110// Supports reordering based on "preferred_endpoint" in auth metadata/attributes.
111// For IDC auth method, automatically uses CodeWhisperer endpoint with CLI origin.
112func getKiroEndpointConfigs(auth *cliproxyauth.Auth) []kiroEndpointConfig {
113 if auth == nil {
114 return kiroEndpointConfigs
115 }
116
117 // For IDC auth, use CodeWhisperer endpoint with AI_EDITOR origin (same as Social auth)
118 // Based on kiro2api analysis: IDC tokens work with CodeWhisperer endpoint using Bearer auth
119 // The difference is only in how tokens are refreshed (OIDC with clientId/clientSecret for IDC)
120 // NOT in how API calls are made - both Social and IDC use the same endpoint/origin
121 if auth.Metadata != nil {
122 authMethod, _ := auth.Metadata["auth_method"].(string)
123 if authMethod == "idc" {
124 log.Debugf("kiro: IDC auth, using CodeWhisperer endpoint")
125 return kiroEndpointConfigs
126 }
127 }
128
129 // Check for preference
130 var preference string
131 if auth.Metadata != nil {
132 if p, ok := auth.Metadata["preferred_endpoint"].(string); ok {
133 preference = p
134 }
135 }
136 // Check attributes as fallback (e.g. from HTTP headers)
137 if preference == "" && auth.Attributes != nil {
138 preference = auth.Attributes["preferred_endpoint"]
139 }
140
141 if preference == "" {
142 return kiroEndpointConfigs
143 }
144
145 preference = strings.ToLower(strings.TrimSpace(preference))
146
147 // Create new slice to avoid modifying global state
148 var sorted []kiroEndpointConfig
149 var remaining []kiroEndpointConfig
150
151 for _, cfg := range kiroEndpointConfigs {
152 name := strings.ToLower(cfg.Name)
153 // Check for matches
154 // CodeWhisperer aliases: codewhisperer, ide
155 // AmazonQ aliases: amazonq, q, cli
156 isMatch := false
157 if (preference == "codewhisperer" || preference == "ide") && name == "codewhisperer" {
158 isMatch = true
159 } else if (preference == "amazonq" || preference == "q" || preference == "cli") && name == "amazonq" {
160 isMatch = true
161 }
162
163 if isMatch {
164 sorted = append(sorted, cfg)
165 } else {
166 remaining = append(remaining, cfg)
167 }
168 }
169
170 // If preference didn't match anything, return default
171 if len(sorted) == 0 {
172 return kiroEndpointConfigs
173 }
174
175 // Combine: preferred first, then others
176 return append(sorted, remaining...)
177}
178
179// KiroExecutor handles requests to AWS CodeWhisperer (Kiro) API.
180type KiroExecutor struct {
181 cfg *config.Config
182 refreshMu sync.Mutex // Serializes token refresh operations to prevent race conditions
183}
184
185// isIDCAuth checks if the auth uses IDC (Identity Center) authentication method.
186func isIDCAuth(auth *cliproxyauth.Auth) bool {
187 if auth == nil || auth.Metadata == nil {
188 return false
189 }
190 authMethod, _ := auth.Metadata["auth_method"].(string)
191 return authMethod == "idc"
192}
193
194// buildKiroPayloadForFormat builds the Kiro API payload based on the source format.
195// This is critical because OpenAI and Claude formats have different tool structures:
196// - OpenAI: tools[].function.name, tools[].function.description
197// - Claude: tools[].name, tools[].description
198// headers parameter allows checking Anthropic-Beta header for thinking mode detection.
199// Returns the serialized JSON payload and a boolean indicating whether thinking mode was injected.
200func buildKiroPayloadForFormat(body []byte, modelID, profileArn, origin string, isAgentic, isChatOnly bool, sourceFormat sdktranslator.Format, headers http.Header) ([]byte, bool) {
201 switch sourceFormat.String() {
202 case "openai":
203 log.Debugf("kiro: using OpenAI payload builder for source format: %s", sourceFormat.String())
204 return kiroopenai.BuildKiroPayloadFromOpenAI(body, modelID, profileArn, origin, isAgentic, isChatOnly, headers, nil)
205 default:
206 // Default to Claude format (also handles "claude", "kiro", etc.)
207 log.Debugf("kiro: using Claude payload builder for source format: %s", sourceFormat.String())
208 return kiroclaude.BuildKiroPayload(body, modelID, profileArn, origin, isAgentic, isChatOnly, headers, nil)
209 }
210}
211
212// NewKiroExecutor creates a new Kiro executor instance.
213func NewKiroExecutor(cfg *config.Config) *KiroExecutor {
214 return &KiroExecutor{cfg: cfg}
215}
216
217// Identifier returns the unique identifier for this executor.
218func (e *KiroExecutor) Identifier() string { return "kiro" }
219
220// PrepareRequest prepares the HTTP request before execution.
221func (e *KiroExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil }
222
223// Execute sends the request to Kiro API and returns the response.
224// Supports automatic token refresh on 401/403 errors.
225func (e *KiroExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
226 accessToken, profileArn := kiroCredentials(auth)
227 if accessToken == "" {
228 return resp, fmt.Errorf("kiro: access token not found in auth")
229 }
230
231 reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
232 defer reporter.trackFailure(ctx, &err)
233
234 // Check if token is expired before making request
235 if e.isTokenExpired(accessToken) {
236 log.Infof("kiro: access token expired, attempting refresh before request")
237 refreshedAuth, refreshErr := e.Refresh(ctx, auth)
238 if refreshErr != nil {
239 log.Warnf("kiro: pre-request token refresh failed: %v", refreshErr)
240 } else if refreshedAuth != nil {
241 auth = refreshedAuth
242 // Persist the refreshed auth to file so subsequent requests use it
243 if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
244 log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
245 }
246 accessToken, profileArn = kiroCredentials(auth)
247 log.Infof("kiro: token refreshed successfully before request")
248 }
249 }
250
251 from := opts.SourceFormat
252 to := sdktranslator.FromString("kiro")
253 body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
254
255 kiroModelID := e.mapModelToKiro(req.Model)
256
257 // Determine agentic mode and effective profile ARN using helper functions
258 isAgentic, isChatOnly := determineAgenticMode(req.Model)
259 effectiveProfileArn := getEffectiveProfileArnWithWarning(auth, profileArn)
260
261 // Execute with retry on 401/403 and 429 (quota exhausted)
262 // Note: currentOrigin and kiroPayload are built inside executeWithRetry for each endpoint
263 resp, err = e.executeWithRetry(ctx, auth, req, opts, accessToken, effectiveProfileArn, nil, body, from, to, reporter, "", kiroModelID, isAgentic, isChatOnly)
264 return resp, err
265}
266
267// executeWithRetry performs the actual HTTP request with automatic retry on auth errors.
268// Supports automatic fallback between endpoints with different quotas:
269// - Amazon Q endpoint (CLI origin) uses Amazon Q Developer quota
270// - CodeWhisperer endpoint (AI_EDITOR origin) uses Kiro IDE quota
271// Also supports multi-endpoint fallback similar to Antigravity implementation.
272func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, accessToken, profileArn string, kiroPayload, body []byte, from, to sdktranslator.Format, reporter *usageReporter, currentOrigin, kiroModelID string, isAgentic, isChatOnly bool) (cliproxyexecutor.Response, error) {
273 var resp cliproxyexecutor.Response
274 maxRetries := 2 // Allow retries for token refresh + endpoint fallback
275 endpointConfigs := getKiroEndpointConfigs(auth)
276 var last429Err error
277
278 for endpointIdx := 0; endpointIdx < len(endpointConfigs); endpointIdx++ {
279 endpointConfig := endpointConfigs[endpointIdx]
280 url := endpointConfig.URL
281 // Use this endpoint's compatible Origin (critical for avoiding 403 errors)
282 currentOrigin = endpointConfig.Origin
283
284 // Rebuild payload with the correct origin for this endpoint
285 // Each endpoint requires its matching Origin value in the request body
286 kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)
287
288 log.Debugf("kiro: trying endpoint %d/%d: %s (Name: %s, Origin: %s)",
289 endpointIdx+1, len(endpointConfigs), url, endpointConfig.Name, currentOrigin)
290
291 for attempt := 0; attempt <= maxRetries; attempt++ {
292 httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(kiroPayload))
293 if err != nil {
294 return resp, err
295 }
296
297 httpReq.Header.Set("Content-Type", kiroContentType)
298 httpReq.Header.Set("Accept", kiroAcceptStream)
299 // Use endpoint-specific X-Amz-Target (critical for avoiding 403 errors)
300 httpReq.Header.Set("X-Amz-Target", endpointConfig.AmzTarget)
301
302 // Use different headers based on auth type
303 // IDC auth uses Kiro IDE style headers (from kiro2api)
304 // Other auth types use Amazon Q CLI style headers
305 if isIDCAuth(auth) {
306 httpReq.Header.Set("User-Agent", kiroIDEUserAgent)
307 httpReq.Header.Set("X-Amz-User-Agent", kiroIDEAmzUserAgent)
308 httpReq.Header.Set("x-amzn-kiro-agent-mode", kiroIDEAgentModeSpec)
309 log.Debugf("kiro: using Kiro IDE headers for IDC auth")
310 } else {
311 httpReq.Header.Set("User-Agent", kiroUserAgent)
312 httpReq.Header.Set("X-Amz-User-Agent", kiroFullUserAgent)
313 }
314 httpReq.Header.Set("Amz-Sdk-Request", "attempt=1; max=3")
315 httpReq.Header.Set("Amz-Sdk-Invocation-Id", uuid.New().String())
316
317 // Bearer token authentication for all auth types (Builder ID, IDC, social, etc.)
318 httpReq.Header.Set("Authorization", "Bearer "+accessToken)
319
320 var attrs map[string]string
321 if auth != nil {
322 attrs = auth.Attributes
323 }
324 util.ApplyCustomHeadersFromAttrs(httpReq, attrs)
325
326 var authID, authLabel, authType, authValue string
327 if auth != nil {
328 authID = auth.ID
329 authLabel = auth.Label
330 authType, authValue = auth.AccountInfo()
331 }
332 recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
333 URL: url,
334 Method: http.MethodPost,
335 Headers: httpReq.Header.Clone(),
336 Body: kiroPayload,
337 Provider: e.Identifier(),
338 AuthID: authID,
339 AuthLabel: authLabel,
340 AuthType: authType,
341 AuthValue: authValue,
342 })
343
344 httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 120*time.Second)
345 httpResp, err := httpClient.Do(httpReq)
346 if err != nil {
347 recordAPIResponseError(ctx, e.cfg, err)
348 return resp, err
349 }
350 recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
351
352 // Handle 429 errors (quota exhausted) - try next endpoint
353 // Each endpoint has its own quota pool, so we can try different endpoints
354 if httpResp.StatusCode == 429 {
355 respBody, _ := io.ReadAll(httpResp.Body)
356 _ = httpResp.Body.Close()
357 appendAPIResponseChunk(ctx, e.cfg, respBody)
358
359 // Preserve last 429 so callers can correctly backoff when all endpoints are exhausted
360 last429Err = statusErr{code: httpResp.StatusCode, msg: string(respBody)}
361
362 log.Warnf("kiro: %s endpoint quota exhausted (429), will try next endpoint, body: %s",
363 endpointConfig.Name, summarizeErrorBody(httpResp.Header.Get("Content-Type"), respBody))
364
365 // Break inner retry loop to try next endpoint (which has different quota)
366 break
367 }
368
369 // Handle 5xx server errors with exponential backoff retry
370 if httpResp.StatusCode >= 500 && httpResp.StatusCode < 600 {
371 respBody, _ := io.ReadAll(httpResp.Body)
372 _ = httpResp.Body.Close()
373 appendAPIResponseChunk(ctx, e.cfg, respBody)
374
375 if attempt < maxRetries {
376 // Exponential backoff: 1s, 2s, 4s... (max 30s)
377 backoff := time.Duration(1<<attempt) * time.Second
378 if backoff > 30*time.Second {
379 backoff = 30 * time.Second
380 }
381 log.Warnf("kiro: server error %d, retrying in %v (attempt %d/%d)", httpResp.StatusCode, backoff, attempt+1, maxRetries)
382 time.Sleep(backoff)
383 continue
384 }
385 log.Errorf("kiro: server error %d after %d retries", httpResp.StatusCode, maxRetries)
386 return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
387 }
388
389 // Handle 401 errors with token refresh and retry
390 // 401 = Unauthorized (token expired/invalid) - refresh token
391 if httpResp.StatusCode == 401 {
392 respBody, _ := io.ReadAll(httpResp.Body)
393 _ = httpResp.Body.Close()
394 appendAPIResponseChunk(ctx, e.cfg, respBody)
395
396 if attempt < maxRetries {
397 log.Warnf("kiro: received 401 error, attempting token refresh and retry (attempt %d/%d)", attempt+1, maxRetries+1)
398
399 refreshedAuth, refreshErr := e.Refresh(ctx, auth)
400 if refreshErr != nil {
401 log.Errorf("kiro: token refresh failed: %v", refreshErr)
402 return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
403 }
404
405 if refreshedAuth != nil {
406 auth = refreshedAuth
407 // Persist the refreshed auth to file so subsequent requests use it
408 if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
409 log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
410 // Continue anyway - the token is valid for this request
411 }
412 accessToken, profileArn = kiroCredentials(auth)
413 // Rebuild payload with new profile ARN if changed
414 kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)
415 log.Infof("kiro: token refreshed successfully, retrying request")
416 continue
417 }
418 }
419
420 log.Warnf("kiro request error, status: 401, body: %s", summarizeErrorBody(httpResp.Header.Get("Content-Type"), respBody))
421 return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
422 }
423
424 // Handle 402 errors - Monthly Limit Reached
425 if httpResp.StatusCode == 402 {
426 respBody, _ := io.ReadAll(httpResp.Body)
427 _ = httpResp.Body.Close()
428 appendAPIResponseChunk(ctx, e.cfg, respBody)
429
430 log.Warnf("kiro: received 402 (monthly limit). Upstream body: %s", string(respBody))
431
432 // Return upstream error body directly
433 return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
434 }
435
436 // Handle 403 errors - Access Denied / Token Expired
437 // Do NOT switch endpoints for 403 errors
438 if httpResp.StatusCode == 403 {
439 respBody, _ := io.ReadAll(httpResp.Body)
440 _ = httpResp.Body.Close()
441 appendAPIResponseChunk(ctx, e.cfg, respBody)
442
443 // Log the 403 error details for debugging
444 log.Warnf("kiro: received 403 error (attempt %d/%d), body: %s", attempt+1, maxRetries+1, summarizeErrorBody(httpResp.Header.Get("Content-Type"), respBody))
445
446 respBodyStr := string(respBody)
447
448 // Check for SUSPENDED status - return immediately without retry
449 if strings.Contains(respBodyStr, "SUSPENDED") || strings.Contains(respBodyStr, "TEMPORARILY_SUSPENDED") {
450 log.Errorf("kiro: account is suspended, cannot proceed")
451 return resp, statusErr{code: httpResp.StatusCode, msg: "account suspended: " + string(respBody)}
452 }
453
454 // Check if this looks like a token-related 403 (some APIs return 403 for expired tokens)
455 isTokenRelated := strings.Contains(respBodyStr, "token") ||
456 strings.Contains(respBodyStr, "expired") ||
457 strings.Contains(respBodyStr, "invalid") ||
458 strings.Contains(respBodyStr, "unauthorized")
459
460 if isTokenRelated && attempt < maxRetries {
461 log.Warnf("kiro: 403 appears token-related, attempting token refresh")
462 refreshedAuth, refreshErr := e.Refresh(ctx, auth)
463 if refreshErr != nil {
464 log.Errorf("kiro: token refresh failed: %v", refreshErr)
465 // Token refresh failed - return error immediately
466 return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
467 }
468 if refreshedAuth != nil {
469 auth = refreshedAuth
470 // Persist the refreshed auth to file so subsequent requests use it
471 if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
472 log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
473 // Continue anyway - the token is valid for this request
474 }
475 accessToken, profileArn = kiroCredentials(auth)
476 kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)
477 log.Infof("kiro: token refreshed for 403, retrying request")
478 continue
479 }
480 }
481
482 // For non-token 403 or after max retries, return error immediately
483 // Do NOT switch endpoints for 403 errors
484 log.Warnf("kiro: 403 error, returning immediately (no endpoint switch)")
485 return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
486 }
487
488 if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
489 b, _ := io.ReadAll(httpResp.Body)
490 appendAPIResponseChunk(ctx, e.cfg, b)
491 log.Debugf("kiro request error, status: %d, body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
492 err = statusErr{code: httpResp.StatusCode, msg: string(b)}
493 if errClose := httpResp.Body.Close(); errClose != nil {
494 log.Errorf("response body close error: %v", errClose)
495 }
496 return resp, err
497 }
498
499 defer func() {
500 if errClose := httpResp.Body.Close(); errClose != nil {
501 log.Errorf("response body close error: %v", errClose)
502 }
503 }()
504
505 content, toolUses, usageInfo, stopReason, err := e.parseEventStream(httpResp.Body)
506 if err != nil {
507 recordAPIResponseError(ctx, e.cfg, err)
508 return resp, err
509 }
510
511 // Fallback for usage if missing from upstream
512 if usageInfo.TotalTokens == 0 {
513 if enc, encErr := getTokenizer(req.Model); encErr == nil {
514 if inp, countErr := countOpenAIChatTokens(enc, opts.OriginalRequest); countErr == nil {
515 usageInfo.InputTokens = inp
516 }
517 }
518 if len(content) > 0 {
519 // Use tiktoken for more accurate output token calculation
520 if enc, encErr := getTokenizer(req.Model); encErr == nil {
521 if tokenCount, countErr := enc.Count(content); countErr == nil {
522 usageInfo.OutputTokens = int64(tokenCount)
523 }
524 }
525 // Fallback to character count estimation if tiktoken fails
526 if usageInfo.OutputTokens == 0 {
527 usageInfo.OutputTokens = int64(len(content) / 4)
528 if usageInfo.OutputTokens == 0 {
529 usageInfo.OutputTokens = 1
530 }
531 }
532 }
533 usageInfo.TotalTokens = usageInfo.InputTokens + usageInfo.OutputTokens
534 }
535
536 appendAPIResponseChunk(ctx, e.cfg, []byte(content))
537 reporter.publish(ctx, usageInfo)
538
539 // Build response in Claude format for Kiro translator
540 // stopReason is extracted from upstream response by parseEventStream
541 kiroResponse := kiroclaude.BuildClaudeResponse(content, toolUses, req.Model, usageInfo, stopReason)
542 out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, kiroResponse, nil)
543 resp = cliproxyexecutor.Response{Payload: []byte(out)}
544 return resp, nil
545 }
546 // Inner retry loop exhausted for this endpoint, try next endpoint
547 // Note: This code is unreachable because all paths in the inner loop
548 // either return or continue. Kept as comment for documentation.
549 }
550
551 // All endpoints exhausted
552 if last429Err != nil {
553 return resp, last429Err
554 }
555 return resp, fmt.Errorf("kiro: all endpoints exhausted")
556}
557
558// ExecuteStream handles streaming requests to Kiro API.
559// Supports automatic token refresh on 401/403 errors and quota fallback on 429.
560func (e *KiroExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
561 accessToken, profileArn := kiroCredentials(auth)
562 if accessToken == "" {
563 return nil, fmt.Errorf("kiro: access token not found in auth")
564 }
565
566 reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
567 defer reporter.trackFailure(ctx, &err)
568
569 // Check if token is expired before making request
570 if e.isTokenExpired(accessToken) {
571 log.Infof("kiro: access token expired, attempting refresh before stream request")
572 refreshedAuth, refreshErr := e.Refresh(ctx, auth)
573 if refreshErr != nil {
574 log.Warnf("kiro: pre-request token refresh failed: %v", refreshErr)
575 } else if refreshedAuth != nil {
576 auth = refreshedAuth
577 // Persist the refreshed auth to file so subsequent requests use it
578 if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
579 log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
580 }
581 accessToken, profileArn = kiroCredentials(auth)
582 log.Infof("kiro: token refreshed successfully before stream request")
583 }
584 }
585
586 from := opts.SourceFormat
587 to := sdktranslator.FromString("kiro")
588 body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
589
590 kiroModelID := e.mapModelToKiro(req.Model)
591
592 // Determine agentic mode and effective profile ARN using helper functions
593 isAgentic, isChatOnly := determineAgenticMode(req.Model)
594 effectiveProfileArn := getEffectiveProfileArnWithWarning(auth, profileArn)
595
596 // Execute stream with retry on 401/403 and 429 (quota exhausted)
597 // Note: currentOrigin and kiroPayload are built inside executeStreamWithRetry for each endpoint
598 return e.executeStreamWithRetry(ctx, auth, req, opts, accessToken, effectiveProfileArn, nil, body, from, reporter, "", kiroModelID, isAgentic, isChatOnly)
599}
600
601// executeStreamWithRetry performs the streaming HTTP request with automatic retry on auth errors.
602// Supports automatic fallback between endpoints with different quotas:
603// - Amazon Q endpoint (CLI origin) uses Amazon Q Developer quota
604// - CodeWhisperer endpoint (AI_EDITOR origin) uses Kiro IDE quota
605// Also supports multi-endpoint fallback similar to Antigravity implementation.
606func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, accessToken, profileArn string, kiroPayload, body []byte, from sdktranslator.Format, reporter *usageReporter, currentOrigin, kiroModelID string, isAgentic, isChatOnly bool) (<-chan cliproxyexecutor.StreamChunk, error) {
607 maxRetries := 2 // Allow retries for token refresh + endpoint fallback
608 endpointConfigs := getKiroEndpointConfigs(auth)
609 var last429Err error
610
611 for endpointIdx := 0; endpointIdx < len(endpointConfigs); endpointIdx++ {
612 endpointConfig := endpointConfigs[endpointIdx]
613 url := endpointConfig.URL
614 // Use this endpoint's compatible Origin (critical for avoiding 403 errors)
615 currentOrigin = endpointConfig.Origin
616
617 // Rebuild payload with the correct origin for this endpoint
618 // Each endpoint requires its matching Origin value in the request body
619 kiroPayload, thinkingEnabled := buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)
620
621 log.Debugf("kiro: stream trying endpoint %d/%d: %s (Name: %s, Origin: %s)",
622 endpointIdx+1, len(endpointConfigs), url, endpointConfig.Name, currentOrigin)
623
624 for attempt := 0; attempt <= maxRetries; attempt++ {
625 httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(kiroPayload))
626 if err != nil {
627 return nil, err
628 }
629
630 httpReq.Header.Set("Content-Type", kiroContentType)
631 httpReq.Header.Set("Accept", kiroAcceptStream)
632 // Use endpoint-specific X-Amz-Target (critical for avoiding 403 errors)
633 httpReq.Header.Set("X-Amz-Target", endpointConfig.AmzTarget)
634
635 // Use different headers based on auth type
636 // IDC auth uses Kiro IDE style headers (from kiro2api)
637 // Other auth types use Amazon Q CLI style headers
638 if isIDCAuth(auth) {
639 httpReq.Header.Set("User-Agent", kiroIDEUserAgent)
640 httpReq.Header.Set("X-Amz-User-Agent", kiroIDEAmzUserAgent)
641 httpReq.Header.Set("x-amzn-kiro-agent-mode", kiroIDEAgentModeSpec)
642 log.Debugf("kiro: using Kiro IDE headers for IDC auth")
643 } else {
644 httpReq.Header.Set("User-Agent", kiroUserAgent)
645 httpReq.Header.Set("X-Amz-User-Agent", kiroFullUserAgent)
646 }
647 httpReq.Header.Set("Amz-Sdk-Request", "attempt=1; max=3")
648 httpReq.Header.Set("Amz-Sdk-Invocation-Id", uuid.New().String())
649
650 // Bearer token authentication for all auth types (Builder ID, IDC, social, etc.)
651 httpReq.Header.Set("Authorization", "Bearer "+accessToken)
652
653 var attrs map[string]string
654 if auth != nil {
655 attrs = auth.Attributes
656 }
657 util.ApplyCustomHeadersFromAttrs(httpReq, attrs)
658
659 var authID, authLabel, authType, authValue string
660 if auth != nil {
661 authID = auth.ID
662 authLabel = auth.Label
663 authType, authValue = auth.AccountInfo()
664 }
665 recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
666 URL: url,
667 Method: http.MethodPost,
668 Headers: httpReq.Header.Clone(),
669 Body: kiroPayload,
670 Provider: e.Identifier(),
671 AuthID: authID,
672 AuthLabel: authLabel,
673 AuthType: authType,
674 AuthValue: authValue,
675 })
676
677 httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
678 httpResp, err := httpClient.Do(httpReq)
679 if err != nil {
680 recordAPIResponseError(ctx, e.cfg, err)
681 return nil, err
682 }
683 recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
684
685 // Handle 429 errors (quota exhausted) - try next endpoint
686 // Each endpoint has its own quota pool, so we can try different endpoints
687 if httpResp.StatusCode == 429 {
688 respBody, _ := io.ReadAll(httpResp.Body)
689 _ = httpResp.Body.Close()
690 appendAPIResponseChunk(ctx, e.cfg, respBody)
691
692 // Preserve last 429 so callers can correctly backoff when all endpoints are exhausted
693 last429Err = statusErr{code: httpResp.StatusCode, msg: string(respBody)}
694
695 log.Warnf("kiro: stream %s endpoint quota exhausted (429), will try next endpoint, body: %s",
696 endpointConfig.Name, summarizeErrorBody(httpResp.Header.Get("Content-Type"), respBody))
697
698 // Break inner retry loop to try next endpoint (which has different quota)
699 break
700 }
701
702 // Handle 5xx server errors with exponential backoff retry
703 if httpResp.StatusCode >= 500 && httpResp.StatusCode < 600 {
704 respBody, _ := io.ReadAll(httpResp.Body)
705 _ = httpResp.Body.Close()
706 appendAPIResponseChunk(ctx, e.cfg, respBody)
707
708 if attempt < maxRetries {
709 // Exponential backoff: 1s, 2s, 4s... (max 30s)
710 backoff := time.Duration(1<<attempt) * time.Second
711 if backoff > 30*time.Second {
712 backoff = 30 * time.Second
713 }
714 log.Warnf("kiro: stream server error %d, retrying in %v (attempt %d/%d)", httpResp.StatusCode, backoff, attempt+1, maxRetries)
715 time.Sleep(backoff)
716 continue
717 }
718 log.Errorf("kiro: stream server error %d after %d retries", httpResp.StatusCode, maxRetries)
719 return nil, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
720 }
721
722 // Handle 400 errors - Credential/Validation issues
723 // Do NOT switch endpoints - return error immediately
724 if httpResp.StatusCode == 400 {
725 respBody, _ := io.ReadAll(httpResp.Body)
726 _ = httpResp.Body.Close()
727 appendAPIResponseChunk(ctx, e.cfg, respBody)
728
729 log.Warnf("kiro: received 400 error (attempt %d/%d), body: %s", attempt+1, maxRetries+1, summarizeErrorBody(httpResp.Header.Get("Content-Type"), respBody))
730
731 // 400 errors indicate request validation issues - return immediately without retry
732 return nil, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
733 }
734
735 // Handle 401 errors with token refresh and retry
736 // 401 = Unauthorized (token expired/invalid) - refresh token
737 if httpResp.StatusCode == 401 {
738 respBody, _ := io.ReadAll(httpResp.Body)
739 _ = httpResp.Body.Close()
740 appendAPIResponseChunk(ctx, e.cfg, respBody)
741
742 if attempt < maxRetries {
743 log.Warnf("kiro: stream received 401 error, attempting token refresh and retry (attempt %d/%d)", attempt+1, maxRetries+1)
744
745 refreshedAuth, refreshErr := e.Refresh(ctx, auth)
746 if refreshErr != nil {
747 log.Errorf("kiro: token refresh failed: %v", refreshErr)
748 return nil, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
749 }
750
751 if refreshedAuth != nil {
752 auth = refreshedAuth
753 // Persist the refreshed auth to file so subsequent requests use it
754 if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
755 log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
756 // Continue anyway - the token is valid for this request
757 }
758 accessToken, profileArn = kiroCredentials(auth)
759 // Rebuild payload with new profile ARN if changed
760 kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)
761 log.Infof("kiro: token refreshed successfully, retrying stream request")
762 continue
763 }
764 }
765
766 log.Warnf("kiro stream error, status: 401, body: %s", string(respBody))
767 return nil, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
768 }
769
770 // Handle 402 errors - Monthly Limit Reached
771 if httpResp.StatusCode == 402 {
772 respBody, _ := io.ReadAll(httpResp.Body)
773 _ = httpResp.Body.Close()
774 appendAPIResponseChunk(ctx, e.cfg, respBody)
775
776 log.Warnf("kiro: stream received 402 (monthly limit). Upstream body: %s", string(respBody))
777
778 // Return upstream error body directly
779 return nil, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
780 }
781
782 // Handle 403 errors - Access Denied / Token Expired
783 // Do NOT switch endpoints for 403 errors
784 if httpResp.StatusCode == 403 {
785 respBody, _ := io.ReadAll(httpResp.Body)
786 _ = httpResp.Body.Close()
787 appendAPIResponseChunk(ctx, e.cfg, respBody)
788
789 // Log the 403 error details for debugging
790 log.Warnf("kiro: stream received 403 error (attempt %d/%d), body: %s", attempt+1, maxRetries+1, string(respBody))
791
792 respBodyStr := string(respBody)
793
794 // Check for SUSPENDED status - return immediately without retry
795 if strings.Contains(respBodyStr, "SUSPENDED") || strings.Contains(respBodyStr, "TEMPORARILY_SUSPENDED") {
796 log.Errorf("kiro: account is suspended, cannot proceed")
797 return nil, statusErr{code: httpResp.StatusCode, msg: "account suspended: " + string(respBody)}
798 }
799
800 // Check if this looks like a token-related 403 (some APIs return 403 for expired tokens)
801 isTokenRelated := strings.Contains(respBodyStr, "token") ||
802 strings.Contains(respBodyStr, "expired") ||
803 strings.Contains(respBodyStr, "invalid") ||
804 strings.Contains(respBodyStr, "unauthorized")
805
806 if isTokenRelated && attempt < maxRetries {
807 log.Warnf("kiro: 403 appears token-related, attempting token refresh")
808 refreshedAuth, refreshErr := e.Refresh(ctx, auth)
809 if refreshErr != nil {
810 log.Errorf("kiro: token refresh failed: %v", refreshErr)
811 // Token refresh failed - return error immediately
812 return nil, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
813 }
814 if refreshedAuth != nil {
815 auth = refreshedAuth
816 // Persist the refreshed auth to file so subsequent requests use it
817 if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
818 log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
819 // Continue anyway - the token is valid for this request
820 }
821 accessToken, profileArn = kiroCredentials(auth)
822 kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)
823 log.Infof("kiro: token refreshed for 403, retrying stream request")
824 continue
825 }
826 }
827
828 // For non-token 403 or after max retries, return error immediately
829 // Do NOT switch endpoints for 403 errors
830 log.Warnf("kiro: 403 error, returning immediately (no endpoint switch)")
831 return nil, statusErr{code: httpResp.StatusCode, msg: string(respBody)}
832 }
833
834 if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
835 b, _ := io.ReadAll(httpResp.Body)
836 appendAPIResponseChunk(ctx, e.cfg, b)
837 log.Debugf("kiro stream error, status: %d, body: %s", httpResp.StatusCode, string(b))
838 if errClose := httpResp.Body.Close(); errClose != nil {
839 log.Errorf("response body close error: %v", errClose)
840 }
841 return nil, statusErr{code: httpResp.StatusCode, msg: string(b)}
842 }
843
844 out := make(chan cliproxyexecutor.StreamChunk)
845
846 go func(resp *http.Response, thinkingEnabled bool) {
847 defer close(out)
848 defer func() {
849 if r := recover(); r != nil {
850 log.Errorf("kiro: panic in stream handler: %v", r)
851 out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("internal error: %v", r)}
852 }
853 }()
854 defer func() {
855 if errClose := resp.Body.Close(); errClose != nil {
856 log.Errorf("response body close error: %v", errClose)
857 }
858 }()
859
860 // Kiro API always returns <thinking> tags regardless of request parameters
861 // So we always enable thinking parsing for Kiro responses
862 log.Debugf("kiro: stream thinkingEnabled = %v (always true for Kiro)", thinkingEnabled)
863
864 e.streamToChannel(ctx, resp.Body, out, from, req.Model, opts.OriginalRequest, body, reporter, thinkingEnabled)
865 }(httpResp, thinkingEnabled)
866
867 return out, nil
868 }
869 // Inner retry loop exhausted for this endpoint, try next endpoint
870 // Note: This code is unreachable because all paths in the inner loop
871 // either return or continue. Kept as comment for documentation.
872 }
873
874 // All endpoints exhausted
875 if last429Err != nil {
876 return nil, last429Err
877 }
878 return nil, fmt.Errorf("kiro: stream all endpoints exhausted")
879}
880
881// kiroCredentials extracts access token and profile ARN from auth.
882func kiroCredentials(auth *cliproxyauth.Auth) (accessToken, profileArn string) {
883 if auth == nil {
884 return "", ""
885 }
886
887 // Try Metadata first (wrapper format)
888 if auth.Metadata != nil {
889 if token, ok := auth.Metadata["access_token"].(string); ok {
890 accessToken = token
891 }
892 if arn, ok := auth.Metadata["profile_arn"].(string); ok {
893 profileArn = arn
894 }
895 }
896
897 // Try Attributes
898 if accessToken == "" && auth.Attributes != nil {
899 accessToken = auth.Attributes["access_token"]
900 profileArn = auth.Attributes["profile_arn"]
901 }
902
903 // Try direct fields from flat JSON format (new AWS Builder ID format)
904 if accessToken == "" && auth.Metadata != nil {
905 if token, ok := auth.Metadata["accessToken"].(string); ok {
906 accessToken = token
907 }
908 if arn, ok := auth.Metadata["profileArn"].(string); ok {
909 profileArn = arn
910 }
911 }
912
913 return accessToken, profileArn
914}
915
916// findRealThinkingEndTag finds the real </thinking> end tag, skipping false positives.
917// Returns -1 if no real end tag is found.
918//
919// Real </thinking> tags from Kiro API have specific characteristics:
920// - Usually preceded by newline (.\n</thinking>)
921// - Usually followed by newline (\n\n)
922// - Not inside code blocks or inline code
923//
924// False positives (discussion text) have characteristics:
925// - In the middle of a sentence
926// - Preceded by discussion words like "标签", "tag", "returns"
927// - Inside code blocks or inline code
928//
929// Parameters:
930// - content: the content to search in
931// - alreadyInCodeBlock: whether we're already inside a code block from previous chunks
932// - alreadyInInlineCode: whether we're already inside inline code from previous chunks
933func findRealThinkingEndTag(content string, alreadyInCodeBlock, alreadyInInlineCode bool) int {
934 searchStart := 0
935 for {
936 endIdx := strings.Index(content[searchStart:], kirocommon.ThinkingEndTag)
937 if endIdx < 0 {
938 return -1
939 }
940 endIdx += searchStart // Adjust to absolute position
941
942 textBeforeEnd := content[:endIdx]
943 textAfterEnd := content[endIdx+len(kirocommon.ThinkingEndTag):]
944
945 // Check 1: Is it inside inline code?
946 // Count backticks in current content and add state from previous chunks
947 backtickCount := strings.Count(textBeforeEnd, "`")
948 effectiveInInlineCode := alreadyInInlineCode
949 if backtickCount%2 == 1 {
950 effectiveInInlineCode = !effectiveInInlineCode
951 }
952 if effectiveInInlineCode {
953 log.Debugf("kiro: found </thinking> inside inline code at pos %d, skipping", endIdx)
954 searchStart = endIdx + len(kirocommon.ThinkingEndTag)
955 continue
956 }
957
958 // Check 2: Is it inside a code block?
959 // Count fences in current content and add state from previous chunks
960 fenceCount := strings.Count(textBeforeEnd, "```")
961 altFenceCount := strings.Count(textBeforeEnd, "~~~")
962 effectiveInCodeBlock := alreadyInCodeBlock
963 if fenceCount%2 == 1 || altFenceCount%2 == 1 {
964 effectiveInCodeBlock = !effectiveInCodeBlock
965 }
966 if effectiveInCodeBlock {
967 log.Debugf("kiro: found </thinking> inside code block at pos %d, skipping", endIdx)
968 searchStart = endIdx + len(kirocommon.ThinkingEndTag)
969 continue
970 }
971
972 // Check 3: Real </thinking> tags are usually preceded by newline or at start
973 // and followed by newline or at end. Check the format.
974 charBeforeTag := byte(0)
975 if endIdx > 0 {
976 charBeforeTag = content[endIdx-1]
977 }
978 charAfterTag := byte(0)
979 if len(textAfterEnd) > 0 {
980 charAfterTag = textAfterEnd[0]
981 }
982
983 // Real end tag format: preceded by newline OR end of sentence (. ! ?)
984 // and followed by newline OR end of content
985 isPrecededByNewlineOrSentenceEnd := charBeforeTag == '\n' || charBeforeTag == '.' ||
986 charBeforeTag == '!' || charBeforeTag == '?' || charBeforeTag == 0
987 isFollowedByNewlineOrEnd := charAfterTag == '\n' || charAfterTag == 0
988
989 // If the tag has proper formatting (newline before/after), it's likely real
990 if isPrecededByNewlineOrSentenceEnd && isFollowedByNewlineOrEnd {
991 log.Debugf("kiro: found properly formatted </thinking> at pos %d", endIdx)
992 return endIdx
993 }
994
995 // Check 4: Is the tag preceded by discussion keywords on the same line?
996 lastNewlineIdx := strings.LastIndex(textBeforeEnd, "\n")
997 lineBeforeTag := textBeforeEnd
998 if lastNewlineIdx >= 0 {
999 lineBeforeTag = textBeforeEnd[lastNewlineIdx+1:]
1000 }
1001 lineBeforeTagLower := strings.ToLower(lineBeforeTag)
1002
1003 // Discussion patterns - if found, this is likely discussion text
1004 discussionPatterns := []string{
1005 "标签", "返回", "输出", "包含", "使用", "解析", "转换", "生成", // Chinese
1006 "tag", "return", "output", "contain", "use", "parse", "emit", "convert", "generate", // English
1007 "<thinking>", // discussing both tags together
1008 "`</thinking>`", // explicitly in inline code
1009 }
1010 isDiscussion := false
1011 for _, pattern := range discussionPatterns {
1012 if strings.Contains(lineBeforeTagLower, pattern) {
1013 isDiscussion = true
1014 break
1015 }
1016 }
1017 if isDiscussion {
1018 log.Debugf("kiro: found </thinking> after discussion text at pos %d, skipping", endIdx)
1019 searchStart = endIdx + len(kirocommon.ThinkingEndTag)
1020 continue
1021 }
1022
1023 // Check 5: Is there text immediately after on the same line?
1024 // Real end tags don't have text immediately after on the same line
1025 if len(textAfterEnd) > 0 && charAfterTag != '\n' && charAfterTag != 0 {
1026 // Find the next newline
1027 nextNewline := strings.Index(textAfterEnd, "\n")
1028 var textOnSameLine string
1029 if nextNewline >= 0 {
1030 textOnSameLine = textAfterEnd[:nextNewline]
1031 } else {
1032 textOnSameLine = textAfterEnd
1033 }
1034 // If there's non-whitespace text on the same line after the tag, it's discussion
1035 if strings.TrimSpace(textOnSameLine) != "" {
1036 log.Debugf("kiro: found </thinking> with text after on same line at pos %d, skipping", endIdx)
1037 searchStart = endIdx + len(kirocommon.ThinkingEndTag)
1038 continue
1039 }
1040 }
1041
1042 // Check 6: Is there another <thinking> tag after this </thinking>?
1043 if strings.Contains(textAfterEnd, kirocommon.ThinkingStartTag) {
1044 nextStartIdx := strings.Index(textAfterEnd, kirocommon.ThinkingStartTag)
1045 textBeforeNextStart := textAfterEnd[:nextStartIdx]
1046 nextBacktickCount := strings.Count(textBeforeNextStart, "`")
1047 nextFenceCount := strings.Count(textBeforeNextStart, "```")
1048 nextAltFenceCount := strings.Count(textBeforeNextStart, "~~~")
1049
1050 // If the next <thinking> is NOT in code, then this </thinking> is discussion text
1051 if nextBacktickCount%2 == 0 && nextFenceCount%2 == 0 && nextAltFenceCount%2 == 0 {
1052 log.Debugf("kiro: found </thinking> followed by <thinking> at pos %d, likely discussion text, skipping", endIdx)
1053 searchStart = endIdx + len(kirocommon.ThinkingEndTag)
1054 continue
1055 }
1056 }
1057
1058 // This looks like a real end tag
1059 return endIdx
1060 }
1061}
1062
1063// determineAgenticMode determines if the model is an agentic or chat-only variant.
1064// Returns (isAgentic, isChatOnly) based on model name suffixes.
1065func determineAgenticMode(model string) (isAgentic, isChatOnly bool) {
1066 isAgentic = strings.HasSuffix(model, "-agentic")
1067 isChatOnly = strings.HasSuffix(model, "-chat")
1068 return isAgentic, isChatOnly
1069}
1070
1071// getEffectiveProfileArn determines if profileArn should be included based on auth method.
1072// profileArn is only needed for social auth (Google OAuth), not for builder-id (AWS SSO).
1073func getEffectiveProfileArn(auth *cliproxyauth.Auth, profileArn string) string {
1074 if auth != nil && auth.Metadata != nil {
1075 if authMethod, ok := auth.Metadata["auth_method"].(string); ok && authMethod == "builder-id" {
1076 return "" // Don't include profileArn for builder-id auth
1077 }
1078 }
1079 return profileArn
1080}
1081
1082// getEffectiveProfileArnWithWarning determines if profileArn should be included based on auth method,
1083// and logs a warning if profileArn is missing for non-builder-id auth.
1084// This consolidates the auth_method check that was previously done separately.
1085func getEffectiveProfileArnWithWarning(auth *cliproxyauth.Auth, profileArn string) string {
1086 if auth != nil && auth.Metadata != nil {
1087 if authMethod, ok := auth.Metadata["auth_method"].(string); ok && (authMethod == "builder-id" || authMethod == "idc") {
1088 // builder-id and idc auth don't need profileArn
1089 return ""
1090 }
1091 }
1092 // For non-builder-id/idc auth (social auth), profileArn is required
1093 if profileArn == "" {
1094 log.Warnf("kiro: profile ARN not found in auth, API calls may fail")
1095 }
1096 return profileArn
1097}
1098
1099// mapModelToKiro maps external model names to Kiro model IDs.
1100// Supports both Kiro and Amazon Q prefixes since they use the same API.
1101// Agentic variants (-agentic suffix) map to the same backend model IDs.
1102func (e *KiroExecutor) mapModelToKiro(model string) string {
1103 modelMap := map[string]string{
1104 // Amazon Q format (amazonq- prefix) - same API as Kiro
1105 "amazonq-auto": "auto",
1106 "amazonq-claude-opus-4-5": "claude-opus-4.5",
1107 "amazonq-claude-sonnet-4-5": "claude-sonnet-4.5",
1108 "amazonq-claude-sonnet-4-5-20250929": "claude-sonnet-4.5",
1109 "amazonq-claude-sonnet-4": "claude-sonnet-4",
1110 "amazonq-claude-sonnet-4-20250514": "claude-sonnet-4",
1111 "amazonq-claude-haiku-4-5": "claude-haiku-4.5",
1112 // Kiro format (kiro- prefix) - valid model names that should be preserved
1113 "kiro-claude-opus-4-5": "claude-opus-4.5",
1114 "kiro-claude-sonnet-4-5": "claude-sonnet-4.5",
1115 "kiro-claude-sonnet-4-5-20250929": "claude-sonnet-4.5",
1116 "kiro-claude-sonnet-4": "claude-sonnet-4",
1117 "kiro-claude-sonnet-4-20250514": "claude-sonnet-4",
1118 "kiro-claude-haiku-4-5": "claude-haiku-4.5",
1119 "kiro-auto": "auto",
1120 // Native format (no prefix) - used by Kiro IDE directly
1121 "claude-opus-4-5": "claude-opus-4.5",
1122 "claude-opus-4.5": "claude-opus-4.5",
1123 "claude-haiku-4-5": "claude-haiku-4.5",
1124 "claude-haiku-4.5": "claude-haiku-4.5",
1125 "claude-sonnet-4-5": "claude-sonnet-4.5",
1126 "claude-sonnet-4-5-20250929": "claude-sonnet-4.5",
1127 "claude-sonnet-4.5": "claude-sonnet-4.5",
1128 "claude-sonnet-4": "claude-sonnet-4",
1129 "claude-sonnet-4-20250514": "claude-sonnet-4",
1130 "auto": "auto",
1131 // Agentic variants (same backend model IDs, but with special system prompt)
1132 "claude-opus-4.5-agentic": "claude-opus-4.5",
1133 "claude-sonnet-4.5-agentic": "claude-sonnet-4.5",
1134 "claude-sonnet-4-agentic": "claude-sonnet-4",
1135 "claude-haiku-4.5-agentic": "claude-haiku-4.5",
1136 "kiro-claude-opus-4-5-agentic": "claude-opus-4.5",
1137 "kiro-claude-sonnet-4-5-agentic": "claude-sonnet-4.5",
1138 "kiro-claude-sonnet-4-agentic": "claude-sonnet-4",
1139 "kiro-claude-haiku-4-5-agentic": "claude-haiku-4.5",
1140 }
1141 if kiroID, ok := modelMap[model]; ok {
1142 return kiroID
1143 }
1144
1145 // Smart fallback: try to infer model type from name patterns
1146 modelLower := strings.ToLower(model)
1147
1148 // Check for Haiku variants
1149 if strings.Contains(modelLower, "haiku") {
1150 log.Debugf("kiro: unknown Haiku model '%s', mapping to claude-haiku-4.5", model)
1151 return "claude-haiku-4.5"
1152 }
1153
1154 // Check for Sonnet variants
1155 if strings.Contains(modelLower, "sonnet") {
1156 // Check for specific version patterns
1157 if strings.Contains(modelLower, "3-7") || strings.Contains(modelLower, "3.7") {
1158 log.Debugf("kiro: unknown Sonnet 3.7 model '%s', mapping to claude-3-7-sonnet-20250219", model)
1159 return "claude-3-7-sonnet-20250219"
1160 }
1161 if strings.Contains(modelLower, "4-5") || strings.Contains(modelLower, "4.5") {
1162 log.Debugf("kiro: unknown Sonnet 4.5 model '%s', mapping to claude-sonnet-4.5", model)
1163 return "claude-sonnet-4.5"
1164 }
1165 // Default to Sonnet 4
1166 log.Debugf("kiro: unknown Sonnet model '%s', mapping to claude-sonnet-4", model)
1167 return "claude-sonnet-4"
1168 }
1169
1170 // Check for Opus variants
1171 if strings.Contains(modelLower, "opus") {
1172 log.Debugf("kiro: unknown Opus model '%s', mapping to claude-opus-4.5", model)
1173 return "claude-opus-4.5"
1174 }
1175
1176 // Final fallback to Sonnet 4.5 (most commonly used model)
1177 log.Warnf("kiro: unknown model '%s', falling back to claude-sonnet-4.5", model)
1178 return "claude-sonnet-4.5"
1179}
1180
1181// EventStreamError represents an Event Stream processing error
1182type EventStreamError struct {
1183 Type string // "fatal", "malformed"
1184 Message string
1185 Cause error
1186}
1187
1188func (e *EventStreamError) Error() string {
1189 if e.Cause != nil {
1190 return fmt.Sprintf("event stream %s: %s: %v", e.Type, e.Message, e.Cause)
1191 }
1192 return fmt.Sprintf("event stream %s: %s", e.Type, e.Message)
1193}
1194
1195// eventStreamMessage represents a parsed AWS Event Stream message
1196type eventStreamMessage struct {
1197 EventType string // Event type from headers (e.g., "assistantResponseEvent")
1198 Payload []byte // JSON payload of the message
1199}
1200
1201// NOTE: Request building functions moved to internal/translator/kiro/claude/kiro_claude_request.go
1202// The executor now uses kiroclaude.BuildKiroPayload() instead
1203
1204// parseEventStream parses AWS Event Stream binary format.
1205// Extracts text content, tool uses, and stop_reason from the response.
1206// Supports embedded [Called ...] tool calls and input buffering for toolUseEvent.
1207// Returns: content, toolUses, usageInfo, stopReason, error
1208func (e *KiroExecutor) parseEventStream(body io.Reader) (string, []kiroclaude.KiroToolUse, usage.Detail, string, error) {
1209 var content strings.Builder
1210 var toolUses []kiroclaude.KiroToolUse
1211 var usageInfo usage.Detail
1212 var stopReason string // Extracted from upstream response
1213 reader := bufio.NewReader(body)
1214
1215 // Tool use state tracking for input buffering and deduplication
1216 processedIDs := make(map[string]bool)
1217 var currentToolUse *kiroclaude.ToolUseState
1218
1219 // Upstream usage tracking - Kiro API returns credit usage and context percentage
1220 var upstreamContextPercentage float64 // Context usage percentage from upstream (e.g., 78.56)
1221
1222 for {
1223 msg, eventErr := e.readEventStreamMessage(reader)
1224 if eventErr != nil {
1225 log.Errorf("kiro: parseEventStream error: %v", eventErr)
1226 return content.String(), toolUses, usageInfo, stopReason, eventErr
1227 }
1228 if msg == nil {
1229 // Normal end of stream (EOF)
1230 break
1231 }
1232
1233 eventType := msg.EventType
1234 payload := msg.Payload
1235 if len(payload) == 0 {
1236 continue
1237 }
1238
1239 var event map[string]interface{}
1240 if err := json.Unmarshal(payload, &event); err != nil {
1241 log.Debugf("kiro: skipping malformed event: %v", err)
1242 continue
1243 }
1244
1245 // Check for error/exception events in the payload (Kiro API may return errors with HTTP 200)
1246 // These can appear as top-level fields or nested within the event
1247 if errType, hasErrType := event["_type"].(string); hasErrType {
1248 // AWS-style error: {"_type": "com.amazon.aws.codewhisperer#ValidationException", "message": "..."}
1249 errMsg := ""
1250 if msg, ok := event["message"].(string); ok {
1251 errMsg = msg
1252 }
1253 log.Errorf("kiro: received AWS error in event stream: type=%s, message=%s", errType, errMsg)
1254 return "", nil, usageInfo, stopReason, fmt.Errorf("kiro API error: %s - %s", errType, errMsg)
1255 }
1256 if errType, hasErrType := event["type"].(string); hasErrType && (errType == "error" || errType == "exception") {
1257 // Generic error event
1258 errMsg := ""
1259 if msg, ok := event["message"].(string); ok {
1260 errMsg = msg
1261 } else if errObj, ok := event["error"].(map[string]interface{}); ok {
1262 if msg, ok := errObj["message"].(string); ok {
1263 errMsg = msg
1264 }
1265 }
1266 log.Errorf("kiro: received error event in stream: type=%s, message=%s", errType, errMsg)
1267 return "", nil, usageInfo, stopReason, fmt.Errorf("kiro API error: %s", errMsg)
1268 }
1269
1270 // Extract stop_reason from various event formats
1271 // Kiro/Amazon Q API may include stop_reason in different locations
1272 if sr := kirocommon.GetString(event, "stop_reason"); sr != "" {
1273 stopReason = sr
1274 log.Debugf("kiro: parseEventStream found stop_reason (top-level): %s", stopReason)
1275 }
1276 if sr := kirocommon.GetString(event, "stopReason"); sr != "" {
1277 stopReason = sr
1278 log.Debugf("kiro: parseEventStream found stopReason (top-level): %s", stopReason)
1279 }
1280
1281 // Handle different event types
1282 switch eventType {
1283 case "followupPromptEvent":
1284 // Filter out followupPrompt events - these are UI suggestions, not content
1285 log.Debugf("kiro: parseEventStream ignoring followupPrompt event")
1286 continue
1287
1288 case "assistantResponseEvent":
1289 if assistantResp, ok := event["assistantResponseEvent"].(map[string]interface{}); ok {
1290 if contentText, ok := assistantResp["content"].(string); ok {
1291 content.WriteString(contentText)
1292 }
1293 // Extract stop_reason from assistantResponseEvent
1294 if sr := kirocommon.GetString(assistantResp, "stop_reason"); sr != "" {
1295 stopReason = sr
1296 log.Debugf("kiro: parseEventStream found stop_reason in assistantResponseEvent: %s", stopReason)
1297 }
1298 if sr := kirocommon.GetString(assistantResp, "stopReason"); sr != "" {
1299 stopReason = sr
1300 log.Debugf("kiro: parseEventStream found stopReason in assistantResponseEvent: %s", stopReason)
1301 }
1302 // Extract tool uses from response
1303 if toolUsesRaw, ok := assistantResp["toolUses"].([]interface{}); ok {
1304 for _, tuRaw := range toolUsesRaw {
1305 if tu, ok := tuRaw.(map[string]interface{}); ok {
1306 toolUseID := kirocommon.GetStringValue(tu, "toolUseId")
1307 // Check for duplicate
1308 if processedIDs[toolUseID] {
1309 log.Debugf("kiro: skipping duplicate tool use from assistantResponse: %s", toolUseID)
1310 continue
1311 }
1312 processedIDs[toolUseID] = true
1313
1314 toolUse := kiroclaude.KiroToolUse{
1315 ToolUseID: toolUseID,
1316 Name: kirocommon.GetStringValue(tu, "name"),
1317 }
1318 if input, ok := tu["input"].(map[string]interface{}); ok {
1319 toolUse.Input = input
1320 }
1321 toolUses = append(toolUses, toolUse)
1322 }
1323 }
1324 }
1325 }
1326 // Also try direct format
1327 if contentText, ok := event["content"].(string); ok {
1328 content.WriteString(contentText)
1329 }
1330 // Direct tool uses
1331 if toolUsesRaw, ok := event["toolUses"].([]interface{}); ok {
1332 for _, tuRaw := range toolUsesRaw {
1333 if tu, ok := tuRaw.(map[string]interface{}); ok {
1334 toolUseID := kirocommon.GetStringValue(tu, "toolUseId")
1335 // Check for duplicate
1336 if processedIDs[toolUseID] {
1337 log.Debugf("kiro: skipping duplicate direct tool use: %s", toolUseID)
1338 continue
1339 }
1340 processedIDs[toolUseID] = true
1341
1342 toolUse := kiroclaude.KiroToolUse{
1343 ToolUseID: toolUseID,
1344 Name: kirocommon.GetStringValue(tu, "name"),
1345 }
1346 if input, ok := tu["input"].(map[string]interface{}); ok {
1347 toolUse.Input = input
1348 }
1349 toolUses = append(toolUses, toolUse)
1350 }
1351 }
1352 }
1353
1354 case "toolUseEvent":
1355 // Handle dedicated tool use events with input buffering
1356 completedToolUses, newState := kiroclaude.ProcessToolUseEvent(event, currentToolUse, processedIDs)
1357 currentToolUse = newState
1358 toolUses = append(toolUses, completedToolUses...)
1359
1360 case "supplementaryWebLinksEvent":
1361 if inputTokens, ok := event["inputTokens"].(float64); ok {
1362 usageInfo.InputTokens = int64(inputTokens)
1363 }
1364 if outputTokens, ok := event["outputTokens"].(float64); ok {
1365 usageInfo.OutputTokens = int64(outputTokens)
1366 }
1367
1368 case "messageStopEvent", "message_stop":
1369 // Handle message stop events which may contain stop_reason
1370 if sr := kirocommon.GetString(event, "stop_reason"); sr != "" {
1371 stopReason = sr
1372 log.Debugf("kiro: parseEventStream found stop_reason in messageStopEvent: %s", stopReason)
1373 }
1374 if sr := kirocommon.GetString(event, "stopReason"); sr != "" {
1375 stopReason = sr
1376 log.Debugf("kiro: parseEventStream found stopReason in messageStopEvent: %s", stopReason)
1377 }
1378
1379 case "messageMetadataEvent", "metadataEvent":
1380 // Handle message metadata events which contain token counts
1381 // Official format: { tokenUsage: { outputTokens, totalTokens, uncachedInputTokens, cacheReadInputTokens, cacheWriteInputTokens, contextUsagePercentage } }
1382 var metadata map[string]interface{}
1383 if m, ok := event["messageMetadataEvent"].(map[string]interface{}); ok {
1384 metadata = m
1385 } else if m, ok := event["metadataEvent"].(map[string]interface{}); ok {
1386 metadata = m
1387 } else {
1388 metadata = event // event itself might be the metadata
1389 }
1390
1391 // Check for nested tokenUsage object (official format)
1392 if tokenUsage, ok := metadata["tokenUsage"].(map[string]interface{}); ok {
1393 // outputTokens - precise output token count
1394 if outputTokens, ok := tokenUsage["outputTokens"].(float64); ok {
1395 usageInfo.OutputTokens = int64(outputTokens)
1396 log.Infof("kiro: parseEventStream found precise outputTokens in tokenUsage: %d", usageInfo.OutputTokens)
1397 }
1398 // totalTokens - precise total token count
1399 if totalTokens, ok := tokenUsage["totalTokens"].(float64); ok {
1400 usageInfo.TotalTokens = int64(totalTokens)
1401 log.Infof("kiro: parseEventStream found precise totalTokens in tokenUsage: %d", usageInfo.TotalTokens)
1402 }
1403 // uncachedInputTokens - input tokens not from cache
1404 if uncachedInputTokens, ok := tokenUsage["uncachedInputTokens"].(float64); ok {
1405 usageInfo.InputTokens = int64(uncachedInputTokens)
1406 log.Infof("kiro: parseEventStream found uncachedInputTokens in tokenUsage: %d", usageInfo.InputTokens)
1407 }
1408 // cacheReadInputTokens - tokens read from cache
1409 if cacheReadTokens, ok := tokenUsage["cacheReadInputTokens"].(float64); ok {
1410 // Add to input tokens if we have uncached tokens, otherwise use as input
1411 if usageInfo.InputTokens > 0 {
1412 usageInfo.InputTokens += int64(cacheReadTokens)
1413 } else {
1414 usageInfo.InputTokens = int64(cacheReadTokens)
1415 }
1416 log.Debugf("kiro: parseEventStream found cacheReadInputTokens in tokenUsage: %d", int64(cacheReadTokens))
1417 }
1418 // contextUsagePercentage - can be used as fallback for input token estimation
1419 if ctxPct, ok := tokenUsage["contextUsagePercentage"].(float64); ok {
1420 upstreamContextPercentage = ctxPct
1421 log.Debugf("kiro: parseEventStream found contextUsagePercentage in tokenUsage: %.2f%%", ctxPct)
1422 }
1423 }
1424
1425 // Fallback: check for direct fields in metadata (legacy format)
1426 if usageInfo.InputTokens == 0 {
1427 if inputTokens, ok := metadata["inputTokens"].(float64); ok {
1428 usageInfo.InputTokens = int64(inputTokens)
1429 log.Debugf("kiro: parseEventStream found inputTokens in messageMetadataEvent: %d", usageInfo.InputTokens)
1430 }
1431 }
1432 if usageInfo.OutputTokens == 0 {
1433 if outputTokens, ok := metadata["outputTokens"].(float64); ok {
1434 usageInfo.OutputTokens = int64(outputTokens)
1435 log.Debugf("kiro: parseEventStream found outputTokens in messageMetadataEvent: %d", usageInfo.OutputTokens)
1436 }
1437 }
1438 if usageInfo.TotalTokens == 0 {
1439 if totalTokens, ok := metadata["totalTokens"].(float64); ok {
1440 usageInfo.TotalTokens = int64(totalTokens)
1441 log.Debugf("kiro: parseEventStream found totalTokens in messageMetadataEvent: %d", usageInfo.TotalTokens)
1442 }
1443 }
1444
1445 case "usageEvent", "usage":
1446 // Handle dedicated usage events
1447 if inputTokens, ok := event["inputTokens"].(float64); ok {
1448 usageInfo.InputTokens = int64(inputTokens)
1449 log.Debugf("kiro: parseEventStream found inputTokens in usageEvent: %d", usageInfo.InputTokens)
1450 }
1451 if outputTokens, ok := event["outputTokens"].(float64); ok {
1452 usageInfo.OutputTokens = int64(outputTokens)
1453 log.Debugf("kiro: parseEventStream found outputTokens in usageEvent: %d", usageInfo.OutputTokens)
1454 }
1455 if totalTokens, ok := event["totalTokens"].(float64); ok {
1456 usageInfo.TotalTokens = int64(totalTokens)
1457 log.Debugf("kiro: parseEventStream found totalTokens in usageEvent: %d", usageInfo.TotalTokens)
1458 }
1459 // Also check nested usage object
1460 if usageObj, ok := event["usage"].(map[string]interface{}); ok {
1461 if inputTokens, ok := usageObj["input_tokens"].(float64); ok {
1462 usageInfo.InputTokens = int64(inputTokens)
1463 } else if inputTokens, ok := usageObj["prompt_tokens"].(float64); ok {
1464 usageInfo.InputTokens = int64(inputTokens)
1465 }
1466 if outputTokens, ok := usageObj["output_tokens"].(float64); ok {
1467 usageInfo.OutputTokens = int64(outputTokens)
1468 } else if outputTokens, ok := usageObj["completion_tokens"].(float64); ok {
1469 usageInfo.OutputTokens = int64(outputTokens)
1470 }
1471 if totalTokens, ok := usageObj["total_tokens"].(float64); ok {
1472 usageInfo.TotalTokens = int64(totalTokens)
1473 }
1474 log.Debugf("kiro: parseEventStream found usage object: input=%d, output=%d, total=%d",
1475 usageInfo.InputTokens, usageInfo.OutputTokens, usageInfo.TotalTokens)
1476 }
1477
1478 case "metricsEvent":
1479 // Handle metrics events which may contain usage data
1480 if metrics, ok := event["metricsEvent"].(map[string]interface{}); ok {
1481 if inputTokens, ok := metrics["inputTokens"].(float64); ok {
1482 usageInfo.InputTokens = int64(inputTokens)
1483 }
1484 if outputTokens, ok := metrics["outputTokens"].(float64); ok {
1485 usageInfo.OutputTokens = int64(outputTokens)
1486 }
1487 log.Debugf("kiro: parseEventStream found metricsEvent: input=%d, output=%d",
1488 usageInfo.InputTokens, usageInfo.OutputTokens)
1489 }
1490
1491 case "meteringEvent":
1492 // Handle metering events from Kiro API (usage billing information)
1493 // Official format: { unit: string, unitPlural: string, usage: number }
1494 if metering, ok := event["meteringEvent"].(map[string]interface{}); ok {
1495 unit := ""
1496 if u, ok := metering["unit"].(string); ok {
1497 unit = u
1498 }
1499 usageVal := 0.0
1500 if u, ok := metering["usage"].(float64); ok {
1501 usageVal = u
1502 }
1503 log.Infof("kiro: parseEventStream received meteringEvent: usage=%.2f %s", usageVal, unit)
1504 // Store metering info for potential billing/statistics purposes
1505 // Note: This is separate from token counts - it's AWS billing units
1506 } else {
1507 // Try direct fields
1508 unit := ""
1509 if u, ok := event["unit"].(string); ok {
1510 unit = u
1511 }
1512 usageVal := 0.0
1513 if u, ok := event["usage"].(float64); ok {
1514 usageVal = u
1515 }
1516 if unit != "" || usageVal > 0 {
1517 log.Infof("kiro: parseEventStream received meteringEvent (direct): usage=%.2f %s", usageVal, unit)
1518 }
1519 }
1520
1521 case "error", "exception", "internalServerException", "invalidStateEvent":
1522 // Handle error events from Kiro API stream
1523 errMsg := ""
1524 errType := eventType
1525
1526 // Try to extract error message from various formats
1527 if msg, ok := event["message"].(string); ok {
1528 errMsg = msg
1529 } else if errObj, ok := event[eventType].(map[string]interface{}); ok {
1530 if msg, ok := errObj["message"].(string); ok {
1531 errMsg = msg
1532 }
1533 if t, ok := errObj["type"].(string); ok {
1534 errType = t
1535 }
1536 } else if errObj, ok := event["error"].(map[string]interface{}); ok {
1537 if msg, ok := errObj["message"].(string); ok {
1538 errMsg = msg
1539 }
1540 if t, ok := errObj["type"].(string); ok {
1541 errType = t
1542 }
1543 }
1544
1545 // Check for specific error reasons
1546 if reason, ok := event["reason"].(string); ok {
1547 errMsg = fmt.Sprintf("%s (reason: %s)", errMsg, reason)
1548 }
1549
1550 log.Errorf("kiro: parseEventStream received error event: type=%s, message=%s", errType, errMsg)
1551
1552 // For invalidStateEvent, we may want to continue processing other events
1553 if eventType == "invalidStateEvent" {
1554 log.Warnf("kiro: invalidStateEvent received, continuing stream processing")
1555 continue
1556 }
1557
1558 // For other errors, return the error
1559 if errMsg != "" {
1560 return "", nil, usageInfo, stopReason, fmt.Errorf("kiro API error (%s): %s", errType, errMsg)
1561 }
1562
1563 default:
1564 // Check for contextUsagePercentage in any event
1565 if ctxPct, ok := event["contextUsagePercentage"].(float64); ok {
1566 upstreamContextPercentage = ctxPct
1567 log.Debugf("kiro: parseEventStream received context usage: %.2f%%", upstreamContextPercentage)
1568 }
1569 // Log unknown event types for debugging (to discover new event formats)
1570 log.Debugf("kiro: parseEventStream unknown event type: %s, payload: %s", eventType, string(payload))
1571 }
1572
1573 // Check for direct token fields in any event (fallback)
1574 if usageInfo.InputTokens == 0 {
1575 if inputTokens, ok := event["inputTokens"].(float64); ok {
1576 usageInfo.InputTokens = int64(inputTokens)
1577 log.Debugf("kiro: parseEventStream found direct inputTokens: %d", usageInfo.InputTokens)
1578 }
1579 }
1580 if usageInfo.OutputTokens == 0 {
1581 if outputTokens, ok := event["outputTokens"].(float64); ok {
1582 usageInfo.OutputTokens = int64(outputTokens)
1583 log.Debugf("kiro: parseEventStream found direct outputTokens: %d", usageInfo.OutputTokens)
1584 }
1585 }
1586
1587 // Check for usage object in any event (OpenAI format)
1588 if usageInfo.InputTokens == 0 || usageInfo.OutputTokens == 0 {
1589 if usageObj, ok := event["usage"].(map[string]interface{}); ok {
1590 if usageInfo.InputTokens == 0 {
1591 if inputTokens, ok := usageObj["input_tokens"].(float64); ok {
1592 usageInfo.InputTokens = int64(inputTokens)
1593 } else if inputTokens, ok := usageObj["prompt_tokens"].(float64); ok {
1594 usageInfo.InputTokens = int64(inputTokens)
1595 }
1596 }
1597 if usageInfo.OutputTokens == 0 {
1598 if outputTokens, ok := usageObj["output_tokens"].(float64); ok {
1599 usageInfo.OutputTokens = int64(outputTokens)
1600 } else if outputTokens, ok := usageObj["completion_tokens"].(float64); ok {
1601 usageInfo.OutputTokens = int64(outputTokens)
1602 }
1603 }
1604 if usageInfo.TotalTokens == 0 {
1605 if totalTokens, ok := usageObj["total_tokens"].(float64); ok {
1606 usageInfo.TotalTokens = int64(totalTokens)
1607 }
1608 }
1609 log.Debugf("kiro: parseEventStream found usage object (fallback): input=%d, output=%d, total=%d",
1610 usageInfo.InputTokens, usageInfo.OutputTokens, usageInfo.TotalTokens)
1611 }
1612 }
1613
1614 // Also check nested supplementaryWebLinksEvent
1615 if usageEvent, ok := event["supplementaryWebLinksEvent"].(map[string]interface{}); ok {
1616 if inputTokens, ok := usageEvent["inputTokens"].(float64); ok {
1617 usageInfo.InputTokens = int64(inputTokens)
1618 }
1619 if outputTokens, ok := usageEvent["outputTokens"].(float64); ok {
1620 usageInfo.OutputTokens = int64(outputTokens)
1621 }
1622 }
1623 }
1624
1625 // Parse embedded tool calls from content (e.g., [Called tool_name with args: {...}])
1626 contentStr := content.String()
1627 cleanedContent, embeddedToolUses := kiroclaude.ParseEmbeddedToolCalls(contentStr, processedIDs)
1628 toolUses = append(toolUses, embeddedToolUses...)
1629
1630 // Deduplicate all tool uses
1631 toolUses = kiroclaude.DeduplicateToolUses(toolUses)
1632
1633 // Apply fallback logic for stop_reason if not provided by upstream
1634 // Priority: upstream stopReason > tool_use detection > end_turn default
1635 if stopReason == "" {
1636 if len(toolUses) > 0 {
1637 stopReason = "tool_use"
1638 log.Debugf("kiro: parseEventStream using fallback stop_reason: tool_use (detected %d tool uses)", len(toolUses))
1639 } else {
1640 stopReason = "end_turn"
1641 log.Debugf("kiro: parseEventStream using fallback stop_reason: end_turn")
1642 }
1643 }
1644
1645 // Log warning if response was truncated due to max_tokens
1646 if stopReason == "max_tokens" {
1647 log.Warnf("kiro: response truncated due to max_tokens limit")
1648 }
1649
1650 // Use contextUsagePercentage to calculate more accurate input tokens
1651 // Kiro model has 200k max context, contextUsagePercentage represents the percentage used
1652 // Formula: input_tokens = contextUsagePercentage * 200000 / 100
1653 if upstreamContextPercentage > 0 {
1654 calculatedInputTokens := int64(upstreamContextPercentage * 200000 / 100)
1655 if calculatedInputTokens > 0 {
1656 localEstimate := usageInfo.InputTokens
1657 usageInfo.InputTokens = calculatedInputTokens
1658 usageInfo.TotalTokens = usageInfo.InputTokens + usageInfo.OutputTokens
1659 log.Infof("kiro: parseEventStream using contextUsagePercentage (%.2f%%) to calculate input tokens: %d (local estimate was: %d)",
1660 upstreamContextPercentage, calculatedInputTokens, localEstimate)
1661 }
1662 }
1663
1664 return cleanedContent, toolUses, usageInfo, stopReason, nil
1665}
1666
1667// readEventStreamMessage reads and validates a single AWS Event Stream message.
1668// Returns the parsed message or a structured error for different failure modes.
1669// This function implements boundary protection and detailed error classification.
1670//
1671// AWS Event Stream binary format:
1672// - Prelude (12 bytes): total_length (4) + headers_length (4) + prelude_crc (4)
1673// - Headers (variable): header entries
1674// - Payload (variable): JSON data
1675// - Message CRC (4 bytes): CRC32C of entire message (not validated, just skipped)
1676func (e *KiroExecutor) readEventStreamMessage(reader *bufio.Reader) (*eventStreamMessage, *EventStreamError) {
1677 // Read prelude (first 12 bytes: total_len + headers_len + prelude_crc)
1678 prelude := make([]byte, 12)
1679 _, err := io.ReadFull(reader, prelude)
1680 if err == io.EOF {
1681 return nil, nil // Normal end of stream
1682 }
1683 if err != nil {
1684 return nil, &EventStreamError{
1685 Type: ErrStreamFatal,
1686 Message: "failed to read prelude",
1687 Cause: err,
1688 }
1689 }
1690
1691 totalLength := binary.BigEndian.Uint32(prelude[0:4])
1692 headersLength := binary.BigEndian.Uint32(prelude[4:8])
1693 // Note: prelude[8:12] is prelude_crc - we read it but don't validate (no CRC check per requirements)
1694
1695 // Boundary check: minimum frame size
1696 if totalLength < minEventStreamFrameSize {
1697 return nil, &EventStreamError{
1698 Type: ErrStreamMalformed,
1699 Message: fmt.Sprintf("invalid message length: %d (minimum is %d)", totalLength, minEventStreamFrameSize),
1700 }
1701 }
1702
1703 // Boundary check: maximum message size
1704 if totalLength > maxEventStreamMsgSize {
1705 return nil, &EventStreamError{
1706 Type: ErrStreamMalformed,
1707 Message: fmt.Sprintf("message too large: %d bytes (maximum is %d)", totalLength, maxEventStreamMsgSize),
1708 }
1709 }
1710
1711 // Boundary check: headers length within message bounds
1712 // Message structure: prelude(12) + headers(headersLength) + payload + message_crc(4)
1713 // So: headersLength must be <= totalLength - 16 (12 for prelude + 4 for message_crc)
1714 if headersLength > totalLength-16 {
1715 return nil, &EventStreamError{
1716 Type: ErrStreamMalformed,
1717 Message: fmt.Sprintf("headers length %d exceeds message bounds (total: %d)", headersLength, totalLength),
1718 }
1719 }
1720
1721 // Read the rest of the message (total - 12 bytes already read)
1722 remaining := make([]byte, totalLength-12)
1723 _, err = io.ReadFull(reader, remaining)
1724 if err != nil {
1725 return nil, &EventStreamError{
1726 Type: ErrStreamFatal,
1727 Message: "failed to read message body",
1728 Cause: err,
1729 }
1730 }
1731
1732 // Extract event type from headers
1733 // Headers start at beginning of 'remaining', length is headersLength
1734 var eventType string
1735 if headersLength > 0 && headersLength <= uint32(len(remaining)) {
1736 eventType = e.extractEventTypeFromBytes(remaining[:headersLength])
1737 }
1738
1739 // Calculate payload boundaries
1740 // Payload starts after headers, ends before message_crc (last 4 bytes)
1741 payloadStart := headersLength
1742 payloadEnd := uint32(len(remaining)) - 4 // Skip message_crc at end
1743
1744 // Validate payload boundaries
1745 if payloadStart >= payloadEnd {
1746 // No payload, return empty message
1747 return &eventStreamMessage{
1748 EventType: eventType,
1749 Payload: nil,
1750 }, nil
1751 }
1752
1753 payload := remaining[payloadStart:payloadEnd]
1754
1755 return &eventStreamMessage{
1756 EventType: eventType,
1757 Payload: payload,
1758 }, nil
1759}
1760
1761func skipEventStreamHeaderValue(headers []byte, offset int, valueType byte) (int, bool) {
1762 switch valueType {
1763 case 0, 1: // bool true / bool false
1764 return offset, true
1765 case 2: // byte
1766 if offset+1 > len(headers) {
1767 return offset, false
1768 }
1769 return offset + 1, true
1770 case 3: // short
1771 if offset+2 > len(headers) {
1772 return offset, false
1773 }
1774 return offset + 2, true
1775 case 4: // int
1776 if offset+4 > len(headers) {
1777 return offset, false
1778 }
1779 return offset + 4, true
1780 case 5: // long
1781 if offset+8 > len(headers) {
1782 return offset, false
1783 }
1784 return offset + 8, true
1785 case 6: // byte array (2-byte length + data)
1786 if offset+2 > len(headers) {
1787 return offset, false
1788 }
1789 valueLen := int(binary.BigEndian.Uint16(headers[offset : offset+2]))
1790 offset += 2
1791 if offset+valueLen > len(headers) {
1792 return offset, false
1793 }
1794 return offset + valueLen, true
1795 case 8: // timestamp
1796 if offset+8 > len(headers) {
1797 return offset, false
1798 }
1799 return offset + 8, true
1800 case 9: // uuid
1801 if offset+16 > len(headers) {
1802 return offset, false
1803 }
1804 return offset + 16, true
1805 default:
1806 return offset, false
1807 }
1808}
1809
1810// extractEventTypeFromBytes extracts the event type from raw header bytes (without prelude CRC prefix)
1811func (e *KiroExecutor) extractEventTypeFromBytes(headers []byte) string {
1812 offset := 0
1813 for offset < len(headers) {
1814 nameLen := int(headers[offset])
1815 offset++
1816 if offset+nameLen > len(headers) {
1817 break
1818 }
1819 name := string(headers[offset : offset+nameLen])
1820 offset += nameLen
1821
1822 if offset >= len(headers) {
1823 break
1824 }
1825 valueType := headers[offset]
1826 offset++
1827
1828 if valueType == 7 { // String type
1829 if offset+2 > len(headers) {
1830 break
1831 }
1832 valueLen := int(binary.BigEndian.Uint16(headers[offset : offset+2]))
1833 offset += 2
1834 if offset+valueLen > len(headers) {
1835 break
1836 }
1837 value := string(headers[offset : offset+valueLen])
1838 offset += valueLen
1839
1840 if name == ":event-type" {
1841 return value
1842 }
1843 continue
1844 }
1845
1846 nextOffset, ok := skipEventStreamHeaderValue(headers, offset, valueType)
1847 if !ok {
1848 break
1849 }
1850 offset = nextOffset
1851 }
1852 return ""
1853}
1854
1855
1856// NOTE: Response building functions moved to internal/translator/kiro/claude/kiro_claude_response.go
1857// The executor now uses kiroclaude.BuildClaudeResponse() and kiroclaude.ExtractThinkingFromContent() instead
1858
1859// streamToChannel converts AWS Event Stream to channel-based streaming.
1860// Supports tool calling - emits tool_use content blocks when tools are used.
1861// Includes embedded [Called ...] tool call parsing and input buffering for toolUseEvent.
1862// Implements duplicate content filtering using lastContentEvent detection (based on AIClient-2-API).
1863// Extracts stop_reason from upstream events when available.
1864// thinkingEnabled controls whether <thinking> tags are parsed - only parse when request enabled thinking.
1865func (e *KiroExecutor) streamToChannel(ctx context.Context, body io.Reader, out chan<- cliproxyexecutor.StreamChunk, targetFormat sdktranslator.Format, model string, originalReq, claudeBody []byte, reporter *usageReporter, thinkingEnabled bool) {
1866 reader := bufio.NewReaderSize(body, 20*1024*1024) // 20MB buffer to match other providers
1867 var totalUsage usage.Detail
1868 var hasToolUses bool // Track if any tool uses were emitted
1869 var upstreamStopReason string // Track stop_reason from upstream events
1870
1871 // Tool use state tracking for input buffering and deduplication
1872 processedIDs := make(map[string]bool)
1873 var currentToolUse *kiroclaude.ToolUseState
1874
1875 // NOTE: Duplicate content filtering removed - it was causing legitimate repeated
1876 // content (like consecutive newlines) to be incorrectly filtered out.
1877 // The previous implementation compared lastContentEvent == contentDelta which
1878 // is too aggressive for streaming scenarios.
1879
1880 // Streaming token calculation - accumulate content for real-time token counting
1881 // Based on AIClient-2-API implementation
1882 var accumulatedContent strings.Builder
1883 accumulatedContent.Grow(4096) // Pre-allocate 4KB capacity to reduce reallocations
1884
1885 // Real-time usage estimation state
1886 // These track when to send periodic usage updates during streaming
1887 var lastUsageUpdateLen int // Last accumulated content length when usage was sent
1888 var lastUsageUpdateTime = time.Now() // Last time usage update was sent
1889 var lastReportedOutputTokens int64 // Last reported output token count
1890
1891 // Upstream usage tracking - Kiro API returns credit usage and context percentage
1892 var upstreamCreditUsage float64 // Credit usage from upstream (e.g., 1.458)
1893 var upstreamContextPercentage float64 // Context usage percentage from upstream (e.g., 78.56)
1894 var hasUpstreamUsage bool // Whether we received usage from upstream
1895
1896 // Translator param for maintaining tool call state across streaming events
1897 // IMPORTANT: This must persist across all TranslateStream calls
1898 var translatorParam any
1899
1900 // Thinking mode state tracking - tag-based parsing for <thinking> tags in content
1901 inThinkBlock := false // Whether we're currently inside a <thinking> block
1902 isThinkingBlockOpen := false // Track if thinking content block SSE event is open
1903 thinkingBlockIndex := -1 // Index of the thinking content block
1904 var accumulatedThinkingContent strings.Builder // Accumulate thinking content for token counting
1905
1906 // Buffer for handling partial tag matches at chunk boundaries
1907 var pendingContent strings.Builder // Buffer content that might be part of a tag
1908
1909 // Pre-calculate input tokens from request if possible
1910 // Kiro uses Claude format, so try Claude format first, then OpenAI format, then fallback
1911 if enc, err := getTokenizer(model); err == nil {
1912 var inputTokens int64
1913 var countMethod string
1914
1915 // Try Claude format first (Kiro uses Claude API format)
1916 if inp, err := countClaudeChatTokens(enc, claudeBody); err == nil && inp > 0 {
1917 inputTokens = inp
1918 countMethod = "claude"
1919 } else if inp, err := countOpenAIChatTokens(enc, originalReq); err == nil && inp > 0 {
1920 // Fallback to OpenAI format (for OpenAI-compatible requests)
1921 inputTokens = inp
1922 countMethod = "openai"
1923 } else {
1924 // Final fallback: estimate from raw request size (roughly 4 chars per token)
1925 inputTokens = int64(len(claudeBody) / 4)
1926 if inputTokens == 0 && len(claudeBody) > 0 {
1927 inputTokens = 1
1928 }
1929 countMethod = "estimate"
1930 }
1931
1932 totalUsage.InputTokens = inputTokens
1933 log.Debugf("kiro: streamToChannel pre-calculated input tokens: %d (method: %s, claude body: %d bytes, original req: %d bytes)",
1934 totalUsage.InputTokens, countMethod, len(claudeBody), len(originalReq))
1935 }
1936
1937 contentBlockIndex := -1
1938 messageStartSent := false
1939 isTextBlockOpen := false
1940 var outputLen int
1941
1942 // Ensure usage is published even on early return
1943 defer func() {
1944 reporter.publish(ctx, totalUsage)
1945 }()
1946
1947 for {
1948 select {
1949 case <-ctx.Done():
1950 return
1951 default:
1952 }
1953
1954 msg, eventErr := e.readEventStreamMessage(reader)
1955 if eventErr != nil {
1956 // Log the error
1957 log.Errorf("kiro: streamToChannel error: %v", eventErr)
1958
1959 // Send error to channel for client notification
1960 out <- cliproxyexecutor.StreamChunk{Err: eventErr}
1961 return
1962 }
1963 if msg == nil {
1964 // Normal end of stream (EOF)
1965 // Flush any incomplete tool use before ending stream
1966 if currentToolUse != nil && !processedIDs[currentToolUse.ToolUseID] {
1967 log.Warnf("kiro: flushing incomplete tool use at EOF: %s (ID: %s)", currentToolUse.Name, currentToolUse.ToolUseID)
1968 fullInput := currentToolUse.InputBuffer.String()
1969 repairedJSON := kiroclaude.RepairJSON(fullInput)
1970 var finalInput map[string]interface{}
1971 if err := json.Unmarshal([]byte(repairedJSON), &finalInput); err != nil {
1972 log.Warnf("kiro: failed to parse incomplete tool input at EOF: %v", err)
1973 finalInput = make(map[string]interface{})
1974 }
1975
1976 processedIDs[currentToolUse.ToolUseID] = true
1977 contentBlockIndex++
1978
1979 // Send tool_use content block
1980 blockStart := kiroclaude.BuildClaudeContentBlockStartEvent(contentBlockIndex, "tool_use", currentToolUse.ToolUseID, currentToolUse.Name)
1981 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStart, &translatorParam)
1982 for _, chunk := range sseData {
1983 if chunk != "" {
1984 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
1985 }
1986 }
1987
1988 // Send tool input as delta
1989 inputBytes, _ := json.Marshal(finalInput)
1990 inputDelta := kiroclaude.BuildClaudeInputJsonDeltaEvent(string(inputBytes), contentBlockIndex)
1991 sseData = sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, inputDelta, &translatorParam)
1992 for _, chunk := range sseData {
1993 if chunk != "" {
1994 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
1995 }
1996 }
1997
1998 // Close block
1999 blockStop := kiroclaude.BuildClaudeContentBlockStopEvent(contentBlockIndex)
2000 sseData = sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStop, &translatorParam)
2001 for _, chunk := range sseData {
2002 if chunk != "" {
2003 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2004 }
2005 }
2006
2007 hasToolUses = true
2008 currentToolUse = nil
2009 }
2010
2011 // DISABLED: Tag-based pending character flushing
2012 // This code block was used for tag-based thinking detection which has been
2013 // replaced by reasoningContentEvent handling. No pending tag chars to flush.
2014 // Original code preserved in git history.
2015 break
2016 }
2017
2018 eventType := msg.EventType
2019 payload := msg.Payload
2020 if len(payload) == 0 {
2021 continue
2022 }
2023 appendAPIResponseChunk(ctx, e.cfg, payload)
2024
2025 var event map[string]interface{}
2026 if err := json.Unmarshal(payload, &event); err != nil {
2027 log.Warnf("kiro: failed to unmarshal event payload: %v, raw: %s", err, string(payload))
2028 continue
2029 }
2030
2031 // Check for error/exception events in the payload (Kiro API may return errors with HTTP 200)
2032 // These can appear as top-level fields or nested within the event
2033 if errType, hasErrType := event["_type"].(string); hasErrType {
2034 // AWS-style error: {"_type": "com.amazon.aws.codewhisperer#ValidationException", "message": "..."}
2035 errMsg := ""
2036 if msg, ok := event["message"].(string); ok {
2037 errMsg = msg
2038 }
2039 log.Errorf("kiro: received AWS error in stream: type=%s, message=%s", errType, errMsg)
2040 out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("kiro API error: %s - %s", errType, errMsg)}
2041 return
2042 }
2043 if errType, hasErrType := event["type"].(string); hasErrType && (errType == "error" || errType == "exception") {
2044 // Generic error event
2045 errMsg := ""
2046 if msg, ok := event["message"].(string); ok {
2047 errMsg = msg
2048 } else if errObj, ok := event["error"].(map[string]interface{}); ok {
2049 if msg, ok := errObj["message"].(string); ok {
2050 errMsg = msg
2051 }
2052 }
2053 log.Errorf("kiro: received error event in stream: type=%s, message=%s", errType, errMsg)
2054 out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("kiro API error: %s", errMsg)}
2055 return
2056 }
2057
2058 // Extract stop_reason from various event formats (streaming)
2059 // Kiro/Amazon Q API may include stop_reason in different locations
2060 if sr := kirocommon.GetString(event, "stop_reason"); sr != "" {
2061 upstreamStopReason = sr
2062 log.Debugf("kiro: streamToChannel found stop_reason (top-level): %s", upstreamStopReason)
2063 }
2064 if sr := kirocommon.GetString(event, "stopReason"); sr != "" {
2065 upstreamStopReason = sr
2066 log.Debugf("kiro: streamToChannel found stopReason (top-level): %s", upstreamStopReason)
2067 }
2068
2069 // Send message_start on first event
2070 if !messageStartSent {
2071 msgStart := kiroclaude.BuildClaudeMessageStartEvent(model, totalUsage.InputTokens)
2072 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, msgStart, &translatorParam)
2073 for _, chunk := range sseData {
2074 if chunk != "" {
2075 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2076 }
2077 }
2078 messageStartSent = true
2079 }
2080
2081 switch eventType {
2082 case "followupPromptEvent":
2083 // Filter out followupPrompt events - these are UI suggestions, not content
2084 log.Debugf("kiro: streamToChannel ignoring followupPrompt event")
2085 continue
2086
2087 case "messageStopEvent", "message_stop":
2088 // Handle message stop events which may contain stop_reason
2089 if sr := kirocommon.GetString(event, "stop_reason"); sr != "" {
2090 upstreamStopReason = sr
2091 log.Debugf("kiro: streamToChannel found stop_reason in messageStopEvent: %s", upstreamStopReason)
2092 }
2093 if sr := kirocommon.GetString(event, "stopReason"); sr != "" {
2094 upstreamStopReason = sr
2095 log.Debugf("kiro: streamToChannel found stopReason in messageStopEvent: %s", upstreamStopReason)
2096 }
2097
2098 case "meteringEvent":
2099 // Handle metering events from Kiro API (usage billing information)
2100 // Official format: { unit: string, unitPlural: string, usage: number }
2101 if metering, ok := event["meteringEvent"].(map[string]interface{}); ok {
2102 unit := ""
2103 if u, ok := metering["unit"].(string); ok {
2104 unit = u
2105 }
2106 usageVal := 0.0
2107 if u, ok := metering["usage"].(float64); ok {
2108 usageVal = u
2109 }
2110 upstreamCreditUsage = usageVal
2111 hasUpstreamUsage = true
2112 log.Infof("kiro: streamToChannel received meteringEvent: usage=%.4f %s", usageVal, unit)
2113 } else {
2114 // Try direct fields (event is meteringEvent itself)
2115 if unit, ok := event["unit"].(string); ok {
2116 if usage, ok := event["usage"].(float64); ok {
2117 upstreamCreditUsage = usage
2118 hasUpstreamUsage = true
2119 log.Infof("kiro: streamToChannel received meteringEvent (direct): usage=%.4f %s", usage, unit)
2120 }
2121 }
2122 }
2123
2124 case "error", "exception", "internalServerException":
2125 // Handle error events from Kiro API stream
2126 errMsg := ""
2127 errType := eventType
2128
2129 // Try to extract error message from various formats
2130 if msg, ok := event["message"].(string); ok {
2131 errMsg = msg
2132 } else if errObj, ok := event[eventType].(map[string]interface{}); ok {
2133 if msg, ok := errObj["message"].(string); ok {
2134 errMsg = msg
2135 }
2136 if t, ok := errObj["type"].(string); ok {
2137 errType = t
2138 }
2139 } else if errObj, ok := event["error"].(map[string]interface{}); ok {
2140 if msg, ok := errObj["message"].(string); ok {
2141 errMsg = msg
2142 }
2143 }
2144
2145 log.Errorf("kiro: streamToChannel received error event: type=%s, message=%s", errType, errMsg)
2146
2147 // Send error to the stream and exit
2148 if errMsg != "" {
2149 out <- cliproxyexecutor.StreamChunk{
2150 Err: fmt.Errorf("kiro API error (%s): %s", errType, errMsg),
2151 }
2152 return
2153 }
2154
2155 case "invalidStateEvent":
2156 // Handle invalid state events - log and continue (non-fatal)
2157 errMsg := ""
2158 if msg, ok := event["message"].(string); ok {
2159 errMsg = msg
2160 } else if stateEvent, ok := event["invalidStateEvent"].(map[string]interface{}); ok {
2161 if msg, ok := stateEvent["message"].(string); ok {
2162 errMsg = msg
2163 }
2164 }
2165 log.Warnf("kiro: streamToChannel received invalidStateEvent: %s, continuing", errMsg)
2166 continue
2167
2168 default:
2169 // Check for upstream usage events from Kiro API
2170 // Format: {"unit":"credit","unitPlural":"credits","usage":1.458}
2171 if unit, ok := event["unit"].(string); ok && unit == "credit" {
2172 if usage, ok := event["usage"].(float64); ok {
2173 upstreamCreditUsage = usage
2174 hasUpstreamUsage = true
2175 log.Debugf("kiro: received upstream credit usage: %.4f", upstreamCreditUsage)
2176 }
2177 }
2178 // Format: {"contextUsagePercentage":78.56}
2179 if ctxPct, ok := event["contextUsagePercentage"].(float64); ok {
2180 upstreamContextPercentage = ctxPct
2181 log.Debugf("kiro: received upstream context usage: %.2f%%", upstreamContextPercentage)
2182 }
2183
2184 // Check for token counts in unknown events
2185 if inputTokens, ok := event["inputTokens"].(float64); ok {
2186 totalUsage.InputTokens = int64(inputTokens)
2187 hasUpstreamUsage = true
2188 log.Debugf("kiro: streamToChannel found inputTokens in event %s: %d", eventType, totalUsage.InputTokens)
2189 }
2190 if outputTokens, ok := event["outputTokens"].(float64); ok {
2191 totalUsage.OutputTokens = int64(outputTokens)
2192 hasUpstreamUsage = true
2193 log.Debugf("kiro: streamToChannel found outputTokens in event %s: %d", eventType, totalUsage.OutputTokens)
2194 }
2195 if totalTokens, ok := event["totalTokens"].(float64); ok {
2196 totalUsage.TotalTokens = int64(totalTokens)
2197 log.Debugf("kiro: streamToChannel found totalTokens in event %s: %d", eventType, totalUsage.TotalTokens)
2198 }
2199
2200 // Check for usage object in unknown events (OpenAI/Claude format)
2201 if usageObj, ok := event["usage"].(map[string]interface{}); ok {
2202 if inputTokens, ok := usageObj["input_tokens"].(float64); ok {
2203 totalUsage.InputTokens = int64(inputTokens)
2204 hasUpstreamUsage = true
2205 } else if inputTokens, ok := usageObj["prompt_tokens"].(float64); ok {
2206 totalUsage.InputTokens = int64(inputTokens)
2207 hasUpstreamUsage = true
2208 }
2209 if outputTokens, ok := usageObj["output_tokens"].(float64); ok {
2210 totalUsage.OutputTokens = int64(outputTokens)
2211 hasUpstreamUsage = true
2212 } else if outputTokens, ok := usageObj["completion_tokens"].(float64); ok {
2213 totalUsage.OutputTokens = int64(outputTokens)
2214 hasUpstreamUsage = true
2215 }
2216 if totalTokens, ok := usageObj["total_tokens"].(float64); ok {
2217 totalUsage.TotalTokens = int64(totalTokens)
2218 }
2219 log.Debugf("kiro: streamToChannel found usage object in event %s: input=%d, output=%d, total=%d",
2220 eventType, totalUsage.InputTokens, totalUsage.OutputTokens, totalUsage.TotalTokens)
2221 }
2222
2223 // Log unknown event types for debugging (to discover new event formats)
2224 if eventType != "" {
2225 log.Debugf("kiro: streamToChannel unknown event type: %s, payload: %s", eventType, string(payload))
2226 }
2227
2228 case "assistantResponseEvent":
2229 var contentDelta string
2230 var toolUses []map[string]interface{}
2231
2232 if assistantResp, ok := event["assistantResponseEvent"].(map[string]interface{}); ok {
2233 if c, ok := assistantResp["content"].(string); ok {
2234 contentDelta = c
2235 }
2236 // Extract stop_reason from assistantResponseEvent
2237 if sr := kirocommon.GetString(assistantResp, "stop_reason"); sr != "" {
2238 upstreamStopReason = sr
2239 log.Debugf("kiro: streamToChannel found stop_reason in assistantResponseEvent: %s", upstreamStopReason)
2240 }
2241 if sr := kirocommon.GetString(assistantResp, "stopReason"); sr != "" {
2242 upstreamStopReason = sr
2243 log.Debugf("kiro: streamToChannel found stopReason in assistantResponseEvent: %s", upstreamStopReason)
2244 }
2245 // Extract tool uses from response
2246 if tus, ok := assistantResp["toolUses"].([]interface{}); ok {
2247 for _, tuRaw := range tus {
2248 if tu, ok := tuRaw.(map[string]interface{}); ok {
2249 toolUses = append(toolUses, tu)
2250 }
2251 }
2252 }
2253 }
2254 if contentDelta == "" {
2255 if c, ok := event["content"].(string); ok {
2256 contentDelta = c
2257 }
2258 }
2259 // Direct tool uses
2260 if tus, ok := event["toolUses"].([]interface{}); ok {
2261 for _, tuRaw := range tus {
2262 if tu, ok := tuRaw.(map[string]interface{}); ok {
2263 toolUses = append(toolUses, tu)
2264 }
2265 }
2266 }
2267
2268 // Handle text content with thinking mode support
2269 if contentDelta != "" {
2270 // NOTE: Duplicate content filtering was removed because it incorrectly
2271 // filtered out legitimate repeated content (like consecutive newlines "\n\n").
2272 // Streaming naturally can have identical chunks that are valid content.
2273
2274 outputLen += len(contentDelta)
2275 // Accumulate content for streaming token calculation
2276 accumulatedContent.WriteString(contentDelta)
2277
2278 // Real-time usage estimation: Check if we should send a usage update
2279 // This helps clients track context usage during long thinking sessions
2280 shouldSendUsageUpdate := false
2281 if accumulatedContent.Len()-lastUsageUpdateLen >= usageUpdateCharThreshold {
2282 shouldSendUsageUpdate = true
2283 } else if time.Since(lastUsageUpdateTime) >= usageUpdateTimeInterval && accumulatedContent.Len() > lastUsageUpdateLen {
2284 shouldSendUsageUpdate = true
2285 }
2286
2287 if shouldSendUsageUpdate {
2288 // Calculate current output tokens using tiktoken
2289 var currentOutputTokens int64
2290 if enc, encErr := getTokenizer(model); encErr == nil {
2291 if tokenCount, countErr := enc.Count(accumulatedContent.String()); countErr == nil {
2292 currentOutputTokens = int64(tokenCount)
2293 }
2294 }
2295 // Fallback to character estimation if tiktoken fails
2296 if currentOutputTokens == 0 {
2297 currentOutputTokens = int64(accumulatedContent.Len() / 4)
2298 if currentOutputTokens == 0 {
2299 currentOutputTokens = 1
2300 }
2301 }
2302
2303 // Only send update if token count has changed significantly (at least 10 tokens)
2304 if currentOutputTokens > lastReportedOutputTokens+10 {
2305 // Send ping event with usage information
2306 // This is a non-blocking update that clients can optionally process
2307 pingEvent := kiroclaude.BuildClaudePingEventWithUsage(totalUsage.InputTokens, currentOutputTokens)
2308 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, pingEvent, &translatorParam)
2309 for _, chunk := range sseData {
2310 if chunk != "" {
2311 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2312 }
2313 }
2314
2315 lastReportedOutputTokens = currentOutputTokens
2316 log.Debugf("kiro: sent real-time usage update - input: %d, output: %d (accumulated: %d chars)",
2317 totalUsage.InputTokens, currentOutputTokens, accumulatedContent.Len())
2318 }
2319
2320 lastUsageUpdateLen = accumulatedContent.Len()
2321 lastUsageUpdateTime = time.Now()
2322 }
2323
2324 // TAG-BASED THINKING PARSING: Parse <thinking> tags from content
2325 // Combine pending content with new content for processing
2326 pendingContent.WriteString(contentDelta)
2327 processContent := pendingContent.String()
2328 pendingContent.Reset()
2329
2330 // Process content looking for thinking tags
2331 for len(processContent) > 0 {
2332 if inThinkBlock {
2333 // We're inside a thinking block, look for </thinking>
2334 endIdx := strings.Index(processContent, kirocommon.ThinkingEndTag)
2335 if endIdx >= 0 {
2336 // Found end tag - emit thinking content before the tag
2337 thinkingText := processContent[:endIdx]
2338 if thinkingText != "" {
2339 // Ensure thinking block is open
2340 if !isThinkingBlockOpen {
2341 contentBlockIndex++
2342 thinkingBlockIndex = contentBlockIndex
2343 isThinkingBlockOpen = true
2344 blockStart := kiroclaude.BuildClaudeContentBlockStartEvent(thinkingBlockIndex, "thinking", "", "")
2345 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStart, &translatorParam)
2346 for _, chunk := range sseData {
2347 if chunk != "" {
2348 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2349 }
2350 }
2351 }
2352 // Send thinking delta
2353 thinkingEvent := kiroclaude.BuildClaudeThinkingDeltaEvent(thinkingText, thinkingBlockIndex)
2354 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, thinkingEvent, &translatorParam)
2355 for _, chunk := range sseData {
2356 if chunk != "" {
2357 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2358 }
2359 }
2360 accumulatedThinkingContent.WriteString(thinkingText)
2361 }
2362 // Close thinking block
2363 if isThinkingBlockOpen {
2364 blockStop := kiroclaude.BuildClaudeThinkingBlockStopEvent(thinkingBlockIndex)
2365 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStop, &translatorParam)
2366 for _, chunk := range sseData {
2367 if chunk != "" {
2368 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2369 }
2370 }
2371 isThinkingBlockOpen = false
2372 }
2373 inThinkBlock = false
2374 processContent = processContent[endIdx+len(kirocommon.ThinkingEndTag):]
2375 log.Debugf("kiro: closed thinking block, remaining content: %d chars", len(processContent))
2376 } else {
2377 // No end tag found - check for partial match at end
2378 partialMatch := false
2379 for i := 1; i < len(kirocommon.ThinkingEndTag) && i <= len(processContent); i++ {
2380 if strings.HasSuffix(processContent, kirocommon.ThinkingEndTag[:i]) {
2381 // Possible partial tag at end, buffer it
2382 pendingContent.WriteString(processContent[len(processContent)-i:])
2383 processContent = processContent[:len(processContent)-i]
2384 partialMatch = true
2385 break
2386 }
2387 }
2388 if !partialMatch || len(processContent) > 0 {
2389 // Emit all as thinking content
2390 if processContent != "" {
2391 if !isThinkingBlockOpen {
2392 contentBlockIndex++
2393 thinkingBlockIndex = contentBlockIndex
2394 isThinkingBlockOpen = true
2395 blockStart := kiroclaude.BuildClaudeContentBlockStartEvent(thinkingBlockIndex, "thinking", "", "")
2396 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStart, &translatorParam)
2397 for _, chunk := range sseData {
2398 if chunk != "" {
2399 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2400 }
2401 }
2402 }
2403 thinkingEvent := kiroclaude.BuildClaudeThinkingDeltaEvent(processContent, thinkingBlockIndex)
2404 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, thinkingEvent, &translatorParam)
2405 for _, chunk := range sseData {
2406 if chunk != "" {
2407 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2408 }
2409 }
2410 accumulatedThinkingContent.WriteString(processContent)
2411 }
2412 }
2413 processContent = ""
2414 }
2415 } else {
2416 // Not in thinking block, look for <thinking>
2417 startIdx := strings.Index(processContent, kirocommon.ThinkingStartTag)
2418 if startIdx >= 0 {
2419 // Found start tag - emit text content before the tag
2420 textBefore := processContent[:startIdx]
2421 if textBefore != "" {
2422 // Close thinking block if open
2423 if isThinkingBlockOpen {
2424 blockStop := kiroclaude.BuildClaudeThinkingBlockStopEvent(thinkingBlockIndex)
2425 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStop, &translatorParam)
2426 for _, chunk := range sseData {
2427 if chunk != "" {
2428 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2429 }
2430 }
2431 isThinkingBlockOpen = false
2432 }
2433 // Ensure text block is open
2434 if !isTextBlockOpen {
2435 contentBlockIndex++
2436 isTextBlockOpen = true
2437 blockStart := kiroclaude.BuildClaudeContentBlockStartEvent(contentBlockIndex, "text", "", "")
2438 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStart, &translatorParam)
2439 for _, chunk := range sseData {
2440 if chunk != "" {
2441 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2442 }
2443 }
2444 }
2445 // Send text delta
2446 claudeEvent := kiroclaude.BuildClaudeStreamEvent(textBefore, contentBlockIndex)
2447 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, claudeEvent, &translatorParam)
2448 for _, chunk := range sseData {
2449 if chunk != "" {
2450 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2451 }
2452 }
2453 }
2454 // Close text block before entering thinking
2455 if isTextBlockOpen {
2456 blockStop := kiroclaude.BuildClaudeContentBlockStopEvent(contentBlockIndex)
2457 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStop, &translatorParam)
2458 for _, chunk := range sseData {
2459 if chunk != "" {
2460 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2461 }
2462 }
2463 isTextBlockOpen = false
2464 }
2465 inThinkBlock = true
2466 processContent = processContent[startIdx+len(kirocommon.ThinkingStartTag):]
2467 log.Debugf("kiro: entered thinking block")
2468 } else {
2469 // No start tag found - check for partial match at end
2470 partialMatch := false
2471 for i := 1; i < len(kirocommon.ThinkingStartTag) && i <= len(processContent); i++ {
2472 if strings.HasSuffix(processContent, kirocommon.ThinkingStartTag[:i]) {
2473 // Possible partial tag at end, buffer it
2474 pendingContent.WriteString(processContent[len(processContent)-i:])
2475 processContent = processContent[:len(processContent)-i]
2476 partialMatch = true
2477 break
2478 }
2479 }
2480 if !partialMatch || len(processContent) > 0 {
2481 // Emit all as text content
2482 if processContent != "" {
2483 if !isTextBlockOpen {
2484 contentBlockIndex++
2485 isTextBlockOpen = true
2486 blockStart := kiroclaude.BuildClaudeContentBlockStartEvent(contentBlockIndex, "text", "", "")
2487 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStart, &translatorParam)
2488 for _, chunk := range sseData {
2489 if chunk != "" {
2490 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2491 }
2492 }
2493 }
2494 claudeEvent := kiroclaude.BuildClaudeStreamEvent(processContent, contentBlockIndex)
2495 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, claudeEvent, &translatorParam)
2496 for _, chunk := range sseData {
2497 if chunk != "" {
2498 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2499 }
2500 }
2501 }
2502 }
2503 processContent = ""
2504 }
2505 }
2506 }
2507 }
2508
2509 // Handle tool uses in response (with deduplication)
2510 for _, tu := range toolUses {
2511 toolUseID := kirocommon.GetString(tu, "toolUseId")
2512 toolName := kirocommon.GetString(tu, "name")
2513
2514 // Check for duplicate
2515 if processedIDs[toolUseID] {
2516 log.Debugf("kiro: skipping duplicate tool use in stream: %s", toolUseID)
2517 continue
2518 }
2519 processedIDs[toolUseID] = true
2520
2521 hasToolUses = true
2522 // Close text block if open before starting tool_use block
2523 if isTextBlockOpen && contentBlockIndex >= 0 {
2524 blockStop := kiroclaude.BuildClaudeContentBlockStopEvent(contentBlockIndex)
2525 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStop, &translatorParam)
2526 for _, chunk := range sseData {
2527 if chunk != "" {
2528 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2529 }
2530 }
2531 isTextBlockOpen = false
2532 }
2533
2534 // Emit tool_use content block
2535 contentBlockIndex++
2536
2537 blockStart := kiroclaude.BuildClaudeContentBlockStartEvent(contentBlockIndex, "tool_use", toolUseID, toolName)
2538 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStart, &translatorParam)
2539 for _, chunk := range sseData {
2540 if chunk != "" {
2541 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2542 }
2543 }
2544
2545 // Send input_json_delta with the tool input
2546 if input, ok := tu["input"].(map[string]interface{}); ok {
2547 inputJSON, err := json.Marshal(input)
2548 if err != nil {
2549 log.Debugf("kiro: failed to marshal tool input: %v", err)
2550 // Don't continue - still need to close the block
2551 } else {
2552 inputDelta := kiroclaude.BuildClaudeInputJsonDeltaEvent(string(inputJSON), contentBlockIndex)
2553 sseData = sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, inputDelta, &translatorParam)
2554 for _, chunk := range sseData {
2555 if chunk != "" {
2556 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2557 }
2558 }
2559 }
2560 }
2561
2562 // Close tool_use block (always close even if input marshal failed)
2563 blockStop := kiroclaude.BuildClaudeContentBlockStopEvent(contentBlockIndex)
2564 sseData = sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStop, &translatorParam)
2565 for _, chunk := range sseData {
2566 if chunk != "" {
2567 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2568 }
2569 }
2570 }
2571
2572 case "reasoningContentEvent":
2573 // Handle official reasoningContentEvent from Kiro API
2574 // This replaces tag-based thinking detection with the proper event type
2575 // Official format: { text: string, signature?: string, redactedContent?: base64 }
2576 var thinkingText string
2577 var signature string
2578
2579 if re, ok := event["reasoningContentEvent"].(map[string]interface{}); ok {
2580 if text, ok := re["text"].(string); ok {
2581 thinkingText = text
2582 }
2583 if sig, ok := re["signature"].(string); ok {
2584 signature = sig
2585 if len(sig) > 20 {
2586 log.Debugf("kiro: reasoningContentEvent has signature: %s...", sig[:20])
2587 } else {
2588 log.Debugf("kiro: reasoningContentEvent has signature: %s", sig)
2589 }
2590 }
2591 } else {
2592 // Try direct fields
2593 if text, ok := event["text"].(string); ok {
2594 thinkingText = text
2595 }
2596 if sig, ok := event["signature"].(string); ok {
2597 signature = sig
2598 }
2599 }
2600
2601 if thinkingText != "" {
2602 // Close text block if open before starting thinking block
2603 if isTextBlockOpen && contentBlockIndex >= 0 {
2604 blockStop := kiroclaude.BuildClaudeContentBlockStopEvent(contentBlockIndex)
2605 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStop, &translatorParam)
2606 for _, chunk := range sseData {
2607 if chunk != "" {
2608 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2609 }
2610 }
2611 isTextBlockOpen = false
2612 }
2613
2614 // Start thinking block if not already open
2615 if !isThinkingBlockOpen {
2616 contentBlockIndex++
2617 thinkingBlockIndex = contentBlockIndex
2618 isThinkingBlockOpen = true
2619 blockStart := kiroclaude.BuildClaudeContentBlockStartEvent(thinkingBlockIndex, "thinking", "", "")
2620 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStart, &translatorParam)
2621 for _, chunk := range sseData {
2622 if chunk != "" {
2623 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2624 }
2625 }
2626 }
2627
2628 // Send thinking content
2629 thinkingEvent := kiroclaude.BuildClaudeThinkingDeltaEvent(thinkingText, thinkingBlockIndex)
2630 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, thinkingEvent, &translatorParam)
2631 for _, chunk := range sseData {
2632 if chunk != "" {
2633 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2634 }
2635 }
2636
2637 // Accumulate for token counting
2638 accumulatedThinkingContent.WriteString(thinkingText)
2639 log.Debugf("kiro: received reasoningContentEvent, text length: %d, has signature: %v", len(thinkingText), signature != "")
2640 }
2641
2642 // Note: We don't close the thinking block here - it will be closed when we see
2643 // the next assistantResponseEvent or at the end of the stream
2644 _ = signature // Signature can be used for verification if needed
2645
2646 case "toolUseEvent":
2647 // Handle dedicated tool use events with input buffering
2648 completedToolUses, newState := kiroclaude.ProcessToolUseEvent(event, currentToolUse, processedIDs)
2649 currentToolUse = newState
2650
2651 // Emit completed tool uses
2652 for _, tu := range completedToolUses {
2653 hasToolUses = true
2654
2655 // Close text block if open
2656 if isTextBlockOpen && contentBlockIndex >= 0 {
2657 blockStop := kiroclaude.BuildClaudeContentBlockStopEvent(contentBlockIndex)
2658 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStop, &translatorParam)
2659 for _, chunk := range sseData {
2660 if chunk != "" {
2661 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2662 }
2663 }
2664 isTextBlockOpen = false
2665 }
2666
2667 contentBlockIndex++
2668
2669 blockStart := kiroclaude.BuildClaudeContentBlockStartEvent(contentBlockIndex, "tool_use", tu.ToolUseID, tu.Name)
2670 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStart, &translatorParam)
2671 for _, chunk := range sseData {
2672 if chunk != "" {
2673 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2674 }
2675 }
2676
2677 if tu.Input != nil {
2678 inputJSON, err := json.Marshal(tu.Input)
2679 if err != nil {
2680 log.Debugf("kiro: failed to marshal tool input in toolUseEvent: %v", err)
2681 } else {
2682 inputDelta := kiroclaude.BuildClaudeInputJsonDeltaEvent(string(inputJSON), contentBlockIndex)
2683 sseData = sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, inputDelta, &translatorParam)
2684 for _, chunk := range sseData {
2685 if chunk != "" {
2686 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2687 }
2688 }
2689 }
2690 }
2691
2692 blockStop := kiroclaude.BuildClaudeContentBlockStopEvent(contentBlockIndex)
2693 sseData = sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStop, &translatorParam)
2694 for _, chunk := range sseData {
2695 if chunk != "" {
2696 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2697 }
2698 }
2699 }
2700
2701 case "supplementaryWebLinksEvent":
2702 if inputTokens, ok := event["inputTokens"].(float64); ok {
2703 totalUsage.InputTokens = int64(inputTokens)
2704 }
2705 if outputTokens, ok := event["outputTokens"].(float64); ok {
2706 totalUsage.OutputTokens = int64(outputTokens)
2707 }
2708
2709 case "messageMetadataEvent", "metadataEvent":
2710 // Handle message metadata events which contain token counts
2711 // Official format: { tokenUsage: { outputTokens, totalTokens, uncachedInputTokens, cacheReadInputTokens, cacheWriteInputTokens, contextUsagePercentage } }
2712 var metadata map[string]interface{}
2713 if m, ok := event["messageMetadataEvent"].(map[string]interface{}); ok {
2714 metadata = m
2715 } else if m, ok := event["metadataEvent"].(map[string]interface{}); ok {
2716 metadata = m
2717 } else {
2718 metadata = event // event itself might be the metadata
2719 }
2720
2721 // Check for nested tokenUsage object (official format)
2722 if tokenUsage, ok := metadata["tokenUsage"].(map[string]interface{}); ok {
2723 // outputTokens - precise output token count
2724 if outputTokens, ok := tokenUsage["outputTokens"].(float64); ok {
2725 totalUsage.OutputTokens = int64(outputTokens)
2726 hasUpstreamUsage = true
2727 log.Infof("kiro: streamToChannel found precise outputTokens in tokenUsage: %d", totalUsage.OutputTokens)
2728 }
2729 // totalTokens - precise total token count
2730 if totalTokens, ok := tokenUsage["totalTokens"].(float64); ok {
2731 totalUsage.TotalTokens = int64(totalTokens)
2732 log.Infof("kiro: streamToChannel found precise totalTokens in tokenUsage: %d", totalUsage.TotalTokens)
2733 }
2734 // uncachedInputTokens - input tokens not from cache
2735 if uncachedInputTokens, ok := tokenUsage["uncachedInputTokens"].(float64); ok {
2736 totalUsage.InputTokens = int64(uncachedInputTokens)
2737 hasUpstreamUsage = true
2738 log.Infof("kiro: streamToChannel found uncachedInputTokens in tokenUsage: %d", totalUsage.InputTokens)
2739 }
2740 // cacheReadInputTokens - tokens read from cache
2741 if cacheReadTokens, ok := tokenUsage["cacheReadInputTokens"].(float64); ok {
2742 // Add to input tokens if we have uncached tokens, otherwise use as input
2743 if totalUsage.InputTokens > 0 {
2744 totalUsage.InputTokens += int64(cacheReadTokens)
2745 } else {
2746 totalUsage.InputTokens = int64(cacheReadTokens)
2747 }
2748 hasUpstreamUsage = true
2749 log.Debugf("kiro: streamToChannel found cacheReadInputTokens in tokenUsage: %d", int64(cacheReadTokens))
2750 }
2751 // contextUsagePercentage - can be used as fallback for input token estimation
2752 if ctxPct, ok := tokenUsage["contextUsagePercentage"].(float64); ok {
2753 upstreamContextPercentage = ctxPct
2754 log.Debugf("kiro: streamToChannel found contextUsagePercentage in tokenUsage: %.2f%%", ctxPct)
2755 }
2756 }
2757
2758 // Fallback: check for direct fields in metadata (legacy format)
2759 if totalUsage.InputTokens == 0 {
2760 if inputTokens, ok := metadata["inputTokens"].(float64); ok {
2761 totalUsage.InputTokens = int64(inputTokens)
2762 hasUpstreamUsage = true
2763 log.Debugf("kiro: streamToChannel found inputTokens in messageMetadataEvent: %d", totalUsage.InputTokens)
2764 }
2765 }
2766 if totalUsage.OutputTokens == 0 {
2767 if outputTokens, ok := metadata["outputTokens"].(float64); ok {
2768 totalUsage.OutputTokens = int64(outputTokens)
2769 hasUpstreamUsage = true
2770 log.Debugf("kiro: streamToChannel found outputTokens in messageMetadataEvent: %d", totalUsage.OutputTokens)
2771 }
2772 }
2773 if totalUsage.TotalTokens == 0 {
2774 if totalTokens, ok := metadata["totalTokens"].(float64); ok {
2775 totalUsage.TotalTokens = int64(totalTokens)
2776 log.Debugf("kiro: streamToChannel found totalTokens in messageMetadataEvent: %d", totalUsage.TotalTokens)
2777 }
2778 }
2779
2780 case "usageEvent", "usage":
2781 // Handle dedicated usage events
2782 if inputTokens, ok := event["inputTokens"].(float64); ok {
2783 totalUsage.InputTokens = int64(inputTokens)
2784 log.Debugf("kiro: streamToChannel found inputTokens in usageEvent: %d", totalUsage.InputTokens)
2785 }
2786 if outputTokens, ok := event["outputTokens"].(float64); ok {
2787 totalUsage.OutputTokens = int64(outputTokens)
2788 log.Debugf("kiro: streamToChannel found outputTokens in usageEvent: %d", totalUsage.OutputTokens)
2789 }
2790 if totalTokens, ok := event["totalTokens"].(float64); ok {
2791 totalUsage.TotalTokens = int64(totalTokens)
2792 log.Debugf("kiro: streamToChannel found totalTokens in usageEvent: %d", totalUsage.TotalTokens)
2793 }
2794 // Also check nested usage object
2795 if usageObj, ok := event["usage"].(map[string]interface{}); ok {
2796 if inputTokens, ok := usageObj["input_tokens"].(float64); ok {
2797 totalUsage.InputTokens = int64(inputTokens)
2798 } else if inputTokens, ok := usageObj["prompt_tokens"].(float64); ok {
2799 totalUsage.InputTokens = int64(inputTokens)
2800 }
2801 if outputTokens, ok := usageObj["output_tokens"].(float64); ok {
2802 totalUsage.OutputTokens = int64(outputTokens)
2803 } else if outputTokens, ok := usageObj["completion_tokens"].(float64); ok {
2804 totalUsage.OutputTokens = int64(outputTokens)
2805 }
2806 if totalTokens, ok := usageObj["total_tokens"].(float64); ok {
2807 totalUsage.TotalTokens = int64(totalTokens)
2808 }
2809 log.Debugf("kiro: streamToChannel found usage object: input=%d, output=%d, total=%d",
2810 totalUsage.InputTokens, totalUsage.OutputTokens, totalUsage.TotalTokens)
2811 }
2812
2813 case "metricsEvent":
2814 // Handle metrics events which may contain usage data
2815 if metrics, ok := event["metricsEvent"].(map[string]interface{}); ok {
2816 if inputTokens, ok := metrics["inputTokens"].(float64); ok {
2817 totalUsage.InputTokens = int64(inputTokens)
2818 }
2819 if outputTokens, ok := metrics["outputTokens"].(float64); ok {
2820 totalUsage.OutputTokens = int64(outputTokens)
2821 }
2822 log.Debugf("kiro: streamToChannel found metricsEvent: input=%d, output=%d",
2823 totalUsage.InputTokens, totalUsage.OutputTokens)
2824 }
2825 }
2826
2827 // Check nested usage event
2828 if usageEvent, ok := event["supplementaryWebLinksEvent"].(map[string]interface{}); ok {
2829 if inputTokens, ok := usageEvent["inputTokens"].(float64); ok {
2830 totalUsage.InputTokens = int64(inputTokens)
2831 }
2832 if outputTokens, ok := usageEvent["outputTokens"].(float64); ok {
2833 totalUsage.OutputTokens = int64(outputTokens)
2834 }
2835 }
2836
2837 // Check for direct token fields in any event (fallback)
2838 if totalUsage.InputTokens == 0 {
2839 if inputTokens, ok := event["inputTokens"].(float64); ok {
2840 totalUsage.InputTokens = int64(inputTokens)
2841 log.Debugf("kiro: streamToChannel found direct inputTokens: %d", totalUsage.InputTokens)
2842 }
2843 }
2844 if totalUsage.OutputTokens == 0 {
2845 if outputTokens, ok := event["outputTokens"].(float64); ok {
2846 totalUsage.OutputTokens = int64(outputTokens)
2847 log.Debugf("kiro: streamToChannel found direct outputTokens: %d", totalUsage.OutputTokens)
2848 }
2849 }
2850
2851 // Check for usage object in any event (OpenAI format)
2852 if totalUsage.InputTokens == 0 || totalUsage.OutputTokens == 0 {
2853 if usageObj, ok := event["usage"].(map[string]interface{}); ok {
2854 if totalUsage.InputTokens == 0 {
2855 if inputTokens, ok := usageObj["input_tokens"].(float64); ok {
2856 totalUsage.InputTokens = int64(inputTokens)
2857 } else if inputTokens, ok := usageObj["prompt_tokens"].(float64); ok {
2858 totalUsage.InputTokens = int64(inputTokens)
2859 }
2860 }
2861 if totalUsage.OutputTokens == 0 {
2862 if outputTokens, ok := usageObj["output_tokens"].(float64); ok {
2863 totalUsage.OutputTokens = int64(outputTokens)
2864 } else if outputTokens, ok := usageObj["completion_tokens"].(float64); ok {
2865 totalUsage.OutputTokens = int64(outputTokens)
2866 }
2867 }
2868 if totalUsage.TotalTokens == 0 {
2869 if totalTokens, ok := usageObj["total_tokens"].(float64); ok {
2870 totalUsage.TotalTokens = int64(totalTokens)
2871 }
2872 }
2873 log.Debugf("kiro: streamToChannel found usage object (fallback): input=%d, output=%d, total=%d",
2874 totalUsage.InputTokens, totalUsage.OutputTokens, totalUsage.TotalTokens)
2875 }
2876 }
2877 }
2878
2879 // Close content block if open
2880 if isTextBlockOpen && contentBlockIndex >= 0 {
2881 blockStop := kiroclaude.BuildClaudeContentBlockStopEvent(contentBlockIndex)
2882 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, blockStop, &translatorParam)
2883 for _, chunk := range sseData {
2884 if chunk != "" {
2885 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2886 }
2887 }
2888 }
2889
2890 // Streaming token calculation - calculate output tokens from accumulated content
2891 // Only use local estimation if server didn't provide usage (server-side usage takes priority)
2892 if totalUsage.OutputTokens == 0 && accumulatedContent.Len() > 0 {
2893 // Try to use tiktoken for accurate counting
2894 if enc, err := getTokenizer(model); err == nil {
2895 if tokenCount, countErr := enc.Count(accumulatedContent.String()); countErr == nil {
2896 totalUsage.OutputTokens = int64(tokenCount)
2897 log.Debugf("kiro: streamToChannel calculated output tokens using tiktoken: %d", totalUsage.OutputTokens)
2898 } else {
2899 // Fallback on count error: estimate from character count
2900 totalUsage.OutputTokens = int64(accumulatedContent.Len() / 4)
2901 if totalUsage.OutputTokens == 0 {
2902 totalUsage.OutputTokens = 1
2903 }
2904 log.Debugf("kiro: streamToChannel tiktoken count failed, estimated from chars: %d", totalUsage.OutputTokens)
2905 }
2906 } else {
2907 // Fallback: estimate from character count (roughly 4 chars per token)
2908 totalUsage.OutputTokens = int64(accumulatedContent.Len() / 4)
2909 if totalUsage.OutputTokens == 0 {
2910 totalUsage.OutputTokens = 1
2911 }
2912 log.Debugf("kiro: streamToChannel estimated output tokens from chars: %d (content len: %d)", totalUsage.OutputTokens, accumulatedContent.Len())
2913 }
2914 } else if totalUsage.OutputTokens == 0 && outputLen > 0 {
2915 // Legacy fallback using outputLen
2916 totalUsage.OutputTokens = int64(outputLen / 4)
2917 if totalUsage.OutputTokens == 0 {
2918 totalUsage.OutputTokens = 1
2919 }
2920 }
2921
2922 // Use contextUsagePercentage to calculate more accurate input tokens
2923 // Kiro model has 200k max context, contextUsagePercentage represents the percentage used
2924 // Formula: input_tokens = contextUsagePercentage * 200000 / 100
2925 // Note: The effective input context is ~170k (200k - 30k reserved for output)
2926 if upstreamContextPercentage > 0 {
2927 // Calculate input tokens from context percentage
2928 // Using 200k as the base since that's what Kiro reports against
2929 calculatedInputTokens := int64(upstreamContextPercentage * 200000 / 100)
2930
2931 // Only use calculated value if it's significantly different from local estimate
2932 // This provides more accurate token counts based on upstream data
2933 if calculatedInputTokens > 0 {
2934 localEstimate := totalUsage.InputTokens
2935 totalUsage.InputTokens = calculatedInputTokens
2936 log.Debugf("kiro: using contextUsagePercentage (%.2f%%) to calculate input tokens: %d (local estimate was: %d)",
2937 upstreamContextPercentage, calculatedInputTokens, localEstimate)
2938 }
2939 }
2940
2941 totalUsage.TotalTokens = totalUsage.InputTokens + totalUsage.OutputTokens
2942
2943 // Log upstream usage information if received
2944 if hasUpstreamUsage {
2945 log.Debugf("kiro: upstream usage - credits: %.4f, context: %.2f%%, final tokens - input: %d, output: %d, total: %d",
2946 upstreamCreditUsage, upstreamContextPercentage,
2947 totalUsage.InputTokens, totalUsage.OutputTokens, totalUsage.TotalTokens)
2948 }
2949
2950 // Determine stop reason: prefer upstream, then detect tool_use, default to end_turn
2951 stopReason := upstreamStopReason
2952 if stopReason == "" {
2953 if hasToolUses {
2954 stopReason = "tool_use"
2955 log.Debugf("kiro: streamToChannel using fallback stop_reason: tool_use")
2956 } else {
2957 stopReason = "end_turn"
2958 log.Debugf("kiro: streamToChannel using fallback stop_reason: end_turn")
2959 }
2960 }
2961
2962 // Log warning if response was truncated due to max_tokens
2963 if stopReason == "max_tokens" {
2964 log.Warnf("kiro: response truncated due to max_tokens limit (streamToChannel)")
2965 }
2966
2967 // Send message_delta event
2968 msgDelta := kiroclaude.BuildClaudeMessageDeltaEvent(stopReason, totalUsage)
2969 sseData := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, msgDelta, &translatorParam)
2970 for _, chunk := range sseData {
2971 if chunk != "" {
2972 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2973 }
2974 }
2975
2976 // Send message_stop event separately
2977 msgStop := kiroclaude.BuildClaudeMessageStopOnlyEvent()
2978 sseData = sdktranslator.TranslateStream(ctx, sdktranslator.FromString("kiro"), targetFormat, model, originalReq, claudeBody, msgStop, &translatorParam)
2979 for _, chunk := range sseData {
2980 if chunk != "" {
2981 out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunk + "\n\n")}
2982 }
2983 }
2984 // reporter.publish is called via defer
2985}
2986
2987// NOTE: Claude SSE event builders moved to internal/translator/kiro/claude/kiro_claude_stream.go
2988// The executor now uses kiroclaude.BuildClaude*Event() functions instead
2989
2990// CountTokens counts tokens locally using tiktoken since Kiro API doesn't expose a token counting endpoint.
2991// This provides approximate token counts for client requests.
2992func (e *KiroExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
2993 // Use tiktoken for local token counting
2994 enc, err := getTokenizer(req.Model)
2995 if err != nil {
2996 log.Warnf("kiro: CountTokens failed to get tokenizer: %v, falling back to estimate", err)
2997 // Fallback: estimate from payload size (roughly 4 chars per token)
2998 estimatedTokens := len(req.Payload) / 4
2999 if estimatedTokens == 0 && len(req.Payload) > 0 {
3000 estimatedTokens = 1
3001 }
3002 return cliproxyexecutor.Response{
3003 Payload: []byte(fmt.Sprintf(`{"count":%d}`, estimatedTokens)),
3004 }, nil
3005 }
3006
3007 // Try to count tokens from the request payload
3008 var totalTokens int64
3009
3010 // Try OpenAI chat format first
3011 if tokens, countErr := countOpenAIChatTokens(enc, req.Payload); countErr == nil && tokens > 0 {
3012 totalTokens = tokens
3013 log.Debugf("kiro: CountTokens counted %d tokens using OpenAI chat format", totalTokens)
3014 } else {
3015 // Fallback: count raw payload tokens
3016 if tokenCount, countErr := enc.Count(string(req.Payload)); countErr == nil {
3017 totalTokens = int64(tokenCount)
3018 log.Debugf("kiro: CountTokens counted %d tokens from raw payload", totalTokens)
3019 } else {
3020 // Final fallback: estimate from payload size
3021 totalTokens = int64(len(req.Payload) / 4)
3022 if totalTokens == 0 && len(req.Payload) > 0 {
3023 totalTokens = 1
3024 }
3025 log.Debugf("kiro: CountTokens estimated %d tokens from payload size", totalTokens)
3026 }
3027 }
3028
3029 return cliproxyexecutor.Response{
3030 Payload: []byte(fmt.Sprintf(`{"count":%d}`, totalTokens)),
3031 }, nil
3032}
3033
3034// Refresh refreshes the Kiro OAuth token.
3035// Supports both AWS Builder ID (SSO OIDC) and Google OAuth (social login).
3036// Uses mutex to prevent race conditions when multiple concurrent requests try to refresh.
3037func (e *KiroExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
3038 // Serialize token refresh operations to prevent race conditions
3039 e.refreshMu.Lock()
3040 defer e.refreshMu.Unlock()
3041
3042 var authID string
3043 if auth != nil {
3044 authID = auth.ID
3045 } else {
3046 authID = "<nil>"
3047 }
3048 log.Debugf("kiro executor: refresh called for auth %s", authID)
3049 if auth == nil {
3050 return nil, fmt.Errorf("kiro executor: auth is nil")
3051 }
3052
3053 // Double-check: After acquiring lock, verify token still needs refresh
3054 // Another goroutine may have already refreshed while we were waiting
3055 // NOTE: This check has a design limitation - it reads from the auth object passed in,
3056 // not from persistent storage. If another goroutine returns a new Auth object (via Clone),
3057 // this check won't see those updates. The mutex still prevents truly concurrent refreshes,
3058 // but queued goroutines may still attempt redundant refreshes. This is acceptable as
3059 // the refresh operation is idempotent and the extra API calls are infrequent.
3060 if auth.Metadata != nil {
3061 if lastRefresh, ok := auth.Metadata["last_refresh"].(string); ok {
3062 if refreshTime, err := time.Parse(time.RFC3339, lastRefresh); err == nil {
3063 // If token was refreshed within the last 30 seconds, skip refresh
3064 if time.Since(refreshTime) < 30*time.Second {
3065 log.Debugf("kiro executor: token was recently refreshed by another goroutine, skipping")
3066 return auth, nil
3067 }
3068 }
3069 }
3070 // Also check if expires_at is now in the future with sufficient buffer
3071 if expiresAt, ok := auth.Metadata["expires_at"].(string); ok {
3072 if expTime, err := time.Parse(time.RFC3339, expiresAt); err == nil {
3073 // If token expires more than 5 minutes from now, it's still valid
3074 if time.Until(expTime) > 5*time.Minute {
3075 log.Debugf("kiro executor: token is still valid (expires in %v), skipping refresh", time.Until(expTime))
3076 // CRITICAL FIX: Set NextRefreshAfter to prevent frequent refresh checks
3077 // Without this, shouldRefresh() will return true again in 5 seconds
3078 updated := auth.Clone()
3079 // Set next refresh to 5 minutes before expiry, or at least 30 seconds from now
3080 nextRefresh := expTime.Add(-5 * time.Minute)
3081 minNextRefresh := time.Now().Add(30 * time.Second)
3082 if nextRefresh.Before(minNextRefresh) {
3083 nextRefresh = minNextRefresh
3084 }
3085 updated.NextRefreshAfter = nextRefresh
3086 log.Debugf("kiro executor: setting NextRefreshAfter to %v (in %v)", nextRefresh.Format(time.RFC3339), time.Until(nextRefresh))
3087 return updated, nil
3088 }
3089 }
3090 }
3091 }
3092
3093 var refreshToken string
3094 var clientID, clientSecret string
3095 var authMethod string
3096 var region, startURL string
3097
3098 if auth.Metadata != nil {
3099 if rt, ok := auth.Metadata["refresh_token"].(string); ok {
3100 refreshToken = rt
3101 }
3102 if cid, ok := auth.Metadata["client_id"].(string); ok {
3103 clientID = cid
3104 }
3105 if cs, ok := auth.Metadata["client_secret"].(string); ok {
3106 clientSecret = cs
3107 }
3108 if am, ok := auth.Metadata["auth_method"].(string); ok {
3109 authMethod = am
3110 }
3111 if r, ok := auth.Metadata["region"].(string); ok {
3112 region = r
3113 }
3114 if su, ok := auth.Metadata["start_url"].(string); ok {
3115 startURL = su
3116 }
3117 }
3118
3119 if refreshToken == "" {
3120 return nil, fmt.Errorf("kiro executor: refresh token not found")
3121 }
3122
3123 var tokenData *kiroauth.KiroTokenData
3124 var err error
3125
3126 ssoClient := kiroauth.NewSSOOIDCClient(e.cfg)
3127
3128 // Use SSO OIDC refresh for AWS Builder ID or IDC, otherwise use Kiro's OAuth refresh endpoint
3129 switch {
3130 case clientID != "" && clientSecret != "" && authMethod == "idc" && region != "":
3131 // IDC refresh with region-specific endpoint
3132 log.Debugf("kiro executor: using SSO OIDC refresh for IDC (region=%s)", region)
3133 tokenData, err = ssoClient.RefreshTokenWithRegion(ctx, clientID, clientSecret, refreshToken, region, startURL)
3134 case clientID != "" && clientSecret != "" && authMethod == "builder-id":
3135 // Builder ID refresh with default endpoint
3136 log.Debugf("kiro executor: using SSO OIDC refresh for AWS Builder ID")
3137 tokenData, err = ssoClient.RefreshToken(ctx, clientID, clientSecret, refreshToken)
3138 default:
3139 // Fallback to Kiro's OAuth refresh endpoint (for social auth: Google/GitHub)
3140 log.Debugf("kiro executor: using Kiro OAuth refresh endpoint")
3141 oauth := kiroauth.NewKiroOAuth(e.cfg)
3142 tokenData, err = oauth.RefreshToken(ctx, refreshToken)
3143 }
3144
3145 if err != nil {
3146 return nil, fmt.Errorf("kiro executor: token refresh failed: %w", err)
3147 }
3148
3149 updated := auth.Clone()
3150 now := time.Now()
3151 updated.UpdatedAt = now
3152 updated.LastRefreshedAt = now
3153
3154 if updated.Metadata == nil {
3155 updated.Metadata = make(map[string]any)
3156 }
3157 updated.Metadata["access_token"] = tokenData.AccessToken
3158 updated.Metadata["refresh_token"] = tokenData.RefreshToken
3159 updated.Metadata["expires_at"] = tokenData.ExpiresAt
3160 updated.Metadata["last_refresh"] = now.Format(time.RFC3339)
3161 if tokenData.ProfileArn != "" {
3162 updated.Metadata["profile_arn"] = tokenData.ProfileArn
3163 }
3164 if tokenData.AuthMethod != "" {
3165 updated.Metadata["auth_method"] = tokenData.AuthMethod
3166 }
3167 if tokenData.Provider != "" {
3168 updated.Metadata["provider"] = tokenData.Provider
3169 }
3170 // Preserve client credentials for future refreshes (AWS Builder ID)
3171 if tokenData.ClientID != "" {
3172 updated.Metadata["client_id"] = tokenData.ClientID
3173 }
3174 if tokenData.ClientSecret != "" {
3175 updated.Metadata["client_secret"] = tokenData.ClientSecret
3176 }
3177
3178 if updated.Attributes == nil {
3179 updated.Attributes = make(map[string]string)
3180 }
3181 updated.Attributes["access_token"] = tokenData.AccessToken
3182 if tokenData.ProfileArn != "" {
3183 updated.Attributes["profile_arn"] = tokenData.ProfileArn
3184 }
3185
3186 // NextRefreshAfter is aligned with RefreshLead (5min)
3187 if expiresAt, parseErr := time.Parse(time.RFC3339, tokenData.ExpiresAt); parseErr == nil {
3188 updated.NextRefreshAfter = expiresAt.Add(-5 * time.Minute)
3189 }
3190
3191 log.Infof("kiro executor: token refreshed successfully, expires at %s", tokenData.ExpiresAt)
3192 return updated, nil
3193}
3194
3195// persistRefreshedAuth persists a refreshed auth record to disk.
3196// This ensures token refreshes from inline retry are saved to the auth file.
3197func (e *KiroExecutor) persistRefreshedAuth(auth *cliproxyauth.Auth) error {
3198 if auth == nil || auth.Metadata == nil {
3199 return fmt.Errorf("kiro executor: cannot persist nil auth or metadata")
3200 }
3201
3202 // Determine the file path from auth attributes or filename
3203 var authPath string
3204 if auth.Attributes != nil {
3205 if p := strings.TrimSpace(auth.Attributes["path"]); p != "" {
3206 authPath = p
3207 }
3208 }
3209 if authPath == "" {
3210 fileName := strings.TrimSpace(auth.FileName)
3211 if fileName == "" {
3212 return fmt.Errorf("kiro executor: auth has no file path or filename")
3213 }
3214 if filepath.IsAbs(fileName) {
3215 authPath = fileName
3216 } else if e.cfg != nil && e.cfg.AuthDir != "" {
3217 authPath = filepath.Join(e.cfg.AuthDir, fileName)
3218 } else {
3219 return fmt.Errorf("kiro executor: cannot determine auth file path")
3220 }
3221 }
3222
3223 // Marshal metadata to JSON
3224 raw, err := json.Marshal(auth.Metadata)
3225 if err != nil {
3226 return fmt.Errorf("kiro executor: marshal metadata failed: %w", err)
3227 }
3228
3229 // Write to temp file first, then rename (atomic write)
3230 tmp := authPath + ".tmp"
3231 if err := os.WriteFile(tmp, raw, 0o600); err != nil {
3232 return fmt.Errorf("kiro executor: write temp auth file failed: %w", err)
3233 }
3234 if err := os.Rename(tmp, authPath); err != nil {
3235 return fmt.Errorf("kiro executor: rename auth file failed: %w", err)
3236 }
3237
3238 log.Debugf("kiro executor: persisted refreshed auth to %s", authPath)
3239 return nil
3240}
3241
3242// isTokenExpired checks if a JWT access token has expired.
3243// Returns true if the token is expired or cannot be parsed.
3244func (e *KiroExecutor) isTokenExpired(accessToken string) bool {
3245 if accessToken == "" {
3246 return true
3247 }
3248
3249 // JWT tokens have 3 parts separated by dots
3250 parts := strings.Split(accessToken, ".")
3251 if len(parts) != 3 {
3252 // Not a JWT token, assume not expired
3253 return false
3254 }
3255
3256 // Decode the payload (second part)
3257 // JWT uses base64url encoding without padding (RawURLEncoding)
3258 payload := parts[1]
3259 decoded, err := base64.RawURLEncoding.DecodeString(payload)
3260 if err != nil {
3261 // Try with padding added as fallback
3262 switch len(payload) % 4 {
3263 case 2:
3264 payload += "=="
3265 case 3:
3266 payload += "="
3267 }
3268 decoded, err = base64.URLEncoding.DecodeString(payload)
3269 if err != nil {
3270 log.Debugf("kiro: failed to decode JWT payload: %v", err)
3271 return false
3272 }
3273 }
3274
3275 var claims struct {
3276 Exp int64 `json:"exp"`
3277 }
3278 if err := json.Unmarshal(decoded, &claims); err != nil {
3279 log.Debugf("kiro: failed to parse JWT claims: %v", err)
3280 return false
3281 }
3282
3283 if claims.Exp == 0 {
3284 // No expiration claim, assume not expired
3285 return false
3286 }
3287
3288 expTime := time.Unix(claims.Exp, 0)
3289 now := time.Now()
3290
3291 // Consider token expired if it expires within 1 minute (buffer for clock skew)
3292 isExpired := now.After(expTime) || expTime.Sub(now) < time.Minute
3293 if isExpired {
3294 log.Debugf("kiro: token expired at %s (now: %s)", expTime.Format(time.RFC3339), now.Format(time.RFC3339))
3295 }
3296
3297 return isExpired
3298}
3299
3300// NOTE: Message merging functions moved to internal/translator/kiro/common/message_merge.go
3301// NOTE: Tool calling support functions moved to internal/translator/kiro/claude/kiro_claude_tools.go
3302// The executor now uses kiroclaude.* and kirocommon.* functions instead
internal/runtime/executor/proxy_helpers.go+44−5
@@ -6,6 +6,7 @@ import ( · 5 unmodified lines
6 "net/http"6 "net/http"
7 "net/url"7 "net/url"
8 "strings"8 "strings"
9 "sync"
9 "time"10 "time"
10 11
11 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"12 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
@@ -14,11 +15,19 @@ import ( · 2 unmodified lines
14 "golang.org/x/net/proxy"15 "golang.org/x/net/proxy"
15)16)
16 17
18// httpClientCache caches HTTP clients by proxy URL to enable connection reuse
19var (
20 httpClientCache = make(map[string]*http.Client)
21 httpClientCacheMutex sync.RWMutex
22)
23
17// newProxyAwareHTTPClient creates an HTTP client with proper proxy configuration priority:24// newProxyAwareHTTPClient creates an HTTP client with proper proxy configuration priority:
18// 1. Use auth.ProxyURL if configured (highest priority)25// 1. Use auth.ProxyURL if configured (highest priority)
19// 2. Use cfg.ProxyURL if auth proxy is not configured26// 2. Use cfg.ProxyURL if auth proxy is not configured
20// 3. Use RoundTripper from context if neither are configured27// 3. Use RoundTripper from context if neither are configured
21//28//
29// This function caches HTTP clients by proxy URL to enable TCP/TLS connection reuse.
30//
22// Parameters:31// Parameters:
23// - ctx: The context containing optional RoundTripper32// - ctx: The context containing optional RoundTripper
24// - cfg: The application configuration33// - cfg: The application configuration
@@ -28,11 +37,6 @@ import ( · 3 unmodified lines
28// Returns:37// Returns:
29// - *http.Client: An HTTP client with configured proxy or transport38// - *http.Client: An HTTP client with configured proxy or transport
30func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client {39func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client {
31 httpClient := &http.Client{}
32 if timeout > 0 {
33 httpClient.Timeout = timeout
34 }
35
36 // Priority 1: Use auth.ProxyURL if configured40 // Priority 1: Use auth.ProxyURL if configured
37 var proxyURL string41 var proxyURL string
38 if auth != nil {42 if auth != nil {
@@ -44,11 +48,39 @@ func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *clip · 5 unmodified lines
44 proxyURL = strings.TrimSpace(cfg.ProxyURL)48 proxyURL = strings.TrimSpace(cfg.ProxyURL)
45 }49 }
46 50
51 // Build cache key from proxy URL (empty string for no proxy)
52 cacheKey := proxyURL
53
54 // Check cache first
55 httpClientCacheMutex.RLock()
56 if cachedClient, ok := httpClientCache[cacheKey]; ok {
57 httpClientCacheMutex.RUnlock()
58 // Return a wrapper with the requested timeout but shared transport
59 if timeout > 0 {
60 return &http.Client{
61 Transport: cachedClient.Transport,
62 Timeout: timeout,
63 }
64 }
65 return cachedClient
66 }
67 httpClientCacheMutex.RUnlock()
68
69 // Create new client
70 httpClient := &http.Client{}
71 if timeout > 0 {
72 httpClient.Timeout = timeout
73 }
74
47 // If we have a proxy URL configured, set up the transport75 // If we have a proxy URL configured, set up the transport
48 if proxyURL != "" {76 if proxyURL != "" {
49 transport := buildProxyTransport(proxyURL)77 transport := buildProxyTransport(proxyURL)
50 if transport != nil {78 if transport != nil {
51 httpClient.Transport = transport79 httpClient.Transport = transport
80 // Cache the client
81 httpClientCacheMutex.Lock()
82 httpClientCache[cacheKey] = httpClient
83 httpClientCacheMutex.Unlock()
52 return httpClient84 return httpClient
53 }85 }
54 // If proxy setup failed, log and fall through to context RoundTripper86 // If proxy setup failed, log and fall through to context RoundTripper
@@ -60,6 +92,13 @@ func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *clip · 5 unmodified lines
60 httpClient.Transport = rt92 httpClient.Transport = rt
61 }93 }
62 94
95 // Cache the client for no-proxy case
96 if proxyURL == "" {
97 httpClientCacheMutex.Lock()
98 httpClientCache[cacheKey] = httpClient
99 httpClientCacheMutex.Unlock()
100 }
101
63 return httpClient102 return httpClient
64}103}
65 104
internal/runtime/executor/token_helpers.go+276−15
@@ -2,43 +2,109 @@ package executor · 1 unmodified line
2 2
3import (3import (
4 "fmt"4 "fmt"
5 "regexp"
6 "strconv"
5 "strings"7 "strings"
8 "sync"
6 9
7 "github.com/tidwall/gjson"10 "github.com/tidwall/gjson"
8 "github.com/tiktoken-go/tokenizer"11 "github.com/tiktoken-go/tokenizer"
9)12)
10 13
14// tokenizerCache stores tokenizer instances to avoid repeated creation
15var tokenizerCache sync.Map
16
17// TokenizerWrapper wraps a tokenizer codec with an adjustment factor for models
18// where tiktoken may not accurately estimate token counts (e.g., Claude models)
19type TokenizerWrapper struct {
20 Codec tokenizer.Codec
21 AdjustmentFactor float64 // 1.0 means no adjustment, >1.0 means tiktoken underestimates
22}
23
24// Count returns the token count with adjustment factor applied
25func (tw *TokenizerWrapper) Count(text string) (int, error) {
26 count, err := tw.Codec.Count(text)
27 if err != nil {
28 return 0, err
29 }
30 if tw.AdjustmentFactor != 1.0 && tw.AdjustmentFactor > 0 {
31 return int(float64(count) * tw.AdjustmentFactor), nil
32 }
33 return count, nil
34}
35
36// getTokenizer returns a cached tokenizer for the given model.
37// This improves performance by avoiding repeated tokenizer creation.
38func getTokenizer(model string) (*TokenizerWrapper, error) {
39 // Check cache first
40 if cached, ok := tokenizerCache.Load(model); ok {
41 return cached.(*TokenizerWrapper), nil
42 }
43
44 // Cache miss, create new tokenizer
45 wrapper, err := tokenizerForModel(model)
46 if err != nil {
47 return nil, err
48 }
49
50 // Store in cache (use LoadOrStore to handle race conditions)
51 actual, _ := tokenizerCache.LoadOrStore(model, wrapper)
52 return actual.(*TokenizerWrapper), nil
53}
54
11// tokenizerForModel returns a tokenizer codec suitable for an OpenAI-style model id.55// tokenizerForModel returns a tokenizer codec suitable for an OpenAI-style model id.
12func tokenizerForModel(model string) (tokenizer.Codec, error) {56// For Claude models, applies a 1.1 adjustment factor since tiktoken may underestimate.
57func tokenizerForModel(model string) (*TokenizerWrapper, error) {
13 sanitized := strings.ToLower(strings.TrimSpace(model))58 sanitized := strings.ToLower(strings.TrimSpace(model))
59
60 // Claude models use cl100k_base with 1.1 adjustment factor
61 // because tiktoken may underestimate Claude's actual token count
62 if strings.Contains(sanitized, "claude") || strings.HasPrefix(sanitized, "kiro-") || strings.HasPrefix(sanitized, "amazonq-") {
63 enc, err := tokenizer.Get(tokenizer.Cl100kBase)
64 if err != nil {
65 return nil, err
66 }
67 return &TokenizerWrapper{Codec: enc, AdjustmentFactor: 1.1}, nil
68 }
69
70 var enc tokenizer.Codec
71 var err error
72
14 switch {73 switch {
15 case sanitized == "":74 case sanitized == "":
16 return tokenizer.Get(tokenizer.Cl100kBase)75 enc, err = tokenizer.Get(tokenizer.Cl100kBase)
17 case strings.HasPrefix(sanitized, "gpt-5"):76 case strings.HasPrefix(sanitized, "gpt-5.2"):
18 return tokenizer.ForModel(tokenizer.GPT5)77 enc, err = tokenizer.ForModel(tokenizer.GPT5)
19 case strings.HasPrefix(sanitized, "gpt-5.1"):78 case strings.HasPrefix(sanitized, "gpt-5.1"):
20 return tokenizer.ForModel(tokenizer.GPT5)79 enc, err = tokenizer.ForModel(tokenizer.GPT5)
80 case strings.HasPrefix(sanitized, "gpt-5"):
81 enc, err = tokenizer.ForModel(tokenizer.GPT5)
21 case strings.HasPrefix(sanitized, "gpt-4.1"):82 case strings.HasPrefix(sanitized, "gpt-4.1"):
22 return tokenizer.ForModel(tokenizer.GPT41)83 enc, err = tokenizer.ForModel(tokenizer.GPT41)
23 case strings.HasPrefix(sanitized, "gpt-4o"):84 case strings.HasPrefix(sanitized, "gpt-4o"):
24 return tokenizer.ForModel(tokenizer.GPT4o)85 enc, err = tokenizer.ForModel(tokenizer.GPT4o)
25 case strings.HasPrefix(sanitized, "gpt-4"):86 case strings.HasPrefix(sanitized, "gpt-4"):
26 return tokenizer.ForModel(tokenizer.GPT4)87 enc, err = tokenizer.ForModel(tokenizer.GPT4)
27 case strings.HasPrefix(sanitized, "gpt-3.5"), strings.HasPrefix(sanitized, "gpt-3"):88 case strings.HasPrefix(sanitized, "gpt-3.5"), strings.HasPrefix(sanitized, "gpt-3"):
28 return tokenizer.ForModel(tokenizer.GPT35Turbo)89 enc, err = tokenizer.ForModel(tokenizer.GPT35Turbo)
29 case strings.HasPrefix(sanitized, "o1"):90 case strings.HasPrefix(sanitized, "o1"):
30 return tokenizer.ForModel(tokenizer.O1)91 enc, err = tokenizer.ForModel(tokenizer.O1)
31 case strings.HasPrefix(sanitized, "o3"):92 case strings.HasPrefix(sanitized, "o3"):
32 return tokenizer.ForModel(tokenizer.O3)93 enc, err = tokenizer.ForModel(tokenizer.O3)
33 case strings.HasPrefix(sanitized, "o4"):94 case strings.HasPrefix(sanitized, "o4"):
34 return tokenizer.ForModel(tokenizer.O4Mini)95 enc, err = tokenizer.ForModel(tokenizer.O4Mini)
35 default:96 default:
36 return tokenizer.Get(tokenizer.O200kBase)97 enc, err = tokenizer.Get(tokenizer.O200kBase)
37 }98 }
99
100 if err != nil {
101 return nil, err
102 }
103 return &TokenizerWrapper{Codec: enc, AdjustmentFactor: 1.0}, nil
38}104}
39 105
40// countOpenAIChatTokens approximates prompt tokens for OpenAI chat completions payloads.106// countOpenAIChatTokens approximates prompt tokens for OpenAI chat completions payloads.
41func countOpenAIChatTokens(enc tokenizer.Codec, payload []byte) (int64, error) {107func countOpenAIChatTokens(enc *TokenizerWrapper, payload []byte) (int64, error) {
42 if enc == nil {108 if enc == nil {
43 return 0, fmt.Errorf("encoder is nil")109 return 0, fmt.Errorf("encoder is nil")
44 }110 }
@@ -62,11 +128,206 @@ func countOpenAIChatTokens(enc tokenizer.Codec, payload []byte) (int64, error) { · 17 unmodified lines
62 return 0, nil128 return 0, nil
63 }129 }
64 130
131 // Count text tokens
65 count, err := enc.Count(joined)132 count, err := enc.Count(joined)
66 if err != nil {133 if err != nil {
67 return 0, err134 return 0, err
68 }135 }
69 return int64(count), nil136
137 // Extract and add image tokens from placeholders
138 imageTokens := extractImageTokens(joined)
139
140 return int64(count) + int64(imageTokens), nil
141}
142
143// countClaudeChatTokens approximates prompt tokens for Claude API chat completions payloads.
144// This handles Claude's message format with system, messages, and tools.
145// Image tokens are estimated based on image dimensions when available.
146func countClaudeChatTokens(enc *TokenizerWrapper, payload []byte) (int64, error) {
147 if enc == nil {
148 return 0, fmt.Errorf("encoder is nil")
149 }
150 if len(payload) == 0 {
151 return 0, nil
152 }
153
154 root := gjson.ParseBytes(payload)
155 segments := make([]string, 0, 32)
156
157 // Collect system prompt (can be string or array of content blocks)
158 collectClaudeSystem(root.Get("system"), &segments)
159
160 // Collect messages
161 collectClaudeMessages(root.Get("messages"), &segments)
162
163 // Collect tools
164 collectClaudeTools(root.Get("tools"), &segments)
165
166 joined := strings.TrimSpace(strings.Join(segments, "\n"))
167 if joined == "" {
168 return 0, nil
169 }
170
171 // Count text tokens
172 count, err := enc.Count(joined)
173 if err != nil {
174 return 0, err
175 }
176
177 // Extract and add image tokens from placeholders
178 imageTokens := extractImageTokens(joined)
179
180 return int64(count) + int64(imageTokens), nil
181}
182
183// imageTokenPattern matches [IMAGE:xxx tokens] format for extracting estimated image tokens
184var imageTokenPattern = regexp.MustCompile(`\[IMAGE:(\d+) tokens\]`)
185
186// extractImageTokens extracts image token estimates from placeholder text.
187// Placeholders are in the format [IMAGE:xxx tokens] where xxx is the estimated token count.
188func extractImageTokens(text string) int {
189 matches := imageTokenPattern.FindAllStringSubmatch(text, -1)
190 total := 0
191 for _, match := range matches {
192 if len(match) > 1 {
193 if tokens, err := strconv.Atoi(match[1]); err == nil {
194 total += tokens
195 }
196 }
197 }
198 return total
199}
200
201// estimateImageTokens calculates estimated tokens for an image based on dimensions.
202// Based on Claude's image token calculation: tokens ≈ (width * height) / 750
203// Minimum 85 tokens, maximum 1590 tokens (for 1568x1568 images).
204func estimateImageTokens(width, height float64) int {
205 if width <= 0 || height <= 0 {
206 // No valid dimensions, use default estimate (medium-sized image)
207 return 1000
208 }
209
210 tokens := int(width * height / 750)
211
212 // Apply bounds
213 if tokens < 85 {
214 tokens = 85
215 }
216 if tokens > 1590 {
217 tokens = 1590
218 }
219
220 return tokens
221}
222
223// collectClaudeSystem extracts text from Claude's system field.
224// System can be a string or an array of content blocks.
225func collectClaudeSystem(system gjson.Result, segments *[]string) {
226 if !system.Exists() {
227 return
228 }
229 if system.Type == gjson.String {
230 addIfNotEmpty(segments, system.String())
231 return
232 }
233 if system.IsArray() {
234 system.ForEach(func(_, block gjson.Result) bool {
235 blockType := block.Get("type").String()
236 if blockType == "text" || blockType == "" {
237 addIfNotEmpty(segments, block.Get("text").String())
238 }
239 // Also handle plain string blocks
240 if block.Type == gjson.String {
241 addIfNotEmpty(segments, block.String())
242 }
243 return true
244 })
245 }
246}
247
248// collectClaudeMessages extracts text from Claude's messages array.
249func collectClaudeMessages(messages gjson.Result, segments *[]string) {
250 if !messages.Exists() || !messages.IsArray() {
251 return
252 }
253 messages.ForEach(func(_, message gjson.Result) bool {
254 addIfNotEmpty(segments, message.Get("role").String())
255 collectClaudeContent(message.Get("content"), segments)
256 return true
257 })
258}
259
260// collectClaudeContent extracts text from Claude's content field.
261// Content can be a string or an array of content blocks.
262// For images, estimates token count based on dimensions when available.
263func collectClaudeContent(content gjson.Result, segments *[]string) {
264 if !content.Exists() {
265 return
266 }
267 if content.Type == gjson.String {
268 addIfNotEmpty(segments, content.String())
269 return
270 }
271 if content.IsArray() {
272 content.ForEach(func(_, part gjson.Result) bool {
273 partType := part.Get("type").String()
274 switch partType {
275 case "text":
276 addIfNotEmpty(segments, part.Get("text").String())
277 case "image":
278 // Estimate image tokens based on dimensions if available
279 source := part.Get("source")
280 if source.Exists() {
281 width := source.Get("width").Float()
282 height := source.Get("height").Float()
283 if width > 0 && height > 0 {
284 tokens := estimateImageTokens(width, height)
285 addIfNotEmpty(segments, fmt.Sprintf("[IMAGE:%d tokens]", tokens))
286 } else {
287 // No dimensions available, use default estimate
288 addIfNotEmpty(segments, "[IMAGE:1000 tokens]")
289 }
290 } else {
291 // No source info, use default estimate
292 addIfNotEmpty(segments, "[IMAGE:1000 tokens]")
293 }
294 case "tool_use":
295 addIfNotEmpty(segments, part.Get("id").String())
296 addIfNotEmpty(segments, part.Get("name").String())
297 if input := part.Get("input"); input.Exists() {
298 addIfNotEmpty(segments, input.Raw)
299 }
300 case "tool_result":
301 addIfNotEmpty(segments, part.Get("tool_use_id").String())
302 collectClaudeContent(part.Get("content"), segments)
303 case "thinking":
304 addIfNotEmpty(segments, part.Get("thinking").String())
305 default:
306 // For unknown types, try to extract any text content
307 if part.Type == gjson.String {
308 addIfNotEmpty(segments, part.String())
309 } else if part.Type == gjson.JSON {
310 addIfNotEmpty(segments, part.Raw)
311 }
312 }
313 return true
314 })
315 }
316}
317
318// collectClaudeTools extracts text from Claude's tools array.
319func collectClaudeTools(tools gjson.Result, segments *[]string) {
320 if !tools.Exists() || !tools.IsArray() {
321 return
322 }
323 tools.ForEach(func(_, tool gjson.Result) bool {
324 addIfNotEmpty(segments, tool.Get("name").String())
325 addIfNotEmpty(segments, tool.Get("description").String())
326 if inputSchema := tool.Get("input_schema"); inputSchema.Exists() {
327 addIfNotEmpty(segments, inputSchema.Raw)
328 }
329 return true
330 })
70}331}
71 332
72// buildOpenAIUsageJSON returns a minimal usage structure understood by downstream translators.333// buildOpenAIUsageJSON returns a minimal usage structure understood by downstream translators.
internal/translator/claude/openai/chat-completions/claude_openai_response.go+4−0
@@ -50,6 +50,10 @@ type ToolCallAccumulator struct { · 49 unmodified lines
50// Returns:50// Returns:
51// - []string: A slice of strings, each containing an OpenAI-compatible JSON response51// - []string: A slice of strings, each containing an OpenAI-compatible JSON response
52func ConvertClaudeResponseToOpenAI(_ context.Context, modelName string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []string {52func ConvertClaudeResponseToOpenAI(_ context.Context, modelName string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []string {
53 var localParam any
54 if param == nil {
55 param = &localParam
56 }
53 if *param == nil {57 if *param == nil {
54 *param = &ConvertAnthropicResponseToOpenAIParams{58 *param = &ConvertAnthropicResponseToOpenAIParams{
55 CreatedAt: 0,59 CreatedAt: 0,
internal/translator/init.go+3−0
@@ -33,4 +33,7 @@ import ( · 32 unmodified lines
33 _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/gemini"33 _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/gemini"
34 _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/openai/chat-completions"34 _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/openai/chat-completions"
35 _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/openai/responses"35 _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/openai/responses"
36
37 _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/claude"
38 _ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/openai"
36)39)
internal/translator/kiro/claude/init.go+20−0
@@ -0,0 +1,20 @@
1// Package claude provides translation between Kiro and Claude formats.
2package claude
3
4import (
5 . "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
6 "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
7 "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/translator"
8)
9
10func init() {
11 translator.Register(
12 Claude,
13 Kiro,
14 ConvertClaudeRequestToKiro,
15 interfaces.TranslateResponse{
16 Stream: ConvertKiroStreamToClaude,
17 NonStream: ConvertKiroNonStreamToClaude,
18 },
19 )
20}
internal/translator/kiro/claude/kiro_claude.go+21−0
@@ -0,0 +1,21 @@
1// Package claude provides translation between Kiro and Claude formats.
2// Since Kiro executor generates Claude-compatible SSE format internally (with event: prefix),
3// translations are pass-through for streaming, but responses need proper formatting.
4package claude
5
6import (
7 "context"
8)
9
10// ConvertKiroStreamToClaude converts Kiro streaming response to Claude format.
11// Kiro executor already generates complete SSE format with "event:" prefix,
12// so this is a simple pass-through.
13func ConvertKiroStreamToClaude(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) []string {
14 return []string{string(rawResponse)}
15}
16
17// ConvertKiroNonStreamToClaude converts Kiro non-streaming response to Claude format.
18// The response is already in Claude format, so this is a pass-through.
19func ConvertKiroNonStreamToClaude(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) string {
20 return string(rawResponse)
21}
internal/translator/kiro/claude/kiro_claude_request.go+809−0
@@ -0,0 +1,809 @@
1// Package claude provides request translation functionality for Claude API to Kiro format.
2// It handles parsing and transforming Claude API requests into the Kiro/Amazon Q API format,
3// extracting model information, system instructions, message contents, and tool declarations.
4package claude
5
6import (
7 "encoding/json"
8 "fmt"
9 "net/http"
10 "strings"
11 "time"
12 "unicode/utf8"
13
14 "github.com/google/uuid"
15 kirocommon "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/common"
16 log "github.com/sirupsen/logrus"
17 "github.com/tidwall/gjson"
18)
19
20
21// Kiro API request structs - field order determines JSON key order
22
23// KiroPayload is the top-level request structure for Kiro API
24type KiroPayload struct {
25 ConversationState KiroConversationState `json:"conversationState"`
26 ProfileArn string `json:"profileArn,omitempty"`
27 InferenceConfig *KiroInferenceConfig `json:"inferenceConfig,omitempty"`
28}
29
30// KiroInferenceConfig contains inference parameters for the Kiro API.
31type KiroInferenceConfig struct {
32 MaxTokens int `json:"maxTokens,omitempty"`
33 Temperature float64 `json:"temperature,omitempty"`
34 TopP float64 `json:"topP,omitempty"`
35}
36
37
38// KiroConversationState holds the conversation context
39type KiroConversationState struct {
40 ChatTriggerType string `json:"chatTriggerType"` // Required: "MANUAL" - must be first field
41 ConversationID string `json:"conversationId"`
42 CurrentMessage KiroCurrentMessage `json:"currentMessage"`
43 History []KiroHistoryMessage `json:"history,omitempty"`
44}
45
46// KiroCurrentMessage wraps the current user message
47type KiroCurrentMessage struct {
48 UserInputMessage KiroUserInputMessage `json:"userInputMessage"`
49}
50
51// KiroHistoryMessage represents a message in the conversation history
52type KiroHistoryMessage struct {
53 UserInputMessage *KiroUserInputMessage `json:"userInputMessage,omitempty"`
54 AssistantResponseMessage *KiroAssistantResponseMessage `json:"assistantResponseMessage,omitempty"`
55}
56
57// KiroImage represents an image in Kiro API format
58type KiroImage struct {
59 Format string `json:"format"`
60 Source KiroImageSource `json:"source"`
61}
62
63// KiroImageSource contains the image data
64type KiroImageSource struct {
65 Bytes string `json:"bytes"` // base64 encoded image data
66}
67
68// KiroUserInputMessage represents a user message
69type KiroUserInputMessage struct {
70 Content string `json:"content"`
71 ModelID string `json:"modelId"`
72 Origin string `json:"origin"`
73 Images []KiroImage `json:"images,omitempty"`
74 UserInputMessageContext *KiroUserInputMessageContext `json:"userInputMessageContext,omitempty"`
75}
76
77// KiroUserInputMessageContext contains tool-related context
78type KiroUserInputMessageContext struct {
79 ToolResults []KiroToolResult `json:"toolResults,omitempty"`
80 Tools []KiroToolWrapper `json:"tools,omitempty"`
81}
82
83// KiroToolResult represents a tool execution result
84type KiroToolResult struct {
85 Content []KiroTextContent `json:"content"`
86 Status string `json:"status"`
87 ToolUseID string `json:"toolUseId"`
88}
89
90// KiroTextContent represents text content
91type KiroTextContent struct {
92 Text string `json:"text"`
93}
94
95// KiroToolWrapper wraps a tool specification
96type KiroToolWrapper struct {
97 ToolSpecification KiroToolSpecification `json:"toolSpecification"`
98}
99
100// KiroToolSpecification defines a tool's schema
101type KiroToolSpecification struct {
102 Name string `json:"name"`
103 Description string `json:"description"`
104 InputSchema KiroInputSchema `json:"inputSchema"`
105}
106
107// KiroInputSchema wraps the JSON schema for tool input
108type KiroInputSchema struct {
109 JSON interface{} `json:"json"`
110}
111
112// KiroAssistantResponseMessage represents an assistant message
113type KiroAssistantResponseMessage struct {
114 Content string `json:"content"`
115 ToolUses []KiroToolUse `json:"toolUses,omitempty"`
116}
117
118// KiroToolUse represents a tool invocation by the assistant
119type KiroToolUse struct {
120 ToolUseID string `json:"toolUseId"`
121 Name string `json:"name"`
122 Input map[string]interface{} `json:"input"`
123}
124
125// ConvertClaudeRequestToKiro converts a Claude API request to Kiro format.
126// This is the main entry point for request translation.
127func ConvertClaudeRequestToKiro(modelName string, inputRawJSON []byte, stream bool) []byte {
128 // For Kiro, we pass through the Claude format since buildKiroPayload
129 // expects Claude format and does the conversion internally.
130 // The actual conversion happens in the executor when building the HTTP request.
131 return inputRawJSON
132}
133
134// BuildKiroPayload constructs the Kiro API request payload from Claude format.
135// Supports tool calling - tools are passed via userInputMessageContext.
136// origin parameter determines which quota to use: "CLI" for Amazon Q, "AI_EDITOR" for Kiro IDE.
137// isAgentic parameter enables chunked write optimization prompt for -agentic model variants.
138// isChatOnly parameter disables tool calling for -chat model variants (pure conversation mode).
139// headers parameter allows checking Anthropic-Beta header for thinking mode detection.
140// metadata parameter is kept for API compatibility but no longer used for thinking configuration.
141// Supports thinking mode - when enabled, injects thinking tags into system prompt.
142// Returns the payload and a boolean indicating whether thinking mode was injected.
143func BuildKiroPayload(claudeBody []byte, modelID, profileArn, origin string, isAgentic, isChatOnly bool, headers http.Header, metadata map[string]any) ([]byte, bool) {
144 // Extract max_tokens for potential use in inferenceConfig
145 // Handle -1 as "use maximum" (Kiro max output is ~32000 tokens)
146 const kiroMaxOutputTokens = 32000
147 var maxTokens int64
148 if mt := gjson.GetBytes(claudeBody, "max_tokens"); mt.Exists() {
149 maxTokens = mt.Int()
150 if maxTokens == -1 {
151 maxTokens = kiroMaxOutputTokens
152 log.Debugf("kiro: max_tokens=-1 converted to %d", kiroMaxOutputTokens)
153 }
154 }
155
156 // Extract temperature if specified
157 var temperature float64
158 var hasTemperature bool
159 if temp := gjson.GetBytes(claudeBody, "temperature"); temp.Exists() {
160 temperature = temp.Float()
161 hasTemperature = true
162 }
163
164 // Extract top_p if specified
165 var topP float64
166 var hasTopP bool
167 if tp := gjson.GetBytes(claudeBody, "top_p"); tp.Exists() {
168 topP = tp.Float()
169 hasTopP = true
170 log.Debugf("kiro: extracted top_p: %.2f", topP)
171 }
172
173 // Normalize origin value for Kiro API compatibility
174 origin = normalizeOrigin(origin)
175 log.Debugf("kiro: normalized origin value: %s", origin)
176
177 messages := gjson.GetBytes(claudeBody, "messages")
178
179 // For chat-only mode, don't include tools
180 var tools gjson.Result
181 if !isChatOnly {
182 tools = gjson.GetBytes(claudeBody, "tools")
183 }
184
185 // Extract system prompt
186 systemPrompt := extractSystemPrompt(claudeBody)
187
188 // Check for thinking mode using the comprehensive IsThinkingEnabledWithHeaders function
189 // This supports Claude API format, OpenAI reasoning_effort, AMP/Cursor format, and Anthropic-Beta header
190 thinkingEnabled := IsThinkingEnabledWithHeaders(claudeBody, headers)
191
192 // Inject timestamp context
193 timestamp := time.Now().Format("2006-01-02 15:04:05 MST")
194 timestampContext := fmt.Sprintf("[Context: Current time is %s]", timestamp)
195 if systemPrompt != "" {
196 systemPrompt = timestampContext + "\n\n" + systemPrompt
197 } else {
198 systemPrompt = timestampContext
199 }
200 log.Debugf("kiro: injected timestamp context: %s", timestamp)
201
202 // Inject agentic optimization prompt for -agentic model variants
203 if isAgentic {
204 if systemPrompt != "" {
205 systemPrompt += "\n"
206 }
207 systemPrompt += kirocommon.KiroAgenticSystemPrompt
208 }
209
210 // Handle tool_choice parameter - Kiro doesn't support it natively, so we inject system prompt hints
211 // Claude tool_choice values: {"type": "auto/any/tool", "name": "..."}
212 toolChoiceHint := extractClaudeToolChoiceHint(claudeBody)
213 if toolChoiceHint != "" {
214 if systemPrompt != "" {
215 systemPrompt += "\n"
216 }
217 systemPrompt += toolChoiceHint
218 log.Debugf("kiro: injected tool_choice hint into system prompt")
219 }
220
221 // Convert Claude tools to Kiro format
222 kiroTools := convertClaudeToolsToKiro(tools)
223
224 // Thinking mode implementation:
225 // Kiro API supports official thinking/reasoning mode via <thinking_mode> tag.
226 // When set to "enabled", Kiro returns reasoning content as official reasoningContentEvent
227 // rather than inline <thinking> tags in assistantResponseEvent.
228 // We use a high max_thinking_length to allow extensive reasoning.
229 if thinkingEnabled {
230 thinkingHint := `<thinking_mode>enabled</thinking_mode>
231<max_thinking_length>200000</max_thinking_length>`
232 if systemPrompt != "" {
233 systemPrompt = thinkingHint + "\n\n" + systemPrompt
234 } else {
235 systemPrompt = thinkingHint
236 }
237 log.Infof("kiro: injected thinking prompt (official mode), has_tools: %v", len(kiroTools) > 0)
238 }
239
240 // Process messages and build history
241 history, currentUserMsg, currentToolResults := processMessages(messages, modelID, origin)
242
243 // Build content with system prompt
244 if currentUserMsg != nil {
245 currentUserMsg.Content = buildFinalContent(currentUserMsg.Content, systemPrompt, currentToolResults)
246
247 // Deduplicate currentToolResults
248 currentToolResults = deduplicateToolResults(currentToolResults)
249
250 // Build userInputMessageContext with tools and tool results
251 if len(kiroTools) > 0 || len(currentToolResults) > 0 {
252 currentUserMsg.UserInputMessageContext = &KiroUserInputMessageContext{
253 Tools: kiroTools,
254 ToolResults: currentToolResults,
255 }
256 }
257 }
258
259 // Build payload
260 var currentMessage KiroCurrentMessage
261 if currentUserMsg != nil {
262 currentMessage = KiroCurrentMessage{UserInputMessage: *currentUserMsg}
263 } else {
264 fallbackContent := ""
265 if systemPrompt != "" {
266 fallbackContent = "--- SYSTEM PROMPT ---\n" + systemPrompt + "\n--- END SYSTEM PROMPT ---\n"
267 }
268 currentMessage = KiroCurrentMessage{UserInputMessage: KiroUserInputMessage{
269 Content: fallbackContent,
270 ModelID: modelID,
271 Origin: origin,
272 }}
273 }
274
275 // Build inferenceConfig if we have any inference parameters
276 // Note: Kiro API doesn't actually use max_tokens for thinking budget
277 var inferenceConfig *KiroInferenceConfig
278 if maxTokens > 0 || hasTemperature || hasTopP {
279 inferenceConfig = &KiroInferenceConfig{}
280 if maxTokens > 0 {
281 inferenceConfig.MaxTokens = int(maxTokens)
282 }
283 if hasTemperature {
284 inferenceConfig.Temperature = temperature
285 }
286 if hasTopP {
287 inferenceConfig.TopP = topP
288 }
289 }
290
291 payload := KiroPayload{
292 ConversationState: KiroConversationState{
293 ChatTriggerType: "MANUAL",
294 ConversationID: uuid.New().String(),
295 CurrentMessage: currentMessage,
296 History: history,
297 },
298 ProfileArn: profileArn,
299 InferenceConfig: inferenceConfig,
300 }
301
302 result, err := json.Marshal(payload)
303 if err != nil {
304 log.Debugf("kiro: failed to marshal payload: %v", err)
305 return nil, false
306 }
307
308 return result, thinkingEnabled
309}
310
311// normalizeOrigin normalizes origin value for Kiro API compatibility
312func normalizeOrigin(origin string) string {
313 switch origin {
314 case "KIRO_CLI":
315 return "CLI"
316 case "KIRO_AI_EDITOR":
317 return "AI_EDITOR"
318 case "AMAZON_Q":
319 return "CLI"
320 case "KIRO_IDE":
321 return "AI_EDITOR"
322 default:
323 return origin
324 }
325}
326
327// extractSystemPrompt extracts system prompt from Claude request
328func extractSystemPrompt(claudeBody []byte) string {
329 systemField := gjson.GetBytes(claudeBody, "system")
330 if systemField.IsArray() {
331 var sb strings.Builder
332 for _, block := range systemField.Array() {
333 if block.Get("type").String() == "text" {
334 sb.WriteString(block.Get("text").String())
335 } else if block.Type == gjson.String {
336 sb.WriteString(block.String())
337 }
338 }
339 return sb.String()
340 }
341 return systemField.String()
342}
343
344// checkThinkingMode checks if thinking mode is enabled in the Claude request
345func checkThinkingMode(claudeBody []byte) (bool, int64) {
346 thinkingEnabled := false
347 var budgetTokens int64 = 24000
348
349 thinkingField := gjson.GetBytes(claudeBody, "thinking")
350 if thinkingField.Exists() {
351 thinkingType := thinkingField.Get("type").String()
352 if thinkingType == "enabled" {
353 thinkingEnabled = true
354 if bt := thinkingField.Get("budget_tokens"); bt.Exists() {
355 budgetTokens = bt.Int()
356 if budgetTokens <= 0 {
357 thinkingEnabled = false
358 log.Debugf("kiro: thinking mode disabled via budget_tokens <= 0")
359 }
360 }
361 if thinkingEnabled {
362 log.Debugf("kiro: thinking mode enabled via Claude API parameter, budget_tokens: %d", budgetTokens)
363 }
364 }
365 }
366
367 return thinkingEnabled, budgetTokens
368}
369
370// hasThinkingTagInBody checks if the request body already contains thinking configuration tags.
371// This is used to prevent duplicate injection when client (e.g., AMP/Cursor) already includes thinking config.
372func hasThinkingTagInBody(body []byte) bool {
373 bodyStr := string(body)
374 return strings.Contains(bodyStr, "<thinking_mode>") || strings.Contains(bodyStr, "<max_thinking_length>")
375}
376
377
378// IsThinkingEnabledFromHeader checks if thinking mode is enabled via Anthropic-Beta header.
379// Claude CLI uses "Anthropic-Beta: interleaved-thinking-2025-05-14" to enable thinking.
380func IsThinkingEnabledFromHeader(headers http.Header) bool {
381 if headers == nil {
382 return false
383 }
384 betaHeader := headers.Get("Anthropic-Beta")
385 if betaHeader == "" {
386 return false
387 }
388 // Check for interleaved-thinking beta feature
389 if strings.Contains(betaHeader, "interleaved-thinking") {
390 log.Debugf("kiro: thinking mode enabled via Anthropic-Beta header: %s", betaHeader)
391 return true
392 }
393 return false
394}
395
396// IsThinkingEnabled is a public wrapper to check if thinking mode is enabled.
397// This is used by the executor to determine whether to parse <thinking> tags in responses.
398// When thinking is NOT enabled in the request, <thinking> tags in responses should be
399// treated as regular text content, not as thinking blocks.
400//
401// Supports multiple formats:
402// - Claude API format: thinking.type = "enabled"
403// - OpenAI format: reasoning_effort parameter
404// - AMP/Cursor format: <thinking_mode>interleaved</thinking_mode> in system prompt
405func IsThinkingEnabled(body []byte) bool {
406 return IsThinkingEnabledWithHeaders(body, nil)
407}
408
409// IsThinkingEnabledWithHeaders checks if thinking mode is enabled from body or headers.
410// This is the comprehensive check that supports all thinking detection methods:
411// - Claude API format: thinking.type = "enabled"
412// - OpenAI format: reasoning_effort parameter
413// - AMP/Cursor format: <thinking_mode>interleaved</thinking_mode> in system prompt
414// - Anthropic-Beta header: interleaved-thinking-2025-05-14
415func IsThinkingEnabledWithHeaders(body []byte, headers http.Header) bool {
416 // Check Anthropic-Beta header first (Claude Code uses this)
417 if IsThinkingEnabledFromHeader(headers) {
418 return true
419 }
420
421 // Check Claude API format first (thinking.type = "enabled")
422 enabled, _ := checkThinkingMode(body)
423 if enabled {
424 log.Debugf("kiro: IsThinkingEnabled returning true (Claude API format)")
425 return true
426 }
427
428 // Check OpenAI format: reasoning_effort parameter
429 // Valid values: "low", "medium", "high", "auto" (not "none")
430 reasoningEffort := gjson.GetBytes(body, "reasoning_effort")
431 if reasoningEffort.Exists() {
432 effort := reasoningEffort.String()
433 if effort != "" && effort != "none" {
434 log.Debugf("kiro: thinking mode enabled via OpenAI reasoning_effort: %s", effort)
435 return true
436 }
437 }
438
439 // Check AMP/Cursor format: <thinking_mode>interleaved</thinking_mode> in system prompt
440 // This is how AMP client passes thinking configuration
441 bodyStr := string(body)
442 if strings.Contains(bodyStr, "<thinking_mode>") && strings.Contains(bodyStr, "</thinking_mode>") {
443 // Extract thinking mode value
444 startTag := "<thinking_mode>"
445 endTag := "</thinking_mode>"
446 startIdx := strings.Index(bodyStr, startTag)
447 if startIdx >= 0 {
448 startIdx += len(startTag)
449 endIdx := strings.Index(bodyStr[startIdx:], endTag)
450 if endIdx >= 0 {
451 thinkingMode := bodyStr[startIdx : startIdx+endIdx]
452 if thinkingMode == "interleaved" || thinkingMode == "enabled" {
453 log.Debugf("kiro: thinking mode enabled via AMP/Cursor format: %s", thinkingMode)
454 return true
455 }
456 }
457 }
458 }
459
460 // Check OpenAI format: max_completion_tokens with reasoning (o1-style)
461 // Some clients use this to indicate reasoning mode
462 if gjson.GetBytes(body, "max_completion_tokens").Exists() {
463 // If max_completion_tokens is set, check if model name suggests reasoning
464 model := gjson.GetBytes(body, "model").String()
465 if strings.Contains(strings.ToLower(model), "thinking") ||
466 strings.Contains(strings.ToLower(model), "reason") {
467 log.Debugf("kiro: thinking mode enabled via model name hint: %s", model)
468 return true
469 }
470 }
471
472 log.Debugf("kiro: IsThinkingEnabled returning false (no thinking mode detected)")
473 return false
474}
475
476// shortenToolNameIfNeeded shortens tool names that exceed 64 characters.
477// MCP tools often have long names like "mcp__server-name__tool-name".
478// This preserves the "mcp__" prefix and last segment when possible.
479func shortenToolNameIfNeeded(name string) string {
480 const limit = 64
481 if len(name) <= limit {
482 return name
483 }
484 // For MCP tools, try to preserve prefix and last segment
485 if strings.HasPrefix(name, "mcp__") {
486 idx := strings.LastIndex(name, "__")
487 if idx > 0 {
488 cand := "mcp__" + name[idx+2:]
489 if len(cand) > limit {
490 return cand[:limit]
491 }
492 return cand
493 }
494 }
495 return name[:limit]
496}
497
498// convertClaudeToolsToKiro converts Claude tools to Kiro format
499func convertClaudeToolsToKiro(tools gjson.Result) []KiroToolWrapper {
500 var kiroTools []KiroToolWrapper
501 if !tools.IsArray() {
502 return kiroTools
503 }
504
505 for _, tool := range tools.Array() {
506 name := tool.Get("name").String()
507 description := tool.Get("description").String()
508 inputSchema := tool.Get("input_schema").Value()
509
510 // Shorten tool name if it exceeds 64 characters (common with MCP tools)
511 originalName := name
512 name = shortenToolNameIfNeeded(name)
513 if name != originalName {
514 log.Debugf("kiro: shortened tool name from '%s' to '%s'", originalName, name)
515 }
516
517 // CRITICAL FIX: Kiro API requires non-empty description
518 if strings.TrimSpace(description) == "" {
519 description = fmt.Sprintf("Tool: %s", name)
520 log.Debugf("kiro: tool '%s' has empty description, using default: %s", name, description)
521 }
522
523 // Truncate long descriptions
524 if len(description) > kirocommon.KiroMaxToolDescLen {
525 truncLen := kirocommon.KiroMaxToolDescLen - 30
526 for truncLen > 0 && !utf8.RuneStart(description[truncLen]) {
527 truncLen--
528 }
529 description = description[:truncLen] + "... (description truncated)"
530 }
531
532 kiroTools = append(kiroTools, KiroToolWrapper{
533 ToolSpecification: KiroToolSpecification{
534 Name: name,
535 Description: description,
536 InputSchema: KiroInputSchema{JSON: inputSchema},
537 },
538 })
539 }
540
541 return kiroTools
542}
543
544// processMessages processes Claude messages and builds Kiro history
545func processMessages(messages gjson.Result, modelID, origin string) ([]KiroHistoryMessage, *KiroUserInputMessage, []KiroToolResult) {
546 var history []KiroHistoryMessage
547 var currentUserMsg *KiroUserInputMessage
548 var currentToolResults []KiroToolResult
549
550 // Merge adjacent messages with the same role
551 messagesArray := kirocommon.MergeAdjacentMessages(messages.Array())
552 for i, msg := range messagesArray {
553 role := msg.Get("role").String()
554 isLastMessage := i == len(messagesArray)-1
555
556 if role == "user" {
557 userMsg, toolResults := BuildUserMessageStruct(msg, modelID, origin)
558 if isLastMessage {
559 currentUserMsg = &userMsg
560 currentToolResults = toolResults
561 } else {
562 // CRITICAL: Kiro API requires content to be non-empty for history messages too
563 if strings.TrimSpace(userMsg.Content) == "" {
564 if len(toolResults) > 0 {
565 userMsg.Content = "Tool results provided."
566 } else {
567 userMsg.Content = "Continue"
568 }
569 }
570 // For history messages, embed tool results in context
571 if len(toolResults) > 0 {
572 userMsg.UserInputMessageContext = &KiroUserInputMessageContext{
573 ToolResults: toolResults,
574 }
575 }
576 history = append(history, KiroHistoryMessage{
577 UserInputMessage: &userMsg,
578 })
579 }
580 } else if role == "assistant" {
581 assistantMsg := BuildAssistantMessageStruct(msg)
582 if isLastMessage {
583 history = append(history, KiroHistoryMessage{
584 AssistantResponseMessage: &assistantMsg,
585 })
586 // Create a "Continue" user message as currentMessage
587 currentUserMsg = &KiroUserInputMessage{
588 Content: "Continue",
589 ModelID: modelID,
590 Origin: origin,
591 }
592 } else {
593 history = append(history, KiroHistoryMessage{
594 AssistantResponseMessage: &assistantMsg,
595 })
596 }
597 }
598 }
599
600 return history, currentUserMsg, currentToolResults
601}
602
603// buildFinalContent builds the final content with system prompt
604func buildFinalContent(content, systemPrompt string, toolResults []KiroToolResult) string {
605 var contentBuilder strings.Builder
606
607 if systemPrompt != "" {
608 contentBuilder.WriteString("--- SYSTEM PROMPT ---\n")
609 contentBuilder.WriteString(systemPrompt)
610 contentBuilder.WriteString("\n--- END SYSTEM PROMPT ---\n\n")
611 }
612
613 contentBuilder.WriteString(content)
614 finalContent := contentBuilder.String()
615
616 // CRITICAL: Kiro API requires content to be non-empty
617 if strings.TrimSpace(finalContent) == "" {
618 if len(toolResults) > 0 {
619 finalContent = "Tool results provided."
620 } else {
621 finalContent = "Continue"
622 }
623 log.Debugf("kiro: content was empty, using default: %s", finalContent)
624 }
625
626 return finalContent
627}
628
629// deduplicateToolResults removes duplicate tool results
630func deduplicateToolResults(toolResults []KiroToolResult) []KiroToolResult {
631 if len(toolResults) == 0 {
632 return toolResults
633 }
634
635 seenIDs := make(map[string]bool)
636 unique := make([]KiroToolResult, 0, len(toolResults))
637 for _, tr := range toolResults {
638 if !seenIDs[tr.ToolUseID] {
639 seenIDs[tr.ToolUseID] = true
640 unique = append(unique, tr)
641 } else {
642 log.Debugf("kiro: skipping duplicate toolResult in currentMessage: %s", tr.ToolUseID)
643 }
644 }
645 return unique
646}
647
648// extractClaudeToolChoiceHint extracts tool_choice from Claude request and returns a system prompt hint.
649// Claude tool_choice values:
650// - {"type": "auto"}: Model decides (default, no hint needed)
651// - {"type": "any"}: Must use at least one tool
652// - {"type": "tool", "name": "..."}: Must use specific tool
653func extractClaudeToolChoiceHint(claudeBody []byte) string {
654 toolChoice := gjson.GetBytes(claudeBody, "tool_choice")
655 if !toolChoice.Exists() {
656 return ""
657 }
658
659 toolChoiceType := toolChoice.Get("type").String()
660 switch toolChoiceType {
661 case "any":
662 return "[INSTRUCTION: You MUST use at least one of the available tools to respond. Do not respond with text only - always make a tool call.]"
663 case "tool":
664 toolName := toolChoice.Get("name").String()
665 if toolName != "" {
666 return fmt.Sprintf("[INSTRUCTION: You MUST use the tool named '%s' to respond. Do not use any other tool or respond with text only.]", toolName)
667 }
668 case "auto":
669 // Default behavior, no hint needed
670 return ""
671 }
672
673 return ""
674}
675
676// BuildUserMessageStruct builds a user message and extracts tool results
677func BuildUserMessageStruct(msg gjson.Result, modelID, origin string) (KiroUserInputMessage, []KiroToolResult) {
678 content := msg.Get("content")
679 var contentBuilder strings.Builder
680 var toolResults []KiroToolResult
681 var images []KiroImage
682
683 // Track seen toolUseIds to deduplicate
684 seenToolUseIDs := make(map[string]bool)
685
686 if content.IsArray() {
687 for _, part := range content.Array() {
688 partType := part.Get("type").String()
689 switch partType {
690 case "text":
691 contentBuilder.WriteString(part.Get("text").String())
692 case "image":
693 mediaType := part.Get("source.media_type").String()
694 data := part.Get("source.data").String()
695
696 format := ""
697 if idx := strings.LastIndex(mediaType, "/"); idx != -1 {
698 format = mediaType[idx+1:]
699 }
700
701 if format != "" && data != "" {
702 images = append(images, KiroImage{
703 Format: format,
704 Source: KiroImageSource{
705 Bytes: data,
706 },
707 })
708 }
709 case "tool_result":
710 toolUseID := part.Get("tool_use_id").String()
711
712 // Skip duplicate toolUseIds
713 if seenToolUseIDs[toolUseID] {
714 log.Debugf("kiro: skipping duplicate tool_result with toolUseId: %s", toolUseID)
715 continue
716 }
717 seenToolUseIDs[toolUseID] = true
718
719 isError := part.Get("is_error").Bool()
720 resultContent := part.Get("content")
721
722 var textContents []KiroTextContent
723 if resultContent.IsArray() {
724 for _, item := range resultContent.Array() {
725 if item.Get("type").String() == "text" {
726 textContents = append(textContents, KiroTextContent{Text: item.Get("text").String()})
727 } else if item.Type == gjson.String {
728 textContents = append(textContents, KiroTextContent{Text: item.String()})
729 }
730 }
731 } else if resultContent.Type == gjson.String {
732 textContents = append(textContents, KiroTextContent{Text: resultContent.String()})
733 }
734
735 if len(textContents) == 0 {
736 textContents = append(textContents, KiroTextContent{Text: "Tool use was cancelled by the user"})
737 }
738
739 status := "success"
740 if isError {
741 status = "error"
742 }
743
744 toolResults = append(toolResults, KiroToolResult{
745 ToolUseID: toolUseID,
746 Content: textContents,
747 Status: status,
748 })
749 }
750 }
751 } else {
752 contentBuilder.WriteString(content.String())
753 }
754
755 userMsg := KiroUserInputMessage{
756 Content: contentBuilder.String(),
757 ModelID: modelID,
758 Origin: origin,
759 }
760
761 if len(images) > 0 {
762 userMsg.Images = images
763 }
764
765 return userMsg, toolResults
766}
767
768// BuildAssistantMessageStruct builds an assistant message with tool uses
769func BuildAssistantMessageStruct(msg gjson.Result) KiroAssistantResponseMessage {
770 content := msg.Get("content")
771 var contentBuilder strings.Builder
772 var toolUses []KiroToolUse
773
774 if content.IsArray() {
775 for _, part := range content.Array() {
776 partType := part.Get("type").String()
777 switch partType {
778 case "text":
779 contentBuilder.WriteString(part.Get("text").String())
780 case "tool_use":
781 toolUseID := part.Get("id").String()
782 toolName := part.Get("name").String()
783 toolInput := part.Get("input")
784
785 var inputMap map[string]interface{}
786 if toolInput.IsObject() {
787 inputMap = make(map[string]interface{})
788 toolInput.ForEach(func(key, value gjson.Result) bool {
789 inputMap[key.String()] = value.Value()
790 return true
791 })
792 }
793
794 toolUses = append(toolUses, KiroToolUse{
795 ToolUseID: toolUseID,
796 Name: toolName,
797 Input: inputMap,
798 })
799 }
800 }
801 } else {
802 contentBuilder.WriteString(content.String())
803 }
804
805 return KiroAssistantResponseMessage{
806 Content: contentBuilder.String(),
807 ToolUses: toolUses,
808 }
809}
internal/translator/kiro/claude/kiro_claude_response.go+204−0
@@ -0,0 +1,204 @@
1// Package claude provides response translation functionality for Kiro API to Claude format.
2// This package handles the conversion of Kiro API responses into Claude-compatible format,
3// including support for thinking blocks and tool use.
4package claude
5
6import (
7 "crypto/sha256"
8 "encoding/base64"
9 "encoding/json"
10 "strings"
11
12 "github.com/google/uuid"
13 "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
14 log "github.com/sirupsen/logrus"
15
16 kirocommon "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/common"
17)
18
19// generateThinkingSignature generates a signature for thinking content.
20// This is required by Claude API for thinking blocks in non-streaming responses.
21// The signature is a base64-encoded hash of the thinking content.
22func generateThinkingSignature(thinkingContent string) string {
23 if thinkingContent == "" {
24 return ""
25 }
26 // Generate a deterministic signature based on content hash
27 hash := sha256.Sum256([]byte(thinkingContent))
28 return base64.StdEncoding.EncodeToString(hash[:])
29}
30
31// Local references to kirocommon constants for thinking block parsing
32var (
33 thinkingStartTag = kirocommon.ThinkingStartTag
34 thinkingEndTag = kirocommon.ThinkingEndTag
35)
36
37// BuildClaudeResponse constructs a Claude-compatible response.
38// Supports tool_use blocks when tools are present in the response.
39// Supports thinking blocks - parses <thinking> tags and converts to Claude thinking content blocks.
40// stopReason is passed from upstream; fallback logic applied if empty.
41func BuildClaudeResponse(content string, toolUses []KiroToolUse, model string, usageInfo usage.Detail, stopReason string) []byte {
42 var contentBlocks []map[string]interface{}
43
44 // Extract thinking blocks and text from content
45 if content != "" {
46 blocks := ExtractThinkingFromContent(content)
47 contentBlocks = append(contentBlocks, blocks...)
48
49 // Log if thinking blocks were extracted
50 for _, block := range blocks {
51 if block["type"] == "thinking" {
52 thinkingContent := block["thinking"].(string)
53 log.Infof("kiro: buildClaudeResponse extracted thinking block (len: %d)", len(thinkingContent))
54 }
55 }
56 }
57
58 // Add tool_use blocks
59 for _, toolUse := range toolUses {
60 contentBlocks = append(contentBlocks, map[string]interface{}{
61 "type": "tool_use",
62 "id": toolUse.ToolUseID,
63 "name": toolUse.Name,
64 "input": toolUse.Input,
65 })
66 }
67
68 // Ensure at least one content block (Claude API requires non-empty content)
69 if len(contentBlocks) == 0 {
70 contentBlocks = append(contentBlocks, map[string]interface{}{
71 "type": "text",
72 "text": "",
73 })
74 }
75
76 // Use upstream stopReason; apply fallback logic if not provided
77 if stopReason == "" {
78 stopReason = "end_turn"
79 if len(toolUses) > 0 {
80 stopReason = "tool_use"
81 }
82 log.Debugf("kiro: buildClaudeResponse using fallback stop_reason: %s", stopReason)
83 }
84
85 // Log warning if response was truncated due to max_tokens
86 if stopReason == "max_tokens" {
87 log.Warnf("kiro: response truncated due to max_tokens limit (buildClaudeResponse)")
88 }
89
90 response := map[string]interface{}{
91 "id": "msg_" + uuid.New().String()[:24],
92 "type": "message",
93 "role": "assistant",
94 "model": model,
95 "content": contentBlocks,
96 "stop_reason": stopReason,
97 "usage": map[string]interface{}{
98 "input_tokens": usageInfo.InputTokens,
99 "output_tokens": usageInfo.OutputTokens,
100 },
101 }
102 result, _ := json.Marshal(response)
103 return result
104}
105
106// ExtractThinkingFromContent parses content to extract thinking blocks and text.
107// Returns a list of content blocks in the order they appear in the content.
108// Handles interleaved thinking and text blocks correctly.
109func ExtractThinkingFromContent(content string) []map[string]interface{} {
110 var blocks []map[string]interface{}
111
112 if content == "" {
113 return blocks
114 }
115
116 // Check if content contains thinking tags at all
117 if !strings.Contains(content, thinkingStartTag) {
118 // No thinking tags, return as plain text
119 return []map[string]interface{}{
120 {
121 "type": "text",
122 "text": content,
123 },
124 }
125 }
126
127 log.Debugf("kiro: extractThinkingFromContent - found thinking tags in content (len: %d)", len(content))
128
129 remaining := content
130
131 for len(remaining) > 0 {
132 // Look for <thinking> tag
133 startIdx := strings.Index(remaining, thinkingStartTag)
134
135 if startIdx == -1 {
136 // No more thinking tags, add remaining as text
137 if strings.TrimSpace(remaining) != "" {
138 blocks = append(blocks, map[string]interface{}{
139 "type": "text",
140 "text": remaining,
141 })
142 }
143 break
144 }
145
146 // Add text before thinking tag (if any meaningful content)
147 if startIdx > 0 {
148 textBefore := remaining[:startIdx]
149 if strings.TrimSpace(textBefore) != "" {
150 blocks = append(blocks, map[string]interface{}{
151 "type": "text",
152 "text": textBefore,
153 })
154 }
155 }
156
157 // Move past the opening tag
158 remaining = remaining[startIdx+len(thinkingStartTag):]
159
160 // Find closing tag
161 endIdx := strings.Index(remaining, thinkingEndTag)
162
163 if endIdx == -1 {
164 // No closing tag found, treat rest as thinking content (incomplete response)
165 if strings.TrimSpace(remaining) != "" {
166 // Generate signature for thinking content (required by Claude API)
167 signature := generateThinkingSignature(remaining)
168 blocks = append(blocks, map[string]interface{}{
169 "type": "thinking",
170 "thinking": remaining,
171 "signature": signature,
172 })
173 log.Warnf("kiro: extractThinkingFromContent - missing closing </thinking> tag")
174 }
175 break
176 }
177
178 // Extract thinking content between tags
179 thinkContent := remaining[:endIdx]
180 if strings.TrimSpace(thinkContent) != "" {
181 // Generate signature for thinking content (required by Claude API)
182 signature := generateThinkingSignature(thinkContent)
183 blocks = append(blocks, map[string]interface{}{
184 "type": "thinking",
185 "thinking": thinkContent,
186 "signature": signature,
187 })
188 log.Debugf("kiro: extractThinkingFromContent - extracted thinking block (len: %d)", len(thinkContent))
189 }
190
191 // Move past the closing tag
192 remaining = remaining[endIdx+len(thinkingEndTag):]
193 }
194
195 // If no blocks were created (all whitespace), return empty text block
196 if len(blocks) == 0 {
197 blocks = append(blocks, map[string]interface{}{
198 "type": "text",
199 "text": "",
200 })
201 }
202
203 return blocks
204}
internal/translator/kiro/claude/kiro_claude_stream.go+186−0
@@ -0,0 +1,186 @@
1// Package claude provides streaming SSE event building for Claude format.
2// This package handles the construction of Claude-compatible Server-Sent Events (SSE)
3// for streaming responses from Kiro API.
4package claude
5
6import (
7 "encoding/json"
8
9 "github.com/google/uuid"
10 "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
11)
12
13// BuildClaudeMessageStartEvent creates the message_start SSE event
14func BuildClaudeMessageStartEvent(model string, inputTokens int64) []byte {
15 event := map[string]interface{}{
16 "type": "message_start",
17 "message": map[string]interface{}{
18 "id": "msg_" + uuid.New().String()[:24],
19 "type": "message",
20 "role": "assistant",
21 "content": []interface{}{},
22 "model": model,
23 "stop_reason": nil,
24 "stop_sequence": nil,
25 "usage": map[string]interface{}{"input_tokens": inputTokens, "output_tokens": 0},
26 },
27 }
28 result, _ := json.Marshal(event)
29 return []byte("event: message_start\ndata: " + string(result))
30}
31
32// BuildClaudeContentBlockStartEvent creates a content_block_start SSE event
33func BuildClaudeContentBlockStartEvent(index int, blockType, toolUseID, toolName string) []byte {
34 var contentBlock map[string]interface{}
35 switch blockType {
36 case "tool_use":
37 contentBlock = map[string]interface{}{
38 "type": "tool_use",
39 "id": toolUseID,
40 "name": toolName,
41 "input": map[string]interface{}{},
42 }
43 case "thinking":
44 contentBlock = map[string]interface{}{
45 "type": "thinking",
46 "thinking": "",
47 }
48 default:
49 contentBlock = map[string]interface{}{
50 "type": "text",
51 "text": "",
52 }
53 }
54
55 event := map[string]interface{}{
56 "type": "content_block_start",
57 "index": index,
58 "content_block": contentBlock,
59 }
60 result, _ := json.Marshal(event)
61 return []byte("event: content_block_start\ndata: " + string(result))
62}
63
64// BuildClaudeStreamEvent creates a text_delta content_block_delta SSE event
65func BuildClaudeStreamEvent(contentDelta string, index int) []byte {
66 event := map[string]interface{}{
67 "type": "content_block_delta",
68 "index": index,
69 "delta": map[string]interface{}{
70 "type": "text_delta",
71 "text": contentDelta,
72 },
73 }
74 result, _ := json.Marshal(event)
75 return []byte("event: content_block_delta\ndata: " + string(result))
76}
77
78// BuildClaudeInputJsonDeltaEvent creates an input_json_delta event for tool use streaming
79func BuildClaudeInputJsonDeltaEvent(partialJSON string, index int) []byte {
80 event := map[string]interface{}{
81 "type": "content_block_delta",
82 "index": index,
83 "delta": map[string]interface{}{
84 "type": "input_json_delta",
85 "partial_json": partialJSON,
86 },
87 }
88 result, _ := json.Marshal(event)
89 return []byte("event: content_block_delta\ndata: " + string(result))
90}
91
92// BuildClaudeContentBlockStopEvent creates a content_block_stop SSE event
93func BuildClaudeContentBlockStopEvent(index int) []byte {
94 event := map[string]interface{}{
95 "type": "content_block_stop",
96 "index": index,
97 }
98 result, _ := json.Marshal(event)
99 return []byte("event: content_block_stop\ndata: " + string(result))
100}
101
102// BuildClaudeThinkingBlockStopEvent creates a content_block_stop SSE event for thinking blocks.
103func BuildClaudeThinkingBlockStopEvent(index int) []byte {
104 event := map[string]interface{}{
105 "type": "content_block_stop",
106 "index": index,
107 }
108 result, _ := json.Marshal(event)
109 return []byte("event: content_block_stop\ndata: " + string(result))
110}
111
112// BuildClaudeMessageDeltaEvent creates the message_delta event with stop_reason and usage
113func BuildClaudeMessageDeltaEvent(stopReason string, usageInfo usage.Detail) []byte {
114 deltaEvent := map[string]interface{}{
115 "type": "message_delta",
116 "delta": map[string]interface{}{
117 "stop_reason": stopReason,
118 "stop_sequence": nil,
119 },
120 "usage": map[string]interface{}{
121 "input_tokens": usageInfo.InputTokens,
122 "output_tokens": usageInfo.OutputTokens,
123 },
124 }
125 deltaResult, _ := json.Marshal(deltaEvent)
126 return []byte("event: message_delta\ndata: " + string(deltaResult))
127}
128
129// BuildClaudeMessageStopOnlyEvent creates only the message_stop event
130func BuildClaudeMessageStopOnlyEvent() []byte {
131 stopEvent := map[string]interface{}{
132 "type": "message_stop",
133 }
134 stopResult, _ := json.Marshal(stopEvent)
135 return []byte("event: message_stop\ndata: " + string(stopResult))
136}
137
138// BuildClaudePingEventWithUsage creates a ping event with embedded usage information.
139// This is used for real-time usage estimation during streaming.
140func BuildClaudePingEventWithUsage(inputTokens, outputTokens int64) []byte {
141 event := map[string]interface{}{
142 "type": "ping",
143 "usage": map[string]interface{}{
144 "input_tokens": inputTokens,
145 "output_tokens": outputTokens,
146 "total_tokens": inputTokens + outputTokens,
147 "estimated": true,
148 },
149 }
150 result, _ := json.Marshal(event)
151 return []byte("event: ping\ndata: " + string(result))
152}
153
154// BuildClaudeThinkingDeltaEvent creates a thinking_delta event for Claude API compatibility.
155// This is used when streaming thinking content wrapped in <thinking> tags.
156func BuildClaudeThinkingDeltaEvent(thinkingDelta string, index int) []byte {
157 event := map[string]interface{}{
158 "type": "content_block_delta",
159 "index": index,
160 "delta": map[string]interface{}{
161 "type": "thinking_delta",
162 "thinking": thinkingDelta,
163 },
164 }
165 result, _ := json.Marshal(event)
166 return []byte("event: content_block_delta\ndata: " + string(result))
167}
168
169// PendingTagSuffix detects if the buffer ends with a partial prefix of the given tag.
170// Returns the length of the partial match (0 if no match).
171// Based on amq2api implementation for handling cross-chunk tag boundaries.
172func PendingTagSuffix(buffer, tag string) int {
173 if buffer == "" || tag == "" {
174 return 0
175 }
176 maxLen := len(buffer)
177 if maxLen > len(tag)-1 {
178 maxLen = len(tag) - 1
179 }
180 for length := maxLen; length > 0; length-- {
181 if len(buffer) >= length && buffer[len(buffer)-length:] == tag[:length] {
182 return length
183 }
184 }
185 return 0
186}
internal/translator/kiro/claude/kiro_claude_tools.go+522−0
@@ -0,0 +1,522 @@
1// Package claude provides tool calling support for Kiro to Claude translation.
2// This package handles parsing embedded tool calls, JSON repair, and deduplication.
3package claude
4
5import (
6 "encoding/json"
7 "regexp"
8 "strings"
9
10 "github.com/google/uuid"
11 kirocommon "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/common"
12 log "github.com/sirupsen/logrus"
13)
14
15// ToolUseState tracks the state of an in-progress tool use during streaming.
16type ToolUseState struct {
17 ToolUseID string
18 Name string
19 InputBuffer strings.Builder
20 IsComplete bool
21}
22
23// Pre-compiled regex patterns for performance
24var (
25 // embeddedToolCallPattern matches [Called tool_name with args: {...}] format
26 embeddedToolCallPattern = regexp.MustCompile(`\[Called\s+([A-Za-z0-9_.-]+)\s+with\s+args:\s*`)
27 // trailingCommaPattern matches trailing commas before closing braces/brackets
28 trailingCommaPattern = regexp.MustCompile(`,\s*([}\]])`)
29)
30
31// ParseEmbeddedToolCalls extracts [Called tool_name with args: {...}] format from text.
32// Kiro sometimes embeds tool calls in text content instead of using toolUseEvent.
33// Returns the cleaned text (with tool calls removed) and extracted tool uses.
34func ParseEmbeddedToolCalls(text string, processedIDs map[string]bool) (string, []KiroToolUse) {
35 if !strings.Contains(text, "[Called") {
36 return text, nil
37 }
38
39 var toolUses []KiroToolUse
40 cleanText := text
41
42 // Find all [Called markers
43 matches := embeddedToolCallPattern.FindAllStringSubmatchIndex(text, -1)
44 if len(matches) == 0 {
45 return text, nil
46 }
47
48 // Process matches in reverse order to maintain correct indices
49 for i := len(matches) - 1; i >= 0; i-- {
50 matchStart := matches[i][0]
51 toolNameStart := matches[i][2]
52 toolNameEnd := matches[i][3]
53
54 if toolNameStart < 0 || toolNameEnd < 0 {
55 continue
56 }
57
58 toolName := text[toolNameStart:toolNameEnd]
59
60 // Find the JSON object start (after "with args:")
61 jsonStart := matches[i][1]
62 if jsonStart >= len(text) {
63 continue
64 }
65
66 // Skip whitespace to find the opening brace
67 for jsonStart < len(text) && (text[jsonStart] == ' ' || text[jsonStart] == '\t') {
68 jsonStart++
69 }
70
71 if jsonStart >= len(text) || text[jsonStart] != '{' {
72 continue
73 }
74
75 // Find matching closing bracket
76 jsonEnd := findMatchingBracket(text, jsonStart)
77 if jsonEnd < 0 {
78 continue
79 }
80
81 // Extract JSON and find the closing bracket of [Called ...]
82 jsonStr := text[jsonStart : jsonEnd+1]
83
84 // Find the closing ] after the JSON
85 closingBracket := jsonEnd + 1
86 for closingBracket < len(text) && text[closingBracket] != ']' {
87 closingBracket++
88 }
89 if closingBracket >= len(text) {
90 continue
91 }
92
93 // End index of the full tool call (closing ']' inclusive)
94 matchEnd := closingBracket + 1
95
96 // Repair and parse JSON
97 repairedJSON := RepairJSON(jsonStr)
98 var inputMap map[string]interface{}
99 if err := json.Unmarshal([]byte(repairedJSON), &inputMap); err != nil {
100 log.Debugf("kiro: failed to parse embedded tool call JSON: %v, raw: %s", err, jsonStr)
101 continue
102 }
103
104 // Generate unique tool ID
105 toolUseID := "toolu_" + uuid.New().String()[:12]
106
107 // Check for duplicates using name+input as key
108 dedupeKey := toolName + ":" + repairedJSON
109 if processedIDs != nil {
110 if processedIDs[dedupeKey] {
111 log.Debugf("kiro: skipping duplicate embedded tool call: %s", toolName)
112 // Still remove from text even if duplicate
113 if matchStart >= 0 && matchEnd <= len(cleanText) && matchStart <= matchEnd {
114 cleanText = cleanText[:matchStart] + cleanText[matchEnd:]
115 }
116 continue
117 }
118 processedIDs[dedupeKey] = true
119 }
120
121 toolUses = append(toolUses, KiroToolUse{
122 ToolUseID: toolUseID,
123 Name: toolName,
124 Input: inputMap,
125 })
126
127 log.Infof("kiro: extracted embedded tool call: %s (ID: %s)", toolName, toolUseID)
128
129 // Remove from clean text (index-based removal to avoid deleting the wrong occurrence)
130 if matchStart >= 0 && matchEnd <= len(cleanText) && matchStart <= matchEnd {
131 cleanText = cleanText[:matchStart] + cleanText[matchEnd:]
132 }
133 }
134
135 return cleanText, toolUses
136}
137
138// findMatchingBracket finds the index of the closing brace/bracket that matches
139// the opening one at startPos. Handles nested objects and strings correctly.
140func findMatchingBracket(text string, startPos int) int {
141 if startPos >= len(text) {
142 return -1
143 }
144
145 openChar := text[startPos]
146 var closeChar byte
147 switch openChar {
148 case '{':
149 closeChar = '}'
150 case '[':
151 closeChar = ']'
152 default:
153 return -1
154 }
155
156 depth := 1
157 inString := false
158 escapeNext := false
159
160 for i := startPos + 1; i < len(text); i++ {
161 char := text[i]
162
163 if escapeNext {
164 escapeNext = false
165 continue
166 }
167
168 if char == '\\' && inString {
169 escapeNext = true
170 continue
171 }
172
173 if char == '"' {
174 inString = !inString
175 continue
176 }
177
178 if !inString {
179 if char == openChar {
180 depth++
181 } else if char == closeChar {
182 depth--
183 if depth == 0 {
184 return i
185 }
186 }
187 }
188 }
189
190 return -1
191}
192
193// RepairJSON attempts to fix common JSON issues that may occur in tool call arguments.
194// Conservative repair strategy:
195// 1. First try to parse JSON directly - if valid, return as-is
196// 2. Only attempt repair if parsing fails
197// 3. After repair, validate the result - if still invalid, return original
198func RepairJSON(jsonString string) string {
199 // Handle empty or invalid input
200 if jsonString == "" {
201 return "{}"
202 }
203
204 str := strings.TrimSpace(jsonString)
205 if str == "" {
206 return "{}"
207 }
208
209 // CONSERVATIVE STRATEGY: First try to parse directly
210 var testParse interface{}
211 if err := json.Unmarshal([]byte(str), &testParse); err == nil {
212 log.Debugf("kiro: repairJSON - JSON is already valid, returning unchanged")
213 return str
214 }
215
216 log.Debugf("kiro: repairJSON - JSON parse failed, attempting repair")
217 originalStr := str
218
219 // First, escape unescaped newlines/tabs within JSON string values
220 str = escapeNewlinesInStrings(str)
221 // Remove trailing commas before closing braces/brackets
222 str = trailingCommaPattern.ReplaceAllString(str, "$1")
223
224 // Calculate bracket balance
225 braceCount := 0
226 bracketCount := 0
227 inString := false
228 escape := false
229 lastValidIndex := -1
230
231 for i := 0; i < len(str); i++ {
232 char := str[i]
233
234 if escape {
235 escape = false
236 continue
237 }
238
239 if char == '\\' {
240 escape = true
241 continue
242 }
243
244 if char == '"' {
245 inString = !inString
246 continue
247 }
248
249 if inString {
250 continue
251 }
252
253 switch char {
254 case '{':
255 braceCount++
256 case '}':
257 braceCount--
258 case '[':
259 bracketCount++
260 case ']':
261 bracketCount--
262 }
263
264 if braceCount >= 0 && bracketCount >= 0 {
265 lastValidIndex = i
266 }
267 }
268
269 // If brackets are unbalanced, try to repair
270 if braceCount > 0 || bracketCount > 0 {
271 if lastValidIndex > 0 && lastValidIndex < len(str)-1 {
272 truncated := str[:lastValidIndex+1]
273 // Recount brackets after truncation
274 braceCount = 0
275 bracketCount = 0
276 inString = false
277 escape = false
278 for i := 0; i < len(truncated); i++ {
279 char := truncated[i]
280 if escape {
281 escape = false
282 continue
283 }
284 if char == '\\' {
285 escape = true
286 continue
287 }
288 if char == '"' {
289 inString = !inString
290 continue
291 }
292 if inString {
293 continue
294 }
295 switch char {
296 case '{':
297 braceCount++
298 case '}':
299 braceCount--
300 case '[':
301 bracketCount++
302 case ']':
303 bracketCount--
304 }
305 }
306 str = truncated
307 }
308
309 // Add missing closing brackets
310 for braceCount > 0 {
311 str += "}"
312 braceCount--
313 }
314 for bracketCount > 0 {
315 str += "]"
316 bracketCount--
317 }
318 }
319
320 // Validate repaired JSON
321 if err := json.Unmarshal([]byte(str), &testParse); err != nil {
322 log.Warnf("kiro: repairJSON - repair failed to produce valid JSON, returning original")
323 return originalStr
324 }
325
326 log.Debugf("kiro: repairJSON - successfully repaired JSON")
327 return str
328}
329
330// escapeNewlinesInStrings escapes literal newlines, tabs, and other control characters
331// that appear inside JSON string values.
332func escapeNewlinesInStrings(raw string) string {
333 var result strings.Builder
334 result.Grow(len(raw) + 100)
335
336 inString := false
337 escaped := false
338
339 for i := 0; i < len(raw); i++ {
340 c := raw[i]
341
342 if escaped {
343 result.WriteByte(c)
344 escaped = false
345 continue
346 }
347
348 if c == '\\' && inString {
349 result.WriteByte(c)
350 escaped = true
351 continue
352 }
353
354 if c == '"' {
355 inString = !inString
356 result.WriteByte(c)
357 continue
358 }
359
360 if inString {
361 switch c {
362 case '\n':
363 result.WriteString("\\n")
364 case '\r':
365 result.WriteString("\\r")
366 case '\t':
367 result.WriteString("\\t")
368 default:
369 result.WriteByte(c)
370 }
371 } else {
372 result.WriteByte(c)
373 }
374 }
375
376 return result.String()
377}
378
379// ProcessToolUseEvent handles a toolUseEvent from the Kiro stream.
380// It accumulates input fragments and emits tool_use blocks when complete.
381// Returns events to emit and updated state.
382func ProcessToolUseEvent(event map[string]interface{}, currentToolUse *ToolUseState, processedIDs map[string]bool) ([]KiroToolUse, *ToolUseState) {
383 var toolUses []KiroToolUse
384
385 // Extract from nested toolUseEvent or direct format
386 tu := event
387 if nested, ok := event["toolUseEvent"].(map[string]interface{}); ok {
388 tu = nested
389 }
390
391 toolUseID := kirocommon.GetString(tu, "toolUseId")
392 toolName := kirocommon.GetString(tu, "name")
393 isStop := false
394 if stop, ok := tu["stop"].(bool); ok {
395 isStop = stop
396 }
397
398 // Get input - can be string (fragment) or object (complete)
399 var inputFragment string
400 var inputMap map[string]interface{}
401
402 if inputRaw, ok := tu["input"]; ok {
403 switch v := inputRaw.(type) {
404 case string:
405 inputFragment = v
406 case map[string]interface{}:
407 inputMap = v
408 }
409 }
410
411 // New tool use starting
412 if toolUseID != "" && toolName != "" {
413 if currentToolUse != nil && currentToolUse.ToolUseID != toolUseID {
414 log.Warnf("kiro: interleaved tool use detected - new ID %s arrived while %s in progress, completing previous",
415 toolUseID, currentToolUse.ToolUseID)
416 if !processedIDs[currentToolUse.ToolUseID] {
417 incomplete := KiroToolUse{
418 ToolUseID: currentToolUse.ToolUseID,
419 Name: currentToolUse.Name,
420 }
421 if currentToolUse.InputBuffer.Len() > 0 {
422 raw := currentToolUse.InputBuffer.String()
423 repaired := RepairJSON(raw)
424
425 var input map[string]interface{}
426 if err := json.Unmarshal([]byte(repaired), &input); err != nil {
427 log.Warnf("kiro: failed to parse interleaved tool input: %v, raw: %s", err, raw)
428 input = make(map[string]interface{})
429 }
430 incomplete.Input = input
431 }
432 toolUses = append(toolUses, incomplete)
433 processedIDs[currentToolUse.ToolUseID] = true
434 }
435 currentToolUse = nil
436 }
437
438 if currentToolUse == nil {
439 if processedIDs != nil && processedIDs[toolUseID] {
440 log.Debugf("kiro: skipping duplicate toolUseEvent: %s", toolUseID)
441 return nil, nil
442 }
443
444 currentToolUse = &ToolUseState{
445 ToolUseID: toolUseID,
446 Name: toolName,
447 }
448 log.Infof("kiro: starting new tool use: %s (ID: %s)", toolName, toolUseID)
449 }
450 }
451
452 // Accumulate input fragments
453 if currentToolUse != nil && inputFragment != "" {
454 currentToolUse.InputBuffer.WriteString(inputFragment)
455 log.Debugf("kiro: accumulated input fragment, total length: %d", currentToolUse.InputBuffer.Len())
456 }
457
458 // If complete input object provided directly
459 if currentToolUse != nil && inputMap != nil {
460 inputBytes, _ := json.Marshal(inputMap)
461 currentToolUse.InputBuffer.Reset()
462 currentToolUse.InputBuffer.Write(inputBytes)
463 }
464
465 // Tool use complete
466 if isStop && currentToolUse != nil {
467 fullInput := currentToolUse.InputBuffer.String()
468
469 // Repair and parse the accumulated JSON
470 repairedJSON := RepairJSON(fullInput)
471 var finalInput map[string]interface{}
472 if err := json.Unmarshal([]byte(repairedJSON), &finalInput); err != nil {
473 log.Warnf("kiro: failed to parse accumulated tool input: %v, raw: %s", err, fullInput)
474 finalInput = make(map[string]interface{})
475 }
476
477 toolUse := KiroToolUse{
478 ToolUseID: currentToolUse.ToolUseID,
479 Name: currentToolUse.Name,
480 Input: finalInput,
481 }
482 toolUses = append(toolUses, toolUse)
483
484 if processedIDs != nil {
485 processedIDs[currentToolUse.ToolUseID] = true
486 }
487
488 log.Infof("kiro: completed tool use: %s (ID: %s)", currentToolUse.Name, currentToolUse.ToolUseID)
489 return toolUses, nil
490 }
491
492 return toolUses, currentToolUse
493}
494
495// DeduplicateToolUses removes duplicate tool uses based on toolUseId and content.
496func DeduplicateToolUses(toolUses []KiroToolUse) []KiroToolUse {
497 seenIDs := make(map[string]bool)
498 seenContent := make(map[string]bool)
499 var unique []KiroToolUse
500
501 for _, tu := range toolUses {
502 if seenIDs[tu.ToolUseID] {
503 log.Debugf("kiro: removing ID-duplicate tool use: %s (name: %s)", tu.ToolUseID, tu.Name)
504 continue
505 }
506
507 inputJSON, _ := json.Marshal(tu.Input)
508 contentKey := tu.Name + ":" + string(inputJSON)
509
510 if seenContent[contentKey] {
511 log.Debugf("kiro: removing content-duplicate tool use: %s (id: %s)", tu.Name, tu.ToolUseID)
512 continue
513 }
514
515 seenIDs[tu.ToolUseID] = true
516 seenContent[contentKey] = true
517 unique = append(unique, tu)
518 }
519
520 return unique
521}
522
internal/translator/kiro/common/constants.go+75−0
@@ -0,0 +1,75 @@
1// Package common provides shared constants and utilities for Kiro translator.
2package common
3
4const (
5 // KiroMaxToolDescLen is the maximum description length for Kiro API tools.
6 // Kiro API limit is 10240 bytes, leave room for "..."
7 KiroMaxToolDescLen = 10237
8
9 // ThinkingStartTag is the start tag for thinking blocks in responses.
10 ThinkingStartTag = "<thinking>"
11
12 // ThinkingEndTag is the end tag for thinking blocks in responses.
13 ThinkingEndTag = "</thinking>"
14
15 // CodeFenceMarker is the markdown code fence marker.
16 CodeFenceMarker = "```"
17
18 // AltCodeFenceMarker is the alternative markdown code fence marker.
19 AltCodeFenceMarker = "~~~"
20
21 // InlineCodeMarker is the markdown inline code marker (backtick).
22 InlineCodeMarker = "`"
23
24 // KiroAgenticSystemPrompt is injected only for -agentic models to prevent timeouts on large writes.
25 // AWS Kiro API has a 2-3 minute timeout for large file write operations.
26 KiroAgenticSystemPrompt = `
27# CRITICAL: CHUNKED WRITE PROTOCOL (MANDATORY)
28
29You MUST follow these rules for ALL file operations. Violation causes server timeouts and task failure.
30
31## ABSOLUTE LIMITS
32- **MAXIMUM 350 LINES** per single write/edit operation - NO EXCEPTIONS
33- **RECOMMENDED 300 LINES** or less for optimal performance
34- **NEVER** write entire files in one operation if >300 lines
35
36## MANDATORY CHUNKED WRITE STRATEGY
37
38### For NEW FILES (>300 lines total):
391. FIRST: Write initial chunk (first 250-300 lines) using write_to_file/fsWrite
402. THEN: Append remaining content in 250-300 line chunks using file append operations
413. REPEAT: Continue appending until complete
42
43### For EDITING EXISTING FILES:
441. Use surgical edits (apply_diff/targeted edits) - change ONLY what's needed
452. NEVER rewrite entire files - use incremental modifications
463. Split large refactors into multiple small, focused edits
47
48### For LARGE CODE GENERATION:
491. Generate in logical sections (imports, types, functions separately)
502. Write each section as a separate operation
513. Use append operations for subsequent sections
52
53## EXAMPLES OF CORRECT BEHAVIOR
54
55✅ CORRECT: Writing a 600-line file
56- Operation 1: Write lines 1-300 (initial file creation)
57- Operation 2: Append lines 301-600
58
59✅ CORRECT: Editing multiple functions
60- Operation 1: Edit function A
61- Operation 2: Edit function B
62- Operation 3: Edit function C
63
64❌ WRONG: Writing 500 lines in single operation → TIMEOUT
65❌ WRONG: Rewriting entire file to change 5 lines → TIMEOUT
66❌ WRONG: Generating massive code blocks without chunking → TIMEOUT
67
68## WHY THIS MATTERS
69- Server has 2-3 minute timeout for operations
70- Large writes exceed timeout and FAIL completely
71- Chunked writes are FASTER and more RELIABLE
72- Failed writes waste time and require retry
73
74REMEMBER: When in doubt, write LESS per operation. Multiple small operations > one large operation.`
75)
internal/translator/kiro/common/message_merge.go+132−0
@@ -0,0 +1,132 @@
1// Package common provides shared utilities for Kiro translators.
2package common
3
4import (
5 "encoding/json"
6
7 "github.com/tidwall/gjson"
8)
9
10// MergeAdjacentMessages merges adjacent messages with the same role.
11// This reduces API call complexity and improves compatibility.
12// Based on AIClient-2-API implementation.
13// NOTE: Tool messages are NOT merged because each has a unique tool_call_id that must be preserved.
14func MergeAdjacentMessages(messages []gjson.Result) []gjson.Result {
15 if len(messages) <= 1 {
16 return messages
17 }
18
19 var merged []gjson.Result
20 for _, msg := range messages {
21 if len(merged) == 0 {
22 merged = append(merged, msg)
23 continue
24 }
25
26 lastMsg := merged[len(merged)-1]
27 currentRole := msg.Get("role").String()
28 lastRole := lastMsg.Get("role").String()
29
30 // Don't merge tool messages - each has a unique tool_call_id
31 if currentRole == "tool" || lastRole == "tool" {
32 merged = append(merged, msg)
33 continue
34 }
35
36 if currentRole == lastRole {
37 // Merge content from current message into last message
38 mergedContent := mergeMessageContent(lastMsg, msg)
39 // Create a new merged message JSON
40 mergedMsg := createMergedMessage(lastRole, mergedContent)
41 merged[len(merged)-1] = gjson.Parse(mergedMsg)
42 } else {
43 merged = append(merged, msg)
44 }
45 }
46
47 return merged
48}
49
50// mergeMessageContent merges the content of two messages with the same role.
51// Handles both string content and array content (with text, tool_use, tool_result blocks).
52func mergeMessageContent(msg1, msg2 gjson.Result) string {
53 content1 := msg1.Get("content")
54 content2 := msg2.Get("content")
55
56 // Extract content blocks from both messages
57 var blocks1, blocks2 []map[string]interface{}
58
59 if content1.IsArray() {
60 for _, block := range content1.Array() {
61 blocks1 = append(blocks1, blockToMap(block))
62 }
63 } else if content1.Type == gjson.String {
64 blocks1 = append(blocks1, map[string]interface{}{
65 "type": "text",
66 "text": content1.String(),
67 })
68 }
69
70 if content2.IsArray() {
71 for _, block := range content2.Array() {
72 blocks2 = append(blocks2, blockToMap(block))
73 }
74 } else if content2.Type == gjson.String {
75 blocks2 = append(blocks2, map[string]interface{}{
76 "type": "text",
77 "text": content2.String(),
78 })
79 }
80
81 // Merge text blocks if both end/start with text
82 if len(blocks1) > 0 && len(blocks2) > 0 {
83 if blocks1[len(blocks1)-1]["type"] == "text" && blocks2[0]["type"] == "text" {
84 // Merge the last text block of msg1 with the first text block of msg2
85 text1 := blocks1[len(blocks1)-1]["text"].(string)
86 text2 := blocks2[0]["text"].(string)
87 blocks1[len(blocks1)-1]["text"] = text1 + "\n" + text2
88 blocks2 = blocks2[1:] // Remove the merged block from blocks2
89 }
90 }
91
92 // Combine all blocks
93 allBlocks := append(blocks1, blocks2...)
94
95 // Convert to JSON
96 result, _ := json.Marshal(allBlocks)
97 return string(result)
98}
99
100// blockToMap converts a gjson.Result block to a map[string]interface{}
101func blockToMap(block gjson.Result) map[string]interface{} {
102 result := make(map[string]interface{})
103 block.ForEach(func(key, value gjson.Result) bool {
104 if value.IsObject() {
105 result[key.String()] = blockToMap(value)
106 } else if value.IsArray() {
107 var arr []interface{}
108 for _, item := range value.Array() {
109 if item.IsObject() {
110 arr = append(arr, blockToMap(item))
111 } else {
112 arr = append(arr, item.Value())
113 }
114 }
115 result[key.String()] = arr
116 } else {
117 result[key.String()] = value.Value()
118 }
119 return true
120 })
121 return result
122}
123
124// createMergedMessage creates a JSON string for a merged message
125func createMergedMessage(role string, content string) string {
126 msg := map[string]interface{}{
127 "role": role,
128 "content": json.RawMessage(content),
129 }
130 result, _ := json.Marshal(msg)
131 return string(result)
132}
internal/translator/kiro/common/utils.go+16−0
@@ -0,0 +1,16 @@
1// Package common provides shared constants and utilities for Kiro translator.
2package common
3
4// GetString safely extracts a string from a map.
5// Returns empty string if the key doesn't exist or the value is not a string.
6func GetString(m map[string]interface{}, key string) string {
7 if v, ok := m[key].(string); ok {
8 return v
9 }
10 return ""
11}
12
13// GetStringValue is an alias for GetString for backward compatibility.
14func GetStringValue(m map[string]interface{}, key string) string {
15 return GetString(m, key)
16}
internal/translator/kiro/openai/init.go+20−0
@@ -0,0 +1,20 @@
1// Package openai provides translation between OpenAI Chat Completions and Kiro formats.
2package openai
3
4import (
5 . "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
6 "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
7 "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/translator"
8)
9
10func init() {
11 translator.Register(
12 OpenAI, // source format
13 Kiro, // target format
14 ConvertOpenAIRequestToKiro,
15 interfaces.TranslateResponse{
16 Stream: ConvertKiroStreamToOpenAI,
17 NonStream: ConvertKiroNonStreamToOpenAI,
18 },
19 )
20}
internal/translator/kiro/openai/kiro_openai.go+371−0
@@ -0,0 +1,371 @@
1// Package openai provides translation between OpenAI Chat Completions and Kiro formats.
2// This package enables direct OpenAI → Kiro translation, bypassing the Claude intermediate layer.
3//
4// The Kiro executor generates Claude-compatible SSE format internally, so the streaming response
5// translation converts from Claude SSE format to OpenAI SSE format.
6package openai
7
8import (
9 "bytes"
10 "context"
11 "encoding/json"
12 "strings"
13
14 kirocommon "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/common"
15 "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
16 log "github.com/sirupsen/logrus"
17 "github.com/tidwall/gjson"
18)
19
20// ConvertKiroStreamToOpenAI converts Kiro streaming response to OpenAI format.
21// The Kiro executor emits Claude-compatible SSE events, so this function translates
22// from Claude SSE format to OpenAI SSE format.
23//
24// Claude SSE format:
25// - event: message_start\ndata: {...}
26// - event: content_block_start\ndata: {...}
27// - event: content_block_delta\ndata: {...}
28// - event: content_block_stop\ndata: {...}
29// - event: message_delta\ndata: {...}
30// - event: message_stop\ndata: {...}
31//
32// OpenAI SSE format:
33// - data: {"id":"...","object":"chat.completion.chunk",...}
34// - data: [DONE]
35func ConvertKiroStreamToOpenAI(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) []string {
36 // Initialize state if needed
37 if *param == nil {
38 *param = NewOpenAIStreamState(model)
39 }
40 state := (*param).(*OpenAIStreamState)
41
42 // Parse the Claude SSE event
43 responseStr := string(rawResponse)
44
45 // Handle raw event format (event: xxx\ndata: {...})
46 var eventType string
47 var eventData string
48
49 if strings.HasPrefix(responseStr, "event:") {
50 // Parse event type and data
51 lines := strings.SplitN(responseStr, "\n", 2)
52 if len(lines) >= 1 {
53 eventType = strings.TrimSpace(strings.TrimPrefix(lines[0], "event:"))
54 }
55 if len(lines) >= 2 && strings.HasPrefix(lines[1], "data:") {
56 eventData = strings.TrimSpace(strings.TrimPrefix(lines[1], "data:"))
57 }
58 } else if strings.HasPrefix(responseStr, "data:") {
59 // Just data line
60 eventData = strings.TrimSpace(strings.TrimPrefix(responseStr, "data:"))
61 } else {
62 // Try to parse as raw JSON
63 eventData = strings.TrimSpace(responseStr)
64 }
65
66 if eventData == "" {
67 return []string{}
68 }
69
70 // Parse the event data as JSON
71 eventJSON := gjson.Parse(eventData)
72 if !eventJSON.Exists() {
73 return []string{}
74 }
75
76 // Determine event type from JSON if not already set
77 if eventType == "" {
78 eventType = eventJSON.Get("type").String()
79 }
80
81 var results []string
82
83 switch eventType {
84 case "message_start":
85 // Send first chunk with role
86 firstChunk := BuildOpenAISSEFirstChunk(state)
87 results = append(results, firstChunk)
88
89 case "content_block_start":
90 // Check block type
91 blockType := eventJSON.Get("content_block.type").String()
92 switch blockType {
93 case "text":
94 // Text block starting - nothing to emit yet
95 case "thinking":
96 // Thinking block starting - nothing to emit yet for OpenAI
97 case "tool_use":
98 // Tool use block starting
99 toolUseID := eventJSON.Get("content_block.id").String()
100 toolName := eventJSON.Get("content_block.name").String()
101 chunk := BuildOpenAISSEToolCallStart(state, toolUseID, toolName)
102 results = append(results, chunk)
103 state.ToolCallIndex++
104 }
105
106 case "content_block_delta":
107 deltaType := eventJSON.Get("delta.type").String()
108 switch deltaType {
109 case "text_delta":
110 textDelta := eventJSON.Get("delta.text").String()
111 if textDelta != "" {
112 chunk := BuildOpenAISSETextDelta(state, textDelta)
113 results = append(results, chunk)
114 }
115 case "thinking_delta":
116 // Convert thinking to reasoning_content for o1-style compatibility
117 thinkingDelta := eventJSON.Get("delta.thinking").String()
118 if thinkingDelta != "" {
119 chunk := BuildOpenAISSEReasoningDelta(state, thinkingDelta)
120 results = append(results, chunk)
121 }
122 case "input_json_delta":
123 // Tool call arguments delta
124 partialJSON := eventJSON.Get("delta.partial_json").String()
125 if partialJSON != "" {
126 // Get the tool index from content block index
127 blockIndex := int(eventJSON.Get("index").Int())
128 chunk := BuildOpenAISSEToolCallArgumentsDelta(state, partialJSON, blockIndex-1) // Adjust for 0-based tool index
129 results = append(results, chunk)
130 }
131 }
132
133 case "content_block_stop":
134 // Content block ended - nothing to emit for OpenAI
135
136 case "message_delta":
137 // Message delta with stop_reason
138 stopReason := eventJSON.Get("delta.stop_reason").String()
139 finishReason := mapKiroStopReasonToOpenAI(stopReason)
140 if finishReason != "" {
141 chunk := BuildOpenAISSEFinish(state, finishReason)
142 results = append(results, chunk)
143 }
144
145 // Extract usage if present
146 if eventJSON.Get("usage").Exists() {
147 inputTokens := eventJSON.Get("usage.input_tokens").Int()
148 outputTokens := eventJSON.Get("usage.output_tokens").Int()
149 usageInfo := usage.Detail{
150 InputTokens: inputTokens,
151 OutputTokens: outputTokens,
152 TotalTokens: inputTokens + outputTokens,
153 }
154 chunk := BuildOpenAISSEUsage(state, usageInfo)
155 results = append(results, chunk)
156 }
157
158 case "message_stop":
159 // Final event - do NOT emit [DONE] here
160 // The handler layer (openai_handlers.go) will send [DONE] when the stream closes
161 // Emitting [DONE] here would cause duplicate [DONE] markers
162
163 case "ping":
164 // Ping event with usage - optionally emit usage chunk
165 if eventJSON.Get("usage").Exists() {
166 inputTokens := eventJSON.Get("usage.input_tokens").Int()
167 outputTokens := eventJSON.Get("usage.output_tokens").Int()
168 usageInfo := usage.Detail{
169 InputTokens: inputTokens,
170 OutputTokens: outputTokens,
171 TotalTokens: inputTokens + outputTokens,
172 }
173 chunk := BuildOpenAISSEUsage(state, usageInfo)
174 results = append(results, chunk)
175 }
176 }
177
178 return results
179}
180
181// ConvertKiroNonStreamToOpenAI converts Kiro non-streaming response to OpenAI format.
182// The Kiro executor returns Claude-compatible JSON responses, so this function translates
183// from Claude format to OpenAI format.
184func ConvertKiroNonStreamToOpenAI(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) string {
185 // Parse the Claude-format response
186 response := gjson.ParseBytes(rawResponse)
187
188 // Extract content
189 var content string
190 var reasoningContent string
191 var toolUses []KiroToolUse
192 var stopReason string
193
194 // Get stop_reason
195 stopReason = response.Get("stop_reason").String()
196
197 // Process content blocks
198 contentBlocks := response.Get("content")
199 if contentBlocks.IsArray() {
200 for _, block := range contentBlocks.Array() {
201 blockType := block.Get("type").String()
202 switch blockType {
203 case "text":
204 content += block.Get("text").String()
205 case "thinking":
206 // Convert thinking blocks to reasoning_content for OpenAI format
207 reasoningContent += block.Get("thinking").String()
208 case "tool_use":
209 toolUseID := block.Get("id").String()
210 toolName := block.Get("name").String()
211 toolInput := block.Get("input")
212
213 var inputMap map[string]interface{}
214 if toolInput.IsObject() {
215 inputMap = make(map[string]interface{})
216 toolInput.ForEach(func(key, value gjson.Result) bool {
217 inputMap[key.String()] = value.Value()
218 return true
219 })
220 }
221
222 toolUses = append(toolUses, KiroToolUse{
223 ToolUseID: toolUseID,
224 Name: toolName,
225 Input: inputMap,
226 })
227 }
228 }
229 }
230
231 // Extract usage
232 usageInfo := usage.Detail{
233 InputTokens: response.Get("usage.input_tokens").Int(),
234 OutputTokens: response.Get("usage.output_tokens").Int(),
235 }
236 usageInfo.TotalTokens = usageInfo.InputTokens + usageInfo.OutputTokens
237
238 // Build OpenAI response with reasoning_content support
239 openaiResponse := BuildOpenAIResponseWithReasoning(content, reasoningContent, toolUses, model, usageInfo, stopReason)
240 return string(openaiResponse)
241}
242
243// ParseClaudeEvent parses a Claude SSE event and returns the event type and data
244func ParseClaudeEvent(rawEvent []byte) (eventType string, eventData []byte) {
245 lines := bytes.Split(rawEvent, []byte("\n"))
246 for _, line := range lines {
247 line = bytes.TrimSpace(line)
248 if bytes.HasPrefix(line, []byte("event:")) {
249 eventType = string(bytes.TrimSpace(bytes.TrimPrefix(line, []byte("event:"))))
250 } else if bytes.HasPrefix(line, []byte("data:")) {
251 eventData = bytes.TrimSpace(bytes.TrimPrefix(line, []byte("data:")))
252 }
253 }
254 return eventType, eventData
255}
256
257// ExtractThinkingFromContent parses content to extract thinking blocks.
258// Returns cleaned content (without thinking tags) and whether thinking was found.
259func ExtractThinkingFromContent(content string) (string, string, bool) {
260 if !strings.Contains(content, kirocommon.ThinkingStartTag) {
261 return content, "", false
262 }
263
264 var cleanedContent strings.Builder
265 var thinkingContent strings.Builder
266 hasThinking := false
267 remaining := content
268
269 for len(remaining) > 0 {
270 startIdx := strings.Index(remaining, kirocommon.ThinkingStartTag)
271 if startIdx == -1 {
272 cleanedContent.WriteString(remaining)
273 break
274 }
275
276 // Add content before thinking tag
277 cleanedContent.WriteString(remaining[:startIdx])
278
279 // Move past opening tag
280 remaining = remaining[startIdx+len(kirocommon.ThinkingStartTag):]
281
282 // Find closing tag
283 endIdx := strings.Index(remaining, kirocommon.ThinkingEndTag)
284 if endIdx == -1 {
285 // No closing tag - treat rest as thinking
286 thinkingContent.WriteString(remaining)
287 hasThinking = true
288 break
289 }
290
291 // Extract thinking content
292 thinkingContent.WriteString(remaining[:endIdx])
293 hasThinking = true
294 remaining = remaining[endIdx+len(kirocommon.ThinkingEndTag):]
295 }
296
297 return strings.TrimSpace(cleanedContent.String()), strings.TrimSpace(thinkingContent.String()), hasThinking
298}
299
300// ConvertOpenAIToolsToKiroFormat is a helper that converts OpenAI tools format to Kiro format
301func ConvertOpenAIToolsToKiroFormat(tools []map[string]interface{}) []KiroToolWrapper {
302 var kiroTools []KiroToolWrapper
303
304 for _, tool := range tools {
305 toolType, _ := tool["type"].(string)
306 if toolType != "function" {
307 continue
308 }
309
310 fn, ok := tool["function"].(map[string]interface{})
311 if !ok {
312 continue
313 }
314
315 name := kirocommon.GetString(fn, "name")
316 description := kirocommon.GetString(fn, "description")
317 parameters := fn["parameters"]
318
319 if name == "" {
320 continue
321 }
322
323 if description == "" {
324 description = "Tool: " + name
325 }
326
327 kiroTools = append(kiroTools, KiroToolWrapper{
328 ToolSpecification: KiroToolSpecification{
329 Name: name,
330 Description: description,
331 InputSchema: KiroInputSchema{JSON: parameters},
332 },
333 })
334 }
335
336 return kiroTools
337}
338
339// OpenAIStreamParams holds parameters for OpenAI streaming conversion
340type OpenAIStreamParams struct {
341 State *OpenAIStreamState
342 ThinkingState *ThinkingTagState
343 ToolCallsEmitted map[string]bool
344}
345
346// NewOpenAIStreamParams creates new streaming parameters
347func NewOpenAIStreamParams(model string) *OpenAIStreamParams {
348 return &OpenAIStreamParams{
349 State: NewOpenAIStreamState(model),
350 ThinkingState: NewThinkingTagState(),
351 ToolCallsEmitted: make(map[string]bool),
352 }
353}
354
355// ConvertClaudeToolUseToOpenAI converts a Claude tool_use block to OpenAI tool_calls format
356func ConvertClaudeToolUseToOpenAI(toolUseID, toolName string, input map[string]interface{}) map[string]interface{} {
357 inputJSON, _ := json.Marshal(input)
358 return map[string]interface{}{
359 "id": toolUseID,
360 "type": "function",
361 "function": map[string]interface{}{
362 "name": toolName,
363 "arguments": string(inputJSON),
364 },
365 }
366}
367
368// LogStreamEvent logs a streaming event for debugging
369func LogStreamEvent(eventType, data string) {
370 log.Debugf("kiro-openai: stream event type=%s, data_len=%d", eventType, len(data))
371}
internal/translator/kiro/openai/kiro_openai_request.go+863−0
@@ -0,0 +1,863 @@
1// Package openai provides request translation from OpenAI Chat Completions to Kiro format.
2// It handles parsing and transforming OpenAI API requests into the Kiro/Amazon Q API format,
3// extracting model information, system instructions, message contents, and tool declarations.
4package openai
5
6import (
7 "encoding/json"
8 "fmt"
9 "net/http"
10 "strings"
11 "time"
12 "unicode/utf8"
13
14 "github.com/google/uuid"
15 kiroclaude "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/claude"
16 kirocommon "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/common"
17 log "github.com/sirupsen/logrus"
18 "github.com/tidwall/gjson"
19)
20
21// Kiro API request structs - reuse from kiroclaude package structure
22
23// KiroPayload is the top-level request structure for Kiro API
24type KiroPayload struct {
25 ConversationState KiroConversationState `json:"conversationState"`
26 ProfileArn string `json:"profileArn,omitempty"`
27 InferenceConfig *KiroInferenceConfig `json:"inferenceConfig,omitempty"`
28}
29
30// KiroInferenceConfig contains inference parameters for the Kiro API.
31type KiroInferenceConfig struct {
32 MaxTokens int `json:"maxTokens,omitempty"`
33 Temperature float64 `json:"temperature,omitempty"`
34 TopP float64 `json:"topP,omitempty"`
35}
36
37// KiroConversationState holds the conversation context
38type KiroConversationState struct {
39 ChatTriggerType string `json:"chatTriggerType"` // Required: "MANUAL"
40 ConversationID string `json:"conversationId"`
41 CurrentMessage KiroCurrentMessage `json:"currentMessage"`
42 History []KiroHistoryMessage `json:"history,omitempty"`
43}
44
45// KiroCurrentMessage wraps the current user message
46type KiroCurrentMessage struct {
47 UserInputMessage KiroUserInputMessage `json:"userInputMessage"`
48}
49
50// KiroHistoryMessage represents a message in the conversation history
51type KiroHistoryMessage struct {
52 UserInputMessage *KiroUserInputMessage `json:"userInputMessage,omitempty"`
53 AssistantResponseMessage *KiroAssistantResponseMessage `json:"assistantResponseMessage,omitempty"`
54}
55
56// KiroImage represents an image in Kiro API format
57type KiroImage struct {
58 Format string `json:"format"`
59 Source KiroImageSource `json:"source"`
60}
61
62// KiroImageSource contains the image data
63type KiroImageSource struct {
64 Bytes string `json:"bytes"` // base64 encoded image data
65}
66
67// KiroUserInputMessage represents a user message
68type KiroUserInputMessage struct {
69 Content string `json:"content"`
70 ModelID string `json:"modelId"`
71 Origin string `json:"origin"`
72 Images []KiroImage `json:"images,omitempty"`
73 UserInputMessageContext *KiroUserInputMessageContext `json:"userInputMessageContext,omitempty"`
74}
75
76// KiroUserInputMessageContext contains tool-related context
77type KiroUserInputMessageContext struct {
78 ToolResults []KiroToolResult `json:"toolResults,omitempty"`
79 Tools []KiroToolWrapper `json:"tools,omitempty"`
80}
81
82// KiroToolResult represents a tool execution result
83type KiroToolResult struct {
84 Content []KiroTextContent `json:"content"`
85 Status string `json:"status"`
86 ToolUseID string `json:"toolUseId"`
87}
88
89// KiroTextContent represents text content
90type KiroTextContent struct {
91 Text string `json:"text"`
92}
93
94// KiroToolWrapper wraps a tool specification
95type KiroToolWrapper struct {
96 ToolSpecification KiroToolSpecification `json:"toolSpecification"`
97}
98
99// KiroToolSpecification defines a tool's schema
100type KiroToolSpecification struct {
101 Name string `json:"name"`
102 Description string `json:"description"`
103 InputSchema KiroInputSchema `json:"inputSchema"`
104}
105
106// KiroInputSchema wraps the JSON schema for tool input
107type KiroInputSchema struct {
108 JSON interface{} `json:"json"`
109}
110
111// KiroAssistantResponseMessage represents an assistant message
112type KiroAssistantResponseMessage struct {
113 Content string `json:"content"`
114 ToolUses []KiroToolUse `json:"toolUses,omitempty"`
115}
116
117// KiroToolUse represents a tool invocation by the assistant
118type KiroToolUse struct {
119 ToolUseID string `json:"toolUseId"`
120 Name string `json:"name"`
121 Input map[string]interface{} `json:"input"`
122}
123
124// ConvertOpenAIRequestToKiro converts an OpenAI Chat Completions request to Kiro format.
125// This is the main entry point for request translation.
126// Note: The actual payload building happens in the executor, this just passes through
127// the OpenAI format which will be converted by BuildKiroPayloadFromOpenAI.
128func ConvertOpenAIRequestToKiro(modelName string, inputRawJSON []byte, stream bool) []byte {
129 // Pass through the OpenAI format - actual conversion happens in BuildKiroPayloadFromOpenAI
130 return inputRawJSON
131}
132
133// BuildKiroPayloadFromOpenAI constructs the Kiro API request payload from OpenAI format.
134// Supports tool calling - tools are passed via userInputMessageContext.
135// origin parameter determines which quota to use: "CLI" for Amazon Q, "AI_EDITOR" for Kiro IDE.
136// isAgentic parameter enables chunked write optimization prompt for -agentic model variants.
137// isChatOnly parameter disables tool calling for -chat model variants (pure conversation mode).
138// headers parameter allows checking Anthropic-Beta header for thinking mode detection.
139// metadata parameter is kept for API compatibility but no longer used for thinking configuration.
140// Returns the payload and a boolean indicating whether thinking mode was injected.
141func BuildKiroPayloadFromOpenAI(openaiBody []byte, modelID, profileArn, origin string, isAgentic, isChatOnly bool, headers http.Header, metadata map[string]any) ([]byte, bool) {
142 // Extract max_tokens for potential use in inferenceConfig
143 // Handle -1 as "use maximum" (Kiro max output is ~32000 tokens)
144 const kiroMaxOutputTokens = 32000
145 var maxTokens int64
146 if mt := gjson.GetBytes(openaiBody, "max_tokens"); mt.Exists() {
147 maxTokens = mt.Int()
148 if maxTokens == -1 {
149 maxTokens = kiroMaxOutputTokens
150 log.Debugf("kiro-openai: max_tokens=-1 converted to %d", kiroMaxOutputTokens)
151 }
152 }
153
154 // Extract temperature if specified
155 var temperature float64
156 var hasTemperature bool
157 if temp := gjson.GetBytes(openaiBody, "temperature"); temp.Exists() {
158 temperature = temp.Float()
159 hasTemperature = true
160 }
161
162 // Extract top_p if specified
163 var topP float64
164 var hasTopP bool
165 if tp := gjson.GetBytes(openaiBody, "top_p"); tp.Exists() {
166 topP = tp.Float()
167 hasTopP = true
168 log.Debugf("kiro-openai: extracted top_p: %.2f", topP)
169 }
170
171 // Normalize origin value for Kiro API compatibility
172 origin = normalizeOrigin(origin)
173 log.Debugf("kiro-openai: normalized origin value: %s", origin)
174
175 messages := gjson.GetBytes(openaiBody, "messages")
176
177 // For chat-only mode, don't include tools
178 var tools gjson.Result
179 if !isChatOnly {
180 tools = gjson.GetBytes(openaiBody, "tools")
181 }
182
183 // Extract system prompt from messages
184 systemPrompt := extractSystemPromptFromOpenAI(messages)
185
186 // Inject timestamp context
187 timestamp := time.Now().Format("2006-01-02 15:04:05 MST")
188 timestampContext := fmt.Sprintf("[Context: Current time is %s]", timestamp)
189 if systemPrompt != "" {
190 systemPrompt = timestampContext + "\n\n" + systemPrompt
191 } else {
192 systemPrompt = timestampContext
193 }
194 log.Debugf("kiro-openai: injected timestamp context: %s", timestamp)
195
196 // Inject agentic optimization prompt for -agentic model variants
197 if isAgentic {
198 if systemPrompt != "" {
199 systemPrompt += "\n"
200 }
201 systemPrompt += kirocommon.KiroAgenticSystemPrompt
202 }
203
204 // Handle tool_choice parameter - Kiro doesn't support it natively, so we inject system prompt hints
205 // OpenAI tool_choice values: "none", "auto", "required", or {"type":"function","function":{"name":"..."}}
206 toolChoiceHint := extractToolChoiceHint(openaiBody)
207 if toolChoiceHint != "" {
208 if systemPrompt != "" {
209 systemPrompt += "\n"
210 }
211 systemPrompt += toolChoiceHint
212 log.Debugf("kiro-openai: injected tool_choice hint into system prompt")
213 }
214
215 // Handle response_format parameter - Kiro doesn't support it natively, so we inject system prompt hints
216 // OpenAI response_format: {"type": "json_object"} or {"type": "json_schema", "json_schema": {...}}
217 responseFormatHint := extractResponseFormatHint(openaiBody)
218 if responseFormatHint != "" {
219 if systemPrompt != "" {
220 systemPrompt += "\n"
221 }
222 systemPrompt += responseFormatHint
223 log.Debugf("kiro-openai: injected response_format hint into system prompt")
224 }
225
226 // Check for thinking mode
227 // Supports OpenAI reasoning_effort parameter, model name hints, and Anthropic-Beta header
228 thinkingEnabled := checkThinkingModeFromOpenAIWithHeaders(openaiBody, headers)
229
230 // Convert OpenAI tools to Kiro format
231 kiroTools := convertOpenAIToolsToKiro(tools)
232
233 // Thinking mode implementation:
234 // Kiro API supports official thinking/reasoning mode via <thinking_mode> tag.
235 // When set to "enabled", Kiro returns reasoning content as official reasoningContentEvent
236 // rather than inline <thinking> tags in assistantResponseEvent.
237 // We use a high max_thinking_length to allow extensive reasoning.
238 if thinkingEnabled {
239 thinkingHint := `<thinking_mode>enabled</thinking_mode>
240<max_thinking_length>200000</max_thinking_length>`
241 if systemPrompt != "" {
242 systemPrompt = thinkingHint + "\n\n" + systemPrompt
243 } else {
244 systemPrompt = thinkingHint
245 }
246 log.Debugf("kiro-openai: injected thinking prompt (official mode)")
247 }
248
249 // Process messages and build history
250 history, currentUserMsg, currentToolResults := processOpenAIMessages(messages, modelID, origin)
251
252 // Build content with system prompt
253 if currentUserMsg != nil {
254 currentUserMsg.Content = buildFinalContent(currentUserMsg.Content, systemPrompt, currentToolResults)
255
256 // Deduplicate currentToolResults
257 currentToolResults = deduplicateToolResults(currentToolResults)
258
259 // Build userInputMessageContext with tools and tool results
260 if len(kiroTools) > 0 || len(currentToolResults) > 0 {
261 currentUserMsg.UserInputMessageContext = &KiroUserInputMessageContext{
262 Tools: kiroTools,
263 ToolResults: currentToolResults,
264 }
265 }
266 }
267
268 // Build payload
269 var currentMessage KiroCurrentMessage
270 if currentUserMsg != nil {
271 currentMessage = KiroCurrentMessage{UserInputMessage: *currentUserMsg}
272 } else {
273 fallbackContent := ""
274 if systemPrompt != "" {
275 fallbackContent = "--- SYSTEM PROMPT ---\n" + systemPrompt + "\n--- END SYSTEM PROMPT ---\n"
276 }
277 currentMessage = KiroCurrentMessage{UserInputMessage: KiroUserInputMessage{
278 Content: fallbackContent,
279 ModelID: modelID,
280 Origin: origin,
281 }}
282 }
283
284 // Build inferenceConfig if we have any inference parameters
285 // Note: Kiro API doesn't actually use max_tokens for thinking budget
286 var inferenceConfig *KiroInferenceConfig
287 if maxTokens > 0 || hasTemperature || hasTopP {
288 inferenceConfig = &KiroInferenceConfig{}
289 if maxTokens > 0 {
290 inferenceConfig.MaxTokens = int(maxTokens)
291 }
292 if hasTemperature {
293 inferenceConfig.Temperature = temperature
294 }
295 if hasTopP {
296 inferenceConfig.TopP = topP
297 }
298 }
299
300 payload := KiroPayload{
301 ConversationState: KiroConversationState{
302 ChatTriggerType: "MANUAL",
303 ConversationID: uuid.New().String(),
304 CurrentMessage: currentMessage,
305 History: history,
306 },
307 ProfileArn: profileArn,
308 InferenceConfig: inferenceConfig,
309 }
310
311 result, err := json.Marshal(payload)
312 if err != nil {
313 log.Debugf("kiro-openai: failed to marshal payload: %v", err)
314 return nil, false
315 }
316
317 return result, thinkingEnabled
318}
319
320// normalizeOrigin normalizes origin value for Kiro API compatibility
321func normalizeOrigin(origin string) string {
322 switch origin {
323 case "KIRO_CLI":
324 return "CLI"
325 case "KIRO_AI_EDITOR":
326 return "AI_EDITOR"
327 case "AMAZON_Q":
328 return "CLI"
329 case "KIRO_IDE":
330 return "AI_EDITOR"
331 default:
332 return origin
333 }
334}
335
336// extractSystemPromptFromOpenAI extracts system prompt from OpenAI messages
337func extractSystemPromptFromOpenAI(messages gjson.Result) string {
338 if !messages.IsArray() {
339 return ""
340 }
341
342 var systemParts []string
343 for _, msg := range messages.Array() {
344 if msg.Get("role").String() == "system" {
345 content := msg.Get("content")
346 if content.Type == gjson.String {
347 systemParts = append(systemParts, content.String())
348 } else if content.IsArray() {
349 // Handle array content format
350 for _, part := range content.Array() {
351 if part.Get("type").String() == "text" {
352 systemParts = append(systemParts, part.Get("text").String())
353 }
354 }
355 }
356 }
357 }
358
359 return strings.Join(systemParts, "\n")
360}
361
362// shortenToolNameIfNeeded shortens tool names that exceed 64 characters.
363// MCP tools often have long names like "mcp__server-name__tool-name".
364// This preserves the "mcp__" prefix and last segment when possible.
365func shortenToolNameIfNeeded(name string) string {
366 const limit = 64
367 if len(name) <= limit {
368 return name
369 }
370 // For MCP tools, try to preserve prefix and last segment
371 if strings.HasPrefix(name, "mcp__") {
372 idx := strings.LastIndex(name, "__")
373 if idx > 0 {
374 cand := "mcp__" + name[idx+2:]
375 if len(cand) > limit {
376 return cand[:limit]
377 }
378 return cand
379 }
380 }
381 return name[:limit]
382}
383
384// convertOpenAIToolsToKiro converts OpenAI tools to Kiro format
385func convertOpenAIToolsToKiro(tools gjson.Result) []KiroToolWrapper {
386 var kiroTools []KiroToolWrapper
387 if !tools.IsArray() {
388 return kiroTools
389 }
390
391 for _, tool := range tools.Array() {
392 // OpenAI tools have type "function" with function definition inside
393 if tool.Get("type").String() != "function" {
394 continue
395 }
396
397 fn := tool.Get("function")
398 if !fn.Exists() {
399 continue
400 }
401
402 name := fn.Get("name").String()
403 description := fn.Get("description").String()
404 parameters := fn.Get("parameters").Value()
405
406 // Shorten tool name if it exceeds 64 characters (common with MCP tools)
407 originalName := name
408 name = shortenToolNameIfNeeded(name)
409 if name != originalName {
410 log.Debugf("kiro-openai: shortened tool name from '%s' to '%s'", originalName, name)
411 }
412
413 // CRITICAL FIX: Kiro API requires non-empty description
414 if strings.TrimSpace(description) == "" {
415 description = fmt.Sprintf("Tool: %s", name)
416 log.Debugf("kiro-openai: tool '%s' has empty description, using default: %s", name, description)
417 }
418
419 // Truncate long descriptions
420 if len(description) > kirocommon.KiroMaxToolDescLen {
421 truncLen := kirocommon.KiroMaxToolDescLen - 30
422 for truncLen > 0 && !utf8.RuneStart(description[truncLen]) {
423 truncLen--
424 }
425 description = description[:truncLen] + "... (description truncated)"
426 }
427
428 kiroTools = append(kiroTools, KiroToolWrapper{
429 ToolSpecification: KiroToolSpecification{
430 Name: name,
431 Description: description,
432 InputSchema: KiroInputSchema{JSON: parameters},
433 },
434 })
435 }
436
437 return kiroTools
438}
439
440// processOpenAIMessages processes OpenAI messages and builds Kiro history
441func processOpenAIMessages(messages gjson.Result, modelID, origin string) ([]KiroHistoryMessage, *KiroUserInputMessage, []KiroToolResult) {
442 var history []KiroHistoryMessage
443 var currentUserMsg *KiroUserInputMessage
444 var currentToolResults []KiroToolResult
445
446 if !messages.IsArray() {
447 return history, currentUserMsg, currentToolResults
448 }
449
450 // Merge adjacent messages with the same role
451 messagesArray := kirocommon.MergeAdjacentMessages(messages.Array())
452
453 // Track pending tool results that should be attached to the next user message
454 // This is critical for LiteLLM-translated requests where tool results appear
455 // as separate "tool" role messages between assistant and user messages
456 var pendingToolResults []KiroToolResult
457
458 for i, msg := range messagesArray {
459 role := msg.Get("role").String()
460 isLastMessage := i == len(messagesArray)-1
461
462 switch role {
463 case "system":
464 // System messages are handled separately via extractSystemPromptFromOpenAI
465 continue
466
467 case "user":
468 userMsg, toolResults := buildUserMessageFromOpenAI(msg, modelID, origin)
469 // Merge any pending tool results from preceding "tool" role messages
470 toolResults = append(pendingToolResults, toolResults...)
471 pendingToolResults = nil // Reset pending tool results
472
473 if isLastMessage {
474 currentUserMsg = &userMsg
475 currentToolResults = toolResults
476 } else {
477 // CRITICAL: Kiro API requires content to be non-empty for history messages
478 if strings.TrimSpace(userMsg.Content) == "" {
479 if len(toolResults) > 0 {
480 userMsg.Content = "Tool results provided."
481 } else {
482 userMsg.Content = "Continue"
483 }
484 }
485 // For history messages, embed tool results in context
486 if len(toolResults) > 0 {
487 userMsg.UserInputMessageContext = &KiroUserInputMessageContext{
488 ToolResults: toolResults,
489 }
490 }
491 history = append(history, KiroHistoryMessage{
492 UserInputMessage: &userMsg,
493 })
494 }
495
496 case "assistant":
497 assistantMsg := buildAssistantMessageFromOpenAI(msg)
498
499 // If there are pending tool results, we need to insert a synthetic user message
500 // before this assistant message to maintain proper conversation structure
501 if len(pendingToolResults) > 0 {
502 syntheticUserMsg := KiroUserInputMessage{
503 Content: "Tool results provided.",
504 ModelID: modelID,
505 Origin: origin,
506 UserInputMessageContext: &KiroUserInputMessageContext{
507 ToolResults: pendingToolResults,
508 },
509 }
510 history = append(history, KiroHistoryMessage{
511 UserInputMessage: &syntheticUserMsg,
512 })
513 pendingToolResults = nil
514 }
515
516 if isLastMessage {
517 history = append(history, KiroHistoryMessage{
518 AssistantResponseMessage: &assistantMsg,
519 })
520 // Create a "Continue" user message as currentMessage
521 currentUserMsg = &KiroUserInputMessage{
522 Content: "Continue",
523 ModelID: modelID,
524 Origin: origin,
525 }
526 } else {
527 history = append(history, KiroHistoryMessage{
528 AssistantResponseMessage: &assistantMsg,
529 })
530 }
531
532 case "tool":
533 // Tool messages in OpenAI format provide results for tool_calls
534 // These are typically followed by user or assistant messages
535 // Collect them as pending and attach to the next user message
536 toolCallID := msg.Get("tool_call_id").String()
537 content := msg.Get("content").String()
538
539 if toolCallID != "" {
540 toolResult := KiroToolResult{
541 ToolUseID: toolCallID,
542 Content: []KiroTextContent{{Text: content}},
543 Status: "success",
544 }
545 // Collect pending tool results to attach to the next user message
546 pendingToolResults = append(pendingToolResults, toolResult)
547 }
548 }
549 }
550
551 // Handle case where tool results are at the end with no following user message
552 if len(pendingToolResults) > 0 {
553 currentToolResults = append(currentToolResults, pendingToolResults...)
554 // If there's no current user message, create a synthetic one for the tool results
555 if currentUserMsg == nil {
556 currentUserMsg = &KiroUserInputMessage{
557 Content: "Tool results provided.",
558 ModelID: modelID,
559 Origin: origin,
560 }
561 }
562 }
563
564 return history, currentUserMsg, currentToolResults
565}
566
567// buildUserMessageFromOpenAI builds a user message from OpenAI format and extracts tool results
568func buildUserMessageFromOpenAI(msg gjson.Result, modelID, origin string) (KiroUserInputMessage, []KiroToolResult) {
569 content := msg.Get("content")
570 var contentBuilder strings.Builder
571 var toolResults []KiroToolResult
572 var images []KiroImage
573
574 if content.IsArray() {
575 for _, part := range content.Array() {
576 partType := part.Get("type").String()
577 switch partType {
578 case "text":
579 contentBuilder.WriteString(part.Get("text").String())
580 case "image_url":
581 imageURL := part.Get("image_url.url").String()
582 if strings.HasPrefix(imageURL, "data:") {
583 // Parse data URL: data:image/png;base64,xxxxx
584 if idx := strings.Index(imageURL, ";base64,"); idx != -1 {
585 mediaType := imageURL[5:idx] // Skip "data:"
586 data := imageURL[idx+8:] // Skip ";base64,"
587
588 format := ""
589 if lastSlash := strings.LastIndex(mediaType, "/"); lastSlash != -1 {
590 format = mediaType[lastSlash+1:]
591 }
592
593 if format != "" && data != "" {
594 images = append(images, KiroImage{
595 Format: format,
596 Source: KiroImageSource{
597 Bytes: data,
598 },
599 })
600 }
601 }
602 }
603 }
604 }
605 } else if content.Type == gjson.String {
606 contentBuilder.WriteString(content.String())
607 }
608
609 userMsg := KiroUserInputMessage{
610 Content: contentBuilder.String(),
611 ModelID: modelID,
612 Origin: origin,
613 }
614
615 if len(images) > 0 {
616 userMsg.Images = images
617 }
618
619 return userMsg, toolResults
620}
621
622// buildAssistantMessageFromOpenAI builds an assistant message from OpenAI format
623func buildAssistantMessageFromOpenAI(msg gjson.Result) KiroAssistantResponseMessage {
624 content := msg.Get("content")
625 var contentBuilder strings.Builder
626 var toolUses []KiroToolUse
627
628 // Handle content
629 if content.Type == gjson.String {
630 contentBuilder.WriteString(content.String())
631 } else if content.IsArray() {
632 for _, part := range content.Array() {
633 if part.Get("type").String() == "text" {
634 contentBuilder.WriteString(part.Get("text").String())
635 }
636 }
637 }
638
639 // Handle tool_calls
640 toolCalls := msg.Get("tool_calls")
641 if toolCalls.IsArray() {
642 for _, tc := range toolCalls.Array() {
643 if tc.Get("type").String() != "function" {
644 continue
645 }
646
647 toolUseID := tc.Get("id").String()
648 toolName := tc.Get("function.name").String()
649 toolArgs := tc.Get("function.arguments").String()
650
651 var inputMap map[string]interface{}
652 if err := json.Unmarshal([]byte(toolArgs), &inputMap); err != nil {
653 log.Debugf("kiro-openai: failed to parse tool arguments: %v", err)
654 inputMap = make(map[string]interface{})
655 }
656
657 toolUses = append(toolUses, KiroToolUse{
658 ToolUseID: toolUseID,
659 Name: toolName,
660 Input: inputMap,
661 })
662 }
663 }
664
665 return KiroAssistantResponseMessage{
666 Content: contentBuilder.String(),
667 ToolUses: toolUses,
668 }
669}
670
671// buildFinalContent builds the final content with system prompt
672func buildFinalContent(content, systemPrompt string, toolResults []KiroToolResult) string {
673 var contentBuilder strings.Builder
674
675 if systemPrompt != "" {
676 contentBuilder.WriteString("--- SYSTEM PROMPT ---\n")
677 contentBuilder.WriteString(systemPrompt)
678 contentBuilder.WriteString("\n--- END SYSTEM PROMPT ---\n\n")
679 }
680
681 contentBuilder.WriteString(content)
682 finalContent := contentBuilder.String()
683
684 // CRITICAL: Kiro API requires content to be non-empty
685 if strings.TrimSpace(finalContent) == "" {
686 if len(toolResults) > 0 {
687 finalContent = "Tool results provided."
688 } else {
689 finalContent = "Continue"
690 }
691 log.Debugf("kiro-openai: content was empty, using default: %s", finalContent)
692 }
693
694 return finalContent
695}
696
697// checkThinkingModeFromOpenAI checks if thinking mode is enabled in the OpenAI request.
698// Returns thinkingEnabled.
699// Supports:
700// - reasoning_effort parameter (low/medium/high/auto)
701// - Model name containing "thinking" or "reason"
702// - <thinking_mode> tag in system prompt (AMP/Cursor format)
703func checkThinkingModeFromOpenAI(openaiBody []byte) bool {
704 return checkThinkingModeFromOpenAIWithHeaders(openaiBody, nil)
705}
706
707// checkThinkingModeFromOpenAIWithHeaders checks if thinking mode is enabled in the OpenAI request.
708// Returns thinkingEnabled.
709// Supports:
710// - Anthropic-Beta header with interleaved-thinking (Claude CLI)
711// - reasoning_effort parameter (low/medium/high/auto)
712// - Model name containing "thinking" or "reason"
713// - <thinking_mode> tag in system prompt (AMP/Cursor format)
714func checkThinkingModeFromOpenAIWithHeaders(openaiBody []byte, headers http.Header) bool {
715 // Check Anthropic-Beta header first (Claude CLI uses this)
716 if kiroclaude.IsThinkingEnabledFromHeader(headers) {
717 log.Debugf("kiro-openai: thinking mode enabled via Anthropic-Beta header")
718 return true
719 }
720
721 // Check OpenAI format: reasoning_effort parameter
722 // Valid values: "low", "medium", "high", "auto" (not "none")
723 reasoningEffort := gjson.GetBytes(openaiBody, "reasoning_effort")
724 if reasoningEffort.Exists() {
725 effort := reasoningEffort.String()
726 if effort != "" && effort != "none" {
727 log.Debugf("kiro-openai: thinking mode enabled via reasoning_effort: %s", effort)
728 return true
729 }
730 }
731
732 // Check AMP/Cursor format: <thinking_mode>interleaved</thinking_mode> in system prompt
733 bodyStr := string(openaiBody)
734 if strings.Contains(bodyStr, "<thinking_mode>") && strings.Contains(bodyStr, "</thinking_mode>") {
735 startTag := "<thinking_mode>"
736 endTag := "</thinking_mode>"
737 startIdx := strings.Index(bodyStr, startTag)
738 if startIdx >= 0 {
739 startIdx += len(startTag)
740 endIdx := strings.Index(bodyStr[startIdx:], endTag)
741 if endIdx >= 0 {
742 thinkingMode := bodyStr[startIdx : startIdx+endIdx]
743 if thinkingMode == "interleaved" || thinkingMode == "enabled" {
744 log.Debugf("kiro-openai: thinking mode enabled via AMP/Cursor format: %s", thinkingMode)
745 return true
746 }
747 }
748 }
749 }
750
751 // Check model name for thinking hints
752 model := gjson.GetBytes(openaiBody, "model").String()
753 modelLower := strings.ToLower(model)
754 if strings.Contains(modelLower, "thinking") || strings.Contains(modelLower, "-reason") {
755 log.Debugf("kiro-openai: thinking mode enabled via model name hint: %s", model)
756 return true
757 }
758
759 log.Debugf("kiro-openai: no thinking mode detected in OpenAI request")
760 return false
761}
762
763// hasThinkingTagInBody checks if the request body already contains thinking configuration tags.
764// This is used to prevent duplicate injection when client (e.g., AMP/Cursor) already includes thinking config.
765func hasThinkingTagInBody(body []byte) bool {
766 bodyStr := string(body)
767 return strings.Contains(bodyStr, "<thinking_mode>") || strings.Contains(bodyStr, "<max_thinking_length>")
768}
769
770
771// extractToolChoiceHint extracts tool_choice from OpenAI request and returns a system prompt hint.
772// OpenAI tool_choice values:
773// - "none": Don't use any tools
774// - "auto": Model decides (default, no hint needed)
775// - "required": Must use at least one tool
776// - {"type":"function","function":{"name":"..."}} : Must use specific tool
777func extractToolChoiceHint(openaiBody []byte) string {
778 toolChoice := gjson.GetBytes(openaiBody, "tool_choice")
779 if !toolChoice.Exists() {
780 return ""
781 }
782
783 // Handle string values
784 if toolChoice.Type == gjson.String {
785 switch toolChoice.String() {
786 case "none":
787 // Note: When tool_choice is "none", we should ideally not pass tools at all
788 // But since we can't modify tool passing here, we add a strong hint
789 return "[INSTRUCTION: Do NOT use any tools. Respond with text only.]"
790 case "required":
791 return "[INSTRUCTION: You MUST use at least one of the available tools to respond. Do not respond with text only - always make a tool call.]"
792 case "auto":
793 // Default behavior, no hint needed
794 return ""
795 }
796 }
797
798 // Handle object value: {"type":"function","function":{"name":"..."}}
799 if toolChoice.IsObject() {
800 if toolChoice.Get("type").String() == "function" {
801 toolName := toolChoice.Get("function.name").String()
802 if toolName != "" {
803 return fmt.Sprintf("[INSTRUCTION: You MUST use the tool named '%s' to respond. Do not use any other tool or respond with text only.]", toolName)
804 }
805 }
806 }
807
808 return ""
809}
810
811// extractResponseFormatHint extracts response_format from OpenAI request and returns a system prompt hint.
812// OpenAI response_format values:
813// - {"type": "text"}: Default, no hint needed
814// - {"type": "json_object"}: Must respond with valid JSON
815// - {"type": "json_schema", "json_schema": {...}}: Must respond with JSON matching schema
816func extractResponseFormatHint(openaiBody []byte) string {
817 responseFormat := gjson.GetBytes(openaiBody, "response_format")
818 if !responseFormat.Exists() {
819 return ""
820 }
821
822 formatType := responseFormat.Get("type").String()
823 switch formatType {
824 case "json_object":
825 return "[INSTRUCTION: You MUST respond with valid JSON only. Do not include any text before or after the JSON. Do not wrap the JSON in markdown code blocks. Output raw JSON directly.]"
826 case "json_schema":
827 // Extract schema if provided
828 schema := responseFormat.Get("json_schema.schema")
829 if schema.Exists() {
830 schemaStr := schema.Raw
831 // Truncate if too long
832 if len(schemaStr) > 500 {
833 schemaStr = schemaStr[:500] + "..."
834 }
835 return fmt.Sprintf("[INSTRUCTION: You MUST respond with valid JSON that matches this schema: %s. Do not include any text before or after the JSON. Do not wrap the JSON in markdown code blocks. Output raw JSON directly.]", schemaStr)
836 }
837 return "[INSTRUCTION: You MUST respond with valid JSON only. Do not include any text before or after the JSON. Do not wrap the JSON in markdown code blocks. Output raw JSON directly.]"
838 case "text":
839 // Default behavior, no hint needed
840 return ""
841 }
842
843 return ""
844}
845
846// deduplicateToolResults removes duplicate tool results
847func deduplicateToolResults(toolResults []KiroToolResult) []KiroToolResult {
848 if len(toolResults) == 0 {
849 return toolResults
850 }
851
852 seenIDs := make(map[string]bool)
853 unique := make([]KiroToolResult, 0, len(toolResults))
854 for _, tr := range toolResults {
855 if !seenIDs[tr.ToolUseID] {
856 seenIDs[tr.ToolUseID] = true
857 unique = append(unique, tr)
858 } else {
859 log.Debugf("kiro-openai: skipping duplicate toolResult: %s", tr.ToolUseID)
860 }
861 }
862 return unique
863}
internal/translator/kiro/openai/kiro_openai_request_test.go+386−0
@@ -0,0 +1,386 @@
1package openai
2
3import (
4 "encoding/json"
5 "testing"
6)
7
8// TestToolResultsAttachedToCurrentMessage verifies that tool results from "tool" role messages
9// are properly attached to the current user message (the last message in the conversation).
10// This is critical for LiteLLM-translated requests where tool results appear as separate messages.
11func TestToolResultsAttachedToCurrentMessage(t *testing.T) {
12 // OpenAI format request simulating LiteLLM's translation from Anthropic format
13 // Sequence: user -> assistant (with tool_calls) -> tool (result) -> user
14 // The last user message should have the tool results attached
15 input := []byte(`{
16 "model": "kiro-claude-opus-4-5-agentic",
17 "messages": [
18 {"role": "user", "content": "Hello, can you read a file for me?"},
19 {
20 "role": "assistant",
21 "content": "I'll read that file for you.",
22 "tool_calls": [
23 {
24 "id": "call_abc123",
25 "type": "function",
26 "function": {
27 "name": "Read",
28 "arguments": "{\"file_path\": \"/tmp/test.txt\"}"
29 }
30 }
31 ]
32 },
33 {
34 "role": "tool",
35 "tool_call_id": "call_abc123",
36 "content": "File contents: Hello World!"
37 },
38 {"role": "user", "content": "What did the file say?"}
39 ]
40 }`)
41
42 result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
43
44 var payload KiroPayload
45 if err := json.Unmarshal(result, &payload); err != nil {
46 t.Fatalf("Failed to unmarshal result: %v", err)
47 }
48
49 // The last user message becomes currentMessage
50 // History should have: user (first), assistant (with tool_calls)
51 t.Logf("History count: %d", len(payload.ConversationState.History))
52 if len(payload.ConversationState.History) != 2 {
53 t.Errorf("Expected 2 history entries (user + assistant), got %d", len(payload.ConversationState.History))
54 }
55
56 // Tool results should be attached to currentMessage (the last user message)
57 ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext
58 if ctx == nil {
59 t.Fatal("Expected currentMessage to have UserInputMessageContext with tool results")
60 }
61
62 if len(ctx.ToolResults) != 1 {
63 t.Fatalf("Expected 1 tool result in currentMessage, got %d", len(ctx.ToolResults))
64 }
65
66 tr := ctx.ToolResults[0]
67 if tr.ToolUseID != "call_abc123" {
68 t.Errorf("Expected toolUseId 'call_abc123', got '%s'", tr.ToolUseID)
69 }
70 if len(tr.Content) == 0 || tr.Content[0].Text != "File contents: Hello World!" {
71 t.Errorf("Tool result content mismatch, got: %+v", tr.Content)
72 }
73}
74
75// TestToolResultsInHistoryUserMessage verifies that when there are multiple user messages
76// after tool results, the tool results are attached to the correct user message in history.
77func TestToolResultsInHistoryUserMessage(t *testing.T) {
78 // Sequence: user -> assistant (with tool_calls) -> tool (result) -> user -> assistant -> user
79 // The first user after tool should have tool results in history
80 input := []byte(`{
81 "model": "kiro-claude-opus-4-5-agentic",
82 "messages": [
83 {"role": "user", "content": "Hello"},
84 {
85 "role": "assistant",
86 "content": "I'll read the file.",
87 "tool_calls": [
88 {
89 "id": "call_1",
90 "type": "function",
91 "function": {
92 "name": "Read",
93 "arguments": "{}"
94 }
95 }
96 ]
97 },
98 {
99 "role": "tool",
100 "tool_call_id": "call_1",
101 "content": "File result"
102 },
103 {"role": "user", "content": "Thanks for the file"},
104 {"role": "assistant", "content": "You're welcome"},
105 {"role": "user", "content": "Bye"}
106 ]
107 }`)
108
109 result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
110
111 var payload KiroPayload
112 if err := json.Unmarshal(result, &payload); err != nil {
113 t.Fatalf("Failed to unmarshal result: %v", err)
114 }
115
116 // History should have: user, assistant, user (with tool results), assistant
117 // CurrentMessage should be: last user "Bye"
118 t.Logf("History count: %d", len(payload.ConversationState.History))
119
120 // Find the user message in history with tool results
121 foundToolResults := false
122 for i, h := range payload.ConversationState.History {
123 if h.UserInputMessage != nil {
124 t.Logf("History[%d]: user message content=%q", i, h.UserInputMessage.Content)
125 if h.UserInputMessage.UserInputMessageContext != nil {
126 if len(h.UserInputMessage.UserInputMessageContext.ToolResults) > 0 {
127 foundToolResults = true
128 t.Logf(" Found %d tool results", len(h.UserInputMessage.UserInputMessageContext.ToolResults))
129 tr := h.UserInputMessage.UserInputMessageContext.ToolResults[0]
130 if tr.ToolUseID != "call_1" {
131 t.Errorf("Expected toolUseId 'call_1', got '%s'", tr.ToolUseID)
132 }
133 }
134 }
135 }
136 if h.AssistantResponseMessage != nil {
137 t.Logf("History[%d]: assistant message content=%q", i, h.AssistantResponseMessage.Content)
138 }
139 }
140
141 if !foundToolResults {
142 t.Error("Tool results were not attached to any user message in history")
143 }
144}
145
146// TestToolResultsWithMultipleToolCalls verifies handling of multiple tool calls
147func TestToolResultsWithMultipleToolCalls(t *testing.T) {
148 input := []byte(`{
149 "model": "kiro-claude-opus-4-5-agentic",
150 "messages": [
151 {"role": "user", "content": "Read two files for me"},
152 {
153 "role": "assistant",
154 "content": "I'll read both files.",
155 "tool_calls": [
156 {
157 "id": "call_1",
158 "type": "function",
159 "function": {
160 "name": "Read",
161 "arguments": "{\"file_path\": \"/tmp/file1.txt\"}"
162 }
163 },
164 {
165 "id": "call_2",
166 "type": "function",
167 "function": {
168 "name": "Read",
169 "arguments": "{\"file_path\": \"/tmp/file2.txt\"}"
170 }
171 }
172 ]
173 },
174 {
175 "role": "tool",
176 "tool_call_id": "call_1",
177 "content": "Content of file 1"
178 },
179 {
180 "role": "tool",
181 "tool_call_id": "call_2",
182 "content": "Content of file 2"
183 },
184 {"role": "user", "content": "What do they say?"}
185 ]
186 }`)
187
188 result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
189
190 var payload KiroPayload
191 if err := json.Unmarshal(result, &payload); err != nil {
192 t.Fatalf("Failed to unmarshal result: %v", err)
193 }
194
195 t.Logf("History count: %d", len(payload.ConversationState.History))
196 t.Logf("CurrentMessage content: %q", payload.ConversationState.CurrentMessage.UserInputMessage.Content)
197
198 // Check if there are any tool results anywhere
199 var totalToolResults int
200 for i, h := range payload.ConversationState.History {
201 if h.UserInputMessage != nil && h.UserInputMessage.UserInputMessageContext != nil {
202 count := len(h.UserInputMessage.UserInputMessageContext.ToolResults)
203 t.Logf("History[%d] user message has %d tool results", i, count)
204 totalToolResults += count
205 }
206 }
207
208 ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext
209 if ctx != nil {
210 t.Logf("CurrentMessage has %d tool results", len(ctx.ToolResults))
211 totalToolResults += len(ctx.ToolResults)
212 } else {
213 t.Logf("CurrentMessage has no UserInputMessageContext")
214 }
215
216 if totalToolResults != 2 {
217 t.Errorf("Expected 2 tool results total, got %d", totalToolResults)
218 }
219}
220
221// TestToolResultsAtEndOfConversation verifies tool results are handled when
222// the conversation ends with tool results (no following user message)
223func TestToolResultsAtEndOfConversation(t *testing.T) {
224 input := []byte(`{
225 "model": "kiro-claude-opus-4-5-agentic",
226 "messages": [
227 {"role": "user", "content": "Read a file"},
228 {
229 "role": "assistant",
230 "content": "Reading the file.",
231 "tool_calls": [
232 {
233 "id": "call_end",
234 "type": "function",
235 "function": {
236 "name": "Read",
237 "arguments": "{\"file_path\": \"/tmp/test.txt\"}"
238 }
239 }
240 ]
241 },
242 {
243 "role": "tool",
244 "tool_call_id": "call_end",
245 "content": "File contents here"
246 }
247 ]
248 }`)
249
250 result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
251
252 var payload KiroPayload
253 if err := json.Unmarshal(result, &payload); err != nil {
254 t.Fatalf("Failed to unmarshal result: %v", err)
255 }
256
257 // When the last message is a tool result, a synthetic user message is created
258 // and tool results should be attached to it
259 ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext
260 if ctx == nil || len(ctx.ToolResults) == 0 {
261 t.Error("Expected tool results to be attached to current message when conversation ends with tool result")
262 } else {
263 if ctx.ToolResults[0].ToolUseID != "call_end" {
264 t.Errorf("Expected toolUseId 'call_end', got '%s'", ctx.ToolResults[0].ToolUseID)
265 }
266 }
267}
268
269// TestToolResultsFollowedByAssistant verifies handling when tool results are followed
270// by an assistant message (no intermediate user message).
271// This is the pattern from LiteLLM translation of Anthropic format where:
272// user message has ONLY tool_result blocks -> LiteLLM creates tool messages
273// then the next message is assistant
274func TestToolResultsFollowedByAssistant(t *testing.T) {
275 // Sequence: user -> assistant (with tool_calls) -> tool -> tool -> assistant -> user
276 // This simulates LiteLLM's translation of:
277 // user: "Read files"
278 // assistant: [tool_use, tool_use]
279 // user: [tool_result, tool_result] <- becomes multiple "tool" role messages
280 // assistant: "I've read them"
281 // user: "What did they say?"
282 input := []byte(`{
283 "model": "kiro-claude-opus-4-5-agentic",
284 "messages": [
285 {"role": "user", "content": "Read two files for me"},
286 {
287 "role": "assistant",
288 "content": "I'll read both files.",
289 "tool_calls": [
290 {
291 "id": "call_1",
292 "type": "function",
293 "function": {
294 "name": "Read",
295 "arguments": "{\"file_path\": \"/tmp/a.txt\"}"
296 }
297 },
298 {
299 "id": "call_2",
300 "type": "function",
301 "function": {
302 "name": "Read",
303 "arguments": "{\"file_path\": \"/tmp/b.txt\"}"
304 }
305 }
306 ]
307 },
308 {
309 "role": "tool",
310 "tool_call_id": "call_1",
311 "content": "Contents of file A"
312 },
313 {
314 "role": "tool",
315 "tool_call_id": "call_2",
316 "content": "Contents of file B"
317 },
318 {
319 "role": "assistant",
320 "content": "I've read both files."
321 },
322 {"role": "user", "content": "What did they say?"}
323 ]
324 }`)
325
326 result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
327
328 var payload KiroPayload
329 if err := json.Unmarshal(result, &payload); err != nil {
330 t.Fatalf("Failed to unmarshal result: %v", err)
331 }
332
333 t.Logf("History count: %d", len(payload.ConversationState.History))
334
335 // Tool results should be attached to a synthetic user message or the history should be valid
336 var totalToolResults int
337 for i, h := range payload.ConversationState.History {
338 if h.UserInputMessage != nil {
339 t.Logf("History[%d]: user message content=%q", i, h.UserInputMessage.Content)
340 if h.UserInputMessage.UserInputMessageContext != nil {
341 count := len(h.UserInputMessage.UserInputMessageContext.ToolResults)
342 t.Logf(" Has %d tool results", count)
343 totalToolResults += count
344 }
345 }
346 if h.AssistantResponseMessage != nil {
347 t.Logf("History[%d]: assistant message content=%q", i, h.AssistantResponseMessage.Content)
348 }
349 }
350
351 ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext
352 if ctx != nil {
353 t.Logf("CurrentMessage has %d tool results", len(ctx.ToolResults))
354 totalToolResults += len(ctx.ToolResults)
355 }
356
357 if totalToolResults != 2 {
358 t.Errorf("Expected 2 tool results total, got %d", totalToolResults)
359 }
360}
361
362// TestAssistantEndsConversation verifies handling when assistant is the last message
363func TestAssistantEndsConversation(t *testing.T) {
364 input := []byte(`{
365 "model": "kiro-claude-opus-4-5-agentic",
366 "messages": [
367 {"role": "user", "content": "Hello"},
368 {
369 "role": "assistant",
370 "content": "Hi there!"
371 }
372 ]
373 }`)
374
375 result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
376
377 var payload KiroPayload
378 if err := json.Unmarshal(result, &payload); err != nil {
379 t.Fatalf("Failed to unmarshal result: %v", err)
380 }
381
382 // When assistant is last, a "Continue" user message should be created
383 if payload.ConversationState.CurrentMessage.UserInputMessage.Content == "" {
384 t.Error("Expected a 'Continue' message to be created when assistant is last")
385 }
386}
internal/translator/kiro/openai/kiro_openai_response.go+277−0
@@ -0,0 +1,277 @@
1// Package openai provides response translation from Kiro to OpenAI format.
2// This package handles the conversion of Kiro API responses into OpenAI Chat Completions-compatible
3// JSON format, transforming streaming events and non-streaming responses.
4package openai
5
6import (
7 "encoding/json"
8 "fmt"
9 "sync/atomic"
10 "time"
11
12 "github.com/google/uuid"
13 "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
14 log "github.com/sirupsen/logrus"
15)
16
17// functionCallIDCounter provides a process-wide unique counter for function call identifiers.
18var functionCallIDCounter uint64
19
20// BuildOpenAIResponse constructs an OpenAI Chat Completions-compatible response.
21// Supports tool_calls when tools are present in the response.
22// stopReason is passed from upstream; fallback logic applied if empty.
23func BuildOpenAIResponse(content string, toolUses []KiroToolUse, model string, usageInfo usage.Detail, stopReason string) []byte {
24 return BuildOpenAIResponseWithReasoning(content, "", toolUses, model, usageInfo, stopReason)
25}
26
27// BuildOpenAIResponseWithReasoning constructs an OpenAI Chat Completions-compatible response with reasoning_content support.
28// Supports tool_calls when tools are present in the response.
29// reasoningContent is included as reasoning_content field in the message when present.
30// stopReason is passed from upstream; fallback logic applied if empty.
31func BuildOpenAIResponseWithReasoning(content, reasoningContent string, toolUses []KiroToolUse, model string, usageInfo usage.Detail, stopReason string) []byte {
32 // Build the message object
33 message := map[string]interface{}{
34 "role": "assistant",
35 "content": content,
36 }
37
38 // Add reasoning_content if present (for thinking/reasoning models)
39 if reasoningContent != "" {
40 message["reasoning_content"] = reasoningContent
41 }
42
43 // Add tool_calls if present
44 if len(toolUses) > 0 {
45 var toolCalls []map[string]interface{}
46 for i, tu := range toolUses {
47 inputJSON, _ := json.Marshal(tu.Input)
48 toolCalls = append(toolCalls, map[string]interface{}{
49 "id": tu.ToolUseID,
50 "type": "function",
51 "index": i,
52 "function": map[string]interface{}{
53 "name": tu.Name,
54 "arguments": string(inputJSON),
55 },
56 })
57 }
58 message["tool_calls"] = toolCalls
59 // When tool_calls are present, content should be null according to OpenAI spec
60 if content == "" {
61 message["content"] = nil
62 }
63 }
64
65 // Use upstream stopReason; apply fallback logic if not provided
66 finishReason := mapKiroStopReasonToOpenAI(stopReason)
67 if finishReason == "" {
68 finishReason = "stop"
69 if len(toolUses) > 0 {
70 finishReason = "tool_calls"
71 }
72 log.Debugf("kiro-openai: buildOpenAIResponse using fallback finish_reason: %s", finishReason)
73 }
74
75 response := map[string]interface{}{
76 "id": "chatcmpl-" + uuid.New().String()[:24],
77 "object": "chat.completion",
78 "created": time.Now().Unix(),
79 "model": model,
80 "choices": []map[string]interface{}{
81 {
82 "index": 0,
83 "message": message,
84 "finish_reason": finishReason,
85 },
86 },
87 "usage": map[string]interface{}{
88 "prompt_tokens": usageInfo.InputTokens,
89 "completion_tokens": usageInfo.OutputTokens,
90 "total_tokens": usageInfo.InputTokens + usageInfo.OutputTokens,
91 },
92 }
93
94 result, _ := json.Marshal(response)
95 return result
96}
97
98// mapKiroStopReasonToOpenAI converts Kiro/Claude stop_reason to OpenAI finish_reason
99func mapKiroStopReasonToOpenAI(stopReason string) string {
100 switch stopReason {
101 case "end_turn":
102 return "stop"
103 case "stop_sequence":
104 return "stop"
105 case "tool_use":
106 return "tool_calls"
107 case "max_tokens":
108 return "length"
109 case "content_filtered":
110 return "content_filter"
111 default:
112 return stopReason
113 }
114}
115
116// BuildOpenAIStreamChunk constructs an OpenAI Chat Completions streaming chunk.
117// This is the delta format used in streaming responses.
118func BuildOpenAIStreamChunk(model string, deltaContent string, deltaToolCalls []map[string]interface{}, finishReason string, index int) []byte {
119 delta := map[string]interface{}{}
120
121 // First chunk should include role
122 if index == 0 && deltaContent == "" && len(deltaToolCalls) == 0 {
123 delta["role"] = "assistant"
124 delta["content"] = ""
125 } else if deltaContent != "" {
126 delta["content"] = deltaContent
127 }
128
129 // Add tool_calls delta if present
130 if len(deltaToolCalls) > 0 {
131 delta["tool_calls"] = deltaToolCalls
132 }
133
134 choice := map[string]interface{}{
135 "index": 0,
136 "delta": delta,
137 }
138
139 if finishReason != "" {
140 choice["finish_reason"] = finishReason
141 } else {
142 choice["finish_reason"] = nil
143 }
144
145 chunk := map[string]interface{}{
146 "id": "chatcmpl-" + uuid.New().String()[:12],
147 "object": "chat.completion.chunk",
148 "created": time.Now().Unix(),
149 "model": model,
150 "choices": []map[string]interface{}{choice},
151 }
152
153 result, _ := json.Marshal(chunk)
154 return result
155}
156
157// BuildOpenAIStreamChunkWithToolCallStart creates a stream chunk for tool call start
158func BuildOpenAIStreamChunkWithToolCallStart(model string, toolUseID, toolName string, toolIndex int) []byte {
159 toolCall := map[string]interface{}{
160 "index": toolIndex,
161 "id": toolUseID,
162 "type": "function",
163 "function": map[string]interface{}{
164 "name": toolName,
165 "arguments": "",
166 },
167 }
168
169 delta := map[string]interface{}{
170 "tool_calls": []map[string]interface{}{toolCall},
171 }
172
173 choice := map[string]interface{}{
174 "index": 0,
175 "delta": delta,
176 "finish_reason": nil,
177 }
178
179 chunk := map[string]interface{}{
180 "id": "chatcmpl-" + uuid.New().String()[:12],
181 "object": "chat.completion.chunk",
182 "created": time.Now().Unix(),
183 "model": model,
184 "choices": []map[string]interface{}{choice},
185 }
186
187 result, _ := json.Marshal(chunk)
188 return result
189}
190
191// BuildOpenAIStreamChunkWithToolCallDelta creates a stream chunk for tool call arguments delta
192func BuildOpenAIStreamChunkWithToolCallDelta(model string, argumentsDelta string, toolIndex int) []byte {
193 toolCall := map[string]interface{}{
194 "index": toolIndex,
195 "function": map[string]interface{}{
196 "arguments": argumentsDelta,
197 },
198 }
199
200 delta := map[string]interface{}{
201 "tool_calls": []map[string]interface{}{toolCall},
202 }
203
204 choice := map[string]interface{}{
205 "index": 0,
206 "delta": delta,
207 "finish_reason": nil,
208 }
209
210 chunk := map[string]interface{}{
211 "id": "chatcmpl-" + uuid.New().String()[:12],
212 "object": "chat.completion.chunk",
213 "created": time.Now().Unix(),
214 "model": model,
215 "choices": []map[string]interface{}{choice},
216 }
217
218 result, _ := json.Marshal(chunk)
219 return result
220}
221
222// BuildOpenAIStreamDoneChunk creates the final [DONE] stream event
223func BuildOpenAIStreamDoneChunk() []byte {
224 return []byte("data: [DONE]")
225}
226
227// BuildOpenAIStreamFinishChunk creates the final chunk with finish_reason
228func BuildOpenAIStreamFinishChunk(model string, finishReason string) []byte {
229 choice := map[string]interface{}{
230 "index": 0,
231 "delta": map[string]interface{}{},
232 "finish_reason": finishReason,
233 }
234
235 chunk := map[string]interface{}{
236 "id": "chatcmpl-" + uuid.New().String()[:12],
237 "object": "chat.completion.chunk",
238 "created": time.Now().Unix(),
239 "model": model,
240 "choices": []map[string]interface{}{choice},
241 }
242
243 result, _ := json.Marshal(chunk)
244 return result
245}
246
247// BuildOpenAIStreamUsageChunk creates a chunk with usage information (optional, for stream_options.include_usage)
248func BuildOpenAIStreamUsageChunk(model string, usageInfo usage.Detail) []byte {
249 chunk := map[string]interface{}{
250 "id": "chatcmpl-" + uuid.New().String()[:12],
251 "object": "chat.completion.chunk",
252 "created": time.Now().Unix(),
253 "model": model,
254 "choices": []map[string]interface{}{},
255 "usage": map[string]interface{}{
256 "prompt_tokens": usageInfo.InputTokens,
257 "completion_tokens": usageInfo.OutputTokens,
258 "total_tokens": usageInfo.InputTokens + usageInfo.OutputTokens,
259 },
260 }
261
262 result, _ := json.Marshal(chunk)
263 return result
264}
265
266// GenerateToolCallID generates a unique tool call ID in OpenAI format
267func GenerateToolCallID(toolName string) string {
268 return fmt.Sprintf("call_%s_%d_%d", toolName[:min(8, len(toolName))], time.Now().UnixNano(), atomic.AddUint64(&functionCallIDCounter, 1))
269}
270
271// min returns the minimum of two integers
272func min(a, b int) int {
273 if a < b {
274 return a
275 }
276 return b
277}
internal/translator/kiro/openai/kiro_openai_stream.go+212−0
@@ -0,0 +1,212 @@
1// Package openai provides streaming SSE event building for OpenAI format.
2// This package handles the construction of OpenAI-compatible Server-Sent Events (SSE)
3// for streaming responses from Kiro API.
4package openai
5
6import (
7 "encoding/json"
8 "time"
9
10 "github.com/google/uuid"
11 "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
12)
13
14// OpenAIStreamState tracks the state of streaming response conversion
15type OpenAIStreamState struct {
16 ChunkIndex int
17 ToolCallIndex int
18 HasSentFirstChunk bool
19 Model string
20 ResponseID string
21 Created int64
22}
23
24// NewOpenAIStreamState creates a new stream state for tracking
25func NewOpenAIStreamState(model string) *OpenAIStreamState {
26 return &OpenAIStreamState{
27 ChunkIndex: 0,
28 ToolCallIndex: 0,
29 HasSentFirstChunk: false,
30 Model: model,
31 ResponseID: "chatcmpl-" + uuid.New().String()[:24],
32 Created: time.Now().Unix(),
33 }
34}
35
36// FormatSSEEvent formats a JSON payload for SSE streaming.
37// Note: This returns raw JSON data without "data:" prefix.
38// The SSE "data:" prefix is added by the Handler layer (e.g., openai_handlers.go)
39// to maintain architectural consistency and avoid double-prefix issues.
40func FormatSSEEvent(data []byte) string {
41 return string(data)
42}
43
44// BuildOpenAISSETextDelta creates an SSE event for text content delta
45func BuildOpenAISSETextDelta(state *OpenAIStreamState, textDelta string) string {
46 delta := map[string]interface{}{
47 "content": textDelta,
48 }
49
50 // Include role in first chunk
51 if !state.HasSentFirstChunk {
52 delta["role"] = "assistant"
53 state.HasSentFirstChunk = true
54 }
55
56 chunk := buildBaseChunk(state, delta, nil)
57 result, _ := json.Marshal(chunk)
58 state.ChunkIndex++
59 return FormatSSEEvent(result)
60}
61
62// BuildOpenAISSEToolCallStart creates an SSE event for tool call start
63func BuildOpenAISSEToolCallStart(state *OpenAIStreamState, toolUseID, toolName string) string {
64 toolCall := map[string]interface{}{
65 "index": state.ToolCallIndex,
66 "id": toolUseID,
67 "type": "function",
68 "function": map[string]interface{}{
69 "name": toolName,
70 "arguments": "",
71 },
72 }
73
74 delta := map[string]interface{}{
75 "tool_calls": []map[string]interface{}{toolCall},
76 }
77
78 // Include role in first chunk if not sent yet
79 if !state.HasSentFirstChunk {
80 delta["role"] = "assistant"
81 state.HasSentFirstChunk = true
82 }
83
84 chunk := buildBaseChunk(state, delta, nil)
85 result, _ := json.Marshal(chunk)
86 state.ChunkIndex++
87 return FormatSSEEvent(result)
88}
89
90// BuildOpenAISSEToolCallArgumentsDelta creates an SSE event for tool call arguments delta
91func BuildOpenAISSEToolCallArgumentsDelta(state *OpenAIStreamState, argumentsDelta string, toolIndex int) string {
92 toolCall := map[string]interface{}{
93 "index": toolIndex,
94 "function": map[string]interface{}{
95 "arguments": argumentsDelta,
96 },
97 }
98
99 delta := map[string]interface{}{
100 "tool_calls": []map[string]interface{}{toolCall},
101 }
102
103 chunk := buildBaseChunk(state, delta, nil)
104 result, _ := json.Marshal(chunk)
105 state.ChunkIndex++
106 return FormatSSEEvent(result)
107}
108
109// BuildOpenAISSEFinish creates an SSE event with finish_reason
110func BuildOpenAISSEFinish(state *OpenAIStreamState, finishReason string) string {
111 chunk := buildBaseChunk(state, map[string]interface{}{}, &finishReason)
112 result, _ := json.Marshal(chunk)
113 state.ChunkIndex++
114 return FormatSSEEvent(result)
115}
116
117// BuildOpenAISSEUsage creates an SSE event with usage information
118func BuildOpenAISSEUsage(state *OpenAIStreamState, usageInfo usage.Detail) string {
119 chunk := map[string]interface{}{
120 "id": state.ResponseID,
121 "object": "chat.completion.chunk",
122 "created": state.Created,
123 "model": state.Model,
124 "choices": []map[string]interface{}{},
125 "usage": map[string]interface{}{
126 "prompt_tokens": usageInfo.InputTokens,
127 "completion_tokens": usageInfo.OutputTokens,
128 "total_tokens": usageInfo.InputTokens + usageInfo.OutputTokens,
129 },
130 }
131 result, _ := json.Marshal(chunk)
132 return FormatSSEEvent(result)
133}
134
135// BuildOpenAISSEDone creates the final [DONE] SSE event.
136// Note: This returns raw "[DONE]" without "data:" prefix.
137// The SSE "data:" prefix is added by the Handler layer (e.g., openai_handlers.go)
138// to maintain architectural consistency and avoid double-prefix issues.
139func BuildOpenAISSEDone() string {
140 return "[DONE]"
141}
142
143// buildBaseChunk creates a base chunk structure for streaming
144func buildBaseChunk(state *OpenAIStreamState, delta map[string]interface{}, finishReason *string) map[string]interface{} {
145 choice := map[string]interface{}{
146 "index": 0,
147 "delta": delta,
148 }
149
150 if finishReason != nil {
151 choice["finish_reason"] = *finishReason
152 } else {
153 choice["finish_reason"] = nil
154 }
155
156 return map[string]interface{}{
157 "id": state.ResponseID,
158 "object": "chat.completion.chunk",
159 "created": state.Created,
160 "model": state.Model,
161 "choices": []map[string]interface{}{choice},
162 }
163}
164
165// BuildOpenAISSEReasoningDelta creates an SSE event for reasoning content delta
166// This is used for o1/o3 style models that expose reasoning tokens
167func BuildOpenAISSEReasoningDelta(state *OpenAIStreamState, reasoningDelta string) string {
168 delta := map[string]interface{}{
169 "reasoning_content": reasoningDelta,
170 }
171
172 // Include role in first chunk
173 if !state.HasSentFirstChunk {
174 delta["role"] = "assistant"
175 state.HasSentFirstChunk = true
176 }
177
178 chunk := buildBaseChunk(state, delta, nil)
179 result, _ := json.Marshal(chunk)
180 state.ChunkIndex++
181 return FormatSSEEvent(result)
182}
183
184// BuildOpenAISSEFirstChunk creates the first chunk with role only
185func BuildOpenAISSEFirstChunk(state *OpenAIStreamState) string {
186 delta := map[string]interface{}{
187 "role": "assistant",
188 "content": "",
189 }
190
191 state.HasSentFirstChunk = true
192 chunk := buildBaseChunk(state, delta, nil)
193 result, _ := json.Marshal(chunk)
194 state.ChunkIndex++
195 return FormatSSEEvent(result)
196}
197
198// ThinkingTagState tracks state for thinking tag detection in streaming
199type ThinkingTagState struct {
200 InThinkingBlock bool
201 PendingStartChars int
202 PendingEndChars int
203}
204
205// NewThinkingTagState creates a new thinking tag state
206func NewThinkingTagState() *ThinkingTagState {
207 return &ThinkingTagState{
208 InThinkingBlock: false,
209 PendingStartChars: 0,
210 PendingEndChars: 0,
211 }
212}
internal/watcher/events.go+67−1
@@ -13,6 +13,7 @@ import ( · 12 unmodified lines
13 "time"13 "time"
14 14
15 "github.com/fsnotify/fsnotify"15 "github.com/fsnotify/fsnotify"
16 kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"
16 log "github.com/sirupsen/logrus"17 log "github.com/sirupsen/logrus"
17)18)
18 19
@@ -39,12 +40,35 @@ func (w *Watcher) start(ctx context.Context) error { · 20 unmodified lines
39 }40 }
40 log.Debugf("watching auth directory: %s", w.authDir)41 log.Debugf("watching auth directory: %s", w.authDir)
41 42
43 w.watchKiroIDETokenFile()
44
42 go w.processEvents(ctx)45 go w.processEvents(ctx)
43 46
44 w.reloadClients(true, nil, false)47 w.reloadClients(true, nil, false)
45 return nil48 return nil
46}49}
47 50
51func (w *Watcher) watchKiroIDETokenFile() {
52 homeDir, err := os.UserHomeDir()
53 if err != nil {
54 log.Debugf("failed to get home directory for Kiro IDE token watch: %v", err)
55 return
56 }
57
58 kiroTokenDir := filepath.Join(homeDir, ".aws", "sso", "cache")
59
60 if _, statErr := os.Stat(kiroTokenDir); os.IsNotExist(statErr) {
61 log.Debugf("Kiro IDE token directory does not exist: %s", kiroTokenDir)
62 return
63 }
64
65 if errAdd := w.watcher.Add(kiroTokenDir); errAdd != nil {
66 log.Debugf("failed to watch Kiro IDE token directory %s: %v", kiroTokenDir, errAdd)
67 return
68 }
69 log.Debugf("watching Kiro IDE token directory: %s", kiroTokenDir)
70}
71
48func (w *Watcher) processEvents(ctx context.Context) {72func (w *Watcher) processEvents(ctx context.Context) {
49 for {73 for {
50 select {74 select {
@@ -73,11 +97,17 @@ func (w *Watcher) handleEvent(event fsnotify.Event) { · 22 unmodified lines
73 isConfigEvent := normalizedName == normalizedConfigPath && event.Op&configOps != 097 isConfigEvent := normalizedName == normalizedConfigPath && event.Op&configOps != 0
74 authOps := fsnotify.Create | fsnotify.Write | fsnotify.Remove | fsnotify.Rename98 authOps := fsnotify.Create | fsnotify.Write | fsnotify.Remove | fsnotify.Rename
75 isAuthJSON := strings.HasPrefix(normalizedName, normalizedAuthDir) && strings.HasSuffix(normalizedName, ".json") && event.Op&authOps != 099 isAuthJSON := strings.HasPrefix(normalizedName, normalizedAuthDir) && strings.HasSuffix(normalizedName, ".json") && event.Op&authOps != 0
76 if !isConfigEvent && !isAuthJSON {100 isKiroIDEToken := w.isKiroIDETokenFile(event.Name) && event.Op&authOps != 0
101 if !isConfigEvent && !isAuthJSON && !isKiroIDEToken {
77 // Ignore unrelated files (e.g., cookie snapshots *.cookie) and other noise.102 // Ignore unrelated files (e.g., cookie snapshots *.cookie) and other noise.
78 return103 return
79 }104 }
80 105
106 if isKiroIDEToken {
107 w.handleKiroIDETokenChange(event)
108 return
109 }
110
81 now := time.Now()111 now := time.Now()
82 log.Debugf("file system event detected: %s %s", event.Op.String(), event.Name)112 log.Debugf("file system event detected: %s %s", event.Op.String(), event.Name)
83 113
@@ -124,6 +154,42 @@ func (w *Watcher) handleEvent(event fsnotify.Event) { · 40 unmodified lines
124 }154 }
125}155}
126 156
157func (w *Watcher) isKiroIDETokenFile(path string) bool {
158 normalized := filepath.ToSlash(path)
159 return strings.HasSuffix(normalized, "kiro-auth-token.json") && strings.Contains(normalized, ".aws/sso/cache")
160}
161
162func (w *Watcher) handleKiroIDETokenChange(event fsnotify.Event) {
163 log.Debugf("Kiro IDE token file event detected: %s %s", event.Op.String(), event.Name)
164
165 if event.Op&(fsnotify.Remove|fsnotify.Rename) != 0 {
166 time.Sleep(replaceCheckDelay)
167 if _, statErr := os.Stat(event.Name); statErr != nil {
168 log.Debugf("Kiro IDE token file removed: %s", event.Name)
169 return
170 }
171 }
172
173 tokenData, err := kiroauth.LoadKiroIDEToken()
174 if err != nil {
175 log.Debugf("failed to load Kiro IDE token after change: %v", err)
176 return
177 }
178
179 log.Infof("Kiro IDE token file updated, access token refreshed (provider: %s)", tokenData.Provider)
180
181 w.refreshAuthState(true)
182
183 w.clientsMutex.RLock()
184 cfg := w.config
185 w.clientsMutex.RUnlock()
186
187 if w.reloadCallback != nil && cfg != nil {
188 log.Debugf("triggering server update callback after Kiro IDE token change")
189 w.reloadCallback(cfg)
190 }
191}
192
127func (w *Watcher) authFileUnchanged(path string) (bool, error) {193func (w *Watcher) authFileUnchanged(path string) (bool, error) {
128 data, errRead := os.ReadFile(path)194 data, errRead := os.ReadFile(path)
129 if errRead != nil {195 if errRead != nil {
internal/watcher/synthesizer/config.go+97−0
@@ -4,8 +4,10 @@ import ( · 3 unmodified lines
4 "fmt"4 "fmt"
5 "strings"5 "strings"
6 6
7 kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"
7 "github.com/router-for-me/CLIProxyAPI/v6/internal/watcher/diff"8 "github.com/router-for-me/CLIProxyAPI/v6/internal/watcher/diff"
8 coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"9 coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
10 log "github.com/sirupsen/logrus"
9)11)
10 12
11// ConfigSynthesizer generates Auth entries from configuration API keys.13// ConfigSynthesizer generates Auth entries from configuration API keys.
@@ -30,6 +32,8 @@ func (s *ConfigSynthesizer) Synthesize(ctx *SynthesisContext) ([]*coreauth.Auth, · 18 unmodified lines
30 out = append(out, s.synthesizeClaudeKeys(ctx)...)32 out = append(out, s.synthesizeClaudeKeys(ctx)...)
31 // Codex API Keys33 // Codex API Keys
32 out = append(out, s.synthesizeCodexKeys(ctx)...)34 out = append(out, s.synthesizeCodexKeys(ctx)...)
35 // Kiro (AWS CodeWhisperer)
36 out = append(out, s.synthesizeKiroKeys(ctx)...)
33 // OpenAI-compat37 // OpenAI-compat
34 out = append(out, s.synthesizeOpenAICompat(ctx)...)38 out = append(out, s.synthesizeOpenAICompat(ctx)...)
35 // Vertex-compat39 // Vertex-compat
@@ -298,3 +302,96 @@ func (s *ConfigSynthesizer) synthesizeVertexCompat(ctx *SynthesisContext) []*cor · 262 unmodified lines
298 }302 }
299 return out303 return out
300}304}
305
306// synthesizeKiroKeys creates Auth entries for Kiro (AWS CodeWhisperer) tokens.
307func (s *ConfigSynthesizer) synthesizeKiroKeys(ctx *SynthesisContext) []*coreauth.Auth {
308 cfg := ctx.Config
309 now := ctx.Now
310 idGen := ctx.IDGenerator
311
312 if len(cfg.KiroKey) == 0 {
313 return nil
314 }
315
316 out := make([]*coreauth.Auth, 0, len(cfg.KiroKey))
317 kAuth := kiroauth.NewKiroAuth(cfg)
318
319 for i := range cfg.KiroKey {
320 kk := cfg.KiroKey[i]
321 var accessToken, profileArn, refreshToken string
322
323 // Try to load from token file first
324 if kk.TokenFile != "" && kAuth != nil {
325 tokenData, err := kAuth.LoadTokenFromFile(kk.TokenFile)
326 if err != nil {
327 log.Warnf("failed to load kiro token file %s: %v", kk.TokenFile, err)
328 } else {
329 accessToken = tokenData.AccessToken
330 profileArn = tokenData.ProfileArn
331 refreshToken = tokenData.RefreshToken
332 }
333 }
334
335 // Override with direct config values if provided
336 if kk.AccessToken != "" {
337 accessToken = kk.AccessToken
338 }
339 if kk.ProfileArn != "" {
340 profileArn = kk.ProfileArn
341 }
342 if kk.RefreshToken != "" {
343 refreshToken = kk.RefreshToken
344 }
345
346 if accessToken == "" {
347 log.Warnf("kiro config[%d] missing access_token, skipping", i)
348 continue
349 }
350
351 // profileArn is optional for AWS Builder ID users
352 id, token := idGen.Next("kiro:token", accessToken, profileArn)
353 attrs := map[string]string{
354 "source": fmt.Sprintf("config:kiro[%s]", token),
355 "access_token": accessToken,
356 }
357 if profileArn != "" {
358 attrs["profile_arn"] = profileArn
359 }
360 if kk.Region != "" {
361 attrs["region"] = kk.Region
362 }
363 if kk.AgentTaskType != "" {
364 attrs["agent_task_type"] = kk.AgentTaskType
365 }
366 if kk.PreferredEndpoint != "" {
367 attrs["preferred_endpoint"] = kk.PreferredEndpoint
368 } else if cfg.KiroPreferredEndpoint != "" {
369 // Apply global default if not overridden by specific key
370 attrs["preferred_endpoint"] = cfg.KiroPreferredEndpoint
371 }
372 if refreshToken != "" {
373 attrs["refresh_token"] = refreshToken
374 }
375 proxyURL := strings.TrimSpace(kk.ProxyURL)
376 a := &coreauth.Auth{
377 ID: id,
378 Provider: "kiro",
379 Label: "kiro-token",
380 Status: coreauth.StatusActive,
381 ProxyURL: proxyURL,
382 Attributes: attrs,
383 CreatedAt: now,
384 UpdatedAt: now,
385 }
386
387 if refreshToken != "" {
388 if a.Metadata == nil {
389 a.Metadata = make(map[string]any)
390 }
391 a.Metadata["refresh_token"] = refreshToken
392 }
393
394 out = append(out, a)
395 }
396 return out
397}
sdk/api/handlers/handlers.go+2−1
@@ -175,10 +175,11 @@ type BaseAPIHandler struct { · 174 unmodified lines
175// Returns:175// Returns:
176// - *BaseAPIHandler: A new API handlers instance176// - *BaseAPIHandler: A new API handlers instance
177func NewBaseAPIHandlers(cfg *config.SDKConfig, authManager *coreauth.Manager) *BaseAPIHandler {177func NewBaseAPIHandlers(cfg *config.SDKConfig, authManager *coreauth.Manager) *BaseAPIHandler {
178 return &BaseAPIHandler{178 h := &BaseAPIHandler{
179 Cfg: cfg,179 Cfg: cfg,
180 AuthManager: authManager,180 AuthManager: authManager,
181 }181 }
182 return h
182}183}
183 184
184// UpdateClients updates the handlers' client list and configuration.185// UpdateClients updates the handlers' client list and configuration.
sdk/auth/github_copilot.go+129−0
@@ -0,0 +1,129 @@
1package auth
2
3import (
4 "context"
5 "fmt"
6 "time"
7
8 "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/copilot"
9 "github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
10 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
11 coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
12 log "github.com/sirupsen/logrus"
13)
14
15// GitHubCopilotAuthenticator implements the OAuth device flow login for GitHub Copilot.
16type GitHubCopilotAuthenticator struct{}
17
18// NewGitHubCopilotAuthenticator constructs a new GitHub Copilot authenticator.
19func NewGitHubCopilotAuthenticator() Authenticator {
20 return &GitHubCopilotAuthenticator{}
21}
22
23// Provider returns the provider key for github-copilot.
24func (GitHubCopilotAuthenticator) Provider() string {
25 return "github-copilot"
26}
27
28// RefreshLead returns nil since GitHub OAuth tokens don't expire in the traditional sense.
29// The token remains valid until the user revokes it or the Copilot subscription expires.
30func (GitHubCopilotAuthenticator) RefreshLead() *time.Duration {
31 return nil
32}
33
34// Login initiates the GitHub device flow authentication for Copilot access.
35func (a GitHubCopilotAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
36 if cfg == nil {
37 return nil, fmt.Errorf("cliproxy auth: configuration is required")
38 }
39 if opts == nil {
40 opts = &LoginOptions{}
41 }
42
43 authSvc := copilot.NewCopilotAuth(cfg)
44
45 // Start the device flow
46 fmt.Println("Starting GitHub Copilot authentication...")
47 deviceCode, err := authSvc.StartDeviceFlow(ctx)
48 if err != nil {
49 return nil, fmt.Errorf("github-copilot: failed to start device flow: %w", err)
50 }
51
52 // Display the user code and verification URL
53 fmt.Printf("\nTo authenticate, please visit: %s\n", deviceCode.VerificationURI)
54 fmt.Printf("And enter the code: %s\n\n", deviceCode.UserCode)
55
56 // Try to open the browser automatically
57 if !opts.NoBrowser {
58 if browser.IsAvailable() {
59 if errOpen := browser.OpenURL(deviceCode.VerificationURI); errOpen != nil {
60 log.Warnf("Failed to open browser automatically: %v", errOpen)
61 }
62 }
63 }
64
65 fmt.Println("Waiting for GitHub authorization...")
66 fmt.Printf("(This will timeout in %d seconds if not authorized)\n", deviceCode.ExpiresIn)
67
68 // Wait for user authorization
69 authBundle, err := authSvc.WaitForAuthorization(ctx, deviceCode)
70 if err != nil {
71 errMsg := copilot.GetUserFriendlyMessage(err)
72 return nil, fmt.Errorf("github-copilot: %s", errMsg)
73 }
74
75 // Verify the token can get a Copilot API token
76 fmt.Println("Verifying Copilot access...")
77 apiToken, err := authSvc.GetCopilotAPIToken(ctx, authBundle.TokenData.AccessToken)
78 if err != nil {
79 return nil, fmt.Errorf("github-copilot: failed to verify Copilot access - you may not have an active Copilot subscription: %w", err)
80 }
81
82 // Create the token storage
83 tokenStorage := authSvc.CreateTokenStorage(authBundle)
84
85 // Build metadata with token information for the executor
86 metadata := map[string]any{
87 "type": "github-copilot",
88 "username": authBundle.Username,
89 "access_token": authBundle.TokenData.AccessToken,
90 "token_type": authBundle.TokenData.TokenType,
91 "scope": authBundle.TokenData.Scope,
92 "timestamp": time.Now().UnixMilli(),
93 }
94
95 if apiToken.ExpiresAt > 0 {
96 metadata["api_token_expires_at"] = apiToken.ExpiresAt
97 }
98
99 fileName := fmt.Sprintf("github-copilot-%s.json", authBundle.Username)
100
101 fmt.Printf("\nGitHub Copilot authentication successful for user: %s\n", authBundle.Username)
102
103 return &coreauth.Auth{
104 ID: fileName,
105 Provider: a.Provider(),
106 FileName: fileName,
107 Label: authBundle.Username,
108 Storage: tokenStorage,
109 Metadata: metadata,
110 }, nil
111}
112
113// RefreshGitHubCopilotToken validates and returns the current token status.
114// GitHub OAuth tokens don't need traditional refresh - we just validate they still work.
115func RefreshGitHubCopilotToken(ctx context.Context, cfg *config.Config, storage *copilot.CopilotTokenStorage) error {
116 if storage == nil || storage.AccessToken == "" {
117 return fmt.Errorf("no token available")
118 }
119
120 authSvc := copilot.NewCopilotAuth(cfg)
121
122 // Validate the token can still get a Copilot API token
123 _, err := authSvc.GetCopilotAPIToken(ctx, storage.AccessToken)
124 if err != nil {
125 return fmt.Errorf("token validation failed: %w", err)
126 }
127
128 return nil
129}
sdk/auth/kiro.go+470−0
@@ -0,0 +1,470 @@
1package auth
2
3import (
4 "context"
5 "fmt"
6 "strings"
7 "time"
8
9 kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"
10 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
11 coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
12)
13
14// extractKiroIdentifier extracts a meaningful identifier for file naming.
15// Returns account name if provided, otherwise profile ARN ID.
16// All extracted values are sanitized to prevent path injection attacks.
17func extractKiroIdentifier(accountName, profileArn string) string {
18 // Priority 1: Use account name if provided
19 if accountName != "" {
20 return kiroauth.SanitizeEmailForFilename(accountName)
21 }
22
23 // Priority 2: Use profile ARN ID part (sanitized to prevent path injection)
24 if profileArn != "" {
25 parts := strings.Split(profileArn, "/")
26 if len(parts) >= 2 {
27 // Sanitize the ARN component to prevent path traversal
28 return kiroauth.SanitizeEmailForFilename(parts[len(parts)-1])
29 }
30 }
31
32 // Fallback: timestamp
33 return fmt.Sprintf("%d", time.Now().UnixNano()%100000)
34}
35
36// KiroAuthenticator implements OAuth authentication for Kiro with Google login.
37type KiroAuthenticator struct{}
38
39// NewKiroAuthenticator constructs a Kiro authenticator.
40func NewKiroAuthenticator() *KiroAuthenticator {
41 return &KiroAuthenticator{}
42}
43
44// Provider returns the provider key for the authenticator.
45func (a *KiroAuthenticator) Provider() string {
46 return "kiro"
47}
48
49// RefreshLead indicates how soon before expiry a refresh should be attempted.
50// Set to 5 minutes to match Antigravity and avoid frequent refresh checks while still ensuring timely token refresh.
51func (a *KiroAuthenticator) RefreshLead() *time.Duration {
52 d := 5 * time.Minute
53 return &d
54}
55
56// createAuthRecord creates an auth record from token data.
57func (a *KiroAuthenticator) createAuthRecord(tokenData *kiroauth.KiroTokenData, source string) (*coreauth.Auth, error) {
58 // Parse expires_at
59 expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
60 if err != nil {
61 expiresAt = time.Now().Add(1 * time.Hour)
62 }
63
64 // Extract identifier for file naming
65 idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
66
67 // Determine label based on auth method
68 label := fmt.Sprintf("kiro-%s", source)
69 if tokenData.AuthMethod == "idc" {
70 label = "kiro-idc"
71 }
72
73 now := time.Now()
74 fileName := fmt.Sprintf("%s-%s.json", label, idPart)
75
76 metadata := map[string]any{
77 "type": "kiro",
78 "access_token": tokenData.AccessToken,
79 "refresh_token": tokenData.RefreshToken,
80 "profile_arn": tokenData.ProfileArn,
81 "expires_at": tokenData.ExpiresAt,
82 "auth_method": tokenData.AuthMethod,
83 "provider": tokenData.Provider,
84 "client_id": tokenData.ClientID,
85 "client_secret": tokenData.ClientSecret,
86 "email": tokenData.Email,
87 }
88
89 // Add IDC-specific fields if present
90 if tokenData.StartURL != "" {
91 metadata["start_url"] = tokenData.StartURL
92 }
93 if tokenData.Region != "" {
94 metadata["region"] = tokenData.Region
95 }
96
97 attributes := map[string]string{
98 "profile_arn": tokenData.ProfileArn,
99 "source": source,
100 "email": tokenData.Email,
101 }
102
103 // Add IDC-specific attributes if present
104 if tokenData.AuthMethod == "idc" {
105 attributes["source"] = "aws-idc"
106 if tokenData.StartURL != "" {
107 attributes["start_url"] = tokenData.StartURL
108 }
109 if tokenData.Region != "" {
110 attributes["region"] = tokenData.Region
111 }
112 }
113
114 record := &coreauth.Auth{
115 ID: fileName,
116 Provider: "kiro",
117 FileName: fileName,
118 Label: label,
119 Status: coreauth.StatusActive,
120 CreatedAt: now,
121 UpdatedAt: now,
122 Metadata: metadata,
123 Attributes: attributes,
124 // NextRefreshAfter is aligned with RefreshLead (5min)
125 NextRefreshAfter: expiresAt.Add(-5 * time.Minute),
126 }
127
128 if tokenData.Email != "" {
129 fmt.Printf("\n✓ Kiro authentication completed successfully! (Account: %s)\n", tokenData.Email)
130 } else {
131 fmt.Println("\n✓ Kiro authentication completed successfully!")
132 }
133
134 return record, nil
135}
136
137// Login performs OAuth login for Kiro with AWS (Builder ID or IDC).
138// This shows a method selection prompt and handles both flows.
139func (a *KiroAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
140 if cfg == nil {
141 return nil, fmt.Errorf("kiro auth: configuration is required")
142 }
143
144 // Use the unified method selection flow (Builder ID or IDC)
145 ssoClient := kiroauth.NewSSOOIDCClient(cfg)
146 tokenData, err := ssoClient.LoginWithMethodSelection(ctx)
147 if err != nil {
148 return nil, fmt.Errorf("login failed: %w", err)
149 }
150
151 return a.createAuthRecord(tokenData, "aws")
152}
153
154// LoginWithAuthCode performs OAuth login for Kiro with AWS Builder ID using authorization code flow.
155// This provides a better UX than device code flow as it uses automatic browser callback.
156func (a *KiroAuthenticator) LoginWithAuthCode(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
157 if cfg == nil {
158 return nil, fmt.Errorf("kiro auth: configuration is required")
159 }
160
161 oauth := kiroauth.NewKiroOAuth(cfg)
162
163 // Use AWS Builder ID authorization code flow
164 tokenData, err := oauth.LoginWithBuilderIDAuthCode(ctx)
165 if err != nil {
166 return nil, fmt.Errorf("login failed: %w", err)
167 }
168
169 // Parse expires_at
170 expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
171 if err != nil {
172 expiresAt = time.Now().Add(1 * time.Hour)
173 }
174
175 // Extract identifier for file naming
176 idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
177
178 now := time.Now()
179 fileName := fmt.Sprintf("kiro-aws-%s.json", idPart)
180
181 record := &coreauth.Auth{
182 ID: fileName,
183 Provider: "kiro",
184 FileName: fileName,
185 Label: "kiro-aws",
186 Status: coreauth.StatusActive,
187 CreatedAt: now,
188 UpdatedAt: now,
189 Metadata: map[string]any{
190 "type": "kiro",
191 "access_token": tokenData.AccessToken,
192 "refresh_token": tokenData.RefreshToken,
193 "profile_arn": tokenData.ProfileArn,
194 "expires_at": tokenData.ExpiresAt,
195 "auth_method": tokenData.AuthMethod,
196 "provider": tokenData.Provider,
197 "client_id": tokenData.ClientID,
198 "client_secret": tokenData.ClientSecret,
199 "email": tokenData.Email,
200 },
201 Attributes: map[string]string{
202 "profile_arn": tokenData.ProfileArn,
203 "source": "aws-builder-id-authcode",
204 "email": tokenData.Email,
205 },
206 // NextRefreshAfter is aligned with RefreshLead (5min)
207 NextRefreshAfter: expiresAt.Add(-5 * time.Minute),
208 }
209
210 if tokenData.Email != "" {
211 fmt.Printf("\n✓ Kiro authentication completed successfully! (Account: %s)\n", tokenData.Email)
212 } else {
213 fmt.Println("\n✓ Kiro authentication completed successfully!")
214 }
215
216 return record, nil
217}
218
219// LoginWithGoogle performs OAuth login for Kiro with Google.
220// This uses a custom protocol handler (kiro://) to receive the callback.
221func (a *KiroAuthenticator) LoginWithGoogle(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
222 if cfg == nil {
223 return nil, fmt.Errorf("kiro auth: configuration is required")
224 }
225
226 oauth := kiroauth.NewKiroOAuth(cfg)
227
228 // Use Google OAuth flow with protocol handler
229 tokenData, err := oauth.LoginWithGoogle(ctx)
230 if err != nil {
231 return nil, fmt.Errorf("google login failed: %w", err)
232 }
233
234 // Parse expires_at
235 expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
236 if err != nil {
237 expiresAt = time.Now().Add(1 * time.Hour)
238 }
239
240 // Extract identifier for file naming
241 idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
242
243 now := time.Now()
244 fileName := fmt.Sprintf("kiro-google-%s.json", idPart)
245
246 record := &coreauth.Auth{
247 ID: fileName,
248 Provider: "kiro",
249 FileName: fileName,
250 Label: "kiro-google",
251 Status: coreauth.StatusActive,
252 CreatedAt: now,
253 UpdatedAt: now,
254 Metadata: map[string]any{
255 "type": "kiro",
256 "access_token": tokenData.AccessToken,
257 "refresh_token": tokenData.RefreshToken,
258 "profile_arn": tokenData.ProfileArn,
259 "expires_at": tokenData.ExpiresAt,
260 "auth_method": tokenData.AuthMethod,
261 "provider": tokenData.Provider,
262 "email": tokenData.Email,
263 },
264 Attributes: map[string]string{
265 "profile_arn": tokenData.ProfileArn,
266 "source": "google-oauth",
267 "email": tokenData.Email,
268 },
269 // NextRefreshAfter is aligned with RefreshLead (5min)
270 NextRefreshAfter: expiresAt.Add(-5 * time.Minute),
271 }
272
273 if tokenData.Email != "" {
274 fmt.Printf("\n✓ Kiro Google authentication completed successfully! (Account: %s)\n", tokenData.Email)
275 } else {
276 fmt.Println("\n✓ Kiro Google authentication completed successfully!")
277 }
278
279 return record, nil
280}
281
282// LoginWithGitHub performs OAuth login for Kiro with GitHub.
283// This uses a custom protocol handler (kiro://) to receive the callback.
284func (a *KiroAuthenticator) LoginWithGitHub(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
285 if cfg == nil {
286 return nil, fmt.Errorf("kiro auth: configuration is required")
287 }
288
289 oauth := kiroauth.NewKiroOAuth(cfg)
290
291 // Use GitHub OAuth flow with protocol handler
292 tokenData, err := oauth.LoginWithGitHub(ctx)
293 if err != nil {
294 return nil, fmt.Errorf("github login failed: %w", err)
295 }
296
297 // Parse expires_at
298 expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
299 if err != nil {
300 expiresAt = time.Now().Add(1 * time.Hour)
301 }
302
303 // Extract identifier for file naming
304 idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
305
306 now := time.Now()
307 fileName := fmt.Sprintf("kiro-github-%s.json", idPart)
308
309 record := &coreauth.Auth{
310 ID: fileName,
311 Provider: "kiro",
312 FileName: fileName,
313 Label: "kiro-github",
314 Status: coreauth.StatusActive,
315 CreatedAt: now,
316 UpdatedAt: now,
317 Metadata: map[string]any{
318 "type": "kiro",
319 "access_token": tokenData.AccessToken,
320 "refresh_token": tokenData.RefreshToken,
321 "profile_arn": tokenData.ProfileArn,
322 "expires_at": tokenData.ExpiresAt,
323 "auth_method": tokenData.AuthMethod,
324 "provider": tokenData.Provider,
325 "email": tokenData.Email,
326 },
327 Attributes: map[string]string{
328 "profile_arn": tokenData.ProfileArn,
329 "source": "github-oauth",
330 "email": tokenData.Email,
331 },
332 // NextRefreshAfter is aligned with RefreshLead (5min)
333 NextRefreshAfter: expiresAt.Add(-5 * time.Minute),
334 }
335
336 if tokenData.Email != "" {
337 fmt.Printf("\n✓ Kiro GitHub authentication completed successfully! (Account: %s)\n", tokenData.Email)
338 } else {
339 fmt.Println("\n✓ Kiro GitHub authentication completed successfully!")
340 }
341
342 return record, nil
343}
344
345// ImportFromKiroIDE imports token from Kiro IDE's token file.
346func (a *KiroAuthenticator) ImportFromKiroIDE(ctx context.Context, cfg *config.Config) (*coreauth.Auth, error) {
347 tokenData, err := kiroauth.LoadKiroIDEToken()
348 if err != nil {
349 return nil, fmt.Errorf("failed to load Kiro IDE token: %w", err)
350 }
351
352 // Parse expires_at
353 expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
354 if err != nil {
355 expiresAt = time.Now().Add(1 * time.Hour)
356 }
357
358 // Extract email from JWT if not already set (for imported tokens)
359 if tokenData.Email == "" {
360 tokenData.Email = kiroauth.ExtractEmailFromJWT(tokenData.AccessToken)
361 }
362
363 // Extract identifier for file naming
364 idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
365 // Sanitize provider to prevent path traversal (defense-in-depth)
366 provider := kiroauth.SanitizeEmailForFilename(strings.ToLower(strings.TrimSpace(tokenData.Provider)))
367 if provider == "" {
368 provider = "imported" // Fallback for legacy tokens without provider
369 }
370
371 now := time.Now()
372 fileName := fmt.Sprintf("kiro-%s-%s.json", provider, idPart)
373
374 record := &coreauth.Auth{
375 ID: fileName,
376 Provider: "kiro",
377 FileName: fileName,
378 Label: fmt.Sprintf("kiro-%s", provider),
379 Status: coreauth.StatusActive,
380 CreatedAt: now,
381 UpdatedAt: now,
382 Metadata: map[string]any{
383 "type": "kiro",
384 "access_token": tokenData.AccessToken,
385 "refresh_token": tokenData.RefreshToken,
386 "profile_arn": tokenData.ProfileArn,
387 "expires_at": tokenData.ExpiresAt,
388 "auth_method": tokenData.AuthMethod,
389 "provider": tokenData.Provider,
390 "email": tokenData.Email,
391 },
392 Attributes: map[string]string{
393 "profile_arn": tokenData.ProfileArn,
394 "source": "kiro-ide-import",
395 "email": tokenData.Email,
396 },
397 // NextRefreshAfter is aligned with RefreshLead (5min)
398 NextRefreshAfter: expiresAt.Add(-5 * time.Minute),
399 }
400
401 // Display the email if extracted
402 if tokenData.Email != "" {
403 fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s, Account: %s)\n", tokenData.Provider, tokenData.Email)
404 } else {
405 fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s)\n", tokenData.Provider)
406 }
407
408 return record, nil
409}
410
411// Refresh refreshes an expired Kiro token using AWS SSO OIDC.
412func (a *KiroAuthenticator) Refresh(ctx context.Context, cfg *config.Config, auth *coreauth.Auth) (*coreauth.Auth, error) {
413 if auth == nil || auth.Metadata == nil {
414 return nil, fmt.Errorf("invalid auth record")
415 }
416
417 refreshToken, ok := auth.Metadata["refresh_token"].(string)
418 if !ok || refreshToken == "" {
419 return nil, fmt.Errorf("refresh token not found")
420 }
421
422 clientID, _ := auth.Metadata["client_id"].(string)
423 clientSecret, _ := auth.Metadata["client_secret"].(string)
424 authMethod, _ := auth.Metadata["auth_method"].(string)
425 startURL, _ := auth.Metadata["start_url"].(string)
426 region, _ := auth.Metadata["region"].(string)
427
428 var tokenData *kiroauth.KiroTokenData
429 var err error
430
431 ssoClient := kiroauth.NewSSOOIDCClient(cfg)
432
433 // Use SSO OIDC refresh for AWS Builder ID or IDC, otherwise use Kiro's OAuth refresh endpoint
434 switch {
435 case clientID != "" && clientSecret != "" && authMethod == "idc" && region != "":
436 // IDC refresh with region-specific endpoint
437 tokenData, err = ssoClient.RefreshTokenWithRegion(ctx, clientID, clientSecret, refreshToken, region, startURL)
438 case clientID != "" && clientSecret != "" && authMethod == "builder-id":
439 // Builder ID refresh with default endpoint
440 tokenData, err = ssoClient.RefreshToken(ctx, clientID, clientSecret, refreshToken)
441 default:
442 // Fallback to Kiro's refresh endpoint (for social auth: Google/GitHub)
443 oauth := kiroauth.NewKiroOAuth(cfg)
444 tokenData, err = oauth.RefreshToken(ctx, refreshToken)
445 }
446
447 if err != nil {
448 return nil, fmt.Errorf("token refresh failed: %w", err)
449 }
450
451 // Parse expires_at
452 expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
453 if err != nil {
454 expiresAt = time.Now().Add(1 * time.Hour)
455 }
456
457 // Clone auth to avoid mutating the input parameter
458 updated := auth.Clone()
459 now := time.Now()
460 updated.UpdatedAt = now
461 updated.LastRefreshedAt = now
462 updated.Metadata["access_token"] = tokenData.AccessToken
463 updated.Metadata["refresh_token"] = tokenData.RefreshToken
464 updated.Metadata["expires_at"] = tokenData.ExpiresAt
465 updated.Metadata["last_refresh"] = now.Format(time.RFC3339) // For double-check optimization
466 // NextRefreshAfter is aligned with RefreshLead (5min)
467 updated.NextRefreshAfter = expiresAt.Add(-5 * time.Minute)
468
469 return updated, nil
470}
sdk/auth/manager.go+13−0
@@ -74,3 +74,16 @@ func (m *Manager) Login(ctx context.Context, provider string, cfg *config.Config · 73 unmodified lines
74 }74 }
75 return record, savedPath, nil75 return record, savedPath, nil
76}76}
77
78// SaveAuth persists an auth record directly without going through the login flow.
79func (m *Manager) SaveAuth(record *coreauth.Auth, cfg *config.Config) (string, error) {
80 if m.store == nil {
81 return "", fmt.Errorf("no store configured")
82 }
83 if cfg != nil {
84 if dirSetter, ok := m.store.(interface{ SetBaseDir(string) }); ok {
85 dirSetter.SetBaseDir(cfg.AuthDir)
86 }
87 }
88 return m.store.Save(context.Background(), record)
89}
sdk/auth/refresh_registry.go+2−0
@@ -14,6 +14,8 @@ func init() { · 13 unmodified lines
14 registerRefreshLead("gemini", func() Authenticator { return NewGeminiAuthenticator() })14 registerRefreshLead("gemini", func() Authenticator { return NewGeminiAuthenticator() })
15 registerRefreshLead("gemini-cli", func() Authenticator { return NewGeminiAuthenticator() })15 registerRefreshLead("gemini-cli", func() Authenticator { return NewGeminiAuthenticator() })
16 registerRefreshLead("antigravity", func() Authenticator { return NewAntigravityAuthenticator() })16 registerRefreshLead("antigravity", func() Authenticator { return NewAntigravityAuthenticator() })
17 registerRefreshLead("kiro", func() Authenticator { return NewKiroAuthenticator() })
18 registerRefreshLead("github-copilot", func() Authenticator { return NewGitHubCopilotAuthenticator() })
17}19}
18 20
19func registerRefreshLead(provider string, factory func() Authenticator) {21func registerRefreshLead(provider string, factory func() Authenticator) {
sdk/cliproxy/auth/conductor.go+4−2
@@ -41,7 +41,7 @@ type RefreshEvaluator interface { · 40 unmodified lines
41const (41const (
42 refreshCheckInterval = 5 * time.Second42 refreshCheckInterval = 5 * time.Second
43 refreshPendingBackoff = time.Minute43 refreshPendingBackoff = time.Minute
44 refreshFailureBackoff = 5 * time.Minute44 refreshFailureBackoff = 1 * time.Minute
45 quotaBackoffBase = time.Second45 quotaBackoffBase = time.Second
46 quotaBackoffMax = 30 * time.Minute46 quotaBackoffMax = 30 * time.Minute
47)47)
@@ -1569,7 +1569,9 @@ func (m *Manager) refreshAuth(ctx context.Context, id string) { · 1521 unmodified lines
1569 updated.Runtime = auth.Runtime1569 updated.Runtime = auth.Runtime
1570 }1570 }
1571 updated.LastRefreshedAt = now1571 updated.LastRefreshedAt = now
1572 updated.NextRefreshAfter = time.Time{}1572 // Preserve NextRefreshAfter set by the Authenticator
1573 // If the Authenticator set a reasonable refresh time, it should not be overwritten
1574 // If the Authenticator did not set it (zero value), shouldRefresh will use default logic
1573 updated.LastError = nil1575 updated.LastError = nil
1574 updated.UpdatedAt = now1576 updated.UpdatedAt = now
1575 _, _ = m.Update(ctx, updated)1577 _, _ = m.Update(ctx, updated)
sdk/cliproxy/service.go+9−0
@@ -379,6 +379,10 @@ func (s *Service) ensureExecutorsForAuth(a *coreauth.Auth) { · 378 unmodified lines
379 s.coreManager.RegisterExecutor(executor.NewQwenExecutor(s.cfg))379 s.coreManager.RegisterExecutor(executor.NewQwenExecutor(s.cfg))
380 case "iflow":380 case "iflow":
381 s.coreManager.RegisterExecutor(executor.NewIFlowExecutor(s.cfg))381 s.coreManager.RegisterExecutor(executor.NewIFlowExecutor(s.cfg))
382 case "kiro":
383 s.coreManager.RegisterExecutor(executor.NewKiroExecutor(s.cfg))
384 case "github-copilot":
385 s.coreManager.RegisterExecutor(executor.NewGitHubCopilotExecutor(s.cfg))
382 default:386 default:
383 providerKey := strings.ToLower(strings.TrimSpace(a.Provider))387 providerKey := strings.ToLower(strings.TrimSpace(a.Provider))
384 if providerKey == "" {388 if providerKey == "" {
@@ -765,6 +769,11 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) { · 380 unmodified lines
765 models = applyExcludedModels(models, excluded)769 models = applyExcludedModels(models, excluded)
766 case "iflow":770 case "iflow":
767 models = registry.GetIFlowModels()771 models = registry.GetIFlowModels()
772 case "github-copilot":
773 models = registry.GetGitHubCopilotModels()
774 models = applyExcludedModels(models, excluded)
775 case "kiro":
776 models = registry.GetKiroModels()
768 models = applyExcludedModels(models, excluded)777 models = applyExcludedModels(models, excluded)
769 default:778 default:
770 // Handle OpenAI-compatibility providers by name using config779 // Handle OpenAI-compatibility providers by name using config

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t