diff --git a/.env.example b/.env.example index a45f46cd..1ced0b8a 100644 --- a/.env.example +++ b/.env.example @@ -6,7 +6,7 @@ HOST=127.0.0.1 BIND_ADDRESS=127.0.0.1 # Subscription browser OAuth uses provider-fixed loopback callbacks: -# Codex localhost:1455 and Claude localhost:54545. +# Codex localhost:1455, Claude localhost:54545, and Antigravity localhost:51121. OAUTH_CALLBACK_BIND_ADDRESS=127.0.0.1 PORT=3001 GRACEFUL_SHUTDOWN_TIMEOUT=10 diff --git a/Dockerfile b/Dockerfile index 4f28b36e..e5f411fc 100644 --- a/Dockerfile +++ b/Dockerfile @@ -52,6 +52,6 @@ COPY --from=go-builder /build/gpt-load . COPY LICENSE THIRD_PARTY_NOTICES.md /app/licenses/ COPY LICENSES/Apache-2.0.txt /app/licenses/Apache-2.0.txt COPY LICENSES/MIT.txt /app/licenses/MIT.txt -EXPOSE 3001 1455 54545 +EXPOSE 3001 1455 54545 51121 USER 10001:10001 ENTRYPOINT ["/app/gpt-load"] diff --git a/README.md b/README.md index e0db8f82..450e223f 100644 --- a/README.md +++ b/README.md @@ -94,7 +94,7 @@ docker compose exec gpt-load sh -c 'cat /app/data/auth.key' The named volume preserves the default SQLite database, `auth.key`, and `encryption.key`. Production deployments should inject explicit `AUTH_KEY` and `ENCRYPTION_KEY` values through protected secret handling. Never commit real secrets to `.env`, logs, or issues. A MySQL or PostgreSQL `DATABASE_DSN` is operator-managed and should be supplied through the deployment secret/configuration system. -Compose listens on all interfaces only inside the container while publishing to host loopback by default. It also publishes the provider-fixed OAuth callbacks for Codex on `127.0.0.1:1455` and Claude on `127.0.0.1:54545`; keep both on loopback unless the network boundary is explicitly controlled. Because these callback ports are fixed by their OAuth clients, only one default Compose instance on a host can expose them at a time. In Docker, SSH, or remote-browser deployments where browser `localhost` does not reach GPT-Load, copy the complete callback URL into the authorization dialog. Setting `BIND_ADDRESS=0.0.0.0`, `OAUTH_CALLBACK_BIND_ADDRESS=0.0.0.0`, or running a native binary with `HOST=0.0.0.0`, is an explicit opt-in. In production, expose either only behind a controlled network boundary with a TLS reverse proxy and ACL/firewall controls. +Compose listens on all interfaces only inside the container while publishing to host loopback by default. It also publishes the provider-fixed OAuth callbacks for Codex on `127.0.0.1:1455`, Claude on `127.0.0.1:54545`, and Antigravity on `127.0.0.1:51121`; keep all of them on loopback unless the network boundary is explicitly controlled. Because these callback ports are fixed by their OAuth clients, only one default Compose instance on a host can expose them at a time. In Docker, SSH, or remote-browser deployments where browser `localhost` does not reach GPT-Load, copy the complete callback URL into the authorization dialog. Setting `BIND_ADDRESS=0.0.0.0`, `OAUTH_CALLBACK_BIND_ADDRESS=0.0.0.0`, or running a native binary with `HOST=0.0.0.0`, is an explicit opt-in. In production, expose either only behind a controlled network boundary with a TLS reverse proxy and ACL/firewall controls. ### Native binary @@ -214,7 +214,7 @@ Usage/Cost quality boundaries: |---|---|---| | `HOST` | `127.0.0.1` | Native HTTP listen address; `0.0.0.0` is an explicit opt-in. The release container overrides this to `0.0.0.0` internally | | `BIND_ADDRESS` | `127.0.0.1` | Compose host-side publish address; not a process setting | -| `OAUTH_CALLBACK_BIND_ADDRESS` | `127.0.0.1` | Compose host-side address for the fixed Codex (`1455`) and Claude (`54545`) OAuth callback ports | +| `OAUTH_CALLBACK_BIND_ADDRESS` | `127.0.0.1` | Compose host-side address for the fixed Codex (`1455`), Claude (`54545`), and Antigravity (`51121`) OAuth callback ports | | `PORT` | `3001` | HTTP listen port | | `DATA_DIR` | `./data` | Native persistent directory; the container overrides it to `/app/data` | | `DATABASE_DSN` | empty → `${DATA_DIR}/gpt-load.db` | Empty selects managed SQLite. Non-empty values use one of `sqlite`, `mysql`, or `postgres` URL forms and are operator-managed | diff --git a/README_CN.md b/README_CN.md index d98f434c..6a7845d3 100644 --- a/README_CN.md +++ b/README_CN.md @@ -93,7 +93,7 @@ docker compose exec gpt-load sh -c 'cat /app/data/auth.key' 默认 named volume 会保存 SQLite、`auth.key` 和 `encryption.key`。生产环境建议通过受保护的 secret 注入显式 `AUTH_KEY` 与 `ENCRYPTION_KEY`；不要把真实 secret 提交到 `.env`、日志或 issue。MySQL 或 PostgreSQL 的 `DATABASE_DSN` 属于 operator 管理配置，应通过部署 secret/configuration system 提供。 -Compose 仅在容器内部监听所有接口，默认仍只发布到宿主 loopback；同时会把 Codex 与 Claude OAuth 客户端固定的回调端口分别发布到 `127.0.0.1:1455` 和 `127.0.0.1:54545`。除非已有明确的受控网络边界，否则必须保持回环地址。由于端口由 OAuth 客户端固定，同一宿主机同一时刻只能有一个默认 Compose 实例占用它们。Docker、SSH 或远程浏览器场景中，如果浏览器的 `localhost` 无法到达 GPT-Load，可将完整回调 URL 复制到授权弹窗完成流程。设置 `BIND_ADDRESS=0.0.0.0`、`OAUTH_CALLBACK_BIND_ADDRESS=0.0.0.0`，或为原生进程设置 `HOST=0.0.0.0`，都属于显式 opt-in；生产环境只能在受控网络边界、TLS reverse proxy 及 ACL/firewall 保护下暴露。 +Compose 仅在容器内部监听所有接口，默认仍只发布到宿主 loopback；同时会把 Codex、Claude 与 Antigravity OAuth 客户端固定的回调端口分别发布到 `127.0.0.1:1455`、`127.0.0.1:54545` 和 `127.0.0.1:51121`。除非已有明确的受控网络边界，否则必须保持回环地址。由于端口由 OAuth 客户端固定，同一宿主机同一时刻只能有一个默认 Compose 实例占用它们。Docker、SSH 或远程浏览器场景中，如果浏览器的 `localhost` 无法到达 GPT-Load，可将完整回调 URL 复制到授权弹窗完成流程。设置 `BIND_ADDRESS=0.0.0.0`、`OAUTH_CALLBACK_BIND_ADDRESS=0.0.0.0`，或为原生进程设置 `HOST=0.0.0.0`，都属于显式 opt-in；生产环境只能在受控网络边界、TLS reverse proxy 及 ACL/firewall 保护下暴露。 ### 原生二进制 @@ -213,7 +213,7 @@ Usage/Cost 质量边界： |---|---|---| | `HOST` | `127.0.0.1` | 原生 HTTP 监听地址；`0.0.0.0` 是显式 opt-in，发布容器仅在内部覆盖为 `0.0.0.0` | | `BIND_ADDRESS` | `127.0.0.1` | Compose 宿主机侧发布地址，不是进程配置 | -| `OAUTH_CALLBACK_BIND_ADDRESS` | `127.0.0.1` | Compose 中固定 Codex（`1455`）与 Claude（`54545`）OAuth 回调端口的宿主机发布地址 | +| `OAUTH_CALLBACK_BIND_ADDRESS` | `127.0.0.1` | Compose 中固定 Codex（`1455`）、Claude（`54545`）与 Antigravity（`51121`）OAuth 回调端口的宿主机发布地址 | | `PORT` | `3001` | HTTP 监听端口 | | `DATA_DIR` | `./data` | 原生持久目录；容器内覆盖为 `/app/data` | | `DATABASE_DSN` | 空 → `${DATA_DIR}/gpt-load.db` | 空值选择 managed SQLite；非空值必须使用统一的 `sqlite`、`mysql` 或 `postgres` URL，并由 operator 管理 | diff --git a/README_JP.md b/README_JP.md index 5f1ad96e..cea18ab8 100644 --- a/README_JP.md +++ b/README_JP.md @@ -94,7 +94,7 @@ docker compose exec gpt-load sh -c 'cat /app/data/auth.key' デフォルトのnamed volumeにはSQLite、`auth.key`、`encryption.key`が保持されます。本番環境では、保護されたsecret処理を通じて明示的な`AUTH_KEY`と`ENCRYPTION_KEY`を注入してください。実際のsecretを`.env`、ログ、issueへコミットしないでください。MySQLまたはPostgreSQLの`DATABASE_DSN`はoperator管理の設定として、デプロイのsecret/configuration systemから渡してください。 -Composeはコンテナ内でのみ全インターフェースをlistenし、ホスト側はデフォルトでloopbackにだけ公開します。OAuthクライアント固定のcallback portも、Codexは`127.0.0.1:1455`、Claudeは`127.0.0.1:54545`に公開します。これらのportは固定されているため、同一ホストで公開できるデフォルトComposeインスタンスは同時に1つだけです。Docker、SSH、またはリモートブラウザー環境でブラウザーの`localhost`からGPT-Loadへ到達できない場合は、完全なcallback URLを認証ダイアログに貼り付けて完了できます。`BIND_ADDRESS=0.0.0.0`、`OAUTH_CALLBACK_BIND_ADDRESS=0.0.0.0`、またはネイティブプロセスの`HOST=0.0.0.0`は明示的なopt-inです。本番では、TLS reverse proxyとACL/firewallを備えた管理下のネットワーク境界の内側でのみ公開してください。 +Composeはコンテナ内でのみ全インターフェースをlistenし、ホスト側はデフォルトでloopbackにだけ公開します。OAuthクライアント固定のcallback portも、Codexは`127.0.0.1:1455`、Claudeは`127.0.0.1:54545`、Antigravityは`127.0.0.1:51121`に公開します。これらのportは固定されているため、同一ホストで公開できるデフォルトComposeインスタンスは同時に1つだけです。Docker、SSH、またはリモートブラウザー環境でブラウザーの`localhost`からGPT-Loadへ到達できない場合は、完全なcallback URLを認証ダイアログに貼り付けて完了できます。`BIND_ADDRESS=0.0.0.0`、`OAUTH_CALLBACK_BIND_ADDRESS=0.0.0.0`、またはネイティブプロセスの`HOST=0.0.0.0`は明示的なopt-inです。本番では、TLS reverse proxyとACL/firewallを備えた管理下のネットワーク境界の内側でのみ公開してください。 ### ネイティブバイナリ @@ -211,7 +211,7 @@ Usage/Costの品質境界： |---|---|---| | `HOST` | `127.0.0.1` | ネイティブHTTPリッスンアドレス。`0.0.0.0`は明示的なopt-inで、リリースコンテナは内部だけ`0.0.0.0`に上書き | | `BIND_ADDRESS` | `127.0.0.1` | Composeのホスト側公開アドレス。プロセス設定ではない | -| `OAUTH_CALLBACK_BIND_ADDRESS` | `127.0.0.1` | Codex（`1455`）とClaude（`54545`）の固定OAuth callback portを公開するComposeホスト側アドレス | +| `OAUTH_CALLBACK_BIND_ADDRESS` | `127.0.0.1` | Codex（`1455`）、Claude（`54545`）、Antigravity（`51121`）の固定OAuth callback portを公開するComposeホスト側アドレス | | `PORT` | `3001` | HTTPリッスンポート | | `DATA_DIR` | `./data` | ネイティブの永続ディレクトリ。コンテナ内では`/app/data`に上書き | | `DATABASE_DSN` | 空 → `${DATA_DIR}/gpt-load.db` | 空ならmanaged SQLiteを選択。非空値は統一された`sqlite`、`mysql`、`postgres` URLで指定し、operatorが管理 | diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 2ad7c1fe..991187ad 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -21,12 +21,12 @@ The complete Apache License 2.0 text is distributed in ## CLIProxyAPI - Module: `github.com/router-for-me/CLIProxyAPI/v7` -- Version: `v7.2.133` +- Version: `v7.2.135` - Copyright: 2025-2005.9 Luis Pater; 2025.9-present Router-For.ME - License: MIT License GPT-Load uses a pinned, execution-only embedded adapter around CLIProxyAPI's -Codex and Claude OAuth and HTTP executor code. GPT-Load retains ownership of +Codex, Claude, and Antigravity OAuth and HTTP executor code. GPT-Load retains ownership of credential storage, account selection, retry, health, affinity, logging, and usage policy; the embedded adapter does not use CLIProxyAPI's manager, pool, file store, WebSocket executor, fallback, or automatic retry. diff --git a/docker-compose.yml b/docker-compose.yml index bcfef960..58b7d703 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -8,6 +8,7 @@ services: - "${BIND_ADDRESS:-127.0.0.1}:${PORT:-3001}:${PORT:-3001}" - "${OAUTH_CALLBACK_BIND_ADDRESS:-127.0.0.1}:1455:1455" - "${OAUTH_CALLBACK_BIND_ADDRESS:-127.0.0.1}:54545:54545" + - "${OAUTH_CALLBACK_BIND_ADDRESS:-127.0.0.1}:51121:51121" env_file: - .env environment: diff --git a/go.mod b/go.mod index 617ef4c5..151d7af5 100644 --- a/go.mod +++ b/go.mod @@ -26,7 +26,7 @@ require ( github.com/pierrec/xxHash v0.1.5 // indirect github.com/redis/go-redis/v9 v9.19.0 // indirect github.com/refraction-networking/utls v1.8.2 // indirect - github.com/router-for-me/CLIProxyAPI/v7 v7.2.133 // indirect + github.com/router-for-me/CLIProxyAPI/v7 v7.2.135 // indirect github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 // indirect github.com/tiktoken-go/tokenizer v0.8.1 // indirect go.uber.org/atomic v1.11.0 // indirect diff --git a/go.sum b/go.sum index 1a19da79..4edbfe5d 100644 --- a/go.sum +++ b/go.sum @@ -207,8 +207,8 @@ github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94 github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/rogpeppe/go-internal v1.15.0 h1:D0RCU5rMAp+SpgkiNdrjfJ+LX4J1M32V2NeCY7EJ6hc= github.com/rogpeppe/go-internal v1.15.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs= -github.com/router-for-me/CLIProxyAPI/v7 v7.2.133 h1:TeFrEctzhSuPRNdYkiAFpFj4W1IICcvN+ot9HVnaEgw= -github.com/router-for-me/CLIProxyAPI/v7 v7.2.133/go.mod h1:lTHwMAGajc1wKGQiRtDvYbwV0FWsM7sy+N0ZU5/gxJQ= +github.com/router-for-me/CLIProxyAPI/v7 v7.2.135 h1:Vzc5xJwNWq319t5++7AXCx/CnQ3U94xOjteeFtusSsk= +github.com/router-for-me/CLIProxyAPI/v7 v7.2.135/go.mod h1:lTHwMAGajc1wKGQiRtDvYbwV0FWsM7sy+N0ZU5/gxJQ= github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0= github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY= github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ= diff --git a/internal/antigravity/antigravity.go b/internal/antigravity/antigravity.go new file mode 100644 index 00000000..4ec9634b --- /dev/null +++ b/internal/antigravity/antigravity.go @@ -0,0 +1,240 @@ +// Package antigravity owns GPT-Load's Antigravity subscription contract and +// prevents CPA internal types from escaping into the application. +package antigravity + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + cpaembedded "github.com/router-for-me/CLIProxyAPI/v7/gptload-embedded/embedded" +) + +const ( + Provider = cpaembedded.ProviderAntigravity + RedirectURI = cpaembedded.AntigravityRedirectURI +) + +var ( + ErrInvalidState = errors.New("antigravity oauth state mismatch") + ErrCredentialIdentityChanged = errors.New("refreshed antigravity credential identity changed") +) + +// Credential is GPT-Load's encrypted canonical Antigravity credential. +type Credential struct { + Type string `json:"type"` + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + AccountID string `json:"account_id"` + Email string `json:"email"` + ProjectID string `json:"project_id"` + ExpiresIn int64 `json:"expires_in,omitempty"` + Timestamp int64 `json:"timestamp,omitempty"` + Expire string `json:"expired"` + LastRefresh string `json:"last_refresh,omitempty"` +} + +func ParseCredentialJSON(raw []byte) (Credential, error) { + value, err := cpaembedded.ParseAntigravityCredentialJSON(raw) + if err != nil { + return Credential{}, normalizeError(err) + } + return credentialFromBridge(value), nil +} + +func MarshalCredential(value Credential) ([]byte, error) { + raw, err := cpaembedded.MarshalAntigravityCredential(credentialToBridge(value)) + if err != nil { + return nil, normalizeError(err) + } + return raw, nil +} + +func CredentialExpiresAt(value Credential) (time.Time, bool) { + return cpaembedded.AntigravityCredentialExpiresAt(credentialToBridge(value)) +} + +func (value Credential) SecretValues() []string { + return credentialToBridge(value).SecretValues() +} + +type BrowserAuthorization struct { + AuthorizationURL string + State string + ExpiresAt time.Time +} + +type BrowserAuthorizationCompletion struct { + ExpectedState string + ReturnedState string + Code string +} + +type TokenEndpointError struct { + StatusCode int + Code string +} + +func (err *TokenEndpointError) Error() string { + if err == nil { + return "Antigravity token endpoint failed" + } + return fmt.Sprintf("Antigravity token endpoint returned status %d", err.StatusCode) +} + +func (err *TokenEndpointError) HTTPStatusCode() int { + if err == nil { + return 0 + } + return err.StatusCode +} + +type UpstreamHTTPError struct { + Operation string + StatusCode int +} + +func (err *UpstreamHTTPError) Error() string { + if err == nil { + return "Antigravity upstream request failed" + } + return fmt.Sprintf("Antigravity %s endpoint returned status %d", err.Operation, err.StatusCode) +} + +func (err *UpstreamHTTPError) HTTPStatusCode() int { + if err == nil { + return 0 + } + return err.StatusCode +} + +type Model struct { + ID string + DisplayName string + MaxTokens int64 + MaxOutputTokens int64 +} + +type GoogleOneAICredit struct { + Amount float64 + MinimumAmount float64 +} + +type AccountObservation struct { + PlanID string + CurrentTierID string + GoogleOneAICredits *GoogleOneAICredit +} + +func IsDefinitiveRefreshRejection(code string) bool { + return cpaembedded.IsDefinitiveAntigravityRefreshRejection(code) +} + +func BeginBrowserAuthorization() (BrowserAuthorization, error) { + value, err := cpaembedded.BeginAntigravityBrowserAuthorization() + if err != nil { + return BrowserAuthorization{}, normalizeError(err) + } + return BrowserAuthorization{ + AuthorizationURL: value.AuthorizationURL, State: value.State, ExpiresAt: value.ExpiresAt, + }, nil +} + +func CompleteBrowserAuthorization( + ctx context.Context, + completion BrowserAuthorizationCompletion, +) (Credential, error) { + value, err := cpaembedded.CompleteAntigravityBrowserAuthorization(ctx, cpaembedded.BrowserAuthorizationCompletion{ + ExpectedState: completion.ExpectedState, ReturnedState: completion.ReturnedState, Code: completion.Code, + }, cpaembedded.AntigravityOptions{}) + if err != nil { + return Credential{}, normalizeError(err) + } + return credentialFromBridge(value), nil +} + +func RefreshCredentialOnce(ctx context.Context, current Credential) (Credential, error) { + value, err := cpaembedded.RefreshAntigravityCredentialOnce(ctx, credentialToBridge(current), cpaembedded.AntigravityOptions{}) + if err != nil { + return Credential{}, normalizeError(err) + } + return credentialFromBridge(value), nil +} + +func ImportCredential(ctx context.Context, raw []byte) (Credential, error) { + value, err := cpaembedded.ImportAntigravityCredential(ctx, raw, cpaembedded.AntigravityOptions{}) + if err != nil { + return Credential{}, normalizeError(err) + } + return credentialFromBridge(value), nil +} + +func ListModels(ctx context.Context, credential Credential) ([]Model, error) { + values, err := cpaembedded.DiscoverAntigravityModels(ctx, credentialToBridge(credential), cpaembedded.AntigravityOptions{}) + if err != nil { + return nil, normalizeError(err) + } + result := make([]Model, 0, len(values)) + for _, value := range values { + result = append(result, Model{ + ID: value.ID, DisplayName: value.DisplayName, + MaxTokens: value.MaxTokens, MaxOutputTokens: value.MaxOutputTokens, + }) + } + return result, nil +} + +func ObserveAccount(ctx context.Context, credential Credential) (AccountObservation, error) { + value, err := cpaembedded.ObserveAntigravityAccount(ctx, credentialToBridge(credential), cpaembedded.AntigravityOptions{}) + if err != nil { + return AccountObservation{}, normalizeError(err) + } + result := AccountObservation{PlanID: value.PlanID, CurrentTierID: value.CurrentTierID} + if value.GoogleOneAICredits != nil { + result.GoogleOneAICredits = &GoogleOneAICredit{ + Amount: value.GoogleOneAICredits.Amount, MinimumAmount: value.GoogleOneAICredits.MinimumAmount, + } + } + return result, nil +} + +func credentialFromBridge(value cpaembedded.AntigravityCredential) Credential { + return Credential{ + Type: value.Type, AccessToken: value.AccessToken, RefreshToken: value.RefreshToken, + AccountID: value.AccountID, Email: value.Email, ProjectID: value.ProjectID, + ExpiresIn: value.ExpiresIn, Timestamp: value.Timestamp, Expire: value.Expire, + LastRefresh: value.LastRefresh, + } +} + +func credentialToBridge(value Credential) cpaembedded.AntigravityCredential { + return cpaembedded.AntigravityCredential{ + Type: value.Type, AccessToken: value.AccessToken, RefreshToken: value.RefreshToken, + AccountID: value.AccountID, Email: value.Email, ProjectID: value.ProjectID, + ExpiresIn: value.ExpiresIn, Timestamp: value.Timestamp, Expire: value.Expire, + LastRefresh: value.LastRefresh, + } +} + +func normalizeError(err error) error { + if err == nil { + return nil + } + if errors.Is(err, cpaembedded.ErrAntigravityInvalidState) { + return ErrInvalidState + } + if errors.Is(err, cpaembedded.ErrAntigravityCredentialIdentityChanged) { + return ErrCredentialIdentityChanged + } + var tokenErr *cpaembedded.AntigravityTokenEndpointError + if errors.As(err, &tokenErr) { + return &TokenEndpointError{StatusCode: tokenErr.StatusCode, Code: strings.TrimSpace(tokenErr.Code)} + } + var upstream *cpaembedded.AntigravityUpstreamHTTPError + if errors.As(err, &upstream) { + return &UpstreamHTTPError{Operation: strings.TrimSpace(upstream.Operation), StatusCode: upstream.StatusCode} + } + return err +} diff --git a/internal/antigravity/antigravity_test.go b/internal/antigravity/antigravity_test.go new file mode 100644 index 00000000..e67fa3b4 --- /dev/null +++ b/internal/antigravity/antigravity_test.go @@ -0,0 +1,29 @@ +package antigravity + +import ( + "reflect" + "testing" +) + +func TestParseCredentialJSONProducesCanonicalAntigravityCredential(t *testing.T) { + credential, err := ParseCredentialJSON([]byte(`{ + "type":"antigravity", + "access_token":"access-secret", + "refresh_token":"refresh-secret", + "account_id":"google-account-one", + "email":"owner@example.com", + "project_id":"project-one", + "expired":"2030-01-01T00:00:00Z" + }`)) + if err != nil { + t.Fatal(err) + } + if credential.AccountID != "google-account-one" || credential.ProjectID != "project-one" { + t.Fatalf("credential = %#v", credential) + } + if got := credential.SecretValues(); !reflect.DeepEqual(got, []string{ + "access-secret", "refresh-secret", "google-account-one", "owner@example.com", "project-one", + }) { + t.Fatalf("SecretValues() = %q", got) + } +} diff --git a/internal/antigravity/executor.go b/internal/antigravity/executor.go new file mode 100644 index 00000000..c93008d4 --- /dev/null +++ b/internal/antigravity/executor.go @@ -0,0 +1,165 @@ +package antigravity + +import ( + "context" + "net/http" + "time" + + cpaembedded "github.com/router-for-me/CLIProxyAPI/v7/gptload-embedded/embedded" +) + +// ExecuteRequest is the canonical request accepted by the embedded +// Antigravity execution bridge. +type ExecuteRequest struct { + AttemptID string + Model string + Payload []byte + Format string + Headers http.Header + OriginalRequest []byte + ContinuityKey string +} + +type ExecuteResponse struct { + Payload []byte + Headers http.Header + AppliedReasoningEffort string +} + +type ExecuteStreamResponse struct { + Headers http.Header + Chunks <-chan ExecuteStreamChunk + AppliedReasoningEffort string +} + +type ExecuteStreamChunk struct { + Payload []byte + Err error +} + +type ExecutionError struct { + status int + typeValue string + codeValue string + summary string + retryAfter *time.Duration + requestScoped bool +} + +func (err *ExecutionError) Error() string { + if err == nil || err.summary == "" { + return "Antigravity upstream request failed" + } + return err.summary +} + +func (err *ExecutionError) StatusCode() int { + if err == nil { + return 0 + } + return err.status +} + +func (err *ExecutionError) ErrorType() string { + if err == nil { + return "" + } + return err.typeValue +} + +func (err *ExecutionError) ErrorCode() string { + if err == nil { + return "" + } + return err.codeValue +} + +func (err *ExecutionError) RetryAfter() *time.Duration { + if err == nil || err.retryAfter == nil { + return nil + } + value := *err.retryAfter + return &value +} + +func (err *ExecutionError) IsRequestScoped() bool { + return err != nil && err.requestScoped +} + +// Executor is the isolated CPA execution surface consumed by GPT-Load. +type Executor interface { + Execute(context.Context, string, Credential, ExecuteRequest) (ExecuteResponse, error) + CountTokens(context.Context, string, Credential, ExecuteRequest) (ExecuteResponse, error) + ExecuteStream(context.Context, string, Credential, ExecuteRequest) (*ExecuteStreamResponse, error) +} + +type executor struct { + bridge cpaembedded.AntigravityHTTPExecutor +} + +func NewExecutor() Executor { return &executor{bridge: cpaembedded.NewAntigravityHTTPExecutor()} } + +func (executor *executor) Execute(ctx context.Context, credentialID string, credential Credential, request ExecuteRequest) (ExecuteResponse, error) { + response, err := executor.bridge.ExecuteCanonical(ctx, credentialID, credentialToBridge(credential), executeRequestToBridge(request)) + return ExecuteResponse{ + Payload: append([]byte(nil), response.Payload...), Headers: response.Headers.Clone(), + AppliedReasoningEffort: response.AppliedReasoningEffort, + }, normalizeExecutionError(err) +} + +func (executor *executor) CountTokens(ctx context.Context, credentialID string, credential Credential, request ExecuteRequest) (ExecuteResponse, error) { + response, err := executor.bridge.CountTokensCanonical(ctx, credentialID, credentialToBridge(credential), executeRequestToBridge(request)) + return ExecuteResponse{Payload: append([]byte(nil), response.Payload...), Headers: response.Headers.Clone()}, normalizeExecutionError(err) +} + +func (executor *executor) ExecuteStream(ctx context.Context, credentialID string, credential Credential, request ExecuteRequest) (*ExecuteStreamResponse, error) { + response, err := executor.bridge.ExecuteStreamCanonical(ctx, credentialID, credentialToBridge(credential), executeRequestToBridge(request)) + if response == nil { + return nil, normalizeExecutionError(err) + } + if err != nil { + return &ExecuteStreamResponse{Headers: response.Headers.Clone(), AppliedReasoningEffort: response.AppliedReasoningEffort}, normalizeExecutionError(err) + } + chunks := make(chan ExecuteStreamChunk) + go func() { + defer close(chunks) + for chunk := range response.Chunks { + converted := ExecuteStreamChunk{Payload: append([]byte(nil), chunk.Payload...), Err: normalizeExecutionError(chunk.Err)} + select { + case chunks <- converted: + case <-ctx.Done(): + return + } + } + }() + return &ExecuteStreamResponse{Headers: response.Headers.Clone(), Chunks: chunks, AppliedReasoningEffort: response.AppliedReasoningEffort}, nil +} + +func executeRequestToBridge(value ExecuteRequest) cpaembedded.ExecuteRequest { + return cpaembedded.ExecuteRequest{ + AttemptID: value.AttemptID, Model: value.Model, Payload: append([]byte(nil), value.Payload...), Format: value.Format, + Headers: value.Headers.Clone(), OriginalRequest: append([]byte(nil), value.OriginalRequest...), + ContinuityKey: value.ContinuityKey, + } +} + +func normalizeExecutionError(err error) error { + if err == nil { + return nil + } + bridge, ok := err.(*cpaembedded.AntigravityExecutionError) + if !ok || bridge == nil { + return err + } + result := &ExecutionError{ + status: bridge.StatusCode(), typeValue: bridge.ErrorType(), codeValue: bridge.ErrorCode(), + summary: bridge.Error(), requestScoped: bridge.IsRequestScoped(), + } + if retryAfter := bridge.RetryAfter(); retryAfter != nil && *retryAfter > 0 { + value := *retryAfter + result.retryAfter = &value + } + return result +} + +var _ Executor = (*executor)(nil) diff --git a/internal/channel/antigravity_module_test.go b/internal/channel/antigravity_module_test.go new file mode 100644 index 00000000..12fc054c --- /dev/null +++ b/internal/channel/antigravity_module_test.go @@ -0,0 +1,81 @@ +package channel + +import ( + "reflect" + "testing" + + "gpt-load/internal/channel/modules" + "gpt-load/internal/channel/spec" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +func TestAntigravityModuleDeclaresSubscriptionContract(t *testing.T) { + t.Parallel() + + definitions, err := compileBuiltInModules([]spec.Module{modules.Antigravity()}) + if err != nil { + t.Fatal(err) + } + registry, err := newRegistry(definitions) + if err != nil { + t.Fatal(err) + } + descriptor, ok := registry.Get(Antigravity) + if !ok { + t.Fatal("Antigravity descriptor is missing") + } + if descriptor.Connection.Type != string(spec.ConnectionSubscription) || + !reflect.DeepEqual(descriptor.Connection.AuthorizationMethods, []AuthorizationMethod{ + AuthorizationBrowserOAuth, + AuthorizationOAuthFile, + }) { + t.Fatalf("Antigravity connection = %#v", descriptor.Connection) + } + if !reflect.DeepEqual(descriptor.Notices, []NoticeDescriptor{{ + ID: NoticeAntigravityOAuthRisk, Tone: NoticeToneWarning, + }}) { + t.Fatalf("Antigravity notices = %#v", descriptor.Notices) + } + target, err := registry.Resolve(Antigravity, nil) + if err != nil { + t.Fatal(err) + } + if target.ProviderKind != ProviderAntigravity || target.CatalogProviderID != "" { + t.Fatalf("Antigravity target = %#v", target) + } + bindings, ok := registry.CapabilityBindings(Antigravity) + if !ok { + t.Fatal("Antigravity capability bindings are missing") + } + if bindings.SubscriptionDriver != modules.AntigravitySubscriptionDriver || + bindings.ModelDiscovery != modules.AntigravityModelDiscovery || + bindings.QuotaObservation != modules.AntigravityQuotaObservation || + bindings.ResetCreditAction != "" { + t.Fatalf("Antigravity capabilities = %#v", bindings) + } + policy, ok := registry.SchedulingPolicy(Antigravity) + if !ok || policy.QuotaPriority { + t.Fatal("Antigravity unexpectedly enables quota-priority scheduling") + } + for _, test := range []struct { + clientProtocol protocol.Protocol + operation execution.Operation + mode RouteMode + }{ + {protocol.Gemini, execution.OperationChatCompletion, RouteNative}, + {protocol.Gemini, execution.OperationCountTokens, RouteNative}, + {protocol.Anthropic, execution.OperationChatCompletion, RouteConverted}, + {protocol.Anthropic, execution.OperationCountTokens, RouteConverted}, + {protocol.OpenAICompletions, execution.OperationChatCompletion, RouteConverted}, + {protocol.OpenAIResponses, execution.OperationResponsesCreate, RouteConverted}, + {protocol.OpenAIResponses, execution.OperationResponsesInputTokens, RouteConverted}, + } { + if mode, exists := target.Mode(test.clientProtocol, test.operation); !exists || mode != test.mode { + t.Fatalf("Antigravity %q/%q mode = %q, %t; want %q", test.clientProtocol, test.operation, mode, exists, test.mode) + } + } + if _, exists := target.Mode(protocol.OpenAIResponses, execution.OperationResponsesCompact); exists { + t.Fatal("Antigravity unexpectedly declares responses compact") + } +} diff --git a/internal/channel/builtins.go b/internal/channel/builtins.go index 9bfc19c3..f3126fee 100644 --- a/internal/channel/builtins.go +++ b/internal/channel/builtins.go @@ -12,6 +12,7 @@ func builtInModules() []spec.Module { modules.OpenAI(), modules.Codex(), modules.Claude(), + modules.Antigravity(), modules.Anthropic(), modules.Gemini(), modules.AzureOpenAI(), diff --git a/internal/channel/channel.go b/internal/channel/channel.go index 72f96a91..f3de8c9c 100644 --- a/internal/channel/channel.go +++ b/internal/channel/channel.go @@ -24,6 +24,7 @@ const ( OpenAI = spec.OpenAI Codex = spec.Codex Claude = spec.Claude + Antigravity = spec.Antigravity Anthropic = spec.Anthropic Gemini = spec.Gemini AzureOpenAI = spec.AzureOpenAI @@ -91,7 +92,10 @@ type CapabilityDescriptor struct { // NoticeID identifies one frontend-localized channel notice. type NoticeID = spec.NoticeID -const NoticeClaudeOAuthRisk = spec.NoticeClaudeOAuthRisk +const ( + NoticeClaudeOAuthRisk = spec.NoticeClaudeOAuthRisk + NoticeAntigravityOAuthRisk = spec.NoticeAntigravityOAuthRisk +) // NoticeTone is the bounded presentation style of a channel notice. type NoticeTone = spec.NoticeTone @@ -180,6 +184,7 @@ const ( ProviderOpenAI = spec.ProviderOpenAI ProviderCodex = spec.ProviderCodex ProviderClaude = spec.ProviderClaude + ProviderAntigravity = spec.ProviderAntigravity ProviderAnthropic = spec.ProviderAnthropic ProviderGemini = spec.ProviderGemini ProviderOpenAICompatible = spec.ProviderOpenAICompatible diff --git a/internal/channel/claude_module_test.go b/internal/channel/claude_module_test.go index a2a290be..abfd81d0 100644 --- a/internal/channel/claude_module_test.go +++ b/internal/channel/claude_module_test.go @@ -103,6 +103,9 @@ func TestBuiltInCountTokensRouteMatrix(t *testing.T) { {Codex, protocol.OpenAIResponses, execution.OperationResponsesInputTokens, RouteNative, true}, {Codex, protocol.Anthropic, countTokens, RouteConverted, true}, {Codex, protocol.Gemini, countTokens, RouteConverted, true}, + {Antigravity, protocol.Gemini, countTokens, RouteNative, true}, + {Antigravity, protocol.Anthropic, countTokens, RouteConverted, true}, + {Antigravity, protocol.OpenAIResponses, execution.OperationResponsesInputTokens, RouteConverted, true}, } for _, test := range tests { t.Run(string(test.channelID)+"/"+string(test.clientProtocol)+"/"+string(test.operation), func(t *testing.T) { diff --git a/internal/channel/final_provider_contract_test.go b/internal/channel/final_provider_contract_test.go index dabfa015..8cfde493 100644 --- a/internal/channel/final_provider_contract_test.go +++ b/internal/channel/final_provider_contract_test.go @@ -15,6 +15,7 @@ func TestFinalRegistryContainsOnlyApprovedChannels(t *testing.T) { OpenAI, Codex, Claude, + Antigravity, Anthropic, Gemini, AzureOpenAI, diff --git a/internal/channel/modules/antigravity.go b/internal/channel/modules/antigravity.go new file mode 100644 index 00000000..84b0dcf5 --- /dev/null +++ b/internal/channel/modules/antigravity.go @@ -0,0 +1,59 @@ +package modules + +import ( + "gpt-load/internal/channel/spec" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +const ( + AntigravitySubscriptionDriver spec.SubscriptionDriverID = "antigravity" + AntigravityModelDiscovery spec.UtilityID = "antigravity_models" + AntigravityQuotaObservation spec.UtilityID = "antigravity_quota" +) + +// Antigravity declares the Google Antigravity subscription channel. +func Antigravity() spec.Module { + return spec.Module{ + Definition: spec.Definition{ + ID: spec.Antigravity, + Name: "Antigravity", + Mark: "AG", + Icon: "antigravity", + SearchTerms: []string{"subscription", "oauth", "google", "antigravity"}, + Description: "Google Antigravity subscription", + Notices: []spec.Notice{{ + ID: spec.NoticeAntigravityOAuthRisk, Tone: spec.NoticeToneWarning, + }}, + Connection: spec.Connection{ + Type: spec.ConnectionSubscription, + CredentialInput: "authorization", + AuthorizationMethods: []spec.AuthorizationMethod{ + spec.AuthorizationBrowserOAuth, + spec.AuthorizationOAuthFile, + }, + }, + Params: []spec.Field{}, + Credentials: []spec.Field{}, + Provider: spec.ProviderBinding{ + ProviderKind: spec.ProviderAntigravity, + EndpointPolicy: spec.EndpointNone, + }, + Routes: []spec.Route{ + spec.NewRoute(protocol.Gemini, execution.OperationChatCompletion, execution.RouteNative), + spec.NewRoute(protocol.Gemini, execution.OperationCountTokens, execution.RouteNative), + spec.NewRoute(protocol.Anthropic, execution.OperationChatCompletion, execution.RouteConverted), + spec.NewRoute(protocol.Anthropic, execution.OperationCountTokens, execution.RouteConverted), + spec.NewRoute(protocol.OpenAICompletions, execution.OperationChatCompletion, execution.RouteConverted), + spec.NewRoute(protocol.OpenAIResponses, execution.OperationResponsesCreate, execution.RouteConverted), + spec.NewRoute(protocol.OpenAIResponses, execution.OperationResponsesInputTokens, execution.RouteConverted), + }, + Capabilities: spec.CapabilityBindings{ + SubscriptionDriver: AntigravitySubscriptionDriver, + ModelDiscovery: AntigravityModelDiscovery, + QuotaObservation: AntigravityQuotaObservation, + }, + Scheduling: spec.SchedulingPolicy{QuotaPriority: false}, + }, + } +} diff --git a/internal/channel/registry_test.go b/internal/channel/registry_test.go index ec85369b..6265aba7 100644 --- a/internal/channel/registry_test.go +++ b/internal/channel/registry_test.go @@ -16,6 +16,7 @@ func TestRegistryHasStableBuiltInOrderAndSearch(t *testing.T) { OpenAI, Codex, Claude, + Antigravity, Anthropic, Gemini, ID("azure_openai"), @@ -38,8 +39,8 @@ func TestRegistryHasStableBuiltInOrderAndSearch(t *testing.T) { if got := descriptorIDs(registry.Search("GeMiNi")); !reflect.DeepEqual(got, []ID{Gemini}) { t.Fatalf("Search(gemini) IDs = %v", got) } - if got := descriptorIDs(registry.Search("subscription")); !reflect.DeepEqual(got, []ID{Codex, Claude}) { - t.Fatalf("Search(subscription) IDs = %v, want [codex claude]", got) + if got := descriptorIDs(registry.Search("subscription")); !reflect.DeepEqual(got, []ID{Codex, Claude, Antigravity}) { + t.Fatalf("Search(subscription) IDs = %v, want [codex claude antigravity]", got) } if got := descriptorIDs(registry.Search("compatible")); !reflect.DeepEqual(got, []ID{OpenAICompatible}) { t.Fatalf("Search(compatible) IDs = %v", got) @@ -63,7 +64,7 @@ func TestRegistryHasStableBuiltInOrderAndSearch(t *testing.T) { first := registry.List() first[0].ClientProtocols[0] = protocol.Protocol("mutated") first[0].ParamFields = append(first[0].ParamFields, FieldDescriptor{Key: "mutated"}) - vertexIndex := 7 + vertexIndex := 8 if first[vertexIndex].ID != GoogleVertex || first[vertexIndex].ParamFields[0].DefaultValue == nil { t.Fatalf("unexpected Vertex descriptor = %#v", first[vertexIndex]) } diff --git a/internal/channel/route_golden_test.go b/internal/channel/route_golden_test.go index 722a2735..33cc2ad3 100644 --- a/internal/channel/route_golden_test.go +++ b/internal/channel/route_golden_test.go @@ -41,7 +41,7 @@ func TestBuiltInRouteGolden(t *testing.T) { } digest := fmt.Sprintf("%x", sha256.Sum256([]byte(strings.Join(got, "\n")))) - const wantDigest = "df79673b4729a8a47f65786b42660dd063ec80ee7249b1c6f7090a836018aac2" + const wantDigest = "fb123f1bc2a0afdb64a29e525eab265933a4c7ab6ba8057d73303293569887f2" if digest != wantDigest { t.Fatalf("built-in routes changed: digest = %s, want %s\n%s", digest, wantDigest, strings.Join(got, "\n")) } diff --git a/internal/channel/spec/definition.go b/internal/channel/spec/definition.go index 2d50a9df..a2dca98d 100644 --- a/internal/channel/spec/definition.go +++ b/internal/channel/spec/definition.go @@ -16,6 +16,7 @@ const ( OpenAI ID = "openai" Codex ID = "codex" Claude ID = "claude" + Antigravity ID = "antigravity" Anthropic ID = "anthropic" Gemini ID = "gemini" AzureOpenAI ID = "azure_openai" @@ -40,6 +41,7 @@ const ( ProviderOpenAI ProviderKind = "openai" ProviderCodex ProviderKind = "codex" ProviderClaude ProviderKind = "claude" + ProviderAntigravity ProviderKind = "antigravity" ProviderAnthropic ProviderKind = "anthropic" ProviderGemini ProviderKind = "gemini" ProviderOpenAICompatible ProviderKind = "openai_compatible" @@ -55,11 +57,14 @@ const ( // NoticeID identifies one code-owned, frontend-localized channel notice. type NoticeID string -const NoticeClaudeOAuthRisk NoticeID = "claude_oauth_risk" +const ( + NoticeClaudeOAuthRisk NoticeID = "claude_oauth_risk" + NoticeAntigravityOAuthRisk NoticeID = "antigravity_oauth_risk" +) // Valid reports whether the notice ID is part of the public channel contract. func (id NoticeID) Valid() bool { - return id == NoticeClaudeOAuthRisk + return id == NoticeClaudeOAuthRisk || id == NoticeAntigravityOAuthRisk } // NoticeTone controls the bounded presentation style of one channel notice. @@ -84,6 +89,7 @@ func (kind ProviderKind) Valid() bool { case ProviderOpenAI, ProviderCodex, ProviderClaude, + ProviderAntigravity, ProviderAnthropic, ProviderGemini, ProviderOpenAICompatible, diff --git a/internal/codex/dependency_test.go b/internal/codex/dependency_test.go index 25b200db..548d7312 100644 --- a/internal/codex/dependency_test.go +++ b/internal/codex/dependency_test.go @@ -17,7 +17,8 @@ func TestBridgeDependencyStaysBehindSubscriptionProviderBoundaries(t *testing.T) for _, line := range strings.Split(strings.TrimSpace(string(output)), "\n") { importer, imports, found := strings.Cut(line, "|") if !found || importer == "gpt-load/internal/codex" || strings.HasPrefix(importer, "gpt-load/internal/codex/") || - importer == "gpt-load/internal/claude" || strings.HasPrefix(importer, "gpt-load/internal/claude/") { + importer == "gpt-load/internal/claude" || strings.HasPrefix(importer, "gpt-load/internal/claude/") || + importer == "gpt-load/internal/antigravity" || strings.HasPrefix(importer, "gpt-load/internal/antigravity/") { continue } for _, dependency := range strings.Split(imports, ",") { diff --git a/internal/container/container.go b/internal/container/container.go index 4723c71b..6760387f 100644 --- a/internal/container/container.go +++ b/internal/container/container.go @@ -208,6 +208,7 @@ func newSubscriptionRuntime(registry *channel.Registry) (*subscriptionruntime.Ru registry, subscriptionruntime.CodexImplementations(), subscriptionruntime.ClaudeImplementations(), + subscriptionruntime.AntigravityImplementations(), ) } @@ -230,6 +231,7 @@ func newProviderAdapterRegistry( {ProviderKind: channel.ProviderXAI, Adapter: bifrost}, {ProviderKind: channel.ProviderCodex, Adapter: cpa}, {ProviderKind: channel.ProviderClaude, Adapter: cpa}, + {ProviderKind: channel.ProviderAntigravity, Adapter: cpa}, } return provideradapter.NewRegistry(channels, bindings) } diff --git a/internal/control/catalog_discovery_merge_test.go b/internal/control/catalog_discovery_merge_test.go index 04db461e..b218061b 100644 --- a/internal/control/catalog_discovery_merge_test.go +++ b/internal/control/catalog_discovery_merge_test.go @@ -6,6 +6,7 @@ import ( "reflect" "testing" + "gpt-load/internal/antigravity" "gpt-load/internal/catalog" "gpt-load/internal/channel" "gpt-load/internal/codex" @@ -102,6 +103,68 @@ func TestClaudeDiscoveryUsesOnlySubscriptionModelsAndReferencePrices(t *testing. } } +func TestAntigravityDiscoveryUsesOnlySubscriptionModelsAndReferencePrices(t *testing.T) { + fixture := newServiceFixture(t) + fixture.catalogRuntime.Publish(&catalog.Snapshot{Providers: map[string]catalog.Provider{ + "google": { + ID: "google", Name: "Google", Models: map[string]catalog.Model{ + "gemini-antigravity": { + ID: "gemini-antigravity", Name: "Google catalog name", + Cost: &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(1)}}, + }, + "google-catalog-only": { + ID: "google-catalog-only", Name: "Google catalog only", + Cost: &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(2)}}, + }, + }, + }, + }}) + canonical, err := antigravity.MarshalCredential(antigravity.Credential{ + Type: "antigravity", AccessToken: "access", RefreshToken: "refresh", AccountID: "google-account", + Email: "antigravity@example.com", ProjectID: "project-one", Expire: "2030-01-01T00:00:00Z", + }) + if err != nil { + t.Fatal(err) + } + driver, ok := fixture.service.subscriptions.Driver(channel.Antigravity) + if !ok { + t.Fatal("Antigravity driver is unavailable") + } + credential, err := driver.Parse(canonical) + if err != nil { + t.Fatal(err) + } + stage, err := fixture.service.persistReadyCredentialStage(t.Context(), channel.Antigravity, "oauth_file", credential) + if err != nil { + t.Fatal(err) + } + fixture.service.discoverSubscriptionModels = func( + _ context.Context, + channelID channel.ID, + _ subscriptionruntime.Credential, + ) ([]string, error) { + if channelID != channel.Antigravity { + t.Fatalf("channel = %q, want Antigravity", channelID) + } + return []string{"gemini-antigravity"}, nil + } + + got, err := fixture.service.DiscoverModels(t.Context(), ModelDiscoveryRequest{ + ChannelID: channel.Antigravity, StagedCredentialID: stage.StageID, ConnectionType: "subscription", + }) + if err != nil { + t.Fatal(err) + } + google := "Google" + want := []ModelCandidate{{ + ID: "gemini-antigravity", Name: "gemini-antigravity", Sources: []string{"live"}, + PricingStatus: PricingStatusConfigured, PricingSource: &google, + }} + if !reflect.DeepEqual(got.Models, want) { + t.Fatalf("Antigravity candidates = %#v, want %#v", got.Models, want) + } +} + func TestDraftDiscoveryMergesLiveAndLocalCatalogByExactIDWithoutURLInference(t *testing.T) { fixture := newServiceFixture(t) fixture.catalogRuntime.Publish(&catalog.Snapshot{Providers: map[string]catalog.Provider{ diff --git a/internal/control/credential_stage_routes_test.go b/internal/control/credential_stage_routes_test.go index df5a6c3a..546de82e 100644 --- a/internal/control/credential_stage_routes_test.go +++ b/internal/control/credential_stage_routes_test.go @@ -221,6 +221,7 @@ func TestOAuthCallbackServerStartsIndependentDriverEndpoints(t *testing.T) { callbacks := []subscriptionruntime.LocalCallbackSpec{ {RedirectURI: "http://localhost:1455/auth/callback"}, {RedirectURI: "http://localhost:54545/callback"}, + {RedirectURI: "http://localhost:51121/oauth-callback"}, } for _, callback := range callbacks { if err := fixture.service.oauthCallback.EnsureStarted(callback); err != nil { @@ -229,16 +230,23 @@ func TestOAuthCallbackServerStartsIndependentDriverEndpoints(t *testing.T) { } t.Cleanup(func() { _ = fixture.service.oauthCallback.Stop(t.Context()) }) - firstAddr := fixture.service.oauthCallback.Addr(callbacks[0]) - secondAddr := fixture.service.oauthCallback.Addr(callbacks[1]) - if firstAddr == "" || secondAddr == "" || firstAddr == secondAddr { - t.Fatalf("callback addresses = %q and %q, want independent listeners", firstAddr, secondAddr) + addresses := make([]string, len(callbacks)) + seen := make(map[string]struct{}, len(callbacks)) + for index, callback := range callbacks { + addresses[index] = fixture.service.oauthCallback.Addr(callback) + if addresses[index] == "" { + t.Fatalf("callback %d has no listener address", index) + } + if _, duplicate := seen[addresses[index]]; duplicate { + t.Fatalf("callback listeners share address %q", addresses[index]) + } + seen[addresses[index]] = struct{}{} } for index, callback := range callbacks { if err := fixture.service.oauthCallback.EnsureStarted(callback); err != nil { t.Fatal(err) } - if got := fixture.service.oauthCallback.Addr(callback); got != []string{firstAddr, secondAddr}[index] { + if got := fixture.service.oauthCallback.Addr(callback); got != addresses[index] { t.Fatalf("idempotent callback address = %q", got) } } diff --git a/internal/control/credential_stages.go b/internal/control/credential_stages.go index 33b191e3..d0cfb894 100644 --- a/internal/control/credential_stages.go +++ b/internal/control/credential_stages.go @@ -5,6 +5,8 @@ import ( "encoding/json" "errors" "fmt" + "net" + "net/http" "sort" "strings" "time" @@ -26,6 +28,34 @@ const ( stagedSubscriptionSchemaV2 = 2 ) +func credentialImportContext(ctx context.Context) (context.Context, context.CancelFunc) { + return context.WithTimeout(ctx, defaultSubscriptionControlTimeout) +} + +func credentialImportAPIError(err error) error { + if errors.Is(err, context.Canceled) { + return context.Canceled + } + if errors.Is(err, context.DeadlineExceeded) { + return app_errors.ErrBadGateway + } + var networkError net.Error + if errors.As(err, &networkError) { + return app_errors.ErrBadGateway + } + var upstream interface{ HTTPStatusCode() int } + if errors.As(err, &upstream) && upstream != nil { + switch status := upstream.HTTPStatusCode(); { + case status == http.StatusRequestTimeout, + status == http.StatusTooEarly, + status == http.StatusTooManyRequests, + status >= http.StatusInternalServerError: + return app_errors.ErrBadGateway + } + } + return app_errors.ErrOAuthFileInvalid +} + type CredentialStageAccount struct { EmailMask string `json:"email_mask,omitempty"` ExpiresAtMS *int64 `json:"expires_at_ms,omitempty"` @@ -115,9 +145,11 @@ func (s *Service) ImportCredentialStage( if len(raw) > maxOAuthFileBytes { return CredentialStageResult{}, app_errors.ErrOAuthFileTooLarge } - credential, err := driver.Parse(raw) + importContext, cancelImport := credentialImportContext(ctx) + defer cancelImport() + credential, err := s.subscriptions.ImportCredential(importContext, channelID, raw) if err != nil { - return CredentialStageResult{}, app_errors.ErrOAuthFileInvalid + return CredentialStageResult{}, credentialImportAPIError(err) } credential, err = s.prepareTransientSubscriptionCredential(ctx, channelID, driver, credential) if err != nil { diff --git a/internal/control/credential_stages_test.go b/internal/control/credential_stages_test.go index ca44e8e9..347a16c4 100644 --- a/internal/control/credential_stages_test.go +++ b/internal/control/credential_stages_test.go @@ -16,6 +16,47 @@ import ( "gpt-load/internal/storage/models" ) +type credentialImportStatusError struct{ status int } + +func (err credentialImportStatusError) Error() string { return "credential import upstream failed" } +func (err credentialImportStatusError) HTTPStatusCode() int { return err.status } + +type credentialImportNetworkError struct{} + +func (credentialImportNetworkError) Error() string { return "credential import network failed" } +func (credentialImportNetworkError) Timeout() bool { return true } +func (credentialImportNetworkError) Temporary() bool { return true } + +func TestCredentialImportUsesBoundedContextAndClassifiesTransientFailures(t *testing.T) { + ctx, cancel := credentialImportContext(t.Context()) + defer cancel() + deadline, known := ctx.Deadline() + remaining := time.Until(deadline) + if !known || remaining <= 0 || remaining > defaultSubscriptionControlTimeout { + t.Fatalf("import context deadline = %v, %t", deadline, known) + } + + for _, test := range []struct { + name string + err error + want error + }{ + {name: "invalid file", err: errors.New("invalid schema"), want: app_errors.ErrOAuthFileInvalid}, + {name: "canceled", err: context.Canceled, want: context.Canceled}, + {name: "deadline", err: context.DeadlineExceeded, want: app_errors.ErrBadGateway}, + {name: "network", err: credentialImportNetworkError{}, want: app_errors.ErrBadGateway}, + {name: "upstream unavailable", err: credentialImportStatusError{status: http.StatusServiceUnavailable}, want: app_errors.ErrBadGateway}, + {name: "upstream throttled", err: credentialImportStatusError{status: http.StatusTooManyRequests}, want: app_errors.ErrBadGateway}, + {name: "token rejected", err: credentialImportStatusError{status: http.StatusBadRequest}, want: app_errors.ErrOAuthFileInvalid}, + } { + t.Run(test.name, func(t *testing.T) { + if got := credentialImportAPIError(fmt.Errorf("import credential: %w", test.err)); !errors.Is(got, test.want) { + t.Fatalf("credentialImportAPIError() = %v, want %v", got, test.want) + } + }) + } +} + func TestImportCodexOAuthJSONCreatesEncryptedReadyStage(t *testing.T) { t.Parallel() diff --git a/internal/control/price_reference_test.go b/internal/control/price_reference_test.go index 0737a6bc..d988bc99 100644 --- a/internal/control/price_reference_test.go +++ b/internal/control/price_reference_test.go @@ -110,6 +110,23 @@ func TestResolveAutomaticPriceForIdentityTreatsClaudeAnthropicPriceAsReference(t } } +func TestResolveAutomaticPriceForIdentityTreatsAntigravityGooglePriceAsReference(t *testing.T) { + googleCost := &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(1)}} + snapshot := &catalog.Snapshot{Providers: map[string]catalog.Provider{ + "google": {ID: "google", Models: map[string]catalog.Model{ + "gemini-antigravity": {ID: "gemini-antigravity", Cost: googleCost}, + }}, + }} + + match, ok := resolveAutomaticPriceForIdentity(snapshot, pricing.Identity{ + ChannelID: string(channel.Antigravity), ModelID: "gemini-antigravity", + }) + if !ok || match.cost != googleCost || match.providerID != "google" || + match.source != ModelPriceMatchSourceProviderPriorityFallback { + t.Fatalf("Antigravity price reference = %#v, %t", match, ok) + } +} + func TestCompatibleAutomaticPriceUsesGlobalPriority(t *testing.T) { openAICost := &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(1)}} alphaCost := &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(2)}} diff --git a/internal/control/test_helpers_test.go b/internal/control/test_helpers_test.go index 422a7024..d3b79214 100644 --- a/internal/control/test_helpers_test.go +++ b/internal/control/test_helpers_test.go @@ -201,6 +201,7 @@ func newServiceFixtureWithDatabase(t *testing.T, db *gorm.DB) serviceFixture { channelRegistry, subscriptionruntime.CodexImplementations(), subscriptionruntime.ClaudeImplementations(), + subscriptionruntime.AntigravityImplementations(), ) if err != nil { t.Fatalf("subscriptionruntime.NewRuntime() error = %v", err) diff --git a/internal/execution/contracts.go b/internal/execution/contracts.go index 5eae6b28..a0efa3e9 100644 --- a/internal/execution/contracts.go +++ b/internal/execution/contracts.go @@ -195,6 +195,9 @@ type AttemptSpec struct { // ForceCredentialRefresh is set only by GPT-Load after a provider explicitly // rejects this selected subscription credential before processing. ForceCredentialRefresh bool `json:"force_credential_refresh,omitempty"` + // ContinuityKey is an opaque tenant-scoped value for provider-private + // thinking/tool continuity. It is never persisted, logged, or exposed. + ContinuityKey string `json:"-"` // TargetConfig is non-secret configuration resolved by the channel registry. TargetConfig json.RawMessage `json:"target_config,omitempty"` Timeouts AttemptTimeouts `json:"timeouts"` @@ -283,6 +286,7 @@ const ( FailureHintReauthorizationRequired FailureHint = "reauthorization_required" FailureHintRateLimited FailureHint = "rate_limited" FailureHintRequestRejected FailureHint = "request_rejected" + FailureHintCandidateUnavailable FailureHint = "candidate_unavailable" FailureHintModelUnavailable FailureHint = "model_unavailable" FailureHintHostError FailureHint = "host_error" ) @@ -292,7 +296,8 @@ func (h FailureHint) Valid() bool { switch h { case "", FailureHintInvalidCredential, FailureHintRefreshRequired, FailureHintReauthorizationRequired, FailureHintRateLimited, - FailureHintRequestRejected, FailureHintModelUnavailable, FailureHintHostError: + FailureHintRequestRejected, FailureHintCandidateUnavailable, + FailureHintModelUnavailable, FailureHintHostError: return true default: return false diff --git a/internal/execution/contracts_test.go b/internal/execution/contracts_test.go index 4a8e2328..0b3753ba 100644 --- a/internal/execution/contracts_test.go +++ b/internal/execution/contracts_test.go @@ -517,6 +517,15 @@ func TestValidationAcceptsValidContractsAndRejectsInvalidFields(t *testing.T) { }).Validate(); err != nil { t.Fatalf("request-rejected error evidence Validate() error = %v", err) } + if err := (ErrorEvidence{ + Kind: ErrorKindHTTP, + Hint: FailureHintCandidateUnavailable, + StatusCode: http.StatusForbidden, + ReplaySafety: ReplaySafetyRejectedBeforeProcessing, + Summary: "candidate is unavailable", + }).Validate(); err != nil { + t.Fatalf("candidate-unavailable error evidence Validate() error = %v", err) + } } func TestStreamEventValidation(t *testing.T) { diff --git a/internal/execution/cpa/adapter.go b/internal/execution/cpa/adapter.go index cc8687c2..6962acc8 100644 --- a/internal/execution/cpa/adapter.go +++ b/internal/execution/cpa/adapter.go @@ -54,7 +54,11 @@ func NewAdapter(credentials *subscription.CredentialManager, channels *channel.R return &Adapter{ credentials: credentials, channels: channels, - providers: indexProviderBridges(newCodexProviderBridge(), newClaudeProviderBridge()), + providers: indexProviderBridges( + newCodexProviderBridge(), + newClaudeProviderBridge(), + newAntigravityProviderBridge(), + ), } } @@ -81,6 +85,16 @@ func (a *Adapter) Execute(ctx context.Context, spec execution.AttemptSpec) execu return unaryNotSent(execution.ErrorKindInvalidRequest, "unsupported subscription request", "", err) } request := bridgeRequest(spec) + if validator, ok := provider.(providerRequestValidator); ok { + if err := validator.ValidateRequest(request); err != nil { + return unaryNotSent( + execution.ErrorKindInvalidRequest, + "subscription request input is not supported", + "unsupported_subscription_input", + err, + ) + } + } if countTokensOperation(spec.Operation) { if local, ok := provider.(providerLocalTokenCounter); ok { if err := local.ValidateLocalTokenCount(request); err != nil { @@ -193,6 +207,12 @@ func (a *Adapter) ExecuteStream(ctx context.Context, spec execution.AttemptSpec, if countTokensOperation(spec.Operation) { return streamNotSent(execution.ErrorKindInvalidRequest, "count tokens does not support streaming", "") } + request := bridgeRequest(spec) + if validator, ok := provider.(providerRequestValidator); ok { + if err := validator.ValidateRequest(request); err != nil { + return streamNotSent(execution.ErrorKindInvalidRequest, "subscription request input is not supported", "unsupported_subscription_input") + } + } preparedCredential, evidence := a.credentials.Prepare(ctx, channel.ID(spec.ChannelID), spec.Credential, spec.ForceCredentialRefresh) if evidence != nil { return execution.StreamResult{DispatchState: execution.DispatchNotSent, Error: evidence} @@ -209,7 +229,7 @@ func (a *Adapter) ExecuteStream(ctx context.Context, spec execution.AttemptSpec, defer cancelStream(context.Canceled) firstByte := startFirstByteGate(spec.Timeouts.FirstByte, cancelStream) defer firstByte.stop() - response, err := provider.ExecuteStream(streamCtx, strconv.FormatUint(uint64(spec.Credential.ID), 10), credential, bridgeRequest(spec)) + response, err := provider.ExecuteStream(streamCtx, strconv.FormatUint(uint64(spec.Credential.ID), 10), credential, request) if err != nil { result := unaryExecutionError(streamCtx, provider, err, credential) var applied *reasoning.Config @@ -340,9 +360,10 @@ func (a *Adapter) validateSpec(spec execution.AttemptSpec) (providerBridge, erro func bridgeRequest(spec execution.AttemptSpec) providerRequest { return providerRequest{ - Model: spec.UpstreamModel, Payload: append([]byte(nil), spec.Body...), + AttemptID: spec.AttemptID, Model: spec.UpstreamModel, Payload: append([]byte(nil), spec.Body...), Format: formatFor(spec.ClientProtocol), Headers: spec.Header.Clone(), OriginalRequest: append([]byte(nil), spec.Body...), + ContinuityKey: spec.ContinuityKey, } } diff --git a/internal/execution/cpa/adapter_test.go b/internal/execution/cpa/adapter_test.go index 7051a1bc..e3e1185d 100644 --- a/internal/execution/cpa/adapter_test.go +++ b/internal/execution/cpa/adapter_test.go @@ -15,6 +15,7 @@ import ( "github.com/glebarez/sqlite" "gorm.io/gorm" + "gpt-load/internal/antigravity" "gpt-load/internal/channel" "gpt-load/internal/codex" "gpt-load/internal/execution" @@ -323,6 +324,33 @@ func TestAdapterRejectsUnsupportedLocalCodexTokenCountInput(t *testing.T) { } } +func TestAdapterRejectsUnsupportedAntigravityInputBeforeCredentialPreparation(t *testing.T) { + canonical, err := antigravity.MarshalCredential(antigravity.Credential{ + Type: "antigravity", AccessToken: "access", RefreshToken: "refresh", AccountID: "google-account-one", + Email: "owner@example.com", ProjectID: "project-one", Expire: time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }) + if err != nil { + t.Fatal(err) + } + adapter, _, _, keyService, row := newSubscriptionAdapterFixture(t, string(channel.Antigravity), canonical, "google-account-one") + preparer := &fakeCredentialPreparer{delegate: adapter.credentials} + adapter.credentials = preparer + spec := validSpec(t, row, keyService) + spec.ChannelID = string(channel.Antigravity) + spec.ClientProtocol = protocol.OpenAIResponses + spec.Operation = execution.OperationResponsesCreate + spec.RouteMode = execution.RouteConverted + spec.ClientModel, spec.UpstreamModel = "gemini-live", "gemini-live" + spec.Body = []byte(`{"input":[{"role":"user","content":[{"type":"input_image","image_url":"https://example.test/image.png"}]}]}`) + + result := adapter.Execute(t.Context(), spec) + if result.DispatchState != execution.DispatchNotSent || result.ResponseStarted || result.Error == nil || + result.Error.Kind != execution.ErrorKindInvalidRequest || result.Error.Code != "unsupported_subscription_input" || + preparer.calls != 0 { + t.Fatalf("result=%#v prepareCalls=%d", result, preparer.calls) + } +} + func TestAdapterMarksLocalCodexTokenCountFailureNotSent(t *testing.T) { adapter, _, _, keyService, row := newAdapterFixture(t, credentialJSON("access", "refresh", time.Now().Add(time.Hour))) preparer := &fakeCredentialPreparer{delegate: adapter.credentials} @@ -775,6 +803,7 @@ func newSubscriptionAdapterFixture( channels, subscriptionruntime.CodexImplementations(), subscriptionruntime.ClaudeImplementations(), + subscriptionruntime.AntigravityImplementations(), ) if err != nil { t.Fatal(err) diff --git a/internal/execution/cpa/antigravity_provider.go b/internal/execution/cpa/antigravity_provider.go new file mode 100644 index 00000000..e3d0b833 --- /dev/null +++ b/internal/execution/cpa/antigravity_provider.go @@ -0,0 +1,406 @@ +package cpa + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net" + "net/http" + "strings" + "time" + + "gpt-load/internal/antigravity" + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/execution/responsealias" + "gpt-load/internal/protocol" +) + +type antigravityProviderCredential struct{ value antigravity.Credential } + +func (credential antigravityProviderCredential) redactionValues() []string { + return credential.value.SecretValues() +} + +type antigravityProviderBridge struct{ executor antigravity.Executor } + +func newAntigravityProviderBridge() *antigravityProviderBridge { + return &antigravityProviderBridge{executor: antigravity.NewExecutor()} +} + +func (*antigravityProviderBridge) ProviderKind() channel.ProviderKind { + return channel.ProviderAntigravity +} + +func (*antigravityProviderBridge) UpstreamProtocol() protocol.Protocol { return protocol.Gemini } + +func (*antigravityProviderBridge) ValidateRouteCapability(route channel.RouteDescriptor) error { + valid := route.ClientProtocol == protocol.Gemini && + (route.Operation == execution.OperationChatCompletion || route.Operation == execution.OperationCountTokens) && + route.RouteMode == execution.RouteNative + if route.ClientProtocol == protocol.Anthropic { + valid = (route.Operation == execution.OperationChatCompletion || route.Operation == execution.OperationCountTokens) && + route.RouteMode == execution.RouteConverted + } + if route.ClientProtocol == protocol.OpenAICompletions { + valid = route.Operation == execution.OperationChatCompletion && route.RouteMode == execution.RouteConverted + } + if route.ClientProtocol == protocol.OpenAIResponses { + valid = (route.Operation == execution.OperationResponsesCreate || route.Operation == execution.OperationResponsesInputTokens) && + route.RouteMode == execution.RouteConverted + } + if !valid { + return fmt.Errorf("route is not implemented by Antigravity") + } + return nil +} + +func (*antigravityProviderBridge) ParseCredential(raw []byte) (providerCredential, error) { + credential, err := antigravity.ParseCredentialJSON(raw) + if err != nil { + return nil, err + } + return antigravityProviderCredential{value: credential}, nil +} + +// ValidateRequest rejects only inputs CPA is known to silently omit or change +// for Antigravity. General protocol validation stays with the existing +// dialect layer so this does not become a second request schema. +func (*antigravityProviderBridge) ValidateRequest(request providerRequest) error { + if strings.TrimSpace(request.Format) == "openai-response" && + strings.Contains(strings.ToLower(strings.TrimSpace(request.Model)), "image") { + return errors.New("Antigravity does not support Responses image output") + } + var payload any + if err := json.Unmarshal(request.Payload, &payload); err != nil { + return errors.New("Antigravity request payload is invalid") + } + if strings.TrimSpace(request.Format) == "openai-response" { + if root, ok := payload.(map[string]any); ok { + if tools, ok := root["tools"].([]any); ok { + for _, rawTool := range tools { + tool, ok := rawTool.(map[string]any) + if !ok { + continue + } + toolType, _ := tool["type"].(string) + if toolType = strings.TrimSpace(toolType); toolType != "" && toolType != "function" { + return errors.New("Antigravity does not support Responses built-in tools") + } + } + } + } + } + if antigravityPayloadContainsRemoteImage(payload, strings.TrimSpace(request.Format)) { + return errors.New("Antigravity does not support remote image URLs") + } + return nil +} + +func antigravityPayloadContainsRemoteImage(value any, format string) bool { + switch current := value.(type) { + case []any: + for _, item := range current { + if antigravityPayloadContainsRemoteImage(item, format) { + return true + } + } + case map[string]any: + itemType, _ := current["type"].(string) + itemType = strings.TrimSpace(itemType) + if format == "openai" || format == "openai-response" { + if itemType == "image_url" || itemType == "input_image" { + if antigravityRemoteImageURL(current["image_url"]) { + return true + } + } + } + if format == "claude" && itemType == "image" { + if source, ok := current["source"].(map[string]any); ok { + sourceType, _ := source["type"].(string) + if strings.EqualFold(strings.TrimSpace(sourceType), "url") && antigravityRemoteImageURL(source["url"]) { + return true + } + } + } + for _, child := range current { + if antigravityPayloadContainsRemoteImage(child, format) { + return true + } + } + } + return false +} + +func antigravityRemoteImageURL(value any) bool { + urlValue, ok := value.(string) + if !ok { + if object, objectOK := value.(map[string]any); objectOK { + urlValue, ok = object["url"].(string) + } + } + urlValue = strings.ToLower(strings.TrimSpace(urlValue)) + return ok && (strings.HasPrefix(urlValue, "https://") || strings.HasPrefix(urlValue, "http://")) +} + +func (bridge *antigravityProviderBridge) Execute( + ctx context.Context, + credentialID string, + credential providerCredential, + request providerRequest, +) (providerResponse, error) { + value, ok := credential.(antigravityProviderCredential) + if !ok || bridge == nil || bridge.executor == nil { + return providerResponse{}, errors.New("Antigravity provider bridge credential mismatch") + } + response, err := bridge.executor.Execute(ctx, credentialID, value.value, antigravity.ExecuteRequest{ + AttemptID: request.AttemptID, Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: request.Format, + Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), + ContinuityKey: antigravityContinuityScope(request.ContinuityKey, credentialID, request.Model, request.AttemptID), + }) + if err == nil { + response.Payload, err = normalizeAntigravityResponseModel(request.Format, response.Payload, request.Model) + } + return providerResponse{ + Payload: append([]byte(nil), response.Payload...), Headers: response.Headers.Clone(), + AppliedReasoningEffort: response.AppliedReasoningEffort, + }, err +} + +func (bridge *antigravityProviderBridge) CountTokens( + ctx context.Context, + credentialID string, + credential providerCredential, + request providerRequest, +) (providerResponse, error) { + value, ok := credential.(antigravityProviderCredential) + if !ok || bridge == nil || bridge.executor == nil { + return providerResponse{}, errors.New("Antigravity provider bridge credential mismatch") + } + response, err := bridge.executor.CountTokens(ctx, credentialID, value.value, antigravity.ExecuteRequest{ + AttemptID: request.AttemptID, Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: request.Format, + Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), + ContinuityKey: antigravityContinuityScope(request.ContinuityKey, credentialID, request.Model, request.AttemptID), + }) + return providerResponse{Payload: append([]byte(nil), response.Payload...), Headers: response.Headers.Clone()}, err +} + +func (bridge *antigravityProviderBridge) ExecuteStream( + ctx context.Context, + credentialID string, + credential providerCredential, + request providerRequest, +) (*providerStreamResponse, error) { + value, ok := credential.(antigravityProviderCredential) + if !ok || bridge == nil || bridge.executor == nil { + return nil, errors.New("Antigravity provider bridge credential mismatch") + } + response, err := bridge.executor.ExecuteStream(ctx, credentialID, value.value, antigravity.ExecuteRequest{ + AttemptID: request.AttemptID, Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: request.Format, + Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), + ContinuityKey: antigravityContinuityScope(request.ContinuityKey, credentialID, request.Model, request.AttemptID), + }) + if response == nil { + return nil, err + } + if err != nil { + return &providerStreamResponse{ + Headers: response.Headers.Clone(), AppliedReasoningEffort: response.AppliedReasoningEffort, + }, err + } + chunks := make(chan providerStreamChunk) + go func() { + defer close(chunks) + for chunk := range response.Chunks { + converted := providerStreamChunk{Payload: append([]byte(nil), chunk.Payload...), Err: chunk.Err} + if converted.Err == nil { + converted.Payload, converted.Err = normalizeAntigravityResponseModel(request.Format, converted.Payload, request.Model) + } + select { + case chunks <- converted: + case <-ctx.Done(): + return + } + } + }() + return &providerStreamResponse{ + Headers: response.Headers.Clone(), Chunks: chunks, AppliedReasoningEffort: response.AppliedReasoningEffort, + }, nil +} + +func normalizeAntigravityResponseModel(format string, payload []byte, model string) ([]byte, error) { + if len(payload) == 0 || strings.TrimSpace(model) == "" { + return bytes.Clone(payload), nil + } + if !antigravityResponseUsesPlaceholderModel(payload) { + return bytes.Clone(payload), nil + } + var clientProtocol protocol.Protocol + switch strings.TrimSpace(format) { + case "openai": + clientProtocol = protocol.OpenAICompletions + case "openai-response": + clientProtocol = protocol.OpenAIResponses + case "claude": + clientProtocol = protocol.Anthropic + case "gemini": + clientProtocol = protocol.Gemini + default: + return nil, fmt.Errorf("unsupported Antigravity response format") + } + trimmed := bytes.TrimSpace(payload) + if json.Valid(trimmed) { + return responsealias.RewriteJSON(clientProtocol, payload, model) + } + if bytes.Equal(trimmed, []byte("[DONE]")) { + return bytes.Clone(payload), nil + } + return responsealias.RewriteSSE(clientProtocol, payload, model) +} + +func antigravityResponseUsesPlaceholderModel(payload []byte) bool { + trimmed := bytes.TrimSpace(payload) + if json.Valid(trimmed) { + return antigravityJSONUsesPlaceholderModel(trimmed) + } + for _, line := range bytes.Split(payload, []byte{'\n'}) { + line = bytes.TrimSpace(line) + if !bytes.HasPrefix(line, []byte("data:")) { + continue + } + data := bytes.TrimSpace(bytes.TrimPrefix(line, []byte("data:"))) + if json.Valid(data) && antigravityJSONUsesPlaceholderModel(data) { + return true + } + } + return false +} + +func antigravityJSONUsesPlaceholderModel(payload []byte) bool { + var root map[string]json.RawMessage + if json.Unmarshal(payload, &root) != nil || root == nil { + return false + } + if antigravityRawModelIsPlaceholder(root["model"]) || antigravityRawModelIsPlaceholder(root["modelVersion"]) { + return true + } + for _, field := range []string{"message", "response"} { + var nested map[string]json.RawMessage + if json.Unmarshal(root[field], &nested) == nil && antigravityRawModelIsPlaceholder(nested["model"]) { + return true + } + } + return false +} + +func antigravityRawModelIsPlaceholder(raw json.RawMessage) bool { + var value string + return json.Unmarshal(raw, &value) == nil && strings.EqualFold(strings.TrimSpace(value), "gemini-default") +} + +func (*antigravityProviderBridge) ClassifyError( + ctx context.Context, + err error, + credential providerCredential, +) (int, *execution.ErrorEvidence) { + if err == nil { + return 0, nil + } + status := 0 + var statusError interface{ StatusCode() int } + if errors.As(err, &statusError) && statusError != nil { + status = statusError.StatusCode() + } + kind := execution.ErrorKindTransport + if errors.Is(err, context.DeadlineExceeded) || ctx != nil && errors.Is(context.Cause(ctx), context.DeadlineExceeded) { + kind = execution.ErrorKindTimeout + } else if errors.Is(err, context.Canceled) || ctx != nil && errors.Is(context.Cause(ctx), context.Canceled) { + kind = execution.ErrorKindCanceled + } else if status != 0 { + kind = execution.ErrorKindHTTP + } else if _, ok := err.(net.Error); ok { + kind = execution.ErrorKindTransport + } + typeValue, codeValue := antigravityErrorTypeCode(err) + evidence := &execution.ErrorEvidence{ + Kind: kind, StatusCode: status, Type: typeValue, Code: codeValue, + Summary: antigravityProviderErrorSummary(status, typeValue, codeValue), + } + var retry interface{ RetryAfter() *time.Duration } + if errors.As(err, &retry) && retry != nil { + if value := retry.RetryAfter(); value != nil && *value > 0 { + evidence.RetryAfter = *value + } + } + switch { + case status == http.StatusUnauthorized: + evidence.Hint = execution.FailureHintRefreshRequired + evidence.ReplaySafety = execution.ReplaySafetyRejectedBeforeProcessing + case status == http.StatusForbidden || + status == http.StatusTooManyRequests && strings.EqualFold(codeValue, "INSUFFICIENT_G1_CREDITS_BALANCE"): + evidence.Hint = execution.FailureHintCandidateUnavailable + evidence.ReplaySafety = execution.ReplaySafetyRejectedBeforeProcessing + case requestScopedFailure(err): + evidence.Hint = execution.FailureHintRequestRejected + case status == http.StatusTooManyRequests: + evidence.Hint = execution.FailureHintRateLimited + case status >= http.StatusInternalServerError: + evidence.Hint = execution.FailureHintHostError + } + return status, evidence +} + +func antigravityProviderErrorSummary(status int, typeValue, codeValue string) string { + switch { + case status == http.StatusUnauthorized: + return "Antigravity authorization was rejected" + case status == http.StatusForbidden: + return "Antigravity access was denied" + case status == http.StatusTooManyRequests && strings.EqualFold(codeValue, "INSUFFICIENT_G1_CREDITS_BALANCE"): + return "Antigravity Google One AI credits are unavailable" + case status == http.StatusTooManyRequests: + return "Antigravity upstream rate limit was reached" + case status >= http.StatusInternalServerError: + return "Antigravity upstream service failed" + case status >= http.StatusBadRequest: + return "Antigravity upstream request was rejected" + case typeValue != "": + return "Antigravity upstream request failed" + default: + return "Antigravity upstream request failed" + } +} + +func antigravityErrorTypeCode(err error) (string, string) { + var typed interface { + ErrorType() string + ErrorCode() string + } + if errors.As(err, &typed) && typed != nil { + return safeScalar(typed.ErrorType()), safeScalar(typed.ErrorCode()) + } + return errorTypeCode(strings.TrimSpace(err.Error())) +} + +func antigravityContinuityScope(base, credentialID, model, attemptID string) string { + base = strings.TrimSpace(base) + credentialID = strings.TrimSpace(credentialID) + model = strings.TrimSpace(model) + attemptID = strings.TrimSpace(attemptID) + if credentialID == "" || model == "" { + return "" + } + if base == "" { + if attemptID == "" { + return "" + } + return strings.Join([]string{"gpt-load-antigravity-attempt", attemptID, credentialID, model}, "\x00") + } + return strings.Join([]string{"gpt-load-antigravity", base, credentialID, model}, "\x00") +} + +var _ providerBridge = (*antigravityProviderBridge)(nil) +var _ providerTokenCounter = (*antigravityProviderBridge)(nil) +var _ providerRequestValidator = (*antigravityProviderBridge)(nil) diff --git a/internal/execution/cpa/antigravity_provider_test.go b/internal/execution/cpa/antigravity_provider_test.go new file mode 100644 index 00000000..920cab4f --- /dev/null +++ b/internal/execution/cpa/antigravity_provider_test.go @@ -0,0 +1,313 @@ +package cpa + +import ( + "context" + "strings" + "testing" + "time" + + "gpt-load/internal/antigravity" + "gpt-load/internal/execution" +) + +type recordingAntigravityExecutor struct { + request antigravity.ExecuteRequest + response []byte + streamChunks [][]byte +} + +type antigravityClassifiedTestError struct { + status int + typeID string + code string +} + +type antigravityRetryAfterTestError struct { + antigravityClassifiedTestError + retryAfter time.Duration +} + +func (err antigravityRetryAfterTestError) RetryAfter() *time.Duration { + return &err.retryAfter +} + +func (err antigravityClassifiedTestError) Error() string { return "provider body must not escape" } +func (err antigravityClassifiedTestError) StatusCode() int { return err.status } +func (err antigravityClassifiedTestError) ErrorType() string { + return err.typeID +} +func (err antigravityClassifiedTestError) ErrorCode() string { return err.code } + +func (executor *recordingAntigravityExecutor) Execute( + _ context.Context, + _ string, + _ antigravity.Credential, + request antigravity.ExecuteRequest, +) (antigravity.ExecuteResponse, error) { + executor.request = request + payload := executor.response + if payload == nil { + payload = []byte(`{"ok":true}`) + } + return antigravity.ExecuteResponse{Payload: append([]byte(nil), payload...)}, nil +} + +func (executor *recordingAntigravityExecutor) CountTokens( + ctx context.Context, + credentialID string, + credential antigravity.Credential, + request antigravity.ExecuteRequest, +) (antigravity.ExecuteResponse, error) { + return executor.Execute(ctx, credentialID, credential, request) +} + +func (executor *recordingAntigravityExecutor) ExecuteStream( + ctx context.Context, + _ string, + _ antigravity.Credential, + request antigravity.ExecuteRequest, +) (*antigravity.ExecuteStreamResponse, error) { + executor.request = request + chunks := make(chan antigravity.ExecuteStreamChunk) + go func() { + defer close(chunks) + for _, payload := range executor.streamChunks { + select { + case chunks <- antigravity.ExecuteStreamChunk{Payload: append([]byte(nil), payload...)}: + case <-ctx.Done(): + return + } + } + }() + return &antigravity.ExecuteStreamResponse{Chunks: chunks}, nil +} + +func antigravityProviderTestCredential() antigravityProviderCredential { + return antigravityProviderCredential{value: antigravity.Credential{ + Type: "antigravity", AccessToken: "access", RefreshToken: "refresh", AccountID: "account", + Email: "owner@example.com", ProjectID: "project", Expire: "2030-01-01T00:00:00Z", + }} +} + +func TestAntigravityProviderNormalizesPlaceholderResponseModels(t *testing.T) { + tests := []struct { + name string + format string + body string + }{ + {name: "OpenAI Chat", format: "openai", body: `{"model":"gemini-default","choices":[]}`}, + {name: "OpenAI Responses", format: "openai-response", body: `{"model":"gemini-default","object":"response"}`}, + {name: "Anthropic", format: "claude", body: `{"model":"gemini-default","type":"message"}`}, + {name: "Gemini", format: "gemini", body: `{"modelVersion":"gemini-default","candidates":[]}`}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + executor := &recordingAntigravityExecutor{response: []byte(test.body)} + response, err := (&antigravityProviderBridge{executor: executor}).Execute( + t.Context(), "17", antigravityProviderTestCredential(), + providerRequest{Model: "gemini-3-flash-agent", Payload: []byte(`{}`), Format: test.format}, + ) + if err != nil { + t.Fatal(err) + } + if got := string(response.Payload); !strings.Contains(got, "gemini-3-flash-agent") || strings.Contains(got, "gemini-default") { + t.Fatalf("response payload = %s", got) + } + }) + } +} + +func TestAntigravityProviderPreservesActualReportedModel(t *testing.T) { + executor := &recordingAntigravityExecutor{response: []byte(`{"model":"routed-model","choices":[]}`)} + response, err := (&antigravityProviderBridge{executor: executor}).Execute( + t.Context(), "17", antigravityProviderTestCredential(), + providerRequest{Model: "requested-model", Payload: []byte(`{}`), Format: "openai"}, + ) + if err != nil { + t.Fatal(err) + } + if got := string(response.Payload); !strings.Contains(got, "routed-model") || strings.Contains(got, "requested-model") { + t.Fatalf("response payload = %s", got) + } +} + +func TestAntigravityProviderNormalizesPlaceholderStreamModels(t *testing.T) { + tests := []struct { + name string + format string + chunk string + }{ + {name: "OpenAI Chat", format: "openai", chunk: "data: {\"model\":\"gemini-default\",\"choices\":[]}\n\n"}, + {name: "OpenAI Responses", format: "openai-response", chunk: "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"model\":\"gemini-default\"}}\n\n"}, + {name: "Anthropic", format: "claude", chunk: "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"model\":\"gemini-default\"}}\n\n"}, + {name: "Gemini", format: "gemini", chunk: `{"modelVersion":"gemini-default","candidates":[]}`}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + executor := &recordingAntigravityExecutor{streamChunks: [][]byte{[]byte(test.chunk)}} + response, err := (&antigravityProviderBridge{executor: executor}).ExecuteStream( + t.Context(), "17", antigravityProviderTestCredential(), + providerRequest{Model: "gemini-3-flash-agent", Payload: []byte(`{}`), Format: test.format}, + ) + if err != nil { + t.Fatal(err) + } + var payload strings.Builder + for chunk := range response.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + _, _ = payload.Write(chunk.Payload) + } + if value := payload.String(); !strings.Contains(value, "gemini-3-flash-agent") || strings.Contains(value, "gemini-default") { + t.Fatalf("stream payload = %s", value) + } + }) + } +} + +func TestAntigravityProviderPreservesActualReportedStreamModel(t *testing.T) { + executor := &recordingAntigravityExecutor{streamChunks: [][]byte{ + []byte("data: {\"model\":\"routed-model\",\"choices\":[]}\n\n"), + }} + response, err := (&antigravityProviderBridge{executor: executor}).ExecuteStream( + t.Context(), "17", antigravityProviderTestCredential(), + providerRequest{Model: "requested-model", Payload: []byte(`{}`), Format: "openai"}, + ) + if err != nil { + t.Fatal(err) + } + var payload strings.Builder + for chunk := range response.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + _, _ = payload.Write(chunk.Payload) + } + if got := payload.String(); !strings.Contains(got, "routed-model") || strings.Contains(got, "requested-model") { + t.Fatalf("stream payload = %s", got) + } +} + +func TestAntigravityProviderScopesContinuityPerCredentialAndModel(t *testing.T) { + executor := &recordingAntigravityExecutor{} + bridge := &antigravityProviderBridge{executor: executor} + _, err := bridge.Execute(context.Background(), "17", antigravityProviderCredential{value: antigravity.Credential{ + Type: "antigravity", AccessToken: "access", RefreshToken: "refresh", AccountID: "account", + Email: "owner@example.com", ProjectID: "project", Expire: "2030-01-01T00:00:00Z", + }}, providerRequest{Model: "gemini-live", Payload: []byte(`{}`), Format: "gemini", ContinuityKey: "tenant-hmac"}) + if err != nil { + t.Fatal(err) + } + if executor.request.ContinuityKey == "tenant-hmac" || !strings.Contains(executor.request.ContinuityKey, "tenant-hmac") || + !strings.Contains(executor.request.ContinuityKey, "17") || !strings.Contains(executor.request.ContinuityKey, "gemini-live") { + t.Fatalf("continuity key = %q", executor.request.ContinuityKey) + } +} + +func TestAntigravityProviderRejectsKnownUnsupportedInputs(t *testing.T) { + bridge := &antigravityProviderBridge{} + validator, ok := any(bridge).(providerRequestValidator) + if !ok { + t.Fatal("Antigravity bridge must validate provider-specific unsupported inputs") + } + + tests := []struct { + name string + format string + payload string + wantErr bool + }{ + { + name: "OpenAI chat remote image", format: "openai", wantErr: true, + payload: `{"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":"https://example.test/image.png"}}]}]}`, + }, + { + name: "Responses remote image", format: "openai-response", wantErr: true, + payload: `{"input":[{"role":"user","content":[{"type":"input_image","image_url":"https://example.test/image.png"}]}]}`, + }, + { + name: "Responses builtin tool", format: "openai-response", wantErr: true, + payload: `{"tools":[{"type":"web_search_preview"}],"input":"hello"}`, + }, + { + name: "Responses image output model", format: "openai-response", wantErr: true, + payload: `{"input":"draw a logo"}`, + }, + { + name: "Anthropic remote image", format: "claude", wantErr: true, + payload: `{"messages":[{"role":"user","content":[{"type":"image","source":{"type":"url","url":"https://example.test/image.png"}}]}]}`, + }, + { + name: "Supported data URL and function tool", format: "openai-response", + payload: `{"tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}],"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}]}`, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + model := "gemini-live" + if test.name == "Responses image output model" { + model = "gemini-3.1-flash-image" + } + err := validator.ValidateRequest(providerRequest{Model: model, Format: test.format, Payload: []byte(test.payload)}) + if (err != nil) != test.wantErr { + t.Fatalf("ValidateRequest() error = %v, want error = %t", err, test.wantErr) + } + }) + } +} + +func TestAntigravityProviderTreatsForbiddenAsCandidateUnavailable(t *testing.T) { + bridge := &antigravityProviderBridge{} + _, evidence := bridge.ClassifyError(t.Context(), statusError{ + status: 403, message: `{"error":{"status":"PERMISSION_DENIED"}}`, + }, antigravityProviderCredential{value: antigravity.Credential{ + Type: "antigravity", AccessToken: "access", RefreshToken: "refresh", AccountID: "account", + Email: "owner@example.com", ProjectID: "project", Expire: "2030-01-01T00:00:00Z", + }}) + if evidence == nil || evidence.Hint != execution.FailureHintCandidateUnavailable || + evidence.ReplaySafety != execution.ReplaySafetyRejectedBeforeProcessing || + strings.Contains(evidence.Summary, "PERMISSION_DENIED") { + t.Fatalf("error evidence = %#v", evidence) + } +} + +func TestAntigravityProviderClassifiesOAuthAndPaidCreditErrorsWithoutCredentialPenalty(t *testing.T) { + bridge := &antigravityProviderBridge{} + credential := antigravityProviderCredential{value: antigravity.Credential{ + Type: "antigravity", AccessToken: "access", RefreshToken: "refresh", AccountID: "account", + Email: "owner@example.com", ProjectID: "project", Expire: "2030-01-01T00:00:00Z", + }} + tests := []struct { + name string + err error + wantHint execution.FailureHint + wantReplay execution.ReplaySafety + wantRetry time.Duration + }{ + { + name: "OAuth 401 refreshes once", err: antigravityClassifiedTestError{status: 401, typeID: "UNAUTHENTICATED"}, + wantHint: execution.FailureHintRefreshRequired, wantReplay: execution.ReplaySafetyRejectedBeforeProcessing, + }, + { + name: "paid credit balance advances to another credential", err: antigravityClassifiedTestError{status: 429, typeID: "RESOURCE_EXHAUSTED", code: "INSUFFICIENT_G1_CREDITS_BALANCE"}, + wantHint: execution.FailureHintCandidateUnavailable, wantReplay: execution.ReplaySafetyRejectedBeforeProcessing, + }, + { + name: "rate limit retains fractional retry delay", err: antigravityRetryAfterTestError{ + antigravityClassifiedTestError: antigravityClassifiedTestError{status: 429, typeID: "RESOURCE_EXHAUSTED", code: "RATE_LIMIT_EXCEEDED"}, + retryAfter: 450 * time.Millisecond, + }, + wantHint: execution.FailureHintRateLimited, wantRetry: 450 * time.Millisecond, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, evidence := bridge.ClassifyError(t.Context(), test.err, credential) + if evidence == nil || evidence.Hint != test.wantHint || evidence.ReplaySafety != test.wantReplay || + strings.Contains(evidence.Summary, "provider body") || evidence.RetryAfter != test.wantRetry { + t.Fatalf("error evidence = %#v", evidence) + } + }) + } +} diff --git a/internal/execution/cpa/provider.go b/internal/execution/cpa/provider.go index 157d8c26..4f2c912e 100644 --- a/internal/execution/cpa/provider.go +++ b/internal/execution/cpa/provider.go @@ -46,11 +46,15 @@ func credentialScopedFailure(err error) (bool, bool) { } type providerRequest struct { + AttemptID string Model string Payload []byte Format string Headers http.Header OriginalRequest []byte + // ContinuityKey is a private, tenant-scoped key used only by providers + // whose tool/thinking protocol needs an isolated multi-request replay lane. + ContinuityKey string } type providerResponse struct { @@ -89,6 +93,14 @@ type providerTokenCounter interface { CountTokens(context.Context, string, providerCredential, providerRequest) (providerResponse, error) } +// providerRequestValidator is an optional, provider-specific pre-dispatch +// check. It is intentionally narrow: the shared adapter only needs to reject +// inputs that a declared route cannot faithfully represent before preparing a +// subscription credential or sending anything upstream. +type providerRequestValidator interface { + ValidateRequest(providerRequest) error +} + // providerLocalTokenCounter is a deliberately narrow contract for providers // whose CountTokens implementation never contacts an upstream service. type providerLocalTokenCounter interface { diff --git a/internal/gateway/affinity.go b/internal/gateway/affinity.go index af9ae74c..3f7c4795 100644 --- a/internal/gateway/affinity.go +++ b/internal/gateway/affinity.go @@ -11,6 +11,7 @@ type requestAffinity struct { key affinity.Key observation affinity.Observation preferredCredentialID uint + continuityKey string } func (handler *Handler) resolveRequestAffinity( @@ -20,12 +21,7 @@ func (handler *Handler) resolveRequestAffinity( prefix []byte, allowedCredentialRefs map[uint]state.CredentialRef, ) requestAffinity { - if handler == nil || handler.affinityCache == nil || snapshot == nil || - !handler.affinityCache.Configure( - snapshot.Revision, - snapshot.Settings.AffinityCapacity, - snapshot.Settings.AffinityTTL, - ) { + if handler == nil || snapshot == nil { return requestAffinity{} } key := affinity.DeriveKey( @@ -34,11 +30,22 @@ func (handler *Handler) resolveRequestAffinity( clientProtocol, prefix, ) - if !key.Valid() { - return requestAffinity{} + result := requestAffinity{continuityKey: string(key)} + if handler.affinityCache == nil || + !handler.affinityCache.Configure( + snapshot.Revision, + snapshot.Settings.AffinityCapacity, + snapshot.Settings.AffinityTTL, + ) { + return result } + if !key.Valid() { + return result + } + result.key = key observation := handler.affinityCache.Lookup(key) - resolved := requestAffinity{key: key, observation: observation} + resolved := result + resolved.observation = observation if !observation.Found() { return resolved } diff --git a/internal/gateway/affinity_test.go b/internal/gateway/affinity_test.go index 62a72bd3..8765d3e6 100644 --- a/internal/gateway/affinity_test.go +++ b/internal/gateway/affinity_test.go @@ -301,6 +301,22 @@ func TestHandlerIgnoresAffinityAfterCredentialIdentityChanges(t *testing.T) { } } +func TestHandlerDerivesPrivateContinuityWithoutReenablingDisabledAffinity(t *testing.T) { + handler, manager, _ := newHandlerForTest(t, &scriptedForwarder{}, "sk-one") + snapshot := manager.Current() + snapshot.Settings.AffinityCapacity = 0 + resolved := handler.resolveRequestAffinity( + snapshot, + 1, + protocol.OpenAICompletions, + []byte(`{"v":1,"user":["hello"]}`), + map[uint]state.CredentialRef{1: {ID: 1, GroupID: 1, IdentityGeneration: 1}}, + ) + if resolved.key.Valid() || resolved.preferredCredentialID != 0 || resolved.continuityKey == "" { + t.Fatalf("disabled affinity resolution = %#v", resolved) + } +} + func successfulAffinityResults(count int) []UpstreamResult { results := make([]UpstreamResult, count) for index := range results { diff --git a/internal/gateway/execution_boundary_test.go b/internal/gateway/execution_boundary_test.go index 7340b92b..4ffe5cce 100644 --- a/internal/gateway/execution_boundary_test.go +++ b/internal/gateway/execution_boundary_test.go @@ -84,6 +84,7 @@ func testCredentialRuntimes(t *testing.T) (*channel.Registry, *subscriptionrunti channels, subscriptionruntime.CodexImplementations(), subscriptionruntime.ClaudeImplementations(), + subscriptionruntime.AntigravityImplementations(), ) if err != nil { t.Fatal(err) diff --git a/internal/gateway/execution_forward.go b/internal/gateway/execution_forward.go index 1d00b732..25cb01bd 100644 --- a/internal/gateway/execution_forward.go +++ b/internal/gateway/execution_forward.go @@ -581,6 +581,7 @@ func newExecutionAttemptSpec(input ForwardInput) (execution.AttemptSpec, error) Body: input.Request.Body, IncludeUsage: input.ObserveUsage && input.Group.InjectUsageOptions, ForceCredentialRefresh: input.ForceCredentialRefresh, + ContinuityKey: input.ContinuityKey, TargetConfig: input.TargetConfig, Timeouts: execution.AttemptTimeouts{ FirstByte: input.Group.Timeouts.FirstByte, diff --git a/internal/gateway/execution_forward_test.go b/internal/gateway/execution_forward_test.go index f231bd8a..3c5b4fa1 100644 --- a/internal/gateway/execution_forward_test.go +++ b/internal/gateway/execution_forward_test.go @@ -1,7 +1,9 @@ package gateway import ( + "bytes" "context" + "encoding/json" "errors" "net/http" "net/http/httptest" @@ -801,6 +803,25 @@ func executionForwardInput() ForwardInput { } } +func TestNewExecutionAttemptSpecKeepsContinuityPrivate(t *testing.T) { + input := executionForwardInput() + input.ContinuityKey = "tenant-scoped-hmac" + spec, err := newExecutionAttemptSpec(input) + if err != nil { + t.Fatal(err) + } + if spec.ContinuityKey != input.ContinuityKey { + t.Fatalf("ContinuityKey = %q, want %q", spec.ContinuityKey, input.ContinuityKey) + } + raw, err := json.Marshal(spec) + if err != nil { + t.Fatal(err) + } + if bytes.Contains(raw, []byte(input.ContinuityKey)) { + t.Fatalf("AttemptSpec JSON exposed continuity key: %s", raw) + } +} + func responsesExecutionForwardInput() ForwardInput { input := executionForwardInput() input.Dialect = dialect.NewOpenAIResponses() diff --git a/internal/gateway/forward.go b/internal/gateway/forward.go index 9ff324c5..cec39972 100644 --- a/internal/gateway/forward.go +++ b/internal/gateway/forward.go @@ -49,6 +49,9 @@ type ForwardInput struct { // ForceCredentialRefresh preserves one explicit provider-auth retry as a // globally counted Attempt on the same selected credential. ForceCredentialRefresh bool + // ContinuityKey is an opaque per-tenant replay boundary for provider-private + // thinking and tool state. It never crosses the gateway DTO boundary. + ContinuityKey string } // UpstreamResult is the gateway's stable view of one logical execution diff --git a/internal/gateway/handler.go b/internal/gateway/handler.go index a8bddd8d..df876e01 100644 --- a/internal/gateway/handler.go +++ b/internal/gateway/handler.go @@ -149,6 +149,7 @@ func NewHandler( channels, subscriptionruntime.CodexImplementations(), subscriptionruntime.ClaudeImplementations(), + subscriptionruntime.AntigravityImplementations(), ) return &Handler{ manager: manager, channels: channels, subscriptions: subscriptions, registry: registry, encryption: encryptionService, @@ -690,6 +691,7 @@ func (handler *Handler) executeAttempts( normalizedCredential.payload, ), ForceCredentialRefresh: forceCredentialRefresh, + ContinuityKey: requestAffinity.continuityKey, OnFirstResponse: func() { recorder.recordFirstResponse() }, diff --git a/internal/gateway/handler_test.go b/internal/gateway/handler_test.go index 62b2b531..516c195f 100644 --- a/internal/gateway/handler_test.go +++ b/internal/gateway/handler_test.go @@ -981,6 +981,54 @@ func TestHandlerDoesNotRotateOrPenalizeRequestRejected429(t *testing.T) { } } +func TestHandlerRetriesCandidateUnavailableWithoutCredentialPenalty(t *testing.T) { + now := time.Date(2026, time.August, 18, 9, 0, 0, 0, time.UTC) + forwarder := &scriptedForwarder{results: []UpstreamResult{ + { + DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + StatusCode: http.StatusForbidden, + Header: http.Header{"Content-Type": {"application/json"}}, + Body: []byte(`{"error":{"type":"PERMISSION_DENIED"}}`), + ExecutionError: &execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, + Hint: execution.FailureHintCandidateUnavailable, + StatusCode: http.StatusForbidden, + ReplaySafety: execution.ReplaySafetyRejectedBeforeProcessing, + Summary: "candidate is unavailable", + }, + }, + {StatusCode: http.StatusOK, Header: make(http.Header), Body: []byte(`{"ok":true}`), RequestWritten: true}, + }} + engine, handler, registry, stats := newStatsHandlerTestRuntime(t, forwarder, "sk-first", "sk-second") + recording := &recordingRuntimeRegistry{CredentialRegistry: registry} + handler.registry = recording + handler.newRandom = func() *rand.Rand { return rand.New(zeroSource{}) } + handler.now = func() time.Time { return now } + + request := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + bytes.NewBufferString(`{"model":"gpt-4o"}`), + ) + request.Header.Set("Authorization", "Bearer gl-client") + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + + if recorder.Code != http.StatusOK || recorder.Body.String() != `{"ok":true}` || len(forwarder.inputs) != 2 { + t.Fatalf("response/attempts = %d %s / %d", recorder.Code, recorder.Body.String(), len(forwarder.inputs)) + } + if recording.cooldownCalls != 0 || recording.incrFailureCalls != 0 || recording.blacklistCalls != 0 { + t.Fatalf("credential mutations = cooldown:%d failure:%d blacklist:%d, want none", + recording.cooldownCalls, recording.incrFailureCalls, recording.blacklistCalls) + } + if got := stats.Snapshot(1, now); got != (health.CredentialStats{}) { + t.Fatalf("first credential stats = %#v, want empty", got) + } + if got := stats.Snapshot(2, now); got != (health.CredentialStats{Success: 1}) { + t.Fatalf("second credential stats = %#v, want success", got) + } +} + func TestHandlerRecordsStreamSuccessOnlyAfterCleanTerminal(t *testing.T) { now := time.Date(2026, time.July, 22, 10, 0, 0, 0, time.UTC) tests := []struct { @@ -3307,6 +3355,22 @@ func TestCrossCandidateRetryRespectsReplaySafety(t *testing.T) { }, decision: retry, }, + { + name: "candidate rejection safely advances to another credential", operation: execution.OperationChatCompletion, + method: http.MethodPost, + result: UpstreamResult{ + DispatchState: execution.DispatchMaybeSent, + StatusCode: http.StatusForbidden, + ExecutionError: &execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, + Hint: execution.FailureHintCandidateUnavailable, + StatusCode: http.StatusForbidden, + ReplaySafety: execution.ReplaySafetyRejectedBeforeProcessing, + }, + }, + decision: health.Result{Category: health.FailureCategoryModelUnavailable, Action: health.ActionRetry}, + want: true, + }, { name: "unsupported count tokens is terminal despite idempotent operation", operation: execution.OperationCountTokens, method: http.MethodPost, diff --git a/internal/health/execution_judge.go b/internal/health/execution_judge.go index 3260d666..bcc6aded 100644 --- a/internal/health/execution_judge.go +++ b/internal/health/execution_judge.go @@ -104,6 +104,8 @@ func classifyExecutionEvidence(attempt ExecutionAttempt) FailureCategory { return FailureCategoryRateLimited case execution.FailureHintRequestRejected: return FailureCategoryClientError + case execution.FailureHintCandidateUnavailable: + return FailureCategoryModelUnavailable case execution.FailureHintModelUnavailable: return FailureCategoryModelUnavailable case execution.FailureHintHostError: @@ -156,6 +158,9 @@ func structuredUnsupportedModelClientError(statusCode int, evidence *execution.E } func resultForExecutionCategory(category FailureCategory, attempt ExecutionAttempt) Result { + if attempt.Evidence != nil && attempt.Evidence.Hint == execution.FailureHintCandidateUnavailable { + return Result{Category: category, Action: ActionRetry} + } switch category { case FailureCategoryRateLimited: if attempt.Evidence != nil && attempt.Evidence.RetryAfter > 0 { diff --git a/internal/health/execution_judge_test.go b/internal/health/execution_judge_test.go index 5da1ba1c..a8a446a3 100644 --- a/internal/health/execution_judge_test.go +++ b/internal/health/execution_judge_test.go @@ -118,6 +118,22 @@ func TestJudgeExecutionUsesNeutralEvidenceAndReplayBoundary(t *testing.T) { }, want: Result{Category: FailureCategoryClientError, Action: ActionTerminate}, }, + { + name: "candidate rejection retries without credential penalty", + attempt: ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + StatusCode: http.StatusForbidden, + Now: now, + Evidence: &execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, + Hint: execution.FailureHintCandidateUnavailable, + StatusCode: http.StatusForbidden, + ReplaySafety: execution.ReplaySafetyRejectedBeforeProcessing, + Summary: "candidate is unavailable", + }, + }, + want: Result{Category: FailureCategoryModelUnavailable, Action: ActionRetry}, + }, { name: "unsupported upstream operation terminates without credential penalty", attempt: ExecutionAttempt{ diff --git a/internal/provideradapter/registry_test.go b/internal/provideradapter/registry_test.go index eaec088d..efe14d29 100644 --- a/internal/provideradapter/registry_test.go +++ b/internal/provideradapter/registry_test.go @@ -128,5 +128,6 @@ func completeBindings(bifrost execution.Executor, codex execution.Executor) []Bi {ProviderKind: channel.ProviderXAI, Adapter: bifrost}, {ProviderKind: channel.ProviderCodex, Adapter: codex}, {ProviderKind: channel.ProviderClaude, Adapter: codex}, + {ProviderKind: channel.ProviderAntigravity, Adapter: codex}, } } diff --git a/internal/state/loader/channel_loader_test.go b/internal/state/loader/channel_loader_test.go index b2169533..21ebe6db 100644 --- a/internal/state/loader/channel_loader_test.go +++ b/internal/state/loader/channel_loader_test.go @@ -106,6 +106,7 @@ func testSubscriptionRuntime(t *testing.T) (*channel.Registry, *subscriptionrunt channels, subscriptionruntime.CodexImplementations(), subscriptionruntime.ClaudeImplementations(), + subscriptionruntime.AntigravityImplementations(), ) if err != nil { t.Fatal(err) diff --git a/internal/subscription/credential_manager_test.go b/internal/subscription/credential_manager_test.go index a6c3398d..0d5a61b5 100644 --- a/internal/subscription/credential_manager_test.go +++ b/internal/subscription/credential_manager_test.go @@ -327,6 +327,7 @@ func newCredentialManagerFixture( channel.NewRegistry(), subscriptionruntime.CodexImplementations(), subscriptionruntime.ClaudeImplementations(), + subscriptionruntime.AntigravityImplementations(), ) if err != nil { t.Fatal(err) diff --git a/internal/subscription/runtime/antigravity_driver.go b/internal/subscription/runtime/antigravity_driver.go new file mode 100644 index 00000000..986efb6a --- /dev/null +++ b/internal/subscription/runtime/antigravity_driver.go @@ -0,0 +1,194 @@ +package subscriptionruntime + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + "gpt-load/internal/antigravity" + "gpt-load/internal/channel/modules" + "gpt-load/internal/channel/spec" +) + +type antigravityDriver struct{} + +func newAntigravityDriver() *antigravityDriver { return &antigravityDriver{} } + +// AntigravityImplementations returns the concrete Antigravity subscription +// behavior assembled by the application composition root. +func AntigravityImplementations() Implementations { + driver := newAntigravityDriver() + return Implementations{ + Drivers: []Driver{driver}, + ModelDiscoveries: []ModelDiscovery{driver.modelDiscovery()}, + QuotaObservations: []QuotaObservation{driver.quotaObservation()}, + } +} + +func (*antigravityDriver) ID() spec.SubscriptionDriverID { + return modules.AntigravitySubscriptionDriver +} + +func (*antigravityDriver) Parse(raw []byte) (Credential, error) { + value, err := antigravity.ParseCredentialJSON(raw) + if err != nil { + return Credential{}, err + } + canonical, err := antigravity.MarshalCredential(value) + if err != nil { + return Credential{}, err + } + return antigravityRuntimeCredential(value, canonical), nil +} + +func (*antigravityDriver) Refresh(ctx context.Context, current Credential) (Credential, error) { + value, err := antigravity.ParseCredentialJSON(current.canonical) + if err != nil { + return Credential{}, err + } + refreshed, err := antigravity.RefreshCredentialOnce(ctx, value) + if err != nil { + return Credential{}, err + } + canonical, err := antigravity.MarshalCredential(refreshed) + if err != nil { + return Credential{}, err + } + return antigravityRuntimeCredential(refreshed, canonical), nil +} + +func (*antigravityDriver) ClassifyRefreshFailure(err error) RefreshFailure { + if errors.Is(err, antigravity.ErrCredentialIdentityChanged) { + return RefreshFailureIdentityChanged + } + var tokenErr *antigravity.TokenEndpointError + if errors.As(err, &tokenErr) && antigravity.IsDefinitiveRefreshRejection(tokenErr.Code) { + return RefreshFailureReauthorizationRequired + } + return RefreshFailureOutcomeUnknown +} + +func (*antigravityDriver) BeginAuthorization() (Authorization, error) { + value, err := antigravity.BeginBrowserAuthorization() + if err != nil { + return Authorization{}, err + } + return Authorization{URL: value.AuthorizationURL, State: value.State, ExpiresAt: value.ExpiresAt}, nil +} + +func (*antigravityDriver) LocalCallback() (LocalCallbackSpec, bool) { + return LocalCallbackSpec{RedirectURI: antigravity.RedirectURI}, true +} + +func (*antigravityDriver) CompleteAuthorization(ctx context.Context, completion AuthorizationCompletion) (Credential, error) { + value, err := antigravity.CompleteBrowserAuthorization(ctx, antigravity.BrowserAuthorizationCompletion{ + ExpectedState: completion.ExpectedState, ReturnedState: completion.ReturnedState, Code: completion.Code, + }) + if err != nil { + return Credential{}, err + } + canonical, err := antigravity.MarshalCredential(value) + if err != nil { + return Credential{}, err + } + return antigravityRuntimeCredential(value, canonical), nil +} + +func (*antigravityDriver) AuthorizationFailureDefinitive(err error) bool { + var tokenErr *antigravity.TokenEndpointError + return errors.As(err, &tokenErr) && antigravity.IsDefinitiveRefreshRejection(tokenErr.Code) +} + +func (*antigravityDriver) ImportCredential(ctx context.Context, raw []byte) (Credential, error) { + value, err := antigravity.ImportCredential(ctx, raw) + if err != nil { + return Credential{}, err + } + canonical, err := antigravity.MarshalCredential(value) + if err != nil { + return Credential{}, err + } + return antigravityRuntimeCredential(value, canonical), nil +} + +func (*antigravityDriver) DiscoverModels(ctx context.Context, credential Credential) ([]string, error) { + value, err := antigravity.ParseCredentialJSON(credential.canonical) + if err != nil { + return nil, err + } + models, err := antigravity.ListModels(ctx, value) + if err != nil { + return nil, err + } + result := make([]string, 0, len(models)) + for _, model := range models { + result = append(result, model.ID) + } + return result, nil +} + +func (*antigravityDriver) Observe(ctx context.Context, credential Credential) (Observation, error) { + value, err := antigravity.ParseCredentialJSON(credential.canonical) + if err != nil { + return Observation{}, err + } + observed, err := antigravity.ObserveAccount(ctx, value) + if err != nil { + var upstream *antigravity.UpstreamHTTPError + if errors.As(err, &upstream) { + return Observation{}, &UpstreamHTTPError{StatusCode: upstream.StatusCode} + } + return Observation{}, err + } + accountObserved, quotaObserved, partial, err := antigravityObservationCompleteness(observed) + if err != nil { + return Observation{}, err + } + normalized, err := NormalizeAntigravityObservation(value.Email, observed) + if err != nil { + return Observation{}, fmt.Errorf("%w: %v", ErrObservationPayloadInvalid, err) + } + return Observation{ + Payload: normalized, Partial: partial, + AccountObserved: accountObserved, QuotaObserved: quotaObserved, + }, nil +} + +func antigravityObservationCompleteness(observed antigravity.AccountObservation) (bool, bool, bool, error) { + accountObserved := strings.TrimSpace(observed.PlanID) != "" + // A successfully observed plan with no Google One AI credit entry is an + // authoritative empty quota result, not a partial observation. + quotaObserved := accountObserved || observed.GoogleOneAICredits != nil + if !accountObserved && !quotaObserved { + return false, false, false, fmt.Errorf("%w: Antigravity account observation has no usable fields", ErrObservationPayloadInvalid) + } + return accountObserved, quotaObserved, !accountObserved || !quotaObserved, nil +} + +type antigravityModelDiscovery struct{ *antigravityDriver } +type antigravityQuotaObservation struct{ *antigravityDriver } + +func (driver *antigravityDriver) modelDiscovery() ModelDiscovery { + return antigravityModelDiscovery{driver} +} + +func (driver *antigravityDriver) quotaObservation() QuotaObservation { + return antigravityQuotaObservation{driver} +} + +func (antigravityModelDiscovery) ID() spec.UtilityID { return modules.AntigravityModelDiscovery } +func (antigravityQuotaObservation) ID() spec.UtilityID { return modules.AntigravityQuotaObservation } + +func antigravityRuntimeCredential(value antigravity.Credential, canonical []byte) Credential { + expiresAt, expires := antigravity.CredentialExpiresAt(value) + account := Account{Email: strings.TrimSpace(value.Email), ExpiresAt: expiresAt, ExpiresAtKnown: expires} + if refreshed, err := time.Parse(time.RFC3339, strings.TrimSpace(value.LastRefresh)); err == nil { + account.LastRefresh, account.LastRefreshKnown = refreshed, true + } + return newCredential(canonical, strings.TrimSpace(value.AccountID), account, expiresAt, expires, value.SecretValues()) +} + +var _ BrowserAuthorizationDriver = (*antigravityDriver)(nil) +var _ CredentialFileImporter = (*antigravityDriver)(nil) diff --git a/internal/subscription/runtime/antigravity_driver_test.go b/internal/subscription/runtime/antigravity_driver_test.go new file mode 100644 index 00000000..629de7e7 --- /dev/null +++ b/internal/subscription/runtime/antigravity_driver_test.go @@ -0,0 +1,144 @@ +package subscriptionruntime + +import ( + "encoding/json" + "errors" + "reflect" + "testing" + + "gpt-load/internal/antigravity" + "gpt-load/internal/channel/modules" +) + +func TestAntigravityDriverProducesProviderNeutralCredential(t *testing.T) { + driver := newAntigravityDriver() + credential, err := driver.Parse([]byte(`{ + "type":"antigravity", + "access_token":"access-secret", + "refresh_token":"refresh-secret", + "account_id":"google-account-one", + "email":"owner@example.com", + "project_id":"project-one", + "expired":"2030-01-01T00:00:00Z" + }`)) + if err != nil { + t.Fatal(err) + } + if credential.Identity() != "google-account-one" || credential.Account().Email != "owner@example.com" || + !credential.Account().ExpiresAtKnown { + t.Fatalf("credential metadata = %q %#v", credential.Identity(), credential.Account()) + } + if got := credential.SecretValues(); !reflect.DeepEqual(got, []string{ + "access-secret", "refresh-secret", "google-account-one", "owner@example.com", "project-one", + }) { + t.Fatalf("SecretValues() = %q", got) + } +} + +func TestAntigravityDriverDeclaresCallbackAndImporter(t *testing.T) { + driver := newAntigravityDriver() + callback, local := driver.LocalCallback() + if !local || callback.RedirectURI != "http://localhost:51121/oauth-callback" { + t.Fatalf("callback = %#v/%t", callback, local) + } + if _, ok := any(driver).(CredentialFileImporter); !ok { + t.Fatal("Antigravity driver does not implement CredentialFileImporter") + } +} + +func TestAntigravityDriverClassifiesRefreshFailures(t *testing.T) { + driver := newAntigravityDriver() + if got := driver.ClassifyRefreshFailure(antigravity.ErrCredentialIdentityChanged); got != RefreshFailureIdentityChanged { + t.Fatalf("identity changed classification = %v", got) + } + if got := driver.ClassifyRefreshFailure(&antigravity.TokenEndpointError{StatusCode: 400, Code: "invalid_grant"}); got != RefreshFailureReauthorizationRequired { + t.Fatalf("invalid_grant classification = %v", got) + } + if got := driver.ClassifyRefreshFailure(errors.New("network unavailable")); got != RefreshFailureOutcomeUnknown { + t.Fatalf("network classification = %v", got) + } +} + +func TestAntigravityImplementationsExposeCompleteReadOnlyCapabilities(t *testing.T) { + implementations := AntigravityImplementations() + if len(implementations.Drivers) != 1 || implementations.Drivers[0].ID() != modules.AntigravitySubscriptionDriver || + len(implementations.ModelDiscoveries) != 1 || + implementations.ModelDiscoveries[0].ID() != modules.AntigravityModelDiscovery || + len(implementations.QuotaObservations) != 1 || + implementations.QuotaObservations[0].ID() != modules.AntigravityQuotaObservation || + len(implementations.ResetCreditActions) != 0 { + t.Fatalf("implementations = %#v", implementations) + } +} + +func TestNormalizeAntigravityObservationShowsCreditsWithoutInventingQuota(t *testing.T) { + raw, err := NormalizeAntigravityObservation("owner@example.com", antigravity.AccountObservation{ + PlanID: "google-one-ai", + GoogleOneAICredits: &antigravity.GoogleOneAICredit{ + Amount: 25000, MinimumAmount: 50, + }, + }) + if err != nil { + t.Fatal(err) + } + var snapshot quotaSnapshot + if err := json.Unmarshal(raw, &snapshot); err != nil { + t.Fatal(err) + } + if snapshot.Plan.Name != "google-one-ai" || snapshot.Account == nil || snapshot.Account.Email != "owner@example.com" || + len(snapshot.QuotaWindows) != 1 { + t.Fatalf("snapshot = %#v", snapshot) + } + window := snapshot.QuotaWindows[0] + if window.ID != "google_one_ai" || window.Scope != "credits" || window.Unit != "credits" || + window.Remaining == nil || *window.Remaining != 25000 || window.State != "available" || + window.Limit != nil || window.Used != nil || window.Utilization != nil || window.ResetAtMS != nil { + t.Fatalf("credit window = %#v", window) + } + + empty, err := NormalizeAntigravityObservation("owner@example.com", antigravity.AccountObservation{PlanID: "free-tier"}) + if err != nil { + t.Fatal(err) + } + if err := json.Unmarshal(empty, &snapshot); err != nil || len(snapshot.QuotaWindows) != 0 { + t.Fatalf("missing credit snapshot = %#v, %v", snapshot, err) + } +} + +func TestAntigravityObservationCompletenessProtectsLastKnownGoodSections(t *testing.T) { + credit := &antigravity.GoogleOneAICredit{Amount: 100, MinimumAmount: 10} + tests := []struct { + name string + observation antigravity.AccountObservation + wantAccount bool + wantQuota bool + wantPartial bool + wantError bool + }{ + { + name: "complete", observation: antigravity.AccountObservation{ + PlanID: "google-one-ai", GoogleOneAICredits: credit, + }, wantAccount: true, wantQuota: true, + }, + { + name: "plan with no paid credits is complete", observation: antigravity.AccountObservation{PlanID: "free-tier"}, + wantAccount: true, wantQuota: true, + }, + { + name: "credits only is partial", observation: antigravity.AccountObservation{GoogleOneAICredits: credit}, + wantQuota: true, wantPartial: true, + }, + {name: "empty is payload failure", wantError: true}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + account, quota, partial, err := antigravityObservationCompleteness(test.observation) + if (err != nil) != test.wantError || account != test.wantAccount || quota != test.wantQuota || partial != test.wantPartial { + t.Fatalf("completeness = %t/%t/%t, %v", account, quota, partial, err) + } + if test.wantError && !errors.Is(err, ErrObservationPayloadInvalid) { + t.Fatalf("error = %v, want ErrObservationPayloadInvalid", err) + } + }) + } +} diff --git a/internal/subscription/runtime/antigravity_quota.go b/internal/subscription/runtime/antigravity_quota.go new file mode 100644 index 00000000..06702cae --- /dev/null +++ b/internal/subscription/runtime/antigravity_quota.go @@ -0,0 +1,38 @@ +package subscriptionruntime + +import ( + "encoding/json" + "fmt" + "math" + "strings" + + "gpt-load/internal/antigravity" +) + +// NormalizeAntigravityObservation creates the same provider-neutral snapshot +// shape used by other subscription channels. Google One AI credits are shown +// as a balance only; no total, utilization, reset window, or payment action is +// inferred from the upstream payload. +func NormalizeAntigravityObservation(email string, observation antigravity.AccountObservation) ([]byte, error) { + result := quotaSnapshot{ + Plan: quotaPlanSummary{Name: strings.TrimSpace(observation.PlanID)}, + Account: "aAccountSummary{Email: strings.TrimSpace(email)}, + QuotaWindows: []quotaWindow{}, + } + if credit := observation.GoogleOneAICredits; credit != nil { + if math.IsNaN(credit.Amount) || math.IsInf(credit.Amount, 0) || credit.Amount < 0 || + math.IsNaN(credit.MinimumAmount) || math.IsInf(credit.MinimumAmount, 0) || credit.MinimumAmount < 0 { + return nil, fmt.Errorf("Antigravity Google One AI credit is invalid") + } + remaining := credit.Amount + state := "available" + if credit.Amount < credit.MinimumAmount { + state = "exhausted" + } + result.QuotaWindows = append(result.QuotaWindows, quotaWindow{ + ID: "google_one_ai", Label: "Google One AI", Scope: "credits", Unit: "credits", + Remaining: &remaining, State: state, + }) + } + return json.Marshal(result) +} diff --git a/internal/subscription/runtime/runtime.go b/internal/subscription/runtime/runtime.go index f7f4605c..7b7092f7 100644 --- a/internal/subscription/runtime/runtime.go +++ b/internal/subscription/runtime/runtime.go @@ -102,6 +102,13 @@ type Driver interface { ClassifyRefreshFailure(error) RefreshFailure } +// CredentialFileImporter is an optional, narrow preprocessing capability for +// OAuth file formats that cannot safely provide a canonical credential without +// one provider-bound verification step. Existing drivers continue to use Parse. +type CredentialFileImporter interface { + ImportCredential(context.Context, []byte) (Credential, error) +} + // BrowserAuthorizationDriver is implemented only by subscription channels // which support interactive browser authorization. type BrowserAuthorizationDriver interface { @@ -393,6 +400,25 @@ func (runtime *Runtime) CanonicalCredential(channelID channel.ID, raw []byte) ([ return credential.Canonical(), nil } +// ImportCredential prepares one OAuth-file credential for a ready stage. Most +// channels use Parse directly; a provider that lacks a stable identity in its +// native file may implement CredentialFileImporter for one bounded enrichment +// step before any persistence or duplicate check occurs. +func (runtime *Runtime) ImportCredential( + ctx context.Context, + channelID channel.ID, + raw []byte, +) (Credential, error) { + driver, ok := runtime.Driver(channelID) + if !ok { + return Credential{}, fmt.Errorf("subscription driver for channel %q is unavailable", channelID) + } + if importer, ok := driver.(CredentialFileImporter); ok { + return importer.ImportCredential(ctx, raw) + } + return driver.Parse(raw) +} + func (runtime *Runtime) ModelDiscovery(channelID channel.ID) (ModelDiscovery, bool) { if runtime == nil { return nil, false diff --git a/internal/subscription/runtime/runtime_test.go b/internal/subscription/runtime/runtime_test.go index c4c8ca4d..45946cd9 100644 --- a/internal/subscription/runtime/runtime_test.go +++ b/internal/subscription/runtime/runtime_test.go @@ -11,11 +11,11 @@ import ( ) func TestRuntimeCompilesSubscriptionCapabilitiesFromChannelBindings(t *testing.T) { - runtime, err := NewRuntime(channel.NewRegistry(), CodexImplementations(), ClaudeImplementations()) + runtime, err := NewRuntime(channel.NewRegistry(), CodexImplementations(), ClaudeImplementations(), AntigravityImplementations()) if err != nil { t.Fatal(err) } - if got, want := runtime.ChannelIDs(), []channel.ID{channel.Claude, channel.Codex}; !reflect.DeepEqual(got, want) { + if got, want := runtime.ChannelIDs(), []channel.ID{channel.Antigravity, channel.Claude, channel.Codex}; !reflect.DeepEqual(got, want) { t.Fatalf("ChannelIDs() = %v, want %v", got, want) } driver, ok := runtime.Driver(channel.Codex) @@ -38,6 +38,16 @@ func TestRuntimeCompilesSubscriptionCapabilitiesFromChannelBindings(t *testing.T if _, ok := runtime.ResetCreditAction(channel.Claude); ok { t.Fatal("ResetCreditAction(claude) unexpectedly resolved") } + antigravityDriver, ok := runtime.Driver(channel.Antigravity) + if !ok || antigravityDriver.ID() != modules.AntigravitySubscriptionDriver { + t.Fatalf("Driver(antigravity) = %#v, %t", antigravityDriver, ok) + } + if discovery, ok := runtime.ModelDiscovery(channel.Antigravity); !ok || discovery.ID() != modules.AntigravityModelDiscovery { + t.Fatalf("ModelDiscovery(antigravity) = %#v, %t", discovery, ok) + } + if observation, ok := runtime.QuotaObservation(channel.Antigravity); !ok || observation.ID() != modules.AntigravityQuotaObservation { + t.Fatalf("QuotaObservation(antigravity) = %#v, %t", observation, ok) + } discovery, ok := runtime.ModelDiscovery(channel.Codex) if !ok || discovery.ID() != modules.CodexModelDiscovery { t.Fatalf("ModelDiscovery(codex) = %#v, %t", discovery, ok) @@ -130,3 +140,44 @@ func TestRuntimeRejectsDuplicateImplementationIDs(t *testing.T) { t.Fatal("compileRuntime() accepted duplicate driver IDs") } } + +type importingAntigravityDriver struct { + *antigravityDriver + imported bool +} + +func (driver *importingAntigravityDriver) ImportCredential(context.Context, []byte) (Credential, error) { + driver.imported = true + return driver.antigravityDriver.Parse([]byte(`{ + "type":"antigravity", + "access_token":"access-secret", + "refresh_token":"refresh-secret", + "account_id":"google-account-one", + "email":"owner@example.com", + "project_id":"project-one", + "expired":"2030-01-01T00:00:00Z" + }`)) +} + +func TestRuntimeImportsCredentialThroughOptionalImporter(t *testing.T) { + importer := &importingAntigravityDriver{antigravityDriver: newAntigravityDriver()} + codex := newCodexDriver() + claude := newClaudeDriver() + runtime, err := compileRuntime( + channel.NewRegistry(), + []Driver{codex, claude, importer}, + []ModelDiscovery{codex.modelDiscovery(), claude.modelDiscovery(), importer.modelDiscovery()}, + []QuotaObservation{codex.quotaObservation(), claude.quotaObservation(), importer.quotaObservation()}, + []ResetCreditAction{codex.resetCreditAction()}, + ) + if err != nil { + t.Fatal(err) + } + credential, err := runtime.ImportCredential(t.Context(), channel.Antigravity, []byte(`{"not":"canonical"}`)) + if err != nil { + t.Fatal(err) + } + if !importer.imported || credential.Identity() != "google-account-one" { + t.Fatalf("imported=%t credential=%#v", importer.imported, credential) + } +} diff --git a/internal/webui/container_contract_test.go b/internal/webui/container_contract_test.go index 43fca6aa..22101259 100644 --- a/internal/webui/container_contract_test.go +++ b/internal/webui/container_contract_test.go @@ -58,7 +58,7 @@ func TestComposeShellPortOverridesDotEnvEverywhere(t *testing.T) { if service.Environment["HOST"] != "0.0.0.0" { t.Fatalf("resolved application HOST = %q, want 0.0.0.0", service.Environment["HOST"]) } - if len(service.Ports) != 3 || + if len(service.Ports) != 4 || service.Ports[0].Target != 41234 || service.Ports[0].Published != "41234" || service.Ports[0].HostIP != "127.0.0.1" || @@ -67,8 +67,11 @@ func TestComposeShellPortOverridesDotEnvEverywhere(t *testing.T) { service.Ports[1].HostIP != "127.0.0.1" || service.Ports[2].Target != 54545 || service.Ports[2].Published != "54545" || - service.Ports[2].HostIP != "127.0.0.1" { - t.Fatalf("resolved ports = %#v, want application 41234 and loopback OAuth callbacks 1455/54545", service.Ports) + service.Ports[2].HostIP != "127.0.0.1" || + service.Ports[3].Target != 51121 || + service.Ports[3].Published != "51121" || + service.Ports[3].HostIP != "127.0.0.1" { + t.Fatalf("resolved ports = %#v, want application 41234 and loopback OAuth callbacks 1455/54545/51121", service.Ports) } if len(service.Healthcheck.Test) != 2 || !strings.Contains(service.Healthcheck.Test[1], "localhost:41234/health") { @@ -83,8 +86,8 @@ func TestDockerfileFinalStageDeclaresNonRootPersistentRuntime(t *testing.T) { t.Fatal("Dockerfile does not contain a final stage") } finalStage := content[finalStageIndex:] - if !strings.Contains(finalStage, "EXPOSE 3001 1455 54545") { - t.Fatal("Dockerfile final stage does not expose the application and both fixed OAuth callback ports") + if !strings.Contains(finalStage, "EXPOSE 3001 1455 54545 51121") { + t.Fatal("Dockerfile final stage does not expose the application and all fixed OAuth callback ports") } orderedBeforeUser := []string{ @@ -239,7 +242,7 @@ func TestComposeBindsLoopbackAndConfiguresContainerAllInterfaces(t *testing.T) { if service.Environment["HOST"] != "0.0.0.0" { t.Fatalf("resolved application HOST = %q, want 0.0.0.0", service.Environment["HOST"]) } - if len(service.Ports) != 3 || + if len(service.Ports) != 4 || service.Ports[0].Target != 3001 || service.Ports[0].Published != "3001" || service.Ports[0].HostIP != "127.0.0.1" || @@ -248,8 +251,11 @@ func TestComposeBindsLoopbackAndConfiguresContainerAllInterfaces(t *testing.T) { service.Ports[1].HostIP != "127.0.0.1" || service.Ports[2].Target != 54545 || service.Ports[2].Published != "54545" || - service.Ports[2].HostIP != "127.0.0.1" { - t.Fatalf("resolved ports = %#v, want application 3001 and loopback OAuth callbacks 1455/54545", service.Ports) + service.Ports[2].HostIP != "127.0.0.1" || + service.Ports[3].Target != 51121 || + service.Ports[3].Published != "51121" || + service.Ports[3].HostIP != "127.0.0.1" { + t.Fatalf("resolved ports = %#v, want application 3001 and loopback OAuth callbacks 1455/54545/51121", service.Ports) } } @@ -322,7 +328,7 @@ func TestComposeProjectsHaveIndependentNamesApplicationPortsAndVolumes(t *testin if service.ContainerName != "" { t.Fatalf("resolved project %s fixes container_name to %q", item.projectName, service.ContainerName) } - if len(service.Ports) != 3 || + if len(service.Ports) != 4 || service.Ports[0].Target != item.targetPort || service.Ports[0].Published != item.publishedPort || service.Ports[0].HostIP != "127.0.0.1" || @@ -331,7 +337,10 @@ func TestComposeProjectsHaveIndependentNamesApplicationPortsAndVolumes(t *testin service.Ports[1].HostIP != "127.0.0.1" || service.Ports[2].Target != 54545 || service.Ports[2].Published != "54545" || - service.Ports[2].HostIP != "127.0.0.1" { + service.Ports[2].HostIP != "127.0.0.1" || + service.Ports[3].Target != 51121 || + service.Ports[3].Published != "51121" || + service.Ports[3].HostIP != "127.0.0.1" { t.Fatalf("resolved project %s ports = %#v", item.projectName, service.Ports) } wantVolume := item.projectName + "_gpt-load-data" diff --git a/third_party/cpaembedded/README.md b/third_party/cpaembedded/README.md index b80c038e..47ec34db 100644 --- a/third_party/cpaembedded/README.md +++ b/third_party/cpaembedded/README.md @@ -1,7 +1,7 @@ # CPA Embedded Bridge This nested Go module is GPT-Load's deliberately small bridge to -CLIProxyAPI (CPA). It exists because the CPA Codex and Claude executors and +CLIProxyAPI (CPA). It exists because the CPA Codex, Claude, and Antigravity executors and OAuth helpers needed by GPT-Load are implemented in CPA `internal` packages and cannot be imported from the root `gpt-load` module directly. @@ -20,6 +20,11 @@ quota policy. This bridge only exposes: - one-shot, context-aware Claude token refresh; - Claude Code account profile, model entitlement, and usage observation; - the stateless Claude HTTP executor, upstream CountTokens request, and supported protocol translators. +- Antigravity browser OAuth, strict native-file enrichment, stable Google account identity, + model discovery, and plan/Google One AI credits observation; +- the execution-only Antigravity HTTP executor, upstream CountTokens request, and supported + protocol translators, with CPA refresh, fallback, cooldown, paid-credit fallback, and global + signature cache disabled. It intentionally excludes CPA Manager, selector, pool, file store, server, watcher, WebSocket/Auto executors, fallback, and internal retry loops. @@ -27,7 +32,7 @@ watcher, WebSocket/Auto executors, fallback, and internal retry loops. ## Pinned upstream - Module: `github.com/router-for-me/CLIProxyAPI/v7` -- Version: `v7.2.133` +- Version: `v7.2.135` The root module consumes this bridge through a local `replace`; releases still resolve CPA itself at the exact version recorded in both `go.mod` files and @@ -38,9 +43,8 @@ resolve CPA itself at the exact version recorded in both `go.mod` files and CPA upgrades are deliberate compatibility work, not automatic dependency bumps: -1. Review upstream changes to Codex and Claude OAuth, token, HTTP executor, - translation, headers, device identity, model discovery, and usage observation - code. +1. Review upstream changes to Codex, Claude, and Antigravity OAuth, token, HTTP + executor, translation, headers, identity, model discovery, and usage observation code. 2. Update the CPA version in this module and run `go mod tidy` here. 3. Fix only bridge compatibility issues; keep the execution-only boundary and do not adopt CPA Manager, retry, WebSocket, fallback, or file persistence. @@ -77,5 +81,12 @@ CPA_LIVE_CLAUDE_MODEL=optional-claude-model-id \ This live test deliberately does not complete interactive browser OAuth, rotate a refresh token, or force real 401/429 responses. Those gates require a disposable -account and an explicitly supervised run; deterministic bridge tests cover their -local classification contracts, but do not constitute real-provider evidence. + account and an explicitly supervised run; deterministic bridge tests cover their + local classification contracts, but do not constitute real-provider evidence. + +The Antigravity contract requires a disposable credential whose Google account is +authorized for the service. It verifies dynamic models, account/credits observation, +all declared unary/streaming protocol routes, and the three upstream CountTokens +routes. It must also confirm that no paid Google One AI credit type is injected. +Browser OAuth, refresh-token rotation, deliberate 401/429 responses, and provider +policy changes remain supervised live gates rather than default test behavior. diff --git a/third_party/cpaembedded/embedded/antigravity.go b/third_party/cpaembedded/embedded/antigravity.go new file mode 100644 index 00000000..052474e4 --- /dev/null +++ b/third_party/cpaembedded/embedded/antigravity.go @@ -0,0 +1,126 @@ +package embedded + +import ( + "encoding/json" + "fmt" + "strings" + "time" + + antigravityauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/antigravity" +) + +const ( + ProviderAntigravity = "antigravity" + AntigravityRedirectURI = "http://localhost:51121/oauth-callback" + antigravityLoginTimeout = 5 * time.Minute + antigravityExecutionBase = "https://daily-cloudcode-pa.googleapis.com" +) + +// AntigravityCredential is the strict canonical credential retained by +// GPT-Load. account_id comes from Google UserInfo and never from a CPA file +// name or project identifier. +type AntigravityCredential struct { + Type string `json:"type"` + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + AccountID string `json:"account_id"` + Email string `json:"email"` + ProjectID string `json:"project_id"` + ExpiresIn int64 `json:"expires_in,omitempty"` + Timestamp int64 `json:"timestamp,omitempty"` + Expire string `json:"expired"` + LastRefresh string `json:"last_refresh,omitempty"` +} + +// SecretValues returns every credential value that must be redacted before a +// provider error can cross the embedded boundary. +func (credential AntigravityCredential) SecretValues() []string { + return []string{ + credential.AccessToken, + credential.RefreshToken, + credential.AccountID, + credential.Email, + credential.ProjectID, + } +} + +// ParseAntigravityCredentialJSON accepts only GPT-Load's canonical schema. +// CPA's native JSON lacks account_id and must first pass through the explicit +// importing/enrichment flow. +func ParseAntigravityCredentialJSON(raw []byte) (AntigravityCredential, error) { + if len(raw) == 0 || len(raw) > maxCredentialBytes { + return AntigravityCredential{}, fmt.Errorf("credential JSON size is invalid") + } + allowed := map[string]struct{}{ + "type": {}, "access_token": {}, "refresh_token": {}, "account_id": {}, + "email": {}, "project_id": {}, "expires_in": {}, "timestamp": {}, + "expired": {}, "last_refresh": {}, + } + if err := validateClaudeCredentialObject(raw, allowed); err != nil { + return AntigravityCredential{}, err + } + var credential AntigravityCredential + if err := json.Unmarshal(raw, &credential); err != nil { + return AntigravityCredential{}, fmt.Errorf("decode credential: %w", err) + } + normalizeAntigravityCredential(&credential) + if err := validateAntigravityCredential(credential); err != nil { + return AntigravityCredential{}, err + } + return credential, nil +} + +// AntigravityCredentialExpiresAt returns the canonical access-token expiry. +func AntigravityCredentialExpiresAt(credential AntigravityCredential) (time.Time, bool) { + if strings.TrimSpace(credential.Expire) == "" { + return time.Time{}, false + } + value, err := time.Parse(time.RFC3339, credential.Expire) + return value, err == nil +} + +func normalizeAntigravityCredential(credential *AntigravityCredential) { + credential.Type = strings.ToLower(strings.TrimSpace(credential.Type)) + credential.AccessToken = strings.TrimSpace(credential.AccessToken) + credential.RefreshToken = strings.TrimSpace(credential.RefreshToken) + credential.AccountID = strings.TrimSpace(credential.AccountID) + credential.Email = strings.TrimSpace(credential.Email) + credential.ProjectID = strings.TrimSpace(credential.ProjectID) + credential.Expire = strings.TrimSpace(credential.Expire) + credential.LastRefresh = strings.TrimSpace(credential.LastRefresh) +} + +func validateAntigravityCredential(credential AntigravityCredential) error { + if credential.Type != ProviderAntigravity { + return fmt.Errorf("credential type must be antigravity") + } + for field, value := range map[string]string{ + "access_token": credential.AccessToken, + "refresh_token": credential.RefreshToken, + "account_id": credential.AccountID, + "email": credential.Email, + "project_id": credential.ProjectID, + } { + if value == "" { + return fmt.Errorf("credential %s is required", field) + } + if len(value) > 16*1024 || strings.ContainsAny(value, "\r\n\x00") { + return fmt.Errorf("credential %s is invalid", field) + } + } + if credential.ExpiresIn < 0 || credential.Timestamp < 0 { + return fmt.Errorf("credential token lifetime is invalid") + } + if credential.Expire == "" { + return fmt.Errorf("credential expired is required") + } + if err := validateTimestamp("expired", credential.Expire); err != nil { + return err + } + if err := validateTimestamp("last_refresh", credential.LastRefresh); err != nil { + return err + } + return nil +} + +var _ = antigravityauth.CallbackPort diff --git a/third_party/cpaembedded/embedded/antigravity_control.go b/third_party/cpaembedded/embedded/antigravity_control.go new file mode 100644 index 00000000..986dd149 --- /dev/null +++ b/third_party/cpaembedded/embedded/antigravity_control.go @@ -0,0 +1,248 @@ +package embedded + +import ( + "context" + "encoding/json" + "fmt" + "math" + "sort" + "strconv" + "strings" + "unicode" + "unicode/utf8" + + antigravityauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/antigravity" +) + +const maxAntigravityModels = 20_000 + +// AntigravityModel is one model that the authenticated account actually +// exposes. It is intentionally not backed by CPA's pinned static registry. +type AntigravityModel struct { + ID string + DisplayName string + MaxTokens int64 + MaxOutputTokens int64 +} + +// AntigravityCredit is a display-only Google One AI credit balance. Amount is +// remaining credit; MinimumAmount is the upstream activation threshold, not a +// total quota or utilization denominator. +type AntigravityCredit struct { + Amount float64 + MinimumAmount float64 +} + +// AntigravityAccountObservation is the bounded, control-plane account view +// consumed by GPT-Load's provider-neutral observation normalizer. +type AntigravityAccountObservation struct { + PlanID string + CurrentTierID string + GoogleOneAICredits *AntigravityCredit +} + +var antigravityExcludedModelIDs = map[string]struct{}{ + "chat_20706": {}, + "chat_23310": {}, + "tab_flash_lite_preview": {}, + "tab_jump_flash_lite_preview": {}, + "gemini-2.5-flash-thinking": {}, + "gemini-2.5-pro": {}, +} + +// DiscoverAntigravityModels retrieves the live, account-bound model map. +func DiscoverAntigravityModels( + ctx context.Context, + credential AntigravityCredential, + options AntigravityOptions, +) ([]AntigravityModel, error) { + if err := validateAntigravityCredential(credential); err != nil { + return nil, err + } + endpoint := strings.TrimSpace(options.FetchModelsURL) + if endpoint == "" { + endpoint = antigravityExecutionBase + "/" + antigravityauth.APIVersion + ":fetchAvailableModels" + } + payload, err := json.Marshal(map[string]string{"project": credential.ProjectID}) + if err != nil { + return nil, err + } + defer clear(payload) + body, err := fetchAntigravityJSON(ctx, endpoint, credential.AccessToken, payload, "fetchAvailableModels", options) + if err != nil { + return nil, err + } + defer clear(body) + var response struct { + Models map[string]struct { + DisplayName string `json:"displayName"` + MaxTokens int64 `json:"maxTokens"` + MaxOutputTokens int64 `json:"maxOutputTokens"` + } `json:"models"` + } + if err := json.Unmarshal(body, &response); err != nil || response.Models == nil { + return nil, fmt.Errorf("decode Antigravity models response") + } + if len(response.Models) > maxAntigravityModels { + return nil, fmt.Errorf("Antigravity models response exceeds limit") + } + models := make([]AntigravityModel, 0, len(response.Models)) + for id, details := range response.Models { + id = strings.TrimSpace(id) + if !validAntigravityModelID(id) { + continue + } + if _, excluded := antigravityExcludedModelIDs[id]; excluded { + continue + } + models = append(models, AntigravityModel{ + ID: id, + DisplayName: boundedAntigravityDisplayName(details.DisplayName), + MaxTokens: positiveAntigravityLimit(details.MaxTokens), + MaxOutputTokens: positiveAntigravityLimit(details.MaxOutputTokens), + }) + } + sort.Slice(models, func(i, j int) bool { return models[i].ID < models[j].ID }) + return models, nil +} + +// ObserveAntigravityAccount reads only plan and display-only Google One AI +// credit information. It never invents reset windows or utilization values. +func ObserveAntigravityAccount( + ctx context.Context, + credential AntigravityCredential, + options AntigravityOptions, +) (AntigravityAccountObservation, error) { + if err := validateAntigravityCredential(credential); err != nil { + return AntigravityAccountObservation{}, err + } + endpoint := strings.TrimSpace(options.LoadCodeAssistURL) + if endpoint == "" { + endpoint = antigravityauth.APIEndpoint + "/" + antigravityauth.APIVersion + ":loadCodeAssist" + } + payload, err := json.Marshal(map[string]any{"metadata": map[string]string{"ideType": "ANTIGRAVITY"}}) + if err != nil { + return AntigravityAccountObservation{}, err + } + defer clear(payload) + body, err := fetchAntigravityJSON(ctx, endpoint, credential.AccessToken, payload, "loadCodeAssist", options) + if err != nil { + return AntigravityAccountObservation{}, err + } + defer clear(body) + var response map[string]any + if err := json.Unmarshal(body, &response); err != nil || response == nil { + return AntigravityAccountObservation{}, fmt.Errorf("decode Antigravity loadCodeAssist response") + } + currentTierID := antigravityNestedID(response, "currentTier") + paidTier, paidTierPresent := response["paidTier"].(map[string]any) + planID := "" + if paidTierPresent { + planID = antigravityMapString(paidTier, "id") + } + if planID == "" { + planID = currentTierID + } + credits, err := antigravityGoogleOneAICredits(paidTier) + if err != nil { + return AntigravityAccountObservation{}, err + } + return AntigravityAccountObservation{ + PlanID: planID, + CurrentTierID: currentTierID, + GoogleOneAICredits: credits, + }, nil +} + +func validAntigravityModelID(value string) bool { + if value == "" || len(value) > 512 || !utf8.ValidString(value) || strings.ContainsAny(value, "\r\n\x00") { + return false + } + for _, character := range value { + if unicode.IsControl(character) { + return false + } + } + return true +} + +func boundedAntigravityDisplayName(value string) string { + value = strings.TrimSpace(value) + if len(value) > 512 || strings.ContainsAny(value, "\r\n\x00") { + return "" + } + return value +} + +func positiveAntigravityLimit(value int64) int64 { + if value < 0 { + return 0 + } + return value +} + +func antigravityNestedID(payload map[string]any, key string) string { + value, _ := payload[key].(map[string]any) + return antigravityMapString(value, "id") +} + +func antigravityMapString(payload map[string]any, key string) string { + if payload == nil { + return "" + } + value, _ := payload[key].(string) + return strings.TrimSpace(value) +} + +func antigravityGoogleOneAICredits(paidTier map[string]any) (*AntigravityCredit, error) { + if paidTier == nil { + return nil, nil + } + rawCredits, exists := paidTier["availableCredits"] + if !exists || rawCredits == nil { + return nil, nil + } + credits, ok := rawCredits.([]any) + if !ok { + return nil, fmt.Errorf("decode Antigravity Google One AI credits") + } + for _, raw := range credits { + credit, ok := raw.(map[string]any) + if !ok || !strings.EqualFold(antigravityMapString(credit, "creditType"), "GOOGLE_ONE_AI") { + continue + } + amount, amountOK := antigravityFloat(credit["creditAmount"]) + minimum, minimumOK := antigravityFloat(credit["minimumCreditAmountForUsage"]) + if !amountOK || !minimumOK || amount < 0 || minimum < 0 { + return nil, fmt.Errorf("decode Antigravity Google One AI credits") + } + return &AntigravityCredit{Amount: amount, MinimumAmount: minimum}, nil + } + return nil, nil +} + +func antigravityFloat(value any) (float64, bool) { + var result float64 + switch typed := value.(type) { + case string: + parsed, err := strconv.ParseFloat(strings.TrimSpace(typed), 64) + if err != nil { + return 0, false + } + result = parsed + case float64: + result = typed + case json.Number: + parsed, err := typed.Float64() + if err != nil { + return 0, false + } + result = parsed + default: + return 0, false + } + if math.IsNaN(result) || math.IsInf(result, 0) { + return 0, false + } + return result, true +} diff --git a/third_party/cpaembedded/embedded/antigravity_executor.go b/third_party/cpaembedded/embedded/antigravity_executor.go new file mode 100644 index 00000000..943041f8 --- /dev/null +++ b/third_party/cpaembedded/embedded/antigravity_executor.go @@ -0,0 +1,490 @@ +package embedded + +import ( + "context" + "crypto/tls" + "encoding/json" + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "sync" + "time" + "unicode/utf8" + + internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + internalexecutor "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +// AntigravityExecutionError is the bounded error shape allowed to cross the +// embedded boundary. The upstream response body is parsed locally then dropped. +type AntigravityExecutionError struct { + status int + typeValue string + codeValue string + summary string + retryAfter time.Duration + requestScoped bool +} + +func (err *AntigravityExecutionError) Error() string { + if err == nil || err.summary == "" { + return "Antigravity upstream request failed" + } + return err.summary +} + +func (err *AntigravityExecutionError) StatusCode() int { + if err == nil { + return 0 + } + return err.status +} + +func (err *AntigravityExecutionError) ErrorType() string { + if err == nil { + return "" + } + return err.typeValue +} + +func (err *AntigravityExecutionError) ErrorCode() string { + if err == nil { + return "" + } + return err.codeValue +} + +func (err *AntigravityExecutionError) RetryAfter() *time.Duration { + if err == nil || err.retryAfter <= 0 { + return nil + } + value := err.retryAfter + return &value +} + +func (err *AntigravityExecutionError) IsRequestScoped() bool { + return err != nil && err.requestScoped +} + +// AntigravityHTTPExecutor is the execution-only CPA surface consumed by +// GPT-Load. It deliberately excludes CPA's manager, auth refresh persistence, +// selector, cooldown ownership, and credits fallback. +type AntigravityHTTPExecutor interface { + ExecuteCanonical(context.Context, string, AntigravityCredential, ExecuteRequest) (ExecuteResponse, error) + CountTokensCanonical(context.Context, string, AntigravityCredential, ExecuteRequest) (ExecuteResponse, error) + ExecuteStreamCanonical(context.Context, string, AntigravityCredential, ExecuteRequest) (*ExecuteStreamResponse, error) +} + +type antigravityHTTPExecutor struct { + cfg *internalconfig.Config + inner *internalexecutor.AntigravityExecutor + baseURL string +} + +// NewAntigravityHTTPExecutor creates the production execution-only facade. +func NewAntigravityHTTPExecutor() AntigravityHTTPExecutor { + return newAntigravityHTTPExecutor(antigravityExecutionBase) +} + +func newAntigravityHTTPExecutor(baseURL string) *antigravityHTTPExecutor { + baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/") + if baseURL == "" { + baseURL = antigravityExecutionBase + } + // RequestRetry=0 and a fixed auth base_url guarantee that CPA never retries + // within a GPT-Load attempt or falls back between daily and production hosts. + // QuotaExceeded remains zero-valued, so no credits flag can trigger paid use. + // CPA's global signature cache is only used by its Antigravity translator and + // is not scoped by GPT-Load tenant or credential. Keep it disabled; the + // isolated reasoning replay lane below remains available through ContinuityKey. + internalcache.SetSignatureCacheEnabled(false) + cfg := &internalconfig.Config{RequestRetry: 0} + return &antigravityHTTPExecutor{cfg: cfg, inner: internalexecutor.NewAntigravityExecutor(cfg), baseURL: baseURL} +} + +// NewAntigravityAuth creates a transient CPA auth object. credentialID is +// intentionally not used as Auth.ID: CPA's private 429 cooldown map is keyed +// by that field, while GPT-Load owns all credential health and cooldown state. +func NewAntigravityAuth(_ string, credential AntigravityCredential, baseURL string) *cliproxyauth.Auth { + if strings.TrimSpace(baseURL) == "" { + baseURL = antigravityExecutionBase + } + // CPA refreshes at a 50-minute skew. The runtime has already prepared this + // token and owns durable refresh, so use an ephemeral far-future expiry and + // never pass a refresh token into this execution-only auth. + now := time.Now().UTC() + return &cliproxyauth.Auth{ + Provider: ProviderAntigravity, + Attributes: map[string]string{ + "base_url": strings.TrimRight(strings.TrimSpace(baseURL), "/"), + }, + Metadata: map[string]any{ + "type": ProviderAntigravity, + "access_token": credential.AccessToken, + "email": credential.Email, + "project_id": credential.ProjectID, + "timestamp": now.UnixMilli(), + "expires_in": int64(24 * 60 * 60), + "expired": now.Add(24 * time.Hour).Format(time.RFC3339), + }, + } +} + +func (executor *antigravityHTTPExecutor) ExecuteCanonical( + ctx context.Context, + credentialID string, + credential AntigravityCredential, + request ExecuteRequest, +) (ExecuteResponse, error) { + if ctx == nil { + ctx = context.Background() + } + if err := validateAntigravityCredential(credential); err != nil { + return ExecuteResponse{}, err + } + request = prepareAntigravityExecutionRequest(request) + format := sdktranslator.FromString(request.Format) + auth := NewAntigravityAuth(credentialID, credential, executor.baseURL) + observation := newProviderExecutionObservation(request, ProviderAntigravity) + response, err := executor.inner.Execute(executor.executionContext(ctx, observation), auth, cliproxyexecutor.Request{ + Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: format, + }, antigravityExecutorOptions(request, format, false)) + if err != nil { + return ExecuteResponse{AppliedReasoningEffort: observation.reasoningEffort()}, normalizeAntigravityExecutionError(err) + } + return ExecuteResponse{ + Payload: normalizeAntigravityConvertedUsage(request.Format, false, response.Payload), Headers: response.Headers.Clone(), + AppliedReasoningEffort: observation.reasoningEffort(), + }, nil +} + +func (executor *antigravityHTTPExecutor) CountTokensCanonical( + ctx context.Context, + credentialID string, + credential AntigravityCredential, + request ExecuteRequest, +) (ExecuteResponse, error) { + if ctx == nil { + ctx = context.Background() + } + if err := validateAntigravityCredential(credential); err != nil { + return ExecuteResponse{}, err + } + request = prepareAntigravityExecutionRequest(request) + format := sdktranslator.FromString(request.Format) + auth := NewAntigravityAuth(credentialID, credential, executor.baseURL) + response, err := executor.inner.CountTokens(executor.executionContext(ctx, nil), auth, cliproxyexecutor.Request{ + Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: format, + }, antigravityExecutorOptions(request, format, false)) + if err != nil { + return ExecuteResponse{}, normalizeAntigravityExecutionError(err) + } + return ExecuteResponse{ + Payload: normalizeAntigravityCountTokens(request.Format, response.Payload), Headers: response.Headers.Clone(), + }, nil +} + +func (executor *antigravityHTTPExecutor) ExecuteStreamCanonical( + ctx context.Context, + credentialID string, + credential AntigravityCredential, + request ExecuteRequest, +) (*ExecuteStreamResponse, error) { + if ctx == nil { + ctx = context.Background() + } + if err := validateAntigravityCredential(credential); err != nil { + return nil, err + } + request = prepareAntigravityExecutionRequest(request) + format := sdktranslator.FromString(request.Format) + auth := NewAntigravityAuth(credentialID, credential, executor.baseURL) + observation := newProviderExecutionObservation(request, ProviderAntigravity) + response, err := executor.inner.ExecuteStream(executor.executionContext(ctx, observation), auth, cliproxyexecutor.Request{ + Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: format, + }, antigravityExecutorOptions(request, format, true)) + if err != nil { + return &ExecuteStreamResponse{AppliedReasoningEffort: observation.reasoningEffort()}, normalizeAntigravityExecutionError(err) + } + chunks := make(chan ExecuteStreamChunk) + go func() { + defer close(chunks) + for chunk := range response.Chunks { + converted := ExecuteStreamChunk{Payload: normalizeAntigravityConvertedUsage(request.Format, true, chunk.Payload)} + if chunk.Err != nil { + converted.Err = normalizeAntigravityExecutionError(chunk.Err) + } + select { + case chunks <- converted: + case <-ctx.Done(): + return + } + } + }() + return &ExecuteStreamResponse{ + Headers: response.Headers.Clone(), Chunks: chunks, AppliedReasoningEffort: observation.reasoningEffort(), + }, nil +} + +func antigravityExecutorOptions( + request ExecuteRequest, + format sdktranslator.Format, + stream bool, +) cliproxyexecutor.Options { + options := cliproxyexecutor.Options{ + Stream: stream, Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), + SourceFormat: format, ResponseFormat: format, + } + if scope := strings.TrimSpace(request.ContinuityKey); scope != "" { + options.Metadata = map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: scope} + } + return options +} + +func prepareAntigravityExecutionRequest(request ExecuteRequest) ExecuteRequest { + if strings.TrimSpace(request.ContinuityKey) == "" && strings.TrimSpace(request.AttemptID) != "" { + request.ContinuityKey = "gpt-load-antigravity-attempt\x00" + strings.TrimSpace(request.AttemptID) + } + request.Payload = stripAntigravityReplayInputs(request.Payload) + request.OriginalRequest = stripAntigravityReplayInputs(request.OriginalRequest) + request.Headers = request.Headers.Clone() + stripAntigravityReplayHeaders(request.Headers) + return request +} + +func stripAntigravityReplayInputs(raw []byte) []byte { + if len(raw) == 0 { + return nil + } + result := append([]byte(nil), raw...) + for _, path := range []string{ + "session_id", "sessionId", "prompt_cache_key", + "metadata.session_id", "metadata.sessionId", "metadata.user_id", + "request.sessionId", "request.session_id", + } { + updated, err := sjson.DeleteBytes(result, path) + if err != nil { + continue + } + result = updated + } + return result +} + +func stripAntigravityReplayHeaders(headers http.Header) { + for name := range headers { + switch { + case strings.EqualFold(name, "Session-Id"), + strings.EqualFold(name, "X-Claude-Code-Session-Id"), + strings.EqualFold(name, "X-Claude-Code-Agent-Id"): + delete(headers, name) + } + } +} + +func normalizeAntigravityConvertedUsage(format string, stream bool, raw []byte) []byte { + result := append([]byte(nil), raw...) + format = strings.TrimSpace(format) + switch format { + case "openai": + return addAntigravityUsageIntegers(result, "usage.completion_tokens", "usage.completion_tokens_details.reasoning_tokens") + case "openai-response": + return addAntigravityUsageIntegers(result, "usage.output_tokens", "usage.output_tokens_details.reasoning_tokens") + case "claude": + if stream { + return result + } + return subtractAntigravityUsageIntegers(result, "usage.input_tokens", "usage.cache_read_input_tokens") + default: + return result + } +} + +func addAntigravityUsageIntegers(raw []byte, totalPath, extraPath string) []byte { + total, totalOK := antigravityUsageInteger(raw, totalPath) + extra, extraOK := antigravityUsageInteger(raw, extraPath) + if !totalOK || !extraOK || extra > 0 && total > int64(^uint64(0)>>1)-extra { + return raw + } + updated, err := sjson.SetBytes(raw, totalPath, total+extra) + if err != nil { + return raw + } + clear(raw) + return updated +} + +func subtractAntigravityUsageIntegers(raw []byte, totalPath, cachedPath string) []byte { + total, totalOK := antigravityUsageInteger(raw, totalPath) + cached, cachedOK := antigravityUsageInteger(raw, cachedPath) + if !totalOK || !cachedOK { + return raw + } + uncached := total - cached + if uncached < 0 { + uncached = 0 + } + updated, err := sjson.SetBytes(raw, totalPath, uncached) + if err != nil { + return raw + } + clear(raw) + return updated +} + +func antigravityUsageInteger(raw []byte, path string) (int64, bool) { + value := gjson.GetBytes(raw, path) + if !value.Exists() || value.Type != gjson.Number { + return 0, false + } + parsed, err := strconv.ParseInt(value.Raw, 10, 64) + if err != nil || parsed < 0 { + return 0, false + } + return parsed, true +} + +func normalizeAntigravityCountTokens(format string, raw []byte) []byte { + if strings.TrimSpace(format) != "openai-response" { + return append([]byte(nil), raw...) + } + count, ok := antigravityUsageInteger(raw, "totalTokens") + if !ok { + return append([]byte(nil), raw...) + } + encoded, err := json.Marshal(struct { + Object string `json:"object"` + InputTokens int64 `json:"input_tokens"` + }{Object: "response.input_tokens", InputTokens: count}) + if err != nil { + return append([]byte(nil), raw...) + } + return encoded +} + +var ( + antigravityExecutionTransportOnce sync.Once + antigravityExecutionTransport *http.Transport +) + +type antigravityExecutionContext struct{ context.Context } + +func (ctx antigravityExecutionContext) Value(key any) any { + if key == "gin" { + return nil + } + return ctx.Context.Value(key) +} + +func (executor *antigravityHTTPExecutor) executionContext(ctx context.Context, observation *executionObservation) context.Context { + ctx = antigravityExecutionContext{Context: ctx} + antigravityExecutionTransportOnce.Do(func() { + base, ok := http.DefaultTransport.(*http.Transport) + if !ok || base == nil { + base = &http.Transport{} + } + transport := base.Clone() + transport.ForceAttemptHTTP2 = false + transport.TLSNextProto = map[string]func(string, *tls.Conn) http.RoundTripper{} + if transport.TLSClientConfig == nil { + transport.TLSClientConfig = &tls.Config{} + } else { + transport.TLSClientConfig = transport.TLSClientConfig.Clone() + } + transport.TLSClientConfig.NextProtos = []string{"http/1.1"} + antigravityExecutionTransport = transport + }) + return context.WithValue(ctx, "cliproxy.roundtripper", noRedirectRoundTripper{ + base: antigravityExecutionTransport, observation: observation, + }) +} + +func normalizeAntigravityExecutionError(err error) error { + if err == nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return err + } + status := 0 + var statusError interface{ StatusCode() int } + if errors.As(err, &statusError) && statusError != nil { + status = statusError.StatusCode() + } + if status == 0 { + return err + } + retryAfter := time.Duration(0) + var retry interface{ RetryAfter() *time.Duration } + if errors.As(err, &retry) && retry != nil { + if value := retry.RetryAfter(); value != nil && *value > 0 { + retryAfter = *value + } + } + typeValue, codeValue := antigravityExecutionErrorFields([]byte(err.Error())) + return &AntigravityExecutionError{ + status: status, typeValue: typeValue, codeValue: codeValue, retryAfter: retryAfter, + requestScoped: status >= http.StatusBadRequest && status < http.StatusInternalServerError && + status != http.StatusUnauthorized && status != http.StatusForbidden && status != http.StatusTooManyRequests, + summary: antigravityExecutionSummary(status, typeValue, codeValue), + } +} + +func antigravityExecutionErrorFields(raw []byte) (string, string) { + defer clear(raw) + var payload struct { + Error struct { + Status string `json:"status"` + Message string `json:"message"` + Details []struct { + Reason string `json:"reason"` + } `json:"details"` + } `json:"error"` + } + if json.Unmarshal(raw, &payload) != nil { + return "", "" + } + code := "" + for _, detail := range payload.Error.Details { + if code = boundedAntigravityErrorScalar(detail.Reason); code != "" { + break + } + } + return boundedAntigravityErrorScalar(payload.Error.Status), code +} + +func boundedAntigravityErrorScalar(value string) string { + value = strings.TrimSpace(value) + if value == "" || len(value) > 128 || !utf8.ValidString(value) || strings.ContainsAny(value, "\r\n\x00") { + return "" + } + return value +} + +func antigravityExecutionSummary(status int, typeValue, codeValue string) string { + switch { + case status == http.StatusUnauthorized: + return "Antigravity authorization was rejected" + case status == http.StatusForbidden: + return "Antigravity access was denied" + case status == http.StatusTooManyRequests && strings.EqualFold(codeValue, "INSUFFICIENT_G1_CREDITS_BALANCE"): + return "Antigravity Google One AI credits are unavailable" + case status == http.StatusTooManyRequests: + return "Antigravity upstream rate limit was reached" + case status >= http.StatusInternalServerError: + return "Antigravity upstream service failed" + case typeValue != "": + return fmt.Sprintf("Antigravity upstream request failed: %s", typeValue) + default: + return fmt.Sprintf("Antigravity upstream request failed with status %d", status) + } +} + +var _ AntigravityHTTPExecutor = (*antigravityHTTPExecutor)(nil) +var _ cliproxyexecutor.RequestScopedError = (*AntigravityExecutionError)(nil) +var _ cliproxyexecutor.StatusError = (*AntigravityExecutionError)(nil) diff --git a/third_party/cpaembedded/embedded/antigravity_oauth.go b/third_party/cpaembedded/embedded/antigravity_oauth.go new file mode 100644 index 00000000..df2fa7ce --- /dev/null +++ b/third_party/cpaembedded/embedded/antigravity_oauth.go @@ -0,0 +1,669 @@ +package embedded + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/url" + "strings" + "time" + + antigravityauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/antigravity" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" +) + +var ( + ErrAntigravityInvalidState = errors.New("antigravity oauth state mismatch") + ErrAntigravityCredentialIdentityChanged = errors.New("refreshed antigravity credential identity changed") +) + +// AntigravityTokenEndpointError retains only the bounded error code needed by +// the subscription lifecycle. Provider response bodies never leave this package. +type AntigravityTokenEndpointError struct { + StatusCode int + Code string +} + +func (err *AntigravityTokenEndpointError) Error() string { + if err == nil { + return "Antigravity token endpoint failed" + } + return fmt.Sprintf("Antigravity token endpoint returned status %d", err.StatusCode) +} + +// AntigravityUpstreamHTTPError retains safe HTTP evidence for control-plane +// requests such as UserInfo, model discovery, and account observation. +type AntigravityUpstreamHTTPError struct { + Operation string + StatusCode int +} + +func (err *AntigravityUpstreamHTTPError) Error() string { + if err == nil { + return "Antigravity upstream request failed" + } + return fmt.Sprintf("Antigravity %s endpoint returned status %d", err.Operation, err.StatusCode) +} + +// AntigravityOptions supplies testable HTTP endpoints and time boundaries. It +// never changes the fixed OAuth client identity or callback URI. +type AntigravityOptions struct { + TokenURL string + UserInfoURL string + LoadCodeAssistURL string + OnboardUserURL string + FetchModelsURL string + HTTPClient *http.Client + Now func() time.Time +} + +type antigravityTokenResponse struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + ExpiresIn int64 `json:"expires_in"` +} + +type antigravityUserInfo struct { + ID string `json:"id"` + Email string `json:"email"` + VerifiedEmail *bool `json:"verified_email"` +} + +type antigravityImportedCredential struct { + Type string `json:"type"` + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + AccountID string `json:"account_id,omitempty"` + Email string `json:"email"` + ProjectID string `json:"project_id,omitempty"` + ExpiresIn int64 `json:"expires_in,omitempty"` + Timestamp int64 `json:"timestamp,omitempty"` + Expire string `json:"expired,omitempty"` + LastRefresh string `json:"last_refresh,omitempty"` + Disabled bool `json:"disabled,omitempty"` +} + +// BeginAntigravityBrowserAuthorization creates a browser challenge without +// opening a browser, binding a listener, or writing a credential file. +func BeginAntigravityBrowserAuthorization() (BrowserAuthorization, error) { + state, err := misc.GenerateRandomState() + if err != nil { + return BrowserAuthorization{}, fmt.Errorf("generate Antigravity OAuth state: %w", err) + } + authorizationURL := antigravityauth.NewAntigravityAuth(&internalconfig.Config{}, nil). + BuildAuthURL(state, AntigravityRedirectURI) + return BrowserAuthorization{ + AuthorizationURL: authorizationURL, + State: state, + ExpiresAt: time.Now().UTC().Add(antigravityLoginTimeout), + }, nil +} + +// CompleteAntigravityBrowserAuthorization verifies the stage-bound state, +// exchanges one authorization code, and enriches the result with Google user +// identity plus the execution project required by Antigravity. +func CompleteAntigravityBrowserAuthorization( + ctx context.Context, + completion BrowserAuthorizationCompletion, + options AntigravityOptions, +) (AntigravityCredential, error) { + if ctx == nil { + ctx = context.Background() + } + if !constantTimeEqual(completion.ExpectedState, completion.ReturnedState) { + return AntigravityCredential{}, ErrAntigravityInvalidState + } + if strings.TrimSpace(completion.Code) == "" { + return AntigravityCredential{}, fmt.Errorf("authorization code is required") + } + token, err := requestAntigravityToken(ctx, url.Values{ + "grant_type": {"authorization_code"}, + "code": {strings.TrimSpace(completion.Code)}, + "redirect_uri": {AntigravityRedirectURI}, + }, options) + if err != nil { + return AntigravityCredential{}, fmt.Errorf("exchange Antigravity authorization code: %w", err) + } + identity, err := fetchAntigravityUserInfo(ctx, token.AccessToken, options) + if err != nil { + return AntigravityCredential{}, fmt.Errorf("fetch Antigravity user identity: %w", err) + } + projectID, _, err := discoverAntigravityProject(ctx, token.AccessToken, options) + if err != nil { + return AntigravityCredential{}, fmt.Errorf("discover Antigravity project: %w", err) + } + now := antigravityNow(options) + credential := AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: token.AccessToken, RefreshToken: token.RefreshToken, + AccountID: identity.ID, Email: identity.Email, ProjectID: projectID, + ExpiresIn: token.ExpiresIn, Timestamp: now.UnixMilli(), LastRefresh: now.Format(time.RFC3339), + } + if token.ExpiresIn > 0 { + credential.Expire = now.Add(time.Duration(token.ExpiresIn) * time.Second).Format(time.RFC3339) + } + normalizeAntigravityCredential(&credential) + if err := validateAntigravityCredential(credential); err != nil { + return AntigravityCredential{}, fmt.Errorf("exchange Antigravity authorization code: %w", err) + } + return credential, nil +} + +// RefreshAntigravityCredentialOnce refreshes exactly once under the caller's +// context, then verifies that Google still identifies the same account. It +// deliberately bypasses CPA's executor refresh path so refreshed tokens always +// return to GPT-Load for CAS persistence. +func RefreshAntigravityCredentialOnce( + ctx context.Context, + current AntigravityCredential, + options AntigravityOptions, +) (AntigravityCredential, error) { + if ctx == nil { + ctx = context.Background() + } + if err := validateAntigravityCredential(current); err != nil { + return AntigravityCredential{}, err + } + token, err := requestAntigravityToken(ctx, url.Values{ + "grant_type": {"refresh_token"}, + "refresh_token": {current.RefreshToken}, + }, options) + if err != nil { + return AntigravityCredential{}, fmt.Errorf("refresh Antigravity credential: %w", err) + } + identity, err := fetchAntigravityUserInfo(ctx, token.AccessToken, options) + if err != nil { + return AntigravityCredential{}, fmt.Errorf("verify refreshed Antigravity identity: %w", err) + } + if identity.ID != current.AccountID { + return AntigravityCredential{}, ErrAntigravityCredentialIdentityChanged + } + refreshed := current + refreshed.AccessToken = token.AccessToken + if token.RefreshToken != "" { + refreshed.RefreshToken = token.RefreshToken + } + refreshed.Email = identity.Email + refreshed.ExpiresIn = token.ExpiresIn + now := antigravityNow(options) + refreshed.Timestamp = now.UnixMilli() + refreshed.Expire = now.Add(time.Duration(token.ExpiresIn) * time.Second).Format(time.RFC3339) + refreshed.LastRefresh = now.Format(time.RFC3339) + normalizeAntigravityCredential(&refreshed) + if err := validateAntigravityCredential(refreshed); err != nil { + return AntigravityCredential{}, fmt.Errorf("refresh Antigravity credential: %w", err) + } + return refreshed, nil +} + +// ImportAntigravityCredential enriches CPA's native OAuth JSON or revalidates +// a GPT-Load canonical download before it enters the encrypted-stage lifecycle. +// disabled is intentionally discarded because it is source-instance state. +func ImportAntigravityCredential( + ctx context.Context, + raw []byte, + options AntigravityOptions, +) (AntigravityCredential, error) { + if ctx == nil { + ctx = context.Background() + } + imported, err := parseAntigravityImportedCredential(raw) + if err != nil { + return AntigravityCredential{}, err + } + now := antigravityNow(options) + refreshed := false + if expiresAt, known := imported.expiresAt(); !known || !expiresAt.After(now.Add(5*time.Minute)) { + token, refreshErr := requestAntigravityToken(ctx, url.Values{ + "grant_type": {"refresh_token"}, + "refresh_token": {imported.RefreshToken}, + }, options) + if refreshErr != nil { + return AntigravityCredential{}, fmt.Errorf("refresh imported Antigravity credential: %w", refreshErr) + } + imported.applyToken(token, now) + refreshed = true + } + identity, err := fetchAntigravityUserInfo(ctx, imported.AccessToken, options) + if _, ok := antigravityUnauthorized(err); ok && !refreshed { + token, refreshErr := requestAntigravityToken(ctx, url.Values{ + "grant_type": {"refresh_token"}, + "refresh_token": {imported.RefreshToken}, + }, options) + if refreshErr != nil { + return AntigravityCredential{}, fmt.Errorf("refresh imported Antigravity credential: %w", refreshErr) + } + imported.applyToken(token, now) + identity, err = fetchAntigravityUserInfo(ctx, imported.AccessToken, options) + } + if err != nil { + return AntigravityCredential{}, fmt.Errorf("verify imported Antigravity identity: %w", err) + } + if !strings.EqualFold(imported.Email, identity.Email) { + return AntigravityCredential{}, fmt.Errorf("imported Antigravity email does not match userinfo") + } + if imported.AccountID != "" && imported.AccountID != identity.ID { + return AntigravityCredential{}, fmt.Errorf("imported Antigravity account does not match userinfo") + } + projectID, _, err := discoverAntigravityProject(ctx, imported.AccessToken, options) + if err != nil { + return AntigravityCredential{}, fmt.Errorf("discover imported Antigravity project: %w", err) + } + credential := AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: imported.AccessToken, RefreshToken: imported.RefreshToken, + AccountID: identity.ID, Email: identity.Email, ProjectID: projectID, + ExpiresIn: imported.ExpiresIn, Timestamp: imported.Timestamp, Expire: imported.Expire, + LastRefresh: imported.LastRefresh, + } + if credential.LastRefresh == "" { + credential.LastRefresh = now.Format(time.RFC3339) + } + normalizeAntigravityCredential(&credential) + if err := validateAntigravityCredential(credential); err != nil { + return AntigravityCredential{}, fmt.Errorf("import Antigravity credential: %w", err) + } + return credential, nil +} + +// MarshalAntigravityCredential returns a revalidated canonical JSON document +// suitable for GPT-Load encryption. +func MarshalAntigravityCredential(credential AntigravityCredential) ([]byte, error) { + raw, err := json.Marshal(credential) + if err != nil { + return nil, err + } + parsed, err := ParseAntigravityCredentialJSON(raw) + clear(raw) + if err != nil { + return nil, err + } + return json.Marshal(parsed) +} + +func IsDefinitiveAntigravityRefreshRejection(code string) bool { + switch strings.ToLower(strings.TrimSpace(code)) { + case "invalid_grant", "invalid_token", "invalid_request": + return true + default: + return false + } +} + +func antigravityNow(options AntigravityOptions) time.Time { + if options.Now != nil { + return options.Now().UTC() + } + return time.Now().UTC() +} + +func antigravityOAuthClient(options AntigravityOptions) *http.Client { + return clientWithoutRedirects(options.HTTPClient) +} + +func requestAntigravityToken( + ctx context.Context, + values url.Values, + options AntigravityOptions, +) (antigravityTokenResponse, error) { + values = cloneURLValues(values) + values.Set("client_id", antigravityauth.ClientID) + values.Set("client_secret", antigravityauth.ClientSecret) + endpoint := strings.TrimSpace(options.TokenURL) + if endpoint == "" { + endpoint = antigravityauth.TokenEndpoint + } + request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(values.Encode())) + if err != nil { + return antigravityTokenResponse{}, err + } + request.Header.Set("Content-Type", "application/x-www-form-urlencoded") + response, err := antigravityOAuthClient(options).Do(request) + if err != nil { + return antigravityTokenResponse{}, err + } + defer func() { _ = response.Body.Close() }() + body, err := readClaudeOAuthBody(response) + if err != nil { + return antigravityTokenResponse{}, err + } + defer clear(body) + if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { + return antigravityTokenResponse{}, &AntigravityTokenEndpointError{ + StatusCode: response.StatusCode, + Code: antigravityTokenEndpointErrorCode(body), + } + } + var token antigravityTokenResponse + if err := json.Unmarshal(body, &token); err != nil { + return antigravityTokenResponse{}, fmt.Errorf("decode Antigravity token response: %w", err) + } + token.AccessToken = strings.TrimSpace(token.AccessToken) + token.RefreshToken = strings.TrimSpace(token.RefreshToken) + if token.AccessToken == "" || token.ExpiresIn <= 0 { + return antigravityTokenResponse{}, fmt.Errorf("Antigravity token response is incomplete") + } + if values.Get("grant_type") == "authorization_code" && token.RefreshToken == "" { + return antigravityTokenResponse{}, fmt.Errorf("Antigravity token response is incomplete") + } + return token, nil +} + +func fetchAntigravityUserInfo( + ctx context.Context, + accessToken string, + options AntigravityOptions, +) (antigravityUserInfo, error) { + endpoint := strings.TrimSpace(options.UserInfoURL) + if endpoint == "" { + endpoint = antigravityauth.UserInfoEndpoint + } + body, err := fetchAntigravityJSON(ctx, endpoint, accessToken, nil, "userinfo", options) + if err != nil { + return antigravityUserInfo{}, err + } + defer clear(body) + var identity antigravityUserInfo + if err := json.Unmarshal(body, &identity); err != nil { + return antigravityUserInfo{}, fmt.Errorf("decode Antigravity userinfo response: %w", err) + } + identity.ID = strings.TrimSpace(identity.ID) + identity.Email = strings.TrimSpace(identity.Email) + if identity.ID == "" || identity.Email == "" || identity.VerifiedEmail == nil || !*identity.VerifiedEmail { + return antigravityUserInfo{}, fmt.Errorf("Antigravity userinfo response is incomplete") + } + return identity, nil +} + +func discoverAntigravityProject( + ctx context.Context, + accessToken string, + options AntigravityOptions, +) (string, map[string]any, error) { + endpoint := strings.TrimSpace(options.LoadCodeAssistURL) + if endpoint == "" { + endpoint = antigravityauth.APIEndpoint + "/" + antigravityauth.APIVersion + ":loadCodeAssist" + } + body, err := json.Marshal(map[string]any{"metadata": map[string]string{"ideType": "ANTIGRAVITY"}}) + if err != nil { + return "", nil, err + } + defer clear(body) + response, err := fetchAntigravityJSON(ctx, endpoint, accessToken, body, "loadCodeAssist", options) + if err != nil { + return "", nil, err + } + defer clear(response) + var payload map[string]any + if err := json.Unmarshal(response, &payload); err != nil || payload == nil { + return "", nil, fmt.Errorf("decode Antigravity loadCodeAssist response") + } + if projectID := antigravityProjectID(payload); projectID != "" { + return projectID, payload, nil + } + projectID, err := onboardAntigravityUser(ctx, accessToken, antigravityDefaultTierID(payload), options) + if err != nil { + return "", nil, err + } + if projectID == "" { + return "", nil, fmt.Errorf("Antigravity project ID is missing") + } + return projectID, payload, nil +} + +func fetchAntigravityJSON( + ctx context.Context, + endpoint string, + accessToken string, + payload []byte, + operation string, + options AntigravityOptions, +) ([]byte, error) { + var reader *bytes.Reader + method := http.MethodGet + if payload != nil { + method = http.MethodPost + reader = bytes.NewReader(payload) + } else { + reader = bytes.NewReader(nil) + } + request, err := http.NewRequestWithContext(ctx, method, endpoint, reader) + if err != nil { + return nil, err + } + request.Header.Set("Authorization", "Bearer "+strings.TrimSpace(accessToken)) + request.Header.Set("Accept", "application/json") + request.Header.Set("User-Agent", misc.AntigravityRequestUserAgent("")) + switch operation { + case "loadCodeAssist": + request.Header.Set("Accept", "*/*") + case "onboardUser": + request.Header.Set("Accept", "*/*") + request.Header.Set("User-Agent", misc.AntigravityOnboardUserUserAgent("")) + request.Header.Set("X-Goog-Api-Client", misc.AntigravityGoogAPIClientUA) + } + if payload != nil { + request.Header.Set("Content-Type", "application/json") + } + response, err := antigravityOAuthClient(options).Do(request) + if err != nil { + return nil, err + } + defer func() { _ = response.Body.Close() }() + body, err := readClaudeOAuthBody(response) + if err != nil { + return nil, err + } + if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { + clear(body) + return nil, &AntigravityUpstreamHTTPError{Operation: operation, StatusCode: response.StatusCode} + } + return body, nil +} + +func onboardAntigravityUser( + ctx context.Context, + accessToken string, + tierID string, + options AntigravityOptions, +) (string, error) { + endpoint := strings.TrimSpace(options.OnboardUserURL) + if endpoint == "" { + endpoint = antigravityauth.DailyAPIEndpoint + "/" + antigravityauth.APIVersion + ":onboardUser" + } + userAgent := misc.AntigravityOnboardUserUserAgent("") + payload, err := json.Marshal(map[string]any{ + "tier_id": tierID, + "metadata": map[string]string{ + "ide_type": "ANTIGRAVITY", "ide_name": "antigravity", + "ide_version": misc.AntigravityVersionFromUserAgent(userAgent), + }, + }) + if err != nil { + return "", err + } + defer clear(payload) + for attempt := 0; attempt < 5; attempt++ { + body, err := fetchAntigravityJSON(ctx, endpoint, accessToken, payload, "onboardUser", options) + if err != nil { + return "", err + } + var result struct { + Done bool `json:"done"` + Response map[string]any `json:"response"` + } + err = json.Unmarshal(body, &result) + clear(body) + if err != nil { + return "", fmt.Errorf("decode Antigravity onboardUser response") + } + if result.Done { + return antigravityProjectID(result.Response), nil + } + if attempt == 4 { + break + } + timer := time.NewTimer(2 * time.Second) + select { + case <-ctx.Done(): + timer.Stop() + return "", ctx.Err() + case <-timer.C: + } + } + return "", fmt.Errorf("Antigravity onboarding did not complete") +} + +func antigravityProjectID(payload map[string]any) string { + for _, key := range []string{"cloudaicompanionProject", "projectId", "project"} { + switch value := payload[key].(type) { + case string: + if value = strings.TrimSpace(value); value != "" { + return value + } + case map[string]any: + if id, ok := value["id"].(string); ok { + if id = strings.TrimSpace(id); id != "" { + return id + } + } + } + } + return "" +} + +func antigravityDefaultTierID(payload map[string]any) string { + if tiers, ok := payload["allowedTiers"].([]any); ok { + for _, raw := range tiers { + tier, ok := raw.(map[string]any) + if !ok || tier["isDefault"] != true { + continue + } + if id, ok := tier["id"].(string); ok && strings.TrimSpace(id) != "" { + return strings.TrimSpace(id) + } + } + } + if tier, ok := payload["currentTier"].(map[string]any); ok { + if id, ok := tier["id"].(string); ok && strings.TrimSpace(id) != "" { + return strings.TrimSpace(id) + } + } + return "free-tier" +} + +func antigravityTokenEndpointErrorCode(body []byte) string { + var payload struct { + Error string `json:"error"` + ErrorDescription string `json:"error_description"` + } + if json.Unmarshal(body, &payload) != nil { + return "" + } + value := strings.ToLower(strings.TrimSpace(payload.Error)) + if value == "" || len(value) > 64 || strings.ContainsAny(value, "\r\n\x00") { + return "" + } + return value +} + +func cloneURLValues(values url.Values) url.Values { + result := make(url.Values, len(values)) + for key, value := range values { + result[key] = append([]string(nil), value...) + } + return result +} + +func parseAntigravityImportedCredential(raw []byte) (antigravityImportedCredential, error) { + if len(raw) == 0 || len(raw) > maxCredentialBytes { + return antigravityImportedCredential{}, fmt.Errorf("credential JSON size is invalid") + } + allowed := map[string]struct{}{ + "type": {}, "access_token": {}, "refresh_token": {}, "account_id": {}, "email": {}, "project_id": {}, + "expires_in": {}, "timestamp": {}, "expired": {}, "last_refresh": {}, "disabled": {}, + } + if err := validateClaudeCredentialObject(raw, allowed); err != nil { + return antigravityImportedCredential{}, err + } + if err := validateAntigravityImportedDisabled(raw); err != nil { + return antigravityImportedCredential{}, err + } + var credential antigravityImportedCredential + if err := json.Unmarshal(raw, &credential); err != nil { + return antigravityImportedCredential{}, fmt.Errorf("decode imported Antigravity credential: %w", err) + } + credential.Type = strings.ToLower(strings.TrimSpace(credential.Type)) + credential.AccessToken = strings.TrimSpace(credential.AccessToken) + credential.RefreshToken = strings.TrimSpace(credential.RefreshToken) + credential.AccountID = strings.TrimSpace(credential.AccountID) + credential.Email = strings.TrimSpace(credential.Email) + credential.ProjectID = strings.TrimSpace(credential.ProjectID) + credential.Expire = strings.TrimSpace(credential.Expire) + credential.LastRefresh = strings.TrimSpace(credential.LastRefresh) + if credential.Type != ProviderAntigravity || credential.AccessToken == "" || credential.RefreshToken == "" || credential.Email == "" { + return antigravityImportedCredential{}, fmt.Errorf("imported Antigravity credential is incomplete") + } + for _, value := range []string{credential.AccessToken, credential.RefreshToken, credential.AccountID, credential.Email, credential.ProjectID} { + if len(value) > 16*1024 || strings.ContainsAny(value, "\r\n\x00") { + return antigravityImportedCredential{}, fmt.Errorf("imported Antigravity credential is invalid") + } + } + if credential.ExpiresIn < 0 || credential.Timestamp < 0 || validateTimestamp("expired", credential.Expire) != nil || + validateTimestamp("last_refresh", credential.LastRefresh) != nil { + return antigravityImportedCredential{}, fmt.Errorf("imported Antigravity credential token lifetime is invalid") + } + if credential.Expire == "" && (credential.ExpiresIn <= 0 || credential.Timestamp <= 0) { + return antigravityImportedCredential{}, fmt.Errorf("imported Antigravity credential expiry is missing") + } + if credential.Expire == "" { + credential.Expire = time.UnixMilli(credential.Timestamp). + Add(time.Duration(credential.ExpiresIn) * time.Second).UTC().Format(time.RFC3339) + } + return credential, nil +} + +func validateAntigravityImportedDisabled(raw []byte) error { + var fields map[string]json.RawMessage + if err := json.Unmarshal(raw, &fields); err != nil { + return fmt.Errorf("decode imported Antigravity credential: %w", err) + } + value, present := fields["disabled"] + if !present { + return nil + } + if strings.EqualFold(strings.TrimSpace(string(value)), "null") { + return fmt.Errorf("imported Antigravity credential disabled is invalid") + } + var disabled bool + if err := json.Unmarshal(value, &disabled); err != nil { + return fmt.Errorf("imported Antigravity credential disabled is invalid") + } + return nil +} + +func (credential antigravityImportedCredential) expiresAt() (time.Time, bool) { + value, err := time.Parse(time.RFC3339, credential.Expire) + return value, err == nil +} + +func (credential *antigravityImportedCredential) applyToken(token antigravityTokenResponse, now time.Time) { + credential.AccessToken = token.AccessToken + if token.RefreshToken != "" { + credential.RefreshToken = token.RefreshToken + } + credential.ExpiresIn = token.ExpiresIn + credential.Timestamp = now.UnixMilli() + credential.Expire = now.Add(time.Duration(token.ExpiresIn) * time.Second).Format(time.RFC3339) + credential.LastRefresh = now.Format(time.RFC3339) +} + +func antigravityUnauthorized(err error) (*AntigravityUpstreamHTTPError, bool) { + var upstream *AntigravityUpstreamHTTPError + if errors.As(err, &upstream) && upstream != nil && upstream.StatusCode == http.StatusUnauthorized { + return upstream, true + } + return nil, false +} diff --git a/third_party/cpaembedded/embedded/antigravity_test.go b/third_party/cpaembedded/embedded/antigravity_test.go new file mode 100644 index 00000000..c76836a8 --- /dev/null +++ b/third_party/cpaembedded/embedded/antigravity_test.go @@ -0,0 +1,860 @@ +package embedded + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" + "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type antigravityStatusTestError struct { + status int + body string +} + +func (err antigravityStatusTestError) Error() string { return err.body } +func (err antigravityStatusTestError) StatusCode() int { return err.status } + +func TestParseAntigravityCredentialJSONRequiresCanonicalStableIdentity(t *testing.T) { + credential, err := ParseAntigravityCredentialJSON([]byte(`{ + "type":"antigravity", + "access_token":"access-secret", + "refresh_token":"refresh-secret", + "account_id":"google-account-one", + "email":"owner@example.com", + "project_id":"project-one", + "expires_in":3600, + "timestamp":1770000000000, + "expired":"2026-08-17T10:00:00Z", + "last_refresh":"2026-08-17T09:00:00Z" + }`)) + if err != nil { + t.Fatal(err) + } + if credential.AccountID != "google-account-one" || credential.ProjectID != "project-one" || + credential.Email != "owner@example.com" { + t.Fatalf("credential = %#v", credential) + } + + for _, raw := range [][]byte{ + []byte(`{"type":"antigravity","access_token":"a","refresh_token":"r","email":"owner@example.com","project_id":"project-one","expired":"2026-08-17T10:00:00Z"}`), + []byte(`{"type":"antigravity","access_token":"a","refresh_token":"r","account_id":"account-one","email":"owner@example.com","project_id":"project-one","expired":"2026-08-17T10:00:00Z","proxy":"http://example.invalid"}`), + []byte(`{"type":"antigravity","access_token":"a","access_token":"b","refresh_token":"r","account_id":"account-one","email":"owner@example.com","project_id":"project-one","expired":"2026-08-17T10:00:00Z"}`), + } { + if _, err := ParseAntigravityCredentialJSON(raw); err == nil { + t.Fatalf("ParseAntigravityCredentialJSON(%s) error = nil", raw) + } + } +} + +func TestAntigravityCredentialSecretValuesAreRedactable(t *testing.T) { + credential, err := ParseAntigravityCredentialJSON([]byte(`{ + "type":"antigravity", + "access_token":"access-secret", + "refresh_token":"refresh-secret", + "account_id":"google-account-one", + "email":"owner@example.com", + "project_id":"project-one", + "expired":"2026-08-17T10:00:00Z" + }`)) + if err != nil { + t.Fatal(err) + } + for _, secret := range credential.SecretValues() { + if secret == "" || !strings.Contains("access-secret refresh-secret google-account-one owner@example.com project-one", secret) { + t.Fatalf("unexpected secret value %q", secret) + } + } +} + +func TestCompleteAntigravityBrowserAuthorizationEnrichesIdentityAndProject(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/token": + if err := request.ParseForm(); err != nil || request.Form.Get("grant_type") != "authorization_code" || + request.Form.Get("code") != "authorization-code" || request.Form.Get("redirect_uri") != AntigravityRedirectURI { + t.Fatalf("token form = %#v, err=%v", request.Form, err) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"access_token":"access-secret","refresh_token":"refresh-secret","expires_in":3600}`)) + case "/userinfo": + if got := request.Header.Get("Authorization"); got != "Bearer access-secret" { + t.Fatalf("userinfo authorization = %q", got) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"id":"google-account-one","email":"owner@example.com","verified_email":true}`)) + case "/load": + if got := request.Header.Get("Authorization"); got != "Bearer access-secret" { + t.Fatalf("load authorization = %q", got) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"cloudaicompanionProject":"project-one"}`)) + default: + http.NotFound(writer, request) + } + })) + defer server.Close() + + credential, err := CompleteAntigravityBrowserAuthorization(context.Background(), BrowserAuthorizationCompletion{ + ExpectedState: "expected-state", ReturnedState: "expected-state", Code: "authorization-code", + }, AntigravityOptions{ + TokenURL: server.URL + "/token", UserInfoURL: server.URL + "/userinfo", + LoadCodeAssistURL: server.URL + "/load", HTTPClient: server.Client(), + }) + if err != nil { + t.Fatal(err) + } + if credential.AccountID != "google-account-one" || credential.Email != "owner@example.com" || + credential.ProjectID != "project-one" || credential.AccessToken != "access-secret" || + credential.RefreshToken != "refresh-secret" { + t.Fatalf("credential = %#v", credential) + } +} + +func TestDiscoverAntigravityProjectUsesCPAControlPlaneHeaders(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/load": + if request.Header.Get("Accept") != "*/*" || request.Header.Get("User-Agent") != misc.AntigravityRequestUserAgent("") || + request.Header.Get("X-Goog-Api-Client") != "" { + t.Fatalf("loadCodeAssist headers = %#v", request.Header) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"allowedTiers":[{"id":"free-tier","isDefault":true}]}`)) + case "/onboard": + if request.Header.Get("Accept") != "*/*" || request.Header.Get("User-Agent") != misc.AntigravityOnboardUserUserAgent("") || + request.Header.Get("X-Goog-Api-Client") != misc.AntigravityGoogAPIClientUA { + t.Fatalf("onboardUser headers = %#v", request.Header) + } + var body map[string]any + err := json.NewDecoder(request.Body).Decode(&body) + metadata, metadataOK := body["metadata"].(map[string]any) + if err != nil || body["tier_id"] != "free-tier" || + !metadataOK || metadata["ide_version"] != misc.AntigravityVersionFromUserAgent(misc.AntigravityOnboardUserUserAgent("")) { + t.Fatalf("onboardUser body = %#v, error = %v", body, err) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"done":true,"response":{"projectId":"project-one"}}`)) + default: + http.NotFound(writer, request) + } + })) + defer server.Close() + + projectID, _, err := discoverAntigravityProject(t.Context(), "access-secret", AntigravityOptions{ + LoadCodeAssistURL: server.URL + "/load", OnboardUserURL: server.URL + "/onboard", HTTPClient: server.Client(), + }) + if err != nil || projectID != "project-one" { + t.Fatalf("discoverAntigravityProject() = %q, %v", projectID, err) + } +} + +func TestBeginAntigravityBrowserAuthorizationUsesFixedCallback(t *testing.T) { + t.Parallel() + + authorization, err := BeginAntigravityBrowserAuthorization() + if err != nil { + t.Fatal(err) + } + parsed, err := url.Parse(authorization.AuthorizationURL) + if err != nil { + t.Fatal(err) + } + if authorization.State == "" || parsed.Query().Get("redirect_uri") != AntigravityRedirectURI || + parsed.Query().Get("state") != authorization.State { + t.Fatalf("authorization = %#v", authorization) + } +} + +func TestRefreshAntigravityCredentialOnceKeepsIdentityAndExistingProject(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/token": + if err := request.ParseForm(); err != nil || request.Form.Get("grant_type") != "refresh_token" || + request.Form.Get("refresh_token") != "refresh-secret" { + t.Errorf("refresh form = %#v, err=%v", request.Form, err) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"access_token":"new-access-secret","expires_in":3600}`)) + case "/userinfo": + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"id":"google-account-one","email":"renamed@example.com","verified_email":true}`)) + default: + http.NotFound(writer, request) + } + })) + defer server.Close() + + refreshed, err := RefreshAntigravityCredentialOnce(context.Background(), AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: "old-access-secret", RefreshToken: "refresh-secret", + AccountID: "google-account-one", Email: "owner@example.com", ProjectID: "project-one", + Expire: "2026-08-17T10:00:00Z", + }, AntigravityOptions{ + TokenURL: server.URL + "/token", UserInfoURL: server.URL + "/userinfo", HTTPClient: server.Client(), + }) + if err != nil { + t.Fatal(err) + } + if refreshed.AccessToken != "new-access-secret" || refreshed.RefreshToken != "refresh-secret" || + refreshed.AccountID != "google-account-one" || refreshed.Email != "renamed@example.com" || + refreshed.ProjectID != "project-one" { + t.Fatalf("refreshed = %#v", refreshed) + } +} + +func TestImportAntigravityCredentialEnrichesNativeCPAFile(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/userinfo": + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"id":"google-account-one","email":"owner@example.com","verified_email":true}`)) + case "/load": + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"projectId":"verified-project"}`)) + default: + http.NotFound(writer, request) + } + })) + defer server.Close() + + credential, err := ImportAntigravityCredential(context.Background(), []byte(`{ + "type":"antigravity", + "access_token":"access-secret", + "refresh_token":"refresh-secret", + "email":"owner@example.com", + "project_id":"unverified-project", + "expires_in":3600, + "timestamp":1770000000000, + "expired":"2030-01-01T00:00:00Z", + "disabled":false + }`), AntigravityOptions{ + UserInfoURL: server.URL + "/userinfo", LoadCodeAssistURL: server.URL + "/load", HTTPClient: server.Client(), + }) + if err != nil { + t.Fatal(err) + } + if credential.AccountID != "google-account-one" || credential.ProjectID != "verified-project" || + credential.Email != "owner@example.com" { + t.Fatalf("credential = %#v", credential) + } +} + +func TestImportAntigravityCredentialAcceptsDownloadedCanonicalFile(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/userinfo": + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"id":"google-account-one","email":"owner@example.com","verified_email":true}`)) + case "/load": + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"projectId":"verified-project"}`)) + default: + http.NotFound(writer, request) + } + })) + defer server.Close() + + raw := `{ + "type":"antigravity", + "access_token":"access-secret", + "refresh_token":"refresh-secret", + "account_id":"google-account-one", + "email":"owner@example.com", + "project_id":"downloaded-project", + "expires_in":3600, + "timestamp":1770000000000, + "expired":"2030-01-01T00:00:00Z", + "last_refresh":"2029-12-31T23:00:00Z" + }` + options := AntigravityOptions{ + UserInfoURL: server.URL + "/userinfo", LoadCodeAssistURL: server.URL + "/load", HTTPClient: server.Client(), + } + credential, err := ImportAntigravityCredential(t.Context(), []byte(raw), options) + if err != nil { + t.Fatal(err) + } + if credential.AccountID != "google-account-one" || credential.ProjectID != "verified-project" { + t.Fatalf("credential = %#v", credential) + } + if _, err := ImportAntigravityCredential( + t.Context(), + []byte(strings.Replace(raw, `"account_id":"google-account-one"`, `"account_id":"different-account"`, 1)), + options, + ); err == nil { + t.Fatal("ImportAntigravityCredential() accepted a mismatched canonical account ID") + } +} + +func TestImportAntigravityCredentialRefreshesAtMostOnce(t *testing.T) { + t.Parallel() + + var tokenCalls int + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/token": + tokenCalls++ + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"access_token":"new-access-secret","expires_in":3600}`)) + case "/userinfo": + writer.WriteHeader(http.StatusUnauthorized) + default: + http.NotFound(writer, request) + } + })) + defer server.Close() + + _, err := ImportAntigravityCredential(t.Context(), []byte(`{ + "type":"antigravity", + "access_token":"expired-access-secret", + "refresh_token":"refresh-secret", + "email":"owner@example.com", + "expired":"2020-01-01T00:00:00Z" + }`), AntigravityOptions{ + TokenURL: server.URL + "/token", UserInfoURL: server.URL + "/userinfo", HTTPClient: server.Client(), + }) + if err == nil { + t.Fatal("ImportAntigravityCredential() error = nil, want userinfo rejection") + } + if tokenCalls != 1 { + t.Fatalf("token calls = %d, want one refresh attempt", tokenCalls) + } +} + +func TestParseAntigravityImportedCredentialRequiresBooleanDisabled(t *testing.T) { + base := `{"type":"antigravity","access_token":"access-secret","refresh_token":"refresh-secret","email":"owner@example.com","expires_in":3600,"timestamp":1770000000000}` + for _, raw := range []string{ + strings.TrimSuffix(base, "}") + `,"disabled":null}`, + strings.TrimSuffix(base, "}") + `,"disabled":"false"}`, + } { + if _, err := parseAntigravityImportedCredential([]byte(raw)); err == nil { + t.Fatalf("parseAntigravityImportedCredential(%s) error = nil", raw) + } + } +} + +func TestDiscoverAntigravityModelsUsesOnlyLiveAccountModels(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != "/models" { + http.NotFound(writer, request) + return + } + if got := request.Header.Get("Authorization"); got != "Bearer access-secret" { + t.Errorf("models authorization = %q", got) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{ + "models":{ + "gemini-live":{"displayName":"Gemini Live"}, + "chat_20706":{"displayName":"Internal"}, + "bad\nmodel":{"displayName":"Invalid"}, + "bad\u0001model":{"displayName":"Invalid"} + } + }`)) + })) + defer server.Close() + + models, err := DiscoverAntigravityModels(context.Background(), AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: "access-secret", RefreshToken: "refresh-secret", + AccountID: "google-account-one", Email: "owner@example.com", ProjectID: "project-one", + Expire: "2030-01-01T00:00:00Z", + }, AntigravityOptions{FetchModelsURL: server.URL + "/models", HTTPClient: server.Client()}) + if err != nil { + t.Fatal(err) + } + if len(models) != 1 || models[0].ID != "gemini-live" || models[0].DisplayName != "Gemini Live" { + t.Fatalf("models = %#v", models) + } +} + +func TestDiscoverAntigravityModelsRejectsExcessiveModelMap(t *testing.T) { + t.Parallel() + const oversizedModelCount = 20_001 + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", "application/json") + models := make(map[string]any, oversizedModelCount) + for index := 0; index < oversizedModelCount; index++ { + models[fmt.Sprintf("gemini-test-%d", index)] = map[string]any{} + } + _ = json.NewEncoder(writer).Encode(map[string]any{"models": models}) + })) + defer server.Close() + + _, err := DiscoverAntigravityModels(t.Context(), AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: "access-secret", RefreshToken: "refresh-secret", + AccountID: "google-account-one", Email: "owner@example.com", ProjectID: "project-one", + Expire: "2030-01-01T00:00:00Z", + }, AntigravityOptions{FetchModelsURL: server.URL, HTTPClient: server.Client()}) + if err == nil { + t.Fatal("DiscoverAntigravityModels() error = nil, want oversized live model map rejection") + } +} + +func TestObserveAntigravityAccountDoesNotInventQuota(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != "/load" { + http.NotFound(writer, request) + return + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{ + "currentTier":{"id":"free-tier"}, + "paidTier":{"id":"google-one-ai","availableCredits":[ + {"creditType":"GOOGLE_ONE_AI","creditAmount":"25000","minimumCreditAmountForUsage":"50"} + ]} + }`)) + })) + defer server.Close() + + observation, err := ObserveAntigravityAccount(context.Background(), AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: "access-secret", RefreshToken: "refresh-secret", + AccountID: "google-account-one", Email: "owner@example.com", ProjectID: "project-one", + Expire: "2030-01-01T00:00:00Z", + }, AntigravityOptions{LoadCodeAssistURL: server.URL + "/load", HTTPClient: server.Client()}) + if err != nil { + t.Fatal(err) + } + if observation.PlanID != "google-one-ai" || observation.GoogleOneAICredits == nil || + observation.GoogleOneAICredits.Amount != 25000 || observation.GoogleOneAICredits.MinimumAmount != 50 { + t.Fatalf("observation = %#v", observation) + } +} + +func TestObserveAntigravityAccountRejectsMalformedGoogleOneAICredits(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{ + "currentTier":{"id":"g1-pro-tier"}, + "paidTier":{"id":"g1-pro-tier","availableCredits":[ + {"creditType":"GOOGLE_ONE_AI","minimumCreditAmountForUsage":"50"} + ]} + }`)) + })) + defer server.Close() + + _, err := ObserveAntigravityAccount(t.Context(), AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: "access-secret", RefreshToken: "refresh-secret", + AccountID: "google-account-one", Email: "owner@example.com", ProjectID: "project-one", + Expire: "2030-01-01T00:00:00Z", + }, AntigravityOptions{LoadCodeAssistURL: server.URL, HTTPClient: server.Client()}) + if err == nil { + t.Fatal("ObserveAntigravityAccount() error = nil, want malformed credits rejection") + } +} + +func TestAntigravityExecutionOnlyBridgeUsesOneFixedUpstreamDispatch(t *testing.T) { + var calls int + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls++ + body, err := io.ReadAll(request.Body) + if err != nil { + t.Errorf("read request body: %v", err) + } + if strings.Contains(string(body), `"enabledCreditTypes"`) { + t.Errorf("execution-only request must not opt into Google One AI credits: %s", body) + } + if request.Header.Get("Authorization") != "Bearer access-secret" { + t.Errorf("authorization = %q", request.Header.Get("Authorization")) + } + switch request.URL.Path { + case "/v1internal:generateContent": + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"response":{"candidates":[{"content":{"parts":[{"text":"ok"}]}}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":1,"totalTokenCount":3}}}`)) + case "/v1internal:countTokens": + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"totalTokens":3}`)) + default: + http.NotFound(writer, request) + } + })) + defer server.Close() + + executor := newAntigravityHTTPExecutor(server.URL) + credential := AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: "access-secret", RefreshToken: "refresh-secret", + AccountID: "google-account-one", Email: "owner@example.com", ProjectID: "project-one", + Expire: "2030-01-01T00:00:00Z", + } + request := ExecuteRequest{ + Model: "gemini-live", Format: "gemini", + Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`), + OriginalRequest: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`), + } + response, err := executor.ExecuteCanonical(context.Background(), "credential-one", credential, request) + if err != nil { + t.Fatal(err) + } + if calls != 1 || !strings.Contains(string(response.Payload), `"text":"ok"`) { + t.Fatalf("calls=%d response=%s", calls, response.Payload) + } + count, err := executor.CountTokensCanonical(context.Background(), "credential-one", credential, request) + if err != nil { + t.Fatal(err) + } + if calls != 2 || !strings.Contains(string(count.Payload), "3") { + t.Fatalf("calls=%d count=%s", calls, count.Payload) + } +} + +func TestAntigravityExecutionOnlyBridgeDoesNotRetryOrKeepPrivateCooldown(t *testing.T) { + var calls int + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls++ + writer.Header().Set("Content-Type", "application/json") + writer.WriteHeader(http.StatusTooManyRequests) + _, _ = writer.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","details":[{"reason":"RATE_LIMIT_EXCEEDED"},{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"0.45s"}]}}`)) + })) + defer server.Close() + + executor := newAntigravityHTTPExecutor(server.URL) + credential := AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: "access-secret", RefreshToken: "refresh-secret", + AccountID: "google-account-one", Email: "owner@example.com", ProjectID: "project-one", + Expire: "2030-01-01T00:00:00Z", + } + request := ExecuteRequest{ + Model: "gemini-live", Format: "gemini", + Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`), + OriginalRequest: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`), + } + for attempt := 0; attempt < 2; attempt++ { + _, err := executor.ExecuteCanonical(context.Background(), "credential-one", credential, request) + var upstream *AntigravityExecutionError + if !errors.As(err, &upstream) || upstream.StatusCode() != http.StatusTooManyRequests || + upstream.ErrorType() != "RESOURCE_EXHAUSTED" || upstream.ErrorCode() != "RATE_LIMIT_EXCEEDED" { + t.Fatalf("attempt %d error = %v", attempt, err) + } + if retryAfter := upstream.RetryAfter(); retryAfter == nil || *retryAfter != 450*time.Millisecond { + t.Fatalf("attempt %d retry after = %v", attempt, retryAfter) + } + } + if calls != 2 { + t.Fatalf("upstream calls = %d, want one dispatch for each GPT-Load attempt", calls) + } +} + +func TestAntigravityForbiddenIsNotMarkedRequestScoped(t *testing.T) { + err := normalizeAntigravityExecutionError(antigravityStatusTestError{ + status: http.StatusForbidden, + body: `{"error":{"status":"PERMISSION_DENIED"}}`, + }) + var scoped interface{ IsRequestScoped() bool } + if !errors.As(err, &scoped) || scoped.IsRequestScoped() { + t.Fatalf("normalized error = %#v, want non-request-scoped 403", err) + } +} + +func TestAntigravityExecutionOnlyBridgeNeverRefreshesPreparedNearExpiryToken(t *testing.T) { + var calls int + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls++ + if request.URL.Path != "/v1internal:generateContent" { + http.NotFound(writer, request) + return + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"response":{"candidates":[{"content":{"parts":[{"text":"ok"}]}}]}}`)) + })) + defer server.Close() + + response, err := newAntigravityHTTPExecutor(server.URL).ExecuteCanonical(t.Context(), "credential-one", AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: "access-secret", RefreshToken: "refresh-secret", + AccountID: "google-account-one", Email: "owner@example.com", ProjectID: "project-one", + Expire: time.Now().UTC().Add(time.Second).Format(time.RFC3339), + }, ExecuteRequest{ + Model: "gemini-live", Format: "gemini", + Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`), + OriginalRequest: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`), + }) + if err != nil || !strings.Contains(string(response.Payload), `"candidates"`) || calls != 1 { + t.Fatalf("response=%s calls=%d error=%v", response.Payload, calls, err) + } +} + +func TestAntigravityExecutionOnlyBridgeConvertsDeclaredUnaryProtocols(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != "/v1internal:generateContent" { + http.NotFound(writer, request) + return + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"cachedContentTokenCount":2,"candidatesTokenCount":3,"thoughtsTokenCount":4,"totalTokenCount":12}}`)) + })) + defer server.Close() + + executor := newAntigravityHTTPExecutor(server.URL) + credential := AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: "access-secret", RefreshToken: "refresh-secret", + AccountID: "google-account-one", Email: "owner@example.com", ProjectID: "project-one", + Expire: "2030-01-01T00:00:00Z", + } + tests := []struct { + name string + format string + body string + want string + wantUsage string + }{ + { + name: "Gemini", format: "gemini", want: `"candidates"`, wantUsage: `"thoughtsTokenCount":4`, + body: `{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`, + }, + { + name: "Anthropic", format: "claude", want: `"content"`, wantUsage: `"input_tokens":3,"output_tokens":7,"cache_read_input_tokens":2`, + body: `{"model":"gemini-live","max_tokens":64,"messages":[{"role":"user","content":"hello"}]}`, + }, + { + name: "OpenAI Chat", format: "openai", want: `"choices"`, wantUsage: `"completion_tokens":7`, + body: `{"model":"gemini-live","messages":[{"role":"user","content":"hello"}]}`, + }, + { + name: "OpenAI Responses", format: "openai-response", want: `"output"`, wantUsage: `"output_tokens":7`, + body: `{"model":"gemini-live","input":"hello"}`, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + response, err := executor.ExecuteCanonical(t.Context(), "credential-one", credential, ExecuteRequest{ + Model: "gemini-live", Format: test.format, Payload: []byte(test.body), OriginalRequest: []byte(test.body), + }) + if err != nil || !json.Valid(response.Payload) || !strings.Contains(string(response.Payload), test.want) || + !strings.Contains(string(response.Payload), test.wantUsage) { + t.Fatalf("response=%s error=%v", response.Payload, err) + } + }) + } +} + +func TestAntigravityExecutionOnlyBridgeConvertsDeclaredStreamingProtocols(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != "/v1internal:streamGenerateContent" { + http.NotFound(writer, request) + return + } + writer.Header().Set("Content-Type", "text/event-stream") + _, _ = writer.Write([]byte(`{"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"cachedContentTokenCount":2,"candidatesTokenCount":3,"thoughtsTokenCount":4,"totalTokenCount":12}}` + "\n")) + })) + defer server.Close() + + executor := newAntigravityHTTPExecutor(server.URL) + credential := AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: "access-secret", RefreshToken: "refresh-secret", + AccountID: "google-account-one", Email: "owner@example.com", ProjectID: "project-one", + Expire: "2030-01-01T00:00:00Z", + } + tests := []struct { + name string + format string + body string + want string + }{ + { + name: "Gemini", format: "gemini", want: `"candidates"`, + body: `{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`, + }, + { + name: "Anthropic", format: "claude", want: `"content_block"`, + body: `{"model":"gemini-live","max_tokens":64,"messages":[{"role":"user","content":"hello"}]}`, + }, + { + name: "OpenAI Chat", format: "openai", want: `"choices"`, + body: `{"model":"gemini-live","messages":[{"role":"user","content":"hello"}]}`, + }, + { + name: "OpenAI Responses", format: "openai-response", want: "response.", + body: `{"model":"gemini-live","input":"hello"}`, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + response, err := executor.ExecuteStreamCanonical(t.Context(), "credential-one", credential, ExecuteRequest{ + Model: "gemini-live", Format: test.format, Payload: []byte(test.body), OriginalRequest: []byte(test.body), + }) + if err != nil || response == nil { + t.Fatalf("ExecuteStreamCanonical() = %#v, %v", response, err) + } + var wire strings.Builder + for chunk := range response.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + wire.Write(chunk.Payload) + } + if !strings.Contains(wire.String(), test.want) { + t.Fatalf("stream wire = %q, want %q", wire.String(), test.want) + } + }) + } +} + +func TestAntigravityExecutionOnlyBridgeCountsTokensUpstreamForDeclaredProtocols(t *testing.T) { + var calls int + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls++ + if request.URL.Path != "/v1internal:countTokens" { + http.NotFound(writer, request) + return + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"totalTokens":42}`)) + })) + defer server.Close() + + executor := newAntigravityHTTPExecutor(server.URL) + credential := AntigravityCredential{ + Type: ProviderAntigravity, AccessToken: "access-secret", RefreshToken: "refresh-secret", + AccountID: "google-account-one", Email: "owner@example.com", ProjectID: "project-one", + Expire: "2030-01-01T00:00:00Z", + } + tests := []struct { + name string + format string + body string + want string + }{ + { + name: "Gemini", format: "gemini", want: `"totalTokens":42`, + body: `{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`, + }, + { + name: "Anthropic", format: "claude", want: `"input_tokens":42`, + body: `{"model":"gemini-live","messages":[{"role":"user","content":"hello"}]}`, + }, + { + name: "OpenAI Responses", format: "openai-response", want: `"object":"response.input_tokens","input_tokens":42`, + body: `{"model":"gemini-live","input":"hello"}`, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + response, err := executor.CountTokensCanonical(t.Context(), "credential-one", credential, ExecuteRequest{ + Model: "gemini-live", Format: test.format, Payload: []byte(test.body), OriginalRequest: []byte(test.body), + }) + if err != nil || !strings.Contains(string(response.Payload), test.want) { + t.Fatalf("response=%s error=%v", response.Payload, err) + } + }) + } + if calls != len(tests) { + t.Fatalf("upstream countTokens calls = %d, want %d", calls, len(tests)) + } +} + +func TestNewAntigravityHTTPExecutorDisablesGlobalSignatureCache(t *testing.T) { + previous := internalcache.SignatureCacheEnabled() + t.Cleanup(func() { internalcache.SetSignatureCacheEnabled(previous) }) + internalcache.SetSignatureCacheEnabled(true) + + _ = NewAntigravityHTTPExecutor() + if internalcache.SignatureCacheEnabled() { + t.Fatal("Antigravity execution must not use CPA's global signature cache") + } +} + +func TestAntigravityExecutionRequestUsesOnlyPrivateContinuityScope(t *testing.T) { + prepared := prepareAntigravityExecutionRequest(ExecuteRequest{ + Payload: []byte(`{"contents":[{"parts":[{"text":"keep"}]}],"session_id":"client-session","prompt_cache_key":"client-cache","metadata":{"session_id":"nested"},"request":{"sessionId":"nested-request","session_id":"nested-request-snake"}}`), + OriginalRequest: []byte(`{"contents":[{"parts":[{"text":"keep"}]}],"session_id":"client-session","prompt_cache_key":"client-cache","metadata":{"session_id":"nested"},"request":{"sessionId":"nested-request","session_id":"nested-request-snake"}}`), + Headers: http.Header{ + "Session-Id": {"client-session"}, + "X-Claude-Code-Session-Id": {"client-claude-session"}, + "X-Claude-Code-Agent-Id": {"client-claude-agent"}, + }, + ContinuityKey: "opaque-tenant-scope", + }) + if strings.Contains(string(prepared.Payload), "client-session") || strings.Contains(string(prepared.OriginalRequest), "client-cache") || + strings.Contains(string(prepared.Payload), "nested-request") || strings.Contains(string(prepared.OriginalRequest), "nested-request-snake") || + prepared.Headers.Get("Session-Id") != "" || prepared.Headers.Get("X-Claude-Code-Session-Id") != "" || + prepared.Headers.Get("X-Claude-Code-Agent-Id") != "" { + t.Fatalf("prepared request retained untrusted replay inputs: %#v", prepared) + } + if !json.Valid(prepared.Payload) || !json.Valid(prepared.OriginalRequest) || + !strings.Contains(string(prepared.Payload), `"text":"keep"`) || + !strings.Contains(string(prepared.OriginalRequest), `"text":"keep"`) { + t.Fatalf("prepared request was corrupted: payload=%q original=%q", prepared.Payload, prepared.OriginalRequest) + } + contextWithGin := context.WithValue(t.Context(), "gin", "downstream-request-context") + if got := newAntigravityHTTPExecutor("").executionContext(contextWithGin, nil).Value("gin"); got != nil { + t.Fatalf("execution context retained downstream gin value: %#v", got) + } + options := antigravityExecutorOptions(prepared, "gemini", false) + if options.Metadata[cliproxyexecutor.ExecutionSessionMetadataKey] != "opaque-tenant-scope" { + t.Fatalf("options metadata = %#v", options.Metadata) + } + fallback := prepareAntigravityExecutionRequest(ExecuteRequest{AttemptID: "attempt-one"}) + if fallback.ContinuityKey == "" || antigravityExecutorOptions(fallback, "gemini", false).Metadata[cliproxyexecutor.ExecutionSessionMetadataKey] == "" { + t.Fatalf("fallback continuity = %#v", fallback) + } +} + +func TestNormalizeAntigravityConvertedUsagePreservesReasoningAndCacheSemantics(t *testing.T) { + tests := []struct { + name string + format string + stream bool + body string + want string + }{ + { + name: "OpenAI chat adds reasoning tokens", format: "openai", + body: `{"usage":{"prompt_tokens":10,"completion_tokens":4,"completion_tokens_details":{"reasoning_tokens":6}}}`, + want: `{"usage":{"prompt_tokens":10,"completion_tokens":10,"completion_tokens_details":{"reasoning_tokens":6}}}`, + }, + { + name: "Responses adds reasoning tokens", format: "openai-response", + body: `{"usage":{"input_tokens":10,"output_tokens":4,"output_tokens_details":{"reasoning_tokens":6}}}`, + want: `{"usage":{"input_tokens":10,"output_tokens":10,"output_tokens_details":{"reasoning_tokens":6}}}`, + }, + { + name: "Anthropic unary subtracts cached input", format: "claude", + body: `{"usage":{"input_tokens":10,"cache_read_input_tokens":4,"output_tokens":3}}`, + want: `{"usage":{"input_tokens":6,"cache_read_input_tokens":4,"output_tokens":3}}`, + }, + { + name: "Anthropic stream already has uncached input", format: "claude", stream: true, + body: `{"usage":{"input_tokens":6,"cache_read_input_tokens":4,"output_tokens":3}}`, + want: `{"usage":{"input_tokens":6,"cache_read_input_tokens":4,"output_tokens":3}}`, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if got := string(normalizeAntigravityConvertedUsage(test.format, test.stream, []byte(test.body))); got != test.want { + t.Fatalf("usage = %s, want %s", got, test.want) + } + }) + } +} + +func TestNormalizeAntigravityResponsesCountTokensContract(t *testing.T) { + if got := string(normalizeAntigravityCountTokens("openai-response", []byte(`{"totalTokens":42}`))); got != `{"object":"response.input_tokens","input_tokens":42}` { + t.Fatalf("Responses CountTokens = %s", got) + } + if got := string(normalizeAntigravityCountTokens("gemini", []byte(`{"totalTokens":42}`))); got != `{"totalTokens":42}` { + t.Fatalf("Gemini CountTokens = %s", got) + } +} diff --git a/third_party/cpaembedded/embedded/embedded.go b/third_party/cpaembedded/embedded/embedded.go index e39c2798..5d7fefc3 100644 --- a/third_party/cpaembedded/embedded/embedded.go +++ b/third_party/cpaembedded/embedded/embedded.go @@ -142,11 +142,13 @@ type HTTPExecutor interface { } type ExecuteRequest struct { + AttemptID string Model string Payload []byte Format string Headers http.Header OriginalRequest []byte + ContinuityKey string } type ExecuteResponse struct { diff --git a/third_party/cpaembedded/go.mod b/third_party/cpaembedded/go.mod index 33ca1727..4536f85f 100644 --- a/third_party/cpaembedded/go.mod +++ b/third_party/cpaembedded/go.mod @@ -4,7 +4,9 @@ go 1.26.0 require ( github.com/andybalholm/brotli v1.0.6 - github.com/router-for-me/CLIProxyAPI/v7 v7.2.133 + github.com/router-for-me/CLIProxyAPI/v7 v7.2.135 + github.com/tidwall/gjson v1.18.0 + github.com/tidwall/sjson v1.2.5 ) require ( @@ -37,10 +39,8 @@ require ( github.com/refraction-networking/utls v1.8.2 // indirect github.com/sirupsen/logrus v1.9.3 // indirect github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 // indirect - github.com/tidwall/gjson v1.18.0 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.0 // indirect - github.com/tidwall/sjson v1.2.5 // indirect github.com/tiktoken-go/tokenizer v0.8.1 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/ugorji/go/codec v1.2.12 // indirect diff --git a/third_party/cpaembedded/go.sum b/third_party/cpaembedded/go.sum index 783e7d94..901429fb 100644 --- a/third_party/cpaembedded/go.sum +++ b/third_party/cpaembedded/go.sum @@ -71,8 +71,8 @@ github.com/redis/go-redis/v9 v9.19.0 h1:XPVaaPSnG6RhYf7p+rmSa9zZfeVAnWsH5h3lxthO github.com/redis/go-redis/v9 v9.19.0/go.mod h1:v/M13XI1PVCDcm01VtPFOADfZtHf8YW3baQf57KlIkA= github.com/refraction-networking/utls v1.8.2 h1:j4Q1gJj0xngdeH+Ox/qND11aEfhpgoEvV+S9iJ2IdQo= github.com/refraction-networking/utls v1.8.2/go.mod h1:jkSOEkLqn+S/jtpEHPOsVv/4V4EVnelwbMQl4vCWXAM= -github.com/router-for-me/CLIProxyAPI/v7 v7.2.133 h1:TeFrEctzhSuPRNdYkiAFpFj4W1IICcvN+ot9HVnaEgw= -github.com/router-for-me/CLIProxyAPI/v7 v7.2.133/go.mod h1:lTHwMAGajc1wKGQiRtDvYbwV0FWsM7sy+N0ZU5/gxJQ= +github.com/router-for-me/CLIProxyAPI/v7 v7.2.135 h1:Vzc5xJwNWq319t5++7AXCx/CnQ3U94xOjteeFtusSsk= +github.com/router-for-me/CLIProxyAPI/v7 v7.2.135/go.mod h1:lTHwMAGajc1wKGQiRtDvYbwV0FWsM7sy+N0ZU5/gxJQ= github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 h1:JIAuq3EEf9cgbU6AtGPK4CTG3Zf6CKMNqf0MHTggAUA= diff --git a/web/src/app/resources/channels.ts b/web/src/app/resources/channels.ts index b17f1997..db073d83 100644 --- a/web/src/app/resources/channels.ts +++ b/web/src/app/resources/channels.ts @@ -20,7 +20,7 @@ export type ChannelFieldInputKind = 'text' | 'url' | 'secret' export type ChannelConnectionType = 'api_key' | 'subscription' export type ChannelAuthorizationMethod = 'browser_oauth' | 'oauth_file' export type ChannelCredentialAction = 'reset_credit' -export type ChannelNoticeID = 'claude_oauth_risk' +export type ChannelNoticeID = 'claude_oauth_risk' | 'antigravity_oauth_risk' export type ChannelNoticeTone = 'warning' export interface ChannelNoticeDto { @@ -127,7 +127,7 @@ const connectionFields = ['type', 'credential_input', 'authorization_methods'] a const capabilityFields = ['model_discovery', 'quota_observation', 'credential_actions'] as const const credentialActions = ['reset_credit'] as const const noticeFields = ['id', 'tone'] as const -const noticeIDs = ['claude_oauth_risk'] as const +const noticeIDs = ['claude_oauth_risk', 'antigravity_oauth_risk'] as const const noticeTones = ['warning'] as const const routeFields = [ 'client_protocol', diff --git a/web/src/assets/channels/antigravity.svg b/web/src/assets/channels/antigravity.svg new file mode 100644 index 00000000..6468839d --- /dev/null +++ b/web/src/assets/channels/antigravity.svg @@ -0,0 +1 @@ +Antigravity diff --git a/web/src/features/groups/credentials/SubscriptionAccountCard.vue b/web/src/features/groups/credentials/SubscriptionAccountCard.vue index b9fd1d9b..847aab7e 100644 --- a/web/src/features/groups/credentials/SubscriptionAccountCard.vue +++ b/web/src/features/groups/credentials/SubscriptionAccountCard.vue @@ -302,6 +302,11 @@ function quotaValueLabel(window: CredentialQuotaWindowDto): string { limit: n(window.limit), }) } + if (window.remaining !== undefined) { + return t('group.credentials.subscription.remainingAmount', { + remaining: n(window.remaining), + }) + } return t('group.credentials.subscription.unknown') } diff --git a/web/src/i18n/locales/en-US/group.ts b/web/src/i18n/locales/en-US/group.ts index 50fdabc2..811c6cd7 100644 --- a/web/src/i18n/locales/en-US/group.ts +++ b/web/src/i18n/locales/en-US/group.ts @@ -402,6 +402,7 @@ export default { modelConstraints: 'Model-specific limits', unknown: 'Not provided', remaining: '{remaining} / {limit} remaining', + remainingAmount: '{remaining} remaining', remainingPercent: '{value}% left', quotaPeriod: 'Current period\n{start} → {end}', quotaExhaustedHint: diff --git a/web/src/i18n/locales/en-US/import.ts b/web/src/i18n/locales/en-US/import.ts index cb773726..6defa23d 100644 --- a/web/src/i18n/locales/en-US/import.ts +++ b/web/src/i18n/locales/en-US/import.ts @@ -111,6 +111,8 @@ export default { channelNotice: { claude_oauth_risk: 'Claude OAuth relies on a Claude Code compatibility protocol rather than a stable public API contract. Upstream changes may interrupt access or affect the account; evaluate this risk before authorizing.', + antigravity_oauth_risk: + 'Antigravity OAuth depends on Google OAuth and Cloud Code interfaces that may change. Upstream policy, permission, or implementation changes can invalidate authorization, interrupt access, or affect the account. Connect only accounts you are authorized to manage and evaluate this risk first.', }, securityNotice: { create: diff --git a/web/src/i18n/locales/ja-JP/group.ts b/web/src/i18n/locales/ja-JP/group.ts index 26007142..c36673ea 100644 --- a/web/src/i18n/locales/ja-JP/group.ts +++ b/web/src/i18n/locales/ja-JP/group.ts @@ -405,6 +405,7 @@ export default { modelConstraints: 'モデル固有の制限', unknown: '未提供', remaining: '残り {remaining} / {limit}', + remainingAmount: '残り {remaining}', remainingPercent: '残り {value}%', quotaPeriod: '現在の期間\n{start} → {end}', quotaExhaustedHint: diff --git a/web/src/i18n/locales/ja-JP/import.ts b/web/src/i18n/locales/ja-JP/import.ts index 40f78b63..1e7d41ce 100644 --- a/web/src/i18n/locales/ja-JP/import.ts +++ b/web/src/i18n/locales/ja-JP/import.ts @@ -109,6 +109,8 @@ export default { channelNotice: { claude_oauth_risk: 'Claude OAuth は安定した公開 API 契約ではなく、Claude Code 互換プロトコルに依存します。上流の変更による接続停止やアカウントへの影響を、認証前に評価してください。', + antigravity_oauth_risk: + 'Antigravity OAuth は Google OAuth と変更される可能性がある Cloud Code インターフェースに依存します。上流のポリシー、権限、実装の変更により認証失効、接続中断、アカウントへの影響が生じる場合があります。管理権限のあるアカウントだけを接続し、事前にリスクを評価してください。', }, securityNotice: { create: diff --git a/web/src/i18n/locales/zh-CN/group.ts b/web/src/i18n/locales/zh-CN/group.ts index 68bc01cb..b0b802bb 100644 --- a/web/src/i18n/locales/zh-CN/group.ts +++ b/web/src/i18n/locales/zh-CN/group.ts @@ -392,6 +392,7 @@ export default { modelConstraints: '模型专项限制', unknown: '上游未提供', remaining: '剩余 {remaining} / {limit}', + remainingAmount: '剩余 {remaining}', remainingPercent: '剩 {value}%', quotaPeriod: '当前周期\n{start} → {end}', quotaExhaustedHint: '额度恢复前会自动跳过此账号；信息过期时回退到现有调度。', diff --git a/web/src/i18n/locales/zh-CN/import.ts b/web/src/i18n/locales/zh-CN/import.ts index b2a1a192..041e9d28 100644 --- a/web/src/i18n/locales/zh-CN/import.ts +++ b/web/src/i18n/locales/zh-CN/import.ts @@ -105,6 +105,8 @@ export default { channelNotice: { claude_oauth_risk: 'Claude OAuth 依赖 Claude Code 的兼容协议，并非稳定的公开 API 契约；上游变化可能导致接入中断或影响账号，请在授权前自行评估风险。', + antigravity_oauth_risk: + 'Antigravity OAuth 依赖 Google OAuth 与可能变化的 Cloud Code 接口；上游策略、权限或实现变化可能导致授权失效、接入中断或账号受限。请仅连接你有权管理的账号，并自行评估风险。', }, securityNotice: { create: '浏览器不会保存访问凭据；服务端加密暂存，创建分组时才正式写入。',