Compare commits

...

160 Commits

Author SHA1 Message Date
Luis Pater
c353606860 Merge pull request #14 from router-for-me/plus
v6.5.59
2025-12-09 17:51:46 +08:00
Luis Pater
2ba31ecc2d Merge branch 'main' into plus 2025-12-09 17:51:18 +08:00
Luis Pater
39b6b3b289 Fixed: #463
fix(antigravity): remove `$ref` and `$defs` from JSON during key deletion
2025-12-09 17:32:17 +08:00
Luis Pater
c600519fa4 refactor(logging): replace log.Fatalf with log.Errorf and add error handling paths 2025-12-09 17:16:30 +08:00
Luis Pater
36380846a7 Merge branch 'router-for-me:main' into main 2025-12-09 10:03:16 +08:00
Luis Pater
92df0cada9 Merge pull request #461 from router-for-me/aistudio
feat(aistudio): normalize thinking budget in request translation
2025-12-09 08:41:46 +08:00
hkfires
96b55acff8 feat(aistudio): normalize thinking budget in request translation 2025-12-09 08:27:44 +08:00
Luis Pater
608cd8ee3d Merge pull request #13 from router-for-me/v6.5.57
v6.5.57
2025-12-08 23:33:52 +08:00
Luis Pater
9f41894573 Merge branch 'main' into v6.5.57 2025-12-08 23:33:39 +08:00
Luis Pater
bb45fee1cf Merge remote-tracking branch 'origin/dev' into dev 2025-12-08 23:28:22 +08:00
Luis Pater
af00304b0c fix(antigravity): remove exclusiveMaximum from JSON during key deletion 2025-12-08 23:28:01 +08:00
vuonglv(Andy)
5c3a013cd1 feat(config): add configurable host binding for server (#454)
* feat(config): add configurable host binding for server
2025-12-08 23:16:39 +08:00
Luis Pater
ab9e9442ec v6.5.56 (#12)
* feat(api): add comprehensive ampcode management endpoints

Add new REST API endpoints under /v0/management/ampcode for managing
ampcode configuration including upstream URL, API key, localhost
restriction, model mappings, and force model mappings settings.

- Move force-model-mappings from config_basic to config_lists
- Add GET/PUT/PATCH/DELETE endpoints for all ampcode settings
- Support model mapping CRUD with upsert (PATCH) capability
- Add comprehensive test coverage for all ampcode endpoints

* refactor(api): simplify request body parsing in ampcode handlers

* feat(logging): add upstream API request/response capture to streaming logs

* style(logging): remove redundant separator line from response section

* feat(antigravity): enforce thinking budget limits for Claude models

* refactor(logging): remove unused variable in `ensureAttempt` and redundant function call

---------

Co-authored-by: hkfires <10558748+hkfires@users.noreply.github.com>
2025-12-08 22:32:29 +08:00
Luis Pater
6ad188921c refactor(logging): remove unused variable in ensureAttempt and redundant function call 2025-12-08 22:25:58 +08:00
Luis Pater
15ed98d6a9 Merge pull request #458 from router-for-me/agry
feat(antigravity): enforce thinking budget limits for Claude models
2025-12-08 20:55:52 +08:00
hkfires
a283545b6b feat(antigravity): enforce thinking budget limits for Claude models 2025-12-08 20:36:17 +08:00
Luis Pater
3efbd865a8 Merge pull request #457 from router-for-me/requestlog
style(logging): remove redundant separator line from response section
2025-12-08 18:21:24 +08:00
hkfires
aee659fb66 style(logging): remove redundant separator line from response section 2025-12-08 18:18:33 +08:00
Luis Pater
5aa386d8b9 Merge pull request #453 from router-for-me/amp
add ampcode management api
2025-12-08 17:42:13 +08:00
Luis Pater
0adc0ee6aa Merge pull request #455 from router-for-me/requestlog
feat(logging): add upstream API request/response capture to streaming logs
2025-12-08 17:40:10 +08:00
hkfires
92f13fc316 feat(logging): add upstream API request/response capture to streaming logs 2025-12-08 17:21:58 +08:00
hkfires
05cfa16e5f refactor(api): simplify request body parsing in ampcode handlers 2025-12-08 14:45:35 +08:00
hkfires
93a6e2d920 feat(api): add comprehensive ampcode management endpoints
Add new REST API endpoints under /v0/management/ampcode for managing
ampcode configuration including upstream URL, API key, localhost
restriction, model mappings, and force model mappings settings.

- Move force-model-mappings from config_basic to config_lists
- Add GET/PUT/PATCH/DELETE endpoints for all ampcode settings
- Support model mapping CRUD with upsert (PATCH) capability
- Add comprehensive test coverage for all ampcode endpoints
2025-12-08 12:03:00 +08:00
Luis Pater
6b49580716 Merge branch 'router-for-me:main' into main 2025-12-08 10:52:39 +08:00
Luis Pater
de77903915 Merge pull request #450 from router-for-me/amp
refactor(config): rename prioritize-model-mappings to force-model-mappings
2025-12-08 10:51:32 +08:00
hkfires
56ed0d8d90 refactor(config): rename prioritize-model-mappings to force-model-mappings 2025-12-08 10:44:39 +08:00
Luis Pater
0d4f32a881 Merge branch 'router-for-me:main' into main 2025-12-08 10:20:08 +08:00
Luis Pater
42e818ce05 Merge pull request #435 from heyhuynhgiabuu/fix/amp-model-mapping-priority
fix: prioritize model mappings over local providers for Amp CLI
2025-12-08 10:17:19 +08:00
Luis Pater
2d4c54ba54 Merge pull request #448 from router-for-me/iflow
Iflow
2025-12-08 09:50:05 +08:00
hkfires
e9eb4db8bb feat(auth): refresh API key during cookie authentication 2025-12-08 09:48:31 +08:00
Luis Pater
d26ed069fa Merge pull request #441 from huynguyen03dev/fix/claude-to-openai-whitespace-text
fix: filter whitespace-only text in Claude to OpenAI translation
2025-12-08 09:43:44 +08:00
huynhgiabuu
afcab5efda feat: add prioritize-model-mappings config option
Add a configuration option to control whether model mappings take
precedence over local API keys for Amp CLI requests.

- Add PrioritizeModelMappings field to AmpCode config struct
- When false (default): Local API keys take precedence (original behavior)
- When true: Model mappings take precedence over local API keys
- Add management API endpoints GET/PUT /prioritize-model-mappings

This allows users who want mapping priority to enable it explicitly
while preserving backward compatibility.

Config example:
  ampcode:
    model-mappings:
      - from: claude-opus-4-5-20251101
        to: gemini-claude-opus-4-5-thinking
    prioritize-model-mappings: true
2025-12-07 22:47:43 +07:00
Ravens2121
a0c6cffb0d fix(kiro):修复 base64 图片格式转换问题 (#10) 2025-12-07 21:55:13 +08:00
Luis Pater
f81ff16022 Merge pull request #8 from Ravens2121/main
feat: 添加Kiro渠道图片支持功能,借鉴justlovemaki/AIClient-2-API实现
2025-12-07 20:18:31 +08:00
Luis Pater
6cf1d8a947 Merge pull request #444 from router-for-me/agry
feat(registry): add explicit thinking support config for antigravity models
2025-12-07 19:38:43 +08:00
hkfires
a174d015f2 feat(openai): handle thinking.budget_tokens from Anthropic-style requests 2025-12-07 19:14:05 +08:00
hkfires
9c09128e00 feat(registry): add explicit thinking support config for antigravity models 2025-12-07 19:12:55 +08:00
Your Name
68cbe20664 feat: 添加Kiro渠道图片支持功能,借鉴justlovemaki/AIClient-2-API实现 2025-12-07 17:56:06 +08:00
huynguyen03.dev
549c0c2c5a fix: filter whitespace-only text content in Claude to OpenAI translation
Remove redundant existence check since TrimSpace handles empty strings
2025-12-07 16:08:12 +07:00
huynguyen03.dev
f092801b61 fix: filter whitespace-only text in Claude to OpenAI translation
Skip text content blocks that are empty or contain only whitespace
when translating Claude messages to OpenAI format. This fixes GLM-4.6
and other strict OpenAI-compatible providers that reject empty text
with error 'text cannot be empty'.
2025-12-07 15:39:58 +07:00
Luis Pater
15353a6b6a Merge branch 'router-for-me:main' into main 2025-12-07 13:37:43 +08:00
Luis Pater
1b638b3629 Merge pull request #432 from huynguyen03dev/fix/amp-gemini-model-mapping
fix(amp): pass mapped model to gemini bridge via context
2025-12-07 13:33:28 +08:00
Luis Pater
6f5f81753d Merge pull request #439 from router-for-me/log
feat(logging): add version info to request log output
2025-12-07 13:31:06 +08:00
Luis Pater
76af454034 **feat(antigravity): enhance handling of "thinking" content and refine Claude model response processing** 2025-12-07 13:19:12 +08:00
hkfires
e54d2f6b2a feat(logging): add version info to request log output 2025-12-07 12:49:14 +08:00
huynguyen03.dev
bfc738b76a refactor: remove duplicate provider check in gemini v1beta1 route
Simplifies routing logic by delegating all provider/mapping/proxy
decisions to FallbackHandler. Previously, the route checked for
provider/mapping availability before calling the handler, then
FallbackHandler performed the same checks again.

Changes:
- Remove model extraction and provider checking from route (lines 182-201)
- Route now only checks if request is POST with /models/ path
- FallbackHandler handles provider -> mapping -> proxy fallback
- Remove unused internal/util import

Benefits:
- Eliminates duplicate checks (addresses PR review feedback #2)
- Centralizes all provider/mapping logic in FallbackHandler
- Reduces routing code by ~20 lines
- Aligns with how other /api/provider routes work

Performance: No impact (checks still happen once in FallbackHandler)
2025-12-07 10:54:58 +07:00
huynguyen03.dev
396899a530 refactor: improve gemini bridge testability and code quality
- Change createGeminiBridgeHandler to accept gin.HandlerFunc instead of *gemini.GeminiAPIHandler
  This allows tests to inject mock handlers instead of duplicating bridge logic
- Replace magic number 8 with len(modelsPrefix) for better maintainability
- Remove redundant test case that doesn't test edge case in production
- Update routes.go to pass geminiHandlers.GeminiHandler directly

Addresses PR review feedback on test architecture and code clarity.

Amp-Thread-ID: https://ampcode.com/threads/T-1ae2c691-e434-4b99-a49a-10cabd3544db
2025-12-07 10:15:42 +07:00
Luis Pater
04f0070a80 Merge branch 'router-for-me:main' into main 2025-12-07 02:38:32 +08:00
Luis Pater
f383840cf9 fix(antigravity): update toolNode role from "tool" to "user" in chat completions 2025-12-07 02:37:46 +08:00
Luis Pater
239fc4a8c4 Merge branch 'router-for-me:main' into main 2025-12-07 01:58:03 +08:00
Luis Pater
fd29ab418a Fixed: #424
**feat(antigravity): add support for maxOutputTokens and refine Claude model handling**
2025-12-07 01:55:57 +08:00
Luis Pater
df91408919 Merge branch 'router-for-me:main' into main 2025-12-07 01:49:20 +08:00
Luis Pater
7a628426dc Fixed: #433
refactor(translator): normalize finish reason casing across all OpenAI response handlers
2025-12-07 01:48:24 +08:00
Luis Pater
aa6c7facab Merge pull request #5 from router-for-me/plus
Sync
2025-12-07 01:23:53 +08:00
Luis Pater
8ba4c7c7ed Merge branch 'main' into plus 2025-12-07 01:23:36 +08:00
Luis Pater
56b4d7a76e docs(readme): add ProxyPal CLIProxyAPI GUI to project list 2025-12-07 01:13:30 +08:00
Luis Pater
b211c3546d Merge pull request #429 from heyhuynhgiabuu/feature/add-proxypal
docs: add ProxyPal to 'Who is with us?' section
2025-12-07 01:10:44 +08:00
huynguyen03.dev
edc654edf9 refactor: simplify provider check logic in amp routes
Amp-Thread-ID: https://ampcode.com/threads/T-a18fd71c-32ce-4c29-93d7-09f082740e51
2025-12-06 22:07:40 +07:00
huynguyen03.dev
08586334af fix(amp): pass mapped model to gemini bridge via context
Gemini handler extracts model from URL path, not JSON body, so
rewriting the request body alone wasn't sufficient for model mapping.

- Add MappedModelContextKey constant for context passing
- Update routes.go to use NewFallbackHandlerWithMapper
- Add check for valid mapping before routing to local handler
- Add tests for gemini bridge model mapping
2025-12-06 18:59:44 +07:00
Luis Pater
a4804b358f Merge branch 'router-for-me:main' into main 2025-12-06 15:15:27 +08:00
Luis Pater
7ea14479fb Merge pull request #428 from router-for-me/amp
feat(amp): add response rewriter for model name substitution in responses
2025-12-06 15:14:10 +08:00
hkfires
54af96d321 fix(amp): restore request body before fallback handler execution 2025-12-06 15:09:25 +08:00
hkfires
22579155c5 refactor(amp): consolidate and simplify model mapping debug logs 2025-12-06 14:54:38 +08:00
Huynh Gia Buu
c04c3832a4 Update README.md
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2025-12-06 13:48:08 +07:00
huynhgiabuu
5ffbd54755 docs: add ProxyPal to 'Who is with us?' section 2025-12-06 13:45:49 +07:00
hkfires
5d12d4ce33 feat(amp): add response rewriter for model name substitution in responses 2025-12-06 14:15:44 +08:00
Luis Pater
b73e53d6c4 docs: update README to fix formatting for GitHub Copilot and Kiro support sections 2025-12-06 12:14:23 +08:00
Luis Pater
b06463c6d9 docs: update README to include Kiro (AWS CodeWhisperer) support integration details 2025-12-06 12:04:36 +08:00
Luis Pater
5eb8453e91 Merge pull request #3 from fuko2935/feature/kiro-integration
Feature/kiro integration
2025-12-06 12:00:48 +08:00
Luis Pater
f77c22e6ff Merge branch 'main' into feature/kiro-integration 2025-12-06 11:52:59 +08:00
Mansi
df83ba877f refactor: replace custom stringContains with strings.Contains
- Remove custom stringContains and findSubstring helper functions
- Use standard library strings.Contains for better maintainability
- No functional change, just cleaner code

Addresses Gemini Code Assist review feedback
2025-12-05 23:16:48 +03:00
Mansi
9583f6b1c5 refactor: extract setKiroIncognitoMode helper to reduce code duplication
- Added setKiroIncognitoMode() helper function to handle Kiro auth incognito mode setting
- Replaced 3 duplicate code blocks (21 lines) with single function calls (3 lines)
- Kiro auth defaults to incognito mode for multi-account support
- Users can override with --incognito or --no-incognito flags

This addresses the code duplication noted in PR #1 review.
2025-12-05 23:06:07 +03:00
Mansi
02d8a1cfec feat(kiro): add AWS Builder ID authentication support
- Add --kiro-aws-login flag for AWS Builder ID device code flow
- Add DoKiroAWSLogin function for AWS SSO OIDC authentication
- Complete Kiro integration with AWS, Google OAuth, and social auth
- Add kiro executor, translator, and SDK components
- Update browser support for Kiro authentication flows
2025-12-05 22:46:24 +03:00
Luis Pater
92f033dec0 Merge branch 'router-for-me:main' into main 2025-12-06 01:33:34 +08:00
Luis Pater
0ebabf5152 feat(antigravity): add FetchAntigravityProjectID function and integrate project ID retrieval 2025-12-06 01:32:12 +08:00
Luis Pater
4b01ecba2e Merge branch 'router-for-me:main' into main 2025-12-06 01:12:43 +08:00
Luis Pater
d7564173dd fix(antigravity): restore production base URL in the executor 2025-12-06 01:11:37 +08:00
Luis Pater
f241124599 Merge branch 'router-for-me:main' into main 2025-12-06 00:43:02 +08:00
Luis Pater
c44c46dd80 Fixed: #421
feat(antigravity): implement project ID retrieval and integration in payload processing
2025-12-06 00:40:55 +08:00
Luis Pater
aa810ee719 Merge branch 'router-for-me:main' into main 2025-12-05 23:06:46 +08:00
Luis Pater
412148af0e feat(antigravity): add function ID to FunctionCall and FunctionResponse models 2025-12-05 23:05:35 +08:00
Luis Pater
5d2baf6058 Merge branch 'router-for-me:main' into main 2025-12-05 21:37:55 +08:00
Luis Pater
d28258501a Merge pull request #423 from router-for-me/amp
fix(amp): suppress ErrAbortHandler panics in reverse proxy handler
2025-12-05 21:36:01 +08:00
hkfires
55cd31fb96 fix(amp): suppress ErrAbortHandler panics in reverse proxy handler 2025-12-05 21:28:58 +08:00
Luis Pater
d138df07bf Merge branch 'router-for-me:main' into main 2025-12-05 21:28:32 +08:00
Luis Pater
c5df8e7897 Merge pull request #422 from router-for-me/amp
Amp
2025-12-05 21:25:43 +08:00
Luis Pater
d4d529833d **refactor(antigravity): handle anyOf property, remove exclusiveMinimum, and comment unused prod URL** 2025-12-05 21:24:12 +08:00
hkfires
caa48e7c6f fix(amp): improve proxy state management and request logging behavior 2025-12-05 21:09:53 +08:00
hkfires
acdfb3bceb feat(amp): add root-level /threads routes for CLI compatibility 2025-12-05 18:14:10 +08:00
hkfires
89d68962b1 fix(amp): filter amp request logging to only provider endpoint 2025-12-05 18:14:09 +08:00
Luis Pater
691cdb6bdf **fix(api): update GitHub release URL and user agent for CLIProxyAPIPlus** 2025-12-05 10:32:28 +08:00
Luis Pater
8064cba288 Merge branch 'router-for-me:main' into main 2025-12-05 10:31:30 +08:00
Luis Pater
361443db10 **feat(api): add GetLatestVersion endpoint to fetch latest release version from GitHub** 2025-12-05 10:29:12 +08:00
Luis Pater
d6352dd4d4 **feat(util): add DeleteKey function and update antigravity executor for Claude model compatibility** 2025-12-05 01:55:45 +08:00
Luis Pater
f8aba62860 Merge branch 'router-for-me:main' into main 2025-12-05 00:45:51 +08:00
Luis Pater
a7eeb06f3d Merge pull request #418 from router-for-me/amp
Amp
2025-12-05 00:43:15 +08:00
hkfires
9426be7a5c fix(amp): update log message wording for disabled proxy state 2025-12-04 21:36:16 +08:00
hkfires
4a135f1986 feat(amp): add hot-reload support for upstream URL and localhost restriction 2025-12-04 21:30:59 +08:00
hkfires
c4c02f4ad0 feat(amp): add partial reload support with config change detection 2025-12-04 21:30:59 +08:00
Luis Pater
b87b9b455f Merge pull request #416 from router-for-me/amp
Amp
2025-12-04 20:52:33 +08:00
hkfires
db03ae9663 feat(watcher): add AmpCode config change detection 2025-12-04 19:50:54 +08:00
hkfires
969ff6bb68 fix(amp): update explicit API key on config change 2025-12-04 19:32:44 +08:00
Luis Pater
e24af6e545 Merge branch 'router-for-me:main' into main 2025-12-04 18:03:42 +08:00
Luis Pater
bceecfb2e3 Fixed: #414
**refactor(gemini): comment out unused CLI preview entry**
2025-12-04 17:55:13 +08:00
Luis Pater
48dd987867 Merge branch 'router-for-me:main' into main 2025-12-04 16:14:52 +08:00
Luis Pater
6a2906e3e5 **feat(antigravity): add support for Claude-Opus-4-5-Thinking model** 2025-12-04 16:13:13 +08:00
Luis Pater
d72886c801 Merge pull request #405 from thurstonsand/fix/amp-missing-proxy-routes
fix(amp): add missing /auth/* and /api/tab/* proxy routes for AMP CLI
2025-12-03 22:23:41 +08:00
Luis Pater
6efba3d829 Merge pull request #406 from router-for-me/api
refactor(api): remove legacy generative-language-api-key endpoints and duplicate GetConfigYAML
2025-12-03 22:21:15 +08:00
Luis Pater
6b60bdd139 Merge branch 'router-for-me:main' into main 2025-12-03 21:55:29 +08:00
Luis Pater
897c40bed8 feat(registry): add DeepSeek-V3.2-Chat model definition
Add new DeepSeek-V3.2-Chat model to the registry with standard chat configuration, positioned before the experimental variant for better organization.
2025-12-03 21:34:50 +08:00
Thurston Sandberg
373ea8d7e4 fix(logging): handle nil caller in LogFormatter to prevent panic 2025-12-03 05:54:38 -05:00
hkfires
b5de004c01 refactor(api): remove legacy generative-language-api-key endpoints and duplicate GetConfigYAML 2025-12-03 18:35:08 +08:00
Thurston Sandberg
94ec772521 test(amp): add tests for /auth/* and /api/tab/* routes 2025-12-03 05:03:25 -05:00
Thurston Sandberg
e216d26731 fix(amp): add missing /auth/* and /api/tab/* proxy routes 2025-12-03 04:40:32 -05:00
Luis Pater
8a98506a15 Merge branch 'router-for-me:main' into main 2025-12-03 16:20:02 +08:00
Luis Pater
6eb94dac33 Merge pull request #404 from router-for-me/config
Legacy Config Migration and Amp Consolidation
2025-12-03 16:11:06 +08:00
hkfires
c4a5be6edf style(amp): standardize log message capitalization 2025-12-03 13:53:18 +08:00
hkfires
651179a642 refactor(config): add detailed logging for legacy configuration migration 2025-12-03 13:39:10 +08:00
hkfires
8c42b21e66 refactor(config): improve OpenAI compatibility target matching logic 2025-12-03 12:41:17 +08:00
hkfires
b693d632d2 docs(config): comment out example API key configurations 2025-12-03 12:31:41 +08:00
hkfires
b5033c22d8 refactor(config): auto-persist migrated legacy configuration fields 2025-12-03 12:26:04 +08:00
hkfires
df0fd1add1 refactor(config): remove deprecated AMP configuration keys during save 2025-12-03 11:42:15 +08:00
hkfires
b6bdbe78ef refactor(config): relocate legacy migration helpers to end of file 2025-12-03 11:23:11 +08:00
hkfires
06c0d2bab2 refactor(config): remove deprecated legacy API key fields 2025-12-03 11:01:56 +08:00
hkfires
bd1678457b refactor(config): consolidate Amp settings into AmpCode struct 2025-12-03 10:42:28 +08:00
hkfires
559b7df404 refactor(config): restructure and uncomment example configuration 2025-12-03 10:29:36 +08:00
Luis Pater
dcd0ff834e Merge branch 'router-for-me:main' into main 2025-12-03 07:24:42 +08:00
Luis Pater
8b13c91132 **docs(internal): add Codex instruction guides for GPT-5 CLI**
- Added `gpt_5_1_prompt.md` and `gpt_5_codex_prompt.md` to document Codex instruction guidelines.
- These detail the behavior, constraints, and execution policies for GPT-5-based Codex agents in the CLI environment.
2025-12-03 07:23:01 +08:00
Luis Pater
e93f87294a refactor(antigravity): uncomment prod environment URL in fallback chain 2025-12-02 22:47:18 +08:00
Luis Pater
a6fa622aca Merge branch 'router-for-me:main' into main 2025-12-02 22:41:23 +08:00
Luis Pater
a67b6811d1 Fixed: #397
fix(auth): use proxy HTTP client for Gemini CLI token requests
2025-12-02 22:39:01 +08:00
Luis Pater
6ee95b18f7 Merge branch 'router-for-me:main' into main 2025-12-02 22:29:53 +08:00
Luis Pater
35fdc4cfd3 fix some bugs (#399)
* feat(config): add pruning of stale YAML mapping keys during config save

* Revert watcher.go in "fix: enable hot reload for amp-model-mappings config"
2025-12-02 22:28:30 +08:00
hkfires
3ebbab0a9a Revert watcher.go in "fix: enable hot reload for amp-model-mappings config" 2025-12-02 22:17:54 +08:00
hkfires
480cd714b2 feat(config): add pruning of stale YAML mapping keys during config save 2025-12-02 21:38:54 +08:00
Luis Pater
8afcc710d4 Merge branch 'router-for-me:main' into main 2025-12-02 18:34:13 +08:00
Luis Pater
41ee44432d **fix(translator): rename responseSchema key for generationConfig**
- Renamed `generationConfig.responseSchema` to `generationConfig.responseJsonSchema` in Gemini request transformation to align with updated schema expectations.
2025-12-02 18:32:23 +08:00
Luis Pater
4419b71699 Merge pull request #2 from router-for-me/v6.5.32
v6.5.32
2025-12-02 11:46:42 +08:00
Luis Pater
43cac7b5f6 Merge branch 'main' into v6.5.32 2025-12-02 11:46:05 +08:00
Luis Pater
1434bc38e5 **refactor(registry): remove Qwen3-Coder from model definitions** 2025-12-02 11:34:38 +08:00
Luis Pater
0fd2abbc3b **refactor(cliproxy, config): remove vertex-compat flow, streamline Vertex API key handling**
- Removed `vertex-compat` executor and related configuration.
- Consolidated Vertex compatibility checks into `vertex` handling with `apikey`-based model resolution.
- Streamlined model generation logic for Vertex API key entries.
2025-12-02 09:18:24 +08:00
Aero
0ebb654019 feat: Add support for VertexAI compatible service (#375)
feat: consolidate Vertex AI compatibility with API key support in Gemini
2025-12-02 08:14:22 +08:00
Luis Pater
08a1d2edf9 Merge pull request #390 from NguyenSiTrung/main
feat(amp): add model mapping support for routing unavailable models to alternatives
2025-12-02 08:07:56 +08:00
NguyenSiTrung
3409f4e336 fix: enable hot reload for amp-model-mappings config
- Store ampModule in Server struct to access it during config updates
- Call ampModule.OnConfigUpdated() in UpdateClients() for hot reload
- Watch config directory instead of file to handle atomic saves (vim, VSCode, etc.)
- Improve config file event detection with basename matching
- Add diagnostic logging for config reload tracing
2025-12-01 13:34:49 +07:00
NguyenSiTrung
9354b87e54 Merge branch 'router-for-me:main' into main 2025-12-01 08:12:29 +07:00
Luis Pater
c6ecdf4a7e Merge pull request #1 from router-for-me/Plus
v6.5.31
2025-12-01 09:10:55 +08:00
Luis Pater
838f47a7e9 Merge branch 'main' into Plus 2025-12-01 09:10:29 +08:00
Luis Pater
54e24110ec Merge pull request #386 from auroraflux/feat/dedupe-thinking-metadata-helpers
refactor(executor): dedupe thinking metadata helpers across Gemini executors
2025-12-01 09:00:27 +08:00
Luis Pater
717c703bff docs(readme): add CCS (Claude Code Switch) to projects list 2025-12-01 07:22:42 +08:00
auroraflux
1c6f4be8ae refactor(executor): dedupe thinking metadata helpers across Gemini executors
Extract applyThinkingMetadata and applyThinkingMetadataCLI helpers to
payload_helpers.go and use them across all four Gemini-based executors:
- gemini_executor.go (Execute, ExecuteStream, CountTokens)
- gemini_cli_executor.go (Execute, ExecuteStream, CountTokens)
- aistudio_executor.go (translateRequest)
- antigravity_executor.go (Execute, ExecuteStream)

This eliminates code duplication introduced in the -reasoning suffix PR
and centralizes the thinking config application logic.

Net reduction: 28 lines of code.
2025-11-30 15:20:15 -08:00
Luis Pater
0de2560cee Merge pull request #379 from kaitranntt/docs/add-ccs-project
docs: add CCS (Claude Code Switch) to projects list
2025-12-01 07:20:04 +08:00
Kai (Tam Nhu) Tran
85eb926482 fix: change AGY to Antigravity 2025-11-30 12:43:12 -05:00
Kai (Tam Nhu) Tran
c52ef08e67 docs: add CCS to projects list 2025-11-30 12:40:35 -05:00
Luis Pater
5b01eba943 Merge branch 'router-for-me:main' into main 2025-11-30 21:30:49 +08:00
Luis Pater
cb580cd083 Merge pull request #377 from router-for-me/gemini
feat(registry): add thinking support to gemini models
2025-11-30 21:27:54 +08:00
hkfires
75e278c7a5 feat(registry): add thinking support to gemini models 2025-11-30 20:56:29 +08:00
Luis Pater
73208c4e55 Merge pull request #376 from auroraflux/feat/reasoning-suffix-support
feat(util): add -reasoning suffix support for Gemini models
2025-11-30 20:55:38 +08:00
auroraflux
32d3809f8c **feat(util): add -reasoning suffix support for Gemini models**
Adds support for the `-reasoning` model name suffix which enables
thinking/reasoning mode with dynamic budget. This allows clients to
request reasoning-enabled inference using model names like
`gemini-2.5-flash-reasoning` without explicit configuration.

The suffix is normalized to the base model (e.g., gemini-2.5-flash)
with thinkingBudget=-1 (dynamic) and include_thoughts=true.

Follows the existing pattern established by -nothinking and
-thinking-N suffixes.
2025-11-30 01:18:57 -08:00
Trung Nguyen
33a5656235 docs: add model mapping documentation for Amp CLI integration
- Add model mapping feature to README.md Amp CLI section
- Add detailed Model Mapping Configuration section to amp-cli-integration.md
- Update architecture diagram to show model mapping flow
- Update Model Fallback Behavior to include mapping step
- Add Table of Contents entry for model mapping
2025-11-29 12:51:03 +07:00
Trung Nguyen
2cd59806e2 feat(amp): add model mapping support for routing unavailable models to alternatives
- Add AmpModelMapping config to route models like 'claude-opus-4.5' to 'claude-sonnet-4'
- Add ModelMapper interface and DefaultModelMapper implementation with hot-reload support
- Enhance FallbackHandler to apply model mappings before falling back to ampcode.com
- Add structured logging for routing decisions (local provider, mapping, amp credits)
- Update config.example.yaml with amp-model-mappings documentation
2025-11-29 12:44:09 +07:00
97 changed files with 12497 additions and 976 deletions

7
.gitignore vendored
View File

@@ -1,5 +1,6 @@
# Binaries
cli-proxy-api
cliproxy
*.exe
# Configuration
@@ -15,6 +16,7 @@ pgstore/*
gitstore/*
objectstore/*
static/*
refs/*
# Authentication data
auths/*
@@ -30,3 +32,8 @@ GEMINI.md
.vscode/*
.claude/*
.serena/*
.mcp/cache/
# macOS
.DS_Store
._*

View File

@@ -10,7 +10,8 @@ The Plus release stays in lockstep with the mainline features.
## Differences from the Mainline
- Added GitHub Copilot support (OAuth login), provided by [em4gp](https://github.com/em4go/CLIProxyAPI/tree/feature/github-copilot-auth)
- Added GitHub Copilot support (OAuth login), provided by [em4go](https://github.com/em4go/CLIProxyAPI/tree/feature/github-copilot-auth)
- Added Kiro (AWS CodeWhisperer) support (OAuth login), provided by [fuko2935](https://github.com/fuko2935/CLIProxyAPI/tree/feature/kiro-integration)
## Contributing

View File

@@ -10,7 +10,8 @@
## 与主线版本版本差异
- 新增 GitHub Copilot 支持OAuth 登录),由[em4gp](https://github.com/em4go/CLIProxyAPI/tree/feature/github-copilot-auth)提供
- 新增 GitHub Copilot 支持OAuth 登录),由[em4go](https://github.com/em4go/CLIProxyAPI/tree/feature/github-copilot-auth)提供
- 新增 Kiro (AWS CodeWhisperer) 支持 (OAuth 登录), 由[fuko2935](https://github.com/fuko2935/CLIProxyAPI/tree/feature/kiro-integration)提供
## 贡献

View File

@@ -47,6 +47,19 @@ func init() {
buildinfo.BuildDate = BuildDate
}
// setKiroIncognitoMode sets the incognito browser mode for Kiro authentication.
// Kiro defaults to incognito mode for multi-account support.
// Users can explicitly override with --incognito or --no-incognito flags.
func setKiroIncognitoMode(cfg *config.Config, useIncognito, noIncognito bool) {
if useIncognito {
cfg.IncognitoBrowser = true
} else if noIncognito {
cfg.IncognitoBrowser = false
} else {
cfg.IncognitoBrowser = true // Kiro default
}
}
// main is the entry point of the application.
// It parses command-line flags, loads configuration, and starts the appropriate
// service based on the provided flags (login, codex-login, or server mode).
@@ -62,11 +75,17 @@ func main() {
var iflowCookie bool
var noBrowser bool
var antigravityLogin bool
var kiroLogin bool
var kiroGoogleLogin bool
var kiroAWSLogin bool
var kiroImport bool
var githubCopilotLogin bool
var projectID string
var vertexImport string
var configPath string
var password string
var noIncognito bool
var useIncognito bool
// Define command-line flags for different operation modes.
flag.BoolVar(&login, "login", false, "Login Google Account")
@@ -76,7 +95,13 @@ func main() {
flag.BoolVar(&iflowLogin, "iflow-login", false, "Login to iFlow using OAuth")
flag.BoolVar(&iflowCookie, "iflow-cookie", false, "Login to iFlow using Cookie")
flag.BoolVar(&noBrowser, "no-browser", false, "Don't open browser automatically for OAuth")
flag.BoolVar(&useIncognito, "incognito", false, "Open browser in incognito/private mode for OAuth (useful for multiple accounts)")
flag.BoolVar(&noIncognito, "no-incognito", false, "Force disable incognito mode (uses existing browser session)")
flag.BoolVar(&antigravityLogin, "antigravity-login", false, "Login to Antigravity using OAuth")
flag.BoolVar(&kiroLogin, "kiro-login", false, "Login to Kiro using Google OAuth")
flag.BoolVar(&kiroGoogleLogin, "kiro-google-login", false, "Login to Kiro using Google OAuth (same as --kiro-login)")
flag.BoolVar(&kiroAWSLogin, "kiro-aws-login", false, "Login to Kiro using AWS Builder ID (device code flow)")
flag.BoolVar(&kiroImport, "kiro-import", false, "Import Kiro token from Kiro IDE (~/.aws/sso/cache/kiro-auth-token.json)")
flag.BoolVar(&githubCopilotLogin, "github-copilot-login", false, "Login to GitHub Copilot using device flow")
flag.StringVar(&projectID, "project_id", "", "Project ID (Gemini only, not required)")
flag.StringVar(&configPath, "config", DefaultConfigPath, "Configure File Path")
@@ -141,7 +166,8 @@ func main() {
wd, err := os.Getwd()
if err != nil {
log.Fatalf("failed to get working directory: %v", err)
log.Errorf("failed to get working directory: %v", err)
return
}
// Load environment variables from .env if present.
@@ -235,13 +261,15 @@ func main() {
})
cancel()
if err != nil {
log.Fatalf("failed to initialize postgres token store: %v", err)
log.Errorf("failed to initialize postgres token store: %v", err)
return
}
examplePath := filepath.Join(wd, "config.example.yaml")
ctx, cancel = context.WithTimeout(context.Background(), 30*time.Second)
if errBootstrap := pgStoreInst.Bootstrap(ctx, examplePath); errBootstrap != nil {
cancel()
log.Fatalf("failed to bootstrap postgres-backed config: %v", errBootstrap)
log.Errorf("failed to bootstrap postgres-backed config: %v", errBootstrap)
return
}
cancel()
configFilePath = pgStoreInst.ConfigPath()
@@ -264,7 +292,8 @@ func main() {
if strings.Contains(resolvedEndpoint, "://") {
parsed, errParse := url.Parse(resolvedEndpoint)
if errParse != nil {
log.Fatalf("failed to parse object store endpoint %q: %v", objectStoreEndpoint, errParse)
log.Errorf("failed to parse object store endpoint %q: %v", objectStoreEndpoint, errParse)
return
}
switch strings.ToLower(parsed.Scheme) {
case "http":
@@ -272,10 +301,12 @@ func main() {
case "https":
useSSL = true
default:
log.Fatalf("unsupported object store scheme %q (only http and https are allowed)", parsed.Scheme)
log.Errorf("unsupported object store scheme %q (only http and https are allowed)", parsed.Scheme)
return
}
if parsed.Host == "" {
log.Fatalf("object store endpoint %q is missing host information", objectStoreEndpoint)
log.Errorf("object store endpoint %q is missing host information", objectStoreEndpoint)
return
}
resolvedEndpoint = parsed.Host
if parsed.Path != "" && parsed.Path != "/" {
@@ -294,13 +325,15 @@ func main() {
}
objectStoreInst, err = store.NewObjectTokenStore(objCfg)
if err != nil {
log.Fatalf("failed to initialize object token store: %v", err)
log.Errorf("failed to initialize object token store: %v", err)
return
}
examplePath := filepath.Join(wd, "config.example.yaml")
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
if errBootstrap := objectStoreInst.Bootstrap(ctx, examplePath); errBootstrap != nil {
cancel()
log.Fatalf("failed to bootstrap object-backed config: %v", errBootstrap)
log.Errorf("failed to bootstrap object-backed config: %v", errBootstrap)
return
}
cancel()
configFilePath = objectStoreInst.ConfigPath()
@@ -325,7 +358,8 @@ func main() {
gitStoreInst = store.NewGitTokenStore(gitStoreRemoteURL, gitStoreUser, gitStorePassword)
gitStoreInst.SetBaseDir(authDir)
if errRepo := gitStoreInst.EnsureRepository(); errRepo != nil {
log.Fatalf("failed to prepare git token store: %v", errRepo)
log.Errorf("failed to prepare git token store: %v", errRepo)
return
}
configFilePath = gitStoreInst.ConfigPath()
if configFilePath == "" {
@@ -334,17 +368,21 @@ func main() {
if _, statErr := os.Stat(configFilePath); errors.Is(statErr, fs.ErrNotExist) {
examplePath := filepath.Join(wd, "config.example.yaml")
if _, errExample := os.Stat(examplePath); errExample != nil {
log.Fatalf("failed to find template config file: %v", errExample)
log.Errorf("failed to find template config file: %v", errExample)
return
}
if errCopy := misc.CopyConfigTemplate(examplePath, configFilePath); errCopy != nil {
log.Fatalf("failed to bootstrap git-backed config: %v", errCopy)
log.Errorf("failed to bootstrap git-backed config: %v", errCopy)
return
}
if errCommit := gitStoreInst.PersistConfig(context.Background()); errCommit != nil {
log.Fatalf("failed to commit initial git-backed config: %v", errCommit)
log.Errorf("failed to commit initial git-backed config: %v", errCommit)
return
}
log.Infof("git-backed config initialized from template: %s", configFilePath)
} else if statErr != nil {
log.Fatalf("failed to inspect git-backed config: %v", statErr)
log.Errorf("failed to inspect git-backed config: %v", statErr)
return
}
cfg, err = config.LoadConfigOptional(configFilePath, isCloudDeploy)
if err == nil {
@@ -357,13 +395,15 @@ func main() {
} else {
wd, err = os.Getwd()
if err != nil {
log.Fatalf("failed to get working directory: %v", err)
log.Errorf("failed to get working directory: %v", err)
return
}
configFilePath = filepath.Join(wd, "config.yaml")
cfg, err = config.LoadConfigOptional(configFilePath, isCloudDeploy)
}
if err != nil {
log.Fatalf("failed to load config: %v", err)
log.Errorf("failed to load config: %v", err)
return
}
if cfg == nil {
cfg = &config.Config{}
@@ -393,7 +433,8 @@ func main() {
coreauth.SetQuotaCooldownDisabled(cfg.DisableCooling)
if err = logging.ConfigureLogOutput(cfg.LoggingToFile); err != nil {
log.Fatalf("failed to configure log output: %v", err)
log.Errorf("failed to configure log output: %v", err)
return
}
log.Infof("CLIProxyAPI Version: %s, Commit: %s, BuiltAt: %s", buildinfo.Version, buildinfo.Commit, buildinfo.BuildDate)
@@ -402,7 +443,8 @@ func main() {
util.SetLogLevel(cfg)
if resolvedAuthDir, errResolveAuthDir := util.ResolveAuthDir(cfg.AuthDir); errResolveAuthDir != nil {
log.Fatalf("failed to resolve auth directory: %v", errResolveAuthDir)
log.Errorf("failed to resolve auth directory: %v", errResolveAuthDir)
return
} else {
cfg.AuthDir = resolvedAuthDir
}
@@ -453,6 +495,26 @@ func main() {
cmd.DoIFlowLogin(cfg, options)
} else if iflowCookie {
cmd.DoIFlowCookieAuth(cfg, options)
} else if kiroLogin {
// For Kiro auth, default to incognito mode for multi-account support
// Users can explicitly override with --no-incognito
// Note: This config mutation is safe - auth commands exit after completion
// and don't share config with StartService (which is in the else branch)
setKiroIncognitoMode(cfg, useIncognito, noIncognito)
cmd.DoKiroLogin(cfg, options)
} else if kiroGoogleLogin {
// For Kiro auth, default to incognito mode for multi-account support
// Users can explicitly override with --no-incognito
// Note: This config mutation is safe - auth commands exit after completion
setKiroIncognitoMode(cfg, useIncognito, noIncognito)
cmd.DoKiroGoogleLogin(cfg, options)
} else if kiroAWSLogin {
// For Kiro auth, default to incognito mode for multi-account support
// Users can explicitly override with --no-incognito
setKiroIncognitoMode(cfg, useIncognito, noIncognito)
cmd.DoKiroAWSLogin(cfg, options)
} else if kiroImport {
cmd.DoKiroImport(cfg, options)
} else {
// In cloud deploy mode without config file, just wait for shutdown signals
if isCloudDeploy && !configFileExists {

View File

@@ -1,3 +1,7 @@
# Server host/interface to bind to. Default is empty ("") to bind all interfaces (IPv4 + IPv6).
# Use "127.0.0.1" or "localhost" to restrict access to local machine only.
host: ""
# Server port
port: 8317
@@ -32,6 +36,11 @@ api-keys:
# Enable debug logging
debug: false
# Open OAuth URLs in incognito/private browser mode.
# Useful when you want to login with a different account without logging out from your current session.
# Default: false (but Kiro auth defaults to true for multi-account support)
incognito-browser: true
# When true, write application logs to rotating files instead of stdout
logging-to-file: false
@@ -55,106 +64,142 @@ quota-exceeded:
# When true, enable authentication for the WebSocket API (/v1/ws).
ws-auth: false
# Gemini API keys (preferred)
#gemini-api-key:
# - api-key: "AIzaSy...01"
# base-url: "https://generativelanguage.googleapis.com"
# headers:
# X-Custom-Header: "custom-value"
# proxy-url: "socks5://proxy.example.com:1080"
# excluded-models:
# - "gemini-2.5-pro" # exclude specific models from this provider (exact match)
# - "gemini-2.5-*" # wildcard matching prefix (e.g. gemini-2.5-flash, gemini-2.5-pro)
# - "*-preview" # wildcard matching suffix (e.g. gemini-3-pro-preview)
# - "*flash*" # wildcard matching substring (e.g. gemini-2.5-flash-lite)
# - api-key: "AIzaSy...02"
# API keys for official Generative Language API (legacy compatibility)
#generative-language-api-key:
# - "AIzaSy...01"
# - "AIzaSy...02"
# Gemini API keys
# gemini-api-key:
# - api-key: "AIzaSy...01"
# base-url: "https://generativelanguage.googleapis.com"
# headers:
# X-Custom-Header: "custom-value"
# proxy-url: "socks5://proxy.example.com:1080"
# excluded-models:
# - "gemini-2.5-pro" # exclude specific models from this provider (exact match)
# - "gemini-2.5-*" # wildcard matching prefix (e.g. gemini-2.5-flash, gemini-2.5-pro)
# - "*-preview" # wildcard matching suffix (e.g. gemini-3-pro-preview)
# - "*flash*" # wildcard matching substring (e.g. gemini-2.5-flash-lite)
# - api-key: "AIzaSy...02"
# Codex API keys
#codex-api-key:
# - api-key: "sk-atSM..."
# base-url: "https://www.example.com" # use the custom codex API endpoint
# headers:
# X-Custom-Header: "custom-value"
# proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override
# excluded-models:
# - "gpt-5.1" # exclude specific models (exact match)
# - "gpt-5-*" # wildcard matching prefix (e.g. gpt-5-medium, gpt-5-codex)
# - "*-mini" # wildcard matching suffix (e.g. gpt-5-codex-mini)
# - "*codex*" # wildcard matching substring (e.g. gpt-5-codex-low)
# codex-api-key:
# - api-key: "sk-atSM..."
# base-url: "https://www.example.com" # use the custom codex API endpoint
# headers:
# X-Custom-Header: "custom-value"
# proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override
# excluded-models:
# - "gpt-5.1" # exclude specific models (exact match)
# - "gpt-5-*" # wildcard matching prefix (e.g. gpt-5-medium, gpt-5-codex)
# - "*-mini" # wildcard matching suffix (e.g. gpt-5-codex-mini)
# - "*codex*" # wildcard matching substring (e.g. gpt-5-codex-low)
# Claude API keys
#claude-api-key:
# - api-key: "sk-atSM..." # use the official claude API key, no need to set the base url
# - api-key: "sk-atSM..."
# base-url: "https://www.example.com" # use the custom claude API endpoint
# headers:
# X-Custom-Header: "custom-value"
# proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override
# models:
# - name: "claude-3-5-sonnet-20241022" # upstream model name
# alias: "claude-sonnet-latest" # client alias mapped to the upstream model
# excluded-models:
# - "claude-opus-4-5-20251101" # exclude specific models (exact match)
# - "claude-3-*" # wildcard matching prefix (e.g. claude-3-7-sonnet-20250219)
# - "*-think" # wildcard matching suffix (e.g. claude-opus-4-5-thinking)
# - "*haiku*" # wildcard matching substring (e.g. claude-3-5-haiku-20241022)
# claude-api-key:
# - api-key: "sk-atSM..." # use the official claude API key, no need to set the base url
# - api-key: "sk-atSM..."
# base-url: "https://www.example.com" # use the custom claude API endpoint
# headers:
# X-Custom-Header: "custom-value"
# proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override
# models:
# - name: "claude-3-5-sonnet-20241022" # upstream model name
# alias: "claude-sonnet-latest" # client alias mapped to the upstream model
# excluded-models:
# - "claude-opus-4-5-20251101" # exclude specific models (exact match)
# - "claude-3-*" # wildcard matching prefix (e.g. claude-3-7-sonnet-20250219)
# - "*-think" # wildcard matching suffix (e.g. claude-opus-4-5-thinking)
# - "*haiku*" # wildcard matching substring (e.g. claude-3-5-haiku-20241022)
# Kiro (AWS CodeWhisperer) configuration
# Note: Kiro API currently only operates in us-east-1 region
#kiro:
# - token-file: "~/.aws/sso/cache/kiro-auth-token.json" # path to Kiro token file
# agent-task-type: "" # optional: "vibe" or empty (API default)
# - access-token: "aoaAAAAA..." # or provide tokens directly
# refresh-token: "aorAAAAA..."
# profile-arn: "arn:aws:codewhisperer:us-east-1:..."
# proxy-url: "socks5://proxy.example.com:1080" # optional: proxy override
# OpenAI compatibility providers
#openai-compatibility:
# - name: "openrouter" # The name of the provider; it will be used in the user agent and other places.
# base-url: "https://openrouter.ai/api/v1" # The base URL of the provider.
# headers:
# X-Custom-Header: "custom-value"
# # New format with per-key proxy support (recommended):
# api-key-entries:
# - api-key: "sk-or-v1-...b780"
# proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override
# - api-key: "sk-or-v1-...b781" # without proxy-url
# # Legacy format (still supported, but cannot specify proxy per key):
# # api-keys:
# # - "sk-or-v1-...b780"
# # - "sk-or-v1-...b781"
# models: # The models supported by the provider.
# - name: "moonshotai/kimi-k2:free" # The actual model name.
# alias: "kimi-k2" # The alias used in the API.
# openai-compatibility:
# - name: "openrouter" # The name of the provider; it will be used in the user agent and other places.
# base-url: "https://openrouter.ai/api/v1" # The base URL of the provider.
# headers:
# X-Custom-Header: "custom-value"
# api-key-entries:
# - api-key: "sk-or-v1-...b780"
# proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override
# - api-key: "sk-or-v1-...b781" # without proxy-url
# models: # The models supported by the provider.
# - name: "moonshotai/kimi-k2:free" # The actual model name.
# alias: "kimi-k2" # The alias used in the API.
#payload: # Optional payload configuration
# default: # Default rules only set parameters when they are missing in the payload.
# - models:
# - name: "gemini-2.5-pro" # Supports wildcards (e.g., "gemini-*")
# protocol: "gemini" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex
# params: # JSON path (gjson/sjson syntax) -> value
# "generationConfig.thinkingConfig.thinkingBudget": 32768
# override: # Override rules always set parameters, overwriting any existing values.
# - models:
# - name: "gpt-*" # Supports wildcards (e.g., "gpt-*")
# protocol: "codex" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex
# params: # JSON path (gjson/sjson syntax) -> value
# "reasoning.effort": "high"
# Vertex API keys (Vertex-compatible endpoints, use API key + base URL)
# vertex-api-key:
# - api-key: "vk-123..." # x-goog-api-key header
# base-url: "https://example.com/api" # e.g. https://zenmux.ai/api
# proxy-url: "socks5://proxy.example.com:1080" # optional per-key proxy override
# headers:
# X-Custom-Header: "custom-value"
# models: # optional: map aliases to upstream model names
# - name: "gemini-2.0-flash" # upstream model name
# alias: "vertex-flash" # client-visible alias
# - name: "gemini-1.5-pro"
# alias: "vertex-pro"
# Amp Integration
# ampcode:
# # Configure upstream URL for Amp CLI OAuth and management features
# upstream-url: "https://ampcode.com"
# # Optional: Override API key for Amp upstream (otherwise uses env or file)
# upstream-api-key: ""
# # Restrict Amp management routes (/api/auth, /api/user, etc.) to localhost only (recommended)
# restrict-management-to-localhost: true
# # Force model mappings to run before checking local API keys (default: false)
# force-model-mappings: false
# # Amp Model Mappings
# # Route unavailable Amp models to alternative models available in your local proxy.
# # Useful when Amp CLI requests models you don't have access to (e.g., Claude Opus 4.5)
# # but you have a similar model available (e.g., Claude Sonnet 4).
# model-mappings:
# - from: "claude-opus-4.5" # Model requested by Amp CLI
# to: "claude-sonnet-4" # Route to this available model instead
# - from: "gpt-5"
# to: "gemini-2.5-pro"
# - from: "claude-3-opus-20240229"
# to: "claude-3-5-sonnet-20241022"
# OAuth provider excluded models
#oauth-excluded-models:
# gemini-cli:
# - "gemini-2.5-pro" # exclude specific models (exact match)
# - "gemini-2.5-*" # wildcard matching prefix (e.g. gemini-2.5-flash, gemini-2.5-pro)
# - "*-preview" # wildcard matching suffix (e.g. gemini-3-pro-preview)
# - "*flash*" # wildcard matching substring (e.g. gemini-2.5-flash-lite)
# vertex:
# - "gemini-3-pro-preview"
# aistudio:
# - "gemini-3-pro-preview"
# antigravity:
# - "gemini-3-pro-preview"
# claude:
# - "claude-3-5-haiku-20241022"
# codex:
# - "gpt-5-codex-mini"
# qwen:
# - "vision-model"
# iflow:
# - "tstars2.0"
# oauth-excluded-models:
# gemini-cli:
# - "gemini-2.5-pro" # exclude specific models (exact match)
# - "gemini-2.5-*" # wildcard matching prefix (e.g. gemini-2.5-flash, gemini-2.5-pro)
# - "*-preview" # wildcard matching suffix (e.g. gemini-3-pro-preview)
# - "*flash*" # wildcard matching substring (e.g. gemini-2.5-flash-lite)
# vertex:
# - "gemini-3-pro-preview"
# aistudio:
# - "gemini-3-pro-preview"
# antigravity:
# - "gemini-3-pro-preview"
# claude:
# - "claude-3-5-haiku-20241022"
# codex:
# - "gpt-5-codex-mini"
# qwen:
# - "vision-model"
# iflow:
# - "tstars2.0"
# Optional payload configuration
# payload:
# default: # Default rules only set parameters when they are missing in the payload.
# - models:
# - name: "gemini-2.5-pro" # Supports wildcards (e.g., "gemini-*")
# protocol: "gemini" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex
# params: # JSON path (gjson/sjson syntax) -> value
# "generationConfig.thinkingConfig.thinkingBudget": 32768
# override: # Override rules always set parameters, overwriting any existing values.
# - models:
# - name: "gpt-*" # Supports wildcards (e.g., "gpt-*")
# protocol: "codex" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex
# params: # JSON path (gjson/sjson syntax) -> value
# "reasoning.effort": "high"

View File

@@ -8,6 +8,7 @@ This guide explains how to use CLIProxyAPI with Amp CLI and Amp IDE extensions,
- [Which Providers Should You Authenticate?](#which-providers-should-you-authenticate)
- [Architecture](#architecture)
- [Configuration](#configuration)
- [Model Mapping Configuration](#model-mapping-configuration)
- [Setup](#setup)
- [Usage](#usage)
- [Troubleshooting](#troubleshooting)
@@ -21,6 +22,7 @@ The Amp CLI integration adds specialized routing to support Amp's API patterns w
- **Provider route aliases**: Maps Amp's `/api/provider/{provider}/v1...` patterns to CLIProxyAPI handlers
- **Management proxy**: Forwards OAuth and account management requests to Amp's control plane
- **Smart fallback**: Automatically routes unconfigured models to ampcode.com
- **Model mapping**: Route unavailable models to alternatives you have access to (e.g., `claude-opus-4.5``claude-sonnet-4`)
- **Secret management**: Configurable precedence (config > env > file) with 5-minute caching
- **Security-first**: Management routes restricted to localhost by default
- **Automatic gzip handling**: Decompresses responses from Amp upstream
@@ -75,7 +77,10 @@ Amp CLI/IDE
│ ↓
│ ├─ Model configured locally?
│ │ YES → Use local OAuth tokens (OpenAI/Claude/Gemini handlers)
│ │ NO → Forward to ampcode.com (reverse proxy)
│ │ NO
│ │ ├─ Model mapping configured?
│ │ │ YES → Rewrite model → Use local handler (free)
│ │ │ NO → Forward to ampcode.com (uses Amp credits)
│ ↓
│ Response
@@ -115,6 +120,49 @@ amp-upstream-url: "https://ampcode.com"
amp-restrict-management-to-localhost: true
```
### Model Mapping Configuration
When Amp CLI requests a model that you don't have access to, you can configure mappings to route those requests to alternative models that you DO have available. This avoids consuming Amp credits for models you could handle locally.
```yaml
# Route unavailable models to alternatives
amp-model-mappings:
# Example: Route Claude Opus 4.5 requests to Claude Sonnet 4
- from: "claude-opus-4.5"
to: "claude-sonnet-4"
# Example: Route GPT-5 requests to Gemini 2.5 Pro
- from: "gpt-5"
to: "gemini-2.5-pro"
# Example: Map older model names to newer versions
- from: "claude-3-opus-20240229"
to: "claude-3-5-sonnet-20241022"
```
**How it works:**
1. Amp CLI requests a model (e.g., `claude-opus-4.5`)
2. CLIProxyAPI checks if a local provider is available for that model
3. If not available, it checks the model mappings
4. If a mapping exists, the request is rewritten to use the target model
5. The request is then handled locally (free, using your OAuth subscription)
**Benefits:**
- **Save Amp credits**: Use your local subscriptions instead of forwarding to ampcode.com
- **Hot-reload**: Mappings can be updated without restarting the proxy
- **Structured logging**: Clear logs show when mappings are applied
**Routing Decision Logs:**
The proxy logs each routing decision with structured fields:
```
[AMP] Using local provider for model: gemini-2.5-pro # Local provider (free)
[AMP] Model mapped: claude-opus-4.5 -> claude-sonnet-4 # Mapping applied (free)
[AMP] Forwarding to ampcode.com (uses Amp credits) - model_id: gpt-5 # Fallback (costs credits)
```
### Secret Resolution Precedence
The Amp module resolves API keys using this precedence order:
@@ -301,11 +349,14 @@ When Amp requests a model:
1. **Check local configuration**: Does CLIProxyAPI have OAuth tokens for this model's provider?
2. **If YES**: Route to local handler (use your OAuth subscription)
3. **If NO**: Forward to ampcode.com (use Amp's default routing)
3. **If NO**: Check if a model mapping exists
4. **If mapping exists**: Rewrite request to mapped model → Route to local handler (free)
5. **If no mapping**: Forward to ampcode.com (uses Amp credits)
This enables seamless mixed usage:
- Models you've configured (Gemini, ChatGPT, Claude) → Your OAuth subscriptions
- Models you haven't configured → Amp's default providers
- Models with mappings configured → Routed to alternative local models (free)
- Models you haven't configured and have no mapping → Amp's default providers (uses credits)
### Example API Calls

3
go.mod
View File

@@ -13,14 +13,15 @@ require (
github.com/joho/godotenv v1.5.1
github.com/klauspost/compress v1.17.4
github.com/minio/minio-go/v7 v7.0.66
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
github.com/sirupsen/logrus v1.9.3
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966
github.com/tidwall/gjson v1.18.0
github.com/tidwall/sjson v1.2.5
github.com/tiktoken-go/tokenizer v0.7.0
golang.org/x/crypto v0.43.0
golang.org/x/net v0.46.0
golang.org/x/oauth2 v0.30.0
golang.org/x/term v0.36.0
gopkg.in/natefinch/lumberjack.v2 v2.2.1
gopkg.in/yaml.v3 v3.0.1
)

5
go.sum
View File

@@ -116,6 +116,8 @@ github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6
github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
github.com/pjbgf/sha1cd v0.5.0 h1:a+UkboSi1znleCDUNT3M5YxjOnN1fz2FhN48FlwCxs0=
github.com/pjbgf/sha1cd v0.5.0/go.mod h1:lhpGlyHLpQZoxMv8HcgXvZEhcGs0PG/vsZnEJ7H0iCM=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.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=
github.com/sergi/go-diff v1.4.0/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4=
github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 h1:JIAuq3EEf9cgbU6AtGPK4CTG3Zf6CKMNqf0MHTggAUA=
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966/go.mod h1:sUM3LWHvSMaG192sy56D9F7CNvL7jUJVXoqM1QKLnog=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
github.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
golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug=
golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=
golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=

View File

@@ -36,9 +36,32 @@ import (
)
var (
oauthStatus = make(map[string]string)
oauthStatus = make(map[string]string)
oauthStatusMutex sync.RWMutex
)
// getOAuthStatus safely retrieves an OAuth status
func getOAuthStatus(key string) (string, bool) {
oauthStatusMutex.RLock()
defer oauthStatusMutex.RUnlock()
status, ok := oauthStatus[key]
return status, ok
}
// setOAuthStatus safely sets an OAuth status
func setOAuthStatus(key string, status string) {
oauthStatusMutex.Lock()
defer oauthStatusMutex.Unlock()
oauthStatus[key] = status
}
// deleteOAuthStatus safely deletes an OAuth status
func deleteOAuthStatus(key string) {
oauthStatusMutex.Lock()
defer oauthStatusMutex.Unlock()
delete(oauthStatus, key)
}
var lastRefreshKeys = []string{"last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"}
const (
@@ -713,14 +736,16 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
// Generate PKCE codes
pkceCodes, err := claude.GeneratePKCECodes()
if err != nil {
log.Fatalf("Failed to generate PKCE codes: %v", err)
log.Errorf("Failed to generate PKCE codes: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate PKCE codes"})
return
}
// Generate random state parameter
state, err := misc.GenerateRandomState()
if err != nil {
log.Fatalf("Failed to generate state parameter: %v", err)
log.Errorf("Failed to generate state parameter: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"})
return
}
@@ -730,7 +755,8 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
// Generate authorization URL (then override redirect_uri to reuse server port)
authURL, state, err := anthropicAuth.GenerateAuthURL(state, pkceCodes)
if err != nil {
log.Fatalf("Failed to generate authorization URL: %v", err)
log.Errorf("Failed to generate authorization URL: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"})
return
}
@@ -760,7 +786,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
deadline := time.Now().Add(timeout)
for {
if time.Now().After(deadline) {
oauthStatus[state] = "Timeout waiting for OAuth callback"
setOAuthStatus(state, "Timeout waiting for OAuth callback")
return nil, fmt.Errorf("timeout waiting for OAuth callback")
}
data, errRead := os.ReadFile(path)
@@ -785,13 +811,13 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
if errStr := resultMap["error"]; errStr != "" {
oauthErr := claude.NewOAuthError(errStr, "", http.StatusBadRequest)
log.Error(claude.GetUserFriendlyMessage(oauthErr))
oauthStatus[state] = "Bad request"
setOAuthStatus(state, "Bad request")
return
}
if resultMap["state"] != state {
authErr := claude.NewAuthenticationError(claude.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, resultMap["state"]))
log.Error(claude.GetUserFriendlyMessage(authErr))
oauthStatus[state] = "State code error"
setOAuthStatus(state, "State code error")
return
}
@@ -824,7 +850,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
if errDo != nil {
authErr := claude.NewAuthenticationError(claude.ErrCodeExchangeFailed, errDo)
log.Errorf("Failed to exchange authorization code for tokens: %v", authErr)
oauthStatus[state] = "Failed to exchange authorization code for tokens"
setOAuthStatus(state, "Failed to exchange authorization code for tokens")
return
}
defer func() {
@@ -835,7 +861,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
respBody, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK {
log.Errorf("token exchange failed with status %d: %s", resp.StatusCode, string(respBody))
oauthStatus[state] = fmt.Sprintf("token exchange failed with status %d", resp.StatusCode)
setOAuthStatus(state, fmt.Sprintf("token exchange failed with status %d", resp.StatusCode))
return
}
var tResp struct {
@@ -848,7 +874,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
}
if errU := json.Unmarshal(respBody, &tResp); errU != nil {
log.Errorf("failed to parse token response: %v", errU)
oauthStatus[state] = "Failed to parse token response"
setOAuthStatus(state, "Failed to parse token response")
return
}
bundle := &claude.ClaudeAuthBundle{
@@ -872,8 +898,8 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
}
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
log.Fatalf("Failed to save authentication tokens: %v", errSave)
oauthStatus[state] = "Failed to save authentication tokens"
log.Errorf("Failed to save authentication tokens: %v", errSave)
setOAuthStatus(state, "Failed to save authentication tokens")
return
}
@@ -882,15 +908,17 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
fmt.Println("API key obtained and saved")
}
fmt.Println("You can now use Claude services through this CLI")
delete(oauthStatus, state)
deleteOAuthStatus(state)
}()
oauthStatus[state] = ""
setOAuthStatus(state, "")
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
}
func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
ctx := context.Background()
proxyHTTPClient := util.SetProxy(&h.cfg.SDKConfig, &http.Client{})
ctx = context.WithValue(ctx, oauth2.HTTPClient, proxyHTTPClient)
// Optional project ID from query
projectID := c.Query("project_id")
@@ -942,7 +970,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
for {
if time.Now().After(deadline) {
log.Error("oauth flow timed out")
oauthStatus[state] = "OAuth flow timed out"
setOAuthStatus(state, "OAuth flow timed out")
return
}
if data, errR := os.ReadFile(waitFile); errR == nil {
@@ -951,13 +979,13 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
_ = os.Remove(waitFile)
if errStr := m["error"]; errStr != "" {
log.Errorf("Authentication failed: %s", errStr)
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
return
}
authCode = m["code"]
if authCode == "" {
log.Errorf("Authentication failed: code not found")
oauthStatus[state] = "Authentication failed: code not found"
setOAuthStatus(state, "Authentication failed: code not found")
return
}
break
@@ -969,27 +997,27 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
token, err := conf.Exchange(ctx, authCode)
if err != nil {
log.Errorf("Failed to exchange token: %v", err)
oauthStatus[state] = "Failed to exchange token"
setOAuthStatus(state, "Failed to exchange token")
return
}
requestedProjectID := strings.TrimSpace(projectID)
// Create token storage (mirrors internal/auth/gemini createTokenStorage)
httpClient := conf.Client(ctx, token)
authHTTPClient := conf.Client(ctx, token)
req, errNewRequest := http.NewRequestWithContext(ctx, "GET", "https://www.googleapis.com/oauth2/v1/userinfo?alt=json", nil)
if errNewRequest != nil {
log.Errorf("Could not get user info: %v", errNewRequest)
oauthStatus[state] = "Could not get user info"
setOAuthStatus(state, "Could not get user info")
return
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken))
resp, errDo := httpClient.Do(req)
resp, errDo := authHTTPClient.Do(req)
if errDo != nil {
log.Errorf("Failed to execute request: %v", errDo)
oauthStatus[state] = "Failed to execute request"
setOAuthStatus(state, "Failed to execute request")
return
}
defer func() {
@@ -1001,7 +1029,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
bodyBytes, _ := io.ReadAll(resp.Body)
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
log.Errorf("Get user info request failed with status %d: %s", resp.StatusCode, string(bodyBytes))
oauthStatus[state] = fmt.Sprintf("Get user info request failed with status %d", resp.StatusCode)
setOAuthStatus(state, fmt.Sprintf("Get user info request failed with status %d", resp.StatusCode))
return
}
@@ -1010,7 +1038,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
fmt.Printf("Authenticated user email: %s\n", email)
} else {
fmt.Println("Failed to get user email from token")
oauthStatus[state] = "Failed to get user email from token"
setOAuthStatus(state, "Failed to get user email from token")
}
// Marshal/unmarshal oauth2.Token to generic map and enrich fields
@@ -1018,7 +1046,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
jsonData, _ := json.Marshal(token)
if errUnmarshal := json.Unmarshal(jsonData, &ifToken); errUnmarshal != nil {
log.Errorf("Failed to unmarshal token: %v", errUnmarshal)
oauthStatus[state] = "Failed to unmarshal token"
setOAuthStatus(state, "Failed to unmarshal token")
return
}
@@ -1043,8 +1071,8 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
gemAuth := geminiAuth.NewGeminiAuth()
gemClient, errGetClient := gemAuth.GetAuthenticatedClient(ctx, &ts, h.cfg, true)
if errGetClient != nil {
log.Fatalf("failed to get authenticated client: %v", errGetClient)
oauthStatus[state] = "Failed to get authenticated client"
log.Errorf("failed to get authenticated client: %v", errGetClient)
setOAuthStatus(state, "Failed to get authenticated client")
return
}
fmt.Println("Authentication successful.")
@@ -1054,12 +1082,12 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
projects, errAll := onboardAllGeminiProjects(ctx, gemClient, &ts)
if errAll != nil {
log.Errorf("Failed to complete Gemini CLI onboarding: %v", errAll)
oauthStatus[state] = "Failed to complete Gemini CLI onboarding"
setOAuthStatus(state, "Failed to complete Gemini CLI onboarding")
return
}
if errVerify := ensureGeminiProjectsEnabled(ctx, gemClient, projects); errVerify != nil {
log.Errorf("Failed to verify Cloud AI API status: %v", errVerify)
oauthStatus[state] = "Failed to verify Cloud AI API status"
setOAuthStatus(state, "Failed to verify Cloud AI API status")
return
}
ts.ProjectID = strings.Join(projects, ",")
@@ -1067,26 +1095,26 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
} else {
if errEnsure := ensureGeminiProjectAndOnboard(ctx, gemClient, &ts, requestedProjectID); errEnsure != nil {
log.Errorf("Failed to complete Gemini CLI onboarding: %v", errEnsure)
oauthStatus[state] = "Failed to complete Gemini CLI onboarding"
setOAuthStatus(state, "Failed to complete Gemini CLI onboarding")
return
}
if strings.TrimSpace(ts.ProjectID) == "" {
log.Error("Onboarding did not return a project ID")
oauthStatus[state] = "Failed to resolve project ID"
setOAuthStatus(state, "Failed to resolve project ID")
return
}
isChecked, errCheck := checkCloudAPIIsEnabled(ctx, gemClient, ts.ProjectID)
if errCheck != nil {
log.Errorf("Failed to verify Cloud AI API status: %v", errCheck)
oauthStatus[state] = "Failed to verify Cloud AI API status"
setOAuthStatus(state, "Failed to verify Cloud AI API status")
return
}
ts.Checked = isChecked
if !isChecked {
log.Error("Cloud AI API is not enabled for the selected project")
oauthStatus[state] = "Cloud AI API not enabled"
setOAuthStatus(state, "Cloud AI API not enabled")
return
}
}
@@ -1108,16 +1136,16 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
}
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
log.Fatalf("Failed to save token to file: %v", errSave)
oauthStatus[state] = "Failed to save token to file"
log.Errorf("Failed to save token to file: %v", errSave)
setOAuthStatus(state, "Failed to save token to file")
return
}
delete(oauthStatus, state)
deleteOAuthStatus(state)
fmt.Printf("You can now use Gemini CLI services through this CLI; token saved to %s\n", savedPath)
}()
oauthStatus[state] = ""
setOAuthStatus(state, "")
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
}
@@ -1129,14 +1157,16 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
// Generate PKCE codes
pkceCodes, err := codex.GeneratePKCECodes()
if err != nil {
log.Fatalf("Failed to generate PKCE codes: %v", err)
log.Errorf("Failed to generate PKCE codes: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate PKCE codes"})
return
}
// Generate random state parameter
state, err := misc.GenerateRandomState()
if err != nil {
log.Fatalf("Failed to generate state parameter: %v", err)
log.Errorf("Failed to generate state parameter: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"})
return
}
@@ -1146,7 +1176,8 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
// Generate authorization URL
authURL, err := openaiAuth.GenerateAuthURL(state, pkceCodes)
if err != nil {
log.Fatalf("Failed to generate authorization URL: %v", err)
log.Errorf("Failed to generate authorization URL: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"})
return
}
@@ -1178,7 +1209,7 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
if time.Now().After(deadline) {
authErr := codex.NewAuthenticationError(codex.ErrCallbackTimeout, fmt.Errorf("timeout waiting for OAuth callback"))
log.Error(codex.GetUserFriendlyMessage(authErr))
oauthStatus[state] = "Timeout waiting for OAuth callback"
setOAuthStatus(state, "Timeout waiting for OAuth callback")
return
}
if data, errR := os.ReadFile(waitFile); errR == nil {
@@ -1188,12 +1219,12 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
if errStr := m["error"]; errStr != "" {
oauthErr := codex.NewOAuthError(errStr, "", http.StatusBadRequest)
log.Error(codex.GetUserFriendlyMessage(oauthErr))
oauthStatus[state] = "Bad Request"
setOAuthStatus(state, "Bad Request")
return
}
if m["state"] != state {
authErr := codex.NewAuthenticationError(codex.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, m["state"]))
oauthStatus[state] = "State code error"
setOAuthStatus(state, "State code error")
log.Error(codex.GetUserFriendlyMessage(authErr))
return
}
@@ -1224,14 +1255,14 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
resp, errDo := httpClient.Do(req)
if errDo != nil {
authErr := codex.NewAuthenticationError(codex.ErrCodeExchangeFailed, errDo)
oauthStatus[state] = "Failed to exchange authorization code for tokens"
setOAuthStatus(state, "Failed to exchange authorization code for tokens")
log.Errorf("Failed to exchange authorization code for tokens: %v", authErr)
return
}
defer func() { _ = resp.Body.Close() }()
respBody, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK {
oauthStatus[state] = fmt.Sprintf("Token exchange failed with status %d", resp.StatusCode)
setOAuthStatus(state, fmt.Sprintf("Token exchange failed with status %d", resp.StatusCode))
log.Errorf("token exchange failed with status %d: %s", resp.StatusCode, string(respBody))
return
}
@@ -1242,7 +1273,7 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
ExpiresIn int `json:"expires_in"`
}
if errU := json.Unmarshal(respBody, &tokenResp); errU != nil {
oauthStatus[state] = "Failed to parse token response"
setOAuthStatus(state, "Failed to parse token response")
log.Errorf("failed to parse token response: %v", errU)
return
}
@@ -1280,8 +1311,8 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
}
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
oauthStatus[state] = "Failed to save authentication tokens"
log.Fatalf("Failed to save authentication tokens: %v", errSave)
log.Errorf("Failed to save authentication tokens: %v", errSave)
setOAuthStatus(state, "Failed to save authentication tokens")
return
}
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
@@ -1289,10 +1320,10 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
fmt.Println("API key obtained and saved")
}
fmt.Println("You can now use Codex services through this CLI")
delete(oauthStatus, state)
deleteOAuthStatus(state)
}()
oauthStatus[state] = ""
setOAuthStatus(state, "")
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
}
@@ -1316,7 +1347,8 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
state, errState := misc.GenerateRandomState()
if errState != nil {
log.Fatalf("Failed to generate state parameter: %v", errState)
log.Errorf("Failed to generate state parameter: %v", errState)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"})
return
}
@@ -1358,7 +1390,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
for {
if time.Now().After(deadline) {
log.Error("oauth flow timed out")
oauthStatus[state] = "OAuth flow timed out"
setOAuthStatus(state, "OAuth flow timed out")
return
}
if data, errReadFile := os.ReadFile(waitFile); errReadFile == nil {
@@ -1367,18 +1399,18 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
_ = os.Remove(waitFile)
if errStr := strings.TrimSpace(payload["error"]); errStr != "" {
log.Errorf("Authentication failed: %s", errStr)
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
return
}
if payloadState := strings.TrimSpace(payload["state"]); payloadState != "" && payloadState != state {
log.Errorf("Authentication failed: state mismatch")
oauthStatus[state] = "Authentication failed: state mismatch"
setOAuthStatus(state, "Authentication failed: state mismatch")
return
}
authCode = strings.TrimSpace(payload["code"])
if authCode == "" {
log.Error("Authentication failed: code not found")
oauthStatus[state] = "Authentication failed: code not found"
setOAuthStatus(state, "Authentication failed: code not found")
return
}
break
@@ -1397,7 +1429,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
req, errNewRequest := http.NewRequestWithContext(ctx, http.MethodPost, "https://oauth2.googleapis.com/token", strings.NewReader(form.Encode()))
if errNewRequest != nil {
log.Errorf("Failed to build token request: %v", errNewRequest)
oauthStatus[state] = "Failed to build token request"
setOAuthStatus(state, "Failed to build token request")
return
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
@@ -1405,7 +1437,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
resp, errDo := httpClient.Do(req)
if errDo != nil {
log.Errorf("Failed to execute token request: %v", errDo)
oauthStatus[state] = "Failed to exchange token"
setOAuthStatus(state, "Failed to exchange token")
return
}
defer func() {
@@ -1417,7 +1449,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
bodyBytes, _ := io.ReadAll(resp.Body)
log.Errorf("Antigravity token exchange failed with status %d: %s", resp.StatusCode, string(bodyBytes))
oauthStatus[state] = fmt.Sprintf("Token exchange failed: %d", resp.StatusCode)
setOAuthStatus(state, fmt.Sprintf("Token exchange failed: %d", resp.StatusCode))
return
}
@@ -1429,7 +1461,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
}
if errDecode := json.NewDecoder(resp.Body).Decode(&tokenResp); errDecode != nil {
log.Errorf("Failed to parse token response: %v", errDecode)
oauthStatus[state] = "Failed to parse token response"
setOAuthStatus(state, "Failed to parse token response")
return
}
@@ -1438,7 +1470,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
infoReq, errInfoReq := http.NewRequestWithContext(ctx, http.MethodGet, "https://www.googleapis.com/oauth2/v1/userinfo?alt=json", nil)
if errInfoReq != nil {
log.Errorf("Failed to build user info request: %v", errInfoReq)
oauthStatus[state] = "Failed to build user info request"
setOAuthStatus(state, "Failed to build user info request")
return
}
infoReq.Header.Set("Authorization", "Bearer "+tokenResp.AccessToken)
@@ -1446,7 +1478,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
infoResp, errInfo := httpClient.Do(infoReq)
if errInfo != nil {
log.Errorf("Failed to execute user info request: %v", errInfo)
oauthStatus[state] = "Failed to execute user info request"
setOAuthStatus(state, "Failed to execute user info request")
return
}
defer func() {
@@ -1465,11 +1497,22 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
} else {
bodyBytes, _ := io.ReadAll(infoResp.Body)
log.Errorf("User info request failed with status %d: %s", infoResp.StatusCode, string(bodyBytes))
oauthStatus[state] = fmt.Sprintf("User info request failed: %d", infoResp.StatusCode)
setOAuthStatus(state, fmt.Sprintf("User info request failed: %d", infoResp.StatusCode))
return
}
}
projectID := ""
if strings.TrimSpace(tokenResp.AccessToken) != "" {
fetchedProjectID, errProject := sdkAuth.FetchAntigravityProjectID(ctx, tokenResp.AccessToken, httpClient)
if errProject != nil {
log.Warnf("antigravity: failed to fetch project ID: %v", errProject)
} else {
projectID = fetchedProjectID
log.Infof("antigravity: obtained project ID %s", projectID)
}
}
now := time.Now()
metadata := map[string]any{
"type": "antigravity",
@@ -1482,6 +1525,9 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
if email != "" {
metadata["email"] = email
}
if projectID != "" {
metadata["project_id"] = projectID
}
fileName := sanitizeAntigravityFileName(email)
label := strings.TrimSpace(email)
@@ -1498,17 +1544,20 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
}
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
log.Fatalf("Failed to save token to file: %v", errSave)
oauthStatus[state] = "Failed to save token to file"
log.Errorf("Failed to save token to file: %v", errSave)
setOAuthStatus(state, "Failed to save token to file")
return
}
delete(oauthStatus, state)
deleteOAuthStatus(state)
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
if projectID != "" {
fmt.Printf("Using GCP project: %s\n", projectID)
}
fmt.Println("You can now use Antigravity services through this CLI")
}()
oauthStatus[state] = ""
setOAuthStatus(state, "")
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
}
@@ -1524,7 +1573,8 @@ func (h *Handler) RequestQwenToken(c *gin.Context) {
// Generate authorization URL
deviceFlow, err := qwenAuth.InitiateDeviceFlow(ctx)
if err != nil {
log.Fatalf("Failed to generate authorization URL: %v", err)
log.Errorf("Failed to generate authorization URL: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"})
return
}
authURL := deviceFlow.VerificationURIComplete
@@ -1533,7 +1583,7 @@ func (h *Handler) RequestQwenToken(c *gin.Context) {
fmt.Println("Waiting for authentication...")
tokenData, errPollForToken := qwenAuth.PollForToken(deviceFlow.DeviceCode, deviceFlow.CodeVerifier)
if errPollForToken != nil {
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
fmt.Printf("Authentication failed: %v\n", errPollForToken)
return
}
@@ -1551,17 +1601,17 @@ func (h *Handler) RequestQwenToken(c *gin.Context) {
}
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
log.Fatalf("Failed to save authentication tokens: %v", errSave)
oauthStatus[state] = "Failed to save authentication tokens"
log.Errorf("Failed to save authentication tokens: %v", errSave)
setOAuthStatus(state, "Failed to save authentication tokens")
return
}
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
fmt.Println("You can now use Qwen services through this CLI")
delete(oauthStatus, state)
deleteOAuthStatus(state)
}()
oauthStatus[state] = ""
setOAuthStatus(state, "")
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
}
@@ -1600,7 +1650,7 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
var resultMap map[string]string
for {
if time.Now().After(deadline) {
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
fmt.Println("Authentication failed: timeout waiting for callback")
return
}
@@ -1613,26 +1663,26 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
}
if errStr := strings.TrimSpace(resultMap["error"]); errStr != "" {
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
fmt.Printf("Authentication failed: %s\n", errStr)
return
}
if resultState := strings.TrimSpace(resultMap["state"]); resultState != state {
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
fmt.Println("Authentication failed: state mismatch")
return
}
code := strings.TrimSpace(resultMap["code"])
if code == "" {
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
fmt.Println("Authentication failed: code missing")
return
}
tokenData, errExchange := authSvc.ExchangeCodeForTokens(ctx, code, redirectURI)
if errExchange != nil {
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
fmt.Printf("Authentication failed: %v\n", errExchange)
return
}
@@ -1654,8 +1704,8 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
oauthStatus[state] = "Failed to save authentication tokens"
log.Fatalf("Failed to save authentication tokens: %v", errSave)
log.Errorf("Failed to save authentication tokens: %v", errSave)
setOAuthStatus(state, "Failed to save authentication tokens")
return
}
@@ -1664,10 +1714,10 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
fmt.Println("API key obtained and saved")
}
fmt.Println("You can now use iFlow services through this CLI")
delete(oauthStatus, state)
deleteOAuthStatus(state)
}()
oauthStatus[state] = ""
setOAuthStatus(state, "")
c.JSON(http.StatusOK, gin.H{"status": "ok", "url": authURL, "state": state})
}
@@ -2084,6 +2134,7 @@ func checkCloudAPIIsEnabled(ctx context.Context, httpClient *http.Client, projec
continue
}
}
_ = resp.Body.Close()
return false, fmt.Errorf("project activation required: %s", errMessage)
}
return true, nil
@@ -2091,7 +2142,7 @@ func checkCloudAPIIsEnabled(ctx context.Context, httpClient *http.Client, projec
func (h *Handler) GetAuthStatus(c *gin.Context) {
state := c.Query("state")
if err, ok := oauthStatus[state]; ok {
if err, ok := getOAuthStatus(state); ok {
if err != "" {
c.JSON(200, gin.H{"status": "error", "error": err})
} else {
@@ -2101,5 +2152,5 @@ func (h *Handler) GetAuthStatus(c *gin.Context) {
} else {
c.JSON(200, gin.H{"status": "ok"})
}
delete(oauthStatus, state)
deleteOAuthStatus(state)
}

View File

@@ -1,43 +1,95 @@
package management
import (
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"strings"
"time"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
log "github.com/sirupsen/logrus"
"gopkg.in/yaml.v3"
)
const (
latestReleaseURL = "https://api.github.com/repos/router-for-me/CLIProxyAPIPlus/releases/latest"
latestReleaseUserAgent = "CLIProxyAPIPlus"
)
func (h *Handler) GetConfig(c *gin.Context) {
if h == nil || h.cfg == nil {
c.JSON(200, gin.H{})
return
}
cfgCopy := *h.cfg
cfgCopy.GlAPIKey = geminiKeyStringsFromConfig(h.cfg)
c.JSON(200, &cfgCopy)
}
func (h *Handler) GetConfigYAML(c *gin.Context) {
data, err := os.ReadFile(h.configFilePath)
type releaseInfo struct {
TagName string `json:"tag_name"`
Name string `json:"name"`
}
// GetLatestVersion returns the latest release version from GitHub without downloading assets.
func (h *Handler) GetLatestVersion(c *gin.Context) {
client := &http.Client{Timeout: 10 * time.Second}
proxyURL := ""
if h != nil && h.cfg != nil {
proxyURL = strings.TrimSpace(h.cfg.ProxyURL)
}
if proxyURL != "" {
sdkCfg := &sdkconfig.SDKConfig{ProxyURL: proxyURL}
util.SetProxy(sdkCfg, client)
}
req, err := http.NewRequestWithContext(c.Request.Context(), http.MethodGet, latestReleaseURL, nil)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "read_failed", "message": err.Error()})
c.JSON(http.StatusInternalServerError, gin.H{"error": "request_create_failed", "message": err.Error()})
return
}
var node yaml.Node
if err = yaml.Unmarshal(data, &node); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "parse_failed", "message": err.Error()})
req.Header.Set("Accept", "application/vnd.github+json")
req.Header.Set("User-Agent", latestReleaseUserAgent)
resp, err := client.Do(req)
if err != nil {
c.JSON(http.StatusBadGateway, gin.H{"error": "request_failed", "message": err.Error()})
return
}
c.Header("Content-Type", "application/yaml; charset=utf-8")
c.Header("Vary", "format, Accept")
enc := yaml.NewEncoder(c.Writer)
enc.SetIndent(2)
_ = enc.Encode(&node)
_ = enc.Close()
defer func() {
if errClose := resp.Body.Close(); errClose != nil {
log.WithError(errClose).Debug("failed to close latest version response body")
}
}()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
c.JSON(http.StatusBadGateway, gin.H{"error": "unexpected_status", "message": fmt.Sprintf("status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))})
return
}
var info releaseInfo
if errDecode := json.NewDecoder(resp.Body).Decode(&info); errDecode != nil {
c.JSON(http.StatusBadGateway, gin.H{"error": "decode_failed", "message": errDecode.Error()})
return
}
version := strings.TrimSpace(info.TagName)
if version == "" {
version = strings.TrimSpace(info.Name)
}
if version == "" {
c.JSON(http.StatusBadGateway, gin.H{"error": "invalid_response", "message": "missing release version"})
return
}
c.JSON(http.StatusOK, gin.H{"latest-version": version})
}
func WriteConfig(path string, data []byte) error {
@@ -111,9 +163,9 @@ func (h *Handler) PutConfigYAML(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true, "changed": []string{"config"}})
}
// GetConfigFile returns the raw config.yaml file bytes without re-encoding.
// GetConfigYAML returns the raw config.yaml file bytes without re-encoding.
// It preserves comments and original formatting/styles.
func (h *Handler) GetConfigFile(c *gin.Context) {
func (h *Handler) GetConfigYAML(c *gin.Context) {
data, err := os.ReadFile(h.configFilePath)
if err != nil {
if os.IsNotExist(err) {

View File

@@ -104,53 +104,6 @@ func (h *Handler) deleteFromStringList(c *gin.Context, target *[]string, after f
c.JSON(400, gin.H{"error": "missing index or value"})
}
func sanitizeStringSlice(in []string) []string {
out := make([]string, 0, len(in))
for i := range in {
if trimmed := strings.TrimSpace(in[i]); trimmed != "" {
out = append(out, trimmed)
}
}
return out
}
func geminiKeyStringsFromConfig(cfg *config.Config) []string {
if cfg == nil || len(cfg.GeminiKey) == 0 {
return nil
}
out := make([]string, 0, len(cfg.GeminiKey))
for i := range cfg.GeminiKey {
if key := strings.TrimSpace(cfg.GeminiKey[i].APIKey); key != "" {
out = append(out, key)
}
}
return out
}
func (h *Handler) applyLegacyKeys(keys []string) {
if h == nil || h.cfg == nil {
return
}
sanitized := sanitizeStringSlice(keys)
existing := make(map[string]config.GeminiKey, len(h.cfg.GeminiKey))
for _, entry := range h.cfg.GeminiKey {
if key := strings.TrimSpace(entry.APIKey); key != "" {
existing[key] = entry
}
}
newList := make([]config.GeminiKey, 0, len(sanitized))
for _, key := range sanitized {
if entry, ok := existing[key]; ok {
newList = append(newList, entry)
} else {
newList = append(newList, config.GeminiKey{APIKey: key})
}
}
h.cfg.GeminiKey = newList
h.cfg.GlAPIKey = sanitized
h.cfg.SanitizeGeminiKeys()
}
// api-keys
func (h *Handler) GetAPIKeys(c *gin.Context) { c.JSON(200, gin.H{"api-keys": h.cfg.APIKeys}) }
func (h *Handler) PutAPIKeys(c *gin.Context) {
@@ -166,24 +119,6 @@ func (h *Handler) DeleteAPIKeys(c *gin.Context) {
h.deleteFromStringList(c, &h.cfg.APIKeys, func() { h.cfg.Access.Providers = nil })
}
// generative-language-api-key
func (h *Handler) GetGlKeys(c *gin.Context) {
c.JSON(200, gin.H{"generative-language-api-key": geminiKeyStringsFromConfig(h.cfg)})
}
func (h *Handler) PutGlKeys(c *gin.Context) {
h.putStringList(c, func(v []string) {
h.applyLegacyKeys(v)
}, nil)
}
func (h *Handler) PatchGlKeys(c *gin.Context) {
target := append([]string(nil), geminiKeyStringsFromConfig(h.cfg)...)
h.patchStringList(c, &target, func() { h.applyLegacyKeys(target) })
}
func (h *Handler) DeleteGlKeys(c *gin.Context) {
target := append([]string(nil), geminiKeyStringsFromConfig(h.cfg)...)
h.deleteFromStringList(c, &target, func() { h.applyLegacyKeys(target) })
}
// gemini-api-key: []GeminiKey
func (h *Handler) GetGeminiKeys(c *gin.Context) {
c.JSON(200, gin.H{"gemini-api-key": h.cfg.GeminiKey})
@@ -409,15 +344,14 @@ func (h *Handler) PutOpenAICompat(c *gin.Context) {
}
arr = obj.Items
}
arr = migrateLegacyOpenAICompatibilityKeys(arr)
// Filter out providers with empty base-url -> remove provider entirely
filtered := make([]config.OpenAICompatibility, 0, len(arr))
for i := range arr {
normalizeOpenAICompatibilityEntry(&arr[i])
if strings.TrimSpace(arr[i].BaseURL) != "" {
filtered = append(filtered, arr[i])
}
}
h.cfg.OpenAICompatibility = migrateLegacyOpenAICompatibilityKeys(filtered)
h.cfg.OpenAICompatibility = filtered
h.cfg.SanitizeOpenAICompatibility()
h.persist(c)
}
@@ -431,7 +365,6 @@ func (h *Handler) PatchOpenAICompat(c *gin.Context) {
c.JSON(400, gin.H{"error": "invalid body"})
return
}
h.cfg.OpenAICompatibility = migrateLegacyOpenAICompatibilityKeys(h.cfg.OpenAICompatibility)
normalizeOpenAICompatibilityEntry(body.Value)
// If base-url becomes empty, delete the provider instead of updating
if strings.TrimSpace(body.Value.BaseURL) == "" {
@@ -731,28 +664,6 @@ func normalizeOpenAICompatibilityEntry(entry *config.OpenAICompatibility) {
existing[trimmed] = struct{}{}
}
}
if len(entry.APIKeys) == 0 {
return
}
for _, legacyKey := range entry.APIKeys {
trimmed := strings.TrimSpace(legacyKey)
if trimmed == "" {
continue
}
if _, ok := existing[trimmed]; ok {
continue
}
entry.APIKeyEntries = append(entry.APIKeyEntries, config.OpenAICompatibilityAPIKey{APIKey: trimmed})
existing[trimmed] = struct{}{}
}
entry.APIKeys = nil
}
func migrateLegacyOpenAICompatibilityKeys(entries []config.OpenAICompatibility) []config.OpenAICompatibility {
for i := range entries {
normalizeOpenAICompatibilityEntry(&entries[i])
}
return entries
}
func normalizedOpenAICompatibilityEntries(entries []config.OpenAICompatibility) []config.OpenAICompatibility {
@@ -765,9 +676,6 @@ func normalizedOpenAICompatibilityEntries(entries []config.OpenAICompatibility)
if len(copyEntry.APIKeyEntries) > 0 {
copyEntry.APIKeyEntries = append([]config.OpenAICompatibilityAPIKey(nil), copyEntry.APIKeyEntries...)
}
if len(copyEntry.APIKeys) > 0 {
copyEntry.APIKeys = append([]string(nil), copyEntry.APIKeys...)
}
normalizeOpenAICompatibilityEntry(&copyEntry)
out[i] = copyEntry
}
@@ -798,3 +706,155 @@ func normalizeClaudeKey(entry *config.ClaudeKey) {
}
entry.Models = normalized
}
// GetAmpCode returns the complete ampcode configuration.
func (h *Handler) GetAmpCode(c *gin.Context) {
if h == nil || h.cfg == nil {
c.JSON(200, gin.H{"ampcode": config.AmpCode{}})
return
}
c.JSON(200, gin.H{"ampcode": h.cfg.AmpCode})
}
// GetAmpUpstreamURL returns the ampcode upstream URL.
func (h *Handler) GetAmpUpstreamURL(c *gin.Context) {
if h == nil || h.cfg == nil {
c.JSON(200, gin.H{"upstream-url": ""})
return
}
c.JSON(200, gin.H{"upstream-url": h.cfg.AmpCode.UpstreamURL})
}
// PutAmpUpstreamURL updates the ampcode upstream URL.
func (h *Handler) PutAmpUpstreamURL(c *gin.Context) {
h.updateStringField(c, func(v string) { h.cfg.AmpCode.UpstreamURL = strings.TrimSpace(v) })
}
// DeleteAmpUpstreamURL clears the ampcode upstream URL.
func (h *Handler) DeleteAmpUpstreamURL(c *gin.Context) {
h.cfg.AmpCode.UpstreamURL = ""
h.persist(c)
}
// GetAmpUpstreamAPIKey returns the ampcode upstream API key.
func (h *Handler) GetAmpUpstreamAPIKey(c *gin.Context) {
if h == nil || h.cfg == nil {
c.JSON(200, gin.H{"upstream-api-key": ""})
return
}
c.JSON(200, gin.H{"upstream-api-key": h.cfg.AmpCode.UpstreamAPIKey})
}
// PutAmpUpstreamAPIKey updates the ampcode upstream API key.
func (h *Handler) PutAmpUpstreamAPIKey(c *gin.Context) {
h.updateStringField(c, func(v string) { h.cfg.AmpCode.UpstreamAPIKey = strings.TrimSpace(v) })
}
// DeleteAmpUpstreamAPIKey clears the ampcode upstream API key.
func (h *Handler) DeleteAmpUpstreamAPIKey(c *gin.Context) {
h.cfg.AmpCode.UpstreamAPIKey = ""
h.persist(c)
}
// GetAmpRestrictManagementToLocalhost returns the localhost restriction setting.
func (h *Handler) GetAmpRestrictManagementToLocalhost(c *gin.Context) {
if h == nil || h.cfg == nil {
c.JSON(200, gin.H{"restrict-management-to-localhost": true})
return
}
c.JSON(200, gin.H{"restrict-management-to-localhost": h.cfg.AmpCode.RestrictManagementToLocalhost})
}
// PutAmpRestrictManagementToLocalhost updates the localhost restriction setting.
func (h *Handler) PutAmpRestrictManagementToLocalhost(c *gin.Context) {
h.updateBoolField(c, func(v bool) { h.cfg.AmpCode.RestrictManagementToLocalhost = v })
}
// GetAmpModelMappings returns the ampcode model mappings.
func (h *Handler) GetAmpModelMappings(c *gin.Context) {
if h == nil || h.cfg == nil {
c.JSON(200, gin.H{"model-mappings": []config.AmpModelMapping{}})
return
}
c.JSON(200, gin.H{"model-mappings": h.cfg.AmpCode.ModelMappings})
}
// PutAmpModelMappings replaces all ampcode model mappings.
func (h *Handler) PutAmpModelMappings(c *gin.Context) {
var body struct {
Value []config.AmpModelMapping `json:"value"`
}
if err := c.ShouldBindJSON(&body); err != nil {
c.JSON(400, gin.H{"error": "invalid body"})
return
}
h.cfg.AmpCode.ModelMappings = body.Value
h.persist(c)
}
// PatchAmpModelMappings adds or updates model mappings.
func (h *Handler) PatchAmpModelMappings(c *gin.Context) {
var body struct {
Value []config.AmpModelMapping `json:"value"`
}
if err := c.ShouldBindJSON(&body); err != nil {
c.JSON(400, gin.H{"error": "invalid body"})
return
}
existing := make(map[string]int)
for i, m := range h.cfg.AmpCode.ModelMappings {
existing[strings.TrimSpace(m.From)] = i
}
for _, newMapping := range body.Value {
from := strings.TrimSpace(newMapping.From)
if idx, ok := existing[from]; ok {
h.cfg.AmpCode.ModelMappings[idx] = newMapping
} else {
h.cfg.AmpCode.ModelMappings = append(h.cfg.AmpCode.ModelMappings, newMapping)
existing[from] = len(h.cfg.AmpCode.ModelMappings) - 1
}
}
h.persist(c)
}
// DeleteAmpModelMappings removes specified model mappings by "from" field.
func (h *Handler) DeleteAmpModelMappings(c *gin.Context) {
var body struct {
Value []string `json:"value"`
}
if err := c.ShouldBindJSON(&body); err != nil || len(body.Value) == 0 {
h.cfg.AmpCode.ModelMappings = nil
h.persist(c)
return
}
toRemove := make(map[string]bool)
for _, from := range body.Value {
toRemove[strings.TrimSpace(from)] = true
}
newMappings := make([]config.AmpModelMapping, 0, len(h.cfg.AmpCode.ModelMappings))
for _, m := range h.cfg.AmpCode.ModelMappings {
if !toRemove[strings.TrimSpace(m.From)] {
newMappings = append(newMappings, m)
}
}
h.cfg.AmpCode.ModelMappings = newMappings
h.persist(c)
}
// GetAmpForceModelMappings returns whether model mappings are forced.
func (h *Handler) GetAmpForceModelMappings(c *gin.Context) {
if h == nil || h.cfg == nil {
c.JSON(200, gin.H{"force-model-mappings": false})
return
}
c.JSON(200, gin.H{"force-model-mappings": h.cfg.AmpCode.ForceModelMappings})
}
// PutAmpForceModelMappings updates the force model mappings setting.
func (h *Handler) PutAmpForceModelMappings(c *gin.Context) {
h.updateBoolField(c, func(v bool) { h.cfg.AmpCode.ForceModelMappings = v })
}

View File

@@ -240,16 +240,6 @@ func (h *Handler) updateBoolField(c *gin.Context, set func(bool)) {
Value *bool `json:"value"`
}
if err := c.ShouldBindJSON(&body); err != nil || body.Value == nil {
var m map[string]any
if err2 := c.ShouldBindJSON(&m); err2 == nil {
for _, v := range m {
if b, ok := v.(bool); ok {
set(b)
h.persist(c)
return
}
}
}
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
return
}

View File

@@ -112,5 +112,10 @@ func shouldLogRequest(path string) bool {
if strings.HasPrefix(path, "/v0/management") || strings.HasPrefix(path, "/management") {
return false
}
if strings.HasPrefix(path, "/api") {
return strings.HasPrefix(path, "/api/provider")
}
return true
}

View File

@@ -232,7 +232,16 @@ func (w *ResponseWriterWrapper) Finalize(c *gin.Context) error {
w.streamDone = nil
}
// Write API Request and Response to the streaming log before closing
if w.streamWriter != nil {
apiRequest := w.extractAPIRequest(c)
if len(apiRequest) > 0 {
_ = w.streamWriter.WriteAPIRequest(apiRequest)
}
apiResponse := w.extractAPIResponse(c)
if len(apiResponse) > 0 {
_ = w.streamWriter.WriteAPIResponse(apiResponse)
}
if err := w.streamWriter.Close(); err != nil {
w.streamWriter = nil
return err

View File

@@ -23,13 +23,24 @@ type Option func(*AmpModule)
// - Reverse proxy to Amp control plane for OAuth/management
// - Provider-specific route aliases (/api/provider/{provider}/...)
// - Automatic gzip decompression for misconfigured upstreams
// - Model mapping for routing unavailable models to alternatives
type AmpModule struct {
secretSource SecretSource
proxy *httputil.ReverseProxy
proxyMu sync.RWMutex // protects proxy for hot-reload
accessManager *sdkaccess.Manager
authMiddleware_ gin.HandlerFunc
modelMapper *DefaultModelMapper
enabled bool
registerOnce sync.Once
// restrictToLocalhost controls localhost-only access for management routes (hot-reloadable)
restrictToLocalhost bool
restrictMu sync.RWMutex
// configMu protects lastConfig for partial reload comparison
configMu sync.RWMutex
lastConfig *config.AmpCode
}
// New creates a new Amp routing module with the given options.
@@ -89,11 +100,22 @@ func (m *AmpModule) Name() string {
return "amp-routing"
}
// forceModelMappings returns whether model mappings should take precedence over local API keys
func (m *AmpModule) forceModelMappings() bool {
m.configMu.RLock()
defer m.configMu.RUnlock()
if m.lastConfig == nil {
return false
}
return m.lastConfig.ForceModelMappings
}
// Register sets up Amp routes if configured.
// This implements the RouteModuleV2 interface with Context.
// Routes are registered only once via sync.Once for idempotent behavior.
func (m *AmpModule) Register(ctx modules.Context) error {
upstreamURL := strings.TrimSpace(ctx.Config.AmpUpstreamURL)
settings := ctx.Config.AmpCode
upstreamURL := strings.TrimSpace(settings.UpstreamURL)
// Determine auth middleware (from module or context)
auth := m.getAuthMiddleware(ctx)
@@ -101,40 +123,36 @@ func (m *AmpModule) Register(ctx modules.Context) error {
// Use registerOnce to ensure routes are only registered once
var regErr error
m.registerOnce.Do(func() {
// Initialize model mapper from config (for routing unavailable models to alternatives)
m.modelMapper = NewModelMapper(settings.ModelMappings)
// Store initial config for partial reload comparison
settingsCopy := settings
m.lastConfig = &settingsCopy
// Initialize localhost restriction setting (hot-reloadable)
m.setRestrictToLocalhost(settings.RestrictManagementToLocalhost)
// Always register provider aliases - these work without an upstream
m.registerProviderAliases(ctx.Engine, ctx.BaseHandler, auth)
// Register management proxy routes once; middleware will gate access when upstream is unavailable.
m.registerManagementRoutes(ctx.Engine, ctx.BaseHandler)
// If no upstream URL, skip proxy routes but provider aliases are still available
if upstreamURL == "" {
log.Debug("Amp upstream proxy disabled (no upstream URL configured)")
log.Debug("Amp provider alias routes registered")
log.Debug("amp upstream proxy disabled (no upstream URL configured)")
log.Debug("amp provider alias routes registered")
m.enabled = false
return
}
// Create secret source with precedence: config > env > file
// Cache secrets for 5 minutes to reduce file I/O
if m.secretSource == nil {
m.secretSource = NewMultiSourceSecret(ctx.Config.AmpUpstreamAPIKey, 0 /* default 5min */)
}
// Create reverse proxy with gzip handling via ModifyResponse
proxy, err := createReverseProxy(upstreamURL, m.secretSource)
if err != nil {
if err := m.enableUpstreamProxy(upstreamURL, &settings); err != nil {
regErr = fmt.Errorf("failed to create amp proxy: %w", err)
return
}
m.proxy = proxy
m.enabled = true
// Register management proxy routes (requires upstream)
// Restrict to localhost by default for security (prevents drive-by browser attacks)
handler := proxyHandler(proxy)
m.registerManagementRoutes(ctx.Engine, ctx.BaseHandler, handler, ctx.Config.AmpRestrictManagementToLocalhost)
log.Infof("Amp upstream proxy enabled for: %s", upstreamURL)
log.Debug("Amp provider alias routes registered")
log.Debug("amp provider alias routes registered")
})
return regErr
@@ -150,34 +168,175 @@ func (m *AmpModule) getAuthMiddleware(ctx modules.Context) gin.HandlerFunc {
return ctx.AuthMiddleware
}
// Fallback: no authentication (should not happen in production)
log.Warn("Amp module: no auth middleware provided, allowing all requests")
log.Warn("amp module: no auth middleware provided, allowing all requests")
return func(c *gin.Context) {
c.Next()
}
}
// OnConfigUpdated handles configuration updates.
// Currently requires restart for URL changes (could be enhanced for dynamic updates).
// OnConfigUpdated handles configuration updates with partial reload support.
// Only updates components that have actually changed to avoid unnecessary work.
// Supports hot-reload for: model-mappings, upstream-api-key, upstream-url, restrict-management-to-localhost.
func (m *AmpModule) OnConfigUpdated(cfg *config.Config) error {
if !m.enabled {
log.Debug("Amp routing not enabled, skipping config update")
return nil
}
newSettings := cfg.AmpCode
upstreamURL := strings.TrimSpace(cfg.AmpUpstreamURL)
if upstreamURL == "" {
log.Warn("Amp upstream URL removed from config, restart required to disable")
return nil
}
// Get previous config for comparison
m.configMu.RLock()
oldSettings := m.lastConfig
m.configMu.RUnlock()
// If API key changed, invalidate the cache
if m.secretSource != nil {
if ms, ok := m.secretSource.(*MultiSourceSecret); ok {
ms.InvalidateCache()
log.Debug("Amp secret cache invalidated due to config update")
if oldSettings != nil && oldSettings.RestrictManagementToLocalhost != newSettings.RestrictManagementToLocalhost {
m.setRestrictToLocalhost(newSettings.RestrictManagementToLocalhost)
if !newSettings.RestrictManagementToLocalhost {
log.Warnf("amp management routes now accessible from any IP - this is insecure!")
}
}
log.Debug("Amp config updated (restart required for URL changes)")
newUpstreamURL := strings.TrimSpace(newSettings.UpstreamURL)
oldUpstreamURL := ""
if oldSettings != nil {
oldUpstreamURL = strings.TrimSpace(oldSettings.UpstreamURL)
}
if !m.enabled && newUpstreamURL != "" {
if err := m.enableUpstreamProxy(newUpstreamURL, &newSettings); err != nil {
log.Errorf("amp config: failed to enable upstream proxy for %s: %v", newUpstreamURL, err)
}
}
// Check model mappings change
modelMappingsChanged := m.hasModelMappingsChanged(oldSettings, &newSettings)
if modelMappingsChanged {
if m.modelMapper != nil {
m.modelMapper.UpdateMappings(newSettings.ModelMappings)
} else if m.enabled {
log.Warnf("amp model mapper not initialized, skipping model mapping update")
}
}
if m.enabled {
// Check upstream URL change - now supports hot-reload
if newUpstreamURL == "" && oldUpstreamURL != "" {
m.setProxy(nil)
m.enabled = false
} else if oldUpstreamURL != "" && newUpstreamURL != oldUpstreamURL && newUpstreamURL != "" {
// Recreate proxy with new URL
proxy, err := createReverseProxy(newUpstreamURL, m.secretSource)
if err != nil {
log.Errorf("amp config: failed to create proxy for new upstream URL %s: %v", newUpstreamURL, err)
} else {
m.setProxy(proxy)
}
}
// Check API key change
apiKeyChanged := m.hasAPIKeyChanged(oldSettings, &newSettings)
if apiKeyChanged {
if m.secretSource != nil {
if ms, ok := m.secretSource.(*MultiSourceSecret); ok {
ms.UpdateExplicitKey(newSettings.UpstreamAPIKey)
ms.InvalidateCache()
}
}
}
}
// Store current config for next comparison
m.configMu.Lock()
settingsCopy := newSettings // copy struct
m.lastConfig = &settingsCopy
m.configMu.Unlock()
return nil
}
func (m *AmpModule) enableUpstreamProxy(upstreamURL string, settings *config.AmpCode) error {
if m.secretSource == nil {
m.secretSource = NewMultiSourceSecret(settings.UpstreamAPIKey, 0 /* default 5min */)
} else if ms, ok := m.secretSource.(*MultiSourceSecret); ok {
ms.UpdateExplicitKey(settings.UpstreamAPIKey)
ms.InvalidateCache()
}
proxy, err := createReverseProxy(upstreamURL, m.secretSource)
if err != nil {
return err
}
m.setProxy(proxy)
m.enabled = true
log.Infof("amp upstream proxy enabled for: %s", upstreamURL)
return nil
}
// hasModelMappingsChanged compares old and new model mappings.
func (m *AmpModule) hasModelMappingsChanged(old *config.AmpCode, new *config.AmpCode) bool {
if old == nil {
return len(new.ModelMappings) > 0
}
if len(old.ModelMappings) != len(new.ModelMappings) {
return true
}
// Build map for efficient comparison
oldMap := make(map[string]string, len(old.ModelMappings))
for _, mapping := range old.ModelMappings {
oldMap[strings.TrimSpace(mapping.From)] = strings.TrimSpace(mapping.To)
}
for _, mapping := range new.ModelMappings {
from := strings.TrimSpace(mapping.From)
to := strings.TrimSpace(mapping.To)
if oldTo, exists := oldMap[from]; !exists || oldTo != to {
return true
}
}
return false
}
// hasAPIKeyChanged compares old and new API keys.
func (m *AmpModule) hasAPIKeyChanged(old *config.AmpCode, new *config.AmpCode) bool {
oldKey := ""
if old != nil {
oldKey = strings.TrimSpace(old.UpstreamAPIKey)
}
newKey := strings.TrimSpace(new.UpstreamAPIKey)
return oldKey != newKey
}
// GetModelMapper returns the model mapper instance (for testing/debugging).
func (m *AmpModule) GetModelMapper() *DefaultModelMapper {
return m.modelMapper
}
// getProxy returns the current proxy instance (thread-safe for hot-reload).
func (m *AmpModule) getProxy() *httputil.ReverseProxy {
m.proxyMu.RLock()
defer m.proxyMu.RUnlock()
return m.proxy
}
// setProxy updates the proxy instance (thread-safe for hot-reload).
func (m *AmpModule) setProxy(proxy *httputil.ReverseProxy) {
m.proxyMu.Lock()
defer m.proxyMu.Unlock()
m.proxy = proxy
}
// IsRestrictedToLocalhost returns whether management routes are restricted to localhost.
func (m *AmpModule) IsRestrictedToLocalhost() bool {
m.restrictMu.RLock()
defer m.restrictMu.RUnlock()
return m.restrictToLocalhost
}
// setRestrictToLocalhost updates the localhost restriction setting.
func (m *AmpModule) setRestrictToLocalhost(restrict bool) {
m.restrictMu.Lock()
defer m.restrictMu.Unlock()
m.restrictToLocalhost = restrict
}

View File

@@ -56,8 +56,10 @@ func TestAmpModule_Register_WithUpstream(t *testing.T) {
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
cfg := &config.Config{
AmpUpstreamURL: upstream.URL,
AmpUpstreamAPIKey: "test-key",
AmpCode: config.AmpCode{
UpstreamURL: upstream.URL,
UpstreamAPIKey: "test-key",
},
}
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
@@ -86,7 +88,9 @@ func TestAmpModule_Register_WithoutUpstream(t *testing.T) {
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
cfg := &config.Config{
AmpUpstreamURL: "", // No upstream
AmpCode: config.AmpCode{
UpstreamURL: "", // No upstream
},
}
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
@@ -121,7 +125,9 @@ func TestAmpModule_Register_InvalidUpstream(t *testing.T) {
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
cfg := &config.Config{
AmpUpstreamURL: "://invalid-url",
AmpCode: config.AmpCode{
UpstreamURL: "://invalid-url",
},
}
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
@@ -151,7 +157,7 @@ func TestAmpModule_OnConfigUpdated_CacheInvalidation(t *testing.T) {
}
// Update config - should invalidate cache
if err := m.OnConfigUpdated(&config.Config{AmpUpstreamURL: "http://x"}); err != nil {
if err := m.OnConfigUpdated(&config.Config{AmpCode: config.AmpCode{UpstreamURL: "http://x"}}); err != nil {
t.Fatal(err)
}
@@ -175,7 +181,7 @@ func TestAmpModule_OnConfigUpdated_URLRemoved(t *testing.T) {
m.secretSource = ms
// Config update with empty URL - should log warning but not error
cfg := &config.Config{AmpUpstreamURL: ""}
cfg := &config.Config{AmpCode: config.AmpCode{UpstreamURL: ""}}
if err := m.OnConfigUpdated(cfg); err != nil {
t.Fatalf("unexpected error: %v", err)
@@ -187,7 +193,7 @@ func TestAmpModule_OnConfigUpdated_NonMultiSourceSecret(t *testing.T) {
m := &AmpModule{enabled: true}
m.secretSource = NewStaticSecretSource("static-key")
cfg := &config.Config{AmpUpstreamURL: "http://example.com"}
cfg := &config.Config{AmpCode: config.AmpCode{UpstreamURL: "http://example.com"}}
// Should not error or panic
if err := m.OnConfigUpdated(cfg); err != nil {
@@ -240,8 +246,10 @@ func TestAmpModule_SecretSource_FromConfig(t *testing.T) {
// Config with explicit API key
cfg := &config.Config{
AmpUpstreamURL: upstream.URL,
AmpUpstreamAPIKey: "config-key",
AmpCode: config.AmpCode{
UpstreamURL: upstream.URL,
UpstreamAPIKey: "config-key",
},
}
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
@@ -283,7 +291,7 @@ func TestAmpModule_ProviderAliasesAlwaysRegistered(t *testing.T) {
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
cfg := &config.Config{AmpUpstreamURL: scenario.configURL}
cfg := &config.Config{AmpCode: config.AmpCode{UpstreamURL: scenario.configURL}}
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
if err := m.Register(ctx); err != nil && scenario.configURL != "" {

View File

@@ -2,34 +2,118 @@ package amp
import (
"bytes"
"encoding/json"
"io"
"net/http/httputil"
"strings"
"time"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// AmpRouteType represents the type of routing decision made for an Amp request
type AmpRouteType string
const (
// RouteTypeLocalProvider indicates the request is handled by a local OAuth provider (free)
RouteTypeLocalProvider AmpRouteType = "LOCAL_PROVIDER"
// RouteTypeModelMapping indicates the request was remapped to another available model (free)
RouteTypeModelMapping AmpRouteType = "MODEL_MAPPING"
// RouteTypeAmpCredits indicates the request is forwarded to ampcode.com (uses Amp credits)
RouteTypeAmpCredits AmpRouteType = "AMP_CREDITS"
// RouteTypeNoProvider indicates no provider or fallback available
RouteTypeNoProvider AmpRouteType = "NO_PROVIDER"
)
// MappedModelContextKey is the Gin context key for passing mapped model names.
const MappedModelContextKey = "mapped_model"
// logAmpRouting logs the routing decision for an Amp request with structured fields
func logAmpRouting(routeType AmpRouteType, requestedModel, resolvedModel, provider, path string) {
fields := log.Fields{
"component": "amp-routing",
"route_type": string(routeType),
"requested_model": requestedModel,
"path": path,
"timestamp": time.Now().Format(time.RFC3339),
}
if resolvedModel != "" && resolvedModel != requestedModel {
fields["resolved_model"] = resolvedModel
}
if provider != "" {
fields["provider"] = provider
}
switch routeType {
case RouteTypeLocalProvider:
fields["cost"] = "free"
fields["source"] = "local_oauth"
log.WithFields(fields).Debugf("amp using local provider for model: %s", requestedModel)
case RouteTypeModelMapping:
fields["cost"] = "free"
fields["source"] = "local_oauth"
fields["mapping"] = requestedModel + " -> " + resolvedModel
// model mapping already logged in mapper; avoid duplicate here
case RouteTypeAmpCredits:
fields["cost"] = "amp_credits"
fields["source"] = "ampcode.com"
fields["model_id"] = requestedModel // Explicit model_id for easy config reference
log.WithFields(fields).Warnf("forwarding to ampcode.com (uses amp credits) - model_id: %s | To use local proxy, add to config: amp-model-mappings: [{from: \"%s\", to: \"<your-local-model>\"}]", requestedModel, requestedModel)
case RouteTypeNoProvider:
fields["cost"] = "none"
fields["source"] = "error"
fields["model_id"] = requestedModel // Explicit model_id for easy config reference
log.WithFields(fields).Warnf("no provider available for model_id: %s", requestedModel)
}
}
// FallbackHandler wraps a standard handler with fallback logic to ampcode.com
// when the model's provider is not available in CLIProxyAPI
type FallbackHandler struct {
getProxy func() *httputil.ReverseProxy
getProxy func() *httputil.ReverseProxy
modelMapper ModelMapper
forceModelMappings func() bool
}
// NewFallbackHandler creates a new fallback handler wrapper
// The getProxy function allows lazy evaluation of the proxy (useful when proxy is created after routes)
func NewFallbackHandler(getProxy func() *httputil.ReverseProxy) *FallbackHandler {
return &FallbackHandler{
getProxy: getProxy,
getProxy: getProxy,
forceModelMappings: func() bool { return false },
}
}
// NewFallbackHandlerWithMapper creates a new fallback handler with model mapping support
func NewFallbackHandlerWithMapper(getProxy func() *httputil.ReverseProxy, mapper ModelMapper, forceModelMappings func() bool) *FallbackHandler {
if forceModelMappings == nil {
forceModelMappings = func() bool { return false }
}
return &FallbackHandler{
getProxy: getProxy,
modelMapper: mapper,
forceModelMappings: forceModelMappings,
}
}
// SetModelMapper sets the model mapper for this handler (allows late binding)
func (fh *FallbackHandler) SetModelMapper(mapper ModelMapper) {
fh.modelMapper = mapper
}
// WrapHandler wraps a gin.HandlerFunc with fallback logic
// If the model's provider is not configured in CLIProxyAPI, it forwards to ampcode.com
func (fh *FallbackHandler) WrapHandler(handler gin.HandlerFunc) gin.HandlerFunc {
return func(c *gin.Context) {
requestPath := c.Request.URL.Path
// Read the request body to extract the model name
bodyBytes, err := io.ReadAll(c.Request.Body)
if err != nil {
@@ -52,15 +136,69 @@ func (fh *FallbackHandler) WrapHandler(handler gin.HandlerFunc) gin.HandlerFunc
// Normalize model (handles Gemini thinking suffixes)
normalizedModel, _ := util.NormalizeGeminiThinkingModel(modelName)
// Check if we have providers for this model
providers := util.GetProviderName(normalizedModel)
// Track resolved model for logging (may change if mapping is applied)
resolvedModel := normalizedModel
usedMapping := false
var providers []string
// Check if model mappings should be forced ahead of local API keys
forceMappings := fh.forceModelMappings != nil && fh.forceModelMappings()
if forceMappings {
// FORCE MODE: Check model mappings FIRST (takes precedence over local API keys)
// This allows users to route Amp requests to their preferred OAuth providers
if fh.modelMapper != nil {
if mappedModel := fh.modelMapper.MapModel(normalizedModel); mappedModel != "" {
// Mapping found - check if we have a provider for the mapped model
mappedProviders := util.GetProviderName(mappedModel)
if len(mappedProviders) > 0 {
// Mapping found and provider available - rewrite the model in request body
bodyBytes = rewriteModelInRequest(bodyBytes, mappedModel)
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
// Store mapped model in context for handlers that check it (like gemini bridge)
c.Set(MappedModelContextKey, mappedModel)
resolvedModel = mappedModel
usedMapping = true
providers = mappedProviders
}
}
}
// If no mapping applied, check for local providers
if !usedMapping {
providers = util.GetProviderName(normalizedModel)
}
} else {
// DEFAULT MODE: Check local providers first, then mappings as fallback
providers = util.GetProviderName(normalizedModel)
if len(providers) == 0 {
// No providers configured - check if we have a model mapping
if fh.modelMapper != nil {
if mappedModel := fh.modelMapper.MapModel(normalizedModel); mappedModel != "" {
// Mapping found - check if we have a provider for the mapped model
mappedProviders := util.GetProviderName(mappedModel)
if len(mappedProviders) > 0 {
// Mapping found and provider available - rewrite the model in request body
bodyBytes = rewriteModelInRequest(bodyBytes, mappedModel)
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
// Store mapped model in context for handlers that check it (like gemini bridge)
c.Set(MappedModelContextKey, mappedModel)
resolvedModel = mappedModel
usedMapping = true
providers = mappedProviders
}
}
}
}
}
// If no providers available, fallback to ampcode.com
if len(providers) == 0 {
// No providers configured - check if we have a proxy for fallback
proxy := fh.getProxy()
if proxy != nil {
// Fallback to ampcode.com
log.Infof("amp fallback: model %s has no configured provider, forwarding to ampcode.com", modelName)
// Log: Forwarding to ampcode.com (uses Amp credits)
logAmpRouting(RouteTypeAmpCredits, modelName, "", "", requestPath)
// Restore body again for the proxy
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
@@ -71,35 +209,73 @@ func (fh *FallbackHandler) WrapHandler(handler gin.HandlerFunc) gin.HandlerFunc
}
// No proxy available, let the normal handler return the error
log.Debugf("amp fallback: model %s has no configured provider and no proxy available", modelName)
logAmpRouting(RouteTypeNoProvider, modelName, "", "", requestPath)
}
// Providers available or no proxy for fallback, restore body and use normal handler
// Filter Anthropic-Beta header to remove features requiring special subscription
// This is needed when using local providers (bypassing the Amp proxy)
if betaHeader := c.Request.Header.Get("Anthropic-Beta"); betaHeader != "" {
filtered := filterBetaFeatures(betaHeader, "context-1m-2025-08-07")
if filtered != "" {
c.Request.Header.Set("Anthropic-Beta", filtered)
} else {
c.Request.Header.Del("Anthropic-Beta")
}
// Log the routing decision
providerName := ""
if len(providers) > 0 {
providerName = providers[0]
}
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
handler(c)
if usedMapping {
// Log: Model was mapped to another model
log.Debugf("amp model mapping: request %s -> %s", normalizedModel, resolvedModel)
logAmpRouting(RouteTypeModelMapping, modelName, resolvedModel, providerName, requestPath)
rewriter := NewResponseRewriter(c.Writer, normalizedModel)
c.Writer = rewriter
// Filter Anthropic-Beta header only for local handling paths
filterAntropicBetaHeader(c)
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
handler(c)
rewriter.Flush()
log.Debugf("amp model mapping: response %s -> %s", resolvedModel, normalizedModel)
} else if len(providers) > 0 {
// Log: Using local provider (free)
logAmpRouting(RouteTypeLocalProvider, modelName, resolvedModel, providerName, requestPath)
// Filter Anthropic-Beta header only for local handling paths
filterAntropicBetaHeader(c)
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
handler(c)
} else {
// No provider, no mapping, no proxy: fall back to the wrapped handler so it can return an error response
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
handler(c)
}
}
}
// filterAntropicBetaHeader filters Anthropic-Beta header to remove features requiring special subscription
// This is needed when using local providers (bypassing the Amp proxy)
func filterAntropicBetaHeader(c *gin.Context) {
if betaHeader := c.Request.Header.Get("Anthropic-Beta"); betaHeader != "" {
if filtered := filterBetaFeatures(betaHeader, "context-1m-2025-08-07"); filtered != "" {
c.Request.Header.Set("Anthropic-Beta", filtered)
} else {
c.Request.Header.Del("Anthropic-Beta")
}
}
}
// rewriteModelInRequest replaces the model name in a JSON request body
func rewriteModelInRequest(body []byte, newModel string) []byte {
if !gjson.GetBytes(body, "model").Exists() {
return body
}
result, err := sjson.SetBytes(body, "model", newModel)
if err != nil {
log.Warnf("amp model mapping: failed to rewrite model in request body: %v", err)
return body
}
return result
}
// extractModelFromRequest attempts to extract the model name from various request formats
func extractModelFromRequest(body []byte, c *gin.Context) string {
// First try to parse from JSON body (OpenAI, Claude, etc.)
var payload map[string]interface{}
if err := json.Unmarshal(body, &payload); err == nil {
// Check common model field names
if model, ok := payload["model"].(string); ok {
return model
}
// Check common model field names
if result := gjson.GetBytes(body, "model"); result.Exists() && result.Type == gjson.String {
return result.String()
}
// For Gemini requests, model is in the URL path

View File

@@ -4,7 +4,6 @@ import (
"strings"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/gemini"
)
// createGeminiBridgeHandler creates a handler that bridges AMP CLI's non-standard Gemini paths
@@ -15,16 +14,31 @@ import (
//
// This extracts the model+method from the AMP path and sets it as the :action parameter
// so the standard Gemini handler can process it.
func createGeminiBridgeHandler(geminiHandler *gemini.GeminiAPIHandler) gin.HandlerFunc {
//
// The handler parameter should be a Gemini-compatible handler that expects the :action param.
func createGeminiBridgeHandler(handler gin.HandlerFunc) gin.HandlerFunc {
return func(c *gin.Context) {
// Get the full path from the catch-all parameter
path := c.Param("path")
// Extract model:method from AMP CLI path format
// Example: /publishers/google/models/gemini-3-pro-preview:streamGenerateContent
if idx := strings.Index(path, "/models/"); idx >= 0 {
// Extract everything after "/models/"
actionPart := path[idx+8:] // Skip "/models/"
const modelsPrefix = "/models/"
if idx := strings.Index(path, modelsPrefix); idx >= 0 {
// Extract everything after modelsPrefix
actionPart := path[idx+len(modelsPrefix):]
// Check if model was mapped by FallbackHandler
if mappedModel, exists := c.Get(MappedModelContextKey); exists {
if strModel, ok := mappedModel.(string); ok && strModel != "" {
// Replace the model part in the action
// actionPart is like "model-name:method"
if colonIdx := strings.Index(actionPart, ":"); colonIdx > 0 {
method := actionPart[colonIdx:] // ":method"
actionPart = strModel + method
}
}
}
// Set this as the :action parameter that the Gemini handler expects
c.Params = append(c.Params, gin.Param{
@@ -32,8 +46,8 @@ func createGeminiBridgeHandler(geminiHandler *gemini.GeminiAPIHandler) gin.Handl
Value: actionPart,
})
// Call the standard Gemini handler
geminiHandler.GeminiHandler(c)
// Call the handler
handler(c)
return
}

View File

@@ -0,0 +1,93 @@
package amp
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
)
func TestCreateGeminiBridgeHandler_ActionParameterExtraction(t *testing.T) {
gin.SetMode(gin.TestMode)
tests := []struct {
name string
path string
mappedModel string // empty string means no mapping
expectedAction string
}{
{
name: "no_mapping_uses_url_model",
path: "/publishers/google/models/gemini-pro:generateContent",
mappedModel: "",
expectedAction: "gemini-pro:generateContent",
},
{
name: "mapped_model_replaces_url_model",
path: "/publishers/google/models/gemini-exp:generateContent",
mappedModel: "gemini-2.0-flash",
expectedAction: "gemini-2.0-flash:generateContent",
},
{
name: "mapping_preserves_method",
path: "/publishers/google/models/gemini-2.5-preview:streamGenerateContent",
mappedModel: "gemini-flash",
expectedAction: "gemini-flash:streamGenerateContent",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
var capturedAction string
mockGeminiHandler := func(c *gin.Context) {
capturedAction = c.Param("action")
c.JSON(http.StatusOK, gin.H{"captured": capturedAction})
}
// Use the actual createGeminiBridgeHandler function
bridgeHandler := createGeminiBridgeHandler(mockGeminiHandler)
r := gin.New()
if tt.mappedModel != "" {
r.Use(func(c *gin.Context) {
c.Set(MappedModelContextKey, tt.mappedModel)
c.Next()
})
}
r.POST("/api/provider/google/v1beta1/*path", bridgeHandler)
req := httptest.NewRequest(http.MethodPost, "/api/provider/google/v1beta1"+tt.path, nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("Expected status 200, got %d", w.Code)
}
if capturedAction != tt.expectedAction {
t.Errorf("Expected action '%s', got '%s'", tt.expectedAction, capturedAction)
}
})
}
}
func TestCreateGeminiBridgeHandler_InvalidPath(t *testing.T) {
gin.SetMode(gin.TestMode)
mockHandler := func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
}
bridgeHandler := createGeminiBridgeHandler(mockHandler)
r := gin.New()
r.POST("/api/provider/google/v1beta1/*path", bridgeHandler)
req := httptest.NewRequest(http.MethodPost, "/api/provider/google/v1beta1/invalid/path", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("Expected status 400 for invalid path, got %d", w.Code)
}
}

View File

@@ -0,0 +1,112 @@
// Package amp provides model mapping functionality for routing Amp CLI requests
// to alternative models when the requested model is not available locally.
package amp
import (
"strings"
"sync"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
log "github.com/sirupsen/logrus"
)
// ModelMapper provides model name mapping/aliasing for Amp CLI requests.
// When an Amp request comes in for a model that isn't available locally,
// this mapper can redirect it to an alternative model that IS available.
type ModelMapper interface {
// MapModel returns the target model name if a mapping exists and the target
// model has available providers. Returns empty string if no mapping applies.
MapModel(requestedModel string) string
// UpdateMappings refreshes the mapping configuration (for hot-reload).
UpdateMappings(mappings []config.AmpModelMapping)
}
// DefaultModelMapper implements ModelMapper with thread-safe mapping storage.
type DefaultModelMapper struct {
mu sync.RWMutex
mappings map[string]string // from -> to (normalized lowercase keys)
}
// NewModelMapper creates a new model mapper with the given initial mappings.
func NewModelMapper(mappings []config.AmpModelMapping) *DefaultModelMapper {
m := &DefaultModelMapper{
mappings: make(map[string]string),
}
m.UpdateMappings(mappings)
return m
}
// MapModel checks if a mapping exists for the requested model and if the
// target model has available local providers. Returns the mapped model name
// or empty string if no valid mapping exists.
func (m *DefaultModelMapper) MapModel(requestedModel string) string {
if requestedModel == "" {
return ""
}
m.mu.RLock()
defer m.mu.RUnlock()
// Normalize the requested model for lookup
normalizedRequest := strings.ToLower(strings.TrimSpace(requestedModel))
// Check for direct mapping
targetModel, exists := m.mappings[normalizedRequest]
if !exists {
return ""
}
// Verify target model has available providers
providers := util.GetProviderName(targetModel)
if len(providers) == 0 {
log.Debugf("amp model mapping: target model %s has no available providers, skipping mapping", targetModel)
return ""
}
// Note: Detailed routing log is handled by logAmpRouting in fallback_handlers.go
return targetModel
}
// UpdateMappings refreshes the mapping configuration from config.
// This is called during initialization and on config hot-reload.
func (m *DefaultModelMapper) UpdateMappings(mappings []config.AmpModelMapping) {
m.mu.Lock()
defer m.mu.Unlock()
// Clear and rebuild mappings
m.mappings = make(map[string]string, len(mappings))
for _, mapping := range mappings {
from := strings.TrimSpace(mapping.From)
to := strings.TrimSpace(mapping.To)
if from == "" || to == "" {
log.Warnf("amp model mapping: skipping invalid mapping (from=%q, to=%q)", from, to)
continue
}
// Store with normalized lowercase key for case-insensitive lookup
normalizedFrom := strings.ToLower(from)
m.mappings[normalizedFrom] = to
log.Debugf("amp model mapping registered: %s -> %s", from, to)
}
if len(m.mappings) > 0 {
log.Infof("amp model mapping: loaded %d mapping(s)", len(m.mappings))
}
}
// GetMappings returns a copy of current mappings (for debugging/status).
func (m *DefaultModelMapper) GetMappings() map[string]string {
m.mu.RLock()
defer m.mu.RUnlock()
result := make(map[string]string, len(m.mappings))
for k, v := range m.mappings {
result[k] = v
}
return result
}

View File

@@ -0,0 +1,186 @@
package amp
import (
"testing"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
)
func TestNewModelMapper(t *testing.T) {
mappings := []config.AmpModelMapping{
{From: "claude-opus-4.5", To: "claude-sonnet-4"},
{From: "gpt-5", To: "gemini-2.5-pro"},
}
mapper := NewModelMapper(mappings)
if mapper == nil {
t.Fatal("Expected non-nil mapper")
}
result := mapper.GetMappings()
if len(result) != 2 {
t.Errorf("Expected 2 mappings, got %d", len(result))
}
}
func TestNewModelMapper_Empty(t *testing.T) {
mapper := NewModelMapper(nil)
if mapper == nil {
t.Fatal("Expected non-nil mapper")
}
result := mapper.GetMappings()
if len(result) != 0 {
t.Errorf("Expected 0 mappings, got %d", len(result))
}
}
func TestModelMapper_MapModel_NoProvider(t *testing.T) {
mappings := []config.AmpModelMapping{
{From: "claude-opus-4.5", To: "claude-sonnet-4"},
}
mapper := NewModelMapper(mappings)
// Without a registered provider for the target, mapping should return empty
result := mapper.MapModel("claude-opus-4.5")
if result != "" {
t.Errorf("Expected empty result when target has no provider, got %s", result)
}
}
func TestModelMapper_MapModel_WithProvider(t *testing.T) {
// Register a mock provider for the target model
reg := registry.GetGlobalRegistry()
reg.RegisterClient("test-client", "claude", []*registry.ModelInfo{
{ID: "claude-sonnet-4", OwnedBy: "anthropic", Type: "claude"},
})
defer reg.UnregisterClient("test-client")
mappings := []config.AmpModelMapping{
{From: "claude-opus-4.5", To: "claude-sonnet-4"},
}
mapper := NewModelMapper(mappings)
// With a registered provider, mapping should work
result := mapper.MapModel("claude-opus-4.5")
if result != "claude-sonnet-4" {
t.Errorf("Expected claude-sonnet-4, got %s", result)
}
}
func TestModelMapper_MapModel_CaseInsensitive(t *testing.T) {
reg := registry.GetGlobalRegistry()
reg.RegisterClient("test-client2", "claude", []*registry.ModelInfo{
{ID: "claude-sonnet-4", OwnedBy: "anthropic", Type: "claude"},
})
defer reg.UnregisterClient("test-client2")
mappings := []config.AmpModelMapping{
{From: "Claude-Opus-4.5", To: "claude-sonnet-4"},
}
mapper := NewModelMapper(mappings)
// Should match case-insensitively
result := mapper.MapModel("claude-opus-4.5")
if result != "claude-sonnet-4" {
t.Errorf("Expected claude-sonnet-4, got %s", result)
}
}
func TestModelMapper_MapModel_NotFound(t *testing.T) {
mappings := []config.AmpModelMapping{
{From: "claude-opus-4.5", To: "claude-sonnet-4"},
}
mapper := NewModelMapper(mappings)
// Unknown model should return empty
result := mapper.MapModel("unknown-model")
if result != "" {
t.Errorf("Expected empty for unknown model, got %s", result)
}
}
func TestModelMapper_MapModel_EmptyInput(t *testing.T) {
mappings := []config.AmpModelMapping{
{From: "claude-opus-4.5", To: "claude-sonnet-4"},
}
mapper := NewModelMapper(mappings)
result := mapper.MapModel("")
if result != "" {
t.Errorf("Expected empty for empty input, got %s", result)
}
}
func TestModelMapper_UpdateMappings(t *testing.T) {
mapper := NewModelMapper(nil)
// Initially empty
if len(mapper.GetMappings()) != 0 {
t.Error("Expected 0 initial mappings")
}
// Update with new mappings
mapper.UpdateMappings([]config.AmpModelMapping{
{From: "model-a", To: "model-b"},
{From: "model-c", To: "model-d"},
})
result := mapper.GetMappings()
if len(result) != 2 {
t.Errorf("Expected 2 mappings after update, got %d", len(result))
}
// Update again should replace, not append
mapper.UpdateMappings([]config.AmpModelMapping{
{From: "model-x", To: "model-y"},
})
result = mapper.GetMappings()
if len(result) != 1 {
t.Errorf("Expected 1 mapping after second update, got %d", len(result))
}
}
func TestModelMapper_UpdateMappings_SkipsInvalid(t *testing.T) {
mapper := NewModelMapper(nil)
mapper.UpdateMappings([]config.AmpModelMapping{
{From: "", To: "model-b"}, // Invalid: empty from
{From: "model-a", To: ""}, // Invalid: empty to
{From: " ", To: "model-b"}, // Invalid: whitespace from
{From: "model-c", To: "model-d"}, // Valid
})
result := mapper.GetMappings()
if len(result) != 1 {
t.Errorf("Expected 1 valid mapping, got %d", len(result))
}
}
func TestModelMapper_GetMappings_ReturnsCopy(t *testing.T) {
mappings := []config.AmpModelMapping{
{From: "model-a", To: "model-b"},
}
mapper := NewModelMapper(mappings)
// Get mappings and modify the returned map
result := mapper.GetMappings()
result["new-key"] = "new-value"
// Original should be unchanged
original := mapper.GetMappings()
if len(original) != 1 {
t.Errorf("Expected original to have 1 mapping, got %d", len(original))
}
if _, exists := original["new-key"]; exists {
t.Error("Original map was modified")
}
}

View File

@@ -0,0 +1,98 @@
package amp
import (
"bytes"
"net/http"
"strings"
"github.com/gin-gonic/gin"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// ResponseRewriter wraps a gin.ResponseWriter to intercept and modify the response body
// It's used to rewrite model names in responses when model mapping is used
type ResponseRewriter struct {
gin.ResponseWriter
body *bytes.Buffer
originalModel string
isStreaming bool
}
// NewResponseRewriter creates a new response rewriter for model name substitution
func NewResponseRewriter(w gin.ResponseWriter, originalModel string) *ResponseRewriter {
return &ResponseRewriter{
ResponseWriter: w,
body: &bytes.Buffer{},
originalModel: originalModel,
}
}
// Write intercepts response writes and buffers them for model name replacement
func (rw *ResponseRewriter) Write(data []byte) (int, error) {
// Detect streaming on first write
if rw.body.Len() == 0 && !rw.isStreaming {
contentType := rw.Header().Get("Content-Type")
rw.isStreaming = strings.Contains(contentType, "text/event-stream") ||
strings.Contains(contentType, "stream")
}
if rw.isStreaming {
return rw.ResponseWriter.Write(rw.rewriteStreamChunk(data))
}
return rw.body.Write(data)
}
// Flush writes the buffered response with model names rewritten
func (rw *ResponseRewriter) Flush() {
if rw.isStreaming {
if flusher, ok := rw.ResponseWriter.(http.Flusher); ok {
flusher.Flush()
}
return
}
if rw.body.Len() > 0 {
if _, err := rw.ResponseWriter.Write(rw.rewriteModelInResponse(rw.body.Bytes())); err != nil {
log.Warnf("amp response rewriter: failed to write rewritten response: %v", err)
}
}
}
// modelFieldPaths lists all JSON paths where model name may appear
var modelFieldPaths = []string{"model", "modelVersion", "response.modelVersion", "message.model"}
// rewriteModelInResponse replaces all occurrences of the mapped model with the original model in JSON
func (rw *ResponseRewriter) rewriteModelInResponse(data []byte) []byte {
if rw.originalModel == "" {
return data
}
for _, path := range modelFieldPaths {
if gjson.GetBytes(data, path).Exists() {
data, _ = sjson.SetBytes(data, path, rw.originalModel)
}
}
return data
}
// rewriteStreamChunk rewrites model names in SSE stream chunks
func (rw *ResponseRewriter) rewriteStreamChunk(chunk []byte) []byte {
if rw.originalModel == "" {
return chunk
}
// SSE format: "data: {json}\n\n"
lines := bytes.Split(chunk, []byte("\n"))
for i, line := range lines {
if bytes.HasPrefix(line, []byte("data: ")) {
jsonData := bytes.TrimPrefix(line, []byte("data: "))
if len(jsonData) > 0 && jsonData[0] == '{' {
// Rewrite JSON in the data line
rewritten := rw.rewriteModelInResponse(jsonData)
lines[i] = append([]byte("data: "), rewritten...)
}
}
}
return bytes.Join(lines, []byte("\n"))
}

View File

@@ -1,12 +1,14 @@
package amp
import (
"errors"
"net"
"net/http"
"net/http/httputil"
"strings"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
"github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/claude"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/gemini"
@@ -14,15 +16,16 @@ import (
log "github.com/sirupsen/logrus"
)
// localhostOnlyMiddleware restricts access to localhost (127.0.0.1, ::1) only.
// Returns 403 Forbidden for non-localhost clients.
//
// Security: Uses RemoteAddr (actual TCP connection) instead of ClientIP() to prevent
// header spoofing attacks via X-Forwarded-For or similar headers. This means the
// middleware will not work correctly behind reverse proxies - users deploying behind
// nginx/Cloudflare should disable this feature and use firewall rules instead.
func localhostOnlyMiddleware() gin.HandlerFunc {
// localhostOnlyMiddleware returns a middleware that dynamically checks the module's
// localhost restriction setting. This allows hot-reload of the restriction without restarting.
func (m *AmpModule) localhostOnlyMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
// Check current setting (hot-reloadable)
if !m.IsRestrictedToLocalhost() {
c.Next()
return
}
// Use actual TCP connection address (RemoteAddr) to prevent header spoofing
// This cannot be forged by X-Forwarded-For or other client-controlled headers
remoteAddr := c.Request.RemoteAddr
@@ -37,7 +40,7 @@ func localhostOnlyMiddleware() gin.HandlerFunc {
// Parse the IP to handle both IPv4 and IPv6
ip := net.ParseIP(host)
if ip == nil {
log.Warnf("Amp management: invalid RemoteAddr %s, denying access", remoteAddr)
log.Warnf("amp management: invalid RemoteAddr %s, denying access", remoteAddr)
c.AbortWithStatusJSON(403, gin.H{
"error": "Access denied: management routes restricted to localhost",
})
@@ -46,7 +49,7 @@ func localhostOnlyMiddleware() gin.HandlerFunc {
// Check if IP is loopback (127.0.0.1 or ::1)
if !ip.IsLoopback() {
log.Warnf("Amp management: non-localhost connection from %s attempted access, denying", remoteAddr)
log.Warnf("amp management: non-localhost connection from %s attempted access, denying", remoteAddr)
c.AbortWithStatusJSON(403, gin.H{
"error": "Access denied: management routes restricted to localhost",
})
@@ -77,21 +80,56 @@ func noCORSMiddleware() gin.HandlerFunc {
}
}
// managementAvailabilityMiddleware short-circuits management routes when the upstream
// proxy is disabled, preventing noisy localhost warnings and accidental exposure.
func (m *AmpModule) managementAvailabilityMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
if m.getProxy() == nil {
logging.SkipGinRequestLogging(c)
c.AbortWithStatusJSON(http.StatusServiceUnavailable, gin.H{
"error": "amp upstream proxy not available",
})
return
}
c.Next()
}
}
// registerManagementRoutes registers Amp management proxy routes
// These routes proxy through to the Amp control plane for OAuth, user management, etc.
// If restrictToLocalhost is true, routes will only accept connections from 127.0.0.1/::1.
func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *handlers.BaseAPIHandler, proxyHandler gin.HandlerFunc, restrictToLocalhost bool) {
// Uses dynamic middleware and proxy getter for hot-reload support.
func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *handlers.BaseAPIHandler) {
ampAPI := engine.Group("/api")
// Always disable CORS for management routes to prevent browser-based attacks
ampAPI.Use(noCORSMiddleware())
ampAPI.Use(m.managementAvailabilityMiddleware(), noCORSMiddleware())
// Apply localhost-only restriction if configured
if restrictToLocalhost {
ampAPI.Use(localhostOnlyMiddleware())
log.Info("Amp management routes restricted to localhost only (CORS disabled)")
} else {
log.Warn("⚠️ Amp management routes are NOT restricted to localhost - this is insecure!")
// Apply dynamic localhost-only restriction (hot-reloadable via m.IsRestrictedToLocalhost())
ampAPI.Use(m.localhostOnlyMiddleware())
if !m.IsRestrictedToLocalhost() {
log.Warn("amp management routes are NOT restricted to localhost - this is insecure!")
}
// Dynamic proxy handler that uses m.getProxy() for hot-reload support
proxyHandler := func(c *gin.Context) {
// Swallow ErrAbortHandler panics from ReverseProxy copyResponse to avoid noisy stack traces
defer func() {
if rec := recover(); rec != nil {
if err, ok := rec.(error); ok && errors.Is(err, http.ErrAbortHandler) {
// Upstream already wrote the status (often 404) before the client/stream ended.
return
}
panic(rec)
}
}()
proxy := m.getProxy()
if proxy == nil {
c.JSON(503, gin.H{"error": "amp upstream proxy not available"})
return
}
proxy.ServeHTTP(c.Writer, c.Request)
}
// Management routes - these are proxied directly to Amp upstream
@@ -110,44 +148,42 @@ func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *ha
ampAPI.Any("/threads/*path", proxyHandler)
ampAPI.Any("/otel", proxyHandler)
ampAPI.Any("/otel/*path", proxyHandler)
ampAPI.Any("/tab", proxyHandler)
ampAPI.Any("/tab/*path", proxyHandler)
// Root-level routes that AMP CLI expects without /api prefix
// These need the same security middleware as the /api/* routes
rootMiddleware := []gin.HandlerFunc{noCORSMiddleware()}
if restrictToLocalhost {
rootMiddleware = append(rootMiddleware, localhostOnlyMiddleware())
}
// These need the same security middleware as the /api/* routes (dynamic for hot-reload)
rootMiddleware := []gin.HandlerFunc{m.managementAvailabilityMiddleware(), noCORSMiddleware(), m.localhostOnlyMiddleware()}
engine.GET("/threads/*path", append(rootMiddleware, proxyHandler)...)
engine.GET("/threads.rss", append(rootMiddleware, proxyHandler)...)
// Root-level auth routes for CLI login flow
// Amp uses multiple auth routes: /auth/cli-login, /auth/callback, /auth/sign-in, /auth/logout
// We proxy all /auth/* to support the complete OAuth flow
engine.Any("/auth", append(rootMiddleware, proxyHandler)...)
engine.Any("/auth/*path", append(rootMiddleware, proxyHandler)...)
// Google v1beta1 passthrough with OAuth fallback
// AMP CLI uses non-standard paths like /publishers/google/models/...
// We bridge these to our standard Gemini handler to enable local OAuth.
// If no local OAuth is available, falls back to ampcode.com proxy.
geminiHandlers := gemini.NewGeminiAPIHandler(baseHandler)
geminiBridge := createGeminiBridgeHandler(geminiHandlers)
geminiV1Beta1Fallback := NewFallbackHandler(func() *httputil.ReverseProxy {
return m.proxy
})
geminiBridge := createGeminiBridgeHandler(geminiHandlers.GeminiHandler)
geminiV1Beta1Fallback := NewFallbackHandlerWithMapper(func() *httputil.ReverseProxy {
return m.getProxy()
}, m.modelMapper, m.forceModelMappings)
geminiV1Beta1Handler := geminiV1Beta1Fallback.WrapHandler(geminiBridge)
// Route POST model calls through Gemini bridge when a local provider exists, otherwise proxy.
// Route POST model calls through Gemini bridge with FallbackHandler.
// FallbackHandler checks provider -> mapping -> proxy fallback automatically.
// All other methods (e.g., GET model listing) always proxy to upstream to preserve Amp CLI behavior.
ampAPI.Any("/provider/google/v1beta1/*path", func(c *gin.Context) {
if c.Request.Method == "POST" {
// Attempt to extract the model name from the AMP-style path
if path := c.Param("path"); strings.Contains(path, "/models/") {
modelPart := path[strings.Index(path, "/models/")+len("/models/"):]
if colonIdx := strings.Index(modelPart, ":"); colonIdx > 0 {
modelPart = modelPart[:colonIdx]
}
if modelPart != "" {
normalized, _ := util.NormalizeGeminiThinkingModel(modelPart)
// Only handle locally when we have a provider; otherwise fall back to proxy
if providers := util.GetProviderName(normalized); len(providers) > 0 {
geminiV1Beta1Handler(c)
return
}
}
// POST with /models/ path -> use Gemini bridge with fallback handler
// FallbackHandler will check provider/mapping and proxy if needed
geminiV1Beta1Handler(c)
return
}
}
// Non-POST or no local provider available -> proxy upstream
@@ -169,10 +205,11 @@ func (m *AmpModule) registerProviderAliases(engine *gin.Engine, baseHandler *han
openaiResponsesHandlers := openai.NewOpenAIResponsesAPIHandler(baseHandler)
// Create fallback handler wrapper that forwards to ampcode.com when provider not found
// Uses lazy evaluation to access proxy (which is created after routes are registered)
fallbackHandler := NewFallbackHandler(func() *httputil.ReverseProxy {
return m.proxy
})
// Uses m.getProxy() for hot-reload support (proxy can be updated at runtime)
// Also includes model mapping support for routing unavailable models to alternatives
fallbackHandler := NewFallbackHandlerWithMapper(func() *httputil.ReverseProxy {
return m.getProxy()
}, m.modelMapper, m.forceModelMappings)
// Provider-specific routes under /api/provider/:provider
ampProviders := engine.Group("/api/provider")

View File

@@ -13,16 +13,26 @@ func TestRegisterManagementRoutes(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
// Spy to track if proxy handler was called
proxyCalled := false
proxyHandler := func(c *gin.Context) {
proxyCalled = true
c.String(200, "proxied")
// Create module with proxy for testing
m := &AmpModule{
restrictToLocalhost: false, // disable localhost restriction for tests
}
m := &AmpModule{}
// Create a mock proxy that tracks calls
proxyCalled := false
mockProxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
proxyCalled = true
w.WriteHeader(200)
w.Write([]byte("proxied"))
}))
defer mockProxy.Close()
// Create real proxy to mock server
proxy, _ := createReverseProxy(mockProxy.URL, NewStaticSecretSource(""))
m.setProxy(proxy)
base := &handlers.BaseAPIHandler{}
m.registerManagementRoutes(r, base, proxyHandler, false) // false = don't restrict to localhost in tests
m.registerManagementRoutes(r, base)
managementPaths := []struct {
path string
@@ -37,8 +47,14 @@ func TestRegisterManagementRoutes(t *testing.T) {
{"/api/meta", http.MethodGet},
{"/api/telemetry", http.MethodGet},
{"/api/threads", http.MethodGet},
{"/threads/", http.MethodGet},
{"/threads.rss", http.MethodGet}, // Root-level route (no /api prefix)
{"/api/otel", http.MethodGet},
{"/api/tab", http.MethodGet},
{"/api/tab/some/path", http.MethodGet},
{"/auth", http.MethodGet}, // Root-level auth route
{"/auth/cli-login", http.MethodGet}, // CLI login flow
{"/auth/callback", http.MethodGet}, // OAuth callback
// Google v1beta1 bridge should still proxy non-model requests (GET) and allow POST
{"/api/provider/google/v1beta1/models", http.MethodGet},
{"/api/provider/google/v1beta1/models", http.MethodPost},
@@ -226,8 +242,13 @@ func TestLocalhostOnlyMiddleware_PreventsSpoofing(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
// Apply localhost-only middleware
r.Use(localhostOnlyMiddleware())
// Create module with localhost restriction enabled
m := &AmpModule{
restrictToLocalhost: true,
}
// Apply dynamic localhost-only middleware
r.Use(m.localhostOnlyMiddleware())
r.GET("/test", func(c *gin.Context) {
c.String(http.StatusOK, "ok")
})
@@ -300,3 +321,53 @@ func TestLocalhostOnlyMiddleware_PreventsSpoofing(t *testing.T) {
})
}
}
func TestLocalhostOnlyMiddleware_HotReload(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
// Create module with localhost restriction initially enabled
m := &AmpModule{
restrictToLocalhost: true,
}
// Apply dynamic localhost-only middleware
r.Use(m.localhostOnlyMiddleware())
r.GET("/test", func(c *gin.Context) {
c.String(http.StatusOK, "ok")
})
// Test 1: Remote IP should be blocked when restriction is enabled
req := httptest.NewRequest(http.MethodGet, "/test", nil)
req.RemoteAddr = "192.168.1.100:12345"
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusForbidden {
t.Errorf("Expected 403 when restriction enabled, got %d", w.Code)
}
// Test 2: Hot-reload - disable restriction
m.setRestrictToLocalhost(false)
req = httptest.NewRequest(http.MethodGet, "/test", nil)
req.RemoteAddr = "192.168.1.100:12345"
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Errorf("Expected 200 after disabling restriction, got %d", w.Code)
}
// Test 3: Hot-reload - re-enable restriction
m.setRestrictToLocalhost(true)
req = httptest.NewRequest(http.MethodGet, "/test", nil)
req.RemoteAddr = "192.168.1.100:12345"
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusForbidden {
t.Errorf("Expected 403 after re-enabling restriction, got %d", w.Code)
}
}

View File

@@ -139,6 +139,17 @@ func (s *MultiSourceSecret) InvalidateCache() {
s.cache = nil
}
// UpdateExplicitKey refreshes the config-provided key and clears cache.
func (s *MultiSourceSecret) UpdateExplicitKey(key string) {
if s == nil {
return
}
s.mu.Lock()
s.explicitKey = strings.TrimSpace(key)
s.cache = nil
s.mu.Unlock()
}
// StaticSecretSource returns a fixed API key (for testing)
type StaticSecretSource struct {
key string

View File

@@ -150,6 +150,9 @@ type Server struct {
// management handler
mgmt *managementHandlers.Handler
// ampModule is the Amp routing module for model mapping hot-reload
ampModule *ampmodule.AmpModule
// managementRoutesRegistered tracks whether the management routes have been attached to the engine.
managementRoutesRegistered atomic.Bool
// managementRoutesEnabled controls whether management endpoints serve real handlers.
@@ -268,14 +271,14 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
s.setupRoutes()
// Register Amp module using V2 interface with Context
ampModule := ampmodule.NewLegacy(accessManager, AuthMiddleware(accessManager))
s.ampModule = ampmodule.NewLegacy(accessManager, AuthMiddleware(accessManager))
ctx := modules.Context{
Engine: engine,
BaseHandler: s.handlers,
Config: cfg,
AuthMiddleware: AuthMiddleware(accessManager),
}
if err := modules.RegisterModule(ctx, ampModule); err != nil {
if err := modules.RegisterModule(ctx, s.ampModule); err != nil {
log.Errorf("Failed to register Amp module: %v", err)
}
@@ -297,7 +300,7 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
// Create HTTP server
s.server = &http.Server{
Addr: fmt.Sprintf(":%d", cfg.Port),
Addr: fmt.Sprintf("%s:%d", cfg.Host, cfg.Port),
Handler: engine,
}
@@ -467,8 +470,9 @@ func (s *Server) registerManagementRoutes() {
{
mgmt.GET("/usage", s.mgmt.GetUsageStatistics)
mgmt.GET("/config", s.mgmt.GetConfig)
mgmt.GET("/config.yaml", s.mgmt.GetConfigYAML)
mgmt.PUT("/config.yaml", s.mgmt.PutConfigYAML)
mgmt.GET("/config.yaml", s.mgmt.GetConfigFile)
mgmt.GET("/latest-version", s.mgmt.GetLatestVersion)
mgmt.GET("/debug", s.mgmt.GetDebug)
mgmt.PUT("/debug", s.mgmt.PutDebug)
@@ -500,11 +504,6 @@ func (s *Server) registerManagementRoutes() {
mgmt.PATCH("/api-keys", s.mgmt.PatchAPIKeys)
mgmt.DELETE("/api-keys", s.mgmt.DeleteAPIKeys)
mgmt.GET("/generative-language-api-key", s.mgmt.GetGlKeys)
mgmt.PUT("/generative-language-api-key", s.mgmt.PutGlKeys)
mgmt.PATCH("/generative-language-api-key", s.mgmt.PatchGlKeys)
mgmt.DELETE("/generative-language-api-key", s.mgmt.DeleteGlKeys)
mgmt.GET("/gemini-api-key", s.mgmt.GetGeminiKeys)
mgmt.PUT("/gemini-api-key", s.mgmt.PutGeminiKeys)
mgmt.PATCH("/gemini-api-key", s.mgmt.PatchGeminiKey)
@@ -521,6 +520,26 @@ func (s *Server) registerManagementRoutes() {
mgmt.PUT("/ws-auth", s.mgmt.PutWebsocketAuth)
mgmt.PATCH("/ws-auth", s.mgmt.PutWebsocketAuth)
mgmt.GET("/ampcode", s.mgmt.GetAmpCode)
mgmt.GET("/ampcode/upstream-url", s.mgmt.GetAmpUpstreamURL)
mgmt.PUT("/ampcode/upstream-url", s.mgmt.PutAmpUpstreamURL)
mgmt.PATCH("/ampcode/upstream-url", s.mgmt.PutAmpUpstreamURL)
mgmt.DELETE("/ampcode/upstream-url", s.mgmt.DeleteAmpUpstreamURL)
mgmt.GET("/ampcode/upstream-api-key", s.mgmt.GetAmpUpstreamAPIKey)
mgmt.PUT("/ampcode/upstream-api-key", s.mgmt.PutAmpUpstreamAPIKey)
mgmt.PATCH("/ampcode/upstream-api-key", s.mgmt.PutAmpUpstreamAPIKey)
mgmt.DELETE("/ampcode/upstream-api-key", s.mgmt.DeleteAmpUpstreamAPIKey)
mgmt.GET("/ampcode/restrict-management-to-localhost", s.mgmt.GetAmpRestrictManagementToLocalhost)
mgmt.PUT("/ampcode/restrict-management-to-localhost", s.mgmt.PutAmpRestrictManagementToLocalhost)
mgmt.PATCH("/ampcode/restrict-management-to-localhost", s.mgmt.PutAmpRestrictManagementToLocalhost)
mgmt.GET("/ampcode/model-mappings", s.mgmt.GetAmpModelMappings)
mgmt.PUT("/ampcode/model-mappings", s.mgmt.PutAmpModelMappings)
mgmt.PATCH("/ampcode/model-mappings", s.mgmt.PatchAmpModelMappings)
mgmt.DELETE("/ampcode/model-mappings", s.mgmt.DeleteAmpModelMappings)
mgmt.GET("/ampcode/force-model-mappings", s.mgmt.GetAmpForceModelMappings)
mgmt.PUT("/ampcode/force-model-mappings", s.mgmt.PutAmpForceModelMappings)
mgmt.PATCH("/ampcode/force-model-mappings", s.mgmt.PutAmpForceModelMappings)
mgmt.GET("/request-retry", s.mgmt.GetRequestRetry)
mgmt.PUT("/request-retry", s.mgmt.PutRequestRetry)
mgmt.PATCH("/request-retry", s.mgmt.PutRequestRetry)
@@ -903,7 +922,7 @@ func (s *Server) UpdateClients(cfg *config.Config) {
for _, p := range cfg.OpenAICompatibility {
providerNames = append(providerNames, p.Name)
}
s.handlers.OpenAICompatProviders = providerNames
s.handlers.SetOpenAICompatProviders(providerNames)
s.handlers.UpdateClients(&cfg.SDKConfig)
@@ -916,28 +935,36 @@ func (s *Server) UpdateClients(cfg *config.Config) {
s.mgmt.SetAuthManager(s.handlers.AuthManager)
}
// Notify Amp module of config changes (for model mapping hot-reload)
if s.ampModule != nil {
log.Debugf("triggering amp module config update")
if err := s.ampModule.OnConfigUpdated(cfg); err != nil {
log.Errorf("failed to update Amp module config: %v", err)
}
} else {
log.Warnf("amp module is nil, skipping config update")
}
// Count client sources from configuration and auth directory
authFiles := util.CountAuthFiles(cfg.AuthDir)
geminiAPIKeyCount := len(cfg.GeminiKey)
claudeAPIKeyCount := len(cfg.ClaudeKey)
codexAPIKeyCount := len(cfg.CodexKey)
vertexAICompatCount := len(cfg.VertexCompatAPIKey)
openAICompatCount := 0
for i := range cfg.OpenAICompatibility {
entry := cfg.OpenAICompatibility[i]
if len(entry.APIKeyEntries) > 0 {
openAICompatCount += len(entry.APIKeyEntries)
continue
}
openAICompatCount += len(entry.APIKeys)
openAICompatCount += len(entry.APIKeyEntries)
}
total := authFiles + geminiAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + openAICompatCount
fmt.Printf("server clients and configuration updated: %d clients (%d auth files + %d Gemini API keys + %d Claude API keys + %d Codex keys + %d OpenAI-compat)\n",
total := authFiles + geminiAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + vertexAICompatCount + openAICompatCount
fmt.Printf("server clients and configuration updated: %d clients (%d auth files + %d Gemini API keys + %d Claude API keys + %d Codex keys + %d Vertex-compat + %d OpenAI-compat)\n",
total,
authFiles,
geminiAPIKeyCount,
claudeAPIKeyCount,
codexAPIKeyCount,
vertexAICompatCount,
openAICompatCount,
)
}

View File

@@ -242,6 +242,11 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {
platformURL = "https://console.anthropic.com/"
}
// Validate platformURL to prevent XSS - only allow http/https URLs
if !isValidURL(platformURL) {
platformURL = "https://console.anthropic.com/"
}
// Generate success page HTML with dynamic content
successHTML := s.generateSuccessHTML(setupRequired, platformURL)
@@ -251,6 +256,12 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {
}
}
// isValidURL checks if the URL is a valid http/https URL to prevent XSS
func isValidURL(urlStr string) bool {
urlStr = strings.TrimSpace(urlStr)
return strings.HasPrefix(urlStr, "https://") || strings.HasPrefix(urlStr, "http://")
}
// generateSuccessHTML creates the HTML content for the success page.
// It customizes the page based on whether additional setup is required
// and includes a link to the platform.

View File

@@ -239,6 +239,11 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {
platformURL = "https://platform.openai.com"
}
// Validate platformURL to prevent XSS - only allow http/https URLs
if !isValidURL(platformURL) {
platformURL = "https://platform.openai.com"
}
// Generate success page HTML with dynamic content
successHTML := s.generateSuccessHTML(setupRequired, platformURL)
@@ -248,6 +253,12 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {
}
}
// isValidURL checks if the URL is a valid http/https URL to prevent XSS
func isValidURL(urlStr string) bool {
urlStr = strings.TrimSpace(urlStr)
return strings.HasPrefix(urlStr, "https://") || strings.HasPrefix(urlStr, "http://")
}
// generateSuccessHTML creates the HTML content for the success page.
// It customizes the page based on whether additional setup is required
// and includes a link to the platform.

View File

@@ -76,7 +76,8 @@ func (g *GeminiAuth) GetAuthenticatedClient(ctx context.Context, ts *GeminiToken
auth := &proxy.Auth{User: username, Password: password}
dialer, errSOCKS5 := proxy.SOCKS5("tcp", proxyURL.Host, auth, proxy.Direct)
if errSOCKS5 != nil {
log.Fatalf("create SOCKS5 dialer failed: %v", errSOCKS5)
log.Errorf("create SOCKS5 dialer failed: %v", errSOCKS5)
return nil, fmt.Errorf("create SOCKS5 dialer failed: %w", errSOCKS5)
}
transport = &http.Transport{
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
@@ -238,7 +239,11 @@ func (g *GeminiAuth) getTokenFromWeb(ctx context.Context, config *oauth2.Config,
// Start the server in a goroutine.
go func() {
if err := server.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
log.Fatalf("ListenAndServe(): %v", err)
log.Errorf("ListenAndServe(): %v", err)
select {
case errChan <- err:
default:
}
}
}()

View File

@@ -9,6 +9,7 @@ import (
"io"
"net/http"
"net/url"
"os"
"strings"
"time"
@@ -28,10 +29,21 @@ const (
iFlowAPIKeyEndpoint = "https://platform.iflow.cn/api/openapi/apikey"
// Client credentials provided by iFlow for the Code Assist integration.
iFlowOAuthClientID = "10009311001"
iFlowOAuthClientSecret = "4Z3YjXycVsQvyGF1etiNlIBB4RsqSDtW"
iFlowOAuthClientID = "10009311001"
// Default client secret (can be overridden via IFLOW_CLIENT_SECRET env var)
defaultIFlowClientSecret = "4Z3YjXycVsQvyGF1etiNlIBB4RsqSDtW"
)
// getIFlowClientSecret returns the iFlow OAuth client secret.
// It first checks the IFLOW_CLIENT_SECRET environment variable,
// falling back to the default value if not set.
func getIFlowClientSecret() string {
if secret := os.Getenv("IFLOW_CLIENT_SECRET"); secret != "" {
return secret
}
return defaultIFlowClientSecret
}
// DefaultAPIBaseURL is the canonical chat completions endpoint.
const DefaultAPIBaseURL = "https://apis.iflow.cn/v1"
@@ -72,7 +84,7 @@ func (ia *IFlowAuth) ExchangeCodeForTokens(ctx context.Context, code, redirectUR
form.Set("code", code)
form.Set("redirect_uri", redirectURI)
form.Set("client_id", iFlowOAuthClientID)
form.Set("client_secret", iFlowOAuthClientSecret)
form.Set("client_secret", getIFlowClientSecret())
req, err := ia.newTokenRequest(ctx, form)
if err != nil {
@@ -88,7 +100,7 @@ func (ia *IFlowAuth) RefreshTokens(ctx context.Context, refreshToken string) (*I
form.Set("grant_type", "refresh_token")
form.Set("refresh_token", refreshToken)
form.Set("client_id", iFlowOAuthClientID)
form.Set("client_secret", iFlowOAuthClientSecret)
form.Set("client_secret", getIFlowClientSecret())
req, err := ia.newTokenRequest(ctx, form)
if err != nil {
@@ -104,7 +116,7 @@ func (ia *IFlowAuth) newTokenRequest(ctx context.Context, form url.Values) (*htt
return nil, fmt.Errorf("iflow token: create request failed: %w", err)
}
basic := base64.StdEncoding.EncodeToString([]byte(iFlowOAuthClientID + ":" + iFlowOAuthClientSecret))
basic := base64.StdEncoding.EncodeToString([]byte(iFlowOAuthClientID + ":" + getIFlowClientSecret()))
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
req.Header.Set("Accept", "application/json")
req.Header.Set("Authorization", "Basic "+basic)
@@ -309,17 +321,23 @@ func (ia *IFlowAuth) AuthenticateWithCookie(ctx context.Context, cookie string)
return nil, fmt.Errorf("iflow cookie authentication: cookie is empty")
}
// First, get initial API key information using GET request
// First, get initial API key information using GET request to obtain the name
keyInfo, err := ia.fetchAPIKeyInfo(ctx, cookie)
if err != nil {
return nil, fmt.Errorf("iflow cookie authentication: fetch initial API key info failed: %w", err)
}
// Convert to token data format
// Refresh the API key using POST request
refreshedKeyInfo, err := ia.RefreshAPIKey(ctx, cookie, keyInfo.Name)
if err != nil {
return nil, fmt.Errorf("iflow cookie authentication: refresh API key failed: %w", err)
}
// Convert to token data format using refreshed key
data := &IFlowTokenData{
APIKey: keyInfo.APIKey,
Expire: keyInfo.ExpireTime,
Email: keyInfo.Name,
APIKey: refreshedKeyInfo.APIKey,
Expire: refreshedKeyInfo.ExpireTime,
Email: refreshedKeyInfo.Name,
Cookie: cookie,
}

301
internal/auth/kiro/aws.go Normal file
View File

@@ -0,0 +1,301 @@
// Package kiro provides authentication functionality for AWS CodeWhisperer (Kiro) API.
// It includes interfaces and implementations for token storage and authentication methods.
package kiro
import (
"encoding/base64"
"encoding/json"
"fmt"
"os"
"path/filepath"
"strings"
)
// PKCECodes holds PKCE verification codes for OAuth2 PKCE flow
type PKCECodes struct {
// CodeVerifier is the cryptographically random string used to correlate
// the authorization request to the token request
CodeVerifier string `json:"code_verifier"`
// CodeChallenge is the SHA256 hash of the code verifier, base64url-encoded
CodeChallenge string `json:"code_challenge"`
}
// KiroTokenData holds OAuth token information from AWS CodeWhisperer (Kiro)
type KiroTokenData struct {
// AccessToken is the OAuth2 access token for API access
AccessToken string `json:"accessToken"`
// RefreshToken is used to obtain new access tokens
RefreshToken string `json:"refreshToken"`
// ProfileArn is the AWS CodeWhisperer profile ARN
ProfileArn string `json:"profileArn"`
// ExpiresAt is the timestamp when the token expires
ExpiresAt string `json:"expiresAt"`
// AuthMethod indicates the authentication method used (e.g., "builder-id", "social")
AuthMethod string `json:"authMethod"`
// Provider indicates the OAuth provider (e.g., "AWS", "Google")
Provider string `json:"provider"`
// ClientID is the OIDC client ID (needed for token refresh)
ClientID string `json:"clientId,omitempty"`
// ClientSecret is the OIDC client secret (needed for token refresh)
ClientSecret string `json:"clientSecret,omitempty"`
// Email is the user's email address (used for file naming)
Email string `json:"email,omitempty"`
}
// KiroAuthBundle aggregates authentication data after OAuth flow completion
type KiroAuthBundle struct {
// TokenData contains the OAuth tokens from the authentication flow
TokenData KiroTokenData `json:"token_data"`
// LastRefresh is the timestamp of the last token refresh
LastRefresh string `json:"last_refresh"`
}
// KiroUsageInfo represents usage information from CodeWhisperer API
type KiroUsageInfo struct {
// SubscriptionTitle is the subscription plan name (e.g., "KIRO FREE")
SubscriptionTitle string `json:"subscription_title"`
// CurrentUsage is the current credit usage
CurrentUsage float64 `json:"current_usage"`
// UsageLimit is the maximum credit limit
UsageLimit float64 `json:"usage_limit"`
// NextReset is the timestamp of the next usage reset
NextReset string `json:"next_reset"`
}
// KiroModel represents a model available through the CodeWhisperer API
type KiroModel struct {
// ModelID is the unique identifier for the model
ModelID string `json:"modelId"`
// ModelName is the human-readable name
ModelName string `json:"modelName"`
// Description is the model description
Description string `json:"description"`
// RateMultiplier is the credit multiplier for this model
RateMultiplier float64 `json:"rateMultiplier"`
// RateUnit is the unit for rate calculation (e.g., "credit")
RateUnit string `json:"rateUnit"`
// MaxInputTokens is the maximum input token limit
MaxInputTokens int `json:"maxInputTokens,omitempty"`
}
// KiroIDETokenFile is the default path to Kiro IDE's token file
const KiroIDETokenFile = ".aws/sso/cache/kiro-auth-token.json"
// LoadKiroIDEToken loads token data from Kiro IDE's token file.
func LoadKiroIDEToken() (*KiroTokenData, error) {
homeDir, err := os.UserHomeDir()
if err != nil {
return nil, fmt.Errorf("failed to get home directory: %w", err)
}
tokenPath := filepath.Join(homeDir, KiroIDETokenFile)
data, err := os.ReadFile(tokenPath)
if err != nil {
return nil, fmt.Errorf("failed to read Kiro IDE token file (%s): %w", tokenPath, err)
}
var token KiroTokenData
if err := json.Unmarshal(data, &token); err != nil {
return nil, fmt.Errorf("failed to parse Kiro IDE token: %w", err)
}
if token.AccessToken == "" {
return nil, fmt.Errorf("access token is empty in Kiro IDE token file")
}
return &token, nil
}
// LoadKiroTokenFromPath loads token data from a custom path.
// This supports multiple accounts by allowing different token files.
func LoadKiroTokenFromPath(tokenPath string) (*KiroTokenData, error) {
// Expand ~ to home directory
if len(tokenPath) > 0 && tokenPath[0] == '~' {
homeDir, err := os.UserHomeDir()
if err != nil {
return nil, fmt.Errorf("failed to get home directory: %w", err)
}
tokenPath = filepath.Join(homeDir, tokenPath[1:])
}
data, err := os.ReadFile(tokenPath)
if err != nil {
return nil, fmt.Errorf("failed to read token file (%s): %w", tokenPath, err)
}
var token KiroTokenData
if err := json.Unmarshal(data, &token); err != nil {
return nil, fmt.Errorf("failed to parse token file: %w", err)
}
if token.AccessToken == "" {
return nil, fmt.Errorf("access token is empty in token file")
}
return &token, nil
}
// ListKiroTokenFiles lists all Kiro token files in the cache directory.
// This supports multiple accounts by finding all token files.
func ListKiroTokenFiles() ([]string, error) {
homeDir, err := os.UserHomeDir()
if err != nil {
return nil, fmt.Errorf("failed to get home directory: %w", err)
}
cacheDir := filepath.Join(homeDir, ".aws", "sso", "cache")
// Check if directory exists
if _, err := os.Stat(cacheDir); os.IsNotExist(err) {
return nil, nil // No token files
}
entries, err := os.ReadDir(cacheDir)
if err != nil {
return nil, fmt.Errorf("failed to read cache directory: %w", err)
}
var tokenFiles []string
for _, entry := range entries {
if entry.IsDir() {
continue
}
name := entry.Name()
// Look for kiro token files only (avoid matching unrelated AWS SSO cache files)
if strings.HasSuffix(name, ".json") && strings.HasPrefix(name, "kiro") {
tokenFiles = append(tokenFiles, filepath.Join(cacheDir, name))
}
}
return tokenFiles, nil
}
// LoadAllKiroTokens loads all Kiro tokens from the cache directory.
// This supports multiple accounts.
func LoadAllKiroTokens() ([]*KiroTokenData, error) {
files, err := ListKiroTokenFiles()
if err != nil {
return nil, err
}
var tokens []*KiroTokenData
for _, file := range files {
token, err := LoadKiroTokenFromPath(file)
if err != nil {
// Skip invalid token files
continue
}
tokens = append(tokens, token)
}
return tokens, nil
}
// JWTClaims represents the claims we care about from a JWT token.
// JWT tokens from Kiro/AWS contain user information in the payload.
type JWTClaims struct {
Email string `json:"email,omitempty"`
Sub string `json:"sub,omitempty"`
PreferredUser string `json:"preferred_username,omitempty"`
Name string `json:"name,omitempty"`
Iss string `json:"iss,omitempty"`
}
// ExtractEmailFromJWT extracts the user's email from a JWT access token.
// JWT tokens typically have format: header.payload.signature
// The payload is base64url-encoded JSON containing user claims.
func ExtractEmailFromJWT(accessToken string) string {
if accessToken == "" {
return ""
}
// JWT format: header.payload.signature
parts := strings.Split(accessToken, ".")
if len(parts) != 3 {
return ""
}
// Decode the payload (second part)
payload := parts[1]
// Add padding if needed (base64url requires padding)
switch len(payload) % 4 {
case 2:
payload += "=="
case 3:
payload += "="
}
decoded, err := base64.URLEncoding.DecodeString(payload)
if err != nil {
// Try RawURLEncoding (no padding)
decoded, err = base64.RawURLEncoding.DecodeString(parts[1])
if err != nil {
return ""
}
}
var claims JWTClaims
if err := json.Unmarshal(decoded, &claims); err != nil {
return ""
}
// Return email if available
if claims.Email != "" {
return claims.Email
}
// Fallback to preferred_username (some providers use this)
if claims.PreferredUser != "" && strings.Contains(claims.PreferredUser, "@") {
return claims.PreferredUser
}
// Fallback to sub if it looks like an email
if claims.Sub != "" && strings.Contains(claims.Sub, "@") {
return claims.Sub
}
return ""
}
// SanitizeEmailForFilename sanitizes an email address for use in a filename.
// Replaces special characters with underscores and prevents path traversal attacks.
// Also handles URL-encoded characters to prevent encoded path traversal attempts.
func SanitizeEmailForFilename(email string) string {
if email == "" {
return ""
}
result := email
// First, handle URL-encoded path traversal attempts (%2F, %2E, %5C, etc.)
// This prevents encoded characters from bypassing the sanitization.
// Note: We replace % last to catch any remaining encodings including double-encoding (%252F)
result = strings.ReplaceAll(result, "%2F", "_") // /
result = strings.ReplaceAll(result, "%2f", "_")
result = strings.ReplaceAll(result, "%5C", "_") // \
result = strings.ReplaceAll(result, "%5c", "_")
result = strings.ReplaceAll(result, "%2E", "_") // .
result = strings.ReplaceAll(result, "%2e", "_")
result = strings.ReplaceAll(result, "%00", "_") // null byte
result = strings.ReplaceAll(result, "%", "_") // Catch remaining % to prevent double-encoding attacks
// Replace characters that are problematic in filenames
// Keep @ and . in middle but replace other special characters
for _, char := range []string{"/", "\\", ":", "*", "?", "\"", "<", ">", "|", " ", "\x00"} {
result = strings.ReplaceAll(result, char, "_")
}
// Prevent path traversal: replace leading dots in each path component
// This handles cases like "../../../etc/passwd" → "_.._.._.._etc_passwd"
parts := strings.Split(result, "_")
for i, part := range parts {
for strings.HasPrefix(part, ".") {
part = "_" + part[1:]
}
parts[i] = part
}
result = strings.Join(parts, "_")
return result
}

View File

@@ -0,0 +1,314 @@
// Package kiro provides OAuth2 authentication functionality for AWS CodeWhisperer (Kiro) API.
// This package implements token loading, refresh, and API communication with CodeWhisperer.
package kiro
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"strings"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
log "github.com/sirupsen/logrus"
)
const (
// awsKiroEndpoint is used for CodeWhisperer management APIs (GetUsageLimits, ListProfiles, etc.)
// Note: This is different from the Amazon Q streaming endpoint (q.us-east-1.amazonaws.com)
// used in kiro_executor.go for GenerateAssistantResponse. Both endpoints are correct
// for their respective API operations.
awsKiroEndpoint = "https://codewhisperer.us-east-1.amazonaws.com"
defaultTokenFile = "~/.aws/sso/cache/kiro-auth-token.json"
targetGetUsage = "AmazonCodeWhispererService.GetUsageLimits"
targetListModels = "AmazonCodeWhispererService.ListAvailableModels"
targetGenerateChat = "AmazonCodeWhispererStreamingService.GenerateAssistantResponse"
)
// KiroAuth handles AWS CodeWhisperer authentication and API communication.
// It provides methods for loading tokens, refreshing expired tokens,
// and communicating with the CodeWhisperer API.
type KiroAuth struct {
httpClient *http.Client
endpoint string
}
// NewKiroAuth creates a new Kiro authentication service.
// It initializes the HTTP client with proxy settings from the configuration.
//
// Parameters:
// - cfg: The application configuration containing proxy settings
//
// Returns:
// - *KiroAuth: A new Kiro authentication service instance
func NewKiroAuth(cfg *config.Config) *KiroAuth {
return &KiroAuth{
httpClient: util.SetProxy(&cfg.SDKConfig, &http.Client{Timeout: 120 * time.Second}),
endpoint: awsKiroEndpoint,
}
}
// LoadTokenFromFile loads token data from a file path.
// This method reads and parses the token file, expanding ~ to the home directory.
//
// Parameters:
// - tokenFile: Path to the token file (supports ~ expansion)
//
// Returns:
// - *KiroTokenData: The parsed token data
// - error: An error if file reading or parsing fails
func (k *KiroAuth) LoadTokenFromFile(tokenFile string) (*KiroTokenData, error) {
// Expand ~ to home directory
if strings.HasPrefix(tokenFile, "~") {
home, err := os.UserHomeDir()
if err != nil {
return nil, fmt.Errorf("failed to get home directory: %w", err)
}
tokenFile = filepath.Join(home, tokenFile[1:])
}
data, err := os.ReadFile(tokenFile)
if err != nil {
return nil, fmt.Errorf("failed to read token file: %w", err)
}
var tokenData KiroTokenData
if err := json.Unmarshal(data, &tokenData); err != nil {
return nil, fmt.Errorf("failed to parse token file: %w", err)
}
return &tokenData, nil
}
// IsTokenExpired checks if the token has expired.
// This method parses the expiration timestamp and compares it with the current time.
//
// Parameters:
// - tokenData: The token data to check
//
// Returns:
// - bool: True if the token has expired, false otherwise
func (k *KiroAuth) IsTokenExpired(tokenData *KiroTokenData) bool {
if tokenData.ExpiresAt == "" {
return true
}
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
if err != nil {
// Try alternate format
expiresAt, err = time.Parse("2006-01-02T15:04:05.000Z", tokenData.ExpiresAt)
if err != nil {
return true
}
}
return time.Now().After(expiresAt)
}
// makeRequest sends a request to the CodeWhisperer API.
// This is an internal method for making authenticated API calls.
//
// Parameters:
// - ctx: The context for the request
// - target: The API target (e.g., "AmazonCodeWhispererService.GetUsageLimits")
// - accessToken: The OAuth access token
// - payload: The request payload
//
// Returns:
// - []byte: The response body
// - error: An error if the request fails
func (k *KiroAuth) makeRequest(ctx context.Context, target string, accessToken string, payload interface{}) ([]byte, error) {
jsonBody, err := json.Marshal(payload)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, k.endpoint, strings.NewReader(string(jsonBody)))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/x-amz-json-1.0")
req.Header.Set("x-amz-target", target)
req.Header.Set("Authorization", "Bearer "+accessToken)
req.Header.Set("Accept", "application/json")
resp, err := k.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("request failed: %w", err)
}
defer func() {
if errClose := resp.Body.Close(); errClose != nil {
log.Errorf("failed to close response body: %v", errClose)
}
}()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(body))
}
return body, nil
}
// GetUsageLimits retrieves usage information from the CodeWhisperer API.
// This method fetches the current usage statistics and subscription information.
//
// Parameters:
// - ctx: The context for the request
// - tokenData: The token data containing access token and profile ARN
//
// Returns:
// - *KiroUsageInfo: The usage information
// - error: An error if the request fails
func (k *KiroAuth) GetUsageLimits(ctx context.Context, tokenData *KiroTokenData) (*KiroUsageInfo, error) {
payload := map[string]interface{}{
"origin": "AI_EDITOR",
"profileArn": tokenData.ProfileArn,
"resourceType": "AGENTIC_REQUEST",
}
body, err := k.makeRequest(ctx, targetGetUsage, tokenData.AccessToken, payload)
if err != nil {
return nil, err
}
var result struct {
SubscriptionInfo struct {
SubscriptionTitle string `json:"subscriptionTitle"`
} `json:"subscriptionInfo"`
UsageBreakdownList []struct {
CurrentUsageWithPrecision float64 `json:"currentUsageWithPrecision"`
UsageLimitWithPrecision float64 `json:"usageLimitWithPrecision"`
} `json:"usageBreakdownList"`
NextDateReset float64 `json:"nextDateReset"`
}
if err := json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse usage response: %w", err)
}
usage := &KiroUsageInfo{
SubscriptionTitle: result.SubscriptionInfo.SubscriptionTitle,
NextReset: fmt.Sprintf("%v", result.NextDateReset),
}
if len(result.UsageBreakdownList) > 0 {
usage.CurrentUsage = result.UsageBreakdownList[0].CurrentUsageWithPrecision
usage.UsageLimit = result.UsageBreakdownList[0].UsageLimitWithPrecision
}
return usage, nil
}
// ListAvailableModels retrieves available models from the CodeWhisperer API.
// This method fetches the list of AI models available for the authenticated user.
//
// Parameters:
// - ctx: The context for the request
// - tokenData: The token data containing access token and profile ARN
//
// Returns:
// - []*KiroModel: The list of available models
// - error: An error if the request fails
func (k *KiroAuth) ListAvailableModels(ctx context.Context, tokenData *KiroTokenData) ([]*KiroModel, error) {
payload := map[string]interface{}{
"origin": "AI_EDITOR",
"profileArn": tokenData.ProfileArn,
}
body, err := k.makeRequest(ctx, targetListModels, tokenData.AccessToken, payload)
if err != nil {
return nil, err
}
var result struct {
Models []struct {
ModelID string `json:"modelId"`
ModelName string `json:"modelName"`
Description string `json:"description"`
RateMultiplier float64 `json:"rateMultiplier"`
RateUnit string `json:"rateUnit"`
TokenLimits struct {
MaxInputTokens int `json:"maxInputTokens"`
} `json:"tokenLimits"`
} `json:"models"`
}
if err := json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse models response: %w", err)
}
models := make([]*KiroModel, 0, len(result.Models))
for _, m := range result.Models {
models = append(models, &KiroModel{
ModelID: m.ModelID,
ModelName: m.ModelName,
Description: m.Description,
RateMultiplier: m.RateMultiplier,
RateUnit: m.RateUnit,
MaxInputTokens: m.TokenLimits.MaxInputTokens,
})
}
return models, nil
}
// CreateTokenStorage creates a new KiroTokenStorage from token data.
// This method converts the token data into a storage structure suitable for persistence.
//
// Parameters:
// - tokenData: The token data to convert
//
// Returns:
// - *KiroTokenStorage: A new token storage instance
func (k *KiroAuth) CreateTokenStorage(tokenData *KiroTokenData) *KiroTokenStorage {
return &KiroTokenStorage{
AccessToken: tokenData.AccessToken,
RefreshToken: tokenData.RefreshToken,
ProfileArn: tokenData.ProfileArn,
ExpiresAt: tokenData.ExpiresAt,
AuthMethod: tokenData.AuthMethod,
Provider: tokenData.Provider,
LastRefresh: time.Now().Format(time.RFC3339),
}
}
// ValidateToken checks if the token is valid by making a test API call.
// This method verifies the token by attempting to fetch usage limits.
//
// Parameters:
// - ctx: The context for the request
// - tokenData: The token data to validate
//
// Returns:
// - error: An error if the token is invalid
func (k *KiroAuth) ValidateToken(ctx context.Context, tokenData *KiroTokenData) error {
_, err := k.GetUsageLimits(ctx, tokenData)
return err
}
// UpdateTokenStorage updates an existing token storage with new token data.
// This method refreshes the token storage with newly obtained access and refresh tokens.
//
// Parameters:
// - storage: The existing token storage to update
// - tokenData: The new token data to apply
func (k *KiroAuth) UpdateTokenStorage(storage *KiroTokenStorage, tokenData *KiroTokenData) {
storage.AccessToken = tokenData.AccessToken
storage.RefreshToken = tokenData.RefreshToken
storage.ProfileArn = tokenData.ProfileArn
storage.ExpiresAt = tokenData.ExpiresAt
storage.AuthMethod = tokenData.AuthMethod
storage.Provider = tokenData.Provider
storage.LastRefresh = time.Now().Format(time.RFC3339)
}

View File

@@ -0,0 +1,161 @@
package kiro
import (
"encoding/base64"
"encoding/json"
"testing"
)
func TestExtractEmailFromJWT(t *testing.T) {
tests := []struct {
name string
token string
expected string
}{
{
name: "Empty token",
token: "",
expected: "",
},
{
name: "Invalid token format",
token: "not.a.valid.jwt",
expected: "",
},
{
name: "Invalid token - not base64",
token: "xxx.yyy.zzz",
expected: "",
},
{
name: "Valid JWT with email",
token: createTestJWT(map[string]any{"email": "test@example.com", "sub": "user123"}),
expected: "test@example.com",
},
{
name: "JWT without email but with preferred_username",
token: createTestJWT(map[string]any{"preferred_username": "user@domain.com", "sub": "user123"}),
expected: "user@domain.com",
},
{
name: "JWT with email-like sub",
token: createTestJWT(map[string]any{"sub": "another@test.com"}),
expected: "another@test.com",
},
{
name: "JWT without any email fields",
token: createTestJWT(map[string]any{"sub": "user123", "name": "Test User"}),
expected: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := ExtractEmailFromJWT(tt.token)
if result != tt.expected {
t.Errorf("ExtractEmailFromJWT() = %q, want %q", result, tt.expected)
}
})
}
}
func TestSanitizeEmailForFilename(t *testing.T) {
tests := []struct {
name string
email string
expected string
}{
{
name: "Empty email",
email: "",
expected: "",
},
{
name: "Simple email",
email: "user@example.com",
expected: "user@example.com",
},
{
name: "Email with space",
email: "user name@example.com",
expected: "user_name@example.com",
},
{
name: "Email with special chars",
email: "user:name@example.com",
expected: "user_name@example.com",
},
{
name: "Email with multiple special chars",
email: "user/name:test@example.com",
expected: "user_name_test@example.com",
},
{
name: "Path traversal attempt",
email: "../../../etc/passwd",
expected: "_.__.__._etc_passwd",
},
{
name: "Path traversal with backslash",
email: `..\..\..\..\windows\system32`,
expected: "_.__.__.__._windows_system32",
},
{
name: "Null byte injection attempt",
email: "user\x00@evil.com",
expected: "user_@evil.com",
},
// URL-encoded path traversal tests
{
name: "URL-encoded slash",
email: "user%2Fpath@example.com",
expected: "user_path@example.com",
},
{
name: "URL-encoded backslash",
email: "user%5Cpath@example.com",
expected: "user_path@example.com",
},
{
name: "URL-encoded dot",
email: "%2E%2E%2Fetc%2Fpasswd",
expected: "___etc_passwd",
},
{
name: "URL-encoded null",
email: "user%00@evil.com",
expected: "user_@evil.com",
},
{
name: "Double URL-encoding attack",
email: "%252F%252E%252E",
expected: "_252F_252E_252E", // % replaced with _, remaining chars preserved (safe)
},
{
name: "Mixed case URL-encoding",
email: "%2f%2F%5c%5C",
expected: "____",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := SanitizeEmailForFilename(tt.email)
if result != tt.expected {
t.Errorf("SanitizeEmailForFilename() = %q, want %q", result, tt.expected)
}
})
}
}
// createTestJWT creates a test JWT token with the given claims
func createTestJWT(claims map[string]any) string {
header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"RS256","typ":"JWT"}`))
payloadBytes, _ := json.Marshal(claims)
payload := base64.RawURLEncoding.EncodeToString(payloadBytes)
signature := base64.RawURLEncoding.EncodeToString([]byte("fake-signature"))
return header + "." + payload + "." + signature
}

296
internal/auth/kiro/oauth.go Normal file
View File

@@ -0,0 +1,296 @@
// Package kiro provides OAuth2 authentication for Kiro using native Google login.
package kiro
import (
"context"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"fmt"
"html"
"io"
"net"
"net/http"
"strings"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
log "github.com/sirupsen/logrus"
)
const (
// Kiro auth endpoint
kiroAuthEndpoint = "https://prod.us-east-1.auth.desktop.kiro.dev"
// Default callback port
defaultCallbackPort = 9876
// Auth timeout
authTimeout = 10 * time.Minute
)
// KiroTokenResponse represents the response from Kiro token endpoint.
type KiroTokenResponse struct {
AccessToken string `json:"accessToken"`
RefreshToken string `json:"refreshToken"`
ProfileArn string `json:"profileArn"`
ExpiresIn int `json:"expiresIn"`
}
// KiroOAuth handles the OAuth flow for Kiro authentication.
type KiroOAuth struct {
httpClient *http.Client
cfg *config.Config
}
// NewKiroOAuth creates a new Kiro OAuth handler.
func NewKiroOAuth(cfg *config.Config) *KiroOAuth {
client := &http.Client{Timeout: 30 * time.Second}
if cfg != nil {
client = util.SetProxy(&cfg.SDKConfig, client)
}
return &KiroOAuth{
httpClient: client,
cfg: cfg,
}
}
// generateCodeVerifier generates a random code verifier for PKCE.
func generateCodeVerifier() (string, error) {
b := make([]byte, 32)
if _, err := rand.Read(b); err != nil {
return "", err
}
return base64.RawURLEncoding.EncodeToString(b), nil
}
// generateCodeChallenge generates the code challenge from verifier.
func generateCodeChallenge(verifier string) string {
h := sha256.Sum256([]byte(verifier))
return base64.RawURLEncoding.EncodeToString(h[:])
}
// generateState generates a random state parameter.
func generateState() (string, error) {
b := make([]byte, 16)
if _, err := rand.Read(b); err != nil {
return "", err
}
return base64.RawURLEncoding.EncodeToString(b), nil
}
// AuthResult contains the authorization code and state from callback.
type AuthResult struct {
Code string
State string
Error string
}
// startCallbackServer starts a local HTTP server to receive the OAuth callback.
func (o *KiroOAuth) startCallbackServer(ctx context.Context, expectedState string) (string, <-chan AuthResult, error) {
// Try to find an available port - use localhost like Kiro does
listener, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", defaultCallbackPort))
if err != nil {
// Try with dynamic port (RFC 8252 allows dynamic ports for native apps)
log.Warnf("kiro oauth: default port %d is busy, falling back to dynamic port", defaultCallbackPort)
listener, err = net.Listen("tcp", "localhost:0")
if err != nil {
return "", nil, fmt.Errorf("failed to start callback server: %w", err)
}
}
port := listener.Addr().(*net.TCPAddr).Port
// Use http scheme for local callback server
redirectURI := fmt.Sprintf("http://localhost:%d/oauth/callback", port)
resultChan := make(chan AuthResult, 1)
server := &http.Server{
ReadHeaderTimeout: 10 * time.Second,
}
mux := http.NewServeMux()
mux.HandleFunc("/oauth/callback", func(w http.ResponseWriter, r *http.Request) {
code := r.URL.Query().Get("code")
state := r.URL.Query().Get("state")
errParam := r.URL.Query().Get("error")
if errParam != "" {
w.Header().Set("Content-Type", "text/html")
w.WriteHeader(http.StatusBadRequest)
fmt.Fprintf(w, `<html><body><h1>Login Failed</h1><p>%s</p><p>You can close this window.</p></body></html>`, html.EscapeString(errParam))
resultChan <- AuthResult{Error: errParam}
return
}
if state != expectedState {
w.Header().Set("Content-Type", "text/html")
w.WriteHeader(http.StatusBadRequest)
fmt.Fprint(w, `<html><body><h1>Login Failed</h1><p>Invalid state parameter</p><p>You can close this window.</p></body></html>`)
resultChan <- AuthResult{Error: "state mismatch"}
return
}
w.Header().Set("Content-Type", "text/html")
fmt.Fprint(w, `<html><body><h1>Login Successful!</h1><p>You can close this window and return to the terminal.</p></body></html>`)
resultChan <- AuthResult{Code: code, State: state}
})
server.Handler = mux
go func() {
if err := server.Serve(listener); err != nil && err != http.ErrServerClosed {
log.Debugf("callback server error: %v", err)
}
}()
go func() {
select {
case <-ctx.Done():
case <-time.After(authTimeout):
case <-resultChan:
}
_ = server.Shutdown(context.Background())
}()
return redirectURI, resultChan, nil
}
// LoginWithBuilderID performs OAuth login with AWS Builder ID using device code flow.
func (o *KiroOAuth) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, error) {
ssoClient := NewSSOOIDCClient(o.cfg)
return ssoClient.LoginWithBuilderID(ctx)
}
// exchangeCodeForToken exchanges the authorization code for tokens.
func (o *KiroOAuth) exchangeCodeForToken(ctx context.Context, code, codeVerifier, redirectURI string) (*KiroTokenData, error) {
payload := map[string]string{
"code": code,
"code_verifier": codeVerifier,
"redirect_uri": redirectURI,
}
body, err := json.Marshal(payload)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
tokenURL := kiroAuthEndpoint + "/oauth/token"
req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(string(body)))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
resp, err := o.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("token request failed: %w", err)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
log.Debugf("token exchange failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("token exchange failed (status %d)", resp.StatusCode)
}
var tokenResp KiroTokenResponse
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
return nil, fmt.Errorf("failed to parse token response: %w", err)
}
// Validate ExpiresIn - use default 1 hour if invalid
expiresIn := tokenResp.ExpiresIn
if expiresIn <= 0 {
expiresIn = 3600
}
expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
return &KiroTokenData{
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
ProfileArn: tokenResp.ProfileArn,
ExpiresAt: expiresAt.Format(time.RFC3339),
AuthMethod: "social",
Provider: "", // Caller should preserve original provider
}, nil
}
// RefreshToken refreshes an expired access token.
func (o *KiroOAuth) RefreshToken(ctx context.Context, refreshToken string) (*KiroTokenData, error) {
payload := map[string]string{
"refreshToken": refreshToken,
}
body, err := json.Marshal(payload)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
refreshURL := kiroAuthEndpoint + "/refreshToken"
req, err := http.NewRequestWithContext(ctx, http.MethodPost, refreshURL, strings.NewReader(string(body)))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
resp, err := o.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("refresh request failed: %w", err)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
}
var tokenResp KiroTokenResponse
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
return nil, fmt.Errorf("failed to parse token response: %w", err)
}
// Validate ExpiresIn - use default 1 hour if invalid
expiresIn := tokenResp.ExpiresIn
if expiresIn <= 0 {
expiresIn = 3600
}
expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
return &KiroTokenData{
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
ProfileArn: tokenResp.ProfileArn,
ExpiresAt: expiresAt.Format(time.RFC3339),
AuthMethod: "social",
Provider: "", // Caller should preserve original provider
}, nil
}
// LoginWithGoogle performs OAuth login with Google using Kiro's social auth.
// This uses a custom protocol handler (kiro://) to receive the callback.
func (o *KiroOAuth) LoginWithGoogle(ctx context.Context) (*KiroTokenData, error) {
socialClient := NewSocialAuthClient(o.cfg)
return socialClient.LoginWithGoogle(ctx)
}
// LoginWithGitHub performs OAuth login with GitHub using Kiro's social auth.
// This uses a custom protocol handler (kiro://) to receive the callback.
func (o *KiroOAuth) LoginWithGitHub(ctx context.Context) (*KiroTokenData, error) {
socialClient := NewSocialAuthClient(o.cfg)
return socialClient.LoginWithGitHub(ctx)
}

View File

@@ -0,0 +1,725 @@
// Package kiro provides custom protocol handler registration for Kiro OAuth.
// This enables the CLI to intercept kiro:// URIs for social authentication (Google/GitHub).
package kiro
import (
"context"
"fmt"
"html"
"net"
"net/http"
"net/url"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"sync"
"time"
log "github.com/sirupsen/logrus"
)
const (
// KiroProtocol is the custom URI scheme used by Kiro
KiroProtocol = "kiro"
// KiroAuthority is the URI authority for authentication callbacks
KiroAuthority = "kiro.kiroAgent"
// KiroAuthPath is the path for successful authentication
KiroAuthPath = "/authenticate-success"
// KiroRedirectURI is the full redirect URI for social auth
KiroRedirectURI = "kiro://kiro.kiroAgent/authenticate-success"
// DefaultHandlerPort is the default port for the local callback server
DefaultHandlerPort = 19876
// HandlerTimeout is how long to wait for the OAuth callback
HandlerTimeout = 10 * time.Minute
)
// ProtocolHandler manages the custom kiro:// protocol handler for OAuth callbacks.
type ProtocolHandler struct {
port int
server *http.Server
listener net.Listener
resultChan chan *AuthCallback
stopChan chan struct{}
mu sync.Mutex
running bool
}
// AuthCallback contains the OAuth callback parameters.
type AuthCallback struct {
Code string
State string
Error string
}
// NewProtocolHandler creates a new protocol handler.
func NewProtocolHandler() *ProtocolHandler {
return &ProtocolHandler{
port: DefaultHandlerPort,
resultChan: make(chan *AuthCallback, 1),
stopChan: make(chan struct{}),
}
}
// Start starts the local callback server that receives redirects from the protocol handler.
func (h *ProtocolHandler) Start(ctx context.Context) (int, error) {
h.mu.Lock()
defer h.mu.Unlock()
if h.running {
return h.port, nil
}
// Drain any stale results from previous runs
select {
case <-h.resultChan:
default:
}
// Reset stopChan for reuse - close old channel first to unblock any waiting goroutines
if h.stopChan != nil {
select {
case <-h.stopChan:
// Already closed
default:
close(h.stopChan)
}
}
h.stopChan = make(chan struct{})
// Try ports in known range (must match handler script port range)
var listener net.Listener
var err error
portRange := []int{DefaultHandlerPort, DefaultHandlerPort + 1, DefaultHandlerPort + 2, DefaultHandlerPort + 3, DefaultHandlerPort + 4}
for _, port := range portRange {
listener, err = net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
if err == nil {
break
}
log.Debugf("kiro protocol handler: port %d busy, trying next", port)
}
if listener == nil {
return 0, fmt.Errorf("failed to start callback server: all ports %d-%d are busy", DefaultHandlerPort, DefaultHandlerPort+4)
}
h.listener = listener
h.port = listener.Addr().(*net.TCPAddr).Port
mux := http.NewServeMux()
mux.HandleFunc("/oauth/callback", h.handleCallback)
h.server = &http.Server{
Handler: mux,
ReadHeaderTimeout: 10 * time.Second,
}
go func() {
if err := h.server.Serve(listener); err != nil && err != http.ErrServerClosed {
log.Debugf("kiro protocol handler server error: %v", err)
}
}()
h.running = true
log.Debugf("kiro protocol handler started on port %d", h.port)
// Auto-shutdown after context done, timeout, or explicit stop
// Capture references to prevent race with new Start() calls
currentStopChan := h.stopChan
currentServer := h.server
currentListener := h.listener
go func() {
select {
case <-ctx.Done():
case <-time.After(HandlerTimeout):
case <-currentStopChan:
return // Already stopped, exit goroutine
}
// Only stop if this is still the current server/listener instance
h.mu.Lock()
if h.server == currentServer && h.listener == currentListener {
h.mu.Unlock()
h.Stop()
} else {
h.mu.Unlock()
}
}()
return h.port, nil
}
// Stop stops the callback server.
func (h *ProtocolHandler) Stop() {
h.mu.Lock()
defer h.mu.Unlock()
if !h.running {
return
}
// Signal the auto-shutdown goroutine to exit.
// This select pattern is safe because stopChan is only modified while holding h.mu,
// and we hold the lock here. The select prevents panic from double-close.
select {
case <-h.stopChan:
// Already closed
default:
close(h.stopChan)
}
if h.server != nil {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_ = h.server.Shutdown(ctx)
}
h.running = false
log.Debug("kiro protocol handler stopped")
}
// WaitForCallback waits for the OAuth callback and returns the result.
func (h *ProtocolHandler) WaitForCallback(ctx context.Context) (*AuthCallback, error) {
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-time.After(HandlerTimeout):
return nil, fmt.Errorf("timeout waiting for OAuth callback")
case result := <-h.resultChan:
return result, nil
}
}
// GetPort returns the port the handler is listening on.
func (h *ProtocolHandler) GetPort() int {
return h.port
}
// handleCallback processes the OAuth callback from the protocol handler script.
func (h *ProtocolHandler) handleCallback(w http.ResponseWriter, r *http.Request) {
code := r.URL.Query().Get("code")
state := r.URL.Query().Get("state")
errParam := r.URL.Query().Get("error")
result := &AuthCallback{
Code: code,
State: state,
Error: errParam,
}
// Send result
select {
case h.resultChan <- result:
default:
// Channel full, ignore duplicate callbacks
}
// Send success response
w.Header().Set("Content-Type", "text/html; charset=utf-8")
if errParam != "" {
w.WriteHeader(http.StatusBadRequest)
fmt.Fprintf(w, `<!DOCTYPE html>
<html>
<head><title>Login Failed</title></head>
<body>
<h1>Login Failed</h1>
<p>Error: %s</p>
<p>You can close this window.</p>
</body>
</html>`, html.EscapeString(errParam))
} else {
fmt.Fprint(w, `<!DOCTYPE html>
<html>
<head><title>Login Successful</title></head>
<body>
<h1>Login Successful!</h1>
<p>You can close this window and return to the terminal.</p>
<script>window.close();</script>
</body>
</html>`)
}
}
// IsProtocolHandlerInstalled checks if the kiro:// protocol handler is installed.
func IsProtocolHandlerInstalled() bool {
switch runtime.GOOS {
case "linux":
return isLinuxHandlerInstalled()
case "windows":
return isWindowsHandlerInstalled()
case "darwin":
return isDarwinHandlerInstalled()
default:
return false
}
}
// InstallProtocolHandler installs the kiro:// protocol handler for the current platform.
func InstallProtocolHandler(handlerPort int) error {
switch runtime.GOOS {
case "linux":
return installLinuxHandler(handlerPort)
case "windows":
return installWindowsHandler(handlerPort)
case "darwin":
return installDarwinHandler(handlerPort)
default:
return fmt.Errorf("unsupported platform: %s", runtime.GOOS)
}
}
// UninstallProtocolHandler removes the kiro:// protocol handler.
func UninstallProtocolHandler() error {
switch runtime.GOOS {
case "linux":
return uninstallLinuxHandler()
case "windows":
return uninstallWindowsHandler()
case "darwin":
return uninstallDarwinHandler()
default:
return fmt.Errorf("unsupported platform: %s", runtime.GOOS)
}
}
// --- Linux Implementation ---
func getLinuxDesktopPath() string {
homeDir, _ := os.UserHomeDir()
return filepath.Join(homeDir, ".local", "share", "applications", "kiro-oauth-handler.desktop")
}
func getLinuxHandlerScriptPath() string {
homeDir, _ := os.UserHomeDir()
return filepath.Join(homeDir, ".local", "bin", "kiro-oauth-handler")
}
func isLinuxHandlerInstalled() bool {
desktopPath := getLinuxDesktopPath()
_, err := os.Stat(desktopPath)
return err == nil
}
func installLinuxHandler(handlerPort int) error {
// Create directories
homeDir, err := os.UserHomeDir()
if err != nil {
return err
}
binDir := filepath.Join(homeDir, ".local", "bin")
appDir := filepath.Join(homeDir, ".local", "share", "applications")
if err := os.MkdirAll(binDir, 0755); err != nil {
return fmt.Errorf("failed to create bin directory: %w", err)
}
if err := os.MkdirAll(appDir, 0755); err != nil {
return fmt.Errorf("failed to create applications directory: %w", err)
}
// Create handler script - tries multiple ports to handle dynamic port allocation
scriptPath := getLinuxHandlerScriptPath()
scriptContent := fmt.Sprintf(`#!/bin/bash
# Kiro OAuth Protocol Handler
# Handles kiro:// URIs - tries CLI first, then forwards to Kiro IDE
URL="$1"
# Check curl availability
if ! command -v curl &> /dev/null; then
echo "Error: curl is required for Kiro OAuth handler" >&2
exit 1
fi
# Extract code and state from URL
[[ "$URL" =~ code=([^&]+) ]] && CODE="${BASH_REMATCH[1]}"
[[ "$URL" =~ state=([^&]+) ]] && STATE="${BASH_REMATCH[1]}"
[[ "$URL" =~ error=([^&]+) ]] && ERROR="${BASH_REMATCH[1]}"
# Try CLI proxy on multiple possible ports (default + dynamic range)
CLI_OK=0
for PORT in %d %d %d %d %d; do
if [ -n "$ERROR" ]; then
curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?error=$ERROR" && CLI_OK=1 && break
elif [ -n "$CODE" ] && [ -n "$STATE" ]; then
curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?code=$CODE&state=$STATE" && CLI_OK=1 && break
fi
done
# If CLI not available, forward to Kiro IDE
if [ $CLI_OK -eq 0 ] && [ -x "/usr/share/kiro/kiro" ]; then
/usr/share/kiro/kiro --open-url "$URL" &
fi
`, handlerPort, handlerPort+1, handlerPort+2, handlerPort+3, handlerPort+4)
if err := os.WriteFile(scriptPath, []byte(scriptContent), 0755); err != nil {
return fmt.Errorf("failed to write handler script: %w", err)
}
// Create .desktop file
desktopPath := getLinuxDesktopPath()
desktopContent := fmt.Sprintf(`[Desktop Entry]
Name=Kiro OAuth Handler
Comment=Handle kiro:// protocol for CLI Proxy API authentication
Exec=%s %%u
Type=Application
Terminal=false
NoDisplay=true
MimeType=x-scheme-handler/kiro;
Categories=Utility;
`, scriptPath)
if err := os.WriteFile(desktopPath, []byte(desktopContent), 0644); err != nil {
return fmt.Errorf("failed to write desktop file: %w", err)
}
// Register handler with xdg-mime
cmd := exec.Command("xdg-mime", "default", "kiro-oauth-handler.desktop", "x-scheme-handler/kiro")
if err := cmd.Run(); err != nil {
log.Warnf("xdg-mime registration failed (may need manual setup): %v", err)
}
// Update desktop database
cmd = exec.Command("update-desktop-database", appDir)
_ = cmd.Run() // Ignore errors, not critical
log.Info("Kiro protocol handler installed for Linux")
return nil
}
func uninstallLinuxHandler() error {
desktopPath := getLinuxDesktopPath()
scriptPath := getLinuxHandlerScriptPath()
if err := os.Remove(desktopPath); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("failed to remove desktop file: %w", err)
}
if err := os.Remove(scriptPath); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("failed to remove handler script: %w", err)
}
log.Info("Kiro protocol handler uninstalled")
return nil
}
// --- Windows Implementation ---
func isWindowsHandlerInstalled() bool {
// Check registry key existence
cmd := exec.Command("reg", "query", `HKCU\Software\Classes\kiro`, "/ve")
return cmd.Run() == nil
}
func installWindowsHandler(handlerPort int) error {
homeDir, err := os.UserHomeDir()
if err != nil {
return err
}
// Create handler script (PowerShell)
scriptDir := filepath.Join(homeDir, ".cliproxyapi")
if err := os.MkdirAll(scriptDir, 0755); err != nil {
return fmt.Errorf("failed to create script directory: %w", err)
}
scriptPath := filepath.Join(scriptDir, "kiro-oauth-handler.ps1")
scriptContent := fmt.Sprintf(`# Kiro OAuth Protocol Handler for Windows
param([string]$url)
# Load required assembly for HttpUtility
Add-Type -AssemblyName System.Web
# Parse URL parameters
$uri = [System.Uri]$url
$query = [System.Web.HttpUtility]::ParseQueryString($uri.Query)
$code = $query["code"]
$state = $query["state"]
$errorParam = $query["error"]
# Try multiple ports (default + dynamic range)
$ports = @(%d, %d, %d, %d, %d)
$success = $false
foreach ($port in $ports) {
if ($success) { break }
$callbackUrl = "http://127.0.0.1:$port/oauth/callback"
try {
if ($errorParam) {
$fullUrl = $callbackUrl + "?error=" + $errorParam
Invoke-WebRequest -Uri $fullUrl -UseBasicParsing -TimeoutSec 1 -ErrorAction Stop | Out-Null
$success = $true
} elseif ($code -and $state) {
$fullUrl = $callbackUrl + "?code=" + $code + "&state=" + $state
Invoke-WebRequest -Uri $fullUrl -UseBasicParsing -TimeoutSec 1 -ErrorAction Stop | Out-Null
$success = $true
}
} catch {
# Try next port
}
}
`, handlerPort, handlerPort+1, handlerPort+2, handlerPort+3, handlerPort+4)
if err := os.WriteFile(scriptPath, []byte(scriptContent), 0644); err != nil {
return fmt.Errorf("failed to write handler script: %w", err)
}
// Create batch wrapper
batchPath := filepath.Join(scriptDir, "kiro-oauth-handler.bat")
batchContent := fmt.Sprintf("@echo off\npowershell -ExecutionPolicy Bypass -File \"%s\" \"%%1\"\n", scriptPath)
if err := os.WriteFile(batchPath, []byte(batchContent), 0644); err != nil {
return fmt.Errorf("failed to write batch wrapper: %w", err)
}
// Register in Windows registry
commands := [][]string{
{"reg", "add", `HKCU\Software\Classes\kiro`, "/ve", "/d", "URL:Kiro Protocol", "/f"},
{"reg", "add", `HKCU\Software\Classes\kiro`, "/v", "URL Protocol", "/d", "", "/f"},
{"reg", "add", `HKCU\Software\Classes\kiro\shell`, "/f"},
{"reg", "add", `HKCU\Software\Classes\kiro\shell\open`, "/f"},
{"reg", "add", `HKCU\Software\Classes\kiro\shell\open\command`, "/ve", "/d", fmt.Sprintf("\"%s\" \"%%1\"", batchPath), "/f"},
}
for _, args := range commands {
cmd := exec.Command(args[0], args[1:]...)
if err := cmd.Run(); err != nil {
return fmt.Errorf("failed to run registry command: %w", err)
}
}
log.Info("Kiro protocol handler installed for Windows")
return nil
}
func uninstallWindowsHandler() error {
// Remove registry keys
cmd := exec.Command("reg", "delete", `HKCU\Software\Classes\kiro`, "/f")
if err := cmd.Run(); err != nil {
log.Warnf("failed to remove registry key: %v", err)
}
// Remove scripts
homeDir, _ := os.UserHomeDir()
scriptDir := filepath.Join(homeDir, ".cliproxyapi")
_ = os.Remove(filepath.Join(scriptDir, "kiro-oauth-handler.ps1"))
_ = os.Remove(filepath.Join(scriptDir, "kiro-oauth-handler.bat"))
log.Info("Kiro protocol handler uninstalled")
return nil
}
// --- macOS Implementation ---
func getDarwinAppPath() string {
homeDir, _ := os.UserHomeDir()
return filepath.Join(homeDir, "Applications", "KiroOAuthHandler.app")
}
func isDarwinHandlerInstalled() bool {
appPath := getDarwinAppPath()
_, err := os.Stat(appPath)
return err == nil
}
func installDarwinHandler(handlerPort int) error {
// Create app bundle structure
appPath := getDarwinAppPath()
contentsPath := filepath.Join(appPath, "Contents")
macOSPath := filepath.Join(contentsPath, "MacOS")
if err := os.MkdirAll(macOSPath, 0755); err != nil {
return fmt.Errorf("failed to create app bundle: %w", err)
}
// Create Info.plist
plistPath := filepath.Join(contentsPath, "Info.plist")
plistContent := `<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
<key>CFBundleIdentifier</key>
<string>com.cliproxyapi.kiro-oauth-handler</string>
<key>CFBundleName</key>
<string>KiroOAuthHandler</string>
<key>CFBundleExecutable</key>
<string>kiro-oauth-handler</string>
<key>CFBundleVersion</key>
<string>1.0</string>
<key>CFBundleURLTypes</key>
<array>
<dict>
<key>CFBundleURLName</key>
<string>Kiro Protocol</string>
<key>CFBundleURLSchemes</key>
<array>
<string>kiro</string>
</array>
</dict>
</array>
<key>LSBackgroundOnly</key>
<true/>
</dict>
</plist>`
if err := os.WriteFile(plistPath, []byte(plistContent), 0644); err != nil {
return fmt.Errorf("failed to write Info.plist: %w", err)
}
// Create executable script - tries multiple ports to handle dynamic port allocation
execPath := filepath.Join(macOSPath, "kiro-oauth-handler")
execContent := fmt.Sprintf(`#!/bin/bash
# Kiro OAuth Protocol Handler for macOS
URL="$1"
# Check curl availability (should always exist on macOS)
if [ ! -x /usr/bin/curl ]; then
echo "Error: curl is required for Kiro OAuth handler" >&2
exit 1
fi
# Extract code and state from URL
[[ "$URL" =~ code=([^&]+) ]] && CODE="${BASH_REMATCH[1]}"
[[ "$URL" =~ state=([^&]+) ]] && STATE="${BASH_REMATCH[1]}"
[[ "$URL" =~ error=([^&]+) ]] && ERROR="${BASH_REMATCH[1]}"
# Try multiple ports (default + dynamic range)
for PORT in %d %d %d %d %d; do
if [ -n "$ERROR" ]; then
/usr/bin/curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?error=$ERROR" && exit 0
elif [ -n "$CODE" ] && [ -n "$STATE" ]; then
/usr/bin/curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?code=$CODE&state=$STATE" && exit 0
fi
done
`, handlerPort, handlerPort+1, handlerPort+2, handlerPort+3, handlerPort+4)
if err := os.WriteFile(execPath, []byte(execContent), 0755); err != nil {
return fmt.Errorf("failed to write executable: %w", err)
}
// Register the app with Launch Services
cmd := exec.Command("/System/Library/Frameworks/CoreServices.framework/Frameworks/LaunchServices.framework/Support/lsregister",
"-f", appPath)
if err := cmd.Run(); err != nil {
log.Warnf("lsregister failed (handler may still work): %v", err)
}
log.Info("Kiro protocol handler installed for macOS")
return nil
}
func uninstallDarwinHandler() error {
appPath := getDarwinAppPath()
// Unregister from Launch Services
cmd := exec.Command("/System/Library/Frameworks/CoreServices.framework/Frameworks/LaunchServices.framework/Support/lsregister",
"-u", appPath)
_ = cmd.Run()
// Remove app bundle
if err := os.RemoveAll(appPath); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("failed to remove app bundle: %w", err)
}
log.Info("Kiro protocol handler uninstalled")
return nil
}
// ParseKiroURI parses a kiro:// URI and extracts the callback parameters.
func ParseKiroURI(rawURI string) (*AuthCallback, error) {
u, err := url.Parse(rawURI)
if err != nil {
return nil, fmt.Errorf("invalid URI: %w", err)
}
if u.Scheme != KiroProtocol {
return nil, fmt.Errorf("invalid scheme: expected %s, got %s", KiroProtocol, u.Scheme)
}
if u.Host != KiroAuthority {
return nil, fmt.Errorf("invalid authority: expected %s, got %s", KiroAuthority, u.Host)
}
query := u.Query()
return &AuthCallback{
Code: query.Get("code"),
State: query.Get("state"),
Error: query.Get("error"),
}, nil
}
// GetHandlerInstructions returns platform-specific instructions for manual handler setup.
func GetHandlerInstructions() string {
switch runtime.GOOS {
case "linux":
return `To manually set up the Kiro protocol handler on Linux:
1. Create ~/.local/share/applications/kiro-oauth-handler.desktop:
[Desktop Entry]
Name=Kiro OAuth Handler
Exec=~/.local/bin/kiro-oauth-handler %u
Type=Application
Terminal=false
MimeType=x-scheme-handler/kiro;
2. Create ~/.local/bin/kiro-oauth-handler (make it executable):
#!/bin/bash
URL="$1"
# ... (see generated script for full content)
3. Run: xdg-mime default kiro-oauth-handler.desktop x-scheme-handler/kiro`
case "windows":
return `To manually set up the Kiro protocol handler on Windows:
1. Open Registry Editor (regedit.exe)
2. Create key: HKEY_CURRENT_USER\Software\Classes\kiro
3. Set default value to: URL:Kiro Protocol
4. Create string value "URL Protocol" with empty data
5. Create subkey: shell\open\command
6. Set default value to: "C:\path\to\handler.bat" "%1"`
case "darwin":
return `To manually set up the Kiro protocol handler on macOS:
1. Create ~/Applications/KiroOAuthHandler.app bundle
2. Add Info.plist with CFBundleURLTypes containing "kiro" scheme
3. Create executable in Contents/MacOS/
4. Run: /System/Library/.../lsregister -f ~/Applications/KiroOAuthHandler.app`
default:
return "Protocol handler setup is not supported on this platform."
}
}
// SetupProtocolHandlerIfNeeded checks and installs the protocol handler if needed.
func SetupProtocolHandlerIfNeeded(handlerPort int) error {
if IsProtocolHandlerInstalled() {
log.Debug("Kiro protocol handler already installed")
return nil
}
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
fmt.Println("║ Kiro Protocol Handler Setup Required ║")
fmt.Println("╚══════════════════════════════════════════════════════════╝")
fmt.Println("\nTo enable Google/GitHub login, we need to install a protocol handler.")
fmt.Println("This allows your browser to redirect back to the CLI after authentication.")
fmt.Println("\nInstalling protocol handler...")
if err := InstallProtocolHandler(handlerPort); err != nil {
fmt.Printf("\n⚠ Automatic installation failed: %v\n", err)
fmt.Println("\nManual setup instructions:")
fmt.Println(strings.Repeat("-", 60))
fmt.Println(GetHandlerInstructions())
return err
}
fmt.Println("\n✓ Protocol handler installed successfully!")
return nil
}

View File

@@ -0,0 +1,403 @@
// Package kiro provides social authentication (Google/GitHub) for Kiro via AuthServiceClient.
package kiro
import (
"bufio"
"context"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"os"
"os/exec"
"runtime"
"strings"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
log "github.com/sirupsen/logrus"
"golang.org/x/term"
)
const (
// Kiro AuthService endpoint
kiroAuthServiceEndpoint = "https://prod.us-east-1.auth.desktop.kiro.dev"
// OAuth timeout
socialAuthTimeout = 10 * time.Minute
)
// SocialProvider represents the social login provider.
type SocialProvider string
const (
// ProviderGoogle is Google OAuth provider
ProviderGoogle SocialProvider = "Google"
// ProviderGitHub is GitHub OAuth provider
ProviderGitHub SocialProvider = "Github"
// Note: AWS Builder ID is NOT supported by Kiro's auth service.
// It only supports: Google, Github, Cognito
// AWS Builder ID must use device code flow via SSO OIDC.
)
// CreateTokenRequest is sent to Kiro's /oauth/token endpoint.
type CreateTokenRequest struct {
Code string `json:"code"`
CodeVerifier string `json:"code_verifier"`
RedirectURI string `json:"redirect_uri"`
InvitationCode string `json:"invitation_code,omitempty"`
}
// SocialTokenResponse from Kiro's /oauth/token endpoint for social auth.
type SocialTokenResponse struct {
AccessToken string `json:"accessToken"`
RefreshToken string `json:"refreshToken"`
ProfileArn string `json:"profileArn"`
ExpiresIn int `json:"expiresIn"`
}
// RefreshTokenRequest is sent to Kiro's /refreshToken endpoint.
type RefreshTokenRequest struct {
RefreshToken string `json:"refreshToken"`
}
// SocialAuthClient handles social authentication with Kiro.
type SocialAuthClient struct {
httpClient *http.Client
cfg *config.Config
protocolHandler *ProtocolHandler
}
// NewSocialAuthClient creates a new social auth client.
func NewSocialAuthClient(cfg *config.Config) *SocialAuthClient {
client := &http.Client{Timeout: 30 * time.Second}
if cfg != nil {
client = util.SetProxy(&cfg.SDKConfig, client)
}
return &SocialAuthClient{
httpClient: client,
cfg: cfg,
protocolHandler: NewProtocolHandler(),
}
}
// generatePKCE generates PKCE code verifier and challenge.
func generatePKCE() (verifier, challenge string, err error) {
// Generate 32 bytes of random data for verifier
b := make([]byte, 32)
if _, err := rand.Read(b); err != nil {
return "", "", fmt.Errorf("failed to generate random bytes: %w", err)
}
verifier = base64.RawURLEncoding.EncodeToString(b)
// Generate SHA256 hash of verifier for challenge
h := sha256.Sum256([]byte(verifier))
challenge = base64.RawURLEncoding.EncodeToString(h[:])
return verifier, challenge, nil
}
// generateState generates a random state parameter.
func generateStateParam() (string, error) {
b := make([]byte, 16)
if _, err := rand.Read(b); err != nil {
return "", err
}
return base64.RawURLEncoding.EncodeToString(b), nil
}
// buildLoginURL constructs the Kiro OAuth login URL.
// The login endpoint expects a GET request with query parameters.
// Format: /login?idp=Google&redirect_uri=...&code_challenge=...&code_challenge_method=S256&state=...&prompt=select_account
// The prompt=select_account parameter forces the account selection screen even if already logged in.
func (c *SocialAuthClient) buildLoginURL(provider, redirectURI, codeChallenge, state string) string {
return fmt.Sprintf("%s/login?idp=%s&redirect_uri=%s&code_challenge=%s&code_challenge_method=S256&state=%s&prompt=select_account",
kiroAuthServiceEndpoint,
provider,
url.QueryEscape(redirectURI),
codeChallenge,
state,
)
}
// createToken exchanges the authorization code for tokens.
func (c *SocialAuthClient) createToken(ctx context.Context, req *CreateTokenRequest) (*SocialTokenResponse, error) {
body, err := json.Marshal(req)
if err != nil {
return nil, fmt.Errorf("failed to marshal token request: %w", err)
}
tokenURL := kiroAuthServiceEndpoint + "/oauth/token"
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(string(body)))
if err != nil {
return nil, fmt.Errorf("failed to create token request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
resp, err := c.httpClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("token request failed: %w", err)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read token response: %w", err)
}
if resp.StatusCode != http.StatusOK {
log.Debugf("token exchange failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("token exchange failed (status %d)", resp.StatusCode)
}
var tokenResp SocialTokenResponse
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
return nil, fmt.Errorf("failed to parse token response: %w", err)
}
return &tokenResp, nil
}
// RefreshSocialToken refreshes an expired social auth token.
func (c *SocialAuthClient) RefreshSocialToken(ctx context.Context, refreshToken string) (*KiroTokenData, error) {
body, err := json.Marshal(&RefreshTokenRequest{RefreshToken: refreshToken})
if err != nil {
return nil, fmt.Errorf("failed to marshal refresh request: %w", err)
}
refreshURL := kiroAuthServiceEndpoint + "/refreshToken"
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, refreshURL, strings.NewReader(string(body)))
if err != nil {
return nil, fmt.Errorf("failed to create refresh request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
resp, err := c.httpClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("refresh request failed: %w", err)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read refresh response: %w", err)
}
if resp.StatusCode != http.StatusOK {
log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
}
var tokenResp SocialTokenResponse
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
return nil, fmt.Errorf("failed to parse refresh response: %w", err)
}
// Validate ExpiresIn - use default 1 hour if invalid
expiresIn := tokenResp.ExpiresIn
if expiresIn <= 0 {
expiresIn = 3600 // Default 1 hour
}
expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
return &KiroTokenData{
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
ProfileArn: tokenResp.ProfileArn,
ExpiresAt: expiresAt.Format(time.RFC3339),
AuthMethod: "social",
Provider: "", // Caller should preserve original provider
}, nil
}
// LoginWithSocial performs OAuth login with Google.
func (c *SocialAuthClient) LoginWithSocial(ctx context.Context, provider SocialProvider) (*KiroTokenData, error) {
providerName := string(provider)
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
fmt.Printf("║ Kiro Authentication (%s) ║\n", providerName)
fmt.Println("╚══════════════════════════════════════════════════════════╝")
// Step 1: Setup protocol handler
fmt.Println("\nSetting up authentication...")
// Start the local callback server
handlerPort, err := c.protocolHandler.Start(ctx)
if err != nil {
return nil, fmt.Errorf("failed to start callback server: %w", err)
}
defer c.protocolHandler.Stop()
// Ensure protocol handler is installed and set as default
if err := SetupProtocolHandlerIfNeeded(handlerPort); err != nil {
fmt.Println("\n⚠ Protocol handler setup failed. Trying alternative method...")
fmt.Println(" If you see a browser 'Open with' dialog, select your default browser.")
fmt.Println(" For manual setup instructions, run: cliproxy kiro --help-protocol")
log.Debugf("kiro: protocol handler setup error: %v", err)
// Continue anyway - user might have set it up manually or select browser manually
} else {
// Force set our handler as default (prevents "Open with" dialog)
forceDefaultProtocolHandler()
}
// Step 2: Generate PKCE codes
codeVerifier, codeChallenge, err := generatePKCE()
if err != nil {
return nil, fmt.Errorf("failed to generate PKCE: %w", err)
}
// Step 3: Generate state
state, err := generateStateParam()
if err != nil {
return nil, fmt.Errorf("failed to generate state: %w", err)
}
// Step 4: Build the login URL (Kiro uses GET request with query params)
authURL := c.buildLoginURL(providerName, KiroRedirectURI, codeChallenge, state)
// Set incognito mode based on config (defaults to true for Kiro, can be overridden with --no-incognito)
// Incognito mode enables multi-account support by bypassing cached sessions
if c.cfg != nil {
browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
if !c.cfg.IncognitoBrowser {
log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.")
} else {
log.Debug("kiro: using incognito mode for multi-account support")
}
} else {
browser.SetIncognitoMode(true) // Default to incognito if no config
log.Debug("kiro: using incognito mode for multi-account support (default)")
}
// Step 5: Open browser for user authentication
fmt.Println("\n════════════════════════════════════════════════════════════")
fmt.Printf(" Opening browser for %s authentication...\n", providerName)
fmt.Println("════════════════════════════════════════════════════════════")
fmt.Printf("\n URL: %s\n\n", authURL)
if err := browser.OpenURL(authURL); err != nil {
log.Warnf("Could not open browser automatically: %v", err)
fmt.Println(" ⚠ Could not open browser automatically.")
fmt.Println(" Please open the URL above in your browser manually.")
} else {
fmt.Println(" (Browser opened automatically)")
}
fmt.Println("\n Waiting for authentication callback...")
// Step 6: Wait for callback
callback, err := c.protocolHandler.WaitForCallback(ctx)
if err != nil {
return nil, fmt.Errorf("failed to receive callback: %w", err)
}
if callback.Error != "" {
return nil, fmt.Errorf("authentication error: %s", callback.Error)
}
if callback.State != state {
// Log state values for debugging, but don't expose in user-facing error
log.Debugf("kiro: OAuth state mismatch - expected %s, got %s", state, callback.State)
return nil, fmt.Errorf("OAuth state validation failed - please try again")
}
if callback.Code == "" {
return nil, fmt.Errorf("no authorization code received")
}
fmt.Println("\n✓ Authorization received!")
// Step 7: Exchange code for tokens
fmt.Println("Exchanging code for tokens...")
tokenReq := &CreateTokenRequest{
Code: callback.Code,
CodeVerifier: codeVerifier,
RedirectURI: KiroRedirectURI,
}
tokenResp, err := c.createToken(ctx, tokenReq)
if err != nil {
return nil, fmt.Errorf("failed to exchange code for tokens: %w", err)
}
fmt.Println("\n✓ Authentication successful!")
// Close the browser window
if err := browser.CloseBrowser(); err != nil {
log.Debugf("Failed to close browser: %v", err)
}
// Validate ExpiresIn - use default 1 hour if invalid
expiresIn := tokenResp.ExpiresIn
if expiresIn <= 0 {
expiresIn = 3600
}
expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
// Try to extract email from JWT access token first
email := ExtractEmailFromJWT(tokenResp.AccessToken)
// If no email in JWT, ask user for account label (only in interactive mode)
if email == "" && isInteractiveTerminal() {
fmt.Print("\n Enter account label for file naming (optional, press Enter to skip): ")
reader := bufio.NewReader(os.Stdin)
var err error
email, err = reader.ReadString('\n')
if err != nil {
log.Debugf("Failed to read account label: %v", err)
}
email = strings.TrimSpace(email)
}
return &KiroTokenData{
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
ProfileArn: tokenResp.ProfileArn,
ExpiresAt: expiresAt.Format(time.RFC3339),
AuthMethod: "social",
Provider: providerName,
Email: email, // JWT email or user-provided label
}, nil
}
// LoginWithGoogle performs OAuth login with Google.
func (c *SocialAuthClient) LoginWithGoogle(ctx context.Context) (*KiroTokenData, error) {
return c.LoginWithSocial(ctx, ProviderGoogle)
}
// LoginWithGitHub performs OAuth login with GitHub.
func (c *SocialAuthClient) LoginWithGitHub(ctx context.Context) (*KiroTokenData, error) {
return c.LoginWithSocial(ctx, ProviderGitHub)
}
// forceDefaultProtocolHandler sets our protocol handler as the default for kiro:// URLs.
// This prevents the "Open with" dialog from appearing on Linux.
// On non-Linux platforms, this is a no-op as they use different mechanisms.
func forceDefaultProtocolHandler() {
if runtime.GOOS != "linux" {
return // Non-Linux platforms use different handler mechanisms
}
// Set our handler as default using xdg-mime
cmd := exec.Command("xdg-mime", "default", "kiro-oauth-handler.desktop", "x-scheme-handler/kiro")
if err := cmd.Run(); err != nil {
log.Warnf("Failed to set default protocol handler: %v. You may see a handler selection dialog.", err)
}
}
// isInteractiveTerminal checks if stdin is connected to an interactive terminal.
// Returns false in CI/automated environments or when stdin is piped.
func isInteractiveTerminal() bool {
return term.IsTerminal(int(os.Stdin.Fd()))
}

View File

@@ -0,0 +1,527 @@
// Package kiro provides AWS SSO OIDC authentication for Kiro.
package kiro
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
log "github.com/sirupsen/logrus"
)
const (
// AWS SSO OIDC endpoints
ssoOIDCEndpoint = "https://oidc.us-east-1.amazonaws.com"
// Kiro's start URL for Builder ID
builderIDStartURL = "https://view.awsapps.com/start"
// Polling interval
pollInterval = 5 * time.Second
)
// SSOOIDCClient handles AWS SSO OIDC authentication.
type SSOOIDCClient struct {
httpClient *http.Client
cfg *config.Config
}
// NewSSOOIDCClient creates a new SSO OIDC client.
func NewSSOOIDCClient(cfg *config.Config) *SSOOIDCClient {
client := &http.Client{Timeout: 30 * time.Second}
if cfg != nil {
client = util.SetProxy(&cfg.SDKConfig, client)
}
return &SSOOIDCClient{
httpClient: client,
cfg: cfg,
}
}
// RegisterClientResponse from AWS SSO OIDC.
type RegisterClientResponse struct {
ClientID string `json:"clientId"`
ClientSecret string `json:"clientSecret"`
ClientIDIssuedAt int64 `json:"clientIdIssuedAt"`
ClientSecretExpiresAt int64 `json:"clientSecretExpiresAt"`
}
// StartDeviceAuthResponse from AWS SSO OIDC.
type StartDeviceAuthResponse struct {
DeviceCode string `json:"deviceCode"`
UserCode string `json:"userCode"`
VerificationURI string `json:"verificationUri"`
VerificationURIComplete string `json:"verificationUriComplete"`
ExpiresIn int `json:"expiresIn"`
Interval int `json:"interval"`
}
// CreateTokenResponse from AWS SSO OIDC.
type CreateTokenResponse struct {
AccessToken string `json:"accessToken"`
TokenType string `json:"tokenType"`
ExpiresIn int `json:"expiresIn"`
RefreshToken string `json:"refreshToken"`
}
// RegisterClient registers a new OIDC client with AWS.
func (c *SSOOIDCClient) RegisterClient(ctx context.Context) (*RegisterClientResponse, error) {
// Generate unique client name for each registration to support multiple accounts
clientName := fmt.Sprintf("CLI-Proxy-API-%d", time.Now().UnixNano())
payload := map[string]interface{}{
"clientName": clientName,
"clientType": "public",
"scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations"},
}
body, err := json.Marshal(payload)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/client/register", strings.NewReader(string(body)))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
log.Debugf("register client failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode)
}
var result RegisterClientResponse
if err := json.Unmarshal(respBody, &result); err != nil {
return nil, err
}
return &result, nil
}
// StartDeviceAuthorization starts the device authorization flow.
func (c *SSOOIDCClient) StartDeviceAuthorization(ctx context.Context, clientID, clientSecret string) (*StartDeviceAuthResponse, error) {
payload := map[string]string{
"clientId": clientID,
"clientSecret": clientSecret,
"startUrl": builderIDStartURL,
}
body, err := json.Marshal(payload)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/device_authorization", strings.NewReader(string(body)))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
log.Debugf("start device auth failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("start device auth failed (status %d)", resp.StatusCode)
}
var result StartDeviceAuthResponse
if err := json.Unmarshal(respBody, &result); err != nil {
return nil, err
}
return &result, nil
}
// CreateToken polls for the access token after user authorization.
func (c *SSOOIDCClient) CreateToken(ctx context.Context, clientID, clientSecret, deviceCode string) (*CreateTokenResponse, error) {
payload := map[string]string{
"clientId": clientID,
"clientSecret": clientSecret,
"deviceCode": deviceCode,
"grantType": "urn:ietf:params:oauth:grant-type:device_code",
}
body, err := json.Marshal(payload)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body)))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
// Check for pending authorization
if resp.StatusCode == http.StatusBadRequest {
var errResp struct {
Error string `json:"error"`
}
if json.Unmarshal(respBody, &errResp) == nil {
if errResp.Error == "authorization_pending" {
return nil, fmt.Errorf("authorization_pending")
}
if errResp.Error == "slow_down" {
return nil, fmt.Errorf("slow_down")
}
}
log.Debugf("create token failed: %s", string(respBody))
return nil, fmt.Errorf("create token failed")
}
if resp.StatusCode != http.StatusOK {
log.Debugf("create token failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode)
}
var result CreateTokenResponse
if err := json.Unmarshal(respBody, &result); err != nil {
return nil, err
}
return &result, nil
}
// RefreshToken refreshes an access token using the refresh token.
func (c *SSOOIDCClient) RefreshToken(ctx context.Context, clientID, clientSecret, refreshToken string) (*KiroTokenData, error) {
payload := map[string]string{
"clientId": clientID,
"clientSecret": clientSecret,
"refreshToken": refreshToken,
"grantType": "refresh_token",
}
body, err := json.Marshal(payload)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body)))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
}
var result CreateTokenResponse
if err := json.Unmarshal(respBody, &result); err != nil {
return nil, err
}
expiresAt := time.Now().Add(time.Duration(result.ExpiresIn) * time.Second)
return &KiroTokenData{
AccessToken: result.AccessToken,
RefreshToken: result.RefreshToken,
ExpiresAt: expiresAt.Format(time.RFC3339),
AuthMethod: "builder-id",
Provider: "AWS",
ClientID: clientID,
ClientSecret: clientSecret,
}, nil
}
// LoginWithBuilderID performs the full device code flow for AWS Builder ID.
func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, error) {
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
fmt.Println("║ Kiro Authentication (AWS Builder ID) ║")
fmt.Println("╚══════════════════════════════════════════════════════════╝")
// Step 1: Register client
fmt.Println("\nRegistering client...")
regResp, err := c.RegisterClient(ctx)
if err != nil {
return nil, fmt.Errorf("failed to register client: %w", err)
}
log.Debugf("Client registered: %s", regResp.ClientID)
// Step 2: Start device authorization
fmt.Println("Starting device authorization...")
authResp, err := c.StartDeviceAuthorization(ctx, regResp.ClientID, regResp.ClientSecret)
if err != nil {
return nil, fmt.Errorf("failed to start device auth: %w", err)
}
// Step 3: Show user the verification URL
fmt.Printf("\n")
fmt.Println("════════════════════════════════════════════════════════════")
fmt.Printf(" Open this URL in your browser:\n")
fmt.Printf(" %s\n", authResp.VerificationURIComplete)
fmt.Println("════════════════════════════════════════════════════════════")
fmt.Printf("\n Or go to: %s\n", authResp.VerificationURI)
fmt.Printf(" And enter code: %s\n\n", authResp.UserCode)
// Set incognito mode based on config (defaults to true for Kiro, can be overridden with --no-incognito)
// Incognito mode enables multi-account support by bypassing cached sessions
if c.cfg != nil {
browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
if !c.cfg.IncognitoBrowser {
log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.")
} else {
log.Debug("kiro: using incognito mode for multi-account support")
}
} else {
browser.SetIncognitoMode(true) // Default to incognito if no config
log.Debug("kiro: using incognito mode for multi-account support (default)")
}
// Open browser using cross-platform browser package
if err := browser.OpenURL(authResp.VerificationURIComplete); err != nil {
log.Warnf("Could not open browser automatically: %v", err)
fmt.Println(" Please open the URL manually in your browser.")
} else {
fmt.Println(" (Browser opened automatically)")
}
// Step 4: Poll for token
fmt.Println("Waiting for authorization...")
interval := pollInterval
if authResp.Interval > 0 {
interval = time.Duration(authResp.Interval) * time.Second
}
deadline := time.Now().Add(time.Duration(authResp.ExpiresIn) * time.Second)
for time.Now().Before(deadline) {
select {
case <-ctx.Done():
browser.CloseBrowser() // Cleanup on cancel
return nil, ctx.Err()
case <-time.After(interval):
tokenResp, err := c.CreateToken(ctx, regResp.ClientID, regResp.ClientSecret, authResp.DeviceCode)
if err != nil {
errStr := err.Error()
if strings.Contains(errStr, "authorization_pending") {
fmt.Print(".")
continue
}
if strings.Contains(errStr, "slow_down") {
interval += 5 * time.Second
continue
}
// Close browser on error before returning
browser.CloseBrowser()
return nil, fmt.Errorf("token creation failed: %w", err)
}
fmt.Println("\n\n✓ Authorization successful!")
// Close the browser window
if err := browser.CloseBrowser(); err != nil {
log.Debugf("Failed to close browser: %v", err)
}
// Step 5: Get profile ARN from CodeWhisperer API
fmt.Println("Fetching profile information...")
profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken)
// Extract email from JWT access token
email := ExtractEmailFromJWT(tokenResp.AccessToken)
if email != "" {
fmt.Printf(" Logged in as: %s\n", email)
}
expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
return &KiroTokenData{
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
ProfileArn: profileArn,
ExpiresAt: expiresAt.Format(time.RFC3339),
AuthMethod: "builder-id",
Provider: "AWS",
ClientID: regResp.ClientID,
ClientSecret: regResp.ClientSecret,
Email: email,
}, nil
}
}
// Close browser on timeout for better UX
if err := browser.CloseBrowser(); err != nil {
log.Debugf("Failed to close browser on timeout: %v", err)
}
return nil, fmt.Errorf("authorization timed out")
}
// fetchProfileArn retrieves the profile ARN from CodeWhisperer API.
// This is needed for file naming since AWS SSO OIDC doesn't return profile info.
func (c *SSOOIDCClient) fetchProfileArn(ctx context.Context, accessToken string) string {
// Try ListProfiles API first
profileArn := c.tryListProfiles(ctx, accessToken)
if profileArn != "" {
return profileArn
}
// Fallback: Try ListAvailableCustomizations
return c.tryListCustomizations(ctx, accessToken)
}
func (c *SSOOIDCClient) tryListProfiles(ctx context.Context, accessToken string) string {
payload := map[string]interface{}{
"origin": "AI_EDITOR",
}
body, err := json.Marshal(payload)
if err != nil {
return ""
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body)))
if err != nil {
return ""
}
req.Header.Set("Content-Type", "application/x-amz-json-1.0")
req.Header.Set("x-amz-target", "AmazonCodeWhispererService.ListProfiles")
req.Header.Set("Authorization", "Bearer "+accessToken)
req.Header.Set("Accept", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return ""
}
defer resp.Body.Close()
respBody, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK {
log.Debugf("ListProfiles failed (status %d): %s", resp.StatusCode, string(respBody))
return ""
}
log.Debugf("ListProfiles response: %s", string(respBody))
var result struct {
Profiles []struct {
Arn string `json:"arn"`
} `json:"profiles"`
ProfileArn string `json:"profileArn"`
}
if err := json.Unmarshal(respBody, &result); err != nil {
return ""
}
if result.ProfileArn != "" {
return result.ProfileArn
}
if len(result.Profiles) > 0 {
return result.Profiles[0].Arn
}
return ""
}
func (c *SSOOIDCClient) tryListCustomizations(ctx context.Context, accessToken string) string {
payload := map[string]interface{}{
"origin": "AI_EDITOR",
}
body, err := json.Marshal(payload)
if err != nil {
return ""
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body)))
if err != nil {
return ""
}
req.Header.Set("Content-Type", "application/x-amz-json-1.0")
req.Header.Set("x-amz-target", "AmazonCodeWhispererService.ListAvailableCustomizations")
req.Header.Set("Authorization", "Bearer "+accessToken)
req.Header.Set("Accept", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return ""
}
defer resp.Body.Close()
respBody, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK {
log.Debugf("ListAvailableCustomizations failed (status %d): %s", resp.StatusCode, string(respBody))
return ""
}
log.Debugf("ListAvailableCustomizations response: %s", string(respBody))
var result struct {
Customizations []struct {
Arn string `json:"arn"`
} `json:"customizations"`
ProfileArn string `json:"profileArn"`
}
if err := json.Unmarshal(respBody, &result); err != nil {
return ""
}
if result.ProfileArn != "" {
return result.ProfileArn
}
if len(result.Customizations) > 0 {
return result.Customizations[0].Arn
}
return ""
}

View File

@@ -0,0 +1,72 @@
package kiro
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
)
// KiroTokenStorage holds the persistent token data for Kiro authentication.
type KiroTokenStorage struct {
// AccessToken is the OAuth2 access token for API access
AccessToken string `json:"access_token"`
// RefreshToken is used to obtain new access tokens
RefreshToken string `json:"refresh_token"`
// ProfileArn is the AWS CodeWhisperer profile ARN
ProfileArn string `json:"profile_arn"`
// ExpiresAt is the timestamp when the token expires
ExpiresAt string `json:"expires_at"`
// AuthMethod indicates the authentication method used
AuthMethod string `json:"auth_method"`
// Provider indicates the OAuth provider
Provider string `json:"provider"`
// LastRefresh is the timestamp of the last token refresh
LastRefresh string `json:"last_refresh"`
}
// SaveTokenToFile persists the token storage to the specified file path.
func (s *KiroTokenStorage) SaveTokenToFile(authFilePath string) error {
dir := filepath.Dir(authFilePath)
if err := os.MkdirAll(dir, 0700); err != nil {
return fmt.Errorf("failed to create directory: %w", err)
}
data, err := json.MarshalIndent(s, "", " ")
if err != nil {
return fmt.Errorf("failed to marshal token storage: %w", err)
}
if err := os.WriteFile(authFilePath, data, 0600); err != nil {
return fmt.Errorf("failed to write token file: %w", err)
}
return nil
}
// LoadFromFile loads token storage from the specified file path.
func LoadFromFile(authFilePath string) (*KiroTokenStorage, error) {
data, err := os.ReadFile(authFilePath)
if err != nil {
return nil, fmt.Errorf("failed to read token file: %w", err)
}
var storage KiroTokenStorage
if err := json.Unmarshal(data, &storage); err != nil {
return nil, fmt.Errorf("failed to parse token file: %w", err)
}
return &storage, nil
}
// ToTokenData converts storage to KiroTokenData for API use.
func (s *KiroTokenStorage) ToTokenData() *KiroTokenData {
return &KiroTokenData{
AccessToken: s.AccessToken,
RefreshToken: s.RefreshToken,
ProfileArn: s.ProfileArn,
ExpiresAt: s.ExpiresAt,
AuthMethod: s.AuthMethod,
Provider: s.Provider,
}
}

View File

@@ -6,14 +6,49 @@ import (
"fmt"
"os/exec"
"runtime"
"strings"
"sync"
pkgbrowser "github.com/pkg/browser"
log "github.com/sirupsen/logrus"
"github.com/skratchdot/open-golang/open"
)
// incognitoMode controls whether to open URLs in incognito/private mode.
// This is useful for OAuth flows where you want to use a different account.
var incognitoMode bool
// lastBrowserProcess stores the last opened browser process for cleanup
var lastBrowserProcess *exec.Cmd
var browserMutex sync.Mutex
// SetIncognitoMode enables or disables incognito/private browsing mode.
func SetIncognitoMode(enabled bool) {
incognitoMode = enabled
}
// IsIncognitoMode returns whether incognito mode is enabled.
func IsIncognitoMode() bool {
return incognitoMode
}
// CloseBrowser closes the last opened browser process.
func CloseBrowser() error {
browserMutex.Lock()
defer browserMutex.Unlock()
if lastBrowserProcess == nil || lastBrowserProcess.Process == nil {
return nil
}
err := lastBrowserProcess.Process.Kill()
lastBrowserProcess = nil
return err
}
// OpenURL opens the specified URL in the default web browser.
// It first attempts to use a platform-agnostic library and falls back to
// platform-specific commands if that fails.
// It uses the pkg/browser library which provides robust cross-platform support
// for Windows, macOS, and Linux.
// If incognito mode is enabled, it will open in a private/incognito window.
//
// Parameters:
// - url: The URL to open.
@@ -21,16 +56,22 @@ import (
// Returns:
// - An error if the URL cannot be opened, otherwise nil.
func OpenURL(url string) error {
fmt.Printf("Attempting to open URL in browser: %s\n", url)
log.Debugf("Opening URL in browser: %s (incognito=%v)", url, incognitoMode)
// Try using the open-golang library first
err := open.Run(url)
// If incognito mode is enabled, use platform-specific incognito commands
if incognitoMode {
log.Debug("Using incognito mode")
return openURLIncognito(url)
}
// Use pkg/browser for cross-platform support
err := pkgbrowser.OpenURL(url)
if err == nil {
log.Debug("Successfully opened URL using open-golang library")
log.Debug("Successfully opened URL using pkg/browser library")
return nil
}
log.Debugf("open-golang failed: %v, trying platform-specific commands", err)
log.Debugf("pkg/browser failed: %v, trying platform-specific commands", err)
// Fallback to platform-specific commands
return openURLPlatformSpecific(url)
@@ -78,18 +119,379 @@ func openURLPlatformSpecific(url string) error {
return nil
}
// openURLIncognito opens a URL in incognito/private browsing mode.
// It first tries to detect the default browser and use its incognito flag.
// Falls back to a chain of known browsers if detection fails.
//
// Parameters:
// - url: The URL to open.
//
// Returns:
// - An error if the URL cannot be opened, otherwise nil.
func openURLIncognito(url string) error {
// First, try to detect and use the default browser
if cmd := tryDefaultBrowserIncognito(url); cmd != nil {
log.Debugf("Using detected default browser: %s %v", cmd.Path, cmd.Args[1:])
if err := cmd.Start(); err == nil {
storeBrowserProcess(cmd)
log.Debug("Successfully opened URL in default browser's incognito mode")
return nil
}
log.Debugf("Failed to start default browser, trying fallback chain")
}
// Fallback to known browser chain
cmd := tryFallbackBrowsersIncognito(url)
if cmd == nil {
log.Warn("No browser with incognito support found, falling back to normal mode")
return openURLPlatformSpecific(url)
}
log.Debugf("Running incognito command: %s %v", cmd.Path, cmd.Args[1:])
err := cmd.Start()
if err != nil {
log.Warnf("Failed to open incognito browser: %v, falling back to normal mode", err)
return openURLPlatformSpecific(url)
}
storeBrowserProcess(cmd)
log.Debug("Successfully opened URL in incognito/private mode")
return nil
}
// storeBrowserProcess safely stores the browser process for later cleanup.
func storeBrowserProcess(cmd *exec.Cmd) {
browserMutex.Lock()
lastBrowserProcess = cmd
browserMutex.Unlock()
}
// tryDefaultBrowserIncognito attempts to detect the default browser and return
// an exec.Cmd configured with the appropriate incognito flag.
func tryDefaultBrowserIncognito(url string) *exec.Cmd {
switch runtime.GOOS {
case "darwin":
return tryDefaultBrowserMacOS(url)
case "windows":
return tryDefaultBrowserWindows(url)
case "linux":
return tryDefaultBrowserLinux(url)
}
return nil
}
// tryDefaultBrowserMacOS detects the default browser on macOS.
func tryDefaultBrowserMacOS(url string) *exec.Cmd {
// Try to get default browser from Launch Services
out, err := exec.Command("defaults", "read", "com.apple.LaunchServices/com.apple.launchservices.secure", "LSHandlers").Output()
if err != nil {
return nil
}
output := string(out)
var browserName string
// Parse the output to find the http/https handler
if containsBrowserID(output, "com.google.chrome") {
browserName = "chrome"
} else if containsBrowserID(output, "org.mozilla.firefox") {
browserName = "firefox"
} else if containsBrowserID(output, "com.apple.safari") {
browserName = "safari"
} else if containsBrowserID(output, "com.brave.browser") {
browserName = "brave"
} else if containsBrowserID(output, "com.microsoft.edgemac") {
browserName = "edge"
}
return createMacOSIncognitoCmd(browserName, url)
}
// containsBrowserID checks if the LaunchServices output contains a browser ID.
func containsBrowserID(output, bundleID string) bool {
return strings.Contains(output, bundleID)
}
// createMacOSIncognitoCmd creates the appropriate incognito command for macOS browsers.
func createMacOSIncognitoCmd(browserName, url string) *exec.Cmd {
switch browserName {
case "chrome":
// Try direct path first
chromePath := "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"
if _, err := exec.LookPath(chromePath); err == nil {
return exec.Command(chromePath, "--incognito", url)
}
return exec.Command("open", "-na", "Google Chrome", "--args", "--incognito", url)
case "firefox":
return exec.Command("open", "-na", "Firefox", "--args", "--private-window", url)
case "safari":
// Safari doesn't have CLI incognito, try AppleScript
return tryAppleScriptSafariPrivate(url)
case "brave":
return exec.Command("open", "-na", "Brave Browser", "--args", "--incognito", url)
case "edge":
return exec.Command("open", "-na", "Microsoft Edge", "--args", "--inprivate", url)
}
return nil
}
// tryAppleScriptSafariPrivate attempts to open Safari in private browsing mode using AppleScript.
func tryAppleScriptSafariPrivate(url string) *exec.Cmd {
// AppleScript to open a new private window in Safari
script := fmt.Sprintf(`
tell application "Safari"
activate
tell application "System Events"
keystroke "n" using {command down, shift down}
delay 0.5
end tell
set URL of document 1 to "%s"
end tell
`, url)
cmd := exec.Command("osascript", "-e", script)
// Test if this approach works by checking if Safari is available
if _, err := exec.LookPath("/Applications/Safari.app/Contents/MacOS/Safari"); err != nil {
log.Debug("Safari not found, AppleScript private window not available")
return nil
}
log.Debug("Attempting Safari private window via AppleScript")
return cmd
}
// tryDefaultBrowserWindows detects the default browser on Windows via registry.
func tryDefaultBrowserWindows(url string) *exec.Cmd {
// Query registry for default browser
out, err := exec.Command("reg", "query",
`HKEY_CURRENT_USER\Software\Microsoft\Windows\Shell\Associations\UrlAssociations\http\UserChoice`,
"/v", "ProgId").Output()
if err != nil {
return nil
}
output := string(out)
var browserName string
// Map ProgId to browser name
if strings.Contains(output, "ChromeHTML") {
browserName = "chrome"
} else if strings.Contains(output, "FirefoxURL") {
browserName = "firefox"
} else if strings.Contains(output, "MSEdgeHTM") {
browserName = "edge"
} else if strings.Contains(output, "BraveHTML") {
browserName = "brave"
}
return createWindowsIncognitoCmd(browserName, url)
}
// createWindowsIncognitoCmd creates the appropriate incognito command for Windows browsers.
func createWindowsIncognitoCmd(browserName, url string) *exec.Cmd {
switch browserName {
case "chrome":
paths := []string{
"chrome",
`C:\Program Files\Google\Chrome\Application\chrome.exe`,
`C:\Program Files (x86)\Google\Chrome\Application\chrome.exe`,
}
for _, p := range paths {
if _, err := exec.LookPath(p); err == nil {
return exec.Command(p, "--incognito", url)
}
}
case "firefox":
if path, err := exec.LookPath("firefox"); err == nil {
return exec.Command(path, "--private-window", url)
}
case "edge":
paths := []string{
"msedge",
`C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe`,
`C:\Program Files\Microsoft\Edge\Application\msedge.exe`,
}
for _, p := range paths {
if _, err := exec.LookPath(p); err == nil {
return exec.Command(p, "--inprivate", url)
}
}
case "brave":
paths := []string{
`C:\Program Files\BraveSoftware\Brave-Browser\Application\brave.exe`,
`C:\Program Files (x86)\BraveSoftware\Brave-Browser\Application\brave.exe`,
}
for _, p := range paths {
if _, err := exec.LookPath(p); err == nil {
return exec.Command(p, "--incognito", url)
}
}
}
return nil
}
// tryDefaultBrowserLinux detects the default browser on Linux using xdg-settings.
func tryDefaultBrowserLinux(url string) *exec.Cmd {
out, err := exec.Command("xdg-settings", "get", "default-web-browser").Output()
if err != nil {
return nil
}
desktop := string(out)
var browserName string
// Map .desktop file to browser name
if strings.Contains(desktop, "google-chrome") || strings.Contains(desktop, "chrome") {
browserName = "chrome"
} else if strings.Contains(desktop, "firefox") {
browserName = "firefox"
} else if strings.Contains(desktop, "chromium") {
browserName = "chromium"
} else if strings.Contains(desktop, "brave") {
browserName = "brave"
} else if strings.Contains(desktop, "microsoft-edge") || strings.Contains(desktop, "msedge") {
browserName = "edge"
}
return createLinuxIncognitoCmd(browserName, url)
}
// createLinuxIncognitoCmd creates the appropriate incognito command for Linux browsers.
func createLinuxIncognitoCmd(browserName, url string) *exec.Cmd {
switch browserName {
case "chrome":
paths := []string{"google-chrome", "google-chrome-stable"}
for _, p := range paths {
if path, err := exec.LookPath(p); err == nil {
return exec.Command(path, "--incognito", url)
}
}
case "firefox":
paths := []string{"firefox", "firefox-esr"}
for _, p := range paths {
if path, err := exec.LookPath(p); err == nil {
return exec.Command(path, "--private-window", url)
}
}
case "chromium":
paths := []string{"chromium", "chromium-browser"}
for _, p := range paths {
if path, err := exec.LookPath(p); err == nil {
return exec.Command(path, "--incognito", url)
}
}
case "brave":
if path, err := exec.LookPath("brave-browser"); err == nil {
return exec.Command(path, "--incognito", url)
}
case "edge":
if path, err := exec.LookPath("microsoft-edge"); err == nil {
return exec.Command(path, "--inprivate", url)
}
}
return nil
}
// tryFallbackBrowsersIncognito tries a chain of known browsers as fallback.
func tryFallbackBrowsersIncognito(url string) *exec.Cmd {
switch runtime.GOOS {
case "darwin":
return tryFallbackBrowsersMacOS(url)
case "windows":
return tryFallbackBrowsersWindows(url)
case "linux":
return tryFallbackBrowsersLinuxChain(url)
}
return nil
}
// tryFallbackBrowsersMacOS tries known browsers on macOS.
func tryFallbackBrowsersMacOS(url string) *exec.Cmd {
// Try Chrome
chromePath := "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"
if _, err := exec.LookPath(chromePath); err == nil {
return exec.Command(chromePath, "--incognito", url)
}
// Try Firefox
if _, err := exec.LookPath("/Applications/Firefox.app/Contents/MacOS/firefox"); err == nil {
return exec.Command("open", "-na", "Firefox", "--args", "--private-window", url)
}
// Try Brave
if _, err := exec.LookPath("/Applications/Brave Browser.app/Contents/MacOS/Brave Browser"); err == nil {
return exec.Command("open", "-na", "Brave Browser", "--args", "--incognito", url)
}
// Try Edge
if _, err := exec.LookPath("/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge"); err == nil {
return exec.Command("open", "-na", "Microsoft Edge", "--args", "--inprivate", url)
}
// Last resort: try Safari with AppleScript
if cmd := tryAppleScriptSafariPrivate(url); cmd != nil {
log.Info("Using Safari with AppleScript for private browsing (may require accessibility permissions)")
return cmd
}
return nil
}
// tryFallbackBrowsersWindows tries known browsers on Windows.
func tryFallbackBrowsersWindows(url string) *exec.Cmd {
// Chrome
chromePaths := []string{
"chrome",
`C:\Program Files\Google\Chrome\Application\chrome.exe`,
`C:\Program Files (x86)\Google\Chrome\Application\chrome.exe`,
}
for _, p := range chromePaths {
if _, err := exec.LookPath(p); err == nil {
return exec.Command(p, "--incognito", url)
}
}
// Firefox
if path, err := exec.LookPath("firefox"); err == nil {
return exec.Command(path, "--private-window", url)
}
// Edge (usually available on Windows 10+)
edgePaths := []string{
"msedge",
`C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe`,
`C:\Program Files\Microsoft\Edge\Application\msedge.exe`,
}
for _, p := range edgePaths {
if _, err := exec.LookPath(p); err == nil {
return exec.Command(p, "--inprivate", url)
}
}
return nil
}
// tryFallbackBrowsersLinuxChain tries known browsers on Linux.
func tryFallbackBrowsersLinuxChain(url string) *exec.Cmd {
type browserConfig struct {
name string
flag string
}
browsers := []browserConfig{
{"google-chrome", "--incognito"},
{"google-chrome-stable", "--incognito"},
{"chromium", "--incognito"},
{"chromium-browser", "--incognito"},
{"firefox", "--private-window"},
{"firefox-esr", "--private-window"},
{"brave-browser", "--incognito"},
{"microsoft-edge", "--inprivate"},
}
for _, b := range browsers {
if path, err := exec.LookPath(b.name); err == nil {
return exec.Command(path, b.flag, url)
}
}
return nil
}
// IsAvailable checks if the system has a command available to open a web browser.
// It verifies the presence of necessary commands for the current operating system.
//
// Returns:
// - true if a browser can be opened, false otherwise.
func IsAvailable() bool {
// First check if open-golang can work
testErr := open.Run("about:blank")
if testErr == nil {
return true
}
// Check platform-specific commands
switch runtime.GOOS {
case "darwin":

View File

@@ -19,6 +19,7 @@ func newAuthManager() *sdkAuth.Manager {
sdkAuth.NewQwenAuthenticator(),
sdkAuth.NewIFlowAuthenticator(),
sdkAuth.NewAntigravityAuthenticator(),
sdkAuth.NewKiroAuthenticator(),
sdkAuth.NewGitHubCopilotAuthenticator(),
)
return manager

160
internal/cmd/kiro_login.go Normal file
View File

@@ -0,0 +1,160 @@
package cmd
import (
"context"
"fmt"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
log "github.com/sirupsen/logrus"
)
// DoKiroLogin triggers the Kiro authentication flow with Google OAuth.
// This is the default login method (same as --kiro-google-login).
//
// Parameters:
// - cfg: The application configuration
// - options: Login options including Prompt field
func DoKiroLogin(cfg *config.Config, options *LoginOptions) {
// Use Google login as default
DoKiroGoogleLogin(cfg, options)
}
// DoKiroGoogleLogin triggers Kiro authentication with Google OAuth.
// This uses a custom protocol handler (kiro://) to receive the callback.
//
// Parameters:
// - cfg: The application configuration
// - options: Login options including prompts
func DoKiroGoogleLogin(cfg *config.Config, options *LoginOptions) {
if options == nil {
options = &LoginOptions{}
}
// Note: Kiro defaults to incognito mode for multi-account support.
// Users can override with --no-incognito if they want to use existing browser sessions.
manager := newAuthManager()
// Use KiroAuthenticator with Google login
authenticator := sdkAuth.NewKiroAuthenticator()
record, err := authenticator.LoginWithGoogle(context.Background(), cfg, &sdkAuth.LoginOptions{
NoBrowser: options.NoBrowser,
Metadata: map[string]string{},
Prompt: options.Prompt,
})
if err != nil {
log.Errorf("Kiro Google authentication failed: %v", err)
fmt.Println("\nTroubleshooting:")
fmt.Println("1. Make sure the protocol handler is installed")
fmt.Println("2. Complete the Google login in the browser")
fmt.Println("3. If callback fails, try: --kiro-import (after logging in via Kiro IDE)")
return
}
// Save the auth record
savedPath, err := manager.SaveAuth(record, cfg)
if err != nil {
log.Errorf("Failed to save auth: %v", err)
return
}
if savedPath != "" {
fmt.Printf("Authentication saved to %s\n", savedPath)
}
if record != nil && record.Label != "" {
fmt.Printf("Authenticated as %s\n", record.Label)
}
fmt.Println("Kiro Google authentication successful!")
}
// DoKiroAWSLogin triggers Kiro authentication with AWS Builder ID.
// This uses the device code flow for AWS SSO OIDC authentication.
//
// Parameters:
// - cfg: The application configuration
// - options: Login options including prompts
func DoKiroAWSLogin(cfg *config.Config, options *LoginOptions) {
if options == nil {
options = &LoginOptions{}
}
// Note: Kiro defaults to incognito mode for multi-account support.
// Users can override with --no-incognito if they want to use existing browser sessions.
manager := newAuthManager()
// Use KiroAuthenticator with AWS Builder ID login (device code flow)
authenticator := sdkAuth.NewKiroAuthenticator()
record, err := authenticator.Login(context.Background(), cfg, &sdkAuth.LoginOptions{
NoBrowser: options.NoBrowser,
Metadata: map[string]string{},
Prompt: options.Prompt,
})
if err != nil {
log.Errorf("Kiro AWS authentication failed: %v", err)
fmt.Println("\nTroubleshooting:")
fmt.Println("1. Make sure you have an AWS Builder ID")
fmt.Println("2. Complete the authorization in the browser")
fmt.Println("3. If callback fails, try: --kiro-import (after logging in via Kiro IDE)")
return
}
// Save the auth record
savedPath, err := manager.SaveAuth(record, cfg)
if err != nil {
log.Errorf("Failed to save auth: %v", err)
return
}
if savedPath != "" {
fmt.Printf("Authentication saved to %s\n", savedPath)
}
if record != nil && record.Label != "" {
fmt.Printf("Authenticated as %s\n", record.Label)
}
fmt.Println("Kiro AWS authentication successful!")
}
// DoKiroImport imports Kiro token from Kiro IDE's token file.
// This is useful for users who have already logged in via Kiro IDE
// and want to use the same credentials in CLI Proxy API.
//
// Parameters:
// - cfg: The application configuration
// - options: Login options (currently unused for import)
func DoKiroImport(cfg *config.Config, options *LoginOptions) {
if options == nil {
options = &LoginOptions{}
}
manager := newAuthManager()
// Use ImportFromKiroIDE instead of Login
authenticator := sdkAuth.NewKiroAuthenticator()
record, err := authenticator.ImportFromKiroIDE(context.Background(), cfg)
if err != nil {
log.Errorf("Kiro token import failed: %v", err)
fmt.Println("\nMake sure you have logged in to Kiro IDE first:")
fmt.Println("1. Open Kiro IDE")
fmt.Println("2. Click 'Sign in with Google' (or GitHub)")
fmt.Println("3. Complete the login process")
fmt.Println("4. Run this command again")
return
}
// Save the imported auth record
savedPath, err := manager.SaveAuth(record, cfg)
if err != nil {
log.Errorf("Failed to save auth: %v", err)
return
}
if savedPath != "" {
fmt.Printf("Authentication saved to %s\n", savedPath)
}
if record != nil && record.Label != "" {
fmt.Printf("Imported as %s\n", record.Label)
}
fmt.Println("Kiro token import successful!")
}

View File

@@ -65,20 +65,20 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {
authenticator := sdkAuth.NewGeminiAuthenticator()
record, errLogin := authenticator.Login(ctx, cfg, loginOpts)
if errLogin != nil {
log.Fatalf("Gemini authentication failed: %v", errLogin)
log.Errorf("Gemini authentication failed: %v", errLogin)
return
}
storage, okStorage := record.Storage.(*gemini.GeminiTokenStorage)
if !okStorage || storage == nil {
log.Fatal("Gemini authentication failed: unsupported token storage")
log.Error("Gemini authentication failed: unsupported token storage")
return
}
geminiAuth := gemini.NewGeminiAuth()
httpClient, errClient := geminiAuth.GetAuthenticatedClient(ctx, storage, cfg, options.NoBrowser)
if errClient != nil {
log.Fatalf("Gemini authentication failed: %v", errClient)
log.Errorf("Gemini authentication failed: %v", errClient)
return
}
@@ -86,7 +86,7 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {
projects, errProjects := fetchGCPProjects(ctx, httpClient)
if errProjects != nil {
log.Fatalf("Failed to get project list: %v", errProjects)
log.Errorf("Failed to get project list: %v", errProjects)
return
}
@@ -98,11 +98,11 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {
selectedProjectID := promptForProjectSelection(projects, strings.TrimSpace(projectID), promptFn)
projectSelections, errSelection := resolveProjectSelections(selectedProjectID, projects)
if errSelection != nil {
log.Fatalf("Invalid project selection: %v", errSelection)
log.Errorf("Invalid project selection: %v", errSelection)
return
}
if len(projectSelections) == 0 {
log.Fatal("No project selected; aborting login.")
log.Error("No project selected; aborting login.")
return
}
@@ -116,7 +116,7 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {
showProjectSelectionHelp(storage.Email, projects)
return
}
log.Fatalf("Failed to complete user setup: %v", errSetup)
log.Errorf("Failed to complete user setup: %v", errSetup)
return
}
finalID := strings.TrimSpace(storage.ProjectID)
@@ -133,11 +133,11 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {
for _, pid := range activatedProjects {
isChecked, errCheck := checkCloudAPIIsEnabled(ctx, httpClient, pid)
if errCheck != nil {
log.Fatalf("Failed to check if Cloud AI API is enabled for %s: %v", pid, errCheck)
log.Errorf("Failed to check if Cloud AI API is enabled for %s: %v", pid, errCheck)
return
}
if !isChecked {
log.Fatalf("Failed to check if Cloud AI API is enabled for project %s. If you encounter an error message, please create an issue.", pid)
log.Errorf("Failed to check if Cloud AI API is enabled for project %s. If you encounter an error message, please create an issue.", pid)
return
}
}
@@ -153,7 +153,7 @@ func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {
savedPath, errSave := store.Save(ctx, record)
if errSave != nil {
log.Fatalf("Failed to save token to file: %v", errSave)
log.Errorf("Failed to save token to file: %v", errSave)
return
}
@@ -555,6 +555,7 @@ func checkCloudAPIIsEnabled(ctx context.Context, httpClient *http.Client, projec
continue
}
}
_ = resp.Body.Close()
return false, fmt.Errorf("project activation required: %s", errMessage)
}
return true, nil

View File

@@ -45,12 +45,13 @@ func StartService(cfg *config.Config, configPath string, localPassword string) {
service, err := builder.Build()
if err != nil {
log.Fatalf("failed to build proxy service: %v", err)
log.Errorf("failed to build proxy service: %v", err)
return
}
err = service.Run(runCtx)
if err != nil && !errors.Is(err, context.Canceled) {
log.Fatalf("proxy service exited with error: %v", err)
log.Errorf("proxy service exited with error: %v", err)
}
}

View File

@@ -29,30 +29,30 @@ func DoVertexImport(cfg *config.Config, keyPath string) {
}
rawPath := strings.TrimSpace(keyPath)
if rawPath == "" {
log.Fatalf("vertex-import: missing service account key path")
log.Errorf("vertex-import: missing service account key path")
return
}
data, errRead := os.ReadFile(rawPath)
if errRead != nil {
log.Fatalf("vertex-import: read file failed: %v", errRead)
log.Errorf("vertex-import: read file failed: %v", errRead)
return
}
var sa map[string]any
if errUnmarshal := json.Unmarshal(data, &sa); errUnmarshal != nil {
log.Fatalf("vertex-import: invalid service account json: %v", errUnmarshal)
log.Errorf("vertex-import: invalid service account json: %v", errUnmarshal)
return
}
// Validate and normalize private_key before saving
normalizedSA, errFix := vertex.NormalizeServiceAccountMap(sa)
if errFix != nil {
log.Fatalf("vertex-import: %v", errFix)
log.Errorf("vertex-import: %v", errFix)
return
}
sa = normalizedSA
email, _ := sa["client_email"].(string)
projectID, _ := sa["project_id"].(string)
if strings.TrimSpace(projectID) == "" {
log.Fatalf("vertex-import: project_id missing in service account json")
log.Errorf("vertex-import: project_id missing in service account json")
return
}
if strings.TrimSpace(email) == "" {
@@ -92,7 +92,7 @@ func DoVertexImport(cfg *config.Config, keyPath string) {
}
path, errSave := store.Save(context.Background(), record)
if errSave != nil {
log.Fatalf("vertex-import: save credential failed: %v", errSave)
log.Errorf("vertex-import: save credential failed: %v", errSave)
return
}
fmt.Printf("Vertex credentials imported: %s\n", path)

View File

@@ -20,22 +20,17 @@ import (
// Config represents the application's configuration, loaded from a YAML file.
type Config struct {
config.SDKConfig `yaml:",inline"`
// Host is the network host/interface on which the API server will bind.
// Default is empty ("") to bind all interfaces (IPv4 + IPv6). Use "127.0.0.1" or "localhost" for local-only access.
Host string `yaml:"host" json:"-"`
// Port is the network port on which the API server will listen.
Port int `yaml:"port" json:"-"`
// TLS config controls HTTPS server settings.
TLS TLSConfig `yaml:"tls" json:"tls"`
// AmpUpstreamURL defines the upstream Amp control plane used for non-provider calls.
AmpUpstreamURL string `yaml:"amp-upstream-url" json:"amp-upstream-url"`
// AmpUpstreamAPIKey optionally overrides the Authorization header when proxying Amp upstream calls.
AmpUpstreamAPIKey string `yaml:"amp-upstream-api-key" json:"amp-upstream-api-key"`
// AmpRestrictManagementToLocalhost restricts Amp management routes (/api/user, /api/threads, etc.)
// to only accept connections from localhost (127.0.0.1, ::1). When true, prevents drive-by
// browser attacks and remote access to management endpoints. Default: true (recommended).
AmpRestrictManagementToLocalhost bool `yaml:"amp-restrict-management-to-localhost" json:"amp-restrict-management-to-localhost"`
// RemoteManagement nests management-related options under 'remote-management'.
RemoteManagement RemoteManagement `yaml:"remote-management" json:"-"`
// AuthDir is the directory where authentication token files are stored.
AuthDir string `yaml:"auth-dir" json:"-"`
@@ -52,40 +47,51 @@ type Config struct {
// DisableCooling disables quota cooldown scheduling when true.
DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"`
// RequestRetry defines the retry times when the request failed.
RequestRetry int `yaml:"request-retry" json:"request-retry"`
// MaxRetryInterval defines the maximum wait time in seconds before retrying a cooled-down credential.
MaxRetryInterval int `yaml:"max-retry-interval" json:"max-retry-interval"`
// QuotaExceeded defines the behavior when a quota is exceeded.
QuotaExceeded QuotaExceeded `yaml:"quota-exceeded" json:"quota-exceeded"`
// WebsocketAuth enables or disables authentication for the WebSocket API.
WebsocketAuth bool `yaml:"ws-auth" json:"ws-auth"`
// GlAPIKey exposes the legacy generative language API key list for backward compatibility.
GlAPIKey []string `yaml:"generative-language-api-key" json:"generative-language-api-key"`
// GeminiKey defines Gemini API key configurations with optional routing overrides.
GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`
// RequestRetry defines the retry times when the request failed.
RequestRetry int `yaml:"request-retry" json:"request-retry"`
// MaxRetryInterval defines the maximum wait time in seconds before retrying a cooled-down credential.
MaxRetryInterval int `yaml:"max-retry-interval" json:"max-retry-interval"`
// ClaudeKey defines a list of Claude API key configurations as specified in the YAML configuration file.
ClaudeKey []ClaudeKey `yaml:"claude-api-key" json:"claude-api-key"`
// KiroKey defines a list of Kiro (AWS CodeWhisperer) configurations.
KiroKey []KiroKey `yaml:"kiro" json:"kiro"`
// Codex defines a list of Codex API key configurations as specified in the YAML configuration file.
CodexKey []CodexKey `yaml:"codex-api-key" json:"codex-api-key"`
// ClaudeKey defines a list of Claude API key configurations as specified in the YAML configuration file.
ClaudeKey []ClaudeKey `yaml:"claude-api-key" json:"claude-api-key"`
// OpenAICompatibility defines OpenAI API compatibility configurations for external providers.
OpenAICompatibility []OpenAICompatibility `yaml:"openai-compatibility" json:"openai-compatibility"`
// RemoteManagement nests management-related options under 'remote-management'.
RemoteManagement RemoteManagement `yaml:"remote-management" json:"-"`
// VertexCompatAPIKey defines Vertex AI-compatible API key configurations for third-party providers.
// Used for services that use Vertex AI-style paths but with simple API key authentication.
VertexCompatAPIKey []VertexCompatKey `yaml:"vertex-api-key" json:"vertex-api-key"`
// AmpCode contains Amp CLI upstream configuration, management restrictions, and model mappings.
AmpCode AmpCode `yaml:"ampcode" json:"ampcode"`
// OAuthExcludedModels defines per-provider global model exclusions applied to OAuth/file-backed auth entries.
OAuthExcludedModels map[string][]string `yaml:"oauth-excluded-models,omitempty" json:"oauth-excluded-models,omitempty"`
// Payload defines default and override rules for provider payload parameters.
Payload PayloadConfig `yaml:"payload" json:"payload"`
// OAuthExcludedModels defines per-provider global model exclusions applied to OAuth/file-backed auth entries.
OAuthExcludedModels map[string][]string `yaml:"oauth-excluded-models,omitempty" json:"oauth-excluded-models,omitempty"`
// IncognitoBrowser enables opening OAuth URLs in incognito/private browsing mode.
// This is useful when you want to login with a different account without logging out
// from your current session. Default: false.
IncognitoBrowser bool `yaml:"incognito-browser" json:"incognito-browser"`
legacyMigrationPending bool `yaml:"-" json:"-"`
}
// TLSConfig holds HTTPS server settings.
@@ -118,6 +124,42 @@ type QuotaExceeded struct {
SwitchPreviewModel bool `yaml:"switch-preview-model" json:"switch-preview-model"`
}
// AmpModelMapping defines a model name mapping for Amp CLI requests.
// When Amp requests a model that isn't available locally, this mapping
// allows routing to an alternative model that IS available.
type AmpModelMapping struct {
// From is the model name that Amp CLI requests (e.g., "claude-opus-4.5").
From string `yaml:"from" json:"from"`
// To is the target model name to route to (e.g., "claude-sonnet-4").
// The target model must have available providers in the registry.
To string `yaml:"to" json:"to"`
}
// AmpCode groups Amp CLI integration settings including upstream routing,
// optional overrides, management route restrictions, and model fallback mappings.
type AmpCode struct {
// UpstreamURL defines the upstream Amp control plane used for non-provider calls.
UpstreamURL string `yaml:"upstream-url" json:"upstream-url"`
// UpstreamAPIKey optionally overrides the Authorization header when proxying Amp upstream calls.
UpstreamAPIKey string `yaml:"upstream-api-key" json:"upstream-api-key"`
// RestrictManagementToLocalhost restricts Amp management routes (/api/user, /api/threads, etc.)
// to only accept connections from localhost (127.0.0.1, ::1). When true, prevents drive-by
// browser attacks and remote access to management endpoints. Default: true (recommended).
RestrictManagementToLocalhost bool `yaml:"restrict-management-to-localhost" json:"restrict-management-to-localhost"`
// ModelMappings defines model name mappings for Amp CLI requests.
// When Amp requests a model that isn't available locally, these mappings
// allow routing to an alternative model that IS available.
ModelMappings []AmpModelMapping `yaml:"model-mappings" json:"model-mappings"`
// ForceModelMappings when true, model mappings take precedence over local API keys.
// When false (default), local API keys are used first if available.
ForceModelMappings bool `yaml:"force-model-mappings" json:"force-model-mappings"`
}
// PayloadConfig defines default and override parameter rules applied to provider payloads.
type PayloadConfig struct {
// Default defines rules that only set parameters when they are missing in the payload.
@@ -213,6 +255,31 @@ type GeminiKey struct {
ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"`
}
// KiroKey represents the configuration for Kiro (AWS CodeWhisperer) authentication.
type KiroKey struct {
// TokenFile is the path to the Kiro token file (default: ~/.aws/sso/cache/kiro-auth-token.json)
TokenFile string `yaml:"token-file,omitempty" json:"token-file,omitempty"`
// AccessToken is the OAuth access token for direct configuration.
AccessToken string `yaml:"access-token,omitempty" json:"access-token,omitempty"`
// RefreshToken is the OAuth refresh token for token renewal.
RefreshToken string `yaml:"refresh-token,omitempty" json:"refresh-token,omitempty"`
// ProfileArn is the AWS CodeWhisperer profile ARN.
ProfileArn string `yaml:"profile-arn,omitempty" json:"profile-arn,omitempty"`
// Region is the AWS region (default: us-east-1).
Region string `yaml:"region,omitempty" json:"region,omitempty"`
// ProxyURL optionally overrides the global proxy for this configuration.
ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"`
// AgentTaskType sets the Kiro API task type. Known values: "vibe", "dev", "chat".
// Leave empty to let API use defaults. Different values may inject different system prompts.
AgentTaskType string `yaml:"agent-task-type,omitempty" json:"agent-task-type,omitempty"`
}
// OpenAICompatibility represents the configuration for OpenAI API compatibility
// with external providers, allowing model aliases to be routed through OpenAI API format.
type OpenAICompatibility struct {
@@ -222,10 +289,6 @@ type OpenAICompatibility struct {
// BaseURL is the base URL for the external OpenAI-compatible API endpoint.
BaseURL string `yaml:"base-url" json:"base-url"`
// APIKeys are the authentication keys for accessing the external API services.
// Deprecated: Use APIKeyEntries instead to support per-key proxy configuration.
APIKeys []string `yaml:"api-keys,omitempty" json:"api-keys,omitempty"`
// APIKeyEntries defines API keys with optional per-key proxy configuration.
APIKeyEntries []OpenAICompatibilityAPIKey `yaml:"api-key-entries,omitempty" json:"api-key-entries,omitempty"`
@@ -293,10 +356,12 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
// Unmarshal the YAML data into the Config struct.
var cfg Config
// Set defaults before unmarshal so that absent keys keep defaults.
cfg.Host = "" // Default empty: binds to all interfaces (IPv4 + IPv6)
cfg.LoggingToFile = false
cfg.UsageStatisticsEnabled = false
cfg.DisableCooling = false
cfg.AmpRestrictManagementToLocalhost = true // Default to secure: only localhost access
cfg.AmpCode.RestrictManagementToLocalhost = true // Default to secure: only localhost access
cfg.IncognitoBrowser = false // Default to normal browser (AWS uses incognito by force)
if err = yaml.Unmarshal(data, &cfg); err != nil {
if optional {
// In cloud deploy mode, if YAML parsing fails, return empty config instead of error.
@@ -305,6 +370,19 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
return nil, fmt.Errorf("failed to parse config file: %w", err)
}
var legacy legacyConfigData
if errLegacy := yaml.Unmarshal(data, &legacy); errLegacy == nil {
if cfg.migrateLegacyGeminiKeys(legacy.LegacyGeminiKeys) {
cfg.legacyMigrationPending = true
}
if cfg.migrateLegacyOpenAICompatibilityKeys(legacy.OpenAICompat) {
cfg.legacyMigrationPending = true
}
if cfg.migrateLegacyAmpConfig(&legacy) {
cfg.legacyMigrationPending = true
}
}
// Hash remote management key if plaintext is detected (nested)
// We consider a value to be already hashed if it looks like a bcrypt hash ($2a$, $2b$, or $2y$ prefix).
if cfg.RemoteManagement.SecretKey != "" && !looksLikeBcrypt(cfg.RemoteManagement.SecretKey) {
@@ -325,18 +403,36 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
// Sanitize Gemini API key configuration and migrate legacy entries.
cfg.SanitizeGeminiKeys()
// Sanitize Vertex-compatible API keys: drop entries without base-url
cfg.SanitizeVertexCompatKeys()
// Sanitize Codex keys: drop entries without base-url
cfg.SanitizeCodexKeys()
// Sanitize Claude key headers
cfg.SanitizeClaudeKeys()
// Sanitize Kiro keys: trim whitespace from credential fields
cfg.SanitizeKiroKeys()
// Sanitize OpenAI compatibility providers: drop entries without base-url
cfg.SanitizeOpenAICompatibility()
// Normalize OAuth provider model exclusion map.
cfg.OAuthExcludedModels = NormalizeOAuthExcludedModels(cfg.OAuthExcludedModels)
if cfg.legacyMigrationPending {
fmt.Println("Detected legacy configuration keys, attempting to persist the normalized config...")
if !optional && configFile != "" {
if err := SaveConfigPreserveComments(configFile, &cfg); err != nil {
return nil, fmt.Errorf("failed to persist migrated legacy config: %w", err)
}
fmt.Println("Legacy configuration normalized and persisted.")
} else {
fmt.Println("Legacy configuration normalized in memory; persistence skipped.")
}
}
// Return the populated configuration struct.
return &cfg, nil
}
@@ -395,6 +491,22 @@ func (cfg *Config) SanitizeClaudeKeys() {
}
}
// SanitizeKiroKeys trims whitespace from Kiro credential fields.
func (cfg *Config) SanitizeKiroKeys() {
if cfg == nil || len(cfg.KiroKey) == 0 {
return
}
for i := range cfg.KiroKey {
entry := &cfg.KiroKey[i]
entry.TokenFile = strings.TrimSpace(entry.TokenFile)
entry.AccessToken = strings.TrimSpace(entry.AccessToken)
entry.RefreshToken = strings.TrimSpace(entry.RefreshToken)
entry.ProfileArn = strings.TrimSpace(entry.ProfileArn)
entry.Region = strings.TrimSpace(entry.Region)
entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
}
}
// SanitizeGeminiKeys deduplicates and normalizes Gemini credentials.
func (cfg *Config) SanitizeGeminiKeys() {
if cfg == nil {
@@ -420,22 +532,6 @@ func (cfg *Config) SanitizeGeminiKeys() {
out = append(out, entry)
}
cfg.GeminiKey = out
if len(cfg.GlAPIKey) > 0 {
for _, raw := range cfg.GlAPIKey {
key := strings.TrimSpace(raw)
if key == "" {
continue
}
if _, exists := seen[key]; exists {
continue
}
cfg.GeminiKey = append(cfg.GeminiKey, GeminiKey{APIKey: key})
seen[key] = struct{}{}
}
}
cfg.GlAPIKey = nil
}
func syncInlineAccessProvider(cfg *Config) {
@@ -571,9 +667,13 @@ func SaveConfigPreserveComments(configFile string, cfg *Config) error {
return fmt.Errorf("expected generated root mapping node")
}
// Remove deprecated auth block before merging to avoid persisting it again.
removeMapKey(original.Content[0], "auth")
// Remove deprecated sections before merging back the sanitized config.
removeLegacyAuthBlock(original.Content[0])
removeLegacyOpenAICompatAPIKeys(original.Content[0])
removeLegacyAmpKeys(original.Content[0])
removeLegacyGenerativeLanguageKeys(original.Content[0])
pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "oauth-excluded-models")
// Merge generated into original in-place, preserving comments/order of existing nodes.
mergeMappingPreserve(original.Content[0], generated.Content[0])
@@ -772,6 +872,10 @@ func mergeNodePreserve(dst, src *yaml.Node) {
continue
}
mergeNodePreserve(dst.Content[i], src.Content[i])
if dst.Content[i] != nil && src.Content[i] != nil &&
dst.Content[i].Kind == yaml.MappingNode && src.Content[i].Kind == yaml.MappingNode {
pruneMissingMapKeys(dst.Content[i], src.Content[i])
}
}
// Append any extra items from src
for i := len(dst.Content); i < len(src.Content); i++ {
@@ -813,6 +917,7 @@ func shouldSkipEmptyCollectionOnPersist(key string, node *yaml.Node) bool {
switch key {
case "generative-language-api-key",
"gemini-api-key",
"vertex-api-key",
"claude-api-key",
"codex-api-key",
"openai-compatibility":
@@ -1030,22 +1135,70 @@ func removeMapKey(mapNode *yaml.Node, key string) {
}
}
func removeLegacyOpenAICompatAPIKeys(root *yaml.Node) {
if root == nil || root.Kind != yaml.MappingNode {
func pruneMappingToGeneratedKeys(dstRoot, srcRoot *yaml.Node, key string) {
if key == "" || dstRoot == nil || srcRoot == nil {
return
}
idx := findMapKeyIndex(root, "openai-compatibility")
if idx < 0 || idx+1 >= len(root.Content) {
if dstRoot.Kind != yaml.MappingNode || srcRoot.Kind != yaml.MappingNode {
return
}
seq := root.Content[idx+1]
if seq == nil || seq.Kind != yaml.SequenceNode {
dstIdx := findMapKeyIndex(dstRoot, key)
if dstIdx < 0 || dstIdx+1 >= len(dstRoot.Content) {
return
}
for i := range seq.Content {
if seq.Content[i] != nil && seq.Content[i].Kind == yaml.MappingNode {
removeMapKey(seq.Content[i], "api-keys")
srcIdx := findMapKeyIndex(srcRoot, key)
if srcIdx < 0 {
removeMapKey(dstRoot, key)
return
}
if srcIdx+1 >= len(srcRoot.Content) {
return
}
srcVal := srcRoot.Content[srcIdx+1]
dstVal := dstRoot.Content[dstIdx+1]
if srcVal == nil {
dstRoot.Content[dstIdx+1] = nil
return
}
if srcVal.Kind != yaml.MappingNode {
dstRoot.Content[dstIdx+1] = deepCopyNode(srcVal)
return
}
if dstVal == nil || dstVal.Kind != yaml.MappingNode {
dstRoot.Content[dstIdx+1] = deepCopyNode(srcVal)
return
}
pruneMissingMapKeys(dstVal, srcVal)
}
func pruneMissingMapKeys(dstMap, srcMap *yaml.Node) {
if dstMap == nil || srcMap == nil || dstMap.Kind != yaml.MappingNode || srcMap.Kind != yaml.MappingNode {
return
}
keep := make(map[string]struct{}, len(srcMap.Content)/2)
for i := 0; i+1 < len(srcMap.Content); i += 2 {
keyNode := srcMap.Content[i]
if keyNode == nil {
continue
}
key := strings.TrimSpace(keyNode.Value)
if key == "" {
continue
}
keep[key] = struct{}{}
}
for i := 0; i+1 < len(dstMap.Content); {
keyNode := dstMap.Content[i]
if keyNode == nil {
i += 2
continue
}
key := strings.TrimSpace(keyNode.Value)
if _, ok := keep[key]; !ok {
dstMap.Content = append(dstMap.Content[:i], dstMap.Content[i+2:]...)
continue
}
i += 2
}
}
@@ -1075,3 +1228,194 @@ func normalizeCollectionNodeStyles(node *yaml.Node) {
// Scalars keep their existing style to preserve quoting
}
}
// Legacy migration helpers (move deprecated config keys into structured fields).
type legacyConfigData struct {
LegacyGeminiKeys []string `yaml:"generative-language-api-key"`
OpenAICompat []legacyOpenAICompatibility `yaml:"openai-compatibility"`
AmpUpstreamURL string `yaml:"amp-upstream-url"`
AmpUpstreamAPIKey string `yaml:"amp-upstream-api-key"`
AmpRestrictManagement *bool `yaml:"amp-restrict-management-to-localhost"`
AmpModelMappings []AmpModelMapping `yaml:"amp-model-mappings"`
}
type legacyOpenAICompatibility struct {
Name string `yaml:"name"`
BaseURL string `yaml:"base-url"`
APIKeys []string `yaml:"api-keys"`
}
func (cfg *Config) migrateLegacyGeminiKeys(legacy []string) bool {
if cfg == nil || len(legacy) == 0 {
return false
}
changed := false
seen := make(map[string]struct{}, len(cfg.GeminiKey))
for i := range cfg.GeminiKey {
key := strings.TrimSpace(cfg.GeminiKey[i].APIKey)
if key == "" {
continue
}
seen[key] = struct{}{}
}
for _, raw := range legacy {
key := strings.TrimSpace(raw)
if key == "" {
continue
}
if _, exists := seen[key]; exists {
continue
}
cfg.GeminiKey = append(cfg.GeminiKey, GeminiKey{APIKey: key})
seen[key] = struct{}{}
changed = true
}
return changed
}
func (cfg *Config) migrateLegacyOpenAICompatibilityKeys(legacy []legacyOpenAICompatibility) bool {
if cfg == nil || len(cfg.OpenAICompatibility) == 0 || len(legacy) == 0 {
return false
}
changed := false
for _, legacyEntry := range legacy {
if len(legacyEntry.APIKeys) == 0 {
continue
}
target := findOpenAICompatTarget(cfg.OpenAICompatibility, legacyEntry.Name, legacyEntry.BaseURL)
if target == nil {
continue
}
if mergeLegacyOpenAICompatAPIKeys(target, legacyEntry.APIKeys) {
changed = true
}
}
return changed
}
func mergeLegacyOpenAICompatAPIKeys(entry *OpenAICompatibility, keys []string) bool {
if entry == nil || len(keys) == 0 {
return false
}
changed := false
existing := make(map[string]struct{}, len(entry.APIKeyEntries))
for i := range entry.APIKeyEntries {
key := strings.TrimSpace(entry.APIKeyEntries[i].APIKey)
if key == "" {
continue
}
existing[key] = struct{}{}
}
for _, raw := range keys {
key := strings.TrimSpace(raw)
if key == "" {
continue
}
if _, ok := existing[key]; ok {
continue
}
entry.APIKeyEntries = append(entry.APIKeyEntries, OpenAICompatibilityAPIKey{APIKey: key})
existing[key] = struct{}{}
changed = true
}
return changed
}
func findOpenAICompatTarget(entries []OpenAICompatibility, legacyName, legacyBase string) *OpenAICompatibility {
nameKey := strings.ToLower(strings.TrimSpace(legacyName))
baseKey := strings.ToLower(strings.TrimSpace(legacyBase))
if nameKey != "" && baseKey != "" {
for i := range entries {
if strings.ToLower(strings.TrimSpace(entries[i].Name)) == nameKey &&
strings.ToLower(strings.TrimSpace(entries[i].BaseURL)) == baseKey {
return &entries[i]
}
}
}
if baseKey != "" {
for i := range entries {
if strings.ToLower(strings.TrimSpace(entries[i].BaseURL)) == baseKey {
return &entries[i]
}
}
}
if nameKey != "" {
for i := range entries {
if strings.ToLower(strings.TrimSpace(entries[i].Name)) == nameKey {
return &entries[i]
}
}
}
return nil
}
func (cfg *Config) migrateLegacyAmpConfig(legacy *legacyConfigData) bool {
if cfg == nil || legacy == nil {
return false
}
changed := false
if cfg.AmpCode.UpstreamURL == "" {
if val := strings.TrimSpace(legacy.AmpUpstreamURL); val != "" {
cfg.AmpCode.UpstreamURL = val
changed = true
}
}
if cfg.AmpCode.UpstreamAPIKey == "" {
if val := strings.TrimSpace(legacy.AmpUpstreamAPIKey); val != "" {
cfg.AmpCode.UpstreamAPIKey = val
changed = true
}
}
if legacy.AmpRestrictManagement != nil {
cfg.AmpCode.RestrictManagementToLocalhost = *legacy.AmpRestrictManagement
changed = true
}
if len(cfg.AmpCode.ModelMappings) == 0 && len(legacy.AmpModelMappings) > 0 {
cfg.AmpCode.ModelMappings = append([]AmpModelMapping(nil), legacy.AmpModelMappings...)
changed = true
}
return changed
}
func removeLegacyOpenAICompatAPIKeys(root *yaml.Node) {
if root == nil || root.Kind != yaml.MappingNode {
return
}
idx := findMapKeyIndex(root, "openai-compatibility")
if idx < 0 || idx+1 >= len(root.Content) {
return
}
seq := root.Content[idx+1]
if seq == nil || seq.Kind != yaml.SequenceNode {
return
}
for i := range seq.Content {
if seq.Content[i] != nil && seq.Content[i].Kind == yaml.MappingNode {
removeMapKey(seq.Content[i], "api-keys")
}
}
}
func removeLegacyAmpKeys(root *yaml.Node) {
if root == nil || root.Kind != yaml.MappingNode {
return
}
removeMapKey(root, "amp-upstream-url")
removeMapKey(root, "amp-upstream-api-key")
removeMapKey(root, "amp-restrict-management-to-localhost")
removeMapKey(root, "amp-model-mappings")
}
func removeLegacyGenerativeLanguageKeys(root *yaml.Node) {
if root == nil || root.Kind != yaml.MappingNode {
return
}
removeMapKey(root, "generative-language-api-key")
}
func removeLegacyAuthBlock(root *yaml.Node) {
if root == nil || root.Kind != yaml.MappingNode {
return
}
removeMapKey(root, "auth")
}

View File

@@ -0,0 +1,84 @@
package config
import "strings"
// VertexCompatKey represents the configuration for Vertex AI-compatible API keys.
// This supports third-party services that use Vertex AI-style endpoint paths
// (/publishers/google/models/{model}:streamGenerateContent) but authenticate
// with simple API keys instead of Google Cloud service account credentials.
//
// Example services: zenmux.ai and similar Vertex-compatible providers.
type VertexCompatKey struct {
// APIKey is the authentication key for accessing the Vertex-compatible API.
// Maps to the x-goog-api-key header.
APIKey string `yaml:"api-key" json:"api-key"`
// BaseURL is the base URL for the Vertex-compatible API endpoint.
// The executor will append "/v1/publishers/google/models/{model}:action" to this.
// Example: "https://zenmux.ai/api" becomes "https://zenmux.ai/api/v1/publishers/google/models/..."
BaseURL string `yaml:"base-url,omitempty" json:"base-url,omitempty"`
// ProxyURL optionally overrides the global proxy for this API key.
ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"`
// Headers optionally adds extra HTTP headers for requests sent with this key.
// Commonly used for cookies, user-agent, and other authentication headers.
Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"`
// Models defines the model configurations including aliases for routing.
Models []VertexCompatModel `yaml:"models,omitempty" json:"models,omitempty"`
}
// VertexCompatModel represents a model configuration for Vertex compatibility,
// including the actual model name and its alias for API routing.
type VertexCompatModel struct {
// Name is the actual model name used by the external provider.
Name string `yaml:"name" json:"name"`
// Alias is the model name alias that clients will use to reference this model.
Alias string `yaml:"alias" json:"alias"`
}
// SanitizeVertexCompatKeys deduplicates and normalizes Vertex-compatible API key credentials.
func (cfg *Config) SanitizeVertexCompatKeys() {
if cfg == nil {
return
}
seen := make(map[string]struct{}, len(cfg.VertexCompatAPIKey))
out := cfg.VertexCompatAPIKey[:0]
for i := range cfg.VertexCompatAPIKey {
entry := cfg.VertexCompatAPIKey[i]
entry.APIKey = strings.TrimSpace(entry.APIKey)
if entry.APIKey == "" {
continue
}
entry.BaseURL = strings.TrimSpace(entry.BaseURL)
if entry.BaseURL == "" {
// BaseURL is required for Vertex API key entries
continue
}
entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
entry.Headers = NormalizeHeaders(entry.Headers)
// Sanitize models: remove entries without valid alias
sanitizedModels := make([]VertexCompatModel, 0, len(entry.Models))
for _, model := range entry.Models {
model.Alias = strings.TrimSpace(model.Alias)
model.Name = strings.TrimSpace(model.Name)
if model.Alias != "" && model.Name != "" {
sanitizedModels = append(sanitizedModels, model)
}
}
entry.Models = sanitizedModels
// Use API key + base URL as uniqueness key
uniqueKey := entry.APIKey + "|" + entry.BaseURL
if _, exists := seen[uniqueKey]; exists {
continue
}
seen[uniqueKey] = struct{}{}
out = append(out, entry)
}
cfg.VertexCompatAPIKey = out
}

View File

@@ -24,4 +24,7 @@ const (
// Antigravity represents the Antigravity response format identifier.
Antigravity = "antigravity"
// Kiro represents the AWS CodeWhisperer (Kiro) provider identifier.
Kiro = "kiro"
)

View File

@@ -56,6 +56,8 @@ type Content struct {
// Part represents a distinct piece of content within a message.
// A part can be text, inline data (like an image), a function call, or a function response.
type Part struct {
Thought bool `json:"thought,omitempty"`
// Text contains plain text content.
Text string `json:"text,omitempty"`
@@ -85,6 +87,9 @@ type InlineData struct {
// FunctionCall represents a tool call requested by the model.
// It includes the function name and its arguments that the model wants to execute.
type FunctionCall struct {
// ID is the identifier of the function to be called.
ID string `json:"id,omitempty"`
// Name is the identifier of the function to be called.
Name string `json:"name"`
@@ -95,6 +100,9 @@ type FunctionCall struct {
// FunctionResponse represents the result of a tool execution.
// This is sent back to the model after a tool call has been processed.
type FunctionResponse struct {
// ID is the identifier of the function to be called.
ID string `json:"id,omitempty"`
// Name is the identifier of the function that was called.
Name string `json:"name"`

View File

@@ -14,6 +14,8 @@ import (
log "github.com/sirupsen/logrus"
)
const skipGinLogKey = "__gin_skip_request_logging__"
// GinLogrusLogger returns a Gin middleware handler that logs HTTP requests and responses
// using logrus. It captures request details including method, path, status code, latency,
// client IP, and any error messages, formatting them in a Gin-style log format.
@@ -28,6 +30,10 @@ func GinLogrusLogger() gin.HandlerFunc {
c.Next()
if shouldSkipGinRequestLogging(c) {
return
}
if raw != "" {
path = path + "?" + raw
}
@@ -77,3 +83,24 @@ func GinLogrusRecovery() gin.HandlerFunc {
c.AbortWithStatus(http.StatusInternalServerError)
})
}
// SkipGinRequestLogging marks the provided Gin context so that GinLogrusLogger
// will skip emitting a log line for the associated request.
func SkipGinRequestLogging(c *gin.Context) {
if c == nil {
return
}
c.Set(skipGinLogKey, true)
}
func shouldSkipGinRequestLogging(c *gin.Context) bool {
if c == nil {
return false
}
val, exists := c.Get(skipGinLogKey)
if !exists {
return false
}
flag, ok := val.(bool)
return ok && flag
}

View File

@@ -38,7 +38,16 @@ func (m *LogFormatter) Format(entry *log.Entry) ([]byte, error) {
timestamp := entry.Time.Format("2006-01-02 15:04:05")
message := strings.TrimRight(entry.Message, "\r\n")
formatted := fmt.Sprintf("[%s] [%s] [%s:%d] %s\n", timestamp, entry.Level, filepath.Base(entry.Caller.File), entry.Caller.Line, message)
// Handle nil Caller (can happen with some log entries)
callerFile := "unknown"
callerLine := 0
if entry.Caller != nil {
callerFile = filepath.Base(entry.Caller.File)
callerLine = entry.Caller.Line
}
formatted := fmt.Sprintf("[%s] [%s] [%s:%d] %s\n", timestamp, entry.Level, callerFile, callerLine, message)
buffer.WriteString(formatted)
return buffer.Bytes(), nil
@@ -49,6 +58,7 @@ func (m *LogFormatter) Format(entry *log.Entry) ([]byte, error) {
func SetupBaseLogger() {
setupOnce.Do(func() {
log.SetOutput(os.Stdout)
log.SetLevel(log.InfoLevel)
log.SetReportCaller(true)
log.SetFormatter(&LogFormatter{})

View File

@@ -20,6 +20,7 @@ import (
"github.com/klauspost/compress/zstd"
log "github.com/sirupsen/logrus"
"github.com/router-for-me/CLIProxyAPI/v6/internal/buildinfo"
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
)
@@ -83,6 +84,26 @@ type StreamingLogWriter interface {
// - error: An error if writing fails, nil otherwise
WriteStatus(status int, headers map[string][]string) error
// WriteAPIRequest writes the upstream API request details to the log.
// This should be called before WriteStatus to maintain proper log ordering.
//
// Parameters:
// - apiRequest: The API request data (typically includes URL, headers, body sent upstream)
//
// Returns:
// - error: An error if writing fails, nil otherwise
WriteAPIRequest(apiRequest []byte) error
// WriteAPIResponse writes the upstream API response details to the log.
// This should be called after the streaming response is complete.
//
// Parameters:
// - apiResponse: The API response data
//
// Returns:
// - error: An error if writing fails, nil otherwise
WriteAPIResponse(apiResponse []byte) error
// Close finalizes the log file and cleans up resources.
//
// Returns:
@@ -247,10 +268,11 @@ func (l *FileRequestLogger) LogStreamingRequest(url, method string, headers map[
// Create streaming writer
writer := &FileStreamingLogWriter{
file: file,
chunkChan: make(chan []byte, 100), // Buffered channel for async writes
closeChan: make(chan struct{}),
errorChan: make(chan error, 1),
file: file,
chunkChan: make(chan []byte, 100), // Buffered channel for async writes
closeChan: make(chan struct{}),
errorChan: make(chan error, 1),
bufferedChunks: &bytes.Buffer{},
}
// Start async writer goroutine
@@ -603,6 +625,7 @@ func (l *FileRequestLogger) formatRequestInfo(url, method string, headers map[st
var content strings.Builder
content.WriteString("=== REQUEST INFO ===\n")
content.WriteString(fmt.Sprintf("Version: %s\n", buildinfo.Version))
content.WriteString(fmt.Sprintf("URL: %s\n", url))
content.WriteString(fmt.Sprintf("Method: %s\n", method))
content.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))
@@ -626,11 +649,12 @@ func (l *FileRequestLogger) formatRequestInfo(url, method string, headers map[st
// FileStreamingLogWriter implements StreamingLogWriter for file-based streaming logs.
// It handles asynchronous writing of streaming response chunks to a file.
// All data is buffered and written in the correct order when Close is called.
type FileStreamingLogWriter struct {
// file is the file where log data is written.
file *os.File
// chunkChan is a channel for receiving response chunks to write.
// chunkChan is a channel for receiving response chunks to buffer.
chunkChan chan []byte
// closeChan is a channel for signaling when the writer is closed.
@@ -639,8 +663,23 @@ type FileStreamingLogWriter struct {
// errorChan is a channel for reporting errors during writing.
errorChan chan error
// statusWritten indicates whether the response status has been written.
// bufferedChunks stores the response chunks in order.
bufferedChunks *bytes.Buffer
// responseStatus stores the HTTP status code.
responseStatus int
// statusWritten indicates whether a non-zero status was recorded.
statusWritten bool
// responseHeaders stores the response headers.
responseHeaders map[string][]string
// apiRequest stores the upstream API request data.
apiRequest []byte
// apiResponse stores the upstream API response data.
apiResponse []byte
}
// WriteChunkAsync writes a response chunk asynchronously (non-blocking).
@@ -664,39 +703,65 @@ func (w *FileStreamingLogWriter) WriteChunkAsync(chunk []byte) {
}
}
// WriteStatus writes the response status and headers to the log.
// WriteStatus buffers the response status and headers for later writing.
//
// Parameters:
// - status: The response status code
// - headers: The response headers
//
// Returns:
// - error: An error if writing fails, nil otherwise
// - error: Always returns nil (buffering cannot fail)
func (w *FileStreamingLogWriter) WriteStatus(status int, headers map[string][]string) error {
if w.file == nil || w.statusWritten {
if status == 0 {
return nil
}
var content strings.Builder
content.WriteString("========================================\n")
content.WriteString("=== RESPONSE ===\n")
content.WriteString(fmt.Sprintf("Status: %d\n", status))
for key, values := range headers {
for _, value := range values {
content.WriteString(fmt.Sprintf("%s: %s\n", key, value))
w.responseStatus = status
if headers != nil {
w.responseHeaders = make(map[string][]string, len(headers))
for key, values := range headers {
headerValues := make([]string, len(values))
copy(headerValues, values)
w.responseHeaders[key] = headerValues
}
}
content.WriteString("\n")
w.statusWritten = true
return nil
}
_, err := w.file.WriteString(content.String())
if err == nil {
w.statusWritten = true
// WriteAPIRequest buffers the upstream API request details for later writing.
//
// Parameters:
// - apiRequest: The API request data (typically includes URL, headers, body sent upstream)
//
// Returns:
// - error: Always returns nil (buffering cannot fail)
func (w *FileStreamingLogWriter) WriteAPIRequest(apiRequest []byte) error {
if len(apiRequest) == 0 {
return nil
}
return err
w.apiRequest = bytes.Clone(apiRequest)
return nil
}
// WriteAPIResponse buffers the upstream API response details for later writing.
//
// Parameters:
// - apiResponse: The API response data
//
// Returns:
// - error: Always returns nil (buffering cannot fail)
func (w *FileStreamingLogWriter) WriteAPIResponse(apiResponse []byte) error {
if len(apiResponse) == 0 {
return nil
}
w.apiResponse = bytes.Clone(apiResponse)
return nil
}
// Close finalizes the log file and cleans up resources.
// It writes all buffered data to the file in the correct order:
// API REQUEST -> API RESPONSE -> RESPONSE (status, headers, body chunks)
//
// Returns:
// - error: An error if closing fails, nil otherwise
@@ -705,27 +770,84 @@ func (w *FileStreamingLogWriter) Close() error {
close(w.chunkChan)
}
// Wait for async writer to finish
// Wait for async writer to finish buffering chunks
if w.closeChan != nil {
<-w.closeChan
w.chunkChan = nil
}
if w.file != nil {
return w.file.Close()
if w.file == nil {
return nil
}
return nil
// Write all content in the correct order
var content strings.Builder
// 1. Write API REQUEST section
if len(w.apiRequest) > 0 {
if bytes.HasPrefix(w.apiRequest, []byte("=== API REQUEST")) {
content.Write(w.apiRequest)
if !bytes.HasSuffix(w.apiRequest, []byte("\n")) {
content.WriteString("\n")
}
} else {
content.WriteString("=== API REQUEST ===\n")
content.Write(w.apiRequest)
content.WriteString("\n")
}
content.WriteString("\n")
}
// 2. Write API RESPONSE section
if len(w.apiResponse) > 0 {
if bytes.HasPrefix(w.apiResponse, []byte("=== API RESPONSE")) {
content.Write(w.apiResponse)
if !bytes.HasSuffix(w.apiResponse, []byte("\n")) {
content.WriteString("\n")
}
} else {
content.WriteString("=== API RESPONSE ===\n")
content.Write(w.apiResponse)
content.WriteString("\n")
}
content.WriteString("\n")
}
// 3. Write RESPONSE section (status, headers, buffered chunks)
content.WriteString("=== RESPONSE ===\n")
if w.statusWritten {
content.WriteString(fmt.Sprintf("Status: %d\n", w.responseStatus))
}
for key, values := range w.responseHeaders {
for _, value := range values {
content.WriteString(fmt.Sprintf("%s: %s\n", key, value))
}
}
content.WriteString("\n")
// Write buffered response body chunks
if w.bufferedChunks != nil && w.bufferedChunks.Len() > 0 {
content.Write(w.bufferedChunks.Bytes())
}
// Write the complete content to file
if _, err := w.file.WriteString(content.String()); err != nil {
_ = w.file.Close()
return err
}
return w.file.Close()
}
// asyncWriter runs in a goroutine to handle async chunk writing.
// It continuously reads chunks from the channel and writes them to the file.
// asyncWriter runs in a goroutine to buffer chunks from the channel.
// It continuously reads chunks from the channel and buffers them for later writing.
func (w *FileStreamingLogWriter) asyncWriter() {
defer close(w.closeChan)
for chunk := range w.chunkChan {
if w.file != nil {
_, _ = w.file.Write(chunk)
if w.bufferedChunks != nil {
w.bufferedChunks.Write(chunk)
}
}
}
@@ -752,6 +874,28 @@ func (w *NoOpStreamingLogWriter) WriteStatus(_ int, _ map[string][]string) error
return nil
}
// WriteAPIRequest is a no-op implementation that does nothing and always returns nil.
//
// Parameters:
// - apiRequest: The API request data (ignored)
//
// Returns:
// - error: Always returns nil
func (w *NoOpStreamingLogWriter) WriteAPIRequest(_ []byte) error {
return nil
}
// WriteAPIResponse is a no-op implementation that does nothing and always returns nil.
//
// Parameters:
// - apiResponse: The API response data (ignored)
//
// Returns:
// - error: Always returns nil
func (w *NoOpStreamingLogWriter) WriteAPIResponse(_ []byte) error {
return nil
}
// Close is a no-op implementation that does nothing and always returns nil.
//
// Returns:

View File

@@ -0,0 +1,368 @@
You are GPT-5.1 running in the Codex CLI, a terminal-based coding assistant. Codex CLI is an open source project led by OpenAI. You are expected to be precise, safe, and helpful.
Your capabilities:
- Receive user prompts and other context provided by the harness, such as files in the workspace.
- Communicate with the user by streaming thinking & responses, and by making & updating plans.
- Emit function calls to run terminal commands and apply patches. Depending on how this specific run is configured, you can request that these function calls be escalated to the user for approval before running. More on this in the "Sandbox and approvals" section.
Within this context, Codex refers to the open-source agentic coding interface (not the old Codex language model built by OpenAI).
# How you work
## Personality
Your default personality and tone is concise, direct, and friendly. You communicate efficiently, always keeping the user clearly informed about ongoing actions without unnecessary detail. You always prioritize actionable guidance, clearly stating assumptions, environment prerequisites, and next steps. Unless explicitly asked, you avoid excessively verbose explanations about your work.
# AGENTS.md spec
- Repos often contain AGENTS.md files. These files can appear anywhere within the repository.
- These files are a way for humans to give you (the agent) instructions or tips for working within the container.
- Some examples might be: coding conventions, info about how code is organized, or instructions for how to run or test code.
- Instructions in AGENTS.md files:
- The scope of an AGENTS.md file is the entire directory tree rooted at the folder that contains it.
- For every file you touch in the final patch, you must obey instructions in any AGENTS.md file whose scope includes that file.
- Instructions about code style, structure, naming, etc. apply only to code within the AGENTS.md file's scope, unless the file states otherwise.
- More-deeply-nested AGENTS.md files take precedence in the case of conflicting instructions.
- Direct system/developer/user instructions (as part of a prompt) take precedence over AGENTS.md instructions.
- The contents of the AGENTS.md file at the root of the repo and any directories from the CWD up to the root are included with the developer message and don't need to be re-read. When working in a subdirectory of CWD, or a directory outside the CWD, check for any AGENTS.md files that may be applicable.
## Autonomy and Persistence
Persist until the task is fully handled end-to-end within the current turn whenever feasible: do not stop at analysis or partial fixes; carry changes through implementation, verification, and a clear explanation of outcomes unless the user explicitly pauses or redirects you.
Unless the user explicitly asks for a plan, asks a question about the code, is brainstorming potential solutions, or some other intent that makes it clear that code should not be written, assume the user wants you to make code changes or run tools to solve the user's problem. In these cases, it's bad to output your proposed solution in a message, you should go ahead and actually implement the change. If you encounter challenges or blockers, you should attempt to resolve them yourself.
## Responsiveness
### User Updates Spec
You'll work for stretches with tool calls — it's critical to keep the user updated as you work.
Frequency & Length:
- Send short updates (12 sentences) whenever there is a meaningful, important insight you need to share with the user to keep them informed.
- If you expect a longer headsdown stretch, post a brief headsdown note with why and when you'll report back; when you resume, summarize what you learned.
- Only the initial plan, plan updates, and final recap can be longer, with multiple bullets and paragraphs
Tone:
- Friendly, confident, senior-engineer energy. Positive, collaborative, humble; fix mistakes quickly.
Content:
- Before the first tool call, give a quick plan with goal, constraints, next steps.
- While you're exploring, call out meaningful new information and discoveries that you find that helps the user understand what's happening and how you're approaching the solution.
- If you change the plan (e.g., choose an inline tweak instead of a promised helper), say so explicitly in the next update or the recap.
**Examples:**
- “Ive explored the repo; now checking the API route definitions.”
- “Next, Ill patch the config and update the related tests.”
- “Im about to scaffold the CLI commands and helper functions.”
- “Ok cool, so Ive wrapped my head around the repo. Now digging into the API routes.”
- “Configs looking tidy. Next up is patching helpers to keep things in sync.”
- “Finished poking at the DB gateway. I will now chase down error handling.”
- “Alright, build pipeline order is interesting. Checking how it reports failures.”
- “Spotted a clever caching util; now hunting where it gets used.”
## Planning
You have access to an `update_plan` tool which tracks steps and progress and renders them to the user. Using the tool helps demonstrate that you've understood the task and convey how you're approaching it. Plans can help to make complex, ambiguous, or multi-phase work clearer and more collaborative for the user. A good plan should break the task into meaningful, logically ordered steps that are easy to verify as you go.
Note that plans are not for padding out simple work with filler steps or stating the obvious. The content of your plan should not involve doing anything that you aren't capable of doing (i.e. don't try to test things that you can't test). Do not use plans for simple or single-step queries that you can just do or answer immediately.
Do not repeat the full contents of the plan after an `update_plan` call — the harness already displays it. Instead, summarize the change made and highlight any important context or next step.
Before running a command, consider whether or not you have completed the previous step, and make sure to mark it as completed before moving on to the next step. It may be the case that you complete all steps in your plan after a single pass of implementation. If this is the case, you can simply mark all the planned steps as completed. Sometimes, you may need to change plans in the middle of a task: call `update_plan` with the updated plan and make sure to provide an `explanation` of the rationale when doing so.
Maintain statuses in the tool: exactly one item in_progress at a time; mark items complete when done; post timely status transitions. Do not jump an item from pending to completed: always set it to in_progress first. Do not batch-complete multiple items after the fact. Finish with all items completed or explicitly canceled/deferred before ending the turn. Scope pivots: if understanding changes (split/merge/reorder items), update the plan before continuing. Do not let the plan go stale while coding.
Use a plan when:
- The task is non-trivial and will require multiple actions over a long time horizon.
- There are logical phases or dependencies where sequencing matters.
- The work has ambiguity that benefits from outlining high-level goals.
- You want intermediate checkpoints for feedback and validation.
- When the user asked you to do more than one thing in a single prompt
- The user has asked you to use the plan tool (aka "TODOs")
- You generate additional steps while working, and plan to do them before yielding to the user
### Examples
**High-quality plans**
Example 1:
1. Add CLI entry with file args
2. Parse Markdown via CommonMark library
3. Apply semantic HTML template
4. Handle code blocks, images, links
5. Add error handling for invalid files
Example 2:
1. Define CSS variables for colors
2. Add toggle with localStorage state
3. Refactor components to use variables
4. Verify all views for readability
5. Add smooth theme-change transition
Example 3:
1. Set up Node.js + WebSocket server
2. Add join/leave broadcast events
3. Implement messaging with timestamps
4. Add usernames + mention highlighting
5. Persist messages in lightweight DB
6. Add typing indicators + unread count
**Low-quality plans**
Example 1:
1. Create CLI tool
2. Add Markdown parser
3. Convert to HTML
Example 2:
1. Add dark mode toggle
2. Save preference
3. Make styles look good
Example 3:
1. Create single-file HTML game
2. Run quick sanity check
3. Summarize usage instructions
If you need to write a plan, only write high quality plans, not low quality ones.
## Task execution
You are a coding agent. You must keep going until the query or task is completely resolved, before ending your turn and yielding back to the user. Persist until the task is fully handled end-to-end within the current turn whenever feasible and persevere even when function calls fail. Only terminate your turn when you are sure that the problem is solved. Autonomously resolve the query to the best of your ability, using the tools available to you, before coming back to the user. Do NOT guess or make up an answer.
You MUST adhere to the following criteria when solving queries:
- Working on the repo(s) in the current environment is allowed, even if they are proprietary.
- Analyzing code for vulnerabilities is allowed.
- Showing user code and tool call details is allowed.
- Use the `apply_patch` tool to edit files (NEVER try `applypatch` or `apply-patch`, only `apply_patch`). This is a FREEFORM tool, so do not wrap the patch in JSON.
If completing the user's task requires writing or modifying files, your code and final answer should follow these coding guidelines, though user instructions (i.e. AGENTS.md) may override these guidelines:
- Fix the problem at the root cause rather than applying surface-level patches, when possible.
- Avoid unneeded complexity in your solution.
- Do not attempt to fix unrelated bugs or broken tests. It is not your responsibility to fix them. (You may mention them to the user in your final message though.)
- Update documentation as necessary.
- Keep changes consistent with the style of the existing codebase. Changes should be minimal and focused on the task.
- Use `git log` and `git blame` to search the history of the codebase if additional context is required.
- NEVER add copyright or license headers unless specifically requested.
- Do not waste tokens by re-reading files after calling `apply_patch` on them. The tool call will fail if it didn't work. The same goes for making folders, deleting folders, etc.
- Do not `git commit` your changes or create new git branches unless explicitly requested.
- Do not add inline comments within code unless explicitly requested.
- Do not use one-letter variable names unless explicitly requested.
- NEVER output inline citations like "【F:README.md†L5-L14】" in your outputs. The CLI is not able to render these so they will just be broken in the UI. Instead, if you output valid filepaths, users will be able to click on them to open the files in their editor.
## Codex CLI harness, sandboxing, and approvals
The Codex CLI harness supports several different configurations for sandboxing and escalation approvals that the user can choose from.
Filesystem sandboxing defines which files can be read or written. The options for `sandbox_mode` are:
- **read-only**: The sandbox only permits reading files.
- **workspace-write**: The sandbox permits reading files, and editing files in `cwd` and `writable_roots`. Editing files in other directories requires approval.
- **danger-full-access**: No filesystem sandboxing - all commands are permitted.
Network sandboxing defines whether network can be accessed without approval. Options for `network_access` are:
- **restricted**: Requires approval
- **enabled**: No approval needed
Approvals are your mechanism to get user consent to run shell commands without the sandbox. Possible configuration options for `approval_policy` are
- **untrusted**: The harness will escalate most commands for user approval, apart from a limited allowlist of safe "read" commands.
- **on-failure**: The harness will allow all commands to run in the sandbox (if enabled), and failures will be escalated to the user for approval to run again without the sandbox.
- **on-request**: Commands will be run in the sandbox by default, and you can specify in your tool call if you want to escalate a command to run without sandboxing. (Note that this mode is not always available. If it is, you'll see parameters for escalating in the tool definition.)
- **never**: This is a non-interactive mode where you may NEVER ask the user for approval to run commands. Instead, you must always persist and work around constraints to solve the task for the user. You MUST do your utmost best to finish the task and validate your work before yielding. If this mode is paired with `danger-full-access`, take advantage of it to deliver the best outcome for the user. Further, in this mode, your default testing philosophy is overridden: Even if you don't see local patterns for testing, you may add tests and scripts to validate your work. Just remove them before yielding.
When you are running with `approval_policy == on-request`, and sandboxing enabled, here are scenarios where you'll need to request approval:
- You need to run a command that writes to a directory that requires it (e.g. running tests that write to /var)
- You need to run a GUI app (e.g., open/xdg-open/osascript) to open browsers or files.
- You are running sandboxed and need to run a command that requires network access (e.g. installing packages)
- If you run a command that is important to solving the user's query, but it fails because of sandboxing, rerun the command with approval. ALWAYS proceed to use the `with_escalated_permissions` and `justification` parameters. Within this harness, prefer requesting approval via the tool over asking in natural language.
- You are about to take a potentially destructive action such as an `rm` or `git reset` that the user did not explicitly ask for
- (for all of these, you should weigh alternative paths that do not require approval)
When `sandbox_mode` is set to read-only, you'll need to request approval for any command that isn't a read.
You will be told what filesystem sandboxing, network sandboxing, and approval mode are active in a developer or user message. If you are not told about this, assume that you are running with workspace-write, network sandboxing enabled, and approval on-failure.
Although they introduce friction to the user because your work is paused until the user responds, you should leverage them when necessary to accomplish important work. If the completing the task requires escalated permissions, Do not let these settings or the sandbox deter you from attempting to accomplish the user's task unless it is set to "never", in which case never ask for approvals.
When requesting approval to execute a command that will require escalated privileges:
- Provide the `with_escalated_permissions` parameter with the boolean value true
- Include a short, 1 sentence explanation for why you need to enable `with_escalated_permissions` in the justification parameter
## Validating your work
If the codebase has tests or the ability to build or run, consider using them to verify changes once your work is complete.
When testing, your philosophy should be to start as specific as possible to the code you changed so that you can catch issues efficiently, then make your way to broader tests as you build confidence. If there's no test for the code you changed, and if the adjacent patterns in the codebases show that there's a logical place for you to add a test, you may do so. However, do not add tests to codebases with no tests.
Similarly, once you're confident in correctness, you can suggest or use formatting commands to ensure that your code is well formatted. If there are issues you can iterate up to 3 times to get formatting right, but if you still can't manage it's better to save the user time and present them a correct solution where you call out the formatting in your final message. If the codebase does not have a formatter configured, do not add one.
For all of testing, running, building, and formatting, do not attempt to fix unrelated bugs. It is not your responsibility to fix them. (You may mention them to the user in your final message though.)
Be mindful of whether to run validation commands proactively. In the absence of behavioral guidance:
- When running in non-interactive approval modes like **never** or **on-failure**, you can proactively run tests, lint and do whatever you need to ensure you've completed the task. If you are unable to run tests, you must still do your utmost best to complete the task.
- When working in interactive approval modes like **untrusted**, or **on-request**, hold off on running tests or lint commands until the user is ready for you to finalize your output, because these commands take time to run and slow down iteration. Instead suggest what you want to do next, and let the user confirm first.
- When working on test-related tasks, such as adding tests, fixing tests, or reproducing a bug to verify behavior, you may proactively run tests regardless of approval mode. Use your judgement to decide whether this is a test-related task.
## Ambition vs. precision
For tasks that have no prior context (i.e. the user is starting something brand new), you should feel free to be ambitious and demonstrate creativity with your implementation.
If you're operating in an existing codebase, you should make sure you do exactly what the user asks with surgical precision. Treat the surrounding codebase with respect, and don't overstep (i.e. changing filenames or variables unnecessarily). You should balance being sufficiently ambitious and proactive when completing tasks of this nature.
You should use judicious initiative to decide on the right level of detail and complexity to deliver based on the user's needs. This means showing good judgment that you're capable of doing the right extras without gold-plating. This might be demonstrated by high-value, creative touches when scope of the task is vague; while being surgical and targeted when scope is tightly specified.
## Sharing progress updates
For especially longer tasks that you work on (i.e. requiring many tool calls, or a plan with multiple steps), you should provide progress updates back to the user at reasonable intervals. These updates should be structured as a concise sentence or two (no more than 8-10 words long) recapping progress so far in plain language: this update demonstrates your understanding of what needs to be done, progress so far (i.e. files explores, subtasks complete), and where you're going next.
Before doing large chunks of work that may incur latency as experienced by the user (i.e. writing a new file), you should send a concise message to the user with an update indicating what you're about to do to ensure they know what you're spending time on. Don't start editing or writing large files before informing the user what you are doing and why.
The messages you send before tool calls should describe what is immediately about to be done next in very concise language. If there was previous work done, this preamble message should also include a note about the work done so far to bring the user along.
## Presenting your work and final message
Your final message should read naturally, like an update from a concise teammate. For casual conversation, brainstorming tasks, or quick questions from the user, respond in a friendly, conversational tone. You should ask questions, suggest ideas, and adapt to the users style. If you've finished a large amount of work, when describing what you've done to the user, you should follow the final answer formatting guidelines to communicate substantive changes. You don't need to add structured formatting for one-word answers, greetings, or purely conversational exchanges.
You can skip heavy formatting for single, simple actions or confirmations. In these cases, respond in plain sentences with any relevant next step or quick option. Reserve multi-section structured responses for results that need grouping or explanation.
The user is working on the same computer as you, and has access to your work. As such there's no need to show the contents of files you have already written unless the user explicitly asks for them. Similarly, if you've created or modified files using `apply_patch`, there's no need to tell users to "save the file" or "copy the code into a file"—just reference the file path.
If there's something that you think you could help with as a logical next step, concisely ask the user if they want you to do so. Good examples of this are running tests, committing changes, or building out the next logical component. If theres something that you couldn't do (even with approval) but that the user might want to do (such as verifying changes by running the app), include those instructions succinctly.
Brevity is very important as a default. You should be very concise (i.e. no more than 10 lines), but can relax this requirement for tasks where additional detail and comprehensiveness is important for the user's understanding.
### Final answer structure and style guidelines
You are producing plain text that will later be styled by the CLI. Follow these rules exactly. Formatting should make results easy to scan, but not feel mechanical. Use judgment to decide how much structure adds value.
**Section Headers**
- Use only when they improve clarity — they are not mandatory for every answer.
- Choose descriptive names that fit the content
- Keep headers short (13 words) and in `**Title Case**`. Always start headers with `**` and end with `**`
- Leave no blank line before the first bullet under a header.
- Section headers should only be used where they genuinely improve scanability; avoid fragmenting the answer.
**Bullets**
- Use `-` followed by a space for every bullet.
- Merge related points when possible; avoid a bullet for every trivial detail.
- Keep bullets to one line unless breaking for clarity is unavoidable.
- Group into short lists (46 bullets) ordered by importance.
- Use consistent keyword phrasing and formatting across sections.
**Monospace**
- Wrap all commands, file paths, env vars, code identifiers, and code samples in backticks (`` `...` ``).
- Apply to inline examples and to bullet keywords if the keyword itself is a literal file/command.
- Never mix monospace and bold markers; choose one based on whether its a keyword (`**`) or inline code/path (`` ` ``).
**File References**
When referencing files in your response, make sure to include the relevant start line and always follow the below rules:
* Use inline code to make file paths clickable.
* Each reference should have a stand alone path. Even if it's the same file.
* Accepted: absolute, workspacerelative, a/ or b/ diff prefixes, or bare filename/suffix.
* Line/column (1based, optional): :line[:column] or #Lline[Ccolumn] (column defaults to 1).
* Do not use URIs like file://, vscode://, or https://.
* Do not provide range of lines
* Examples: src/app.ts, src/app.ts:42, b/server/index.js#L10, C:\repo\project\main.rs:12:5
**Structure**
- Place related bullets together; dont mix unrelated concepts in the same section.
- Order sections from general → specific → supporting info.
- For subsections (e.g., “Binaries” under “Rust Workspace”), introduce with a bolded keyword bullet, then list items under it.
- Match structure to complexity:
- Multi-part or detailed results → use clear headers and grouped bullets.
- Simple results → minimal headers, possibly just a short list or paragraph.
**Tone**
- Keep the voice collaborative and natural, like a coding partner handing off work.
- Be concise and factual — no filler or conversational commentary and avoid unnecessary repetition
- Use present tense and active voice (e.g., “Runs tests” not “This will run tests”).
- Keep descriptions self-contained; dont refer to “above” or “below”.
- Use parallel structure in lists for consistency.
**Verbosity**
- Final answer compactness rules (enforced):
- Tiny/small single-file change (≤ ~10 lines): 25 sentences or ≤3 bullets. No headings. 01 short snippet (≤3 lines) only if essential.
- Medium change (single area or a few files): ≤6 bullets or 610 sentences. At most 12 short snippets total (≤8 lines each).
- Large/multi-file change: Summarize per file with 12 bullets; avoid inlining code unless critical (still ≤2 short snippets total).
- Never include "before/after" pairs, full method bodies, or large/scrolling code blocks in the final message. Prefer referencing file/symbol names instead.
**Dont**
- Dont use literal words “bold” or “monospace” in the content.
- Dont nest bullets or create deep hierarchies.
- Dont output ANSI escape codes directly — the CLI renderer applies them.
- Dont cram unrelated keywords into a single bullet; split for clarity.
- Dont let keyword lists run long — wrap or reformat for scanability.
Generally, ensure your final answers adapt their shape and depth to the request. For example, answers to code explanations should have a precise, structured explanation with code references that answer the question directly. For tasks with a simple implementation, lead with the outcome and supplement only with whats needed for clarity. Larger changes can be presented as a logical walkthrough of your approach, grouping related steps, explaining rationale where it adds value, and highlighting next actions to accelerate the user. Your answers should provide the right level of detail while being easily scannable.
For casual greetings, acknowledgements, or other one-off conversational messages that are not delivering substantive information or structured results, respond naturally without section headers or bullet formatting.
# Tool Guidelines
## Shell commands
When using the shell, you must adhere to the following guidelines:
- When searching for text or files, prefer using `rg` or `rg --files` respectively because `rg` is much faster than alternatives like `grep`. (If the `rg` command is not found, then use alternatives.)
- Read files in chunks with a max chunk size of 250 lines. Do not use python scripts to attempt to output larger chunks of a file. Command line output will be truncated after 10 kilobytes or 256 lines of output, regardless of the command used.
## apply_patch
Use the `apply_patch` tool to edit files. Your patch language is a strippeddown, fileoriented diff format designed to be easy to parse and safe to apply. You can think of it as a highlevel envelope:
*** Begin Patch
[ one or more file sections ]
*** End Patch
Within that envelope, you get a sequence of file operations.
You MUST include a header to specify the action you are taking.
Each operation starts with one of three headers:
*** Add File: <path> - create a new file. Every following line is a + line (the initial contents).
*** Delete File: <path> - remove an existing file. Nothing follows.
*** Update File: <path> - patch an existing file in place (optionally with a rename).
Example patch:
```
*** Begin Patch
*** Add File: hello.txt
+Hello world
*** Update File: src/app.py
*** Move to: src/main.py
@@ def greet():
-print("Hi")
+print("Hello, world!")
*** Delete File: obsolete.txt
*** End Patch
```
It is important to remember:
- You must include a header with your intended action (Add/Delete/Update)
- You must prefix new lines with `+` even when creating a new file
## `update_plan`
A tool named `update_plan` is available to you. You can use it to keep an uptodate, stepbystep plan for the task.
To create a new plan, call `update_plan` with a short list of 1sentence steps (no more than 5-7 words each) with a `status` for each step (`pending`, `in_progress`, or `completed`).
When steps have been completed, use `update_plan` to mark each finished step as `completed` and the next step you are working on as `in_progress`. There should always be exactly one `in_progress` step until everything is done. You can mark multiple items as complete in a single `update_plan` call.
If all steps are complete, ensure you call `update_plan` to mark all steps as `completed`.

View File

@@ -0,0 +1,105 @@
You are Codex, based on GPT-5. You are running as a coding agent in the Codex CLI on a user's computer.
## General
- When searching for text or files, prefer using `rg` or `rg --files` respectively because `rg` is much faster than alternatives like `grep`. (If the `rg` command is not found, then use alternatives.)
## Editing constraints
- Default to ASCII when editing or creating files. Only introduce non-ASCII or other Unicode characters when there is a clear justification and the file already uses them.
- Add succinct code comments that explain what is going on if code is not self-explanatory. You should not add comments like "Assigns the value to the variable", but a brief comment might be useful ahead of a complex code block that the user would otherwise have to spend time parsing out. Usage of these comments should be rare.
- Try to use apply_patch for single file edits, but it is fine to explore other options to make the edit if it does not work well. Do not use apply_patch for changes that are auto-generated (i.e. generating package.json or running a lint or format command like gofmt) or when scripting is more efficient (such as search and replacing a string across a codebase).
- You may be in a dirty git worktree.
* NEVER revert existing changes you did not make unless explicitly requested, since these changes were made by the user.
* If asked to make a commit or code edits and there are unrelated changes to your work or changes that you didn't make in those files, don't revert those changes.
* If the changes are in files you've touched recently, you should read carefully and understand how you can work with the changes rather than reverting them.
* If the changes are in unrelated files, just ignore them and don't revert them.
- Do not amend a commit unless explicitly requested to do so.
- While you are working, you might notice unexpected changes that you didn't make. If this happens, STOP IMMEDIATELY and ask the user how they would like to proceed.
- **NEVER** use destructive commands like `git reset --hard` or `git checkout --` unless specifically requested or approved by the user.
## Plan tool
When using the planning tool:
- Skip using the planning tool for straightforward tasks (roughly the easiest 25%).
- Do not make single-step plans.
- When you made a plan, update it after having performed one of the sub-tasks that you shared on the plan.
## Codex CLI harness, sandboxing, and approvals
The Codex CLI harness supports several different configurations for sandboxing and escalation approvals that the user can choose from.
Filesystem sandboxing defines which files can be read or written. The options for `sandbox_mode` are:
- **read-only**: The sandbox only permits reading files.
- **workspace-write**: The sandbox permits reading files, and editing files in `cwd` and `writable_roots`. Editing files in other directories requires approval.
- **danger-full-access**: No filesystem sandboxing - all commands are permitted.
Network sandboxing defines whether network can be accessed without approval. Options for `network_access` are:
- **restricted**: Requires approval
- **enabled**: No approval needed
Approvals are your mechanism to get user consent to run shell commands without the sandbox. Possible configuration options for `approval_policy` are
- **untrusted**: The harness will escalate most commands for user approval, apart from a limited allowlist of safe "read" commands.
- **on-failure**: The harness will allow all commands to run in the sandbox (if enabled), and failures will be escalated to the user for approval to run again without the sandbox.
- **on-request**: Commands will be run in the sandbox by default, and you can specify in your tool call if you want to escalate a command to run without sandboxing. (Note that this mode is not always available. If it is, you'll see parameters for it in the `shell` command description.)
- **never**: This is a non-interactive mode where you may NEVER ask the user for approval to run commands. Instead, you must always persist and work around constraints to solve the task for the user. You MUST do your utmost best to finish the task and validate your work before yielding. If this mode is paired with `danger-full-access`, take advantage of it to deliver the best outcome for the user. Further, in this mode, your default testing philosophy is overridden: Even if you don't see local patterns for testing, you may add tests and scripts to validate your work. Just remove them before yielding.
When you are running with `approval_policy == on-request`, and sandboxing enabled, here are scenarios where you'll need to request approval:
- You need to run a command that writes to a directory that requires it (e.g. running tests that write to /var)
- You need to run a GUI app (e.g., open/xdg-open/osascript) to open browsers or files.
- You are running sandboxed and need to run a command that requires network access (e.g. installing packages)
- If you run a command that is important to solving the user's query, but it fails because of sandboxing, rerun the command with approval. ALWAYS proceed to use the `with_escalated_permissions` and `justification` parameters - do not message the user before requesting approval for the command.
- You are about to take a potentially destructive action such as an `rm` or `git reset` that the user did not explicitly ask for
- (for all of these, you should weigh alternative paths that do not require approval)
When `sandbox_mode` is set to read-only, you'll need to request approval for any command that isn't a read.
You will be told what filesystem sandboxing, network sandboxing, and approval mode are active in a developer or user message. If you are not told about this, assume that you are running with workspace-write, network sandboxing enabled, and approval on-failure.
Although they introduce friction to the user because your work is paused until the user responds, you should leverage them when necessary to accomplish important work. If the completing the task requires escalated permissions, Do not let these settings or the sandbox deter you from attempting to accomplish the user's task unless it is set to "never", in which case never ask for approvals.
When requesting approval to execute a command that will require escalated privileges:
- Provide the `with_escalated_permissions` parameter with the boolean value true
- Include a short, 1 sentence explanation for why you need to enable `with_escalated_permissions` in the justification parameter
## Special user requests
- If the user makes a simple request (such as asking for the time) which you can fulfill by running a terminal command (such as `date`), you should do so.
- If the user asks for a "review", default to a code review mindset: prioritise identifying bugs, risks, behavioural regressions, and missing tests. Findings must be the primary focus of the response - keep summaries or overviews brief and only after enumerating the issues. Present findings first (ordered by severity with file/line references), follow with open questions or assumptions, and offer a change-summary only as a secondary detail. If no findings are discovered, state that explicitly and mention any residual risks or testing gaps.
## Presenting your work and final message
You are producing plain text that will later be styled by the CLI. Follow these rules exactly. Formatting should make results easy to scan, but not feel mechanical. Use judgment to decide how much structure adds value.
- Default: be very concise; friendly coding teammate tone.
- Ask only when needed; suggest ideas; mirror the user's style.
- For substantial work, summarize clearly; follow finalanswer formatting.
- Skip heavy formatting for simple confirmations.
- Don't dump large files you've written; reference paths only.
- No "save/copy this file" - User is on the same machine.
- Offer logical next steps (tests, commits, build) briefly; add verify steps if you couldn't do something.
- For code changes:
* Lead with a quick explanation of the change, and then give more details on the context covering where and why a change was made. Do not start this explanation with "summary", just jump right in.
* If there are natural next steps the user may want to take, suggest them at the end of your response. Do not make suggestions if there are no natural next steps.
* When suggesting multiple options, use numeric lists for the suggestions so the user can quickly respond with a single number.
- The user does not command execution outputs. When asked to show the output of a command (e.g. `git show`), relay the important details in your answer or summarize the key lines so the user understands the result.
### Final answer structure and style guidelines
- Plain text; CLI handles styling. Use structure only when it helps scanability.
- Headers: optional; short Title Case (1-3 words) wrapped in **…**; no blank line before the first bullet; add only if they truly help.
- Bullets: use - ; merge related points; keep to one line when possible; 46 per list ordered by importance; keep phrasing consistent.
- Monospace: backticks for commands/paths/env vars/code ids and inline examples; use for literal keyword bullets; never combine with **.
- Code samples or multi-line snippets should be wrapped in fenced code blocks; include an info string as often as possible.
- Structure: group related bullets; order sections general → specific → supporting; for subsections, start with a bolded keyword bullet, then items; match complexity to the task.
- Tone: collaborative, concise, factual; present tense, active voice; selfcontained; no "above/below"; parallel wording.
- Don'ts: no nested bullets/hierarchies; no ANSI codes; don't cram unrelated keywords; keep keyword lists short—wrap/reformat if long; avoid naming formatting styles in answers.
- Adaptation: code explanations → precise, structured with code refs; simple tasks → lead with outcome; big changes → logical walkthrough + rationale + next actions; casual one-offs → plain sentences, no headers/bullets.
- File References: When referencing files in your response, make sure to include the relevant start line and always follow the below rules:
* Use inline code to make file paths clickable.
* Each reference should have a stand alone path. Even if it's the same file.
* Accepted: absolute, workspacerelative, a/ or b/ diff prefixes, or bare filename/suffix.
* Line/column (1based, optional): :line[:column] or #Lline[Ccolumn] (column defaults to 1).
* Do not use URIs like file://, vscode://, or https://.
* Do not provide range of lines
* Examples: src/app.ts, src/app.ts:42, b/server/index.js#L10, C:\repo\project\main.rs:12:5

View File

@@ -222,6 +222,7 @@ func GetGeminiModels() []*ModelInfo {
InputTokenLimit: 1048576,
OutputTokenLimit: 65536,
SupportedGenerationMethods: []string{"generateContent", "countTokens", "createCachedContent", "batchGenerateContent"},
Thinking: &ThinkingSupport{Min: 128, Max: 32768, ZeroAllowed: false, DynamicAllowed: true},
},
}
}
@@ -301,6 +302,7 @@ func GetGeminiVertexModels() []*ModelInfo {
InputTokenLimit: 1048576,
OutputTokenLimit: 65536,
SupportedGenerationMethods: []string{"generateContent", "countTokens", "createCachedContent", "batchGenerateContent"},
Thinking: &ThinkingSupport{Min: 128, Max: 32768, ZeroAllowed: false, DynamicAllowed: true},
},
}
}
@@ -942,7 +944,6 @@ func GetQwenModels() []*ModelInfo {
}
// GetIFlowModels returns supported models for iFlow OAuth accounts.
func GetIFlowModels() []*ModelInfo {
entries := []struct {
ID string
@@ -952,7 +953,6 @@ func GetIFlowModels() []*ModelInfo {
}{
{ID: "tstars2.0", DisplayName: "TStars-2.0", Description: "iFlow TStars-2.0 multimodal assistant", Created: 1746489600},
{ID: "qwen3-coder-plus", DisplayName: "Qwen3-Coder-Plus", Description: "Qwen3 Coder Plus code generation", Created: 1753228800},
{ID: "qwen3-coder", DisplayName: "Qwen3-Coder-480B-A35B", Description: "Qwen3 Coder 480B A35B", Created: 1753228800},
{ID: "qwen3-max", DisplayName: "Qwen3-Max", Description: "Qwen3 flagship model", Created: 1758672000},
{ID: "qwen3-vl-plus", DisplayName: "Qwen3-VL-Plus", Description: "Qwen3 multimodal vision-language", Created: 1758672000},
{ID: "qwen3-max-preview", DisplayName: "Qwen3-Max-Preview", Description: "Qwen3 Max preview build", Created: 1757030400},
@@ -960,6 +960,7 @@ func GetIFlowModels() []*ModelInfo {
{ID: "glm-4.6", DisplayName: "GLM-4.6", Description: "Zhipu GLM 4.6 general model", Created: 1759190400},
{ID: "kimi-k2", DisplayName: "Kimi-K2", Description: "Moonshot Kimi K2 general model", Created: 1752192000},
{ID: "kimi-k2-thinking", DisplayName: "Kimi-K2-Thinking", Description: "Moonshot Kimi K2 general model", Created: 1762387200},
{ID: "deepseek-v3.2-chat", DisplayName: "DeepSeek-V3.2", Description: "DeepSeek V3.2", Created: 1764576000},
{ID: "deepseek-v3.2", DisplayName: "DeepSeek-V3.2-Exp", Description: "DeepSeek V3.2 experimental", Created: 1759104000},
{ID: "deepseek-v3.1", DisplayName: "DeepSeek-V3.1-Terminus", Description: "DeepSeek V3.1 Terminus", Created: 1756339200},
{ID: "deepseek-r1", DisplayName: "DeepSeek-R1", Description: "DeepSeek reasoning model R1", Created: 1737331200},
@@ -985,6 +986,25 @@ func GetIFlowModels() []*ModelInfo {
return models
}
// AntigravityModelConfig captures static antigravity model overrides, including
// Thinking budget limits and provider max completion tokens.
type AntigravityModelConfig struct {
Thinking *ThinkingSupport
MaxCompletionTokens int
}
// GetAntigravityModelConfig returns static configuration for antigravity models.
// Keys use the ALIASED model names (after modelName2Alias conversion) for direct lookup.
func GetAntigravityModelConfig() map[string]*AntigravityModelConfig {
return map[string]*AntigravityModelConfig{
"gemini-2.5-flash": {Thinking: &ThinkingSupport{Min: 0, Max: 24576, ZeroAllowed: true, DynamicAllowed: true}},
"gemini-2.5-flash-lite": {Thinking: &ThinkingSupport{Min: 0, Max: 24576, ZeroAllowed: true, DynamicAllowed: true}},
"gemini-3-pro-preview": {Thinking: &ThinkingSupport{Min: 128, Max: 32768, ZeroAllowed: false, DynamicAllowed: true}},
"gemini-claude-sonnet-4-5-thinking": {Thinking: &ThinkingSupport{Min: 1024, Max: 200000, ZeroAllowed: false, DynamicAllowed: true}, MaxCompletionTokens: 64000},
"gemini-claude-opus-4-5-thinking": {Thinking: &ThinkingSupport{Min: 1024, Max: 200000, ZeroAllowed: false, DynamicAllowed: true}, MaxCompletionTokens: 64000},
}
}
// GetGitHubCopilotModels returns the available models for GitHub Copilot.
// These models are available through the GitHub Copilot API at api.githubcopilot.com.
func GetGitHubCopilotModels() []*ModelInfo {
@@ -1168,3 +1188,161 @@ func GetGitHubCopilotModels() []*ModelInfo {
},
}
}
// GetKiroModels returns the Kiro (AWS CodeWhisperer) model definitions
func GetKiroModels() []*ModelInfo {
return []*ModelInfo{
{
ID: "kiro-auto",
Object: "model",
Created: 1732752000, // 2024-11-28
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Auto",
Description: "Automatic model selection by AWS CodeWhisperer",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "kiro-claude-opus-4.5",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Opus 4.5",
Description: "Claude Opus 4.5 via Kiro (2.2x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "kiro-claude-sonnet-4.5",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Sonnet 4.5",
Description: "Claude Sonnet 4.5 via Kiro (1.3x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "kiro-claude-sonnet-4",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Sonnet 4",
Description: "Claude Sonnet 4 via Kiro (1.3x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "kiro-claude-haiku-4.5",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Haiku 4.5",
Description: "Claude Haiku 4.5 via Kiro (0.4x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
// --- Chat Variant (No tool calling, for pure conversation) ---
{
ID: "kiro-claude-opus-4.5-chat",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Opus 4.5 (Chat)",
Description: "Claude Opus 4.5 for chat only (no tool calling)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
// --- Agentic Variants (Optimized for coding agents with chunked writes) ---
{
ID: "kiro-claude-opus-4.5-agentic",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Opus 4.5 (Agentic)",
Description: "Claude Opus 4.5 optimized for coding agents (chunked writes)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "kiro-claude-sonnet-4.5-agentic",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Sonnet 4.5 (Agentic)",
Description: "Claude Sonnet 4.5 optimized for coding agents (chunked writes)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
}
}
// GetAmazonQModels returns the Amazon Q (AWS CodeWhisperer) model definitions.
// These models use the same API as Kiro and share the same executor.
func GetAmazonQModels() []*ModelInfo {
return []*ModelInfo{
{
ID: "amazonq-auto",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro", // Uses Kiro executor - same API
DisplayName: "Amazon Q Auto",
Description: "Automatic model selection by Amazon Q",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "amazonq-claude-opus-4.5",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Amazon Q Claude Opus 4.5",
Description: "Claude Opus 4.5 via Amazon Q (2.2x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "amazonq-claude-sonnet-4.5",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Amazon Q Claude Sonnet 4.5",
Description: "Claude Sonnet 4.5 via Amazon Q (1.3x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "amazonq-claude-sonnet-4",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Amazon Q Claude Sonnet 4",
Description: "Claude Sonnet 4 via Amazon Q (1.3x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "amazonq-claude-haiku-4.5",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Amazon Q Claude Haiku 4.5",
Description: "Claude Haiku 4.5 via Amazon Q (0.4x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
}
}

View File

@@ -308,14 +308,12 @@ func (e *AIStudioExecutor) translateRequest(req cliproxyexecutor.Request, opts c
from := opts.SourceFormat
to := sdktranslator.FromString("gemini")
payload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), stream)
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
if budgetOverride != nil {
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
budgetOverride = &norm
}
payload = util.ApplyGeminiThinkingConfig(payload, budgetOverride, includeOverride)
}
payload = applyThinkingMetadata(payload, req.Metadata, req.Model)
payload = util.ConvertThinkingLevelToBudget(payload)
if budget := gjson.GetBytes(payload, "generationConfig.thinkingConfig.thinkingBudget"); budget.Exists() {
normalized := util.NormalizeThinkingBudget(req.Model, int(budget.Int()))
payload, _ = sjson.SetBytes(payload, "generationConfig.thinkingConfig.thinkingBudget", normalized)
}
payload = util.StripThinkingConfigIfUnsupported(req.Model, payload)
payload = fixGeminiImageAspectRatio(req.Model, payload)
payload = applyPayloadConfig(e.cfg, req.Model, payload)

View File

@@ -12,11 +12,13 @@ import (
"net/url"
"strconv"
"strings"
"sync"
"time"
"github.com/google/uuid"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
@@ -40,7 +42,10 @@ const (
streamScannerBuffer int = 20_971_520
)
var randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
var (
randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
randSourceMutex sync.Mutex
)
// AntigravityExecutor proxies requests to the antigravity upstream.
type AntigravityExecutor struct {
@@ -75,6 +80,9 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
to := sdktranslator.FromString("antigravity")
translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)
translated = normalizeAntigravityThinking(req.Model, translated)
baseURLs := antigravityBaseURLFallbackOrder(auth)
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
@@ -166,6 +174,9 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
to := sdktranslator.FromString("antigravity")
translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)
translated = normalizeAntigravityThinking(req.Model, translated)
baseURLs := antigravityBaseURLFallbackOrder(auth)
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
@@ -361,28 +372,29 @@ func FetchAntigravityModels(ctx context.Context, auth *cliproxyauth.Auth, cfg *c
}
now := time.Now().Unix()
modelConfig := registry.GetAntigravityModelConfig()
models := make([]*registry.ModelInfo, 0, len(result.Map()))
for id := range result.Map() {
id = modelName2Alias(id)
if id != "" {
for originalName := range result.Map() {
aliasName := modelName2Alias(originalName)
if aliasName != "" {
modelInfo := &registry.ModelInfo{
ID: id,
Name: id,
Description: id,
DisplayName: id,
Version: id,
ID: aliasName,
Name: aliasName,
Description: aliasName,
DisplayName: aliasName,
Version: aliasName,
Object: "model",
Created: now,
OwnedBy: antigravityAuthType,
Type: antigravityAuthType,
}
// Add Thinking support for thinking models
if strings.HasSuffix(id, "-thinking") || strings.Contains(id, "-thinking-") {
modelInfo.Thinking = &registry.ThinkingSupport{
Min: 1024,
Max: 100000,
ZeroAllowed: false,
DynamicAllowed: true,
// Look up Thinking support from static config using alias name
if cfg, ok := modelConfig[aliasName]; ok {
if cfg.Thinking != nil {
modelInfo.Thinking = cfg.Thinking
}
if cfg.MaxCompletionTokens > 0 {
modelInfo.MaxCompletionTokens = cfg.MaxCompletionTokens
}
}
models = append(models, modelInfo)
@@ -504,8 +516,49 @@ func (e *AntigravityExecutor) buildRequest(ctx context.Context, auth *cliproxyau
requestURL.WriteString(url.QueryEscape(alt))
}
payload = geminiToAntigravity(modelName, payload)
// Extract project_id from auth metadata if available
projectID := ""
if auth != nil && auth.Metadata != nil {
if pid, ok := auth.Metadata["project_id"].(string); ok {
projectID = strings.TrimSpace(pid)
}
}
payload = geminiToAntigravity(modelName, payload, projectID)
payload, _ = sjson.SetBytes(payload, "model", alias2ModelName(modelName))
if strings.Contains(modelName, "claude") {
strJSON := string(payload)
paths := make([]string, 0)
util.Walk(gjson.ParseBytes(payload), "", "parametersJsonSchema", &paths)
for _, p := range paths {
strJSON, _ = util.RenameKey(strJSON, p, p[:len(p)-len("parametersJsonSchema")]+"parameters")
}
strJSON = util.DeleteKey(strJSON, "$schema")
strJSON = util.DeleteKey(strJSON, "maxItems")
strJSON = util.DeleteKey(strJSON, "minItems")
strJSON = util.DeleteKey(strJSON, "minLength")
strJSON = util.DeleteKey(strJSON, "maxLength")
strJSON = util.DeleteKey(strJSON, "exclusiveMinimum")
strJSON = util.DeleteKey(strJSON, "exclusiveMaximum")
strJSON = util.DeleteKey(strJSON, "$ref")
strJSON = util.DeleteKey(strJSON, "$defs")
paths = make([]string, 0)
util.Walk(gjson.Parse(strJSON), "", "anyOf", &paths)
for _, p := range paths {
anyOf := gjson.Get(strJSON, p)
if anyOf.IsArray() {
anyOfItems := anyOf.Array()
if len(anyOfItems) > 0 {
strJSON, _ = sjson.SetRaw(strJSON, p[:len(p)-len(".anyOf")], anyOfItems[0].Raw)
}
}
}
payload = []byte(strJSON)
}
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(payload))
if errReq != nil {
return nil, errReq
@@ -642,7 +695,7 @@ func antigravityBaseURLFallbackOrder(auth *cliproxyauth.Auth) []string {
return []string{
antigravityBaseURLDaily,
antigravityBaseURLAutopush,
// antigravityBaseURLProd,
antigravityBaseURLProd,
}
}
@@ -666,16 +719,22 @@ func resolveCustomAntigravityBaseURL(auth *cliproxyauth.Auth) string {
return ""
}
func geminiToAntigravity(modelName string, payload []byte) []byte {
func geminiToAntigravity(modelName string, payload []byte, projectID string) []byte {
template, _ := sjson.Set(string(payload), "model", modelName)
template, _ = sjson.Set(template, "userAgent", "antigravity")
template, _ = sjson.Set(template, "project", generateProjectID())
// Use real project ID from auth if available, otherwise generate random (legacy fallback)
if projectID != "" {
template, _ = sjson.Set(template, "project", projectID)
} else {
template, _ = sjson.Set(template, "project", generateProjectID())
}
template, _ = sjson.Set(template, "requestId", generateRequestID())
template, _ = sjson.Set(template, "request.sessionId", generateSessionID())
template, _ = sjson.Delete(template, "request.safetySettings")
template, _ = sjson.Set(template, "request.toolConfig.functionCallingConfig.mode", "VALIDATED")
template, _ = sjson.Delete(template, "request.generationConfig.maxOutputTokens")
if !strings.HasPrefix(modelName, "gemini-3-") {
if thinkingLevel := gjson.Get(template, "request.generationConfig.thinkingConfig.thinkingLevel"); thinkingLevel.Exists() {
template, _ = sjson.Delete(template, "request.generationConfig.thinkingConfig.thinkingLevel")
@@ -683,7 +742,7 @@ func geminiToAntigravity(modelName string, payload []byte) []byte {
}
}
if strings.HasPrefix(modelName, "claude-sonnet-") {
if strings.Contains(modelName, "claude") {
gjson.Get(template, "request.tools").ForEach(func(key, tool gjson.Result) bool {
tool.Get("functionDeclarations").ForEach(func(funKey, funcDecl gjson.Result) bool {
if funcDecl.Get("parametersJsonSchema").Exists() {
@@ -695,6 +754,8 @@ func geminiToAntigravity(modelName string, payload []byte) []byte {
})
return true
})
} else {
template, _ = sjson.Delete(template, "request.generationConfig.maxOutputTokens")
}
return []byte(template)
@@ -705,15 +766,19 @@ func generateRequestID() string {
}
func generateSessionID() string {
randSourceMutex.Lock()
n := randSource.Int63n(9_000_000_000_000_000_000)
randSourceMutex.Unlock()
return "-" + strconv.FormatInt(n, 10)
}
func generateProjectID() string {
adjectives := []string{"useful", "bright", "swift", "calm", "bold"}
nouns := []string{"fuze", "wave", "spark", "flow", "core"}
randSourceMutex.Lock()
adj := adjectives[randSource.Intn(len(adjectives))]
noun := nouns[randSource.Intn(len(nouns))]
randSourceMutex.Unlock()
randomPart := strings.ToLower(uuid.NewString())[:5]
return adj + "-" + noun + "-" + randomPart
}
@@ -730,6 +795,8 @@ func modelName2Alias(modelName string) string {
return "gemini-claude-sonnet-4-5"
case "claude-sonnet-4-5-thinking":
return "gemini-claude-sonnet-4-5-thinking"
case "claude-opus-4-5-thinking":
return "gemini-claude-opus-4-5-thinking"
case "chat_20706", "chat_23310", "gemini-2.5-flash-thinking", "gemini-3-pro-low", "gemini-2.5-pro":
return ""
default:
@@ -749,7 +816,59 @@ func alias2ModelName(modelName string) string {
return "claude-sonnet-4-5"
case "gemini-claude-sonnet-4-5-thinking":
return "claude-sonnet-4-5-thinking"
case "gemini-claude-opus-4-5-thinking":
return "claude-opus-4-5-thinking"
default:
return modelName
}
}
// normalizeAntigravityThinking clamps or removes thinking config based on model support.
// For Claude models, it additionally ensures thinking budget < max_tokens.
func normalizeAntigravityThinking(model string, payload []byte) []byte {
payload = util.StripThinkingConfigIfUnsupported(model, payload)
if !util.ModelSupportsThinking(model) {
return payload
}
budget := gjson.GetBytes(payload, "request.generationConfig.thinkingConfig.thinkingBudget")
if !budget.Exists() {
return payload
}
raw := int(budget.Int())
normalized := util.NormalizeThinkingBudget(model, raw)
isClaude := strings.Contains(strings.ToLower(model), "claude")
if isClaude {
effectiveMax, setDefaultMax := antigravityEffectiveMaxTokens(model, payload)
if effectiveMax > 0 && normalized >= effectiveMax {
normalized = effectiveMax - 1
if normalized < 1 {
normalized = 1
}
}
if setDefaultMax {
if res, errSet := sjson.SetBytes(payload, "request.generationConfig.maxOutputTokens", effectiveMax); errSet == nil {
payload = res
}
}
}
updated, err := sjson.SetBytes(payload, "request.generationConfig.thinkingConfig.thinkingBudget", normalized)
if err != nil {
return payload
}
return updated
}
// antigravityEffectiveMaxTokens returns the max tokens to cap thinking:
// prefer request-provided maxOutputTokens; otherwise fall back to model default.
// The boolean indicates whether the value came from the model default (and thus should be written back).
func antigravityEffectiveMaxTokens(model string, payload []byte) (max int, fromModel bool) {
if maxTok := gjson.GetBytes(payload, "request.generationConfig.maxOutputTokens"); maxTok.Exists() && maxTok.Int() > 0 {
return int(maxTok.Int()), false
}
if modelInfo := registry.GetGlobalRegistry().GetModelInfo(model); modelInfo != nil && modelInfo.MaxCompletionTokens > 0 {
return modelInfo.MaxCompletionTokens, true
}
return 0, false
}

View File

@@ -1,10 +1,38 @@
package executor
import "time"
import (
"sync"
"time"
)
type codexCache struct {
ID string
Expire time.Time
}
var codexCacheMap = map[string]codexCache{}
var (
codexCacheMap = map[string]codexCache{}
codexCacheMutex sync.RWMutex
)
// getCodexCache safely retrieves a cache entry
func getCodexCache(key string) (codexCache, bool) {
codexCacheMutex.RLock()
defer codexCacheMutex.RUnlock()
cache, ok := codexCacheMap[key]
return cache, ok
}
// setCodexCache safely sets a cache entry
func setCodexCache(key string, cache codexCache) {
codexCacheMutex.Lock()
defer codexCacheMutex.Unlock()
codexCacheMap[key] = cache
}
// deleteCodexCache safely deletes a cache entry
func deleteCodexCache(key string) {
codexCacheMutex.Lock()
defer codexCacheMutex.Unlock()
delete(codexCacheMap, key)
}

View File

@@ -506,12 +506,12 @@ func (e *CodexExecutor) cacheHelper(ctx context.Context, from sdktranslator.Form
if userIDResult.Exists() {
var hasKey bool
key := fmt.Sprintf("%s-%s", req.Model, userIDResult.String())
if cache, hasKey = codexCacheMap[key]; !hasKey || cache.Expire.Before(time.Now()) {
if cache, hasKey = getCodexCache(key); !hasKey || cache.Expire.Before(time.Now()) {
cache = codexCache{
ID: uuid.New().String(),
Expire: time.Now().Add(1 * time.Hour),
}
codexCacheMap[key] = cache
setCodexCache(key, cache)
}
}
} else if from == "openai-response" {

View File

@@ -62,15 +62,8 @@ func (e *GeminiCLIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth
from := opts.SourceFormat
to := sdktranslator.FromString("gemini-cli")
budgetOverride, includeOverride, hasOverride := util.GeminiThinkingFromMetadata(req.Metadata)
basePayload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
if hasOverride && util.ModelSupportsThinking(req.Model) {
if budgetOverride != nil {
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
budgetOverride = &norm
}
basePayload = util.ApplyGeminiCLIThinkingConfig(basePayload, budgetOverride, includeOverride)
}
basePayload = applyThinkingMetadataCLI(basePayload, req.Metadata, req.Model)
basePayload = util.StripThinkingConfigIfUnsupported(req.Model, basePayload)
basePayload = fixGeminiCLIImageAspectRatio(req.Model, basePayload)
basePayload = applyPayloadConfigWithRoot(e.cfg, req.Model, "gemini", "request", basePayload)
@@ -204,15 +197,8 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
from := opts.SourceFormat
to := sdktranslator.FromString("gemini-cli")
budgetOverride, includeOverride, hasOverride := util.GeminiThinkingFromMetadata(req.Metadata)
basePayload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
if hasOverride && util.ModelSupportsThinking(req.Model) {
if budgetOverride != nil {
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
budgetOverride = &norm
}
basePayload = util.ApplyGeminiCLIThinkingConfig(basePayload, budgetOverride, includeOverride)
}
basePayload = applyThinkingMetadataCLI(basePayload, req.Metadata, req.Model)
basePayload = util.StripThinkingConfigIfUnsupported(req.Model, basePayload)
basePayload = fixGeminiCLIImageAspectRatio(req.Model, basePayload)
basePayload = applyPayloadConfigWithRoot(e.cfg, req.Model, "gemini", "request", basePayload)
@@ -408,16 +394,9 @@ func (e *GeminiCLIExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.
var lastStatus int
var lastBody []byte
budgetOverride, includeOverride, hasOverride := util.GeminiThinkingFromMetadata(req.Metadata)
for _, attemptModel := range models {
payload := sdktranslator.TranslateRequest(from, to, attemptModel, bytes.Clone(req.Payload), false)
if hasOverride && util.ModelSupportsThinking(req.Model) {
if budgetOverride != nil {
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
budgetOverride = &norm
}
payload = util.ApplyGeminiCLIThinkingConfig(payload, budgetOverride, includeOverride)
}
payload = applyThinkingMetadataCLI(payload, req.Metadata, req.Model)
payload = deleteJSONField(payload, "project")
payload = deleteJSONField(payload, "model")
payload = deleteJSONField(payload, "request.safetySettings")
@@ -688,7 +667,7 @@ func cliPreviewFallbackOrder(model string) []string {
case "gemini-2.5-pro":
return []string{
// "gemini-2.5-pro-preview-05-06",
"gemini-2.5-pro-preview-06-05",
// "gemini-2.5-pro-preview-06-05",
}
case "gemini-2.5-flash":
return []string{

View File

@@ -79,13 +79,7 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
from := opts.SourceFormat
to := sdktranslator.FromString("gemini")
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
if budgetOverride != nil {
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
budgetOverride = &norm
}
body = util.ApplyGeminiThinkingConfig(body, budgetOverride, includeOverride)
}
body = applyThinkingMetadata(body, req.Metadata, req.Model)
body = util.StripThinkingConfigIfUnsupported(req.Model, body)
body = fixGeminiImageAspectRatio(req.Model, body)
body = applyPayloadConfig(e.cfg, req.Model, body)
@@ -174,13 +168,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
from := opts.SourceFormat
to := sdktranslator.FromString("gemini")
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
if budgetOverride != nil {
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
budgetOverride = &norm
}
body = util.ApplyGeminiThinkingConfig(body, budgetOverride, includeOverride)
}
body = applyThinkingMetadata(body, req.Metadata, req.Model)
body = util.StripThinkingConfigIfUnsupported(req.Model, body)
body = fixGeminiImageAspectRatio(req.Model, body)
body = applyPayloadConfig(e.cfg, req.Model, body)
@@ -288,13 +276,7 @@ func (e *GeminiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut
from := opts.SourceFormat
to := sdktranslator.FromString("gemini")
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
if budgetOverride != nil {
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
budgetOverride = &norm
}
translatedReq = util.ApplyGeminiThinkingConfig(translatedReq, budgetOverride, includeOverride)
}
translatedReq = applyThinkingMetadata(translatedReq, req.Metadata, req.Model)
translatedReq = util.StripThinkingConfigIfUnsupported(req.Model, translatedReq)
translatedReq = fixGeminiImageAspectRatio(req.Model, translatedReq)
respCtx := context.WithValue(ctx, "alt", opts.Alt)

View File

@@ -51,11 +51,238 @@ func (e *GeminiVertexExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.A
// Execute handles non-streaming requests.
func (e *GeminiVertexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
projectID, location, saJSON, errCreds := vertexCreds(auth)
if errCreds != nil {
return resp, errCreds
// Try API key authentication first
apiKey, baseURL := vertexAPICreds(auth)
// If no API key found, fall back to service account authentication
if apiKey == "" {
projectID, location, saJSON, errCreds := vertexCreds(auth)
if errCreds != nil {
return resp, errCreds
}
return e.executeWithServiceAccount(ctx, auth, req, opts, projectID, location, saJSON)
}
// Use API key authentication
return e.executeWithAPIKey(ctx, auth, req, opts, apiKey, baseURL)
}
// ExecuteStream handles SSE streaming for Vertex.
func (e *GeminiVertexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
// Try API key authentication first
apiKey, baseURL := vertexAPICreds(auth)
// If no API key found, fall back to service account authentication
if apiKey == "" {
projectID, location, saJSON, errCreds := vertexCreds(auth)
if errCreds != nil {
return nil, errCreds
}
return e.executeStreamWithServiceAccount(ctx, auth, req, opts, projectID, location, saJSON)
}
// Use API key authentication
return e.executeStreamWithAPIKey(ctx, auth, req, opts, apiKey, baseURL)
}
// CountTokens calls Vertex countTokens endpoint.
func (e *GeminiVertexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
// Try API key authentication first
apiKey, baseURL := vertexAPICreds(auth)
// If no API key found, fall back to service account authentication
if apiKey == "" {
projectID, location, saJSON, errCreds := vertexCreds(auth)
if errCreds != nil {
return cliproxyexecutor.Response{}, errCreds
}
return e.countTokensWithServiceAccount(ctx, auth, req, opts, projectID, location, saJSON)
}
// Use API key authentication
return e.countTokensWithAPIKey(ctx, auth, req, opts, apiKey, baseURL)
}
// countTokensWithServiceAccount handles token counting using service account credentials.
func (e *GeminiVertexExecutor) countTokensWithServiceAccount(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, projectID, location string, saJSON []byte) (cliproxyexecutor.Response, error) {
from := opts.SourceFormat
to := sdktranslator.FromString("gemini")
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
if budgetOverride != nil {
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
budgetOverride = &norm
}
translatedReq = util.ApplyGeminiThinkingConfig(translatedReq, budgetOverride, includeOverride)
}
translatedReq = util.StripThinkingConfigIfUnsupported(req.Model, translatedReq)
translatedReq = fixGeminiImageAspectRatio(req.Model, translatedReq)
respCtx := context.WithValue(ctx, "alt", opts.Alt)
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
baseURL := vertexBaseURL(location)
url := fmt.Sprintf("%s/%s/projects/%s/locations/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, projectID, location, req.Model, "countTokens")
httpReq, errNewReq := http.NewRequestWithContext(respCtx, http.MethodPost, url, bytes.NewReader(translatedReq))
if errNewReq != nil {
return cliproxyexecutor.Response{}, errNewReq
}
httpReq.Header.Set("Content-Type", "application/json")
if token, errTok := vertexAccessToken(ctx, e.cfg, auth, saJSON); errTok == nil && token != "" {
httpReq.Header.Set("Authorization", "Bearer "+token)
} else if errTok != nil {
log.Errorf("vertex executor: access token error: %v", errTok)
return cliproxyexecutor.Response{}, statusErr{code: 500, msg: "internal server error"}
}
applyGeminiHeaders(httpReq, auth)
var authID, authLabel, authType, authValue string
if auth != nil {
authID = auth.ID
authLabel = auth.Label
authType, authValue = auth.AccountInfo()
}
recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
URL: url,
Method: http.MethodPost,
Headers: httpReq.Header.Clone(),
Body: translatedReq,
Provider: e.Identifier(),
AuthID: authID,
AuthLabel: authLabel,
AuthType: authType,
AuthValue: authValue,
})
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
httpResp, errDo := httpClient.Do(httpReq)
if errDo != nil {
recordAPIResponseError(ctx, e.cfg, errDo)
return cliproxyexecutor.Response{}, errDo
}
defer func() {
if errClose := httpResp.Body.Close(); errClose != nil {
log.Errorf("vertex executor: close response body error: %v", errClose)
}
}()
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
b, _ := io.ReadAll(httpResp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(b)}
}
data, errRead := io.ReadAll(httpResp.Body)
if errRead != nil {
recordAPIResponseError(ctx, e.cfg, errRead)
return cliproxyexecutor.Response{}, errRead
}
appendAPIResponseChunk(ctx, e.cfg, data)
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(data)}
}
count := gjson.GetBytes(data, "totalTokens").Int()
out := sdktranslator.TranslateTokenCount(ctx, to, from, count, data)
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
}
// countTokensWithAPIKey handles token counting using API key credentials.
func (e *GeminiVertexExecutor) countTokensWithAPIKey(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, apiKey, baseURL string) (cliproxyexecutor.Response, error) {
from := opts.SourceFormat
to := sdktranslator.FromString("gemini")
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
if budgetOverride != nil {
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
budgetOverride = &norm
}
translatedReq = util.ApplyGeminiThinkingConfig(translatedReq, budgetOverride, includeOverride)
}
translatedReq = util.StripThinkingConfigIfUnsupported(req.Model, translatedReq)
translatedReq = fixGeminiImageAspectRatio(req.Model, translatedReq)
respCtx := context.WithValue(ctx, "alt", opts.Alt)
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
// For API key auth, use simpler URL format without project/location
if baseURL == "" {
baseURL = "https://generativelanguage.googleapis.com"
}
url := fmt.Sprintf("%s/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, req.Model, "countTokens")
httpReq, errNewReq := http.NewRequestWithContext(respCtx, http.MethodPost, url, bytes.NewReader(translatedReq))
if errNewReq != nil {
return cliproxyexecutor.Response{}, errNewReq
}
httpReq.Header.Set("Content-Type", "application/json")
if apiKey != "" {
httpReq.Header.Set("x-goog-api-key", apiKey)
}
applyGeminiHeaders(httpReq, auth)
var authID, authLabel, authType, authValue string
if auth != nil {
authID = auth.ID
authLabel = auth.Label
authType, authValue = auth.AccountInfo()
}
recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
URL: url,
Method: http.MethodPost,
Headers: httpReq.Header.Clone(),
Body: translatedReq,
Provider: e.Identifier(),
AuthID: authID,
AuthLabel: authLabel,
AuthType: authType,
AuthValue: authValue,
})
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
httpResp, errDo := httpClient.Do(httpReq)
if errDo != nil {
recordAPIResponseError(ctx, e.cfg, errDo)
return cliproxyexecutor.Response{}, errDo
}
defer func() {
if errClose := httpResp.Body.Close(); errClose != nil {
log.Errorf("vertex executor: close response body error: %v", errClose)
}
}()
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
b, _ := io.ReadAll(httpResp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(b)}
}
data, errRead := io.ReadAll(httpResp.Body)
if errRead != nil {
recordAPIResponseError(ctx, e.cfg, errRead)
return cliproxyexecutor.Response{}, errRead
}
appendAPIResponseChunk(ctx, e.cfg, data)
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(data)}
}
count := gjson.GetBytes(data, "totalTokens").Int()
out := sdktranslator.TranslateTokenCount(ctx, to, from, count, data)
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
}
// Refresh is a no-op for service account based credentials.
func (e *GeminiVertexExecutor) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
return auth, nil
}
// executeWithServiceAccount handles authentication using service account credentials.
// This method contains the original service account authentication logic.
func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, projectID, location string, saJSON []byte) (resp cliproxyexecutor.Response, err error) {
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
defer reporter.trackFailure(ctx, &err)
@@ -149,13 +376,104 @@ func (e *GeminiVertexExecutor) Execute(ctx context.Context, auth *cliproxyauth.A
return resp, nil
}
// ExecuteStream handles SSE streaming for Vertex.
func (e *GeminiVertexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
projectID, location, saJSON, errCreds := vertexCreds(auth)
if errCreds != nil {
return nil, errCreds
// executeWithAPIKey handles authentication using API key credentials.
func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, apiKey, baseURL string) (resp cliproxyexecutor.Response, err error) {
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
defer reporter.trackFailure(ctx, &err)
from := opts.SourceFormat
to := sdktranslator.FromString("gemini")
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
if budgetOverride != nil {
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
budgetOverride = &norm
}
body = util.ApplyGeminiThinkingConfig(body, budgetOverride, includeOverride)
}
body = util.StripThinkingConfigIfUnsupported(req.Model, body)
body = fixGeminiImageAspectRatio(req.Model, body)
body = applyPayloadConfig(e.cfg, req.Model, body)
action := "generateContent"
if req.Metadata != nil {
if a, _ := req.Metadata["action"].(string); a == "countTokens" {
action = "countTokens"
}
}
// For API key auth, use simpler URL format without project/location
if baseURL == "" {
baseURL = "https://generativelanguage.googleapis.com"
}
url := fmt.Sprintf("%s/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, req.Model, action)
if opts.Alt != "" && action != "countTokens" {
url = url + fmt.Sprintf("?$alt=%s", opts.Alt)
}
body, _ = sjson.DeleteBytes(body, "session_id")
httpReq, errNewReq := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if errNewReq != nil {
return resp, errNewReq
}
httpReq.Header.Set("Content-Type", "application/json")
if apiKey != "" {
httpReq.Header.Set("x-goog-api-key", apiKey)
}
applyGeminiHeaders(httpReq, auth)
var authID, authLabel, authType, authValue string
if auth != nil {
authID = auth.ID
authLabel = auth.Label
authType, authValue = auth.AccountInfo()
}
recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
URL: url,
Method: http.MethodPost,
Headers: httpReq.Header.Clone(),
Body: body,
Provider: e.Identifier(),
AuthID: authID,
AuthLabel: authLabel,
AuthType: authType,
AuthValue: authValue,
})
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
httpResp, errDo := httpClient.Do(httpReq)
if errDo != nil {
recordAPIResponseError(ctx, e.cfg, errDo)
return resp, errDo
}
defer func() {
if errClose := httpResp.Body.Close(); errClose != nil {
log.Errorf("vertex executor: close response body error: %v", errClose)
}
}()
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
b, _ := io.ReadAll(httpResp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
err = statusErr{code: httpResp.StatusCode, msg: string(b)}
return resp, err
}
data, errRead := io.ReadAll(httpResp.Body)
if errRead != nil {
recordAPIResponseError(ctx, e.cfg, errRead)
return resp, errRead
}
appendAPIResponseChunk(ctx, e.cfg, data)
reporter.publish(ctx, parseGeminiUsage(data))
var param any
out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, data, &param)
resp = cliproxyexecutor.Response{Payload: []byte(out)}
return resp, nil
}
// executeStreamWithServiceAccount handles streaming authentication using service account credentials.
func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, projectID, location string, saJSON []byte) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
defer reporter.trackFailure(ctx, &err)
@@ -266,42 +584,44 @@ func (e *GeminiVertexExecutor) ExecuteStream(ctx context.Context, auth *cliproxy
return stream, nil
}
// CountTokens calls Vertex countTokens endpoint.
func (e *GeminiVertexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
projectID, location, saJSON, errCreds := vertexCreds(auth)
if errCreds != nil {
return cliproxyexecutor.Response{}, errCreds
}
// executeStreamWithAPIKey handles streaming authentication using API key credentials.
func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, apiKey, baseURL string) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
defer reporter.trackFailure(ctx, &err)
from := opts.SourceFormat
to := sdktranslator.FromString("gemini")
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
if budgetOverride != nil {
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
budgetOverride = &norm
}
translatedReq = util.ApplyGeminiThinkingConfig(translatedReq, budgetOverride, includeOverride)
body = util.ApplyGeminiThinkingConfig(body, budgetOverride, includeOverride)
}
translatedReq = util.StripThinkingConfigIfUnsupported(req.Model, translatedReq)
translatedReq = fixGeminiImageAspectRatio(req.Model, translatedReq)
respCtx := context.WithValue(ctx, "alt", opts.Alt)
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
body = util.StripThinkingConfigIfUnsupported(req.Model, body)
body = fixGeminiImageAspectRatio(req.Model, body)
body = applyPayloadConfig(e.cfg, req.Model, body)
baseURL := vertexBaseURL(location)
url := fmt.Sprintf("%s/%s/projects/%s/locations/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, projectID, location, req.Model, "countTokens")
// For API key auth, use simpler URL format without project/location
if baseURL == "" {
baseURL = "https://generativelanguage.googleapis.com"
}
url := fmt.Sprintf("%s/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, req.Model, "streamGenerateContent")
if opts.Alt == "" {
url = url + "?alt=sse"
} else {
url = url + fmt.Sprintf("?$alt=%s", opts.Alt)
}
body, _ = sjson.DeleteBytes(body, "session_id")
httpReq, errNewReq := http.NewRequestWithContext(respCtx, http.MethodPost, url, bytes.NewReader(translatedReq))
httpReq, errNewReq := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if errNewReq != nil {
return cliproxyexecutor.Response{}, errNewReq
return nil, errNewReq
}
httpReq.Header.Set("Content-Type", "application/json")
if token, errTok := vertexAccessToken(ctx, e.cfg, auth, saJSON); errTok == nil && token != "" {
httpReq.Header.Set("Authorization", "Bearer "+token)
} else if errTok != nil {
log.Errorf("vertex executor: access token error: %v", errTok)
return cliproxyexecutor.Response{}, statusErr{code: 500, msg: "internal server error"}
if apiKey != "" {
httpReq.Header.Set("x-goog-api-key", apiKey)
}
applyGeminiHeaders(httpReq, auth)
@@ -315,7 +635,7 @@ func (e *GeminiVertexExecutor) CountTokens(ctx context.Context, auth *cliproxyau
URL: url,
Method: http.MethodPost,
Headers: httpReq.Header.Clone(),
Body: translatedReq,
Body: body,
Provider: e.Identifier(),
AuthID: authID,
AuthLabel: authLabel,
@@ -327,38 +647,53 @@ func (e *GeminiVertexExecutor) CountTokens(ctx context.Context, auth *cliproxyau
httpResp, errDo := httpClient.Do(httpReq)
if errDo != nil {
recordAPIResponseError(ctx, e.cfg, errDo)
return cliproxyexecutor.Response{}, errDo
return nil, errDo
}
defer func() {
if errClose := httpResp.Body.Close(); errClose != nil {
log.Errorf("vertex executor: close response body error: %v", errClose)
}
}()
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
b, _ := io.ReadAll(httpResp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(b)}
if errClose := httpResp.Body.Close(); errClose != nil {
log.Errorf("vertex executor: close response body error: %v", errClose)
}
return nil, statusErr{code: httpResp.StatusCode, msg: string(b)}
}
data, errRead := io.ReadAll(httpResp.Body)
if errRead != nil {
recordAPIResponseError(ctx, e.cfg, errRead)
return cliproxyexecutor.Response{}, errRead
}
appendAPIResponseChunk(ctx, e.cfg, data)
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(data)}
}
count := gjson.GetBytes(data, "totalTokens").Int()
out := sdktranslator.TranslateTokenCount(ctx, to, from, count, data)
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
}
// Refresh is a no-op for service account based credentials.
func (e *GeminiVertexExecutor) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
return auth, nil
out := make(chan cliproxyexecutor.StreamChunk)
stream = out
go func() {
defer close(out)
defer func() {
if errClose := httpResp.Body.Close(); errClose != nil {
log.Errorf("vertex executor: close response body error: %v", errClose)
}
}()
scanner := bufio.NewScanner(httpResp.Body)
scanner.Buffer(nil, 20_971_520)
var param any
for scanner.Scan() {
line := scanner.Bytes()
appendAPIResponseChunk(ctx, e.cfg, line)
if detail, ok := parseGeminiStreamUsage(line); ok {
reporter.publish(ctx, detail)
}
lines := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, bytes.Clone(line), &param)
for i := range lines {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(lines[i])}
}
}
lines := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, []byte("[DONE]"), &param)
for i := range lines {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(lines[i])}
}
if errScan := scanner.Err(); errScan != nil {
recordAPIResponseError(ctx, e.cfg, errScan)
reporter.publishFailure(ctx)
out <- cliproxyexecutor.StreamChunk{Err: errScan}
}
}()
return stream, nil
}
// vertexCreds extracts project, location and raw service account JSON from auth metadata.
@@ -401,6 +736,23 @@ func vertexCreds(a *cliproxyauth.Auth) (projectID, location string, serviceAccou
return projectID, location, saJSON, nil
}
// vertexAPICreds extracts API key and base URL from auth attributes following the claudeCreds pattern.
func vertexAPICreds(a *cliproxyauth.Auth) (apiKey, baseURL string) {
if a == nil {
return "", ""
}
if a.Attributes != nil {
apiKey = a.Attributes["api_key"]
baseURL = a.Attributes["base_url"]
}
if apiKey == "" && a.Metadata != nil {
if v, ok := a.Metadata["access_token"].(string); ok {
apiKey = v
}
}
return
}
func vertexBaseURL(location string) string {
loc := strings.TrimSpace(location)
if loc == "" {

File diff suppressed because it is too large Load Diff

View File

@@ -157,7 +157,7 @@ func appendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byt
if ginCtx == nil {
return
}
attempts, attempt := ensureAttempt(ginCtx)
_, attempt := ensureAttempt(ginCtx)
ensureResponseIntro(attempt)
if !attempt.headersWritten {
@@ -175,8 +175,6 @@ func appendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byt
}
attempt.response.WriteString(string(data))
attempt.bodyHasContent = true
updateAggregatedResponse(ginCtx, attempts)
}
func ginContextFrom(ctx context.Context) *gin.Context {

View File

@@ -4,10 +4,42 @@ import (
"strings"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// applyThinkingMetadata applies thinking config from model suffix metadata (e.g., -reasoning, -thinking-N)
// for standard Gemini format payloads. It normalizes the budget when the model supports thinking.
func applyThinkingMetadata(payload []byte, metadata map[string]any, model string) []byte {
budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(metadata)
if !ok {
return payload
}
if !util.ModelSupportsThinking(model) {
return payload
}
if budgetOverride != nil {
norm := util.NormalizeThinkingBudget(model, *budgetOverride)
budgetOverride = &norm
}
return util.ApplyGeminiThinkingConfig(payload, budgetOverride, includeOverride)
}
// applyThinkingMetadataCLI applies thinking config from model suffix metadata (e.g., -reasoning, -thinking-N)
// for Gemini CLI format payloads (nested under "request"). It normalizes the budget when the model supports thinking.
func applyThinkingMetadataCLI(payload []byte, metadata map[string]any, model string) []byte {
budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(metadata)
if !ok {
return payload
}
if budgetOverride != nil && util.ModelSupportsThinking(model) {
norm := util.NormalizeThinkingBudget(model, *budgetOverride)
budgetOverride = &norm
}
return util.ApplyGeminiCLIThinkingConfig(payload, budgetOverride, includeOverride)
}
// applyPayloadConfig applies payload default and override rules from configuration
// to the given JSON payload for the specified model.
// Defaults only fill missing fields, while overrides always overwrite existing values.

View File

@@ -6,6 +6,7 @@ import (
"net/http"
"net/url"
"strings"
"sync"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
@@ -14,11 +15,19 @@ import (
"golang.org/x/net/proxy"
)
// httpClientCache caches HTTP clients by proxy URL to enable connection reuse
var (
httpClientCache = make(map[string]*http.Client)
httpClientCacheMutex sync.RWMutex
)
// newProxyAwareHTTPClient creates an HTTP client with proper proxy configuration priority:
// 1. Use auth.ProxyURL if configured (highest priority)
// 2. Use cfg.ProxyURL if auth proxy is not configured
// 3. Use RoundTripper from context if neither are configured
//
// This function caches HTTP clients by proxy URL to enable TCP/TLS connection reuse.
//
// Parameters:
// - ctx: The context containing optional RoundTripper
// - cfg: The application configuration
@@ -28,11 +37,6 @@ import (
// Returns:
// - *http.Client: An HTTP client with configured proxy or transport
func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client {
httpClient := &http.Client{}
if timeout > 0 {
httpClient.Timeout = timeout
}
// Priority 1: Use auth.ProxyURL if configured
var proxyURL string
if auth != nil {
@@ -44,11 +48,39 @@ func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *clip
proxyURL = strings.TrimSpace(cfg.ProxyURL)
}
// Build cache key from proxy URL (empty string for no proxy)
cacheKey := proxyURL
// Check cache first
httpClientCacheMutex.RLock()
if cachedClient, ok := httpClientCache[cacheKey]; ok {
httpClientCacheMutex.RUnlock()
// Return a wrapper with the requested timeout but shared transport
if timeout > 0 {
return &http.Client{
Transport: cachedClient.Transport,
Timeout: timeout,
}
}
return cachedClient
}
httpClientCacheMutex.RUnlock()
// Create new client
httpClient := &http.Client{}
if timeout > 0 {
httpClient.Timeout = timeout
}
// If we have a proxy URL configured, set up the transport
if proxyURL != "" {
transport := buildProxyTransport(proxyURL)
if transport != nil {
httpClient.Transport = transport
// Cache the client
httpClientCacheMutex.Lock()
httpClientCache[cacheKey] = httpClient
httpClientCacheMutex.Unlock()
return httpClient
}
// 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
httpClient.Transport = rt
}
// Cache the client for no-proxy case
if proxyURL == "" {
httpClientCacheMutex.Lock()
httpClientCache[cacheKey] = httpClient
httpClientCacheMutex.Unlock()
}
return httpClient
}

View File

@@ -83,18 +83,33 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
for j := 0; j < len(contentResults); j++ {
contentResult := contentResults[j]
contentTypeResult := contentResult.Get("type")
if contentTypeResult.Type == gjson.String && contentTypeResult.String() == "text" {
if contentTypeResult.Type == gjson.String && contentTypeResult.String() == "thinking" {
prompt := contentResult.Get("thinking").String()
signatureResult := contentResult.Get("signature")
signature := geminiCLIClaudeThoughtSignature
if signatureResult.Exists() {
signature = signatureResult.String()
}
clientContent.Parts = append(clientContent.Parts, client.Part{Text: prompt, Thought: true, ThoughtSignature: signature})
} else if contentTypeResult.Type == gjson.String && contentTypeResult.String() == "text" {
prompt := contentResult.Get("text").String()
clientContent.Parts = append(clientContent.Parts, client.Part{Text: prompt})
} else if contentTypeResult.Type == gjson.String && contentTypeResult.String() == "tool_use" {
functionName := contentResult.Get("name").String()
functionArgs := contentResult.Get("input").String()
functionID := contentResult.Get("id").String()
var args map[string]any
if err := json.Unmarshal([]byte(functionArgs), &args); err == nil {
clientContent.Parts = append(clientContent.Parts, client.Part{
FunctionCall: &client.FunctionCall{Name: functionName, Args: args},
ThoughtSignature: geminiCLIClaudeThoughtSignature,
})
if strings.Contains(modelName, "claude") {
clientContent.Parts = append(clientContent.Parts, client.Part{
FunctionCall: &client.FunctionCall{ID: functionID, Name: functionName, Args: args},
})
} else {
clientContent.Parts = append(clientContent.Parts, client.Part{
FunctionCall: &client.FunctionCall{ID: functionID, Name: functionName, Args: args},
ThoughtSignature: geminiCLIClaudeThoughtSignature,
})
}
}
} else if contentTypeResult.Type == gjson.String && contentTypeResult.String() == "tool_result" {
toolCallID := contentResult.Get("tool_use_id").String()
@@ -105,7 +120,7 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
funcName = strings.Join(toolCallIDs[0:len(toolCallIDs)-1], "-")
}
responseData := contentResult.Get("content").Raw
functionResponse := client.FunctionResponse{Name: funcName, Response: map[string]interface{}{"result": responseData}}
functionResponse := client.FunctionResponse{ID: toolCallID, Name: funcName, Response: map[string]interface{}{"result": responseData}}
clientContent.Parts = append(clientContent.Parts, client.Part{FunctionResponse: &functionResponse})
}
}
@@ -180,6 +195,9 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
if v := gjson.GetBytes(rawJSON, "top_k"); v.Exists() && v.Type == gjson.Number {
out, _ = sjson.Set(out, "request.generationConfig.topK", v.Num)
}
if v := gjson.GetBytes(rawJSON, "max_tokens"); v.Exists() && v.Type == gjson.Number {
out, _ = sjson.Set(out, "request.generationConfig.maxOutputTokens", v.Num)
}
outBytes := []byte(out)
outBytes = common.AttachDefaultSafetySettings(outBytes, "request.safetySettings")

View File

@@ -111,8 +111,11 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq
if partTextResult.Exists() {
// Process thinking content (internal reasoning)
if partResult.Get("thought").Bool() {
// Continue existing thinking block if already in thinking state
if params.ResponseType == 2 {
if thoughtSignature := partResult.Get("thoughtSignature"); thoughtSignature.Exists() && thoughtSignature.String() != "" {
output = output + "event: content_block_delta\n"
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":""}}`, params.ResponseIndex), "delta.signature", thoughtSignature.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
} else if params.ResponseType == 2 { // Continue existing thinking block if already in thinking state
output = output + "event: content_block_delta\n"
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"thinking_delta","thinking":""}}`, params.ResponseIndex), "delta.thinking", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
@@ -141,35 +144,39 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq
params.ResponseType = 2 // Set state to thinking
}
} else {
// Process regular text content (user-visible output)
// Continue existing text block if already in content state
if params.ResponseType == 1 {
output = output + "event: content_block_delta\n"
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"text_delta","text":""}}`, params.ResponseIndex), "delta.text", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
} else {
// Transition from another state to text content
// First, close any existing content block
if params.ResponseType != 0 {
if params.ResponseType == 2 {
// output = output + "event: content_block_delta\n"
// output = output + fmt.Sprintf(`data: {"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":null}}`, params.ResponseIndex)
// output = output + "\n\n\n"
finishReasonResult := gjson.GetBytes(rawJSON, "response.candidates.0.finishReason")
if partTextResult.String() != "" || !finishReasonResult.Exists() {
// Process regular text content (user-visible output)
// Continue existing text block if already in content state
if params.ResponseType == 1 {
output = output + "event: content_block_delta\n"
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"text_delta","text":""}}`, params.ResponseIndex), "delta.text", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
} else {
// Transition from another state to text content
// First, close any existing content block
if params.ResponseType != 0 {
if params.ResponseType == 2 {
// output = output + "event: content_block_delta\n"
// output = output + fmt.Sprintf(`data: {"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":null}}`, params.ResponseIndex)
// output = output + "\n\n\n"
}
output = output + "event: content_block_stop\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, params.ResponseIndex)
output = output + "\n\n\n"
params.ResponseIndex++
}
if partTextResult.String() != "" {
// Start a new text content block
output = output + "event: content_block_start\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"text","text":""}}`, params.ResponseIndex)
output = output + "\n\n\n"
output = output + "event: content_block_delta\n"
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"text_delta","text":""}}`, params.ResponseIndex), "delta.text", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
params.ResponseType = 1 // Set state to content
}
output = output + "event: content_block_stop\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, params.ResponseIndex)
output = output + "\n\n\n"
params.ResponseIndex++
}
// Start a new text content block
output = output + "event: content_block_start\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"text","text":""}}`, params.ResponseIndex)
output = output + "\n\n\n"
output = output + "event: content_block_delta\n"
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"text_delta","text":""}}`, params.ResponseIndex), "delta.text", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
params.ResponseType = 1 // Set state to content
}
}
} else if functionCallResult.Exists() {

View File

@@ -88,6 +88,20 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
}
}
// Claude/Anthropic API format: thinking.type == "enabled" with budget_tokens
// This allows Claude Code and other Claude API clients to pass thinking configuration
if !gjson.GetBytes(out, "request.generationConfig.thinkingConfig").Exists() && util.ModelSupportsThinking(modelName) {
if t := gjson.GetBytes(rawJSON, "thinking"); t.Exists() && t.IsObject() {
if t.Get("type").String() == "enabled" {
if b := t.Get("budget_tokens"); b.Exists() && b.Type == gjson.Number {
budget := util.NormalizeThinkingBudget(modelName, int(b.Int()))
out, _ = sjson.SetBytes(out, "request.generationConfig.thinkingConfig.thinkingBudget", budget)
out, _ = sjson.SetBytes(out, "request.generationConfig.thinkingConfig.include_thoughts", true)
}
}
}
}
// For gemini-3-pro-preview, always send default thinkingConfig when none specified.
// This matches the official Gemini CLI behavior which always sends:
// { thinkingBudget: -1, includeThoughts: true }
@@ -97,7 +111,7 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
out, _ = sjson.SetBytes(out, "request.generationConfig.thinkingConfig.include_thoughts", true)
}
// Temperature/top_p/top_k
// Temperature/top_p/top_k/max_tokens
if tr := gjson.GetBytes(rawJSON, "temperature"); tr.Exists() && tr.Type == gjson.Number {
out, _ = sjson.SetBytes(out, "request.generationConfig.temperature", tr.Num)
}
@@ -107,6 +121,9 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
if tkr := gjson.GetBytes(rawJSON, "top_k"); tkr.Exists() && tkr.Type == gjson.Number {
out, _ = sjson.SetBytes(out, "request.generationConfig.topK", tkr.Num)
}
if maxTok := gjson.GetBytes(rawJSON, "max_tokens"); maxTok.Exists() && maxTok.Type == gjson.Number {
out, _ = sjson.SetBytes(out, "request.generationConfig.maxOutputTokens", maxTok.Num)
}
// Map OpenAI modalities -> Gemini CLI request.generationConfig.responseModalities
// e.g. "modalities": ["image", "text"] -> ["IMAGE", "TEXT"]
@@ -251,6 +268,7 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
fid := tc.Get("id").String()
fname := tc.Get("function.name").String()
fargs := tc.Get("function.arguments").String()
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".functionCall.id", fid)
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".functionCall.name", fname)
node, _ = sjson.SetRawBytes(node, "parts."+itoa(p)+".functionCall.args", []byte(fargs))
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".thoughtSignature", geminiCLIFunctionThoughtSignature)
@@ -262,10 +280,11 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
out, _ = sjson.SetRawBytes(out, "request.contents.-1", node)
// Append a single tool content combining name + response per function
toolNode := []byte(`{"role":"tool","parts":[]}`)
toolNode := []byte(`{"role":"user","parts":[]}`)
pp := 0
for _, fid := range fIDs {
if name, ok := tcID2Name[fid]; ok {
toolNode, _ = sjson.SetBytes(toolNode, "parts."+itoa(pp)+".functionResponse.id", fid)
toolNode, _ = sjson.SetBytes(toolNode, "parts."+itoa(pp)+".functionResponse.name", name)
resp := toolResponses[fid]
if resp == "" {

View File

@@ -10,6 +10,7 @@ import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
. "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/gemini/openai/chat-completions"
@@ -75,8 +76,8 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq
// Extract and set the finish reason.
if finishReasonResult := gjson.GetBytes(rawJSON, "response.candidates.0.finishReason"); finishReasonResult.Exists() {
template, _ = sjson.Set(template, "choices.0.finish_reason", finishReasonResult.String())
template, _ = sjson.Set(template, "choices.0.native_finish_reason", finishReasonResult.String())
template, _ = sjson.Set(template, "choices.0.finish_reason", strings.ToLower(finishReasonResult.String()))
template, _ = sjson.Set(template, "choices.0.native_finish_reason", strings.ToLower(finishReasonResult.String()))
}
// Extract and set usage metadata (token counts).

View File

@@ -331,8 +331,9 @@ func ConvertClaudeResponseToGeminiNonStream(_ context.Context, modelName string,
streamingEvents := make([][]byte, 0)
scanner := bufio.NewScanner(bytes.NewReader(rawJSON))
buffer := make([]byte, 20_971_520)
scanner.Buffer(buffer, 20_971_520)
// Use a smaller initial buffer (64KB) that can grow up to 20MB if needed
// This prevents allocating 20MB for every request regardless of size
scanner.Buffer(make([]byte, 64*1024), 20_971_520)
for scanner.Scan() {
line := scanner.Bytes()
// log.Debug(string(line))

View File

@@ -50,6 +50,10 @@ type ToolCallAccumulator struct {
// Returns:
// - []string: A slice of strings, each containing an OpenAI-compatible JSON response
func ConvertClaudeResponseToOpenAI(_ context.Context, modelName string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []string {
var localParam any
if param == nil {
param = &localParam
}
if *param == nil {
*param = &ConvertAnthropicResponseToOpenAIParams{
CreatedAt: 0,

View File

@@ -327,7 +327,7 @@ func buildReverseMapFromGeminiOriginal(original []byte) map[string]string {
func mustMarshalJSON(v interface{}) string {
data, err := json.Marshal(v)
if err != nil {
panic(err)
return ""
}
return string(data)
}

View File

@@ -60,12 +60,12 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
// Track whether tools are being used in this response chunk
usedTool := false
output := ""
var sb strings.Builder
// Initialize the streaming session with a message_start event
// This is only sent for the very first response chunk to establish the streaming session
if !(*param).(*Params).HasFirstResponse {
output = "event: message_start\n"
sb.WriteString("event: message_start\n")
// Create the initial message structure with default values according to Claude Code API specification
// This follows the Claude Code API specification for streaming message initialization
@@ -78,7 +78,7 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
if responseIDResult := gjson.GetBytes(rawJSON, "response.responseId"); responseIDResult.Exists() {
messageStartTemplate, _ = sjson.Set(messageStartTemplate, "message.id", responseIDResult.String())
}
output = output + fmt.Sprintf("data: %s\n\n\n", messageStartTemplate)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", messageStartTemplate))
(*param).(*Params).HasFirstResponse = true
}
@@ -101,62 +101,52 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
if partResult.Get("thought").Bool() {
// Continue existing thinking block if already in thinking state
if (*param).(*Params).ResponseType == 2 {
output = output + "event: content_block_delta\n"
sb.WriteString("event: content_block_delta\n")
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"thinking_delta","thinking":""}}`, (*param).(*Params).ResponseIndex), "delta.thinking", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
} else {
// Transition from another state to thinking
// First, close any existing content block
if (*param).(*Params).ResponseType != 0 {
if (*param).(*Params).ResponseType == 2 {
// output = output + "event: content_block_delta\n"
// output = output + fmt.Sprintf(`data: {"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":null}}`, (*param).(*Params).ResponseIndex)
// output = output + "\n\n\n"
}
output = output + "event: content_block_stop\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
sb.WriteString("event: content_block_stop\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
(*param).(*Params).ResponseIndex++
}
// Start a new thinking content block
output = output + "event: content_block_start\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"thinking","thinking":""}}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
output = output + "event: content_block_delta\n"
sb.WriteString("event: content_block_start\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"thinking","thinking":""}}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
sb.WriteString("event: content_block_delta\n")
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"thinking_delta","thinking":""}}`, (*param).(*Params).ResponseIndex), "delta.thinking", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
(*param).(*Params).ResponseType = 2 // Set state to thinking
}
} else {
// Process regular text content (user-visible output)
// Continue existing text block if already in content state
if (*param).(*Params).ResponseType == 1 {
output = output + "event: content_block_delta\n"
sb.WriteString("event: content_block_delta\n")
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"text_delta","text":""}}`, (*param).(*Params).ResponseIndex), "delta.text", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
} else {
// Transition from another state to text content
// First, close any existing content block
if (*param).(*Params).ResponseType != 0 {
if (*param).(*Params).ResponseType == 2 {
// output = output + "event: content_block_delta\n"
// output = output + fmt.Sprintf(`data: {"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":null}}`, (*param).(*Params).ResponseIndex)
// output = output + "\n\n\n"
}
output = output + "event: content_block_stop\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
sb.WriteString("event: content_block_stop\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
(*param).(*Params).ResponseIndex++
}
// Start a new text content block
output = output + "event: content_block_start\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"text","text":""}}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
output = output + "event: content_block_delta\n"
sb.WriteString("event: content_block_start\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"text","text":""}}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
sb.WriteString("event: content_block_delta\n")
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"text_delta","text":""}}`, (*param).(*Params).ResponseIndex), "delta.text", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
(*param).(*Params).ResponseType = 1 // Set state to content
}
}
@@ -169,42 +159,35 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
// Handle state transitions when switching to function calls
// Close any existing function call block first
if (*param).(*Params).ResponseType == 3 {
output = output + "event: content_block_stop\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
sb.WriteString("event: content_block_stop\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
(*param).(*Params).ResponseIndex++
(*param).(*Params).ResponseType = 0
}
// Special handling for thinking state transition
if (*param).(*Params).ResponseType == 2 {
// output = output + "event: content_block_delta\n"
// output = output + fmt.Sprintf(`data: {"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":null}}`, (*param).(*Params).ResponseIndex)
// output = output + "\n\n\n"
}
// Close any other existing content block
if (*param).(*Params).ResponseType != 0 {
output = output + "event: content_block_stop\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
sb.WriteString("event: content_block_stop\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
(*param).(*Params).ResponseIndex++
}
// Start a new tool use content block
// This creates the structure for a function call in Claude Code format
output = output + "event: content_block_start\n"
sb.WriteString("event: content_block_start\n")
// Create the tool use block with unique ID and function details
data := fmt.Sprintf(`{"type":"content_block_start","index":%d,"content_block":{"type":"tool_use","id":"","name":"","input":{}}}`, (*param).(*Params).ResponseIndex)
data, _ = sjson.Set(data, "content_block.id", fmt.Sprintf("%s-%d", fcName, time.Now().UnixNano()))
data, _ = sjson.Set(data, "content_block.name", fcName)
output = output + fmt.Sprintf("data: %s\n\n\n", data)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
if fcArgsResult := functionCallResult.Get("args"); fcArgsResult.Exists() {
output = output + "event: content_block_delta\n"
sb.WriteString("event: content_block_delta\n")
data, _ = sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"input_json_delta","partial_json":""}}`, (*param).(*Params).ResponseIndex), "delta.partial_json", fcArgsResult.Raw)
output = output + fmt.Sprintf("data: %s\n\n\n", data)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
}
(*param).(*Params).ResponseType = 3
}
@@ -216,13 +199,13 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
if usageResult.Exists() && bytes.Contains(rawJSON, []byte(`"finishReason"`)) {
if candidatesTokenCountResult := usageResult.Get("candidatesTokenCount"); candidatesTokenCountResult.Exists() {
// Close the final content block
output = output + "event: content_block_stop\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
sb.WriteString("event: content_block_stop\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
// Send the final message delta with usage information and stop reason
output = output + "event: message_delta\n"
output = output + `data: `
sb.WriteString("event: message_delta\n")
sb.WriteString(`data: `)
// Create the message delta template with appropriate stop reason
template := `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":0,"output_tokens":0}}`
@@ -236,11 +219,11 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
template, _ = sjson.Set(template, "usage.output_tokens", candidatesTokenCountResult.Int()+thoughtsTokenCount)
template, _ = sjson.Set(template, "usage.input_tokens", usageResult.Get("promptTokenCount").Int())
output = output + template + "\n\n\n"
sb.WriteString(template + "\n\n\n")
}
}
return []string{output}
return []string{sb.String()}
}
// ConvertGeminiCLIResponseToClaudeNonStream converts a non-streaming Gemini CLI response to a non-streaming Claude response.

View File

@@ -10,6 +10,7 @@ import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
. "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/gemini/openai/chat-completions"
@@ -75,8 +76,8 @@ func ConvertCliResponseToOpenAI(_ context.Context, _ string, originalRequestRawJ
// Extract and set the finish reason.
if finishReasonResult := gjson.GetBytes(rawJSON, "response.candidates.0.finishReason"); finishReasonResult.Exists() {
template, _ = sjson.Set(template, "choices.0.finish_reason", finishReasonResult.String())
template, _ = sjson.Set(template, "choices.0.native_finish_reason", finishReasonResult.String())
template, _ = sjson.Set(template, "choices.0.finish_reason", strings.ToLower(finishReasonResult.String()))
template, _ = sjson.Set(template, "choices.0.native_finish_reason", strings.ToLower(finishReasonResult.String()))
}
// Extract and set usage metadata (token counts).

View File

@@ -91,6 +91,11 @@ func ConvertGeminiRequestToGemini(_ string, inputRawJSON []byte, _ bool) []byte
return true
})
if gjson.GetBytes(rawJSON, "generationConfig.responseSchema").Exists() {
strJson, _ := util.RenameKey(string(out), "generationConfig.responseSchema", "generationConfig.responseJsonSchema")
out = []byte(strJson)
}
out = common.AttachDefaultSafetySettings(out, "safetySettings")
return out
}

View File

@@ -10,6 +10,7 @@ import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
"github.com/tidwall/gjson"
@@ -78,8 +79,8 @@ func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestR
// Extract and set the finish reason.
if finishReasonResult := gjson.GetBytes(rawJSON, "candidates.0.finishReason"); finishReasonResult.Exists() {
template, _ = sjson.Set(template, "choices.0.finish_reason", finishReasonResult.String())
template, _ = sjson.Set(template, "choices.0.native_finish_reason", finishReasonResult.String())
template, _ = sjson.Set(template, "choices.0.finish_reason", strings.ToLower(finishReasonResult.String()))
template, _ = sjson.Set(template, "choices.0.native_finish_reason", strings.ToLower(finishReasonResult.String()))
}
// Extract and set usage metadata (token counts).
@@ -230,8 +231,8 @@ func ConvertGeminiResponseToOpenAINonStream(_ context.Context, _ string, origina
}
if finishReasonResult := gjson.GetBytes(rawJSON, "candidates.0.finishReason"); finishReasonResult.Exists() {
template, _ = sjson.Set(template, "choices.0.finish_reason", finishReasonResult.String())
template, _ = sjson.Set(template, "choices.0.native_finish_reason", finishReasonResult.String())
template, _ = sjson.Set(template, "choices.0.finish_reason", strings.ToLower(finishReasonResult.String()))
template, _ = sjson.Set(template, "choices.0.native_finish_reason", strings.ToLower(finishReasonResult.String()))
}
if usageResult := gjson.GetBytes(rawJSON, "usageMetadata"); usageResult.Exists() {

View File

@@ -249,6 +249,7 @@ func ConvertOpenAIResponsesRequestToGemini(modelName string, inputRawJSON []byte
functionCall := `{"functionCall":{"name":"","args":{}}}`
functionCall, _ = sjson.Set(functionCall, "functionCall.name", name)
functionCall, _ = sjson.Set(functionCall, "thoughtSignature", geminiResponsesThoughtSignature)
functionCall, _ = sjson.Set(functionCall, "functionCall.id", item.Get("call_id").String())
// Parse arguments JSON string and set as args object
if arguments != "" {
@@ -285,6 +286,7 @@ func ConvertOpenAIResponsesRequestToGemini(modelName string, inputRawJSON []byte
}
functionResponse, _ = sjson.Set(functionResponse, "functionResponse.name", functionName)
functionResponse, _ = sjson.Set(functionResponse, "functionResponse.id", callID)
// Set the raw JSON output directly (preserves string encoding)
if outputRaw != "" && outputRaw != "null" {

View File

@@ -33,4 +33,7 @@ import (
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/gemini"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/openai/chat-completions"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/openai/responses"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/claude"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/openai/chat-completions"
)

View File

@@ -0,0 +1,19 @@
package claude
import (
. "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
"github.com/router-for-me/CLIProxyAPI/v6/internal/translator/translator"
)
func init() {
translator.Register(
Claude,
Kiro,
ConvertClaudeRequestToKiro,
interfaces.TranslateResponse{
Stream: ConvertKiroResponseToClaude,
NonStream: ConvertKiroResponseToClaudeNonStream,
},
)
}

View File

@@ -0,0 +1,24 @@
// Package claude provides translation between Kiro and Claude formats.
// Since Kiro uses Claude-compatible format internally, translations are mostly pass-through.
package claude
import (
"bytes"
"context"
)
// ConvertClaudeRequestToKiro converts Claude request to Kiro format.
// Since Kiro uses Claude format internally, this is mostly a pass-through.
func ConvertClaudeRequestToKiro(modelName string, inputRawJSON []byte, stream bool) []byte {
return bytes.Clone(inputRawJSON)
}
// ConvertKiroResponseToClaude converts Kiro streaming response to Claude format.
func ConvertKiroResponseToClaude(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) []string {
return []string{string(rawResponse)}
}
// ConvertKiroResponseToClaudeNonStream converts Kiro non-streaming response to Claude format.
func ConvertKiroResponseToClaudeNonStream(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) string {
return string(rawResponse)
}

View File

@@ -0,0 +1,19 @@
package chat_completions
import (
. "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
"github.com/router-for-me/CLIProxyAPI/v6/internal/translator/translator"
)
func init() {
translator.Register(
OpenAI,
Kiro,
ConvertOpenAIRequestToKiro,
interfaces.TranslateResponse{
Stream: ConvertKiroResponseToOpenAI,
NonStream: ConvertKiroResponseToOpenAINonStream,
},
)
}

View File

@@ -0,0 +1,319 @@
// Package chat_completions provides request translation from OpenAI to Kiro format.
package chat_completions
import (
"bytes"
"encoding/json"
"strings"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// ConvertOpenAIRequestToKiro transforms an OpenAI Chat Completions API request into Kiro (Claude) format.
// Kiro uses Claude-compatible format internally, so we primarily pass through to Claude format.
// Supports tool calling: OpenAI tools -> Claude tools, tool_calls -> tool_use, tool messages -> tool_result.
func ConvertOpenAIRequestToKiro(modelName string, inputRawJSON []byte, stream bool) []byte {
rawJSON := bytes.Clone(inputRawJSON)
root := gjson.ParseBytes(rawJSON)
// Build Claude-compatible request
out := `{"model":"","max_tokens":32000,"messages":[]}`
// Set model
out, _ = sjson.Set(out, "model", modelName)
// Copy max_tokens if present
if v := root.Get("max_tokens"); v.Exists() {
out, _ = sjson.Set(out, "max_tokens", v.Int())
}
// Copy temperature if present
if v := root.Get("temperature"); v.Exists() {
out, _ = sjson.Set(out, "temperature", v.Float())
}
// Copy top_p if present
if v := root.Get("top_p"); v.Exists() {
out, _ = sjson.Set(out, "top_p", v.Float())
}
// Convert OpenAI tools to Claude tools format
if tools := root.Get("tools"); tools.Exists() && tools.IsArray() {
claudeTools := make([]interface{}, 0)
for _, tool := range tools.Array() {
if tool.Get("type").String() == "function" {
fn := tool.Get("function")
claudeTool := map[string]interface{}{
"name": fn.Get("name").String(),
"description": fn.Get("description").String(),
}
// Convert parameters to input_schema
if params := fn.Get("parameters"); params.Exists() {
claudeTool["input_schema"] = params.Value()
} else {
claudeTool["input_schema"] = map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{},
}
}
claudeTools = append(claudeTools, claudeTool)
}
}
if len(claudeTools) > 0 {
out, _ = sjson.Set(out, "tools", claudeTools)
}
}
// Process messages
messages := root.Get("messages")
if messages.Exists() && messages.IsArray() {
claudeMessages := make([]interface{}, 0)
var systemPrompt string
// Track pending tool results to merge with next user message
var pendingToolResults []map[string]interface{}
for _, msg := range messages.Array() {
role := msg.Get("role").String()
content := msg.Get("content")
if role == "system" {
// Extract system message
if content.IsArray() {
for _, part := range content.Array() {
if part.Get("type").String() == "text" {
systemPrompt += part.Get("text").String() + "\n"
}
}
} else {
systemPrompt = content.String()
}
continue
}
if role == "tool" {
// OpenAI tool message -> Claude tool_result content block
toolCallID := msg.Get("tool_call_id").String()
toolContent := content.String()
toolResult := map[string]interface{}{
"type": "tool_result",
"tool_use_id": toolCallID,
}
// Handle content - can be string or structured
if content.IsArray() {
contentParts := make([]interface{}, 0)
for _, part := range content.Array() {
if part.Get("type").String() == "text" {
contentParts = append(contentParts, map[string]interface{}{
"type": "text",
"text": part.Get("text").String(),
})
}
}
toolResult["content"] = contentParts
} else {
toolResult["content"] = toolContent
}
pendingToolResults = append(pendingToolResults, toolResult)
continue
}
claudeMsg := map[string]interface{}{
"role": role,
}
// Handle assistant messages with tool_calls
if role == "assistant" && msg.Get("tool_calls").Exists() {
contentParts := make([]interface{}, 0)
// Add text content if present
if content.Exists() && content.String() != "" {
contentParts = append(contentParts, map[string]interface{}{
"type": "text",
"text": content.String(),
})
}
// Convert tool_calls to tool_use blocks
for _, toolCall := range msg.Get("tool_calls").Array() {
toolUseID := toolCall.Get("id").String()
fnName := toolCall.Get("function.name").String()
fnArgs := toolCall.Get("function.arguments").String()
// Parse arguments JSON
var argsMap map[string]interface{}
if err := json.Unmarshal([]byte(fnArgs), &argsMap); err != nil {
argsMap = map[string]interface{}{"raw": fnArgs}
}
contentParts = append(contentParts, map[string]interface{}{
"type": "tool_use",
"id": toolUseID,
"name": fnName,
"input": argsMap,
})
}
claudeMsg["content"] = contentParts
claudeMessages = append(claudeMessages, claudeMsg)
continue
}
// Handle user messages - may need to include pending tool results
if role == "user" && len(pendingToolResults) > 0 {
contentParts := make([]interface{}, 0)
// Add pending tool results first
for _, tr := range pendingToolResults {
contentParts = append(contentParts, tr)
}
pendingToolResults = nil
// Add user content
if content.IsArray() {
for _, part := range content.Array() {
partType := part.Get("type").String()
if partType == "text" {
contentParts = append(contentParts, map[string]interface{}{
"type": "text",
"text": part.Get("text").String(),
})
} else if partType == "image_url" {
imageURL := part.Get("image_url.url").String()
// Check if it's base64 format (data:image/png;base64,xxxxx)
if strings.HasPrefix(imageURL, "data:") {
// Parse data URL format
// Format: data:image/png;base64,xxxxx
commaIdx := strings.Index(imageURL, ",")
if commaIdx != -1 {
// Extract media_type (e.g., "image/png")
header := imageURL[5:commaIdx] // Remove "data:" prefix
mediaType := header
if semiIdx := strings.Index(header, ";"); semiIdx != -1 {
mediaType = header[:semiIdx]
}
// Extract base64 data
base64Data := imageURL[commaIdx+1:]
contentParts = append(contentParts, map[string]interface{}{
"type": "image",
"source": map[string]interface{}{
"type": "base64",
"media_type": mediaType,
"data": base64Data,
},
})
}
} else {
// Regular URL format - keep original logic
contentParts = append(contentParts, map[string]interface{}{
"type": "image",
"source": map[string]interface{}{
"type": "url",
"url": imageURL,
},
})
}
}
}
} else if content.String() != "" {
contentParts = append(contentParts, map[string]interface{}{
"type": "text",
"text": content.String(),
})
}
claudeMsg["content"] = contentParts
claudeMessages = append(claudeMessages, claudeMsg)
continue
}
// Handle regular content
if content.IsArray() {
contentParts := make([]interface{}, 0)
for _, part := range content.Array() {
partType := part.Get("type").String()
if partType == "text" {
contentParts = append(contentParts, map[string]interface{}{
"type": "text",
"text": part.Get("text").String(),
})
} else if partType == "image_url" {
imageURL := part.Get("image_url.url").String()
// Check if it's base64 format (data:image/png;base64,xxxxx)
if strings.HasPrefix(imageURL, "data:") {
// Parse data URL format
// Format: data:image/png;base64,xxxxx
commaIdx := strings.Index(imageURL, ",")
if commaIdx != -1 {
// Extract media_type (e.g., "image/png")
header := imageURL[5:commaIdx] // Remove "data:" prefix
mediaType := header
if semiIdx := strings.Index(header, ";"); semiIdx != -1 {
mediaType = header[:semiIdx]
}
// Extract base64 data
base64Data := imageURL[commaIdx+1:]
contentParts = append(contentParts, map[string]interface{}{
"type": "image",
"source": map[string]interface{}{
"type": "base64",
"media_type": mediaType,
"data": base64Data,
},
})
}
} else {
// Regular URL format - keep original logic
contentParts = append(contentParts, map[string]interface{}{
"type": "image",
"source": map[string]interface{}{
"type": "url",
"url": imageURL,
},
})
}
}
}
claudeMsg["content"] = contentParts
} else {
claudeMsg["content"] = content.String()
}
claudeMessages = append(claudeMessages, claudeMsg)
}
// If there are pending tool results without a following user message,
// create a user message with just the tool results
if len(pendingToolResults) > 0 {
contentParts := make([]interface{}, 0)
for _, tr := range pendingToolResults {
contentParts = append(contentParts, tr)
}
claudeMessages = append(claudeMessages, map[string]interface{}{
"role": "user",
"content": contentParts,
})
}
out, _ = sjson.Set(out, "messages", claudeMessages)
if systemPrompt != "" {
out, _ = sjson.Set(out, "system", systemPrompt)
}
}
// Set stream
out, _ = sjson.Set(out, "stream", stream)
return []byte(out)
}

View File

@@ -0,0 +1,316 @@
// Package chat_completions provides response translation from Kiro to OpenAI format.
package chat_completions
import (
"context"
"encoding/json"
"time"
"github.com/google/uuid"
"github.com/tidwall/gjson"
)
// ConvertKiroResponseToOpenAI converts Kiro streaming response to OpenAI SSE format.
// Handles Claude SSE events: content_block_start, content_block_delta, input_json_delta,
// content_block_stop, message_delta, and message_stop.
func ConvertKiroResponseToOpenAI(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) []string {
root := gjson.ParseBytes(rawResponse)
var results []string
eventType := root.Get("type").String()
switch eventType {
case "message_start":
// Initial message event - could emit initial chunk if needed
return results
case "content_block_start":
// Start of a content block (text or tool_use)
blockType := root.Get("content_block.type").String()
index := int(root.Get("index").Int())
if blockType == "tool_use" {
// Start of tool_use block
toolUseID := root.Get("content_block.id").String()
toolName := root.Get("content_block.name").String()
toolCall := map[string]interface{}{
"index": index,
"id": toolUseID,
"type": "function",
"function": map[string]interface{}{
"name": toolName,
"arguments": "",
},
}
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion.chunk",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"delta": map[string]interface{}{
"tool_calls": []map[string]interface{}{toolCall},
},
"finish_reason": nil,
},
},
}
result, _ := json.Marshal(response)
results = append(results, string(result))
}
return results
case "content_block_delta":
index := int(root.Get("index").Int())
deltaType := root.Get("delta.type").String()
if deltaType == "text_delta" {
// Text content delta
contentDelta := root.Get("delta.text").String()
if contentDelta != "" {
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion.chunk",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"delta": map[string]interface{}{
"content": contentDelta,
},
"finish_reason": nil,
},
},
}
result, _ := json.Marshal(response)
results = append(results, string(result))
}
} else if deltaType == "input_json_delta" {
// Tool input delta (streaming arguments)
partialJSON := root.Get("delta.partial_json").String()
if partialJSON != "" {
toolCall := map[string]interface{}{
"index": index,
"function": map[string]interface{}{
"arguments": partialJSON,
},
}
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion.chunk",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"delta": map[string]interface{}{
"tool_calls": []map[string]interface{}{toolCall},
},
"finish_reason": nil,
},
},
}
result, _ := json.Marshal(response)
results = append(results, string(result))
}
}
return results
case "content_block_stop":
// End of content block - no output needed for OpenAI format
return results
case "message_delta":
// Final message delta with stop_reason
stopReason := root.Get("delta.stop_reason").String()
if stopReason != "" {
finishReason := "stop"
if stopReason == "tool_use" {
finishReason = "tool_calls"
} else if stopReason == "end_turn" {
finishReason = "stop"
} else if stopReason == "max_tokens" {
finishReason = "length"
}
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion.chunk",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"delta": map[string]interface{}{},
"finish_reason": finishReason,
},
},
}
result, _ := json.Marshal(response)
results = append(results, string(result))
}
return results
case "message_stop":
// End of message - could emit [DONE] marker
return results
}
// Fallback: handle raw content for backward compatibility
var contentDelta string
if delta := root.Get("delta.text"); delta.Exists() {
contentDelta = delta.String()
} else if content := root.Get("content"); content.Exists() && root.Get("type").String() == "" {
contentDelta = content.String()
}
if contentDelta != "" {
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion.chunk",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"delta": map[string]interface{}{
"content": contentDelta,
},
"finish_reason": nil,
},
},
}
result, _ := json.Marshal(response)
results = append(results, string(result))
}
// Handle tool_use content blocks (Claude format) - fallback
toolUses := root.Get("delta.tool_use")
if !toolUses.Exists() {
toolUses = root.Get("tool_use")
}
if toolUses.Exists() && toolUses.IsObject() {
inputJSON := toolUses.Get("input").String()
if inputJSON == "" {
if inputObj := toolUses.Get("input"); inputObj.Exists() {
inputBytes, _ := json.Marshal(inputObj.Value())
inputJSON = string(inputBytes)
}
}
toolCall := map[string]interface{}{
"index": 0,
"id": toolUses.Get("id").String(),
"type": "function",
"function": map[string]interface{}{
"name": toolUses.Get("name").String(),
"arguments": inputJSON,
},
}
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion.chunk",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"delta": map[string]interface{}{
"tool_calls": []map[string]interface{}{toolCall},
},
"finish_reason": nil,
},
},
}
result, _ := json.Marshal(response)
results = append(results, string(result))
}
return results
}
// ConvertKiroResponseToOpenAINonStream converts Kiro non-streaming response to OpenAI format.
func ConvertKiroResponseToOpenAINonStream(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) string {
root := gjson.ParseBytes(rawResponse)
var content string
var toolCalls []map[string]interface{}
contentArray := root.Get("content")
if contentArray.IsArray() {
for _, item := range contentArray.Array() {
itemType := item.Get("type").String()
if itemType == "text" {
content += item.Get("text").String()
} else if itemType == "tool_use" {
// Convert Claude tool_use to OpenAI tool_calls format
inputJSON := item.Get("input").String()
if inputJSON == "" {
// If input is an object, marshal it
if inputObj := item.Get("input"); inputObj.Exists() {
inputBytes, _ := json.Marshal(inputObj.Value())
inputJSON = string(inputBytes)
}
}
toolCall := map[string]interface{}{
"id": item.Get("id").String(),
"type": "function",
"function": map[string]interface{}{
"name": item.Get("name").String(),
"arguments": inputJSON,
},
}
toolCalls = append(toolCalls, toolCall)
}
}
} else {
content = root.Get("content").String()
}
inputTokens := root.Get("usage.input_tokens").Int()
outputTokens := root.Get("usage.output_tokens").Int()
message := map[string]interface{}{
"role": "assistant",
"content": content,
}
// Add tool_calls if present
if len(toolCalls) > 0 {
message["tool_calls"] = toolCalls
}
finishReason := "stop"
if len(toolCalls) > 0 {
finishReason = "tool_calls"
}
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"message": message,
"finish_reason": finishReason,
},
},
"usage": map[string]interface{}{
"prompt_tokens": inputTokens,
"completion_tokens": outputTokens,
"total_tokens": inputTokens + outputTokens,
},
}
result, _ := json.Marshal(response)
return string(result)
}

View File

@@ -8,6 +8,7 @@ package claude
import (
"bytes"
"encoding/json"
"strings"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
@@ -242,11 +243,12 @@ func convertClaudeContentPart(part gjson.Result) (string, bool) {
switch partType {
case "text":
if !part.Get("text").Exists() {
text := part.Get("text").String()
if strings.TrimSpace(text) == "" {
return "", false
}
textContent := `{"type":"text","text":""}`
textContent, _ = sjson.Set(textContent, "text", part.Get("text").String())
textContent, _ = sjson.Set(textContent, "text", text)
return textContent, true
case "image":

View File

@@ -34,6 +34,15 @@ func ParseGeminiThinkingSuffix(model string) (string, *int, *bool, bool) {
return base, &budgetValue, &include, true
}
// Handle "-reasoning" suffix: enables thinking with dynamic budget (-1)
// Maps: gemini-2.5-flash-reasoning -> gemini-2.5-flash with thinkingBudget=-1
if strings.HasSuffix(lower, "-reasoning") {
base := model[:len(model)-len("-reasoning")]
budgetValue := -1 // Dynamic budget
include := true
return base, &budgetValue, &include, true
}
idx := strings.LastIndex(lower, "-thinking-")
if idx == -1 {
return model, nil, nil, false

View File

@@ -79,6 +79,15 @@ func RenameKey(jsonStr, oldKeyPath, newKeyPath string) (string, error) {
return finalJson, nil
}
func DeleteKey(jsonStr, keyName string) string {
paths := make([]string, 0)
Walk(gjson.Parse(jsonStr), "", keyName, &paths)
for _, p := range paths {
jsonStr, _ = sjson.Delete(jsonStr, p)
}
return jsonStr
}
// FixJSON converts non-standard JSON that uses single quotes for strings into
// RFC 8259-compliant JSON by converting those single-quoted strings to
// double-quoted strings with proper escaping.

View File

@@ -21,6 +21,7 @@ import (
"github.com/fsnotify/fsnotify"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"
"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/geminicli"
"gopkg.in/yaml.v3"
@@ -176,6 +177,9 @@ func (w *Watcher) Start(ctx context.Context) error {
}
log.Debugf("watching auth directory: %s", w.authDir)
// Watch Kiro IDE token file directory for automatic token updates
w.watchKiroIDETokenFile()
// Start the event processing goroutine
go w.processEvents(ctx)
@@ -184,6 +188,31 @@ func (w *Watcher) Start(ctx context.Context) error {
return nil
}
// watchKiroIDETokenFile adds the Kiro IDE token file directory to the watcher.
// This enables automatic detection of token updates from Kiro IDE.
func (w *Watcher) watchKiroIDETokenFile() {
homeDir, err := os.UserHomeDir()
if err != nil {
log.Debugf("failed to get home directory for Kiro IDE token watch: %v", err)
return
}
// Kiro IDE stores tokens in ~/.aws/sso/cache/
kiroTokenDir := filepath.Join(homeDir, ".aws", "sso", "cache")
// Check if directory exists
if _, statErr := os.Stat(kiroTokenDir); os.IsNotExist(statErr) {
log.Debugf("Kiro IDE token directory does not exist: %s", kiroTokenDir)
return
}
if errAdd := w.watcher.Add(kiroTokenDir); errAdd != nil {
log.Debugf("failed to watch Kiro IDE token directory %s: %v", kiroTokenDir, errAdd)
return
}
log.Debugf("watching Kiro IDE token directory: %s", kiroTokenDir)
}
// Stop stops the file watcher
func (w *Watcher) Stop() error {
w.stopDispatch()
@@ -496,6 +525,18 @@ func computeOpenAICompatModelsHash(models []config.OpenAICompatibilityModel) str
return hex.EncodeToString(sum[:])
}
func computeVertexCompatModelsHash(models []config.VertexCompatModel) string {
if len(models) == 0 {
return ""
}
data, err := json.Marshal(models)
if err != nil || len(data) == 0 {
return ""
}
sum := sha256.Sum256(data)
return hex.EncodeToString(sum[:])
}
// computeClaudeModelsHash returns a stable hash for Claude model aliases.
func computeClaudeModelsHash(models []config.ClaudeModel) string {
if len(models) == 0 {
@@ -558,6 +599,35 @@ func summarizeExcludedModels(list []string) excludedModelsSummary {
}
}
type ampModelMappingsSummary struct {
hash string
count int
}
func summarizeAmpModelMappings(mappings []config.AmpModelMapping) ampModelMappingsSummary {
if len(mappings) == 0 {
return ampModelMappingsSummary{}
}
entries := make([]string, 0, len(mappings))
for _, mapping := range mappings {
from := strings.TrimSpace(mapping.From)
to := strings.TrimSpace(mapping.To)
if from == "" && to == "" {
continue
}
entries = append(entries, from+"->"+to)
}
if len(entries) == 0 {
return ampModelMappingsSummary{}
}
sort.Strings(entries)
sum := sha256.Sum256([]byte(strings.Join(entries, "|")))
return ampModelMappingsSummary{
hash: hex.EncodeToString(sum[:]),
count: len(entries),
}
}
func summarizeOAuthExcludedModels(entries map[string][]string) map[string]excludedModelsSummary {
if len(entries) == 0 {
return nil
@@ -703,10 +773,20 @@ func (w *Watcher) handleEvent(event fsnotify.Event) {
isConfigEvent := event.Name == w.configPath && event.Op&configOps != 0
authOps := fsnotify.Create | fsnotify.Write | fsnotify.Remove | fsnotify.Rename
isAuthJSON := strings.HasPrefix(event.Name, w.authDir) && strings.HasSuffix(event.Name, ".json") && event.Op&authOps != 0
if !isConfigEvent && !isAuthJSON {
// Check for Kiro IDE token file changes
isKiroIDEToken := w.isKiroIDETokenFile(event.Name) && event.Op&authOps != 0
if !isConfigEvent && !isAuthJSON && !isKiroIDEToken {
// Ignore unrelated files (e.g., cookie snapshots *.cookie) and other noise.
return
}
// Handle Kiro IDE token file changes
if isKiroIDEToken {
w.handleKiroIDETokenChange(event)
return
}
now := time.Now()
log.Debugf("file system event detected: %s %s", event.Op.String(), event.Name)
@@ -764,6 +844,51 @@ func (w *Watcher) scheduleConfigReload() {
})
}
// isKiroIDETokenFile checks if the given path is the Kiro IDE token file.
func (w *Watcher) isKiroIDETokenFile(path string) bool {
// Check if it's the kiro-auth-token.json file in ~/.aws/sso/cache/
// Use filepath.ToSlash to ensure consistent separators across platforms (Windows uses backslashes)
normalized := filepath.ToSlash(path)
return strings.HasSuffix(normalized, "kiro-auth-token.json") && strings.Contains(normalized, ".aws/sso/cache")
}
// handleKiroIDETokenChange processes changes to the Kiro IDE token file.
// When the token file is updated by Kiro IDE, this triggers a reload of Kiro auth.
func (w *Watcher) handleKiroIDETokenChange(event fsnotify.Event) {
log.Debugf("Kiro IDE token file event detected: %s %s", event.Op.String(), event.Name)
if event.Op&(fsnotify.Remove|fsnotify.Rename) != 0 {
// Token file removed - wait briefly for potential atomic replace
time.Sleep(replaceCheckDelay)
if _, statErr := os.Stat(event.Name); statErr != nil {
log.Debugf("Kiro IDE token file removed: %s", event.Name)
return
}
}
// Try to load the updated token
tokenData, err := kiroauth.LoadKiroIDEToken()
if err != nil {
log.Debugf("failed to load Kiro IDE token after change: %v", err)
return
}
log.Infof("Kiro IDE token file updated, access token refreshed (provider: %s)", tokenData.Provider)
// Trigger auth state refresh to pick up the new token
w.refreshAuthState()
// Notify callback if set
w.clientsMutex.RLock()
cfg := w.config
w.clientsMutex.RUnlock()
if w.reloadCallback != nil && cfg != nil {
log.Debugf("triggering server update callback after Kiro IDE token change")
w.reloadCallback(cfg)
}
}
func (w *Watcher) reloadConfigIfChanged() {
data, err := os.ReadFile(w.configPath)
if err != nil {
@@ -902,8 +1027,8 @@ func (w *Watcher) reloadClients(rescanAuth bool, affectedOAuthProviders []string
// no legacy clients to unregister
// Create new API key clients based on the new config
geminiAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount := BuildAPIKeyClients(cfg)
totalAPIKeyClients := geminiAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + openAICompatCount
geminiAPIKeyCount, vertexCompatAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount := BuildAPIKeyClients(cfg)
totalAPIKeyClients := geminiAPIKeyCount + vertexCompatAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + openAICompatCount
log.Debugf("loaded %d API key clients", totalAPIKeyClients)
var authFileCount int
@@ -946,7 +1071,7 @@ func (w *Watcher) reloadClients(rescanAuth bool, affectedOAuthProviders []string
w.clientsMutex.Unlock()
}
totalNewClients := authFileCount + geminiAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + openAICompatCount
totalNewClients := authFileCount + geminiAPIKeyCount + vertexCompatAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + openAICompatCount
// Ensure consumers observe the new configuration before auth updates dispatch.
if w.reloadCallback != nil {
@@ -956,10 +1081,11 @@ func (w *Watcher) reloadClients(rescanAuth bool, affectedOAuthProviders []string
w.refreshAuthState()
log.Infof("full client load complete - %d clients (%d auth files + %d Gemini API keys + %d Claude API keys + %d Codex keys + %d OpenAI-compat)",
log.Infof("full client load complete - %d clients (%d auth files + %d Gemini API keys + %d Vertex API keys + %d Claude API keys + %d Codex keys + %d OpenAI-compat)",
totalNewClients,
authFileCount,
geminiAPIKeyCount,
vertexCompatAPIKeyCount,
claudeAPIKeyCount,
codexAPIKeyCount,
openAICompatCount,
@@ -1074,6 +1200,7 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
applyAuthExcludedModelsMeta(a, cfg, entry.ExcludedModels, "apikey")
out = append(out, a)
}
// Claude API keys -> synthesize auths
for i := range cfg.ClaudeKey {
ck := cfg.ClaudeKey[i]
@@ -1138,6 +1265,67 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
applyAuthExcludedModelsMeta(a, cfg, ck.ExcludedModels, "apikey")
out = append(out, a)
}
// Kiro (AWS CodeWhisperer) -> synthesize auths
var kAuth *kiroauth.KiroAuth
if len(cfg.KiroKey) > 0 {
kAuth = kiroauth.NewKiroAuth(cfg)
}
for i := range cfg.KiroKey {
kk := cfg.KiroKey[i]
var accessToken, profileArn string
// Try to load from token file first
if kk.TokenFile != "" && kAuth != nil {
tokenData, err := kAuth.LoadTokenFromFile(kk.TokenFile)
if err != nil {
log.Warnf("failed to load kiro token file %s: %v", kk.TokenFile, err)
} else {
accessToken = tokenData.AccessToken
profileArn = tokenData.ProfileArn
}
}
// Override with direct config values if provided
if kk.AccessToken != "" {
accessToken = kk.AccessToken
}
if kk.ProfileArn != "" {
profileArn = kk.ProfileArn
}
if accessToken == "" {
log.Warnf("kiro config[%d] missing access_token, skipping", i)
continue
}
// profileArn is optional for AWS Builder ID users
id, token := idGen.next("kiro:token", accessToken, profileArn)
attrs := map[string]string{
"source": fmt.Sprintf("config:kiro[%s]", token),
"access_token": accessToken,
}
if profileArn != "" {
attrs["profile_arn"] = profileArn
}
if kk.Region != "" {
attrs["region"] = kk.Region
}
if kk.AgentTaskType != "" {
attrs["agent_task_type"] = kk.AgentTaskType
}
proxyURL := strings.TrimSpace(kk.ProxyURL)
a := &coreauth.Auth{
ID: id,
Provider: "kiro",
Label: "kiro-token",
Status: coreauth.StatusActive,
ProxyURL: proxyURL,
Attributes: attrs,
CreatedAt: now,
UpdatedAt: now,
}
out = append(out, a)
}
for i := range cfg.OpenAICompatibility {
compat := &cfg.OpenAICompatibility[i]
providerName := strings.ToLower(strings.TrimSpace(compat.Name))
@@ -1148,71 +1336,37 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
// Handle new APIKeyEntries format (preferred)
createdEntries := 0
if len(compat.APIKeyEntries) > 0 {
for j := range compat.APIKeyEntries {
entry := &compat.APIKeyEntries[j]
key := strings.TrimSpace(entry.APIKey)
proxyURL := strings.TrimSpace(entry.ProxyURL)
idKind := fmt.Sprintf("openai-compatibility:%s", providerName)
id, token := idGen.next(idKind, key, base, proxyURL)
attrs := map[string]string{
"source": fmt.Sprintf("config:%s[%s]", providerName, token),
"base_url": base,
"compat_name": compat.Name,
"provider_key": providerName,
}
if key != "" {
attrs["api_key"] = key
}
if hash := computeOpenAICompatModelsHash(compat.Models); hash != "" {
attrs["models_hash"] = hash
}
addConfigHeadersToAttrs(compat.Headers, attrs)
a := &coreauth.Auth{
ID: id,
Provider: providerName,
Label: compat.Name,
Status: coreauth.StatusActive,
ProxyURL: proxyURL,
Attributes: attrs,
CreatedAt: now,
UpdatedAt: now,
}
out = append(out, a)
createdEntries++
for j := range compat.APIKeyEntries {
entry := &compat.APIKeyEntries[j]
key := strings.TrimSpace(entry.APIKey)
proxyURL := strings.TrimSpace(entry.ProxyURL)
idKind := fmt.Sprintf("openai-compatibility:%s", providerName)
id, token := idGen.next(idKind, key, base, proxyURL)
attrs := map[string]string{
"source": fmt.Sprintf("config:%s[%s]", providerName, token),
"base_url": base,
"compat_name": compat.Name,
"provider_key": providerName,
}
} else {
// Handle legacy APIKeys format for backward compatibility
for j := range compat.APIKeys {
key := strings.TrimSpace(compat.APIKeys[j])
if key == "" {
continue
}
idKind := fmt.Sprintf("openai-compatibility:%s", providerName)
id, token := idGen.next(idKind, key, base)
attrs := map[string]string{
"source": fmt.Sprintf("config:%s[%s]", providerName, token),
"base_url": base,
"compat_name": compat.Name,
"provider_key": providerName,
}
if key != "" {
attrs["api_key"] = key
if hash := computeOpenAICompatModelsHash(compat.Models); hash != "" {
attrs["models_hash"] = hash
}
addConfigHeadersToAttrs(compat.Headers, attrs)
a := &coreauth.Auth{
ID: id,
Provider: providerName,
Label: compat.Name,
Status: coreauth.StatusActive,
Attributes: attrs,
CreatedAt: now,
UpdatedAt: now,
}
out = append(out, a)
createdEntries++
}
if hash := computeOpenAICompatModelsHash(compat.Models); hash != "" {
attrs["models_hash"] = hash
}
addConfigHeadersToAttrs(compat.Headers, attrs)
a := &coreauth.Auth{
ID: id,
Provider: providerName,
Label: compat.Name,
Status: coreauth.StatusActive,
ProxyURL: proxyURL,
Attributes: attrs,
CreatedAt: now,
UpdatedAt: now,
}
out = append(out, a)
createdEntries++
}
if createdEntries == 0 {
idKind := fmt.Sprintf("openai-compatibility:%s", providerName)
@@ -1240,8 +1394,50 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
}
}
}
// Process Vertex API key providers (Vertex-compatible endpoints)
for i := range cfg.VertexCompatAPIKey {
compat := &cfg.VertexCompatAPIKey[i]
providerName := "vertex"
base := strings.TrimSpace(compat.BaseURL)
key := strings.TrimSpace(compat.APIKey)
proxyURL := strings.TrimSpace(compat.ProxyURL)
idKind := fmt.Sprintf("vertex:apikey:%s", base)
id, token := idGen.next(idKind, key, base, proxyURL)
attrs := map[string]string{
"source": fmt.Sprintf("config:vertex-apikey[%s]", token),
"base_url": base,
"provider_key": providerName,
}
if key != "" {
attrs["api_key"] = key
}
if hash := computeVertexCompatModelsHash(compat.Models); hash != "" {
attrs["models_hash"] = hash
}
addConfigHeadersToAttrs(compat.Headers, attrs)
a := &coreauth.Auth{
ID: id,
Provider: providerName,
Label: "vertex-apikey",
Status: coreauth.StatusActive,
ProxyURL: proxyURL,
Attributes: attrs,
CreatedAt: now,
UpdatedAt: now,
}
applyAuthExcludedModelsMeta(a, cfg, nil, "apikey")
out = append(out, a)
}
// Also synthesize auth entries directly from auth files (for OAuth/file-backed providers)
entries, _ := os.ReadDir(w.authDir)
log.Debugf("SnapshotCoreAuths: scanning auth directory: %s", w.authDir)
entries, readErr := os.ReadDir(w.authDir)
if readErr != nil {
log.Errorf("SnapshotCoreAuths: failed to read auth directory %s: %v", w.authDir, readErr)
}
log.Debugf("SnapshotCoreAuths: found %d entries in auth directory", len(entries))
for _, e := range entries {
if e.IsDir() {
continue
@@ -1260,9 +1456,20 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
continue
}
t, _ := metadata["type"].(string)
// Detect Kiro auth files by auth_method field (they don't have "type" field)
if t == "" {
if authMethod, _ := metadata["auth_method"].(string); authMethod == "builder-id" || authMethod == "social" {
t = "kiro"
log.Debugf("SnapshotCoreAuths: detected Kiro auth by auth_method: %s", name)
}
}
if t == "" {
log.Debugf("SnapshotCoreAuths: skipping file without type: %s", name)
continue
}
log.Debugf("SnapshotCoreAuths: processing auth file: %s (type=%s)", name, t)
provider := strings.ToLower(t)
if provider == "gemini" {
provider = "gemini-cli"
@@ -1271,6 +1478,12 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
if email, _ := metadata["email"].(string); email != "" {
label = email
}
// For Kiro, use provider field as label if available
if provider == "kiro" {
if kiroProvider, _ := metadata["provider"].(string); kiroProvider != "" {
label = fmt.Sprintf("kiro-%s", strings.ToLower(kiroProvider))
}
}
// Use relative path under authDir as ID to stay consistent with the file-based token store
id := full
if rel, errRel := filepath.Rel(w.authDir, full); errRel == nil && rel != "" {
@@ -1296,6 +1509,16 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
CreatedAt: now,
UpdatedAt: now,
}
// Set NextRefreshAfter for Kiro auth based on expires_at
if provider == "kiro" {
if expiresAtStr, ok := metadata["expires_at"].(string); ok && expiresAtStr != "" {
if expiresAt, parseErr := time.Parse(time.RFC3339, expiresAtStr); parseErr == nil {
// Refresh 30 minutes before expiry
a.NextRefreshAfter = expiresAt.Add(-30 * time.Minute)
}
}
}
applyAuthExcludedModelsMeta(a, cfg, nil, "oauth")
if provider == "gemini-cli" {
if virtuals := synthesizeGeminiVirtualAuths(a, metadata, now); len(virtuals) > 0 {
@@ -1456,8 +1679,9 @@ func (w *Watcher) loadFileClients(cfg *config.Config) int {
return authFileCount
}
func BuildAPIKeyClients(cfg *config.Config) (int, int, int, int) {
func BuildAPIKeyClients(cfg *config.Config) (int, int, int, int, int) {
geminiAPIKeyCount := 0
vertexCompatAPIKeyCount := 0
claudeAPIKeyCount := 0
codexAPIKeyCount := 0
openAICompatCount := 0
@@ -1466,6 +1690,9 @@ func BuildAPIKeyClients(cfg *config.Config) (int, int, int, int) {
// Stateless executor handles Gemini API keys; avoid constructing legacy clients.
geminiAPIKeyCount += len(cfg.GeminiKey)
}
if len(cfg.VertexCompatAPIKey) > 0 {
vertexCompatAPIKeyCount += len(cfg.VertexCompatAPIKey)
}
if len(cfg.ClaudeKey) > 0 {
claudeAPIKeyCount += len(cfg.ClaudeKey)
}
@@ -1475,15 +1702,10 @@ func BuildAPIKeyClients(cfg *config.Config) (int, int, int, int) {
if len(cfg.OpenAICompatibility) > 0 {
// Do not construct legacy clients for OpenAI-compat providers; these are handled by the stateless executor.
for _, compatConfig := range cfg.OpenAICompatibility {
// Count from new APIKeyEntries format if present, otherwise fall back to legacy APIKeys
if len(compatConfig.APIKeyEntries) > 0 {
openAICompatCount += len(compatConfig.APIKeyEntries)
} else {
openAICompatCount += len(compatConfig.APIKeys)
}
openAICompatCount += len(compatConfig.APIKeyEntries)
}
}
return geminiAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount
return geminiAPIKeyCount, vertexCompatAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount
}
func diffOpenAICompatibility(oldList, newList []config.OpenAICompatibility) []string {
@@ -1557,24 +1779,9 @@ func describeOpenAICompatibilityUpdate(oldEntry, newEntry config.OpenAICompatibi
}
func countAPIKeys(entry config.OpenAICompatibility) int {
// Prefer new APIKeyEntries format
if len(entry.APIKeyEntries) > 0 {
count := 0
for _, keyEntry := range entry.APIKeyEntries {
if strings.TrimSpace(keyEntry.APIKey) != "" {
count++
}
}
return count
}
// Fall back to legacy APIKeys format
return countNonEmptyStrings(entry.APIKeys)
}
func countNonEmptyStrings(values []string) int {
count := 0
for _, value := range values {
if strings.TrimSpace(value) != "" {
for _, keyEntry := range entry.APIKeyEntries {
if strings.TrimSpace(keyEntry.APIKey) != "" {
count++
}
}
@@ -1699,9 +1906,6 @@ func buildConfigChangeDetails(oldCfg, newCfg *config.Config) []string {
changes = append(changes, fmt.Sprintf("gemini[%d].excluded-models: updated (%d -> %d entries)", i, oldExcluded.count, newExcluded.count))
}
}
if !reflect.DeepEqual(trimStrings(oldCfg.GlAPIKey), trimStrings(newCfg.GlAPIKey)) {
changes = append(changes, "generative-language-api-key: values updated (legacy view, redacted)")
}
}
// Claude keys (do not print key material)
@@ -1764,6 +1968,31 @@ func buildConfigChangeDetails(oldCfg, newCfg *config.Config) []string {
}
}
// AmpCode settings (redacted where needed)
oldAmpURL := strings.TrimSpace(oldCfg.AmpCode.UpstreamURL)
newAmpURL := strings.TrimSpace(newCfg.AmpCode.UpstreamURL)
if oldAmpURL != newAmpURL {
changes = append(changes, fmt.Sprintf("ampcode.upstream-url: %s -> %s", oldAmpURL, newAmpURL))
}
oldAmpKey := strings.TrimSpace(oldCfg.AmpCode.UpstreamAPIKey)
newAmpKey := strings.TrimSpace(newCfg.AmpCode.UpstreamAPIKey)
switch {
case oldAmpKey == "" && newAmpKey != "":
changes = append(changes, "ampcode.upstream-api-key: added")
case oldAmpKey != "" && newAmpKey == "":
changes = append(changes, "ampcode.upstream-api-key: removed")
case oldAmpKey != newAmpKey:
changes = append(changes, "ampcode.upstream-api-key: updated")
}
if oldCfg.AmpCode.RestrictManagementToLocalhost != newCfg.AmpCode.RestrictManagementToLocalhost {
changes = append(changes, fmt.Sprintf("ampcode.restrict-management-to-localhost: %t -> %t", oldCfg.AmpCode.RestrictManagementToLocalhost, newCfg.AmpCode.RestrictManagementToLocalhost))
}
oldMappings := summarizeAmpModelMappings(oldCfg.AmpCode.ModelMappings)
newMappings := summarizeAmpModelMappings(newCfg.AmpCode.ModelMappings)
if oldMappings.hash != newMappings.hash {
changes = append(changes, fmt.Sprintf("ampcode.model-mappings: updated (%d -> %d entries)", oldMappings.count, newMappings.count))
}
if entries, _ := diffOAuthExcludedModelChanges(oldCfg.OAuthExcludedModels, newCfg.OAuthExcludedModels); len(entries) > 0 {
changes = append(changes, entries...)
}

View File

@@ -8,6 +8,7 @@ import (
"fmt"
"net/http"
"strings"
"sync"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
@@ -50,7 +51,25 @@ type BaseAPIHandler struct {
Cfg *config.SDKConfig
// OpenAICompatProviders is a list of provider names for OpenAI compatibility.
OpenAICompatProviders []string
openAICompatProviders []string
openAICompatMutex sync.RWMutex
}
// GetOpenAICompatProviders safely returns a copy of the provider names
func (h *BaseAPIHandler) GetOpenAICompatProviders() []string {
h.openAICompatMutex.RLock()
defer h.openAICompatMutex.RUnlock()
result := make([]string, len(h.openAICompatProviders))
copy(result, h.openAICompatProviders)
return result
}
// SetOpenAICompatProviders safely sets the provider names
func (h *BaseAPIHandler) SetOpenAICompatProviders(providers []string) {
h.openAICompatMutex.Lock()
defer h.openAICompatMutex.Unlock()
h.openAICompatProviders = make([]string, len(providers))
copy(h.openAICompatProviders, providers)
}
// NewBaseAPIHandlers creates a new API handlers instance.
@@ -63,11 +82,12 @@ type BaseAPIHandler struct {
// Returns:
// - *BaseAPIHandler: A new API handlers instance
func NewBaseAPIHandlers(cfg *config.SDKConfig, authManager *coreauth.Manager, openAICompatProviders []string) *BaseAPIHandler {
return &BaseAPIHandler{
Cfg: cfg,
AuthManager: authManager,
OpenAICompatProviders: openAICompatProviders,
h := &BaseAPIHandler{
Cfg: cfg,
AuthManager: authManager,
}
h.SetOpenAICompatProviders(openAICompatProviders)
return h
}
// UpdateClients updates the handlers' client list and configuration.
@@ -363,7 +383,7 @@ func (h *BaseAPIHandler) parseDynamicModel(modelName string) (providerName, mode
}
// Check if the provider is a configured openai-compatibility provider
for _, pName := range h.OpenAICompatProviders {
for _, pName := range h.GetOpenAICompatProviders() {
if pName == providerPart {
return providerPart, modelPart, true
}

View File

@@ -4,6 +4,7 @@ import (
"context"
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"net/url"
@@ -127,6 +128,18 @@ func (AntigravityAuthenticator) Login(ctx context.Context, cfg *config.Config, o
}
}
// Fetch project ID via loadCodeAssist (same approach as Gemini CLI)
projectID := ""
if tokenResp.AccessToken != "" {
fetchedProjectID, errProject := fetchAntigravityProjectID(ctx, tokenResp.AccessToken, httpClient)
if errProject != nil {
log.Warnf("antigravity: failed to fetch project ID: %v", errProject)
} else {
projectID = fetchedProjectID
log.Infof("antigravity: obtained project ID %s", projectID)
}
}
now := time.Now()
metadata := map[string]any{
"type": "antigravity",
@@ -139,6 +152,9 @@ func (AntigravityAuthenticator) Login(ctx context.Context, cfg *config.Config, o
if email != "" {
metadata["email"] = email
}
if projectID != "" {
metadata["project_id"] = projectID
}
fileName := sanitizeAntigravityFileName(email)
label := email
@@ -147,6 +163,9 @@ func (AntigravityAuthenticator) Login(ctx context.Context, cfg *config.Config, o
}
fmt.Println("Antigravity authentication successful")
if projectID != "" {
fmt.Printf("Using GCP project: %s\n", projectID)
}
return &coreauth.Auth{
ID: fileName,
Provider: "antigravity",
@@ -291,3 +310,89 @@ func sanitizeAntigravityFileName(email string) string {
replacer := strings.NewReplacer("@", "_", ".", "_")
return fmt.Sprintf("antigravity-%s.json", replacer.Replace(email))
}
// Antigravity API constants for project discovery
const (
antigravityAPIEndpoint = "https://cloudcode-pa.googleapis.com"
antigravityAPIVersion = "v1internal"
antigravityAPIUserAgent = "google-api-nodejs-client/9.15.1"
antigravityAPIClient = "google-cloud-sdk vscode_cloudshelleditor/0.1"
antigravityClientMetadata = `{"ideType":"IDE_UNSPECIFIED","platform":"PLATFORM_UNSPECIFIED","pluginType":"GEMINI"}`
)
// FetchAntigravityProjectID exposes project discovery for external callers.
func FetchAntigravityProjectID(ctx context.Context, accessToken string, httpClient *http.Client) (string, error) {
return fetchAntigravityProjectID(ctx, accessToken, httpClient)
}
// fetchAntigravityProjectID retrieves the project ID for the authenticated user via loadCodeAssist.
// This uses the same approach as Gemini CLI to get the cloudaicompanionProject.
func fetchAntigravityProjectID(ctx context.Context, accessToken string, httpClient *http.Client) (string, error) {
// Call loadCodeAssist to get the project
loadReqBody := map[string]any{
"metadata": map[string]string{
"ideType": "IDE_UNSPECIFIED",
"platform": "PLATFORM_UNSPECIFIED",
"pluginType": "GEMINI",
},
}
rawBody, errMarshal := json.Marshal(loadReqBody)
if errMarshal != nil {
return "", fmt.Errorf("marshal request body: %w", errMarshal)
}
endpointURL := fmt.Sprintf("%s/%s:loadCodeAssist", antigravityAPIEndpoint, antigravityAPIVersion)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, strings.NewReader(string(rawBody)))
if err != nil {
return "", fmt.Errorf("create request: %w", err)
}
req.Header.Set("Authorization", "Bearer "+accessToken)
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", antigravityAPIUserAgent)
req.Header.Set("X-Goog-Api-Client", antigravityAPIClient)
req.Header.Set("Client-Metadata", antigravityClientMetadata)
resp, errDo := httpClient.Do(req)
if errDo != nil {
return "", fmt.Errorf("execute request: %w", errDo)
}
defer func() {
if errClose := resp.Body.Close(); errClose != nil {
log.Errorf("antigravity loadCodeAssist: close body error: %v", errClose)
}
}()
bodyBytes, errRead := io.ReadAll(resp.Body)
if errRead != nil {
return "", fmt.Errorf("read response: %w", errRead)
}
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
return "", fmt.Errorf("request failed with status %d: %s", resp.StatusCode, strings.TrimSpace(string(bodyBytes)))
}
var loadResp map[string]any
if errDecode := json.Unmarshal(bodyBytes, &loadResp); errDecode != nil {
return "", fmt.Errorf("decode response: %w", errDecode)
}
// Extract projectID from response
projectID := ""
if id, ok := loadResp["cloudaicompanionProject"].(string); ok {
projectID = strings.TrimSpace(id)
}
if projectID == "" {
if projectMap, ok := loadResp["cloudaicompanionProject"].(map[string]any); ok {
if id, okID := projectMap["id"].(string); okID {
projectID = strings.TrimSpace(id)
}
}
}
if projectID == "" {
return "", fmt.Errorf("no cloudaicompanionProject in response")
}
return projectID, nil
}

357
sdk/auth/kiro.go Normal file
View File

@@ -0,0 +1,357 @@
package auth
import (
"context"
"fmt"
"strings"
"time"
kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
)
// extractKiroIdentifier extracts a meaningful identifier for file naming.
// Returns account name if provided, otherwise profile ARN ID.
// All extracted values are sanitized to prevent path injection attacks.
func extractKiroIdentifier(accountName, profileArn string) string {
// Priority 1: Use account name if provided
if accountName != "" {
return kiroauth.SanitizeEmailForFilename(accountName)
}
// Priority 2: Use profile ARN ID part (sanitized to prevent path injection)
if profileArn != "" {
parts := strings.Split(profileArn, "/")
if len(parts) >= 2 {
// Sanitize the ARN component to prevent path traversal
return kiroauth.SanitizeEmailForFilename(parts[len(parts)-1])
}
}
// Fallback: timestamp
return fmt.Sprintf("%d", time.Now().UnixNano()%100000)
}
// KiroAuthenticator implements OAuth authentication for Kiro with Google login.
type KiroAuthenticator struct{}
// NewKiroAuthenticator constructs a Kiro authenticator.
func NewKiroAuthenticator() *KiroAuthenticator {
return &KiroAuthenticator{}
}
// Provider returns the provider key for the authenticator.
func (a *KiroAuthenticator) Provider() string {
return "kiro"
}
// RefreshLead indicates how soon before expiry a refresh should be attempted.
func (a *KiroAuthenticator) RefreshLead() *time.Duration {
d := 30 * time.Minute
return &d
}
// Login performs OAuth login for Kiro with AWS Builder ID.
func (a *KiroAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
if cfg == nil {
return nil, fmt.Errorf("kiro auth: configuration is required")
}
oauth := kiroauth.NewKiroOAuth(cfg)
// Use AWS Builder ID device code flow
tokenData, err := oauth.LoginWithBuilderID(ctx)
if err != nil {
return nil, fmt.Errorf("login failed: %w", err)
}
// Parse expires_at
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
if err != nil {
expiresAt = time.Now().Add(1 * time.Hour)
}
// Extract identifier for file naming
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
now := time.Now()
fileName := fmt.Sprintf("kiro-aws-%s.json", idPart)
record := &coreauth.Auth{
ID: fileName,
Provider: "kiro",
FileName: fileName,
Label: "kiro-aws",
Status: coreauth.StatusActive,
CreatedAt: now,
UpdatedAt: now,
Metadata: map[string]any{
"type": "kiro",
"access_token": tokenData.AccessToken,
"refresh_token": tokenData.RefreshToken,
"profile_arn": tokenData.ProfileArn,
"expires_at": tokenData.ExpiresAt,
"auth_method": tokenData.AuthMethod,
"provider": tokenData.Provider,
"client_id": tokenData.ClientID,
"client_secret": tokenData.ClientSecret,
"email": tokenData.Email,
},
Attributes: map[string]string{
"profile_arn": tokenData.ProfileArn,
"source": "aws-builder-id",
"email": tokenData.Email,
},
NextRefreshAfter: expiresAt.Add(-30 * time.Minute),
}
if tokenData.Email != "" {
fmt.Printf("\n✓ Kiro authentication completed successfully! (Account: %s)\n", tokenData.Email)
} else {
fmt.Println("\n✓ Kiro authentication completed successfully!")
}
return record, nil
}
// LoginWithGoogle performs OAuth login for Kiro with Google.
// This uses a custom protocol handler (kiro://) to receive the callback.
func (a *KiroAuthenticator) LoginWithGoogle(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
if cfg == nil {
return nil, fmt.Errorf("kiro auth: configuration is required")
}
oauth := kiroauth.NewKiroOAuth(cfg)
// Use Google OAuth flow with protocol handler
tokenData, err := oauth.LoginWithGoogle(ctx)
if err != nil {
return nil, fmt.Errorf("google login failed: %w", err)
}
// Parse expires_at
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
if err != nil {
expiresAt = time.Now().Add(1 * time.Hour)
}
// Extract identifier for file naming
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
now := time.Now()
fileName := fmt.Sprintf("kiro-google-%s.json", idPart)
record := &coreauth.Auth{
ID: fileName,
Provider: "kiro",
FileName: fileName,
Label: "kiro-google",
Status: coreauth.StatusActive,
CreatedAt: now,
UpdatedAt: now,
Metadata: map[string]any{
"type": "kiro",
"access_token": tokenData.AccessToken,
"refresh_token": tokenData.RefreshToken,
"profile_arn": tokenData.ProfileArn,
"expires_at": tokenData.ExpiresAt,
"auth_method": tokenData.AuthMethod,
"provider": tokenData.Provider,
"email": tokenData.Email,
},
Attributes: map[string]string{
"profile_arn": tokenData.ProfileArn,
"source": "google-oauth",
"email": tokenData.Email,
},
NextRefreshAfter: expiresAt.Add(-30 * time.Minute),
}
if tokenData.Email != "" {
fmt.Printf("\n✓ Kiro Google authentication completed successfully! (Account: %s)\n", tokenData.Email)
} else {
fmt.Println("\n✓ Kiro Google authentication completed successfully!")
}
return record, nil
}
// LoginWithGitHub performs OAuth login for Kiro with GitHub.
// This uses a custom protocol handler (kiro://) to receive the callback.
func (a *KiroAuthenticator) LoginWithGitHub(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
if cfg == nil {
return nil, fmt.Errorf("kiro auth: configuration is required")
}
oauth := kiroauth.NewKiroOAuth(cfg)
// Use GitHub OAuth flow with protocol handler
tokenData, err := oauth.LoginWithGitHub(ctx)
if err != nil {
return nil, fmt.Errorf("github login failed: %w", err)
}
// Parse expires_at
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
if err != nil {
expiresAt = time.Now().Add(1 * time.Hour)
}
// Extract identifier for file naming
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
now := time.Now()
fileName := fmt.Sprintf("kiro-github-%s.json", idPart)
record := &coreauth.Auth{
ID: fileName,
Provider: "kiro",
FileName: fileName,
Label: "kiro-github",
Status: coreauth.StatusActive,
CreatedAt: now,
UpdatedAt: now,
Metadata: map[string]any{
"type": "kiro",
"access_token": tokenData.AccessToken,
"refresh_token": tokenData.RefreshToken,
"profile_arn": tokenData.ProfileArn,
"expires_at": tokenData.ExpiresAt,
"auth_method": tokenData.AuthMethod,
"provider": tokenData.Provider,
"email": tokenData.Email,
},
Attributes: map[string]string{
"profile_arn": tokenData.ProfileArn,
"source": "github-oauth",
"email": tokenData.Email,
},
NextRefreshAfter: expiresAt.Add(-30 * time.Minute),
}
if tokenData.Email != "" {
fmt.Printf("\n✓ Kiro GitHub authentication completed successfully! (Account: %s)\n", tokenData.Email)
} else {
fmt.Println("\n✓ Kiro GitHub authentication completed successfully!")
}
return record, nil
}
// ImportFromKiroIDE imports token from Kiro IDE's token file.
func (a *KiroAuthenticator) ImportFromKiroIDE(ctx context.Context, cfg *config.Config) (*coreauth.Auth, error) {
tokenData, err := kiroauth.LoadKiroIDEToken()
if err != nil {
return nil, fmt.Errorf("failed to load Kiro IDE token: %w", err)
}
// Parse expires_at
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
if err != nil {
expiresAt = time.Now().Add(1 * time.Hour)
}
// Extract email from JWT if not already set (for imported tokens)
if tokenData.Email == "" {
tokenData.Email = kiroauth.ExtractEmailFromJWT(tokenData.AccessToken)
}
// Extract identifier for file naming
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
// Sanitize provider to prevent path traversal (defense-in-depth)
provider := kiroauth.SanitizeEmailForFilename(strings.ToLower(strings.TrimSpace(tokenData.Provider)))
if provider == "" {
provider = "imported" // Fallback for legacy tokens without provider
}
now := time.Now()
fileName := fmt.Sprintf("kiro-%s-%s.json", provider, idPart)
record := &coreauth.Auth{
ID: fileName,
Provider: "kiro",
FileName: fileName,
Label: fmt.Sprintf("kiro-%s", provider),
Status: coreauth.StatusActive,
CreatedAt: now,
UpdatedAt: now,
Metadata: map[string]any{
"type": "kiro",
"access_token": tokenData.AccessToken,
"refresh_token": tokenData.RefreshToken,
"profile_arn": tokenData.ProfileArn,
"expires_at": tokenData.ExpiresAt,
"auth_method": tokenData.AuthMethod,
"provider": tokenData.Provider,
"email": tokenData.Email,
},
Attributes: map[string]string{
"profile_arn": tokenData.ProfileArn,
"source": "kiro-ide-import",
"email": tokenData.Email,
},
NextRefreshAfter: expiresAt.Add(-30 * time.Minute),
}
// Display the email if extracted
if tokenData.Email != "" {
fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s, Account: %s)\n", tokenData.Provider, tokenData.Email)
} else {
fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s)\n", tokenData.Provider)
}
return record, nil
}
// Refresh refreshes an expired Kiro token using AWS SSO OIDC.
func (a *KiroAuthenticator) Refresh(ctx context.Context, cfg *config.Config, auth *coreauth.Auth) (*coreauth.Auth, error) {
if auth == nil || auth.Metadata == nil {
return nil, fmt.Errorf("invalid auth record")
}
refreshToken, ok := auth.Metadata["refresh_token"].(string)
if !ok || refreshToken == "" {
return nil, fmt.Errorf("refresh token not found")
}
clientID, _ := auth.Metadata["client_id"].(string)
clientSecret, _ := auth.Metadata["client_secret"].(string)
authMethod, _ := auth.Metadata["auth_method"].(string)
var tokenData *kiroauth.KiroTokenData
var err error
// Use SSO OIDC refresh for AWS Builder ID, otherwise use Kiro's OAuth refresh endpoint
if clientID != "" && clientSecret != "" && authMethod == "builder-id" {
ssoClient := kiroauth.NewSSOOIDCClient(cfg)
tokenData, err = ssoClient.RefreshToken(ctx, clientID, clientSecret, refreshToken)
} else {
// Fallback to Kiro's refresh endpoint (for social auth: Google/GitHub)
oauth := kiroauth.NewKiroOAuth(cfg)
tokenData, err = oauth.RefreshToken(ctx, refreshToken)
}
if err != nil {
return nil, fmt.Errorf("token refresh failed: %w", err)
}
// Parse expires_at
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
if err != nil {
expiresAt = time.Now().Add(1 * time.Hour)
}
// Clone auth to avoid mutating the input parameter
updated := auth.Clone()
now := time.Now()
updated.UpdatedAt = now
updated.LastRefreshedAt = now
updated.Metadata["access_token"] = tokenData.AccessToken
updated.Metadata["refresh_token"] = tokenData.RefreshToken
updated.Metadata["expires_at"] = tokenData.ExpiresAt
updated.Metadata["last_refresh"] = now.Format(time.RFC3339) // For double-check optimization
updated.NextRefreshAfter = expiresAt.Add(-30 * time.Minute)
return updated, nil
}

View File

@@ -74,3 +74,16 @@ func (m *Manager) Login(ctx context.Context, provider string, cfg *config.Config
}
return record, savedPath, nil
}
// SaveAuth persists an auth record directly without going through the login flow.
func (m *Manager) SaveAuth(record *coreauth.Auth, cfg *config.Config) (string, error) {
if m.store == nil {
return "", fmt.Errorf("no store configured")
}
if cfg != nil {
if dirSetter, ok := m.store.(interface{ SetBaseDir(string) }); ok {
dirSetter.SetBaseDir(cfg.AuthDir)
}
}
return m.store.Save(context.Background(), record)
}

View File

@@ -14,6 +14,7 @@ func init() {
registerRefreshLead("gemini", func() Authenticator { return NewGeminiAuthenticator() })
registerRefreshLead("gemini-cli", func() Authenticator { return NewGeminiAuthenticator() })
registerRefreshLead("antigravity", func() Authenticator { return NewAntigravityAuthenticator() })
registerRefreshLead("kiro", func() Authenticator { return NewKiroAuthenticator() })
registerRefreshLead("github-copilot", func() Authenticator { return NewGitHubCopilotAuthenticator() })
}

View File

@@ -29,7 +29,7 @@ func NewAPIKeyClientProvider() APIKeyClientProvider {
type apiKeyClientProvider struct{}
func (p *apiKeyClientProvider) Load(ctx context.Context, cfg *config.Config) (*APIKeyClientResult, error) {
geminiCount, claudeCount, codexCount, openAICompat := watcher.BuildAPIKeyClients(cfg)
geminiCount, vertexCompatCount, claudeCount, codexCount, openAICompat := watcher.BuildAPIKeyClients(cfg)
if ctx != nil {
select {
case <-ctx.Done():
@@ -38,9 +38,10 @@ func (p *apiKeyClientProvider) Load(ctx context.Context, cfg *config.Config) (*A
}
}
return &APIKeyClientResult{
GeminiKeyCount: geminiCount,
ClaudeKeyCount: claudeCount,
CodexKeyCount: codexCount,
OpenAICompatCount: openAICompat,
GeminiKeyCount: geminiCount,
VertexCompatKeyCount: vertexCompatCount,
ClaudeKeyCount: claudeCount,
CodexKeyCount: codexCount,
OpenAICompatCount: openAICompat,
}, nil
}

View File

@@ -324,7 +324,7 @@ func openAICompatInfoFromAuth(a *coreauth.Auth) (providerKey string, compatName
if len(a.Attributes) > 0 {
providerKey = strings.TrimSpace(a.Attributes["provider_key"])
compatName = strings.TrimSpace(a.Attributes["compat_name"])
if providerKey != "" || compatName != "" {
if compatName != "" {
if providerKey == "" {
providerKey = compatName
}
@@ -379,6 +379,8 @@ func (s *Service) ensureExecutorsForAuth(a *coreauth.Auth) {
s.coreManager.RegisterExecutor(executor.NewQwenExecutor(s.cfg))
case "iflow":
s.coreManager.RegisterExecutor(executor.NewIFlowExecutor(s.cfg))
case "kiro":
s.coreManager.RegisterExecutor(executor.NewKiroExecutor(s.cfg))
case "github-copilot":
s.coreManager.RegisterExecutor(executor.NewGitHubCopilotExecutor(s.cfg))
default:
@@ -500,7 +502,7 @@ func (s *Service) Run(ctx context.Context) error {
}()
time.Sleep(100 * time.Millisecond)
fmt.Println("API server started successfully")
fmt.Printf("API server started successfully on: %s:%d\n", s.cfg.Host, s.cfg.Port)
if s.hooks.OnAfterStart != nil {
s.hooks.OnAfterStart(s)
@@ -681,6 +683,11 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) {
case "vertex":
// Vertex AI Gemini supports the same model identifiers as Gemini.
models = registry.GetGeminiVertexModels()
if authKind == "apikey" {
if entry := s.resolveConfigVertexCompatKey(a); entry != nil && len(entry.Models) > 0 {
models = buildVertexCompatConfigModels(entry)
}
}
models = applyExcludedModels(models, excluded)
case "gemini-cli":
models = registry.GetGeminiCLIModels()
@@ -720,6 +727,9 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) {
case "github-copilot":
models = registry.GetGitHubCopilotModels()
models = applyExcludedModels(models, excluded)
case "kiro":
models = registry.GetKiroModels()
models = applyExcludedModels(models, excluded)
default:
// Handle OpenAI-compatibility providers by name using config
if s.cfg != nil {
@@ -878,6 +888,40 @@ func (s *Service) resolveConfigGeminiKey(auth *coreauth.Auth) *config.GeminiKey
return nil
}
func (s *Service) resolveConfigVertexCompatKey(auth *coreauth.Auth) *config.VertexCompatKey {
if auth == nil || s.cfg == nil {
return nil
}
var attrKey, attrBase string
if auth.Attributes != nil {
attrKey = strings.TrimSpace(auth.Attributes["api_key"])
attrBase = strings.TrimSpace(auth.Attributes["base_url"])
}
for i := range s.cfg.VertexCompatAPIKey {
entry := &s.cfg.VertexCompatAPIKey[i]
cfgKey := strings.TrimSpace(entry.APIKey)
cfgBase := strings.TrimSpace(entry.BaseURL)
if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
return entry
}
continue
}
if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
return entry
}
}
if attrKey != "" {
for i := range s.cfg.VertexCompatAPIKey {
entry := &s.cfg.VertexCompatAPIKey[i]
if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) {
return entry
}
}
}
return nil
}
func (s *Service) resolveConfigCodexKey(auth *coreauth.Auth) *config.CodexKey {
if auth == nil || s.cfg == nil {
return nil
@@ -996,6 +1040,44 @@ func matchWildcard(pattern, value string) bool {
return true
}
func buildVertexCompatConfigModels(entry *config.VertexCompatKey) []*ModelInfo {
if entry == nil || len(entry.Models) == 0 {
return nil
}
now := time.Now().Unix()
out := make([]*ModelInfo, 0, len(entry.Models))
seen := make(map[string]struct{}, len(entry.Models))
for i := range entry.Models {
model := entry.Models[i]
name := strings.TrimSpace(model.Name)
alias := strings.TrimSpace(model.Alias)
if alias == "" {
alias = name
}
if alias == "" {
continue
}
key := strings.ToLower(alias)
if _, exists := seen[key]; exists {
continue
}
seen[key] = struct{}{}
display := name
if display == "" {
display = alias
}
out = append(out, &ModelInfo{
ID: alias,
Object: "model",
Created: now,
OwnedBy: "vertex",
Type: "vertex",
DisplayName: display,
})
}
return out
}
func buildClaudeConfigModels(entry *config.ClaudeKey) []*ModelInfo {
if entry == nil || len(entry.Models) == 0 {
return nil

View File

@@ -49,19 +49,21 @@ type APIKeyClientProvider interface {
Load(ctx context.Context, cfg *config.Config) (*APIKeyClientResult, error)
}
// APIKeyClientResult contains API key based clients along with type counts.
// It provides metadata about the number of clients loaded for each provider type.
// APIKeyClientResult is returned by APIKeyClientProvider.Load()
type APIKeyClientResult struct {
// GeminiKeyCount is the number of Gemini API key clients loaded.
// GeminiKeyCount is the number of Gemini API keys loaded
GeminiKeyCount int
// ClaudeKeyCount is the number of Claude API key clients loaded.
// VertexCompatKeyCount is the number of Vertex-compatible API keys loaded
VertexCompatKeyCount int
// ClaudeKeyCount is the number of Claude API keys loaded
ClaudeKeyCount int
// CodexKeyCount is the number of Codex API key clients loaded.
// CodexKeyCount is the number of Codex API keys loaded
CodexKeyCount int
// OpenAICompatCount is the number of OpenAI-compatible API key clients loaded.
// OpenAICompatCount is the number of OpenAI compatibility API keys loaded
OpenAICompatCount int
}

827
test/amp_management_test.go Normal file
View File

@@ -0,0 +1,827 @@
package test
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"testing"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/handlers/management"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
)
func init() {
gin.SetMode(gin.TestMode)
}
// newAmpTestHandler creates a test handler with default ampcode configuration.
func newAmpTestHandler(t *testing.T) (*management.Handler, string) {
t.Helper()
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.yaml")
cfg := &config.Config{
AmpCode: config.AmpCode{
UpstreamURL: "https://example.com",
UpstreamAPIKey: "test-api-key-12345",
RestrictManagementToLocalhost: true,
ForceModelMappings: false,
ModelMappings: []config.AmpModelMapping{
{From: "gpt-4", To: "gemini-pro"},
},
},
}
if err := os.WriteFile(configPath, []byte("port: 8080\n"), 0644); err != nil {
t.Fatalf("failed to write config file: %v", err)
}
h := management.NewHandler(cfg, configPath, nil)
return h, configPath
}
// setupAmpRouter creates a test router with all ampcode management endpoints.
func setupAmpRouter(h *management.Handler) *gin.Engine {
r := gin.New()
mgmt := r.Group("/v0/management")
{
mgmt.GET("/ampcode", h.GetAmpCode)
mgmt.GET("/ampcode/upstream-url", h.GetAmpUpstreamURL)
mgmt.PUT("/ampcode/upstream-url", h.PutAmpUpstreamURL)
mgmt.DELETE("/ampcode/upstream-url", h.DeleteAmpUpstreamURL)
mgmt.GET("/ampcode/upstream-api-key", h.GetAmpUpstreamAPIKey)
mgmt.PUT("/ampcode/upstream-api-key", h.PutAmpUpstreamAPIKey)
mgmt.DELETE("/ampcode/upstream-api-key", h.DeleteAmpUpstreamAPIKey)
mgmt.GET("/ampcode/restrict-management-to-localhost", h.GetAmpRestrictManagementToLocalhost)
mgmt.PUT("/ampcode/restrict-management-to-localhost", h.PutAmpRestrictManagementToLocalhost)
mgmt.GET("/ampcode/model-mappings", h.GetAmpModelMappings)
mgmt.PUT("/ampcode/model-mappings", h.PutAmpModelMappings)
mgmt.PATCH("/ampcode/model-mappings", h.PatchAmpModelMappings)
mgmt.DELETE("/ampcode/model-mappings", h.DeleteAmpModelMappings)
mgmt.GET("/ampcode/force-model-mappings", h.GetAmpForceModelMappings)
mgmt.PUT("/ampcode/force-model-mappings", h.PutAmpForceModelMappings)
}
return r
}
// TestGetAmpCode verifies GET /v0/management/ampcode returns full ampcode config.
func TestGetAmpCode(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
var resp map[string]config.AmpCode
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal response: %v", err)
}
ampcode := resp["ampcode"]
if ampcode.UpstreamURL != "https://example.com" {
t.Errorf("expected upstream-url %q, got %q", "https://example.com", ampcode.UpstreamURL)
}
if len(ampcode.ModelMappings) != 1 {
t.Errorf("expected 1 model mapping, got %d", len(ampcode.ModelMappings))
}
}
// TestGetAmpUpstreamURL verifies GET /v0/management/ampcode/upstream-url returns the upstream URL.
func TestGetAmpUpstreamURL(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-url", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
var resp map[string]string
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal response: %v", err)
}
if resp["upstream-url"] != "https://example.com" {
t.Errorf("expected %q, got %q", "https://example.com", resp["upstream-url"])
}
}
// TestPutAmpUpstreamURL verifies PUT /v0/management/ampcode/upstream-url updates the upstream URL.
func TestPutAmpUpstreamURL(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": "https://new-upstream.com"}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/upstream-url", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d: %s", http.StatusOK, w.Code, w.Body.String())
}
}
// TestDeleteAmpUpstreamURL verifies DELETE /v0/management/ampcode/upstream-url clears the upstream URL.
func TestDeleteAmpUpstreamURL(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/upstream-url", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
}
// TestGetAmpUpstreamAPIKey verifies GET /v0/management/ampcode/upstream-api-key returns the API key.
func TestGetAmpUpstreamAPIKey(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-api-key", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
var resp map[string]any
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal response: %v", err)
}
key := resp["upstream-api-key"].(string)
if key != "test-api-key-12345" {
t.Errorf("expected key %q, got %q", "test-api-key-12345", key)
}
}
// TestPutAmpUpstreamAPIKey verifies PUT /v0/management/ampcode/upstream-api-key updates the API key.
func TestPutAmpUpstreamAPIKey(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": "new-secret-key"}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/upstream-api-key", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
}
// TestDeleteAmpUpstreamAPIKey verifies DELETE /v0/management/ampcode/upstream-api-key clears the API key.
func TestDeleteAmpUpstreamAPIKey(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/upstream-api-key", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
}
// TestGetAmpRestrictManagementToLocalhost verifies GET returns the localhost restriction setting.
func TestGetAmpRestrictManagementToLocalhost(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/restrict-management-to-localhost", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
var resp map[string]bool
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal response: %v", err)
}
if resp["restrict-management-to-localhost"] != true {
t.Error("expected restrict-management-to-localhost to be true")
}
}
// TestPutAmpRestrictManagementToLocalhost verifies PUT updates the localhost restriction setting.
func TestPutAmpRestrictManagementToLocalhost(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": false}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/restrict-management-to-localhost", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
}
// TestGetAmpModelMappings verifies GET /v0/management/ampcode/model-mappings returns all mappings.
func TestGetAmpModelMappings(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
var resp map[string][]config.AmpModelMapping
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal response: %v", err)
}
mappings := resp["model-mappings"]
if len(mappings) != 1 {
t.Fatalf("expected 1 mapping, got %d", len(mappings))
}
if mappings[0].From != "gpt-4" || mappings[0].To != "gemini-pro" {
t.Errorf("unexpected mapping: %+v", mappings[0])
}
}
// TestPutAmpModelMappings verifies PUT /v0/management/ampcode/model-mappings replaces all mappings.
func TestPutAmpModelMappings(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": [{"from": "claude-3", "to": "gpt-4o"}, {"from": "gemini", "to": "claude"}]}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d: %s", http.StatusOK, w.Code, w.Body.String())
}
}
// TestPatchAmpModelMappings verifies PATCH updates existing mappings and adds new ones.
func TestPatchAmpModelMappings(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": [{"from": "gpt-4", "to": "updated-model"}, {"from": "new-model", "to": "target"}]}`
req := httptest.NewRequest(http.MethodPatch, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d: %s", http.StatusOK, w.Code, w.Body.String())
}
}
// TestDeleteAmpModelMappings_Specific verifies DELETE removes specified mappings by "from" field.
func TestDeleteAmpModelMappings_Specific(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": ["gpt-4"]}`
req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
}
// TestDeleteAmpModelMappings_All verifies DELETE with empty body removes all mappings.
func TestDeleteAmpModelMappings_All(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/model-mappings", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
}
// TestGetAmpForceModelMappings verifies GET returns the force-model-mappings setting.
func TestGetAmpForceModelMappings(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/force-model-mappings", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
var resp map[string]bool
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal response: %v", err)
}
if resp["force-model-mappings"] != false {
t.Error("expected force-model-mappings to be false")
}
}
// TestPutAmpForceModelMappings verifies PUT updates the force-model-mappings setting.
func TestPutAmpForceModelMappings(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": true}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/force-model-mappings", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
}
// TestPutAmpModelMappings_VerifyState verifies PUT replaces mappings and state is persisted.
func TestPutAmpModelMappings_VerifyState(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": [{"from": "model-a", "to": "model-b"}, {"from": "model-c", "to": "model-d"}, {"from": "model-e", "to": "model-f"}]}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("PUT failed: status %d, body: %s", w.Code, w.Body.String())
}
req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
var resp map[string][]config.AmpModelMapping
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
mappings := resp["model-mappings"]
if len(mappings) != 3 {
t.Fatalf("expected 3 mappings, got %d", len(mappings))
}
expected := map[string]string{"model-a": "model-b", "model-c": "model-d", "model-e": "model-f"}
for _, m := range mappings {
if expected[m.From] != m.To {
t.Errorf("mapping %q -> expected %q, got %q", m.From, expected[m.From], m.To)
}
}
}
// TestPatchAmpModelMappings_VerifyState verifies PATCH merges mappings correctly.
func TestPatchAmpModelMappings_VerifyState(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": [{"from": "gpt-4", "to": "updated-target"}, {"from": "new-model", "to": "new-target"}]}`
req := httptest.NewRequest(http.MethodPatch, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("PATCH failed: status %d", w.Code)
}
req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
var resp map[string][]config.AmpModelMapping
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
mappings := resp["model-mappings"]
if len(mappings) != 2 {
t.Fatalf("expected 2 mappings (1 updated + 1 new), got %d", len(mappings))
}
found := make(map[string]string)
for _, m := range mappings {
found[m.From] = m.To
}
if found["gpt-4"] != "updated-target" {
t.Errorf("gpt-4 should map to updated-target, got %q", found["gpt-4"])
}
if found["new-model"] != "new-target" {
t.Errorf("new-model should map to new-target, got %q", found["new-model"])
}
}
// TestDeleteAmpModelMappings_VerifyState verifies DELETE removes specific mappings and keeps others.
func TestDeleteAmpModelMappings_VerifyState(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
putBody := `{"value": [{"from": "a", "to": "1"}, {"from": "b", "to": "2"}, {"from": "c", "to": "3"}]}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(putBody))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
delBody := `{"value": ["a", "c"]}`
req = httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(delBody))
req.Header.Set("Content-Type", "application/json")
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("DELETE failed: status %d", w.Code)
}
req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
var resp map[string][]config.AmpModelMapping
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
mappings := resp["model-mappings"]
if len(mappings) != 1 {
t.Fatalf("expected 1 mapping remaining, got %d", len(mappings))
}
if mappings[0].From != "b" || mappings[0].To != "2" {
t.Errorf("expected b->2, got %s->%s", mappings[0].From, mappings[0].To)
}
}
// TestDeleteAmpModelMappings_NonExistent verifies DELETE with non-existent mapping doesn't affect existing ones.
func TestDeleteAmpModelMappings_NonExistent(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
delBody := `{"value": ["non-existent-model"]}`
req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(delBody))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
var resp map[string][]config.AmpModelMapping
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
if len(resp["model-mappings"]) != 1 {
t.Errorf("original mapping should remain, got %d mappings", len(resp["model-mappings"]))
}
}
// TestPutAmpModelMappings_Empty verifies PUT with empty array clears all mappings.
func TestPutAmpModelMappings_Empty(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": []}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
var resp map[string][]config.AmpModelMapping
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
if len(resp["model-mappings"]) != 0 {
t.Errorf("expected 0 mappings, got %d", len(resp["model-mappings"]))
}
}
// TestPutAmpUpstreamURL_VerifyState verifies PUT updates upstream URL and persists state.
func TestPutAmpUpstreamURL_VerifyState(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": "https://new-api.example.com"}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/upstream-url", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("PUT failed: status %d", w.Code)
}
req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-url", nil)
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
var resp map[string]string
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
if resp["upstream-url"] != "https://new-api.example.com" {
t.Errorf("expected %q, got %q", "https://new-api.example.com", resp["upstream-url"])
}
}
// TestDeleteAmpUpstreamURL_VerifyState verifies DELETE clears upstream URL.
func TestDeleteAmpUpstreamURL_VerifyState(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/upstream-url", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("DELETE failed: status %d", w.Code)
}
req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-url", nil)
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
var resp map[string]string
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
if resp["upstream-url"] != "" {
t.Errorf("expected empty string, got %q", resp["upstream-url"])
}
}
// TestPutAmpUpstreamAPIKey_VerifyState verifies PUT updates API key and persists state.
func TestPutAmpUpstreamAPIKey_VerifyState(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": "new-secret-api-key-xyz"}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/upstream-api-key", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("PUT failed: status %d", w.Code)
}
req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-api-key", nil)
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
var resp map[string]string
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
if resp["upstream-api-key"] != "new-secret-api-key-xyz" {
t.Errorf("expected %q, got %q", "new-secret-api-key-xyz", resp["upstream-api-key"])
}
}
// TestDeleteAmpUpstreamAPIKey_VerifyState verifies DELETE clears API key.
func TestDeleteAmpUpstreamAPIKey_VerifyState(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/upstream-api-key", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("DELETE failed: status %d", w.Code)
}
req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/upstream-api-key", nil)
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
var resp map[string]string
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
if resp["upstream-api-key"] != "" {
t.Errorf("expected empty string, got %q", resp["upstream-api-key"])
}
}
// TestPutAmpRestrictManagementToLocalhost_VerifyState verifies PUT updates localhost restriction.
func TestPutAmpRestrictManagementToLocalhost_VerifyState(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": false}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/restrict-management-to-localhost", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("PUT failed: status %d", w.Code)
}
req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/restrict-management-to-localhost", nil)
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
var resp map[string]bool
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
if resp["restrict-management-to-localhost"] != false {
t.Error("expected false after update")
}
}
// TestPutAmpForceModelMappings_VerifyState verifies PUT updates force-model-mappings setting.
func TestPutAmpForceModelMappings_VerifyState(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{"value": true}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/force-model-mappings", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("PUT failed: status %d", w.Code)
}
req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/force-model-mappings", nil)
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
var resp map[string]bool
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
if resp["force-model-mappings"] != true {
t.Error("expected true after update")
}
}
// TestPutBoolField_EmptyObject verifies PUT with empty object returns 400.
func TestPutBoolField_EmptyObject(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
body := `{}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/force-model-mappings", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Fatalf("expected status %d for empty object, got %d", http.StatusBadRequest, w.Code)
}
}
// TestComplexMappingsWorkflow tests a full workflow: PUT, PATCH, DELETE, and GET.
func TestComplexMappingsWorkflow(t *testing.T) {
h, _ := newAmpTestHandler(t)
r := setupAmpRouter(h)
putBody := `{"value": [{"from": "m1", "to": "t1"}, {"from": "m2", "to": "t2"}, {"from": "m3", "to": "t3"}, {"from": "m4", "to": "t4"}]}`
req := httptest.NewRequest(http.MethodPut, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(putBody))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
patchBody := `{"value": [{"from": "m2", "to": "t2-updated"}, {"from": "m5", "to": "t5"}]}`
req = httptest.NewRequest(http.MethodPatch, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(patchBody))
req.Header.Set("Content-Type", "application/json")
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
delBody := `{"value": ["m1", "m3"]}`
req = httptest.NewRequest(http.MethodDelete, "/v0/management/ampcode/model-mappings", bytes.NewBufferString(delBody))
req.Header.Set("Content-Type", "application/json")
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
req = httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
w = httptest.NewRecorder()
r.ServeHTTP(w, req)
var resp map[string][]config.AmpModelMapping
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
mappings := resp["model-mappings"]
if len(mappings) != 3 {
t.Fatalf("expected 3 mappings (m2, m4, m5), got %d", len(mappings))
}
expected := map[string]string{"m2": "t2-updated", "m4": "t4", "m5": "t5"}
found := make(map[string]string)
for _, m := range mappings {
found[m.From] = m.To
}
for from, to := range expected {
if found[from] != to {
t.Errorf("mapping %s: expected %q, got %q", from, to, found[from])
}
}
}
// TestNilHandlerGetAmpCode verifies handler works with empty config.
func TestNilHandlerGetAmpCode(t *testing.T) {
cfg := &config.Config{}
h := management.NewHandler(cfg, "", nil)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
}
// TestEmptyConfigGetAmpModelMappings verifies GET returns empty array for fresh config.
func TestEmptyConfigGetAmpModelMappings(t *testing.T) {
cfg := &config.Config{}
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.yaml")
if err := os.WriteFile(configPath, []byte("port: 8080\n"), 0644); err != nil {
t.Fatalf("failed to write config: %v", err)
}
h := management.NewHandler(cfg, configPath, nil)
r := setupAmpRouter(h)
req := httptest.NewRequest(http.MethodGet, "/v0/management/ampcode/model-mappings", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected status %d, got %d", http.StatusOK, w.Code)
}
var resp map[string][]config.AmpModelMapping
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to unmarshal: %v", err)
}
if len(resp["model-mappings"]) != 0 {
t.Errorf("expected 0 mappings, got %d", len(resp["model-mappings"]))
}
}

View File

@@ -0,0 +1,195 @@
package test
import (
"os"
"path/filepath"
"strings"
"testing"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
)
func TestLegacyConfigMigration(t *testing.T) {
t.Run("onlyLegacyFields", func(t *testing.T) {
path := writeConfig(t, `
port: 8080
generative-language-api-key:
- "legacy-gemini-1"
openai-compatibility:
- name: "legacy-provider"
base-url: "https://example.com"
api-keys:
- "legacy-openai-1"
amp-upstream-url: "https://amp.example.com"
amp-upstream-api-key: "amp-legacy-key"
amp-restrict-management-to-localhost: false
amp-model-mappings:
- from: "old-model"
to: "new-model"
`)
cfg, err := config.LoadConfig(path)
if err != nil {
t.Fatalf("load legacy config: %v", err)
}
if got := len(cfg.GeminiKey); got != 1 || cfg.GeminiKey[0].APIKey != "legacy-gemini-1" {
t.Fatalf("gemini migration mismatch: %+v", cfg.GeminiKey)
}
if got := len(cfg.OpenAICompatibility); got != 1 {
t.Fatalf("expected 1 openai-compat provider, got %d", got)
}
if entries := cfg.OpenAICompatibility[0].APIKeyEntries; len(entries) != 1 || entries[0].APIKey != "legacy-openai-1" {
t.Fatalf("openai-compat migration mismatch: %+v", entries)
}
if cfg.AmpCode.UpstreamURL != "https://amp.example.com" || cfg.AmpCode.UpstreamAPIKey != "amp-legacy-key" {
t.Fatalf("amp migration failed: %+v", cfg.AmpCode)
}
if cfg.AmpCode.RestrictManagementToLocalhost {
t.Fatalf("expected amp restriction to be false after migration")
}
if got := len(cfg.AmpCode.ModelMappings); got != 1 || cfg.AmpCode.ModelMappings[0].From != "old-model" {
t.Fatalf("amp mappings migration mismatch: %+v", cfg.AmpCode.ModelMappings)
}
updated := readFile(t, path)
if strings.Contains(updated, "generative-language-api-key") {
t.Fatalf("legacy gemini key still present:\n%s", updated)
}
if strings.Contains(updated, "amp-upstream-url") || strings.Contains(updated, "amp-restrict-management-to-localhost") {
t.Fatalf("legacy amp keys still present:\n%s", updated)
}
if strings.Contains(updated, "\n api-keys:") {
t.Fatalf("legacy openai compat keys still present:\n%s", updated)
}
})
t.Run("mixedLegacyAndNewFields", func(t *testing.T) {
path := writeConfig(t, `
gemini-api-key:
- api-key: "new-gemini"
generative-language-api-key:
- "new-gemini"
- "legacy-gemini-only"
openai-compatibility:
- name: "mixed-provider"
base-url: "https://mixed.example.com"
api-key-entries:
- api-key: "new-entry"
api-keys:
- "legacy-entry"
- "new-entry"
`)
cfg, err := config.LoadConfig(path)
if err != nil {
t.Fatalf("load mixed config: %v", err)
}
if got := len(cfg.GeminiKey); got != 2 {
t.Fatalf("expected 2 gemini entries, got %d: %+v", got, cfg.GeminiKey)
}
seen := make(map[string]struct{}, len(cfg.GeminiKey))
for _, entry := range cfg.GeminiKey {
if _, exists := seen[entry.APIKey]; exists {
t.Fatalf("duplicate gemini key %q after migration", entry.APIKey)
}
seen[entry.APIKey] = struct{}{}
}
provider := cfg.OpenAICompatibility[0]
if got := len(provider.APIKeyEntries); got != 2 {
t.Fatalf("expected 2 openai entries, got %d: %+v", got, provider.APIKeyEntries)
}
entrySeen := make(map[string]struct{}, len(provider.APIKeyEntries))
for _, entry := range provider.APIKeyEntries {
if _, ok := entrySeen[entry.APIKey]; ok {
t.Fatalf("duplicate openai key %q after migration", entry.APIKey)
}
entrySeen[entry.APIKey] = struct{}{}
}
})
t.Run("onlyNewFields", func(t *testing.T) {
path := writeConfig(t, `
gemini-api-key:
- api-key: "new-only"
openai-compatibility:
- name: "new-only-provider"
base-url: "https://new-only.example.com"
api-key-entries:
- api-key: "new-only-entry"
ampcode:
upstream-url: "https://amp.new"
upstream-api-key: "new-amp-key"
restrict-management-to-localhost: true
model-mappings:
- from: "a"
to: "b"
`)
cfg, err := config.LoadConfig(path)
if err != nil {
t.Fatalf("load new config: %v", err)
}
if len(cfg.GeminiKey) != 1 || cfg.GeminiKey[0].APIKey != "new-only" {
t.Fatalf("unexpected gemini entries: %+v", cfg.GeminiKey)
}
if len(cfg.OpenAICompatibility) != 1 || len(cfg.OpenAICompatibility[0].APIKeyEntries) != 1 {
t.Fatalf("unexpected openai compat entries: %+v", cfg.OpenAICompatibility)
}
if cfg.AmpCode.UpstreamURL != "https://amp.new" || cfg.AmpCode.UpstreamAPIKey != "new-amp-key" {
t.Fatalf("unexpected amp config: %+v", cfg.AmpCode)
}
})
t.Run("duplicateNamesDifferentBase", func(t *testing.T) {
path := writeConfig(t, `
openai-compatibility:
- name: "dup-provider"
base-url: "https://provider-a"
api-keys:
- "key-a"
- name: "dup-provider"
base-url: "https://provider-b"
api-keys:
- "key-b"
`)
cfg, err := config.LoadConfig(path)
if err != nil {
t.Fatalf("load duplicate config: %v", err)
}
if len(cfg.OpenAICompatibility) != 2 {
t.Fatalf("expected 2 providers, got %d", len(cfg.OpenAICompatibility))
}
for _, entry := range cfg.OpenAICompatibility {
if len(entry.APIKeyEntries) != 1 {
t.Fatalf("expected 1 key entry per provider: %+v", entry)
}
switch entry.BaseURL {
case "https://provider-a":
if entry.APIKeyEntries[0].APIKey != "key-a" {
t.Fatalf("provider-a key mismatch: %+v", entry.APIKeyEntries)
}
case "https://provider-b":
if entry.APIKeyEntries[0].APIKey != "key-b" {
t.Fatalf("provider-b key mismatch: %+v", entry.APIKeyEntries)
}
default:
t.Fatalf("unexpected provider base url: %s", entry.BaseURL)
}
}
})
}
func writeConfig(t *testing.T, content string) string {
t.Helper()
dir := t.TempDir()
path := filepath.Join(dir, "config.yaml")
if err := os.WriteFile(path, []byte(strings.TrimSpace(content)+"\n"), 0o644); err != nil {
t.Fatalf("write temp config: %v", err)
}
return path
}
func readFile(t *testing.T, path string) string {
t.Helper()
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("read temp config: %v", err)
}
return string(data)
}