From d1fd07ce3e6007d299571ad2ca1cb4f2427d0ef6 Mon Sep 17 00:00:00 2001 From: tbphp Date: Tue, 18 Aug 2026 21:22:13 +0800 Subject: [PATCH] =?UTF-8?q?feat(subscription):=20=E6=8E=A5=E5=85=A5=20xAI?= =?UTF-8?q?=20OAuth=20=E8=AE=A2=E9=98=85=E6=B8=A0=E9=81=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 1 + README.md | 4 +- README_CN.md | 4 +- README_JP.md | 4 +- THIRD_PARTY_NOTICES.md | 2 +- internal/channel/builtins.go | 1 + internal/channel/channel.go | 3 + internal/channel/compiler_test.go | 14 + .../channel/final_provider_contract_test.go | 1 + internal/channel/grok_module_test.go | 67 ++ internal/channel/modules/grok.go | 52 ++ internal/channel/registry_test.go | 20 +- internal/channel/route_golden_test.go | 2 +- internal/channel/spec/definition.go | 6 +- internal/container/container.go | 1 + .../control/catalog_discovery_merge_test.go | 63 ++ .../control/credential_stage_routes_test.go | 55 ++ internal/control/credential_stages.go | 355 ++++++++++- internal/control/credential_stages_test.go | 249 ++++++++ internal/control/http_routes.go | 7 + internal/control/price_reference_test.go | 17 + internal/control/server.go | 13 + internal/control/service.go | 23 + internal/execution/cpa/adapter.go | 1 + internal/execution/cpa/grok_provider.go | 265 ++++++++ internal/execution/cpa/grok_provider_test.go | 105 ++++ internal/execution/cpa/provider_test.go | 14 + internal/provideradapter/registry_test.go | 1 + internal/storage/db_test.go | 1 + internal/storage/migration.go | 6 + internal/storage/migration_test.go | 1 + .../storage/migrations/0009_device_oauth.go | 190 ++++++ .../migrations/0009_device_oauth_test.go | 44 ++ .../subscription/providers/grok/driver.go | 159 +++++ .../providers/grok/driver_test.go | 71 +++ .../subscription/providers/grok/executor.go | 93 +++ internal/subscription/providers/grok/grok.go | 245 ++++++++ internal/subscription/providers/providers.go | 2 + .../subscription/providers/providers_test.go | 11 +- internal/subscription/runtime/runtime.go | 90 ++- internal/subscription/runtime/runtime_test.go | 59 ++ third_party/cpaembedded/README.md | 20 +- third_party/cpaembedded/embedded/grok.go | 280 +++++++++ .../cpaembedded/embedded/grok_executor.go | 297 +++++++++ .../embedded/grok_executor_test.go | 308 ++++++++++ .../cpaembedded/embedded/grok_models.go | 78 +++ .../cpaembedded/embedded/grok_oauth.go | 577 ++++++++++++++++++ third_party/cpaembedded/embedded/grok_test.go | 279 +++++++++ third_party/cpaembedded/embedded/live_test.go | 116 ++++ web/src/app/resources/channels.ts | 2 +- web/src/app/resources/credential-stages.ts | 37 ++ .../import/SubscriptionCredentialStager.vue | 115 +++- web/src/i18n/locales/en-US/import.ts | 8 + web/src/i18n/locales/ja-JP/import.ts | 9 + web/src/i18n/locales/zh-CN/import.ts | 6 + 55 files changed, 4382 insertions(+), 72 deletions(-) create mode 100644 internal/channel/grok_module_test.go create mode 100644 internal/channel/modules/grok.go create mode 100644 internal/execution/cpa/grok_provider.go create mode 100644 internal/execution/cpa/grok_provider_test.go create mode 100644 internal/storage/migrations/0009_device_oauth.go create mode 100644 internal/storage/migrations/0009_device_oauth_test.go create mode 100644 internal/subscription/providers/grok/driver.go create mode 100644 internal/subscription/providers/grok/driver_test.go create mode 100644 internal/subscription/providers/grok/executor.go create mode 100644 internal/subscription/providers/grok/grok.go create mode 100644 third_party/cpaembedded/embedded/grok.go create mode 100644 third_party/cpaembedded/embedded/grok_executor.go create mode 100644 third_party/cpaembedded/embedded/grok_executor_test.go create mode 100644 third_party/cpaembedded/embedded/grok_models.go create mode 100644 third_party/cpaembedded/embedded/grok_oauth.go create mode 100644 third_party/cpaembedded/embedded/grok_test.go diff --git a/.env.example b/.env.example index 1ced0b8a..f07c4134 100644 --- a/.env.example +++ b/.env.example @@ -7,6 +7,7 @@ HOST=127.0.0.1 BIND_ADDRESS=127.0.0.1 # Subscription browser OAuth uses provider-fixed loopback callbacks: # Codex localhost:1455, Claude localhost:54545, and Antigravity localhost:51121. +# Grok uses device OAuth and does not require a callback port. OAUTH_CALLBACK_BIND_ADDRESS=127.0.0.1 PORT=3001 GRACEFUL_SHUTDOWN_TIMEOUT=10 diff --git a/README.md b/README.md index c739136a..8a3eebf7 100644 --- a/README.md +++ b/README.md @@ -45,14 +45,14 @@ For the maintained 1.4.x release documentation, visit the [official documentatio - **Four client protocols:** OpenAI Completions, OpenAI Responses, Anthropic Messages, and Gemini requests keep their public endpoints. Each channel declares its supported native and converted operations; unsupported feature combinations fail before dispatch. - **Channel and traffic management:** Users select a searchable code-defined channel, then enter models and encrypted credentials. GPT-Load owns AccessKey filtering, cross-Group credential scheduling, retry decisions, health state, cooldown, blacklist, and automatic weights. - **Provider execution:** The official Bifrost Core Go SDK maintains provider-specific authentication, request/response conversion, streaming, normalized usage, and provider errors. GPT-Load pins one selected credential per logical attempt and disables Bifrost's configured retry and fallback policy. -- **Codex and Claude subscriptions:** OpenAI Groups can use API keys or Codex browser OAuth/CPA JSON accounts. The separate Claude channel accepts Claude browser OAuth/CPA JSON accounts, while the existing Anthropic channel continues to use API keys. Both subscription channels share GPT-Load's scheduling, retry, health, weight, and affinity behavior, and their credentials stay encrypted. +- **Subscription channels:** Codex, Claude, Antigravity, and Grok are independent subscription-backed channels alongside the existing API-key channels. Codex, Claude, and Antigravity use browser OAuth/CPA JSON; Grok uses xAI device OAuth/CPA JSON. All reuse GPT-Load's encrypted credential lifecycle, scheduling, retry, health, weight, and affinity behavior. - **Control and observability:** runtime settings, route inspection, health views, RequestLog, and a Chinese, English, and Japanese admin UI. - **Usage and estimated cost:** usage extraction for the four protocols where the endpoint returns generation usage, 24-hour/30-day reports, per-request quality states, exact four-slot model prices synchronized from Models.dev where available, and user-managed prices. The M3 control-plane UI and M4 usage/pricing scope are present in the local candidate, but their formal exit and public release are unfinished. Prices and costs are best-effort **estimates** derived from upstream usage and the active pricing rules. They are not a billing ledger, invoice, or provider bill, and historical requests are not repriced. > [!WARNING] -> Claude subscription OAuth and execution rely on a pinned, compatibility-sensitive CPA contract rather than the Anthropic API-key contract. Upstream changes may require a reviewed CPA bridge update; validate the channel with an authorized account before production use. +> Subscription OAuth execution relies on pinned, compatibility-sensitive CPA contracts rather than the corresponding API-key contracts. Upstream changes may require a reviewed CPA bridge update; validate each enabled channel with an authorized account before production use. ## 2.0.0 support boundaries diff --git a/README_CN.md b/README_CN.md index 429e4134..8fa3dc19 100644 --- a/README_CN.md +++ b/README_CN.md @@ -45,14 +45,14 @@ GPT-Load 是一个用 Go 构建的自托管多渠道 AI 网关。它通过单个 - **四种客户端协议**：继续提供 OpenAI Completions、OpenAI Responses、Anthropic Messages 与 Gemini 入口。每个渠道明确声明可原生执行或转换的 operation；不支持的能力组合会在发送前拒绝。 - **渠道与流量管理**：用户只需搜索并选择代码内置渠道，再填写模型和加密凭据。GPT-Load 继续负责 AccessKey 筛选、跨 Group 的凭据调度、重试决策、健康状态、cooldown、blacklist 与自动权重。 - **Provider 执行**：官方 Bifrost Core Go SDK 负责不同 Provider 的认证、请求/响应转换、流式、usage 归一化和错误解析。每个逻辑 attempt 由 GPT-Load 固定一个凭据，并关闭 Bifrost 可配置的重试和 fallback。 -- **Codex 与 Claude 订阅**：OpenAI Group 可选择 API Key，或通过浏览器 OAuth/CPA JSON 接入 Codex 订阅账号；独立的 Claude 渠道通过浏览器 OAuth/CPA JSON 接入 Claude 订阅账号，原 Anthropic 渠道继续使用 API Key。两类订阅渠道统一复用 GPT-Load 的调度、重试、健康、权重和亲和机制，凭据始终加密。 +- **订阅类型渠道**：Codex、Claude、Antigravity 与 Grok 都是独立于现有 API Key 渠道的订阅渠道。Codex、Claude、Antigravity 使用浏览器 OAuth/CPA JSON，Grok 使用 xAI Device OAuth/CPA JSON；统一复用 GPT-Load 的加密凭据生命周期、调度、重试、健康、权重和亲和机制。 - **控制与可观测性**：运行设置、路由检查、健康视图、RequestLog，以及中文、英文、日文管理 UI。 - **用量与估算成本**：对四种协议中会返回生成 usage 的接口进行采集，提供 24 小时/30 天汇总、明细质量状态、可用时从 Models.dev 同步的精确四槽模型价格，以及用户管理的价格。 M3 控制面 UI 与 M4 用量/定价范围已经进入本地候选，但正式出口与公开发布尚未完成。价格和成本是基于上游返回 usage 与当前价格规则的 best-effort **估算**，不是 billing ledger、发票或供应商账单，也不会对历史请求重新计价。 > [!WARNING] -> Claude 订阅 OAuth 与执行依赖固定版本 CPA 提供的兼容性协议，不属于 Anthropic API Key 契约。上游变更可能需要审查并更新 CPA bridge；生产使用前应通过已授权账号完成真实合同验证。 +> 订阅 OAuth 执行依赖固定版本 CPA 提供的兼容性协议，不属于同品牌 API Key 契约。上游变更可能需要审查并更新 CPA bridge；生产使用前应通过已授权账号完成各渠道真实合同验证。 ## 2.0.0 支持边界 diff --git a/README_JP.md b/README_JP.md index b5b3ad9a..ebf63cde 100644 --- a/README_JP.md +++ b/README_JP.md @@ -45,14 +45,14 @@ GPT-Loadは、Goで構築されたセルフホスト型のマルチチャネルA - **4つのクライアントプロトコル**：OpenAI Completions、OpenAI Responses、Anthropic Messages、Geminiの公開エンドポイントを維持します。各チャネルはネイティブまたは変換可能なoperationを明示し、未対応の機能組み合わせは送信前に拒否します。 - **チャネルとトラフィックの管理**：ユーザーは検索可能なコード定義チャネルを選び、モデルと暗号化認証情報を入力します。GPT-LoadがAccessKeyフィルター、Groupをまたぐ認証情報スケジューリング、リトライ判断、ヘルス、cooldown、blacklist、自動重み付けを所有します。 - **Provider実行**：公式Bifrost Core Go SDKがProvider固有の認証、request/response変換、streaming、usage正規化、エラー解析を担当します。GPT-Loadは論理attemptごとに1つの認証情報を固定し、Bifrostの設定可能なretryとfallbackを無効にします。 -- **CodexとClaudeのサブスクリプション**：OpenAI GroupではAPI Key、またはブラウザーOAuth/CPA JSONによるCodexアカウントを選択できます。独立したClaudeチャネルはブラウザーOAuth/CPA JSONでClaudeアカウントを受け付け、既存のAnthropicチャネルは引き続きAPI Keyを使用します。両サブスクリプションチャネルはGPT-Loadのスケジューリング、retry、health、weight、affinityを共有し、認証情報は暗号化されます。 +- **サブスクリプションチャネル**：Codex、Claude、Antigravity、Grokは、既存のAPI Keyチャネルとは独立したサブスクリプションチャネルです。Codex、Claude、AntigravityはブラウザーOAuth/CPA JSON、GrokはxAI Device OAuth/CPA JSONを使用し、暗号化された認証情報ライフサイクル、スケジューリング、retry、health、weight、affinityを共通利用します。 - **制御と可観測性**：ランタイム設定、ルート検査、ヘルス表示、RequestLog、中国語・英語・日本語の管理UI。 - **使用量と推定コスト**：4プロトコルのうち生成usageを返すエンドポイントからusageを取得し、24時間/30日レポート、リクエスト単位の品質状態、利用可能な場合にModels.devから同期する完全一致の4価格スロット、ユーザー管理価格を提供します。 M3のコントロールプレーンUIとM4のusage/pricing範囲はローカル候補に含まれていますが、正式なリリース判断・承認と公開は未完了です。価格とコストは、上流から返されたusageと現在の価格ルールに基づくbest-effortの**推定値**です。billing ledger、請求書、プロバイダー請求ではなく、過去のリクエストを再計算することもありません。 > [!WARNING] -> ClaudeサブスクリプションのOAuthと実行は、Anthropic API Key契約ではなく、固定バージョンCPAの互換性に敏感な契約に依存します。上流変更時はCPA bridgeのレビューと更新が必要になる可能性があるため、本番利用前に許可されたアカウントで実契約を検証してください。 +> サブスクリプションOAuthの実行は、同ブランドのAPI Key契約ではなく、固定バージョンCPAの互換性に敏感な契約に依存します。上流変更時はCPA bridgeのレビューと更新が必要になる可能性があるため、本番利用前に許可されたアカウントで各チャネルの実契約を検証してください。 ## 2.0.0のサポート境界 diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 991187ad..e8507c11 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -26,7 +26,7 @@ The complete Apache License 2.0 text is distributed in - License: MIT License GPT-Load uses a pinned, execution-only embedded adapter around CLIProxyAPI's -Codex, Claude, and Antigravity OAuth and HTTP executor code. GPT-Load retains ownership of +Codex, Claude, Antigravity, and xAI 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/internal/channel/builtins.go b/internal/channel/builtins.go index f3126fee..d117d693 100644 --- a/internal/channel/builtins.go +++ b/internal/channel/builtins.go @@ -13,6 +13,7 @@ func builtInModules() []spec.Module { modules.Codex(), modules.Claude(), modules.Antigravity(), + modules.Grok(), modules.Anthropic(), modules.Gemini(), modules.AzureOpenAI(), diff --git a/internal/channel/channel.go b/internal/channel/channel.go index f3de8c9c..02c1dd5e 100644 --- a/internal/channel/channel.go +++ b/internal/channel/channel.go @@ -25,6 +25,7 @@ const ( Codex = spec.Codex Claude = spec.Claude Antigravity = spec.Antigravity + Grok = spec.Grok Anthropic = spec.Anthropic Gemini = spec.Gemini AzureOpenAI = spec.AzureOpenAI @@ -74,6 +75,7 @@ type AuthorizationMethod = spec.AuthorizationMethod const ( AuthorizationBrowserOAuth = spec.AuthorizationBrowserOAuth + AuthorizationDeviceOAuth = spec.AuthorizationDeviceOAuth AuthorizationOAuthFile = spec.AuthorizationOAuthFile ) @@ -185,6 +187,7 @@ const ( ProviderCodex = spec.ProviderCodex ProviderClaude = spec.ProviderClaude ProviderAntigravity = spec.ProviderAntigravity + ProviderGrok = spec.ProviderGrok ProviderAnthropic = spec.ProviderAnthropic ProviderGemini = spec.ProviderGemini ProviderOpenAICompatible = spec.ProviderOpenAICompatible diff --git a/internal/channel/compiler_test.go b/internal/channel/compiler_test.go index a37d9d94..5a48ed77 100644 --- a/internal/channel/compiler_test.go +++ b/internal/channel/compiler_test.go @@ -50,6 +50,20 @@ func TestCompilerRejectsUnknownAuthorizationMethod(t *testing.T) { } } +func TestCompilerAcceptsDeviceOAuthAuthorizationMethod(t *testing.T) { + t.Parallel() + + codex := findModule(t, builtInModules(), Codex) + codex.Definition.Connection.AuthorizationMethods = []spec.AuthorizationMethod{ + "device_oauth", + spec.AuthorizationOAuthFile, + } + + if _, err := compileBuiltInModules([]spec.Module{codex}); err != nil { + t.Fatalf("compileBuiltInModules() rejected device OAuth: %v", err) + } +} + func TestCompilerRejectsQuotaPriorityWithoutObservation(t *testing.T) { t.Parallel() diff --git a/internal/channel/final_provider_contract_test.go b/internal/channel/final_provider_contract_test.go index 8cfde493..388a7df1 100644 --- a/internal/channel/final_provider_contract_test.go +++ b/internal/channel/final_provider_contract_test.go @@ -16,6 +16,7 @@ func TestFinalRegistryContainsOnlyApprovedChannels(t *testing.T) { Codex, Claude, Antigravity, + Grok, Anthropic, Gemini, AzureOpenAI, diff --git a/internal/channel/grok_module_test.go b/internal/channel/grok_module_test.go new file mode 100644 index 00000000..93046416 --- /dev/null +++ b/internal/channel/grok_module_test.go @@ -0,0 +1,67 @@ +package channel + +import ( + "reflect" + "testing" + + "gpt-load/internal/channel/modules" + "gpt-load/internal/channel/spec" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +func TestGrokModuleDeclaresCompleteSubscriptionContract(t *testing.T) { + t.Parallel() + + definitions, err := compileBuiltInModules([]spec.Module{modules.Grok()}) + if err != nil { + t.Fatal(err) + } + registry, err := newRegistry(definitions) + if err != nil { + t.Fatal(err) + } + descriptor, ok := registry.Get(Grok) + if !ok { + t.Fatal("Grok descriptor is missing") + } + if descriptor.Connection.Type != string(spec.ConnectionSubscription) || + !reflect.DeepEqual(descriptor.Connection.AuthorizationMethods, []AuthorizationMethod{ + AuthorizationDeviceOAuth, + AuthorizationOAuthFile, + }) { + t.Fatalf("Grok connection = %#v", descriptor.Connection) + } + target, err := registry.Resolve(Grok, nil) + if err != nil { + t.Fatal(err) + } + if target.ProviderKind != ProviderGrok || target.CatalogProviderID != "" { + t.Fatalf("Grok target = %#v", target) + } + bindings, ok := registry.CapabilityBindings(Grok) + if !ok || bindings.SubscriptionDriver != modules.GrokSubscriptionDriver || + bindings.ModelDiscovery != modules.GrokModelDiscovery || bindings.QuotaObservation != "" { + t.Fatalf("Grok capabilities = %#v, %t", bindings, ok) + } + for _, test := range []struct { + clientProtocol protocol.Protocol + operation execution.Operation + mode RouteMode + }{ + {protocol.OpenAIResponses, execution.OperationResponsesCreate, RouteNative}, + {protocol.OpenAIResponses, execution.OperationResponsesInputTokens, RouteNative}, + {protocol.OpenAICompletions, execution.OperationChatCompletion, RouteConverted}, + {protocol.Anthropic, execution.OperationChatCompletion, RouteConverted}, + {protocol.Anthropic, execution.OperationCountTokens, RouteConverted}, + {protocol.Gemini, execution.OperationChatCompletion, RouteConverted}, + {protocol.Gemini, execution.OperationCountTokens, RouteConverted}, + } { + if mode, exists := target.Mode(test.clientProtocol, test.operation); !exists || mode != test.mode { + t.Fatalf("Grok %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("Grok unexpectedly declares Responses Compact") + } +} diff --git a/internal/channel/modules/grok.go b/internal/channel/modules/grok.go new file mode 100644 index 00000000..7d809585 --- /dev/null +++ b/internal/channel/modules/grok.go @@ -0,0 +1,52 @@ +package modules + +import ( + "gpt-load/internal/channel/spec" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +const ( + GrokSubscriptionDriver spec.SubscriptionDriverID = "grok" + GrokModelDiscovery spec.UtilityID = "grok_models" +) + +// Grok declares the subscription-backed xAI OAuth channel. +func Grok() spec.Module { + return spec.Module{Definition: spec.Definition{ + ID: spec.Grok, + Name: "Grok", + Mark: "GR", + Icon: "xai", + SearchTerms: []string{"subscription", "oauth", "xai", "grok"}, + Description: "xAI Grok subscription", + Connection: spec.Connection{ + Type: spec.ConnectionSubscription, + CredentialInput: "authorization", + AuthorizationMethods: []spec.AuthorizationMethod{ + spec.AuthorizationDeviceOAuth, + spec.AuthorizationOAuthFile, + }, + }, + Params: []spec.Field{}, + Credentials: []spec.Field{}, + Provider: spec.ProviderBinding{ + ProviderKind: spec.ProviderGrok, + EndpointPolicy: spec.EndpointNone, + }, + Routes: []spec.Route{ + spec.NewRoute(protocol.OpenAIResponses, execution.OperationResponsesCreate, execution.RouteNative), + spec.NewRoute(protocol.OpenAIResponses, execution.OperationResponsesInputTokens, execution.RouteNative), + spec.NewRoute(protocol.OpenAICompletions, execution.OperationChatCompletion, execution.RouteConverted), + spec.NewRoute(protocol.Anthropic, execution.OperationChatCompletion, execution.RouteConverted), + spec.NewRoute(protocol.Anthropic, execution.OperationCountTokens, execution.RouteConverted), + spec.NewRoute(protocol.Gemini, execution.OperationChatCompletion, execution.RouteConverted), + spec.NewRoute(protocol.Gemini, execution.OperationCountTokens, execution.RouteConverted), + }, + Capabilities: spec.CapabilityBindings{ + SubscriptionDriver: GrokSubscriptionDriver, + ModelDiscovery: GrokModelDiscovery, + }, + Scheduling: spec.SchedulingPolicy{QuotaPriority: false}, + }} +} diff --git a/internal/channel/registry_test.go b/internal/channel/registry_test.go index 482b7fab..0912488f 100644 --- a/internal/channel/registry_test.go +++ b/internal/channel/registry_test.go @@ -18,6 +18,7 @@ func TestRegistryHasStableBuiltInOrderAndSearch(t *testing.T) { Codex, Claude, Antigravity, + Grok, Anthropic, Gemini, ID("azure_openai"), @@ -40,8 +41,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, Antigravity}) { - t.Fatalf("Search(subscription) IDs = %v, want [codex claude antigravity]", got) + if got := descriptorIDs(registry.Search("subscription")); !reflect.DeepEqual(got, []ID{Codex, Claude, Antigravity, Grok}) { + t.Fatalf("Search(subscription) IDs = %v, want [codex claude antigravity grok]", got) } if got := descriptorIDs(registry.Search("compatible")); !reflect.DeepEqual(got, []ID{OpenAICompatible}) { t.Fatalf("Search(compatible) IDs = %v", got) @@ -65,7 +66,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 := 8 + vertexIndex := 9 if first[vertexIndex].ID != GoogleVertex || first[vertexIndex].ParamFields[0].DefaultValue == nil { t.Fatalf("unexpected Vertex descriptor = %#v", first[vertexIndex]) } @@ -250,12 +251,14 @@ func TestRegistryPublicDescriptorUsesSingularConnectionAndExplicitRoutes(t *test } } -func TestOnlyCodexSupportsSubscription(t *testing.T) { +func TestSubscriptionChannelsRemainSeparateFromAPIKeyChannels(t *testing.T) { t.Parallel() registry := NewRegistry() - if !registry.SupportsConnectionType(Codex, "subscription") { - t.Fatal("Codex subscription is not supported") + for _, id := range []ID{Codex, Claude, Antigravity, Grok} { + if !registry.SupportsConnectionType(id, "subscription") { + t.Fatalf("%s subscription is not supported", id) + } } if registry.SupportsConnectionType(OpenAI, "subscription") || registry.SupportsConnectionType(Anthropic, "subscription") || @@ -273,7 +276,7 @@ func TestOnlyCodexSupportsSubscription(t *testing.T) { func TestRegistryReturnsExactCatalogProviderMappingWithoutResolvingParams(t *testing.T) { registry := NewRegistry() for id, want := range map[ID]string{ - OpenAI: "openai", Codex: "", Claude: "", Anthropic: "anthropic", Gemini: "google", + OpenAI: "openai", Codex: "", Claude: "", Antigravity: "", Grok: "", Anthropic: "anthropic", Gemini: "google", ID("azure_openai"): "azure", ID("aws_bedrock"): "amazon-bedrock", ID("google_vertex"): "google-vertex", OpenAICompatible: "", } { @@ -295,6 +298,9 @@ func TestRegistryReturnsProviderKindWithoutExposingItInDescriptor(t *testing.T) if got, ok := registry.ProviderKind(Codex); !ok || string(got) != "codex" { t.Fatalf("ProviderKind(codex) = %q, %t, want codex, true", got, ok) } + if got, ok := registry.ProviderKind(Grok); !ok || got != ProviderGrok { + t.Fatalf("ProviderKind(grok) = %q, %t", got, ok) + } if got, ok := registry.ProviderKind(ID("missing")); ok || got != "" { t.Fatalf("ProviderKind(missing) = %q, %t", got, ok) } diff --git a/internal/channel/route_golden_test.go b/internal/channel/route_golden_test.go index 33cc2ad3..fdbf1773 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 = "fb123f1bc2a0afdb64a29e525eab265933a4c7ab6ba8057d73303293569887f2" + const wantDigest = "4fb788d9792b871c7e2ea00d073756fb3464841370b85dcbfbbb9a904496dea8" 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 a2dca98d..9d8f66af 100644 --- a/internal/channel/spec/definition.go +++ b/internal/channel/spec/definition.go @@ -17,6 +17,7 @@ const ( Codex ID = "codex" Claude ID = "claude" Antigravity ID = "antigravity" + Grok ID = "grok" Anthropic ID = "anthropic" Gemini ID = "gemini" AzureOpenAI ID = "azure_openai" @@ -42,6 +43,7 @@ const ( ProviderCodex ProviderKind = "codex" ProviderClaude ProviderKind = "claude" ProviderAntigravity ProviderKind = "antigravity" + ProviderGrok ProviderKind = "grok" ProviderAnthropic ProviderKind = "anthropic" ProviderGemini ProviderKind = "gemini" ProviderOpenAICompatible ProviderKind = "openai_compatible" @@ -90,6 +92,7 @@ func (kind ProviderKind) Valid() bool { ProviderCodex, ProviderClaude, ProviderAntigravity, + ProviderGrok, ProviderAnthropic, ProviderGemini, ProviderOpenAICompatible, @@ -143,13 +146,14 @@ type AuthorizationMethod string const ( AuthorizationBrowserOAuth AuthorizationMethod = "browser_oauth" + AuthorizationDeviceOAuth AuthorizationMethod = "device_oauth" AuthorizationOAuthFile AuthorizationMethod = "oauth_file" ) // Valid reports whether the authorization method is part of the public channel contract. func (method AuthorizationMethod) Valid() bool { switch method { - case AuthorizationBrowserOAuth, AuthorizationOAuthFile: + case AuthorizationBrowserOAuth, AuthorizationDeviceOAuth, AuthorizationOAuthFile: return true default: return false diff --git a/internal/container/container.go b/internal/container/container.go index 1fd47b9e..a4e16e8e 100644 --- a/internal/container/container.go +++ b/internal/container/container.go @@ -228,6 +228,7 @@ func newProviderAdapterRegistry( {ProviderKind: channel.ProviderCodex, Adapter: cpa}, {ProviderKind: channel.ProviderClaude, Adapter: cpa}, {ProviderKind: channel.ProviderAntigravity, Adapter: cpa}, + {ProviderKind: channel.ProviderGrok, 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 10e89be8..59af5c51 100644 --- a/internal/control/catalog_discovery_merge_test.go +++ b/internal/control/catalog_discovery_merge_test.go @@ -14,6 +14,7 @@ import ( "gpt-load/internal/storage/models" "gpt-load/internal/subscription/providers/antigravity" "gpt-load/internal/subscription/providers/codex" + "gpt-load/internal/subscription/providers/grok" subscriptionruntime "gpt-load/internal/subscription/runtime" ) @@ -165,6 +166,68 @@ func TestAntigravityDiscoveryUsesOnlySubscriptionModelsAndReferencePrices(t *tes } } +func TestGrokDiscoveryUsesOnlySubscriptionModelsAndReferencePrices(t *testing.T) { + fixture := newServiceFixture(t) + fixture.catalogRuntime.Publish(&catalog.Snapshot{Providers: map[string]catalog.Provider{ + "xai": { + ID: "xai", Name: "xAI", Models: map[string]catalog.Model{ + "grok-4.3": { + ID: "grok-4.3", Name: "xAI catalog name", + Cost: &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(1)}}, + }, + "xai-catalog-only": { + ID: "xai-catalog-only", Name: "xAI catalog only", + Cost: &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(2)}}, + }, + }, + }, + }}) + canonical, err := grok.MarshalCredential(grok.Credential{ + Type: grok.Provider, AccessToken: "access", RefreshToken: "refresh", AccountID: "xai-account", + Email: "grok@example.com", Expire: "2030-01-01T00:00:00Z", TokenEndpoint: "https://auth.x.ai/oauth/token", + }) + if err != nil { + t.Fatal(err) + } + driver, ok := fixture.service.subscriptions.Driver(channel.Grok) + if !ok { + t.Fatal("Grok driver is unavailable") + } + credential, err := driver.Parse(canonical) + if err != nil { + t.Fatal(err) + } + stage, err := fixture.service.persistReadyCredentialStage(t.Context(), channel.Grok, "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.Grok { + t.Fatalf("channel = %q, want Grok", channelID) + } + return []string{"grok-4.3"}, nil + } + + got, err := fixture.service.DiscoverModels(t.Context(), ModelDiscoveryRequest{ + ChannelID: channel.Grok, StagedCredentialID: stage.StageID, ConnectionType: "subscription", + }) + if err != nil { + t.Fatal(err) + } + xai := "xAI" + want := []ModelCandidate{{ + ID: "grok-4.3", Name: "grok-4.3", Sources: []string{"live"}, + PricingStatus: PricingStatusConfigured, PricingSource: &xai, + }} + if !reflect.DeepEqual(got.Models, want) { + t.Fatalf("Grok 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 0adc98f4..38867ad1 100644 --- a/internal/control/credential_stage_routes_test.go +++ b/internal/control/credential_stage_routes_test.go @@ -13,6 +13,7 @@ import ( "net/url" "strings" "testing" + "time" "github.com/gin-gonic/gin" @@ -77,6 +78,60 @@ func TestCredentialStageRoutesRequireAuthAndNeverReturnSecrets(t *testing.T) { } } +func TestDeviceAuthorizationRouteReturnsSafeChallengeAndPollsByPOST(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + now := fixture.service.now().UTC() + fixture.service.now = func() time.Time { return now } + fixture.service.beginDeviceAuthorization = func(context.Context, channel.ID) (subscriptionruntime.DeviceAuthorization, error) { + return subscriptionruntime.DeviceAuthorization{ + VerificationURL: "https://auth.x.ai/device", UserCode: "ABCD-EFGH", + DriverState: []byte(`{"device_code":"device-secret"}`), ExpiresAt: now.Add(20 * time.Minute), + PollInterval: time.Second, + }, nil + } + fixture.service.pollDeviceAuthorization = func(context.Context, channel.ID, []byte) (subscriptionruntime.DeviceAuthorizationPoll, error) { + return subscriptionruntime.DeviceAuthorizationPoll{ + Status: subscriptionruntime.DeviceAuthorizationPending, + DriverState: []byte(`{"device_code":"next-secret"}`), PollInterval: 5 * time.Second, + }, nil + } + listenerCalled := false + fixture.service.oauthCallback.listen = func(string, string) (net.Listener, error) { + listenerCalled = true + return nil, fmt.Errorf("device OAuth must not start a listener") + } + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + + request := httptest.NewRequest(http.MethodPost, "/api/credential-stages/authorizations", strings.NewReader(`{"channel_id":"grok"}`)) + request.Header.Set("Authorization", "Bearer test-auth-key") + request.Header.Set("Content-Type", "application/json") + response := httptest.NewRecorder() + engine.ServeHTTP(response, request) + if response.Code != http.StatusOK || listenerCalled || strings.Contains(response.Body.String(), "device-secret") { + t.Fatalf("device begin response = %d %s, listener=%t", response.Code, response.Body, listenerCalled) + } + var envelope struct { + Data CredentialStageResult `json:"data"` + } + if err := json.Unmarshal(response.Body.Bytes(), &envelope); err != nil { + t.Fatal(err) + } + if envelope.Data.AuthorizationMethod != string(channel.AuthorizationDeviceOAuth) || envelope.Data.UserCode != "ABCD-EFGH" { + t.Fatalf("device challenge = %#v", envelope.Data) + } + + now = now.Add(time.Second) + poll := httptest.NewRequest(http.MethodPost, "/api/credential-stages/"+envelope.Data.StageID+"/device-poll", nil) + poll.Header.Set("Authorization", "Bearer test-auth-key") + pollResponse := httptest.NewRecorder() + engine.ServeHTTP(pollResponse, poll) + if pollResponse.Code != http.StatusOK || strings.Contains(pollResponse.Body.String(), "next-secret") { + t.Fatalf("device poll response = %d %s", pollResponse.Code, pollResponse.Body) + } +} + func TestBeginCredentialAuthorizationAllowsManualCallbackWhenListenerIsUnavailable(t *testing.T) { initControlI18n(t) fixture := newServiceFixture(t) diff --git a/internal/control/credential_stages.go b/internal/control/credential_stages.go index fc9fa70e..bdf4bde2 100644 --- a/internal/control/credential_stages.go +++ b/internal/control/credential_stages.go @@ -7,6 +7,7 @@ import ( "fmt" "net" "net/http" + "net/url" "sort" "strings" "time" @@ -23,8 +24,11 @@ import ( const ( credentialStageReadyTTL = 30 * time.Minute credentialStageAuthTTL = 5 * time.Minute + credentialStageDeviceMaxTTL = 30 * time.Minute credentialStageTombstoneTTL = 24 * time.Hour maxOAuthFileBytes = 64 * 1024 + maxDeviceAuthorizationURL = 4096 + maxDeviceAuthorizationCode = 128 stagedSubscriptionSchemaV2 = 2 ) @@ -63,13 +67,27 @@ type CredentialStageAccount struct { } type CredentialStageResult struct { - StageID string `json:"stage_id"` - Status string `json:"status"` - AuthorizationURL string `json:"authorization_url,omitempty"` - RedirectURI string `json:"redirect_uri,omitempty"` - Account CredentialStageAccount `json:"account"` - ExpiresAtMS int64 `json:"expires_at_ms"` - ErrorCode string `json:"error_code,omitempty"` + StageID string `json:"stage_id"` + Status string `json:"status"` + AuthorizationMethod string `json:"authorization_method,omitempty"` + AuthorizationURL string `json:"authorization_url,omitempty"` + RedirectURI string `json:"redirect_uri,omitempty"` + UserCode string `json:"user_code,omitempty"` + NextPollAtMS int64 `json:"next_poll_at_ms,omitempty"` + Account CredentialStageAccount `json:"account"` + ExpiresAtMS int64 `json:"expires_at_ms"` + ErrorCode string `json:"error_code,omitempty"` +} + +type credentialStageAuthorizationSummary struct { + AuthorizationURL string `json:"authorization_url"` + UserCode string `json:"user_code"` + NextPollAtMS int64 `json:"next_poll_at_ms"` + PollIntervalMS int64 `json:"poll_interval_ms"` +} + +type credentialStageSafeSummary struct { + Authorization *credentialStageAuthorizationSummary `json:"authorization,omitempty"` } type stagedSubscriptionPayload struct { @@ -349,13 +367,25 @@ func (s *Service) BeginCredentialAuthorization( if s == nil || s.db == nil || s.encryption == nil { return CredentialStageResult{}, app_errors.ErrValidation } - if s.channelRegistry == nil || - !s.channelRegistry.SupportsAuthorizationMethod(channelID, channel.AuthorizationBrowserOAuth) { + if s.channelRegistry == nil { return CredentialStageResult{}, app_errors.ErrValidation } if _, err := s.subscriptionDriver(channelID); err != nil { return CredentialStageResult{}, err } + if s.channelRegistry.SupportsAuthorizationMethod(channelID, channel.AuthorizationBrowserOAuth) { + return s.beginBrowserCredentialAuthorization(ctx, channelID) + } + if s.channelRegistry.SupportsAuthorizationMethod(channelID, channel.AuthorizationDeviceOAuth) { + return s.beginDeviceCredentialAuthorization(ctx, channelID) + } + return CredentialStageResult{}, app_errors.ErrValidation +} + +func (s *Service) beginBrowserCredentialAuthorization( + ctx context.Context, + channelID channel.ID, +) (CredentialStageResult, error) { if s.beginSubscriptionAuthorization == nil { return CredentialStageResult{}, app_errors.ErrAuthorizationUnavailable } @@ -390,12 +420,98 @@ func (s *Service) BeginCredentialAuthorization( return CredentialStageResult{}, app_errors.ParseDBError(err) } return CredentialStageResult{ - StageID: row.ID, Status: string(row.Status), AuthorizationURL: authorization.URL, - RedirectURI: authorization.RedirectURI, - Account: CredentialStageAccount{}, ExpiresAtMS: row.ExpiresAtMS, + StageID: row.ID, Status: string(row.Status), AuthorizationMethod: row.AuthorizationMethod, + AuthorizationURL: authorization.URL, + RedirectURI: authorization.RedirectURI, + Account: CredentialStageAccount{}, ExpiresAtMS: row.ExpiresAtMS, }, nil } +func (s *Service) beginDeviceCredentialAuthorization( + ctx context.Context, + channelID channel.ID, +) (CredentialStageResult, error) { + if s.beginDeviceAuthorization == nil { + return CredentialStageResult{}, app_errors.ErrAuthorizationUnavailable + } + beginContext, cancelBegin := context.WithTimeout(ctx, defaultSubscriptionControlTimeout) + defer cancelBegin() + authorization, err := s.beginDeviceAuthorization(beginContext, channelID) + if err != nil { + return CredentialStageResult{}, app_errors.ErrAuthorizationUnavailable + } + now := s.now().UTC() + verificationURL := strings.TrimSpace(authorization.VerificationURL) + userCode := strings.TrimSpace(authorization.UserCode) + parsedURL, parseErr := url.Parse(verificationURL) + pollInterval, intervalOK := normalizeDevicePollInterval(authorization.PollInterval) + if parseErr != nil || !parsedURL.IsAbs() || parsedURL.Opaque != "" || parsedURL.Host == "" || + !strings.EqualFold(parsedURL.Scheme, "https") || parsedURL.User != nil || parsedURL.Fragment != "" || + len(verificationURL) > maxDeviceAuthorizationURL || !validDeviceAuthorizationCode(userCode) || + len(authorization.DriverState) == 0 || len(authorization.DriverState) > maxOAuthFileBytes || + !intervalOK || !authorization.ExpiresAt.After(now) || + authorization.ExpiresAt.After(now.Add(credentialStageDeviceMaxTTL)) { + return CredentialStageResult{}, app_errors.ErrAuthorizationUnavailable + } + nextPollAtMS := now.Add(pollInterval).UnixMilli() + summaryJSON, err := json.Marshal(credentialStageSafeSummary{Authorization: &credentialStageAuthorizationSummary{ + AuthorizationURL: verificationURL, + UserCode: userCode, + NextPollAtMS: nextPollAtMS, + PollIntervalMS: pollInterval.Milliseconds(), + }}) + if err != nil { + return CredentialStageResult{}, app_errors.ErrInternalServer + } + payload, err := json.Marshal(stagedSubscriptionPayload{DriverState: authorization.DriverState}) + if err != nil { + return CredentialStageResult{}, app_errors.ErrInternalServer + } + ciphertext, err := s.encryption.Encrypt(string(payload)) + clear(payload) + if err != nil { + return CredentialStageResult{}, app_errors.ErrInternalServer + } + stageID, err := newOperationID(s.random) + if err != nil { + return CredentialStageResult{}, app_errors.ErrInternalServer + } + row := models.CredentialStage{ + ID: stageID, ChannelID: string(channelID), ConnectionType: models.ConnectionTypeSubscription, + AuthorizationMethod: string(channel.AuthorizationDeviceOAuth), Status: models.CredentialStagePendingAuthorization, + EncryptedPayload: ciphertext, PayloadSchemaVersion: stagedSubscriptionSchemaV2, + SafeSummaryJSON: models.JSON(summaryJSON), ExpiresAtMS: authorization.ExpiresAt.UTC().UnixMilli(), + CreatedAtMS: now.UnixMilli(), UpdatedAtMS: now.UnixMilli(), + } + if err := s.db.WithContext(ctx).Create(&row).Error; err != nil { + return CredentialStageResult{}, app_errors.ParseDBError(err) + } + return CredentialStageResult{ + StageID: row.ID, Status: string(row.Status), AuthorizationMethod: row.AuthorizationMethod, + AuthorizationURL: verificationURL, UserCode: userCode, NextPollAtMS: nextPollAtMS, + Account: CredentialStageAccount{}, ExpiresAtMS: row.ExpiresAtMS, + }, nil +} + +func normalizeDevicePollInterval(value time.Duration) (time.Duration, bool) { + if value < time.Second || value > time.Minute { + return 0, false + } + return value, true +} + +func validDeviceAuthorizationCode(value string) bool { + if value == "" || len(value) > maxDeviceAuthorizationCode { + return false + } + for index := 0; index < len(value); index++ { + if value[index] < 0x20 || value[index] > 0x7e { + return false + } + } + return true +} + // CompleteCredentialAuthorization atomically claims one pending state before // issuing the one-shot OAuth code exchange. func (s *Service) CompleteCredentialAuthorization( @@ -655,6 +771,189 @@ func (s *Service) CompleteCredentialAuthorizationCallback( ) } +// PollCredentialDeviceAuthorization performs at most one provider token poll. +// Calling it before the advertised interval returns the current Stage without +// dispatching upstream. +func (s *Service) PollCredentialDeviceAuthorization( + ctx context.Context, + stageID string, +) (CredentialStageResult, error) { + if s == nil || s.db == nil || s.encryption == nil || s.pollDeviceAuthorization == nil || + strings.TrimSpace(stageID) == "" { + return CredentialStageResult{}, app_errors.ErrAuthorizationStateInvalid + } + row, err := s.loadCredentialStage(ctx, stageID) + if err != nil { + return CredentialStageResult{}, err + } + if row.AuthorizationMethod != string(channel.AuthorizationDeviceOAuth) { + return CredentialStageResult{}, app_errors.ErrAuthorizationStateInvalid + } + if row.Status == models.CredentialStageExchanging { + return s.GetCredentialStage(ctx, stageID) + } + if row.Status != models.CredentialStagePendingAuthorization { + return CredentialStageResult{}, app_errors.ErrAuthorizationStateInvalid + } + now := s.now().UTC() + if now.UnixMilli() >= row.ExpiresAtMS { + if err := s.expireCredentialStage(ctx, &row); err != nil { + return CredentialStageResult{}, app_errors.ParseDBError(err) + } + return s.GetCredentialStage(ctx, stageID) + } + _, authorization := decodeCredentialStageSafeSummary(row.SafeSummaryJSON) + if authorization == nil { + return CredentialStageResult{}, app_errors.ErrAuthorizationStateInvalid + } + if now.UnixMilli() < authorization.NextPollAtMS { + return s.GetCredentialStage(ctx, stageID) + } + claim := s.db.WithContext(ctx).Model(&models.CredentialStage{}). + Where( + "id = ? AND authorization_method = ? AND status = ? AND expires_at_ms > ?", + row.ID, string(channel.AuthorizationDeviceOAuth), models.CredentialStagePendingAuthorization, now.UnixMilli(), + ). + Updates(map[string]any{ + "status": models.CredentialStageExchanging, + "expires_at_ms": now.Add(defaultSubscriptionControlTimeout).UnixMilli(), + "updated_at_ms": now.UnixMilli(), + }) + if claim.Error != nil { + return CredentialStageResult{}, app_errors.ParseDBError(claim.Error) + } + if claim.RowsAffected != 1 { + return s.GetCredentialStage(ctx, stageID) + } + plaintext, err := s.encryption.Decrypt(row.EncryptedPayload) + if err != nil { + _ = s.markCredentialStageOutcomeUnknown(ctx, row.ID) + return CredentialStageResult{}, app_errors.ErrAuthorizationExchangeFailed + } + payload, err := decodeStagedAuthorizationPayload(row.PayloadSchemaVersion, []byte(plaintext)) + if err != nil || len(payload.DriverState) == 0 { + _ = s.markCredentialStageOutcomeUnknown(ctx, row.ID) + return CredentialStageResult{}, app_errors.ErrAuthorizationExchangeFailed + } + pollContext, cancelPoll := context.WithTimeout(context.WithoutCancel(ctx), defaultSubscriptionControlTimeout) + defer cancelPoll() + poll, err := s.pollDeviceAuthorization(pollContext, channel.ID(row.ChannelID), payload.DriverState) + if err != nil { + if markErr := s.markCredentialStageOutcomeUnknown(ctx, row.ID); markErr != nil { + return CredentialStageResult{}, app_errors.ErrInternalServer + } + return CredentialStageResult{}, app_errors.ErrAuthorizationExchangeFailed + } + switch poll.Status { + case subscriptionruntime.DeviceAuthorizationPending: + return s.finishPendingDeviceAuthorization(ctx, row, payload.DriverState, authorization, poll) + case subscriptionruntime.DeviceAuthorizationAuthorized: + finalizeContext, cancelFinalize := credentialStageFinalizeContext(ctx) + defer cancelFinalize() + return s.finishCredentialStageExchange(finalizeContext, row, poll.Credential) + case subscriptionruntime.DeviceAuthorizationDenied: + if err := s.finishDeviceAuthorizationFailure(ctx, row.ID, models.CredentialStageFailed, "authorization_denied"); err != nil { + return CredentialStageResult{}, err + } + return s.GetCredentialStage(ctx, row.ID) + case subscriptionruntime.DeviceAuthorizationExpired: + if err := s.finishDeviceAuthorizationFailure(ctx, row.ID, models.CredentialStageExpired, "authorization_expired"); err != nil { + return CredentialStageResult{}, err + } + return s.GetCredentialStage(ctx, row.ID) + default: + if err := s.markCredentialStageOutcomeUnknown(ctx, row.ID); err != nil { + return CredentialStageResult{}, app_errors.ErrInternalServer + } + return CredentialStageResult{}, app_errors.ErrAuthorizationExchangeFailed + } +} + +func (s *Service) finishPendingDeviceAuthorization( + ctx context.Context, + row models.CredentialStage, + previousDriverState []byte, + authorization *credentialStageAuthorizationSummary, + poll subscriptionruntime.DeviceAuthorizationPoll, +) (CredentialStageResult, error) { + driverState := poll.DriverState + if len(driverState) == 0 { + driverState = previousDriverState + } + if len(driverState) == 0 || len(driverState) > maxOAuthFileBytes { + _ = s.markCredentialStageOutcomeUnknown(ctx, row.ID) + return CredentialStageResult{}, app_errors.ErrAuthorizationExchangeFailed + } + pollInterval := poll.PollInterval + if pollInterval == 0 { + pollInterval = time.Duration(authorization.PollIntervalMS) * time.Millisecond + } + var intervalOK bool + pollInterval, intervalOK = normalizeDevicePollInterval(pollInterval) + if !intervalOK { + _ = s.markCredentialStageOutcomeUnknown(ctx, row.ID) + return CredentialStageResult{}, app_errors.ErrAuthorizationExchangeFailed + } + payload, err := json.Marshal(stagedSubscriptionPayload{DriverState: driverState}) + if err != nil { + _ = s.markCredentialStageOutcomeUnknown(ctx, row.ID) + return CredentialStageResult{}, app_errors.ErrInternalServer + } + ciphertext, err := s.encryption.Encrypt(string(payload)) + clear(payload) + if err != nil { + _ = s.markCredentialStageOutcomeUnknown(ctx, row.ID) + return CredentialStageResult{}, app_errors.ErrInternalServer + } + now := s.now().UTC() + authorization.NextPollAtMS = now.Add(pollInterval).UnixMilli() + authorization.PollIntervalMS = pollInterval.Milliseconds() + summaryJSON, err := json.Marshal(credentialStageSafeSummary{Authorization: authorization}) + if err != nil { + _ = s.markCredentialStageOutcomeUnknown(ctx, row.ID) + return CredentialStageResult{}, app_errors.ErrInternalServer + } + finalizeContext, cancelFinalize := credentialStageFinalizeContext(ctx) + defer cancelFinalize() + result := s.db.WithContext(finalizeContext).Model(&models.CredentialStage{}). + Where("id = ? AND status = ?", row.ID, models.CredentialStageExchanging). + Updates(map[string]any{ + "status": models.CredentialStagePendingAuthorization, "encrypted_payload": ciphertext, + "safe_summary_json": models.JSON(summaryJSON), "expires_at_ms": row.ExpiresAtMS, + "updated_at_ms": now.UnixMilli(), + }) + if result.Error != nil { + return CredentialStageResult{}, app_errors.ParseDBError(result.Error) + } + if result.RowsAffected != 1 { + return CredentialStageResult{}, app_errors.ErrCredentialAuthOutcomeUnknown + } + return s.GetCredentialStage(ctx, row.ID) +} + +func (s *Service) finishDeviceAuthorizationFailure( + ctx context.Context, + stageID string, + status models.CredentialStageStatus, + errorCode string, +) error { + finalizeContext, cancelFinalize := credentialStageFinalizeContext(ctx) + defer cancelFinalize() + result := s.db.WithContext(finalizeContext).Model(&models.CredentialStage{}). + Where("id = ? AND status = ?", stageID, models.CredentialStageExchanging). + Updates(map[string]any{ + "status": status, "encrypted_payload": "", "error_code": errorCode, + "updated_at_ms": s.now().UnixMilli(), + }) + if result.Error != nil { + return app_errors.ParseDBError(result.Error) + } + if result.RowsAffected != 1 { + return app_errors.ErrCredentialAuthOutcomeUnknown + } + return nil +} + func (s *Service) GetCredentialStage(ctx context.Context, stageID string) (CredentialStageResult, error) { row, err := s.loadCredentialStage(ctx, stageID) if err != nil { @@ -667,14 +966,17 @@ func (s *Service) GetCredentialStage(ctx context.Context, stageID string) (Crede return CredentialStageResult{}, app_errors.ParseDBError(err) } } - var account CredentialStageAccount - if len(row.SafeSummaryJSON) > 0 { - _ = json.Unmarshal(row.SafeSummaryJSON, &account) - } + account, authorization := decodeCredentialStageSafeSummary(row.SafeSummaryJSON) result := CredentialStageResult{ - StageID: row.ID, Status: string(row.Status), Account: account, ExpiresAtMS: row.ExpiresAtMS, + StageID: row.ID, Status: string(row.Status), AuthorizationMethod: row.AuthorizationMethod, + Account: account, ExpiresAtMS: row.ExpiresAtMS, ErrorCode: stageResultErrorCode(row.Status, row.ErrorCode), } + if authorization != nil { + result.AuthorizationURL = authorization.AuthorizationURL + result.UserCode = authorization.UserCode + result.NextPollAtMS = authorization.NextPollAtMS + } if row.Status == models.CredentialStagePendingAuthorization || row.Status == models.CredentialStageExchanging { if callback, ok := s.subscriptions.LocalCallback(channel.ID(row.ChannelID)); ok { result.RedirectURI = callback.RedirectURI @@ -683,6 +985,19 @@ func (s *Service) GetCredentialStage(ctx context.Context, stageID string) (Crede return result, nil } +func decodeCredentialStageSafeSummary(raw []byte) (CredentialStageAccount, *credentialStageAuthorizationSummary) { + if len(raw) == 0 { + return CredentialStageAccount{}, nil + } + var safe credentialStageSafeSummary + if json.Unmarshal(raw, &safe) == nil && safe.Authorization != nil { + return CredentialStageAccount{}, safe.Authorization + } + var account CredentialStageAccount + _ = json.Unmarshal(raw, &account) + return account, nil +} + func (s *Service) CancelCredentialStage(ctx context.Context, stageID string) error { if strings.TrimSpace(stageID) == "" { return app_errors.ErrBadRequest @@ -741,7 +1056,8 @@ func (s *Service) persistReadyCredentialStage( return CredentialStageResult{}, app_errors.ParseDBError(err) } return CredentialStageResult{ - StageID: row.ID, Status: string(row.Status), Account: summary, ExpiresAtMS: row.ExpiresAtMS, + StageID: row.ID, Status: string(row.Status), AuthorizationMethod: row.AuthorizationMethod, + Account: summary, ExpiresAtMS: row.ExpiresAtMS, }, nil } @@ -782,7 +1098,8 @@ func (s *Service) finishCredentialStageExchange( return CredentialStageResult{}, app_errors.ErrAuthorizationStateInvalid } return CredentialStageResult{ - StageID: row.ID, Status: string(models.CredentialStageReady), Account: summary, ExpiresAtMS: expiresAtMS, + StageID: row.ID, Status: string(models.CredentialStageReady), AuthorizationMethod: row.AuthorizationMethod, + Account: summary, ExpiresAtMS: expiresAtMS, }, nil } diff --git a/internal/control/credential_stages_test.go b/internal/control/credential_stages_test.go index 78c26ad3..3a8c5e1f 100644 --- a/internal/control/credential_stages_test.go +++ b/internal/control/credential_stages_test.go @@ -7,6 +7,7 @@ import ( "fmt" "net/http" "strings" + "sync/atomic" "testing" "time" @@ -14,6 +15,7 @@ import ( app_errors "gpt-load/internal/platform/errors" "gpt-load/internal/storage/models" "gpt-load/internal/subscription/providers/codex" + subscriptionruntime "gpt-load/internal/subscription/runtime" ) type credentialImportStatusError struct{ status int } @@ -182,6 +184,253 @@ func TestBeginBrowserAuthorizationCreatesPendingStage(t *testing.T) { } } +func TestBeginDeviceAuthorizationCreatesEncryptedPendingStage(t *testing.T) { + fixture := newServiceFixture(t) + now := time.UnixMilli(1_800_000_000_000).UTC() + fixture.service.now = func() time.Time { return now } + fixture.service.beginDeviceAuthorization = func(ctx context.Context, channelID channel.ID) (subscriptionruntime.DeviceAuthorization, error) { + if ctx.Err() != nil || channelID != channel.Grok { + t.Fatalf("begin device input = %v, %q", ctx.Err(), channelID) + } + return subscriptionruntime.DeviceAuthorization{ + VerificationURL: "https://auth.x.ai/device", UserCode: "ABCD-EFGH", + DriverState: []byte(`{"device_code":"device-secret"}`), + ExpiresAt: now.Add(20 * time.Minute), PollInterval: 5 * time.Second, + }, nil + } + + result, err := fixture.service.BeginCredentialAuthorization(t.Context(), channel.Grok) + if err != nil { + t.Fatalf("BeginCredentialAuthorization() error = %v", err) + } + if result.Status != string(models.CredentialStagePendingAuthorization) || + result.AuthorizationMethod != string(channel.AuthorizationDeviceOAuth) || + result.AuthorizationURL != "https://auth.x.ai/device" || result.UserCode != "ABCD-EFGH" || + result.NextPollAtMS != now.Add(5*time.Second).UnixMilli() || + result.ExpiresAtMS != now.Add(20*time.Minute).UnixMilli() { + t.Fatalf("device stage result = %#v", result) + } + var row models.CredentialStage + if err := fixture.db.Take(&row, "id = ?", result.StageID).Error; err != nil { + t.Fatal(err) + } + if row.AuthorizationMethod != string(channel.AuthorizationDeviceOAuth) || row.OAuthStateHash != nil || + row.EncryptedPayload == "" || strings.Contains(row.EncryptedPayload, "device-secret") || + strings.Contains(string(row.SafeSummaryJSON), "device-secret") { + t.Fatalf("stored device stage = %#v", row) + } +} + +func TestBeginDeviceAuthorizationRejectsUnsafeChallenge(t *testing.T) { + now := time.UnixMilli(1_800_000_000_000).UTC() + for _, test := range []struct { + name string + verification string + userCode string + driverState []byte + }{ + {name: "control character", verification: "https://auth.x.ai/device", userCode: "ABCD\nEFGH", driverState: []byte(`{"device_code":"secret"}`)}, + {name: "oversized URL", verification: "https://auth.x.ai/" + strings.Repeat("a", 4096), userCode: "ABCD-EFGH", driverState: []byte(`{"device_code":"secret"}`)}, + {name: "oversized state", verification: "https://auth.x.ai/device", userCode: "ABCD-EFGH", driverState: make([]byte, maxOAuthFileBytes+1)}, + } { + t.Run(test.name, func(t *testing.T) { + fixture := newServiceFixture(t) + fixture.service.now = func() time.Time { return now } + fixture.service.beginDeviceAuthorization = func(context.Context, channel.ID) (subscriptionruntime.DeviceAuthorization, error) { + return subscriptionruntime.DeviceAuthorization{ + VerificationURL: test.verification, + UserCode: test.userCode, + DriverState: test.driverState, + ExpiresAt: now.Add(20 * time.Minute), + PollInterval: 5 * time.Second, + }, nil + } + if _, err := fixture.service.BeginCredentialAuthorization(t.Context(), channel.Grok); !errors.Is(err, app_errors.ErrAuthorizationUnavailable) { + t.Fatalf("BeginCredentialAuthorization() error = %v", err) + } + }) + } +} + +func TestPollDeviceAuthorizationHonorsIntervalAndPersistsPendingState(t *testing.T) { + fixture := newServiceFixture(t) + now := time.UnixMilli(1_800_000_000_000).UTC() + fixture.service.now = func() time.Time { return now } + fixture.service.beginDeviceAuthorization = func(context.Context, channel.ID) (subscriptionruntime.DeviceAuthorization, error) { + return subscriptionruntime.DeviceAuthorization{ + VerificationURL: "https://auth.x.ai/device", UserCode: "ABCD-EFGH", + DriverState: []byte(`{"device_code":"initial"}`), ExpiresAt: now.Add(20 * time.Minute), + PollInterval: 5 * time.Second, + }, nil + } + started, err := fixture.service.BeginCredentialAuthorization(t.Context(), channel.Grok) + if err != nil { + t.Fatal(err) + } + pollCalls := 0 + fixture.service.pollDeviceAuthorization = func(_ context.Context, channelID channel.ID, state []byte) (subscriptionruntime.DeviceAuthorizationPoll, error) { + pollCalls++ + var claimed models.CredentialStage + if err := fixture.db.Take(&claimed, "id = ?", started.StageID).Error; err != nil { + t.Fatal(err) + } + if claimed.Status != models.CredentialStageExchanging || + claimed.ExpiresAtMS != now.Add(defaultSubscriptionControlTimeout).UnixMilli() { + t.Fatalf("claimed device stage = %#v", claimed) + } + if channelID != channel.Grok || !strings.Contains(string(state), "initial") { + t.Fatalf("poll input = %q, %s", channelID, state) + } + return subscriptionruntime.DeviceAuthorizationPoll{ + Status: subscriptionruntime.DeviceAuthorizationPending, DriverState: []byte(`{"device_code":"next"}`), + PollInterval: 10 * time.Second, + }, nil + } + + if _, err := fixture.service.PollCredentialDeviceAuthorization(t.Context(), started.StageID); err != nil { + t.Fatal(err) + } + if pollCalls != 0 { + t.Fatalf("early poll calls = %d", pollCalls) + } + now = now.Add(5 * time.Second) + pending, err := fixture.service.PollCredentialDeviceAuthorization(t.Context(), started.StageID) + if err != nil { + t.Fatal(err) + } + if pending.Status != string(models.CredentialStagePendingAuthorization) || + pending.NextPollAtMS != now.Add(10*time.Second).UnixMilli() || pollCalls != 1 { + t.Fatalf("pending poll = %#v, calls = %d", pending, pollCalls) + } + var restored models.CredentialStage + if err := fixture.db.Take(&restored, "id = ?", started.StageID).Error; err != nil { + t.Fatal(err) + } + if restored.ExpiresAtMS != started.ExpiresAtMS { + t.Fatalf("pending device expiry = %d, want %d", restored.ExpiresAtMS, started.ExpiresAtMS) + } +} + +func TestPollDeviceAuthorizationConcurrentRequestDoesNotDispatchTwice(t *testing.T) { + fixture := newServiceFixture(t) + now := time.UnixMilli(1_800_000_000_000).UTC() + fixture.service.now = func() time.Time { return now } + fixture.service.beginDeviceAuthorization = func(context.Context, channel.ID) (subscriptionruntime.DeviceAuthorization, error) { + return subscriptionruntime.DeviceAuthorization{ + VerificationURL: "https://auth.x.ai/device", UserCode: "ABCD-EFGH", + DriverState: []byte(`{"device_code":"initial"}`), ExpiresAt: now.Add(20 * time.Minute), + PollInterval: time.Second, + }, nil + } + started, err := fixture.service.BeginCredentialAuthorization(t.Context(), channel.Grok) + if err != nil { + t.Fatal(err) + } + now = now.Add(time.Second) + entered := make(chan struct{}) + release := make(chan struct{}) + var calls atomic.Int32 + fixture.service.pollDeviceAuthorization = func(context.Context, channel.ID, []byte) (subscriptionruntime.DeviceAuthorizationPoll, error) { + calls.Add(1) + close(entered) + <-release + return subscriptionruntime.DeviceAuthorizationPoll{ + Status: subscriptionruntime.DeviceAuthorizationPending, PollInterval: 5 * time.Second, + }, nil + } + firstDone := make(chan error, 1) + go func() { + _, pollErr := fixture.service.PollCredentialDeviceAuthorization(t.Context(), started.StageID) + firstDone <- pollErr + }() + <-entered + concurrent, err := fixture.service.PollCredentialDeviceAuthorization(t.Context(), started.StageID) + if err != nil || concurrent.Status != string(models.CredentialStageExchanging) || calls.Load() != 1 { + t.Fatalf("concurrent poll = %#v, %v, calls=%d", concurrent, err, calls.Load()) + } + close(release) + if err := <-firstDone; err != nil { + t.Fatal(err) + } +} + +func TestPollDeviceAuthorizationCompletesReadyStage(t *testing.T) { + fixture := newServiceFixture(t) + now := time.UnixMilli(1_800_000_000_000).UTC() + fixture.service.now = func() time.Time { return now } + fixture.service.beginDeviceAuthorization = func(context.Context, channel.ID) (subscriptionruntime.DeviceAuthorization, error) { + return subscriptionruntime.DeviceAuthorization{ + VerificationURL: "https://auth.x.ai/device", UserCode: "ABCD-EFGH", + DriverState: []byte(`{"device_code":"initial"}`), ExpiresAt: now.Add(20 * time.Minute), + PollInterval: time.Second, + }, nil + } + started, err := fixture.service.BeginCredentialAuthorization(t.Context(), channel.Grok) + if err != nil { + t.Fatal(err) + } + fixture.service.pollDeviceAuthorization = func(context.Context, channel.ID, []byte) (subscriptionruntime.DeviceAuthorizationPoll, error) { + credential := subscriptionruntime.NewCredential( + []byte(`{"type":"grok","access_token":"access","refresh_token":"refresh","account_id":"account"}`), + "account", subscriptionruntime.Account{Email: "owner@example.com"}, now.Add(time.Hour), true, + []string{"access", "refresh"}, + ) + return subscriptionruntime.DeviceAuthorizationPoll{Status: subscriptionruntime.DeviceAuthorizationAuthorized, Credential: credential}, nil + } + now = now.Add(time.Second) + ready, err := fixture.service.PollCredentialDeviceAuthorization(t.Context(), started.StageID) + if err != nil { + t.Fatal(err) + } + if ready.Status != string(models.CredentialStageReady) || ready.Account.EmailMask == "" { + t.Fatalf("ready poll = %#v", ready) + } +} + +func TestPollDeviceAuthorizationPersistsTerminalStatesAndClearsSecrets(t *testing.T) { + for _, test := range []struct { + name string + poll subscriptionruntime.DeviceAuthorizationStatus + wantState models.CredentialStageStatus + wantCode string + }{ + {name: "denied", poll: subscriptionruntime.DeviceAuthorizationDenied, wantState: models.CredentialStageFailed, wantCode: "authorization_denied"}, + {name: "expired", poll: subscriptionruntime.DeviceAuthorizationExpired, wantState: models.CredentialStageExpired, wantCode: "authorization_expired"}, + } { + t.Run(test.name, func(t *testing.T) { + fixture := newServiceFixture(t) + now := time.UnixMilli(1_800_000_000_000).UTC() + fixture.service.now = func() time.Time { return now } + fixture.service.beginDeviceAuthorization = func(context.Context, channel.ID) (subscriptionruntime.DeviceAuthorization, error) { + return subscriptionruntime.DeviceAuthorization{ + VerificationURL: "https://auth.x.ai/device", UserCode: "ABCD-EFGH", + DriverState: []byte(`{"device_code":"secret"}`), ExpiresAt: now.Add(20 * time.Minute), + PollInterval: time.Second, + }, nil + } + started, err := fixture.service.BeginCredentialAuthorization(t.Context(), channel.Grok) + if err != nil { + t.Fatal(err) + } + fixture.service.pollDeviceAuthorization = func(context.Context, channel.ID, []byte) (subscriptionruntime.DeviceAuthorizationPoll, error) { + return subscriptionruntime.DeviceAuthorizationPoll{Status: test.poll}, nil + } + now = now.Add(time.Second) + result, err := fixture.service.PollCredentialDeviceAuthorization(t.Context(), started.StageID) + if err != nil || result.Status != string(test.wantState) { + t.Fatalf("poll result = %#v, %v", result, err) + } + var row models.CredentialStage + if err := fixture.db.Take(&row, "id = ?", started.StageID).Error; err != nil { + t.Fatal(err) + } + if row.Status != test.wantState || row.ErrorCode != test.wantCode || row.EncryptedPayload != "" { + t.Fatalf("terminal device stage = %#v", row) + } + }) + } +} + func TestCancelCredentialStageClearsPayload(t *testing.T) { t.Parallel() diff --git a/internal/control/http_routes.go b/internal/control/http_routes.go index 35b45cb1..1a1c5712 100644 --- a/internal/control/http_routes.go +++ b/internal/control/http_routes.go @@ -65,6 +65,13 @@ func (s *Server) HTTPModule() httproute.Module { s.auditMutation(newMutationDescriptor("credential_stage_oauth_callback", "credential_stage", credentialStageMutationLocator)), s.handleCredentialAuthorizationCallback, ), + controlRoute( + "control.credential-stages.device-poll", + http.MethodPost, + "/credential-stages/:stage_id/device-poll", + s.auditMutation(newMutationDescriptor("credential_stage_device_poll", "credential_stage", credentialStageMutationLocator)), + s.handlePollCredentialDeviceAuthorization, + ), controlRoute( "control.credential-stages.cancel", http.MethodDelete, diff --git a/internal/control/price_reference_test.go b/internal/control/price_reference_test.go index d988bc99..b2784e30 100644 --- a/internal/control/price_reference_test.go +++ b/internal/control/price_reference_test.go @@ -127,6 +127,23 @@ func TestResolveAutomaticPriceForIdentityTreatsAntigravityGooglePriceAsReference } } +func TestResolveAutomaticPriceForIdentityTreatsGrokXAIPriceAsReference(t *testing.T) { + xaiCost := &catalog.ModelCost{Prices: pricing.Prices{Input: priceTestValue(1)}} + snapshot := &catalog.Snapshot{Providers: map[string]catalog.Provider{ + "xai": {ID: "xai", Models: map[string]catalog.Model{ + "grok-4.3": {ID: "grok-4.3", Cost: xaiCost}, + }}, + }} + + match, ok := resolveAutomaticPriceForIdentity(snapshot, pricing.Identity{ + ChannelID: string(channel.Grok), ModelID: "grok-4.3", + }) + if !ok || match.cost != xaiCost || match.providerID != "xai" || + match.source != ModelPriceMatchSourceProviderPriorityFallback { + t.Fatalf("Grok 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/server.go b/internal/control/server.go index 0c200181..69c7ce67 100644 --- a/internal/control/server.go +++ b/internal/control/server.go @@ -127,6 +127,19 @@ func (s *Server) handleCredentialAuthorizationCallback(c *gin.Context) { response.SuccessI18n(c, "common.success", result) } +func (s *Server) handlePollCredentialDeviceAuthorization(c *gin.Context) { + result, err := s.service.PollCredentialDeviceAuthorization( + c.Request.Context(), strings.TrimSpace(c.Param("stage_id")), + ) + if err != nil { + writeServiceError(c, "poll_credential_device_authorization", err) + return + } + setMutationResourceLocator(c, "credential-stage:"+result.StageID) + setSecretResponseHeaders(c) + response.SuccessI18n(c, "common.success", result) +} + func (s *Server) handleImportCredentialStage(c *gin.Context) { reader, err := c.Request.MultipartReader() if err != nil { diff --git a/internal/control/service.go b/internal/control/service.go index 23c910a6..27548e3d 100644 --- a/internal/control/service.go +++ b/internal/control/service.go @@ -61,6 +61,8 @@ type Service struct { operationRandom io.Reader beginSubscriptionAuthorization func(channel.ID) (subscriptionruntime.Authorization, error) completeSubscriptionAuthorization func(context.Context, channel.ID, subscriptionruntime.AuthorizationCompletion) (subscriptionruntime.Credential, error) + beginDeviceAuthorization func(context.Context, channel.ID) (subscriptionruntime.DeviceAuthorization, error) + pollDeviceAuthorization func(context.Context, channel.ID, []byte) (subscriptionruntime.DeviceAuthorizationPoll, error) refreshSubscriptionCredential func(context.Context, channel.ID, subscriptionruntime.Credential) (subscriptionruntime.Credential, error) prepareSubscriptionCredential func(context.Context, channel.ID, execution.CredentialSnapshot, bool) (subscriptionruntime.Credential, *execution.ErrorEvidence) discoverSubscriptionModels func(context.Context, channel.ID, subscriptionruntime.Credential) ([]string, error) @@ -196,6 +198,20 @@ func NewService( } return browser.CompleteAuthorization(ctx, completion) }, + beginDeviceAuthorization: func(ctx context.Context, channelID channel.ID) (subscriptionruntime.DeviceAuthorization, error) { + device, ok := subscriptionsDevice(subscriptions, channelID) + if !ok { + return subscriptionruntime.DeviceAuthorization{}, app_errors.ErrAuthorizationUnavailable + } + return device.BeginDeviceAuthorization(ctx) + }, + pollDeviceAuthorization: func(ctx context.Context, channelID channel.ID, state []byte) (subscriptionruntime.DeviceAuthorizationPoll, error) { + device, ok := subscriptionsDevice(subscriptions, channelID) + if !ok { + return subscriptionruntime.DeviceAuthorizationPoll{}, app_errors.ErrAuthorizationUnavailable + } + return device.PollDeviceAuthorization(ctx, state) + }, refreshSubscriptionCredential: func(ctx context.Context, channelID channel.ID, credential subscriptionruntime.Credential) (subscriptionruntime.Credential, error) { driver, ok := subscriptionsDriver(subscriptions, channelID) if !ok { @@ -274,6 +290,13 @@ func subscriptionsBrowser(runtime *subscriptionruntime.Runtime, channelID channe return runtime.BrowserAuthorization(channelID) } +func subscriptionsDevice(runtime *subscriptionruntime.Runtime, channelID channel.ID) (subscriptionruntime.DeviceAuthorizationDriver, bool) { + if runtime == nil { + return nil, false + } + return runtime.DeviceAuthorization(channelID) +} + type configMutationPublication struct { ConfigInput state.CompileInput PriceTable *pricing.Table diff --git a/internal/execution/cpa/adapter.go b/internal/execution/cpa/adapter.go index 6962acc8..31029572 100644 --- a/internal/execution/cpa/adapter.go +++ b/internal/execution/cpa/adapter.go @@ -58,6 +58,7 @@ func NewAdapter(credentials *subscription.CredentialManager, channels *channel.R newCodexProviderBridge(), newClaudeProviderBridge(), newAntigravityProviderBridge(), + newGrokProviderBridge(), ), } } diff --git a/internal/execution/cpa/grok_provider.go b/internal/execution/cpa/grok_provider.go new file mode 100644 index 00000000..cefec137 --- /dev/null +++ b/internal/execution/cpa/grok_provider.go @@ -0,0 +1,265 @@ +package cpa + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net" + "net/http" + "strings" + "time" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" + "gpt-load/internal/subscription/providers/grok" +) + +type grokProviderCredential struct{ value grok.Credential } + +func (credential grokProviderCredential) redactionValues() []string { + return credential.value.SecretValues() +} + +type grokProviderBridge struct{ executor grok.Executor } + +func newGrokProviderBridge() *grokProviderBridge { + return &grokProviderBridge{executor: grok.NewExecutor()} +} + +func (*grokProviderBridge) ProviderKind() channel.ProviderKind { return channel.ProviderGrok } +func (*grokProviderBridge) UpstreamProtocol() protocol.Protocol { return protocol.OpenAIResponses } + +func (*grokProviderBridge) ValidateRouteCapability(route channel.RouteDescriptor) error { + valid := route.ClientProtocol == protocol.OpenAIResponses && + (route.Operation == execution.OperationResponsesCreate || + route.Operation == execution.OperationResponsesInputTokens) && + route.RouteMode == execution.RouteNative + if route.ClientProtocol == protocol.OpenAICompletions { + valid = route.Operation == execution.OperationChatCompletion && route.RouteMode == execution.RouteConverted + } + if route.ClientProtocol == protocol.Anthropic || route.ClientProtocol == protocol.Gemini { + valid = (route.Operation == execution.OperationChatCompletion || route.Operation == execution.OperationCountTokens) && + route.RouteMode == execution.RouteConverted + } + if !valid { + return fmt.Errorf("route is not implemented by Grok") + } + return nil +} + +func (*grokProviderBridge) ParseCredential(raw []byte) (providerCredential, error) { + credential, err := grok.ParseCredentialJSON(raw) + if err != nil { + return nil, err + } + return grokProviderCredential{value: credential}, nil +} + +func (bridge *grokProviderBridge) Execute( + ctx context.Context, + credentialID string, + credential providerCredential, + request providerRequest, +) (providerResponse, error) { + value, ok := credential.(grokProviderCredential) + if !ok || bridge == nil || bridge.executor == nil { + return providerResponse{}, errors.New("Grok provider bridge credential mismatch") + } + response, err := bridge.executor.Execute(ctx, credentialID, value.value, grokRequest(request, credentialID)) + return providerResponse{ + Payload: append([]byte(nil), response.Payload...), Headers: response.Headers.Clone(), + AppliedReasoningEffort: response.AppliedReasoningEffort, + }, err +} + +func (bridge *grokProviderBridge) ExecuteStream( + ctx context.Context, + credentialID string, + credential providerCredential, + request providerRequest, +) (*providerStreamResponse, error) { + value, ok := credential.(grokProviderCredential) + if !ok || bridge == nil || bridge.executor == nil { + return nil, errors.New("Grok provider bridge credential mismatch") + } + response, err := bridge.executor.ExecuteStream(ctx, credentialID, value.value, grokRequest(request, credentialID)) + 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} + select { + case chunks <- converted: + case <-ctx.Done(): + return + } + } + }() + return &providerStreamResponse{ + Headers: response.Headers.Clone(), Chunks: chunks, + AppliedReasoningEffort: response.AppliedReasoningEffort, + }, nil +} + +func (bridge *grokProviderBridge) ValidateLocalTokenCount(request providerRequest) error { + if !grokLocalTokenCountModelSupported(request.Model) { + return errors.New("Grok local token count tokenizer is unavailable for model") + } + var root map[string]any + if err := json.Unmarshal(request.Payload, &root); err != nil || root == nil { + return errors.New("Grok local token count request is invalid") + } + switch request.Format { + case "openai-response": + return validateLocalResponsesTokenCount(root) + case "claude": + return validateLocalClaudeTokenCount(root) + case "gemini": + return validateLocalGeminiTokenCount(root) + default: + return errors.New("Grok local token count format is unsupported") + } +} + +func grokLocalTokenCountModelSupported(model string) bool { + return strings.HasPrefix(strings.ToLower(strings.TrimSpace(model)), "grok-") +} + +func (bridge *grokProviderBridge) CountTokensLocal( + ctx context.Context, + request providerRequest, +) (providerResponse, error) { + if bridge == nil || bridge.executor == nil { + return providerResponse{}, errors.New("Grok provider bridge is unavailable") + } + response, err := bridge.executor.CountTokens(ctx, grok.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...), + }) + headers := response.Headers.Clone() + if headers == nil { + headers = make(http.Header) + } + headers.Set(localTokenCountHeader, "local-estimate") + return providerResponse{ + Payload: append([]byte(nil), response.Payload...), Headers: headers, + AppliedReasoningEffort: response.AppliedReasoningEffort, + }, err +} + +func grokRequest(request providerRequest, credentialID string) grok.ExecuteRequest { + return grok.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: grokContinuityScope(request.ContinuityKey, credentialID, request.Model, request.AttemptID), + } +} + +func grokContinuityScope(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-grok-attempt", attemptID, credentialID, model}, "\x00") + } + return strings.Join([]string{"gpt-load-grok", base, credentialID, model}, "\x00") +} + +func (*grokProviderBridge) 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 + } + code := "" + var coded interface{ ErrorCode() string } + if errors.As(err, &coded) && coded != nil { + code = safeScalar(coded.ErrorCode()) + } + evidence := &execution.ErrorEvidence{ + Kind: kind, StatusCode: status, Code: code, + Summary: grokProviderErrorSummary(status, code), + } + 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: + evidence.Hint = execution.FailureHintCandidateUnavailable + evidence.ReplaySafety = execution.ReplaySafetyRejectedBeforeProcessing + case status == http.StatusBadRequest: + evidence.Hint = execution.FailureHintRequestRejected + case status == http.StatusTooManyRequests: + evidence.Hint = execution.FailureHintRateLimited + case status >= http.StatusInternalServerError: + evidence.Hint = execution.FailureHintHostError + } + if evidence.Summary == "" { + evidence.Summary = safeErrorSummary(err, credential.redactionValues()) + } + return status, evidence +} + +func grokProviderErrorSummary(status int, code string) string { + switch { + case status == http.StatusUnauthorized: + return "Grok authorization was rejected" + case status == http.StatusForbidden: + return "Grok access was denied" + case status == http.StatusTooManyRequests && strings.Contains(strings.ToLower(code), "free-usage-exhausted"): + return "Grok included free usage is exhausted" + case status == http.StatusTooManyRequests: + return "Grok upstream rate limit was reached" + case status >= http.StatusInternalServerError: + return "Grok upstream service failed" + case status >= http.StatusBadRequest: + return "Grok upstream request was rejected" + default: + return "Grok upstream request failed" + } +} + +var _ providerBridge = (*grokProviderBridge)(nil) +var _ providerLocalTokenCounter = (*grokProviderBridge)(nil) diff --git a/internal/execution/cpa/grok_provider_test.go b/internal/execution/cpa/grok_provider_test.go new file mode 100644 index 00000000..4b1721b1 --- /dev/null +++ b/internal/execution/cpa/grok_provider_test.go @@ -0,0 +1,105 @@ +package cpa + +import ( + "context" + "net/http" + "strings" + "testing" + "time" + + "gpt-load/internal/execution" +) + +type grokProviderTestError struct { + status int + code string + retry time.Duration +} + +func (err grokProviderTestError) Error() string { return "Grok provider test error" } +func (err grokProviderTestError) StatusCode() int { return err.status } +func (err grokProviderTestError) ErrorCode() string { return err.code } +func (err grokProviderTestError) RetryAfter() *time.Duration { + if err.retry <= 0 { + return nil + } + value := err.retry + return &value +} + +func TestGrokProviderClassifiesOAuthAndQuotaFailures(t *testing.T) { + bridge := newGrokProviderBridge() + credential := grokProviderCredential{} + for _, test := range []struct { + name string + err error + hint execution.FailureHint + replay execution.ReplaySafety + }{ + {name: "unauthorized", err: grokProviderTestError{status: http.StatusUnauthorized}, hint: execution.FailureHintRefreshRequired, replay: execution.ReplaySafetyRejectedBeforeProcessing}, + {name: "forbidden", err: grokProviderTestError{status: http.StatusForbidden}, hint: execution.FailureHintCandidateUnavailable, replay: execution.ReplaySafetyRejectedBeforeProcessing}, + {name: "invalid request", err: grokProviderTestError{status: http.StatusBadRequest}, hint: execution.FailureHintRequestRejected}, + {name: "free usage", err: grokProviderTestError{status: http.StatusTooManyRequests, code: "subscription:free-usage-exhausted", retry: 24 * time.Hour}, hint: execution.FailureHintRateLimited}, + {name: "host", err: grokProviderTestError{status: http.StatusServiceUnavailable}, hint: execution.FailureHintHostError}, + } { + t.Run(test.name, func(t *testing.T) { + status, evidence := bridge.ClassifyError(t.Context(), test.err, credential) + if status != test.err.(grokProviderTestError).status || evidence == nil || + evidence.Hint != test.hint || evidence.ReplaySafety != test.replay { + t.Fatalf("classification = %d/%#v", status, evidence) + } + if test.name == "free usage" && evidence.RetryAfter != 24*time.Hour { + t.Fatalf("RetryAfter = %v", evidence.RetryAfter) + } + }) + } +} + +func TestGrokLocalTokenCountRejectsStatefulAndMultimodalInputs(t *testing.T) { + bridge := newGrokProviderBridge() + for _, test := range []struct { + name string + payload string + wantErr bool + }{ + {name: "text", payload: `{"model":"grok-4.3","input":"hello"}`}, + {name: "stateful", payload: `{"model":"grok-4.3","previous_response_id":"resp_1","input":"hello"}`, wantErr: true}, + {name: "image", payload: `{"model":"grok-4.3","input":[{"type":"message","role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}]}`, wantErr: true}, + } { + t.Run(test.name, func(t *testing.T) { + err := bridge.ValidateLocalTokenCount(providerRequest{ + Model: "grok-4.3", Format: "openai-response", Payload: []byte(test.payload), + }) + if (err != nil) != test.wantErr { + t.Fatalf("ValidateLocalTokenCount() error = %v", err) + } + }) + } +} + +func TestGrokLocalTokenCountDoesNotRequireCredential(t *testing.T) { + bridge := newGrokProviderBridge() + response, err := bridge.CountTokensLocal(context.Background(), providerRequest{ + Model: "grok-4.3", Format: "openai-response", + Payload: []byte(`{"model":"grok-4.3","input":"hello"}`), + }) + if err != nil { + t.Fatal(err) + } + if len(response.Payload) == 0 || response.Headers.Get(localTokenCountHeader) != "local-estimate" { + t.Fatalf("response = %#v", response) + } + if !strings.Contains(string(response.Payload), `"object":"response.input_tokens"`) || + !strings.Contains(string(response.Payload), `"input_tokens":`) { + t.Fatalf("Responses token count = %s", response.Payload) + } +} + +func TestGrokContinuityScopeIsCredentialAndModelScoped(t *testing.T) { + left := grokContinuityScope("tenant", "credential-1", "grok-4.3", "attempt") + right := grokContinuityScope("tenant", "credential-2", "grok-4.3", "attempt") + otherModel := grokContinuityScope("tenant", "credential-1", "grok-code-fast-1", "attempt") + if left == "" || left == right || left == otherModel { + t.Fatalf("continuity scopes = %q / %q / %q", left, right, otherModel) + } +} diff --git a/internal/execution/cpa/provider_test.go b/internal/execution/cpa/provider_test.go index da73c046..261f4ff1 100644 --- a/internal/execution/cpa/provider_test.go +++ b/internal/execution/cpa/provider_test.go @@ -64,6 +64,20 @@ func TestAdapterDelegatesRouteValidationByProviderKind(t *testing.T) { } } +func TestAdapterValidatesAllDeclaredGrokRoutes(t *testing.T) { + registry := channel.NewRegistry() + adapter := NewAdapter(nil, registry) + descriptor, ok := registry.Get(channel.Grok) + if !ok { + t.Fatal("Grok channel is missing") + } + for _, route := range descriptor.Routes { + if err := adapter.ValidateRouteCapability(channel.ProviderGrok, route); err != nil { + t.Fatalf("ValidateRouteCapability(%#v) error = %v", route, err) + } + } +} + func TestRequestScopedFailureSurvivesErrorWrapping(t *testing.T) { if !requestScopedFailure(errors.New("unrelated: " + requestScopedTestError{}.Error())) { // A string with the same text is deliberately not sufficient. diff --git a/internal/provideradapter/registry_test.go b/internal/provideradapter/registry_test.go index efe14d29..ca6c3d7f 100644 --- a/internal/provideradapter/registry_test.go +++ b/internal/provideradapter/registry_test.go @@ -129,5 +129,6 @@ func completeBindings(bifrost execution.Executor, codex execution.Executor) []Bi {ProviderKind: channel.ProviderCodex, Adapter: codex}, {ProviderKind: channel.ProviderClaude, Adapter: codex}, {ProviderKind: channel.ProviderAntigravity, Adapter: codex}, + {ProviderKind: channel.ProviderGrok, Adapter: codex}, } } diff --git a/internal/storage/db_test.go b/internal/storage/db_test.go index e9a1c644..a9398517 100644 --- a/internal/storage/db_test.go +++ b/internal/storage/db_test.go @@ -624,6 +624,7 @@ func TestAutoMigrateCreatesUsageJournalAndMigrationLedger(t *testing.T) { "0006_upstream_protocol", "0007_model_price_mode_schedules", "0008_observation_auth_refresh", + "0009_device_oauth", } if !reflect.DeepEqual(migrationIDs, wantMigrationIDs) { t.Fatalf("schema_migrations IDs = %v, want %v", migrationIDs, wantMigrationIDs) diff --git a/internal/storage/migration.go b/internal/storage/migration.go index 0cf227e7..e50818ef 100644 --- a/internal/storage/migration.go +++ b/internal/storage/migration.go @@ -82,6 +82,12 @@ var migrations = []migration{ Validate: migrationfiles.Validate0008, ValidateRecoverable: migrationfiles.ValidateRecoverable0008, }, + { + ID: migrationfiles.ID0009, + Up: migrationfiles.Up0009, + Validate: migrationfiles.Validate0009, + ValidateRecoverable: migrationfiles.ValidateRecoverable0009, + }, } func applyMigrations(db *gorm.DB) error { diff --git a/internal/storage/migration_test.go b/internal/storage/migration_test.go index eba786ce..0ee7e58f 100644 --- a/internal/storage/migration_test.go +++ b/internal/storage/migration_test.go @@ -20,6 +20,7 @@ func TestMigrationRegistryContainsOrderedMigrations(t *testing.T) { migrationfiles.ID0006, migrationfiles.ID0007, migrationfiles.ID0008, + migrationfiles.ID0009, } if len(migrations) != len(wantIDs) { t.Fatalf("migration registry length = %d, want %d", len(migrations), len(wantIDs)) diff --git a/internal/storage/migrations/0009_device_oauth.go b/internal/storage/migrations/0009_device_oauth.go new file mode 100644 index 00000000..8ad0715f --- /dev/null +++ b/internal/storage/migrations/0009_device_oauth.go @@ -0,0 +1,190 @@ +package migrations + +import ( + "database/sql" + "fmt" + "strings" + + "gorm.io/gorm" +) + +const ID0009 = "0009_device_oauth" + +const credentialStageAuthorizationConstraint0009 = "chk_credential_stage_authorization_method" + +type credentialStageAuthorization0009 struct { + AuthorizationMethod string `gorm:"column:authorization_method;check:chk_credential_stage_authorization_method,authorization_method IN ('browser_oauth','device_oauth','oauth_file')"` +} + +func (credentialStageAuthorization0009) TableName() string { return "credential_stages" } + +// Up0009 extends the durable authorization method contract for RFC 8628 device +// OAuth while retaining all existing credential stages and indexes. +func Up0009(db *gorm.DB) error { + if err := validateCredentialStagesTable0009(db); err != nil { + return err + } + if definition, err := credentialStageAuthorizationConstraintDefinition0009(db); err == nil && + strings.Contains(strings.ToLower(definition), "device_oauth") { + return Validate0009(db) + } + if strings.EqualFold(db.Dialector.Name(), "sqlite") { + if err := rebuildCredentialStagesSQLite0009(db); err != nil { + return err + } + return Validate0009(db) + } + model := &credentialStageAuthorization0009{} + if db.Migrator().HasConstraint(model, credentialStageAuthorizationConstraint0009) { + if err := db.Migrator().DropConstraint(model, credentialStageAuthorizationConstraint0009); err != nil { + return fmt.Errorf("drop credential stage authorization constraint: %w", err) + } + } + if err := db.Migrator().CreateConstraint(model, credentialStageAuthorizationConstraint0009); err != nil { + return fmt.Errorf("create credential stage authorization constraint: %w", err) + } + return Validate0009(db) +} + +func rebuildCredentialStagesSQLite0009(db *gorm.DB) error { + const temporaryTable = "credential_stages_0009" + if err := db.Exec(`DROP TABLE IF EXISTS credential_stages_0009`).Error; err != nil { + return fmt.Errorf("drop stale SQLite credential stages table: %w", err) + } + if err := db.Exec(`CREATE TABLE credential_stages_0009 ( + id varchar(36) NOT NULL, + channel_id varchar(64) NOT NULL, + connection_type varchar(32) NOT NULL, + authorization_method varchar(32) NOT NULL, + status varchar(32) NOT NULL, + encrypted_payload text NOT NULL, + payload_schema_version integer NOT NULL DEFAULT 1, + safe_summary_json json NOT NULL, + identity_fingerprint varchar(128) NOT NULL DEFAULT '', + oauth_state_hash varchar(128), + expires_at_ms integer NOT NULL, + consumed_at_ms integer, + consumed_group_id integer, + error_code varchar(64) NOT NULL DEFAULT '', + created_at_ms integer NOT NULL, + updated_at_ms integer NOT NULL, + PRIMARY KEY (id), + CONSTRAINT chk_credential_stage_updated_at CHECK (updated_at_ms >= 0), + CONSTRAINT chk_credential_stage_payload_schema CHECK (payload_schema_version > 0), + CONSTRAINT chk_credential_stage_consumed_at CHECK (consumed_at_ms IS NULL OR consumed_at_ms >= 0), + CONSTRAINT chk_credential_stage_expires_at CHECK (expires_at_ms >= 0), + CONSTRAINT chk_credential_stage_created_at CHECK (created_at_ms >= 0), + CONSTRAINT chk_credential_stage_connection_type CHECK (connection_type = 'subscription'), + CONSTRAINT chk_credential_stage_authorization_method CHECK (authorization_method IN ('browser_oauth','device_oauth','oauth_file')), + CONSTRAINT chk_credential_stage_status CHECK (status IN ('pending_authorization','exchanging','ready','consumed','failed','cancelled','expired','outcome_unknown')) + )`).Error; err != nil { + return fmt.Errorf("create SQLite credential stages replacement: %w", err) + } + columns := `id, channel_id, connection_type, authorization_method, status, + encrypted_payload, payload_schema_version, safe_summary_json, identity_fingerprint, + oauth_state_hash, expires_at_ms, consumed_at_ms, consumed_group_id, error_code, + created_at_ms, updated_at_ms` + if err := db.Exec(fmt.Sprintf( + "INSERT INTO %s (%s) SELECT %s FROM credential_stages", + temporaryTable, columns, columns, + )).Error; err != nil { + return fmt.Errorf("copy SQLite credential stages: %w", err) + } + if err := db.Exec(`DROP TABLE credential_stages`).Error; err != nil { + return fmt.Errorf("drop old SQLite credential stages: %w", err) + } + if err := db.Exec(`ALTER TABLE credential_stages_0009 RENAME TO credential_stages`).Error; err != nil { + return fmt.Errorf("rename SQLite credential stages replacement: %w", err) + } + for _, statement := range []string{ + `CREATE UNIQUE INDEX idx_credential_stages_oauth_state ON credential_stages (oauth_state_hash)`, + `CREATE INDEX idx_credential_stages_status_expires ON credential_stages (status, expires_at_ms, updated_at_ms)`, + } { + if err := db.Exec(statement).Error; err != nil { + return fmt.Errorf("restore SQLite credential stage index: %w", err) + } + } + return nil +} + +func ValidateRecoverable0009(db *gorm.DB) error { + if err := validateCredentialStagesTable0009(db); err != nil { + return err + } + if !db.Migrator().HasConstraint(&credentialStageAuthorization0009{}, credentialStageAuthorizationConstraint0009) { + return nil + } + definition, err := credentialStageAuthorizationConstraintDefinition0009(db) + if err != nil { + return err + } + if strings.Contains(strings.ToLower(definition), "device_oauth") { + return nil + } + // The frozen 0002 constraint is a safe, known predecessor which Up0009 can replace. + lower := strings.ToLower(definition) + if strings.Contains(lower, "browser_oauth") && strings.Contains(lower, "oauth_file") { + return nil + } + return fmt.Errorf("credential stage authorization constraint is not recoverable") +} + +func Validate0009(db *gorm.DB) error { + if err := validateCredentialStagesTable0009(db); err != nil { + return err + } + if !db.Migrator().HasConstraint(&credentialStageAuthorization0009{}, credentialStageAuthorizationConstraint0009) { + return fmt.Errorf("credential stage authorization constraint is missing") + } + definition, err := credentialStageAuthorizationConstraintDefinition0009(db) + if err != nil { + return err + } + lower := strings.ToLower(definition) + for _, method := range []string{"browser_oauth", "device_oauth", "oauth_file"} { + if !strings.Contains(lower, method) { + return fmt.Errorf("credential stage authorization constraint does not allow %q", method) + } + } + return nil +} + +func validateCredentialStagesTable0009(db *gorm.DB) error { + if db == nil || db.Dialector == nil || !db.Migrator().HasTable("credential_stages") { + return fmt.Errorf("credential stages table is missing") + } + return nil +} + +func credentialStageAuthorizationConstraintDefinition0009(db *gorm.DB) (string, error) { + var definition sql.NullString + switch strings.ToLower(db.Dialector.Name()) { + case "sqlite": + if err := db.Raw("SELECT sql FROM sqlite_master WHERE type = 'table' AND name = ?", "credential_stages").Scan(&definition).Error; err != nil { + return "", fmt.Errorf("inspect SQLite credential stage constraint: %w", err) + } + case "mysql": + if err := db.Raw( + "SELECT CHECK_CLAUSE FROM information_schema.CHECK_CONSTRAINTS WHERE CONSTRAINT_SCHEMA = DATABASE() AND CONSTRAINT_NAME = ? LIMIT 1", + credentialStageAuthorizationConstraint0009, + ).Scan(&definition).Error; err != nil { + return "", fmt.Errorf("inspect MySQL credential stage constraint: %w", err) + } + case "postgres", "postgresql": + if err := db.Raw(` +SELECT pg_get_constraintdef(c.oid) +FROM pg_constraint c +JOIN pg_class t ON t.oid = c.conrelid +JOIN pg_namespace n ON n.oid = t.relnamespace +WHERE c.conname = ? AND t.relname = ? AND n.nspname = current_schema() +LIMIT 1`, credentialStageAuthorizationConstraint0009, "credential_stages").Scan(&definition).Error; err != nil { + return "", fmt.Errorf("inspect PostgreSQL credential stage constraint: %w", err) + } + default: + return "", fmt.Errorf("unsupported database driver %q", db.Dialector.Name()) + } + if !definition.Valid || strings.TrimSpace(definition.String) == "" { + return "", fmt.Errorf("credential stage authorization constraint definition is missing") + } + return definition.String, nil +} diff --git a/internal/storage/migrations/0009_device_oauth_test.go b/internal/storage/migrations/0009_device_oauth_test.go new file mode 100644 index 00000000..9ecdd031 --- /dev/null +++ b/internal/storage/migrations/0009_device_oauth_test.go @@ -0,0 +1,44 @@ +package migrations_test + +import ( + "testing" + + "gorm.io/gorm" + + "gpt-load/internal/storage/migrations" +) + +func TestDeviceOAuthMigrationPreservesStagesAndExtendsConstraint(t *testing.T) { + db := openInitialTestDatabase(t) + for _, up := range []func(*gorm.DB) error{migrations.Up0001, migrations.Up0002} { + if err := up(db); err != nil { + t.Fatal(err) + } + } + insertCredentialStage0009(t, db, "existing", "browser_oauth") + if err := migrations.Up0009(db); err != nil { + t.Fatal(err) + } + insertCredentialStage0009(t, db, "device", "device_oauth") + if err := db.Exec(`UPDATE credential_stages SET authorization_method = 'unsupported' WHERE id = 'device'`).Error; err == nil { + t.Fatal("unsupported authorization method was accepted") + } + var count int64 + if err := db.Table("credential_stages").Where("id IN ?", []string{"existing", "device"}).Count(&count).Error; err != nil || count != 2 { + t.Fatalf("preserved stages = %d, %v", count, err) + } + if err := migrations.Up0009(db); err != nil { + t.Fatalf("idempotent Up0009() error = %v", err) + } +} + +func insertCredentialStage0009(t *testing.T, db *gorm.DB, id, method string) { + t.Helper() + if err := db.Exec(`INSERT INTO credential_stages ( + id, channel_id, connection_type, authorization_method, status, + encrypted_payload, payload_schema_version, safe_summary_json, + identity_fingerprint, expires_at_ms, error_code, created_at_ms, updated_at_ms + ) VALUES (?, 'grok', 'subscription', ?, 'pending_authorization', 'encrypted', 2, '{}', '', 100, '', 1, 1)`, id, method).Error; err != nil { + t.Fatalf("insert stage %q/%q: %v", id, method, err) + } +} diff --git a/internal/subscription/providers/grok/driver.go b/internal/subscription/providers/grok/driver.go new file mode 100644 index 00000000..4fb8f7d6 --- /dev/null +++ b/internal/subscription/providers/grok/driver.go @@ -0,0 +1,159 @@ +package grok + +import ( + "context" + "errors" + "strings" + "time" + + "gpt-load/internal/channel/modules" + "gpt-load/internal/channel/spec" + subscriptionruntime "gpt-load/internal/subscription/runtime" +) + +type grokDriver struct{} + +func newGrokDriver() *grokDriver { return &grokDriver{} } + +func Implementations() subscriptionruntime.Implementations { + driver := newGrokDriver() + return subscriptionruntime.Implementations{ + Drivers: []subscriptionruntime.Driver{driver}, + ModelDiscoveries: []subscriptionruntime.ModelDiscovery{grokModelDiscovery{driver}}, + } +} + +func (*grokDriver) ID() spec.SubscriptionDriverID { return modules.GrokSubscriptionDriver } + +func (*grokDriver) Parse(raw []byte) (subscriptionruntime.Credential, error) { + value, err := ParseCredentialJSON(raw) + if err != nil { + return subscriptionruntime.Credential{}, err + } + canonical, err := MarshalCredential(value) + if err != nil { + return subscriptionruntime.Credential{}, err + } + return grokRuntimeCredential(value, canonical), nil +} + +func (*grokDriver) Refresh(ctx context.Context, current subscriptionruntime.Credential) (subscriptionruntime.Credential, error) { + value, err := ParseCredentialJSON(current.Canonical()) + if err != nil { + return subscriptionruntime.Credential{}, err + } + refreshed, err := RefreshCredentialOnce(ctx, value) + if err != nil { + return subscriptionruntime.Credential{}, err + } + canonical, err := MarshalCredential(refreshed) + if err != nil { + return subscriptionruntime.Credential{}, err + } + return grokRuntimeCredential(refreshed, canonical), nil +} + +func (*grokDriver) ClassifyRefreshFailure(err error) subscriptionruntime.RefreshFailure { + if errors.Is(err, ErrCredentialIdentityChanged) { + return subscriptionruntime.RefreshFailureIdentityChanged + } + var tokenErr *TokenEndpointError + if errors.As(err, &tokenErr) && IsDefinitiveRefreshRejection(tokenErr.Code) { + return subscriptionruntime.RefreshFailureReauthorizationRequired + } + return subscriptionruntime.RefreshFailureOutcomeUnknown +} + +func (*grokDriver) BeginDeviceAuthorization(ctx context.Context) (subscriptionruntime.DeviceAuthorization, error) { + value, err := BeginDeviceAuthorization(ctx) + if err != nil { + return subscriptionruntime.DeviceAuthorization{}, err + } + state, err := marshalDeviceState(value.State) + if err != nil { + return subscriptionruntime.DeviceAuthorization{}, err + } + return subscriptionruntime.DeviceAuthorization{ + VerificationURL: value.VerificationURL, UserCode: value.UserCode, DriverState: state, + ExpiresAt: value.ExpiresAt, PollInterval: value.PollInterval, + }, nil +} + +func (*grokDriver) PollDeviceAuthorization(ctx context.Context, raw []byte) (subscriptionruntime.DeviceAuthorizationPoll, error) { + state, err := unmarshalDeviceState(raw) + if err != nil { + return subscriptionruntime.DeviceAuthorizationPoll{}, err + } + value, err := PollDeviceAuthorizationOnce(ctx, state) + if err != nil { + return subscriptionruntime.DeviceAuthorizationPoll{}, err + } + result := subscriptionruntime.DeviceAuthorizationPoll{PollInterval: value.PollInterval} + switch value.Status { + case DeviceAuthorizationPending: + result.Status = subscriptionruntime.DeviceAuthorizationPending + result.DriverState, err = marshalDeviceState(value.State) + case DeviceAuthorizationAuthorized: + result.Status = subscriptionruntime.DeviceAuthorizationAuthorized + canonical, marshalErr := MarshalCredential(value.Credential) + if marshalErr != nil { + return subscriptionruntime.DeviceAuthorizationPoll{}, marshalErr + } + result.Credential = grokRuntimeCredential(value.Credential, canonical) + case DeviceAuthorizationDenied: + result.Status = subscriptionruntime.DeviceAuthorizationDenied + case DeviceAuthorizationExpired: + result.Status = subscriptionruntime.DeviceAuthorizationExpired + default: + return subscriptionruntime.DeviceAuthorizationPoll{}, errors.New("unknown Grok device authorization status") + } + return result, err +} + +func (*grokDriver) ImportCredential(ctx context.Context, raw []byte) (subscriptionruntime.Credential, error) { + value, err := ImportCredential(ctx, raw) + if err != nil { + return subscriptionruntime.Credential{}, err + } + canonical, err := MarshalCredential(value) + if err != nil { + return subscriptionruntime.Credential{}, err + } + return grokRuntimeCredential(value, canonical), nil +} + +type grokModelDiscovery struct{ *grokDriver } + +func (grokModelDiscovery) ID() spec.UtilityID { return modules.GrokModelDiscovery } + +func (*grokDriver) DiscoverModels(ctx context.Context, credential subscriptionruntime.Credential) ([]string, error) { + value, err := ParseCredentialJSON(credential.Canonical()) + if err != nil { + return nil, err + } + models, err := ListModels(ctx, value) + if err != nil { + var upstream *UpstreamHTTPError + if errors.As(err, &upstream) { + return nil, &subscriptionruntime.UpstreamHTTPError{StatusCode: upstream.StatusCode} + } + return nil, err + } + return models, nil +} + +func grokRuntimeCredential(value Credential, canonical []byte) subscriptionruntime.Credential { + expiresAt, expires := CredentialExpiresAt(value) + account := subscriptionruntime.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 subscriptionruntime.NewCredential( + canonical, strings.TrimSpace(value.AccountID), account, expiresAt, expires, value.SecretValues(), + ) +} + +var _ subscriptionruntime.DeviceAuthorizationDriver = (*grokDriver)(nil) +var _ subscriptionruntime.CredentialFileImporter = (*grokDriver)(nil) diff --git a/internal/subscription/providers/grok/driver_test.go b/internal/subscription/providers/grok/driver_test.go new file mode 100644 index 00000000..178ca054 --- /dev/null +++ b/internal/subscription/providers/grok/driver_test.go @@ -0,0 +1,71 @@ +package grok + +import ( + "errors" + "reflect" + "testing" + + "gpt-load/internal/channel/modules" + subscriptionruntime "gpt-load/internal/subscription/runtime" +) + +func TestGrokDriverProducesProviderNeutralCredential(t *testing.T) { + driver := newGrokDriver() + credential, err := driver.Parse([]byte(`{ + "type":"grok","access_token":"access-secret","refresh_token":"refresh-secret", + "id_token":"id-secret","account_id":"account-1","email":"owner@example.com", + "expired":"2030-01-01T00:00:00Z","token_endpoint":"https://auth.x.ai/oauth/token" + }`)) + if err != nil { + t.Fatal(err) + } + if credential.Identity() != "account-1" || credential.Account().Email != "owner@example.com" || + !credential.Account().ExpiresAtKnown { + t.Fatalf("credential = %#v / %#v", credential.Identity(), credential.Account()) + } + if got := credential.SecretValues(); !reflect.DeepEqual(got, []string{ + "access-secret", "refresh-secret", "id-secret", "account-1", "owner@example.com", + }) { + t.Fatalf("SecretValues() = %q", got) + } +} + +func TestGrokDriverDeclaresDeviceOAuthImporterAndModelDiscovery(t *testing.T) { + driver := newGrokDriver() + if _, ok := any(driver).(subscriptionruntime.DeviceAuthorizationDriver); !ok { + t.Fatal("Grok driver does not implement DeviceAuthorizationDriver") + } + if _, ok := any(driver).(subscriptionruntime.CredentialFileImporter); !ok { + t.Fatal("Grok driver does not implement CredentialFileImporter") + } + implementations := Implementations() + if len(implementations.Drivers) != 1 || implementations.Drivers[0].ID() != modules.GrokSubscriptionDriver || + len(implementations.ModelDiscoveries) != 1 || implementations.ModelDiscoveries[0].ID() != modules.GrokModelDiscovery || + len(implementations.QuotaObservations) != 0 || len(implementations.ResetCreditActions) != 0 { + t.Fatalf("implementations = %#v", implementations) + } +} + +func TestGrokDriverClassifiesRefreshFailures(t *testing.T) { + driver := newGrokDriver() + if got := driver.ClassifyRefreshFailure(ErrCredentialIdentityChanged); got != subscriptionruntime.RefreshFailureIdentityChanged { + t.Fatalf("identity classification = %v", got) + } + if got := driver.ClassifyRefreshFailure(&TokenEndpointError{StatusCode: 400, Code: "invalid_grant"}); got != subscriptionruntime.RefreshFailureReauthorizationRequired { + t.Fatalf("invalid_grant classification = %v", got) + } + if got := driver.ClassifyRefreshFailure(errors.New("network unavailable")); got != subscriptionruntime.RefreshFailureOutcomeUnknown { + t.Fatalf("network classification = %v", got) + } +} + +func TestGrokDeviceStateRejectsUnknownOrIncompleteFields(t *testing.T) { + for _, raw := range [][]byte{ + []byte(`{"device_code":"code","token_endpoint":"https://auth.x.ai/token","userinfo_endpoint":"https://auth.x.ai/userinfo","poll_interval_seconds":5,"proxy":"bad"}`), + []byte(`{"device_code":"code","poll_interval_seconds":5}`), + } { + if _, err := unmarshalDeviceState(raw); err == nil { + t.Fatalf("unmarshalDeviceState(%s) succeeded", raw) + } + } +} diff --git a/internal/subscription/providers/grok/executor.go b/internal/subscription/providers/grok/executor.go new file mode 100644 index 00000000..b64bc895 --- /dev/null +++ b/internal/subscription/providers/grok/executor.go @@ -0,0 +1,93 @@ +package grok + +import ( + "context" + "net/http" + + cpaembedded "github.com/router-for-me/CLIProxyAPI/v7/gptload-embedded/embedded" +) + +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 ExecuteStreamChunk struct { + Payload []byte + Err error +} + +type ExecuteStreamResponse struct { + Headers http.Header + Chunks <-chan ExecuteStreamChunk + AppliedReasoningEffort string +} + +type Executor interface { + Execute(context.Context, string, Credential, ExecuteRequest) (ExecuteResponse, error) + CountTokens(context.Context, ExecuteRequest) (ExecuteResponse, error) + ExecuteStream(context.Context, string, Credential, ExecuteRequest) (*ExecuteStreamResponse, error) +} + +type executor struct{ bridge cpaembedded.GrokHTTPExecutor } + +func NewExecutor() Executor { return &executor{bridge: cpaembedded.NewGrokHTTPExecutor()} } + +func (value *executor) Execute(ctx context.Context, credentialID string, credential Credential, request ExecuteRequest) (ExecuteResponse, error) { + response, err := value.bridge.ExecuteCanonical(ctx, credentialID, credentialToBridge(credential), requestToBridge(request)) + return responseFromBridge(response), normalizeError(err) +} + +func (value *executor) CountTokens(ctx context.Context, request ExecuteRequest) (ExecuteResponse, error) { + response, err := value.bridge.CountTokensCanonical(ctx, requestToBridge(request)) + return responseFromBridge(response), normalizeError(err) +} + +func (value *executor) ExecuteStream(ctx context.Context, credentialID string, credential Credential, request ExecuteRequest) (*ExecuteStreamResponse, error) { + response, err := value.bridge.ExecuteStreamCanonical(ctx, credentialID, credentialToBridge(credential), requestToBridge(request)) + if response == nil { + return nil, normalizeError(err) + } + chunks := make(chan ExecuteStreamChunk) + go func() { + defer close(chunks) + for chunk := range response.Chunks { + converted := ExecuteStreamChunk{Payload: append([]byte(nil), chunk.Payload...), Err: normalizeError(chunk.Err)} + select { + case chunks <- converted: + case <-ctx.Done(): + return + } + } + }() + return &ExecuteStreamResponse{ + Headers: response.Headers.Clone(), Chunks: chunks, + AppliedReasoningEffort: response.AppliedReasoningEffort, + }, normalizeError(err) +} + +func requestToBridge(request ExecuteRequest) cpaembedded.ExecuteRequest { + return cpaembedded.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: request.ContinuityKey, + } +} + +func responseFromBridge(response cpaembedded.ExecuteResponse) ExecuteResponse { + return ExecuteResponse{ + Payload: append([]byte(nil), response.Payload...), Headers: response.Headers.Clone(), + AppliedReasoningEffort: response.AppliedReasoningEffort, + } +} diff --git a/internal/subscription/providers/grok/grok.go b/internal/subscription/providers/grok/grok.go new file mode 100644 index 00000000..110429a2 --- /dev/null +++ b/internal/subscription/providers/grok/grok.go @@ -0,0 +1,245 @@ +// Package grok owns GPT-Load's Grok subscription contract and isolates the +// embedded CLIProxyAPI xAI bridge from the rest of the application. +package grok + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "time" + + cpaembedded "github.com/router-for-me/CLIProxyAPI/v7/gptload-embedded/embedded" +) + +const Provider = cpaembedded.ProviderGrok + +var ErrCredentialIdentityChanged = errors.New("refreshed grok credential identity changed") + +type Credential struct { + Type string `json:"type"` + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + IDToken string `json:"id_token,omitempty"` + TokenType string `json:"token_type,omitempty"` + ExpiresIn int `json:"expires_in,omitempty"` + Expire string `json:"expired"` + LastRefresh string `json:"last_refresh,omitempty"` + AccountID string `json:"account_id"` + Email string `json:"email"` + TokenEndpoint string `json:"token_endpoint"` +} + +type DeviceAuthorization struct { + VerificationURL string + UserCode string + State DeviceAuthorizationState + ExpiresAt time.Time + PollInterval time.Duration +} + +type DeviceAuthorizationState struct { + DeviceCode string `json:"device_code"` + TokenEndpoint string `json:"token_endpoint"` + UserInfoEndpoint string `json:"userinfo_endpoint"` + PollIntervalSeconds int `json:"poll_interval_seconds"` +} + +type DeviceAuthorizationStatus string + +const ( + DeviceAuthorizationPending DeviceAuthorizationStatus = "pending" + DeviceAuthorizationAuthorized DeviceAuthorizationStatus = "authorized" + DeviceAuthorizationDenied DeviceAuthorizationStatus = "denied" + DeviceAuthorizationExpired DeviceAuthorizationStatus = "expired" +) + +type DeviceAuthorizationPoll struct { + Status DeviceAuthorizationStatus + State DeviceAuthorizationState + Credential Credential + PollInterval time.Duration +} + +type TokenEndpointError struct { + StatusCode int + Code string +} + +func (err *TokenEndpointError) Error() string { + if err == nil { + return "Grok token endpoint failed" + } + return fmt.Sprintf("Grok 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 "Grok upstream request failed" + } + return fmt.Sprintf("Grok %s endpoint returned status %d", err.Operation, err.StatusCode) +} + +func (err *UpstreamHTTPError) HTTPStatusCode() int { + if err == nil { + return 0 + } + return err.StatusCode +} + +func ParseCredentialJSON(raw []byte) (Credential, error) { + value, err := cpaembedded.ParseGrokCredentialJSON(raw) + if err != nil { + return Credential{}, normalizeError(err) + } + return credentialFromBridge(value), nil +} + +func MarshalCredential(value Credential) ([]byte, error) { + raw, err := cpaembedded.MarshalGrokCredential(credentialToBridge(value)) + if err != nil { + return nil, normalizeError(err) + } + return raw, nil +} + +func CredentialExpiresAt(value Credential) (time.Time, bool) { + return cpaembedded.GrokCredentialExpiresAt(credentialToBridge(value)) +} + +func (value Credential) SecretValues() []string { return credentialToBridge(value).SecretValues() } + +func BeginDeviceAuthorization(ctx context.Context) (DeviceAuthorization, error) { + value, err := cpaembedded.BeginGrokDeviceAuthorization(ctx, cpaembedded.GrokOptions{}) + if err != nil { + return DeviceAuthorization{}, normalizeError(err) + } + return DeviceAuthorization{ + VerificationURL: value.VerificationURL, UserCode: value.UserCode, + State: stateFromBridge(value.State), ExpiresAt: value.ExpiresAt, PollInterval: value.PollInterval, + }, nil +} + +func PollDeviceAuthorizationOnce(ctx context.Context, state DeviceAuthorizationState) (DeviceAuthorizationPoll, error) { + value, err := cpaembedded.PollGrokDeviceAuthorizationOnce(ctx, stateToBridge(state), cpaembedded.GrokOptions{}) + if err != nil { + return DeviceAuthorizationPoll{}, normalizeError(err) + } + return DeviceAuthorizationPoll{ + Status: DeviceAuthorizationStatus(value.Status), State: stateFromBridge(value.State), + Credential: credentialFromBridge(value.Credential), PollInterval: value.PollInterval, + }, nil +} + +func RefreshCredentialOnce(ctx context.Context, current Credential) (Credential, error) { + value, err := cpaembedded.RefreshGrokCredentialOnce(ctx, credentialToBridge(current), cpaembedded.GrokOptions{}) + if err != nil { + return Credential{}, normalizeError(err) + } + return credentialFromBridge(value), nil +} + +func ImportCredential(ctx context.Context, raw []byte) (Credential, error) { + value, err := cpaembedded.ImportGrokCredential(ctx, raw, cpaembedded.GrokOptions{}) + if err != nil { + return Credential{}, normalizeError(err) + } + return credentialFromBridge(value), nil +} + +func ListModels(ctx context.Context, credential Credential) ([]string, error) { + values, err := cpaembedded.DiscoverGrokModels(ctx, credentialToBridge(credential), cpaembedded.GrokOptions{}) + if err != nil { + return nil, normalizeError(err) + } + return append([]string(nil), values...), nil +} + +func IsDefinitiveRefreshRejection(code string) bool { + switch strings.ToLower(strings.TrimSpace(code)) { + case "invalid_grant", "refresh_token_expired", "refresh_token_revoked", "access_denied": + return true + default: + return false + } +} + +func credentialFromBridge(value cpaembedded.GrokCredential) Credential { + return Credential{ + Type: value.Type, AccessToken: value.AccessToken, RefreshToken: value.RefreshToken, + IDToken: value.IDToken, TokenType: value.TokenType, ExpiresIn: value.ExpiresIn, + Expire: value.Expire, LastRefresh: value.LastRefresh, AccountID: value.AccountID, + Email: value.Email, TokenEndpoint: value.TokenEndpoint, + } +} + +func credentialToBridge(value Credential) cpaembedded.GrokCredential { + return cpaembedded.GrokCredential{ + Type: value.Type, AccessToken: value.AccessToken, RefreshToken: value.RefreshToken, + IDToken: value.IDToken, TokenType: value.TokenType, ExpiresIn: value.ExpiresIn, + Expire: value.Expire, LastRefresh: value.LastRefresh, AccountID: value.AccountID, + Email: value.Email, TokenEndpoint: value.TokenEndpoint, + } +} + +func stateFromBridge(value cpaembedded.GrokDeviceState) DeviceAuthorizationState { + return DeviceAuthorizationState{ + DeviceCode: value.DeviceCode, TokenEndpoint: value.TokenEndpoint, + UserInfoEndpoint: value.UserInfoEndpoint, PollIntervalSeconds: value.PollIntervalSeconds, + } +} + +func stateToBridge(value DeviceAuthorizationState) cpaembedded.GrokDeviceState { + return cpaembedded.GrokDeviceState{ + DeviceCode: value.DeviceCode, TokenEndpoint: value.TokenEndpoint, + UserInfoEndpoint: value.UserInfoEndpoint, PollIntervalSeconds: value.PollIntervalSeconds, + } +} + +func normalizeError(err error) error { + if err == nil { + return nil + } + if errors.Is(err, cpaembedded.ErrGrokCredentialIdentityChanged) { + return ErrCredentialIdentityChanged + } + var tokenErr *cpaembedded.GrokTokenEndpointError + if errors.As(err, &tokenErr) { + return &TokenEndpointError{StatusCode: tokenErr.StatusCode, Code: strings.TrimSpace(tokenErr.Code)} + } + var upstream *cpaembedded.GrokUpstreamHTTPError + if errors.As(err, &upstream) { + return &UpstreamHTTPError{Operation: strings.TrimSpace(upstream.Operation), StatusCode: upstream.StatusCode} + } + return err +} + +func marshalDeviceState(value DeviceAuthorizationState) ([]byte, error) { return json.Marshal(value) } + +func unmarshalDeviceState(raw []byte) (DeviceAuthorizationState, error) { + var value DeviceAuthorizationState + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&value); err != nil { + return DeviceAuthorizationState{}, err + } + if strings.TrimSpace(value.DeviceCode) == "" || strings.TrimSpace(value.TokenEndpoint) == "" || + strings.TrimSpace(value.UserInfoEndpoint) == "" || value.PollIntervalSeconds < 1 || value.PollIntervalSeconds > 60 { + return DeviceAuthorizationState{}, errors.New("invalid Grok device authorization state") + } + return value, nil +} diff --git a/internal/subscription/providers/providers.go b/internal/subscription/providers/providers.go index b401dd64..2e139e17 100644 --- a/internal/subscription/providers/providers.go +++ b/internal/subscription/providers/providers.go @@ -7,6 +7,7 @@ import ( "gpt-load/internal/subscription/providers/antigravity" "gpt-load/internal/subscription/providers/claude" "gpt-load/internal/subscription/providers/codex" + "gpt-load/internal/subscription/providers/grok" subscriptionruntime "gpt-load/internal/subscription/runtime" ) @@ -15,5 +16,6 @@ func Implementations() []subscriptionruntime.Implementations { codex.Implementations(), claude.Implementations(), antigravity.Implementations(), + grok.Implementations(), } } diff --git a/internal/subscription/providers/providers_test.go b/internal/subscription/providers/providers_test.go index 7b5b9d25..30e27043 100644 --- a/internal/subscription/providers/providers_test.go +++ b/internal/subscription/providers/providers_test.go @@ -16,7 +16,7 @@ func TestRuntimeCompilesAllSubscriptionProviderCapabilities(t *testing.T) { if err != nil { t.Fatal(err) } - if got, want := runtime.ChannelIDs(), []channel.ID{channel.Antigravity, channel.Claude, channel.Codex}; !reflect.DeepEqual(got, want) { + if got, want := runtime.ChannelIDs(), []channel.ID{channel.Antigravity, channel.Claude, channel.Codex, channel.Grok}; !reflect.DeepEqual(got, want) { t.Fatalf("ChannelIDs() = %v, want %v", got, want) } tests := []struct { @@ -29,6 +29,7 @@ func TestRuntimeCompilesAllSubscriptionProviderCapabilities(t *testing.T) { {channel.Codex, string(modules.CodexSubscriptionDriver), string(modules.CodexModelDiscovery), string(modules.CodexQuotaObservation), string(modules.CodexResetCreditAction)}, {channel.Claude, string(modules.ClaudeSubscriptionDriver), string(modules.ClaudeModelDiscovery), string(modules.ClaudeQuotaObservation), ""}, {channel.Antigravity, string(modules.AntigravitySubscriptionDriver), string(modules.AntigravityModelDiscovery), string(modules.AntigravityQuotaObservation), ""}, + {channel.Grok, string(modules.GrokSubscriptionDriver), string(modules.GrokModelDiscovery), "", ""}, } for _, test := range tests { t.Run(string(test.channelID), func(t *testing.T) { @@ -41,7 +42,7 @@ func TestRuntimeCompilesAllSubscriptionProviderCapabilities(t *testing.T) { t.Fatalf("discovery = %#v/%t", discovery, ok) } observation, ok := runtime.QuotaObservation(test.channelID) - if !ok || string(observation.ID()) != test.observationID { + if test.observationID == "" && ok || test.observationID != "" && (!ok || string(observation.ID()) != test.observationID) { t.Fatalf("observation = %#v/%t", observation, ok) } action, ok := runtime.ResetCreditAction(test.channelID) @@ -53,6 +54,12 @@ func TestRuntimeCompilesAllSubscriptionProviderCapabilities(t *testing.T) { if _, ok := runtime.Driver(channel.OpenAI); ok { t.Fatal("Driver(openai) unexpectedly resolved") } + if device, ok := runtime.DeviceAuthorization(channel.Grok); !ok || device.ID() != modules.GrokSubscriptionDriver { + t.Fatalf("Grok device authorization = %#v/%t", device, ok) + } + if _, ok := runtime.BrowserAuthorization(channel.Grok); ok { + t.Fatal("Grok unexpectedly resolves browser OAuth") + } } type browserlessDriver struct{ subscriptionruntime.Driver } diff --git a/internal/subscription/runtime/runtime.go b/internal/subscription/runtime/runtime.go index 6826e773..4986149b 100644 --- a/internal/subscription/runtime/runtime.go +++ b/internal/subscription/runtime/runtime.go @@ -119,6 +119,43 @@ type BrowserAuthorizationDriver interface { LocalCallback() (LocalCallbackSpec, bool) } +// DeviceAuthorization describes one RFC 8628 device challenge. DriverState +// contains secrets and is persisted only inside the encrypted Stage payload. +type DeviceAuthorization struct { + VerificationURL string + UserCode string + DriverState []byte + ExpiresAt time.Time + PollInterval time.Duration +} + +// DeviceAuthorizationStatus is the bounded result of one token-endpoint poll. +type DeviceAuthorizationStatus string + +const ( + DeviceAuthorizationPending DeviceAuthorizationStatus = "pending" + DeviceAuthorizationAuthorized DeviceAuthorizationStatus = "authorized" + DeviceAuthorizationDenied DeviceAuthorizationStatus = "denied" + DeviceAuthorizationExpired DeviceAuthorizationStatus = "expired" +) + +// DeviceAuthorizationPoll contains the result of exactly one token request. +// Pending results may replace DriverState and adjust PollInterval (slow_down). +type DeviceAuthorizationPoll struct { + Status DeviceAuthorizationStatus + Credential Credential + DriverState []byte + PollInterval time.Duration +} + +// DeviceAuthorizationDriver is implemented by subscription channels whose +// interactive flow uses a device code rather than an OAuth redirect callback. +type DeviceAuthorizationDriver interface { + Driver + BeginDeviceAuthorization(context.Context) (DeviceAuthorization, error) + PollDeviceAuthorization(context.Context, []byte) (DeviceAuthorizationPoll, error) +} + // ModelDiscovery is a narrow optional subscription capability. type ModelDiscovery interface { ID() spec.UtilityID @@ -174,6 +211,7 @@ func (err *UpstreamHTTPError) Error() string { type channelRuntime struct { driver Driver browser BrowserAuthorizationDriver + device DeviceAuthorizationDriver callback *LocalCallbackSpec discovery ModelDiscovery observation QuotaObservation @@ -279,27 +317,38 @@ func compileRuntime( return nil, fmt.Errorf("compile subscription runtime: channel %q references unknown driver %q", descriptor.ID, bindings.SubscriptionDriver) } compiled := channelRuntime{driver: driver} + interactiveMethods := 0 for _, method := range descriptor.Connection.AuthorizationMethods { - if method != channel.AuthorizationBrowserOAuth { - continue - } - compiled.browser, ok = driver.(BrowserAuthorizationDriver) - if !ok { - return nil, fmt.Errorf("compile subscription runtime: channel %q declares browser OAuth without driver support", descriptor.ID) - } - callback, local := compiled.browser.LocalCallback() - if local { - callback, err = validateLocalCallbackSpec(callback) - if err != nil { - return nil, fmt.Errorf("compile subscription runtime: channel %q has invalid local callback: %w", descriptor.ID, err) + switch method { + case channel.AuthorizationBrowserOAuth: + interactiveMethods++ + compiled.browser, ok = driver.(BrowserAuthorizationDriver) + if !ok { + return nil, fmt.Errorf("compile subscription runtime: channel %q declares browser OAuth without driver support", descriptor.ID) } - if owner, duplicate := callbackChannels[callback.RedirectURI]; duplicate { - return nil, fmt.Errorf("compile subscription runtime: channels %q and %q share local callback %q", owner, descriptor.ID, callback.RedirectURI) + callback, local := compiled.browser.LocalCallback() + if local { + callback, err = validateLocalCallbackSpec(callback) + if err != nil { + return nil, fmt.Errorf("compile subscription runtime: channel %q has invalid local callback: %w", descriptor.ID, err) + } + if owner, duplicate := callbackChannels[callback.RedirectURI]; duplicate { + return nil, fmt.Errorf("compile subscription runtime: channels %q and %q share local callback %q", owner, descriptor.ID, callback.RedirectURI) + } + callbackChannels[callback.RedirectURI] = descriptor.ID + compiled.callback = &callback + } + case channel.AuthorizationDeviceOAuth: + interactiveMethods++ + compiled.device, ok = driver.(DeviceAuthorizationDriver) + if !ok { + return nil, fmt.Errorf("compile subscription runtime: channel %q declares device OAuth without driver support", descriptor.ID) } - callbackChannels[callback.RedirectURI] = descriptor.ID - compiled.callback = &callback } } + if interactiveMethods > 1 { + return nil, fmt.Errorf("compile subscription runtime: channel %q declares multiple interactive authorization methods", descriptor.ID) + } if bindings.ModelDiscovery != "" { compiled.discovery, ok = discoveryByID[spec.ExtensionID(bindings.ModelDiscovery)] if !ok { @@ -376,6 +425,15 @@ func (runtime *Runtime) BrowserAuthorization(channelID channel.ID) (BrowserAutho return value.browser, ok && value.browser != nil } +// DeviceAuthorization resolves the optional device authorization driver. +func (runtime *Runtime) DeviceAuthorization(channelID channel.ID) (DeviceAuthorizationDriver, bool) { + if runtime == nil { + return nil, false + } + value, ok := runtime.byChannel[channelID] + return value.device, ok && value.device != nil +} + // LocalCallback returns the startup-validated fixed callback for a channel. func (runtime *Runtime) LocalCallback(channelID channel.ID) (LocalCallbackSpec, bool) { if runtime == nil { diff --git a/internal/subscription/runtime/runtime_test.go b/internal/subscription/runtime/runtime_test.go index 97be604d..ea1c017f 100644 --- a/internal/subscription/runtime/runtime_test.go +++ b/internal/subscription/runtime/runtime_test.go @@ -2,6 +2,7 @@ package subscriptionruntime import ( "context" + "strings" "testing" "gpt-load/internal/channel" @@ -15,6 +16,64 @@ func TestRuntimeRequiresExplicitImplementations(t *testing.T) { } } +type testDiscovery struct{ id spec.UtilityID } + +func (value testDiscovery) ID() spec.UtilityID { return value.id } +func (testDiscovery) DiscoverModels(context.Context, Credential) ([]string, error) { + return nil, nil +} + +type testObservation struct{ id spec.UtilityID } + +func (value testObservation) ID() spec.UtilityID { return value.id } +func (testObservation) Observe(context.Context, Credential) (Observation, error) { + return Observation{}, nil +} + +type testResetCredit struct{ id spec.ActionID } + +func (value testResetCredit) ID() spec.ActionID { return value.id } +func (testResetCredit) Consume(context.Context, Credential, string) (ResetCreditResult, error) { + return ResetCreditResult{}, nil +} + +type testBrowserDriver struct{ duplicateDriver } + +func (testBrowserDriver) BeginAuthorization() (Authorization, error) { return Authorization{}, nil } +func (testBrowserDriver) CompleteAuthorization(context.Context, AuthorizationCompletion) (Credential, error) { + return Credential{}, nil +} +func (testBrowserDriver) AuthorizationFailureDefinitive(error) bool { return false } +func (testBrowserDriver) LocalCallback() (LocalCallbackSpec, bool) { return LocalCallbackSpec{}, false } + +func completeTestRuntimeImplementations() ([]Driver, []ModelDiscovery, []QuotaObservation, []ResetCreditAction) { + return []Driver{ + testBrowserDriver{duplicateDriver{id: modules.CodexSubscriptionDriver}}, + testBrowserDriver{duplicateDriver{id: modules.ClaudeSubscriptionDriver}}, + testBrowserDriver{duplicateDriver{id: modules.AntigravitySubscriptionDriver}}, + duplicateDriver{id: modules.GrokSubscriptionDriver}, + }, []ModelDiscovery{ + testDiscovery{id: modules.CodexModelDiscovery}, + testDiscovery{id: modules.ClaudeModelDiscovery}, + testDiscovery{id: modules.AntigravityModelDiscovery}, + testDiscovery{id: modules.GrokModelDiscovery}, + }, []QuotaObservation{ + testObservation{id: modules.CodexQuotaObservation}, + testObservation{id: modules.ClaudeQuotaObservation}, + testObservation{id: modules.AntigravityQuotaObservation}, + }, []ResetCreditAction{ + testResetCredit{id: modules.CodexResetCreditAction}, + } +} + +func TestRuntimeRejectsDeviceOAuthWithoutDriverSupport(t *testing.T) { + drivers, discoveries, observations, actions := completeTestRuntimeImplementations() + _, err := compileRuntime(channel.NewRegistry(), drivers, discoveries, observations, actions) + if err == nil || !strings.Contains(err.Error(), "device OAuth") { + t.Fatalf("compileRuntime() error = %v, want missing device OAuth support", err) + } +} + func TestRuntimeFailsClosedWhenBoundImplementationIsMissing(t *testing.T) { _, err := compileRuntime(channel.NewRegistry(), nil, nil, nil, nil) if err == nil { diff --git a/third_party/cpaembedded/README.md b/third_party/cpaembedded/README.md index b3b0fff4..6281867d 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, Claude, and Antigravity executors and +CLIProxyAPI (CPA). It exists because the CPA Codex, Claude, Antigravity, and xAI 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. @@ -25,6 +25,11 @@ quota policy. This bridge only exposes: - 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. +- xAI OIDC discovery, one-shot device-code begin/poll, strict native/canonical file enrichment, + stable OIDC subject identity, and one-shot context-aware refresh; +- the execution-only Grok HTTP executor, live OAuth model discovery, explicit local CountTokens + estimator, and supported protocol translators, without CPA manager, refresh, retry, fallback, + WebSocket, image, or video execution paths. It intentionally excludes CPA Manager, selector, pool, file store, server, watcher, WebSocket/Auto executors, fallback, and internal retry loops. @@ -43,7 +48,7 @@ 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, Claude, and Antigravity OAuth, token, HTTP +1. Review upstream changes to Codex, Claude, Antigravity, and xAI 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 @@ -97,3 +102,14 @@ CPA_LIVE_ANTIGRAVITY_MODEL=optional-antigravity-model-id \ Browser OAuth, refresh-token rotation, deliberate 401/429 responses, and provider policy changes remain supervised live gates rather than default test behavior. + +The Grok contract requires a prepared canonical xAI OAuth credential. It verifies live OAuth +models, all four unary/streaming protocol routes, and the three explicit local CountTokens +representations. It does not invent an active plan or quota observation because CPA and the +verified xAI execution surface expose no stable proactive quota endpoint. + +```bash +CPA_LIVE_GROK_CREDENTIAL_FILE=/absolute/path/to/grok.json \ +CPA_LIVE_GROK_MODEL=optional-grok-model-id \ + go test -count=1 -run '^TestLiveGrokContract$' ./embedded +``` diff --git a/third_party/cpaembedded/embedded/grok.go b/third_party/cpaembedded/embedded/grok.go new file mode 100644 index 00000000..6fa4a5f5 --- /dev/null +++ b/third_party/cpaembedded/embedded/grok.go @@ -0,0 +1,280 @@ +package embedded + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "time" +) + +const ProviderGrok = "grok" + +var ErrGrokCredentialIdentityChanged = errors.New("refreshed grok credential identity changed") + +// GrokCredential is GPT-Load's canonical, execution-only xAI OAuth schema. +type GrokCredential struct { + Type string `json:"type"` + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + IDToken string `json:"id_token,omitempty"` + TokenType string `json:"token_type,omitempty"` + ExpiresIn int `json:"expires_in,omitempty"` + Expire string `json:"expired"` + LastRefresh string `json:"last_refresh,omitempty"` + AccountID string `json:"account_id"` + Email string `json:"email"` + TokenEndpoint string `json:"token_endpoint"` +} + +// GrokOptions provides test-only endpoint/clock seams. Production callers use +// zero values so all endpoints come from xAI's fixed issuer discovery. +type GrokOptions struct { + DiscoveryURL string + UserInfoURL string + ModelsURL string + HTTPClient *http.Client + Now func() time.Time +} + +type GrokDeviceState struct { + DeviceCode string `json:"device_code"` + TokenEndpoint string `json:"token_endpoint"` + UserInfoEndpoint string `json:"userinfo_endpoint"` + PollIntervalSeconds int `json:"poll_interval_seconds"` +} + +type GrokDeviceAuthorization struct { + VerificationURL string + UserCode string + State GrokDeviceState + ExpiresAt time.Time + PollInterval time.Duration +} + +type GrokDeviceStatus string + +const ( + GrokDevicePending GrokDeviceStatus = "pending" + GrokDeviceAuthorized GrokDeviceStatus = "authorized" + GrokDeviceDenied GrokDeviceStatus = "denied" + GrokDeviceExpired GrokDeviceStatus = "expired" +) + +type GrokDevicePoll struct { + Status GrokDeviceStatus + State GrokDeviceState + Credential GrokCredential + PollInterval time.Duration +} + +type GrokTokenEndpointError struct { + StatusCode int + Code string +} + +func (err *GrokTokenEndpointError) Error() string { + if err == nil { + return "Grok token endpoint failed" + } + return fmt.Sprintf("Grok token endpoint returned status %d", err.StatusCode) +} + +func (err *GrokTokenEndpointError) HTTPStatusCode() int { + if err == nil { + return 0 + } + return err.StatusCode +} + +type GrokUpstreamHTTPError struct { + Operation string + StatusCode int +} + +func (err *GrokUpstreamHTTPError) Error() string { + if err == nil { + return "Grok upstream request failed" + } + return fmt.Sprintf("Grok %s endpoint returned status %d", err.Operation, err.StatusCode) +} + +func (err *GrokUpstreamHTTPError) HTTPStatusCode() int { + if err == nil { + return 0 + } + return err.StatusCode +} + +func ParseGrokCredentialJSON(raw []byte) (GrokCredential, error) { + if len(raw) == 0 || len(raw) > maxCredentialBytes { + return GrokCredential{}, fmt.Errorf("credential JSON size is invalid") + } + allowed := grokCanonicalFields() + allowCPAAuthFileControlFields(allowed) + if err := validateClaudeCredentialObject(raw, allowed); err != nil { + return GrokCredential{}, err + } + if err := validateCPAAuthFileControlMetadata(raw); err != nil { + return GrokCredential{}, err + } + var credential GrokCredential + if err := json.Unmarshal(raw, &credential); err != nil { + return GrokCredential{}, fmt.Errorf("decode credential: %w", err) + } + normalizeGrokCredential(&credential) + if err := validateGrokCredential(credential); err != nil { + return GrokCredential{}, err + } + return credential, nil +} + +func MarshalGrokCredential(credential GrokCredential) ([]byte, error) { + raw, err := json.Marshal(credential) + if err != nil { + return nil, err + } + parsed, err := ParseGrokCredentialJSON(raw) + clear(raw) + if err != nil { + return nil, err + } + return json.Marshal(parsed) +} + +func grokCanonicalFields() map[string]struct{} { + return map[string]struct{}{ + "type": {}, "access_token": {}, "refresh_token": {}, "id_token": {}, + "token_type": {}, "expires_in": {}, "expired": {}, "last_refresh": {}, + "account_id": {}, "email": {}, "token_endpoint": {}, + } +} + +func normalizeGrokCredential(credential *GrokCredential) { + credential.Type = strings.ToLower(strings.TrimSpace(credential.Type)) + credential.AccessToken = strings.TrimSpace(credential.AccessToken) + credential.RefreshToken = strings.TrimSpace(credential.RefreshToken) + credential.IDToken = strings.TrimSpace(credential.IDToken) + credential.TokenType = strings.TrimSpace(credential.TokenType) + if credential.TokenType == "" { + credential.TokenType = "Bearer" + } + credential.Expire = strings.TrimSpace(credential.Expire) + credential.LastRefresh = strings.TrimSpace(credential.LastRefresh) + credential.AccountID = strings.TrimSpace(credential.AccountID) + credential.Email = strings.TrimSpace(credential.Email) + credential.TokenEndpoint = strings.TrimSpace(credential.TokenEndpoint) +} + +func validateGrokCredential(credential GrokCredential) error { + return validateGrokCredentialWithOptions(credential, GrokOptions{}) +} + +func validateGrokCredentialWithOptions(credential GrokCredential, options GrokOptions) error { + if credential.Type != ProviderGrok { + return fmt.Errorf("credential type must be grok") + } + for field, value := range map[string]string{ + "access_token": credential.AccessToken, "refresh_token": credential.RefreshToken, + "account_id": credential.AccountID, "email": credential.Email, + "expired": credential.Expire, "token_endpoint": credential.TokenEndpoint, + } { + if value == "" { + return fmt.Errorf("credential %s is required", field) + } + } + if credential.ExpiresIn < 0 { + return fmt.Errorf("credential expires_in is invalid") + } + if !strings.EqualFold(credential.TokenType, "Bearer") { + return fmt.Errorf("credential token_type must be Bearer") + } + if err := validateTimestamp("expired", credential.Expire); err != nil { + return err + } + if err := validateTimestamp("last_refresh", credential.LastRefresh); err != nil { + return err + } + for field, value := range map[string]string{ + "account_id": credential.AccountID, "email": credential.Email, + "token_type": credential.TokenType, "token_endpoint": credential.TokenEndpoint, + } { + if len(value) > 2048 || strings.ContainsAny(value, "\r\n\x00") { + return fmt.Errorf("credential %s is invalid", field) + } + } + if _, err := validateGrokOAuthEndpoint( + credential.TokenEndpoint, + "token_endpoint", + grokAllowsTestEndpoints(options), + ); err != nil { + return err + } + return nil +} + +func grokAllowsTestEndpoints(options GrokOptions) bool { + return strings.TrimSpace(options.DiscoveryURL) != "" || + strings.TrimSpace(options.UserInfoURL) != "" || + strings.TrimSpace(options.ModelsURL) != "" || + options.HTTPClient != nil +} + +func GrokCredentialExpiresAt(credential GrokCredential) (time.Time, bool) { + parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(credential.Expire)) + return parsed, err == nil +} + +func (credential GrokCredential) SecretValues() []string { + return []string{ + credential.AccessToken, credential.RefreshToken, credential.IDToken, + credential.AccountID, credential.Email, + } +} + +func BeginGrokDeviceAuthorization(ctx context.Context, options GrokOptions) (GrokDeviceAuthorization, error) { + return beginGrokDeviceAuthorization(ctx, options) +} + +func PollGrokDeviceAuthorizationOnce(ctx context.Context, state GrokDeviceState, options GrokOptions) (GrokDevicePoll, error) { + return pollGrokDeviceAuthorizationOnce(ctx, state, options) +} + +func ImportGrokCredential(ctx context.Context, raw []byte, options GrokOptions) (GrokCredential, error) { + return importGrokCredential(ctx, raw, options) +} + +func RefreshGrokCredentialOnce(ctx context.Context, current GrokCredential, options GrokOptions) (GrokCredential, error) { + return refreshGrokCredentialOnce(ctx, current, options) +} + +func grokNow(options GrokOptions) time.Time { + if options.Now != nil { + return options.Now().UTC() + } + return time.Now().UTC() +} + +func grokHTTPClient(options GrokOptions) *http.Client { + if options.HTTPClient != nil { + return clientWithoutRedirects(options.HTTPClient) + } + return clientWithoutRedirects(http.DefaultClient) +} + +func decodeGrokDeviceState(raw []byte) (GrokDeviceState, error) { + var state GrokDeviceState + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&state); err != nil { + return GrokDeviceState{}, err + } + if strings.TrimSpace(state.DeviceCode) == "" || strings.TrimSpace(state.TokenEndpoint) == "" || + strings.TrimSpace(state.UserInfoEndpoint) == "" || state.PollIntervalSeconds < 1 || state.PollIntervalSeconds > 60 { + return GrokDeviceState{}, fmt.Errorf("Grok device authorization state is invalid") + } + return state, nil +} diff --git a/third_party/cpaembedded/embedded/grok_executor.go b/third_party/cpaembedded/embedded/grok_executor.go new file mode 100644 index 00000000..a8aa3d5b --- /dev/null +++ b/third_party/cpaembedded/embedded/grok_executor.go @@ -0,0 +1,297 @@ +package embedded + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + internalexecutor "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + 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/sjson" +) + +const grokCPAProvider = "xai" + +type GrokExecutionError struct { + status int + code string + summary string + retryAfter time.Duration +} + +func (err *GrokExecutionError) Error() string { + if err == nil || err.summary == "" { + return "Grok upstream request failed" + } + return err.summary +} + +func (err *GrokExecutionError) StatusCode() int { + if err == nil { + return 0 + } + return err.status +} + +func (err *GrokExecutionError) ErrorCode() string { + if err == nil { + return "" + } + return err.code +} + +func (err *GrokExecutionError) RetryAfter() *time.Duration { + if err == nil || err.retryAfter <= 0 { + return nil + } + value := err.retryAfter + return &value +} + +type GrokHTTPExecutor interface { + ExecuteCanonical(context.Context, string, GrokCredential, ExecuteRequest) (ExecuteResponse, error) + CountTokensCanonical(context.Context, ExecuteRequest) (ExecuteResponse, error) + ExecuteStreamCanonical(context.Context, string, GrokCredential, ExecuteRequest) (*ExecuteStreamResponse, error) +} + +type grokHTTPExecutor struct { + cfg *internalconfig.Config + inner *internalexecutor.XAIExecutor + baseURL string +} + +func NewGrokHTTPExecutor() GrokHTTPExecutor { + cfg := &internalconfig.Config{} + return &grokHTTPExecutor{cfg: cfg, inner: internalexecutor.NewXAIExecutor(cfg)} +} + +func NewGrokAuth(id string, credential GrokCredential, baseURL string) *cliproxyauth.Auth { + metadata := map[string]any{ + "type": ProviderGrok, "access_token": credential.AccessToken, + "auth_kind": "oauth", "using_api": false, + } + attributes := map[string]string{ + "auth_kind": "oauth", "using_api": "false", + } + if value := strings.TrimSpace(baseURL); value != "" { + attributes["base_url"] = value + } + return &cliproxyauth.Auth{ + ID: strings.TrimSpace(id), Provider: grokCPAProvider, Attributes: attributes, Metadata: metadata, + } +} + +func (executor *grokHTTPExecutor) ExecuteCanonical( + ctx context.Context, + credentialID string, + credential GrokCredential, + request ExecuteRequest, +) (ExecuteResponse, error) { + if ctx == nil { + ctx = context.Background() + } + if err := validateGrokCredential(credential); err != nil { + return ExecuteResponse{}, err + } + request = prepareGrokExecutionRequest(request) + format := sdktranslator.FromString(request.Format) + auth := NewGrokAuth(credentialID, credential, executor.baseURL) + observation := newProviderExecutionObservation(request, grokCPAProvider) + response, err := executor.inner.Execute(executor.executionContext(ctx, auth, observation), auth, cliproxyexecutor.Request{ + Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: format, + }, grokExecutorOptions(request, format, false)) + if err != nil { + return ExecuteResponse{AppliedReasoningEffort: observation.reasoningEffort()}, normalizeGrokExecutionError(err) + } + return ExecuteResponse{ + Payload: append([]byte(nil), response.Payload...), Headers: response.Headers.Clone(), + AppliedReasoningEffort: observation.reasoningEffort(), + }, nil +} + +func (executor *grokHTTPExecutor) CountTokensCanonical( + ctx context.Context, + request ExecuteRequest, +) (ExecuteResponse, error) { + if ctx == nil { + ctx = context.Background() + } + request = prepareGrokExecutionRequest(request) + format := sdktranslator.FromString(request.Format) + response, err := executor.inner.CountTokens(ctx, NewGrokAuth("local-token-count", GrokCredential{}, executor.baseURL), cliproxyexecutor.Request{ + Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: format, + }, grokExecutorOptions(request, format, false)) + if err != nil { + return ExecuteResponse{}, normalizeGrokExecutionError(err) + } + payload := append([]byte(nil), response.Payload...) + if format == sdktranslator.FormatOpenAIResponse { + payload, err = normalizeCodexResponsesTokenCount(payload) + if err != nil { + return ExecuteResponse{}, err + } + } + return ExecuteResponse{Payload: payload, Headers: response.Headers.Clone()}, nil +} + +func (executor *grokHTTPExecutor) ExecuteStreamCanonical( + ctx context.Context, + credentialID string, + credential GrokCredential, + request ExecuteRequest, +) (*ExecuteStreamResponse, error) { + if ctx == nil { + ctx = context.Background() + } + if err := validateGrokCredential(credential); err != nil { + return nil, err + } + request = prepareGrokExecutionRequest(request) + format := sdktranslator.FromString(request.Format) + auth := NewGrokAuth(credentialID, credential, executor.baseURL) + observation := newProviderExecutionObservation(request, grokCPAProvider) + response, err := executor.inner.ExecuteStream(executor.executionContext(ctx, auth, observation), auth, cliproxyexecutor.Request{ + Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: format, + }, grokExecutorOptions(request, format, true)) + if err != nil { + return &ExecuteStreamResponse{AppliedReasoningEffort: observation.reasoningEffort()}, normalizeGrokExecutionError(err) + } + chunks := make(chan ExecuteStreamChunk) + go func() { + defer close(chunks) + for chunk := range response.Chunks { + converted := ExecuteStreamChunk{Payload: append([]byte(nil), chunk.Payload...)} + if chunk.Err != nil { + converted.Err = normalizeGrokExecutionError(chunk.Err) + } + select { + case chunks <- converted: + case <-ctx.Done(): + return + } + } + }() + return &ExecuteStreamResponse{ + Headers: response.Headers.Clone(), Chunks: chunks, + AppliedReasoningEffort: observation.reasoningEffort(), + }, nil +} + +func grokExecutorOptions(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 prepareGrokExecutionRequest(request ExecuteRequest) ExecuteRequest { + if strings.TrimSpace(request.ContinuityKey) == "" && strings.TrimSpace(request.AttemptID) != "" { + request.ContinuityKey = "gpt-load-grok-attempt\x00" + strings.TrimSpace(request.AttemptID) + } + request.Payload = stripGrokCallerContinuity(request.Payload) + request.OriginalRequest = stripGrokCallerContinuity(request.OriginalRequest) + request.Headers = request.Headers.Clone() + for name := range request.Headers { + if strings.EqualFold(name, "Session-Id") || strings.EqualFold(name, "X-Grok-Conv-Id") { + delete(request.Headers, name) + } + } + return request +} + +func stripGrokCallerContinuity(raw []byte) []byte { + result := append([]byte(nil), raw...) + for _, path := range []string{ + "session_id", "sessionId", "prompt_cache_key", "metadata.session_id", "metadata.sessionId", + } { + updated, err := sjson.DeleteBytes(result, path) + if err == nil { + result = updated + } + } + return result +} + +func (executor *grokHTTPExecutor) executionContext( + ctx context.Context, + auth *cliproxyauth.Auth, + observation *executionObservation, +) context.Context { + baseClient := helps.NewProxyAwareHTTPClient(ctx, executor.cfg, auth, 0) + transport := baseClient.Transport + if transport == nil { + transport = http.DefaultTransport + } + return context.WithValue(ctx, "cliproxy.roundtripper", noRedirectRoundTripper{ + base: transport, observation: observation, + }) +} + +func normalizeGrokExecutionError(err error) error { + if err == nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return err + } + var statusError interface { + error + StatusCode() int + } + if !errors.As(err, &statusError) || statusError == nil || statusError.StatusCode() == 0 { + return err + } + status := statusError.StatusCode() + 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 + } + } + code := boundedGrokExecutionCode([]byte(statusError.Error())) + return &GrokExecutionError{ + status: status, code: code, retryAfter: retryAfter, + summary: grokExecutionSummary(status, code), + } +} + +func boundedGrokExecutionCode(raw []byte) string { + var payload struct { + Code string `json:"code"` + Error any `json:"error"` + } + if json.Unmarshal(raw, &payload) != nil { + return "" + } + code := strings.TrimSpace(payload.Code) + if code == "" { + if object, ok := payload.Error.(map[string]any); ok { + code, _ = object["code"].(string) + } + } + code = strings.TrimSpace(code) + if code == "" || len(code) > 128 || strings.ContainsAny(code, "\r\n\x00") { + return "" + } + return code +} + +func grokExecutionSummary(status int, code string) string { + if strings.Contains(strings.ToLower(code), "free-usage-exhausted") { + return "Grok included free usage is exhausted." + } + return fmt.Sprintf("Grok upstream returned status %d", status) +} + +var _ GrokHTTPExecutor = (*grokHTTPExecutor)(nil) diff --git a/third_party/cpaembedded/embedded/grok_executor_test.go b/third_party/cpaembedded/embedded/grok_executor_test.go new file mode 100644 index 00000000..8144f22f --- /dev/null +++ b/third_party/cpaembedded/embedded/grok_executor_test.go @@ -0,0 +1,308 @@ +package embedded + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + internalexecutor "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor" + "github.com/tidwall/gjson" +) + +func TestGrokExecutorConvertsFourProtocolsUnaryAndStream(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/responses" || r.Header.Get("Authorization") != "Bearer access-secret" { + t.Fatalf("request = %s %s, auth %q", r.Method, r.URL.Path, r.Header.Get("Authorization")) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("event: response.completed\n")) + _, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","object":"response","created_at":0,"status":"completed","model":"grok-4.3","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":2,"output_tokens":1,"total_tokens":3}}}` + "\n\n")) + })) + defer server.Close() + + cfg := &internalconfig.Config{} + executor := &grokHTTPExecutor{cfg: cfg, inner: internalexecutor.NewXAIExecutor(cfg), baseURL: server.URL} + credential := testGrokExecutionCredential() + for _, test := range []struct { + name string + format string + payload string + usage string + }{ + {name: "responses", format: "openai-response", payload: `{"model":"grok-4.3","input":"hi"}`, usage: "usage.input_tokens"}, + {name: "chat", format: "openai", payload: `{"model":"grok-4.3","messages":[{"role":"user","content":"hi"}]}`, usage: "usage.prompt_tokens"}, + {name: "anthropic", format: "claude", payload: `{"model":"grok-4.3","max_tokens":64,"messages":[{"role":"user","content":"hi"}]}`, usage: "usage.input_tokens"}, + {name: "gemini", format: "gemini", payload: `{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`, usage: "usageMetadata.promptTokenCount"}, + } { + t.Run(test.name+" unary", func(t *testing.T) { + response, err := executor.ExecuteCanonical(t.Context(), "credential-1", credential, ExecuteRequest{ + AttemptID: "attempt-1", Model: "grok-4.3", Format: test.format, + Payload: []byte(test.payload), OriginalRequest: []byte(test.payload), ContinuityKey: "tenant-scope", + }) + if err != nil { + t.Fatal(err) + } + if !bytes.Contains(response.Payload, []byte("ok")) || gjson.GetBytes(response.Payload, test.usage).Int() != 2 { + t.Fatalf("response = %s", response.Payload) + } + }) + t.Run(test.name+" stream", func(t *testing.T) { + response, err := executor.ExecuteStreamCanonical(t.Context(), "credential-1", credential, ExecuteRequest{ + AttemptID: "attempt-1", Model: "grok-4.3", Format: test.format, + Payload: []byte(test.payload), OriginalRequest: []byte(test.payload), ContinuityKey: "tenant-scope", + }) + if err != nil { + t.Fatal(err) + } + var output []byte + for chunk := range response.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + output = append(output, chunk.Payload...) + } + if len(output) == 0 { + t.Fatal("stream returned no translated chunks") + } + }) + } +} + +func TestGrokExecutorCountsTokensLocally(t *testing.T) { + executor := NewGrokHTTPExecutor() + for _, test := range []struct { + format string + payload string + path string + }{ + {format: "openai-response", payload: `{"model":"grok-4.3","input":"hello world"}`, path: "input_tokens"}, + {format: "claude", payload: `{"model":"grok-4.3","messages":[{"role":"user","content":"hello world"}]}`, path: "input_tokens"}, + {format: "gemini", payload: `{"contents":[{"role":"user","parts":[{"text":"hello world"}]}]}`, path: "totalTokens"}, + } { + response, err := executor.CountTokensCanonical(context.Background(), ExecuteRequest{ + Model: "grok-4.3", Format: test.format, Payload: []byte(test.payload), OriginalRequest: []byte(test.payload), + }) + if err != nil { + t.Fatal(err) + } + if gjson.GetBytes(response.Payload, test.path).Int() <= 0 || strings.Contains(string(response.Payload), "access-secret") { + t.Fatalf("count response = %s", response.Payload) + } + if test.format == "openai-response" && gjson.GetBytes(response.Payload, "object").String() != "response.input_tokens" { + t.Fatalf("Responses token count = %s", response.Payload) + } + } +} + +func TestGrokExecutorPreservesRichChatRequestAndToolResponse(t *testing.T) { + var upstreamBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + if err != nil { + t.Fatal(err) + } + upstreamBody = append([]byte(nil), body...) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("event: response.completed\n")) + _, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_tool","object":"response","status":"completed","model":"grok-4.3","output":[{"type":"reasoning","summary":[{"type":"summary_text","text":"considering"}]},{"type":"function_call","id":"fc_1","call_id":"call_1","name":"lookup","arguments":"{\"q\":\"x\"}","status":"completed"}],"usage":{"input_tokens":5,"output_tokens":3,"total_tokens":8}}}` + "\n\n")) + })) + defer server.Close() + + payload := []byte(`{ + "model":"grok-4.3","reasoning_effort":"high", + "messages":[ + {"role":"system","content":"follow the rule"}, + {"role":"user","content":[ + {"type":"text","text":"inspect"}, + {"type":"image_url","image_url":{"url":"data:image/png;base64,AAAA"}} + ]} + ], + "tools":[{"type":"function","function":{"name":"lookup","description":"lookup","parameters":{"type":"object","properties":{"q":{"type":"string"}}}}}], + "tool_choice":"auto" + }`) + response, err := testGrokHTTPExecutor(server.URL).ExecuteCanonical( + t.Context(), "credential-1", testGrokExecutionCredential(), ExecuteRequest{ + AttemptID: "attempt-1", Model: "grok-4.3", Format: "openai", + Payload: payload, OriginalRequest: payload, ContinuityKey: "tenant-scope", + }, + ) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(upstreamBody), "follow the rule") || + gjson.GetBytes(upstreamBody, "tools.0.name").String() != "lookup" || + gjson.GetBytes(upstreamBody, "reasoning.effort").String() != "high" || + !strings.Contains(string(upstreamBody), "data:image/png;base64,AAAA") { + t.Fatalf("rich upstream request = %s", upstreamBody) + } + if gjson.GetBytes(response.Payload, "choices.0.message.tool_calls.0.function.name").String() != "lookup" || + gjson.GetBytes(response.Payload, "choices.0.finish_reason").String() != "tool_calls" || + gjson.GetBytes(response.Payload, "usage.prompt_tokens").Int() != 5 || + gjson.GetBytes(response.Payload, "usage.completion_tokens").Int() != 3 || + response.AppliedReasoningEffort != "high" { + t.Fatalf("rich downstream response = %s, applied=%q", response.Payload, response.AppliedReasoningEffort) + } +} + +func TestGrokExecutionOnlyBridgeUsesOneDispatchAndRejectsRedirects(t *testing.T) { + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls.Add(1) + w.Header().Set("Location", "https://example.invalid/redirected") + w.WriteHeader(http.StatusTemporaryRedirect) + })) + defer server.Close() + + executor := testGrokHTTPExecutor(server.URL) + _, err := executor.ExecuteCanonical(t.Context(), "credential-1", testGrokExecutionCredential(), testGrokRequest()) + if !errors.Is(err, ErrRedirectNotAllowed) || calls.Load() != 1 { + t.Fatalf("redirect result = %v, calls = %d", err, calls.Load()) + } +} + +func TestGrokExecutionOnlyBridgeDoesNotRetryOrRefreshPreparedCredential(t *testing.T) { + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls.Add(1) + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"code":"subscription:free-usage-exhausted","error":"exhausted"}`)) + })) + defer server.Close() + + credential := testGrokExecutionCredential() + credential.Expire = time.Now().UTC().Add(time.Second).Format(time.RFC3339) + _, err := testGrokHTTPExecutor(server.URL).ExecuteCanonical(t.Context(), "credential-1", credential, testGrokRequest()) + var upstream *GrokExecutionError + if !errors.As(err, &upstream) || upstream.StatusCode() != http.StatusTooManyRequests || + upstream.ErrorCode() != "subscription:free-usage-exhausted" || calls.Load() != 1 { + t.Fatalf("execution result = %#v, calls = %d", err, calls.Load()) + } + if retry := upstream.RetryAfter(); retry == nil || *retry != 24*time.Hour { + t.Fatalf("RetryAfter = %v", retry) + } +} + +func TestGrokExecutionOnlyBridgeMapsBadCredentialsToUnauthorized(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"code":"unauthenticated:bad-credentials","error":"invalid OAuth access token"}`)) + })) + defer server.Close() + + _, err := testGrokHTTPExecutor(server.URL).ExecuteCanonical( + t.Context(), "credential-1", testGrokExecutionCredential(), testGrokRequest(), + ) + var upstream *GrokExecutionError + if !errors.As(err, &upstream) || upstream.StatusCode() != http.StatusUnauthorized || + upstream.ErrorCode() != "unauthenticated:bad-credentials" { + t.Fatalf("bad credentials error = %#v", err) + } +} + +func TestGrokExecutionRequestRemovesCallerContinuity(t *testing.T) { + prepared := prepareGrokExecutionRequest(ExecuteRequest{ + AttemptID: "attempt-1", ContinuityKey: "private-scope", + Payload: []byte(`{"input":"keep","session_id":"caller","prompt_cache_key":"caller-cache","metadata":{"sessionId":"nested"}}`), + OriginalRequest: []byte(`{"input":"keep","sessionId":"caller","prompt_cache_key":"caller-cache","metadata":{"session_id":"nested"}}`), + Headers: http.Header{"Session-Id": {"caller"}, "X-Grok-Conv-Id": {"caller-conversation"}}, + }) + for _, raw := range [][]byte{prepared.Payload, prepared.OriginalRequest} { + if !json.Valid(raw) || strings.Contains(string(raw), "caller") || !strings.Contains(string(raw), "keep") { + t.Fatalf("prepared payload = %s", raw) + } + } + if prepared.Headers.Get("Session-Id") != "" || prepared.Headers.Get("X-Grok-Conv-Id") != "" || + prepared.ContinuityKey != "private-scope" { + t.Fatalf("prepared request = %#v", prepared) + } +} + +func TestGrokExecutionOnlyBridgePropagatesCancellation(t *testing.T) { + started := make(chan struct{}) + release := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + close(started) + select { + case <-r.Context().Done(): + case <-release: + } + })) + defer func() { + close(release) + server.Close() + }() + + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { + _, err := testGrokHTTPExecutor(server.URL).ExecuteCanonical(ctx, "credential-1", testGrokExecutionCredential(), testGrokRequest()) + done <- err + }() + <-started + cancel() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) { + t.Fatalf("cancellation error = %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("canceled Grok request did not return") + } +} + +type grokWrappedStatusError struct { + status int + body string +} + +func (err grokWrappedStatusError) Error() string { return err.body } +func (err grokWrappedStatusError) StatusCode() int { return err.status } +func (err grokWrappedStatusError) RetryAfter() *time.Duration { + value := 3 * time.Second + return &value +} + +func TestNormalizeGrokExecutionErrorPreservesWrappedMetadata(t *testing.T) { + err := normalizeGrokExecutionError(fmt.Errorf("execute Grok request: %w", grokWrappedStatusError{ + status: http.StatusTooManyRequests, + body: `{"code":"subscription:free-usage-exhausted"}`, + })) + var normalized *GrokExecutionError + if !errors.As(err, &normalized) || normalized.StatusCode() != http.StatusTooManyRequests || + normalized.ErrorCode() != "subscription:free-usage-exhausted" || normalized.RetryAfter() == nil || + *normalized.RetryAfter() != 3*time.Second { + t.Fatalf("normalized error = %#v", err) + } +} + +func testGrokHTTPExecutor(baseURL string) *grokHTTPExecutor { + cfg := &internalconfig.Config{} + return &grokHTTPExecutor{cfg: cfg, inner: internalexecutor.NewXAIExecutor(cfg), baseURL: baseURL} +} + +func testGrokRequest() ExecuteRequest { + payload := []byte(`{"model":"grok-4.3","input":"hello"}`) + return ExecuteRequest{ + AttemptID: "attempt-1", Model: "grok-4.3", Format: "openai-response", + Payload: payload, OriginalRequest: append([]byte(nil), payload...), ContinuityKey: "tenant-scope", + } +} + +func testGrokExecutionCredential() GrokCredential { + return GrokCredential{ + Type: ProviderGrok, AccessToken: "access-secret", RefreshToken: "refresh-secret", + TokenType: "Bearer", AccountID: "account-1", Email: "owner@example.com", + TokenEndpoint: "https://auth.x.ai/oauth/token", + Expire: time.Now().UTC().Add(time.Hour).Format(time.RFC3339), + } +} diff --git a/third_party/cpaembedded/embedded/grok_models.go b/third_party/cpaembedded/embedded/grok_models.go new file mode 100644 index 00000000..d43d1745 --- /dev/null +++ b/third_party/cpaembedded/embedded/grok_models.go @@ -0,0 +1,78 @@ +package embedded + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "sort" + "strings" + + xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" +) + +const ( + grokClientVersion = "0.2.120" + grokModelsURL = xaiauth.CLIChatProxyBaseURL + "/models" +) + +func DiscoverGrokModels(ctx context.Context, credential GrokCredential, options GrokOptions) ([]string, error) { + endpoint := strings.TrimSpace(options.ModelsURL) + if endpoint == "" { + endpoint = grokModelsURL + } + request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) + if err != nil { + return nil, err + } + request.Header.Set("Accept", "application/json") + request.Header.Set("Authorization", "Bearer "+strings.TrimSpace(credential.AccessToken)) + request.Header.Set("X-XAI-Token-Auth", "xai-grok-cli") + request.Header.Set("x-grok-client-version", grokClientVersion) + request.Header.Set("User-Agent", "xai-grok-workspace/"+grokClientVersion) + request.Header.Set("x-grok-client-identifier", "grok-shell") + request.Header.Set("x-authenticateresponse", "authenticate-response") + response, err := grokHTTPClient(options).Do(request) + if err != nil { + return nil, err + } + defer func() { _ = response.Body.Close() }() + body, err := io.ReadAll(io.LimitReader(response.Body, maxTokenResponse+1)) + if err != nil { + return nil, err + } + defer clear(body) + if len(body) > maxTokenResponse { + return nil, fmt.Errorf("Grok models response is too large") + } + if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { + return nil, &GrokUpstreamHTTPError{Operation: "models", StatusCode: response.StatusCode} + } + var payload struct { + Data []struct { + ID string `json:"id"` + } `json:"data"` + } + if err := json.Unmarshal(body, &payload); err != nil { + return nil, fmt.Errorf("decode Grok models: %w", err) + } + seen := make(map[string]struct{}, len(payload.Data)) + models := make([]string, 0, len(payload.Data)) + for _, item := range payload.Data { + id := strings.TrimSpace(item.ID) + if id == "" || len(id) > 256 || strings.ContainsAny(id, "\r\n\x00") { + continue + } + if _, duplicate := seen[id]; duplicate { + continue + } + seen[id] = struct{}{} + models = append(models, id) + } + if len(models) == 0 { + return nil, fmt.Errorf("Grok models response has no usable model IDs") + } + sort.Strings(models) + return models, nil +} diff --git a/third_party/cpaembedded/embedded/grok_oauth.go b/third_party/cpaembedded/embedded/grok_oauth.go new file mode 100644 index 00000000..73e82cfa --- /dev/null +++ b/third_party/cpaembedded/embedded/grok_oauth.go @@ -0,0 +1,577 @@ +package embedded + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" + + xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" +) + +const grokDefaultPollInterval = 5 * time.Second + +type grokDiscovery struct { + DeviceAuthorizationEndpoint string `json:"device_authorization_endpoint"` + TokenEndpoint string `json:"token_endpoint"` + UserInfoEndpoint string `json:"userinfo_endpoint"` +} + +type grokTokenResponse struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + IDToken string `json:"id_token"` + TokenType string `json:"token_type"` + ExpiresIn int `json:"expires_in"` + Error string `json:"error"` +} + +type grokUserInfo struct { + Subject string `json:"sub"` + Email string `json:"email"` + EmailVerified bool `json:"email_verified"` +} + +type grokImportFile struct { + Type string `json:"type"` + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + IDToken string `json:"id_token"` + TokenType string `json:"token_type"` + ExpiresIn int `json:"expires_in"` + Expire string `json:"expired"` + LastRefresh string `json:"last_refresh"` + AccountID string `json:"account_id"` + Subject string `json:"sub"` + Email string `json:"email"` + TokenEndpoint string `json:"token_endpoint"` + BaseURL string `json:"base_url"` + RedirectURI string `json:"redirect_uri"` + AuthKind string `json:"auth_kind"` +} + +func beginGrokDeviceAuthorization(ctx context.Context, options GrokOptions) (GrokDeviceAuthorization, error) { + discovery, err := discoverGrokOAuth(ctx, options) + if err != nil { + return GrokDeviceAuthorization{}, err + } + form := url.Values{"client_id": {xaiauth.ClientID}, "scope": {xaiauth.Scope}} + body, status, err := doGrokForm(ctx, options, discovery.DeviceAuthorizationEndpoint, form) + if err != nil { + return GrokDeviceAuthorization{}, fmt.Errorf("begin Grok device authorization: %w", err) + } + defer clear(body) + if status != http.StatusOK { + return GrokDeviceAuthorization{}, &GrokUpstreamHTTPError{Operation: "device authorization", StatusCode: status} + } + var payload struct { + DeviceCode string `json:"device_code"` + UserCode string `json:"user_code"` + VerificationURI string `json:"verification_uri"` + VerificationURIComplete string `json:"verification_uri_complete"` + ExpiresIn int `json:"expires_in"` + Interval int `json:"interval"` + } + if err := json.Unmarshal(body, &payload); err != nil { + return GrokDeviceAuthorization{}, fmt.Errorf("decode Grok device authorization: %w", err) + } + verificationURL := strings.TrimSpace(payload.VerificationURIComplete) + if verificationURL == "" { + verificationURL = strings.TrimSpace(payload.VerificationURI) + } + if strings.TrimSpace(payload.DeviceCode) == "" || strings.TrimSpace(payload.UserCode) == "" || + verificationURL == "" || payload.ExpiresIn <= 0 { + return GrokDeviceAuthorization{}, fmt.Errorf("Grok device authorization response is incomplete") + } + interval := time.Duration(payload.Interval) * time.Second + if interval < grokDefaultPollInterval { + interval = grokDefaultPollInterval + } + if interval > time.Minute { + return GrokDeviceAuthorization{}, fmt.Errorf("Grok device authorization interval is invalid") + } + expiresIn := time.Duration(payload.ExpiresIn) * time.Second + if expiresIn > xaiauth.MaxPollDuration { + expiresIn = xaiauth.MaxPollDuration + } + return GrokDeviceAuthorization{ + VerificationURL: verificationURL, + UserCode: strings.TrimSpace(payload.UserCode), + State: GrokDeviceState{ + DeviceCode: strings.TrimSpace(payload.DeviceCode), TokenEndpoint: discovery.TokenEndpoint, + UserInfoEndpoint: discovery.UserInfoEndpoint, PollIntervalSeconds: int(interval / time.Second), + }, + ExpiresAt: grokNow(options).Add(expiresIn), PollInterval: interval, + }, nil +} + +func pollGrokDeviceAuthorizationOnce(ctx context.Context, state GrokDeviceState, options GrokOptions) (GrokDevicePoll, error) { + raw, err := json.Marshal(state) + if err != nil { + return GrokDevicePoll{}, err + } + state, err = decodeGrokDeviceState(raw) + clear(raw) + if err != nil { + return GrokDevicePoll{}, err + } + form := url.Values{ + "grant_type": {xaiauth.DeviceCodeGrantType}, "device_code": {state.DeviceCode}, "client_id": {xaiauth.ClientID}, + } + body, status, err := doGrokForm(ctx, options, state.TokenEndpoint, form) + if err != nil { + return GrokDevicePoll{}, fmt.Errorf("poll Grok device authorization: %w", err) + } + defer clear(body) + var token grokTokenResponse + if err := json.Unmarshal(body, &token); err != nil { + return GrokDevicePoll{}, fmt.Errorf("decode Grok device token response: %w", err) + } + code := strings.ToLower(boundedGrokOAuthCode(token.Error)) + interval := time.Duration(state.PollIntervalSeconds) * time.Second + if code != "" { + switch code { + case "authorization_pending": + return GrokDevicePoll{Status: GrokDevicePending, State: state, PollInterval: interval}, nil + case "slow_down": + interval += grokDefaultPollInterval + if interval > time.Minute { + interval = time.Minute + } + state.PollIntervalSeconds = int(interval / time.Second) + return GrokDevicePoll{Status: GrokDevicePending, State: state, PollInterval: interval}, nil + case "access_denied": + return GrokDevicePoll{Status: GrokDeviceDenied, State: state}, nil + case "expired_token": + return GrokDevicePoll{Status: GrokDeviceExpired, State: state}, nil + default: + return GrokDevicePoll{}, &GrokTokenEndpointError{StatusCode: status, Code: code} + } + } + if status != http.StatusOK { + return GrokDevicePoll{}, &GrokTokenEndpointError{StatusCode: status} + } + credential, err := grokCredentialFromToken(ctx, token, "", state.TokenEndpoint, state.UserInfoEndpoint, options, true) + if err != nil { + return GrokDevicePoll{}, err + } + return GrokDevicePoll{Status: GrokDeviceAuthorized, State: state, Credential: credential}, nil +} + +func importGrokCredential(ctx context.Context, raw []byte, options GrokOptions) (GrokCredential, error) { + if len(raw) == 0 || len(raw) > maxCredentialBytes { + return GrokCredential{}, fmt.Errorf("credential JSON size is invalid") + } + allowed := grokCanonicalFields() + for _, field := range []string{"sub", "base_url", "redirect_uri", "auth_kind"} { + allowed[field] = struct{}{} + } + allowCPAAuthFileControlFields(allowed) + if err := validateClaudeCredentialObject(raw, allowed); err != nil { + return GrokCredential{}, err + } + if err := validateCPAAuthFileControlMetadata(raw); err != nil { + return GrokCredential{}, err + } + var imported grokImportFile + if err := json.Unmarshal(raw, &imported); err != nil { + return GrokCredential{}, fmt.Errorf("decode Grok credential: %w", err) + } + imported.Type = strings.ToLower(strings.TrimSpace(imported.Type)) + if imported.Type != ProviderGrok && imported.Type != "xai" { + return GrokCredential{}, fmt.Errorf("credential type must be grok or xai") + } + if imported.AuthKind != "" && !strings.EqualFold(strings.TrimSpace(imported.AuthKind), "oauth") { + return GrokCredential{}, fmt.Errorf("credential auth_kind must be oauth") + } + if imported.TokenType != "" && !strings.EqualFold(strings.TrimSpace(imported.TokenType), "Bearer") { + return GrokCredential{}, fmt.Errorf("credential token_type must be Bearer") + } + if imported.ExpiresIn < 0 { + return GrokCredential{}, fmt.Errorf("credential expires_in is invalid") + } + if strings.TrimSpace(imported.AccountID) != "" && strings.TrimSpace(imported.Subject) != "" && + strings.TrimSpace(imported.AccountID) != strings.TrimSpace(imported.Subject) { + return GrokCredential{}, ErrGrokCredentialIdentityChanged + } + if imported.BaseURL != "" { + baseURL := strings.TrimRight(strings.TrimSpace(imported.BaseURL), "/") + if baseURL != strings.TrimRight(xaiauth.DefaultAPIBaseURL, "/") && baseURL != strings.TrimRight(xaiauth.CLIChatProxyBaseURL, "/") { + return GrokCredential{}, fmt.Errorf("credential base_url is not supported") + } + } + if imported.RedirectURI != "" { + if _, err := validateGrokOAuthEndpoint(imported.RedirectURI, "redirect_uri", false); err != nil { + return GrokCredential{}, fmt.Errorf("credential redirect_uri is invalid") + } + } + if strings.TrimSpace(imported.AccessToken) == "" || strings.TrimSpace(imported.RefreshToken) == "" { + return GrokCredential{}, fmt.Errorf("credential tokens are required") + } + discovery, err := resolveGrokTokenAndUserInfo(ctx, imported.TokenEndpoint, options) + if err != nil { + return GrokCredential{}, err + } + candidate := GrokCredential{ + Type: ProviderGrok, AccessToken: imported.AccessToken, RefreshToken: imported.RefreshToken, + IDToken: imported.IDToken, TokenType: imported.TokenType, ExpiresIn: imported.ExpiresIn, + Expire: imported.Expire, LastRefresh: imported.LastRefresh, + AccountID: firstNonEmptyString(imported.AccountID, imported.Subject), Email: imported.Email, + TokenEndpoint: discovery.TokenEndpoint, + } + normalizeGrokCredential(&candidate) + refreshed := false + if expiresAt, ok := GrokCredentialExpiresAt(candidate); !ok || !expiresAt.After(grokNow(options).Add(5*time.Minute)) { + candidate, err = refreshGrokTokensOnly(ctx, candidate, options) + if err != nil { + return GrokCredential{}, err + } + refreshed = true + } + profile, err := fetchGrokUserInfo(ctx, discovery.UserInfoEndpoint, candidate.AccessToken, options) + if err != nil && !refreshed { + var upstream *GrokUpstreamHTTPError + if errors.As(err, &upstream) && upstream.StatusCode == http.StatusUnauthorized { + candidate, err = refreshGrokTokensOnly(ctx, candidate, options) + if err == nil { + profile, err = fetchGrokUserInfo(ctx, discovery.UserInfoEndpoint, candidate.AccessToken, options) + } + } + } + if err != nil { + return GrokCredential{}, err + } + expectedID := firstNonEmptyString(imported.AccountID, imported.Subject) + if expectedID != "" && profile.Subject != expectedID { + return GrokCredential{}, ErrGrokCredentialIdentityChanged + } + if expectedID == "" && strings.TrimSpace(imported.Email) != "" && !strings.EqualFold(strings.TrimSpace(imported.Email), profile.Email) { + return GrokCredential{}, ErrGrokCredentialIdentityChanged + } + candidate.AccountID = profile.Subject + candidate.Email = profile.Email + if candidate.Expire == "" { + return GrokCredential{}, fmt.Errorf("credential expired is required") + } + normalizeGrokCredential(&candidate) + if err := validateGrokCredentialWithOptions(candidate, options); err != nil { + return GrokCredential{}, err + } + return candidate, nil +} + +func refreshGrokCredentialOnce(ctx context.Context, current GrokCredential, options GrokOptions) (GrokCredential, error) { + normalizeGrokCredential(¤t) + if err := validateGrokCredentialWithOptions(current, options); err != nil { + return GrokCredential{}, err + } + refreshed, err := refreshGrokTokensOnly(ctx, current, options) + if err != nil { + return GrokCredential{}, err + } + if subject := grokJWTSubject(refreshed.IDToken); subject != "" && subject != current.AccountID { + return GrokCredential{}, ErrGrokCredentialIdentityChanged + } + discovery, discoveryErr := resolveGrokTokenAndUserInfo(ctx, refreshed.TokenEndpoint, options) + if discoveryErr == nil { + if profile, profileErr := fetchGrokUserInfo(ctx, discovery.UserInfoEndpoint, refreshed.AccessToken, options); profileErr == nil { + if profile.Subject != current.AccountID { + return GrokCredential{}, ErrGrokCredentialIdentityChanged + } + refreshed.Email = profile.Email + } + } + refreshed.AccountID = current.AccountID + normalizeGrokCredential(&refreshed) + if err := validateGrokCredentialWithOptions(refreshed, options); err != nil { + return GrokCredential{}, err + } + return refreshed, nil +} + +func refreshGrokTokensOnly(ctx context.Context, current GrokCredential, options GrokOptions) (GrokCredential, error) { + if strings.TrimSpace(current.RefreshToken) == "" { + return GrokCredential{}, fmt.Errorf("Grok refresh token is required") + } + discovery, err := resolveGrokTokenAndUserInfo(ctx, current.TokenEndpoint, options) + if err != nil { + return GrokCredential{}, err + } + form := url.Values{ + "grant_type": {"refresh_token"}, "client_id": {xaiauth.ClientID}, "refresh_token": {strings.TrimSpace(current.RefreshToken)}, + } + body, status, err := doGrokForm(ctx, options, discovery.TokenEndpoint, form) + if err != nil { + return GrokCredential{}, fmt.Errorf("refresh Grok credential: %w", err) + } + defer clear(body) + var token grokTokenResponse + if err := json.Unmarshal(body, &token); err != nil { + return GrokCredential{}, fmt.Errorf("decode Grok refresh response: %w", err) + } + if status != http.StatusOK { + return GrokCredential{}, &GrokTokenEndpointError{StatusCode: status, Code: boundedGrokOAuthCode(token.Error)} + } + if strings.TrimSpace(token.AccessToken) == "" { + return GrokCredential{}, fmt.Errorf("Grok refresh response has no access token") + } + now := grokNow(options) + refreshed := current + refreshed.Type = ProviderGrok + refreshed.AccessToken = strings.TrimSpace(token.AccessToken) + if strings.TrimSpace(token.RefreshToken) != "" { + refreshed.RefreshToken = strings.TrimSpace(token.RefreshToken) + } + if strings.TrimSpace(token.IDToken) != "" { + refreshed.IDToken = strings.TrimSpace(token.IDToken) + } + if strings.TrimSpace(token.TokenType) != "" { + refreshed.TokenType = strings.TrimSpace(token.TokenType) + } + if token.ExpiresIn > 0 { + refreshed.ExpiresIn = token.ExpiresIn + refreshed.Expire = now.Add(time.Duration(token.ExpiresIn) * time.Second).Format(time.RFC3339) + } + refreshed.LastRefresh = now.Format(time.RFC3339) + refreshed.TokenEndpoint = discovery.TokenEndpoint + return refreshed, nil +} + +func grokCredentialFromToken( + ctx context.Context, + token grokTokenResponse, + expectedID string, + tokenEndpoint string, + userInfoEndpoint string, + options GrokOptions, + requireRefreshToken bool, +) (GrokCredential, error) { + if strings.TrimSpace(token.AccessToken) == "" || (requireRefreshToken && strings.TrimSpace(token.RefreshToken) == "") { + return GrokCredential{}, fmt.Errorf("Grok token response is incomplete") + } + profile, err := fetchGrokUserInfo(ctx, userInfoEndpoint, token.AccessToken, options) + if err != nil { + return GrokCredential{}, err + } + if expectedID != "" && profile.Subject != expectedID { + return GrokCredential{}, ErrGrokCredentialIdentityChanged + } + if subject := grokJWTSubject(token.IDToken); subject != "" && subject != profile.Subject { + return GrokCredential{}, ErrGrokCredentialIdentityChanged + } + now := grokNow(options) + credential := GrokCredential{ + Type: ProviderGrok, AccessToken: token.AccessToken, RefreshToken: token.RefreshToken, + IDToken: token.IDToken, TokenType: token.TokenType, ExpiresIn: token.ExpiresIn, + LastRefresh: now.Format(time.RFC3339), AccountID: profile.Subject, Email: profile.Email, + TokenEndpoint: tokenEndpoint, + } + if token.ExpiresIn > 0 { + credential.Expire = now.Add(time.Duration(token.ExpiresIn) * time.Second).Format(time.RFC3339) + } + normalizeGrokCredential(&credential) + if err := validateGrokCredentialWithOptions(credential, options); err != nil { + return GrokCredential{}, err + } + return credential, nil +} + +func discoverGrokOAuth(ctx context.Context, options GrokOptions) (grokDiscovery, error) { + discoveryURL := strings.TrimSpace(options.DiscoveryURL) + allowTestEndpoint := discoveryURL != "" + if discoveryURL == "" { + discoveryURL = xaiauth.DiscoveryURL + } + request, err := http.NewRequestWithContext(ctx, http.MethodGet, discoveryURL, nil) + if err != nil { + return grokDiscovery{}, err + } + request.Header.Set("Accept", "application/json") + response, err := grokHTTPClient(options).Do(request) + if err != nil { + return grokDiscovery{}, err + } + defer func() { _ = response.Body.Close() }() + body, err := io.ReadAll(io.LimitReader(response.Body, maxTokenResponse+1)) + if err != nil { + return grokDiscovery{}, err + } + defer clear(body) + if len(body) > maxTokenResponse { + return grokDiscovery{}, fmt.Errorf("Grok discovery response is too large") + } + if response.StatusCode != http.StatusOK { + return grokDiscovery{}, &GrokUpstreamHTTPError{Operation: "discovery", StatusCode: response.StatusCode} + } + var discovery grokDiscovery + if err := json.Unmarshal(body, &discovery); err != nil { + return grokDiscovery{}, fmt.Errorf("decode Grok discovery: %w", err) + } + for field, raw := range map[string]string{ + "device_authorization_endpoint": discovery.DeviceAuthorizationEndpoint, + "token_endpoint": discovery.TokenEndpoint, + "userinfo_endpoint": discovery.UserInfoEndpoint, + } { + validated, err := validateGrokOAuthEndpoint(raw, field, allowTestEndpoint) + if err != nil { + return grokDiscovery{}, err + } + switch field { + case "device_authorization_endpoint": + discovery.DeviceAuthorizationEndpoint = validated + case "token_endpoint": + discovery.TokenEndpoint = validated + case "userinfo_endpoint": + discovery.UserInfoEndpoint = validated + } + } + return discovery, nil +} + +func resolveGrokTokenAndUserInfo(ctx context.Context, tokenEndpoint string, options GrokOptions) (grokDiscovery, error) { + userInfo := strings.TrimSpace(options.UserInfoURL) + allowTestEndpoint := strings.TrimSpace(options.DiscoveryURL) != "" || userInfo != "" + if strings.TrimSpace(tokenEndpoint) != "" && userInfo != "" { + validatedToken, err := validateGrokOAuthEndpoint(tokenEndpoint, "token_endpoint", allowTestEndpoint) + if err != nil { + return grokDiscovery{}, err + } + validatedUserInfo, err := validateGrokOAuthEndpoint(userInfo, "userinfo_endpoint", allowTestEndpoint) + if err != nil { + return grokDiscovery{}, err + } + return grokDiscovery{TokenEndpoint: validatedToken, UserInfoEndpoint: validatedUserInfo}, nil + } + discovery, err := discoverGrokOAuth(ctx, options) + if err != nil { + return grokDiscovery{}, err + } + if strings.TrimSpace(tokenEndpoint) != "" { + validated, err := validateGrokOAuthEndpoint(tokenEndpoint, "token_endpoint", allowTestEndpoint) + if err != nil { + return grokDiscovery{}, err + } + discovery.TokenEndpoint = validated + } + if userInfo != "" { + validated, err := validateGrokOAuthEndpoint(userInfo, "userinfo_endpoint", allowTestEndpoint) + if err != nil { + return grokDiscovery{}, err + } + discovery.UserInfoEndpoint = validated + } + return discovery, nil +} + +func validateGrokOAuthEndpoint(raw, field string, allowTestEndpoint bool) (string, error) { + raw = strings.TrimSpace(raw) + parsed, err := url.Parse(raw) + if err != nil || !parsed.IsAbs() || parsed.User != nil || parsed.Fragment != "" { + return "", fmt.Errorf("Grok %s is invalid", field) + } + if allowTestEndpoint { + return raw, nil + } + return xaiauth.ValidateOAuthEndpoint(raw, field) +} + +func fetchGrokUserInfo(ctx context.Context, endpoint, accessToken string, options GrokOptions) (grokUserInfo, error) { + request, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimSpace(endpoint), nil) + if err != nil { + return grokUserInfo{}, err + } + request.Header.Set("Accept", "application/json") + request.Header.Set("Authorization", "Bearer "+strings.TrimSpace(accessToken)) + response, err := grokHTTPClient(options).Do(request) + if err != nil { + return grokUserInfo{}, err + } + defer func() { _ = response.Body.Close() }() + body, err := io.ReadAll(io.LimitReader(response.Body, maxTokenResponse+1)) + if err != nil { + return grokUserInfo{}, err + } + defer clear(body) + if len(body) > maxTokenResponse { + return grokUserInfo{}, fmt.Errorf("Grok userinfo response is too large") + } + if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { + return grokUserInfo{}, &GrokUpstreamHTTPError{Operation: "userinfo", StatusCode: response.StatusCode} + } + var profile grokUserInfo + if err := json.Unmarshal(body, &profile); err != nil { + return grokUserInfo{}, fmt.Errorf("decode Grok userinfo: %w", err) + } + profile.Subject = strings.TrimSpace(profile.Subject) + profile.Email = strings.TrimSpace(profile.Email) + if profile.Subject == "" || profile.Email == "" || !profile.EmailVerified || + strings.ContainsAny(profile.Subject+profile.Email, "\r\n\x00") { + return grokUserInfo{}, fmt.Errorf("Grok userinfo identity is incomplete") + } + return profile, nil +} + +func doGrokForm(ctx context.Context, options GrokOptions, endpoint string, form url.Values) ([]byte, int, error) { + request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimSpace(endpoint), strings.NewReader(form.Encode())) + if err != nil { + return nil, 0, err + } + request.Header.Set("Content-Type", "application/x-www-form-urlencoded") + request.Header.Set("Accept", "application/json") + response, err := grokHTTPClient(options).Do(request) + if err != nil { + return nil, 0, err + } + defer func() { _ = response.Body.Close() }() + body, err := io.ReadAll(io.LimitReader(response.Body, maxTokenResponse+1)) + if err != nil { + return nil, response.StatusCode, err + } + if len(body) > maxTokenResponse { + clear(body) + return nil, response.StatusCode, fmt.Errorf("Grok OAuth response is too large") + } + return body, response.StatusCode, nil +} + +func grokJWTSubject(token string) string { + parts := strings.Split(strings.TrimSpace(token), ".") + if len(parts) < 2 { + return "" + } + payload, err := base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + return "" + } + defer clear(payload) + var claims struct { + Subject string `json:"sub"` + } + if json.Unmarshal(payload, &claims) != nil { + return "" + } + return strings.TrimSpace(claims.Subject) +} + +func firstNonEmptyString(values ...string) string { + for _, value := range values { + if trimmed := strings.TrimSpace(value); trimmed != "" { + return trimmed + } + } + return "" +} + +func boundedGrokOAuthCode(value string) string { + value = strings.TrimSpace(value) + if value == "" || len(value) > 128 || strings.ContainsAny(value, "\r\n\x00") { + return "" + } + return value +} diff --git a/third_party/cpaembedded/embedded/grok_test.go b/third_party/cpaembedded/embedded/grok_test.go new file mode 100644 index 00000000..c7de04aa --- /dev/null +++ b/third_party/cpaembedded/embedded/grok_test.go @@ -0,0 +1,279 @@ +package embedded + +import ( + "encoding/base64" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" +) + +func TestBeginGrokDeviceAuthorizationUsesDiscoveryAndReturnsOneChallenge(t *testing.T) { + now := time.Date(2026, time.August, 18, 8, 0, 0, 0, time.UTC) + var deviceCalls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/discovery": + _ = json.NewEncoder(w).Encode(map[string]string{ + "device_authorization_endpoint": serverURL(r) + "/device", + "token_endpoint": serverURL(r) + "/token", + "userinfo_endpoint": serverURL(r) + "/userinfo", + }) + case "/device": + deviceCalls.Add(1) + if err := r.ParseForm(); err != nil { + t.Fatal(err) + } + if r.Form.Get("client_id") == "" || !strings.Contains(r.Form.Get("scope"), "offline_access") { + t.Fatalf("device form = %v", r.Form) + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "device_code": "device-secret", "user_code": "ABCD-EFGH", + "verification_uri": "https://auth.x.ai/device", + "verification_uri_complete": "https://auth.x.ai/device?user_code=ABCD-EFGH", + "expires_in": 1200, "interval": 5, + }) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + challenge, err := BeginGrokDeviceAuthorization(t.Context(), GrokOptions{ + DiscoveryURL: server.URL + "/discovery", HTTPClient: server.Client(), Now: func() time.Time { return now }, + }) + if err != nil { + t.Fatal(err) + } + if challenge.VerificationURL != "https://auth.x.ai/device?user_code=ABCD-EFGH" || challenge.UserCode != "ABCD-EFGH" || + challenge.ExpiresAt != now.Add(20*time.Minute) || challenge.PollInterval != 5*time.Second || + challenge.State.DeviceCode != "device-secret" || deviceCalls.Load() != 1 { + t.Fatalf("challenge = %#v, calls = %d", challenge, deviceCalls.Load()) + } +} + +func TestParseGrokCredentialRejectsUntrustedTokenEndpoint(t *testing.T) { + raw := []byte(`{ + "type":"grok","access_token":"access","refresh_token":"refresh", + "account_id":"account-1","email":"owner@example.com", + "expired":"2028-01-01T00:00:00Z","token_endpoint":"http://127.0.0.1/token" + }`) + if _, err := ParseGrokCredentialJSON(raw); err == nil { + t.Fatal("untrusted token endpoint was accepted") + } +} + +func TestPollGrokDeviceAuthorizationPerformsExactlyOneRequest(t *testing.T) { + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error":"slow_down"}`)) + })) + defer server.Close() + + result, err := PollGrokDeviceAuthorizationOnce(t.Context(), GrokDeviceState{ + DeviceCode: "device-secret", TokenEndpoint: server.URL, UserInfoEndpoint: server.URL + "/userinfo", + PollIntervalSeconds: 5, + }, GrokOptions{HTTPClient: server.Client()}) + if err != nil { + t.Fatal(err) + } + if result.Status != GrokDevicePending || result.PollInterval != 10*time.Second || + result.State.PollIntervalSeconds != 10 || calls.Load() != 1 { + t.Fatalf("poll = %#v, calls = %d", result, calls.Load()) + } +} + +func TestPollGrokDeviceAuthorizationMapsTerminalDeviceStates(t *testing.T) { + for _, test := range []struct { + name string + code string + status GrokDeviceStatus + }{ + {name: "denied", code: "access_denied", status: GrokDeviceDenied}, + {name: "expired", code: "expired_token", status: GrokDeviceExpired}, + } { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusBadRequest) + _ = json.NewEncoder(w).Encode(map[string]string{"error": test.code}) + })) + defer server.Close() + + result, err := PollGrokDeviceAuthorizationOnce(t.Context(), GrokDeviceState{ + DeviceCode: "device-secret", TokenEndpoint: server.URL, + UserInfoEndpoint: server.URL + "/userinfo", PollIntervalSeconds: 5, + }, GrokOptions{HTTPClient: server.Client()}) + if err != nil || result.Status != test.status { + t.Fatalf("poll = %#v, %v", result, err) + } + }) + } +} + +func TestPollGrokDeviceAuthorizationVerifiesUserInfoIdentity(t *testing.T) { + now := time.Date(2026, time.August, 18, 8, 0, 0, 0, time.UTC) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/token": + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "access-secret", "refresh_token": "refresh-secret", + "id_token": testUnsignedJWT(map[string]any{"sub": "account-1"}), + "token_type": "Bearer", "expires_in": 3600, + }) + case "/userinfo": + if r.Header.Get("Authorization") != "Bearer access-secret" { + t.Fatalf("Authorization = %q", r.Header.Get("Authorization")) + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "sub": "account-1", "email": "owner@example.com", "email_verified": true, + }) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + result, err := PollGrokDeviceAuthorizationOnce(t.Context(), GrokDeviceState{ + DeviceCode: "device-secret", TokenEndpoint: server.URL + "/token", + UserInfoEndpoint: server.URL + "/userinfo", PollIntervalSeconds: 5, + }, GrokOptions{HTTPClient: server.Client(), Now: func() time.Time { return now }}) + if err != nil { + t.Fatal(err) + } + if result.Status != GrokDeviceAuthorized || result.Credential.AccountID != "account-1" || + result.Credential.Email != "owner@example.com" || result.Credential.Type != ProviderGrok || + result.Credential.Expire != now.Add(time.Hour).Format(time.RFC3339) { + t.Fatalf("poll = %#v", result) + } +} + +func TestImportGrokCanonicalUsesAccountIDAndUpdatesEmail(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]any{ + "sub": "account-1", "email": "new@example.com", "email_verified": true, + }) + })) + defer server.Close() + raw := []byte(`{ + "type":"grok","access_token":"access","refresh_token":"refresh", + "account_id":"account-1","email":"old@example.com","expired":"2028-01-01T00:00:00Z", + "token_endpoint":"https://auth.x.ai/oauth/token" + }`) + credential, err := ImportGrokCredential(t.Context(), raw, GrokOptions{ + UserInfoURL: server.URL, HTTPClient: server.Client(), + }) + if err != nil { + t.Fatal(err) + } + if credential.AccountID != "account-1" || credential.Email != "new@example.com" { + t.Fatalf("credential = %#v", credential) + } +} + +func TestImportGrokNativeCPAFileVerifiesSubjectAndStripsControlFields(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "Bearer access" { + t.Fatalf("Authorization = %q", r.Header.Get("Authorization")) + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "sub": "account-1", "email": "owner@example.com", "email_verified": true, + }) + })) + defer server.Close() + raw := []byte(`{ + "type":"xai","access_token":"access","refresh_token":"refresh", + "sub":"account-1","email":"owner@example.com","expired":"2028-01-01T00:00:00Z", + "token_endpoint":"https://auth.x.ai/oauth2/token","base_url":"https://api.x.ai/v1", + "auth_kind":"oauth","disabled":false,"websockets":true,"prefix":"legacy","note":"backup" + }`) + credential, err := ImportGrokCredential(t.Context(), raw, GrokOptions{ + UserInfoURL: server.URL, HTTPClient: server.Client(), + }) + if err != nil { + t.Fatal(err) + } + canonical, err := MarshalGrokCredential(credential) + if err != nil { + t.Fatal(err) + } + for _, discarded := range []string{"base_url", "auth_kind", "disabled", "websockets", "prefix", "note", `"sub"`} { + if strings.Contains(string(canonical), discarded) { + t.Fatalf("canonical credential retained %q: %s", discarded, canonical) + } + } +} + +func TestRefreshGrokCredentialOncePreservesRotatingTokensAndIdentity(t *testing.T) { + now := time.Date(2026, time.August, 18, 8, 0, 0, 0, time.UTC) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/token": + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "new-access", "refresh_token": "new-refresh", + "token_type": "Bearer", "expires_in": 3600, + }) + case "/userinfo": + _ = json.NewEncoder(w).Encode(map[string]any{ + "sub": "account-1", "email": "new@example.com", "email_verified": true, + }) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + current := GrokCredential{ + Type: ProviderGrok, AccessToken: "old-access", RefreshToken: "old-refresh", IDToken: "old-id", + AccountID: "account-1", Email: "old@example.com", Expire: "2026-08-18T07:00:00Z", + TokenEndpoint: server.URL + "/token", + } + refreshed, err := RefreshGrokCredentialOnce(t.Context(), current, GrokOptions{ + UserInfoURL: server.URL + "/userinfo", HTTPClient: server.Client(), Now: func() time.Time { return now }, + }) + if err != nil { + t.Fatal(err) + } + if refreshed.AccessToken != "new-access" || refreshed.RefreshToken != "new-refresh" || + refreshed.IDToken != "old-id" || refreshed.AccountID != "account-1" || refreshed.Email != "new@example.com" { + t.Fatalf("refreshed = %#v", refreshed) + } +} + +func TestDiscoverGrokModelsUsesOAuthExecutionIdentity(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "Bearer access-secret" || + r.Header.Get("X-XAI-Token-Auth") != "xai-grok-cli" || + r.Header.Get("x-grok-client-identifier") != "grok-shell" { + t.Fatalf("model headers = %#v", r.Header) + } + _, _ = w.Write([]byte(`{"data":[{"id":"grok-4.3"},{"id":"grok-code-fast-1"},{"id":"grok-4.3"}]}`)) + })) + defer server.Close() + models, err := DiscoverGrokModels(t.Context(), GrokCredential{AccessToken: "access-secret"}, GrokOptions{ + ModelsURL: server.URL, HTTPClient: server.Client(), + }) + if err != nil { + t.Fatal(err) + } + if len(models) != 2 || models[0] != "grok-4.3" || models[1] != "grok-code-fast-1" { + t.Fatalf("models = %v", models) + } +} + +func serverURL(r *http.Request) string { + return "http://" + r.Host +} + +func testUnsignedJWT(claims map[string]any) string { + header, _ := json.Marshal(map[string]string{"alg": "none"}) + payload, _ := json.Marshal(claims) + return base64RawURL(header) + "." + base64RawURL(payload) + "." +} + +func base64RawURL(value []byte) string { + return base64.RawURLEncoding.EncodeToString(value) +} diff --git a/third_party/cpaembedded/embedded/live_test.go b/third_party/cpaembedded/embedded/live_test.go index 2d132942..71be0329 100644 --- a/third_party/cpaembedded/embedded/live_test.go +++ b/third_party/cpaembedded/embedded/live_test.go @@ -419,6 +419,122 @@ func selectLiveAntigravityModel(models []AntigravityModel) string { return models[0].ID } +// TestLiveGrokContract is opt-in because it consumes a real Grok subscription. +// It uses a prepared canonical credential and never refreshes or logs it. +func TestLiveGrokContract(t *testing.T) { + credentialFile := strings.TrimSpace(os.Getenv("CPA_LIVE_GROK_CREDENTIAL_FILE")) + model := strings.TrimSpace(os.Getenv("CPA_LIVE_GROK_MODEL")) + if credentialFile == "" { + t.Skip("CPA_LIVE_GROK_CREDENTIAL_FILE is not set") + } + raw, err := os.ReadFile(credentialFile) + if err != nil { + t.Fatalf("read live Grok credential file: %v", err) + } + credential, err := ParseGrokCredentialJSON(raw) + clear(raw) + if err != nil { + t.Fatalf("parse live Grok credential: %v", err) + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + models, err := DiscoverGrokModels(ctx, credential, GrokOptions{}) + if err != nil { + t.Fatalf("discover live Grok models: %v", err) + } + if len(models) == 0 { + t.Fatal("live Grok model list is empty") + } + if model == "" { + model = models[0] + } + payloads := []struct { + name string + format string + body map[string]any + }{ + {name: "OpenAI Responses", format: "openai-response", body: map[string]any{ + "model": model, "input": "Reply exactly OK.", "store": false, + }}, + {name: "OpenAI Chat Completions", format: "openai", body: map[string]any{ + "model": model, "messages": []map[string]any{{"role": "user", "content": "Reply exactly OK."}}, + }}, + {name: "Anthropic Messages", format: "claude", body: map[string]any{ + "model": model, "max_tokens": 16, + "messages": []map[string]any{{"role": "user", "content": "Reply exactly OK."}}, + }}, + {name: "Gemini GenerateContent", format: "gemini", body: map[string]any{ + "contents": []map[string]any{{"role": "user", "parts": []map[string]any{{"text": "Reply exactly OK."}}}}, + }}, + } + executor := NewGrokHTTPExecutor() + for _, request := range payloads { + payload, marshalErr := json.Marshal(request.body) + if marshalErr != nil { + t.Fatal(marshalErr) + } + t.Run(request.name+" unary", func(t *testing.T) { + response, executeErr := executor.ExecuteCanonical(ctx, "live-grok-contract", credential, ExecuteRequest{ + Model: model, Format: request.format, Payload: payload, OriginalRequest: payload, + }) + if executeErr != nil { + t.Fatalf("execute live request: %v", executeErr) + } + if !json.Valid(response.Payload) { + t.Fatal("live response is not JSON") + } + }) + t.Run(request.name+" stream", func(t *testing.T) { + streamBody := make(map[string]any, len(request.body)+1) + for key, value := range request.body { + streamBody[key] = value + } + streamBody["stream"] = true + streamPayload, marshalErr := json.Marshal(streamBody) + if marshalErr != nil { + t.Fatal(marshalErr) + } + stream, streamErr := executor.ExecuteStreamCanonical(ctx, "live-grok-contract", credential, ExecuteRequest{ + Model: model, Format: request.format, Payload: streamPayload, OriginalRequest: streamPayload, + }) + if streamErr != nil { + t.Fatalf("start live stream: %v", streamErr) + } + chunks := 0 + for chunk := range stream.Chunks { + if chunk.Err != nil { + t.Fatalf("read live stream: %v", chunk.Err) + } + if len(chunk.Payload) > 0 { + chunks++ + } + } + if chunks == 0 { + t.Fatal("live stream returned no chunks") + } + }) + } + for _, request := range []struct { + format string + body map[string]any + }{ + {format: "openai-response", body: payloads[0].body}, + {format: "claude", body: payloads[2].body}, + {format: "gemini", body: payloads[3].body}, + } { + payload, marshalErr := json.Marshal(request.body) + if marshalErr != nil { + t.Fatal(marshalErr) + } + response, countErr := executor.CountTokensCanonical(ctx, ExecuteRequest{ + Model: model, Format: request.format, Payload: payload, OriginalRequest: payload, + }) + if countErr != nil || !json.Valid(response.Payload) { + t.Fatalf("count live Grok tokens: %v / %s", countErr, response.Payload) + } + } +} + func selectLiveClaudeModel(models []ClaudeModel) string { for _, preferred := range []string{"claude-sonnet-4-6", "claude-opus-4-6", "claude-haiku-4-5-20251001"} { for _, model := range models { diff --git a/web/src/app/resources/channels.ts b/web/src/app/resources/channels.ts index db073d83..a6b92f52 100644 --- a/web/src/app/resources/channels.ts +++ b/web/src/app/resources/channels.ts @@ -18,7 +18,7 @@ import { export type ChannelFieldInputKind = 'text' | 'url' | 'secret' export type ChannelConnectionType = 'api_key' | 'subscription' -export type ChannelAuthorizationMethod = 'browser_oauth' | 'oauth_file' +export type ChannelAuthorizationMethod = 'browser_oauth' | 'device_oauth' | 'oauth_file' export type ChannelCredentialAction = 'reset_credit' export type ChannelNoticeID = 'claude_oauth_risk' | 'antigravity_oauth_risk' export type ChannelNoticeTone = 'warning' diff --git a/web/src/app/resources/credential-stages.ts b/web/src/app/resources/credential-stages.ts index 6674e844..59aea29f 100644 --- a/web/src/app/resources/credential-stages.ts +++ b/web/src/app/resources/credential-stages.ts @@ -30,8 +30,11 @@ export interface CredentialStageAccount { export interface CredentialStage { stage_id: string status: CredentialStageStatus + authorization_method?: 'browser_oauth' | 'device_oauth' | 'oauth_file' authorization_url?: string redirect_uri?: string + user_code?: string + next_poll_at_ms?: number account: CredentialStageAccount expires_at_ms: number error_code?: string @@ -46,8 +49,11 @@ export interface CredentialConnectResult { const stageFields = [ 'stage_id', 'status', + 'authorization_method', 'authorization_url', 'redirect_uri', + 'user_code', + 'next_poll_at_ms', 'account', 'expires_at_ms', 'error_code', @@ -63,6 +69,7 @@ const stageStatuses = [ 'expired', 'outcome_unknown', ] as const +const authorizationMethods = ['browser_oauth', 'device_oauth', 'oauth_file'] as const function invalidResponse(): never { throw new InvalidResponseError() @@ -105,11 +112,27 @@ export function projectCredentialStage(value: unknown): CredentialStage { record.authorization_url === undefined ? undefined : projectHTTPURL(record.authorization_url) const redirectURI = record.redirect_uri === undefined ? undefined : projectHTTPURL(record.redirect_uri) + const authorizationMethod = + record.authorization_method === undefined + ? undefined + : projectEnum(record.authorization_method, authorizationMethods) + const userCode = record.user_code === undefined ? undefined : projectString(record.user_code) + if ( + userCode !== undefined && + (!/^[\x21-\x7e ]{1,128}$/u.test(userCode) || userCode.trim() !== userCode) + ) { + invalidResponse() + } return { stage_id: projectStageID(record.stage_id), status: projectEnum(record.status, stageStatuses), + ...(authorizationMethod === undefined ? {} : { authorization_method: authorizationMethod }), ...(authorizationURL === undefined ? {} : { authorization_url: authorizationURL }), ...(redirectURI === undefined ? {} : { redirect_uri: redirectURI }), + ...(userCode === undefined ? {} : { user_code: userCode }), + ...(record.next_poll_at_ms === undefined + ? {} + : { next_poll_at_ms: projectEpochMilliseconds(record.next_poll_at_ms) }), account: projectAccount(record.account), expires_at_ms: projectEpochMilliseconds(record.expires_at_ms), ...(record.error_code === undefined @@ -158,6 +181,20 @@ export async function completeCredentialAuthorization( ) } +export async function pollCredentialDeviceAuthorization( + client: ApiClient, + stageID: string, + signal?: AbortSignal, +): Promise { + const id = projectStageID(stageID) + return projectCredentialStage( + await client.request(`/api/credential-stages/${id}/device-poll`, { + method: 'POST', + signal, + }), + ) +} + export async function importCredentialStage( client: ApiClient, channelID: string, diff --git a/web/src/features/import/SubscriptionCredentialStager.vue b/web/src/features/import/SubscriptionCredentialStager.vue index c21beed7..49623a7b 100644 --- a/web/src/features/import/SubscriptionCredentialStager.vue +++ b/web/src/features/import/SubscriptionCredentialStager.vue @@ -10,6 +10,7 @@ import { completeCredentialAuthorization, getCredentialStage, importCredentialStage, + pollCredentialDeviceAuthorization, type CredentialStage, } from '@/app/resources/credential-stages' import type { ChannelAuthorizationMethod, ChannelNoticeDto } from '@/app/resources/channels' @@ -76,8 +77,12 @@ const readyCount = computed( const hasAccounts = computed(() => props.modelValue.length > 0) const entryBusy = computed(() => props.disabled || Boolean(busyAction.value)) const supportsBrowserOAuth = computed(() => props.authorizationMethods.includes('browser_oauth')) +const supportsDeviceOAuth = computed(() => props.authorizationMethods.includes('device_oauth')) const supportsOAuthFile = computed(() => props.authorizationMethods.includes('oauth_file')) -const hasEntryMethod = computed(() => supportsBrowserOAuth.value || supportsOAuthFile.value) +const supportsInteractiveOAuth = computed( + () => supportsBrowserOAuth.value || supportsDeviceOAuth.value, +) +const hasEntryMethod = computed(() => supportsInteractiveOAuth.value || supportsOAuthFile.value) const channelLabel = computed(() => props.channelName.trim() || props.channelId) function replaceStage(stage: CredentialStage): void { @@ -87,6 +92,9 @@ function replaceStage(stage: CredentialStage): void { ...stage, authorization_url: stage.authorization_url ?? existing.authorization_url, redirect_uri: stage.redirect_uri ?? existing.redirect_uri, + authorization_method: stage.authorization_method ?? existing.authorization_method, + user_code: stage.user_code ?? existing.user_code, + next_poll_at_ms: stage.next_poll_at_ms ?? existing.next_poll_at_ms, } : stage const index = props.modelValue.findIndex((item) => item.stage_id === stage.stage_id) @@ -174,8 +182,15 @@ function schedulePoll(stage: CredentialStage): void { const poll = async () => { const controller = new AbortController() state.controller = controller + let polledStage: CredentialStage | undefined try { - const next = await getCredentialStage(client, stage.stage_id, controller.signal) + const current = props.modelValue.find((item) => item.stage_id === stage.stage_id) ?? stage + const next = + current.authorization_method === 'device_oauth' && + current.status === 'pending_authorization' + ? await pollCredentialDeviceAuthorization(client, stage.stage_id, controller.signal) + : await getCredentialStage(client, stage.stage_id, controller.signal) + polledStage = next state.failures = 0 if (feedbackKey.value === 'import.subscription.pollFailed') feedbackKey.value = '' replaceStage(next) @@ -199,10 +214,23 @@ function schedulePoll(stage: CredentialStage): void { state.controller = undefined } if (polling.has(stage.stage_id)) { - state.timer = window.setTimeout(poll, 1_200 * Math.min(state.failures + 1, 4)) + const latest = + polledStage ?? props.modelValue.find((item) => item.stage_id === stage.stage_id) + const providerDelay = + latest?.authorization_method === 'device_oauth' && latest.next_poll_at_ms + ? Math.max(250, latest.next_poll_at_ms - Date.now()) + : 1_200 + state.timer = window.setTimeout( + poll, + state.failures ? 1_200 * Math.min(state.failures + 1, 4) : providerDelay, + ) } } - state.timer = window.setTimeout(poll, 800) + const initialDelay = + stage.authorization_method === 'device_oauth' && stage.next_poll_at_ms + ? Math.max(250, stage.next_poll_at_ms - Date.now()) + : 800 + state.timer = window.setTimeout(poll, initialDelay) } watch( @@ -229,7 +257,7 @@ function openAuthorizationPopup(): Window | null { } async function beginAuthorization(existingPopup?: Window | null): Promise { - if (!supportsBrowserOAuth.value || props.disabled || busyAction.value) { + if (!supportsInteractiveOAuth.value || props.disabled || busyAction.value) { existingPopup?.close() return } @@ -324,7 +352,7 @@ function callbackPlaceholder(stage: CredentialStage): string { } async function restartAuthorization(stage: CredentialStage): Promise { - if (!supportsBrowserOAuth.value || props.disabled || busyAction.value) return + if (!supportsInteractiveOAuth.value || props.disabled || busyAction.value) return const popup = openAuthorizationPopup() await removeStage(stage) await nextTick() @@ -380,6 +408,7 @@ function stageErrorKey(code: string): string { const known: Readonly> = { authorization_denied: 'import.subscription.stageError.authorizationDenied', authorization_failed: 'import.subscription.stageError.authorizationFailed', + authorization_expired: 'import.subscription.stageError.authorizationExpired', authorization_exchange_rejected: 'import.subscription.stageError.exchangeRejected', authorization_exchange_unknown: 'import.subscription.stageError.exchangeUnknown', authorization_exchange_interrupted: 'import.subscription.stageError.exchangeInterrupted', @@ -464,10 +493,18 @@ onBeforeUnmount(() => { {{ stage.status === 'exchanging' ? t('import.subscription.exchanging') - : t('import.subscription.waiting') + : stage.authorization_method === 'device_oauth' + ? t('import.subscription.deviceWaiting') + : t('import.subscription.waiting') }} - {{ t('import.subscription.waitingHelp', { channel: channelLabel }) }} + + {{ + stage.authorization_method === 'device_oauth' + ? t('import.subscription.deviceWaitingHelp', { channel: channelLabel }) + : t('import.subscription.waitingHelp', { channel: channelLabel }) + }} + { -
+
+ + {{ t('import.subscription.userCode') }} + +
+ {{ stage.user_code }} + +
+
+ + { {
-
+
{
{ font-weight: 600; text-decoration: none; } +.subscription-stager__device-code { + display: grid; + grid-template-columns: minmax(0, max-content) auto; + align-items: center; + justify-content: start; + gap: var(--space-1); +} +.subscription-stager__device-code code { + border: 1px solid var(--color-border-control); + border-radius: var(--radius-control); + background: var(--color-surface-sunken); + color: var(--color-text); + padding: 9px 12px; + font-family: var(--font-mono); + font-size: var(--text-body); + font-weight: 700; + letter-spacing: 0.08em; + white-space: nowrap; +} .subscription-stager__callback { display: grid; grid-template-columns: minmax(0, 1fr) auto; diff --git a/web/src/i18n/locales/en-US/import.ts b/web/src/i18n/locales/en-US/import.ts index 6defa23d..7a2fdc47 100644 --- a/web/src/i18n/locales/en-US/import.ts +++ b/web/src/i18n/locales/en-US/import.ts @@ -129,6 +129,9 @@ export default { waiting: 'Waiting for the browser to finish authorization', waitingHelp: '{channel} opened in a new window; this page continues automatically when it is done', + deviceWaiting: 'Waiting for device authorization', + deviceWaitingHelp: + 'Enter the user code below on the {channel} authorization page; this page continues automatically', exchanging: 'Finishing authorization…', sessionRemaining: 'Authorization session time left', restart: 'Start a new authorization', @@ -138,9 +141,13 @@ export default { }, manualHint: 'A local deployment finishes on its own. For a remote deployment where the browser cannot reach {redirectUri}, authorize with the link below and paste back the final browser URL.', + deviceInstructions: + 'Open the authorization link and enter the user code. Keep this page open; it continues automatically when authorization finishes.', authorizationLink: 'Authorization link', copyAuthorization: 'Copy link', openAuthorization: 'Open', + userCode: 'User code', + copyUserCode: 'Copy user code', callbackLabel: 'URL your browser ended on', callbackEndpointFallback: 'the local OAuth callback', callbackPlaceholder: 'http://localhost:/callback?code=...&state=...', @@ -171,6 +178,7 @@ export default { stageError: { authorizationDenied: '{channel} authorization was cancelled; start again', authorizationFailed: '{channel} could not finish authorization; start again', + authorizationExpired: 'The device authorization code expired; start again', exchangeRejected: '{channel} rejected the authorization code; authorize again', exchangeUnknown: 'The authorization exchange outcome is unknown; authorize again', exchangeInterrupted: 'The authorization exchange was interrupted; authorize again', diff --git a/web/src/i18n/locales/ja-JP/import.ts b/web/src/i18n/locales/ja-JP/import.ts index 1e7d41ce..e46c85e9 100644 --- a/web/src/i18n/locales/ja-JP/import.ts +++ b/web/src/i18n/locales/ja-JP/import.ts @@ -126,6 +126,9 @@ export default { '新しいウィンドウがブロックされました。下の認証リンクから手動で開いてください。', waiting: 'ブラウザーでの認証完了を待っています', waitingHelp: '新しいウィンドウで {channel} を開きました。完了すると自動で続行します', + deviceWaiting: 'デバイス認証を待っています', + deviceWaitingHelp: + '{channel} の認証ページで下のユーザーコードを入力してください。完了すると自動で続行します', exchanging: '認証を完了しています…', sessionRemaining: '認証セッションの残り時間', restart: '認証をやり直す', @@ -135,9 +138,13 @@ export default { }, manualHint: 'ローカル配置では自動的に完了します。リモート配置でブラウザーから {redirectUri} に到達できない場合は、下のリンクで認証し、ブラウザーが最後に表示した完全な URL を貼り付けてください。', + deviceInstructions: + '認証リンクを開いてユーザーコードを入力してください。このページを開いたままにすると、認証完了後に自動で続行します。', authorizationLink: '認証リンク', copyAuthorization: 'リンクをコピー', openAuthorization: '開く', + userCode: 'ユーザーコード', + copyUserCode: 'ユーザーコードをコピー', callbackLabel: 'ブラウザーが最後に表示した URL', callbackEndpointFallback: 'ローカル OAuth コールバック', callbackPlaceholder: 'http://localhost:/callback?code=...&state=...', @@ -166,6 +173,8 @@ export default { stageError: { authorizationDenied: '{channel} 認証がキャンセルされました。最初からやり直してください', authorizationFailed: '{channel} 認証を完了できませんでした。最初からやり直してください', + authorizationExpired: + 'デバイス認証コードの有効期限が切れました。最初からやり直してください', exchangeRejected: '{channel} が認証コードを拒否しました。再認証してください', exchangeUnknown: '認証交換の結果を確認できません。再認証してください', exchangeInterrupted: '認証交換が中断されました。再認証してください', diff --git a/web/src/i18n/locales/zh-CN/import.ts b/web/src/i18n/locales/zh-CN/import.ts index 041e9d28..a0e37405 100644 --- a/web/src/i18n/locales/zh-CN/import.ts +++ b/web/src/i18n/locales/zh-CN/import.ts @@ -119,6 +119,8 @@ export default { popupBlocked: '新窗口被浏览器拦截了，请用下面的授权链接手动打开', waiting: '等待浏览器完成授权', waitingHelp: '已在新窗口打开 {channel}，完成后这里会自动继续', + deviceWaiting: '等待设备授权', + deviceWaitingHelp: '请在 {channel} 授权页输入下方用户码，完成后这里会自动继续', exchanging: '正在完成授权…', sessionRemaining: '授权会话剩余时间', restart: '重新开始一次授权', @@ -128,9 +130,12 @@ export default { }, manualHint: '本机部署会自动完成。远程部署若浏览器无法访问 {redirectUri}，请打开下面的链接授权，再粘贴浏览器最终停留的完整网址。', + deviceInstructions: '打开授权链接并输入用户码；请保持此页面开启，授权完成后会自动继续。', authorizationLink: '授权链接', copyAuthorization: '复制链接', openAuthorization: '打开', + userCode: '用户码', + copyUserCode: '复制用户码', callbackLabel: '授权后浏览器停在的网址', callbackEndpointFallback: '本地 OAuth 回调地址', callbackPlaceholder: 'http://localhost:<端口>/callback?code=...&state=...', @@ -158,6 +163,7 @@ export default { stageError: { authorizationDenied: '用户取消了 {channel} 授权，请重新开始', authorizationFailed: '{channel} 未能完成授权，请重新开始', + authorizationExpired: '设备授权码已过期，请重新开始', exchangeRejected: '{channel} 拒绝了授权码，请重新授权', exchangeUnknown: '授权交换结果无法确认，请重新授权', exchangeInterrupted: '授权交换被中断，请重新授权',