mirror of
https://github.com/router-for-me/CLIProxyAPIPlus.git
synced 2026-03-30 09:18:12 +00:00
Compare commits
426 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1fa094dac6 | ||
|
|
f55754621f | ||
|
|
ac26e7db43 | ||
|
|
10b824fcac | ||
|
|
7dccc7ba2f | ||
|
|
70c90687fd | ||
|
|
8144ffd5c8 | ||
|
|
6b45d311ec | ||
|
|
7386a70724 | ||
|
|
1821bf7051 | ||
|
|
d42b5d4e78 | ||
|
|
1b7447b682 | ||
|
|
40dee4453a | ||
|
|
8902e1cccb | ||
|
|
de5fe71478 | ||
|
|
dcfbec2990 | ||
|
|
c95620f90e | ||
|
|
9613f0b3f9 | ||
|
|
274f29e26b | ||
|
|
c8e79c3787 | ||
|
|
8afef43887 | ||
|
|
c1083cbfc6 | ||
|
|
1e6bc81cfd | ||
|
|
1a149475e0 | ||
|
|
e5166841db | ||
|
|
19c52bcb60 | ||
|
|
bb9b2d1758 | ||
|
|
7fa527193c | ||
|
|
ed0eb51b4d | ||
|
|
0e4f669c8b | ||
|
|
76c064c729 | ||
|
|
d2f652f436 | ||
|
|
6a452a54d5 | ||
|
|
9e5693e74f | ||
|
|
528b1a2307 | ||
|
|
0cc978ec1d | ||
|
|
d312422ab4 | ||
|
|
fee736933b | ||
|
|
09c92aa0b5 | ||
|
|
8c67b3ae64 | ||
|
|
000e4ceb4e | ||
|
|
5c99846ecf | ||
|
|
d475aaba96 | ||
|
|
1dc4ecb1b8 | ||
|
|
1315f710f5 | ||
|
|
96f55570f7 | ||
|
|
0906aeca87 | ||
|
|
7333619f15 | ||
|
|
97c0487add | ||
|
|
2db8df8e38 | ||
|
|
a576088d5f | ||
|
|
66ff916838 | ||
|
|
7b0453074e | ||
|
|
a000eb523d | ||
|
|
18a4fedc7f | ||
|
|
5d6cdccda0 | ||
|
|
1b7f4ac3e1 | ||
|
|
afc1a5b814 | ||
|
|
7ed38db54f | ||
|
|
28c10f4e69 | ||
|
|
6e12441a3b | ||
|
|
65c439c18d | ||
|
|
0ed2d16596 | ||
|
|
db335ac616 | ||
|
|
f3c59165d7 | ||
|
|
e6690cb447 | ||
|
|
35907416b8 | ||
|
|
e8bb350467 | ||
|
|
5331d51f27 | ||
|
|
755ca75879 | ||
|
|
2398ebad55 | ||
|
|
c1bf298216 | ||
|
|
e005208d76 | ||
|
|
d1df70d02f | ||
|
|
f81acd0760 | ||
|
|
636da4c932 | ||
|
|
cccb77b552 | ||
|
|
2bd646ad70 | ||
|
|
52c1fa025e | ||
|
|
680105f84d | ||
|
|
f7069e9548 | ||
|
|
7275e99b41 | ||
|
|
c28b65f849 | ||
|
|
793840cdb4 | ||
|
|
8f421de532 | ||
|
|
be2dd60ee7 | ||
|
|
ea3e0b713e | ||
|
|
8179d5a8a4 | ||
|
|
6fa7abe434 | ||
|
|
5135c22cd6 | ||
|
|
1e27990561 | ||
|
|
e1e9fc43c1 | ||
|
|
b2921518ac | ||
|
|
dd64adbeeb | ||
|
|
616d41c06a | ||
|
|
e0e337aeb9 | ||
|
|
d52839fced | ||
|
|
4022e69651 | ||
|
|
56073ded69 | ||
|
|
9738a53f49 | ||
|
|
be3f8dbf7e | ||
|
|
9c6c3612a8 | ||
|
|
19e1a4447a | ||
|
|
7c2ad4cda2 | ||
|
|
fb95813fbf | ||
|
|
db63f9b5d6 | ||
|
|
25f6c4a250 | ||
|
|
b24ae74216 | ||
|
|
59ad8f40dc | ||
|
|
ff03dc6a2c | ||
|
|
dc7187ca5b | ||
|
|
b1dcff778c | ||
|
|
cef2aeeb08 | ||
|
|
bcd1e8cc34 | ||
|
|
198b3f4a40 | ||
|
|
9fee7f488e | ||
|
|
1b46d39b8b | ||
|
|
c1241a98e2 | ||
|
|
8d8f5970ee | ||
|
|
f90120f846 | ||
|
|
0b94d36c4a | ||
|
|
152c310bb7 | ||
|
|
f6bbca35ab | ||
|
|
c8cee6a209 | ||
|
|
b5701f416b | ||
|
|
4b1a404fcb | ||
|
|
b93cce5412 | ||
|
|
c6cb24039d | ||
|
|
5382408489 | ||
|
|
67669196ed | ||
|
|
58fd9bf964 | ||
|
|
7b3dfc67bc | ||
|
|
cdd24052d3 | ||
|
|
733fd8edab | ||
|
|
af27f2b8bc | ||
|
|
2e1925d762 | ||
|
|
77254bd074 | ||
|
|
5b6342e6ac | ||
|
|
3960c93d51 | ||
|
|
339a81b650 | ||
|
|
560c020477 | ||
|
|
aec65e3be3 | ||
|
|
f44f0702f8 | ||
|
|
b76b79068f | ||
|
|
34c8ccb961 | ||
|
|
d08e164af3 | ||
|
|
8178efaeda | ||
|
|
86d5db472a | ||
|
|
020d36f6e8 | ||
|
|
1db23979e8 | ||
|
|
c3d5dbe96f | ||
|
|
5484489406 | ||
|
|
0ac52da460 | ||
|
|
817cebb321 | ||
|
|
683f3709d6 | ||
|
|
dbd42a42b2 | ||
|
|
ec24baf757 | ||
|
|
dea3e74d35 | ||
|
|
a6c3042e34 | ||
|
|
861537c9bd | ||
|
|
8c92cb0883 | ||
|
|
89d7be9525 | ||
|
|
2b79d7f22f | ||
|
|
2bb686f594 | ||
|
|
163fe287ce | ||
|
|
70988d387b | ||
|
|
52058a1659 | ||
|
|
df5595a0c9 | ||
|
|
ddaa9d2436 | ||
|
|
7b7b258c38 | ||
|
|
a00f774f5a | ||
|
|
9daf1ba8b5 | ||
|
|
76f2359637 | ||
|
|
dcb1c9be8a | ||
|
|
a24f4ace78 | ||
|
|
c631df8c3b | ||
|
|
54c3eb1b1e | ||
|
|
bb28cd26ad | ||
|
|
046865461e | ||
|
|
cf74ed2f0c | ||
|
|
c3762328a5 | ||
|
|
e333fbea3d | ||
|
|
efbe36d1d4 | ||
|
|
8553cfa40e | ||
|
|
30d5c95b26 | ||
|
|
d1e3195e6f | ||
|
|
05a35662ae | ||
|
|
ce53d3a287 | ||
|
|
4cc99e7449 | ||
|
|
71773fe032 | ||
|
|
a1e0fa0f39 | ||
|
|
fc2f0b6983 | ||
|
|
5c9997cdac | ||
|
|
6f81046730 | ||
|
|
0687472d01 | ||
|
|
7739738fb3 | ||
|
|
99d1ce247b | ||
|
|
f5941a411c | ||
|
|
ba672bbd07 | ||
|
|
d9c6627a53 | ||
|
|
2e9907c3ac | ||
|
|
90afb9cb73 | ||
|
|
d0cc0cd9a5 | ||
|
|
338321e553 | ||
|
|
182b31963a | ||
|
|
4f48e5254a | ||
|
|
15dd5db1d7 | ||
|
|
424711b718 | ||
|
|
91a2b1f0b4 | ||
|
|
2b134fc378 | ||
|
|
b9153719b0 | ||
|
|
631e5c8331 | ||
|
|
e9c60a0a67 | ||
|
|
98a1bb5a7f | ||
|
|
ca90487a8c | ||
|
|
1042489f85 | ||
|
|
38277c1ea6 | ||
|
|
ee0c24628f | ||
|
|
3a18f6fcca | ||
|
|
099e734a02 | ||
|
|
a52da26b5d | ||
|
|
522a68a4ea | ||
|
|
a02eda54d0 | ||
|
|
97ef633c57 | ||
|
|
dae8463ba1 | ||
|
|
7c1299922e | ||
|
|
ddcf1f279d | ||
|
|
7e6bb8fdc5 | ||
|
|
9cee8ef87b | ||
|
|
93fb841bcb | ||
|
|
0c05131aeb | ||
|
|
5ebc58fab4 | ||
|
|
2b609dd891 | ||
|
|
a8cbc68c3e | ||
|
|
11a795a01c | ||
|
|
89c428216e | ||
|
|
2695a99623 | ||
|
|
242aecd924 | ||
|
|
ce8cc1ba33 | ||
|
|
ad5253bd2b | ||
|
|
97fdd2e088 | ||
|
|
9397f7049f | ||
|
|
a14d19b92c | ||
|
|
8ae0c05ea6 | ||
|
|
8822f20d17 | ||
|
|
553d6f50ea | ||
|
|
f0e5a5a367 | ||
|
|
f6dfea9357 | ||
|
|
cc8dc7f62c | ||
|
|
a3846ea513 | ||
|
|
8d44be858e | ||
|
|
0e6bb076e9 | ||
|
|
ac135fc7cb | ||
|
|
4e1d09809d | ||
|
|
9e855f8100 | ||
|
|
25680a8259 | ||
|
|
13c93e8cfd | ||
|
|
88aa1b9fd1 | ||
|
|
352cb98ff0 | ||
|
|
ac95e92829 | ||
|
|
8526c2da25 | ||
|
|
68a6cabf8b | ||
|
|
ac0e387da1 | ||
|
|
7fe1d102cb | ||
|
|
5850492a93 | ||
|
|
fdbd4041ca | ||
|
|
ebef1fae2a | ||
|
|
c51851689b | ||
|
|
419bf784ab | ||
|
|
4bbeb92e9a | ||
|
|
b436dad8bc | ||
|
|
6ae15d6c44 | ||
|
|
0468bde0d6 | ||
|
|
1d7329e797 | ||
|
|
48ffc4dee7 | ||
|
|
7ebd8f0c44 | ||
|
|
b680c146c1 | ||
|
|
7d6660d181 | ||
|
|
d8e3d4e2b6 | ||
|
|
d26ad8224d | ||
|
|
5c84d69d42 | ||
|
|
527e4b7f26 | ||
|
|
b48485b42b | ||
|
|
79009bb3d4 | ||
|
|
26fc611f86 | ||
|
|
b43743d4f1 | ||
|
|
179e5434b1 | ||
|
|
9f95b31158 | ||
|
|
5da07eae4c | ||
|
|
835ae178d4 | ||
|
|
c80ab8bf0d | ||
|
|
ce87714ef1 | ||
|
|
0452b869e8 | ||
|
|
d2e5857b82 | ||
|
|
f9b005f21f | ||
|
|
532107b4fa | ||
|
|
c44793789b | ||
|
|
4e99525279 | ||
|
|
7547d1d0b3 | ||
|
|
68934942d0 | ||
|
|
09fec34e1c | ||
|
|
9229708b6c | ||
|
|
914db94e79 | ||
|
|
660bd7eff5 | ||
|
|
b907d21851 | ||
|
|
dd44413ba5 | ||
|
|
10fa0f2062 | ||
|
|
d6cc976d1f | ||
|
|
8aa2cce8c5 | ||
|
|
bf9b2c49df | ||
|
|
77b42c6165 | ||
|
|
446150a747 | ||
|
|
1cbc4834e1 | ||
|
|
30338ecec4 | ||
|
|
9a37defed3 | ||
|
|
c83a057996 | ||
|
|
a8a5d03c33 | ||
|
|
76aa917882 | ||
|
|
6ac9b31e4e | ||
|
|
0ad3e8457f | ||
|
|
444a47ae63 | ||
|
|
725f4fdff4 | ||
|
|
c23e46f45d | ||
|
|
b148820c35 | ||
|
|
134f41496d | ||
|
|
c5838dd58d | ||
|
|
b6ca5ef7ce | ||
|
|
1ae994b4aa | ||
|
|
84e9793e61 | ||
|
|
32e64dacfd | ||
|
|
cc1d8f6629 | ||
|
|
5446cd2b02 | ||
|
|
8de0885b7d | ||
|
|
16243f18fd | ||
|
|
a6ce5f36e6 | ||
|
|
e73cf42e28 | ||
|
|
b45343e812 | ||
|
|
8599b1560e | ||
|
|
8bde8c37c0 | ||
|
|
82df5bf88a | ||
|
|
acb1066de8 | ||
|
|
27c68f5bb2 | ||
|
|
68dd2bfe82 | ||
|
|
65a87815e7 | ||
|
|
b80793ca82 | ||
|
|
601550f238 | ||
|
|
41b1cf2273 | ||
|
|
2baf35b3ef | ||
|
|
846e75b893 | ||
|
|
fc0257d6d9 | ||
|
|
f3c164d345 | ||
|
|
4040b1e766 | ||
|
|
3b4f9f43db | ||
|
|
37a09ecb23 | ||
|
|
0da34d3c2d | ||
|
|
74bf7eda8f | ||
|
|
9032042cfa | ||
|
|
030bf5e6c7 | ||
|
|
d3100085b0 | ||
|
|
f481d25133 | ||
|
|
8c6c90da74 | ||
|
|
24bcfd9c03 | ||
|
|
816fb4c5da | ||
|
|
c1bb77c7c9 | ||
|
|
6bcac3a55a | ||
|
|
fc346f4537 | ||
|
|
43e531a3b6 | ||
|
|
d24ea4ce2a | ||
|
|
2c30c981ae | ||
|
|
aa1da8a858 | ||
|
|
f1e9a787d7 | ||
|
|
4eeec297de | ||
|
|
77cc4ce3a0 | ||
|
|
37dfea1d3f | ||
|
|
e6626c672a | ||
|
|
c66cb0afd2 | ||
|
|
fb48eee973 | ||
|
|
bb44e5ec44 | ||
|
|
c785c1a3ca | ||
|
|
514ae341c8 | ||
|
|
0659ffab75 | ||
|
|
8ce07f38dd | ||
|
|
7cb398d167 | ||
|
|
c3e12c5e58 | ||
|
|
1825fc7503 | ||
|
|
48732ba05e | ||
|
|
acf483c9e6 | ||
|
|
3b3e0d1141 | ||
|
|
7acd428507 | ||
|
|
0aaf177640 | ||
|
|
450d1227bd | ||
|
|
b7588428c5 | ||
|
|
492b9c46f0 | ||
|
|
6e634fe3f9 | ||
|
|
4e26182d14 | ||
|
|
8f97a5f77c | ||
|
|
eb7571936c | ||
|
|
5382764d8a | ||
|
|
49c8ec69d0 | ||
|
|
3b421c8181 | ||
|
|
2a4d3e60f3 | ||
|
|
8b5af2ab84 | ||
|
|
d887716ebd | ||
|
|
5dc1848466 | ||
|
|
9491517b26 | ||
|
|
9370b5bd04 | ||
|
|
abb51a0d93 | ||
|
|
c8d809131b | ||
|
|
dd71c73a9f | ||
|
|
afc8a0f9be | ||
|
|
a99522224f | ||
|
|
f5d46b9ca2 | ||
|
|
d693d7993b | ||
|
|
5936f9895c | ||
|
|
0cbfe7f457 | ||
|
|
b9ae4ab803 | ||
|
|
65debb874f | ||
|
|
3caadac003 | ||
|
|
6a9e3a6b84 | ||
|
|
269972440a | ||
|
|
cce13e6ad2 | ||
|
|
8a565dcad8 | ||
|
|
d536110404 | ||
|
|
48e957ddff | ||
|
|
94563d622c | ||
|
|
2615f489d6 |
@@ -31,6 +31,7 @@ bin/*
|
||||
.agent/*
|
||||
.agents/*
|
||||
.opencode/*
|
||||
.idea/*
|
||||
.bmad/*
|
||||
_bmad/*
|
||||
_bmad-output/*
|
||||
|
||||
14
.github/workflows/docker-image.yml
vendored
14
.github/workflows/docker-image.yml
vendored
@@ -16,6 +16,10 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
- name: Refresh models catalog
|
||||
run: |
|
||||
git fetch --depth 1 https://github.com/router-for-me/models.git main
|
||||
git show FETCH_HEAD:models.json > internal/registry/models/models.json
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
- name: Login to DockerHub
|
||||
@@ -25,7 +29,7 @@ jobs:
|
||||
password: ${{ secrets.DOCKERHUB_TOKEN }}
|
||||
- name: Generate Build Metadata
|
||||
run: |
|
||||
echo VERSION=`git describe --tags --always --dirty` >> $GITHUB_ENV
|
||||
echo "VERSION=${GITHUB_REF_NAME}" >> $GITHUB_ENV
|
||||
echo COMMIT=`git rev-parse --short HEAD` >> $GITHUB_ENV
|
||||
echo BUILD_DATE=`date -u +%Y-%m-%dT%H:%M:%SZ` >> $GITHUB_ENV
|
||||
- name: Build and push (amd64)
|
||||
@@ -47,6 +51,10 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
- name: Refresh models catalog
|
||||
run: |
|
||||
git fetch --depth 1 https://github.com/router-for-me/models.git main
|
||||
git show FETCH_HEAD:models.json > internal/registry/models/models.json
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
- name: Login to DockerHub
|
||||
@@ -56,7 +64,7 @@ jobs:
|
||||
password: ${{ secrets.DOCKERHUB_TOKEN }}
|
||||
- name: Generate Build Metadata
|
||||
run: |
|
||||
echo VERSION=`git describe --tags --always --dirty` >> $GITHUB_ENV
|
||||
echo "VERSION=${GITHUB_REF_NAME}" >> $GITHUB_ENV
|
||||
echo COMMIT=`git rev-parse --short HEAD` >> $GITHUB_ENV
|
||||
echo BUILD_DATE=`date -u +%Y-%m-%dT%H:%M:%SZ` >> $GITHUB_ENV
|
||||
- name: Build and push (arm64)
|
||||
@@ -90,7 +98,7 @@ jobs:
|
||||
password: ${{ secrets.DOCKERHUB_TOKEN }}
|
||||
- name: Generate Build Metadata
|
||||
run: |
|
||||
echo VERSION=`git describe --tags --always --dirty` >> $GITHUB_ENV
|
||||
echo "VERSION=${GITHUB_REF_NAME}" >> $GITHUB_ENV
|
||||
echo COMMIT=`git rev-parse --short HEAD` >> $GITHUB_ENV
|
||||
echo BUILD_DATE=`date -u +%Y-%m-%dT%H:%M:%SZ` >> $GITHUB_ENV
|
||||
- name: Create and push multi-arch manifests
|
||||
|
||||
4
.github/workflows/pr-test-build.yml
vendored
4
.github/workflows/pr-test-build.yml
vendored
@@ -12,6 +12,10 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
- name: Refresh models catalog
|
||||
run: |
|
||||
git fetch --depth 1 https://github.com/router-for-me/models.git main
|
||||
git show FETCH_HEAD:models.json > internal/registry/models/models.json
|
||||
- name: Set up Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
|
||||
9
.github/workflows/release.yaml
vendored
9
.github/workflows/release.yaml
vendored
@@ -16,6 +16,10 @@ jobs:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- name: Refresh models catalog
|
||||
run: |
|
||||
git fetch --depth 1 https://github.com/router-for-me/models.git main
|
||||
git show FETCH_HEAD:models.json > internal/registry/models/models.json
|
||||
- run: git fetch --force --tags
|
||||
- uses: actions/setup-go@v4
|
||||
with:
|
||||
@@ -23,15 +27,14 @@ jobs:
|
||||
cache: true
|
||||
- name: Generate Build Metadata
|
||||
run: |
|
||||
VERSION=$(git describe --tags --always --dirty)
|
||||
echo "VERSION=${VERSION}" >> $GITHUB_ENV
|
||||
echo "VERSION=${GITHUB_REF_NAME}" >> $GITHUB_ENV
|
||||
echo COMMIT=`git rev-parse --short HEAD` >> $GITHUB_ENV
|
||||
echo BUILD_DATE=`date -u +%Y-%m-%dT%H:%M:%SZ` >> $GITHUB_ENV
|
||||
- uses: goreleaser/goreleaser-action@v4
|
||||
with:
|
||||
distribution: goreleaser
|
||||
version: latest
|
||||
args: release --clean
|
||||
args: release --clean --skip=validate
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
VERSION: ${{ env.VERSION }}
|
||||
|
||||
2
.gitignore
vendored
2
.gitignore
vendored
@@ -1,6 +1,7 @@
|
||||
# Binaries
|
||||
cli-proxy-api
|
||||
cliproxy
|
||||
/server
|
||||
*.exe
|
||||
|
||||
|
||||
@@ -44,6 +45,7 @@ GEMINI.md
|
||||
.agents/*
|
||||
.agents/*
|
||||
.opencode/*
|
||||
.idea/*
|
||||
.bmad/*
|
||||
_bmad/*
|
||||
_bmad-output/*
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
version: 2
|
||||
|
||||
builds:
|
||||
- id: "cli-proxy-api-plus"
|
||||
env:
|
||||
@@ -6,6 +8,7 @@ builds:
|
||||
- linux
|
||||
- windows
|
||||
- darwin
|
||||
- freebsd
|
||||
goarch:
|
||||
- amd64
|
||||
- arm64
|
||||
|
||||
81
README.md
81
README.md
@@ -8,87 +8,6 @@ All third-party provider support is maintained by community contributors; CLIPro
|
||||
|
||||
The Plus release stays in lockstep with the mainline features.
|
||||
|
||||
## Differences from the Mainline
|
||||
|
||||
- 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), [Ravens2121](https://github.com/Ravens2121/CLIProxyAPIPlus/)
|
||||
|
||||
## New Features (Plus Enhanced)
|
||||
|
||||
- **OAuth Web Authentication**: Browser-based OAuth login for Kiro with beautiful web UI
|
||||
- **Rate Limiter**: Built-in request rate limiting to prevent API abuse
|
||||
- **Background Token Refresh**: Automatic token refresh 10 minutes before expiration
|
||||
- **Metrics & Monitoring**: Request metrics collection for monitoring and debugging
|
||||
- **Device Fingerprint**: Device fingerprint generation for enhanced security
|
||||
- **Cooldown Management**: Smart cooldown mechanism for API rate limits
|
||||
- **Usage Checker**: Real-time usage monitoring and quota management
|
||||
- **Model Converter**: Unified model name conversion across providers
|
||||
- **UTF-8 Stream Processing**: Improved streaming response handling
|
||||
|
||||
## Kiro Authentication
|
||||
|
||||
### Web-based OAuth Login
|
||||
|
||||
Access the Kiro OAuth web interface at:
|
||||
|
||||
```
|
||||
http://your-server:8080/v0/oauth/kiro
|
||||
```
|
||||
|
||||
This provides a browser-based OAuth flow for Kiro (AWS CodeWhisperer) authentication with:
|
||||
- AWS Builder ID login
|
||||
- AWS Identity Center (IDC) login
|
||||
- Token import from Kiro IDE
|
||||
|
||||
## Quick Deployment with Docker
|
||||
|
||||
### One-Command Deployment
|
||||
|
||||
```bash
|
||||
# Create deployment directory
|
||||
mkdir -p ~/cli-proxy && cd ~/cli-proxy
|
||||
|
||||
# Create docker-compose.yml
|
||||
cat > docker-compose.yml << 'EOF'
|
||||
services:
|
||||
cli-proxy-api:
|
||||
image: eceasy/cli-proxy-api-plus:latest
|
||||
container_name: cli-proxy-api-plus
|
||||
ports:
|
||||
- "8317:8317"
|
||||
volumes:
|
||||
- ./config.yaml:/CLIProxyAPI/config.yaml
|
||||
- ./auths:/root/.cli-proxy-api
|
||||
- ./logs:/CLIProxyAPI/logs
|
||||
restart: unless-stopped
|
||||
EOF
|
||||
|
||||
# Download example config
|
||||
curl -o config.yaml https://raw.githubusercontent.com/router-for-me/CLIProxyAPIPlus/main/config.example.yaml
|
||||
|
||||
# Pull and start
|
||||
docker compose pull && docker compose up -d
|
||||
```
|
||||
|
||||
### Configuration
|
||||
|
||||
Edit `config.yaml` before starting:
|
||||
|
||||
```yaml
|
||||
# Basic configuration example
|
||||
server:
|
||||
port: 8317
|
||||
|
||||
# Add your provider configurations here
|
||||
```
|
||||
|
||||
### Update to Latest Version
|
||||
|
||||
```bash
|
||||
cd ~/cli-proxy
|
||||
docker compose pull && docker compose up -d
|
||||
```
|
||||
|
||||
## Contributing
|
||||
|
||||
This project only accepts pull requests that relate to third-party provider support. Any pull requests unrelated to third-party provider support will be rejected.
|
||||
|
||||
87
README_CN.md
87
README_CN.md
@@ -1,94 +1,11 @@
|
||||
# CLIProxyAPI Plus
|
||||
|
||||
[English](README.md) | 中文
|
||||
[English](README.md) | 中文 | [日本語](README_JA.md)
|
||||
|
||||
这是 [CLIProxyAPI](https://github.com/router-for-me/CLIProxyAPI) 的 Plus 版本,在原有基础上增加了第三方供应商的支持。
|
||||
|
||||
所有的第三方供应商支持都由第三方社区维护者提供,CLIProxyAPI 不提供技术支持。如需取得支持,请与对应的社区维护者联系。
|
||||
|
||||
该 Plus 版本的主线功能与主线功能强制同步。
|
||||
|
||||
## 与主线版本版本差异
|
||||
|
||||
- 新增 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)、[Ravens2121](https://github.com/Ravens2121/CLIProxyAPIPlus/)提供
|
||||
|
||||
## 新增功能 (Plus 增强版)
|
||||
|
||||
- **OAuth Web 认证**: 基于浏览器的 Kiro OAuth 登录,提供美观的 Web UI
|
||||
- **请求限流器**: 内置请求限流,防止 API 滥用
|
||||
- **后台令牌刷新**: 过期前 10 分钟自动刷新令牌
|
||||
- **监控指标**: 请求指标收集,用于监控和调试
|
||||
- **设备指纹**: 设备指纹生成,增强安全性
|
||||
- **冷却管理**: 智能冷却机制,应对 API 速率限制
|
||||
- **用量检查器**: 实时用量监控和配额管理
|
||||
- **模型转换器**: 跨供应商的统一模型名称转换
|
||||
- **UTF-8 流处理**: 改进的流式响应处理
|
||||
|
||||
## Kiro 认证
|
||||
|
||||
### 网页端 OAuth 登录
|
||||
|
||||
访问 Kiro OAuth 网页认证界面:
|
||||
|
||||
```
|
||||
http://your-server:8080/v0/oauth/kiro
|
||||
```
|
||||
|
||||
提供基于浏览器的 Kiro (AWS CodeWhisperer) OAuth 认证流程,支持:
|
||||
- AWS Builder ID 登录
|
||||
- AWS Identity Center (IDC) 登录
|
||||
- 从 Kiro IDE 导入令牌
|
||||
|
||||
## Docker 快速部署
|
||||
|
||||
### 一键部署
|
||||
|
||||
```bash
|
||||
# 创建部署目录
|
||||
mkdir -p ~/cli-proxy && cd ~/cli-proxy
|
||||
|
||||
# 创建 docker-compose.yml
|
||||
cat > docker-compose.yml << 'EOF'
|
||||
services:
|
||||
cli-proxy-api:
|
||||
image: eceasy/cli-proxy-api-plus:latest
|
||||
container_name: cli-proxy-api-plus
|
||||
ports:
|
||||
- "8317:8317"
|
||||
volumes:
|
||||
- ./config.yaml:/CLIProxyAPI/config.yaml
|
||||
- ./auths:/root/.cli-proxy-api
|
||||
- ./logs:/CLIProxyAPI/logs
|
||||
restart: unless-stopped
|
||||
EOF
|
||||
|
||||
# 下载示例配置
|
||||
curl -o config.yaml https://raw.githubusercontent.com/router-for-me/CLIProxyAPIPlus/main/config.example.yaml
|
||||
|
||||
# 拉取并启动
|
||||
docker compose pull && docker compose up -d
|
||||
```
|
||||
|
||||
### 配置说明
|
||||
|
||||
启动前请编辑 `config.yaml`:
|
||||
|
||||
```yaml
|
||||
# 基本配置示例
|
||||
server:
|
||||
port: 8317
|
||||
|
||||
# 在此添加你的供应商配置
|
||||
```
|
||||
|
||||
### 更新到最新版本
|
||||
|
||||
```bash
|
||||
cd ~/cli-proxy
|
||||
docker compose pull && docker compose up -d
|
||||
```
|
||||
|
||||
## 贡献
|
||||
|
||||
该项目仅接受第三方供应商支持的 Pull Request。任何非第三方供应商支持的 Pull Request 都将被拒绝。
|
||||
@@ -97,4 +14,4 @@ docker compose pull && docker compose up -d
|
||||
|
||||
## 许可证
|
||||
|
||||
此项目根据 MIT 许可证授权 - 有关详细信息,请参阅 [LICENSE](LICENSE) 文件。
|
||||
此项目根据 MIT 许可证授权 - 有关详细信息,请参阅 [LICENSE](LICENSE) 文件。
|
||||
|
||||
187
README_JA.md
Normal file
187
README_JA.md
Normal file
@@ -0,0 +1,187 @@
|
||||
# CLI Proxy API
|
||||
|
||||
[English](README.md) | [中文](README_CN.md) | 日本語
|
||||
|
||||
CLI向けのOpenAI/Gemini/Claude/Codex互換APIインターフェースを提供するプロキシサーバーです。
|
||||
|
||||
OAuth経由でOpenAI Codex(GPTモデル)およびClaude Codeもサポートしています。
|
||||
|
||||
ローカルまたはマルチアカウントのCLIアクセスを、OpenAI(Responses含む)/Gemini/Claude互換のクライアントやSDKで利用できます。
|
||||
|
||||
## スポンサー
|
||||
|
||||
[](https://z.ai/subscribe?ic=8JVLJQFSKB)
|
||||
|
||||
本プロジェクトはZ.aiにスポンサーされており、GLM CODING PLANの提供を受けています。
|
||||
|
||||
GLM CODING PLANはAIコーディング向けに設計されたサブスクリプションサービスで、月額わずか$10から利用可能です。フラッグシップのGLM-4.7および(GLM-5はProユーザーのみ利用可能)モデルを10以上の人気AIコーディングツール(Claude Code、Cline、Roo Codeなど)で利用でき、開発者にトップクラスの高速かつ安定したコーディング体験を提供します。
|
||||
|
||||
GLM CODING PLANを10%割引で取得:https://z.ai/subscribe?ic=8JVLJQFSKB
|
||||
|
||||
---
|
||||
|
||||
<table>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td width="180"><a href="https://www.packyapi.com/register?aff=cliproxyapi"><img src="./assets/packycode.png" alt="PackyCode" width="150"></a></td>
|
||||
<td>PackyCodeのスポンサーシップに感謝します!PackyCodeは信頼性が高く効率的なAPIリレーサービスプロバイダーで、Claude Code、Codex、Geminiなどのリレーサービスを提供しています。PackyCodeは当ソフトウェアのユーザーに特別割引を提供しています:<a href="https://www.packyapi.com/register?aff=cliproxyapi">こちらのリンク</a>から登録し、チャージ時にプロモーションコード「cliproxyapi」を入力すると10%割引になります。</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td width="180"><a href="https://www.aicodemirror.com/register?invitecode=TJNAIF"><img src="./assets/aicodemirror.png" alt="AICodeMirror" width="150"></a></td>
|
||||
<td>AICodeMirrorのスポンサーシップに感謝します!AICodeMirrorはClaude Code / Codex / Gemini CLI向けの公式高安定性リレーサービスを提供しており、エンタープライズグレードの同時接続、迅速な請求書発行、24時間365日の専任技術サポートを備えています。Claude Code / Codex / Geminiの公式チャネルが元の価格の38% / 2% / 9%で利用でき、チャージ時にはさらに割引があります!CLIProxyAPIユーザー向けの特別特典:<a href="https://www.aicodemirror.com/register?invitecode=TJNAIF">こちらのリンク</a>から登録すると、初回チャージが20%割引になり、エンタープライズのお客様は最大25%割引を受けられます!</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td width="180"><a href="https://shop.bmoplus.com/?utm_source=github"><img src="./assets/bmoplus.png" alt="BmoPlus" width="150"></a></td>
|
||||
<td>本プロジェクトにご支援いただいた BmoPlus に感謝いたします!BmoPlusは、AIサブスクリプションのヘビーユーザー向けに特化した信頼性の高いAIアカウントサービスプロバイダーであり、安定した ChatGPT Plus / ChatGPT Pro (完全保証) / Claude Pro / Super Grok / Gemini Pro の公式代行チャージおよび即納アカウントを提供しています。こちらの<a href="https://shop.bmoplus.com/?utm_source=github">BmoPlus AIアカウント専門店/代行チャージ</a>経由でご登録・ご注文いただいたユーザー様は、GPTを <b>公式サイト価格の約1割(90% OFF)</b> という驚異的な価格でご利用いただけます!</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
## 概要
|
||||
|
||||
- CLIモデル向けのOpenAI/Gemini/Claude互換APIエンドポイント
|
||||
- OAuthログインによるOpenAI Codexサポート(GPTモデル)
|
||||
- OAuthログインによるClaude Codeサポート
|
||||
- OAuthログインによるQwen Codeサポート
|
||||
- OAuthログインによるiFlowサポート
|
||||
- プロバイダールーティングによるAmp CLIおよびIDE拡張機能のサポート
|
||||
- ストリーミングおよび非ストリーミングレスポンス
|
||||
- 関数呼び出し/ツールのサポート
|
||||
- マルチモーダル入力サポート(テキストと画像)
|
||||
- ラウンドロビン負荷分散による複数アカウント対応(Gemini、OpenAI、Claude、QwenおよびiFlow)
|
||||
- シンプルなCLI認証フロー(Gemini、OpenAI、Claude、QwenおよびiFlow)
|
||||
- Generative Language APIキーのサポート
|
||||
- AI Studioビルドのマルチアカウント負荷分散
|
||||
- Gemini CLIのマルチアカウント負荷分散
|
||||
- Claude Codeのマルチアカウント負荷分散
|
||||
- Qwen Codeのマルチアカウント負荷分散
|
||||
- iFlowのマルチアカウント負荷分散
|
||||
- OpenAI Codexのマルチアカウント負荷分散
|
||||
- 設定によるOpenAI互換アップストリームプロバイダー(例:OpenRouter)
|
||||
- プロキシ埋め込み用の再利用可能なGo SDK(`docs/sdk-usage.md`を参照)
|
||||
|
||||
## はじめに
|
||||
|
||||
CLIProxyAPIガイド:[https://help.router-for.me/](https://help.router-for.me/)
|
||||
|
||||
## 管理API
|
||||
|
||||
[MANAGEMENT_API.md](https://help.router-for.me/management/api)を参照
|
||||
|
||||
## Amp CLIサポート
|
||||
|
||||
CLIProxyAPIは[Amp CLI](https://ampcode.com)およびAmp IDE拡張機能の統合サポートを含んでおり、Google/ChatGPT/ClaudeのOAuthサブスクリプションをAmpのコーディングツールで使用できます:
|
||||
|
||||
- Ampの APIパターン用のプロバイダールートエイリアス(`/api/provider/{provider}/v1...`)
|
||||
- OAuth認証およびアカウント機能用の管理プロキシ
|
||||
- 自動ルーティングによるスマートモデルフォールバック
|
||||
- 利用できないモデルを代替モデルにルーティングする**モデルマッピング**(例:`claude-opus-4.5` → `claude-sonnet-4`)
|
||||
- localhostのみの管理エンドポイントによるセキュリティファーストの設計
|
||||
|
||||
**→ [Amp CLI統合ガイドの完全版](https://help.router-for.me/agent-client/amp-cli.html)**
|
||||
|
||||
## SDKドキュメント
|
||||
|
||||
- 使い方:[docs/sdk-usage.md](docs/sdk-usage.md)
|
||||
- 上級(エグゼキューターとトランスレーター):[docs/sdk-advanced.md](docs/sdk-advanced.md)
|
||||
- アクセス:[docs/sdk-access.md](docs/sdk-access.md)
|
||||
- ウォッチャー:[docs/sdk-watcher.md](docs/sdk-watcher.md)
|
||||
- カスタムプロバイダーの例:`examples/custom-provider`
|
||||
|
||||
## コントリビューション
|
||||
|
||||
コントリビューションを歓迎します!お気軽にPull Requestを送ってください。
|
||||
|
||||
1. リポジトリをフォーク
|
||||
2. フィーチャーブランチを作成(`git checkout -b feature/amazing-feature`)
|
||||
3. 変更をコミット(`git commit -m 'Add some amazing feature'`)
|
||||
4. ブランチにプッシュ(`git push origin feature/amazing-feature`)
|
||||
5. Pull Requestを作成
|
||||
|
||||
## 関連プロジェクト
|
||||
|
||||
CLIProxyAPIをベースにした以下のプロジェクトがあります:
|
||||
|
||||
### [vibeproxy](https://github.com/automazeio/vibeproxy)
|
||||
|
||||
macOSネイティブのメニューバーアプリで、Claude CodeとChatGPTのサブスクリプションをAIコーディングツールで使用可能 - APIキー不要
|
||||
|
||||
### [Subtitle Translator](https://github.com/VjayC/SRT-Subtitle-Translator-Validator)
|
||||
|
||||
CLIProxyAPI経由でGeminiサブスクリプションを使用してSRT字幕を翻訳するブラウザベースのツール。自動検証/エラー修正機能付き - APIキー不要
|
||||
|
||||
### [CCS (Claude Code Switch)](https://github.com/kaitranntt/ccs)
|
||||
|
||||
CLIProxyAPI OAuthを使用して複数のClaudeアカウントや代替モデル(Gemini、Codex、Antigravity)を即座に切り替えるCLIラッパー - APIキー不要
|
||||
|
||||
### [ProxyPal](https://github.com/heyhuynhgiabuu/proxypal)
|
||||
|
||||
CLIProxyAPI管理用のmacOSネイティブGUI:OAuth経由でプロバイダー、モデルマッピング、エンドポイントを設定 - APIキー不要
|
||||
|
||||
### [Quotio](https://github.com/nguyenphutrong/quotio)
|
||||
|
||||
Claude、Gemini、OpenAI、Qwen、Antigravityのサブスクリプションを統合し、リアルタイムのクォータ追跡とスマート自動フェイルオーバーを備えたmacOSネイティブのメニューバーアプリ。Claude Code、OpenCode、Droidなどのコーディングツール向け - APIキー不要
|
||||
|
||||
### [CodMate](https://github.com/loocor/CodMate)
|
||||
|
||||
CLI AIセッション(Codex、Claude Code、Gemini CLI)を管理するmacOS SwiftUIネイティブアプリ。統合プロバイダー管理、Gitレビュー、プロジェクト整理、グローバル検索、ターミナル統合機能を搭載。CLIProxyAPIと統合し、Codex、Claude、Gemini、Antigravity、Qwen CodeのOAuth認証を提供。単一のプロキシエンドポイントを通じた組み込みおよびサードパーティプロバイダーの再ルーティングに対応 - OAuthプロバイダーではAPIキー不要
|
||||
|
||||
### [ProxyPilot](https://github.com/Finesssee/ProxyPilot)
|
||||
|
||||
TUI、システムトレイ、マルチプロバイダーOAuthを備えたWindows向けCLIProxyAPIフォーク - AIコーディングツール用、APIキー不要
|
||||
|
||||
### [Claude Proxy VSCode](https://github.com/uzhao/claude-proxy-vscode)
|
||||
|
||||
Claude Codeモデルを素早く切り替えるVSCode拡張機能。バックエンドとしてCLIProxyAPIを統合し、バックグラウンドでの自動ライフサイクル管理を搭載
|
||||
|
||||
### [ZeroLimit](https://github.com/0xtbug/zero-limit)
|
||||
|
||||
CLIProxyAPIを使用してAIコーディングアシスタントのクォータを監視するTauri + React製のWindowsデスクトップアプリ。Gemini、Claude、OpenAI Codex、Antigravityアカウントの使用量をリアルタイムダッシュボード、システムトレイ統合、ワンクリックプロキシコントロールで追跡 - APIキー不要
|
||||
|
||||
### [CPA-XXX Panel](https://github.com/ferretgeek/CPA-X)
|
||||
|
||||
CLIProxyAPI向けの軽量Web管理パネル。ヘルスチェック、リソース監視、リアルタイムログ、自動更新、リクエスト統計、料金表示機能を搭載。ワンクリックインストールとsystemdサービスに対応
|
||||
|
||||
### [CLIProxyAPI Tray](https://github.com/kitephp/CLIProxyAPI_Tray)
|
||||
|
||||
PowerShellスクリプトで実装されたWindowsトレイアプリケーション。サードパーティライブラリに依存せず、ショートカットの自動作成、サイレント実行、パスワード管理、チャネル切り替え(Main / Plus)、自動ダウンロードおよび自動更新に対応
|
||||
|
||||
### [霖君](https://github.com/wangdabaoqq/LinJun)
|
||||
|
||||
霖君はAIプログラミングアシスタントを管理するクロスプラットフォームデスクトップアプリケーションで、macOS、Windows、Linuxシステムに対応。Claude Code、Gemini CLI、OpenAI Codex、Qwen Codeなどのコーディングツールを統合管理し、ローカルプロキシによるマルチアカウントクォータ追跡とワンクリック設定が可能
|
||||
|
||||
### [CLIProxyAPI Dashboard](https://github.com/itsmylife44/cliproxyapi-dashboard)
|
||||
|
||||
Next.js、React、PostgreSQLで構築されたCLIProxyAPI用のモダンなWebベース管理ダッシュボード。リアルタイムログストリーミング、構造化された設定編集、APIキー管理、Claude/Gemini/Codex向けOAuthプロバイダー統合、使用量分析、コンテナ管理、コンパニオンプラグインによるOpenCodeとの設定同期機能を搭載 - 手動でのYAML編集は不要
|
||||
|
||||
### [All API Hub](https://github.com/qixing-jk/all-api-hub)
|
||||
|
||||
New API互換リレーサイトアカウントをワンストップで管理するブラウザ拡張機能。残高と使用量のダッシュボード、自動チェックイン、一般的なアプリへのワンクリックキーエクスポート、ページ内API可用性テスト、チャネル/モデルの同期とリダイレクト機能を搭載。Management APIを通じてCLIProxyAPIと統合し、ワンクリックでプロバイダーのインポートと設定同期が可能
|
||||
|
||||
### [Shadow AI](https://github.com/HEUDavid/shadow-ai)
|
||||
|
||||
Shadow AIは制限された環境向けに特別に設計されたAIアシスタントツールです。ウィンドウや痕跡のないステルス動作モードを提供し、LAN(ローカルエリアネットワーク)を介したクロスデバイスAI質疑応答のインタラクションと制御を可能にします。本質的には「画面/音声キャプチャ + AI推論 + 低摩擦デリバリー」の自動化コラボレーションレイヤーであり、制御されたデバイスや制限された環境でアプリケーション横断的にAIアシスタントを没入的に使用できるようユーザーを支援します。
|
||||
|
||||
> [!NOTE]
|
||||
> CLIProxyAPIをベースにプロジェクトを開発した場合は、PRを送ってこのリストに追加してください。
|
||||
|
||||
## その他の選択肢
|
||||
|
||||
以下のプロジェクトはCLIProxyAPIの移植版またはそれに触発されたものです:
|
||||
|
||||
### [9Router](https://github.com/decolua/9router)
|
||||
|
||||
CLIProxyAPIに触発されたNext.js実装。インストールと使用が簡単で、フォーマット変換(OpenAI/Claude/Gemini/Ollama)、自動フォールバック付きコンボシステム、指数バックオフ付きマルチアカウント管理、Next.js Webダッシュボード、CLIツール(Cursor、Claude Code、Cline、RooCode)のサポートをゼロから構築 - APIキー不要
|
||||
|
||||
### [OmniRoute](https://github.com/diegosouzapw/OmniRoute)
|
||||
|
||||
コーディングを止めない。無料および低コストのAIモデルへのスマートルーティングと自動フォールバック。
|
||||
|
||||
OmniRouteはマルチプロバイダーLLM向けのAIゲートウェイです:スマートルーティング、負荷分散、リトライ、フォールバックを備えたOpenAI互換エンドポイント。ポリシー、レート制限、キャッシュ、可観測性を追加して、信頼性が高くコストを意識した推論を実現します。
|
||||
|
||||
> [!NOTE]
|
||||
> CLIProxyAPIの移植版またはそれに触発されたプロジェクトを開発した場合は、PRを送ってこのリストに追加してください。
|
||||
|
||||
## ライセンス
|
||||
|
||||
本プロジェクトはMITライセンスの下でライセンスされています - 詳細は[LICENSE](LICENSE)ファイルを参照してください。
|
||||
BIN
assets/bmoplus.png
Normal file
BIN
assets/bmoplus.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 28 KiB |
275
cmd/fetch_antigravity_models/main.go
Normal file
275
cmd/fetch_antigravity_models/main.go
Normal file
@@ -0,0 +1,275 @@
|
||||
// Command fetch_antigravity_models connects to the Antigravity API using the
|
||||
// stored auth credentials and saves the dynamically fetched model list to a
|
||||
// JSON file for inspection or offline use.
|
||||
//
|
||||
// Usage:
|
||||
//
|
||||
// go run ./cmd/fetch_antigravity_models [flags]
|
||||
//
|
||||
// Flags:
|
||||
//
|
||||
// --auths-dir <path> Directory containing auth JSON files (default: "auths")
|
||||
// --output <path> Output JSON file path (default: "antigravity_models.json")
|
||||
// --pretty Pretty-print the output JSON (default: true)
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
|
||||
sdkauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/proxyutil"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
const (
|
||||
antigravityBaseURLDaily = "https://daily-cloudcode-pa.googleapis.com"
|
||||
antigravitySandboxBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com"
|
||||
antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com"
|
||||
antigravityModelsPath = "/v1internal:fetchAvailableModels"
|
||||
)
|
||||
|
||||
func init() {
|
||||
logging.SetupBaseLogger()
|
||||
log.SetLevel(log.InfoLevel)
|
||||
}
|
||||
|
||||
// modelOutput wraps the fetched model list with fetch metadata.
|
||||
type modelOutput struct {
|
||||
Models []modelEntry `json:"models"`
|
||||
}
|
||||
|
||||
// modelEntry contains only the fields we want to keep for static model definitions.
|
||||
type modelEntry struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
OwnedBy string `json:"owned_by"`
|
||||
Type string `json:"type"`
|
||||
DisplayName string `json:"display_name"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
ContextLength int `json:"context_length,omitempty"`
|
||||
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
|
||||
}
|
||||
|
||||
func main() {
|
||||
var authsDir string
|
||||
var outputPath string
|
||||
var pretty bool
|
||||
|
||||
flag.StringVar(&authsDir, "auths-dir", "auths", "Directory containing auth JSON files")
|
||||
flag.StringVar(&outputPath, "output", "antigravity_models.json", "Output JSON file path")
|
||||
flag.BoolVar(&pretty, "pretty", true, "Pretty-print the output JSON")
|
||||
flag.Parse()
|
||||
|
||||
// Resolve relative paths against the working directory.
|
||||
wd, err := os.Getwd()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "error: cannot get working directory: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if !filepath.IsAbs(authsDir) {
|
||||
authsDir = filepath.Join(wd, authsDir)
|
||||
}
|
||||
if !filepath.IsAbs(outputPath) {
|
||||
outputPath = filepath.Join(wd, outputPath)
|
||||
}
|
||||
|
||||
fmt.Printf("Scanning auth files in: %s\n", authsDir)
|
||||
|
||||
// Load all auth records from the directory.
|
||||
fileStore := sdkauth.NewFileTokenStore()
|
||||
fileStore.SetBaseDir(authsDir)
|
||||
|
||||
ctx := context.Background()
|
||||
auths, err := fileStore.List(ctx)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "error: failed to list auth files: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if len(auths) == 0 {
|
||||
fmt.Fprintf(os.Stderr, "error: no auth files found in %s\n", authsDir)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Find the first enabled antigravity auth.
|
||||
var chosen *coreauth.Auth
|
||||
for _, a := range auths {
|
||||
if a == nil || a.Disabled {
|
||||
continue
|
||||
}
|
||||
if strings.EqualFold(strings.TrimSpace(a.Provider), "antigravity") {
|
||||
chosen = a
|
||||
break
|
||||
}
|
||||
}
|
||||
if chosen == nil {
|
||||
fmt.Fprintf(os.Stderr, "error: no enabled antigravity auth found in %s\n", authsDir)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
fmt.Printf("Using auth: id=%s label=%s\n", chosen.ID, chosen.Label)
|
||||
|
||||
// Fetch models from the upstream Antigravity API.
|
||||
fmt.Println("Fetching Antigravity model list from upstream...")
|
||||
|
||||
fetchCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
defer cancel()
|
||||
|
||||
models := fetchModels(fetchCtx, chosen)
|
||||
if len(models) == 0 {
|
||||
fmt.Fprintln(os.Stderr, "warning: no models returned (API may be unavailable or token expired)")
|
||||
} else {
|
||||
fmt.Printf("Fetched %d models.\n", len(models))
|
||||
}
|
||||
|
||||
// Build the output payload.
|
||||
out := modelOutput{
|
||||
Models: models,
|
||||
}
|
||||
|
||||
// Marshal to JSON.
|
||||
var raw []byte
|
||||
if pretty {
|
||||
raw, err = json.MarshalIndent(out, "", " ")
|
||||
} else {
|
||||
raw, err = json.Marshal(out)
|
||||
}
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "error: failed to marshal JSON: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
if err = os.WriteFile(outputPath, raw, 0o644); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "error: failed to write output file %s: %v\n", outputPath, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
fmt.Printf("Model list saved to: %s\n", outputPath)
|
||||
}
|
||||
|
||||
func fetchModels(ctx context.Context, auth *coreauth.Auth) []modelEntry {
|
||||
accessToken := metaStringValue(auth.Metadata, "access_token")
|
||||
if accessToken == "" {
|
||||
fmt.Fprintln(os.Stderr, "error: no access token found in auth")
|
||||
return nil
|
||||
}
|
||||
|
||||
baseURLs := []string{antigravityBaseURLProd, antigravityBaseURLDaily, antigravitySandboxBaseURLDaily}
|
||||
|
||||
for _, baseURL := range baseURLs {
|
||||
modelsURL := baseURL + antigravityModelsPath
|
||||
|
||||
var payload []byte
|
||||
if auth != nil && auth.Metadata != nil {
|
||||
if pid, ok := auth.Metadata["project_id"].(string); ok && strings.TrimSpace(pid) != "" {
|
||||
payload = []byte(fmt.Sprintf(`{"project": "%s"}`, strings.TrimSpace(pid)))
|
||||
}
|
||||
}
|
||||
if len(payload) == 0 {
|
||||
payload = []byte(`{}`)
|
||||
}
|
||||
|
||||
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, modelsURL, strings.NewReader(string(payload)))
|
||||
if errReq != nil {
|
||||
continue
|
||||
}
|
||||
httpReq.Close = true
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
httpReq.Header.Set("User-Agent", "antigravity/1.19.6 darwin/arm64")
|
||||
|
||||
httpClient := &http.Client{Timeout: 30 * time.Second}
|
||||
if transport, _, errProxy := proxyutil.BuildHTTPTransport(auth.ProxyURL); errProxy == nil && transport != nil {
|
||||
httpClient.Transport = transport
|
||||
}
|
||||
httpResp, errDo := httpClient.Do(httpReq)
|
||||
if errDo != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
bodyBytes, errRead := io.ReadAll(httpResp.Body)
|
||||
httpResp.Body.Close()
|
||||
if errRead != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices {
|
||||
continue
|
||||
}
|
||||
|
||||
result := gjson.GetBytes(bodyBytes, "models")
|
||||
if !result.Exists() {
|
||||
continue
|
||||
}
|
||||
|
||||
var models []modelEntry
|
||||
|
||||
for originalName, modelData := range result.Map() {
|
||||
modelID := strings.TrimSpace(originalName)
|
||||
if modelID == "" {
|
||||
continue
|
||||
}
|
||||
// Skip internal/experimental models
|
||||
switch modelID {
|
||||
case "chat_20706", "chat_23310", "tab_flash_lite_preview", "tab_jump_flash_lite_preview", "gemini-2.5-flash-thinking", "gemini-2.5-pro":
|
||||
continue
|
||||
}
|
||||
|
||||
displayName := modelData.Get("displayName").String()
|
||||
if displayName == "" {
|
||||
displayName = modelID
|
||||
}
|
||||
|
||||
entry := modelEntry{
|
||||
ID: modelID,
|
||||
Object: "model",
|
||||
OwnedBy: "antigravity",
|
||||
Type: "antigravity",
|
||||
DisplayName: displayName,
|
||||
Name: modelID,
|
||||
Description: displayName,
|
||||
}
|
||||
|
||||
if maxTok := modelData.Get("maxTokens").Int(); maxTok > 0 {
|
||||
entry.ContextLength = int(maxTok)
|
||||
}
|
||||
if maxOut := modelData.Get("maxOutputTokens").Int(); maxOut > 0 {
|
||||
entry.MaxCompletionTokens = int(maxOut)
|
||||
}
|
||||
|
||||
models = append(models, entry)
|
||||
}
|
||||
|
||||
return models
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func metaStringValue(m map[string]interface{}, key string) string {
|
||||
if m == nil {
|
||||
return ""
|
||||
}
|
||||
v, ok := m[key]
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
switch val := v.(type) {
|
||||
case string:
|
||||
return val
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
20
cmd/mcpdebug/main.go
Normal file
20
cmd/mcpdebug/main.go
Normal file
@@ -0,0 +1,20 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
cursorproto "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/cursor/proto"
|
||||
)
|
||||
|
||||
func main() {
|
||||
// Encode MCP result with empty execId
|
||||
resultBytes := cursorproto.EncodeExecMcpResult(1, "", `{"test": "data"}`, false)
|
||||
fmt.Printf("Result protobuf hex: %s\n", hex.EncodeToString(resultBytes))
|
||||
fmt.Printf("Result length: %d bytes\n", len(resultBytes))
|
||||
|
||||
// Write to file for analysis
|
||||
os.WriteFile("mcp_result.bin", resultBytes)
|
||||
fmt.Println("Wrote mcp_result.bin")
|
||||
}
|
||||
32
cmd/protocheck/main.go
Normal file
32
cmd/protocheck/main.go
Normal file
@@ -0,0 +1,32 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
cursorproto "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/cursor/proto"
|
||||
)
|
||||
|
||||
func main() {
|
||||
ecm := cursorproto.NewMsg("ExecClientMessage")
|
||||
|
||||
// Try different field names
|
||||
names := []string{
|
||||
"mcp_result", "mcpResult", "McpResult", "MCP_RESULT",
|
||||
"shell_result", "shellResult",
|
||||
}
|
||||
|
||||
for _, name := range names {
|
||||
fd := ecm.Descriptor().Fields().ByName(name)
|
||||
if fd != nil {
|
||||
fmt.Printf("Found field %q: number=%d, kind=%s\n", name, fd.Number(), fd.Kind())
|
||||
} else {
|
||||
fmt.Printf("Field %q NOT FOUND\n", name)
|
||||
}
|
||||
}
|
||||
|
||||
// List all fields
|
||||
fmt.Println("\nAll fields in ExecClientMessage:")
|
||||
for i := 0; i < ecm.Descriptor().Fields().Len(); i++ {
|
||||
f := ecm.Descriptor().Fields().Get(i)
|
||||
fmt.Printf(" %d: %q (number=%d)\n", i, f.Name(), f.Number())
|
||||
}
|
||||
}
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/store"
|
||||
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/tui"
|
||||
@@ -72,21 +73,30 @@ func main() {
|
||||
// Command-line flags to control the application's behavior.
|
||||
var login bool
|
||||
var codexLogin bool
|
||||
var codexDeviceLogin bool
|
||||
var claudeLogin bool
|
||||
var qwenLogin bool
|
||||
var kiloLogin bool
|
||||
var iflowLogin bool
|
||||
var iflowCookie bool
|
||||
var gitlabLogin bool
|
||||
var gitlabTokenLogin bool
|
||||
var noBrowser bool
|
||||
var oauthCallbackPort int
|
||||
var antigravityLogin bool
|
||||
var kimiLogin bool
|
||||
var cursorLogin bool
|
||||
var kiroLogin bool
|
||||
var kiroGoogleLogin bool
|
||||
var kiroAWSLogin bool
|
||||
var kiroAWSAuthCode bool
|
||||
var kiroImport bool
|
||||
var kiroIDCLogin bool
|
||||
var kiroIDCStartURL string
|
||||
var kiroIDCRegion string
|
||||
var kiroIDCFlow string
|
||||
var githubCopilotLogin bool
|
||||
var codeBuddyLogin bool
|
||||
var projectID string
|
||||
var vertexImport string
|
||||
var configPath string
|
||||
@@ -95,33 +105,44 @@ func main() {
|
||||
var standalone bool
|
||||
var noIncognito bool
|
||||
var useIncognito bool
|
||||
var localModel bool
|
||||
|
||||
// Define command-line flags for different operation modes.
|
||||
flag.BoolVar(&login, "login", false, "Login Google Account")
|
||||
flag.BoolVar(&codexLogin, "codex-login", false, "Login to Codex using OAuth")
|
||||
flag.BoolVar(&codexDeviceLogin, "codex-device-login", false, "Login to Codex using device code flow")
|
||||
flag.BoolVar(&claudeLogin, "claude-login", false, "Login to Claude using OAuth")
|
||||
flag.BoolVar(&qwenLogin, "qwen-login", false, "Login to Qwen using OAuth")
|
||||
flag.BoolVar(&kiloLogin, "kilo-login", false, "Login to Kilo AI using device flow")
|
||||
flag.BoolVar(&iflowLogin, "iflow-login", false, "Login to iFlow using OAuth")
|
||||
flag.BoolVar(&iflowCookie, "iflow-cookie", false, "Login to iFlow using Cookie")
|
||||
flag.BoolVar(&gitlabLogin, "gitlab-login", false, "Login to GitLab Duo using OAuth")
|
||||
flag.BoolVar(&gitlabTokenLogin, "gitlab-token-login", false, "Login to GitLab Duo using a personal access token")
|
||||
flag.BoolVar(&noBrowser, "no-browser", false, "Don't open browser automatically for OAuth")
|
||||
flag.IntVar(&oauthCallbackPort, "oauth-callback-port", 0, "Override OAuth callback port (defaults to provider-specific port)")
|
||||
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(&kimiLogin, "kimi-login", false, "Login to Kimi using OAuth")
|
||||
flag.BoolVar(&cursorLogin, "cursor-login", false, "Login to Cursor 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(&kiroAWSAuthCode, "kiro-aws-authcode", false, "Login to Kiro using AWS Builder ID (authorization code flow, better UX)")
|
||||
flag.BoolVar(&kiroImport, "kiro-import", false, "Import Kiro token from Kiro IDE (~/.aws/sso/cache/kiro-auth-token.json)")
|
||||
flag.BoolVar(&kiroIDCLogin, "kiro-idc-login", false, "Login to Kiro using IAM Identity Center (IDC)")
|
||||
flag.StringVar(&kiroIDCStartURL, "kiro-idc-start-url", "", "IDC start URL (required with --kiro-idc-login)")
|
||||
flag.StringVar(&kiroIDCRegion, "kiro-idc-region", "", "IDC region (default: us-east-1)")
|
||||
flag.StringVar(&kiroIDCFlow, "kiro-idc-flow", "", "IDC flow type: authcode (default) or device")
|
||||
flag.BoolVar(&githubCopilotLogin, "github-copilot-login", false, "Login to GitHub Copilot using device flow")
|
||||
flag.BoolVar(&codeBuddyLogin, "codebuddy-login", false, "Login to CodeBuddy using browser OAuth flow")
|
||||
flag.StringVar(&projectID, "project_id", "", "Project ID (Gemini only, not required)")
|
||||
flag.StringVar(&configPath, "config", DefaultConfigPath, "Configure File Path")
|
||||
flag.StringVar(&vertexImport, "vertex-import", "", "Import Vertex service account key JSON file")
|
||||
flag.StringVar(&password, "password", "", "")
|
||||
flag.BoolVar(&tuiMode, "tui", false, "Start with terminal management UI")
|
||||
flag.BoolVar(&standalone, "standalone", false, "In TUI mode, start an embedded local server")
|
||||
flag.BoolVar(&localModel, "local-model", false, "Use embedded model catalog only, skip remote model fetching")
|
||||
|
||||
flag.CommandLine.Usage = func() {
|
||||
out := flag.CommandLine.Output()
|
||||
@@ -499,9 +520,15 @@ func main() {
|
||||
} else if githubCopilotLogin {
|
||||
// Handle GitHub Copilot login
|
||||
cmd.DoGitHubCopilotLogin(cfg, options)
|
||||
} else if codeBuddyLogin {
|
||||
// Handle CodeBuddy login
|
||||
cmd.DoCodeBuddyLogin(cfg, options)
|
||||
} else if codexLogin {
|
||||
// Handle Codex login
|
||||
cmd.DoCodexLogin(cfg, options)
|
||||
} else if codexDeviceLogin {
|
||||
// Handle Codex device-code login
|
||||
cmd.DoCodexDeviceLogin(cfg, options)
|
||||
} else if claudeLogin {
|
||||
// Handle Claude login
|
||||
cmd.DoClaudeLogin(cfg, options)
|
||||
@@ -513,32 +540,48 @@ func main() {
|
||||
cmd.DoIFlowLogin(cfg, options)
|
||||
} else if iflowCookie {
|
||||
cmd.DoIFlowCookieAuth(cfg, options)
|
||||
} else if gitlabLogin {
|
||||
cmd.DoGitLabLogin(cfg, options)
|
||||
} else if gitlabTokenLogin {
|
||||
cmd.DoGitLabTokenLogin(cfg, options)
|
||||
} else if kimiLogin {
|
||||
cmd.DoKimiLogin(cfg, options)
|
||||
} else if cursorLogin {
|
||||
cmd.DoCursorLogin(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)
|
||||
kiro.InitFingerprintConfig(cfg)
|
||||
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)
|
||||
kiro.InitFingerprintConfig(cfg)
|
||||
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)
|
||||
kiro.InitFingerprintConfig(cfg)
|
||||
cmd.DoKiroAWSLogin(cfg, options)
|
||||
} else if kiroAWSAuthCode {
|
||||
// For Kiro auth with authorization code flow (better UX)
|
||||
setKiroIncognitoMode(cfg, useIncognito, noIncognito)
|
||||
kiro.InitFingerprintConfig(cfg)
|
||||
cmd.DoKiroAWSAuthCodeLogin(cfg, options)
|
||||
} else if kiroImport {
|
||||
kiro.InitFingerprintConfig(cfg)
|
||||
cmd.DoKiroImport(cfg, options)
|
||||
} else if kiroIDCLogin {
|
||||
// For Kiro IDC auth, default to incognito mode for multi-account support
|
||||
setKiroIncognitoMode(cfg, useIncognito, noIncognito)
|
||||
kiro.InitFingerprintConfig(cfg)
|
||||
cmd.DoKiroIDCLogin(cfg, options, kiroIDCStartURL, kiroIDCRegion, kiroIDCFlow)
|
||||
} else {
|
||||
// In cloud deploy mode without config file, just wait for shutdown signals
|
||||
if isCloudDeploy && !configFileExists {
|
||||
@@ -546,10 +589,16 @@ func main() {
|
||||
cmd.WaitForCloudDeploy()
|
||||
return
|
||||
}
|
||||
if localModel && (!tuiMode || standalone) {
|
||||
log.Info("Local model mode: using embedded model catalog, remote model updates disabled")
|
||||
}
|
||||
if tuiMode {
|
||||
if standalone {
|
||||
// Standalone mode: start an embedded local server and connect TUI client to it.
|
||||
managementasset.StartAutoUpdater(context.Background(), configFilePath)
|
||||
if !localModel {
|
||||
registry.StartModelsUpdater(context.Background())
|
||||
}
|
||||
hook := tui.NewLogHook(2000)
|
||||
hook.SetFormatter(&logging.LogFormatter{})
|
||||
log.AddHook(hook)
|
||||
@@ -620,15 +669,18 @@ func main() {
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Start the main proxy service
|
||||
managementasset.StartAutoUpdater(context.Background(), configFilePath)
|
||||
// Start the main proxy service
|
||||
managementasset.StartAutoUpdater(context.Background(), configFilePath)
|
||||
if !localModel {
|
||||
registry.StartModelsUpdater(context.Background())
|
||||
}
|
||||
|
||||
if cfg.AuthDir != "" {
|
||||
kiro.InitializeAndStart(cfg.AuthDir, cfg)
|
||||
defer kiro.StopGlobalRefreshManager()
|
||||
}
|
||||
if cfg.AuthDir != "" {
|
||||
kiro.InitializeAndStart(cfg.AuthDir, cfg)
|
||||
defer kiro.StopGlobalRefreshManager()
|
||||
}
|
||||
|
||||
cmd.StartService(cfg, configFilePath, password)
|
||||
cmd.StartService(cfg, configFilePath, password)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -25,6 +25,10 @@ remote-management:
|
||||
# Disable the bundled management control panel asset download and HTTP route when true.
|
||||
disable-control-panel: false
|
||||
|
||||
# Disable automatic periodic background updates of the management panel from GitHub (default: false).
|
||||
# When enabled, the panel is only downloaded on first access if missing, and never auto-updated afterward.
|
||||
# disable-auto-update-panel: false
|
||||
|
||||
# GitHub repository for the management control panel. Accepts a repository URL or releases API URL.
|
||||
panel-github-repository: 'https://github.com/router-for-me/Cli-Proxy-API-Management-Center'
|
||||
|
||||
@@ -68,7 +72,8 @@ error-logs-max-files: 10
|
||||
usage-statistics-enabled: false
|
||||
|
||||
# Proxy URL. Supports socks5/http/https protocols. Example: socks5://user:pass@192.168.1.1:1080/
|
||||
proxy-url: ''
|
||||
# Per-entry proxy-url also supports "direct" or "none" to bypass both the global proxy-url and environment proxies explicitly.
|
||||
proxy-url: ""
|
||||
|
||||
# When true, unprefixed model requests only use credentials without a prefix (except when prefix == model name).
|
||||
force-model-prefix: false
|
||||
@@ -80,6 +85,10 @@ passthrough-headers: false
|
||||
# Number of times to retry a request. Retries will occur if the HTTP response code is 403, 408, 500, 502, 503, or 504.
|
||||
request-retry: 3
|
||||
|
||||
# Maximum number of different credentials to try for one failed request.
|
||||
# Set to 0 to keep legacy behavior (try all available credentials).
|
||||
max-retry-credentials: 0
|
||||
|
||||
# Maximum wait time in seconds for a cooled-down credential before triggering a retry.
|
||||
max-retry-interval: 30
|
||||
|
||||
@@ -111,6 +120,7 @@ nonstream-keepalive-interval: 0
|
||||
# headers:
|
||||
# X-Custom-Header: "custom-value"
|
||||
# proxy-url: "socks5://proxy.example.com:1080"
|
||||
# # proxy-url: "direct" # optional: explicit direct connect for this credential
|
||||
# models:
|
||||
# - name: "gemini-2.5-flash" # upstream model name
|
||||
# alias: "gemini-flash" # client alias mapped to the upstream model
|
||||
@@ -129,6 +139,7 @@ nonstream-keepalive-interval: 0
|
||||
# headers:
|
||||
# X-Custom-Header: "custom-value"
|
||||
# proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override
|
||||
# # proxy-url: "direct" # optional: explicit direct connect for this credential
|
||||
# models:
|
||||
# - name: "gpt-5-codex" # upstream model name
|
||||
# alias: "codex-latest" # client alias mapped to the upstream model
|
||||
@@ -147,6 +158,7 @@ nonstream-keepalive-interval: 0
|
||||
# headers:
|
||||
# X-Custom-Header: "custom-value"
|
||||
# proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override
|
||||
# # proxy-url: "direct" # optional: explicit direct connect for this credential
|
||||
# models:
|
||||
# - name: "claude-3-5-sonnet-20241022" # upstream model name
|
||||
# alias: "claude-sonnet-latest" # client alias mapped to the upstream model
|
||||
@@ -167,18 +179,35 @@ nonstream-keepalive-interval: 0
|
||||
# cache-user-id: true # optional: default is false; set true to reuse cached user_id per API key instead of generating a random one each request
|
||||
|
||||
# Default headers for Claude API requests. Update when Claude Code releases new versions.
|
||||
# These are used as fallbacks when the client does not send its own headers.
|
||||
# In legacy mode, user-agent/package-version/runtime-version/timeout are used as fallbacks
|
||||
# when the client omits them, while OS/arch remain runtime-derived. When
|
||||
# stabilize-device-profile is enabled, OS/arch stay pinned to the baseline values below,
|
||||
# while user-agent/package-version/runtime-version seed a software fingerprint that can
|
||||
# still upgrade to newer official Claude client versions.
|
||||
# claude-header-defaults:
|
||||
# user-agent: "claude-cli/2.1.44 (external, sdk-cli)"
|
||||
# package-version: "0.74.0"
|
||||
# runtime-version: "v24.3.0"
|
||||
# os: "MacOS"
|
||||
# arch: "arm64"
|
||||
# timeout: "600"
|
||||
# stabilize-device-profile: false # optional, default false; set true to enable per-auth/API-key fingerprint pinning
|
||||
|
||||
# Default headers for Codex OAuth model requests.
|
||||
# These are used only for file-backed/OAuth Codex requests when the client
|
||||
# does not send the header. `user-agent` applies to HTTP and websocket requests;
|
||||
# `beta-features` only applies to websocket requests. They do not apply to codex-api-key entries.
|
||||
# codex-header-defaults:
|
||||
# user-agent: "codex_cli_rs/0.114.0 (Mac OS 14.2.0; x86_64) vscode/1.111.0"
|
||||
# beta-features: "multi_agent"
|
||||
|
||||
# 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)
|
||||
# start-url: "https://your-company.awsapps.com/start" # optional: IDC start URL (preset for login)
|
||||
# region: "us-east-1" # optional: OIDC region for IDC login and token refresh
|
||||
# - access-token: "aoaAAAAA..." # or provide tokens directly
|
||||
# refresh-token: "aorAAAAA..."
|
||||
# profile-arn: "arn:aws:codewhisperer:us-east-1:..."
|
||||
@@ -209,17 +238,32 @@ nonstream-keepalive-interval: 0
|
||||
# api-key-entries:
|
||||
# - api-key: "sk-or-v1-...b780"
|
||||
# proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override
|
||||
# # proxy-url: "direct" # optional: explicit direct connect for this credential
|
||||
# - 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.
|
||||
# alias: "kimi-k2" # The alias used in the API.
|
||||
# thinking: # optional: omit to default to levels ["low","medium","high"]
|
||||
# levels: ["low", "medium", "high"]
|
||||
# # You may repeat the same alias to build an internal model pool.
|
||||
# # The client still sees only one alias in the model list.
|
||||
# # Requests to that alias will round-robin across the upstream names below,
|
||||
# # and if the chosen upstream fails before producing output, the request will
|
||||
# # continue with the next upstream model in the same alias pool.
|
||||
# - name: "qwen3.5-plus"
|
||||
# alias: "claude-opus-4.66"
|
||||
# - name: "glm-5"
|
||||
# alias: "claude-opus-4.66"
|
||||
# - name: "kimi-k2.5"
|
||||
# alias: "claude-opus-4.66"
|
||||
|
||||
# Vertex API keys (Vertex-compatible endpoints, use API key + base URL)
|
||||
# Vertex API keys (Vertex-compatible endpoints, base-url is optional)
|
||||
# vertex-api-key:
|
||||
# - api-key: "vk-123..." # x-goog-api-key header
|
||||
# prefix: "test" # optional: require calls like "test/vertex-pro" to target this credential
|
||||
# base-url: "https://example.com/api" # e.g. https://zenmux.ai/api
|
||||
# base-url: "https://example.com/api" # optional, e.g. https://zenmux.ai/api; falls back to Google Vertex when omitted
|
||||
# proxy-url: "socks5://proxy.example.com:1080" # optional per-key proxy override
|
||||
# # proxy-url: "direct" # optional: explicit direct connect for this credential
|
||||
# headers:
|
||||
# X-Custom-Header: "custom-value"
|
||||
# models: # optional: map aliases to upstream model names
|
||||
@@ -227,6 +271,9 @@ nonstream-keepalive-interval: 0
|
||||
# alias: "vertex-flash" # client-visible alias
|
||||
# - name: "gemini-2.5-pro"
|
||||
# alias: "vertex-pro"
|
||||
# excluded-models: # optional: models to exclude from listing
|
||||
# - "imagen-3.0-generate-002"
|
||||
# - "imagen-*"
|
||||
|
||||
# Amp Integration
|
||||
# ampcode:
|
||||
|
||||
115
docs/gitlab-duo.md
Normal file
115
docs/gitlab-duo.md
Normal file
@@ -0,0 +1,115 @@
|
||||
# GitLab Duo guide
|
||||
|
||||
CLIProxyAPI can now use GitLab Duo as a first-class provider instead of treating it as a plain text wrapper.
|
||||
|
||||
It supports:
|
||||
|
||||
- OAuth login
|
||||
- personal access token login
|
||||
- automatic refresh of GitLab `direct_access` metadata
|
||||
- dynamic model discovery from GitLab metadata
|
||||
- native GitLab AI gateway routing for Anthropic and OpenAI/Codex managed models
|
||||
- Claude-compatible and OpenAI-compatible downstream APIs
|
||||
|
||||
## What this means
|
||||
|
||||
If GitLab Duo returns an Anthropic-managed model, CLIProxyAPI routes requests through the GitLab AI gateway Anthropic proxy and uses the existing Claude executor path.
|
||||
|
||||
If GitLab Duo returns an OpenAI-managed model, CLIProxyAPI routes requests through the GitLab AI gateway OpenAI proxy and uses the existing Codex/OpenAI executor path.
|
||||
|
||||
That gives GitLab Duo much closer runtime behavior to the built-in `codex` provider:
|
||||
|
||||
- Claude-compatible clients can use GitLab Duo models through `/v1/messages`
|
||||
- OpenAI-compatible clients can use GitLab Duo models through `/v1/chat/completions`
|
||||
- OpenAI Responses clients can use GitLab Duo models through `/v1/responses`
|
||||
|
||||
The model list is not hardcoded. CLIProxyAPI reads the current model metadata from GitLab `direct_access` and registers:
|
||||
|
||||
- a stable alias: `gitlab-duo`
|
||||
- any discovered managed model names, such as `claude-sonnet-4-5` or `gpt-5-codex`
|
||||
|
||||
## Login
|
||||
|
||||
OAuth login:
|
||||
|
||||
```bash
|
||||
./CLIProxyAPI -gitlab-login
|
||||
```
|
||||
|
||||
PAT login:
|
||||
|
||||
```bash
|
||||
./CLIProxyAPI -gitlab-token-login
|
||||
```
|
||||
|
||||
You can also provide inputs through environment variables:
|
||||
|
||||
```bash
|
||||
export GITLAB_BASE_URL=https://gitlab.com
|
||||
export GITLAB_OAUTH_CLIENT_ID=your-client-id
|
||||
export GITLAB_OAUTH_CLIENT_SECRET=your-client-secret
|
||||
export GITLAB_PERSONAL_ACCESS_TOKEN=glpat-...
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- OAuth requires a GitLab OAuth application.
|
||||
- PAT login requires a personal access token that can call the GitLab APIs used by Duo. In practice, `api` scope is the safe baseline.
|
||||
- Self-managed GitLab instances are supported through `GITLAB_BASE_URL`.
|
||||
|
||||
## Using the models
|
||||
|
||||
After login, start CLIProxyAPI normally and point your client at the local proxy.
|
||||
|
||||
You can select:
|
||||
|
||||
- `gitlab-duo` to use the current Duo-managed model for that account
|
||||
- the discovered provider model name if you want to pin it explicitly
|
||||
|
||||
Examples:
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8080/v1/models
|
||||
```
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8080/v1/chat/completions \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model": "gitlab-duo",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Write a Go HTTP middleware for request IDs."}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
If the GitLab account is currently mapped to an Anthropic model, Claude-compatible clients can use the same account through the Claude handler path. If the account is currently mapped to an OpenAI/Codex model, OpenAI-compatible clients can use `/v1/chat/completions` or `/v1/responses`.
|
||||
|
||||
## How model freshness works
|
||||
|
||||
CLIProxyAPI does not ship a fixed GitLab Duo model catalog.
|
||||
|
||||
Instead, it refreshes GitLab `direct_access` metadata and uses the returned `model_details` and any discovered model list entries to keep the local registry aligned with the current GitLab-managed model assignment.
|
||||
|
||||
This matches GitLab's current public contract better than hardcoding model names.
|
||||
|
||||
## Current scope
|
||||
|
||||
The GitLab Duo provider now has:
|
||||
|
||||
- OAuth and PAT auth flows
|
||||
- runtime refresh of Duo gateway credentials
|
||||
- native Anthropic gateway routing
|
||||
- native OpenAI/Codex gateway routing
|
||||
- handler-level smoke tests for Claude-compatible and OpenAI-compatible paths
|
||||
|
||||
Still out of scope today:
|
||||
|
||||
- websocket or session-specific parity beyond the current HTTP APIs
|
||||
- GitLab-specific IDE features that are not exposed through the public gateway contract
|
||||
|
||||
## References
|
||||
|
||||
- GitLab Code Suggestions API: https://docs.gitlab.com/api/code_suggestions/
|
||||
- GitLab Agent Assistant and managed credentials: https://docs.gitlab.com/user/duo_agent_platform/agent_assistant/
|
||||
- GitLab Duo model selection: https://docs.gitlab.com/user/gitlab_duo/model_selection/
|
||||
115
docs/gitlab-duo_CN.md
Normal file
115
docs/gitlab-duo_CN.md
Normal file
@@ -0,0 +1,115 @@
|
||||
# GitLab Duo 使用说明
|
||||
|
||||
CLIProxyAPI 现在可以把 GitLab Duo 当作一等 Provider 来使用,而不是仅仅把它当成简单的文本补全封装。
|
||||
|
||||
当前支持:
|
||||
|
||||
- OAuth 登录
|
||||
- personal access token 登录
|
||||
- 自动刷新 GitLab `direct_access` 元数据
|
||||
- 根据 GitLab 返回的元数据动态发现模型
|
||||
- 针对 Anthropic 和 OpenAI/Codex 托管模型的 GitLab AI gateway 原生路由
|
||||
- Claude 兼容与 OpenAI 兼容下游 API
|
||||
|
||||
## 这意味着什么
|
||||
|
||||
如果 GitLab Duo 返回的是 Anthropic 托管模型,CLIProxyAPI 会通过 GitLab AI gateway 的 Anthropic 代理转发,并复用现有的 Claude executor 路径。
|
||||
|
||||
如果 GitLab Duo 返回的是 OpenAI 托管模型,CLIProxyAPI 会通过 GitLab AI gateway 的 OpenAI 代理转发,并复用现有的 Codex/OpenAI executor 路径。
|
||||
|
||||
这让 GitLab Duo 的运行时行为更接近内置的 `codex` Provider:
|
||||
|
||||
- Claude 兼容客户端可以通过 `/v1/messages` 使用 GitLab Duo 模型
|
||||
- OpenAI 兼容客户端可以通过 `/v1/chat/completions` 使用 GitLab Duo 模型
|
||||
- OpenAI Responses 客户端可以通过 `/v1/responses` 使用 GitLab Duo 模型
|
||||
|
||||
模型列表不是硬编码的。CLIProxyAPI 会从 GitLab `direct_access` 中读取当前模型元数据,并注册:
|
||||
|
||||
- 一个稳定别名:`gitlab-duo`
|
||||
- GitLab 当前发现到的托管模型名,例如 `claude-sonnet-4-5` 或 `gpt-5-codex`
|
||||
|
||||
## 登录
|
||||
|
||||
OAuth 登录:
|
||||
|
||||
```bash
|
||||
./CLIProxyAPI -gitlab-login
|
||||
```
|
||||
|
||||
PAT 登录:
|
||||
|
||||
```bash
|
||||
./CLIProxyAPI -gitlab-token-login
|
||||
```
|
||||
|
||||
也可以通过环境变量提供输入:
|
||||
|
||||
```bash
|
||||
export GITLAB_BASE_URL=https://gitlab.com
|
||||
export GITLAB_OAUTH_CLIENT_ID=your-client-id
|
||||
export GITLAB_OAUTH_CLIENT_SECRET=your-client-secret
|
||||
export GITLAB_PERSONAL_ACCESS_TOKEN=glpat-...
|
||||
```
|
||||
|
||||
说明:
|
||||
|
||||
- OAuth 方式需要一个 GitLab OAuth application。
|
||||
- PAT 登录需要一个能够调用 GitLab Duo 相关 API 的 personal access token。实践上,`api` scope 是最稳妥的基线。
|
||||
- 自建 GitLab 实例可以通过 `GITLAB_BASE_URL` 接入。
|
||||
|
||||
## 如何使用模型
|
||||
|
||||
登录完成后,正常启动 CLIProxyAPI,并让客户端连接到本地代理。
|
||||
|
||||
你可以选择:
|
||||
|
||||
- `gitlab-duo`,始终使用该账号当前的 Duo 托管模型
|
||||
- GitLab 当前发现到的 provider 模型名,如果你想显式固定模型
|
||||
|
||||
示例:
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8080/v1/models
|
||||
```
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8080/v1/chat/completions \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model": "gitlab-duo",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Write a Go HTTP middleware for request IDs."}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
如果该 GitLab 账号当前绑定的是 Anthropic 模型,Claude 兼容客户端可以通过 Claude handler 路径直接使用它。如果当前绑定的是 OpenAI/Codex 模型,OpenAI 兼容客户端可以通过 `/v1/chat/completions` 或 `/v1/responses` 使用它。
|
||||
|
||||
## 模型如何保持最新
|
||||
|
||||
CLIProxyAPI 不内置固定的 GitLab Duo 模型清单。
|
||||
|
||||
它会刷新 GitLab `direct_access` 元数据,并使用返回的 `model_details` 以及可能存在的模型列表字段,让本地 registry 尽量与 GitLab 当前分配的托管模型保持一致。
|
||||
|
||||
这比硬编码模型名更符合 GitLab 当前公开 API 的实际契约。
|
||||
|
||||
## 当前覆盖范围
|
||||
|
||||
GitLab Duo Provider 目前已经具备:
|
||||
|
||||
- OAuth 和 PAT 登录流程
|
||||
- Duo gateway 凭据的运行时刷新
|
||||
- Anthropic gateway 原生路由
|
||||
- OpenAI/Codex gateway 原生路由
|
||||
- Claude 兼容和 OpenAI 兼容路径的 handler 级 smoke 测试
|
||||
|
||||
当前仍未覆盖:
|
||||
|
||||
- websocket 或 session 级别的完全对齐
|
||||
- GitLab 公开 gateway 契约之外的 IDE 专有能力
|
||||
|
||||
## 参考资料
|
||||
|
||||
- GitLab Code Suggestions API: https://docs.gitlab.com/api/code_suggestions/
|
||||
- GitLab Agent Assistant 与 managed credentials: https://docs.gitlab.com/user/duo_agent_platform/agent_assistant/
|
||||
- GitLab Duo 模型选择: https://docs.gitlab.com/user/gitlab_duo/model_selection/
|
||||
@@ -52,11 +52,11 @@ func init() {
|
||||
sdktr.Register(fOpenAI, fMyProv,
|
||||
func(model string, raw []byte, stream bool) []byte { return raw },
|
||||
sdktr.ResponseTransform{
|
||||
Stream: func(ctx context.Context, model string, originalReq, translatedReq, raw []byte, param *any) []string {
|
||||
return []string{string(raw)}
|
||||
Stream: func(ctx context.Context, model string, originalReq, translatedReq, raw []byte, param *any) [][]byte {
|
||||
return [][]byte{raw}
|
||||
},
|
||||
NonStream: func(ctx context.Context, model string, originalReq, translatedReq, raw []byte, param *any) string {
|
||||
return string(raw)
|
||||
NonStream: func(ctx context.Context, model string, originalReq, translatedReq, raw []byte, param *any) []byte {
|
||||
return raw
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
278
gitlab-duo-codex-parity-plan.md
Normal file
278
gitlab-duo-codex-parity-plan.md
Normal file
@@ -0,0 +1,278 @@
|
||||
# Plan: GitLab Duo Codex Parity
|
||||
|
||||
**Generated**: 2026-03-10
|
||||
**Estimated Complexity**: High
|
||||
|
||||
## Overview
|
||||
Bring GitLab Duo support from the current "auth + basic executor" stage to the same practical level as `codex` inside `CLIProxyAPI`: a user logs in once, points external clients such as Claude Code at `CLIProxyAPI`, selects GitLab Duo-backed models, and gets stable streaming, multi-turn behavior, tool calling compatibility, and predictable model routing without manual provider-specific workarounds.
|
||||
|
||||
The core architectural shift is to stop treating GitLab Duo as only two REST wrappers (`/api/v4/chat/completions` and `/api/v4/code_suggestions/completions`) and instead use GitLab's `direct_access` contract as the primary runtime entrypoint wherever possible. Official GitLab docs confirm that `direct_access` returns AI gateway connection details, headers, token, and expiry; that contract is the closest path to codex-like provider behavior.
|
||||
|
||||
## Prerequisites
|
||||
- Official GitLab Duo API references confirmed during implementation:
|
||||
- `POST /api/v4/code_suggestions/direct_access`
|
||||
- `POST /api/v4/code_suggestions/completions`
|
||||
- `POST /api/v4/chat/completions`
|
||||
- Access to at least one real GitLab Duo account for manual verification.
|
||||
- One downstream client target for acceptance testing:
|
||||
- Claude Code against Claude-compatible endpoint
|
||||
- OpenAI-compatible client against `/v1/chat/completions` and `/v1/responses`
|
||||
- Existing PR branch as starting point:
|
||||
- `feat/gitlab-duo-auth`
|
||||
- PR [#2028](https://github.com/router-for-me/CLIProxyAPI/pull/2028)
|
||||
|
||||
## Definition Of Done
|
||||
- GitLab Duo models can be used via `CLIProxyAPI` from the same client surfaces that already work for `codex`.
|
||||
- Upstream streaming is real passthrough or faithful chunked forwarding, not synthetic whole-response replay.
|
||||
- Tool/function calling survives translation layers without dropping fields or corrupting names.
|
||||
- Multi-turn and session semantics are stable across `chat/completions`, `responses`, and Claude-compatible routes.
|
||||
- Model exposure stays current from GitLab metadata or gateway discovery without hardcoded stale model tables.
|
||||
- `go test ./...` stays green and at least one real manual end-to-end client flow is documented.
|
||||
|
||||
## Sprint 1: Contract And Gap Closure
|
||||
**Goal**: Replace assumptions with a hard compatibility contract between current `codex` behavior and what GitLab Duo can actually support.
|
||||
|
||||
**Demo/Validation**:
|
||||
- Written matrix showing `codex` features vs current GitLab Duo behavior.
|
||||
- One checked-in developer note or test fixture for real GitLab Duo payload examples.
|
||||
|
||||
### Task 1.1: Freeze Codex Parity Checklist
|
||||
- **Location**: [internal/runtime/executor/codex_executor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/codex_executor.go), [internal/runtime/executor/codex_websockets_executor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/codex_websockets_executor.go), [sdk/api/handlers/openai/openai_responses_handlers.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/sdk/api/handlers/openai/openai_responses_handlers.go), [sdk/api/handlers/openai/openai_responses_websocket.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/sdk/api/handlers/openai/openai_responses_websocket.go)
|
||||
- **Description**: Produce a concrete feature matrix for `codex`: HTTP execute, SSE execute, `/v1/responses`, websocket downstream path, tool calling, request IDs, session close semantics, and model registration behavior.
|
||||
- **Dependencies**: None
|
||||
- **Acceptance Criteria**:
|
||||
- A checklist exists in repo docs or issue notes.
|
||||
- Each capability is marked `required`, `optional`, or `not possible` for GitLab Duo.
|
||||
- **Validation**:
|
||||
- Review against current `codex` code paths.
|
||||
|
||||
### Task 1.2: Lock GitLab Duo Runtime Contract
|
||||
- **Location**: [internal/auth/gitlab/gitlab.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/auth/gitlab/gitlab.go), [internal/runtime/executor/gitlab_executor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/gitlab_executor.go)
|
||||
- **Description**: Validate the exact upstream contract we can rely on:
|
||||
- `direct_access` fields and refresh cadence
|
||||
- whether AI gateway path is usable directly
|
||||
- when `chat/completions` is available vs when fallback is required
|
||||
- what streaming shape is returned by `code_suggestions/completions?stream=true`
|
||||
- **Dependencies**: Task 1.1
|
||||
- **Acceptance Criteria**:
|
||||
- GitLab transport decision is explicit: `gateway-first`, `REST-first`, or `hybrid`.
|
||||
- Unknown areas are isolated behind feature flags, not spread across executor logic.
|
||||
- **Validation**:
|
||||
- Official docs + captured real responses from a Duo account.
|
||||
|
||||
### Task 1.3: Define Client-Facing Compatibility Targets
|
||||
- **Location**: [README.md](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/README.md), [gitlab-duo-codex-parity-plan.md](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/gitlab-duo-codex-parity-plan.md)
|
||||
- **Description**: Define exactly which external flows must work to call GitLab Duo support "like codex".
|
||||
- **Dependencies**: Task 1.2
|
||||
- **Acceptance Criteria**:
|
||||
- Required surfaces are listed:
|
||||
- Claude-compatible route
|
||||
- OpenAI `chat/completions`
|
||||
- OpenAI `responses`
|
||||
- optional downstream websocket path
|
||||
- Non-goals are explicit if GitLab upstream cannot support them.
|
||||
- **Validation**:
|
||||
- Maintainer review of stated scope.
|
||||
|
||||
## Sprint 2: Primary Transport Parity
|
||||
**Goal**: Move GitLab Duo execution onto a transport that supports codex-like runtime behavior.
|
||||
|
||||
**Demo/Validation**:
|
||||
- A GitLab Duo model works over real streaming through `/v1/chat/completions`.
|
||||
- No synthetic "collect full body then fake stream" path remains on the primary flow.
|
||||
|
||||
### Task 2.1: Refactor GitLab Executor Into Strategy Layers
|
||||
- **Location**: [internal/runtime/executor/gitlab_executor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/gitlab_executor.go)
|
||||
- **Description**: Split current executor into explicit strategies:
|
||||
- auth refresh/direct access refresh
|
||||
- gateway transport
|
||||
- GitLab REST fallback transport
|
||||
- downstream translation helpers
|
||||
- **Dependencies**: Sprint 1
|
||||
- **Acceptance Criteria**:
|
||||
- Executor no longer mixes discovery, refresh, fallback selection, and response synthesis in one path.
|
||||
- Transport choice is testable in isolation.
|
||||
- **Validation**:
|
||||
- Unit tests for strategy selection and fallback boundaries.
|
||||
|
||||
### Task 2.2: Implement Real Streaming Path
|
||||
- **Location**: [internal/runtime/executor/gitlab_executor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/gitlab_executor.go), [internal/runtime/executor/gitlab_executor_test.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/gitlab_executor_test.go)
|
||||
- **Description**: Replace synthetic streaming with true upstream incremental forwarding:
|
||||
- use gateway stream if available
|
||||
- otherwise consume GitLab Code Suggestions streaming response and map chunks incrementally
|
||||
- **Dependencies**: Task 2.1
|
||||
- **Acceptance Criteria**:
|
||||
- `ExecuteStream` emits chunks before upstream completion.
|
||||
- error handling preserves status and early failure semantics.
|
||||
- **Validation**:
|
||||
- tests with chunked upstream server
|
||||
- manual curl check against `/v1/chat/completions` with `stream=true`
|
||||
|
||||
### Task 2.3: Preserve Upstream Auth And Headers Correctly
|
||||
- **Location**: [internal/runtime/executor/gitlab_executor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/gitlab_executor.go), [internal/auth/gitlab/gitlab.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/auth/gitlab/gitlab.go)
|
||||
- **Description**: Use `direct_access` connection details as first-class transport state:
|
||||
- gateway token
|
||||
- expiry
|
||||
- mandatory forwarded headers
|
||||
- model metadata
|
||||
- **Dependencies**: Task 2.1
|
||||
- **Acceptance Criteria**:
|
||||
- executor stops ignoring gateway headers/token when transport requires them
|
||||
- refresh logic never over-fetches `direct_access`
|
||||
- **Validation**:
|
||||
- tests verifying propagated headers and refresh interval behavior
|
||||
|
||||
## Sprint 3: Request/Response Semantics Parity
|
||||
**Goal**: Make GitLab Duo behave correctly under the same request shapes that current `codex` consumers send.
|
||||
|
||||
**Demo/Validation**:
|
||||
- OpenAI and Claude-compatible clients can do non-streaming and streaming conversations without losing structure.
|
||||
|
||||
### Task 3.1: Normalize Multi-Turn Message Mapping
|
||||
- **Location**: [internal/runtime/executor/gitlab_executor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/gitlab_executor.go), [sdk/translator](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/sdk/translator)
|
||||
- **Description**: Replace the current "flatten prompt into one instruction" behavior with stable multi-turn mapping:
|
||||
- preserve system context
|
||||
- preserve user/assistant ordering
|
||||
- maintain bounded context truncation
|
||||
- **Dependencies**: Sprint 2
|
||||
- **Acceptance Criteria**:
|
||||
- multi-turn requests are not collapsed into a lossy single string unless fallback mode explicitly requires it
|
||||
- truncation policy is deterministic and tested
|
||||
- **Validation**:
|
||||
- golden tests for request mapping
|
||||
|
||||
### Task 3.2: Tool Calling Compatibility Layer
|
||||
- **Location**: [internal/runtime/executor/gitlab_executor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/gitlab_executor.go), [sdk/api/handlers/openai/openai_responses_handlers.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/sdk/api/handlers/openai/openai_responses_handlers.go)
|
||||
- **Description**: Decide and implement one of two paths:
|
||||
- native pass-through if GitLab gateway supports tool/function structures
|
||||
- strict downgrade path with explicit unsupported errors instead of silent field loss
|
||||
- **Dependencies**: Task 3.1
|
||||
- **Acceptance Criteria**:
|
||||
- tool-related fields are either preserved correctly or rejected explicitly
|
||||
- no silent corruption of tool names, tool calls, or tool results
|
||||
- **Validation**:
|
||||
- table-driven tests for tool payloads
|
||||
- one manual client scenario using tools
|
||||
|
||||
### Task 3.3: Token Counting And Usage Reporting Fidelity
|
||||
- **Location**: [internal/runtime/executor/gitlab_executor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/gitlab_executor.go), [internal/runtime/executor/usage_helpers.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/usage_helpers.go)
|
||||
- **Description**: Improve token/usage reporting so GitLab models behave like first-class providers in logs and scheduling.
|
||||
- **Dependencies**: Sprint 2
|
||||
- **Acceptance Criteria**:
|
||||
- `CountTokens` uses the closest supported estimation path
|
||||
- usage logging distinguishes prompt vs completion when possible
|
||||
- **Validation**:
|
||||
- unit tests for token estimation outputs
|
||||
|
||||
## Sprint 4: Responses And Session Parity
|
||||
**Goal**: Reach codex-level support for OpenAI Responses clients and long-lived sessions where GitLab upstream permits it.
|
||||
|
||||
**Demo/Validation**:
|
||||
- `/v1/responses` works with GitLab Duo in a realistic client flow.
|
||||
- If websocket parity is not possible, the code explicitly declines it and keeps HTTP paths stable.
|
||||
|
||||
### Task 4.1: Make GitLab Compatible With `/v1/responses`
|
||||
- **Location**: [sdk/api/handlers/openai/openai_responses_handlers.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/sdk/api/handlers/openai/openai_responses_handlers.go), [internal/runtime/executor/gitlab_executor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/gitlab_executor.go)
|
||||
- **Description**: Ensure GitLab transport can safely back the Responses API path, including compact responses if applicable.
|
||||
- **Dependencies**: Sprint 3
|
||||
- **Acceptance Criteria**:
|
||||
- GitLab Duo can be selected behind `/v1/responses`
|
||||
- response IDs and follow-up semantics are defined
|
||||
- **Validation**:
|
||||
- handler tests analogous to codex/openai responses tests
|
||||
|
||||
### Task 4.2: Evaluate Downstream Websocket Parity
|
||||
- **Location**: [sdk/api/handlers/openai/openai_responses_websocket.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/sdk/api/handlers/openai/openai_responses_websocket.go), [internal/runtime/executor/gitlab_executor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/gitlab_executor.go)
|
||||
- **Description**: Decide whether GitLab Duo can support downstream websocket sessions like codex:
|
||||
- if yes, add session-aware execution path
|
||||
- if no, mark GitLab auth as websocket-ineligible and keep HTTP routes first-class
|
||||
- **Dependencies**: Task 4.1
|
||||
- **Acceptance Criteria**:
|
||||
- websocket behavior is explicit, not accidental
|
||||
- no route claims websocket support when the upstream cannot honor it
|
||||
- **Validation**:
|
||||
- websocket handler tests or explicit capability tests
|
||||
|
||||
### Task 4.3: Add Session Cleanup And Failure Recovery Semantics
|
||||
- **Location**: [internal/runtime/executor/gitlab_executor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/gitlab_executor.go), [sdk/cliproxy/auth/conductor.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/sdk/cliproxy/auth/conductor.go)
|
||||
- **Description**: Add codex-like session cleanup, retry boundaries, and model suspension/resume behavior for GitLab failures and quota events.
|
||||
- **Dependencies**: Sprint 2
|
||||
- **Acceptance Criteria**:
|
||||
- auth/model cooldown behavior is predictable on GitLab 4xx/5xx/quota responses
|
||||
- executor cleans up per-session resources if any are introduced
|
||||
- **Validation**:
|
||||
- tests for quota and retry behavior
|
||||
|
||||
## Sprint 5: Client UX, Model UX, And Manual E2E
|
||||
**Goal**: Make GitLab Duo feel like a normal built-in provider to operators and downstream clients.
|
||||
|
||||
**Demo/Validation**:
|
||||
- A documented setup exists for "login once, point Claude Code at CLIProxyAPI, use GitLab Duo-backed model".
|
||||
|
||||
### Task 5.1: Model Alias And Provider UX Cleanup
|
||||
- **Location**: [sdk/cliproxy/service.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/sdk/cliproxy/service.go), [README.md](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/README.md)
|
||||
- **Description**: Normalize what users see:
|
||||
- stable alias such as `gitlab-duo`
|
||||
- discovered upstream model names
|
||||
- optional prefix behavior
|
||||
- account labels that clearly distinguish OAuth vs PAT
|
||||
- **Dependencies**: Sprint 3
|
||||
- **Acceptance Criteria**:
|
||||
- users can select a stable GitLab alias even when upstream model changes
|
||||
- dynamic model discovery does not cause confusing model churn
|
||||
- **Validation**:
|
||||
- registry tests and manual `/v1/models` inspection
|
||||
|
||||
### Task 5.2: Add Real End-To-End Acceptance Tests
|
||||
- **Location**: [internal/runtime/executor/gitlab_executor_test.go](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/internal/runtime/executor/gitlab_executor_test.go), [sdk/api/handlers/openai](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/sdk/api/handlers/openai)
|
||||
- **Description**: Add higher-level tests covering the actual proxy surfaces:
|
||||
- OpenAI `chat/completions`
|
||||
- OpenAI `responses`
|
||||
- Claude-compatible request path if GitLab is routed there
|
||||
- **Dependencies**: Sprint 4
|
||||
- **Acceptance Criteria**:
|
||||
- tests fail if streaming regresses into synthetic buffering again
|
||||
- tests cover at least one tool-related request and one multi-turn request
|
||||
- **Validation**:
|
||||
- `go test ./...`
|
||||
|
||||
### Task 5.3: Publish Operator Documentation
|
||||
- **Location**: [README.md](/home/luxvtz/projects/cliproxyapi/CLIProxyAPI/README.md)
|
||||
- **Description**: Document:
|
||||
- OAuth setup requirements
|
||||
- PAT requirements
|
||||
- current capability matrix
|
||||
- known limitations if websocket/tool parity is partial
|
||||
- **Dependencies**: Sprint 5.1
|
||||
- **Acceptance Criteria**:
|
||||
- setup instructions are enough for a new user to reproduce the GitLab Duo flow
|
||||
- limitations are explicit
|
||||
- **Validation**:
|
||||
- dry-run docs review from a clean environment
|
||||
|
||||
## Testing Strategy
|
||||
- Keep `go test ./...` green after every committable task.
|
||||
- Add table-driven tests first for request mapping, refresh behavior, and dynamic model registration.
|
||||
- Add transport tests with `httptest.Server` for:
|
||||
- real chunked streaming
|
||||
- header propagation from `direct_access`
|
||||
- upstream fallback rules
|
||||
- Add at least one manual acceptance checklist:
|
||||
- login via OAuth
|
||||
- login via PAT
|
||||
- list models
|
||||
- run one streaming prompt via OpenAI route
|
||||
- run one prompt from the target downstream client
|
||||
|
||||
## Potential Risks & Gotchas
|
||||
- GitLab public docs expose `direct_access`, but do not fully document every possible AI gateway path. We should isolate any empirically discovered gateway assumptions behind one transport layer and feature flags.
|
||||
- `chat/completions` availability differs by GitLab offering and version. The executor must not assume it always exists.
|
||||
- Code Suggestions is completion-oriented; lossy mapping from rich chat/tool payloads will make GitLab Duo feel worse than codex unless explicitly handled.
|
||||
- Synthetic streaming is not good enough for codex parity and will cause regressions in interactive clients.
|
||||
- Dynamic model discovery can create unstable UX if the stable alias and discovered model IDs are not separated cleanly.
|
||||
- PAT auth may validate successfully while still lacking effective Duo permissions. Error reporting must surface this explicitly.
|
||||
|
||||
## Rollback Plan
|
||||
- Keep the current basic GitLab executor behind a fallback mode until the new transport path is stable.
|
||||
- If parity work destabilizes existing providers, revert only GitLab-specific executor changes and leave auth support intact.
|
||||
- Preserve the stable `gitlab-duo` alias so rollback does not break client configuration.
|
||||
2
go.mod
2
go.mod
@@ -91,8 +91,8 @@ require (
|
||||
github.com/tidwall/pretty v1.2.0 // indirect
|
||||
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
|
||||
github.com/ugorji/go/codec v1.2.12 // indirect
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
|
||||
github.com/x448/float16 v0.8.4 // indirect
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
|
||||
golang.org/x/arch v0.8.0 // indirect
|
||||
golang.org/x/sys v0.38.0 // indirect
|
||||
golang.org/x/text v0.31.0 // indirect
|
||||
|
||||
@@ -6,7 +6,6 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
@@ -14,13 +13,12 @@ import (
|
||||
|
||||
"github.com/fxamacker/cbor/v2"
|
||||
"github.com/gin-gonic/gin"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/net/proxy"
|
||||
"golang.org/x/oauth2"
|
||||
"golang.org/x/oauth2/google"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/geminicli"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/proxyutil"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/oauth2"
|
||||
"golang.org/x/oauth2/google"
|
||||
)
|
||||
|
||||
const defaultAPICallTimeout = 60 * time.Second
|
||||
@@ -725,47 +723,12 @@ func (h *Handler) apiCallTransport(auth *coreauth.Auth) http.RoundTripper {
|
||||
}
|
||||
|
||||
func buildProxyTransport(proxyStr string) *http.Transport {
|
||||
proxyStr = strings.TrimSpace(proxyStr)
|
||||
if proxyStr == "" {
|
||||
transport, _, errBuild := proxyutil.BuildHTTPTransport(proxyStr)
|
||||
if errBuild != nil {
|
||||
log.WithError(errBuild).Debug("build proxy transport failed")
|
||||
return nil
|
||||
}
|
||||
|
||||
proxyURL, errParse := url.Parse(proxyStr)
|
||||
if errParse != nil {
|
||||
log.WithError(errParse).Debug("parse proxy URL failed")
|
||||
return nil
|
||||
}
|
||||
if proxyURL.Scheme == "" || proxyURL.Host == "" {
|
||||
log.Debug("proxy URL missing scheme/host")
|
||||
return nil
|
||||
}
|
||||
|
||||
if proxyURL.Scheme == "socks5" {
|
||||
var proxyAuth *proxy.Auth
|
||||
if proxyURL.User != nil {
|
||||
username := proxyURL.User.Username()
|
||||
password, _ := proxyURL.User.Password()
|
||||
proxyAuth = &proxy.Auth{User: username, Password: password}
|
||||
}
|
||||
dialer, errSOCKS5 := proxy.SOCKS5("tcp", proxyURL.Host, proxyAuth, proxy.Direct)
|
||||
if errSOCKS5 != nil {
|
||||
log.WithError(errSOCKS5).Debug("create SOCKS5 dialer failed")
|
||||
return nil
|
||||
}
|
||||
return &http.Transport{
|
||||
Proxy: nil,
|
||||
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||
return dialer.Dial(network, addr)
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
if proxyURL.Scheme == "http" || proxyURL.Scheme == "https" {
|
||||
return &http.Transport{Proxy: http.ProxyURL(proxyURL)}
|
||||
}
|
||||
|
||||
log.Debugf("unsupported proxy scheme: %s", proxyURL.Scheme)
|
||||
return nil
|
||||
return transport
|
||||
}
|
||||
|
||||
// headerContainsValue checks whether a header map contains a target value (case-insensitive key and value).
|
||||
|
||||
@@ -2,172 +2,112 @@ package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
|
||||
)
|
||||
|
||||
type memoryAuthStore struct {
|
||||
mu sync.Mutex
|
||||
items map[string]*coreauth.Auth
|
||||
}
|
||||
func TestAPICallTransportDirectBypassesGlobalProxy(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
func (s *memoryAuthStore) List(ctx context.Context) ([]*coreauth.Auth, error) {
|
||||
_ = ctx
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := make([]*coreauth.Auth, 0, len(s.items))
|
||||
for _, a := range s.items {
|
||||
out = append(out, a.Clone())
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *memoryAuthStore) Save(ctx context.Context, auth *coreauth.Auth) (string, error) {
|
||||
_ = ctx
|
||||
if auth == nil {
|
||||
return "", nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
if s.items == nil {
|
||||
s.items = make(map[string]*coreauth.Auth)
|
||||
}
|
||||
s.items[auth.ID] = auth.Clone()
|
||||
s.mu.Unlock()
|
||||
return auth.ID, nil
|
||||
}
|
||||
|
||||
func (s *memoryAuthStore) Delete(ctx context.Context, id string) error {
|
||||
_ = ctx
|
||||
s.mu.Lock()
|
||||
delete(s.items, id)
|
||||
s.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestResolveTokenForAuth_Antigravity_RefreshesExpiredToken(t *testing.T) {
|
||||
var callCount int
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
callCount++
|
||||
if r.Method != http.MethodPost {
|
||||
t.Fatalf("expected POST, got %s", r.Method)
|
||||
}
|
||||
if ct := r.Header.Get("Content-Type"); !strings.HasPrefix(ct, "application/x-www-form-urlencoded") {
|
||||
t.Fatalf("unexpected content-type: %s", ct)
|
||||
}
|
||||
bodyBytes, _ := io.ReadAll(r.Body)
|
||||
_ = r.Body.Close()
|
||||
values, err := url.ParseQuery(string(bodyBytes))
|
||||
if err != nil {
|
||||
t.Fatalf("parse form: %v", err)
|
||||
}
|
||||
if values.Get("grant_type") != "refresh_token" {
|
||||
t.Fatalf("unexpected grant_type: %s", values.Get("grant_type"))
|
||||
}
|
||||
if values.Get("refresh_token") != "rt" {
|
||||
t.Fatalf("unexpected refresh_token: %s", values.Get("refresh_token"))
|
||||
}
|
||||
if values.Get("client_id") != antigravityOAuthClientID {
|
||||
t.Fatalf("unexpected client_id: %s", values.Get("client_id"))
|
||||
}
|
||||
if values.Get("client_secret") != antigravityOAuthClientSecret {
|
||||
t.Fatalf("unexpected client_secret")
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"access_token": "new-token",
|
||||
"refresh_token": "rt2",
|
||||
"expires_in": int64(3600),
|
||||
"token_type": "Bearer",
|
||||
})
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
originalURL := antigravityOAuthTokenURL
|
||||
antigravityOAuthTokenURL = srv.URL
|
||||
t.Cleanup(func() { antigravityOAuthTokenURL = originalURL })
|
||||
|
||||
store := &memoryAuthStore{}
|
||||
manager := coreauth.NewManager(store, nil, nil)
|
||||
|
||||
auth := &coreauth.Auth{
|
||||
ID: "antigravity-test.json",
|
||||
FileName: "antigravity-test.json",
|
||||
Provider: "antigravity",
|
||||
Metadata: map[string]any{
|
||||
"type": "antigravity",
|
||||
"access_token": "old-token",
|
||||
"refresh_token": "rt",
|
||||
"expires_in": int64(3600),
|
||||
"timestamp": time.Now().Add(-2 * time.Hour).UnixMilli(),
|
||||
"expired": time.Now().Add(-1 * time.Hour).Format(time.RFC3339),
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
SDKConfig: sdkconfig.SDKConfig{ProxyURL: "http://global-proxy.example.com:8080"},
|
||||
},
|
||||
}
|
||||
if _, err := manager.Register(context.Background(), auth); err != nil {
|
||||
t.Fatalf("register auth: %v", err)
|
||||
|
||||
transport := h.apiCallTransport(&coreauth.Auth{ProxyURL: "direct"})
|
||||
httpTransport, ok := transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("transport type = %T, want *http.Transport", transport)
|
||||
}
|
||||
if httpTransport.Proxy != nil {
|
||||
t.Fatal("expected direct transport to disable proxy function")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPICallTransportInvalidAuthFallsBackToGlobalProxy(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := &Handler{
|
||||
cfg: &config.Config{
|
||||
SDKConfig: sdkconfig.SDKConfig{ProxyURL: "http://global-proxy.example.com:8080"},
|
||||
},
|
||||
}
|
||||
|
||||
transport := h.apiCallTransport(&coreauth.Auth{ProxyURL: "bad-value"})
|
||||
httpTransport, ok := transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("transport type = %T, want *http.Transport", transport)
|
||||
}
|
||||
|
||||
req, errRequest := http.NewRequest(http.MethodGet, "https://example.com", nil)
|
||||
if errRequest != nil {
|
||||
t.Fatalf("http.NewRequest returned error: %v", errRequest)
|
||||
}
|
||||
|
||||
proxyURL, errProxy := httpTransport.Proxy(req)
|
||||
if errProxy != nil {
|
||||
t.Fatalf("httpTransport.Proxy returned error: %v", errProxy)
|
||||
}
|
||||
if proxyURL == nil || proxyURL.String() != "http://global-proxy.example.com:8080" {
|
||||
t.Fatalf("proxy URL = %v, want http://global-proxy.example.com:8080", proxyURL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthByIndexDistinguishesSharedAPIKeysAcrossProviders(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
geminiAuth := &coreauth.Auth{
|
||||
ID: "gemini:apikey:123",
|
||||
Provider: "gemini",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "shared-key",
|
||||
},
|
||||
}
|
||||
compatAuth := &coreauth.Auth{
|
||||
ID: "openai-compatibility:bohe:456",
|
||||
Provider: "bohe",
|
||||
Label: "bohe",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "shared-key",
|
||||
"compat_name": "bohe",
|
||||
"provider_key": "bohe",
|
||||
},
|
||||
}
|
||||
|
||||
if _, errRegister := manager.Register(context.Background(), geminiAuth); errRegister != nil {
|
||||
t.Fatalf("register gemini auth: %v", errRegister)
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), compatAuth); errRegister != nil {
|
||||
t.Fatalf("register compat auth: %v", errRegister)
|
||||
}
|
||||
|
||||
geminiIndex := geminiAuth.EnsureIndex()
|
||||
compatIndex := compatAuth.EnsureIndex()
|
||||
if geminiIndex == compatIndex {
|
||||
t.Fatalf("shared api key produced duplicate auth_index %q", geminiIndex)
|
||||
}
|
||||
|
||||
h := &Handler{authManager: manager}
|
||||
token, err := h.resolveTokenForAuth(context.Background(), auth)
|
||||
if err != nil {
|
||||
t.Fatalf("resolveTokenForAuth: %v", err)
|
||||
|
||||
gotGemini := h.authByIndex(geminiIndex)
|
||||
if gotGemini == nil {
|
||||
t.Fatal("expected gemini auth by index")
|
||||
}
|
||||
if token != "new-token" {
|
||||
t.Fatalf("expected refreshed token, got %q", token)
|
||||
}
|
||||
if callCount != 1 {
|
||||
t.Fatalf("expected 1 refresh call, got %d", callCount)
|
||||
if gotGemini.ID != geminiAuth.ID {
|
||||
t.Fatalf("authByIndex(gemini) returned %q, want %q", gotGemini.ID, geminiAuth.ID)
|
||||
}
|
||||
|
||||
updated, ok := manager.GetByID(auth.ID)
|
||||
if !ok || updated == nil {
|
||||
t.Fatalf("expected auth in manager after update")
|
||||
gotCompat := h.authByIndex(compatIndex)
|
||||
if gotCompat == nil {
|
||||
t.Fatal("expected compat auth by index")
|
||||
}
|
||||
if got := tokenValueFromMetadata(updated.Metadata); got != "new-token" {
|
||||
t.Fatalf("expected manager metadata updated, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveTokenForAuth_Antigravity_SkipsRefreshWhenTokenValid(t *testing.T) {
|
||||
var callCount int
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
callCount++
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
originalURL := antigravityOAuthTokenURL
|
||||
antigravityOAuthTokenURL = srv.URL
|
||||
t.Cleanup(func() { antigravityOAuthTokenURL = originalURL })
|
||||
|
||||
auth := &coreauth.Auth{
|
||||
ID: "antigravity-valid.json",
|
||||
FileName: "antigravity-valid.json",
|
||||
Provider: "antigravity",
|
||||
Metadata: map[string]any{
|
||||
"type": "antigravity",
|
||||
"access_token": "ok-token",
|
||||
"expired": time.Now().Add(30 * time.Minute).Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
h := &Handler{}
|
||||
token, err := h.resolveTokenForAuth(context.Background(), auth)
|
||||
if err != nil {
|
||||
t.Fatalf("resolveTokenForAuth: %v", err)
|
||||
}
|
||||
if token != "ok-token" {
|
||||
t.Fatalf("expected existing token, got %q", token)
|
||||
}
|
||||
if callCount != 0 {
|
||||
t.Fatalf("expected no refresh calls, got %d", callCount)
|
||||
if gotCompat.ID != compatAuth.ID {
|
||||
t.Fatalf("authByIndex(compat) returned %q, want %q", gotCompat.ID, compatAuth.ID)
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
197
internal/api/handlers/management/auth_files_batch_test.go
Normal file
197
internal/api/handlers/management/auth_files_batch_test.go
Normal file
@@ -0,0 +1,197 @@
|
||||
package management
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestUploadAuthFile_BatchMultipart(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
authDir := t.TempDir()
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
|
||||
files := []struct {
|
||||
name string
|
||||
content string
|
||||
}{
|
||||
{name: "alpha.json", content: `{"type":"codex","email":"alpha@example.com"}`},
|
||||
{name: "beta.json", content: `{"type":"claude","email":"beta@example.com"}`},
|
||||
}
|
||||
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
for _, file := range files {
|
||||
part, err := writer.CreateFormFile("file", file.name)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create multipart file: %v", err)
|
||||
}
|
||||
if _, err = part.Write([]byte(file.content)); err != nil {
|
||||
t.Fatalf("failed to write multipart content: %v", err)
|
||||
}
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
t.Fatalf("failed to close multipart writer: %v", err)
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v0/management/auth-files", &body)
|
||||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
ctx.Request = req
|
||||
|
||||
h.UploadAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected upload status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
|
||||
t.Fatalf("failed to decode response: %v", err)
|
||||
}
|
||||
if got, ok := payload["uploaded"].(float64); !ok || int(got) != len(files) {
|
||||
t.Fatalf("expected uploaded=%d, got %#v", len(files), payload["uploaded"])
|
||||
}
|
||||
|
||||
for _, file := range files {
|
||||
fullPath := filepath.Join(authDir, file.name)
|
||||
data, err := os.ReadFile(fullPath)
|
||||
if err != nil {
|
||||
t.Fatalf("expected uploaded file %s to exist: %v", file.name, err)
|
||||
}
|
||||
if string(data) != file.content {
|
||||
t.Fatalf("expected file %s content %q, got %q", file.name, file.content, string(data))
|
||||
}
|
||||
}
|
||||
|
||||
auths := manager.List()
|
||||
if len(auths) != len(files) {
|
||||
t.Fatalf("expected %d auth entries, got %d", len(files), len(auths))
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadAuthFile_BatchMultipart_InvalidJSONDoesNotOverwriteExistingFile(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
authDir := t.TempDir()
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
|
||||
existingName := "alpha.json"
|
||||
existingContent := `{"type":"codex","email":"alpha@example.com"}`
|
||||
if err := os.WriteFile(filepath.Join(authDir, existingName), []byte(existingContent), 0o600); err != nil {
|
||||
t.Fatalf("failed to seed existing auth file: %v", err)
|
||||
}
|
||||
|
||||
files := []struct {
|
||||
name string
|
||||
content string
|
||||
}{
|
||||
{name: existingName, content: `{"type":"codex"`},
|
||||
{name: "beta.json", content: `{"type":"claude","email":"beta@example.com"}`},
|
||||
}
|
||||
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
for _, file := range files {
|
||||
part, err := writer.CreateFormFile("file", file.name)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create multipart file: %v", err)
|
||||
}
|
||||
if _, err = part.Write([]byte(file.content)); err != nil {
|
||||
t.Fatalf("failed to write multipart content: %v", err)
|
||||
}
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
t.Fatalf("failed to close multipart writer: %v", err)
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v0/management/auth-files", &body)
|
||||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
ctx.Request = req
|
||||
|
||||
h.UploadAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusMultiStatus {
|
||||
t.Fatalf("expected upload status %d, got %d with body %s", http.StatusMultiStatus, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(filepath.Join(authDir, existingName))
|
||||
if err != nil {
|
||||
t.Fatalf("expected existing auth file to remain readable: %v", err)
|
||||
}
|
||||
if string(data) != existingContent {
|
||||
t.Fatalf("expected existing auth file to remain %q, got %q", existingContent, string(data))
|
||||
}
|
||||
|
||||
betaData, err := os.ReadFile(filepath.Join(authDir, "beta.json"))
|
||||
if err != nil {
|
||||
t.Fatalf("expected valid auth file to be created: %v", err)
|
||||
}
|
||||
if string(betaData) != files[1].content {
|
||||
t.Fatalf("expected beta auth file content %q, got %q", files[1].content, string(betaData))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteAuthFile_BatchQuery(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
authDir := t.TempDir()
|
||||
files := []string{"alpha.json", "beta.json"}
|
||||
for _, name := range files {
|
||||
if err := os.WriteFile(filepath.Join(authDir, name), []byte(`{"type":"codex"}`), 0o600); err != nil {
|
||||
t.Fatalf("failed to write auth file %s: %v", name, err)
|
||||
}
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
h.tokenStore = &memoryAuthStore{}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(
|
||||
http.MethodDelete,
|
||||
"/v0/management/auth-files?name="+url.QueryEscape(files[0])+"&name="+url.QueryEscape(files[1]),
|
||||
nil,
|
||||
)
|
||||
ctx.Request = req
|
||||
|
||||
h.DeleteAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected delete status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
|
||||
t.Fatalf("failed to decode response: %v", err)
|
||||
}
|
||||
if got, ok := payload["deleted"].(float64); !ok || int(got) != len(files) {
|
||||
t.Fatalf("expected deleted=%d, got %#v", len(files), payload["deleted"])
|
||||
}
|
||||
|
||||
for _, name := range files {
|
||||
if _, err := os.Stat(filepath.Join(authDir, name)); !os.IsNotExist(err) {
|
||||
t.Fatalf("expected auth file %s to be removed, stat err: %v", name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
129
internal/api/handlers/management/auth_files_delete_test.go
Normal file
129
internal/api/handlers/management/auth_files_delete_test.go
Normal file
@@ -0,0 +1,129 @@
|
||||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestDeleteAuthFile_UsesAuthPathFromManager(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tempDir := t.TempDir()
|
||||
authDir := filepath.Join(tempDir, "auth")
|
||||
externalDir := filepath.Join(tempDir, "external")
|
||||
if errMkdirAuth := os.MkdirAll(authDir, 0o700); errMkdirAuth != nil {
|
||||
t.Fatalf("failed to create auth dir: %v", errMkdirAuth)
|
||||
}
|
||||
if errMkdirExternal := os.MkdirAll(externalDir, 0o700); errMkdirExternal != nil {
|
||||
t.Fatalf("failed to create external dir: %v", errMkdirExternal)
|
||||
}
|
||||
|
||||
fileName := "codex-user@example.com-plus.json"
|
||||
shadowPath := filepath.Join(authDir, fileName)
|
||||
realPath := filepath.Join(externalDir, fileName)
|
||||
if errWriteShadow := os.WriteFile(shadowPath, []byte(`{"type":"codex","email":"shadow@example.com"}`), 0o600); errWriteShadow != nil {
|
||||
t.Fatalf("failed to write shadow file: %v", errWriteShadow)
|
||||
}
|
||||
if errWriteReal := os.WriteFile(realPath, []byte(`{"type":"codex","email":"real@example.com"}`), 0o600); errWriteReal != nil {
|
||||
t.Fatalf("failed to write real file: %v", errWriteReal)
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
record := &coreauth.Auth{
|
||||
ID: "legacy/" + fileName,
|
||||
FileName: fileName,
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusError,
|
||||
Unavailable: true,
|
||||
Attributes: map[string]string{
|
||||
"path": realPath,
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"type": "codex",
|
||||
"email": "real@example.com",
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), record); errRegister != nil {
|
||||
t.Fatalf("failed to register auth record: %v", errRegister)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
h.tokenStore = &memoryAuthStore{}
|
||||
|
||||
deleteRec := httptest.NewRecorder()
|
||||
deleteCtx, _ := gin.CreateTestContext(deleteRec)
|
||||
deleteReq := httptest.NewRequest(http.MethodDelete, "/v0/management/auth-files?name="+url.QueryEscape(fileName), nil)
|
||||
deleteCtx.Request = deleteReq
|
||||
h.DeleteAuthFile(deleteCtx)
|
||||
|
||||
if deleteRec.Code != http.StatusOK {
|
||||
t.Fatalf("expected delete status %d, got %d with body %s", http.StatusOK, deleteRec.Code, deleteRec.Body.String())
|
||||
}
|
||||
if _, errStatReal := os.Stat(realPath); !os.IsNotExist(errStatReal) {
|
||||
t.Fatalf("expected managed auth file to be removed, stat err: %v", errStatReal)
|
||||
}
|
||||
if _, errStatShadow := os.Stat(shadowPath); errStatShadow != nil {
|
||||
t.Fatalf("expected shadow auth file to remain, stat err: %v", errStatShadow)
|
||||
}
|
||||
|
||||
listRec := httptest.NewRecorder()
|
||||
listCtx, _ := gin.CreateTestContext(listRec)
|
||||
listReq := httptest.NewRequest(http.MethodGet, "/v0/management/auth-files", nil)
|
||||
listCtx.Request = listReq
|
||||
h.ListAuthFiles(listCtx)
|
||||
|
||||
if listRec.Code != http.StatusOK {
|
||||
t.Fatalf("expected list status %d, got %d with body %s", http.StatusOK, listRec.Code, listRec.Body.String())
|
||||
}
|
||||
var listPayload map[string]any
|
||||
if errUnmarshal := json.Unmarshal(listRec.Body.Bytes(), &listPayload); errUnmarshal != nil {
|
||||
t.Fatalf("failed to decode list payload: %v", errUnmarshal)
|
||||
}
|
||||
filesRaw, ok := listPayload["files"].([]any)
|
||||
if !ok {
|
||||
t.Fatalf("expected files array, payload: %#v", listPayload)
|
||||
}
|
||||
if len(filesRaw) != 0 {
|
||||
t.Fatalf("expected removed auth to be hidden from list, got %d entries", len(filesRaw))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteAuthFile_FallbackToAuthDirPath(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
authDir := t.TempDir()
|
||||
fileName := "fallback-user.json"
|
||||
filePath := filepath.Join(authDir, fileName)
|
||||
if errWrite := os.WriteFile(filePath, []byte(`{"type":"codex"}`), 0o600); errWrite != nil {
|
||||
t.Fatalf("failed to write auth file: %v", errWrite)
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, manager)
|
||||
h.tokenStore = &memoryAuthStore{}
|
||||
|
||||
deleteRec := httptest.NewRecorder()
|
||||
deleteCtx, _ := gin.CreateTestContext(deleteRec)
|
||||
deleteReq := httptest.NewRequest(http.MethodDelete, "/v0/management/auth-files?name="+url.QueryEscape(fileName), nil)
|
||||
deleteCtx.Request = deleteReq
|
||||
h.DeleteAuthFile(deleteCtx)
|
||||
|
||||
if deleteRec.Code != http.StatusOK {
|
||||
t.Fatalf("expected delete status %d, got %d with body %s", http.StatusOK, deleteRec.Code, deleteRec.Body.String())
|
||||
}
|
||||
if _, errStat := os.Stat(filePath); !os.IsNotExist(errStat) {
|
||||
t.Fatalf("expected auth file to be removed from auth dir, stat err: %v", errStat)
|
||||
}
|
||||
}
|
||||
62
internal/api/handlers/management/auth_files_download_test.go
Normal file
62
internal/api/handlers/management/auth_files_download_test.go
Normal file
@@ -0,0 +1,62 @@
|
||||
package management
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
)
|
||||
|
||||
func TestDownloadAuthFile_ReturnsFile(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
authDir := t.TempDir()
|
||||
fileName := "download-user.json"
|
||||
expected := []byte(`{"type":"codex"}`)
|
||||
if err := os.WriteFile(filepath.Join(authDir, fileName), expected, 0o600); err != nil {
|
||||
t.Fatalf("failed to write auth file: %v", err)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files/download?name="+url.QueryEscape(fileName), nil)
|
||||
h.DownloadAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected download status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
if got := rec.Body.Bytes(); string(got) != string(expected) {
|
||||
t.Fatalf("unexpected download content: %q", string(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadAuthFile_RejectsPathSeparators(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, nil)
|
||||
|
||||
for _, name := range []string{
|
||||
"../external/secret.json",
|
||||
`..\\external\\secret.json`,
|
||||
"nested/secret.json",
|
||||
`nested\\secret.json`,
|
||||
} {
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/auth-files/download?name="+url.QueryEscape(name), nil)
|
||||
h.DownloadAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("expected %d for name %q, got %d with body %s", http.StatusBadRequest, name, rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
//go:build windows
|
||||
|
||||
package management
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
)
|
||||
|
||||
func TestDownloadAuthFile_PreventsWindowsSlashTraversal(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tempDir := t.TempDir()
|
||||
authDir := filepath.Join(tempDir, "auth")
|
||||
externalDir := filepath.Join(tempDir, "external")
|
||||
if err := os.MkdirAll(authDir, 0o700); err != nil {
|
||||
t.Fatalf("failed to create auth dir: %v", err)
|
||||
}
|
||||
if err := os.MkdirAll(externalDir, 0o700); err != nil {
|
||||
t.Fatalf("failed to create external dir: %v", err)
|
||||
}
|
||||
|
||||
secretName := "secret.json"
|
||||
secretPath := filepath.Join(externalDir, secretName)
|
||||
if err := os.WriteFile(secretPath, []byte(`{"secret":true}`), 0o600); err != nil {
|
||||
t.Fatalf("failed to write external file: %v", err)
|
||||
}
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(
|
||||
http.MethodGet,
|
||||
"/v0/management/auth-files/download?name="+url.QueryEscape("../external/"+secretName),
|
||||
nil,
|
||||
)
|
||||
h.DownloadAuthFile(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusBadRequest, rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
164
internal/api/handlers/management/auth_files_gitlab_test.go
Normal file
164
internal/api/handlers/management/auth_files_gitlab_test.go
Normal file
@@ -0,0 +1,164 @@
|
||||
package management
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
func TestRequestGitLabPATToken_SavesAuthRecord(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if got := r.Header.Get("Authorization"); got != "Bearer glpat-test-token" {
|
||||
t.Fatalf("authorization header = %q, want Bearer glpat-test-token", got)
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
switch r.URL.Path {
|
||||
case "/api/v4/user":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"id": 42,
|
||||
"username": "gitlab-user",
|
||||
"name": "GitLab User",
|
||||
"email": "gitlab@example.com",
|
||||
})
|
||||
case "/api/v4/personal_access_tokens/self":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"id": 7,
|
||||
"name": "management-center",
|
||||
"scopes": []string{"api", "read_user"},
|
||||
"user_id": 42,
|
||||
})
|
||||
case "/api/v4/code_suggestions/direct_access":
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"base_url": "https://cloud.gitlab.example.com",
|
||||
"token": "gateway-token",
|
||||
"expires_at": 1893456000,
|
||||
"headers": map[string]string{
|
||||
"X-Gitlab-Realm": "saas",
|
||||
},
|
||||
"model_details": map[string]any{
|
||||
"model_provider": "anthropic",
|
||||
"model_name": "claude-sonnet-4-5",
|
||||
},
|
||||
})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
store := &memoryAuthStore{}
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, coreauth.NewManager(nil, nil, nil))
|
||||
h.tokenStore = store
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v0/management/gitlab-auth-url", strings.NewReader(`{"base_url":"`+upstream.URL+`","personal_access_token":"glpat-test-token"}`))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h.RequestGitLabPATToken(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
var resp map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("decode response: %v", err)
|
||||
}
|
||||
if got := resp["status"]; got != "ok" {
|
||||
t.Fatalf("status = %#v, want ok", got)
|
||||
}
|
||||
if got := resp["model_provider"]; got != "anthropic" {
|
||||
t.Fatalf("model_provider = %#v, want anthropic", got)
|
||||
}
|
||||
if got := resp["model_name"]; got != "claude-sonnet-4-5" {
|
||||
t.Fatalf("model_name = %#v, want claude-sonnet-4-5", got)
|
||||
}
|
||||
|
||||
store.mu.Lock()
|
||||
defer store.mu.Unlock()
|
||||
if len(store.items) != 1 {
|
||||
t.Fatalf("expected 1 saved auth record, got %d", len(store.items))
|
||||
}
|
||||
var saved *coreauth.Auth
|
||||
for _, item := range store.items {
|
||||
saved = item
|
||||
}
|
||||
if saved == nil {
|
||||
t.Fatal("expected saved auth record")
|
||||
}
|
||||
if saved.Provider != "gitlab" {
|
||||
t.Fatalf("provider = %q, want gitlab", saved.Provider)
|
||||
}
|
||||
if got := saved.Metadata["auth_kind"]; got != "personal_access_token" {
|
||||
t.Fatalf("auth_kind = %#v, want personal_access_token", got)
|
||||
}
|
||||
if got := saved.Metadata["model_provider"]; got != "anthropic" {
|
||||
t.Fatalf("saved model_provider = %#v, want anthropic", got)
|
||||
}
|
||||
if got := saved.Metadata["duo_gateway_token"]; got != "gateway-token" {
|
||||
t.Fatalf("saved duo_gateway_token = %#v, want gateway-token", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostOAuthCallback_GitLabWritesPendingCallbackFile(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
authDir := t.TempDir()
|
||||
state := "gitlab-state-123"
|
||||
RegisterOAuthSession(state, "gitlab")
|
||||
t.Cleanup(func() { CompleteOAuthSession(state) })
|
||||
|
||||
h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, coreauth.NewManager(nil, nil, nil))
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v0/management/oauth-callback", strings.NewReader(`{"provider":"gitlab","redirect_url":"http://localhost:17171/auth/callback?code=test-code&state=`+state+`"}`))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h.PostOAuthCallback(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
filePath := filepath.Join(authDir, ".oauth-gitlab-"+state+".oauth")
|
||||
data, err := os.ReadFile(filePath)
|
||||
if err != nil {
|
||||
t.Fatalf("read callback file: %v", err)
|
||||
}
|
||||
|
||||
var payload map[string]string
|
||||
if err := json.Unmarshal(data, &payload); err != nil {
|
||||
t.Fatalf("decode callback payload: %v", err)
|
||||
}
|
||||
if got := payload["code"]; got != "test-code" {
|
||||
t.Fatalf("callback code = %q, want test-code", got)
|
||||
}
|
||||
if got := payload["state"]; got != state {
|
||||
t.Fatalf("callback state = %q, want %q", got, state)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeOAuthProvider_GitLab(t *testing.T) {
|
||||
provider, err := NormalizeOAuthProvider("gitlab")
|
||||
if err != nil {
|
||||
t.Fatalf("NormalizeOAuthProvider returned error: %v", err)
|
||||
}
|
||||
if provider != "gitlab" {
|
||||
t.Fatalf("provider = %q, want gitlab", provider)
|
||||
}
|
||||
}
|
||||
@@ -509,19 +509,24 @@ func (h *Handler) PutVertexCompatKeys(c *gin.Context) {
|
||||
}
|
||||
for i := range arr {
|
||||
normalizeVertexCompatKey(&arr[i])
|
||||
if arr[i].APIKey == "" {
|
||||
c.JSON(400, gin.H{"error": fmt.Sprintf("vertex-api-key[%d].api-key is required", i)})
|
||||
return
|
||||
}
|
||||
}
|
||||
h.cfg.VertexCompatAPIKey = arr
|
||||
h.cfg.VertexCompatAPIKey = append([]config.VertexCompatKey(nil), arr...)
|
||||
h.cfg.SanitizeVertexCompatKeys()
|
||||
h.persist(c)
|
||||
}
|
||||
func (h *Handler) PatchVertexCompatKey(c *gin.Context) {
|
||||
type vertexCompatPatch struct {
|
||||
APIKey *string `json:"api-key"`
|
||||
Prefix *string `json:"prefix"`
|
||||
BaseURL *string `json:"base-url"`
|
||||
ProxyURL *string `json:"proxy-url"`
|
||||
Headers *map[string]string `json:"headers"`
|
||||
Models *[]config.VertexCompatModel `json:"models"`
|
||||
APIKey *string `json:"api-key"`
|
||||
Prefix *string `json:"prefix"`
|
||||
BaseURL *string `json:"base-url"`
|
||||
ProxyURL *string `json:"proxy-url"`
|
||||
Headers *map[string]string `json:"headers"`
|
||||
Models *[]config.VertexCompatModel `json:"models"`
|
||||
ExcludedModels *[]string `json:"excluded-models"`
|
||||
}
|
||||
var body struct {
|
||||
Index *int `json:"index"`
|
||||
@@ -585,6 +590,9 @@ func (h *Handler) PatchVertexCompatKey(c *gin.Context) {
|
||||
if body.Value.Models != nil {
|
||||
entry.Models = append([]config.VertexCompatModel(nil), (*body.Value.Models)...)
|
||||
}
|
||||
if body.Value.ExcludedModels != nil {
|
||||
entry.ExcludedModels = config.NormalizeExcludedModels(*body.Value.ExcludedModels)
|
||||
}
|
||||
normalizeVertexCompatKey(&entry)
|
||||
h.cfg.VertexCompatAPIKey[targetIndex] = entry
|
||||
h.cfg.SanitizeVertexCompatKeys()
|
||||
@@ -1029,6 +1037,7 @@ func normalizeVertexCompatKey(entry *config.VertexCompatKey) {
|
||||
entry.BaseURL = strings.TrimSpace(entry.BaseURL)
|
||||
entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
|
||||
entry.Headers = config.NormalizeHeaders(entry.Headers)
|
||||
entry.ExcludedModels = config.NormalizeExcludedModels(entry.ExcludedModels)
|
||||
if len(entry.Models) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -47,6 +47,7 @@ type Handler struct {
|
||||
allowRemoteOverride bool
|
||||
envSecret string
|
||||
logDir string
|
||||
postAuthHook coreauth.PostAuthHook
|
||||
}
|
||||
|
||||
// NewHandler creates a new management handler instance.
|
||||
@@ -128,6 +129,11 @@ func (h *Handler) SetLogDirectory(dir string) {
|
||||
h.logDir = dir
|
||||
}
|
||||
|
||||
// SetPostAuthHook registers a hook to be called after auth record creation but before persistence.
|
||||
func (h *Handler) SetPostAuthHook(hook coreauth.PostAuthHook) {
|
||||
h.postAuthHook = hook
|
||||
}
|
||||
|
||||
// Middleware enforces access control for management endpoints.
|
||||
// All requests (local and remote) require a valid management key.
|
||||
// Additionally, remote access requires allow-remote-management=true.
|
||||
|
||||
@@ -228,6 +228,8 @@ func NormalizeOAuthProvider(provider string) (string, error) {
|
||||
return "anthropic", nil
|
||||
case "codex", "openai":
|
||||
return "codex", nil
|
||||
case "gitlab":
|
||||
return "gitlab", nil
|
||||
case "gemini", "google":
|
||||
return "gemini", nil
|
||||
case "iflow", "i-flow":
|
||||
|
||||
49
internal/api/handlers/management/test_store_test.go
Normal file
49
internal/api/handlers/management/test_store_test.go
Normal file
@@ -0,0 +1,49 @@
|
||||
package management
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
)
|
||||
|
||||
type memoryAuthStore struct {
|
||||
mu sync.Mutex
|
||||
items map[string]*coreauth.Auth
|
||||
}
|
||||
|
||||
func (s *memoryAuthStore) List(_ context.Context) ([]*coreauth.Auth, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
out := make([]*coreauth.Auth, 0, len(s.items))
|
||||
for _, item := range s.items {
|
||||
out = append(out, item)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *memoryAuthStore) Save(_ context.Context, auth *coreauth.Auth) (string, error) {
|
||||
if auth == nil {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.items == nil {
|
||||
s.items = make(map[string]*coreauth.Auth)
|
||||
}
|
||||
s.items[auth.ID] = auth
|
||||
return auth.ID, nil
|
||||
}
|
||||
|
||||
func (s *memoryAuthStore) Delete(_ context.Context, id string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
delete(s.items, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *memoryAuthStore) SetBaseDir(string) {}
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
@@ -77,6 +78,9 @@ func createReverseProxy(upstreamURL string, secretSource SecretSource) (*httputi
|
||||
req.Header.Del("X-Api-Key")
|
||||
req.Header.Del("X-Goog-Api-Key")
|
||||
|
||||
// Remove proxy, client identity, and browser fingerprint headers
|
||||
misc.ScrubProxyAndFingerprintHeaders(req)
|
||||
|
||||
// Remove query-based credentials if they match the authenticated client API key.
|
||||
// This prevents leaking client auth material to the Amp upstream while avoiding
|
||||
// breaking unrelated upstream query parameters.
|
||||
|
||||
@@ -52,6 +52,7 @@ type serverOptionConfig struct {
|
||||
keepAliveEnabled bool
|
||||
keepAliveTimeout time.Duration
|
||||
keepAliveOnTimeout func()
|
||||
postAuthHook auth.PostAuthHook
|
||||
}
|
||||
|
||||
// ServerOption customises HTTP server construction.
|
||||
@@ -59,10 +60,8 @@ type ServerOption func(*serverOptionConfig)
|
||||
|
||||
func defaultRequestLoggerFactory(cfg *config.Config, configPath string) logging.RequestLogger {
|
||||
configDir := filepath.Dir(configPath)
|
||||
if base := util.WritablePath(); base != "" {
|
||||
return logging.NewFileRequestLogger(cfg.RequestLog, filepath.Join(base, "logs"), configDir, cfg.ErrorLogsMaxFiles)
|
||||
}
|
||||
return logging.NewFileRequestLogger(cfg.RequestLog, "logs", configDir, cfg.ErrorLogsMaxFiles)
|
||||
logsDir := logging.ResolveLogDirectory(cfg)
|
||||
return logging.NewFileRequestLogger(cfg.RequestLog, logsDir, configDir, cfg.ErrorLogsMaxFiles)
|
||||
}
|
||||
|
||||
// WithMiddleware appends additional Gin middleware during server construction.
|
||||
@@ -112,6 +111,13 @@ func WithRequestLoggerFactory(factory func(*config.Config, string) logging.Reque
|
||||
}
|
||||
}
|
||||
|
||||
// WithPostAuthHook registers a hook to be called after auth record creation.
|
||||
func WithPostAuthHook(hook auth.PostAuthHook) ServerOption {
|
||||
return func(cfg *serverOptionConfig) {
|
||||
cfg.postAuthHook = hook
|
||||
}
|
||||
}
|
||||
|
||||
// Server represents the main API server.
|
||||
// It encapsulates the Gin engine, HTTP server, handlers, and configuration.
|
||||
type Server struct {
|
||||
@@ -252,7 +258,7 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
|
||||
s.oldConfigYaml, _ = yaml.Marshal(cfg)
|
||||
s.applyAccessConfig(nil, cfg)
|
||||
if authManager != nil {
|
||||
authManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second)
|
||||
authManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second, cfg.MaxRetryCredentials)
|
||||
}
|
||||
managementasset.SetCurrentConfig(cfg)
|
||||
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
|
||||
@@ -263,6 +269,9 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
|
||||
}
|
||||
logDir := logging.ResolveLogDirectory(cfg)
|
||||
s.mgmt.SetLogDirectory(logDir)
|
||||
if optionState.postAuthHook != nil {
|
||||
s.mgmt.SetPostAuthHook(optionState.postAuthHook)
|
||||
}
|
||||
s.localPassword = optionState.localPassword
|
||||
|
||||
// Setup routes
|
||||
@@ -394,6 +403,20 @@ func (s *Server) setupRoutes() {
|
||||
c.String(http.StatusOK, oauthCallbackSuccessHTML)
|
||||
})
|
||||
|
||||
s.engine.GET("/gitlab/callback", func(c *gin.Context) {
|
||||
code := c.Query("code")
|
||||
state := c.Query("state")
|
||||
errStr := c.Query("error")
|
||||
if errStr == "" {
|
||||
errStr = c.Query("error_description")
|
||||
}
|
||||
if state != "" {
|
||||
_, _ = managementHandlers.WriteOAuthCallbackFileForPendingSession(s.cfg.AuthDir, "gitlab", state, code, errStr)
|
||||
}
|
||||
c.Header("Content-Type", "text/html; charset=utf-8")
|
||||
c.String(http.StatusOK, oauthCallbackSuccessHTML)
|
||||
})
|
||||
|
||||
s.engine.GET("/google/callback", func(c *gin.Context) {
|
||||
code := c.Query("code")
|
||||
state := c.Query("state")
|
||||
@@ -649,6 +672,8 @@ func (s *Server) registerManagementRoutes() {
|
||||
|
||||
mgmt.GET("/anthropic-auth-url", s.mgmt.RequestAnthropicToken)
|
||||
mgmt.GET("/codex-auth-url", s.mgmt.RequestCodexToken)
|
||||
mgmt.GET("/gitlab-auth-url", s.mgmt.RequestGitLabToken)
|
||||
mgmt.POST("/gitlab-auth-url", s.mgmt.RequestGitLabPATToken)
|
||||
mgmt.GET("/gemini-cli-auth-url", s.mgmt.RequestGeminiCLIToken)
|
||||
mgmt.GET("/antigravity-auth-url", s.mgmt.RequestAntigravityToken)
|
||||
mgmt.GET("/qwen-auth-url", s.mgmt.RequestQwenToken)
|
||||
@@ -657,6 +682,7 @@ func (s *Server) registerManagementRoutes() {
|
||||
mgmt.GET("/iflow-auth-url", s.mgmt.RequestIFlowToken)
|
||||
mgmt.POST("/iflow-auth-url", s.mgmt.RequestIFlowCookieToken)
|
||||
mgmt.GET("/kiro-auth-url", s.mgmt.RequestKiroToken)
|
||||
mgmt.GET("/cursor-auth-url", s.mgmt.RequestCursorToken)
|
||||
mgmt.GET("/github-auth-url", s.mgmt.RequestGitHubToken)
|
||||
mgmt.POST("/oauth-callback", s.mgmt.PostOAuthCallback)
|
||||
mgmt.GET("/get-auth-status", s.mgmt.GetAuthStatus)
|
||||
@@ -935,7 +961,7 @@ func (s *Server) UpdateClients(cfg *config.Config) {
|
||||
}
|
||||
|
||||
if s.handlers != nil && s.handlers.AuthManager != nil {
|
||||
s.handlers.AuthManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second)
|
||||
s.handlers.AuthManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second, cfg.MaxRetryCredentials)
|
||||
}
|
||||
|
||||
// Update log level dynamically when debug flag changes
|
||||
|
||||
@@ -7,9 +7,11 @@ import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
gin "github.com/gin-gonic/gin"
|
||||
proxyconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
internallogging "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
|
||||
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
|
||||
@@ -109,3 +111,100 @@ func TestAmpProviderModelRoutes(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultRequestLoggerFactory_UsesResolvedLogDirectory(t *testing.T) {
|
||||
t.Setenv("WRITABLE_PATH", "")
|
||||
t.Setenv("writable_path", "")
|
||||
|
||||
originalWD, errGetwd := os.Getwd()
|
||||
if errGetwd != nil {
|
||||
t.Fatalf("failed to get current working directory: %v", errGetwd)
|
||||
}
|
||||
|
||||
tmpDir := t.TempDir()
|
||||
if errChdir := os.Chdir(tmpDir); errChdir != nil {
|
||||
t.Fatalf("failed to switch working directory: %v", errChdir)
|
||||
}
|
||||
defer func() {
|
||||
if errChdirBack := os.Chdir(originalWD); errChdirBack != nil {
|
||||
t.Fatalf("failed to restore working directory: %v", errChdirBack)
|
||||
}
|
||||
}()
|
||||
|
||||
// Force ResolveLogDirectory to fallback to auth-dir/logs by making ./logs not a writable directory.
|
||||
if errWriteFile := os.WriteFile(filepath.Join(tmpDir, "logs"), []byte("not-a-directory"), 0o644); errWriteFile != nil {
|
||||
t.Fatalf("failed to create blocking logs file: %v", errWriteFile)
|
||||
}
|
||||
|
||||
configDir := filepath.Join(tmpDir, "config")
|
||||
if errMkdirConfig := os.MkdirAll(configDir, 0o755); errMkdirConfig != nil {
|
||||
t.Fatalf("failed to create config dir: %v", errMkdirConfig)
|
||||
}
|
||||
configPath := filepath.Join(configDir, "config.yaml")
|
||||
|
||||
authDir := filepath.Join(tmpDir, "auth")
|
||||
if errMkdirAuth := os.MkdirAll(authDir, 0o700); errMkdirAuth != nil {
|
||||
t.Fatalf("failed to create auth dir: %v", errMkdirAuth)
|
||||
}
|
||||
|
||||
cfg := &proxyconfig.Config{
|
||||
SDKConfig: proxyconfig.SDKConfig{
|
||||
RequestLog: false,
|
||||
},
|
||||
AuthDir: authDir,
|
||||
ErrorLogsMaxFiles: 10,
|
||||
}
|
||||
|
||||
logger := defaultRequestLoggerFactory(cfg, configPath)
|
||||
fileLogger, ok := logger.(*internallogging.FileRequestLogger)
|
||||
if !ok {
|
||||
t.Fatalf("expected *FileRequestLogger, got %T", logger)
|
||||
}
|
||||
|
||||
errLog := fileLogger.LogRequestWithOptions(
|
||||
"/v1/chat/completions",
|
||||
http.MethodPost,
|
||||
map[string][]string{"Content-Type": []string{"application/json"}},
|
||||
[]byte(`{"input":"hello"}`),
|
||||
http.StatusBadGateway,
|
||||
map[string][]string{"Content-Type": []string{"application/json"}},
|
||||
[]byte(`{"error":"upstream failure"}`),
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
true,
|
||||
"issue-1711",
|
||||
time.Now(),
|
||||
time.Now(),
|
||||
)
|
||||
if errLog != nil {
|
||||
t.Fatalf("failed to write forced error request log: %v", errLog)
|
||||
}
|
||||
|
||||
authLogsDir := filepath.Join(authDir, "logs")
|
||||
authEntries, errReadAuthDir := os.ReadDir(authLogsDir)
|
||||
if errReadAuthDir != nil {
|
||||
t.Fatalf("failed to read auth logs dir %s: %v", authLogsDir, errReadAuthDir)
|
||||
}
|
||||
foundErrorLogInAuthDir := false
|
||||
for _, entry := range authEntries {
|
||||
if strings.HasPrefix(entry.Name(), "error-") && strings.HasSuffix(entry.Name(), ".log") {
|
||||
foundErrorLogInAuthDir = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !foundErrorLogInAuthDir {
|
||||
t.Fatalf("expected forced error log in auth fallback dir %s, got entries: %+v", authLogsDir, authEntries)
|
||||
}
|
||||
|
||||
configLogsDir := filepath.Join(configDir, "logs")
|
||||
configEntries, errReadConfigDir := os.ReadDir(configLogsDir)
|
||||
if errReadConfigDir != nil && !os.IsNotExist(errReadConfigDir) {
|
||||
t.Fatalf("failed to inspect config logs dir %s: %v", configLogsDir, errReadConfigDir)
|
||||
}
|
||||
for _, entry := range configEntries {
|
||||
if strings.HasPrefix(entry.Name(), "error-") && strings.HasSuffix(entry.Name(), ".log") {
|
||||
t.Fatalf("unexpected forced error log in config dir %s", configLogsDir)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -36,11 +36,21 @@ type ClaudeTokenStorage struct {
|
||||
|
||||
// Expire is the timestamp when the current access token expires.
|
||||
Expire string `json:"expired"`
|
||||
|
||||
// Metadata holds arbitrary key-value pairs injected via hooks.
|
||||
// It is not exported to JSON directly to allow flattening during serialization.
|
||||
Metadata map[string]any `json:"-"`
|
||||
}
|
||||
|
||||
// SetMetadata allows external callers to inject metadata into the storage before saving.
|
||||
func (ts *ClaudeTokenStorage) SetMetadata(meta map[string]any) {
|
||||
ts.Metadata = meta
|
||||
}
|
||||
|
||||
// SaveTokenToFile serializes the Claude token storage to a JSON file.
|
||||
// This method creates the necessary directory structure and writes the token
|
||||
// data in JSON format to the specified file path for persistent storage.
|
||||
// It merges any injected metadata into the top-level JSON object.
|
||||
//
|
||||
// Parameters:
|
||||
// - authFilePath: The full path where the token file should be saved
|
||||
@@ -65,8 +75,14 @@ func (ts *ClaudeTokenStorage) SaveTokenToFile(authFilePath string) error {
|
||||
_ = f.Close()
|
||||
}()
|
||||
|
||||
// Merge metadata using helper
|
||||
data, errMerge := misc.MergeMetadata(ts, ts.Metadata)
|
||||
if errMerge != nil {
|
||||
return fmt.Errorf("failed to merge metadata: %w", errMerge)
|
||||
}
|
||||
|
||||
// Encode and write the token data as JSON
|
||||
if err = json.NewEncoder(f).Encode(ts); err != nil {
|
||||
if err = json.NewEncoder(f).Encode(data); err != nil {
|
||||
return fmt.Errorf("failed to write token to file: %w", err)
|
||||
}
|
||||
return nil
|
||||
|
||||
@@ -4,18 +4,18 @@ package claude
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
tls "github.com/refraction-networking/utls"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/proxyutil"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/net/http2"
|
||||
"golang.org/x/net/proxy"
|
||||
)
|
||||
|
||||
// utlsRoundTripper implements http.RoundTripper using utls with Firefox fingerprint
|
||||
// utlsRoundTripper implements http.RoundTripper using utls with Chrome fingerprint
|
||||
// to bypass Cloudflare's TLS fingerprinting on Anthropic domains.
|
||||
type utlsRoundTripper struct {
|
||||
// mu protects the connections map and pending map
|
||||
@@ -31,17 +31,12 @@ type utlsRoundTripper struct {
|
||||
// newUtlsRoundTripper creates a new utls-based round tripper with optional proxy support
|
||||
func newUtlsRoundTripper(cfg *config.SDKConfig) *utlsRoundTripper {
|
||||
var dialer proxy.Dialer = proxy.Direct
|
||||
if cfg != nil && cfg.ProxyURL != "" {
|
||||
proxyURL, err := url.Parse(cfg.ProxyURL)
|
||||
if err != nil {
|
||||
log.Errorf("failed to parse proxy URL %q: %v", cfg.ProxyURL, err)
|
||||
} else {
|
||||
pDialer, err := proxy.FromURL(proxyURL, proxy.Direct)
|
||||
if err != nil {
|
||||
log.Errorf("failed to create proxy dialer for %q: %v", cfg.ProxyURL, err)
|
||||
} else {
|
||||
dialer = pDialer
|
||||
}
|
||||
if cfg != nil {
|
||||
proxyDialer, mode, errBuild := proxyutil.BuildDialer(cfg.ProxyURL)
|
||||
if errBuild != nil {
|
||||
log.Errorf("failed to configure proxy dialer for %q: %v", cfg.ProxyURL, errBuild)
|
||||
} else if mode != proxyutil.ModeInherit && proxyDialer != nil {
|
||||
dialer = proxyDialer
|
||||
}
|
||||
}
|
||||
|
||||
@@ -100,7 +95,9 @@ func (t *utlsRoundTripper) getOrCreateConnection(host, addr string) (*http2.Clie
|
||||
return h2Conn, nil
|
||||
}
|
||||
|
||||
// createConnection creates a new HTTP/2 connection with Firefox TLS fingerprint
|
||||
// createConnection creates a new HTTP/2 connection with Chrome TLS fingerprint.
|
||||
// Chrome's TLS fingerprint is closer to Node.js/OpenSSL (which real Claude Code uses)
|
||||
// than Firefox, reducing the mismatch between TLS layer and HTTP headers.
|
||||
func (t *utlsRoundTripper) createConnection(host, addr string) (*http2.ClientConn, error) {
|
||||
conn, err := t.dialer.Dial("tcp", addr)
|
||||
if err != nil {
|
||||
@@ -108,7 +105,7 @@ func (t *utlsRoundTripper) createConnection(host, addr string) (*http2.ClientCon
|
||||
}
|
||||
|
||||
tlsConfig := &tls.Config{ServerName: host}
|
||||
tlsConn := tls.UClient(conn, tlsConfig, tls.HelloFirefox_Auto)
|
||||
tlsConn := tls.UClient(conn, tlsConfig, tls.HelloChrome_Auto)
|
||||
|
||||
if err := tlsConn.Handshake(); err != nil {
|
||||
conn.Close()
|
||||
@@ -156,7 +153,7 @@ func (t *utlsRoundTripper) RoundTrip(req *http.Request) (*http.Response, error)
|
||||
}
|
||||
|
||||
// NewAnthropicHttpClient creates an HTTP client that bypasses TLS fingerprinting
|
||||
// for Anthropic domains by using utls with Firefox fingerprint.
|
||||
// for Anthropic domains by using utls with Chrome fingerprint.
|
||||
// It accepts optional SDK configuration for proxy settings.
|
||||
func NewAnthropicHttpClient(cfg *config.SDKConfig) *http.Client {
|
||||
return &http.Client{
|
||||
|
||||
335
internal/auth/codebuddy/codebuddy_auth.go
Normal file
335
internal/auth/codebuddy/codebuddy_auth.go
Normal file
@@ -0,0 +1,335 @@
|
||||
package codebuddy
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
)
|
||||
|
||||
const (
|
||||
BaseURL = "https://copilot.tencent.com"
|
||||
DefaultDomain = "www.codebuddy.cn"
|
||||
UserAgent = "CLI/2.63.2 CodeBuddy/2.63.2"
|
||||
|
||||
codeBuddyStatePath = "/v2/plugin/auth/state"
|
||||
codeBuddyTokenPath = "/v2/plugin/auth/token"
|
||||
codeBuddyRefreshPath = "/v2/plugin/auth/token/refresh"
|
||||
pollInterval = 5 * time.Second
|
||||
maxPollDuration = 5 * time.Minute
|
||||
codeLoginPending = 11217
|
||||
codeSuccess = 0
|
||||
)
|
||||
|
||||
type CodeBuddyAuth struct {
|
||||
httpClient *http.Client
|
||||
cfg *config.Config
|
||||
baseURL string
|
||||
}
|
||||
|
||||
func NewCodeBuddyAuth(cfg *config.Config) *CodeBuddyAuth {
|
||||
httpClient := &http.Client{Timeout: 30 * time.Second}
|
||||
if cfg != nil {
|
||||
httpClient = util.SetProxy(&cfg.SDKConfig, httpClient)
|
||||
}
|
||||
return &CodeBuddyAuth{httpClient: httpClient, cfg: cfg, baseURL: BaseURL}
|
||||
}
|
||||
|
||||
// AuthState holds the state and auth URL returned by the auth state API.
|
||||
type AuthState struct {
|
||||
State string
|
||||
AuthURL string
|
||||
}
|
||||
|
||||
// FetchAuthState calls POST /v2/plugin/auth/state?platform=CLI to get the state and login URL.
|
||||
func (a *CodeBuddyAuth) FetchAuthState(ctx context.Context) (*AuthState, error) {
|
||||
stateURL := fmt.Sprintf("%s%s?platform=CLI", a.baseURL, codeBuddyStatePath)
|
||||
body := []byte("{}")
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, stateURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("codebuddy: failed to create auth state request: %w", err)
|
||||
}
|
||||
|
||||
requestID := uuid.NewString()
|
||||
req.Header.Set("Accept", "application/json, text/plain, */*")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("X-Requested-With", "XMLHttpRequest")
|
||||
req.Header.Set("X-Domain", "copilot.tencent.com")
|
||||
req.Header.Set("X-No-Authorization", "true")
|
||||
req.Header.Set("X-No-User-Id", "true")
|
||||
req.Header.Set("X-No-Enterprise-Id", "true")
|
||||
req.Header.Set("X-No-Department-Info", "true")
|
||||
req.Header.Set("X-Product", "SaaS")
|
||||
req.Header.Set("User-Agent", UserAgent)
|
||||
req.Header.Set("X-Request-ID", requestID)
|
||||
|
||||
resp, err := a.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("codebuddy: auth state request failed: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if errClose := resp.Body.Close(); errClose != nil {
|
||||
log.Errorf("codebuddy auth state: close body error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
|
||||
bodyBytes, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("codebuddy: failed to read auth state response: %w", err)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("codebuddy: auth state request returned status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
var result struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg"`
|
||||
Data *struct {
|
||||
State string `json:"state"`
|
||||
AuthURL string `json:"authUrl"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err = json.Unmarshal(bodyBytes, &result); err != nil {
|
||||
return nil, fmt.Errorf("codebuddy: failed to parse auth state response: %w", err)
|
||||
}
|
||||
if result.Code != codeSuccess {
|
||||
return nil, fmt.Errorf("codebuddy: auth state request failed with code %d: %s", result.Code, result.Msg)
|
||||
}
|
||||
if result.Data == nil || result.Data.State == "" || result.Data.AuthURL == "" {
|
||||
return nil, fmt.Errorf("codebuddy: auth state response missing state or authUrl")
|
||||
}
|
||||
|
||||
return &AuthState{
|
||||
State: result.Data.State,
|
||||
AuthURL: result.Data.AuthURL,
|
||||
}, nil
|
||||
}
|
||||
|
||||
type pollResponse struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg"`
|
||||
RequestID string `json:"requestId"`
|
||||
Data *struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
RefreshToken string `json:"refreshToken"`
|
||||
ExpiresIn int64 `json:"expiresIn"`
|
||||
TokenType string `json:"tokenType"`
|
||||
Domain string `json:"domain"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
// doPollRequest performs a single polling request, safely reading and closing the response body
|
||||
func (a *CodeBuddyAuth) doPollRequest(ctx context.Context, pollURL string) ([]byte, int, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, pollURL, nil)
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("%w: %v", ErrTokenFetchFailed, err)
|
||||
}
|
||||
a.applyPollHeaders(req)
|
||||
|
||||
resp, err := a.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer func() {
|
||||
if errClose := resp.Body.Close(); errClose != nil {
|
||||
log.Errorf("codebuddy poll: close body error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, resp.StatusCode, fmt.Errorf("codebuddy poll: failed to read response body: %w", err)
|
||||
}
|
||||
return body, resp.StatusCode, nil
|
||||
}
|
||||
|
||||
// PollForToken polls until the user completes browser authorization and returns auth data.
|
||||
func (a *CodeBuddyAuth) PollForToken(ctx context.Context, state string) (*CodeBuddyTokenStorage, error) {
|
||||
deadline := time.Now().Add(maxPollDuration)
|
||||
pollURL := fmt.Sprintf("%s%s?state=%s", a.baseURL, codeBuddyTokenPath, url.QueryEscape(state))
|
||||
|
||||
for time.Now().Before(deadline) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case <-time.After(pollInterval):
|
||||
}
|
||||
|
||||
body, statusCode, err := a.doPollRequest(ctx, pollURL)
|
||||
if err != nil {
|
||||
log.Debugf("codebuddy poll: request error: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
if statusCode != http.StatusOK {
|
||||
log.Debugf("codebuddy poll: unexpected status %d", statusCode)
|
||||
continue
|
||||
}
|
||||
|
||||
var result pollResponse
|
||||
if err := json.Unmarshal(body, &result); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
switch result.Code {
|
||||
case codeSuccess:
|
||||
if result.Data == nil {
|
||||
return nil, fmt.Errorf("%w: empty data in response", ErrTokenFetchFailed)
|
||||
}
|
||||
userID, _ := a.DecodeUserID(result.Data.AccessToken)
|
||||
return &CodeBuddyTokenStorage{
|
||||
AccessToken: result.Data.AccessToken,
|
||||
RefreshToken: result.Data.RefreshToken,
|
||||
ExpiresIn: result.Data.ExpiresIn,
|
||||
TokenType: result.Data.TokenType,
|
||||
Domain: result.Data.Domain,
|
||||
UserID: userID,
|
||||
Type: "codebuddy",
|
||||
}, nil
|
||||
case codeLoginPending:
|
||||
// continue polling
|
||||
default:
|
||||
// TODO: when the CodeBuddy API error code for user denial is known,
|
||||
// return ErrAccessDenied here instead of ErrTokenFetchFailed.
|
||||
return nil, fmt.Errorf("%w: server returned code %d: %s", ErrTokenFetchFailed, result.Code, result.Msg)
|
||||
}
|
||||
}
|
||||
return nil, ErrPollingTimeout
|
||||
}
|
||||
|
||||
// DecodeUserID decodes the sub field from a JWT access token as the user ID.
|
||||
func (a *CodeBuddyAuth) DecodeUserID(accessToken string) (string, error) {
|
||||
parts := strings.Split(accessToken, ".")
|
||||
if len(parts) < 2 {
|
||||
return "", ErrJWTDecodeFailed
|
||||
}
|
||||
payload, err := base64.RawURLEncoding.DecodeString(parts[1])
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: %v", ErrJWTDecodeFailed, err)
|
||||
}
|
||||
var claims struct {
|
||||
Sub string `json:"sub"`
|
||||
}
|
||||
if err := json.Unmarshal(payload, &claims); err != nil {
|
||||
return "", fmt.Errorf("%w: %v", ErrJWTDecodeFailed, err)
|
||||
}
|
||||
if claims.Sub == "" {
|
||||
return "", fmt.Errorf("%w: sub claim is empty", ErrJWTDecodeFailed)
|
||||
}
|
||||
return claims.Sub, nil
|
||||
}
|
||||
|
||||
// RefreshToken exchanges a refresh token for a new access token.
|
||||
// It calls POST /v2/plugin/auth/token/refresh with the required headers.
|
||||
func (a *CodeBuddyAuth) RefreshToken(ctx context.Context, accessToken, refreshToken, userID, domain string) (*CodeBuddyTokenStorage, error) {
|
||||
if domain == "" {
|
||||
domain = DefaultDomain
|
||||
}
|
||||
refreshURL := fmt.Sprintf("%s%s", a.baseURL, codeBuddyRefreshPath)
|
||||
body := []byte("{}")
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, refreshURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("codebuddy: failed to create refresh request: %w", err)
|
||||
}
|
||||
|
||||
requestID := strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
req.Header.Set("Accept", "application/json, text/plain, */*")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("X-Requested-With", "XMLHttpRequest")
|
||||
req.Header.Set("X-Domain", domain)
|
||||
req.Header.Set("X-Refresh-Token", refreshToken)
|
||||
req.Header.Set("X-Auth-Refresh-Source", "plugin")
|
||||
req.Header.Set("X-Request-ID", requestID)
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
req.Header.Set("X-User-Id", userID)
|
||||
req.Header.Set("X-Product", "SaaS")
|
||||
req.Header.Set("User-Agent", UserAgent)
|
||||
|
||||
resp, err := a.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("codebuddy: refresh request failed: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if errClose := resp.Body.Close(); errClose != nil {
|
||||
log.Errorf("codebuddy refresh: close body error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
|
||||
bodyBytes, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("codebuddy: failed to read refresh response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
|
||||
return nil, fmt.Errorf("codebuddy: refresh token rejected (status %d)", resp.StatusCode)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("codebuddy: refresh failed with status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
var result struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg"`
|
||||
Data *struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
RefreshToken string `json:"refreshToken"`
|
||||
ExpiresIn int64 `json:"expiresIn"`
|
||||
RefreshExpiresIn int64 `json:"refreshExpiresIn"`
|
||||
TokenType string `json:"tokenType"`
|
||||
Domain string `json:"domain"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err = json.Unmarshal(bodyBytes, &result); err != nil {
|
||||
return nil, fmt.Errorf("codebuddy: failed to parse refresh response: %w", err)
|
||||
}
|
||||
if result.Code != codeSuccess {
|
||||
return nil, fmt.Errorf("codebuddy: refresh failed with code %d: %s", result.Code, result.Msg)
|
||||
}
|
||||
if result.Data == nil {
|
||||
return nil, fmt.Errorf("codebuddy: empty data in refresh response")
|
||||
}
|
||||
|
||||
newUserID, _ := a.DecodeUserID(result.Data.AccessToken)
|
||||
if newUserID == "" {
|
||||
newUserID = userID
|
||||
}
|
||||
tokenDomain := result.Data.Domain
|
||||
if tokenDomain == "" {
|
||||
tokenDomain = domain
|
||||
}
|
||||
|
||||
return &CodeBuddyTokenStorage{
|
||||
AccessToken: result.Data.AccessToken,
|
||||
RefreshToken: result.Data.RefreshToken,
|
||||
ExpiresIn: result.Data.ExpiresIn,
|
||||
RefreshExpiresIn: result.Data.RefreshExpiresIn,
|
||||
TokenType: result.Data.TokenType,
|
||||
Domain: tokenDomain,
|
||||
UserID: newUserID,
|
||||
Type: "codebuddy",
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (a *CodeBuddyAuth) applyPollHeaders(req *http.Request) {
|
||||
req.Header.Set("Accept", "application/json, text/plain, */*")
|
||||
req.Header.Set("User-Agent", UserAgent)
|
||||
req.Header.Set("X-Requested-With", "XMLHttpRequest")
|
||||
req.Header.Set("X-No-Authorization", "true")
|
||||
req.Header.Set("X-No-User-Id", "true")
|
||||
req.Header.Set("X-No-Enterprise-Id", "true")
|
||||
req.Header.Set("X-No-Department-Info", "true")
|
||||
req.Header.Set("X-Product", "SaaS")
|
||||
}
|
||||
285
internal/auth/codebuddy/codebuddy_auth_http_test.go
Normal file
285
internal/auth/codebuddy/codebuddy_auth_http_test.go
Normal file
@@ -0,0 +1,285 @@
|
||||
package codebuddy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// newTestAuth creates a CodeBuddyAuth pointing at the given test server.
|
||||
func newTestAuth(serverURL string) *CodeBuddyAuth {
|
||||
return &CodeBuddyAuth{
|
||||
httpClient: http.DefaultClient,
|
||||
baseURL: serverURL,
|
||||
}
|
||||
}
|
||||
|
||||
// fakeJWT builds a minimal JWT with the given sub claim for testing.
|
||||
func fakeJWT(sub string) string {
|
||||
header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"RS256"}`))
|
||||
payload, _ := json.Marshal(map[string]any{"sub": sub, "iat": 1234567890})
|
||||
encodedPayload := base64.RawURLEncoding.EncodeToString(payload)
|
||||
return header + "." + encodedPayload + ".sig"
|
||||
}
|
||||
|
||||
// --- FetchAuthState tests ---
|
||||
|
||||
func TestFetchAuthState_Success(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
t.Errorf("expected POST, got %s", r.Method)
|
||||
}
|
||||
if got := r.URL.Path; got != codeBuddyStatePath {
|
||||
t.Errorf("expected path %s, got %s", codeBuddyStatePath, got)
|
||||
}
|
||||
if got := r.URL.Query().Get("platform"); got != "CLI" {
|
||||
t.Errorf("expected platform=CLI, got %s", got)
|
||||
}
|
||||
if got := r.Header.Get("User-Agent"); got != UserAgent {
|
||||
t.Errorf("expected User-Agent %s, got %s", UserAgent, got)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"code": 0,
|
||||
"msg": "ok",
|
||||
"data": map[string]any{
|
||||
"state": "test-state-abc",
|
||||
"authUrl": "https://example.com/login?state=test-state-abc",
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
auth := newTestAuth(srv.URL)
|
||||
result, err := auth.FetchAuthState(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if result.State != "test-state-abc" {
|
||||
t.Errorf("expected state 'test-state-abc', got '%s'", result.State)
|
||||
}
|
||||
if result.AuthURL != "https://example.com/login?state=test-state-abc" {
|
||||
t.Errorf("unexpected authURL: %s", result.AuthURL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchAuthState_NonOKStatus(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = w.Write([]byte("internal error"))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
auth := newTestAuth(srv.URL)
|
||||
_, err := auth.FetchAuthState(context.Background())
|
||||
if err == nil {
|
||||
t.Fatal("expected error for non-200 status")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchAuthState_APIErrorCode(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"code": 10001,
|
||||
"msg": "rate limited",
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
auth := newTestAuth(srv.URL)
|
||||
_, err := auth.FetchAuthState(context.Background())
|
||||
if err == nil {
|
||||
t.Fatal("expected error for non-zero code")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchAuthState_MissingData(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"code": 0,
|
||||
"msg": "ok",
|
||||
"data": map[string]any{
|
||||
"state": "",
|
||||
"authUrl": "",
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
auth := newTestAuth(srv.URL)
|
||||
_, err := auth.FetchAuthState(context.Background())
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty state/authUrl")
|
||||
}
|
||||
}
|
||||
|
||||
// --- RefreshToken tests ---
|
||||
|
||||
func TestRefreshToken_Success(t *testing.T) {
|
||||
newAccessToken := fakeJWT("refreshed-user-456")
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
t.Errorf("expected POST, got %s", r.Method)
|
||||
}
|
||||
if got := r.URL.Path; got != codeBuddyRefreshPath {
|
||||
t.Errorf("expected path %s, got %s", codeBuddyRefreshPath, got)
|
||||
}
|
||||
if got := r.Header.Get("X-Refresh-Token"); got != "old-refresh-token" {
|
||||
t.Errorf("expected X-Refresh-Token 'old-refresh-token', got '%s'", got)
|
||||
}
|
||||
if got := r.Header.Get("Authorization"); got != "Bearer old-access-token" {
|
||||
t.Errorf("expected Authorization 'Bearer old-access-token', got '%s'", got)
|
||||
}
|
||||
if got := r.Header.Get("X-User-Id"); got != "user-123" {
|
||||
t.Errorf("expected X-User-Id 'user-123', got '%s'", got)
|
||||
}
|
||||
if got := r.Header.Get("X-Domain"); got != "custom.domain.com" {
|
||||
t.Errorf("expected X-Domain 'custom.domain.com', got '%s'", got)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"code": 0,
|
||||
"msg": "ok",
|
||||
"data": map[string]any{
|
||||
"accessToken": newAccessToken,
|
||||
"refreshToken": "new-refresh-token",
|
||||
"expiresIn": 3600,
|
||||
"refreshExpiresIn": 86400,
|
||||
"tokenType": "bearer",
|
||||
"domain": "custom.domain.com",
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
auth := newTestAuth(srv.URL)
|
||||
storage, err := auth.RefreshToken(context.Background(), "old-access-token", "old-refresh-token", "user-123", "custom.domain.com")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if storage.AccessToken != newAccessToken {
|
||||
t.Errorf("expected new access token, got '%s'", storage.AccessToken)
|
||||
}
|
||||
if storage.RefreshToken != "new-refresh-token" {
|
||||
t.Errorf("expected 'new-refresh-token', got '%s'", storage.RefreshToken)
|
||||
}
|
||||
if storage.UserID != "refreshed-user-456" {
|
||||
t.Errorf("expected userID 'refreshed-user-456', got '%s'", storage.UserID)
|
||||
}
|
||||
if storage.ExpiresIn != 3600 {
|
||||
t.Errorf("expected expiresIn 3600, got %d", storage.ExpiresIn)
|
||||
}
|
||||
if storage.RefreshExpiresIn != 86400 {
|
||||
t.Errorf("expected refreshExpiresIn 86400, got %d", storage.RefreshExpiresIn)
|
||||
}
|
||||
if storage.Domain != "custom.domain.com" {
|
||||
t.Errorf("expected domain 'custom.domain.com', got '%s'", storage.Domain)
|
||||
}
|
||||
if storage.Type != "codebuddy" {
|
||||
t.Errorf("expected type 'codebuddy', got '%s'", storage.Type)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshToken_DefaultDomain(t *testing.T) {
|
||||
var receivedDomain string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
receivedDomain = r.Header.Get("X-Domain")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"code": 0,
|
||||
"msg": "ok",
|
||||
"data": map[string]any{
|
||||
"accessToken": fakeJWT("user-1"),
|
||||
"refreshToken": "rt",
|
||||
"expiresIn": 3600,
|
||||
"tokenType": "bearer",
|
||||
"domain": DefaultDomain,
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
auth := newTestAuth(srv.URL)
|
||||
_, err := auth.RefreshToken(context.Background(), "at", "rt", "uid", "")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if receivedDomain != DefaultDomain {
|
||||
t.Errorf("expected default domain '%s', got '%s'", DefaultDomain, receivedDomain)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshToken_Unauthorized(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
auth := newTestAuth(srv.URL)
|
||||
_, err := auth.RefreshToken(context.Background(), "at", "rt", "uid", "d")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for 401 response")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshToken_Forbidden(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusForbidden)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
auth := newTestAuth(srv.URL)
|
||||
_, err := auth.RefreshToken(context.Background(), "at", "rt", "uid", "d")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for 403 response")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshToken_APIErrorCode(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"code": 40001,
|
||||
"msg": "invalid refresh token",
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
auth := newTestAuth(srv.URL)
|
||||
_, err := auth.RefreshToken(context.Background(), "at", "rt", "uid", "d")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for non-zero API code")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshToken_FallbackUserIDAndDomain(t *testing.T) {
|
||||
// When the new access token cannot be decoded for userID, it should fall back to the provided one.
|
||||
// When the response domain is empty, it should fall back to the request domain.
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"code": 0,
|
||||
"msg": "ok",
|
||||
"data": map[string]any{
|
||||
"accessToken": "not-a-valid-jwt",
|
||||
"refreshToken": "new-rt",
|
||||
"expiresIn": 7200,
|
||||
"tokenType": "bearer",
|
||||
"domain": "",
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
auth := newTestAuth(srv.URL)
|
||||
storage, err := auth.RefreshToken(context.Background(), "at", "rt", "original-uid", "original.domain.com")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if storage.UserID != "original-uid" {
|
||||
t.Errorf("expected fallback userID 'original-uid', got '%s'", storage.UserID)
|
||||
}
|
||||
if storage.Domain != "original.domain.com" {
|
||||
t.Errorf("expected fallback domain 'original.domain.com', got '%s'", storage.Domain)
|
||||
}
|
||||
}
|
||||
22
internal/auth/codebuddy/codebuddy_auth_test.go
Normal file
22
internal/auth/codebuddy/codebuddy_auth_test.go
Normal file
@@ -0,0 +1,22 @@
|
||||
package codebuddy_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/auth/codebuddy"
|
||||
)
|
||||
|
||||
func TestDecodeUserID_ValidJWT(t *testing.T) {
|
||||
// JWT payload: {"sub":"test-user-id-123","iat":1234567890}
|
||||
// base64url encode: eyJzdWIiOiJ0ZXN0LXVzZXItaWQtMTIzIiwiaWF0IjoxMjM0NTY3ODkwfQ
|
||||
token := "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ0ZXN0LXVzZXItaWQtMTIzIiwiaWF0IjoxMjM0NTY3ODkwfQ.sig"
|
||||
auth := codebuddy.NewCodeBuddyAuth(nil)
|
||||
userID, err := auth.DecodeUserID(token)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if userID != "test-user-id-123" {
|
||||
t.Errorf("expected 'test-user-id-123', got '%s'", userID)
|
||||
}
|
||||
}
|
||||
|
||||
25
internal/auth/codebuddy/errors.go
Normal file
25
internal/auth/codebuddy/errors.go
Normal file
@@ -0,0 +1,25 @@
|
||||
package codebuddy
|
||||
|
||||
import "errors"
|
||||
|
||||
var (
|
||||
ErrPollingTimeout = errors.New("codebuddy: polling timeout, user did not authorize in time")
|
||||
ErrAccessDenied = errors.New("codebuddy: access denied by user")
|
||||
ErrTokenFetchFailed = errors.New("codebuddy: failed to fetch token from server")
|
||||
ErrJWTDecodeFailed = errors.New("codebuddy: failed to decode JWT token")
|
||||
)
|
||||
|
||||
func GetUserFriendlyMessage(err error) string {
|
||||
switch {
|
||||
case errors.Is(err, ErrPollingTimeout):
|
||||
return "Authentication timed out. Please try again."
|
||||
case errors.Is(err, ErrAccessDenied):
|
||||
return "Access denied. Please try again and approve the login request."
|
||||
case errors.Is(err, ErrJWTDecodeFailed):
|
||||
return "Failed to decode token. Please try logging in again."
|
||||
case errors.Is(err, ErrTokenFetchFailed):
|
||||
return "Failed to fetch token from server. Please try again."
|
||||
default:
|
||||
return "Authentication failed: " + err.Error()
|
||||
}
|
||||
}
|
||||
65
internal/auth/codebuddy/token.go
Normal file
65
internal/auth/codebuddy/token.go
Normal file
@@ -0,0 +1,65 @@
|
||||
// Package codebuddy provides authentication and token management functionality
|
||||
// for CodeBuddy AI services. It handles OAuth2 token storage, serialization,
|
||||
// and retrieval for maintaining authenticated sessions with the CodeBuddy API.
|
||||
package codebuddy
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
|
||||
)
|
||||
|
||||
// CodeBuddyTokenStorage stores OAuth token information for CodeBuddy API authentication.
|
||||
// It maintains compatibility with the existing auth system while adding CodeBuddy-specific fields
|
||||
// for managing access tokens and user account information.
|
||||
type CodeBuddyTokenStorage struct {
|
||||
// AccessToken is the OAuth2 access token used for authenticating API requests.
|
||||
AccessToken string `json:"access_token"`
|
||||
// RefreshToken is the OAuth2 refresh token used to obtain new access tokens.
|
||||
RefreshToken string `json:"refresh_token"`
|
||||
// ExpiresIn is the number of seconds until the access token expires.
|
||||
ExpiresIn int64 `json:"expires_in"`
|
||||
// RefreshExpiresIn is the number of seconds until the refresh token expires.
|
||||
RefreshExpiresIn int64 `json:"refresh_expires_in,omitempty"`
|
||||
// TokenType is the type of token, typically "bearer".
|
||||
TokenType string `json:"token_type"`
|
||||
// Domain is the CodeBuddy service domain/region.
|
||||
Domain string `json:"domain"`
|
||||
// UserID is the user ID associated with this token.
|
||||
UserID string `json:"user_id"`
|
||||
// Type indicates the authentication provider type, always "codebuddy" for this storage.
|
||||
Type string `json:"type"`
|
||||
}
|
||||
|
||||
// SaveTokenToFile serializes the CodeBuddy token storage to a JSON file.
|
||||
// This method creates the necessary directory structure and writes the token
|
||||
// data in JSON format to the specified file path for persistent storage.
|
||||
//
|
||||
// Parameters:
|
||||
// - authFilePath: The full path where the token file should be saved
|
||||
//
|
||||
// Returns:
|
||||
// - error: An error if the operation fails, nil otherwise
|
||||
func (s *CodeBuddyTokenStorage) SaveTokenToFile(authFilePath string) error {
|
||||
misc.LogSavingCredentials(authFilePath)
|
||||
s.Type = "codebuddy"
|
||||
if err := os.MkdirAll(filepath.Dir(authFilePath), 0700); err != nil {
|
||||
return fmt.Errorf("failed to create directory: %w", err)
|
||||
}
|
||||
|
||||
f, err := os.OpenFile(authFilePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create token file: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
_ = f.Close()
|
||||
}()
|
||||
|
||||
if err = json.NewEncoder(f).Encode(s); err != nil {
|
||||
return fmt.Errorf("failed to write token to file: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -71,16 +71,26 @@ func (o *CodexAuth) GenerateAuthURL(state string, pkceCodes *PKCECodes) (string,
|
||||
// It performs an HTTP POST request to the OpenAI token endpoint with the provided
|
||||
// authorization code and PKCE verifier.
|
||||
func (o *CodexAuth) ExchangeCodeForTokens(ctx context.Context, code string, pkceCodes *PKCECodes) (*CodexAuthBundle, error) {
|
||||
return o.ExchangeCodeForTokensWithRedirect(ctx, code, RedirectURI, pkceCodes)
|
||||
}
|
||||
|
||||
// ExchangeCodeForTokensWithRedirect exchanges an authorization code for tokens using
|
||||
// a caller-provided redirect URI. This supports alternate auth flows such as device
|
||||
// login while preserving the existing token parsing and storage behavior.
|
||||
func (o *CodexAuth) ExchangeCodeForTokensWithRedirect(ctx context.Context, code, redirectURI string, pkceCodes *PKCECodes) (*CodexAuthBundle, error) {
|
||||
if pkceCodes == nil {
|
||||
return nil, fmt.Errorf("PKCE codes are required for token exchange")
|
||||
}
|
||||
if strings.TrimSpace(redirectURI) == "" {
|
||||
return nil, fmt.Errorf("redirect URI is required for token exchange")
|
||||
}
|
||||
|
||||
// Prepare token exchange request
|
||||
data := url.Values{
|
||||
"grant_type": {"authorization_code"},
|
||||
"client_id": {ClientID},
|
||||
"code": {code},
|
||||
"redirect_uri": {RedirectURI},
|
||||
"redirect_uri": {strings.TrimSpace(redirectURI)},
|
||||
"code_verifier": {pkceCodes.CodeVerifier},
|
||||
}
|
||||
|
||||
@@ -266,6 +276,10 @@ func (o *CodexAuth) RefreshTokensWithRetry(ctx context.Context, refreshToken str
|
||||
if err == nil {
|
||||
return tokenData, nil
|
||||
}
|
||||
if isNonRetryableRefreshErr(err) {
|
||||
log.Warnf("Token refresh attempt %d failed with non-retryable error: %v", attempt+1, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
lastErr = err
|
||||
log.Warnf("Token refresh attempt %d failed: %v", attempt+1, err)
|
||||
@@ -274,6 +288,14 @@ func (o *CodexAuth) RefreshTokensWithRetry(ctx context.Context, refreshToken str
|
||||
return nil, fmt.Errorf("token refresh failed after %d attempts: %w", maxRetries, lastErr)
|
||||
}
|
||||
|
||||
func isNonRetryableRefreshErr(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
raw := strings.ToLower(err.Error())
|
||||
return strings.Contains(raw, "refresh_token_reused")
|
||||
}
|
||||
|
||||
// UpdateTokenStorage updates an existing CodexTokenStorage with new token data.
|
||||
// This is typically called after a successful token refresh to persist the new credentials.
|
||||
func (o *CodexAuth) UpdateTokenStorage(storage *CodexTokenStorage, tokenData *CodexTokenData) {
|
||||
|
||||
44
internal/auth/codex/openai_auth_test.go
Normal file
44
internal/auth/codex/openai_auth_test.go
Normal file
@@ -0,0 +1,44 @@
|
||||
package codex
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return f(req)
|
||||
}
|
||||
|
||||
func TestRefreshTokensWithRetry_NonRetryableOnlyAttemptsOnce(t *testing.T) {
|
||||
var calls int32
|
||||
auth := &CodexAuth{
|
||||
httpClient: &http.Client{
|
||||
Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
atomic.AddInt32(&calls, 1)
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":"invalid_grant","code":"refresh_token_reused"}`)),
|
||||
Header: make(http.Header),
|
||||
Request: req,
|
||||
}, nil
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
_, err := auth.RefreshTokensWithRetry(context.Background(), "dummy_refresh_token", 3)
|
||||
if err == nil {
|
||||
t.Fatalf("expected error for non-retryable refresh failure")
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(err.Error()), "refresh_token_reused") {
|
||||
t.Fatalf("expected refresh_token_reused in error, got: %v", err)
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != 1 {
|
||||
t.Fatalf("expected 1 refresh attempt, got %d", got)
|
||||
}
|
||||
}
|
||||
@@ -32,11 +32,21 @@ type CodexTokenStorage struct {
|
||||
Type string `json:"type"`
|
||||
// Expire is the timestamp when the current access token expires.
|
||||
Expire string `json:"expired"`
|
||||
|
||||
// Metadata holds arbitrary key-value pairs injected via hooks.
|
||||
// It is not exported to JSON directly to allow flattening during serialization.
|
||||
Metadata map[string]any `json:"-"`
|
||||
}
|
||||
|
||||
// SetMetadata allows external callers to inject metadata into the storage before saving.
|
||||
func (ts *CodexTokenStorage) SetMetadata(meta map[string]any) {
|
||||
ts.Metadata = meta
|
||||
}
|
||||
|
||||
// SaveTokenToFile serializes the Codex token storage to a JSON file.
|
||||
// This method creates the necessary directory structure and writes the token
|
||||
// data in JSON format to the specified file path for persistent storage.
|
||||
// It merges any injected metadata into the top-level JSON object.
|
||||
//
|
||||
// Parameters:
|
||||
// - authFilePath: The full path where the token file should be saved
|
||||
@@ -58,7 +68,13 @@ func (ts *CodexTokenStorage) SaveTokenToFile(authFilePath string) error {
|
||||
_ = f.Close()
|
||||
}()
|
||||
|
||||
if err = json.NewEncoder(f).Encode(ts); err != nil {
|
||||
// Merge metadata using helper
|
||||
data, errMerge := misc.MergeMetadata(ts, ts.Metadata)
|
||||
if errMerge != nil {
|
||||
return fmt.Errorf("failed to merge metadata: %w", errMerge)
|
||||
}
|
||||
|
||||
if err = json.NewEncoder(f).Encode(data); err != nil {
|
||||
return fmt.Errorf("failed to write token to file: %w", err)
|
||||
}
|
||||
return nil
|
||||
|
||||
@@ -8,6 +8,8 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
@@ -82,15 +84,21 @@ func (c *CopilotAuth) WaitForAuthorization(ctx context.Context, deviceCode *Devi
|
||||
}
|
||||
|
||||
// Fetch the GitHub username
|
||||
username, err := c.deviceClient.FetchUserInfo(ctx, tokenData.AccessToken)
|
||||
userInfo, err := c.deviceClient.FetchUserInfo(ctx, tokenData.AccessToken)
|
||||
if err != nil {
|
||||
log.Warnf("copilot: failed to fetch user info: %v", err)
|
||||
username = "unknown"
|
||||
}
|
||||
|
||||
username := userInfo.Login
|
||||
if username == "" {
|
||||
username = "github-user"
|
||||
}
|
||||
|
||||
return &CopilotAuthBundle{
|
||||
TokenData: tokenData,
|
||||
Username: username,
|
||||
Email: userInfo.Email,
|
||||
Name: userInfo.Name,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -150,12 +158,12 @@ func (c *CopilotAuth) ValidateToken(ctx context.Context, accessToken string) (bo
|
||||
return false, "", nil
|
||||
}
|
||||
|
||||
username, err := c.deviceClient.FetchUserInfo(ctx, accessToken)
|
||||
userInfo, err := c.deviceClient.FetchUserInfo(ctx, accessToken)
|
||||
if err != nil {
|
||||
return false, "", err
|
||||
}
|
||||
|
||||
return true, username, nil
|
||||
return true, userInfo.Login, nil
|
||||
}
|
||||
|
||||
// CreateTokenStorage creates a new CopilotTokenStorage from auth bundle.
|
||||
@@ -165,6 +173,8 @@ func (c *CopilotAuth) CreateTokenStorage(bundle *CopilotAuthBundle) *CopilotToke
|
||||
TokenType: bundle.TokenData.TokenType,
|
||||
Scope: bundle.TokenData.Scope,
|
||||
Username: bundle.Username,
|
||||
Email: bundle.Email,
|
||||
Name: bundle.Name,
|
||||
Type: "github-copilot",
|
||||
}
|
||||
}
|
||||
@@ -214,6 +224,97 @@ func (c *CopilotAuth) MakeAuthenticatedRequest(ctx context.Context, method, url
|
||||
return req, nil
|
||||
}
|
||||
|
||||
// CopilotModelEntry represents a single model entry returned by the Copilot /models API.
|
||||
type CopilotModelEntry struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
OwnedBy string `json:"owned_by"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Version string `json:"version,omitempty"`
|
||||
Capabilities map[string]any `json:"capabilities,omitempty"`
|
||||
}
|
||||
|
||||
// CopilotModelsResponse represents the response from the Copilot /models endpoint.
|
||||
type CopilotModelsResponse struct {
|
||||
Data []CopilotModelEntry `json:"data"`
|
||||
Object string `json:"object"`
|
||||
}
|
||||
|
||||
// maxModelsResponseSize is the maximum allowed response size from the /models endpoint (2 MB).
|
||||
const maxModelsResponseSize = 2 * 1024 * 1024
|
||||
|
||||
// allowedCopilotAPIHosts is the set of hosts that are considered safe for Copilot API requests.
|
||||
var allowedCopilotAPIHosts = map[string]bool{
|
||||
"api.githubcopilot.com": true,
|
||||
"api.individual.githubcopilot.com": true,
|
||||
"api.business.githubcopilot.com": true,
|
||||
"copilot-proxy.githubusercontent.com": true,
|
||||
}
|
||||
|
||||
// ListModels fetches the list of available models from the Copilot API.
|
||||
// It requires a valid Copilot API token (not the GitHub access token).
|
||||
func (c *CopilotAuth) ListModels(ctx context.Context, apiToken *CopilotAPIToken) ([]CopilotModelEntry, error) {
|
||||
if apiToken == nil || apiToken.Token == "" {
|
||||
return nil, fmt.Errorf("copilot: api token is required for listing models")
|
||||
}
|
||||
|
||||
// Build models URL, validating the endpoint host to prevent SSRF.
|
||||
modelsURL := copilotAPIEndpoint + "/models"
|
||||
if ep := strings.TrimRight(apiToken.Endpoints.API, "/"); ep != "" {
|
||||
parsed, err := url.Parse(ep)
|
||||
if err == nil && parsed.Scheme == "https" && allowedCopilotAPIHosts[parsed.Host] {
|
||||
modelsURL = ep + "/models"
|
||||
} else {
|
||||
log.Warnf("copilot: ignoring untrusted API endpoint %q, using default", ep)
|
||||
}
|
||||
}
|
||||
|
||||
req, err := c.MakeAuthenticatedRequest(ctx, http.MethodGet, modelsURL, nil, apiToken)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("copilot: failed to create models request: %w", err)
|
||||
}
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("copilot: models request failed: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if errClose := resp.Body.Close(); errClose != nil {
|
||||
log.Errorf("copilot list models: close body error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
|
||||
// Limit response body to prevent memory exhaustion.
|
||||
limitedReader := io.LimitReader(resp.Body, maxModelsResponseSize)
|
||||
bodyBytes, err := io.ReadAll(limitedReader)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("copilot: failed to read models response: %w", err)
|
||||
}
|
||||
|
||||
if !isHTTPSuccess(resp.StatusCode) {
|
||||
return nil, fmt.Errorf("copilot: list models failed with status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
var modelsResp CopilotModelsResponse
|
||||
if err = json.Unmarshal(bodyBytes, &modelsResp); err != nil {
|
||||
return nil, fmt.Errorf("copilot: failed to parse models response: %w", err)
|
||||
}
|
||||
|
||||
return modelsResp.Data, nil
|
||||
}
|
||||
|
||||
// ListModelsWithGitHubToken is a convenience method that exchanges a GitHub access token
|
||||
// for a Copilot API token and then fetches the available models.
|
||||
func (c *CopilotAuth) ListModelsWithGitHubToken(ctx context.Context, githubAccessToken string) ([]CopilotModelEntry, error) {
|
||||
apiToken, err := c.GetCopilotAPIToken(ctx, githubAccessToken)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("copilot: failed to get API token for model listing: %w", err)
|
||||
}
|
||||
|
||||
return c.ListModels(ctx, apiToken)
|
||||
}
|
||||
|
||||
// buildChatCompletionURL builds the URL for chat completions API.
|
||||
func buildChatCompletionURL() string {
|
||||
return copilotAPIEndpoint + "/chat/completions"
|
||||
|
||||
@@ -53,7 +53,7 @@ func NewDeviceFlowClient(cfg *config.Config) *DeviceFlowClient {
|
||||
func (c *DeviceFlowClient) RequestDeviceCode(ctx context.Context) (*DeviceCodeResponse, error) {
|
||||
data := url.Values{}
|
||||
data.Set("client_id", copilotClientID)
|
||||
data.Set("scope", "user:email")
|
||||
data.Set("scope", "read:user user:email")
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, copilotDeviceCodeURL, strings.NewReader(data.Encode()))
|
||||
if err != nil {
|
||||
@@ -211,15 +211,25 @@ func (c *DeviceFlowClient) exchangeDeviceCode(ctx context.Context, deviceCode st
|
||||
}, nil
|
||||
}
|
||||
|
||||
// FetchUserInfo retrieves the GitHub username for the authenticated user.
|
||||
func (c *DeviceFlowClient) FetchUserInfo(ctx context.Context, accessToken string) (string, error) {
|
||||
// GitHubUserInfo holds GitHub user profile information.
|
||||
type GitHubUserInfo struct {
|
||||
// Login is the GitHub username.
|
||||
Login string
|
||||
// Email is the primary email address (may be empty if not public).
|
||||
Email string
|
||||
// Name is the display name.
|
||||
Name string
|
||||
}
|
||||
|
||||
// FetchUserInfo retrieves the GitHub user profile for the authenticated user.
|
||||
func (c *DeviceFlowClient) FetchUserInfo(ctx context.Context, accessToken string) (GitHubUserInfo, error) {
|
||||
if accessToken == "" {
|
||||
return "", NewAuthenticationError(ErrUserInfoFailed, fmt.Errorf("access token is empty"))
|
||||
return GitHubUserInfo{}, NewAuthenticationError(ErrUserInfoFailed, fmt.Errorf("access token is empty"))
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, copilotUserInfoURL, nil)
|
||||
if err != nil {
|
||||
return "", NewAuthenticationError(ErrUserInfoFailed, err)
|
||||
return GitHubUserInfo{}, NewAuthenticationError(ErrUserInfoFailed, err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
@@ -227,7 +237,7 @@ func (c *DeviceFlowClient) FetchUserInfo(ctx context.Context, accessToken string
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", NewAuthenticationError(ErrUserInfoFailed, err)
|
||||
return GitHubUserInfo{}, NewAuthenticationError(ErrUserInfoFailed, err)
|
||||
}
|
||||
defer func() {
|
||||
if errClose := resp.Body.Close(); errClose != nil {
|
||||
@@ -237,19 +247,25 @@ func (c *DeviceFlowClient) FetchUserInfo(ctx context.Context, accessToken string
|
||||
|
||||
if !isHTTPSuccess(resp.StatusCode) {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
return "", NewAuthenticationError(ErrUserInfoFailed, fmt.Errorf("status %d: %s", resp.StatusCode, string(bodyBytes)))
|
||||
return GitHubUserInfo{}, NewAuthenticationError(ErrUserInfoFailed, fmt.Errorf("status %d: %s", resp.StatusCode, string(bodyBytes)))
|
||||
}
|
||||
|
||||
var userInfo struct {
|
||||
var raw struct {
|
||||
Login string `json:"login"`
|
||||
Email string `json:"email"`
|
||||
Name string `json:"name"`
|
||||
}
|
||||
if err = json.NewDecoder(resp.Body).Decode(&userInfo); err != nil {
|
||||
return "", NewAuthenticationError(ErrUserInfoFailed, err)
|
||||
if err = json.NewDecoder(resp.Body).Decode(&raw); err != nil {
|
||||
return GitHubUserInfo{}, NewAuthenticationError(ErrUserInfoFailed, err)
|
||||
}
|
||||
|
||||
if userInfo.Login == "" {
|
||||
return "", NewAuthenticationError(ErrUserInfoFailed, fmt.Errorf("empty username"))
|
||||
if raw.Login == "" {
|
||||
return GitHubUserInfo{}, NewAuthenticationError(ErrUserInfoFailed, fmt.Errorf("empty username"))
|
||||
}
|
||||
|
||||
return userInfo.Login, nil
|
||||
return GitHubUserInfo{
|
||||
Login: raw.Login,
|
||||
Email: raw.Email,
|
||||
Name: raw.Name,
|
||||
}, nil
|
||||
}
|
||||
|
||||
213
internal/auth/copilot/oauth_test.go
Normal file
213
internal/auth/copilot/oauth_test.go
Normal file
@@ -0,0 +1,213 @@
|
||||
package copilot
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// roundTripFunc lets us inject a custom transport for testing.
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
|
||||
|
||||
// newTestClient returns an *http.Client whose requests are redirected to the given test server,
|
||||
// regardless of the original URL host.
|
||||
func newTestClient(srv *httptest.Server) *http.Client {
|
||||
return &http.Client{
|
||||
Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
req2 := req.Clone(req.Context())
|
||||
req2.URL.Scheme = "http"
|
||||
req2.URL.Host = strings.TrimPrefix(srv.URL, "http://")
|
||||
return srv.Client().Transport.RoundTrip(req2)
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// TestFetchUserInfo_FullProfile verifies that FetchUserInfo returns login, email, and name.
|
||||
func TestFetchUserInfo_FullProfile(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if !strings.HasPrefix(r.Header.Get("Authorization"), "Bearer ") {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]string{
|
||||
"login": "octocat",
|
||||
"email": "octocat@github.com",
|
||||
"name": "The Octocat",
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := &DeviceFlowClient{httpClient: newTestClient(srv)}
|
||||
info, err := client.FetchUserInfo(context.Background(), "test-token")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if info.Login != "octocat" {
|
||||
t.Errorf("Login: got %q, want %q", info.Login, "octocat")
|
||||
}
|
||||
if info.Email != "octocat@github.com" {
|
||||
t.Errorf("Email: got %q, want %q", info.Email, "octocat@github.com")
|
||||
}
|
||||
if info.Name != "The Octocat" {
|
||||
t.Errorf("Name: got %q, want %q", info.Name, "The Octocat")
|
||||
}
|
||||
}
|
||||
|
||||
// TestFetchUserInfo_EmptyEmail verifies graceful handling when email is absent (private account).
|
||||
func TestFetchUserInfo_EmptyEmail(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
// GitHub returns null for private emails.
|
||||
_, _ = w.Write([]byte(`{"login":"privateuser","email":null,"name":"Private User"}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := &DeviceFlowClient{httpClient: newTestClient(srv)}
|
||||
info, err := client.FetchUserInfo(context.Background(), "test-token")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if info.Login != "privateuser" {
|
||||
t.Errorf("Login: got %q, want %q", info.Login, "privateuser")
|
||||
}
|
||||
if info.Email != "" {
|
||||
t.Errorf("Email: got %q, want empty string", info.Email)
|
||||
}
|
||||
if info.Name != "Private User" {
|
||||
t.Errorf("Name: got %q, want %q", info.Name, "Private User")
|
||||
}
|
||||
}
|
||||
|
||||
// TestFetchUserInfo_EmptyToken verifies error is returned for empty access token.
|
||||
func TestFetchUserInfo_EmptyToken(t *testing.T) {
|
||||
client := &DeviceFlowClient{httpClient: http.DefaultClient}
|
||||
_, err := client.FetchUserInfo(context.Background(), "")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty token, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestFetchUserInfo_EmptyLogin verifies error is returned when API returns no login.
|
||||
func TestFetchUserInfo_EmptyLogin(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"email":"someone@example.com","name":"No Login"}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := &DeviceFlowClient{httpClient: newTestClient(srv)}
|
||||
_, err := client.FetchUserInfo(context.Background(), "test-token")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty login, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestFetchUserInfo_HTTPError verifies error is returned on non-2xx response.
|
||||
func TestFetchUserInfo_HTTPError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
_, _ = w.Write([]byte(`{"message":"Bad credentials"}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := &DeviceFlowClient{httpClient: newTestClient(srv)}
|
||||
_, err := client.FetchUserInfo(context.Background(), "bad-token")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for 401 response, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestCopilotTokenStorage_EmailNameFields verifies Email and Name serialise correctly.
|
||||
func TestCopilotTokenStorage_EmailNameFields(t *testing.T) {
|
||||
ts := &CopilotTokenStorage{
|
||||
AccessToken: "ghu_abc",
|
||||
TokenType: "bearer",
|
||||
Scope: "read:user user:email",
|
||||
Username: "octocat",
|
||||
Email: "octocat@github.com",
|
||||
Name: "The Octocat",
|
||||
Type: "github-copilot",
|
||||
}
|
||||
|
||||
data, err := json.Marshal(ts)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal error: %v", err)
|
||||
}
|
||||
|
||||
var out map[string]any
|
||||
if err = json.Unmarshal(data, &out); err != nil {
|
||||
t.Fatalf("unmarshal error: %v", err)
|
||||
}
|
||||
|
||||
for _, key := range []string{"access_token", "username", "email", "name", "type"} {
|
||||
if _, ok := out[key]; !ok {
|
||||
t.Errorf("expected key %q in JSON output, not found", key)
|
||||
}
|
||||
}
|
||||
if out["email"] != "octocat@github.com" {
|
||||
t.Errorf("email: got %v, want %q", out["email"], "octocat@github.com")
|
||||
}
|
||||
if out["name"] != "The Octocat" {
|
||||
t.Errorf("name: got %v, want %q", out["name"], "The Octocat")
|
||||
}
|
||||
}
|
||||
|
||||
// TestCopilotTokenStorage_OmitEmptyEmailName verifies email/name are omitted when empty (omitempty).
|
||||
func TestCopilotTokenStorage_OmitEmptyEmailName(t *testing.T) {
|
||||
ts := &CopilotTokenStorage{
|
||||
AccessToken: "ghu_abc",
|
||||
Username: "octocat",
|
||||
Type: "github-copilot",
|
||||
}
|
||||
|
||||
data, err := json.Marshal(ts)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal error: %v", err)
|
||||
}
|
||||
|
||||
var out map[string]any
|
||||
if err = json.Unmarshal(data, &out); err != nil {
|
||||
t.Fatalf("unmarshal error: %v", err)
|
||||
}
|
||||
|
||||
if _, ok := out["email"]; ok {
|
||||
t.Error("email key should be omitted when empty (omitempty), but was present")
|
||||
}
|
||||
if _, ok := out["name"]; ok {
|
||||
t.Error("name key should be omitted when empty (omitempty), but was present")
|
||||
}
|
||||
}
|
||||
|
||||
// TestCopilotAuthBundle_EmailNameFields verifies bundle carries email and name through the pipeline.
|
||||
func TestCopilotAuthBundle_EmailNameFields(t *testing.T) {
|
||||
bundle := &CopilotAuthBundle{
|
||||
TokenData: &CopilotTokenData{AccessToken: "ghu_abc"},
|
||||
Username: "octocat",
|
||||
Email: "octocat@github.com",
|
||||
Name: "The Octocat",
|
||||
}
|
||||
if bundle.Email != "octocat@github.com" {
|
||||
t.Errorf("bundle.Email: got %q, want %q", bundle.Email, "octocat@github.com")
|
||||
}
|
||||
if bundle.Name != "The Octocat" {
|
||||
t.Errorf("bundle.Name: got %q, want %q", bundle.Name, "The Octocat")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGitHubUserInfo_Struct verifies the exported GitHubUserInfo struct fields are accessible.
|
||||
func TestGitHubUserInfo_Struct(t *testing.T) {
|
||||
info := GitHubUserInfo{
|
||||
Login: "octocat",
|
||||
Email: "octocat@github.com",
|
||||
Name: "The Octocat",
|
||||
}
|
||||
if info.Login == "" || info.Email == "" || info.Name == "" {
|
||||
t.Error("GitHubUserInfo fields should not be empty")
|
||||
}
|
||||
}
|
||||
@@ -26,6 +26,10 @@ type CopilotTokenStorage struct {
|
||||
ExpiresAt string `json:"expires_at,omitempty"`
|
||||
// Username is the GitHub username associated with this token.
|
||||
Username string `json:"username"`
|
||||
// Email is the GitHub email address associated with this token.
|
||||
Email string `json:"email,omitempty"`
|
||||
// Name is the GitHub display name associated with this token.
|
||||
Name string `json:"name,omitempty"`
|
||||
// Type indicates the authentication provider type, always "github-copilot" for this storage.
|
||||
Type string `json:"type"`
|
||||
}
|
||||
@@ -46,6 +50,10 @@ type CopilotAuthBundle struct {
|
||||
TokenData *CopilotTokenData
|
||||
// Username is the GitHub username.
|
||||
Username string
|
||||
// Email is the GitHub email address.
|
||||
Email string
|
||||
// Name is the GitHub display name.
|
||||
Name string
|
||||
}
|
||||
|
||||
// DeviceCodeResponse represents GitHub's device code response.
|
||||
|
||||
33
internal/auth/cursor/filename.go
Normal file
33
internal/auth/cursor/filename.go
Normal file
@@ -0,0 +1,33 @@
|
||||
package cursor
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// CredentialFileName returns the filename used to persist Cursor credentials.
|
||||
// Priority: explicit label > auto-generated from JWT sub hash.
|
||||
// If both label and subHash are empty, falls back to "cursor.json".
|
||||
func CredentialFileName(label, subHash string) string {
|
||||
label = strings.TrimSpace(label)
|
||||
subHash = strings.TrimSpace(subHash)
|
||||
if label != "" {
|
||||
return fmt.Sprintf("cursor.%s.json", label)
|
||||
}
|
||||
if subHash != "" {
|
||||
return fmt.Sprintf("cursor.%s.json", subHash)
|
||||
}
|
||||
return "cursor.json"
|
||||
}
|
||||
|
||||
// DisplayLabel returns a human-readable label for the Cursor account.
|
||||
func DisplayLabel(label, subHash string) string {
|
||||
label = strings.TrimSpace(label)
|
||||
if label != "" {
|
||||
return "Cursor " + label
|
||||
}
|
||||
if subHash != "" {
|
||||
return "Cursor " + subHash
|
||||
}
|
||||
return "Cursor User"
|
||||
}
|
||||
249
internal/auth/cursor/oauth.go
Normal file
249
internal/auth/cursor/oauth.go
Normal file
@@ -0,0 +1,249 @@
|
||||
// Package cursor implements Cursor OAuth PKCE authentication and token refresh.
|
||||
package cursor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
CursorLoginURL = "https://cursor.com/loginDeepControl"
|
||||
CursorPollURL = "https://api2.cursor.sh/auth/poll"
|
||||
CursorRefreshURL = "https://api2.cursor.sh/auth/exchange_user_api_key"
|
||||
|
||||
pollMaxAttempts = 150
|
||||
pollBaseDelay = 1 * time.Second
|
||||
pollMaxDelay = 10 * time.Second
|
||||
pollBackoffMultiply = 1.2
|
||||
maxConsecutiveErrors = 10
|
||||
)
|
||||
|
||||
// AuthParams holds the PKCE parameters for Cursor login.
|
||||
type AuthParams struct {
|
||||
Verifier string
|
||||
Challenge string
|
||||
UUID string
|
||||
LoginURL string
|
||||
}
|
||||
|
||||
// TokenPair holds the access and refresh tokens from Cursor.
|
||||
type TokenPair struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
RefreshToken string `json:"refreshToken"`
|
||||
}
|
||||
|
||||
// GeneratePKCE creates a PKCE verifier and challenge pair.
|
||||
func GeneratePKCE() (verifier, challenge string, err error) {
|
||||
verifierBytes := make([]byte, 96)
|
||||
if _, err = rand.Read(verifierBytes); err != nil {
|
||||
return "", "", fmt.Errorf("cursor: failed to generate PKCE verifier: %w", err)
|
||||
}
|
||||
verifier = base64.RawURLEncoding.EncodeToString(verifierBytes)
|
||||
|
||||
h := sha256.Sum256([]byte(verifier))
|
||||
challenge = base64.RawURLEncoding.EncodeToString(h[:])
|
||||
return verifier, challenge, nil
|
||||
}
|
||||
|
||||
// GenerateAuthParams creates the full set of auth params for Cursor login.
|
||||
func GenerateAuthParams() (*AuthParams, error) {
|
||||
verifier, challenge, err := GeneratePKCE()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
uuidBytes := make([]byte, 16)
|
||||
if _, err = rand.Read(uuidBytes); err != nil {
|
||||
return nil, fmt.Errorf("cursor: failed to generate UUID: %w", err)
|
||||
}
|
||||
uuid := fmt.Sprintf("%x-%x-%x-%x-%x",
|
||||
uuidBytes[0:4], uuidBytes[4:6], uuidBytes[6:8], uuidBytes[8:10], uuidBytes[10:16])
|
||||
|
||||
loginURL := fmt.Sprintf("%s?challenge=%s&uuid=%s&mode=login&redirectTarget=cli",
|
||||
CursorLoginURL, challenge, uuid)
|
||||
|
||||
return &AuthParams{
|
||||
Verifier: verifier,
|
||||
Challenge: challenge,
|
||||
UUID: uuid,
|
||||
LoginURL: loginURL,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// PollForAuth polls the Cursor auth endpoint until the user completes login.
|
||||
func PollForAuth(ctx context.Context, uuid, verifier string) (*TokenPair, error) {
|
||||
delay := pollBaseDelay
|
||||
consecutiveErrors := 0
|
||||
|
||||
client := &http.Client{Timeout: 10 * time.Second}
|
||||
|
||||
for attempt := 0; attempt < pollMaxAttempts; attempt++ {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case <-time.After(delay):
|
||||
}
|
||||
|
||||
url := fmt.Sprintf("%s?uuid=%s&verifier=%s", CursorPollURL, uuid, verifier)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cursor: failed to create poll request: %w", err)
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
consecutiveErrors++
|
||||
if consecutiveErrors >= maxConsecutiveErrors {
|
||||
return nil, fmt.Errorf("cursor: too many consecutive poll errors (last: %v)", err)
|
||||
}
|
||||
delay = minDuration(time.Duration(float64(delay)*pollBackoffMultiply), pollMaxDelay)
|
||||
continue
|
||||
}
|
||||
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
|
||||
if resp.StatusCode == http.StatusNotFound {
|
||||
// Still waiting for user to authorize
|
||||
consecutiveErrors = 0
|
||||
delay = minDuration(time.Duration(float64(delay)*pollBackoffMultiply), pollMaxDelay)
|
||||
continue
|
||||
}
|
||||
|
||||
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||||
var tokens TokenPair
|
||||
if err := json.Unmarshal(body, &tokens); err != nil {
|
||||
return nil, fmt.Errorf("cursor: failed to parse auth response: %w", err)
|
||||
}
|
||||
return &tokens, nil
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("cursor: poll failed with status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("cursor: authentication polling timeout (waited ~%.0f seconds)",
|
||||
float64(pollMaxAttempts)*pollMaxDelay.Seconds()/2)
|
||||
}
|
||||
|
||||
// RefreshToken refreshes a Cursor access token using the refresh token.
|
||||
func RefreshToken(ctx context.Context, refreshToken string) (*TokenPair, error) {
|
||||
client := &http.Client{Timeout: 10 * time.Second}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, CursorRefreshURL,
|
||||
strings.NewReader("{}"))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cursor: failed to create refresh request: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+refreshToken)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cursor: token refresh request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("cursor: token refresh failed (status %d): %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var tokens TokenPair
|
||||
if err := json.Unmarshal(body, &tokens); err != nil {
|
||||
return nil, fmt.Errorf("cursor: failed to parse refresh response: %w", err)
|
||||
}
|
||||
|
||||
// Keep original refresh token if not returned
|
||||
if tokens.RefreshToken == "" {
|
||||
tokens.RefreshToken = refreshToken
|
||||
}
|
||||
|
||||
return &tokens, nil
|
||||
}
|
||||
|
||||
// ParseJWTSub extracts the "sub" claim from a Cursor JWT access token.
|
||||
// Cursor JWTs contain "sub" like "auth0|user_XXXX" which uniquely identifies
|
||||
// the account. Returns empty string if parsing fails.
|
||||
func ParseJWTSub(token string) string {
|
||||
decoded := decodeJWTPayload(token)
|
||||
if decoded == nil {
|
||||
return ""
|
||||
}
|
||||
var claims struct {
|
||||
Sub string `json:"sub"`
|
||||
}
|
||||
if err := json.Unmarshal(decoded, &claims); err != nil {
|
||||
return ""
|
||||
}
|
||||
return claims.Sub
|
||||
}
|
||||
|
||||
// SubToShortHash converts a JWT sub claim to a short hex hash for use in filenames.
|
||||
// e.g. "auth0|user_2x..." → "a3f8b2c1"
|
||||
func SubToShortHash(sub string) string {
|
||||
if sub == "" {
|
||||
return ""
|
||||
}
|
||||
h := sha256.Sum256([]byte(sub))
|
||||
return fmt.Sprintf("%x", h[:4]) // 8 hex chars
|
||||
}
|
||||
|
||||
// decodeJWTPayload decodes the payload (middle) part of a JWT.
|
||||
func decodeJWTPayload(token string) []byte {
|
||||
parts := strings.Split(token, ".")
|
||||
if len(parts) != 3 {
|
||||
return nil
|
||||
}
|
||||
payload := parts[1]
|
||||
switch len(payload) % 4 {
|
||||
case 2:
|
||||
payload += "=="
|
||||
case 3:
|
||||
payload += "="
|
||||
}
|
||||
payload = strings.ReplaceAll(payload, "-", "+")
|
||||
payload = strings.ReplaceAll(payload, "_", "/")
|
||||
decoded, err := base64.StdEncoding.DecodeString(payload)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return decoded
|
||||
}
|
||||
|
||||
// GetTokenExpiry extracts the JWT expiry from an access token with a 5-minute safety margin.
|
||||
// Falls back to 1 hour from now if the token can't be parsed.
|
||||
func GetTokenExpiry(token string) time.Time {
|
||||
decoded := decodeJWTPayload(token)
|
||||
if decoded == nil {
|
||||
return time.Now().Add(1 * time.Hour)
|
||||
}
|
||||
|
||||
var claims struct {
|
||||
Exp float64 `json:"exp"`
|
||||
}
|
||||
if err := json.Unmarshal(decoded, &claims); err != nil || claims.Exp == 0 {
|
||||
return time.Now().Add(1 * time.Hour)
|
||||
}
|
||||
|
||||
sec, frac := math.Modf(claims.Exp)
|
||||
expiry := time.Unix(int64(sec), int64(frac*1e9))
|
||||
// Subtract 5-minute safety margin
|
||||
return expiry.Add(-5 * time.Minute)
|
||||
}
|
||||
|
||||
func minDuration(a, b time.Duration) time.Duration {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
84
internal/auth/cursor/proto/connect.go
Normal file
84
internal/auth/cursor/proto/connect.go
Normal file
@@ -0,0 +1,84 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
const (
|
||||
// ConnectEndStreamFlag marks the end-of-stream frame (trailers).
|
||||
ConnectEndStreamFlag byte = 0x02
|
||||
// ConnectCompressionFlag indicates the payload is compressed (not supported).
|
||||
ConnectCompressionFlag byte = 0x01
|
||||
// ConnectFrameHeaderSize is the fixed 5-byte frame header.
|
||||
ConnectFrameHeaderSize = 5
|
||||
)
|
||||
|
||||
// FrameConnectMessage wraps a protobuf payload in a Connect frame.
|
||||
// Frame format: [1 byte flags][4 bytes payload length (big-endian)][payload]
|
||||
func FrameConnectMessage(data []byte, flags byte) []byte {
|
||||
frame := make([]byte, ConnectFrameHeaderSize+len(data))
|
||||
frame[0] = flags
|
||||
binary.BigEndian.PutUint32(frame[1:5], uint32(len(data)))
|
||||
copy(frame[5:], data)
|
||||
return frame
|
||||
}
|
||||
|
||||
// ParseConnectFrame extracts one frame from a buffer.
|
||||
// Returns (flags, payload, bytesConsumed, ok).
|
||||
// ok is false when the buffer is too short for a complete frame.
|
||||
func ParseConnectFrame(buf []byte) (flags byte, payload []byte, consumed int, ok bool) {
|
||||
if len(buf) < ConnectFrameHeaderSize {
|
||||
return 0, nil, 0, false
|
||||
}
|
||||
flags = buf[0]
|
||||
length := binary.BigEndian.Uint32(buf[1:5])
|
||||
total := ConnectFrameHeaderSize + int(length)
|
||||
if len(buf) < total {
|
||||
return 0, nil, 0, false
|
||||
}
|
||||
return flags, buf[5:total], total, true
|
||||
}
|
||||
|
||||
// ConnectError is a structured error from the Connect protocol end-of-stream trailer.
|
||||
// The Code field contains the server-defined error code (e.g. gRPC standard codes
|
||||
// like "resource_exhausted", "unauthenticated", "permission_denied", "unavailable").
|
||||
type ConnectError struct {
|
||||
Code string // server-defined error code
|
||||
Message string // human-readable error description
|
||||
}
|
||||
|
||||
func (e *ConnectError) Error() string {
|
||||
return fmt.Sprintf("Connect error %s: %s", e.Code, e.Message)
|
||||
}
|
||||
|
||||
// ParseConnectEndStream parses a Connect end-of-stream frame payload (JSON).
|
||||
// Returns nil if there is no error in the trailer.
|
||||
// On error, returns a *ConnectError with the server's error code and message.
|
||||
func ParseConnectEndStream(data []byte) error {
|
||||
if len(data) == 0 {
|
||||
return nil
|
||||
}
|
||||
var trailer struct {
|
||||
Error *struct {
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &trailer); err != nil {
|
||||
return fmt.Errorf("failed to parse Connect end stream: %w", err)
|
||||
}
|
||||
if trailer.Error != nil {
|
||||
code := trailer.Error.Code
|
||||
if code == "" {
|
||||
code = "unknown"
|
||||
}
|
||||
msg := trailer.Error.Message
|
||||
if msg == "" {
|
||||
msg = "Unknown error"
|
||||
}
|
||||
return &ConnectError{Code: code, Message: msg}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
564
internal/auth/cursor/proto/decode.go
Normal file
564
internal/auth/cursor/proto/decode.go
Normal file
@@ -0,0 +1,564 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"google.golang.org/protobuf/encoding/protowire"
|
||||
)
|
||||
|
||||
// ServerMessageType identifies the kind of decoded server message.
|
||||
type ServerMessageType int
|
||||
|
||||
const (
|
||||
ServerMsgUnknown ServerMessageType = iota
|
||||
ServerMsgTextDelta // Text content delta
|
||||
ServerMsgThinkingDelta // Thinking/reasoning delta
|
||||
ServerMsgThinkingCompleted // Thinking completed
|
||||
ServerMsgKvGetBlob // Server wants a blob
|
||||
ServerMsgKvSetBlob // Server wants to store a blob
|
||||
ServerMsgExecRequestCtx // Server requests context (tools, etc.)
|
||||
ServerMsgExecMcpArgs // Server wants MCP tool execution
|
||||
ServerMsgExecShellArgs // Rejected: shell command
|
||||
ServerMsgExecReadArgs // Rejected: file read
|
||||
ServerMsgExecWriteArgs // Rejected: file write
|
||||
ServerMsgExecDeleteArgs // Rejected: file delete
|
||||
ServerMsgExecLsArgs // Rejected: directory listing
|
||||
ServerMsgExecGrepArgs // Rejected: grep search
|
||||
ServerMsgExecFetchArgs // Rejected: HTTP fetch
|
||||
ServerMsgExecDiagnostics // Respond with empty diagnostics
|
||||
ServerMsgExecShellStream // Rejected: shell stream
|
||||
ServerMsgExecBgShellSpawn // Rejected: background shell
|
||||
ServerMsgExecWriteShellStdin // Rejected: write shell stdin
|
||||
ServerMsgExecOther // Other exec types (respond with empty)
|
||||
ServerMsgTurnEnded // Turn has ended (no more output)
|
||||
ServerMsgHeartbeat // Server heartbeat
|
||||
ServerMsgTokenDelta // Token usage delta
|
||||
ServerMsgCheckpoint // Conversation checkpoint update
|
||||
)
|
||||
|
||||
// DecodedServerMessage holds parsed data from an AgentServerMessage.
|
||||
type DecodedServerMessage struct {
|
||||
Type ServerMessageType
|
||||
|
||||
// For text/thinking deltas
|
||||
Text string
|
||||
|
||||
// For KV messages
|
||||
KvId uint32
|
||||
BlobId []byte // hex-encoded blob ID
|
||||
BlobData []byte // for setBlobArgs
|
||||
|
||||
// For exec messages
|
||||
ExecMsgId uint32
|
||||
ExecId string
|
||||
|
||||
// For MCP args
|
||||
McpToolName string
|
||||
McpToolCallId string
|
||||
McpArgs map[string][]byte // arg name -> protobuf-encoded value
|
||||
|
||||
// For rejection context
|
||||
Path string
|
||||
Command string
|
||||
WorkingDirectory string
|
||||
Url string
|
||||
|
||||
// For other exec - the raw field number for building a response
|
||||
ExecFieldNumber int
|
||||
|
||||
// For TokenDeltaUpdate
|
||||
TokenDelta int64
|
||||
|
||||
// For conversation checkpoint update (raw bytes, not decoded)
|
||||
CheckpointData []byte
|
||||
}
|
||||
|
||||
// DecodeAgentServerMessage parses an AgentServerMessage and returns
|
||||
// a structured representation of the first meaningful message found.
|
||||
func DecodeAgentServerMessage(data []byte) (*DecodedServerMessage, error) {
|
||||
msg := &DecodedServerMessage{Type: ServerMsgUnknown}
|
||||
|
||||
for len(data) > 0 {
|
||||
num, typ, n := protowire.ConsumeTag(data)
|
||||
if n < 0 {
|
||||
return msg, fmt.Errorf("invalid tag")
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
switch typ {
|
||||
case protowire.BytesType:
|
||||
val, n := protowire.ConsumeBytes(data)
|
||||
if n < 0 {
|
||||
return msg, fmt.Errorf("invalid bytes field %d", num)
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
// Debug: log top-level ASM fields
|
||||
log.Debugf("DecodeAgentServerMessage: found ASM field %d, len=%d", num, len(val))
|
||||
|
||||
switch num {
|
||||
case ASM_InteractionUpdate:
|
||||
log.Debugf("DecodeAgentServerMessage: calling decodeInteractionUpdate")
|
||||
decodeInteractionUpdate(val, msg)
|
||||
case ASM_ExecServerMessage:
|
||||
log.Debugf("DecodeAgentServerMessage: calling decodeExecServerMessage")
|
||||
decodeExecServerMessage(val, msg)
|
||||
case ASM_KvServerMessage:
|
||||
decodeKvServerMessage(val, msg)
|
||||
case ASM_ConversationCheckpoint:
|
||||
msg.Type = ServerMsgCheckpoint
|
||||
msg.CheckpointData = append([]byte(nil), val...) // copy raw bytes
|
||||
log.Debugf("DecodeAgentServerMessage: captured checkpoint %d bytes", len(val))
|
||||
}
|
||||
|
||||
case protowire.VarintType:
|
||||
_, n := protowire.ConsumeVarint(data)
|
||||
if n < 0 {
|
||||
return msg, fmt.Errorf("invalid varint field %d", num)
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
default:
|
||||
// Skip unknown wire types
|
||||
n := protowire.ConsumeFieldValue(num, typ, data)
|
||||
if n < 0 {
|
||||
return msg, fmt.Errorf("invalid field %d", num)
|
||||
}
|
||||
data = data[n:]
|
||||
}
|
||||
}
|
||||
|
||||
return msg, nil
|
||||
}
|
||||
|
||||
func decodeInteractionUpdate(data []byte, msg *DecodedServerMessage) {
|
||||
log.Debugf("decodeInteractionUpdate: input len=%d, hex=%x", len(data), data)
|
||||
for len(data) > 0 {
|
||||
num, typ, n := protowire.ConsumeTag(data)
|
||||
if n < 0 {
|
||||
log.Debugf("decodeInteractionUpdate: invalid tag, remaining=%x", data)
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
log.Debugf("decodeInteractionUpdate: field=%d wire=%d remaining=%d bytes", num, typ, len(data))
|
||||
|
||||
if typ == protowire.BytesType {
|
||||
val, n := protowire.ConsumeBytes(data)
|
||||
if n < 0 {
|
||||
log.Debugf("decodeInteractionUpdate: invalid bytes field %d", num)
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
log.Debugf("decodeInteractionUpdate: field %d content len=%d, first 20 bytes: %x", num, len(val), val[:min(20, len(val))])
|
||||
|
||||
switch num {
|
||||
case IU_TextDelta:
|
||||
msg.Type = ServerMsgTextDelta
|
||||
msg.Text = decodeStringField(val, TDU_Text)
|
||||
log.Debugf("decodeInteractionUpdate: TextDelta text=%q", msg.Text)
|
||||
case IU_ThinkingDelta:
|
||||
msg.Type = ServerMsgThinkingDelta
|
||||
msg.Text = decodeStringField(val, TKD_Text)
|
||||
log.Debugf("decodeInteractionUpdate: ThinkingDelta text=%q", msg.Text)
|
||||
case IU_ThinkingCompleted:
|
||||
msg.Type = ServerMsgThinkingCompleted
|
||||
log.Debugf("decodeInteractionUpdate: ThinkingCompleted")
|
||||
case 2:
|
||||
// tool_call_started - ignore but log
|
||||
log.Debugf("decodeInteractionUpdate: ToolCallStarted (ignored)")
|
||||
case 3:
|
||||
// tool_call_completed - ignore but log
|
||||
log.Debugf("decodeInteractionUpdate: ToolCallCompleted (ignored)")
|
||||
case 8:
|
||||
// token_delta - extract token count
|
||||
msg.Type = ServerMsgTokenDelta
|
||||
msg.TokenDelta = decodeVarintField(val, 1)
|
||||
log.Debugf("decodeInteractionUpdate: TokenDeltaUpdate tokens=%d", msg.TokenDelta)
|
||||
case 13:
|
||||
// heartbeat from server
|
||||
msg.Type = ServerMsgHeartbeat
|
||||
case 14:
|
||||
// turn_ended - critical: model finished generating
|
||||
msg.Type = ServerMsgTurnEnded
|
||||
log.Debugf("decodeInteractionUpdate: TurnEndedUpdate - stream should end")
|
||||
case 16:
|
||||
// step_started - ignore
|
||||
log.Debugf("decodeInteractionUpdate: StepStartedUpdate (ignored)")
|
||||
case 17:
|
||||
// step_completed - ignore
|
||||
log.Debugf("decodeInteractionUpdate: StepCompletedUpdate (ignored)")
|
||||
default:
|
||||
log.Debugf("decodeInteractionUpdate: unknown field %d", num)
|
||||
}
|
||||
} else {
|
||||
n := protowire.ConsumeFieldValue(num, typ, data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func decodeKvServerMessage(data []byte, msg *DecodedServerMessage) {
|
||||
for len(data) > 0 {
|
||||
num, typ, n := protowire.ConsumeTag(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
switch typ {
|
||||
case protowire.VarintType:
|
||||
val, n := protowire.ConsumeVarint(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
if num == KSM_Id {
|
||||
msg.KvId = uint32(val)
|
||||
}
|
||||
|
||||
case protowire.BytesType:
|
||||
val, n := protowire.ConsumeBytes(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
switch num {
|
||||
case KSM_GetBlobArgs:
|
||||
msg.Type = ServerMsgKvGetBlob
|
||||
msg.BlobId = decodeBytesField(val, GBA_BlobId)
|
||||
case KSM_SetBlobArgs:
|
||||
msg.Type = ServerMsgKvSetBlob
|
||||
decodeSetBlobArgs(val, msg)
|
||||
}
|
||||
|
||||
default:
|
||||
n := protowire.ConsumeFieldValue(num, typ, data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func decodeSetBlobArgs(data []byte, msg *DecodedServerMessage) {
|
||||
for len(data) > 0 {
|
||||
num, typ, n := protowire.ConsumeTag(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
if typ == protowire.BytesType {
|
||||
val, n := protowire.ConsumeBytes(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
switch num {
|
||||
case SBA_BlobId:
|
||||
msg.BlobId = val
|
||||
case SBA_BlobData:
|
||||
msg.BlobData = val
|
||||
}
|
||||
} else {
|
||||
n := protowire.ConsumeFieldValue(num, typ, data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func decodeExecServerMessage(data []byte, msg *DecodedServerMessage) {
|
||||
for len(data) > 0 {
|
||||
num, typ, n := protowire.ConsumeTag(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
switch typ {
|
||||
case protowire.VarintType:
|
||||
val, n := protowire.ConsumeVarint(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
if num == ESM_Id {
|
||||
msg.ExecMsgId = uint32(val)
|
||||
log.Debugf("decodeExecServerMessage: ESM_Id = %d", val)
|
||||
}
|
||||
|
||||
case protowire.BytesType:
|
||||
val, n := protowire.ConsumeBytes(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
// Debug: log all fields found in ExecServerMessage
|
||||
log.Debugf("decodeExecServerMessage: found field %d, len=%d, first 20 bytes: %x", num, len(val), val[:min(20, len(val))])
|
||||
|
||||
switch num {
|
||||
case ESM_ExecId:
|
||||
msg.ExecId = string(val)
|
||||
log.Debugf("decodeExecServerMessage: ESM_ExecId = %q", msg.ExecId)
|
||||
case ESM_RequestContextArgs:
|
||||
msg.Type = ServerMsgExecRequestCtx
|
||||
case ESM_McpArgs:
|
||||
msg.Type = ServerMsgExecMcpArgs
|
||||
decodeMcpArgs(val, msg)
|
||||
case ESM_ShellArgs:
|
||||
msg.Type = ServerMsgExecShellArgs
|
||||
decodeShellArgs(val, msg)
|
||||
case ESM_ShellStreamArgs:
|
||||
msg.Type = ServerMsgExecShellStream
|
||||
decodeShellArgs(val, msg)
|
||||
case ESM_ReadArgs:
|
||||
msg.Type = ServerMsgExecReadArgs
|
||||
msg.Path = decodeStringField(val, RA_Path)
|
||||
case ESM_WriteArgs:
|
||||
msg.Type = ServerMsgExecWriteArgs
|
||||
msg.Path = decodeStringField(val, WA_Path)
|
||||
case ESM_DeleteArgs:
|
||||
msg.Type = ServerMsgExecDeleteArgs
|
||||
msg.Path = decodeStringField(val, DA_Path)
|
||||
case ESM_LsArgs:
|
||||
msg.Type = ServerMsgExecLsArgs
|
||||
msg.Path = decodeStringField(val, LA_Path)
|
||||
case ESM_GrepArgs:
|
||||
msg.Type = ServerMsgExecGrepArgs
|
||||
case ESM_FetchArgs:
|
||||
msg.Type = ServerMsgExecFetchArgs
|
||||
msg.Url = decodeStringField(val, FA_Url)
|
||||
case ESM_DiagnosticsArgs:
|
||||
msg.Type = ServerMsgExecDiagnostics
|
||||
case ESM_BackgroundShellSpawn:
|
||||
msg.Type = ServerMsgExecBgShellSpawn
|
||||
decodeShellArgs(val, msg) // same structure
|
||||
case ESM_WriteShellStdinArgs:
|
||||
msg.Type = ServerMsgExecWriteShellStdin
|
||||
default:
|
||||
// Unknown exec types - only set if we haven't identified the type yet
|
||||
// (other fields like span_context (19) come after the exec type field)
|
||||
if msg.Type == ServerMsgUnknown {
|
||||
msg.Type = ServerMsgExecOther
|
||||
msg.ExecFieldNumber = int(num)
|
||||
}
|
||||
}
|
||||
|
||||
default:
|
||||
n := protowire.ConsumeFieldValue(num, typ, data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func decodeMcpArgs(data []byte, msg *DecodedServerMessage) {
|
||||
msg.McpArgs = make(map[string][]byte)
|
||||
for len(data) > 0 {
|
||||
num, typ, n := protowire.ConsumeTag(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
if typ == protowire.BytesType {
|
||||
val, n := protowire.ConsumeBytes(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
switch num {
|
||||
case MCA_Name:
|
||||
msg.McpToolName = string(val)
|
||||
case MCA_Args:
|
||||
// Map entries are encoded as submessages with key=1, value=2
|
||||
decodeMapEntry(val, msg.McpArgs)
|
||||
case MCA_ToolCallId:
|
||||
msg.McpToolCallId = string(val)
|
||||
case MCA_ToolName:
|
||||
// ToolName takes precedence if present
|
||||
if msg.McpToolName == "" || string(val) != "" {
|
||||
msg.McpToolName = string(val)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
n := protowire.ConsumeFieldValue(num, typ, data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func decodeMapEntry(data []byte, m map[string][]byte) {
|
||||
var key string
|
||||
var value []byte
|
||||
for len(data) > 0 {
|
||||
num, typ, n := protowire.ConsumeTag(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
if typ == protowire.BytesType {
|
||||
val, n := protowire.ConsumeBytes(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
if num == 1 {
|
||||
key = string(val)
|
||||
} else if num == 2 {
|
||||
value = append([]byte(nil), val...)
|
||||
}
|
||||
} else {
|
||||
n := protowire.ConsumeFieldValue(num, typ, data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
}
|
||||
}
|
||||
if key != "" {
|
||||
m[key] = value
|
||||
}
|
||||
}
|
||||
|
||||
func decodeShellArgs(data []byte, msg *DecodedServerMessage) {
|
||||
for len(data) > 0 {
|
||||
num, typ, n := protowire.ConsumeTag(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
if typ == protowire.BytesType {
|
||||
val, n := protowire.ConsumeBytes(data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
switch num {
|
||||
case SHA_Command:
|
||||
msg.Command = string(val)
|
||||
case SHA_WorkingDirectory:
|
||||
msg.WorkingDirectory = string(val)
|
||||
}
|
||||
} else {
|
||||
n := protowire.ConsumeFieldValue(num, typ, data)
|
||||
if n < 0 {
|
||||
return
|
||||
}
|
||||
data = data[n:]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- Helper decoders ---
|
||||
|
||||
// decodeStringField extracts a string from the first matching field in a submessage.
|
||||
func decodeStringField(data []byte, targetField protowire.Number) string {
|
||||
for len(data) > 0 {
|
||||
num, typ, n := protowire.ConsumeTag(data)
|
||||
if n < 0 {
|
||||
return ""
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
if typ == protowire.BytesType {
|
||||
val, n := protowire.ConsumeBytes(data)
|
||||
if n < 0 {
|
||||
return ""
|
||||
}
|
||||
data = data[n:]
|
||||
if num == targetField {
|
||||
return string(val)
|
||||
}
|
||||
} else {
|
||||
n := protowire.ConsumeFieldValue(num, typ, data)
|
||||
if n < 0 {
|
||||
return ""
|
||||
}
|
||||
data = data[n:]
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// decodeBytesField extracts bytes from the first matching field in a submessage.
|
||||
func decodeBytesField(data []byte, targetField protowire.Number) []byte {
|
||||
for len(data) > 0 {
|
||||
num, typ, n := protowire.ConsumeTag(data)
|
||||
if n < 0 {
|
||||
return nil
|
||||
}
|
||||
data = data[n:]
|
||||
|
||||
if typ == protowire.BytesType {
|
||||
val, n := protowire.ConsumeBytes(data)
|
||||
if n < 0 {
|
||||
return nil
|
||||
}
|
||||
data = data[n:]
|
||||
if num == targetField {
|
||||
return append([]byte(nil), val...)
|
||||
}
|
||||
} else {
|
||||
n := protowire.ConsumeFieldValue(num, typ, data)
|
||||
if n < 0 {
|
||||
return nil
|
||||
}
|
||||
data = data[n:]
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// decodeVarintField extracts an int64 from the first matching varint field in a submessage.
|
||||
func decodeVarintField(data []byte, targetField protowire.Number) int64 {
|
||||
for len(data) > 0 {
|
||||
num, typ, n := protowire.ConsumeTag(data)
|
||||
if n < 0 {
|
||||
return 0
|
||||
}
|
||||
data = data[n:]
|
||||
if typ == protowire.VarintType {
|
||||
val, n := protowire.ConsumeVarint(data)
|
||||
if n < 0 {
|
||||
return 0
|
||||
}
|
||||
data = data[n:]
|
||||
if num == targetField {
|
||||
return int64(val)
|
||||
}
|
||||
} else {
|
||||
n := protowire.ConsumeFieldValue(num, typ, data)
|
||||
if n < 0 {
|
||||
return 0
|
||||
}
|
||||
data = data[n:]
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// BlobIdHex returns the hex string of a blob ID for use as a map key.
|
||||
func BlobIdHex(blobId []byte) string {
|
||||
return hex.EncodeToString(blobId)
|
||||
}
|
||||
|
||||
1244
internal/auth/cursor/proto/descriptor.go
Normal file
1244
internal/auth/cursor/proto/descriptor.go
Normal file
File diff suppressed because it is too large
Load Diff
664
internal/auth/cursor/proto/encode.go
Normal file
664
internal/auth/cursor/proto/encode.go
Normal file
@@ -0,0 +1,664 @@
|
||||
// Package proto provides protobuf encoding for Cursor's gRPC API,
|
||||
// using dynamicpb with the embedded FileDescriptorProto from agent.proto.
|
||||
// This mirrors the cursor-auth TS plugin's use of @bufbuild/protobuf create()+toBinary().
|
||||
package proto
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"google.golang.org/protobuf/encoding/protowire"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/reflect/protoreflect"
|
||||
"google.golang.org/protobuf/types/dynamicpb"
|
||||
"google.golang.org/protobuf/types/known/structpb"
|
||||
)
|
||||
|
||||
// --- Public types ---
|
||||
|
||||
// RunRequestParams holds all data needed to build an AgentRunRequest.
|
||||
type RunRequestParams struct {
|
||||
ModelId string
|
||||
SystemPrompt string
|
||||
UserText string
|
||||
MessageId string
|
||||
ConversationId string
|
||||
Images []ImageData
|
||||
Turns []TurnData
|
||||
McpTools []McpToolDef
|
||||
BlobStore map[string][]byte // hex(sha256) -> data, populated during encoding
|
||||
RawCheckpoint []byte // if non-nil, use as conversation_state directly (from server checkpoint)
|
||||
}
|
||||
|
||||
type ImageData struct {
|
||||
MimeType string
|
||||
Data []byte
|
||||
}
|
||||
|
||||
type TurnData struct {
|
||||
UserText string
|
||||
AssistantText string
|
||||
}
|
||||
|
||||
type McpToolDef struct {
|
||||
Name string
|
||||
Description string
|
||||
InputSchema json.RawMessage
|
||||
}
|
||||
|
||||
// --- Helper: create a dynamic message and set fields ---
|
||||
|
||||
func newMsg(name string) *dynamicpb.Message {
|
||||
return dynamicpb.NewMessage(Msg(name))
|
||||
}
|
||||
|
||||
func field(msg *dynamicpb.Message, name string) protoreflect.FieldDescriptor {
|
||||
return msg.Descriptor().Fields().ByName(protoreflect.Name(name))
|
||||
}
|
||||
|
||||
func setStr(msg *dynamicpb.Message, name, val string) {
|
||||
if val != "" {
|
||||
msg.Set(field(msg, name), protoreflect.ValueOfString(val))
|
||||
}
|
||||
}
|
||||
|
||||
func setBytes(msg *dynamicpb.Message, name string, val []byte) {
|
||||
if len(val) > 0 {
|
||||
msg.Set(field(msg, name), protoreflect.ValueOfBytes(val))
|
||||
}
|
||||
}
|
||||
|
||||
func setUint32(msg *dynamicpb.Message, name string, val uint32) {
|
||||
msg.Set(field(msg, name), protoreflect.ValueOfUint32(val))
|
||||
}
|
||||
|
||||
func setBool(msg *dynamicpb.Message, name string, val bool) {
|
||||
msg.Set(field(msg, name), protoreflect.ValueOfBool(val))
|
||||
}
|
||||
|
||||
func setMsg(msg *dynamicpb.Message, name string, sub *dynamicpb.Message) {
|
||||
msg.Set(field(msg, name), protoreflect.ValueOfMessage(sub.ProtoReflect()))
|
||||
}
|
||||
|
||||
func marshal(msg *dynamicpb.Message) []byte {
|
||||
b, err := proto.Marshal(msg)
|
||||
if err != nil {
|
||||
panic("cursor proto marshal: " + err.Error())
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// --- Encode functions mirroring cursor-fetch.ts ---
|
||||
|
||||
// EncodeHeartbeat returns an encoded AgentClientMessage with clientHeartbeat.
|
||||
// Mirrors: create(AgentClientMessageSchema, { message: { case: 'clientHeartbeat', value: create(ClientHeartbeatSchema, {}) } })
|
||||
func EncodeHeartbeat() []byte {
|
||||
hb := newMsg("ClientHeartbeat")
|
||||
acm := newMsg("AgentClientMessage")
|
||||
setMsg(acm, "client_heartbeat", hb)
|
||||
return marshal(acm)
|
||||
}
|
||||
|
||||
// EncodeRunRequest builds a full AgentClientMessage wrapping an AgentRunRequest.
|
||||
// Mirrors buildCursorRequest() in cursor-fetch.ts.
|
||||
// If p.RawCheckpoint is set, it is used directly as the conversation_state bytes
|
||||
// (from a previous conversation_checkpoint_update), skipping manual turn construction.
|
||||
func EncodeRunRequest(p *RunRequestParams) []byte {
|
||||
if p.RawCheckpoint != nil {
|
||||
return encodeRunRequestWithCheckpoint(p)
|
||||
}
|
||||
|
||||
if p.BlobStore == nil {
|
||||
p.BlobStore = make(map[string][]byte)
|
||||
}
|
||||
|
||||
// --- Conversation turns ---
|
||||
// Each turn is serialized as bytes (ConversationTurnStructure → bytes)
|
||||
var turnBytes [][]byte
|
||||
for _, turn := range p.Turns {
|
||||
// UserMessage for this turn
|
||||
um := newMsg("UserMessage")
|
||||
setStr(um, "text", turn.UserText)
|
||||
setStr(um, "message_id", generateId())
|
||||
umBytes := marshal(um)
|
||||
|
||||
// Steps (assistant response)
|
||||
var stepBytes [][]byte
|
||||
if turn.AssistantText != "" {
|
||||
am := newMsg("AssistantMessage")
|
||||
setStr(am, "text", turn.AssistantText)
|
||||
step := newMsg("ConversationStep")
|
||||
setMsg(step, "assistant_message", am)
|
||||
stepBytes = append(stepBytes, marshal(step))
|
||||
}
|
||||
|
||||
// AgentConversationTurnStructure (fields are bytes, not submessages)
|
||||
agentTurn := newMsg("AgentConversationTurnStructure")
|
||||
setBytes(agentTurn, "user_message", umBytes)
|
||||
for _, sb := range stepBytes {
|
||||
stepsField := field(agentTurn, "steps")
|
||||
list := agentTurn.Mutable(stepsField).List()
|
||||
list.Append(protoreflect.ValueOfBytes(sb))
|
||||
}
|
||||
|
||||
// ConversationTurnStructure (oneof turn → agentConversationTurn)
|
||||
cts := newMsg("ConversationTurnStructure")
|
||||
setMsg(cts, "agent_conversation_turn", agentTurn)
|
||||
turnBytes = append(turnBytes, marshal(cts))
|
||||
}
|
||||
|
||||
// --- System prompt blob ---
|
||||
systemJSON, _ := json.Marshal(map[string]string{"role": "system", "content": p.SystemPrompt})
|
||||
blobId := sha256Sum(systemJSON)
|
||||
p.BlobStore[hex.EncodeToString(blobId)] = systemJSON
|
||||
|
||||
// --- ConversationStateStructure ---
|
||||
css := newMsg("ConversationStateStructure")
|
||||
// rootPromptMessagesJson: repeated bytes
|
||||
rootField := field(css, "root_prompt_messages_json")
|
||||
rootList := css.Mutable(rootField).List()
|
||||
rootList.Append(protoreflect.ValueOfBytes(blobId))
|
||||
// turns: repeated bytes (field 8) + turns_old (field 2) for compatibility
|
||||
turnsField := field(css, "turns")
|
||||
turnsList := css.Mutable(turnsField).List()
|
||||
for _, tb := range turnBytes {
|
||||
turnsList.Append(protoreflect.ValueOfBytes(tb))
|
||||
}
|
||||
turnsOldField := field(css, "turns_old")
|
||||
if turnsOldField != nil {
|
||||
turnsOldList := css.Mutable(turnsOldField).List()
|
||||
for _, tb := range turnBytes {
|
||||
turnsOldList.Append(protoreflect.ValueOfBytes(tb))
|
||||
}
|
||||
}
|
||||
|
||||
// --- UserMessage (current) ---
|
||||
userMessage := newMsg("UserMessage")
|
||||
setStr(userMessage, "text", p.UserText)
|
||||
setStr(userMessage, "message_id", p.MessageId)
|
||||
|
||||
// Images via SelectedContext
|
||||
if len(p.Images) > 0 {
|
||||
sc := newMsg("SelectedContext")
|
||||
imgsField := field(sc, "selected_images")
|
||||
imgsList := sc.Mutable(imgsField).List()
|
||||
for _, img := range p.Images {
|
||||
si := newMsg("SelectedImage")
|
||||
setStr(si, "uuid", generateId())
|
||||
setStr(si, "mime_type", img.MimeType)
|
||||
setBytes(si, "data", img.Data)
|
||||
imgsList.Append(protoreflect.ValueOfMessage(si.ProtoReflect()))
|
||||
}
|
||||
setMsg(userMessage, "selected_context", sc)
|
||||
}
|
||||
|
||||
// --- UserMessageAction ---
|
||||
uma := newMsg("UserMessageAction")
|
||||
setMsg(uma, "user_message", userMessage)
|
||||
|
||||
// --- ConversationAction ---
|
||||
ca := newMsg("ConversationAction")
|
||||
setMsg(ca, "user_message_action", uma)
|
||||
|
||||
// --- ModelDetails ---
|
||||
md := newMsg("ModelDetails")
|
||||
setStr(md, "model_id", p.ModelId)
|
||||
setStr(md, "display_model_id", p.ModelId)
|
||||
setStr(md, "display_name", p.ModelId)
|
||||
|
||||
// --- AgentRunRequest ---
|
||||
arr := newMsg("AgentRunRequest")
|
||||
setMsg(arr, "conversation_state", css)
|
||||
setMsg(arr, "action", ca)
|
||||
setMsg(arr, "model_details", md)
|
||||
setStr(arr, "conversation_id", p.ConversationId)
|
||||
|
||||
// McpTools
|
||||
if len(p.McpTools) > 0 {
|
||||
mcpTools := newMsg("McpTools")
|
||||
toolsField := field(mcpTools, "mcp_tools")
|
||||
toolsList := mcpTools.Mutable(toolsField).List()
|
||||
for _, tool := range p.McpTools {
|
||||
td := newMsg("McpToolDefinition")
|
||||
setStr(td, "name", tool.Name)
|
||||
setStr(td, "description", tool.Description)
|
||||
if len(tool.InputSchema) > 0 {
|
||||
setBytes(td, "input_schema", jsonToProtobufValueBytes(tool.InputSchema))
|
||||
}
|
||||
setStr(td, "provider_identifier", "proxy")
|
||||
setStr(td, "tool_name", tool.Name)
|
||||
toolsList.Append(protoreflect.ValueOfMessage(td.ProtoReflect()))
|
||||
}
|
||||
setMsg(arr, "mcp_tools", mcpTools)
|
||||
}
|
||||
|
||||
// --- AgentClientMessage ---
|
||||
acm := newMsg("AgentClientMessage")
|
||||
setMsg(acm, "run_request", arr)
|
||||
|
||||
return marshal(acm)
|
||||
}
|
||||
|
||||
// encodeRunRequestWithCheckpoint builds an AgentClientMessage using a raw checkpoint
|
||||
// as conversation_state. The checkpoint bytes are embedded directly without deserialization.
|
||||
func encodeRunRequestWithCheckpoint(p *RunRequestParams) []byte {
|
||||
// Build UserMessage
|
||||
userMessage := newMsg("UserMessage")
|
||||
setStr(userMessage, "text", p.UserText)
|
||||
setStr(userMessage, "message_id", p.MessageId)
|
||||
if len(p.Images) > 0 {
|
||||
sc := newMsg("SelectedContext")
|
||||
imgsField := field(sc, "selected_images")
|
||||
imgsList := sc.Mutable(imgsField).List()
|
||||
for _, img := range p.Images {
|
||||
si := newMsg("SelectedImage")
|
||||
setStr(si, "uuid", generateId())
|
||||
setStr(si, "mime_type", img.MimeType)
|
||||
setBytes(si, "data", img.Data)
|
||||
imgsList.Append(protoreflect.ValueOfMessage(si.ProtoReflect()))
|
||||
}
|
||||
setMsg(userMessage, "selected_context", sc)
|
||||
}
|
||||
|
||||
// Build ConversationAction with UserMessageAction
|
||||
uma := newMsg("UserMessageAction")
|
||||
setMsg(uma, "user_message", userMessage)
|
||||
ca := newMsg("ConversationAction")
|
||||
setMsg(ca, "user_message_action", uma)
|
||||
caBytes := marshal(ca)
|
||||
|
||||
// Build ModelDetails
|
||||
md := newMsg("ModelDetails")
|
||||
setStr(md, "model_id", p.ModelId)
|
||||
setStr(md, "display_model_id", p.ModelId)
|
||||
setStr(md, "display_name", p.ModelId)
|
||||
mdBytes := marshal(md)
|
||||
|
||||
// Build McpTools
|
||||
var mcpToolsBytes []byte
|
||||
if len(p.McpTools) > 0 {
|
||||
mcpTools := newMsg("McpTools")
|
||||
toolsField := field(mcpTools, "mcp_tools")
|
||||
toolsList := mcpTools.Mutable(toolsField).List()
|
||||
for _, tool := range p.McpTools {
|
||||
td := newMsg("McpToolDefinition")
|
||||
setStr(td, "name", tool.Name)
|
||||
setStr(td, "description", tool.Description)
|
||||
if len(tool.InputSchema) > 0 {
|
||||
setBytes(td, "input_schema", jsonToProtobufValueBytes(tool.InputSchema))
|
||||
}
|
||||
setStr(td, "provider_identifier", "proxy")
|
||||
setStr(td, "tool_name", tool.Name)
|
||||
toolsList.Append(protoreflect.ValueOfMessage(td.ProtoReflect()))
|
||||
}
|
||||
mcpToolsBytes = marshal(mcpTools)
|
||||
}
|
||||
|
||||
// Manually assemble AgentRunRequest using protowire to embed raw checkpoint
|
||||
var arrBuf []byte
|
||||
// field 1: conversation_state = raw checkpoint bytes (length-delimited)
|
||||
arrBuf = protowire.AppendTag(arrBuf, ARR_ConversationState, protowire.BytesType)
|
||||
arrBuf = protowire.AppendBytes(arrBuf, p.RawCheckpoint)
|
||||
// field 2: action = ConversationAction
|
||||
arrBuf = protowire.AppendTag(arrBuf, ARR_Action, protowire.BytesType)
|
||||
arrBuf = protowire.AppendBytes(arrBuf, caBytes)
|
||||
// field 3: model_details = ModelDetails
|
||||
arrBuf = protowire.AppendTag(arrBuf, ARR_ModelDetails, protowire.BytesType)
|
||||
arrBuf = protowire.AppendBytes(arrBuf, mdBytes)
|
||||
// field 4: mcp_tools = McpTools
|
||||
if len(mcpToolsBytes) > 0 {
|
||||
arrBuf = protowire.AppendTag(arrBuf, ARR_McpTools, protowire.BytesType)
|
||||
arrBuf = protowire.AppendBytes(arrBuf, mcpToolsBytes)
|
||||
}
|
||||
// field 5: conversation_id = string
|
||||
if p.ConversationId != "" {
|
||||
arrBuf = protowire.AppendTag(arrBuf, ARR_ConversationId, protowire.BytesType)
|
||||
arrBuf = protowire.AppendString(arrBuf, p.ConversationId)
|
||||
}
|
||||
|
||||
// Wrap in AgentClientMessage field 1 (run_request)
|
||||
var acmBuf []byte
|
||||
acmBuf = protowire.AppendTag(acmBuf, ACM_RunRequest, protowire.BytesType)
|
||||
acmBuf = protowire.AppendBytes(acmBuf, arrBuf)
|
||||
|
||||
log.Debugf("cursor encode: built RunRequest with checkpoint (%d bytes), total=%d bytes", len(p.RawCheckpoint), len(acmBuf))
|
||||
return acmBuf
|
||||
}
|
||||
|
||||
// ResumeRequestParams holds data for a ResumeAction request.
|
||||
type ResumeRequestParams struct {
|
||||
ModelId string
|
||||
ConversationId string
|
||||
McpTools []McpToolDef
|
||||
}
|
||||
|
||||
// EncodeResumeRequest builds an AgentClientMessage with ResumeAction.
|
||||
// Used to resume a conversation by conversation_id without re-sending full history.
|
||||
func EncodeResumeRequest(p *ResumeRequestParams) []byte {
|
||||
// RequestContext with tools
|
||||
rc := newMsg("RequestContext")
|
||||
if len(p.McpTools) > 0 {
|
||||
toolsField := field(rc, "tools")
|
||||
toolsList := rc.Mutable(toolsField).List()
|
||||
for _, tool := range p.McpTools {
|
||||
td := newMsg("McpToolDefinition")
|
||||
setStr(td, "name", tool.Name)
|
||||
setStr(td, "description", tool.Description)
|
||||
if len(tool.InputSchema) > 0 {
|
||||
setBytes(td, "input_schema", jsonToProtobufValueBytes(tool.InputSchema))
|
||||
}
|
||||
setStr(td, "provider_identifier", "proxy")
|
||||
setStr(td, "tool_name", tool.Name)
|
||||
toolsList.Append(protoreflect.ValueOfMessage(td.ProtoReflect()))
|
||||
}
|
||||
}
|
||||
|
||||
// ResumeAction
|
||||
ra := newMsg("ResumeAction")
|
||||
setMsg(ra, "request_context", rc)
|
||||
|
||||
// ConversationAction with resume_action
|
||||
ca := newMsg("ConversationAction")
|
||||
setMsg(ca, "resume_action", ra)
|
||||
|
||||
// ModelDetails
|
||||
md := newMsg("ModelDetails")
|
||||
setStr(md, "model_id", p.ModelId)
|
||||
setStr(md, "display_model_id", p.ModelId)
|
||||
setStr(md, "display_name", p.ModelId)
|
||||
|
||||
// AgentRunRequest — no conversation_state needed for resume
|
||||
arr := newMsg("AgentRunRequest")
|
||||
setMsg(arr, "action", ca)
|
||||
setMsg(arr, "model_details", md)
|
||||
setStr(arr, "conversation_id", p.ConversationId)
|
||||
|
||||
// McpTools at top level
|
||||
if len(p.McpTools) > 0 {
|
||||
mcpTools := newMsg("McpTools")
|
||||
toolsField := field(mcpTools, "mcp_tools")
|
||||
toolsList := mcpTools.Mutable(toolsField).List()
|
||||
for _, tool := range p.McpTools {
|
||||
td := newMsg("McpToolDefinition")
|
||||
setStr(td, "name", tool.Name)
|
||||
setStr(td, "description", tool.Description)
|
||||
if len(tool.InputSchema) > 0 {
|
||||
setBytes(td, "input_schema", jsonToProtobufValueBytes(tool.InputSchema))
|
||||
}
|
||||
setStr(td, "provider_identifier", "proxy")
|
||||
setStr(td, "tool_name", tool.Name)
|
||||
toolsList.Append(protoreflect.ValueOfMessage(td.ProtoReflect()))
|
||||
}
|
||||
setMsg(arr, "mcp_tools", mcpTools)
|
||||
}
|
||||
|
||||
acm := newMsg("AgentClientMessage")
|
||||
setMsg(acm, "run_request", arr)
|
||||
return marshal(acm)
|
||||
}
|
||||
|
||||
// --- KV response encoders ---
|
||||
// Mirrors handleKvMessage() in cursor-fetch.ts
|
||||
|
||||
// EncodeKvGetBlobResult responds to a getBlobArgs request.
|
||||
func EncodeKvGetBlobResult(kvId uint32, blobData []byte) []byte {
|
||||
result := newMsg("GetBlobResult")
|
||||
if blobData != nil {
|
||||
setBytes(result, "blob_data", blobData)
|
||||
}
|
||||
|
||||
kvc := newMsg("KvClientMessage")
|
||||
setUint32(kvc, "id", kvId)
|
||||
setMsg(kvc, "get_blob_result", result)
|
||||
|
||||
acm := newMsg("AgentClientMessage")
|
||||
setMsg(acm, "kv_client_message", kvc)
|
||||
return marshal(acm)
|
||||
}
|
||||
|
||||
// EncodeKvSetBlobResult responds to a setBlobArgs request.
|
||||
func EncodeKvSetBlobResult(kvId uint32) []byte {
|
||||
result := newMsg("SetBlobResult")
|
||||
|
||||
kvc := newMsg("KvClientMessage")
|
||||
setUint32(kvc, "id", kvId)
|
||||
setMsg(kvc, "set_blob_result", result)
|
||||
|
||||
acm := newMsg("AgentClientMessage")
|
||||
setMsg(acm, "kv_client_message", kvc)
|
||||
return marshal(acm)
|
||||
}
|
||||
|
||||
// --- Exec response encoders ---
|
||||
// Mirrors handleExecMessage() and sendExec() in cursor-fetch.ts
|
||||
|
||||
// EncodeExecRequestContextResult responds to requestContextArgs with tool definitions.
|
||||
func EncodeExecRequestContextResult(execMsgId uint32, execId string, tools []McpToolDef) []byte {
|
||||
// RequestContext with tools
|
||||
rc := newMsg("RequestContext")
|
||||
if len(tools) > 0 {
|
||||
toolsField := field(rc, "tools")
|
||||
toolsList := rc.Mutable(toolsField).List()
|
||||
for _, tool := range tools {
|
||||
td := newMsg("McpToolDefinition")
|
||||
setStr(td, "name", tool.Name)
|
||||
setStr(td, "description", tool.Description)
|
||||
if len(tool.InputSchema) > 0 {
|
||||
setBytes(td, "input_schema", jsonToProtobufValueBytes(tool.InputSchema))
|
||||
}
|
||||
setStr(td, "provider_identifier", "proxy")
|
||||
setStr(td, "tool_name", tool.Name)
|
||||
toolsList.Append(protoreflect.ValueOfMessage(td.ProtoReflect()))
|
||||
}
|
||||
}
|
||||
|
||||
// RequestContextSuccess
|
||||
rcs := newMsg("RequestContextSuccess")
|
||||
setMsg(rcs, "request_context", rc)
|
||||
|
||||
// RequestContextResult (oneof success)
|
||||
rcr := newMsg("RequestContextResult")
|
||||
setMsg(rcr, "success", rcs)
|
||||
|
||||
return encodeExecClientMsg(execMsgId, execId, "request_context_result", rcr)
|
||||
}
|
||||
|
||||
// EncodeExecMcpResult responds with MCP tool result.
|
||||
func EncodeExecMcpResult(execMsgId uint32, execId string, content string, isError bool) []byte {
|
||||
textContent := newMsg("McpTextContent")
|
||||
setStr(textContent, "text", content)
|
||||
|
||||
contentItem := newMsg("McpToolResultContentItem")
|
||||
setMsg(contentItem, "text", textContent)
|
||||
|
||||
success := newMsg("McpSuccess")
|
||||
contentField := field(success, "content")
|
||||
contentList := success.Mutable(contentField).List()
|
||||
contentList.Append(protoreflect.ValueOfMessage(contentItem.ProtoReflect()))
|
||||
setBool(success, "is_error", isError)
|
||||
|
||||
result := newMsg("McpResult")
|
||||
setMsg(result, "success", success)
|
||||
|
||||
return encodeExecClientMsg(execMsgId, execId, "mcp_result", result)
|
||||
}
|
||||
|
||||
// EncodeExecMcpError responds with MCP error.
|
||||
func EncodeExecMcpError(execMsgId uint32, execId string, errMsg string) []byte {
|
||||
mcpErr := newMsg("McpError")
|
||||
setStr(mcpErr, "error", errMsg)
|
||||
|
||||
result := newMsg("McpResult")
|
||||
setMsg(result, "error", mcpErr)
|
||||
|
||||
return encodeExecClientMsg(execMsgId, execId, "mcp_result", result)
|
||||
}
|
||||
|
||||
// --- Rejection encoders (mirror handleExecMessage rejections) ---
|
||||
|
||||
func EncodeExecReadRejected(execMsgId uint32, execId string, path, reason string) []byte {
|
||||
rej := newMsg("ReadRejected")
|
||||
setStr(rej, "path", path)
|
||||
setStr(rej, "reason", reason)
|
||||
result := newMsg("ReadResult")
|
||||
setMsg(result, "rejected", rej)
|
||||
return encodeExecClientMsg(execMsgId, execId, "read_result", result)
|
||||
}
|
||||
|
||||
func EncodeExecShellRejected(execMsgId uint32, execId string, command, workDir, reason string) []byte {
|
||||
rej := newMsg("ShellRejected")
|
||||
setStr(rej, "command", command)
|
||||
setStr(rej, "working_directory", workDir)
|
||||
setStr(rej, "reason", reason)
|
||||
result := newMsg("ShellResult")
|
||||
setMsg(result, "rejected", rej)
|
||||
return encodeExecClientMsg(execMsgId, execId, "shell_result", result)
|
||||
}
|
||||
|
||||
func EncodeExecWriteRejected(execMsgId uint32, execId string, path, reason string) []byte {
|
||||
rej := newMsg("WriteRejected")
|
||||
setStr(rej, "path", path)
|
||||
setStr(rej, "reason", reason)
|
||||
result := newMsg("WriteResult")
|
||||
setMsg(result, "rejected", rej)
|
||||
return encodeExecClientMsg(execMsgId, execId, "write_result", result)
|
||||
}
|
||||
|
||||
func EncodeExecDeleteRejected(execMsgId uint32, execId string, path, reason string) []byte {
|
||||
rej := newMsg("DeleteRejected")
|
||||
setStr(rej, "path", path)
|
||||
setStr(rej, "reason", reason)
|
||||
result := newMsg("DeleteResult")
|
||||
setMsg(result, "rejected", rej)
|
||||
return encodeExecClientMsg(execMsgId, execId, "delete_result", result)
|
||||
}
|
||||
|
||||
func EncodeExecLsRejected(execMsgId uint32, execId string, path, reason string) []byte {
|
||||
rej := newMsg("LsRejected")
|
||||
setStr(rej, "path", path)
|
||||
setStr(rej, "reason", reason)
|
||||
result := newMsg("LsResult")
|
||||
setMsg(result, "rejected", rej)
|
||||
return encodeExecClientMsg(execMsgId, execId, "ls_result", result)
|
||||
}
|
||||
|
||||
func EncodeExecGrepError(execMsgId uint32, execId string, errMsg string) []byte {
|
||||
grepErr := newMsg("GrepError")
|
||||
setStr(grepErr, "error", errMsg)
|
||||
result := newMsg("GrepResult")
|
||||
setMsg(result, "error", grepErr)
|
||||
return encodeExecClientMsg(execMsgId, execId, "grep_result", result)
|
||||
}
|
||||
|
||||
func EncodeExecFetchError(execMsgId uint32, execId string, url, errMsg string) []byte {
|
||||
fetchErr := newMsg("FetchError")
|
||||
setStr(fetchErr, "url", url)
|
||||
setStr(fetchErr, "error", errMsg)
|
||||
result := newMsg("FetchResult")
|
||||
setMsg(result, "error", fetchErr)
|
||||
return encodeExecClientMsg(execMsgId, execId, "fetch_result", result)
|
||||
}
|
||||
|
||||
func EncodeExecDiagnosticsResult(execMsgId uint32, execId string) []byte {
|
||||
result := newMsg("DiagnosticsResult")
|
||||
return encodeExecClientMsg(execMsgId, execId, "diagnostics_result", result)
|
||||
}
|
||||
|
||||
func EncodeExecBackgroundShellSpawnRejected(execMsgId uint32, execId string, command, workDir, reason string) []byte {
|
||||
rej := newMsg("ShellRejected")
|
||||
setStr(rej, "command", command)
|
||||
setStr(rej, "working_directory", workDir)
|
||||
setStr(rej, "reason", reason)
|
||||
result := newMsg("BackgroundShellSpawnResult")
|
||||
setMsg(result, "rejected", rej)
|
||||
return encodeExecClientMsg(execMsgId, execId, "background_shell_spawn_result", result)
|
||||
}
|
||||
|
||||
func EncodeExecWriteShellStdinError(execMsgId uint32, execId string, errMsg string) []byte {
|
||||
wsErr := newMsg("WriteShellStdinError")
|
||||
setStr(wsErr, "error", errMsg)
|
||||
result := newMsg("WriteShellStdinResult")
|
||||
setMsg(result, "error", wsErr)
|
||||
return encodeExecClientMsg(execMsgId, execId, "write_shell_stdin_result", result)
|
||||
}
|
||||
|
||||
// encodeExecClientMsg wraps an exec result in AgentClientMessage.
|
||||
// Mirrors sendExec() in cursor-fetch.ts.
|
||||
func encodeExecClientMsg(id uint32, execId string, resultFieldName string, resultMsg *dynamicpb.Message) []byte {
|
||||
ecm := newMsg("ExecClientMessage")
|
||||
setUint32(ecm, "id", id)
|
||||
// Force set exec_id even if empty - Cursor requires this field to be set
|
||||
ecm.Set(field(ecm, "exec_id"), protoreflect.ValueOfString(execId))
|
||||
|
||||
// Debug: check if field exists
|
||||
fd := field(ecm, resultFieldName)
|
||||
if fd == nil {
|
||||
panic(fmt.Sprintf("field %q NOT FOUND in ExecClientMessage! Available fields: %v", resultFieldName, listFields(ecm)))
|
||||
}
|
||||
|
||||
// Debug: log the actual field being set
|
||||
log.Debugf("encodeExecClientMsg: setting field %q (number=%d, kind=%s)", fd.Name(), fd.Number(), fd.Kind())
|
||||
|
||||
ecm.Set(fd, protoreflect.ValueOfMessage(resultMsg.ProtoReflect()))
|
||||
|
||||
acm := newMsg("AgentClientMessage")
|
||||
setMsg(acm, "exec_client_message", ecm)
|
||||
return marshal(acm)
|
||||
}
|
||||
|
||||
func listFields(msg *dynamicpb.Message) []string {
|
||||
var names []string
|
||||
for i := 0; i < msg.Descriptor().Fields().Len(); i++ {
|
||||
names = append(names, string(msg.Descriptor().Fields().Get(i).Name()))
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
// --- Utilities ---
|
||||
|
||||
// jsonToProtobufValueBytes converts a JSON schema (json.RawMessage) to protobuf Value binary.
|
||||
// This mirrors the TS pattern: toBinary(ValueSchema, fromJson(ValueSchema, jsonSchema))
|
||||
func jsonToProtobufValueBytes(jsonData json.RawMessage) []byte {
|
||||
if len(jsonData) == 0 {
|
||||
return nil
|
||||
}
|
||||
var v interface{}
|
||||
if err := json.Unmarshal(jsonData, &v); err != nil {
|
||||
return jsonData // fallback to raw JSON if parsing fails
|
||||
}
|
||||
pbVal, err := structpb.NewValue(v)
|
||||
if err != nil {
|
||||
return jsonData // fallback
|
||||
}
|
||||
b, err := proto.Marshal(pbVal)
|
||||
if err != nil {
|
||||
return jsonData // fallback
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// ProtobufValueBytesToJSON converts protobuf Value binary back to JSON.
|
||||
// This mirrors the TS pattern: toJson(ValueSchema, fromBinary(ValueSchema, value))
|
||||
func ProtobufValueBytesToJSON(data []byte) (interface{}, error) {
|
||||
val := &structpb.Value{}
|
||||
if err := proto.Unmarshal(data, val); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return val.AsInterface(), nil
|
||||
}
|
||||
|
||||
func sha256Sum(data []byte) []byte {
|
||||
h := sha256.Sum256(data)
|
||||
return h[:]
|
||||
}
|
||||
|
||||
var idCounter uint64
|
||||
|
||||
func generateId() string {
|
||||
idCounter++
|
||||
h := sha256.Sum256([]byte{byte(idCounter), byte(idCounter >> 8), byte(idCounter >> 16)})
|
||||
return hex.EncodeToString(h[:16])
|
||||
}
|
||||
332
internal/auth/cursor/proto/fieldnumbers.go
Normal file
332
internal/auth/cursor/proto/fieldnumbers.go
Normal file
@@ -0,0 +1,332 @@
|
||||
// Package proto provides hand-rolled protobuf encode/decode for Cursor's gRPC API.
|
||||
// Field numbers are extracted from the TypeScript generated proto/agent_pb.ts in alma-plugins/cursor-auth.
|
||||
package proto
|
||||
|
||||
// AgentClientMessage (msg 118) oneof "message"
|
||||
const (
|
||||
ACM_RunRequest = 1 // AgentRunRequest
|
||||
ACM_ExecClientMessage = 2 // ExecClientMessage
|
||||
ACM_KvClientMessage = 3 // KvClientMessage
|
||||
ACM_ConversationAction = 4 // ConversationAction
|
||||
ACM_ExecClientControlMsg = 5 // ExecClientControlMessage
|
||||
ACM_InteractionResponse = 6 // InteractionResponse
|
||||
ACM_ClientHeartbeat = 7 // ClientHeartbeat
|
||||
)
|
||||
|
||||
// AgentServerMessage (msg 119) oneof "message"
|
||||
const (
|
||||
ASM_InteractionUpdate = 1 // InteractionUpdate
|
||||
ASM_ExecServerMessage = 2 // ExecServerMessage
|
||||
ASM_ConversationCheckpoint = 3 // ConversationStateStructure
|
||||
ASM_KvServerMessage = 4 // KvServerMessage
|
||||
ASM_ExecServerControlMessage = 5 // ExecServerControlMessage
|
||||
ASM_InteractionQuery = 7 // InteractionQuery
|
||||
)
|
||||
|
||||
// AgentRunRequest (msg 91)
|
||||
const (
|
||||
ARR_ConversationState = 1 // ConversationStateStructure
|
||||
ARR_Action = 2 // ConversationAction
|
||||
ARR_ModelDetails = 3 // ModelDetails
|
||||
ARR_McpTools = 4 // McpTools
|
||||
ARR_ConversationId = 5 // string (optional)
|
||||
)
|
||||
|
||||
// ConversationStateStructure (msg 83)
|
||||
const (
|
||||
CSS_RootPromptMessagesJson = 1 // repeated bytes
|
||||
CSS_TurnsOld = 2 // repeated bytes (deprecated)
|
||||
CSS_Todos = 3 // repeated bytes
|
||||
CSS_PendingToolCalls = 4 // repeated string
|
||||
CSS_Turns = 8 // repeated bytes (CURRENT field for turns)
|
||||
CSS_PreviousWorkspaceUris = 9 // repeated string
|
||||
CSS_SelfSummaryCount = 17 // uint32
|
||||
CSS_ReadPaths = 18 // repeated string
|
||||
)
|
||||
|
||||
// ConversationAction (msg 54) oneof "action"
|
||||
const (
|
||||
CA_UserMessageAction = 1 // UserMessageAction
|
||||
)
|
||||
|
||||
// UserMessageAction (msg 55)
|
||||
const (
|
||||
UMA_UserMessage = 1 // UserMessage
|
||||
)
|
||||
|
||||
// UserMessage (msg 63)
|
||||
const (
|
||||
UM_Text = 1 // string
|
||||
UM_MessageId = 2 // string
|
||||
UM_SelectedContext = 3 // SelectedContext (optional)
|
||||
)
|
||||
|
||||
// SelectedContext
|
||||
const (
|
||||
SC_SelectedImages = 1 // repeated SelectedImage
|
||||
)
|
||||
|
||||
// SelectedImage
|
||||
const (
|
||||
SI_BlobId = 1 // bytes (oneof dataOrBlobId)
|
||||
SI_Uuid = 2 // string
|
||||
SI_Path = 3 // string
|
||||
SI_MimeType = 7 // string
|
||||
SI_Data = 8 // bytes (oneof dataOrBlobId)
|
||||
)
|
||||
|
||||
// ModelDetails (msg 88)
|
||||
const (
|
||||
MD_ModelId = 1 // string
|
||||
MD_ThinkingDetails = 2 // ThinkingDetails (optional)
|
||||
MD_DisplayModelId = 3 // string
|
||||
MD_DisplayName = 4 // string
|
||||
)
|
||||
|
||||
// McpTools (msg 307)
|
||||
const (
|
||||
MT_McpTools = 1 // repeated McpToolDefinition
|
||||
)
|
||||
|
||||
// McpToolDefinition (msg 306)
|
||||
const (
|
||||
MTD_Name = 1 // string
|
||||
MTD_Description = 2 // string
|
||||
MTD_InputSchema = 3 // bytes
|
||||
MTD_ProviderIdentifier = 4 // string
|
||||
MTD_ToolName = 5 // string
|
||||
)
|
||||
|
||||
// ConversationTurnStructure (msg 70) oneof "turn"
|
||||
const (
|
||||
CTS_AgentConversationTurn = 1 // AgentConversationTurnStructure
|
||||
)
|
||||
|
||||
// AgentConversationTurnStructure (msg 72)
|
||||
const (
|
||||
ACTS_UserMessage = 1 // bytes (serialized UserMessage)
|
||||
ACTS_Steps = 2 // repeated bytes (serialized ConversationStep)
|
||||
)
|
||||
|
||||
// ConversationStep (msg 53) oneof "message"
|
||||
const (
|
||||
CS_AssistantMessage = 1 // AssistantMessage
|
||||
)
|
||||
|
||||
// AssistantMessage
|
||||
const (
|
||||
AM_Text = 1 // string
|
||||
)
|
||||
|
||||
// --- Server-side message fields ---
|
||||
|
||||
// InteractionUpdate oneof "message"
|
||||
const (
|
||||
IU_TextDelta = 1 // TextDeltaUpdate
|
||||
IU_ThinkingDelta = 4 // ThinkingDeltaUpdate
|
||||
IU_ThinkingCompleted = 5 // ThinkingCompletedUpdate
|
||||
)
|
||||
|
||||
// TextDeltaUpdate (msg 92)
|
||||
const (
|
||||
TDU_Text = 1 // string
|
||||
)
|
||||
|
||||
// ThinkingDeltaUpdate (msg 97)
|
||||
const (
|
||||
TKD_Text = 1 // string
|
||||
)
|
||||
|
||||
// KvServerMessage (msg 271)
|
||||
const (
|
||||
KSM_Id = 1 // uint32
|
||||
KSM_GetBlobArgs = 2 // GetBlobArgs
|
||||
KSM_SetBlobArgs = 3 // SetBlobArgs
|
||||
)
|
||||
|
||||
// GetBlobArgs (msg 267)
|
||||
const (
|
||||
GBA_BlobId = 1 // bytes
|
||||
)
|
||||
|
||||
// SetBlobArgs (msg 269)
|
||||
const (
|
||||
SBA_BlobId = 1 // bytes
|
||||
SBA_BlobData = 2 // bytes
|
||||
)
|
||||
|
||||
// KvClientMessage (msg 272)
|
||||
const (
|
||||
KCM_Id = 1 // uint32
|
||||
KCM_GetBlobResult = 2 // GetBlobResult
|
||||
KCM_SetBlobResult = 3 // SetBlobResult
|
||||
)
|
||||
|
||||
// GetBlobResult (msg 268)
|
||||
const (
|
||||
GBR_BlobData = 1 // bytes (optional)
|
||||
)
|
||||
|
||||
// ExecServerMessage
|
||||
const (
|
||||
ESM_Id = 1 // uint32
|
||||
ESM_ExecId = 15 // string
|
||||
// oneof message:
|
||||
ESM_ShellArgs = 2 // ShellArgs
|
||||
ESM_WriteArgs = 3 // WriteArgs
|
||||
ESM_DeleteArgs = 4 // DeleteArgs
|
||||
ESM_GrepArgs = 5 // GrepArgs
|
||||
ESM_ReadArgs = 7 // ReadArgs (NOTE: 6 is skipped)
|
||||
ESM_LsArgs = 8 // LsArgs
|
||||
ESM_DiagnosticsArgs = 9 // DiagnosticsArgs
|
||||
ESM_RequestContextArgs = 10 // RequestContextArgs
|
||||
ESM_McpArgs = 11 // McpArgs
|
||||
ESM_ShellStreamArgs = 14 // ShellArgs (stream variant)
|
||||
ESM_BackgroundShellSpawn = 16 // BackgroundShellSpawnArgs
|
||||
ESM_FetchArgs = 20 // FetchArgs
|
||||
ESM_WriteShellStdinArgs = 23 // WriteShellStdinArgs
|
||||
)
|
||||
|
||||
// ExecClientMessage
|
||||
const (
|
||||
ECM_Id = 1 // uint32
|
||||
ECM_ExecId = 15 // string
|
||||
// oneof message (mirrors server fields):
|
||||
ECM_ShellResult = 2
|
||||
ECM_WriteResult = 3
|
||||
ECM_DeleteResult = 4
|
||||
ECM_GrepResult = 5
|
||||
ECM_ReadResult = 7
|
||||
ECM_LsResult = 8
|
||||
ECM_DiagnosticsResult = 9
|
||||
ECM_RequestContextResult = 10
|
||||
ECM_McpResult = 11
|
||||
ECM_ShellStream = 14
|
||||
ECM_BackgroundShellSpawnRes = 16
|
||||
ECM_FetchResult = 20
|
||||
ECM_WriteShellStdinResult = 23
|
||||
)
|
||||
|
||||
// McpArgs
|
||||
const (
|
||||
MCA_Name = 1 // string
|
||||
MCA_Args = 2 // map<string, bytes>
|
||||
MCA_ToolCallId = 3 // string
|
||||
MCA_ProviderIdentifier = 4 // string
|
||||
MCA_ToolName = 5 // string
|
||||
)
|
||||
|
||||
// RequestContextResult oneof "result"
|
||||
const (
|
||||
RCR_Success = 1 // RequestContextSuccess
|
||||
RCR_Error = 2 // RequestContextError
|
||||
)
|
||||
|
||||
// RequestContextSuccess (msg 337)
|
||||
const (
|
||||
RCS_RequestContext = 1 // RequestContext
|
||||
)
|
||||
|
||||
// RequestContext
|
||||
const (
|
||||
RC_Rules = 2 // repeated CursorRule
|
||||
RC_Tools = 7 // repeated McpToolDefinition
|
||||
)
|
||||
|
||||
// McpResult oneof "result"
|
||||
const (
|
||||
MCR_Success = 1 // McpSuccess
|
||||
MCR_Error = 2 // McpError
|
||||
MCR_Rejected = 3 // McpRejected
|
||||
)
|
||||
|
||||
// McpSuccess (msg 290)
|
||||
const (
|
||||
MCS_Content = 1 // repeated McpToolResultContentItem
|
||||
MCS_IsError = 2 // bool
|
||||
)
|
||||
|
||||
// McpToolResultContentItem oneof "content"
|
||||
const (
|
||||
MTRCI_Text = 1 // McpTextContent
|
||||
)
|
||||
|
||||
// McpTextContent (msg 287)
|
||||
const (
|
||||
MTC_Text = 1 // string
|
||||
)
|
||||
|
||||
// McpError (msg 291)
|
||||
const (
|
||||
MCE_Error = 1 // string
|
||||
)
|
||||
|
||||
// --- Rejection messages ---
|
||||
|
||||
// ReadRejected: path=1, reason=2
|
||||
// ShellRejected: command=1, workingDirectory=2, reason=3, isReadonly=4
|
||||
// WriteRejected: path=1, reason=2
|
||||
// DeleteRejected: path=1, reason=2
|
||||
// LsRejected: path=1, reason=2
|
||||
// GrepError: error=1
|
||||
// FetchError: url=1, error=2
|
||||
// WriteShellStdinError: error=1
|
||||
|
||||
// ReadResult oneof: success=1, error=2, rejected=3
|
||||
// ShellResult oneof: success=1 (+ various), rejected=?
|
||||
// The TS code uses specific result field numbers from the oneof:
|
||||
const (
|
||||
RR_Rejected = 3 // ReadResult.rejected
|
||||
SR_Rejected = 5 // ShellResult.rejected (from TS: ShellResult has success/various/rejected)
|
||||
WR_Rejected = 5 // WriteResult.rejected
|
||||
DR_Rejected = 3 // DeleteResult.rejected
|
||||
LR_Rejected = 3 // LsResult.rejected
|
||||
GR_Error = 2 // GrepResult.error
|
||||
FR_Error = 2 // FetchResult.error
|
||||
BSSR_Rejected = 2 // BackgroundShellSpawnResult.rejected (error field)
|
||||
WSSR_Error = 2 // WriteShellStdinResult.error
|
||||
)
|
||||
|
||||
// --- Rejection struct fields ---
|
||||
const (
|
||||
REJ_Path = 1
|
||||
REJ_Reason = 2
|
||||
SREJ_Command = 1
|
||||
SREJ_WorkingDir = 2
|
||||
SREJ_Reason = 3
|
||||
SREJ_IsReadonly = 4
|
||||
GERR_Error = 1
|
||||
FERR_Url = 1
|
||||
FERR_Error = 2
|
||||
)
|
||||
|
||||
// ReadArgs
|
||||
const (
|
||||
RA_Path = 1 // string
|
||||
)
|
||||
|
||||
// WriteArgs
|
||||
const (
|
||||
WA_Path = 1 // string
|
||||
)
|
||||
|
||||
// DeleteArgs
|
||||
const (
|
||||
DA_Path = 1 // string
|
||||
)
|
||||
|
||||
// LsArgs
|
||||
const (
|
||||
LA_Path = 1 // string
|
||||
)
|
||||
|
||||
// ShellArgs
|
||||
const (
|
||||
SHA_Command = 1 // string
|
||||
SHA_WorkingDirectory = 2 // string
|
||||
)
|
||||
|
||||
// FetchArgs
|
||||
const (
|
||||
FA_Url = 1 // string
|
||||
)
|
||||
313
internal/auth/cursor/proto/h2stream.go
Normal file
313
internal/auth/cursor/proto/h2stream.go
Normal file
@@ -0,0 +1,313 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/net/http2"
|
||||
"golang.org/x/net/http2/hpack"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultInitialWindowSize = 65535 // HTTP/2 default
|
||||
maxFramePayload = 16384 // HTTP/2 default max frame size
|
||||
)
|
||||
|
||||
// H2Stream provides bidirectional HTTP/2 streaming for the Connect protocol.
|
||||
// Go's net/http does not support full-duplex HTTP/2, so we use the low-level framer.
|
||||
type H2Stream struct {
|
||||
framer *http2.Framer
|
||||
conn net.Conn
|
||||
streamID uint32
|
||||
mu sync.Mutex
|
||||
id string // unique identifier for debugging
|
||||
frameNum int64 // sequential frame counter for debugging
|
||||
|
||||
dataCh chan []byte
|
||||
doneCh chan struct{}
|
||||
err error
|
||||
|
||||
// Send-side flow control
|
||||
sendWindow int32 // available bytes we can send on this stream
|
||||
connWindow int32 // available bytes on the connection level
|
||||
windowCond *sync.Cond // signaled when window is updated
|
||||
windowMu sync.Mutex // protects sendWindow, connWindow
|
||||
}
|
||||
|
||||
// ID returns the unique identifier for this stream (for logging).
|
||||
func (s *H2Stream) ID() string { return s.id }
|
||||
|
||||
// FrameNum returns the current frame number for debugging.
|
||||
func (s *H2Stream) FrameNum() int64 {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.frameNum
|
||||
}
|
||||
|
||||
// DialH2Stream establishes a TLS+HTTP/2 connection and opens a new stream.
|
||||
func DialH2Stream(host string, headers map[string]string) (*H2Stream, error) {
|
||||
tlsConn, err := tls.Dial("tcp", host+":443", &tls.Config{
|
||||
NextProtos: []string{"h2"},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("h2: TLS dial failed: %w", err)
|
||||
}
|
||||
if tlsConn.ConnectionState().NegotiatedProtocol != "h2" {
|
||||
tlsConn.Close()
|
||||
return nil, fmt.Errorf("h2: server did not negotiate h2")
|
||||
}
|
||||
|
||||
framer := http2.NewFramer(tlsConn, tlsConn)
|
||||
|
||||
// Client connection preface
|
||||
if _, err := tlsConn.Write([]byte(http2.ClientPreface)); err != nil {
|
||||
tlsConn.Close()
|
||||
return nil, fmt.Errorf("h2: preface write failed: %w", err)
|
||||
}
|
||||
|
||||
// Send initial SETTINGS (tell server how much WE can receive)
|
||||
if err := framer.WriteSettings(
|
||||
http2.Setting{ID: http2.SettingInitialWindowSize, Val: 4 * 1024 * 1024},
|
||||
http2.Setting{ID: http2.SettingMaxConcurrentStreams, Val: 100},
|
||||
); err != nil {
|
||||
tlsConn.Close()
|
||||
return nil, fmt.Errorf("h2: settings write failed: %w", err)
|
||||
}
|
||||
|
||||
// Connection-level window update (for receiving)
|
||||
if err := framer.WriteWindowUpdate(0, 3*1024*1024); err != nil {
|
||||
tlsConn.Close()
|
||||
return nil, fmt.Errorf("h2: window update failed: %w", err)
|
||||
}
|
||||
|
||||
// Read and handle initial server frames (SETTINGS, WINDOW_UPDATE)
|
||||
// Track server's initial window size (how much WE can send)
|
||||
serverInitialWindowSize := int32(defaultInitialWindowSize)
|
||||
connWindowSize := int32(defaultInitialWindowSize) // connection-level send window
|
||||
for i := 0; i < 10; i++ {
|
||||
f, err := framer.ReadFrame()
|
||||
if err != nil {
|
||||
tlsConn.Close()
|
||||
return nil, fmt.Errorf("h2: initial frame read failed: %w", err)
|
||||
}
|
||||
switch sf := f.(type) {
|
||||
case *http2.SettingsFrame:
|
||||
if !sf.IsAck() {
|
||||
sf.ForeachSetting(func(s http2.Setting) error {
|
||||
if s.ID == http2.SettingInitialWindowSize {
|
||||
serverInitialWindowSize = int32(s.Val)
|
||||
log.Debugf("h2: server initial window size: %d", s.Val)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
framer.WriteSettingsAck()
|
||||
} else {
|
||||
goto handshakeDone
|
||||
}
|
||||
case *http2.WindowUpdateFrame:
|
||||
if sf.StreamID == 0 {
|
||||
connWindowSize += int32(sf.Increment)
|
||||
log.Debugf("h2: initial conn window update: +%d, total=%d", sf.Increment, connWindowSize)
|
||||
}
|
||||
default:
|
||||
// unexpected but continue
|
||||
}
|
||||
}
|
||||
handshakeDone:
|
||||
|
||||
// Build HEADERS
|
||||
streamID := uint32(1)
|
||||
var hdrBuf []byte
|
||||
enc := hpack.NewEncoder(&sliceWriter{buf: &hdrBuf})
|
||||
enc.WriteField(hpack.HeaderField{Name: ":method", Value: "POST"})
|
||||
enc.WriteField(hpack.HeaderField{Name: ":scheme", Value: "https"})
|
||||
enc.WriteField(hpack.HeaderField{Name: ":authority", Value: host})
|
||||
if p, ok := headers[":path"]; ok {
|
||||
enc.WriteField(hpack.HeaderField{Name: ":path", Value: p})
|
||||
}
|
||||
for k, v := range headers {
|
||||
if len(k) > 0 && k[0] == ':' {
|
||||
continue
|
||||
}
|
||||
enc.WriteField(hpack.HeaderField{Name: k, Value: v})
|
||||
}
|
||||
|
||||
if err := framer.WriteHeaders(http2.HeadersFrameParam{
|
||||
StreamID: streamID,
|
||||
BlockFragment: hdrBuf,
|
||||
EndStream: false,
|
||||
EndHeaders: true,
|
||||
}); err != nil {
|
||||
tlsConn.Close()
|
||||
return nil, fmt.Errorf("h2: headers write failed: %w", err)
|
||||
}
|
||||
|
||||
s := &H2Stream{
|
||||
framer: framer,
|
||||
conn: tlsConn,
|
||||
streamID: streamID,
|
||||
dataCh: make(chan []byte, 256),
|
||||
doneCh: make(chan struct{}),
|
||||
id: fmt.Sprintf("%d-%s", streamID, time.Now().Format("150405.000")),
|
||||
frameNum: 0,
|
||||
sendWindow: serverInitialWindowSize,
|
||||
connWindow: connWindowSize,
|
||||
}
|
||||
s.windowCond = sync.NewCond(&s.windowMu)
|
||||
go s.readLoop()
|
||||
return s, nil
|
||||
}
|
||||
|
||||
// Write sends a DATA frame on the stream, respecting flow control.
|
||||
func (s *H2Stream) Write(data []byte) error {
|
||||
for len(data) > 0 {
|
||||
chunk := data
|
||||
if len(chunk) > maxFramePayload {
|
||||
chunk = data[:maxFramePayload]
|
||||
}
|
||||
|
||||
// Wait for flow control window
|
||||
s.windowMu.Lock()
|
||||
for s.sendWindow <= 0 || s.connWindow <= 0 {
|
||||
s.windowCond.Wait()
|
||||
}
|
||||
// Limit chunk to available window
|
||||
allowed := int(s.sendWindow)
|
||||
if int(s.connWindow) < allowed {
|
||||
allowed = int(s.connWindow)
|
||||
}
|
||||
if len(chunk) > allowed {
|
||||
chunk = chunk[:allowed]
|
||||
}
|
||||
s.sendWindow -= int32(len(chunk))
|
||||
s.connWindow -= int32(len(chunk))
|
||||
s.windowMu.Unlock()
|
||||
|
||||
s.mu.Lock()
|
||||
err := s.framer.WriteData(s.streamID, false, chunk)
|
||||
s.mu.Unlock()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
data = data[len(chunk):]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Data returns the channel of received data chunks.
|
||||
func (s *H2Stream) Data() <-chan []byte { return s.dataCh }
|
||||
|
||||
// Done returns a channel closed when the stream ends.
|
||||
func (s *H2Stream) Done() <-chan struct{} { return s.doneCh }
|
||||
|
||||
// Err returns the error (if any) that caused the stream to close.
|
||||
// Returns nil for a clean shutdown (EOF / StreamEnded).
|
||||
func (s *H2Stream) Err() error { return s.err }
|
||||
|
||||
// Close tears down the connection.
|
||||
func (s *H2Stream) Close() {
|
||||
s.conn.Close()
|
||||
// Unblock any writers waiting on flow control
|
||||
s.windowCond.Broadcast()
|
||||
}
|
||||
|
||||
func (s *H2Stream) readLoop() {
|
||||
defer close(s.doneCh)
|
||||
defer close(s.dataCh)
|
||||
|
||||
for {
|
||||
f, err := s.framer.ReadFrame()
|
||||
if err != nil {
|
||||
if err != io.EOF {
|
||||
s.err = err
|
||||
log.Debugf("h2stream[%s]: readLoop error: %v", s.id, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Increment frame counter
|
||||
s.mu.Lock()
|
||||
s.frameNum++
|
||||
s.mu.Unlock()
|
||||
|
||||
switch frame := f.(type) {
|
||||
case *http2.DataFrame:
|
||||
if frame.StreamID == s.streamID && len(frame.Data()) > 0 {
|
||||
cp := make([]byte, len(frame.Data()))
|
||||
copy(cp, frame.Data())
|
||||
s.dataCh <- cp
|
||||
|
||||
// Flow control: send WINDOW_UPDATE for received data
|
||||
s.mu.Lock()
|
||||
s.framer.WriteWindowUpdate(0, uint32(len(cp)))
|
||||
s.framer.WriteWindowUpdate(s.streamID, uint32(len(cp)))
|
||||
s.mu.Unlock()
|
||||
}
|
||||
if frame.StreamEnded() {
|
||||
return
|
||||
}
|
||||
|
||||
case *http2.HeadersFrame:
|
||||
if frame.StreamEnded() {
|
||||
return
|
||||
}
|
||||
|
||||
case *http2.RSTStreamFrame:
|
||||
s.err = fmt.Errorf("h2: RST_STREAM code=%d", frame.ErrCode)
|
||||
log.Debugf("h2stream[%s]: received RST_STREAM code=%d", s.id, frame.ErrCode)
|
||||
return
|
||||
|
||||
case *http2.GoAwayFrame:
|
||||
s.err = fmt.Errorf("h2: GOAWAY code=%d", frame.ErrCode)
|
||||
return
|
||||
|
||||
case *http2.PingFrame:
|
||||
if !frame.IsAck() {
|
||||
s.mu.Lock()
|
||||
s.framer.WritePing(true, frame.Data)
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
case *http2.SettingsFrame:
|
||||
if !frame.IsAck() {
|
||||
// Check for window size changes
|
||||
frame.ForeachSetting(func(setting http2.Setting) error {
|
||||
if setting.ID == http2.SettingInitialWindowSize {
|
||||
s.windowMu.Lock()
|
||||
delta := int32(setting.Val) - s.sendWindow
|
||||
s.sendWindow += delta
|
||||
s.windowMu.Unlock()
|
||||
s.windowCond.Broadcast()
|
||||
}
|
||||
return nil
|
||||
})
|
||||
s.mu.Lock()
|
||||
s.framer.WriteSettingsAck()
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
case *http2.WindowUpdateFrame:
|
||||
// Update send-side flow control window
|
||||
s.windowMu.Lock()
|
||||
if frame.StreamID == 0 {
|
||||
s.connWindow += int32(frame.Increment)
|
||||
} else if frame.StreamID == s.streamID {
|
||||
s.sendWindow += int32(frame.Increment)
|
||||
}
|
||||
s.windowMu.Unlock()
|
||||
s.windowCond.Broadcast()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type sliceWriter struct{ buf *[]byte }
|
||||
|
||||
func (w *sliceWriter) Write(p []byte) (int, error) {
|
||||
*w.buf = append(*w.buf, p...)
|
||||
return len(p), nil
|
||||
}
|
||||
@@ -10,9 +10,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/auth/codex"
|
||||
@@ -20,9 +18,9 @@ import (
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/proxyutil"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
"golang.org/x/net/proxy"
|
||||
|
||||
"golang.org/x/oauth2"
|
||||
"golang.org/x/oauth2/google"
|
||||
@@ -80,36 +78,16 @@ func (g *GeminiAuth) GetAuthenticatedClient(ctx context.Context, ts *GeminiToken
|
||||
}
|
||||
callbackURL := fmt.Sprintf("http://localhost:%d/oauth2callback", callbackPort)
|
||||
|
||||
// Configure proxy settings for the HTTP client if a proxy URL is provided.
|
||||
proxyURL, err := url.Parse(cfg.ProxyURL)
|
||||
if err == nil {
|
||||
var transport *http.Transport
|
||||
if proxyURL.Scheme == "socks5" {
|
||||
// Handle SOCKS5 proxy.
|
||||
username := proxyURL.User.Username()
|
||||
password, _ := proxyURL.User.Password()
|
||||
auth := &proxy.Auth{User: username, Password: password}
|
||||
dialer, errSOCKS5 := proxy.SOCKS5("tcp", proxyURL.Host, auth, proxy.Direct)
|
||||
if errSOCKS5 != nil {
|
||||
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) {
|
||||
return dialer.Dial(network, addr)
|
||||
},
|
||||
}
|
||||
} else if proxyURL.Scheme == "http" || proxyURL.Scheme == "https" {
|
||||
// Handle HTTP/HTTPS proxy.
|
||||
transport = &http.Transport{Proxy: http.ProxyURL(proxyURL)}
|
||||
}
|
||||
|
||||
if transport != nil {
|
||||
proxyClient := &http.Client{Transport: transport}
|
||||
ctx = context.WithValue(ctx, oauth2.HTTPClient, proxyClient)
|
||||
}
|
||||
transport, _, errBuild := proxyutil.BuildHTTPTransport(cfg.ProxyURL)
|
||||
if errBuild != nil {
|
||||
log.Errorf("%v", errBuild)
|
||||
} else if transport != nil {
|
||||
proxyClient := &http.Client{Transport: transport}
|
||||
ctx = context.WithValue(ctx, oauth2.HTTPClient, proxyClient)
|
||||
}
|
||||
|
||||
var err error
|
||||
|
||||
// Configure the OAuth2 client.
|
||||
conf := &oauth2.Config{
|
||||
ClientID: ClientID,
|
||||
@@ -327,6 +305,9 @@ func (g *GeminiAuth) getTokenFromWeb(ctx context.Context, config *oauth2.Config,
|
||||
defer manualPromptTimer.Stop()
|
||||
}
|
||||
|
||||
var manualInputCh <-chan string
|
||||
var manualInputErrCh <-chan error
|
||||
|
||||
waitForCallback:
|
||||
for {
|
||||
select {
|
||||
@@ -348,13 +329,14 @@ waitForCallback:
|
||||
return nil, err
|
||||
default:
|
||||
}
|
||||
input, err := opts.Prompt("Paste the Gemini callback URL (or press Enter to keep waiting): ")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
parsed, err := misc.ParseOAuthCallback(input)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
manualInputCh, manualInputErrCh = misc.AsyncPrompt(opts.Prompt, "Paste the Gemini callback URL (or press Enter to keep waiting): ")
|
||||
continue
|
||||
case input := <-manualInputCh:
|
||||
manualInputCh = nil
|
||||
manualInputErrCh = nil
|
||||
parsed, errParse := misc.ParseOAuthCallback(input)
|
||||
if errParse != nil {
|
||||
return nil, errParse
|
||||
}
|
||||
if parsed == nil {
|
||||
continue
|
||||
@@ -367,6 +349,8 @@ waitForCallback:
|
||||
}
|
||||
authCode = parsed.Code
|
||||
break waitForCallback
|
||||
case errManual := <-manualInputErrCh:
|
||||
return nil, errManual
|
||||
case <-timeoutTimer.C:
|
||||
return nil, fmt.Errorf("oauth flow timed out")
|
||||
}
|
||||
|
||||
@@ -35,11 +35,21 @@ type GeminiTokenStorage struct {
|
||||
|
||||
// Type indicates the authentication provider type, always "gemini" for this storage.
|
||||
Type string `json:"type"`
|
||||
|
||||
// Metadata holds arbitrary key-value pairs injected via hooks.
|
||||
// It is not exported to JSON directly to allow flattening during serialization.
|
||||
Metadata map[string]any `json:"-"`
|
||||
}
|
||||
|
||||
// SetMetadata allows external callers to inject metadata into the storage before saving.
|
||||
func (ts *GeminiTokenStorage) SetMetadata(meta map[string]any) {
|
||||
ts.Metadata = meta
|
||||
}
|
||||
|
||||
// SaveTokenToFile serializes the Gemini token storage to a JSON file.
|
||||
// This method creates the necessary directory structure and writes the token
|
||||
// data in JSON format to the specified file path for persistent storage.
|
||||
// It merges any injected metadata into the top-level JSON object.
|
||||
//
|
||||
// Parameters:
|
||||
// - authFilePath: The full path where the token file should be saved
|
||||
@@ -49,6 +59,11 @@ type GeminiTokenStorage struct {
|
||||
func (ts *GeminiTokenStorage) SaveTokenToFile(authFilePath string) error {
|
||||
misc.LogSavingCredentials(authFilePath)
|
||||
ts.Type = "gemini"
|
||||
// Merge metadata using helper
|
||||
data, errMerge := misc.MergeMetadata(ts, ts.Metadata)
|
||||
if errMerge != nil {
|
||||
return fmt.Errorf("failed to merge metadata: %w", errMerge)
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(authFilePath), 0700); err != nil {
|
||||
return fmt.Errorf("failed to create directory: %v", err)
|
||||
}
|
||||
@@ -63,7 +78,9 @@ func (ts *GeminiTokenStorage) SaveTokenToFile(authFilePath string) error {
|
||||
}
|
||||
}()
|
||||
|
||||
if err = json.NewEncoder(f).Encode(ts); err != nil {
|
||||
enc := json.NewEncoder(f)
|
||||
enc.SetIndent("", " ")
|
||||
if err := enc.Encode(data); err != nil {
|
||||
return fmt.Errorf("failed to write token to file: %w", err)
|
||||
}
|
||||
return nil
|
||||
|
||||
492
internal/auth/gitlab/gitlab.go
Normal file
492
internal/auth/gitlab/gitlab.go
Normal file
@@ -0,0 +1,492 @@
|
||||
package gitlab
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"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 (
|
||||
DefaultBaseURL = "https://gitlab.com"
|
||||
DefaultCallbackPort = 17171
|
||||
defaultOAuthScope = "api read_user"
|
||||
)
|
||||
|
||||
type PKCECodes struct {
|
||||
CodeVerifier string
|
||||
CodeChallenge string
|
||||
}
|
||||
|
||||
type OAuthResult struct {
|
||||
Code string
|
||||
State string
|
||||
Error string
|
||||
}
|
||||
|
||||
type OAuthServer struct {
|
||||
server *http.Server
|
||||
port int
|
||||
resultChan chan *OAuthResult
|
||||
errorChan chan error
|
||||
mu sync.Mutex
|
||||
running bool
|
||||
}
|
||||
|
||||
type TokenResponse struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
TokenType string `json:"token_type"`
|
||||
RefreshToken string `json:"refresh_token"`
|
||||
Scope string `json:"scope"`
|
||||
CreatedAt int64 `json:"created_at"`
|
||||
ExpiresIn int `json:"expires_in"`
|
||||
}
|
||||
|
||||
type User struct {
|
||||
ID int64 `json:"id"`
|
||||
Username string `json:"username"`
|
||||
Name string `json:"name"`
|
||||
Email string `json:"email"`
|
||||
PublicEmail string `json:"public_email"`
|
||||
}
|
||||
|
||||
type PersonalAccessTokenSelf struct {
|
||||
ID int64 `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Scopes []string `json:"scopes"`
|
||||
UserID int64 `json:"user_id"`
|
||||
}
|
||||
|
||||
type ModelDetails struct {
|
||||
ModelProvider string `json:"model_provider"`
|
||||
ModelName string `json:"model_name"`
|
||||
}
|
||||
|
||||
type DirectAccessResponse struct {
|
||||
BaseURL string `json:"base_url"`
|
||||
Token string `json:"token"`
|
||||
ExpiresAt int64 `json:"expires_at"`
|
||||
Headers map[string]string `json:"headers"`
|
||||
ModelDetails *ModelDetails `json:"model_details,omitempty"`
|
||||
}
|
||||
|
||||
type DiscoveredModel struct {
|
||||
ModelProvider string
|
||||
ModelName string
|
||||
}
|
||||
|
||||
type AuthClient struct {
|
||||
httpClient *http.Client
|
||||
}
|
||||
|
||||
func NewAuthClient(cfg *config.Config) *AuthClient {
|
||||
client := &http.Client{}
|
||||
if cfg != nil {
|
||||
client = util.SetProxy(&cfg.SDKConfig, client)
|
||||
}
|
||||
return &AuthClient{httpClient: client}
|
||||
}
|
||||
|
||||
func NormalizeBaseURL(raw string) string {
|
||||
value := strings.TrimSpace(raw)
|
||||
if value == "" {
|
||||
return DefaultBaseURL
|
||||
}
|
||||
if !strings.Contains(value, "://") {
|
||||
value = "https://" + value
|
||||
}
|
||||
value = strings.TrimRight(value, "/")
|
||||
return value
|
||||
}
|
||||
|
||||
func TokenExpiry(now time.Time, token *TokenResponse) time.Time {
|
||||
if token == nil {
|
||||
return time.Time{}
|
||||
}
|
||||
if token.CreatedAt > 0 && token.ExpiresIn > 0 {
|
||||
return time.Unix(token.CreatedAt+int64(token.ExpiresIn), 0).UTC()
|
||||
}
|
||||
if token.ExpiresIn > 0 {
|
||||
return now.UTC().Add(time.Duration(token.ExpiresIn) * time.Second)
|
||||
}
|
||||
return time.Time{}
|
||||
}
|
||||
|
||||
func GeneratePKCECodes() (*PKCECodes, error) {
|
||||
verifierBytes := make([]byte, 32)
|
||||
if _, err := rand.Read(verifierBytes); err != nil {
|
||||
return nil, fmt.Errorf("gitlab pkce generation failed: %w", err)
|
||||
}
|
||||
verifier := base64.RawURLEncoding.EncodeToString(verifierBytes)
|
||||
sum := sha256.Sum256([]byte(verifier))
|
||||
challenge := base64.RawURLEncoding.EncodeToString(sum[:])
|
||||
return &PKCECodes{
|
||||
CodeVerifier: verifier,
|
||||
CodeChallenge: challenge,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func NewOAuthServer(port int) *OAuthServer {
|
||||
return &OAuthServer{
|
||||
port: port,
|
||||
resultChan: make(chan *OAuthResult, 1),
|
||||
errorChan: make(chan error, 1),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OAuthServer) Start() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.running {
|
||||
return fmt.Errorf("gitlab oauth server already running")
|
||||
}
|
||||
if !s.isPortAvailable() {
|
||||
return fmt.Errorf("port %d is already in use", s.port)
|
||||
}
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/auth/callback", s.handleCallback)
|
||||
|
||||
s.server = &http.Server{
|
||||
Addr: fmt.Sprintf(":%d", s.port),
|
||||
Handler: mux,
|
||||
ReadTimeout: 10 * time.Second,
|
||||
WriteTimeout: 10 * time.Second,
|
||||
}
|
||||
s.running = true
|
||||
|
||||
go func() {
|
||||
if err := s.server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
s.errorChan <- err
|
||||
}
|
||||
}()
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *OAuthServer) Stop(ctx context.Context) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if !s.running || s.server == nil {
|
||||
return nil
|
||||
}
|
||||
defer func() {
|
||||
s.running = false
|
||||
s.server = nil
|
||||
}()
|
||||
return s.server.Shutdown(ctx)
|
||||
}
|
||||
|
||||
func (s *OAuthServer) WaitForCallback(timeout time.Duration) (*OAuthResult, error) {
|
||||
select {
|
||||
case result := <-s.resultChan:
|
||||
return result, nil
|
||||
case err := <-s.errorChan:
|
||||
return nil, err
|
||||
case <-time.After(timeout):
|
||||
return nil, fmt.Errorf("timeout waiting for OAuth callback")
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OAuthServer) handleCallback(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
query := r.URL.Query()
|
||||
if errParam := strings.TrimSpace(query.Get("error")); errParam != "" {
|
||||
s.sendResult(&OAuthResult{Error: errParam})
|
||||
http.Error(w, errParam, http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
code := strings.TrimSpace(query.Get("code"))
|
||||
state := strings.TrimSpace(query.Get("state"))
|
||||
if code == "" || state == "" {
|
||||
s.sendResult(&OAuthResult{Error: "missing_code_or_state"})
|
||||
http.Error(w, "missing code or state", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
s.sendResult(&OAuthResult{Code: code, State: state})
|
||||
_, _ = w.Write([]byte("GitLab authentication received. You can close this tab."))
|
||||
}
|
||||
|
||||
func (s *OAuthServer) sendResult(result *OAuthResult) {
|
||||
select {
|
||||
case s.resultChan <- result:
|
||||
default:
|
||||
log.Debug("gitlab oauth result channel full, dropping callback result")
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OAuthServer) isPortAvailable() bool {
|
||||
listener, err := net.Listen("tcp", fmt.Sprintf(":%d", s.port))
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
_ = listener.Close()
|
||||
return true
|
||||
}
|
||||
|
||||
func RedirectURL(port int) string {
|
||||
return fmt.Sprintf("http://localhost:%d/auth/callback", port)
|
||||
}
|
||||
|
||||
func (c *AuthClient) GenerateAuthURL(baseURL, clientID, redirectURI, state string, pkce *PKCECodes) (string, error) {
|
||||
if pkce == nil {
|
||||
return "", fmt.Errorf("gitlab auth URL generation failed: PKCE codes are required")
|
||||
}
|
||||
if strings.TrimSpace(clientID) == "" {
|
||||
return "", fmt.Errorf("gitlab auth URL generation failed: client ID is required")
|
||||
}
|
||||
baseURL = NormalizeBaseURL(baseURL)
|
||||
params := url.Values{
|
||||
"client_id": {strings.TrimSpace(clientID)},
|
||||
"response_type": {"code"},
|
||||
"redirect_uri": {strings.TrimSpace(redirectURI)},
|
||||
"scope": {defaultOAuthScope},
|
||||
"state": {strings.TrimSpace(state)},
|
||||
"code_challenge": {pkce.CodeChallenge},
|
||||
"code_challenge_method": {"S256"},
|
||||
}
|
||||
return fmt.Sprintf("%s/oauth/authorize?%s", baseURL, params.Encode()), nil
|
||||
}
|
||||
|
||||
func (c *AuthClient) ExchangeCodeForTokens(ctx context.Context, baseURL, clientID, clientSecret, redirectURI, code, codeVerifier string) (*TokenResponse, error) {
|
||||
form := url.Values{
|
||||
"grant_type": {"authorization_code"},
|
||||
"client_id": {strings.TrimSpace(clientID)},
|
||||
"code": {strings.TrimSpace(code)},
|
||||
"redirect_uri": {strings.TrimSpace(redirectURI)},
|
||||
"code_verifier": {strings.TrimSpace(codeVerifier)},
|
||||
}
|
||||
if secret := strings.TrimSpace(clientSecret); secret != "" {
|
||||
form.Set("client_secret", secret)
|
||||
}
|
||||
return c.postToken(ctx, NormalizeBaseURL(baseURL)+"/oauth/token", form)
|
||||
}
|
||||
|
||||
func (c *AuthClient) RefreshTokens(ctx context.Context, baseURL, clientID, clientSecret, refreshToken string) (*TokenResponse, error) {
|
||||
form := url.Values{
|
||||
"grant_type": {"refresh_token"},
|
||||
"refresh_token": {strings.TrimSpace(refreshToken)},
|
||||
}
|
||||
if clientID = strings.TrimSpace(clientID); clientID != "" {
|
||||
form.Set("client_id", clientID)
|
||||
}
|
||||
if secret := strings.TrimSpace(clientSecret); secret != "" {
|
||||
form.Set("client_secret", secret)
|
||||
}
|
||||
return c.postToken(ctx, NormalizeBaseURL(baseURL)+"/oauth/token", form)
|
||||
}
|
||||
|
||||
func (c *AuthClient) postToken(ctx context.Context, tokenURL string, form url.Values) (*TokenResponse, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(form.Encode()))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gitlab token request failed: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gitlab token request failed: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gitlab token response read failed: %w", err)
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("gitlab token request failed with status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
|
||||
}
|
||||
var token TokenResponse
|
||||
if err := json.Unmarshal(body, &token); err != nil {
|
||||
return nil, fmt.Errorf("gitlab token response decode failed: %w", err)
|
||||
}
|
||||
return &token, nil
|
||||
}
|
||||
|
||||
func (c *AuthClient) GetCurrentUser(ctx context.Context, baseURL, token string) (*User, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, NormalizeBaseURL(baseURL)+"/api/v4/user", nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gitlab user request failed: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(token))
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gitlab user request failed: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gitlab user response read failed: %w", err)
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("gitlab user request failed with status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
|
||||
}
|
||||
|
||||
var user User
|
||||
if err := json.Unmarshal(body, &user); err != nil {
|
||||
return nil, fmt.Errorf("gitlab user response decode failed: %w", err)
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (c *AuthClient) GetPersonalAccessTokenSelf(ctx context.Context, baseURL, token string) (*PersonalAccessTokenSelf, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, NormalizeBaseURL(baseURL)+"/api/v4/personal_access_tokens/self", nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gitlab PAT self request failed: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(token))
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gitlab PAT self request failed: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gitlab PAT self response read failed: %w", err)
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("gitlab PAT self request failed with status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
|
||||
}
|
||||
|
||||
var pat PersonalAccessTokenSelf
|
||||
if err := json.Unmarshal(body, &pat); err != nil {
|
||||
return nil, fmt.Errorf("gitlab PAT self response decode failed: %w", err)
|
||||
}
|
||||
return &pat, nil
|
||||
}
|
||||
|
||||
func (c *AuthClient) FetchDirectAccess(ctx context.Context, baseURL, token string) (*DirectAccessResponse, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, NormalizeBaseURL(baseURL)+"/api/v4/code_suggestions/direct_access", nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gitlab direct access request failed: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(token))
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gitlab direct access request failed: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gitlab direct access response read failed: %w", err)
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("gitlab direct access request failed with status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
|
||||
}
|
||||
|
||||
var direct DirectAccessResponse
|
||||
if err := json.Unmarshal(body, &direct); err != nil {
|
||||
return nil, fmt.Errorf("gitlab direct access response decode failed: %w", err)
|
||||
}
|
||||
if direct.Headers == nil {
|
||||
direct.Headers = make(map[string]string)
|
||||
}
|
||||
return &direct, nil
|
||||
}
|
||||
|
||||
func ExtractDiscoveredModels(metadata map[string]any) []DiscoveredModel {
|
||||
if len(metadata) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
models := make([]DiscoveredModel, 0, 4)
|
||||
seen := make(map[string]struct{})
|
||||
appendModel := func(provider, name string) {
|
||||
provider = strings.TrimSpace(provider)
|
||||
name = strings.TrimSpace(name)
|
||||
if name == "" {
|
||||
return
|
||||
}
|
||||
key := strings.ToLower(name)
|
||||
if _, ok := seen[key]; ok {
|
||||
return
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
models = append(models, DiscoveredModel{
|
||||
ModelProvider: provider,
|
||||
ModelName: name,
|
||||
})
|
||||
}
|
||||
|
||||
if raw, ok := metadata["model_details"]; ok {
|
||||
appendDiscoveredModels(raw, appendModel)
|
||||
}
|
||||
appendModel(stringValue(metadata["model_provider"]), stringValue(metadata["model_name"]))
|
||||
|
||||
for _, key := range []string{"models", "supported_models", "discovered_models"} {
|
||||
if raw, ok := metadata[key]; ok {
|
||||
appendDiscoveredModels(raw, appendModel)
|
||||
}
|
||||
}
|
||||
|
||||
return models
|
||||
}
|
||||
|
||||
func appendDiscoveredModels(raw any, appendModel func(provider, name string)) {
|
||||
switch typed := raw.(type) {
|
||||
case map[string]any:
|
||||
appendModel(stringValue(typed["model_provider"]), stringValue(typed["model_name"]))
|
||||
appendModel(stringValue(typed["provider"]), stringValue(typed["name"]))
|
||||
if nested, ok := typed["models"]; ok {
|
||||
appendDiscoveredModels(nested, appendModel)
|
||||
}
|
||||
case []any:
|
||||
for _, item := range typed {
|
||||
appendDiscoveredModels(item, appendModel)
|
||||
}
|
||||
case []string:
|
||||
for _, item := range typed {
|
||||
appendModel("", item)
|
||||
}
|
||||
case string:
|
||||
appendModel("", typed)
|
||||
}
|
||||
}
|
||||
|
||||
func stringValue(raw any) string {
|
||||
switch typed := raw.(type) {
|
||||
case string:
|
||||
return strings.TrimSpace(typed)
|
||||
case fmt.Stringer:
|
||||
return strings.TrimSpace(typed.String())
|
||||
case json.Number:
|
||||
return typed.String()
|
||||
case int:
|
||||
return strconv.Itoa(typed)
|
||||
case int64:
|
||||
return strconv.FormatInt(typed, 10)
|
||||
case float64:
|
||||
return strconv.FormatInt(int64(typed), 10)
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
138
internal/auth/gitlab/gitlab_test.go
Normal file
138
internal/auth/gitlab/gitlab_test.go
Normal file
@@ -0,0 +1,138 @@
|
||||
package gitlab
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAuthClientGenerateAuthURLIncludesPKCE(t *testing.T) {
|
||||
client := NewAuthClient(nil)
|
||||
pkce, err := GeneratePKCECodes()
|
||||
if err != nil {
|
||||
t.Fatalf("GeneratePKCECodes() error = %v", err)
|
||||
}
|
||||
|
||||
rawURL, err := client.GenerateAuthURL("https://gitlab.example.com", "client-id", RedirectURL(17171), "state-123", pkce)
|
||||
if err != nil {
|
||||
t.Fatalf("GenerateAuthURL() error = %v", err)
|
||||
}
|
||||
|
||||
parsed, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
t.Fatalf("Parse(authURL) error = %v", err)
|
||||
}
|
||||
if got := parsed.Path; got != "/oauth/authorize" {
|
||||
t.Fatalf("expected /oauth/authorize path, got %q", got)
|
||||
}
|
||||
query := parsed.Query()
|
||||
if got := query.Get("client_id"); got != "client-id" {
|
||||
t.Fatalf("expected client_id, got %q", got)
|
||||
}
|
||||
if got := query.Get("scope"); got != defaultOAuthScope {
|
||||
t.Fatalf("expected scope %q, got %q", defaultOAuthScope, got)
|
||||
}
|
||||
if got := query.Get("code_challenge_method"); got != "S256" {
|
||||
t.Fatalf("expected PKCE method S256, got %q", got)
|
||||
}
|
||||
if got := query.Get("code_challenge"); got == "" {
|
||||
t.Fatal("expected non-empty code_challenge")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthClientExchangeCodeForTokens(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/oauth/token" {
|
||||
t.Fatalf("unexpected path %q", r.URL.Path)
|
||||
}
|
||||
if err := r.ParseForm(); err != nil {
|
||||
t.Fatalf("ParseForm() error = %v", err)
|
||||
}
|
||||
if got := r.Form.Get("grant_type"); got != "authorization_code" {
|
||||
t.Fatalf("expected authorization_code grant, got %q", got)
|
||||
}
|
||||
if got := r.Form.Get("code_verifier"); got != "verifier-123" {
|
||||
t.Fatalf("expected code_verifier, got %q", got)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"access_token": "oauth-access",
|
||||
"refresh_token": "oauth-refresh",
|
||||
"token_type": "Bearer",
|
||||
"scope": "api read_user",
|
||||
"created_at": 1710000000,
|
||||
"expires_in": 3600,
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := NewAuthClient(nil)
|
||||
token, err := client.ExchangeCodeForTokens(context.Background(), srv.URL, "client-id", "client-secret", RedirectURL(17171), "auth-code", "verifier-123")
|
||||
if err != nil {
|
||||
t.Fatalf("ExchangeCodeForTokens() error = %v", err)
|
||||
}
|
||||
if token.AccessToken != "oauth-access" {
|
||||
t.Fatalf("expected access token, got %q", token.AccessToken)
|
||||
}
|
||||
if token.RefreshToken != "oauth-refresh" {
|
||||
t.Fatalf("expected refresh token, got %q", token.RefreshToken)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractDiscoveredModels(t *testing.T) {
|
||||
models := ExtractDiscoveredModels(map[string]any{
|
||||
"model_details": map[string]any{
|
||||
"model_provider": "anthropic",
|
||||
"model_name": "claude-sonnet-4-5",
|
||||
},
|
||||
"supported_models": []any{
|
||||
map[string]any{"model_provider": "openai", "model_name": "gpt-4.1"},
|
||||
"claude-sonnet-4-5",
|
||||
},
|
||||
})
|
||||
if len(models) != 2 {
|
||||
t.Fatalf("expected 2 unique models, got %d", len(models))
|
||||
}
|
||||
if models[0].ModelName != "claude-sonnet-4-5" {
|
||||
t.Fatalf("unexpected first model %q", models[0].ModelName)
|
||||
}
|
||||
if models[1].ModelName != "gpt-4.1" {
|
||||
t.Fatalf("unexpected second model %q", models[1].ModelName)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchDirectAccessDecodesModelDetails(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/api/v4/code_suggestions/direct_access" {
|
||||
t.Fatalf("unexpected path %q", r.URL.Path)
|
||||
}
|
||||
if got := r.Header.Get("Authorization"); !strings.Contains(got, "token-123") {
|
||||
t.Fatalf("expected bearer token, got %q", got)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"base_url": "https://cloud.gitlab.example.com",
|
||||
"token": "gateway-token",
|
||||
"expires_at": 1710003600,
|
||||
"headers": map[string]string{
|
||||
"X-Gitlab-Realm": "saas",
|
||||
},
|
||||
"model_details": map[string]any{
|
||||
"model_provider": "anthropic",
|
||||
"model_name": "claude-sonnet-4-5",
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := NewAuthClient(nil)
|
||||
direct, err := client.FetchDirectAccess(context.Background(), srv.URL, "token-123")
|
||||
if err != nil {
|
||||
t.Fatalf("FetchDirectAccess() error = %v", err)
|
||||
}
|
||||
if direct.ModelDetails == nil || direct.ModelDetails.ModelName != "claude-sonnet-4-5" {
|
||||
t.Fatalf("expected model details, got %+v", direct.ModelDetails)
|
||||
}
|
||||
}
|
||||
@@ -21,6 +21,15 @@ type IFlowTokenStorage struct {
|
||||
Scope string `json:"scope"`
|
||||
Cookie string `json:"cookie"`
|
||||
Type string `json:"type"`
|
||||
|
||||
// Metadata holds arbitrary key-value pairs injected via hooks.
|
||||
// It is not exported to JSON directly to allow flattening during serialization.
|
||||
Metadata map[string]any `json:"-"`
|
||||
}
|
||||
|
||||
// SetMetadata allows external callers to inject metadata into the storage before saving.
|
||||
func (ts *IFlowTokenStorage) SetMetadata(meta map[string]any) {
|
||||
ts.Metadata = meta
|
||||
}
|
||||
|
||||
// SaveTokenToFile serialises the token storage to disk.
|
||||
@@ -37,7 +46,13 @@ func (ts *IFlowTokenStorage) SaveTokenToFile(authFilePath string) error {
|
||||
}
|
||||
defer func() { _ = f.Close() }()
|
||||
|
||||
if err = json.NewEncoder(f).Encode(ts); err != nil {
|
||||
// Merge metadata using helper
|
||||
data, errMerge := misc.MergeMetadata(ts, ts.Metadata)
|
||||
if errMerge != nil {
|
||||
return fmt.Errorf("failed to merge metadata: %w", errMerge)
|
||||
}
|
||||
|
||||
if err = json.NewEncoder(f).Encode(data); err != nil {
|
||||
return fmt.Errorf("iflow token: encode token failed: %w", err)
|
||||
}
|
||||
return nil
|
||||
|
||||
@@ -29,6 +29,15 @@ type KimiTokenStorage struct {
|
||||
Expired string `json:"expired,omitempty"`
|
||||
// Type indicates the authentication provider type, always "kimi" for this storage.
|
||||
Type string `json:"type"`
|
||||
|
||||
// Metadata holds arbitrary key-value pairs injected via hooks.
|
||||
// It is not exported to JSON directly to allow flattening during serialization.
|
||||
Metadata map[string]any `json:"-"`
|
||||
}
|
||||
|
||||
// SetMetadata allows external callers to inject metadata into the storage before saving.
|
||||
func (ts *KimiTokenStorage) SetMetadata(meta map[string]any) {
|
||||
ts.Metadata = meta
|
||||
}
|
||||
|
||||
// KimiTokenData holds the raw OAuth token response from Kimi.
|
||||
@@ -86,9 +95,15 @@ func (ts *KimiTokenStorage) SaveTokenToFile(authFilePath string) error {
|
||||
_ = f.Close()
|
||||
}()
|
||||
|
||||
// Merge metadata using helper
|
||||
data, errMerge := misc.MergeMetadata(ts, ts.Metadata)
|
||||
if errMerge != nil {
|
||||
return fmt.Errorf("failed to merge metadata: %w", errMerge)
|
||||
}
|
||||
|
||||
encoder := json.NewEncoder(f)
|
||||
encoder.SetIndent("", " ")
|
||||
if err = encoder.Encode(ts); err != nil {
|
||||
if err = encoder.Encode(data); err != nil {
|
||||
return fmt.Errorf("failed to write token to file: %w", err)
|
||||
}
|
||||
return nil
|
||||
|
||||
@@ -7,10 +7,13 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// PKCECodes holds PKCE verification codes for OAuth2 PKCE flow
|
||||
@@ -47,7 +50,7 @@ type KiroTokenData struct {
|
||||
Email string `json:"email,omitempty"`
|
||||
// StartURL is the IDC/Identity Center start URL (only for IDC auth method)
|
||||
StartURL string `json:"startUrl,omitempty"`
|
||||
// Region is the AWS region for IDC authentication (only for IDC auth method)
|
||||
// Region is the OIDC region for IDC login and token refresh
|
||||
Region string `json:"region,omitempty"`
|
||||
}
|
||||
|
||||
@@ -520,3 +523,159 @@ func GenerateTokenFileName(tokenData *KiroTokenData) string {
|
||||
// Priority 3: Fallback to authMethod only with sequence
|
||||
return fmt.Sprintf("kiro-%s-%05d.json", authMethod, seq)
|
||||
}
|
||||
|
||||
// DefaultKiroRegion is the fallback region when none is specified.
|
||||
const DefaultKiroRegion = "us-east-1"
|
||||
|
||||
// GetCodeWhispererLegacyEndpoint returns the legacy CodeWhisperer JSON-RPC endpoint.
|
||||
// This endpoint supports JSON-RPC style requests with x-amz-target headers.
|
||||
// The Q endpoint (q.{region}.amazonaws.com) does NOT support JSON-RPC style.
|
||||
func GetCodeWhispererLegacyEndpoint(region string) string {
|
||||
if region == "" {
|
||||
region = DefaultKiroRegion
|
||||
}
|
||||
return "https://codewhisperer." + region + ".amazonaws.com"
|
||||
}
|
||||
|
||||
// ProfileARN represents a parsed AWS CodeWhisperer profile ARN.
|
||||
// ARN format: arn:partition:service:region:account-id:resource-type/resource-id
|
||||
// Example: arn:aws:codewhisperer:us-east-1:123456789012:profile/ABCDEFGHIJKL
|
||||
type ProfileARN struct {
|
||||
// Raw is the original ARN string
|
||||
Raw string
|
||||
// Partition is the AWS partition (aws)
|
||||
Partition string
|
||||
// Service is the AWS service name (codewhisperer)
|
||||
Service string
|
||||
// Region is the AWS region (us-east-1, ap-southeast-1, etc.)
|
||||
Region string
|
||||
// AccountID is the AWS account ID
|
||||
AccountID string
|
||||
// ResourceType is the resource type (profile)
|
||||
ResourceType string
|
||||
// ResourceID is the resource identifier (e.g., ABCDEFGHIJKL)
|
||||
ResourceID string
|
||||
}
|
||||
|
||||
// ParseProfileARN parses an AWS ARN string into a ProfileARN struct.
|
||||
// Returns nil if the ARN is empty, invalid, or not a codewhisperer ARN.
|
||||
func ParseProfileARN(arn string) *ProfileARN {
|
||||
if arn == "" {
|
||||
return nil
|
||||
}
|
||||
// ARN format: arn:partition:service:region:account-id:resource
|
||||
// Minimum 6 parts separated by ":"
|
||||
parts := strings.Split(arn, ":")
|
||||
if len(parts) < 6 {
|
||||
log.Warnf("invalid ARN format: %s", arn)
|
||||
return nil
|
||||
}
|
||||
// Validate ARN prefix
|
||||
if parts[0] != "arn" {
|
||||
return nil
|
||||
}
|
||||
// Validate partition
|
||||
partition := parts[1]
|
||||
if partition == "" {
|
||||
return nil
|
||||
}
|
||||
// Validate service is codewhisperer
|
||||
service := parts[2]
|
||||
if service != "codewhisperer" {
|
||||
return nil
|
||||
}
|
||||
// Validate region format (must contain "-")
|
||||
region := parts[3]
|
||||
if region == "" || !strings.Contains(region, "-") {
|
||||
return nil
|
||||
}
|
||||
// Account ID
|
||||
accountID := parts[4]
|
||||
|
||||
// Parse resource (format: resource-type/resource-id)
|
||||
// Join remaining parts in case resource contains ":"
|
||||
resource := strings.Join(parts[5:], ":")
|
||||
resourceType := ""
|
||||
resourceID := ""
|
||||
if idx := strings.Index(resource, "/"); idx > 0 {
|
||||
resourceType = resource[:idx]
|
||||
resourceID = resource[idx+1:]
|
||||
} else {
|
||||
resourceType = resource
|
||||
}
|
||||
|
||||
return &ProfileARN{
|
||||
Raw: arn,
|
||||
Partition: partition,
|
||||
Service: service,
|
||||
Region: region,
|
||||
AccountID: accountID,
|
||||
ResourceType: resourceType,
|
||||
ResourceID: resourceID,
|
||||
}
|
||||
}
|
||||
|
||||
// GetKiroAPIEndpoint returns the Q API endpoint for the specified region.
|
||||
// If region is empty, defaults to us-east-1.
|
||||
func GetKiroAPIEndpoint(region string) string {
|
||||
if region == "" {
|
||||
region = DefaultKiroRegion
|
||||
}
|
||||
return "https://q." + region + ".amazonaws.com"
|
||||
}
|
||||
|
||||
// GetKiroAPIEndpointFromProfileArn extracts region from profileArn and returns the endpoint.
|
||||
// Returns default us-east-1 endpoint if region cannot be extracted.
|
||||
func GetKiroAPIEndpointFromProfileArn(profileArn string) string {
|
||||
region := ExtractRegionFromProfileArn(profileArn)
|
||||
return GetKiroAPIEndpoint(region)
|
||||
}
|
||||
|
||||
// ExtractRegionFromProfileArn extracts the AWS region from a ProfileARN string.
|
||||
// Returns empty string if ARN is invalid or region cannot be extracted.
|
||||
func ExtractRegionFromProfileArn(profileArn string) string {
|
||||
parsed := ParseProfileARN(profileArn)
|
||||
if parsed == nil {
|
||||
return ""
|
||||
}
|
||||
return parsed.Region
|
||||
}
|
||||
|
||||
// ExtractRegionFromMetadata extracts API region from auth metadata.
|
||||
// Priority: api_region > profile_arn > DefaultKiroRegion
|
||||
func ExtractRegionFromMetadata(metadata map[string]interface{}) string {
|
||||
if metadata == nil {
|
||||
return DefaultKiroRegion
|
||||
}
|
||||
|
||||
// Priority 1: Explicit api_region override
|
||||
if r, ok := metadata["api_region"].(string); ok && r != "" {
|
||||
return r
|
||||
}
|
||||
|
||||
// Priority 2: Extract from ProfileARN
|
||||
if profileArn, ok := metadata["profile_arn"].(string); ok && profileArn != "" {
|
||||
if region := ExtractRegionFromProfileArn(profileArn); region != "" {
|
||||
return region
|
||||
}
|
||||
}
|
||||
|
||||
return DefaultKiroRegion
|
||||
}
|
||||
|
||||
func buildURL(endpoint, path string, queryParams map[string]string) string {
|
||||
fullURL := fmt.Sprintf("%s/%s", endpoint, path)
|
||||
if len(queryParams) > 0 {
|
||||
values := url.Values{}
|
||||
for key, value := range queryParams {
|
||||
if value == "" {
|
||||
continue
|
||||
}
|
||||
values.Set(key, value)
|
||||
}
|
||||
if encoded := values.Encode(); encoded != "" {
|
||||
fullURL = fullURL + "?" + encoded
|
||||
}
|
||||
}
|
||||
return fullURL
|
||||
}
|
||||
|
||||
@@ -19,15 +19,8 @@ import (
|
||||
)
|
||||
|
||||
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"
|
||||
pathGetUsageLimits = "getUsageLimits"
|
||||
pathListAvailableModels = "ListAvailableModels"
|
||||
)
|
||||
|
||||
// KiroAuth handles AWS CodeWhisperer authentication and API communication.
|
||||
@@ -35,7 +28,6 @@ const (
|
||||
// and communicating with the CodeWhisperer API.
|
||||
type KiroAuth struct {
|
||||
httpClient *http.Client
|
||||
endpoint string
|
||||
}
|
||||
|
||||
// NewKiroAuth creates a new Kiro authentication service.
|
||||
@@ -49,7 +41,6 @@ type KiroAuth struct {
|
||||
func NewKiroAuth(cfg *config.Config) *KiroAuth {
|
||||
return &KiroAuth{
|
||||
httpClient: util.SetProxy(&cfg.SDKConfig, &http.Client{Timeout: 120 * time.Second}),
|
||||
endpoint: awsKiroEndpoint,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -110,33 +101,30 @@ func (k *KiroAuth) IsTokenExpired(tokenData *KiroTokenData) bool {
|
||||
return time.Now().After(expiresAt)
|
||||
}
|
||||
|
||||
// makeRequest sends a request to the CodeWhisperer API.
|
||||
// This is an internal method for making authenticated API calls.
|
||||
// makeRequest sends a REST-style GET request to the CodeWhisperer API.
|
||||
//
|
||||
// Parameters:
|
||||
// - ctx: The context for the request
|
||||
// - target: The API target (e.g., "AmazonCodeWhispererService.GetUsageLimits")
|
||||
// - accessToken: The OAuth access token
|
||||
// - payload: The request payload
|
||||
// - path: The API path (e.g., "getUsageLimits")
|
||||
// - tokenData: The token data containing access token, refresh token, and profile ARN
|
||||
// - queryParams: Query parameters to add to the URL
|
||||
//
|
||||
// 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)
|
||||
}
|
||||
func (k *KiroAuth) makeRequest(ctx context.Context, path string, tokenData *KiroTokenData, queryParams map[string]string) ([]byte, error) {
|
||||
// Get endpoint from profileArn (defaults to us-east-1 if empty)
|
||||
profileArn := queryParams["profileArn"]
|
||||
endpoint := GetKiroAPIEndpointFromProfileArn(profileArn)
|
||||
url := buildURL(endpoint, path, queryParams)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, k.endpoint, strings.NewReader(string(jsonBody)))
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
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")
|
||||
accountKey := GetAccountKey(tokenData.ClientID, tokenData.RefreshToken)
|
||||
setRuntimeHeaders(req, tokenData.AccessToken, accountKey)
|
||||
|
||||
resp, err := k.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -171,13 +159,13 @@ func (k *KiroAuth) makeRequest(ctx context.Context, target string, accessToken s
|
||||
// - *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{}{
|
||||
queryParams := map[string]string{
|
||||
"origin": "AI_EDITOR",
|
||||
"profileArn": tokenData.ProfileArn,
|
||||
"resourceType": "AGENTIC_REQUEST",
|
||||
}
|
||||
|
||||
body, err := k.makeRequest(ctx, targetGetUsage, tokenData.AccessToken, payload)
|
||||
body, err := k.makeRequest(ctx, pathGetUsageLimits, tokenData, queryParams)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -221,12 +209,12 @@ func (k *KiroAuth) GetUsageLimits(ctx context.Context, tokenData *KiroTokenData)
|
||||
// - []*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{}{
|
||||
queryParams := map[string]string{
|
||||
"origin": "AI_EDITOR",
|
||||
"profileArn": tokenData.ProfileArn,
|
||||
}
|
||||
|
||||
body, err := k.makeRequest(ctx, targetListModels, tokenData.AccessToken, payload)
|
||||
body, err := k.makeRequest(ctx, pathListAvailableModels, tokenData, queryParams)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ package kiro
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
@@ -217,7 +218,8 @@ func TestGenerateTokenFileName(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
tokenData *KiroTokenData
|
||||
expected string
|
||||
exact string // exact match (for cases with email)
|
||||
prefix string // prefix match (for cases without email, where sequence is appended)
|
||||
}{
|
||||
{
|
||||
name: "IDC with email",
|
||||
@@ -226,7 +228,7 @@ func TestGenerateTokenFileName(t *testing.T) {
|
||||
Email: "user@example.com",
|
||||
StartURL: "https://d-1234567890.awsapps.com/start",
|
||||
},
|
||||
expected: "kiro-idc-user-example-com.json",
|
||||
exact: "kiro-idc-user-example-com.json",
|
||||
},
|
||||
{
|
||||
name: "IDC without email but with startUrl",
|
||||
@@ -235,7 +237,7 @@ func TestGenerateTokenFileName(t *testing.T) {
|
||||
Email: "",
|
||||
StartURL: "https://d-1234567890.awsapps.com/start",
|
||||
},
|
||||
expected: "kiro-idc-d-1234567890.json",
|
||||
prefix: "kiro-idc-d-1234567890-",
|
||||
},
|
||||
{
|
||||
name: "IDC with company name in startUrl",
|
||||
@@ -244,7 +246,7 @@ func TestGenerateTokenFileName(t *testing.T) {
|
||||
Email: "",
|
||||
StartURL: "https://my-company.awsapps.com/start",
|
||||
},
|
||||
expected: "kiro-idc-my-company.json",
|
||||
prefix: "kiro-idc-my-company-",
|
||||
},
|
||||
{
|
||||
name: "IDC without email and without startUrl",
|
||||
@@ -253,7 +255,7 @@ func TestGenerateTokenFileName(t *testing.T) {
|
||||
Email: "",
|
||||
StartURL: "",
|
||||
},
|
||||
expected: "kiro-idc.json",
|
||||
prefix: "kiro-idc-",
|
||||
},
|
||||
{
|
||||
name: "Builder ID with email",
|
||||
@@ -262,7 +264,7 @@ func TestGenerateTokenFileName(t *testing.T) {
|
||||
Email: "user@gmail.com",
|
||||
StartURL: "https://view.awsapps.com/start",
|
||||
},
|
||||
expected: "kiro-builder-id-user-gmail-com.json",
|
||||
exact: "kiro-builder-id-user-gmail-com.json",
|
||||
},
|
||||
{
|
||||
name: "Builder ID without email",
|
||||
@@ -271,7 +273,7 @@ func TestGenerateTokenFileName(t *testing.T) {
|
||||
Email: "",
|
||||
StartURL: "https://view.awsapps.com/start",
|
||||
},
|
||||
expected: "kiro-builder-id.json",
|
||||
prefix: "kiro-builder-id-",
|
||||
},
|
||||
{
|
||||
name: "Social auth with email",
|
||||
@@ -279,7 +281,7 @@ func TestGenerateTokenFileName(t *testing.T) {
|
||||
AuthMethod: "google",
|
||||
Email: "user@gmail.com",
|
||||
},
|
||||
expected: "kiro-google-user-gmail-com.json",
|
||||
exact: "kiro-google-user-gmail-com.json",
|
||||
},
|
||||
{
|
||||
name: "Empty auth method",
|
||||
@@ -287,7 +289,7 @@ func TestGenerateTokenFileName(t *testing.T) {
|
||||
AuthMethod: "",
|
||||
Email: "",
|
||||
},
|
||||
expected: "kiro-unknown.json",
|
||||
prefix: "kiro-unknown-",
|
||||
},
|
||||
{
|
||||
name: "Email with special characters",
|
||||
@@ -296,16 +298,454 @@ func TestGenerateTokenFileName(t *testing.T) {
|
||||
Email: "user.name+tag@sub.example.com",
|
||||
StartURL: "https://d-1234567890.awsapps.com/start",
|
||||
},
|
||||
expected: "kiro-idc-user-name+tag-sub-example-com.json",
|
||||
exact: "kiro-idc-user-name+tag-sub-example-com.json",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := GenerateTokenFileName(tt.tokenData)
|
||||
if result != tt.expected {
|
||||
t.Errorf("GenerateTokenFileName() = %q, want %q", result, tt.expected)
|
||||
if tt.exact != "" {
|
||||
if result != tt.exact {
|
||||
t.Errorf("GenerateTokenFileName() = %q, want %q", result, tt.exact)
|
||||
}
|
||||
} else if tt.prefix != "" {
|
||||
if !strings.HasPrefix(result, tt.prefix) || !strings.HasSuffix(result, ".json") {
|
||||
t.Errorf("GenerateTokenFileName() = %q, want prefix %q with .json suffix", result, tt.prefix)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseProfileARN(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
arn string
|
||||
expected *ProfileARN
|
||||
}{
|
||||
{
|
||||
name: "Empty ARN",
|
||||
arn: "",
|
||||
expected: nil,
|
||||
},
|
||||
{
|
||||
name: "Invalid format - too few parts",
|
||||
arn: "arn:aws:codewhisperer",
|
||||
expected: nil,
|
||||
},
|
||||
{
|
||||
name: "Invalid prefix - not arn",
|
||||
arn: "notarn:aws:codewhisperer:us-east-1:123456789012:profile/ABC",
|
||||
expected: nil,
|
||||
},
|
||||
{
|
||||
name: "Invalid service - not codewhisperer",
|
||||
arn: "arn:aws:s3:us-east-1:123456789012:bucket/mybucket",
|
||||
expected: nil,
|
||||
},
|
||||
{
|
||||
name: "Invalid region - no hyphen",
|
||||
arn: "arn:aws:codewhisperer:useast1:123456789012:profile/ABC",
|
||||
expected: nil,
|
||||
},
|
||||
{
|
||||
name: "Empty partition",
|
||||
arn: "arn::codewhisperer:us-east-1:123456789012:profile/ABC",
|
||||
expected: nil,
|
||||
},
|
||||
{
|
||||
name: "Empty region",
|
||||
arn: "arn:aws:codewhisperer::123456789012:profile/ABC",
|
||||
expected: nil,
|
||||
},
|
||||
{
|
||||
name: "Valid ARN - us-east-1",
|
||||
arn: "arn:aws:codewhisperer:us-east-1:123456789012:profile/ABCDEFGHIJKL",
|
||||
expected: &ProfileARN{
|
||||
Raw: "arn:aws:codewhisperer:us-east-1:123456789012:profile/ABCDEFGHIJKL",
|
||||
Partition: "aws",
|
||||
Service: "codewhisperer",
|
||||
Region: "us-east-1",
|
||||
AccountID: "123456789012",
|
||||
ResourceType: "profile",
|
||||
ResourceID: "ABCDEFGHIJKL",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Valid ARN - ap-southeast-1",
|
||||
arn: "arn:aws:codewhisperer:ap-southeast-1:987654321098:profile/ZYXWVUTSRQ",
|
||||
expected: &ProfileARN{
|
||||
Raw: "arn:aws:codewhisperer:ap-southeast-1:987654321098:profile/ZYXWVUTSRQ",
|
||||
Partition: "aws",
|
||||
Service: "codewhisperer",
|
||||
Region: "ap-southeast-1",
|
||||
AccountID: "987654321098",
|
||||
ResourceType: "profile",
|
||||
ResourceID: "ZYXWVUTSRQ",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Valid ARN - eu-west-1",
|
||||
arn: "arn:aws:codewhisperer:eu-west-1:111222333444:profile/PROFILE123",
|
||||
expected: &ProfileARN{
|
||||
Raw: "arn:aws:codewhisperer:eu-west-1:111222333444:profile/PROFILE123",
|
||||
Partition: "aws",
|
||||
Service: "codewhisperer",
|
||||
Region: "eu-west-1",
|
||||
AccountID: "111222333444",
|
||||
ResourceType: "profile",
|
||||
ResourceID: "PROFILE123",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Valid ARN - aws-cn partition",
|
||||
arn: "arn:aws-cn:codewhisperer:cn-north-1:123456789012:profile/CHINAID",
|
||||
expected: &ProfileARN{
|
||||
Raw: "arn:aws-cn:codewhisperer:cn-north-1:123456789012:profile/CHINAID",
|
||||
Partition: "aws-cn",
|
||||
Service: "codewhisperer",
|
||||
Region: "cn-north-1",
|
||||
AccountID: "123456789012",
|
||||
ResourceType: "profile",
|
||||
ResourceID: "CHINAID",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Valid ARN - resource without slash",
|
||||
arn: "arn:aws:codewhisperer:us-west-2:123456789012:profile",
|
||||
expected: &ProfileARN{
|
||||
Raw: "arn:aws:codewhisperer:us-west-2:123456789012:profile",
|
||||
Partition: "aws",
|
||||
Service: "codewhisperer",
|
||||
Region: "us-west-2",
|
||||
AccountID: "123456789012",
|
||||
ResourceType: "profile",
|
||||
ResourceID: "",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Valid ARN - resource with colon",
|
||||
arn: "arn:aws:codewhisperer:us-east-1:123456789012:profile/ABC:extra",
|
||||
expected: &ProfileARN{
|
||||
Raw: "arn:aws:codewhisperer:us-east-1:123456789012:profile/ABC:extra",
|
||||
Partition: "aws",
|
||||
Service: "codewhisperer",
|
||||
Region: "us-east-1",
|
||||
AccountID: "123456789012",
|
||||
ResourceType: "profile",
|
||||
ResourceID: "ABC:extra",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := ParseProfileARN(tt.arn)
|
||||
if tt.expected == nil {
|
||||
if result != nil {
|
||||
t.Errorf("ParseProfileARN(%q) = %+v, want nil", tt.arn, result)
|
||||
}
|
||||
return
|
||||
}
|
||||
if result == nil {
|
||||
t.Errorf("ParseProfileARN(%q) = nil, want %+v", tt.arn, tt.expected)
|
||||
return
|
||||
}
|
||||
if result.Raw != tt.expected.Raw {
|
||||
t.Errorf("Raw = %q, want %q", result.Raw, tt.expected.Raw)
|
||||
}
|
||||
if result.Partition != tt.expected.Partition {
|
||||
t.Errorf("Partition = %q, want %q", result.Partition, tt.expected.Partition)
|
||||
}
|
||||
if result.Service != tt.expected.Service {
|
||||
t.Errorf("Service = %q, want %q", result.Service, tt.expected.Service)
|
||||
}
|
||||
if result.Region != tt.expected.Region {
|
||||
t.Errorf("Region = %q, want %q", result.Region, tt.expected.Region)
|
||||
}
|
||||
if result.AccountID != tt.expected.AccountID {
|
||||
t.Errorf("AccountID = %q, want %q", result.AccountID, tt.expected.AccountID)
|
||||
}
|
||||
if result.ResourceType != tt.expected.ResourceType {
|
||||
t.Errorf("ResourceType = %q, want %q", result.ResourceType, tt.expected.ResourceType)
|
||||
}
|
||||
if result.ResourceID != tt.expected.ResourceID {
|
||||
t.Errorf("ResourceID = %q, want %q", result.ResourceID, tt.expected.ResourceID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractRegionFromProfileArn(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
profileArn string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "Empty ARN",
|
||||
profileArn: "",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "Invalid ARN",
|
||||
profileArn: "invalid-arn",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "Valid ARN - us-east-1",
|
||||
profileArn: "arn:aws:codewhisperer:us-east-1:123456789012:profile/ABC",
|
||||
expected: "us-east-1",
|
||||
},
|
||||
{
|
||||
name: "Valid ARN - ap-southeast-1",
|
||||
profileArn: "arn:aws:codewhisperer:ap-southeast-1:123456789012:profile/ABC",
|
||||
expected: "ap-southeast-1",
|
||||
},
|
||||
{
|
||||
name: "Valid ARN - eu-central-1",
|
||||
profileArn: "arn:aws:codewhisperer:eu-central-1:123456789012:profile/ABC",
|
||||
expected: "eu-central-1",
|
||||
},
|
||||
{
|
||||
name: "Non-codewhisperer ARN",
|
||||
profileArn: "arn:aws:s3:us-east-1:123456789012:bucket/mybucket",
|
||||
expected: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := ExtractRegionFromProfileArn(tt.profileArn)
|
||||
if result != tt.expected {
|
||||
t.Errorf("ExtractRegionFromProfileArn(%q) = %q, want %q", tt.profileArn, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetKiroAPIEndpoint(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
region string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "Empty region - defaults to us-east-1",
|
||||
region: "",
|
||||
expected: "https://q.us-east-1.amazonaws.com",
|
||||
},
|
||||
{
|
||||
name: "us-east-1",
|
||||
region: "us-east-1",
|
||||
expected: "https://q.us-east-1.amazonaws.com",
|
||||
},
|
||||
{
|
||||
name: "us-west-2",
|
||||
region: "us-west-2",
|
||||
expected: "https://q.us-west-2.amazonaws.com",
|
||||
},
|
||||
{
|
||||
name: "ap-southeast-1",
|
||||
region: "ap-southeast-1",
|
||||
expected: "https://q.ap-southeast-1.amazonaws.com",
|
||||
},
|
||||
{
|
||||
name: "eu-west-1",
|
||||
region: "eu-west-1",
|
||||
expected: "https://q.eu-west-1.amazonaws.com",
|
||||
},
|
||||
{
|
||||
name: "cn-north-1",
|
||||
region: "cn-north-1",
|
||||
expected: "https://q.cn-north-1.amazonaws.com",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := GetKiroAPIEndpoint(tt.region)
|
||||
if result != tt.expected {
|
||||
t.Errorf("GetKiroAPIEndpoint(%q) = %q, want %q", tt.region, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetKiroAPIEndpointFromProfileArn(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
profileArn string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "Empty ARN - defaults to us-east-1",
|
||||
profileArn: "",
|
||||
expected: "https://q.us-east-1.amazonaws.com",
|
||||
},
|
||||
{
|
||||
name: "Invalid ARN - defaults to us-east-1",
|
||||
profileArn: "invalid-arn",
|
||||
expected: "https://q.us-east-1.amazonaws.com",
|
||||
},
|
||||
{
|
||||
name: "Valid ARN - us-east-1",
|
||||
profileArn: "arn:aws:codewhisperer:us-east-1:123456789012:profile/ABC",
|
||||
expected: "https://q.us-east-1.amazonaws.com",
|
||||
},
|
||||
{
|
||||
name: "Valid ARN - ap-southeast-1",
|
||||
profileArn: "arn:aws:codewhisperer:ap-southeast-1:123456789012:profile/ABC",
|
||||
expected: "https://q.ap-southeast-1.amazonaws.com",
|
||||
},
|
||||
{
|
||||
name: "Valid ARN - eu-central-1",
|
||||
profileArn: "arn:aws:codewhisperer:eu-central-1:123456789012:profile/ABC",
|
||||
expected: "https://q.eu-central-1.amazonaws.com",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := GetKiroAPIEndpointFromProfileArn(tt.profileArn)
|
||||
if result != tt.expected {
|
||||
t.Errorf("GetKiroAPIEndpointFromProfileArn(%q) = %q, want %q", tt.profileArn, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetCodeWhispererLegacyEndpoint(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
region string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "Empty region - defaults to us-east-1",
|
||||
region: "",
|
||||
expected: "https://codewhisperer.us-east-1.amazonaws.com",
|
||||
},
|
||||
{
|
||||
name: "us-east-1",
|
||||
region: "us-east-1",
|
||||
expected: "https://codewhisperer.us-east-1.amazonaws.com",
|
||||
},
|
||||
{
|
||||
name: "us-west-2",
|
||||
region: "us-west-2",
|
||||
expected: "https://codewhisperer.us-west-2.amazonaws.com",
|
||||
},
|
||||
{
|
||||
name: "ap-northeast-1",
|
||||
region: "ap-northeast-1",
|
||||
expected: "https://codewhisperer.ap-northeast-1.amazonaws.com",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := GetCodeWhispererLegacyEndpoint(tt.region)
|
||||
if result != tt.expected {
|
||||
t.Errorf("GetCodeWhispererLegacyEndpoint(%q) = %q, want %q", tt.region, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractRegionFromMetadata(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
metadata map[string]interface{}
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "Nil metadata - defaults to us-east-1",
|
||||
metadata: nil,
|
||||
expected: "us-east-1",
|
||||
},
|
||||
{
|
||||
name: "Empty metadata - defaults to us-east-1",
|
||||
metadata: map[string]interface{}{},
|
||||
expected: "us-east-1",
|
||||
},
|
||||
{
|
||||
name: "Priority 1: api_region override",
|
||||
metadata: map[string]interface{}{
|
||||
"api_region": "eu-west-1",
|
||||
"profile_arn": "arn:aws:codewhisperer:us-east-1:123456789012:profile/ABC",
|
||||
},
|
||||
expected: "eu-west-1",
|
||||
},
|
||||
{
|
||||
name: "Priority 2: profile_arn when api_region is empty",
|
||||
metadata: map[string]interface{}{
|
||||
"api_region": "",
|
||||
"profile_arn": "arn:aws:codewhisperer:ap-southeast-1:123456789012:profile/ABC",
|
||||
},
|
||||
expected: "ap-southeast-1",
|
||||
},
|
||||
{
|
||||
name: "Priority 2: profile_arn when api_region is missing",
|
||||
metadata: map[string]interface{}{
|
||||
"profile_arn": "arn:aws:codewhisperer:eu-central-1:123456789012:profile/ABC",
|
||||
},
|
||||
expected: "eu-central-1",
|
||||
},
|
||||
{
|
||||
name: "Fallback: default when profile_arn is invalid",
|
||||
metadata: map[string]interface{}{
|
||||
"profile_arn": "invalid-arn",
|
||||
},
|
||||
expected: "us-east-1",
|
||||
},
|
||||
{
|
||||
name: "Fallback: default when profile_arn is empty",
|
||||
metadata: map[string]interface{}{
|
||||
"profile_arn": "",
|
||||
},
|
||||
expected: "us-east-1",
|
||||
},
|
||||
{
|
||||
name: "OIDC region is NOT used for API region",
|
||||
metadata: map[string]interface{}{
|
||||
"region": "ap-northeast-2", // OIDC region - should be ignored
|
||||
},
|
||||
expected: "us-east-1",
|
||||
},
|
||||
{
|
||||
name: "api_region takes precedence over OIDC region",
|
||||
metadata: map[string]interface{}{
|
||||
"api_region": "us-west-2",
|
||||
"region": "ap-northeast-2", // OIDC region - should be ignored
|
||||
},
|
||||
expected: "us-west-2",
|
||||
},
|
||||
{
|
||||
name: "Non-string api_region is ignored",
|
||||
metadata: map[string]interface{}{
|
||||
"api_region": 123, // wrong type
|
||||
"profile_arn": "arn:aws:codewhisperer:ap-south-1:123456789012:profile/ABC",
|
||||
},
|
||||
expected: "ap-south-1",
|
||||
},
|
||||
{
|
||||
name: "Non-string profile_arn is ignored",
|
||||
metadata: map[string]interface{}{
|
||||
"profile_arn": 123, // wrong type
|
||||
},
|
||||
expected: "us-east-1",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := ExtractRegionFromMetadata(tt.metadata)
|
||||
if result != tt.expected {
|
||||
t.Errorf("ExtractRegionFromMetadata(%v) = %q, want %q", tt.metadata, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -9,30 +9,23 @@ import (
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
codeWhispererAPI = "https://codewhisperer.us-east-1.amazonaws.com"
|
||||
kiroVersion = "0.6.18"
|
||||
)
|
||||
|
||||
// CodeWhispererClient handles CodeWhisperer API calls.
|
||||
type CodeWhispererClient struct {
|
||||
httpClient *http.Client
|
||||
machineID string
|
||||
}
|
||||
|
||||
// UsageLimitsResponse represents the getUsageLimits API response.
|
||||
type UsageLimitsResponse struct {
|
||||
DaysUntilReset *int `json:"daysUntilReset,omitempty"`
|
||||
NextDateReset *float64 `json:"nextDateReset,omitempty"`
|
||||
UserInfo *UserInfo `json:"userInfo,omitempty"`
|
||||
SubscriptionInfo *SubscriptionInfo `json:"subscriptionInfo,omitempty"`
|
||||
UsageBreakdownList []UsageBreakdown `json:"usageBreakdownList,omitempty"`
|
||||
DaysUntilReset *int `json:"daysUntilReset,omitempty"`
|
||||
NextDateReset *float64 `json:"nextDateReset,omitempty"`
|
||||
UserInfo *UserInfo `json:"userInfo,omitempty"`
|
||||
SubscriptionInfo *SubscriptionInfo `json:"subscriptionInfo,omitempty"`
|
||||
UsageBreakdownList []UsageBreakdown `json:"usageBreakdownList,omitempty"`
|
||||
}
|
||||
|
||||
// UserInfo contains user information from the API.
|
||||
@@ -49,13 +42,13 @@ type SubscriptionInfo struct {
|
||||
|
||||
// UsageBreakdown contains usage details.
|
||||
type UsageBreakdown struct {
|
||||
UsageLimit *int `json:"usageLimit,omitempty"`
|
||||
CurrentUsage *int `json:"currentUsage,omitempty"`
|
||||
UsageLimitWithPrecision *float64 `json:"usageLimitWithPrecision,omitempty"`
|
||||
CurrentUsageWithPrecision *float64 `json:"currentUsageWithPrecision,omitempty"`
|
||||
NextDateReset *float64 `json:"nextDateReset,omitempty"`
|
||||
DisplayName string `json:"displayName,omitempty"`
|
||||
ResourceType string `json:"resourceType,omitempty"`
|
||||
UsageLimit *int `json:"usageLimit,omitempty"`
|
||||
CurrentUsage *int `json:"currentUsage,omitempty"`
|
||||
UsageLimitWithPrecision *float64 `json:"usageLimitWithPrecision,omitempty"`
|
||||
CurrentUsageWithPrecision *float64 `json:"currentUsageWithPrecision,omitempty"`
|
||||
NextDateReset *float64 `json:"nextDateReset,omitempty"`
|
||||
DisplayName string `json:"displayName,omitempty"`
|
||||
ResourceType string `json:"resourceType,omitempty"`
|
||||
}
|
||||
|
||||
// NewCodeWhispererClient creates a new CodeWhisperer client.
|
||||
@@ -64,40 +57,34 @@ func NewCodeWhispererClient(cfg *config.Config, machineID string) *CodeWhisperer
|
||||
if cfg != nil {
|
||||
client = util.SetProxy(&cfg.SDKConfig, client)
|
||||
}
|
||||
if machineID == "" {
|
||||
machineID = uuid.New().String()
|
||||
}
|
||||
return &CodeWhispererClient{
|
||||
httpClient: client,
|
||||
machineID: machineID,
|
||||
}
|
||||
}
|
||||
|
||||
// generateInvocationID generates a unique invocation ID.
|
||||
func generateInvocationID() string {
|
||||
return uuid.New().String()
|
||||
}
|
||||
|
||||
// GetUsageLimits fetches usage limits and user info from CodeWhisperer API.
|
||||
// This is the recommended way to get user email after login.
|
||||
func (c *CodeWhispererClient) GetUsageLimits(ctx context.Context, accessToken string) (*UsageLimitsResponse, error) {
|
||||
url := fmt.Sprintf("%s/getUsageLimits?isEmailRequired=true&origin=AI_EDITOR&resourceType=AGENTIC_REQUEST", codeWhispererAPI)
|
||||
func (c *CodeWhispererClient) GetUsageLimits(ctx context.Context, accessToken, clientID, refreshToken, profileArn string) (*UsageLimitsResponse, error) {
|
||||
queryParams := map[string]string{
|
||||
"origin": "AI_EDITOR",
|
||||
"resourceType": "AGENTIC_REQUEST",
|
||||
}
|
||||
// Determine endpoint based on profileArn region
|
||||
endpoint := GetKiroAPIEndpointFromProfileArn(profileArn)
|
||||
if profileArn != "" {
|
||||
queryParams["profileArn"] = profileArn
|
||||
} else {
|
||||
queryParams["isEmailRequired"] = "true"
|
||||
}
|
||||
url := buildURL(endpoint, pathGetUsageLimits, queryParams)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
// Set headers to match Kiro IDE
|
||||
xAmzUserAgent := fmt.Sprintf("aws-sdk-js/1.0.0 KiroIDE-%s-%s", kiroVersion, c.machineID)
|
||||
userAgent := fmt.Sprintf("aws-sdk-js/1.0.0 ua/2.1 os/windows lang/js md/nodejs#20.16.0 api/codewhispererruntime#1.0.0 m/E KiroIDE-%s-%s", kiroVersion, c.machineID)
|
||||
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
req.Header.Set("x-amz-user-agent", xAmzUserAgent)
|
||||
req.Header.Set("User-Agent", userAgent)
|
||||
req.Header.Set("amz-sdk-invocation-id", generateInvocationID())
|
||||
req.Header.Set("amz-sdk-request", "attempt=1; max=1")
|
||||
req.Header.Set("Connection", "close")
|
||||
accountKey := GetAccountKey(clientID, refreshToken)
|
||||
setRuntimeHeaders(req, accessToken, accountKey)
|
||||
|
||||
log.Debugf("codewhisperer: GET %s", url)
|
||||
|
||||
@@ -128,8 +115,8 @@ func (c *CodeWhispererClient) GetUsageLimits(ctx context.Context, accessToken st
|
||||
|
||||
// FetchUserEmailFromAPI fetches user email using CodeWhisperer getUsageLimits API.
|
||||
// This is more reliable than JWT parsing as it uses the official API.
|
||||
func (c *CodeWhispererClient) FetchUserEmailFromAPI(ctx context.Context, accessToken string) string {
|
||||
resp, err := c.GetUsageLimits(ctx, accessToken)
|
||||
func (c *CodeWhispererClient) FetchUserEmailFromAPI(ctx context.Context, accessToken, clientID, refreshToken string) string {
|
||||
resp, err := c.GetUsageLimits(ctx, accessToken, clientID, refreshToken, "")
|
||||
if err != nil {
|
||||
log.Debugf("codewhisperer: failed to get usage limits: %v", err)
|
||||
return ""
|
||||
@@ -146,10 +133,10 @@ func (c *CodeWhispererClient) FetchUserEmailFromAPI(ctx context.Context, accessT
|
||||
|
||||
// FetchUserEmailWithFallback fetches user email with multiple fallback methods.
|
||||
// Priority: 1. CodeWhisperer API 2. userinfo endpoint 3. JWT parsing
|
||||
func FetchUserEmailWithFallback(ctx context.Context, cfg *config.Config, accessToken string) string {
|
||||
func FetchUserEmailWithFallback(ctx context.Context, cfg *config.Config, accessToken, clientID, refreshToken string) string {
|
||||
// Method 1: Try CodeWhisperer API (most reliable)
|
||||
cwClient := NewCodeWhispererClient(cfg, "")
|
||||
email := cwClient.FetchUserEmailFromAPI(ctx, accessToken)
|
||||
email := cwClient.FetchUserEmailFromAPI(ctx, accessToken, clientID, refreshToken)
|
||||
if email != "" {
|
||||
return email
|
||||
}
|
||||
|
||||
@@ -2,77 +2,105 @@ package kiro
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"net/http"
|
||||
"runtime"
|
||||
"slices"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// Fingerprint 多维度指纹信息
|
||||
// Fingerprint holds multi-dimensional fingerprint data for runtime request disguise.
|
||||
type Fingerprint struct {
|
||||
SDKVersion string // 1.0.20-1.0.27
|
||||
OIDCSDKVersion string // 3.7xx (AWS SDK JS)
|
||||
RuntimeSDKVersion string // 1.0.x (runtime API)
|
||||
StreamingSDKVersion string // 1.0.x (streaming API)
|
||||
OSType string // darwin/windows/linux
|
||||
OSVersion string // 10.0.22621
|
||||
NodeVersion string // 18.x/20.x/22.x
|
||||
KiroVersion string // 0.3.x-0.8.x
|
||||
OSVersion string
|
||||
NodeVersion string
|
||||
KiroVersion string
|
||||
KiroHash string // SHA256
|
||||
AcceptLanguage string
|
||||
ScreenResolution string // 1920x1080
|
||||
ColorDepth int // 24
|
||||
HardwareConcurrency int // CPU 核心数
|
||||
TimezoneOffset int
|
||||
}
|
||||
|
||||
// FingerprintManager 指纹管理器
|
||||
// FingerprintConfig holds external fingerprint overrides.
|
||||
type FingerprintConfig struct {
|
||||
OIDCSDKVersion string
|
||||
RuntimeSDKVersion string
|
||||
StreamingSDKVersion string
|
||||
OSType string
|
||||
OSVersion string
|
||||
NodeVersion string
|
||||
KiroVersion string
|
||||
KiroHash string
|
||||
}
|
||||
|
||||
// FingerprintManager manages per-account fingerprint generation and caching.
|
||||
type FingerprintManager struct {
|
||||
mu sync.RWMutex
|
||||
fingerprints map[string]*Fingerprint // tokenKey -> fingerprint
|
||||
rng *rand.Rand
|
||||
config *FingerprintConfig // External config (Optional)
|
||||
}
|
||||
|
||||
var (
|
||||
sdkVersions = []string{
|
||||
"1.0.20", "1.0.21", "1.0.22", "1.0.23",
|
||||
"1.0.24", "1.0.25", "1.0.26", "1.0.27",
|
||||
// SDK versions
|
||||
oidcSDKVersions = []string{
|
||||
"3.980.0", "3.975.0", "3.972.0", "3.808.0",
|
||||
"3.738.0", "3.737.0", "3.736.0", "3.735.0",
|
||||
}
|
||||
// SDKVersions for getUsageLimits/ListAvailableModels/GetProfile (runtime API)
|
||||
runtimeSDKVersions = []string{"1.0.0"}
|
||||
// SDKVersions for generateAssistantResponse (streaming API)
|
||||
streamingSDKVersions = []string{"1.0.27"}
|
||||
// Valid OS types
|
||||
osTypes = []string{"darwin", "windows", "linux"}
|
||||
// OS versions
|
||||
osVersions = map[string][]string{
|
||||
"darwin": {"14.0", "14.1", "14.2", "14.3", "14.4", "14.5", "15.0", "15.1"},
|
||||
"windows": {"10.0.19041", "10.0.19042", "10.0.19043", "10.0.19044", "10.0.22621", "10.0.22631"},
|
||||
"linux": {"5.15.0", "6.1.0", "6.2.0", "6.5.0", "6.6.0", "6.8.0"},
|
||||
"darwin": {"25.2.0", "25.1.0", "25.0.0", "24.5.0", "24.4.0", "24.3.0"},
|
||||
"windows": {"10.0.26200", "10.0.26100", "10.0.22631", "10.0.22621", "10.0.19045"},
|
||||
"linux": {"6.12.0", "6.11.0", "6.8.0", "6.6.0", "6.5.0", "6.1.0"},
|
||||
}
|
||||
// Node versions
|
||||
nodeVersions = []string{
|
||||
"18.17.0", "18.18.0", "18.19.0", "18.20.0",
|
||||
"20.9.0", "20.10.0", "20.11.0", "20.12.0", "20.13.0",
|
||||
"22.0.0", "22.1.0", "22.2.0", "22.3.0",
|
||||
"22.21.1", "22.21.0", "22.20.0", "22.19.0", "22.18.0",
|
||||
"20.18.0", "20.17.0", "20.16.0",
|
||||
}
|
||||
// Kiro IDE versions
|
||||
kiroVersions = []string{
|
||||
"0.3.0", "0.3.1", "0.4.0", "0.4.1", "0.5.0", "0.5.1",
|
||||
"0.6.0", "0.6.1", "0.7.0", "0.7.1", "0.8.0", "0.8.1",
|
||||
"0.10.32", "0.10.16", "0.10.10",
|
||||
"0.9.47", "0.9.40", "0.9.2",
|
||||
"0.8.206", "0.8.140", "0.8.135", "0.8.86",
|
||||
}
|
||||
acceptLanguages = []string{
|
||||
"en-US,en;q=0.9",
|
||||
"en-GB,en;q=0.9",
|
||||
"zh-CN,zh;q=0.9,en;q=0.8",
|
||||
"zh-TW,zh;q=0.9,en;q=0.8",
|
||||
"ja-JP,ja;q=0.9,en;q=0.8",
|
||||
"ko-KR,ko;q=0.9,en;q=0.8",
|
||||
"de-DE,de;q=0.9,en;q=0.8",
|
||||
"fr-FR,fr;q=0.9,en;q=0.8",
|
||||
}
|
||||
screenResolutions = []string{
|
||||
"1920x1080", "2560x1440", "3840x2160",
|
||||
"1366x768", "1440x900", "1680x1050",
|
||||
"2560x1600", "3440x1440",
|
||||
}
|
||||
colorDepths = []int{24, 32}
|
||||
hardwareConcurrencies = []int{4, 6, 8, 10, 12, 16, 20, 24, 32}
|
||||
timezoneOffsets = []int{-480, -420, -360, -300, -240, 0, 60, 120, 480, 540}
|
||||
// Global singleton
|
||||
globalFingerprintManager *FingerprintManager
|
||||
globalFingerprintManagerOnce sync.Once
|
||||
)
|
||||
|
||||
// NewFingerprintManager 创建指纹管理器
|
||||
func GlobalFingerprintManager() *FingerprintManager {
|
||||
globalFingerprintManagerOnce.Do(func() {
|
||||
globalFingerprintManager = NewFingerprintManager()
|
||||
})
|
||||
return globalFingerprintManager
|
||||
}
|
||||
|
||||
func SetGlobalFingerprintConfig(cfg *FingerprintConfig) {
|
||||
GlobalFingerprintManager().SetConfig(cfg)
|
||||
}
|
||||
|
||||
// SetConfig applies the config and clears the fingerprint cache.
|
||||
func (fm *FingerprintManager) SetConfig(cfg *FingerprintConfig) {
|
||||
fm.mu.Lock()
|
||||
defer fm.mu.Unlock()
|
||||
fm.config = cfg
|
||||
// Clear cached fingerprints so they regenerate with the new config
|
||||
fm.fingerprints = make(map[string]*Fingerprint)
|
||||
}
|
||||
|
||||
func NewFingerprintManager() *FingerprintManager {
|
||||
return &FingerprintManager{
|
||||
fingerprints: make(map[string]*Fingerprint),
|
||||
@@ -80,7 +108,7 @@ func NewFingerprintManager() *FingerprintManager {
|
||||
}
|
||||
}
|
||||
|
||||
// GetFingerprint 获取或生成 Token 关联的指纹
|
||||
// GetFingerprint returns the fingerprint for tokenKey, creating one if it doesn't exist.
|
||||
func (fm *FingerprintManager) GetFingerprint(tokenKey string) *Fingerprint {
|
||||
fm.mu.RLock()
|
||||
if fp, exists := fm.fingerprints[tokenKey]; exists {
|
||||
@@ -101,97 +129,150 @@ func (fm *FingerprintManager) GetFingerprint(tokenKey string) *Fingerprint {
|
||||
return fp
|
||||
}
|
||||
|
||||
// generateFingerprint 生成新的指纹
|
||||
func (fm *FingerprintManager) generateFingerprint(tokenKey string) *Fingerprint {
|
||||
osType := fm.randomChoice(osTypes)
|
||||
osVersion := fm.randomChoice(osVersions[osType])
|
||||
kiroVersion := fm.randomChoice(kiroVersions)
|
||||
if fm.config != nil {
|
||||
return fm.generateFromConfig(tokenKey)
|
||||
}
|
||||
return fm.generateRandom(tokenKey)
|
||||
}
|
||||
|
||||
fp := &Fingerprint{
|
||||
SDKVersion: fm.randomChoice(sdkVersions),
|
||||
OSType: osType,
|
||||
OSVersion: osVersion,
|
||||
NodeVersion: fm.randomChoice(nodeVersions),
|
||||
KiroVersion: kiroVersion,
|
||||
AcceptLanguage: fm.randomChoice(acceptLanguages),
|
||||
ScreenResolution: fm.randomChoice(screenResolutions),
|
||||
ColorDepth: fm.randomIntChoice(colorDepths),
|
||||
HardwareConcurrency: fm.randomIntChoice(hardwareConcurrencies),
|
||||
TimezoneOffset: fm.randomIntChoice(timezoneOffsets),
|
||||
// generateFromConfig uses config values, falling back to random for empty fields.
|
||||
func (fm *FingerprintManager) generateFromConfig(tokenKey string) *Fingerprint {
|
||||
cfg := fm.config
|
||||
|
||||
// Helper: config value or random selection
|
||||
configOrRandom := func(configVal string, choices []string) string {
|
||||
if configVal != "" {
|
||||
return configVal
|
||||
}
|
||||
return choices[fm.rng.Intn(len(choices))]
|
||||
}
|
||||
|
||||
fp.KiroHash = fm.generateKiroHash(tokenKey, kiroVersion, osType)
|
||||
return fp
|
||||
osType := cfg.OSType
|
||||
if osType == "" {
|
||||
osType = runtime.GOOS
|
||||
if !slices.Contains(osTypes, osType) {
|
||||
osType = osTypes[fm.rng.Intn(len(osTypes))]
|
||||
}
|
||||
}
|
||||
|
||||
osVersion := cfg.OSVersion
|
||||
if osVersion == "" {
|
||||
if versions, ok := osVersions[osType]; ok {
|
||||
osVersion = versions[fm.rng.Intn(len(versions))]
|
||||
}
|
||||
}
|
||||
|
||||
kiroHash := cfg.KiroHash
|
||||
if kiroHash == "" {
|
||||
hash := sha256.Sum256([]byte(tokenKey))
|
||||
kiroHash = hex.EncodeToString(hash[:])
|
||||
}
|
||||
|
||||
return &Fingerprint{
|
||||
OIDCSDKVersion: configOrRandom(cfg.OIDCSDKVersion, oidcSDKVersions),
|
||||
RuntimeSDKVersion: configOrRandom(cfg.RuntimeSDKVersion, runtimeSDKVersions),
|
||||
StreamingSDKVersion: configOrRandom(cfg.StreamingSDKVersion, streamingSDKVersions),
|
||||
OSType: osType,
|
||||
OSVersion: osVersion,
|
||||
NodeVersion: configOrRandom(cfg.NodeVersion, nodeVersions),
|
||||
KiroVersion: configOrRandom(cfg.KiroVersion, kiroVersions),
|
||||
KiroHash: kiroHash,
|
||||
}
|
||||
}
|
||||
|
||||
// generateKiroHash 生成 Kiro Hash
|
||||
func (fm *FingerprintManager) generateKiroHash(tokenKey, kiroVersion, osType string) string {
|
||||
data := fmt.Sprintf("%s:%s:%s:%d", tokenKey, kiroVersion, osType, time.Now().UnixNano())
|
||||
hash := sha256.Sum256([]byte(data))
|
||||
return hex.EncodeToString(hash[:])
|
||||
// generateRandom generates a deterministic fingerprint seeded by accountKey hash.
|
||||
func (fm *FingerprintManager) generateRandom(accountKey string) *Fingerprint {
|
||||
// Use accountKey hash as seed for deterministic random selection
|
||||
hash := sha256.Sum256([]byte(accountKey))
|
||||
seed := int64(binary.BigEndian.Uint64(hash[:8]))
|
||||
rng := rand.New(rand.NewSource(seed))
|
||||
|
||||
osType := runtime.GOOS
|
||||
if !slices.Contains(osTypes, osType) {
|
||||
osType = osTypes[rng.Intn(len(osTypes))]
|
||||
}
|
||||
osVersion := osVersions[osType][rng.Intn(len(osVersions[osType]))]
|
||||
|
||||
return &Fingerprint{
|
||||
OIDCSDKVersion: oidcSDKVersions[rng.Intn(len(oidcSDKVersions))],
|
||||
RuntimeSDKVersion: runtimeSDKVersions[rng.Intn(len(runtimeSDKVersions))],
|
||||
StreamingSDKVersion: streamingSDKVersions[rng.Intn(len(streamingSDKVersions))],
|
||||
OSType: osType,
|
||||
OSVersion: osVersion,
|
||||
NodeVersion: nodeVersions[rng.Intn(len(nodeVersions))],
|
||||
KiroVersion: kiroVersions[rng.Intn(len(kiroVersions))],
|
||||
KiroHash: hex.EncodeToString(hash[:]),
|
||||
}
|
||||
}
|
||||
|
||||
// randomChoice 随机选择字符串
|
||||
func (fm *FingerprintManager) randomChoice(choices []string) string {
|
||||
return choices[fm.rng.Intn(len(choices))]
|
||||
// GenerateAccountKey returns a 16-char hex key derived from SHA256(seed).
|
||||
func GenerateAccountKey(seed string) string {
|
||||
hash := sha256.Sum256([]byte(seed))
|
||||
return hex.EncodeToString(hash[:8])
|
||||
}
|
||||
|
||||
// randomIntChoice 随机选择整数
|
||||
func (fm *FingerprintManager) randomIntChoice(choices []int) int {
|
||||
return choices[fm.rng.Intn(len(choices))]
|
||||
// GetAccountKey derives an account key from clientID > refreshToken > random UUID.
|
||||
func GetAccountKey(clientID, refreshToken string) string {
|
||||
// 1. Prefer ClientID
|
||||
if clientID != "" {
|
||||
return GenerateAccountKey(clientID)
|
||||
}
|
||||
|
||||
// 2. Fallback to RefreshToken
|
||||
if refreshToken != "" {
|
||||
return GenerateAccountKey(refreshToken)
|
||||
}
|
||||
|
||||
// 3. Random fallback
|
||||
return GenerateAccountKey(uuid.New().String())
|
||||
}
|
||||
|
||||
// ApplyToRequest 将指纹信息应用到 HTTP 请求头
|
||||
func (fp *Fingerprint) ApplyToRequest(req *http.Request) {
|
||||
req.Header.Set("X-Kiro-SDK-Version", fp.SDKVersion)
|
||||
req.Header.Set("X-Kiro-OS-Type", fp.OSType)
|
||||
req.Header.Set("X-Kiro-OS-Version", fp.OSVersion)
|
||||
req.Header.Set("X-Kiro-Node-Version", fp.NodeVersion)
|
||||
req.Header.Set("X-Kiro-Version", fp.KiroVersion)
|
||||
req.Header.Set("X-Kiro-Hash", fp.KiroHash)
|
||||
req.Header.Set("Accept-Language", fp.AcceptLanguage)
|
||||
req.Header.Set("X-Screen-Resolution", fp.ScreenResolution)
|
||||
req.Header.Set("X-Color-Depth", fmt.Sprintf("%d", fp.ColorDepth))
|
||||
req.Header.Set("X-Hardware-Concurrency", fmt.Sprintf("%d", fp.HardwareConcurrency))
|
||||
req.Header.Set("X-Timezone-Offset", fmt.Sprintf("%d", fp.TimezoneOffset))
|
||||
}
|
||||
|
||||
// RemoveFingerprint 移除 Token 关联的指纹
|
||||
func (fm *FingerprintManager) RemoveFingerprint(tokenKey string) {
|
||||
fm.mu.Lock()
|
||||
defer fm.mu.Unlock()
|
||||
delete(fm.fingerprints, tokenKey)
|
||||
}
|
||||
|
||||
// Count 返回当前管理的指纹数量
|
||||
func (fm *FingerprintManager) Count() int {
|
||||
fm.mu.RLock()
|
||||
defer fm.mu.RUnlock()
|
||||
return len(fm.fingerprints)
|
||||
}
|
||||
|
||||
// BuildUserAgent 构建 User-Agent 字符串 (Kiro IDE 风格)
|
||||
// 格式: aws-sdk-js/{SDKVersion} ua/2.1 os/{OSType}#{OSVersion} lang/js md/nodejs#{NodeVersion} api/codewhispererstreaming#{SDKVersion} m/E KiroIDE-{KiroVersion}-{KiroHash}
|
||||
// BuildUserAgent format: aws-sdk-js/{SDKVersion} ua/2.1 os/{OSType}#{OSVersion} lang/js md/nodejs#{NodeVersion} api/codewhispererstreaming#{SDKVersion} m/E KiroIDE-{KiroVersion}-{KiroHash}
|
||||
func (fp *Fingerprint) BuildUserAgent() string {
|
||||
return fmt.Sprintf(
|
||||
"aws-sdk-js/%s ua/2.1 os/%s#%s lang/js md/nodejs#%s api/codewhispererstreaming#%s m/E KiroIDE-%s-%s",
|
||||
fp.SDKVersion,
|
||||
fp.StreamingSDKVersion,
|
||||
fp.OSType,
|
||||
fp.OSVersion,
|
||||
fp.NodeVersion,
|
||||
fp.SDKVersion,
|
||||
fp.StreamingSDKVersion,
|
||||
fp.KiroVersion,
|
||||
fp.KiroHash,
|
||||
)
|
||||
}
|
||||
|
||||
// BuildAmzUserAgent 构建 X-Amz-User-Agent 字符串
|
||||
// 格式: aws-sdk-js/{SDKVersion} KiroIDE-{KiroVersion}-{KiroHash}
|
||||
// BuildAmzUserAgent format: aws-sdk-js/{SDKVersion} KiroIDE-{KiroVersion}-{KiroHash}
|
||||
func (fp *Fingerprint) BuildAmzUserAgent() string {
|
||||
return fmt.Sprintf(
|
||||
"aws-sdk-js/%s KiroIDE-%s-%s",
|
||||
fp.SDKVersion,
|
||||
fp.StreamingSDKVersion,
|
||||
fp.KiroVersion,
|
||||
fp.KiroHash,
|
||||
)
|
||||
}
|
||||
|
||||
func SetOIDCHeaders(req *http.Request) {
|
||||
fp := GlobalFingerprintManager().GetFingerprint("oidc-session")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("x-amz-user-agent", fmt.Sprintf("aws-sdk-js/%s KiroIDE", fp.OIDCSDKVersion))
|
||||
req.Header.Set("User-Agent", fmt.Sprintf(
|
||||
"aws-sdk-js/%s ua/2.1 os/%s#%s lang/js md/nodejs#%s api/%s#%s m/E KiroIDE",
|
||||
fp.OIDCSDKVersion, fp.OSType, fp.OSVersion, fp.NodeVersion, "sso-oidc", fp.OIDCSDKVersion))
|
||||
req.Header.Set("amz-sdk-invocation-id", uuid.New().String())
|
||||
req.Header.Set("amz-sdk-request", "attempt=1; max=4")
|
||||
}
|
||||
|
||||
func setRuntimeHeaders(req *http.Request, accessToken string, accountKey string) {
|
||||
fp := GlobalFingerprintManager().GetFingerprint(accountKey)
|
||||
machineID := fp.KiroHash
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
req.Header.Set("x-amz-user-agent", fmt.Sprintf("aws-sdk-js/%s KiroIDE-%s-%s",
|
||||
fp.RuntimeSDKVersion, fp.KiroVersion, machineID))
|
||||
req.Header.Set("User-Agent", fmt.Sprintf(
|
||||
"aws-sdk-js/%s ua/2.1 os/%s#%s lang/js md/nodejs#%s api/codewhispererruntime#%s m/N,E KiroIDE-%s-%s",
|
||||
fp.RuntimeSDKVersion, fp.OSType, fp.OSVersion, fp.NodeVersion, fp.RuntimeSDKVersion,
|
||||
fp.KiroVersion, machineID))
|
||||
req.Header.Set("amz-sdk-invocation-id", uuid.New().String())
|
||||
req.Header.Set("amz-sdk-request", "attempt=1; max=1")
|
||||
}
|
||||
|
||||
@@ -2,6 +2,8 @@ package kiro
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
@@ -26,8 +28,14 @@ func TestGetFingerprint_NewToken(t *testing.T) {
|
||||
if fp == nil {
|
||||
t.Fatal("expected non-nil Fingerprint")
|
||||
}
|
||||
if fp.SDKVersion == "" {
|
||||
t.Error("expected non-empty SDKVersion")
|
||||
if fp.OIDCSDKVersion == "" {
|
||||
t.Error("expected non-empty OIDCSDKVersion")
|
||||
}
|
||||
if fp.RuntimeSDKVersion == "" {
|
||||
t.Error("expected non-empty RuntimeSDKVersion")
|
||||
}
|
||||
if fp.StreamingSDKVersion == "" {
|
||||
t.Error("expected non-empty StreamingSDKVersion")
|
||||
}
|
||||
if fp.OSType == "" {
|
||||
t.Error("expected non-empty OSType")
|
||||
@@ -44,18 +52,6 @@ func TestGetFingerprint_NewToken(t *testing.T) {
|
||||
if fp.KiroHash == "" {
|
||||
t.Error("expected non-empty KiroHash")
|
||||
}
|
||||
if fp.AcceptLanguage == "" {
|
||||
t.Error("expected non-empty AcceptLanguage")
|
||||
}
|
||||
if fp.ScreenResolution == "" {
|
||||
t.Error("expected non-empty ScreenResolution")
|
||||
}
|
||||
if fp.ColorDepth == 0 {
|
||||
t.Error("expected non-zero ColorDepth")
|
||||
}
|
||||
if fp.HardwareConcurrency == 0 {
|
||||
t.Error("expected non-zero HardwareConcurrency")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetFingerprint_SameTokenReturnsSameFingerprint(t *testing.T) {
|
||||
@@ -78,72 +74,18 @@ func TestGetFingerprint_DifferentTokens(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemoveFingerprint(t *testing.T) {
|
||||
fm := NewFingerprintManager()
|
||||
fm.GetFingerprint("token1")
|
||||
if fm.Count() != 1 {
|
||||
t.Fatalf("expected count 1, got %d", fm.Count())
|
||||
}
|
||||
|
||||
fm.RemoveFingerprint("token1")
|
||||
if fm.Count() != 0 {
|
||||
t.Errorf("expected count 0, got %d", fm.Count())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemoveFingerprint_NonExistent(t *testing.T) {
|
||||
fm := NewFingerprintManager()
|
||||
fm.RemoveFingerprint("nonexistent")
|
||||
if fm.Count() != 0 {
|
||||
t.Errorf("expected count 0, got %d", fm.Count())
|
||||
}
|
||||
}
|
||||
|
||||
func TestCount(t *testing.T) {
|
||||
fm := NewFingerprintManager()
|
||||
if fm.Count() != 0 {
|
||||
t.Errorf("expected count 0, got %d", fm.Count())
|
||||
}
|
||||
|
||||
fm.GetFingerprint("token1")
|
||||
fm.GetFingerprint("token2")
|
||||
fm.GetFingerprint("token3")
|
||||
|
||||
if fm.Count() != 3 {
|
||||
t.Errorf("expected count 3, got %d", fm.Count())
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyToRequest(t *testing.T) {
|
||||
func TestBuildUserAgent(t *testing.T) {
|
||||
fm := NewFingerprintManager()
|
||||
fp := fm.GetFingerprint("token1")
|
||||
|
||||
req, _ := http.NewRequest("GET", "http://example.com", nil)
|
||||
fp.ApplyToRequest(req)
|
||||
ua := fp.BuildUserAgent()
|
||||
if ua == "" {
|
||||
t.Error("expected non-empty User-Agent")
|
||||
}
|
||||
|
||||
if req.Header.Get("X-Kiro-SDK-Version") != fp.SDKVersion {
|
||||
t.Error("X-Kiro-SDK-Version header mismatch")
|
||||
}
|
||||
if req.Header.Get("X-Kiro-OS-Type") != fp.OSType {
|
||||
t.Error("X-Kiro-OS-Type header mismatch")
|
||||
}
|
||||
if req.Header.Get("X-Kiro-OS-Version") != fp.OSVersion {
|
||||
t.Error("X-Kiro-OS-Version header mismatch")
|
||||
}
|
||||
if req.Header.Get("X-Kiro-Node-Version") != fp.NodeVersion {
|
||||
t.Error("X-Kiro-Node-Version header mismatch")
|
||||
}
|
||||
if req.Header.Get("X-Kiro-Version") != fp.KiroVersion {
|
||||
t.Error("X-Kiro-Version header mismatch")
|
||||
}
|
||||
if req.Header.Get("X-Kiro-Hash") != fp.KiroHash {
|
||||
t.Error("X-Kiro-Hash header mismatch")
|
||||
}
|
||||
if req.Header.Get("Accept-Language") != fp.AcceptLanguage {
|
||||
t.Error("Accept-Language header mismatch")
|
||||
}
|
||||
if req.Header.Get("X-Screen-Resolution") != fp.ScreenResolution {
|
||||
t.Error("X-Screen-Resolution header mismatch")
|
||||
amzUA := fp.BuildAmzUserAgent()
|
||||
if amzUA == "" {
|
||||
t.Error("expected non-empty X-Amz-User-Agent")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -166,6 +108,33 @@ func TestGetFingerprint_OSVersionMatchesOSType(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateFromConfig_OSTypeFromRuntimeGOOS(t *testing.T) {
|
||||
fm := NewFingerprintManager()
|
||||
|
||||
// Set config with empty OSType to trigger runtime.GOOS fallback
|
||||
fm.SetConfig(&FingerprintConfig{
|
||||
OIDCSDKVersion: "3.738.0", // Set other fields to use config path
|
||||
})
|
||||
|
||||
fp := fm.GetFingerprint("test-token")
|
||||
|
||||
// Expected OS type based on runtime.GOOS mapping
|
||||
var expectedOS string
|
||||
switch runtime.GOOS {
|
||||
case "darwin":
|
||||
expectedOS = "darwin"
|
||||
case "windows":
|
||||
expectedOS = "windows"
|
||||
default:
|
||||
expectedOS = "linux"
|
||||
}
|
||||
|
||||
if fp.OSType != expectedOS {
|
||||
t.Errorf("expected OSType '%s' from runtime.GOOS '%s', got '%s'",
|
||||
expectedOS, runtime.GOOS, fp.OSType)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFingerprintManager_ConcurrentAccess(t *testing.T) {
|
||||
fm := NewFingerprintManager()
|
||||
const numGoroutines = 100
|
||||
@@ -174,22 +143,18 @@ func TestFingerprintManager_ConcurrentAccess(t *testing.T) {
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(numGoroutines)
|
||||
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
for i := range numGoroutines {
|
||||
go func(id int) {
|
||||
defer wg.Done()
|
||||
for j := 0; j < numOperations; j++ {
|
||||
for j := range numOperations {
|
||||
tokenKey := "token" + string(rune('a'+id%26))
|
||||
switch j % 4 {
|
||||
switch j % 2 {
|
||||
case 0:
|
||||
fm.GetFingerprint(tokenKey)
|
||||
case 1:
|
||||
fm.Count()
|
||||
case 2:
|
||||
fp := fm.GetFingerprint(tokenKey)
|
||||
req, _ := http.NewRequest("GET", "http://example.com", nil)
|
||||
fp.ApplyToRequest(req)
|
||||
case 3:
|
||||
fm.RemoveFingerprint(tokenKey)
|
||||
_ = fp.BuildUserAgent()
|
||||
_ = fp.BuildAmzUserAgent()
|
||||
}
|
||||
}
|
||||
}(i)
|
||||
@@ -198,16 +163,20 @@ func TestFingerprintManager_ConcurrentAccess(t *testing.T) {
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestKiroHashUniqueness(t *testing.T) {
|
||||
func TestKiroHashStability(t *testing.T) {
|
||||
fm := NewFingerprintManager()
|
||||
hashes := make(map[string]bool)
|
||||
|
||||
for i := 0; i < 100; i++ {
|
||||
fp := fm.GetFingerprint("token" + string(rune(i)))
|
||||
if hashes[fp.KiroHash] {
|
||||
t.Errorf("duplicate KiroHash detected: %s", fp.KiroHash)
|
||||
}
|
||||
hashes[fp.KiroHash] = true
|
||||
// Same token should always return same hash
|
||||
fp1 := fm.GetFingerprint("token1")
|
||||
fp2 := fm.GetFingerprint("token1")
|
||||
if fp1.KiroHash != fp2.KiroHash {
|
||||
t.Errorf("same token should have same hash: %s vs %s", fp1.KiroHash, fp2.KiroHash)
|
||||
}
|
||||
|
||||
// Different tokens should have different hashes
|
||||
fp3 := fm.GetFingerprint("token2")
|
||||
if fp1.KiroHash == fp3.KiroHash {
|
||||
t.Errorf("different tokens should have different hashes")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -220,8 +189,590 @@ func TestKiroHashFormat(t *testing.T) {
|
||||
}
|
||||
|
||||
for _, c := range fp.KiroHash {
|
||||
if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f')) {
|
||||
if (c < '0' || c > '9') && (c < 'a' || c > 'f') {
|
||||
t.Errorf("invalid hex character in KiroHash: %c", c)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGlobalFingerprintManager(t *testing.T) {
|
||||
fm1 := GlobalFingerprintManager()
|
||||
fm2 := GlobalFingerprintManager()
|
||||
|
||||
if fm1 == nil {
|
||||
t.Fatal("expected non-nil GlobalFingerprintManager")
|
||||
}
|
||||
if fm1 != fm2 {
|
||||
t.Error("expected GlobalFingerprintManager to return same instance")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetOIDCHeaders(t *testing.T) {
|
||||
req, _ := http.NewRequest("GET", "http://example.com", nil)
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
if req.Header.Get("Content-Type") != "application/json" {
|
||||
t.Error("expected Content-Type header to be set")
|
||||
}
|
||||
|
||||
amzUA := req.Header.Get("x-amz-user-agent")
|
||||
if amzUA == "" {
|
||||
t.Error("expected x-amz-user-agent header to be set")
|
||||
}
|
||||
if !strings.Contains(amzUA, "aws-sdk-js/") {
|
||||
t.Errorf("x-amz-user-agent should contain aws-sdk-js: %s", amzUA)
|
||||
}
|
||||
if !strings.Contains(amzUA, "KiroIDE") {
|
||||
t.Errorf("x-amz-user-agent should contain KiroIDE: %s", amzUA)
|
||||
}
|
||||
|
||||
ua := req.Header.Get("User-Agent")
|
||||
if ua == "" {
|
||||
t.Error("expected User-Agent header to be set")
|
||||
}
|
||||
if !strings.Contains(ua, "api/sso-oidc") {
|
||||
t.Errorf("User-Agent should contain api name: %s", ua)
|
||||
}
|
||||
|
||||
if req.Header.Get("amz-sdk-invocation-id") == "" {
|
||||
t.Error("expected amz-sdk-invocation-id header to be set")
|
||||
}
|
||||
if req.Header.Get("amz-sdk-request") != "attempt=1; max=4" {
|
||||
t.Errorf("unexpected amz-sdk-request header: %s", req.Header.Get("amz-sdk-request"))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildURL(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
endpoint string
|
||||
path string
|
||||
queryParams map[string]string
|
||||
want string
|
||||
wantContains []string
|
||||
}{
|
||||
{
|
||||
name: "no query params",
|
||||
endpoint: "https://api.example.com",
|
||||
path: "getUsageLimits",
|
||||
queryParams: nil,
|
||||
want: "https://api.example.com/getUsageLimits",
|
||||
},
|
||||
{
|
||||
name: "empty query params",
|
||||
endpoint: "https://api.example.com",
|
||||
path: "getUsageLimits",
|
||||
queryParams: map[string]string{},
|
||||
want: "https://api.example.com/getUsageLimits",
|
||||
},
|
||||
{
|
||||
name: "single query param",
|
||||
endpoint: "https://api.example.com",
|
||||
path: "getUsageLimits",
|
||||
queryParams: map[string]string{
|
||||
"origin": "AI_EDITOR",
|
||||
},
|
||||
want: "https://api.example.com/getUsageLimits?origin=AI_EDITOR",
|
||||
},
|
||||
{
|
||||
name: "multiple query params",
|
||||
endpoint: "https://api.example.com",
|
||||
path: "getUsageLimits",
|
||||
queryParams: map[string]string{
|
||||
"origin": "AI_EDITOR",
|
||||
"resourceType": "AGENTIC_REQUEST",
|
||||
"profileArn": "arn:aws:codewhisperer:us-east-1:123456789012:profile/ABCDEF",
|
||||
},
|
||||
wantContains: []string{
|
||||
"https://api.example.com/getUsageLimits?",
|
||||
"origin=AI_EDITOR",
|
||||
"profileArn=arn%3Aaws%3Acodewhisperer%3Aus-east-1%3A123456789012%3Aprofile%2FABCDEF",
|
||||
"resourceType=AGENTIC_REQUEST",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "omit empty params",
|
||||
endpoint: "https://api.example.com",
|
||||
path: "getUsageLimits",
|
||||
queryParams: map[string]string{
|
||||
"origin": "AI_EDITOR",
|
||||
"profileArn": "",
|
||||
},
|
||||
want: "https://api.example.com/getUsageLimits?origin=AI_EDITOR",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := buildURL(tt.endpoint, tt.path, tt.queryParams)
|
||||
if tt.want != "" {
|
||||
if got != tt.want {
|
||||
t.Errorf("buildURL() = %v, want %v", got, tt.want)
|
||||
}
|
||||
}
|
||||
if tt.wantContains != nil {
|
||||
for _, substr := range tt.wantContains {
|
||||
if !strings.Contains(got, substr) {
|
||||
t.Errorf("buildURL() = %v, want to contain %v", got, substr)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildUserAgentFormat(t *testing.T) {
|
||||
fm := NewFingerprintManager()
|
||||
fp := fm.GetFingerprint("token1")
|
||||
|
||||
ua := fp.BuildUserAgent()
|
||||
requiredParts := []string{
|
||||
"aws-sdk-js/",
|
||||
"ua/2.1",
|
||||
"os/",
|
||||
"lang/js",
|
||||
"md/nodejs#",
|
||||
"api/codewhispererstreaming#",
|
||||
"m/E",
|
||||
"KiroIDE-",
|
||||
}
|
||||
for _, part := range requiredParts {
|
||||
if !strings.Contains(ua, part) {
|
||||
t.Errorf("User-Agent missing required part %q: %s", part, ua)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAmzUserAgentFormat(t *testing.T) {
|
||||
fm := NewFingerprintManager()
|
||||
fp := fm.GetFingerprint("token1")
|
||||
|
||||
amzUA := fp.BuildAmzUserAgent()
|
||||
requiredParts := []string{
|
||||
"aws-sdk-js/",
|
||||
"KiroIDE-",
|
||||
}
|
||||
for _, part := range requiredParts {
|
||||
if !strings.Contains(amzUA, part) {
|
||||
t.Errorf("X-Amz-User-Agent missing required part %q: %s", part, amzUA)
|
||||
}
|
||||
}
|
||||
|
||||
// Amz-User-Agent should be shorter than User-Agent
|
||||
ua := fp.BuildUserAgent()
|
||||
if len(amzUA) >= len(ua) {
|
||||
t.Error("X-Amz-User-Agent should be shorter than User-Agent")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetRuntimeHeaders(t *testing.T) {
|
||||
req, _ := http.NewRequest("GET", "http://example.com", nil)
|
||||
accessToken := "test-access-token-1234567890"
|
||||
clientID := "test-client-id-12345"
|
||||
accountKey := GenerateAccountKey(clientID)
|
||||
fp := GlobalFingerprintManager().GetFingerprint(accountKey)
|
||||
machineID := fp.KiroHash
|
||||
|
||||
setRuntimeHeaders(req, accessToken, accountKey)
|
||||
|
||||
// Check Authorization header
|
||||
if req.Header.Get("Authorization") != "Bearer "+accessToken {
|
||||
t.Errorf("expected Authorization header 'Bearer %s', got '%s'", accessToken, req.Header.Get("Authorization"))
|
||||
}
|
||||
|
||||
// Check x-amz-user-agent header
|
||||
amzUA := req.Header.Get("x-amz-user-agent")
|
||||
if amzUA == "" {
|
||||
t.Error("expected x-amz-user-agent header to be set")
|
||||
}
|
||||
if !strings.Contains(amzUA, "aws-sdk-js/") {
|
||||
t.Errorf("x-amz-user-agent should contain aws-sdk-js: %s", amzUA)
|
||||
}
|
||||
if !strings.Contains(amzUA, "KiroIDE-") {
|
||||
t.Errorf("x-amz-user-agent should contain KiroIDE: %s", amzUA)
|
||||
}
|
||||
if !strings.Contains(amzUA, machineID) {
|
||||
t.Errorf("x-amz-user-agent should contain machineID: %s", amzUA)
|
||||
}
|
||||
|
||||
// Check User-Agent header
|
||||
ua := req.Header.Get("User-Agent")
|
||||
if ua == "" {
|
||||
t.Error("expected User-Agent header to be set")
|
||||
}
|
||||
if !strings.Contains(ua, "api/codewhispererruntime#") {
|
||||
t.Errorf("User-Agent should contain api/codewhispererruntime: %s", ua)
|
||||
}
|
||||
if !strings.Contains(ua, "m/N,E") {
|
||||
t.Errorf("User-Agent should contain m/N,E: %s", ua)
|
||||
}
|
||||
|
||||
// Check amz-sdk-invocation-id (should be a UUID)
|
||||
invocationID := req.Header.Get("amz-sdk-invocation-id")
|
||||
if invocationID == "" {
|
||||
t.Error("expected amz-sdk-invocation-id header to be set")
|
||||
}
|
||||
if len(invocationID) != 36 {
|
||||
t.Errorf("expected amz-sdk-invocation-id to be UUID (36 chars), got %d", len(invocationID))
|
||||
}
|
||||
|
||||
// Check amz-sdk-request
|
||||
if req.Header.Get("amz-sdk-request") != "attempt=1; max=1" {
|
||||
t.Errorf("unexpected amz-sdk-request header: %s", req.Header.Get("amz-sdk-request"))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSDKVersionsAreValid(t *testing.T) {
|
||||
// Verify all OIDC SDK versions match expected format (3.xxx.x)
|
||||
for _, v := range oidcSDKVersions {
|
||||
if !strings.HasPrefix(v, "3.") {
|
||||
t.Errorf("OIDC SDK version should start with 3.: %s", v)
|
||||
}
|
||||
parts := strings.Split(v, ".")
|
||||
if len(parts) != 3 {
|
||||
t.Errorf("OIDC SDK version should have 3 parts: %s", v)
|
||||
}
|
||||
}
|
||||
|
||||
for _, v := range runtimeSDKVersions {
|
||||
parts := strings.Split(v, ".")
|
||||
if len(parts) != 3 {
|
||||
t.Errorf("Runtime SDK version should have 3 parts: %s", v)
|
||||
}
|
||||
}
|
||||
|
||||
for _, v := range streamingSDKVersions {
|
||||
parts := strings.Split(v, ".")
|
||||
if len(parts) != 3 {
|
||||
t.Errorf("Streaming SDK version should have 3 parts: %s", v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestKiroVersionsAreValid(t *testing.T) {
|
||||
// Verify all Kiro versions match expected format (0.x.xxx)
|
||||
for _, v := range kiroVersions {
|
||||
if !strings.HasPrefix(v, "0.") {
|
||||
t.Errorf("Kiro version should start with 0.: %s", v)
|
||||
}
|
||||
parts := strings.Split(v, ".")
|
||||
if len(parts) != 3 {
|
||||
t.Errorf("Kiro version should have 3 parts: %s", v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNodeVersionsAreValid(t *testing.T) {
|
||||
// Verify all Node versions match expected format (xx.xx.x)
|
||||
for _, v := range nodeVersions {
|
||||
parts := strings.Split(v, ".")
|
||||
if len(parts) != 3 {
|
||||
t.Errorf("Node version should have 3 parts: %s", v)
|
||||
}
|
||||
// Should be Node 20.x or 22.x
|
||||
if !strings.HasPrefix(v, "20.") && !strings.HasPrefix(v, "22.") {
|
||||
t.Errorf("Node version should be 20.x or 22.x LTS: %s", v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFingerprintManager_SetConfig(t *testing.T) {
|
||||
fm := NewFingerprintManager()
|
||||
|
||||
// Without config, should generate random fingerprint
|
||||
fp1 := fm.GetFingerprint("token1")
|
||||
if fp1 == nil {
|
||||
t.Fatal("expected non-nil fingerprint")
|
||||
}
|
||||
|
||||
// Set config with all fields
|
||||
cfg := &FingerprintConfig{
|
||||
OIDCSDKVersion: "3.999.0",
|
||||
RuntimeSDKVersion: "9.9.9",
|
||||
StreamingSDKVersion: "8.8.8",
|
||||
OSType: "darwin",
|
||||
OSVersion: "99.0.0",
|
||||
NodeVersion: "99.99.99",
|
||||
KiroVersion: "9.9.999",
|
||||
KiroHash: "customhash123",
|
||||
}
|
||||
fm.SetConfig(cfg)
|
||||
|
||||
// After setting config, should use config values
|
||||
fp2 := fm.GetFingerprint("token2")
|
||||
if fp2.OIDCSDKVersion != "3.999.0" {
|
||||
t.Errorf("expected OIDCSDKVersion '3.999.0', got '%s'", fp2.OIDCSDKVersion)
|
||||
}
|
||||
if fp2.RuntimeSDKVersion != "9.9.9" {
|
||||
t.Errorf("expected RuntimeSDKVersion '9.9.9', got '%s'", fp2.RuntimeSDKVersion)
|
||||
}
|
||||
if fp2.StreamingSDKVersion != "8.8.8" {
|
||||
t.Errorf("expected StreamingSDKVersion '8.8.8', got '%s'", fp2.StreamingSDKVersion)
|
||||
}
|
||||
if fp2.OSType != "darwin" {
|
||||
t.Errorf("expected OSType 'darwin', got '%s'", fp2.OSType)
|
||||
}
|
||||
if fp2.OSVersion != "99.0.0" {
|
||||
t.Errorf("expected OSVersion '99.0.0', got '%s'", fp2.OSVersion)
|
||||
}
|
||||
if fp2.NodeVersion != "99.99.99" {
|
||||
t.Errorf("expected NodeVersion '99.99.99', got '%s'", fp2.NodeVersion)
|
||||
}
|
||||
if fp2.KiroVersion != "9.9.999" {
|
||||
t.Errorf("expected KiroVersion '9.9.999', got '%s'", fp2.KiroVersion)
|
||||
}
|
||||
if fp2.KiroHash != "customhash123" {
|
||||
t.Errorf("expected KiroHash 'customhash123', got '%s'", fp2.KiroHash)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFingerprintManager_SetConfig_PartialFields(t *testing.T) {
|
||||
fm := NewFingerprintManager()
|
||||
|
||||
// Set config with only some fields
|
||||
cfg := &FingerprintConfig{
|
||||
KiroVersion: "1.2.345",
|
||||
KiroHash: "myhash",
|
||||
// Other fields empty - should use random
|
||||
}
|
||||
fm.SetConfig(cfg)
|
||||
|
||||
fp := fm.GetFingerprint("token1")
|
||||
|
||||
// Configured fields should use config values
|
||||
if fp.KiroVersion != "1.2.345" {
|
||||
t.Errorf("expected KiroVersion '1.2.345', got '%s'", fp.KiroVersion)
|
||||
}
|
||||
if fp.KiroHash != "myhash" {
|
||||
t.Errorf("expected KiroHash 'myhash', got '%s'", fp.KiroHash)
|
||||
}
|
||||
|
||||
// Empty fields should be randomly selected (non-empty)
|
||||
if fp.OIDCSDKVersion == "" {
|
||||
t.Error("expected non-empty OIDCSDKVersion")
|
||||
}
|
||||
if fp.OSType == "" {
|
||||
t.Error("expected non-empty OSType")
|
||||
}
|
||||
if fp.NodeVersion == "" {
|
||||
t.Error("expected non-empty NodeVersion")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFingerprintManager_SetConfig_ClearsCache(t *testing.T) {
|
||||
fm := NewFingerprintManager()
|
||||
|
||||
// Get fingerprint before config
|
||||
fp1 := fm.GetFingerprint("token1")
|
||||
originalHash := fp1.KiroHash
|
||||
|
||||
// Set config
|
||||
cfg := &FingerprintConfig{
|
||||
KiroHash: "newcustomhash",
|
||||
}
|
||||
fm.SetConfig(cfg)
|
||||
|
||||
// Same token should now return different fingerprint (cache cleared)
|
||||
fp2 := fm.GetFingerprint("token1")
|
||||
if fp2.KiroHash == originalHash {
|
||||
t.Error("expected cache to be cleared after SetConfig")
|
||||
}
|
||||
if fp2.KiroHash != "newcustomhash" {
|
||||
t.Errorf("expected KiroHash 'newcustomhash', got '%s'", fp2.KiroHash)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateAccountKey(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
seed string
|
||||
check func(t *testing.T, result string)
|
||||
}{
|
||||
{
|
||||
name: "Empty seed",
|
||||
seed: "",
|
||||
check: func(t *testing.T, result string) {
|
||||
if result == "" {
|
||||
t.Error("expected non-empty result for empty seed")
|
||||
}
|
||||
if len(result) != 16 {
|
||||
t.Errorf("expected 16 char hex string, got %d chars", len(result))
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Simple seed",
|
||||
seed: "test-client-id",
|
||||
check: func(t *testing.T, result string) {
|
||||
if len(result) != 16 {
|
||||
t.Errorf("expected 16 char hex string, got %d chars", len(result))
|
||||
}
|
||||
// Verify it's valid hex
|
||||
for _, c := range result {
|
||||
if (c < '0' || c > '9') && (c < 'a' || c > 'f') {
|
||||
t.Errorf("invalid hex character: %c", c)
|
||||
}
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Same seed produces same result",
|
||||
seed: "deterministic-seed",
|
||||
check: func(t *testing.T, result string) {
|
||||
result2 := GenerateAccountKey("deterministic-seed")
|
||||
if result != result2 {
|
||||
t.Errorf("same seed should produce same result: %s vs %s", result, result2)
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Different seeds produce different results",
|
||||
seed: "seed-one",
|
||||
check: func(t *testing.T, result string) {
|
||||
result2 := GenerateAccountKey("seed-two")
|
||||
if result == result2 {
|
||||
t.Errorf("different seeds should produce different results: %s vs %s", result, result2)
|
||||
}
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := GenerateAccountKey(tt.seed)
|
||||
tt.check(t, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetAccountKey(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
clientID string
|
||||
refreshToken string
|
||||
check func(t *testing.T, result string)
|
||||
}{
|
||||
{
|
||||
name: "Priority 1: clientID when both provided",
|
||||
clientID: "client-id-123",
|
||||
refreshToken: "refresh-token-456",
|
||||
check: func(t *testing.T, result string) {
|
||||
expected := GenerateAccountKey("client-id-123")
|
||||
if result != expected {
|
||||
t.Errorf("expected clientID-based key %s, got %s", expected, result)
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Priority 2: refreshToken when clientID is empty",
|
||||
clientID: "",
|
||||
refreshToken: "refresh-token-789",
|
||||
check: func(t *testing.T, result string) {
|
||||
expected := GenerateAccountKey("refresh-token-789")
|
||||
if result != expected {
|
||||
t.Errorf("expected refreshToken-based key %s, got %s", expected, result)
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Priority 3: random when both empty",
|
||||
clientID: "",
|
||||
refreshToken: "",
|
||||
check: func(t *testing.T, result string) {
|
||||
if len(result) != 16 {
|
||||
t.Errorf("expected 16 char key, got %d chars", len(result))
|
||||
}
|
||||
// Should be different each time (random UUID)
|
||||
result2 := GetAccountKey("", "")
|
||||
if result == result2 {
|
||||
t.Log("warning: random keys are the same (possible but unlikely)")
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "clientID only",
|
||||
clientID: "solo-client-id",
|
||||
refreshToken: "",
|
||||
check: func(t *testing.T, result string) {
|
||||
expected := GenerateAccountKey("solo-client-id")
|
||||
if result != expected {
|
||||
t.Errorf("expected clientID-based key %s, got %s", expected, result)
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "refreshToken only",
|
||||
clientID: "",
|
||||
refreshToken: "solo-refresh-token",
|
||||
check: func(t *testing.T, result string) {
|
||||
expected := GenerateAccountKey("solo-refresh-token")
|
||||
if result != expected {
|
||||
t.Errorf("expected refreshToken-based key %s, got %s", expected, result)
|
||||
}
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := GetAccountKey(tt.clientID, tt.refreshToken)
|
||||
tt.check(t, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetAccountKey_Deterministic(t *testing.T) {
|
||||
// Verify that GetAccountKey produces deterministic results for same inputs
|
||||
clientID := "test-client-id-abc"
|
||||
refreshToken := "test-refresh-token-xyz"
|
||||
|
||||
// Call multiple times with same inputs
|
||||
results := make([]string, 10)
|
||||
for i := range 10 {
|
||||
results[i] = GetAccountKey(clientID, refreshToken)
|
||||
}
|
||||
|
||||
// All results should be identical
|
||||
for i := 1; i < 10; i++ {
|
||||
if results[i] != results[0] {
|
||||
t.Errorf("GetAccountKey should be deterministic: got %s and %s", results[0], results[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFingerprintDeterministic(t *testing.T) {
|
||||
// Verify that fingerprints are deterministic based on accountKey
|
||||
fm := NewFingerprintManager()
|
||||
|
||||
accountKey := GenerateAccountKey("test-client-id")
|
||||
|
||||
// Get fingerprint multiple times
|
||||
fp1 := fm.GetFingerprint(accountKey)
|
||||
fp2 := fm.GetFingerprint(accountKey)
|
||||
|
||||
// Should be the same pointer (cached)
|
||||
if fp1 != fp2 {
|
||||
t.Error("expected same fingerprint pointer for same key")
|
||||
}
|
||||
|
||||
// Create new manager and verify same values
|
||||
fm2 := NewFingerprintManager()
|
||||
fp3 := fm2.GetFingerprint(accountKey)
|
||||
|
||||
// Values should be identical (deterministic generation)
|
||||
if fp1.KiroHash != fp3.KiroHash {
|
||||
t.Errorf("KiroHash should be deterministic: %s vs %s", fp1.KiroHash, fp3.KiroHash)
|
||||
}
|
||||
if fp1.OSType != fp3.OSType {
|
||||
t.Errorf("OSType should be deterministic: %s vs %s", fp1.OSType, fp3.OSType)
|
||||
}
|
||||
if fp1.OSVersion != fp3.OSVersion {
|
||||
t.Errorf("OSVersion should be deterministic: %s vs %s", fp1.OSVersion, fp3.OSVersion)
|
||||
}
|
||||
if fp1.KiroVersion != fp3.KiroVersion {
|
||||
t.Errorf("KiroVersion should be deterministic: %s vs %s", fp1.KiroVersion, fp3.KiroVersion)
|
||||
}
|
||||
if fp1.NodeVersion != fp3.NodeVersion {
|
||||
t.Errorf("NodeVersion should be deterministic: %s vs %s", fp1.NodeVersion, fp3.NodeVersion)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -23,10 +23,10 @@ import (
|
||||
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
|
||||
)
|
||||
@@ -41,8 +41,10 @@ type KiroTokenResponse struct {
|
||||
|
||||
// KiroOAuth handles the OAuth flow for Kiro authentication.
|
||||
type KiroOAuth struct {
|
||||
httpClient *http.Client
|
||||
cfg *config.Config
|
||||
httpClient *http.Client
|
||||
cfg *config.Config
|
||||
machineID string
|
||||
kiroVersion string
|
||||
}
|
||||
|
||||
// NewKiroOAuth creates a new Kiro OAuth handler.
|
||||
@@ -51,9 +53,12 @@ func NewKiroOAuth(cfg *config.Config) *KiroOAuth {
|
||||
if cfg != nil {
|
||||
client = util.SetProxy(&cfg.SDKConfig, client)
|
||||
}
|
||||
fp := GlobalFingerprintManager().GetFingerprint("login")
|
||||
return &KiroOAuth{
|
||||
httpClient: client,
|
||||
cfg: cfg,
|
||||
httpClient: client,
|
||||
cfg: cfg,
|
||||
machineID: fp.KiroHash,
|
||||
kiroVersion: fp.KiroVersion,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -190,7 +195,8 @@ func (o *KiroOAuth) exchangeCodeForToken(ctx context.Context, code, codeVerifier
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", "KiroIDE-0.7.45-cli-proxy-api")
|
||||
req.Header.Set("User-Agent", fmt.Sprintf("KiroIDE-%s-%s", o.kiroVersion, o.machineID))
|
||||
req.Header.Set("Accept", "application/json, text/plain, */*")
|
||||
|
||||
resp, err := o.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -256,11 +262,8 @@ func (o *KiroOAuth) RefreshTokenWithFingerprint(ctx context.Context, refreshToke
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
// Use KiroIDE-style User-Agent to match official Kiro IDE behavior
|
||||
// This helps avoid 403 errors from server-side User-Agent validation
|
||||
userAgent := buildKiroUserAgent(tokenKey)
|
||||
req.Header.Set("User-Agent", userAgent)
|
||||
req.Header.Set("User-Agent", fmt.Sprintf("KiroIDE-%s-%s", o.kiroVersion, o.machineID))
|
||||
req.Header.Set("Accept", "application/json, text/plain, */*")
|
||||
|
||||
resp, err := o.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -301,19 +304,6 @@ func (o *KiroOAuth) RefreshTokenWithFingerprint(ctx context.Context, refreshToke
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildKiroUserAgent builds a KiroIDE-style User-Agent string.
|
||||
// If tokenKey is provided, uses fingerprint manager for consistent fingerprint.
|
||||
// Otherwise generates a simple KiroIDE User-Agent.
|
||||
func buildKiroUserAgent(tokenKey string) string {
|
||||
if tokenKey != "" {
|
||||
fm := NewFingerprintManager()
|
||||
fp := fm.GetFingerprint(tokenKey)
|
||||
return fmt.Sprintf("KiroIDE-%s-%s", fp.KiroVersion, fp.KiroHash[:16])
|
||||
}
|
||||
// Default KiroIDE User-Agent matching kiro-openai-gateway format
|
||||
return "KiroIDE-0.7.45-cli-proxy-api"
|
||||
}
|
||||
|
||||
// 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) {
|
||||
|
||||
@@ -35,35 +35,35 @@ const (
|
||||
)
|
||||
|
||||
type webAuthSession struct {
|
||||
stateID string
|
||||
deviceCode string
|
||||
userCode string
|
||||
authURL string
|
||||
verificationURI string
|
||||
expiresIn int
|
||||
interval int
|
||||
status authSessionStatus
|
||||
startedAt time.Time
|
||||
completedAt time.Time
|
||||
expiresAt time.Time
|
||||
error string
|
||||
tokenData *KiroTokenData
|
||||
ssoClient *SSOOIDCClient
|
||||
clientID string
|
||||
clientSecret string
|
||||
region string
|
||||
cancelFunc context.CancelFunc
|
||||
authMethod string // "google", "github", "builder-id", "idc"
|
||||
startURL string // Used for IDC
|
||||
codeVerifier string // Used for social auth PKCE
|
||||
codeChallenge string // Used for social auth PKCE
|
||||
stateID string
|
||||
deviceCode string
|
||||
userCode string
|
||||
authURL string
|
||||
verificationURI string
|
||||
expiresIn int
|
||||
interval int
|
||||
status authSessionStatus
|
||||
startedAt time.Time
|
||||
completedAt time.Time
|
||||
expiresAt time.Time
|
||||
error string
|
||||
tokenData *KiroTokenData
|
||||
ssoClient *SSOOIDCClient
|
||||
clientID string
|
||||
clientSecret string
|
||||
region string
|
||||
cancelFunc context.CancelFunc
|
||||
authMethod string // "google", "github", "builder-id", "idc"
|
||||
startURL string // Used for IDC
|
||||
codeVerifier string // Used for social auth PKCE
|
||||
codeChallenge string // Used for social auth PKCE
|
||||
}
|
||||
|
||||
type OAuthWebHandler struct {
|
||||
cfg *config.Config
|
||||
sessions map[string]*webAuthSession
|
||||
mu sync.RWMutex
|
||||
onTokenObtained func(*KiroTokenData)
|
||||
cfg *config.Config
|
||||
sessions map[string]*webAuthSession
|
||||
mu sync.RWMutex
|
||||
onTokenObtained func(*KiroTokenData)
|
||||
}
|
||||
|
||||
func NewOAuthWebHandler(cfg *config.Config) *OAuthWebHandler {
|
||||
@@ -104,7 +104,7 @@ func (h *OAuthWebHandler) handleSelect(c *gin.Context) {
|
||||
|
||||
func (h *OAuthWebHandler) handleStart(c *gin.Context) {
|
||||
method := c.Query("method")
|
||||
|
||||
|
||||
if method == "" {
|
||||
c.Redirect(http.StatusFound, "/v0/oauth/kiro")
|
||||
return
|
||||
@@ -138,7 +138,7 @@ func (h *OAuthWebHandler) startSocialAuth(c *gin.Context, method string) {
|
||||
}
|
||||
|
||||
socialClient := NewSocialAuthClient(h.cfg)
|
||||
|
||||
|
||||
var provider string
|
||||
if method == "google" {
|
||||
provider = string(ProviderGoogle)
|
||||
@@ -373,22 +373,28 @@ func (h *OAuthWebHandler) pollForToken(ctx context.Context, session *webAuthSess
|
||||
}
|
||||
|
||||
expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
|
||||
profileArn := session.ssoClient.fetchProfileArn(ctx, tokenResp.AccessToken)
|
||||
email := FetchUserEmailWithFallback(ctx, h.cfg, tokenResp.AccessToken)
|
||||
|
||||
// Fetch profileArn for IDC
|
||||
var profileArn string
|
||||
if session.authMethod == "idc" {
|
||||
profileArn = session.ssoClient.FetchProfileArn(ctx, tokenResp.AccessToken, session.clientID, tokenResp.RefreshToken)
|
||||
}
|
||||
|
||||
email := FetchUserEmailWithFallback(ctx, h.cfg, tokenResp.AccessToken, session.clientID, tokenResp.RefreshToken)
|
||||
|
||||
tokenData := &KiroTokenData{
|
||||
AccessToken: tokenResp.AccessToken,
|
||||
RefreshToken: tokenResp.RefreshToken,
|
||||
ProfileArn: profileArn,
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: session.authMethod,
|
||||
Provider: "AWS",
|
||||
ClientID: session.clientID,
|
||||
ClientSecret: session.clientSecret,
|
||||
Email: email,
|
||||
Region: session.region,
|
||||
StartURL: session.startURL,
|
||||
}
|
||||
AccessToken: tokenResp.AccessToken,
|
||||
RefreshToken: tokenResp.RefreshToken,
|
||||
ProfileArn: profileArn,
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: session.authMethod,
|
||||
Provider: "AWS",
|
||||
ClientID: session.clientID,
|
||||
ClientSecret: session.clientSecret,
|
||||
Email: email,
|
||||
Region: session.region,
|
||||
StartURL: session.startURL,
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
session.status = statusSuccess
|
||||
@@ -442,7 +448,7 @@ func (h *OAuthWebHandler) saveTokenToFile(tokenData *KiroTokenData) {
|
||||
fileName := GenerateTokenFileName(tokenData)
|
||||
|
||||
authFilePath := filepath.Join(authDir, fileName)
|
||||
|
||||
|
||||
// Convert to storage format and save
|
||||
storage := &KiroTokenStorage{
|
||||
Type: "kiro",
|
||||
@@ -459,12 +465,12 @@ func (h *OAuthWebHandler) saveTokenToFile(tokenData *KiroTokenData) {
|
||||
StartURL: tokenData.StartURL,
|
||||
Email: tokenData.Email,
|
||||
}
|
||||
|
||||
|
||||
if err := storage.SaveTokenToFile(authFilePath); err != nil {
|
||||
log.Errorf("OAuth Web: failed to save token to file: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
log.Infof("OAuth Web: token saved to %s", authFilePath)
|
||||
}
|
||||
|
||||
|
||||
@@ -10,14 +10,14 @@ import (
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// RefreshManager 是后台刷新器的单例管理器
|
||||
// RefreshManager is a singleton manager for background token refreshing.
|
||||
type RefreshManager struct {
|
||||
mu sync.Mutex
|
||||
refresher *BackgroundRefresher
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
started bool
|
||||
onTokenRefreshed func(tokenID string, tokenData *KiroTokenData) // 刷新成功回调
|
||||
onTokenRefreshed func(tokenID string, tokenData *KiroTokenData)
|
||||
}
|
||||
|
||||
var (
|
||||
@@ -25,7 +25,7 @@ var (
|
||||
managerOnce sync.Once
|
||||
)
|
||||
|
||||
// GetRefreshManager 获取全局刷新管理器实例
|
||||
// GetRefreshManager returns the global RefreshManager singleton.
|
||||
func GetRefreshManager() *RefreshManager {
|
||||
managerOnce.Do(func() {
|
||||
globalRefreshManager = &RefreshManager{}
|
||||
@@ -33,9 +33,7 @@ func GetRefreshManager() *RefreshManager {
|
||||
return globalRefreshManager
|
||||
}
|
||||
|
||||
// Initialize 初始化后台刷新器
|
||||
// baseDir: token 文件所在的目录
|
||||
// cfg: 应用配置
|
||||
// Initialize sets up the background refresher.
|
||||
func (m *RefreshManager) Initialize(baseDir string, cfg *config.Config) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
@@ -58,18 +56,16 @@ func (m *RefreshManager) Initialize(baseDir string, cfg *config.Config) error {
|
||||
baseDir = resolvedBaseDir
|
||||
}
|
||||
|
||||
// 创建 token 存储库
|
||||
repo := NewFileTokenRepository(baseDir)
|
||||
|
||||
// 创建后台刷新器,配置参数
|
||||
opts := []RefresherOption{
|
||||
WithInterval(time.Minute), // 每分钟检查一次
|
||||
WithBatchSize(50), // 每批最多处理 50 个 token
|
||||
WithConcurrency(10), // 最多 10 个并发刷新
|
||||
WithConfig(cfg), // 设置 OAuth 和 SSO 客户端
|
||||
WithInterval(time.Minute),
|
||||
WithBatchSize(50),
|
||||
WithConcurrency(10),
|
||||
WithConfig(cfg),
|
||||
}
|
||||
|
||||
// 如果已设置回调,传递给 BackgroundRefresher
|
||||
// Pass callback to BackgroundRefresher if already set
|
||||
if m.onTokenRefreshed != nil {
|
||||
opts = append(opts, WithOnTokenRefreshed(m.onTokenRefreshed))
|
||||
}
|
||||
@@ -80,7 +76,7 @@ func (m *RefreshManager) Initialize(baseDir string, cfg *config.Config) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Start 启动后台刷新
|
||||
// Start begins background token refreshing.
|
||||
func (m *RefreshManager) Start() {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
@@ -102,7 +98,7 @@ func (m *RefreshManager) Start() {
|
||||
log.Info("refresh manager: background refresh started")
|
||||
}
|
||||
|
||||
// Stop 停止后台刷新
|
||||
// Stop halts background token refreshing.
|
||||
func (m *RefreshManager) Stop() {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
@@ -123,14 +119,14 @@ func (m *RefreshManager) Stop() {
|
||||
log.Info("refresh manager: background refresh stopped")
|
||||
}
|
||||
|
||||
// IsRunning 检查后台刷新是否正在运行
|
||||
// IsRunning reports whether background refreshing is active.
|
||||
func (m *RefreshManager) IsRunning() bool {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return m.started
|
||||
}
|
||||
|
||||
// UpdateBaseDir 更新 token 目录(用于运行时配置更改)
|
||||
// UpdateBaseDir changes the token directory at runtime.
|
||||
func (m *RefreshManager) UpdateBaseDir(baseDir string) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
@@ -143,16 +139,15 @@ func (m *RefreshManager) UpdateBaseDir(baseDir string) {
|
||||
}
|
||||
}
|
||||
|
||||
// SetOnTokenRefreshed 设置 token 刷新成功后的回调函数
|
||||
// 可以在任何时候调用,支持运行时更新回调
|
||||
// callback: 回调函数,接收 tokenID(文件名)和新的 token 数据
|
||||
// SetOnTokenRefreshed registers a callback invoked after a successful token refresh.
|
||||
// Can be called at any time; supports runtime callback updates.
|
||||
func (m *RefreshManager) SetOnTokenRefreshed(callback func(tokenID string, tokenData *KiroTokenData)) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
m.onTokenRefreshed = callback
|
||||
|
||||
// 如果 refresher 已经创建,使用并发安全的方式更新它的回调
|
||||
// Update the refresher's callback in a thread-safe manner if already created
|
||||
if m.refresher != nil {
|
||||
m.refresher.callbackMu.Lock()
|
||||
m.refresher.onTokenRefreshed = callback
|
||||
@@ -162,8 +157,11 @@ func (m *RefreshManager) SetOnTokenRefreshed(callback func(tokenID string, token
|
||||
log.Debug("refresh manager: token refresh callback registered")
|
||||
}
|
||||
|
||||
// InitializeAndStart 初始化并启动后台刷新(便捷方法)
|
||||
// InitializeAndStart initializes and starts background refreshing (convenience method).
|
||||
func InitializeAndStart(baseDir string, cfg *config.Config) {
|
||||
// Initialize global fingerprint config
|
||||
initGlobalFingerprintConfig(cfg)
|
||||
|
||||
manager := GetRefreshManager()
|
||||
if err := manager.Initialize(baseDir, cfg); err != nil {
|
||||
log.Errorf("refresh manager: initialization failed: %v", err)
|
||||
@@ -172,7 +170,31 @@ func InitializeAndStart(baseDir string, cfg *config.Config) {
|
||||
manager.Start()
|
||||
}
|
||||
|
||||
// StopGlobalRefreshManager 停止全局刷新管理器
|
||||
// initGlobalFingerprintConfig loads fingerprint settings from application config.
|
||||
func initGlobalFingerprintConfig(cfg *config.Config) {
|
||||
if cfg == nil || cfg.KiroFingerprint == nil {
|
||||
return
|
||||
}
|
||||
fpCfg := cfg.KiroFingerprint
|
||||
SetGlobalFingerprintConfig(&FingerprintConfig{
|
||||
OIDCSDKVersion: fpCfg.OIDCSDKVersion,
|
||||
RuntimeSDKVersion: fpCfg.RuntimeSDKVersion,
|
||||
StreamingSDKVersion: fpCfg.StreamingSDKVersion,
|
||||
OSType: fpCfg.OSType,
|
||||
OSVersion: fpCfg.OSVersion,
|
||||
NodeVersion: fpCfg.NodeVersion,
|
||||
KiroVersion: fpCfg.KiroVersion,
|
||||
KiroHash: fpCfg.KiroHash,
|
||||
})
|
||||
log.Debug("kiro: global fingerprint config loaded")
|
||||
}
|
||||
|
||||
// InitFingerprintConfig initializes the global fingerprint config from application config.
|
||||
func InitFingerprintConfig(cfg *config.Config) {
|
||||
initGlobalFingerprintConfig(cfg)
|
||||
}
|
||||
|
||||
// StopGlobalRefreshManager stops the global refresh manager.
|
||||
func StopGlobalRefreshManager() {
|
||||
if globalRefreshManager != nil {
|
||||
globalRefreshManager.Stop()
|
||||
|
||||
@@ -84,6 +84,8 @@ type SocialAuthClient struct {
|
||||
httpClient *http.Client
|
||||
cfg *config.Config
|
||||
protocolHandler *ProtocolHandler
|
||||
machineID string
|
||||
kiroVersion string
|
||||
}
|
||||
|
||||
// NewSocialAuthClient creates a new social auth client.
|
||||
@@ -92,10 +94,13 @@ func NewSocialAuthClient(cfg *config.Config) *SocialAuthClient {
|
||||
if cfg != nil {
|
||||
client = util.SetProxy(&cfg.SDKConfig, client)
|
||||
}
|
||||
fp := GlobalFingerprintManager().GetFingerprint("login")
|
||||
return &SocialAuthClient{
|
||||
httpClient: client,
|
||||
cfg: cfg,
|
||||
protocolHandler: NewProtocolHandler(),
|
||||
machineID: fp.KiroHash,
|
||||
kiroVersion: fp.KiroVersion,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -229,7 +234,8 @@ func (c *SocialAuthClient) CreateToken(ctx context.Context, req *CreateTokenRequ
|
||||
}
|
||||
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("User-Agent", "KiroIDE-0.7.45-cli-proxy-api")
|
||||
httpReq.Header.Set("User-Agent", fmt.Sprintf("KiroIDE-%s-%s", c.kiroVersion, c.machineID))
|
||||
httpReq.Header.Set("Accept", "application/json, text/plain, */*")
|
||||
|
||||
resp, err := c.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
@@ -269,7 +275,8 @@ func (c *SocialAuthClient) RefreshSocialToken(ctx context.Context, refreshToken
|
||||
}
|
||||
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
|
||||
httpReq.Header.Set("User-Agent", fmt.Sprintf("KiroIDE-%s-%s", c.kiroVersion, c.machineID))
|
||||
httpReq.Header.Set("Accept", "application/json, text/plain, */*")
|
||||
|
||||
resp, err := c.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
@@ -466,7 +473,7 @@ 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 {
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -40,21 +41,13 @@ const (
|
||||
// Authorization code flow callback
|
||||
authCodeCallbackPath = "/oauth/callback"
|
||||
authCodeCallbackPort = 19877
|
||||
|
||||
// User-Agent to match official Kiro IDE
|
||||
kiroUserAgent = "KiroIDE"
|
||||
|
||||
// IDC token refresh headers (matching Kiro IDE behavior)
|
||||
idcAmzUserAgent = "aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE"
|
||||
)
|
||||
|
||||
// Sentinel errors for OIDC token polling
|
||||
var (
|
||||
ErrAuthorizationPending = errors.New("authorization_pending")
|
||||
ErrSlowDown = errors.New("slow_down")
|
||||
)
|
||||
|
||||
// SSOOIDCClient handles AWS SSO OIDC authentication.
|
||||
type SSOOIDCClient struct {
|
||||
httpClient *http.Client
|
||||
cfg *config.Config
|
||||
@@ -74,10 +67,10 @@ func NewSSOOIDCClient(cfg *config.Config) *SSOOIDCClient {
|
||||
|
||||
// RegisterClientResponse from AWS SSO OIDC.
|
||||
type RegisterClientResponse struct {
|
||||
ClientID string `json:"clientId"`
|
||||
ClientSecret string `json:"clientSecret"`
|
||||
ClientIDIssuedAt int64 `json:"clientIdIssuedAt"`
|
||||
ClientSecretExpiresAt int64 `json:"clientSecretExpiresAt"`
|
||||
ClientID string `json:"clientId"`
|
||||
ClientSecret string `json:"clientSecret"`
|
||||
ClientIDIssuedAt int64 `json:"clientIdIssuedAt"`
|
||||
ClientSecretExpiresAt int64 `json:"clientSecretExpiresAt"`
|
||||
}
|
||||
|
||||
// StartDeviceAuthResponse from AWS SSO OIDC.
|
||||
@@ -174,8 +167,7 @@ func (c *SSOOIDCClient) RegisterClientWithRegion(ctx context.Context, region str
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", kiroUserAgent)
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -220,8 +212,7 @@ func (c *SSOOIDCClient) StartDeviceAuthorizationWithIDC(ctx context.Context, cli
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", kiroUserAgent)
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -267,8 +258,7 @@ func (c *SSOOIDCClient) CreateTokenWithRegion(ctx context.Context, clientID, cli
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", kiroUserAgent)
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -311,8 +301,11 @@ func (c *SSOOIDCClient) CreateTokenWithRegion(ctx context.Context, clientID, cli
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// RefreshTokenWithRegion refreshes an access token using the refresh token with a specific region.
|
||||
// RefreshTokenWithRegion refreshes an access token using the refresh token with a specific OIDC region.
|
||||
func (c *SSOOIDCClient) RefreshTokenWithRegion(ctx context.Context, clientID, clientSecret, refreshToken, region, startURL string) (*KiroTokenData, error) {
|
||||
if region == "" {
|
||||
region = defaultIDCRegion
|
||||
}
|
||||
endpoint := getOIDCEndpoint(region)
|
||||
|
||||
payload := map[string]string{
|
||||
@@ -331,18 +324,7 @@ func (c *SSOOIDCClient) RefreshTokenWithRegion(ctx context.Context, clientID, cl
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Set headers matching kiro2api's IDC token refresh
|
||||
// These headers are required for successful IDC token refresh
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Host", fmt.Sprintf("oidc.%s.amazonaws.com", region))
|
||||
req.Header.Set("Connection", "keep-alive")
|
||||
req.Header.Set("x-amz-user-agent", idcAmzUserAgent)
|
||||
req.Header.Set("Accept", "*/*")
|
||||
req.Header.Set("Accept-Language", "*")
|
||||
req.Header.Set("sec-fetch-mode", "cors")
|
||||
req.Header.Set("User-Agent", "node")
|
||||
req.Header.Set("Accept-Encoding", "br, gzip, deflate")
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -469,10 +451,10 @@ func (c *SSOOIDCClient) LoginWithIDC(ctx context.Context, startURL, region strin
|
||||
|
||||
// Step 5: Get profile ARN from CodeWhisperer API
|
||||
fmt.Println("Fetching profile information...")
|
||||
profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken)
|
||||
profileArn := c.FetchProfileArn(ctx, tokenResp.AccessToken, regResp.ClientID, tokenResp.RefreshToken)
|
||||
|
||||
// Fetch user email
|
||||
email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken)
|
||||
email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken, regResp.ClientID, tokenResp.RefreshToken)
|
||||
if email != "" {
|
||||
fmt.Printf(" Logged in as: %s\n", email)
|
||||
}
|
||||
@@ -502,12 +484,36 @@ func (c *SSOOIDCClient) LoginWithIDC(ctx context.Context, startURL, region strin
|
||||
return nil, fmt.Errorf("authorization timed out")
|
||||
}
|
||||
|
||||
// IDCLoginOptions holds optional parameters for IDC login.
|
||||
type IDCLoginOptions struct {
|
||||
StartURL string // Pre-configured start URL (skips prompt if set)
|
||||
Region string // OIDC region for login and token refresh (defaults to us-east-1)
|
||||
UseDeviceCode bool // Use Device Code flow instead of Auth Code flow
|
||||
}
|
||||
|
||||
// LoginWithMethodSelection prompts the user to select between Builder ID and IDC, then performs the login.
|
||||
func (c *SSOOIDCClient) LoginWithMethodSelection(ctx context.Context) (*KiroTokenData, error) {
|
||||
// Options can be provided to pre-configure IDC parameters (startURL, region).
|
||||
// If StartURL is provided in opts, IDC flow is used directly without prompting.
|
||||
func (c *SSOOIDCClient) LoginWithMethodSelection(ctx context.Context, opts *IDCLoginOptions) (*KiroTokenData, error) {
|
||||
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
|
||||
fmt.Println("║ Kiro Authentication (AWS) ║")
|
||||
fmt.Println("╚══════════════════════════════════════════════════════════╝")
|
||||
|
||||
// If IDC options with StartURL are provided, skip method selection and use IDC directly
|
||||
if opts != nil && opts.StartURL != "" {
|
||||
region := opts.Region
|
||||
if region == "" {
|
||||
region = defaultIDCRegion
|
||||
}
|
||||
fmt.Printf("\n Using IDC with Start URL: %s\n", opts.StartURL)
|
||||
fmt.Printf(" Region: %s\n", region)
|
||||
|
||||
if opts.UseDeviceCode {
|
||||
return c.LoginWithIDCAndOptions(ctx, opts.StartURL, region)
|
||||
}
|
||||
return c.LoginWithIDCAuthCode(ctx, opts.StartURL, region)
|
||||
}
|
||||
|
||||
// Prompt for login method
|
||||
options := []string{
|
||||
"Use with Builder ID (personal AWS account)",
|
||||
@@ -520,15 +526,41 @@ func (c *SSOOIDCClient) LoginWithMethodSelection(ctx context.Context) (*KiroToke
|
||||
return c.LoginWithBuilderID(ctx)
|
||||
}
|
||||
|
||||
// IDC flow - prompt for start URL and region
|
||||
fmt.Println()
|
||||
startURL := promptInput("? Enter Start URL", "")
|
||||
if startURL == "" {
|
||||
return nil, fmt.Errorf("start URL is required for IDC login")
|
||||
// IDC flow - use pre-configured values or prompt
|
||||
var startURL, region string
|
||||
|
||||
if opts != nil {
|
||||
startURL = opts.StartURL
|
||||
region = opts.Region
|
||||
}
|
||||
|
||||
region := promptInput("? Enter Region", defaultIDCRegion)
|
||||
fmt.Println()
|
||||
|
||||
// Use pre-configured startURL or prompt
|
||||
if startURL == "" {
|
||||
startURL = promptInput("? Enter Start URL", "")
|
||||
if startURL == "" {
|
||||
return nil, fmt.Errorf("start URL is required for IDC login")
|
||||
}
|
||||
} else {
|
||||
fmt.Printf(" Using pre-configured Start URL: %s\n", startURL)
|
||||
}
|
||||
|
||||
// Use pre-configured region or prompt
|
||||
if region == "" {
|
||||
region = promptInput("? Enter Region", defaultIDCRegion)
|
||||
} else {
|
||||
fmt.Printf(" Using pre-configured Region: %s\n", region)
|
||||
}
|
||||
|
||||
if opts != nil && opts.UseDeviceCode {
|
||||
return c.LoginWithIDCAndOptions(ctx, startURL, region)
|
||||
}
|
||||
return c.LoginWithIDCAuthCode(ctx, startURL, region)
|
||||
}
|
||||
|
||||
// LoginWithIDCAndOptions performs IDC login with the specified region.
|
||||
func (c *SSOOIDCClient) LoginWithIDCAndOptions(ctx context.Context, startURL, region string) (*KiroTokenData, error) {
|
||||
return c.LoginWithIDC(ctx, startURL, region)
|
||||
}
|
||||
|
||||
@@ -550,8 +582,7 @@ func (c *SSOOIDCClient) RegisterClient(ctx context.Context) (*RegisterClientResp
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", kiroUserAgent)
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -594,8 +625,7 @@ func (c *SSOOIDCClient) StartDeviceAuthorization(ctx context.Context, clientID,
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", kiroUserAgent)
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -639,8 +669,7 @@ func (c *SSOOIDCClient) CreateToken(ctx context.Context, clientID, clientSecret,
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", kiroUserAgent)
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -702,13 +731,7 @@ func (c *SSOOIDCClient) RefreshToken(ctx context.Context, clientID, clientSecret
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Set headers matching Kiro IDE behavior for better compatibility
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Host", "oidc.us-east-1.amazonaws.com")
|
||||
req.Header.Set("x-amz-user-agent", idcAmzUserAgent)
|
||||
req.Header.Set("User-Agent", "node")
|
||||
req.Header.Set("Accept", "*/*")
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -835,12 +858,8 @@ func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData,
|
||||
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)
|
||||
|
||||
// Fetch user email (tries CodeWhisperer API first, then userinfo endpoint, then JWT parsing)
|
||||
email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken)
|
||||
email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken, regResp.ClientID, tokenResp.RefreshToken)
|
||||
if email != "" {
|
||||
fmt.Printf(" Logged in as: %s\n", email)
|
||||
}
|
||||
@@ -850,7 +869,7 @@ func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData,
|
||||
return &KiroTokenData{
|
||||
AccessToken: tokenResp.AccessToken,
|
||||
RefreshToken: tokenResp.RefreshToken,
|
||||
ProfileArn: profileArn,
|
||||
ProfileArn: "", // Builder ID has no profile
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: "builder-id",
|
||||
Provider: "AWS",
|
||||
@@ -859,15 +878,15 @@ func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData,
|
||||
Email: email,
|
||||
Region: defaultIDCRegion,
|
||||
}, 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")
|
||||
}
|
||||
// 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")
|
||||
}
|
||||
|
||||
// FetchUserEmail retrieves the user's email from AWS SSO OIDC userinfo endpoint.
|
||||
// Falls back to JWT parsing if userinfo fails.
|
||||
@@ -931,20 +950,64 @@ func (c *SSOOIDCClient) tryUserInfoEndpoint(ctx context.Context, accessToken str
|
||||
return ""
|
||||
}
|
||||
|
||||
// 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)
|
||||
// FetchProfileArn fetches the profile ARN from ListAvailableProfiles API.
|
||||
// This is used to get profileArn for imported accounts that may not have it.
|
||||
func (c *SSOOIDCClient) FetchProfileArn(ctx context.Context, accessToken, clientID, refreshToken string) string {
|
||||
profileArn := c.tryListAvailableProfiles(ctx, accessToken, clientID, refreshToken)
|
||||
if profileArn != "" {
|
||||
return profileArn
|
||||
}
|
||||
|
||||
// Fallback: Try ListAvailableCustomizations
|
||||
return c.tryListCustomizations(ctx, accessToken)
|
||||
return c.tryListProfilesLegacy(ctx, accessToken)
|
||||
}
|
||||
|
||||
func (c *SSOOIDCClient) tryListProfiles(ctx context.Context, accessToken string) string {
|
||||
func (c *SSOOIDCClient) tryListAvailableProfiles(ctx context.Context, accessToken, clientID, refreshToken string) string {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, GetKiroAPIEndpoint("")+"/ListAvailableProfiles", strings.NewReader("{}"))
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
accountKey := GetAccountKey(clientID, refreshToken)
|
||||
setRuntimeHeaders(req, accessToken, accountKey)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
log.Debugf("ListAvailableProfiles request failed: %v", err)
|
||||
return ""
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("ListAvailableProfiles failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return ""
|
||||
}
|
||||
|
||||
log.Debugf("ListAvailableProfiles response: %s", string(respBody))
|
||||
|
||||
var result struct {
|
||||
Profiles []struct {
|
||||
Arn string `json:"arn"`
|
||||
ProfileName string `json:"profileName"`
|
||||
} `json:"profiles"`
|
||||
NextToken *string `json:"nextToken"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
log.Debugf("ListAvailableProfiles parse error: %v", err)
|
||||
return ""
|
||||
}
|
||||
|
||||
if len(result.Profiles) > 0 {
|
||||
log.Debugf("Found profile: %s (%s)", result.Profiles[0].ProfileName, result.Profiles[0].Arn)
|
||||
return result.Profiles[0].Arn
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
func (c *SSOOIDCClient) tryListProfilesLegacy(ctx context.Context, accessToken string) string {
|
||||
payload := map[string]interface{}{
|
||||
"origin": "AI_EDITOR",
|
||||
}
|
||||
@@ -954,7 +1017,9 @@ func (c *SSOOIDCClient) tryListProfiles(ctx context.Context, accessToken string)
|
||||
return ""
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body)))
|
||||
// Use the legacy CodeWhisperer endpoint for JSON-RPC style requests.
|
||||
// The Q endpoint (q.{region}.amazonaws.com) does NOT support x-amz-target headers.
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, GetCodeWhispererLegacyEndpoint(""), strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
@@ -973,11 +1038,11 @@ func (c *SSOOIDCClient) tryListProfiles(ctx context.Context, accessToken string)
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("ListProfiles failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
log.Debugf("ListProfiles (legacy) failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return ""
|
||||
}
|
||||
|
||||
log.Debugf("ListProfiles response: %s", string(respBody))
|
||||
log.Debugf("ListProfiles (legacy) response: %s", string(respBody))
|
||||
|
||||
var result struct {
|
||||
Profiles []struct {
|
||||
@@ -1001,63 +1066,6 @@ func (c *SSOOIDCClient) tryListProfiles(ctx context.Context, accessToken string)
|
||||
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 ""
|
||||
}
|
||||
|
||||
// RegisterClientForAuthCode registers a new OIDC client for authorization code flow.
|
||||
func (c *SSOOIDCClient) RegisterClientForAuthCode(ctx context.Context, redirectURI string) (*RegisterClientResponse, error) {
|
||||
payload := map[string]interface{}{
|
||||
@@ -1078,8 +1086,7 @@ func (c *SSOOIDCClient) RegisterClientForAuthCode(ctx context.Context, redirectU
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", kiroUserAgent)
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -1105,6 +1112,53 @@ func (c *SSOOIDCClient) RegisterClientForAuthCode(ctx context.Context, redirectU
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
func (c *SSOOIDCClient) RegisterClientForAuthCodeWithIDC(ctx context.Context, redirectURI, issuerUrl, region string) (*RegisterClientResponse, error) {
|
||||
endpoint := getOIDCEndpoint(region)
|
||||
|
||||
payload := map[string]interface{}{
|
||||
"clientName": "Kiro IDE",
|
||||
"clientType": "public",
|
||||
"scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"},
|
||||
"grantTypes": []string{"authorization_code", "refresh_token"},
|
||||
"redirectUris": []string{redirectURI},
|
||||
"issuerUrl": issuerUrl,
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/client/register", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
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 for auth code with IDC 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
|
||||
}
|
||||
|
||||
// AuthCodeCallbackResult contains the result from authorization code callback.
|
||||
type AuthCodeCallbackResult struct {
|
||||
Code string
|
||||
@@ -1128,6 +1182,7 @@ func (c *SSOOIDCClient) startAuthCodeCallbackServer(ctx context.Context, expecte
|
||||
port := listener.Addr().(*net.TCPAddr).Port
|
||||
redirectURI := fmt.Sprintf("http://127.0.0.1:%d%s", port, authCodeCallbackPath)
|
||||
resultChan := make(chan AuthCodeCallbackResult, 1)
|
||||
doneChan := make(chan struct{})
|
||||
|
||||
server := &http.Server{
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
@@ -1147,6 +1202,7 @@ func (c *SSOOIDCClient) startAuthCodeCallbackServer(ctx context.Context, expecte
|
||||
<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))
|
||||
resultChan <- AuthCodeCallbackResult{Error: errParam}
|
||||
close(doneChan)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1156,6 +1212,7 @@ func (c *SSOOIDCClient) startAuthCodeCallbackServer(ctx context.Context, expecte
|
||||
<html><head><title>Login Failed</title></head>
|
||||
<body><h1>Login Failed</h1><p>Invalid state parameter</p><p>You can close this window.</p></body></html>`)
|
||||
resultChan <- AuthCodeCallbackResult{Error: "state mismatch"}
|
||||
close(doneChan)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1164,6 +1221,7 @@ func (c *SSOOIDCClient) startAuthCodeCallbackServer(ctx context.Context, expecte
|
||||
<body><h1>Login Successful!</h1><p>You can close this window and return to the terminal.</p>
|
||||
<script>window.close();</script></body></html>`)
|
||||
resultChan <- AuthCodeCallbackResult{Code: code, State: state}
|
||||
close(doneChan)
|
||||
})
|
||||
|
||||
server.Handler = mux
|
||||
@@ -1178,7 +1236,7 @@ func (c *SSOOIDCClient) startAuthCodeCallbackServer(ctx context.Context, expecte
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-time.After(10 * time.Minute):
|
||||
case <-resultChan:
|
||||
case <-doneChan:
|
||||
}
|
||||
_ = server.Shutdown(context.Background())
|
||||
}()
|
||||
@@ -1227,8 +1285,54 @@ func (c *SSOOIDCClient) CreateTokenWithAuthCode(ctx context.Context, clientID, c
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", kiroUserAgent)
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
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("create token with auth code 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
|
||||
}
|
||||
|
||||
func (c *SSOOIDCClient) CreateTokenWithAuthCodeAndRegion(ctx context.Context, clientID, clientSecret, code, codeVerifier, redirectURI, region string) (*CreateTokenResponse, error) {
|
||||
endpoint := getOIDCEndpoint(region)
|
||||
|
||||
payload := map[string]string{
|
||||
"clientId": clientID,
|
||||
"clientSecret": clientSecret,
|
||||
"code": code,
|
||||
"codeVerifier": codeVerifier,
|
||||
"redirectUri": redirectURI,
|
||||
"grantType": "authorization_code",
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/token", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
SetOIDCHeaders(req)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -1352,12 +1456,118 @@ func (c *SSOOIDCClient) LoginWithBuilderIDAuthCode(ctx context.Context) (*KiroTo
|
||||
|
||||
fmt.Println("\n✓ Authentication successful!")
|
||||
|
||||
// Step 8: Get profile ARN
|
||||
fmt.Println("Fetching profile information...")
|
||||
profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken)
|
||||
|
||||
// Fetch user email (tries CodeWhisperer API first, then userinfo endpoint, then JWT parsing)
|
||||
email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken)
|
||||
email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken, regResp.ClientID, tokenResp.RefreshToken)
|
||||
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: "", // Builder ID has no profile
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: "builder-id",
|
||||
Provider: "AWS",
|
||||
ClientID: regResp.ClientID,
|
||||
ClientSecret: regResp.ClientSecret,
|
||||
Email: email,
|
||||
Region: defaultIDCRegion,
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (c *SSOOIDCClient) LoginWithIDCAuthCode(ctx context.Context, startURL, region string) (*KiroTokenData, error) {
|
||||
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
|
||||
fmt.Println("║ Kiro Authentication (AWS IDC - Auth Code) ║")
|
||||
fmt.Println("╚══════════════════════════════════════════════════════════╝")
|
||||
|
||||
if region == "" {
|
||||
region = defaultIDCRegion
|
||||
}
|
||||
|
||||
codeVerifier, codeChallenge, err := generatePKCEForAuthCode()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to generate PKCE: %w", err)
|
||||
}
|
||||
|
||||
state, err := generateStateForAuthCode()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to generate state: %w", err)
|
||||
}
|
||||
|
||||
fmt.Println("\nStarting callback server...")
|
||||
redirectURI, resultChan, err := c.startAuthCodeCallbackServer(ctx, state)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to start callback server: %w", err)
|
||||
}
|
||||
log.Debugf("Callback server started, redirect URI: %s", redirectURI)
|
||||
|
||||
fmt.Println("Registering client...")
|
||||
regResp, err := c.RegisterClientForAuthCodeWithIDC(ctx, redirectURI, startURL, region)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to register client: %w", err)
|
||||
}
|
||||
log.Debugf("Client registered: %s", regResp.ClientID)
|
||||
|
||||
endpoint := getOIDCEndpoint(region)
|
||||
scopes := "codewhisperer:completions,codewhisperer:analysis,codewhisperer:conversations,codewhisperer:transformations,codewhisperer:taskassist"
|
||||
authURL := buildAuthorizationURL(endpoint, regResp.ClientID, redirectURI, scopes, state, codeChallenge)
|
||||
|
||||
fmt.Println("\n════════════════════════════════════════════════════════════")
|
||||
fmt.Println(" Opening browser for authentication...")
|
||||
fmt.Println("════════════════════════════════════════════════════════════")
|
||||
fmt.Printf("\n URL: %s\n\n", authURL)
|
||||
|
||||
if c.cfg != nil {
|
||||
browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
|
||||
} else {
|
||||
browser.SetIncognitoMode(true)
|
||||
}
|
||||
|
||||
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 authorization callback...")
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
browser.CloseBrowser()
|
||||
return nil, ctx.Err()
|
||||
case <-time.After(10 * time.Minute):
|
||||
browser.CloseBrowser()
|
||||
return nil, fmt.Errorf("authorization timed out")
|
||||
case result := <-resultChan:
|
||||
if result.Error != "" {
|
||||
browser.CloseBrowser()
|
||||
return nil, fmt.Errorf("authorization failed: %s", result.Error)
|
||||
}
|
||||
|
||||
fmt.Println("\n✓ Authorization received!")
|
||||
|
||||
if err := browser.CloseBrowser(); err != nil {
|
||||
log.Debugf("Failed to close browser: %v", err)
|
||||
}
|
||||
|
||||
fmt.Println("Exchanging code for tokens...")
|
||||
tokenResp, err := c.CreateTokenWithAuthCodeAndRegion(ctx, regResp.ClientID, regResp.ClientSecret, result.Code, codeVerifier, redirectURI, region)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to exchange code for tokens: %w", err)
|
||||
}
|
||||
|
||||
fmt.Println("\n✓ Authentication successful!")
|
||||
|
||||
fmt.Println("Fetching profile information...")
|
||||
profileArn := c.FetchProfileArn(ctx, tokenResp.AccessToken, regResp.ClientID, tokenResp.RefreshToken)
|
||||
|
||||
email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken, regResp.ClientID, tokenResp.RefreshToken)
|
||||
if email != "" {
|
||||
fmt.Printf(" Logged in as: %s\n", email)
|
||||
}
|
||||
@@ -1369,12 +1579,25 @@ func (c *SSOOIDCClient) LoginWithBuilderIDAuthCode(ctx context.Context) (*KiroTo
|
||||
RefreshToken: tokenResp.RefreshToken,
|
||||
ProfileArn: profileArn,
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: "builder-id",
|
||||
AuthMethod: "idc",
|
||||
Provider: "AWS",
|
||||
ClientID: regResp.ClientID,
|
||||
ClientSecret: regResp.ClientSecret,
|
||||
Email: email,
|
||||
Region: defaultIDCRegion,
|
||||
StartURL: startURL,
|
||||
Region: region,
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
|
||||
func buildAuthorizationURL(endpoint, clientID, redirectURI, scopes, state, codeChallenge string) string {
|
||||
params := url.Values{}
|
||||
params.Set("response_type", "code")
|
||||
params.Set("client_id", clientID)
|
||||
params.Set("redirect_uri", redirectURI)
|
||||
params.Set("scopes", scopes)
|
||||
params.Set("state", state)
|
||||
params.Set("code_challenge", codeChallenge)
|
||||
params.Set("code_challenge_method", "S256")
|
||||
return fmt.Sprintf("%s/authorize?%s", endpoint, params.Encode())
|
||||
}
|
||||
|
||||
261
internal/auth/kiro/sso_oidc_test.go
Normal file
261
internal/auth/kiro/sso_oidc_test.go
Normal file
@@ -0,0 +1,261 @@
|
||||
package kiro
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type recordingRoundTripper struct {
|
||||
lastReq *http.Request
|
||||
}
|
||||
|
||||
func (rt *recordingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
rt.lastReq = req
|
||||
body := `{"nextToken":null,"profiles":[{"arn":"arn:aws:codewhisperer:us-east-1:123456789012:profile/ABC","profileName":"test"}]}`
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func TestTryListAvailableProfiles_UsesClientIDForAccountKey(t *testing.T) {
|
||||
rt := &recordingRoundTripper{}
|
||||
client := &SSOOIDCClient{
|
||||
httpClient: &http.Client{Transport: rt},
|
||||
}
|
||||
|
||||
profileArn := client.tryListAvailableProfiles(context.Background(), "access-token", "client-id-123", "refresh-token-456")
|
||||
if profileArn == "" {
|
||||
t.Fatal("expected profileArn, got empty result")
|
||||
}
|
||||
|
||||
accountKey := GetAccountKey("client-id-123", "refresh-token-456")
|
||||
fp := GlobalFingerprintManager().GetFingerprint(accountKey)
|
||||
expected := fmt.Sprintf("aws-sdk-js/%s KiroIDE-%s-%s", fp.RuntimeSDKVersion, fp.KiroVersion, fp.KiroHash)
|
||||
got := rt.lastReq.Header.Get("X-Amz-User-Agent")
|
||||
if got != expected {
|
||||
t.Errorf("X-Amz-User-Agent = %q, want %q", got, expected)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTryListAvailableProfiles_UsesRefreshTokenWhenClientIDMissing(t *testing.T) {
|
||||
rt := &recordingRoundTripper{}
|
||||
client := &SSOOIDCClient{
|
||||
httpClient: &http.Client{Transport: rt},
|
||||
}
|
||||
|
||||
profileArn := client.tryListAvailableProfiles(context.Background(), "access-token", "", "refresh-token-789")
|
||||
if profileArn == "" {
|
||||
t.Fatal("expected profileArn, got empty result")
|
||||
}
|
||||
|
||||
accountKey := GetAccountKey("", "refresh-token-789")
|
||||
fp := GlobalFingerprintManager().GetFingerprint(accountKey)
|
||||
expected := fmt.Sprintf("aws-sdk-js/%s KiroIDE-%s-%s", fp.RuntimeSDKVersion, fp.KiroVersion, fp.KiroHash)
|
||||
got := rt.lastReq.Header.Get("X-Amz-User-Agent")
|
||||
if got != expected {
|
||||
t.Errorf("X-Amz-User-Agent = %q, want %q", got, expected)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterClientForAuthCodeWithIDC(t *testing.T) {
|
||||
var capturedReq struct {
|
||||
Method string
|
||||
Path string
|
||||
Headers http.Header
|
||||
Body map[string]interface{}
|
||||
}
|
||||
|
||||
mockResp := RegisterClientResponse{
|
||||
ClientID: "test-client-id",
|
||||
ClientSecret: "test-client-secret",
|
||||
ClientIDIssuedAt: 1700000000,
|
||||
ClientSecretExpiresAt: 1700086400,
|
||||
}
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
capturedReq.Method = r.Method
|
||||
capturedReq.Path = r.URL.Path
|
||||
capturedReq.Headers = r.Header.Clone()
|
||||
|
||||
bodyBytes, _ := io.ReadAll(r.Body)
|
||||
json.Unmarshal(bodyBytes, &capturedReq.Body)
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(mockResp)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
// Extract host to build a region that resolves to our test server.
|
||||
// Override getOIDCEndpoint by passing region="" and patching the endpoint.
|
||||
// Since getOIDCEndpoint builds "https://oidc.{region}.amazonaws.com", we
|
||||
// instead inject the test server URL directly via a custom HTTP client transport.
|
||||
client := &SSOOIDCClient{
|
||||
httpClient: ts.Client(),
|
||||
}
|
||||
|
||||
// We need to route the request to our test server. Use a transport that rewrites the URL.
|
||||
client.httpClient.Transport = &rewriteTransport{
|
||||
base: ts.Client().Transport,
|
||||
targetURL: ts.URL,
|
||||
}
|
||||
|
||||
resp, err := client.RegisterClientForAuthCodeWithIDC(
|
||||
context.Background(),
|
||||
"http://127.0.0.1:19877/oauth/callback",
|
||||
"https://my-idc-instance.awsapps.com/start",
|
||||
"us-east-1",
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// Verify request method and path
|
||||
if capturedReq.Method != http.MethodPost {
|
||||
t.Errorf("method = %q, want POST", capturedReq.Method)
|
||||
}
|
||||
if capturedReq.Path != "/client/register" {
|
||||
t.Errorf("path = %q, want /client/register", capturedReq.Path)
|
||||
}
|
||||
|
||||
// Verify headers
|
||||
if ct := capturedReq.Headers.Get("Content-Type"); ct != "application/json" {
|
||||
t.Errorf("Content-Type = %q, want application/json", ct)
|
||||
}
|
||||
ua := capturedReq.Headers.Get("User-Agent")
|
||||
if !strings.Contains(ua, "KiroIDE") {
|
||||
t.Errorf("User-Agent %q does not contain KiroIDE", ua)
|
||||
}
|
||||
if !strings.Contains(ua, "sso-oidc") {
|
||||
t.Errorf("User-Agent %q does not contain sso-oidc", ua)
|
||||
}
|
||||
xua := capturedReq.Headers.Get("X-Amz-User-Agent")
|
||||
if !strings.Contains(xua, "KiroIDE") {
|
||||
t.Errorf("x-amz-user-agent %q does not contain KiroIDE", xua)
|
||||
}
|
||||
|
||||
// Verify body fields
|
||||
if v, _ := capturedReq.Body["clientName"].(string); v != "Kiro IDE" {
|
||||
t.Errorf("clientName = %q, want %q", v, "Kiro IDE")
|
||||
}
|
||||
if v, _ := capturedReq.Body["clientType"].(string); v != "public" {
|
||||
t.Errorf("clientType = %q, want %q", v, "public")
|
||||
}
|
||||
if v, _ := capturedReq.Body["issuerUrl"].(string); v != "https://my-idc-instance.awsapps.com/start" {
|
||||
t.Errorf("issuerUrl = %q, want %q", v, "https://my-idc-instance.awsapps.com/start")
|
||||
}
|
||||
|
||||
// Verify scopes array
|
||||
scopesRaw, ok := capturedReq.Body["scopes"].([]interface{})
|
||||
if !ok || len(scopesRaw) != 5 {
|
||||
t.Fatalf("scopes: got %v, want 5-element array", capturedReq.Body["scopes"])
|
||||
}
|
||||
expectedScopes := []string{
|
||||
"codewhisperer:completions", "codewhisperer:analysis",
|
||||
"codewhisperer:conversations", "codewhisperer:transformations",
|
||||
"codewhisperer:taskassist",
|
||||
}
|
||||
for i, s := range expectedScopes {
|
||||
if scopesRaw[i].(string) != s {
|
||||
t.Errorf("scopes[%d] = %q, want %q", i, scopesRaw[i], s)
|
||||
}
|
||||
}
|
||||
|
||||
// Verify grantTypes
|
||||
grantTypesRaw, ok := capturedReq.Body["grantTypes"].([]interface{})
|
||||
if !ok || len(grantTypesRaw) != 2 {
|
||||
t.Fatalf("grantTypes: got %v, want 2-element array", capturedReq.Body["grantTypes"])
|
||||
}
|
||||
if grantTypesRaw[0].(string) != "authorization_code" || grantTypesRaw[1].(string) != "refresh_token" {
|
||||
t.Errorf("grantTypes = %v, want [authorization_code, refresh_token]", grantTypesRaw)
|
||||
}
|
||||
|
||||
// Verify redirectUris
|
||||
redirectRaw, ok := capturedReq.Body["redirectUris"].([]interface{})
|
||||
if !ok || len(redirectRaw) != 1 {
|
||||
t.Fatalf("redirectUris: got %v, want 1-element array", capturedReq.Body["redirectUris"])
|
||||
}
|
||||
if redirectRaw[0].(string) != "http://127.0.0.1:19877/oauth/callback" {
|
||||
t.Errorf("redirectUris[0] = %q, want %q", redirectRaw[0], "http://127.0.0.1:19877/oauth/callback")
|
||||
}
|
||||
|
||||
// Verify response parsing
|
||||
if resp.ClientID != "test-client-id" {
|
||||
t.Errorf("ClientID = %q, want %q", resp.ClientID, "test-client-id")
|
||||
}
|
||||
if resp.ClientSecret != "test-client-secret" {
|
||||
t.Errorf("ClientSecret = %q, want %q", resp.ClientSecret, "test-client-secret")
|
||||
}
|
||||
if resp.ClientIDIssuedAt != 1700000000 {
|
||||
t.Errorf("ClientIDIssuedAt = %d, want %d", resp.ClientIDIssuedAt, 1700000000)
|
||||
}
|
||||
if resp.ClientSecretExpiresAt != 1700086400 {
|
||||
t.Errorf("ClientSecretExpiresAt = %d, want %d", resp.ClientSecretExpiresAt, 1700086400)
|
||||
}
|
||||
}
|
||||
|
||||
// rewriteTransport redirects all requests to the test server URL.
|
||||
type rewriteTransport struct {
|
||||
base http.RoundTripper
|
||||
targetURL string
|
||||
}
|
||||
|
||||
func (t *rewriteTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
target, _ := url.Parse(t.targetURL)
|
||||
req.URL.Scheme = target.Scheme
|
||||
req.URL.Host = target.Host
|
||||
if t.base != nil {
|
||||
return t.base.RoundTrip(req)
|
||||
}
|
||||
return http.DefaultTransport.RoundTrip(req)
|
||||
}
|
||||
|
||||
func TestBuildAuthorizationURL(t *testing.T) {
|
||||
scopes := "codewhisperer:completions,codewhisperer:analysis,codewhisperer:conversations,codewhisperer:transformations,codewhisperer:taskassist"
|
||||
endpoint := "https://oidc.us-east-1.amazonaws.com"
|
||||
redirectURI := "http://127.0.0.1:19877/oauth/callback"
|
||||
|
||||
authURL := buildAuthorizationURL(endpoint, "test-client-id", redirectURI, scopes, "random-state", "test-challenge")
|
||||
|
||||
// Verify colons and commas in scopes are percent-encoded
|
||||
if !strings.Contains(authURL, "codewhisperer%3Acompletions") {
|
||||
t.Errorf("expected colons in scopes to be percent-encoded, got: %s", authURL)
|
||||
}
|
||||
if !strings.Contains(authURL, "completions%2Ccodewhisperer") {
|
||||
t.Errorf("expected commas in scopes to be percent-encoded, got: %s", authURL)
|
||||
}
|
||||
|
||||
// Parse back and verify all parameters round-trip correctly
|
||||
parsed, err := url.Parse(authURL)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to parse auth URL: %v", err)
|
||||
}
|
||||
|
||||
if !strings.HasPrefix(authURL, endpoint+"/authorize?") {
|
||||
t.Errorf("expected URL to start with %s/authorize?, got: %s", endpoint, authURL)
|
||||
}
|
||||
|
||||
q := parsed.Query()
|
||||
checks := map[string]string{
|
||||
"response_type": "code",
|
||||
"client_id": "test-client-id",
|
||||
"redirect_uri": redirectURI,
|
||||
"scopes": scopes,
|
||||
"state": "random-state",
|
||||
"code_challenge": "test-challenge",
|
||||
"code_challenge_method": "S256",
|
||||
}
|
||||
for key, want := range checks {
|
||||
if got := q.Get(key); got != want {
|
||||
t.Errorf("%s = %q, want %q", key, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -29,7 +29,7 @@ type KiroTokenStorage struct {
|
||||
ClientID string `json:"client_id,omitempty"`
|
||||
// ClientSecret is the OAuth client secret (required for token refresh)
|
||||
ClientSecret string `json:"client_secret,omitempty"`
|
||||
// Region is the AWS region
|
||||
// Region is the OIDC region for IDC login and token refresh
|
||||
Region string `json:"region,omitempty"`
|
||||
// StartURL is the AWS Identity Center start URL (for IDC auth)
|
||||
StartURL string `json:"start_url,omitempty"`
|
||||
|
||||
@@ -200,36 +200,22 @@ func (r *FileTokenRepository) readTokenFile(path string) (*Token, error) {
|
||||
}
|
||||
|
||||
// 解析各字段
|
||||
if v, ok := metadata["access_token"].(string); ok {
|
||||
token.AccessToken = v
|
||||
}
|
||||
if v, ok := metadata["refresh_token"].(string); ok {
|
||||
token.RefreshToken = v
|
||||
}
|
||||
if v, ok := metadata["client_id"].(string); ok {
|
||||
token.ClientID = v
|
||||
}
|
||||
if v, ok := metadata["client_secret"].(string); ok {
|
||||
token.ClientSecret = v
|
||||
}
|
||||
if v, ok := metadata["region"].(string); ok {
|
||||
token.Region = v
|
||||
}
|
||||
if v, ok := metadata["start_url"].(string); ok {
|
||||
token.StartURL = v
|
||||
}
|
||||
if v, ok := metadata["provider"].(string); ok {
|
||||
token.Provider = v
|
||||
}
|
||||
token.AccessToken, _ = metadata["access_token"].(string)
|
||||
token.RefreshToken, _ = metadata["refresh_token"].(string)
|
||||
token.ClientID, _ = metadata["client_id"].(string)
|
||||
token.ClientSecret, _ = metadata["client_secret"].(string)
|
||||
token.Region, _ = metadata["region"].(string)
|
||||
token.StartURL, _ = metadata["start_url"].(string)
|
||||
token.Provider, _ = metadata["provider"].(string)
|
||||
|
||||
// 解析时间字段
|
||||
if v, ok := metadata["expires_at"].(string); ok {
|
||||
if t, err := time.Parse(time.RFC3339, v); err == nil {
|
||||
if expiresAtStr, ok := metadata["expires_at"].(string); ok && expiresAtStr != "" {
|
||||
if t, err := time.Parse(time.RFC3339, expiresAtStr); err == nil {
|
||||
token.ExpiresAt = t
|
||||
}
|
||||
}
|
||||
if v, ok := metadata["last_refresh"].(string); ok {
|
||||
if t, err := time.Parse(time.RFC3339, v); err == nil {
|
||||
if lastRefreshStr, ok := metadata["last_refresh"].(string); ok && lastRefreshStr != "" {
|
||||
if t, err := time.Parse(time.RFC3339, lastRefreshStr); err == nil {
|
||||
token.LastVerified = t
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,7 +8,6 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
@@ -51,14 +50,12 @@ type QuotaStatus struct {
|
||||
// UsageChecker provides methods for checking token quota usage.
|
||||
type UsageChecker struct {
|
||||
httpClient *http.Client
|
||||
endpoint string
|
||||
}
|
||||
|
||||
// NewUsageChecker creates a new UsageChecker instance.
|
||||
func NewUsageChecker(cfg *config.Config) *UsageChecker {
|
||||
return &UsageChecker{
|
||||
httpClient: util.SetProxy(&cfg.SDKConfig, &http.Client{Timeout: 30 * time.Second}),
|
||||
endpoint: awsKiroEndpoint,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -66,7 +63,6 @@ func NewUsageChecker(cfg *config.Config) *UsageChecker {
|
||||
func NewUsageCheckerWithClient(client *http.Client) *UsageChecker {
|
||||
return &UsageChecker{
|
||||
httpClient: client,
|
||||
endpoint: awsKiroEndpoint,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -80,26 +76,23 @@ func (c *UsageChecker) CheckUsage(ctx context.Context, tokenData *KiroTokenData)
|
||||
return nil, fmt.Errorf("access token is empty")
|
||||
}
|
||||
|
||||
payload := map[string]interface{}{
|
||||
queryParams := map[string]string{
|
||||
"origin": "AI_EDITOR",
|
||||
"profileArn": tokenData.ProfileArn,
|
||||
"resourceType": "AGENTIC_REQUEST",
|
||||
}
|
||||
|
||||
jsonBody, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
// Use endpoint from profileArn if available
|
||||
endpoint := GetKiroAPIEndpointFromProfileArn(tokenData.ProfileArn)
|
||||
url := buildURL(endpoint, pathGetUsageLimits, queryParams)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint, strings.NewReader(string(jsonBody)))
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
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", targetGetUsage)
|
||||
req.Header.Set("Authorization", "Bearer "+tokenData.AccessToken)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
accountKey := GetAccountKey(tokenData.ClientID, tokenData.RefreshToken)
|
||||
setRuntimeHeaders(req, tokenData.AccessToken, accountKey)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
|
||||
@@ -30,11 +30,21 @@ type QwenTokenStorage struct {
|
||||
Type string `json:"type"`
|
||||
// Expire is the timestamp when the current access token expires.
|
||||
Expire string `json:"expired"`
|
||||
|
||||
// Metadata holds arbitrary key-value pairs injected via hooks.
|
||||
// It is not exported to JSON directly to allow flattening during serialization.
|
||||
Metadata map[string]any `json:"-"`
|
||||
}
|
||||
|
||||
// SetMetadata allows external callers to inject metadata into the storage before saving.
|
||||
func (ts *QwenTokenStorage) SetMetadata(meta map[string]any) {
|
||||
ts.Metadata = meta
|
||||
}
|
||||
|
||||
// SaveTokenToFile serializes the Qwen token storage to a JSON file.
|
||||
// This method creates the necessary directory structure and writes the token
|
||||
// data in JSON format to the specified file path for persistent storage.
|
||||
// It merges any injected metadata into the top-level JSON object.
|
||||
//
|
||||
// Parameters:
|
||||
// - authFilePath: The full path where the token file should be saved
|
||||
@@ -56,7 +66,13 @@ func (ts *QwenTokenStorage) SaveTokenToFile(authFilePath string) error {
|
||||
_ = f.Close()
|
||||
}()
|
||||
|
||||
if err = json.NewEncoder(f).Encode(ts); err != nil {
|
||||
// Merge metadata using helper
|
||||
data, errMerge := misc.MergeMetadata(ts, ts.Metadata)
|
||||
if errMerge != nil {
|
||||
return fmt.Errorf("failed to merge metadata: %w", errMerge)
|
||||
}
|
||||
|
||||
if err = json.NewEncoder(f).Encode(data); err != nil {
|
||||
return fmt.Errorf("failed to write token to file: %w", err)
|
||||
}
|
||||
return nil
|
||||
|
||||
@@ -5,8 +5,7 @@ import (
|
||||
)
|
||||
|
||||
// newAuthManager creates a new authentication manager instance with all supported
|
||||
// authenticators and a file-based token store. It initializes authenticators for
|
||||
// Gemini, Codex, Claude, Qwen, IFlow, Antigravity, and GitHub Copilot providers.
|
||||
// authenticators and a file-based token store.
|
||||
//
|
||||
// Returns:
|
||||
// - *sdkAuth.Manager: A configured authentication manager instance
|
||||
@@ -23,6 +22,9 @@ func newAuthManager() *sdkAuth.Manager {
|
||||
sdkAuth.NewKiroAuthenticator(),
|
||||
sdkAuth.NewGitHubCopilotAuthenticator(),
|
||||
sdkAuth.NewKiloAuthenticator(),
|
||||
sdkAuth.NewGitLabAuthenticator(),
|
||||
sdkAuth.NewCodeBuddyAuthenticator(),
|
||||
sdkAuth.NewCursorAuthenticator(),
|
||||
)
|
||||
return manager
|
||||
}
|
||||
|
||||
43
internal/cmd/codebuddy_login.go
Normal file
43
internal/cmd/codebuddy_login.go
Normal file
@@ -0,0 +1,43 @@
|
||||
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"
|
||||
)
|
||||
|
||||
// DoCodeBuddyLogin triggers the browser OAuth polling flow for CodeBuddy and saves tokens.
|
||||
// It initiates the OAuth authentication, displays the user code for the user to enter
|
||||
// at the CodeBuddy verification URL, and waits for authorization before saving the tokens.
|
||||
//
|
||||
// Parameters:
|
||||
// - cfg: The application configuration containing proxy and auth directory settings
|
||||
// - options: Login options including browser behavior settings
|
||||
func DoCodeBuddyLogin(cfg *config.Config, options *LoginOptions) {
|
||||
if options == nil {
|
||||
options = &LoginOptions{}
|
||||
}
|
||||
|
||||
manager := newAuthManager()
|
||||
authOpts := &sdkAuth.LoginOptions{
|
||||
NoBrowser: options.NoBrowser,
|
||||
Metadata: map[string]string{},
|
||||
}
|
||||
|
||||
record, savedPath, err := manager.Login(context.Background(), "codebuddy", cfg, authOpts)
|
||||
if err != nil {
|
||||
log.Errorf("CodeBuddy authentication failed: %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("CodeBuddy authentication successful!")
|
||||
}
|
||||
37
internal/cmd/cursor_login.go
Normal file
37
internal/cmd/cursor_login.go
Normal file
@@ -0,0 +1,37 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// DoCursorLogin triggers the OAuth PKCE flow for Cursor and saves tokens.
|
||||
func DoCursorLogin(cfg *config.Config, options *LoginOptions) {
|
||||
if options == nil {
|
||||
options = &LoginOptions{}
|
||||
}
|
||||
|
||||
manager := newAuthManager()
|
||||
authOpts := &sdkAuth.LoginOptions{
|
||||
NoBrowser: options.NoBrowser,
|
||||
Metadata: map[string]string{},
|
||||
Prompt: options.Prompt,
|
||||
}
|
||||
|
||||
record, savedPath, err := manager.Login(context.Background(), "cursor", cfg, authOpts)
|
||||
if err != nil {
|
||||
log.Errorf("Cursor authentication failed: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if savedPath != "" {
|
||||
log.Infof("Authentication saved to %s", savedPath)
|
||||
}
|
||||
if record != nil && record.Label != "" {
|
||||
log.Infof("Authenticated as %s", record.Label)
|
||||
}
|
||||
log.Info("Cursor authentication successful!")
|
||||
}
|
||||
69
internal/cmd/gitlab_login.go
Normal file
69
internal/cmd/gitlab_login.go
Normal file
@@ -0,0 +1,69 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
|
||||
)
|
||||
|
||||
func DoGitLabLogin(cfg *config.Config, options *LoginOptions) {
|
||||
if options == nil {
|
||||
options = &LoginOptions{}
|
||||
}
|
||||
|
||||
promptFn := options.Prompt
|
||||
if promptFn == nil {
|
||||
promptFn = defaultProjectPrompt()
|
||||
}
|
||||
|
||||
manager := newAuthManager()
|
||||
authOpts := &sdkAuth.LoginOptions{
|
||||
NoBrowser: options.NoBrowser,
|
||||
CallbackPort: options.CallbackPort,
|
||||
Metadata: map[string]string{
|
||||
"login_mode": "oauth",
|
||||
},
|
||||
Prompt: promptFn,
|
||||
}
|
||||
|
||||
_, savedPath, err := manager.Login(context.Background(), "gitlab", cfg, authOpts)
|
||||
if err != nil {
|
||||
fmt.Printf("GitLab Duo authentication failed: %v\n", err)
|
||||
return
|
||||
}
|
||||
if savedPath != "" {
|
||||
fmt.Printf("Authentication saved to %s\n", savedPath)
|
||||
}
|
||||
fmt.Println("GitLab Duo authentication successful!")
|
||||
}
|
||||
|
||||
func DoGitLabTokenLogin(cfg *config.Config, options *LoginOptions) {
|
||||
if options == nil {
|
||||
options = &LoginOptions{}
|
||||
}
|
||||
|
||||
promptFn := options.Prompt
|
||||
if promptFn == nil {
|
||||
promptFn = defaultProjectPrompt()
|
||||
}
|
||||
|
||||
manager := newAuthManager()
|
||||
authOpts := &sdkAuth.LoginOptions{
|
||||
Metadata: map[string]string{
|
||||
"login_mode": "pat",
|
||||
},
|
||||
Prompt: promptFn,
|
||||
}
|
||||
|
||||
_, savedPath, err := manager.Login(context.Background(), "gitlab", cfg, authOpts)
|
||||
if err != nil {
|
||||
fmt.Printf("GitLab Duo PAT authentication failed: %v\n", err)
|
||||
return
|
||||
}
|
||||
if savedPath != "" {
|
||||
fmt.Printf("Authentication saved to %s\n", savedPath)
|
||||
}
|
||||
fmt.Println("GitLab Duo PAT authentication successful!")
|
||||
}
|
||||
@@ -206,3 +206,52 @@ func DoKiroImport(cfg *config.Config, options *LoginOptions) {
|
||||
}
|
||||
fmt.Println("Kiro token import successful!")
|
||||
}
|
||||
|
||||
func DoKiroIDCLogin(cfg *config.Config, options *LoginOptions, startURL, region, flow string) {
|
||||
if options == nil {
|
||||
options = &LoginOptions{}
|
||||
}
|
||||
|
||||
if startURL == "" {
|
||||
log.Errorf("Kiro IDC login requires --kiro-idc-start-url")
|
||||
fmt.Println("\nUsage: --kiro-idc-login --kiro-idc-start-url https://d-xxx.awsapps.com/start")
|
||||
return
|
||||
}
|
||||
|
||||
manager := newAuthManager()
|
||||
|
||||
authenticator := sdkAuth.NewKiroAuthenticator()
|
||||
metadata := map[string]string{
|
||||
"start-url": startURL,
|
||||
"region": region,
|
||||
"flow": flow,
|
||||
}
|
||||
|
||||
record, err := authenticator.Login(context.Background(), cfg, &sdkAuth.LoginOptions{
|
||||
NoBrowser: options.NoBrowser,
|
||||
Metadata: metadata,
|
||||
Prompt: options.Prompt,
|
||||
})
|
||||
if err != nil {
|
||||
log.Errorf("Kiro IDC authentication failed: %v", err)
|
||||
fmt.Println("\nTroubleshooting:")
|
||||
fmt.Println("1. Make sure your IDC Start URL is correct")
|
||||
fmt.Println("2. Complete the authorization in the browser")
|
||||
fmt.Println("3. If auth code flow fails, try: --kiro-idc-flow device")
|
||||
return
|
||||
}
|
||||
|
||||
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 IDC authentication successful!")
|
||||
}
|
||||
|
||||
@@ -20,6 +20,7 @@ import (
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/auth/gemini"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
|
||||
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
|
||||
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
log "github.com/sirupsen/logrus"
|
||||
@@ -27,11 +28,8 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
geminiCLIEndpoint = "https://cloudcode-pa.googleapis.com"
|
||||
geminiCLIVersion = "v1internal"
|
||||
geminiCLIUserAgent = "google-api-nodejs-client/9.15.1"
|
||||
geminiCLIApiClient = "gl-node/22.17.0"
|
||||
geminiCLIClientMetadata = "ideType=IDE_UNSPECIFIED,platform=PLATFORM_UNSPECIFIED,pluginType=GEMINI"
|
||||
geminiCLIEndpoint = "https://cloudcode-pa.googleapis.com"
|
||||
geminiCLIVersion = "v1internal"
|
||||
)
|
||||
|
||||
type projectSelectionRequiredError struct{}
|
||||
@@ -409,9 +407,7 @@ func callGeminiCLI(ctx context.Context, httpClient *http.Client, endpoint string
|
||||
return fmt.Errorf("create request: %w", errRequest)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", geminiCLIUserAgent)
|
||||
req.Header.Set("X-Goog-Api-Client", geminiCLIApiClient)
|
||||
req.Header.Set("Client-Metadata", geminiCLIClientMetadata)
|
||||
req.Header.Set("User-Agent", misc.GeminiCLIUserAgent(""))
|
||||
|
||||
resp, errDo := httpClient.Do(req)
|
||||
if errDo != nil {
|
||||
@@ -630,7 +626,7 @@ func checkCloudAPIIsEnabled(ctx context.Context, httpClient *http.Client, projec
|
||||
return false, fmt.Errorf("failed to create request: %w", errRequest)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", geminiCLIUserAgent)
|
||||
req.Header.Set("User-Agent", misc.GeminiCLIUserAgent(""))
|
||||
resp, errDo := httpClient.Do(req)
|
||||
if errDo != nil {
|
||||
return false, fmt.Errorf("failed to execute request: %w", errDo)
|
||||
@@ -651,7 +647,7 @@ func checkCloudAPIIsEnabled(ctx context.Context, httpClient *http.Client, projec
|
||||
return false, fmt.Errorf("failed to create request: %w", errRequest)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", geminiCLIUserAgent)
|
||||
req.Header.Set("User-Agent", misc.GeminiCLIUserAgent(""))
|
||||
resp, errDo = httpClient.Do(req)
|
||||
if errDo != nil {
|
||||
return false, fmt.Errorf("failed to execute request: %w", errDo)
|
||||
|
||||
60
internal/cmd/openai_device_login.go
Normal file
60
internal/cmd/openai_device_login.go
Normal file
@@ -0,0 +1,60 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/auth/codex"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
codexLoginModeMetadataKey = "codex_login_mode"
|
||||
codexLoginModeDevice = "device"
|
||||
)
|
||||
|
||||
// DoCodexDeviceLogin triggers the Codex device-code flow while keeping the
|
||||
// existing codex-login OAuth callback flow intact.
|
||||
func DoCodexDeviceLogin(cfg *config.Config, options *LoginOptions) {
|
||||
if options == nil {
|
||||
options = &LoginOptions{}
|
||||
}
|
||||
|
||||
promptFn := options.Prompt
|
||||
if promptFn == nil {
|
||||
promptFn = defaultProjectPrompt()
|
||||
}
|
||||
|
||||
manager := newAuthManager()
|
||||
|
||||
authOpts := &sdkAuth.LoginOptions{
|
||||
NoBrowser: options.NoBrowser,
|
||||
CallbackPort: options.CallbackPort,
|
||||
Metadata: map[string]string{
|
||||
codexLoginModeMetadataKey: codexLoginModeDevice,
|
||||
},
|
||||
Prompt: promptFn,
|
||||
}
|
||||
|
||||
_, savedPath, err := manager.Login(context.Background(), "codex", cfg, authOpts)
|
||||
if err != nil {
|
||||
if authErr, ok := errors.AsType[*codex.AuthenticationError](err); ok {
|
||||
log.Error(codex.GetUserFriendlyMessage(authErr))
|
||||
if authErr.Type == codex.ErrPortInUse.Type {
|
||||
os.Exit(codex.ErrPortInUse.Code)
|
||||
}
|
||||
return
|
||||
}
|
||||
fmt.Printf("Codex device authentication failed: %v\n", err)
|
||||
return
|
||||
}
|
||||
|
||||
if savedPath != "" {
|
||||
fmt.Printf("Authentication saved to %s\n", savedPath)
|
||||
}
|
||||
fmt.Println("Codex device authentication successful!")
|
||||
}
|
||||
55
internal/config/claude_header_defaults_test.go
Normal file
55
internal/config/claude_header_defaults_test.go
Normal file
@@ -0,0 +1,55 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLoadConfigOptional_ClaudeHeaderDefaults(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
configPath := filepath.Join(dir, "config.yaml")
|
||||
configYAML := []byte(`
|
||||
claude-header-defaults:
|
||||
user-agent: " claude-cli/2.1.70 (external, cli) "
|
||||
package-version: " 0.80.0 "
|
||||
runtime-version: " v24.5.0 "
|
||||
os: " MacOS "
|
||||
arch: " arm64 "
|
||||
timeout: " 900 "
|
||||
stabilize-device-profile: false
|
||||
`)
|
||||
if err := os.WriteFile(configPath, configYAML, 0o600); err != nil {
|
||||
t.Fatalf("failed to write config: %v", err)
|
||||
}
|
||||
|
||||
cfg, err := LoadConfigOptional(configPath, false)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfigOptional() error = %v", err)
|
||||
}
|
||||
|
||||
if got := cfg.ClaudeHeaderDefaults.UserAgent; got != "claude-cli/2.1.70 (external, cli)" {
|
||||
t.Fatalf("UserAgent = %q, want %q", got, "claude-cli/2.1.70 (external, cli)")
|
||||
}
|
||||
if got := cfg.ClaudeHeaderDefaults.PackageVersion; got != "0.80.0" {
|
||||
t.Fatalf("PackageVersion = %q, want %q", got, "0.80.0")
|
||||
}
|
||||
if got := cfg.ClaudeHeaderDefaults.RuntimeVersion; got != "v24.5.0" {
|
||||
t.Fatalf("RuntimeVersion = %q, want %q", got, "v24.5.0")
|
||||
}
|
||||
if got := cfg.ClaudeHeaderDefaults.OS; got != "MacOS" {
|
||||
t.Fatalf("OS = %q, want %q", got, "MacOS")
|
||||
}
|
||||
if got := cfg.ClaudeHeaderDefaults.Arch; got != "arm64" {
|
||||
t.Fatalf("Arch = %q, want %q", got, "arm64")
|
||||
}
|
||||
if got := cfg.ClaudeHeaderDefaults.Timeout; got != "900" {
|
||||
t.Fatalf("Timeout = %q, want %q", got, "900")
|
||||
}
|
||||
if cfg.ClaudeHeaderDefaults.StabilizeDeviceProfile == nil {
|
||||
t.Fatal("StabilizeDeviceProfile = nil, want non-nil")
|
||||
}
|
||||
if got := *cfg.ClaudeHeaderDefaults.StabilizeDeviceProfile; got {
|
||||
t.Fatalf("StabilizeDeviceProfile = %v, want false", got)
|
||||
}
|
||||
}
|
||||
32
internal/config/codex_websocket_header_defaults_test.go
Normal file
32
internal/config/codex_websocket_header_defaults_test.go
Normal file
@@ -0,0 +1,32 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLoadConfigOptional_CodexHeaderDefaults(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
configPath := filepath.Join(dir, "config.yaml")
|
||||
configYAML := []byte(`
|
||||
codex-header-defaults:
|
||||
user-agent: " my-codex-client/1.0 "
|
||||
beta-features: " feature-a,feature-b "
|
||||
`)
|
||||
if err := os.WriteFile(configPath, configYAML, 0o600); err != nil {
|
||||
t.Fatalf("failed to write config: %v", err)
|
||||
}
|
||||
|
||||
cfg, err := LoadConfigOptional(configPath, false)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfigOptional() error = %v", err)
|
||||
}
|
||||
|
||||
if got := cfg.CodexHeaderDefaults.UserAgent; got != "my-codex-client/1.0" {
|
||||
t.Fatalf("UserAgent = %q, want %q", got, "my-codex-client/1.0")
|
||||
}
|
||||
if got := cfg.CodexHeaderDefaults.BetaFeatures; got != "feature-a,feature-b" {
|
||||
t.Fatalf("BetaFeatures = %q, want %q", got, "feature-a,feature-b")
|
||||
}
|
||||
}
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"strings"
|
||||
"syscall"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
"gopkg.in/yaml.v3"
|
||||
@@ -69,6 +70,9 @@ type Config struct {
|
||||
|
||||
// RequestRetry defines the retry times when the request failed.
|
||||
RequestRetry int `yaml:"request-retry" json:"request-retry"`
|
||||
// MaxRetryCredentials defines the maximum number of credentials to try for a failed request.
|
||||
// Set to 0 or a negative value to keep trying all available credentials (legacy behavior).
|
||||
MaxRetryCredentials int `yaml:"max-retry-credentials" json:"max-retry-credentials"`
|
||||
// MaxRetryInterval defines the maximum wait time in seconds before retrying a cooled-down credential.
|
||||
MaxRetryInterval int `yaml:"max-retry-interval" json:"max-retry-interval"`
|
||||
|
||||
@@ -87,6 +91,10 @@ type Config struct {
|
||||
// KiroKey defines a list of Kiro (AWS CodeWhisperer) configurations.
|
||||
KiroKey []KiroKey `yaml:"kiro" json:"kiro"`
|
||||
|
||||
// KiroFingerprint defines a global fingerprint configuration for all Kiro requests.
|
||||
// When set, all Kiro requests will use this fixed fingerprint instead of random generation.
|
||||
KiroFingerprint *KiroFingerprintConfig `yaml:"kiro-fingerprint,omitempty" json:"kiro-fingerprint,omitempty"`
|
||||
|
||||
// KiroPreferredEndpoint sets the global default preferred endpoint for all Kiro providers.
|
||||
// Values: "ide" (default, CodeWhisperer) or "cli" (Amazon Q).
|
||||
KiroPreferredEndpoint string `yaml:"kiro-preferred-endpoint" json:"kiro-preferred-endpoint"`
|
||||
@@ -94,6 +102,10 @@ type Config struct {
|
||||
// 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"`
|
||||
|
||||
// CodexHeaderDefaults configures fallback headers for Codex OAuth model requests.
|
||||
// These are used only when the client does not send its own headers.
|
||||
CodexHeaderDefaults CodexHeaderDefaults `yaml:"codex-header-defaults" json:"codex-header-defaults"`
|
||||
|
||||
// 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"`
|
||||
|
||||
@@ -134,13 +146,27 @@ type Config struct {
|
||||
legacyMigrationPending bool `yaml:"-" json:"-"`
|
||||
}
|
||||
|
||||
// ClaudeHeaderDefaults configures default header values injected into Claude API requests
|
||||
// when the client does not send them. Update these when Claude Code releases a new version.
|
||||
// ClaudeHeaderDefaults configures default header values injected into Claude API requests.
|
||||
// In legacy mode, UserAgent/PackageVersion/RuntimeVersion/Timeout act as fallbacks when
|
||||
// the client omits them, while OS/Arch remain runtime-derived. When stabilized device
|
||||
// profiles are enabled, OS/Arch become the pinned platform baseline, while
|
||||
// UserAgent/PackageVersion/RuntimeVersion seed the upgradeable software fingerprint.
|
||||
type ClaudeHeaderDefaults struct {
|
||||
UserAgent string `yaml:"user-agent" json:"user-agent"`
|
||||
PackageVersion string `yaml:"package-version" json:"package-version"`
|
||||
RuntimeVersion string `yaml:"runtime-version" json:"runtime-version"`
|
||||
Timeout string `yaml:"timeout" json:"timeout"`
|
||||
UserAgent string `yaml:"user-agent" json:"user-agent"`
|
||||
PackageVersion string `yaml:"package-version" json:"package-version"`
|
||||
RuntimeVersion string `yaml:"runtime-version" json:"runtime-version"`
|
||||
OS string `yaml:"os" json:"os"`
|
||||
Arch string `yaml:"arch" json:"arch"`
|
||||
Timeout string `yaml:"timeout" json:"timeout"`
|
||||
StabilizeDeviceProfile *bool `yaml:"stabilize-device-profile,omitempty" json:"stabilize-device-profile,omitempty"`
|
||||
}
|
||||
|
||||
// CodexHeaderDefaults configures fallback header values injected into Codex
|
||||
// model requests for OAuth/file-backed auth when the client omits them.
|
||||
// UserAgent applies to HTTP and websocket requests; BetaFeatures only applies to websockets.
|
||||
type CodexHeaderDefaults struct {
|
||||
UserAgent string `yaml:"user-agent" json:"user-agent"`
|
||||
BetaFeatures string `yaml:"beta-features" json:"beta-features"`
|
||||
}
|
||||
|
||||
// TLSConfig holds HTTPS server settings.
|
||||
@@ -169,6 +195,9 @@ type RemoteManagement struct {
|
||||
SecretKey string `yaml:"secret-key"`
|
||||
// DisableControlPanel skips serving and syncing the bundled management UI when true.
|
||||
DisableControlPanel bool `yaml:"disable-control-panel"`
|
||||
// DisableAutoUpdatePanel disables automatic periodic background updates of the management panel asset from GitHub.
|
||||
// When false (the default), the background updater remains enabled; when true, the panel is only downloaded on first access if missing.
|
||||
DisableAutoUpdatePanel bool `yaml:"disable-auto-update-panel"`
|
||||
// PanelGitHubRepository overrides the GitHub repository used to fetch the management panel asset.
|
||||
// Accepts either a repository URL (https://github.com/org/repo) or an API releases endpoint.
|
||||
PanelGitHubRepository string `yaml:"panel-github-repository"`
|
||||
@@ -477,6 +506,9 @@ type KiroKey struct {
|
||||
// Region is the AWS region (default: us-east-1).
|
||||
Region string `yaml:"region,omitempty" json:"region,omitempty"`
|
||||
|
||||
// StartURL is the IAM Identity Center (IDC) start URL for SSO login.
|
||||
StartURL string `yaml:"start-url,omitempty" json:"start-url,omitempty"`
|
||||
|
||||
// ProxyURL optionally overrides the global proxy for this configuration.
|
||||
ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"`
|
||||
|
||||
@@ -489,6 +521,20 @@ type KiroKey struct {
|
||||
PreferredEndpoint string `yaml:"preferred-endpoint,omitempty" json:"preferred-endpoint,omitempty"`
|
||||
}
|
||||
|
||||
// KiroFingerprintConfig defines a global fingerprint configuration for Kiro requests.
|
||||
// When configured, all Kiro requests will use this fixed fingerprint instead of random generation.
|
||||
// Empty fields will fall back to random selection from built-in pools.
|
||||
type KiroFingerprintConfig struct {
|
||||
OIDCSDKVersion string `yaml:"oidc-sdk-version,omitempty" json:"oidc-sdk-version,omitempty"`
|
||||
RuntimeSDKVersion string `yaml:"runtime-sdk-version,omitempty" json:"runtime-sdk-version,omitempty"`
|
||||
StreamingSDKVersion string `yaml:"streaming-sdk-version,omitempty" json:"streaming-sdk-version,omitempty"`
|
||||
OSType string `yaml:"os-type,omitempty" json:"os-type,omitempty"`
|
||||
OSVersion string `yaml:"os-version,omitempty" json:"os-version,omitempty"`
|
||||
NodeVersion string `yaml:"node-version,omitempty" json:"node-version,omitempty"`
|
||||
KiroVersion string `yaml:"kiro-version,omitempty" json:"kiro-version,omitempty"`
|
||||
KiroHash string `yaml:"kiro-hash,omitempty" json:"kiro-hash,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 {
|
||||
@@ -532,6 +578,10 @@ type OpenAICompatibilityModel struct {
|
||||
|
||||
// Alias is the model name alias that clients will use to reference this model.
|
||||
Alias string `yaml:"alias" json:"alias"`
|
||||
|
||||
// Thinking configures the thinking/reasoning capability for this model.
|
||||
// If nil, the model defaults to level-based reasoning with levels ["low", "medium", "high"].
|
||||
Thinking *registry.ThinkingSupport `yaml:"thinking,omitempty" json:"thinking,omitempty"`
|
||||
}
|
||||
|
||||
func (m OpenAICompatibilityModel) GetName() string { return m.Name }
|
||||
@@ -555,16 +605,6 @@ func LoadConfig(configFile string) (*Config, error) {
|
||||
// If optional is true and the file is missing, it returns an empty Config.
|
||||
// If optional is true and the file is empty or invalid, it returns an empty Config.
|
||||
func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
|
||||
// NOTE: Startup oauth-model-alias migration is intentionally disabled.
|
||||
// Reason: avoid mutating config.yaml during server startup.
|
||||
// Re-enable the block below if automatic startup migration is needed again.
|
||||
// if migrated, err := MigrateOAuthModelAlias(configFile); err != nil {
|
||||
// // Log warning but don't fail - config loading should still work
|
||||
// fmt.Printf("Warning: oauth-model-alias migration failed: %v\n", err)
|
||||
// } else if migrated {
|
||||
// fmt.Println("Migrated oauth-model-mappings to oauth-model-alias")
|
||||
// }
|
||||
|
||||
// Read the entire configuration file into memory.
|
||||
data, err := os.ReadFile(configFile)
|
||||
if err != nil {
|
||||
@@ -652,15 +692,25 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
|
||||
cfg.ErrorLogsMaxFiles = 10
|
||||
}
|
||||
|
||||
if cfg.MaxRetryCredentials < 0 {
|
||||
cfg.MaxRetryCredentials = 0
|
||||
}
|
||||
|
||||
// Sanitize Gemini API key configuration and migrate legacy entries.
|
||||
cfg.SanitizeGeminiKeys()
|
||||
|
||||
// Sanitize Vertex-compatible API keys: drop entries without base-url
|
||||
// Sanitize Vertex-compatible API keys.
|
||||
cfg.SanitizeVertexCompatKeys()
|
||||
|
||||
// Sanitize Codex keys: drop entries without base-url
|
||||
cfg.SanitizeCodexKeys()
|
||||
|
||||
// Sanitize Codex header defaults.
|
||||
cfg.SanitizeCodexHeaderDefaults()
|
||||
|
||||
// Sanitize Claude header defaults.
|
||||
cfg.SanitizeClaudeHeaderDefaults()
|
||||
|
||||
// Sanitize Claude key headers
|
||||
cfg.SanitizeClaudeKeys()
|
||||
|
||||
@@ -753,6 +803,30 @@ func payloadRawString(value any) ([]byte, bool) {
|
||||
}
|
||||
}
|
||||
|
||||
// SanitizeCodexHeaderDefaults trims surrounding whitespace from the
|
||||
// configured Codex header fallback values.
|
||||
func (cfg *Config) SanitizeCodexHeaderDefaults() {
|
||||
if cfg == nil {
|
||||
return
|
||||
}
|
||||
cfg.CodexHeaderDefaults.UserAgent = strings.TrimSpace(cfg.CodexHeaderDefaults.UserAgent)
|
||||
cfg.CodexHeaderDefaults.BetaFeatures = strings.TrimSpace(cfg.CodexHeaderDefaults.BetaFeatures)
|
||||
}
|
||||
|
||||
// SanitizeClaudeHeaderDefaults trims surrounding whitespace from the
|
||||
// configured Claude fingerprint baseline values.
|
||||
func (cfg *Config) SanitizeClaudeHeaderDefaults() {
|
||||
if cfg == nil {
|
||||
return
|
||||
}
|
||||
cfg.ClaudeHeaderDefaults.UserAgent = strings.TrimSpace(cfg.ClaudeHeaderDefaults.UserAgent)
|
||||
cfg.ClaudeHeaderDefaults.PackageVersion = strings.TrimSpace(cfg.ClaudeHeaderDefaults.PackageVersion)
|
||||
cfg.ClaudeHeaderDefaults.RuntimeVersion = strings.TrimSpace(cfg.ClaudeHeaderDefaults.RuntimeVersion)
|
||||
cfg.ClaudeHeaderDefaults.OS = strings.TrimSpace(cfg.ClaudeHeaderDefaults.OS)
|
||||
cfg.ClaudeHeaderDefaults.Arch = strings.TrimSpace(cfg.ClaudeHeaderDefaults.Arch)
|
||||
cfg.ClaudeHeaderDefaults.Timeout = strings.TrimSpace(cfg.ClaudeHeaderDefaults.Timeout)
|
||||
}
|
||||
|
||||
// SanitizeOAuthModelAlias normalizes and deduplicates global OAuth model name aliases.
|
||||
// It trims whitespace, normalizes channel keys to lower-case, drops empty entries,
|
||||
// allows multiple aliases per upstream name, and ensures aliases are unique within each channel.
|
||||
@@ -1648,9 +1722,6 @@ func pruneMappingToGeneratedKeys(dstRoot, srcRoot *yaml.Node, key string) {
|
||||
srcIdx := findMapKeyIndex(srcRoot, key)
|
||||
if srcIdx < 0 {
|
||||
// Keep an explicit empty mapping for oauth-model-alias when it was previously present.
|
||||
//
|
||||
// Rationale: LoadConfig runs MigrateOAuthModelAlias before unmarshalling. If the
|
||||
// oauth-model-alias key is missing, migration will add the default antigravity aliases.
|
||||
// When users delete the last channel from oauth-model-alias via the management API,
|
||||
// we want that deletion to persist across hot reloads and restarts.
|
||||
if key == "oauth-model-alias" {
|
||||
|
||||
61
internal/config/oauth_model_alias_defaults.go
Normal file
61
internal/config/oauth_model_alias_defaults.go
Normal file
@@ -0,0 +1,61 @@
|
||||
package config
|
||||
|
||||
import "strings"
|
||||
|
||||
// defaultKiroAliases returns default oauth-model-alias entries for Kiro.
|
||||
// These aliases expose standard Claude IDs for Kiro-prefixed upstream models.
|
||||
func defaultKiroAliases() []OAuthModelAlias {
|
||||
return []OAuthModelAlias{
|
||||
// Sonnet 4.6
|
||||
{Name: "kiro-claude-sonnet-4-6", Alias: "claude-sonnet-4-6", Fork: true},
|
||||
// Sonnet 4.5
|
||||
{Name: "kiro-claude-sonnet-4-5", Alias: "claude-sonnet-4-5-20250929", Fork: true},
|
||||
{Name: "kiro-claude-sonnet-4-5", Alias: "claude-sonnet-4-5", Fork: true},
|
||||
// Sonnet 4
|
||||
{Name: "kiro-claude-sonnet-4", Alias: "claude-sonnet-4-20250514", Fork: true},
|
||||
{Name: "kiro-claude-sonnet-4", Alias: "claude-sonnet-4", Fork: true},
|
||||
// Opus 4.6
|
||||
{Name: "kiro-claude-opus-4-6", Alias: "claude-opus-4-6", Fork: true},
|
||||
// Opus 4.5
|
||||
{Name: "kiro-claude-opus-4-5", Alias: "claude-opus-4-5-20251101", Fork: true},
|
||||
{Name: "kiro-claude-opus-4-5", Alias: "claude-opus-4-5", Fork: true},
|
||||
// Haiku 4.5
|
||||
{Name: "kiro-claude-haiku-4-5", Alias: "claude-haiku-4-5-20251001", Fork: true},
|
||||
{Name: "kiro-claude-haiku-4-5", Alias: "claude-haiku-4-5", Fork: true},
|
||||
}
|
||||
}
|
||||
|
||||
// defaultGitHubCopilotAliases returns default oauth-model-alias entries for
|
||||
// GitHub Copilot Claude models. It exposes hyphen-style IDs used by clients.
|
||||
func defaultGitHubCopilotAliases() []OAuthModelAlias {
|
||||
return []OAuthModelAlias{
|
||||
{Name: "claude-haiku-4.5", Alias: "claude-haiku-4-5", Fork: true},
|
||||
{Name: "claude-opus-4.1", Alias: "claude-opus-4-1", Fork: true},
|
||||
{Name: "claude-opus-4.5", Alias: "claude-opus-4-5", Fork: true},
|
||||
{Name: "claude-opus-4.6", Alias: "claude-opus-4-6", Fork: true},
|
||||
{Name: "claude-sonnet-4.5", Alias: "claude-sonnet-4-5", Fork: true},
|
||||
{Name: "claude-sonnet-4.6", Alias: "claude-sonnet-4-6", Fork: true},
|
||||
}
|
||||
}
|
||||
|
||||
// GitHubCopilotAliasesFromModels generates oauth-model-alias entries from a dynamic
|
||||
// list of model IDs fetched from the Copilot API. It auto-creates aliases for
|
||||
// models whose ID contains a dot (e.g. "claude-opus-4.6" → "claude-opus-4-6"),
|
||||
// which is the pattern used by Claude models on Copilot.
|
||||
func GitHubCopilotAliasesFromModels(modelIDs []string) []OAuthModelAlias {
|
||||
var aliases []OAuthModelAlias
|
||||
seen := make(map[string]struct{})
|
||||
for _, id := range modelIDs {
|
||||
if !strings.Contains(id, ".") {
|
||||
continue
|
||||
}
|
||||
hyphenID := strings.ReplaceAll(id, ".", "-")
|
||||
key := id + "→" + hyphenID
|
||||
if _, ok := seen[key]; ok {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
aliases = append(aliases, OAuthModelAlias{Name: id, Alias: hyphenID, Fork: true})
|
||||
}
|
||||
return aliases
|
||||
}
|
||||
@@ -1,314 +0,0 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// antigravityModelConversionTable maps old built-in aliases to actual model names
|
||||
// for the antigravity channel during migration.
|
||||
var antigravityModelConversionTable = map[string]string{
|
||||
"gemini-2.5-computer-use-preview-10-2025": "rev19-uic3-1p",
|
||||
"gemini-3-pro-image-preview": "gemini-3-pro-image",
|
||||
"gemini-3-pro-preview": "gemini-3-pro-high",
|
||||
"gemini-3-flash-preview": "gemini-3-flash",
|
||||
"gemini-claude-sonnet-4-5": "claude-sonnet-4-5",
|
||||
"gemini-claude-sonnet-4-5-thinking": "claude-sonnet-4-5-thinking",
|
||||
"gemini-claude-opus-4-5-thinking": "claude-opus-4-5-thinking",
|
||||
"gemini-claude-opus-4-6-thinking": "claude-opus-4-6-thinking",
|
||||
}
|
||||
|
||||
// defaultKiroAliases returns the default oauth-model-alias configuration
|
||||
// for the kiro channel. Maps kiro-prefixed model names to standard Claude model
|
||||
// names so that clients like Claude Code can use standard names directly.
|
||||
func defaultKiroAliases() []OAuthModelAlias {
|
||||
return []OAuthModelAlias{
|
||||
// Sonnet 4.5
|
||||
{Name: "kiro-claude-sonnet-4-5", Alias: "claude-sonnet-4-5-20250929", Fork: true},
|
||||
{Name: "kiro-claude-sonnet-4-5", Alias: "claude-sonnet-4-5", Fork: true},
|
||||
// Sonnet 4
|
||||
{Name: "kiro-claude-sonnet-4", Alias: "claude-sonnet-4-20250514", Fork: true},
|
||||
{Name: "kiro-claude-sonnet-4", Alias: "claude-sonnet-4", Fork: true},
|
||||
// Opus 4.6
|
||||
{Name: "kiro-claude-opus-4-6", Alias: "claude-opus-4-6", Fork: true},
|
||||
// Opus 4.5
|
||||
{Name: "kiro-claude-opus-4-5", Alias: "claude-opus-4-5-20251101", Fork: true},
|
||||
{Name: "kiro-claude-opus-4-5", Alias: "claude-opus-4-5", Fork: true},
|
||||
// Haiku 4.5
|
||||
{Name: "kiro-claude-haiku-4-5", Alias: "claude-haiku-4-5-20251001", Fork: true},
|
||||
{Name: "kiro-claude-haiku-4-5", Alias: "claude-haiku-4-5", Fork: true},
|
||||
}
|
||||
}
|
||||
|
||||
// defaultGitHubCopilotAliases returns default oauth-model-alias entries that
|
||||
// expose Claude hyphen-style IDs for GitHub Copilot Claude models.
|
||||
// This keeps compatibility with clients (e.g. Claude Code) that use
|
||||
// Anthropic-style model IDs like "claude-opus-4-6".
|
||||
func defaultGitHubCopilotAliases() []OAuthModelAlias {
|
||||
return []OAuthModelAlias{
|
||||
{Name: "claude-haiku-4.5", Alias: "claude-haiku-4-5", Fork: true},
|
||||
{Name: "claude-opus-4.1", Alias: "claude-opus-4-1", Fork: true},
|
||||
{Name: "claude-opus-4.5", Alias: "claude-opus-4-5", Fork: true},
|
||||
{Name: "claude-opus-4.6", Alias: "claude-opus-4-6", Fork: true},
|
||||
{Name: "claude-sonnet-4.5", Alias: "claude-sonnet-4-5", Fork: true},
|
||||
{Name: "claude-sonnet-4.6", Alias: "claude-sonnet-4-6", Fork: true},
|
||||
}
|
||||
}
|
||||
|
||||
// defaultAntigravityAliases returns the default oauth-model-alias configuration
|
||||
// for the antigravity channel when neither field exists.
|
||||
func defaultAntigravityAliases() []OAuthModelAlias {
|
||||
return []OAuthModelAlias{
|
||||
{Name: "rev19-uic3-1p", Alias: "gemini-2.5-computer-use-preview-10-2025"},
|
||||
{Name: "gemini-3-pro-image", Alias: "gemini-3-pro-image-preview"},
|
||||
{Name: "gemini-3-pro-high", Alias: "gemini-3-pro-preview"},
|
||||
{Name: "gemini-3-flash", Alias: "gemini-3-flash-preview"},
|
||||
{Name: "claude-sonnet-4-5", Alias: "gemini-claude-sonnet-4-5"},
|
||||
{Name: "claude-sonnet-4-5-thinking", Alias: "gemini-claude-sonnet-4-5-thinking"},
|
||||
{Name: "claude-opus-4-5-thinking", Alias: "gemini-claude-opus-4-5-thinking"},
|
||||
{Name: "claude-opus-4-6-thinking", Alias: "gemini-claude-opus-4-6-thinking"},
|
||||
}
|
||||
}
|
||||
|
||||
// MigrateOAuthModelAlias checks for and performs migration from oauth-model-mappings
|
||||
// to oauth-model-alias at startup. Returns true if migration was performed.
|
||||
//
|
||||
// Migration flow:
|
||||
// 1. Check if oauth-model-alias exists -> skip migration
|
||||
// 2. Check if oauth-model-mappings exists -> convert and migrate
|
||||
// - For antigravity channel, convert old built-in aliases to actual model names
|
||||
//
|
||||
// 3. Neither exists -> add default antigravity config
|
||||
func MigrateOAuthModelAlias(configFile string) (bool, error) {
|
||||
data, err := os.ReadFile(configFile)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return false, nil
|
||||
}
|
||||
return false, err
|
||||
}
|
||||
if len(data) == 0 {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// Parse YAML into node tree to preserve structure
|
||||
var root yaml.Node
|
||||
if err := yaml.Unmarshal(data, &root); err != nil {
|
||||
return false, nil
|
||||
}
|
||||
if root.Kind != yaml.DocumentNode || len(root.Content) == 0 {
|
||||
return false, nil
|
||||
}
|
||||
rootMap := root.Content[0]
|
||||
if rootMap == nil || rootMap.Kind != yaml.MappingNode {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// Check if oauth-model-alias already exists
|
||||
if findMapKeyIndex(rootMap, "oauth-model-alias") >= 0 {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// Check if oauth-model-mappings exists
|
||||
oldIdx := findMapKeyIndex(rootMap, "oauth-model-mappings")
|
||||
if oldIdx >= 0 {
|
||||
// Migrate from old field
|
||||
return migrateFromOldField(configFile, &root, rootMap, oldIdx)
|
||||
}
|
||||
|
||||
// Neither field exists - add default antigravity config
|
||||
return addDefaultAntigravityConfig(configFile, &root, rootMap)
|
||||
}
|
||||
|
||||
// migrateFromOldField converts oauth-model-mappings to oauth-model-alias
|
||||
func migrateFromOldField(configFile string, root *yaml.Node, rootMap *yaml.Node, oldIdx int) (bool, error) {
|
||||
if oldIdx+1 >= len(rootMap.Content) {
|
||||
return false, nil
|
||||
}
|
||||
oldValue := rootMap.Content[oldIdx+1]
|
||||
if oldValue == nil || oldValue.Kind != yaml.MappingNode {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// Parse the old aliases
|
||||
oldAliases := parseOldAliasNode(oldValue)
|
||||
if len(oldAliases) == 0 {
|
||||
// Remove the old field and write
|
||||
removeMapKeyByIndex(rootMap, oldIdx)
|
||||
return writeYAMLNode(configFile, root)
|
||||
}
|
||||
|
||||
// Convert model names for antigravity channel
|
||||
newAliases := make(map[string][]OAuthModelAlias, len(oldAliases))
|
||||
for channel, entries := range oldAliases {
|
||||
converted := make([]OAuthModelAlias, 0, len(entries))
|
||||
for _, entry := range entries {
|
||||
newEntry := OAuthModelAlias{
|
||||
Name: entry.Name,
|
||||
Alias: entry.Alias,
|
||||
Fork: entry.Fork,
|
||||
}
|
||||
// Convert model names for antigravity channel
|
||||
if strings.EqualFold(channel, "antigravity") {
|
||||
if actual, ok := antigravityModelConversionTable[entry.Name]; ok {
|
||||
newEntry.Name = actual
|
||||
}
|
||||
}
|
||||
converted = append(converted, newEntry)
|
||||
}
|
||||
newAliases[channel] = converted
|
||||
}
|
||||
|
||||
// For antigravity channel, supplement missing default aliases
|
||||
if antigravityEntries, exists := newAliases["antigravity"]; exists {
|
||||
// Build a set of already configured model names (upstream names)
|
||||
configuredModels := make(map[string]bool, len(antigravityEntries))
|
||||
for _, entry := range antigravityEntries {
|
||||
configuredModels[entry.Name] = true
|
||||
}
|
||||
|
||||
// Add missing default aliases
|
||||
for _, defaultAlias := range defaultAntigravityAliases() {
|
||||
if !configuredModels[defaultAlias.Name] {
|
||||
antigravityEntries = append(antigravityEntries, defaultAlias)
|
||||
}
|
||||
}
|
||||
newAliases["antigravity"] = antigravityEntries
|
||||
}
|
||||
|
||||
// Build new node
|
||||
newNode := buildOAuthModelAliasNode(newAliases)
|
||||
|
||||
// Replace old key with new key and value
|
||||
rootMap.Content[oldIdx].Value = "oauth-model-alias"
|
||||
rootMap.Content[oldIdx+1] = newNode
|
||||
|
||||
return writeYAMLNode(configFile, root)
|
||||
}
|
||||
|
||||
// addDefaultAntigravityConfig adds the default antigravity configuration
|
||||
func addDefaultAntigravityConfig(configFile string, root *yaml.Node, rootMap *yaml.Node) (bool, error) {
|
||||
defaults := map[string][]OAuthModelAlias{
|
||||
"antigravity": defaultAntigravityAliases(),
|
||||
}
|
||||
newNode := buildOAuthModelAliasNode(defaults)
|
||||
|
||||
// Add new key-value pair
|
||||
keyNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: "oauth-model-alias"}
|
||||
rootMap.Content = append(rootMap.Content, keyNode, newNode)
|
||||
|
||||
return writeYAMLNode(configFile, root)
|
||||
}
|
||||
|
||||
// parseOldAliasNode parses the old oauth-model-mappings node structure
|
||||
func parseOldAliasNode(node *yaml.Node) map[string][]OAuthModelAlias {
|
||||
if node == nil || node.Kind != yaml.MappingNode {
|
||||
return nil
|
||||
}
|
||||
result := make(map[string][]OAuthModelAlias)
|
||||
for i := 0; i+1 < len(node.Content); i += 2 {
|
||||
channelNode := node.Content[i]
|
||||
entriesNode := node.Content[i+1]
|
||||
if channelNode == nil || entriesNode == nil {
|
||||
continue
|
||||
}
|
||||
channel := strings.ToLower(strings.TrimSpace(channelNode.Value))
|
||||
if channel == "" || entriesNode.Kind != yaml.SequenceNode {
|
||||
continue
|
||||
}
|
||||
entries := make([]OAuthModelAlias, 0, len(entriesNode.Content))
|
||||
for _, entryNode := range entriesNode.Content {
|
||||
if entryNode == nil || entryNode.Kind != yaml.MappingNode {
|
||||
continue
|
||||
}
|
||||
entry := parseAliasEntry(entryNode)
|
||||
if entry.Name != "" && entry.Alias != "" {
|
||||
entries = append(entries, entry)
|
||||
}
|
||||
}
|
||||
if len(entries) > 0 {
|
||||
result[channel] = entries
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// parseAliasEntry parses a single alias entry node
|
||||
func parseAliasEntry(node *yaml.Node) OAuthModelAlias {
|
||||
var entry OAuthModelAlias
|
||||
for i := 0; i+1 < len(node.Content); i += 2 {
|
||||
keyNode := node.Content[i]
|
||||
valNode := node.Content[i+1]
|
||||
if keyNode == nil || valNode == nil {
|
||||
continue
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(keyNode.Value)) {
|
||||
case "name":
|
||||
entry.Name = strings.TrimSpace(valNode.Value)
|
||||
case "alias":
|
||||
entry.Alias = strings.TrimSpace(valNode.Value)
|
||||
case "fork":
|
||||
entry.Fork = strings.ToLower(strings.TrimSpace(valNode.Value)) == "true"
|
||||
}
|
||||
}
|
||||
return entry
|
||||
}
|
||||
|
||||
// buildOAuthModelAliasNode creates a YAML node for oauth-model-alias
|
||||
func buildOAuthModelAliasNode(aliases map[string][]OAuthModelAlias) *yaml.Node {
|
||||
node := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
|
||||
for channel, entries := range aliases {
|
||||
channelNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: channel}
|
||||
entriesNode := &yaml.Node{Kind: yaml.SequenceNode, Tag: "!!seq"}
|
||||
for _, entry := range entries {
|
||||
entryNode := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
|
||||
entryNode.Content = append(entryNode.Content,
|
||||
&yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: "name"},
|
||||
&yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: entry.Name},
|
||||
&yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: "alias"},
|
||||
&yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: entry.Alias},
|
||||
)
|
||||
if entry.Fork {
|
||||
entryNode.Content = append(entryNode.Content,
|
||||
&yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: "fork"},
|
||||
&yaml.Node{Kind: yaml.ScalarNode, Tag: "!!bool", Value: "true"},
|
||||
)
|
||||
}
|
||||
entriesNode.Content = append(entriesNode.Content, entryNode)
|
||||
}
|
||||
node.Content = append(node.Content, channelNode, entriesNode)
|
||||
}
|
||||
return node
|
||||
}
|
||||
|
||||
// removeMapKeyByIndex removes a key-value pair from a mapping node by index
|
||||
func removeMapKeyByIndex(mapNode *yaml.Node, keyIdx int) {
|
||||
if mapNode == nil || mapNode.Kind != yaml.MappingNode {
|
||||
return
|
||||
}
|
||||
if keyIdx < 0 || keyIdx+1 >= len(mapNode.Content) {
|
||||
return
|
||||
}
|
||||
mapNode.Content = append(mapNode.Content[:keyIdx], mapNode.Content[keyIdx+2:]...)
|
||||
}
|
||||
|
||||
// writeYAMLNode writes the YAML node tree back to file
|
||||
func writeYAMLNode(configFile string, root *yaml.Node) (bool, error) {
|
||||
f, err := os.Create(configFile)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
enc := yaml.NewEncoder(f)
|
||||
enc.SetIndent(2)
|
||||
if err := enc.Encode(root); err != nil {
|
||||
return false, err
|
||||
}
|
||||
if err := enc.Close(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
@@ -1,245 +0,0 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
func TestMigrateOAuthModelAlias_SkipsIfNewFieldExists(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dir := t.TempDir()
|
||||
configFile := filepath.Join(dir, "config.yaml")
|
||||
|
||||
content := `oauth-model-alias:
|
||||
gemini-cli:
|
||||
- name: "gemini-2.5-pro"
|
||||
alias: "g2.5p"
|
||||
`
|
||||
if err := os.WriteFile(configFile, []byte(content), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
migrated, err := MigrateOAuthModelAlias(configFile)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if migrated {
|
||||
t.Fatal("expected no migration when oauth-model-alias already exists")
|
||||
}
|
||||
|
||||
// Verify file unchanged
|
||||
data, _ := os.ReadFile(configFile)
|
||||
if !strings.Contains(string(data), "oauth-model-alias:") {
|
||||
t.Fatal("file should still contain oauth-model-alias")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrateOAuthModelAlias_MigratesOldField(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dir := t.TempDir()
|
||||
configFile := filepath.Join(dir, "config.yaml")
|
||||
|
||||
content := `oauth-model-mappings:
|
||||
gemini-cli:
|
||||
- name: "gemini-2.5-pro"
|
||||
alias: "g2.5p"
|
||||
fork: true
|
||||
`
|
||||
if err := os.WriteFile(configFile, []byte(content), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
migrated, err := MigrateOAuthModelAlias(configFile)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if !migrated {
|
||||
t.Fatal("expected migration to occur")
|
||||
}
|
||||
|
||||
// Verify new field exists and old field removed
|
||||
data, _ := os.ReadFile(configFile)
|
||||
if strings.Contains(string(data), "oauth-model-mappings:") {
|
||||
t.Fatal("old field should be removed")
|
||||
}
|
||||
if !strings.Contains(string(data), "oauth-model-alias:") {
|
||||
t.Fatal("new field should exist")
|
||||
}
|
||||
|
||||
// Parse and verify structure
|
||||
var root yaml.Node
|
||||
if err := yaml.Unmarshal(data, &root); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrateOAuthModelAlias_ConvertsAntigravityModels(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dir := t.TempDir()
|
||||
configFile := filepath.Join(dir, "config.yaml")
|
||||
|
||||
// Use old model names that should be converted
|
||||
content := `oauth-model-mappings:
|
||||
antigravity:
|
||||
- name: "gemini-2.5-computer-use-preview-10-2025"
|
||||
alias: "computer-use"
|
||||
- name: "gemini-3-pro-preview"
|
||||
alias: "g3p"
|
||||
`
|
||||
if err := os.WriteFile(configFile, []byte(content), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
migrated, err := MigrateOAuthModelAlias(configFile)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if !migrated {
|
||||
t.Fatal("expected migration to occur")
|
||||
}
|
||||
|
||||
// Verify model names were converted
|
||||
data, _ := os.ReadFile(configFile)
|
||||
content = string(data)
|
||||
if !strings.Contains(content, "rev19-uic3-1p") {
|
||||
t.Fatal("expected gemini-2.5-computer-use-preview-10-2025 to be converted to rev19-uic3-1p")
|
||||
}
|
||||
if !strings.Contains(content, "gemini-3-pro-high") {
|
||||
t.Fatal("expected gemini-3-pro-preview to be converted to gemini-3-pro-high")
|
||||
}
|
||||
|
||||
// Verify missing default aliases were supplemented
|
||||
if !strings.Contains(content, "gemini-3-pro-image") {
|
||||
t.Fatal("expected missing default alias gemini-3-pro-image to be added")
|
||||
}
|
||||
if !strings.Contains(content, "gemini-3-flash") {
|
||||
t.Fatal("expected missing default alias gemini-3-flash to be added")
|
||||
}
|
||||
if !strings.Contains(content, "claude-sonnet-4-5") {
|
||||
t.Fatal("expected missing default alias claude-sonnet-4-5 to be added")
|
||||
}
|
||||
if !strings.Contains(content, "claude-sonnet-4-5-thinking") {
|
||||
t.Fatal("expected missing default alias claude-sonnet-4-5-thinking to be added")
|
||||
}
|
||||
if !strings.Contains(content, "claude-opus-4-5-thinking") {
|
||||
t.Fatal("expected missing default alias claude-opus-4-5-thinking to be added")
|
||||
}
|
||||
if !strings.Contains(content, "claude-opus-4-6-thinking") {
|
||||
t.Fatal("expected missing default alias claude-opus-4-6-thinking to be added")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrateOAuthModelAlias_AddsDefaultIfNeitherExists(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dir := t.TempDir()
|
||||
configFile := filepath.Join(dir, "config.yaml")
|
||||
|
||||
content := `debug: true
|
||||
port: 8080
|
||||
`
|
||||
if err := os.WriteFile(configFile, []byte(content), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
migrated, err := MigrateOAuthModelAlias(configFile)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if !migrated {
|
||||
t.Fatal("expected migration to add default config")
|
||||
}
|
||||
|
||||
// Verify default antigravity config was added
|
||||
data, _ := os.ReadFile(configFile)
|
||||
content = string(data)
|
||||
if !strings.Contains(content, "oauth-model-alias:") {
|
||||
t.Fatal("expected oauth-model-alias to be added")
|
||||
}
|
||||
if !strings.Contains(content, "antigravity:") {
|
||||
t.Fatal("expected antigravity channel to be added")
|
||||
}
|
||||
if !strings.Contains(content, "rev19-uic3-1p") {
|
||||
t.Fatal("expected default antigravity aliases to include rev19-uic3-1p")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrateOAuthModelAlias_PreservesOtherConfig(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dir := t.TempDir()
|
||||
configFile := filepath.Join(dir, "config.yaml")
|
||||
|
||||
content := `debug: true
|
||||
port: 8080
|
||||
oauth-model-mappings:
|
||||
gemini-cli:
|
||||
- name: "test"
|
||||
alias: "t"
|
||||
api-keys:
|
||||
- "key1"
|
||||
- "key2"
|
||||
`
|
||||
if err := os.WriteFile(configFile, []byte(content), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
migrated, err := MigrateOAuthModelAlias(configFile)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if !migrated {
|
||||
t.Fatal("expected migration to occur")
|
||||
}
|
||||
|
||||
// Verify other config preserved
|
||||
data, _ := os.ReadFile(configFile)
|
||||
content = string(data)
|
||||
if !strings.Contains(content, "debug: true") {
|
||||
t.Fatal("expected debug field to be preserved")
|
||||
}
|
||||
if !strings.Contains(content, "port: 8080") {
|
||||
t.Fatal("expected port field to be preserved")
|
||||
}
|
||||
if !strings.Contains(content, "api-keys:") {
|
||||
t.Fatal("expected api-keys field to be preserved")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrateOAuthModelAlias_NonexistentFile(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
migrated, err := MigrateOAuthModelAlias("/nonexistent/path/config.yaml")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error for nonexistent file: %v", err)
|
||||
}
|
||||
if migrated {
|
||||
t.Fatal("expected no migration for nonexistent file")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrateOAuthModelAlias_EmptyFile(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dir := t.TempDir()
|
||||
configFile := filepath.Join(dir, "config.yaml")
|
||||
|
||||
if err := os.WriteFile(configFile, []byte(""), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
migrated, err := MigrateOAuthModelAlias(configFile)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if migrated {
|
||||
t.Fatal("expected no migration for empty file")
|
||||
}
|
||||
}
|
||||
@@ -20,9 +20,9 @@ type VertexCompatKey struct {
|
||||
// Prefix optionally namespaces model aliases for this credential (e.g., "teamA/vertex-pro").
|
||||
Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"`
|
||||
|
||||
// BaseURL is the base URL for the Vertex-compatible API endpoint.
|
||||
// BaseURL optionally overrides 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/..."
|
||||
// When empty, requests fall back to the default Vertex API base URL.
|
||||
BaseURL string `yaml:"base-url,omitempty" json:"base-url,omitempty"`
|
||||
|
||||
// ProxyURL optionally overrides the global proxy for this API key.
|
||||
@@ -34,6 +34,9 @@ type VertexCompatKey struct {
|
||||
|
||||
// Models defines the model configurations including aliases for routing.
|
||||
Models []VertexCompatModel `yaml:"models,omitempty" json:"models,omitempty"`
|
||||
|
||||
// ExcludedModels lists model IDs that should be excluded for this provider.
|
||||
ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"`
|
||||
}
|
||||
|
||||
func (k VertexCompatKey) GetAPIKey() string { return k.APIKey }
|
||||
@@ -68,12 +71,9 @@ func (cfg *Config) SanitizeVertexCompatKeys() {
|
||||
}
|
||||
entry.Prefix = normalizeModelPrefix(entry.Prefix)
|
||||
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)
|
||||
entry.ExcludedModels = NormalizeExcludedModels(entry.ExcludedModels)
|
||||
|
||||
// Sanitize models: remove entries without valid alias
|
||||
sanitizedModels := make([]VertexCompatModel, 0, len(entry.Models))
|
||||
|
||||
@@ -31,6 +31,7 @@ const (
|
||||
httpUserAgent = "CLIProxyAPI-management-updater"
|
||||
managementSyncMinInterval = 30 * time.Second
|
||||
updateCheckInterval = 3 * time.Hour
|
||||
maxAssetDownloadSize = 50 << 20 // 10 MB safety limit for management asset downloads
|
||||
)
|
||||
|
||||
// ManagementFileName exposes the control panel asset filename.
|
||||
@@ -88,6 +89,10 @@ func runAutoUpdater(ctx context.Context) {
|
||||
log.Debug("management asset auto-updater skipped: control panel disabled")
|
||||
return
|
||||
}
|
||||
if cfg.RemoteManagement.DisableAutoUpdatePanel {
|
||||
log.Debug("management asset auto-updater skipped: disable-auto-update-panel is enabled")
|
||||
return
|
||||
}
|
||||
|
||||
configPath, _ := schedulerConfigPath.Load().(string)
|
||||
staticDir := StaticDir(configPath)
|
||||
@@ -259,7 +264,8 @@ func EnsureLatestManagementHTML(ctx context.Context, staticDir string, proxyURL
|
||||
}
|
||||
|
||||
if remoteHash != "" && !strings.EqualFold(remoteHash, downloadedHash) {
|
||||
log.Warnf("remote digest mismatch for management asset: expected %s got %s", remoteHash, downloadedHash)
|
||||
log.Errorf("management asset digest mismatch: expected %s got %s — aborting update for safety", remoteHash, downloadedHash)
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
if err = atomicWriteFile(localPath, data); err != nil {
|
||||
@@ -282,6 +288,9 @@ func ensureFallbackManagementHTML(ctx context.Context, client *http.Client, loca
|
||||
return false
|
||||
}
|
||||
|
||||
log.Warnf("management asset downloaded from fallback URL without digest verification (hash=%s) — "+
|
||||
"enable verified GitHub updates by keeping disable-auto-update-panel set to false", downloadedHash)
|
||||
|
||||
if err = atomicWriteFile(localPath, data); err != nil {
|
||||
log.WithError(err).Warn("failed to persist fallback management control panel page")
|
||||
return false
|
||||
@@ -392,10 +401,13 @@ func downloadAsset(ctx context.Context, client *http.Client, downloadURL string)
|
||||
return nil, "", fmt.Errorf("unexpected download status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
|
||||
}
|
||||
|
||||
data, err := io.ReadAll(resp.Body)
|
||||
data, err := io.ReadAll(io.LimitReader(resp.Body, maxAssetDownloadSize+1))
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("read download body: %w", err)
|
||||
}
|
||||
if int64(len(data)) > maxAssetDownloadSize {
|
||||
return nil, "", fmt.Errorf("download exceeds maximum allowed size of %d bytes", maxAssetDownloadSize)
|
||||
}
|
||||
|
||||
sum := sha256.Sum256(data)
|
||||
return data, hex.EncodeToString(sum[:]), nil
|
||||
|
||||
@@ -1 +1 @@
|
||||
[{"type":"text","text":"You are Claude Code, Anthropic's official CLI for Claude.","cache_control":{"type":"ephemeral"}}]
|
||||
[{"type":"text","text":"You are a Claude agent, built on Anthropic's Claude Agent SDK.","cache_control":{"type":"ephemeral","ttl":"1h"}}]
|
||||
@@ -1,6 +1,7 @@
|
||||
package misc
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
@@ -24,3 +25,37 @@ func LogSavingCredentials(path string) {
|
||||
func LogCredentialSeparator() {
|
||||
log.Debug(credentialSeparator)
|
||||
}
|
||||
|
||||
// MergeMetadata serializes the source struct into a map and merges the provided metadata into it.
|
||||
func MergeMetadata(source any, metadata map[string]any) (map[string]any, error) {
|
||||
var data map[string]any
|
||||
|
||||
// Fast path: if source is already a map, just copy it to avoid mutation of original
|
||||
if srcMap, ok := source.(map[string]any); ok {
|
||||
data = make(map[string]any, len(srcMap)+len(metadata))
|
||||
for k, v := range srcMap {
|
||||
data[k] = v
|
||||
}
|
||||
} else {
|
||||
// Slow path: marshal to JSON and back to map to respect JSON tags
|
||||
temp, err := json.Marshal(source)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal source: %w", err)
|
||||
}
|
||||
if err := json.Unmarshal(temp, &data); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal to map: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Merge extra metadata
|
||||
if metadata != nil {
|
||||
if data == nil {
|
||||
data = make(map[string]any)
|
||||
}
|
||||
for k, v := range metadata {
|
||||
data[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
return data, nil
|
||||
}
|
||||
|
||||
@@ -4,10 +4,98 @@
|
||||
package misc
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"runtime"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
// GeminiCLIVersion is the version string reported in the User-Agent for upstream requests.
|
||||
GeminiCLIVersion = "0.31.0"
|
||||
|
||||
// GeminiCLIApiClientHeader is the value for the X-Goog-Api-Client header sent to the Gemini CLI upstream.
|
||||
GeminiCLIApiClientHeader = "google-genai-sdk/1.41.0 gl-node/v22.19.0"
|
||||
)
|
||||
|
||||
// geminiCLIOS maps Go runtime OS names to the Node.js-style platform strings used by Gemini CLI.
|
||||
func geminiCLIOS() string {
|
||||
switch runtime.GOOS {
|
||||
case "windows":
|
||||
return "win32"
|
||||
default:
|
||||
return runtime.GOOS
|
||||
}
|
||||
}
|
||||
|
||||
// geminiCLIArch maps Go runtime architecture names to the Node.js-style arch strings used by Gemini CLI.
|
||||
func geminiCLIArch() string {
|
||||
switch runtime.GOARCH {
|
||||
case "amd64":
|
||||
return "x64"
|
||||
case "386":
|
||||
return "x86"
|
||||
default:
|
||||
return runtime.GOARCH
|
||||
}
|
||||
}
|
||||
|
||||
// GeminiCLIUserAgent returns a User-Agent string that matches the Gemini CLI format.
|
||||
// The model parameter is included in the UA; pass "" or "unknown" when the model is not applicable.
|
||||
func GeminiCLIUserAgent(model string) string {
|
||||
if model == "" {
|
||||
model = "unknown"
|
||||
}
|
||||
return fmt.Sprintf("GeminiCLI/%s/%s (%s; %s)", GeminiCLIVersion, model, geminiCLIOS(), geminiCLIArch())
|
||||
}
|
||||
|
||||
// ScrubProxyAndFingerprintHeaders removes all headers that could reveal
|
||||
// proxy infrastructure, client identity, or browser fingerprints from an
|
||||
// outgoing request. This ensures requests to upstream services look like they
|
||||
// originate directly from a native client rather than a third-party client
|
||||
// behind a reverse proxy.
|
||||
func ScrubProxyAndFingerprintHeaders(req *http.Request) {
|
||||
if req == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// --- Proxy tracing headers ---
|
||||
req.Header.Del("X-Forwarded-For")
|
||||
req.Header.Del("X-Forwarded-Host")
|
||||
req.Header.Del("X-Forwarded-Proto")
|
||||
req.Header.Del("X-Forwarded-Port")
|
||||
req.Header.Del("X-Real-IP")
|
||||
req.Header.Del("Forwarded")
|
||||
req.Header.Del("Via")
|
||||
|
||||
// --- Client identity headers ---
|
||||
req.Header.Del("X-Title")
|
||||
req.Header.Del("X-Stainless-Lang")
|
||||
req.Header.Del("X-Stainless-Package-Version")
|
||||
req.Header.Del("X-Stainless-Os")
|
||||
req.Header.Del("X-Stainless-Arch")
|
||||
req.Header.Del("X-Stainless-Runtime")
|
||||
req.Header.Del("X-Stainless-Runtime-Version")
|
||||
req.Header.Del("Http-Referer")
|
||||
req.Header.Del("Referer")
|
||||
|
||||
// --- Browser / Chromium fingerprint headers ---
|
||||
// These are sent by Electron-based clients (e.g. CherryStudio) using the
|
||||
// Fetch API, but NOT by Node.js https module (which Antigravity uses).
|
||||
req.Header.Del("Sec-Ch-Ua")
|
||||
req.Header.Del("Sec-Ch-Ua-Mobile")
|
||||
req.Header.Del("Sec-Ch-Ua-Platform")
|
||||
req.Header.Del("Sec-Fetch-Mode")
|
||||
req.Header.Del("Sec-Fetch-Site")
|
||||
req.Header.Del("Sec-Fetch-Dest")
|
||||
req.Header.Del("Priority")
|
||||
|
||||
// --- Encoding negotiation ---
|
||||
// Antigravity (Node.js) sends "gzip, deflate, br" by default;
|
||||
// Electron-based clients may add "zstd" which is a fingerprint mismatch.
|
||||
req.Header.Del("Accept-Encoding")
|
||||
}
|
||||
|
||||
// EnsureHeader ensures that a header exists in the target header map by checking
|
||||
// multiple sources in order of priority: source headers, existing target headers,
|
||||
// and finally the default value. It only sets the header if it's not already present
|
||||
|
||||
@@ -30,6 +30,23 @@ type OAuthCallback struct {
|
||||
ErrorDescription string
|
||||
}
|
||||
|
||||
// AsyncPrompt runs a prompt function in a goroutine and returns channels for
|
||||
// the result. The returned channels are buffered (size 1) so the goroutine can
|
||||
// complete even if the caller abandons the channels.
|
||||
func AsyncPrompt(promptFn func(string) (string, error), message string) (<-chan string, <-chan error) {
|
||||
inputCh := make(chan string, 1)
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
input, err := promptFn(message)
|
||||
if err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
inputCh <- input
|
||||
}()
|
||||
return inputCh, errCh
|
||||
}
|
||||
|
||||
// ParseOAuthCallback extracts OAuth parameters from a callback URL.
|
||||
// It returns nil when the input is empty.
|
||||
func ParseOAuthCallback(input string) (*OAuthCallback, error) {
|
||||
|
||||
@@ -1,12 +1,186 @@
|
||||
// Package registry provides model definitions and lookup helpers for various AI providers.
|
||||
// Static model metadata is stored in model_definitions_static_data.go.
|
||||
// Static model metadata is loaded from the embedded models.json file and can be refreshed from network.
|
||||
package registry
|
||||
|
||||
import (
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// staticModelsJSON mirrors the top-level structure of models.json.
|
||||
type staticModelsJSON struct {
|
||||
Claude []*ModelInfo `json:"claude"`
|
||||
Gemini []*ModelInfo `json:"gemini"`
|
||||
Vertex []*ModelInfo `json:"vertex"`
|
||||
GeminiCLI []*ModelInfo `json:"gemini-cli"`
|
||||
AIStudio []*ModelInfo `json:"aistudio"`
|
||||
CodexFree []*ModelInfo `json:"codex-free"`
|
||||
CodexTeam []*ModelInfo `json:"codex-team"`
|
||||
CodexPlus []*ModelInfo `json:"codex-plus"`
|
||||
CodexPro []*ModelInfo `json:"codex-pro"`
|
||||
Qwen []*ModelInfo `json:"qwen"`
|
||||
IFlow []*ModelInfo `json:"iflow"`
|
||||
Kimi []*ModelInfo `json:"kimi"`
|
||||
Antigravity []*ModelInfo `json:"antigravity"`
|
||||
}
|
||||
|
||||
// GetClaudeModels returns the standard Claude model definitions.
|
||||
func GetClaudeModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().Claude)
|
||||
}
|
||||
|
||||
// GetGeminiModels returns the standard Gemini model definitions.
|
||||
func GetGeminiModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().Gemini)
|
||||
}
|
||||
|
||||
// GetGeminiVertexModels returns Gemini model definitions for Vertex AI.
|
||||
func GetGeminiVertexModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().Vertex)
|
||||
}
|
||||
|
||||
// GetGeminiCLIModels returns Gemini model definitions for the Gemini CLI.
|
||||
func GetGeminiCLIModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().GeminiCLI)
|
||||
}
|
||||
|
||||
// GetAIStudioModels returns model definitions for AI Studio.
|
||||
func GetAIStudioModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().AIStudio)
|
||||
}
|
||||
|
||||
// GetCodexFreeModels returns model definitions for the Codex free plan tier.
|
||||
func GetCodexFreeModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().CodexFree)
|
||||
}
|
||||
|
||||
// GetCodexTeamModels returns model definitions for the Codex team plan tier.
|
||||
func GetCodexTeamModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().CodexTeam)
|
||||
}
|
||||
|
||||
// GetCodexPlusModels returns model definitions for the Codex plus plan tier.
|
||||
func GetCodexPlusModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().CodexPlus)
|
||||
}
|
||||
|
||||
// GetCodexProModels returns model definitions for the Codex pro plan tier.
|
||||
func GetCodexProModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().CodexPro)
|
||||
}
|
||||
|
||||
// GetQwenModels returns the standard Qwen model definitions.
|
||||
func GetQwenModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().Qwen)
|
||||
}
|
||||
|
||||
// GetIFlowModels returns the standard iFlow model definitions.
|
||||
func GetIFlowModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().IFlow)
|
||||
}
|
||||
|
||||
// GetKimiModels returns the standard Kimi (Moonshot AI) model definitions.
|
||||
func GetKimiModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().Kimi)
|
||||
}
|
||||
|
||||
// GetAntigravityModels returns the standard Antigravity model definitions.
|
||||
func GetAntigravityModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().Antigravity)
|
||||
}
|
||||
|
||||
// GetCodeBuddyModels returns the available models for CodeBuddy (Tencent).
|
||||
// These models are served through the copilot.tencent.com API.
|
||||
func GetCodeBuddyModels() []*ModelInfo {
|
||||
now := int64(1748044800) // 2025-05-24
|
||||
return []*ModelInfo{
|
||||
{
|
||||
ID: "glm-5.0",
|
||||
Object: "model",
|
||||
Created: now,
|
||||
OwnedBy: "tencent",
|
||||
Type: "codebuddy",
|
||||
DisplayName: "GLM-5.0",
|
||||
Description: "GLM-5.0 via CodeBuddy",
|
||||
ContextLength: 128000,
|
||||
MaxCompletionTokens: 32768,
|
||||
SupportedEndpoints: []string{"/chat/completions"},
|
||||
},
|
||||
{
|
||||
ID: "glm-4.7",
|
||||
Object: "model",
|
||||
Created: now,
|
||||
OwnedBy: "tencent",
|
||||
Type: "codebuddy",
|
||||
DisplayName: "GLM-4.7",
|
||||
Description: "GLM-4.7 via CodeBuddy",
|
||||
ContextLength: 128000,
|
||||
MaxCompletionTokens: 32768,
|
||||
SupportedEndpoints: []string{"/chat/completions"},
|
||||
},
|
||||
{
|
||||
ID: "minimax-m2.5",
|
||||
Object: "model",
|
||||
Created: now,
|
||||
OwnedBy: "tencent",
|
||||
Type: "codebuddy",
|
||||
DisplayName: "MiniMax M2.5",
|
||||
Description: "MiniMax M2.5 via CodeBuddy",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 32768,
|
||||
SupportedEndpoints: []string{"/chat/completions"},
|
||||
},
|
||||
{
|
||||
ID: "kimi-k2.5",
|
||||
Object: "model",
|
||||
Created: now,
|
||||
OwnedBy: "tencent",
|
||||
Type: "codebuddy",
|
||||
DisplayName: "Kimi K2.5",
|
||||
Description: "Kimi K2.5 via CodeBuddy",
|
||||
ContextLength: 128000,
|
||||
MaxCompletionTokens: 32768,
|
||||
SupportedEndpoints: []string{"/chat/completions"},
|
||||
},
|
||||
{
|
||||
ID: "deepseek-v3-2-volc",
|
||||
Object: "model",
|
||||
Created: now,
|
||||
OwnedBy: "tencent",
|
||||
Type: "codebuddy",
|
||||
DisplayName: "DeepSeek V3.2 (Volc)",
|
||||
Description: "DeepSeek V3.2 via CodeBuddy (Volcano Engine)",
|
||||
ContextLength: 128000,
|
||||
MaxCompletionTokens: 32768,
|
||||
SupportedEndpoints: []string{"/chat/completions"},
|
||||
},
|
||||
{
|
||||
ID: "hunyuan-2.0-thinking",
|
||||
Object: "model",
|
||||
Created: now,
|
||||
OwnedBy: "tencent",
|
||||
Type: "codebuddy",
|
||||
DisplayName: "Hunyuan 2.0 Thinking",
|
||||
Description: "Tencent Hunyuan 2.0 Thinking via CodeBuddy",
|
||||
ContextLength: 128000,
|
||||
MaxCompletionTokens: 32768,
|
||||
Thinking: &ThinkingSupport{ZeroAllowed: true},
|
||||
SupportedEndpoints: []string{"/chat/completions"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// cloneModelInfos returns a shallow copy of the slice with each element deep-cloned.
|
||||
func cloneModelInfos(models []*ModelInfo) []*ModelInfo {
|
||||
if len(models) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*ModelInfo, len(models))
|
||||
for i, m := range models {
|
||||
out[i] = cloneModelInfo(m)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// GetStaticModelDefinitionsByChannel returns static model definitions for a given channel/provider.
|
||||
// It returns nil when the channel is unknown.
|
||||
//
|
||||
@@ -20,10 +194,8 @@ import (
|
||||
// - qwen
|
||||
// - iflow
|
||||
// - kimi
|
||||
// - kiro
|
||||
// - kilo
|
||||
// - github-copilot
|
||||
// - kiro
|
||||
// - amazonq
|
||||
// - antigravity (returns static overrides only)
|
||||
func GetStaticModelDefinitionsByChannel(channel string) []*ModelInfo {
|
||||
@@ -40,7 +212,7 @@ func GetStaticModelDefinitionsByChannel(channel string) []*ModelInfo {
|
||||
case "aistudio":
|
||||
return GetAIStudioModels()
|
||||
case "codex":
|
||||
return GetOpenAIModels()
|
||||
return GetCodexProModels()
|
||||
case "qwen":
|
||||
return GetQwenModels()
|
||||
case "iflow":
|
||||
@@ -56,33 +228,28 @@ func GetStaticModelDefinitionsByChannel(channel string) []*ModelInfo {
|
||||
case "amazonq":
|
||||
return GetAmazonQModels()
|
||||
case "antigravity":
|
||||
cfg := GetAntigravityModelConfig()
|
||||
if len(cfg) == 0 {
|
||||
return nil
|
||||
}
|
||||
models := make([]*ModelInfo, 0, len(cfg))
|
||||
for modelID, entry := range cfg {
|
||||
if modelID == "" || entry == nil {
|
||||
continue
|
||||
}
|
||||
models = append(models, &ModelInfo{
|
||||
ID: modelID,
|
||||
Object: "model",
|
||||
OwnedBy: "antigravity",
|
||||
Type: "antigravity",
|
||||
Thinking: entry.Thinking,
|
||||
MaxCompletionTokens: entry.MaxCompletionTokens,
|
||||
})
|
||||
}
|
||||
sort.Slice(models, func(i, j int) bool {
|
||||
return strings.ToLower(models[i].ID) < strings.ToLower(models[j].ID)
|
||||
})
|
||||
return models
|
||||
return GetAntigravityModels()
|
||||
case "codebuddy":
|
||||
return GetCodeBuddyModels()
|
||||
case "cursor":
|
||||
return GetCursorModels()
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// GetCursorModels returns the fallback Cursor model definitions.
|
||||
func GetCursorModels() []*ModelInfo {
|
||||
return []*ModelInfo{
|
||||
{ID: "composer-2", Object: "model", OwnedBy: "cursor", Type: "cursor", DisplayName: "Composer 2", ContextLength: 200000, MaxCompletionTokens: 64000, Thinking: &ThinkingSupport{Max: 50000, DynamicAllowed: true}},
|
||||
{ID: "claude-4-sonnet", Object: "model", OwnedBy: "cursor", Type: "cursor", DisplayName: "Claude 4 Sonnet", ContextLength: 200000, MaxCompletionTokens: 64000, Thinking: &ThinkingSupport{Max: 50000, DynamicAllowed: true}},
|
||||
{ID: "claude-3.5-sonnet", Object: "model", OwnedBy: "cursor", Type: "cursor", DisplayName: "Claude 3.5 Sonnet", ContextLength: 200000, MaxCompletionTokens: 8192},
|
||||
{ID: "gpt-4o", Object: "model", OwnedBy: "cursor", Type: "cursor", DisplayName: "GPT-4o", ContextLength: 128000, MaxCompletionTokens: 16384},
|
||||
{ID: "cursor-small", Object: "model", OwnedBy: "cursor", Type: "cursor", DisplayName: "Cursor Small", ContextLength: 200000, MaxCompletionTokens: 64000},
|
||||
{ID: "gemini-2.5-pro", Object: "model", OwnedBy: "cursor", Type: "cursor", DisplayName: "Gemini 2.5 Pro", ContextLength: 1000000, MaxCompletionTokens: 65536, Thinking: &ThinkingSupport{Max: 50000, DynamicAllowed: true}},
|
||||
}
|
||||
}
|
||||
|
||||
// LookupStaticModelInfo searches all static model definitions for a model by ID.
|
||||
// Returns nil if no matching model is found.
|
||||
func LookupStaticModelInfo(modelID string) *ModelInfo {
|
||||
@@ -90,38 +257,33 @@ func LookupStaticModelInfo(modelID string) *ModelInfo {
|
||||
return nil
|
||||
}
|
||||
|
||||
data := getModels()
|
||||
allModels := [][]*ModelInfo{
|
||||
GetClaudeModels(),
|
||||
GetGeminiModels(),
|
||||
GetGeminiVertexModels(),
|
||||
GetGeminiCLIModels(),
|
||||
GetAIStudioModels(),
|
||||
GetOpenAIModels(),
|
||||
GetQwenModels(),
|
||||
GetIFlowModels(),
|
||||
GetKimiModels(),
|
||||
data.Claude,
|
||||
data.Gemini,
|
||||
data.Vertex,
|
||||
data.GeminiCLI,
|
||||
data.AIStudio,
|
||||
data.CodexPro,
|
||||
data.Qwen,
|
||||
data.IFlow,
|
||||
data.Kimi,
|
||||
data.Antigravity,
|
||||
GetGitHubCopilotModels(),
|
||||
GetKiroModels(),
|
||||
GetKiloModels(),
|
||||
GetAmazonQModels(),
|
||||
GetCodeBuddyModels(),
|
||||
GetCursorModels(),
|
||||
}
|
||||
for _, models := range allModels {
|
||||
for _, m := range models {
|
||||
if m != nil && m.ID == modelID {
|
||||
return m
|
||||
return cloneModelInfo(m)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check Antigravity static config
|
||||
if cfg := GetAntigravityModelConfig()[modelID]; cfg != nil {
|
||||
return &ModelInfo{
|
||||
ID: modelID,
|
||||
Thinking: cfg.Thinking,
|
||||
MaxCompletionTokens: cfg.MaxCompletionTokens,
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -152,6 +314,7 @@ func GetGitHubCopilotModels() []*ModelInfo {
|
||||
Description: "OpenAI GPT-4.1 via GitHub Copilot",
|
||||
ContextLength: 128000,
|
||||
MaxCompletionTokens: 16384,
|
||||
SupportedEndpoints: []string{"/chat/completions", "/responses"},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -166,6 +329,7 @@ func GetGitHubCopilotModels() []*ModelInfo {
|
||||
Description: entry.Description,
|
||||
ContextLength: 128000,
|
||||
MaxCompletionTokens: 16384,
|
||||
SupportedEndpoints: []string{"/chat/completions", "/responses"},
|
||||
})
|
||||
}
|
||||
|
||||
@@ -300,6 +464,19 @@ func GetGitHubCopilotModels() []*ModelInfo {
|
||||
SupportedEndpoints: []string{"/responses"},
|
||||
Thinking: &ThinkingSupport{Levels: []string{"none", "low", "medium", "high", "xhigh"}},
|
||||
},
|
||||
{
|
||||
ID: "gpt-5.4",
|
||||
Object: "model",
|
||||
Created: now,
|
||||
OwnedBy: "github-copilot",
|
||||
Type: "github-copilot",
|
||||
DisplayName: "GPT-5.4",
|
||||
Description: "OpenAI GPT-5.4 via GitHub Copilot",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 32768,
|
||||
SupportedEndpoints: []string{"/responses"},
|
||||
Thinking: &ThinkingSupport{Levels: []string{"none", "low", "medium", "high", "xhigh"}},
|
||||
},
|
||||
{
|
||||
ID: "claude-haiku-4.5",
|
||||
Object: "model",
|
||||
@@ -696,6 +873,42 @@ func GetKiroModels() []*ModelInfo {
|
||||
MaxCompletionTokens: 64000,
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true},
|
||||
},
|
||||
{
|
||||
ID: "kiro-deepseek-3-2-agentic",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Kiro DeepSeek 3.2 (Agentic)",
|
||||
Description: "DeepSeek 3.2 optimized for coding agents (chunked writes)",
|
||||
ContextLength: 128000,
|
||||
MaxCompletionTokens: 32768,
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true},
|
||||
},
|
||||
{
|
||||
ID: "kiro-minimax-m2-1-agentic",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Kiro MiniMax M2.1 (Agentic)",
|
||||
Description: "MiniMax M2.1 optimized for coding agents (chunked writes)",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true},
|
||||
},
|
||||
{
|
||||
ID: "kiro-qwen3-coder-next-agentic",
|
||||
Object: "model",
|
||||
Created: 1732752000,
|
||||
OwnedBy: "aws",
|
||||
Type: "kiro",
|
||||
DisplayName: "Kiro Qwen3 Coder Next (Agentic)",
|
||||
Description: "Qwen3 Coder Next optimized for coding agents (chunked writes)",
|
||||
ContextLength: 128000,
|
||||
MaxCompletionTokens: 32768,
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user