diff --git a/docs/oauth-subscriptions.md b/docs/oauth-subscriptions.md index 6bb70c98d..40d1fce07 100644 --- a/docs/oauth-subscriptions.md +++ b/docs/oauth-subscriptions.md @@ -37,7 +37,7 @@ before. Tokens are stored 0600 (or the OS keyring with Running `/provider` opens a **"How do you want to connect?"** chooser: ```text -❯ Sign in with OAuth One-click browser login (OpenRouter, xAI, ChatGPT, Hugging Face) +❯ Sign in with OAuth No API key to copy — browser login (OpenRouter, xAI, ChatGPT, Hugging Face) or device code (xAI, Kimi Code, Hugging Face) Paste an API key / browse providers Any of 20+ providers, local, or a proxy ``` @@ -46,23 +46,28 @@ Pick **Sign in with OAuth** → the list of providers that do real OAuth → cho ```text ❯ OpenRouter browser sign-in · creates a key xAI (Grok) browser or device code + Kimi Code device code (managed coding endpoint) ChatGPT browser (Codex backend, ChatGPT Plus/Pro) Hugging Face browser or device code ``` - **OpenRouter / xAI / ChatGPT / Hugging Face** are real OAuth: your browser - opens to approve → done (no key to paste). OpenRouter mints a key; xAI / - ChatGPT / Hugging Face store a refreshable bearer. Hugging Face requires a - one-time OAuth-app registration (no secret needed for "public" apps); the - preset pre-fills scopes, endpoints, and the OIDC issuer. The same chooser - appears in first-run onboarding. (xAI uses an opt-in preset — set - `ZERO_OAUTH_ALLOW_PRESETS=1` or your own `ZERO_OAUTH_XAI_*`; see below.) -- **Device code (headless / SSH):** for a provider that supports it (xAI, - Hugging Face), press **d** on the list to get a code to enter on another - device instead of opening a browser. On an SSH session or headless Linux box - (no `DISPLAY`) device code is used automatically; set `ZERO_OAUTH_DEVICE=1` - to force it anywhere. The CLI equivalent is - `zero auth login --device`. + opens to approve → done (no key to paste). OpenRouter mints a key; xAI / + ChatGPT / Hugging Face store a refreshable bearer. Hugging Face requires a + one-time OAuth-app registration (no secret needed for "public" apps); the + preset pre-fills scopes, endpoints, and the OIDC issuer. Kimi Code is also + real OAuth but has no browser flow at all — see the device-code bullet + below. The same chooser appears in first-run onboarding. (xAI and Kimi Code + use an opt-in preset — set `ZERO_OAUTH_ALLOW_PRESETS=1` or your own + `ZERO_OAUTH_XAI_*` / `ZERO_OAUTH_KIMI_CODE_*`; see below.) +- **Device code (headless / SSH):** for a provider that supports it (xAI, Kimi + Code, Hugging Face), press **d** on the list to get a code to enter on + another device instead of opening a browser. On an SSH session or headless + Linux box (no `DISPLAY`) device code is used automatically; set + `ZERO_OAUTH_DEVICE=1` to force it anywhere. The CLI equivalent is + `zero auth login --device`. (Kimi Code is **device-code only** — it + has no loopback/browser flow, so `zero auth kimi` always uses the device + path and pressing plain Enter on it in the wizard does too.) - **ChatGPT / Claude are intentionally not in this list for the proxy path** — use the dedicated `chatgpt-proxy` / `custom-anthropic-compatible` preset (see §2) for subscription-via-proxy. ChatGPT *is* a first-class OAuth @@ -83,9 +88,37 @@ Pick **Sign in with OAuth** → the list of providers that do real OAuth → cho (browser, or `--device` for headless) works one-click; the token is used directly on `api.x.ai/v1`. Without the opt-in, set `ZERO_OAUTH_XAI_CLIENT_ID` (and endpoints, or an issuer) yourself via `ZERO_OAUTH_XAI_*`. Either way the preset is - fully overridable by `ZERO_OAUTH_XAI_*` (env wins), and it requires a - SuperGrok / X Premium+ subscription; the client_id is an undocumented public - Grok-CLI client that may change without notice. + fully overridable by `ZERO_OAUTH_XAI_*` (env wins), and it requires a + SuperGrok / X Premium+ subscription; the client_id is an undocumented public + Grok-CLI client that may change without notice. +- **Kimi Code — built-in preset, device-code only** — `zero auth kimi` (or + `zero auth login kimi-code --device`) runs the RFC 8628 device-code flow + against `https://auth.kimi.com`. You approve on another device and enter the + code; the returned access token is stored and used **directly** as a bearer + on Kimi's managed coding endpoint `https://api.kimi.com/coding/v1` (an + OpenAI-compatible chat-completions endpoint) — no ID-token claim extraction + is needed. Kimi has **no browser/loopback flow**, so the device code is the + only path (it is used automatically, and `--device` is accepted but + redundant). The catalog/provider ID is `kimi-code`, not `kimi` — the + `moonshot` provider already uses `kimi` as an alias for its separate, + API-key-based endpoint, so `zero auth kimi` is CLI sugar that forwards to + `kimi-code` rather than reusing that name. Like xAI, the preset ships the + public kimi-cli client identity (`17e5f671-d194-4dfb-9706-5516cb48c098`). + Unlike the raw `ZERO_OAUTH_XAI_*`-only path, no `ZERO_OAUTH_ALLOW_PRESETS=1` + is needed for Kimi: both `zero auth kimi` and `zero auth login kimi-code` + run through the `auth login` engine, which enables presets unconditionally. + Any field is still overridable with + `ZERO_OAUTH_KIMI_CODE_*`. Kimi's backend also requires a handful of + vendor-identity `X-Msh-*` headers across all applicable OAuth and API calls + (device authorization, polling, code exchange, refresh, and managed + runtime/completions requests). These are reverse-engineered from kimi-cli, + not from public documentation — verify against a real login before relying + on this. + This is distinct from the `moonshot` catalog entry, which is the API-key path at + `https://api.moonshot.ai/v1` (set `MOONSHOT_API_KEY`). To override the managed + endpoint, set `baseURL` on the provider profile; to override the OAuth host, + set `ZERO_OAUTH_KIMI_CODE_ISSUER_URL`/`ZERO_OAUTH_KIMI_CODE_DEVICE_URL`/`ZERO_OAUTH_KIMI_CODE_TOKEN_URL` + (the provider resolves as `kimi-code`, so the env prefix is `KIMI_CODE`). - **ChatGPT (Codex) — opt-in preset** — `zero auth chatgpt` opens a browser, you approve with your ChatGPT Plus/Pro/Business/Enterprise account, and the bearer is stored. The bearer routes to `https://chatgpt.com/backend-api/codex/responses` diff --git a/internal/cli/auth.go b/internal/cli/auth.go index f3ecdcc42..a5b966471 100644 --- a/internal/cli/auth.go +++ b/internal/cli/auth.go @@ -78,6 +78,18 @@ func runAuth(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in return runAuthOpenRouter(args[1:], stdout, stderr, deps) case "chatgpt": return runAuthChatGPT(args[1:], stdout, stderr, deps) + case "kimi": + // Kimi Code is a standard device-code OAuth preset (no bespoke client + // like ChatGPT's Codex flow), so it reuses the generic `auth login` + // engine — which already opts into presets and resolves the baked-in + // kimi-code client_id/endpoints. `zero auth kimi` is sugar for + // `zero auth login kimi-code`, forwarding whatever flags the caller + // passed after "kimi" (--device, --scope, --help) through the real + // parser instead of discarding them: `zero auth kimi --help` must show + // help, not silently start a real device authorization, and an unknown + // flag or extra positional must be rejected the same way `zero auth + // login` rejects one. + return runAuthLogin(append([]string{"kimi-code"}, args[1:]...), stdout, stderr, deps) default: return writeExecUsageError(stderr, fmt.Sprintf("unknown auth subcommand %q", args[0])) } @@ -591,12 +603,15 @@ Commands: refresh [--watch] Force a token refresh (--watch keeps it fresh) openrouter Log in to OpenRouter in the browser; mints an API key chatgpt Log in to ChatGPT in the browser (Codex backend, ChatGPT Plus/Pro) + kimi Log in to Kimi Code via device code (managed coding endpoint) A provider is any OAuth 2.0 / OIDC server. "openrouter" ('zero auth openrouter') -works out of the box. "xai" ('zero auth login xai') uses a built-in preset that is -off by default — enable it with ZERO_OAUTH_ALLOW_PRESETS=1, or set the -ZERO_OAUTH_XAI_* vars yourself. "chatgpt" ('zero auth login chatgpt' or -'zero auth chatgpt') uses a fixed-port loopback flow against the Codex backend. +and "kimi-code" ('zero auth login kimi-code' or 'zero auth kimi') work out of the +box. "xai" ('zero auth login xai') uses a built-in preset that is off by +default — enable it with ZERO_OAUTH_ALLOW_PRESETS=1, or set the +ZERO_OAUTH_XAI_* vars yourself. "chatgpt" ('zero auth +login chatgpt' or 'zero auth chatgpt') uses a fixed-port loopback flow against +the Codex backend. Any preset field is overridable via the env vars below. For a custom provider named , set: ZERO_OAUTH__CLIENT_ID (required) ZERO_OAUTH__CLIENT_SECRET (optional) diff --git a/internal/config/resolver.go b/internal/config/resolver.go index 0d9b11786..7164c9c32 100644 --- a/internal/config/resolver.go +++ b/internal/config/resolver.go @@ -1088,11 +1088,12 @@ func applyCatalogDescriptor(profile *ProviderProfile, descriptor providercatalog merged[key] = value } profile.CustomHeaders = merged - } else if strings.EqualFold(strings.TrimSpace(descriptor.ID), "aimlapi") && !canonicalCatalogEndpoint { - // AIMLAPI attribution is owned by the catalog endpoint. A profile can retain - // those generated headers after its base URL is edited; strip their names - // before sending requests to an arbitrary staging/proxy host while preserving - // unrelated headers explicitly supplied by the user. + } else if len(descriptor.CustomHeaders) > 0 && !canonicalCatalogEndpoint { + // Catalog-owned headers (AIMLAPI attribution, Kimi's X-Msh-* device + // identity, etc.) are only valid against the catalog endpoint. A profile + // can retain those generated headers after its base URL is edited; strip + // their names before sending requests to an arbitrary staging/proxy host + // while preserving unrelated headers explicitly supplied by the user. for profileKey := range profile.CustomHeaders { for catalogKey := range descriptor.CustomHeaders { if strings.EqualFold(profileKey, catalogKey) { diff --git a/internal/config/resolver_test.go b/internal/config/resolver_test.go index 35c8cf872..b954d5eda 100644 --- a/internal/config/resolver_test.go +++ b/internal/config/resolver_test.go @@ -5,6 +5,7 @@ import ( "os" "path/filepath" "reflect" + "runtime" "strings" "testing" @@ -1561,6 +1562,49 @@ func TestApplyCatalogDescriptorStripsAimlapiAttributionFromRetargetedProfile(t * } } +// setKimiUserConfigRoot redirects os.UserConfigDir() to a throwaway temp dir +// before the kimi-code descriptor's RuntimeHeaders (kimiidentity.Headers) run, +// so this test never creates or touches the real kimi-device-id file. +func setKimiUserConfigRoot(t *testing.T) { + t.Helper() + root := t.TempDir() + switch runtime.GOOS { + case "windows": + t.Setenv("APPDATA", root) + case "darwin": + t.Setenv("HOME", root) + default: + t.Setenv("XDG_CONFIG_HOME", root) + } +} + +func TestApplyCatalogDescriptorStripsKimiIdentityFromRetargetedProfile(t *testing.T) { + setKimiUserConfigRoot(t) + descriptor, err := providercatalog.Require("kimi-code") + if err != nil { + t.Fatal(err) + } + profile := ProviderProfile{ + BaseURL: "https://proxy.example.test/v1", + CustomHeaders: map[string]string{ + "x-msh-platform": "kimi_code_cli", + "X-Msh-Device-Id": "persisted-device-id", + "X-Environment": "staging", + }, + } + + applyCatalogDescriptor(&profile, descriptor, true) + + for key := range profile.CustomHeaders { + if strings.HasPrefix(strings.ToLower(key), "x-msh-") { + t.Fatalf("kimi identity headers survived retargeting: %#v", profile.CustomHeaders) + } + } + if profile.CustomHeaders["X-Environment"] != "staging" { + t.Fatalf("user header was removed: %#v", profile.CustomHeaders) + } +} + func TestResolveProviderProfileParseThinkTagsFalseAlias(t *testing.T) { path := writeConfig(t, `{ "activeProvider": "custom", diff --git a/internal/kimiidentity/kimiidentity.go b/internal/kimiidentity/kimiidentity.go new file mode 100644 index 000000000..79effcfbd --- /dev/null +++ b/internal/kimiidentity/kimiidentity.go @@ -0,0 +1,305 @@ +// Package kimiidentity builds the X-Msh-* vendor-identity headers Kimi +// Code's backend requires on every request — OAuth device authorization, +// token polling, refresh, AND managed-endpoint model calls. It exists as its +// own dependency-free package because both internal/oauth (login/refresh) +// and internal/providercatalog (the kimi-code descriptor's CustomHeaders, +// applied to runtime completions) must send the SAME identity: a login +// accepted under one device identity and completions sent under another (or +// under none) is rejected by the backend. +// +// Header names and general shape are reverse-engineered from the +// open-source kimi-cli client (src/kimi_cli/auth/oauth.py, _common_headers); +// Kimi has no published public API documentation for this, so these values +// are a best-effort match, not a verified spec. +package kimiidentity + +import ( + "crypto/rand" + "fmt" + "os" + "path/filepath" + "runtime" + "strings" + "sync" + "time" +) + +// Headers returns the X-Msh-* vendor-identity headers, including the stable +// per-device identifier. +// +// X-Msh-Platform is "kimi_code_cli". That is the value Moonshot's own Kimi +// Code CLI sends (packages/oauth/src/identity.ts, KIMI_CODE_PLATFORM) as of +// its oauth package changelog entry correcting the header from an earlier +// "kimi-code-cli" typo (PR MoonshotAI/kimi-code#52, commit 064343a); the +// older, separate open-source kimi-cli client instead hardcodes "kimi_cli". +// Kimi's coding/v1 endpoint documents a client whitelist ("Kimi CLI, Claude +// Code, Roo Code, ..."); sending the wrong platform value risks the managed +// endpoint rejecting completions even after a successful login. +func Headers() map[string]string { + hostname, err := os.Hostname() + if err != nil || strings.TrimSpace(hostname) == "" { + hostname = "unknown-host" + } + return map[string]string{ + "X-Msh-Platform": "kimi_code_cli", + "X-Msh-Version": "unknown", + "X-Msh-Device-Name": asciiHeaderValue(hostname), + "X-Msh-Device-Model": asciiHeaderValue(runtime.GOOS + " " + runtime.GOARCH), + "X-Msh-Os-Version": runtime.GOOS, + "X-Msh-Device-Id": DeviceID(), + } +} + +var ( + deviceIDMu sync.Mutex + testDeviceIDPath string + cachedDeviceID = sync.OnceValue(func() string { return loadOrCreateDeviceIDAt(deviceIDPath()) }) +) + +// DeviceID returns the persistent device identifier sent as X-Msh-Device-Id. +// Kimi Code's own CLI persists this to ~/.kimi/device_id so the same value +// follows a device across logins, refreshes, and model calls; mirroring +// that, the ID is stored under the user config dir (zero/kimi-device-id) and +// minted once on first use. When the config dir is unavailable the ID is +// still stable for the life of the process. +func DeviceID() string { + deviceIDMu.Lock() + fn := cachedDeviceID + deviceIDMu.Unlock() + return fn() +} + +// ResetDeviceIDForTest resets the process-global DeviceID cache so unit tests +// can inspect or isolate device identity creation without state pollution. +func ResetDeviceIDForTest() { + deviceIDMu.Lock() + cachedDeviceID = sync.OnceValue(func() string { return loadOrCreateDeviceIDAt(deviceIDPath()) }) + deviceIDMu.Unlock() +} + +// SetDeviceIDPathForTest configures an explicit storage path for device ID in +// unit tests, overriding os.UserConfigDir(). It resets the DeviceID cache and +// returns a cleanup function that restores the original path and resets the +// cache. +func SetDeviceIDPathForTest(path string) func() { + deviceIDMu.Lock() + prevPath := testDeviceIDPath + testDeviceIDPath = path + cachedDeviceID = sync.OnceValue(func() string { return loadOrCreateDeviceIDAt(deviceIDPath()) }) + deviceIDMu.Unlock() + return func() { + deviceIDMu.Lock() + testDeviceIDPath = prevPath + cachedDeviceID = sync.OnceValue(func() string { return loadOrCreateDeviceIDAt(deviceIDPath()) }) + deviceIDMu.Unlock() + } +} + +// loadOrCreateDeviceIDAt is the real load-or-create logic behind DeviceID, +// parameterized by the storage path so tests can exercise production code +// directly (env var indirection through os.UserConfigDir is not portable to +// redirect in tests). It reads an existing UUID if present, otherwise mints +// one and persists it exclusively (see the concurrency note below). +func loadOrCreateDeviceIDAt(path string) string { + if path != "" { + if id := readValidDeviceID(path); id != "" { + return id + } + } + id := generateDeviceID() + if path == "" { + return id + } + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return id + } + // Create exclusively rather than os.WriteFile: two processes racing on + // first use must converge on the SAME id, since a login accepted under + // one device identity and completions sent under another are rejected + // by the backend. The loser reads back the winner's file instead of + // overwriting it. After O_EXCL succeeds the winner still has to write + // content, so a concurrent loser may briefly observe an empty file; + // readValidDeviceIDWithRetry waits for the UUID rather than minting a + // second divergent identity. If the winner dies mid-publish leaving an + // empty/invalid file, the next caller removes that abandoned file once + // and retries exclusive create so the identity is repaired instead of + // permanently stuck on a divergent in-memory id. + return createOrAdoptDeviceID(path, id) +} + +// createOrAdoptDeviceID publishes id at path or adopts a concurrent winner. +func createOrAdoptDeviceID(path, id string) string { + f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + if !os.IsExist(err) { + return id + } + if existingID := readValidDeviceIDWithRetry(path); existingID != "" { + return existingID + } + // path exists but never became a valid UUID (abandoned create, + // corrupt file). Repair it rather than removing it ourselves: an + // unlocked remove here could unlink another racer's just-published + // winner between our failed read and the remove call, handing that + // racer back an id that is no longer the one on disk. + return repairAbandonedDeviceID(path, id) + } + _, _ = f.WriteString(id + "\n") + _ = f.Sync() + _ = f.Close() + // Re-read so a racing repair that replaced us still converges. + if existingID := readValidDeviceID(path); existingID != "" { + return existingID + } + return id +} + +// repairAbandonedDeviceID fixes an invalid/empty device-id file left behind +// by a process that exclusive-created path and died before writing a UUID. +// Repair itself is serialized through an exclusive lock file so only one +// racing process ever removes and recreates path: without that, one process +// could unlink another's freshly published replacement and mint a second, +// divergent id, leaving the first process holding an id that stops matching +// what is actually persisted. Callers that lose the lock wait for the holder +// to publish instead of attempting their own repair. If the lock holder dies +// mid-repair without publishing or unlocking, the lock is broken and retried. +func repairAbandonedDeviceID(path, id string) string { + if existingID := readValidDeviceID(path); existingID != "" { + return existingID + } + lockPath := path + ".lock" + lock, err := os.OpenFile(lockPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + if !os.IsExist(err) { + return id + } + if existingID := readValidDeviceIDWithRetry(path); existingID != "" { + return existingID + } + // If lockPath exists but is stale (owner crashed), break lock and retry once. + if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > 2*time.Second { + _ = os.Remove(lockPath) + lock, err = os.OpenFile(lockPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + } + if err != nil { + if existingID := readValidDeviceIDWithRetry(path); existingID != "" { + return existingID + } + return id + } + } + defer func() { + _ = lock.Close() + _ = os.Remove(lockPath) + }() + + if existingID := readValidDeviceID(path); existingID != "" { + return existingID + } + tmpPath := fmt.Sprintf("%s.tmp.%d.%d", path, os.Getpid(), time.Now().UnixNano()) + if err := os.WriteFile(tmpPath, []byte(id+"\n"), 0o600); err != nil { + return id + } + defer func() { _ = os.Remove(tmpPath) }() + + _ = os.Remove(path) + if err := os.Rename(tmpPath, path); err != nil { + if errWrite := os.WriteFile(path, []byte(id+"\n"), 0o600); errWrite == nil { + return id + } + if existingID := readValidDeviceIDWithRetry(path); existingID != "" { + return existingID + } + return id + } + return id +} + +// readValidDeviceID returns a UUID from path, or "" if missing/invalid. +func readValidDeviceID(path string) string { + raw, err := os.ReadFile(path) + if err != nil { + return "" + } + if id := strings.TrimSpace(string(raw)); isUUID(id) { + return id + } + return "" +} + +// readValidDeviceIDWithRetry re-reads path briefly so a process that lost the +// exclusive create can adopt the winner even if it observed the file before +// the winner finished writing the UUID. +func readValidDeviceIDWithRetry(path string) string { + const attempts = 40 + const delay = 5 * time.Millisecond + for i := 0; i < attempts; i++ { + if id := readValidDeviceID(path); id != "" { + return id + } + time.Sleep(delay) + } + return "" +} + +func deviceIDPath() string { + deviceIDMu.Lock() + testPath := testDeviceIDPath + deviceIDMu.Unlock() + if testPath != "" { + return testPath + } + configDir, err := os.UserConfigDir() + if err != nil || strings.TrimSpace(configDir) == "" { + return "" + } + return filepath.Join(configDir, "zero", "kimi-device-id") +} + +func generateDeviceID() string { + raw := make([]byte, 16) + if _, err := rand.Read(raw); err != nil { + return "00000000-0000-0000-0000-000000000000" + } + raw[6] = (raw[6] & 0x0f) | 0x40 // version 4 + raw[8] = (raw[8] & 0x3f) | 0x80 // variant 10 + return fmt.Sprintf("%x-%x-%x-%x-%x", raw[0:4], raw[4:6], raw[6:8], raw[8:10], raw[10:16]) +} + +func isUUID(s string) bool { + if len(s) != 36 { + return false + } + for i, r := range s { + switch i { + case 8, 13, 18, 23: + if r != '-' { + return false + } + default: + if (r < '0' || r > '9') && (r < 'a' || r > 'f') && (r < 'A' || r > 'F') { + return false + } + } + } + return true +} + +// asciiHeaderValue strips anything outside printable ASCII (0x20-0x7e). This +// mirrors a defensive fix kimi-cli itself needed: a raw platform-version +// string containing "#" broke an HTTP client's header validation on Linux +// (MoonshotAI/kimi-cli#1169) because HTTP header values must not contain +// control characters. +func asciiHeaderValue(s string) string { + var b strings.Builder + for _, r := range s { + if r >= 0x20 && r <= 0x7e { + b.WriteRune(r) + } + } + clean := strings.TrimSpace(b.String()) + if clean == "" { + return "unknown" + } + return clean +} diff --git a/internal/kimiidentity/kimiidentity_test.go b/internal/kimiidentity/kimiidentity_test.go new file mode 100644 index 000000000..024b0d7f5 --- /dev/null +++ b/internal/kimiidentity/kimiidentity_test.go @@ -0,0 +1,301 @@ +package kimiidentity + +import ( + "os" + "path/filepath" + "runtime" + "strings" + "sync" + "testing" + "time" +) + +// setUserConfigRoot redirects os.UserConfigDir() to a throwaway temp dir so a +// test that calls the process-global Headers()/DeviceID() never creates or +// touches the real ~/.config/zero (or AppData\zero) kimi-device-id file. +// Mirrors internal/workspacetrust/trust_test.go: os.UserConfigDir reads +// APPDATA on Windows, HOME on darwin, and XDG_CONFIG_HOME on Linux, so a +// single env var isn't portable. Must be called before the first call to +// Headers()/DeviceID() in the process — DeviceID is a sync.OnceValue, so a +// call before the redirect is in place would permanently cache the real path. +func setUserConfigRoot(t *testing.T) { + t.Helper() + root := t.TempDir() + cleanup := SetDeviceIDPathForTest(filepath.Join(root, "zero", "kimi-device-id")) + t.Cleanup(cleanup) + switch runtime.GOOS { + case "windows": + t.Setenv("APPDATA", root) + case "darwin": + t.Setenv("HOME", root) + default: + t.Setenv("XDG_CONFIG_HOME", root) + } +} + +func TestHeadersIncludesDeviceIdentity(t *testing.T) { + setUserConfigRoot(t) + headers := Headers() + for _, key := range []string{ + "X-Msh-Platform", + "X-Msh-Version", + "X-Msh-Device-Name", + "X-Msh-Device-Model", + "X-Msh-Os-Version", + "X-Msh-Device-Id", + } { + if headers[key] == "" { + t.Fatalf("Headers()[%q] empty", key) + } + } + if headers["X-Msh-Platform"] != "kimi_code_cli" { + t.Fatalf("X-Msh-Platform = %q, want kimi_code_cli", headers["X-Msh-Platform"]) + } + if !isUUID(headers["X-Msh-Device-Id"]) { + t.Fatalf("X-Msh-Device-Id = %q, want UUID", headers["X-Msh-Device-Id"]) + } +} + +func TestLoadOrCreateDeviceIDExclusiveCreate(t *testing.T) { + // Exercise the production loader directly via its path-parameterized + // helper. Concurrent first-use must converge on a single persisted ID: + // the O_EXCL loser reads back the winner's file instead of overwriting it. + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + + const workers = 8 + ids := make([]string, workers) + var wg sync.WaitGroup + wg.Add(workers) + for i := range workers { + go func(i int) { + defer wg.Done() + ids[i] = loadOrCreateDeviceIDAt(path) + }(i) + } + wg.Wait() + + winner := "" + for _, id := range ids { + if id == "" { + t.Fatal("worker returned empty id") + } + if winner == "" { + winner = id + continue + } + if id != winner { + t.Fatalf("workers diverged: got %q and %q", winner, id) + } + } + if !isUUID(winner) { + t.Fatalf("winner id %q is not a UUID", winner) + } + // The persisted file carries the winner exactly once. + if raw, err := os.ReadFile(path); err != nil { + t.Fatalf("read persisted id: %v", err) + } else if got := strings.TrimSpace(string(raw)); got != winner { + t.Fatalf("persisted id = %q, want %q", got, winner) + } +} + +func TestLoadOrCreateDeviceIDReadsExisting(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + const existing = "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee" + if err := os.WriteFile(path, []byte(existing+"\n"), 0o600); err != nil { + t.Fatal(err) + } + if got := loadOrCreateDeviceIDAt(path); got != existing { + t.Fatalf("loadOrCreateDeviceIDAt = %q, want existing %q", got, existing) + } +} + +// TestLoadOrCreateDeviceIDAdoptsWinnerAfterEmptyCreate covers the +// multi-process window where the O_EXCL winner has created the file but not +// yet written the UUID. Concurrent callers must wait and adopt that UUID +// rather than each minting a divergent identity. +func TestLoadOrCreateDeviceIDAdoptsWinnerAfterEmptyCreate(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + // Simulate the exclusive-create winner that has not written yet. + f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + t.Fatal(err) + } + const winner = "11111111-2222-4333-8444-555555555555" + done := make(chan struct{}) + go func() { + defer close(done) + time.Sleep(30 * time.Millisecond) + _, _ = f.WriteString(winner + "\n") + _ = f.Sync() + _ = f.Close() + }() + + const workers = 4 + ids := make([]string, workers) + var wg sync.WaitGroup + wg.Add(workers) + for i := range workers { + go func(i int) { + defer wg.Done() + ids[i] = loadOrCreateDeviceIDAt(path) + }(i) + } + wg.Wait() + <-done + + for _, id := range ids { + if id != winner { + t.Fatalf("worker returned %q, want winner %q (all: %v)", id, winner, ids) + } + } +} + +// TestLoadOrCreateDeviceIDRepairsAbandonedEmptyFile covers the case where +// a previous process exclusive-created the path and died before writing a +// UUID. Callers must not permanently diverge: after the retry window the +// empty file is removed and a new exclusive create publishes a valid id. +func TestLoadOrCreateDeviceIDRepairsAbandonedEmptyFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + t.Fatal(err) + } + _ = f.Close() // abandoned: never written + + got := loadOrCreateDeviceIDAt(path) + if !isUUID(got) { + t.Fatalf("repaired id %q is not a UUID", got) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read repaired file: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != got { + t.Fatalf("persisted %q, want repaired %q", persisted, got) + } +} + +// TestLoadOrCreateDeviceIDConcurrentAbandonedFileRepairConverges covers +// multiple racing processes all finding the same abandoned/invalid file at +// once. Repair must be mutually exclusive: only one racer may remove and +// recreate the file, so every caller ends up with the same id and that id is +// exactly what is persisted (no caller returns an in-memory id that a later +// repair silently unlinked and replaced). +func TestLoadOrCreateDeviceIDConcurrentAbandonedFileRepairConverges(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + t.Fatal(err) + } + _ = f.Close() // abandoned: never written + + const workers = 16 + ids := make([]string, workers) + var wg sync.WaitGroup + wg.Add(workers) + for i := range workers { + go func(i int) { + defer wg.Done() + ids[i] = loadOrCreateDeviceIDAt(path) + }(i) + } + wg.Wait() + + winner := "" + for _, id := range ids { + if id == "" { + t.Fatal("worker returned empty id") + } + if winner == "" { + winner = id + continue + } + if id != winner { + t.Fatalf("workers diverged repairing abandoned file: got %q and %q", winner, id) + } + } + if !isUUID(winner) { + t.Fatalf("winner id %q is not a UUID", winner) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read persisted id: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != winner { + t.Fatalf("persisted %q, want winner %q", persisted, winner) + } + if _, err := os.Stat(path + ".lock"); !os.IsNotExist(err) { + t.Fatalf("repair lock file left behind: err=%v", err) + } +} + +func TestLoadOrCreateDeviceIDRepairsStaleRepairLock(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + t.Fatal(err) + } + _ = f.Close() // abandoned target file + + lockPath := path + ".lock" + lockF, err := os.OpenFile(lockPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + t.Fatal(err) + } + _ = lockF.Close() // abandoned lock file + // Backdate lock file mtime to make it stale (> 1s) + oldTime := time.Now().Add(-5 * time.Second) + if err := os.Chtimes(lockPath, oldTime, oldTime); err != nil { + t.Fatal(err) + } + + got := loadOrCreateDeviceIDAt(path) + if !isUUID(got) { + t.Fatalf("repaired id %q is not a UUID", got) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read repaired file: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != got { + t.Fatalf("persisted %q, want repaired %q", persisted, got) + } +} + +func TestAsciiHeaderValueStripsNonPrintable(t *testing.T) { + if got := asciiHeaderValue("linux#6.1"); got != "linux#6.1" { + // printable ASCII including # is kept; the kimi-cli bug was a different + // control character path — ensure we still strip true controls. + t.Fatalf("got %q", got) + } + if got := asciiHeaderValue("a\nb\x00c"); got != "abc" { + t.Fatalf("got %q, want abc", got) + } + if got := asciiHeaderValue("\x01\x02"); got != "unknown" { + t.Fatalf("got %q, want unknown", got) + } +} diff --git a/internal/oauth/device.go b/internal/oauth/device.go index 697de057c..038ff6b2c 100644 --- a/internal/oauth/device.go +++ b/internal/oauth/device.go @@ -71,6 +71,7 @@ func RequestDeviceCode(ctx context.Context, client *http.Client, cfg Config, now } request.Header.Set("Content-Type", "application/x-www-form-urlencoded") request.Header.Set("Accept", "application/json") + applyExtraHeaders(request, cfg.ExtraHeaders) response, err := client.Do(request) if err != nil { return DeviceAuth{}, fmt.Errorf("oauth: device authorization request failed: %w", err) @@ -173,6 +174,7 @@ func pollDeviceOnce(ctx context.Context, client *http.Client, cfg Config, device } request.Header.Set("Content-Type", "application/x-www-form-urlencoded") request.Header.Set("Accept", "application/json") + applyExtraHeaders(request, cfg.ExtraHeaders) response, err := client.Do(request) if err != nil { return Token{}, fmt.Errorf("oauth: device token poll failed: %w", err) diff --git a/internal/oauth/flow.go b/internal/oauth/flow.go index 9ff94cccd..09cc68e60 100644 --- a/internal/oauth/flow.go +++ b/internal/oauth/flow.go @@ -154,7 +154,7 @@ func ExchangeCode(ctx context.Context, client *http.Client, cfg Config, code, ve if secret := trimmed(cfg.ClientSecret); secret != "" { form.Set("client_secret", secret) } - return PostToken(ctx, client, cfg.TokenEndpoint, form, Token{Scopes: cfg.Scopes}, now) + return PostToken(ctx, client, cfg.TokenEndpoint, form, Token{Scopes: cfg.Scopes}, cfg.ExtraHeaders, now) } // Refresh exchanges a refresh token for a fresh access token. A response that @@ -181,7 +181,7 @@ func Refresh(ctx context.Context, client *http.Client, cfg Config, current Token // and PostToken only overwrites TokenType when the response supplies one, so // without seeding it here the type would be silently lost across refreshes (L15). base := Token{Scopes: current.Scopes, RefreshToken: refresh, Account: current.Account, IDToken: current.IDToken, TokenType: current.TokenType} - return PostToken(ctx, client, cfg.TokenEndpoint, form, base, now) + return PostToken(ctx, client, cfg.TokenEndpoint, form, base, cfg.ExtraHeaders, now) } // PostToken performs a token-endpoint POST and maps the response onto a Token. @@ -189,7 +189,7 @@ func Refresh(ctx context.Context, client *http.Client, cfg Config, current Token // (e.g. an existing refresh token or scopes) when the response omits them. // Error messages carry only the server's error/error_description — never the raw // body — so token material in an unexpected payload is not leaked. -func PostToken(ctx context.Context, client *http.Client, tokenEndpoint string, form url.Values, base Token, now func() time.Time) (Token, error) { +func PostToken(ctx context.Context, client *http.Client, tokenEndpoint string, form url.Values, base Token, extraHeaders map[string]string, now func() time.Time) (Token, error) { if err := validateTokenEndpoint(tokenEndpoint); err != nil { return Token{}, err } @@ -205,6 +205,7 @@ func PostToken(ctx context.Context, client *http.Client, tokenEndpoint string, f } request.Header.Set("Content-Type", "application/x-www-form-urlencoded") request.Header.Set("Accept", "application/json") + applyExtraHeaders(request, extraHeaders) response, err := client.Do(request) if err != nil { diff --git a/internal/oauth/flow_test.go b/internal/oauth/flow_test.go index 09ef01251..01345a8c7 100644 --- a/internal/oauth/flow_test.go +++ b/internal/oauth/flow_test.go @@ -153,7 +153,7 @@ func TestPostTokenRefusesRedirect(t *testing.T) { form.Set("grant_type", "authorization_code") form.Set("code", "the-code") form.Set("client_secret", "shh") - _, err := PostToken(context.Background(), http.DefaultClient, endpoint, form, Token{}, nil) + _, err := PostToken(context.Background(), http.DefaultClient, endpoint, form, Token{}, nil, nil) if !errors.Is(err, ErrUnsafeRedirect) { t.Fatalf("PostToken err = %v, want ErrUnsafeRedirect", err) } @@ -257,7 +257,7 @@ func TestExchangeCodeErrorRedactsBody(t *testing.T) { } func TestPostTokenRefusesInsecureEndpoint(t *testing.T) { - _, err := PostToken(context.Background(), http.DefaultClient, "http://auth.example.com/token", url.Values{}, Token{}, nil) + _, err := PostToken(context.Background(), http.DefaultClient, "http://auth.example.com/token", url.Values{}, Token{}, nil, nil) if !errors.Is(err, ErrInsecureTokenEndpoint) { t.Fatalf("err = %v, want ErrInsecureTokenEndpoint", err) } diff --git a/internal/oauth/oauth.go b/internal/oauth/oauth.go index 4938e7bf3..2c5b8fea1 100644 --- a/internal/oauth/oauth.go +++ b/internal/oauth/oauth.go @@ -23,6 +23,7 @@ package oauth import ( "errors" + "net/http" "strings" "time" ) @@ -79,6 +80,14 @@ type Config struct { // ExtraAuthParams are appended to the authorization URL (e.g. login_hint). ExtraAuthParams map[string]string + + // ExtraHeaders are set on every device-authorization, token-poll, code + // exchange, and refresh HTTP request this Config makes. Most providers + // need none; Kimi Code's backend rejects all of those requests with 401 + // unless its vendor-identity X-Msh-* headers are present (see + // kimiExtraHeaders in presets.go), which the generic RFC 8628/OAuth2 form + // bodies built elsewhere in this package have no other way to carry. + ExtraHeaders map[string]string } // Errors returned by the engine. Callers can match these with errors.Is. @@ -105,3 +114,13 @@ var ( // trimmed is a tiny helper used across the package. func trimmed(s string) string { return strings.TrimSpace(s) } + +// applyExtraHeaders sets a Config's provider-specific headers (see +// Config.ExtraHeaders) on an outgoing request, after the caller has already +// set the standard Content-Type/Accept pair so a provider entry could +// (in principle) override them. +func applyExtraHeaders(request *http.Request, headers map[string]string) { + for key, value := range headers { + request.Header.Set(key, value) + } +} diff --git a/internal/oauth/presets.go b/internal/oauth/presets.go index 3f8b7fe79..159019f86 100644 --- a/internal/oauth/presets.go +++ b/internal/oauth/presets.go @@ -3,6 +3,8 @@ package oauth import ( "os" "strings" + + "github.com/Gitlawb/zero/internal/kimiidentity" ) // envWithPresetsAllowed returns an env map that opts into the baked-in presets @@ -99,6 +101,26 @@ var builtinOAuthPresets = map[string]providerPreset{ Scopes: []string{"openid", "profile", "email", "offline_access", "api.connectors.read", "api.connectors.invoke"}, Flow: FlowLoopback, }, + // Kimi Code uses the same public OAuth client identity the open-source + // kimi-cli ships (github.com/MoonshotAI/kimi-cli, + // src/kimi_cli/auth/oauth.py). The flow is device-code only (RFC 8628) + // against auth.kimi.com — there is no loopback/authorize path — and the + // resulting access token is accepted directly as a bearer on the managed + // coding endpoint https://api.kimi.com/coding/v1 (an OpenAI-compatible + // endpoint). No ID-token claim extraction is needed; the bearer is the + // whole credential. The preset (and catalog descriptor) key is + // "kimi-code", not "kimi": the existing `moonshot` entry already aliases + // "kimi" to itself (its API-key path at api.moonshot.ai), and Get() + // matches an exact descriptor ID before it ever reaches another + // descriptor's aliases, so reusing "kimi" here would silently steal that + // alias out from under existing moonshot profiles. + "kimi-code": { + ClientID: "17e5f671-d194-4dfb-9706-5516cb48c098", + DeviceAuthorizationEndpoint: "https://auth.kimi.com/api/oauth/device_authorization", + TokenEndpoint: "https://auth.kimi.com/api/oauth/token", + Scopes: []string{"openid", "profile", "email", "offline_access"}, + Flow: FlowDevice, + }, } // lookupOAuthPreset returns the baked-in preset for a provider name (if any). @@ -128,3 +150,28 @@ func scopesOrPreset(envScopes string, preset []string) []string { // Copy so a caller appending to cfg.Scopes can't mutate the shared preset slice. return append([]string(nil), preset...) } + +// providerExtraHeaders returns the Config.ExtraHeaders a provider's OAuth +// requests need beyond the generic RFC 8628/OAuth2 form bodies this package +// builds. This is a protocol requirement of the provider's OWN backend +// (not tied to whether its preset client_id or an operator-supplied one is in +// use), so unlike the presets above it applies regardless of +// ZERO_OAUTH_ALLOW_PRESETS. +func providerExtraHeaders(name string) map[string]string { + if strings.ToLower(strings.TrimSpace(name)) == "kimi-code" { + return kimiExtraHeaders() + } + return nil +} + +// kimiExtraHeaders returns the X-Msh-* vendor-identity headers Kimi Code's +// OAuth/API backend requires on every device-authorization, poll, and refresh +// request — reported to reject all of them with 401 otherwise. They come from +// the shared kimiidentity package so login/refresh and the catalog's +// managed-endpoint completions (kimi-code CustomHeaders) present the SAME +// persistent device identity; values are reverse-engineered from the +// open-source kimi-cli client, not from published documentation, and should +// be confirmed against a real login before this ships. +func kimiExtraHeaders() map[string]string { + return kimiidentity.Headers() +} diff --git a/internal/oauth/presets_test.go b/internal/oauth/presets_test.go index f73f1e102..aeba07d32 100644 --- a/internal/oauth/presets_test.go +++ b/internal/oauth/presets_test.go @@ -1,8 +1,11 @@ package oauth import ( + "path/filepath" "strings" "testing" + + "github.com/Gitlawb/zero/internal/kimiidentity" ) func TestScopesOrPresetReturnsACopy(t *testing.T) { @@ -179,6 +182,62 @@ func TestResolveConfigHuggingFaceWithEnvClientID(t *testing.T) { } } +// Kimi Code ships a baked-in client_id (the public kimi-cli identity) AND a +// device-code endpoint, so the preset resolves without env. The flow is device +// only (RFC 8628): no loopback/authorize endpoint, no issuer discovery. The +// preset key is "kimi-code", not "kimi": moonshot already aliases "kimi" to +// itself (see TestKimiAliasStillResolvesToMoonshot in the providercatalog +// package's catalog_test.go), so reusing it here would steal that alias from +// existing moonshot profiles. +func TestResolveConfigKimiCodePreset(t *testing.T) { + tempDir := t.TempDir() + cleanup := kimiidentity.SetDeviceIDPathForTest(filepath.Join(tempDir, "kimi-device-id")) + t.Cleanup(cleanup) + t.Setenv("XDG_CONFIG_HOME", tempDir) + t.Setenv("APPDATA", tempDir) + r := NewRegistry() + cfg, flow, err := r.ResolveConfig("kimi-code", map[string]string{"ZERO_OAUTH_ALLOW_PRESETS": "1"}) + if err != nil { + t.Fatalf("ResolveConfig(kimi-code): %v", err) + } + if cfg.ClientID != "17e5f671-d194-4dfb-9706-5516cb48c098" { + t.Fatalf("client_id = %q", cfg.ClientID) + } + if cfg.DeviceAuthorizationEndpoint != "https://auth.kimi.com/api/oauth/device_authorization" { + t.Fatalf("device endpoint = %q", cfg.DeviceAuthorizationEndpoint) + } + if cfg.TokenEndpoint != "https://auth.kimi.com/api/oauth/token" { + t.Fatalf("token = %q", cfg.TokenEndpoint) + } + if cfg.AuthorizationEndpoint != "" { + t.Fatalf("authorize endpoint = %q, want empty (Kimi has no loopback flow)", cfg.AuthorizationEndpoint) + } + if flow != FlowDevice { + t.Fatalf("flow = %q, want device (RFC 8628 only)", flow) + } + if len(cfg.Scopes) == 0 { + t.Fatal("preset scopes should be populated") + } + for _, header := range []string{"X-Msh-Platform", "X-Msh-Version", "X-Msh-Device-Name", "X-Msh-Device-Model", "X-Msh-Os-Version", "X-Msh-Device-Id"} { + if cfg.ExtraHeaders[header] == "" { + t.Fatalf("ExtraHeaders[%q] = %q, want a non-empty vendor-identity header (Kimi's backend rejects requests missing these)", header, cfg.ExtraHeaders[header]) + } + } +} + +// A provider with no ExtraHeaders requirement (e.g. xAI) must not pick up +// Kimi's headers or any other provider's. +func TestResolveConfigWithoutExtraHeadersRequirement(t *testing.T) { + r := NewRegistry() + cfg, _, err := r.ResolveConfig("xai", map[string]string{"ZERO_OAUTH_ALLOW_PRESETS": "1"}) + if err != nil { + t.Fatalf("ResolveConfig(xai): %v", err) + } + if len(cfg.ExtraHeaders) != 0 { + t.Fatalf("ExtraHeaders = %#v, want none for a provider with no header requirement", cfg.ExtraHeaders) + } +} + // ChatGPT (Codex) ships a baked-in client_id (the public Codex CLI identity), // so the preset resolves without env. The flow is loopback because the Codex // backend requires a browser; there is no device-code path. diff --git a/internal/oauth/providers.go b/internal/oauth/providers.go index e3b2807a3..125c68d9d 100644 --- a/internal/oauth/providers.go +++ b/internal/oauth/providers.go @@ -77,6 +77,7 @@ func (r *Registry) ResolveConfig(name string, env map[string]string) (Config, Fl DeviceAuthorizationEndpoint: firstNonEmpty(strings.TrimSpace(envValue(env, envKey(name, "DEVICE_URL"))), preset.DeviceAuthorizationEndpoint), IssuerURL: firstNonEmpty(strings.TrimSpace(envValue(env, envKey(name, "ISSUER_URL"))), preset.IssuerURL), Scopes: scopesOrPreset(envValue(env, envKey(name, "SCOPES")), preset.Scopes), + ExtraHeaders: providerExtraHeaders(name), } if cfg.ClientID == "" { hint := "" diff --git a/internal/providercatalog/catalog.go b/internal/providercatalog/catalog.go index 9108b857d..c8f72c4d3 100644 --- a/internal/providercatalog/catalog.go +++ b/internal/providercatalog/catalog.go @@ -5,6 +5,8 @@ import ( "fmt" "strings" "unicode" + + "github.com/Gitlawb/zero/internal/kimiidentity" ) type Transport string @@ -70,6 +72,19 @@ type Descriptor struct { // OAuthDeviceFlow reports that RFC 8628 device-code login is supported (for // headless / SSH use) in addition to the browser flow. OAuthDeviceFlow bool + // OAuthDeviceOnly reports that device-code is the ONLY OAuth path this + // provider has — there is no browser/loopback authorize endpoint to fall + // back to (Kimi Code). Callers that default a plain Enter/click to the + // browser flow for OAuthDeviceFlow providers must check this first: for a + // device-only provider that generic path has no endpoint to hit at all. + OAuthDeviceOnly bool + + // RuntimeHeaders, if set, lazily produces headers (e.g. a vendor's + // device-identity headers) to attach to CustomHeaders only when + // withRuntimeHeaders is true. Providers that need identity/headers on + // runtime requests (not just OAuth) set this instead of hardcoding an ID + // check inside cloneDescriptor — keeping the clone helper provider-agnostic. + RuntimeHeaders func() map[string]string // Recommended marks a provider surfaced at the top and badged in // catalog-ordered lists and pickers. The recommended descriptors are the first @@ -135,6 +150,38 @@ var descriptors = []Descriptor{ d.RequiresAuth = true return oauthProvider(d, false, false) }(), + // Kimi Code (managed OAuth) — the bearer from a Kimi Code device-code OAuth + // login routes to the managed coding endpoint at https://api.kimi.com/coding/v1 + // (NOT api.moonshot.ai/v1, which is the API-key path for the `moonshot` + // catalog entry). Kimi only supports the RFC 8628 device-code flow against + // auth.kimi.com; the access token is accepted directly as a bearer on the + // managed endpoint, so no client spoofing is involved. The baked-in preset + // ships the public Kimi Code client_id and endpoints (off by default, like + // the other presets); env overrides via ZERO_OAUTH_KIMI_CODE_* win. + // + // Descriptor ID is "kimi-code", not "kimi": the `moonshot` entry below + // already aliases "kimi" to itself, and Get() matches an exact descriptor + // ID before it reaches another descriptor's aliases, so reusing "kimi" + // here would silently steal that alias out from under any existing + // moonshot profile (changing its endpoint, default model, and + // MOONSHOT_API_KEY auth without the user asking for it). + // + // Default model is "kimi-for-coding" — the managed endpoint's standard + // tier, available to every Kimi Code member. "kimi-for-coding-highspeed" + // also exists but requires a higher subscription tier (Allegretto+); a + // user who has it can select it explicitly rather than have a fresh + // `zero auth kimi-code` default to a model their plan may not include. + func() Descriptor { + d := openAICompat("kimi-code", "Kimi Code", "https://api.kimi.com/coding/v1", "kimi-for-coding", nil) + d.RequiresAuth = true + d = oauthProvider(d, false, true) + d.OAuthDeviceOnly = true + // Kimi's managed endpoint requires stable vendor-identity headers on + // every request; lazily mint them (with a persistent device ID) only + // when a runtime request actually needs them, not merely on listing. + d.RuntimeHeaders = kimiidentity.Headers + return d + }(), openAICompat("groq", "Groq", "https://api.groq.com/openai/v1", "llama-3.3-70b-versatile", []string{"GROQ_API_KEY"}), openAICompat("deepseek", "DeepSeek", "https://api.deepseek.com/v1", "deepseek-chat", []string{"DEEPSEEK_API_KEY"}), openAICompat("together", "Together AI", "https://api.together.xyz/v1", "meta-llama/Llama-3.3-70B-Instruct-Turbo", []string{"TOGETHER_API_KEY"}), @@ -182,7 +229,8 @@ var descriptors = []Descriptor{ func All() []Descriptor { copied := make([]Descriptor, 0, len(descriptors)) for _, descriptor := range descriptors { - copied = append(copied, cloneDescriptor(descriptor)) + // Listing must not mint Kimi's on-disk device identity. + copied = append(copied, cloneDescriptor(descriptor, false)) } return copied } @@ -199,11 +247,11 @@ func Get(id string) (Descriptor, bool) { normalized := NormalizeID(id) for _, descriptor := range descriptors { if descriptor.ID == normalized { - return cloneDescriptor(descriptor), true + return cloneDescriptor(descriptor, true), true } for _, alias := range descriptor.Aliases { if NormalizeID(alias) == normalized { - return cloneDescriptor(descriptor), true + return cloneDescriptor(descriptor, true), true } } } @@ -224,7 +272,7 @@ func ListByTransport(transport Transport) []Descriptor { items := make([]Descriptor, 0) for _, descriptor := range descriptors { if descriptor.Transport == normalized { - items = append(items, cloneDescriptor(descriptor)) + items = append(items, cloneDescriptor(descriptor, false)) } } return items @@ -344,7 +392,9 @@ func OAuthProviders() []Descriptor { out := []Descriptor{} for _, descriptor := range descriptors { if descriptor.OAuth { - out = append(out, cloneDescriptor(descriptor)) + // Listing only: runtime headers (and the Kimi device-id file they + // mint) are applied when Get/Require builds a real profile. + out = append(out, cloneDescriptor(descriptor, false)) } } return out @@ -399,12 +449,22 @@ func transportDescriptor(id string, name string, transport Transport, baseURL st } } -func cloneDescriptor(descriptor Descriptor) Descriptor { +// cloneDescriptor returns an independent copy of descriptor. When +// withRuntimeHeaders is true, a descriptor with RuntimeHeaders set (e.g. +// kimi-code) also receives its vendor identity headers (including a persistent +// device ID on disk). Listing paths (All, OAuthProviders, ListByTransport) +// pass false so merely enumerating providers never mints +// ~/.config/zero/kimi-device-id for users who never touch Kimi; Get/Require +// pass true so resolve-time profile building and completions still present the +// same identity the OAuth login used. +func cloneDescriptor(descriptor Descriptor, withRuntimeHeaders bool) Descriptor { descriptor.AuthEnvVars = append([]string{}, descriptor.AuthEnvVars...) descriptor.SupportedAPIFormats = append([]APIFormat{}, descriptor.SupportedAPIFormats...) descriptor.Aliases = append([]string{}, descriptor.Aliases...) if descriptor.CustomHeaders != nil { descriptor.CustomHeaders = copyStringMap(descriptor.CustomHeaders) + } else if withRuntimeHeaders && descriptor.RuntimeHeaders != nil { + descriptor.CustomHeaders = descriptor.RuntimeHeaders() } return descriptor } diff --git a/internal/providercatalog/catalog_test.go b/internal/providercatalog/catalog_test.go index 8fd245eac..3f1bbc657 100644 --- a/internal/providercatalog/catalog_test.go +++ b/internal/providercatalog/catalog_test.go @@ -2,9 +2,12 @@ package providercatalog import ( "errors" + "path/filepath" "reflect" "strings" "testing" + + "github.com/Gitlawb/zero/internal/kimiidentity" ) var expectedCatalogIDs = []string{ @@ -19,6 +22,7 @@ var expectedCatalogIDs = []string{ "openrouter", "huggingface", "chatgpt", + "kimi-code", "groq", "deepseek", "together", @@ -346,7 +350,7 @@ func TestListByTransportPreservesCatalogOrder(t *testing.T) { TransportBedrock: {"bedrock"}, TransportVertex: {"vertex"}, TransportAnthropicCompat: {"minimax", "minimaxi-cn", "opencode-go-anthropic-compatible", "custom-anthropic-compatible"}, - TransportOpenAICompat: {"gitlawb-opengateway", "aimlapi", "ollama-cloud", "ollama", "lmstudio", "openrouter", "huggingface", "chatgpt", "groq", "deepseek", "together", "dashscope", "moonshot", "atlascloud", "longcat", "nvidia-nim", "mistral", "github", "xai", "venice", "xiaomi-mimo", "bankr", "zai", "zai-cn", "kilocode", "opencode", "opencode-go", "atomic-chat", "chatgpt-proxy", "custom-openai-compatible"}, + TransportOpenAICompat: {"gitlawb-opengateway", "aimlapi", "ollama-cloud", "ollama", "lmstudio", "openrouter", "huggingface", "chatgpt", "kimi-code", "groq", "deepseek", "together", "dashscope", "moonshot", "atlascloud", "longcat", "nvidia-nim", "mistral", "github", "xai", "venice", "xiaomi-mimo", "bankr", "zai", "zai-cn", "kilocode", "opencode", "opencode-go", "atomic-chat", "chatgpt-proxy", "custom-openai-compatible"}, } for transport, wantIDs := range cases { @@ -405,7 +409,7 @@ func TestReturnedDescriptorsAreCopies(t *testing.T) { func TestOAuthProviderClassification(t *testing.T) { oauthIDs := descriptorIDs(OAuthProviders()) - if want := []string{"openrouter", "huggingface", "chatgpt", "xai"}; !reflect.DeepEqual(oauthIDs, want) { + if want := []string{"openrouter", "huggingface", "chatgpt", "kimi-code", "xai"}; !reflect.DeepEqual(oauthIDs, want) { t.Fatalf("OAuthProviders() = %#v, want %#v", oauthIDs, want) } if d, _ := Get("openrouter"); !d.OAuthMintsKey { @@ -414,6 +418,9 @@ func TestOAuthProviderClassification(t *testing.T) { if d, _ := Get("xai"); !d.OAuthDeviceFlow { t.Fatal("xai should advertise device-code flow") } + if d, _ := Get("kimi-code"); !d.OAuthDeviceFlow || !d.OAuthDeviceOnly { + t.Fatal("kimi-code should advertise device-only code flow") + } if d, _ := Get("huggingface"); !d.OAuthDeviceFlow { t.Fatal("huggingface should advertise device-code flow") } @@ -422,6 +429,73 @@ func TestOAuthProviderClassification(t *testing.T) { } } +// TestKimiAliasStillResolvesToMoonshot pins the alias-collision fix: moonshot +// already exposes "kimi" as an alias for its API-key endpoint. The Kimi Code +// OAuth descriptor must use a non-conflicting ID ("kimi-code") so resolving +// "kimi" continues to land on moonshot (endpoint, default model, MOONSHOT_API_KEY). +func TestKimiAliasStillResolvesToMoonshot(t *testing.T) { + d, ok := Get("kimi") + if !ok { + t.Fatal(`Get("kimi") returned false`) + } + if d.ID != "moonshot" { + t.Fatalf(`Get("kimi").ID = %q, want "moonshot" (kimi-code must not steal this alias)`, d.ID) + } + if d.DefaultBaseURL != "https://api.moonshot.ai/v1" { + t.Fatalf(`Get("kimi").DefaultBaseURL = %q, want moonshot API-key endpoint`, d.DefaultBaseURL) + } + if len(d.AuthEnvVars) == 0 || d.AuthEnvVars[0] != "MOONSHOT_API_KEY" { + t.Fatalf(`Get("kimi").AuthEnvVars = %#v, want MOONSHOT_API_KEY`, d.AuthEnvVars) + } + if d.OAuth { + t.Fatal(`Get("kimi") must not be OAuth (that is kimi-code, not the moonshot alias)`) + } + + code, ok := Get("kimi-code") + if !ok { + t.Fatal(`Get("kimi-code") returned false`) + } + if code.ID != "kimi-code" { + t.Fatalf(`Get("kimi-code").ID = %q`, code.ID) + } + if code.DefaultBaseURL != "https://api.kimi.com/coding/v1" { + t.Fatalf(`Get("kimi-code").DefaultBaseURL = %q, want managed coding endpoint`, code.DefaultBaseURL) + } + if !code.OAuth || !code.OAuthDeviceOnly { + t.Fatalf(`Get("kimi-code") oauth flags wrong: OAuth=%v OAuthDeviceOnly=%v`, code.OAuth, code.OAuthDeviceOnly) + } +} + +// TestKimiRuntimeHeadersOnlyOnGet ensures listing providers (All / OAuthProviders) +// does not populate kimi-code's CustomHeaders (which mints a device-id file), +// while Get does so resolve-time request building still gets the vendor headers. +func TestKimiRuntimeHeadersOnlyOnGet(t *testing.T) { + tempDir := t.TempDir() + cleanup := kimiidentity.SetDeviceIDPathForTest(filepath.Join(tempDir, "kimi-device-id")) + t.Cleanup(cleanup) + t.Setenv("XDG_CONFIG_HOME", tempDir) + t.Setenv("APPDATA", tempDir) + for _, d := range All() { + if d.ID == "kimi-code" && d.CustomHeaders != nil { + t.Fatalf("All() must not populate kimi-code CustomHeaders: %#v", d.CustomHeaders) + } + } + for _, d := range OAuthProviders() { + if d.ID == "kimi-code" && d.CustomHeaders != nil { + t.Fatalf("OAuthProviders() must not populate kimi-code CustomHeaders: %#v", d.CustomHeaders) + } + } + d, ok := Get("kimi-code") + if !ok { + t.Fatal(`Get("kimi-code") returned false`) + } + for _, header := range []string{"X-Msh-Platform", "X-Msh-Version", "X-Msh-Device-Name", "X-Msh-Device-Model", "X-Msh-Os-Version", "X-Msh-Device-Id"} { + if d.CustomHeaders[header] == "" { + t.Fatalf("Get(kimi-code).CustomHeaders[%q] empty, want vendor-identity header for completions", header) + } + } +} + func descriptorIDs(descriptors []Descriptor) []string { ids := make([]string, 0, len(descriptors)) for _, descriptor := range descriptors { diff --git a/internal/providercatalog/oauth_test.go b/internal/providercatalog/oauth_test.go index 0733f32ae..1888226c5 100644 --- a/internal/providercatalog/oauth_test.go +++ b/internal/providercatalog/oauth_test.go @@ -4,8 +4,8 @@ import "testing" func TestOAuthProviders(t *testing.T) { providers := OAuthProviders() - if len(providers) != 4 { - t.Fatalf("OAuthProviders() = %d, want 4 (openrouter, xai, huggingface, chatgpt)", len(providers)) + if len(providers) != 5 { + t.Fatalf("OAuthProviders() = %d, want 5 (openrouter, xai, kimi-code, huggingface, chatgpt)", len(providers)) } byID := map[string]Descriptor{} for _, d := range providers { @@ -30,6 +30,10 @@ func TestOAuthProviders(t *testing.T) { if !ok || cg.OAuthMintsKey || cg.OAuthDeviceFlow { t.Fatalf("chatgpt oauth flags wrong: %+v", cg) } + kimi, ok := byID["kimi-code"] + if !ok || kimi.OAuthMintsKey || !kimi.OAuthDeviceFlow || !kimi.OAuthDeviceOnly { + t.Fatalf("kimi-code oauth flags wrong: %+v", kimi) + } } func TestOAuthProvidersReturnsIndependentClones(t *testing.T) { @@ -58,7 +62,7 @@ func TestOAuthProvidersReturnsIndependentClones(t *testing.T) { func TestNonOAuthProvidersNotFlagged(t *testing.T) { for _, d := range All() { switch d.ID { - case "openrouter", "xai", "huggingface", "chatgpt": + case "openrouter", "xai", "kimi-code", "huggingface", "chatgpt": continue } if d.OAuth { diff --git a/internal/tui/model.go b/internal/tui/model.go index 280c8ccfb..3a8e9b61e 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -1029,8 +1029,10 @@ func (m model) noBlockingModal() bool { func (m model) quit() (tea.Model, tea.Cmd) { if m.providerWizard != nil { + m.providerWizard.cancelDeviceLogin() m.providerWizard.resetAimlapiOnboard() } + m.setup.cancelDeviceLogin() m.stopPRWatcher() m.stopAllBackgroundTerminalSessions() m.shutdownLSPManager() diff --git a/internal/tui/oauth_device.go b/internal/tui/oauth_device.go index 00bdb557c..e6ec9276b 100644 --- a/internal/tui/oauth_device.go +++ b/internal/tui/oauth_device.go @@ -59,7 +59,13 @@ func oauthDevicePrepare(name string) (oauth.DeviceAuth, oauth.Config, error) { // oauthDeviceComplete polls for the token authorized via oauthDevicePrepare and // stores it under provider: (phase 2). The runtime resolver then attaches // the refreshable token to model calls. -func oauthDeviceComplete(name string, cfg oauth.Config, auth oauth.DeviceAuth) error { +// +// parent must be cancelable by the caller (not context.Background()): the poll +// can run for up to 10 minutes waiting on the user to finish authorizing in +// their browser, and if the caller has since abandoned the flow (e.g. Esc in +// the TUI), canceling parent is the only way to stop this from silently +// completing and persisting a credential the user believed they'd canceled. +func oauthDeviceComplete(parent context.Context, name string, cfg oauth.Config, auth oauth.DeviceAuth) error { store, err := oauth.NewStore(oauth.StoreOptions{}) if err != nil { return err @@ -72,7 +78,7 @@ func oauthDeviceComplete(name string, cfg oauth.Config, auth oauth.DeviceAuth) e if err != nil { return err } - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) + ctx, cancel := context.WithTimeout(parent, 10*time.Minute) defer cancel() _, err = manager.CompleteDeviceLogin(ctx, name, cfg, auth) return err diff --git a/internal/tui/onboarding.go b/internal/tui/onboarding.go index c4b26705c..d8dba54ce 100644 --- a/internal/tui/onboarding.go +++ b/internal/tui/onboarding.go @@ -50,23 +50,34 @@ type setupState struct { oauthMode bool oauthPending bool oauthErr string + // oauthAttemptID identifies the current OAuth attempt so a late phase-one + // or poll-result message from an attempt the user already abandoned (Esc, + // then restarted the same provider) cannot be mistaken for the current + // one — providerID alone is not enough since the provider doesn't change + // across attempts. + oauthAttemptID int // Device-code login (RFC 8628) state while an OAuth login is in flight. oauthDevice bool deviceUserCode string deviceVerificationURI string - stage setupStage - err string - baseURL string - name string - apiKey textinput.Model - models []providerWizardModel - modelIndex int - modelQuery string - modelForID string - modelLoad bool - modelErr string - modelSrc string - modelGen uint64 + // deviceLoginCancel cancels the background context backing an in-flight + // device-code poll (setupDevicePollCmd), so abandoning setup (Esc) actually + // stops the poll instead of leaving it to run for up to 10 minutes and + // silently save a credential the user backed out of. + deviceLoginCancel context.CancelFunc + stage setupStage + err string + baseURL string + name string + apiKey textinput.Model + models []providerWizardModel + modelIndex int + modelQuery string + modelForID string + modelLoad bool + modelErr string + modelSrc string + modelGen uint64 // aimlapi holds the shared aimlapi.com onboarding sub-flow while // setup is on setupStageAimlapi. aimlapi *aimlapiOnboardState @@ -93,6 +104,7 @@ type setupOAuthMsg struct { apiKey string tokenLogin bool providerID string + attemptID int err error } @@ -154,10 +166,18 @@ func (m model) handleSetupKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { if m.setup.oauthPending { switch { case keyCtrl(msg, 'c'): + // Stop a device-code poll actually running in the background instead + // of merely quitting the UI — otherwise completing the login later + // in the browser still silently saves the credential. + m.setup.cancelDeviceLogin() return m, tea.Quit case keyIs(msg, tea.KeyEsc): m.setup.oauthPending = false m.setup.oauthDevice = false + // Stop a device-code poll actually running in the background instead + // of merely dismissing the UI — otherwise completing the login later + // in the browser still silently saves the credential. + m.setup.cancelDeviceLogin() } return m, nil } @@ -531,7 +551,7 @@ func (m *model) moveSetupMethod(delta int) { // setupOAuthCmd runs the chosen provider's browser OAuth login off the UI // goroutine for first-run setup. Mirrors the /provider wizard's flow. -func setupOAuthCmd(provider providercatalog.Descriptor) tea.Cmd { +func setupOAuthCmd(provider providercatalog.Descriptor, attemptID int) tea.Cmd { switch { case provider.OAuthMintsKey: return func() tea.Msg { @@ -539,17 +559,17 @@ func setupOAuthCmd(provider providercatalog.Descriptor) tea.Cmd { OpenBrowser: browser.OpenURL, Timeout: 3 * time.Minute, }) - return setupOAuthMsg{apiKey: key, providerID: provider.ID, err: err} + return setupOAuthMsg{apiKey: key, providerID: provider.ID, attemptID: attemptID, err: err} } case provider.ID == "chatgpt": return func() tea.Msg { err := runProviderChatGPTLogin() - return setupOAuthMsg{tokenLogin: true, providerID: provider.ID, err: err} + return setupOAuthMsg{tokenLogin: true, providerID: provider.ID, attemptID: attemptID, err: err} } default: name := provider.ID return func() tea.Msg { - return setupOAuthMsg{tokenLogin: true, providerID: name, err: runProviderTokenLogin(name)} + return setupOAuthMsg{tokenLogin: true, providerID: name, attemptID: attemptID, err: runProviderTokenLogin(name)} } } } @@ -558,6 +578,7 @@ func setupOAuthCmd(provider providercatalog.Descriptor) tea.Cmd { // for first-run setup: the user_code + verification URI to display. type setupOAuthDeviceMsg struct { providerID string + attemptID int userCode string verifyURL string cfg oauth.Config @@ -565,14 +586,15 @@ type setupOAuthDeviceMsg struct { err error } -func setupDevicePrepareCmd(name string) tea.Cmd { +func setupDevicePrepareCmd(name string, attemptID int) tea.Cmd { return func() tea.Msg { auth, cfg, err := oauthDevicePrepare(name) if err != nil { - return setupOAuthDeviceMsg{providerID: name, err: err} + return setupOAuthDeviceMsg{providerID: name, attemptID: attemptID, err: err} } return setupOAuthDeviceMsg{ providerID: name, + attemptID: attemptID, userCode: auth.UserCode, verifyURL: oauthDeviceVerifyTarget(auth), cfg: cfg, @@ -581,35 +603,66 @@ func setupDevicePrepareCmd(name string) tea.Cmd { } } -func setupDevicePollCmd(name string, cfg oauth.Config, auth oauth.DeviceAuth) tea.Cmd { +// setupDevicePollCmd runs phase 2 (poll for the token + store) off the UI +// goroutine and reports completion as a regular OAuth result. ctx must be +// cancelable by the caller so abandoning setup actually stops the poll. +func setupDevicePollCmd(ctx context.Context, name string, attemptID int, cfg oauth.Config, auth oauth.DeviceAuth) tea.Cmd { return func() tea.Msg { - return setupOAuthMsg{tokenLogin: true, providerID: name, err: oauthDeviceComplete(name, cfg, auth)} + return setupOAuthMsg{tokenLogin: true, providerID: name, attemptID: attemptID, err: oauthDeviceComplete(ctx, name, cfg, auth)} } } +// cancelDeviceLogin stops an in-flight device-code poll, if any, and clears +// the stored cancel func. Safe to call even when no poll is running. +func (setup *setupState) cancelDeviceLogin() { + if setup == nil || setup.deviceLoginCancel == nil { + return + } + setup.deviceLoginCancel() + setup.deviceLoginCancel = nil +} + +// beginSetupOAuthAttempt starts a new first-run OAuth attempt and returns its +// id. A new attempt supersedes any poll left over from a previous one and +// bumps oauthAttemptID so a late phase-one/poll-result message tagged with an +// older id is rejected by setupOAuthResultMatches even though the provider +// (and oauthPending) look the same as the new attempt. +func (setup *setupState) beginSetupOAuthAttempt(device bool) int { + setup.cancelDeviceLogin() + setup.oauthAttemptID++ + setup.oauthPending = true + setup.oauthDevice = device + setup.oauthErr = "" + setup.deviceUserCode = "" + setup.deviceVerificationURI = "" + return setup.oauthAttemptID +} + +// setupOAuthResultMatches reports whether an OAuth result message still +// belongs to the in-flight attempt, rejecting a stale message left over from +// an attempt the user abandoned (Esc) and then restarted against the same +// provider. +func (setup *setupState) setupOAuthResultMatches(providerID string, attemptID int) bool { + if setup == nil || !setup.visible || !setup.oauthPending || strings.TrimSpace(providerID) == "" { + return false + } + return setup.oauthAttemptID == attemptID +} + // startSetupDeviceLogin begins the device-code flow for the selected OAuth // provider during first-run setup (phase 1). func (m model) startSetupDeviceLogin(descriptor providercatalog.Descriptor) (tea.Model, tea.Cmd) { if !descriptor.OAuth || !descriptor.OAuthDeviceFlow { return m, nil } - m.setup.oauthPending = true - m.setup.oauthDevice = true - m.setup.oauthErr = "" - m.setup.deviceUserCode = "" - m.setup.deviceVerificationURI = "" - return m, setupDevicePrepareCmd(descriptor.ID) + attemptID := m.setup.beginSetupOAuthAttempt(true) + return m, setupDevicePrepareCmd(descriptor.ID, attemptID) } // applySetupOAuthDeviceCode handles phase 1 of device-code login: show the code, // then start phase 2 (the token poll). On error the redacted message is shown. func (m model) applySetupOAuthDeviceCode(msg setupOAuthDeviceMsg) (tea.Model, tea.Cmd) { - if !m.setup.visible || !m.setup.oauthPending { - return m, nil - } - // Ignore a stale result from a login the user has since replaced with a - // different provider (an in-flight prepare landing after the switch). - if msg.providerID != "" && msg.providerID != m.setupProviderDescriptor().ID { + if !m.setup.setupOAuthResultMatches(msg.providerID, msg.attemptID) { return m, nil } if msg.err != nil { @@ -620,22 +673,26 @@ func (m model) applySetupOAuthDeviceCode(msg setupOAuthDeviceMsg) (tea.Model, te } m.setup.deviceUserCode = msg.userCode m.setup.deviceVerificationURI = msg.verifyURL - return m, setupDevicePollCmd(msg.providerID, msg.cfg, msg.auth) + // The poll runs off the UI goroutine for up to 10 minutes; give it a + // context setup can cancel (Esc) so abandoning the flow actually stops it + // instead of leaving it to complete in the background. m.ctx is the + // parent so quitting zero entirely also unblocks the poll. + ctx, cancel := context.WithCancel(m.ctx) + m.setup.deviceLoginCancel = cancel + return m, setupDevicePollCmd(ctx, msg.providerID, msg.attemptID, msg.cfg, msg.auth) } // applySetupOAuth folds an OAuth login result into the first-run setup: on success // the credential is captured (minted key) or relied upon (stored token) and setup // jumps to model selection; on failure the redacted error is shown. func (m model) applySetupOAuth(msg setupOAuthMsg) (tea.Model, tea.Cmd) { - if !m.setup.visible || !m.setup.oauthPending { - return m, nil - } - // Ignore a stale result for a provider the user has since switched away from, - // so a late login can't capture a credential against the wrong provider. - if msg.providerID != "" && msg.providerID != m.setupProviderDescriptor().ID { + if !m.setup.setupOAuthResultMatches(msg.providerID, msg.attemptID) { return m, nil } m.setup.oauthPending = false + // The poll (if this was a device-code login) is done either way; release + // its context. + m.setup.cancelDeviceLogin() if msg.err != nil { m.setup.oauthErr = redaction.RedactString(msg.err.Error(), redaction.Options{}) return m, nil @@ -778,13 +835,16 @@ func (m model) advanceSetup() (tea.Model, tea.Cmd) { if descriptor.OAuth { // Headless/SSH boxes can't open a browser — use device code there // by default (the user can also force it with "d" from the list). - if descriptor.OAuthDeviceFlow && oauthPreferDeviceFlow() { + // A device-only provider (no loopback/authorization endpoint) + // must take this path on desktops too: the generic manager + // would still run the device flow, but with no UI its + // verification URL and user code are discarded and the + // "browser login" spinner just times out. + if descriptor.OAuthDeviceFlow && (descriptor.OAuthDeviceOnly || oauthPreferDeviceFlow()) { return m.startSetupDeviceLogin(descriptor) } - m.setup.oauthPending = true - m.setup.oauthDevice = false - m.setup.oauthErr = "" - return m, setupOAuthCmd(descriptor) + attemptID := m.setup.beginSetupOAuthAttempt(false) + return m, setupOAuthCmd(descriptor, attemptID) } } if m.setup.stage == setupStageProvider { diff --git a/internal/tui/onboarding_test.go b/internal/tui/onboarding_test.go index 1f70616fb..cdb37c556 100644 --- a/internal/tui/onboarding_test.go +++ b/internal/tui/onboarding_test.go @@ -5,6 +5,7 @@ import ( "errors" "net/http" "os" + "path/filepath" "strings" "testing" @@ -2036,6 +2037,106 @@ func TestApplySetupOAuthDeviceCodeShowsCodeAndPolls(t *testing.T) { } } +// TestSetupCtrlCCancelsDeviceLoginPoll regression-tests a bug where Ctrl+C +// during a first-run device-code poll (phase 2) quit the whole program +// without canceling the background context the poll command runs on. Since +// the TUI's parent context is context.Background(), the poll (up to 10 +// minutes) kept running after the process later exited via os.Exit, and if +// the user then finished authorizing in the browser, a completed-in-flight +// write could still land. Ctrl+C must cancel the poll before quitting. +func TestSetupCtrlCCancelsDeviceLoginPoll(t *testing.T) { + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + t.Setenv("ZERO_OAUTH_TOKENS_PATH", filepath.Join(t.TempDir(), "oauth-tokens.json")) + + m := setupAtOAuthList(t) + for i, p := range m.setup.providers { + if p.ID == "xai" { + m.setup.selected = i + break + } + } + m.setup.oauthPending = true + m.setup.oauthDevice = true + attemptID := m.setup.oauthAttemptID + + res, cmd := m.applySetupOAuthDeviceCode(setupOAuthDeviceMsg{ + providerID: "xai", attemptID: attemptID, userCode: "WXYZ-9", verifyURL: "https://x.ai/device", + }) + m = res.(model) + if cmd == nil { + t.Fatal("device-code msg should start the poll command") + } + if m.setup.deviceLoginCancel == nil { + t.Fatal("starting the poll should store a cancel func on setup") + } + + updated, _ := m.Update(testKeyCtrl('c')) + m = updated.(model) + if m.setup.deviceLoginCancel != nil { + t.Fatal("Ctrl+C should cancel the in-flight device-code poll") + } + + msg, ok := cmd().(setupOAuthMsg) + if !ok { + t.Fatalf("poll command returned %T, want setupOAuthMsg", msg) + } + if !errors.Is(msg.err, context.Canceled) { + t.Fatalf("poll error = %v, want context.Canceled (Ctrl+C should have canceled the background poll)", msg.err) + } +} + +// TestSetupStaleDeviceCodeAttemptRejected regression-tests a bug where +// abandoning a device-code login with Esc and immediately restarting it for +// the same provider let a late phase-one result from the FIRST attempt +// overwrite the second attempt's displayed code and start polling an +// authorization the user had already backed out of: setupOAuthDeviceMsg only +// carried providerID, which is identical across both attempts. +func TestSetupStaleDeviceCodeAttemptRejected(t *testing.T) { + m := setupAtOAuthList(t) + for i, p := range m.setup.providers { + if p.ID == "xai" { + m.setup.selected = i + break + } + } + + updated, cmd := m.Update(testKeyText("d")) // attempt 1 + m = updated.(model) + if cmd == nil { + t.Fatal("'d' should return the device-prepare command") + } + staleAttemptID := m.setup.oauthAttemptID + + // Abandon attempt 1 with Esc, then immediately restart against the same + // provider (attempt 2). + updated, _ = m.Update(testKey(tea.KeyEsc)) + m = updated.(model) + if m.setup.oauthPending { + t.Fatal("Esc should abandon the pending device login") + } + updated, cmd = m.Update(testKeyText("d")) // attempt 2 + m = updated.(model) + if cmd == nil { + t.Fatal("restarting should return a new device-prepare command") + } + if m.setup.oauthAttemptID == staleAttemptID { + t.Fatal("restarting the device flow should assign a new attempt id") + } + + // The stale phase-one result from attempt 1 lands after attempt 2 is + // already in flight — same provider, so providerID alone can't reject it. + res, pollCmd := m.applySetupOAuthDeviceCode(setupOAuthDeviceMsg{ + providerID: "xai", attemptID: staleAttemptID, userCode: "STALE-1", verifyURL: "https://x.ai/device/stale", + }) + m = res.(model) + if pollCmd != nil { + t.Fatal("stale attempt's phase-one result must not start a poll") + } + if m.setup.deviceUserCode == "STALE-1" { + t.Fatalf("stale attempt overwrote the current attempt's device code: %+v", m.setup) + } +} + func TestApplySetupOAuthSuccessAdvancesToModel(t *testing.T) { m := newModel(context.Background(), Options{ DiscoverProviderModels: func(ctx context.Context, profile config.ProviderProfile) ([]providermodeldiscovery.Model, error) { @@ -2182,3 +2283,58 @@ func TestCompleteSetupExportsActiveProviderEnv(t *testing.T) { t.Fatalf("%s = %q after setup save, want %q (children would spawn on the stale provider)", config.ActiveProviderEnv, got, next.providerName) } } + +// TestSetupEnterStartsDeviceFlowForDeviceOnlyProvider pins the first-run +// onboarding counterpart of the /provider wizard fix: Kimi Code has no +// loopback/authorize endpoint, so a plain desktop Enter must take the +// device-code path (showing the verification URL and user code) instead of +// the generic browser-login command, whose manager would run the device flow +// with a discarded output writer and leave the spinner to time out. +func TestSetupEnterStartsDeviceFlowForDeviceOnlyProvider(t *testing.T) { + // Force a "normal desktop with a browser available" environment: + // oauthPreferDeviceFlow() already picks device flow on headless boxes, + // which would mask the bug this test exists to catch. + t.Setenv("ZERO_OAUTH_DEVICE", "") + t.Setenv("SSH_CONNECTION", "") + t.Setenv("SSH_TTY", "") + t.Setenv("DISPLAY", ":0") + t.Setenv("WAYLAND_DISPLAY", "") + + m := newModel(context.Background(), Options{Setup: SetupOptions{ + Visible: true, + Providers: []SetupProviderOption{ + {ID: "kimi-code", Name: "Kimi Code", RequiresAuth: true}, + {ID: "xai", Name: "xAI", DefaultModel: "grok-4", RequiresAuth: true, EnvVar: "XAI_API_KEY"}, + }, + }}) + m.width = 100 + m.height = 30 + m = pressSetupContinueOnce(m) // Welcome → Method + m.setup.selectedMethod = 0 // Sign in with OAuth + updated, _ := m.Update(testKey(tea.KeyEnter)) + m = updated.(model) + + found := false + for i, p := range m.setup.providers { + if p.ID == "kimi-code" { + m.setup.selected = i + found = true + break + } + } + if !found { + t.Fatalf("kimi-code missing from OAuth provider list: %#v", m.setup.providers) + } + if !m.setupProviderDescriptor().OAuthDeviceOnly { + t.Fatal("test fixture assumes kimi-code is OAuthDeviceOnly") + } + + updated, cmd := m.Update(testKey(tea.KeyEnter)) + m = updated.(model) + if !m.setup.oauthPending || !m.setup.oauthDevice { + t.Fatalf("Enter on a device-only provider should start device login (pending=%v device=%v)", m.setup.oauthPending, m.setup.oauthDevice) + } + if cmd == nil { + t.Fatal("Enter on a device-only provider should return the device-prepare command") + } +} diff --git a/internal/tui/provider_wizard.go b/internal/tui/provider_wizard.go index 8f4ee8bb3..bcb8eba58 100644 --- a/internal/tui/provider_wizard.go +++ b/internal/tui/provider_wizard.go @@ -45,6 +45,9 @@ func (m model) applyProviderWizardOAuth(msg providerWizardOAuthMsg) (model, tea. return m, nil } m.providerWizard.oauthPending = false + // The poll (if this was a device-code login) is done either way; release + // its context. + m.providerWizard.cancelDeviceLogin() if msg.err != nil { m.providerWizard.oauthErr = redaction.ErrorMessage(msg.err, redaction.Options{}) return m, nil @@ -137,7 +140,13 @@ func (m model) applyProviderWizardDeviceCode(msg providerWizardDeviceCodeMsg) (m } m.providerWizard.deviceUserCode = msg.userCode m.providerWizard.deviceVerificationURI = msg.verifyURL - return m, providerWizardDevicePollCmd(msg.providerID, msg.attemptID, msg.cfg, msg.auth) + // The poll runs off the UI goroutine for up to 10 minutes; give it a + // context the wizard can cancel (Esc) so abandoning the flow actually + // stops it instead of leaving it to complete in the background. m.ctx is + // the parent so quitting zero entirely also unblocks the poll. + ctx, cancel := context.WithCancel(m.ctx) + m.providerWizard.deviceLoginCancel = cancel + return m, providerWizardDevicePollCmd(ctx, msg.providerID, msg.attemptID, msg.cfg, msg.auth) } // providerWizardSupportsOAuth reports whether the credential step should offer a @@ -315,10 +324,11 @@ func providerWizardDevicePrepareCmd(name string, attemptID int) tea.Cmd { } // providerWizardDevicePollCmd runs phase 2 (poll for the token + store) off the -// UI goroutine and reports completion as a regular OAuth result. -func providerWizardDevicePollCmd(name string, attemptID int, cfg oauth.Config, auth oauth.DeviceAuth) tea.Cmd { +// UI goroutine and reports completion as a regular OAuth result. ctx must be +// cancelable by the caller so abandoning the wizard actually stops the poll. +func providerWizardDevicePollCmd(ctx context.Context, name string, attemptID int, cfg oauth.Config, auth oauth.DeviceAuth) tea.Cmd { return func() tea.Msg { - return providerWizardOAuthMsg{providerID: name, attemptID: attemptID, tokenLogin: true, err: oauthDeviceComplete(name, cfg, auth)} + return providerWizardOAuthMsg{providerID: name, attemptID: attemptID, tokenLogin: true, err: oauthDeviceComplete(ctx, name, cfg, auth)} } } @@ -376,7 +386,7 @@ func providerWizardMethodOptions() []providerWizardMethodOption { options = append(options, providerWizardMethodOption{ oauth: true, label: "Sign in with OAuth", - subtitle: "One-click browser login, no API key to copy (OpenRouter, xAI, ChatGPT, Hugging Face).", + subtitle: "No API key to copy — one-click browser login (OpenRouter, xAI, ChatGPT, Hugging Face) or device code (Kimi Code).", }) } options = append(options, providerWizardMethodOption{ @@ -439,6 +449,11 @@ type providerWizardState struct { oauthDevice bool deviceUserCode string deviceVerificationURI string + // deviceLoginCancel cancels the background context backing an in-flight + // device-code poll (providerWizardDevicePollCmd), so abandoning the wizard + // (Esc) actually stops the poll instead of leaving it to run for up to 10 + // minutes and silently save a credential the user backed out of. + deviceLoginCancel context.CancelFunc // aimlapi holds the shared aimlapi.com onboarding sub-flow while // the wizard is on providerWizardStepAimlapi. aimlapi *aimlapiOnboardState @@ -500,6 +515,8 @@ func (wizard *providerWizardState) currentProvider() providercatalog.Descriptor } func (wizard *providerWizardState) beginOAuthAttempt(device bool) int { + // A new attempt supersedes any poll left over from a previous one. + wizard.cancelDeviceLogin() wizard.oauthAttemptID++ wizard.oauthPending = true wizard.oauthDevice = device @@ -516,6 +533,16 @@ func (wizard *providerWizardState) oauthResultMatches(providerID string, attempt return wizard.currentProvider().ID == providerID && wizard.oauthAttemptID == attemptID } +// cancelDeviceLogin stops an in-flight device-code poll, if any, and clears +// the stored cancel func. Safe to call even when no poll is running. +func (wizard *providerWizardState) cancelDeviceLogin() { + if wizard == nil || wizard.deviceLoginCancel == nil { + return + } + wizard.deviceLoginCancel() + wizard.deviceLoginCancel = nil +} + // resetAimlapiOnboard drops any in-flight aimlapi.com sub-flow (cancelling a // running top-up stream) when the selected provider changes. func (wizard *providerWizardState) resetAimlapiOnboard() { @@ -591,6 +618,7 @@ func (wizard *providerWizardState) move(delta int) { wizard.modelLoading = false wizard.modelLoadError = "" wizard.oauthPending = false + wizard.cancelDeviceLogin() wizard.oauthErr = "" wizard.resetAimlapiOnboard() wizard.refreshModels() @@ -818,13 +846,17 @@ func (m model) handleProviderWizardKey(msg tea.KeyMsg) (model, tea.Cmd) { // provider list instead of destroying the overlay — and with it the user's // cursor and status — from a step deep in the flow. An in-flight OAuth login // is abandoned the same way; bumping the attempt id makes its late result - // stale so applyProviderWizardOAuth drops it. + // stale so applyProviderWizardOAuth drops it, and cancelDeviceLogin stops a + // device-code poll actually running in the background — otherwise + // completing the login later in the browser still silently saves the + // credential even though the wizard message is discarded. if m.providerWizard.manage && !m.providerWizard.managerStep() && keyIs(msg, tea.KeyEsc) { wizard := m.providerWizard if wizard.oauthPending { wizard.oauthPending = false wizard.oauthAttemptID++ } + wizard.cancelDeviceLogin() wizard.oauthDevice = false wizard.deviceUserCode = "" wizard.deviceVerificationURI = "" @@ -838,9 +870,12 @@ func (m model) handleProviderWizardKey(msg tea.KeyMsg) (model, tea.Cmd) { return m, nil } // While a browser/device OAuth login is in flight, ignore input except Esc, - // which abandons the wizard (the background flow times out and is dropped). + // which abandons the wizard. cancelDeviceLogin stops a device-code poll + // actually running in the background (the browser OAuth flow still relies + // on its own timeout, since it does not persist anything mid-flight). if m.providerWizard.oauthPending { if keyIs(msg, tea.KeyEsc) { + m.providerWizard.cancelDeviceLogin() m.providerWizard = nil } return m, nil @@ -868,10 +903,16 @@ func (m model) handleProviderWizardKey(msg tea.KeyMsg) (model, tea.Cmd) { return m, nil } // On the OAuth provider list, "d" forces device-code login for a device-capable - // provider (xAI) — useful on a desktop when the browser flow won't work. + // provider (xAI) — useful on a desktop when the browser flow won't work. A + // device-ONLY provider (Kimi Code has no loopback/authorize endpoint at + // all) has no browser flow to fall back to, so a plain Enter must also go + // straight to device login: the generic Enter/advanceProviderWizard path + // further below assumes a loopback flow exists and would otherwise hang or + // error against an endpoint that genuinely does not exist. if m.providerWizard.step == providerWizardStepProvider && m.providerWizard.oauthMode && - (keyText(msg) == "d" || keyText(msg) == "D") && - m.providerWizard.currentProvider().OAuthDeviceFlow { + m.providerWizard.currentProvider().OAuthDeviceFlow && + (keyText(msg) == "d" || keyText(msg) == "D" || + (keyIs(msg, tea.KeyEnter) && m.providerWizard.currentProvider().OAuthDeviceOnly)) { return m.startProviderDeviceLogin() } if m.providerWizard.step == providerWizardStepProvider { @@ -1463,6 +1504,9 @@ func (wizard *providerWizardState) footerText() string { case providerWizardStepMethod: return "↑/↓ move Enter/→ continue Esc close" case providerWizardStepProvider: + if wizard.oauthMode && wizard.currentProvider().OAuthDeviceOnly { + return "↑/↓ move Enter device code ← back Esc close" + } if wizard.oauthMode && wizard.currentProvider().OAuthDeviceFlow { return "↑/↓ move Enter sign in d device code ← back Esc close" } @@ -2083,12 +2127,25 @@ func providerWizardProfile(provider providercatalog.Descriptor, model string, ap APIFormat: providerWizardAPIFormat(provider), Model: firstProviderDisplayValue(model, provider.DefaultModel), } - // Catalog custom headers (e.g. aimlapi.com's partner attribution) belong to the - // catalog endpoint: the resolver only applies them when the base URL is the - // default, so a profile built against a staging/proxy/custom override must not - // bake them in either — otherwise attribution leaks to an arbitrary host. + // Catalog custom headers (e.g. aimlapi.com's partner attribution, or Kimi + // Code's X-Msh-* vendor-identity headers) belong to the catalog endpoint: + // the resolver only applies them when the base URL is the default, so a + // profile built against a staging/proxy/custom override must not bake + // them in either — otherwise attribution leaks to an arbitrary host. if sameProviderBaseURL(resolvedBaseURL, provider.DefaultBaseURL) { - profile.CustomHeaders = maps.Clone(provider.CustomHeaders) + // `provider` here normally comes from a listing call (OAuthProviders()/ + // All()), which deliberately leaves RuntimeHeaders-backed CustomHeaders + // (Kimi's device-identity headers) unset so merely browsing providers + // doesn't mint Kimi's on-disk device id. Re-resolve through Get(), which + // runs RuntimeHeaders the same way profile-resolve time does (see + // cloneDescriptor / applyCatalogDescriptor), so the wizard's first + // authenticated /models call and the profile it activates immediately + // afterward carry those headers without requiring a restart. + customHeaders := provider.CustomHeaders + if runtimeDescriptor, ok := providercatalog.Get(provider.ID); ok { + customHeaders = runtimeDescriptor.CustomHeaders + } + profile.CustomHeaders = maps.Clone(customHeaders) if providerWizardIsAimlapi(provider) { profile.CustomHeaders = aimlapi.WithResolvedPartnerHeader(profile.CustomHeaders) } diff --git a/internal/tui/provider_wizard_discovery.go b/internal/tui/provider_wizard_discovery.go index b0dcada69..688c443e4 100644 --- a/internal/tui/provider_wizard_discovery.go +++ b/internal/tui/provider_wizard_discovery.go @@ -43,8 +43,12 @@ func (m model) advanceProviderWizard() (model, tea.Cmd) { if m.providerWizard.step == providerWizardStepProvider && m.providerWizard.oauthMode && m.providerWizard.currentProvider().OAuth { provider := m.providerWizard.currentProvider() // Headless/SSH boxes can't open a browser — use device code there by - // default (the user can also force it with "d" from the list). - if provider.OAuthDeviceFlow && oauthPreferDeviceFlow() { + // default (the user can also force it with "d" from the list). A + // device-ONLY provider (Kimi Code has no loopback/authorize endpoint at + // all) must also go straight to device login here: this is the mouse + // double-click activation path, which bypasses the keyboard Enter + // handler's OAuthDeviceOnly check further up in provider_wizard.go. + if provider.OAuthDeviceFlow && (provider.OAuthDeviceOnly || oauthPreferDeviceFlow()) { return m.startProviderDeviceLogin() } attemptID := m.providerWizard.beginOAuthAttempt(false) diff --git a/internal/tui/provider_wizard_oauth_test.go b/internal/tui/provider_wizard_oauth_test.go index bd4ac4d14..b2cf078cd 100644 --- a/internal/tui/provider_wizard_oauth_test.go +++ b/internal/tui/provider_wizard_oauth_test.go @@ -1,6 +1,7 @@ package tui import ( + "context" "encoding/json" "errors" "os" @@ -8,6 +9,8 @@ import ( "strings" "testing" + tea "charm.land/bubbletea/v2" + "github.com/Gitlawb/zero/internal/config" "github.com/Gitlawb/zero/internal/providercatalog" ) @@ -106,6 +109,72 @@ func TestProviderWizardDeviceShortcutStartsDeviceFlow(t *testing.T) { } } +// TestProviderWizardEnterStartsDeviceFlowForDeviceOnlyProvider pins the fix +// for a device-only provider (Kimi Code has no loopback/authorize endpoint at +// all): the generic Enter path assumes a browser flow exists and would +// otherwise hang or error, so Enter must behave exactly like the "d" shortcut +// for a provider with OAuthDeviceOnly set. +func TestProviderWizardEnterStartsDeviceFlowForDeviceOnlyProvider(t *testing.T) { + // oauthPreferDeviceFlow() already defaults to device flow on a headless + // box (no DISPLAY/WAYLAND_DISPLAY, an SSH session, or ZERO_OAUTH_DEVICE + // set) — exactly the environment this test suite runs in — which would + // mask the bug this test exists to catch. Force a "normal desktop with a + // browser available" environment so Enter actually exercises the + // otherwise-loopback-preferring path. + t.Setenv("ZERO_OAUTH_DEVICE", "") + t.Setenv("SSH_CONNECTION", "") + t.Setenv("SSH_TTY", "") + t.Setenv("DISPLAY", ":0") + t.Setenv("WAYLAND_DISPLAY", "") + + m := mouseTestModel() + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() // → OAuth list + m = selectWizardOAuthProvider(t, next, "kimi-code") + if !m.providerWizard.currentProvider().OAuthDeviceOnly { + t.Fatal("test fixture assumes kimi-code is OAuthDeviceOnly") + } + + out, cmd := m.handleProviderWizardKey(tea.KeyPressMsg(tea.Key{Code: tea.KeyEnter})) + if !out.providerWizard.oauthPending || !out.providerWizard.oauthDevice { + t.Fatalf("Enter on a device-only provider should start device login (pending=%v device=%v)", out.providerWizard.oauthPending, out.providerWizard.oauthDevice) + } + if cmd == nil { + t.Fatal("Enter on a device-only provider should return the device-prepare command") + } +} + +// TestProviderWizardMouseAdvanceStartsDeviceFlowForDeviceOnlyProvider covers +// double-click activation: mouse.go routes to advanceProviderWizard, which +// bypasses the keyboard handler's OAuthDeviceOnly check. On a desktop (where +// oauthPreferDeviceFlow is false), advance must still start device login for +// a device-only provider so the verification URL/user code are not discarded. +func TestProviderWizardMouseAdvanceStartsDeviceFlowForDeviceOnlyProvider(t *testing.T) { + t.Setenv("ZERO_OAUTH_DEVICE", "") + t.Setenv("SSH_CONNECTION", "") + t.Setenv("SSH_TTY", "") + t.Setenv("DISPLAY", ":0") + t.Setenv("WAYLAND_DISPLAY", "") + + m := mouseTestModel() + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() // → OAuth list + m = selectWizardOAuthProvider(t, next, "kimi-code") + if !m.providerWizard.currentProvider().OAuthDeviceOnly { + t.Fatal("test fixture assumes kimi-code is OAuthDeviceOnly") + } + + out, cmd := m.advanceProviderWizard() + if !out.providerWizard.oauthPending || !out.providerWizard.oauthDevice { + t.Fatalf("mouse advance on a device-only provider should start device login (pending=%v device=%v)", out.providerWizard.oauthPending, out.providerWizard.oauthDevice) + } + if cmd == nil { + t.Fatal("mouse advance on a device-only provider should return the device-prepare command") + } +} + func TestProviderWizardDeviceCodeMsgShowsCodeAndPolls(t *testing.T) { m := mouseTestModel() m.providerWizard = m.newProviderWizard() @@ -129,6 +198,98 @@ func TestProviderWizardDeviceCodeMsgShowsCodeAndPolls(t *testing.T) { } } +// TestProviderWizardEscCancelsDeviceLoginPoll regression-tests a bug where +// abandoning the wizard with Esc during a device-code poll (phase 2) never +// canceled the background context the poll command runs on. The wizard +// message was made stale so the UI wouldn't show a stray login, but the +// underlying poll kept running for up to 10 minutes: if the user then +// finished authorizing in their browser, it would still silently succeed and +// save a credential the user believed they had backed out of. Esc must +// actually cancel the context, not just discard the eventual result. +func TestProviderWizardEscCancelsDeviceLoginPoll(t *testing.T) { + // Isolate the oauth token store the poll command touches (see + // managerTestModel), even though the canceled-context path exercised here + // returns before any read/write reaches it. + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + t.Setenv("ZERO_OAUTH_TOKENS_PATH", filepath.Join(t.TempDir(), "oauth-tokens.json")) + + m := mouseTestModel() + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() + m = selectWizardOAuthProvider(t, next, "xai") + providerID, attemptID := beginTestOAuthAttempt(m.providerWizard, true) + + out, cmd := m.applyProviderWizardDeviceCode(providerWizardDeviceCodeMsg{ + providerID: providerID, attemptID: attemptID, userCode: "ABCD-1234", verifyURL: "https://x.ai/device", + }) + if cmd == nil { + t.Fatal("device-code msg should start the poll command") + } + if out.providerWizard.deviceLoginCancel == nil { + t.Fatal("starting the poll should store a cancel func on the wizard") + } + + escaped, _ := out.handleProviderWizardKey(testKey(tea.KeyEsc)) + if escaped.providerWizard != nil { + t.Fatal("Esc while oauthPending should close the wizard") + } + + // The poll command captured the context created for this attempt; run it + // now (after Esc) and confirm CompleteDeviceLogin actually observed + // cancellation instead of running to completion in the background. + msg, ok := cmd().(providerWizardOAuthMsg) + if !ok { + t.Fatalf("poll command returned %T, want providerWizardOAuthMsg", msg) + } + if !errors.Is(msg.err, context.Canceled) { + t.Fatalf("poll error = %v, want context.Canceled (Esc should have canceled the background poll)", msg.err) + } +} + +// TestModelQuitCancelsProviderWizardDeviceLoginPoll regression-tests a bug +// where model.quit() (the path every Ctrl+C-to-exit and "q"-to-exit +// eventually reaches) only reset the aimlapi.com sub-flow, never the +// provider wizard's device-code poll. Since the TUI runs on +// context.Background(), quitting via Ctrl+C left the poll running in the +// background: authorizing the abandoned login afterward could still save a +// credential the user had just quit to avoid. +func TestModelQuitCancelsProviderWizardDeviceLoginPoll(t *testing.T) { + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + t.Setenv("ZERO_OAUTH_TOKENS_PATH", filepath.Join(t.TempDir(), "oauth-tokens.json")) + + m := mouseTestModel() + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() + m = selectWizardOAuthProvider(t, next, "xai") + providerID, attemptID := beginTestOAuthAttempt(m.providerWizard, true) + + out, cmd := m.applyProviderWizardDeviceCode(providerWizardDeviceCodeMsg{ + providerID: providerID, attemptID: attemptID, userCode: "ABCD-1234", verifyURL: "https://x.ai/device", + }) + if cmd == nil { + t.Fatal("device-code msg should start the poll command") + } + if out.providerWizard.deviceLoginCancel == nil { + t.Fatal("starting the poll should store a cancel func on the wizard") + } + + quit, _ := out.quit() + quitModel := quit.(model) + if quitModel.providerWizard != nil && quitModel.providerWizard.deviceLoginCancel != nil { + t.Fatal("quit should cancel the in-flight device-code poll") + } + + msg, ok := cmd().(providerWizardOAuthMsg) + if !ok { + t.Fatalf("poll command returned %T, want providerWizardOAuthMsg", msg) + } + if !errors.Is(msg.err, context.Canceled) { + t.Fatalf("poll error = %v, want context.Canceled (quit should have canceled the background poll)", msg.err) + } +} + // A failed OAuth attempt leaves the wizard on the provider list; the error must be // rendered there (not just on the credential step) so a click isn't a silent // no-op, and Hugging Face gets an actionable client_id hint. @@ -449,3 +610,33 @@ func TestAppendOAuthLoginProfileAddsOnceAndRespectsRenames(t *testing.T) { t.Fatalf("unknown provider must not append, got %+v", got) } } + +// TestProviderWizardProfileAppliesKimiRuntimeHeaders regression-tests a bug +// where the /provider wizard built its profile straight from the descriptor +// OAuthProviders() returns. That listing call deliberately omits +// RuntimeHeaders-backed CustomHeaders (Kimi's X-Msh-* vendor-identity +// headers) so merely browsing providers doesn't mint Kimi's on-disk device +// id — but that meant the wizard's first authenticated /models call and the +// profile it activated immediately after finishing were missing those +// headers until zero was restarted. providerWizardProfile must re-resolve +// through providercatalog.Get (which does run RuntimeHeaders) instead of +// using the listing descriptor's CustomHeaders directly. +func TestProviderWizardProfileAppliesKimiRuntimeHeaders(t *testing.T) { + m := mouseTestModel() + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() // → OAuth list + m = selectWizardOAuthProvider(t, next, "kimi-code") + provider := m.providerWizard.currentProvider() + if len(provider.CustomHeaders) != 0 { + t.Fatalf("test fixture assumes the OAuth listing omits runtime headers, got %#v", provider.CustomHeaders) + } + + profile := providerWizardProfile(provider, provider.DefaultModel, "", "", "") + if profile.CustomHeaders["X-Msh-Platform"] != "kimi_code_cli" { + t.Fatalf("providerWizardProfile CustomHeaders[X-Msh-Platform] = %q, want kimi_code_cli", profile.CustomHeaders["X-Msh-Platform"]) + } + if profile.CustomHeaders["X-Msh-Device-Id"] == "" { + t.Fatal("providerWizardProfile should carry Kimi's device id header") + } +}