mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 20:12:39 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
fd80b08c90 | ||
|
|
dc9216512d | ||
|
|
d0731be2f0 | ||
|
|
a2e1618700 | ||
|
|
31635849e0 | ||
|
|
d03274ad5f | ||
|
|
0c860e81e1 | ||
|
|
daf6cbb5d9 | ||
|
|
912e1c414c | ||
|
|
0a43efb9a4 | ||
|
|
f7f0a56f35 | ||
|
|
ae2735cdb6 | ||
|
|
6624947b62 | ||
|
|
afd069b169 | ||
|
|
6452aeffa1 | ||
|
|
4318a63bd8 | ||
|
|
ca08825935 | ||
|
|
c060aa0031 | ||
|
|
5e244e7382 | ||
|
|
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 | ||
|
|
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 | ||
|
|
1716e99151 | ||
|
|
483cd56924 | ||
|
|
b6c6f1add5 | ||
|
|
3442a6a53e | ||
|
|
9d9e90960e | ||
|
|
fcd07ba0af | ||
|
|
405727805a | ||
|
|
a46493724b | ||
|
|
d8eeb0989b | ||
|
|
a266e50e8e | ||
|
|
5295bd8227 | ||
|
|
7fbc59eb69 | ||
|
|
fa79b1ebc3 | ||
|
|
0c0df031c5 | ||
|
|
85ed17b313 | ||
|
|
b754ff6f1b | ||
|
|
e35123b8ac | ||
|
|
7f1f85cb76 | ||
|
|
935db68517 | ||
|
|
1c0a69c0ce | ||
|
|
27f4398b50 | ||
|
|
e47255715d | ||
|
|
d0ed0eacac | ||
|
|
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 | ||
|
|
1bfe0d3728 | ||
|
|
783a0a9830 | ||
|
|
cd2a4357c2 | ||
|
|
f3a2dcabba | ||
|
|
7bf3d0b585 | ||
|
|
cb2bb6084c | ||
|
|
00591cbe9a | ||
|
|
ba26e543f9 | ||
|
|
f6eeed3d5a | ||
|
|
3b0bb60ea3 | ||
|
|
26c09b7de6 | ||
|
|
4d5a4ab0e5 | ||
|
|
50f79e11fc | ||
|
|
6b09c74e35 | ||
|
|
b252821c59 | ||
|
|
68aa26477f | ||
|
|
7e530c9d4e | ||
|
|
1a9d49e16f | ||
|
|
2d37088bd9 | ||
|
|
4ccdead2a2 | ||
|
|
f7280213f7 | ||
|
|
37238c0985 | ||
|
|
7b282c4f1d | ||
|
|
d054ee9af3 | ||
|
|
9612a70c77 | ||
|
|
24f4736f64 | ||
|
|
9d5c8d961c | ||
|
|
7b4de8b6ae | ||
|
|
4805f069e4 | ||
|
|
cfd2fc779b | ||
|
|
0920e06a42 | ||
|
|
3fdd54ca12 | ||
|
|
cf10d6d01d | ||
|
|
c12d0131cb | ||
|
|
94225b5ed3 | ||
|
|
a9ff663387 | ||
|
|
57b7dbdc85 | ||
|
|
e74c690ac4 | ||
|
|
0952ce3410 | ||
|
|
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.7
|
||||
|
||||
@@ -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,7 +287,8 @@ 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)
|
||||
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)
|
||||
usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig)
|
||||
userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient)
|
||||
userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig)
|
||||
@@ -341,6 +342,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 +352,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 +407,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 +433,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 +581,10 @@ func provideCleanup(
|
||||
codexVersionSync.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"ClaudeCodeVersionSyncService", func() error {
|
||||
claudeCodeVersionSync.Stop()
|
||||
return nil
|
||||
}},
|
||||
{"ProxyExpiryService", func() error {
|
||||
proxyExpiry.Stop()
|
||||
return nil
|
||||
@@ -685,6 +693,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
|
||||
|
||||
@@ -96,6 +96,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 +105,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 控制管理员手动上传的本地进程插件。
|
||||
@@ -1992,6 +2001,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")
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -66,6 +66,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 +78,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 +679,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 +955,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 +1628,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,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,
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
@@ -303,6 +303,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 {
|
||||
|
||||
@@ -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"`
|
||||
@@ -1455,6 +1457,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) {
|
||||
@@ -1768,6 +1779,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 +1834,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
|
||||
}(),
|
||||
@@ -2310,6 +2336,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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -1375,7 +1375,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
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
@@ -1315,3 +1316,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-luna", "gpt-6-sol"}, false, []string{"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-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))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2526,6 +2526,20 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 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 +2797,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,
|
||||
@@ -2893,7 +2912,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)
|
||||
@@ -3340,6 +3359,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 +3379,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)
|
||||
|
||||
@@ -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,9 @@ func newOpenAIWSHandlerTestServer(t *testing.T, h *OpenAIGatewayHandler, subject
|
||||
}
|
||||
|
||||
type openAIResponsesWSUsageLogCase struct {
|
||||
firstPayload string
|
||||
simpleModeRejectAtRead int64
|
||||
closeReason string
|
||||
firstPayload string
|
||||
// midPayload 在首个 turn 完成后发送(如 session.update),上游桩会为它
|
||||
// 回一个 response.completed,客户端按普通事件读取。
|
||||
midPayload string
|
||||
@@ -2963,7 +2979,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,
|
||||
@@ -2998,6 +3020,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
},
|
||||
}
|
||||
h := &OpenAIGatewayHandler{
|
||||
cfg: cfg,
|
||||
gatewayService: gatewaySvc,
|
||||
billingCacheService: billingCacheSvc,
|
||||
apiKeyService: &service.APIKeyService{},
|
||||
@@ -3009,6 +3032,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
|
||||
}
|
||||
@@ -3051,7 +3077,11 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
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")
|
||||
reason := tc.closeReason
|
||||
if reason == "" {
|
||||
reason = "not available for this group"
|
||||
}
|
||||
require.Contains(t, closeErr.Reason, reason)
|
||||
_ = clientConn.CloseNow()
|
||||
return openAIResponsesWSUsageLogResult{}
|
||||
}
|
||||
@@ -3086,7 +3116,11 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
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")
|
||||
reason := tc.closeReason
|
||||
if reason == "" {
|
||||
reason = "not available for this group"
|
||||
}
|
||||
require.Contains(t, closeErr.Reason, reason)
|
||||
_ = clientConn.CloseNow()
|
||||
return openAIResponsesWSUsageLogResult{}
|
||||
}
|
||||
@@ -3199,3 +3233,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",
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -165,7 +165,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())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -49,9 +49,11 @@ func ProvideAdminHandlers(
|
||||
auditLogHandler *admin.AuditLogHandler,
|
||||
upstreamBillingProbe *service.UpstreamBillingProbeService,
|
||||
ollamaCloudUsage *service.OllamaCloudUsageService,
|
||||
opencodeGoUsage *service.OpenCodeGoUsageService,
|
||||
) *AdminHandlers {
|
||||
accountHandler.SetUpstreamBillingProbeService(upstreamBillingProbe)
|
||||
accountHandler.SetOllamaCloudUsageService(ollamaCloudUsage)
|
||||
accountHandler.SetOpenCodeGoUsageService(opencodeGoUsage)
|
||||
return &AdminHandlers{
|
||||
Dashboard: dashboardHandler,
|
||||
User: userHandler,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,37 @@ 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}
|
||||
}
|
||||
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 +292,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 +323,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 +579,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"])
|
||||
}
|
||||
@@ -1819,6 +1819,72 @@ 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 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"
|
||||
|
||||
@@ -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
|
||||
@@ -470,7 +472,7 @@ func boolPtr(v bool) *bool {
|
||||
// All gpt-5.x models are reasoning-only; the Responses API returns
|
||||
// "Unsupported parameter: temperature" if these fields are present.
|
||||
func isReasoningModel(model string) bool {
|
||||
return strings.HasPrefix(model, "gpt-5")
|
||||
return strings.HasPrefix(model, "gpt-5") || openai.IsGPT6SolOrLunaModelSpelling(model)
|
||||
}
|
||||
|
||||
// normalizeToolParameters ensures the tool parameter schema is valid for
|
||||
|
||||
@@ -2,16 +2,32 @@ package apicompat
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"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 +51,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) && (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 +71,10 @@ func AnthropicToResponsesResponse(resp *AnthropicResponse) *ResponsesResponse {
|
||||
})
|
||||
}
|
||||
case "text":
|
||||
if claude.IsOpus55(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 +192,11 @@ 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
|
||||
|
||||
// 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 +278,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)
|
||||
if state.Model == "" {
|
||||
state.Model = evt.Message.Model
|
||||
}
|
||||
@@ -279,7 +310,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 +323,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{
|
||||
@@ -409,7 +442,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,13 +467,19 @@ func anthToResHandleContentBlockStop(evt *AnthropicStreamEvent, state *Anthropic
|
||||
return events
|
||||
|
||||
case "function_call":
|
||||
// Emit function_call_arguments.done + output item done
|
||||
// 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 := []ResponsesStreamEvent{
|
||||
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)...)
|
||||
@@ -546,6 +588,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}}
|
||||
}
|
||||
@@ -560,6 +606,7 @@ func closeCurrentResponsesItem(state *AnthropicEventToResponsesState) []Response
|
||||
state.CurrentContent = nil
|
||||
state.CurrentArgs = ""
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,6 +4,8 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||
)
|
||||
|
||||
type chatMessageContent struct {
|
||||
@@ -27,18 +29,19 @@ func ChatCompletionsToResponses(req *ChatCompletionsRequest) (*ResponsesRequest,
|
||||
}
|
||||
|
||||
out := &ResponsesRequest{
|
||||
Model: req.Model,
|
||||
Instructions: req.Instructions,
|
||||
Input: inputJSON,
|
||||
Stream: true, // upstream always streams
|
||||
Include: []string{"reasoning.encrypted_content"},
|
||||
ServiceTier: req.ServiceTier,
|
||||
ParallelToolCalls: req.ParallelToolCalls,
|
||||
Model: req.Model,
|
||||
Instructions: req.Instructions,
|
||||
Input: inputJSON,
|
||||
Stream: true, // upstream always streams
|
||||
Include: []string{"reasoning.encrypted_content"},
|
||||
ServiceTier: req.ServiceTier,
|
||||
PromptCacheOptions: req.PromptCacheOptions,
|
||||
ParallelToolCalls: req.ParallelToolCalls,
|
||||
}
|
||||
|
||||
// Reasoning models (gpt-5.x) do not accept sampling parameters.
|
||||
// See isReasoningModel in anthropic_to_responses.go.
|
||||
if !isReasoningModel(req.Model) {
|
||||
if !isReasoningModel(req.Model) || (openai.IsGPT6SolOrLunaModelSpelling(req.Model) && req.ReasoningEffort == "none") {
|
||||
out.Temperature = req.Temperature
|
||||
out.TopP = req.TopP
|
||||
}
|
||||
@@ -365,26 +368,29 @@ func convertChatContentPartsToResponses(parts []ChatContentPart) []ResponsesCont
|
||||
for _, p := range parts {
|
||||
switch p.Type {
|
||||
case "text":
|
||||
if p.Text != "" {
|
||||
if p.Text != "" || len(p.PromptCacheBreakpoint) > 0 {
|
||||
responseParts = append(responseParts, ResponsesContentPart{
|
||||
Type: "input_text",
|
||||
Text: p.Text,
|
||||
PromptCacheBreakpoint: p.PromptCacheBreakpoint,
|
||||
Type: "input_text",
|
||||
Text: p.Text,
|
||||
})
|
||||
}
|
||||
case "image_url":
|
||||
if p.ImageURL != nil && p.ImageURL.URL != "" && !isEmptyBase64DataURI(p.ImageURL.URL) {
|
||||
responseParts = append(responseParts, ResponsesContentPart{
|
||||
Type: "input_image",
|
||||
ImageURL: p.ImageURL.URL,
|
||||
PromptCacheBreakpoint: p.PromptCacheBreakpoint,
|
||||
Type: "input_image",
|
||||
ImageURL: p.ImageURL.URL,
|
||||
})
|
||||
}
|
||||
case "file":
|
||||
if p.File != nil && (p.File.FileData != "" || p.File.FileID != "") {
|
||||
responseParts = append(responseParts, ResponsesContentPart{
|
||||
Type: "input_file",
|
||||
Filename: p.File.Filename,
|
||||
FileData: p.File.FileData,
|
||||
FileID: p.File.FileID,
|
||||
PromptCacheBreakpoint: p.PromptCacheBreakpoint,
|
||||
Type: "input_file",
|
||||
Filename: p.File.Filename,
|
||||
FileData: p.File.FileData,
|
||||
FileID: p.File.FileID,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -80,7 +80,7 @@ func TestConvertResponsesInputToAnthropic_DeveloperRole(t *testing.T) {
|
||||
{"role":"user","content":"review this code"}
|
||||
]`
|
||||
|
||||
system, messages, err := convertResponsesInputToAnthropic("", json.RawMessage(input))
|
||||
system, messages, err := convertResponsesInputToAnthropic("", json.RawMessage(input), false)
|
||||
require.NoError(t, err)
|
||||
|
||||
var systemText string
|
||||
@@ -97,7 +97,7 @@ func TestConvertResponsesInputToAnthropic_DeveloperRole(t *testing.T) {
|
||||
{"role":"user","content":"hi"}
|
||||
]`
|
||||
|
||||
_, messages, err := convertResponsesInputToAnthropic("", json.RawMessage(input))
|
||||
_, messages, err := convertResponsesInputToAnthropic("", json.RawMessage(input), false)
|
||||
require.NoError(t, err)
|
||||
|
||||
for _, m := range messages {
|
||||
@@ -116,7 +116,7 @@ func TestConvertResponsesInputToAnthropic_DeveloperRole(t *testing.T) {
|
||||
{"role":"user","content":"hello"}
|
||||
]`
|
||||
|
||||
system, _, err := convertResponsesInputToAnthropic("Main instruction.", json.RawMessage(input))
|
||||
system, _, err := convertResponsesInputToAnthropic("Main instruction.", json.RawMessage(input), false)
|
||||
require.NoError(t, err)
|
||||
|
||||
var systemText string
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
)
|
||||
|
||||
// ResponsesToAnthropicRequest converts a Responses API request into an
|
||||
@@ -11,7 +14,7 @@ import (
|
||||
// enables Anthropic platform groups to accept OpenAI Responses API requests
|
||||
// by converting them to the native /v1/messages format before forwarding upstream.
|
||||
func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, error) {
|
||||
system, messages, err := convertResponsesInputToAnthropic(req.Instructions, req.Input)
|
||||
system, messages, err := convertResponsesInputToAnthropic(req.Instructions, req.Input, claude.IsOpus55(req.Model))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -51,6 +54,34 @@ func ResponsesToAnthropicRequest(req *ResponsesRequest) (*AnthropicRequest, erro
|
||||
out.ToolChoice = tc
|
||||
}
|
||||
|
||||
// Opus 5.5 always uses adaptive thinking. Resolve the upstream model before
|
||||
// conversion, because client aliases need not identify a Claude model.
|
||||
if claude.IsOpus55(req.Model) {
|
||||
var choice struct {
|
||||
Type string `json:"type"`
|
||||
}
|
||||
if len(out.ToolChoice) > 0 {
|
||||
if err := json.Unmarshal(out.ToolChoice, &choice); err != nil {
|
||||
return nil, fmt.Errorf("invalid tool_choice: %w", err)
|
||||
}
|
||||
}
|
||||
if choice.Type == "any" || choice.Type == "tool" {
|
||||
return nil, fmt.Errorf("claude-opus-5-5 does not support forced tool_choice; use auto or none")
|
||||
}
|
||||
effort := "medium"
|
||||
if req.Reasoning != nil && req.Reasoning.Effort != "" {
|
||||
effort = req.Reasoning.Effort
|
||||
}
|
||||
switch effort {
|
||||
case "low", "medium", "high", "xhigh", "max":
|
||||
default:
|
||||
return nil, fmt.Errorf("claude-opus-5-5 does not support reasoning effort %q; use low, medium, high, xhigh or max", effort)
|
||||
}
|
||||
out.Thinking = &AnthropicThinking{Type: "adaptive"}
|
||||
out.OutputConfig = &AnthropicOutputConfig{Effort: effort}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// reasoning.effort → output_config.effort + thinking
|
||||
if req.Reasoning != nil && req.Reasoning.Effort != "" {
|
||||
effort := mapResponsesEffortToAnthropic(req.Reasoning.Effort)
|
||||
@@ -100,7 +131,7 @@ func mapResponsesEffortToAnthropic(effort string) string {
|
||||
// convertResponsesInputToAnthropic extracts system prompt and messages from
|
||||
// a Responses API instructions + input array. Returns the system as raw JSON
|
||||
// (for Anthropic's polymorphic system field) and a list of Anthropic messages.
|
||||
func convertResponsesInputToAnthropic(instructions string, inputRaw json.RawMessage) (json.RawMessage, []AnthropicMessage, error) {
|
||||
func convertResponsesInputToAnthropic(instructions string, inputRaw json.RawMessage, preserveThinking bool) (json.RawMessage, []AnthropicMessage, error) {
|
||||
var systemParts []string
|
||||
if strings.TrimSpace(instructions) != "" {
|
||||
systemParts = append(systemParts, strings.TrimSpace(instructions))
|
||||
@@ -165,11 +196,26 @@ func convertResponsesInputToAnthropic(instructions string, inputRaw json.RawMess
|
||||
})
|
||||
|
||||
case item.Type == "reasoning":
|
||||
// Anthropic 无法摄入 OpenAI 的 reasoning:encrypted_content 是不透明的,
|
||||
// 而 thinking 块的重放需要 Anthropic 自己签发的 signature,无法伪造。
|
||||
// Codex 常见形态(只带 summary + encrypted_content)本来就会被丢弃,
|
||||
// 这里让带 content 数组的形态保持同样行为——否则 reasoning_text 块会被
|
||||
// 原样塞进 Anthropic 请求体,上游直接回 400。
|
||||
// Only decode marked Anthropic bridge envelopes, not arbitrary
|
||||
// OpenAI ciphertext. The upstream remains responsible for signature validation.
|
||||
if preserveThinking && strings.HasPrefix(item.EncryptedContent, anthropicThinkingEnvelopePrefix) {
|
||||
raw, err := base64.RawStdEncoding.DecodeString(strings.TrimPrefix(item.EncryptedContent, anthropicThinkingEnvelopePrefix))
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("invalid Anthropic thinking envelope: %w", err)
|
||||
}
|
||||
var block AnthropicContentBlock
|
||||
if err := json.Unmarshal(raw, &block); err != nil {
|
||||
return nil, nil, fmt.Errorf("invalid Anthropic thinking block: %w", err)
|
||||
}
|
||||
if (block.Type != "thinking" || block.Signature == "") && (block.Type != "redacted_thinking" || block.Data == "") {
|
||||
return nil, nil, fmt.Errorf("invalid Anthropic signed thinking block")
|
||||
}
|
||||
content, err := json.Marshal([]AnthropicContentBlock{block})
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
messages = append(messages, AnthropicMessage{Role: "assistant", Content: content})
|
||||
}
|
||||
|
||||
case item.Role == "user":
|
||||
content, err := convertResponsesUserToAnthropicContent(item.Content)
|
||||
|
||||
@@ -58,7 +58,7 @@ func hasToolResult(blocks []AnthropicContentBlock, toolUseID string) bool {
|
||||
|
||||
func convertAnthropic(t *testing.T, input string) []AnthropicMessage {
|
||||
t.Helper()
|
||||
_, messages, err := convertResponsesInputToAnthropic("", json.RawMessage(input))
|
||||
_, messages, err := convertResponsesInputToAnthropic("", json.RawMessage(input), false)
|
||||
require.NoError(t, err)
|
||||
assertAnthropicPairing(t, messages)
|
||||
return messages
|
||||
|
||||
@@ -66,6 +66,7 @@ type AnthropicContentBlock struct {
|
||||
// Signature carries provider encrypted reasoning (e.g. xAI encrypted_content)
|
||||
// so multi-turn Claude clients can round-trip it back on subsequent turns.
|
||||
Signature string `json:"signature,omitempty"`
|
||||
Data string `json:"data,omitempty"` // redacted_thinking
|
||||
|
||||
// type=image
|
||||
Source *AnthropicImageSource `json:"source,omitempty"`
|
||||
@@ -226,6 +227,7 @@ type AnthropicDelta struct {
|
||||
|
||||
// ResponsesRequest is the request body for POST /v1/responses.
|
||||
type ResponsesRequest struct {
|
||||
PromptCacheOptions json.RawMessage `json:"prompt_cache_options,omitempty"`
|
||||
Model string `json:"model"`
|
||||
Instructions string `json:"instructions,omitempty"`
|
||||
Input json.RawMessage `json:"input"` // string or []ResponsesInputItem
|
||||
@@ -309,9 +311,10 @@ func (i *ResponsesInputItem) UnmarshalJSON(data []byte) error {
|
||||
|
||||
// ResponsesContentPart is a typed content part in a Responses message.
|
||||
type ResponsesContentPart struct {
|
||||
Type string `json:"type"` // "input_text" | "output_text" | "input_image" | "input_file"
|
||||
Text string `json:"text,omitempty"`
|
||||
ImageURL string `json:"image_url,omitempty"` // data URI for input_image
|
||||
PromptCacheBreakpoint json.RawMessage `json:"prompt_cache_breakpoint,omitempty"`
|
||||
Type string `json:"type"` // "input_text" | "output_text" | "input_image" | "input_file"
|
||||
Text string `json:"text,omitempty"`
|
||||
ImageURL string `json:"image_url,omitempty"` // data URI for input_image
|
||||
|
||||
// input_file fields.
|
||||
Filename string `json:"filename,omitempty"`
|
||||
@@ -655,6 +658,7 @@ type ResponsesStreamEvent struct {
|
||||
|
||||
// ChatCompletionsRequest is the request body for POST /v1/chat/completions.
|
||||
type ChatCompletionsRequest struct {
|
||||
PromptCacheOptions json.RawMessage `json:"prompt_cache_options,omitempty"`
|
||||
Model string `json:"model"`
|
||||
Messages []ChatMessage `json:"messages"`
|
||||
Instructions string `json:"instructions,omitempty"` // OpenAI Responses API compat
|
||||
@@ -698,10 +702,11 @@ type ChatMessage struct {
|
||||
|
||||
// ChatContentPart is a typed content part in a multi-modal message.
|
||||
type ChatContentPart struct {
|
||||
Type string `json:"type"` // "text" | "image_url" | "file"
|
||||
Text string `json:"text,omitempty"`
|
||||
ImageURL *ChatImageURL `json:"image_url,omitempty"`
|
||||
File *ChatFile `json:"file,omitempty"`
|
||||
PromptCacheBreakpoint json.RawMessage `json:"prompt_cache_breakpoint,omitempty"`
|
||||
Type string `json:"type"` // "text" | "image_url" | "file"
|
||||
Text string `json:"text,omitempty"`
|
||||
ImageURL *ChatImageURL `json:"image_url,omitempty"`
|
||||
File *ChatFile `json:"file,omitempty"`
|
||||
}
|
||||
|
||||
// ChatImageURL contains the URL for an image content part.
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
package claude
|
||||
|
||||
import "sync/atomic"
|
||||
|
||||
// cliVersionResolver 返回运行期生效的 Claude CLI 版本号(管理员面板手动值 →
|
||||
// 后台同步值)。由 SettingService 在装配时注入(wire.go),解析器内部自带
|
||||
// TTL 缓存,热路径不触库。
|
||||
//
|
||||
// 未注入时(所有不经过 wire 的单元测试)EffectiveCLIVersion() 与 CLIVersion()
|
||||
// 完全等价,行为不变。
|
||||
type cliVersionResolverFunc func() string
|
||||
|
||||
var cliVersionResolver atomic.Pointer[cliVersionResolverFunc]
|
||||
|
||||
// SetCLIVersionResolver 注入运行期版本解析器;传 nil 可撤销注入。
|
||||
func SetCLIVersionResolver(resolver func() string) {
|
||||
if resolver == nil {
|
||||
cliVersionResolver.Store(nil)
|
||||
return
|
||||
}
|
||||
r := cliVersionResolverFunc(resolver)
|
||||
cliVersionResolver.Store(&r)
|
||||
}
|
||||
|
||||
// EffectiveCLIVersion 返回当前生效的 Claude CLI 版本号。
|
||||
// resolver 返回值必须通过 IsSupportedCLIVersion 校验(严格三段 semver 且不低于
|
||||
// 内置基线),否则回退 CLIVersion()(环境变量覆盖 → 内置基线)。
|
||||
// 运行期值恒 >= 内置基线,identity_service 的"只升不降"语义不受影响。
|
||||
func EffectiveCLIVersion() string {
|
||||
if r := cliVersionResolver.Load(); r != nil {
|
||||
if v := (*r)(); IsSupportedCLIVersion(v) {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return CLIVersion()
|
||||
}
|
||||
|
||||
// DefaultUserAgent 返回 Claude Code CLI 伪装 User-Agent。
|
||||
// 所有出站路径都必须经本函数取 UA,保证 User-Agent 头与请求体 billing
|
||||
// attribution 块中的 cc_version 始终源自同一个版本号。
|
||||
func DefaultUserAgent() string {
|
||||
return "claude-cli/" + EffectiveCLIVersion() + " (external, cli)"
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package claude
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// withCLIVersionResolver 注入测试解析器并在用例结束时还原,避免用例间污染。
|
||||
func withCLIVersionResolver(t *testing.T, resolver func() string) {
|
||||
t.Helper()
|
||||
SetCLIVersionResolver(resolver)
|
||||
t.Cleanup(func() { SetCLIVersionResolver(nil) })
|
||||
}
|
||||
|
||||
// a. resolver 未注入时,EffectiveCLIVersion() 必须等于 CLIVersion(),
|
||||
// 且 DefaultHeaders()["User-Agent"] 与改造前的 init 固化值完全一致。
|
||||
func TestEffectiveCLIVersionDefaultsToCLIVersion(t *testing.T) {
|
||||
SetCLIVersionResolver(nil)
|
||||
if got := EffectiveCLIVersion(); got != CLIVersion() {
|
||||
t.Fatalf("EffectiveCLIVersion() = %q, want CLIVersion() = %q", got, CLIVersion())
|
||||
}
|
||||
wantUA := "claude-cli/" + CLIVersion() + " (external, cli)"
|
||||
if got := DefaultHeaders()["User-Agent"]; got != wantUA {
|
||||
t.Fatalf("DefaultHeaders()[User-Agent] = %q, want %q", got, wantUA)
|
||||
}
|
||||
}
|
||||
|
||||
// b. 注入返回更高版本的 resolver 后,DefaultHeaders()["User-Agent"] 跟随变化。
|
||||
func TestEffectiveCLIVersionFollowsResolver(t *testing.T) {
|
||||
const upgraded = "9.9.9"
|
||||
withCLIVersionResolver(t, func() string { return upgraded })
|
||||
|
||||
if got := EffectiveCLIVersion(); got != upgraded {
|
||||
t.Fatalf("EffectiveCLIVersion() = %q, want %q", got, upgraded)
|
||||
}
|
||||
wantUA := "claude-cli/" + upgraded + " (external, cli)"
|
||||
if got := DefaultHeaders()["User-Agent"]; got != wantUA {
|
||||
t.Fatalf("DefaultHeaders()[User-Agent] = %q, want %q", got, wantUA)
|
||||
}
|
||||
if got := DefaultHeaders()["X-App"]; got != "cli" {
|
||||
t.Fatalf("DefaultHeaders()[X-App] = %q, want cli", got)
|
||||
}
|
||||
}
|
||||
|
||||
// c. resolver 返回非法值(空串、非 semver、低于内置基线)时回退 CLIVersion()。
|
||||
func TestEffectiveCLIVersionFallsBackOnInvalidResolverValues(t *testing.T) {
|
||||
for name, invalid := range map[string]string{
|
||||
"empty": "",
|
||||
"not semver": "abc",
|
||||
"below baseline": "1.0.0",
|
||||
"whitespace": " ",
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
withCLIVersionResolver(t, func() string { return invalid })
|
||||
if got := EffectiveCLIVersion(); got != CLIVersion() {
|
||||
t.Fatalf("resolver %q: EffectiveCLIVersion() = %q, want fallback %q", invalid, got, CLIVersion())
|
||||
}
|
||||
ua := DefaultHeaders()["User-Agent"]
|
||||
if !strings.HasPrefix(ua, "claude-cli/"+CLIVersion()+" ") {
|
||||
t.Fatalf("UA %q does not carry fallback version %q", ua, CLIVersion())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 撤销注入(传 nil)后回到 CLIVersion()。
|
||||
func TestSetCLIVersionResolverNilRestoresFallback(t *testing.T) {
|
||||
withCLIVersionResolver(t, func() string { return "9.9.9" })
|
||||
SetCLIVersionResolver(nil)
|
||||
if got := EffectiveCLIVersion(); got != CLIVersion() {
|
||||
t.Fatalf("after nil resolver: EffectiveCLIVersion() = %q, want %q", got, CLIVersion())
|
||||
}
|
||||
}
|
||||
@@ -60,7 +60,7 @@ func TestResolveCLIVersion(t *testing.T) {
|
||||
// 两者由不同代码路径写入同一个请求,不一致会被上游判为非正版客户端。
|
||||
func TestDefaultHeadersUserAgentMatchesCLIVersion(t *testing.T) {
|
||||
want := "claude-cli/" + CLIVersion() + " (external, cli)"
|
||||
if got := DefaultHeaders["User-Agent"]; got != want {
|
||||
if got := DefaultHeaders()["User-Agent"]; got != want {
|
||||
t.Fatalf("DefaultHeaders[User-Agent] = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -108,21 +108,26 @@ func FullClaudeCodeMimicryBetas() []string {
|
||||
}
|
||||
|
||||
// DefaultHeaders 是 Claude Code 客户端默认请求头。
|
||||
var DefaultHeaders = map[string]string{
|
||||
// Keep these in sync with recent Claude CLI traffic to reduce the chance
|
||||
// that Claude Code-scoped OAuth credentials are rejected as "non-CLI" usage.
|
||||
// 版本参考:对齐 Parrot (src/transform/cc_mimicry.py:49) 的 CLI_USER_AGENT。
|
||||
"User-Agent": "claude-cli/" + CLIVersion() + " (external, cli)",
|
||||
"X-Stainless-Lang": "js",
|
||||
"X-Stainless-Package-Version": "0.94.0",
|
||||
"X-Stainless-OS": "Linux",
|
||||
"X-Stainless-Arch": "arm64",
|
||||
"X-Stainless-Runtime": "node",
|
||||
"X-Stainless-Runtime-Version": "v24.3.0",
|
||||
"X-Stainless-Retry-Count": "0",
|
||||
"X-Stainless-Timeout": "600",
|
||||
"X-App": "cli",
|
||||
"Anthropic-Dangerous-Direct-Browser-Access": "true",
|
||||
// 每次调用现构造:User-Agent 走 DefaultUserAgent()(运行期可变版本号),
|
||||
// 不再在包 init 时固化。同一次请求内应只取一次 UA 字符串并在出站头与
|
||||
// billing 两条路径间复用,避免版本缓存翻转瞬间头/体不一致。
|
||||
func DefaultHeaders() map[string]string {
|
||||
return map[string]string{
|
||||
// Keep these in sync with recent Claude CLI traffic to reduce the chance
|
||||
// that Claude Code-scoped OAuth credentials are rejected as "non-CLI" usage.
|
||||
// 版本参考:对齐 Parrot (src/transform/cc_mimicry.py:49) 的 CLI_USER_AGENT。
|
||||
"User-Agent": DefaultUserAgent(),
|
||||
"X-Stainless-Lang": "js",
|
||||
"X-Stainless-Package-Version": "0.94.0",
|
||||
"X-Stainless-OS": "Linux",
|
||||
"X-Stainless-Arch": "arm64",
|
||||
"X-Stainless-Runtime": "node",
|
||||
"X-Stainless-Runtime-Version": "v24.3.0",
|
||||
"X-Stainless-Retry-Count": "0",
|
||||
"X-Stainless-Timeout": "600",
|
||||
"X-App": "cli",
|
||||
"Anthropic-Dangerous-Direct-Browser-Access": "true",
|
||||
}
|
||||
}
|
||||
|
||||
// Model 表示一个 Claude 模型
|
||||
@@ -171,6 +176,12 @@ var DefaultModels = []Model{
|
||||
DisplayName: "Claude Opus 4.8",
|
||||
CreatedAt: "2026-05-29T00:00:00Z",
|
||||
},
|
||||
{
|
||||
ID: "claude-opus-5-5",
|
||||
Type: "model",
|
||||
DisplayName: "Claude Opus 5.5",
|
||||
CreatedAt: "2026-09-22T00:00:00Z",
|
||||
},
|
||||
{
|
||||
ID: "claude-opus-5",
|
||||
Type: "model",
|
||||
|
||||
@@ -21,7 +21,7 @@ const fable51MinCLIVersion = "2.1.251"
|
||||
func TestCLICurrentVersionMatchesDefaultUserAgent(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ua := DefaultHeaders["User-Agent"]
|
||||
ua := DefaultHeaders()["User-Agent"]
|
||||
const prefix = "claude-cli/"
|
||||
require.True(t, strings.HasPrefix(ua, prefix), "unexpected User-Agent: %q", ua)
|
||||
|
||||
|
||||
@@ -18,3 +18,15 @@ func TestDefaultModelsContainsClaudeFable51(t *testing.T) {
|
||||
}
|
||||
t.Fatal("claude-fable-5-1 missing from DefaultModels")
|
||||
}
|
||||
|
||||
func TestDefaultModelsContainsOpus55(t *testing.T) {
|
||||
for _, model := range DefaultModels {
|
||||
if model.ID == "claude-opus-5-5" {
|
||||
if model.DisplayName != "Claude Opus 5.5" || model.CreatedAt != "2026-09-22T00:00:00Z" {
|
||||
t.Fatalf("unexpected Opus 5.5 descriptor: %+v", model)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatal("claude-opus-5-5 missing")
|
||||
}
|
||||
|
||||
@@ -24,6 +24,7 @@ var effortFamilies = []struct {
|
||||
{family: "claude-opus-4-7", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-opus-4-6", levels: effortLowMediumHighMax},
|
||||
{family: "claude-opus-4-5", levels: effortLowMediumHigh},
|
||||
{family: "claude-opus-5-5", levels: effortLowMediumHighXHighMax},
|
||||
{family: "claude-opus-5", levels: effortLowMediumHighXHighMax},
|
||||
}
|
||||
|
||||
@@ -39,6 +40,11 @@ func EffortLevelsForModel(model string) []string {
|
||||
return nil
|
||||
}
|
||||
|
||||
// IsOpus55 identifies the fixed Opus 5.5 ID after provider/local suffix normalization.
|
||||
func IsOpus55(model string) bool {
|
||||
return normalizeEffortModelID(model) == "claude-opus-5-5"
|
||||
}
|
||||
|
||||
func normalizeEffortModelID(model string) string {
|
||||
id := strings.ToLower(strings.TrimSpace(model))
|
||||
id = strings.TrimPrefix(id, "models/")
|
||||
|
||||
@@ -23,6 +23,8 @@ var DefaultModels = []Model{
|
||||
{ID: "gpt-5.6", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 (Sol)"},
|
||||
{ID: "gpt-5.6-terra", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 Terra"},
|
||||
{ID: "gpt-5.6-luna", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 Luna"},
|
||||
{ID: "gpt-6-sol", Object: "model", Created: 1790035200, OwnedBy: "openai", Type: "model", DisplayName: "GPT-6 Sol"},
|
||||
{ID: "gpt-6-luna", Object: "model", Created: 1790035200, OwnedBy: "openai", Type: "model", DisplayName: "GPT-6 Luna"},
|
||||
{ID: "gpt-6-astra", Object: "model", Created: 1788480000, OwnedBy: "openai", Type: "model", DisplayName: "GPT-6 Astra"},
|
||||
{ID: "gpt-5.5", Object: "model", Created: 1776873600, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.5"},
|
||||
{ID: "gpt-5.4", Object: "model", Created: 1738368000, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.4"},
|
||||
@@ -157,3 +159,21 @@ func CodexBaseInstructionsForModel(model string) string {
|
||||
}
|
||||
return latestCodexInstructions()
|
||||
}
|
||||
|
||||
// IsGPT6SolOrLunaModelSpelling recognizes official IDs and existing local effort/compact suffixes.
|
||||
func IsGPT6SolOrLunaModelSpelling(model string) bool {
|
||||
canonical := CanonicalizeOpenAIModelAliasSpelling(model)
|
||||
for _, base := range []string{"gpt-6-sol", "gpt-6-luna"} {
|
||||
if canonical == base {
|
||||
return true
|
||||
}
|
||||
suffix, ok := strings.CutPrefix(canonical, base+"-")
|
||||
if ok {
|
||||
switch suffix {
|
||||
case "none", "low", "medium", "high", "xhigh", "max", "openai-compact":
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -32,3 +32,13 @@ func TestDefaultModelsIncludeGPTImage25(t *testing.T) {
|
||||
require.Contains(t, DefaultModelIDs(), "gpt-image-2.5-flare")
|
||||
require.Contains(t, DefaultModelIDs(), "gpt-image-2.5-sunburst")
|
||||
}
|
||||
|
||||
func TestGPT6SolLunaModelIdentity(t *testing.T) {
|
||||
for _, model := range []string{"gpt-6-sol", "gpt-6-luna"} {
|
||||
require.Contains(t, DefaultModelIDs(), model)
|
||||
require.True(t, IsGPT6SolOrLunaModelSpelling(model))
|
||||
}
|
||||
require.False(t, IsGPT6SolOrLunaModelSpelling("gpt-6-astra"))
|
||||
require.False(t, IsGPT6SolOrLunaModelSpelling("gpt-6-solitude"))
|
||||
require.False(t, IsGPT6SolOrLunaModelSpelling("gpt-6-luna-preview"))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
package typesafe
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type Question struct {
|
||||
Type string `json:"type"`
|
||||
Instructions string `json:"instructions"`
|
||||
}
|
||||
|
||||
type Request struct {
|
||||
Model string `json:"model"`
|
||||
State string `json:"state"`
|
||||
Questions map[string]Question `json:"questions"`
|
||||
}
|
||||
|
||||
type Result struct {
|
||||
Model string
|
||||
Scores map[string]float64
|
||||
Usage Usage
|
||||
}
|
||||
|
||||
type Usage struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
}
|
||||
|
||||
// Evaluate performs one attempt. The caller owns timeouts, retries and key rotation.
|
||||
func Evaluate(ctx context.Context, client *http.Client, baseURL, key string, input Request) (*Result, int, error) {
|
||||
endpoint, err := url.JoinPath(strings.TrimRight(baseURL, "/"), "/v1/systemone")
|
||||
if err != nil {
|
||||
return nil, 0, errors.New("typesafe invalid endpoint")
|
||||
}
|
||||
body, err := json.Marshal(input)
|
||||
if err != nil {
|
||||
return nil, 0, errors.New("typesafe invalid request")
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, 0, errors.New("typesafe invalid request")
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return nil, 0, ctx.Err()
|
||||
}
|
||||
return nil, 0, errors.New("typesafe transport unavailable")
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
// Do not log provider error bodies: they may echo user input or credentials.
|
||||
return nil, resp.StatusCode, fmt.Errorf("typesafe API status %d", resp.StatusCode)
|
||||
}
|
||||
var out struct {
|
||||
Model string `json:"model"`
|
||||
Usage Usage `json:"usage"`
|
||||
Answers map[string]struct {
|
||||
Type string `json:"type"`
|
||||
Noul *float64 `json:"noul"`
|
||||
} `json:"answers"`
|
||||
}
|
||||
if err := json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&out); err != nil || strings.TrimSpace(out.Model) == "" {
|
||||
return nil, resp.StatusCode, errors.New("typesafe invalid response")
|
||||
}
|
||||
result := &Result{Model: out.Model, Usage: out.Usage, Scores: make(map[string]float64, len(input.Questions))}
|
||||
for id := range input.Questions {
|
||||
answer, ok := out.Answers[id]
|
||||
if !ok || answer.Type != "noul" || answer.Noul == nil || math.IsNaN(*answer.Noul) || math.IsInf(*answer.Noul, 0) || *answer.Noul < 0 || *answer.Noul > 1 {
|
||||
return nil, resp.StatusCode, fmt.Errorf("typesafe invalid answer for %s", id)
|
||||
}
|
||||
result.Scores[id] = *answer.Noul
|
||||
}
|
||||
return result, resp.StatusCode, nil
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
package typesafe
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestEvaluateContract(t *testing.T) {
|
||||
questions := map[string]Question{"sexual": {Type: "noul", Instructions: "Evaluate sexual content"}, "illicit": {Type: "noul", Instructions: "Evaluate illegal assistance"}}
|
||||
for _, tc := range []struct {
|
||||
name, response string
|
||||
valid bool
|
||||
}{
|
||||
{"valid", `{"model":"jev-1.13.0","answers":{"sexual":{"type":"noul","noul":0},"illicit":{"type":"noul","noul":1}}}`, true},
|
||||
{"missing_answer", `{"model":"jev-1.13.0","answers":{"sexual":{"type":"noul","noul":0}}}`, false},
|
||||
{"missing_probability", `{"model":"jev-1.13.0","answers":{"sexual":{"type":"noul"},"illicit":{"type":"noul","noul":0}}}`, false},
|
||||
{"negative", `{"model":"jev-1.13.0","answers":{"sexual":{"type":"noul","noul":-0.1},"illicit":{"type":"noul","noul":0}}}`, false},
|
||||
{"above_one", `{"model":"jev-1.13.0","answers":{"sexual":{"type":"noul","noul":1.1},"illicit":{"type":"noul","noul":0}}}`, false},
|
||||
{"wrong_type", `{"model":"jev-1.13.0","answers":{"sexual":{"type":"score","noul":0},"illicit":{"type":"noul","noul":0}}}`, false},
|
||||
{"missing_model", `{"answers":{"sexual":{"type":"noul","noul":0},"illicit":{"type":"noul","noul":0}}}`, false},
|
||||
{"invalid_json", `{bad`, false},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
require.Equal(t, "/v1/systemone", r.URL.Path)
|
||||
require.Equal(t, "Bearer test-secret", r.Header.Get("Authorization"))
|
||||
var input Request
|
||||
require.NoError(t, json.NewDecoder(r.Body).Decode(&input))
|
||||
require.Equal(t, questions, input.Questions)
|
||||
require.Equal(t, "sample", input.State)
|
||||
_, err := w.Write([]byte(tc.response))
|
||||
require.NoError(t, err)
|
||||
}))
|
||||
defer server.Close()
|
||||
result, status, err := Evaluate(context.Background(), server.Client(), server.URL, "test-secret", Request{Model: "jev-latest", State: "sample", Questions: questions})
|
||||
require.Equal(t, 200, status)
|
||||
if tc.valid {
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "jev-1.13.0", result.Model)
|
||||
require.Equal(t, 1.0, result.Scores["illicit"])
|
||||
} else {
|
||||
require.Error(t, err)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvaluateHTTPAndTimeoutDoNotExposePayload(t *testing.T) {
|
||||
for _, status := range []int{401, 429, 500, 529} {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(status)
|
||||
_, err := w.Write([]byte("private request text and test-secret"))
|
||||
require.NoError(t, err)
|
||||
}))
|
||||
_, code, err := Evaluate(context.Background(), server.Client(), server.URL, "test-secret", Request{})
|
||||
require.Equal(t, status, code)
|
||||
require.Error(t, err)
|
||||
require.NotContains(t, err.Error(), "private")
|
||||
require.NotContains(t, err.Error(), "test-secret")
|
||||
server.Close()
|
||||
}
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
_, _, err := Evaluate(ctx, server.Client(), server.URL, "test-secret", Request{})
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
require.False(t, strings.Contains(err.Error(), "test-secret"))
|
||||
}
|
||||
@@ -84,6 +84,7 @@ func (o ModelMappingOptions) defaultText() string {
|
||||
|
||||
var defaultModels = []Model{
|
||||
// Text
|
||||
{ID: "grok-4.7", Object: "model", Type: "model", OwnedBy: "xai", DisplayName: "Grok 4.7"},
|
||||
{ID: "grok-4.6", Object: "model", Type: "model", OwnedBy: "xai", DisplayName: "Grok 4.6"},
|
||||
{ID: "grok-4.5", Object: "model", Type: "model", OwnedBy: "xai", DisplayName: "Grok 4.5"},
|
||||
{ID: "grok-4.3", Object: "model", Type: "model", OwnedBy: "xai", DisplayName: "Grok 4.3"},
|
||||
@@ -106,6 +107,8 @@ var defaultModels = []Model{
|
||||
var grokTextResponsesModelAliases = map[string]string{
|
||||
"grok": DefaultTextModel,
|
||||
"grok-latest": DefaultTextModel,
|
||||
"grok-4.7": "grok-4.7",
|
||||
"grok-4.7-latest": "grok-4.7",
|
||||
"grok-4.6": "grok-4.6",
|
||||
"grok-4.6-latest": "grok-4.6",
|
||||
"grok-4.5": "grok-4.5",
|
||||
|
||||
@@ -68,9 +68,14 @@ func TestIsGrokImagineModel(t *testing.T) {
|
||||
require.False(t, IsGrokImagineModel("grok-build-0.1"))
|
||||
}
|
||||
|
||||
func TestDefaultModelsIncludesGrok46(t *testing.T) {
|
||||
func TestDefaultModelsIncludesGrok46And47(t *testing.T) {
|
||||
t.Parallel()
|
||||
ids := DefaultModelIDs()
|
||||
require.Contains(t, ids, "grok-4.7")
|
||||
require.Equal(t, "grok-4.7", ResolveGrokTextResponsesModelID("grok-4.7"))
|
||||
require.Equal(t, "grok-4.7", ResolveGrokTextResponsesModelID("grok-4.7-latest"))
|
||||
|
||||
// Keep the existing assertions for the previous current model as well.
|
||||
require.Contains(t, ids, "grok-4.6")
|
||||
require.Equal(t, "grok-4.6", ResolveGrokTextResponsesModelID("grok-4.6"))
|
||||
require.Equal(t, "grok-4.6", ResolveGrokTextResponsesModelID("grok-4.6-latest"))
|
||||
|
||||
@@ -65,6 +65,8 @@ var schedulerNeutralExtraKeyPrefixes = []string{
|
||||
|
||||
var schedulerNeutralExtraKeys = map[string]struct{}{
|
||||
"codex_usage_updated_at": {},
|
||||
"codex_credits_snapshot": {},
|
||||
"codex_referral_snapshot": {},
|
||||
"grok_billing_snapshot": {},
|
||||
"session_window_utilization": {},
|
||||
}
|
||||
@@ -646,7 +648,29 @@ func lockAndMergeAccountProbeExtra(
|
||||
extra -> 'upstream_billing_probe',
|
||||
extra -> 'ollama_cloud_usage_session',
|
||||
extra -> 'ollama_cloud_usage_auto_refresh',
|
||||
extra -> 'ollama_cloud_usage_snapshot'
|
||||
extra -> 'ollama_cloud_usage_snapshot',
|
||||
COALESCE(
|
||||
(
|
||||
-- opencode_go 平台分支:新旧都是 opencode_go Go 订阅(account_mode
|
||||
-- 未设置/非 zen 均为 Go,与 GetOpenCodeAccountMode 默认兼容一致),
|
||||
-- 不校验 base_url。
|
||||
(platform = 'opencode_go' AND $2 = 'opencode_go'
|
||||
AND COALESCE(btrim(credentials ->> 'account_mode') <> 'zen', true)
|
||||
AND COALESCE(btrim($4::jsonb ->> 'account_mode') <> 'zen', true))
|
||||
-- 挂载白名单分支:新旧平台都在白名单内且 base_url 均为官方
|
||||
-- OpenCode Go 基址(与 Ollama 同范式,允许白名单内跨平台)。
|
||||
OR (platform IN (`+opencodeGoUsageMountPlatformsSQL+`)
|
||||
AND $2 IN (`+opencodeGoUsageMountPlatformsSQL+`)
|
||||
AND `+opencodeGoBaseURLMatchSQLPrefix+`credentials ->> 'base_url'`+opencodeGoBaseURLMatchSQLSuffix+`
|
||||
AND `+opencodeGoBaseURLMatchSQLPrefix+`$4::jsonb ->> 'base_url'`+opencodeGoBaseURLMatchSQLSuffix+`)
|
||||
)
|
||||
AND type = 'apikey'
|
||||
AND $3 = 'apikey'
|
||||
AND credentials -> 'api_key' IS NOT DISTINCT FROM $4::jsonb -> 'api_key',
|
||||
false
|
||||
),
|
||||
extra -> 'opencode_go_usage_auto_refresh',
|
||||
extra -> 'opencode_go_usage_snapshot'
|
||||
FROM accounts
|
||||
WHERE id = $1 AND deleted_at IS NULL
|
||||
FOR NO KEY UPDATE
|
||||
@@ -663,15 +687,18 @@ func lockAndMergeAccountProbeExtra(
|
||||
}
|
||||
|
||||
var (
|
||||
identityUnchanged bool
|
||||
ollamaGroupIdentityUnchanged bool
|
||||
ollamaProxyIdentityUnchanged bool
|
||||
currentEnabled []byte
|
||||
currentRateSyncEnabled []byte
|
||||
currentSnapshot []byte
|
||||
currentOllamaSession []byte
|
||||
currentOllamaAutoRefresh []byte
|
||||
currentOllamaSnapshot []byte
|
||||
identityUnchanged bool
|
||||
ollamaGroupIdentityUnchanged bool
|
||||
ollamaProxyIdentityUnchanged bool
|
||||
opencodeGroupIdentityUnchanged bool
|
||||
currentEnabled []byte
|
||||
currentRateSyncEnabled []byte
|
||||
currentSnapshot []byte
|
||||
currentOllamaSession []byte
|
||||
currentOllamaAutoRefresh []byte
|
||||
currentOllamaSnapshot []byte
|
||||
currentOpenCodeAutoRefresh []byte
|
||||
currentOpenCodeSnapshot []byte
|
||||
)
|
||||
if err := rows.Scan(
|
||||
&identityUnchanged,
|
||||
@@ -683,6 +710,9 @@ func lockAndMergeAccountProbeExtra(
|
||||
¤tOllamaSession,
|
||||
¤tOllamaAutoRefresh,
|
||||
¤tOllamaSnapshot,
|
||||
&opencodeGroupIdentityUnchanged,
|
||||
¤tOpenCodeAutoRefresh,
|
||||
¤tOpenCodeSnapshot,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -698,6 +728,8 @@ func lockAndMergeAccountProbeExtra(
|
||||
service.OllamaCloudUsageSessionExtraKey,
|
||||
service.OllamaCloudUsageAutoRefreshExtraKey,
|
||||
service.OllamaCloudUsageSnapshotExtraKey,
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey,
|
||||
} {
|
||||
delete(extra, key)
|
||||
}
|
||||
@@ -775,6 +807,25 @@ func lockAndMergeAccountProbeExtra(
|
||||
}
|
||||
}
|
||||
}
|
||||
// OpenCode Go 受管状态与 Ollama 同范式:组身份(opencode_go Go 订阅,或挂载
|
||||
// 白名单平台 + 官方 OpenCode Go 基址,且 api_key 不变)未变时从锁定的 DB 行
|
||||
// 回填,客户端 DTO 脱敏或并发 refresh 写入都无法覆盖/丢失;身份改变或不再
|
||||
// eligible 时既不回填,写出的 extra 即清除旧状态。
|
||||
// 代理不属于组身份,但 snapshot 外呼与 CAS 经代理,故代理变化时快照失效而开关保留。
|
||||
if service.IsOpenCodeGoUsageAccount(account) && opencodeGroupIdentityUnchanged {
|
||||
if value, ok, err := decodeAccountExtraJSON(currentOpenCodeAutoRefresh); err != nil {
|
||||
return nil, err
|
||||
} else if ok {
|
||||
extra[service.OpenCodeGoUsageAutoRefreshExtraKey] = value
|
||||
}
|
||||
if ollamaProxyIdentityUnchanged {
|
||||
if snapshot, ok, err := decodeAccountExtraJSON(currentOpenCodeSnapshot); err != nil {
|
||||
return nil, err
|
||||
} else if ok {
|
||||
extra[service.OpenCodeGoUsageSnapshotExtraKey] = snapshot
|
||||
}
|
||||
}
|
||||
}
|
||||
return extra, nil
|
||||
}
|
||||
|
||||
@@ -817,6 +868,44 @@ func (r *accountRepository) UpdateCredentials(ctx context.Context, id int64, cre
|
||||
SET
|
||||
credentials = $1::jsonb,
|
||||
extra = CASE
|
||||
-- 正确性依赖(非防御):OpenCode 分支必须先于 Ollama 分支求值。两分支
|
||||
-- 的 WHEN 并不互斥:Ollama 分支的守卫是宽谓词——NOT(ollamaMatch(old)
|
||||
-- AND ollamaMatch(new)) 在旧行不匹配 ollama.com 基址时恒真,且两侧
|
||||
-- 平台白名单完全相同,因此挂载行(白名单平台 + 官方 OpenCode Go 基址)
|
||||
-- 会同时满足两分支的 WHEN。若把 Ollama 分支前移,挂载行的 api_key 变化
|
||||
-- 会先命中 Ollama 分支,opencode_go_usage_snapshot /
|
||||
-- opencode_go_usage_auto_refresh 残留,陈旧快照跟着新 api_key 走,
|
||||
-- 造成跨 key 组污染。互斥的只是两侧身份谓词(opencode.ai 基址正则 vs
|
||||
-- ollama.com 基址正则):Ollama 行不会被 OpenCode 分支遮蔽;
|
||||
-- opencode_go 平台行不在 Ollama 白名单内,与 Ollama 分支真互斥。
|
||||
WHEN type = 'apikey'
|
||||
AND credentials IS DISTINCT FROM $1::jsonb
|
||||
AND (
|
||||
-- 旧行按新谓词属于 OpenCode 用量身份
|
||||
(platform = 'opencode_go'
|
||||
AND COALESCE(btrim(credentials ->> 'account_mode') <> 'zen', true))
|
||||
OR (platform IN (`+opencodeGoUsageMountPlatformsSQL+`)
|
||||
AND `+opencodeGoBaseURLMatchSQLPrefix+`credentials ->> 'base_url'`+opencodeGoBaseURLMatchSQLSuffix+`)
|
||||
)
|
||||
AND (
|
||||
credentials -> 'api_key' IS DISTINCT FROM $1::jsonb -> 'api_key'
|
||||
-- 或组身份不再成立:opencode_go 转 Zen(COALESCE 恒非 NULL,
|
||||
-- IS NOT TRUE 等价于判定为假;平台条件必须在 IS NOT TRUE 作用
|
||||
-- 域之外,否则挂载行的 FALSE IS NOT TRUE 会误判为已变化);
|
||||
-- 挂载行 base_url 不再指向官方 Go 基址。NULL-safe:
|
||||
-- regex(NULL) 为 NULL 时 IS NOT TRUE 把 NULL 视为不匹配。
|
||||
OR (platform = 'opencode_go'
|
||||
AND COALESCE(btrim($1::jsonb ->> 'account_mode') <> 'zen', true) IS NOT TRUE)
|
||||
OR (platform IN (`+opencodeGoUsageMountPlatformsSQL+`)
|
||||
AND (`+opencodeGoBaseURLMatchSQLPrefix+`$1::jsonb ->> 'base_url'`+opencodeGoBaseURLMatchSQLSuffix+`) IS NOT TRUE)
|
||||
)
|
||||
THEN COALESCE(extra, '{}'::jsonb)
|
||||
- 'upstream_billing_probe'
|
||||
- 'opencode_go_usage_auto_refresh'
|
||||
- 'opencode_go_usage_snapshot'
|
||||
- 'ollama_cloud_usage_session'
|
||||
- 'ollama_cloud_usage_auto_refresh'
|
||||
- 'ollama_cloud_usage_snapshot'
|
||||
-- 凭证整体未变化 ⇒ Ollama 组身份必然未变化;顶层 DISTINCT 守卫防止
|
||||
-- 非 Ollama 账号的无变化持久化误清探测快照或重写 NULL extra。
|
||||
WHEN platform IN (`+ollamaCloudUsagePlatformsSQL+`)
|
||||
@@ -3007,7 +3096,39 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
" AND "+ollamaCloudBaseURLMatchesSQL(credentialPlaceholder+"::jsonb ->> 'base_url'")+")")
|
||||
}
|
||||
|
||||
if len(updates.Extra) > 0 || len(ollamaGroupIdentityChanges) > 0 || ollamaProxyIdentityChanged != "" || updates.EnsureCodexFingerprintSeed {
|
||||
// OpenCode Go 身份清理与 Ollama 同范式但更精确:仅当旧行按新谓词属于 OpenCode
|
||||
// 用量身份(opencode_go Go 订阅,或挂载白名单平台 + 官方 OpenCode Go 基址)时
|
||||
// 才可能持有受管键,且只在该组身份(api_key / normalized base URL /
|
||||
// account_mode)真正变化或不再 eligible 时清除,避免与 Ollama 分支交叉误清。
|
||||
opencodeGroupIdentityChanges := make([]string, 0, 3)
|
||||
opencodeOldBaseURL := opencodeGoBaseURLMatchSQLPrefix + "credentials ->> 'base_url'" + opencodeGoBaseURLMatchSQLSuffix
|
||||
// 旧行 OpenCode 用量身份谓词(与 opencodeGoUsageEligibleSQL 的行侧条件镜像;
|
||||
// type 由 opencodeEligibleAccount 统一携带)。account_mode 未设置/为 null/非
|
||||
// zen 均视为 Go 订阅,与 GetOpenCodeAccountMode 默认兼容一致。
|
||||
opencodeOldUsageIdentity := "((platform = 'opencode_go' AND COALESCE(btrim(credentials ->> 'account_mode') <> 'zen', true))" +
|
||||
" OR (platform IN (" + opencodeGoUsageMountPlatformsSQL + ") AND " + opencodeOldBaseURL + "))"
|
||||
if _, ok := updates.Credentials["api_key"]; ok {
|
||||
opencodeGroupIdentityChanges = append(opencodeGroupIdentityChanges,
|
||||
opencodeOldUsageIdentity+" AND credentials -> 'api_key' IS DISTINCT FROM "+credentialPlaceholder+"::jsonb -> 'api_key'")
|
||||
}
|
||||
if _, ok := updates.Credentials["base_url"]; ok {
|
||||
// 仅挂载行身份钉在官方基址上(opencode_go 行的 base_url 不参与资格,两个
|
||||
// 官方变体都合法,不由此子句清理)。NULL-safe:新 base_url 缺失/为 null 时
|
||||
// regex(NULL) 为 NULL,NOT NULL 仍为 NULL 会令 WHEN 不命中而残留 OpenCode
|
||||
// 状态;IS NOT TRUE 把 NULL 视为不匹配。
|
||||
opencodeGroupIdentityChanges = append(opencodeGroupIdentityChanges,
|
||||
"platform IN ("+opencodeGoUsageMountPlatformsSQL+") AND "+opencodeOldBaseURL+
|
||||
" AND ("+opencodeGoBaseURLMatchSQLPrefix+credentialPlaceholder+"::jsonb ->> 'base_url'"+opencodeGoBaseURLMatchSQLSuffix+") IS NOT TRUE")
|
||||
}
|
||||
if _, ok := updates.Credentials["account_mode"]; ok {
|
||||
// 仅 opencode_go 行身份钉在 Go 订阅上(挂载行的 account_mode 与 OpenCode
|
||||
// 资格无关,不由此子句清理)。模式转 Zen 即不再 eligible。
|
||||
opencodeGroupIdentityChanges = append(opencodeGroupIdentityChanges,
|
||||
"platform = 'opencode_go' AND COALESCE(btrim(credentials ->> 'account_mode') <> 'zen', true)"+
|
||||
" AND COALESCE(btrim("+credentialPlaceholder+"::jsonb ->> 'account_mode') <> 'zen', true) IS NOT TRUE")
|
||||
}
|
||||
|
||||
if len(updates.Extra) > 0 || len(ollamaGroupIdentityChanges) > 0 || len(opencodeGroupIdentityChanges) > 0 || ollamaProxyIdentityChanged != "" || updates.EnsureCodexFingerprintSeed {
|
||||
extraExpression := "COALESCE(extra, '{}'::jsonb)"
|
||||
if len(updates.Extra) > 0 {
|
||||
payload, err := json.Marshal(updates.Extra)
|
||||
@@ -3038,13 +3159,54 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
snapshotIdentityChanged = "(" + snapshotIdentityChanged + " OR " + proxyChanged + ")"
|
||||
}
|
||||
}
|
||||
// 代理不属于 OpenCode 组身份,但 snapshot 外呼与 CAS 经代理:
|
||||
// 组身份变化清除开关+快照,仅代理变化清除快照而保留开关。
|
||||
// eligible 判定必须包含旧行 OpenCode 身份(opencode_go Go 订阅,或挂载
|
||||
// 白名单平台 + opencode 基址):否则白名单平台的 Ollama 行在代理变化时会
|
||||
// 先命中 OpenCode 分支(CASE 按序求值)而遮蔽 Ollama 分支的快照清理。
|
||||
// 注意 OpenCode 与 Ollama 两套分支的 WHEN 并不互斥:eligibleAccount 只看
|
||||
// platform 白名单 + type='apikey',根本不含 base_url,且与 OpenCode 挂载
|
||||
// 白名单完全相同,因此挂载行(白名单平台 + 官方 OpenCode Go 基址)会同时
|
||||
// 满足两套分支的 WHEN;互斥的只是两侧身份谓词(opencode.ai 基址正则 vs
|
||||
// ollama.com 基址正则,host 不同)。因此下方 caseBranches 必须先拼 OpenCode
|
||||
// 分支再拼 Ollama 分支——这是正确性要求而非防御:若 Ollama 分支在前,
|
||||
// 挂载行的 api_key 变化会先命中 Ollama 分支,opencode_go_usage_snapshot /
|
||||
// opencode_go_usage_auto_refresh 残留,陈旧快照跟着新 api_key 走,跨 key
|
||||
// 组污染。反方向安全:Ollama 行(ollama.com 基址)不满足 OpenCode 身份
|
||||
// 谓词;opencode_go 平台行不在 Ollama 白名单内,与 Ollama 分支真互斥。
|
||||
opencodeEligibleAccount := "type = 'apikey' AND " + opencodeOldUsageIdentity
|
||||
opencodeGroupIdentityChanged := ""
|
||||
if len(opencodeGroupIdentityChanges) > 0 {
|
||||
opencodeGroupIdentityChanged = "(" + opencodeEligibleAccount + " AND (" + joinClauses(opencodeGroupIdentityChanges, " OR ") + "))"
|
||||
}
|
||||
opencodeSnapshotIdentityChanged := opencodeGroupIdentityChanged
|
||||
if ollamaProxyIdentityChanged != "" {
|
||||
opencodeProxyChanged := "(" + opencodeEligibleAccount + " AND " + ollamaProxyIdentityChanged + ")"
|
||||
if opencodeSnapshotIdentityChanged == "" {
|
||||
opencodeSnapshotIdentityChanged = opencodeProxyChanged
|
||||
} else {
|
||||
opencodeSnapshotIdentityChanged = "(" + opencodeSnapshotIdentityChanged + " OR " + opencodeProxyChanged + ")"
|
||||
}
|
||||
}
|
||||
caseBranches := make([]string, 0, 4)
|
||||
if opencodeGroupIdentityChanged != "" {
|
||||
caseBranches = append(caseBranches,
|
||||
" WHEN "+opencodeGroupIdentityChanged+" THEN ("+extraExpression+") - 'opencode_go_usage_auto_refresh' - 'opencode_go_usage_snapshot'")
|
||||
}
|
||||
if opencodeSnapshotIdentityChanged != "" {
|
||||
caseBranches = append(caseBranches,
|
||||
" WHEN "+opencodeSnapshotIdentityChanged+" THEN ("+extraExpression+") - 'opencode_go_usage_snapshot'")
|
||||
}
|
||||
if groupIdentityChanged != "" {
|
||||
extraExpression = "CASE" +
|
||||
" WHEN " + groupIdentityChanged + " THEN (" + extraExpression + ") - 'ollama_cloud_usage_session' - 'ollama_cloud_usage_auto_refresh' - 'ollama_cloud_usage_snapshot'" +
|
||||
" WHEN " + snapshotIdentityChanged + " THEN (" + extraExpression + ") - 'ollama_cloud_usage_snapshot'" +
|
||||
" ELSE " + extraExpression + " END"
|
||||
} else if snapshotIdentityChanged != "" {
|
||||
extraExpression = "CASE WHEN " + snapshotIdentityChanged + " THEN (" + extraExpression + ") - 'ollama_cloud_usage_snapshot' ELSE " + extraExpression + " END"
|
||||
caseBranches = append(caseBranches,
|
||||
" WHEN "+groupIdentityChanged+" THEN ("+extraExpression+") - 'ollama_cloud_usage_session' - 'ollama_cloud_usage_auto_refresh' - 'ollama_cloud_usage_snapshot'")
|
||||
}
|
||||
if snapshotIdentityChanged != "" {
|
||||
caseBranches = append(caseBranches,
|
||||
" WHEN "+snapshotIdentityChanged+" THEN ("+extraExpression+") - 'ollama_cloud_usage_snapshot'")
|
||||
}
|
||||
if len(caseBranches) > 0 {
|
||||
extraExpression = "CASE" + strings.Join(caseBranches, "") + " ELSE " + extraExpression + " END"
|
||||
}
|
||||
if updates.EnsureCodexFingerprintSeed {
|
||||
extraExpression = ensureCodexFingerprintSeedSQL(extraExpression)
|
||||
@@ -3885,8 +4047,12 @@ func (r *accountRepository) ResetQuotaUsedAndClearRateLimitCooldown(ctx context.
|
||||
// 仅当 proxy_fallback_origin_id IS NOT NULL 时执行更新;
|
||||
// 若影响行数为 0,则返回 ErrAccountNotInFallback(账号存在但不在 fallback 状态)。
|
||||
func (r *accountRepository) RevertProxyFallback(ctx context.Context, accountID int64) error {
|
||||
// Probe snapshots belong to the network identity; invalidate only on a real proxy change.
|
||||
res, err := r.sql.ExecContext(ctx, `
|
||||
UPDATE accounts SET proxy_id=proxy_fallback_origin_id, proxy_fallback_origin_id=NULL, updated_at=NOW()
|
||||
UPDATE accounts SET
|
||||
extra=CASE WHEN type='apikey' AND proxy_id IS DISTINCT FROM proxy_fallback_origin_id
|
||||
THEN extra - 'upstream_billing_probe' ELSE extra END,
|
||||
proxy_id=proxy_fallback_origin_id, proxy_fallback_origin_id=NULL, updated_at=NOW()
|
||||
WHERE id=$1 AND proxy_fallback_origin_id IS NOT NULL AND deleted_at IS NULL`, accountID)
|
||||
if err != nil {
|
||||
return err
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"regexp"
|
||||
"testing"
|
||||
|
||||
"entgo.io/ent/dialect"
|
||||
entsql "entgo.io/ent/dialect/sql"
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestUpdateExtraCodexDisplaySnapshotsAvoidSchedulerOutbox(t *testing.T) {
|
||||
for _, key := range []string{"codex_credits_snapshot", "codex_referral_snapshot"} {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
value any
|
||||
schedulingChange bool
|
||||
}{
|
||||
{"refresh", map[string]any{"fetched_at": 1770000000}, false},
|
||||
{"invalidate", nil, false},
|
||||
{"mixed scheduling update", nil, true},
|
||||
} {
|
||||
t.Run(key+"/"+tc.name, func(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db)))
|
||||
t.Cleanup(func() { _ = client.Close() })
|
||||
repo := newAccountRepositoryWithSQL(client, db, nil)
|
||||
updates := map[string]any{key: tc.value}
|
||||
if tc.schedulingChange {
|
||||
updates["openai_compact_enabled"] = true
|
||||
mock.ExpectBegin()
|
||||
}
|
||||
payload, err := json.Marshal(updates)
|
||||
require.NoError(t, err)
|
||||
mock.ExpectExec(regexp.QuoteMeta("UPDATE accounts SET extra = COALESCE(extra, '{}'::jsonb) || $1::jsonb, updated_at = NOW() WHERE id = $2 AND deleted_at IS NULL")).
|
||||
WithArgs(string(payload), int64(27)).WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
if tc.schedulingChange {
|
||||
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).
|
||||
WithArgs(service.SchedulerOutboxEventAccountChanged, int64(27), nil, nil, sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
}
|
||||
// Snapshot-only changes must perform just the UPDATE, without even a
|
||||
// transaction or account_changed event that would rebuild scheduler buckets.
|
||||
require.NoError(t, repo.UpdateExtra(context.Background(), 27, updates))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,516 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/lib/pq"
|
||||
)
|
||||
|
||||
const (
|
||||
// 与 service.isOpenCodeGoBaseURL 对齐:Go 侧接受显式默认端口 :443(parsed.Host ==
|
||||
// hostname+":443")与两个官方基址变体(CC/Responses 的 /zen/go/v1 与 Anthropic
|
||||
// 的 /zen/go),SQL 正则必须同样接受,否则同一 base_url 在 Go 判定 eligible 而
|
||||
// SQL 判定不 eligible,身份清理与组查询会漏行。只做 parity,不扩域名/路径;
|
||||
// Zen 基址(/zen、/zen/v1)必须拒绝。
|
||||
opencodeGoBaseURLRegexSQL = `^[hH][tT][tT][pP][sS]://[oO][pP][eE][nN][cC][oO][dD][eE]\.[aA][iI](:[4][4][3])?/[zZ][eE][nN]/[gG][oO](/[vV]1)?/?$`
|
||||
opencodeGoBaseURLMatchSQLPrefix = "btrim("
|
||||
opencodeGoBaseURLMatchSQLSuffix = ") ~ '" + opencodeGoBaseURLRegexSQL + "'"
|
||||
// opencodeGoUsageMountPlatformsSQL 是 service.isOpenCodeGoUsageMountPlatform 的
|
||||
// SQL 镜像:OpenCode Go key 允许挂在 openai/anthropic 与国产 OpenAI 兼容平台
|
||||
// 下复用,与 ollama 的 ollamaCloudUsagePlatformsSQL 完全一致。opencode_go
|
||||
// 平台本身不在名单内——平台账号走下方 account_mode 分支。所有平台白名单 SQL
|
||||
// 只允许引用本常量,不得各处重写字面量,防止漂移。
|
||||
opencodeGoUsageMountPlatformsSQL = "'openai', 'anthropic', 'kimi', 'zhipu', 'deepseek', 'minimax'"
|
||||
// opencodeGoUsageEligibleSQL 与 service.IsOpenCodeGoUsageAccount 互为镜像:
|
||||
// - opencode_go 平台:account_mode 存储于 credentials(domain/constants.go),
|
||||
// 未设置/为 null/非 "zen" 一律视为 Go 订阅,与 GetOpenCodeAccountMode 的
|
||||
// 默认兼容逻辑一致;COALESCE 把 <> 对 NULL 的结果归一为 true。不校验
|
||||
// base_url(平台字段已是权威来源)。
|
||||
// - 挂载白名单平台:base_url 必须匹配官方 OpenCode Go 基址正则。
|
||||
//
|
||||
// 已知且可接受的差异:btrim 只去空格,Go 侧 strings.TrimSpace 还会去
|
||||
// \t\n\v\f\r 等空白;account_mode 等凭证字段实际不会出现这类空白,且改用
|
||||
// btrim(x, E' \t\n\r\f\v') 会在 parity 关键的 SQL 字符串里引入转义脆弱性,
|
||||
// 故保持现状不改行为。
|
||||
opencodeGoUsageEligibleSQL = `
|
||||
(
|
||||
(platform = 'opencode_go'
|
||||
AND COALESCE(btrim(credentials ->> 'account_mode') <> 'zen', true))
|
||||
OR (platform IN (` + opencodeGoUsageMountPlatformsSQL + `)
|
||||
AND ` + opencodeGoBaseURLMatchSQLPrefix + `credentials ->> 'base_url'` + opencodeGoBaseURLMatchSQLSuffix + `)
|
||||
)
|
||||
AND type = 'apikey'
|
||||
AND jsonb_typeof(credentials -> 'api_key') = 'string'
|
||||
`
|
||||
)
|
||||
|
||||
// ListOpenCodeGoUsageGroupAccounts resolves every sibling for all supplied
|
||||
// identities with one ID query and one batch hydration. API keys are query
|
||||
// parameters only; no derived shared key is persisted.
|
||||
func (r *accountRepository) ListOpenCodeGoUsageGroupAccounts(ctx context.Context, accounts []*service.Account) ([]service.Account, error) {
|
||||
if r == nil || r.sql == nil {
|
||||
return nil, service.ErrOpenCodeGoUsageUnavailable
|
||||
}
|
||||
keys := make([]string, 0, len(accounts))
|
||||
seen := make(map[string]struct{}, len(accounts))
|
||||
for _, account := range accounts {
|
||||
if !service.IsOpenCodeGoUsageAccount(account) || account.Credentials == nil {
|
||||
continue
|
||||
}
|
||||
apiKey, ok := account.Credentials["api_key"].(string)
|
||||
if !ok || apiKey == "" {
|
||||
continue
|
||||
}
|
||||
if _, duplicate := seen[apiKey]; duplicate {
|
||||
continue
|
||||
}
|
||||
seen[apiKey] = struct{}{}
|
||||
keys = append(keys, apiKey)
|
||||
}
|
||||
if len(keys) == 0 {
|
||||
return []service.Account{}, nil
|
||||
}
|
||||
rows, err := r.sql.QueryContext(ctx, `
|
||||
SELECT id
|
||||
FROM accounts
|
||||
WHERE deleted_at IS NULL
|
||||
AND `+opencodeGoUsageEligibleSQL+`
|
||||
AND credentials ->> 'api_key' = ANY($1)
|
||||
ORDER BY id
|
||||
`, pq.Array(keys))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
ids := make([]int64, 0, len(keys))
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hydrated, err := r.GetByIDs(ctx, ids)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result := make([]service.Account, 0, len(hydrated))
|
||||
for _, account := range hydrated {
|
||||
if account != nil {
|
||||
result = append(result, *account)
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// SetOpenCodeGoUsageAutoRefresh persists the group-scoped auto-refresh switch.
|
||||
func (r *accountRepository) SetOpenCodeGoUsageAutoRefresh(ctx context.Context, account *service.Account, enabled bool) error {
|
||||
if account == nil {
|
||||
return service.ErrAccountNilInput
|
||||
}
|
||||
if r == nil || r.client == nil || !service.IsOpenCodeGoUsageAccount(account) {
|
||||
return service.ErrOpenCodeGoUsageUnavailable
|
||||
}
|
||||
return r.updateOpenCodeGoUsageGroup(ctx, account, map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: enabled,
|
||||
}, nil, true)
|
||||
}
|
||||
|
||||
// UpdateOpenCodeGoUsageSnapshot persists the group-scoped usage snapshot.
|
||||
func (r *accountRepository) UpdateOpenCodeGoUsageSnapshot(ctx context.Context, account *service.Account, snapshot *service.OpenCodeGoUsageSnapshot) error {
|
||||
if account == nil || snapshot == nil {
|
||||
return service.ErrAccountNilInput
|
||||
}
|
||||
if r == nil || r.client == nil || !service.IsOpenCodeGoUsageAccount(account) {
|
||||
return service.ErrOpenCodeGoUsageUnavailable
|
||||
}
|
||||
return r.updateOpenCodeGoUsageGroup(ctx, account, map[string]any{
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: snapshot,
|
||||
}, nil, true)
|
||||
}
|
||||
|
||||
// DisableOpenCodeGoUsageAutoRefresh is group-scoped and retains the loaded
|
||||
// identity CAS. It cannot disable a new group after the account changes key.
|
||||
func (r *accountRepository) DisableOpenCodeGoUsageAutoRefresh(ctx context.Context, account *service.Account) error {
|
||||
if account == nil {
|
||||
return service.ErrAccountNilInput
|
||||
}
|
||||
if r == nil || r.client == nil || !service.IsOpenCodeGoUsageAccount(account) {
|
||||
return service.ErrOpenCodeGoUsageUnavailable
|
||||
}
|
||||
payload := openCodeGoUsageManagedPayload(account)
|
||||
payload[service.OpenCodeGoUsageAutoRefreshExtraKey] = false
|
||||
delete(payload, service.OpenCodeGoUsageSnapshotExtraKey)
|
||||
return r.updateOpenCodeGoUsageGroup(ctx, account, payload,
|
||||
[]string{service.OpenCodeGoUsageSnapshotExtraKey}, true)
|
||||
}
|
||||
|
||||
func openCodeGoUsageManagedPayload(account *service.Account) map[string]any {
|
||||
payload := make(map[string]any, 2)
|
||||
if account == nil || account.Extra == nil {
|
||||
return payload
|
||||
}
|
||||
for _, key := range []string{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey,
|
||||
} {
|
||||
if value, ok := account.Extra[key]; ok {
|
||||
payload[key] = value
|
||||
}
|
||||
}
|
||||
return payload
|
||||
}
|
||||
|
||||
type lockedOpenCodeGoUsageMember struct {
|
||||
id int64
|
||||
anchorMatches bool
|
||||
autoJSON string
|
||||
snapshotJSON string
|
||||
}
|
||||
|
||||
// updateOpenCodeGoUsageGroup locks every member of the exact api_key group and
|
||||
// merges the payload onto each member's extra. The merge is pure
|
||||
// (COALESCE(extra,'{}') || payload) unless deleteKeys is non-empty, in which
|
||||
// case those keys are removed first — used only by the group-level disable path
|
||||
// to clear the stale snapshot. Managed keys are never removed on ordinary
|
||||
// writes, so a snapshot write can never wipe the auto-refresh switch or vice versa.
|
||||
func (r *accountRepository) updateOpenCodeGoUsageGroup(
|
||||
ctx context.Context,
|
||||
account *service.Account,
|
||||
payload map[string]any,
|
||||
deleteKeys []string,
|
||||
requireExpectedState bool,
|
||||
) error {
|
||||
if account == nil {
|
||||
return service.ErrAccountNilInput
|
||||
}
|
||||
if r == nil || r.client == nil || !service.IsOpenCodeGoUsageAccount(account) {
|
||||
return service.ErrOpenCodeGoUsageUnavailable
|
||||
}
|
||||
apiKey, ok := account.Credentials["api_key"].(string)
|
||||
if !ok || apiKey == "" {
|
||||
return service.ErrOpenCodeGoUsageAccountInvalid
|
||||
}
|
||||
apply := func(txCtx context.Context, client *dbent.Client) error {
|
||||
matchesProxy, err := lockAndMatchProbeProxyIdentity(txCtx, client, account)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !matchesProxy {
|
||||
return service.ErrOpenCodeGoUsageIdentityChanged
|
||||
}
|
||||
members, err := lockOpenCodeGoUsageGroup(txCtx, client, account, apiKey)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
anchorMatches := false
|
||||
for _, member := range members {
|
||||
anchorMatches = anchorMatches || member.anchorMatches
|
||||
}
|
||||
if !anchorMatches {
|
||||
return service.ErrOpenCodeGoUsageIdentityChanged
|
||||
}
|
||||
if requireExpectedState {
|
||||
expectedAuto, err := canonicalAccountExtraJSON(account, service.OpenCodeGoUsageAutoRefreshExtraKey)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
expectedSnapshot, err := canonicalAccountExtraJSON(account, service.OpenCodeGoUsageSnapshotExtraKey)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
stateMatches := false
|
||||
for _, member := range members {
|
||||
if canonicalJSON(member.autoJSON) == expectedAuto &&
|
||||
canonicalJSON(member.snapshotJSON) == expectedSnapshot {
|
||||
stateMatches = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !stateMatches {
|
||||
return service.ErrOpenCodeGoUsageIdentityChanged
|
||||
}
|
||||
}
|
||||
encoded, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
memberIDs := make([]int64, len(members))
|
||||
for index := range members {
|
||||
memberIDs[index] = members[index].id
|
||||
}
|
||||
mergeExpr := "COALESCE(extra, '{}'::jsonb)"
|
||||
for _, key := range deleteKeys {
|
||||
mergeExpr += " - " + pq.QuoteLiteral(key)
|
||||
}
|
||||
mergeExpr += " || $1::jsonb"
|
||||
result, err := client.ExecContext(txCtx, `
|
||||
UPDATE accounts
|
||||
SET extra = `+mergeExpr+`,
|
||||
updated_at = NOW()
|
||||
WHERE deleted_at IS NULL
|
||||
AND `+opencodeGoUsageEligibleSQL+`
|
||||
AND credentials ->> 'api_key' = $2
|
||||
AND id = ANY($3)
|
||||
`, string(encoded), apiKey, pq.Array(memberIDs))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
affected, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if affected != int64(len(members)) {
|
||||
return service.ErrOpenCodeGoUsageIdentityChanged
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if dbent.TxFromContext(ctx) != nil {
|
||||
return apply(ctx, clientFromContext(ctx, r.client))
|
||||
}
|
||||
tx, err := r.client.Tx(ctx)
|
||||
if errors.Is(err, dbent.ErrTxStarted) {
|
||||
return apply(ctx, r.client)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
txCtx := dbent.NewTxContext(ctx, tx)
|
||||
if err := apply(txCtx, tx.Client()); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func lockOpenCodeGoUsageGroup(
|
||||
ctx context.Context,
|
||||
client *dbent.Client,
|
||||
account *service.Account,
|
||||
apiKey string,
|
||||
) ([]lockedOpenCodeGoUsageMember, error) {
|
||||
credentials, err := json.Marshal(normalizeJSONMap(account.Credentials))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var proxyID any
|
||||
if account.ProxyID != nil {
|
||||
proxyID = *account.ProxyID
|
||||
}
|
||||
rows, err := client.QueryContext(ctx, `
|
||||
SELECT
|
||||
id,
|
||||
id = $2
|
||||
AND platform = $3
|
||||
AND type = $4
|
||||
AND credentials = $5::jsonb
|
||||
AND proxy_id IS NOT DISTINCT FROM $6,
|
||||
COALESCE((extra -> 'opencode_go_usage_auto_refresh')::text, 'null'),
|
||||
COALESCE((extra -> 'opencode_go_usage_snapshot')::text, 'null')
|
||||
FROM accounts
|
||||
WHERE deleted_at IS NULL
|
||||
AND `+opencodeGoUsageEligibleSQL+`
|
||||
AND credentials ->> 'api_key' = $1
|
||||
ORDER BY id
|
||||
FOR NO KEY UPDATE
|
||||
`, apiKey, account.ID, account.Platform, account.Type, string(credentials), proxyID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
members := make([]lockedOpenCodeGoUsageMember, 0, 1)
|
||||
for rows.Next() {
|
||||
var member lockedOpenCodeGoUsageMember
|
||||
if err := rows.Scan(&member.id, &member.anchorMatches, &member.autoJSON, &member.snapshotJSON); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
members = append(members, member)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(members) == 0 {
|
||||
return nil, service.ErrOpenCodeGoUsageIdentityChanged
|
||||
}
|
||||
return members, nil
|
||||
}
|
||||
|
||||
// ListDueOpenCodeGoUsageAccounts returns at most one truly-due activity-driven
|
||||
// candidate per exact API key. Due timing (debounce, max-wait, failure backoff)
|
||||
// is evaluated in SQL before LIMIT so non-due active groups cannot starve due ones.
|
||||
// Account.LastUsedAt is stamped with the group MAX(last_used_at) for a service
|
||||
// pure-function recheck against races between list and refresh.
|
||||
//
|
||||
// Rules mirror service.openCodeGoUsageAutoRefreshDueAt (keep both in sync):
|
||||
// - missing/invalid snapshot or times → fail-open first due
|
||||
// - success: activity after fetched_at;
|
||||
// due_at = GREATEST(LEAST(last_used+debounce, fetched+maxWait), fetched+minFetchInterval)
|
||||
// - failed/unauthorized: activity after last_attempt; activity_due = LEAST(...);
|
||||
// final due_at is not earlier than a valid next_refresh_at (invalid/missing fail-open)
|
||||
func (r *accountRepository) ListDueOpenCodeGoUsageAccounts(
|
||||
ctx context.Context,
|
||||
now time.Time,
|
||||
debounce, maxWait time.Duration,
|
||||
limit int,
|
||||
) ([]service.Account, error) {
|
||||
if limit <= 0 {
|
||||
return []service.Account{}, nil
|
||||
}
|
||||
if r == nil || r.sql == nil {
|
||||
return nil, errors.New("account repository SQL executor not configured")
|
||||
}
|
||||
if debounce <= 0 {
|
||||
debounce = time.Minute
|
||||
}
|
||||
if maxWait <= 0 {
|
||||
maxWait = 15 * time.Minute
|
||||
}
|
||||
debounceSeconds := debounce.Seconds()
|
||||
maxWaitSeconds := maxWait.Seconds()
|
||||
minFetchIntervalSeconds := service.OpenCodeGoUsageMinFetchInterval.Seconds()
|
||||
rows, err := r.sql.QueryContext(ctx, `
|
||||
WITH eligible AS (
|
||||
SELECT id,
|
||||
credentials ->> 'api_key' AS api_key,
|
||||
last_used_at,
|
||||
extra -> 'opencode_go_usage_snapshot' AS snapshot
|
||||
FROM accounts
|
||||
WHERE deleted_at IS NULL
|
||||
AND status = 'active'
|
||||
AND `+opencodeGoUsageEligibleSQL+`
|
||||
AND extra @> '{"opencode_go_usage_auto_refresh": true}'::jsonb
|
||||
), group_activity AS (
|
||||
SELECT credentials ->> 'api_key' AS api_key,
|
||||
MAX(last_used_at) AS group_last_used_at
|
||||
FROM accounts
|
||||
WHERE deleted_at IS NULL
|
||||
AND `+opencodeGoUsageEligibleSQL+`
|
||||
AND jsonb_typeof(credentials -> 'api_key') = 'string'
|
||||
GROUP BY credentials ->> 'api_key'
|
||||
), joined AS (
|
||||
SELECT e.id, e.api_key, e.snapshot, g.group_last_used_at,
|
||||
e.snapshot #>> '{status}' AS status,
|
||||
e.snapshot #>> '{fetched_at}' AS fetched_at,
|
||||
e.snapshot #>> '{last_attempt_at}' AS last_attempt_at,
|
||||
e.snapshot #>> '{next_refresh_at}' AS next_refresh_at
|
||||
FROM eligible e
|
||||
JOIN group_activity g ON g.api_key = e.api_key
|
||||
), parsed AS MATERIALIZED (
|
||||
SELECT id, api_key, snapshot, group_last_used_at, status,
|
||||
`+ollamaCloudUsageParseRFC3339SQL("fetched_at")+` AS parsed_fetched_at,
|
||||
`+ollamaCloudUsageParseRFC3339SQL("last_attempt_at")+` AS parsed_last_attempt_at,
|
||||
`+ollamaCloudUsageParseRFC3339SQL("next_refresh_at")+` AS parsed_next_refresh_at
|
||||
FROM joined
|
||||
), timed AS (
|
||||
SELECT *,
|
||||
CASE
|
||||
WHEN status = 'ok'
|
||||
AND parsed_fetched_at IS NOT NULL
|
||||
AND group_last_used_at IS NOT NULL
|
||||
AND group_last_used_at > parsed_fetched_at::timestamptz
|
||||
THEN GREATEST(
|
||||
LEAST(
|
||||
group_last_used_at + make_interval(secs => $2::double precision),
|
||||
parsed_fetched_at::timestamptz + make_interval(secs => $3::double precision)
|
||||
),
|
||||
parsed_fetched_at::timestamptz + make_interval(secs => $5::double precision)
|
||||
)
|
||||
WHEN status IN ('failed', 'unauthorized')
|
||||
AND parsed_last_attempt_at IS NOT NULL
|
||||
AND group_last_used_at IS NOT NULL
|
||||
AND group_last_used_at > parsed_last_attempt_at::timestamptz
|
||||
THEN GREATEST(
|
||||
LEAST(
|
||||
group_last_used_at + make_interval(secs => $2::double precision),
|
||||
parsed_last_attempt_at::timestamptz + make_interval(secs => $3::double precision)
|
||||
),
|
||||
COALESCE(parsed_next_refresh_at::timestamptz, '-infinity'::timestamptz)
|
||||
)
|
||||
ELSE NULL
|
||||
END AS activity_due_at
|
||||
FROM parsed
|
||||
), candidates AS (
|
||||
SELECT *,
|
||||
CASE
|
||||
WHEN snapshot IS NULL OR snapshot = 'null'::jsonb OR status IS NULL
|
||||
OR status NOT IN ('ok', 'failed', 'unauthorized') THEN 0
|
||||
WHEN status = 'ok' AND parsed_fetched_at IS NULL THEN 0
|
||||
WHEN status IN ('failed', 'unauthorized') AND parsed_last_attempt_at IS NULL THEN 0
|
||||
WHEN activity_due_at IS NOT NULL AND $1 >= activity_due_at THEN 1
|
||||
ELSE NULL
|
||||
END AS due_class,
|
||||
activity_due_at AS due_at
|
||||
FROM timed
|
||||
), ranked AS (
|
||||
SELECT id, api_key, group_last_used_at, due_class, due_at,
|
||||
row_number() OVER (
|
||||
PARTITION BY api_key
|
||||
ORDER BY due_class,
|
||||
due_at NULLS FIRST,
|
||||
id
|
||||
) AS group_rank
|
||||
FROM candidates
|
||||
WHERE due_class IS NOT NULL
|
||||
)
|
||||
SELECT id, group_last_used_at
|
||||
FROM ranked
|
||||
WHERE group_rank = 1
|
||||
ORDER BY due_class, due_at NULLS FIRST, id
|
||||
LIMIT $4
|
||||
`, now.UTC(), debounceSeconds, maxWaitSeconds, limit, minFetchIntervalSeconds)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
type dueRow struct {
|
||||
id int64
|
||||
groupLastUsed *time.Time
|
||||
}
|
||||
rowsOut := make([]dueRow, 0, limit)
|
||||
ids := make([]int64, 0, limit)
|
||||
for rows.Next() {
|
||||
var row dueRow
|
||||
if err := rows.Scan(&row.id, &row.groupLastUsed); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rowsOut = append(rowsOut, row)
|
||||
ids = append(ids, row.id)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
accounts, err := r.GetByIDs(ctx, ids)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
byID := make(map[int64]*service.Account, len(accounts))
|
||||
for _, account := range accounts {
|
||||
if account != nil {
|
||||
byID[account.ID] = account
|
||||
}
|
||||
}
|
||||
result := make([]service.Account, 0, len(rowsOut))
|
||||
for _, row := range rowsOut {
|
||||
account := byID[row.id]
|
||||
if account == nil {
|
||||
continue
|
||||
}
|
||||
// Stamp group MAX(last_used_at) for service due evaluation.
|
||||
if row.groupLastUsed != nil {
|
||||
ts := row.groupLastUsed.UTC()
|
||||
account.LastUsedAt = &ts
|
||||
} else {
|
||||
account.LastUsedAt = nil
|
||||
}
|
||||
result = append(result, *account)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -0,0 +1,371 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// openCodeGoManagedExtra 构造同时携带探测 / Ollama / OpenCode 受管键的 extra,
|
||||
// 用于验证各身份清理分支只清自己该清的键、不跨身份误清。
|
||||
func openCodeGoManagedExtra(now time.Time) map[string]any {
|
||||
return map[string]any{
|
||||
service.UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"},
|
||||
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
|
||||
service.OllamaCloudUsageAutoRefreshExtraKey: true,
|
||||
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OllamaCloudUsageStatusOK, "last_attempt_at": now, "next_refresh_at": now.Add(time.Hour),
|
||||
},
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: true,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OpenCodeGoUsageStatusOK, "last_attempt_at": now, "next_refresh_at": now.Add(time.Hour),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// F1:混合 OpenCode + openai Ollama + anthropic Ollama 行的批量代理变化与
|
||||
// 代理+key 变化,验证 OpenCode/Ollama CASE 分支互斥、各自只清自己该清的键。
|
||||
// 回归点:OpenCode eligible 判定若不含旧行 opencode base URL,OpenAI+Ollama 行
|
||||
// 在代理变化时会先命中 OpenCode 分支而遮蔽 Ollama 快照清理。
|
||||
func TestBulkUpdateOpenCodeGoMixedPlatformIdentityCleanup(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
now := time.Now().UTC()
|
||||
|
||||
create := func(name, platform, baseURL string, extra map[string]any) *service.Account {
|
||||
t.Helper()
|
||||
return mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: name, Platform: platform, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "shared-key", "base_url": baseURL},
|
||||
Extra: extra,
|
||||
})
|
||||
}
|
||||
ollamaExtra := func() map[string]any {
|
||||
return map[string]any{
|
||||
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
|
||||
service.OllamaCloudUsageAutoRefreshExtraKey: true,
|
||||
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OllamaCloudUsageStatusOK, "last_attempt_at": now, "next_refresh_at": now.Add(time.Hour),
|
||||
},
|
||||
}
|
||||
}
|
||||
opencode := create("mixed-opencode", service.PlatformOpenAI, "https://opencode.ai/zen/go/v1", map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: true,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OpenCodeGoUsageStatusOK, "last_attempt_at": now, "next_refresh_at": now.Add(time.Hour),
|
||||
},
|
||||
})
|
||||
openaiOllama := create("mixed-openai-ollama", service.PlatformOpenAI, "https://ollama.com", ollamaExtra())
|
||||
anthropicOllama := create("mixed-anthropic-ollama", service.PlatformAnthropic, "https://ollama.com", ollamaExtra())
|
||||
proxy := mustCreateProxy(t, tx.Client(), &service.Proxy{
|
||||
Name: "mixed-proxy", Protocol: "http", Host: "127.0.0.1", Port: 3128,
|
||||
Username: "user", Password: "pass", Status: service.StatusActive,
|
||||
})
|
||||
proxy2 := mustCreateProxy(t, tx.Client(), &service.Proxy{
|
||||
Name: "mixed-proxy-2", Protocol: "http", Host: "127.0.0.2", Port: 3128,
|
||||
Username: "user", Password: "pass", Status: service.StatusActive,
|
||||
})
|
||||
|
||||
// 只改 proxy:OpenCode 行清快照留开关;Ollama 行清快照留 session/开关。
|
||||
rows, err := repo.BulkUpdate(ctx, []int64{opencode.ID, openaiOllama.ID, anthropicOllama.ID}, service.AccountBulkUpdate{
|
||||
ProxyID: &proxy.ID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(3), rows)
|
||||
|
||||
opencodeLoaded, err := repo.GetByID(ctx, opencode.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, true, opencodeLoaded.Extra[service.OpenCodeGoUsageAutoRefreshExtraKey])
|
||||
require.NotContains(t, opencodeLoaded.Extra, service.OpenCodeGoUsageSnapshotExtraKey)
|
||||
|
||||
for _, id := range []int64{openaiOllama.ID, anthropicOllama.ID} {
|
||||
loaded, err := repo.GetByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "cipher:wos-session=fixture", loaded.Extra[service.OllamaCloudUsageSessionExtraKey])
|
||||
require.Equal(t, true, loaded.Extra[service.OllamaCloudUsageAutoRefreshExtraKey])
|
||||
require.NotContains(t, loaded.Extra, service.OllamaCloudUsageSnapshotExtraKey,
|
||||
"OpenAI+Ollama 行的快照必须由 Ollama 分支清理,不能被 OpenCode 分支遮蔽")
|
||||
}
|
||||
|
||||
// 同时改 proxy + api_key:OpenCode 行清开关+快照;Ollama 行清 session/开关/快照。
|
||||
rows, err = repo.BulkUpdate(ctx, []int64{opencode.ID, openaiOllama.ID, anthropicOllama.ID}, service.AccountBulkUpdate{
|
||||
ProxyID: &proxy2.ID,
|
||||
Credentials: map[string]any{"api_key": "rotated-key"},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(3), rows)
|
||||
|
||||
opencodeLoaded, err = repo.GetByID(ctx, opencode.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, opencodeLoaded.Extra, service.OpenCodeGoUsageAutoRefreshExtraKey)
|
||||
require.NotContains(t, opencodeLoaded.Extra, service.OpenCodeGoUsageSnapshotExtraKey)
|
||||
|
||||
for _, id := range []int64{openaiOllama.ID, anthropicOllama.ID} {
|
||||
loaded, err := repo.GetByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, loaded.Extra, service.OllamaCloudUsageSessionExtraKey)
|
||||
require.NotContains(t, loaded.Extra, service.OllamaCloudUsageAutoRefreshExtraKey)
|
||||
require.NotContains(t, loaded.Extra, service.OllamaCloudUsageSnapshotExtraKey)
|
||||
}
|
||||
}
|
||||
|
||||
// F2:UpdateCredentials 的 OpenCode 分支必须与通用身份变化不变量一致——
|
||||
// 清 probe、OpenCode 两键以及不应跨身份残留的 Ollama session/auto/snapshot。
|
||||
func TestUpdateCredentialsOpenCodeGoClearsAllManagedStateOnIdentityChange(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
now := time.Now().UTC()
|
||||
|
||||
account := mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "opencode-credentials", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "old-key", "base_url": "https://opencode.ai/zen/go/v1"},
|
||||
Extra: openCodeGoManagedExtra(now),
|
||||
})
|
||||
|
||||
require.NoError(t, repo.UpdateCredentials(ctx, account.ID, map[string]any{
|
||||
"api_key": "new-key", "base_url": "https://opencode.ai/zen/go/v1",
|
||||
}))
|
||||
loaded, err := repo.GetByID(ctx, account.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, loaded.Extra, service.UpstreamBillingProbeExtraKey)
|
||||
require.NotContains(t, loaded.Extra, service.OpenCodeGoUsageAutoRefreshExtraKey)
|
||||
require.NotContains(t, loaded.Extra, service.OpenCodeGoUsageSnapshotExtraKey)
|
||||
require.NotContains(t, loaded.Extra, service.OllamaCloudUsageSessionExtraKey)
|
||||
require.NotContains(t, loaded.Extra, service.OllamaCloudUsageAutoRefreshExtraKey)
|
||||
require.NotContains(t, loaded.Extra, service.OllamaCloudUsageSnapshotExtraKey)
|
||||
}
|
||||
|
||||
// F3:新 credentials 缺 base_url(NULL)时 OpenCode 分支必须仍命中(NULL-safe),
|
||||
// 否则 OpenCode 受管状态会残留。api_key 保持不变,专门压 NOT regex(NULL) 路径。
|
||||
func TestUpdateCredentialsOpenCodeGoMissingBaseURLClearsManagedState(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
now := time.Now().UTC()
|
||||
|
||||
account := mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "opencode-missing-base-url", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "old-key", "base_url": "https://opencode.ai/zen/go/v1"},
|
||||
Extra: map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: true,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OpenCodeGoUsageStatusOK, "last_attempt_at": now, "next_refresh_at": now.Add(time.Hour),
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
require.NoError(t, repo.UpdateCredentials(ctx, account.ID, map[string]any{
|
||||
"api_key": "old-key",
|
||||
}))
|
||||
loaded, err := repo.GetByID(ctx, account.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, loaded.Extra, service.OpenCodeGoUsageAutoRefreshExtraKey)
|
||||
require.NotContains(t, loaded.Extra, service.OpenCodeGoUsageSnapshotExtraKey)
|
||||
}
|
||||
|
||||
// F3 延伸:BulkUpdate 移除 base_url(null)时 opencode 受管键必须清除,
|
||||
// 否则 NOT regex(NULL) 为 NULL 会令 WHEN 不命中而残留。
|
||||
func TestBulkUpdateOpenCodeGoRemovedBaseURLClearsManagedState(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
now := time.Now().UTC()
|
||||
|
||||
account := mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "opencode-bulk-remove-base-url", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "old-key", "base_url": "https://opencode.ai/zen/go/v1"},
|
||||
Extra: map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: true,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OpenCodeGoUsageStatusOK, "last_attempt_at": now, "next_refresh_at": now.Add(time.Hour),
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
rows, err := repo.BulkUpdate(ctx, []int64{account.ID}, service.AccountBulkUpdate{
|
||||
Credentials: map[string]any{"base_url": nil},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), rows)
|
||||
loaded, err := repo.GetByID(ctx, account.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, loaded.Extra, service.OpenCodeGoUsageAutoRefreshExtraKey)
|
||||
require.NotContains(t, loaded.Extra, service.OpenCodeGoUsageSnapshotExtraKey)
|
||||
}
|
||||
|
||||
// F6:SQL eligible 判定与 Go 侧对齐,接受显式默认端口 :443。
|
||||
func TestOpenCodeGoUsageEligibleSQLAcceptsDefaultPort443(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
|
||||
account := mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "opencode-port-443", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "port-key", "base_url": "https://opencode.ai:443/zen/go/v1"},
|
||||
Extra: map[string]any{},
|
||||
})
|
||||
|
||||
// ListOpenCodeGoUsageGroupAccounts 走 opencodeGoUsageEligibleSQL,能解析到
|
||||
// :443 行即证明 SQL 正则与 Go 判定一致。
|
||||
groups, err := repo.ListOpenCodeGoUsageGroupAccounts(ctx, []*service.Account{account})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, groups, 1)
|
||||
require.Equal(t, account.ID, groups[0].ID)
|
||||
}
|
||||
|
||||
// 资格泛化回归:opencode_go 平台 Go 订阅账号(不依赖 base_url,未设置
|
||||
// account_mode 默认 Go)与 Anthropic 基址(/zen/go)挂载行都必须命中
|
||||
// opencodeGoUsageEligibleSQL;同 key 的 Zen 模式行必须被 SQL 排除。
|
||||
func TestOpenCodeGoUsageEligibleSQLOpencodePlatformAndAnthropicBaseURL(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
|
||||
platformGo := mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "opencode-platform-go", Platform: service.PlatformOpenCodeGo, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "platform-key"}, // 无 base_url、无 account_mode(默认 Go)
|
||||
Extra: map[string]any{},
|
||||
})
|
||||
// 同 key 的 Zen 行:SQL 必须排除(account_mode = zen)。
|
||||
platformZen := mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "opencode-platform-zen", Platform: service.PlatformOpenCodeGo, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "platform-key", "account_mode": service.AccountModeZen},
|
||||
Extra: map[string]any{},
|
||||
})
|
||||
anthropicMount := mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "opencode-anthropic-mount", Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "mount-key", "base_url": "https://opencode.ai/zen/go"},
|
||||
Extra: map[string]any{},
|
||||
})
|
||||
|
||||
groups, err := repo.ListOpenCodeGoUsageGroupAccounts(ctx, []*service.Account{platformGo, anthropicMount})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, groups, 2)
|
||||
ids := make(map[int64]struct{}, len(groups))
|
||||
for _, account := range groups {
|
||||
ids[account.ID] = struct{}{}
|
||||
}
|
||||
require.Contains(t, ids, platformGo.ID)
|
||||
require.Contains(t, ids, anthropicMount.ID)
|
||||
require.NotContains(t, ids, platformZen.ID, "同 key 的 Zen 模式行必须被 SQL 谓词排除")
|
||||
}
|
||||
|
||||
// 资格泛化回归:UpdateCredentials 的 OpenCode 清理分支里,IS NOT TRUE 只能作用于
|
||||
// 新凭证的检查本身,平台条件必须在作用域外——否则挂载行 (FALSE AND …) IS NOT TRUE
|
||||
// 恒为 TRUE,仅添加无关凭证键也会误清受管键。两类行在 api_key/身份不变、凭证仅
|
||||
// 无关变化时都必须保留 OpenCode 受管键。
|
||||
func TestUpdateCredentialsOpenCodeGoUnrelatedCredentialChangeKeepsManagedState(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
now := time.Now().UTC()
|
||||
|
||||
create := func(name, platform string, credentials map[string]any) *service.Account {
|
||||
return mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: name, Platform: platform, Type: service.AccountTypeAPIKey,
|
||||
Credentials: credentials,
|
||||
Extra: map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: true,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OpenCodeGoUsageStatusOK, "last_attempt_at": now, "next_refresh_at": now.Add(time.Hour),
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
anthropicMount := create("opencode-mount-unrelated", service.PlatformAnthropic, map[string]any{
|
||||
"api_key": "mount-key", "base_url": "https://opencode.ai/zen/go",
|
||||
})
|
||||
platformGo := create("opencode-platform-unrelated", service.PlatformOpenCodeGo, map[string]any{
|
||||
"api_key": "platform-key", "account_mode": service.AccountModeGo,
|
||||
})
|
||||
|
||||
// 挂载行:api_key 与官方基址都未变,仅新增无关凭证键。
|
||||
require.NoError(t, repo.UpdateCredentials(ctx, anthropicMount.ID, map[string]any{
|
||||
"api_key": "mount-key", "base_url": "https://opencode.ai/zen/go", "remark": "unrelated",
|
||||
}))
|
||||
// opencode_go 行:api_key 未变、模式仍为 Go,仅新增无关凭证键。
|
||||
require.NoError(t, repo.UpdateCredentials(ctx, platformGo.ID, map[string]any{
|
||||
"api_key": "platform-key", "account_mode": service.AccountModeGo, "remark": "unrelated",
|
||||
}))
|
||||
|
||||
for _, id := range []int64{anthropicMount.ID, platformGo.ID} {
|
||||
loaded, err := repo.GetByID(ctx, id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, true, loaded.Extra[service.OpenCodeGoUsageAutoRefreshExtraKey])
|
||||
snapshot, ok := loaded.Extra[service.OpenCodeGoUsageSnapshotExtraKey].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, service.OpenCodeGoUsageStatusOK, snapshot["status"])
|
||||
}
|
||||
}
|
||||
|
||||
// F7:生产故障路径——前端脱敏普通编辑(incoming Extra 不含 OpenCode 受管键)经真实
|
||||
// repo.Update 走 lockAndMergeAccountProbeExtra 的 SELECT/Scan/UPDATE 后,custom 生效
|
||||
// 且两个受管键从锁定 DB 行回填;incoming Extra 伪造受管键时以锁定 DB 值为准。
|
||||
// 旧代码(无 opencode 三列与回填分支)下:脱敏编辑会直接抹掉 DB 里的开关与快照,
|
||||
// 伪造编辑会把伪造值原样写入,两条断言都会失败。
|
||||
func TestUpdateOpenCodeGoManagedStateSurvivesDesensitizedEdit(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
now := time.Now().UTC()
|
||||
|
||||
proxy := mustCreateProxy(t, tx.Client(), &service.Proxy{
|
||||
Name: "opencode-edit-proxy", Protocol: "http", Host: "127.0.0.1", Port: 3128,
|
||||
Username: "user", Password: "pass", Status: service.StatusActive,
|
||||
})
|
||||
account := mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: "opencode-desensitized-edit", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "edit-key", "base_url": "https://opencode.ai/zen/go/v1"},
|
||||
ProxyID: &proxy.ID,
|
||||
Extra: map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: true,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: map[string]any{
|
||||
"status": service.OpenCodeGoUsageStatusOK, "last_attempt_at": now, "next_refresh_at": now.Add(time.Hour),
|
||||
},
|
||||
"custom": "original",
|
||||
},
|
||||
})
|
||||
|
||||
// 前端脱敏普通编辑:incoming Extra 不含两个受管键,只改 custom;base_url 用
|
||||
// :443 + 尾斜杠变体证明 normalized 身份判定(同 ID/platform/type/api_key/proxy)。
|
||||
edit, err := repo.GetByID(ctx, account.ID)
|
||||
require.NoError(t, err)
|
||||
edit.Credentials = map[string]any{"api_key": "edit-key", "base_url": "https://opencode.ai:443/zen/go/v1/"}
|
||||
edit.Extra = map[string]any{"custom": "edited"}
|
||||
require.NoError(t, repo.Update(ctx, edit))
|
||||
|
||||
loaded, err := repo.GetByID(ctx, account.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "edited", loaded.Extra["custom"])
|
||||
require.Equal(t, true, loaded.Extra[service.OpenCodeGoUsageAutoRefreshExtraKey],
|
||||
"脱敏编辑不得抹掉 auto_refresh")
|
||||
snapshot, ok := loaded.Extra[service.OpenCodeGoUsageSnapshotExtraKey].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, service.OpenCodeGoUsageStatusOK, snapshot["status"], "脱敏编辑不得抹掉快照")
|
||||
|
||||
// incoming Extra 伪造受管键:锁定 DB 值必须胜出。
|
||||
forged, err := repo.GetByID(ctx, account.ID)
|
||||
require.NoError(t, err)
|
||||
forged.Extra = map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: false,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: map[string]any{"status": "forged"},
|
||||
"custom": "edited-again",
|
||||
}
|
||||
require.NoError(t, repo.Update(ctx, forged))
|
||||
|
||||
loaded, err = repo.GetByID(ctx, account.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "edited-again", loaded.Extra["custom"])
|
||||
require.Equal(t, true, loaded.Extra[service.OpenCodeGoUsageAutoRefreshExtraKey],
|
||||
"伪造 auto_refresh 必须以锁定 DB 值为准")
|
||||
snapshot, ok = loaded.Extra[service.OpenCodeGoUsageSnapshotExtraKey].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, service.OpenCodeGoUsageStatusOK, snapshot["status"], "伪造快照必须以锁定 DB 值为准")
|
||||
}
|
||||
@@ -0,0 +1,526 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func openCodeGoUsageRepositoryAccount() *service.Account {
|
||||
previousAttempt := time.Date(2026, time.August, 15, 8, 0, 0, 0, time.UTC)
|
||||
return &service.Account{
|
||||
ID: 17,
|
||||
Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "key", "base_url": "https://opencode.ai/zen/go/v1"},
|
||||
Extra: map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: true,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: &service.OpenCodeGoUsageSnapshot{
|
||||
Status: service.OpenCodeGoUsageStatusOK,
|
||||
LastAttemptAt: previousAttempt,
|
||||
NextRefreshAt: previousAttempt.Add(time.Hour),
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateOpenCodeGoUsageSnapshotWritesSnapshotOnly(t *testing.T) {
|
||||
client, mock := newOllamaCloudUsageRepositoryTestClient(t)
|
||||
account := openCodeGoUsageRepositoryAccount()
|
||||
previousSnapshotJSON, err := json.Marshal(account.Extra[service.OpenCodeGoUsageSnapshotExtraKey])
|
||||
require.NoError(t, err)
|
||||
|
||||
attemptedAt := time.Date(2026, time.August, 15, 9, 0, 0, 0, time.UTC)
|
||||
snapshot := &service.OpenCodeGoUsageSnapshot{
|
||||
Status: service.OpenCodeGoUsageStatusOK,
|
||||
LastAttemptAt: attemptedAt,
|
||||
NextRefreshAt: attemptedAt.Add(time.Hour),
|
||||
}
|
||||
expectedPayload, err := json.Marshal(map[string]any{
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: snapshot,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, string(expectedPayload), service.OpenCodeGoUsageAutoRefreshExtraKey)
|
||||
|
||||
credentials, err := json.Marshal(normalizeJSONMap(account.Credentials))
|
||||
require.NoError(t, err)
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")).
|
||||
WithArgs("key", account.ID, account.Platform, account.Type, string(credentials), nil).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "anchor_matches", "auto_refresh", "snapshot"}).
|
||||
AddRow(account.ID, true, `true`, string(previousSnapshotJSON)))
|
||||
mock.ExpectExec(`(?s)`+regexp.QuoteMeta("UPDATE accounts")).
|
||||
WithArgs(string(expectedPayload), "key", sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
repo := newAccountRepositoryWithSQL(client, nil, nil)
|
||||
|
||||
err = repo.UpdateOpenCodeGoUsageSnapshot(context.Background(), account, snapshot)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
// lockAndMergeAccountProbeExtra 的 SELECT 现在多出三列(opencode 组身份 / 开关 / 快照),
|
||||
// 供通用 Update 路径做 OpenCode 受管键的原子回填。
|
||||
func openCodeGoMergeMockColumns() []string {
|
||||
return []string{
|
||||
"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged",
|
||||
"enabled", "rate_sync_enabled", "snapshot",
|
||||
"ollama_session", "ollama_auto", "ollama_snapshot",
|
||||
"opencode_group_unchanged", "opencode_auto", "opencode_snapshot",
|
||||
}
|
||||
}
|
||||
|
||||
func openCodeGoSnapshotJSON() string {
|
||||
previousAttempt := time.Date(2026, time.August, 15, 8, 0, 0, 0, time.UTC)
|
||||
raw, err := json.Marshal(&service.OpenCodeGoUsageSnapshot{
|
||||
Status: service.OpenCodeGoUsageStatusOK,
|
||||
LastAttemptAt: previousAttempt,
|
||||
NextRefreshAt: previousAttempt.Add(time.Hour),
|
||||
})
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return string(raw)
|
||||
}
|
||||
|
||||
func TestLockAndMergeAccountProbeExtraPreservesOpenCodeGoManagedState(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
account *service.Account
|
||||
groupUnchanged bool
|
||||
proxyUnchanged bool
|
||||
databaseAuto string
|
||||
databaseSnapshot string
|
||||
inputExtra map[string]any
|
||||
wantAuto any // nil 表示键不存在
|
||||
wantSnapshot any // nil 表示键不存在
|
||||
wantSnapshotCheck bool
|
||||
}{
|
||||
{
|
||||
name: "identity unchanged restores both managed keys over forged input",
|
||||
account: openCodeGoUsageRepositoryAccount(),
|
||||
groupUnchanged: true,
|
||||
proxyUnchanged: true,
|
||||
databaseAuto: "true",
|
||||
databaseSnapshot: openCodeGoSnapshotJSON(),
|
||||
inputExtra: map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: false,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: map[string]any{"status": "forged"},
|
||||
},
|
||||
wantAuto: true,
|
||||
wantSnapshotCheck: true,
|
||||
},
|
||||
{
|
||||
name: "identity unchanged keeps disabled switch from database",
|
||||
account: openCodeGoUsageRepositoryAccount(),
|
||||
groupUnchanged: true,
|
||||
proxyUnchanged: true,
|
||||
databaseAuto: "false",
|
||||
databaseSnapshot: openCodeGoSnapshotJSON(),
|
||||
inputExtra: map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: true,
|
||||
},
|
||||
wantAuto: false,
|
||||
wantSnapshotCheck: true,
|
||||
},
|
||||
{
|
||||
name: "api key changed clears both managed keys",
|
||||
account: openCodeGoUsageRepositoryAccount(),
|
||||
groupUnchanged: false,
|
||||
proxyUnchanged: true,
|
||||
databaseAuto: "true",
|
||||
databaseSnapshot: openCodeGoSnapshotJSON(),
|
||||
inputExtra: map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: true,
|
||||
service.OpenCodeGoUsageSnapshotExtraKey: map[string]any{"status": "stale"},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "base url no longer eligible clears both managed keys",
|
||||
account: openCodeGoUsageRepositoryAccount(),
|
||||
groupUnchanged: false,
|
||||
proxyUnchanged: true,
|
||||
databaseAuto: "true",
|
||||
databaseSnapshot: openCodeGoSnapshotJSON(),
|
||||
inputExtra: map[string]any{"custom": "value"},
|
||||
},
|
||||
{
|
||||
name: "proxy changed invalidates snapshot but keeps auto-refresh switch",
|
||||
account: openCodeGoUsageRepositoryAccount(),
|
||||
groupUnchanged: true,
|
||||
proxyUnchanged: false,
|
||||
databaseAuto: "true",
|
||||
databaseSnapshot: openCodeGoSnapshotJSON(),
|
||||
inputExtra: map[string]any{"custom": "value"},
|
||||
wantAuto: true,
|
||||
},
|
||||
{
|
||||
name: "no longer eligible account never inherits managed keys",
|
||||
account: func() *service.Account {
|
||||
account := openCodeGoUsageRepositoryAccount()
|
||||
account.Type = service.AccountTypeOAuth
|
||||
return account
|
||||
}(),
|
||||
groupUnchanged: false,
|
||||
proxyUnchanged: true,
|
||||
databaseAuto: "true",
|
||||
databaseSnapshot: openCodeGoSnapshotJSON(),
|
||||
inputExtra: map[string]any{
|
||||
service.OpenCodeGoUsageAutoRefreshExtraKey: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
client, mock := newOllamaCloudUsageRepositoryTestClient(t)
|
||||
credentials, err := json.Marshal(normalizeJSONMap(tt.account.Credentials))
|
||||
require.NoError(t, err)
|
||||
mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")).
|
||||
WithArgs(tt.account.ID, tt.account.Platform, tt.account.Type, string(credentials), nil).
|
||||
WillReturnRows(sqlmock.NewRows(openCodeGoMergeMockColumns()).
|
||||
AddRow(false, false, tt.proxyUnchanged, nil, nil, nil, nil, nil, nil, tt.groupUnchanged, tt.databaseAuto, tt.databaseSnapshot))
|
||||
|
||||
got, err := lockAndMergeAccountProbeExtra(context.Background(), client, tt.account, nil, nil)
|
||||
require.NoError(t, err)
|
||||
if tt.wantAuto == nil {
|
||||
require.NotContains(t, got, service.OpenCodeGoUsageAutoRefreshExtraKey)
|
||||
} else {
|
||||
require.Equal(t, tt.wantAuto, got[service.OpenCodeGoUsageAutoRefreshExtraKey])
|
||||
}
|
||||
if tt.wantSnapshotCheck {
|
||||
snapshot, ok := got[service.OpenCodeGoUsageSnapshotExtraKey].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, service.OpenCodeGoUsageStatusOK, snapshot["status"])
|
||||
} else if tt.wantSnapshot == nil {
|
||||
require.NotContains(t, got, service.OpenCodeGoUsageSnapshotExtraKey)
|
||||
} else {
|
||||
require.Equal(t, tt.wantSnapshot, got[service.OpenCodeGoUsageSnapshotExtraKey])
|
||||
}
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateCredentialsOpenCodeGoIdentityChangeClearsManagedExtra(t *testing.T) {
|
||||
client, mock := newOllamaCloudUsageRepositoryTestClient(t)
|
||||
mock.ExpectBegin()
|
||||
// opencode 清理分支必须文本上先于 ollama 分支出现,否则 opencode 行的
|
||||
// api_key/base_url 变化会被先求值的 Ollama 分支遮蔽。
|
||||
mock.ExpectExec(`(?s)UPDATE accounts.*- 'opencode_go_usage_auto_refresh'.*- 'opencode_go_usage_snapshot'.*- 'ollama_cloud_usage_session'`).
|
||||
WithArgs(`{"api_key":"new-key","base_url":"https://opencode.ai/zen/go/v1"}`, int64(17)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).
|
||||
WithArgs(service.SchedulerOutboxEventAccountChanged, int64(17), nil, nil, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
repo := newAccountRepositoryWithSQL(client, nil, nil)
|
||||
|
||||
err := repo.UpdateCredentials(context.Background(), 17, map[string]any{
|
||||
"api_key": "new-key", "base_url": "https://opencode.ai/zen/go/v1",
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
// opencode → ollama 的 base_url 跨域变化同样必须清除 opencode 受管键(不再 eligible)。
|
||||
func TestUpdateCredentialsOpenCodeGoToOllamaCrossOverClearsManagedExtra(t *testing.T) {
|
||||
client, mock := newOllamaCloudUsageRepositoryTestClient(t)
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(`(?s)UPDATE accounts.*opencode_go_usage_auto_refresh.*opencode_go_usage_snapshot`).
|
||||
WithArgs(`{"api_key":"same-key","base_url":"https://ollama.com/v1"}`, int64(17)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).
|
||||
WithArgs(service.SchedulerOutboxEventAccountChanged, int64(17), nil, nil, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
repo := newAccountRepositoryWithSQL(client, nil, nil)
|
||||
|
||||
err := repo.UpdateCredentials(context.Background(), 17, map[string]any{
|
||||
"api_key": "same-key", "base_url": "https://ollama.com/v1",
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
// OpenCode 清理分支必须带顶层 credentials DISTINCT 守卫:凭证未变化时不能误清受管键。
|
||||
func TestUpdateCredentialsOpenCodeGoCleanupRequiresChangedCredentials(t *testing.T) {
|
||||
client, mock := newOllamaCloudUsageRepositoryTestClient(t)
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(`(?s)UPDATE accounts.*CASE.*AND credentials IS DISTINCT FROM \$1::jsonb`).
|
||||
WithArgs(`{"api_key":"same-key","base_url":"https://relay.example.com/v1"}`, int64(17)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).
|
||||
WithArgs(service.SchedulerOutboxEventAccountChanged, int64(17), nil, nil, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
repo := newAccountRepositoryWithSQL(client, nil, nil)
|
||||
|
||||
err := repo.UpdateCredentials(context.Background(), 17, map[string]any{
|
||||
"api_key": "same-key", "base_url": "https://relay.example.com/v1",
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestBulkUpdateOpenCodeGoIdentityCleanupIsValueConditional(t *testing.T) {
|
||||
exec := &recordingSQLExecutor{result: rowsAffectedResult(1)}
|
||||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||||
|
||||
_, err := repo.BulkUpdate(context.Background(), []int64{17}, service.AccountBulkUpdate{
|
||||
Credentials: map[string]any{"api_key": "new-key"},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, exec.execQueries)
|
||||
query := normalizeSQLWhitespace(exec.execQueries[0])
|
||||
// OpenCode 分支必须出现在 Ollama 分支之前,且 eligible 判定按新谓词同时覆盖
|
||||
// opencode_go 平台(Go 订阅)与挂载白名单平台 + opencode 基址。
|
||||
require.Contains(t, query, "platform = 'opencode_go'")
|
||||
require.Contains(t, query, "platform IN ("+opencodeGoUsageMountPlatformsSQL+")")
|
||||
require.Contains(t, query, "- 'opencode_go_usage_auto_refresh' - 'opencode_go_usage_snapshot'")
|
||||
opencodeBranch := strings.Index(query, "opencode_go_usage_auto_refresh")
|
||||
ollamaBranch := strings.Index(query, "ollama_cloud_usage_auto_refresh")
|
||||
require.NotEqual(t, -1, opencodeBranch)
|
||||
require.NotEqual(t, -1, ollamaBranch)
|
||||
require.Less(t, opencodeBranch, ollamaBranch)
|
||||
}
|
||||
|
||||
// F1 回归:OpenCode eligible 判定必须包含旧行 opencode base URL,否则 OpenAI+Ollama
|
||||
// 行在代理变化时会先命中 OpenCode 分支而遮蔽 Ollama 快照清理。这里断言 OpenCode 分支
|
||||
// 的 WHEN 携带 opencode 正则(与 Ollama 正则互斥),真实行为由 integration 测试覆盖。
|
||||
func TestBulkUpdateOpenCodeGoEligiblePredicateIncludesBaseURL(t *testing.T) {
|
||||
exec := &recordingSQLExecutor{result: rowsAffectedResult(1)}
|
||||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||||
|
||||
proxyID := int64(9)
|
||||
_, err := repo.BulkUpdate(context.Background(), []int64{17}, service.AccountBulkUpdate{
|
||||
ProxyID: &proxyID,
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, exec.execQueries)
|
||||
query := normalizeSQLWhitespace(exec.execQueries[0])
|
||||
// 第一个 WHEN 分支是 OpenCode 快照失效分支(代理变化),其 eligible 判定必须
|
||||
// 包含 opencode base URL 正则,使 OpenAI+Ollama 行无法命中该分支。
|
||||
caseStart := strings.Index(query, "CASE")
|
||||
firstThen := strings.Index(query, "THEN")
|
||||
require.NotEqual(t, -1, caseStart)
|
||||
require.NotEqual(t, -1, firstThen)
|
||||
require.Less(t, caseStart, firstThen)
|
||||
firstWhen := query[caseStart:firstThen]
|
||||
// 第一个 WHEN 分支是 OpenCode 快照失效分支(代理变化),其 eligible 判定必须
|
||||
// 覆盖 opencode_go 平台与挂载白名单的 opencode 基址,使 Ollama 行无法命中。
|
||||
require.Contains(t, firstWhen, "platform = 'opencode_go'")
|
||||
require.Contains(t, firstWhen, "[oO][pP][eE][nN][cC][oO][dD][eE]")
|
||||
require.Contains(t, firstWhen, "credentials ->> 'base_url'")
|
||||
require.NotContains(t, firstWhen, "[oO][lL][lL][aA][mM][aA]")
|
||||
}
|
||||
|
||||
// F3 延伸:BulkUpdate 的 opencode base_url 变化子句同样必须 NULL-safe——新 base_url
|
||||
// 缺失/为 null 时 regex(NULL) 为 NULL,NOT NULL 仍为 NULL 会令 WHEN 不命中而残留
|
||||
// OpenCode 状态;IS NOT TRUE 把 NULL 视为不匹配。
|
||||
func TestBulkUpdateOpenCodeGoBaseURLClauseIsNullSafe(t *testing.T) {
|
||||
exec := &recordingSQLExecutor{result: rowsAffectedResult(1)}
|
||||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||||
|
||||
_, err := repo.BulkUpdate(context.Background(), []int64{17}, service.AccountBulkUpdate{
|
||||
Credentials: map[string]any{"base_url": nil},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, exec.execQueries)
|
||||
query := normalizeSQLWhitespace(exec.execQueries[0])
|
||||
require.Contains(t, query, "IS NOT TRUE")
|
||||
require.NotContains(t, query, "AND NOT btrim(credentials ->> 'base_url')")
|
||||
}
|
||||
|
||||
// F6 回归:SQL 正则与 service.isOpenCodeGoBaseURL 对齐,接受显式默认端口 :443 与
|
||||
// 两个官方基址变体(CC/Responses 的 /zen/go/v1、Anthropic 的 /zen/go)。该正则
|
||||
// 同时用于 eligible 判定与身份清理,Go 侧接受而 SQL 侧拒绝会导致漏清/漏组。
|
||||
func TestOpenCodeGoBaseURLRegexSQLAcceptsDefaultPort443(t *testing.T) {
|
||||
re := regexp.MustCompile(opencodeGoBaseURLRegexSQL)
|
||||
for _, url := range []string{
|
||||
"https://opencode.ai/zen/go/v1",
|
||||
"https://opencode.ai/zen/go/v1/",
|
||||
"https://opencode.ai:443/zen/go/v1",
|
||||
"https://opencode.ai:443/zen/go/v1/",
|
||||
"https://opencode.ai/zen/go",
|
||||
"https://opencode.ai/zen/go/",
|
||||
"https://opencode.ai:443/zen/go",
|
||||
"HTTPS://OPENCODE.AI:443/ZEN/GO/V1",
|
||||
} {
|
||||
require.True(t, re.MatchString(url), "SQL regex must accept %s", url)
|
||||
}
|
||||
for _, url := range []string{
|
||||
"https://opencode.ai:444/zen/go/v1",
|
||||
"https://opencode.ai/v1",
|
||||
"https://ollama.com/zen/go/v1",
|
||||
"https://opencode.ai/zen/go/v1?x=1",
|
||||
// zen 基址必须拒绝:按量付费无订阅配额窗口
|
||||
"https://opencode.ai/zen",
|
||||
"https://opencode.ai/zen/",
|
||||
"https://opencode.ai/zen/v1",
|
||||
"https://opencode.ai/zen/v1/",
|
||||
// 非官方 scheme / 子域 / 路径前缀拼接
|
||||
"http://opencode.ai/zen/go/v1",
|
||||
"https://www.opencode.ai/zen/go/v1",
|
||||
"https://opencode.ai/zen/gov1",
|
||||
} {
|
||||
require.False(t, re.MatchString(url), "SQL regex must reject %s", url)
|
||||
}
|
||||
}
|
||||
|
||||
// SQL 挂载平台白名单常量与 service 判定必须互为镜像:对每个挂载候选平台,常量里
|
||||
// 的成员关系都要与 IsOpenCodeGoUsageAccount(apikey + 官方 OpenCode Go 基址)一致,
|
||||
// 防止两侧平台列表各自漂移。opencode_go 走平台 + Go 订阅分支而非挂载白名单:
|
||||
// 不得出现在白名单常量里,且其资格由 account_mode 决定,在此一并钉死。
|
||||
func TestOpenCodeGoUsagePlatformWhitelistMatchesServicePredicate(t *testing.T) {
|
||||
matches := regexp.MustCompile(`'([^']+)'`).FindAllStringSubmatch(opencodeGoUsageMountPlatformsSQL, -1)
|
||||
sqlPlatforms := make(map[string]struct{}, len(matches))
|
||||
for _, match := range matches {
|
||||
sqlPlatforms[match[1]] = struct{}{}
|
||||
}
|
||||
require.Len(t, sqlPlatforms, 6)
|
||||
for _, platform := range []string{
|
||||
service.PlatformOpenAI, service.PlatformAnthropic,
|
||||
service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, service.PlatformMiniMax,
|
||||
service.PlatformGemini, service.PlatformGrok, service.PlatformAntigravity,
|
||||
service.PlatformComposite, "kiro",
|
||||
} {
|
||||
account := openCodeGoUsageRepositoryAccount()
|
||||
account.Platform = platform
|
||||
_, inSQL := sqlPlatforms[platform]
|
||||
require.Equal(t, inSQL, service.IsOpenCodeGoUsageAccount(account), platform)
|
||||
}
|
||||
// opencode_go 不在挂载白名单里,资格来自平台 + Go 订阅分支。
|
||||
_, opencodeInSQL := sqlPlatforms[service.PlatformOpenCodeGo]
|
||||
require.False(t, opencodeInSQL)
|
||||
goAccount := openCodeGoUsageRepositoryAccount()
|
||||
goAccount.Platform = service.PlatformOpenCodeGo
|
||||
require.True(t, service.IsOpenCodeGoUsageAccount(goAccount))
|
||||
zenAccount := openCodeGoUsageRepositoryAccount()
|
||||
zenAccount.Platform = service.PlatformOpenCodeGo
|
||||
zenAccount.Credentials["account_mode"] = service.AccountModeZen
|
||||
require.False(t, service.IsOpenCodeGoUsageAccount(zenAccount))
|
||||
}
|
||||
|
||||
func TestBulkUpdateOpenCodeGoProxyChangeClearsSnapshotOnly(t *testing.T) {
|
||||
exec := &recordingSQLExecutor{result: rowsAffectedResult(1)}
|
||||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||||
|
||||
proxyID := int64(9)
|
||||
_, err := repo.BulkUpdate(context.Background(), []int64{17}, service.AccountBulkUpdate{
|
||||
ProxyID: &proxyID,
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, exec.execQueries)
|
||||
query := normalizeSQLWhitespace(exec.execQueries[0])
|
||||
require.Contains(t, query, "- 'opencode_go_usage_snapshot'")
|
||||
require.NotContains(t, query, "- 'opencode_go_usage_auto_refresh'")
|
||||
}
|
||||
|
||||
func TestInvalidateProxyProbeSnapshotsClearsOpenCodeGoSnapshot(t *testing.T) {
|
||||
client, mock := newOllamaCloudUsageRepositoryTestClient(t)
|
||||
mock.ExpectQuery(`(?s)UPDATE accounts.*opencode_go_usage_snapshot.*RETURNING id`).
|
||||
WithArgs(int64(9)).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(int64(17)))
|
||||
|
||||
ids, err := invalidateProxyProbeSnapshots(context.Background(), client, 9)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []int64{17}, ids)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
// UpdateCredentials 的 CASE 求值顺序是正确性依赖而非防御:Ollama 分支的 WHEN 是
|
||||
// 宽守卫(NOT(ollamaMatch(old) AND ollamaMatch(new)),旧行不匹配 ollama.com 基址
|
||||
// 时恒真,且不含 opencode.ai 正则),而两侧平台白名单完全相同,因此挂载行
|
||||
// (白名单平台 + 官方 OpenCode Go 基址)会同时满足两分支的 WHEN。若把 Ollama
|
||||
// 分支前移,挂载行的 api_key 变化会先命中 Ollama 分支,OpenCode THEN 才清除的
|
||||
// opencode_go_usage_snapshot / opencode_go_usage_auto_refresh 残留,陈旧快照跟着
|
||||
// 新 api_key 走,跨 key 组污染。本测试用正则钉死 OpenCode 分支的 WHEN 标记
|
||||
// (platform = 'opencode_go',只出现在 OpenCode 分支)文本上先于 Ollama 分支的
|
||||
// WHEN 标记(ollama.com 基址正则片段)。
|
||||
func TestUpdateCredentialsOpenCodeBranchPrecedesOllamaBranch(t *testing.T) {
|
||||
client, mock := newOllamaCloudUsageRepositoryTestClient(t)
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(`(?s)UPDATE accounts.*platform = 'opencode_go'.*\[oO\]\[lL\]\[lL\]\[aA\]\[mM\]\[aA\]`).
|
||||
WithArgs(`{"api_key":"new-key","base_url":"https://opencode.ai/zen/go/v1"}`, int64(17)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).
|
||||
WithArgs(service.SchedulerOutboxEventAccountChanged, int64(17), nil, nil, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
repo := newAccountRepositoryWithSQL(client, nil, nil)
|
||||
|
||||
err := repo.UpdateCredentials(context.Background(), 17, map[string]any{
|
||||
"api_key": "new-key", "base_url": "https://opencode.ai/zen/go/v1",
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
// SQL 侧「默认 Go」语义的常量级护栏:opencode_go 平台行 account_mode 缺失/为
|
||||
// JSON null 时,COALESCE 兜底必须落 true(视为 Go 订阅,与 GetOpenCodeAccountMode
|
||||
// 的默认兼容逻辑一致)。若误改成 false,存量 opencode_go 账号在 SQL 侧集体失去
|
||||
// 资格(组查询漏行、身份清理漏清、RunDue 自动刷新停摆),而 Go 侧仍判合格,
|
||||
// 两侧不一致;唯一的行为级覆盖在需要 Docker 的 integration 测试,无 Docker 的
|
||||
// unit 通道此前完全拦不住,这里以文本形态钉死该 COALESCE 表达式。
|
||||
func TestOpenCodeGoUsageEligibleSQLDefaultsMissingAccountModeToGo(t *testing.T) {
|
||||
require.Contains(t, opencodeGoUsageEligibleSQL,
|
||||
"COALESCE(btrim(credentials ->> 'account_mode') <> 'zen', true)")
|
||||
}
|
||||
|
||||
// Go 侧「默认 Go」契约的表驱动钉死:account_mode 只有 trim 后恰好等于 "zen" 才是
|
||||
// Zen,其余取值一律判 Go。SQL 侧(opencodeGoUsageEligibleSQL)以
|
||||
// COALESCE(btrim(...) <> 'zen', true) 声明同一语义:btrim 只去空格,与本表全部
|
||||
// 取值在 Go 侧 strings.TrimSpace 下的结论一致(空白字符集差异见常量注释)。
|
||||
func TestOpenCodeAccountModeOnlyTrimmedZenIsZen(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
accountMode func(map[string]any)
|
||||
wantMode string
|
||||
wantGoPlan bool
|
||||
}{
|
||||
{name: "键缺失", accountMode: func(map[string]any) {}, wantMode: service.AccountModeGo, wantGoPlan: true},
|
||||
{name: "JSON null", accountMode: func(c map[string]any) { c["account_mode"] = nil }, wantMode: service.AccountModeGo, wantGoPlan: true},
|
||||
{name: "空串", accountMode: func(c map[string]any) { c["account_mode"] = "" }, wantMode: service.AccountModeGo, wantGoPlan: true},
|
||||
{name: "go", accountMode: func(c map[string]any) { c["account_mode"] = service.AccountModeGo }, wantMode: service.AccountModeGo, wantGoPlan: true},
|
||||
{name: "zen", accountMode: func(c map[string]any) { c["account_mode"] = service.AccountModeZen }, wantMode: service.AccountModeZen, wantGoPlan: false},
|
||||
{name: "前后空格的 zen", accountMode: func(c map[string]any) { c["account_mode"] = " zen " }, wantMode: service.AccountModeZen, wantGoPlan: false},
|
||||
{name: "大小写不同的 ZEN", accountMode: func(c map[string]any) { c["account_mode"] = "ZEN" }, wantMode: service.AccountModeGo, wantGoPlan: true},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := &service.Account{
|
||||
ID: 17,
|
||||
Platform: service.PlatformOpenCodeGo,
|
||||
Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "key"},
|
||||
}
|
||||
tt.accountMode(account.Credentials)
|
||||
require.Equal(t, tt.wantMode, account.GetOpenCodeAccountMode())
|
||||
require.Equal(t, tt.wantGoPlan, account.IsOpenCodeGoPlan())
|
||||
require.Equal(t, !tt.wantGoPlan, account.IsOpenCodeZen())
|
||||
// opencode_go 平台行的用量资格即 Go 订阅判定,与 SQL 侧
|
||||
// opencodeGoUsageEligibleSQL 的「默认 Go、仅 trim 后等于 zen 排除」
|
||||
// 互为镜像。
|
||||
require.Equal(t, tt.wantGoPlan, service.IsOpenCodeGoUsageAccount(account))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -81,8 +81,8 @@ func TestLockAndMergeAccountProbeExtraUsesCurrentDatabaseSnapshot(t *testing.T)
|
||||
|
||||
mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")).
|
||||
WithArgs(int64(27), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test"}`, nil).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}).
|
||||
AddRow(tt.identityUnchanged, false, true, tt.databaseEnabled, nil, tt.databaseSnapshot, nil, nil, nil))
|
||||
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot", "opencode_group_unchanged", "opencode_auto", "opencode_snapshot"}).
|
||||
AddRow(tt.identityUnchanged, false, true, tt.databaseEnabled, nil, tt.databaseSnapshot, nil, nil, nil, false, nil, nil))
|
||||
|
||||
account := &service.Account{
|
||||
ID: 27,
|
||||
@@ -172,8 +172,8 @@ func TestLockAndMergeAccountProbeExtraNeverInfersProbeFromRateSync(t *testing.T)
|
||||
|
||||
mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")).
|
||||
WithArgs(int64(31), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test"}`, nil).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}).
|
||||
AddRow(true, false, true, tt.databaseEnabled, tt.databaseRateSync, nil, nil, nil, nil))
|
||||
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot", "opencode_group_unchanged", "opencode_auto", "opencode_snapshot"}).
|
||||
AddRow(true, false, true, tt.databaseEnabled, tt.databaseRateSync, nil, nil, nil, nil, false, nil, nil))
|
||||
|
||||
account := &service.Account{
|
||||
ID: 31,
|
||||
@@ -211,8 +211,8 @@ func TestLockAndMergeAccountProbeExtraProtectsOllamaManagedFields(t *testing.T)
|
||||
|
||||
mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")).
|
||||
WithArgs(int64(29), service.PlatformAnthropic, service.AccountTypeAPIKey, `{"api_key":"key","base_url":"https://ollama.com"}`, nil).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}).
|
||||
AddRow(identityUnchanged, identityUnchanged, true, nil, nil, nil, []byte(`"local-ciphertext"`), []byte(`true`), []byte(`{"status":"ok"}`)))
|
||||
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot", "opencode_group_unchanged", "opencode_auto", "opencode_snapshot"}).
|
||||
AddRow(identityUnchanged, identityUnchanged, true, nil, nil, nil, []byte(`"local-ciphertext"`), []byte(`true`), []byte(`{"status":"ok"}`), false, nil, nil))
|
||||
|
||||
account := &service.Account{
|
||||
ID: 29, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey,
|
||||
@@ -371,8 +371,8 @@ func TestUpdateWithAccountBillingSettingsRollsBackWhenOutboxFails(t *testing.T)
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")).
|
||||
WithArgs(int64(27), service.PlatformOpenAI, service.AccountTypeAPIKey, `{"api_key":"sk-test"}`, nil).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot"}).
|
||||
AddRow(true, false, true, []byte(`true`), []byte(`true`), []byte(`{"status":"ok"}`), nil, nil, nil))
|
||||
WillReturnRows(sqlmock.NewRows([]string{"identity_unchanged", "ollama_group_unchanged", "ollama_proxy_unchanged", "enabled", "rate_sync_enabled", "snapshot", "ollama_session", "ollama_auto", "ollama_snapshot", "opencode_group_unchanged", "opencode_auto", "opencode_snapshot"}).
|
||||
AddRow(true, false, true, []byte(`true`), []byte(`true`), []byte(`{"status":"ok"}`), nil, nil, nil, false, nil, nil))
|
||||
mock.ExpectExec(`(?s)UPDATE .*accounts.*SET.*WHERE .*id.*`).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectQuery(`(?s)SELECT .* FROM "accounts" WHERE "id" = \$1`).
|
||||
|
||||
@@ -341,6 +341,171 @@ VALUES ($1, 'transfer', $2, NULL, $3, $4, $5, $6, NOW(), NOW())`,
|
||||
return transferred, newBalance, nil
|
||||
}
|
||||
|
||||
// WithdrawQuota 登记一笔线下提现。operationID 标识一次登记:同一事务内先以
|
||||
// operation_id 唯一约束写入占位流水,同标识的流水已存在时不重复扣减,用户与金额
|
||||
// 一致则返回该流水记录的结果(Replayed=true),不一致返回
|
||||
// ErrIdempotencyKeyConflict。占位成功后先解冻已到期的冻结额度,再以
|
||||
// aff_quota >= amount 为条件原子扣减并补写额度快照。可提取额度不足或用户没有
|
||||
// 返利档案时返回 ErrAffiliateQuotaInsufficient,占位随事务回滚,不占用该标识。
|
||||
func (r *affiliateRepository) WithdrawQuota(ctx context.Context, userID int64, amount float64, operationID string) (*service.AffiliateWithdrawResult, error) {
|
||||
if userID <= 0 {
|
||||
return nil, service.ErrUserNotFound
|
||||
}
|
||||
if amount <= 0 {
|
||||
return nil, service.ErrAffiliateWithdrawAmountInvalid
|
||||
}
|
||||
if operationID == "" {
|
||||
return nil, service.ErrIdempotencyKeyRequired
|
||||
}
|
||||
|
||||
var result *service.AffiliateWithdrawResult
|
||||
err := r.withTx(ctx, func(txCtx context.Context, txClient *dbent.Client) error {
|
||||
ledgerID, claimed, err := claimAffiliateWithdrawLedger(txCtx, txClient, userID, amount, operationID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !claimed {
|
||||
result, err = findAffiliateWithdrawByOperation(txCtx, txClient, userID, amount, operationID)
|
||||
return err
|
||||
}
|
||||
|
||||
if _, err := thawFrozenQuotaTx(txCtx, txClient, userID); err != nil {
|
||||
return fmt.Errorf("thaw before withdraw: %w", err)
|
||||
}
|
||||
|
||||
res, err := txClient.ExecContext(txCtx, `
|
||||
UPDATE user_affiliates
|
||||
SET aff_quota = aff_quota - $1,
|
||||
updated_at = NOW()
|
||||
WHERE user_id = $2
|
||||
AND aff_quota >= $1`, amount, userID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("deduct affiliate quota: %w", err)
|
||||
}
|
||||
affected, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return fmt.Errorf("deduct affiliate quota: %w", err)
|
||||
}
|
||||
if affected == 0 {
|
||||
return service.ErrAffiliateQuotaInsufficient
|
||||
}
|
||||
|
||||
snapshot, err := queryAffiliateTransferSnapshot(txCtx, txClient, userID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if _, err := txClient.ExecContext(txCtx, `
|
||||
UPDATE user_affiliate_ledger
|
||||
SET balance_after = $1,
|
||||
aff_quota_after = $2,
|
||||
aff_frozen_quota_after = $3,
|
||||
aff_history_quota_after = $4,
|
||||
updated_at = NOW()
|
||||
WHERE id = $5`,
|
||||
snapshot.BalanceAfter,
|
||||
snapshot.AvailableQuotaAfter,
|
||||
snapshot.FrozenQuotaAfter,
|
||||
snapshot.HistoryQuotaAfter,
|
||||
ledgerID,
|
||||
); err != nil {
|
||||
return fmt.Errorf("record affiliate withdraw snapshot: %w", err)
|
||||
}
|
||||
|
||||
result = &service.AffiliateWithdrawResult{
|
||||
LedgerID: ledgerID,
|
||||
UserID: userID,
|
||||
Amount: amount,
|
||||
AvailableQuotaAfter: snapshot.AvailableQuotaAfter,
|
||||
FrozenQuotaAfter: snapshot.FrozenQuotaAfter,
|
||||
HistoryQuotaAfter: snapshot.HistoryQuotaAfter,
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// claimAffiliateWithdrawLedger 为有返利档案的用户写入一条带 operation_id 的
|
||||
// withdraw 占位流水。同标识的流水已存在或用户没有返利档案时 claimed 为 false。
|
||||
// 同标识的并发事务在唯一索引上等待先写入者结束:先写入者提交后返回 claimed=false,
|
||||
// 回滚后本事务继续写入。
|
||||
func claimAffiliateWithdrawLedger(ctx context.Context, client affiliateQueryExecer, userID int64, amount float64, operationID string) (int64, bool, error) {
|
||||
rows, err := client.QueryContext(ctx, `
|
||||
INSERT INTO user_affiliate_ledger (user_id, action, amount, operation_id, created_at, updated_at)
|
||||
SELECT ua.user_id, 'withdraw', $2, $3, NOW(), NOW()
|
||||
FROM user_affiliates ua
|
||||
WHERE ua.user_id = $1
|
||||
ON CONFLICT (operation_id) WHERE operation_id IS NOT NULL DO NOTHING
|
||||
RETURNING id`, userID, amount, operationID)
|
||||
if err != nil {
|
||||
return 0, false, fmt.Errorf("claim affiliate withdraw ledger: %w", err)
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
|
||||
if !rows.Next() {
|
||||
if err := rows.Err(); err != nil {
|
||||
return 0, false, fmt.Errorf("claim affiliate withdraw ledger: %w", err)
|
||||
}
|
||||
return 0, false, nil
|
||||
}
|
||||
var ledgerID int64
|
||||
if err := rows.Scan(&ledgerID); err != nil {
|
||||
return 0, false, err
|
||||
}
|
||||
return ledgerID, true, rows.Close()
|
||||
}
|
||||
|
||||
// findAffiliateWithdrawByOperation 读取 operation_id 已对应的线下提现流水,
|
||||
// 用户与金额一致时返回该流水记录的登记结果(Replayed=true),不一致返回
|
||||
// ErrIdempotencyKeyConflict;没有对应流水说明用户没有返利档案,返回
|
||||
// ErrAffiliateQuotaInsufficient。
|
||||
func findAffiliateWithdrawByOperation(ctx context.Context, client affiliateQueryExecer, userID int64, amount float64, operationID string) (*service.AffiliateWithdrawResult, error) {
|
||||
rows, err := client.QueryContext(ctx, `
|
||||
SELECT id,
|
||||
user_id = $2 AND action = 'withdraw' AND amount = CAST($3 AS DECIMAL(20,8)),
|
||||
user_id,
|
||||
amount::double precision,
|
||||
COALESCE(aff_quota_after, 0)::double precision,
|
||||
COALESCE(aff_frozen_quota_after, 0)::double precision,
|
||||
COALESCE(aff_history_quota_after, 0)::double precision
|
||||
FROM user_affiliate_ledger
|
||||
WHERE operation_id = $1`, operationID, userID, amount)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query affiliate withdraw by operation: %w", err)
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
|
||||
if !rows.Next() {
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("query affiliate withdraw by operation: %w", err)
|
||||
}
|
||||
return nil, service.ErrAffiliateQuotaInsufficient
|
||||
}
|
||||
var sameRequest bool
|
||||
result := &service.AffiliateWithdrawResult{Replayed: true}
|
||||
if err := rows.Scan(
|
||||
&result.LedgerID,
|
||||
&sameRequest,
|
||||
&result.UserID,
|
||||
&result.Amount,
|
||||
&result.AvailableQuotaAfter,
|
||||
&result.FrozenQuotaAfter,
|
||||
&result.HistoryQuotaAfter,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !sameRequest {
|
||||
return nil, service.ErrIdempotencyKeyConflict
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (r *affiliateRepository) ListInvitees(ctx context.Context, inviterID int64, limit int) ([]service.AffiliateInvitee, error) {
|
||||
if limit <= 0 {
|
||||
limit = 100
|
||||
@@ -466,13 +631,14 @@ func (r *affiliateRepository) ListAffiliateRebateRecords(ctx context.Context, fi
|
||||
"inviter.email", "inviter.username", "invitee.email", "invitee.username",
|
||||
"po.id::text", "po.out_trade_no", "po.payment_type", "po.status",
|
||||
})
|
||||
// 返利记录列出全部 accrue 流水:非订单来源的流水没有订单,被邀请人账号
|
||||
// 删除后 source_user_id 为 NULL,这两类行都保留。
|
||||
baseJoin := `
|
||||
FROM user_affiliate_ledger ual
|
||||
JOIN payment_orders po ON po.id = ual.source_order_id
|
||||
JOIN users invitee ON invitee.id = ual.source_user_id
|
||||
JOIN users inviter ON inviter.id = ual.user_id
|
||||
WHERE ual.action = 'accrue'
|
||||
AND ual.source_order_id IS NOT NULL`
|
||||
LEFT JOIN users invitee ON invitee.id = ual.source_user_id
|
||||
LEFT JOIN payment_orders po ON po.id = ual.source_order_id
|
||||
WHERE ual.action = 'accrue'`
|
||||
if where != "" {
|
||||
where = strings.Replace(where, "WHERE ", " AND ", 1)
|
||||
}
|
||||
@@ -496,7 +662,7 @@ WHERE ual.action = 'accrue'
|
||||
args = append(args, filter.PageSize, (filter.Page-1)*filter.PageSize)
|
||||
rows, err := client.QueryContext(ctx, `
|
||||
SELECT po.id,
|
||||
po.out_trade_no,
|
||||
COALESCE(po.out_trade_no, ''),
|
||||
ual.user_id,
|
||||
COALESCE(inviter.email, ''),
|
||||
COALESCE(inviter.username, ''),
|
||||
@@ -506,8 +672,8 @@ SELECT po.id,
|
||||
po.amount::double precision,
|
||||
po.pay_amount::double precision,
|
||||
ual.amount::double precision,
|
||||
po.payment_type,
|
||||
po.status,
|
||||
COALESCE(po.payment_type, ''),
|
||||
COALESCE(po.status, ''),
|
||||
ual.created_at
|
||||
`+baseJoin+where+`
|
||||
`+orderBy+`
|
||||
@@ -520,17 +686,21 @@ LIMIT $`+fmt.Sprint(len(args)-1)+` OFFSET $`+fmt.Sprint(len(args)), args...)
|
||||
items := make([]service.AffiliateRebateRecord, 0)
|
||||
for rows.Next() {
|
||||
var item service.AffiliateRebateRecord
|
||||
var orderID sql.NullInt64
|
||||
var inviteeID sql.NullInt64
|
||||
var orderAmount sql.NullFloat64
|
||||
var payAmount sql.NullFloat64
|
||||
if err := rows.Scan(
|
||||
&item.OrderID,
|
||||
&orderID,
|
||||
&item.OutTradeNo,
|
||||
&item.InviterID,
|
||||
&item.InviterEmail,
|
||||
&item.InviterUsername,
|
||||
&item.InviteeID,
|
||||
&inviteeID,
|
||||
&item.InviteeEmail,
|
||||
&item.InviteeUsername,
|
||||
&item.OrderAmount,
|
||||
&item.PayAmount,
|
||||
&orderAmount,
|
||||
&payAmount,
|
||||
&item.RebateAmount,
|
||||
&item.PaymentType,
|
||||
&item.OrderStatus,
|
||||
@@ -538,6 +708,10 @@ LIMIT $`+fmt.Sprint(len(args)-1)+` OFFSET $`+fmt.Sprint(len(args)), args...)
|
||||
); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
item.OrderID = nullableInt64Ptr(orderID)
|
||||
item.InviteeID = nullableInt64Ptr(inviteeID)
|
||||
item.OrderAmount = nullableFloat64Ptr(orderAmount)
|
||||
item.PayAmount = nullableFloat64Ptr(payAmount)
|
||||
items = append(items, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
@@ -554,7 +728,7 @@ func (r *affiliateRepository) ListAffiliateTransferRecords(ctx context.Context,
|
||||
baseJoin := `
|
||||
FROM user_affiliate_ledger ual
|
||||
JOIN users u ON u.id = ual.user_id
|
||||
WHERE ual.action = 'transfer'`
|
||||
WHERE ual.action IN ('transfer', 'withdraw')`
|
||||
if where != "" {
|
||||
where = strings.Replace(where, "WHERE ", " AND ", 1)
|
||||
}
|
||||
@@ -566,6 +740,7 @@ WHERE ual.action = 'transfer'`
|
||||
|
||||
orderBy := buildAffiliateRecordOrderBy(filter, map[string]string{
|
||||
"user": "u.email",
|
||||
"action": "ual.action",
|
||||
"amount": "ual.amount",
|
||||
"balance_after": "ual.balance_after",
|
||||
"available_quota_after": "ual.aff_quota_after",
|
||||
@@ -576,6 +751,7 @@ WHERE ual.action = 'transfer'`
|
||||
args = append(args, filter.PageSize, (filter.Page-1)*filter.PageSize)
|
||||
rows, err := client.QueryContext(ctx, `
|
||||
SELECT ual.id,
|
||||
ual.action,
|
||||
ual.user_id,
|
||||
COALESCE(u.email, ''),
|
||||
COALESCE(u.username, ''),
|
||||
@@ -602,6 +778,7 @@ LIMIT $`+fmt.Sprint(len(args)-1)+` OFFSET $`+fmt.Sprint(len(args)), args...)
|
||||
var historyQuotaAfter sql.NullFloat64
|
||||
if err := rows.Scan(
|
||||
&item.LedgerID,
|
||||
&item.Action,
|
||||
&item.UserID,
|
||||
&item.UserEmail,
|
||||
&item.Username,
|
||||
@@ -958,6 +1135,13 @@ func nullableFloat64Ptr(v sql.NullFloat64) *float64 {
|
||||
return &v.Float64
|
||||
}
|
||||
|
||||
func nullableInt64Ptr(v sql.NullInt64) *int64 {
|
||||
if !v.Valid {
|
||||
return nil
|
||||
}
|
||||
return &v.Int64
|
||||
}
|
||||
|
||||
func generateAffiliateCode() (string, error) {
|
||||
buf := make([]byte, affiliateCodeLength)
|
||||
if _, err := rand.Read(buf); err != nil {
|
||||
|
||||
@@ -4,7 +4,10 @@ package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -417,3 +420,383 @@ func TestAffiliateRepository_ListUsersWithCustomSettings(t *testing.T) {
|
||||
|
||||
require.GreaterOrEqual(t, total, int64(2), "total must include at least our 2 custom rows")
|
||||
}
|
||||
|
||||
func mustCreateAffiliateWithdrawUser(t *testing.T, client *dbent.Client, label string) *service.User {
|
||||
t.Helper()
|
||||
return mustCreateUser(t, client, &service.User{
|
||||
Email: fmt.Sprintf("affiliate-withdraw-%s-%d@example.com", label, time.Now().UnixNano()),
|
||||
PasswordHash: "hash",
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 5.5,
|
||||
Concurrency: 5,
|
||||
})
|
||||
}
|
||||
|
||||
func mustSeedAffiliateQuota(t *testing.T, ctx context.Context, client *dbent.Client, userID int64, quota, frozen, history float64) {
|
||||
t.Helper()
|
||||
affCode := fmt.Sprintf("AFW%d-%06d", userID, time.Now().UnixNano()%1_000_000)
|
||||
_, err := client.ExecContext(ctx, `
|
||||
INSERT INTO user_affiliates (user_id, aff_code, aff_quota, aff_frozen_quota, aff_history_quota, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, $5, NOW(), NOW())`, userID, affCode, quota, frozen, history)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// mustCreateCommittedAffiliateWithdrawUser 在测试事务之外建用户与返利档案,
|
||||
// 测试结束时删除(级联返利档案与流水)。WithdrawQuota 在这些用例里自行开启
|
||||
// 并提交事务,与生产调用一致。
|
||||
func mustCreateCommittedAffiliateWithdrawUser(t *testing.T, label string, quota, frozen, history float64) *service.User {
|
||||
t.Helper()
|
||||
client := testEntClient(t)
|
||||
u := mustCreateAffiliateWithdrawUser(t, client, label)
|
||||
t.Cleanup(func() {
|
||||
_, _ = integrationDB.ExecContext(context.Background(), "DELETE FROM users WHERE id = $1", u.ID)
|
||||
})
|
||||
mustSeedAffiliateQuota(t, context.Background(), client, u.ID, quota, frozen, history)
|
||||
return u
|
||||
}
|
||||
|
||||
func newAffiliateWithdrawOperationID(label string) string {
|
||||
sum := sha256.Sum256([]byte(fmt.Sprintf("%s-%d", label, time.Now().UnixNano())))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
// TestAffiliateRepository_WithdrawQuota_DeductsAndRecordsLedger 覆盖线下提现主路径:
|
||||
// 只扣可提取额度,余额与累计返利不变,流水 action=withdraw 带额度快照,并出现在提取记录中。
|
||||
func TestAffiliateRepository_WithdrawQuota_DeductsAndRecordsLedger(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
txCtx := dbent.NewTxContext(ctx, tx)
|
||||
client := tx.Client()
|
||||
repo := NewAffiliateRepository(client, integrationDB)
|
||||
|
||||
u := mustCreateAffiliateWithdrawUser(t, client, "user")
|
||||
mustSeedAffiliateQuota(t, txCtx, client, u.ID, 20, 0, 30)
|
||||
operationID := newAffiliateWithdrawOperationID("deduct")
|
||||
|
||||
result, err := repo.WithdrawQuota(txCtx, u.ID, 12.5, operationID)
|
||||
require.NoError(t, err)
|
||||
require.Positive(t, result.LedgerID)
|
||||
require.False(t, result.Replayed)
|
||||
require.InDelta(t, 12.5, result.Amount, 1e-9)
|
||||
require.InDelta(t, 7.5, result.AvailableQuotaAfter, 1e-9)
|
||||
require.InDelta(t, 0.0, result.FrozenQuotaAfter, 1e-9)
|
||||
require.InDelta(t, 30.0, result.HistoryQuotaAfter, 1e-9)
|
||||
|
||||
require.InDelta(t, 7.5, querySingleFloat(t, txCtx, client,
|
||||
"SELECT aff_quota::double precision FROM user_affiliates WHERE user_id = $1", u.ID), 1e-9)
|
||||
require.InDelta(t, 30.0, querySingleFloat(t, txCtx, client,
|
||||
"SELECT aff_history_quota::double precision FROM user_affiliates WHERE user_id = $1", u.ID), 1e-9)
|
||||
require.InDelta(t, 5.5, querySingleFloat(t, txCtx, client,
|
||||
"SELECT balance::double precision FROM users WHERE id = $1", u.ID), 1e-9)
|
||||
|
||||
records, total, err := repo.ListAffiliateTransferRecords(txCtx, service.AffiliateRecordFilter{
|
||||
Search: u.Email,
|
||||
Page: 1,
|
||||
PageSize: 20,
|
||||
SortDesc: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 1, total)
|
||||
require.Len(t, records, 1)
|
||||
record := records[0]
|
||||
require.Equal(t, result.LedgerID, record.LedgerID)
|
||||
require.Equal(t, "withdraw", record.Action)
|
||||
require.Equal(t, u.ID, record.UserID)
|
||||
require.InDelta(t, 12.5, record.Amount, 1e-9)
|
||||
require.True(t, record.SnapshotAvailable)
|
||||
require.InDelta(t, 5.5, *record.BalanceAfter, 1e-9)
|
||||
require.InDelta(t, 7.5, *record.AvailableQuotaAfter, 1e-9)
|
||||
require.InDelta(t, 30.0, *record.HistoryQuotaAfter, 1e-9)
|
||||
|
||||
rows, err := client.QueryContext(txCtx, "SELECT operation_id FROM user_affiliate_ledger WHERE id = $1", result.LedgerID)
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = rows.Close() }()
|
||||
require.True(t, rows.Next())
|
||||
var storedOperationID string
|
||||
require.NoError(t, rows.Scan(&storedOperationID))
|
||||
require.Equal(t, operationID, storedOperationID)
|
||||
}
|
||||
|
||||
// TestAffiliateRepository_WithdrawQuota_ThawsMaturedQuotaFirst 验证已过冻结期
|
||||
// 但尚未解冻的额度在扣减前解冻并可被提现。
|
||||
func TestAffiliateRepository_WithdrawQuota_ThawsMaturedQuotaFirst(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
txCtx := dbent.NewTxContext(ctx, tx)
|
||||
client := tx.Client()
|
||||
repo := NewAffiliateRepository(client, integrationDB)
|
||||
|
||||
u := mustCreateAffiliateWithdrawUser(t, client, "thaw")
|
||||
mustSeedAffiliateQuota(t, txCtx, client, u.ID, 0, 10, 10)
|
||||
_, err := client.ExecContext(txCtx, `
|
||||
INSERT INTO user_affiliate_ledger (user_id, action, amount, frozen_until, created_at, updated_at)
|
||||
VALUES ($1, 'accrue', 10, NOW() - INTERVAL '1 hour', NOW(), NOW())`, u.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
result, err := repo.WithdrawQuota(txCtx, u.ID, 10, newAffiliateWithdrawOperationID("thaw"))
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 0.0, result.AvailableQuotaAfter, 1e-9)
|
||||
require.InDelta(t, 0.0, result.FrozenQuotaAfter, 1e-9)
|
||||
require.Equal(t, 1, querySingleInt(t, txCtx, client,
|
||||
"SELECT COUNT(*) FROM user_affiliate_ledger WHERE user_id = $1 AND action = 'withdraw'", u.ID))
|
||||
}
|
||||
|
||||
// TestAffiliateRepository_WithdrawQuota_InsufficientLeavesQuotaUntouched 验证未到期的
|
||||
// 冻结额度不可提现:额度不足时不扣减、不写流水,占位随事务回滚、不占用该
|
||||
// operation_id,同一标识改成可提取范围内的金额仍可登记;没有返利档案的用户
|
||||
// 同样按额度不足处理。
|
||||
func TestAffiliateRepository_WithdrawQuota_InsufficientLeavesQuotaUntouched(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
client := testEntClient(t)
|
||||
repo := NewAffiliateRepository(client, integrationDB)
|
||||
|
||||
u := mustCreateCommittedAffiliateWithdrawUser(t, "insufficient", 5, 10, 15)
|
||||
_, err := integrationDB.ExecContext(ctx, `
|
||||
INSERT INTO user_affiliate_ledger (user_id, action, amount, frozen_until, created_at, updated_at)
|
||||
VALUES ($1, 'accrue', 10, NOW() + INTERVAL '1 hour', NOW(), NOW())`, u.ID)
|
||||
require.NoError(t, err)
|
||||
operationID := newAffiliateWithdrawOperationID("insufficient")
|
||||
|
||||
_, err = repo.WithdrawQuota(ctx, u.ID, 6, operationID)
|
||||
require.ErrorIs(t, err, service.ErrAffiliateQuotaInsufficient)
|
||||
|
||||
require.InDelta(t, 5.0, querySingleFloat(t, ctx, client,
|
||||
"SELECT aff_quota::double precision FROM user_affiliates WHERE user_id = $1", u.ID), 1e-9)
|
||||
require.InDelta(t, 10.0, querySingleFloat(t, ctx, client,
|
||||
"SELECT aff_frozen_quota::double precision FROM user_affiliates WHERE user_id = $1", u.ID), 1e-9)
|
||||
require.Equal(t, 0, querySingleInt(t, ctx, client,
|
||||
"SELECT COUNT(*) FROM user_affiliate_ledger WHERE user_id = $1 AND action = 'withdraw'", u.ID))
|
||||
|
||||
result, err := repo.WithdrawQuota(ctx, u.ID, 5, operationID)
|
||||
require.NoError(t, err)
|
||||
require.False(t, result.Replayed)
|
||||
require.InDelta(t, 0.0, result.AvailableQuotaAfter, 1e-9)
|
||||
|
||||
noProfile := mustCreateAffiliateWithdrawUser(t, client, "no-profile")
|
||||
t.Cleanup(func() {
|
||||
_, _ = integrationDB.ExecContext(context.Background(), "DELETE FROM users WHERE id = $1", noProfile.ID)
|
||||
})
|
||||
_, err = repo.WithdrawQuota(ctx, noProfile.ID, 1, newAffiliateWithdrawOperationID("no-profile"))
|
||||
require.ErrorIs(t, err, service.ErrAffiliateQuotaInsufficient)
|
||||
require.Equal(t, 0, querySingleInt(t, ctx, client,
|
||||
"SELECT COUNT(*) FROM user_affiliate_ledger WHERE user_id = $1", noProfile.ID))
|
||||
}
|
||||
|
||||
// TestAffiliateRepository_WithdrawQuota_RetryAfterCommitReplaysFirstResult 复现
|
||||
// 「提交成功但响应丢失后重试」:首次登记已在独立事务中提交,以同一 operation_id
|
||||
// 重试不重复扣减、不重复写流水,返回首次登记的流水与额度快照。两笔合法的等额打款
|
||||
// 使用不同的 operation_id,各扣一次。
|
||||
func TestAffiliateRepository_WithdrawQuota_RetryAfterCommitReplaysFirstResult(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
client := testEntClient(t)
|
||||
repo := NewAffiliateRepository(client, integrationDB)
|
||||
|
||||
u := mustCreateCommittedAffiliateWithdrawUser(t, "retry", 100, 0, 100)
|
||||
operationID := newAffiliateWithdrawOperationID("retry")
|
||||
|
||||
first, err := repo.WithdrawQuota(ctx, u.ID, 10, operationID)
|
||||
require.NoError(t, err)
|
||||
require.False(t, first.Replayed)
|
||||
|
||||
retry, err := repo.WithdrawQuota(ctx, u.ID, 10, operationID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, retry.Replayed)
|
||||
require.Equal(t, first.LedgerID, retry.LedgerID)
|
||||
require.Equal(t, u.ID, retry.UserID)
|
||||
require.InDelta(t, 10.0, retry.Amount, 1e-9)
|
||||
require.InDelta(t, 90.0, retry.AvailableQuotaAfter, 1e-9)
|
||||
require.InDelta(t, first.FrozenQuotaAfter, retry.FrozenQuotaAfter, 1e-9)
|
||||
require.InDelta(t, 100.0, retry.HistoryQuotaAfter, 1e-9)
|
||||
|
||||
require.InDelta(t, 90.0, querySingleFloat(t, ctx, client,
|
||||
"SELECT aff_quota::double precision FROM user_affiliates WHERE user_id = $1", u.ID), 1e-9)
|
||||
require.Equal(t, 1, querySingleInt(t, ctx, client,
|
||||
"SELECT COUNT(*) FROM user_affiliate_ledger WHERE user_id = $1 AND action = 'withdraw'", u.ID))
|
||||
|
||||
second, err := repo.WithdrawQuota(ctx, u.ID, 10, newAffiliateWithdrawOperationID("retry-second"))
|
||||
require.NoError(t, err)
|
||||
require.False(t, second.Replayed)
|
||||
require.NotEqual(t, first.LedgerID, second.LedgerID)
|
||||
require.InDelta(t, 80.0, querySingleFloat(t, ctx, client,
|
||||
"SELECT aff_quota::double precision FROM user_affiliates WHERE user_id = $1", u.ID), 1e-9)
|
||||
require.Equal(t, 2, querySingleInt(t, ctx, client,
|
||||
"SELECT COUNT(*) FROM user_affiliate_ledger WHERE user_id = $1 AND action = 'withdraw'", u.ID))
|
||||
}
|
||||
|
||||
// TestAffiliateRepository_WithdrawQuota_SameOperationWaitsForInFlightRegistration
|
||||
// 固定「并发同键」的交错:首个登记已占位并扣减但未提交时,同一 operation_id 的
|
||||
// 请求在唯一索引上等待。首个登记提交后,等待的请求不重复扣减,返回同一条流水;
|
||||
// 首个登记回滚后,等待的请求自行完成登记。两种结局都只扣一次。
|
||||
func TestAffiliateRepository_WithdrawQuota_SameOperationWaitsForInFlightRegistration(t *testing.T) {
|
||||
for _, commitFirst := range []bool{true, false} {
|
||||
name := "first rolls back"
|
||||
if commitFirst {
|
||||
name = "first commits"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
client := testEntClient(t)
|
||||
repo := NewAffiliateRepository(client, integrationDB)
|
||||
|
||||
u := mustCreateCommittedAffiliateWithdrawUser(t, "inflight", 100, 0, 100)
|
||||
operationID := newAffiliateWithdrawOperationID("inflight")
|
||||
|
||||
tx, err := client.Tx(ctx)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = tx.Rollback() })
|
||||
first, err := repo.WithdrawQuota(dbent.NewTxContext(ctx, tx), u.ID, 10, operationID)
|
||||
require.NoError(t, err)
|
||||
|
||||
type outcome struct {
|
||||
result *service.AffiliateWithdrawResult
|
||||
err error
|
||||
}
|
||||
waiting := make(chan outcome, 1)
|
||||
go func() {
|
||||
result, err := repo.WithdrawQuota(ctx, u.ID, 10, operationID)
|
||||
waiting <- outcome{result: result, err: err}
|
||||
}()
|
||||
|
||||
select {
|
||||
case got := <-waiting:
|
||||
t.Fatalf("same operation must wait for the in-flight registration, got result=%+v err=%v", got.result, got.err)
|
||||
case <-time.After(300 * time.Millisecond):
|
||||
}
|
||||
|
||||
if commitFirst {
|
||||
require.NoError(t, tx.Commit())
|
||||
} else {
|
||||
require.NoError(t, tx.Rollback())
|
||||
}
|
||||
|
||||
var got outcome
|
||||
select {
|
||||
case got = <-waiting:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("waiting registration did not finish after the in-flight one ended")
|
||||
}
|
||||
require.NoError(t, got.err)
|
||||
if commitFirst {
|
||||
require.True(t, got.result.Replayed)
|
||||
require.Equal(t, first.LedgerID, got.result.LedgerID)
|
||||
} else {
|
||||
require.False(t, got.result.Replayed)
|
||||
require.NotEqual(t, first.LedgerID, got.result.LedgerID)
|
||||
}
|
||||
require.InDelta(t, 90.0, got.result.AvailableQuotaAfter, 1e-9)
|
||||
|
||||
require.InDelta(t, 90.0, querySingleFloat(t, ctx, client,
|
||||
"SELECT aff_quota::double precision FROM user_affiliates WHERE user_id = $1", u.ID), 1e-9)
|
||||
require.Equal(t, 1, querySingleInt(t, ctx, client,
|
||||
"SELECT COUNT(*) FROM user_affiliate_ledger WHERE user_id = $1 AND action = 'withdraw'", u.ID))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAffiliateRepository_WithdrawQuota_ConcurrentSameOperationDeductsOnce 让多个请求
|
||||
// 同时以同一 operation_id 登记:只有一个执行扣减,其余返回同一条流水。
|
||||
func TestAffiliateRepository_WithdrawQuota_ConcurrentSameOperationDeductsOnce(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
client := testEntClient(t)
|
||||
repo := NewAffiliateRepository(client, integrationDB)
|
||||
|
||||
u := mustCreateCommittedAffiliateWithdrawUser(t, "concurrent", 100, 0, 100)
|
||||
operationID := newAffiliateWithdrawOperationID("concurrent")
|
||||
|
||||
const workers = 8
|
||||
results := make([]*service.AffiliateWithdrawResult, workers)
|
||||
errs := make([]error, workers)
|
||||
start := make(chan struct{})
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
results[i], errs[i] = repo.WithdrawQuota(ctx, u.ID, 10, operationID)
|
||||
}(i)
|
||||
}
|
||||
close(start)
|
||||
wg.Wait()
|
||||
|
||||
executed := 0
|
||||
for i := 0; i < workers; i++ {
|
||||
require.NoError(t, errs[i])
|
||||
require.Equal(t, results[0].LedgerID, results[i].LedgerID)
|
||||
require.InDelta(t, 90.0, results[i].AvailableQuotaAfter, 1e-9)
|
||||
if !results[i].Replayed {
|
||||
executed++
|
||||
}
|
||||
}
|
||||
require.Equal(t, 1, executed)
|
||||
require.InDelta(t, 90.0, querySingleFloat(t, ctx, client,
|
||||
"SELECT aff_quota::double precision FROM user_affiliates WHERE user_id = $1", u.ID), 1e-9)
|
||||
require.Equal(t, 1, querySingleInt(t, ctx, client,
|
||||
"SELECT COUNT(*) FROM user_affiliate_ledger WHERE user_id = $1 AND action = 'withdraw'", u.ID))
|
||||
}
|
||||
|
||||
// TestAffiliateRepository_WithdrawQuota_RejectsOperationReuseWithDifferentRequest 验证
|
||||
// 同一 operation_id 换了金额或用户时以 ErrIdempotencyKeyConflict 拒绝,
|
||||
// 不扣减、不写流水。
|
||||
func TestAffiliateRepository_WithdrawQuota_RejectsOperationReuseWithDifferentRequest(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
client := testEntClient(t)
|
||||
repo := NewAffiliateRepository(client, integrationDB)
|
||||
|
||||
u := mustCreateCommittedAffiliateWithdrawUser(t, "conflict", 100, 0, 100)
|
||||
other := mustCreateCommittedAffiliateWithdrawUser(t, "conflict-other", 100, 0, 100)
|
||||
operationID := newAffiliateWithdrawOperationID("conflict")
|
||||
|
||||
_, err := repo.WithdrawQuota(ctx, u.ID, 10, operationID)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = repo.WithdrawQuota(ctx, u.ID, 11, operationID)
|
||||
require.ErrorIs(t, err, service.ErrIdempotencyKeyConflict)
|
||||
_, err = repo.WithdrawQuota(ctx, other.ID, 10, operationID)
|
||||
require.ErrorIs(t, err, service.ErrIdempotencyKeyConflict)
|
||||
|
||||
require.InDelta(t, 90.0, querySingleFloat(t, ctx, client,
|
||||
"SELECT aff_quota::double precision FROM user_affiliates WHERE user_id = $1", u.ID), 1e-9)
|
||||
require.InDelta(t, 100.0, querySingleFloat(t, ctx, client,
|
||||
"SELECT aff_quota::double precision FROM user_affiliates WHERE user_id = $1", other.ID), 1e-9)
|
||||
require.Equal(t, 1, querySingleInt(t, ctx, client,
|
||||
"SELECT COUNT(*) FROM user_affiliate_ledger WHERE user_id IN ($1, $2) AND action = 'withdraw'", u.ID, other.ID))
|
||||
}
|
||||
|
||||
// TestAffiliateRepository_ListAffiliateRebateRecords_IncludesNonOrderAccruals 验证兑换码
|
||||
// 等非订单来源的返利出现在返利记录中,订单字段为空。
|
||||
func TestAffiliateRepository_ListAffiliateRebateRecords_IncludesNonOrderAccruals(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
txCtx := dbent.NewTxContext(ctx, tx)
|
||||
client := tx.Client()
|
||||
repo := NewAffiliateRepository(client, integrationDB)
|
||||
|
||||
inviter := mustCreateAffiliateWithdrawUser(t, client, "inviter")
|
||||
invitee := mustCreateAffiliateWithdrawUser(t, client, "invitee")
|
||||
mustSeedAffiliateQuota(t, txCtx, client, inviter.ID, 0, 0, 0)
|
||||
|
||||
applied, err := repo.AccrueQuota(txCtx, inviter.ID, invitee.ID, 2, 0, nil)
|
||||
require.NoError(t, err)
|
||||
require.True(t, applied)
|
||||
|
||||
records, total, err := repo.ListAffiliateRebateRecords(txCtx, service.AffiliateRecordFilter{
|
||||
Search: invitee.Email,
|
||||
Page: 1,
|
||||
PageSize: 20,
|
||||
SortDesc: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 1, total)
|
||||
require.Len(t, records, 1)
|
||||
record := records[0]
|
||||
require.Nil(t, record.OrderID)
|
||||
require.Nil(t, record.OrderAmount)
|
||||
require.Nil(t, record.PayAmount)
|
||||
require.Empty(t, record.OutTradeNo)
|
||||
require.Equal(t, inviter.ID, record.InviterID)
|
||||
require.NotNil(t, record.InviteeID)
|
||||
require.Equal(t, invitee.ID, *record.InviteeID)
|
||||
require.InDelta(t, 2.0, record.RebateAmount, 1e-9)
|
||||
}
|
||||
|
||||
@@ -20,9 +20,56 @@ func TestAffiliateRecordQueriesUseLedgerAuditFields(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
content := string(source)
|
||||
|
||||
require.Contains(t, content, "JOIN payment_orders po ON po.id = ual.source_order_id")
|
||||
require.Contains(t, content, "LEFT JOIN payment_orders po ON po.id = ual.source_order_id")
|
||||
require.Contains(t, content, "ual.amount::double precision")
|
||||
require.Contains(t, content, "ual.balance_after::double precision")
|
||||
require.NotContains(t, content, "parseAffiliateRebateAmount")
|
||||
require.NotContains(t, content, `"current_balance": "u.balance"`)
|
||||
}
|
||||
|
||||
// TestAffiliateRebateRecordsQueryKeepsNonOrderAccruals 锁定返利记录列出全部
|
||||
// accrue 流水:订单与被邀请人均为 LEFT JOIN,且不按 source_order_id 过滤,
|
||||
// 兑换码、管理员充值来源的返利才会出现在明细里。
|
||||
func TestAffiliateRebateRecordsQueryKeepsNonOrderAccruals(t *testing.T) {
|
||||
source, err := os.ReadFile("affiliate_repo.go")
|
||||
require.NoError(t, err)
|
||||
content := string(source)
|
||||
|
||||
require.Contains(t, content, "LEFT JOIN payment_orders po ON po.id = ual.source_order_id")
|
||||
require.Contains(t, content, "LEFT JOIN users invitee ON invitee.id = ual.source_user_id")
|
||||
require.NotContains(t, content, "\nJOIN payment_orders po ON po.id = ual.source_order_id")
|
||||
require.NotContains(t, content, "\nJOIN users invitee ON invitee.id = ual.source_user_id")
|
||||
require.NotContains(t, content, "AND ual.source_order_id IS NOT NULL")
|
||||
}
|
||||
|
||||
// TestAffiliateTransferRecordsQueryIncludesOfflineWithdrawals 锁定提取记录同时
|
||||
// 列出转入余额与线下提现两类额度流出。
|
||||
func TestAffiliateTransferRecordsQueryIncludesOfflineWithdrawals(t *testing.T) {
|
||||
source, err := os.ReadFile("affiliate_repo.go")
|
||||
require.NoError(t, err)
|
||||
content := string(source)
|
||||
|
||||
require.Contains(t, content, "WHERE ual.action IN ('transfer', 'withdraw')")
|
||||
}
|
||||
|
||||
// TestAffiliateWithdrawClaimsOperationBeforeDeducting 锁定线下提现的幂等形态:
|
||||
// 同一事务内先按 operation_id 唯一约束写入占位流水,冲突时不扣减,
|
||||
// 且唯一约束由迁移建立。
|
||||
func TestAffiliateWithdrawClaimsOperationBeforeDeducting(t *testing.T) {
|
||||
source, err := os.ReadFile("affiliate_repo.go")
|
||||
require.NoError(t, err)
|
||||
content := string(source)
|
||||
|
||||
require.Contains(t, content, "ON CONFLICT (operation_id) WHERE operation_id IS NOT NULL DO NOTHING")
|
||||
withdraw := content[strings.Index(content, "func (r *affiliateRepository) WithdrawQuota("):]
|
||||
withdraw = withdraw[:strings.Index(withdraw, "\n}\n")]
|
||||
claimAt := strings.Index(withdraw, "claimAffiliateWithdrawLedger(")
|
||||
deductAt := strings.Index(withdraw, "SET aff_quota = aff_quota - $1")
|
||||
require.Positive(t, claimAt)
|
||||
require.Greater(t, deductAt, claimAt, "operation claim must precede the quota deduction")
|
||||
|
||||
migration, err := os.ReadFile("../../migrations/240_affiliate_ledger_operation_id.sql")
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, string(migration), "CREATE UNIQUE INDEX IF NOT EXISTS idx_user_affiliate_ledger_operation_id")
|
||||
require.Contains(t, string(migration), "WHERE operation_id IS NOT NULL")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
dbmigrations "github.com/Wei-Shaw/sub2api/migrations"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestMigration239ReasoningEffortMultipliers(t *testing.T) {
|
||||
tx := testTx(t)
|
||||
ctx := context.Background()
|
||||
_, err := tx.ExecContext(ctx, `ALTER TABLE channel_model_pricing DROP COLUMN reasoning_effort_multipliers`)
|
||||
require.NoError(t, err)
|
||||
_, err = tx.ExecContext(ctx, `ALTER TABLE channel_account_stats_model_pricing DROP COLUMN reasoning_effort_multipliers`)
|
||||
require.NoError(t, err)
|
||||
|
||||
var channelID, configuredID, unsetID int64
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `INSERT INTO channels (name) VALUES ('migration-reasoning-multipliers') RETURNING id`).Scan(&channelID))
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
INSERT INTO channel_model_pricing (channel_id, models, max_reasoning_effort_multiplier)
|
||||
VALUES ($1, '["custom-model"]', 2.5) RETURNING id`, channelID).Scan(&configuredID))
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
INSERT INTO channel_model_pricing (channel_id, models)
|
||||
VALUES ($1, '["claude-fable-5-1"]') RETURNING id`, channelID).Scan(&unsetID))
|
||||
|
||||
migrationSQL, err := dbmigrations.FS.ReadFile("239_channel_reasoning_effort_multipliers.sql")
|
||||
require.NoError(t, err)
|
||||
_, err = tx.ExecContext(ctx, string(migrationSQL))
|
||||
require.NoError(t, err)
|
||||
var configured, unset string
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `SELECT reasoning_effort_multipliers::text FROM channel_model_pricing WHERE id = $1`, configuredID).Scan(&configured))
|
||||
require.JSONEq(t, `{"max":2.5}`, configured)
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `SELECT reasoning_effort_multipliers::text FROM channel_model_pricing WHERE id = $1`, unsetID).Scan(&unset))
|
||||
require.JSONEq(t, `{}`, unset)
|
||||
|
||||
// Changing or clearing the generic configuration must survive a migration replay.
|
||||
for _, current := range []string{`{"high":1.5,"max":4}`, `{"low":0.75}`, `{}`} {
|
||||
_, err = tx.ExecContext(ctx, `UPDATE channel_model_pricing SET reasoning_effort_multipliers = $1 WHERE id = $2`, current, configuredID)
|
||||
require.NoError(t, err)
|
||||
_, err = tx.ExecContext(ctx, string(migrationSQL))
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `SELECT reasoning_effort_multipliers::text FROM channel_model_pricing WHERE id = $1`, configuredID).Scan(&configured))
|
||||
require.JSONEq(t, current, configured)
|
||||
}
|
||||
|
||||
// Account statistics use the same empty-by-default shape.
|
||||
var ruleID int64
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `INSERT INTO channel_account_stats_pricing_rules (channel_id) VALUES ($1) RETURNING id`, channelID).Scan(&ruleID))
|
||||
var accountStatsDefault string
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `INSERT INTO channel_account_stats_model_pricing (rule_id) VALUES ($1) RETURNING reasoning_effort_multipliers::text`, ruleID).Scan(&accountStatsDefault))
|
||||
require.JSONEq(t, `{}`, accountStatsDefault)
|
||||
}
|
||||
|
||||
func TestMigration239GroupReasoningEffortMultipliers(t *testing.T) {
|
||||
tx := testTx(t)
|
||||
ctx := context.Background()
|
||||
var groupID int64
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
INSERT INTO groups (name, platform, model_pricing)
|
||||
VALUES ('migration-group-reasoning', 'anthropic', '[
|
||||
{"models":["custom-model"],"input_price":2,"max_reasoning_effort_multiplier":2.5},
|
||||
{"models":["existing-map"],"max_reasoning_effort_multiplier":3,"reasoning_effort_multipliers":{"high":1.5,"max":4}},
|
||||
{"models":["cleared-map"],"max_reasoning_effort_multiplier":3,"reasoning_effort_multipliers":{}},
|
||||
{"models":["claude-fable-5-1"],"max_reasoning_effort_multiplier":null}
|
||||
]'::jsonb) RETURNING id`).Scan(&groupID))
|
||||
|
||||
migrationSQL, err := dbmigrations.FS.ReadFile("239_channel_reasoning_effort_multipliers.sql")
|
||||
require.NoError(t, err)
|
||||
for range 2 {
|
||||
_, err = tx.ExecContext(ctx, string(migrationSQL))
|
||||
require.NoError(t, err)
|
||||
var actual string
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `SELECT model_pricing::text FROM groups WHERE id = $1`, groupID).Scan(&actual))
|
||||
require.JSONEq(t, `[
|
||||
{"models":["custom-model"],"input_price":2,"reasoning_effort_multipliers":{"max":2.5}},
|
||||
{"models":["existing-map"],"reasoning_effort_multipliers":{"high":1.5,"max":4}},
|
||||
{"models":["cleared-map"],"reasoning_effort_multipliers":{}},
|
||||
{"models":["claude-fable-5-1"]}
|
||||
]`, actual)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newReasoningPricingIntegrationChannel(t *testing.T, channel *service.Channel) *channelRepository {
|
||||
t.Helper()
|
||||
repo := &channelRepository{db: integrationDB}
|
||||
channel.Name = t.Name()
|
||||
channel.Status = service.StatusActive
|
||||
channel.BillingModelSource = service.BillingModelSourceRequested
|
||||
require.NoError(t, repo.Create(context.Background(), channel))
|
||||
// Channel writes manage their own transactions, so clean up committed fixtures.
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, repo.Delete(context.Background(), channel.ID))
|
||||
})
|
||||
return repo
|
||||
}
|
||||
|
||||
func requireReasoningPricingIntegrationMaps(t *testing.T, pricing []service.ChannelModelPricing, expected map[string]map[string]float64) {
|
||||
t.Helper()
|
||||
require.Len(t, pricing, len(expected))
|
||||
for _, entry := range pricing {
|
||||
require.Len(t, entry.Models, 1)
|
||||
multipliers, ok := expected[entry.Models[0]]
|
||||
require.True(t, ok, "unexpected model %s", entry.Models[0])
|
||||
if len(multipliers) == 0 {
|
||||
require.Empty(t, entry.ReasoningEffortMultipliers, entry.Models[0])
|
||||
} else {
|
||||
require.Equal(t, multipliers, entry.ReasoningEffortMultipliers, entry.Models[0])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestChannelReasoningPricingCreateUpdateReplace(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
channel := &service.Channel{}
|
||||
repo := newReasoningPricingIntegrationChannel(t, channel)
|
||||
allLevels := map[string]float64{"none": 0.5, "minimal": 0.75, "low": 0.9, "medium": 1.25, "high": 1.5, "xhigh": 2, "max": 3}
|
||||
pricing := service.ChannelModelPricing{
|
||||
ChannelID: channel.ID, Platform: service.PlatformOpenAI, Models: []string{"custom-reasoning-model"},
|
||||
ReasoningEffortMultipliers: allLevels,
|
||||
}
|
||||
require.NoError(t, repo.CreateModelPricing(ctx, &pricing))
|
||||
require.NotZero(t, pricing.ID)
|
||||
assertStored := func(expected map[string]map[string]float64) {
|
||||
t.Helper()
|
||||
listed, err := repo.ListModelPricing(ctx, channel.ID)
|
||||
require.NoError(t, err)
|
||||
requireReasoningPricingIntegrationMaps(t, listed, expected)
|
||||
// Cache refresh uses this batch loading path instead of ListModelPricing.
|
||||
batch, err := repo.batchLoadModelPricing(ctx, []int64{channel.ID})
|
||||
require.NoError(t, err)
|
||||
requireReasoningPricingIntegrationMaps(t, batch[channel.ID], expected)
|
||||
}
|
||||
assertStored(map[string]map[string]float64{"custom-reasoning-model": allLevels})
|
||||
|
||||
pricing.ReasoningEffortMultipliers = map[string]float64{"high": 2.25}
|
||||
require.NoError(t, repo.UpdateModelPricing(ctx, &pricing))
|
||||
assertStored(map[string]map[string]float64{"custom-reasoning-model": {"high": 2.25}})
|
||||
|
||||
pricing.ReasoningEffortMultipliers = nil
|
||||
require.NoError(t, repo.UpdateModelPricing(ctx, &pricing))
|
||||
assertStored(map[string]map[string]float64{"custom-reasoning-model": nil})
|
||||
var stored string
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx,
|
||||
`SELECT reasoning_effort_multipliers::text FROM channel_model_pricing WHERE id = $1`, pricing.ID).Scan(&stored))
|
||||
require.JSONEq(t, `{}`, stored)
|
||||
|
||||
replacement := []service.ChannelModelPricing{
|
||||
{Platform: service.PlatformOpenAI, Models: []string{"custom-reasoning-model"}, ReasoningEffortMultipliers: map[string]float64{"low": 0.8, "max": 4}},
|
||||
{Platform: service.PlatformAnthropic, Models: []string{"claude-fable-5-1"}, ReasoningEffortMultipliers: map[string]float64{}},
|
||||
}
|
||||
require.NoError(t, repo.ReplaceModelPricing(ctx, channel.ID, replacement))
|
||||
expected := map[string]map[string]float64{"custom-reasoning-model": {"low": 0.8, "max": 4}, "claude-fable-5-1": nil}
|
||||
assertStored(expected)
|
||||
|
||||
// A serialization failure after deleting the old rows must roll back the replacement.
|
||||
require.Error(t, repo.ReplaceModelPricing(ctx, channel.ID, []service.ChannelModelPricing{{
|
||||
Models: []string{"invalid"}, ReasoningEffortMultipliers: map[string]float64{"max": math.Inf(1)},
|
||||
}}))
|
||||
assertStored(expected)
|
||||
}
|
||||
|
||||
func TestAccountStatsReasoningPricingRoundTrip(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
channel := &service.Channel{
|
||||
ModelPricing: []service.ChannelModelPricing{{
|
||||
Platform: service.PlatformOpenAI, Models: []string{"custom-reasoning-model"},
|
||||
ReasoningEffortMultipliers: map[string]float64{"max": 4},
|
||||
}},
|
||||
AccountStatsPricingRules: []service.AccountStatsPricingRule{{
|
||||
Name: "independent account statistics", GroupIDs: []int64{}, AccountIDs: []int64{},
|
||||
Pricing: []service.ChannelModelPricing{
|
||||
{Platform: service.PlatformOpenAI, Models: []string{"custom-reasoning-model"}, ReasoningEffortMultipliers: map[string]float64{"none": 0.5, "high": 1.25, "max": 2.5}},
|
||||
{Platform: service.PlatformAnthropic, Models: []string{"claude-fable-5-1"}},
|
||||
},
|
||||
}},
|
||||
}
|
||||
repo := newReasoningPricingIntegrationChannel(t, channel)
|
||||
assertStored := func(multipliers map[string]float64) *service.Channel {
|
||||
t.Helper()
|
||||
loaded, err := repo.GetByID(ctx, channel.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, loaded.AccountStatsPricingRules, 1)
|
||||
expected := map[string]map[string]float64{"custom-reasoning-model": multipliers, "claude-fable-5-1": nil}
|
||||
requireReasoningPricingIntegrationMaps(t, loaded.AccountStatsPricingRules[0].Pricing, expected)
|
||||
batch, err := repo.batchLoadAccountStatsPricingRules(ctx, []int64{channel.ID})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, batch[channel.ID], 1)
|
||||
requireReasoningPricingIntegrationMaps(t, batch[channel.ID][0].Pricing, expected)
|
||||
// Custom statistics prices must not alter the user-facing channel prices.
|
||||
requireReasoningPricingIntegrationMaps(t, loaded.ModelPricing, map[string]map[string]float64{"custom-reasoning-model": {"max": 4}})
|
||||
return loaded
|
||||
}
|
||||
loaded := assertStored(map[string]float64{"none": 0.5, "high": 1.25, "max": 2.5})
|
||||
for _, multipliers := range []map[string]float64{{"low": 0.75, "max": 1.5}, nil, {"high": 2}, {}} {
|
||||
loaded.AccountStatsPricingRules[0].Pricing[0].ReasoningEffortMultipliers = multipliers
|
||||
require.NoError(t, repo.Update(ctx, loaded))
|
||||
loaded = assertStored(multipliers)
|
||||
}
|
||||
var stored string
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx,
|
||||
`SELECT reasoning_effort_multipliers::text FROM channel_account_stats_model_pricing WHERE id = $1`, loaded.AccountStatsPricingRules[0].Pricing[0].ID).Scan(&stored))
|
||||
require.JSONEq(t, `{}`, stored)
|
||||
}
|
||||
|
||||
func TestGroupReasoningPricingRoundTripAndBilling(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newGroupRepositoryWithSQL(tx.Client(), tx)
|
||||
inputPrice, outputPrice := 1e-6, 2e-6
|
||||
allLevels := map[string]float64{"none": 0.5, "minimal": 0.75, "low": 0.9, "medium": 1.25, "high": 1.5, "xhigh": 2, "max": 3}
|
||||
group := &service.Group{
|
||||
Name: t.Name(), Platform: service.PlatformAnthropic, RateMultiplier: 0.8,
|
||||
Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard,
|
||||
ModelPricing: []service.ChannelModelPricing{{
|
||||
Platform: service.PlatformAnthropic, Models: []string{"claude-fable-5-1"}, BillingMode: service.BillingModeToken,
|
||||
InputPrice: &inputPrice, OutputPrice: &outputPrice, ReasoningEffortMultipliers: allLevels,
|
||||
}},
|
||||
}
|
||||
require.NoError(t, repo.Create(ctx, group))
|
||||
billing := service.NewBillingService(&config.Config{}, nil)
|
||||
resolver := service.NewModelPricingResolver(nil, billing)
|
||||
assertStoredAndBilled := func(expected map[string]float64) *service.Group {
|
||||
t.Helper()
|
||||
loaded, err := repo.GetByID(ctx, group.ID)
|
||||
require.NoError(t, err)
|
||||
requireReasoningPricingIntegrationMaps(t, loaded.ModelPricing, map[string]map[string]float64{"claude-fable-5-1": expected})
|
||||
for _, effort := range []string{"none", "minimal", "low", "medium", "high", "xhigh", "max", ""} {
|
||||
multiplier := 1.0
|
||||
if configured, ok := expected[effort]; ok {
|
||||
multiplier = configured
|
||||
}
|
||||
cost, err := billing.CalculateTokenCostForRequest(service.TokenCostRequest{
|
||||
Ctx: ctx, Model: "claude-fable-5-1", Group: loaded, Resolver: resolver,
|
||||
Tokens: service.UsageTokens{InputTokens: 1000, OutputTokens: 200},
|
||||
RateMultiplier: loaded.RateMultiplier, ReasoningEffort: effort,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 0.0014*multiplier, cost.TotalCost, 1e-12, "effort=%q", effort)
|
||||
require.InDelta(t, 0.0014*multiplier*0.8, cost.ActualCost, 1e-12, "effort=%q", effort)
|
||||
}
|
||||
return loaded
|
||||
}
|
||||
loaded := assertStoredAndBilled(allLevels)
|
||||
for _, multipliers := range []map[string]float64{{"high": 2.25, "max": 4}, nil, {"high": 1.75}, {}} {
|
||||
loaded.ModelPricing[0].ReasoningEffortMultipliers = multipliers
|
||||
require.NoError(t, repo.Update(ctx, loaded))
|
||||
loaded = assertStoredAndBilled(multipliers)
|
||||
}
|
||||
}
|
||||
@@ -67,7 +67,7 @@ func (r *channelRepository) batchLoadAccountStatsModelPricing(ctx context.Contex
|
||||
|
||||
rows, err := r.db.QueryContext(ctx,
|
||||
`SELECT id, rule_id, platform, models, billing_mode, input_price, output_price,
|
||||
cache_write_price, cache_write_1h_price, cache_read_price, image_output_price, per_request_price, created_at, updated_at
|
||||
cache_write_price, cache_write_1h_price, cache_read_price, reasoning_effort_multipliers, image_output_price, per_request_price, created_at, updated_at
|
||||
FROM channel_account_stats_model_pricing WHERE rule_id = ANY($1) ORDER BY rule_id, id`,
|
||||
pq.Array(ruleIDs),
|
||||
)
|
||||
@@ -81,16 +81,22 @@ func (r *channelRepository) batchLoadAccountStatsModelPricing(ctx context.Contex
|
||||
var p service.ChannelModelPricing
|
||||
var ruleID int64
|
||||
var modelsJSON []byte
|
||||
var reasoningEffortMultipliersJSON []byte
|
||||
if err := rows.Scan(
|
||||
&p.ID, &ruleID, &p.Platform, &modelsJSON, &p.BillingMode,
|
||||
&p.InputPrice, &p.OutputPrice, &p.CacheWritePrice, &p.CacheWrite1hPrice, &p.CacheReadPrice,
|
||||
&p.ImageOutputPrice, &p.PerRequestPrice, &p.CreatedAt, &p.UpdatedAt,
|
||||
&reasoningEffortMultipliersJSON, &p.ImageOutputPrice, &p.PerRequestPrice, &p.CreatedAt, &p.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, fmt.Errorf("scan account stats model pricing: %w", err)
|
||||
}
|
||||
if err := json.Unmarshal(modelsJSON, &p.Models); err != nil {
|
||||
p.Models = []string{}
|
||||
}
|
||||
multipliers, err := unmarshalReasoningEffortMultipliers(reasoningEffortMultipliersJSON)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
p.ReasoningEffortMultipliers = multipliers
|
||||
pricingMap[ruleID] = append(pricingMap[ruleID], p)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
@@ -172,17 +178,21 @@ func createAccountStatsModelPricingTx(ctx context.Context, tx *sql.Tx, ruleID in
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal models: %w", err)
|
||||
}
|
||||
reasoningEffortMultipliersJSON, err := marshalReasoningEffortMultipliers(pricing.ReasoningEffortMultipliers)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
billingMode := pricing.BillingMode
|
||||
if billingMode == "" {
|
||||
billingMode = service.BillingModeToken
|
||||
}
|
||||
platform := pricing.Platform
|
||||
err = tx.QueryRowContext(ctx,
|
||||
`INSERT INTO channel_account_stats_model_pricing (rule_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, image_output_price, per_request_price)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) RETURNING id, created_at, updated_at`,
|
||||
`INSERT INTO channel_account_stats_model_pricing (rule_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, reasoning_effort_multipliers, image_output_price, per_request_price)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) RETURNING id, created_at, updated_at`,
|
||||
ruleID, platform, modelsJSON, billingMode,
|
||||
pricing.InputPrice, pricing.OutputPrice, pricing.CacheWritePrice, pricing.CacheWrite1hPrice, pricing.CacheReadPrice,
|
||||
pricing.ImageOutputPrice, pricing.PerRequestPrice,
|
||||
reasoningEffortMultipliersJSON, pricing.ImageOutputPrice, pricing.PerRequestPrice,
|
||||
).Scan(&pricing.ID, &pricing.CreatedAt, &pricing.UpdatedAt)
|
||||
if err != nil {
|
||||
return fmt.Errorf("insert account stats model pricing: %w", err)
|
||||
|
||||
@@ -16,7 +16,7 @@ import (
|
||||
|
||||
func (r *channelRepository) ListModelPricing(ctx context.Context, channelID int64) ([]service.ChannelModelPricing, error) {
|
||||
rows, err := r.db.QueryContext(ctx,
|
||||
`SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, max_reasoning_effort_multiplier, image_input_price, image_output_price, per_request_price, time_pricing, created_at, updated_at
|
||||
`SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, reasoning_effort_multipliers, image_input_price, image_output_price, per_request_price, time_pricing, created_at, updated_at
|
||||
FROM channel_model_pricing WHERE channel_id = $1 ORDER BY id`, channelID,
|
||||
)
|
||||
if err != nil {
|
||||
@@ -51,6 +51,10 @@ func (r *channelRepository) UpdateModelPricing(ctx context.Context, pricing *ser
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal models: %w", err)
|
||||
}
|
||||
reasoningEffortMultipliersJSON, err := marshalReasoningEffortMultipliers(pricing.ReasoningEffortMultipliers)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
timePricingJSON, err := marshalChannelTimePricing(pricing.TimePricing)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -61,10 +65,10 @@ func (r *channelRepository) UpdateModelPricing(ctx context.Context, pricing *ser
|
||||
}
|
||||
result, err := r.db.ExecContext(ctx,
|
||||
`UPDATE channel_model_pricing
|
||||
SET models = $1, billing_mode = $2, input_price = $3, output_price = $4, cache_write_price = $5, cache_write_1h_price = $6, cache_read_price = $7, fast_multiplier = $8, flex_multiplier = $9, max_reasoning_effort_multiplier = $10, image_input_price = $11, image_output_price = $12, per_request_price = $13, time_pricing = $14, platform = $15, updated_at = NOW()
|
||||
SET models = $1, billing_mode = $2, input_price = $3, output_price = $4, cache_write_price = $5, cache_write_1h_price = $6, cache_read_price = $7, fast_multiplier = $8, flex_multiplier = $9, reasoning_effort_multipliers = $10, image_input_price = $11, image_output_price = $12, per_request_price = $13, time_pricing = $14, platform = $15, updated_at = NOW()
|
||||
WHERE id = $16`,
|
||||
modelsJSON, billingMode, pricing.InputPrice, pricing.OutputPrice, pricing.CacheWritePrice, pricing.CacheWrite1hPrice, pricing.CacheReadPrice,
|
||||
pricing.FastMultiplier, pricing.FlexMultiplier, pricing.MaxReasoningEffortMultiplier, pricing.ImageInputPrice, pricing.ImageOutputPrice, pricing.PerRequestPrice,
|
||||
pricing.FastMultiplier, pricing.FlexMultiplier, reasoningEffortMultipliersJSON, pricing.ImageInputPrice, pricing.ImageOutputPrice, pricing.PerRequestPrice,
|
||||
timePricingJSON, pricing.Platform, pricing.ID,
|
||||
)
|
||||
if err != nil {
|
||||
@@ -96,7 +100,7 @@ func (r *channelRepository) ReplaceModelPricing(ctx context.Context, channelID i
|
||||
// batchLoadModelPricing 批量加载多个渠道的模型定价(含区间)
|
||||
func (r *channelRepository) batchLoadModelPricing(ctx context.Context, channelIDs []int64) (map[int64][]service.ChannelModelPricing, error) {
|
||||
rows, err := r.db.QueryContext(ctx,
|
||||
`SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, max_reasoning_effort_multiplier, image_input_price, image_output_price, per_request_price, time_pricing, created_at, updated_at
|
||||
`SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, reasoning_effort_multipliers, image_input_price, image_output_price, per_request_price, time_pricing, created_at, updated_at
|
||||
FROM channel_model_pricing WHERE channel_id = ANY($1) ORDER BY channel_id, id`,
|
||||
pq.Array(channelIDs),
|
||||
)
|
||||
@@ -177,10 +181,11 @@ func scanModelPricingRows(rows *sql.Rows) ([]service.ChannelModelPricing, []int6
|
||||
var p service.ChannelModelPricing
|
||||
var modelsJSON []byte
|
||||
var timePricingJSON []byte
|
||||
var reasoningEffortMultipliersJSON []byte
|
||||
if err := rows.Scan(
|
||||
&p.ID, &p.ChannelID, &p.Platform, &modelsJSON, &p.BillingMode,
|
||||
&p.InputPrice, &p.OutputPrice, &p.CacheWritePrice, &p.CacheWrite1hPrice, &p.CacheReadPrice,
|
||||
&p.FastMultiplier, &p.FlexMultiplier, &p.MaxReasoningEffortMultiplier,
|
||||
&p.FastMultiplier, &p.FlexMultiplier, &reasoningEffortMultipliersJSON,
|
||||
&p.ImageInputPrice, &p.ImageOutputPrice, &p.PerRequestPrice, &timePricingJSON, &p.CreatedAt, &p.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, nil, fmt.Errorf("scan model pricing: %w", err)
|
||||
@@ -188,6 +193,11 @@ func scanModelPricingRows(rows *sql.Rows) ([]service.ChannelModelPricing, []int6
|
||||
if err := json.Unmarshal(modelsJSON, &p.Models); err != nil {
|
||||
p.Models = []string{}
|
||||
}
|
||||
multipliers, err := unmarshalReasoningEffortMultipliers(reasoningEffortMultipliersJSON)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
p.ReasoningEffortMultipliers = multipliers
|
||||
timePricing, err := unmarshalChannelTimePricing(timePricingJSON)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
@@ -234,6 +244,10 @@ func createModelPricingExec(ctx context.Context, exec dbExec, pricing *service.C
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal models: %w", err)
|
||||
}
|
||||
reasoningEffortMultipliersJSON, err := marshalReasoningEffortMultipliers(pricing.ReasoningEffortMultipliers)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
timePricingJSON, err := marshalChannelTimePricing(pricing.TimePricing)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -247,11 +261,11 @@ func createModelPricingExec(ctx context.Context, exec dbExec, pricing *service.C
|
||||
platform = "anthropic"
|
||||
}
|
||||
err = exec.QueryRowContext(ctx,
|
||||
`INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, max_reasoning_effort_multiplier, image_input_price, image_output_price, per_request_price, time_pricing)
|
||||
`INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, reasoning_effort_multipliers, image_input_price, image_output_price, per_request_price, time_pricing)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16) RETURNING id, created_at, updated_at`,
|
||||
pricing.ChannelID, platform, modelsJSON, billingMode,
|
||||
pricing.InputPrice, pricing.OutputPrice, pricing.CacheWritePrice, pricing.CacheWrite1hPrice, pricing.CacheReadPrice,
|
||||
pricing.FastMultiplier, pricing.FlexMultiplier, pricing.MaxReasoningEffortMultiplier, pricing.ImageInputPrice, pricing.ImageOutputPrice,
|
||||
pricing.FastMultiplier, pricing.FlexMultiplier, reasoningEffortMultipliersJSON, pricing.ImageInputPrice, pricing.ImageOutputPrice,
|
||||
pricing.PerRequestPrice, timePricingJSON,
|
||||
).Scan(&pricing.ID, &pricing.CreatedAt, &pricing.UpdatedAt)
|
||||
if err != nil {
|
||||
@@ -268,6 +282,31 @@ func createModelPricingExec(ctx context.Context, exec dbExec, pricing *service.C
|
||||
return nil
|
||||
}
|
||||
|
||||
func marshalReasoningEffortMultipliers(multipliers map[string]float64) (string, error) {
|
||||
if len(multipliers) == 0 {
|
||||
return "{}", nil
|
||||
}
|
||||
data, err := json.Marshal(multipliers)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("marshal reasoning effort multipliers: %w", err)
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
func unmarshalReasoningEffortMultipliers(data []byte) (map[string]float64, error) {
|
||||
multipliers := make(map[string]float64)
|
||||
if len(data) == 0 {
|
||||
return multipliers, nil
|
||||
}
|
||||
if err := json.Unmarshal(data, &multipliers); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal reasoning effort multipliers: %w", err)
|
||||
}
|
||||
if multipliers == nil {
|
||||
multipliers = make(map[string]float64)
|
||||
}
|
||||
return multipliers, nil
|
||||
}
|
||||
|
||||
func marshalChannelTimePricing(config *service.ChannelTimePricing) (any, error) {
|
||||
if config == nil || len(config.Periods) == 0 {
|
||||
return nil, nil
|
||||
|
||||
@@ -0,0 +1,175 @@
|
||||
//go:build unit
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql/driver"
|
||||
"math"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/lib/pq"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type capturedReasoningMultipliersJSON struct {
|
||||
value string
|
||||
}
|
||||
|
||||
func (c *capturedReasoningMultipliersJSON) Match(value driver.Value) bool {
|
||||
data, ok := value.(string)
|
||||
if ok {
|
||||
c.value = data
|
||||
}
|
||||
return ok
|
||||
}
|
||||
|
||||
func reasoningPricingRow(multipliers any) *sqlmock.Rows {
|
||||
return sqlmock.NewRows(channelModelPricingTimePricingColumns).AddRow(
|
||||
int64(11), int64(7), "openai", `["custom-model"]`, service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, multipliers, nil, nil, nil, nil,
|
||||
time.Time{}, time.Time{},
|
||||
)
|
||||
}
|
||||
|
||||
func TestChannelReasoningEffortMultipliersRoundTrip(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
multipliers map[string]float64
|
||||
wantJSON string
|
||||
}{
|
||||
{"custom levels", map[string]float64{"low": 0.75, "high": 1.5, "max": 4}, `{"low":0.75,"high":1.5,"max":4}`},
|
||||
{"nil clears configuration", nil, `{}`},
|
||||
{"empty clears configuration", map[string]float64{}, `{}`},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
for _, operation := range []string{"create", "update", "replace"} {
|
||||
t.Run(tt.name+"/"+operation, func(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
ctx := context.Background()
|
||||
pricing := &service.ChannelModelPricing{
|
||||
ID: 11, ChannelID: 7, Platform: "openai", Models: []string{"custom-model"},
|
||||
ReasoningEffortMultipliers: tt.multipliers,
|
||||
}
|
||||
stored := &capturedReasoningMultipliersJSON{}
|
||||
if operation == "update" {
|
||||
mock.ExpectExec(`(?s)UPDATE channel_model_pricing.*reasoning_effort_multipliers = \$10.*WHERE id = \$16`).
|
||||
WithArgs([]byte(`["custom-model"]`), service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, stored, nil, nil, nil, nil, "openai", int64(11)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
require.NoError(t, repo.UpdateModelPricing(ctx, pricing))
|
||||
} else {
|
||||
if operation == "replace" {
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(`DELETE FROM channel_model_pricing WHERE channel_id = \$1`).
|
||||
WithArgs(int64(7)).WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
}
|
||||
mock.ExpectQuery(`INSERT INTO channel_model_pricing .*reasoning_effort_multipliers`).
|
||||
WithArgs(int64(7), "openai", []byte(`["custom-model"]`), service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, stored, nil, nil, nil, nil).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "created_at", "updated_at"}).AddRow(int64(11), time.Time{}, time.Time{}))
|
||||
if operation == "replace" {
|
||||
mock.ExpectCommit()
|
||||
require.NoError(t, repo.ReplaceModelPricing(ctx, 7, []service.ChannelModelPricing{*pricing}))
|
||||
} else {
|
||||
require.NoError(t, repo.CreateModelPricing(ctx, pricing))
|
||||
}
|
||||
}
|
||||
require.JSONEq(t, tt.wantJSON, stored.value)
|
||||
|
||||
mock.ExpectQuery(`(?s)SELECT .*reasoning_effort_multipliers.*FROM channel_model_pricing.*channel_id = \$1`).
|
||||
WithArgs(int64(7)).WillReturnRows(reasoningPricingRow(stored.value))
|
||||
expectEmptyModelPricingIntervals(mock)
|
||||
loaded, err := repo.ListModelPricing(ctx, 7)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, loaded, 1)
|
||||
if len(tt.multipliers) == 0 {
|
||||
require.Empty(t, loaded[0].ReasoningEffortMultipliers)
|
||||
} else {
|
||||
require.Equal(t, tt.multipliers, loaded[0].ReasoningEffortMultipliers)
|
||||
}
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestChannelReasoningEffortMultipliersBatchLoad(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
mock.ExpectQuery(`(?s)SELECT .*reasoning_effort_multipliers.*FROM channel_model_pricing.*channel_id = ANY`).
|
||||
WithArgs(pq.Array([]int64{7})).WillReturnRows(reasoningPricingRow(`{"medium":1.25,"high":2}`))
|
||||
expectEmptyModelPricingIntervals(mock)
|
||||
pricing, err := repo.batchLoadModelPricing(context.Background(), []int64{7})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, pricing[7], 1)
|
||||
require.Equal(t, map[string]float64{"medium": 1.25, "high": 2}, pricing[7][0].ReasoningEffortMultipliers)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestChannelReasoningEffortMultipliersInvalidJSON(t *testing.T) {
|
||||
for _, data := range []string{`{"high":`, `{"high":"2"}`, `[]`} {
|
||||
t.Run(data, func(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
mock.ExpectQuery(`SELECT .*reasoning_effort_multipliers`).WithArgs(int64(7)).WillReturnRows(reasoningPricingRow(data))
|
||||
_, err := repo.ListModelPricing(context.Background(), 7)
|
||||
require.ErrorContains(t, err, "unmarshal reasoning effort multipliers")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
}
|
||||
for _, operation := range []string{"create", "update"} {
|
||||
t.Run(operation+" invalid number", func(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
pricing := &service.ChannelModelPricing{ReasoningEffortMultipliers: map[string]float64{"max": math.NaN()}}
|
||||
var err error
|
||||
if operation == "create" {
|
||||
err = repo.CreateModelPricing(context.Background(), pricing)
|
||||
} else {
|
||||
err = repo.UpdateModelPricing(context.Background(), pricing)
|
||||
}
|
||||
require.ErrorContains(t, err, "marshal reasoning effort multipliers")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccountStatsReasoningEffortMultipliersRoundTrip(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
ctx := context.Background()
|
||||
pricing := &service.ChannelModelPricing{
|
||||
Platform: "anthropic", Models: []string{"custom-model"},
|
||||
ReasoningEffortMultipliers: map[string]float64{"medium": 0.5, "high": 2},
|
||||
}
|
||||
stored := &capturedReasoningMultipliersJSON{}
|
||||
mock.ExpectBegin()
|
||||
tx, err := repo.db.BeginTx(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
mock.ExpectQuery(`INSERT INTO channel_account_stats_model_pricing .*reasoning_effort_multipliers`).
|
||||
WithArgs(int64(9), "anthropic", []byte(`["custom-model"]`), service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, stored, nil, nil).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "created_at", "updated_at"}).AddRow(int64(11), time.Time{}, time.Time{}))
|
||||
require.NoError(t, createAccountStatsModelPricingTx(ctx, tx, 9, pricing))
|
||||
mock.ExpectCommit()
|
||||
require.NoError(t, tx.Commit())
|
||||
require.JSONEq(t, `{"medium":0.5,"high":2}`, stored.value)
|
||||
|
||||
columns := []string{
|
||||
"id", "rule_id", "platform", "models", "billing_mode", "input_price", "output_price",
|
||||
"cache_write_price", "cache_write_1h_price", "cache_read_price", "reasoning_effort_multipliers",
|
||||
"image_output_price", "per_request_price", "created_at", "updated_at",
|
||||
}
|
||||
mock.ExpectQuery(`(?s)SELECT .*reasoning_effort_multipliers.*FROM channel_account_stats_model_pricing`).
|
||||
WithArgs(pq.Array([]int64{9})).WillReturnRows(sqlmock.NewRows(columns).AddRow(
|
||||
int64(11), int64(9), "anthropic", `["custom-model"]`, service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, stored.value, nil, nil, time.Time{}, time.Time{},
|
||||
))
|
||||
mock.ExpectQuery(`(?s)SELECT .*FROM channel_account_stats_pricing_intervals`).
|
||||
WithArgs(pq.Array([]int64{11})).WillReturnRows(sqlmock.NewRows([]string{"id"}))
|
||||
loaded, err := repo.batchLoadAccountStatsModelPricing(ctx, []int64{9})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, loaded[9], 1)
|
||||
require.Equal(t, pricing.ReasoningEffortMultipliers, loaded[9][0].ReasoningEffortMultipliers)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
@@ -16,7 +16,7 @@ import (
|
||||
|
||||
var channelModelPricingTimePricingColumns = []string{
|
||||
"id", "channel_id", "platform", "models", "billing_mode", "input_price", "output_price",
|
||||
"cache_write_price", "cache_write_1h_price", "cache_read_price", "fast_multiplier", "flex_multiplier", "max_reasoning_effort_multiplier", "image_input_price", "image_output_price",
|
||||
"cache_write_price", "cache_write_1h_price", "cache_read_price", "fast_multiplier", "flex_multiplier", "reasoning_effort_multipliers", "image_input_price", "image_output_price",
|
||||
"per_request_price", "time_pricing", "created_at", "updated_at",
|
||||
}
|
||||
|
||||
@@ -33,7 +33,7 @@ func newChannelModelPricingTimePricingRepo(t *testing.T) (*channelRepository, sq
|
||||
func modelPricingTimePricingRow(timePricing any) *sqlmock.Rows {
|
||||
return sqlmock.NewRows(channelModelPricingTimePricingColumns).AddRow(
|
||||
int64(11), int64(7), "openai", `["gpt-5"]`, service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, timePricing,
|
||||
nil, nil, nil, nil, nil, nil, nil, "{}", nil, nil, nil, timePricing,
|
||||
time.Date(2026, 8, 17, 0, 0, 0, 0, time.UTC), time.Date(2026, 8, 17, 1, 0, 0, 0, time.UTC),
|
||||
)
|
||||
}
|
||||
@@ -105,10 +105,10 @@ func TestChannelModelPricingTimePricingCreateAndUpdateRoundTrip(t *testing.T) {
|
||||
|
||||
t.Run("create writes JSON", func(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
mock.ExpectQuery(regexp.QuoteMeta("INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, max_reasoning_effort_multiplier, image_input_price, image_output_price, per_request_price, time_pricing)")).
|
||||
mock.ExpectQuery(regexp.QuoteMeta("INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, reasoning_effort_multipliers, image_input_price, image_output_price, per_request_price, time_pricing)")).
|
||||
WithArgs(
|
||||
int64(7), "openai", []byte(`["gpt-5"]`), service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, channelModelPricingTimePricingJSON,
|
||||
nil, nil, nil, nil, nil, nil, nil, "{}", nil, nil, nil, channelModelPricingTimePricingJSON,
|
||||
).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "created_at", "updated_at"}).AddRow(int64(11), time.Time{}, time.Time{}))
|
||||
|
||||
@@ -121,7 +121,7 @@ func TestChannelModelPricingTimePricingCreateAndUpdateRoundTrip(t *testing.T) {
|
||||
mock.ExpectExec(`(?s)UPDATE channel_model_pricing.*per_request_price = \$13, time_pricing = \$14, platform = \$15.*WHERE id = \$16`).
|
||||
WithArgs(
|
||||
[]byte(`["gpt-5"]`), service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, channelModelPricingTimePricingJSON, "openai", int64(11),
|
||||
nil, nil, nil, nil, nil, nil, nil, "{}", nil, nil, nil, channelModelPricingTimePricingJSON, "openai", int64(11),
|
||||
).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
@@ -153,10 +153,10 @@ func TestChannelModelPricingTimePricingCreateAndUpdateWriteNullWhenDisabled(t *t
|
||||
|
||||
t.Run("create writes SQL NULL", func(t *testing.T) {
|
||||
repo, mock := newChannelModelPricingTimePricingRepo(t)
|
||||
mock.ExpectQuery(regexp.QuoteMeta("INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, max_reasoning_effort_multiplier, image_input_price, image_output_price, per_request_price, time_pricing)")).
|
||||
mock.ExpectQuery(regexp.QuoteMeta("INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, reasoning_effort_multipliers, image_input_price, image_output_price, per_request_price, time_pricing)")).
|
||||
WithArgs(
|
||||
int64(7), "openai", []byte(`["gpt-5"]`), service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
nil, nil, nil, nil, nil, nil, nil, "{}", nil, nil, nil, nil,
|
||||
).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "created_at", "updated_at"}).AddRow(int64(11), time.Time{}, time.Time{}))
|
||||
|
||||
@@ -169,7 +169,7 @@ func TestChannelModelPricingTimePricingCreateAndUpdateWriteNullWhenDisabled(t *t
|
||||
mock.ExpectExec(`(?s)UPDATE channel_model_pricing.*per_request_price = \$13, time_pricing = \$14, platform = \$15.*WHERE id = \$16`).
|
||||
WithArgs(
|
||||
[]byte(`["gpt-5"]`), service.BillingModeToken,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, "openai", int64(11),
|
||||
nil, nil, nil, nil, nil, nil, nil, "{}", nil, nil, nil, nil, "openai", int64(11),
|
||||
).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
|
||||
@@ -32,6 +32,14 @@ func (r *contentModerationRepository) CreateLog(ctx context.Context, log *servic
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal moderation thresholds: %w", err)
|
||||
}
|
||||
var engineMeta any
|
||||
if log.EngineMeta != nil {
|
||||
raw, err := json.Marshal(log.EngineMeta)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal moderation engine metadata: %w", err)
|
||||
}
|
||||
engineMeta = string(raw)
|
||||
}
|
||||
var userID any
|
||||
if log.UserID != nil {
|
||||
userID = *log.UserID
|
||||
@@ -53,17 +61,17 @@ INSERT INTO content_moderation_logs (
|
||||
request_id, user_id, user_email, api_key_id, api_key_name, group_id, group_name,
|
||||
endpoint, provider, model, mode, action, flagged, highest_category, highest_score,
|
||||
category_scores, threshold_snapshot, input_excerpt, upstream_latency_ms, error,
|
||||
violation_count, auto_banned, email_sent, queue_delay_ms, matched_keyword
|
||||
violation_count, auto_banned, email_sent, queue_delay_ms, matched_keyword, engine_meta
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7,
|
||||
$8, $9, $10, $11, $12, $13, $14, $15,
|
||||
$16::jsonb, $17::jsonb, $18, $19, $20,
|
||||
$21, $22, $23, $24, $25
|
||||
$21, $22, $23, $24, $25, $26::jsonb
|
||||
) RETURNING id, created_at`,
|
||||
log.RequestID, userID, log.UserEmail, apiKeyID, log.APIKeyName, groupID, log.GroupName,
|
||||
log.Endpoint, log.Provider, log.Model, log.Mode, log.Action, log.Flagged, log.HighestCategory, log.HighestScore,
|
||||
string(categoryScores), string(thresholdSnapshot), log.InputExcerpt, latency, log.Error,
|
||||
log.ViolationCount, log.AutoBanned, log.EmailSent, nullableIntPtr(log.QueueDelayMS), log.MatchedKeyword,
|
||||
log.ViolationCount, log.AutoBanned, log.EmailSent, nullableIntPtr(log.QueueDelayMS), log.MatchedKeyword, engineMeta,
|
||||
).Scan(&log.ID, &log.CreatedAt)
|
||||
if err != nil {
|
||||
return fmt.Errorf("insert content moderation log: %w", err)
|
||||
@@ -97,7 +105,7 @@ SELECT
|
||||
l.id, l.request_id, l.user_id, l.user_email, l.api_key_id, l.api_key_name, l.group_id, l.group_name,
|
||||
l.endpoint, l.provider, l.model, l.mode, l.action, l.flagged, l.highest_category, l.highest_score,
|
||||
l.category_scores, l.threshold_snapshot, l.input_excerpt, l.upstream_latency_ms, l.error,
|
||||
l.violation_count, l.auto_banned, l.email_sent, COALESCE(u.status, ''), l.queue_delay_ms, l.matched_keyword, l.created_at
|
||||
l.violation_count, l.auto_banned, l.email_sent, COALESCE(u.status, ''), l.queue_delay_ms, l.matched_keyword, l.created_at, l.engine_meta
|
||||
FROM content_moderation_logs l
|
||||
LEFT JOIN users u ON u.id = l.user_id `+whereSQL+`
|
||||
ORDER BY l.created_at DESC, l.id DESC
|
||||
@@ -113,7 +121,7 @@ LIMIT $`+fmt.Sprint(len(queryArgs)-1)+` OFFSET $`+fmt.Sprint(len(queryArgs)),
|
||||
for rows.Next() {
|
||||
var item service.ContentModerationLog
|
||||
var userID, apiKeyID, groupID, latency, queueDelay sql.NullInt64
|
||||
var scoresRaw, thresholdsRaw []byte
|
||||
var scoresRaw, thresholdsRaw, engineRaw []byte
|
||||
if err := rows.Scan(
|
||||
&item.ID,
|
||||
&item.RequestID,
|
||||
@@ -143,6 +151,7 @@ LIMIT $`+fmt.Sprint(len(queryArgs)-1)+` OFFSET $`+fmt.Sprint(len(queryArgs)),
|
||||
&queueDelay,
|
||||
&item.MatchedKeyword,
|
||||
&item.CreatedAt,
|
||||
&engineRaw,
|
||||
); err != nil {
|
||||
return nil, nil, fmt.Errorf("scan content moderation log: %w", err)
|
||||
}
|
||||
@@ -170,6 +179,11 @@ LIMIT $`+fmt.Sprint(len(queryArgs)-1)+` OFFSET $`+fmt.Sprint(len(queryArgs)),
|
||||
_ = json.Unmarshal(scoresRaw, &item.CategoryScores)
|
||||
item.ThresholdSnapshot = map[string]float64{}
|
||||
_ = json.Unmarshal(thresholdsRaw, &item.ThresholdSnapshot)
|
||||
if len(engineRaw) > 0 {
|
||||
if err := json.Unmarshal(engineRaw, &item.EngineMeta); err != nil {
|
||||
return nil, nil, fmt.Errorf("decode moderation engine metadata: %w", err)
|
||||
}
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
|
||||
@@ -2,6 +2,7 @@ package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql/driver"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
@@ -12,6 +13,51 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestContentModerationRepositoryEngineMetaInsert(t *testing.T) {
|
||||
for _, meta := range []*service.ContentModerationEngineMeta{nil, {Engine: "typesafe", Model: "jev-fixed", RulesVersion: "rules-v1", SkippedImages: 1}} {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
args := make([]driver.Value, 26)
|
||||
for i := range args {
|
||||
args[i] = sqlmock.AnyArg()
|
||||
}
|
||||
if meta == nil {
|
||||
args[25] = nil
|
||||
} else {
|
||||
args[25] = `{"engine":"typesafe","model":"jev-fixed","rules_version":"rules-v1","skipped_images":1}`
|
||||
}
|
||||
mock.ExpectQuery("INSERT INTO content_moderation_logs").WithArgs(args...).WillReturnRows(sqlmock.NewRows([]string{"id", "created_at"}).AddRow(1, time.Now()))
|
||||
err = NewContentModerationRepository(db).CreateLog(context.Background(), &service.ContentModerationLog{EngineMeta: meta})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
mock.ExpectClose()
|
||||
require.NoError(t, db.Close())
|
||||
}
|
||||
}
|
||||
|
||||
func TestContentModerationRepositoryEngineMetaRead(t *testing.T) {
|
||||
for _, meta := range []any{nil, `{"engine":"typesafe","model":"jev-fixed","rules_version":"rules-v1","skipped_images":1}`} {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
mock.ExpectQuery("SELECT COUNT").WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
|
||||
columns := []string{"id", "request_id", "user_id", "user_email", "api_key_id", "api_key_name", "group_id", "group_name", "endpoint", "provider", "model", "mode", "action", "flagged", "highest_category", "highest_score", "category_scores", "threshold_snapshot", "input_excerpt", "upstream_latency_ms", "error", "violation_count", "auto_banned", "email_sent", "status", "queue_delay_ms", "matched_keyword", "created_at", "engine_meta"}
|
||||
mock.ExpectQuery("SELECT[\\s\\S]*l.engine_meta").WillReturnRows(sqlmock.NewRows(columns).AddRow(1, "req", nil, "", nil, "", nil, "", "/v1/responses", "business", "gpt-model", "observe", "allow", false, "sexual", 0.1, `{"sexual":0.1}`, `{"sexual":0.8}`, "sample", 10, "", 0, false, false, "active", nil, "", time.Now(), meta))
|
||||
logs, _, err := NewContentModerationRepository(db).ListLogs(context.Background(), service.ContentModerationLogFilter{})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, logs, 1)
|
||||
require.Equal(t, "gpt-model", logs[0].Model)
|
||||
if meta == nil {
|
||||
require.Nil(t, logs[0].EngineMeta)
|
||||
} else {
|
||||
require.Equal(t, "jev-fixed", logs[0].EngineMeta.Model)
|
||||
require.Equal(t, 1, logs[0].EngineMeta.SkippedImages)
|
||||
}
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
mock.ExpectClose()
|
||||
require.NoError(t, db.Close())
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildContentModerationLogWhere_BlockedIncludesAllBlockActions(t *testing.T) {
|
||||
where, args := buildContentModerationLogWhere(service.ContentModerationLogFilter{Result: "blocked"})
|
||||
|
||||
|
||||
@@ -249,6 +249,18 @@ func TestDashboardAggregationRepositoryCleanupUsageLogsPartitionedSortsAndInvali
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
mock.ExpectCommit()
|
||||
}
|
||||
// The boundary partition is pruned by exact timestamp. tableoid is required
|
||||
// because ctid alone can also identify a retained row in another partition.
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery(`(?s)SELECT tableoid, ctid.*WHERE created_at < \$1.*WHERE \(tableoid, ctid\) IN.*RETURNING created_at`).
|
||||
WithArgs(cutoff, usageLogsCleanupBatchSize).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(cutoff.Add(-time.Hour)))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(cutoff.Add(-time.Hour), "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
@@ -271,7 +283,7 @@ func TestDashboardAggregationRepositoryCleanupUsageLogsNonPartitionedFailureRoll
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery(`(?s)SELECT ctid.*ORDER BY created_at ASC, id ASC.*DELETE FROM usage_logs.*RETURNING created_at`).
|
||||
mock.ExpectQuery(`(?s)SELECT tableoid, ctid.*ORDER BY created_at ASC, id ASC.*DELETE FROM usage_logs.*RETURNING created_at`).
|
||||
WithArgs(cutoff, usageLogsCleanupBatchSize).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(deletedAt))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
|
||||
@@ -239,7 +239,10 @@ func (r *dashboardAggregationRepository) CleanupUsageLogs(ctx context.Context, c
|
||||
if err := r.dropUsageLogsPartitions(ctx, cutoff); err != nil {
|
||||
return err
|
||||
}
|
||||
} else if err := r.cleanupUsageLogsBatches(ctx, cutoff); err != nil {
|
||||
}
|
||||
// Remove the expired part of the boundary partition as well, so a retention
|
||||
// window measured in days does not silently round up to a whole month.
|
||||
if err := r.cleanupUsageLogsBatches(ctx, cutoff); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.SyncGroupUsageRollups(ctx, service.GroupUsageTodayStart(r.now()))
|
||||
@@ -261,14 +264,14 @@ func (r *dashboardAggregationRepository) cleanupUsageLogsBatches(ctx context.Con
|
||||
|
||||
res, err := r.sql.ExecContext(ctx, `
|
||||
WITH victims AS (
|
||||
SELECT ctid
|
||||
SELECT tableoid, ctid
|
||||
FROM usage_logs
|
||||
WHERE created_at < $1
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $2
|
||||
)
|
||||
DELETE FROM usage_logs
|
||||
WHERE ctid IN (SELECT ctid FROM victims)
|
||||
WHERE (tableoid, ctid) IN (SELECT tableoid, ctid FROM victims)
|
||||
`, cutoff.UTC(), usageLogsCleanupBatchSize)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -298,14 +301,14 @@ func cleanupUsageLogsBatchWithRollupInvalidation(ctx context.Context, db *sql.DB
|
||||
}
|
||||
rows, err := tx.QueryContext(ctx, `
|
||||
WITH victims AS (
|
||||
SELECT ctid
|
||||
SELECT tableoid, ctid
|
||||
FROM usage_logs
|
||||
WHERE created_at < $1
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $2
|
||||
)
|
||||
DELETE FROM usage_logs
|
||||
WHERE ctid IN (SELECT ctid FROM victims)
|
||||
WHERE (tableoid, ctid) IN (SELECT tableoid, ctid FROM victims)
|
||||
RETURNING created_at
|
||||
`, cutoff.UTC(), usageLogsCleanupBatchSize)
|
||||
if err != nil {
|
||||
|
||||
@@ -157,21 +157,25 @@ func InitEnt(cfg *config.Config) (*ent.Client, *sql.DB, error) {
|
||||
return nil, nil, fmt.Errorf("validate config after secret bootstrap: %w", err)
|
||||
}
|
||||
|
||||
// SIMPLE 模式:启动时补齐各平台默认分组。
|
||||
// - anthropic/openai/gemini: 确保存在 <platform>-default
|
||||
// - antigravity: 仅要求存在 >=2 个未软删除分组(用于 claude/gemini 混合调度场景)
|
||||
if cfg.RunMode == config.RunModeSimple {
|
||||
seedCtx, seedCancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer seedCancel()
|
||||
if err := ensureSimpleModeDefaultGroups(seedCtx, client); err != nil {
|
||||
_ = client.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
if err := ensureSimpleModeAdminConcurrency(seedCtx, client); err != nil {
|
||||
_ = client.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
seedCtx, seedCancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer seedCancel()
|
||||
if err := ensureSimpleModeStartup(seedCtx, client, cfg); err != nil {
|
||||
_ = client.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
return client, drv.DB(), nil
|
||||
}
|
||||
|
||||
// ensureSimpleModeStartup keeps admin concurrency setup independent of group seeding.
|
||||
func ensureSimpleModeStartup(ctx context.Context, client *ent.Client, cfg *config.Config) error {
|
||||
if cfg.RunMode != config.RunModeSimple {
|
||||
return nil
|
||||
}
|
||||
if cfg.SimpleMode.AutoCreateDefaultGroups {
|
||||
if err := ensureSimpleModeDefaultGroups(ctx, client); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return ensureSimpleModeAdminConcurrency(ctx, client)
|
||||
}
|
||||
|
||||
@@ -76,6 +76,9 @@ const (
|
||||
defaultOpenAIHTTP2FallbackErrorThreshold = 2
|
||||
defaultOpenAIHTTP2FallbackWindow = 60 * time.Second
|
||||
defaultOpenAIHTTP2FallbackTTL = 10 * time.Minute
|
||||
// OpenAI 保留原有的探测与应答期限,避免中转站的延迟 PING 应答被长流策略提前判死。
|
||||
openAIHTTP2ReadIdleTimeout = 15 * time.Second
|
||||
openAIHTTP2PingTimeout = 15 * time.Second
|
||||
// 长流 HTTP/2 连接健康探测:池化连接被代理/NAT
|
||||
// 静默掐断会成为“死连接”(两端都以为存活),请求落上去会挂到 TCP 重传超时
|
||||
// (分钟级)。Go 的 http2.Transport 默认 ReadIdleTimeout=0(不发健康 PING),
|
||||
@@ -216,7 +219,7 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i
|
||||
// 执行请求
|
||||
client := s.httpClientForUpstreamRequest(entry.client, req)
|
||||
client = httpClientWithGrokAccessDeniedFallback(client)
|
||||
resp, err := servertiming.Do(client, req)
|
||||
resp, err := doUpstreamRequest(client, req)
|
||||
if err != nil {
|
||||
s.recordOpenAIHTTP2Failure(profile, entry.protocolMode, entry.proxyKey, err)
|
||||
// 请求失败,立即减少计数
|
||||
@@ -226,9 +229,6 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i
|
||||
}
|
||||
s.recordOpenAIHTTP2Success(profile, entry.protocolMode, entry.proxyKey)
|
||||
|
||||
// 如果上游返回了压缩内容,解压后再交给业务层
|
||||
decompressResponseBody(resp)
|
||||
|
||||
// 包装响应体,在关闭时自动减少计数并更新时间戳
|
||||
// 这确保了流式响应(如 SSE)在完全读取前不会被淘汰
|
||||
resp.Body = wrapTrackedBody(resp.Body, func() {
|
||||
@@ -280,7 +280,7 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco
|
||||
|
||||
client := s.httpClientForUpstreamRequest(entry.client, req)
|
||||
client = httpClientWithGrokAccessDeniedFallback(client)
|
||||
resp, err := servertiming.Do(client, req)
|
||||
resp, err := doUpstreamRequest(client, req)
|
||||
if err != nil {
|
||||
atomic.AddInt64(&entry.inFlight, -1)
|
||||
atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
|
||||
@@ -288,8 +288,6 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco
|
||||
return nil, err
|
||||
}
|
||||
|
||||
decompressResponseBody(resp)
|
||||
|
||||
resp.Body = wrapTrackedBody(resp.Body, func() {
|
||||
atomic.AddInt64(&entry.inFlight, -1)
|
||||
atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
|
||||
@@ -298,6 +296,56 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// doUpstreamRequest owns cancellation for one attempt, without cancelling the
|
||||
// caller's context (which may be detached for billing or reused for retries).
|
||||
func doUpstreamRequest(client *http.Client, req *http.Request) (*http.Response, error) {
|
||||
ctx, cancel := context.WithCancel(req.Context())
|
||||
resp, err := servertiming.Do(client, req.WithContext(ctx))
|
||||
if err != nil {
|
||||
cancel()
|
||||
return resp, err
|
||||
}
|
||||
decompressResponseBody(resp)
|
||||
resp.Body = &cancelOnCloseBody{ReadCloser: resp.Body, cancel: cancel}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
type cancelOnCloseBody struct {
|
||||
io.ReadCloser
|
||||
cancel context.CancelFunc
|
||||
once sync.Once
|
||||
err error
|
||||
readMu sync.Mutex
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (b *cancelOnCloseBody) Read(p []byte) (int, error) {
|
||||
b.readMu.Lock()
|
||||
defer b.readMu.Unlock()
|
||||
if b.closed {
|
||||
return 0, http.ErrBodyReadAfterClose
|
||||
}
|
||||
return b.ReadCloser.Read(p)
|
||||
}
|
||||
|
||||
func (b *cancelOnCloseBody) Close() error {
|
||||
b.once.Do(func() {
|
||||
// Cancel before closing, including before closing a decompressor. In Go
|
||||
// 1.27 an early HTTP/1 close concurrent with Read can otherwise leave an
|
||||
// EOF waiter on a reused connection and stall subsequent responses.
|
||||
// Fully consumed responses have already released their transport request,
|
||||
// so cancelling here preserves normal keep-alive reuse.
|
||||
b.cancel()
|
||||
// Wait for an active read to observe cancellation before touching the
|
||||
// body or decompressor. Do not hold readMu while cancelling the request.
|
||||
b.readMu.Lock()
|
||||
defer b.readMu.Unlock()
|
||||
b.closed = true
|
||||
b.err = b.ReadCloser.Close()
|
||||
})
|
||||
return b.err
|
||||
}
|
||||
|
||||
// httpClientForUpstreamRequest 按请求上下文的标记派生客户端:禁用重定向,或对重定向的每一跳做主机校验。
|
||||
// 派生的克隆与缓存客户端共享 Transport;未打标记时原样返回。
|
||||
func (s *httpUpstreamService) httpClientForUpstreamRequest(client *http.Client, req *http.Request) *http.Client {
|
||||
@@ -1340,7 +1388,7 @@ func buildUpstreamTransport(settings poolSettings, proxyURL *url.URL, protocolMo
|
||||
transport.ForceAttemptHTTP2 = true
|
||||
// 显式配置 http2 并启用 PING 健康探测,剔除代理/NAT 静默掐断的死连接,
|
||||
// 避免请求挂在死连接上直到 TCP 重传超时(分钟级)。
|
||||
if _, err := enableHTTP2KeepAlive(transport); err != nil {
|
||||
if _, err := enableHTTP2KeepAlive(transport, protocolMode); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
case upstreamProtocolModeOpenAIH1:
|
||||
@@ -1361,7 +1409,7 @@ func buildUpstreamTransport(settings poolSettings, proxyURL *url.URL, protocolMo
|
||||
// Go 默认惰性配置 http2 且 ReadIdleTimeout=0(不发健康 PING),无法检测被代理/NAT
|
||||
// 静默掐断的死连接。此处主动设置 ReadIdleTimeout/PingTimeout,让死连接被提前 PING
|
||||
// 出并关闭,请求得以重建连接而非挂到 TCP 重传超时。返回底层 *http2.Transport 便于测试。
|
||||
func enableHTTP2KeepAlive(transport *http.Transport) (*http2.Transport, error) {
|
||||
func enableHTTP2KeepAlive(transport *http.Transport, protocolMode string) (*http2.Transport, error) {
|
||||
h2, err := http2.ConfigureTransports(transport)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -1369,6 +1417,10 @@ func enableHTTP2KeepAlive(transport *http.Transport) (*http2.Transport, error) {
|
||||
if h2 != nil {
|
||||
h2.ReadIdleTimeout = longStreamHTTP2ReadIdleTimeout
|
||||
h2.PingTimeout = longStreamHTTP2PingTimeout
|
||||
if protocolMode == upstreamProtocolModeOpenAIH2 {
|
||||
h2.ReadIdleTimeout = openAIHTTP2ReadIdleTimeout
|
||||
h2.PingTimeout = openAIHTTP2PingTimeout
|
||||
}
|
||||
}
|
||||
return h2, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type lifecycleUsageLogRepo struct {
|
||||
service.UsageLogRepository
|
||||
logs []*service.UsageLog
|
||||
}
|
||||
|
||||
func (r *lifecycleUsageLogRepo) Create(ctx context.Context, log *service.UsageLog) (bool, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
r.logs = append(r.logs, log)
|
||||
return true, nil
|
||||
}
|
||||
|
||||
type lifecycleBillingRepo struct {
|
||||
service.UsageBillingRepository
|
||||
commands []*service.UsageBillingCommand
|
||||
}
|
||||
|
||||
func (r *lifecycleBillingRepo) Apply(ctx context.Context, cmd *service.UsageBillingCommand) (*service.UsageBillingApplyResult, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
r.commands = append(r.commands, cmd)
|
||||
return &service.UsageBillingApplyResult{Applied: true}, nil
|
||||
}
|
||||
|
||||
type lifecycleDisconnectedWriter struct {
|
||||
gin.ResponseWriter
|
||||
disconnect context.CancelFunc
|
||||
}
|
||||
|
||||
func (w *lifecycleDisconnectedWriter) Write([]byte) (int, error) {
|
||||
w.disconnect()
|
||||
return 0, errors.New("client disconnected")
|
||||
}
|
||||
|
||||
// Exercise the real forwarding and RecordUsage services with the real HTTP
|
||||
// upstream. Only the persistence boundary is replaced: closing an attempt must
|
||||
// not cancel the detached stream before its final usage reaches billing.
|
||||
func TestHTTPUpstreamForwardDrainsUsageAfterClientDisconnect(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
clientCtx, disconnect := context.WithCancel(t.Context())
|
||||
defer disconnect()
|
||||
releaseUsage := make(chan struct{})
|
||||
var releaseOnce sync.Once
|
||||
release := func() { releaseOnce.Do(func() { close(releaseUsage) }) }
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("X-Request-Id", "lifecycle-billing")
|
||||
_, _ = io.WriteString(w, "data: {\"type\":\"response.output_text.delta\",\"delta\":\"hello\"}\n\n")
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
select {
|
||||
case <-releaseUsage:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
_, _ = io.WriteString(w, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_lifecycle\",\"status\":\"completed\",\"usage\":{\"input_tokens\":17,\"output_tokens\":8,\"total_tokens\":25}}}\n\n")
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
defer release()
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.Default.RateMultiplier = 1
|
||||
cfg.Security.URLAllowlist.AllowPrivateHosts = true
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
|
||||
cfg.Gateway.StreamKeepaliveInterval = 1
|
||||
upstream := NewHTTPUpstream(cfg)
|
||||
usageRepo := &lifecycleUsageLogRepo{}
|
||||
billingRepo := &lifecycleBillingRepo{}
|
||||
svc := service.NewOpenAIGatewayService(
|
||||
nil, usageRepo, billingRepo, nil, nil, nil, nil, cfg, nil, nil,
|
||||
service.NewBillingService(cfg, nil), nil, nil, upstream,
|
||||
&service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
account := &service.Account{
|
||||
ID: 1, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{"base_url": srv.URL, "api_key": "test-key"},
|
||||
}
|
||||
body := []byte(`{"model":"gpt-5.1","stream":true,"input":"hello"}`)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(string(body))).WithContext(clientCtx)
|
||||
c.Writer = &lifecycleDisconnectedWriter{ResponseWriter: c.Writer, disconnect: disconnect}
|
||||
|
||||
type forwardResult struct {
|
||||
result *service.OpenAIForwardResult
|
||||
err error
|
||||
}
|
||||
done := make(chan forwardResult, 1)
|
||||
go func() {
|
||||
result, err := svc.Forward(clientCtx, c, account, body)
|
||||
done <- forwardResult{result, err}
|
||||
}()
|
||||
select {
|
||||
case <-clientCtx.Done():
|
||||
case got := <-done:
|
||||
t.Fatalf("forward returned before downstream disconnect: %v", got.err)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("downstream did not receive the first event")
|
||||
}
|
||||
release()
|
||||
select {
|
||||
case got := <-done:
|
||||
require.NoError(t, got.err)
|
||||
require.NotNil(t, got.result)
|
||||
require.Equal(t, 17, got.result.Usage.InputTokens)
|
||||
require.Equal(t, 8, got.result.Usage.OutputTokens)
|
||||
// RecordUsage must still work with the cancelled downstream context.
|
||||
require.NoError(t, svc.RecordUsage(clientCtx, &service.OpenAIRecordUsageInput{
|
||||
Result: got.result, APIKey: &service.APIKey{ID: 2},
|
||||
User: &service.User{ID: 3}, Account: account,
|
||||
}))
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("forward did not finish draining usage")
|
||||
}
|
||||
require.Len(t, billingRepo.commands, 1)
|
||||
require.Len(t, usageRepo.logs, 1)
|
||||
require.Equal(t, 17, usageRepo.logs[0].InputTokens)
|
||||
require.Equal(t, 8, usageRepo.logs[0].OutputTokens)
|
||||
require.Greater(t, usageRepo.logs[0].ActualCost, 0.0)
|
||||
s, ok := upstream.(*httpUpstreamService)
|
||||
require.True(t, ok)
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
for _, entry := range s.clients {
|
||||
entry.client.CloseIdleConnections()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,377 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/http/httptrace"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// Exercise both public entry points. The TLS case uses the test server's trusted
|
||||
// transport so the test remains local and does not depend on system trust roots.
|
||||
func lifecycleClient(t *testing.T, srv *httptest.Server, withTLS bool) (*httpUpstreamService, func(*http.Request) (*http.Response, error)) {
|
||||
t.Helper()
|
||||
s, ok := NewHTTPUpstream(nil).(*httpUpstreamService)
|
||||
require.True(t, ok)
|
||||
profile := &tlsfingerprint.Profile{Name: "body-lifecycle-test"}
|
||||
var entry *upstreamClientEntry
|
||||
var err error
|
||||
if withTLS {
|
||||
entry, err = s.getClientEntryWithTLS("", 1, 1, profile, service.HTTPUpstreamProfileDefault, false, true)
|
||||
} else {
|
||||
entry, err = s.acquireClientWithProfile("", 1, 1, service.HTTPUpstreamProfileDefault)
|
||||
if err == nil {
|
||||
atomic.AddInt64(&entry.inFlight, -1)
|
||||
}
|
||||
}
|
||||
require.NoError(t, err)
|
||||
baseTransport, ok := srv.Client().Transport.(*http.Transport)
|
||||
require.True(t, ok)
|
||||
tr := baseTransport.Clone()
|
||||
tr.MaxConnsPerHost = 1
|
||||
tr.MaxIdleConnsPerHost = 1
|
||||
tr.IdleConnTimeout = 10 * time.Second
|
||||
entry.client = &http.Client{Transport: tr}
|
||||
t.Cleanup(tr.CloseIdleConnections)
|
||||
return s, func(req *http.Request) (*http.Response, error) {
|
||||
if withTLS {
|
||||
return s.DoWithTLS(req, "", 1, 1, profile)
|
||||
}
|
||||
return s.Do(req, "", 1, 1)
|
||||
}
|
||||
}
|
||||
|
||||
func requireNoUpstreamInFlight(t *testing.T, s *httpUpstreamService) {
|
||||
t.Helper()
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
for _, entry := range s.clients {
|
||||
require.Zero(t, atomic.LoadInt64(&entry.inFlight))
|
||||
}
|
||||
}
|
||||
|
||||
type notifyReadCloser struct {
|
||||
io.ReadCloser
|
||||
started chan<- struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (r *notifyReadCloser) Read(p []byte) (int, error) {
|
||||
r.once.Do(func() { close(r.started) })
|
||||
return r.ReadCloser.Read(p)
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamConcurrentEarlyCloseDoesNotPoisonNextResponse(t *testing.T) {
|
||||
for _, withTLS := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("tls=%v", withTLS), func(t *testing.T) {
|
||||
finishFirst := make(chan struct{})
|
||||
firstFinished := make(chan struct{})
|
||||
var requests atomic.Int32
|
||||
srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = io.WriteString(w, "data: begin\n\n")
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
if requests.Add(1) == 1 {
|
||||
defer close(firstFinished)
|
||||
select {
|
||||
case <-finishFirst:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
_, _ = io.WriteString(w, "data: complete\n\n")
|
||||
}))
|
||||
if withTLS {
|
||||
srv.StartTLS()
|
||||
} else {
|
||||
srv.Start()
|
||||
}
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
s, do := lifecycleClient(t, srv, withTLS)
|
||||
ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
|
||||
require.NoError(t, err)
|
||||
resp, err := do(req)
|
||||
require.NoError(t, err)
|
||||
prefix := make([]byte, len("data: begin\n\n"))
|
||||
_, err = io.ReadFull(resp.Body, prefix)
|
||||
require.NoError(t, err)
|
||||
readerDone := make(chan struct{})
|
||||
readStarted := make(chan struct{})
|
||||
resp.Body = ¬ifyReadCloser{ReadCloser: resp.Body, started: readStarted}
|
||||
go func() { _, _ = io.Copy(io.Discard, resp.Body); close(readerDone) }()
|
||||
// Start Close only after the background reader has entered Read;
|
||||
// the server is still holding the response open at this point.
|
||||
<-readStarted
|
||||
closed := make(chan error, 1)
|
||||
go func() { closed <- resp.Body.Close() }()
|
||||
select {
|
||||
case err := <-closed:
|
||||
require.NoError(t, err)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("Close blocked behind the active reader")
|
||||
}
|
||||
// Let the early-close drain finish only after Close has returned.
|
||||
// Go 1.27 can reuse that connection while its old reader awaits EOF.
|
||||
close(finishFirst)
|
||||
<-firstFinished
|
||||
for i := 0; i < 8; i++ {
|
||||
resp, err := do(req)
|
||||
require.NoError(t, err)
|
||||
done := make(chan error, 1)
|
||||
go func() { _, err := io.Copy(io.Discard, resp.Body); _ = resp.Body.Close(); done <- err }()
|
||||
select {
|
||||
case err := <-done:
|
||||
require.NoError(t, err)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("response EOF waits for a subsequent request")
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-readerDone:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("early-closed reader leaked")
|
||||
}
|
||||
require.NoError(t, req.Context().Err(), "closing an attempt must not cancel its caller or retries")
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamCompletedBodyPreservesKeepAlive(t *testing.T) {
|
||||
for _, encoding := range []string{"", "gzip"} {
|
||||
t.Run("encoding="+encoding, func(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if encoding == "gzip" {
|
||||
w.Header().Set("Content-Encoding", encoding)
|
||||
zw := gzip.NewWriter(w)
|
||||
_, _ = io.WriteString(zw, "complete response")
|
||||
_ = zw.Close()
|
||||
return
|
||||
}
|
||||
_, _ = io.WriteString(w, "complete response")
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
s, do := lifecycleClient(t, srv, false)
|
||||
for i := 0; i < 20; i++ {
|
||||
var reused bool
|
||||
ctx := httptrace.WithClientTrace(t.Context(), &httptrace.ClientTrace{GotConn: func(info httptrace.GotConnInfo) { reused = info.Reused }})
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
|
||||
require.NoError(t, err)
|
||||
req.Header.Set("Accept-Encoding", encoding)
|
||||
resp, err := do(req)
|
||||
require.NoError(t, err)
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "complete response", string(body))
|
||||
var wg sync.WaitGroup
|
||||
for j := 0; j < 4; j++ {
|
||||
wg.Add(1)
|
||||
go func() { defer wg.Done(); _ = resp.Body.Close() }()
|
||||
}
|
||||
wg.Wait()
|
||||
if i > 0 {
|
||||
require.True(t, reused, "fully consumed responses must keep their connection")
|
||||
}
|
||||
}
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamCancellationReleasesAttempt(t *testing.T) {
|
||||
for _, beforeHeaders := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("before_headers=%v", beforeHeaders), func(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if !beforeHeaders {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
}
|
||||
<-r.Context().Done()
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
s, do := lifecycleClient(t, srv, false)
|
||||
ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
|
||||
require.NoError(t, err)
|
||||
resp, err := do(req)
|
||||
if beforeHeaders {
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
_, err = io.ReadAll(resp.Body)
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
require.NoError(t, resp.Body.Close())
|
||||
}
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamCompressedEarlyCloseStopsReader(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Encoding", "gzip")
|
||||
zw := gzip.NewWriter(w)
|
||||
_, _ = io.WriteString(zw, "data: begin\n\n")
|
||||
_ = zw.Flush()
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
<-r.Context().Done()
|
||||
_ = zw.Close()
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
s, do := lifecycleClient(t, srv, false)
|
||||
req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, srv.URL, nil)
|
||||
require.NoError(t, err)
|
||||
req.Header.Set("Accept-Encoding", "gzip")
|
||||
resp, err := do(req)
|
||||
require.NoError(t, err)
|
||||
prefix := make([]byte, len("data: begin\n\n"))
|
||||
_, err = io.ReadFull(resp.Body, prefix)
|
||||
require.NoError(t, err)
|
||||
done := make(chan error, 1)
|
||||
readStarted := make(chan struct{})
|
||||
resp.Body = ¬ifyReadCloser{ReadCloser: resp.Body, started: readStarted}
|
||||
go func() { _, err := io.Copy(io.Discard, resp.Body); done <- err }()
|
||||
<-readStarted
|
||||
closed := make(chan error, 1)
|
||||
go func() { closed <- resp.Body.Close() }()
|
||||
select {
|
||||
case err := <-closed:
|
||||
require.NoError(t, err)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("compressed body Close blocked")
|
||||
}
|
||||
select {
|
||||
case err := <-done:
|
||||
require.Error(t, err)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("compressed body reader leaked")
|
||||
}
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamHTTP2CloseDoesNotCancelOtherStream(t *testing.T) {
|
||||
srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/held" {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
<-r.Context().Done()
|
||||
return
|
||||
}
|
||||
_, _ = io.WriteString(w, "complete response")
|
||||
}))
|
||||
srv.EnableHTTP2 = true
|
||||
srv.StartTLS()
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
s, do := lifecycleClient(t, srv, false)
|
||||
req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, srv.URL+"/held", nil)
|
||||
require.NoError(t, err)
|
||||
held, err := do(req)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, held.ProtoMajor)
|
||||
for i := 0; i < 2; i++ {
|
||||
var reused bool
|
||||
ctx := httptrace.WithClientTrace(t.Context(), &httptrace.ClientTrace{GotConn: func(info httptrace.GotConnInfo) { reused = info.Reused }})
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
|
||||
require.NoError(t, err)
|
||||
resp, err := do(req)
|
||||
require.NoError(t, err)
|
||||
if i == 0 {
|
||||
require.NoError(t, held.Body.Close())
|
||||
}
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "complete response", string(body))
|
||||
require.NoError(t, resp.Body.Close())
|
||||
require.True(t, reused, "closing one HTTP/2 stream must retain the shared connection")
|
||||
}
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamDetachedRequestDrainsAfterClientDisconnect(t *testing.T) {
|
||||
disconnectCtx, disconnect := context.WithCancel(t.Context())
|
||||
defer disconnect()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.WriteString(w, "begin ")
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
<-disconnectCtx.Done()
|
||||
_, _ = io.WriteString(w, "usage and completion")
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
s, do := lifecycleClient(t, srv, false)
|
||||
req, err := http.NewRequestWithContext(context.WithoutCancel(disconnectCtx), http.MethodGet, srv.URL, nil)
|
||||
require.NoError(t, err)
|
||||
resp, err := do(req)
|
||||
require.NoError(t, err)
|
||||
disconnect()
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "begin usage and completion", string(body))
|
||||
require.NoError(t, resp.Body.Close())
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
}
|
||||
|
||||
// Keep the same benchmark runnable against the base revision: it measures a
|
||||
// complete public Do/read/Close lifecycle without network latency hiding the
|
||||
// per-attempt context and per-Read synchronization costs.
|
||||
func BenchmarkHTTPUpstreamBodyLifecycle(b *testing.B) {
|
||||
for _, size := range []int{256, 64 * 1024} {
|
||||
b.Run(fmt.Sprintf("bytes=%d", size), func(b *testing.B) {
|
||||
s, ok := NewHTTPUpstream(nil).(*httpUpstreamService)
|
||||
require.True(b, ok)
|
||||
entry, err := s.getOrCreateClient("", 1, 1)
|
||||
require.NoError(b, err)
|
||||
payload := strings.Repeat("x", size)
|
||||
entry.client = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK, Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(payload)),
|
||||
}, nil
|
||||
})}
|
||||
req, err := http.NewRequestWithContext(b.Context(), http.MethodGet, "https://example.com", nil)
|
||||
require.NoError(b, err)
|
||||
buf := make([]byte, 256)
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(int64(size))
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
resp, err := s.Do(req, "", 1, 1)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
for {
|
||||
_, err = resp.Body.Read(buf)
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
if err != io.EOF {
|
||||
b.Fatal(err)
|
||||
}
|
||||
if err := resp.Body.Close(); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user