From dc917ae6d2682ac61ef09520c2aad12b2d7d540a Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Mon, 14 Sep 2026 18:16:03 -0600 Subject: [PATCH 1/2] feat(ai): support Grok subscription authentication (fixes #68) --- .pylon/features.yaml | 14 ++ .pylon/upstream-review.md | 7 + .../ai/.changes/eng-6059-xai-subscription.md | 1 + packages/ai/README.md | 13 +- .../src/providers/openai-responses-shared.ts | 54 ++++- packages/ai/src/providers/openai-responses.ts | 1 + packages/ai/src/utils/oauth/index.ts | 4 +- packages/ai/src/utils/oauth/xai.ts | 227 ++++++++++++++++++ packages/ai/test/xai-oauth.test.ts | 139 +++++++++++ packages/ai/test/xai-responses.test.ts | 195 +++++++++++++++ .../.changes/eng-6059-xai-subscription.md | 1 + packages/coding-agent/README.md | 3 + packages/coding-agent/docs/providers.md | 13 + .../coding-agent/src/core/agent-session.ts | 32 ++- .../coding-agent/src/core/auth-storage.ts | 6 +- .../coding-agent/src/core/model-registry.ts | 84 ++++++- .../src/core/provider-display-names.ts | 2 +- packages/coding-agent/src/core/sdk.ts | 7 +- .../daemon-agent-connection.ts | 2 - .../in-process-agent-connection.ts | 10 +- .../src/modes/daemon/daemon-mode.ts | 18 +- .../src/modes/interactive/interactive-mode.ts | 18 ++ .../test/agent-connection-daemon.test.ts | 38 +++ .../test/agent-connection-in-process.test.ts | 5 + packages/coding-agent/test/auth-flows.test.ts | 15 +- .../test/configuration-menu.test.ts | 26 +- .../test/interactive-mode-status.test.ts | 3 + .../coding-agent/test/model-registry.test.ts | 36 ++- .../4575-model-auth-selection.test.ts | 37 ++- .../68-xai-subscription-auth.test.ts | 155 ++++++++++++ .../702-codex-client-version.test.ts | 15 +- packages/coding-agent/test/xai-auth.test.ts | 124 ++++++++++ 32 files changed, 1246 insertions(+), 59 deletions(-) create mode 100644 packages/ai/.changes/eng-6059-xai-subscription.md create mode 100644 packages/ai/src/utils/oauth/xai.ts create mode 100644 packages/ai/test/xai-oauth.test.ts create mode 100644 packages/ai/test/xai-responses.test.ts create mode 100644 packages/coding-agent/.changes/eng-6059-xai-subscription.md create mode 100644 packages/coding-agent/test/suite/regressions/68-xai-subscription-auth.test.ts create mode 100644 packages/coding-agent/test/xai-auth.test.ts diff --git a/.pylon/features.yaml b/.pylon/features.yaml index 78f47f6ace..3c9063cac0 100644 --- a/.pylon/features.yaml +++ b/.pylon/features.yaml @@ -434,3 +434,17 @@ decisions: upstream_support: Absorb kernel pipe errors, park blocked idle waiters, and retry stale Codex continuation IDs once before output; retain Pylon lifecycle, cancellation, scoped identity, recovery cursors and ownership without wire changes. revisit_when: - The next complete upstream integration includes these commits; reconcile the selective patches without duplicating behavior. + + grok-subscription-authentication: + area: provider-authentication + state: candidate + owner: shared + decision: hybridize + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/68 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/pull/2252 + fork_change: grok-subscription-with-owned-provider-routing + upstream_support: Adopt device-code authentication and auth-specific Responses routing while retaining Pylon scoped provider hooks, recovery cursors and the three-tab configuration menu. + revisit_when: + - The complete upstream integration includes equivalent subscription routing and preserves Pylon ownership contracts. diff --git a/.pylon/upstream-review.md b/.pylon/upstream-review.md index a50295bcc0..e878343f38 100644 --- a/.pylon/upstream-review.md +++ b/.pylon/upstream-review.md @@ -334,3 +334,10 @@ Follow-up: Task10 builds/packs the exact merged tree into a private prefix and r - Adopt upstream #2314 (`7d1913969d33fa5a31366321e3f5db0070dcf2bf`): bind Codex continuation state to the producing connection and retry `previous_response_not_found` once with full context only before output. Preserve scoped provider hooks and avoid duplicate stream events or prompt replay. - Compatibility: no daemon commands, events, response shapes, protocol/schema revisions or public capability tokens change. No credentials, managed runtime, live session or release is modified. Existing installations require a separately verified immutable managed build to receive these fixes. - Validation: 125 focused coding-agent tests (including the real isolated kernel pipe regression) and 25 Codex stream tests pass. `npm run check` passes Biome, TypeScript, installer and browser bundle checks. Required hosted validation and final review are recorded in the owning PR. + +## 2026-09-14 — Grok subscription authentication + +- Tracking #68. Hybridize upstream #2252 (`ca67580b524800a971f92420d4c665bd83df1dfa`), reviewed against current upstream `f5859162c`. Retain scoped provider hooks, owned-session recovery and Pylon's three-tab configuration menu. The full integration checkpoint remains `1eee2938b4eeb7a4d72e17035adda669a89b63de`. +- Device-code OAuth and subscription refresh select xAI Responses request models; API keys retain their original models. Authentication changes refresh active, maintenance and daemon model views. Failed subscription refresh does not silently fall back to an API key. Shared response streaming handles interleaved tool arguments and encrypted reasoning while preserving Codex behavior. +- No daemon commands, events, response shapes, schema revision or negotiated capabilities change. Existing model/catalog refresh commands carry the updated values. No credentials or installed runtime are modified by this adoption. +- Validation: 580 focused tests pass across AI streaming/OAuth and coding-agent authentication, model selection, maintenance and daemon/in-process connections; four optional AI tests are skipped. `npm run check` passes formatting, types, installer and browser smoke. Live paid-provider access is not part of these faux-provider tests. Required hosted checks and final review accompany the PR. diff --git a/packages/ai/.changes/eng-6059-xai-subscription.md b/packages/ai/.changes/eng-6059-xai-subscription.md new file mode 100644 index 0000000000..280c02ec15 --- /dev/null +++ b/packages/ai/.changes/eng-6059-xai-subscription.md @@ -0,0 +1 @@ +- Added xAI subscription device-code authentication and Responses support for all bundled Grok tool models while preserving API-key access. diff --git a/packages/ai/README.md b/packages/ai/README.md index 90d4819bc6..b79788b032 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -76,7 +76,7 @@ Unified LLM API with automatic model discovery, provider configuration, token an - **Cerebras** - **Cloudflare AI Gateway** - **Cloudflare Workers AI** -- **xAI** +- **xAI** (API key or Grok/X subscription) - **OpenRouter** - **Vercel AI Gateway** - **MiniMax** @@ -1126,14 +1126,23 @@ const key = getEnvApiKey('openai'); // checks OPENAI_API_KEY ## OAuth Providers -Several providers require OAuth authentication instead of static API keys: +OAuth authentication is available for these providers and subscriptions: - **Anthropic** (Claude Pro/Max subscription) - **OpenAI Codex** (ChatGPT Plus/Pro subscription, access to GPT-5.x Codex models) - **GitHub Copilot** (Copilot subscription) +- **xAI** (Grok/X subscription, device-code login) For paid Cloud Code Assist subscriptions, set `GOOGLE_CLOUD_PROJECT` or `GOOGLE_CLOUD_PROJECT_ID` to your project ID. +### xAI subscription + +In Prime Agent, run `/login` and select the **xAI subscription** entry. Open the displayed HTTPS URL and enter the device code. `XAI_API_KEY` and API-key login remain supported. Stored subscription credentials take priority over `XAI_API_KEY`; an explicit `--api-key` override uses API-key routing. + +All bundled xAI tool-capable language models can use subscription authentication through the Responses API. Each model keeps its own reasoning and input capabilities; only verified reasoning-effort controls are sent. Account eligibility, available models, and usage limits are controlled by xAI; signing in does not guarantee access or unlimited usage. + +SDK callers can use `loginXai(callbacks)` and `getOAuthApiKey("xai", credentials)` from `prime-agent-ai/oauth`. Persist refreshed credentials securely. When the effective credential is an xAI subscription, pass `getXaiSubscriptionModel(getModel("xai", "grok-4.5"))` to `stream` or `complete` with the resolved access token. The helper accepts configured xAI model descriptors and returns `undefined` for other providers. Do not apply this projection to API keys: the generated API-key models remain unchanged. Subscription requests use `https://api.x.ai/v1`. + ### Vertex AI Vertex AI models support either a Google Cloud API key or Application Default Credentials (ADC): diff --git a/packages/ai/src/providers/openai-responses-shared.ts b/packages/ai/src/providers/openai-responses-shared.ts index ae312d9457..8eee761941 100644 --- a/packages/ai/src/providers/openai-responses-shared.ts +++ b/packages/ai/src/providers/openai-responses-shared.ts @@ -164,6 +164,7 @@ export function convertResponsesMessages( if (block.type === "thinking") { if (block.thinkingSignature) { const reasoningItem = JSON.parse(block.thinkingSignature) as ResponseReasoningItem; + if (model.provider === "xai") delete reasoningItem.status; output.push(reasoningItem); } } else if (block.type === "text") { @@ -274,13 +275,29 @@ export async function processResponsesStream( ): Promise { let currentItem: ResponseReasoningItem | ResponseOutputMessage | ResponseFunctionToolCall | null = null; let currentBlock: ThinkingContent | TextContent | (ToolCall & { partialJson: string }) | null = null; - const blocks = output.content; - const blockIndex = () => blocks.length - 1; + let currentContentIndex = -1; + let sawTerminalResponse = false; + const slots = new Map< + number, + { + item: ResponseReasoningItem | ResponseOutputMessage | ResponseFunctionToolCall; + block: ThinkingContent | TextContent | (ToolCall & { partialJson: string }); + contentIndex: number; + } + >(); + const blockIndex = () => currentContentIndex; for await (const event of openaiStream) { + if ("output_index" in event) { + const slot = slots.get(event.output_index); + currentItem = slot?.item ?? null; + currentBlock = slot?.block ?? null; + currentContentIndex = slot?.contentIndex ?? output.content.length; + } if (event.type === "response.created") { output.responseId = event.response.id; } else if (event.type === "response.output_item.added") { + currentContentIndex = output.content.length; const item = event.item; if (item.type === "reasoning") { currentItem = item; @@ -304,6 +321,13 @@ export async function processResponsesStream( output.content.push(currentBlock); stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output }); } + if (currentItem && currentBlock) { + slots.set(event.output_index, { + item: currentItem, + block: currentBlock, + contentIndex: currentContentIndex, + }); + } } else if (event.type === "response.reasoning_summary_part.added") { if (currentItem && currentItem.type === "reasoning") { currentItem.summary = currentItem.summary || []; @@ -420,6 +444,7 @@ export async function processResponsesStream( } } } else if (event.type === "response.output_item.done") { + slots.delete(event.output_index); const item = event.item; if (item.type === "reasoning" && currentBlock?.type === "thinking") { @@ -445,10 +470,9 @@ export async function processResponsesStream( }); currentBlock = null; } else if (item.type === "function_call") { - const args = - currentBlock?.type === "toolCall" && currentBlock.partialJson - ? parseStreamingJson(currentBlock.partialJson) - : parseStreamingJson(item.arguments || "{}"); + const args = parseStreamingJson( + item.arguments || (currentBlock?.type === "toolCall" ? currentBlock.partialJson : "") || "{}", + ); let toolCall: ToolCall; if (currentBlock?.type === "toolCall") { @@ -469,8 +493,21 @@ export async function processResponsesStream( currentBlock = null; stream.push({ type: "toolcall_end", contentIndex: blockIndex(), toolCall, partial: output }); } - } else if (event.type === "response.completed") { + } else if (event.type === "response.completed" || event.type === "response.incomplete") { + sawTerminalResponse = true; const response = event.response; + if (model.provider === "xai") { + for (const item of response.output ?? []) { + if (item.type !== "reasoning" || !item.encrypted_content) continue; + for (const block of output.content) { + if (block.type !== "thinking" || !block.thinkingSignature) continue; + const stored = JSON.parse(block.thinkingSignature) as ResponseReasoningItem; + if (stored.id === item.id && !stored.encrypted_content) { + block.thinkingSignature = JSON.stringify({ ...stored, encrypted_content: item.encrypted_content }); + } + } + } + } if (response?.id) { output.responseId = response.id; } @@ -520,6 +557,9 @@ export async function processResponsesStream( }); } } + if (model.provider === "xai" && !sawTerminalResponse) { + throw new StreamFailureError("xAI Responses stream ended before a terminal response event", { kind: "unknown" }); + } } function mapStopReason(status: OpenAI.Responses.ResponseStatus | undefined): StopReason { diff --git a/packages/ai/src/providers/openai-responses.ts b/packages/ai/src/providers/openai-responses.ts index 5186d1a25e..945d889471 100644 --- a/packages/ai/src/providers/openai-responses.ts +++ b/packages/ai/src/providers/openai-responses.ts @@ -264,6 +264,7 @@ function buildParams(model: Model<"openai-responses">, context: Context, options effort: (model.thinkingLevelMap?.off ?? "none") as NonNullable["effort"], }; } + if (model.provider === "xai") params.include = ["reasoning.encrypted_content"]; } return params; diff --git a/packages/ai/src/utils/oauth/index.ts b/packages/ai/src/utils/oauth/index.ts index 6c0b356ed3..90b5bc7158 100644 --- a/packages/ai/src/utils/oauth/index.ts +++ b/packages/ai/src/utils/oauth/index.ts @@ -16,18 +16,20 @@ export { refreshGitHubCopilotToken, } from "./github-copilot.js"; export { loginOpenAICodex, openaiCodexOAuthProvider, refreshOpenAICodexToken } from "./openai-codex.js"; - export * from "./types.js"; +export { getXaiSubscriptionModel, loginXai, refreshXaiToken, xaiOAuthProvider } from "./xai.js"; import { anthropicOAuthProvider } from "./anthropic.js"; import { githubCopilotOAuthProvider } from "./github-copilot.js"; import { openaiCodexOAuthProvider } from "./openai-codex.js"; import type { OAuthCredentials, OAuthProviderId, OAuthProviderInfo, OAuthProviderInterface } from "./types.js"; +import { xaiOAuthProvider } from "./xai.js"; const BUILT_IN_OAUTH_PROVIDERS: OAuthProviderInterface[] = [ anthropicOAuthProvider, githubCopilotOAuthProvider, openaiCodexOAuthProvider, + xaiOAuthProvider, ]; const oauthProviderRegistry = new Map( diff --git a/packages/ai/src/utils/oauth/xai.ts b/packages/ai/src/utils/oauth/xai.ts new file mode 100644 index 0000000000..9cb4a46575 --- /dev/null +++ b/packages/ai/src/utils/oauth/xai.ts @@ -0,0 +1,227 @@ +import type { Api, Model } from "../../types.js"; +import type { OAuthCredentials, OAuthLoginCallbacks, OAuthProviderInterface } from "./types.js"; + +const CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828"; +const SCOPE = "openid profile email offline_access grok-cli:access api:access"; +const DEVICE_CODE_URL = "https://auth.x.ai/oauth2/device/code"; +const TOKEN_URL = "https://auth.x.ai/oauth2/token"; +const REQUEST_TIMEOUT_MS = 30_000; +const REFRESH_SKEW_MS = 5 * 60 * 1000; + +type JsonObject = Record; +type OAuthResponse = { ok: boolean; status: number; body: JsonObject }; + +function requiredString(body: JsonObject, field: string): string { + const value = body[field]; + if (typeof value !== "string" || !value.trim()) throw new Error(`Invalid xAI OAuth response field: ${field}`); + return value; +} + +function positiveSeconds(value: unknown, field: string): number { + if ( + typeof value !== "number" || + !Number.isFinite(value) || + value <= 0 || + value * 1000 > Number.MAX_SAFE_INTEGER - Date.now() + ) { + throw new Error(`Invalid xAI OAuth response field: ${field}`); + } + return value; +} + +function verificationUri(raw: string): string { + let url: URL; + try { + url = new URL(raw); + } catch { + throw new Error("Untrusted verification URI in xAI OAuth response"); + } + if (url.protocol !== "https:" || url.username || url.password || /[\u0000-\u0020\u007f-\u009f]/.test(raw)) { + throw new Error("Untrusted verification URI in xAI OAuth response"); + } + return url.href; +} + +function checkCancelled(signal?: AbortSignal): void { + if (signal?.aborted) throw new Error("Login cancelled"); +} + +async function postForm( + url: string, + fields: Record, + signal?: AbortSignal, + timeoutMs = REQUEST_TIMEOUT_MS, +): Promise { + checkCancelled(signal); + const controller = new AbortController(); + const onAbort = () => controller.abort(); + signal?.addEventListener("abort", onAbort, { once: true }); + const timeout = setTimeout(() => controller.abort(), Math.min(timeoutMs, REQUEST_TIMEOUT_MS)); + try { + const response = await fetch(url, { + method: "POST", + headers: { Accept: "application/json", "Content-Type": "application/x-www-form-urlencoded" }, + body: new URLSearchParams(fields), + signal: controller.signal, + redirect: "error", + }); + let parsed: unknown; + try { + parsed = await response.json(); + } catch { + if (controller.signal.aborted) throw new Error("Request aborted"); + throw new Error(`xAI OAuth returned invalid JSON (HTTP ${response.status})`); + } + checkCancelled(signal); + if (controller.signal.aborted) throw new Error("Request aborted"); + return { + ok: response.ok, + status: response.status, + body: parsed && typeof parsed === "object" && !Array.isArray(parsed) ? (parsed as JsonObject) : {}, + }; + } catch (error) { + checkCancelled(signal); + if (controller.signal.aborted) throw new Error("xAI OAuth request timed out. Try signing in again."); + if (error instanceof Error && error.message.startsWith("xAI OAuth returned invalid JSON")) throw error; + throw new Error("xAI OAuth request failed. Check your connection and try again."); + } finally { + clearTimeout(timeout); + signal?.removeEventListener("abort", onAbort); + } +} + +function requestFailure(action: string, response: OAuthResponse): Error { + // Do not print arbitrary provider response bodies: they may echo credentials. + const code = response.body.error === "invalid_grant" ? ": authorization expired or revoked; sign in again" : ""; + return new Error(`xAI OAuth ${action} failed (HTTP ${response.status})${code}`); +} + +function credentialsFromResponse(body: JsonObject, previousRefresh?: string): OAuthCredentials { + const access = requiredString(body, "access_token"); + const refresh = + body.refresh_token === undefined && previousRefresh ? previousRefresh : requiredString(body, "refresh_token"); + const lifetimeMs = positiveSeconds(body.expires_in === undefined ? 3600 : body.expires_in, "expires_in") * 1000; + return { access, refresh, expires: Date.now() + lifetimeMs - Math.min(REFRESH_SKEW_MS, lifetimeMs / 2) }; +} + +function wait(ms: number, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(new Error("Login cancelled")); + return; + } + const onAbort = () => { + clearTimeout(timeout); + reject(new Error("Login cancelled")); + }; + const timeout = setTimeout(() => { + signal?.removeEventListener("abort", onAbort); + resolve(); + }, ms); + signal?.addEventListener("abort", onAbort, { once: true }); + }); +} + +export async function loginXai(callbacks: OAuthLoginCallbacks): Promise { + const response = await postForm( + DEVICE_CODE_URL, + { client_id: CLIENT_ID, scope: SCOPE, referrer: "pi" }, + callbacks.signal, + ); + if (!response.ok) throw requestFailure("device authorization", response); + const deviceCode = requiredString(response.body, "device_code"); + const userCode = requiredString(response.body, "user_code"); + if (!/^[A-Za-z0-9-]+$/.test(userCode)) throw new Error("Invalid xAI OAuth response field: user_code"); + const url = verificationUri(requiredString(response.body, "verification_uri")); + const deadline = Date.now() + positiveSeconds(response.body.expires_in, "expires_in") * 1000; + const interval = response.body.interval; + let intervalMs = + typeof interval === "number" && Number.isFinite(interval) && interval > 0 + ? Math.max(1000, interval * 1000) + : 5000; + callbacks.onAuth({ url, instructions: `Enter code: ${userCode}` }); + while (Date.now() < deadline) { + await wait(Math.min(intervalMs, deadline - Date.now(), 2_147_483_647), callbacks.signal); + checkCancelled(callbacks.signal); + if (Date.now() >= deadline) break; + const token = await postForm( + TOKEN_URL, + { + grant_type: "urn:ietf:params:oauth:grant-type:device_code", + client_id: CLIENT_ID, + device_code: deviceCode, + }, + callbacks.signal, + deadline - Date.now(), + ); + if (token.ok) return credentialsFromResponse(token.body); + if (token.body.error === "authorization_pending") continue; + if (token.body.error === "slow_down") { + const next = token.body.interval; + intervalMs = + typeof next === "number" && Number.isFinite(next) && next > 0 + ? Math.max(intervalMs + 5000, next * 1000) + : intervalMs + 5000; + continue; + } + if (token.body.error === "access_denied" || token.body.error === "authorization_denied") + throw new Error("xAI device authorization was denied"); + if (token.body.error === "expired_token") throw new Error("xAI device code expired; sign in again"); + throw requestFailure("device token polling", token); + } + throw new Error("xAI device code expired; sign in again"); +} + +export async function refreshXaiToken(refreshToken: string, signal?: AbortSignal): Promise { + const response = await postForm( + TOKEN_URL, + { grant_type: "refresh_token", client_id: CLIENT_ID, refresh_token: refreshToken }, + signal, + ); + if (!response.ok) throw requestFailure("token refresh", response); + return credentialsFromResponse(response.body, refreshToken); +} + +export function getXaiSubscriptionModel(model: Model): Model<"openai-responses"> | undefined { + if (model.provider !== "xai") return undefined; + let thinkingLevelMap = model.thinkingLevelMap; + if (!thinkingLevelMap) { + switch (model.id) { + case "grok-4.3": + thinkingLevelMap = { off: "none", minimal: null }; + break; + case "grok-4.5": + thinkingLevelMap = { off: null, minimal: null }; + break; + case "grok-4.6": + thinkingLevelMap = { off: null, minimal: null, xhigh: "xhigh" }; + break; + default: + // Keep reasoning output without sending unverified effort controls. + thinkingLevelMap = { + off: null, + minimal: null, + low: null, + medium: null, + high: null, + xhigh: null, + max: null, + }; + } + } + return { + ...model, + api: "openai-responses", + baseUrl: "https://api.x.ai/v1", + thinkingLevelMap, + compat: { supportsLongCacheRetention: false }, + }; +} + +export const xaiOAuthProvider: OAuthProviderInterface = { + id: "xai", + name: "xAI (Grok)", + login: loginXai, + refreshToken: (credentials) => refreshXaiToken(credentials.refresh), + getApiKey: (credentials) => credentials.access, +}; diff --git a/packages/ai/test/xai-oauth.test.ts b/packages/ai/test/xai-oauth.test.ts new file mode 100644 index 0000000000..b2c850a7df --- /dev/null +++ b/packages/ai/test/xai-oauth.test.ts @@ -0,0 +1,139 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { loginXai, refreshXaiToken } from "../src/utils/oauth/xai.js"; + +const DEVICE_URL = "https://auth.x.ai/oauth2/device/code"; +const TOKEN_URL = "https://auth.x.ai/oauth2/token"; +const CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828"; +const device = { + device_code: "secret-device", + user_code: "ABCD-1234", + verification_uri: "https://accounts.x.ai/oauth2/device", + expires_in: 900, + interval: 5, +}; +const token = { access_token: "secret-access", refresh_token: "secret-refresh", expires_in: 21600 }; +function json(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { status, headers: { "Content-Type": "application/json" } }); +} +function login(onAuth: Parameters[0]["onAuth"] = vi.fn(), signal?: AbortSignal) { + return loginXai({ onAuth, onPrompt: vi.fn(), signal }); +} +function pendingFetch(_input: unknown, init?: RequestInit): Promise { + return new Promise((_resolve, reject) => { + init?.signal?.addEventListener("abort", () => reject(new Error("abort")), { once: true }); + }); +} + +describe("xAI device OAuth", () => { + beforeEach(() => { + vi.useFakeTimers(); + vi.setSystemTime(0); + }); + afterEach(() => { + vi.restoreAllMocks(); + vi.unstubAllGlobals(); + vi.useRealTimers(); + }); + + it("uses the device grant, first-poll delay, pending and slow_down", async () => { + const times: number[] = []; + const replies = [ + json({ error: "authorization_pending" }, 400), + json({ error: "slow_down", interval: 10 }, 400), + json(token), + ]; + vi.stubGlobal( + "fetch", + vi.fn(async (url: string, init: RequestInit) => { + expect(init.redirect).toBe("error"); + const form = new URLSearchParams(String(init.body)); + expect(form.get("client_id")).toBe(CLIENT_ID); + if (url === DEVICE_URL) { + expect(form.get("scope")).toBe("openid profile email offline_access grok-cli:access api:access"); + expect(form.get("referrer")).toBe("pi"); + return json(device); + } + expect(url).toBe(TOKEN_URL); + expect(form.get("grant_type")).toBe("urn:ietf:params:oauth:grant-type:device_code"); + expect(form.get("device_code")).toBe("secret-device"); + times.push(Date.now()); + return replies.shift()!; + }), + ); + const onAuth = vi.fn(); + const result = login(onAuth); + await vi.advanceTimersByTimeAsync(0); + expect(onAuth).toHaveBeenCalledWith({ url: device.verification_uri, instructions: "Enter code: ABCD-1234" }); + expect(times).toEqual([]); + await vi.advanceTimersByTimeAsync(20000); + expect(times).toEqual([5000, 10000, 20000]); + expect(await result).toEqual({ + access: token.access_token, + refresh: token.refresh_token, + expires: 20000 + 21600000 - 300000, + }); + }); + + it("cancels an in-flight request", async () => { + const controller = new AbortController(); + vi.stubGlobal("fetch", vi.fn(pendingFetch)); + const result = login(vi.fn(), controller.signal); + controller.abort(); + await expect(result).rejects.toThrow("Login cancelled"); + }); + + it("bounds token polling by the remaining device lifetime", async () => { + const fetchMock = vi + .fn() + .mockResolvedValueOnce(json({ ...device, expires_in: 6 })) + .mockImplementation(pendingFetch); + vi.stubGlobal("fetch", fetchMock); + const assertion = expect(login()).rejects.toThrow("timed out"); + await vi.advanceTimersByTimeAsync(6000); + await assertion; + expect(fetchMock).toHaveBeenCalledTimes(2); + }); + + it("refreshes with the device client and preserves an omitted refresh token", async () => { + const fetchMock = vi.fn(async (url: string, init: RequestInit) => { + expect(url).toBe(TOKEN_URL); + expect(Object.fromEntries(new URLSearchParams(String(init.body)))).toEqual({ + grant_type: "refresh_token", + client_id: CLIENT_ID, + refresh_token: "old-refresh", + }); + return json({ access_token: "new-access" }); + }); + vi.stubGlobal("fetch", fetchMock); + expect(await refreshXaiToken("old-refresh")).toEqual({ + access: "new-access", + refresh: "old-refresh", + expires: 3300000, + }); + fetchMock.mockResolvedValue(json(token)); + expect((await refreshXaiToken("old-refresh")).refresh).toBe("secret-refresh"); + }); + + it("rejects an unsafe verification URL before exposing it to the UI", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(json({ ...device, verification_uri: "file:///etc/passwd" }))); + const onAuth = vi.fn(); + await expect(login(onAuth)).rejects.toThrow("Untrusted verification URI"); + expect(onAuth).not.toHaveBeenCalled(); + }); + + it("rejects malformed tokens and redacts provider failures", async () => { + vi.stubGlobal( + "fetch", + vi + .fn() + .mockResolvedValueOnce(json({ ...token, access_token: "" })) + .mockResolvedValueOnce( + json({ error: "invalid_grant", error_description: "secret-refresh\u001b[31m" }, 400), + ), + ); + await expect(refreshXaiToken("secret-refresh")).rejects.toThrow("Invalid xAI OAuth response field"); + await expect(refreshXaiToken("secret-refresh")).rejects.toEqual( + new Error("xAI OAuth token refresh failed (HTTP 400): authorization expired or revoked; sign in again"), + ); + }); +}); diff --git a/packages/ai/test/xai-responses.test.ts b/packages/ai/test/xai-responses.test.ts new file mode 100644 index 0000000000..d8d137963b --- /dev/null +++ b/packages/ai/test/xai-responses.test.ts @@ -0,0 +1,195 @@ +import { Type } from "typebox"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { getModel, getSupportedThinkingLevels } from "../src/models.js"; +import { streamSimpleOpenAIResponses } from "../src/providers/openai-responses.js"; +import type { Context } from "../src/types.js"; +import { getXaiSubscriptionModel } from "../src/utils/oauth/xai.js"; + +const model = getXaiSubscriptionModel(getModel("xai", "grok-4.5"))!; +const reasoning = { + type: "reasoning", + id: "rs_grok", + summary: [], + content: [{ type: "reasoning_text", text: "Inspect the file." }], + encrypted_content: "encrypted-reasoning", + status: "completed", +}; +const call = { + type: "function_call", + id: "fc_grok", + call_id: "call_grok", + name: "ipython", + arguments: '{"code":"1+1"}', + status: "completed", +}; +const text = { + type: "message", + id: "msg_grok", + role: "assistant", + status: "completed", + content: [{ type: "output_text", text: "2", annotations: [] }], +}; +function sse(events: unknown[]): Response { + return new Response(`${events.map((event) => `data: ${JSON.stringify(event)}\n\n`).join("")}data: [DONE]\n\n`, { + headers: { "Content-Type": "text/event-stream" }, + }); +} +function terminal(output: unknown[] = [], status = "completed") { + return { + type: `response.${status}`, + response: { + id: "resp_grok", + status, + output, + usage: { input_tokens: 20, output_tokens: 8, total_tokens: 28, input_tokens_details: { cached_tokens: 5 } }, + }, + }; +} +function toolEvents(): Record[] { + return [ + { type: "response.created", response: { id: "resp_grok" } }, + { type: "response.output_item.added", output_index: 0, item: { type: "reasoning", id: "rs_grok", summary: [] } }, + { type: "response.output_item.added", output_index: 1, item: { ...call, arguments: "" } }, + { type: "response.reasoning_text.delta", output_index: 0, delta: "Inspect the file." }, + { type: "response.function_call_arguments.delta", output_index: 1, delta: '{"code":' }, + { type: "response.function_call_arguments.done", output_index: 1, arguments: call.arguments }, + { type: "response.output_item.done", output_index: 0, item: { ...reasoning, encrypted_content: undefined } }, + { type: "response.output_item.done", output_index: 1, item: call }, + terminal([reasoning, call]), + ]; +} +function textEvents() { + return [ + { type: "response.output_item.added", output_index: 0, item: { ...text, content: [] } }, + { + type: "response.content_part.added", + output_index: 0, + part: { type: "output_text", text: "", annotations: [] }, + }, + { type: "response.output_text.delta", output_index: 0, delta: "2" }, + { type: "response.output_item.done", output_index: 0, item: text }, + terminal([text]), + ]; +} + +describe("xAI subscription Responses", () => { + afterEach(() => { + vi.restoreAllMocks(); + vi.unstubAllGlobals(); + }); + + it("preserves model capabilities and sends only verified reasoning efforts", async () => { + let body: Record = {}; + vi.stubGlobal( + "fetch", + vi.fn(async (input: Parameters[0], init?: RequestInit) => { + body = JSON.parse(await new Request(input, init).text()); + return sse(textEvents()); + }), + ); + for (const [id, requested, effort] of [ + ["grok-4.3", undefined, "none"], + ["grok-4.3", "low", "low"], + ["grok-4.5", "xhigh", "high"], + ["grok-4.6", "xhigh", "xhigh"], + ["grok-4.20-0309-reasoning", "high", undefined], + ["grok-4.20-0309-non-reasoning", "high", undefined], + ["grok-build-0.1", "high", undefined], + ] as const) { + const source = getModel("xai", id); + const adapted = getXaiSubscriptionModel(source)!; + expect(adapted).toMatchObject({ ...source, api: "openai-responses" }); + const result = await streamSimpleOpenAIResponses( + adapted, + { + messages: [{ role: "user", content: "hello", timestamp: 1 }], + }, + { apiKey: "test-subscription-token", reasoning: requested }, + ).result(); + expect(result.stopReason, result.errorMessage).toBe("stop"); + expect(body.model).toBe(id); + if (effort) expect(body.reasoning).toMatchObject({ effort }); + else expect(body).not.toHaveProperty("reasoning"); + expect(body.include).toEqual(source.reasoning ? ["reasoning.encrypted_content"] : undefined); + } + expect(getSupportedThinkingLevels(getXaiSubscriptionModel(getModel("xai", "grok-4.6"))!)).toEqual([ + "low", + "medium", + "high", + "xhigh", + ]); + expect(getSupportedThinkingLevels(getXaiSubscriptionModel(getModel("xai", "grok-build-0.1"))!)).toEqual([]); + }); + + it("streams interleaved thinking/tool calls and replays a complete second turn", async () => { + const requests: { url: string; headers: Headers; body: Record }[] = []; + vi.stubGlobal( + "fetch", + vi.fn(async (input: Parameters[0], init?: RequestInit) => { + const request = new Request(input, init); + requests.push({ url: request.url, headers: request.headers, body: JSON.parse(await request.text()) }); + return sse(requests.length === 1 ? toolEvents() : textEvents()); + }), + ); + const context: Context = { + systemPrompt: "Use the kernel.", + messages: [{ role: "user", content: "Calculate 1+1", timestamp: 1 }], + tools: [{ name: "ipython", description: "Execute Python", parameters: Type.Object({ code: Type.String() }) }], + }; + const stream = streamSimpleOpenAIResponses(model, context, { + apiKey: "test-subscription-token", + reasoning: "medium", + cacheRetention: "long", + sessionId: "session-grok", + }); + const events = []; + for await (const event of stream) events.push(event); + const first = await stream.result(); + expect(first.stopReason, first.errorMessage).toBe("toolUse"); + const thinking = first.content.find((block) => block.type === "thinking"); + expect(thinking?.thinking).toBe("Inspect the file."); + expect(JSON.parse(thinking?.thinkingSignature ?? "{}")).toEqual(reasoning); + const tool = first.content.find((block) => block.type === "toolCall"); + if (!tool) throw new Error("Missing tool call"); + expect(tool.arguments).toEqual({ code: "1+1" }); + expect(events).toEqual( + expect.arrayContaining([ + expect.objectContaining({ type: "thinking_delta", contentIndex: 0 }), + expect.objectContaining({ type: "toolcall_delta", contentIndex: 1 }), + ]), + ); + context.messages.push(first, { + role: "toolResult", + toolCallId: tool.id, + toolName: tool.name, + content: [{ type: "text", text: "2" }], + isError: false, + timestamp: 2, + }); + const second = await streamSimpleOpenAIResponses(model, context, { + apiKey: "test-subscription-token", + }).result(); + expect(second.stopReason, second.errorMessage).toBe("stop"); + expect(second.content[0]).toMatchObject({ type: "text", text: "2" }); + expect(requests[0].url).toBe("https://api.x.ai/v1/responses"); + expect(requests[0].headers.get("authorization")).toBe("Bearer test-subscription-token"); + expect(requests[0].body).toMatchObject({ + model: "grok-4.5", + store: false, + stream: true, + reasoning: { effort: "medium" }, + include: ["reasoning.encrypted_content"], + }); + expect(requests[0].body).not.toHaveProperty("prompt_cache_retention"); + expect(requests[1].body.include).toEqual(["reasoning.encrypted_content"]); + const replay = requests[1].body.input as Record[]; + const replayReasoning = replay.find((item) => item.type === "reasoning"); + expect(replayReasoning).toMatchObject({ id: "rs_grok", encrypted_content: "encrypted-reasoning" }); + expect(replayReasoning).not.toHaveProperty("status"); + const replayCall = replay.find((item) => item.type === "function_call"); + const replayResult = replay.find((item) => item.type === "function_call_output"); + expect(replayCall?.call_id).toBeTruthy(); + expect(replayResult?.call_id).toBe(replayCall?.call_id); + expect(replayResult?.output).toBe("2"); + }); +}); diff --git a/packages/coding-agent/.changes/eng-6059-xai-subscription.md b/packages/coding-agent/.changes/eng-6059-xai-subscription.md new file mode 100644 index 0000000000..78504e3784 --- /dev/null +++ b/packages/coding-agent/.changes/eng-6059-xai-subscription.md @@ -0,0 +1 @@ +- Added xAI Grok subscription authentication through the existing `/login` menu for all bundled tool models, with auth changes applied to the current session. diff --git a/packages/coding-agent/README.md b/packages/coding-agent/README.md index 5f6bd39454..5d1036cc44 100644 --- a/packages/coding-agent/README.md +++ b/packages/coding-agent/README.md @@ -80,6 +80,9 @@ For each built-in provider, Prime Agent maintains a list of tool-capable models, - Anthropic Claude Pro/Max - OpenAI ChatGPT Plus/Pro (Codex) - GitHub Copilot +- xAI Grok (eligible subscriptions) + +Select the xAI subscription entry in `/login` to sign in. Model access depends on your xAI account entitlement. See [xAI setup](docs/providers.md#xai-grok). **API keys:** - Anthropic diff --git a/packages/coding-agent/docs/providers.md b/packages/coding-agent/docs/providers.md index 1aaa6a2957..c5d704d8d0 100644 --- a/packages/coding-agent/docs/providers.md +++ b/packages/coding-agent/docs/providers.md @@ -18,6 +18,7 @@ Use `/login` in interactive mode, then select a provider: - ChatGPT Plus/Pro (Codex) - Claude Pro/Max - GitHub Copilot +- xAI Grok (eligible subscriptions) Use `/logout` to clear credentials. Tokens are stored in `~/.prime/agent/auth.json` and auto-refresh when expired. @@ -35,6 +36,18 @@ Anthropic subscription auth is active for Claude Pro/Max accounts. Third-party h - Press Enter for github.com, or enter your GitHub Enterprise Server domain - If you get "model not supported", enable it in VS Code: Copilot Chat → model selector → select model → "Enable" +### xAI Grok + +Use `/login` and select the xAI subscription entry to open browser sign-in. Complete the authorization flow for an eligible Grok subscription. The existing xAI API-key entry still accepts a key, and `XAI_API_KEY` remains supported. + +Both methods use provider ID `xai` and `https://api.x.ai/v1`. Subscription requests use `/responses`; API-key requests keep the existing `/chat/completions` route and model defaults. Changing authentication updates the current session without requiring model reselection. + +Choose any bundled xAI tool-capable language model with `/model` after initial setup. The same catalog is shown for subscription and API-key login. Subscription requests preserve each model’s reasoning and input capabilities; reasoning-effort controls are limited to verified options. Access and usage limits depend on your account entitlement; listing a model does not guarantee a successful request. If a model is unavailable or authorization fails, check your plan or use an API key. + +`grok-code-fast-1` remains a legacy alias for `grok-build-0.1`, not a separate model. Image/video generators and the multi-agent model are not included because they do not support the agent’s custom function tools. + +`/logout` removes saved xAI authentication, but does not unset `XAI_API_KEY`; an environment key can remain active after logout. + ## API Keys ### Environment Variables or Auth File diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 8d67ea6d03..ff97bbf4cf 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -1788,6 +1788,7 @@ export class AgentSession { private async _getRequiredRequestAuth(model: Model): Promise<{ apiKey: string; headers?: Record; + requestModel: Model; }> { const result = await this._modelRegistry.getApiKeyAndHeaders(model); if (!result.ok) { @@ -1797,7 +1798,7 @@ export class AgentSession { throw new Error(result.error); } if (result.apiKey) { - return { apiKey: result.apiKey, headers: result.headers }; + return { apiKey: result.apiKey, headers: result.headers, requestModel: result.requestModel ?? model }; } const isOAuth = this._modelRegistry.isUsingOAuth(model); @@ -8361,9 +8362,9 @@ export class AgentSession { throw new Error(formatNoModelSelectedMessage()); } - const { apiKey, headers } = await this._getRequiredRequestAuth(this.model); + const { apiKey, headers, requestModel } = await this._getRequiredRequestAuth(this.model); const result = await this._performCompaction({ - model: this.model, + model: requestModel, apiKey, headers, customInstructions, @@ -9038,12 +9039,12 @@ export class AgentSession { if (!model) { return { shouldRefine: false, rationale: "No model selected." }; } - const { apiKey, headers } = await this._getRequiredRequestAuth(model); + const { apiKey, headers, requestModel } = await this._getRequiredRequestAuth(model); return reviewAutoRefine( this.agent.state.messages, this._loadMergedHarnessState(), this._loadRefinementHistory(), - model, + requestModel, apiKey, context, headers, @@ -9285,8 +9286,7 @@ export class AgentSession { throw new Error(formatNoModelSelectedMessage()); } - const model = this.model; - const { apiKey, headers } = await this._getRequiredRequestAuth(model); + const { apiKey, headers, requestModel: model } = await this._getRequiredRequestAuth(this.model); const globalHarnessStateDir = getGlobalHarnessStateDir(); const localHarnessStateDir = this._localHarnessStateDir(); const requestedScope = options.global ? "global" : "local"; @@ -9760,7 +9760,7 @@ export class AgentSession { } const result = await this._performCompaction({ - model: this.model, + model: authResult.requestModel ?? this.model, apiKey: authResult.apiKey, headers: authResult.headers, customInstructions, @@ -9957,6 +9957,19 @@ export class AgentSession { : undefined; } + refreshModelMetadata(): void { + if (this.model?.provider === "xai") { + this.agent.state.model = this._modelRegistry.getModelForCurrentAuth(this.model); + this.setThinkingLevel(this.thinkingLevel); + this._clampServiceTierForModel(); + } + this._scopedModels = this._scopedModels.map((scoped) => + scoped.model.provider === "xai" + ? { ...scoped, model: this._modelRegistry.getModelForCurrentAuth(scoped.model) } + : scoped, + ); + } + private _refreshCurrentModelFromRegistry(): void { const currentModel = this.model; if (!currentModel) { @@ -12887,8 +12900,7 @@ export class AgentSession { let summaryDetails: unknown; let summaryUsage: Usage | undefined; if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) { - const model = this.model!; - const { apiKey, headers } = await this._getRequiredRequestAuth(model); + const { apiKey, headers, requestModel: model } = await this._getRequiredRequestAuth(this.model!); const branchSummarySettings = this.settingsManager.getBranchSummarySettings(); const result = await generateBranchSummary(entriesToSummarize, { model, diff --git a/packages/coding-agent/src/core/auth-storage.ts b/packages/coding-agent/src/core/auth-storage.ts index 1faabefd5c..5f2acb9089 100644 --- a/packages/coding-agent/src/core/auth-storage.ts +++ b/packages/coding-agent/src/core/auth-storage.ts @@ -99,6 +99,7 @@ type AuthSourceCandidate = { type AuthApiKeyResult = { apiKey?: string; sourceToken?: AuthSourceToken; + credentialType?: AuthCredential["type"]; }; export interface AuthStorageBackend { @@ -924,7 +925,7 @@ export class AuthStorage { storedCandidate) : storedCandidate, ); - return { apiKey, sourceToken }; + return { apiKey, sourceToken, credentialType: "api_key" }; } } @@ -945,6 +946,7 @@ export class AuthStorage { const refreshedCandidate = this.getStoredAuthCandidate(providerId); return { apiKey: result.apiKey, + credentialType: "oauth", sourceToken: refreshedCandidate ? this.getAuthSourceTokenForCandidate(providerId, refreshedCandidate) : undefined, @@ -960,6 +962,7 @@ export class AuthStorage { const updatedCandidate = this.getStoredAuthCandidate(providerId); return { apiKey: provider.getApiKey(updatedCred), + credentialType: "oauth", sourceToken: updatedCandidate ? this.getAuthSourceTokenForCandidate(providerId, updatedCandidate) : undefined, @@ -972,6 +975,7 @@ export class AuthStorage { } else { return { apiKey: provider.getApiKey(cred), + credentialType: "oauth", sourceToken: this.getAuthSourceTokenForCandidate(providerId, storedCandidate), }; } diff --git a/packages/coding-agent/src/core/model-registry.ts b/packages/coding-agent/src/core/model-registry.ts index f0906daff5..46776bc48b 100644 --- a/packages/coding-agent/src/core/model-registry.ts +++ b/packages/coding-agent/src/core/model-registry.ts @@ -3,7 +3,7 @@ */ import { Buffer } from "node:buffer"; -import { createHash } from "node:crypto"; +import { createHash, createHmac } from "node:crypto"; import { type AnthropicMessagesCompat, type Api, @@ -22,7 +22,7 @@ import { type SimpleStreamOptions, } from "@earendil-works/pi-ai"; import { registerBuiltinMcpOAuthProviders } from "@earendil-works/pi-ai/mcp"; -import { registerOAuthProvider, resetOAuthProviders } from "@earendil-works/pi-ai/oauth"; +import { getXaiSubscriptionModel, registerOAuthProvider, resetOAuthProviders } from "@earendil-works/pi-ai/oauth"; import { existsSync, readFileSync } from "fs"; import { dirname, join } from "path"; import { type Static, type TProperties, Type } from "typebox"; @@ -272,6 +272,7 @@ export type ResolvedRequestAuth = ok: true; apiKey?: string; headers?: Record; + requestModel?: Model; } | { ok: false; @@ -432,7 +433,11 @@ interface PrivatePrimeAuthorizationCache { } function privatePrimeAuthorizationFingerprint(apiKey: string, teamId: string): string { - return createHash("sha256").update(apiKey).update("\0").update(teamId).digest("hex"); + // Use the bearer token as a MAC key, not a password to hash. Keep the scope stable for disk cache reuse. + return createHmac("sha256", apiKey) + .update("prime-agent:private-prime-authorization:v1\0") + .update(teamId) + .digest("hex"); } function isOfflineModeEnabled(): boolean { @@ -446,6 +451,7 @@ function isOfflineModeEnabled(): boolean { */ export class ModelRegistry { private models: Model[] = []; + private xaiModelSources = new WeakMap, Model>(); private providerRequestConfigs: Map = new Map(); private staleProviderRequestAuthSources: Map = new Map(); private lastProviderAuthSourceTokens: Map = new Map(); @@ -798,7 +804,28 @@ export class ModelRegistry { * If models.json had errors, returns only built-in models. */ getAll(): Model[] { - return this.models; + return this.models.map((model) => this.getModelForCurrentAuth(model)); + } + + getModelForCurrentAuth(model: Model): Model { + if (model.provider !== "xai") return model; + const source = this.xaiModelSources.get(model) ?? model; + return this.isUsingXaiSubscription(source) ? (this.getXaiSubscriptionModel(source) ?? source) : source; + } + + private getXaiSubscriptionModel(model: Model): Model | undefined { + const source = this.xaiModelSources.get(model) ?? model; + const subscription = getXaiSubscriptionModel(source); + if (subscription) this.xaiModelSources.set(subscription, source); + return subscription; + } + + private isUsingXaiSubscription(model: Model): boolean { + return ( + model.provider === "xai" && + this.authStorage.get("xai")?.type === "oauth" && + this.getProviderAuthStatus("xai").source === "stored" + ); } /** @@ -806,7 +833,7 @@ export class ModelRegistry { * This is a fast check that doesn't refresh OAuth tokens. */ getAvailable(): Model[] { - return this.models.filter((model) => { + return this.getAll().filter((model) => { if (isPrivatePrimeInferenceModel(model) && !this.isAuthorizedPrivatePrimeInferenceModel(model)) { return false; } @@ -1067,7 +1094,7 @@ export class ModelRegistry { availableModels.filter(isPrivatePrimeInferenceModel).map((model) => `${model.provider}/${model.id}`), ); return { - models: this.models.filter( + models: this.getAll().filter( (model) => !isPrivatePrimeInferenceModel(model) || availablePrivateModels.has(`${model.provider}/${model.id}`), ), @@ -1111,7 +1138,9 @@ export class ModelRegistry { if (!auth.ok || !auth.apiKey) { return availableModels.filter((model) => model.provider !== "openai-codex"); } - const authFingerprint = createHash("sha256").update(auth.apiKey).digest("hex"); + const authFingerprint = createHmac("sha256", auth.apiKey) + .update("prime-agent:openai-codex-models:v1") + .digest("hex"); const cached = this.openAICodexModelsCache; if (cached?.authFingerprint === authFingerprint && Date.now() - cached.refreshedAt < 300_000) { return availableModels.filter((model) => model.provider !== "openai-codex" || cached.modelIds.has(model.id)); @@ -1151,7 +1180,8 @@ export class ModelRegistry { * Find a model by provider and ID. */ find(provider: string, modelId: string): Model | undefined { - return this.models.find((m) => m.provider === provider && m.id === modelId); + const model = this.models.find((m) => m.provider === provider && m.id === modelId); + return model ? this.getModelForCurrentAuth(model) : undefined; } /** @@ -1434,7 +1464,7 @@ export class ModelRegistry { /** * Get API key and request headers for a model. */ - async getApiKeyAndHeaders(model: Model): Promise { + async getApiKeyAndHeaders(model: Model, requestHeaders?: Record): Promise { try { const providerConfig = this.providerRequestConfigs.get(model.provider); const authStorageAuth = await this.authStorage.getApiKeyWithSourceToken(model.provider, { @@ -1459,6 +1489,25 @@ export class ModelRegistry { } this.setLastProviderAuthSourceToken(model.provider, apiKey === undefined ? undefined : authSourceToken); + let requestModel: Model | undefined; + if (model.provider === "xai") { + if (!apiKey) { + return { + ok: false, + error: "No usable xAI credential. Use /login and select an xAI entry to sign in again or configure an API key.", + }; + } + const configuredModel = this.xaiModelSources.get(model) ?? model; + const subscription = authStorageAuth.credentialType === "oauth" && authSourceToken?.source === "stored"; + requestModel = subscription ? this.getXaiSubscriptionModel(configuredModel) : configuredModel; + if (!requestModel) { + return { + ok: false, + error: `Cannot configure "${model.id}" for Grok subscription requests. Select an xAI model or use /login and select the xAI API-key entry.`, + }; + } + } + const providerHeaders = resolveHeadersOrThrow(providerConfig?.headers, `provider "${model.provider}"`); const authStorageHeaders = this.authStorage.getProviderHeaders(model.provider); const modelHeaders = resolveHeadersOrThrow( @@ -1478,10 +1527,26 @@ export class ModelRegistry { headers = { ...headers, Authorization: `Bearer ${apiKey}` }; } + if (requestHeaders) headers = { ...headers, ...requestHeaders }; + if ( + model.provider === "xai" && + authStorageAuth.credentialType === "oauth" && + authSourceToken?.source === "stored" + ) { + for (const [name, value] of Object.entries(headers ?? {})) { + if (name.toLowerCase() === "authorization" && value !== `Bearer ${apiKey}`) { + return { + ok: false, + error: "Grok subscription cannot use a custom Authorization header. Remove the header or use /login and select the xAI API-key entry.", + }; + } + } + } return { ok: true, apiKey, headers: headers && Object.keys(headers).length > 0 ? headers : undefined, + ...(requestModel ? { requestModel } : {}), }; } catch (error) { return { @@ -1569,6 +1634,7 @@ export class ModelRegistry { * Check if a model is using OAuth credentials (subscription). */ isUsingOAuth(model: Model): boolean { + if (model.provider === "xai") return this.isUsingXaiSubscription(model); const cred = this.authStorage.get(model.provider); return cred?.type === "oauth"; } diff --git a/packages/coding-agent/src/core/provider-display-names.ts b/packages/coding-agent/src/core/provider-display-names.ts index 9cae04a8d7..5ea16a5fd0 100644 --- a/packages/coding-agent/src/core/provider-display-names.ts +++ b/packages/coding-agent/src/core/provider-display-names.ts @@ -24,7 +24,7 @@ export const BUILT_IN_PROVIDER_DISPLAY_NAMES: Record = { "prime-agent-traces": "Prime Agent Traces", "prime-inference": "Prime Inference", "vercel-ai-gateway": "Vercel AI Gateway", - xai: "xAI", + xai: "xAI (Grok)", zai: "ZAI", xiaomi: "Xiaomi MiMo", "xiaomi-token-plan-cn": "Xiaomi MiMo Token Plan (China)", diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 9989e27585..27c7bf276a 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -286,16 +286,17 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} }, convertToLlm: convertToLlmWithBlockImages, streamFn: async (model, context, options) => { - const auth = await modelRegistry.getApiKeyAndHeaders(model); + const auth = await modelRegistry.getApiKeyAndHeaders(model, options?.headers); if (!auth.ok) { throw new Error(auth.error); } const providerRetrySettings = settingsManager.getProviderRetrySettings(); - return streamSimple(model, context, { + const requestModel = auth.requestModel ?? model; + return streamSimple(requestModel, context, { ...options, apiKey: auth.apiKey, timeoutMs: options?.timeoutMs ?? providerRetrySettings.timeoutMs, - headers: auth.headers || options?.headers ? { ...auth.headers, ...options?.headers } : undefined, + headers: auth.headers, }); }, onPayload: providerHooks.onPayload, diff --git a/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts b/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts index e3e6dbe243..3221666d7e 100644 --- a/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts +++ b/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts @@ -5417,8 +5417,6 @@ function invalidatesCachedSnapshot(commandType: DaemonCommandBody["type"]): bool case "get_session_stats": case "get_commands": case "get_resource_snapshot": - case "get_model_catalog": - case "get_available_models": case "get_queue": case "cron_list": case "heartbeats_list": diff --git a/packages/coding-agent/src/modes/agent-connection/in-process-agent-connection.ts b/packages/coding-agent/src/modes/agent-connection/in-process-agent-connection.ts index 5e39f2db44..47fdb50d7f 100644 --- a/packages/coding-agent/src/modes/agent-connection/in-process-agent-connection.ts +++ b/packages/coding-agent/src/modes/agent-connection/in-process-agent-connection.ts @@ -185,11 +185,17 @@ export class InProcessAgentConnection implements AgentConnection { } async getAvailableModels(): Promise { - return this.session.modelRegistry.refreshAvailableModels(); + const session = this.session; + const models = await session.modelRegistry.refreshAvailableModels(); + session.refreshModelMetadata(); + return models; } async getModelCatalog(): Promise { - return this.session.modelRegistry.refreshModelCatalog(); + const session = this.session; + const catalog = await session.modelRegistry.refreshModelCatalog(); + session.refreshModelMetadata(); + return catalog; } async getSessionStats(): Promise { diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index dcf4ea78c3..23eb26661a 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -5266,19 +5266,17 @@ export class AgentDaemon { } case "get_available_models": { - const state = this.getSessionState(command.activeSessionId); - return success(command.id, "get_available_models", { - models: await state.runtime.session.modelRegistry.refreshAvailableModels(), - }); + const { session } = this.getSessionState(command.activeSessionId).runtime; + const models = await session.modelRegistry.refreshAvailableModels(); + session.refreshModelMetadata(); + return success(command.id, "get_available_models", { models }); } case "get_model_catalog": { - const state = this.getSessionState(command.activeSessionId); - return success( - command.id, - "get_model_catalog", - await state.runtime.session.modelRegistry.refreshModelCatalog(), - ); + const { session } = this.getSessionState(command.activeSessionId).runtime; + const catalog = await session.modelRegistry.refreshModelCatalog(); + session.refreshModelMetadata(); + return success(command.id, "get_model_catalog", catalog); } case "get_queue": { diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index e1927c6f6d..c9b69ac794 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -7934,8 +7934,26 @@ export class InteractiveMode { } private async refreshConnectionModelsAfterAuthChange(): Promise { + const connection = this.agentConnection; + const sessionId = this.connectionState?.sessionId; this.invalidateConnectionModels(); await this.getConnectionAvailableModels(); + const state = await connection.getState(); + if ( + this.agentConnection !== connection || + this.connectionState?.sessionId !== sessionId || + (sessionId !== undefined && state.sessionId !== sessionId) + ) { + return; + } + this.patchConnectionState({ + model: state.model, + scopedModels: state.scopedModels, + serviceTier: state.serviceTier, + availableThinkingLevels: state.availableThinkingLevels, + }); + this.subagentSummaryLine.invalidate(); + this.setupAutocompleteProvider(); } private async getModelCandidates(): Promise { diff --git a/packages/coding-agent/test/agent-connection-daemon.test.ts b/packages/coding-agent/test/agent-connection-daemon.test.ts index fa85e9c54a..d85cf43eb4 100644 --- a/packages/coding-agent/test/agent-connection-daemon.test.ts +++ b/packages/coding-agent/test/agent-connection-daemon.test.ts @@ -8943,6 +8943,44 @@ describe("DaemonAgentConnection", () => { await connection.dispose(); }); + it.each(["catalog", "legacyCatalog", "available"] as const)( + "fetches fresh model state after a %s refresh invalidates the attach snapshot", + async (refresh) => { + const fakeClient = new FakeDaemonClient(); + if (refresh === "catalog") fakeClient.serverCapabilities.add("model_catalog"); + const model = getModel("xai", "grok-4.5"); + const initialState = { ...createConnectionState("active-1", "session-current"), model }; + fakeClient.attachResultFactory = (command) => + createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12, { + state: initialState, + }); + const connection = await DaemonAgentConnection.attach(asDaemonClient(fakeClient), "active-1"); + try { + const subscriptionModel = { ...model, api: "openai-responses" as const }; + const freshState = { + ...initialState, + model: subscriptionModel, + scopedModels: [{ model: subscriptionModel }], + availableThinkingLevels: ["low", "medium", "high"] as AgentConnectionState["availableThinkingLevels"], + }; + fakeClient.connectionStateFactory = () => freshState; + expect(await connection.getState()).toBe(initialState); + expect(fakeClient.requests.filter((request) => request.type === "get_connection_state")).toHaveLength(0); + + if (refresh === "available") await connection.getAvailableModels(); + else await connection.getModelCatalog(); + + expect(await connection.getState()).toBe(freshState); + expect(fakeClient.requests.slice(-2).map((request) => request.type)).toEqual([ + refresh === "catalog" ? "get_model_catalog" : "get_available_models", + "get_connection_state", + ]); + } finally { + await connection.dispose(); + } + }, + ); + it("loads the full model catalog through the daemon protocol", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("model_catalog"); diff --git a/packages/coding-agent/test/agent-connection-in-process.test.ts b/packages/coding-agent/test/agent-connection-in-process.test.ts index 5ac8be1ade..60f09ca6df 100644 --- a/packages/coding-agent/test/agent-connection-in-process.test.ts +++ b/packages/coding-agent/test/agent-connection-in-process.test.ts @@ -105,7 +105,9 @@ function createFakeSession(id: string, messages: AgentMessage[]): FakeSessionCon goalState: emptyGoalState(), modelRegistry: { refreshModelCatalog: async () => ({ models: model ? [model] : [], configuredProviders: ["openai"] }), + refreshAvailableModels: async () => (model ? [model] : []), }, + refreshModelMetadata: vi.fn(), scopedModels: [], getActiveToolNames: () => ["ipython"], getContextUsage: () => undefined, @@ -291,6 +293,9 @@ describe("InProcessAgentConnection", () => { expect(catalog.configuredProviders).toEqual(["openai"]); expect(catalog.models).toHaveLength(1); expect(catalog.models[0]).toMatchObject({ provider: "openai", id: "gpt-5.1" }); + expect(session.session.refreshModelMetadata).toHaveBeenCalledOnce(); + await connection.getAvailableModels(); + expect(session.session.refreshModelMetadata).toHaveBeenCalledTimes(2); }); it("exposes serializable tool metadata without local execution or renderer callbacks", async () => { diff --git a/packages/coding-agent/test/auth-flows.test.ts b/packages/coding-agent/test/auth-flows.test.ts index 43c3f7328d..6c6fff6746 100644 --- a/packages/coding-agent/test/auth-flows.test.ts +++ b/packages/coding-agent/test/auth-flows.test.ts @@ -5,7 +5,7 @@ import type { Component, OverlayHandle, TUI } from "@earendil-works/pi-tui"; import stripAnsi from "strip-ansi"; import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.js"; -import type { ModelRegistry } from "../src/core/model-registry.js"; +import { ModelRegistry } from "../src/core/model-registry.js"; import { PRIME_INFERENCE_PROVIDER_ID } from "../src/core/prime-inference-auth.js"; import { ProviderAuthFlows, type ProviderAuthFlowsHost } from "../src/modes/interactive/auth-flows.js"; import { initTheme } from "../src/modes/interactive/theme/theme.js"; @@ -207,6 +207,19 @@ describe("ProviderAuthFlows", () => { await expect(logoutResult).resolves.toBeNull(); }); + it("includes xAI subscription and API-key entries from the provider registry", () => { + const { host } = createHost(AuthStorage.inMemory()); + const flows = new ProviderAuthFlows({ + ...host, + modelRegistry: ModelRegistry.inMemory(host.modelRegistry.authStorage), + }); + + expect(flows.getLoginProviderOptions().filter((provider) => provider.id === "xai")).toEqual([ + { id: "xai", name: "xAI (Grok)", authType: "oauth" }, + { id: "xai", name: "xAI (Grok)", authType: "api_key" }, + ]); + }); + it("opens login on the requested MCP Connections category", async () => { const authStorage = AuthStorage.create(authJsonPath, { usePrimeCliConfig: false }); const { host, overlays } = createHost(authStorage); diff --git a/packages/coding-agent/test/configuration-menu.test.ts b/packages/coding-agent/test/configuration-menu.test.ts index e63c750d31..b4e617d006 100644 --- a/packages/coding-agent/test/configuration-menu.test.ts +++ b/packages/coding-agent/test/configuration-menu.test.ts @@ -1,9 +1,12 @@ +import { xaiOAuthProvider } from "@earendil-works/pi-ai/oauth"; import { setKeybindings, type TUI, visibleWidth } from "@earendil-works/pi-tui"; import stripAnsi from "strip-ansi"; import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import { KeybindingsManager } from "../src/core/keybindings.js"; +import { BUILT_IN_PROVIDER_DISPLAY_NAMES } from "../src/core/provider-display-names.js"; import { ConfigurationMenuComponent, + type ConfigurationMenuOptions, type ConfigurationMenuTab, } from "../src/modes/interactive/components/configuration-menu.js"; import { initTheme } from "../src/modes/interactive/theme/theme.js"; @@ -21,6 +24,7 @@ describe("ConfigurationMenuComponent", () => { async function createMenu( options: { initialTab?: ConfigurationMenuTab; + providerOptions?: ConfigurationMenuOptions["providerOptions"]; getRows?: () => number; requestRender?: () => void; onSelectProvider?: () => void; @@ -36,7 +40,7 @@ describe("ConfigurationMenuComponent", () => { initialTab: options.initialTab ?? "providers", tui: createFakeTui(), authStorage: harness.session.modelRegistry.authStorage, - providerOptions: [ + providerOptions: options.providerOptions ?? [ { id: "anthropic", name: "Anthropic", authType: "oauth" }, { id: "serper", @@ -74,6 +78,26 @@ describe("ConfigurationMenuComponent", () => { } }); + it.each(["grok", "xai"])("finds both xAI auth entries when searching %s", async (query) => { + const selectProvider = vi.fn(); + const menu = await createMenu({ + providerOptions: [ + { id: xaiOAuthProvider.id, name: xaiOAuthProvider.name, authType: "oauth" }, + { id: "xai", name: BUILT_IN_PROVIDER_DISPLAY_NAMES.xai, authType: "api_key" }, + ], + onSelectProvider: selectProvider, + }); + menu.handleInput(query); + expect(stripAnsi(menu.render(120).join("\n")).match(/xAI \(Grok\)/g)).toHaveLength(2); + menu.handleInput("\r"); + menu.handleInput("\x1b[B"); + menu.handleInput("\r"); + expect(selectProvider.mock.calls.map(([provider]) => [provider.id, provider.authType])).toEqual([ + ["xai", "oauth"], + ["xai", "api_key"], + ]); + }); + it("uses one clearly delineated three-tab menu and keeps each tab body mounted", async () => { const requestRender = vi.fn(); const selectProvider = vi.fn(); diff --git a/packages/coding-agent/test/interactive-mode-status.test.ts b/packages/coding-agent/test/interactive-mode-status.test.ts index a013a0df67..2b9b0ec6e6 100644 --- a/packages/coding-agent/test/interactive-mode-status.test.ts +++ b/packages/coding-agent/test/interactive-mode-status.test.ts @@ -3104,6 +3104,7 @@ function createFakeConnectionSession(commandName: string): AgentSessionRuntime[" getAgentsFiles: () => ({ agentsFiles: [] }), }, modelRegistry: { refreshModelCatalog: async () => ({ models: [], configuredProviders: [] }) }, + refreshModelMetadata: vi.fn(), sessionManager: { getCwd: () => "/tmp/project", getSessionDir: () => "/tmp/sessions", @@ -3221,10 +3222,12 @@ describe("InteractiveMode session switch command catalog", () => { if (operation === "switchSession") await connection.switchSession("/target/session.jsonl"); else if (operation === "newSession") await connection.newSession(); else await connection.fork("entry-1"); + expect(fakeThis.showError).not.toHaveBeenCalled(); await interactiveHarness.renderCurrentSessionState(); expect(calls).toEqual(["reset", "catalog", "render", "reset", "render"]); expect(getCommands).toHaveBeenCalledTimes(2); // initial catalog + one replacement refresh + expect(targetSession.refreshModelMetadata).toHaveBeenCalledOnce(); }, ); diff --git a/packages/coding-agent/test/model-registry.test.ts b/packages/coding-agent/test/model-registry.test.ts index 87721fa085..b8b4f84f7a 100644 --- a/packages/coding-agent/test/model-registry.test.ts +++ b/packages/coding-agent/test/model-registry.test.ts @@ -22,6 +22,7 @@ describe("ModelRegistry", () => { afterEach(() => { vi.unstubAllGlobals(); + vi.unstubAllEnvs(); if (tempDir && existsSync(tempDir)) { rmSync(tempDir, { recursive: true }); } @@ -650,7 +651,8 @@ describe("ModelRegistry", () => { expect(getModelsForProvider(registry, "openrouter")).toHaveLength(getModels("openrouter").length); }); - test("restores cached authorized deployment metadata without waiting for the network", async () => { + test("restores cached private metadata only for matching credentials and team", async () => { + vi.stubEnv("PI_OFFLINE", "0"); const privateRoute = { id: "vendor/model:deployment", display_name: "Private Deployment", @@ -662,11 +664,12 @@ describe("ModelRegistry", () => { supports_reasoning: false, }, }; - authStorage.set("prime-inference", { - type: "api_key", + const credential = { + type: "api_key" as const, key: "prime-key", primeTeam: { teamId: "research-team", name: "Research" }, - }); + }; + authStorage.set("prime-inference", credential); vi.stubGlobal( "fetch", vi.fn( @@ -687,10 +690,33 @@ describe("ModelRegistry", () => { throw new Error("offline"); }), ); - const restoredRegistry = ModelRegistry.create(authStorage, modelsJsonPath); + vi.stubEnv("PI_OFFLINE", "1"); + const restoredRegistry = ModelRegistry.create(AuthStorage.create(join(tempDir, "auth.json")), modelsJsonPath); expect( (await restoredRegistry.refreshAvailableModels()).find((model) => model.id === privateRoute.id), ).toMatchObject({ name: "Private Deployment", contextWindow: 200_000 }); + + for (const changed of [ + { ...credential, key: "different-prime-key" }, + { ...credential, primeTeam: { teamId: "other-team", name: "Other" } }, + ]) { + authStorage.set("prime-inference", changed); + const changedRegistry = ModelRegistry.create(authStorage, modelsJsonPath); + expect((await changedRegistry.refreshAvailableModels()).some((model) => model.id === privateRoute.id)).toBe( + false, + ); + } + + authStorage.set("prime-inference", credential); + const cachePath = join(tempDir, "prime-inference-private-models.json"); + const cache = JSON.parse(readFileSync(cachePath, "utf8")); + // Pre-HMAC SHA256("prime-key\0research-team") cache entries must miss safely on upgrade. + cache.fingerprint = "9ffd3740e055c8cc8923a1d2653c6d02a4a9c95e6ab151bc179aeaa94dd046b4"; + writeFileSync(cachePath, JSON.stringify(cache)); + const legacyRegistry = ModelRegistry.create(authStorage, modelsJsonPath); + expect((await legacyRegistry.refreshAvailableModels()).some((model) => model.id === privateRoute.id)).toBe( + false, + ); }); }); diff --git a/packages/coding-agent/test/suite/regressions/4575-model-auth-selection.test.ts b/packages/coding-agent/test/suite/regressions/4575-model-auth-selection.test.ts index 93b41ed17e..30eeee5967 100644 --- a/packages/coding-agent/test/suite/regressions/4575-model-auth-selection.test.ts +++ b/packages/coding-agent/test/suite/regressions/4575-model-auth-selection.test.ts @@ -2,7 +2,11 @@ import { type AutocompleteProvider, setKeybindings, type TUI } from "@earendil-w import stripAnsi from "strip-ansi"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { KeybindingsManager } from "../../../src/core/keybindings.js"; -import type { AgentConnectionModel, AgentConnectionModelCatalog } from "../../../src/modes/agent-connection/types.js"; +import type { + AgentConnectionModel, + AgentConnectionModelCatalog, + AgentConnectionState, +} from "../../../src/modes/agent-connection/types.js"; import { ModelSelectorComponent } from "../../../src/modes/interactive/components/model-selector.js"; import { InteractiveMode } from "../../../src/modes/interactive/interactive-mode.js"; import { initTheme } from "../../../src/modes/interactive/theme/theme.js"; @@ -10,7 +14,14 @@ import { getModelArgumentCompletions } from "../../../src/modes/model-autocomple import { createHarness, type Harness } from "../harness.js"; interface ConnectionAuthRefreshHarness { - agentConnection: { getModelCatalog(): Promise }; + agentConnection: { + getModelCatalog(): Promise; + getState?(): Promise>; + }; + connectionState?: Partial; + patchConnectionState?: (patch: Partial) => void; + subagentSummaryLine?: { invalidate(): void }; + setupAutocompleteProvider?: () => void; connectionModelCatalog: AgentConnectionModel[]; connectionConfiguredProviders: Set; connectionModelsFetchedAt: number; @@ -123,7 +134,19 @@ describe("ENG-4575 model authentication", () => { const model = { ...harness.getModel("base")!, provider: "openai" } as AgentConnectionModel; const getModelCatalog = vi.fn(async () => ({ models: [model], configuredProviders: [] })); const fakeThis = Object.create(InteractiveMode.prototype) as ConnectionAuthRefreshHarness; - fakeThis.agentConnection = { getModelCatalog }; + const state: Partial = { + sessionId: harness.session.sessionId, + model, + scopedModels: [{ model }], + serviceTier: "default", + availableThinkingLevels: ["low", "medium", "high"], + }; + const getState = vi.fn(async () => state); + fakeThis.agentConnection = { getModelCatalog, getState }; + fakeThis.connectionState = { sessionId: state.sessionId }; + fakeThis.patchConnectionState = vi.fn(); + fakeThis.subagentSummaryLine = { invalidate: vi.fn() }; + fakeThis.setupAutocompleteProvider = vi.fn(); fakeThis.connectionModelCatalog = [model]; fakeThis.connectionConfiguredProviders = new Set([model.provider]); fakeThis.connectionModelsFetchedAt = Date.now(); @@ -133,6 +156,14 @@ describe("ENG-4575 model authentication", () => { await fakeThis.refreshConnectionModelsAfterAuthChange(); expect(getModelCatalog).toHaveBeenCalledOnce(); + expect(getState).toHaveBeenCalledOnce(); + expect(getModelCatalog.mock.invocationCallOrder[0]).toBeLessThan(getState.mock.invocationCallOrder[0]); + expect(fakeThis.patchConnectionState).toHaveBeenCalledWith({ + model, + scopedModels: state.scopedModels, + serviceTier: state.serviceTier, + availableThinkingLevels: state.availableThinkingLevels, + }); expect(fakeThis.connectionConfiguredProviders).toEqual(new Set()); expect(fakeThis.getAvailableConnectionModels()).toEqual([]); expect(fakeThis.connectionModelCatalog).toEqual([model]); diff --git a/packages/coding-agent/test/suite/regressions/68-xai-subscription-auth.test.ts b/packages/coding-agent/test/suite/regressions/68-xai-subscription-auth.test.ts new file mode 100644 index 0000000000..36aca4ab5e --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/68-xai-subscription-auth.test.ts @@ -0,0 +1,155 @@ +import { join } from "node:path"; +import { + type Api, + type ApiStreamSimpleFunction, + fauxAssistantMessage, + getApiProvider, + getModel, + type Model, + registerApiProvider, +} from "@earendil-works/pi-ai"; +import { xaiOAuthProvider } from "@earendil-works/pi-ai/oauth"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import type { AgentSession } from "../../../src/core/agent-session.js"; +import { AuthStorage } from "../../../src/core/auth-storage.js"; +import { ModelRegistry } from "../../../src/core/model-registry.js"; +import { createAgentSession } from "../../../src/core/sdk.js"; +import { SessionManager } from "../../../src/core/session-manager.js"; +import { InProcessAgentConnection } from "../../../src/modes/agent-connection/in-process-agent-connection.js"; +import { createTestResourceLoader } from "../../utilities.js"; +import { createHarness, type Harness } from "../harness.js"; + +describe("ENG-6059 xAI subscription dispatch", () => { + const harnesses: Harness[] = []; + const sessions: AgentSession[] = []; + beforeEach(() => { + vi.stubEnv("XAI_API_KEY", "environment-key"); + vi.stubEnv("PI_OFFLINE", "1"); + }); + afterEach(() => { + for (const session of sessions.splice(0)) session.dispose(); + for (const harness of harnesses.splice(0)) harness.cleanup(); + vi.restoreAllMocks(); + vi.unstubAllEnvs(); + }); + + async function setup() { + const harness = await createHarness({ + api: "openai-completions", + provider: "xai", + models: [{ id: "grok-4.5" }], + settings: { compaction: { enabled: false, keepRecentTokens: 1 } }, + }); + harnesses.push(harness); + const faux = getApiProvider("openai-completions")!; + const requests: Array<{ model: Model; key?: string }> = []; + const stream: ApiStreamSimpleFunction = (model, context, options) => { + requests.push({ model, key: options?.apiKey }); + return faux.streamSimple({ ...model, api: "openai-completions" }, context, options); + }; + const installTransports = () => { + for (const api of ["openai-completions", "openai-responses"] as const) { + registerApiProvider({ api, stream, streamSimple: stream }); + } + }; + const path = join(harness.tempDir, "auth.json"); + const registry = ModelRegistry.inMemory(AuthStorage.create(path)); + const model = { + ...getModel("xai", "grok-4.5"), + baseUrl: "https://caller.invalid/v2", + headers: { "X-Caller": "kept" }, + }; + const { session } = await createAgentSession({ + cwd: harness.tempDir, + agentDir: harness.tempDir, + modelRegistry: registry, + authStorage: registry.authStorage, + settingsManager: harness.settingsManager, + sessionManager: SessionManager.inMemory(), + resourceLoader: createTestResourceLoader(), + model, + scopedModels: [{ model, thinkingLevel: "high" }], + tools: [], + noTools: "all", + }); + sessions.push(session); + const writer = AuthStorage.create(path); + const login = (expires = Date.now() + 60_000) => + writer.set("xai", { + type: "oauth", + access: "subscription-key", + refresh: "refresh-key", + expires, + }); + const connection = Object.create(InProcessAgentConnection.prototype) as InProcessAgentConnection; + Object.defineProperty(connection, "runtimeHost", { value: { session } }); + const refresh = async () => { + await connection.getModelCatalog(); + installTransports(); + }; + installTransports(); + return { harness, session, registry, writer, model, requests, login, refresh, installTransports }; + } + + test("switches an existing custom API-key session to subscription and back, including compaction", async () => { + const { harness, session, registry, writer, model, requests, login, refresh, installTransports } = await setup(); + harness.setResponses([ + fauxAssistantMessage("one"), + fauxAssistantMessage("two"), + fauxAssistantMessage("summary"), + fauxAssistantMessage("turn summary"), + fauxAssistantMessage("restored"), + ]); + await session.prompt("API key turn"); + expect(requests.at(-1)?.model).toBe(model); + login(); + registry.refresh(); + installTransports(); + // The active object is still the API-key model: dispatch must resolve the new credential and route together. + expect(session.model).toBe(model); + await session.prompt("subscription turn"); + expect(requests.at(-1)).toMatchObject({ key: "subscription-key", model: { api: "openai-responses" } }); + login(Date.now() - 1000); + await refresh(); + expect(session.model?.api).toBe("openai-responses"); + expect(session.scopedModels[0]?.model.api).toBe("openai-responses"); + vi.spyOn(xaiOAuthProvider, "refreshToken").mockResolvedValue({ + access: "rotated", + refresh: "rotated-refresh", + expires: Date.now() + 60_000, + }); + const beforeCompaction = requests.length; + await session.compact(); + expect(requests.length).toBeGreaterThan(beforeCompaction); + for (const request of requests.slice(beforeCompaction)) { + expect(request).toMatchObject({ key: "rotated", model: { api: "openai-responses" } }); + } + writer.set("xai", { type: "api_key", key: "saved-key" }); + await refresh(); + expect(session.model).toBe(model); + expect(session.scopedModels[0]?.model).toBe(model); + await session.prompt("custom API key turn again"); + expect(requests.at(-1)?.model).toBe(model); + expect(requests.at(-1)?.key).toBe("saved-key"); + }); + + test("rejects final per-request conflicting Authorization before SDK dispatch", async () => { + const { session, requests, login, refresh } = await setup(); + login(); + await refresh(); + await expect( + session.agent.streamFn(session.model!, { messages: [] }, { headers: { authorization: "conflicting-secret" } }), + ).rejects.toThrow("Remove the header"); + expect(requests).toHaveLength(0); + }); + + test("does not silently use XAI_API_KEY after subscription refresh fails", async () => { + const { session, requests, login, refresh } = await setup(); + login(Date.now() - 1000); + await refresh(); + vi.spyOn(xaiOAuthProvider, "refreshToken").mockRejectedValue(new Error("revoked")); + await session.prompt("do not silently charge API key"); + expect(requests).toHaveLength(0); + expect(session.state.errorMessage).toContain("/login and select"); + }); +}); diff --git a/packages/coding-agent/test/suite/regressions/702-codex-client-version.test.ts b/packages/coding-agent/test/suite/regressions/702-codex-client-version.test.ts index e143f56b6c..b49c65a1ab 100644 --- a/packages/coding-agent/test/suite/regressions/702-codex-client-version.test.ts +++ b/packages/coding-agent/test/suite/regressions/702-codex-client-version.test.ts @@ -43,7 +43,8 @@ describe("issue #702 codex model discovery client version", () => { }), ); - const registry = ModelRegistry.create(AuthStorage.create(authPath), join(tempDir, "models.json")); + const authStorage = AuthStorage.create(authPath); + const registry = ModelRegistry.create(authStorage, join(tempDir, "models.json")); const codexModels = registry.getAvailable().filter((model) => model.provider === "openai-codex"); expect(codexModels.length).toBeGreaterThan(0); @@ -70,5 +71,17 @@ describe("issue #702 codex model discovery client version", () => { expect((major ?? 0) > 0 || (minor ?? 0) >= 153).toBe(true); expect(executable.some((model) => model.provider === "openai-codex")).toBe(true); + await registry.getExecutableModels(); + expect(requestedUrls.filter((url) => url.includes("/codex/models"))).toHaveLength(1); + + authStorage.set("openai-codex", { + type: "oauth", + access: `${codexAccessToken("account-123")}-rotated`, + refresh: "rotated-refresh-token", + expires: Date.now() + 60 * 60 * 1000, + accountId: "account-123", + }); + await registry.getExecutableModels(); + expect(requestedUrls.filter((url) => url.includes("/codex/models"))).toHaveLength(2); }); }); diff --git a/packages/coding-agent/test/xai-auth.test.ts b/packages/coding-agent/test/xai-auth.test.ts new file mode 100644 index 0000000000..b2114f3c6d --- /dev/null +++ b/packages/coding-agent/test/xai-auth.test.ts @@ -0,0 +1,124 @@ +import { getModel, getModels } from "@earendil-works/pi-ai"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { AuthStorage, type OAuthCredential } from "../src/core/auth-storage.js"; +import { ModelRegistry } from "../src/core/model-registry.js"; + +const oauth = (): OAuthCredential => ({ + type: "oauth", + access: "subscription-access", + refresh: "subscription-refresh", + expires: Date.now() + 60_000, +}); + +describe("xAI credential source and request model", () => { + let storage: AuthStorage; + let registry: ModelRegistry; + beforeEach(() => { + vi.stubEnv("XAI_API_KEY", "environment-key"); + storage = AuthStorage.inMemory(); + registry = ModelRegistry.inMemory(storage); + }); + afterEach(() => { + vi.restoreAllMocks(); + vi.unstubAllEnvs(); + }); + + test("routes the selected OAuth, runtime or stale-fallback credential without changing API-key models", async () => { + registry.registerProvider("xai", { baseUrl: "https://example.invalid/custom", apiKey: "config-key" }); + const original = registry.find("xai", "grok-4.5")!; + const snapshot = structuredClone(original); + storage.set("xai", oauth()); + const cachedSubscription = registry.find("xai", "grok-4.5")!; + const expectRoute = async (apiKey: string, api: string) => { + expect(await registry.getApiKeyAndHeaders(cachedSubscription)).toMatchObject({ + ok: true, + apiKey, + requestModel: { api, baseUrl: api === "openai-responses" ? "https://api.x.ai/v1" : original.baseUrl }, + }); + }; + await expectRoute("subscription-access", "openai-responses"); + storage.setRuntimeApiKey("xai", "runtime-key"); + await expectRoute("runtime-key", "openai-completions"); + storage.removeRuntimeApiKey("xai"); + storage.markAuthStale("xai"); + await expectRoute("environment-key", "openai-completions"); + vi.stubEnv("XAI_API_KEY", ""); + await expectRoute("config-key", "openai-completions"); + expect(original).toEqual(snapshot); + expect(getModel("xai", "grok-4.5").api).toBe("openai-completions"); + }); + + test("validates final Authorization case-insensitively only for subscription credentials", async () => { + storage.set("xai", oauth()); + const model = { ...getModel("xai", "grok-4.5"), headers: { Authorization: "custom-secret" } }; + const rejected = await registry.getApiKeyAndHeaders(model); + expect(rejected).toMatchObject({ ok: false, error: expect.stringContaining("Remove the header") }); + expect(JSON.stringify(rejected)).not.toContain("custom-secret"); + expect(await registry.getApiKeyAndHeaders(model, { Authorization: "Bearer subscription-access" })).toMatchObject({ + ok: true, + }); + expect(await registry.getApiKeyAndHeaders(model, { authorization: "Bearer subscription-access" })).toMatchObject({ + ok: false, + }); + storage.setRuntimeApiKey("xai", "runtime-key"); + expect(await registry.getApiKeyAndHeaders(model)).toMatchObject({ + ok: true, + apiKey: "runtime-key", + headers: { Authorization: "custom-secret" }, + }); + }); + + test("keeps all configured xAI tool models selectable for subscription and API-key auth", async () => { + const models = getModels("xai"); + const customModel = { + ...getModel("xai", "grok-4.6"), + id: "custom-grok", + baseUrl: "https://example.invalid/custom", + input: ["text"] as ["text"], + contextWindow: 1234, + maxTokens: 512, + thinkingLevelMap: { off: null, minimal: null, low: "low", medium: "medium", high: "high" }, + }; + const ids = models.map((model) => model.id); + for (const credential of [oauth(), { type: "api_key" as const, key: "api-key" }]) { + storage.set("xai", credential); + expect( + registry + .getAvailable() + .filter((model) => model.provider === "xai") + .map((model) => model.id), + ).toEqual(ids); + expect( + (await registry.refreshModelCatalog()).models + .filter((model) => model.provider === "xai") + .map((model) => model.id), + ).toEqual(ids); + for (const model of [...models, customModel]) { + await expect(registry.canUseModel(model)).resolves.toBe(true); + expect(await registry.getApiKeyAndHeaders(model)).toMatchObject({ + ok: true, + requestModel: { + ...model, + api: credential.type === "oauth" ? "openai-responses" : model.api, + baseUrl: credential.type === "oauth" ? "https://api.x.ai/v1" : model.baseUrl, + }, + }); + } + } + }); + + test("uses the credential type returned with the key even if storage changes before dispatch", async () => { + storage.set("xai", oauth()); + const getAuth = storage.getApiKeyWithSourceToken.bind(storage); + vi.spyOn(storage, "getApiKeyWithSourceToken").mockImplementation(async (...args) => { + const result = await getAuth(...args); + storage.set("xai", { type: "api_key", key: "later-key" }); + return result; + }); + expect(await registry.getApiKeyAndHeaders(registry.find("xai", "grok-4.5")!)).toMatchObject({ + ok: true, + apiKey: "subscription-access", + requestModel: { api: "openai-responses" }, + }); + }); +}); From 72e4e9fbb644c5ef05117b819f6b80b9b1ae9a5b Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Mon, 14 Sep 2026 18:49:15 -0600 Subject: [PATCH 2/2] fix(coding-agent): preserve custom message timestamps (fixes #76) --- .pylon/features.yaml | 15 ++++ .pylon/upstream-review.md | 7 ++ .../pylon-custom-message-timestamps.md | 1 + .../coding-agent/src/core/agent-session.ts | 15 +++- .../coding-agent/src/core/session-manager.ts | 11 ++- .../76-custom-message-timestamps.test.ts | 68 +++++++++++++++++++ 6 files changed, 114 insertions(+), 3 deletions(-) create mode 100644 packages/coding-agent/.changes/pylon-custom-message-timestamps.md create mode 100644 packages/coding-agent/test/suite/regressions/76-custom-message-timestamps.test.ts diff --git a/.pylon/features.yaml b/.pylon/features.yaml index 78f47f6ace..37d769fb46 100644 --- a/.pylon/features.yaml +++ b/.pylon/features.yaml @@ -434,3 +434,18 @@ decisions: upstream_support: Absorb kernel pipe errors, park blocked idle waiters, and retry stale Codex continuation IDs once before output; retain Pylon lifecycle, cancellation, scoped identity, recovery cursors and ownership without wire changes. revisit_when: - The next complete upstream integration includes these commits; reconcile the selective patches without duplicating behavior. + + custom-message-persistence-identity: + area: runtime-reliability + state: candidate + owner: pylon-fork + decision: retain + pylon_refs: + - https://github.com/pylon-code/prime-agent/issues/76 + - https://github.com/pylon-code/pylon/issues/552 + upstream_refs: + - https://github.com/PrimeIntellect-ai/prime-agent/tree/f5859162c + fork_change: preserve-custom-message-timestamps + upstream_support: Current upstream allocates a second timestamp when persisting custom messages, so reloading or compacting changes live transcript identities. Preserve the original message timestamp at every AgentSession persistence boundary. + revisit_when: + - Upstream preserves exact custom-message identities across live emission, persistence and reconstruction. diff --git a/.pylon/upstream-review.md b/.pylon/upstream-review.md index a50295bcc0..f10dc27d38 100644 --- a/.pylon/upstream-review.md +++ b/.pylon/upstream-review.md @@ -334,3 +334,10 @@ Follow-up: Task10 builds/packs the exact merged tree into a private prefix and r - Adopt upstream #2314 (`7d1913969d33fa5a31366321e3f5db0070dcf2bf`): bind Codex continuation state to the producing connection and retry `previous_response_not_found` once with full context only before output. Preserve scoped provider hooks and avoid duplicate stream events or prompt replay. - Compatibility: no daemon commands, events, response shapes, protocol/schema revisions or public capability tokens change. No credentials, managed runtime, live session or release is modified. Existing installations require a separately verified immutable managed build to receive these fixes. - Validation: 125 focused coding-agent tests (including the real isolated kernel pipe regression) and 25 Codex stream tests pass. `npm run check` passes Biome, TypeScript, installer and browser bundle checks. Required hosted validation and final review are recorded in the owning PR. + +## 2026-09-14 — preserve custom message timestamps + +- Tracking #76 and Pylon #552. **Retain** a native persistence correction after a real isolated faux-provider compaction exposed different timestamps for the same live and persisted harness digest. Current upstream `f5859162c` still mints a new timestamp in `appendCustomMessageEntry`; its timestamp-related history has no equivalent correction. The complete integration checkpoint remains `1eee2938b4eeb7a4d72e17035adda669a89b63de`. +- Add an optional original timestamp to both custom append methods and pass it from every AgentSession custom-message write: delivered/queued messages, direct messages, session commands, Python state, harness digests, refinement and compaction outcomes. Calls that originate durable markers without a live message retain their existing current-time default. Invalid or lossy timestamps fail before publication. +- No daemon command, event, response shape, protocol/schema revision, artifact recipe or capability changes. Historical files are not rewritten. An older custom record whose live and persisted timestamps already disagree remains unproved; Pylon must not relax history matching to accept it. New managed bytes are required for the correction. +- Validation: 165 focused tests pass across the timestamp regression, session persistence/rollback, compaction and input queue. The regression deliberately advances the clock at the persistence boundary and verifies the exact live, pre-compaction tree, current context and reopened session messages. `npm run check` and required hosted checks accompany the PR. diff --git a/packages/coding-agent/.changes/pylon-custom-message-timestamps.md b/packages/coding-agent/.changes/pylon-custom-message-timestamps.md new file mode 100644 index 0000000000..56f6a92513 --- /dev/null +++ b/packages/coding-agent/.changes/pylon-custom-message-timestamps.md @@ -0,0 +1 @@ +- Fixed custom message timestamps changing when sessions are compacted or reloaded. diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 8d67ea6d03..b099de7950 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -2518,6 +2518,7 @@ export class AgentSession { message.content, message.display, message.details, + message.timestamp, ); this._emit({ type: "message_start", message }); this._emit({ type: "message_end", message }); @@ -4175,6 +4176,7 @@ export class AgentSession { event.message.content, event.message.display, event.message.details, + event.message.timestamp, ); } else if ( event.message.role === "user" || @@ -7076,6 +7078,7 @@ export class AgentSession { message.content, message.display, message.details, + message.timestamp, ); this.agent.state.messages.push(message); this._emit({ type: "message_start", message }); @@ -7161,6 +7164,7 @@ export class AgentSession { message.content, message.display, message.details, + appMessage.timestamp, ); this._emit({ type: "message_start", message: appMessage }); this._emit({ type: "message_end", message: appMessage }); @@ -8292,7 +8296,13 @@ export class AgentSession { } else { messages.push(message); } - this.sessionManager.appendCustomMessageEntry(message.customType, message.content, message.display, undefined); + this.sessionManager.appendCustomMessageEntry( + message.customType, + message.content, + message.display, + undefined, + message.timestamp, + ); this._emit({ type: "message_start", message }); this._emit({ type: "message_end", message }); } @@ -9087,6 +9097,7 @@ export class AgentSession { message.content, message.display, message.details, + message.timestamp, ); } catch { // Unpersisted session: context-only injection. @@ -9379,6 +9390,7 @@ export class AgentSession { message.content, message.display, message.details, + message.timestamp, ); } catch { // Not in the session file, so context rebuilds would drop the outcome. @@ -9688,6 +9700,7 @@ export class AgentSession { outcomeMessage.content, outcomeMessage.display, outcomeMessage.details, + outcomeMessage.timestamp, ); } catch (error) { const persistenceError = error instanceof Error ? error.message : String(error); diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 227e54b012..2403578d06 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -1998,7 +1998,11 @@ export class SessionManager { content: string | (TextContent | ImageContent)[], display: boolean, details?: T, + messageTimestamp?: number, ): string { + if (messageTimestamp !== undefined && !Number.isSafeInteger(messageTimestamp)) { + throw new Error("Custom message timestamp must be an integer millisecond value"); + } const entry: CustomMessageEntry = { type: "custom_message", customType, @@ -2007,7 +2011,7 @@ export class SessionManager { details, id: generateId(this.byId), parentId: this.leafId, - timestamp: new Date().toISOString(), + timestamp: new Date(messageTimestamp ?? Date.now()).toISOString(), }; this._appendEntry(entry); return entry.id; @@ -2022,8 +2026,11 @@ export class SessionManager { content: string | (TextContent | ImageContent)[], display: boolean, details?: T, + messageTimestamp?: number, ): string { - return this._appendEntryWithRollback(() => this.appendCustomMessageEntry(customType, content, display, details)); + return this._appendEntryWithRollback(() => + this.appendCustomMessageEntry(customType, content, display, details, messageTimestamp), + ); } private _appendEntryWithRollback(append: () => string): string { diff --git a/packages/coding-agent/test/suite/regressions/76-custom-message-timestamps.test.ts b/packages/coding-agent/test/suite/regressions/76-custom-message-timestamps.test.ts new file mode 100644 index 0000000000..d9237d442c --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/76-custom-message-timestamps.test.ts @@ -0,0 +1,68 @@ +import { fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { afterEach, expect, it, vi } from "vitest"; +import { buildSessionContext, SessionManager } from "../../../src/core/session-manager.js"; +import { createHarness, type Harness } from "../harness.js"; + +const harnesses: Harness[] = []; +afterEach(() => { + while (harnesses.length) harnesses.pop()!.cleanup(); + vi.restoreAllMocks(); + vi.useRealTimers(); +}); + +it("preserves live custom identities through delayed persistence, compaction and reload", async () => { + vi.useFakeTimers({ toFake: ["Date"] }); + const harness = await createHarness({ + persistSession: true, + settings: { compaction: { enabled: false, keepRecentTokens: 1 }, autoRefine: { enabled: false } }, + }); + harnesses.push(harness); + const append = harness.sessionManager.appendCustomMessageEntry.bind(harness.sessionManager); + vi.spyOn(harness.sessionManager, "appendCustomMessageEntry").mockImplementation((...args) => { + vi.setSystemTime(Date.now() + 50); + return append(...args); + }); + harness.setResponses(["one", "two", "summary", "turn summary"].map((text) => fauxAssistantMessage(text))); + await harness.session.prompt("one"); + await harness.session.sendCustomMessage({ + customType: "async_bash_completion", + content: "done", + display: true, + details: { pid: 1, command: "fixture", exitCode: 0 }, + }); + await harness.session.sendCustomMessage( + { customType: "ipython_state_restored", content: "restored", display: false, details: { restored: true } }, + { deliverAs: "nextTurn" }, + ); + await harness.session.prompt("two"); + const before = structuredClone(harness.session.messages); + expect(harness.sessionManager.buildSessionContext().messages).toEqual(before); + await harness.session.compact(); + const compaction = harness.sessionManager + .getEntries() + .slice() + .reverse() + .find((entry) => entry.type === "compaction"); + expect(compaction).toBeDefined(); + expect(buildSessionContext(harness.sessionManager.getEntries(), compaction!.parentId).messages).toEqual(before); + expect(harness.sessionManager.buildSessionContext().messages).toEqual(harness.session.messages); + const reopened = SessionManager.open(harness.sessionManager.getSessionFile()!); + expect(reopened.buildSessionContext().messages).toEqual(harness.session.messages); +}); + +it("preserves an explicit timestamp through rollback-backed writes and rejects lossy timestamps", async () => { + const harness = await createHarness({ persistSession: true }); + harnesses.push(harness); + const manager = harness.sessionManager; + manager.appendCustomMessageEntryWithRollback("refinement_notice", "notice", false, { source: "auto" }, 1234); + expect(manager.buildSessionContext().messages.at(-1)?.timestamp).toBe(1234); + const before = manager.getEntries(); + for (const timestamp of [NaN, Infinity, -Infinity, 1.5, Number.MAX_SAFE_INTEGER]) { + expect(() => + manager.appendCustomMessageEntryWithRollback("refinement_notice", "invalid", false, {}, timestamp), + ).toThrow(); + expect(manager.getEntries()).toEqual(before); + } + manager.appendCustomMessageEntry("default-time", "notice", false); + expect(manager.buildSessionContext().messages.at(-1)?.timestamp).toBeGreaterThan(1234); +});