@gajae-code/ai 0.4.5 → 0.5.1

package/src/providers/openai-completions.ts

@@ -1,4 +1,4 @@
-import { $env, $inheritedEnv, extractHttpStatusFromError } from "@gajae-code/utils";
+import { $credentialEnv, $env, $inheritedEnv, extractHttpStatusFromError, logger } from "@gajae-code/utils";
 import OpenAI from "openai";
 import type {
 	ChatCompletionAssistantMessageParam,
@@ -61,6 +61,12 @@ import { adaptSchemaForStrict, NO_STRICT, toolWireSchema } from "../utils/schema
 import { wrapFetchForSseDebug } from "../utils/sse-debug";
 import { type HealedToolCall, modelMayLeakKimiToolCalls, ToolCallHealer } from "../utils/tool-call-healing";
 import { isForcedToolChoice, mapToOpenAICompletionsToolChoice } from "../utils/tool-choice";
+import {
+	isForcedToolChoiceUnsupportedError,
+	markToolChoiceIncapability,
+	resolveToolChoice,
+} from "../utils/tool-choice-capability";
+import { COMPOSER_EDIT_DISCIPLINE_PROMPT, isComposerHarnessModel } from "./composer-discipline";
 import {
 	buildCopilotDynamicHeaders,
 	hasCopilotVisionInput,
@@ -493,7 +499,25 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions"> = (
 				});
 			} catch (error) {
 				const capturedErrorResponse = getCapturedErrorResponse();
-				if (
+				const sentForcedToolChoice = isForcedToolChoice(
+					(rawRequestDump?.body as { tool_choice?: unknown } | undefined)?.tool_choice,
+				);
+				if (firstTokenTime === undefined && isForcedToolChoiceUnsupportedError(error, sentForcedToolChoice)) {
+					const reason = await finalizeErrorMessage(error, rawRequestDump, capturedErrorResponse);
+					markToolChoiceIncapability(model, "auto", reason);
+					const resolvedToolChoice = resolveToolChoice(model, options?.toolChoice);
+					stream.push({
+						type: "toolChoiceIncapability",
+						api: model.api,
+						provider: model.provider,
+						model: model.id,
+						requestedLevel: resolvedToolChoice.requestedLevel,
+						resolvedLevel: "auto",
+						reason,
+						registryKey: resolvedToolChoice.registryKey,
+					});
+					openaiStream = await createCompletionsStream();
+				} else if (
 					isOpenRouterAnthropicModel(model) &&
 					!disableStrictTools &&
 					isCompiledGrammarTooLargeStrictError(error, capturedErrorResponse)
@@ -928,12 +952,12 @@ async function createClient(
 	clearCapturedErrorResponse: () => void;
 }> {
 	if (!apiKey) {
-		if (!$env.OPENAI_API_KEY) {
+		apiKey = $credentialEnv("OPENAI_API_KEY");
+		if (!apiKey) {
 			throw new Error(
 				"OpenAI API key is required. Set OPENAI_API_KEY environment variable or pass it as an argument.",
 			);
 		}
-		apiKey = $env.OPENAI_API_KEY;
 	}
 	const rawApiKey = apiKey;
 
@@ -1166,8 +1190,19 @@ function buildParams(
 		params.tools = [];
 	}
 
-	if (options?.toolChoice && compat.supportsToolChoice) {
-		params.tool_choice = mapToOpenAICompletionsToolChoice(options.toolChoice);
+	if (options?.toolChoice) {
+		const toolChoice = resolveToolChoice(model, options.toolChoice, compat);
+		if (toolChoice.degraded && toolChoice.supportSource === "runtime") {
+			logger.debug("openai-completions: degraded tool_choice after runtime capability discovery", {
+				model: model.id,
+				requestedLevel: toolChoice.requestedLevel,
+				resolvedLevel: toolChoice.resolvedLevel,
+				reason: toolChoice.reason,
+			});
+		}
+		if (toolChoice.resolvedChoice !== undefined) {
+			params.tool_choice = mapToOpenAICompletionsToolChoice(toolChoice.resolvedChoice);
+		}
 	}
 
 	if (params.tool_choice === "none" && (!Array.isArray(params.tools) || params.tools.length === 0)) {
@@ -1430,6 +1465,11 @@ export function convertMessages(
 	};
 
 	const systemPrompts = normalizeSystemPrompts(context.systemPrompt);
+	// Composer-harness models need anchor/edit discipline pinned ahead of the
+	// host prompt (see composer-discipline.ts for the observed failure modes).
+	if (systemPrompts.length > 0 && isComposerHarnessModel(model.id)) {
+		systemPrompts.unshift(COMPOSER_EDIT_DISCIPLINE_PROMPT);
+	}
 	if (systemPrompts.length > 0) {
 		const useDeveloperRole = model.reasoning && compat.supportsDeveloperRole;
 		const role = useDeveloperRole ? "developer" : "system";

package/src/providers/openai-request-transform.ts

@@ -39,6 +39,7 @@ const OPENAI_PROXY_STRIP_HEADERS = [
 	"x-stainless-helper-method",
 	"openai-organization",
 	"openai-project",
+	"openai-beta",
 ] as const;
 
 function resolveRequestTransform(

package/src/providers/openai-responses.ts

@@ -1,4 +1,11 @@
-import { $env, $inheritedEnv, extractHttpStatusFromError, structuredCloneJSON } from "@gajae-code/utils";
+import {
+	$credentialEnv,
+	$env,
+	$inheritedEnv,
+	extractHttpStatusFromError,
+	logger,
+	structuredCloneJSON,
+} from "@gajae-code/utils";
 import OpenAI from "openai";
 import type {
 	Tool as OpenAITool,
@@ -46,6 +53,11 @@ import { resolveRetryBudget } from "../utils/retry-budget";
 import { adaptSchemaForStrict, NO_STRICT, sanitizeSchemaForOpenAIResponses, toolWireSchema } from "../utils/schema";
 import { wrapFetchForSseDebug } from "../utils/sse-debug";
 import { mapToOpenAIResponsesToolChoice, type OpenAIResponsesToolChoice } from "../utils/tool-choice";
+import {
+	isForcedToolChoiceUnsupportedError,
+	markToolChoiceIncapability,
+	resolveToolChoice,
+} from "../utils/tool-choice-capability";
 import {
 	buildCopilotDynamicHeaders,
 	hasCopilotVisionInput,
@@ -278,7 +290,31 @@ export const streamOpenAIResponses: StreamFunction<"openai-responses"> = (
 					return data;
 				},
 				{ provider: model.provider, signal: requestSignal },
-			);
+			).catch(async error => {
+				if (!isForcedToolChoiceUnsupportedError(error, isForcedOpenAIResponsesToolChoice(params.tool_choice))) {
+					throw error;
+				}
+				const reason = await finalizeErrorMessage(error, rawRequestDump);
+				markToolChoiceIncapability(model, "auto", reason);
+				const resolvedToolChoice = resolveToolChoice(model, options?.toolChoice);
+				stream.push({
+					type: "toolChoiceIncapability",
+					api: model.api,
+					provider: model.provider,
+					model: model.id,
+					requestedLevel: resolvedToolChoice.requestedLevel,
+					resolvedLevel: "auto",
+					reason,
+					registryKey: resolvedToolChoice.registryKey,
+				});
+				delete params.tool_choice;
+				if (rawRequestDump) rawRequestDump.body = params;
+				const { data, response, request_id } = await client.responses
+					.create(params, { signal: requestSignal })
+					.withResponse();
+				await notifyProviderResponse(options, response, model, request_id);
+				return data;
+			});
 			const firstEventWatchdog = createWatchdog(
 				options?.streamFirstEventTimeoutMs ?? getStreamFirstEventTimeoutMs(idleTimeoutMs),
 				() => abortTracker.abortLocally(firstEventTimeoutAbortError),
@@ -363,12 +399,12 @@ function createClient(
 	baseUrl: string | undefined;
 } {
 	if (!apiKey) {
-		if (!$env.OPENAI_API_KEY) {
+		apiKey = $credentialEnv("OPENAI_API_KEY");
+		if (!apiKey) {
 			throw new Error(
 				"OpenAI API key is required. Set OPENAI_API_KEY environment variable or pass it as an argument.",
 			);
 		}
-		apiKey = $env.OPENAI_API_KEY;
 	}
 	const rawApiKey = apiKey;
 
@@ -490,7 +526,16 @@ function buildParams(
 	if (context.tools) {
 		params.tools = convertTools(context.tools, supportsStrictMode(model), model);
 		if (options?.toolChoice) {
-			params.tool_choice = mapOpenAIResponsesToolChoiceForTools(options.toolChoice, context.tools, model);
+			const toolChoice = resolveToolChoice(model, options.toolChoice);
+			if (toolChoice.degraded && toolChoice.supportSource === "runtime") {
+				logger.debug("openai-responses: degraded tool_choice after runtime capability discovery", {
+					model: model.id,
+					requestedLevel: toolChoice.requestedLevel,
+					resolvedLevel: toolChoice.resolvedLevel,
+					reason: toolChoice.reason,
+				});
+			}
+			params.tool_choice = mapOpenAIResponsesToolChoiceForTools(toolChoice.resolvedChoice, context.tools, model);
 		}
 		// The apply_patch spec §1 marks only `apply_patch` itself as
 		// `supports_parallel_tool_calls = false`. OpenAI's Responses API
@@ -651,6 +696,10 @@ export function mapOpenAIResponsesToolChoiceForTools(
 	return customTool ? { type: "custom", name: customTool.customWireName ?? customTool.name } : mapped;
 }
 
+function isForcedOpenAIResponsesToolChoice(choice: unknown): boolean {
+	return !!choice && choice !== "none" && choice !== "auto";
+}
+
 /** @internal Exported for tests. */
 export function convertTools(tools: Tool[], strictMode: boolean, model: Model<"openai-responses">): OpenAITool[] {
 	const allowFreeform = supportsFreeformApplyPatch(model);

package/src/providers/register-builtins.ts

@@ -6,9 +6,9 @@
  * openai) at startup. The loaded module promise is cached so subsequent calls
  * reuse the same import.
  *
- * NOTE: stream.ts currently imports providers directly, so this file is not yet
- * wired into the main streaming path. It provides the infrastructure for lazy
- * loading that can be integrated when stream.ts is refactored.
+ * stream.ts imports its provider stream functions from this module (see the
+ * lazy wrappers below), so this file IS the main streaming path's provider
+ * loader: heavy SDKs stay out of the CLI startup parse graph.
  */
 import type {
 	Api,
@@ -390,9 +390,8 @@ function loadBedrockProviderModule(): Promise<LazyProviderModule<"bedrock-conver
 // ---------------------------------------------------------------------------
 // Lazy stream function exports
 //
-// These use the same names as the direct provider stream functions. When
-// stream.ts is updated to import from this module instead of individual
-// providers, the lazy loading will take effect on the main code path.
+// Provider registry code imports these wrappers so the concrete provider modules
+// are loaded on first use instead of during package initialization.
 // ---------------------------------------------------------------------------
 
 export const streamAnthropic = createLazyStream(loadAnthropicProviderModule);

package/src/rate-limit-utils.ts

@@ -45,7 +45,16 @@ export function parseRateLimitReason(errorMessage: string): RateLimitReason {
 		return "RATE_LIMIT_EXCEEDED";
 	}
 
-	if (lower.includes("exhausted") || lower.includes("quota") || lower.includes("usage limit")) {
+	if (
+		lower.includes("exhausted") ||
+		lower.includes("quota") ||
+		lower.includes("usage limit") ||
+		lower.includes("model limit") ||
+		lower.includes("model_limit") ||
+		lower.includes("message limit") ||
+		lower.includes("message_limit") ||
+		lower.includes("limit for this model")
+	) {
 		return "QUOTA_EXHAUSTED";
 	}
 
@@ -77,7 +86,7 @@ export function calculateRateLimitBackoffMs(reason: RateLimitReason): number {
 
 /** Detect usage/quota limit errors in error messages (persistent, requires credential switch). */
 const USAGE_LIMIT_PATTERN =
-	/usage.?limit|usage_limit_reached|usage_not_included|limit_reached|quota.?exceeded|resource.?exhausted/i;
+	/usage.?limit|usage_limit_reached|usage_not_included|limit_reached|model.?limit|model_limit_reached|message.?limit|message_limit_reached|limit for this model|quota.?exceeded|resource has been exhausted[^\n]*(?:quota|limit)/i;
 
 export function isUsageLimitError(errorMessage: string): boolean {
 	return USAGE_LIMIT_PATTERN.test(errorMessage);

package/src/types.ts

@@ -161,6 +161,18 @@ export type ToolChoice =
 	| { type: "function"; function: { name: string } }
 	| { type: "tool"; name: string };
 
+export type ToolChoiceSupport = "none" | "auto" | "required" | "named";
+export type ToolChoiceSupportSource = "static" | "derived" | "runtime";
+
+export interface ToolChoiceCompat {
+	/** Maximum supported tool_choice level. */
+	toolChoiceSupport?: ToolChoiceSupport;
+	/** Legacy flag for accepting the tool_choice parameter. */
+	supportsToolChoice?: boolean;
+	/** Legacy flag for forced tool_choice support. */
+	supportsForcedToolChoice?: boolean;
+}
+
 // Base options all providers share
 export type CacheRetention = "none" | "short" | "long";
 
@@ -705,13 +717,24 @@ export type AssistantMessageEvent =
 			contentIndex?: undefined;
 			reason: Extract<StopReason, "aborted" | "error">;
 			error: AssistantMessage;
+	  }
+	| {
+			type: "toolChoiceIncapability";
+			contentIndex?: undefined;
+			api: string;
+			provider: string;
+			model: string;
+			requestedLevel: ToolChoiceSupport;
+			resolvedLevel: ToolChoiceSupport;
+			reason: string;
+			registryKey: string;
 	  };
 
 /**
  * Compatibility settings for openai-completions API.
  * Use this to override URL-based auto-detection for custom providers.
  */
-export interface OpenAICompat {
+export interface OpenAICompat extends ToolChoiceCompat {
 	/** Whether the provider supports the `store` field. Default: auto-detected from URL. */
 	supportsStore?: boolean;
 	/** Whether the provider supports the `developer` role (vs `system`). Default: auto-detected from URL. */
@@ -757,6 +780,8 @@ export interface OpenAICompat {
 	requiresAssistantContentForToolCalls?: boolean;
 	/** Whether the provider supports the `tool_choice` parameter. Default: true. */
 	supportsToolChoice?: boolean;
+	/** Whether `tool_choice` may force a tool (`required` / named tool). Default: true. */
+	supportsForcedToolChoice?: boolean;
 	/**
 	 * Drop reasoning fields (`reasoning_effort`, OpenRouter `reasoning`) for
 	 * the request when `tool_choice` forces a tool call. Mirrors the Anthropic
@@ -789,7 +814,7 @@ export interface OpenAICompat {
  * Use this to disable features that strict-by-default Anthropic accepts but
  * that proxy gateways (Vertex AI, AWS Bedrock-style fronts, etc.) reject.
  */
-export interface AnthropicCompat {
+export interface AnthropicCompat extends ToolChoiceCompat {
 	/**
 	 * Drop the top-level `strict: true` field on tool definitions. Vertex AI's
 	 * Anthropic-compatible endpoint rejects unknown tool fields with
@@ -911,7 +936,16 @@ export interface Model<TApi extends Api = any> {
 		? OpenAICompat
 		: TApi extends "anthropic-messages"
 			? AnthropicCompat
-			: never;
+			: TApi extends
+						| "bedrock-converse-stream"
+						| "google-generative-ai"
+						| "google-gemini-cli"
+						| "google-vertex"
+						| "ollama-chat"
+						| "azure-openai-responses"
+						| "openai-codex-responses"
+				? ToolChoiceCompat
+				: never;
 	/**
 	 * Which shape to use when exposing the OpenAI code backend `apply_patch` tool to this model.
 	 * Generated catalog policy sets `"freeform"` for first-party GPT-5 Responses

package/src/usage/grok-cli.ts

@@ -0,0 +1,163 @@
+import type {
+	CredentialRankingStrategy,
+	UsageFetchContext,
+	UsageFetchParams,
+	UsageLimit,
+	UsageProvider,
+	UsageReport,
+} from "../usage";
+
+interface BillingUsage {
+	monthlyLimit: number;
+	used: number;
+	billingPeriodEnd: string;
+}
+const DEFAULT_GROK_BUILD_BASE_URL = "https://cli-chat-proxy.grok.com/v1";
+const ALLOWED_GROK_BUILD_HOSTS = new Set(["cli-chat-proxy.grok.com"]);
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return !!value && typeof value === "object";
+}
+
+function finiteNumber(value: unknown): number | undefined {
+	return typeof value === "number" && Number.isFinite(value) ? value : undefined;
+}
+
+function parseValNumber(value: unknown): number | undefined {
+	return isRecord(value) ? finiteNumber(value.val) : undefined;
+}
+
+export function parseGrokCliBillingUsage(payload: unknown): BillingUsage {
+	if (!isRecord(payload) || !isRecord(payload.config)) {
+		throw new Error("invalid Grok CLI billing payload");
+	}
+	const monthlyLimit = parseValNumber(payload.config.monthlyLimit);
+	const used = parseValNumber(payload.config.used);
+	const billingPeriodEnd = payload.config.billingPeriodEnd;
+	if (
+		monthlyLimit === undefined ||
+		used === undefined ||
+		typeof billingPeriodEnd !== "string" ||
+		!Number.isFinite(new Date(billingPeriodEnd).getTime())
+	) {
+		throw new Error("invalid Grok CLI billing payload");
+	}
+	return { monthlyLimit, used, billingPeriodEnd };
+}
+
+function isAllowedGrokCredentialHost(baseUrl: string): boolean {
+	try {
+		const url = new URL(baseUrl);
+		return url.protocol === "https:" && ALLOWED_GROK_BUILD_HOSTS.has(url.hostname.toLowerCase());
+	} catch {
+		return false;
+	}
+}
+
+function normalizeGrokBaseUrl(baseUrl?: string): string {
+	const normalized = (baseUrl?.trim() || DEFAULT_GROK_BUILD_BASE_URL).replace(/\/+$/, "");
+	return isAllowedGrokCredentialHost(normalized) ? normalized : DEFAULT_GROK_BUILD_BASE_URL;
+}
+function isUnsafeGrokBaseUrlOverride(baseUrl?: string): boolean {
+	const normalized = baseUrl?.trim().replace(/\/+$/, "");
+	return !!normalized && !isAllowedGrokCredentialHost(normalized);
+}
+
+function resolveAccessToken(params: UsageFetchParams): string | undefined {
+	const token = params.credential.accessToken ?? params.credential.apiKey ?? process.env.GROK_CLI_OAUTH_TOKEN;
+	return token?.trim() || undefined;
+}
+
+function buildMonthlyUsageLimit(usage: BillingUsage, nowMs: number): UsageLimit {
+	const usedFraction = usage.monthlyLimit > 0 ? usage.used / usage.monthlyLimit : 0;
+	const percent = usedFraction * 100;
+	const resetsAt = new Date(usage.billingPeriodEnd).getTime();
+	return {
+		id: "grok-build:7d",
+		label: "SuperGrok monthly credits",
+		scope: { provider: "grok-build", shared: true, windowId: "7d" },
+		window: {
+			id: "7d",
+			label: "Monthly credits",
+			resetsAt,
+		},
+		amount: {
+			unit: "percent",
+			used: percent,
+			limit: 100,
+			remaining: Math.max(0, 100 - percent),
+			usedFraction,
+			remainingFraction: Math.max(0, 1 - usedFraction),
+		},
+		status: percent >= 95 ? "exhausted" : percent >= 80 ? "warning" : "ok",
+		notes: [
+			`${usage.used}/${usage.monthlyLimit} credits used`,
+			`resets in ${Math.max(0, Math.round((resetsAt - nowMs) / 3_600_000))}h`,
+		],
+	};
+}
+
+export const grokCliUsageProvider: UsageProvider = {
+	id: "grok-build",
+
+	supports(params) {
+		return params.provider === "grok-build";
+	},
+
+	async fetchUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise<UsageReport | null> {
+		const accessToken = resolveAccessToken(params);
+		if (!accessToken) {
+			ctx.logger?.warn("Grok Build usage: no access token", { provider: params.provider });
+			return null;
+		}
+
+		if (isUnsafeGrokBaseUrlOverride(params.baseUrl)) {
+			ctx.logger?.warn("Grok Build usage: ignoring unsafe base URL override for credential safety", {
+				provider: params.provider,
+			});
+		}
+		const billingBaseUrl = normalizeGrokBaseUrl(params.baseUrl);
+		const response = await ctx.fetch(`${billingBaseUrl}/billing`, {
+			headers: {
+				Authorization: `Bearer ${accessToken}`,
+				"x-xai-token-auth": "xai-grok-cli",
+				accept: "application/json",
+			},
+			signal: params.signal,
+		});
+		if (!response.ok) {
+			ctx.logger?.warn("Grok Build billing request failed", { status: response.status, provider: params.provider });
+			return null;
+		}
+
+		const payload = (await response.json()) as unknown;
+		let billing: BillingUsage;
+		try {
+			billing = parseGrokCliBillingUsage(payload);
+		} catch (error) {
+			ctx.logger?.warn("Grok Build billing parse failed", { error: String(error) });
+			return null;
+		}
+
+		const nowMs = Date.now();
+		return {
+			provider: "grok-build",
+			fetchedAt: nowMs,
+			limits: [buildMonthlyUsageLimit(billing, nowMs)],
+			metadata: {
+				email: params.credential.email,
+				accountId: params.credential.accountId,
+				subscription: true,
+			},
+			raw: payload,
+		};
+	},
+};
+
+export const grokCliRankingStrategy: CredentialRankingStrategy = {
+	findWindowLimits(report) {
+		const monthly = report.limits.find(limit => limit.id === "grok-build:7d");
+		return { secondary: monthly };
+	},
+	windowDefaults: { primaryMs: 5 * 60 * 60 * 1000, secondaryMs: 30 * 24 * 60 * 60 * 1000 },
+};

package/src/utils/event-stream.ts

@@ -2,7 +2,8 @@ import type { AssistantMessage, AssistantMessageEvent } from "../types";
 
 // Generic event stream class for async iteration
 export class EventStream<T, R = T> implements AsyncIterable<T> {
-	queue: T[] = [];
+	#queue: T[] = [];
+	#queueHead = 0;
 	waiting: Array<{ resolve: (value: IteratorResult<T>) => void; reject: (err: unknown) => void }> = [];
 	done = false;
 	#failed = false;
@@ -25,6 +26,35 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
 		this.extractResult = extractResult;
 	}
 
+	#enqueue(event: T): void {
+		this.#queue.push(event);
+	}
+
+	#dequeue(): T | undefined {
+		if (this.#queueHead >= this.#queue.length) return undefined;
+		const event = this.#queue[this.#queueHead]!;
+		this.#queue[this.#queueHead] = undefined as T;
+		this.#queueHead++;
+		if (this.#queueHead > 1024 && this.#queueHead * 2 >= this.#queue.length) {
+			this.#queue = this.#queue.slice(this.#queueHead);
+			this.#queueHead = 0;
+		}
+		return event;
+	}
+
+	get #queueLength(): number {
+		return this.#queue.length - this.#queueHead;
+	}
+
+	/**
+	 * Read-only snapshot of the not-yet-consumed events. Always a fresh copy:
+	 * external code can never mutate internal queue state or observe head-index
+	 * tombstones, so the deque cannot desynchronize.
+	 */
+	get queue(): T[] {
+		return this.#queue.slice(this.#queueHead);
+	}
+
 	push(event: T): void {
 		if (this.done) return;
 
@@ -38,7 +68,7 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
 		if (waiter) {
 			waiter.resolve({ value: event, done: false });
 		} else {
-			this.queue.push(event);
+			this.#enqueue(event);
 		}
 	}
 
@@ -47,7 +77,7 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
 		if (waiter) {
 			waiter.resolve({ value: event, done: false });
 		} else {
-			this.queue.push(event);
+			this.#enqueue(event);
 		}
 	}
 
@@ -84,8 +114,8 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
 
 	async *[Symbol.asyncIterator](): AsyncIterator<T> {
 		while (true) {
-			if (this.queue.length > 0) {
-				yield this.queue.shift()!;
+			if (this.#queueLength > 0) {
+				yield this.#dequeue()!;
 			} else if (this.#failed) {
 				throw this.#error;
 			} else if (this.done) {

package/src/utils/oauth/xai.ts

@@ -30,6 +30,15 @@ interface XaiTokenPayload {
 	token_type?: unknown;
 }
 
+export interface XaiOAuthFlowOptions {
+	extraAuthorizeParams?: Readonly<Record<string, string>>;
+}
+
+export interface XaiOAuthRefreshOptions {
+	signal?: AbortSignal;
+	extraTokenParams?: Readonly<Record<string, string>>;
+}
+
 interface XaiJwtPayload {
 	sub?: unknown;
 	email?: unknown;
@@ -41,6 +50,28 @@ function requestSignal(signal: AbortSignal | undefined): AbortSignal {
 	return signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal;
 }
 
+function addNonOverridingParams(
+	target: URLSearchParams | Record<string, string>,
+	params: Readonly<Record<string, string>>,
+): void {
+	for (const [key, value] of Object.entries(params)) {
+		if (key.length === 0 || value.length === 0) continue;
+		if (target instanceof URLSearchParams) {
+			if (!target.has(key)) target.set(key, value);
+		} else if (!(key in target)) {
+			target[key] = value;
+		}
+	}
+}
+
+function isAbortSignal(value: AbortSignal | XaiOAuthRefreshOptions | undefined): value is AbortSignal {
+	return value instanceof AbortSignal;
+}
+
+function resolveRefreshOptions(options: AbortSignal | XaiOAuthRefreshOptions | undefined): XaiOAuthRefreshOptions {
+	return isAbortSignal(options) ? { signal: options } : (options ?? {});
+}
+
 function validateXaiEndpoint(rawUrl: string): string {
 	const parsed = new URL(rawUrl);
 	const host = parsed.hostname.toLowerCase();
@@ -136,8 +167,9 @@ function credentialsFromTokenPayload(payload: XaiTokenPayload, refreshFallback =
 export class XaiOAuthFlow extends OAuthCallbackFlow {
 	#verifier = "";
 	#discovery: XaiDiscovery | undefined;
+	#extraAuthorizeParams: Readonly<Record<string, string>>;
 
-	constructor(ctrl: OAuthController) {
+	constructor(ctrl: OAuthController, options: XaiOAuthFlowOptions = {}) {
 		super(ctrl, {
 			preferredPort: XAI_OAUTH_CALLBACK_PORT,
 			callbackPath: XAI_OAUTH_CALLBACK_PATH,
@@ -145,6 +177,7 @@ export class XaiOAuthFlow extends OAuthCallbackFlow {
 			callbackBindHostname: "127.0.0.1",
 			redirectUri: `http://127.0.0.1:${XAI_OAUTH_CALLBACK_PORT}${XAI_OAUTH_CALLBACK_PATH}`,
 		} satisfies OAuthCallbackFlowOptions);
+		this.#extraAuthorizeParams = options.extraAuthorizeParams ?? {};
 	}
 
 	async generateAuthUrl(state: string, redirectUri: string): Promise<{ url: string; instructions?: string }> {
@@ -161,6 +194,7 @@ export class XaiOAuthFlow extends OAuthCallbackFlow {
 			state,
 			nonce: crypto.randomUUID(),
 		});
+		addNonOverridingParams(params, this.#extraAuthorizeParams);
 		return {
 			url: `${this.#discovery.authorizationEndpoint}?${params.toString()}`,
 			instructions:
@@ -188,23 +222,25 @@ export class XaiOAuthFlow extends OAuthCallbackFlow {
 	}
 }
 
-export async function loginXai(ctrl: OAuthController): Promise<OAuthCredentials> {
-	return new XaiOAuthFlow(ctrl).login();
+export async function loginXai(ctrl: OAuthController, options?: XaiOAuthFlowOptions): Promise<OAuthCredentials> {
+	return new XaiOAuthFlow(ctrl, options).login();
 }
 
-export async function refreshXaiToken(refreshToken: string, signal?: AbortSignal): Promise<OAuthCredentials> {
+export async function refreshXaiToken(
+	refreshToken: string,
+	options?: AbortSignal | XaiOAuthRefreshOptions,
+): Promise<OAuthCredentials> {
 	if (!refreshToken) {
 		throw new Error("xAI credentials are expired and do not include a refresh token");
 	}
+	const { signal, extraTokenParams = {} } = resolveRefreshOptions(options);
 	const discovery = await discoverXaiOAuthEndpoints(signal);
-	const tokenPayload = await postXaiToken(
-		discovery.tokenEndpoint,
-		{
-			grant_type: "refresh_token",
-			client_id: XAI_OAUTH_CLIENT_ID,
-			refresh_token: refreshToken,
-		},
-		signal,
-	);
+	const body = {
+		grant_type: "refresh_token",
+		client_id: XAI_OAUTH_CLIENT_ID,
+		refresh_token: refreshToken,
+	};
+	addNonOverridingParams(body, extraTokenParams);
+	const tokenPayload = await postXaiToken(discovery.tokenEndpoint, body, signal);
 	return credentialsFromTokenPayload(tokenPayload, refreshToken);
 }