Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion scripts/test-layout/layout.json
Original file line number Diff line number Diff line change
Expand Up @@ -1558,7 +1558,8 @@
"gui-codex-usage-score-parity.test.ts": "gui",
"web-search-sidecar-429.test.ts": "web-search",
"management-google-tool-schema-policy.test.ts": "server",
"codex-shim-destroyed-probe.test.ts": "codex-integration"
"codex-shim-destroyed-probe.test.ts": "codex-integration",
"client-runtime.test.ts": "clients"
},
"migrated": [
"adapters",
Expand Down
6 changes: 2 additions & 4 deletions src/adapters/openai-chat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -297,10 +297,8 @@ export function createOpenAIChatAdapter(provider: OcxProviderConfig): ProviderAd
};
};
if (hasShrinkableOpenAIChatImages(messages)) {
return normalizeOpenAIChatImages(messages, {
tierBias: incoming?.imageTierBias,
abortSignal: incoming?.abortSignal,
}).then(finish, error => {
const imageOptions = { tierBias: incoming?.imageTierBias, abortSignal: incoming?.abortSignal };
return normalizeOpenAIChatImages(messages, imageOptions).then(finish, error => {
if (incoming?.abortSignal?.aborted) throw error;
return finish();
});
Expand Down
121 changes: 3 additions & 118 deletions src/oauth/index.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import type { KiroOAuthMetadata, OAuthController, OAuthCredentials } from "./types";
import { initializeProviderModelSelection } from "../providers/initial-model-selection";
import { parseCallbackInput } from "./callback-server";
import type { OcxConfig, OcxProviderConfig, RefreshPolicy } from "../types";
import { ConfigMutationLockError, loadConfig, mutatePersistedConfig, saveConfig } from "../config";
import { resolveProviderApiKey } from "../providers/key-store";
Expand Down Expand Up @@ -51,8 +50,9 @@ import { providerModelsUrl, resolveProviderModelDiscoveryUrl } from "../provider
import { resolveProviderTransport } from "../providers/xai-transport";
import { detectClaudeCodeToken, detectGrokCliToken, hasComparableGrokIdentity, isSameGrokIdentity, shouldAdoptGrokGeneration } from "./local-token-detect";
import { logOAuthEvent } from "./log";
import { captureConfigGeneration, sweepExpiredOnWrite, type GenerationContext } from "../lib/state-store-sweeper";
import { retainedUtf8Bytes } from "../lib/admission";
import { captureConfigGeneration, sweepExpiredOnWrite } from "../lib/state-store-sweeper";
import { clearManualCodeSlot, ensureManualCodeSlot, kiroLoginSettling, loginAbort, loginState, waitForManualLoginCode } from "./login-flow-state";
export { reconcileOAuthFlowState, submitManualLoginCode } from "./login-flow-state";
import { randomUUID } from "node:crypto";
export {
CODEX_HEALTH_AUTH_FAILED_NOTE,
Expand Down Expand Up @@ -1710,121 +1710,6 @@ export async function runLogin(
* localhost), the GUI can POST the final redirect URL or authorization code via
* submitManualLoginCode(), which feeds OAuthController.onManualCodeInput.
*/
const loginState = new Map<string, { error?: string; done: boolean }>();
const loginAbort = new Map<string, AbortController>();
const kiroLoginSettling = new Set<string>();

/** Pending paste for a login in progress: either a waiter or a stashed early submission. */
interface ManualCodeSlot {
pendingInput?: string;
resolve?: (value: string) => void;
/** Registered by the callback flow so submits can validate state synchronously. */
expectedState?: string;
}
const loginManual = new Map<string, ManualCodeSlot>();
const OAUTH_PENDING_CODE_MAX_BYTES = 4 * 1024;
let lastOAuthFlowReconciledGeneration = 0;

export function reconcileOAuthFlowState(context: GenerationContext): number {
if (context.generation <= lastOAuthFlowReconciledGeneration) return 0;
let removed = 0;
for (const [provider, state] of loginState) {
if (context.providerNames.has(provider) || !state.done || loginAbort.has(provider)) continue;
if (loginState.delete(provider)) removed += 1;
if (loginManual.delete(provider)) removed += 1;
if (loginAbort.delete(provider)) removed += 1;
}
lastOAuthFlowReconciledGeneration = context.generation;
return removed;
}

function clearManualCodeSlot(provider: string): void {
loginManual.delete(provider);
}

function ensureManualCodeSlot(provider: string): ManualCodeSlot {
let slot = loginManual.get(provider);
if (!slot) {
slot = {};
loginManual.set(provider, slot);
}
return slot;
}

/** Wait for a GUI/CLI paste of the OAuth redirect URL or code (or return a stashed early submit). */
function waitForManualLoginCode(provider: string, signal: AbortSignal, expectedState?: string): Promise<string> {
if (signal.aborted) {
return Promise.reject(new Error(`OAuth callback cancelled: ${signal.reason}`));
}
const slot = ensureManualCodeSlot(provider);
if (expectedState !== undefined) slot.expectedState = expectedState;
if (slot.pendingInput !== undefined) {
const value = slot.pendingInput;
slot.pendingInput = undefined;
return Promise.resolve(value);
}
return new Promise<string>((resolve, reject) => {
const onAbort = () => {
if (slot.resolve === resolve) slot.resolve = undefined;
reject(new Error(`OAuth callback cancelled: ${signal.reason}`));
};
signal.addEventListener("abort", onAbort, { once: true });
slot.resolve = (value: string) => {
signal.removeEventListener("abort", onAbort);
if (slot.resolve === resolve) slot.resolve = undefined;
resolve(value);
};
});
}

/**
* Feed a pasted redirect URL or authorization code into an in-progress GUI login.
* Returns ok:false when no login is waiting (or input is empty). Invalid pastes are accepted
* here and re-prompted by the OAuth callback loop if they cannot be parsed / fail state checks.
*/
export function submitManualLoginCode(provider: string, input: string): { ok: true } | { ok: false; error: string } {
const trimmed = input.trim();
if (!trimmed) return { ok: false, error: "empty code" };
if (retainedUtf8Bytes(trimmed) > OAUTH_PENDING_CODE_MAX_BYTES) return { ok: false, error: "code too large" };
const st = loginState.get(provider);
if (!st || st.done) return { ok: false, error: "no login in progress" };
const slot = ensureManualCodeSlot(provider);
// Synchronous validation (validated request/ack): reject un-parseable input and
// authorization responses (url/query kind) whose state is missing or mismatched
// once the flow has registered its expected state. Raw codes stay in-session-PKCE
// protected — but a raw paste carrying an explicit #state suffix is state-bearing
// and checked too. Early posts (flow not yet waiting, no expectedState) are
// stashed and re-validated by the callback loop.
const parsed = parseCallbackInput(trimmed);
// Command Code's manual fallback accepts a pasted JSON callback payload
// (`{ apiKey, state, ... }`) which has no `code` param. Let it through the
// shared gate so the provider-specific parser can validate it.
const isCommandCodeJson = provider === "command-code" && trimmed.startsWith("{") && !parsed.code;
if (!parsed.code && !isCommandCodeJson) return { ok: false, error: "no authorization code found in input" };
// A raw paste carrying an explicit code#state suffix is state-bearing too: it
// must match the expected state rather than bypass validation.
const stateBearing = parsed.kind !== "raw" || parsed.state !== undefined;
if (stateBearing && slot.expectedState !== undefined) {
if (parsed.state === undefined) return { ok: false, error: "redirect URL is missing the state parameter" };
if (parsed.state !== slot.expectedState) {
return {
ok: false,
error: parsed.kind === "raw"
? "state mismatch — paste the bare code, or the correct code#state from THIS login attempt"
: "state mismatch — paste the redirect URL from THIS login attempt",
};
}
}
if (slot.resolve) {
const resolve = slot.resolve;
slot.resolve = undefined;
resolve(trimmed);
} else {
// Race: GUI may POST before the flow reaches onManualCodeInput — stash for the waiter.
slot.pendingInput = trimmed;
}
return { ok: true };
}

export interface OAuthAccountSummary {
id: string;
Expand Down
127 changes: 127 additions & 0 deletions src/oauth/login-flow-state.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,127 @@
import { parseCallbackInput } from "./callback-server";
import { retainedUtf8Bytes } from "../lib/admission";
import type { GenerationContext } from "../lib/state-store-sweeper";

/**
* In-flight login flow state and the manual paste slot that feeds it.
*
* Split out of `index.ts` when that file crossed the repository file-size ratchet at 2009
* lines. The cut follows a seam rather than a line count: everything here is bookkeeping for
* a login that has started and not yet settled, and none of it reads or writes a stored
Comment on lines +8 to +10

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Update the owned OAuth architecture documents

This creates a new owner for in-flight OAuth login state, but none of the four architecture documents mapped to src/oauth/ (structure/runtime.md, structure/transports/inventory.md, structure/providers-and-adapters.md, and structure/providers/xai-grok.md) is updated. Add the new module boundary to the relevant sections of every mapped document so the maintained architecture reference continues to identify where this security-sensitive state lives.

AGENTS.md reference: src/AGENTS.md:L10-L11

Useful? React with 👍 / 👎.

* credential. `index.ts` re-exports the two public names, so existing importers are unaffected.
*/
export const loginState = new Map<string, { error?: string; done: boolean }>();
export const loginAbort = new Map<string, AbortController>();
export const kiroLoginSettling = new Set<string>();

/** Pending paste for a login in progress: either a waiter or a stashed early submission. */
export interface ManualCodeSlot {
pendingInput?: string;
resolve?: (value: string) => void;
/** Registered by the callback flow so submits can validate state synchronously. */
expectedState?: string;
}
const loginManual = new Map<string, ManualCodeSlot>();
const OAUTH_PENDING_CODE_MAX_BYTES = 4 * 1024;
let lastOAuthFlowReconciledGeneration = 0;

export function reconcileOAuthFlowState(context: GenerationContext): number {
if (context.generation <= lastOAuthFlowReconciledGeneration) return 0;
let removed = 0;
for (const [provider, state] of loginState) {
if (context.providerNames.has(provider) || !state.done || loginAbort.has(provider)) continue;
if (loginState.delete(provider)) removed += 1;
if (loginManual.delete(provider)) removed += 1;
if (loginAbort.delete(provider)) removed += 1;
}
lastOAuthFlowReconciledGeneration = context.generation;
return removed;
}

export function clearManualCodeSlot(provider: string): void {
loginManual.delete(provider);
}

export function ensureManualCodeSlot(provider: string): ManualCodeSlot {
let slot = loginManual.get(provider);
if (!slot) {
slot = {};
loginManual.set(provider, slot);
}
return slot;
}

/** Wait for a GUI/CLI paste of the OAuth redirect URL or code (or return a stashed early submit). */
export function waitForManualLoginCode(provider: string, signal: AbortSignal, expectedState?: string): Promise<string> {
if (signal.aborted) {
return Promise.reject(new Error(`OAuth callback cancelled: ${signal.reason}`));
}
const slot = ensureManualCodeSlot(provider);
if (expectedState !== undefined) slot.expectedState = expectedState;
if (slot.pendingInput !== undefined) {
const value = slot.pendingInput;
slot.pendingInput = undefined;
return Promise.resolve(value);
}
return new Promise<string>((resolve, reject) => {
const onAbort = () => {
if (slot.resolve === resolve) slot.resolve = undefined;
reject(new Error(`OAuth callback cancelled: ${signal.reason}`));
};
signal.addEventListener("abort", onAbort, { once: true });
slot.resolve = (value: string) => {
signal.removeEventListener("abort", onAbort);
if (slot.resolve === resolve) slot.resolve = undefined;
resolve(value);
};
});
}

/**
* Feed a pasted redirect URL or authorization code into an in-progress GUI login.
* Returns ok:false when no login is waiting (or input is empty). Invalid pastes are accepted
* here and re-prompted by the OAuth callback loop if they cannot be parsed / fail state checks.
*/
export function submitManualLoginCode(provider: string, input: string): { ok: true } | { ok: false; error: string } {
const trimmed = input.trim();
if (!trimmed) return { ok: false, error: "empty code" };
if (retainedUtf8Bytes(trimmed) > OAUTH_PENDING_CODE_MAX_BYTES) return { ok: false, error: "code too large" };
const st = loginState.get(provider);
if (!st || st.done) return { ok: false, error: "no login in progress" };
const slot = ensureManualCodeSlot(provider);
// Synchronous validation (validated request/ack): reject un-parseable input and
// authorization responses (url/query kind) whose state is missing or mismatched
// once the flow has registered its expected state. Raw codes stay in-session-PKCE
// protected — but a raw paste carrying an explicit #state suffix is state-bearing
// and checked too. Early posts (flow not yet waiting, no expectedState) are
// stashed and re-validated by the callback loop.
const parsed = parseCallbackInput(trimmed);
// Command Code's manual fallback accepts a pasted JSON callback payload
// (`{ apiKey, state, ... }`) which has no `code` param. Let it through the
// shared gate so the provider-specific parser can validate it.
const isCommandCodeJson = provider === "command-code" && trimmed.startsWith("{") && !parsed.code;
if (!parsed.code && !isCommandCodeJson) return { ok: false, error: "no authorization code found in input" };
// A raw paste carrying an explicit code#state suffix is state-bearing too: it
// must match the expected state rather than bypass validation.
const stateBearing = parsed.kind !== "raw" || parsed.state !== undefined;
if (stateBearing && slot.expectedState !== undefined) {
if (parsed.state === undefined) return { ok: false, error: "redirect URL is missing the state parameter" };
if (parsed.state !== slot.expectedState) {
return {
ok: false,
error: parsed.kind === "raw"
? "state mismatch — paste the bare code, or the correct code#state from THIS login attempt"
: "state mismatch — paste the redirect URL from THIS login attempt",
};
}
}
if (slot.resolve) {
const resolve = slot.resolve;
slot.resolve = undefined;
resolve(trimmed);
} else {
// Race: GUI may POST before the flow reaches onManualCodeInput — stash for the waiter.
slot.pendingInput = trimmed;
}
return { ok: true };
}
3 changes: 2 additions & 1 deletion tests/fixtures/test-layout-expected.json
Original file line number Diff line number Diff line change
Expand Up @@ -1390,5 +1390,6 @@
"responses-code-mode-patch-compile.test.ts": "responses",
"gui-codex-usage-score-parity.test.ts": "gui",
"web-search-sidecar-429.test.ts": "web-search",
"codex-shim-destroyed-probe.test.ts": "codex-integration"
"codex-shim-destroyed-probe.test.ts": "codex-integration",
"client-runtime.test.ts": "clients"
}
Loading