mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Compare commits
379
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
96f4c115c9 | ||
|
|
327c322184 | ||
|
|
48b7109588 | ||
|
|
4f0c079e27 | ||
|
|
b0cfff9826 | ||
|
|
017bbcb901 | ||
|
|
9ecb340822 | ||
|
|
a0f41f95a0 | ||
|
|
b5efbe3f4c | ||
|
|
e6d191a83f | ||
|
|
c91bb6124a | ||
|
|
9688571a83 | ||
|
|
514ad0aa8e | ||
|
|
dc2efc6b9a | ||
|
|
1872d44b58 | ||
|
|
196b47757a | ||
|
|
a60a29549f | ||
|
|
2f3fed2fdb | ||
|
|
fe9dadfc7d | ||
|
|
41dcfec489 | ||
|
|
052e08e764 | ||
|
|
1f955af353 | ||
|
|
ccb027b515 | ||
|
|
f50b99b66c | ||
|
|
883a53fea0 | ||
|
|
b35a611538 | ||
|
|
1dca86f287 | ||
|
|
dd6cb410e6 | ||
|
|
27a9421e6a | ||
|
|
860782e6fb | ||
|
|
35232c3949 | ||
|
|
421da3f0d4 | ||
|
|
48970f61b4 | ||
|
|
f8984b8a93 | ||
|
|
6a8f921650 | ||
|
|
2cc459e24f | ||
|
|
475c0839c6 | ||
|
|
b886fb8edb | ||
|
|
8490a81867 | ||
|
|
652a7d74f9 | ||
|
|
52c123e6f4 | ||
|
|
43f88c99aa | ||
|
|
6073704fd7 | ||
|
|
6a69dd051a | ||
|
|
2840cdeced | ||
|
|
aa3a3ea787 | ||
|
|
bd48b2b2d2 | ||
|
|
9a62841fd1 | ||
|
|
4c00df2e01 | ||
|
|
b5298fdde8 | ||
|
|
5a350108c4 | ||
|
|
1df5c39311 | ||
|
|
b3c7038cf4 | ||
|
|
8616d4e036 | ||
|
|
93c3da03eb | ||
|
|
68c2f5ae6d | ||
|
|
ff6b85e30b | ||
|
|
79c18ec836 | ||
|
|
790723e040 | ||
|
|
1ff2bd6d95 | ||
|
|
f592b336a3 | ||
|
|
37f35b12ea | ||
|
|
e4b34446af | ||
|
|
90ae81e9f8 | ||
|
|
9375e288ba | ||
|
|
06f70a6232 | ||
|
|
6906062187 | ||
|
|
dc4874a2f1 | ||
|
|
11608c51f0 | ||
|
|
8aa7fb8c2d | ||
|
|
00ae7c38e2 | ||
|
|
51a10e42f7 | ||
|
|
6655f4ec0e | ||
|
|
be64b553f9 | ||
|
|
66f07efd26 | ||
|
|
89e9976950 | ||
|
|
99d86e41ac | ||
|
|
ec334fda2c | ||
|
|
35d2ef0cbd | ||
|
|
512aa8bc93 | ||
|
|
f45126c8ec | ||
|
|
82d2552360 | ||
|
|
f652d0ed38 | ||
|
|
6b40449c3f | ||
|
|
02088be2d1 | ||
|
|
14483c9254 | ||
|
|
fd7af1d462 | ||
|
|
8b95dc94f8 | ||
|
|
7626110d0b | ||
|
|
3c5ea297ff | ||
|
|
994048e6a4 | ||
|
|
c802c45e37 | ||
|
|
f4f8ff04db | ||
|
|
510b680fca | ||
|
|
afbe51abf8 | ||
|
|
52c2f4cb10 | ||
|
|
630639f747 | ||
|
|
11ffee5543 | ||
|
|
85810365a6 | ||
|
|
f0ebc183e1 | ||
|
|
ca2746f3bb | ||
|
|
b31b435091 | ||
|
|
1d126c833b | ||
|
|
9461196864 | ||
|
|
a2007db707 | ||
|
|
97bdde3135 | ||
|
|
090b2cc000 | ||
|
|
82b092ffc2 | ||
|
|
5979521309 | ||
|
|
1ca7b2a519 | ||
|
|
00f6f4e4a8 | ||
|
|
e2c35cf394 | ||
|
|
a3eb7ef302 | ||
|
|
2fc91663d2 | ||
|
|
fd80b08c90 | ||
|
|
453b796df0 | ||
|
|
e878ab710f | ||
|
|
13fcfc0f3b | ||
|
|
0416c9ede4 | ||
|
|
dc9216512d | ||
|
|
d0731be2f0 | ||
|
|
a2e1618700 | ||
|
|
31635849e0 | ||
|
|
d03274ad5f | ||
|
|
0c860e81e1 | ||
|
|
daf6cbb5d9 | ||
|
|
912e1c414c | ||
|
|
0a43efb9a4 | ||
|
|
f7f0a56f35 | ||
|
|
ae2735cdb6 | ||
|
|
6624947b62 | ||
|
|
afd069b169 | ||
|
|
6452aeffa1 | ||
|
|
4318a63bd8 | ||
|
|
ca08825935 | ||
|
|
c060aa0031 | ||
|
|
ecba49f711 | ||
|
|
17ce21c038 | ||
|
|
5e244e7382 | ||
|
|
b5b5f1f94b | ||
|
|
5b5f315b87 | ||
|
|
c4c2e66070 | ||
|
|
d8d5e1dd35 | ||
|
|
59a00631ed | ||
|
|
b2d6b954a1 | ||
|
|
1947add842 | ||
|
|
6b71d75bde | ||
|
|
f53252e292 | ||
|
|
61e7d31605 | ||
|
|
033047b7d4 | ||
|
|
d68a68fd0d | ||
|
|
31a6e74770 | ||
|
|
7c0a2a556c | ||
|
|
0dd71286f2 | ||
|
|
b350e079f8 | ||
|
|
24872fda7c | ||
|
|
5ad7cf4fb9 | ||
|
|
60b9bf7558 | ||
|
|
d6e8b44bcb | ||
|
|
fa3f526785 | ||
|
|
7d5615998c | ||
|
|
d7e4bba0e7 | ||
|
|
d3d0f653e5 | ||
|
|
9e0e146984 | ||
|
|
fc4465f78b | ||
|
|
0eaa7c3c84 | ||
|
|
5e4968492f | ||
|
|
a81547f244 | ||
|
|
bdb9a91db0 | ||
|
|
3b336b2eca | ||
|
|
9d7c636db9 | ||
|
|
c19204289d | ||
|
|
8f8358cc67 | ||
|
|
20a94fbb56 | ||
|
|
da25b18db7 | ||
|
|
02d9901b49 | ||
|
|
13be6ca27c | ||
|
|
e26abaef7d | ||
|
|
583398d186 | ||
|
|
5cc6ca6f5b | ||
|
|
b12a187d4f | ||
|
|
974a819ba1 | ||
|
|
d27fbb3b7d | ||
|
|
ea34a2356c | ||
|
|
fa9f104a93 | ||
|
|
51f73840a1 | ||
|
|
2b12de14ee | ||
|
|
7865051fdf | ||
|
|
4b0adb6863 | ||
|
|
ba5737fe14 | ||
|
|
7b13cb410d | ||
|
|
bb1f40e35f | ||
|
|
7be1628a74 | ||
|
|
3eb4bae9ab | ||
|
|
1f5cfe8973 | ||
|
|
5fc4a06af6 | ||
|
|
a1eaad6033 | ||
|
|
1af3269adb | ||
|
|
1ca54bcc25 | ||
|
|
7ba6efe805 | ||
|
|
4d667f1b08 | ||
|
|
5c38d372c1 | ||
|
|
9e4010695b | ||
|
|
4daea093f9 | ||
|
|
1d640c40e5 | ||
|
|
b0e32fe5be | ||
|
|
65f086d2d0 | ||
|
|
de6388aeca | ||
|
|
55108ee877 | ||
|
|
231e72fdde | ||
|
|
a079a6596a | ||
|
|
ed2360d7ea | ||
|
|
b717ac7f0b | ||
|
|
c619ff8474 | ||
|
|
2448a38b55 | ||
|
|
5fd3461142 | ||
|
|
c1319917cf | ||
|
|
a5015e1fc4 | ||
|
|
738ef1c73c | ||
|
|
5af29a3b93 | ||
|
|
62a8919a1c | ||
|
|
1ea00258d0 | ||
|
|
023b38354f | ||
|
|
a765b3c176 | ||
|
|
ca47fa352c | ||
|
|
680a992fc9 | ||
|
|
f2e55bf4d7 | ||
|
|
c96d36ec95 | ||
|
|
e4f97a53ec | ||
|
|
9e175bd491 | ||
|
|
e7d348868c | ||
|
|
c0b23ff791 | ||
|
|
322cc5e5f1 | ||
|
|
23dc073e65 | ||
|
|
a5d8db2445 | ||
|
|
de28eea11f | ||
|
|
66843cc95c | ||
|
|
f867e3e2ec | ||
|
|
1716e99151 | ||
|
|
483cd56924 | ||
|
|
b6c6f1add5 | ||
|
|
3442a6a53e | ||
|
|
9d9e90960e | ||
|
|
fcd07ba0af | ||
|
|
405727805a | ||
|
|
a46493724b | ||
|
|
d8eeb0989b | ||
|
|
a266e50e8e | ||
|
|
5295bd8227 | ||
|
|
7fbc59eb69 | ||
|
|
fa79b1ebc3 | ||
|
|
0c0df031c5 | ||
|
|
85ed17b313 | ||
|
|
b754ff6f1b | ||
|
|
e35123b8ac | ||
|
|
7f1f85cb76 | ||
|
|
935db68517 | ||
|
|
1c0a69c0ce | ||
|
|
27f4398b50 | ||
|
|
e47255715d | ||
|
|
d0ed0eacac | ||
|
|
ece1ecccca | ||
|
|
b13200d7a1 | ||
|
|
386f321d3c | ||
|
|
60a64568bf | ||
|
|
de98bdff1e | ||
|
|
992c4f4b89 | ||
|
|
7c700729c2 | ||
|
|
87957a608b | ||
|
|
0d02000592 | ||
|
|
8d88b66668 | ||
|
|
763443abf2 | ||
|
|
7f18f3e9af | ||
|
|
b86849e442 | ||
|
|
319299519b | ||
|
|
23a1d38187 | ||
|
|
1582347204 | ||
|
|
d1556a7a96 | ||
|
|
d2416a9b68 | ||
|
|
fb58dfbd93 | ||
|
|
b8d52fad34 | ||
|
|
d2e319b2a1 | ||
|
|
0892ef3a9a | ||
|
|
dd292e3b00 | ||
|
|
95134201b2 | ||
|
|
acdf3c54da | ||
|
|
fc96132e72 | ||
|
|
8f6bbb59c7 | ||
|
|
fbb9006ade | ||
|
|
e009ea3036 | ||
|
|
63c079d873 | ||
|
|
e509937382 | ||
|
|
5090ffe05e | ||
|
|
bbdcfbac0a | ||
|
|
7403a01177 | ||
|
|
53b4bbe736 | ||
|
|
8bb48622a3 | ||
|
|
f28adb6ddb | ||
|
|
19794bc46a | ||
|
|
29eead917f | ||
|
|
fabdfb8a5b | ||
|
|
1411b7eb7e | ||
|
|
ef51690b5a | ||
|
|
1bfe0d3728 | ||
|
|
783a0a9830 | ||
|
|
cd2a4357c2 | ||
|
|
f3a2dcabba | ||
|
|
7bf3d0b585 | ||
|
|
cb2bb6084c | ||
|
|
00591cbe9a | ||
|
|
ba26e543f9 | ||
|
|
f6eeed3d5a | ||
|
|
3b0bb60ea3 | ||
|
|
26c09b7de6 | ||
|
|
4d5a4ab0e5 | ||
|
|
50f79e11fc | ||
|
|
6b09c74e35 | ||
|
|
b252821c59 | ||
|
|
68aa26477f | ||
|
|
7e530c9d4e | ||
|
|
1a9d49e16f | ||
|
|
aea725f2ea | ||
|
|
74431ffad8 | ||
|
|
0c9e83b59c | ||
|
|
09f88865e0 | ||
|
|
1b45d21328 | ||
|
|
d5cea617a5 | ||
|
|
4275047d16 | ||
|
|
4dbcdce43c | ||
|
|
2d37088bd9 | ||
|
|
a8a1a3a958 | ||
|
|
4ccdead2a2 | ||
|
|
f7280213f7 | ||
|
|
db8692d679 | ||
|
|
37238c0985 | ||
|
|
7b282c4f1d | ||
|
|
d054ee9af3 | ||
|
|
9612a70c77 | ||
|
|
24f4736f64 | ||
|
|
9d5c8d961c | ||
|
|
7b4de8b6ae | ||
|
|
bcc73f8d40 | ||
|
|
4805f069e4 | ||
|
|
cfd2fc779b | ||
|
|
0920e06a42 | ||
|
|
3fdd54ca12 | ||
|
|
cf10d6d01d | ||
|
|
a9c7c6e8bf | ||
|
|
c12d0131cb | ||
|
|
0f4d8acaa8 | ||
|
|
94225b5ed3 | ||
|
|
a9ff663387 | ||
|
|
57b7dbdc85 | ||
|
|
f79b8bf96c | ||
|
|
e74c690ac4 | ||
|
|
0952ce3410 | ||
|
|
d099c06945 | ||
|
|
aba34524f1 | ||
|
|
1e231ee2ae | ||
|
|
4f6f206458 | ||
|
|
6582260bc8 | ||
|
|
0e1329b520 | ||
|
|
29bce46a2a | ||
|
|
7dace2fecb | ||
|
|
a985de062c | ||
|
|
7511e4de59 | ||
|
|
b18e4ce455 | ||
|
|
69d49744cf | ||
|
|
5d781ee6b7 | ||
|
|
c24da29f6b | ||
|
|
baa39f39e8 | ||
|
|
9fedf899e2 | ||
|
|
8e23e13158 | ||
|
|
f744375a4c | ||
|
|
d4e8ddb0c5 | ||
|
|
ac770d98fb | ||
|
|
6f9fde093f | ||
|
|
277aa1411d | ||
|
|
89ba561a6f |
@@ -0,0 +1 @@
|
||||
__pycache__/
|
||||
@@ -0,0 +1,30 @@
|
||||
# Release matrix
|
||||
|
||||
The release workflow builds the frontend once, then runs each configured Go target on its own Linux runner. `CGO_ENABLED=0` permits cross-compilation. The target matrix is read from `.goreleaser.yaml`, including its exclusions; simple releases select Linux amd64 only.
|
||||
|
||||
Each build uses GoReleaser OSS in snapshot mode with the selected release version and one target. Archive naming, bundled files, Go flags and release templates remain in the existing GoReleaser configurations. Every archive is accompanied by its source commit, target, version and SHA256. The publishing job verifies the complete matrix before building images or publishing. It uses GoReleaser's `extra_files` support to publish existing archives and checksums, with builds disabled. No Pro license is needed.
|
||||
|
||||
Go caches are isolated by target and refreshed on each source commit, with fallback to the preceding target cache. Save uses the original restore key, even if a build hook changes `go.sum`. Matrix jobs upload uniquely named artifacts. The publishing job extracts only the regular Linux binary from each verified archive and restores its executable permission before constructing Docker contexts. QEMU remains limited to runtime-image instructions. DockerHub images are omitted when its credentials are absent; GHCR is always retained. Simple mode still publishes only the amd64 GHCR image and the simple release description.
|
||||
|
||||
All build jobs use the commit resolved by `prepare`, including a manual release's selected tag. Helper scripts come from the workflow revision and are passed as a run-local artifact, so older application tags do not need to contain the new scripts. The workflow serializes release runs to prevent simultaneous updates to moving image tags.
|
||||
|
||||
## Validate without publication
|
||||
|
||||
From a branch containing this workflow:
|
||||
|
||||
```bash
|
||||
gh workflow run release.yml --ref <branch> \
|
||||
-f tag=<branch> -f dry_run=true -f simple_release=false
|
||||
```
|
||||
|
||||
A dry run builds all selected archives and both runtime images, verifies artifact provenance and produces the final checksum file. It exports images locally as OCI archives instead of pushing them. It skips registry logins, GitHub Release publication, DockerHub description updates, Telegram notifications and VERSION synchronization. Test the simple path separately with `simple_release=true`.
|
||||
|
||||
Dry-run artifacts are available in the Actions run, including `release-dry-run-report`. Compare job start/end times, GoReleaser's build duration and cache restore results. Do not present an initial cold-cache run as a warmed-cache benchmark; publishing network time is not measured by dry runs.
|
||||
|
||||
Helper checks:
|
||||
|
||||
```bash
|
||||
python -m pip install -r .github/release-tools/requirements-release.txt
|
||||
python -m unittest discover -s .github/release-tools -p 'test_release_matrix.py'
|
||||
bash -n .github/release-tools/release-images.sh
|
||||
```
|
||||
@@ -0,0 +1,38 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
: "${RELEASE_VERSION:?}" "${RELEASE_SHA:?}" "${GITHUB_REPOSITORY:?}" "${RUNNER_TEMP:?}"
|
||||
owner=${GITHUB_REPOSITORY%%/*}
|
||||
registries=("ghcr.io/${owner,,}/sub2api")
|
||||
if [[ ${SIMPLE_RELEASE:-false} != true && ${DOCKERHUB_USERNAME:-skip} != skip ]]; then
|
||||
registries+=("${DOCKERHUB_USERNAME}/sub2api")
|
||||
fi
|
||||
arches=(amd64 arm64)
|
||||
if [[ ${SIMPLE_RELEASE:-false} == true ]]; then arches=(amd64); fi
|
||||
for arch in "${arches[@]}"; do
|
||||
args=(--platform "linux/$arch" --file ".release-context/$arch/Dockerfile"
|
||||
--label "org.opencontainers.image.version=$RELEASE_VERSION"
|
||||
--label "org.opencontainers.image.revision=$RELEASE_SHA"
|
||||
--label "org.opencontainers.image.source=https://github.com/$GITHUB_REPOSITORY")
|
||||
for registry in "${registries[@]}"; do
|
||||
args+=(--tag "$registry:$RELEASE_VERSION-$arch")
|
||||
if [[ ${SIMPLE_RELEASE:-false} == true ]]; then
|
||||
args+=(--tag "$registry:$RELEASE_VERSION" --tag "$registry:latest")
|
||||
fi
|
||||
done
|
||||
if [[ ${DRY_RUN:-false} == true ]]; then
|
||||
args+=(--output "type=oci,dest=$RUNNER_TEMP/sub2api-$arch.oci.tar")
|
||||
else
|
||||
args+=(--push)
|
||||
fi
|
||||
docker buildx build "${args[@]}" ".release-context/$arch"
|
||||
done
|
||||
if [[ ${DRY_RUN:-false} != true && ${SIMPLE_RELEASE:-false} != true ]]; then
|
||||
major=${RELEASE_VERSION%%.*}
|
||||
minor=${RELEASE_VERSION#*.}; minor=${minor%%.*}
|
||||
for registry in "${registries[@]}"; do
|
||||
docker buildx imagetools create \
|
||||
--tag "$registry:$RELEASE_VERSION" --tag "$registry:latest" \
|
||||
--tag "$registry:$major.$minor" --tag "$registry:$major" \
|
||||
"$registry:$RELEASE_VERSION-amd64" "$registry:$RELEASE_VERSION-arm64"
|
||||
done
|
||||
fi
|
||||
@@ -0,0 +1,189 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Split GoReleaser builds across runners without requiring GoReleaser Pro."""
|
||||
import argparse
|
||||
import hashlib
|
||||
import itertools
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import tarfile
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import yaml
|
||||
|
||||
FULL_CONFIG = Path('.goreleaser.yaml')
|
||||
SIMPLE_CONFIG = Path('.goreleaser.simple.yaml')
|
||||
VERSION_FILE = Path('backend/cmd/server/VERSION')
|
||||
VERSION_RE = re.compile(r'\d+\.\d+\.\d+(?:-[0-9A-Za-z.-]+)?')
|
||||
|
||||
|
||||
def config(simple=False):
|
||||
return yaml.safe_load((SIMPLE_CONFIG if simple else FULL_CONFIG).read_text())
|
||||
|
||||
|
||||
def targets(simple=False):
|
||||
build = config()['builds'][0]
|
||||
result = []
|
||||
for goos, goarch in itertools.product(build['goos'], build['goarch']):
|
||||
item = {'goos': goos, 'goarch': goarch}
|
||||
if any(all(item.get(k) == v for k, v in rule.items()) for rule in build.get('ignore', [])):
|
||||
continue
|
||||
if not simple or item == {'goos': 'linux', 'goarch': 'amd64'}:
|
||||
result.append(item)
|
||||
if not result:
|
||||
raise ValueError('empty release target matrix')
|
||||
return result
|
||||
|
||||
|
||||
def archive_name(version, target):
|
||||
if not VERSION_RE.fullmatch(version) or target not in targets():
|
||||
raise ValueError('invalid release version or target')
|
||||
suffix = 'zip' if target['goos'] == 'windows' else 'tar.gz'
|
||||
return f"sub2api_{version}_{target['goos']}_{target['goarch']}.{suffix}"
|
||||
|
||||
|
||||
def sha256(path):
|
||||
with path.open('rb') as stream:
|
||||
return hashlib.file_digest(stream, 'sha256').hexdigest()
|
||||
|
||||
|
||||
def plan(args):
|
||||
sha = subprocess.check_output(['git', 'rev-parse', 'HEAD'], text=True).strip()
|
||||
if args.dry_run:
|
||||
version = VERSION_FILE.read_text().strip()
|
||||
tag = 'v' + version
|
||||
else:
|
||||
tag = args.ref
|
||||
version = tag.removeprefix('v')
|
||||
if not tag.startswith('v') or not VERSION_RE.fullmatch(version):
|
||||
raise ValueError('publishing requires a v-prefixed release version tag')
|
||||
tagged_sha = subprocess.check_output(['git', 'rev-parse', '--verify', f'refs/tags/{tag}^{{commit}}'], text=True).strip()
|
||||
if sha != tagged_sha:
|
||||
raise ValueError('checkout does not match the selected release tag')
|
||||
if not VERSION_RE.fullmatch(version):
|
||||
raise ValueError('invalid VERSION')
|
||||
VERSION_FILE.write_text(version + '\n')
|
||||
result = {'sha': sha, 'tag': tag, 'version': version,
|
||||
'owner_lower': os.environ.get('GITHUB_REPOSITORY_OWNER', '').lower(),
|
||||
'simple': str(args.simple).lower(), 'dry_run': str(args.dry_run).lower(),
|
||||
'date': datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ'),
|
||||
'matrix': json.dumps({'include': targets(args.simple)}, separators=(',', ':'))}
|
||||
with Path(os.environ['GITHUB_OUTPUT']).open('a') as output:
|
||||
for key, value in result.items():
|
||||
output.write(f'{key}={value}\n')
|
||||
|
||||
|
||||
def generate_config(args):
|
||||
data = config(args.simple if args.mode == 'publish' else False)
|
||||
data['snapshot'] = {'version_template': '{{ .Env.RELEASE_VERSION }}'}
|
||||
data['dockers'] = []
|
||||
data['docker_manifests'] = []
|
||||
if args.mode == 'build':
|
||||
target = {'goos': args.goos, 'goarch': args.goarch}
|
||||
if target not in targets():
|
||||
raise ValueError('unsupported build target')
|
||||
for build in data['builds']:
|
||||
build['goos'], build['goarch'], build['ignore'] = [args.goos], [args.goarch], []
|
||||
build['ldflags'] = [re.sub(r'{{\s*\.Date\s*}}', '{{ .Env.RELEASE_DATE }}', flag)
|
||||
for flag in build.get('ldflags', [])]
|
||||
else:
|
||||
# Artifacts are supplied through the OSS extra_files mechanism. No build
|
||||
# is repeated on the publishing runner, and release templates stay intact.
|
||||
data['before'] = {'hooks': []}
|
||||
data['builds'] = [{'id': 'sub2api', 'skip': True}]
|
||||
data['archives'] = []
|
||||
extra = [{'glob': 'release-input/sub2api_*.tar.gz'}, {'glob': 'release-input/sub2api_*.zip'}]
|
||||
if args.simple:
|
||||
data['checksum'] = {'disable': True}
|
||||
else:
|
||||
data['release']['extra_files'] = extra
|
||||
data['checksum'] = {'name_template': 'checksums.txt', 'algorithm': 'sha256', 'extra_files': extra}
|
||||
Path(args.output).write_text(yaml.safe_dump(data, sort_keys=False, allow_unicode=True))
|
||||
|
||||
|
||||
def collect(args):
|
||||
target = {'goos': args.goos, 'goarch': args.goarch}
|
||||
name = archive_name(args.version, target)
|
||||
source = Path('dist') / name
|
||||
checksums = {line.split()[1].lstrip('*'): line.split()[0] for line in Path('dist/checksums.txt').read_text().splitlines()}
|
||||
digest = sha256(source)
|
||||
if checksums.get(name) != digest:
|
||||
raise ValueError('archive does not match the build checksum')
|
||||
output = Path(args.output)
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(source, output / name)
|
||||
manifest = {'sha': args.sha, 'version': args.version, 'target': target, 'archive': name, 'sha256': digest}
|
||||
(output / f"manifest-{args.goos}-{args.goarch}.json").write_text(json.dumps(manifest) + '\n')
|
||||
|
||||
|
||||
def verify(args):
|
||||
directory = Path(args.input)
|
||||
expected = set()
|
||||
for target in targets(args.simple):
|
||||
name = archive_name(args.version, target)
|
||||
manifest_name = f"manifest-{target['goos']}-{target['goarch']}.json"
|
||||
expected.update((name, manifest_name))
|
||||
manifest = json.loads((directory / manifest_name).read_text())
|
||||
if manifest != {'sha': args.sha, 'version': args.version, 'target': target,
|
||||
'archive': name, 'sha256': sha256(directory / name)}:
|
||||
raise ValueError(f'build provenance or checksum mismatch: {name}')
|
||||
if {p.name for p in directory.iterdir()} != expected:
|
||||
raise ValueError('missing or unexpected release artifacts')
|
||||
|
||||
|
||||
def contexts(args):
|
||||
verify(args)
|
||||
for target in targets(args.simple):
|
||||
if target['goos'] != 'linux':
|
||||
continue
|
||||
dest = Path(args.output) / target['goarch']
|
||||
dest.mkdir(parents=True, exist_ok=True)
|
||||
with tarfile.open(Path(args.input) / archive_name(args.version, target), 'r:gz') as archive:
|
||||
members = [member for member in archive.getmembers() if member.name in ('sub2api', './sub2api')]
|
||||
if len(members) != 1 or not members[0].isfile():
|
||||
raise ValueError('archive must contain one regular sub2api binary')
|
||||
with archive.extractfile(members[0]) as source, (dest / 'sub2api').open('wb') as output:
|
||||
shutil.copyfileobj(source, output)
|
||||
(dest / 'sub2api').chmod(0o755)
|
||||
shutil.copy2('Dockerfile.goreleaser', dest / 'Dockerfile')
|
||||
(dest / 'deploy').mkdir(exist_ok=True)
|
||||
shutil.copy2('deploy/docker-entrypoint.sh', dest / 'deploy/docker-entrypoint.sh')
|
||||
shutil.copytree('backend/resources', dest / 'backend/resources', dirs_exist_ok=True)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
commands = parser.add_subparsers(dest='command', required=True)
|
||||
p = commands.add_parser('plan')
|
||||
p.add_argument('--ref', required=True)
|
||||
p.add_argument('--simple', action='store_true')
|
||||
p.add_argument('--dry-run', action='store_true')
|
||||
p.set_defaults(run=plan)
|
||||
p = commands.add_parser('config')
|
||||
p.add_argument('mode', choices=['build', 'publish'])
|
||||
p.add_argument('--goos')
|
||||
p.add_argument('--goarch')
|
||||
p.add_argument('--simple', action='store_true')
|
||||
p.add_argument('--output', required=True)
|
||||
p.set_defaults(run=generate_config)
|
||||
p = commands.add_parser('collect')
|
||||
for arg in ('version', 'sha', 'goos', 'goarch', 'output'):
|
||||
p.add_argument('--' + arg, required=True)
|
||||
p.set_defaults(run=collect)
|
||||
for command, handler in [('verify', verify), ('contexts', contexts)]:
|
||||
p = commands.add_parser(command)
|
||||
for arg in ('version', 'sha', 'input'):
|
||||
p.add_argument('--' + arg, required=True)
|
||||
p.add_argument('--simple', action='store_true')
|
||||
if command == 'contexts':
|
||||
p.add_argument('--output', required=True)
|
||||
p.set_defaults(run=handler)
|
||||
args = parser.parse_args()
|
||||
args.run(args)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1 @@
|
||||
PyYAML==6.0.3
|
||||
@@ -0,0 +1,191 @@
|
||||
import argparse
|
||||
import hashlib
|
||||
import importlib.util
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import subprocess
|
||||
import tarfile
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import yaml
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
spec = importlib.util.spec_from_file_location('release_matrix', Path(__file__).with_name('release_matrix.py'))
|
||||
release = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(release)
|
||||
|
||||
|
||||
class ReleaseMatrixTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.previous = Path.cwd()
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.temp.cleanup)
|
||||
os.chdir(self.temp.name)
|
||||
self.addCleanup(os.chdir, self.previous)
|
||||
for name in ('.goreleaser.yaml', '.goreleaser.simple.yaml'):
|
||||
shutil.copyfile(ROOT / name, name)
|
||||
Path('backend/cmd/server').mkdir(parents=True)
|
||||
release.VERSION_FILE.write_text('9.8.7\n')
|
||||
|
||||
def fixture_artifacts(self, simple=False):
|
||||
directory = Path('release-input')
|
||||
directory.mkdir()
|
||||
for target in release.targets(simple):
|
||||
name = release.archive_name('9.8.7', target)
|
||||
archive = directory / name
|
||||
if target['goos'] == 'linux':
|
||||
with tarfile.open(archive, 'w:gz') as out:
|
||||
info = tarfile.TarInfo('sub2api')
|
||||
info.size = 7
|
||||
info.mode = 0o755
|
||||
out.addfile(info, io.BytesIO(b'fixture'))
|
||||
else:
|
||||
archive.write_bytes(b'fixture archive')
|
||||
metadata = {'version': '9.8.7', 'sha': 'a' * 40, 'target': target,
|
||||
'archive': name, 'sha256': release.sha256(archive)}
|
||||
(directory / f"manifest-{target['goos']}-{target['goarch']}.json").write_text(json.dumps(metadata))
|
||||
return argparse.Namespace(input='release-input', version='9.8.7', sha='a' * 40, simple=simple, output='contexts')
|
||||
|
||||
def test_full_and_simple_matrix_match_existing_targets(self):
|
||||
full = release.targets()
|
||||
self.assertEqual(len(full), 5)
|
||||
self.assertNotIn({'goos': 'windows', 'goarch': 'arm64'}, full)
|
||||
self.assertEqual(release.targets(True), [{'goos': 'linux', 'goarch': 'amd64'}])
|
||||
|
||||
def test_leaf_keeps_packaging_and_selects_only_one_target(self):
|
||||
original = release.config()
|
||||
release.generate_config(argparse.Namespace(mode='build', simple=False, goos='darwin', goarch='arm64', output='leaf.yaml'))
|
||||
leaf = yaml.safe_load(Path('leaf.yaml').read_text())
|
||||
self.assertEqual(leaf['builds'][0]['goos'], ['darwin'])
|
||||
self.assertEqual(leaf['builds'][0]['goarch'], ['arm64'])
|
||||
self.assertEqual(leaf['builds'][0]['ignore'], [])
|
||||
self.assertEqual(leaf['archives'], original['archives'])
|
||||
self.assertEqual(leaf['release'], original['release'])
|
||||
self.assertFalse(leaf['dockers'])
|
||||
self.assertIn('{{ .Env.RELEASE_DATE }}', '\n'.join(leaf['builds'][0]['ldflags']))
|
||||
|
||||
def test_publication_config_has_no_compilation_or_docker_work(self):
|
||||
for simple in (False, True):
|
||||
with self.subTest(simple=simple):
|
||||
original = release.config(simple)
|
||||
release.generate_config(argparse.Namespace(mode='publish', simple=simple, output='publisher.yaml'))
|
||||
data = yaml.safe_load(Path('publisher.yaml').read_text())
|
||||
self.assertTrue(data['builds'][0]['skip'])
|
||||
self.assertFalse(data['archives'])
|
||||
self.assertFalse(data['dockers'])
|
||||
self.assertEqual(data['release']['header'], original['release']['header'])
|
||||
self.assertEqual(data['release']['footer'], original['release']['footer'])
|
||||
if simple:
|
||||
self.assertTrue(data['checksum']['disable'])
|
||||
self.assertTrue(data['release']['skip_upload'])
|
||||
else:
|
||||
self.assertEqual(data['checksum']['extra_files'], data['release']['extra_files'])
|
||||
|
||||
def test_collect_and_verify_hash_and_source_binding(self):
|
||||
args = self.fixture_artifacts()
|
||||
release.verify(args)
|
||||
file = next(Path(args.input).glob('*.tar.gz'))
|
||||
file.write_bytes(b'corrupted')
|
||||
with self.assertRaisesRegex(ValueError, 'checksum mismatch'):
|
||||
release.verify(args)
|
||||
|
||||
def test_missing_extra_and_wrong_commit_artifacts_are_rejected(self):
|
||||
args = self.fixture_artifacts(True)
|
||||
args.sha = 'b' * 40
|
||||
with self.assertRaises(ValueError):
|
||||
release.verify(args)
|
||||
args.sha = 'a' * 40
|
||||
Path('release-input/unexpected').write_text('not an asset')
|
||||
with self.assertRaises(ValueError):
|
||||
release.verify(args)
|
||||
Path('release-input/unexpected').unlink()
|
||||
next(Path('release-input').glob('*.tar.gz')).unlink()
|
||||
with self.assertRaises(FileNotFoundError):
|
||||
release.verify(args)
|
||||
|
||||
def test_linux_context_preserves_binary_executable_mode(self):
|
||||
args = self.fixture_artifacts()
|
||||
Path('Dockerfile.goreleaser').write_text('FROM scratch\nCOPY sub2api /sub2api\n')
|
||||
Path('deploy').mkdir()
|
||||
Path('deploy/docker-entrypoint.sh').write_text('#!/bin/sh\nexec /app/sub2api\n')
|
||||
Path('backend/resources').mkdir()
|
||||
Path('backend/resources/data').write_text('fixture')
|
||||
release.contexts(args)
|
||||
for arch in ('amd64', 'arm64'):
|
||||
binary = Path('contexts') / arch / 'sub2api'
|
||||
self.assertEqual(binary.read_bytes(), b'fixture')
|
||||
self.assertEqual(binary.stat().st_mode & 0o777, 0o755)
|
||||
|
||||
def test_plan_requires_a_tag_for_publication(self):
|
||||
args = argparse.Namespace(ref='main', dry_run=False, simple=False)
|
||||
with patch.object(subprocess, 'check_output', return_value='a' * 40 + '\n'):
|
||||
with self.assertRaisesRegex(ValueError, 'version tag'):
|
||||
release.plan(args)
|
||||
args.ref = 'v9.8.7'
|
||||
with patch.object(subprocess, 'check_output', side_effect=['a' * 40 + '\n', 'b' * 40 + '\n']):
|
||||
with self.assertRaisesRegex(ValueError, 'does not match'):
|
||||
release.plan(args)
|
||||
|
||||
def test_dry_run_plan_resolves_matrix_without_a_new_tag(self):
|
||||
with patch.dict(os.environ, {'GITHUB_OUTPUT': 'outputs', 'GITHUB_REPOSITORY_OWNER': 'ExampleOwner'}), patch.object(subprocess, 'check_output', return_value='a' * 40 + '\n'):
|
||||
release.plan(argparse.Namespace(ref='feature/matrix', dry_run=True, simple=False))
|
||||
output = dict(line.split('=', 1) for line in Path('outputs').read_text().splitlines())
|
||||
self.assertEqual(output['dry_run'], 'true')
|
||||
self.assertEqual(output['owner_lower'], 'exampleowner')
|
||||
self.assertEqual(len(json.loads(output['matrix'])['include']), 5)
|
||||
|
||||
def test_docker_commands_do_not_publish_during_dry_run(self):
|
||||
fake_bin = Path('bin')
|
||||
fake_bin.mkdir()
|
||||
docker = fake_bin / 'docker'
|
||||
docker.write_text('#!/bin/sh\nprintf "%s\\n" "$*" >> "$DOCKER_LOG"\n')
|
||||
docker.chmod(0o755)
|
||||
env = {**os.environ, 'PATH': str(fake_bin.resolve()) + os.pathsep + os.environ['PATH'],
|
||||
'DOCKER_LOG': str(Path('docker.log').resolve()), 'RUNNER_TEMP': self.temp.name,
|
||||
'RELEASE_VERSION': '9.8.7', 'RELEASE_SHA': 'a' * 40, 'GITHUB_REPOSITORY': 'ExampleOwner/sub2api',
|
||||
'DRY_RUN': 'true', 'SIMPLE_RELEASE': 'false', 'DOCKERHUB_USERNAME': 'skip'}
|
||||
subprocess.run(['bash', str(ROOT / '.github/release-tools/release-images.sh')], env=env, check=True)
|
||||
log = Path('docker.log').read_text()
|
||||
self.assertEqual(log.count('buildx build'), 2)
|
||||
self.assertIn('linux/arm64', log)
|
||||
self.assertNotIn('--push', log)
|
||||
self.assertNotIn('imagetools', log)
|
||||
self.assertNotIn('skip/sub2api', log)
|
||||
self.assertIn('ghcr.io/exampleowner/sub2api', log)
|
||||
|
||||
|
||||
def test_published_full_and_simple_image_tags(self):
|
||||
fake_bin = Path('bin')
|
||||
fake_bin.mkdir()
|
||||
docker = fake_bin / 'docker'
|
||||
docker.write_text('#!/bin/sh\nprintf "%s\\n" "$*" >> "$DOCKER_LOG"\n')
|
||||
docker.chmod(0o755)
|
||||
for simple in (False, True):
|
||||
with self.subTest(simple=simple):
|
||||
log_path = Path(f'docker-{simple}.log').resolve()
|
||||
env = {**os.environ, 'PATH': str(fake_bin.resolve()) + os.pathsep + os.environ['PATH'],
|
||||
'DOCKER_LOG': str(log_path), 'RUNNER_TEMP': self.temp.name,
|
||||
'RELEASE_VERSION': '9.8.7', 'RELEASE_SHA': 'a' * 40, 'GITHUB_REPOSITORY': 'ExampleOwner/sub2api',
|
||||
'DRY_RUN': 'false', 'SIMPLE_RELEASE': str(simple).lower(), 'DOCKERHUB_USERNAME': 'fixturehub'}
|
||||
subprocess.run(['bash', str(ROOT / '.github/release-tools/release-images.sh')], env=env, check=True)
|
||||
log = log_path.read_text()
|
||||
self.assertIn('--push', log)
|
||||
self.assertEqual(log.count('buildx build'), 1 if simple else 2)
|
||||
if simple:
|
||||
self.assertNotIn('fixturehub', log)
|
||||
self.assertNotIn('imagetools', log)
|
||||
self.assertIn('ghcr.io/exampleowner/sub2api:latest', log)
|
||||
else:
|
||||
self.assertEqual(log.count('imagetools create'), 2)
|
||||
self.assertIn('fixturehub/sub2api:9.8', log)
|
||||
self.assertIn('ghcr.io/exampleowner/sub2api:9', log)
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -25,6 +25,8 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- name: Check Docker Compose simple mode environment
|
||||
run: /bin/sh deploy/tests/docker-compose-simple-mode-env-test.sh
|
||||
- uses: actions/setup-go@v6
|
||||
with:
|
||||
go-version-file: backend/go.mod
|
||||
@@ -80,3 +82,14 @@ jobs:
|
||||
version: v2.13
|
||||
args: --timeout=30m
|
||||
working-directory: backend
|
||||
|
||||
release-helpers:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.12'
|
||||
- run: python -m pip install -r .github/release-tools/requirements-release.txt
|
||||
- run: python -m unittest discover -s .github/release-tools -p 'test_release_matrix.py'
|
||||
- run: bash -n .github/release-tools/release-images.sh
|
||||
|
||||
+216
-114
@@ -7,7 +7,7 @@ on:
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
tag:
|
||||
description: 'Tag to release (e.g., v1.0.0)'
|
||||
description: 'Tag to release (or a branch/ref for dry run)'
|
||||
required: true
|
||||
type: string
|
||||
simple_release:
|
||||
@@ -15,179 +15,286 @@ on:
|
||||
required: false
|
||||
type: boolean
|
||||
default: false
|
||||
dry_run:
|
||||
description: 'Build and verify all selected artifacts without publishing'
|
||||
required: false
|
||||
type: boolean
|
||||
default: false
|
||||
|
||||
# 环境变量:合并 workflow_dispatch 输入和 repository variable
|
||||
# tag push 触发时读取 vars.SIMPLE_RELEASE,workflow_dispatch 时使用输入参数
|
||||
env:
|
||||
SIMPLE_RELEASE: ${{ github.event.inputs.simple_release == 'true' || vars.SIMPLE_RELEASE == 'true' }}
|
||||
concurrency:
|
||||
group: release-${{ github.repository }}
|
||||
cancel-in-progress: false
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
packages: write
|
||||
contents: read
|
||||
|
||||
env:
|
||||
SIMPLE_RELEASE: ${{ inputs.simple_release == true || vars.SIMPLE_RELEASE == 'true' }}
|
||||
DRY_RUN: ${{ github.event_name == 'workflow_dispatch' && inputs.dry_run == true }}
|
||||
|
||||
jobs:
|
||||
# Update VERSION file with tag version
|
||||
update-version:
|
||||
prepare:
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
sha: ${{ steps.plan.outputs.sha }}
|
||||
owner_lower: ${{ steps.plan.outputs.owner_lower }}
|
||||
tag: ${{ steps.plan.outputs.tag }}
|
||||
version: ${{ steps.plan.outputs.version }}
|
||||
date: ${{ steps.plan.outputs.date }}
|
||||
matrix: ${{ steps.plan.outputs.matrix }}
|
||||
simple: ${{ steps.plan.outputs.simple }}
|
||||
dry_run: ${{ steps.plan.outputs.dry_run }}
|
||||
steps:
|
||||
- name: Checkout
|
||||
- name: Checkout workflow tooling
|
||||
uses: actions/checkout@v6
|
||||
|
||||
- name: Update VERSION file
|
||||
with:
|
||||
ref: ${{ github.sha }}
|
||||
- name: Preserve tooling for releases of older tags
|
||||
run: cp -R .github/release-tools "$RUNNER_TEMP/release-tools"
|
||||
- name: Checkout selected application source
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.tag || github.ref }}
|
||||
fetch-depth: 0
|
||||
- uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.12'
|
||||
- run: python -m pip install -r "$RUNNER_TEMP/release-tools/requirements-release.txt"
|
||||
- name: Resolve one source commit and target matrix
|
||||
id: plan
|
||||
env:
|
||||
REQUESTED_REF: ${{ inputs.tag || github.ref_name }}
|
||||
run: |
|
||||
if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then
|
||||
VERSION=${{ github.event.inputs.tag }}
|
||||
VERSION=${VERSION#v}
|
||||
else
|
||||
VERSION=${GITHUB_REF#refs/tags/v}
|
||||
fi
|
||||
echo "$VERSION" > backend/cmd/server/VERSION
|
||||
echo "Updated VERSION file to: $VERSION"
|
||||
|
||||
- name: Upload VERSION artifact
|
||||
uses: actions/upload-artifact@v7
|
||||
args=(--ref "$REQUESTED_REF")
|
||||
if [[ "$SIMPLE_RELEASE" == true ]]; then args+=(--simple); fi
|
||||
if [[ "$DRY_RUN" == true ]]; then args+=(--dry-run); fi
|
||||
python "$RUNNER_TEMP/release-tools/release_matrix.py" plan "${args[@]}"
|
||||
- uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: release-tooling
|
||||
path: ${{ runner.temp }}/release-tools/
|
||||
retention-days: 1
|
||||
- uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: version-file
|
||||
path: backend/cmd/server/VERSION
|
||||
retention-days: 1
|
||||
|
||||
build-frontend:
|
||||
needs: prepare
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v6
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v6
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ needs.prepare.outputs.sha }}
|
||||
- uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 9
|
||||
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@v6
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '20'
|
||||
cache: 'pnpm'
|
||||
cache: pnpm
|
||||
cache-dependency-path: frontend/pnpm-lock.yaml
|
||||
|
||||
- name: Install dependencies
|
||||
run: pnpm install --frozen-lockfile
|
||||
- run: pnpm install --frozen-lockfile
|
||||
working-directory: frontend
|
||||
|
||||
- name: Build frontend
|
||||
run: pnpm run build
|
||||
- run: pnpm run build
|
||||
working-directory: frontend
|
||||
|
||||
- name: Upload frontend artifact
|
||||
uses: actions/upload-artifact@v7
|
||||
- uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: frontend-dist
|
||||
path: backend/internal/web/dist/
|
||||
retention-days: 1
|
||||
|
||||
release:
|
||||
needs: [update-version, build-frontend]
|
||||
build-binaries:
|
||||
needs: [prepare, build-frontend]
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix: ${{ fromJSON(needs.prepare.outputs.matrix) }}
|
||||
env:
|
||||
RELEASE_VERSION: ${{ needs.prepare.outputs.version }}
|
||||
RELEASE_DATE: ${{ needs.prepare.outputs.date }}
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ needs.prepare.outputs.sha }}
|
||||
fetch-depth: 0
|
||||
ref: ${{ github.event.inputs.tag || github.ref }}
|
||||
|
||||
- name: Download VERSION artifact
|
||||
uses: actions/download-artifact@v8
|
||||
- uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: release-tooling
|
||||
path: ${{ runner.temp }}/release-tools/
|
||||
- uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: version-file
|
||||
path: backend/cmd/server/
|
||||
|
||||
- name: Download frontend artifact
|
||||
uses: actions/download-artifact@v8
|
||||
- uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: frontend-dist
|
||||
path: backend/internal/web/dist/
|
||||
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v6
|
||||
- uses: actions/setup-go@v6
|
||||
with:
|
||||
go-version-file: backend/go.mod
|
||||
check-latest: false
|
||||
cache-dependency-path: backend/go.sum
|
||||
|
||||
cache: false
|
||||
- name: Verify Go version
|
||||
run: |
|
||||
go version | grep -q 'go1.27.0'
|
||||
run: go version | grep -q 'go1.27.0'
|
||||
- name: Restore target-specific Go cache
|
||||
id: go-cache
|
||||
uses: actions/cache/restore@v6
|
||||
with:
|
||||
path: |
|
||||
~/go/pkg/mod
|
||||
~/.cache/go-build
|
||||
key: release-go-${{ runner.os }}-${{ runner.arch }}-${{ matrix.goos }}-${{ matrix.goarch }}-${{ hashFiles('backend/go.mod', 'backend/go.sum') }}-${{ needs.prepare.outputs.sha }}
|
||||
restore-keys: |
|
||||
release-go-${{ runner.os }}-${{ runner.arch }}-${{ matrix.goos }}-${{ matrix.goarch }}-${{ hashFiles('backend/go.mod', 'backend/go.sum') }}-
|
||||
release-go-${{ runner.os }}-${{ runner.arch }}-${{ matrix.goos }}-${{ matrix.goarch }}-
|
||||
- uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.12'
|
||||
- run: python -m pip install -r "$RUNNER_TEMP/release-tools/requirements-release.txt"
|
||||
- name: Select one GoReleaser target
|
||||
env:
|
||||
TARGET_OS: ${{ matrix.goos }}
|
||||
TARGET_ARCH: ${{ matrix.goarch }}
|
||||
run: python "$RUNNER_TEMP/release-tools/release_matrix.py" config build --goos "$TARGET_OS" --goarch "$TARGET_ARCH" --output "$RUNNER_TEMP/goreleaser.matrix.yaml"
|
||||
- name: Build and archive without publishing
|
||||
uses: goreleaser/goreleaser-action@v7
|
||||
with:
|
||||
version: '~> v2'
|
||||
args: release --snapshot --clean --skip=validate,publish,docker,announce --config=${{ runner.temp }}/goreleaser.matrix.yaml
|
||||
env:
|
||||
GITHUB_REPO_OWNER: ${{ github.repository_owner }}
|
||||
GITHUB_REPO_OWNER_LOWER: ${{ needs.prepare.outputs.owner_lower }}
|
||||
GITHUB_REPO_NAME: ${{ github.event.repository.name }}
|
||||
DOCKERHUB_USERNAME: skip
|
||||
TAG_MESSAGE: ''
|
||||
- name: Verify and collect the target archive
|
||||
env:
|
||||
TARGET_OS: ${{ matrix.goos }}
|
||||
TARGET_ARCH: ${{ matrix.goarch }}
|
||||
RELEASE_SHA: ${{ needs.prepare.outputs.sha }}
|
||||
run: python "$RUNNER_TEMP/release-tools/release_matrix.py" collect --version "$RELEASE_VERSION" --sha "$RELEASE_SHA" --goos "$TARGET_OS" --goarch "$TARGET_ARCH" --output matrix-output
|
||||
- uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: release-binary-${{ matrix.goos }}-${{ matrix.goarch }}
|
||||
path: matrix-output/
|
||||
if-no-files-found: error
|
||||
retention-days: 1
|
||||
- name: Save updated target cache
|
||||
if: ${{ steps.go-cache.outputs.cache-hit != 'true' }}
|
||||
uses: actions/cache/save@v6
|
||||
with:
|
||||
path: |
|
||||
~/go/pkg/mod
|
||||
~/.cache/go-build
|
||||
key: ${{ steps.go-cache.outputs.cache-primary-key }}
|
||||
|
||||
# Docker setup for GoReleaser
|
||||
release:
|
||||
needs: [prepare, build-binaries]
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
packages: write
|
||||
env:
|
||||
SIMPLE_RELEASE: ${{ needs.prepare.outputs.simple }}
|
||||
DRY_RUN: ${{ needs.prepare.outputs.dry_run }}
|
||||
RELEASE_TAG: ${{ needs.prepare.outputs.tag }}
|
||||
RELEASE_VERSION: ${{ needs.prepare.outputs.version }}
|
||||
RELEASE_SHA: ${{ needs.prepare.outputs.sha }}
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ needs.prepare.outputs.sha }}
|
||||
fetch-depth: 0
|
||||
- uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: release-tooling
|
||||
path: ${{ runner.temp }}/release-tools/
|
||||
- uses: actions/download-artifact@v8
|
||||
with:
|
||||
pattern: release-binary-*
|
||||
path: release-input
|
||||
merge-multiple: true
|
||||
- uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.12'
|
||||
- run: python -m pip install -r "$RUNNER_TEMP/release-tools/requirements-release.txt"
|
||||
- name: Verify complete matrix and prepare Docker contexts
|
||||
run: |
|
||||
args=(--version "$RELEASE_VERSION" --sha "$RELEASE_SHA" --input release-input --output .release-context)
|
||||
if [[ "$SIMPLE_RELEASE" == true ]]; then args+=(--simple); fi
|
||||
python "$RUNNER_TEMP/release-tools/release_matrix.py" contexts "${args[@]}"
|
||||
- uses: actions/setup-go@v6
|
||||
with:
|
||||
go-version-file: backend/go.mod
|
||||
cache: false
|
||||
- name: Set up QEMU
|
||||
uses: docker/setup-qemu-action@v3
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
|
||||
- name: Login to DockerHub
|
||||
if: ${{ env.DOCKERHUB_USERNAME != '' }}
|
||||
if: ${{ env.DRY_RUN != 'true' && env.SIMPLE_RELEASE != 'true' && env.DOCKERHUB_USERNAME != '' }}
|
||||
uses: docker/login-action@v3
|
||||
env:
|
||||
DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||
with:
|
||||
username: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||
password: ${{ secrets.DOCKERHUB_TOKEN }}
|
||||
|
||||
- name: Login to GitHub Container Registry
|
||||
if: ${{ env.DRY_RUN != 'true' }}
|
||||
uses: docker/login-action@v3
|
||||
with:
|
||||
registry: ghcr.io
|
||||
username: ${{ github.repository_owner }}
|
||||
password: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
- name: Fetch tags with annotations
|
||||
run: |
|
||||
# 确保获取完整的 annotated tag 信息
|
||||
git fetch --tags --force
|
||||
|
||||
- name: Get tag message
|
||||
- name: Build images and publish manifests
|
||||
env:
|
||||
DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_USERNAME || 'skip' }}
|
||||
run: bash "$RUNNER_TEMP/release-tools/release-images.sh"
|
||||
- name: Get annotated tag notes
|
||||
id: tag_message
|
||||
run: |
|
||||
if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then
|
||||
TAG_NAME=${{ github.event.inputs.tag }}
|
||||
else
|
||||
TAG_NAME=${GITHUB_REF#refs/tags/}
|
||||
fi
|
||||
echo "Processing tag: $TAG_NAME"
|
||||
|
||||
# 获取完整的 tag message(跳过第一行标题)
|
||||
TAG_MESSAGE=$(git tag -l --format='%(contents:body)' "$TAG_NAME")
|
||||
|
||||
# 调试输出
|
||||
echo "Tag message length: ${#TAG_MESSAGE}"
|
||||
echo "Tag message preview:"
|
||||
echo "$TAG_MESSAGE" | head -10
|
||||
|
||||
# 使用 EOF 分隔符处理多行内容
|
||||
echo "message<<EOF" >> $GITHUB_OUTPUT
|
||||
echo "$TAG_MESSAGE" >> $GITHUB_OUTPUT
|
||||
echo "EOF" >> $GITHUB_OUTPUT
|
||||
|
||||
- name: Set lowercase owner for GHCR
|
||||
id: lowercase
|
||||
run: echo "owner=$(echo '${{ github.repository_owner }}' | tr '[:upper:]' '[:lower:]')" >> $GITHUB_OUTPUT
|
||||
|
||||
- name: Run GoReleaser
|
||||
TAG_MESSAGE=$(git tag -l --format='%(contents:body)' "$RELEASE_TAG")
|
||||
DELIMITER="TAG_MESSAGE_$(openssl rand -hex 16)"
|
||||
{
|
||||
echo "message<<$DELIMITER"
|
||||
echo "$TAG_MESSAGE"
|
||||
echo "$DELIMITER"
|
||||
} >> "$GITHUB_OUTPUT"
|
||||
- name: Prepare publication-only GoReleaser config
|
||||
run: |
|
||||
args=(publish --output "$RUNNER_TEMP/goreleaser.publisher.yaml")
|
||||
if [[ "$SIMPLE_RELEASE" == true ]]; then args+=(--simple); fi
|
||||
python "$RUNNER_TEMP/release-tools/release_matrix.py" config "${args[@]}"
|
||||
- name: Publish existing archives and release notes
|
||||
uses: goreleaser/goreleaser-action@v7
|
||||
with:
|
||||
version: '~> v2'
|
||||
args: release --clean --skip=validate ${{ env.SIMPLE_RELEASE == 'true' && '--config=.goreleaser.simple.yaml' || '' }}
|
||||
args: release --clean --skip=validate,before,docker ${{ env.DRY_RUN == 'true' && '--snapshot --skip=publish,announce' || '' }} --config=${{ runner.temp }}/goreleaser.publisher.yaml
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
GORELEASER_CURRENT_TAG: ${{ needs.prepare.outputs.tag }}
|
||||
TAG_MESSAGE: ${{ steps.tag_message.outputs.message }}
|
||||
GITHUB_REPO_OWNER: ${{ github.repository_owner }}
|
||||
GITHUB_REPO_OWNER_LOWER: ${{ steps.lowercase.outputs.owner }}
|
||||
GITHUB_REPO_NAME: ${{ github.event.repository.name }}
|
||||
GITHUB_REPO_OWNER_LOWER: ${{ needs.prepare.outputs.owner_lower }}
|
||||
DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_USERNAME || 'skip' }}
|
||||
- name: Upload dry-run checksums and metadata
|
||||
if: ${{ env.DRY_RUN == 'true' }}
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: release-dry-run-report
|
||||
path: |
|
||||
dist/checksums.txt
|
||||
dist/artifacts.json
|
||||
dist/metadata.json
|
||||
if-no-files-found: error
|
||||
retention-days: 3
|
||||
|
||||
# Update DockerHub description
|
||||
- name: Update DockerHub description
|
||||
if: ${{ env.SIMPLE_RELEASE != 'true' && env.DOCKERHUB_USERNAME != '' }}
|
||||
if: ${{ env.DRY_RUN != 'true' && env.SIMPLE_RELEASE != 'true' && env.DOCKERHUB_USERNAME != '' }}
|
||||
uses: peter-evans/dockerhub-description@v5
|
||||
env:
|
||||
DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||
@@ -200,8 +307,9 @@ jobs:
|
||||
|
||||
# Send Telegram notification
|
||||
- name: Send Telegram Notification
|
||||
if: ${{ env.SIMPLE_RELEASE != 'true' }}
|
||||
if: ${{ env.DRY_RUN != 'true' && env.SIMPLE_RELEASE != 'true' }}
|
||||
env:
|
||||
TAG_MESSAGE: ${{ steps.tag_message.outputs.message }}
|
||||
TELEGRAM_BOT_TOKEN: ${{ secrets.TELEGRAM_BOT_TOKEN }}
|
||||
TELEGRAM_CHAT_ID: ${{ secrets.TELEGRAM_CHAT_ID }}
|
||||
DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||
@@ -213,17 +321,12 @@ jobs:
|
||||
exit 0
|
||||
fi
|
||||
|
||||
if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then
|
||||
TAG_NAME=${{ github.event.inputs.tag }}
|
||||
else
|
||||
TAG_NAME=${GITHUB_REF#refs/tags/}
|
||||
fi
|
||||
TAG_NAME="$RELEASE_TAG"
|
||||
VERSION=${TAG_NAME#v}
|
||||
REPO="${{ github.repository }}"
|
||||
GHCR_IMAGE="ghcr.io/${REPO,,}" # ${,,} converts to lowercase
|
||||
|
||||
# 获取 tag message 内容并转义 Markdown 特殊字符
|
||||
TAG_MESSAGE='${{ steps.tag_message.outputs.message }}'
|
||||
TAG_MESSAGE=$(echo "$TAG_MESSAGE" | sed 's/\([_*`\[]\)/\\\1/g')
|
||||
|
||||
# 限制消息长度(Telegram 消息限制 4096 字符,预留空间给头尾固定内容)
|
||||
@@ -273,9 +376,13 @@ jobs:
|
||||
}')"
|
||||
|
||||
sync-version-file:
|
||||
needs: [release]
|
||||
if: ${{ needs.release.result == 'success' }}
|
||||
needs: [prepare, release]
|
||||
if: ${{ needs.release.result == 'success' && needs.prepare.outputs.dry_run != 'true' }}
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
env:
|
||||
RELEASE_VERSION: ${{ needs.prepare.outputs.version }}
|
||||
steps:
|
||||
- name: Checkout default branch
|
||||
uses: actions/checkout@v6
|
||||
@@ -284,12 +391,7 @@ jobs:
|
||||
|
||||
- name: Sync VERSION file to released tag
|
||||
run: |
|
||||
if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then
|
||||
VERSION=${{ github.event.inputs.tag }}
|
||||
VERSION=${VERSION#v}
|
||||
else
|
||||
VERSION=${GITHUB_REF#refs/tags/v}
|
||||
fi
|
||||
VERSION="$RELEASE_VERSION"
|
||||
|
||||
CURRENT_VERSION=$(tr -d '\r\n' < backend/cmd/server/VERSION || true)
|
||||
if [ "$CURRENT_VERSION" = "$VERSION" ]; then
|
||||
|
||||
+4
-1
@@ -133,7 +133,6 @@ deploy/docker-compose.override.yml
|
||||
.gocache/
|
||||
vite.config.js
|
||||
docs/*
|
||||
!docs/ANTIGRAVITY_ATTRIBUTION_429.md
|
||||
!docs/PAYMENT.md
|
||||
!docs/PAYMENT_CN.md
|
||||
!docs/ADMIN_PAYMENT_INTEGRATION_API.md
|
||||
@@ -142,17 +141,21 @@ docs/*
|
||||
!docs/COMPOSITE_GROUPS.md
|
||||
!docs/PLUGIN_DEVELOPMENT.md
|
||||
!docs/channel-monitor-v2-safe-defaults.md
|
||||
!docs/codex-credits.md
|
||||
!docs/legal/
|
||||
!docs/legal/*.md
|
||||
!docs/screenshots/
|
||||
docs/screenshots/*
|
||||
!docs/screenshots/mobile-account-actions-menu.png
|
||||
!docs/screenshots/codex-referral-dialog.png
|
||||
!docs/screenshots/codex-referral-success.png
|
||||
.serena/
|
||||
.codex/
|
||||
frontend/coverage/
|
||||
aicodex
|
||||
output/
|
||||
.codegraph/
|
||||
issues/
|
||||
|
||||
# Vitest / Vite cache at repo root
|
||||
.vite/
|
||||
|
||||
@@ -5,7 +5,11 @@ FRONTEND_CRITICAL_VITEST := \
|
||||
src/api/__tests__/client.spec.ts \
|
||||
src/api/__tests__/tokenRefresh.spec.ts \
|
||||
src/api/__tests__/keys.bulkUpdate.spec.ts \
|
||||
src/components/account/__tests__/OpenAIReferralCell.spec.ts \
|
||||
src/components/account/__tests__/OpenAIReferralCell.transport.spec.ts \
|
||||
src/components/account/__tests__/OpenAIQuotaResetCell.spark_shadow.spec.ts \
|
||||
src/components/keys/__tests__/BulkEditKeysModal.spec.ts \
|
||||
src/components/admin/user/__tests__/UserPlatformQuotaModal.spec.ts \
|
||||
src/views/user/__tests__/KeysView.spec.ts \
|
||||
src/api/__tests__/channelMonitorV2.spec.ts \
|
||||
src/views/auth/__tests__/LinuxDoCallbackView.spec.ts \
|
||||
|
||||
@@ -725,7 +725,10 @@ go generate ./cmd/server
|
||||
Simple Mode is designed for individual developers or internal teams who want quick access without full SaaS features.
|
||||
|
||||
- Enable: Set environment variable `RUN_MODE=simple`
|
||||
- Default groups are seeded on each startup. Set `SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS=false` (or YAML `simple_mode.auto_create_default_groups: false`) to manage groups yourself. The default is `true`; disabling it does not delete existing groups or change runtime auto-binding or admin concurrency setup.
|
||||
- Difference: Hides SaaS-related features and skips billing process
|
||||
- Optional key windows: Set `SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED=true` to enforce each API key's configured 5-hour, daily, and 7-day spending windows. The default is `false`; balance and subscription debit remain bypassed when enabled.
|
||||
- Window enforcement uses the database as its source of truth and records only API-key window usage. It is a post-request soft cap, so concurrent in-flight requests can overshoot by their final costs. Historical simple-mode usage is not backfilled.
|
||||
- Security note: In production, you must also set `SIMPLE_MODE_CONFIRM=true` to allow startup
|
||||
|
||||
---
|
||||
|
||||
@@ -735,7 +735,10 @@ OAuth / Setup Token 图片请求使用 Responses 主控模型调用 `image_gener
|
||||
简易模式适合个人开发者或内部团队快速使用,不依赖完整 SaaS 功能。
|
||||
|
||||
- 启用方式:设置环境变量 `RUN_MODE=simple`
|
||||
- 默认每次启动会补齐默认分组。设置 `SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS=false`(或 YAML `simple_mode.auto_create_default_groups: false`)可自行管理分组。默认值为 `true`;关闭后不删除已有分组,不改变运行时自动绑定或管理员并发设置。
|
||||
- 功能差异:隐藏 SaaS 相关功能,跳过计费流程
|
||||
- 可选密钥窗口:设置 `SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED=true` 后,按每个 API Key 配置的 5 小时、1 天、7 天消费窗口进行限制,默认值为 `false`;启用后仍跳过余额和订阅扣费。
|
||||
- 窗口限制以数据库为准,只记录 API Key 窗口用量。它在请求完成后记账,并发中的请求可能以各自最终费用超过窗口上限。启用前的历史用量不会自动补算。
|
||||
- 安全注意事项:生产环境需同时设置 `SIMPLE_MODE_CONFIRM=true` 才允许启动
|
||||
|
||||
---
|
||||
|
||||
@@ -1 +1 @@
|
||||
0.2.5
|
||||
0.2.10
|
||||
|
||||
@@ -99,6 +99,7 @@ func provideCleanup(
|
||||
accountExpiry *service.AccountExpiryService,
|
||||
cnProviderBalanceCheck *service.CNProviderBalanceCheckService,
|
||||
codexVersionSync *service.OpenAICodexVersionSyncService,
|
||||
claudeCodeVersionSync *service.ClaudeCodeVersionSyncService,
|
||||
proxyExpiry *service.ProxyExpiryService,
|
||||
subscriptionExpiry *service.SubscriptionExpiryService,
|
||||
usageCleanup *service.UsageCleanupService,
|
||||
@@ -124,6 +125,7 @@ func provideCleanup(
|
||||
quotaFlusher *service.UserPlatformQuotaUsageFlusher,
|
||||
upstreamBillingProbe *service.UpstreamBillingProbeService,
|
||||
ollamaCloudUsage *service.OllamaCloudUsageService,
|
||||
opencodeGoUsage *service.OpenCodeGoUsageService,
|
||||
auditLog *service.AuditLogService,
|
||||
openAIAutoReset *service.OpenAIQuotaAutoResetService,
|
||||
promptAudit *securityaudit.PromptService,
|
||||
@@ -272,6 +274,10 @@ func provideCleanup(
|
||||
codexVersionSync.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"ClaudeCodeVersionSyncService", func() error {
|
||||
claudeCodeVersionSync.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"ProxyExpiryService", func() error {
|
||||
proxyExpiry.Stop()
|
||||
return nil
|
||||
@@ -380,6 +386,12 @@ func provideCleanup(
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"OpenCodeGoUsageService", func() error {
|
||||
if opencodeGoUsage != nil {
|
||||
opencodeGoUsage.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
}
|
||||
|
||||
infraSteps := []cleanupStep{
|
||||
|
||||
@@ -185,7 +185,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
dashboardAggregationRepository := repository.NewDashboardAggregationRepository(db)
|
||||
dashboardStatsCache := repository.NewDashboardCache(redisClient, configConfig)
|
||||
dashboardService := service.NewDashboardService(usageLogRepository, dashboardAggregationRepository, dashboardStatsCache, configConfig)
|
||||
dashboardAggregationService := service.ProvideDashboardAggregationService(dashboardAggregationRepository, timingWheelService, leaderLockCache, db, configConfig)
|
||||
dashboardAggregationService := service.ProvideDashboardAggregationService(dashboardAggregationRepository, timingWheelService, leaderLockCache, db, configConfig, settingRepository)
|
||||
dashboardHandler := admin.NewDashboardHandler(dashboardService, dashboardAggregationService)
|
||||
adminGroupRepository := repository.NewAdminGroupRepository(client, db)
|
||||
adminAccountRepository := repository.NewAdminAccountRepository(client, db, schedulerCache)
|
||||
@@ -199,14 +199,14 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository, configConfig)
|
||||
grokQuotaFetcher := service.NewGrokQuotaFetcher()
|
||||
grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream, configConfig, usageLogRepository, settingService)
|
||||
openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory, openAIGatewayService)
|
||||
openAIReferralClient := repository.NewOpenAIReferralClient(privacyClientFactory)
|
||||
openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory, openAIReferralClient, openAIGatewayService)
|
||||
usageCache := service.NewUsageCache()
|
||||
accountUsageService := service.ProvideAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, grokQuotaService, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService, openAIGatewayService)
|
||||
pluginRepository := repository.NewPluginRepository(db)
|
||||
pluginHostInfo := providePluginHostInfo(buildInfo)
|
||||
pluginKVStore := repository.NewPluginKVStore(redisClient)
|
||||
pluginManager := service.NewPluginManager(pluginRepository, secretEncryptor, configConfig, pluginHostInfo, pluginKVStore)
|
||||
pluginManager.SetAccountDirectory(openAIGatewayService)
|
||||
pluginManager := service.ProvidePluginManager(pluginRepository, secretEncryptor, configConfig, pluginHostInfo, pluginKVStore, openAIGatewayService)
|
||||
accountTestService := service.ProvideAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService, openAIGatewayService, settingService, pluginManager)
|
||||
crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig)
|
||||
accountHandler := admin.ProvideAccountHandler(configConfig, adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator, grokQuotaService)
|
||||
@@ -287,14 +287,17 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
auditLogService := service.ProvideAuditLogService(auditLogRepository, settingService)
|
||||
auditLogHandler := admin.NewAuditLogHandler(auditLogService, totpService)
|
||||
upstreamBillingProbeService := service.ProvideUpstreamBillingProbeService(accountRepository, accountTestService, settingService, leaderLockCache, db)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, pluginHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService)
|
||||
openCodeGoUsageService := service.ProvideOpenCodeGoUsageService(accountRepository, httpUpstream, settingService, leaderLockCache, db)
|
||||
idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig)
|
||||
claudeResetCreditService := service.ProvideClaudeResetCreditService(accountRepository, claudeTokenProvider, proxyRepository, settingService, idempotencyCoordinator, leaderLockCache)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, pluginHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, claudeResetCreditService)
|
||||
usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig)
|
||||
userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient)
|
||||
userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig)
|
||||
legacyEngine := securityaudit.NewLegacyModerationAdapter(contentModerationService)
|
||||
coordinator := securityaudit.NewCoordinator(legacyEngine, promptService)
|
||||
gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator)
|
||||
openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, pluginManager, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator)
|
||||
openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, pluginManager, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator, compositeRouteResolver)
|
||||
handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService)
|
||||
totpHandler := handler.NewTotpHandler(totpService)
|
||||
passkeyRepository := repository.NewPasskeyRepository(db)
|
||||
@@ -320,7 +323,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
batchImageDownloadService := service.NewBatchImageDownloadService(batchImageRepository, accountRepository, batchImageDownloadLimiter, configConfig)
|
||||
batchImageCleanupService := service.ProvideBatchImageCleanupService(batchImageRepository, accountRepository, configConfig)
|
||||
batchImageHandler := handler.ProvideBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService, openAIGatewayHandler)
|
||||
idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig)
|
||||
idempotencyCleanupService := service.ProvideIdempotencyCleanupService(idempotencyRepository, configConfig)
|
||||
openAIQuotaAutoResetService := service.ProvideOpenAIQuotaAutoResetService(accountRepository, openAIQuotaService, rateLimitService, idempotencyCoordinator, auditLogService, settingService, leaderLockCache)
|
||||
handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, channelMonitorUserHandler, channelMonitorV2Handler, adminHandlers, gatewayHandler, openAIGatewayHandler, handlerSettingHandler, totpHandler, passkeyHandler, handlerPaymentHandler, paymentWebhookHandler, availableChannelHandler, modelPlazaHandler, asyncImageHandler, batchImageHandler, idempotencyCoordinator, idempotencyCleanupService, openAIQuotaAutoResetService)
|
||||
@@ -341,6 +343,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
accountExpiryService := service.ProvideAccountExpiryService(accountRepository)
|
||||
cnProviderBalanceCheckService := service.ProvideCNProviderBalanceCheckService(accountRepository, cnProviderBalanceService, cnProviderQuotaService, configConfig)
|
||||
openAICodexVersionSyncService := service.ProvideOpenAICodexVersionSyncService(settingRepository, settingService, gitHubReleaseClient)
|
||||
claudeCodeVersionSyncService := service.ProvideClaudeCodeVersionSyncService(settingRepository, settingService, gitHubReleaseClient)
|
||||
proxyExpiryService := service.ProvideProxyExpiryService(proxyRepository)
|
||||
subscriptionExpiryService := service.ProvideSubscriptionExpiryService(userSubscriptionRepository, settingRepository, notificationEmailService, leaderLockCache, db)
|
||||
batchImageWorkerRuntime := service.ProvideBatchImageWorkerRuntime(batchImageRepository, accountRepository, batchImageQueue, usageBillingRepository, usageLogRepository, batchImageModelPricingResolver, apiKeyAuthCacheInvalidator, configConfig)
|
||||
@@ -350,7 +353,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService, channelMonitorQuotaFetcher)
|
||||
channelMonitorV2Aggregator := service.ProvideChannelMonitorV2Aggregator(channelMonitorV2Repository, db, settingService)
|
||||
userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService)
|
||||
v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, cnProviderBalanceCheckService, openAICodexVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, auditLogService, openAIQuotaAutoResetService, promptService, pluginManager)
|
||||
v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, cnProviderBalanceCheckService, openAICodexVersionSyncService, claudeCodeVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, auditLogService, openAIQuotaAutoResetService, promptService, pluginManager)
|
||||
application := &Application{
|
||||
Server: httpServer,
|
||||
PromptAudit: promptService,
|
||||
@@ -405,6 +408,7 @@ func provideCleanup(
|
||||
accountExpiry *service.AccountExpiryService,
|
||||
cnProviderBalanceCheck *service.CNProviderBalanceCheckService,
|
||||
codexVersionSync *service.OpenAICodexVersionSyncService,
|
||||
claudeCodeVersionSync *service.ClaudeCodeVersionSyncService,
|
||||
proxyExpiry *service.ProxyExpiryService,
|
||||
subscriptionExpiry *service.SubscriptionExpiryService,
|
||||
usageCleanup *service.UsageCleanupService,
|
||||
@@ -430,6 +434,7 @@ func provideCleanup(
|
||||
quotaFlusher *service.UserPlatformQuotaUsageFlusher,
|
||||
upstreamBillingProbe *service.UpstreamBillingProbeService,
|
||||
ollamaCloudUsage *service.OllamaCloudUsageService,
|
||||
opencodeGoUsage *service.OpenCodeGoUsageService,
|
||||
auditLog *service.AuditLogService,
|
||||
openAIAutoReset *service.OpenAIQuotaAutoResetService,
|
||||
promptAudit *securityaudit.PromptService,
|
||||
@@ -577,6 +582,10 @@ func provideCleanup(
|
||||
codexVersionSync.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"ClaudeCodeVersionSyncService", func() error {
|
||||
claudeCodeVersionSync.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"ProxyExpiryService", func() error {
|
||||
proxyExpiry.Stop()
|
||||
return nil
|
||||
@@ -685,6 +694,12 @@ func provideCleanup(
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"OpenCodeGoUsageService", func() error {
|
||||
if opencodeGoUsage != nil {
|
||||
opencodeGoUsage.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
}
|
||||
|
||||
infraSteps := []cleanupStep{
|
||||
|
||||
@@ -41,6 +41,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) {
|
||||
)
|
||||
accountExpirySvc := service.NewAccountExpiryService(nil, time.Second)
|
||||
codexVersionSyncSvc := service.NewOpenAICodexVersionSyncService(nil, nil, nil, time.Second)
|
||||
claudeCodeVersionSyncSvc := service.NewClaudeCodeVersionSyncService(nil, nil, nil, time.Second)
|
||||
proxyExpirySvc := service.NewProxyExpiryService(nil, time.Second)
|
||||
subscriptionExpirySvc := service.NewSubscriptionExpiryService(nil, time.Second)
|
||||
pricingSvc := service.NewPricingService(cfg, nil)
|
||||
@@ -68,6 +69,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) {
|
||||
accountExpirySvc,
|
||||
nil, // cnProviderBalanceCheck
|
||||
codexVersionSyncSvc,
|
||||
claudeCodeVersionSyncSvc,
|
||||
proxyExpirySvc,
|
||||
subscriptionExpirySvc,
|
||||
&service.UsageCleanupService{},
|
||||
@@ -93,6 +95,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) {
|
||||
nil, // quotaFlusher
|
||||
nil, // upstreamBillingProbe
|
||||
nil, // ollamaCloudUsage
|
||||
nil, // opencodeGoUsage
|
||||
nil, // auditLog
|
||||
nil, // openAIAutoReset
|
||||
nil, // promptAudit
|
||||
|
||||
@@ -89,6 +89,7 @@ type Config struct {
|
||||
Pricing PricingConfig `mapstructure:"pricing"`
|
||||
Gateway GatewayConfig `mapstructure:"gateway"`
|
||||
APIKeyAuth APIKeyAuthCacheConfig `mapstructure:"api_key_auth_cache"`
|
||||
APIKeyCreate APIKeyCreateConfig `mapstructure:"api_key_create"`
|
||||
SubscriptionCache SubscriptionCacheConfig `mapstructure:"subscription_cache"`
|
||||
SubscriptionMaintenance SubscriptionMaintenanceConfig `mapstructure:"subscription_maintenance"`
|
||||
Dashboard DashboardCacheConfig `mapstructure:"dashboard_cache"`
|
||||
@@ -96,6 +97,7 @@ type Config struct {
|
||||
UsageCleanup UsageCleanupConfig `mapstructure:"usage_cleanup"`
|
||||
Concurrency ConcurrencyConfig `mapstructure:"concurrency"`
|
||||
TokenRefresh TokenRefreshConfig `mapstructure:"token_refresh"`
|
||||
SimpleMode SimpleModeConfig `mapstructure:"simple_mode" yaml:"simple_mode"`
|
||||
RunMode string `mapstructure:"run_mode" yaml:"run_mode"`
|
||||
Timezone string `mapstructure:"timezone"` // e.g. "Asia/Shanghai", "UTC"
|
||||
Gemini GeminiConfig `mapstructure:"gemini"`
|
||||
@@ -104,6 +106,14 @@ type Config struct {
|
||||
BatchImage BatchImageConfig `mapstructure:"batch_image"`
|
||||
ImageStorage ImageStorageConfig `mapstructure:"image_storage"`
|
||||
Plugins PluginConfig `mapstructure:"plugins"`
|
||||
|
||||
// Enforce only API-key spending windows in simple mode.
|
||||
SimpleModeKeyRateLimitEnabled bool `mapstructure:"simple_mode_key_rate_limit_enabled" yaml:"simple_mode_key_rate_limit_enabled"`
|
||||
}
|
||||
|
||||
// SimpleModeConfig controls startup behavior in simple mode.
|
||||
type SimpleModeConfig struct {
|
||||
AutoCreateDefaultGroups bool `mapstructure:"auto_create_default_groups" yaml:"auto_create_default_groups"`
|
||||
}
|
||||
|
||||
// PluginConfig 控制管理员手动上传的本地进程插件。
|
||||
@@ -912,6 +922,29 @@ type BillingConfig struct {
|
||||
// UserPlatformQuotaSentinelTTLSeconds sentinel(无 limit 占位)entry 的 TTL,
|
||||
// 显著短于 quota cache 默认 86400s 以控 Redis 内存;默认 3600=1h。
|
||||
UserPlatformQuotaSentinelTTLSeconds int `mapstructure:"user_platform_quota_sentinel_ttl_seconds"`
|
||||
// InflightReservation 余额模式在途请求预留(Redis),防止并发请求在预检时看到同一份余额而集体透支。
|
||||
InflightReservation InflightReservationConfig `mapstructure:"inflight_reservation"`
|
||||
}
|
||||
|
||||
// InflightReservationConfig 余额模式在途预留配置。
|
||||
// 准入时按 输入估算 + 输出单价 × max_tokens 估算单请求费用,在 Redis 中原子地
|
||||
// 校验 缓存余额 - 在途预留合计 >= 估算 后登记预留,请求结束(任意路径)释放。
|
||||
// 估算失败或 Redis 不可用时 fail-open,退回旧的仅余额 > 阈值检查。
|
||||
type InflightReservationConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"`
|
||||
// TTLSeconds 单条预留的最长存活时间;进程崩溃等泄漏的预留到期自动失效。
|
||||
TTLSeconds int `mapstructure:"ttl_seconds"`
|
||||
// DefaultMaxTokens 请求未携带 max_tokens 时用于估算的输出 token 数。
|
||||
DefaultMaxTokens int `mapstructure:"default_max_tokens"`
|
||||
// MaxOutputTokens 估算输出 token 的上限(max_tokens 超出时截断)。
|
||||
MaxOutputTokens int `mapstructure:"max_output_tokens"`
|
||||
// MaxInputTokens 输入 token 估算(请求体字节数 / 4)的上限。
|
||||
MaxInputTokens int `mapstructure:"max_input_tokens"`
|
||||
// MaxReservationUSD 单请求预留金额上限;0 表示不设上限。
|
||||
MaxReservationUSD float64 `mapstructure:"max_reservation_usd"`
|
||||
// FailClosedOnUnpriced 无法为请求估算费用(模型/分组/渠道均无定价)时是否拒绝请求。
|
||||
// 默认 false:放行且不预留(fail-open,节流告警日志)。
|
||||
FailClosedOnUnpriced bool `mapstructure:"fail_closed_on_unpriced"`
|
||||
}
|
||||
|
||||
type CircuitBreakerConfig struct {
|
||||
@@ -1687,6 +1720,14 @@ type APIKeyAuthCacheConfig struct {
|
||||
InvalidAbuse InvalidAuthAbuseConfig `mapstructure:"invalid_abuse"`
|
||||
}
|
||||
|
||||
// APIKeyCreateConfig 用户创建 API Key 的防滥用限制(0 表示不限制)
|
||||
type APIKeyCreateConfig struct {
|
||||
// MaxActivePerUser 单个用户同时存在(未删除)的 API Key 上限
|
||||
MaxActivePerUser int `mapstructure:"max_active_per_user"`
|
||||
// MaxPerUserPerHour 单个用户每小时可创建的 API Key 次数(删除不返还次数)
|
||||
MaxPerUserPerHour int `mapstructure:"max_per_user_per_hour"`
|
||||
}
|
||||
|
||||
type InvalidAuthAbuseConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"`
|
||||
Threshold int `mapstructure:"threshold"`
|
||||
@@ -1992,6 +2033,8 @@ func configureConfigSource(setConfigFile, addConfigPath func(string)) {
|
||||
|
||||
func setDefaults() {
|
||||
viper.SetDefault("run_mode", RunModeStandard)
|
||||
viper.SetDefault("simple_mode.auto_create_default_groups", true)
|
||||
viper.SetDefault("simple_mode_key_rate_limit_enabled", false)
|
||||
|
||||
// Server
|
||||
viper.SetDefault("server.host", "0.0.0.0")
|
||||
@@ -2082,6 +2125,13 @@ func setDefaults() {
|
||||
viper.SetDefault("billing.minimum_balance_reserve", 0.000001)
|
||||
viper.SetDefault("billing.user_platform_quota_cache_ttl_seconds", 86400)
|
||||
viper.SetDefault("billing.user_platform_quota_sentinel_ttl_seconds", 3600)
|
||||
viper.SetDefault("billing.inflight_reservation.enabled", true)
|
||||
viper.SetDefault("billing.inflight_reservation.ttl_seconds", 900)
|
||||
viper.SetDefault("billing.inflight_reservation.default_max_tokens", 8192)
|
||||
viper.SetDefault("billing.inflight_reservation.max_output_tokens", 128000)
|
||||
viper.SetDefault("billing.inflight_reservation.max_input_tokens", 200000)
|
||||
viper.SetDefault("billing.inflight_reservation.max_reservation_usd", 0)
|
||||
viper.SetDefault("billing.inflight_reservation.fail_closed_on_unpriced", false)
|
||||
|
||||
// Turnstile
|
||||
viper.SetDefault("turnstile.required", false)
|
||||
@@ -2321,6 +2371,8 @@ func setDefaults() {
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.window_seconds", 60)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.block_seconds", 60)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.capacity", 16384)
|
||||
viper.SetDefault("api_key_create.max_active_per_user", 200)
|
||||
viper.SetDefault("api_key_create.max_per_user_per_hour", 60)
|
||||
|
||||
// Subscription auth L1 cache
|
||||
viper.SetDefault("subscription_cache.l1_size", 16384)
|
||||
@@ -2700,6 +2752,12 @@ func (c *Config) Validate() error {
|
||||
return fmt.Errorf("server.h2c.max_upload_buffer_per_stream must be positive")
|
||||
}
|
||||
}
|
||||
if c.APIKeyCreate.MaxActivePerUser < 0 {
|
||||
return fmt.Errorf("api_key_create.max_active_per_user must be non-negative")
|
||||
}
|
||||
if c.APIKeyCreate.MaxPerUserPerHour < 0 {
|
||||
return fmt.Errorf("api_key_create.max_per_user_per_hour must be non-negative")
|
||||
}
|
||||
if c.APIKeyAuth.InvalidAbuse.Enabled {
|
||||
if c.APIKeyAuth.InvalidAbuse.Threshold < 10 {
|
||||
return fmt.Errorf("api_key_auth_cache.invalid_abuse.threshold must be at least 10")
|
||||
@@ -3051,6 +3109,11 @@ func (c *Config) Validate() error {
|
||||
if c.Billing.MinimumBalanceReserve < 0 {
|
||||
return fmt.Errorf("billing.minimum_balance_reserve must be non-negative")
|
||||
}
|
||||
if c.Billing.InflightReservation.TTLSeconds < 0 || c.Billing.InflightReservation.DefaultMaxTokens < 0 ||
|
||||
c.Billing.InflightReservation.MaxOutputTokens < 0 || c.Billing.InflightReservation.MaxInputTokens < 0 ||
|
||||
c.Billing.InflightReservation.MaxReservationUSD < 0 {
|
||||
return fmt.Errorf("billing.inflight_reservation values must be non-negative")
|
||||
}
|
||||
if c.Database.MaxOpenConns <= 0 {
|
||||
return fmt.Errorf("database.max_open_conns must be positive")
|
||||
}
|
||||
|
||||
@@ -79,6 +79,15 @@ func TestLoadServerTimingConfig(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestLoadSimpleModeKeyRateLimitEnabledFromEnvironment(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED", "true")
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.True(t, cfg.SimpleModeKeyRateLimitEnabled)
|
||||
}
|
||||
|
||||
func TestLoadRedisUsernameFromEnvironment(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("REDIS_USERNAME", "app-user")
|
||||
@@ -2630,3 +2639,42 @@ func TestLoad_DefaultGatewayImageStreamConfig(t *testing.T) {
|
||||
t.Fatalf("image stream timeout = %d, want greater than ordinary stream timeout %d", cfg.Gateway.ImageStreamDataIntervalTimeout, cfg.Gateway.StreamDataIntervalTimeout)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadSimpleModeAutoCreateDefaultGroups(t *testing.T) {
|
||||
for _, loader := range []struct {
|
||||
name string
|
||||
load func() (*Config, error)
|
||||
}{{"Load", Load}, {"LoadForBootstrap", LoadForBootstrap}} {
|
||||
t.Run(loader.name, func(t *testing.T) {
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
yaml string
|
||||
env string
|
||||
unset bool
|
||||
want bool
|
||||
}{
|
||||
{name: "unset env uses application default", unset: true, want: true},
|
||||
{name: "empty env uses application default", want: true},
|
||||
{name: "unset env preserves yaml false", unset: true, yaml: "simple_mode:\n auto_create_default_groups: false\n", want: false},
|
||||
{name: "empty env preserves yaml false", yaml: "simple_mode:\n auto_create_default_groups: false\n", want: false},
|
||||
{name: "env false overrides yaml true", yaml: "simple_mode:\n auto_create_default_groups: true\n", env: "false", want: false},
|
||||
{name: "env false without yaml", env: "false", want: false},
|
||||
{name: "env true overrides yaml false", yaml: "simple_mode:\n auto_create_default_groups: false\n", env: "true", want: true},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS", tt.env)
|
||||
if tt.unset {
|
||||
require.NoError(t, os.Unsetenv("SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS"))
|
||||
}
|
||||
path := filepath.Join(t.TempDir(), "config.yaml")
|
||||
require.NoError(t, os.WriteFile(path, []byte("run_mode: simple\n"+tt.yaml), 0o600))
|
||||
t.Setenv("CONFIG_FILE", path)
|
||||
cfg, err := loader.load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tt.want, cfg.SimpleMode.AutoCreateDefaultGroups)
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -190,6 +190,8 @@ var DefaultBedrockModelMapping = map[string]string{
|
||||
"claude-opus-4-1": "us.anthropic.claude-opus-4-1-20250805-v1:0",
|
||||
"claude-opus-4-20250514": "us.anthropic.claude-opus-4-20250514-v1:0",
|
||||
// Claude Sonnet
|
||||
// Sonnet 5.5 is available on bedrock-runtime through Global inference only.
|
||||
"claude-sonnet-5-5": "global.anthropic.claude-sonnet-5-5",
|
||||
"claude-sonnet-5": "us.anthropic.claude-sonnet-5-v1",
|
||||
"claude-sonnet-4-6-thinking": "us.anthropic.claude-sonnet-4-6",
|
||||
"claude-sonnet-4-6": "us.anthropic.claude-sonnet-4-6",
|
||||
|
||||
@@ -109,9 +109,10 @@ func TestDefaultBedrockModelMapping_ContainsNewClaudeModels(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := map[string]string{
|
||||
"claude-fable-5-1": "anthropic.claude-fable-5-1",
|
||||
"claude-fable-5": "anthropic.claude-fable-5",
|
||||
"claude-opus-4-8": "us.anthropic.claude-opus-4-8-v1",
|
||||
"claude-fable-5-1": "anthropic.claude-fable-5-1",
|
||||
"claude-fable-5": "anthropic.claude-fable-5",
|
||||
"claude-opus-4-8": "us.anthropic.claude-opus-4-8-v1",
|
||||
"claude-sonnet-5-5": "global.anthropic.claude-sonnet-5-5",
|
||||
}
|
||||
for from, want := range cases {
|
||||
got, ok := DefaultBedrockModelMapping[from]
|
||||
|
||||
@@ -48,6 +48,7 @@ func NewOAuthHandler(oauthService *service.OAuthService) *OAuthHandler {
|
||||
|
||||
// AccountHandler handles admin account management
|
||||
type AccountHandler struct {
|
||||
claudeResetCredits claudeResetReader
|
||||
adminService service.AdminService
|
||||
oauthService *service.OAuthService
|
||||
openaiOAuthService *service.OpenAIOAuthService
|
||||
@@ -66,6 +67,7 @@ type AccountHandler struct {
|
||||
upstreamBillingProbe *service.UpstreamBillingProbeService
|
||||
ollamaCloudUsage *service.OllamaCloudUsageService
|
||||
cfg *config.Config
|
||||
opencodeGoUsage *service.OpenCodeGoUsageService
|
||||
}
|
||||
|
||||
// SetUpstreamBillingProbeService attaches the optional remote billing probe service.
|
||||
@@ -77,6 +79,10 @@ func (h *AccountHandler) SetOllamaCloudUsageService(usage *service.OllamaCloudUs
|
||||
h.ollamaCloudUsage = usage
|
||||
}
|
||||
|
||||
func (h *AccountHandler) SetOpenCodeGoUsageService(usage *service.OpenCodeGoUsageService) {
|
||||
h.opencodeGoUsage = usage
|
||||
}
|
||||
|
||||
// NewAccountHandler creates a new admin account handler
|
||||
func NewAccountHandler(
|
||||
adminService service.AdminService,
|
||||
@@ -674,14 +680,22 @@ func (h *AccountHandler) List(c *gin.Context) {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
if h.ollamaCloudUsage != nil && len(accounts) > 0 {
|
||||
if len(accounts) > 0 {
|
||||
accountPointers := make([]*service.Account, len(accounts))
|
||||
for index := range accounts {
|
||||
accountPointers[index] = &accounts[index]
|
||||
}
|
||||
if err := h.ollamaCloudUsage.ResolveAccounts(c.Request.Context(), accountPointers); err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
if h.ollamaCloudUsage != nil {
|
||||
if err := h.ollamaCloudUsage.ResolveAccounts(c.Request.Context(), accountPointers); err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if h.opencodeGoUsage != nil {
|
||||
if err := h.opencodeGoUsage.ResolveOpenCodeGoUsageAccounts(c.Request.Context(), accountPointers); err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -942,6 +956,12 @@ func (h *AccountHandler) GetByID(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
}
|
||||
if h.opencodeGoUsage != nil {
|
||||
if err := h.opencodeGoUsage.ResolveOpenCodeGoUsageAccounts(c.Request.Context(), []*service.Account{account}); err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, h.buildAccountResponseWithRuntime(c.Request.Context(), account))
|
||||
}
|
||||
@@ -1609,6 +1629,10 @@ func (h *AccountHandler) ApplyOAuthCredentials(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// Re-auth only returns authentication fields. Preserve account configuration
|
||||
// stored alongside them (for example model_mapping), while allowing the new
|
||||
// OAuth values to replace their existing counterparts.
|
||||
req.Credentials = service.MergeCredentials(existing.Credentials, req.Credentials)
|
||||
// Drop SSO/password residue; re-auth must leave only OAuth tokens on disk.
|
||||
req.Credentials = service.SanitizeStoredCredentials(existing.Platform, req.Credentials)
|
||||
|
||||
|
||||
@@ -142,6 +142,44 @@ func TestApplyOAuthCredentialsRejectsMalformedOpenAILongContextBillingBeforeMuta
|
||||
require.Zero(t, stub.updateAccountExtraCalls)
|
||||
}
|
||||
|
||||
func TestApplyOAuthCredentialsPreservesExistingNonAuthCredentials(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
stub := newStubAdminService()
|
||||
stub.getAccountResult = &service.Account{
|
||||
ID: 1,
|
||||
Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "old-token",
|
||||
"refresh_token": "old-refresh-token",
|
||||
"model_mapping": map[string]any{"gpt-5": "gpt-5"},
|
||||
"account_id": "existing-account-id",
|
||||
"password": "must-not-survive",
|
||||
"sso_token": "must-not-survive",
|
||||
"cookie": "must-not-survive",
|
||||
},
|
||||
}
|
||||
handler := NewAccountHandler(stub, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
router := gin.New()
|
||||
router.POST("/accounts/:id/apply-oauth-credentials", handler.ApplyOAuthCredentials)
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPost, "/accounts/1/apply-oauth-credentials", bytes.NewBufferString(
|
||||
`{"type":"oauth","credentials":{"access_token":"new-token","refresh_token":"new-refresh-token"}}`,
|
||||
))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.Equal(t, 1, stub.updateAccountCalls)
|
||||
require.Equal(t, map[string]any{
|
||||
"access_token": "new-token",
|
||||
"refresh_token": "new-refresh-token",
|
||||
"model_mapping": map[string]any{"gpt-5": "gpt-5"},
|
||||
"account_id": "existing-account-id",
|
||||
}, stub.lastUpdateAccountInput.Credentials)
|
||||
}
|
||||
|
||||
func TestOpenAIOAuthCodexPATBoundaryRejectsMalformedOpenAILongContextBillingValueBeforeTokenValidation(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
handler := NewOpenAIOAuthHandler(nil, newStubAdminService(), nil, nil)
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type openCodeGoUsageAutoRefreshRequest struct {
|
||||
Enabled *bool `json:"enabled" binding:"required"`
|
||||
}
|
||||
|
||||
func (h *AccountHandler) GetOpenCodeGoUsageSettings(c *gin.Context) {
|
||||
if h.opencodeGoUsage == nil {
|
||||
response.ErrorFrom(c, service.ErrOpenCodeGoUsageUnavailable)
|
||||
return
|
||||
}
|
||||
settings, err := h.opencodeGoUsage.GetSettings(c.Request.Context())
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, settings)
|
||||
}
|
||||
|
||||
func (h *AccountHandler) UpdateOpenCodeGoUsageSettings(c *gin.Context) {
|
||||
if h.opencodeGoUsage == nil {
|
||||
response.ErrorFrom(c, service.ErrOpenCodeGoUsageUnavailable)
|
||||
return
|
||||
}
|
||||
var req service.OpenCodeGoUsageSettings
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
if err := h.opencodeGoUsage.UpdateSettings(c.Request.Context(), &req); err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
settings, err := h.opencodeGoUsage.GetSettings(c.Request.Context())
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, settings)
|
||||
}
|
||||
|
||||
func (h *AccountHandler) GetOpenCodeGoUsage(c *gin.Context) {
|
||||
if !h.requireOpenCodeGoUsage(c) {
|
||||
return
|
||||
}
|
||||
accountID, ok := openCodeGoUsageAccountID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
state, err := h.opencodeGoUsage.GetState(c.Request.Context(), accountID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, state)
|
||||
}
|
||||
|
||||
func (h *AccountHandler) SetOpenCodeGoUsageAutoRefresh(c *gin.Context) {
|
||||
if !h.requireOpenCodeGoUsage(c) {
|
||||
return
|
||||
}
|
||||
accountID, ok := openCodeGoUsageAccountID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var req openCodeGoUsageAutoRefreshRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
state, err := h.opencodeGoUsage.SetAutoRefresh(c.Request.Context(), accountID, *req.Enabled)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, state)
|
||||
}
|
||||
|
||||
func (h *AccountHandler) RefreshOpenCodeGoUsage(c *gin.Context) {
|
||||
if !h.requireOpenCodeGoUsage(c) {
|
||||
return
|
||||
}
|
||||
accountID, ok := openCodeGoUsageAccountID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
state, err := h.opencodeGoUsage.Refresh(c.Request.Context(), accountID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, state)
|
||||
}
|
||||
|
||||
func (h *AccountHandler) requireOpenCodeGoUsage(c *gin.Context) bool {
|
||||
if h != nil && h.opencodeGoUsage != nil {
|
||||
return true
|
||||
}
|
||||
response.ErrorFrom(c, service.ErrOpenCodeGoUsageUnavailable)
|
||||
return false
|
||||
}
|
||||
|
||||
func openCodeGoUsageAccountID(c *gin.Context) (int64, bool) {
|
||||
if c == nil {
|
||||
return 0, false
|
||||
}
|
||||
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil || accountID <= 0 {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return 0, false
|
||||
}
|
||||
return accountID, true
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type openCodeGoUsageHandlerTestRepo struct {
|
||||
service.AccountRepository
|
||||
account *service.Account
|
||||
accounts []*service.Account
|
||||
}
|
||||
|
||||
func (r *openCodeGoUsageHandlerTestRepo) GetByID(_ context.Context, id int64) (*service.Account, error) {
|
||||
if r.account != nil && r.account.ID == id {
|
||||
return r.account, nil
|
||||
}
|
||||
for _, account := range r.accounts {
|
||||
if account.ID == id {
|
||||
return account, nil
|
||||
}
|
||||
}
|
||||
return nil, service.ErrAccountNotFound
|
||||
}
|
||||
|
||||
func (r *openCodeGoUsageHandlerTestRepo) ListOpenCodeGoUsageGroupAccounts(_ context.Context, anchors []*service.Account) ([]service.Account, error) {
|
||||
wanted := make(map[string]struct{}, len(anchors))
|
||||
for _, anchor := range anchors {
|
||||
if apiKey, ok := openCodeGoUsageHandlerTestAPIKey(anchor); ok {
|
||||
wanted[apiKey] = struct{}{}
|
||||
}
|
||||
}
|
||||
result := make([]service.Account, 0, len(r.accounts)+1)
|
||||
if r.account != nil {
|
||||
if apiKey, ok := openCodeGoUsageHandlerTestAPIKey(r.account); ok {
|
||||
if _, match := wanted[apiKey]; match {
|
||||
result = append(result, *r.account)
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, account := range r.accounts {
|
||||
if apiKey, ok := openCodeGoUsageHandlerTestAPIKey(account); ok {
|
||||
if _, match := wanted[apiKey]; match {
|
||||
result = append(result, *account)
|
||||
}
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func openCodeGoUsageHandlerTestAPIKey(account *service.Account) (string, bool) {
|
||||
if account == nil || account.Credentials == nil {
|
||||
return "", false
|
||||
}
|
||||
apiKey, ok := account.Credentials["api_key"].(string)
|
||||
return apiKey, ok && apiKey != ""
|
||||
}
|
||||
|
||||
func (r *openCodeGoUsageHandlerTestRepo) SetOpenCodeGoUsageAutoRefresh(context.Context, *service.Account, bool) error {
|
||||
return nil
|
||||
}
|
||||
func (r *openCodeGoUsageHandlerTestRepo) UpdateOpenCodeGoUsageSnapshot(context.Context, *service.Account, *service.OpenCodeGoUsageSnapshot) error {
|
||||
return nil
|
||||
}
|
||||
func (r *openCodeGoUsageHandlerTestRepo) DisableOpenCodeGoUsageAutoRefresh(context.Context, *service.Account) error {
|
||||
return nil
|
||||
}
|
||||
func (r *openCodeGoUsageHandlerTestRepo) ListDueOpenCodeGoUsageAccounts(context.Context, time.Time, time.Duration, time.Duration, int) ([]service.Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func newOpenCodeGoUsageHandlerTestService(t *testing.T) *service.OpenCodeGoUsageService {
|
||||
t.Helper()
|
||||
svc := service.NewOpenCodeGoUsageService(nil, nil, nil)
|
||||
t.Cleanup(svc.Stop)
|
||||
return svc
|
||||
}
|
||||
|
||||
func newOpenCodeGoUsageHandlerContext(method, target, body, id string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(method, target, bytes.NewBufferString(body))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = request
|
||||
if id != "" {
|
||||
ctx.Params = gin.Params{{Key: "id", Value: id}}
|
||||
}
|
||||
return ctx, recorder
|
||||
}
|
||||
|
||||
func TestOpenCodeGoUsageHandlersValidateRequestsAndDependencies(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newOpenCodeGoUsageHandlerTestService(t)
|
||||
|
||||
t.Run("invalid account id", func(t *testing.T) {
|
||||
ctx, recorder := newOpenCodeGoUsageHandlerContext(http.MethodGet, "/admin/accounts/not-an-id/opencode-go-usage", "", "not-an-id")
|
||||
(&AccountHandler{opencodeGoUsage: svc}).GetOpenCodeGoUsage(ctx)
|
||||
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
||||
})
|
||||
|
||||
t.Run("missing enabled", func(t *testing.T) {
|
||||
ctx, recorder := newOpenCodeGoUsageHandlerContext(http.MethodPut, "/admin/accounts/7/opencode-go-usage/auto-refresh", `{}`, "7")
|
||||
(&AccountHandler{opencodeGoUsage: svc}).SetOpenCodeGoUsageAutoRefresh(ctx)
|
||||
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
||||
})
|
||||
|
||||
t.Run("service unavailable", func(t *testing.T) {
|
||||
ctx, recorder := newOpenCodeGoUsageHandlerContext(http.MethodGet, "/admin/accounts/7/opencode-go-usage", "", "7")
|
||||
(&AccountHandler{}).GetOpenCodeGoUsage(ctx)
|
||||
require.Equal(t, http.StatusServiceUnavailable, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "OPENCODE_GO_USAGE_UNAVAILABLE")
|
||||
})
|
||||
}
|
||||
|
||||
func TestGetOpenCodeGoUsageSettingsHandlerSuccess(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
ctx, recorder := newOpenCodeGoUsageHandlerContext(http.MethodGet, "/admin/accounts/opencode-go-usage/settings", "", "")
|
||||
handler := &AccountHandler{opencodeGoUsage: newOpenCodeGoUsageHandlerTestService(t)}
|
||||
|
||||
handler.GetOpenCodeGoUsageSettings(ctx)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), `"enabled":false`)
|
||||
require.Contains(t, recorder.Body.String(), `"interval_minutes":15`)
|
||||
}
|
||||
|
||||
func TestOpenCodeGoUsageStateEmbeddedInListAndDetail(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
now := time.Now().UTC()
|
||||
account := &service.Account{
|
||||
ID: 7, Name: "opencode", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"base_url": "https://opencode.ai/zen/go/v1", "api_key": "test-key"},
|
||||
Extra: map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: true,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: &service.OpenCodeGoUsageSnapshot{
|
||||
Status: service.OpenCodeGoUsageStatusOK, Data: &service.OpenCodeGoUsageData{
|
||||
Rolling: service.OpenCodeGoUsageWindow{Status: "ok", Percent: 6},
|
||||
},
|
||||
LastAttemptAt: now, NextRefreshAt: now.Add(time.Hour),
|
||||
},
|
||||
},
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
repo := &openCodeGoUsageHandlerTestRepo{account: account}
|
||||
adminService := newStubAdminService()
|
||||
adminService.accounts = []service.Account{*account}
|
||||
adminService.getAccountResult = account
|
||||
usageService := service.NewOpenCodeGoUsageService(repo, nil, nil)
|
||||
t.Cleanup(usageService.Stop)
|
||||
handler := NewAccountHandler(adminService, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
handler.SetOpenCodeGoUsageService(usageService)
|
||||
router := gin.New()
|
||||
router.GET("/accounts", handler.List)
|
||||
router.GET("/accounts/:id", handler.GetByID)
|
||||
router.GET("/accounts/:id/opencode-go-usage", handler.GetOpenCodeGoUsage)
|
||||
|
||||
listRecorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(listRecorder, httptest.NewRequest(http.MethodGet, "/accounts?page=1&page_size=20", nil))
|
||||
require.Equal(t, http.StatusOK, listRecorder.Code)
|
||||
var listPayload struct {
|
||||
Data struct {
|
||||
Items []struct {
|
||||
OpenCodeGoUsage *service.OpenCodeGoUsageState `json:"opencode_go_usage"`
|
||||
} `json:"items"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(listRecorder.Body.Bytes(), &listPayload))
|
||||
require.Len(t, listPayload.Data.Items, 1)
|
||||
require.NotNil(t, listPayload.Data.Items[0].OpenCodeGoUsage)
|
||||
require.True(t, listPayload.Data.Items[0].OpenCodeGoUsage.AutoRefreshEnabled)
|
||||
require.Equal(t, 6.0, listPayload.Data.Items[0].OpenCodeGoUsage.Snapshot.Data.Rolling.Percent)
|
||||
|
||||
detailRecorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(detailRecorder, httptest.NewRequest(http.MethodGet, "/accounts/7", nil))
|
||||
require.Equal(t, http.StatusOK, detailRecorder.Code)
|
||||
var detailPayload struct {
|
||||
Data struct {
|
||||
OpenCodeGoUsage *service.OpenCodeGoUsageState `json:"opencode_go_usage"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(detailRecorder.Body.Bytes(), &detailPayload))
|
||||
require.NotNil(t, detailPayload.Data.OpenCodeGoUsage)
|
||||
require.Equal(t, 6.0, detailPayload.Data.OpenCodeGoUsage.Snapshot.Data.Rolling.Percent)
|
||||
|
||||
stateRecorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(stateRecorder, httptest.NewRequest(http.MethodGet, "/accounts/7/opencode-go-usage", nil))
|
||||
require.Equal(t, http.StatusOK, stateRecorder.Code)
|
||||
var statePayload struct {
|
||||
Data service.OpenCodeGoUsageState `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(stateRecorder.Body.Bytes(), &statePayload))
|
||||
require.Equal(t, statePayload.Data.Snapshot, detailPayload.Data.OpenCodeGoUsage.Snapshot)
|
||||
|
||||
for _, body := range []string{listRecorder.Body.String(), detailRecorder.Body.String(), stateRecorder.Body.String()} {
|
||||
require.NotContains(t, body, "test-key")
|
||||
}
|
||||
}
|
||||
@@ -201,6 +201,42 @@ func (h *AffiliateHandler) GetUserOverview(c *gin.Context) {
|
||||
response.Success(c, overview)
|
||||
}
|
||||
|
||||
// WithdrawQuotaRequest 是登记线下提现的请求体。Amount 是已在站外打款给
|
||||
// 用户的金额,由 AffiliateService 校验。
|
||||
type WithdrawQuotaRequest struct {
|
||||
Amount float64 `json:"amount"`
|
||||
}
|
||||
|
||||
// WithdrawQuota records an offline withdrawal paid outside the site and
|
||||
// deducts it from the user's available affiliate quota. The required
|
||||
// Idempotency-Key header identifies one registration: a repeated request with
|
||||
// the same key deducts nothing and returns the first result with
|
||||
// X-Idempotency-Replayed: true.
|
||||
// POST /api/v1/admin/affiliates/users/:user_id/withdraw
|
||||
func (h *AffiliateHandler) WithdrawQuota(c *gin.Context) {
|
||||
userID, err := strconv.ParseInt(c.Param("user_id"), 10, 64)
|
||||
if err != nil || userID <= 0 {
|
||||
response.BadRequest(c, "Invalid user_id")
|
||||
return
|
||||
}
|
||||
|
||||
var req WithdrawQuotaRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
result, err := h.affiliateService.AdminWithdrawQuota(c.Request.Context(), userID, req.Amount, c.GetHeader("Idempotency-Key"))
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
if result != nil && result.Replayed {
|
||||
c.Header("X-Idempotency-Replayed", "true")
|
||||
}
|
||||
response.Success(c, result)
|
||||
}
|
||||
|
||||
// ListInviteRecords returns all inviter-invitee relationships.
|
||||
// GET /api/v1/admin/affiliates/invites
|
||||
func (h *AffiliateHandler) ListInviteRecords(c *gin.Context) {
|
||||
@@ -214,7 +250,8 @@ func (h *AffiliateHandler) ListInviteRecords(c *gin.Context) {
|
||||
response.Paginated(c, items, total, filter.Page, filter.PageSize)
|
||||
}
|
||||
|
||||
// ListRebateRecords returns all order-level affiliate rebate records.
|
||||
// ListRebateRecords returns every affiliate rebate accrual, with order details
|
||||
// when the rebate came from a payment order.
|
||||
// GET /api/v1/admin/affiliates/rebates
|
||||
func (h *AffiliateHandler) ListRebateRecords(c *gin.Context) {
|
||||
page, pageSize := response.ParsePagination(c)
|
||||
@@ -227,7 +264,8 @@ func (h *AffiliateHandler) ListRebateRecords(c *gin.Context) {
|
||||
response.Paginated(c, items, total, filter.Page, filter.PageSize)
|
||||
}
|
||||
|
||||
// ListTransferRecords returns all affiliate quota-to-balance transfer records.
|
||||
// ListTransferRecords returns affiliate quota outflows: user transfers into
|
||||
// balance and admin-recorded offline withdrawals.
|
||||
// GET /api/v1/admin/affiliates/transfers
|
||||
func (h *AffiliateHandler) ListTransferRecords(c *gin.Context) {
|
||||
page, pageSize := response.ParsePagination(c)
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
const testWithdrawIdempotencyKey = "affiliate-withdraw-1-5d0c7a2e"
|
||||
|
||||
type withdrawQuotaCall struct {
|
||||
userID int64
|
||||
amount float64
|
||||
operationID string
|
||||
}
|
||||
|
||||
type withdrawQuotaRepoStub struct {
|
||||
service.AffiliateRepository
|
||||
calls []withdrawQuotaCall
|
||||
result *service.AffiliateWithdrawResult
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *withdrawQuotaRepoStub) WithdrawQuota(_ context.Context, userID int64, amount float64, operationID string) (*service.AffiliateWithdrawResult, error) {
|
||||
s.calls = append(s.calls, withdrawQuotaCall{userID: userID, amount: amount, operationID: operationID})
|
||||
return s.result, s.err
|
||||
}
|
||||
|
||||
type withdrawQuotaResponse struct {
|
||||
Code int `json:"code"`
|
||||
Reason string `json:"reason"`
|
||||
Data json.RawMessage `json:"data"`
|
||||
}
|
||||
|
||||
func performWithdrawQuota(t *testing.T, repo *withdrawQuotaRepoStub, userID, idempotencyKey, body string) (*httptest.ResponseRecorder, withdrawQuotaResponse) {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
handler := NewAffiliateHandler(service.NewAffiliateService(repo, nil, nil, nil), nil)
|
||||
router.POST("/api/v1/admin/affiliates/users/:user_id/withdraw", handler.WithdrawQuota)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/affiliates/users/"+userID+"/withdraw", bytes.NewBufferString(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if idempotencyKey != "" {
|
||||
req.Header.Set("Idempotency-Key", idempotencyKey)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
var resp withdrawQuotaResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
return rec, resp
|
||||
}
|
||||
|
||||
// TestAffiliateHandlerWithdrawQuota_Success 验证登记线下提现把路径中的用户、
|
||||
// 舍入后的金额与幂等键派生的 operation_id 交给仓储,并返回登记结果。
|
||||
func TestAffiliateHandlerWithdrawQuota_Success(t *testing.T) {
|
||||
repo := &withdrawQuotaRepoStub{result: &service.AffiliateWithdrawResult{LedgerID: 9, UserID: 42, Amount: 12.34567891}}
|
||||
rec, resp := performWithdrawQuota(t, repo, "42", testWithdrawIdempotencyKey, `{"amount":12.3456789149}`)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, 0, resp.Code)
|
||||
var data service.AffiliateWithdrawResult
|
||||
require.NoError(t, json.Unmarshal(resp.Data, &data))
|
||||
require.Equal(t, int64(9), data.LedgerID)
|
||||
require.Empty(t, rec.Header().Get("X-Idempotency-Replayed"))
|
||||
|
||||
require.Len(t, repo.calls, 1)
|
||||
require.Equal(t, int64(42), repo.calls[0].userID)
|
||||
require.Equal(t, 12.34567891, repo.calls[0].amount)
|
||||
require.Regexp(t, `^[0-9a-f]{64}$`, repo.calls[0].operationID)
|
||||
}
|
||||
|
||||
// TestAffiliateHandlerWithdrawQuota_SameKeyMapsToSameOperation 验证重试时带同一个
|
||||
// Idempotency-Key 的请求落到同一个 operation_id,换一个键则是另一笔登记。
|
||||
func TestAffiliateHandlerWithdrawQuota_SameKeyMapsToSameOperation(t *testing.T) {
|
||||
repo := &withdrawQuotaRepoStub{result: &service.AffiliateWithdrawResult{LedgerID: 9, UserID: 42, Amount: 10}}
|
||||
for _, key := range []string{testWithdrawIdempotencyKey, testWithdrawIdempotencyKey, testWithdrawIdempotencyKey + "-next"} {
|
||||
rec, _ := performWithdrawQuota(t, repo, "42", key, `{"amount":10}`)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
}
|
||||
|
||||
require.Len(t, repo.calls, 3)
|
||||
require.Equal(t, repo.calls[0].operationID, repo.calls[1].operationID)
|
||||
require.NotEqual(t, repo.calls[0].operationID, repo.calls[2].operationID)
|
||||
}
|
||||
|
||||
// TestAffiliateHandlerWithdrawQuota_ReplayReturnsFirstResult 验证命中已有登记时
|
||||
// 原样返回首次登记的结果,并以 X-Idempotency-Replayed 标明是重放。
|
||||
func TestAffiliateHandlerWithdrawQuota_ReplayReturnsFirstResult(t *testing.T) {
|
||||
repo := &withdrawQuotaRepoStub{result: &service.AffiliateWithdrawResult{
|
||||
LedgerID: 9,
|
||||
UserID: 42,
|
||||
Amount: 10,
|
||||
AvailableQuotaAfter: 90,
|
||||
HistoryQuotaAfter: 100,
|
||||
Replayed: true,
|
||||
}}
|
||||
rec, resp := performWithdrawQuota(t, repo, "42", testWithdrawIdempotencyKey, `{"amount":10}`)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "true", rec.Header().Get("X-Idempotency-Replayed"))
|
||||
require.JSONEq(t, `{
|
||||
"ledger_id": 9,
|
||||
"user_id": 42,
|
||||
"amount": 10,
|
||||
"available_quota_after": 90,
|
||||
"frozen_quota_after": 0,
|
||||
"history_quota_after": 100
|
||||
}`, string(resp.Data))
|
||||
}
|
||||
|
||||
// TestAffiliateHandlerWithdrawQuota_RejectsBadRequests 验证非法用户 ID、非法 JSON、
|
||||
// 缺少或非法的幂等键与非法金额在进入仓储前以 400 拒绝。
|
||||
func TestAffiliateHandlerWithdrawQuota_RejectsBadRequests(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
userID string
|
||||
key string
|
||||
body string
|
||||
wantReason string
|
||||
}{
|
||||
{name: "non-numeric user id", userID: "abc", key: testWithdrawIdempotencyKey, body: `{"amount":1}`},
|
||||
{name: "non-positive user id", userID: "0", key: testWithdrawIdempotencyKey, body: `{"amount":1}`},
|
||||
{name: "malformed json", userID: "42", key: testWithdrawIdempotencyKey, body: `{"amount":`},
|
||||
{name: "missing idempotency key", userID: "42", body: `{"amount":1}`, wantReason: "IDEMPOTENCY_KEY_REQUIRED"},
|
||||
{name: "invalid idempotency key", userID: "42", key: "has space", body: `{"amount":1}`, wantReason: "IDEMPOTENCY_KEY_INVALID"},
|
||||
{name: "zero amount", userID: "42", key: testWithdrawIdempotencyKey, body: `{"amount":0}`, wantReason: "AFFILIATE_WITHDRAW_AMOUNT_INVALID"},
|
||||
{name: "missing amount", userID: "42", key: testWithdrawIdempotencyKey, body: `{}`, wantReason: "AFFILIATE_WITHDRAW_AMOUNT_INVALID"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
repo := &withdrawQuotaRepoStub{}
|
||||
rec, resp := performWithdrawQuota(t, repo, tc.userID, tc.key, tc.body)
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
require.Equal(t, tc.wantReason, resp.Reason)
|
||||
require.Empty(t, repo.calls)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAffiliateHandlerWithdrawQuota_RepositoryRejections 验证额度不足以 400、
|
||||
// 同一幂等键换了用户或金额以 409 返回,前端据错误码提示。
|
||||
func TestAffiliateHandlerWithdrawQuota_RepositoryRejections(t *testing.T) {
|
||||
cases := []struct {
|
||||
err error
|
||||
wantStatus int
|
||||
wantReason string
|
||||
}{
|
||||
{err: service.ErrAffiliateQuotaInsufficient, wantStatus: http.StatusBadRequest, wantReason: "AFFILIATE_QUOTA_INSUFFICIENT"},
|
||||
{err: service.ErrIdempotencyKeyConflict, wantStatus: http.StatusConflict, wantReason: "IDEMPOTENCY_KEY_CONFLICT"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.wantReason, func(t *testing.T) {
|
||||
repo := &withdrawQuotaRepoStub{err: tc.err}
|
||||
rec, resp := performWithdrawQuota(t, repo, "42", testWithdrawIdempotencyKey, `{"amount":5}`)
|
||||
|
||||
require.Equal(t, tc.wantStatus, rec.Code)
|
||||
require.Equal(t, tc.wantReason, resp.Reason)
|
||||
require.Empty(t, rec.Header().Get("X-Idempotency-Replayed"))
|
||||
require.Len(t, repo.calls, 1)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -140,7 +140,13 @@ func (h *BackupHandler) DeleteBackup(c *gin.Context) {
|
||||
response.BadRequest(c, "backup ID is required")
|
||||
return
|
||||
}
|
||||
if err := h.backupService.DeleteBackup(c.Request.Context(), backupID); err != nil {
|
||||
var err error
|
||||
if c.Query("delete_archived") == "true" {
|
||||
err = h.backupService.DeleteArchivedBackup(c.Request.Context(), backupID)
|
||||
} else {
|
||||
err = h.backupService.DeleteBackup(c.Request.Context(), backupID)
|
||||
}
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -57,22 +57,22 @@ type updateChannelRequest struct {
|
||||
}
|
||||
|
||||
type channelModelPricingRequest struct {
|
||||
Platform string `json:"platform" binding:"omitempty,max=50"`
|
||||
Models []string `json:"models" binding:"required,min=1,max=100"`
|
||||
BillingMode string `json:"billing_mode" binding:"omitempty,oneof=token per_request image"`
|
||||
InputPrice *float64 `json:"input_price" binding:"omitempty,min=0"`
|
||||
OutputPrice *float64 `json:"output_price" binding:"omitempty,min=0"`
|
||||
CacheWritePrice *float64 `json:"cache_write_price" binding:"omitempty,min=0"`
|
||||
CacheWrite1hPrice *float64 `json:"cache_write_1h_price" binding:"omitempty,min=0"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price" binding:"omitempty,min=0"`
|
||||
FastMultiplier *float64 `json:"fast_multiplier" binding:"omitempty,gt=0"`
|
||||
FlexMultiplier *float64 `json:"flex_multiplier" binding:"omitempty,gt=0"`
|
||||
MaxReasoningEffortMultiplier *float64 `json:"max_reasoning_effort_multiplier" binding:"omitempty,gt=0"`
|
||||
ImageInputPrice *float64 `json:"image_input_price" binding:"omitempty,min=0"`
|
||||
ImageOutputPrice *float64 `json:"image_output_price" binding:"omitempty,min=0"`
|
||||
PerRequestPrice *float64 `json:"per_request_price" binding:"omitempty,min=0"`
|
||||
Intervals []pricingIntervalRequest `json:"intervals"`
|
||||
TimePricing *channelTimePricingRequest `json:"time_pricing"`
|
||||
Platform string `json:"platform" binding:"omitempty,max=50"`
|
||||
Models []string `json:"models" binding:"required,min=1,max=100"`
|
||||
BillingMode string `json:"billing_mode" binding:"omitempty,oneof=token per_request image"`
|
||||
InputPrice *float64 `json:"input_price" binding:"omitempty,min=0"`
|
||||
OutputPrice *float64 `json:"output_price" binding:"omitempty,min=0"`
|
||||
CacheWritePrice *float64 `json:"cache_write_price" binding:"omitempty,min=0"`
|
||||
CacheWrite1hPrice *float64 `json:"cache_write_1h_price" binding:"omitempty,min=0"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price" binding:"omitempty,min=0"`
|
||||
FastMultiplier *float64 `json:"fast_multiplier" binding:"omitempty,gt=0"`
|
||||
FlexMultiplier *float64 `json:"flex_multiplier" binding:"omitempty,gt=0"`
|
||||
ReasoningEffortMultipliers map[string]float64 `json:"reasoning_effort_multipliers"`
|
||||
ImageInputPrice *float64 `json:"image_input_price" binding:"omitempty,min=0"`
|
||||
ImageOutputPrice *float64 `json:"image_output_price" binding:"omitempty,min=0"`
|
||||
PerRequestPrice *float64 `json:"per_request_price" binding:"omitempty,min=0"`
|
||||
Intervals []pricingIntervalRequest `json:"intervals"`
|
||||
TimePricing *channelTimePricingRequest `json:"time_pricing"`
|
||||
}
|
||||
|
||||
type channelTimePricingRequest struct {
|
||||
@@ -130,23 +130,23 @@ type channelResponse struct {
|
||||
}
|
||||
|
||||
type channelModelPricingResponse struct {
|
||||
ID int64 `json:"id"`
|
||||
Platform string `json:"platform"`
|
||||
Models []string `json:"models"`
|
||||
BillingMode string `json:"billing_mode"`
|
||||
InputPrice *float64 `json:"input_price"`
|
||||
OutputPrice *float64 `json:"output_price"`
|
||||
CacheWritePrice *float64 `json:"cache_write_price"`
|
||||
CacheWrite1hPrice *float64 `json:"cache_write_1h_price"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price"`
|
||||
FastMultiplier *float64 `json:"fast_multiplier"`
|
||||
FlexMultiplier *float64 `json:"flex_multiplier"`
|
||||
MaxReasoningEffortMultiplier *float64 `json:"max_reasoning_effort_multiplier"`
|
||||
ImageInputPrice *float64 `json:"image_input_price"`
|
||||
ImageOutputPrice *float64 `json:"image_output_price"`
|
||||
PerRequestPrice *float64 `json:"per_request_price"`
|
||||
Intervals []pricingIntervalResponse `json:"intervals"`
|
||||
TimePricing *channelTimePricingResponse `json:"time_pricing"`
|
||||
ID int64 `json:"id"`
|
||||
Platform string `json:"platform"`
|
||||
Models []string `json:"models"`
|
||||
BillingMode string `json:"billing_mode"`
|
||||
InputPrice *float64 `json:"input_price"`
|
||||
OutputPrice *float64 `json:"output_price"`
|
||||
CacheWritePrice *float64 `json:"cache_write_price"`
|
||||
CacheWrite1hPrice *float64 `json:"cache_write_1h_price"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price"`
|
||||
FastMultiplier *float64 `json:"fast_multiplier"`
|
||||
FlexMultiplier *float64 `json:"flex_multiplier"`
|
||||
ReasoningEffortMultipliers map[string]float64 `json:"reasoning_effort_multipliers"`
|
||||
ImageInputPrice *float64 `json:"image_input_price"`
|
||||
ImageOutputPrice *float64 `json:"image_output_price"`
|
||||
PerRequestPrice *float64 `json:"per_request_price"`
|
||||
Intervals []pricingIntervalResponse `json:"intervals"`
|
||||
TimePricing *channelTimePricingResponse `json:"time_pricing"`
|
||||
}
|
||||
|
||||
type channelTimePricingResponse struct {
|
||||
@@ -260,23 +260,23 @@ func pricingToResponse(p *service.ChannelModelPricing) channelModelPricingRespon
|
||||
intervals = append(intervals, intervalToResponse(iv))
|
||||
}
|
||||
return channelModelPricingResponse{
|
||||
ID: p.ID,
|
||||
Platform: platform,
|
||||
Models: models,
|
||||
BillingMode: billingMode,
|
||||
InputPrice: p.InputPrice,
|
||||
OutputPrice: p.OutputPrice,
|
||||
CacheWritePrice: p.CacheWritePrice,
|
||||
CacheWrite1hPrice: p.CacheWrite1hPrice,
|
||||
CacheReadPrice: p.CacheReadPrice,
|
||||
FastMultiplier: p.FastMultiplier,
|
||||
FlexMultiplier: p.FlexMultiplier,
|
||||
MaxReasoningEffortMultiplier: p.MaxReasoningEffortMultiplier,
|
||||
ImageInputPrice: p.ImageInputPrice,
|
||||
ImageOutputPrice: p.ImageOutputPrice,
|
||||
PerRequestPrice: p.PerRequestPrice,
|
||||
Intervals: intervals,
|
||||
TimePricing: timePricingToResponse(p.TimePricing),
|
||||
ID: p.ID,
|
||||
Platform: platform,
|
||||
Models: models,
|
||||
BillingMode: billingMode,
|
||||
InputPrice: p.InputPrice,
|
||||
OutputPrice: p.OutputPrice,
|
||||
CacheWritePrice: p.CacheWritePrice,
|
||||
CacheWrite1hPrice: p.CacheWrite1hPrice,
|
||||
CacheReadPrice: p.CacheReadPrice,
|
||||
FastMultiplier: p.FastMultiplier,
|
||||
FlexMultiplier: p.FlexMultiplier,
|
||||
ReasoningEffortMultipliers: p.ReasoningEffortMultipliers,
|
||||
ImageInputPrice: p.ImageInputPrice,
|
||||
ImageOutputPrice: p.ImageOutputPrice,
|
||||
PerRequestPrice: p.PerRequestPrice,
|
||||
Intervals: intervals,
|
||||
TimePricing: timePricingToResponse(p.TimePricing),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -353,29 +353,28 @@ func pricingRequestToService(reqs []channelModelPricingRequest, allowChannelMult
|
||||
SortOrder: iv.SortOrder,
|
||||
})
|
||||
}
|
||||
var fastMultiplier, flexMultiplier, maxReasoningEffortMultiplier *float64
|
||||
var fastMultiplier, flexMultiplier *float64
|
||||
if allowChannelMultipliers {
|
||||
fastMultiplier = r.FastMultiplier
|
||||
flexMultiplier = r.FlexMultiplier
|
||||
maxReasoningEffortMultiplier = r.MaxReasoningEffortMultiplier
|
||||
}
|
||||
result = append(result, service.ChannelModelPricing{
|
||||
Platform: platform,
|
||||
Models: r.Models,
|
||||
BillingMode: billingMode,
|
||||
InputPrice: r.InputPrice,
|
||||
OutputPrice: r.OutputPrice,
|
||||
CacheWritePrice: r.CacheWritePrice,
|
||||
CacheWrite1hPrice: r.CacheWrite1hPrice,
|
||||
CacheReadPrice: r.CacheReadPrice,
|
||||
FastMultiplier: fastMultiplier,
|
||||
FlexMultiplier: flexMultiplier,
|
||||
MaxReasoningEffortMultiplier: maxReasoningEffortMultiplier,
|
||||
ImageInputPrice: r.ImageInputPrice,
|
||||
ImageOutputPrice: r.ImageOutputPrice,
|
||||
PerRequestPrice: r.PerRequestPrice,
|
||||
Intervals: intervals,
|
||||
TimePricing: timePricingRequestToService(r.TimePricing),
|
||||
Platform: platform,
|
||||
Models: r.Models,
|
||||
BillingMode: billingMode,
|
||||
InputPrice: r.InputPrice,
|
||||
OutputPrice: r.OutputPrice,
|
||||
CacheWritePrice: r.CacheWritePrice,
|
||||
CacheWrite1hPrice: r.CacheWrite1hPrice,
|
||||
CacheReadPrice: r.CacheReadPrice,
|
||||
FastMultiplier: fastMultiplier,
|
||||
FlexMultiplier: flexMultiplier,
|
||||
ReasoningEffortMultipliers: r.ReasoningEffortMultipliers,
|
||||
ImageInputPrice: r.ImageInputPrice,
|
||||
ImageOutputPrice: r.ImageOutputPrice,
|
||||
PerRequestPrice: r.PerRequestPrice,
|
||||
Intervals: intervals,
|
||||
TimePricing: timePricingRequestToService(r.TimePricing),
|
||||
})
|
||||
}
|
||||
return result
|
||||
@@ -620,15 +619,15 @@ func (h *ChannelHandler) GetModelDefaultPricing(c *gin.Context) {
|
||||
}
|
||||
|
||||
response.Success(c, gin.H{
|
||||
"found": true,
|
||||
"input_price": pricing.InputPricePerToken,
|
||||
"output_price": pricing.OutputPricePerToken,
|
||||
"cache_write_price": cacheWritePrice,
|
||||
"cache_write_1h_price": cacheWrite1hPrice,
|
||||
"cache_read_price": pricing.CacheReadPricePerToken,
|
||||
"max_reasoning_effort_multiplier": pricing.MaxReasoningEffortMultiplier,
|
||||
"image_input_price": pricing.ImageInputPricePerToken,
|
||||
"image_output_price": pricing.ImageOutputPricePerToken,
|
||||
"found": true,
|
||||
"input_price": pricing.InputPricePerToken,
|
||||
"output_price": pricing.OutputPricePerToken,
|
||||
"cache_write_price": cacheWritePrice,
|
||||
"cache_write_1h_price": cacheWrite1hPrice,
|
||||
"cache_read_price": pricing.CacheReadPricePerToken,
|
||||
"reasoning_effort_multipliers": pricing.ReasoningEffortMultipliers,
|
||||
"image_input_price": pricing.ImageInputPricePerToken,
|
||||
"image_output_price": pricing.ImageOutputPricePerToken,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -449,15 +449,14 @@ func TestPricingRequestToService_TimePricingNil(t *testing.T) {
|
||||
require.Nil(t, got[0].TimePricing)
|
||||
}
|
||||
|
||||
// 账号成本统计规则不支持倍率:allowChannelMultipliers=false 时必须丢弃,
|
||||
// 避免渠道倍率意外污染账号成本口径。
|
||||
// Fast/Flex 和区间倍率只用于渠道售价;思考等级倍率可在账号统计规则中独立配置。
|
||||
func TestPricingRequestToService_MultipliersGatedByFlag(t *testing.T) {
|
||||
req := channelModelPricingRequest{
|
||||
Models: []string{"gpt-5"},
|
||||
BillingMode: "token",
|
||||
FastMultiplier: float64Ptr(2.5),
|
||||
FlexMultiplier: float64Ptr(0.5),
|
||||
MaxReasoningEffortMultiplier: float64Ptr(3),
|
||||
Models: []string{"gpt-5"},
|
||||
BillingMode: "token",
|
||||
FastMultiplier: float64Ptr(2.5),
|
||||
FlexMultiplier: float64Ptr(0.5),
|
||||
ReasoningEffortMultipliers: map[string]float64{"high": 1.5, "max": 3},
|
||||
Intervals: []pricingIntervalRequest{{
|
||||
MinTokens: 272000,
|
||||
InputMultiplier: float64Ptr(2),
|
||||
@@ -470,7 +469,7 @@ func TestPricingRequestToService_MultipliersGatedByFlag(t *testing.T) {
|
||||
allowed := pricingRequestToService([]channelModelPricingRequest{req}, true)
|
||||
require.Equal(t, float64Ptr(2.5), allowed[0].FastMultiplier)
|
||||
require.Equal(t, float64Ptr(0.5), allowed[0].FlexMultiplier)
|
||||
require.Equal(t, float64Ptr(3), allowed[0].MaxReasoningEffortMultiplier)
|
||||
require.Equal(t, req.ReasoningEffortMultipliers, allowed[0].ReasoningEffortMultipliers)
|
||||
require.Equal(t, float64Ptr(2), allowed[0].Intervals[0].InputMultiplier)
|
||||
require.Equal(t, float64Ptr(1.5), allowed[0].Intervals[0].OutputMultiplier)
|
||||
require.Equal(t, float64Ptr(2), allowed[0].Intervals[0].CacheWriteMultiplier)
|
||||
@@ -479,7 +478,7 @@ func TestPricingRequestToService_MultipliersGatedByFlag(t *testing.T) {
|
||||
dropped := pricingRequestToService([]channelModelPricingRequest{req}, false)
|
||||
require.Nil(t, dropped[0].FastMultiplier)
|
||||
require.Nil(t, dropped[0].FlexMultiplier)
|
||||
require.Nil(t, dropped[0].MaxReasoningEffortMultiplier)
|
||||
require.Equal(t, req.ReasoningEffortMultipliers, dropped[0].ReasoningEffortMultipliers)
|
||||
require.Nil(t, dropped[0].Intervals[0].InputMultiplier)
|
||||
require.Nil(t, dropped[0].Intervals[0].OutputMultiplier)
|
||||
require.Nil(t, dropped[0].Intervals[0].CacheWriteMultiplier)
|
||||
@@ -511,6 +510,23 @@ func TestPricingToResponse_TimePricingNil(t *testing.T) {
|
||||
require.Nil(t, got.TimePricing)
|
||||
}
|
||||
|
||||
func TestPricingRequestAndResponse_ReasoningEffortMultipliers(t *testing.T) {
|
||||
var req channelModelPricingRequest
|
||||
require.NoError(t, json.Unmarshal([]byte(`{"models":["custom-model"],"reasoning_effort_multipliers":{"none":0.5,"high":1.5,"max":3}}`), &req))
|
||||
pricing := pricingRequestToService([]channelModelPricingRequest{req}, true)
|
||||
got := pricingToResponse(&pricing[0])
|
||||
require.Equal(t, map[string]float64{"none": 0.5, "high": 1.5, "max": 3}, got.ReasoningEffortMultipliers)
|
||||
data, err := json.Marshal(got)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, string(data), `"reasoning_effort_multipliers":{"high":1.5,"max":3,"none":0.5}`)
|
||||
require.NotContains(t, string(data), "max_reasoning_effort_multiplier")
|
||||
|
||||
req = channelModelPricingRequest{}
|
||||
require.NoError(t, json.Unmarshal([]byte(`{"models":["custom-model"],"reasoning_effort_multipliers":{}}`), &req))
|
||||
pricing = pricingRequestToService([]channelModelPricingRequest{req}, true)
|
||||
require.Empty(t, pricing[0].ReasoningEffortMultipliers)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 3. SyncPricingModels handler
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -583,10 +599,10 @@ func TestGetModelDefaultPricing_ReturnsFable51CacheTTLs(t *testing.T) {
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
var body struct {
|
||||
Data struct {
|
||||
Found bool `json:"found"`
|
||||
CacheWritePrice float64 `json:"cache_write_price"`
|
||||
CacheWrite1hPrice *float64 `json:"cache_write_1h_price"`
|
||||
MaxReasoningEffortMultiplier *float64 `json:"max_reasoning_effort_multiplier"`
|
||||
Found bool `json:"found"`
|
||||
CacheWritePrice float64 `json:"cache_write_price"`
|
||||
CacheWrite1hPrice *float64 `json:"cache_write_1h_price"`
|
||||
ReasoningEffortMultipliers map[string]float64 `json:"reasoning_effort_multipliers"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
@@ -594,8 +610,7 @@ func TestGetModelDefaultPricing_ReturnsFable51CacheTTLs(t *testing.T) {
|
||||
require.InDelta(t, 12.5e-6, body.Data.CacheWritePrice, 1e-12)
|
||||
require.NotNil(t, body.Data.CacheWrite1hPrice)
|
||||
require.InDelta(t, 20e-6, *body.Data.CacheWrite1hPrice, 1e-12)
|
||||
require.NotNil(t, body.Data.MaxReasoningEffortMultiplier)
|
||||
require.Equal(t, 3.0, *body.Data.MaxReasoningEffortMultiplier)
|
||||
require.Empty(t, body.Data.ReasoningEffortMultipliers)
|
||||
}
|
||||
|
||||
func TestGetModelDefaultPricing_OmitsUnsupportedCache1hPrice(t *testing.T) {
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
type claudeResetReader interface {
|
||||
Query(context.Context, int64) (*service.ClaudeResetCredits, error)
|
||||
Redeem(context.Context, int64, string) (*service.ClaudeResetOutcome, error)
|
||||
}
|
||||
|
||||
func (h *AccountHandler) SetClaudeResetCreditService(s *service.ClaudeResetCreditService) {
|
||||
h.claudeResetCredits = s
|
||||
}
|
||||
func (h *AccountHandler) ClaudeResetCredits(c *gin.Context) {
|
||||
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
if h.claudeResetCredits == nil {
|
||||
response.Error(c, 503, "Claude reset service unavailable")
|
||||
return
|
||||
}
|
||||
status, err := h.claudeResetCredits.Query(c.Request.Context(), id)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, status)
|
||||
}
|
||||
|
||||
// RedeemClaudeResetCredit consumes the upstream next reset credit. The request has
|
||||
// no body: the server picks the grant; the Idempotency-Key header identifies one
|
||||
// operator confirmation and replays its outcome.
|
||||
func (h *AccountHandler) RedeemClaudeResetCredit(c *gin.Context) {
|
||||
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
if h.claudeResetCredits == nil {
|
||||
response.Error(c, 503, "Claude reset service unavailable")
|
||||
return
|
||||
}
|
||||
result, err := h.claudeResetCredits.Redeem(c.Request.Context(), id, c.GetHeader("Idempotency-Key"))
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, result)
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type claudeResetHandlerStub struct {
|
||||
key string
|
||||
calls int
|
||||
}
|
||||
|
||||
func (s *claudeResetHandlerStub) Query(context.Context, int64) (*service.ClaudeResetCredits, error) {
|
||||
return &service.ClaudeResetCredits{Credits: []service.ClaudeResetCredit{}}, nil
|
||||
}
|
||||
|
||||
func (s *claudeResetHandlerStub) Redeem(_ context.Context, _ int64, key string) (*service.ClaudeResetOutcome, error) {
|
||||
s.calls++
|
||||
s.key = key
|
||||
if key == "" {
|
||||
return nil, service.ErrIdempotencyKeyRequired
|
||||
}
|
||||
return &service.ClaudeResetOutcome{Outcome: "reset", Cleared: []string{"five_hour"}}, nil
|
||||
}
|
||||
|
||||
func TestClaudeResetHandlerRedeemContract(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
stub := &claudeResetHandlerStub{}
|
||||
h := &AccountHandler{claudeResetCredits: stub}
|
||||
r := gin.New()
|
||||
r.POST("/accounts/:id/claude/reset-credits/redeem", h.RedeemClaudeResetCredit)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/accounts/1/claude/reset-credits/redeem", nil)
|
||||
req.Header.Set("Idempotency-Key", "same-operation")
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, "same-operation", stub.key)
|
||||
require.Contains(t, w.Body.String(), `"outcome":"reset"`)
|
||||
|
||||
w = httptest.NewRecorder()
|
||||
r.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/accounts/1/claude/reset-credits/redeem", nil))
|
||||
require.Equal(t, http.StatusBadRequest, w.Code)
|
||||
require.Contains(t, w.Body.String(), "IDEMPOTENCY_KEY_REQUIRED")
|
||||
|
||||
w = httptest.NewRecorder()
|
||||
r.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/accounts/abc/claude/reset-credits/redeem", nil))
|
||||
require.Equal(t, http.StatusBadRequest, w.Code)
|
||||
require.Equal(t, 2, stub.calls)
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type moderationEngineSettings struct {
|
||||
service.SettingRepository
|
||||
value string
|
||||
}
|
||||
|
||||
func (r *moderationEngineSettings) GetValue(context.Context, string) (string, error) {
|
||||
return r.value, nil
|
||||
}
|
||||
func (r *moderationEngineSettings) Set(_ context.Context, _ string, value string) error {
|
||||
r.value = value
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestContentModerationEngineHandlerRoundTrip(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
settings := &moderationEngineSettings{value: `{"api_keys":["legacy-secret"],"model":"omni-moderation-latest"}`}
|
||||
svc := service.NewContentModerationService(settings, nil, nil, nil, nil, nil, nil, nil)
|
||||
h := NewContentModerationHandler(svc)
|
||||
router := gin.New()
|
||||
router.PUT("/config", h.UpdateConfig)
|
||||
router.GET("/config", h.GetConfig)
|
||||
body := `{"engine":"typesafe","engine_configs":{"typesafe":{"base_url":"https://api.typesafe.ai","model":"jev-latest","api_keys":["new-secret"],"thresholds":{"sexual":0.91}}}}`
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPut, "/config", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(w, req)
|
||||
require.Equal(t, 200, w.Code)
|
||||
require.Contains(t, w.Body.String(), `"engine":"typesafe"`)
|
||||
require.Contains(t, w.Body.String(), `"sexual":0.91`)
|
||||
require.NotContains(t, w.Body.String(), "new-secret")
|
||||
require.NotContains(t, w.Body.String(), "legacy-secret")
|
||||
w = httptest.NewRecorder()
|
||||
router.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/config", nil))
|
||||
require.Equal(t, 200, w.Code)
|
||||
require.Contains(t, w.Body.String(), `"engine_configs"`)
|
||||
require.Contains(t, w.Body.String(), `"model":"jev-latest"`)
|
||||
require.Contains(t, settings.value, "legacy-secret")
|
||||
require.Contains(t, settings.value, "new-secret")
|
||||
}
|
||||
@@ -20,10 +20,12 @@ func NewContentModerationHandler(svc *service.ContentModerationService) *Content
|
||||
}
|
||||
|
||||
type contentModerationConfigRequest struct {
|
||||
Enabled *bool `json:"enabled"`
|
||||
Mode *string `json:"mode"`
|
||||
BaseURL *string `json:"base_url"`
|
||||
Model *string `json:"model"`
|
||||
Engine *string `json:"engine"`
|
||||
EngineConfigs map[string]service.UpdateContentModerationEngineInput `json:"engine_configs"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
Mode *string `json:"mode"`
|
||||
BaseURL *string `json:"base_url"`
|
||||
Model *string `json:"model"`
|
||||
// 审计请求使用的代理服务器:null 不修改;0 清除(直连);>0 指定代理。
|
||||
ProxyID *int64 `json:"proxy_id"`
|
||||
APIKey *string `json:"api_key"`
|
||||
@@ -58,13 +60,15 @@ type contentModerationConfigRequest struct {
|
||||
}
|
||||
|
||||
type contentModerationAPIKeyTestRequest struct {
|
||||
APIKeys []string `json:"api_keys"`
|
||||
BaseURL string `json:"base_url"`
|
||||
Model string `json:"model"`
|
||||
TimeoutMS int `json:"timeout_ms"`
|
||||
ProxyID *int64 `json:"proxy_id"`
|
||||
Prompt string `json:"prompt"`
|
||||
Images []string `json:"images"`
|
||||
Engine string `json:"engine"`
|
||||
Thresholds *map[string]float64 `json:"thresholds"`
|
||||
APIKeys []string `json:"api_keys"`
|
||||
BaseURL string `json:"base_url"`
|
||||
Model string `json:"model"`
|
||||
TimeoutMS int `json:"timeout_ms"`
|
||||
ProxyID *int64 `json:"proxy_id"`
|
||||
Prompt string `json:"prompt"`
|
||||
Images []string `json:"images"`
|
||||
}
|
||||
|
||||
type contentModerationHashRequest struct {
|
||||
@@ -87,6 +91,7 @@ func (h *ContentModerationHandler) UpdateConfig(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
cfg, err := h.service.UpdateConfig(c.Request.Context(), service.UpdateContentModerationConfigInput{
|
||||
Engine: req.Engine, EngineConfigs: req.EngineConfigs,
|
||||
Enabled: req.Enabled,
|
||||
Mode: req.Mode,
|
||||
BaseURL: req.BaseURL,
|
||||
@@ -134,6 +139,7 @@ func (h *ContentModerationHandler) TestAPIKeys(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
result, err := h.service.TestAPIKeys(c.Request.Context(), service.TestContentModerationAPIKeysInput{
|
||||
Engine: req.Engine, Thresholds: req.Thresholds,
|
||||
APIKeys: req.APIKeys,
|
||||
BaseURL: req.BaseURL,
|
||||
Model: req.Model,
|
||||
|
||||
@@ -500,7 +500,11 @@ func (h *DashboardHandler) GetUserUsageTrend(c *gin.Context) {
|
||||
limit = 12
|
||||
}
|
||||
|
||||
trend, hit, err := h.getUserUsageTrendCached(c.Request.Context(), startTime, endTime, granularity, limit)
|
||||
metric := c.DefaultQuery("metric", "tokens")
|
||||
if metric != "tokens" && metric != "actual_cost" {
|
||||
metric = "tokens"
|
||||
}
|
||||
trend, hit, err := h.getUserUsageTrendCached(c.Request.Context(), startTime, endTime, granularity, limit, metric)
|
||||
if err != nil {
|
||||
response.Error(c, 500, "Failed to get user usage trend")
|
||||
return
|
||||
|
||||
@@ -45,6 +45,7 @@ func (r *dashboardUsageRepoCacheProbe) GetUserUsageTrend(
|
||||
startTime, endTime time.Time,
|
||||
granularity string,
|
||||
limit int,
|
||||
metric string,
|
||||
) ([]usagestats.UserUsageTrendPoint, error) {
|
||||
r.usersTrendCalls.Add(1)
|
||||
return []usagestats.UserUsageTrendPoint{{
|
||||
@@ -115,4 +116,20 @@ func TestDashboardHandler_GetUserUsageTrend_UsesCache(t *testing.T) {
|
||||
require.Equal(t, http.StatusOK, rec2.Code)
|
||||
require.Equal(t, "hit", rec2.Header().Get("X-Snapshot-Cache"))
|
||||
require.Equal(t, int32(1), repo.usersTrendCalls.Load())
|
||||
|
||||
for _, tc := range []struct {
|
||||
metric, cache string
|
||||
calls int32
|
||||
}{
|
||||
{"actual_cost", "miss", 2},
|
||||
{"actual_cost", "hit", 2},
|
||||
{"tokens", "hit", 2},
|
||||
} {
|
||||
req := httptest.NewRequest(http.MethodGet, "/admin/dashboard/users-trend?start_date=2026-03-01&end_date=2026-03-07&granularity=day&limit=8&metric="+tc.metric, nil)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, tc.cache, rec.Header().Get("X-Snapshot-Cache"))
|
||||
require.Equal(t, tc.calls, repo.usersTrendCalls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -53,6 +53,7 @@ type dashboardEntityTrendCacheKey struct {
|
||||
EndTime string `json:"end_time"`
|
||||
Granularity string `json:"granularity"`
|
||||
Limit int `json:"limit"`
|
||||
Metric string `json:"metric,omitempty"`
|
||||
}
|
||||
|
||||
func cacheStatusValue(hit bool) string {
|
||||
@@ -213,15 +214,16 @@ func (h *DashboardHandler) getAPIKeyUsageTrendCached(ctx context.Context, startT
|
||||
return trend, hit, err
|
||||
}
|
||||
|
||||
func (h *DashboardHandler) getUserUsageTrendCached(ctx context.Context, startTime, endTime time.Time, granularity string, limit int) ([]usagestats.UserUsageTrendPoint, bool, error) {
|
||||
func (h *DashboardHandler) getUserUsageTrendCached(ctx context.Context, startTime, endTime time.Time, granularity string, limit int, metric string) ([]usagestats.UserUsageTrendPoint, bool, error) {
|
||||
key := mustMarshalDashboardCacheKey(dashboardEntityTrendCacheKey{
|
||||
StartTime: startTime.UTC().Format(time.RFC3339),
|
||||
EndTime: endTime.UTC().Format(time.RFC3339),
|
||||
Granularity: granularity,
|
||||
Limit: limit,
|
||||
Metric: metric,
|
||||
})
|
||||
entry, hit, err := dashboardUsersTrendCache.GetOrLoad(key, func() (any, error) {
|
||||
return h.dashboardService.GetUserUsageTrend(ctx, startTime, endTime, granularity, limit)
|
||||
return h.dashboardService.GetUserUsageTrend(ctx, startTime, endTime, granularity, limit, metric)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, hit, err
|
||||
|
||||
@@ -243,7 +243,7 @@ func (h *DashboardHandler) buildSnapshotV2Response(
|
||||
}
|
||||
|
||||
if includeUsersTrend {
|
||||
usersTrend, _, err := h.getUserUsageTrendCached(ctx, startTime, endTime, granularity, usersTrendLimit)
|
||||
usersTrend, _, err := h.getUserUsageTrendCached(ctx, startTime, endTime, granularity, usersTrendLimit, "tokens")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to get user usage trend")
|
||||
}
|
||||
|
||||
@@ -21,12 +21,14 @@ type OpenAIOAuthHandler struct {
|
||||
openaiOAuthService *service.OpenAIOAuthService
|
||||
adminService service.AdminService
|
||||
quotaService openAIQuotaService
|
||||
referralService openAIReferralService
|
||||
rateLimitService openAIAccountStateRecoverer
|
||||
}
|
||||
|
||||
type openAIQuotaService interface {
|
||||
QueryUsage(ctx context.Context, accountID int64) (*service.OpenAIQuotaUsage, error)
|
||||
CacheResetCreditsSnapshot(ctx context.Context, accountID int64, credits *service.OpenAIRateLimitResetCredits) error
|
||||
CacheCreditsSnapshot(ctx context.Context, accountID int64, usage *service.OpenAIQuotaUsage) error
|
||||
CachePostResetSnapshot(ctx context.Context, accountID int64, usage *service.OpenAIQuotaUsage) error
|
||||
ResetCredit(ctx context.Context, accountID int64) (*service.OpenAIQuotaResetResult, error)
|
||||
}
|
||||
@@ -57,7 +59,8 @@ type openAIQuotaResetResponse struct {
|
||||
// failed display-cache write must never discard a successful upstream read.
|
||||
type openAIQuotaRefreshResponse struct {
|
||||
service.OpenAIQuotaUsage
|
||||
CachePersisted bool `json:"cache_persisted"`
|
||||
CachePersisted bool `json:"cache_persisted"`
|
||||
CreditsCachePersisted bool `json:"credits_cache_persisted"`
|
||||
}
|
||||
|
||||
// openAIQuotaResetPostProcessContext detaches the post-reset bookkeeping from the
|
||||
@@ -92,6 +95,7 @@ func NewOpenAIOAuthHandler(
|
||||
// `== nil` capability guards below and panic instead of returning 400.
|
||||
if quotaService != nil {
|
||||
h.quotaService = quotaService
|
||||
h.referralService = quotaService
|
||||
}
|
||||
if rateLimitService != nil {
|
||||
h.rateLimitService = rateLimitService
|
||||
@@ -522,6 +526,11 @@ func (h *OpenAIOAuthHandler) RefreshQuota(c *gin.Context) {
|
||||
service.NotifyOpenAIAutoResetCredit(accountID)
|
||||
|
||||
refreshResponse := openAIQuotaRefreshResponse{OpenAIQuotaUsage: *usage}
|
||||
if err := h.quotaService.CacheCreditsSnapshot(c.Request.Context(), accountID, usage); err != nil {
|
||||
slog.Warn("openai_quota_credits_cache_persist_failed", "account_id", accountID, "error", err)
|
||||
} else {
|
||||
refreshResponse.CreditsCachePersisted = true
|
||||
}
|
||||
// A failed snapshot write leaves the previous cache intact — report it as a
|
||||
// partial success instead of discarding the usage payload we just fetched,
|
||||
// which would leave the card without a credit count at all.
|
||||
|
||||
@@ -17,11 +17,14 @@ import (
|
||||
)
|
||||
|
||||
type openAIQuotaWorkflowStub struct {
|
||||
resetResult *service.OpenAIQuotaResetResult
|
||||
resetErr error
|
||||
queryResult *service.OpenAIQuotaUsage
|
||||
queryErr error
|
||||
cacheErr error
|
||||
resetResult *service.OpenAIQuotaResetResult
|
||||
resetErr error
|
||||
queryResult *service.OpenAIQuotaUsage
|
||||
queryErr error
|
||||
cacheErr error
|
||||
creditsCacheErr error
|
||||
creditsCacheCalls int
|
||||
cachedCreditsUsage *service.OpenAIQuotaUsage
|
||||
|
||||
resetCalls int
|
||||
queryCalls int
|
||||
@@ -48,6 +51,12 @@ func (s *openAIQuotaWorkflowStub) CacheResetCreditsSnapshot(ctx context.Context,
|
||||
return s.cacheErr
|
||||
}
|
||||
|
||||
func (s *openAIQuotaWorkflowStub) CacheCreditsSnapshot(_ context.Context, _ int64, usage *service.OpenAIQuotaUsage) error {
|
||||
s.creditsCacheCalls++
|
||||
s.cachedCreditsUsage = usage
|
||||
return s.creditsCacheErr
|
||||
}
|
||||
|
||||
func (s *openAIQuotaWorkflowStub) CachePostResetSnapshot(ctx context.Context, _ int64, _ *service.OpenAIQuotaUsage) error {
|
||||
s.cacheCalls++
|
||||
s.cacheCtxErr = ctx.Err()
|
||||
@@ -410,6 +419,34 @@ func TestOpenAIRefreshQuota_PersistFailureStillReturnsUsage(t *testing.T) {
|
||||
require.Equal(t, 1, quota.cacheCalls)
|
||||
}
|
||||
|
||||
func TestOpenAIRefreshQuota_CreditsPersistIndependently(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
resetErr error
|
||||
creditsErr error
|
||||
}{
|
||||
{name: "both saved"},
|
||||
{name: "reset details missing", resetErr: errors.New("missing expirations")},
|
||||
{name: "points cache failed", creditsErr: errors.New("write failed")},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
quota := successfulOpenAIQuotaWorkflowStub()
|
||||
balance := "1250.75"
|
||||
quota.queryResult.Credits = &service.OpenAICredits{HasCredits: true, Balance: &balance}
|
||||
quota.cacheErr = tc.resetErr
|
||||
quota.creditsCacheErr = tc.creditsErr
|
||||
status, envelope := performOpenAIQuotaRefreshRequest(t, &OpenAIOAuthHandler{quotaService: quota})
|
||||
require.Equal(t, http.StatusOK, status)
|
||||
require.Equal(t, tc.resetErr == nil, envelope.Data.CachePersisted)
|
||||
require.Equal(t, tc.creditsErr == nil, envelope.Data.CreditsCachePersisted)
|
||||
require.Equal(t, quota.queryResult.Credits, envelope.Data.Credits)
|
||||
require.Equal(t, quota.queryResult, quota.cachedCreditsUsage)
|
||||
require.Equal(t, 1, quota.creditsCacheCalls)
|
||||
require.Zero(t, quota.resetCalls)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// An empty-but-successful upstream read must not be dereferenced blindly.
|
||||
func TestOpenAIQuotaEmptyUsageIsHandledWithoutPanic(t *testing.T) {
|
||||
t.Run("refresh reports an internal error", func(t *testing.T) {
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type openAIReferralService interface {
|
||||
QueryReferralEligibility(context.Context, int64) (*service.OpenAIReferralEligibility, error)
|
||||
CacheReferralSnapshot(context.Context, int64, *service.OpenAIReferralEligibility) error
|
||||
SendReferralInvite(context.Context, int64, service.OpenAIReferralSendRequest) (*service.OpenAIReferralSendResult, error)
|
||||
}
|
||||
|
||||
type openAIReferralRefreshResponse struct {
|
||||
Eligibility *service.OpenAIReferralEligibility `json:"eligibility"`
|
||||
CachePersisted bool `json:"cache_persisted"`
|
||||
}
|
||||
|
||||
type openAIReferralSendResponse struct {
|
||||
service.OpenAIReferralSendResult
|
||||
openAIReferralRefreshResponse
|
||||
RefreshFailed bool `json:"refresh_failed"`
|
||||
}
|
||||
|
||||
func (h *OpenAIOAuthHandler) referralAccountID(c *gin.Context) (int64, bool) {
|
||||
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return 0, false
|
||||
}
|
||||
if h.referralService == nil {
|
||||
response.BadRequest(c, "OpenAI referral service is not enabled")
|
||||
return 0, false
|
||||
}
|
||||
return id, true
|
||||
}
|
||||
|
||||
// RefreshReferrals persists a display snapshot, hence POST and admin audit.
|
||||
func (h *OpenAIOAuthHandler) RefreshReferrals(c *gin.Context) {
|
||||
id, ok := h.referralAccountID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
eligibility, err := h.referralService.QueryReferralEligibility(c.Request.Context(), id)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
if eligibility == nil {
|
||||
response.Error(c, http.StatusBadGateway, "Empty invitation eligibility response")
|
||||
return
|
||||
}
|
||||
cacheErr := h.referralService.CacheReferralSnapshot(c.Request.Context(), id, eligibility)
|
||||
response.Success(c, openAIReferralRefreshResponse{Eligibility: eligibility, CachePersisted: cacheErr == nil})
|
||||
}
|
||||
|
||||
func (h *OpenAIOAuthHandler) SendReferralInvite(c *gin.Context) {
|
||||
id, ok := h.referralAccountID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input service.OpenAIReferralSendRequest
|
||||
if err := c.ShouldBindJSON(&input); err != nil {
|
||||
response.BadRequest(c, "Invalid invitation request")
|
||||
return
|
||||
}
|
||||
result, err := h.referralService.SendReferralInvite(c.Request.Context(), id, input)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
if result == nil || !result.Sent {
|
||||
response.Error(c, http.StatusBadGateway, "Invitation outcome is unknown; check Codex before sending again")
|
||||
return
|
||||
}
|
||||
// The email is already sent. Refresh failure must not turn it into a failed
|
||||
// submission and encourage a duplicate send, even if the browser disconnects.
|
||||
baseCtx := context.WithoutCancel(c.Request.Context())
|
||||
refreshCtx, cancelRefresh := context.WithTimeout(baseCtx, 8*time.Second)
|
||||
eligibility, refreshErr := h.referralService.QueryReferralEligibility(refreshCtx, id)
|
||||
cancelRefresh()
|
||||
if refreshErr != nil {
|
||||
eligibility = nil
|
||||
}
|
||||
// Refresh may exhaust its entire deadline. Give persistence (including
|
||||
// invalidating a stale snapshot) a fresh, independent deadline.
|
||||
cacheCtx, cancelCache := context.WithTimeout(baseCtx, 3*time.Second)
|
||||
defer cancelCache()
|
||||
cacheErr := h.referralService.CacheReferralSnapshot(cacheCtx, id, eligibility)
|
||||
response.Success(c, openAIReferralSendResponse{
|
||||
OpenAIReferralSendResult: *result,
|
||||
openAIReferralRefreshResponse: openAIReferralRefreshResponse{Eligibility: eligibility, CachePersisted: cacheErr == nil},
|
||||
RefreshFailed: eligibility == nil,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
//go:build unit
|
||||
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"testing/synctest"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type referralHandlerStub struct {
|
||||
query func(context.Context) (*service.OpenAIReferralEligibility, error)
|
||||
cache func(context.Context, *service.OpenAIReferralEligibility) error
|
||||
eligibility *service.OpenAIReferralEligibility
|
||||
queryErr, cacheErr error
|
||||
input service.OpenAIReferralSendRequest
|
||||
cached *service.OpenAIReferralEligibility
|
||||
sends int
|
||||
}
|
||||
|
||||
func (s *referralHandlerStub) QueryReferralEligibility(ctx context.Context, _ int64) (*service.OpenAIReferralEligibility, error) {
|
||||
if s.query != nil {
|
||||
return s.query(ctx)
|
||||
}
|
||||
return s.eligibility, s.queryErr
|
||||
}
|
||||
func (s *referralHandlerStub) CacheReferralSnapshot(ctx context.Context, _ int64, e *service.OpenAIReferralEligibility) error {
|
||||
if s.cache != nil {
|
||||
return s.cache(ctx, e)
|
||||
}
|
||||
s.cached = e
|
||||
return s.cacheErr
|
||||
}
|
||||
func (s *referralHandlerStub) SendReferralInvite(_ context.Context, _ int64, input service.OpenAIReferralSendRequest) (*service.OpenAIReferralSendResult, error) {
|
||||
s.sends++
|
||||
s.input = input
|
||||
return &service.OpenAIReferralSendResult{Email: input.Email, Sent: true}, nil
|
||||
}
|
||||
|
||||
func referralHandlerRequest(t *testing.T, stub *referralHandlerStub, path, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
h := &OpenAIOAuthHandler{referralService: stub}
|
||||
router := gin.New()
|
||||
router.POST("/:id/referrals/refresh", h.RefreshReferrals)
|
||||
router.POST("/:id/referrals/invite", h.SendReferralInvite)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, path, strings.NewReader(body)))
|
||||
return rec
|
||||
}
|
||||
|
||||
func TestOpenAIReferralHandlerRefresh(t *testing.T) {
|
||||
count := 2
|
||||
stub := &referralHandlerStub{eligibility: &service.OpenAIReferralEligibility{ShouldShow: true, AvailableInvites: &count}, cacheErr: errors.New("cache failed")}
|
||||
rec := referralHandlerRequest(t, stub, "/100/referrals/refresh", "")
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var body struct {
|
||||
Data openAIReferralRefreshResponse `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body))
|
||||
require.Equal(t, 2, *body.Data.Eligibility.AvailableInvites)
|
||||
require.False(t, body.Data.CachePersisted)
|
||||
require.Zero(t, stub.sends)
|
||||
}
|
||||
|
||||
func TestOpenAIReferralHandlerSentSurvivesRefreshFailure(t *testing.T) {
|
||||
stub := &referralHandlerStub{queryErr: errors.New("upstream timed out")}
|
||||
rec := referralHandlerRequest(t, stub, "/100/referrals/invite", `{"email":"friend@example.com","program_id":"codex_referral_consumer","confirmed":true}`)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var body struct {
|
||||
Data openAIReferralSendResponse `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body))
|
||||
require.True(t, body.Data.Sent)
|
||||
require.Equal(t, "friend@example.com", body.Data.Email)
|
||||
require.True(t, body.Data.RefreshFailed)
|
||||
require.True(t, body.Data.CachePersisted)
|
||||
require.Nil(t, stub.cached, "stale remaining count must be invalidated")
|
||||
require.True(t, stub.input.Confirmed)
|
||||
require.Equal(t, 1, stub.sends)
|
||||
}
|
||||
|
||||
func TestOpenAIReferralHandlerRejectsInvalidInput(t *testing.T) {
|
||||
stub := &referralHandlerStub{}
|
||||
for _, path := range []string{"/invalid/referrals/invite", "/0/referrals/invite", "/100/referrals/invite"} {
|
||||
rec := referralHandlerRequest(t, stub, path, "not-json")
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
}
|
||||
require.Zero(t, stub.sends)
|
||||
}
|
||||
|
||||
func TestOpenAIReferralHandlerInvalidatesCacheAfterRefreshDeadline(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
cached := false
|
||||
stub := &referralHandlerStub{
|
||||
query: func(ctx context.Context) (*service.OpenAIReferralEligibility, error) {
|
||||
<-ctx.Done()
|
||||
require.ErrorIs(t, ctx.Err(), context.DeadlineExceeded)
|
||||
return nil, ctx.Err()
|
||||
},
|
||||
cache: func(ctx context.Context, e *service.OpenAIReferralEligibility) error {
|
||||
require.NoError(t, ctx.Err(), "cache invalidation needs a fresh deadline")
|
||||
deadline, ok := ctx.Deadline()
|
||||
require.True(t, ok)
|
||||
require.Equal(t, 3*time.Second, time.Until(deadline))
|
||||
require.Nil(t, e, "expired count must be cleared")
|
||||
cached = true
|
||||
return nil
|
||||
},
|
||||
}
|
||||
rec := referralHandlerRequest(t, stub, "/100/referrals/invite", `{"email":"friend@example.com","program_id":"codex_referral_consumer","confirmed":true}`)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var body struct {
|
||||
Data openAIReferralSendResponse `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body))
|
||||
require.True(t, cached)
|
||||
require.True(t, body.Data.Sent)
|
||||
require.True(t, body.Data.RefreshFailed)
|
||||
require.True(t, body.Data.CachePersisted)
|
||||
require.Equal(t, 1, stub.sends)
|
||||
})
|
||||
}
|
||||
@@ -265,6 +265,7 @@ func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
DefaultBalance: settings.DefaultBalance,
|
||||
RiskControlEnabled: settings.RiskControlEnabled,
|
||||
CyberSessionBlockEnabled: settings.CyberSessionBlockEnabled,
|
||||
CyberPolicyUserAllowlist: settings.CyberPolicyUserAllowlist,
|
||||
CyberSessionBlockTTLSeconds: settings.CyberSessionBlockTTLSeconds,
|
||||
AffiliateRebateRate: settings.AffiliateRebateRate,
|
||||
AffiliateRebateFreezeHours: settings.AffiliateRebateFreezeHours,
|
||||
@@ -303,6 +304,9 @@ func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
OpenAICodexClientVersion: settings.OpenAICodexClientVersion,
|
||||
OpenAICodexClientVersionSynced: settings.OpenAICodexClientVersionSynced,
|
||||
OpenAICodexVersionAutoSyncEnabled: settings.OpenAICodexVersionAutoSyncEnabled,
|
||||
ClaudeCodeClientVersion: settings.ClaudeCodeClientVersion,
|
||||
ClaudeCodeClientVersionSynced: settings.ClaudeCodeClientVersionSynced,
|
||||
ClaudeCodeVersionAutoSyncEnabled: settings.ClaudeCodeVersionAutoSyncEnabled,
|
||||
MinCodexVersion: settings.MinCodexVersion,
|
||||
MaxCodexVersion: settings.MaxCodexVersion,
|
||||
CodexCLIOnlyBlacklist: settings.CodexCLIOnlyBlacklist,
|
||||
|
||||
@@ -491,6 +491,12 @@ func diffSettings(before *service.SystemSettings, after *service.SystemSettings,
|
||||
if before.OpenAICodexVersionAutoSyncEnabled != after.OpenAICodexVersionAutoSyncEnabled {
|
||||
changed = append(changed, "openai_codex_version_auto_sync_enabled")
|
||||
}
|
||||
if before.ClaudeCodeClientVersion != after.ClaudeCodeClientVersion {
|
||||
changed = append(changed, "claude_code_client_version")
|
||||
}
|
||||
if before.ClaudeCodeVersionAutoSyncEnabled != after.ClaudeCodeVersionAutoSyncEnabled {
|
||||
changed = append(changed, "claude_code_version_auto_sync_enabled")
|
||||
}
|
||||
if before.PaymentVisibleMethodAlipaySource != after.PaymentVisibleMethodAlipaySource {
|
||||
changed = append(changed, "payment_visible_method_alipay_source")
|
||||
}
|
||||
@@ -506,7 +512,7 @@ func diffSettings(before *service.SystemSettings, after *service.SystemSettings,
|
||||
if before.OpenAILowUpstreamRatePriorityEnabled != after.OpenAILowUpstreamRatePriorityEnabled {
|
||||
changed = append(changed, "openai_low_upstream_rate_priority_enabled")
|
||||
}
|
||||
if before.OpenAIOAuthSchedulingRateMultiplier != after.OpenAIOAuthSchedulingRateMultiplier {
|
||||
if !equalNullableFloat(before.OpenAIOAuthSchedulingRateMultiplier, after.OpenAIOAuthSchedulingRateMultiplier) {
|
||||
changed = append(changed, "openai_oauth_scheduling_rate_multiplier")
|
||||
}
|
||||
if before.OpenAIAdvancedSchedulerEnabled != after.OpenAIAdvancedSchedulerEnabled {
|
||||
@@ -597,6 +603,9 @@ func diffSettings(before *service.SystemSettings, after *service.SystemSettings,
|
||||
if before.RiskControlEnabled != after.RiskControlEnabled {
|
||||
changed = append(changed, "risk_control_enabled")
|
||||
}
|
||||
if before.CyberPolicyUserAllowlist != after.CyberPolicyUserAllowlist {
|
||||
changed = append(changed, "cyber_policy_user_allowlist")
|
||||
}
|
||||
if before.CyberSessionBlockEnabled != after.CyberSessionBlockEnabled {
|
||||
changed = append(changed, "cyber_session_block_enabled")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSettingHandler_OAuthSchedulingRateRoundTrip(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
initial map[string]string
|
||||
body string
|
||||
stored string
|
||||
value any
|
||||
status int
|
||||
}{
|
||||
{"absent keeps default", nil, `{}`, "1", float64(1), http.StatusOK},
|
||||
{"omitted preserves override", map[string]string{service.SettingKeyOpenAIOAuthSchedulingRateMultiplier: "0.7"}, `{}`, "0.7", 0.7, http.StatusOK},
|
||||
{"clear override", map[string]string{service.SettingKeyOpenAIOAuthSchedulingRateMultiplier: "0.7"}, `{"openai_oauth_scheduling_rate_multiplier":null}`, "", nil, http.StatusOK},
|
||||
{"omitted preserves cleared", map[string]string{service.SettingKeyOpenAIOAuthSchedulingRateMultiplier: ""}, `{}`, "", nil, http.StatusOK},
|
||||
{"set override", nil, `{"openai_oauth_scheduling_rate_multiplier":0.7}`, "0.7", 0.7, http.StatusOK},
|
||||
{"zero is explicit", nil, `{"openai_oauth_scheduling_rate_multiplier":0}`, "0", float64(0), http.StatusOK},
|
||||
{"negative is rejected", nil, `{"openai_oauth_scheduling_rate_multiplier":-1}`, "", nil, http.StatusBadRequest},
|
||||
{"string is rejected", nil, `{"openai_oauth_scheduling_rate_multiplier":"invalid"}`, "", nil, http.StatusBadRequest},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
repo := &settingHandlerRepoStub{values: tt.initial}
|
||||
svc := service.NewSettingService(repo, &config.Config{})
|
||||
handler := NewSettingHandler(svc, nil, nil, nil, nil, nil, nil)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPut, "/api/v1/admin/settings", bytes.NewBufferString(tt.body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
handler.UpdateSettings(c)
|
||||
require.Equal(t, tt.status, rec.Code, rec.Body.String())
|
||||
if tt.status != http.StatusOK {
|
||||
require.Empty(t, repo.lastUpdates)
|
||||
return
|
||||
}
|
||||
require.Equal(t, tt.stored, repo.values[service.SettingKeyOpenAIOAuthSchedulingRateMultiplier])
|
||||
var body struct {
|
||||
Data map[string]any `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body))
|
||||
require.Contains(t, body.Data, service.SettingKeyOpenAIOAuthSchedulingRateMultiplier)
|
||||
require.Equal(t, tt.value, body.Data[service.SettingKeyOpenAIOAuthSchedulingRateMultiplier])
|
||||
|
||||
getRec := httptest.NewRecorder()
|
||||
getContext, _ := gin.CreateTestContext(getRec)
|
||||
getContext.Request = httptest.NewRequest(http.MethodGet, "/api/v1/admin/settings", nil)
|
||||
handler.GetSettings(getContext)
|
||||
require.Equal(t, http.StatusOK, getRec.Code, getRec.Body.String())
|
||||
require.NoError(t, json.Unmarshal(getRec.Body.Bytes(), &body))
|
||||
require.Contains(t, body.Data, service.SettingKeyOpenAIOAuthSchedulingRateMultiplier)
|
||||
require.Equal(t, tt.value, body.Data[service.SettingKeyOpenAIOAuthSchedulingRateMultiplier])
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSettingsAuditTracksOAuthRateValueNotPointer(t *testing.T) {
|
||||
beforeRate, afterRate := 0.7, 0.7
|
||||
before := &service.SystemSettings{OpenAIOAuthSchedulingRateMultiplier: &beforeRate}
|
||||
after := &service.SystemSettings{OpenAIOAuthSchedulingRateMultiplier: &afterRate}
|
||||
require.NotContains(t, diffSettings(before, after, nil, nil, UpdateSettingsRequest{}), service.SettingKeyOpenAIOAuthSchedulingRateMultiplier)
|
||||
after.OpenAIOAuthSchedulingRateMultiplier = nil
|
||||
require.Contains(t, diffSettings(before, after, nil, nil, UpdateSettingsRequest{}), service.SettingKeyOpenAIOAuthSchedulingRateMultiplier)
|
||||
}
|
||||
@@ -257,6 +257,8 @@ type UpdateSettingsRequest struct {
|
||||
OpenAICodexUserAgent *string `json:"openai_codex_user_agent"`
|
||||
OpenAICodexClientVersion *string `json:"openai_codex_client_version"`
|
||||
OpenAICodexVersionAutoSyncEnabled *bool `json:"openai_codex_version_auto_sync_enabled"`
|
||||
ClaudeCodeClientVersion *string `json:"claude_code_client_version"`
|
||||
ClaudeCodeVersionAutoSyncEnabled *bool `json:"claude_code_version_auto_sync_enabled"`
|
||||
|
||||
// codex_cli_only 加固(global-only)
|
||||
MinCodexVersion string `json:"min_codex_version"`
|
||||
@@ -362,8 +364,9 @@ type UpdateSettingsRequest struct {
|
||||
RiskControlEnabled *bool `json:"risk_control_enabled"`
|
||||
|
||||
// cyber 会话屏蔽开关 + TTL
|
||||
CyberSessionBlockEnabled *bool `json:"cyber_session_block_enabled"`
|
||||
CyberSessionBlockTTLSeconds *int `json:"cyber_session_block_ttl_seconds"`
|
||||
CyberSessionBlockEnabled *bool `json:"cyber_session_block_enabled"`
|
||||
CyberPolicyUserAllowlist *string `json:"cyber_policy_user_allowlist"`
|
||||
CyberSessionBlockTTLSeconds *int `json:"cyber_session_block_ttl_seconds"`
|
||||
|
||||
// OpenAI fast/flex policy (optional, only updated when provided)
|
||||
OpenAIFastPolicySettings *dto.OpenAIFastPolicySettings `json:"openai_fast_policy_settings,omitempty"`
|
||||
@@ -1455,6 +1458,15 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
}
|
||||
req.OpenAICodexClientVersion = &normalized
|
||||
}
|
||||
if req.ClaudeCodeClientVersion != nil {
|
||||
// 该值会被拼进出站 User-Agent 与 billing attribution,必须是合法版本号;空串表示跟随自动同步。
|
||||
normalized := strings.TrimSpace(*req.ClaudeCodeClientVersion)
|
||||
if normalized != "" && service.NormalizeClaudeCodeClientVersion(normalized) == "" {
|
||||
response.Error(c, http.StatusBadRequest, "claude_code_client_version must be empty or a valid version (e.g. 2.1.258)")
|
||||
return
|
||||
}
|
||||
req.ClaudeCodeClientVersion = &normalized
|
||||
}
|
||||
|
||||
// codex_cli_only 加固:最低/最高 Codex 版本(空=禁用,或合法 semver;max>=min)
|
||||
if req.MinCodexVersion != "" && !semverPattern.MatchString(req.MinCodexVersion) {
|
||||
@@ -1492,6 +1504,13 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
if req.CyberPolicyUserAllowlist != nil {
|
||||
if _, err := service.ParseCyberPolicyUserAllowlist(*req.CyberPolicyUserAllowlist); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// cyber 会话屏蔽 TTL 校验:提供时必须 > 0
|
||||
if req.CyberSessionBlockTTLSeconds != nil && *req.CyberSessionBlockTTLSeconds <= 0 {
|
||||
response.BadRequest(c, "cyber_session_block_ttl_seconds must be > 0")
|
||||
@@ -1768,6 +1787,20 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
}
|
||||
return previousSettings.OpenAICodexVersionAutoSyncEnabled
|
||||
}(),
|
||||
ClaudeCodeClientVersion: func() string {
|
||||
if req.ClaudeCodeClientVersion != nil {
|
||||
return *req.ClaudeCodeClientVersion
|
||||
}
|
||||
return previousSettings.ClaudeCodeClientVersion
|
||||
}(),
|
||||
// 同步值由自动同步任务独占写入,面板保存时原样带回,避免被清空。
|
||||
ClaudeCodeClientVersionSynced: previousSettings.ClaudeCodeClientVersionSynced,
|
||||
ClaudeCodeVersionAutoSyncEnabled: func() bool {
|
||||
if req.ClaudeCodeVersionAutoSyncEnabled != nil {
|
||||
return *req.ClaudeCodeVersionAutoSyncEnabled
|
||||
}
|
||||
return previousSettings.ClaudeCodeVersionAutoSyncEnabled
|
||||
}(),
|
||||
MinCodexVersion: strings.TrimSpace(req.MinCodexVersion),
|
||||
MaxCodexVersion: strings.TrimSpace(req.MaxCodexVersion),
|
||||
CodexCLIOnlyBlacklist: strings.TrimSpace(req.CodexCLIOnlyBlacklist),
|
||||
@@ -1809,9 +1842,10 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
}
|
||||
return previousSettings.OpenAILowUpstreamRatePriorityEnabled
|
||||
}(),
|
||||
OpenAIOAuthSchedulingRateMultiplier: func() float64 {
|
||||
if req.OpenAIOAuthSchedulingRateMultiplier != nil {
|
||||
return *req.OpenAIOAuthSchedulingRateMultiplier
|
||||
OpenAIOAuthSchedulingRateMultiplier: func() *float64 {
|
||||
// Omitted fields preserve the override; explicit null clears it.
|
||||
if _, sent := sentFields[service.SettingKeyOpenAIOAuthSchedulingRateMultiplier]; sent {
|
||||
return req.OpenAIOAuthSchedulingRateMultiplier
|
||||
}
|
||||
return previousSettings.OpenAIOAuthSchedulingRateMultiplier
|
||||
}(),
|
||||
@@ -1982,6 +2016,12 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
}
|
||||
return previousSettings.RiskControlEnabled
|
||||
}(),
|
||||
CyberPolicyUserAllowlist: func() string {
|
||||
if req.CyberPolicyUserAllowlist != nil {
|
||||
return *req.CyberPolicyUserAllowlist
|
||||
}
|
||||
return previousSettings.CyberPolicyUserAllowlist
|
||||
}(),
|
||||
CyberSessionBlockEnabled: func() bool {
|
||||
if req.CyberSessionBlockEnabled != nil {
|
||||
return *req.CyberSessionBlockEnabled
|
||||
@@ -2310,6 +2350,9 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
OpenAICodexClientVersion: updatedSettings.OpenAICodexClientVersion,
|
||||
OpenAICodexClientVersionSynced: updatedSettings.OpenAICodexClientVersionSynced,
|
||||
OpenAICodexVersionAutoSyncEnabled: updatedSettings.OpenAICodexVersionAutoSyncEnabled,
|
||||
ClaudeCodeClientVersion: updatedSettings.ClaudeCodeClientVersion,
|
||||
ClaudeCodeClientVersionSynced: updatedSettings.ClaudeCodeClientVersionSynced,
|
||||
ClaudeCodeVersionAutoSyncEnabled: updatedSettings.ClaudeCodeVersionAutoSyncEnabled,
|
||||
MinCodexVersion: updatedSettings.MinCodexVersion,
|
||||
MaxCodexVersion: updatedSettings.MaxCodexVersion,
|
||||
CodexCLIOnlyBlacklist: updatedSettings.CodexCLIOnlyBlacklist,
|
||||
@@ -2400,6 +2443,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
|
||||
RiskControlEnabled: updatedSettings.RiskControlEnabled,
|
||||
CyberSessionBlockEnabled: updatedSettings.CyberSessionBlockEnabled,
|
||||
CyberPolicyUserAllowlist: updatedSettings.CyberPolicyUserAllowlist,
|
||||
CyberSessionBlockTTLSeconds: updatedSettings.CyberSessionBlockTTLSeconds,
|
||||
AccountSchedulingThresholds: updatedSettings.AccountSchedulingThresholds,
|
||||
AllowUserViewErrorRequests: updatedSettings.AllowUserViewErrorRequests,
|
||||
|
||||
@@ -407,6 +407,10 @@ func (r *oauthEmailAffiliateRepoStub) TransferQuotaToBalance(context.Context, in
|
||||
panic("unexpected TransferQuotaToBalance call")
|
||||
}
|
||||
|
||||
func (r *oauthEmailAffiliateRepoStub) WithdrawQuota(context.Context, int64, float64, string) (*service.AffiliateWithdrawResult, error) {
|
||||
panic("unexpected WithdrawQuota call")
|
||||
}
|
||||
|
||||
func (r *oauthEmailAffiliateRepoStub) ListInvitees(context.Context, int64, int) ([]service.AffiliateInvitee, error) {
|
||||
panic("unexpected ListInvitees call")
|
||||
}
|
||||
|
||||
@@ -66,17 +66,17 @@ type userAvailableGroup struct {
|
||||
|
||||
// userSupportedModelPricing 用户可见的定价字段白名单。
|
||||
type userSupportedModelPricing struct {
|
||||
BillingMode string `json:"billing_mode"`
|
||||
InputPrice *float64 `json:"input_price"`
|
||||
OutputPrice *float64 `json:"output_price"`
|
||||
CacheWritePrice *float64 `json:"cache_write_price"`
|
||||
CacheWrite1hPrice *float64 `json:"cache_write_1h_price"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price"`
|
||||
MaxReasoningEffortMultiplier *float64 `json:"max_reasoning_effort_multiplier,omitempty"`
|
||||
ImageInputPrice *float64 `json:"image_input_price"`
|
||||
ImageOutputPrice *float64 `json:"image_output_price"`
|
||||
PerRequestPrice *float64 `json:"per_request_price"`
|
||||
Intervals []userPricingIntervalDTO `json:"intervals"`
|
||||
BillingMode string `json:"billing_mode"`
|
||||
InputPrice *float64 `json:"input_price"`
|
||||
OutputPrice *float64 `json:"output_price"`
|
||||
CacheWritePrice *float64 `json:"cache_write_price"`
|
||||
CacheWrite1hPrice *float64 `json:"cache_write_1h_price"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price"`
|
||||
ReasoningEffortMultipliers map[string]float64 `json:"reasoning_effort_multipliers,omitempty"`
|
||||
ImageInputPrice *float64 `json:"image_input_price"`
|
||||
ImageOutputPrice *float64 `json:"image_output_price"`
|
||||
PerRequestPrice *float64 `json:"per_request_price"`
|
||||
Intervals []userPricingIntervalDTO `json:"intervals"`
|
||||
}
|
||||
|
||||
// userPricingIntervalDTO 定价区间白名单(去掉内部 ID、SortOrder 等前端不渲染的字段)。
|
||||
@@ -332,16 +332,16 @@ func toUserPricing(p *service.ChannelModelPricing) *userSupportedModelPricing {
|
||||
billingMode = string(service.BillingModeToken)
|
||||
}
|
||||
return &userSupportedModelPricing{
|
||||
BillingMode: billingMode,
|
||||
InputPrice: p.InputPrice,
|
||||
OutputPrice: p.OutputPrice,
|
||||
CacheWritePrice: p.CacheWritePrice,
|
||||
CacheWrite1hPrice: p.CacheWrite1hPrice,
|
||||
CacheReadPrice: p.CacheReadPrice,
|
||||
MaxReasoningEffortMultiplier: p.MaxReasoningEffortMultiplier,
|
||||
ImageInputPrice: p.ImageInputPrice,
|
||||
ImageOutputPrice: p.ImageOutputPrice,
|
||||
PerRequestPrice: p.PerRequestPrice,
|
||||
Intervals: intervals,
|
||||
BillingMode: billingMode,
|
||||
InputPrice: p.InputPrice,
|
||||
OutputPrice: p.OutputPrice,
|
||||
CacheWritePrice: p.CacheWritePrice,
|
||||
CacheWrite1hPrice: p.CacheWrite1hPrice,
|
||||
CacheReadPrice: p.CacheReadPrice,
|
||||
ReasoningEffortMultipliers: p.ReasoningEffortMultipliers,
|
||||
ImageInputPrice: p.ImageInputPrice,
|
||||
ImageOutputPrice: p.ImageOutputPrice,
|
||||
PerRequestPrice: p.PerRequestPrice,
|
||||
Intervals: intervals,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -118,7 +118,8 @@ func TestUserAvailableChannel_FieldWhitelist(t *testing.T) {
|
||||
cacheWriteMultiplier := 2.0
|
||||
cacheReadMultiplier := 2.0
|
||||
pricing := toUserPricing(&service.ChannelModelPricing{
|
||||
BillingMode: service.BillingModeToken,
|
||||
BillingMode: service.BillingModeToken,
|
||||
ReasoningEffortMultipliers: map[string]float64{"high": 1.5, "max": 3},
|
||||
Intervals: []service.PricingInterval{
|
||||
{
|
||||
ID: 7, MinTokens: 0, MaxTokens: nil, SortOrder: 3,
|
||||
@@ -128,6 +129,7 @@ func TestUserAvailableChannel_FieldWhitelist(t *testing.T) {
|
||||
},
|
||||
})
|
||||
require.NotNil(t, pricing)
|
||||
require.Equal(t, map[string]float64{"high": 1.5, "max": 3}, pricing.ReasoningEffortMultipliers)
|
||||
require.Len(t, pricing.Intervals, 1)
|
||||
rawIv, err := json.Marshal(pricing.Intervals[0])
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -2,10 +2,13 @@ package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
@@ -73,3 +76,65 @@ func TestConcurrencyErrorResponse(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGoogleConcurrencyError(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
err error
|
||||
wantStatus int
|
||||
wantGStatus string
|
||||
wantMessage string
|
||||
}{
|
||||
{
|
||||
name: "client cancellation is 499",
|
||||
err: context.Canceled,
|
||||
wantStatus: statusClientClosedRequest,
|
||||
wantGStatus: "CANCELLED",
|
||||
wantMessage: "context canceled",
|
||||
},
|
||||
{
|
||||
name: "full wait queue is 429",
|
||||
err: &WaitQueueFullError{SlotType: "user"},
|
||||
wantStatus: http.StatusTooManyRequests,
|
||||
wantGStatus: "RESOURCE_EXHAUSTED",
|
||||
wantMessage: "Too many pending requests, please retry later",
|
||||
},
|
||||
{
|
||||
name: "slot wait timeout is 429",
|
||||
err: &ConcurrencyError{SlotType: "user", IsTimeout: true},
|
||||
wantStatus: http.StatusTooManyRequests,
|
||||
wantGStatus: "RESOURCE_EXHAUSTED",
|
||||
wantMessage: "Concurrency limit exceeded for user, please retry later",
|
||||
},
|
||||
{
|
||||
name: "acquire backend error is 503",
|
||||
err: errors.New("redis unavailable"),
|
||||
wantStatus: http.StatusServiceUnavailable,
|
||||
wantGStatus: "INTERNAL",
|
||||
wantMessage: "Service temporarily unavailable, please retry later",
|
||||
},
|
||||
}
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
|
||||
googleConcurrencyError(c, tt.err, "user")
|
||||
|
||||
require.Equal(t, tt.wantStatus, rec.Code)
|
||||
var body struct {
|
||||
Error struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Status string `json:"status"`
|
||||
} `json:"error"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body))
|
||||
require.Equal(t, tt.wantStatus, body.Error.Code)
|
||||
require.Equal(t, tt.wantGStatus, body.Error.Status)
|
||||
require.Equal(t, tt.wantMessage, body.Error.Message)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
package dto
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// 账号列表页固定传 lite=1,lite 投影(AccountListItemFromAccount)必须与完整
|
||||
// Account DTO 一样携带 opencode_go_usage,否则前端用量单元格会退回「余额 --」。
|
||||
func TestAccountListItemFromAccount_CarriesOpenCodeGoUsage(t *testing.T) {
|
||||
now := time.Now().UTC()
|
||||
// 生产账号 164 的形态:deepseek 平台 + OpenCode Go 官方基址 + apikey
|
||||
eligible := &service.Account{
|
||||
ID: 164, Name: "DeepSeek OpenCode", Platform: service.PlatformDeepseek, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"base_url": "https://opencode.ai/zen/go/v1", "api_key": "test-key"},
|
||||
Extra: map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: true,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: &service.OpenCodeGoUsageSnapshot{
|
||||
Status: service.OpenCodeGoUsageStatusOK,
|
||||
Data: &service.OpenCodeGoUsageData{Rolling: service.OpenCodeGoUsageWindow{Status: "ok", Percent: 6}},
|
||||
LastAttemptAt: now, NextRefreshAt: now.Add(time.Hour),
|
||||
},
|
||||
},
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
|
||||
item := AccountListItemFromAccount(AccountFromServiceShallow(eligible))
|
||||
require.NotNil(t, item.OpenCodeGoUsage)
|
||||
require.True(t, item.OpenCodeGoUsage.Eligible)
|
||||
require.True(t, item.OpenCodeGoUsage.AutoRefreshEnabled)
|
||||
require.NotNil(t, item.OpenCodeGoUsage.Snapshot)
|
||||
require.NotNil(t, item.OpenCodeGoUsage.Snapshot.Data)
|
||||
require.Equal(t, 6.0, item.OpenCodeGoUsage.Snapshot.Data.Rolling.Percent)
|
||||
|
||||
raw, err := json.Marshal(item)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, string(raw), `"opencode_go_usage"`)
|
||||
require.Contains(t, string(raw), `"eligible":true`)
|
||||
|
||||
// 对照:base_url 指向 api.deepseek.com 的普通 deepseek 账号不参与用量窗口
|
||||
ineligible := &service.Account{
|
||||
ID: 165, Name: "DeepSeek Direct", Platform: service.PlatformDeepseek, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"base_url": "https://api.deepseek.com", "api_key": "test-key-2"},
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
|
||||
plainItem := AccountListItemFromAccount(AccountFromServiceShallow(ineligible))
|
||||
require.Nil(t, plainItem.OpenCodeGoUsage)
|
||||
|
||||
plainRaw, err := json.Marshal(plainItem)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, string(plainRaw), `"opencode_go_usage"`)
|
||||
}
|
||||
@@ -241,6 +241,10 @@ func AccountFromServiceShallow(a *service.Account) *Account {
|
||||
if state := service.OllamaCloudUsageStateFromAccount(a); state.Eligible {
|
||||
ollamaCloudUsage = state
|
||||
}
|
||||
var openCodeGoUsage *service.OpenCodeGoUsageState
|
||||
if state := service.OpenCodeGoUsageStateFromAccount(a); state.Eligible {
|
||||
openCodeGoUsage = state
|
||||
}
|
||||
out := &Account{
|
||||
ID: a.ID,
|
||||
Name: a.Name,
|
||||
@@ -251,6 +255,7 @@ func AccountFromServiceShallow(a *service.Account) *Account {
|
||||
CredentialsStatus: credsStatus,
|
||||
Extra: extra,
|
||||
OllamaCloudUsage: ollamaCloudUsage,
|
||||
OpenCodeGoUsage: openCodeGoUsage,
|
||||
ProxyID: a.ProxyID,
|
||||
ProxyFallbackOriginID: a.ProxyFallbackOriginID,
|
||||
ProxyFallbackOriginName: a.ProxyFallbackOriginName,
|
||||
@@ -417,7 +422,9 @@ func redactAccountManagedExtra(extra map[string]any) map[string]any {
|
||||
switch key {
|
||||
case service.OllamaCloudUsageSessionExtraKey,
|
||||
service.OllamaCloudUsageAutoRefreshExtraKey,
|
||||
service.OllamaCloudUsageSnapshotExtraKey:
|
||||
service.OllamaCloudUsageSnapshotExtraKey,
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey:
|
||||
continue
|
||||
default:
|
||||
redacted[key] = value
|
||||
@@ -458,8 +465,8 @@ func AccountListItemFromAccount(a *Account) *AccountListItem {
|
||||
return &AccountListItem{
|
||||
ID: a.ID, Name: a.Name, Notes: a.Notes, Platform: a.Platform, Type: a.Type,
|
||||
Credentials: a.Credentials, CredentialsStatus: a.CredentialsStatus, Extra: a.Extra,
|
||||
OllamaCloudUsage: a.OllamaCloudUsage,
|
||||
ProxyID: a.ProxyID, ProxyFallbackOriginID: a.ProxyFallbackOriginID, ProxyFallbackOriginName: a.ProxyFallbackOriginName,
|
||||
OllamaCloudUsage: a.OllamaCloudUsage, OpenCodeGoUsage: a.OpenCodeGoUsage,
|
||||
ProxyID: a.ProxyID, ProxyFallbackOriginID: a.ProxyFallbackOriginID, ProxyFallbackOriginName: a.ProxyFallbackOriginName,
|
||||
Concurrency: a.Concurrency, LoadFactor: a.LoadFactor, Priority: a.Priority, RateMultiplier: a.RateMultiplier,
|
||||
Status: a.Status, ErrorMessage: a.ErrorMessage, LastUsedAt: a.LastUsedAt, ExpiresAt: a.ExpiresAt,
|
||||
AutoPauseOnExpired: a.AutoPauseOnExpired, CreatedAt: a.CreatedAt, UpdatedAt: a.UpdatedAt,
|
||||
|
||||
@@ -217,6 +217,9 @@ type SystemSettings struct {
|
||||
OpenAICodexClientVersion string `json:"openai_codex_client_version"`
|
||||
OpenAICodexClientVersionSynced string `json:"openai_codex_client_version_synced"`
|
||||
OpenAICodexVersionAutoSyncEnabled bool `json:"openai_codex_version_auto_sync_enabled"`
|
||||
ClaudeCodeClientVersion string `json:"claude_code_client_version"`
|
||||
ClaudeCodeClientVersionSynced string `json:"claude_code_client_version_synced"`
|
||||
ClaudeCodeVersionAutoSyncEnabled bool `json:"claude_code_version_auto_sync_enabled"`
|
||||
|
||||
// codex_cli_only 加固
|
||||
MinCodexVersion string `json:"min_codex_version"`
|
||||
@@ -236,33 +239,33 @@ type SystemSettings struct {
|
||||
PaymentVisibleMethodWxpayEnabled bool `json:"payment_visible_method_wxpay_enabled"`
|
||||
|
||||
// OpenAI account scheduling
|
||||
OpenAILowUpstreamRatePriorityEnabled bool `json:"openai_low_upstream_rate_priority_enabled"`
|
||||
OpenAIOAuthSchedulingRateMultiplier float64 `json:"openai_oauth_scheduling_rate_multiplier"`
|
||||
OpenAIAdvancedSchedulerEnabled bool `json:"openai_advanced_scheduler_enabled"`
|
||||
OpenAIAdvancedSchedulerStickyWeightedEnabled bool `json:"openai_advanced_scheduler_sticky_weighted_enabled"`
|
||||
OpenAIAdvancedSchedulerSubscriptionPriorityEnabled bool `json:"openai_advanced_scheduler_subscription_priority_enabled"`
|
||||
OpenAIAdvancedSchedulerLBTopK string `json:"openai_advanced_scheduler_lb_top_k"`
|
||||
OpenAIAdvancedSchedulerWeightPriority string `json:"openai_advanced_scheduler_weight_priority"`
|
||||
OpenAIAdvancedSchedulerWeightLoad string `json:"openai_advanced_scheduler_weight_load"`
|
||||
OpenAIAdvancedSchedulerWeightQueue string `json:"openai_advanced_scheduler_weight_queue"`
|
||||
OpenAIAdvancedSchedulerWeightErrorRate string `json:"openai_advanced_scheduler_weight_error_rate"`
|
||||
OpenAIAdvancedSchedulerWeightTTFT string `json:"openai_advanced_scheduler_weight_ttft"`
|
||||
OpenAIAdvancedSchedulerWeightReset string `json:"openai_advanced_scheduler_weight_reset"`
|
||||
OpenAIAdvancedSchedulerWeightQuotaHeadroom string `json:"openai_advanced_scheduler_weight_quota_headroom"`
|
||||
OpenAIAdvancedSchedulerWeightUpstreamCost string `json:"openai_advanced_scheduler_weight_upstream_cost"`
|
||||
OpenAIAdvancedSchedulerWeightPreviousResponse string `json:"openai_advanced_scheduler_weight_previous_response"`
|
||||
OpenAIAdvancedSchedulerWeightSessionSticky string `json:"openai_advanced_scheduler_weight_session_sticky"`
|
||||
OpenAIAdvancedSchedulerEffectiveLBTopK string `json:"openai_advanced_scheduler_effective_lb_top_k"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightPriority string `json:"openai_advanced_scheduler_effective_weight_priority"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightLoad string `json:"openai_advanced_scheduler_effective_weight_load"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightQueue string `json:"openai_advanced_scheduler_effective_weight_queue"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightErrorRate string `json:"openai_advanced_scheduler_effective_weight_error_rate"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightTTFT string `json:"openai_advanced_scheduler_effective_weight_ttft"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightReset string `json:"openai_advanced_scheduler_effective_weight_reset"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightQuotaHeadroom string `json:"openai_advanced_scheduler_effective_weight_quota_headroom"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightUpstreamCost string `json:"openai_advanced_scheduler_effective_weight_upstream_cost"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightPreviousResponse string `json:"openai_advanced_scheduler_effective_weight_previous_response"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightSessionSticky string `json:"openai_advanced_scheduler_effective_weight_session_sticky"`
|
||||
OpenAILowUpstreamRatePriorityEnabled bool `json:"openai_low_upstream_rate_priority_enabled"`
|
||||
OpenAIOAuthSchedulingRateMultiplier *float64 `json:"openai_oauth_scheduling_rate_multiplier"`
|
||||
OpenAIAdvancedSchedulerEnabled bool `json:"openai_advanced_scheduler_enabled"`
|
||||
OpenAIAdvancedSchedulerStickyWeightedEnabled bool `json:"openai_advanced_scheduler_sticky_weighted_enabled"`
|
||||
OpenAIAdvancedSchedulerSubscriptionPriorityEnabled bool `json:"openai_advanced_scheduler_subscription_priority_enabled"`
|
||||
OpenAIAdvancedSchedulerLBTopK string `json:"openai_advanced_scheduler_lb_top_k"`
|
||||
OpenAIAdvancedSchedulerWeightPriority string `json:"openai_advanced_scheduler_weight_priority"`
|
||||
OpenAIAdvancedSchedulerWeightLoad string `json:"openai_advanced_scheduler_weight_load"`
|
||||
OpenAIAdvancedSchedulerWeightQueue string `json:"openai_advanced_scheduler_weight_queue"`
|
||||
OpenAIAdvancedSchedulerWeightErrorRate string `json:"openai_advanced_scheduler_weight_error_rate"`
|
||||
OpenAIAdvancedSchedulerWeightTTFT string `json:"openai_advanced_scheduler_weight_ttft"`
|
||||
OpenAIAdvancedSchedulerWeightReset string `json:"openai_advanced_scheduler_weight_reset"`
|
||||
OpenAIAdvancedSchedulerWeightQuotaHeadroom string `json:"openai_advanced_scheduler_weight_quota_headroom"`
|
||||
OpenAIAdvancedSchedulerWeightUpstreamCost string `json:"openai_advanced_scheduler_weight_upstream_cost"`
|
||||
OpenAIAdvancedSchedulerWeightPreviousResponse string `json:"openai_advanced_scheduler_weight_previous_response"`
|
||||
OpenAIAdvancedSchedulerWeightSessionSticky string `json:"openai_advanced_scheduler_weight_session_sticky"`
|
||||
OpenAIAdvancedSchedulerEffectiveLBTopK string `json:"openai_advanced_scheduler_effective_lb_top_k"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightPriority string `json:"openai_advanced_scheduler_effective_weight_priority"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightLoad string `json:"openai_advanced_scheduler_effective_weight_load"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightQueue string `json:"openai_advanced_scheduler_effective_weight_queue"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightErrorRate string `json:"openai_advanced_scheduler_effective_weight_error_rate"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightTTFT string `json:"openai_advanced_scheduler_effective_weight_ttft"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightReset string `json:"openai_advanced_scheduler_effective_weight_reset"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightQuotaHeadroom string `json:"openai_advanced_scheduler_effective_weight_quota_headroom"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightUpstreamCost string `json:"openai_advanced_scheduler_effective_weight_upstream_cost"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightPreviousResponse string `json:"openai_advanced_scheduler_effective_weight_previous_response"`
|
||||
OpenAIAdvancedSchedulerEffectiveWeightSessionSticky string `json:"openai_advanced_scheduler_effective_weight_session_sticky"`
|
||||
|
||||
// Payment configuration
|
||||
PaymentEnabled bool `json:"payment_enabled"`
|
||||
@@ -332,8 +335,9 @@ type SystemSettings struct {
|
||||
RiskControlEnabled bool `json:"risk_control_enabled"`
|
||||
|
||||
// cyber 会话屏蔽开关 + TTL
|
||||
CyberSessionBlockEnabled bool `json:"cyber_session_block_enabled"`
|
||||
CyberSessionBlockTTLSeconds int `json:"cyber_session_block_ttl_seconds"`
|
||||
CyberSessionBlockEnabled bool `json:"cyber_session_block_enabled"`
|
||||
CyberPolicyUserAllowlist string `json:"cyber_policy_user_allowlist"`
|
||||
CyberSessionBlockTTLSeconds int `json:"cyber_session_block_ttl_seconds"`
|
||||
|
||||
// Affiliate (邀请返利) feature switch
|
||||
AffiliateEnabled bool `json:"affiliate_enabled"`
|
||||
|
||||
@@ -216,6 +216,7 @@ type Account struct {
|
||||
CredentialsStatus map[string]bool `json:"credentials_status,omitempty"`
|
||||
Extra map[string]any `json:"extra"`
|
||||
OllamaCloudUsage *service.OllamaCloudUsageState `json:"ollama_cloud_usage,omitempty"`
|
||||
OpenCodeGoUsage *service.OpenCodeGoUsageState `json:"opencode_go_usage,omitempty"`
|
||||
ProxyID *int64 `json:"proxy_id"`
|
||||
ProxyFallbackOriginID *int64 `json:"proxy_fallback_origin_id"`
|
||||
ProxyFallbackOriginName *string `json:"proxy_fallback_origin_name,omitempty"`
|
||||
@@ -340,6 +341,7 @@ type AccountListItem struct {
|
||||
CredentialsStatus map[string]bool `json:"credentials_status,omitempty"`
|
||||
Extra map[string]any `json:"extra,omitempty"`
|
||||
OllamaCloudUsage *service.OllamaCloudUsageState `json:"ollama_cloud_usage,omitempty"`
|
||||
OpenCodeGoUsage *service.OpenCodeGoUsageState `json:"opencode_go_usage,omitempty"`
|
||||
|
||||
ProxyID *int64 `json:"proxy_id"`
|
||||
ProxyFallbackOriginID *int64 `json:"proxy_fallback_origin_id"`
|
||||
|
||||
@@ -29,6 +29,7 @@ const (
|
||||
EndpointVideosEdits = "/v1/videos/edits"
|
||||
EndpointVideosExtensions = "/v1/videos/extensions"
|
||||
EndpointVideos = "/v1/videos"
|
||||
EndpointSeedanceTasks = "/api/v3/contents/generations/tasks"
|
||||
EndpointGeminiModels = "/v1beta/models"
|
||||
)
|
||||
|
||||
@@ -81,6 +82,8 @@ const (
|
||||
func NormalizeInboundEndpoint(path string) string {
|
||||
path = strings.TrimSpace(path)
|
||||
switch {
|
||||
case strings.Contains(path, "/contents/generations/tasks"):
|
||||
return EndpointSeedanceTasks
|
||||
case strings.Contains(path, EndpointResponsesInputTokens) || isResponsesInputTokensAliasPath(path):
|
||||
return EndpointResponsesInputTokens
|
||||
case strings.Contains(path, EndpointEmbeddings):
|
||||
|
||||
@@ -263,6 +263,19 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gateway.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// 设置请求所属分组 ID(用于渠道级功能判断,如 WebSearch 模拟)
|
||||
parsedReq.GroupID = apiKey.GroupID
|
||||
|
||||
@@ -330,6 +343,10 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
for {
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("gateway.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel, service.PlatformGemini)
|
||||
if !cls.ModelNotFound {
|
||||
@@ -660,6 +677,10 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
)
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), currentAPIKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, parsedReq.MetadataUserID, subject.UserID)
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("gateway.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, currentAPIKey, reqModel, reqModel, platform)
|
||||
if !cls.ModelNotFound {
|
||||
@@ -1375,7 +1396,7 @@ func writeGrokModelsList(c *gin.Context, modelIDs []string) {
|
||||
|
||||
func grokModelSupportsConfigurableReasoning(modelID string) bool {
|
||||
switch strings.ToLower(strings.TrimSpace(modelID)) {
|
||||
case "grok-4.6", "grok-4.6-latest", "grok-4.5", "grok-4.5-latest", "grok", "grok-latest", "grok-build", "grok-build-latest", "grok-build-0.1":
|
||||
case "grok-4.7", "grok-4.7-latest", "grok-4.6", "grok-4.6-latest", "grok-4.5", "grok-4.5-latest", "grok", "grok-latest", "grok-build", "grok-build-latest", "grok-build-0.1":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
@@ -1971,10 +1992,15 @@ func (h *GatewayHandler) handleStreamingAwareErrorWithCode(c *gin.Context, statu
|
||||
// Writer 已被写过时(ping 已 flush)走 streamStarted 分支,
|
||||
// 让 handleStreamingAwareError 通过 SSE 发协议合规的终止事件,
|
||||
// 否则下游收到的就是 silent EOF。
|
||||
// 客户端已断开时不补写,响应未提交则标记 499(见 failoverClientGone)。
|
||||
func (h *GatewayHandler) ensureForwardErrorResponse(c *gin.Context, streamStarted bool) bool {
|
||||
if c == nil || c.Writer == nil {
|
||||
return false
|
||||
}
|
||||
if c.Request != nil && errors.Is(c.Request.Context().Err(), context.Canceled) {
|
||||
failoverClientGone(c)
|
||||
return false
|
||||
}
|
||||
if service.IsResponseCommitted(c) {
|
||||
return false
|
||||
}
|
||||
@@ -2501,9 +2527,12 @@ func (h *GatewayHandler) submitUsageRecordTask(parent context.Context, task serv
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
task, abandon := wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped {
|
||||
if mode.Dropped() {
|
||||
abandon()
|
||||
}
|
||||
return
|
||||
}
|
||||
// 池已停止(进程关停窗口):计费任务不能静默丢失,降级为内联同步执行。
|
||||
@@ -2531,7 +2560,7 @@ func (h *GatewayHandler) submitMandatoryUsageRecordTask(parent context.Context,
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
task, _ = wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() {
|
||||
return
|
||||
|
||||
@@ -93,6 +93,8 @@ func TestGatewayHandlerPreCancelledCompatibleRequestsDoNotSelectAccount(t *testi
|
||||
require.Zero(t, schedulerCache.snapshotCalls.Load(), "a cancelled request must stop before the account selector")
|
||||
_, selected := c.Get(opsAccountIDKey)
|
||||
require.False(t, selected)
|
||||
require.Equal(t, statusClientClosedRequest, c.Writer.Status(), "an uncommitted cancelled request is marked 499")
|
||||
require.Zero(t, recorder.Body.Len())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -101,8 +101,11 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
// 解析渠道级模型映射
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, reqModel)
|
||||
|
||||
// Claude Code only restriction
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly {
|
||||
// Claude Code only restriction: /v1/chat/completions is never a Claude Code
|
||||
// endpoint. With a fallback group the request continues and account selection
|
||||
// (checkClaudeCodeRestriction) schedules it in the fallback group; without one
|
||||
// it is rejected here.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly && apiKey.Group.FallbackGroupID == nil {
|
||||
h.chatCompletionsErrorResponse(c, http.StatusForbidden, "permission_error",
|
||||
"This group is restricted to Claude Code clients (/v1/messages only)")
|
||||
return
|
||||
@@ -144,6 +147,19 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gateway.cc.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.chatCompletionsErrorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// Parse request for session hash
|
||||
bodyRef := service.NewRequestBodyRef(body)
|
||||
parsedReq, _ := service.ParseGatewayRequest(bodyRef, "chat_completions")
|
||||
@@ -168,11 +184,15 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
}
|
||||
|
||||
for {
|
||||
if c.Request.Context().Err() != nil {
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, selectionSessionHash, reqModel, fs.FailedAccountIDs, "", int64(0))
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("gateway.cc.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel, groupPlatform)
|
||||
cls = classifySelectionFailureError(err, cls)
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// groupScopedSchedulerCache 按桶的分组只返回该分组的成员账号,并记录被查询过的分组。
|
||||
type groupScopedSchedulerCache struct {
|
||||
*fakeSchedulerCache
|
||||
mu sync.Mutex
|
||||
groupIDs []int64
|
||||
}
|
||||
|
||||
func (c *groupScopedSchedulerCache) GetSnapshot(_ context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) {
|
||||
c.mu.Lock()
|
||||
c.groupIDs = append(c.groupIDs, bucket.GroupID)
|
||||
c.mu.Unlock()
|
||||
var members []*service.Account
|
||||
for _, account := range c.accounts {
|
||||
for _, ag := range account.AccountGroups {
|
||||
if ag.GroupID == bucket.GroupID {
|
||||
members = append(members, account)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return members, true, nil
|
||||
}
|
||||
|
||||
func (c *groupScopedSchedulerCache) queriedGroupIDs() []int64 {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return append([]int64(nil), c.groupIDs...)
|
||||
}
|
||||
|
||||
type groupMapRepo struct {
|
||||
*fakeGroupRepo
|
||||
groups map[int64]*service.Group
|
||||
}
|
||||
|
||||
func (r *groupMapRepo) GetByID(_ context.Context, id int64) (*service.Group, error) {
|
||||
if group, ok := r.groups[id]; ok {
|
||||
return group, nil
|
||||
}
|
||||
return nil, service.ErrGroupNotFound
|
||||
}
|
||||
|
||||
func (r *groupMapRepo) GetByIDLite(ctx context.Context, id int64) (*service.Group, error) {
|
||||
return r.GetByID(ctx, id)
|
||||
}
|
||||
|
||||
func TestGatewayOpenAICompatibleHandlersClaudeCodeOnlyFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
const (
|
||||
primaryGroupID = int64(9300)
|
||||
fallbackGroupID = int64(9301)
|
||||
primaryAccountID = int64(9310)
|
||||
fallbackAccountID = int64(9311)
|
||||
)
|
||||
|
||||
// 账号不带 api_key:转发在取令牌时失败,请求不会触达上游,断言只看选号结果。
|
||||
newAccount := func(id, groupID int64) *service.Account {
|
||||
return &service.Account{
|
||||
ID: id, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive, Schedulable: true, Concurrency: 1,
|
||||
AccountGroups: []service.AccountGroup{{AccountID: id, GroupID: groupID}},
|
||||
}
|
||||
}
|
||||
|
||||
endpoints := []struct {
|
||||
name string
|
||||
path string
|
||||
body string
|
||||
call func(*GatewayHandler, *gin.Context)
|
||||
}{
|
||||
{
|
||||
name: "responses", path: "/v1/responses",
|
||||
body: `{"model":"claude-sonnet-4-5","input":"hello","stream":false}`,
|
||||
call: (*GatewayHandler).Responses,
|
||||
},
|
||||
{
|
||||
name: "chat completions", path: "/v1/chat/completions",
|
||||
body: `{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hello"}],"stream":false}`,
|
||||
call: (*GatewayHandler).ChatCompletions,
|
||||
},
|
||||
}
|
||||
|
||||
for _, ep := range endpoints {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
hasFallback bool
|
||||
}{
|
||||
{name: "with fallback group", hasFallback: true},
|
||||
{name: "without fallback group", hasFallback: false},
|
||||
} {
|
||||
t.Run(ep.name+"/"+tc.name, func(t *testing.T) {
|
||||
primary := &service.Group{
|
||||
ID: primaryGroupID, Hydrated: true, Platform: service.PlatformAnthropic,
|
||||
Status: service.StatusActive, ClaudeCodeOnly: true,
|
||||
}
|
||||
if tc.hasFallback {
|
||||
fallbackID := fallbackGroupID
|
||||
primary.FallbackGroupID = &fallbackID
|
||||
}
|
||||
fallback := &service.Group{
|
||||
ID: fallbackGroupID, Hydrated: true, Platform: service.PlatformAnthropic,
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
|
||||
schedulerCache := &groupScopedSchedulerCache{fakeSchedulerCache: &fakeSchedulerCache{accounts: []*service.Account{
|
||||
newAccount(primaryAccountID, primaryGroupID),
|
||||
newAccount(fallbackAccountID, fallbackGroupID),
|
||||
}}}
|
||||
gatewayService := service.NewGatewayService(
|
||||
nil, &groupMapRepo{fakeGroupRepo: &fakeGroupRepo{}, groups: map[int64]*service.Group{
|
||||
primaryGroupID: primary,
|
||||
fallbackGroupID: fallback,
|
||||
}}, nil, nil, nil, nil, nil, nil, nil,
|
||||
service.NewSchedulerSnapshotService(schedulerCache, nil, nil, nil, nil),
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
billingCacheService := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billingCacheService.Stop)
|
||||
h := &GatewayHandler{
|
||||
gatewayService: gatewayService,
|
||||
billingCacheService: billingCacheService,
|
||||
concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatClaude, 0),
|
||||
maxAccountSwitches: 1,
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
primaryGroupIDRef := primaryGroupID
|
||||
apiKey := &service.APIKey{
|
||||
ID: 9320, UserID: 9330, GroupID: &primaryGroupIDRef, Group: primary, Status: service.StatusActive,
|
||||
User: &service.User{ID: 9330, Concurrency: 10, Balance: 100},
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
ctx := context.WithValue(context.Background(), ctxkey.Group, primary)
|
||||
req := httptest.NewRequest(http.MethodPost, ep.path, bytes.NewBufferString(ep.body)).WithContext(ctx)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.Request = req
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10})
|
||||
|
||||
ep.call(h, c)
|
||||
|
||||
selected, reachedSelection := c.Get(opsAccountIDKey)
|
||||
if !tc.hasFallback {
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "This group is restricted to Claude Code clients")
|
||||
require.False(t, reachedSelection, "a Claude Code only group without fallback must be rejected before account selection")
|
||||
require.Empty(t, schedulerCache.queriedGroupIDs())
|
||||
return
|
||||
}
|
||||
require.NotContains(t, recorder.Body.String(), "restricted to Claude Code clients")
|
||||
require.True(t, reachedSelection, "a Claude Code only group with fallback must reach account selection")
|
||||
require.Equal(t, fallbackAccountID, selected)
|
||||
queried := schedulerCache.queriedGroupIDs()
|
||||
require.Contains(t, queried, fallbackGroupID)
|
||||
require.NotContains(t, queried, primaryGroupID)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
@@ -70,6 +71,43 @@ func TestGatewayEnsureForwardErrorResponse_SkipsCommittedSSEError(t *testing.T)
|
||||
require.Equal(t, 1, strings.Count(w.Body.String(), "event: error"))
|
||||
}
|
||||
|
||||
// 客户端已断开且响应未提交:不补写 502,标记 499 供访问日志与 ops 归类。
|
||||
func TestGatewayEnsureForwardErrorResponse_SkipsCanceledClient(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, EndpointMessages, nil).WithContext(ctx)
|
||||
cancel()
|
||||
|
||||
h := &GatewayHandler{}
|
||||
wrote := h.ensureForwardErrorResponse(c, false)
|
||||
|
||||
require.False(t, wrote)
|
||||
require.Equal(t, statusClientClosedRequest, c.Writer.Status())
|
||||
require.Empty(t, w.Body.String())
|
||||
}
|
||||
|
||||
// 客户端已断开但流已开始:状态码已固化为 200,不再向断开的连接追加错误帧。
|
||||
func TestGatewayEnsureForwardErrorResponse_CanceledClientAfterStreamStartedAppendsNothing(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, EndpointMessages, nil).WithContext(ctx)
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
_, _ = c.Writer.WriteString(":\n\n")
|
||||
cancel()
|
||||
|
||||
h := &GatewayHandler{}
|
||||
wrote := h.ensureForwardErrorResponse(c, true)
|
||||
|
||||
require.False(t, wrote)
|
||||
require.Equal(t, http.StatusOK, c.Writer.Status())
|
||||
require.Equal(t, ":\n\n", w.Body.String())
|
||||
require.Empty(t, service.GetOpsStreamErrors(c))
|
||||
}
|
||||
|
||||
// case B 回归:Anthropic-backed /responses,Writer 已被写过时
|
||||
// ensureForwardErrorResponse 仍要发 response.failed。
|
||||
func TestGatewayEnsureForwardErrorResponse_ResponsesRouteAfterWrittenEmitsResponseFailed(t *testing.T) {
|
||||
|
||||
@@ -105,13 +105,11 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
// 解析渠道级模型映射
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(requestCtx, apiKey.GroupID, reqModel)
|
||||
|
||||
// Claude Code only restriction:
|
||||
// /v1/responses is never a Claude Code endpoint.
|
||||
// When claude_code_only is enabled, this endpoint is rejected.
|
||||
// The existing service-layer checkClaudeCodeRestriction handles degradation
|
||||
// to fallback groups when the Forward path calls SelectAccountForModelWithExclusions.
|
||||
// Here we just reject at handler level since /v1/responses clients can't be Claude Code.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly {
|
||||
// Claude Code only restriction: /v1/responses is never a Claude Code
|
||||
// endpoint. With a fallback group the request continues and account selection
|
||||
// (checkClaudeCodeRestriction) schedules it in the fallback group; without one
|
||||
// it is rejected here.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly && apiKey.Group.FallbackGroupID == nil {
|
||||
h.responsesErrorResponse(c, http.StatusForbidden, "permission_error",
|
||||
"This group is restricted to Claude Code clients (/v1/messages only)")
|
||||
return
|
||||
@@ -153,6 +151,19 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gateway.responses.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.responsesErrorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// Parse request for session hash
|
||||
bodyRef := service.NewRequestBodyRef(body)
|
||||
parsedReq, _ := service.ParseGatewayRequest(bodyRef, "responses")
|
||||
@@ -170,11 +181,15 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
fs := NewFailoverState(h.maxAccountSwitches, false)
|
||||
|
||||
for {
|
||||
if requestCtx.Err() != nil {
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(requestCtx, apiKey.GroupID, sessionHash, reqModel, fs.FailedAccountIDs, "", int64(0))
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("gateway.responses.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel, effectiveAPIKeyPlatform(c, apiKey))
|
||||
cls = classifySelectionFailureError(err, cls)
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// inflightReservationEstimator 估算单请求在途预留金额(USD);false 表示无法定价。
|
||||
type inflightReservationEstimator interface {
|
||||
EstimateInflightReservation(ctx context.Context, apiKey *service.APIKey, req service.InflightEstimateRequest) (float64, bool)
|
||||
}
|
||||
|
||||
// requestMaxOutputTokens 从请求体中提取输出 token 上限(兼容 Anthropic / OpenAI Chat / Responses / Gemini)。
|
||||
func requestMaxOutputTokens(body []byte) int {
|
||||
for _, path := range []string{"max_tokens", "max_completion_tokens", "max_output_tokens", "generationConfig.maxOutputTokens", "generation_config.max_output_tokens"} {
|
||||
if v := gjson.GetBytes(body, path); v.Exists() && v.Type == gjson.Number && v.Int() > 0 {
|
||||
return int(v.Int())
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// tokenInflightEstimate 文本类请求的估算输入。
|
||||
func tokenInflightEstimate(model string, body []byte) service.InflightEstimateRequest {
|
||||
return service.InflightEstimateRequest{
|
||||
Model: model,
|
||||
BodyBytes: len(body),
|
||||
MaxTokens: requestMaxOutputTokens(body),
|
||||
Kind: service.InflightEstimateToken,
|
||||
}
|
||||
}
|
||||
|
||||
func inflightNoop() {}
|
||||
|
||||
// reserveInflightBalance 在 CheckBillingEligibility 之后为余额模式请求登记在途预留。
|
||||
//
|
||||
// 成功时把预留句柄挂到 c.Request 的 context 上:之后通过 submit*UsageRecordTask 提交的
|
||||
// 计费任务会接管一个引用,直到余额缓存实际扣减后才释放。返回的 done 必须 defer 调用
|
||||
// (停止续期并归还 handler 引用;没有提交计费任务时即刻释放)。
|
||||
// 开关关闭、订阅模式、Redis 故障时返回 no-op(fail-open);无法定价时默认 fail-open,
|
||||
// 配置 fail_closed_on_unpriced=true 时返回 ErrInsufficientBalance。
|
||||
func reserveInflightBalance(
|
||||
c *gin.Context,
|
||||
billing *service.BillingCacheService,
|
||||
estimator inflightReservationEstimator,
|
||||
apiKey *service.APIKey,
|
||||
subscription *service.UserSubscription,
|
||||
req service.InflightEstimateRequest,
|
||||
) (func(), error) {
|
||||
if c == nil || c.Request == nil {
|
||||
return inflightNoop, nil
|
||||
}
|
||||
ctx, done, err := reserveInflightBalanceCtx(c.Request.Context(), billing, estimator, apiKey, subscription, req)
|
||||
if err != nil {
|
||||
return inflightNoop, err
|
||||
}
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
return done, nil
|
||||
}
|
||||
|
||||
// reserveInflightBalanceCtx 同 reserveInflightBalance,但返回携带预留句柄的新 context
|
||||
// (供 WebSocket 等自管 context 的路径)。
|
||||
func reserveInflightBalanceCtx(
|
||||
ctx context.Context,
|
||||
billing *service.BillingCacheService,
|
||||
estimator inflightReservationEstimator,
|
||||
apiKey *service.APIKey,
|
||||
subscription *service.UserSubscription,
|
||||
req service.InflightEstimateRequest,
|
||||
) (context.Context, func(), error) {
|
||||
if billing == nil || estimator == nil || apiKey == nil || apiKey.User == nil || !billing.InflightReservationEnabled() {
|
||||
return ctx, inflightNoop, nil
|
||||
}
|
||||
if apiKey.Group != nil && apiKey.Group.IsSubscriptionType() && subscription != nil {
|
||||
return ctx, inflightNoop, nil
|
||||
}
|
||||
estimate, priced := estimator.EstimateInflightReservation(ctx, apiKey, req)
|
||||
if !priced && billing.InflightReservationFailClosedOnUnpriced() {
|
||||
return ctx, inflightNoop, service.ErrInsufficientBalance
|
||||
}
|
||||
if estimate <= 0 {
|
||||
return ctx, inflightNoop, nil
|
||||
}
|
||||
res, err := billing.ReserveInflight(ctx, apiKey.User, apiKey.Group, subscription, estimate)
|
||||
if err != nil {
|
||||
return ctx, inflightNoop, err
|
||||
}
|
||||
if res == nil {
|
||||
return ctx, inflightNoop, nil
|
||||
}
|
||||
return service.WithInflightReservation(ctx, res), res.HandlerDone, nil
|
||||
}
|
||||
|
||||
// grokMediaInflightEstimate 媒体生成请求的估算输入;状态/内容查询返回空模型(不预留:
|
||||
// 查询会为已生成的媒体计费,不能因余额预留而拦截用户取回已付费结果)。
|
||||
func grokMediaInflightEstimate(endpoint service.GrokMediaEndpoint, model string, info service.GrokMediaRequestInfo, body []byte) service.InflightEstimateRequest {
|
||||
if !endpoint.IsGenerationRequest() {
|
||||
return service.InflightEstimateRequest{}
|
||||
}
|
||||
switch endpoint {
|
||||
case service.GrokMediaEndpointImagesGenerations, service.GrokMediaEndpointImagesEdits:
|
||||
return service.InflightEstimateRequest{Model: model, BodyBytes: len(body), Kind: service.InflightEstimateImage, Units: info.N}
|
||||
default:
|
||||
return service.InflightEstimateRequest{
|
||||
Model: model,
|
||||
BodyBytes: len(body),
|
||||
Kind: service.InflightEstimateVideo,
|
||||
Units: 1,
|
||||
VideoResolution: info.Resolution,
|
||||
VideoDurationSeconds: info.DurationSeconds,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// grokVoiceSTTBytesPerSecond STT 时长粗估(~128kbps 压缩音频)。
|
||||
const grokVoiceSTTBytesPerSecond = 16000
|
||||
|
||||
// grokVoiceInflightEstimate 语音 HTTP 接口估算:TTS 按输入字符数(百万字符),STT 按音频字节粗估时长(小时)。
|
||||
// 其他接口(custom-voices)无音频计量,返回的估算为 0(不预留)。
|
||||
func grokVoiceInflightEstimate(endpoint string, body []byte) service.InflightEstimateRequest {
|
||||
req := service.InflightEstimateRequest{Model: endpoint, Kind: service.InflightEstimateAudio, AudioMode: endpoint}
|
||||
switch endpoint {
|
||||
case "tts":
|
||||
req.AudioUnits = float64(len([]rune(extractGrokTTSInputText(body)))) / 1e6
|
||||
case "stt":
|
||||
req.AudioUnits = float64(len(body)) / grokVoiceSTTBytesPerSecond / 3600
|
||||
}
|
||||
return req
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestRequestMaxOutputTokens(t *testing.T) {
|
||||
require.Equal(t, 1024, requestMaxOutputTokens([]byte(`{"max_tokens":1024}`)))
|
||||
require.Equal(t, 2048, requestMaxOutputTokens([]byte(`{"max_completion_tokens":2048}`)))
|
||||
require.Equal(t, 4096, requestMaxOutputTokens([]byte(`{"max_output_tokens":4096}`)))
|
||||
require.Equal(t, 512, requestMaxOutputTokens([]byte(`{"generationConfig":{"maxOutputTokens":512}}`)))
|
||||
require.Equal(t, 0, requestMaxOutputTokens([]byte(`{"max_tokens":"x"}`)))
|
||||
require.Equal(t, 0, requestMaxOutputTokens([]byte(`{}`)))
|
||||
}
|
||||
|
||||
type countingEstimator struct {
|
||||
calls int
|
||||
cost float64
|
||||
priced bool
|
||||
}
|
||||
|
||||
func (e *countingEstimator) EstimateInflightReservation(context.Context, *service.APIKey, service.InflightEstimateRequest) (float64, bool) {
|
||||
e.calls++
|
||||
return e.cost, e.priced
|
||||
}
|
||||
|
||||
func newInflightTestGinContext() *gin.Context {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
return c
|
||||
}
|
||||
|
||||
func TestReserveInflightBalance_SkipsWhenDisabledOrSubscription(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
billing := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billing.Stop)
|
||||
est := &countingEstimator{cost: 1, priced: true}
|
||||
apiKey := &service.APIKey{User: &service.User{ID: 1}}
|
||||
|
||||
done, err := reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, nil, tokenInflightEstimate("m", []byte(`{}`)))
|
||||
require.NoError(t, err)
|
||||
done()
|
||||
require.Equal(t, 0, est.calls, "disabled switch must not even estimate")
|
||||
|
||||
cfg.Billing.InflightReservation.Enabled = true
|
||||
apiKey.Group = &service.Group{SubscriptionType: service.SubscriptionTypeSubscription}
|
||||
done, err = reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, &service.UserSubscription{}, tokenInflightEstimate("m", []byte(`{}`)))
|
||||
require.NoError(t, err)
|
||||
done()
|
||||
require.Equal(t, 0, est.calls, "subscription mode must be unaffected")
|
||||
}
|
||||
|
||||
func TestReserveInflightBalance_UnpricedFailOpenByDefaultFailClosedOptIn(t *testing.T) {
|
||||
cache := newHandlerInflightCache(10)
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60}
|
||||
billing := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billing.Stop)
|
||||
apiKey := &service.APIKey{User: &service.User{ID: 1}}
|
||||
est := &countingEstimator{priced: false}
|
||||
|
||||
done, err := reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, nil, tokenInflightEstimate("unknown", nil))
|
||||
require.NoError(t, err)
|
||||
done()
|
||||
|
||||
cfg.Billing.InflightReservation.FailClosedOnUnpriced = true
|
||||
_, err = reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, nil, tokenInflightEstimate("unknown", nil))
|
||||
require.ErrorIs(t, err, service.ErrInsufficientBalance)
|
||||
}
|
||||
|
||||
// handlerInflightCache 内存版余额缓存 + 在途预留(语义同 Redis Lua)。
|
||||
type handlerInflightCache struct {
|
||||
service.BillingCache
|
||||
mu sync.Mutex
|
||||
balance float64
|
||||
res map[string]float64
|
||||
}
|
||||
|
||||
func newHandlerInflightCache(balance float64) *handlerInflightCache {
|
||||
return &handlerInflightCache{balance: balance, res: map[string]float64{}}
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) GetUserBalance(context.Context, int64) (float64, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return m.balance, nil
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) GetUserPlatformQuotaCache(context.Context, int64, string) (*service.UserPlatformQuotaCacheEntry, bool, error) {
|
||||
return nil, false, nil
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) ReserveInflightBalance(_ context.Context, _ int64, id string, amount, balance float64, _ time.Duration) (bool, float64, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
sum := 0.0
|
||||
for _, v := range m.res {
|
||||
sum += v
|
||||
}
|
||||
if len(m.res) > 0 && balance-sum < amount {
|
||||
return false, sum, nil
|
||||
}
|
||||
m.res[id] = amount
|
||||
return true, sum, nil
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) ReleaseInflightBalance(_ context.Context, _ int64, id string) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
delete(m.res, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) count() int {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return len(m.res)
|
||||
}
|
||||
|
||||
func TestWrapUsageRecordTaskContext_HandsReservationToBillingTask(t *testing.T) {
|
||||
cache := newHandlerInflightCache(1)
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60}
|
||||
billing := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billing.Stop)
|
||||
apiKey := &service.APIKey{User: &service.User{ID: 5}}
|
||||
|
||||
c := newInflightTestGinContext()
|
||||
done, err := reserveInflightBalance(c, billing, &countingEstimator{cost: 0.9, priced: true}, apiKey, nil, tokenInflightEstimate("m", nil))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, cache.count())
|
||||
|
||||
ran := false
|
||||
task, abandon := wrapUsageRecordTaskContext(c.Request.Context(), func(context.Context) { ran = true })
|
||||
done() // handler returns; billing still pending
|
||||
require.Equal(t, 1, cache.count(), "reservation held until the billing task finishes")
|
||||
task(context.Background())
|
||||
require.True(t, ran)
|
||||
require.Equal(t, 0, cache.count())
|
||||
abandon() // idempotent with the task's own done
|
||||
|
||||
// Dropped task: the submitter abandons it and the reservation is released.
|
||||
c2 := newInflightTestGinContext()
|
||||
done2, err := reserveInflightBalance(c2, billing, &countingEstimator{cost: 0.9, priced: true}, apiKey, nil, tokenInflightEstimate("m", nil))
|
||||
require.NoError(t, err)
|
||||
_, abandon2 := wrapUsageRecordTaskContext(c2.Request.Context(), func(context.Context) {})
|
||||
done2()
|
||||
require.Equal(t, 1, cache.count())
|
||||
abandon2()
|
||||
require.Equal(t, 0, cache.count())
|
||||
}
|
||||
|
||||
// 新接入的端点(独立 web_search):在途预留超过余额时拒绝,且不残留预留。
|
||||
func TestWebSearch_RejectsWhenInflightExceedsBalance(t *testing.T) {
|
||||
cache := newHandlerInflightCache(1.5)
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60}
|
||||
billingCache := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billingCache.Stop)
|
||||
billing := service.NewBillingService(cfg, nil)
|
||||
gw := service.NewGatewayService(
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, cfg, nil, nil, billing, nil, nil,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, service.NewModelPricingResolver(nil, billing), nil, nil, nil,
|
||||
)
|
||||
h := &GatewayHandler{gatewayService: gw, billingCacheService: billingCache}
|
||||
|
||||
groupID := int64(3)
|
||||
perK := 1000.0 // $1 per search
|
||||
apiKey := &service.APIKey{
|
||||
ID: 9, User: &service.User{ID: 42, Balance: 1.5}, GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, RateMultiplier: 1, SearchPricePer1k: &perK},
|
||||
}
|
||||
|
||||
// Another in-flight request of this user already holds $1.
|
||||
held, err := billingCache.ReserveInflight(context.Background(), apiKey.User, apiKey.Group, nil, 1.0)
|
||||
require.NoError(t, err)
|
||||
defer held.HandlerDone()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/web_search", bytes.NewBufferString(`{"query":"sub2api"}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), apiKey)
|
||||
|
||||
h.WebSearch(c)
|
||||
|
||||
require.NotEqual(t, http.StatusOK, w.Code)
|
||||
require.Contains(t, w.Body.String(), "balance")
|
||||
require.Equal(t, 1, cache.count(), "rejected request must not leave a reservation behind")
|
||||
}
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
@@ -286,6 +287,18 @@ func TestGatewayModels_UnmappedOpenAIAccountsSupplementMappedModels(t *testing.T
|
||||
accounts: accounts[1:],
|
||||
want: []string{sparkModel, alias},
|
||||
},
|
||||
{
|
||||
// A passthrough account with a stale mapping behaves like an unmapped
|
||||
// one: it adds the defaults but never its own mapping keys, and it no
|
||||
// longer hides the aliases declared on ordinary accounts.
|
||||
name: "passthrough account contributes defaults without hiding mapped aliases",
|
||||
accounts: append([]service.Account{{
|
||||
ID: 5, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{"stale-model": "stale-model"}},
|
||||
Extra: map[string]any{"openai_passthrough": true},
|
||||
}}, accounts[1:]...),
|
||||
want: append(openai.DefaultModelIDs(), alias),
|
||||
},
|
||||
{
|
||||
name: "unmapped accounts from another platform do not add defaults",
|
||||
accounts: append([]service.Account{{ID: 4, Platform: service.PlatformAnthropic}}, accounts[1:]...),
|
||||
@@ -1315,3 +1328,197 @@ func modelIDsForTest(models []gatewayModelItemForTest) []string {
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
// gemini 分组混合调度:组内 antigravity 账号的 gemini-* 映射也应出现在 /v1/models,
|
||||
// 与路由层(listSchedulableAccountsOnce 的 useMixedScheduling)同源;claude-* 不带入,
|
||||
// 非混合调度的异平台账号(anthropic)仍被过滤。
|
||||
func TestGatewayModels_GeminiGroupIncludesAntigravityGeminiMappings(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
groupID := int64(22)
|
||||
h := newGatewayModelsHandlerForTest(
|
||||
&gatewayModelsAccountRepoStub{
|
||||
byGroup: map[int64][]service.Account{
|
||||
groupID: {
|
||||
{
|
||||
ID: 1,
|
||||
Platform: service.PlatformAntigravity,
|
||||
Extra: map[string]any{"mixed_scheduling": true},
|
||||
Credentials: map[string]any{
|
||||
"model_mapping": map[string]any{
|
||||
"gemini-3.8-flash-high": "gemini-3.8-flash-high",
|
||||
"gemini-synced-custom": "gemini-3.8-flash-high",
|
||||
"gemini-3.8-flash-low": "gemini-3.8-flash-low",
|
||||
"claude-sonnet-4-6": "claude-sonnet-4-6",
|
||||
},
|
||||
},
|
||||
},
|
||||
{ID: 3, Platform: service.PlatformAntigravity,
|
||||
Extra: map[string]any{"mixed_scheduling": false},
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{"gemini-disabled-only": "gemini-3.8-flash-high"}}},
|
||||
{ID: 4, Platform: service.PlatformAntigravity,
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{"gemini-unset-only": "gemini-3.8-flash-high"}}},
|
||||
{ID: 5, Platform: service.PlatformGemini,
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{"gemini-native-only": "gemini-2.5-pro"}}},
|
||||
{
|
||||
ID: 2,
|
||||
Platform: service.PlatformAnthropic,
|
||||
Credentials: map[string]any{
|
||||
"model_mapping": map[string]any{
|
||||
"claude-opus-4-6": "claude-opus-4-6",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGemini},
|
||||
})
|
||||
|
||||
h.Models(c)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var got gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
ids := modelIDsForTest(got.Data)
|
||||
// 账号显式配置的 gemini-* 映射必须出现(resolveModelMapping 还会补一组 antigravity
|
||||
// 默认透传别名,这里只断言显式条目与排除项,不锁定补齐后的完整集合)。
|
||||
require.Contains(t, ids, "gemini-3.8-flash-high")
|
||||
require.Contains(t, ids, "gemini-synced-custom")
|
||||
require.Contains(t, ids, "gemini-3.8-flash-low")
|
||||
require.Contains(t, ids, "gemini-native-only")
|
||||
require.NotContains(t, ids, "gemini-disabled-only")
|
||||
require.NotContains(t, ids, "gemini-unset-only")
|
||||
// antigravity 账号的 claude-* 映射与 anthropic 账号的映射都不应出现在 gemini 分组。
|
||||
require.NotContains(t, ids, "claude-sonnet-4-6")
|
||||
require.NotContains(t, ids, "claude-opus-4-6")
|
||||
for _, id := range ids {
|
||||
require.True(t, strings.HasPrefix(id, "gemini-"), "unexpected non-gemini model on gemini group: %s", id)
|
||||
}
|
||||
// 没有回落到 geminicli 静态表。
|
||||
require.NotContains(t, ids, "gemini-2.0-flash")
|
||||
}
|
||||
|
||||
// antigravity 账号未配置 model_mapping 时使用 DefaultAntigravityModelMapping:
|
||||
// gemini 分组应列出其中的 gemini-* 条目,而不是回落到 geminicli 静态表;claude-* 仍不出现。
|
||||
func TestGatewayModels_GeminiGroupUsesAntigravityDefaultMappingWhenUnset(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
groupID := int64(23)
|
||||
h := newGatewayModelsHandlerForTest(
|
||||
&gatewayModelsAccountRepoStub{
|
||||
byGroup: map[int64][]service.Account{
|
||||
groupID: {
|
||||
{ID: 1, Platform: service.PlatformAntigravity, Extra: map[string]any{"mixed_scheduling": true}, Credentials: map[string]any{}},
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGemini},
|
||||
})
|
||||
|
||||
h.Models(c)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var got gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
ids := modelIDsForTest(got.Data)
|
||||
require.Contains(t, ids, "gemini-3.8-flash-high")
|
||||
for _, id := range ids {
|
||||
require.True(t, strings.HasPrefix(id, "gemini-"), "unexpected non-gemini model on gemini group: %s", id)
|
||||
}
|
||||
// geminicli 静态表独有条目不应出现(说明没有回落到默认列表)。
|
||||
require.NotContains(t, ids, "gemini-2.0-flash")
|
||||
}
|
||||
|
||||
// Codex 通过 /models?client_version= 走 CodexModels,同样应看到混合调度账号的 gemini-* 映射。
|
||||
func TestGatewayModels_CodexGeminiGroupListsAntigravityGeminiMappings(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
groupID := int64(24)
|
||||
h := newGatewayModelsHandlerForTest(
|
||||
&gatewayModelsAccountRepoStub{
|
||||
byGroup: map[int64][]service.Account{
|
||||
groupID: {
|
||||
{
|
||||
ID: 1,
|
||||
Platform: service.PlatformAntigravity,
|
||||
Extra: map[string]any{"mixed_scheduling": true},
|
||||
Credentials: map[string]any{
|
||||
"model_mapping": map[string]any{
|
||||
"gemini-3.8-flash-high": "gemini-3.8-flash-high",
|
||||
"gemini-synced-custom": "gemini-3.8-flash-high",
|
||||
"claude-sonnet-4-6": "claude-sonnet-4-6",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/models?client_version=0.147.0", nil)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGemini},
|
||||
})
|
||||
|
||||
h.CodexModels(c)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var got codexModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
slugs := codexModelSlugsForTest(got.Models)
|
||||
require.Contains(t, slugs, "gemini-3.8-flash-high")
|
||||
require.Contains(t, slugs, "gemini-synced-custom")
|
||||
require.NotContains(t, slugs, "claude-sonnet-4-6")
|
||||
require.NotContains(t, slugs, "gemini-2.0-flash")
|
||||
for _, slug := range slugs {
|
||||
require.True(t, strings.HasPrefix(slug, "gemini-"), "unexpected non-gemini model on gemini group: %s", slug)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGatewayModels_GPT6SolLunaDiscoveryRespectsGroupAndAccountRestrictions(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
selected []string
|
||||
restricted bool
|
||||
want []string
|
||||
}{
|
||||
{"selected and ordered", []string{"gpt-6.1-sol", "gpt-6-luna", "gpt-6-sol"}, false, []string{"gpt-6.1-sol", "gpt-6-luna", "gpt-6-sol"}},
|
||||
{"group excludes new models", []string{"gpt-5.6-sol"}, false, []string{"gpt-5.6-sol"}},
|
||||
{"account restricts new models", []string{"gpt-6.1-sol", "gpt-6-sol", "gpt-6-luna", "gpt-5.6-sol"}, true, []string{"gpt-5.6-sol"}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
groupID := int64(25)
|
||||
account := service.Account{ID: 1, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
||||
if tc.restricted {
|
||||
account.Credentials = map[string]any{"model_mapping": map[string]any{"gpt-5.6-sol": "gpt-5.6-sol"}}
|
||||
}
|
||||
h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{groupID: {account}}})
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, ModelAllowlist: service.GroupModelAllowlist{Enabled: true, Models: tc.selected}}})
|
||||
h.Models(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var got gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
require.Equal(t, tc.want, modelIDsForTest(got.Data))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -88,6 +88,18 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: "grok-" + strings.ReplaceAll(searchLabel, "_", "-"), Kind: service.InflightEstimatePerRequest, SearchCalls: 1})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
c.JSON(status, gin.H{"error": gin.H{"type": code, "message": message}})
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
subject, _ := middleware2.GetAuthSubjectFromContext(c)
|
||||
reqLog := requestLogger(c, "handler.gateway.web_search")
|
||||
// Audit user search query before upstream Grok web_search traffic.
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// clientCancelUpstream 模拟上游尚未响应时客户端断开:Do 期间取消请求 context,
|
||||
// 并按 net/http 的形态返回包裹 context.Canceled 的 *url.Error。
|
||||
type clientCancelUpstream struct {
|
||||
service.HTTPUpstream
|
||||
cancel context.CancelFunc
|
||||
calls atomic.Int32
|
||||
}
|
||||
|
||||
func (u *clientCancelUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
u.calls.Add(1)
|
||||
u.cancel()
|
||||
return nil, &url.Error{Op: "Post", URL: req.URL.String(), Err: context.Canceled}
|
||||
}
|
||||
|
||||
type geminiClientCancelFixture struct {
|
||||
handler *GatewayHandler
|
||||
group *service.Group
|
||||
apiKey *service.APIKey
|
||||
upstream *clientCancelUpstream
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func newGeminiClientCancelFixture(t *testing.T) *geminiClientCancelFixture {
|
||||
t.Helper()
|
||||
groupID := int64(9200)
|
||||
accountID := int64(9201)
|
||||
group := &service.Group{ID: groupID, Hydrated: true, Platform: service.PlatformGemini, Status: service.StatusActive}
|
||||
account := &service.Account{
|
||||
ID: accountID,
|
||||
Platform: service.PlatformGemini,
|
||||
Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "test-key"},
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Priority: 1,
|
||||
AccountGroups: []service.AccountGroup{{AccountID: accountID, GroupID: groupID}},
|
||||
}
|
||||
h, cleanup := newTestGatewayHandler(t, group, []*service.Account{account})
|
||||
t.Cleanup(cleanup)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
t.Cleanup(cancel)
|
||||
upstream := &clientCancelUpstream{cancel: cancel}
|
||||
h.geminiCompatService = service.NewGeminiMessagesCompatService(nil, nil, nil, nil, nil, nil, upstream, nil, &config.Config{})
|
||||
|
||||
apiKey := &service.APIKey{
|
||||
ID: 9202, UserID: 9203, GroupID: &groupID, Group: group, Status: service.StatusActive,
|
||||
User: &service.User{ID: 9203, Concurrency: 10, Balance: 100},
|
||||
}
|
||||
return &geminiClientCancelFixture{handler: h, group: group, apiKey: apiKey, upstream: upstream, ctx: ctx}
|
||||
}
|
||||
|
||||
// serve 经 ops 错误日志中间件执行请求,返回响应与 ops 队列中的条目数。
|
||||
func (f *geminiClientCancelFixture) serve(t *testing.T, route, path, body string, call func(*gin.Context)) (*httptest.ResponseRecorder, int64) {
|
||||
t.Helper()
|
||||
setupOpsErrorLogTestQueue(t, 4)
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
router := gin.New()
|
||||
router.Use(OpsErrorLoggerMiddleware(ops))
|
||||
router.POST(route, func(c *gin.Context) {
|
||||
c.Set(string(middleware.ContextKeyAPIKey), f.apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: f.apiKey.UserID, Concurrency: 10})
|
||||
call(c)
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, path, bytes.NewBufferString(body))
|
||||
req = req.WithContext(context.WithValue(f.ctx, ctxkey.Group, f.group))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
return rec, OpsErrorLogQueueLength()
|
||||
}
|
||||
|
||||
// 原生入口:上游响应前客户端断开,响应未提交,标记 499;纯客户端取消不落 ops 错误日志。
|
||||
func TestGeminiV1BetaModels_ClientCancelBeforeUpstreamResponseMarks499(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
f := newGeminiClientCancelFixture(t)
|
||||
|
||||
rec, opsQueued := f.serve(t,
|
||||
"/v1beta/models/*modelAction",
|
||||
"/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
|
||||
`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`,
|
||||
f.handler.GeminiV1BetaModels,
|
||||
)
|
||||
|
||||
require.Equal(t, int32(1), f.upstream.calls.Load())
|
||||
require.Equal(t, statusClientClosedRequest, rec.Code)
|
||||
require.Zero(t, rec.Body.Len())
|
||||
require.Zero(t, opsQueued)
|
||||
}
|
||||
|
||||
// Chat Completions 兼容入口(Gemini 分组):同一场景标记 499,不再补写 502。
|
||||
func TestGatewayChatCompletions_GeminiClientCancelBeforeUpstreamResponseMarks499(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
f := newGeminiClientCancelFixture(t)
|
||||
|
||||
rec, opsQueued := f.serve(t,
|
||||
"/v1/chat/completions",
|
||||
"/v1/chat/completions",
|
||||
`{"model":"gemini-2.5-flash","messages":[{"role":"user","content":"hi"}],"stream":true}`,
|
||||
f.handler.ChatCompletions,
|
||||
)
|
||||
|
||||
require.Equal(t, int32(1), f.upstream.calls.Load())
|
||||
require.Equal(t, statusClientClosedRequest, rec.Code)
|
||||
require.Zero(t, rec.Body.Len())
|
||||
require.Zero(t, opsQueued)
|
||||
}
|
||||
@@ -370,7 +370,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
userReleaseFunc, err := geminiConcurrency.AcquireUserSlotWithWait(c, authSubject.UserID, authSubject.Concurrency, stream, &streamStarted)
|
||||
if err != nil {
|
||||
reqLog.Warn("gemini.user_slot_acquire_failed", zap.Error(err))
|
||||
googleError(c, http.StatusTooManyRequests, err.Error())
|
||||
googleConcurrencyError(c, err, "user")
|
||||
return
|
||||
}
|
||||
// 确保请求取消时也会释放槽位,避免长连接被动中断造成泄漏
|
||||
@@ -390,6 +390,19 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(modelName, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gemini.inflight_reservation_rejected", zap.Error(err))
|
||||
status, _, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleError(c, status, message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// 3) select account (sticky session based on request body)
|
||||
// 优先使用 Gemini CLI 的会话标识(privileged-user-id + tmp 目录哈希)
|
||||
sessionHash := extractGeminiCLISessionHash(c, body)
|
||||
@@ -506,6 +519,10 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
for {
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("gemini.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, modelName, modelName, service.PlatformGemini)
|
||||
if !cls.ModelNotFound {
|
||||
@@ -598,7 +615,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
)
|
||||
if err != nil {
|
||||
reqLog.Warn("gemini.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
googleError(c, http.StatusTooManyRequests, err.Error())
|
||||
googleConcurrencyError(c, err, "account")
|
||||
return
|
||||
}
|
||||
if accountWaitCounted {
|
||||
@@ -674,7 +691,8 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
}
|
||||
// ForwardNative already wrote the response
|
||||
// 转发层已写出错误响应;客户端断开时转发层不写,响应未提交则标记 499。
|
||||
failoverClientGone(c)
|
||||
reqLog.Error("gemini.forward_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
return
|
||||
}
|
||||
@@ -836,6 +854,13 @@ func googleError(c *gin.Context, status int, message string) {
|
||||
})
|
||||
}
|
||||
|
||||
// googleConcurrencyError 以 Google 错误格式回写并发槽获取失败,状态码与文案
|
||||
// 沿用 concurrencyErrorResponse 的统一映射(客户端断开为 499)。
|
||||
func googleConcurrencyError(c *gin.Context, err error, slotType string) {
|
||||
status, _, _, message := concurrencyErrorResponse(err, slotType)
|
||||
googleError(c, status, message)
|
||||
}
|
||||
|
||||
func writeUpstreamResponse(c *gin.Context, res *service.UpstreamHTTPResult) {
|
||||
if res == nil {
|
||||
googleError(c, http.StatusBadGateway, "Empty upstream response")
|
||||
|
||||
@@ -49,6 +49,18 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) {
|
||||
if strings.TrimSpace(model) == "" {
|
||||
model = "grok-voice-latest"
|
||||
}
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: model, Kind: service.InflightEstimateAudio, AudioMode: "realtime", AudioUnits: 1})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
// Keep the HTTP response uncommitted while selecting and probing an account.
|
||||
// Realtime is not an HTTP streaming response; using reqStream=true here would
|
||||
// let the wait queue flush an SSE ping before the WebSocket handshake succeeds.
|
||||
@@ -212,6 +224,18 @@ func (h *OpenAIGatewayHandler) GrokVoice(c *gin.Context, endpoint string) {
|
||||
return
|
||||
}
|
||||
}
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, grokVoiceInflightEstimate(endpoint, body))
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
contentType := c.GetHeader("Content-Type")
|
||||
if strings.TrimSpace(contentType) == "" {
|
||||
contentType = "application/json"
|
||||
|
||||
@@ -52,6 +52,12 @@ func (h *OpenAIGatewayHandler) GrokVideoContent(c *gin.Context) {
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.GrokMediaEndpoint, requestID string) {
|
||||
platform := service.PlatformGrok
|
||||
noAccountCode, noAccountMessage := "grok_media_no_eligible_account", "No eligible Grok media accounts"
|
||||
if endpoint.IsSeedance() {
|
||||
platform = service.PlatformOpenAI
|
||||
noAccountCode, noAccountMessage = "seedance_no_eligible_account", "No eligible Seedance accounts"
|
||||
}
|
||||
streamStarted := false
|
||||
defer h.recoverResponsesPanic(c, &streamStarted)
|
||||
|
||||
@@ -99,6 +105,13 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
|
||||
contentType := c.GetHeader("Content-Type")
|
||||
requestInfo := service.ParseGrokMediaRequest(contentType, body)
|
||||
if endpoint == service.SeedanceEndpointCreate {
|
||||
requestInfo, err = service.ParseSeedanceRequest(body)
|
||||
if err != nil {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
requestModel := requestInfo.Model
|
||||
routingModel := service.NormalizeGrokMediaModelForEndpoint(endpoint, requestModel, requestInfo.HasInputImage())
|
||||
if endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) == "" {
|
||||
@@ -160,6 +173,18 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, grokMediaInflightEstimate(endpoint, routingModel, requestInfo, body))
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
sessionSeed := body
|
||||
if len(sessionSeed) == 0 && strings.TrimSpace(requestID) != "" {
|
||||
sessionSeed = []byte(requestID)
|
||||
@@ -215,8 +240,8 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
var selection *service.AccountSelectionResult
|
||||
var scheduleDecision service.OpenAIAccountScheduleDecision
|
||||
if boundLookupAccountID > 0 {
|
||||
selection, scheduleDecision, err = h.gatewayService.SelectGrokMediaVideoRequestAccount(
|
||||
requestCtx, apiKey.GroupID, sessionHash, boundLookupAccountID, routingModel,
|
||||
selection, scheduleDecision, err = h.gatewayService.SelectMediaVideoRequestAccount(
|
||||
requestCtx, apiKey.GroupID, sessionHash, boundLookupAccountID, routingModel, platform,
|
||||
)
|
||||
} else {
|
||||
selection, scheduleDecision, err = h.gatewayService.SelectAccountWithSchedulerForCapability(
|
||||
@@ -231,7 +256,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
service.PlatformGrok,
|
||||
platform,
|
||||
)
|
||||
}
|
||||
// Own an eagerly acquired slot before any rejection or eligibility probe.
|
||||
@@ -256,11 +281,11 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
if endpoint.IsGenerationRequest() && errors.Is(err, service.ErrNoAvailableAccounts) &&
|
||||
(len(failedAccountIDs) == 0 || (mediaEligibilityRejected && lastFailoverErr == nil)) {
|
||||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, noAccountCode, noAccountMessage)
|
||||
return
|
||||
}
|
||||
if len(failedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, routingModel, service.PlatformGrok)
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, routingModel, platform)
|
||||
if !cls.ModelNotFound {
|
||||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||||
}
|
||||
@@ -277,10 +302,10 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
if selection == nil || selection.Account == nil {
|
||||
if endpoint.IsGenerationRequest() {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, noAccountCode, noAccountMessage)
|
||||
return
|
||||
}
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, routingModel, service.PlatformGrok)
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, routingModel, platform)
|
||||
if !cls.ModelNotFound {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
}
|
||||
@@ -309,7 +334,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
)
|
||||
|
||||
account := selection.Account
|
||||
if endpoint.IsGenerationRequest() {
|
||||
if endpoint.IsGenerationRequest() && !endpoint.IsSeedance() {
|
||||
eligible, eligibilityReason, eligibilityErr := h.ensureGrokMediaAccountEligibility(requestCtx, account)
|
||||
if !eligible {
|
||||
releaseAccount()
|
||||
@@ -360,6 +385,9 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
writerSizeBeforeForward := c.Writer.Size()
|
||||
result, err := func() (*service.OpenAIForwardResult, error) {
|
||||
defer releaseAccount()
|
||||
if endpoint.IsSeedance() {
|
||||
return h.gatewayService.ForwardSeedance(requestCtx, c, account, endpoint, requestID, body)
|
||||
}
|
||||
return h.gatewayService.ForwardGrokMedia(requestCtx, c, account, endpoint, requestID, body, contentType)
|
||||
}()
|
||||
|
||||
@@ -490,7 +518,11 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
}
|
||||
// Status poll OR content download can observe official done+video.url.
|
||||
// Both paths share the same claim key so the customer is charged once.
|
||||
if endpoint == service.GrokMediaEndpointVideoStatus || endpoint == service.GrokMediaEndpointVideoContent {
|
||||
if endpoint == service.SeedanceEndpointStatus {
|
||||
if billResult := prepareSeedanceCompletionBilling(requestCtx, h, apiKey, subject, requestID, result); billResult != nil {
|
||||
recordGrokMediaUsage(c, h, reqLog, apiKey, subject, subscription, account, billResult, billResult.Model, body, requestID)
|
||||
}
|
||||
} else if endpoint == service.GrokMediaEndpointVideoStatus || endpoint == service.GrokMediaEndpointVideoContent {
|
||||
taskID := strings.TrimSpace(requestID)
|
||||
if billResult := prepareGrokVideoCompletionBilling(requestCtx, h, reqLog, apiKey, subject, taskID, result); billResult != nil {
|
||||
recordGrokMediaUsage(c, h, reqLog, apiKey, subject, subscription, account, billResult, billResult.Model, body, taskID)
|
||||
@@ -521,6 +553,9 @@ func (h *OpenAIGatewayHandler) ensureGrokMediaAccountEligibility(ctx context.Con
|
||||
}
|
||||
|
||||
func grokMediaRequiredCapability(endpoint service.GrokMediaEndpoint) service.OpenAIEndpointCapability {
|
||||
if endpoint.IsSeedance() {
|
||||
return service.OpenAIEndpointCapabilitySeedance
|
||||
}
|
||||
if endpoint.IsGenerationRequest() {
|
||||
return service.OpenAIEndpointCapabilityGrokMediaGeneration
|
||||
}
|
||||
@@ -539,7 +574,7 @@ func grokMediaScheduleModel(account *service.Account, routingModel string, resul
|
||||
|
||||
func isGrokVideoCreateEndpoint(endpoint service.GrokMediaEndpoint) bool {
|
||||
switch endpoint {
|
||||
case service.GrokMediaEndpointVideosGenerations,
|
||||
case service.SeedanceEndpointCreate, service.GrokMediaEndpointVideosGenerations,
|
||||
service.GrokMediaEndpointVideosEdits,
|
||||
service.GrokMediaEndpointVideosExtensions:
|
||||
return true
|
||||
@@ -712,7 +747,7 @@ func recordGrokMediaUsage(
|
||||
}
|
||||
// Async video: force durable task request id and release claim if billing fails.
|
||||
videoTaskID := ""
|
||||
if result != nil && result.VideoCount > 0 {
|
||||
if result != nil && (result.VideoCount > 0 || strings.HasPrefix(result.ResponseID, "seedance:")) {
|
||||
videoTaskID = strings.TrimSpace(firstNonEmptyString(requestID, result.ResponseID))
|
||||
if stable := service.StableGrokVideoBillingRequestID(firstNonEmptyString(result.ResponseID, requestID)); stable != "" {
|
||||
result.RequestID = stable
|
||||
|
||||
@@ -120,10 +120,22 @@ func (s *grokMediaSlotsCache) assertReleased(t *testing.T) {
|
||||
|
||||
type grokMediaSlotBindings struct {
|
||||
testutil.StubGatewayCache
|
||||
owner int64
|
||||
writes int
|
||||
key string
|
||||
billed map[string]bool
|
||||
owner int64
|
||||
writes int
|
||||
key string
|
||||
billed map[string]bool
|
||||
pending map[string][]byte
|
||||
}
|
||||
|
||||
func (s *grokMediaSlotBindings) SetGrokVideoPendingBilling(_ context.Context, key string, body []byte, _ time.Duration) error {
|
||||
if s.pending == nil {
|
||||
s.pending = make(map[string][]byte)
|
||||
}
|
||||
s.pending[key] = append([]byte(nil), body...)
|
||||
return nil
|
||||
}
|
||||
func (s *grokMediaSlotBindings) GetGrokVideoPendingBilling(_ context.Context, key string) ([]byte, error) {
|
||||
return s.pending[key], nil
|
||||
}
|
||||
|
||||
func (s *grokMediaSlotBindings) GetSessionAccountID(_ context.Context, groupID int64, key string) (int64, error) {
|
||||
@@ -186,7 +198,7 @@ func (p grokMediaSlotProber) ProbeMediaEligibility(ctx context.Context, id int64
|
||||
return p(ctx, id)
|
||||
}
|
||||
|
||||
func newGrokMediaSlotHandler(t *testing.T, oauth, mismatch bool) (*OpenAIGatewayHandler, *grokMediaSlotsCache, *grokMediaSlotBindings, *grokMediaSlotUpstream) {
|
||||
func newGrokMediaSlotHandler(t *testing.T, oauth, mismatch bool, platforms ...string) (*OpenAIGatewayHandler, *grokMediaSlotsCache, *grokMediaSlotBindings, *grokMediaSlotUpstream) {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
accounts := make([]service.Account, 3)
|
||||
@@ -194,6 +206,11 @@ func newGrokMediaSlotHandler(t *testing.T, oauth, mismatch bool) (*OpenAIGateway
|
||||
accounts[i] = service.Account{ID: int64(i + 1), Platform: service.PlatformGrok, Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive, Schedulable: true, Concurrency: 50, Priority: i,
|
||||
GroupIDs: []int64{24}, Credentials: map[string]any{"api_key": "test-key", "access_token": "test-token"}}
|
||||
if len(platforms) > 0 {
|
||||
accounts[i].Platform = platforms[0]
|
||||
accounts[i].Credentials["base_url"] = "https://ark.cn-beijing.volces.com/api/v3"
|
||||
accounts[i].Credentials["openai_capabilities"] = []string{"seedance"}
|
||||
}
|
||||
if oauth {
|
||||
accounts[i].Type = service.AccountTypeOAuth
|
||||
accounts[i].Credentials["refresh_token"] = "test-refresh"
|
||||
|
||||
@@ -92,6 +92,8 @@ type modelPlazaGroup struct {
|
||||
// 不取分组/用户专属倍率。
|
||||
ImageRateIndependent bool `json:"image_rate_independent"`
|
||||
ImageRateMultiplier float64 `json:"image_rate_multiplier"`
|
||||
VideoRateIndependent bool `json:"video_rate_independent"`
|
||||
VideoRateMultiplier float64 `json:"video_rate_multiplier"`
|
||||
// 分组是否启用长上下文阶梯计费;关闭时模型实付列只展示最低档/基础价。
|
||||
LongContextPricingEnabled bool `json:"long_context_pricing_enabled"`
|
||||
Models []modelPlazaModel `json:"models"`
|
||||
@@ -211,6 +213,8 @@ func toModelPlazaGroupDTO(g *service.PlazaGroup, userRates map[int64]float64) mo
|
||||
IsExclusive: g.IsExclusive,
|
||||
ImageRateIndependent: g.ImageRateIndependent,
|
||||
ImageRateMultiplier: g.ImageRateMultiplier,
|
||||
VideoRateIndependent: g.VideoRateIndependent,
|
||||
VideoRateMultiplier: g.VideoRateMultiplier,
|
||||
LongContextPricingEnabled: g.LongContextPricingEnabled,
|
||||
Models: models,
|
||||
}
|
||||
|
||||
@@ -89,6 +89,7 @@ func TestToModelPlazaGroupDTO_UserRateAndFieldWhitelist(t *testing.T) {
|
||||
g := service.PlazaGroup{
|
||||
ID: 2, Name: "vip", Description: "d", Platform: "anthropic",
|
||||
SubscriptionType: "standard", RateMultiplier: 1, IsExclusive: true,
|
||||
VideoRateIndependent: true, VideoRateMultiplier: 0.7,
|
||||
Models: []service.PlazaModel{{
|
||||
Name: "claude-sonnet",
|
||||
Platform: "anthropic",
|
||||
@@ -115,11 +116,14 @@ func TestToModelPlazaGroupDTO_UserRateAndFieldWhitelist(t *testing.T) {
|
||||
"rate_multiplier", "user_rate_multiplier", "is_exclusive", "models",
|
||||
"peak_rate_enabled", "peak_start", "peak_end", "peak_rate_multiplier",
|
||||
"image_rate_independent", "image_rate_multiplier", "long_context_pricing_enabled",
|
||||
"video_rate_independent", "video_rate_multiplier",
|
||||
} {
|
||||
_, exists := decoded[key]
|
||||
require.Truef(t, exists, "plaza group DTO must expose %q", key)
|
||||
}
|
||||
require.InDelta(t, 0.5, decoded["user_rate_multiplier"].(float64), 1e-9)
|
||||
require.Equal(t, true, decoded["video_rate_independent"])
|
||||
require.Equal(t, 0.7, decoded["video_rate_multiplier"])
|
||||
|
||||
// 模型条目:pricing + official_pricing 并存;official 缺失字段输出 null 而非省略
|
||||
models := decoded["models"].([]any)
|
||||
|
||||
@@ -109,6 +109,18 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: requestedModel, BodyBytes: len(body), MaxTokens: requestMaxOutputTokens(body), Kind: service.InflightEstimateToken, SearchCalls: 1})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
searchID := strings.TrimSpace(gjson.GetBytes(body, "id").String())
|
||||
sessionHash := h.gatewayService.GenerateSessionHashWithFallback(c, nil, searchID)
|
||||
profitVetoCount := 0
|
||||
@@ -197,6 +209,10 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if !errors.As(err, &failoverErr) {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account, openAIAccountScheduleModel(c, account, requestedModel, false, result), false, nil, err)
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai_alpha_search.forward_aborted_client_disconnected", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
return
|
||||
}
|
||||
if c.Writer.Size() == writerSizeBeforeForward {
|
||||
h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
|
||||
}
|
||||
|
||||
@@ -146,6 +146,19 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("openai_chat_completions.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
sessionHash := h.gatewayService.GenerateSessionHash(c, body)
|
||||
promptCacheKey := h.gatewayService.ExtractSessionID(c, body)
|
||||
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func (h *OpenAIGatewayHandler) cyberPolicyLogOnly(c *gin.Context, apiKey *service.APIKey) bool {
|
||||
return h != nil && c != nil && c.Request != nil && h.gatewayService.CyberPolicyLogOnly(c.Request.Context(), apiKey)
|
||||
}
|
||||
|
||||
// Both HTTP and WebSocket admission skip existing blocks for trusted users.
|
||||
func (h *OpenAIGatewayHandler) findBlockedCyberSessionForAPIKey(c *gin.Context, apiKey *service.APIKey, body []byte) string {
|
||||
if apiKey == nil || h.cyberPolicyLogOnly(c, apiKey) {
|
||||
return ""
|
||||
}
|
||||
return findBlockedCyberSessionKey(c.Request.Context(), h.gatewayService, apiKey.ID, c, body)
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestCyberAllowlistedUserBypassesExistingBlocksAndContinuesWebSocket(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := coderws.Accept(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer func() { _ = conn.CloseNow() }()
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
|
||||
defer cancel()
|
||||
for _, response := range []string{
|
||||
`{"type":"response.failed","response":{"id":"resp_cyber","model":"gpt-5.1","error":{"code":"cyber_policy","message":"blocked by upstream"}}}`,
|
||||
`{"type":"response.completed","response":{"id":"resp_ok","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`,
|
||||
} {
|
||||
if _, _, err := conn.Read(ctx); err != nil {
|
||||
return
|
||||
}
|
||||
if err := conn.Write(ctx, coderws.MessageText, []byte(response)); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
// Let the downstream close so the completed frame can drain normally.
|
||||
_, _, _ = conn.Read(ctx)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
harness := newOpenAIWSPassthroughHandlerHarness(t, upstream.URL, map[string]string{
|
||||
service.SettingKeyCyberPolicyUserAllowlist: "1751",
|
||||
})
|
||||
payload := []byte(`{"type":"response.create","model":"gpt-5.1","prompt_cache_key":"trusted-session","input":"test"}`)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(string(payload)))
|
||||
explicitKey := service.CyberSessionExplicitBlockKey(harness.apiKey.ID, c, payload)
|
||||
store, ok := harness.gatewayCache.(service.CyberSessionBlockStore)
|
||||
require.True(t, ok)
|
||||
require.NoError(t, store.SetCyberSessionBlocked(context.Background(), "", []string{explicitKey}, time.Minute))
|
||||
for _, format := range []cyberSessionBlockFormat{cyberBlockFormatResponses, cyberBlockFormatChat, cyberBlockFormatAnthropic} {
|
||||
require.False(t, harness.handler.rejectIfCyberSessionBlocked(c, harness.apiKey, payload, "gpt-5.1", format))
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
require.NoError(t, harness.clientConn.Write(ctx, coderws.MessageText, payload))
|
||||
_, event, err := harness.clientConn.Read(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "cyber_policy", gjson.GetBytes(event, "response.error.code").String())
|
||||
require.Eventually(t, func() bool {
|
||||
logs := harness.moderationRepo.logSnapshot()
|
||||
return len(logs) == 1 && logs[0].Mode == service.ContentModerationModeCyberLogOnly
|
||||
}, 3*time.Second, 10*time.Millisecond)
|
||||
matched, err := store.FindCyberSessionBlocked(ctx, service.CyberSessionTranscriptBlockKeys(harness.apiKey.ID, payload))
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, matched, "the cyber response must not write new transcript blocks")
|
||||
require.NoError(t, harness.clientConn.Write(ctx, coderws.MessageText, payload))
|
||||
_, event, err = harness.clientConn.Read(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
|
||||
_ = harness.clientConn.CloseNow()
|
||||
select {
|
||||
case <-harness.handlerDone:
|
||||
case <-ctx.Done():
|
||||
t.Fatal("handler did not exit")
|
||||
}
|
||||
}
|
||||
@@ -108,6 +108,18 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: reqModel, BodyBytes: len(body), MaxTokens: 1, Kind: service.InflightEstimateToken})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
profitVetoCount := 0
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
|
||||
@@ -34,6 +34,7 @@ import (
|
||||
|
||||
// OpenAIGatewayHandler handles OpenAI API gateway requests
|
||||
type OpenAIGatewayHandler struct {
|
||||
compositeResolver *service.CompositeRouteResolver
|
||||
gatewayService *service.OpenAIGatewayService
|
||||
billingCacheService *service.BillingCacheService
|
||||
apiKeyService *service.APIKeyService
|
||||
@@ -257,13 +258,21 @@ func usageRecordContext(parent context.Context, base context.Context) context.Co
|
||||
return base
|
||||
}
|
||||
|
||||
func wrapUsageRecordTaskContext(parent context.Context, task service.UsageRecordTask) service.UsageRecordTask {
|
||||
// wrapUsageRecordTaskContext 包装计费任务:复制请求级 context 值,并接管请求的在途余额预留引用。
|
||||
// 返回的 abandon 在任务未被执行(被丢弃)时必须调用以归还预留引用;任务执行结束时(含 panic)
|
||||
// 自动归还,此时余额缓存已在计费路径中同步扣减。
|
||||
func wrapUsageRecordTaskContext(parent context.Context, task service.UsageRecordTask) (service.UsageRecordTask, func()) {
|
||||
if task == nil {
|
||||
return nil
|
||||
return nil, func() {}
|
||||
}
|
||||
done := func() {}
|
||||
if parent != nil {
|
||||
done = service.InflightReservationFromContext(parent).Acquire()
|
||||
}
|
||||
return func(ctx context.Context) {
|
||||
defer done()
|
||||
task(usageRecordContext(parent, ctx))
|
||||
}
|
||||
}, done
|
||||
}
|
||||
|
||||
func openAICompatibleRequestPlatform(ctx context.Context, apiKey *service.APIKey) string {
|
||||
@@ -606,6 +615,19 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("openai.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// Generate session hash (header first; fallback to prompt_cache_key)
|
||||
sessionHash := h.gatewayService.GenerateSessionHash(c, sessionHashBody)
|
||||
if h.rejectIfCyberSessionBlocked(c, apiKey, sessionHashBody, reqModel, cyberBlockFormatResponses) {
|
||||
@@ -1242,6 +1264,18 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.anthropicStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
sessionHash := h.gatewayService.GenerateSessionHash(c, body)
|
||||
promptCacheKey := h.gatewayService.ExtractSessionID(c, body)
|
||||
sessionHash, promptCacheKey = resolveOpenAIMessagesMetadataSession(c, sessionHash, promptCacheKey, reqModel, body)
|
||||
@@ -2391,7 +2425,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
// 分组级模型白名单:首帧校验客户端模型,不通过则关闭连接并标记运维原因。
|
||||
// 必须在 ensureCompositeTargetPlatform(合成路由改写)之前执行。
|
||||
// 必须在合成路由解析和上游模型映射之前执行。
|
||||
// 与 HTTP 准入一致:帧内重复 model 键/大小写变体可能被上游按末值绑定,
|
||||
// 全部候选值逐一校验,任一未命中即拒绝。
|
||||
if blocked := blockedModelAllowlistCandidate(apiKey.Group, requestmodel.FromBodyCandidates("", "application/json", firstMessage)); blocked != "" {
|
||||
@@ -2400,7 +2434,21 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, fmt.Sprintf("Model %q is not available for this group", blocked))
|
||||
return
|
||||
}
|
||||
ensureCompositeTargetPlatform(c, apiKey, reqModel)
|
||||
// Keep the client model in the frame for admission, response identity and usage.
|
||||
// Apply the resolved upstream model through MapRequestModel on every turn.
|
||||
wsRouteModel := reqModel
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
decision, resolveErr := h.compositeResolver.Resolve(c.Request.Context(), apiKey.Group.ID, reqModel, service.CompositeRouteEndpointResponses)
|
||||
if resolveErr != nil {
|
||||
reqLog.Error("openai.websocket_composite_route_failed", zap.Error(resolveErr))
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "Failed to resolve composite model route")
|
||||
return
|
||||
}
|
||||
if decision.Matched {
|
||||
c.Request = c.Request.WithContext(service.WithCompositeRouteDecision(c.Request.Context(), decision))
|
||||
wsRouteModel = decision.UpstreamModel
|
||||
}
|
||||
}
|
||||
ctx = c.Request.Context()
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
platform, ok := service.ResolvedTargetPlatformFromContext(ctx)
|
||||
@@ -2440,7 +2488,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
|
||||
// The first response.create frame is available here, so explicit IDs are
|
||||
// checked directly and body-derived sessions use the coarse scope gate.
|
||||
if cyberBlockKey := findBlockedCyberSessionKey(c.Request.Context(), h.gatewayService, apiKey.ID, c, firstMessage); cyberBlockKey != "" {
|
||||
if cyberBlockKey := h.findBlockedCyberSessionForAPIKey(c, apiKey, firstMessage); cyberBlockKey != "" {
|
||||
writeCyberSessionBlockedWSError(c.Request.Context(), wsConn)
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "session blocked by cyber-security policy")
|
||||
h.enqueueCyberSessionBlockedOpsEntry(c, apiKey, reqModel, cyberBlockKey)
|
||||
@@ -2464,8 +2512,8 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 解析渠道级模型映射
|
||||
channelMappingWS, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, reqModel)
|
||||
wsForwardModel := openAIChannelForwardModel(channelMappingWS, reqModel)
|
||||
channelMappingWS, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, wsRouteModel)
|
||||
wsForwardModel := openAIChannelForwardModel(channelMappingWS, wsRouteModel)
|
||||
|
||||
var currentUserRelease func()
|
||||
var currentAccountRelease func()
|
||||
@@ -2526,6 +2574,31 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(会话级):按首帧估算一次,会话期间续期;每轮计费任务接管引用,
|
||||
// 会话结束且所有轮次扣减落地后释放。
|
||||
inflightCtx, inflightDone, inflightErr := reserveInflightBalanceCtx(ctx, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, firstMessage))
|
||||
if inflightErr != nil {
|
||||
reqLog.Info("openai.websocket_inflight_reservation_rejected", zap.Error(inflightErr))
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "billing check failed")
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
ctx = inflightCtx
|
||||
|
||||
// A WebSocket may outlive a key's remaining spending window. Recheck
|
||||
// after acquiring turn slots, including the first account-selection wait.
|
||||
// Restrict this extra check to the opt-in mode so standard-mode RPM checks
|
||||
// are not charged a second time for the same request.
|
||||
checkSimpleModeTurnBilling := func() error {
|
||||
if h.cfg == nil || h.cfg.RunMode != config.RunModeSimple || !h.cfg.SimpleModeKeyRateLimitEnabled {
|
||||
return nil
|
||||
}
|
||||
if err := h.billingCacheService.CheckBillingEligibility(ctx, apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(ctx, apiKey)); err != nil {
|
||||
return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "billing check failed", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
sessionHash := h.gatewayService.GenerateSessionHashWithFallback(
|
||||
c,
|
||||
firstMessage,
|
||||
@@ -2783,6 +2856,11 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
// turn 级定价:首轮回退到 TurnStarted 的所属 turn 时刻;后续 turn 由
|
||||
// BeforeTurn 重新冻结 pricingAt 并按最新门复核当前账号。
|
||||
var turnPricing openAIWSTurnPricing
|
||||
// Passthrough ingress does not invoke BeforeTurn for the first frame.
|
||||
if err := checkSimpleModeTurnBilling(); err != nil {
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "billing check failed")
|
||||
return
|
||||
}
|
||||
hooks := &service.OpenAIWSIngressHooks{
|
||||
ClientLifecycleContext: clientLifecycleCtx,
|
||||
InitialRequestModel: reqModel,
|
||||
@@ -2798,7 +2876,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
// 连接级 cyber session gate 也在 BeforeRequest 先执行,使 native 与
|
||||
// passthrough ingress 都能在 BeforeTurn 及上游写入前无副作用地拒绝。
|
||||
// BeforeTurn 中保留同一检查作为防御式兜底。
|
||||
if cyberBlockedThisConn {
|
||||
if cyberBlockedThisConn && !h.cyberPolicyLogOnly(c, apiKey) {
|
||||
return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, cyberSessionBlockedClientMsg, nil)
|
||||
}
|
||||
if turn == 1 {
|
||||
@@ -2837,7 +2915,16 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
model = reqModel
|
||||
}
|
||||
setOpsRequestContext(c, model, true)
|
||||
mapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, model)
|
||||
routeModel := model
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
// The account, target platform and route context are connection-scoped.
|
||||
// A different public model needs a new connection and fresh resolution.
|
||||
if model != reqModel {
|
||||
return "", newOpenAIWSUnsupportedModelSwitchError(model)
|
||||
}
|
||||
routeModel = wsRouteModel
|
||||
}
|
||||
mapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(ctx, apiKey.GroupID, routeModel)
|
||||
mappedModelUnchanged := false
|
||||
if previous := turnChannelMapping.Load(); previous != nil && previous.turn < turn {
|
||||
mappedModelUnchanged = strings.TrimSpace(previous.mapping.MappedModel) == strings.TrimSpace(mapping.MappedModel)
|
||||
@@ -2850,7 +2937,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
},
|
||||
BeforeTurn: func(turn int) error {
|
||||
// turn==1 的会话屏蔽已由握手层检查覆盖;连接内 flag 只拦截后续 turn。
|
||||
if cyberBlockedThisConn {
|
||||
if cyberBlockedThisConn && !h.cyberPolicyLogOnly(c, apiKey) {
|
||||
return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, cyberSessionBlockedClientMsg, nil)
|
||||
}
|
||||
// 长连接跨峰谷/倍率刷新防护:每个 turn 按当前时刻重装门并复核
|
||||
@@ -2893,7 +2980,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
}
|
||||
currentUserRelease = wrapReleaseOnDone(ctx, userReleaseFunc)
|
||||
currentAccountRelease = wrapReleaseOnDone(ctx, accountReleaseFunc)
|
||||
return nil
|
||||
return checkSimpleModeTurnBilling()
|
||||
},
|
||||
AfterTurn: func(turn int, result *service.OpenAIForwardResult, turnErr error) {
|
||||
turnStart := getTurnStart(turn)
|
||||
@@ -2930,7 +3017,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
cyberBlockedThisConn, cyberBlockPendingAfterFailover = advanceOpenAIWSCyberBlockState(
|
||||
cyberBlockedThisConn,
|
||||
cyberBlockPendingAfterFailover,
|
||||
cyberMarked,
|
||||
cyberMarked && !h.cyberPolicyLogOnly(c, apiKey),
|
||||
turnErr,
|
||||
)
|
||||
if turnErr != nil {
|
||||
@@ -3247,9 +3334,12 @@ func (h *OpenAIGatewayHandler) submitUsageRecordTask(parent context.Context, tas
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
task, abandon := wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped {
|
||||
if mode.Dropped() {
|
||||
abandon()
|
||||
}
|
||||
return
|
||||
}
|
||||
// 池已停止(进程关停窗口):计费任务不能静默丢失,降级为内联同步执行。
|
||||
@@ -3286,7 +3376,7 @@ func (h *OpenAIGatewayHandler) submitMandatoryUsageRecordTask(parent context.Con
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
task, _ = wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() {
|
||||
return
|
||||
@@ -3340,6 +3430,7 @@ func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverE
|
||||
h.handleFailoverExhaustedSimple(c, http.StatusBadGateway, streamStarted)
|
||||
return
|
||||
}
|
||||
copyFailoverRetryAfter(c, failoverErr.ResponseHeaders)
|
||||
if failoverErr.IsOpenAIRequestBodyTooLarge() {
|
||||
service.SetOpsUpstreamError(c, http.StatusRequestEntityTooLarge, service.OpenAIRequestBodyTooLargeClientMessage, "")
|
||||
h.handleStreamingAwareError(
|
||||
@@ -3359,7 +3450,19 @@ func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverE
|
||||
h.handleStreamingAwareError(c, http.StatusBadRequest, "invalid_request_error", message, streamStarted)
|
||||
return
|
||||
}
|
||||
copyFailoverRetryAfter(c, failoverErr.ResponseHeaders)
|
||||
if failoverErr.Reason == service.OpenAIImagesInsufficientBalanceReason {
|
||||
status := failoverErr.ClientStatusCode
|
||||
if status <= 0 {
|
||||
status = http.StatusPaymentRequired
|
||||
}
|
||||
message := strings.TrimSpace(failoverErr.ClientMessage)
|
||||
if message == "" {
|
||||
message = service.OpenAIImagesInsufficientBalanceMessage
|
||||
}
|
||||
service.SetOpsUpstreamError(c, failoverErr.StatusCode, message, "")
|
||||
h.handleStreamingAwareErrorWithCode(c, status, "upstream_error", service.OpenAIImagesInsufficientBalanceCode, message, streamStarted, false)
|
||||
return
|
||||
}
|
||||
if failoverErr.IsCredentialFailure() {
|
||||
status, message := credentialFailoverClientResponse(failoverErr)
|
||||
h.handleStreamingAwareError(c, status, "upstream_error", message, streamStarted)
|
||||
@@ -4013,7 +4116,7 @@ func (h *OpenAIGatewayHandler) rejectIfCyberSessionBlocked(c *gin.Context, apiKe
|
||||
if enabled, _ := h.gatewayService.CyberSessionBlockRuntime(c.Request.Context()); !enabled {
|
||||
return false
|
||||
}
|
||||
key := findBlockedCyberSessionKey(c.Request.Context(), h.gatewayService, apiKey.ID, c, body)
|
||||
key := h.findBlockedCyberSessionForAPIKey(c, apiKey, body)
|
||||
if key == "" {
|
||||
return false
|
||||
}
|
||||
@@ -4210,7 +4313,8 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey
|
||||
ClientIP: clientIPStr,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if gwSvc != nil && apiKey != nil {
|
||||
cyberLogOnly := h.cyberPolicyLogOnly(c, apiKey)
|
||||
if gwSvc != nil && apiKey != nil && !cyberLogOnly {
|
||||
plan := buildCyberSessionBlockWritePlan(apiKey.ID, c, cyberBlockBody)
|
||||
if len(plan.keys) > 0 {
|
||||
blockCtx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
|
||||
@@ -4223,6 +4327,7 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey
|
||||
defer cancel()
|
||||
if cmSvc != nil {
|
||||
cmSvc.RecordCyberPolicyEvent(ctx, service.CyberPolicyRecordInput{
|
||||
LogOnly: cyberLogOnly,
|
||||
RequestID: requestID,
|
||||
UserID: userID,
|
||||
UserEmail: userEmail,
|
||||
|
||||
@@ -381,6 +381,20 @@ func TestOpenAIEnsureForwardErrorResponse_AfterDeltaAppendsSingleValidResponseFa
|
||||
require.Equal(t, 1, errorEvents)
|
||||
}
|
||||
|
||||
func TestOpenAIEnsureForwardErrorResponse_PreservesCommittedStreamError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil)
|
||||
body := "event: error\ndata: {\"type\":\"error\",\"code\":\"upstream_stream_read_error\",\"message\":\"Upstream response stream was interrupted\",\"param\":null,\"sequence_number\":0}\n\n"
|
||||
_, err := c.Writer.WriteString(body)
|
||||
require.NoError(t, err)
|
||||
service.MarkResponseCommitted(c)
|
||||
h := &OpenAIGatewayHandler{}
|
||||
require.False(t, h.ensureForwardErrorResponse(c, true))
|
||||
require.Equal(t, body, w.Body.String(), "preserve the service error without a second terminal event")
|
||||
}
|
||||
|
||||
func TestOpenAIEnsureForwardErrorResponse_CompactKeepaliveOnlyWritesResponseFailed(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
@@ -1916,7 +1930,12 @@ func newOpenAIWSHandlerTestServer(t *testing.T, h *OpenAIGatewayHandler, subject
|
||||
}
|
||||
|
||||
type openAIResponsesWSUsageLogCase struct {
|
||||
firstPayload string
|
||||
simpleModeRejectAtRead int64
|
||||
compositeResolver *service.CompositeRouteResolver
|
||||
accountPlatform string
|
||||
closeReason string
|
||||
closeStatus coderws.StatusCode
|
||||
firstPayload string
|
||||
// midPayload 在首个 turn 完成后发送(如 session.update),上游桩会为它
|
||||
// 回一个 response.completed,客户端按普通事件读取。
|
||||
midPayload string
|
||||
@@ -2862,6 +2881,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
upstreamErrCh := make(chan error, 1)
|
||||
var channelSvc *service.ChannelService
|
||||
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if tc.accountPlatform == service.PlatformGrok {
|
||||
payload, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
upstreamErrCh <- err
|
||||
return
|
||||
}
|
||||
upstreamPayloadCh <- payload
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = fmt.Fprintf(w, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_grok_test\",\"model\":%q,\"usage\":{\"input_tokens\":2,\"output_tokens\":1}}}\n\n", gjson.GetBytes(payload, "model").String())
|
||||
return
|
||||
}
|
||||
conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
|
||||
CompressionMode: coderws.CompressionContextTakeover,
|
||||
})
|
||||
@@ -2929,6 +2959,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
"openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
|
||||
},
|
||||
}
|
||||
if tc.accountPlatform != "" {
|
||||
account.Platform = tc.accountPlatform
|
||||
}
|
||||
if strings.TrimSpace(tc.ingressMode) != "" {
|
||||
account.Extra["openai_apikey_responses_websockets_v2_mode"] = tc.ingressMode
|
||||
}
|
||||
@@ -2963,7 +2996,13 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
}, nil, nil, nil, nil)
|
||||
}
|
||||
|
||||
billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
var keyRepo service.APIKeyRepository
|
||||
if tc.simpleModeRejectAtRead > 0 {
|
||||
cfg.SimpleModeKeyRateLimitEnabled = true
|
||||
keyRepo = &simpleModeWSRateLimitRepo{rejectAt: tc.simpleModeRejectAtRead}
|
||||
}
|
||||
billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, keyRepo, nil, nil, cfg, nil)
|
||||
t.Cleanup(billingCacheSvc.Stop)
|
||||
gatewaySvc := service.NewOpenAIGatewayService(
|
||||
accountRepo,
|
||||
usageRepo,
|
||||
@@ -2978,7 +3017,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
service.NewBillingService(cfg, nil),
|
||||
nil,
|
||||
billingCacheSvc,
|
||||
nil,
|
||||
&compositeWSHTTPUpstream{},
|
||||
&service.DeferredService{},
|
||||
nil,
|
||||
nil,
|
||||
@@ -2998,6 +3037,8 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
},
|
||||
}
|
||||
h := &OpenAIGatewayHandler{
|
||||
cfg: cfg,
|
||||
compositeResolver: tc.compositeResolver,
|
||||
gatewayService: gatewaySvc,
|
||||
billingCacheService: billingCacheSvc,
|
||||
apiKeyService: &service.APIKeyService{},
|
||||
@@ -3009,6 +3050,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
GroupID: &groupID,
|
||||
User: &service.User{ID: 1701, Status: service.StatusActive},
|
||||
}
|
||||
if tc.simpleModeRejectAtRead > 0 {
|
||||
apiKey.RateLimit5h = 1
|
||||
}
|
||||
if tc.group != nil {
|
||||
apiKey.Group = tc.group
|
||||
}
|
||||
@@ -3043,6 +3087,12 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
cancelWrite()
|
||||
require.NoError(t, err)
|
||||
|
||||
if tc.closeReason == "" {
|
||||
tc.closeReason = "not available for this group"
|
||||
}
|
||||
if tc.closeStatus == 0 {
|
||||
tc.closeStatus = coderws.StatusPolicyViolation
|
||||
}
|
||||
if tc.firstFrameCloseExpected {
|
||||
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
_, _, readErr := clientConn.Read(readCtx)
|
||||
@@ -3050,8 +3100,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
require.Error(t, readErr, "first frame should have been rejected with a close")
|
||||
var closeErr coderws.CloseError
|
||||
require.ErrorAs(t, readErr, &closeErr)
|
||||
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
|
||||
require.Contains(t, closeErr.Reason, "not available for this group")
|
||||
status := tc.closeStatus
|
||||
if status == 0 {
|
||||
status = coderws.StatusPolicyViolation
|
||||
}
|
||||
require.Equal(t, status, closeErr.Code)
|
||||
reason := tc.closeReason
|
||||
if reason == "" {
|
||||
reason = "not available for this group"
|
||||
}
|
||||
require.Contains(t, closeErr.Reason, reason)
|
||||
require.Empty(t, upstreamPayloadCh, "rejected first frame must not reach upstream")
|
||||
_ = clientConn.CloseNow()
|
||||
return openAIResponsesWSUsageLogResult{}
|
||||
}
|
||||
@@ -3085,8 +3144,17 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
require.Error(t, readErr, "second turn should have been rejected with a close")
|
||||
var closeErr coderws.CloseError
|
||||
require.ErrorAs(t, readErr, &closeErr)
|
||||
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
|
||||
require.Contains(t, closeErr.Reason, "not available for this group")
|
||||
status := tc.closeStatus
|
||||
if status == 0 {
|
||||
status = coderws.StatusPolicyViolation
|
||||
}
|
||||
require.Equal(t, status, closeErr.Code)
|
||||
reason := tc.closeReason
|
||||
if reason == "" {
|
||||
reason = "not available for this group"
|
||||
}
|
||||
require.Contains(t, closeErr.Reason, reason)
|
||||
require.Len(t, upstreamPayloadCh, turnCount-1, "rejected turn must not reach upstream")
|
||||
_ = clientConn.CloseNow()
|
||||
return openAIResponsesWSUsageLogResult{}
|
||||
}
|
||||
@@ -3115,6 +3183,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
}
|
||||
}
|
||||
|
||||
if tc.accountPlatform == service.PlatformGrok {
|
||||
upstreamErrCh <- nil
|
||||
}
|
||||
select {
|
||||
case upstreamErr := <-upstreamErrCh:
|
||||
require.NoError(t, upstreamErr)
|
||||
@@ -3199,3 +3270,37 @@ data: {"type":"response.failed","error":{"message":"This content was flagged"}}
|
||||
require.False(t, openAIForwardErrorAlreadyCommunicated(c, c.Writer.Size(), errors.New("openai cyber_policy: blocked")))
|
||||
})
|
||||
}
|
||||
|
||||
// The loader simulates window usage reaching its limit while a connection is
|
||||
// waiting for its first account or between completed WebSocket turns.
|
||||
type simpleModeWSRateLimitRepo struct {
|
||||
service.APIKeyRepository
|
||||
reads atomic.Int64
|
||||
rejectAt int64
|
||||
}
|
||||
|
||||
func (r *simpleModeWSRateLimitRepo) GetRateLimitData(context.Context, int64) (*service.APIKeyRateLimitData, error) {
|
||||
now := time.Now()
|
||||
usage := 0.0
|
||||
if r.reads.Add(1) >= r.rejectAt {
|
||||
usage = 1
|
||||
}
|
||||
return &service.APIKeyRateLimitData{Usage5h: usage, Window5hStart: &now}, nil
|
||||
}
|
||||
func TestOpenAIResponsesWebSocketSimpleModeRechecksKeyWindows(t *testing.T) {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeCtxPool} {
|
||||
t.Run(mode+"/after-account-selection", func(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.1"}`,
|
||||
ingressMode: mode, simpleModeRejectAtRead: 2, firstFrameCloseExpected: true, closeReason: "billing check failed",
|
||||
})
|
||||
})
|
||||
t.Run(mode+"/followup-turn", func(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.1"}`,
|
||||
secondPayload: `{"type":"response.create","model":"gpt-5.1"}`,
|
||||
ingressMode: mode, simpleModeRejectAtRead: 3, secondTurnCloseExpected: true, closeReason: "billing check failed",
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type compositeWSRouteRepo struct {
|
||||
service.CompositeModelRouteRepository
|
||||
routes []service.CompositeModelRoute
|
||||
err error
|
||||
}
|
||||
|
||||
func (r *compositeWSRouteRepo) ListByGroup(context.Context, int64, bool) ([]service.CompositeModelRoute, error) {
|
||||
return r.routes, r.err
|
||||
}
|
||||
|
||||
type compositeWSHTTPUpstream struct{ service.HTTPUpstream }
|
||||
|
||||
func (*compositeWSHTTPUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
return http.DefaultClient.Do(req)
|
||||
}
|
||||
|
||||
func compositeWSResolver(platform, endpoint, upstream string) *service.CompositeRouteResolver {
|
||||
return service.NewCompositeRouteResolver(&compositeWSRouteRepo{routes: []service.CompositeModelRoute{{
|
||||
GroupID: 4201, PublicModel: "public-alias", MatchType: service.CompositeRouteMatchExact,
|
||||
TargetPlatform: platform, Endpoint: endpoint, UpstreamModel: upstream, Enabled: true,
|
||||
}}})
|
||||
}
|
||||
func compositeWSGroup(models ...string) *service.Group {
|
||||
g := wsAllowlistGroup(len(models) > 0, models...)
|
||||
g.Platform = service.PlatformComposite
|
||||
return g
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeAlias(t *testing.T) {
|
||||
for _, platform := range []string{service.PlatformOpenAI, service.PlatformGrok} {
|
||||
for _, endpoint := range []string{service.CompositeRouteEndpointResponses, service.CompositeRouteEndpointAny} {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeDedicated} {
|
||||
t.Run(platform+"/"+endpoint+"/"+mode, func(t *testing.T) {
|
||||
upstream := "gpt-5.4"
|
||||
if platform == service.PlatformGrok {
|
||||
upstream = "grok-4.3"
|
||||
}
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias","input":"hi"}`,
|
||||
secondPayload: `{"type":"response.create","input":"again"}`,
|
||||
group: compositeWSGroup("public-alias"), accountPlatform: platform, ingressMode: mode,
|
||||
compositeResolver: compositeWSResolver(platform, endpoint, upstream),
|
||||
})
|
||||
require.Len(t, got.upstreamPayloads, 2)
|
||||
for i, payload := range got.upstreamPayloads {
|
||||
require.Equal(t, upstream, gjson.GetBytes(payload, "model").String())
|
||||
require.Equal(t, "public-alias", gjson.GetBytes(got.clientEvents[i], "response.model").String())
|
||||
require.Equal(t, "public-alias", got.logs[i].RequestedModel)
|
||||
require.NotNil(t, got.logs[i].UpstreamModel)
|
||||
require.Equal(t, upstream, *got.logs[i].UpstreamModel)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeRouteRejections(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, platform, endpoint, model, reason string
|
||||
repoErr error
|
||||
status coderws.StatusCode
|
||||
}{
|
||||
{name: "disallowed platform", platform: service.PlatformAnthropic, endpoint: "responses", model: "public-alias", reason: "only supports OpenAI-compatible"},
|
||||
{name: "wrong endpoint", platform: service.PlatformOpenAI, endpoint: "messages", model: "public-alias", reason: "only supports OpenAI-compatible"},
|
||||
{name: "unknown alias", platform: service.PlatformOpenAI, endpoint: "responses", model: "unknown-alias", reason: "only supports OpenAI-compatible"},
|
||||
{name: "resolver error", model: "gpt-5.4", repoErr: errors.New("database unavailable"), reason: "Failed to resolve composite model route", status: coderws.StatusInternalError},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
resolver := compositeWSResolver(tc.platform, tc.endpoint, "gpt-5.4")
|
||||
if tc.repoErr != nil {
|
||||
resolver = service.NewCompositeRouteResolver(&compositeWSRouteRepo{err: tc.repoErr})
|
||||
}
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"` + tc.model + `"}`, group: compositeWSGroup(), compositeResolver: resolver,
|
||||
firstFrameCloseExpected: true, closeReason: tc.reason, closeStatus: tc.status,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeAdmissionUsesPublicModel(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias"}`, group: compositeWSGroup("gpt-5.4"),
|
||||
compositeResolver: compositeWSResolver(service.PlatformOpenAI, "responses", "gpt-5.4"), firstFrameCloseExpected: true,
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeModelSwitchRequiresReconnect(t *testing.T) {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeDedicated} {
|
||||
for _, model := range []string{"gpt-5.4", "grok-4.3", "unknown-alias"} {
|
||||
t.Run(mode+"/"+model, func(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias"}`,
|
||||
secondPayload: `{"type":"response.create","model":"` + model + `"}`, group: compositeWSGroup(), ingressMode: mode,
|
||||
compositeResolver: compositeWSResolver(service.PlatformOpenAI, "responses", "gpt-5.4"),
|
||||
secondTurnCloseExpected: true, closeReason: "model switch requires reconnect",
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeChannelBilling(t *testing.T) {
|
||||
for _, source := range []string{service.BillingModelSourceRequested, service.BillingModelSourceChannelMapped, service.BillingModelSourceUpstream} {
|
||||
t.Run(source, func(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.6-sol"}`,
|
||||
secondPayload: `{"type":"response.create","model":"gpt-5.6-sol"}`,
|
||||
group: compositeWSGroup("gpt-5.6-sol"),
|
||||
compositeResolver: service.NewCompositeRouteResolver(&compositeWSRouteRepo{routes: []service.CompositeModelRoute{{PublicModel: "gpt-5.6-sol", MatchType: service.CompositeRouteMatchExact, TargetPlatform: service.PlatformOpenAI, Endpoint: service.CompositeRouteEndpointResponses, UpstreamModel: "route-target"}}}),
|
||||
channelMapping: map[string]string{"route-target": "gpt-5.4"}, billingModelSource: source,
|
||||
accountModelMapping: map[string]any{"gpt-5.4": "gpt-5.4"},
|
||||
})
|
||||
for i, log := range got.logs {
|
||||
require.Equal(t, "gpt-5.4", gjson.GetBytes(got.upstreamPayloads[i], "model").String())
|
||||
require.Equal(t, "gpt-5.6-sol", log.RequestedModel)
|
||||
require.Equal(t, "gpt-5.6-sol", log.Model)
|
||||
if source == service.BillingModelSourceRequested {
|
||||
require.InDelta(t, 40e-6, log.TotalCost, 1e-12)
|
||||
} else {
|
||||
require.InDelta(t, 20e-6, log.TotalCost, 1e-12)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeDetectorFallback(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4"}`, group: compositeWSGroup(),
|
||||
compositeResolver: service.NewCompositeRouteResolver(&compositeWSRouteRepo{}),
|
||||
})
|
||||
require.Equal(t, "gpt-5.4", gjson.GetBytes(got.upstreamFirstPayload, "model").String())
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_CompositeSessionModelSwitchRequiresReconnect(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"public-alias"}`,
|
||||
midPayload: `{"type":"session.update","session":{"model":"grok-4.3"}}`,
|
||||
secondPayload: `{"type":"response.create"}`, group: compositeWSGroup(),
|
||||
compositeResolver: compositeWSResolver(service.PlatformOpenAI, "responses", "gpt-5.4"),
|
||||
secondTurnCloseExpected: true, closeReason: "model switch requires reconnect",
|
||||
})
|
||||
}
|
||||
@@ -143,6 +143,18 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: routingModel, BodyBytes: len(body), Kind: service.InflightEstimateImage, Units: parsed.N})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, body)
|
||||
requestCtx := service.WithOpenAIImagesEndpoint(service.WithOpenAIImageGenerationIntent(c.Request.Context()))
|
||||
|
||||
@@ -165,7 +177,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
sessionHash,
|
||||
routingModel,
|
||||
failedAccountIDs,
|
||||
parsed.RequiredCapability,
|
||||
parsed.RequiredCapabilityForModel(channelMapping.MappedModel),
|
||||
)
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestOpenAIImagesBalanceExhaustionPreservesMachineReadableCause(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, streamStarted := range []bool{false, true} {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", nil)
|
||||
(&OpenAIGatewayHandler{}).handleFailoverExhausted(c, &service.UpstreamFailoverError{
|
||||
StatusCode: http.StatusBadGateway,
|
||||
Reason: service.OpenAIImagesInsufficientBalanceReason,
|
||||
ClientStatusCode: http.StatusPaymentRequired,
|
||||
ClientMessage: service.OpenAIImagesInsufficientBalanceMessage,
|
||||
ResponseHeaders: http.Header{"Retry-After": []string{"23"}, "X-Request-Id": []string{"req-images-balance"}},
|
||||
}, streamStarted)
|
||||
require.Equal(t, "23", recorder.Header().Get("Retry-After"))
|
||||
|
||||
if streamStarted {
|
||||
require.Contains(t, recorder.Body.String(), `"code":"insufficient_balance"`)
|
||||
} else {
|
||||
require.Equal(t, http.StatusPaymentRequired, recorder.Code)
|
||||
require.Equal(t, service.OpenAIImagesInsufficientBalanceCode, gjson.GetBytes(recorder.Body.Bytes(), "error.code").String())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
)
|
||||
|
||||
type openAIWSPassthroughHandlerHarness struct {
|
||||
handler *OpenAIGatewayHandler
|
||||
clientConn *coderws.Conn
|
||||
handlerDone <-chan struct{}
|
||||
moderationRepo *contentModerationHandlerTestRepo
|
||||
@@ -26,7 +27,7 @@ type openAIWSPassthroughHandlerHarness struct {
|
||||
apiKey *service.APIKey
|
||||
}
|
||||
|
||||
func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string) *openAIWSPassthroughHandlerHarness {
|
||||
func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string, settings ...map[string]string) *openAIWSPassthroughHandlerHarness {
|
||||
t.Helper()
|
||||
gatewayCache := testutil.NewRedisGatewayCache(t)
|
||||
|
||||
@@ -35,6 +36,11 @@ func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string) *ope
|
||||
service.SettingKeyCyberSessionBlockEnabled: "true",
|
||||
service.SettingKeyCyberSessionBlockTTLSeconds: "60",
|
||||
}}
|
||||
for _, overrides := range settings {
|
||||
for key, value := range overrides {
|
||||
settingRepo.values[key] = value
|
||||
}
|
||||
}
|
||||
moderationRepo := &contentModerationHandlerTestRepo{}
|
||||
moderationSvc := service.NewContentModerationService(settingRepo, moderationRepo, nil, nil, nil, nil, nil, nil)
|
||||
settingSvc := service.NewSettingService(settingRepo, nil)
|
||||
@@ -90,6 +96,7 @@ func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string) *ope
|
||||
|
||||
apiKey := &service.APIKey{
|
||||
ID: 1851,
|
||||
UserID: 1751,
|
||||
Name: "ws-cyber-key",
|
||||
Key: "sk-handler-cyber-test",
|
||||
GroupID: &groupID,
|
||||
@@ -116,6 +123,7 @@ func newOpenAIWSPassthroughHandlerHarness(t *testing.T, upstreamURL string) *ope
|
||||
t.Cleanup(func() { _ = clientConn.CloseNow() })
|
||||
|
||||
return &openAIWSPassthroughHandlerHarness{
|
||||
handler: h,
|
||||
clientConn: clientConn,
|
||||
handlerDone: handlerDone,
|
||||
moderationRepo: moderationRepo,
|
||||
|
||||
@@ -1146,6 +1146,9 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
if shouldSkipOpsErrorLog(c.Request.Context(), ops, parsed.Message, string(body), c.Request.URL.Path) {
|
||||
return
|
||||
}
|
||||
if shouldSkipOpsClientClosed(c, ops, status) {
|
||||
return
|
||||
}
|
||||
|
||||
apiKey := getOpsAPIKey(c)
|
||||
|
||||
@@ -2491,6 +2494,21 @@ func shouldSkipOpsErrorLog(ctx context.Context, ops *service.OpsService, message
|
||||
return false
|
||||
}
|
||||
|
||||
// shouldSkipOpsClientClosed 按 IgnoreContextCanceled 过滤纯客户端取消的 499。
|
||||
// 499 表示客户端在响应提交前断开(见 failoverClientGone),通常不带
|
||||
// "context canceled" 文案,shouldSkipOpsErrorLog 的文本过滤命中不了。
|
||||
// 本次请求未观察到上游错误时是纯客户端取消;已有上游错误的 499 表示上游失败后
|
||||
// 客户端没等到换号结果就离开,仍按上游失败落库。
|
||||
func shouldSkipOpsClientClosed(c *gin.Context, ops *service.OpsService, status int) bool {
|
||||
if status != statusClientClosedRequest || ops == nil {
|
||||
return false
|
||||
}
|
||||
if !ops.OpsAdvancedSettingsSnapshot().IgnoreContextCanceled {
|
||||
return false
|
||||
}
|
||||
return !hasOpsUpstreamErrorContext(c)
|
||||
}
|
||||
|
||||
// shouldSkipOpsErrorLogForCyber:cyber_policy 命中的请求由 recordCyberPolicyIfMarked
|
||||
// 统一落一条 status=403 的错误请求,故中间件跳过自身落库,避免双写。
|
||||
func shouldSkipOpsErrorLogForCyber(c *gin.Context) bool {
|
||||
|
||||
@@ -2168,3 +2168,83 @@ func TestNormalizeOpsErrorType_KeepsGeminiInBandSignalTypes(t *testing.T) {
|
||||
require.Equal(t, errType, normalizeOpsErrorType(errType, "PROHIBITED_CONTENT"), errType)
|
||||
}
|
||||
}
|
||||
|
||||
type opsAdvancedSettingsRepoStub struct {
|
||||
service.SettingRepository
|
||||
advanced string
|
||||
}
|
||||
|
||||
func (r *opsAdvancedSettingsRepoStub) GetValue(context.Context, string) (string, error) {
|
||||
return "", service.ErrSettingNotFound
|
||||
}
|
||||
|
||||
func (r *opsAdvancedSettingsRepoStub) GetMultiple(context.Context, []string) (map[string]string, error) {
|
||||
return map[string]string{service.SettingKeyOpsAdvancedSettings: r.advanced}, nil
|
||||
}
|
||||
|
||||
func (r *opsAdvancedSettingsRepoStub) Set(context.Context, string, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// serveClientClosedRequest 以已取消的请求 context 走 failoverClientGone,复现网关标记 499 的收尾路径。
|
||||
func serveClientClosedRequest(t *testing.T, ops *service.OpsService, prepare func(c *gin.Context)) {
|
||||
t.Helper()
|
||||
router := gin.New()
|
||||
router.Use(OpsErrorLoggerMiddleware(ops))
|
||||
router.POST("/v1/messages", func(c *gin.Context) {
|
||||
if prepare != nil {
|
||||
prepare(c)
|
||||
}
|
||||
failoverClientGone(c)
|
||||
})
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/v1/messages", nil).WithContext(ctx))
|
||||
require.Equal(t, statusClientClosedRequest, recorder.Code)
|
||||
require.Zero(t, recorder.Body.Len())
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_SkipsPureClientClosed(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 2)
|
||||
gin.SetMode(gin.TestMode)
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
require.True(t, ops.OpsAdvancedSettingsSnapshot().IgnoreContextCanceled)
|
||||
|
||||
serveClientClosedRequest(t, ops, nil)
|
||||
|
||||
require.Zero(t, OpsErrorLogQueueLength(), "纯客户端取消的 499 按 IgnoreContextCanceled 跳过")
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_RecordsClientClosedAfterUpstreamFailure(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 2)
|
||||
gin.SetMode(gin.TestMode)
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
serveClientClosedRequest(t, ops, func(c *gin.Context) {
|
||||
c.Set(service.OpsUpstreamErrorsKey, []*service.OpsUpstreamErrorEvent{{
|
||||
AccountID: 7, UpstreamStatusCode: 524, Kind: "failover", Message: "upstream timeout",
|
||||
}})
|
||||
})
|
||||
|
||||
require.Equal(t, int64(1), OpsErrorLogQueueLength(), "上游失败后客户端离开仍按上游失败落库")
|
||||
job := <-opsErrorLogQueue
|
||||
require.Equal(t, statusClientClosedRequest, job.entry.StatusCode)
|
||||
require.Equal(t, "upstream", job.entry.ErrorPhase)
|
||||
require.NotNil(t, job.entry.UpstreamStatusCode)
|
||||
require.Equal(t, 524, *job.entry.UpstreamStatusCode)
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_RecordsClientClosedWhenIgnoreContextCanceledDisabled(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 2)
|
||||
gin.SetMode(gin.TestMode)
|
||||
settings := &opsAdvancedSettingsRepoStub{advanced: `{"ignore_context_canceled":false}`}
|
||||
ops := service.NewOpsService(nil, settings, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
require.False(t, ops.OpsAdvancedSettingsSnapshot().IgnoreContextCanceled)
|
||||
|
||||
serveClientClosedRequest(t, ops, nil)
|
||||
|
||||
require.Equal(t, int64(1), OpsErrorLogQueueLength())
|
||||
job := <-opsErrorLogQueue
|
||||
require.Equal(t, statusClientClosedRequest, job.entry.StatusCode)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"mime"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// SeedanceTasks exposes Ark's native asynchronous video task protocol.
|
||||
func (h *OpenAIGatewayHandler) SeedanceTasks(c *gin.Context) {
|
||||
if c.Request.Method == http.MethodPost && c.GetHeader("Content-Type") != "" {
|
||||
mediaType, _, err := mime.ParseMediaType(c.GetHeader("Content-Type"))
|
||||
if err != nil || mediaType != "application/json" {
|
||||
h.errorResponse(c, http.StatusUnsupportedMediaType, "invalid_request_error", "Seedance requires application/json")
|
||||
return
|
||||
}
|
||||
}
|
||||
key, ok := middleware.GetAPIKeyFromContext(c)
|
||||
if !ok || key.Group == nil || (key.Group.Platform != service.PlatformOpenAI && key.Group.Platform != service.PlatformComposite) {
|
||||
h.errorResponse(c, http.StatusForbidden, "permission_error", "Seedance requires an OpenAI or composite group")
|
||||
return
|
||||
}
|
||||
endpoint := service.SeedanceEndpointCreate
|
||||
setActualUpstreamEndpoint(c, EndpointSeedanceTasks)
|
||||
taskID := ""
|
||||
if c.Request.Method != http.MethodPost {
|
||||
taskID = service.SeedanceTaskKey(c.Param("task_id"))
|
||||
endpoint = service.SeedanceEndpointStatus
|
||||
if c.Request.Method == http.MethodDelete {
|
||||
endpoint = service.SeedanceEndpointDelete
|
||||
}
|
||||
}
|
||||
h.handleGrokMedia(c, endpoint, taskID)
|
||||
}
|
||||
|
||||
// Ark reports actual completion tokens. Never infer tokens from duration or use
|
||||
// Grok's per-second video tariff. Repeated polls share the durable task dedup key.
|
||||
func prepareSeedanceCompletionBilling(ctx context.Context, h *OpenAIGatewayHandler, key *service.APIKey, subject middleware.AuthSubject, taskID string, result *service.OpenAIForwardResult) *service.OpenAIForwardResult {
|
||||
if result == nil || result.Usage.OutputTokens <= 0 {
|
||||
return nil
|
||||
}
|
||||
pending, err := h.gatewayService.LoadGrokVideoPendingBilling(ctx, taskID, subject.UserID, key.ID)
|
||||
if err != nil || pending == nil {
|
||||
return nil
|
||||
}
|
||||
claimed, err := h.gatewayService.ClaimGrokVideoBilling(ctx, taskID, subject.UserID, key.ID)
|
||||
if err != nil || !claimed {
|
||||
return nil
|
||||
}
|
||||
merged := *result
|
||||
merged.Model = pending.Model
|
||||
merged.BillingModel = firstNonEmptyString(pending.BillingModel, pending.Model)
|
||||
merged.UpstreamModel = firstNonEmptyString(pending.UpstreamModel, result.UpstreamModel)
|
||||
merged.RequestID = service.StableGrokVideoBillingRequestID(taskID)
|
||||
merged.ResponseID = taskID
|
||||
merged.Duration = service.GrokVideoE2EDuration(pending.CreatedAt, time.Now())
|
||||
return &merged
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSeedanceHandlerLifecycleAndOwnership(t *testing.T) {
|
||||
h, slots, bindings, upstream := newGrokMediaSlotHandler(t, false, false, service.PlatformOpenAI)
|
||||
var owner int64
|
||||
upstream.call = func(req *http.Request, id int64) (*http.Response, error) {
|
||||
body := `{"id":"task-ark","status":"queued"}`
|
||||
if req.Method == http.MethodPost {
|
||||
owner = id
|
||||
} else {
|
||||
require.Equal(t, owner, id)
|
||||
}
|
||||
return &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(body))}, nil
|
||||
}
|
||||
newContext := func(method string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
c, w := grokMediaSlotContext(context.Background(), method == http.MethodPost)
|
||||
key, _ := middleware.GetAPIKeyFromContext(c)
|
||||
key.Group.Platform = service.PlatformOpenAI
|
||||
body := ""
|
||||
if method == http.MethodPost {
|
||||
body = `{"model":"doubao-seedance","content":[{"type":"text","text":"waves"}]}`
|
||||
}
|
||||
c.Request = httptest.NewRequest(method, "/api/v3/contents/generations/tasks", strings.NewReader(body))
|
||||
c.Params = gin.Params{{Key: "task_id", Value: "task-ark"}}
|
||||
return c, w
|
||||
}
|
||||
c, w := newContext(http.MethodPost)
|
||||
h.SeedanceTasks(c)
|
||||
require.Equal(t, 200, w.Code, w.Body.String())
|
||||
require.Positive(t, owner)
|
||||
require.Len(t, bindings.pending, 1)
|
||||
slots.assertReleased(t)
|
||||
for _, method := range []string{http.MethodGet, http.MethodDelete} {
|
||||
c, w = newContext(method)
|
||||
h.SeedanceTasks(c)
|
||||
require.Equal(t, 200, w.Code, w.Body.String())
|
||||
slots.assertReleased(t)
|
||||
}
|
||||
for _, other := range []string{"user", "key", "group", "task", "provider"} {
|
||||
c, w = newContext(http.MethodGet)
|
||||
key, _ := middleware.GetAPIKeyFromContext(c)
|
||||
switch other {
|
||||
case "user":
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 11, Concurrency: 5})
|
||||
case "key":
|
||||
key.ID = 21
|
||||
case "group":
|
||||
group := int64(25)
|
||||
key.GroupID = &group
|
||||
case "task":
|
||||
c.Params = gin.Params{{Key: "task_id", Value: "other"}}
|
||||
case "provider":
|
||||
c.Params = gin.Params{{Key: "request_id", Value: "task-ark"}}
|
||||
}
|
||||
before := upstream.calls
|
||||
if other == "provider" {
|
||||
h.GrokVideoStatus(c)
|
||||
} else {
|
||||
h.SeedanceTasks(c)
|
||||
}
|
||||
require.Equal(t, 404, w.Code, other+": "+w.Body.String())
|
||||
require.Equal(t, before, upstream.calls)
|
||||
slots.assertReleased(t)
|
||||
}
|
||||
c, _ = newContext(http.MethodGet)
|
||||
key, _ := middleware.GetAPIKeyFromContext(c)
|
||||
subject, _ := middleware.GetAuthSubjectFromContext(c)
|
||||
result := &service.OpenAIForwardResult{Usage: service.OpenAIUsage{OutputTokens: 12345}, ResponseID: "seedance:task-ark"}
|
||||
for i := range 20 {
|
||||
billed := prepareSeedanceCompletionBilling(context.Background(), h, key, subject, result.ResponseID, result)
|
||||
if i == 0 {
|
||||
require.NotNil(t, billed)
|
||||
require.Equal(t, "doubao-seedance", billed.BillingModel)
|
||||
require.Equal(t, 12345, billed.Usage.OutputTokens)
|
||||
require.Zero(t, billed.VideoCount)
|
||||
} else {
|
||||
require.Nil(t, billed)
|
||||
}
|
||||
}
|
||||
require.Len(t, bindings.billed, 1)
|
||||
}
|
||||
@@ -49,9 +49,13 @@ func ProvideAdminHandlers(
|
||||
auditLogHandler *admin.AuditLogHandler,
|
||||
upstreamBillingProbe *service.UpstreamBillingProbeService,
|
||||
ollamaCloudUsage *service.OllamaCloudUsageService,
|
||||
opencodeGoUsage *service.OpenCodeGoUsageService,
|
||||
claudeResetCredits *service.ClaudeResetCreditService,
|
||||
) *AdminHandlers {
|
||||
accountHandler.SetUpstreamBillingProbeService(upstreamBillingProbe)
|
||||
accountHandler.SetOllamaCloudUsageService(ollamaCloudUsage)
|
||||
accountHandler.SetOpenCodeGoUsageService(opencodeGoUsage)
|
||||
accountHandler.SetClaudeResetCreditService(claudeResetCredits)
|
||||
return &AdminHandlers{
|
||||
Dashboard: dashboardHandler,
|
||||
User: userHandler,
|
||||
@@ -130,10 +134,12 @@ func ProvideOpenAIGatewayHandler(
|
||||
grokQuotaService *service.GrokQuotaService,
|
||||
cfg *config.Config,
|
||||
coordinator *securityaudit.Coordinator,
|
||||
compositeResolver *service.CompositeRouteResolver,
|
||||
) *OpenAIGatewayHandler {
|
||||
gatewayService.SetPluginManager(pluginManager)
|
||||
h := NewOpenAIGatewayHandler(gatewayService, concurrencyService, billingCacheService, apiKeyService,
|
||||
usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, cfg)
|
||||
h.compositeResolver = compositeResolver
|
||||
h.securityAuditCoordinator = coordinator
|
||||
h.grokMediaEligibilityProber = grokQuotaService
|
||||
return h
|
||||
|
||||
@@ -86,3 +86,98 @@ func TestTransformClaudeToGemini_AttributionDoesNotHideIdentity(t *testing.T) {
|
||||
require.Equal(t, "You are Antigravity. Keep this identity.", got.Parts[0].Text)
|
||||
require.False(t, strings.Contains(got.Parts[0].Text, "cc_version="))
|
||||
}
|
||||
|
||||
func TestNeutralizeClaudeIdentity(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
text string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "agent sdk opener",
|
||||
text: "You are a Claude agent, built on Anthropic's Claude Agent SDK. You are naming a coding session.",
|
||||
want: "You are an AI agent. You are naming a coding session.",
|
||||
},
|
||||
{
|
||||
name: "cli opener",
|
||||
text: "You are Claude Code, Anthropic's official CLI for Claude. Help with engineering tasks.",
|
||||
want: "You are an AI agent. Help with engineering tasks.",
|
||||
},
|
||||
{
|
||||
name: "leading whitespace",
|
||||
text: "\n You are a Claude agent, built on Anthropic's Claude Agent SDK.\nRest.",
|
||||
want: "You are an AI agent.\nRest.",
|
||||
},
|
||||
// User instructions that merely mention the vendor must stay untouched:
|
||||
// they are content, not a client identity declaration.
|
||||
{
|
||||
name: "mention inside instructions",
|
||||
text: "Follow CLAUDE.md. Prefer claude-sonnet-4-6 for long tasks.",
|
||||
want: "Follow CLAUDE.md. Prefer claude-sonnet-4-6 for long tasks.",
|
||||
},
|
||||
{
|
||||
name: "opener not at start",
|
||||
text: "Example: You are a Claude agent, built on Anthropic's Claude Agent SDK.",
|
||||
want: "Example: You are a Claude agent, built on Anthropic's Claude Agent SDK.",
|
||||
},
|
||||
{name: "ordinary text", text: "Keep these instructions.", want: "Keep these instructions."},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
require.Equal(t, tt.want, neutralizeClaudeIdentity(tt.text))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The attribution line and the identity sentence arrive as separate system
|
||||
// blocks in real Claude Code traffic; both have to be handled.
|
||||
func TestAttributionAndIdentityTogether(t *testing.T) {
|
||||
attribution := "x-anthropic-billing-header: cc_version=2.1.278.0a9; cc_entrypoint=claude-vscode;"
|
||||
identity := "You are a Claude agent, built on Anthropic's Claude Agent SDK."
|
||||
require.Equal(t, "", stripClaudeAttribution(attribution))
|
||||
require.Equal(t, "You are an AI agent.", neutralizeClaudeIdentity(identity))
|
||||
}
|
||||
|
||||
// End-to-end through the transformer: the attribution block and the identity
|
||||
// sentence arrive as separate system blocks in real Claude Code traffic. The
|
||||
// same words inside user content must survive untouched, because there they are
|
||||
// conversation content rather than a client identity declaration.
|
||||
func TestTransformClaudeToGemini_IdentitySystemBlock(t *testing.T) {
|
||||
const (
|
||||
attribution = "x-anthropic-billing-header: cc_version=2.1.278.0a9; cc_entrypoint=claude-vscode;"
|
||||
identity = "You are a Claude agent, built on Anthropic's Claude Agent SDK."
|
||||
task = "You are naming a coding session."
|
||||
userText = "Does Claude Code support claude-sonnet-4-6?"
|
||||
)
|
||||
system := []SystemBlock{
|
||||
{Type: "text", Text: attribution},
|
||||
{Type: "text", Text: identity},
|
||||
{Type: "text", Text: task},
|
||||
}
|
||||
systemJSON, err := json.Marshal(system)
|
||||
require.NoError(t, err)
|
||||
userJSON, err := json.Marshal(userText)
|
||||
require.NoError(t, err)
|
||||
input := &ClaudeRequest{
|
||||
Model: "gemini-3.8-flash-high",
|
||||
System: systemJSON,
|
||||
Messages: []ClaudeMessage{{Role: "user", Content: userJSON}},
|
||||
}
|
||||
body, err := TransformClaudeToGeminiWithOptions(input, "test-project", input.Model, TransformOptions{})
|
||||
require.NoError(t, err)
|
||||
var got V1InternalRequest
|
||||
require.NoError(t, json.Unmarshal(body, &got))
|
||||
|
||||
var texts []string
|
||||
for _, part := range got.Request.SystemInstruction.Parts {
|
||||
if part.Text != "\n--- [SYSTEM_PROMPT_END] ---" {
|
||||
texts = append(texts, part.Text)
|
||||
}
|
||||
}
|
||||
require.Equal(t, []string{"You are an AI agent.", task}, texts)
|
||||
for _, text := range texts {
|
||||
require.NotContains(t, text, "Claude")
|
||||
require.NotContains(t, text, "Anthropic")
|
||||
}
|
||||
require.Equal(t, userText, got.Request.Contents[0].Parts[0].Text)
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"fmt"
|
||||
"log"
|
||||
"math/rand"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -297,6 +298,29 @@ func filterOpenCodePrompt(text string) string {
|
||||
// from Antigravity system text. This is prompt metadata, not an HTTP header; it
|
||||
// can trigger RESOURCE_EXHAUSTED on the Google upstream. Keep this scoped to the
|
||||
// Antigravity transformer: native Anthropic OAuth may require the attribution.
|
||||
// claudeIdentityOpeners 是部分 Anthropic 客户端放在系统块开头的 Claude Agent SDK 身份语句。
|
||||
// 仅移除 x-anthropic-billing-header 行还不够:如果系统文本仍以这些语句开头,
|
||||
// Google 上游仍会返回 429 RESOURCE_EXHAUSTED。
|
||||
//
|
||||
// The patterns are anchored at the start of a system block on purpose, so user
|
||||
// instructions that merely mention Claude or Anthropic are left untouched.
|
||||
var claudeIdentityOpeners = []*regexp.Regexp{
|
||||
regexp.MustCompile(`(?i)^[ \t\r\n]*You are a Claude agent, built on Anthropic'?s Claude Agent SDK\.?`),
|
||||
regexp.MustCompile(`(?i)^[ \t\r\n]*You are Claude Code, Anthropic'?s official CLI for Claude\.?`),
|
||||
}
|
||||
|
||||
// neutralizeClaudeIdentity rewrites a leading vendor identity sentence into a
|
||||
// vendor-neutral one. It keeps the rest of the block, including any user
|
||||
// instructions that follow.
|
||||
func neutralizeClaudeIdentity(text string) string {
|
||||
for _, re := range claudeIdentityOpeners {
|
||||
if loc := re.FindStringIndex(text); loc != nil {
|
||||
return "You are an AI agent." + text[loc[1]:]
|
||||
}
|
||||
}
|
||||
return text
|
||||
}
|
||||
|
||||
func stripClaudeAttribution(text string) string {
|
||||
trimmed := strings.TrimLeft(text, " \t\r\n")
|
||||
if !strings.HasPrefix(trimmed, "x-anthropic-billing-header:") {
|
||||
@@ -325,7 +349,7 @@ func buildSystemInstruction(system json.RawMessage, modelName string, opts Trans
|
||||
// 尝试解析为字符串
|
||||
var sysStr string
|
||||
if err := json.Unmarshal(system, &sysStr); err == nil {
|
||||
sysStr = stripClaudeAttribution(sysStr)
|
||||
sysStr = neutralizeClaudeIdentity(stripClaudeAttribution(sysStr))
|
||||
if strings.TrimSpace(sysStr) != "" {
|
||||
if strings.Contains(sysStr, "You are Antigravity") {
|
||||
userHasAntigravityIdentity = true
|
||||
@@ -341,7 +365,7 @@ func buildSystemInstruction(system json.RawMessage, modelName string, opts Trans
|
||||
var sysBlocks []SystemBlock
|
||||
if err := json.Unmarshal(system, &sysBlocks); err == nil {
|
||||
for _, block := range sysBlocks {
|
||||
block.Text = stripClaudeAttribution(block.Text)
|
||||
block.Text = neutralizeClaudeIdentity(stripClaudeAttribution(block.Text))
|
||||
if block.Type == "text" && strings.TrimSpace(block.Text) != "" {
|
||||
if strings.Contains(block.Text, "You are Antigravity") {
|
||||
userHasAntigravityIdentity = true
|
||||
@@ -508,8 +532,8 @@ func buildParts(content json.RawMessage, toolIDToName map[string]string, allowDu
|
||||
}
|
||||
parts = append(parts, part)
|
||||
|
||||
case "image":
|
||||
if block.Source != nil && block.Source.Type == "base64" {
|
||||
case "image", "document":
|
||||
if block.Source != nil && block.Source.Type == "base64" && strings.TrimSpace(block.Source.Data) != "" {
|
||||
parts = append(parts, GeminiPart{
|
||||
InlineData: &GeminiInlineData{
|
||||
MimeType: block.Source.MediaType,
|
||||
|
||||
@@ -646,6 +646,20 @@ func TestGeminiToolConfig_DropsBuiltinsWhenClientFunctionsPresent(t *testing.T)
|
||||
})
|
||||
}
|
||||
|
||||
func TestBuildParts_DocumentBecomesInlineData(t *testing.T) {
|
||||
content := `[
|
||||
{"type":"text","text":"read this"},
|
||||
{"type":"document","source":{"type":"base64","media_type":"application/pdf","data":"JVBERi0="}}
|
||||
]`
|
||||
parts, stripped, err := buildParts(json.RawMessage(content), map[string]string{}, true)
|
||||
require.NoError(t, err)
|
||||
require.False(t, stripped)
|
||||
require.Len(t, parts, 2)
|
||||
require.NotNil(t, parts[1].InlineData)
|
||||
require.Equal(t, "application/pdf", parts[1].InlineData.MimeType)
|
||||
require.Equal(t, "JVBERi0=", parts[1].InlineData.Data)
|
||||
}
|
||||
|
||||
// TestToolConfigAlwaysPresent ensures toolConfig is always emitted, including for
|
||||
// reasoning models without any tools: upstream rejects requests that omit it.
|
||||
func TestToolConfigAlwaysPresent(t *testing.T) {
|
||||
|
||||
@@ -119,14 +119,32 @@ func cleanJSONSchemaRecursive(value any) any {
|
||||
// 0. [NEW] 合并 allOf
|
||||
mergeAllOf(schemaMap)
|
||||
|
||||
// 0.5 [FIX] 处理 Draft 2020-12 的 prefixItems (元组定义)
|
||||
if prefixItemsVal, hasPrefix := schemaMap["prefixItems"]; hasPrefix {
|
||||
if prefixItems, ok := prefixItemsVal.([]any); ok && len(prefixItems) > 0 {
|
||||
itemsVal, hasItems := schemaMap["items"]
|
||||
_, itemsIsBool := itemsVal.(bool)
|
||||
// 如果没有 items,或者 items 为 nil,或者 items 是布尔值(如 Draft 2020-12 中表示元组闭合的 items: false),
|
||||
// 用选出的最佳元组项填充 items,以满足 Gemini 对 array.items 必须为合法 Schema 对象的规范。
|
||||
if !hasItems || itemsVal == nil || itemsIsBool {
|
||||
best := extractBestSchemaFromUnion(prefixItems)
|
||||
if best == nil {
|
||||
best = map[string]any{"type": "string"}
|
||||
}
|
||||
schemaMap["items"] = best
|
||||
}
|
||||
}
|
||||
// 无论 prefixItems 是否为空数组,都彻底删除 prefixItems 关键字以防 Gemini 400
|
||||
delete(schemaMap, "prefixItems")
|
||||
}
|
||||
|
||||
// 1. [CRITICAL] 深度递归处理子项
|
||||
if props, ok := schemaMap["properties"].(map[string]any); ok {
|
||||
for _, v := range props {
|
||||
cleanJSONSchemaRecursive(v)
|
||||
}
|
||||
// Go 中不需要像 Rust 那样显式处理 nullable_keys remove required,
|
||||
// 因为我们在子项处理中会正确设置 type 和 description
|
||||
} else if items, ok := schemaMap["items"]; ok {
|
||||
}
|
||||
if items, ok := schemaMap["items"]; ok {
|
||||
// [FIX] Gemini 期望 "items" 是单个 Schema 对象(列表验证),而不是数组(元组验证)。
|
||||
if itemsArr, ok := items.([]any); ok {
|
||||
// 策略:将元组 [A, B] 视为 A、B 中的最佳匹配项。
|
||||
@@ -141,8 +159,9 @@ func cleanJSONSchemaRecursive(value any) any {
|
||||
} else {
|
||||
cleanJSONSchemaRecursive(items)
|
||||
}
|
||||
} else {
|
||||
// 遍历所有值递归
|
||||
}
|
||||
// 如果既无 properties 也无 items,遍历其他 map/slice 子项递归
|
||||
if !hasKey(schemaMap, "properties") && !hasKey(schemaMap, "items") {
|
||||
for _, v := range schemaMap {
|
||||
if _, isMap := v.(map[string]any); isMap {
|
||||
cleanJSONSchemaRecursive(v)
|
||||
@@ -205,6 +224,18 @@ func cleanJSONSchemaRecursive(value any) any {
|
||||
schemaMap[k] = deepCopy(v)
|
||||
}
|
||||
}
|
||||
// 清除已合并的 union 关键字
|
||||
delete(schemaMap, "anyOf")
|
||||
delete(schemaMap, "oneOf")
|
||||
// 对合并进来的新属性和 items 进行深度递归清洗,确保合并进来的子结构(如 prefixItems)也能被处理
|
||||
if props, ok := schemaMap["properties"].(map[string]any); ok {
|
||||
for _, pv := range props {
|
||||
cleanJSONSchemaRecursive(pv)
|
||||
}
|
||||
}
|
||||
if items, ok := schemaMap["items"]; ok {
|
||||
cleanJSONSchemaRecursive(items)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -214,14 +245,49 @@ func cleanJSONSchemaRecursive(value any) any {
|
||||
hasKey(schemaMap, "properties") ||
|
||||
hasKey(schemaMap, "items") ||
|
||||
hasKey(schemaMap, "enum") ||
|
||||
hasKey(schemaMap, "const") ||
|
||||
hasKey(schemaMap, "anyOf") ||
|
||||
hasKey(schemaMap, "oneOf") ||
|
||||
hasKey(schemaMap, "allOf")
|
||||
|
||||
if looksLikeSchema {
|
||||
// 4. [ROBUST] 约束迁移
|
||||
// 4. [ROBUST] 约束迁移与 const 规范化
|
||||
migrateConstraints(schemaMap)
|
||||
|
||||
// 规范化 const 关键字 (Draft 6+ 转换为 Gemini 兼容的 enum: [const])
|
||||
if constVal, exists := schemaMap["const"]; exists {
|
||||
if _, hasEnum := schemaMap["enum"]; !hasEnum {
|
||||
schemaMap["enum"] = []any{constVal}
|
||||
} else if constant, ok := constVal.(string); ok {
|
||||
if existing, ok := schemaMap["enum"].([]any); ok {
|
||||
// Both constraints apply; preserve their intersection before dropping const.
|
||||
values := []any{}
|
||||
for _, value := range existing {
|
||||
if text, ok := value.(string); ok && text == constant {
|
||||
values = append(values, constant)
|
||||
break
|
||||
}
|
||||
}
|
||||
schemaMap["enum"] = values
|
||||
}
|
||||
}
|
||||
if _, hasType := schemaMap["type"]; !hasType {
|
||||
switch constVal.(type) {
|
||||
case string:
|
||||
schemaMap["type"] = "string"
|
||||
case int, int32, int64:
|
||||
schemaMap["type"] = "integer"
|
||||
case float32, float64:
|
||||
schemaMap["type"] = "number"
|
||||
case bool:
|
||||
schemaMap["type"] = "boolean"
|
||||
default:
|
||||
schemaMap["type"] = "string"
|
||||
}
|
||||
}
|
||||
delete(schemaMap, "const")
|
||||
}
|
||||
|
||||
// 5. [CRITICAL] 白名单过滤
|
||||
allowedFields := map[string]bool{
|
||||
"type": true,
|
||||
@@ -238,43 +304,7 @@ func cleanJSONSchemaRecursive(value any) any {
|
||||
}
|
||||
}
|
||||
|
||||
// 6. [SAFETY] 处理空 Object
|
||||
if t, _ := schemaMap["type"].(string); t == "object" {
|
||||
hasProps := false
|
||||
if props, ok := schemaMap["properties"].(map[string]any); ok && len(props) > 0 {
|
||||
hasProps = true
|
||||
}
|
||||
if !hasProps {
|
||||
schemaMap["properties"] = map[string]any{
|
||||
"reason": map[string]any{
|
||||
"type": "string",
|
||||
"description": "Reason for calling this tool",
|
||||
},
|
||||
}
|
||||
schemaMap["required"] = []any{"reason"}
|
||||
}
|
||||
}
|
||||
|
||||
// 7. [SAFETY] Required 字段对齐
|
||||
if props, ok := schemaMap["properties"].(map[string]any); ok {
|
||||
if req, ok := schemaMap["required"].([]any); ok {
|
||||
var validReq []any
|
||||
for _, r := range req {
|
||||
if rStr, ok := r.(string); ok {
|
||||
if _, exists := props[rStr]; exists {
|
||||
validReq = append(validReq, r)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(validReq) > 0 {
|
||||
schemaMap["required"] = validReq
|
||||
} else {
|
||||
delete(schemaMap, "required")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 8. 处理 type 字段 (Lowercase + Nullable 提取)
|
||||
// 6. 处理 type 字段 (Lowercase + Nullable 提取)
|
||||
isEffectivelyNullable := false
|
||||
if typeVal, exists := schemaMap["type"]; exists {
|
||||
var selectedType string
|
||||
@@ -305,16 +335,86 @@ func cleanJSONSchemaRecursive(value any) any {
|
||||
}
|
||||
schemaMap["type"] = selectedType
|
||||
} else {
|
||||
// 默认 object 如果有 properties (虽然上面白名单过滤可能删了 type 如果它不在... 但 type 必在 allowlist)
|
||||
// 如果没有 type,但有 properties,补一个
|
||||
// 如果没有 type,但有 properties/items/enum,推断并补全
|
||||
if hasKey(schemaMap, "properties") {
|
||||
schemaMap["type"] = "object"
|
||||
} else if hasKey(schemaMap, "items") {
|
||||
schemaMap["type"] = "array"
|
||||
} else if enumArr, ok := schemaMap["enum"].([]any); ok && len(enumArr) > 0 {
|
||||
// [FIX] 针对仅含 enum 的 schema,按其元素类型推断标量类型,避免被误判为 object 并注入 reason 属性
|
||||
switch enumArr[0].(type) {
|
||||
case int, int32, int64:
|
||||
schemaMap["type"] = "integer"
|
||||
case float32, float64:
|
||||
schemaMap["type"] = "number"
|
||||
case bool:
|
||||
schemaMap["type"] = "boolean"
|
||||
default:
|
||||
schemaMap["type"] = "string"
|
||||
}
|
||||
} else if hasKey(schemaMap, "enum") {
|
||||
schemaMap["type"] = "string"
|
||||
} else {
|
||||
// 默认为 string ? or object? Gemini 通常需要明确 type
|
||||
schemaMap["type"] = "object"
|
||||
}
|
||||
}
|
||||
|
||||
// 7. [SAFETY] 处理空 Object
|
||||
if t, _ := schemaMap["type"].(string); t == "object" {
|
||||
hasProps := false
|
||||
if props, ok := schemaMap["properties"].(map[string]any); ok && len(props) > 0 {
|
||||
hasProps = true
|
||||
}
|
||||
if !hasProps {
|
||||
schemaMap["properties"] = map[string]any{
|
||||
"reason": map[string]any{
|
||||
"type": "string",
|
||||
"description": "Reason for calling this tool",
|
||||
},
|
||||
}
|
||||
schemaMap["required"] = []any{"reason"}
|
||||
}
|
||||
}
|
||||
|
||||
// 8. [SAFETY] 处理缺失 items 的 Array (Gemini 严苛要求 array 必须声明 items Schema)
|
||||
if t, _ := schemaMap["type"].(string); t == "array" {
|
||||
itemsVal, hasItems := schemaMap["items"]
|
||||
needDefaultItems := false
|
||||
if !hasItems || itemsVal == nil {
|
||||
needDefaultItems = true
|
||||
} else if itemsMap, ok := itemsVal.(map[string]any); ok {
|
||||
if len(itemsMap) == 0 {
|
||||
needDefaultItems = true
|
||||
}
|
||||
} else if _, isBool := itemsVal.(bool); isBool {
|
||||
needDefaultItems = true
|
||||
}
|
||||
if needDefaultItems {
|
||||
schemaMap["items"] = map[string]any{
|
||||
"type": "string",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 9. [SAFETY] Required 字段对齐
|
||||
if props, ok := schemaMap["properties"].(map[string]any); ok {
|
||||
if req, ok := schemaMap["required"].([]any); ok {
|
||||
var validReq []any
|
||||
for _, r := range req {
|
||||
if rStr, ok := r.(string); ok {
|
||||
if _, exists := props[rStr]; exists {
|
||||
validReq = append(validReq, r)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(validReq) > 0 {
|
||||
schemaMap["required"] = validReq
|
||||
} else {
|
||||
delete(schemaMap, "required")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if isEffectivelyNullable {
|
||||
desc, _ := schemaMap["description"].(string)
|
||||
if !strings.Contains(desc, "nullable") {
|
||||
@@ -491,6 +591,9 @@ func scoreSchemaOption(val any) int {
|
||||
if hasKey(m, "items") || typeStr == "array" {
|
||||
return 2
|
||||
}
|
||||
if hasKey(m, "enum") || hasKey(m, "const") {
|
||||
return 1
|
||||
}
|
||||
if typeStr != "" && typeStr != "null" {
|
||||
return 1
|
||||
}
|
||||
|
||||
@@ -0,0 +1,242 @@
|
||||
package antigravity
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCleanJSONSchema_ArrayPrefixItems(t *testing.T) {
|
||||
// 模拟 Claude Code 2.1 Artifact 工具的 query.where 参数 Schema (Draft 2020-12 prefixItems 元组)
|
||||
input := map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"query": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"where": map[string]any{
|
||||
"type": "array",
|
||||
"items": map[string]any{
|
||||
"type": "array",
|
||||
"prefixItems": []any{
|
||||
map[string]any{"type": "string"},
|
||||
map[string]any{"type": "string", "enum": []any{"==", "!=", ">", "<"}},
|
||||
map[string]any{},
|
||||
},
|
||||
},
|
||||
"maxItems": float64(10),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cleaned := CleanJSONSchema(input)
|
||||
require.NotNil(t, cleaned)
|
||||
|
||||
query, ok := cleaned["properties"].(map[string]any)["query"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
where, ok := query["properties"].(map[string]any)["where"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
|
||||
assert.Equal(t, "array", where["type"])
|
||||
|
||||
whereItems, ok := where["items"].(map[string]any)
|
||||
require.True(t, ok, "where.items must be an object")
|
||||
assert.Equal(t, "array", whereItems["type"])
|
||||
// prefixItems 应在 whereItems 中被彻底移除
|
||||
assert.Nil(t, whereItems["prefixItems"])
|
||||
|
||||
// 关键验证:where.items.items 必须存在且有效,不能缺失导致 Gemini 400
|
||||
innerItems, ok := whereItems["items"].(map[string]any)
|
||||
require.True(t, ok, "where.items.items must be an object")
|
||||
assert.Equal(t, "string", innerItems["type"])
|
||||
}
|
||||
|
||||
func TestCleanJSONSchema_ArrayMissingItemsFallback(t *testing.T) {
|
||||
// 针对任何缺少 items 的 array,必须兜底注入 items: {type: string}
|
||||
input := map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"tags": map[string]any{
|
||||
"type": "array",
|
||||
},
|
||||
"nested_empty_array": map[string]any{
|
||||
"type": "array",
|
||||
"items": map[string]any{
|
||||
"type": "array",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cleaned := CleanJSONSchema(input)
|
||||
require.NotNil(t, cleaned)
|
||||
|
||||
props, ok := cleaned["properties"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
tags, ok := props["tags"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
tagsItems, ok := tags["items"].(map[string]any)
|
||||
require.True(t, ok, "tags.items must be an object")
|
||||
assert.Equal(t, "string", tagsItems["type"])
|
||||
|
||||
nested, ok := props["nested_empty_array"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
nestedItems, ok := nested["items"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "array", nestedItems["type"])
|
||||
nestedInnerItems, ok := nestedItems["items"].(map[string]any)
|
||||
require.True(t, ok, "nested items.items must be an object")
|
||||
assert.Equal(t, "string", nestedInnerItems["type"])
|
||||
}
|
||||
|
||||
func TestCleanJSONSchema_ArrayExistingItemsPreserved(t *testing.T) {
|
||||
// 正常的 array items 不受影响
|
||||
input := map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"numbers": map[string]any{
|
||||
"type": "array",
|
||||
"items": map[string]any{
|
||||
"type": "integer",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cleaned := CleanJSONSchema(input)
|
||||
require.NotNil(t, cleaned)
|
||||
|
||||
props, ok := cleaned["properties"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
numbers, ok := props["numbers"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "array", numbers["type"])
|
||||
items, ok := numbers["items"].(map[string]any)
|
||||
require.True(t, ok, "numbers.items must be an object")
|
||||
assert.Equal(t, "integer", items["type"])
|
||||
}
|
||||
|
||||
func TestCleanJSONSchema_EnumOnlySchemaInTuple(t *testing.T) {
|
||||
// 测试 enum-only schema 不会被误判为 object,也不应被注入 reason 属性
|
||||
input := map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"choice": map[string]any{
|
||||
"enum": []any{"option_a", "option_b"},
|
||||
},
|
||||
"tuple_with_enum": map[string]any{
|
||||
"type": "array",
|
||||
"prefixItems": []any{
|
||||
map[string]any{
|
||||
"enum": []any{"read", "write"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cleaned := CleanJSONSchema(input)
|
||||
require.NotNil(t, cleaned)
|
||||
|
||||
props, ok := cleaned["properties"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
choice, ok := props["choice"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "string", choice["type"])
|
||||
assert.Nil(t, choice["properties"], "enum-only schema must NOT be treated as object with reason property")
|
||||
assert.Equal(t, []any{"option_a", "option_b"}, choice["enum"])
|
||||
|
||||
tupleArray, ok := props["tuple_with_enum"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
tupleItems, ok := tupleArray["items"].(map[string]any)
|
||||
require.True(t, ok, "tuple_with_enum.items must be an object")
|
||||
assert.Equal(t, "string", tupleItems["type"])
|
||||
assert.Nil(t, tupleItems["properties"])
|
||||
assert.Equal(t, []any{"read", "write"}, tupleItems["enum"])
|
||||
}
|
||||
|
||||
func TestCleanJSONSchema_ConstKeywordConversion(t *testing.T) {
|
||||
// 测试 const 关键字自动转换为 Gemini 兼容的 enum: [const] 且具有合法 type
|
||||
input := map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"action": map[string]any{
|
||||
"const": "ping",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cleaned := CleanJSONSchema(input)
|
||||
require.NotNil(t, cleaned)
|
||||
|
||||
props, ok := cleaned["properties"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
action, ok := props["action"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Nil(t, action["const"], "const keyword must be removed")
|
||||
assert.Equal(t, "string", action["type"])
|
||||
assert.Equal(t, []any{"ping"}, action["enum"])
|
||||
}
|
||||
|
||||
func TestCleanJSONSchema_AnyOfNestedPrefixItems(t *testing.T) {
|
||||
// 测试 anyOf 分支合并进来的嵌套 prefixItems 也被完整深度清洗
|
||||
input := map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"status": map[string]any{"type": "string"},
|
||||
},
|
||||
"anyOf": []any{
|
||||
map[string]any{
|
||||
"properties": map[string]any{
|
||||
"extra_tuple": map[string]any{
|
||||
"type": "array",
|
||||
"prefixItems": []any{
|
||||
map[string]any{"type": "integer"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cleaned := CleanJSONSchema(input)
|
||||
require.NotNil(t, cleaned)
|
||||
|
||||
props, ok := cleaned["properties"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.NotNil(t, props["extra_tuple"])
|
||||
extraTuple, ok := props["extra_tuple"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Nil(t, extraTuple["prefixItems"], "nested prefixItems from anyOf merge must be cleaned")
|
||||
extraItems, ok := extraTuple["items"].(map[string]any)
|
||||
require.True(t, ok, "extra_tuple.items must be an object")
|
||||
assert.Equal(t, "integer", extraItems["type"])
|
||||
}
|
||||
|
||||
func TestCleanJSONSchema_EmptyPrefixItems(t *testing.T) {
|
||||
// 验证空 prefixItems: [] 也被安全删除,不遗留非法关键字
|
||||
input := map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"empty_tuple": map[string]any{
|
||||
"type": "array",
|
||||
"prefixItems": []any{},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cleaned := CleanJSONSchema(input)
|
||||
require.NotNil(t, cleaned)
|
||||
|
||||
props, ok := cleaned["properties"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
emptyTuple, ok := props["empty_tuple"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Nil(t, emptyTuple["prefixItems"], "empty prefixItems array must be removed")
|
||||
emptyTupleItems, ok := emptyTuple["items"].(map[string]any)
|
||||
require.True(t, ok, "empty_tuple.items must be an object")
|
||||
assert.Equal(t, "string", emptyTupleItems["type"])
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package antigravity
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestBuildToolsPreservesStringConst(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
schema string
|
||||
want string
|
||||
}{
|
||||
{"typed", `{"type":"string","const":"browser"}`, `{"type":"string","enum":["browser"]}`},
|
||||
{"inferred", `{"const":"browser"}`, `{"type":"string","enum":["browser"]}`},
|
||||
{"empty", `{"type":"string","const":""}`, `{"type":"string","enum":[""]}`},
|
||||
{"existing enum", `{"type":"string","const":"browser","enum":["browser","shell"]}`, `{"type":"string","enum":["browser"]}`},
|
||||
{"conflicting enum", `{"type":"string","const":"browser","enum":["shell"]}`, `{"type":"string","enum":[]}`},
|
||||
{"ordinary enum", `{"type":"string","enum":["browser","shell"]}`, `{"type":"string","enum":["browser","shell"]}`},
|
||||
{"array items", `{"type":"array","items":{"type":"string","const":"browser"}}`, `{"type":"array","items":{"type":"string","enum":["browser"]}}`},
|
||||
{"nested property named const", `{"type":"object","properties":{"const":{"const":"browser"}}}`, `{"type":"object","properties":{"const":{"type":"string","enum":["browser"]}}}`},
|
||||
{"schema metadata", `{"type":"string","const":"browser","$schema":"https://json-schema.org/draft/2020-12/schema","description":"Tool kind"}`, `{"type":"string","enum":["browser"],"description":"Tool kind"}`},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var property map[string]any
|
||||
require.NoError(t, json.Unmarshal([]byte(tc.schema), &property))
|
||||
tools := buildTools([]ClaudeTool{{
|
||||
Name: "dispatch",
|
||||
InputSchema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{"action": property},
|
||||
},
|
||||
}})
|
||||
require.Len(t, tools, 1)
|
||||
require.Len(t, tools[0].FunctionDeclarations, 1)
|
||||
properties, ok := tools[0].FunctionDeclarations[0].Parameters["properties"].(map[string]any)
|
||||
require.True(t, ok, "tool parameters must contain a properties object")
|
||||
got, err := json.Marshal(properties["action"])
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, tc.want, string(got))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -40,6 +40,7 @@ type StreamingProcessor struct {
|
||||
outputTokens int
|
||||
cacheReadTokens int
|
||||
imageOutputTokens int
|
||||
hasContent bool
|
||||
}
|
||||
|
||||
// NewStreamingProcessor 创建流式响应处理器
|
||||
@@ -175,6 +176,11 @@ func (p *StreamingProcessor) MessageStartSent() bool {
|
||||
return p.messageStartSent
|
||||
}
|
||||
|
||||
// HasContent reports whether any substantive text, thinking, or tool calls were emitted.
|
||||
func (p *StreamingProcessor) HasContent() bool {
|
||||
return p.hasContent
|
||||
}
|
||||
|
||||
// emitMessageStart 发送 message_start 事件
|
||||
func (p *StreamingProcessor) emitMessageStart(v1Resp *V1InternalResponse) []byte {
|
||||
if p.messageStartSent {
|
||||
@@ -296,6 +302,7 @@ func (p *StreamingProcessor) processThinking(text, signature string) []byte {
|
||||
}
|
||||
|
||||
if text != "" {
|
||||
p.hasContent = true
|
||||
_, _ = result.Write(p.emitDelta("thinking_delta", map[string]any{
|
||||
"thinking": text,
|
||||
}))
|
||||
@@ -350,6 +357,7 @@ func (p *StreamingProcessor) processText(text, signature string) []byte {
|
||||
}))
|
||||
}
|
||||
|
||||
p.hasContent = true
|
||||
_, _ = result.Write(p.emitDelta("text_delta", map[string]any{
|
||||
"text": text,
|
||||
}))
|
||||
@@ -362,6 +370,7 @@ func (p *StreamingProcessor) processFunctionCall(fc *GeminiFunctionCall, signatu
|
||||
var result bytes.Buffer
|
||||
|
||||
p.usedTool = true
|
||||
p.hasContent = true
|
||||
|
||||
toolID := fc.ID
|
||||
if toolID == "" {
|
||||
|
||||
@@ -1136,9 +1136,24 @@ func TestAnthropicToResponses_ThinkingDisabled(t *testing.T) {
|
||||
|
||||
resp, err := AnthropicToResponses(req)
|
||||
require.NoError(t, err)
|
||||
// Default effort applies (medium) even when thinking is disabled.
|
||||
require.NotNil(t, resp.Reasoning)
|
||||
assert.Equal(t, "medium", resp.Reasoning.Effort)
|
||||
assert.Equal(t, "none", resp.Reasoning.Effort)
|
||||
assert.Empty(t, resp.Reasoning.Summary)
|
||||
}
|
||||
|
||||
func TestAnthropicToResponses_ThinkingDisabledOverridesOutputEffort(t *testing.T) {
|
||||
req := &AnthropicRequest{
|
||||
Model: "gpt-5.6-sol",
|
||||
MaxTokens: 1024,
|
||||
Messages: []AnthropicMessage{{Role: "user", Content: json.RawMessage(`"Hello"`)}},
|
||||
Thinking: &AnthropicThinking{Type: "disabled"},
|
||||
OutputConfig: &AnthropicOutputConfig{Effort: "max"},
|
||||
}
|
||||
|
||||
resp, err := AnthropicToResponses(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp.Reasoning)
|
||||
assert.Equal(t, "none", resp.Reasoning.Effort)
|
||||
}
|
||||
|
||||
func TestAnthropicToResponses_NoThinking(t *testing.T) {
|
||||
@@ -1819,6 +1834,130 @@ func TestAnthropicEventToResponses_CacheTokensFromMessageDelta(t *testing.T) {
|
||||
assert.Equal(t, 11, completed.Response.Usage.InputTokensDetails.CachedTokens)
|
||||
}
|
||||
|
||||
func TestOpus55ResponsesAdaptiveThinkingAndToolChoice(t *testing.T) {
|
||||
for _, effort := range []string{"", "low", "medium", "high", "xhigh", "max"} {
|
||||
req := &ResponsesRequest{Model: "claude-opus-5-5", Input: json.RawMessage(`"hello"`), Reasoning: &ResponsesReasoning{Effort: effort}}
|
||||
out, err := ResponsesToAnthropicRequest(req)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "adaptive", out.Thinking.Type)
|
||||
require.Zero(t, out.Thinking.BudgetTokens)
|
||||
if effort == "" {
|
||||
effort = "medium"
|
||||
}
|
||||
require.Equal(t, effort, out.OutputConfig.Effort)
|
||||
}
|
||||
for _, choice := range []string{`"required"`, `{"type":"function","name":"lookup"}`} {
|
||||
_, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-opus-5-5", Input: json.RawMessage(`"hello"`), ToolChoice: json.RawMessage(choice)})
|
||||
require.ErrorContains(t, err, "forced tool_choice")
|
||||
}
|
||||
_, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-opus-5-5", Input: json.RawMessage(`"hello"`), Reasoning: &ResponsesReasoning{Effort: "none"}})
|
||||
require.ErrorContains(t, err, "reasoning effort")
|
||||
old, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-opus-5", Input: json.RawMessage(`"hello"`), Reasoning: &ResponsesReasoning{Effort: "xhigh"}})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "max", old.OutputConfig.Effort)
|
||||
require.Equal(t, "enabled", old.Thinking.Type)
|
||||
}
|
||||
|
||||
func TestOpus55SignedThinkingResponsesRoundTrip(t *testing.T) {
|
||||
block := AnthropicContentBlock{Type: "thinking", Thinking: "", Signature: "upstream-signed-block"}
|
||||
response := AnthropicToResponsesResponse(&AnthropicResponse{Model: "claude-opus-5-5", Content: []AnthropicContentBlock{block, {Type: "tool_use", ID: "toolu_1", Name: "lookup", Input: json.RawMessage(`{}`)}}})
|
||||
require.Len(t, response.Output, 2)
|
||||
require.NotEmpty(t, response.Output[0].EncryptedContent)
|
||||
raw, err := json.Marshal(response.Output)
|
||||
require.NoError(t, err)
|
||||
var items []ResponsesInputItem
|
||||
require.NoError(t, json.Unmarshal(raw, &items))
|
||||
items = append(items, ResponsesInputItem{Type: "function_call_output", CallID: response.Output[1].CallID, Output: "ok"})
|
||||
raw, err = json.Marshal(items)
|
||||
require.NoError(t, err)
|
||||
converted, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-opus-5-5", Input: raw})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, converted.Messages, 2)
|
||||
var blocks []AnthropicContentBlock
|
||||
require.NoError(t, json.Unmarshal(converted.Messages[0].Content, &blocks))
|
||||
require.Equal(t, block, blocks[0])
|
||||
require.Equal(t, "tool_use", blocks[1].Type)
|
||||
// Arbitrary OpenAI ciphertext must never be treated as an Anthropic signature.
|
||||
_, _, err = convertResponsesInputToAnthropic("", json.RawMessage(`[{"type":"reasoning","encrypted_content":"anthropic-thinking-v1:!"}]`), true)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestSonnet55ResponsesThinkingAndSampling(t *testing.T) {
|
||||
for _, effort := range []string{"", "low", "medium", "high", "xhigh", "max", "none"} {
|
||||
req := &ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), Reasoning: &ResponsesReasoning{Effort: effort}}
|
||||
out, err := ResponsesToAnthropicRequest(req)
|
||||
require.NoError(t, err, effort)
|
||||
require.Zero(t, out.Thinking.BudgetTokens)
|
||||
if effort == "none" {
|
||||
require.Equal(t, "between_tools", out.Thinking.Type)
|
||||
require.Equal(t, "low", out.OutputConfig.Effort)
|
||||
} else {
|
||||
require.Equal(t, "adaptive", out.Thinking.Type)
|
||||
if effort == "" {
|
||||
effort = "high"
|
||||
}
|
||||
require.Equal(t, effort, out.OutputConfig.Effort)
|
||||
}
|
||||
}
|
||||
for _, choice := range []string{`"required"`, `{"type":"function","name":"lookup"}`} {
|
||||
_, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), ToolChoice: json.RawMessage(choice)})
|
||||
require.ErrorContains(t, err, "forced tool_choice")
|
||||
}
|
||||
temperature := 0.7
|
||||
_, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), Temperature: &temperature})
|
||||
require.ErrorContains(t, err, "temperature")
|
||||
topP := 0.5
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), TopP: &topP})
|
||||
require.ErrorContains(t, err, "top_p")
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), Reasoning: &ResponsesReasoning{Effort: "minimal"}})
|
||||
require.ErrorContains(t, err, "reasoning effort")
|
||||
|
||||
temperature, topP = 1, 0.99
|
||||
_, err = ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: json.RawMessage(`"hello"`), Temperature: &temperature, TopP: &topP})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestSonnet55SignedThinkingResponsesRoundTrip(t *testing.T) {
|
||||
block := AnthropicContentBlock{Type: "thinking", Thinking: "", Signature: "signed-sonnet-block"}
|
||||
response := AnthropicToResponsesResponse(&AnthropicResponse{Model: "claude-sonnet-5-5", Content: []AnthropicContentBlock{block, {Type: "text", Text: "progress"}, {Type: "tool_use", ID: "toolu_1", Name: "lookup", Input: json.RawMessage(`{}`)}}})
|
||||
require.Len(t, response.Output, 3)
|
||||
require.NotEmpty(t, response.Output[0].EncryptedContent)
|
||||
require.Equal(t, "message", response.Output[1].Type)
|
||||
raw, err := json.Marshal(response.Output)
|
||||
require.NoError(t, err)
|
||||
var items []ResponsesInputItem
|
||||
require.NoError(t, json.Unmarshal(raw, &items))
|
||||
items = append(items, ResponsesInputItem{Type: "function_call_output", CallID: response.Output[2].CallID, Output: "ok"})
|
||||
raw, err = json.Marshal(items)
|
||||
require.NoError(t, err)
|
||||
converted, err := ResponsesToAnthropicRequest(&ResponsesRequest{Model: "claude-sonnet-5-5", Input: raw})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, converted.Messages, 2)
|
||||
var blocks []AnthropicContentBlock
|
||||
require.NoError(t, json.Unmarshal(converted.Messages[0].Content, &blocks))
|
||||
require.Equal(t, block, blocks[0])
|
||||
require.Equal(t, "text", blocks[1].Type)
|
||||
require.Equal(t, "tool_use", blocks[2].Type)
|
||||
}
|
||||
|
||||
func TestGPT6ChatSamplingAndCacheFields(t *testing.T) {
|
||||
temperature := 0.7
|
||||
for _, model := range []string{"gpt-6-sol", "gpt-6-luna"} {
|
||||
for _, effort := range []string{"", "none", "medium", "max"} {
|
||||
out, err := ChatCompletionsToResponses(&ChatCompletionsRequest{Model: model, ReasoningEffort: effort, Temperature: &temperature, TopP: &temperature, PromptCacheOptions: json.RawMessage(`{"mode":"explicit","ttl":"30m"}`), Messages: []ChatMessage{{Role: "user", Content: json.RawMessage(`[{"type":"text","text":"hello","prompt_cache_breakpoint":{"mode":"explicit"}}]`)}}})
|
||||
require.NoError(t, err)
|
||||
if effort == "none" {
|
||||
require.NotNil(t, out.Temperature)
|
||||
} else {
|
||||
require.Nil(t, out.Temperature)
|
||||
require.Nil(t, out.TopP)
|
||||
}
|
||||
require.JSONEq(t, `{"mode":"explicit","ttl":"30m"}`, string(out.PromptCacheOptions))
|
||||
require.Contains(t, string(out.Input), "prompt_cache_breakpoint")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageStartSSE_StopReasonIsJSONNull(t *testing.T) {
|
||||
state := NewResponsesEventToAnthropicState()
|
||||
state.Model = "grok-4.5"
|
||||
@@ -1840,3 +1979,16 @@ func TestMessageStartSSE_StopReasonIsJSONNull(t *testing.T) {
|
||||
require.Contains(t, sse, `"stop_reason":null`)
|
||||
require.NotContains(t, sse, `"stop_reason":""`)
|
||||
}
|
||||
|
||||
func TestGPT61SolCacheOptionsAndBreakpointsSurviveChatBridge(t *testing.T) {
|
||||
sampling := 0.7
|
||||
for _, effort := range []string{"low", "medium", "high", "xhigh", "max"} {
|
||||
out, err := ChatCompletionsToResponses(&ChatCompletionsRequest{Model: "gpt-6.1-sol", ReasoningEffort: effort, Temperature: &sampling, TopP: &sampling, PromptCacheOptions: json.RawMessage(`{"ttl":"30m","mode":"explicit"}`), Messages: []ChatMessage{{Role: "user", Content: json.RawMessage(`[{"type":"text","text":"prefix","prompt_cache_breakpoint":{"mode":"explicit"}}]`)}}})
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, out.Temperature)
|
||||
require.Nil(t, out.TopP)
|
||||
require.Equal(t, effort, out.Reasoning.Effort)
|
||||
require.JSONEq(t, `{"ttl":"30m","mode":"explicit"}`, string(out.PromptCacheOptions))
|
||||
require.Contains(t, string(out.Input), "prompt_cache_breakpoint")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,6 +4,8 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||
)
|
||||
|
||||
// AnthropicToResponses converts an Anthropic Messages request directly into
|
||||
@@ -11,6 +13,9 @@ import (
|
||||
// Chat Completions intermediary round-trip (e.g. thinking, cache_control,
|
||||
// structured system prompts).
|
||||
func AnthropicToResponses(req *AnthropicRequest) (*ResponsesRequest, error) {
|
||||
if err := openai.ValidateGPT61SolReasoningEffort(req.Model, anthropicReasoningEffort(req)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
input, err := convertAnthropicToResponsesInput(req.System, req.Messages)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -55,17 +60,16 @@ func AnthropicToResponses(req *AnthropicRequest) (*ResponsesRequest, error) {
|
||||
out.Tools = convertAnthropicToolsToResponses(req.Tools)
|
||||
}
|
||||
|
||||
// Determine reasoning effort: only output_config.effort controls the
|
||||
// level; thinking.type is ignored. Default follows Codex CLI / airgate's
|
||||
// Anthropic bridge shape, which uses medium when unset.
|
||||
// Anthropic levels map 1:1 to OpenAI: low→low, medium→medium, high→high, max→xhigh.
|
||||
effort := "medium"
|
||||
if req.OutputConfig != nil && req.OutputConfig.Effort != "" {
|
||||
effort = req.OutputConfig.Effort
|
||||
// An explicit thinking disable takes precedence over output_config.effort.
|
||||
effort := anthropicReasoningEffort(req)
|
||||
if openai.IsGPT61SolModelSpelling(req.Model) && req.OutputConfig != nil && req.OutputConfig.Effort == "max" && effort != "none" {
|
||||
effort = "max"
|
||||
}
|
||||
out.Reasoning = &ResponsesReasoning{
|
||||
Effort: mapAnthropicEffortToResponses(effort),
|
||||
Summary: "auto",
|
||||
Effort: effort,
|
||||
}
|
||||
if effort != "none" {
|
||||
out.Reasoning.Summary = "auto"
|
||||
}
|
||||
|
||||
// Convert tool_choice
|
||||
@@ -421,17 +425,22 @@ func extractAnthropicTextFromBlocks(blocks []AnthropicContentBlock) string {
|
||||
return strings.Join(parts, "\n\n")
|
||||
}
|
||||
|
||||
// mapAnthropicEffortToResponses converts Anthropic reasoning effort levels to
|
||||
// OpenAI Responses API effort levels.
|
||||
//
|
||||
// Both APIs default to "high". The mapping is 1:1 for shared levels;
|
||||
// only Anthropic's "max" (Opus 4.6 exclusive) maps to OpenAI's "xhigh"
|
||||
// (GPT-5.2+ exclusive) as both represent the highest reasoning tier.
|
||||
//
|
||||
// low → low
|
||||
// medium → medium
|
||||
// high → high
|
||||
// max → xhigh
|
||||
// anthropicReasoningEffort resolves the Anthropic request preference for both
|
||||
// OpenAI bridges. Explicitly disabled thinking overrides output_config.effort;
|
||||
// otherwise the bridge keeps its medium default.
|
||||
func anthropicReasoningEffort(req *AnthropicRequest) string {
|
||||
if req.Thinking != nil && req.Thinking.Type == "disabled" {
|
||||
return "none"
|
||||
}
|
||||
effort := "medium"
|
||||
if req.OutputConfig != nil && req.OutputConfig.Effort != "" {
|
||||
effort = req.OutputConfig.Effort
|
||||
}
|
||||
return mapAnthropicEffortToResponses(effort)
|
||||
}
|
||||
|
||||
// mapAnthropicEffortToResponses maps shared effort levels directly and maps
|
||||
// Anthropic's max to OpenAI's xhigh.
|
||||
func mapAnthropicEffortToResponses(effort string) string {
|
||||
if effort == "max" {
|
||||
return "xhigh"
|
||||
@@ -467,10 +476,38 @@ func boolPtr(v bool) *bool {
|
||||
|
||||
// isReasoningModel reports whether model is a reasoning model that does not
|
||||
// support sampling parameters (temperature, top_p) via the Responses API.
|
||||
// All gpt-5.x models are reasoning-only; the Responses API returns
|
||||
// "Unsupported parameter: temperature" if these fields are present.
|
||||
// GPT-5 and every later generation are reasoning-only; the Responses API
|
||||
// returns "Unsupported parameter: temperature" if these fields are present.
|
||||
//
|
||||
// Keyed on the generation number instead of a "gpt-5" prefix: pinning the
|
||||
// prefix meant each new family (gpt-6-astra and whatever follows) silently
|
||||
// fell through to the sampling branch and failed upstream on every compat
|
||||
// request until someone edited this line.
|
||||
func isReasoningModel(model string) bool {
|
||||
return strings.HasPrefix(model, "gpt-5")
|
||||
major, ok := openAIModelGeneration(model)
|
||||
return (ok && major >= 5) || openai.IsGPT6SolOrLunaModelSpelling(model)
|
||||
}
|
||||
|
||||
// openAIModelGeneration extracts N from a "gpt-N[.M][-suffix]" model id.
|
||||
// ok is false for non-GPT ids and for GPT families that carry no numeric
|
||||
// generation (gpt-image-1, gpt-audio, ...).
|
||||
func openAIModelGeneration(model string) (int, bool) {
|
||||
rest, ok := strings.CutPrefix(strings.ToLower(strings.TrimSpace(model)), "gpt-")
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
major, digits := 0, 0
|
||||
for _, r := range rest {
|
||||
if r < '0' || r > '9' {
|
||||
break
|
||||
}
|
||||
major = major*10 + int(r-'0')
|
||||
digits++
|
||||
}
|
||||
if digits == 0 {
|
||||
return 0, false
|
||||
}
|
||||
return major, true
|
||||
}
|
||||
|
||||
// normalizeToolParameters ensures the tool parameter schema is valid for
|
||||
|
||||
@@ -2,16 +2,33 @@ package apicompat
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Non-streaming: AnthropicResponse → ResponsesResponse
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const anthropicThinkingEnvelopePrefix = "anthropic-thinking-v1:"
|
||||
|
||||
func encodeAnthropicThinking(block AnthropicContentBlock) string {
|
||||
// Only fields belonging to the signed thinking block are retained.
|
||||
payload, _ := json.Marshal(struct {
|
||||
Type string `json:"type"`
|
||||
Thinking string `json:"thinking"`
|
||||
Signature string `json:"signature,omitempty"`
|
||||
Data string `json:"data,omitempty"`
|
||||
}{block.Type, block.Thinking, block.Signature, block.Data})
|
||||
return anthropicThinkingEnvelopePrefix + base64.RawStdEncoding.EncodeToString(payload)
|
||||
}
|
||||
|
||||
// AnthropicToResponsesResponse converts an Anthropic Messages response into a
|
||||
// Responses API response. This is the reverse of ResponsesToAnthropic and
|
||||
// enables Anthropic upstream responses to be returned in OpenAI Responses format.
|
||||
@@ -35,7 +52,15 @@ func AnthropicToResponsesResponse(resp *AnthropicResponse) *ResponsesResponse {
|
||||
|
||||
for _, block := range resp.Content {
|
||||
switch block.Type {
|
||||
case "thinking":
|
||||
case "thinking", "redacted_thinking":
|
||||
if (claude.IsOpus55(resp.Model) || claude.IsSonnet55(resp.Model)) && (block.Signature != "" || block.Data != "") {
|
||||
item := ResponsesOutput{Type: "reasoning", ID: generateItemID(), EncryptedContent: encodeAnthropicThinking(block)}
|
||||
if block.Thinking != "" {
|
||||
item.Summary = []ResponsesSummary{{Type: "summary_text", Text: block.Thinking}}
|
||||
}
|
||||
outputs = append(outputs, item)
|
||||
continue
|
||||
}
|
||||
if block.Thinking != "" {
|
||||
outputs = append(outputs, ResponsesOutput{
|
||||
Type: "reasoning",
|
||||
@@ -47,6 +72,10 @@ func AnthropicToResponsesResponse(resp *AnthropicResponse) *ResponsesResponse {
|
||||
})
|
||||
}
|
||||
case "text":
|
||||
if (claude.IsOpus55(resp.Model) || claude.IsSonnet55(resp.Model)) && block.Text != "" {
|
||||
outputs = append(outputs, ResponsesOutput{Type: "message", ID: generateItemID(), Role: "assistant", Status: "completed", Content: []ResponsesContentPart{{Type: "output_text", Text: block.Text}}})
|
||||
continue
|
||||
}
|
||||
if block.Text != "" {
|
||||
msgParts = append(msgParts, ResponsesContentPart{
|
||||
Type: "output_text",
|
||||
@@ -164,9 +193,17 @@ type AnthropicEventToResponsesState struct {
|
||||
CurrentName string
|
||||
|
||||
// Content of the currently open item, folded into Outputs when it closes.
|
||||
CurrentContent []ResponsesContentPart // message
|
||||
CurrentArgs string // function_call
|
||||
CurrentSummary string // reasoning
|
||||
CurrentContent []ResponsesContentPart // message
|
||||
CurrentArgs string // function_call
|
||||
CurrentSummary string // reasoning
|
||||
CurrentThinking AnthropicContentBlock
|
||||
PreserveThinkingSignatures bool
|
||||
|
||||
// PendingToolInput holds tool arguments that arrived complete on
|
||||
// content_block_start instead of as input_json_delta. It is only consumed at
|
||||
// content_block_stop, and only when no delta ever arrived, so a canonical
|
||||
// Anthropic stream keeps its exact event sequence.
|
||||
PendingToolInput string
|
||||
|
||||
// Outputs accumulates every closed output item so that response.completed
|
||||
// can carry the full output list. The OpenAI SDK's get_final_response()
|
||||
@@ -248,6 +285,7 @@ func ResponsesEventToSSE(evt ResponsesStreamEvent) (string, error) {
|
||||
func anthToResHandleMessageStart(evt *AnthropicStreamEvent, state *AnthropicEventToResponsesState) []ResponsesStreamEvent {
|
||||
if evt.Message != nil {
|
||||
state.ResponseID = evt.Message.ID
|
||||
state.PreserveThinkingSignatures = state.PreserveThinkingSignatures || claude.IsOpus55(evt.Message.Model) || claude.IsSonnet55(evt.Message.Model)
|
||||
if state.Model == "" {
|
||||
state.Model = evt.Message.Model
|
||||
}
|
||||
@@ -279,7 +317,7 @@ func anthToResHandleContentBlockStart(evt *AnthropicStreamEvent, state *Anthropi
|
||||
var events []ResponsesStreamEvent
|
||||
|
||||
switch evt.ContentBlock.Type {
|
||||
case "thinking":
|
||||
case "thinking", "redacted_thinking":
|
||||
// 开新 item 前必须先关掉在开的那个,与下面的 tool_use 分支一致。
|
||||
// 一个 message item 在它的 text 块 content_block_stop 时是刻意保持打开的
|
||||
// (同一 item 里可能还有后续 text 块),所以 thinking 块到来时它仍然开着:
|
||||
@@ -292,6 +330,8 @@ func anthToResHandleContentBlockStart(evt *AnthropicStreamEvent, state *Anthropi
|
||||
|
||||
state.CurrentItemID = generateItemID()
|
||||
state.CurrentItemType = "reasoning"
|
||||
state.CurrentThinking = *evt.ContentBlock
|
||||
state.CurrentSummary = evt.ContentBlock.Thinking
|
||||
state.ContentIndex = 0
|
||||
|
||||
events = append(events, makeResponsesEvent(state, "response.output_item.added", &ResponsesStreamEvent{
|
||||
@@ -349,6 +389,12 @@ func anthToResHandleContentBlockStart(evt *AnthropicStreamEvent, state *Anthropi
|
||||
state.CurrentItemType = "function_call"
|
||||
state.CurrentCallID = toResponsesCallID(evt.ContentBlock.ID)
|
||||
state.CurrentName = evt.ContentBlock.Name
|
||||
// The canonical Anthropic stream leaves input empty here and streams the
|
||||
// arguments as input_json_delta, but Anthropic-compatible relays may put
|
||||
// the complete arguments on this event and never send a delta. Keep them
|
||||
// as a seed rather than emitting now: a delta, if one follows, is
|
||||
// authoritative and must not be concatenated onto this JSON.
|
||||
state.PendingToolInput = seedToolArguments(evt.ContentBlock.Input)
|
||||
|
||||
events = append(events, makeResponsesEvent(state, "response.output_item.added", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
@@ -399,6 +445,9 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi
|
||||
if evt.Delta.PartialJSON == "" {
|
||||
return nil
|
||||
}
|
||||
// A real delta supersedes whatever content_block_start carried; keeping
|
||||
// both would splice two complete JSON documents together.
|
||||
state.PendingToolInput = ""
|
||||
state.CurrentArgs += evt.Delta.PartialJSON
|
||||
return []ResponsesStreamEvent{makeResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
@@ -409,7 +458,10 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi
|
||||
})}
|
||||
|
||||
case "signature_delta":
|
||||
// Anthropic signature deltas have no Responses equivalent; skip
|
||||
// Keep signatures in the opaque bridge envelope, never in visible text.
|
||||
if state.PreserveThinkingSignatures {
|
||||
state.CurrentThinking.Signature += evt.Delta.Signature
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -431,15 +483,35 @@ func anthToResHandleContentBlockStop(evt *AnthropicStreamEvent, state *Anthropic
|
||||
return events
|
||||
|
||||
case "function_call":
|
||||
// Emit function_call_arguments.done + output item done
|
||||
events := []ResponsesStreamEvent{
|
||||
makeResponsesEvent(state, "response.function_call_arguments.done", &ResponsesStreamEvent{
|
||||
var events []ResponsesStreamEvent
|
||||
// No delta ever arrived, so the arguments the upstream put on
|
||||
// content_block_start are all there is. Emit them as one delta here so
|
||||
// the done event below still repeats exactly what the deltas streamed.
|
||||
if state.CurrentArgs == "" && state.PendingToolInput != "" {
|
||||
state.CurrentArgs = state.PendingToolInput
|
||||
events = append(events, makeResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
Delta: state.PendingToolInput,
|
||||
ItemID: state.CurrentItemID,
|
||||
CallID: state.CurrentCallID,
|
||||
Name: state.CurrentName,
|
||||
}),
|
||||
}))
|
||||
}
|
||||
state.PendingToolInput = ""
|
||||
|
||||
// Emit function_call_arguments.done + output item done.
|
||||
// arguments must repeat exactly what the deltas already streamed for this
|
||||
// item: clients reconcile the done event against the accumulated
|
||||
// function_call_arguments.delta payloads and reject the call as
|
||||
// inconsistent_tool_call when the two disagree. Omitting the field left it
|
||||
// empty while the deltas carried the whole JSON.
|
||||
events = append(events, makeResponsesEvent(state, "response.function_call_arguments.done", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
ItemID: state.CurrentItemID,
|
||||
CallID: state.CurrentCallID,
|
||||
Name: state.CurrentName,
|
||||
Arguments: state.CurrentArgs,
|
||||
}))
|
||||
events = append(events, closeCurrentResponsesItem(state)...)
|
||||
return events
|
||||
|
||||
@@ -520,6 +592,19 @@ func anthropicResponsesStreamTerminalState(stopReason string) (string, *Response
|
||||
return "completed", nil
|
||||
}
|
||||
|
||||
// seedToolArguments normalizes a tool_use content block's inline input into a
|
||||
// seed for the streaming converter. Empty, absent and no-argument payloads
|
||||
// return "" so the existing "{}" fallback still applies and no empty delta is
|
||||
// synthesized.
|
||||
func seedToolArguments(input json.RawMessage) string {
|
||||
trimmed := strings.TrimSpace(string(input))
|
||||
switch trimmed {
|
||||
case "", "{}", "null":
|
||||
return ""
|
||||
}
|
||||
return trimmed
|
||||
}
|
||||
|
||||
func closeCurrentResponsesItem(state *AnthropicEventToResponsesState) []ResponsesStreamEvent {
|
||||
if state.CurrentItemType == "" {
|
||||
return nil
|
||||
@@ -546,6 +631,10 @@ func closeCurrentResponsesItem(state *AnthropicEventToResponsesState) []Response
|
||||
}
|
||||
item.Arguments = args
|
||||
case "reasoning":
|
||||
if state.PreserveThinkingSignatures && (state.CurrentThinking.Signature != "" || state.CurrentThinking.Data != "") {
|
||||
state.CurrentThinking.Thinking = state.CurrentSummary
|
||||
item.EncryptedContent = encodeAnthropicThinking(state.CurrentThinking)
|
||||
}
|
||||
if state.CurrentSummary != "" {
|
||||
item.Summary = []ResponsesSummary{{Type: "summary_text", Text: state.CurrentSummary}}
|
||||
}
|
||||
@@ -559,7 +648,9 @@ func closeCurrentResponsesItem(state *AnthropicEventToResponsesState) []Response
|
||||
state.CurrentName = ""
|
||||
state.CurrentContent = nil
|
||||
state.CurrentArgs = ""
|
||||
state.PendingToolInput = ""
|
||||
state.CurrentSummary = ""
|
||||
state.CurrentThinking = AnthropicContentBlock{}
|
||||
state.TextAccum = ""
|
||||
state.OutputIndex++
|
||||
state.ContentIndex = 0
|
||||
|
||||
@@ -394,3 +394,63 @@ func TestAnthropicEventToResponses_ItemLifecycleIsBalanced(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicEventToResponses_ToolCallArgumentsDoneMatchesDeltas pins that
|
||||
// response.function_call_arguments.done repeats exactly the arguments the deltas
|
||||
// already streamed for that item.
|
||||
//
|
||||
// Why: clients reconcile the done event against the accumulated
|
||||
// function_call_arguments.delta payloads and reject the whole call as
|
||||
// inconsistent_tool_call when the two disagree — Qoder surfaces that as a bare
|
||||
// "something went wrong" and never runs the tool. The field used to be left
|
||||
// unset, so done carried "" while the deltas and the item carried the full JSON.
|
||||
// TestAnthropicEventToResponses_ToolCallCompletedCarriesArguments did not catch
|
||||
// it because it only looked at output_item.done and response.completed.
|
||||
func TestAnthropicEventToResponses_ToolCallArgumentsDoneMatchesDeltas(t *testing.T) {
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
state.Model = "claude-sonnet-4-5"
|
||||
|
||||
var events []ResponsesStreamEvent
|
||||
feed := func(evt *AnthropicStreamEvent) {
|
||||
events = append(events, AnthropicEventToResponsesEvents(evt, state)...)
|
||||
}
|
||||
|
||||
idx := 0
|
||||
feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{
|
||||
Type: "tool_use", ID: "toolu_1", Name: "grep_search",
|
||||
}})
|
||||
// Anthropic streams arguments as JSON fragments; their sum is the contract
|
||||
// clients diff the done event against.
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{
|
||||
Type: "input_json_delta", PartialJSON: `{"path":`,
|
||||
}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{
|
||||
Type: "input_json_delta", PartialJSON: `".","pattern":"TODO"}`,
|
||||
}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx})
|
||||
|
||||
var streamed, done string
|
||||
var sawDone bool
|
||||
for _, e := range events {
|
||||
switch e.Type {
|
||||
case "response.function_call_arguments.delta":
|
||||
streamed += e.Delta
|
||||
case "response.function_call_arguments.done":
|
||||
sawDone = true
|
||||
done = e.Arguments
|
||||
}
|
||||
}
|
||||
|
||||
const want = `{"path":".","pattern":"TODO"}`
|
||||
if streamed != want {
|
||||
t.Fatalf("sum of deltas = %q, want %q", streamed, want)
|
||||
}
|
||||
if !sawDone {
|
||||
t.Fatalf("response.function_call_arguments.done was not emitted; got %d events", len(events))
|
||||
}
|
||||
if done != streamed {
|
||||
t.Errorf("function_call_arguments.done arguments = %q, want %q (must equal the streamed deltas)",
|
||||
done, streamed)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,168 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// collectToolCallStreamEvents drives a tool_use block through the Anthropic →
|
||||
// Responses stream converter and returns every emitted event in order.
|
||||
func collectToolCallStreamEvents(t *testing.T, blockInput json.RawMessage, partials []string) []ResponsesStreamEvent {
|
||||
t.Helper()
|
||||
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
var all []ResponsesStreamEvent
|
||||
|
||||
all = append(all, AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
|
||||
Type: "message_start",
|
||||
Message: &AnthropicResponse{
|
||||
ID: "msg_tool_input",
|
||||
Model: "gemini-3.7-flash",
|
||||
Usage: AnthropicUsage{InputTokens: 12},
|
||||
},
|
||||
}, state)...)
|
||||
|
||||
all = append(all, AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
|
||||
Type: "content_block_start",
|
||||
ContentBlock: &AnthropicContentBlock{
|
||||
Type: "tool_use",
|
||||
ID: "toolu_01abc",
|
||||
Name: "eval",
|
||||
Input: blockInput,
|
||||
},
|
||||
}, state)...)
|
||||
|
||||
for _, partial := range partials {
|
||||
all = append(all, AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
|
||||
Type: "content_block_delta",
|
||||
Delta: &AnthropicDelta{Type: "input_json_delta", PartialJSON: partial},
|
||||
}, state)...)
|
||||
}
|
||||
|
||||
all = append(all, AnthropicEventToResponsesEvents(&AnthropicStreamEvent{Type: "content_block_stop"}, state)...)
|
||||
all = append(all, AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
|
||||
Type: "message_delta",
|
||||
Usage: &AnthropicUsage{OutputTokens: 9},
|
||||
}, state)...)
|
||||
all = append(all, AnthropicEventToResponsesEvents(&AnthropicStreamEvent{Type: "message_stop"}, state)...)
|
||||
|
||||
return all
|
||||
}
|
||||
|
||||
func concatArgumentDeltas(events []ResponsesStreamEvent) string {
|
||||
out := ""
|
||||
for _, e := range events {
|
||||
if e.Type == "response.function_call_arguments.delta" {
|
||||
out += e.Delta
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func findFunctionCallOutput(events []ResponsesStreamEvent) *ResponsesOutput {
|
||||
for i := range events {
|
||||
if events[i].Type != "response.output_item.done" || events[i].Item == nil {
|
||||
continue
|
||||
}
|
||||
if events[i].Item.Type == "function_call" {
|
||||
return events[i].Item
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Upstreams that are not the canonical Anthropic API (here: an Anthropic-compatible
|
||||
// relay fronting Gemini) put the complete tool arguments on content_block_start and
|
||||
// never emit an input_json_delta. The converter must not drop them: this repository
|
||||
// already reads ContentBlock.Input on the non-streaming path
|
||||
// (anthropicResponseToResponsesOutputs) and on the ChatCompletions bridge, so the
|
||||
// streaming path has to agree. Dropping it makes every client-side tool call arrive
|
||||
// with `{}` and the agent loop cannot proceed.
|
||||
func TestAnthropicEventToResponses_ToolInputOnContentBlockStart(t *testing.T) {
|
||||
const args = `{"language":"py","code":"print(1)"}`
|
||||
|
||||
events := collectToolCallStreamEvents(t, json.RawMessage(args), nil)
|
||||
|
||||
item := findFunctionCallOutput(events)
|
||||
require.NotNil(t, item, "a function_call output item must be emitted")
|
||||
assert.Equal(t, args, item.Arguments,
|
||||
"arguments carried on content_block_start must survive to output_item.done")
|
||||
assert.Equal(t, "eval", item.Name)
|
||||
|
||||
// Clients that accumulate from the event stream (rather than reading the
|
||||
// terminal item) must see the arguments too.
|
||||
assert.Equal(t, args, concatArgumentDeltas(events),
|
||||
"argument deltas must reconstruct the full arguments JSON")
|
||||
|
||||
done := findEvent(events, "response.function_call_arguments.done")
|
||||
require.NotNil(t, done, "function_call_arguments.done must be emitted")
|
||||
assert.Equal(t, args, done.Arguments,
|
||||
"function_call_arguments.done must carry the complete arguments")
|
||||
|
||||
completed := findEvent(events, "response.completed")
|
||||
require.NotNil(t, completed)
|
||||
require.NotNil(t, completed.Response)
|
||||
require.Len(t, completed.Response.Output, 1)
|
||||
assert.Equal(t, args, completed.Response.Output[0].Arguments,
|
||||
"response.completed must carry the same arguments")
|
||||
}
|
||||
|
||||
// The canonical Anthropic shape (empty input on content_block_start, arguments
|
||||
// streamed as input_json_delta) must keep its existing event sequence exactly.
|
||||
func TestAnthropicEventToResponses_ToolInputFromDeltasUnchanged(t *testing.T) {
|
||||
const args = `{"language":"py","code":"print(1)"}`
|
||||
|
||||
events := collectToolCallStreamEvents(t, json.RawMessage(`{}`),
|
||||
[]string{`{"language":"py",`, `"code":"print(1)"}`})
|
||||
|
||||
item := findFunctionCallOutput(events)
|
||||
require.NotNil(t, item)
|
||||
assert.Equal(t, args, item.Arguments)
|
||||
|
||||
// Exactly the two upstream deltas, not duplicated by the seed.
|
||||
var deltas []string
|
||||
for _, e := range events {
|
||||
if e.Type == "response.function_call_arguments.delta" {
|
||||
deltas = append(deltas, e.Delta)
|
||||
}
|
||||
}
|
||||
assert.Equal(t, []string{`{"language":"py",`, `"code":"print(1)"}`}, deltas,
|
||||
"canonical delta streaming must not gain or lose events")
|
||||
}
|
||||
|
||||
// An upstream that sends both a populated content_block_start and real deltas
|
||||
// must not have the two concatenated into malformed JSON.
|
||||
func TestAnthropicEventToResponses_ToolInputSeedNotDuplicatedByDeltas(t *testing.T) {
|
||||
const args = `{"language":"py","code":"print(1)"}`
|
||||
|
||||
events := collectToolCallStreamEvents(t, json.RawMessage(args),
|
||||
[]string{`{"language":"py",`, `"code":"print(1)"}`})
|
||||
|
||||
item := findFunctionCallOutput(events)
|
||||
require.NotNil(t, item)
|
||||
assert.Equal(t, args, item.Arguments, "deltas win over the content_block_start seed")
|
||||
assert.True(t, json.Valid([]byte(item.Arguments)), "arguments must stay valid JSON")
|
||||
assert.Equal(t, args, concatArgumentDeltas(events), "seed must not be replayed alongside deltas")
|
||||
}
|
||||
|
||||
// An empty or absent input on content_block_start with no deltas keeps the
|
||||
// existing "{}" fallback rather than emitting a spurious delta.
|
||||
func TestAnthropicEventToResponses_ToolInputEmptyStaysEmptyObject(t *testing.T) {
|
||||
for name, input := range map[string]json.RawMessage{
|
||||
"absent": nil,
|
||||
"empty": json.RawMessage(`{}`),
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
events := collectToolCallStreamEvents(t, input, nil)
|
||||
|
||||
item := findFunctionCallOutput(events)
|
||||
require.NotNil(t, item)
|
||||
assert.Equal(t, "{}", item.Arguments)
|
||||
assert.Empty(t, concatArgumentDeltas(events),
|
||||
"no argument delta should be synthesized when there are no arguments")
|
||||
})
|
||||
}
|
||||
}
|
||||
+131
@@ -0,0 +1,131 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// driveToolUseThroughRestorer runs a tool_use block through the full outbound
|
||||
// chain the gateway actually uses: the Anthropic→Responses stream converter
|
||||
// feeding ResponsesClientToolStreamRestorer. Returns the events a client sees.
|
||||
func driveToolUseThroughRestorer(
|
||||
blockInput json.RawMessage,
|
||||
partials []string,
|
||||
mapping ResponsesClientToolMapping,
|
||||
) []ResponsesStreamEvent {
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
restorer := NewResponsesClientToolStreamRestorer(mapping)
|
||||
|
||||
var clientEvents []ResponsesStreamEvent
|
||||
feed := func(evt *AnthropicStreamEvent) {
|
||||
for _, converted := range AnthropicEventToResponsesEvents(evt, state) {
|
||||
clientEvents = append(clientEvents, restorer.Restore(converted)...)
|
||||
}
|
||||
}
|
||||
|
||||
feed(&AnthropicStreamEvent{
|
||||
Type: "message_start",
|
||||
Message: &AnthropicResponse{
|
||||
ID: "msg_chain",
|
||||
Model: "gemini-3.7-flash",
|
||||
Usage: AnthropicUsage{InputTokens: 11},
|
||||
},
|
||||
})
|
||||
feed(&AnthropicStreamEvent{
|
||||
Type: "content_block_start",
|
||||
ContentBlock: &AnthropicContentBlock{
|
||||
Type: "tool_use",
|
||||
ID: "toolu_chain",
|
||||
Name: "eval",
|
||||
Input: blockInput,
|
||||
},
|
||||
})
|
||||
for _, partial := range partials {
|
||||
feed(&AnthropicStreamEvent{
|
||||
Type: "content_block_delta",
|
||||
Delta: &AnthropicDelta{Type: "input_json_delta", PartialJSON: partial},
|
||||
})
|
||||
}
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_stop"})
|
||||
feed(&AnthropicStreamEvent{Type: "message_delta", Usage: &AnthropicUsage{OutputTokens: 6}})
|
||||
feed(&AnthropicStreamEvent{Type: "message_stop"})
|
||||
|
||||
return clientEvents
|
||||
}
|
||||
|
||||
func firstEventOfType(events []ResponsesStreamEvent, typ string) *ResponsesStreamEvent {
|
||||
for i := range events {
|
||||
if events[i].Type == typ {
|
||||
return &events[i]
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// A client-side custom tool whose arguments arrived only on content_block_start
|
||||
// must still reach the client as a populated custom_tool_call. The restorer
|
||||
// suppresses function_call_arguments.delta for adapted tools and rebuilds the
|
||||
// input from the accumulated arguments, so a converter that drops the inline
|
||||
// input yields an empty tool call even though nothing errors.
|
||||
func TestToolInputOnContentBlockStart_SurvivesClientToolRestore(t *testing.T) {
|
||||
mapping := ResponsesClientToolMapping{CustomTools: map[string]bool{"eval": true}}
|
||||
|
||||
events := driveToolUseThroughRestorer(
|
||||
json.RawMessage(`{"input":"print(1)"}`), nil, mapping)
|
||||
|
||||
inputDone := firstEventOfType(events, "response.custom_tool_call_input.done")
|
||||
require.NotNil(t, inputDone, "custom_tool_call_input.done must be emitted")
|
||||
assert.Equal(t, "print(1)", inputDone.Input,
|
||||
"inline tool input must survive conversion and client-tool restoration")
|
||||
|
||||
var itemDone *ResponsesStreamEvent
|
||||
for i := range events {
|
||||
if events[i].Type == "response.output_item.done" && events[i].Item != nil &&
|
||||
events[i].Item.Type == "custom_tool_call" {
|
||||
itemDone = &events[i]
|
||||
}
|
||||
}
|
||||
require.NotNil(t, itemDone, "a custom_tool_call item must be closed")
|
||||
assert.Equal(t, "print(1)", itemDone.Item.Input)
|
||||
}
|
||||
|
||||
// The same block delivered as a plain (non-adapted) function tool must carry
|
||||
// its arguments through untouched.
|
||||
func TestToolInputOnContentBlockStart_PlainFunctionToolChain(t *testing.T) {
|
||||
const args = `{"language":"py","code":"print(1)"}`
|
||||
|
||||
events := driveToolUseThroughRestorer(
|
||||
json.RawMessage(args), nil, ResponsesClientToolMapping{})
|
||||
|
||||
done := firstEventOfType(events, "response.function_call_arguments.done")
|
||||
require.NotNil(t, done)
|
||||
assert.Equal(t, args, done.Arguments)
|
||||
|
||||
var itemDone *ResponsesStreamEvent
|
||||
for i := range events {
|
||||
if events[i].Type == "response.output_item.done" && events[i].Item != nil &&
|
||||
events[i].Item.Type == "function_call" {
|
||||
itemDone = &events[i]
|
||||
}
|
||||
}
|
||||
require.NotNil(t, itemDone)
|
||||
assert.Equal(t, args, itemDone.Item.Arguments)
|
||||
}
|
||||
|
||||
// Canonical delta streaming through the same chain must keep working, so the
|
||||
// seed path cannot be mistaken for the only source of arguments.
|
||||
func TestToolInputFromDeltas_SurvivesClientToolRestore(t *testing.T) {
|
||||
mapping := ResponsesClientToolMapping{CustomTools: map[string]bool{"eval": true}}
|
||||
|
||||
events := driveToolUseThroughRestorer(
|
||||
json.RawMessage(`{}`),
|
||||
[]string{`{"input":"pri`, `nt(1)"}`},
|
||||
mapping)
|
||||
|
||||
inputDone := firstEventOfType(events, "response.custom_tool_call_input.done")
|
||||
require.NotNil(t, inputDone)
|
||||
assert.Equal(t, "print(1)", inputDone.Input)
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user