diff --git a/src/browser/App.tsx b/src/browser/App.tsx index 6fb376dd4c0..a8f12c0d55e 100644 --- a/src/browser/App.tsx +++ b/src/browser/App.tsx @@ -4,6 +4,7 @@ import { useLocation, useNavigate } from "react-router-dom"; import "./styles/globals.css"; import { useWorkspaceContext, toWorkspaceSelection } from "./contexts/WorkspaceContext"; import { useProjectContext } from "./contexts/ProjectContext"; +import { getCodexOauthProjectPath } from "@/common/utils/providers/codexOauthRouting"; import type { WorkspaceSelection } from "./components/ProjectSidebar/ProjectSidebar"; import { LeftSidebar } from "./components/LeftSidebar/LeftSidebar"; import { ProjectCreateModal } from "./components/ProjectCreateModal/ProjectCreateModal"; @@ -197,6 +198,7 @@ function AppInner() { const { userProjects, + getProjectConfig, refreshProjects, removeProject, openProjectCreateModal, @@ -263,6 +265,18 @@ function AppInner() { ) : null; const creationScopeId = creationScope ? getProjectScopeId(creationScope.projectPath) : null; + const accountWorkspaceId = selectedWorkspace?.workspaceId ?? currentWorkspaceId; + const accountProjectPath = getCodexOauthProjectPath( + accountWorkspaceId + ? (workspaceMetadata.get(accountWorkspaceId) ?? selectedWorkspace) + : { + projectPath: creationScope?.projectPath, + subProjectPath: creationScope?.subProjectPath ?? undefined, + } + ); + const codexOauthAccountId = accountProjectPath + ? getProjectConfig(accountProjectPath)?.codexOauthAccountId + : undefined; // History navigation (back/forward) const navigate = useNavigate(); @@ -690,9 +704,17 @@ function AppInner() { const provider = getFastModeProvider(model, { providersConfig, resolvedRouteProvider: getRouteForModel(normalizeToCanonical(model)), + codexOauthAccountId, }); return provider != null && providersConfig[provider]?.serviceTier === "priority"; - }, [creationScopeId, getModelForWorkspace, getRouteForModel, providersConfig, selectedWorkspace]); + }, [ + codexOauthAccountId, + creationScopeId, + getModelForWorkspace, + getRouteForModel, + providersConfig, + selectedWorkspace, + ]); const fastModeToggleInFlightRef = useRef(false); const toggleFastMode = useCallback(async () => { @@ -707,6 +729,7 @@ function AppInner() { const provider = getFastModeProvider(model, { providersConfig, resolvedRouteProvider: getRouteForModel(normalizeToCanonical(model)), + codexOauthAccountId, }); if (provider == null) { fastModeToggleInFlightRef.current = false; @@ -736,6 +759,7 @@ function AppInner() { } }, [ api, + codexOauthAccountId, creationScopeId, getModelForWorkspace, getRouteForModel, @@ -984,6 +1008,7 @@ function AppInner() { workspaceMetadata, selectedWorkspace, creationScopeId, + codexOauthAccountId, themePreference, getThinkingLevel: getThinkingLevelForWorkspace, onSetThinkingLevel: setThinkingLevelFromPalette, diff --git a/src/browser/components/ChatPane/ChatPane.tsx b/src/browser/components/ChatPane/ChatPane.tsx index f1c333bd245..be8ba2d05f6 100644 --- a/src/browser/components/ChatPane/ChatPane.tsx +++ b/src/browser/components/ChatPane/ChatPane.tsx @@ -8,6 +8,7 @@ import React, { useMemo, } from "react"; import { Lightbulb } from "lucide-react"; +import { getCodexOauthProjectPath } from "@/common/utils/providers/codexOauthRouting"; import { MessageListProvider } from "@/browser/features/Messages/MessageListContext"; import { cn } from "@/common/lib/utils"; import { ChatInstructionsChatDecoration } from "@/browser/components/InstructionsTab/AdditionalSystemContextScratchpad"; @@ -99,6 +100,7 @@ import { useReviews } from "@/browser/hooks/useReviews"; import { ReviewsBanner } from "../ReviewsBanner/ReviewsBanner"; import type { ReviewNoteData } from "@/common/types/review"; import { useWorkspaceContext } from "@/browser/contexts/WorkspaceContext"; +import { useProjectContext } from "@/browser/contexts/ProjectContext"; import { useBackgroundBashActions, useBackgroundBashError, @@ -467,6 +469,11 @@ const ChatPaneContent: React.FC = (props) => { loadingOlderHistory, activeBashMonitorCount, } = workspaceState; + const { getProjectConfig } = useProjectContext(); + const accountProjectPath = getCodexOauthProjectPath(meta ?? { projectPath }); + const codexOauthAccountId = accountProjectPath + ? getProjectConfig(accountProjectPath)?.codexOauthAccountId + : undefined; const shouldShowPinnedTodoList = workspaceState.todos.length > 0; const shouldShowReviewsBanner = reviews.reviews.length > 0; const shouldRenderLoadOlderMessagesButton = hasOlderHistory && !isPixelSnapshotEnvironment(); @@ -486,6 +493,7 @@ const ChatPaneContent: React.FC = (props) => { api: api ?? undefined, pendingSendOptions, providersConfig, + codexOauthAccountId, }); // Apply message transformations: @@ -577,9 +585,17 @@ const ChatPaneContent: React.FC = (props) => { use1M, autoCompactionThreshold / 100, undefined, - providersConfig + providersConfig, + { codexOauthAccountId } ), - [workspaceUsage, pendingModel, use1M, providersConfig, autoCompactionThreshold] + [ + workspaceUsage, + pendingModel, + use1M, + providersConfig, + autoCompactionThreshold, + codexOauthAccountId, + ] ); // Show warning when: shouldShowWarning flag is true AND not currently compacting. diff --git a/src/browser/components/ProjectSidebar/ProjectSidebar.test.tsx b/src/browser/components/ProjectSidebar/ProjectSidebar.test.tsx index 36e52ffab41..6c464bd8925 100644 --- a/src/browser/components/ProjectSidebar/ProjectSidebar.test.tsx +++ b/src/browser/components/ProjectSidebar/ProjectSidebar.test.tsx @@ -534,6 +534,8 @@ function installProjectSidebarTestDoubles() { providersExpandedProvider: null, setProvidersExpandedProvider: () => undefined, providersStartCoderLogin: false, + codexAccountAction: null, + setCodexAccountAction: () => undefined, setProvidersStartCoderLogin: () => undefined, runtimesProjectPath: null, setRuntimesProjectPath: () => undefined, diff --git a/src/browser/components/ThinkingSelector/ThinkingSelector.tsx b/src/browser/components/ThinkingSelector/ThinkingSelector.tsx index a7b1c8c644b..0cef3629df6 100644 --- a/src/browser/components/ThinkingSelector/ThinkingSelector.tsx +++ b/src/browser/components/ThinkingSelector/ThinkingSelector.tsx @@ -56,6 +56,7 @@ export interface ThinkingInheritOption { interface ThinkingSelectorControlProps { modelString: string | undefined; + codexOauthAccountId?: string; /** Delegated preferences may inherit a model that is only known at launch. */ modelCapabilitiesDeferred?: boolean; /** Independent of effort/model inheritance; false denotes an explicit mode override. */ @@ -130,12 +131,14 @@ export const ThinkingSelectorControl: React.FC = ( openaiProModeAvailable(props.modelString, { providersConfig, resolvedRouteProvider: resolvedRoute, + codexOauthAccountId: props.codexOauthAccountId, }), fastModeProvider: props.allowFastMode !== false && providersConfig != null ? getFastModeProvider(props.modelString, { providersConfig, resolvedRouteProvider: resolvedRoute, + codexOauthAccountId: props.codexOauthAccountId, }) : null, }; @@ -437,6 +440,7 @@ export const ThinkingSelectorControl: React.FC = ( interface ThinkingSelectorProps { modelString: string; + codexOauthAccountId?: string; /** Some embedded clients cannot resolve route-aware provider options safely. */ allowProMode?: boolean; /** Some embedded clients do not expose provider configuration mutations. */ @@ -451,6 +455,7 @@ export const ThinkingSelector: React.FC = (props) => { return ( void; + codexAccountAction: CodexAccountSettingsIntent | null; + setCodexAccountAction: Dispatch>; + /** One-shot hint for RuntimesSection to pre-select a project scope. */ runtimesProjectPath: string | null; setRuntimesProjectPath: (path: string | null) => void; @@ -68,6 +80,9 @@ export function SettingsProvider(props: { children: ReactNode }) { const router = useRouter(); const [providersExpandedProvider, setProvidersExpandedProvider] = useState(null); const [providersStartCoderLogin, setProvidersStartCoderLogin] = useState(false); + const [codexAccountAction, setCodexAccountAction] = useState( + null + ); const [runtimesProjectPath, setRuntimesProjectPath] = useState(null); const [secretsProjectPath, setSecretsProjectPath] = useState(null); const [instructionsProjectPath, setInstructionsProjectPath] = useState(null); @@ -83,9 +98,14 @@ export function SettingsProvider(props: { children: ReactNode }) { if (nextSection === "providers") { setProvidersExpandedProvider(options?.expandProvider ?? null); setProvidersStartCoderLogin(options?.startCoderLogin ?? false); + // A fresh identity lets repeated commands reach an already-open settings section. + setCodexAccountAction( + options?.codexAccountAction ? { ...options.codexAccountAction } : null + ); } else { setProvidersExpandedProvider(null); setProvidersStartCoderLogin(false); + setCodexAccountAction(null); } if (nextSection === "runtimes") { setRuntimesProjectPath(options?.runtimesProjectPath ?? null); @@ -121,6 +141,7 @@ export function SettingsProvider(props: { children: ReactNode }) { if (wasOpenRef.current && !isOpen) { setProvidersExpandedProvider(null); setProvidersStartCoderLogin(false); + setCodexAccountAction(null); setRuntimesProjectPath(null); setSecretsProjectPath(null); setInstructionsProjectPath(null); @@ -134,6 +155,7 @@ export function SettingsProvider(props: { children: ReactNode }) { const close = useCallback(() => { setProvidersExpandedProvider(null); setProvidersStartCoderLogin(false); + setCodexAccountAction(null); setRuntimesProjectPath(null); setSecretsProjectPath(null); setInstructionsProjectPath(null); @@ -145,6 +167,7 @@ export function SettingsProvider(props: { children: ReactNode }) { if (section !== "providers") { setProvidersExpandedProvider(null); setProvidersStartCoderLogin(false); + setCodexAccountAction(null); } if (section !== "runtimes") { // Runtime scope hints are one-shot and should not persist across section changes. @@ -173,6 +196,8 @@ export function SettingsProvider(props: { children: ReactNode }) { setProvidersExpandedProvider, providersStartCoderLogin, setProvidersStartCoderLogin, + codexAccountAction, + setCodexAccountAction, runtimesProjectPath, setRuntimesProjectPath, secretsProjectPath, @@ -189,6 +214,7 @@ export function SettingsProvider(props: { children: ReactNode }) { registerOnClose, providersExpandedProvider, providersStartCoderLogin, + codexAccountAction, runtimesProjectPath, secretsProjectPath, instructionsProjectPath, diff --git a/src/browser/features/ChatInput/index.tsx b/src/browser/features/ChatInput/index.tsx index 3eee481a2b0..8c266faa977 100644 --- a/src/browser/features/ChatInput/index.tsx +++ b/src/browser/features/ChatInput/index.tsx @@ -165,6 +165,10 @@ import type { import { CreationControls } from "./CreationControls"; import { SEND_DISPATCH_MODES } from "./sendDispatchModes"; import { CodexOauthWarningBanner } from "./CodexOauthWarningBanner"; +import { + getCodexOauthProjectPath, + hasCodexOauthTokens, +} from "@/common/utils/providers/codexOauthRouting"; import { useCreationWorkspace } from "./useCreationWorkspace"; import { useCoderWorkspace } from "@/browser/hooks/useCoderWorkspace"; import { useTutorial } from "@/browser/contexts/TutorialContext"; @@ -234,7 +238,7 @@ const ChatInputInner: React.FC = (props) => { [effectivePolicy] ); const { variant } = props; - const { userProjects } = useProjectContext(); + const { userProjects, getProjectConfig } = useProjectContext(); const creationScope = variant === "creation" ? resolveWorkspaceCreationScope(props.projectPath, userProjects, props.pendingSubProjectPath) @@ -491,7 +495,7 @@ const ChatInputInner: React.FC = (props) => { ); const { open } = useSettings(); - const { selectedWorkspace, beginWorkspaceCreation } = useWorkspaceContext(); + const { selectedWorkspace, workspaceMetadata, beginWorkspaceCreation } = useWorkspaceContext(); const { agentId, currentAgent, agents } = useAgent(); // Use current agent's uiColor, or neutral border until agents load @@ -505,7 +509,6 @@ const ChatInputInner: React.FC = (props) => { ensureModelInSettings, defaultModel, setDefaultModel, - codexOauthSet, requiresCodexOauth, } = useModelsFromSettings(); @@ -601,17 +604,46 @@ const ChatInputInner: React.FC = (props) => { const usage = useWorkspaceUsage(workspaceIdForUsage); const { has1MContext } = useProviderOptions(); const { config: providersConfig } = useProvidersConfig(); + const accountProjectPath = getCodexOauthProjectPath( + variant === "creation" + ? { projectPath: creationParentProjectPath, subProjectPath: creationSubProjectPath } + : (workspaceMetadata.get(props.workspaceId) ?? + (selectedWorkspace?.workspaceId === props.workspaceId ? selectedWorkspace : undefined)) + ); + const codexOauthAccountId = accountProjectPath + ? getProjectConfig(accountProjectPath)?.codexOauthAccountId + : undefined; + const codexOauthSet = + providersConfig == null + ? null + : hasCodexOauthTokens(providersConfig.openai, codexOauthAccountId); const lastUsage = usage?.liveUsage ?? usage?.lastContextUsage; // Token counts come from usage metadata, but context limits/1M eligibility should // follow the currently selected model unless a stream is actively running. - const activeUsageModel = usage?.liveUsage?.model ?? null; + const activeUsageModel = usage?.liveModel ?? usage?.liveUsage?.model ?? null; const contextDisplayModel = activeUsageModel ?? baseModel; const use1M = has1MContext(contextDisplayModel); + const liveContextLimit = usage?.liveContextLimit; const contextUsageData = useMemo(() => { - return lastUsage - ? calculateTokenMeterData(lastUsage, contextDisplayModel, use1M, false, providersConfig) - : { segments: [], totalTokens: 0, totalPercentage: 0 }; - }, [lastUsage, contextDisplayModel, use1M, providersConfig]); + return calculateTokenMeterData( + lastUsage, + contextDisplayModel, + use1M, + false, + providersConfig, + { + codexOauthAccountId, + }, + liveContextLimit + ); + }, [ + lastUsage, + contextDisplayModel, + use1M, + providersConfig, + codexOauthAccountId, + liveContextLimit, + ]); const { threshold: autoCompactThreshold, setThreshold: setAutoCompactThreshold } = useAutoCompactionSettings(workspaceIdForUsage, contextDisplayModel); const autoCompactionProps = useMemo( @@ -2765,7 +2797,10 @@ const ChatInputInner: React.FC = (props) => { className="flex shrink-0 items-center" data-component="ThinkingSelectorGroup" > - + diff --git a/src/browser/features/ChatInput/useCreationWorkspace.ts b/src/browser/features/ChatInput/useCreationWorkspace.ts index 05c198db90b..6790697e95c 100644 --- a/src/browser/features/ChatInput/useCreationWorkspace.ts +++ b/src/browser/features/ChatInput/useCreationWorkspace.ts @@ -370,6 +370,7 @@ export function useCreationWorkspace({ debounceMs: 500, userModel, scopeId: workspaceNameScopeId, + projectPath: (subProjectPath ?? projectPath) || undefined, }); // Destructure name state functions for use in callbacks diff --git a/src/browser/features/RightSidebar/ContextUsageBar.tsx b/src/browser/features/RightSidebar/ContextUsageBar.tsx index 51f84b5109d..5e207216a8e 100644 --- a/src/browser/features/RightSidebar/ContextUsageBar.tsx +++ b/src/browser/features/RightSidebar/ContextUsageBar.tsx @@ -12,6 +12,7 @@ interface ContextUsageBarProps { /** Current model ID — used to show 1M context toggle for supported models */ model?: string; showTitle?: boolean; + showEmpty?: boolean; testId?: string; } @@ -20,6 +21,7 @@ const ContextUsageBarComponent: React.FC = ({ autoCompaction, model, showTitle = true, + showEmpty = false, testId, }) => { const totalDisplay = formatTokens(data.totalTokens); @@ -30,7 +32,7 @@ const ContextUsageBarComponent: React.FC = ({ const showThresholdSlider = Boolean(autoCompaction && data.maxTokens); const contextWarning = autoCompaction?.contextWarning; - if (data.totalTokens === 0) return null; + if (data.totalTokens === 0 && !showEmpty) return null; return (
diff --git a/src/browser/features/RightSidebar/ContextUsageSection.tsx b/src/browser/features/RightSidebar/ContextUsageSection.tsx index 68dcaa26703..732bb9aac12 100644 --- a/src/browser/features/RightSidebar/ContextUsageSection.tsx +++ b/src/browser/features/RightSidebar/ContextUsageSection.tsx @@ -1,4 +1,7 @@ import React from "react"; +import { getCodexOauthProjectPath } from "@/common/utils/providers/codexOauthRouting"; +import { useProjectContext } from "@/browser/contexts/ProjectContext"; +import { useWorkspaceContext } from "@/browser/contexts/WorkspaceContext"; import { useWorkspaceUsage } from "@/browser/stores/WorkspaceStore"; import { usePersistedState } from "@/browser/hooks/usePersistedState"; import { AGENT_AI_DEFAULTS_KEY } from "@/common/constants/storage"; @@ -33,10 +36,17 @@ export const ContextUsageSection: React.FC = ({ worksp const { has1MContext } = useProviderOptions(); const pendingSendOptions = useSendMessageOptions(workspaceId); const { config: providersConfig } = useProvidersConfig(); + const { getProjectConfig } = useProjectContext(); + const { workspaceMetadata } = useWorkspaceContext(); + const projectPath = getCodexOauthProjectPath(workspaceMetadata.get(workspaceId)); + const codexOauthAccountId = projectPath + ? getProjectConfig(projectPath)?.codexOauthAccountId + : undefined; // Token counts come from usage metadata, but context limits/1M eligibility should // follow the currently selected model unless a stream is actively running. - const contextDisplayModel = usage.liveUsage?.model ?? pendingSendOptions.baseModel; + const contextDisplayModel = + usage.liveModel ?? usage.liveUsage?.model ?? pendingSendOptions.baseModel; // Align warning with /compact model resolution so it matches actual compaction behavior. const effectiveCompactionModel = resolveCompactionModel(configuredCompactionModel) ?? contextDisplayModel; @@ -46,7 +56,7 @@ export const ContextUsageSection: React.FC = ({ worksp useAutoCompactionSettings(workspaceId, contextDisplayModel); const contextUsage = usage.liveUsage ?? usage.lastContextUsage; - if (!contextUsage) { + if (!contextUsage && !usage.liveModel) { return null; } @@ -55,7 +65,9 @@ export const ContextUsageSection: React.FC = ({ worksp contextDisplayModel, has1MContext(contextDisplayModel), false, - providersConfig + providersConfig, + { codexOauthAccountId }, + usage.liveContextLimit ); // Warn when the compaction model can't fit the auto-compact threshold to avoid failures. @@ -67,7 +79,8 @@ export const ContextUsageSection: React.FC = ({ worksp const compactionMaxTokens = getEffectiveContextLimit( effectiveCompactionModel, has1MContext(effectiveCompactionModel), - providersConfig + providersConfig, + { codexOauthAccountId } ); if (compactionMaxTokens && compactionMaxTokens < thresholdTokens) { @@ -83,6 +96,7 @@ export const ContextUsageSection: React.FC = ({ worksp > [0]; +type Account = NonNullable[number]; +interface LoginFlow { + flowId: string; + url: string; + userCode?: string; + cancel: () => Promise; +} + +const legacyAccounts: Account[] = [{ id: CODEX_OAUTH_DEFAULT_ACCOUNT_ID, label: "Default" }]; +const noAccounts: Account[] = []; + +const inputClassName = + "bg-background border-border-light text-foreground w-full min-w-0 rounded border px-2 py-1.5 text-xs"; + +function accountSelectionLabel(account: Account, accounts: readonly Account[]): string { + const label = formatCodexAccountLabel(account, accounts); + return account.reconnectRequired ? label + " (Reconnect required)" : label; +} + +function AccountSelect(props: { + label: string; + value: string; + accounts: Account[]; + defaultLabel?: string; + selectRef?: Ref; + disabled: boolean; + onChange: (value: string) => void; +}) { + // Keep missing selections visible. Selecting another account must require a user action. + const missing = + props.value !== "" && !props.accounts.some((account) => account.id === props.value); + return ( + + ); +} + +export function CodexAccounts() { + const { api } = useAPI(); + const { config, loading, refresh } = useProvidersConfig(); + const { codexAccountAction, setCodexAccountAction } = useSettings(); + const { userProjects, refreshProjects } = useProjectContext(); + const [label, setLabel] = useState(""); + const [rename, setRename] = useState(null); + const [busy, setBusy] = useState(false); + const [loginInProgress, setLoginInProgress] = useState(false); + const [error, setError] = useState(null); + const [flow, setFlow] = useState(null); + const newNameRef = useRef(null); + const renameRef = useRef(null); + const defaultSelectRef = useRef(null); + const projectSelectRefs = useRef(new Map()); + const consumedActionRef = useRef(null); + const mountedRef = useRef(false); + const attemptRef = useRef(0); + const flowRef = useRef(null); + const openai = config?.openai; + const accounts: Account[] = + openai?.codexOauthAccounts ?? (openai?.codexOauthSet ? legacyAccounts : noAccounts); + const defaultId = openai?.codexOauthDefaultAccountId ?? CODEX_OAUTH_DEFAULT_ACCOUNT_ID; + const defaultAccount = accounts.find((account) => account.id === defaultId); + const defaultLabel = defaultAccount + ? accountSelectionLabel(defaultAccount, accounts) + : `Missing account (${defaultId})`; + const isDesktop = !!window.api; + const showBrowser = + isDesktop || ["localhost", "127.0.0.1", "::1"].includes(window.location.hostname); + const disabled = !api || busy; + // Keep API-key recovery available after the selected OAuth account disconnects. + const authEditable = openai?.apiKeySet === true || !!openai?.apiKeySource; + + // StrictMode replays mount effects before a pending login start can return. + // Keep that start valid. A real unmount cancels its result when it arrives. + useEffect(() => { + mountedRef.current = true; + return () => { + mountedRef.current = false; + if (flowRef.current) { + attemptRef.current += 1; + flowRef.current.cancel().catch(() => undefined); + } + }; + }, []); + + function runAction(operation: Promise): void { + operation.catch((err: unknown) => setError(getErrorMessage(err))); + } + + /* eslint-disable react-hooks/exhaustive-deps -- React Compiler owns callback memoization. Keep the intent effect dependencies explicit. */ + function startRename(account: Account) { + if (rename?.id === account.id) { + renameRef.current?.focus(); + return; + } + setRename(account); + } + + async function saveName() { + if (!api || !rename) return; + const input = { accountId: rename.id, label: rename.label.trim() }; + if (await mutate(() => api.codexOauth.renameAccount(input))) setRename(null); + } + + async function refreshState() { + await Promise.all([refresh(), refreshProjects()]); + } + + async function mutate(operation: () => Promise>) { + setBusy(true); + setError(null); + try { + const result = await operation(); + if (!result.success) { + setError(result.error); + return false; + } + await refreshState(); + return true; + } catch (err) { + setError(getErrorMessage(err)); + return false; + } finally { + setBusy(false); + } + } + + function disconnect(accountId: string) { + if (!api) return; + runAction(mutate(() => api.codexOauth.disconnect({ accountId }))); + } + + async function connect(device: boolean, input: LoginInput) { + if (!api) return; + const attempt = ++attemptRef.current; + const isCurrent = () => mountedRef.current && attempt === attemptRef.current; + setLoginInProgress(true); + setBusy(true); + setError(null); + try { + let nextFlow: LoginFlow; + if (device || !showBrowser) { + const result = await api.codexOauth.startDeviceFlow(input); + if (!result.success) throw new Error(result.error); + const { flowId, userCode, verifyUrl } = result.data; + nextFlow = { + flowId, + userCode, + url: verifyUrl, + cancel: () => api.codexOauth.cancelDeviceFlow({ flowId }), + }; + } else { + const result = await api.codexOauth.startDesktopFlow(input); + if (!result.success) throw new Error(result.error); + const { flowId, authorizeUrl } = result.data; + nextFlow = { + flowId, + url: authorizeUrl, + cancel: () => api.codexOauth.cancelDesktopFlow({ flowId }), + }; + } + if (!isCurrent()) { + await nextFlow.cancel(); + return; + } + flowRef.current = nextFlow; + setFlow(nextFlow); + const result = + nextFlow.userCode != null + ? await api.codexOauth.waitForDeviceFlow({ flowId: nextFlow.flowId }) + : await api.codexOauth.waitForDesktopFlow({ flowId: nextFlow.flowId }); + if (!isCurrent()) return; + if (!result.success) throw new Error(result.error); + setLabel(""); + await refreshState(); + } catch (err) { + if (isCurrent()) setError(getErrorMessage(err)); + } finally { + if (isCurrent()) { + flowRef.current = null; + setFlow(null); + setLoginInProgress(false); + setBusy(false); + } + } + } + + /* eslint-enable react-hooks/exhaustive-deps */ + + async function cancel() { + attemptRef.current++; + try { + await flowRef.current?.cancel(); + setError(null); + } catch (err) { + setError(getErrorMessage(err)); + } finally { + flowRef.current = null; + setFlow(null); + setLoginInProgress(false); + setBusy(false); + } + } + + // Commands must reach the same controls and handlers after Settings mounts. + // Consume each intent before starting work. Busy commands must not run later. + useEffect(() => { + const action = codexAccountAction; + if (!action || loading || consumedActionRef.current === action) return; + consumedActionRef.current = action; + setCodexAccountAction((current) => (current === action ? null : current)); + if (disabled) { + setError(busy ? "An account operation is in progress. Try again." : "API unavailable."); + return; + } + setError(null); + switch (action.type) { + case "add": + newNameRef.current?.focus(); + break; + case "default": + if (accounts.length === 0) { + setError("No Codex accounts are available. Add an account first."); + return; + } + defaultSelectRef.current?.focus(); + break; + case "project": { + const select = projectSelectRefs.current.get(action.projectPath); + if (!select) { + setError("The project is no longer available. Select another project."); + return; + } + select.focus(); + break; + } + default: { + const account = accounts.find((item) => item.id === action.accountId); + if (!account) { + setError("The account is no longer available. Select another account."); + return; + } + if (action.type === "rename") startRename(account); + else if (action.type === "reconnect") runAction(connect(false, { accountId: account.id })); + else disconnect(account.id); + } + } + }, [ + codexAccountAction, + loading, + setCodexAccountAction, + disabled, + busy, + accounts, + startRename, + connect, + disconnect, + ]); + + const loginInput = label.trim() ? { label: label.trim() } : undefined; + return ( +
+
+

ChatGPT (Codex) OAuth

+

+ {accounts.some((account) => account.reconnectRequired) + ? "Reconnect required" + : openai?.codexOauthSet + ? "Connected" + : "Not connected"} +

+
+
    + {accounts.map((account) => ( +
  • + {rename?.id === account.id ? ( +
    { + event.preventDefault(); + runAction(saveName()); + }} + > + setRename({ ...account, label: event.target.value })} + /> + + +
    + ) : ( + <> +

    + {formatCodexAccountLabel(account, accounts)} + {account.id === defaultId && ( + · Global default + )} +

    + {account.reconnectRequired && ( +

    Reconnect required

    + )} +
    + + {showBrowser && ( + + )} + + +
    + + )} +
    • + ))} +
+
{ + event.preventDefault(); + runAction(connect(false, loginInput)); + }} + > + +
+ {showBrowser && ( + + )} + +
+
+ {busy && ( +
+ + + {flow ? "Waiting for authorization..." : loginInProgress ? "Starting..." : "Saving..."} + + {loginInProgress && ( + + )} +
+ )} + {flow && ( +
+ {flow.userCode && ( + <> +

Enter this code on the OpenAI verification page:

+ + {flow.userCode} + + + )} + +
+ )} + {error && ( +

+ {error} +

+ )} + + api && runAction(mutate(() => api.codexOauth.setDefaultAccount({ accountId }))) + } + /> +
+

Project accounts

+

+ Project selections override the global default. Missing accounts do not use another + account. +

+ {Array.from(userProjects, ([projectPath, project]) => ( + { + if (select) projectSelectRefs.current.set(projectPath, select); + else projectSelectRefs.current.delete(projectPath); + }} + label={project.displayName ?? projectPath} + value={project.codexOauthAccountId ?? ""} + accounts={accounts} + defaultLabel={defaultLabel} + disabled={disabled} + onChange={(accountId) => + api && + runAction( + mutate(() => + api.projects.setCodexOauthAccount({ projectPath, accountId: accountId || null }) + ) + ) + } + /> + ))} +
+ +

+ ChatGPT OAuth costs use API-equivalent estimates. Your plan may include usage or charge + credits. API keys use OpenAI platform billing. +

+ {!authEditable && ( +

Set an OpenAI API key to change this setting.

+ )} +
+ ); +} diff --git a/src/browser/features/Settings/Sections/ProvidersSection.test.tsx b/src/browser/features/Settings/Sections/ProvidersSection.test.tsx index 5e43ef22147..4d2100cb2c6 100644 --- a/src/browser/features/Settings/Sections/ProvidersSection.test.tsx +++ b/src/browser/features/Settings/Sections/ProvidersSection.test.tsx @@ -1,10 +1,12 @@ import type React from "react"; +import { StrictMode } from "react"; import { cleanup, fireEvent, render, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, mock, spyOn, test } from "bun:test"; import { installDom } from "../../../../../tests/ui/dom"; import { createSelectPrimitiveDouble } from "../../../../../tests/ui/selectPrimitiveDouble"; import type { APIClient } from "@/browser/contexts/API"; +import { Err, Ok } from "@/common/types/result"; import * as ActualSelectPrimitiveModule from "@/browser/components/SelectPrimitive/SelectPrimitive"; import * as SettingsContextModule from "@/browser/contexts/SettingsContext"; import type * as WorkspaceStoreModule from "@/browser/stores/WorkspaceStore"; @@ -46,6 +48,7 @@ void mock.module("@/browser/utils/modelPreferenceRepair", () => ({ })); let providersConfigMock: ProvidersConfigMap | null = null; +let codexPolicyBlocked = false; let apiMock: APIClient | null = null; const providersRefreshMock = mock(() => Promise.resolve()); const updateOptimisticallyMock = mock((provider: string, updates: Partial) => { @@ -82,8 +85,8 @@ void mock.module("@/browser/contexts/API", () => ({ void mock.module("@/browser/contexts/PolicyContext", () => ({ usePolicy: () => ({ - status: { state: "disabled" as const }, - policy: null, + status: { state: codexPolicyBlocked ? "enforced" : "disabled" }, + policy: codexPolicyBlocked ? { providerAccess: [] } : null, }), })); @@ -209,6 +212,61 @@ function renderProvidersSection() { return { ...view, ...providerMocks, providersConfig }; } +function CodexIntentHarness(props: { action: SettingsContextModule.CodexAccountSettingsIntent }) { + const settings = SettingsContextModule.useSettings(); + return ( + <> + + + {settings.codexAccountAction?.type ?? "none"} + + + + ); +} + +function setupCodexIntent(action: SettingsContextModule.CodexAccountSettingsIntent) { + providersConfigMock = createProvidersConfig(); + providersConfigMock.openai.codexOauthAccounts = [{ id: "work", label: "Work" }]; + const client = setupSettingsStory({}); + apiMock = client; + const startDesktopFlow = mock(() => + Promise.resolve(Err("Unexpected login")) + ); + const startDeviceFlow = mock(() => + Promise.resolve(Err("Unexpected login")) + ); + const disconnect = mock(() => + Promise.resolve(Ok(undefined)) + ); + client.codexOauth = { + startDesktopFlow, + startDeviceFlow, + disconnect, + waitForDesktopFlow: () => Promise.resolve(Ok(undefined)), + waitForDeviceFlow: () => Promise.resolve(Ok(undefined)), + cancelDesktopFlow: () => Promise.resolve(), + cancelDeviceFlow: () => Promise.resolve(), + renameAccount: () => Promise.resolve(Ok(undefined)), + setDefaultAccount: () => Promise.resolve(Ok(undefined)), + }; + const renderContent = () => ( + client}> + + + + + ); + const view = render(renderContent()); + return { view, client, startDesktopFlow, startDeviceFlow, disconnect, renderContent }; +} + function getProviderCard(button: HTMLElement): HTMLElement { const card = button.parentElement; if (!card) { @@ -232,6 +290,7 @@ describe("ProvidersSection", () => { ); providersConfigMock = null; apiMock = null; + codexPolicyBlocked = false; providersRefreshMock.mockClear(); updateOptimisticallyMock.mockClear(); }); @@ -251,6 +310,131 @@ describe("ProvidersSection", () => { restoreDom = null; }); + test.each(["reconnect", "rename", "disconnect"] as const)( + "consumes a stale Codex %s command without a mutation", + async (type) => { + const { view, startDesktopFlow, startDeviceFlow, disconnect } = setupCodexIntent({ + type, + accountId: "deleted", + }); + fireEvent.click(view.getByRole("button", { name: "Run account command" })); + await waitFor(() => + expect(view.getByRole("alert").textContent).toContain("no longer available") + ); + expect(view.getByTestId("pending-account-command").textContent).toBe("none"); + expect(view.queryByRole("textbox", { name: "Account name" })).toBeNull(); + expect(startDesktopFlow).not.toHaveBeenCalled(); + expect(startDeviceFlow).not.toHaveBeenCalled(); + expect(disconnect).not.toHaveBeenCalled(); + } + ); + + test("consumes a stale project command without selecting another project", async () => { + const { view } = setupCodexIntent({ type: "project", projectPath: "/deleted" }); + fireEvent.click(view.getByRole("button", { name: "Run account command" })); + await waitFor(() => + expect(view.getByRole("alert").textContent).toContain("no longer available") + ); + expect(view.getByTestId("pending-account-command").textContent).toBe("none"); + }); + + test("does not replay a Codex command after policy permits OpenAI", async () => { + codexPolicyBlocked = true; + const { view, startDesktopFlow, startDeviceFlow, renderContent } = setupCodexIntent({ + type: "reconnect", + accountId: "work", + }); + fireEvent.click(view.getByRole("button", { name: "Run account command" })); + await waitFor(() => + expect(view.getByTestId("pending-account-command").textContent).toBe("none") + ); + expect(view.queryByRole("region", { name: "ChatGPT (Codex) accounts" })).toBeNull(); + codexPolicyBlocked = false; + view.rerender(renderContent()); + await view.findByRole("region", { name: "ChatGPT (Codex) accounts" }); + expect(startDesktopFlow).not.toHaveBeenCalled(); + expect(startDeviceFlow).not.toHaveBeenCalled(); + }); + + test("does not replay a Codex command after a busy login ends", async () => { + const { view, client, startDesktopFlow, startDeviceFlow } = setupCodexIntent({ + type: "reconnect", + accountId: "work", + }); + startDesktopFlow.mockResolvedValue({ + success: true, + data: { flowId: "busy", authorizeUrl: "https://auth.openai.com/authorize" }, + }); + startDeviceFlow.mockResolvedValue({ + success: true, + data: { + flowId: "busy", + userCode: "CODE", + verifyUrl: "https://auth.openai.com/codex/device", + intervalSeconds: 5, + }, + }); + let finishLogin: () => void = () => undefined; + const waitForLogin = () => + new Promise<{ success: true; data: undefined }>((resolve) => { + finishLogin = () => resolve({ success: true, data: undefined }); + }); + client.codexOauth.waitForDesktopFlow = waitForLogin; + client.codexOauth.waitForDeviceFlow = waitForLogin; + client.codexOauth.cancelDesktopFlow = () => { + finishLogin(); + return Promise.resolve(); + }; + client.codexOauth.cancelDeviceFlow = client.codexOauth.cancelDesktopFlow; + await userEvent.click(view.getByRole("button", { name: "Run account command" })); + await view.findByText("Waiting for authorization..."); + await userEvent.click(view.getByRole("button", { name: "Run account command" })); + await waitFor(() => expect(view.getByRole("alert").textContent).toContain("in progress")); + expect(view.getByTestId("pending-account-command").textContent).toBe("none"); + await userEvent.click(view.getByRole("button", { name: "Cancel" })); + await waitFor(() => expect(view.queryByText("Waiting for authorization...")).toBeNull()); + expect(startDesktopFlow.mock.calls.length + startDeviceFlow.mock.calls.length).toBe(1); + await userEvent.click(view.getByRole("button", { name: "Run account command" })); + await waitFor(() => + expect(startDesktopFlow.mock.calls.length + startDeviceFlow.mock.calls.length).toBe(2) + ); + await userEvent.click(view.getByRole("button", { name: "Cancel" })); + }); + + test("cancels a late command login result after a real unmount", async () => { + const { view, client, startDesktopFlow, startDeviceFlow } = setupCodexIntent({ + type: "reconnect", + accountId: "work", + }); + const data = { + flowId: "late", + authorizeUrl: "https://auth.openai.com/authorize", + userCode: "CODE", + verifyUrl: "https://auth.openai.com/codex/device", + intervalSeconds: 5, + }; + let finishStart: () => void = () => undefined; + const pendingStart = new Promise<{ success: true; data: typeof data }>((resolve) => { + finishStart = () => resolve({ success: true, data }); + }); + startDesktopFlow.mockReturnValue(pendingStart); + startDeviceFlow.mockReturnValue(pendingStart); + const cancelDesktop = spyOn(client.codexOauth, "cancelDesktopFlow"); + const cancelDevice = spyOn(client.codexOauth, "cancelDeviceFlow"); + const waitDesktop = spyOn(client.codexOauth, "waitForDesktopFlow"); + const waitDevice = spyOn(client.codexOauth, "waitForDeviceFlow"); + await userEvent.click(view.getByRole("button", { name: "Run account command" })); + expect(startDesktopFlow.mock.calls.length + startDeviceFlow.mock.calls.length).toBe(1); + view.unmount(); + finishStart(); + await waitFor(() => + expect(cancelDesktop.mock.calls.length + cancelDevice.mock.calls.length).toBe(1) + ); + expect(waitDesktop).not.toHaveBeenCalled(); + expect(waitDevice).not.toHaveBeenCalled(); + expect(providersRefreshMock).not.toHaveBeenCalled(); + }); + test("renders built-in and custom providers in separate groups", async () => { const view = renderProvidersSection(); @@ -709,6 +893,8 @@ describe("ProvidersSection", () => { providersExpandedProvider: opts.hint ? null : "coder", setProvidersExpandedProvider: () => undefined, providersStartCoderLogin: startCoderLoginHint, + codexAccountAction: null, + setCodexAccountAction: () => undefined, setProvidersStartCoderLogin, runtimesProjectPath: null, setRuntimesProjectPath: () => undefined, diff --git a/src/browser/features/Settings/Sections/ProvidersSection.tsx b/src/browser/features/Settings/Sections/ProvidersSection.tsx index 75ea804c4af..1c4cfa79d8d 100644 --- a/src/browser/features/Settings/Sections/ProvidersSection.tsx +++ b/src/browser/features/Settings/Sections/ProvidersSection.tsx @@ -52,10 +52,6 @@ import { SelectValue, } from "@/browser/components/SelectPrimitive/SelectPrimitive"; import { Switch } from "@/browser/components/Switch/Switch"; -import { - ToggleGroup, - ToggleGroupItem, -} from "@/browser/components/ToggleGroupPrimitive/ToggleGroupPrimitive"; import { HelpIndicator, Tooltip, @@ -80,6 +76,7 @@ import type { AddCustomProviderInput, ProviderConfigInfo } from "@/common/orpc/t import type { ServiceTier, XAIServiceTier } from "@/common/config/schemas/providersConfig"; import type { Result } from "@/common/types/result"; import { CODER_OAUTH_SERVER_START_PATH } from "@/common/constants/coderOAuth"; +import { CodexAccounts } from "./CodexAccounts"; type MuxGatewayLoginStatus = "idle" | "starting" | "waiting" | "success" | "error"; type CodexOauthFlowStatus = "idle" | "starting" | "waiting" | "error"; @@ -104,12 +101,6 @@ function isXAIServiceTier(value: string): value is XAIServiceTier { return value === "default" || value === "priority"; } -interface CodexOauthDeviceFlow { - flowId: string; - userCode: string; - verifyUrl: string; -} - interface OAuthMessage { type?: unknown; state?: unknown; @@ -458,6 +449,8 @@ export function ProvidersSection() { const { providersExpandedProvider, setProvidersExpandedProvider, + codexAccountAction, + setCodexAccountAction, providersStartCoderLogin, setProvidersStartCoderLogin, } = useSettings(); @@ -469,6 +462,17 @@ export function ProvidersSection() { () => getAllowedProvidersForUi(effectivePolicy, config), [effectivePolicy, config] ); + // Hidden providers never mount their account controls. Do not retain commands for a later policy change. + useEffect(() => { + if ( + codexAccountAction && + !configLoading && + (!visibleProviders.includes("openai") || isCustomProviderInfo(config?.openai)) + ) { + setCodexAccountAction((current) => (current === codexAccountAction ? null : current)); + } + }, [codexAccountAction, configLoading, visibleProviders, config, setCodexAccountAction]); + const { data: muxGatewayAccountStatus, error: muxGatewayAccountError, @@ -549,290 +553,9 @@ export function ProvidersSection() { const isDesktop = !!window.api; - // The "Connect (Browser)" OAuth flow requires a redirect back to this origin, - // which only works when the host is the user's local machine. On a remote mux - // server the redirect would land on the server, not the user's browser. - const isRemoteServer = - !isDesktop && !["localhost", "127.0.0.1", "::1"].includes(window.location.hostname); - - const [codexOauthStatus, setCodexOauthStatus] = useState("idle"); - const [codexOauthError, setCodexOauthError] = useState(null); - - const codexOauthAttemptRef = useRef(0); - const [codexOauthDesktopFlowId, setCodexOauthDesktopFlowId] = useState(null); - const [codexOauthDeviceFlow, setCodexOauthDeviceFlow] = useState( - null - ); - const [codexOauthAuthorizeUrl, setCodexOauthAuthorizeUrl] = useState(null); - const codexOauthIsConnected = config?.openai?.codexOauthSet === true; - const openaiApiKeySet = config?.openai?.apiKeySet === true || !!config?.openai?.apiKeySource; - const codexOauthDefaultAuth = - config?.openai?.codexOauthDefaultAuth === "apiKey" ? "apiKey" : "oauth"; - const codexOauthDefaultAuthIsEditable = codexOauthIsConnected && openaiApiKeySet; - - const codexOauthLoginInProgress = - codexOauthStatus === "starting" || codexOauthStatus === "waiting"; - - const startCodexOauthBrowserConnect = async () => { - const attempt = ++codexOauthAttemptRef.current; - - if (!api) { - setCodexOauthStatus("error"); - setCodexOauthError("Xum API not connected."); - return; - } - - // Best-effort: cancel any in-progress flow before starting a new one. - if (codexOauthDesktopFlowId) { - void api.codexOauth.cancelDesktopFlow({ flowId: codexOauthDesktopFlowId }); - } - if (codexOauthDeviceFlow) { - void api.codexOauth.cancelDeviceFlow({ flowId: codexOauthDeviceFlow.flowId }); - } - - setCodexOauthError(null); - setCodexOauthDesktopFlowId(null); - setCodexOauthDeviceFlow(null); - setCodexOauthAuthorizeUrl(null); - - try { - setCodexOauthStatus("starting"); - - if (!isDesktop) { - const startResult = await api.codexOauth.startDeviceFlow(); - - if (attempt !== codexOauthAttemptRef.current) { - if (startResult.success) { - void api.codexOauth.cancelDeviceFlow({ flowId: startResult.data.flowId }); - } - return; - } - - if (!startResult.success) { - setCodexOauthStatus("error"); - setCodexOauthError(startResult.error); - return; - } - - setCodexOauthDeviceFlow({ - flowId: startResult.data.flowId, - userCode: startResult.data.userCode, - verifyUrl: startResult.data.verifyUrl, - }); - setCodexOauthStatus("waiting"); - - // Keep device-code login manual per user request: we only open the - // verification page from the explicit "Copy & Open" action. - const waitResult = await api.codexOauth.waitForDeviceFlow({ - flowId: startResult.data.flowId, - }); - - if (attempt !== codexOauthAttemptRef.current) { - return; - } - - if (!waitResult.success) { - setCodexOauthStatus("error"); - setCodexOauthError(waitResult.error); - return; - } - - setCodexOauthStatus("idle"); - setCodexOauthDeviceFlow(null); - setCodexOauthAuthorizeUrl(null); - await refresh(); - return; - } - - const startResult = await api.codexOauth.startDesktopFlow(); - - if (attempt !== codexOauthAttemptRef.current) { - if (startResult.success) { - void api.codexOauth.cancelDesktopFlow({ flowId: startResult.data.flowId }); - } - return; - } - - if (!startResult.success) { - setCodexOauthStatus("error"); - setCodexOauthError(startResult.error); - return; - } - - const { flowId, authorizeUrl } = startResult.data; - setCodexOauthDesktopFlowId(flowId); - setCodexOauthAuthorizeUrl(authorizeUrl); - setCodexOauthStatus("waiting"); - - const waitResult = await api.codexOauth.waitForDesktopFlow({ flowId }); - - if (attempt !== codexOauthAttemptRef.current) { - return; - } - - if (!waitResult.success) { - setCodexOauthStatus("error"); - setCodexOauthError(waitResult.error); - return; - } - - setCodexOauthStatus("idle"); - setCodexOauthDesktopFlowId(null); - await refresh(); - } catch (err) { - if (attempt !== codexOauthAttemptRef.current) { - return; - } - - setCodexOauthStatus("error"); - setCodexOauthError(getErrorMessage(err)); - } - }; - - const startCodexOauthDeviceConnect = async () => { - const attempt = ++codexOauthAttemptRef.current; - - if (!api) { - setCodexOauthStatus("error"); - setCodexOauthError("Xum API not connected."); - return; - } - - // Best-effort: cancel any in-progress flow before starting a new one. - if (codexOauthDesktopFlowId) { - void api.codexOauth.cancelDesktopFlow({ flowId: codexOauthDesktopFlowId }); - } - if (codexOauthDeviceFlow) { - void api.codexOauth.cancelDeviceFlow({ flowId: codexOauthDeviceFlow.flowId }); - } - - setCodexOauthError(null); - setCodexOauthDesktopFlowId(null); - setCodexOauthDeviceFlow(null); - setCodexOauthAuthorizeUrl(null); - - try { - setCodexOauthStatus("starting"); - const startResult = await api.codexOauth.startDeviceFlow(); - - if (attempt !== codexOauthAttemptRef.current) { - if (startResult.success) { - void api.codexOauth.cancelDeviceFlow({ flowId: startResult.data.flowId }); - } - return; - } - - if (!startResult.success) { - setCodexOauthStatus("error"); - setCodexOauthError(startResult.error); - return; - } - - setCodexOauthDeviceFlow({ - flowId: startResult.data.flowId, - userCode: startResult.data.userCode, - verifyUrl: startResult.data.verifyUrl, - }); - setCodexOauthStatus("waiting"); - - const waitResult = await api.codexOauth.waitForDeviceFlow({ - flowId: startResult.data.flowId, - }); - - if (attempt !== codexOauthAttemptRef.current) { - return; - } - - if (!waitResult.success) { - setCodexOauthStatus("error"); - setCodexOauthError(waitResult.error); - return; - } - - setCodexOauthStatus("idle"); - setCodexOauthDeviceFlow(null); - setCodexOauthAuthorizeUrl(null); - await refresh(); - } catch (err) { - if (attempt !== codexOauthAttemptRef.current) { - return; - } - - setCodexOauthStatus("error"); - setCodexOauthError(getErrorMessage(err)); - } - }; - - const disconnectCodexOauth = async () => { - const attempt = ++codexOauthAttemptRef.current; - - if (!api) { - setCodexOauthStatus("error"); - setCodexOauthError("Xum API not connected."); - return; - } - - // Best-effort: cancel any in-progress flow. - if (codexOauthDesktopFlowId) { - void api.codexOauth.cancelDesktopFlow({ flowId: codexOauthDesktopFlowId }); - } - if (codexOauthDeviceFlow) { - void api.codexOauth.cancelDeviceFlow({ flowId: codexOauthDeviceFlow.flowId }); - } - - setCodexOauthError(null); - setCodexOauthDesktopFlowId(null); - setCodexOauthDeviceFlow(null); - setCodexOauthAuthorizeUrl(null); - - try { - setCodexOauthStatus("starting"); - const result = await api.codexOauth.disconnect(); - - if (attempt !== codexOauthAttemptRef.current) { - return; - } - - if (!result.success) { - setCodexOauthStatus("error"); - setCodexOauthError(result.error); - return; - } - - updateOptimistically("openai", { codexOauthSet: false }); - setCodexOauthStatus("idle"); - await refresh(); - } catch (err) { - if (attempt !== codexOauthAttemptRef.current) { - return; - } - - setCodexOauthStatus("error"); - setCodexOauthError(getErrorMessage(err)); - } - }; const [muxGatewayLoginStatus, setMuxGatewayLoginStatus] = useState("idle"); - const cancelCodexOauth = () => { - codexOauthAttemptRef.current++; - - if (api) { - if (codexOauthDesktopFlowId) { - void api.codexOauth.cancelDesktopFlow({ flowId: codexOauthDesktopFlowId }); - } - if (codexOauthDeviceFlow) { - void api.codexOauth.cancelDeviceFlow({ flowId: codexOauthDeviceFlow.flowId }); - } - } - - setCodexOauthDesktopFlowId(null); - setCodexOauthDeviceFlow(null); - setCodexOauthAuthorizeUrl(null); - setCodexOauthStatus("idle"); - setCodexOauthError(null); - }; - const [muxGatewayLoginError, setMuxGatewayLoginError] = useState(null); const muxGatewayLoginAttemptRef = useRef(0); @@ -2651,185 +2374,7 @@ export function ProvidersSection() { const openAIWebSocketTransportVisible = openAIWireFormat === "responses"; return (
-
- - - {codexOauthStatus === "starting" - ? "Starting..." - : codexOauthStatus === "waiting" - ? "Waiting for login..." - : codexOauthIsConnected - ? "Connected" - : "Not connected"} - -
- -
- {!isRemoteServer && ( - - )} - - - {codexOauthStatus === "waiting" && - !codexOauthDeviceFlow && - codexOauthAuthorizeUrl && ( - - )} - - {codexOauthLoginInProgress && ( - - )} - - {codexOauthIsConnected && ( - - )} -
- - {codexOauthDeviceFlow && ( -
-

- Enter this code on the OpenAI verification page: -

-
- - {codexOauthDeviceFlow.userCode} - - -
-

- - Waiting for authorization... -

-
- )} - - {codexOauthStatus === "waiting" && !codexOauthDeviceFlow && ( -

- - Waiting for authorization... -

- )} - - {codexOauthStatus === "error" && codexOauthError && ( -

{codexOauthError}

- )} - -
-
- -

- Applies to models that support both ChatGPT OAuth and API keys - (e.g. gpt-5.5). -

-
- - { - if (!api) return; - if (next !== "oauth" && next !== "apiKey") { - return; - } - - updateOptimistically("openai", { codexOauthDefaultAuth: next }); - void api.providers.setProviderConfig({ - provider: "openai", - keyPath: ["codexOauthDefaultAuth"], - value: next, - }); - }} - size="sm" - className="h-9" - disabled={!api || !codexOauthDefaultAuthIsEditable} - > - - Use ChatGPT OAuth by default - - - Use OpenAI API key by default - - - -

- ChatGPT OAuth costs use API-equivalent estimates. Your plan may - include usage or charge credits. API keys use OpenAI platform - billing. -

- - {!codexOauthDefaultAuthIsEditable && ( -

- Connect ChatGPT OAuth and set an OpenAI API key to change this - setting. -

- )} -
+
diff --git a/src/browser/hooks/useContextSwitchWarning.ts b/src/browser/hooks/useContextSwitchWarning.ts index a7c48cf73db..0accc42f0d1 100644 --- a/src/browser/hooks/useContextSwitchWarning.ts +++ b/src/browser/hooks/useContextSwitchWarning.ts @@ -37,6 +37,7 @@ interface UseContextSwitchWarningProps { api: RouterClient | undefined; pendingSendOptions: SendMessageOptions; providersConfig: ProvidersConfigMap | null; + codexOauthAccountId?: string; } interface UseContextSwitchWarningResult { @@ -148,11 +149,12 @@ export function useContextSwitchWarning( const checkOptions: ContextSwitchOptions = useMemo( () => ({ providersConfig, + codexOauthAccountId: props.codexOauthAccountId, policy: effectivePolicy, routePriority, routeOverrides, }), - [providersConfig, effectivePolicy, routePriority, routeOverrides] + [providersConfig, props.codexOauthAccountId, effectivePolicy, routePriority, routeOverrides] ); const prevCheckOptionsRef = useRef(checkOptions); @@ -197,14 +199,19 @@ export function useContextSwitchWarning( }); if (suggestion) { - const limit = getEffectiveContextLimit(suggestion.modelId, use1M, providersConfig); + const limit = getEffectiveContextLimit( + suggestion.modelId, + use1M, + providersConfig, + checkOptions + ); if (limit && limit > w.currentTokens) { return { ...w, compactionModel: suggestion.modelId, errorMessage: null }; } } return w; }, - [providersConfig, effectivePolicy, routePriority, routeOverrides, use1M] + [providersConfig, effectivePolicy, routePriority, routeOverrides, use1M, checkOptions] ); const evaluateWarning = useCallback( @@ -413,8 +420,18 @@ export function useContextSwitchWarning( // OFF → ON: may clear warning if context now fits // ON → OFF: may show warning if context no longer fits if (wasEnabled !== use1M) { - const previousLimit = getEffectiveContextLimit(pendingModel, wasEnabled, providersConfig); - const nextLimit = getEffectiveContextLimit(pendingModel, use1M, providersConfig); + const previousLimit = getEffectiveContextLimit( + pendingModel, + wasEnabled, + providersConfig, + checkOptions + ); + const nextLimit = getEffectiveContextLimit( + pendingModel, + use1M, + providersConfig, + checkOptions + ); // Only surface same-model warnings if the effective limit actually changed. if (previousLimit === nextLimit) { @@ -441,7 +458,16 @@ export function useContextSwitchWarning( dispatch({ type: "CLEAR_WARNING" }); } } - }, [use1M, pendingModel, tokens, messages, providersConfig, evaluateWarning, warning]); + }, [ + use1M, + pendingModel, + tokens, + messages, + providersConfig, + evaluateWarning, + warning, + checkOptions, + ]); return { warning, handleModelChange, handleCompact, handleDismiss }; } diff --git a/src/browser/hooks/useWorkspaceName.ts b/src/browser/hooks/useWorkspaceName.ts index 6ebef4260d7..b577496009c 100644 --- a/src/browser/hooks/useWorkspaceName.ts +++ b/src/browser/hooks/useWorkspaceName.ts @@ -49,6 +49,8 @@ export interface UseWorkspaceNameOptions { debounceMs?: number; /** User's selected model to try after preferred models */ userModel?: string; + /** Apply the project account before the workspace exists. */ + projectPath?: string; /** * Optional storage scope for persisting draft name-generation state. * @@ -142,7 +144,7 @@ export function getDisplayTitleFromPersistedState(state: unknown): string { * auto-generation resumes. */ export function useWorkspaceName(options: UseWorkspaceNameOptions): UseWorkspaceNameReturn { - const { message, debounceMs = 500, userModel, scopeId } = options; + const { message, debounceMs = 500, userModel, scopeId, projectPath } = options; const { api } = useAPI(); const candidates = useMemo(() => buildNameGenCandidates(userModel), [userModel]); @@ -235,6 +237,7 @@ export function useWorkspaceName(options: UseWorkspaceNameOptions): UseWorkspace const result = await api.nameGeneration.generate({ message: forMessage, candidates, + projectPath, }); // Check if this request is still current (wasn't cancelled) @@ -288,7 +291,7 @@ export function useWorkspaceName(options: UseWorkspaceNameOptions): UseWorkspace } } }, - [api, setStored, candidates] + [api, setStored, candidates, projectPath] ); // Debounced generation effect diff --git a/src/browser/stores/WorkspaceStore.test.ts b/src/browser/stores/WorkspaceStore.test.ts index f49bed623c1..46cb1e3d674 100644 --- a/src/browser/stores/WorkspaceStore.test.ts +++ b/src/browser/stores/WorkspaceStore.test.ts @@ -1313,6 +1313,62 @@ describe("WorkspaceStore", () => { }); describe("live usage identity pinning", () => { + it("keeps the limit only for live usage and replaces it on the next turn", async () => { + const workspaceId = "live-context-limit"; + createAndAddWorkspace(store, workspaceId); + await tick(10); + const aggregator = store.getAggregator(workspaceId); + if (!aggregator) throw new Error("Missing live-context aggregator"); + const bump = () => + getInternal<{ usageStore: { bump: (id: string) => void } }>(store).usageStore.bump( + workspaceId + ); + const usage = { inputTokens: 1000, outputTokens: 100, totalTokens: 1100 }; + for (const [index, effectiveContextLimit] of [272_000, 500_000, null].entries()) { + const messageId = "live-limit-" + index; + aggregator.handleStreamStart({ + type: "stream-start", + workspaceId, + messageId, + model: "openai:gpt-5.5", + historySequence: index + 1, + startTime: 1000 + index, + effectiveContextLimit, + }); + bump(); + const starting = store.getWorkspaceUsage(workspaceId); + expect(starting.liveUsage).toBeUndefined(); + expect(starting.liveModel).toBe("openai:gpt-5.5"); + expect(starting.liveContextLimit).toBe(effectiveContextLimit); + aggregator.handleUsageDelta({ + type: "usage-delta", + workspaceId, + messageId, + usage, + cumulativeUsage: usage, + effectiveContextLimit, + }); + bump(); + expect(store.getWorkspaceUsage(workspaceId).liveUsage?.effectiveContextLimit).toBe( + effectiveContextLimit + ); + aggregator.handleStreamEnd({ + type: "stream-end", + workspaceId, + messageId, + metadata: { model: "openai:gpt-5.5", usage, contextUsage: usage }, + parts: [{ type: "text", text: "Done" }], + }); + bump(); + const idle = store.getWorkspaceUsage(workspaceId); + expect(idle.liveUsage).toBeUndefined(); + expect(idle.liveModel).toBeUndefined(); + expect(idle.liveContextLimit).toBeUndefined(); + expect(idle.lastContextUsage?.input.tokens).toBe(1000); + expect(idle.lastContextUsage?.effectiveContextLimit).toBeUndefined(); + } + }); + it("prices live Coder usage via the stream's pinned metadataModel", async () => { const workspaceId = "live-coder-usage-pinned"; createAndAddWorkspace(store, workspaceId); diff --git a/src/browser/stores/WorkspaceStore.ts b/src/browser/stores/WorkspaceStore.ts index 41983492018..b33a9f72e74 100644 --- a/src/browser/stores/WorkspaceStore.ts +++ b/src/browser/stores/WorkspaceStore.ts @@ -318,6 +318,9 @@ export interface WorkspaceUsageState { totalTokens: number; /** Live context usage during streaming (last step's inputTokens = current context window) */ liveUsage?: ChatUsageDisplay; + /** Active request metadata exists before the first usage event. */ + liveModel?: string; + liveContextLimit?: number | null; /** Live cost usage during streaming (cumulative across all steps) */ liveCostUsage?: ChatUsageDisplay; /** @@ -910,6 +913,12 @@ export class WorkspaceStore { // fresh recompute regardless of how the previous one was wound down. this.streamingStatsStore.bump(workspaceId); }, + "stream-model-update": (workspaceId, aggregator, data) => { + applyWorkspaceChatEventToAggregator(aggregator, data); + this.states.bump(workspaceId); + this.usageStore.bump(workspaceId); + this.streamingStatsStore.bump(workspaceId); + }, "stream-lifecycle": (workspaceId, aggregator, data) => { applyWorkspaceChatEventToAggregator(aggregator, data); this.states.bump(workspaceId); @@ -2871,6 +2880,10 @@ export class WorkspaceStore { // Live streaming data (unchanged) const activeStreamId = aggregator.getActiveStreamMessageId(); + const liveModel = activeStreamId ? model : undefined; + const liveContextLimit = activeStreamId + ? aggregator.getActiveStreamContextLimit(activeStreamId) + : undefined; // Request-pinned identity stamped by the backend at stream start: a // Coder catalog refresh can remove/retag the instance mid-stream, and // re-resolving the raw model against the refreshed config would price @@ -2892,6 +2905,10 @@ export class WorkspaceStore { ) : undefined; + if (liveUsage) { + liveUsage.effectiveContextLimit = liveContextLimit; + } + const rawCumulativeUsage = activeStreamId ? aggregator.getActiveStreamCumulativeUsage(activeStreamId) : undefined; @@ -2915,6 +2932,8 @@ export class WorkspaceStore { lastContextUsage, totalTokens, liveUsage, + liveModel, + liveContextLimit, liveCostUsage, liveMetadataModel, }; diff --git a/src/browser/stories/App.codexAccounts.stories.tsx b/src/browser/stories/App.codexAccounts.stories.tsx new file mode 100644 index 00000000000..7c0ec485f7d --- /dev/null +++ b/src/browser/stories/App.codexAccounts.stories.tsx @@ -0,0 +1,948 @@ +import { expect, fn, userEvent, waitFor, within } from "@storybook/test"; +import { StrictMode } from "react"; +import { appMeta, AppWithMocks, type AppStory } from "./meta"; +import { + collapseLeftSidebar, + expandRightSidebar, + expandLeftSidebar, + expandProjects, + selectWorkspace, +} from "./helpers/uiState"; +import { createMockORPCClient } from "./mocks/orpc"; +import { createWorkspace, groupWorkspacesByProject } from "./mocks/workspaces"; +import type { APIClient } from "@/browser/contexts/API"; +import type { ProvidersConfigMap, WorkspaceChatMessage } from "@/common/orpc/types"; +import { updatePersistedState } from "@/browser/hooks/usePersistedState"; +import { getModelKey, RIGHT_SIDEBAR_TAB_KEY } from "@/common/constants/storage"; +import { Err, Ok } from "@/common/types/result"; +import { MULTI_PROJECT_CONFIG_KEY } from "@/common/constants/multiProject"; + +export default { ...appMeta, title: "App/CodexAccounts" }; + +const startLogin = fn<(input: Parameters[0]) => void>(); + +const browserLogin = + fn<(input: Parameters[0]) => void>(); +const generateTitle = fn<(input: Parameters[0]) => void>(); + +function setupAccounts( + revokedWork = false, + workLabel = "Work", + onChat?: (workspaceId: string, emit: (event: WorkspaceChatMessage) => void) => void, + workspaceId = "codex-accounts" +) { + expandLeftSidebar(); + startLogin.mockClear(); + browserLogin.mockClear(); + const workspace = createWorkspace({ + id: workspaceId, + name: "main", + projectName: "my-app", + projectPath: "/projects/my-app", + }); + selectWorkspace(workspace); + const projects = groupWorkspacesByProject([workspace]); + const project = projects.get(workspace.projectPath); + if (project && revokedWork) project.codexOauthAccountId = "work"; + const providers: ProvidersConfigMap = { + openai: { + apiKeySet: true, + isEnabled: true, + isConfigured: true, + codexOauthSet: true, + codexOauthDefaultAccountId: revokedWork ? "work" : undefined, + codexOauthAccounts: [ + { id: "default", label: "Personal" }, + { id: "work", label: workLabel, reconnectRequired: revokedWork || undefined }, + ], + }, + }; + let slot = 0; + let reconnectAccountId: string | undefined; + const client = createMockORPCClient({ + projects, + workspaces: [workspace], + providersConfig: providers, + providersList: ["openai"], + onChat, + }); + const start: APIClient["codexOauth"]["startDeviceFlow"] = (input) => { + startLogin(input); + reconnectAccountId = input?.accountId; + if (input?.label) { + providers.openai.codexOauthAccounts?.push({ id: "slot-" + ++slot, label: input.label }); + } + return Promise.resolve( + Ok({ + flowId: "login", + userCode: "CODE-1234", + verifyUrl: "https://auth.openai.com/codex/device", + intervalSeconds: 5, + }) + ); + }; + const finishLogin = () => { + const account = providers.openai.codexOauthAccounts?.find( + (item) => item.id === reconnectAccountId + ); + if (account) delete account.reconnectRequired; + return Promise.resolve(Ok(undefined)); + }; + client.codexOauth = { + startDeviceFlow: start, + startDesktopFlow: async (input) => { + browserLogin(input); + await start(input); + return Ok({ flowId: "login", authorizeUrl: "https://auth.openai.com/authorize" }); + }, + waitForDeviceFlow: finishLogin, + waitForDesktopFlow: finishLogin, + cancelDeviceFlow: () => Promise.resolve(), + cancelDesktopFlow: () => Promise.resolve(), + disconnect: (input) => { + providers.openai.codexOauthAccounts = providers.openai.codexOauthAccounts?.filter( + (account) => account.id !== (input?.accountId ?? "default") + ); + providers.openai.codexOauthSet = !!providers.openai.codexOauthAccounts?.length; + return Promise.resolve(Ok(undefined)); + }, + renameAccount: (input) => { + const account = providers.openai.codexOauthAccounts?.find( + (item) => item.id === input.accountId + ); + if (!account) return Promise.resolve(Err("Account is missing")); + account.label = input.label; + return Promise.resolve(Ok(undefined)); + }, + setDefaultAccount: (input) => { + providers.openai.codexOauthDefaultAccountId = input.accountId; + return Promise.resolve(Ok(undefined)); + }, + }; + client.providers.getConfig = () => Promise.resolve(structuredClone(providers)); + client.providers.setProviderConfig = (input) => { + if ( + input.keyPath[0] === "codexOauthDefaultAuth" && + (input.value === "apiKey" || input.value === "oauth") + ) { + providers.openai.codexOauthDefaultAuth = input.value; + } + return Promise.resolve(Ok(undefined)); + }; + client.projects.setCodexOauthAccount = (input) => { + const project = projects.get(input.projectPath); + if (!project) return Promise.resolve(Err("Project is missing")); + project.codexOauthAccountId = input.accountId ?? undefined; + return Promise.resolve(Ok(undefined)); + }; + return client; +} + +async function openAccounts(canvasElement: HTMLElement) { + const canvas = within(canvasElement); + await userEvent.click(await canvas.findByTestId("settings-button", {}, { timeout: 10000 })); + await userEvent.click(await canvas.findByRole("button", { name: "Providers" })); + const openai = await canvas.findByText("OpenAI", { exact: true }); + await userEvent.click(openai); + const section = await canvas.findByRole("region", { name: "ChatGPT (Codex) accounts" }); + return section; +} + +async function exerciseAccounts(canvasElement: HTMLElement) { + const section = await openAccounts(canvasElement); + const controls = within(section); + const global = controls.getByRole("combobox", { name: "Global default account" }); + const project = controls.getByRole("combobox", { name: "/projects/my-app" }); + await expect(global).toHaveValue("default"); + await expect(project).toHaveValue(""); + await userEvent.selectOptions(global, "work"); + await waitFor(() => expect(global).toHaveValue("work")); + await waitFor(() => expect(project).toBeEnabled()); + await userEvent.selectOptions(project, "work"); + await waitFor(() => expect(project).toHaveValue("work")); + + const work = within(controls.getByRole("listitem", { name: "Work" })); + await userEvent.click(work.getByRole("button", { name: "Rename" })); + const name = controls.getByRole("textbox", { name: "Account name" }); + await userEvent.clear(name); + await userEvent.type(name, "Team{Enter}"); + await controls.findByRole("listitem", { name: "Team" }); + await expect(global).toHaveDisplayValue("Team"); + await expect(project).toHaveDisplayValue("Team"); + + await userEvent.type(controls.getByRole("textbox", { name: "New account name" }), "Lab"); + await userEvent.click(controls.getByRole("button", { name: "Connect (Device)" })); + await controls.findByRole("listitem", { name: "Lab" }); + await expect(startLogin).toHaveBeenLastCalledWith({ label: "Lab" }); + await userEvent.click( + within(controls.getByRole("listitem", { name: "Team" })).getByRole("button", { + name: "Reconnect", + }) + ); + await waitFor(() => expect(global).toBeEnabled()); + await expect(startLogin).toHaveBeenLastCalledWith({ accountId: "work" }); + if (controls.queryByRole("button", { name: "Connect (Browser)" })) { + await expect(browserLogin).toHaveBeenLastCalledWith({ accountId: "work" }); + } + await expect(controls.getAllByRole("listitem")).toHaveLength(3); + + const auth = controls.getByRole("combobox", { name: "Default auth (when both are set)" }); + await userEvent.selectOptions(auth, "apiKey"); + await waitFor(() => expect(auth).toHaveValue("apiKey")); + await userEvent.click( + within(controls.getByRole("listitem", { name: "Team" })).getByRole("button", { + name: "Disconnect", + }) + ); + await waitFor(() => expect(controls.queryByRole("listitem", { name: "Team" })).toBeNull()); + await expect(global).toHaveValue("work"); + await expect(project).toHaveValue("work"); + await expect(global).toHaveDisplayValue(/Missing account/); + await expect(project).toHaveDisplayValue(/Missing account/); + await expect(controls.getAllByRole("listitem")).toHaveLength(2); + await userEvent.selectOptions(global, "default"); + await waitFor(() => expect(global).toHaveValue("default")); + await expect(project).toHaveValue("work"); + await userEvent.selectOptions(project, ""); + await waitFor(() => expect(project).toHaveValue("")); + await expect(project).toHaveDisplayValue(/Personal/); + await userEvent.selectOptions(project, "slot-1"); + await waitFor(() => expect(project).toHaveValue("slot-1")); + section.scrollIntoView({ block: "start" }); + + // The test-runner ignores viewport globals. Pixel and the phone runner enforce these bounds. + if (window.innerWidth < 768) { + await expect(section.getBoundingClientRect().right).toBeLessThanOrEqual(window.innerWidth); + await expect(section.scrollWidth).toBeLessThanOrEqual(section.clientWidth); + } +} + +export const Desktop: AppStory = { + globals: { viewport: { value: "desktop", isRotated: false } }, + parameters: { pixel: { matrix: { themes: ["dark", "light"], viewports: ["desktop"] } } }, + render: () => , + play: async ({ canvasElement }) => exerciseAccounts(canvasElement), +}; + +export const Phone: AppStory = { + ...Desktop, + globals: { viewport: { value: "mobile1", isRotated: false } }, + parameters: { pixel: { matrix: { themes: ["dark", "light"], viewports: ["phone"] } } }, + play: async ({ canvasElement }) => exerciseAccounts(canvasElement), +}; + +async function runAccountCommand(canvasElement: HTMLElement, title: string, choice?: string) { + const canvas = within(canvasElement); + await userEvent.keyboard("{F4}"); + const search = await canvas.findByPlaceholderText(/Switch workspaces or type/); + await userEvent.clear(search); + await userEvent.keyboard(">Codex: " + title); + await canvas.findByRole("option", { name: new RegExp("Codex: " + title) }); + await userEvent.keyboard("{Enter}"); + if (choice) { + const options = await canvas.findByPlaceholderText("Search options…"); + await userEvent.clear(options); + await userEvent.keyboard(choice); + await within(canvas.getByRole("listbox")).findByRole("option", { name: choice }); + await userEvent.keyboard("{Enter}"); + } + await waitFor(() => expect(canvas.queryByPlaceholderText("Search options…")).toBeNull()); +} + +async function exerciseKeyboardAccounts(canvasElement: HTMLElement) { + const canvas = within(canvasElement); + await canvas.findByTestId("settings-button", {}, { timeout: 10000 }); + // Reconnect must survive StrictMode when the command first mounts Settings. + await runAccountCommand(canvasElement, "Reconnect account", "Work"); + const section = await canvas.findByRole("region", { name: "ChatGPT (Codex) accounts" }); + const controls = within(section); + const newName = controls.getByRole("textbox", { name: "New account name" }); + await waitFor(() => expect(startLogin).toHaveBeenCalledTimes(1)); + await waitFor(() => expect(newName).toBeEnabled()); + await runAccountCommand(canvasElement, "Add account"); + await waitFor(() => expect(newName).toHaveFocus()); + await userEvent.keyboard("Lab{Enter}"); + await controls.findByRole("listitem", { name: "Lab" }); + await runAccountCommand(canvasElement, "Add account"); + await waitFor(() => expect(newName).toHaveFocus()); + await userEvent.keyboard("Second{Enter}"); + await controls.findByRole("listitem", { name: "Second" }); + + await runAccountCommand(canvasElement, "Rename account", "Work"); + const name = await controls.findByRole("textbox", { name: "Account name" }); + await waitFor(() => expect(name).toHaveFocus()); + await userEvent.clear(name); + await userEvent.keyboard("Team"); + await runAccountCommand(canvasElement, "Rename account", "Work"); + await waitFor(() => expect(name).toHaveFocus()); + await expect(name).toHaveValue("Team"); + await userEvent.keyboard("{Enter}"); + await controls.findByRole("listitem", { name: "Team" }); + + for (let attempt = 0; attempt < 2; attempt++) { + await runAccountCommand(canvasElement, "Reconnect account", "Team"); + await waitFor(() => expect(startLogin).toHaveBeenCalledTimes(4 + attempt)); + await expect(startLogin).toHaveBeenLastCalledWith({ accountId: "work" }); + await waitFor(() => expect(newName).toBeEnabled()); + } + await expect(controls.getAllByRole("listitem")).toHaveLength(4); + + const global = controls.getByRole("combobox", { name: "Global default account" }); + const project = controls.getByRole("combobox", { name: "/projects/my-app" }); + for (const accountId of ["work", "default"]) { + await runAccountCommand(canvasElement, "Change default account"); + await waitFor(() => expect(global).toHaveFocus()); + // userEvent does not implement native select keyboard actions. + await userEvent.selectOptions(global, accountId); + await waitFor(() => expect(global).toHaveValue(accountId)); + await runAccountCommand(canvasElement, "Change project account", "my-app"); + await waitFor(() => expect(project).toHaveFocus()); + await userEvent.selectOptions(project, accountId); + await waitFor(() => expect(project).toHaveValue(accountId)); + } + for (const account of ["Team", "Lab"]) { + await runAccountCommand(canvasElement, "Disconnect account", account); + await waitFor(() => expect(controls.queryByRole("listitem", { name: account })).toBeNull()); + } + await expect(controls.getAllByRole("listitem")).toHaveLength(2); + section.scrollIntoView({ block: "start" }); + if (window.innerWidth < 768) { + await expect(section.getBoundingClientRect().right).toBeLessThanOrEqual(window.innerWidth); + await expect(section.scrollWidth).toBeLessThanOrEqual(section.clientWidth); + } +} + +export const KeyboardCommands: AppStory = { + ...Desktop, + render: () => ( + + + + ), + play: async ({ canvasElement }) => exerciseKeyboardAccounts(canvasElement), +}; + +export const KeyboardCommandsPhone: AppStory = { + ...Phone, + render: KeyboardCommands.render, + play: KeyboardCommands.play, +}; + +async function exerciseDuplicateAccountLabels(canvasElement: HTMLElement) { + const section = await openAccounts(canvasElement); + const controls = within(section); + const global = controls.getByRole("combobox", { name: "Global default account" }); + const project = controls.getByRole("combobox", { name: "/projects/my-app" }); + const first = within(global).getByRole("option", { name: /Personal.*default/ }); + const second = within(global).getByRole("option", { name: /Personal.*work/ }); + const firstLabel = first.textContent?.trim(); + const secondLabel = second.textContent?.trim(); + if (!firstLabel || !secondLabel) throw new Error("Expected distinct account labels"); + await expect(firstLabel).not.toBe(secondLabel); + await expect(controls.getByRole("listitem", { name: firstLabel })).toBeVisible(); + await expect(controls.getByRole("listitem", { name: secondLabel })).toBeVisible(); + await expect( + within(project).getByRole("option", { name: /Inherit global default/ }) + ).toHaveTextContent(firstLabel); + + // The visible choice must select its stable ID, not the first matching stored name. + await userEvent.selectOptions(global, second); + await waitFor(() => expect(global).toHaveValue("work")); + await waitFor(() => expect(project).toBeEnabled()); + await userEvent.selectOptions(project, within(project).getByRole("option", { name: firstLabel })); + await waitFor(() => expect(project).toHaveValue("default")); + await runAccountCommand(canvasElement, "Reconnect account", secondLabel); + await waitFor(() => expect(startLogin).toHaveBeenLastCalledWith({ accountId: "work" })); + await waitFor(() => expect(global).toBeEnabled()); + + await runAccountCommand(canvasElement, "Rename account", firstLabel); + const name = await controls.findByRole("textbox", { name: "Account name" }); + await expect(name).toHaveValue("Personal"); + await userEvent.clear(name); + await userEvent.type(name, "Home{Enter}"); + await controls.findByRole("listitem", { name: "Home" }); + await expect(global).toHaveDisplayValue("Personal"); + await expect(project).toHaveDisplayValue("Home"); + await expect(global).toHaveValue("work"); + await expect(project).toHaveValue("default"); + + await runAccountCommand(canvasElement, "Rename account", "Home"); + const rename = await controls.findByRole("textbox", { name: "Account name" }); + await userEvent.clear(rename); + await userEvent.type(rename, "Personal{Enter}"); + await controls.findByRole("listitem", { name: firstLabel }); + await expect(global).toHaveDisplayValue(secondLabel); + await expect(project).toHaveDisplayValue(firstLabel); + section.scrollIntoView({ block: "start" }); + if (window.innerWidth < 768) { + await expect(section.getBoundingClientRect().right).toBeLessThanOrEqual(window.innerWidth); + await expect(section.scrollWidth).toBeLessThanOrEqual(section.clientWidth); + } +} + +export const DuplicateAccountLabels: AppStory = { + ...Desktop, + render: () => setupAccounts(false, "Personal")} />, + play: async ({ canvasElement }) => exerciseDuplicateAccountLabels(canvasElement), +}; + +export const DuplicateAccountLabelsPhone: AppStory = { + ...Phone, + render: DuplicateAccountLabels.render, + play: DuplicateAccountLabels.play, +}; + +async function exerciseRevokedSelections(canvasElement: HTMLElement) { + const section = await openAccounts(canvasElement); + const controls = within(section); + const global = controls.getByRole("combobox", { name: "Global default account" }); + const project = controls.getByRole("combobox", { name: "/projects/my-app" }); + const work = within(controls.getByRole("listitem", { name: "Work" })); + await expect(work.getByRole("button", { name: "Reconnect" })).toBeEnabled(); + + for (const select of [global, project]) { + // Stored selections stay visible, but revoked credentials cannot become new selections. + await expect(select).toHaveValue("work"); + await expect(select).toHaveDisplayValue(/Work.*Reconnect required/); + const options = within(select); + await expect(options.getByRole("option", { name: "Personal" })).toBeEnabled(); + await expect(options.getByRole("option", { name: /^Work/ })).toBeDisabled(); + await userEvent.selectOptions(select, "default"); + await waitFor(() => expect(select).toHaveValue("default")); + await waitFor(() => expect(select).toBeEnabled()); + await userEvent.selectOptions(select, "work"); + await expect(select).toHaveValue("default"); + } + await expect(work.getByText("Reconnect required")).toBeVisible(); + await expect(controls.getAllByRole("listitem")).toHaveLength(2); + + await userEvent.click(work.getByRole("button", { name: "Reconnect" })); + await waitFor(() => expect(work.queryByText("Reconnect required")).toBeNull()); + await expect(startLogin).toHaveBeenLastCalledWith({ accountId: "work" }); + for (const select of [global, project]) { + await waitFor(() => expect(select).toBeEnabled()); + await expect(within(select).getByRole("option", { name: "Work" })).toBeEnabled(); + await userEvent.selectOptions(select, "work"); + await waitFor(() => expect(select).toHaveValue("work")); + } + section.scrollIntoView({ block: "start" }); + if (window.innerWidth < 768) { + await expect(section.getBoundingClientRect().right).toBeLessThanOrEqual(window.innerWidth); + await expect(section.scrollWidth).toBeLessThanOrEqual(section.clientWidth); + } +} + +export const RevokedAccountSelections: AppStory = { + ...Desktop, + render: () => setupAccounts(true)} />, + play: async ({ canvasElement }) => exerciseRevokedSelections(canvasElement), +}; + +export const RevokedAccountSelectionsPhone: AppStory = { + ...Phone, + render: RevokedAccountSelections.render, + play: RevokedAccountSelections.play, +}; + +function setupReconnectRequired(apiKeySet = false) { + const client = setupAccounts(); + const providers: ProvidersConfigMap = { + openai: { + apiKeySet, + isConfigured: apiKeySet, + isEnabled: true, + codexOauthSet: false, + codexOauthDefaultAccountId: "work", + codexOauthAccounts: [{ id: "work", label: "Work", reconnectRequired: true }], + }, + }; + client.providers.getConfig = () => Promise.resolve(structuredClone(providers)); + const finishReconnect = () => { + providers.openai.codexOauthAccounts = [{ id: "work", label: "Work" }]; + providers.openai.codexOauthSet = true; + providers.openai.isConfigured = true; + return Promise.resolve(Ok(undefined)); + }; + client.codexOauth.waitForDesktopFlow = finishReconnect; + client.codexOauth.waitForDeviceFlow = finishReconnect; + return client; +} + +async function checkReconnectRequired(canvasElement: HTMLElement) { + const section = await openAccounts(canvasElement); + const controls = within(section); + const work = within(controls.getByRole("listitem", { name: "Work" })); + await expect(work.getByText("Reconnect required")).toBeVisible(); + await expect(controls.queryByText("Connected", { exact: true })).toBeNull(); + await expect(controls.getByRole("combobox", { name: "Global default account" })).toHaveValue( + "work" + ); + await expect(work.getByRole("button", { name: "Reconnect" })).toBeEnabled(); + section.scrollIntoView({ block: "start" }); + if (window.innerWidth < 768) { + await expect(section.getBoundingClientRect().right).toBeLessThanOrEqual(window.innerWidth); + await expect(section.scrollWidth).toBeLessThanOrEqual(section.clientWidth); + } + return controls; +} + +export const ReconnectRequired: AppStory = { + globals: { viewport: { value: "desktop", isRotated: false } }, + parameters: { pixel: { matrix: { themes: ["dark", "light"], viewports: ["desktop"] } } }, + render: () => , + play: async ({ canvasElement }) => { + await checkReconnectRequired(canvasElement); + }, +}; + +export const ReconnectRequiredPhone: AppStory = { + ...ReconnectRequired, + globals: { viewport: { value: "mobile1", isRotated: false } }, + parameters: { pixel: { matrix: { themes: ["dark", "light"], viewports: ["phone"] } } }, +}; + +export const ReconnectRestoresAccount: AppStory = { + render: () => setupReconnectRequired(true)} />, + play: async ({ canvasElement }) => { + const controls = await checkReconnectRequired(canvasElement); + const preference = controls.getByRole("combobox", { name: "Default auth (when both are set)" }); + await expect(preference).toBeEnabled(); + await expect(within(preference).getByRole("option", { name: /ChatGPT OAuth/ })).toBeDisabled(); + await userEvent.click( + within(controls.getByRole("listitem", { name: "Work" })).getByRole("button", { + name: "Reconnect", + }) + ); + await controls.findByText("Connected", { exact: true }); + await expect(preference).toBeEnabled(); + await expect(within(preference).getByRole("option", { name: /ChatGPT OAuth/ })).toBeEnabled(); + await expect(controls.queryByText("Reconnect required")).toBeNull(); + await expect(controls.getAllByRole("listitem")).toHaveLength(1); + await expect(controls.getByRole("combobox", { name: "Global default account" })).toHaveValue( + "work" + ); + await expect(startLogin).toHaveBeenLastCalledWith({ accountId: "work" }); + }, +}; + +async function exerciseDisconnectedDefaultRecovery(canvasElement: HTMLElement) { + const section = await openAccounts(canvasElement); + const controls = within(section); + const global = controls.getByRole("combobox", { name: "Global default account" }); + const preference = controls.getByRole("combobox", { name: "Default auth (when both are set)" }); + await userEvent.selectOptions(global, "work"); + await waitFor(() => expect(global).toHaveValue("work")); + for (const name of ["Personal", "Work"]) { + const account = controls.getByRole("listitem", { name }); + const disconnect = within(account).getByRole("button", { name: "Disconnect" }); + await waitFor(() => expect(disconnect).toBeEnabled()); + await userEvent.click(disconnect); + await waitFor(() => expect(controls.queryByRole("listitem", { name })).toBeNull()); + } + // Disconnect does not select another identity. Recovery requires an explicit auth choice. + await expect(global).toHaveValue("work"); + await expect(global).toBeDisabled(); + await expect(preference).toHaveValue("oauth"); + await waitFor(() => expect(preference).toBeEnabled()); + await userEvent.selectOptions(preference, "apiKey"); + await waitFor(() => expect(preference).toHaveValue("apiKey")); + await waitFor(() => expect(preference).toBeEnabled()); + await expect(within(preference).getByRole("option", { name: /ChatGPT OAuth/ })).toBeDisabled(); + await userEvent.selectOptions(preference, "oauth"); + await expect(preference).toHaveValue("apiKey"); + await expect(global).toHaveValue("work"); + section.scrollIntoView({ block: "start" }); + if (window.innerWidth < 768) { + await expect(section.getBoundingClientRect().right).toBeLessThanOrEqual(window.innerWidth); + await expect(section.scrollWidth).toBeLessThanOrEqual(section.clientWidth); + } +} + +export const DisconnectedDefaultRecovery: AppStory = { + ...Desktop, + play: async ({ canvasElement }) => exerciseDisconnectedDefaultRecovery(canvasElement), +}; + +export const DisconnectedDefaultRecoveryPhone: AppStory = { + ...Phone, + play: DisconnectedDefaultRecovery.play, +}; + +let contextStream: + | { + finish: () => void; + next: () => void; + reportUsage: () => void; + fallback: (model: string, effectiveContextLimit: number | null) => void; + } + | undefined; + +function setupLiveContextLimit(workspaceId = "codex-live-limit") { + contextStream = undefined; + const model = "openai:gpt-5.5"; + const usage = { inputTokens: 100_000, outputTokens: 0, totalTokens: 100_000 }; + const client = setupAccounts( + false, + "Work", + (workspaceId, emit) => { + let turn = 0; + const start = (effectiveContextLimit: number) => { + turn += 1; + const messageId = "context-turn-" + turn; + emit({ + type: "stream-start", + workspaceId, + messageId, + model, + historySequence: turn, + startTime: 1000 + turn, + effectiveContextLimit, + }); + }; + const refusedModels: string[] = []; + let currentModel = model; + contextStream = { + fallback: (nextModel, effectiveContextLimit) => { + refusedModels.push(currentModel); + currentModel = nextModel; + emit({ + type: "stream-model-update", + workspaceId, + messageId: "context-turn-" + turn, + model: nextModel, + metadataModel: nextModel, + effectiveContextLimit, + routedThroughGateway: false, + modelFallback: { requestedModel: model, refusedModels: [...refusedModels] }, + }); + }, + reportUsage: () => + emit({ + type: "usage-delta", + workspaceId, + messageId: "context-turn-" + turn, + usage, + cumulativeUsage: usage, + }), + finish: () => + emit({ + type: "stream-end", + workspaceId, + messageId: "context-turn-" + turn, + metadata: { model, usage, contextUsage: usage }, + parts: [{ type: "text", text: "The first turn is complete." }], + }), + next: () => start(500_000), + }; + queueMicrotask(() => { + emit({ type: "caught-up", hasOlderHistory: false }); + start(272_000); + }); + }, + workspaceId + ); + const getConfig = client.providers.getConfig; + client.providers.getConfig = async () => { + const providers = await getConfig(); + return { + ...providers, + openai: { ...providers.openai, models: [{ id: "gpt-5.5", contextWindowTokens: 500_000 }] }, + }; + }; + updatePersistedState(getModelKey(workspaceId), model); + updatePersistedState(RIGHT_SIDEBAR_TAB_KEY, "costs"); + expandRightSidebar(); + return client; +} + +async function checkContextMeters( + canvasElement: HTMLElement, + limit: string, + percentage: string, + tokens = "100.0k" +) { + const canvas = within(canvasElement); + await waitFor( + async () => { + await expect( + canvas.getByRole("button", { + name: new RegExp("Context usage: " + tokens + " / " + limit), + }) + ).toHaveAccessibleName(expect.stringContaining(percentage)); + await expect(canvas.getByTestId("context-usage")).toHaveTextContent(limit); + await expect(canvas.getByTestId("context-usage")).toHaveTextContent(percentage); + }, + { timeout: 10000 } + ); +} + +async function exerciseLiveContextLimit(canvasElement: HTMLElement) { + const canvas = within(canvasElement); + await checkContextMeters(canvasElement, "272.0k", "0.0%", "0"); + // Settings changes must not alter the accepted limit before the first usage event. + const controls = within(await openAccounts(canvasElement)); + const global = controls.getByRole("combobox", { name: "Global default account" }); + const project = controls.getByRole("combobox", { name: "/projects/my-app" }); + await userEvent.selectOptions(global, "work"); + await waitFor(() => expect(global).toHaveValue("work")); + await waitFor(() => expect(project).toBeEnabled()); + await userEvent.selectOptions(project, "work"); + await waitFor(() => expect(project).toHaveValue("work")); + const preference = controls.getByRole("combobox", { name: "Default auth (when both are set)" }); + await waitFor(() => expect(preference).toBeEnabled()); + await userEvent.selectOptions(preference, "apiKey"); + await waitFor(() => expect(preference).toHaveValue("apiKey")); + await userEvent.click( + canvas.getAllByRole("button", { name: /Close settings|Back to previous page/ })[0] + ); + await checkContextMeters(canvasElement, "272.0k", "0.0%", "0"); + if (!contextStream) throw new Error("The live context stream is missing"); + contextStream.reportUsage(); + await checkContextMeters(canvasElement, "272.0k", "36.8%"); + contextStream.finish(); + await checkContextMeters(canvasElement, "500.0k", "20.0%"); + contextStream.next(); + await checkContextMeters(canvasElement, "500.0k", "20.0%"); + if (window.innerWidth < 768) { + const meter = canvas.getByTestId("context-usage"); + await expect(meter.getBoundingClientRect().right).toBeLessThanOrEqual(window.innerWidth); + } +} + +export const LiveContextLimit: AppStory = { + render: () => , + play: async ({ canvasElement }) => exerciseLiveContextLimit(canvasElement), +}; + +export const LiveContextLimitPhone: AppStory = { + ...LiveContextLimit, + render: () => setupLiveContextLimit("codex-live-limit-phone")} />, + globals: { viewport: { value: "mobile1", isRotated: false } }, + parameters: { pixel: { matrix: { themes: ["dark", "light"], viewports: ["phone"] } } }, +}; + +async function exerciseFallbackContextLimit(canvasElement: HTMLElement) { + const canvas = within(canvasElement); + await checkContextMeters(canvasElement, "272.0k", "0.0%", "0"); + if (!contextStream) throw new Error("The fallback context stream is missing"); + contextStream.reportUsage(); + await checkContextMeters(canvasElement, "272.0k", "36.8%", "100.0k"); + // No usage follows these updates. Both meters must immediately use each accepted fallback limit. + contextStream.fallback("anthropic:claude-sonnet-4-5", 200_000); + await checkContextMeters(canvasElement, "200.0k", "0.0%", "0"); + contextStream.fallback("openai:gpt-5.5", null); + await waitFor(async () => { + await expect( + canvas.getByRole("button", { name: "Context usage: 0 (unknown limit)" }) + ).toBeVisible(); + await expect(within(canvas.getByTestId("context-usage")).queryByRole("slider")).toBeNull(); + await expect(canvas.getByTestId("context-usage")).not.toHaveTextContent("272.0k"); + await expect(canvas.getByTestId("context-usage")).not.toHaveTextContent("200.0k"); + }); + if (window.innerWidth < 768) { + await expect( + canvas.getByTestId("context-usage").getBoundingClientRect().right + ).toBeLessThanOrEqual(window.innerWidth); + } +} + +export const FallbackContextLimit: AppStory = { + render: () => setupLiveContextLimit("codex-fallback-limit")} />, + play: async ({ canvasElement }) => exerciseFallbackContextLimit(canvasElement), +}; + +export const FallbackContextLimitPhone: AppStory = { + ...FallbackContextLimit, + render: () => setupLiveContextLimit("codex-fallback-limit-phone")} />, + play: async ({ canvasElement }) => exerciseFallbackContextLimit(canvasElement), + globals: { viewport: { value: "mobile1", isRotated: false } }, + parameters: { pixel: { matrix: { themes: ["dark", "light"], viewports: ["phone"] } } }, +}; + +function setupLoginFailure() { + const client = setupAccounts(); + let firstAttempt = true; + let finishWait: () => void = () => undefined; + client.codexOauth.startDeviceFlow = () => { + if (firstAttempt) { + firstAttempt = false; + return Promise.resolve(Err("OpenAI login is unavailable. Try again.")); + } + return Promise.resolve( + Ok({ + flowId: "pending-login", + userCode: "CODE-1234", + verifyUrl: "https://auth.openai.com/codex/device", + intervalSeconds: 5, + }) + ); + }; + client.codexOauth.waitForDeviceFlow = () => + new Promise((resolve) => { + finishWait = () => resolve(Err("Login cancelled")); + }); + client.codexOauth.cancelDeviceFlow = () => { + finishWait(); + return Promise.resolve(); + }; + client.codexOauth.setDefaultAccount = () => Promise.resolve(Err("Account update failed")); + return client; +} + +export const LoginFailureAndCancel: AppStory = { + render: () => , + play: async ({ canvasElement }) => { + const section = await openAccounts(canvasElement); + const controls = within(section); + const global = controls.getByRole("combobox", { name: "Global default account" }); + await userEvent.selectOptions(global, "work"); + await controls.findByRole("alert"); + await expect(global).toHaveValue("default"); + const name = controls.getByRole("textbox", { name: "New account name" }); + await userEvent.type(name, "Lab"); + await userEvent.click(controls.getByRole("button", { name: "Connect (Device)" })); + await waitFor(() => + expect(controls.getByRole("alert")).toHaveTextContent(/login is unavailable/) + ); + await expect(controls.getAllByRole("listitem")).toHaveLength(2); + await userEvent.click(controls.getByRole("button", { name: "Connect (Device)" })); + await controls.findByText("CODE-1234"); + await expect(name).toBeDisabled(); + await userEvent.click(controls.getByRole("button", { name: "Cancel" })); + await waitFor(() => expect(name).toBeEnabled()); + await expect(controls.queryByText("CODE-1234")).toBeNull(); + await expect(controls.queryByRole("alert")).toBeNull(); + await expect(controls.getAllByRole("listitem")).toHaveLength(2); + section.scrollIntoView({ block: "start" }); + }, +}; + +function setupScopedAccount(kind: "subproject" | "multi" | "creation") { + const root = "/projects/account-root"; + const subproject = root + "/sub"; + const workspace = createWorkspace({ + id: "scoped-account", + name: "main", + projectName: "account-root", + projectPath: root, + }); + if (kind === "subproject") workspace.subProjectPath = subproject; + if (kind === "multi") { + workspace.projectPath = MULTI_PROJECT_CONFIG_KEY; + workspace.projects = [ + { projectPath: subproject, projectName: "sub" }, + { projectPath: root, projectName: "account-root" }, + ]; + } + selectWorkspace(workspace); + if (kind === "creation") { + expandLeftSidebar(); + expandProjects([root, subproject]); + } else { + collapseLeftSidebar(); + } + const projects = groupWorkspacesByProject([workspace]); + projects.set(root, { + ...projects.get(root), + codexOauthAccountId: "default", + workspaces: projects.get(root)?.workspaces ?? [], + }); + projects.set(subproject, { + parentProjectPath: root, + codexOauthAccountId: "deleted", + workspaces: [], + }); + return createMockORPCClient({ + projects, + workspaces: [workspace], + agentAiDefaults: { + exec: { modelString: "openai:gpt-5.3-codex-spark" }, + plan: { modelString: "openai:gpt-5.3-codex-spark" }, + }, + providersList: ["openai"], + providersConfig: { + openai: { + apiKeySet: true, + isConfigured: true, + isEnabled: true, + codexOauthSet: true, + codexOauthAccounts: [{ id: "default", label: "Personal" }], + }, + }, + }); +} + +async function checkScopedAccountWarning(canvasElement: HTMLElement) { + const warning = await within(canvasElement).findByTestId( + "codex-oauth-warning-banner", + {}, + { timeout: 10000 } + ); + await waitFor(() => expect(warning).toBeVisible()); + if (window.innerWidth < 768) { + await expect(warning.getBoundingClientRect().right).toBeLessThanOrEqual(window.innerWidth); + } +} + +export const SubprojectRouting: AppStory = { + render: () => setupScopedAccount("subproject")} />, + play: async ({ canvasElement }) => checkScopedAccountWarning(canvasElement), +}; + +export const MultiProjectRouting: AppStory = { + render: () => setupScopedAccount("multi")} />, + play: SubprojectRouting.play, +}; + +export const SubprojectRoutingPhone: AppStory = { + ...SubprojectRouting, + globals: { viewport: { value: "mobile1", isRotated: false } }, + parameters: { pixel: { matrix: { themes: ["dark"], viewports: ["phone"] } } }, +}; + +export const CreationSubprojectRouting: AppStory = { + render: () => setupScopedAccount("creation")} />, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await userEvent.click( + await canvas.findByRole("button", { name: "New chat in sub-project" }, { timeout: 10000 }) + ); + await checkScopedAccountWarning(canvasElement); + }, +}; + +function setupProjectTitle() { + const client = setupAccounts(); + generateTitle.mockClear(); + client.nameGeneration.generate = (input) => { + generateTitle(input); + return Promise.resolve( + Ok({ name: "project-title", title: "Project title", modelUsed: "openai:gpt-5.5" }) + ); + }; + return client; +} + +export const TitleUsesProject: AppStory = { + render: () => , + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const newChat = await canvas.findByRole( + "button", + { name: "New chat in my-app" }, + { timeout: 10000 } + ); + newChat.focus(); + await userEvent.keyboard("{Enter}"); + const message = await canvas.findByRole("textbox", { name: "Message Claude" }); + await userEvent.type(message, "Fix the project title"); + await waitFor(() => + expect(generateTitle).toHaveBeenLastCalledWith( + expect.objectContaining({ projectPath: "/projects/my-app" }) + ) + ); + await waitFor(() => + expect(canvasElement.querySelector("#workspace-name")).toHaveValue("project-title") + ); + }, +}; diff --git a/src/browser/utils/codexAccountDisplay.test.ts b/src/browser/utils/codexAccountDisplay.test.ts new file mode 100644 index 00000000000..382dc939288 --- /dev/null +++ b/src/browser/utils/codexAccountDisplay.test.ts @@ -0,0 +1,29 @@ +import { expect, test } from "bun:test"; +import { formatCodexAccountLabel } from "./codexAccountDisplay"; + +const accounts = [ + { id: "slot-one", label: "Personal" }, + { id: "slot-two", label: "Personal" }, + { id: "slot-three", label: "Work" }, +]; + +test("duplicate account labels use stable IDs regardless of list order", () => { + const labels = accounts.map((account) => formatCodexAccountLabel(account, accounts)); + expect(new Set(labels).size).toBe(accounts.length); + for (const account of accounts.slice(0, 2)) { + const label = formatCodexAccountLabel(account, accounts); + expect(label).toContain(account.id); + expect(label).toContain(account.label); + expect(formatCodexAccountLabel(account, accounts.toReversed())).toBe(label); + } + expect(labels[2]).toBe(accounts[2].label); +}); + +test("renaming or removing a duplicate restores the unique label without changing identity", () => { + const [first, second] = accounts; + const renamed = { ...second, label: "Home" }; + expect(formatCodexAccountLabel(first, [first, renamed])).toBe(first.label); + expect(formatCodexAccountLabel(renamed, [first, renamed])).toBe(renamed.label); + expect(formatCodexAccountLabel(first, [first])).toBe(first.label); + expect(first).toEqual({ id: "slot-one", label: "Personal" }); +}); diff --git a/src/browser/utils/codexAccountDisplay.ts b/src/browser/utils/codexAccountDisplay.ts new file mode 100644 index 00000000000..db606d50e6a --- /dev/null +++ b/src/browser/utils/codexAccountDisplay.ts @@ -0,0 +1,11 @@ +import type { ProviderConfigInfo } from "@/common/orpc/types"; + +type Account = Pick[number], "id" | "label">; + +/** Disambiguate duplicate labels without changing stored names or account IDs. */ +export function formatCodexAccountLabel(account: Account, accounts: readonly Account[]): string { + const duplicate = accounts.some( + (other) => other.id !== account.id && other.label === account.label + ); + return duplicate ? account.label + " (" + account.id + ")" : account.label; +} diff --git a/src/browser/utils/commandIds.ts b/src/browser/utils/commandIds.ts index 81ff199ad92..46fb7adbde2 100644 --- a/src/browser/utils/commandIds.ts +++ b/src/browser/utils/commandIds.ts @@ -94,6 +94,9 @@ export const CommandIds = { // Settings commands settingsOpen: () => "settings:open" as const, settingsOpenSection: (section: string) => `settings:open:${section}` as const, + codexAccountAction: ( + action: "add" | "reconnect" | "rename" | "disconnect" | "default" | "project" + ) => `providers:openai:codex:${action}` as const, coderDisconnect: () => "providers:coder:disconnect" as const, coderRefreshModels: () => "providers:coder:refresh-models" as const, diff --git a/src/browser/utils/commands/sources.test.ts b/src/browser/utils/commands/sources.test.ts index b6ae2b3d289..cd31b4362af 100644 --- a/src/browser/utils/commands/sources.test.ts +++ b/src/browser/utils/commands/sources.test.ts @@ -491,6 +491,132 @@ test("multi-project workspace command hides itself when the experiment is disabl expect(onStartMultiProjectWorkspaceCreation).not.toHaveBeenCalled(); }); +test("Codex commands open each operation without making account RPCs", async () => { + const onOpenSettings = mock(); + const actions = getActions({ + api: null, + onOpenSettings, + providersConfig: { + openai: { + apiKeySet: false, + isEnabled: true, + isConfigured: false, + codexOauthSet: false, + codexOauthAccounts: [{ id: "work", label: "Work", reconnectRequired: true }], + }, + }, + }); + for (const type of ["add", "reconnect", "rename", "disconnect", "default", "project"] as const) { + const action = actions.find((candidate) => candidate.id === "providers:openai:codex:" + type); + expect(action).toBeDefined(); + expect(action?.visible?.()).toBe(true); + expect(action?.enabled?.() ?? true).toBe(true); + if (action?.prompt) { + const field = action.prompt.fields[0]; + if (field.type !== "select") throw new Error("Expected a metadata selector"); + const choices = await field.getOptions({}); + expect(choices.map((choice) => choice.id)).toEqual([type === "project" ? "/repo/a" : "work"]); + await action.prompt.onSubmit( + type === "project" ? { projectPath: "/repo/a" } : { accountId: "work" } + ); + } else { + await action?.run(); + } + expect(onOpenSettings).toHaveBeenLastCalledWith("providers", { + expandProvider: "openai", + codexAccountAction: + type === "project" + ? { type, projectPath: "/repo/a" } + : type === "add" || type === "default" + ? { type } + : { type, accountId: "work" }, + }); + } + expect(onOpenSettings).toHaveBeenCalledTimes(6); +}); + +test("Codex account commands distinguish duplicate labels without changing account IDs", async () => { + const accounts = [ + { id: "first-slot", label: "Personal" }, + { id: "second-slot", label: "Personal" }, + { id: "work-slot", label: "Work" }, + ]; + const onOpenSettings = mock(); + const actions = getActions({ + onOpenSettings, + providersConfig: { + openai: { + apiKeySet: false, + isEnabled: true, + isConfigured: true, + codexOauthSet: true, + codexOauthAccounts: accounts, + }, + }, + }); + for (const type of ["reconnect", "rename", "disconnect"] as const) { + const prompt = actions.find((action) => action.id === "providers:openai:codex:" + type)?.prompt; + const field = prompt?.fields[0]; + if (!prompt || field?.type !== "select") throw new Error("Expected account selector"); + const choices = await field.getOptions({}); + expect(new Set(choices.map((choice) => choice.label)).size).toBe(accounts.length); + expect(choices.map((choice) => choice.id)).toEqual(accounts.map((account) => account.id)); + expect(choices[2].label).toBe(accounts[2].label); + const second = choices.find((choice) => choice.label.includes(accounts[1].id)); + expect(second).toBeDefined(); + await prompt.onSubmit({ accountId: second!.id }); + expect(onOpenSettings).toHaveBeenLastCalledWith("providers", { + expandProvider: "openai", + codexAccountAction: { type, accountId: accounts[1].id }, + }); + } +}); + +const hiddenCodexProviders: Array<{ + name: string; + providersConfig: Parameters[0]["providersConfig"]; +}> = [ + { name: "loading metadata", providersConfig: undefined }, + { name: "policy-hidden OpenAI", providersConfig: {} }, + { + name: "a custom OpenAI shadow", + providersConfig: { + openai: { apiKeySet: true, isEnabled: true, isConfigured: true, isCustom: true }, + }, + }, +]; + +test.each(hiddenCodexProviders)("Codex commands stay hidden with $name", ({ providersConfig }) => { + const commands = getActions({ onOpenSettings: mock(), providersConfig }).filter((action) => + action.id.startsWith("providers:openai:codex:") + ); + expect(commands.length).toBeGreaterThan(0); + for (const command of commands) { + expect(command.visible?.()).toBe(false); + } +}); + +test("Codex commands use legacy metadata and disable account operations without slots", async () => { + const openai = { apiKeySet: false, isEnabled: true, isConfigured: true, codexOauthSet: true }; + const legacy = getActions({ onOpenSettings: mock(), providersConfig: { openai } }); + const field = legacy.find((action) => action.id === "providers:openai:codex:reconnect")?.prompt + ?.fields[0]; + if (field?.type !== "select") throw new Error("Expected account selector"); + expect((await field.getOptions({})).map((choice) => choice.id)).toEqual(["default"]); + const disconnected = getActions({ + onOpenSettings: mock(), + providersConfig: { openai: { ...openai, isConfigured: false, codexOauthSet: false } }, + }); + for (const type of ["reconnect", "rename", "disconnect", "default"]) { + expect( + disconnected.find((action) => action.id === "providers:openai:codex:" + type)?.enabled?.() + ).toBe(false); + } + expect( + disconnected.find((action) => action.id === "providers:openai:codex:add")?.visible?.() + ).toBe(true); +}); + test("Login with Coder command opens providers expanded on Coder and starts the login", async () => { // Regression: the command must reach the login operation (expand the Coder // provider and start the OAuth flow via the one-shot hints consumed by diff --git a/src/browser/utils/commands/sources.ts b/src/browser/utils/commands/sources.ts index fba9c4fb94e..d55b57c179c 100644 --- a/src/browser/utils/commands/sources.ts +++ b/src/browser/utils/commands/sources.ts @@ -1,5 +1,10 @@ import { THEME_OPTIONS, type ThemePreference } from "@/browser/contexts/ThemeContext"; -import type { OpenSettingsOptions } from "@/browser/contexts/SettingsContext"; +import type { + CodexAccountSettingsIntent, + OpenSettingsOptions, +} from "@/browser/contexts/SettingsContext"; +import { formatCodexAccountLabel } from "@/browser/utils/codexAccountDisplay"; +import { CODEX_OAUTH_DEFAULT_ACCOUNT_ID } from "@/common/constants/codexOauthAccounts"; import type { CommandAction } from "@/browser/contexts/CommandRegistryContext"; import type { APIClient } from "@/browser/contexts/API"; import type { ConfirmDialogOptions } from "@/browser/contexts/ConfirmDialogContext"; @@ -99,6 +104,7 @@ export interface BuildSourcesParams { } | null; /** Project-scoped preference ID used while a creation composer is active. */ creationScopeId?: string | null; + codexOauthAccountId?: string; streamingModels?: Map; // UI actions getThinkingLevel: (workspaceId: string) => ThinkingLevel; @@ -1304,6 +1310,7 @@ export function buildCoreSources(p: BuildSourcesParams): Array<() => CommandActi getFastModeProvider(providerOptionGateModel ?? "", { providersConfig: p.providersConfig, resolvedRouteProvider: providerOptionRoute, + codexOauthAccountId: p.codexOauthAccountId, }) != null ? { id: CommandIds.toggleFastMode(), @@ -1416,6 +1423,7 @@ export function buildCoreSources(p: BuildSourcesParams): Array<() => CommandActi openaiProModeAvailable(proGateModelString ?? "", { providersConfig: p.providersConfig, resolvedRouteProvider: currentModelRoute, + codexOauthAccountId: p.codexOauthAccountId, }) ) { const proActive = p.getReasoningMode(workspaceId) === "pro"; @@ -1606,6 +1614,97 @@ export function buildCoreSources(p: BuildSourcesParams): Array<() => CommandActi // Settings if (p.onOpenSettings) { const openSettings = p.onOpenSettings; + const openCodexAction = (intent: CodexAccountSettingsIntent) => + openSettings("providers", { expandProvider: "openai", codexAccountAction: intent }); + const codexVisible = () => { + // Policy can omit OpenAI metadata. Custom providers do not expose built-in account controls. + const openai = p.providersConfig?.openai; + return openai != null && openai.isCustom !== true; + }; + const getCodexAccounts = () => { + const openai = p.providersConfig?.openai; + return ( + openai?.codexOauthAccounts ?? + (openai?.codexOauthSet ? [{ id: CODEX_OAUTH_DEFAULT_ACCOUNT_ID, label: "Default" }] : []) + ); + }; + actions.push(() => [ + { + id: CommandIds.codexAccountAction("add"), + title: "Codex: Add account", + section: section.settings, + keywords: ["openai", "chatgpt", "oauth", "connect", "login"], + visible: codexVisible, + run: () => openCodexAction({ type: "add" }), + }, + ...(["reconnect", "rename", "disconnect"] as const).map((type): CommandAction => { + const titles = { reconnect: "Reconnect", rename: "Rename", disconnect: "Disconnect" }; + return { + id: CommandIds.codexAccountAction(type), + title: "Codex: " + titles[type] + " account…", + section: section.settings, + keywords: ["openai", "chatgpt", "oauth", type], + visible: codexVisible, + enabled: () => getCodexAccounts().length > 0, + run: () => undefined, + prompt: { + title: titles[type] + " Codex account", + fields: [ + { + type: "select", + name: "accountId", + label: "Account", + getOptions: () => { + const accounts = getCodexAccounts(); + return accounts.map((account) => ({ + id: account.id, + label: formatCodexAccountLabel(account, accounts), + keywords: [account.id], + })); + }, + }, + ], + onSubmit: (values) => openCodexAction({ type, accountId: values.accountId }), + }, + }; + }), + { + id: CommandIds.codexAccountAction("default"), + title: "Codex: Change default account", + section: section.settings, + keywords: ["openai", "chatgpt", "oauth", "global", "default"], + visible: codexVisible, + enabled: () => getCodexAccounts().length > 0, + run: () => openCodexAction({ type: "default" }), + }, + { + id: CommandIds.codexAccountAction("project"), + title: "Codex: Change project account…", + section: section.settings, + keywords: ["openai", "chatgpt", "oauth", "project", "inherit"], + visible: codexVisible, + enabled: () => p.userProjects.size > 0, + run: () => undefined, + prompt: { + title: "Change project Codex account", + fields: [ + { + type: "select", + name: "projectPath", + label: "Project", + getOptions: () => + Array.from(p.userProjects.keys(), (projectPath) => ({ + id: projectPath, + label: formatProjectHierarchyLabel(projectPath, p.userProjects), + keywords: [projectPath], + })), + }, + ], + onSubmit: (values) => + openCodexAction({ type: "project", projectPath: values.projectPath }), + }, + }, + ]); actions.push(() => [ { id: CommandIds.settingsOpen(), diff --git a/src/browser/utils/compaction/contextSwitchCheck.test.ts b/src/browser/utils/compaction/contextSwitchCheck.test.ts index 1b4efddff8b..8f7a8ea87d6 100644 --- a/src/browser/utils/compaction/contextSwitchCheck.test.ts +++ b/src/browser/utils/compaction/contextSwitchCheck.test.ts @@ -10,6 +10,24 @@ const OPTIONS = { }; describe("checkContextSwitch", () => { + test("caps the context limit for the project account", () => { + const warning = checkContextSwitch(350_000, "openai:gpt-5.5", "google:gemini-2.5-pro", false, { + ...OPTIONS, + codexOauthAccountId: "work", + providersConfig: { + openai: { + apiKeySet: true, + isEnabled: true, + isConfigured: true, + codexOauthSet: true, + codexOauthAccounts: [{ id: "work", label: "Work" }], + }, + }, + }); + expect(warning).not.toBeNull(); + expect(warning?.targetLimit).toBeLessThan(350_000); + }); + test("returns null when target model matches previous model", () => { const targetModel = "openai:gpt-5.2-codex"; const limit = getEffectiveContextLimit(targetModel, false); diff --git a/src/browser/utils/compaction/contextSwitchCheck.ts b/src/browser/utils/compaction/contextSwitchCheck.ts index 8210bbfe60a..c255318ea4e 100644 --- a/src/browser/utils/compaction/contextSwitchCheck.ts +++ b/src/browser/utils/compaction/contextSwitchCheck.ts @@ -41,6 +41,7 @@ export function findPreviousModel(messages: DisplayedMessage[]): string | null { /** Options for accessibility checks in context switch validation */ export interface ContextSwitchOptions extends CompactionRouteOptions { providersConfig: ProvidersConfigMap | null; + codexOauthAccountId?: string; policy: EffectivePolicy | null; } @@ -63,7 +64,9 @@ function resolveCompactionModel( routePriority: options.routePriority, routeOverrides: options.routeOverrides, }); - const limit = getEffectiveContextLimit(preferred, use1M, options.providersConfig); + const limit = getEffectiveContextLimit(preferred, use1M, options.providersConfig, { + codexOauthAccountId: options.codexOauthAccountId, + }); if (accessible && limit && limit > currentTokens) return preferred; } if (previousModel) { @@ -74,7 +77,9 @@ function resolveCompactionModel( routePriority: options.routePriority, routeOverrides: options.routeOverrides, }); - const limit = getEffectiveContextLimit(previousModel, use1M, options.providersConfig); + const limit = getEffectiveContextLimit(previousModel, use1M, options.providersConfig, { + codexOauthAccountId: options.codexOauthAccountId, + }); if (accessible && limit && limit > currentTokens) return previousModel; } return null; @@ -108,7 +113,9 @@ export function checkContextSwitch( return null; } - const targetLimit = getEffectiveContextLimit(targetModel, use1M, options.providersConfig); + const targetLimit = getEffectiveContextLimit(targetModel, use1M, options.providersConfig, { + codexOauthAccountId: options.codexOauthAccountId, + }); // Unknown model or context fits with 10% buffer - no warning if (!targetLimit || currentTokens <= targetLimit * CONTEXT_FIT_THRESHOLD) { diff --git a/src/browser/utils/fastModeServiceTier.test.ts b/src/browser/utils/fastModeServiceTier.test.ts index ba3820602f2..f139ef00e2d 100644 --- a/src/browser/utils/fastModeServiceTier.test.ts +++ b/src/browser/utils/fastModeServiceTier.test.ts @@ -20,6 +20,30 @@ function createWriter() { } describe("fast mode service tier", () => { + test("uses the project account and preserves the API-key preference", () => { + const providersConfig = { + openai: { + apiKeySet: true, + isEnabled: true, + isConfigured: true, + codexOauthSet: true, + codexOauthAccounts: [{ id: "work", label: "Work" }], + }, + }; + expect( + getFastModeProvider("openai:gpt-5.5", { + providersConfig, + codexOauthAccountId: "work", + }) + ).toBeNull(); + expect( + getFastModeProvider("openai:gpt-5.5", { + providersConfig: { openai: { ...providersConfig.openai, codexOauthDefaultAuth: "apiKey" } }, + codexOauthAccountId: "work", + }) + ).toBe("openai"); + }); + test("resolves direct native providers and rejects gateway routes", () => { expect(getFastModeProvider("openai:gpt-5.6-sol", { resolvedRouteProvider: "direct" })).toBe( "openai" diff --git a/src/browser/utils/fastModeServiceTier.ts b/src/browser/utils/fastModeServiceTier.ts index 3b08ea1255a..ddf5a831a29 100644 --- a/src/browser/utils/fastModeServiceTier.ts +++ b/src/browser/utils/fastModeServiceTier.ts @@ -21,6 +21,7 @@ export interface FastModeServiceTierChange { export interface FastModeAvailabilityOptions { resolvedRouteProvider?: string | null; providersConfig?: ProvidersConfigMap | null; + codexOauthAccountId?: string; } type ProviderConfigWriter = Pick; diff --git a/src/browser/utils/messages/StreamingMessageAggregator.test.ts b/src/browser/utils/messages/StreamingMessageAggregator.test.ts index dcaa4070817..f004241a5d3 100644 --- a/src/browser/utils/messages/StreamingMessageAggregator.test.ts +++ b/src/browser/utils/messages/StreamingMessageAggregator.test.ts @@ -3515,6 +3515,123 @@ describe("StreamingMessageAggregator", () => { }); }); + test.each([false, true])( + "fallback metadata preserves content and timing, replay: %s", + (replay) => { + const aggregator = createTestAggregator(); + const start = { + workspaceId: TEST_WORKSPACE_ID, + messageId: "msg-1", + model: "mux-gateway:anthropic/claude-sonnet-4-5", + historySequence: 1, + startTime: 1000, + }; + aggregator.handleStreamStart({ + ...start, + type: "stream-start", + effectiveContextLimit: 200_000, + routedThroughGateway: true, + routeProvider: "mux-gateway", + thinkingLevel: "high", + }); + aggregator.handleStreamDelta({ + type: "stream-delta", + workspaceId: TEST_WORKSPACE_ID, + messageId: "msg-1", + delta: "Partial answer", + tokens: 3, + timestamp: 1100, + }); + startToolCall(aggregator, { + toolCallId: "tool-a", + toolName: "bash", + args: {}, + timestamp: 1200, + }); + aggregator.handleToolCallExecutionStart({ + type: "tool-call-execution-start", + workspaceId: TEST_WORKSPACE_ID, + messageId: "msg-1", + toolCallId: "tool-a", + timestamp: 1250, + }); + aggregator.handleStreamDelta({ + type: "stream-delta", + workspaceId: TEST_WORKSPACE_ID, + messageId: "msg-1", + delta: "More text", + tokens: 2, + timestamp: 1300, + }); + const original = aggregator.getAllMessages()[0]; + const originalTimestamp = original.metadata?.timestamp; + const parts = structuredClone(original.parts); + const timing = aggregator.getActiveStreamTimingStats(); + const refusedModels = [start.model]; + for (const limit of [272_000, null]) { + aggregator.handleUsageDelta({ + type: "usage-delta", + workspaceId: TEST_WORKSPACE_ID, + messageId: "msg-1", + usage: { inputTokens: 1000, outputTokens: 3, totalTokens: 1003 }, + cumulativeUsage: { inputTokens: 1000, outputTokens: 3, totalTokens: 1003 }, + }); + const update = { + ...start, + model: "openai:gpt-5.5", + metadataModel: "openai:gpt-5.5", + effectiveContextLimit: limit, + routedThroughGateway: false, + modelFallback: { requestedModel: start.model, refusedModels: [...refusedModels] }, + }; + if (replay) { + aggregator.handleStreamStart({ ...update, type: "stream-start", replay: true }); + } else { + aggregator.handleStreamModelUpdate({ ...update, type: "stream-model-update" }); + } + const current = aggregator.getAllMessages()[0]; + expect(current.parts).toEqual(parts); + expect(current.metadata?.timestamp).toBe(originalTimestamp); + expect(current.metadata).toMatchObject({ + model: update.model, + metadataModel: update.metadataModel, + routedThroughGateway: false, + modelFallback: update.modelFallback, + }); + expect(current.metadata?.routeProvider).toBeUndefined(); + expect(current.metadata?.thinkingLevel).toBeUndefined(); + expect(aggregator.getActiveStreamContextLimit("msg-1")).toBe(limit); + expect(aggregator.getActiveStreamUsage("msg-1")).toBeUndefined(); + const currentTiming = aggregator.getActiveStreamTimingStats(); + expect(currentTiming?.model).toBe(update.model); + if (replay) { + // Replay translates server timestamps again. The first-token duration must remain unchanged. + expect(currentTiming!.firstTokenTime! - currentTiming!.startTime).toBe( + timing!.firstTokenTime! - timing!.startTime + ); + } else { + expect(currentTiming).toMatchObject({ + startTime: timing?.startTime, + firstTokenTime: timing?.firstTokenTime, + }); + } + expect(aggregator.getActiveStreamMetadataModel()).toBe(update.metadataModel); + expect( + aggregator.getDisplayedMessages().findLast((row) => row.type === "assistant") + ?.streamPresentation + ).toEqual({ source: replay ? "replay" : "live" }); + refusedModels.push(update.model); + } + endToolCall(aggregator, { + toolCallId: "tool-a", + toolName: "bash", + result: {}, + timestamp: 2000, + }); + expect(aggregator.getActiveStreamTimingStats()?.toolExecutionMs).toBe(750); + } + ); + describe("usage-delta handling", () => { test("handleUsageDelta stores usage by messageId", () => { const aggregator = new StreamingMessageAggregator(TEST_CREATED_AT); @@ -3534,6 +3651,36 @@ describe("StreamingMessageAggregator", () => { }); }); + test.each([272_000, null])( + "keeps start-time limits across replay and missing usage fields: %s", + (limit) => { + const aggregator = new StreamingMessageAggregator(TEST_CREATED_AT); + const start = { + type: "stream-start" as const, + workspaceId: "ws-1", + messageId: "msg-1", + model: "openai:gpt-5.5", + historySequence: 1, + startTime: 1000, + }; + aggregator.handleStreamStart({ ...start, effectiveContextLimit: limit }); + expect(aggregator.getActiveStreamUsage("msg-1")).toBeUndefined(); + expect(aggregator.getActiveStreamContextLimit("msg-1")).toBe(limit); + aggregator.handleStreamStart({ ...start, replay: true }); + expect(aggregator.getActiveStreamContextLimit("msg-1")).toBe(limit); + aggregator.handleUsageDelta({ + type: "usage-delta", + workspaceId: "ws-1", + messageId: "msg-1", + usage: { inputTokens: 1000, outputTokens: 0, totalTokens: 1000 }, + cumulativeUsage: { inputTokens: 1000, outputTokens: 0, totalTokens: 1000 }, + }); + expect(aggregator.getActiveStreamContextLimit("msg-1")).toBe(limit); + aggregator.handleStreamStart({ ...start, replay: true, effectiveContextLimit: null }); + expect(aggregator.getActiveStreamContextLimit("msg-1")).toBeNull(); + } + ); + test("clearTokenState removes usage", () => { const aggregator = new StreamingMessageAggregator(TEST_CREATED_AT); diff --git a/src/browser/utils/messages/StreamingMessageAggregator.ts b/src/browser/utils/messages/StreamingMessageAggregator.ts index 37d11203608..2e714c054d2 100644 --- a/src/browser/utils/messages/StreamingMessageAggregator.ts +++ b/src/browser/utils/messages/StreamingMessageAggregator.ts @@ -16,6 +16,7 @@ import { import { copyStreamLifecycleSnapshot, type StreamStartEvent, + type StreamModelUpdateEvent, type StreamDeltaEvent, type UsageDeltaEvent, type StreamEndEvent, @@ -248,6 +249,8 @@ interface StreamingContext { * stream is active. */ metadataModel?: string; + effectiveContextLimit?: number | null; + modelFallback?: MuxMetadata["modelFallback"]; routedThroughGateway?: boolean; routeProvider?: string; @@ -2158,6 +2161,17 @@ export class StreamingMessageAggregator { this.getLatestUnresolvedCompactionRequest()?.parsed.followUpContent?.dispatchOptions ?.source === "internal-resume"; const now = Date.now(); + const existingContext = this.activeStreams.get(data.messageId); + // A reconnect can skip the update event while the backend accepts a fallback attempt. + if ( + data.replay && + existingContext && + (existingContext.model !== data.model || + existingContext.modelFallback?.refusedModels.length !== + data.modelFallback?.refusedModels.length) + ) { + this.activeStreamUsage.delete(data.messageId); + } const context: StreamingContext = { serverStartTime: data.startTime, clockOffsetMs: now - data.startTime, @@ -2169,6 +2183,13 @@ export class StreamingMessageAggregator { isReplay: data.replay === true, model: data.model, metadataModel: data.metadataModel, + modelFallback: data.modelFallback, + effectiveContextLimit: + data.effectiveContextLimit !== undefined + ? data.effectiveContextLimit + : data.replay + ? existingContext?.effectiveContextLimit + : undefined, routedThroughGateway: data.routedThroughGateway, routeProvider, serverFirstTokenTime: null, @@ -2182,7 +2203,6 @@ export class StreamingMessageAggregator { // For incremental replay: stream-start may be re-emitted to re-establish context. // If we already have this message with accumulated parts, don't wipe its content. const existingMessage = this.messages.get(data.messageId); - const existingContext = this.activeStreams.get(data.messageId); if (data.replay && existingMessage && existingMessage.parts.length > 0) { if (existingContext) { // Preserve the highest observed server timestamp across reconnect boundaries. @@ -2205,6 +2225,8 @@ export class StreamingMessageAggregator { this.activeStreams.set(data.messageId, context); if (existingMessage.metadata) { existingMessage.metadata.model = data.model; + existingMessage.metadata.metadataModel = data.metadataModel; + existingMessage.metadata.modelFallback = data.modelFallback; existingMessage.metadata.routedThroughGateway = data.routedThroughGateway; existingMessage.metadata.routeProvider = routeProvider; if (data.agentId != null) { @@ -2226,6 +2248,8 @@ export class StreamingMessageAggregator { historySequence: data.historySequence, timestamp: Date.now(), model: data.model, + metadataModel: data.metadataModel, + modelFallback: data.modelFallback, routedThroughGateway: data.routedThroughGateway, routeProvider, agentId: data.agentId, @@ -2237,6 +2261,32 @@ export class StreamingMessageAggregator { this.markMessageDirty(data.messageId); } + handleStreamModelUpdate(data: StreamModelUpdateEvent): void { + const context = this.activeStreams.get(data.messageId); + const message = this.messages.get(data.messageId); + if (!context || !message) return; + + // Keep parts, timestamps, and tool timing. Usage belongs to the new attempt only after its first step. + context.model = data.model; + context.metadataModel = data.metadataModel; + context.effectiveContextLimit = data.effectiveContextLimit; + context.modelFallback = data.modelFallback; + context.routedThroughGateway = data.routedThroughGateway; + context.routeProvider = resolveRouteProvider(data.routeProvider, data.routedThroughGateway); + context.thinkingLevel = data.thinkingLevel; + message.metadata = { + ...message.metadata, + model: data.model, + metadataModel: data.metadataModel, + modelFallback: data.modelFallback, + routedThroughGateway: data.routedThroughGateway, + routeProvider: context.routeProvider, + thinkingLevel: data.thinkingLevel, + }; + this.activeStreamUsage.delete(data.messageId); + this.markMessageDirty(data.messageId); + } + handleStreamDelta(data: StreamDeltaEvent): void { const message = this.messages.get(data.messageId); if (!message) return; @@ -3965,6 +4015,10 @@ export class StreamingMessageAggregator { * Handle usage-delta event: update usage tracking for active stream */ handleUsageDelta(data: UsageDeltaEvent): void { + const context = this.activeStreams.get(data.messageId); + if (context && data.effectiveContextLimit !== undefined) { + context.effectiveContextLimit = data.effectiveContextLimit; + } this.activeStreamUsage.set(data.messageId, { step: { usage: data.usage, providerMetadata: data.providerMetadata }, cumulative: { @@ -3974,6 +4028,10 @@ export class StreamingMessageAggregator { }); } + getActiveStreamContextLimit(messageId: string): number | null | undefined { + return this.activeStreams.get(messageId)?.effectiveContextLimit; + } + /** * Get active stream usage for context window display (last step's inputTokens = context size) */ diff --git a/src/browser/utils/messages/applyWorkspaceChatEventToAggregator.test.ts b/src/browser/utils/messages/applyWorkspaceChatEventToAggregator.test.ts index 9877f5f4b76..fb860daf673 100644 --- a/src/browser/utils/messages/applyWorkspaceChatEventToAggregator.test.ts +++ b/src/browser/utils/messages/applyWorkspaceChatEventToAggregator.test.ts @@ -10,6 +10,7 @@ import type { StreamDeltaEvent, StreamEndEvent, StreamStartEvent, + StreamModelUpdateEvent, ToolCallDeltaEvent, ToolCallEndEvent, ToolCallExecutionStartEvent, @@ -29,6 +30,10 @@ class StubAggregator implements WorkspaceChatEventAggregator { this.calls.push(`handleStreamStart:${data.messageId}`); } + handleStreamModelUpdate(data: StreamModelUpdateEvent): void { + this.calls.push("handleStreamModelUpdate:" + data.messageId); + } + handleStreamDelta(data: StreamDeltaEvent): void { this.calls.push(`handleStreamDelta:${data.messageId}`); } diff --git a/src/browser/utils/messages/applyWorkspaceChatEventToAggregator.ts b/src/browser/utils/messages/applyWorkspaceChatEventToAggregator.ts index e476b174e20..bab554020ea 100644 --- a/src/browser/utils/messages/applyWorkspaceChatEventToAggregator.ts +++ b/src/browser/utils/messages/applyWorkspaceChatEventToAggregator.ts @@ -35,6 +35,7 @@ import type { StreamDeltaEvent, StreamEndEvent, StreamStartEvent, + StreamModelUpdateEvent, ToolCallDeltaEvent, ToolCallEndEvent, ToolCallExecutionStartEvent, @@ -63,6 +64,7 @@ export interface ApplyWorkspaceChatEventToAggregatorOptions { */ export interface WorkspaceChatEventAggregator { handleStreamStart(data: StreamStartEvent): void; + handleStreamModelUpdate(data: StreamModelUpdateEvent): void; handleStreamDelta(data: StreamDeltaEvent): void; handleStreamEnd(data: StreamEndEvent): void; handleStreamAbort(data: StreamAbortEvent): void; @@ -132,6 +134,11 @@ export function applyWorkspaceChatEventToAggregator( return "immediate"; } + if (event.type === "stream-model-update") { + aggregator.handleStreamModelUpdate(event); + return "immediate"; + } + if (isStreamDelta(event)) { aggregator.handleStreamDelta(event); return "throttled"; diff --git a/src/cli/run.ts b/src/cli/run.ts index b5624fc04bd..143de9c0649 100644 --- a/src/cli/run.ts +++ b/src/cli/run.ts @@ -15,7 +15,11 @@ import { z } from "zod"; import * as path from "path"; import * as fs from "fs/promises"; import { createConfigStores } from "../node/config"; -import { materializeResolvedTrust, replaceRunTrustProjects } from "./trust"; +import { + materializeCodexOauthAccount, + materializeResolvedTrust, + replaceRunTrustProjects, +} from "./trust"; import { runBestEffortCleanup } from "./runCleanup"; import { DisposableTempDir } from "../node/services/tempDir"; import { AgentSession, type AgentSessionChatEvent } from "../node/services/agentSession"; @@ -653,7 +657,6 @@ async function main(): Promise { initStateManager, backgroundProcessManager, mcpServerManager, - providerService, workspaceService, workspaceGoalService, idleDispatcher, @@ -681,11 +684,6 @@ async function main(): Promise { : undefined, }); - // `xum run` uses createCoreServices directly (without ServiceContainer), so wire - // Codex OAuth explicitly to ensure Codex-routed OpenAI requests can load/refresh - // OAuth tokens from providers.jsonc. - const codexOauthService = new CodexOauthService(runProvidersStore, providerService); - turnRequestBuilderBindings.codexOauthService = codexOauthService; // Same for Coder OAuth: coder:* models need per-request token loading/refresh. // Bind it to the REAL config (not the ephemeral tempDir copy): Coder rotates // the refresh token on every use, so persisting rotations only to tempDir @@ -698,6 +696,9 @@ async function main(): Promise { realProvidersStore, realFileLeaseManager ); + // Keep Codex token rotations after the temporary CLI root is removed. + const codexOauthService = new CodexOauthService(realProvidersStore, realProviderService); + turnRequestBuilderBindings.codexOauthService = codexOauthService; const coderOauthService = new CoderOauthService( realProvidersStore, realFileLeaseManager, @@ -887,6 +888,9 @@ async function main(): Promise { projectName: path.basename(projectDir), runtimeConfig, }); + const registered = config.findWorkspace(workspaceId); + assert(registered, "CLI workspace registration must exist"); + await materializeCodexOauthAccount(realConfig, config, projectDir, registered.projectPath); // Note: taskService.initialize() is intentionally NOT called. It resumes tasks from a // previous session, but xum run uses an ephemeral Config (temp dir) with no prior state. diff --git a/src/cli/trust.test.ts b/src/cli/trust.test.ts index 11b969f707f..031a599086c 100644 --- a/src/cli/trust.test.ts +++ b/src/cli/trust.test.ts @@ -1,10 +1,15 @@ import * as fs from "node:fs/promises"; import * as path from "node:path"; -import { describe, expect, test } from "bun:test"; +import { describe, expect, spyOn, test } from "bun:test"; import { DisposableTempDir } from "@/node/services/tempDir"; import { Config } from "@/node/config"; -import { materializeResolvedTrust, replaceRunTrustProjects, resolveProjectDir } from "./trust"; +import { + materializeCodexOauthAccount, + materializeResolvedTrust, + replaceRunTrustProjects, + resolveProjectDir, +} from "./trust"; const BUN_EXECUTABLE = process.execPath; const TRUST_ENTRY = path.join(import.meta.dir, "trust.ts"); @@ -98,6 +103,262 @@ describe("xum trust CLI", () => { expect(trustByPath.get(worktree)).toBe(false); }, 15_000); + test("copies Codex overrides without requiring project trust", async () => { + using tmp = new DisposableTempDir("codex-project-copy"); + const real = new Config(path.join(tmp.path, "real")); + const target = new Config(path.join(tmp.path, "target")); + const projectPath = path.join(tmp.path, "project"); + await real.editConfig((config) => { + config.projects.set(projectPath, { workspaces: [], codexOauthAccountId: "work" }); + return config; + }); + await replaceRunTrustProjects(real, target); + expect(target.loadConfigOrDefault().projects.get(projectPath)).toMatchObject({ + codexOauthAccountId: "work", + workspaces: [], + }); + await real.editConfig((config) => { + config.projects.delete(projectPath); + return config; + }); + await replaceRunTrustProjects(real, target); + expect(target.loadConfigOrDefault().projects.has(projectPath)).toBe(false); + }); + + test("materializes subproject accounts onto temporary CLI workspace projects", async () => { + using tmp = new DisposableTempDir("codex-worktree-copy"); + const real = new Config(path.join(tmp.path, "real")); + const target = new Config(path.join(tmp.path, "target")); + const root = path.join(tmp.path, "project"); + const subproject = path.join(root, "subproject"); + const worktree = path.join(tmp.path, "checkout"); + const cliProject = path.join(tmp.path, "cli-project"); + await real.editConfig((config) => { + config.projects.set(root, { + codexOauthAccountId: "personal", + workspaces: [{ path: worktree, id: "test-account-workspace", subProjectPath: subproject }], + }); + config.projects.set(subproject, { workspaces: [], codexOauthAccountId: "work" }); + return config; + }); + await materializeCodexOauthAccount(real, target, worktree, cliProject); + expect(target.loadConfigOrDefault().projects.get(cliProject)?.codexOauthAccountId).toBe("work"); + await real.editConfig((config) => { + delete config.projects.get(subproject)!.codexOauthAccountId; + return config; + }); + await materializeCodexOauthAccount(real, target, worktree, cliProject); + expect( + target.loadConfigOrDefault().projects.get(cliProject)?.codexOauthAccountId + ).toBeUndefined(); + }); + + test("copies a physical non-git project's account through a requested symlink", async () => { + using tmp = new DisposableTempDir("codex-requested-symlink"); + const real = new Config(path.join(tmp.path, "real")); + const target = new Config(path.join(tmp.path, "target")); + const projectPath = path.join(tmp.path, "project"); + const alias = path.join(tmp.path, "alias"); + await fs.mkdir(projectPath); + await fs.symlink(projectPath, alias, "junction"); + await real.editConfig((config) => { + config.projects.set(projectPath, { workspaces: [], codexOauthAccountId: "work" }); + return config; + }); + await materializeCodexOauthAccount(real, target, alias, alias); + expect(target.loadConfigOrDefault().projects.get(alias)?.codexOauthAccountId).toBe("work"); + }); + + test("resolves linked worktree fallback through a registered repository alias", async () => { + using tmp = new DisposableTempDir("codex-linked-worktree-alias"); + const real = new Config(path.join(tmp.path, "real")); + const target = new Config(path.join(tmp.path, "target")); + const repo = path.join(tmp.path, "repo"); + const repoAlias = path.join(tmp.path, "repo-alias"); + const worktree = path.join(tmp.path, "worktree"); + const worktreeAlias = path.join(tmp.path, "worktree-alias"); + await fs.mkdir(repo); + await Bun.$`git init`.cwd(repo).quiet(); + await Bun.$`git -c user.name=Test -c user.email=test@example.com commit --allow-empty -m init` + .cwd(repo) + .quiet(); + await Bun.$`git worktree add ${worktree} -b feature`.cwd(repo).quiet(); + await fs.mkdir(path.join(worktree, "src")); + await fs.symlink(repo, repoAlias, "junction"); + await fs.symlink(worktree, worktreeAlias, "junction"); + await real.editConfig((config) => { + config.projects.set(repoAlias, { workspaces: [], codexOauthAccountId: "work" }); + return config; + }); + await materializeCodexOauthAccount( + real, + target, + path.join(worktreeAlias, "src"), + worktreeAlias + ); + expect(target.loadConfigOrDefault().projects.get(worktreeAlias)?.codexOauthAccountId).toBe( + "work" + ); + }); + + test.each(["work", undefined])( + "resolves non-git account scopes through directory aliases: %s", + async (accountId) => { + using tmp = new DisposableTempDir("codex-symlink-account"); + const real = new Config(path.join(tmp.path, "real")); + const target = new Config(path.join(tmp.path, "target")); + const root = path.join(tmp.path, "project"); + const subproject = path.join(root, "packages", "api"); + const rootAlias = path.join(tmp.path, "a-long-parent-alias"); + const subprojectAlias = path.join(tmp.path, "child"); + const targetProject = path.join(tmp.path, "cli-project"); + await fs.mkdir(path.join(subproject, "src"), { recursive: true }); + await fs.mkdir(path.join(root, "packages", "api-other"), { recursive: true }); + await fs.symlink(root, rootAlias, "junction"); + await fs.symlink(subproject, subprojectAlias, "junction"); + await real.editConfig((config) => { + // Alias length must not let a parent displace a deeper physical project. + config.projects.set(rootAlias, { workspaces: [], codexOauthAccountId: "personal" }); + config.projects.set(subprojectAlias, { workspaces: [], codexOauthAccountId: accountId }); + return config; + }); + for (const requestedPath of [ + subproject, + subprojectAlias, + path.join(rootAlias, "packages", "api"), + path.join(subproject, "src"), + path.join(rootAlias, "packages", "api", "src"), + ]) { + await target.editConfig((config) => { + config.projects.set(targetProject, { workspaces: [], codexOauthAccountId: "stale" }); + return config; + }); + await materializeCodexOauthAccount(real, target, requestedPath, targetProject); + const project = target.loadConfigOrDefault().projects.get(targetProject); + expect(project?.codexOauthAccountId).toBe(accountId); + expect(Object.hasOwn(project ?? {}, "codexOauthAccountId")).toBe(accountId !== undefined); + } + await materializeCodexOauthAccount( + real, + target, + path.join(root, "packages", "api-other"), + targetProject + ); + expect(target.loadConfigOrDefault().projects.get(targetProject)?.codexOauthAccountId).toBe( + "personal" + ); + // An exact configured path still wins over another spelling of the same directory. + await real.editConfig((config) => { + config.projects.set(subproject, { workspaces: [], codexOauthAccountId: "exact" }); + return config; + }); + await materializeCodexOauthAccount(real, target, subproject, targetProject); + expect(target.loadConfigOrDefault().projects.get(targetProject)?.codexOauthAccountId).toBe( + "exact" + ); + } + ); + + test("matches aliased workspace paths and subproject references", async () => { + using tmp = new DisposableTempDir("codex-symlink-workspace"); + const real = new Config(path.join(tmp.path, "real")); + const target = new Config(path.join(tmp.path, "target")); + const root = path.join(tmp.path, "project"); + const subproject = path.join(root, "subproject"); + const subprojectAlias = path.join(tmp.path, "subproject-alias"); + const checkout = path.join(tmp.path, "checkout"); + const checkoutAlias = path.join(tmp.path, "checkout-alias"); + const cliProject = path.join(tmp.path, "cli-project"); + await fs.mkdir(subproject, { recursive: true }); + await fs.mkdir(checkout); + await fs.symlink(subproject, subprojectAlias, "junction"); + await fs.symlink(checkout, checkoutAlias, "junction"); + await real.editConfig((config) => { + config.projects.set(root, { + codexOauthAccountId: "personal", + workspaces: [{ id: "aliased-workspace", path: checkoutAlias, subProjectPath: subproject }], + }); + config.projects.set(subprojectAlias, { workspaces: [], codexOauthAccountId: "work" }); + return config; + }); + await materializeCodexOauthAccount(real, target, checkout, cliProject); + expect(target.loadConfigOrDefault().projects.get(cliProject)?.codexOauthAccountId).toBe("work"); + await real.editConfig((config) => { + delete config.projects.get(subprojectAlias)!.codexOauthAccountId; + return config; + }); + await materializeCodexOauthAccount(real, target, checkoutAlias, cliProject); + expect( + target.loadConfigOrDefault().projects.get(cliProject)?.codexOauthAccountId + ).toBeUndefined(); + }); + + test.each(["work", undefined])( + "uses the deepest registered account scope: %s", + async (accountId) => { + using tmp = new DisposableTempDir("codex-nested-account"); + const real = new Config(path.join(tmp.path, "real")); + const target = new Config(path.join(tmp.path, "target")); + const root = path.join(tmp.path, "project"); + const subproject = path.join(root, "packages", "api"); + const targetProject = path.join(tmp.path, "cli-project"); + await real.editConfig((config) => { + config.projects.set(subproject, { workspaces: [], codexOauthAccountId: accountId }); + config.projects.set(root, { workspaces: [], codexOauthAccountId: "personal" }); + return config; + }); + await materializeCodexOauthAccount(real, target, path.join(subproject, "src"), targetProject); + expect(target.loadConfigOrDefault().projects.get(targetProject)?.codexOauthAccountId).toBe( + accountId + ); + + await materializeCodexOauthAccount( + real, + target, + path.join(root, "packages", "api-other", "src"), + targetProject + ); + expect(target.loadConfigOrDefault().projects.get(targetProject)?.codexOauthAccountId).toBe( + "personal" + ); + } + ); + + test.each(["work", undefined])( + "rejects a lost account selection write: %s", + async (accountId) => { + using tmp = new DisposableTempDir("codex-account-write-failure"); + const real = new Config(path.join(tmp.path, "real")); + const target = new Config(path.join(tmp.path, "target")); + const projectPath = path.join(tmp.path, "project"); + await real.editConfig((config) => { + config.projects.set(projectPath, { workspaces: [], codexOauthAccountId: accountId }); + return config; + }); + await target.editConfig((config) => { + config.projects.set(projectPath, { workspaces: [], codexOauthAccountId: "personal" }); + return config; + }); + // Simulate a config edit that reports success without writing the selection. + const edit = spyOn(target, "editConfig").mockResolvedValue(undefined); + try { + let error: unknown; + try { + await materializeCodexOauthAccount(real, target, projectPath, projectPath); + } catch (caught) { + error = caught; + } + expect(error).toBeInstanceOf(Error); + expect(String(error)).toContain("Failed to persist Codex OAuth account"); + expect(target.loadConfigOrDefault().projects.get(projectPath)?.codexOauthAccountId).toBe( + "personal" + ); + } finally { + edit.mockRestore(); + } + } + ); + test("replaceRunTrustProjects rebuilds config without foreign settings", async () => { using tmp = new DisposableTempDir("trust-replace-run"); const realConfig = new Config(path.join(tmp.path, "real-root")); diff --git a/src/cli/trust.ts b/src/cli/trust.ts index 5ec85f4b11a..c392e4fbf4e 100644 --- a/src/cli/trust.ts +++ b/src/cli/trust.ts @@ -168,26 +168,94 @@ export async function resolveProjectTrusted( return mainRepoDir != null && isProjectTrusted(realConfig, mainRepoDir); } -/** Replace all run-root trust entries so removed grants cannot survive root reuse. */ +/** Replace run project settings so removed trust grants and account overrides cannot survive root reuse. */ export async function replaceRunTrustProjects( realConfig: Config, targetConfig: Config ): Promise { const trustOnlyProjects = new Map(); for (const [projectPath, projectConfig] of realConfig.loadConfigOrDefault().projects) { - if (projectConfig.trusted === undefined) { + if (projectConfig.trusted === undefined && projectConfig.codexOauthAccountId === undefined) { continue; } trustOnlyProjects.set(projectPath, { workspaces: [], trusted: projectConfig.trusted, + codexOauthAccountId: projectConfig.codexOauthAccountId, }); } await targetConfig.editConfig(() => ({ projects: trustOnlyProjects })); } +/** Copy the source project account onto the CLI workspace project. */ +export async function materializeCodexOauthAccount( + realConfig: Config, + targetConfig: Config, + projectDir: string, + targetProjectPath: string +): Promise { + const projects = realConfig.loadConfigOrDefault().projects; + // Compare physical directories, but retain config keys and exact-path precedence. + const resolvedProjectDir = await realpathOrResolve(projectDir); + const resolvedProjects = new Map(); + for (const projectPath of projects.keys()) { + resolvedProjects.set(projectPath, await realpathOrResolve(projectPath)); + } + const findExactProject = (requestedPath: string, resolvedPath: string): string | undefined => + projects.has(requestedPath) + ? requestedPath + : Array.from(resolvedProjects).find(([, physicalPath]) => physicalPath === resolvedPath)?.[0]; + let sourcePath = findExactProject(projectDir, resolvedProjectDir); + if (sourcePath === undefined) { + for (const [projectPath, project] of projects) { + for (const workspace of project.workspaces) { + if ((await realpathOrResolve(workspace.path)) === resolvedProjectDir) { + sourcePath = + workspace.projects?.[0]?.projectPath ?? workspace.subProjectPath ?? projectPath; + break; + } + } + if (sourcePath !== undefined) break; + } + } + if (sourcePath === undefined) { + // An explicit directory can sit below a registered subproject. Keep its account scope. + let longestMatch = -1; + for (const [projectPath, physicalPath] of resolvedProjects) { + const relative = path.relative(physicalPath, resolvedProjectDir); + if (relative === ".." || relative.startsWith(`..${path.sep}`) || path.isAbsolute(relative)) { + continue; + } + if (physicalPath.length > longestMatch) { + sourcePath = projectPath; + longestMatch = physicalPath.length; + } + } + } + sourcePath ??= + (await findMainRepoDir(projectDir)) ?? (await findGitRoot(projectDir)) ?? projectDir; + const sourceKey = findExactProject(sourcePath, await realpathOrResolve(sourcePath)); + const accountId = + sourceKey === undefined ? undefined : projects.get(sourceKey)?.codexOauthAccountId; + await targetConfig.editConfig((config) => { + const project = config.projects.get(targetProjectPath) ?? { workspaces: [] }; + if (accountId === undefined) { + delete project.codexOauthAccountId; + } else { + project.codexOauthAccountId = accountId; + } + config.projects.set(targetProjectPath, project); + return config; + }); + // Config.saveConfig swallows write errors. Never send with an unintended account. + const persistedProject = targetConfig.loadConfigOrDefault().projects.get(targetProjectPath); + if (!persistedProject || persistedProject.codexOauthAccountId !== accountId) { + throw new Error(`Failed to persist Codex OAuth account for ${targetProjectPath}`); + } +} + /** * Copy effective trust onto the target path because the task-spawn gate uses exact lookups. */ diff --git a/src/cli/workflow.ts b/src/cli/workflow.ts index baccd8b04c6..b95ce46f48f 100644 --- a/src/cli/workflow.ts +++ b/src/cli/workflow.ts @@ -9,7 +9,6 @@ import * as path from "node:path"; import { Command } from "commander"; import { EXPERIMENT_IDS, LEGACY_PTC_EXCLUSIVE_EXPERIMENT_ID } from "@/common/constants/experiments"; -import type { ProjectConfig } from "@/common/types/project"; import { parseRuntimeModeAndHost, RUNTIME_MODE, type RuntimeConfig } from "@/common/types/runtime"; import { DEFAULT_THINKING_LEVEL, @@ -46,7 +45,12 @@ import { hasAnyConfiguredProvider, buildProvidersFromEnv } from "@/node/utils/pr import { runBestEffortCleanup } from "./runCleanup"; import { getParseOptions } from "./argv"; import { exitAfterStdoutFlush } from "./processExit"; -import { resolveProjectDir, resolveProjectTrusted } from "./trust"; +import { + materializeCodexOauthAccount, + replaceRunTrustProjects, + resolveProjectDir, + resolveProjectTrusted, +} from "./trust"; const VALID_EXPERIMENT_IDS = new Set(Object.values(EXPERIMENT_IDS)); const THINKING_LABELS_LIST = [...new Set(Object.values(THINKING_DISPLAY_LABELS))].join(", "); @@ -226,17 +230,7 @@ async function copyPersistentConfig( await runStores.secretsStore.saveSecretsConfig(existingSecrets); } - const existingConfig = realConfig.loadConfigOrDefault(); - const trustOnlyProjects = new Map(); - for (const [projectPath, projectConfig] of existingConfig.projects) { - if (projectConfig.trusted !== undefined) { - trustOnlyProjects.set(projectPath, { workspaces: [], trusted: projectConfig.trusted }); - } - } - if (trustOnlyProjects.size > 0) { - // Config.saveConfig is private (lost-update safety); route through the queue. - await config.editConfig((cfg) => ({ ...cfg, projects: trustOnlyProjects })); - } + await replaceRunTrustProjects(realConfig, config); } function buildExperimentsObject(experimentIds: readonly string[]) { @@ -370,8 +364,6 @@ async function createWorkflowContext(options: { extensionMetadataPath: path.join(tempDir.path, "extensionMetadata.json"), mcpConfig: realConfig, }); - codexOauthService = new CodexOauthService(runProvidersStore, services.providerService); - services.turnRequestBuilderBindings.codexOauthService = codexOauthService; // Bind Coder OAuth to the REAL config (not the ephemeral tempDir copy): // Coder rotates the refresh token on every use, so persisting rotations // only to tempDir would strand ~/.xum/providers.jsonc with a consumed @@ -382,6 +374,9 @@ async function createWorkflowContext(options: { realProvidersStore, realFileLeaseManager ); + // Keep Codex token rotations after the temporary CLI root is removed. + codexOauthService = new CodexOauthService(realProvidersStore, realProviderService); + services.turnRequestBuilderBindings.codexOauthService = codexOauthService; coderOauthService = new CoderOauthService( realProvidersStore, realFileLeaseManager, @@ -427,6 +422,14 @@ async function createWorkflowContext(options: { projectName: path.basename(options.projectDir), runtimeConfig, }); + const registered = config.findWorkspace(workspaceId); + assert(registered, "Workflow workspace registration must exist"); + await materializeCodexOauthAccount( + realConfig, + config, + options.projectDir, + registered.projectPath + ); assert(workspacePath.length > 0, "xum workflow workspace path must be non-empty"); return { diff --git a/src/common/config/schemas/providersConfig.test.ts b/src/common/config/schemas/providersConfig.test.ts index 60e86fc3ca0..4caecc6f568 100644 --- a/src/common/config/schemas/providersConfig.test.ts +++ b/src/common/config/schemas/providersConfig.test.ts @@ -3,6 +3,23 @@ import { describe, expect, it } from "bun:test"; import { ProvidersConfigSchema } from "./providersConfig"; describe("ProvidersConfigSchema", () => { + it("requires named OAuth credentials to use the protected storage field", () => { + const credentials = { type: "oauth", access: "access", refresh: "refresh", expires: 1000 }; + const account = { label: "Work", credentials }; + const document = { openai: { codexOauthAccounts: { work: account } } }; + expect(ProvidersConfigSchema.parse(document)).toEqual(document); + for (const unsafe of [ + { label: "Work", auth: credentials }, + { ...account, auth: credentials }, + ]) { + expect( + ProvidersConfigSchema.safeParse({ + openai: { codexOauthAccounts: { work: unsafe } }, + }).success + ).toBe(false); + } + }); + it("validates a valid providers config with anthropic key", () => { const valid = { anthropic: { apiKey: "sk-ant-123", cacheTtl: "5m" }, diff --git a/src/common/config/schemas/providersConfig.ts b/src/common/config/schemas/providersConfig.ts index 1af5601552e..abb380709cf 100644 --- a/src/common/config/schemas/providersConfig.ts +++ b/src/common/config/schemas/providersConfig.ts @@ -40,6 +40,13 @@ export const OpenAIProviderConfigSchema = BaseProviderConfigSchema.extend({ organization: z.string().optional(), codexOauthDefaultAuth: CodexOauthDefaultAuthSchema.optional(), codexOauth: z.record(z.string(), z.unknown()).optional(), + // Keep named tokens under credentials so older recursive config redactors cannot expose them. + codexOauthAccounts: z + .record( + z.string(), + z.object({ label: z.string(), credentials: z.record(z.string(), z.unknown()) }).strict() + ) + .optional(), defaultModel: z.string().optional(), apiVersion: z.string().optional(), webSocketTransportEnabled: z.boolean().optional(), diff --git a/src/common/constants/codexOauthAccounts.ts b/src/common/constants/codexOauthAccounts.ts new file mode 100644 index 00000000000..c9dcd47a33b --- /dev/null +++ b/src/common/constants/codexOauthAccounts.ts @@ -0,0 +1,10 @@ +// Local slot IDs differ from ChatGPT account IDs in token claims. +export const CODEX_OAUTH_DEFAULT_ACCOUNT_ID = "default"; +export const CODEX_OAUTH_REFRESH_TIMEOUT_MS = 30_000; +export const CODEX_OAUTH_START_TIMEOUT_MS = 30_000; +export const CODEX_OAUTH_REFRESH_LOCK_TIMEOUT_MS = 45_000; +export const CODEX_OAUTH_REFRESH_LOCK_STALE_MS = 60_000; +export const CODEX_OAUTH_ACCOUNT_ID_MAX_LENGTH = 200; +export const CODEX_OAUTH_ACCOUNT_LABEL_MAX_LENGTH = 100; +export const CODEX_OAUTH_ACCOUNT_ID_PATTERN = /^[a-zA-Z0-9_-]+$/; +export const CODEX_OAUTH_RESERVED_ACCOUNT_IDS = new Set(["__proto__", "constructor", "prototype"]); diff --git a/src/common/orpc/schemas.ts b/src/common/orpc/schemas.ts index 5d49b6efd6a..499027343cd 100644 --- a/src/common/orpc/schemas.ts +++ b/src/common/orpc/schemas.ts @@ -138,6 +138,7 @@ export { StreamEndEventSchema, StreamErrorMessageSchema, StreamStartEventSchema, + StreamModelUpdateEventSchema, ToolCallDeltaEventSchema, ToolCallEndEventSchema, ToolCallExecutionStartEventSchema, diff --git a/src/common/orpc/schemas/api.test.ts b/src/common/orpc/schemas/api.test.ts index abee162fbf3..f7099644e1e 100644 --- a/src/common/orpc/schemas/api.test.ts +++ b/src/common/orpc/schemas/api.test.ts @@ -4,6 +4,8 @@ import { ProviderConfigInfoSchema, ProvidersConfigMapSchema, config, + codexOauth, + projects, workspace, } from "./api"; import type { AWSCredentialStatus, ProviderConfigInfo, ProvidersConfigMap } from "../types"; @@ -18,6 +20,35 @@ import type { AWSCredentialStatus, ProviderConfigInfo, ProvidersConfigMap } from * If these tests fail, it means the schema is missing fields that the backend * service returns, which would cause data loss when crossing the IPC boundary. */ +describe("Codex account input validation", () => { + it.each(["__proto__", "constructor", "prototype", "", "work/account"])( + "rejects unsafe slot ID %s at each mutation boundary", + (accountId) => { + expect(codexOauth.startDesktopFlow.input.safeParse({ accountId }).success).toBe(false); + expect(codexOauth.startDeviceFlow.input.safeParse({ accountId }).success).toBe(false); + expect(codexOauth.disconnect.input.safeParse({ accountId }).success).toBe(false); + expect(codexOauth.setDefaultAccount.input.safeParse({ accountId }).success).toBe(false); + expect(codexOauth.renameAccount.input.safeParse({ accountId, label: "Work" }).success).toBe( + false + ); + expect( + projects.setCodexOauthAccount.input.safeParse({ projectPath: "/project", accountId }) + .success + ).toBe(false); + } + ); + + it("allows legacy login and explicit project inheritance", () => { + expect(codexOauth.startDesktopFlow.input.safeParse(undefined).success).toBe(true); + expect(codexOauth.startDeviceFlow.input.safeParse(undefined).success).toBe(true); + expect( + projects.setCodexOauthAccount.input.safeParse({ projectPath: "/project", accountId: null }) + .success + ).toBe(true); + expect(codexOauth.startDesktopFlow.input.safeParse({ label: " " }).success).toBe(false); + }); +}); + describe("ProviderConfigInfoSchema conformance", () => { it("preserves all AWSCredentialStatus fields", () => { const full: AWSCredentialStatus = { @@ -118,6 +149,8 @@ describe("ProviderConfigInfoSchema conformance", () => { cacheTtl: "1h", disableBetaFeatures: true, codexOauthSet: true, + codexOauthAccounts: [{ id: "work", label: "Work" }], + codexOauthDefaultAccountId: "work", codexOauthDefaultAuth: "apiKey", aws: { region: "ap-northeast-1", diff --git a/src/common/orpc/schemas/api.ts b/src/common/orpc/schemas/api.ts index 5b4d8b73f6b..cf7c3e6d12c 100644 --- a/src/common/orpc/schemas/api.ts +++ b/src/common/orpc/schemas/api.ts @@ -6,6 +6,12 @@ import { import { eventIterator } from "@orpc/server"; import { UIModeSchema } from "../../types/mode"; import { z } from "zod"; +import { + CODEX_OAUTH_ACCOUNT_ID_MAX_LENGTH, + CODEX_OAUTH_ACCOUNT_LABEL_MAX_LENGTH, + CODEX_OAUTH_ACCOUNT_ID_PATTERN, + CODEX_OAUTH_RESERVED_ACCOUNT_IDS, +} from "@/common/constants/codexOauthAccounts"; import { CODER_ARCHIVE_BEHAVIORS } from "@/common/config/coderArchiveBehavior"; import { WORKTREE_ARCHIVE_BEHAVIORS } from "@/common/config/worktreeArchiveBehavior"; import { HEARTBEAT_MAX_INTERVAL_MS, HEARTBEAT_MIN_INTERVAL_MS } from "@/constants/heartbeat"; @@ -288,8 +294,15 @@ export const ProviderConfigInfoSchema = z.object({ /** Anthropic-specific fields */ cacheTtl: CacheTtlSchema.optional(), disableBetaFeatures: z.boolean().optional(), - /** OpenAI-only: whether Codex OAuth tokens are present in providers.jsonc */ + /** OpenAI-only: whether usable Codex OAuth credentials exist. */ codexOauthSet: z.boolean().optional(), + /** Account identities remain available for reconnect. Credentials stay in the backend. */ + codexOauthAccounts: z + .array( + z.object({ id: z.string(), label: z.string(), reconnectRequired: z.boolean().optional() }) + ) + .optional(), + codexOauthDefaultAccountId: z.string().optional(), /** * OpenAI-only: default auth precedence to use for Codex-OAuth-allowed models when BOTH * ChatGPT OAuth and an OpenAI API key are configured. @@ -540,9 +553,28 @@ export const muxGovernorOauth = { }; // Codex OAuth (ChatGPT subscription auth) +const CodexOauthAccountIdSchema = z + .string() + .min(1) + .max(CODEX_OAUTH_ACCOUNT_ID_MAX_LENGTH) + .regex(CODEX_OAUTH_ACCOUNT_ID_PATTERN) + .refine((id) => !CODEX_OAUTH_RESERVED_ACCOUNT_IDS.has(id), "Invalid account ID"); +const CodexOauthAccountLabelSchema = z + .string() + .trim() + .min(1) + .max(CODEX_OAUTH_ACCOUNT_LABEL_MAX_LENGTH); +const CodexOauthLoginInputSchema = z + .object({ + label: CodexOauthAccountLabelSchema.optional(), + accountId: CodexOauthAccountIdSchema.optional(), + }) + .strict() + .optional(); + export const codexOauth = { startDesktopFlow: { - input: z.void(), + input: CodexOauthLoginInputSchema, output: ResultSchema(z.object({ flowId: z.string(), authorizeUrl: z.string() }), z.string()), }, waitForDesktopFlow: { @@ -559,7 +591,7 @@ export const codexOauth = { output: z.void(), }, startDeviceFlow: { - input: z.void(), + input: CodexOauthLoginInputSchema, output: ResultSchema( z.object({ flowId: z.string(), @@ -584,7 +616,17 @@ export const codexOauth = { output: z.void(), }, disconnect: { - input: z.void(), + input: z.object({ accountId: CodexOauthAccountIdSchema }).strict().optional(), + output: ResultSchema(z.void(), z.string()), + }, + setDefaultAccount: { + input: z.object({ accountId: CodexOauthAccountIdSchema }).strict(), + output: ResultSchema(z.void(), z.string()), + }, + renameAccount: { + input: z + .object({ accountId: CodexOauthAccountIdSchema, label: CodexOauthAccountLabelSchema }) + .strict(), output: ResultSchema(z.void(), z.string()), }, }; @@ -835,6 +877,12 @@ export const projects = { .passthrough(), output: z.void(), }, + setCodexOauthAccount: { + input: z + .object({ projectPath: z.string(), accountId: CodexOauthAccountIdSchema.nullable() }) + .strict(), + output: ResultSchema(z.void(), z.string()), + }, setCustomInstructions: { input: z .object({ @@ -2279,6 +2327,7 @@ export const nameGeneration = { generate: { input: z.object({ message: z.string(), + projectPath: z.string().optional(), /** Ordered list of model candidates to try (backend resolves gateway routing in createModel) */ candidates: z.array(z.string()), }), diff --git a/src/common/orpc/schemas/chatStats.ts b/src/common/orpc/schemas/chatStats.ts index ebd0cf4496b..58f9ff965b7 100644 --- a/src/common/orpc/schemas/chatStats.ts +++ b/src/common/orpc/schemas/chatStats.ts @@ -32,6 +32,7 @@ export const ChatUsageDisplaySchema = z.object({ output: ChatUsageComponentSchema, reasoning: ChatUsageComponentSchema, model: z.string().optional(), + effectiveContextLimit: z.number().positive().nullable().optional(), costsIncluded: z.boolean().optional(), }); diff --git a/src/common/orpc/schemas/stream.ts b/src/common/orpc/schemas/stream.ts index 63d44f311fa..ece44811a86 100644 --- a/src/common/orpc/schemas/stream.ts +++ b/src/common/orpc/schemas/stream.ts @@ -166,6 +166,9 @@ export const StreamStartEventSchema = z.object({ .optional() .meta({ description: "True when this event is emitted during stream replay" }), model: z.string(), + // Publish the accepted limit before usage arrives. Null means the limit is unknown. + effectiveContextLimit: z.number().positive().nullable().optional(), + modelFallback: ModelFallbackRecordSchema.optional(), metadataModel: z .string() .optional() @@ -198,6 +201,26 @@ export const StreamStartEventSchema = z.object({ .meta({ description: "ACP prompt correlation id for matching stream events" }), }); +// Replaces attempt metadata without starting a new stream. Omitted optional fields clear prior values. +export const StreamModelUpdateEventSchema = StreamStartEventSchema.pick({ + workspaceId: true, + messageId: true, + model: true, + metadataModel: true, + effectiveContextLimit: true, + modelFallback: true, + routedThroughGateway: true, + routeProvider: true, + thinkingLevel: true, +}) + .required({ + metadataModel: true, + effectiveContextLimit: true, + modelFallback: true, + routedThroughGateway: true, + }) + .extend({ type: z.literal("stream-model-update") }); + export const StreamDeltaEventSchema = z.object({ type: z.literal("stream-delta"), workspaceId: z.string(), @@ -564,6 +587,9 @@ export const UsageDeltaEventSchema = z.object({ .optional() .meta({ description: "True when this event is emitted during stream replay" }), + // Safe numeric limit from the accepted request. Never expose its routing snapshot. + effectiveContextLimit: z.number().positive().nullable().optional(), + // Step-level: this step only (for context window display) usage: LanguageModelV2UsageSchema, providerMetadata: z.record(z.string(), z.unknown()).optional(), @@ -687,6 +713,7 @@ export const WorkspaceChatMessageSchema = z.discriminatedUnion("type", [ DeleteMessageSchema, StreamLifecycleEventSchema, StreamStartEventSchema, + StreamModelUpdateEventSchema, StreamDeltaEventSchema, StreamEndEventSchema, StreamAbortEventSchema, diff --git a/src/common/schemas/project.ts b/src/common/schemas/project.ts index bd1c33dc78e..15282f25e0b 100644 --- a/src/common/schemas/project.ts +++ b/src/common/schemas/project.ts @@ -285,6 +285,8 @@ export const WorkspaceConfigSchema = z.object({ }); export const ProjectConfigSchema = z.object({ + /** Use this account for the project. Unset inherits the global default. */ + codexOauthAccountId: z.string().min(1).optional(), displayName: z.string().nullish().meta({ description: "Custom display name for the project", }), diff --git a/src/common/types/stream.ts b/src/common/types/stream.ts index 223e58df8fe..d603b9a4c09 100644 --- a/src/common/types/stream.ts +++ b/src/common/types/stream.ts @@ -19,6 +19,7 @@ import type { StreamDeltaEventSchema, StreamEndEventSchema, StreamStartEventSchema, + StreamModelUpdateEventSchema, ToolCallDeltaEventSchema, ToolCallEndEventSchema, ToolCallExecutionStartEventSchema, @@ -40,6 +41,7 @@ import type { export type CompletedMessagePart = MuxReasoningPart | MuxTextPart | MuxToolPart; export type StreamStartEvent = z.infer; +export type StreamModelUpdateEvent = z.infer; export type StreamDeltaEvent = z.infer; export type StreamEndEvent = z.infer; export type StreamAbortReason = z.infer; diff --git a/src/common/utils/ai/cacheStrategy.ts b/src/common/utils/ai/cacheStrategy.ts index 32ae7920c41..d3c7f03a19c 100644 --- a/src/common/utils/ai/cacheStrategy.ts +++ b/src/common/utils/ai/cacheStrategy.ts @@ -4,7 +4,7 @@ import { isGpt56FamilyModel } from "@/common/types/thinking"; import assert from "@/common/utils/assert"; import { cloneToolPreservingDescriptors } from "@/common/utils/tools/cloneToolPreservingDescriptors"; import { - wouldRouteOpenAIThroughCodexOauth, + resolveCodexOauthRouting, type CodexOauthRoutingOptions, } from "@/common/utils/providers/codexOauthRouting"; import { resolveCoderWireCanonicalModel } from "@/common/constants/coderOAuth"; @@ -250,7 +250,7 @@ function isOfficialOpenAIBaseUrl(baseUrl: string): boolean { * capability target must itself be an OpenAI GPT-5.6-family model; * - the backend-resolved route provider must be exactly "openai" — missing, * legacy, gateway, or unknown route metadata fails closed; - * - Codex OAuth precedence (mirrored by wouldRouteOpenAIThroughCodexOauth) + * - Codex OAuth precedence (mirrored by resolveCodexOauthRouting) * fails closed because the ChatGPT backend strips these fields; * - a configured custom base URL fails closed unless it is the official * endpoint. Transport-level HTTP proxy env vars are not endpoint overrides @@ -297,7 +297,7 @@ export function openaiExplicitPromptCachingAvailable( return false; } - if (wouldRouteOpenAIThroughCodexOauth(normalized, providersConfig, options)) { + if (resolveCodexOauthRouting(normalized, providersConfig, options) !== "other") { return false; } diff --git a/src/common/utils/ai/openaiProviderOptionsAvailability.ts b/src/common/utils/ai/openaiProviderOptionsAvailability.ts index 9c69028fc4d..02f9d12ec71 100644 --- a/src/common/utils/ai/openaiProviderOptionsAvailability.ts +++ b/src/common/utils/ai/openaiProviderOptionsAvailability.ts @@ -6,7 +6,7 @@ import type { ProvidersConfigMap } from "@/common/orpc/types"; import type { OpenAIWireFormat } from "@/common/types/providerOptions"; import { PROVIDER_DEFINITIONS } from "@/common/constants/providers"; import { getExplicitGatewayPrefix, normalizeToCanonical } from "@/common/utils/ai/models"; -import { wouldRouteOpenAIThroughCodexOauth } from "@/common/utils/providers/codexOauthRouting"; +import { resolveCodexOauthRouting } from "@/common/utils/providers/codexOauthRouting"; export interface OpenAIDirectProviderOptionsAvailability { /** Settings-resolved route for the canonical model ("direct" = no gateway). */ @@ -15,6 +15,7 @@ export interface OpenAIDirectProviderOptionsAvailability { providersConfig?: ProvidersConfigMap | null; /** Request-level OpenAI wire format; the stored config value wins when set. */ openaiWireFormat?: OpenAIWireFormat | null; + codexOauthAccountId?: string; } export function openaiDirectProviderOptionsAvailable( @@ -53,8 +54,9 @@ export function openaiDirectProviderOptionsAvailable( // API-only provider options, so toggles for those options must fail closed. return !( options?.providersConfig != null && - wouldRouteOpenAIThroughCodexOauth(normalized, options.providersConfig, { + resolveCodexOauthRouting(normalized, options.providersConfig, { openaiWireFormat: options.openaiWireFormat, - }) + codexOauthAccountId: options.codexOauthAccountId, + }) !== "other" ); } diff --git a/src/common/utils/compaction/contextLimit.ts b/src/common/utils/compaction/contextLimit.ts index 55929e4e7a6..f1effa7b198 100644 --- a/src/common/utils/compaction/contextLimit.ts +++ b/src/common/utils/compaction/contextLimit.ts @@ -10,7 +10,7 @@ import { getCodexOauthContextWindowOverride, } from "@/common/constants/codexOAuth"; import { - wouldRouteOpenAIThroughCodexOauth, + resolveCodexOauthRouting, type CodexOauthRoutingOptions, } from "@/common/utils/providers/codexOauthRouting"; import type { ProvidersConfigMap } from "@/common/orpc/types"; @@ -36,7 +36,7 @@ function getCodexOauthContextLimit( return null; } - return wouldRouteOpenAIThroughCodexOauth(model, providersConfig, options) ? oauthLimit : null; + return resolveCodexOauthRouting(model, providersConfig, options) !== "other" ? oauthLimit : null; } /** diff --git a/src/common/utils/providers/codexOauthRouting.test.ts b/src/common/utils/providers/codexOauthRouting.test.ts new file mode 100644 index 00000000000..0ca53e85c2f --- /dev/null +++ b/src/common/utils/providers/codexOauthRouting.test.ts @@ -0,0 +1,196 @@ +import { describe, expect, it } from "bun:test"; +import type { ProvidersConfigMap } from "@/common/orpc/types"; +import { + getCodexOauthProjectPath, + hasCodexOauthTokens, + resolveCodexOauthRouting, + wouldRouteOpenAIThroughCodexOauth, +} from "./codexOauthRouting"; +import { getEffectiveContextLimit } from "@/common/utils/compaction/contextLimit"; +import { openaiProModeAvailable } from "@/common/utils/ai/proMode"; +import { openaiExplicitPromptCachingAvailable } from "@/common/utils/ai/cacheStrategy"; +import { getFastModeProvider } from "@/browser/utils/fastModeServiceTier"; +import { openaiDirectProviderOptionsAvailable } from "@/common/utils/ai/openaiProviderOptionsAvailability"; + +const auth = { type: "oauth", access: "access", refresh: "refresh", expires: 1000 }; + +describe("Codex OAuth account routing", () => { + it("uses multi-project attribution before subproject and root scope", () => { + const scope = { + projectPath: "/root", + attributionProjectPath: "/attribution", + subProjectPath: "/root/sub", + projects: [ + { projectPath: "/first", projectName: "first" }, + { projectPath: "/second", projectName: "second" }, + ], + }; + expect(getCodexOauthProjectPath(scope)).toBe("/first"); + expect(getCodexOauthProjectPath({ ...scope, projects: [] })).toBe("/root/sub"); + expect(getCodexOauthProjectPath({ ...scope, projects: [], subProjectPath: undefined })).toBe( + "/attribution" + ); + expect(getCodexOauthProjectPath({ projectPath: "/root" })).toBe("/root"); + expect( + getCodexOauthProjectPath({ projectPath: "/root", subProjectPath: "/root/new-subproject" }) + ).toBe("/root/new-subproject"); + expect(getCodexOauthProjectPath()).toBeUndefined(); + }); + + it.each(["project", "global", "implicit"])( + "fails closed for a missing %s OAuth selection", + (source) => { + const providersConfig: ProvidersConfigMap = { + openai: { + apiKeySet: true, + isConfigured: true, + isEnabled: true, + codexOauthAccounts: [{ id: "work", label: "Work" }], + codexOauthSet: true, + ...(source === "global" ? { codexOauthDefaultAccountId: "deleted" } : {}), + }, + }; + const options = { + providersConfig, + ...(source === "project" ? { codexOauthAccountId: "deleted" } : {}), + }; + const model = "openai:gpt-5.6-sol"; + expect(resolveCodexOauthRouting(model, providersConfig, options)).toBe("missing-account"); + expect(wouldRouteOpenAIThroughCodexOauth(model, providersConfig, options)).toBe(false); + expect(openaiProModeAvailable(model, options)).toBe(false); + expect(getFastModeProvider(model, options)).toBeNull(); + expect(openaiExplicitPromptCachingAvailable(model, "openai", providersConfig, options)).toBe( + false + ); + expect(getEffectiveContextLimit(model, false, providersConfig, options)).toBe(372_000); + expect( + resolveCodexOauthRouting(model, providersConfig, { + ...options, + codexOauthAccountId: "work", + }) + ).toBe("oauth"); + + providersConfig.openai.codexOauthDefaultAuth = "apiKey"; + expect(resolveCodexOauthRouting(model, providersConfig, options)).toBe("other"); + expect(openaiProModeAvailable(model, options)).toBe(true); + expect(getFastModeProvider(model, options)).toBe("openai"); + expect(getEffectiveContextLimit(model, false, providersConfig, options)).toBeGreaterThan( + 372_000 + ); + // Required Codex models still reject the missing selected slot. + expect(resolveCodexOauthRouting("openai:gpt-5.3-codex-spark", providersConfig, options)).toBe( + "missing-account" + ); + } + ); + + it("keeps legacy API keys and gateway routes independent from missing selections", () => { + const providersConfig: ProvidersConfigMap = { + openai: { apiKeySet: true, isConfigured: true, isEnabled: true }, + }; + const model = "openai:gpt-5.6-sol"; + expect(resolveCodexOauthRouting(model, providersConfig)).toBe("other"); + const options = { + providersConfig, + codexOauthAccountId: "deleted", + openaiWireFormat: "chatCompletions" as const, + }; + expect(resolveCodexOauthRouting(model, providersConfig, options)).toBe("other"); + expect( + getFastModeProvider(model, { ...options, resolvedRouteProvider: "openrouter" }) + ).toBeNull(); + expect( + resolveCodexOauthRouting("openrouter:openai/gpt-5.6-sol", providersConfig, options) + ).toBe("other"); + expect( + getEffectiveContextLimit("openrouter:openai/gpt-5.6-sol", false, providersConfig, options) + ).toBeGreaterThan(372_000); + providersConfig.openai.codexOauthDefaultAccountId = "deleted"; + expect(resolveCodexOauthRouting(model, providersConfig)).toBe("missing-account"); + }); + + it("keeps legacy metadata and raw token defaults compatible", () => { + for (const config of [{ codexOauthSet: true }, { codexOauth: auth }]) { + expect(hasCodexOauthTokens(config)).toBe(true); + expect(hasCodexOauthTokens(config, "default")).toBe(true); + expect(hasCodexOauthTokens(config, "missing")).toBe(false); + } + }); + + it("keeps an invalid implicit raw slot on the reconnect path", () => { + const providersConfig = { + openai: { + apiKeySet: true, + isConfigured: true, + isEnabled: true, + codexOauth: { ...auth, invalidReason: "invalid_grant" }, + }, + }; + const model = "openai:gpt-5.6-sol"; + expect(hasCodexOauthTokens(providersConfig.openai)).toBe(false); + expect(resolveCodexOauthRouting(model, providersConfig)).toBe("missing-account"); + expect(getEffectiveContextLimit(model, false, providersConfig)).toBe(372_000); + const restoredConfig = { openai: { ...providersConfig.openai, codexOauth: auth } }; + expect(resolveCodexOauthRouting(model, restoredConfig)).toBe("oauth"); + }); + + it("uses metadata account IDs instead of the aggregate connection flag", () => { + const config = { + codexOauthSet: true, + codexOauthAccounts: [{ id: "work", label: "Work" }], + codexOauthDefaultAccountId: "missing", + }; + expect(hasCodexOauthTokens(config)).toBe(false); + expect(hasCodexOauthTokens(config, "default")).toBe(false); + expect(hasCodexOauthTokens(config, "work")).toBe(true); + expect(hasCodexOauthTokens({ ...config, codexOauthDefaultAccountId: "work" })).toBe(true); + expect(hasCodexOauthTokens({ ...config, codexOauthAccounts: [] }, "work")).toBe(false); + }); + + it("selects raw account maps without substituting legacy tokens", () => { + const config = { + codexOauth: auth, + codexOauthAccounts: { + work: { label: "Work", credentials: auth }, + invalid: { label: "Invalid", credentials: { ...auth, refresh: "" } }, + }, + codexOauthDefaultAccountId: "work", + }; + expect(hasCodexOauthTokens(config)).toBe(true); + expect(hasCodexOauthTokens(config, "default")).toBe(true); + expect(hasCodexOauthTokens(config, "missing")).toBe(false); + expect(hasCodexOauthTokens(config, "invalid")).toBe(false); + expect(hasCodexOauthTokens({ ...config, codexOauthAccounts: {} })).toBe(false); + }); + + it("threads selected accounts into direct provider option availability", () => { + const providersConfig: ProvidersConfigMap = { + openai: { + apiKeySet: true, + isConfigured: true, + isEnabled: true, + codexOauthSet: true, + codexOauthDefaultAccountId: "deleted", + codexOauthAccounts: [{ id: "work", label: "Work" }], + }, + }; + const model = "openai:gpt-5.5"; + expect(wouldRouteOpenAIThroughCodexOauth(model, providersConfig)).toBe(false); + expect( + wouldRouteOpenAIThroughCodexOauth(model, providersConfig, { codexOauthAccountId: "work" }) + ).toBe(true); + expect( + openaiDirectProviderOptionsAvailable(model, { providersConfig, codexOauthAccountId: "work" }) + ).toBe(false); + expect( + wouldRouteOpenAIThroughCodexOauth(model, providersConfig, { + codexOauthAccountId: "work", + openaiWireFormat: "chatCompletions", + }) + ).toBe(false); + providersConfig.openai.codexOauthDefaultAuth = "apiKey"; + expect( + wouldRouteOpenAIThroughCodexOauth(model, providersConfig, { codexOauthAccountId: "work" }) + ).toBe(false); + }); +}); diff --git a/src/common/utils/providers/codexOauthRouting.ts b/src/common/utils/providers/codexOauthRouting.ts index ad777860236..03d87725c93 100644 --- a/src/common/utils/providers/codexOauthRouting.ts +++ b/src/common/utils/providers/codexOauthRouting.ts @@ -12,6 +12,8 @@ import { isCodexOauthAllowedModel, isCodexOauthRequiredModel } from "@/common/constants/codexOAuth"; import type { ProvidersConfigMap } from "@/common/orpc/types"; import type { OpenAIWireFormat } from "@/common/types/providerOptions"; +import type { WorkspaceMetadata } from "@/common/types/workspace"; +import { CODEX_OAUTH_DEFAULT_ACCOUNT_ID } from "@/common/constants/codexOauthAccounts"; /** Request-level inputs the stored providers config cannot carry. */ export interface CodexOauthRoutingOptions { @@ -20,6 +22,24 @@ export interface CodexOauthRoutingOptions { * The stored `openai.wireFormat` wins when set, matching providerModelFactory. */ openaiWireFormat?: OpenAIWireFormat | null; + /** Local account slot selected for this request. */ + codexOauthAccountId?: string; +} + +/** Use the same account scope as backend model construction. */ +export function getCodexOauthProjectPath( + scope?: + | (Partial> & { + attributionProjectPath?: string; + }) + | null +): string | undefined { + return ( + scope?.projects?.[0]?.projectPath ?? + scope?.subProjectPath ?? + scope?.attributionProjectPath ?? + scope?.projectPath + ); } function asRecord(value: unknown): Record | null { @@ -33,25 +53,43 @@ function hasNonEmptyString(value: unknown): value is string { return typeof value === "string" && value.trim().length > 0; } -export function hasCodexOauthTokens(config: unknown): boolean { +export function hasCodexOauthTokens(config: unknown, accountId?: string): boolean { const record = asRecord(config); if (!record) { return false; } - if (record.codexOauthSet === true) { + const selectedId = + accountId ?? record.codexOauthDefaultAccountId ?? CODEX_OAUTH_DEFAULT_ACCOUNT_ID; + if (Array.isArray(record.codexOauthAccounts)) { + return record.codexOauthAccounts.some((account: unknown) => { + const entry = asRecord(account); + return entry?.id === selectedId && entry.reconnectRequired !== true; + }); + } + + // Old metadata contains only the legacy connection flag. + if (record.codexOauthSet === true && selectedId === CODEX_OAUTH_DEFAULT_ACCOUNT_ID) { return true; } - // Backend compaction can receive raw providers.jsonc config in older tests/fallback paths. - // Detect the stored token shape without importing node-only OAuth parsing into common code. - const oauth = asRecord(record.codexOauth); + // Raw configs contain tokens. Never substitute another connected account. + const accounts = asRecord(record.codexOauthAccounts); + const selectedAccount = typeof selectedId === "string" ? asRecord(accounts?.[selectedId]) : null; + if (selectedId !== CODEX_OAUTH_DEFAULT_ACCOUNT_ID && !hasNonEmptyString(selectedAccount?.label)) { + return false; + } + const oauth = asRecord( + selectedId === CODEX_OAUTH_DEFAULT_ACCOUNT_ID ? record.codexOauth : selectedAccount?.credentials + ); return ( oauth?.type === "oauth" && + oauth.invalidReason === undefined && hasNonEmptyString(oauth.access) && hasNonEmptyString(oauth.refresh) && typeof oauth.expires === "number" && - Number.isFinite(oauth.expires) + Number.isFinite(oauth.expires) && + (oauth.accountId === undefined || hasNonEmptyString(oauth.accountId)) ); } @@ -69,38 +107,60 @@ export function hasOpenAIApiKey(config: unknown): boolean { return record.apiKeySet === true || hasNonEmptyString(record.apiKey); } -/** - * Would a direct-OpenAI request for this model route through Codex OAuth? - * - * Mirrors providerModelFactory: allowed model + stored OAuth tokens, then - * Chat Completions with an API key never routes OAuth, required models always - * route OAuth; otherwise OAuth wins when no API key is configured or when - * `codexOauthDefaultAuth` prefers OAuth over a present key. - */ -export function wouldRouteOpenAIThroughCodexOauth( +export type CodexOauthRouting = "oauth" | "missing-account" | "other"; + +/** Resolve OAuth routing without treating a missing selection as API-key fallback. */ +export function resolveCodexOauthRouting( model: string, providersConfig: ProvidersConfigMap | null | undefined, options?: CodexOauthRoutingOptions -): boolean { +): CodexOauthRouting { const openAIConfig = providersConfig?.openai; if (!isCodexOauthAllowedModel(model, providersConfig ?? null)) { - return false; + return "other"; } - if (!hasCodexOauthTokens(openAIConfig)) { - return false; + const record = asRecord(openAIConfig); + const hasApiKey = hasOpenAIApiKey(openAIConfig); + const hasSelectedAccount = hasCodexOauthTokens(openAIConfig, options?.codexOauthAccountId); + const required = isCodexOauthRequiredModel(model, providersConfig ?? null); + if (required && !hasSelectedAccount && !hasApiKey) { + return "missing-account"; } - // Codex OAuth serves only the Responses API. With Chat Completions selected, - // the factory falls back to the API key whenever one exists. - const wireFormat = asRecord(openAIConfig)?.wireFormat ?? options?.openaiWireFormat; - if (wireFormat === "chatCompletions" && hasOpenAIApiKey(openAIConfig)) { - return false; + // Chat Completions uses the API key, even when OAuth is preferred. + const wireFormat = record?.wireFormat ?? options?.openaiWireFormat; + if (wireFormat === "chatCompletions" && hasApiKey) { + return "other"; } - if (isCodexOauthRequiredModel(model, providersConfig ?? null)) { - return true; + if (!required && hasApiKey && record?.codexOauthDefaultAuth === "apiKey") { + return "other"; } - if (!hasOpenAIApiKey(openAIConfig)) { - return true; + if (hasSelectedAccount) { + return "oauth"; } - return asRecord(openAIConfig)?.codexOauthDefaultAuth !== "apiKey"; + const hasExplicitSelection = + options?.codexOauthAccountId !== undefined || record?.codexOauthDefaultAccountId !== undefined; + const accounts = record?.codexOauthAccounts; + const hasInvalidStoredAccount = + asRecord(record?.codexOauth)?.invalidReason === "invalid_grant" || + Object.values(asRecord(accounts) ?? {}).some( + (account) => asRecord(asRecord(account)?.credentials)?.invalidReason === "invalid_grant" + ); + const hasAccountSlots = Array.isArray(accounts) + ? accounts.length > 0 + : record?.codexOauthSet === true || + hasInvalidStoredAccount || + hasCodexOauthTokens(openAIConfig, CODEX_OAUTH_DEFAULT_ACCOUNT_ID) || + Object.keys(asRecord(accounts) ?? {}).some((id) => hasCodexOauthTokens(openAIConfig, id)); + // Missing slots must not enable API-only controls or remove the OAuth context cap. + return hasExplicitSelection || hasAccountSlots ? "missing-account" : "other"; +} + +/** Return true only when the selected OAuth account can serve the request. */ +export function wouldRouteOpenAIThroughCodexOauth( + model: string, + providersConfig: ProvidersConfigMap | null | undefined, + options?: CodexOauthRoutingOptions +): boolean { + return resolveCodexOauthRouting(model, providersConfig, options) === "oauth"; } diff --git a/src/common/utils/tokens/tokenMeterUtils.test.ts b/src/common/utils/tokens/tokenMeterUtils.test.ts index 5531640f97c..9e55ae6037c 100644 --- a/src/common/utils/tokens/tokenMeterUtils.test.ts +++ b/src/common/utils/tokens/tokenMeterUtils.test.ts @@ -42,6 +42,66 @@ describe("calculateTokenMeterData", () => { }, }; + test("keeps the accepted limit during live usage and uses current settings when idle", () => { + const changedConfig: ProvidersConfigMap = { + openai: { + apiKeySet: true, + isEnabled: true, + isConfigured: true, + codexOauthDefaultAuth: "apiKey", + models: [{ id: "gpt-5.5", contextWindowTokens: 500_000 }], + }, + }; + const live = calculateTokenMeterData( + SAMPLE_USAGE, + "openai:gpt-5.5", + false, + false, + changedConfig, + { codexOauthAccountId: "different-account" }, + 272_000 + ); + const idle = calculateTokenMeterData( + SAMPLE_USAGE, + "openai:gpt-5.5", + false, + false, + changedConfig + ); + expect(live.maxTokens).toBe(272_000); + expect(live.totalPercentage).toBeCloseTo((11_000 / 272_000) * 100); + expect(idle.maxTokens).toBe(500_000); + expect(idle.totalPercentage).toBeCloseTo(2.2); + }); + + test.each([272_000, null])("keeps the accepted limit before usage arrives: %s", (limit) => { + const result = calculateTokenMeterData( + undefined, + "anthropic:claude-sonnet-4-20250514", + true, + false, + providerConfigWithOverride, + undefined, + limit + ); + expect(result.maxTokens).toBe(limit ?? undefined); + expect(result.totalTokens).toBe(0); + expect(result.totalPercentage).toBe(0); + }); + + test("keeps an unknown accepted limit despite a current model override", () => { + const result = calculateTokenMeterData( + SAMPLE_USAGE, + "anthropic:claude-sonnet-4-20250514", + true, + false, + providerConfigWithOverride, + undefined, + null + ); + expect(result.maxTokens).toBeUndefined(); + }); + test("uses custom context override for beta Sonnet models", () => { const result = calculateTokenMeterData( SAMPLE_USAGE, @@ -76,6 +136,37 @@ describe("calculateTokenMeterData", () => { expect(result.totalPercentage).toBeCloseTo(1.1); }); + test("keeps the OAuth cap for an unavailable global account and a connected project account", () => { + const providersConfig: ProvidersConfigMap = { + openai: { + apiKeySet: true, + isEnabled: true, + isConfigured: true, + codexOauthSet: true, + codexOauthDefaultAccountId: "missing", + codexOauthAccounts: [{ id: "work", label: "Work" }], + }, + }; + const globalMeter = calculateTokenMeterData( + SAMPLE_USAGE, + "openai:gpt-5.5", + false, + false, + providersConfig + ); + const projectMeter = calculateTokenMeterData( + SAMPLE_USAGE, + "openai:gpt-5.5", + false, + false, + providersConfig, + { codexOauthAccountId: "work" } + ); + expect(projectMeter.maxTokens).toBe(272_000); + expect(globalMeter.maxTokens).toBe(projectMeter.maxTokens); + expect(projectMeter.totalPercentage).toBe(globalMeter.totalPercentage); + }); + test("uses the Codex OAuth cap for GPT-5.5 token meter percentages", () => { const result = calculateTokenMeterData(SAMPLE_USAGE, "openai:gpt-5.5", false, false, { openai: { diff --git a/src/common/utils/tokens/tokenMeterUtils.ts b/src/common/utils/tokens/tokenMeterUtils.ts index 77b542d9e43..dff4a185a95 100644 --- a/src/common/utils/tokens/tokenMeterUtils.ts +++ b/src/common/utils/tokens/tokenMeterUtils.ts @@ -1,5 +1,6 @@ import type { ProvidersConfigMap } from "@/common/orpc/types"; import { getEffectiveContextLimit } from "@/common/utils/compaction/contextLimit"; +import type { CodexOauthRoutingOptions } from "@/common/utils/providers/codexOauthRouting"; import type { ChatUsageDisplay } from "./usageAggregator"; // NOTE: Provide theme-matching fallbacks so token meters render consistently @@ -60,11 +61,24 @@ export function calculateTokenMeterData( model: string, use1M: boolean, verticalProportions = false, - providersConfig: ProvidersConfigMap | null = null + providersConfig: ProvidersConfigMap | null = null, + routingOptions?: CodexOauthRoutingOptions, + effectiveContextLimit?: number | null ): TokenMeterData { - if (!usage) return { segments: [], totalTokens: 0, totalPercentage: 0 }; + if (!usage) { + return { + segments: [], + totalTokens: 0, + totalPercentage: 0, + maxTokens: effectiveContextLimit ?? undefined, + }; + } - const maxTokens = getEffectiveContextLimit(model, use1M, providersConfig) ?? undefined; + // Live usage keeps the accepted request limit. Idle callers omit the override. + const maxTokens = + (effectiveContextLimit !== undefined + ? effectiveContextLimit + : getEffectiveContextLimit(model, use1M, providersConfig, routingOptions)) ?? undefined; // Total tokens used in the request. // For Anthropic prompt caching, cacheCreate tokens are reported separately but still diff --git a/src/common/utils/tokens/usageAggregator.ts b/src/common/utils/tokens/usageAggregator.ts index 50ca39a6d94..8c5601d9af5 100644 --- a/src/common/utils/tokens/usageAggregator.ts +++ b/src/common/utils/tokens/usageAggregator.ts @@ -30,6 +30,9 @@ export interface ChatUsageDisplay { // Optional model field for display purposes (context window calculation, etc.) model?: string; + // Backend limit for the active request. Null means the limit is unknown. + effectiveContextLimit?: number | null; + // True if any model in the sum had unknown pricing (costs are partial/incomplete) hasUnknownCosts?: boolean; diff --git a/src/node/config.test.ts b/src/node/config.test.ts index 4da8df62961..de7ad5a142e 100644 --- a/src/node/config.test.ts +++ b/src/node/config.test.ts @@ -337,6 +337,25 @@ describe("Config", () => { }); describe("loadConfigOrDefault customInstructions sanitizing", () => { + it("sanitizes malformed Codex selections but retains disconnected account IDs", () => { + fs.writeFileSync( + path.join(tempDir, "config.json"), + JSON.stringify({ + projects: [ + ["/home/user/number", { workspaces: [], codexOauthAccountId: 42 }], + ["/home/user/blank", { workspaces: [], codexOauthAccountId: " " }], + ["/home/user/disconnected", { workspaces: [], codexOauthAccountId: "removed-account" }], + ], + }) + ); + const loaded = config.loadConfigOrDefault(); + expect(loaded.projects.get("/home/user/number")?.codexOauthAccountId).toBeUndefined(); + expect(loaded.projects.get("/home/user/blank")?.codexOauthAccountId).toBeUndefined(); + expect(loaded.projects.get("/home/user/disconnected")?.codexOauthAccountId).toBe( + "removed-account" + ); + }); + it("discards malformed non-string customInstructions and keeps valid ones", () => { // A malformed value must not survive load: it would fail the // projects.list z.string() output schema and brick the project list. diff --git a/src/node/config/fileLeaseManager.ts b/src/node/config/fileLeaseManager.ts index 9a52bd05cb6..bc63e9d0d3e 100644 --- a/src/node/config/fileLeaseManager.ts +++ b/src/node/config/fileLeaseManager.ts @@ -2,6 +2,10 @@ import * as crypto from "crypto"; import * as fs from "fs"; import * as path from "path"; import { getXumHome } from "@/common/constants/paths"; +import { + CODEX_OAUTH_REFRESH_LOCK_TIMEOUT_MS, + CODEX_OAUTH_REFRESH_LOCK_STALE_MS, +} from "@/common/constants/codexOauthAccounts"; import { log } from "@/node/services/log"; import { ensurePrivateDirSync } from "@/node/utils/fs"; @@ -42,6 +46,18 @@ export class FileLeaseManager { return this.withDirLock(`${this.providersFile}.coder-refresh.lock`, 45_000, 60_000, fn); } + /** Serialize each Codex slot across processes before consuming its refresh token. */ + async withCodexOauthRefreshLock(accountId: string, fn: () => Promise | T): Promise { + // Hash slot IDs to keep lock paths bounded and independent of user input. + const slotKey = crypto.createHash("sha256").update(accountId).digest("hex"); + return this.withDirLock( + `${this.providersFile}.codex-refresh-${slotKey}.lock`, + CODEX_OAUTH_REFRESH_LOCK_TIMEOUT_MS, + CODEX_OAUTH_REFRESH_LOCK_STALE_MS, + fn + ); + } + /** Keeps each desktop login rollback snapshot anchored to committed credentials. */ async withCoderOauthLoginCommitLock(fn: () => Promise | T): Promise { return this.withDirLock(`${this.providersFile}.coder-login.lock`, 15_000, 20_000, fn); diff --git a/src/node/config/index.ts b/src/node/config/index.ts index 385d4a64181..f518eeb7768 100644 --- a/src/node/config/index.ts +++ b/src/node/config/index.ts @@ -722,6 +722,7 @@ function normalizeProjectRuntimeSettings(projectConfig: ProjectConfig): ProjectC projectKind?: unknown; customInstructions?: unknown; codeWorkspaceSyncPath?: unknown; + codexOauthAccountId?: unknown; }; const runtimeEnablement = normalizeRuntimeEnablementOverrides(record.runtimeEnablement); const defaultRuntime = normalizeRuntimeEnablementId(record.defaultRuntime); @@ -766,6 +767,13 @@ function normalizeProjectRuntimeSettings(projectConfig: ProjectConfig): ProjectC delete next.customInstructions; } + // Invalid disk values must not break the project list. Preserve missing account IDs for explicit recovery. + if (typeof record.codexOauthAccountId === "string" && record.codexOauthAccountId.trim()) { + next.codexOauthAccountId = record.codexOauthAccountId; + } else { + delete next.codexOauthAccountId; + } + // Same hand-edit hazard as customInstructions above. if (typeof record.codeWorkspaceSyncPath === "string" && record.codeWorkspaceSyncPath.trim()) { next.codeWorkspaceSyncPath = record.codeWorkspaceSyncPath; @@ -2932,6 +2940,7 @@ export class Config { workspacePath: string; projectPath: string; attributionProjectPath?: string; + subProjectPath?: string; projects?: Workspace["projects"]; workspaceName?: string; parentWorkspaceId?: string; @@ -2951,6 +2960,7 @@ export class Config { // config.projects.get(projectPath), even for multi-project workspaces under _multi. projectPath, attributionProjectPath, + subProjectPath: workspace.subProjectPath, projects: workspace.projects, workspaceName: workspace.name, parentWorkspaceId: workspace.parentWorkspaceId, @@ -2993,6 +3003,7 @@ export class Config { workspacePath: workspace.path, projectPath, attributionProjectPath, + subProjectPath: metadata.subProjectPath ?? workspace.subProjectPath, projects: metadata.projects ?? workspace.projects, workspaceName: undefined, parentWorkspaceId: undefined, @@ -3042,6 +3053,7 @@ export class Config { workspacePath: workspace.path, projectPath, attributionProjectPath, + subProjectPath: legacyMetadata.subProjectPath ?? workspace.subProjectPath, projects: legacyMetadata.projects ?? workspace.projects, workspaceName: undefined, parentWorkspaceId: undefined, @@ -3061,6 +3073,7 @@ export class Config { workspacePath: workspace.path, projectPath, attributionProjectPath, + subProjectPath: workspace.subProjectPath, projects: workspace.projects, workspaceName: undefined, parentWorkspaceId: undefined, diff --git a/src/node/orpc/router.ts b/src/node/orpc/router.ts index 02bfc7ae578..b8f2aa5ee66 100644 --- a/src/node/orpc/router.ts +++ b/src/node/orpc/router.ts @@ -776,8 +776,8 @@ export const router = (authToken?: string) => { .input(schemas.codexOauth.startDesktopFlow.input) .output(schemas.codexOauth.startDesktopFlow.output) .handler( - handlerGen(function* ({ context }) { - return yield* context.codexOauthService.startDesktopFlowEffect(); + handlerGen(function* ({ context }, input) { + return yield* context.codexOauthService.startDesktopFlowEffect(input); }) ), waitForDesktopFlow: t @@ -802,8 +802,8 @@ export const router = (authToken?: string) => { .input(schemas.codexOauth.startDeviceFlow.input) .output(schemas.codexOauth.startDeviceFlow.output) .handler( - handlerGen(function* ({ context }) { - return yield* context.codexOauthService.startDeviceFlowEffect(); + handlerGen(function* ({ context }, input) { + return yield* context.codexOauthService.startDeviceFlowEffect(input); }) ), waitForDeviceFlow: t @@ -828,8 +828,27 @@ export const router = (authToken?: string) => { .input(schemas.codexOauth.disconnect.input) .output(schemas.codexOauth.disconnect.output) .handler( - handlerGen(function* ({ context }) { - return yield* context.codexOauthService.disconnectEffect(); + handlerGen(function* ({ context }, input) { + return yield* context.codexOauthService.disconnectEffect(input?.accountId); + }) + ), + setDefaultAccount: t + .input(schemas.codexOauth.setDefaultAccount.input) + .output(schemas.codexOauth.setDefaultAccount.output) + .handler( + handlerGen(function* ({ context }, input) { + return yield* context.codexOauthService.setDefaultAccountEffect(input.accountId); + }) + ), + renameAccount: t + .input(schemas.codexOauth.renameAccount.input) + .output(schemas.codexOauth.renameAccount.output) + .handler( + handlerGen(function* ({ context }, input) { + return yield* context.codexOauthService.renameAccountEffect( + input.accountId, + input.label + ); }) ), }, @@ -1183,6 +1202,12 @@ export const router = (authToken?: string) => { .handler(({ context, input }) => context.projectService.setColor(input.projectPath, input.color) ), + setCodexOauthAccount: t + .input(schemas.projects.setCodexOauthAccount.input) + .output(schemas.projects.setCodexOauthAccount.output) + .handler(({ context, input }) => + context.projectService.setCodexOauthAccount(input.projectPath, input.accountId) + ), setCustomInstructions: t .input(schemas.projects.setCustomInstructions.input) .output(schemas.projects.setCustomInstructions.output) @@ -1329,7 +1354,16 @@ export const router = (authToken?: string) => { .input(schemas.nameGeneration.generate.input) .output(schemas.nameGeneration.generate.output) .handler(({ context, input }) => - generateWorkspaceIdentity(input.message, input.candidates, context.aiService) + generateWorkspaceIdentity( + input.message, + input.candidates, + context.aiService, + undefined, + undefined, + { + projectPath: input.projectPath, + } + ) ), }, coder: { diff --git a/src/node/services/agentSession.admissionGates.test.ts b/src/node/services/agentSession.admissionGates.test.ts index 523f113e27d..af39a3abd5f 100644 --- a/src/node/services/agentSession.admissionGates.test.ts +++ b/src/node/services/agentSession.admissionGates.test.ts @@ -8,7 +8,10 @@ import type { SendMessageError } from "@/common/types/errors"; import { createMuxMessage } from "@/common/types/message"; import { Ok } from "@/common/types/result"; import { AgentSession, CONTEXT_MUTATION_SEND_BLOCKED_MESSAGE } from "./agentSession"; -import { createStreamLifecycleMocks } from "./agentSession.testHarness"; +import { + createModelRoutingSnapshotMock, + createStreamLifecycleMocks, +} from "./agentSession.testHarness"; import { createTestHistoryService } from "./testHistoryService"; const TEST_MODEL = "anthropic:claude-3-5-sonnet-latest"; @@ -16,6 +19,7 @@ const config = { rootDir: "/tmp", sessionsDir: "/tmp", srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: () => ({}), } as unknown as Config; @@ -36,6 +40,7 @@ describe("AgentSession.sendMessage (admission gates)", () => { const streamMessage = mock(() => Promise.resolve(Ok(undefined))); const aiService = Object.assign(new EventEmitter(), { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), isStreaming: mock((_workspaceId: string) => false), stopStream: mock((_workspaceId: string) => Promise.resolve(Ok(undefined))), streamMessage: streamMessage as unknown as AIService["streamMessage"], diff --git a/src/node/services/agentSession.autoCompaction.test.ts b/src/node/services/agentSession.autoCompaction.test.ts index 0bbf4a9dd75..4f054c66cc5 100644 --- a/src/node/services/agentSession.autoCompaction.test.ts +++ b/src/node/services/agentSession.autoCompaction.test.ts @@ -12,12 +12,14 @@ import { import { GOAL_CONTINUATION_KIND } from "@/constants/goals"; import { Ok, Err } from "@/common/types/result"; import { ProvidersConfigStore, type Config } from "@/node/config"; -import type { AIService } from "@/node/services/aiService"; +import { AIService } from "@/node/services/aiService"; +import { ProviderService } from "./providerService"; import type { BackgroundProcessManager } from "@/node/services/backgroundProcessManager"; import type { InitStateManager } from "@/node/services/initStateManager"; import { AgentSession } from "./agentSession"; import type { CompactionMonitor } from "./compactionMonitor"; import { + createModelRoutingSnapshotMock, createAgentSessionHarness, createStartedTurnHandle, createStreamLifecycleMocks, @@ -402,6 +404,7 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { rootDir: "/tmp", sessionsDir: "/tmp", srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: () => ({ agentAiDefaults: { compact: { modelString: compactionModel } }, }), @@ -661,6 +664,7 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { rootDir: "/tmp", sessionsDir: "/tmp", srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: () => ({ agentAiDefaults: { compact: { modelString: compactionModel } }, }), @@ -713,6 +717,7 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { rootDir: "/tmp", sessionsDir: "/tmp", srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: () => ({ agentAiDefaults: { compact: { modelString: "openai:gpt-5.5", thinkingLevel: "high" }, @@ -792,7 +797,7 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { session.dispose(); }); - test("threads providers config into pre-send and mid-stream compaction checks", async () => { + test("threads provider config and project routing into pre-send and mid-stream checks", async () => { const workspaceId = "ws-auto-compaction-providers-config"; const { config, historyService, cleanup } = await createTestHistoryService(); @@ -802,13 +807,20 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { openai: { models: [ { - id: "openai:gpt-4o", - contextWindow: 222_222, + id: "gpt-4o", + contextWindowTokens: 222_222, }, ], }, }; new ProvidersConfigStore(config.rootDir).saveProvidersConfig(providersConfig); + await config.editConfig((cfg) => { + cfg.projects.set(config.rootDir, { + codexOauthAccountId: "work", + workspaces: [{ id: workspaceId, name: workspaceId, path: config.rootDir }], + }); + return cfg; + }); const aiEmitter = new EventEmitter(); const streamMessage = mock((_history: MuxMessage[]) => { @@ -842,6 +854,7 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { const aiService = Object.assign(aiEmitter, { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), isStreaming: mock((_workspaceId: string) => false), stopStream: mock((_workspaceId: string) => Promise.resolve(Ok(undefined))), streamMessage: streamMessage as unknown as ( @@ -850,6 +863,14 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { }) as unknown as AIService; const initStateManager = new EventEmitter() as unknown as InitStateManager; + const routingService = new AIService( + config, + historyService, + initStateManager, + new ProviderService(config) + ); + aiService.captureModelRoutingSnapshot = + routingService.captureModelRoutingSnapshot.bind(routingService); const backgroundProcessManager = { cleanup: mock((_workspaceId: string) => Promise.resolve()), @@ -892,11 +913,13 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { expect(checkBeforeSend).toHaveBeenCalledTimes(1); expect(checkBeforeSend.mock.calls[0]?.[0]).toMatchObject({ providersConfig, + codexOauthAccountId: "work", }); expect(checkMidStream).toHaveBeenCalledTimes(1); expect(checkMidStream.mock.calls[0]?.[0]).toMatchObject({ providersConfig, + codexOauthAccountId: "work", }); session.dispose(); @@ -957,6 +980,7 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { ); const aiService = Object.assign(aiEmitter, { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), isStreaming: mock((_workspaceId: string) => false), stopStream: mock((_workspaceId: string) => Promise.resolve(Ok(undefined))), streamMessage: streamMessage as unknown as ( @@ -977,6 +1001,7 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { rootDir: "/tmp", sessionsDir: "/tmp", srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: () => ({}), } as unknown as Config; @@ -1069,6 +1094,7 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { const aiService = Object.assign(aiEmitter, { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), isStreaming: mock((_workspaceId: string) => false), stopStream, streamMessage: streamMessage as unknown as ( @@ -1089,6 +1115,7 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { rootDir: "/tmp", sessionsDir: "/tmp", srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: () => ({}), } as unknown as Config; @@ -1220,6 +1247,7 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { const aiService = Object.assign(aiEmitter, { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), isStreaming: mock((_workspaceId: string) => false), stopStream, streamMessage: streamMessage as unknown as ( @@ -1240,6 +1268,7 @@ describe("AgentSession on-send auto-compaction snapshot deferral", () => { rootDir: "/tmp", sessionsDir: "/tmp", srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: () => ({}), } as unknown as Config; diff --git a/src/node/services/agentSession.budgetGate.test.ts b/src/node/services/agentSession.budgetGate.test.ts index 000f2b18680..0008eef56f7 100644 --- a/src/node/services/agentSession.budgetGate.test.ts +++ b/src/node/services/agentSession.budgetGate.test.ts @@ -7,7 +7,10 @@ import type { HistoryService } from "./historyService"; import type { InitStateManager } from "./initStateManager"; import { AgentSession } from "./agentSession"; import { createTestHistoryService } from "./testHistoryService"; -import { createStartedTurnHandle } from "./agentSession.testHarness"; +import { + createModelRoutingSnapshotMock, + createStartedTurnHandle, +} from "./agentSession.testHarness"; import { WorkspaceGoalService } from "./workspaceGoalService"; // Registers a no-op goal-continuation consumer so the in-AS pricing gate // path runs end-to-end (DEREM-52). Bridge registration alone is now @@ -52,6 +55,7 @@ function createAiService(workspaceId: string): AIService { streamMessage: mock((_request: unknown) => Promise.resolve(Ok(createStartedTurnHandle()))), getStreamInfo: mock((_workspaceId: string) => null), getProvidersConfig: mock(() => null), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), getWorkspaceMetadata: mock((_workspaceId: string) => Promise.resolve( Ok({ diff --git a/src/node/services/agentSession.continuousCompaction.test.ts b/src/node/services/agentSession.continuousCompaction.test.ts index 1085dd3d3ba..bfb0792439d 100644 --- a/src/node/services/agentSession.continuousCompaction.test.ts +++ b/src/node/services/agentSession.continuousCompaction.test.ts @@ -8,6 +8,10 @@ import { MockLanguageModelV3, simulateReadableStream } from "ai/test"; import type { LanguageModelV3CallOptions, LanguageModelV3StreamPart } from "@ai-sdk/provider"; import { summarizeContinuousCompaction } from "./continuousCompactionSummary"; import type { SessionUsageService } from "./sessionUsageService"; +import { AIService } from "./aiService"; +import { ProviderService } from "./providerService"; +import { ProvidersConfigStore } from "@/node/config"; +import type { CompactionMonitor } from "./compactionMonitor"; import { EXPERIMENT_IDS } from "@/common/constants/experiments"; import type { ProvidersConfigMap, SendMessageOptions } from "@/common/orpc/types"; import { @@ -24,6 +28,10 @@ import { type AgentSessionHarness, } from "./agentSession.testHarness"; import type { ContinuousCompactor } from "./continuousCompactor"; +import { MULTI_PROJECT_CONFIG_KEY } from "@/common/constants/multiProject"; +import { resolveCodexOauthRouting } from "@/common/utils/providers/codexOauthRouting"; +import { estimateMuxMessageTokens } from "@/common/utils/messages/keepRecentTail"; +import { SUMMARIZER_INPUT_FRACTION } from "@/constants/continuousCompaction"; const workspaceId = "continuous-session"; const model = "openai:gpt-4o"; @@ -34,6 +42,12 @@ const sendOptions: SendMessageOptions = { }; interface SessionInternals { + getContinuousCompactionContext: ( + model: string, + options?: SendMessageOptions + ) => { + contextWindowTokens: number; + }; continuousCompactor: ContinuousCompactor; activeStreamContext?: { modelString: string; @@ -84,6 +98,19 @@ describe("AgentSession continuous compaction wiring", () => { async function setup(usagePercent = 0) { harness = await createAgentSessionHarness({ workspaceId, captureEvents: true }); + const h = harness; + const routingService = new AIService( + h.config, + h.historyService, + h.initStateManager, + new ProviderService(h.config) + ); + spyOn(routingService, "getProvidersConfig").mockImplementation(() => + structuredClone(h.aiService.getProvidersConfig()) + ); + spyOn(h.aiService, "captureModelRoutingSnapshot").mockImplementation((id) => + routingService.captureModelRoutingSnapshot(id) + ); harness.session.setAutoCompactionThreshold(0.7); if (usagePercent > 0) { await harness.historyService.appendToHistory( @@ -102,9 +129,237 @@ describe("AgentSession continuous compaction wiring", () => { }) ); } - return harness; + return { ...harness, routingService }; + } + + async function setWorkspaceAccountScope( + h: AgentSessionHarness, + scope: "root" | "subproject" | "multi" | "inherited" + ) { + const root = h.config.rootDir; + const subproject = root + "/sub"; + await h.config.editConfig((cfg) => { + const workspace = { + id: workspaceId, + name: workspaceId, + path: root, + ...(scope !== "root" ? { subProjectPath: subproject } : {}), + ...(scope === "multi" + ? { + projects: [ + { projectPath: subproject, projectName: "sub" }, + { projectPath: root, projectName: "root" }, + ], + } + : {}), + }; + cfg.projects.set(root, { + codexOauthAccountId: scope === "root" || scope === "inherited" ? "missing" : undefined, + workspaces: scope === "multi" ? [] : [workspace], + }); + cfg.projects.set(subproject, { + codexOauthAccountId: scope === "inherited" ? undefined : "missing-sub", + workspaces: [], + }); + if (scope === "multi") { + cfg.projects.set(MULTI_PROJECT_CONFIG_KEY, { workspaces: [workspace] }); + } + return cfg; + }); } + test.each(["root", "subproject", "multi", "inherited"] as const)( + "continuous limits use the %s account scope without changing API-key preference", + async (scope) => { + const h = await setup(); + await setWorkspaceAccountScope(h, scope); + const providersConfig: ProvidersConfigMap = { + openai: { apiKeySet: true, isEnabled: true, isConfigured: true }, + }; + spyOn(h.aiService, "getProvidersConfig").mockReturnValue(providersConfig); + const scopedModel = "openai:gpt-5.5"; + const accountId = + scope === "inherited" ? undefined : scope === "root" ? "missing" : "missing-sub"; + // This fixture actually changes the route; API-key preference alone cannot select project OAuth. + expect(resolveCodexOauthRouting(scopedModel, providersConfig)).toBe("other"); + expect( + resolveCodexOauthRouting(scopedModel, providersConfig, { codexOauthAccountId: accountId }) + ).toBe(scope === "inherited" ? "other" : "missing-account"); + const readLimit = (options?: SendMessageOptions) => + internals(h.session).getContinuousCompactionContext(scopedModel, options) + .contextWindowTokens; + expect(readLimit()).toBe(scope === "inherited" ? 1_050_000 : 272_000); + expect( + readLimit({ + ...sendOptions, + model: scopedModel, + providerOptions: { openai: { wireFormat: "chatCompletions" } }, + }) + ).toBe(1_050_000); + providersConfig.openai.codexOauthDefaultAuth = "apiKey"; + expect(readLimit()).toBe(1_050_000); + delete providersConfig.openai.codexOauthDefaultAuth; + providersConfig.openai.codexOauthAccounts = [ + { id: accountId ?? "default", label: "Reconnected" }, + ]; + providersConfig.openai.codexOauthSet = true; + expect(readLimit()).toBe(272_000); + } + ); + + test("active compaction keeps its routing snapshot until the next turn", async () => { + const h = await setup(); + await setWorkspaceAccountScope(h, "root"); + const providersConfig: ProvidersConfigMap = { + openai: { apiKeySet: true, isEnabled: true, isConfigured: true }, + }; + spyOn(h.aiService, "getProvidersConfig").mockReturnValue(providersConfig); + const scopedModel = "openai:gpt-5.5"; + const options = { ...sendOptions, model: scopedModel }; + spyOn(h.aiService, "streamMessage").mockImplementation(() => { + startStream(h); + return Promise.resolve(Ok(createStartedTurnHandle())); + }); + const observed = deferred(); + const observe = spyOn(internals(h.session).continuousCompactor, "observe").mockImplementation( + (_usage, context) => { + if (context.phase === "mid-stream") observed.resolve(); + return Promise.resolve("none"); + } + ); + expect((await h.session.sendMessage("Start", options)).success).toBe(true); + await h.config.editConfig((cfg) => { + delete cfg.projects.get(h.config.rootDir)!.codexOauthAccountId; + return cfg; + }); + // The next turn now has an API route. This turn keeps its explicit selection and smaller cap. + expect(resolveCodexOauthRouting(scopedModel, providersConfig)).toBe("other"); + expect( + internals(h.session).getContinuousCompactionContext(scopedModel, options).contextWindowTokens + ).toBe(272_000); + h.aiEmitter.emit("usage-delta", { + type: "usage-delta", + workspaceId, + messageId: "live-assistant", + usage: { inputTokens: 200_000, outputTokens: 1, totalTokens: 200_001 }, + }); + await observed.promise; + expect( + observe.mock.calls.find((call) => call[1].phase === "mid-stream")?.[1].contextWindowTokens + ).toBe(272_000); + endStream(h); + await h.session.waitForIdle(); + expect((await h.session.sendMessage("Next", options)).success).toBe(true); + expect( + internals(h.session).getContinuousCompactionContext(scopedModel, options).contextWindowTokens + ).toBe(1_050_000); + }); + + test.each(["api-key", "wire-format"] as const)( + "active OAuth compaction ignores later %s changes", + async (change) => { + const h = await setup(); + const providersConfig: ProvidersConfigMap = { + openai: { apiKeySet: true, isEnabled: true, isConfigured: true, codexOauthSet: true }, + }; + spyOn(h.aiService, "getProvidersConfig").mockReturnValue(providersConfig); + const scopedModel = "openai:gpt-5.5"; + const options = { ...sendOptions, model: scopedModel }; + spyOn(h.aiService, "streamMessage").mockImplementation(() => { + startStream(h); + return Promise.resolve(Ok(createStartedTurnHandle())); + }); + spyOn(internals(h.session).continuousCompactor, "observe").mockResolvedValue("none"); + expect(resolveCodexOauthRouting(scopedModel, providersConfig)).toBe("oauth"); + expect((await h.session.sendMessage("Start", options)).success).toBe(true); + if (change === "api-key") providersConfig.openai.codexOauthDefaultAuth = "apiKey"; + else providersConfig.openai.wireFormat = "chatCompletions"; + expect(resolveCodexOauthRouting(scopedModel, providersConfig)).toBe("other"); + expect( + internals(h.session).getContinuousCompactionContext(scopedModel, options) + .contextWindowTokens + ).toBe(272_000); + endStream(h); + await h.session.waitForIdle(); + expect((await h.session.sendMessage("Next", options)).success).toBe(true); + expect( + internals(h.session).getContinuousCompactionContext(scopedModel, options) + .contextWindowTokens + ).toBe(1_050_000); + } + ); + + test("pre-send and model construction share the selection before asynchronous preparation", async () => { + const h = await setup(); + await setWorkspaceAccountScope(h, "root"); + spyOn(h.aiService, "getProvidersConfig").mockReturnValue({ + openai: { apiKeySet: true, isEnabled: true, isConfigured: true }, + }); + spyOn(internals(h.session).continuousCompactor, "observe").mockImplementation( + async (_usage, context) => { + if (context.phase === "on-send") + await h.config.editConfig((cfg) => { + delete cfg.projects.get(h.config.rootDir)!.codexOauthAccountId; + return cfg; + }); + return "none"; + } + ); + const stream = spyOn(h.aiService, "streamMessage").mockImplementation(() => { + startStream(h); + return Promise.resolve(Ok(createStartedTurnHandle())); + }); + const options = { ...sendOptions, model: "openai:gpt-5.5" }; + expect((await h.session.sendMessage("Start", options)).success).toBe(true); + expect(stream.mock.calls[0]?.[0].modelRoutingSnapshot?.codexOauthSelection).toEqual({ + accountId: "missing", + explicit: true, + }); + expect( + internals(h.session).getContinuousCompactionContext(options.model, options) + .contextWindowTokens + ).toBe(272_000); + }); + + test("compaction follows the current fallback model within the pinned turn", async () => { + const h = await setup(); + spyOn(h.aiService, "getProvidersConfig").mockReturnValue({ + openai: { apiKeySet: true, isConfigured: true, isEnabled: true, codexOauthSet: true }, + }); + spyOn(h.aiService, "streamMessage").mockImplementation(() => { + startStream(h); + return Promise.resolve(Ok(createStartedTurnHandle())); + }); + const observed = deferred(); + const observe = spyOn(internals(h.session).continuousCompactor, "observe").mockImplementation( + (_usage, context) => { + if (context.phase === "mid-stream") observed.resolve(); + return Promise.resolve("none"); + } + ); + expect( + (await h.session.sendMessage("Start", { ...sendOptions, model: "openai:gpt-5.6-sol" })) + .success + ).toBe(true); + spyOn(h.aiService, "getStreamInfo").mockReturnValue({ + messageId: "live-assistant", + model: "openai:gpt-5.5", + parts: [], + toolCompletionTimestamps: new Map(), + }); + h.aiEmitter.emit("usage-delta", { + type: "usage-delta", + workspaceId, + messageId: "live-assistant", + usage: { inputTokens: 200_000, outputTokens: 1, totalTokens: 200_001 }, + }); + await observed.promise; + expect(observe.mock.calls.find((call) => call[1].phase === "mid-stream")?.[1]).toMatchObject({ + model: "openai:gpt-5.5", + contextWindowTokens: 272_000, + }); + }); + async function rows(h: AgentSessionHarness): Promise { const history = await h.historyService.getHistoryFromLatestBoundary(workspaceId); if (!history.success) throw new Error(history.error); @@ -895,6 +1150,144 @@ describe("AgentSession continuous compaction wiring", () => { return streamMessage; } + test.each([ + { outcome: "applied", change: "account" }, + { outcome: "applied", change: "preference" }, + { outcome: "failed", change: "account" }, + { outcome: "failed", change: "preference" }, + { outcome: "legacy-fallback", change: "account" }, + { outcome: "legacy-fallback", change: "preference" }, + { outcome: "applied", change: "priority" }, + { outcome: "applied", change: "override" }, + { outcome: "failed", change: "priority" }, + { outcome: "failed", change: "override" }, + { outcome: "legacy-fallback", change: "priority" }, + { outcome: "legacy-fallback", change: "override" }, + ] as const)( + "$outcome compaction keeps accepted routing after changing $change", + async ({ outcome, change }) => { + const h = await setup(); + // Exercise real capture and metadata projection for the continued turn. + spyOn(h.routingService, "getProvidersConfig").mockRestore(); + spyOn(h.aiService, "getProvidersConfig").mockImplementation(() => + h.routingService.getProvidersConfig() + ); + const store = new ProvidersConfigStore(h.config.rootDir); + const providersConfig = { + openrouter: { apiKey: "openrouter-key" }, + openai: { + apiKey: "test-api-key", + codexOauthAccounts: { + work: { + label: "Work", + credentials: { + type: "oauth" as const, + access: "work-access", + refresh: "work-refresh", + expires: Date.now() + 60_000, + }, + }, + personal: { + label: "Personal", + credentials: { + type: "oauth" as const, + access: "personal-access", + refresh: "personal-refresh", + expires: Date.now() + 60_000, + }, + }, + }, + }, + }; + store.saveProvidersConfig(providersConfig); + await h.config.editConfig((cfg) => { + cfg.routePriority = ["direct"]; + cfg.routeOverrides = {}; + cfg.projects.set(h.config.rootDir, { + codexOauthAccountId: "work", + workspaces: [{ id: workspaceId, name: workspaceId, path: h.config.rootDir }], + }); + return cfg; + }); + const options = { ...sendOptions, model: "openai:gpt-5.5" }; + spyOn(internals(h.session).continuousCompactor, "observe").mockResolvedValue("none"); + const stream = mockAbortableStream(h); + expect((await h.session.sendMessage("Working", options)).success).toBe(true); + const original = stream.mock.calls[0]?.[0].modelRoutingSnapshot; + expect(original?.codexOauthSelection.accountId).toBe("work"); + const monitor = Reflect.get(h.session, "compactionMonitor") as CompactionMonitor; + const checkBeforeSend = monitor.checkBeforeSend.bind(monitor); + const pressure = spyOn(monitor, "checkBeforeSend").mockImplementation((args) => ({ + ...checkBeforeSend(args), + shouldForceCompact: outcome === "legacy-fallback", + })); + await applyThenFinish(h.session, async (followUp) => { + if (change === "account") { + await h.config.editConfig((cfg) => { + cfg.projects.get(h.config.rootDir)!.codexOauthAccountId = "personal"; + return cfg; + }); + } else if (change === "priority" || change === "override") { + await h.config.editConfig((cfg) => { + if (change === "priority") cfg.routePriority = ["openrouter", "direct"]; + else cfg.routeOverrides = { "openai:gpt-5.5": "openrouter" }; + return cfg; + }); + } else { + store.saveProvidersConfig({ + openai: { ...providersConfig.openai, codexOauthDefaultAuth: "apiKey" }, + }); + } + if (outcome === "applied") await appendBoundary(h, followUp); + return outcome === "applied"; + }); + expect(stream).toHaveBeenCalledTimes(2); + const continuedSnapshot = stream.mock.calls[1]?.[0].modelRoutingSnapshot; + expect(continuedSnapshot).toBe(original); + const continuedModel = await h.routingService.createModel(options.model, undefined, { + modelRoutingSnapshot: continuedSnapshot, + }); + expect(continuedModel.success).toBe(true); + if (continuedModel.success && typeof continuedModel.data !== "string") { + expect(continuedModel.data.modelId).toBe("gpt-5.5"); + } + expect( + internals(h.session).getContinuousCompactionContext(options.model, options) + .contextWindowTokens + ).toBe(272_000); + const history = await rows(h); + expect(history.at(-1)?.metadata?.muxMetadata?.type === "compaction-request").toBe( + outcome === "legacy-fallback" + ); + // Credentials remain transient even when compaction stores a durable follow-up. + expect(JSON.stringify(history)).not.toContain("work-access"); + expect(JSON.stringify(history)).not.toContain("work-refresh"); + pressure.mockRestore(); + endStream(h); + await h.session.waitForIdle(); + expect((await h.session.sendMessage("New user turn", options)).success).toBe(true); + const next = stream.mock.calls.at(-1)?.[0].modelRoutingSnapshot; + expect(next).not.toBe(original); + const nextModel = await h.routingService.createModel(options.model, undefined, { + modelRoutingSnapshot: next, + }); + expect(nextModel.success).toBe(true); + if (nextModel.success && typeof nextModel.data !== "string") { + expect(nextModel.data.modelId).toBe( + change === "priority" || change === "override" ? "openai/gpt-5.5" : "gpt-5.5" + ); + } + expect(next?.codexOauthSelection.accountId).toBe(change === "account" ? "personal" : "work"); + expect(next?.providersConfig.openai?.codexOauthDefaultAuth).toBe( + change === "preference" ? "apiKey" : undefined + ); + expect( + internals(h.session).getContinuousCompactionContext(options.model, options) + .contextWindowTokens + ).toBe(change === "preference" ? 1_050_000 : 272_000); + } + ); + test("abandon during fast apply cannot resume the abandoned turn", async () => { const h = await setup(); spyOn(internals(h.session).continuousCompactor, "observe").mockResolvedValue("none"); @@ -1089,6 +1482,66 @@ describe("AgentSession continuous compaction wiring", () => { expect(result?.model).toBe(model); }); + test("headless summaries reject a head above the workspace OAuth cap", async () => { + const { h, args } = await summarySetup(); + await setWorkspaceAccountScope(h, "subproject"); + spyOn(h.aiService, "getProvidersConfig").mockReturnValue({ + openai: { apiKeySet: true, isEnabled: true, isConfigured: true }, + }); + const head = [ + createMuxMessage("large-head", "user", "Important context to retain. ".repeat(40_000)), + ]; + const headTokens = estimateMuxMessageTokens(head[0]); + expect(headTokens).toBeGreaterThan(272_000 * SUMMARIZER_INPUT_FRACTION); + expect(headTokens).toBeLessThan(1_050_000 * SUMMARIZER_INPUT_FRACTION); + const sdkModel = new MockLanguageModelV3({ + doStream: () => + Promise.resolve({ stream: simulateReadableStream({ chunks: modelChunks() }) }), + }); + const create = spyOn(h.aiService, "createModelWithPinnedMetadata").mockResolvedValue( + Ok({ model: sdkModel, metadataModel: "openai:gpt-5.5" }) + ); + const result = await summarizeContinuousCompaction({ + ...args, + head, + compactOptions: { + ...args.compactOptions, + model: "openai:gpt-5.5", + providerOptions: { openai: { wireFormat: "chatCompletions" } }, + }, + }); + expect(result).toBeNull(); + expect(create).not.toHaveBeenCalled(); + }); + + test("headless fallback does not inherit the active Chat Completions window", async () => { + const { h, args } = await summarySetup(); + spyOn(h.aiService, "getProvidersConfig").mockReturnValue({ + openai: { + apiKeySet: true, + isConfigured: true, + isEnabled: true, + codexOauthSet: true, + models: [{ id: "gpt-4.1-mini", contextWindowTokens: 100 }], + }, + }); + const create = spyOn(h.aiService, "createModelWithPinnedMetadata"); + const result = await summarizeContinuousCompaction({ + ...args, + head: [ + createMuxMessage("large-head", "user", "Important context to retain. ".repeat(40_000)), + ], + context: { ...args.context, model: "openai:gpt-5.5", contextWindowTokens: 1_050_000 }, + baseOptions: { + ...sendOptions, + model: "openai:gpt-5.5", + providerOptions: { openai: { wireFormat: "chatCompletions" } }, + }, + }); + expect(result).toBeNull(); + expect(create).not.toHaveBeenCalled(); + }); + test("returns null without calling a model when neither configured context can fit the head", async () => { const { h, args } = await summarySetup(); spyOn(h.aiService, "getProvidersConfig").mockReturnValue({ diff --git a/src/node/services/agentSession.disposeRace.test.ts b/src/node/services/agentSession.disposeRace.test.ts index a950e5b8398..0031f68cf71 100644 --- a/src/node/services/agentSession.disposeRace.test.ts +++ b/src/node/services/agentSession.disposeRace.test.ts @@ -18,7 +18,11 @@ import { startAbandonedBranchSummaryInBackground, type BranchSummaryAiService, } from "./branchSummary"; -import { createAgentSessionHarness, createStreamLifecycleMocks } from "./agentSession.testHarness"; +import { + createAgentSessionHarness, + createModelRoutingSnapshotMock, + createStreamLifecycleMocks, +} from "./agentSession.testHarness"; import type { StreamMessageOptions } from "./aiService"; import type { TurnCompletion } from "./streamManager"; @@ -41,6 +45,7 @@ describe("AgentSession disposal race conditions", () => { const aiService: AIService = { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on(eventName: string | symbol, listener: (...args: unknown[]) => void) { aiHandlers.set(String(eventName), listener); return this; @@ -135,6 +140,7 @@ describe("AgentSession disposal race conditions", () => { const streamMessage = mock(() => Promise.resolve(Ok(undefined))); const aiService: AIService = { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on(_eventName: string | symbol, _listener: (...args: unknown[]) => void) { return this; }, @@ -187,6 +193,7 @@ describe("AgentSession disposal race conditions", () => { releaseModel = resolve; }); const gatedAiService = { + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), createModelWithPinnedMetadata: async () => { await modelGate; return Err({ type: "api_key_not_found" as const, provider: "anthropic" }); @@ -254,6 +261,7 @@ describe("AgentSession disposal race conditions", () => { const aiService: AIService = { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on(eventName: string | symbol, listener: (...args: unknown[]) => void) { aiHandlers.set(String(eventName), listener); return this; @@ -341,6 +349,7 @@ describe("AgentSession disposal race conditions", () => { const aiService: AIService = { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on(eventName: string | symbol, listener: (...args: unknown[]) => void) { aiHandlers.set(String(eventName), listener); return this; @@ -436,6 +445,7 @@ describe("AgentSession disposal race conditions", () => { test("does not reset auto-retry intent for synthetic or rejected sends", async () => { const aiService: AIService = { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on(_eventName: string | symbol, _listener: (...args: unknown[]) => void) { return this; }, @@ -581,6 +591,7 @@ describe("AgentSession disposal race conditions", () => { test("preserves synthetic flag when flushing queued messages", () => { const aiService: AIService = { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on(_eventName: string | symbol, _listener: (...args: unknown[]) => void) { return this; }, diff --git a/src/node/services/agentSession.editMessageId.test.ts b/src/node/services/agentSession.editMessageId.test.ts index 5c6125fb702..0a0333113f6 100644 --- a/src/node/services/agentSession.editMessageId.test.ts +++ b/src/node/services/agentSession.editMessageId.test.ts @@ -8,7 +8,11 @@ import { createMuxMessage } from "@/common/types/message"; import { Ok } from "@/common/types/result"; import { AgentSession } from "./agentSession"; import { createTestHistoryService } from "./testHistoryService"; -import { createStartedTurnHandle, createStreamLifecycleMocks } from "./agentSession.testHarness"; +import { + createModelRoutingSnapshotMock, + createStartedTurnHandle, + createStreamLifecycleMocks, +} from "./agentSession.testHarness"; type StreamMessageHandler = AIService["streamMessage"]; @@ -18,6 +22,7 @@ const config = { rootDir: "/tmp", sessionsDir: "/tmp", srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: () => ({}), } as unknown as Config; @@ -45,6 +50,7 @@ describe("AgentSession.sendMessage (editMessageId)", () => { const streamMessage = mock(streamHandler); const aiService = Object.assign(new EventEmitter(), { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), isStreaming: mock((_workspaceId: string) => false), stopStream: mock((_workspaceId: string) => Promise.resolve(Ok(undefined))), streamMessage: streamMessage as unknown as AIService["streamMessage"], diff --git a/src/node/services/agentSession.fileChangeNotification.test.ts b/src/node/services/agentSession.fileChangeNotification.test.ts index f18428ce287..87f1e4ac18b 100644 --- a/src/node/services/agentSession.fileChangeNotification.test.ts +++ b/src/node/services/agentSession.fileChangeNotification.test.ts @@ -12,7 +12,11 @@ import type { AIService, StreamMessageOptions } from "./aiService"; import type { BackgroundProcessManager } from "./backgroundProcessManager"; import type { InitStateManager } from "./initStateManager"; import { createTestHistoryService } from "./testHistoryService"; -import { createStartedTurnHandle, createStreamLifecycleMocks } from "./agentSession.testHarness"; +import { + createModelRoutingSnapshotMock, + createStartedTurnHandle, + createStreamLifecycleMocks, +} from "./agentSession.testHarness"; /** * Log purity: externally-edited files must produce a durable @@ -52,6 +56,7 @@ describe("AgentSession file-change notification (turn start)", () => { }); const aiService: AIService = { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on: mock(() => aiService), off: mock(() => aiService), stopStream: mock(() => Promise.resolve(Ok(undefined))), diff --git a/src/node/services/agentSession.goalAutoPause.test.ts b/src/node/services/agentSession.goalAutoPause.test.ts index b3e71c0236d..58fa6c3c4e3 100644 --- a/src/node/services/agentSession.goalAutoPause.test.ts +++ b/src/node/services/agentSession.goalAutoPause.test.ts @@ -21,7 +21,11 @@ import { } from "@/constants/goals"; import { waitForCondition } from "./testDispatchHelpers"; import { IdleDispatcher } from "./idleDispatcher"; -import { createFailedTurnHandle, createStartedTurnHandle } from "./agentSession.testHarness"; +import { + createModelRoutingSnapshotMock, + createFailedTurnHandle, + createStartedTurnHandle, +} from "./agentSession.testHarness"; const PROJECT_PATH = "/tmp/mux-agent-session-goal-test-project"; const SEND_OPTIONS: SendMessageOptions = { model: "openai:gpt-4o", agentId: "exec" }; @@ -56,6 +60,7 @@ function createAiService(workspaceId: string): AIService & EventEmitter { streamMessage: mock((_request: unknown) => Promise.resolve(Ok(createStartedTurnHandle()))), getStreamInfo: mock((_workspaceId: string) => null), getProvidersConfig: mock(() => null), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), getWorkspaceMetadata: mock((_workspaceId: string) => Promise.resolve( Ok({ diff --git a/src/node/services/agentSession.postCompactionRetry.test.ts b/src/node/services/agentSession.postCompactionRetry.test.ts index 2e5c193991d..606fbff994c 100644 --- a/src/node/services/agentSession.postCompactionRetry.test.ts +++ b/src/node/services/agentSession.postCompactionRetry.test.ts @@ -14,6 +14,7 @@ import type { MuxMessage } from "@/common/types/message"; import type { SendMessageOptions } from "@/common/orpc/types"; import { createTestHistoryService } from "./testHistoryService"; import { + createModelRoutingSnapshotMock, createFailedTurnHandle, createStartedTurnHandle, createStreamLifecycleMocks, @@ -120,6 +121,7 @@ describe("AgentSession post-compaction context retry", () => { const aiService: AIService = { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on(eventName: string | symbol, listener: (...args: unknown[]) => void) { aiEmitter.on(String(eventName), listener); return this; @@ -151,6 +153,7 @@ describe("AgentSession post-compaction context retry", () => { rootDir: sessionsDir, sessionsDir, srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: mock(() => ({})), } as unknown as Config; @@ -272,6 +275,7 @@ describe("AgentSession post-compaction context retry", () => { const aiService: AIService = { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on(eventName: string | symbol, listener: (...args: unknown[]) => void) { aiEmitter.on(String(eventName), listener); return this; @@ -304,6 +308,7 @@ describe("AgentSession post-compaction context retry", () => { rootDir: sessionsDir, sessionsDir, srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: mock(() => ({})), } as unknown as Config; @@ -415,6 +420,7 @@ describe("AgentSession post-compaction context retry", () => { const aiService: AIService = { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on(eventName: string | symbol, listener: (...args: unknown[]) => void) { aiEmitter.on(String(eventName), listener); return this; @@ -447,6 +453,7 @@ describe("AgentSession post-compaction context retry", () => { rootDir: sessionsDir, sessionsDir, srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: mock(() => ({})), } as unknown as Config; diff --git a/src/node/services/agentSession.preStreamError.test.ts b/src/node/services/agentSession.preStreamError.test.ts index c94f377ef0e..27d5840eb5a 100644 --- a/src/node/services/agentSession.preStreamError.test.ts +++ b/src/node/services/agentSession.preStreamError.test.ts @@ -16,7 +16,10 @@ import { type WorkspaceChatMessage, } from "@/common/orpc/types"; import { AgentSession } from "./agentSession"; -import { createAgentSessionHarness } from "./agentSession.testHarness"; +import { + createModelRoutingSnapshotMock, + createAgentSessionHarness, +} from "./agentSession.testHarness"; import { createTestHistoryService } from "./testHistoryService"; interface ReplayHarnessStreamInfo { @@ -386,6 +389,7 @@ describe("AgentSession pre-stream errors", () => { ); }); const aiService = Object.assign(aiEmitter, { + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), isStreaming: mock((_workspaceId: string) => false), stopStream: mock((_workspaceId: string) => Promise.resolve(Ok(undefined))), getStreamInfo: mock((_workspaceId: string) => undefined), diff --git a/src/node/services/agentSession.preTurnMessages.test.ts b/src/node/services/agentSession.preTurnMessages.test.ts index 3ee5be6bfae..43944a8bb82 100644 --- a/src/node/services/agentSession.preTurnMessages.test.ts +++ b/src/node/services/agentSession.preTurnMessages.test.ts @@ -8,13 +8,18 @@ import { createMuxMessage } from "@/common/types/message"; import { Err, Ok } from "@/common/types/result"; import { AgentSession } from "./agentSession"; import { createTestHistoryService } from "./testHistoryService"; -import { createStartedTurnHandle, createStreamLifecycleMocks } from "./agentSession.testHarness"; +import { + createModelRoutingSnapshotMock, + createStartedTurnHandle, + createStreamLifecycleMocks, +} from "./agentSession.testHarness"; const TEST_MODEL = "anthropic:claude-3-5-sonnet-latest"; const config = { rootDir: "/tmp", sessionsDir: "/tmp", srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: () => ({}), } as unknown as Config; @@ -32,6 +37,7 @@ describe("AgentSession.sendMessage (preTurnMessages)", () => { const streamMessage = mock(() => Promise.resolve(Ok(createStartedTurnHandle()))); const aiService = Object.assign(new EventEmitter(), { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), isStreaming: mock((_workspaceId: string) => false), stopStream: mock((_workspaceId: string) => Promise.resolve(Ok(undefined))), streamMessage: streamMessage as unknown as AIService["streamMessage"], diff --git a/src/node/services/agentSession.resumeStreamEmptyHistory.test.ts b/src/node/services/agentSession.resumeStreamEmptyHistory.test.ts index 1f2f5fb3b5c..d4561419cc5 100644 --- a/src/node/services/agentSession.resumeStreamEmptyHistory.test.ts +++ b/src/node/services/agentSession.resumeStreamEmptyHistory.test.ts @@ -1,7 +1,10 @@ import { describe, expect, test, mock, afterEach } from "bun:test"; import { AgentSession } from "./agentSession"; -import { createStreamLifecycleMocks } from "./agentSession.testHarness"; +import { + createModelRoutingSnapshotMock, + createStreamLifecycleMocks, +} from "./agentSession.testHarness"; import type { Config } from "@/node/config"; import type { AIService } from "./aiService"; import type { InitStateManager } from "./initStateManager"; @@ -20,6 +23,7 @@ describe("AgentSession.resumeStream", () => { const aiService: AIService = { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on: mock(() => aiService), off: mock(() => aiService), stopStream: mock(() => Promise.resolve(Ok(undefined))), @@ -44,6 +48,7 @@ describe("AgentSession.resumeStream", () => { rootDir: "/tmp", sessionsDir: "/tmp", srcDir: "/tmp", + findWorkspace: () => null, loadConfigOrDefault: mock(() => ({})), } as unknown as Config; diff --git a/src/node/services/agentSession.startupAutoRetry.test.ts b/src/node/services/agentSession.startupAutoRetry.test.ts index 2abd4af325c..65d5199b254 100644 --- a/src/node/services/agentSession.startupAutoRetry.test.ts +++ b/src/node/services/agentSession.startupAutoRetry.test.ts @@ -17,7 +17,9 @@ import type { WorkspaceChatMessage, SendMessageOptions } from "@/common/orpc/typ import { createMuxMessage, pickStartupRetrySendOptions } from "@/common/types/message"; import { DEFAULT_RUNTIME_CONFIG } from "@/common/constants/workspace"; import type { WorkspaceMetadata } from "@/common/types/workspace"; -import { Ok } from "@/common/types/result"; +import { Err, Ok } from "@/common/types/result"; +import type { ModelRoutingSnapshot } from "./modelRoutingSnapshot"; +import { getEffectiveContextLimit } from "@/common/utils/compaction/contextLimit"; import { GOAL_CONTINUATION_KIND } from "@/constants/goals"; import { formatSubagentReportEnvelope } from "@/common/utils/subagentReportEnvelope"; import { WORKSPACE_DEFAULTS } from "@/constants/workspaceDefaults"; @@ -415,6 +417,61 @@ describe("AgentSession startup auto-retry recovery", () => { session.dispose(); }); + test.each(["automatic", "explicit"] as const)( + "%s resume uses the correct provider snapshot after settings change", + async (kind) => { + const workspaceId = "retry-account-" + kind; + const { session, historyService, aiService, cleanup } = + await createSessionBundle(workspaceId); + cleanups.push(cleanup); + const initial: ModelRoutingSnapshot = { + routeConfig: { routePriority: ["direct"], routeOverrides: {} }, + providersConfig: { openai: { apiKey: "key", codexOauthDefaultAuth: "oauth" } }, + metadata: { + openai: { apiKeySet: true, isEnabled: true, isConfigured: true, codexOauthSet: true }, + }, + codexOauthSelection: { accountId: "work", explicit: true }, + }; + let current = initial; + const capture = spyOn(aiService, "captureModelRoutingSnapshot").mockImplementation( + () => current + ); + const stream = spyOn(aiService, "streamMessage") + .mockResolvedValueOnce(Err({ type: "unknown", raw: "Temporary connection failure" })) + .mockResolvedValue(Ok(createStartedTurnHandle())); + const options = { model: "openai:gpt-5.5", agentId: "exec" }; + await historyService.appendToHistory( + workspaceId, + createMuxMessage("user", "user", "Continue") + ); + try { + expect((await session.resumeStream(options)).success).toBe(false); + current = { + routeConfig: { routePriority: ["direct"], routeOverrides: {} }, + providersConfig: { openai: { apiKey: "key", codexOauthDefaultAuth: "apiKey" } }, + metadata: { openai: { apiKeySet: true, isEnabled: true, isConfigured: true } }, + codexOauthSelection: { accountId: "personal", explicit: true }, + }; + if (kind === "automatic") { + await (session as unknown as RetryableSessionForTests).retryActiveStream(); + } else { + expect((await session.resumeStream(options)).success).toBe(true); + } + expect(stream).toHaveBeenCalledTimes(2); + const snapshot = stream.mock.calls[1]?.[0].modelRoutingSnapshot; + expect(snapshot).toBe(kind === "automatic" ? initial : current); + expect(capture).toHaveBeenCalledTimes(kind === "automatic" ? 1 : 2); + expect(getEffectiveContextLimit(options.model, false, snapshot?.metadata)).toBe( + kind === "automatic" ? 272_000 : 1_050_000 + ); + } finally { + stream.mockRestore(); + capture.mockRestore(); + session.dispose(); + } + } + ); + test("startup auto-retry reuses workspace-turn metadata from the retry user message", async () => { const workspaceId = "startup-retry-workspace-turn-metadata"; const { session, historyService, aiService, cleanup } = await createSessionBundle(workspaceId); diff --git a/src/node/services/agentSession.testHarness.ts b/src/node/services/agentSession.testHarness.ts index cf09d0cafea..47b7b7f21e4 100644 --- a/src/node/services/agentSession.testHarness.ts +++ b/src/node/services/agentSession.testHarness.ts @@ -19,6 +19,7 @@ import type { InitStateManager } from "@/node/services/initStateManager"; import type { MCPServerManager } from "@/node/services/mcpServerManager"; import { createTestHistoryService } from "@/node/services/testHistoryService"; import type { StreamErrorType } from "@/common/types/errors"; +import type { ModelRoutingSnapshot } from "./modelRoutingSnapshot"; export function createStartedTurnHandle(messageId = "test-assistant"): TurnStreamHandle { return { messageId, completion: new Promise(() => undefined) }; @@ -68,6 +69,7 @@ function createAgentSessionTestConfig(sessionDir = "/tmp"): Config { rootDir: sessionDir, sessionsDir: sessionDir, srcDir: sessionDir, + findWorkspace: () => null, loadConfigOrDefault: mock(() => ({})), } as unknown as Config; } @@ -98,6 +100,19 @@ export function createStreamLifecycleMocks() { }; } +export function createModelRoutingSnapshotMock( + getProvidersConfig: AgentSessionAIService["getProvidersConfig"] = () => null +) { + return mock( + (_workspaceId: string): ModelRoutingSnapshot => ({ + providersConfig: {}, + routeConfig: { routePriority: ["direct"], routeOverrides: {} }, + metadata: structuredClone(getProvidersConfig()), + codexOauthSelection: { accountId: "default", explicit: false }, + }) + ); +} + function createMockAiService(args?: { emitter?: EventEmitter; overrides?: Partial; @@ -117,6 +132,9 @@ function createMockAiService(args?: { Promise.resolve(Err("Test AI service has no workspace metadata")) ), getProvidersConfig: mock(() => null), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(() => + aiService.getProvidersConfig() + ), isExperimentEnabled: mock((_experimentId) => false), ...createStreamLifecycleMocks(), streamMessage: mock(() => diff --git a/src/node/services/agentSession.ts b/src/node/services/agentSession.ts index 35a8eab89f4..8b4d978c76f 100644 --- a/src/node/services/agentSession.ts +++ b/src/node/services/agentSession.ts @@ -212,10 +212,13 @@ import { injectPostCompactionAttachments } from "@/browser/utils/messages/modelM import { estimateMuxMessageTokens } from "@/common/utils/messages/keepRecentTail"; import { ContinuousCompactor, type ContinuousCompactionContext } from "./continuousCompactor"; import { getEffectiveContextLimit } from "@/common/utils/compaction/contextLimit"; +import type { CodexOauthRoutingOptions } from "@/common/utils/providers/codexOauthRouting"; +import type { ModelRoutingSnapshot } from "./modelRoutingSnapshot"; import { summarizeContinuousCompaction } from "./continuousCompactionSummary"; type SessionCompactionContext = ContinuousCompactionContext & { sendOptions?: SendMessageOptions; + modelRoutingSnapshot?: ModelRoutingSnapshot; }; /** @@ -261,6 +264,8 @@ interface AutoRetryResumeRequest { // ACP correlation/delegation lives in transient send options that are // intentionally omitted from durable startup-recovery snapshots. options: SendMessageOptions; + // Keep billing identity across backoff. Credentials stay in memory, never in startup-recovery metadata. + modelRoutingSnapshot?: ModelRoutingSnapshot; agentInitiated?: boolean; goalKind?: GoalSyntheticMessageKind; /** Goal identity matching goalKind; keeps retried streams goal-scoped. */ @@ -553,6 +558,7 @@ export interface AgentSessionMetadataEvent { interface AgentSessionActiveStreamInfo { messageId: string; + model?: string; startTime?: number; parts: Array< MuxMessage["parts"][number] & { timestamp?: number; workflowRun?: { timestamp?: number } } @@ -597,6 +603,7 @@ export interface AgentSessionAIService extends BranchSummaryAiService { getStreamInfo?(workspaceId: string): AgentSessionActiveStreamInfo | undefined; replayStream?(workspaceId: string, options?: { afterTimestamp?: number }): Promise; getProvidersConfig(): ProvidersConfigMap | null; + captureModelRoutingSnapshot(workspaceId: string): ModelRoutingSnapshot; isExperimentEnabled(experimentId: ExperimentId): boolean; buildMemorySessionContext?( workspaceId: string, @@ -879,6 +886,7 @@ export class AgentSession { agentInitiated?: boolean; openaiTruncationModeOverride?: "auto" | "disabled"; providersConfig: ProvidersConfigMap | null; + modelRoutingSnapshot: ModelRoutingSnapshot; goalKind?: GoalSyntheticMessageKind; /** Goal identity matching goalKind, so mid-stream compaction follow-ups stay goal-scoped. */ goalId?: string; @@ -1033,6 +1041,7 @@ export class AgentSession { context, baseOptions, compactOptions: request.sendOptions, + modelRoutingSnapshot: context.modelRoutingSnapshot, }); }, fastApply: (apply) => this.interruptForContinuousCompaction(apply), @@ -1315,7 +1324,8 @@ export class AgentSession { options: SendMessageOptions | undefined, agentInitiated?: boolean, goalKind?: GoalSyntheticMessageKind, - goalId?: string + goalId?: string, + modelRoutingSnapshot?: ModelRoutingSnapshot ): void { if (!options) { this.lastAutoRetryResumeRequest = undefined; @@ -1324,6 +1334,7 @@ export class AgentSession { this.lastAutoRetryResumeRequest = { options, + ...(modelRoutingSnapshot != null ? { modelRoutingSnapshot } : {}), ...(agentInitiated === true ? { agentInitiated: true } : {}), ...(goalKind != null ? { goalKind } : {}), ...(goalId != null ? { goalId } : {}), @@ -1361,6 +1372,7 @@ export class AgentSession { agentInitiated: request.agentInitiated === true ? true : undefined, goalKind: request.goalKind, goalId: request.goalId, + modelRoutingSnapshot: request.modelRoutingSnapshot, }); if (result.success) { if (!result.data.started) { @@ -3071,6 +3083,8 @@ export class AgentSession { /** A dequeued send keeps its admission owner through acceptance and startup failure. */ turnReservation?: TurnId; synthetic?: boolean; + /** Same-session continuations keep accepted routing. Restart recovery captures current settings. */ + modelRoutingSnapshot?: ModelRoutingSnapshot; agentInitiated?: boolean; goalContinuation?: boolean; goalKind?: GoalSyntheticMessageKind; @@ -3687,6 +3701,9 @@ export class AgentSession { // away would dangle the trigger's message-ID reference. Family sends are // small and bounded, so skip on-send compaction for them; mid-stream // forcing still protects the context limit. + // Share this snapshot with model construction; settings changes apply to the next turn. + const modelRoutingSnapshot = + internal?.modelRoutingSnapshot ?? this.captureModelRoutingSnapshot(); const hasPreTurnMessages = (internal?.preTurnMessages?.length ?? 0) > 0; if (!isCompactionRequest && !editMessageId && !hasPreTurnMessages) { // Seed usage state from persisted history on the first send after restart @@ -3697,7 +3714,7 @@ export class AgentSession { return Ok(undefined); } - const providersConfigForCompaction = this.getProvidersConfigSafe(); + const providersConfigForCompaction = modelRoutingSnapshot.metadata; // Recover before measuring pressure so the old pre-swap usage cannot force another fold. if (await this.continuousCompactor.recover()) this.clearUsageState(); const compactionResult = this.compactionMonitor.checkBeforeSend({ @@ -3709,12 +3726,13 @@ export class AgentSession { providersConfigForCompaction ), providersConfig: providersConfigForCompaction, - openaiWireFormat: optionsForStream.providerOptions?.openai?.wireFormat, + ...this.getCompactionRoutingOptions(optionsForStream, modelRoutingSnapshot), }); const continuousContext = this.getContinuousCompactionContext( modelForStream, - optionsForStream + optionsForStream, + modelRoutingSnapshot ); if (!continuousContext.enabled) this.continuousCompactor.reset("disabled"); const continuousResult = continuousContext.enabled @@ -4106,7 +4124,13 @@ export class AgentSession { // Same-session retry should resume the exact accepted request we just finalized // in history, even if runtime warmup fails before streamWithHistory() starts. - this.setAutoRetryResumeState(optionsForStream, agentInitiated, goalKind, internal?.goalId); + this.setAutoRetryResumeState( + optionsForStream, + agentInitiated, + goalKind, + internal?.goalId, + modelRoutingSnapshot + ); try { await internal?.onAccepted?.(); } catch (error) { @@ -4196,7 +4220,8 @@ export class AgentSession { preparedTurnAbortController.signal, goalKind, internal?.goalId, - turnThinkingOverride + turnThinkingOverride, + modelRoutingSnapshot ); if (streamResult.success && preparedTurnAbortController.signal.aborted) { await notifyAcceptedPreStreamFailure( @@ -4264,7 +4289,12 @@ export class AgentSession { async resumeStream( options: SendMessageOptions, - internal?: { agentInitiated?: boolean; goalKind?: GoalSyntheticMessageKind; goalId?: string } + internal?: { + agentInitiated?: boolean; + goalKind?: GoalSyntheticMessageKind; + goalId?: string; + modelRoutingSnapshot?: ModelRoutingSnapshot; + } ): Promise> { this.assertNotDisposed("resumeStream"); @@ -4303,13 +4333,18 @@ export class AgentSession { return Ok({ started: false }); } + // Automatic retries keep their failed attempt. Explicit resumes and restart recovery capture current settings. + const modelRoutingSnapshot = + internal?.modelRoutingSnapshot ?? this.captureModelRoutingSnapshot(); + // A resumed attempt becomes the latest live resume request as soon as we // accept its options, even if startup fails before the stream fully begins. this.setAutoRetryResumeState( optionsForStream, internal?.agentInitiated, internal?.goalKind, - internal?.goalId + internal?.goalId, + modelRoutingSnapshot ); this.preparingWorkspaceTurnMetadata = getWorkspaceTurnMuxMetadata(optionsForStream.muxMetadata); const preparedTurn = this.coordinator.prepare(); @@ -4330,7 +4365,8 @@ export class AgentSession { undefined, internal?.goalKind, internal?.goalId, - turnThinkingOverride + turnThinkingOverride, + modelRoutingSnapshot ); if (!result.success) { return result; @@ -4376,6 +4412,23 @@ export class AgentSession { return this.lastUsageState; } + private captureModelRoutingSnapshot(): ModelRoutingSnapshot { + return this.aiService.captureModelRoutingSnapshot(this.workspaceId); + } + + private getCompactionRoutingOptions( + options?: SendMessageOptions, + snapshot = this.activeStreamContext?.modelRoutingSnapshot ?? this.captureModelRoutingSnapshot() + ): CodexOauthRoutingOptions { + return { + // Preserve implicit API-only routing instead of synthesizing an explicit default selection. + codexOauthAccountId: snapshot.codexOauthSelection.explicit + ? snapshot.codexOauthSelection.accountId + : undefined, + openaiWireFormat: options?.providerOptions?.openai?.wireFormat, + }; + } + private getProvidersConfigSafe(): ProvidersConfigMap | null { try { // Prefer ProviderService's safe config view: it includes env/file API-key source @@ -4860,11 +4913,18 @@ export class AgentSession { }); } + private syncActiveCompactionModel(): void { + const model = this.streamManager.getStreamInfo(this.workspaceId)?.model; + if (model && this.activeStreamContext) this.activeStreamContext.modelString = model; + } + private getContinuousCompactionContext( model: string, - options?: SendMessageOptions + options?: SendMessageOptions, + modelRoutingSnapshot = this.activeStreamContext?.modelRoutingSnapshot ?? + this.captureModelRoutingSnapshot() ): SessionCompactionContext { - const providersConfig = this.getProvidersConfigSafe(); + const providersConfig = modelRoutingSnapshot.metadata; const enabled = options?.experiments?.continuousCompaction ?? (typeof this.aiService.isExperimentEnabled === "function" && @@ -4884,19 +4944,22 @@ export class AgentSession { getEffectiveContextLimit( model, this.is1MContextEnabledForModel(model, options, providersConfig), - providersConfig + providersConfig, + this.getCompactionRoutingOptions(options, modelRoutingSnapshot) ) ?? 0, thresholdPercent: this.compactionMonitor.getThreshold() * 100, systemMessageTokens: this.streamManager.getStreamInfo(this.workspaceId)?.initialMetadata?.systemMessageTokens ?? this.lastSystemMessageTokens, sendOptions: options, + modelRoutingSnapshot, }; } private async observeContinuousCompactionAtStreamEnd( model: string, - options?: SendMessageOptions + options?: SendMessageOptions, + modelRoutingSnapshot = this.captureModelRoutingSnapshot() ): Promise { // fastApply waits for this handler to reach IDLE; waiting on its latch here // (or re-entering it from the generated Continue send) would deadlock. @@ -4907,7 +4970,7 @@ export class AgentSession { ) return; try { - const context = this.getContinuousCompactionContext(model, options); + const context = this.getContinuousCompactionContext(model, options, modelRoutingSnapshot); if (!context.enabled && !this.continuousCompactor.hasConsumedSwap()) { this.continuousCompactor.reset("disabled"); return; @@ -4918,9 +4981,10 @@ export class AgentSession { use1MContext: this.is1MContextEnabledForModel( model, options, - this.getProvidersConfigSafe() + modelRoutingSnapshot.metadata ), - providersConfig: this.getProvidersConfigSafe(), + providersConfig: modelRoutingSnapshot.metadata, + ...this.getCompactionRoutingOptions(options, modelRoutingSnapshot), }); const result = await this.continuousCompactor.observe(usage.usagePercentage, { ...context, @@ -5081,6 +5145,7 @@ export class AgentSession { context.providersConfig ), providersConfig: context.providersConfig, + ...this.getCompactionRoutingOptions(context.options, context.modelRoutingSnapshot), }); if (pressure.shouldForceCompact) { await eventSpine.run("compaction.prepare", { @@ -5102,6 +5167,7 @@ export class AgentSession { { synthetic: true, agentInitiated: fallback?.agentInitiated ?? context.agentInitiated, + modelRoutingSnapshot: context.modelRoutingSnapshot, goalKind: fallback ? undefined : context.goalKind, goalId: fallback ? undefined : context.goalId, admissionStale: () => this.continuousCompactionAbandoned, @@ -5116,7 +5182,8 @@ export class AgentSession { const summaryId = this.pendingCompactionFollowUpSummaryId; await this.dispatchPendingFollowUp( summaryId ?? undefined, - () => this.continuousCompactionAbandoned + () => this.continuousCompactionAbandoned, + context.modelRoutingSnapshot ); if (this.pendingCompactionFollowUpSummaryId === summaryId) this.pendingCompactionFollowUpSummaryId = null; @@ -5182,7 +5249,11 @@ export class AgentSession { ...autoCompactionRequest.sendOptions, muxMetadata: autoCompactionRequest.metadata, }, - { synthetic: true, agentInitiated: autoCompactionRequest.agentInitiated } + { + synthetic: true, + agentInitiated: autoCompactionRequest.agentInitiated, + modelRoutingSnapshot: streamContext.modelRoutingSnapshot, + } ); if (!sendResult.success) { log.warn("Failed to dispatch mid-stream compaction request", { @@ -5354,7 +5425,8 @@ export class AgentSession { // Session-owned per-turn holder for mid-turn thinking changes. Passed // explicitly (not read from the field) so a preempted turn can never pick // up its replacement's holder. Absent for internal retry paths. - activeTurnThinkingOverride?: ActiveTurnThinkingOverride + activeTurnThinkingOverride?: ActiveTurnThinkingOverride, + modelRoutingSnapshot = this.captureModelRoutingSnapshot() ): Promise> { // Re-read at every pre-stream checkpoint below: dispose or shutdown can land while a // recovery-initiated stream (which carries no abortSignal) awaits commitPartial, file-change @@ -5376,7 +5448,7 @@ export class AgentSession { this.ackPendingPostCompactionStateOnStreamEnd = false; this.activeStreamHadAnyDelta = false; this.activeStreamHadPostCompactionInjection = false; - const providersConfig = this.getProvidersConfigSafe(); + const providersConfig = modelRoutingSnapshot.metadata; this.activeStreamContext = { modelString, options, @@ -5385,6 +5457,7 @@ export class AgentSession { ...(goalKind != null ? { goalKind } : {}), ...(goalId != null ? { goalId } : {}), providersConfig, + modelRoutingSnapshot, }; this.activeStreamUserMessageId = undefined; @@ -5586,6 +5659,7 @@ export class AgentSession { additionalSystemInstructions: options?.additionalSystemInstructions, maxOutputTokens: options?.maxOutputTokens, muxProviderOptions: options?.providerOptions, + modelRoutingSnapshot, agentInitiated, agentId: options?.agentId, acpPromptId, @@ -5846,6 +5920,8 @@ export class AgentSession { const retryAgentInitiated = this.activeStreamContext?.agentInitiated; const retryGoalKind = this.activeStreamContext?.goalKind; const retryGoalId = this.activeStreamContext?.goalId; + const retryRoutingSnapshot = + this.activeStreamContext?.modelRoutingSnapshot ?? this.captureModelRoutingSnapshot(); const retryOptionsForResume = retryOptions ?? { model: context.modelString, agentId: WORKSPACE_DEFAULTS.agentId, @@ -5869,7 +5945,8 @@ export class AgentSession { retryOptionsForResume, retryAgentInitiated, retryGoalKind, - retryGoalId + retryGoalId, + retryRoutingSnapshot ); this.preparingWorkspaceTurnMetadata = getWorkspaceTurnMuxMetadata( retryOptionsForResume.muxMetadata @@ -5886,7 +5963,9 @@ export class AgentSession { retryAgentInitiated, undefined, retryGoalKind, - retryGoalId + retryGoalId, + undefined, + retryRoutingSnapshot ); } finally { if (this.coordinator.isCurrentTurn(preparedTurn)) { @@ -5994,7 +6073,9 @@ export class AgentSession { context.agentInitiated, undefined, context.goalKind, - context.goalId + context.goalId, + undefined, + context.modelRoutingSnapshot ); } finally { if (this.coordinator.isCurrentTurn(preparedTurn)) { @@ -6238,6 +6319,7 @@ export class AgentSession { return; const activeModelForAbort = this.activeStreamContext?.modelString; const activeOptionsForAbort = this.activeStreamContext?.options; + const activeRoutingForAbort = this.activeStreamContext?.modelRoutingSnapshot; this.lastSystemMessageTokens = systemMessageTokens ?? this.lastSystemMessageTokens; if (activeModelForAbort) { this.updateUsageStateFromModelUsage({ @@ -6287,7 +6369,11 @@ export class AgentSession { this.activeCompactionRequest = undefined; this.resetActiveStreamState(); if (!hadCompactionRequest && activeModelForAbort && !this.continuousCompactionAbandoned) { - await this.observeContinuousCompactionAtStreamEnd(activeModelForAbort, activeOptionsForAbort); + await this.observeContinuousCompactionAtStreamEnd( + activeModelForAbort, + activeOptionsForAbort, + activeRoutingForAbort + ); if (!this.coordinator.isCurrentTurn(turn) || !this.coordinator.isCurrentOperation(operation)) return; } @@ -6331,6 +6417,7 @@ export class AgentSession { const streamEndPayload = payload; const activeStreamGoalKind = this.activeStreamContext?.goalKind; const activeStreamOptions = this.activeStreamContext?.options; + const activeRoutingSnapshot = this.activeStreamContext?.modelRoutingSnapshot; let goalContinuationRequest: { sendOptions: SendMessageOptions; @@ -6420,7 +6507,8 @@ export class AgentSession { if (!handled && !completedCompactionRequest) { await this.observeContinuousCompactionAtStreamEnd( streamEndPayload.metadata.model, - activeStreamOptions + activeStreamOptions, + activeRoutingSnapshot ); if ( !this.coordinator.isCurrentTurn(turn) || @@ -6436,7 +6524,11 @@ export class AgentSession { // not the last row, so target it by ID (stashed in onCompactionComplete). const rlmSummaryId = this.pendingCompactionFollowUpSummaryId; this.pendingCompactionFollowUpSummaryId = null; - continuedAfterCompaction = await this.dispatchPendingFollowUp(rlmSummaryId ?? undefined); + continuedAfterCompaction = await this.dispatchPendingFollowUp( + rlmSummaryId ?? undefined, + undefined, + activeRoutingSnapshot + ); if ( !this.coordinator.isCurrentTurn(turn) || !this.coordinator.isCurrentOperation(operation) @@ -6566,6 +6658,9 @@ export class AgentSession { this.emitChatEvent(payload); } }); + forward("stream-model-update", (payload) => { + this.emitChatEvent(payload); + }); forward("stream-delta", (payload) => { this.markActiveStreamHadAnyOutput(); this.emitChatEvent(payload); @@ -6644,6 +6739,7 @@ export class AgentSession { await this.waitForContinuousCompactionObservation(); await this.continuousCompactor.waitForIdle(); await this.waitForContinuousCompactionObservation(); + this.syncActiveCompactionModel(); const context = this.activeStreamContext; if ( !context || @@ -6673,6 +6769,8 @@ export class AgentSession { return; } + // Refusal fallback changes the model without starting a new workspace turn. + this.syncActiveCompactionModel(); const modelForUsage = this.activeStreamContext?.modelString; if (!modelForUsage) { return; @@ -6753,7 +6851,7 @@ export class AgentSession { streamContext?.providersConfig ?? null ), providersConfig: streamContext?.providersConfig ?? null, - openaiWireFormat: streamOptions?.providerOptions?.openai?.wireFormat, + ...this.getCompactionRoutingOptions(streamOptions), }); if (shouldInterruptForCompaction) { @@ -7598,7 +7696,8 @@ export class AgentSession { */ private async dispatchPendingFollowUp( summaryMessageId?: string, - cancelResume?: () => boolean + cancelResume?: () => boolean, + modelRoutingSnapshot?: ModelRoutingSnapshot ): Promise { if (this.coordinator.disposed || this.coordinator.closing) { return false; @@ -7907,7 +8006,8 @@ export class AgentSession { options, followUp.agentInitiated, persistedGoalKind, - persistedGoalId + persistedGoalId, + modelRoutingSnapshot ); // Await sendMessage to ensure the follow-up is persisted before returning. @@ -7917,6 +8017,7 @@ export class AgentSession { // re-enable auto-retry after a user explicitly opted out. const sendResult = await this.sendMessage(finalText, options, { synthetic: true, + modelRoutingSnapshot, agentInitiated: followUp.agentInitiated, goalKind: persistedGoalKind, // Keep the re-dispatched continuation row goal-scoped so a replaced diff --git a/src/node/services/agentStatusService.ts b/src/node/services/agentStatusService.ts index 83abc8b2c04..1f4f870d128 100644 --- a/src/node/services/agentStatusService.ts +++ b/src/node/services/agentStatusService.ts @@ -405,6 +405,7 @@ export class AgentStatusService { // would leak background LLM work past our lifecycle. if (this.stopped) return; const result = await generateWorkspaceStatus(transcript, candidates, this.aiService, { + workspaceId, streaming, recordUsage: async (modelString, usage, usageOptions) => { const recorded = await this.sessionUsageService?.recordHeadlessUsage( diff --git a/src/node/services/aiService.test.ts b/src/node/services/aiService.test.ts index b85a2ea7143..f93d0fcfd4d 100644 --- a/src/node/services/aiService.test.ts +++ b/src/node/services/aiService.test.ts @@ -10,6 +10,7 @@ import { describe, it, expect, beforeEach, afterEach, mock, spyOn } from "bun:te import { resolveModelForMetadata } from "@/common/utils/providers/modelEntries"; import { AIService, resolveMuxProjectRootForHostFs } from "./aiService"; +import { getEffectiveContextLimit } from "@/common/utils/compaction/contextLimit"; import { discoverAvailableSubagentsForToolContext } from "./turnContextAssembler"; import { normalizeAnthropicBaseURL, @@ -31,11 +32,14 @@ import { XUM_APP_ATTRIBUTION_TITLE, XUM_APP_ATTRIBUTION_URL } from "@/constants/ import type { ProviderName } from "@/common/constants/providers"; import { KNOWN_MODELS } from "@/common/constants/knownModels"; import type { CodexOauthService } from "@/node/services/codexOauthService"; +import type { TaskService } from "./taskService"; +import type { WorkflowServiceOptions } from "./workflows/WorkflowService"; import { DEFAULT_RUNTIME_CONFIG } from "@/common/constants/workspace"; import { CODEX_ENDPOINT } from "@/common/constants/codexOAuth"; import { jsonSchema, tool, type LanguageModel, type Tool } from "ai"; import { createMuxMessage } from "@/common/types/message"; +import { Ok } from "@/common/types/result"; import type { ModelMessage } from "@/common/types/message"; import type { XumToolScope } from "@/common/types/toolScope"; import type { WorkspaceMetadata } from "@/common/types/workspace"; @@ -110,12 +114,14 @@ function createBasicAIService( sessionUsageService?: SessionUsageService; devToolsService?: DevToolsService; experimentsService?: ExperimentsService; + providersConfigStore?: ProvidersConfigStore; } ): BasicAIServiceParts { const config = new Config(root); const historyService = new HistoryService(config); const initStateManager = new InitStateManager(config); - const providersConfigStore = new ProvidersConfigStore(config.rootDir); + const providersConfigStore = + options?.providersConfigStore ?? new ProvidersConfigStore(config.rootDir); const providerService = new ProviderService(config, undefined, providersConfigStore); const service = new AIService( config, @@ -302,6 +308,7 @@ function stubCommonStreamMessageDependencies(args: { metadata: WorkspaceMetadata; startStreamCalls?: TurnExecutionOptions[]; routeProvider?: ProviderName; + codexOauthAccountId?: string; allTools?: Record; workspacePathOverride?: string; historySequence?: number; @@ -375,6 +382,7 @@ function stubCommonStreamMessageDependencies(args: { wireProviderName: args.canonicalProviderName ?? providerNameFromModelString(canonicalModelString), routedThroughGateway: false, + codexOauthAccountId: args.codexOauthAccountId, ...(args.routeProvider != null ? { routeProvider: args.routeProvider } : {}), }, }); @@ -1031,7 +1039,351 @@ describe("AIService.createModel (Codex OAuth routing)", () => { }); }); +describe("AIService.captureModelRoutingSnapshot", () => { + afterEach(() => { + mock.restore(); + }); + + it.each([ + { workspaceId: undefined, project: "project", accountId: "project" }, + { workspaceId: "routing-workspace", project: "project", accountId: "workspace" }, + { workspaceId: "routing-workspace", project: undefined, accountId: "workspace" }, + { workspaceId: "unknown-workspace", project: "project", accountId: "project" }, + { workspaceId: undefined, project: undefined, accountId: "global" }, + ])("captures the factory account scope: %j", async (testCase) => { + using xumHome = new DisposableTempDir("snapshot-project-scope"); + const { config, service, providersConfigStore } = createBasicAIService(xumHome.path); + const root = path.join(xumHome.path, "root"); + const subproject = path.join(root, "subproject"); + const project = path.join(xumHome.path, "project"); + await config.editConfig((cfg) => { + cfg.projects.set(root, { + codexOauthAccountId: "root", + workspaces: [ + { + id: "routing-workspace", + name: "routing-workspace", + path: root, + subProjectPath: subproject, + }, + ], + }); + cfg.projects.set(subproject, { codexOauthAccountId: "workspace", workspaces: [] }); + cfg.projects.set(project, { codexOauthAccountId: "project", workspaces: [] }); + return cfg; + }); + providersConfigStore.saveProvidersConfig({ openai: { codexOauthDefaultAccountId: "global" } }); + const context = { + workspaceId: testCase.workspaceId, + projectPath: testCase.project ? project : undefined, + }; + const snapshot = service.captureModelRoutingSnapshot(context.workspaceId, context.projectPath); + expect(snapshot.codexOauthSelection).toEqual({ accountId: testCase.accountId, explicit: true }); + }); + + it.each(["defaults", "priority", "override"] as const)( + "copies captured route rules before later settings edits: %s", + async (change) => { + using xumHome = new DisposableTempDir("snapshot-route-config"); + const { config, service, providersConfigStore } = createBasicAIService(xumHome.path); + providersConfigStore.saveProvidersConfig({ + openai: { apiKey: "openai-key" }, + openrouter: { apiKey: "openrouter-key" }, + }); + const appConfig = config.loadConfigOrDefault(); + delete appConfig.routePriority; + delete appConfig.routeOverrides; + if (change !== "defaults") { + appConfig.routePriority = ["direct"]; + appConfig.routeOverrides = { "openai:gpt-5.5": "direct" }; + } + spyOn(config, "loadConfigOrDefault").mockReturnValue(appConfig); + const getProvidersConfig = service.getProvidersConfig.bind(service); + // Metadata projection must not expose shared mutable routing arrays or maps. + spyOn(service, "getProvidersConfig").mockImplementationOnce((raw) => { + if (change === "priority") appConfig.routePriority!.unshift("openrouter"); + else if (change === "override") appConfig.routeOverrides!["openai:gpt-5.5"] = "openrouter"; + else { + appConfig.routePriority = ["openrouter", "direct"]; + appConfig.routeOverrides = { "openai:gpt-5.5": "openrouter" }; + } + if (change === "priority") delete appConfig.routeOverrides!["openai:gpt-5.5"]; + return getProvidersConfig(raw); + }); + const snapshot = service.captureModelRoutingSnapshot(); + for (const create of [ + service.createModel.bind(service), + async ( + model: string, + _options: undefined, + opts: Parameters[2] + ) => { + const result = await service.createModelWithPinnedMetadata(model, opts); + return result.success ? Ok(result.data.model) : result; + }, + ]) { + const captured = await create("openai:gpt-5.5", undefined, { + modelRoutingSnapshot: snapshot, + }); + expect(captured.success).toBe(true); + if (captured.success && typeof captured.data !== "string") + expect(captured.data.modelId).toBe("gpt-5.5"); + const current = await create("openai:gpt-5.5", undefined, { + modelRoutingSnapshot: service.captureModelRoutingSnapshot(), + }); + expect(current.success).toBe(true); + if (current.success && typeof current.data !== "string") + expect(current.data.modelId).toBe("openai/gpt-5.5"); + } + } + ); + + it("createModel keeps captured config and selection ahead of live settings and raw overrides", async () => { + using xumHome = new DisposableTempDir("snapshot-model-options"); + const { config, service, providersConfigStore } = createBasicAIService(xumHome.path); + const projectPath = path.join(xumHome.path, "project"); + await config.editConfig((cfg) => { + cfg.projects.set(projectPath, { codexOauthAccountId: "work", workspaces: [] }); + return cfg; + }); + const requests: RecordedFetchRequest[] = []; + const fetch = createRecordingOpenAIFetch(requests); + const auth: Record = { + work: { ...TEST_CODEX_OAUTH, access: "work-access", accountId: "work-provider-id" }, + personal: { + ...TEST_CODEX_OAUTH, + access: "personal-access", + accountId: "personal-provider-id", + }, + }; + const providersConfig = { + openai: { + apiKey: "stored-api-key", + codexOauthAccounts: { + work: { label: "Work", credentials: auth.work }, + personal: { label: "Personal", credentials: auth.personal }, + }, + fetch, + }, + }; + const read = spyOn(providersConfigStore, "loadProvidersConfig").mockReturnValue( + providersConfig + ); + const getValidAuth = mock((accountId: string) => + Promise.resolve({ success: true, data: auth[accountId] }) + ); + service.turnRequestBuilderBindings.codexOauthService = { + getValidAuth, + } as unknown as CodexOauthService; + const modelRoutingSnapshot = service.captureModelRoutingSnapshot(undefined, projectPath); + await config.editConfig((cfg) => { + cfg.projects.get(projectPath)!.codexOauthAccountId = "personal"; + return cfg; + }); + read.mockReturnValue({ + openai: { ...providersConfig.openai, codexOauthDefaultAuth: "apiKey" }, + }); + const rawOverride = { openai: { apiKey: "override-api-key", fetch } }; + const pinned = await service.createModel("openai:gpt-5.5", undefined, { + projectPath, + providersConfig: rawOverride, + modelRoutingSnapshot, + }); + expect(pinned.success).toBe(true); + if (!pinned.success || typeof pinned.data === "string") + throw new Error("Expected an SDK model"); + await pinned.data.doGenerate({ + prompt: [{ role: "user", content: [{ type: "text", text: "Hello" }] }], + }); + expect(getValidAuth).toHaveBeenCalledWith("work", expect.any(Object)); + expect(getFetchUrl(requests[0].input)).toBe(CODEX_ENDPOINT); + expect(new Headers(requests[0].init?.headers).get("authorization")).toBe("Bearer work-access"); + expect(new Headers(requests[0].init?.headers).get("chatgpt-account-id")).toBe( + "work-provider-id" + ); + + // Raw overrides retain their existing precedence when no routing snapshot exists. + const overridden = await service.createModel("openai:gpt-5.5", undefined, { + providersConfig: rawOverride, + }); + expect(overridden.success).toBe(true); + if (!overridden.success || typeof overridden.data === "string") + throw new Error("Expected an SDK model"); + await overridden.data.doGenerate({ + prompt: [{ role: "user", content: [{ type: "text", text: "Hello" }] }], + }); + expect(getFetchUrl(requests[1].input)).not.toBe(CODEX_ENDPOINT); + expect(new Headers(requests[1].init?.headers).get("authorization")).toBe( + "Bearer override-api-key" + ); + + const nextSnapshot = service.captureModelRoutingSnapshot(undefined, projectPath); + expect(nextSnapshot.codexOauthSelection.accountId).toBe("personal"); + const next = await service.createModel("openai:gpt-5.5", undefined, { + projectPath, + modelRoutingSnapshot: nextSnapshot, + }); + expect(next.success).toBe(true); + if (!next.success || typeof next.data === "string") throw new Error("Expected an SDK model"); + await next.data.doGenerate({ + prompt: [{ role: "user", content: [{ type: "text", text: "Hello" }] }], + }); + expect(getFetchUrl(requests[2].input)).not.toBe(CODEX_ENDPOINT); + expect(new Headers(requests[2].init?.headers).get("authorization")).toBe( + "Bearer stored-api-key" + ); + expect(requests).toHaveLength(3); + expect(getValidAuth).toHaveBeenCalledTimes(1); + }); + + it("uses permanent provider credentials with temporary workspace routing config", async () => { + using temporaryHome = new DisposableTempDir("snapshot-temporary-routing"); + using permanentHome = new DisposableTempDir("snapshot-permanent-providers"); + const providersConfigStore = new ProvidersConfigStore(permanentHome.path); + const { config, service } = createBasicAIService(temporaryHome.path, { providersConfigStore }); + const workspaceId = "temporary-workspace"; + await config.editConfig((cfg) => { + cfg.projects.set(temporaryHome.path, { + codexOauthAccountId: "work", + workspaces: [{ id: workspaceId, name: workspaceId, path: temporaryHome.path }], + }); + return cfg; + }); + new ProvidersConfigStore(temporaryHome.path).saveProvidersConfig({ + openai: { apiKey: "temporary-api-key" }, + }); + providersConfigStore.saveProvidersConfig({ + openai: { + codexOauthAccounts: { work: { label: "Work", credentials: TEST_CODEX_OAUTH } }, + }, + }); + const requests: RecordedFetchRequest[] = []; + const loadProvidersConfig = providersConfigStore.loadProvidersConfig.bind(providersConfigStore); + const read = spyOn(providersConfigStore, "loadProvidersConfig").mockImplementation(() => { + const raw = loadProvidersConfig()!; + return { ...raw, openai: { ...raw.openai, fetch: createRecordingOpenAIFetch(requests) } }; + }); + const getValidAuth = mock(() => Promise.resolve({ success: true, data: TEST_CODEX_OAUTH })); + service.turnRequestBuilderBindings.codexOauthService = { + getValidAuth, + } as unknown as CodexOauthService; + + const snapshot = service.captureModelRoutingSnapshot(workspaceId); + expect(read).toHaveBeenCalledTimes(1); + expect(snapshot.codexOauthSelection).toEqual({ accountId: "work", explicit: true }); + expect( + getEffectiveContextLimit("openai:gpt-5.5", false, snapshot.metadata, { + codexOauthAccountId: snapshot.codexOauthSelection.accountId, + }) + ).toBe(272_000); + const pinned = await service.createModelWithPinnedMetadata("openai:gpt-5.5", { + workspaceId, + modelRoutingSnapshot: snapshot, + }); + expect(pinned.success).toBe(true); + if (!pinned.success || typeof pinned.data.model === "string") { + throw new Error("Expected a generated model"); + } + await pinned.data.model.doGenerate({ + prompt: [{ role: "user", content: [{ type: "text", text: "Hello" }] }], + }); + expect(getValidAuth).toHaveBeenCalledWith("work", expect.any(Object)); + expect(requests).toHaveLength(1); + expect(getFetchUrl(requests[0].input)).toBe(CODEX_ENDPOINT); + expect(new Headers(requests[0].init?.headers).get("authorization")).toBe( + "Bearer test-access-token" + ); + }); +}); + describe("AIService.createModelWithPinnedMetadata", () => { + it("keeps model and metadata routing together when settings change during construction", async () => { + using xumHome = new DisposableTempDir("metadata-route-snapshot"); + const { config, service, providersConfigStore } = createBasicAIService(xumHome.path); + providersConfigStore.saveProvidersConfig({ + openai: { apiKey: "openai-key" }, + openrouter: { apiKey: "openrouter-key" }, + }); + await config.editConfig((cfg) => { + cfg.routePriority = ["direct"]; + return cfg; + }); + const factory = Reflect.get(service, "providerModelFactory") as ProviderModelFactory; + const createModel = factory.createModel.bind(factory); + const resolveRoute = factory.resolveEffectiveModelString.bind(factory); + const resolvedRoutes: string[] = []; + spyOn(factory, "resolveEffectiveModelString").mockImplementation((...args) => { + const route = resolveRoute(...args); + resolvedRoutes.push(route); + return route; + }); + spyOn(factory, "createModel").mockImplementationOnce(async (...args) => { + const model = await createModel(...args); + await config.editConfig((cfg) => { + cfg.routePriority = ["openrouter", "direct"]; + return cfg; + }); + return model; + }); + + const first = await service.createModelWithPinnedMetadata("openai:gpt-5.5"); + expect(first.success).toBe(true); + if (!first.success || typeof first.data.model === "string") + throw new Error("Expected model construction to succeed"); + expect(first.data.model.modelId).toBe("gpt-5.5"); + expect(resolvedRoutes.length).toBeGreaterThan(1); + expect(new Set(resolvedRoutes)).toEqual(new Set(["openai:gpt-5.5"])); + + resolvedRoutes.length = 0; + const next = await service.createModelWithPinnedMetadata("openai:gpt-5.5"); + expect(next.success).toBe(true); + if (!next.success || typeof next.data.model === "string") + throw new Error("Expected model construction to succeed"); + expect(next.data.model.modelId).toBe("openai/gpt-5.5"); + expect(new Set(resolvedRoutes)).toEqual(new Set(["openrouter:openai/gpt-5.5"])); + }); + + it("uses the summary account snapshot after project selection changes", async () => { + using xumHome = new DisposableTempDir("headless-account-snapshot"); + const { config, service, providersConfigStore } = createBasicAIService(xumHome.path); + const workspaceId = "summary-snapshot"; + await config.editConfig((cfg) => { + cfg.projects.set(xumHome.path, { + codexOauthAccountId: "work", + workspaces: [{ id: workspaceId, name: workspaceId, path: xumHome.path }], + }); + return cfg; + }); + providersConfigStore.saveProvidersConfig({ + openai: { + codexOauthAccounts: { + work: { + label: "Work", + credentials: { + type: "oauth", + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + }, + }, + }, + }, + }); + const modelRoutingSnapshot = service.captureModelRoutingSnapshot(workspaceId); + await config.editConfig((cfg) => { + cfg.projects.get(xumHome.path)!.codexOauthAccountId = "deleted"; + return cfg; + }); + const pinned = await service.createModelWithPinnedMetadata("openai:gpt-5.5", { + workspaceId, + modelRoutingSnapshot, + }); + expect(pinned.success).toBe(true); + const next = await service.createModelWithPinnedMetadata("openai:gpt-5.5", { workspaceId }); + expect(next.success).toBe(false); + if (!next.success) expect(next.error.type).toBe("oauth_not_connected"); + }); + it("derives the pinned identity from the effective route when a coder selection falls away", async () => { using xumHome = new DisposableTempDir("pinned-metadata-fallback-away"); @@ -1087,6 +1439,7 @@ describe("AIService.streamMessage compaction boundary slicing", () => { metadata: WorkspaceMetadata, options?: { routeProvider?: ProviderName; + codexOauthAccountId?: string; allTools?: Record; postPolicyTools?: Record; sessionUsageService?: SessionUsageService; @@ -1095,6 +1448,7 @@ describe("AIService.streamMessage compaction boundary slicing", () => { canonicalModelId?: string; useRequestedModelString?: boolean; experimentsService?: ExperimentsService; + providersConfigStore?: ProvidersConfigStore; } ): StreamMessageHarness { const { config, historyService, initStateManager, service } = createBasicAIService( @@ -1102,6 +1456,7 @@ describe("AIService.streamMessage compaction boundary slicing", () => { { sessionUsageService: options?.sessionUsageService, experimentsService: options?.experimentsService, + providersConfigStore: options?.providersConfigStore, } ); const planPayloadMessageIds: string[][] = []; @@ -1122,6 +1477,7 @@ describe("AIService.streamMessage compaction boundary slicing", () => { metadata, startStreamCalls, routeProvider: options?.routeProvider, + codexOauthAccountId: options?.codexOauthAccountId, allTools: options?.allTools, effectiveModelString: options?.effectiveModelString, canonicalProviderName: options?.canonicalProviderName, @@ -1167,6 +1523,168 @@ describe("AIService.streamMessage compaction boundary slicing", () => { }; } + it("pins the turn context limit to the model account rather than the global default", async () => { + using xumHome = new DisposableTempDir("ai-service-codex-account-snapshot"); + const projectPath = path.join(xumHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + const workspaceId = "account-snapshot"; + const harness = createHarness( + xumHome.path, + createLocalWorkspaceMetadata(workspaceId, projectPath), + { + effectiveModelString: "openai:gpt-5.5", + codexOauthAccountId: "work", + } + ); + new ProvidersConfigStore(harness.config.rootDir).saveProvidersConfig({ + openai: { + apiKey: "test-key", + codexOauthDefaultAccountId: "missing", + codexOauthAccounts: { + work: { + label: "Work", + credentials: { + type: "oauth", + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + }, + }, + }, + }, + }); + const result = await harness.service.streamMessage({ + messages: [createMuxMessage("user", "user", "continue")], + workspaceId, + modelString: "openai:gpt-5.5", + thinkingLevel: "off", + }); + expect(result.success).toBe(true); + const snapshot = harness.startStreamCalls[0]?.providersConfigSnapshot; + expect(snapshot?.openai?.codexOauthDefaultAccountId).toBe("work"); + expect(getEffectiveContextLimit("openai:gpt-5.5", false, snapshot)).toBe(272_000); + expect( + new ProvidersConfigStore(harness.config.rootDir).loadProvidersConfig()?.openai + ?.codexOauthDefaultAccountId + ).toBe("missing"); + }); + + it("keeps context limits and real factory routing on the raw snapshot during provider edits", async () => { + using xumHome = new DisposableTempDir("stream-account-snapshot"); + const projectPath = path.join(xumHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + const workspaceId = "stream-snapshot"; + const store = new ProvidersConfigStore(xumHome.path); + const harness = createHarness( + xumHome.path, + createLocalWorkspaceMetadata(workspaceId, projectPath), + { providersConfigStore: store } + ); + const factory = Reflect.get(harness.service, "providerModelFactory") as ProviderModelFactory; + spyOn(factory, "resolveAndCreateModel").mockRestore(); + const requests: RecordedFetchRequest[] = []; + configureOpenAICodexOAuth(harness.service, store, requests); + const read = spyOn(store, "loadProvidersConfig"); + const getProvidersConfig = harness.service.getProvidersConfig.bind(harness.service); + // Change the backing store after the raw read, before metadata projection. + spyOn(harness.service, "getProvidersConfig").mockImplementationOnce((raw) => { + read.mockReturnValue({ + openai: { + apiKey: "replacement-api-key", + codexOauthDefaultAuth: "apiKey", + fetch: createRecordingOpenAIFetch(requests), + }, + }); + return getProvidersConfig(raw); + }); + const modelRoutingSnapshot = harness.service.captureModelRoutingSnapshot(workspaceId); + expect(read).toHaveBeenCalledTimes(1); + expect(getEffectiveContextLimit("openai:gpt-5.5", false, modelRoutingSnapshot.metadata)).toBe( + 272_000 + ); + const result = await harness.service.streamMessage({ + messages: [createMuxMessage("user", "user", "continue")], + workspaceId, + modelString: "openai:gpt-5.5", + thinkingLevel: "off", + modelRoutingSnapshot, + }); + expect(result.success).toBe(true); + const stream = harness.startStreamCalls[0]; + expect(getEffectiveContextLimit("openai:gpt-5.5", false, stream.providersConfigSnapshot)).toBe( + 272_000 + ); + if (typeof stream.model === "string") throw new Error("Expected a generated model"); + await stream.model.doGenerate({ + prompt: [{ role: "user", content: [{ type: "text", text: "Hello" }] }], + }); + expect(requests).toHaveLength(1); + expect(getFetchUrl(requests[0].input)).toBe(CODEX_ENDPOINT); + expect(new Headers(requests[0].init?.headers).get("authorization")).toBe( + "Bearer test-access-token" + ); + + const nextSnapshot = harness.service.captureModelRoutingSnapshot(workspaceId); + expect(getEffectiveContextLimit("openai:gpt-5.5", false, nextSnapshot.metadata)).toBe( + 1_050_000 + ); + const nextResult = await harness.service.streamMessage({ + messages: [createMuxMessage("next-user", "user", "continue")], + workspaceId, + modelString: "openai:gpt-5.5", + thinkingLevel: "off", + modelRoutingSnapshot: nextSnapshot, + }); + expect(nextResult.success).toBe(true); + const nextStream = harness.startStreamCalls[1]; + expect( + getEffectiveContextLimit("openai:gpt-5.5", false, nextStream.providersConfigSnapshot) + ).toBe(1_050_000); + if (typeof nextStream.model === "string") throw new Error("Expected a generated model"); + await nextStream.model.doGenerate({ + prompt: [{ role: "user", content: [{ type: "text", text: "Hello" }] }], + }); + expect(requests).toHaveLength(2); + expect(getFetchUrl(requests[1].input)).not.toBe(CODEX_ENDPOINT); + expect(new Headers(requests[1].init?.headers).get("authorization")).toBe( + "Bearer replacement-api-key" + ); + }); + + it("preserves implicit API-key routing in a captured turn snapshot", async () => { + using xumHome = new DisposableTempDir("stream-implicit-api-snapshot"); + const projectPath = path.join(xumHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + const workspaceId = "implicit-api-snapshot"; + const harness = createHarness( + xumHome.path, + createLocalWorkspaceMetadata(workspaceId, projectPath), + { + effectiveModelString: "openai:gpt-5.5", + codexOauthAccountId: "default", + } + ); + new ProvidersConfigStore(harness.config.rootDir).saveProvidersConfig({ + openai: { apiKey: "test-key" }, + }); + const modelRoutingSnapshot = harness.service.captureModelRoutingSnapshot(workspaceId); + const result = await harness.service.streamMessage({ + messages: [createMuxMessage("user", "user", "continue")], + workspaceId, + modelString: "openai:gpt-5.5", + thinkingLevel: "off", + modelRoutingSnapshot, + }); + expect(result.success).toBe(true); + expect( + getEffectiveContextLimit( + "openai:gpt-5.5", + false, + harness.startStreamCalls[0]?.providersConfigSnapshot + ) + ).toBe(1_050_000); + }); + interface AdvisorRuntimeForTests { createModel: (modelString: string) => Promise; takeToolCallSnapshot: (toolCallId: string) => @@ -2629,6 +3147,420 @@ describe("AIService.streamMessage compaction boundary slicing", () => { expect(typeof sessionUsageDeltaRecord.timestamp).toBe("number"); }); + it.each(["project-account", "default-account", "priority", "override"] as const)( + "background workflow continuations retain accepted routing after a %s update", + async (change) => { + using xumHome = new DisposableTempDir("workflow-routing-snapshot"); + const projectPath = path.join(xumHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + const workspaceId = "workflow-routing"; + const store = new ProvidersConfigStore(xumHome.path); + const harness = createHarness( + xumHome.path, + createLocalWorkspaceMetadata(workspaceId, projectPath), + { + providersConfigStore: store, + useRequestedModelString: true, + } + ); + await harness.config.editConfig((cfg) => ({ + ...cfg, + routePriority: ["direct"], + routeOverrides: {}, + projects: new Map([ + [ + projectPath, + { + workspaces: [], + ...(change === "project-account" ? { codexOauthAccountId: "work" } : {}), + }, + ], + ]), + })); + const requests: RecordedFetchRequest[] = []; + const auth: Record = { + work: { ...TEST_CODEX_OAUTH, access: "workflow-work-access", accountId: "work-id" }, + personal: { + ...TEST_CODEX_OAUTH, + access: "workflow-personal-access", + accountId: "personal-id", + }, + }; + const providersConfig = { + openai: { + apiKey: "openai-key", + codexOauthDefaultAccountId: "work", + codexOauthAccounts: { + work: { label: "Work", credentials: auth.work }, + personal: { label: "Personal", credentials: auth.personal }, + }, + fetch: createRecordingOpenAIFetch(requests, "gpt-5.5"), + }, + openrouter: { apiKey: "openrouter-key" }, + }; + const read = spyOn(store, "loadProvidersConfig").mockReturnValue(providersConfig); + harness.service.turnRequestBuilderBindings.codexOauthService = { + getValidAuth: (accountId: string) => Promise.resolve(Ok(auth[accountId])), + } as unknown as CodexOauthService; + const bindings = harness.service.turnRequestBuilderBindings; + bindings.taskService = {} as unknown as TaskService; + type Sender = NonNullable; + const sendMessage = mock(() => Promise.resolve(Ok(undefined))); + bindings.workflowResultContinuationSender = { + isWorkflowInvocationCurrent: () => Promise.resolve(true), + sendMessage, + }; + const snapshot = harness.service.captureModelRoutingSnapshot(workspaceId); + expect( + ( + await harness.service.streamMessage({ + messages: [createMuxMessage("user", "user", "Start the workflow")], + workspaceId, + modelString: "openai:gpt-5.5", + thinkingLevel: "off", + experiments: { dynamicWorkflows: true }, + modelRoutingSnapshot: snapshot, + }) + ).success + ).toBe(true); + const workflowService = harness.getToolsForModelSpy.mock.calls.at(-1)?.[1].workflowService; + if (!workflowService) throw new Error("Expected workflow service"); + const onTerminal = Reflect.get(workflowService, "onBackgroundRunTerminal") as NonNullable< + WorkflowServiceOptions["onBackgroundRunTerminal"] + >; + // Exercise the sender fallback when the live task-attention binding is unavailable. + bindings.taskService = undefined; + if (change === "default-account") { + read.mockReturnValue({ + ...providersConfig, + openai: { ...providersConfig.openai, codexOauthDefaultAccountId: "personal" }, + }); + } else { + await harness.config.editConfig((cfg) => { + if (change === "project-account") + cfg.projects.get(projectPath)!.codexOauthAccountId = "personal"; + else if (change === "priority") cfg.routePriority = ["openrouter", "direct"]; + else cfg.routeOverrides = { "openai:gpt-5.5": "openrouter" }; + return cfg; + }); + } + const terminalEvent: Parameters[0] = { + runId: "wfr_routing", + status: "completed", + result: { reportMarkdown: "Done" }, + run: { + id: "wfr_routing", + workspaceId, + status: "completed", + workflow: { + name: "routing", + description: "Test routing", + scope: "project", + executable: true, + }, + source: "return {}", + sourceHash: "test-hash", + args: {}, + events: [], + steps: [], + createdAt: new Date(0).toISOString(), + updatedAt: new Date(0).toISOString(), + }, + }; + const noteWorkflowRunTerminalAttention = mock(() => undefined); + bindings.taskService = { noteWorkflowRunTerminalAttention } as unknown as TaskService; + await onTerminal(terminalEvent); + expect(noteWorkflowRunTerminalAttention).toHaveBeenCalledWith({ + ownerWorkspaceId: workspaceId, + runId: "wfr_routing", + status: "completed", + }); + expect(sendMessage).not.toHaveBeenCalled(); + bindings.taskService = undefined; + await onTerminal(terminalEvent); + expect(sendMessage).toHaveBeenCalledTimes(1); + const [, message, options, internal] = sendMessage.mock.calls[0]; + expect(internal?.modelRoutingSnapshot).toBe(snapshot); + expect(internal?.requireIdle).toBe(true); + expect(JSON.stringify({ message, options })).not.toContain("workflow-work-access"); + const continued = await harness.service.createModel( + options.model, + options.providerOptions, + internal + ); + if (!continued.success || typeof continued.data === "string") + throw new Error("Expected continued SDK model"); + await continued.data.doGenerate({ + prompt: [{ role: "user", content: [{ type: "text", text: "Hello" }] }], + }); + expect(getFetchUrl(requests[0].input)).toBe(CODEX_ENDPOINT); + expect(new Headers(requests[0].init?.headers).get("authorization")).toBe( + "Bearer workflow-work-access" + ); + const factory = Reflect.get(harness.service, "providerModelFactory") as ProviderModelFactory; + spyOn(factory, "resolveAndCreateModel").mockRestore(); + expect( + ( + await harness.service.streamMessage({ + messages: [createMuxMessage("next-user", "user", "New user turn")], + workspaceId, + modelString: options.model, + thinkingLevel: "off", + modelRoutingSnapshot: harness.service.captureModelRoutingSnapshot(workspaceId), + }) + ).success + ).toBe(true); + const nextModel = harness.startStreamCalls.at(-1)?.model; + if (!nextModel || typeof nextModel === "string") throw new Error("Expected new SDK model"); + if (change === "priority" || change === "override") { + expect(nextModel.modelId).toBe("openai/gpt-5.5"); + } else { + await nextModel.doGenerate({ + prompt: [{ role: "user", content: [{ type: "text", text: "Hello" }] }], + }); + expect(new Headers(requests[1].init?.headers).get("authorization")).toBe( + "Bearer workflow-personal-access" + ); + } + } + ); + + it.each([ + { toolName: "advisor", change: "priority" }, + { toolName: "advisor", change: "override" }, + { toolName: "intuition", change: "priority" }, + { toolName: "intuition", change: "override" }, + ] as const)( + "$toolName keeps captured route $change after settings change", + async ({ toolName, change }) => { + using xumHome = new DisposableTempDir("nested-tool-route-config"); + const projectPath = path.join(xumHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + const workspaceId = "nested-route-config"; + const store = new ProvidersConfigStore(xumHome.path); + store.saveProvidersConfig({ + openai: { apiKey: "openai-key" }, + openrouter: { apiKey: "openrouter-key" }, + }); + const experimentsService = new ExperimentsService({ + telemetryService: new TelemetryService(xumHome.path), + xumHome: xumHome.path, + }); + spyOn(experimentsService, "isExperimentEnabled").mockImplementation( + (id) => id === EXPERIMENT_IDS.MEMORY_INTUITION + ); + const harness = createHarness( + xumHome.path, + createLocalWorkspaceMetadata(workspaceId, projectPath), + { + providersConfigStore: store, + experimentsService, + useRequestedModelString: true, + } + ); + const factory = Reflect.get(harness.service, "providerModelFactory") as ProviderModelFactory; + spyOn(factory, "resolveAndCreateModel").mockRestore(); + harness.service.turnRequestBuilderBindings.memoryService = new MemoryService( + harness.config, + new MemoryMetaService(xumHome.path) + ); + await enableAdvisorForHarness(harness, "openai:gpt-5.5"); + await harness.config.editConfig((cfg) => ({ + ...cfg, + routePriority: ["direct"], + routeOverrides: {}, + })); + const snapshot = harness.service.captureModelRoutingSnapshot(workspaceId); + await harness.config.editConfig((cfg) => { + if (change === "priority") cfg.routePriority = ["openrouter", "direct"]; + else cfg.routeOverrides = { "openai:gpt-5.5": "openrouter" }; + return cfg; + }); + const startTurn = (modelRoutingSnapshot: typeof snapshot) => + harness.service.streamMessage({ + messages: [createMuxMessage("user", "user", "Continue")], + workspaceId, + modelString: "openai:gpt-5.5", + thinkingLevel: "off", + experiments: { advisorTool: true, memory: true }, + modelRoutingSnapshot, + }); + const getRuntime = () => { + const config = harness.getToolsForModelSpy.mock.calls.at(-1)?.[1]; + const runtime = toolName === "advisor" ? config?.advisorRuntime : config?.intuitionRuntime; + if (!runtime) throw new Error("Expected tool runtime"); + return runtime; + }; + expect((await startTurn(snapshot)).success).toBe(true); + const capturedParentModel = harness.startStreamCalls.at(-1)?.model; + if (!capturedParentModel || typeof capturedParentModel === "string") { + throw new Error("Expected parent SDK model"); + } + expect(capturedParentModel.modelId).toBe("gpt-5.5"); + const captured = await getRuntime().createModel("openai:gpt-5.5"); + if (typeof captured.model === "string") throw new Error("Expected SDK model"); + expect(captured.model.modelId).toBe("gpt-5.5"); + expect( + (await startTurn(harness.service.captureModelRoutingSnapshot(workspaceId))).success + ).toBe(true); + const parentModel = harness.startStreamCalls.at(-1)?.model; + if (!parentModel || typeof parentModel === "string") + throw new Error("Expected parent SDK model"); + expect(parentModel.modelId).toBe("openai/gpt-5.5"); + const current = await getRuntime().createModel("openai:gpt-5.5"); + if (typeof current.model === "string") throw new Error("Expected SDK model"); + expect(current.model.modelId).toBe("openai/gpt-5.5"); + } + ); + + it.each([ + { toolName: "advisor", change: "project-account" }, + { toolName: "advisor", change: "default-account" }, + { toolName: "advisor", change: "preference" }, + { toolName: "intuition", change: "project-account" }, + { toolName: "intuition", change: "default-account" }, + { toolName: "intuition", change: "preference" }, + ] as const)( + "$toolName keeps accepted routing after a $change update", + async ({ toolName, change }) => { + using xumHome = new DisposableTempDir("nested-tool-routing-snapshot"); + const projectPath = path.join(xumHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + const workspaceId = "nested-tool-routing"; + const providersStore = new ProvidersConfigStore(xumHome.path); + const experimentsService = new ExperimentsService({ + telemetryService: new TelemetryService(xumHome.path), + xumHome: xumHome.path, + }); + spyOn(experimentsService, "isExperimentEnabled").mockImplementation( + (id) => id === EXPERIMENT_IDS.MEMORY_INTUITION + ); + const harness = createHarness( + xumHome.path, + createLocalWorkspaceMetadata(workspaceId, projectPath), + { + providersConfigStore: providersStore, + experimentsService, + useRequestedModelString: true, + codexOauthAccountId: "work", + } + ); + harness.service.turnRequestBuilderBindings.memoryService = new MemoryService( + harness.config, + new MemoryMetaService(xumHome.path) + ); + await enableAdvisorForHarness(harness, "openai:gpt-5.5"); + await harness.config.editConfig((cfg) => { + cfg.projects.set(projectPath, { + workspaces: [], + ...(change !== "default-account" ? { codexOauthAccountId: "work" } : {}), + }); + return cfg; + }); + const requests: RecordedFetchRequest[] = []; + const auth: Record = { + work: { ...TEST_CODEX_OAUTH, access: "work-access", accountId: "work-provider-id" }, + personal: { + ...TEST_CODEX_OAUTH, + access: "personal-access", + accountId: "personal-provider-id", + }, + }; + const providersConfig = { + openai: { + apiKey: "test-api-key", + wireFormat: "responses" as const, + codexOauthDefaultAccountId: "work", + codexOauthAccounts: { + work: { label: "Work", credentials: auth.work }, + personal: { label: "Personal", credentials: auth.personal }, + }, + fetch: createRecordingOpenAIFetch(requests, "gpt-5.5"), + }, + }; + const read = spyOn(providersStore, "loadProvidersConfig").mockReturnValue(providersConfig); + const getValidAuth = mock((accountId: string) => Promise.resolve(Ok(auth[accountId]))); + harness.service.turnRequestBuilderBindings.codexOauthService = { + getValidAuth, + } as unknown as CodexOauthService; + const startTurn = () => + harness.service.streamMessage({ + messages: [createMuxMessage("user", "user", "Continue")], + workspaceId, + modelString: "openai:gpt-5.5", + thinkingLevel: "off", + // Nested tools use their own options, not the parent chat's settings. + muxProviderOptions: { + openai: { wireFormat: "chatCompletions", serviceTier: "priority" }, + }, + experiments: { advisorTool: true, memory: true }, + modelRoutingSnapshot: harness.service.captureModelRoutingSnapshot(workspaceId), + }); + expect((await startTurn()).success).toBe(true); + const snapshot = harness.startStreamCalls[0]?.providersConfigSnapshot; + const getRuntime = () => { + const toolConfig = harness.getToolsForModelSpy.mock.calls.at(-1)?.[1]; + const runtime = + toolName === "advisor" ? toolConfig?.advisorRuntime : toolConfig?.intuitionRuntime; + if (!runtime) throw new Error("Expected " + toolName + " runtime"); + return runtime; + }; + const runtime = getRuntime(); + if (change === "project-account") { + await harness.config.editConfig((cfg) => { + cfg.projects.get(projectPath)!.codexOauthAccountId = "personal"; + return cfg; + }); + } else { + read.mockReturnValue({ + openai: { + ...providersConfig.openai, + ...(change === "default-account" + ? { codexOauthDefaultAccountId: "personal" } + : { codexOauthDefaultAuth: "apiKey" }), + }, + }); + } + const createModel = spyOn(harness.service, "createModel"); + const created = await runtime.createModel("openai:gpt-5.5"); + if (typeof created.model === "string") throw new Error("Expected an SDK model"); + await created.model.doGenerate({ + prompt: [{ role: "user", content: [{ type: "text", text: "Hello" }] }], + }); + expect(getFetchUrl(requests[0].input)).toBe(CODEX_ENDPOINT); + expect(new Headers(requests[0].init?.headers).get("authorization")).toBe( + "Bearer work-access" + ); + expect(new Headers(requests[0].init?.headers).get("chatgpt-account-id")).toBe( + "work-provider-id" + ); + expect(created.optionsProvidersConfig?.openai).toEqual(snapshot?.openai); + expect(created.optionsMuxProviderOptions?.openai?.wireFormat).toBe("responses"); + expect(created.optionsMuxProviderOptions?.openai?.serviceTier).toBeUndefined(); + expect( + createModel.mock.calls[0]?.[2]?.modelRoutingSnapshot?.codexOauthSelection.accountId + ).toBe("work"); + + expect((await startTurn()).success).toBe(true); + const next = await getRuntime().createModel("openai:gpt-5.5"); + if (typeof next.model === "string") throw new Error("Expected an SDK model"); + await next.model.doGenerate({ + prompt: [{ role: "user", content: [{ type: "text", text: "Hello" }] }], + }); + expect(requests).toHaveLength(2); + expect(new Headers(requests[1].init?.headers).get("authorization")).toBe( + change === "preference" ? "Bearer test-api-key" : "Bearer personal-access" + ); + expect(getFetchUrl(requests[1].input) === CODEX_ENDPOINT).toBe(change !== "preference"); + expect(next.optionsProvidersConfig?.openai?.codexOauthDefaultAuth).toBe( + change === "preference" ? "apiKey" : undefined + ); + expect( + createModel.mock.calls[1]?.[2]?.modelRoutingSnapshot?.codexOauthSelection.accountId + ).toBe(change === "preference" ? "work" : "personal"); + } + ); + it.each(["advisor", "intuition"] as const)( "records API-equivalent costs for %s tool usage through Codex OAuth", async (toolName) => { diff --git a/src/node/services/aiService.ts b/src/node/services/aiService.ts index 92aaacf28a2..2c93b15329d 100644 --- a/src/node/services/aiService.ts +++ b/src/node/services/aiService.ts @@ -67,6 +67,9 @@ import type { HistoryService } from "./historyService"; import type { SessionUsageService } from "./sessionUsageService"; import type { ProvidersConfig } from "@/common/config/schemas/providersConfig"; +import type { ModelRoutingSnapshot } from "./modelRoutingSnapshot"; +import { getCodexOauthProjectPath } from "@/common/utils/providers/codexOauthRouting"; +import { getCodexOauthAccountId } from "@/node/utils/codexOauthAuth"; import { getProjects, isMultiProject } from "@/common/utils/multiProject"; import { resolveMemoryProjectIdentity, @@ -313,8 +316,36 @@ export class AIService extends EventEmitter { } } - getProvidersConfig(): ProvidersConfigMap | null { - return this.providerService.getConfig(); + getProvidersConfig(providersConfig?: ProvidersConfig): ProvidersConfigMap | null { + return this.providerService.getConfig(providersConfig); + } + + /** Keep model credentials and compaction limits on one backend-only snapshot for the turn. */ + captureModelRoutingSnapshot(workspaceId?: string, projectPath?: string): ModelRoutingSnapshot { + // CLI routing config can be temporary. Read credentials from the injected provider store. + const providersConfig = this.providersConfigStore.loadProvidersConfig() ?? {}; + const workspace = workspaceId ? this.config.findWorkspace(workspaceId) : undefined; + const routingProjectPath = getCodexOauthProjectPath(workspace) ?? projectPath; + const appConfig = this.config.loadConfigOrDefault(); + const projectAccountId = routingProjectPath + ? appConfig.projects.get(routingProjectPath)?.codexOauthAccountId + : undefined; + // Copy routing rules before metadata reads. Nested models need their own captured overrides. + const routeConfig = { + routePriority: [...(appConfig.routePriority ?? ["direct"])], + routeOverrides: { ...appConfig.routeOverrides }, + }; + return { + providersConfig, + routeConfig, + metadata: this.getProvidersConfig(providersConfig), + codexOauthSelection: { + accountId: getCodexOauthAccountId(providersConfig.openai, projectAccountId), + explicit: + projectAccountId !== undefined || + providersConfig.openai?.codexOauthDefaultAccountId !== undefined, + }, + }; } private emitEngineEvent(event: TurnEngineEvent): void | Promise { @@ -543,11 +574,18 @@ export class AIService extends EventEmitter { opts?: { agentInitiated?: boolean; workspaceId?: string; + projectPath?: string; /** Snapshot pass-through (see ProviderModelFactory.createModel). */ providersConfig?: ProvidersConfig; + modelRoutingSnapshot?: ModelRoutingSnapshot; } ): Promise> { - return this.providerModelFactory.createModel(modelString, muxProviderOptions, opts); + return this.providerModelFactory.createModel(modelString, muxProviderOptions, { + ...opts, + providersConfig: opts?.modelRoutingSnapshot?.providersConfig ?? opts?.providersConfig, + codexOauthSelection: opts?.modelRoutingSnapshot?.codexOauthSelection, + routeConfig: opts?.modelRoutingSnapshot?.routeConfig, + }); } /** @@ -560,12 +598,23 @@ export class AIService extends EventEmitter { */ async createModelWithPinnedMetadata( modelString: string, - opts?: { agentInitiated?: boolean; workspaceId?: string } + opts?: { + agentInitiated?: boolean; + workspaceId?: string; + projectPath?: string; + modelRoutingSnapshot?: ModelRoutingSnapshot; + } ): Promise> { - const providersConfig = this.providersConfigStore.loadProvidersConfig() ?? {}; + // Keep model construction and metadata on one route, including callers without an outer snapshot. + const snapshot = + opts?.modelRoutingSnapshot ?? + this.captureModelRoutingSnapshot(opts?.workspaceId, opts?.projectPath); + const providersConfig = snapshot.providersConfig; const result = await this.providerModelFactory.createModel(modelString, undefined, { ...opts, providersConfig, + codexOauthSelection: snapshot.codexOauthSelection, + routeConfig: snapshot.routeConfig, }); if (!result.success) { return result; @@ -579,7 +628,8 @@ export class AIService extends EventEmitter { const effectiveModelString = this.providerModelFactory.resolveEffectiveModelString( modelString, undefined, - providersConfig + providersConfig, + snapshot.routeConfig ); const metadataSeed = effectiveModelString.startsWith("coder:") ? modelString diff --git a/src/node/services/branchSummary.test.ts b/src/node/services/branchSummary.test.ts index 4de71436580..6ede90c2414 100644 --- a/src/node/services/branchSummary.test.ts +++ b/src/node/services/branchSummary.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, spyOn, test } from "bun:test"; +import { describe, expect, mock, spyOn, test } from "bun:test"; import * as fs from "node:fs/promises"; import * as path from "node:path"; @@ -37,6 +37,7 @@ import { type SideChannelMetadata, } from "./branchSummary"; import { createTestHistoryService } from "./testHistoryService"; +import type { ModelRoutingSnapshot } from "./modelRoutingSnapshot"; function finishChunk(unified: "stop" | "length" = "stop"): LanguageModelV3StreamPart { return { @@ -95,6 +96,12 @@ function fakeAiService( const workspaceModel = opts?.workspaceModel === undefined ? "anthropic:claude-haiku-4-5" : opts.workspaceModel; return { + captureModelRoutingSnapshot: () => ({ + providersConfig: {}, + metadata: null, + routeConfig: { routePriority: ["direct"], routeOverrides: {} }, + codexOauthSelection: { accountId: "default", explicit: false }, + }), createModelWithPinnedMetadata: ((modelString: string) => { opts?.onCreateModel?.(modelString); if (!model) { @@ -750,6 +757,92 @@ describe("maybeAppendAbandonedBranchSummary", () => { } }); + test.each([ + { failure: "creation", change: "account" }, + { failure: "creation", change: "route" }, + { failure: "stream", change: "account" }, + { failure: "stream", change: "route" }, + ] as const)( + "pins summary routing across a $failure failure and $change change", + async ({ failure, change }) => { + const { historyService, cleanup } = await createTestHistoryService(); + try { + const settings: ModelRoutingSnapshot = { + providersConfig: { openai: { apiKey: "summary-secret" } }, + metadata: null, + routeConfig: { routePriority: ["direct"], routeOverrides: {} }, + codexOauthSelection: { accountId: "work", explicit: true }, + }; + const captureModelRoutingSnapshot = mock(() => structuredClone(settings)); + const changeSettings = () => { + if (change === "account") settings.codexOauthSelection.accountId = "personal"; + else settings.routeConfig.routePriority.unshift("openrouter"); + }; + const failedModel = new MockLanguageModelV3({ + doStream: () => { + changeSettings(); + throw new Error("First summary stream fails"); + }, + }); + const successfulModel = summaryModel("The abandoned branch identifies the race."); + let firstAttempt = true; + const createModelWithPinnedMetadata = mock< + BranchSummaryAiService["createModelWithPinnedMetadata"] + >((modelString) => { + if (firstAttempt) { + firstAttempt = false; + if (failure === "creation") { + changeSettings(); + return Promise.resolve(Err({ type: "oauth_not_connected", provider: "openai" })); + } + return Promise.resolve(Ok({ model: failedModel, metadataModel: modelString })); + } + return Promise.resolve(Ok({ model: successfulModel, metadataModel: modelString })); + }); + const aiService: BranchSummaryAiService = { + ...fakeAiService(successfulModel), + captureModelRoutingSnapshot, + createModelWithPinnedMetadata, + }; + const generate = () => + maybeAppendAbandonedBranchSummary({ + historyService, + aiService, + workspaceId: "ws-routing", + abandonedMessages: meatyExchange("routing"), + experiments: RLM_ON, + modelCandidates: ["openai:gpt-5.5", "openai:gpt-5.3-codex"], + }); + expect(await generate()).not.toBeNull(); + expect(captureModelRoutingSnapshot).toHaveBeenCalledTimes(1); + expect(captureModelRoutingSnapshot).toHaveBeenCalledWith("ws-routing"); + expect(createModelWithPinnedMetadata).toHaveBeenCalledTimes(2); + const original = createModelWithPinnedMetadata.mock.calls[0]?.[1]?.modelRoutingSnapshot; + expect(original?.codexOauthSelection.accountId).toBe("work"); + expect(original?.routeConfig.routePriority).toEqual(["direct"]); + expect(createModelWithPinnedMetadata.mock.calls[1]?.[1]?.modelRoutingSnapshot).toBe( + original + ); + expect(await generate()).not.toBeNull(); + expect(captureModelRoutingSnapshot).toHaveBeenCalledTimes(2); + expect(createModelWithPinnedMetadata).toHaveBeenCalledTimes(3); + const next = createModelWithPinnedMetadata.mock.calls[2]?.[1]?.modelRoutingSnapshot; + expect(next).not.toBe(original); + expect(next?.codexOauthSelection.accountId).toBe( + change === "account" ? "personal" : "work" + ); + expect(next?.routeConfig.routePriority).toEqual( + change === "route" ? ["openrouter", "direct"] : ["direct"] + ); + const history = await historyService.getHistoryFromLatestBoundary("ws-routing"); + expect(history.success).toBe(true); + expect(JSON.stringify(history)).not.toContain("summary-secret"); + } finally { + await cleanup(); + } + } + ); + test("generation failure skips the row and never throws", async () => { const { historyService, cleanup } = await createTestHistoryService(); try { @@ -842,6 +935,7 @@ describe("maybeAppendAbandonedBranchSummary", () => { // and workspace removal waits forever on the background drain. const base = fakeAiService(null); const wedgedCreation: BranchSummaryAiService = { + captureModelRoutingSnapshot: base.captureModelRoutingSnapshot, createModelWithPinnedMetadata: () => new Promise(() => undefined), getWorkspaceMetadata: base.getWorkspaceMetadata, }; @@ -1385,6 +1479,7 @@ describe("branch summary placement on fork/truncate flows", () => { }); const model = summaryModel("The abandoned branch context both requests need."); const gatedAiService: BranchSummaryAiService = { + captureModelRoutingSnapshot: fakeAiService(model).captureModelRoutingSnapshot, createModelWithPinnedMetadata: (async (...createArgs) => { await modelGate; return fakeAiService(model).createModelWithPinnedMetadata(...createArgs); @@ -1533,6 +1628,7 @@ describe("branch summary placement on fork/truncate flows", () => { }); const model = summaryModel("A summary that must never land after removal."); const gatedAiService: BranchSummaryAiService = { + captureModelRoutingSnapshot: fakeAiService(model).captureModelRoutingSnapshot, createModelWithPinnedMetadata: (async (...createArgs) => { await modelGate; return fakeAiService(model).createModelWithPinnedMetadata(...createArgs); diff --git a/src/node/services/branchSummary.ts b/src/node/services/branchSummary.ts index d28522f41e3..edbe8225c0e 100644 --- a/src/node/services/branchSummary.ts +++ b/src/node/services/branchSummary.ts @@ -42,6 +42,7 @@ import { } from "@/constants/streamDrain"; import type { AIService } from "./aiService"; +import type { ModelRoutingSnapshot } from "./modelRoutingSnapshot"; import type { HistoryService } from "./historyService"; import { runLanguageModelCleanup } from "./languageModelCleanup"; import { log } from "./log"; @@ -60,7 +61,7 @@ export const BRANCH_SUMMARY_LABEL = "Summary of the abandoned branch:"; */ export type BranchSummaryAiService = Pick< AIService, - "createModelWithPinnedMetadata" | "getWorkspaceMetadata" + "createModelWithPinnedMetadata" | "captureModelRoutingSnapshot" | "getWorkspaceMetadata" >; /** Send-option experiment flags relevant to RLM gating (subset of ExperimentsSchema). */ @@ -325,6 +326,7 @@ export function trackPendingUsageWrite(workspaceId: string, write: Promise async function generateAbandonedBranchSummaryText(input: { aiService: BranchSummaryAiService; + modelRoutingSnapshot: ModelRoutingSnapshot; /** * Routes the side-channel request into the workspace's devtools.jsonl: * model creation installs its API-debug middleware only when a workspaceId @@ -387,6 +389,7 @@ async function generateAbandonedBranchSummaryText(input: { const modelPromise = input.aiService.createModelWithPinnedMetadata(modelString, { agentInitiated: true, workspaceId: input.workspaceId, + modelRoutingSnapshot: input.modelRoutingSnapshot, }); const modelResult = await Promise.race([modelPromise, deadline]); if (modelResult === null) { @@ -728,6 +731,8 @@ export async function maybeAppendAbandonedBranchSummary( return null; } + // One accepted summary keeps its routing across candidate and stream retries. + const modelRoutingSnapshot = input.aiService.captureModelRoutingSnapshot(input.workspaceId); const candidates = input.modelCandidates ?? (await getSideChannelModelCandidates(input.aiService, input.workspaceId)); @@ -738,6 +743,7 @@ export async function maybeAppendAbandonedBranchSummary( const sessionUsageService = input.sessionUsageService; const summaryText = await generateAbandonedBranchSummaryText({ aiService: input.aiService, + modelRoutingSnapshot, workspaceId: input.workspaceId, candidates, system: buildAbandonedBranchSummarySystemPrompt(), diff --git a/src/node/services/codexOauthService.test.ts b/src/node/services/codexOauthService.test.ts index df8209b13e9..d56ce4c1f63 100644 --- a/src/node/services/codexOauthService.test.ts +++ b/src/node/services/codexOauthService.test.ts @@ -1,10 +1,25 @@ -import type { ProvidersConfigStore } from "@/node/config"; -import { describe, it, expect, beforeEach, afterEach } from "bun:test"; - -import type { Result } from "@/common/types/result"; -import { Ok } from "@/common/types/result"; +import { Config, FileLeaseManager, ProvidersConfigStore } from "@/node/config"; +import * as fs from "fs"; +import http from "node:http"; +import * as os from "os"; +import * as path from "path"; +import { describe, it, expect, beforeEach, afterEach, spyOn } from "bun:test"; + +import { Err, Ok } from "@/common/types/result"; +import { Effect } from "effect"; +import { + getCodexOauthAccounts, + getCodexOauthAccountId, + getCodexOauthAuth, +} from "@/node/utils/codexOauthAuth"; +import { createDeferred } from "@/node/utils/oauthUtils"; +import { + CODEX_OAUTH_TOKEN_URL, + CODEX_OAUTH_DEVICE_USERCODE_URL, + CODEX_OAUTH_DEVICE_TOKEN_POLL_URL, +} from "@/common/constants/codexOAuth"; import type { ProvidersConfig } from "@/node/config"; -import type { ProviderService } from "@/node/services/providerService"; +import { ProviderService } from "@/node/services/providerService"; import type { WindowService } from "@/node/services/windowService"; import type { CodexOauthAuth } from "@/node/utils/codexOauthAuth"; import { CodexOauthService } from "./codexOauthService"; @@ -24,6 +39,7 @@ function fakeJwt(claims: Record): string { function validAuth(overrides?: Partial): CodexOauthAuth { return { type: "oauth", + credentialId: "1c9c50b0-d777-4dd2-998c-09c156ba9754", access: fakeJwt({ sub: "user" }), refresh: "rt_test", expires: Date.now() + 3_600_000, // 1h from now @@ -49,13 +65,17 @@ function mockRefreshResponse(body: Record, status = 200): Respo // --------------------------------------------------------------------------- interface MockDeps { + rootDir: string; providersConfig: ProvidersConfig; setConfigValueCalls: Array<{ provider: string; keyPath: string[]; value: unknown }>; focusCalls: number; + onUpdate?: () => void; + policyDenied?: boolean; } function createMockDeps(): MockDeps { return { + rootDir: fs.mkdtempSync(path.join(os.tmpdir(), "xum-codex-oauth-")), providersConfig: {}, setConfigValueCalls: [], focusCalls: 0, @@ -64,33 +84,55 @@ function createMockDeps(): MockDeps { function createMockProvidersConfigStore( deps: MockDeps -): Pick { +): Pick { return { + rootDir: deps.rootDir, loadProvidersConfig: () => deps.providersConfig, }; } -function createMockProviderService(deps: MockDeps): Pick { +function createMockProviderService( + deps: MockDeps +): Pick { + const setConfigValue: ProviderService["setConfigValue"] = (provider, keyPath, value) => { + deps.setConfigValueCalls.push({ provider, keyPath, value }); + deps.providersConfig[provider] ??= {}; + let current = deps.providersConfig[provider] as Record; + for (const key of keyPath.slice(0, -1)) { + current[key] ??= {}; + current = current[key] as Record; + } + const key = keyPath[keyPath.length - 1]; + if (value === undefined) delete current[key]; + else current[key] = value; + return Promise.resolve(Ok(undefined)); + }; return { - setConfigValue: ( - provider: string, - keyPath: string[], - value: unknown - ): Promise> => { - deps.setConfigValueCalls.push({ provider, keyPath, value }); - // Also update the in-memory config so readStoredAuth() sees the write - if (provider === "openai" && keyPath[0] === "codexOauth") { - if (value === undefined) { - const openai = deps.providersConfig.openai; - if (openai) { - delete openai.codexOauth; - } - } else { - deps.providersConfig.openai ??= {}; - deps.providersConfig.openai.codexOauth = value; - } + setConfigValue, + updateProviderSection: (provider, update, options) => { + if (options?.enforcePolicy && deps.policyDenied) + return Promise.resolve(Err("Provider edits are disabled")); + const next = update(deps.providersConfig[provider]); + deps.onUpdate?.(); + if (!next) return Promise.resolve(Ok({ applied: false })); + deps.providersConfig[provider] = next.value; + deps.setConfigValueCalls.push({ provider, keyPath: [], value: next.value }); + return Promise.resolve(Ok({ applied: true })); + }, + updateConfigValue: async (provider, keyPath, update, options) => { + if (options?.enforcePolicy && deps.policyDenied) return Err("Provider edits are disabled"); + let current: unknown = deps.providersConfig[provider]; + for (const key of keyPath) { + current = + current !== null && typeof current === "object" + ? (current as Record)[key] + : undefined; } - return Promise.resolve(Ok(undefined)); + const next = update(current); + deps.onUpdate?.(); + if (!next) return Ok({ applied: false }); + await setConfigValue(provider, keyPath, next.value); + return Ok({ applied: true }); }, }; } @@ -103,14 +145,24 @@ function createMockWindowService(deps: MockDeps): Pick Promise): void { globalThis.fetch = Object.assign(fn, { @@ -137,6 +189,7 @@ describe("CodexOauthService", () => { afterEach(async () => { globalThis.fetch = originalFetch; await service.dispose(); + fs.rmSync(deps.rootDir, { recursive: true, force: true }); }); // ------------------------------------------------------------------------- @@ -152,6 +205,58 @@ describe("CodexOauthService", () => { } }); + for (const credentialId of [undefined, "1c9c50b0-d777-4dd2-998c-09c156ba9754"]) { + it( + "refreshes a pinned " + (credentialId ? "identified" : "legacy") + " credential", + async () => { + const auth = expiredAuth({ credentialId }); + deps.providersConfig = { openai: { codexOauth: auth } }; + mockFetch(() => + Promise.resolve( + mockRefreshResponse({ + access_token: "rotated-access", + refresh_token: "rotated-refresh", + expires_in: 3600, + }) + ) + ); + const result = await service.getValidAuth("default", { credentialId }); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.access).toBe("rotated-access"); + expect(result.data.refresh).toBe("rotated-refresh"); + expect(result.data.credentialId).toBe(credentialId); + } + } + ); + } + + it("does not let an invalid alias authorize a legacy snapshot", async () => { + const auth = validAuth(); + for (const legacyCredentialId of [ + auth.credentialId, + undefined, + null, + "invalid", + 42, + "50e00a32-b964-4ce2-b131-6b53356ce2db", + ]) { + deps.providersConfig = { openai: { codexOauth: { ...auth, legacyCredentialId } } }; + const pinned = await service.getValidAuth("default", { credentialId: undefined }); + expect(pinned.success).toBe(legacyCredentialId === auth.credentialId); + expect( + (await service.getValidAuth("default", { credentialId: auth.credentialId })).success + ).toBe(true); + expect( + ( + await service.getValidAuth("default", { + credentialId: "81df6409-6a24-4a19-950e-7daf7ef47c4e", + }) + ).success + ).toBe(false); + } + }); + it("returns stored auth when token is not expired", async () => { const auth = validAuth(); deps.providersConfig = { openai: { codexOauth: auth } }; @@ -241,73 +346,71 @@ describe("CodexOauthService", () => { // Invalid grant cleanup // ------------------------------------------------------------------------- - describe("invalid grant cleanup", () => { - it("calls disconnect + clears stored auth on invalid_grant response", async () => { - const expired = expiredAuth(); - deps.providersConfig = { openai: { codexOauth: expired } }; - - mockFetch(() => - Promise.resolve( - new Response(JSON.stringify({ error: "invalid_grant" }), { - status: 400, - headers: { "Content-Type": "application/json" }, - }) - ) - ); - - const result = await service.getValidAuth(); - expect(result.success).toBe(false); - - // Should have called setConfigValue to clear auth (disconnect) - const clearCall = deps.setConfigValueCalls.find( - (c) => c.provider === "openai" && c.keyPath[0] === "codexOauth" && c.value === undefined - ); - expect(clearCall).toBeDefined(); - }); - - it("clears auth when error text contains 'revoked'", async () => { - const expired = expiredAuth(); - deps.providersConfig = { openai: { codexOauth: expired } }; - - mockFetch(() => - Promise.resolve( - new Response("Token has been revoked", { - status: 401, - }) - ) - ); - - const result = await service.getValidAuth(); - expect(result.success).toBe(false); - - const clearCall = deps.setConfigValueCalls.find( - (c) => c.provider === "openai" && c.keyPath[0] === "codexOauth" && c.value === undefined - ); - expect(clearCall).toBeDefined(); - }); - - it("subsequent getValidAuth returns error after invalid_grant cleanup", async () => { - const expired = expiredAuth(); - deps.providersConfig = { openai: { codexOauth: expired } }; - - mockFetch(() => - Promise.resolve( - new Response(JSON.stringify({ error: "invalid_grant" }), { - status: 400, - headers: { "Content-Type": "application/json" }, - }) - ) - ); - - // First call triggers disconnect - await service.getValidAuth(); - - // Second call should see no stored auth - const result = await service.getValidAuth(); - expect(result.success).toBe(false); - if (!result.success) { - expect(result.error).toContain("not configured"); + describe("invalid grant marking", () => { + it.each(["invalid_grant", "Token has been revoked"])( + "retains credential identity after %s", + async (error) => { + const expired = expiredAuth(); + deps.providersConfig = { openai: { codexOauth: expired } }; + mockFetch(() => Promise.resolve(mockRefreshResponse({ error }, 400))); + expect((await service.getValidAuth()).success).toBe(false); + expect(getCodexOauthAuth(deps.providersConfig.openai)).toEqual({ + ...expired, + invalidReason: "invalid_grant", + }); + expect(getCodexOauthAccounts(deps.providersConfig.openai)).toHaveLength(1); + } + ); + + it.each([false, true])( + "never returns or refreshes marked credentials (expired=%s)", + async (expired) => { + const auth = expired + ? expiredAuth({ invalidReason: "invalid_grant" }) + : validAuth({ invalidReason: "invalid_grant" }); + const legacy = validAuth({ access: "other" }); + deps.providersConfig = { + openai: { + codexOauth: legacy, + codexOauthAccounts: { work: { label: "Work", credentials: auth } }, + codexOauthDefaultAccountId: "work", + }, + }; + let fetchCount = 0; + mockFetch(() => { + fetchCount++; + return Promise.reject(new Error("Must not refresh")); + }); + expect((await service.getValidAuth()).success).toBe(false); + expect((await service.getValidAuth("work")).success).toBe(false); + expect(await service.getValidAuth("default")).toEqual(Ok(legacy)); + expect(fetchCount).toBe(0); + expect(getCodexOauthAccounts(deps.providersConfig.openai)).toHaveLength(2); + expect(await service.disconnect("work")).toEqual(Ok(undefined)); + expect(getCodexOauthAuth(deps.providersConfig.openai, "work")).toBeNull(); } + ); + + it("rejects a marker after the local refresh mutex without refreshing again", async () => { + deps.providersConfig = { openai: { codexOauth: expiredAuth() } }; + const started = createDeferred(); + const response = createDeferred(); + let fetchCount = 0; + mockFetch(() => { + fetchCount++; + started.resolve(undefined); + return response.promise; + }); + const first = service.getValidAuth(); + await started.promise; + const second = service.getValidAuth(); + response.resolve(mockRefreshResponse({ error: "invalid_grant" }, 400)); + expect((await first).success).toBe(false); + expect((await second).success).toBe(false); + expect((await service.getValidAuth()).success).toBe(false); + expect(fetchCount).toBe(1); + expect(await service.disconnect()).toEqual(Ok(undefined)); + expect(getCodexOauthAuth(deps.providersConfig.openai)).toBeNull(); }); }); @@ -451,4 +554,1541 @@ describe("CodexOauthService", () => { } }); }); + describe("account selection", () => { + it("uses the global selection and an explicit override without falling back", async () => { + const legacy = validAuth({ access: "legacy" }); + const work = validAuth({ access: "work", accountId: "chatgpt-work" }); + deps.providersConfig = { + openai: { + codexOauth: legacy, + codexOauthAccounts: { work: { label: "Work", credentials: work } }, + codexOauthDefaultAccountId: "work", + }, + }; + expect(await service.getValidAuth()).toEqual(Ok(work)); + expect(await service.getValidAuth("default")).toEqual(Ok(legacy)); + expect((await service.getValidAuth("missing")).success).toBe(false); + expect(await service.disconnect("work")).toEqual(Ok(undefined)); + expect((await service.getValidAuth()).success).toBe(false); + expect(await service.getValidAuth("default")).toEqual(Ok(legacy)); + }); + + it("renames each slot and selects defaults without changing credentials", async () => { + const legacy = validAuth(); + const work = validAuth({ refresh: "work" }); + deps.providersConfig = { + openai: { + codexOauth: legacy, + codexOauthAccounts: { work: { label: "Work", credentials: work } }, + }, + }; + expect(await Effect.runPromise(service.renameAccountEffect("default", " Personal "))).toEqual( + Ok(undefined) + ); + expect(await service.renameAccount("work", " Team ")).toEqual(Ok(undefined)); + expect(await Effect.runPromise(service.setDefaultAccountEffect("work"))).toEqual( + Ok(undefined) + ); + expect(getCodexOauthAccounts(deps.providersConfig.openai)).toEqual([ + { id: "default", label: "Personal", auth: legacy }, + { id: "work", label: "Team", auth: work }, + ]); + expect(await service.getValidAuth()).toEqual(Ok(work)); + expect((await service.setDefaultAccount("missing")).success).toBe(false); + expect((await service.renameAccount("missing", "Name")).success).toBe(false); + }); + + it("enforces policy for account edits and token refresh", async () => { + const stored = expiredAuth(); + deps.providersConfig = { openai: { codexOauth: stored } }; + deps.policyDenied = true; + expect((await service.setDefaultAccount("default")).success).toBe(false); + expect((await service.renameAccount("default", "Renamed")).success).toBe(false); + expect((await service.disconnect()).success).toBe(false); + mockFetch(() => + Promise.resolve(mockRefreshResponse({ access_token: "refreshed", expires_in: 3600 })) + ); + expect((await service.getValidAuth()).success).toBe(false); + expect(getCodexOauthAuth(deps.providersConfig.openai)).toEqual(stored); + expect(deps.setConfigValueCalls).toHaveLength(0); + }); + + it("enforces policy when a revoked refresh attempts to mark credentials", async () => { + const stored = expiredAuth(); + deps.providersConfig = { openai: { codexOauth: stored } }; + deps.policyDenied = true; + mockFetch(() => Promise.resolve(mockRefreshResponse({ error: "invalid_grant" }, 400))); + expect((await service.getValidAuth()).success).toBe(false); + expect(getCodexOauthAuth(deps.providersConfig.openai)).toEqual(stored); + expect(deps.setConfigValueCalls).toHaveLength(0); + }); + + it("rejects unsafe slot IDs and invalid labels before storage writes", async () => { + for (const id of ["", "__proto__", "constructor", "prototype", "../bad", "x".repeat(201)]) { + expect((await service.disconnect(id)).success).toBe(false); + expect((await service.setDefaultAccount(id)).success).toBe(false); + expect((await service.renameAccount(id, "Name")).success).toBe(false); + expect((await service.startDeviceFlow({ accountId: id })).success).toBe(false); + expect((await service.startDesktopFlow({ accountId: id })).success).toBe(false); + } + for (const label of [" ", "x".repeat(101)]) { + expect((await service.renameAccount("default", label)).success).toBe(false); + expect((await service.startDeviceFlow({ label })).success).toBe(false); + } + expect((await service.startDeviceFlow({ accountId: "default", label: "Name" })).success).toBe( + false + ); + expect(deps.setConfigValueCalls).toHaveLength(0); + }); + }); + + describe("account refresh isolation", () => { + it.each([ + { change: "rotation", accepted: true }, + { change: "replacement", accepted: false }, + { change: "expired replacement", accepted: false }, + { change: "dropped ID", accepted: false }, + { change: "legacy rotation", accepted: true }, + { change: "legacy stamp", accepted: false }, + { change: "marked", accepted: false }, + { change: "removed", accepted: false }, + ])("checks the pinned credential after lease wait: $change", async ({ change, accepted }) => { + const initial = expiredAuth(change.startsWith("legacy") ? { credentialId: undefined } : {}); + const latest = validAuth({ + ...initial, + access: "updated", + expires: change === "expired replacement" ? Date.now() - 1000 : Date.now() + 3600000, + credentialId: + change.includes("replacement") || change === "legacy stamp" + ? "50e00a32-b964-4ce2-b131-6b53356ce2db" + : change === "dropped ID" + ? undefined + : initial.credentialId, + invalidReason: change === "marked" ? "invalid_grant" : undefined, + }); + const provider = new ProviderService(new Config(deps.rootDir)); + const store = provider.providersConfigStore; + store.saveProvidersConfig({ openai: { codexOauth: initial } }); + const entered = createDeferred(); + const release = createDeferred(); + const attempted = createDeferred(); + const holder = new FileLeaseManager(deps.rootDir).withCodexOauthRefreshLock( + "default", + async () => { + entered.resolve(undefined); + await release.promise; + } + ); + await entered.promise; + class ObservedLeaseManager extends FileLeaseManager { + override withCodexOauthRefreshLock( + accountId: string, + fn: () => Promise | T + ): Promise { + attempted.resolve(undefined); + return super.withCodexOauthRefreshLock(accountId, fn); + } + } + const waiting = new CodexOauthService( + store, + provider, + undefined, + new ObservedLeaseManager(deps.rootDir) + ); + let fetchCount = 0; + mockFetch(() => { + fetchCount++; + return Promise.reject(new Error("Unexpected refresh")); + }); + try { + const pending = waiting.getValidAuth(); + await attempted.promise; + store.saveProvidersConfig({ openai: change === "removed" ? {} : { codexOauth: latest } }); + release.resolve(undefined); + await holder; + const result = await pending; + expect(result.success).toBe(accepted); + if (accepted) expect(result).toEqual(Ok(latest)); + expect(fetchCount).toBe(0); + } finally { + release.resolve(undefined); + await holder; + await waiting.dispose(); + } + }); + + it("does not adopt a replaced credential after waiting for the local mutex", async () => { + deps.providersConfig = { openai: { codexOauth: expiredAuth() } }; + const started = createDeferred(); + const response = createDeferred(); + let fetchCount = 0; + mockFetch(() => { + fetchCount++; + started.resolve(undefined); + return response.promise; + }); + const first = service.getValidAuth(); + await started.promise; + const waiting = service.getValidAuth(); + deps.providersConfig.openai.codexOauth = validAuth({ + credentialId: "50e00a32-b964-4ce2-b131-6b53356ce2db", + }); + response.resolve(mockRefreshResponse({ access_token: "old-rotation", expires_in: 3600 })); + expect((await first).success).toBe(false); + expect((await waiting).success).toBe(false); + expect(fetchCount).toBe(1); + }); + + it("shares a refresh lease across service instances and persists the rotated token", async () => { + const store = new ProvidersConfigStore(deps.rootDir); + store.saveProvidersConfig({ + openai: { + codexOauthAccounts: { + work: { label: "Work", credentials: expiredAuth({ refresh: "old" }) }, + }, + }, + }); + const secondAttempt = createDeferred(); + class ObservedLeaseManager extends FileLeaseManager { + override withCodexOauthRefreshLock( + accountId: string, + fn: () => Promise | T + ): Promise { + secondAttempt.resolve(undefined); + return super.withCodexOauthRefreshLock(accountId, fn); + } + } + const firstService = new CodexOauthService( + store, + new ProviderService(new Config(deps.rootDir)) + ); + const secondService = new CodexOauthService( + new ProvidersConfigStore(deps.rootDir), + new ProviderService(new Config(deps.rootDir)), + undefined, + new ObservedLeaseManager(deps.rootDir) + ); + const refreshStarted = createDeferred(); + const refreshResponse = createDeferred(); + let refreshCount = 0; + mockFetch(() => { + refreshCount++; + refreshStarted.resolve(undefined); + return refreshResponse.promise; + }); + try { + const first = firstService.getValidAuth("work"); + await refreshStarted.promise; + const second = secondService.getValidAuth("work"); + await secondAttempt.promise; + refreshResponse.resolve( + mockRefreshResponse({ + access_token: "refreshed", + refresh_token: "rotated", + expires_in: 3600, + }) + ); + const results = await Promise.all([first, second]); + expect(results[0].success).toBe(true); + expect(results[1]).toEqual(results[0]); + expect(refreshCount).toBe(1); + expect( + getCodexOauthAuth( + new ProvidersConfigStore(deps.rootDir).loadProvidersConfig()?.openai, + "work" + )?.refresh + ).toBe("rotated"); + } finally { + await firstService.dispose(); + await secondService.dispose(); + } + }); + + it("refreshes two accounts concurrently and keeps both writes", async () => { + deps.providersConfig = { + openai: { + codexOauth: expiredAuth({ refresh: "legacy" }), + codexOauthAccounts: { + work: { label: "Work", credentials: expiredAuth({ refresh: "work" }) }, + }, + }, + }; + const bothStarted = createDeferred(); + const release = createDeferred(); + let count = 0; + mockFetch(async (_input, init) => { + if (++count === 2) bothStarted.resolve(undefined); + await release.promise; + const token = new URLSearchParams(requestBody(init)).get("refresh_token"); + return mockRefreshResponse({ access_token: token + "-new", expires_in: 3600 }); + }); + const legacy = service.getValidAuth(); + const work = service.getValidAuth("work"); + await bothStarted.promise; + release.resolve(undefined); + expect((await legacy).success).toBe(true); + expect((await work).success).toBe(true); + expect(getCodexOauthAuth(deps.providersConfig.openai, "default")?.access).toBe("legacy-new"); + expect(getCodexOauthAuth(deps.providersConfig.openai, "work")?.access).toBe("work-new"); + }); + + it("pins the default while a refresh runs and preserves the ChatGPT identity", async () => { + deps.providersConfig = { + openai: { + codexOauth: expiredAuth({ accountId: "chatgpt-original" }), + codexOauthAccounts: { + work: { label: "Work", credentials: validAuth({ access: "work" }) }, + }, + }, + }; + const started = createDeferred(); + const response = createDeferred(); + mockFetch(() => { + started.resolve(undefined); + return response.promise; + }); + const pending = service.getValidAuth(); + await started.promise; + expect(await service.setDefaultAccount("work")).toEqual(Ok(undefined)); + response.resolve( + mockRefreshResponse({ + access_token: fakeJwt({ chatgpt_account_id: "chatgpt-other" }), + expires_in: 3600, + }) + ); + const result = await pending; + expect(result.success).toBe(true); + if (result.success) expect(result.data.accountId).toBe("chatgpt-original"); + expect(getCodexOauthAuth(deps.providersConfig.openai, "work")?.access).toBe("work"); + }); + + it.each([200, 400])( + "does not overwrite a reconnect after an old refresh returns %s", + async (status) => { + deps.providersConfig = { openai: { codexOauth: expiredAuth() } }; + const started = createDeferred(); + const response = createDeferred(); + mockFetch(() => { + started.resolve(undefined); + return response.promise; + }); + const pending = service.getValidAuth(); + await started.promise; + // A separate process changes the same slot without this service's revision map. + const reconnected = validAuth({ access: "reconnected", refresh: "new-session" }); + deps.providersConfig.openai.codexOauth = reconnected; + response.resolve( + mockRefreshResponse( + status === 200 + ? { access_token: "stale", expires_in: 3600 } + : { error: "invalid_grant" }, + status + ) + ); + expect((await pending).success).toBe(false); + expect(await service.getValidAuth()).toEqual(Ok(reconnected)); + } + ); + + it.each(["refresh", "disconnect"] as const)( + "serializes disconnect and refresh when %s persists first", + async (first) => { + deps.providersConfig = { + openai: { codexOauthAccounts: { work: { label: "Work", credentials: expiredAuth() } } }, + }; + const provider = createMockProviderService(deps); + service = createService(deps, provider); + const entered = createDeferred(); + const release = createDeferred(); + const fetchStarted = createDeferred(); + const mutationOrder: string[] = []; + const originalUpdate = provider.updateConfigValue; + let calls = 0; + const update = spyOn(provider, "updateConfigValue").mockImplementation(async (...args) => { + const isFirst = ++calls === 1; + if (isFirst) { + entered.resolve(undefined); + await release.promise; + } + mutationOrder.push(isFirst ? first : first === "refresh" ? "disconnect" : "refresh"); + return originalUpdate(...args); + }); + mockFetch(() => { + fetchStarted.resolve(undefined); + return Promise.resolve( + mockRefreshResponse({ access_token: "rotated", expires_in: 3600 }) + ); + }); + try { + const initial = + first === "refresh" ? service.getValidAuth("work") : service.disconnect("work"); + await entered.promise; + const waiting = + first === "refresh" ? service.disconnect("work") : service.getValidAuth("work"); + await fetchStarted.promise; + release.resolve(undefined); + const [initialResult, waitingResult] = await Promise.all([initial, waiting]); + const disconnected = first === "disconnect" ? initialResult : waitingResult; + expect(disconnected).toEqual(Ok(undefined)); + expect(mutationOrder).toEqual( + first === "refresh" ? ["refresh", "disconnect"] : ["disconnect", "refresh"] + ); + if (first === "disconnect") expect(waitingResult.success).toBe(false); + expect(getCodexOauthAuth(deps.providersConfig.openai, "work")).toBeNull(); + } finally { + release.resolve(undefined); + update.mockRestore(); + } + } + ); + + it("does not restore a disconnected slot after refresh", async () => { + const work = expiredAuth(); + const legacy = validAuth({ access: "legacy" }); + deps.providersConfig = { + openai: { + codexOauth: legacy, + codexOauthAccounts: { work: { label: "Work", credentials: work } }, + }, + }; + const started = createDeferred(); + const response = createDeferred(); + mockFetch(() => { + started.resolve(undefined); + return response.promise; + }); + const pending = service.getValidAuth("work"); + await started.promise; + await service.disconnect("work"); + response.resolve(mockRefreshResponse({ access_token: "stale", expires_in: 3600 })); + expect((await pending).success).toBe(false); + expect(getCodexOauthAuth(deps.providersConfig.openai, "work")).toBeNull(); + expect(await service.getValidAuth("default")).toEqual(Ok(legacy)); + }); + + it("keeps a concurrent rename when refreshed credentials persist", async () => { + deps.providersConfig = { + openai: { codexOauthAccounts: { work: { label: "Work", credentials: expiredAuth() } } }, + }; + const started = createDeferred(); + const response = createDeferred(); + mockFetch(() => { + started.resolve(undefined); + return response.promise; + }); + const pending = service.getValidAuth("work"); + await started.promise; + await service.renameAccount("work", "Team"); + response.resolve(mockRefreshResponse({ access_token: "new", expires_in: 3600 })); + expect((await pending).success).toBe(true); + expect(getCodexOauthAccounts(deps.providersConfig.openai)[0].label).toBe("Team"); + }); + }); + + describe("login destinations", () => { + function activeLoginSelectionCount(): number { + // Inspect retained credentials without exposing internal state through the service API. + // eslint-disable-next-line @typescript-eslint/dot-notation -- Bracket access permits private state checks in tests. + return service["loginSelections"].size; + } + + function deviceFetch(exchange?: (init?: RequestInit) => Promise): void { + let nextCode = 0; + mockFetch(async (input, init) => { + const url = input instanceof Request ? input.url : input.toString(); + if (url === CODEX_OAUTH_DEVICE_USERCODE_URL) { + return mockRefreshResponse({ + device_auth_id: "device-" + ++nextCode, + user_code: "code", + interval: 1, + expires_in: 60, + }); + } + if (url === CODEX_OAUTH_DEVICE_TOKEN_POLL_URL) { + const body = JSON.parse(requestBody(init)) as { device_auth_id: string }; + return mockRefreshResponse({ + authorization_code: body.device_auth_id, + code_verifier: "verifier", + }); + } + if (url === CODEX_OAUTH_TOKEN_URL) { + if (exchange) return exchange(init); + return mockRefreshResponse({ + access_token: "access-" + new URLSearchParams(requestBody(init)).get("code"), + refresh_token: "refresh", + expires_in: 3600, + }); + } + throw new Error("Unexpected fetch URL"); + }); + } + + for (const kind of ["desktop", "device"] as const) { + for (const terminal of ["cancelled", "expired", "failed"] as const) { + it("releases " + terminal + " named " + kind + " login selections", async () => { + deviceFetch(() => Promise.resolve(mockRefreshResponse({ error: "access_denied" }, 400))); + const flow = + kind === "desktop" + ? await service.startDesktopFlow({ label: "Work" }) + : await service.startDeviceFlow({ label: "Work" }); + if (!flow.success) throw new Error(flow.error); + // Terminal flows must release selections because they retain credentials. + expect(activeLoginSelectionCount()).toBe(1); + const flowId = flow.data.flowId; + if (kind === "desktop") { + // eslint-disable-next-line @typescript-eslint/dot-notation -- Await the private cleanup signal without timing assumptions. + const completion = service["desktopFlows"].get(flowId)!.resultDeferred.promise; + if (terminal === "cancelled") await service.cancelDesktopFlow(flowId); + if (terminal === "failed") { + const response = await originalFetch( + "http://localhost:1455/auth/callback?error=access_denied&state=" + flowId + ); + await response.text(); + } + const result = await service.waitForDesktopFlow(flowId, { timeoutMs: 0 }); + expect(result.success).toBe(false); + // The flow manager releases resources in a detached fiber after a waiter timeout. + await completion; + } else { + if (terminal === "cancelled") await service.cancelDeviceFlow(flowId); + const clock = + terminal === "expired" + ? spyOn(Date, "now").mockReturnValue(Date.now() + 120_000) + : undefined; + try { + const result = await service.waitForDeviceFlow(flowId); + expect(result.success).toBe(false); + if (terminal === "expired" && !result.success) + expect(result.error).toContain("expired"); + } finally { + clock?.mockRestore(); + } + } + expect(activeLoginSelectionCount()).toBe(0); + expect(getCodexOauthAccounts(deps.providersConfig.openai)).toEqual([]); + }); + } + + it("keeps the newer selection when a superseded " + kind + " flow terminates", async () => { + deps.providersConfig = { + openai: { codexOauthAccounts: { work: { label: "Work", credentials: validAuth() } } }, + }; + deviceFetch(); + const first = + kind === "desktop" + ? await service.startDesktopFlow({ accountId: "work" }) + : await service.startDeviceFlow({ accountId: "work" }); + if (!first.success) throw new Error(first.error); + const second = await service.startDeviceFlow({ accountId: "work" }); + if (!second.success) throw new Error(second.error); + if (kind === "desktop") await service.cancelDesktopFlow(first.data.flowId); + else await service.cancelDeviceFlow(first.data.flowId); + expect(activeLoginSelectionCount()).toBe(1); + expect(await service.waitForDeviceFlow(second.data.flowId)).toEqual(Ok(undefined)); + expect(activeLoginSelectionCount()).toBe(0); + const auth = await service.getValidAuth("work"); + expect(auth.success).toBe(true); + if (auth.success) + expect(auth.data.access).toBe(kind === "desktop" ? "access-device-1" : "access-device-2"); + }); + } + + it("does not save login credentials when policy denies provider edits", async () => { + deps.policyDenied = true; + deviceFetch(); + const flow = await service.startDeviceFlow({ label: "Work" }); + if (!flow.success) throw new Error(flow.error); + expect((await service.waitForDeviceFlow(flow.data.flowId)).success).toBe(false); + expect(getCodexOauthAccounts(deps.providersConfig.openai)).toEqual([]); + expect(deps.setConfigValueCalls).toHaveLength(0); + }); + + for (const action of ["rename", "refresh", "stamp", "reconnect"] as const) { + it.each([undefined, 42, " ", " Work "])( + "drops unsafe disk fields during " + action + " with label %s", + async (label) => { + const credentials = validAuth({ + credentialId: action === "stamp" ? undefined : validAuth().credentialId, + expires: action === "refresh" ? 0 : Date.now() + 3_600_000, + }); + const store = new ProvidersConfigStore(deps.rootDir); + const document = { + openai: { + codexOauthAccounts: { + work: { label, credentials, auth: { access: "unsafe-copy" } }, + }, + }, + }; + // Manual disk damage bypasses write validation. Loading must still permit account recovery. + await fs.promises.writeFile(store.providersFile, JSON.stringify(document)); + service = new CodexOauthService(store, new ProviderService(new Config(deps.rootDir))); + const expectedLabel = typeof label === "string" && label.trim() ? label.trim() : "work"; + expect(getCodexOauthAccounts(store.loadProvidersConfig()?.openai)).toEqual([ + { id: "work", label: expectedLabel, auth: credentials }, + ]); + deviceFetch(() => + Promise.resolve( + mockRefreshResponse({ + access_token: "updated", + refresh_token: "updated-refresh", + expires_in: 3600, + }) + ) + ); + if (action === "rename") { + expect(await service.renameAccount("work", "Renamed")).toEqual(Ok(undefined)); + } else if (action === "refresh") { + expect(await service.getValidAuth("work")).toMatchObject({ + success: true, + data: { access: "updated" }, + }); + } else { + const flow = await service.startDeviceFlow({ accountId: "work" }); + if (!flow.success) throw new Error(flow.error); + if (action === "stamp") await service.cancelDeviceFlow(flow.data.flowId); + else expect(await service.waitForDeviceFlow(flow.data.flowId)).toEqual(Ok(undefined)); + } + const saved = store.loadProvidersConfig()?.openai; + expect(saved?.codexOauthAccounts).toEqual({ + work: { + label: action === "rename" ? "Renamed" : expectedLabel, + credentials: getCodexOauthAuth(saved, "work"), + }, + }); + expect(await fs.promises.readFile(store.providersFile, "utf8")).not.toContain( + "unsafe-copy" + ); + } + ); + } + + it("creates a named slot and selects the first account globally", async () => { + deviceFetch(); + const flow = await service.startDeviceFlow({ label: " Personal " }); + if (!flow.success) throw new Error(flow.error); + expect(await service.waitForDeviceFlow(flow.data.flowId)).toEqual(Ok(undefined)); + const accounts = getCodexOauthAccounts(deps.providersConfig.openai); + expect(accounts).toHaveLength(1); + expect(accounts[0].id).not.toBe("default"); + expect(accounts[0].label).toBe("Personal"); + // Persist one protected credential object, without a second token copy under auth. + expect(deps.providersConfig.openai?.codexOauthAccounts).toEqual({ + [accounts[0].id]: { label: "Personal", credentials: accounts[0].auth }, + }); + expect(await service.getValidAuth()).toEqual(Ok(accounts[0].auth)); + expect(deps.providersConfig.openai?.codexOauth).toBeUndefined(); + }); + + it("keeps concurrent named logins in separate slots", async () => { + deps.providersConfig = { openai: { codexOauth: validAuth({ access: "legacy" }) } }; + deviceFetch(); + const first = await service.startDeviceFlow({ label: "One" }); + const second = await service.startDeviceFlow({ label: "Two" }); + if (!first.success || !second.success) throw new Error("Login start failed"); + expect(await service.setDefaultAccount("default")).toEqual(Ok(undefined)); + const results = await Promise.all([ + service.waitForDeviceFlow(first.data.flowId), + service.waitForDeviceFlow(second.data.flowId), + ]); + expect(results.every((result) => result.success)).toBe(true); + const accounts = getCodexOauthAccounts(deps.providersConfig.openai); + expect(accounts).toHaveLength(3); + expect(new Set(accounts.map((account) => account.id)).size).toBe(3); + expect(accounts.map((account) => account.auth.access).sort()).toEqual([ + "access-device-1", + "access-device-2", + "legacy", + ]); + expect(getCodexOauthAuth(deps.providersConfig.openai)?.access).toBe("legacy"); + }); + + it("selects one account when the first named logins finish concurrently", async () => { + deviceFetch(); + const first = await service.startDeviceFlow({ label: "One" }); + const second = await service.startDeviceFlow({ label: "Two" }); + if (!first.success || !second.success) throw new Error("Login start failed"); + const results = await Promise.all([ + service.waitForDeviceFlow(first.data.flowId), + service.waitForDeviceFlow(second.data.flowId), + ]); + expect(results.every((result) => result.success)).toBe(true); + expect(getCodexOauthAccounts(deps.providersConfig.openai)).toHaveLength(2); + expect((await service.getValidAuth()).success).toBe(true); + }); + + it("pins a named desktop login while the global default changes", async () => { + const legacy = validAuth({ access: "legacy" }); + deps.providersConfig = { openai: { codexOauth: legacy } }; + mockFetch(() => + Promise.resolve( + mockRefreshResponse({ + access_token: "desktop", + refresh_token: "desktop-refresh", + expires_in: 3600, + }) + ) + ); + const flow = await service.startDesktopFlow({ label: "Desktop" }); + if (!flow.success) throw new Error(flow.error); + await service.setDefaultAccount("default"); + const callback = await originalFetch( + "http://localhost:1455/auth/callback?code=auth-code&state=" + flow.data.flowId + ); + expect(callback.ok).toBe(true); + await callback.text(); + expect(await service.waitForDesktopFlow(flow.data.flowId)).toEqual(Ok(undefined)); + const accounts = getCodexOauthAccounts(deps.providersConfig.openai); + expect(accounts).toHaveLength(2); + expect(accounts[1].auth.access).toBe("desktop"); + expect(accounts[1].label).toBe("Desktop"); + expect(await service.getValidAuth()).toEqual(Ok(legacy)); + }); + + it("keeps no-argument login in the legacy slot despite a named default", async () => { + const work = validAuth({ access: "work" }); + deps.providersConfig = { + openai: { + codexOauthAccounts: { work: { label: "Work", credentials: work } }, + codexOauthDefaultAccountId: "work", + }, + }; + deviceFetch(); + const flow = await service.startDeviceFlow(); + if (!flow.success) throw new Error(flow.error); + expect(await service.waitForDeviceFlow(flow.data.flowId)).toEqual(Ok(undefined)); + expect(getCodexOauthAuth(deps.providersConfig.openai, "default")?.access).toBe( + "access-device-1" + ); + expect(await service.getValidAuth()).toEqual(Ok(work)); + }); + + it.each(["policy", "I/O"] as const)( + "keeps pending refresh and reconnect valid after a failed %s disconnect", + async (failure) => { + const initial = expiredAuth(); + deps.providersConfig = { + openai: { codexOauthAccounts: { work: { label: "Work", credentials: initial } } }, + }; + const provider = createMockProviderService(deps); + service = createService(deps, provider); + const refreshStarted = createDeferred(); + const loginStarted = createDeferred(); + const refreshResponse = createDeferred(); + const loginResponse = createDeferred(); + deviceFetch((init) => { + if (new URLSearchParams(requestBody(init)).get("grant_type") === "refresh_token") { + refreshStarted.resolve(undefined); + return refreshResponse.promise; + } + loginStarted.resolve(undefined); + return loginResponse.promise; + }); + const flow = await service.startDeviceFlow({ accountId: "work" }); + if (!flow.success) throw new Error(flow.error); + const login = service.waitForDeviceFlow(flow.data.flowId); + await loginStarted.promise; + const refresh = service.getValidAuth("work"); + await refreshStarted.promise; + const update = spyOn(provider, "updateConfigValue"); + if (failure === "policy") deps.policyDenied = true; + else update.mockRejectedValueOnce(new Error("Disk write failed")); + const disconnected = await service.disconnect("work"); + const afterFailure = getCodexOauthAuth(deps.providersConfig.openai, "work"); + deps.policyDenied = false; + update.mockRestore(); + refreshResponse.resolve( + mockRefreshResponse({ + access_token: "rotated", + refresh_token: "rotated-refresh", + expires_in: 3600, + }) + ); + const refreshed = await refresh; + loginResponse.resolve( + mockRefreshResponse({ + access_token: "reconnected", + refresh_token: "reconnected-refresh", + expires_in: 3600, + }) + ); + const reconnected = await login; + expect(disconnected).toEqual( + Err(failure === "policy" ? "Provider edits are disabled" : "Disk write failed") + ); + expect(afterFailure).toEqual(initial); + expect(refreshed).toMatchObject({ success: true, data: { access: "rotated" } }); + expect(reconnected).toEqual(Ok(undefined)); + expect(getCodexOauthAuth(deps.providersConfig.openai, "work")?.access).toBe("reconnected"); + } + ); + + it.each(["disconnect", "cancel"])("does not persist an exchange after %s", async (action) => { + deps.providersConfig = { + openai: { codexOauthAccounts: { work: { label: "Work", credentials: validAuth() } } }, + }; + const started = createDeferred(); + const response = createDeferred(); + deviceFetch(() => { + started.resolve(undefined); + return response.promise; + }); + const flow = await service.startDeviceFlow({ accountId: "work" }); + if (!flow.success) throw new Error(flow.error); + const pending = service.waitForDeviceFlow(flow.data.flowId); + await started.promise; + if (action === "disconnect") await service.disconnect("work"); + else await service.cancelDeviceFlow(flow.data.flowId); + const writeCount = deps.setConfigValueCalls.length; + const compared = createDeferred(); + deps.onUpdate = () => compared.resolve(undefined); + response.resolve( + mockRefreshResponse({ access_token: "stale", refresh_token: "stale", expires_in: 3600 }) + ); + expect((await pending).success).toBe(false); + await compared.promise; + deps.onUpdate = undefined; + expect(getCodexOauthAuth(deps.providersConfig.openai, "work")?.access).not.toBe("stale"); + if (action === "disconnect") expect(deps.setConfigValueCalls).toHaveLength(writeCount); + }); + + it("keeps refreshed credentials when a concurrent reconnect is cancelled", async () => { + deps.providersConfig = { + openai: { + codexOauthAccounts: { + work: { label: "Work", credentials: expiredAuth({ refresh: "old" }) }, + }, + }, + }; + const refreshStarted = createDeferred(); + const refreshResponse = createDeferred(); + deviceFetch(() => { + refreshStarted.resolve(undefined); + return refreshResponse.promise; + }); + const refresh = service.getValidAuth("work"); + await refreshStarted.promise; + const flow = await service.startDeviceFlow({ accountId: "work" }); + if (!flow.success) throw new Error(flow.error); + await service.cancelDeviceFlow(flow.data.flowId); + refreshResponse.resolve( + mockRefreshResponse({ + access_token: "refreshed", + refresh_token: "rotated", + expires_in: 3600, + }) + ); + const result = await refresh; + expect(result.success).toBe(true); + expect(getCodexOauthAuth(deps.providersConfig.openai, "work")?.refresh).toBe("rotated"); + expect(await service.getValidAuth("work")).toEqual(result); + }); + + it("keeps refreshed credentials when reconnect startup fails", async () => { + deps.providersConfig = { openai: { codexOauth: expiredAuth({ refresh: "old" }) } }; + const refreshStarted = createDeferred(); + const refreshResponse = createDeferred(); + mockFetch((input) => { + if (input === CODEX_OAUTH_DEVICE_USERCODE_URL) + return Promise.reject(new Error("Request failed")); + refreshStarted.resolve(undefined); + return refreshResponse.promise; + }); + const refresh = service.getValidAuth(); + await refreshStarted.promise; + expect((await service.startDeviceFlow({ accountId: "default" })).success).toBe(false); + refreshResponse.resolve( + mockRefreshResponse({ + access_token: "refreshed", + refresh_token: "rotated", + expires_in: 3600, + }) + ); + expect((await refresh).success).toBe(true); + expect(getCodexOauthAuth(deps.providersConfig.openai)?.refresh).toBe("rotated"); + }); + + it("lets only the latest reconnect attempt replace the same slot", async () => { + deps.providersConfig = { openai: { codexOauth: validAuth() } }; + deviceFetch(); + const first = await service.startDeviceFlow({ accountId: "default" }); + const second = await service.startDeviceFlow({ accountId: "default" }); + if (!first.success || !second.success) throw new Error("Login start failed"); + expect((await service.waitForDeviceFlow(first.data.flowId)).success).toBe(false); + expect(await service.waitForDeviceFlow(second.data.flowId)).toEqual(Ok(undefined)); + expect(getCodexOauthAuth(deps.providersConfig.openai)?.access).toBe("access-device-2"); + }); + + function sharedServices(auth: CodexOauthAuth, accountId = "work") { + const provider = new ProviderService(new Config(deps.rootDir)); + const store = provider.providersConfigStore; + store.saveProvidersConfig({ + openai: + accountId === "default" + ? { codexOauth: auth } + : { codexOauthAccounts: { [accountId]: { label: "Work", credentials: auth } } }, + }); + return { + provider, + store, + first: new CodexOauthService(store, provider), + second: new CodexOauthService( + new ProvidersConfigStore(deps.rootDir), + new ProviderService(new Config(deps.rootDir)) + ), + }; + } + + it.each(["default", "work"])( + "reconnects a %s credential with a malformed durable ID", + async (accountId) => { + const auth = expiredAuth({ credentialId: "damaged-id" }); + deps.providersConfig = { + openai: + accountId === "default" + ? { codexOauth: auth } + : { codexOauthAccounts: { work: { label: "Work", credentials: auth } } }, + }; + deviceFetch(); + const flow = await service.startDeviceFlow({ accountId }); + if (!flow.success) throw new Error(flow.error); + const stamped = getCodexOauthAuth(deps.providersConfig.openai, accountId); + expect(stamped?.credentialId).toBeDefined(); + expect(stamped?.credentialId).not.toBe(auth.credentialId); + expect(await service.waitForDeviceFlow(flow.data.flowId)).toEqual(Ok(undefined)); + const reconnected = await service.getValidAuth(accountId); + expect(reconnected.success).toBe(true); + if (reconnected.success) { + expect(reconnected.data.credentialId).not.toBe(stamped?.credentialId); + } + } + ); + + for (const kind of ["desktop", "device"] as const) { + it.each(["cancelled", "expired", "completed"] as const)( + "preserves pinned legacy identity until " + kind + " reconnect is %s", + async (terminal) => { + const { store, first, second } = sharedServices( + validAuth({ credentialId: undefined }), + "default" + ); + const snapshot = { + credentialId: getCodexOauthAuth(store.loadProvidersConfig()?.openai, "default") + ?.credentialId, + }; + expect(snapshot.credentialId).toBeUndefined(); + deviceFetch(); + try { + expect((await second.getValidAuth("default", snapshot)).success).toBe(true); + const flow = + kind === "desktop" + ? await first.startDesktopFlow({ accountId: "default" }) + : await first.startDeviceFlow({ accountId: "default" }); + if (!flow.success) throw new Error(flow.error); + const duringLogin = await second.getValidAuth("default", snapshot); + if (kind === "desktop") { + if (terminal === "cancelled") await first.cancelDesktopFlow(flow.data.flowId); + else if (terminal === "expired") { + expect( + (await first.waitForDesktopFlow(flow.data.flowId, { timeoutMs: 0 })).success + ).toBe(false); + } else { + const callback = await originalFetch( + "http://localhost:1455/auth/callback?code=code&state=" + flow.data.flowId + ); + await callback.text(); + expect(await first.waitForDesktopFlow(flow.data.flowId)).toEqual(Ok(undefined)); + } + } else if (terminal === "cancelled") { + await first.cancelDeviceFlow(flow.data.flowId); + } else if (terminal === "expired") { + const clock = spyOn(Date, "now").mockReturnValue(Date.now() + 120_000); + try { + expect((await first.waitForDeviceFlow(flow.data.flowId)).success).toBe(false); + } finally { + clock.mockRestore(); + } + } else { + expect(await first.waitForDeviceFlow(flow.data.flowId)).toEqual(Ok(undefined)); + } + expect(duringLogin.success).toBe(true); + const pinned = await second.getValidAuth("default", snapshot); + expect(pinned.success).toBe(terminal !== "completed"); + expect((await second.getValidAuth("default")).success).toBe(true); + } finally { + await first.dispose(); + await second.dispose(); + } + } + ); + } + + it("preserves the legacy snapshot through backfill and cross-process token rotation", async () => { + const { store, first, second } = sharedServices(expiredAuth({ credentialId: undefined })); + const snapshot = { credentialId: undefined }; + deviceFetch(() => + Promise.resolve( + mockRefreshResponse({ + access_token: "rotated", + refresh_token: "rotated-refresh", + expires_in: 3600, + }) + ) + ); + try { + const flow = await first.startDeviceFlow({ accountId: "work" }); + if (!flow.success) throw new Error(flow.error); + await first.cancelDeviceFlow(flow.data.flowId); + const result = await second.getValidAuth("work", snapshot); + expect(result).toMatchObject({ + success: true, + data: { access: "rotated", refresh: "rotated-refresh" }, + }); + expect( + getCodexOauthAuth(store.loadProvidersConfig()?.openai, "work")?.credentialId + ).toBeDefined(); + expect(await first.getValidAuth("work", snapshot)).toEqual(result); + } finally { + await first.dispose(); + await second.dispose(); + } + }); + + it("reconnects a marked legacy credential that has no durable ID", async () => { + deps.providersConfig = { + openai: { + codexOauth: expiredAuth({ credentialId: undefined, invalidReason: "invalid_grant" }), + }, + }; + deviceFetch(); + const flow = await service.startDeviceFlow({ accountId: "default" }); + if (!flow.success) throw new Error(flow.error); + const stamped = getCodexOauthAuth(deps.providersConfig.openai); + expect(stamped?.credentialId).toBeDefined(); + expect(stamped?.invalidReason).toBe("invalid_grant"); + expect(await service.waitForDeviceFlow(flow.data.flowId)).toEqual(Ok(undefined)); + const reconnected = await service.getValidAuth(); + expect(reconnected.success).toBe(true); + if (reconnected.success) { + expect(reconnected.data.invalidReason).toBeUndefined(); + expect(reconnected.data.credentialId).not.toBe(stamped?.credentialId); + } + }); + + it("completes cross-process reconnect over the same credential's invalid marker", async () => { + const initial = expiredAuth(); + const { store, first, second } = sharedServices(initial); + deviceFetch((init) => + Promise.resolve( + new URLSearchParams(requestBody(init)).get("grant_type") === "refresh_token" + ? mockRefreshResponse({ error: "invalid_grant" }, 400) + : mockRefreshResponse({ + access_token: "reconnected", + refresh_token: "new-login", + expires_in: 3600, + }) + ) + ); + try { + const flow = await first.startDeviceFlow({ accountId: "work" }); + if (!flow.success) throw new Error(flow.error); + expect((await second.getValidAuth("work")).success).toBe(false); + expect(getCodexOauthAuth(store.loadProvidersConfig()?.openai, "work")).toEqual({ + ...initial, + invalidReason: "invalid_grant", + }); + expect(await first.waitForDeviceFlow(flow.data.flowId)).toEqual(Ok(undefined)); + const result = await first.getValidAuth("work"); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.credentialId).not.toBe(initial.credentialId); + expect(result.data.invalidReason).toBeUndefined(); + expect(result.data.access).toBe("reconnected"); + } + } finally { + await first.dispose(); + await second.dispose(); + } + }); + + it("does not invalidate a completed reconnect when an older process refresh fails", async () => { + const { store, first, second } = sharedServices(expiredAuth()); + const started = createDeferred(); + const response = createDeferred(); + deviceFetch((init) => { + if (new URLSearchParams(requestBody(init)).get("grant_type") === "refresh_token") { + started.resolve(undefined); + return response.promise; + } + return Promise.resolve( + mockRefreshResponse({ + access_token: "reconnected", + refresh_token: "new-login", + expires_in: 3600, + }) + ); + }); + try { + const refresh = second.getValidAuth("work"); + await started.promise; + const flow = await first.startDeviceFlow({ accountId: "work" }); + if (!flow.success) throw new Error(flow.error); + expect(await first.waitForDeviceFlow(flow.data.flowId)).toEqual(Ok(undefined)); + const reconnected = getCodexOauthAuth(store.loadProvidersConfig()?.openai, "work"); + response.resolve(mockRefreshResponse({ error: "invalid_grant" }, 400)); + expect((await refresh).success).toBe(false); + expect(getCodexOauthAuth(store.loadProvidersConfig()?.openai, "work")).toEqual(reconnected); + expect((await first.getValidAuth("work")).success).toBe(true); + } finally { + await first.dispose(); + await second.dispose(); + } + }); + + it("completes reconnect after another service refreshes the same credential", async () => { + const initial = expiredAuth(); + const { store, first, second } = sharedServices(initial); + deviceFetch((init) => + Promise.resolve( + mockRefreshResponse({ + access_token: + new URLSearchParams(requestBody(init)).get("grant_type") === "refresh_token" + ? "rotated-access" + : "login-access", + refresh_token: "new-refresh", + expires_in: 3600, + }) + ) + ); + try { + const flow = await first.startDeviceFlow({ accountId: "work" }); + if (!flow.success) throw new Error(flow.error); + const refreshed = await second.getValidAuth("work"); + expect(refreshed.success).toBe(true); + expect(getCodexOauthAuth(store.loadProvidersConfig()?.openai, "work")?.credentialId).toBe( + initial.credentialId + ); + expect(await first.waitForDeviceFlow(flow.data.flowId)).toEqual(Ok(undefined)); + const reconnected = getCodexOauthAuth(store.loadProvidersConfig()?.openai, "work"); + expect(reconnected?.access).toBe("login-access"); + expect(reconnected?.credentialId).toBeDefined(); + expect(reconnected?.credentialId).not.toBe(initial.credentialId); + } finally { + await first.dispose(); + await second.dispose(); + } + }); + + it.each([undefined, "1c9c50b0-d777-4dd2-998c-09c156ba9754"])( + "lets the first completed reconnect replace credential %s", + async (credentialId) => { + const { store, first, second } = sharedServices(validAuth({ credentialId })); + deviceFetch(); + try { + const older = await first.startDeviceFlow({ accountId: "work" }); + const newer = await second.startDeviceFlow({ accountId: "work" }); + if (!older.success || !newer.success) throw new Error("Login start failed"); + expect(await first.waitForDeviceFlow(older.data.flowId)).toEqual(Ok(undefined)); + const winner = getCodexOauthAuth(store.loadProvidersConfig()?.openai, "work"); + expect((await second.waitForDeviceFlow(newer.data.flowId)).success).toBe(false); + expect(getCodexOauthAuth(store.loadProvidersConfig()?.openai, "work")).toEqual(winner); + } finally { + await first.dispose(); + await second.dispose(); + } + } + ); + + it.each(["removed", "recreated", "older-client"])( + "rejects stale login after external %s credentials", + async (change) => { + const initial = validAuth(); + const { provider, store, first, second } = sharedServices(initial); + deviceFetch(); + try { + const flow = await first.startDeviceFlow({ accountId: "work" }); + if (!flow.success) throw new Error(flow.error); + if (change !== "older-client") { + await provider.setConfigValue("openai", ["codexOauthAccounts", "work"], undefined); + } + if (change !== "removed") { + await provider.setConfigValue("openai", ["codexOauthAccounts", "work"], { + label: "Work", + credentials: { + ...initial, + credentialId: + change === "recreated" ? "50e00a32-b964-4ce2-b131-6b53356ce2db" : undefined, + }, + }); + } + const current = store.loadProvidersConfig()?.openai; + expect((await first.waitForDeviceFlow(flow.data.flowId)).success).toBe(false); + expect(store.loadProvidersConfig()?.openai).toEqual(current); + } finally { + await first.dispose(); + await second.dispose(); + } + } + ); + + it("waits for an old credential refresh before assigning its durable login ID", async () => { + const initial = expiredAuth({ credentialId: undefined }); + const { store, provider, first, second } = sharedServices(initial); + const leaseAttempt = createDeferred(); + class ObservedLeaseManager extends FileLeaseManager { + override withCodexOauthRefreshLock( + accountId: string, + fn: () => Promise | T + ): Promise { + leaseAttempt.resolve(undefined); + return super.withCodexOauthRefreshLock(accountId, fn); + } + } + const reconnectService = new CodexOauthService( + store, + provider, + undefined, + new ObservedLeaseManager(deps.rootDir) + ); + const refreshStarted = createDeferred(); + const response = createDeferred(); + deviceFetch(() => { + refreshStarted.resolve(undefined); + return response.promise; + }); + try { + const refresh = second.getValidAuth("work"); + await refreshStarted.promise; + const startup = reconnectService.startDeviceFlow({ accountId: "work" }); + await leaseAttempt.promise; + expect( + getCodexOauthAuth(store.loadProvidersConfig()?.openai, "work")?.credentialId + ).toBeUndefined(); + response.resolve( + mockRefreshResponse({ + access_token: "rotated-access", + refresh_token: "rotated-refresh", + expires_in: 3600, + }) + ); + expect((await refresh).success).toBe(true); + const flow = await startup; + if (!flow.success) throw new Error(flow.error); + const stamped = getCodexOauthAuth(store.loadProvidersConfig()?.openai, "work"); + expect(stamped?.access).toBe("rotated-access"); + expect(stamped?.refresh).toBe("rotated-refresh"); + expect(stamped?.credentialId).toBeDefined(); + await reconnectService.cancelDeviceFlow(flow.data.flowId); + expect(await first.getValidAuth("work")).toEqual(Ok(stamped!)); + } finally { + await reconnectService.dispose(); + await first.dispose(); + await second.dispose(); + } + }); + + for (const startup of ["device", "desktop"] as const) { + it.each([undefined, "damaged-id"])( + `keeps legacy requests usable after failed ${startup} startup with ID %s`, + async (credentialId) => { + const auth = validAuth({ credentialId }); + deps.providersConfig = { openai: { codexOauth: auth } }; + if (startup === "device") { + mockFetch(() => Promise.reject(new Error("Device startup failed"))); + } else { + spyOn(http, "createServer").mockImplementationOnce(() => { + throw new Error("Listener startup failed"); + }); + } + const result = + startup === "device" + ? await service.startDeviceFlow({ accountId: "default" }) + : await service.startDesktopFlow({ accountId: "default" }); + expect(result.success).toBe(false); + expect(deps.providersConfig.openai?.codexOauth).toEqual(auth); + expect((await service.getValidAuth("default", { credentialId: undefined })).success).toBe( + true + ); + } + ); + } + + it("closes the listener when the credential changes before its ID write", async () => { + const { store, provider, first, second } = sharedServices( + validAuth({ credentialId: undefined }) + ); + const replacement = validAuth({ refresh: "replacement" }); + const update = provider.updateProviderSection.bind(provider); + const write = spyOn(provider, "updateProviderSection").mockImplementationOnce( + (name, transform, options) => { + store.saveProvidersConfig({ + openai: { codexOauthAccounts: { work: { label: "Work", credentials: replacement } } }, + }); + return update(name, transform, options); + } + ); + try { + expect((await first.startDesktopFlow({ accountId: "work" })).success).toBe(false); + expect(getCodexOauthAuth(store.loadProvidersConfig()?.openai, "work")).toEqual(replacement); + // A second bind proves failed startup releases the listener. + const retry = await first.startDesktopFlow({ accountId: "work" }); + if (!retry.success) throw new Error(retry.error); + await first.cancelDesktopFlow(retry.data.flowId); + } finally { + write.mockRestore(); + await first.dispose(); + await second.dispose(); + } + }); + + it("releases the listener when a legacy ID write fails", async () => { + const auth = validAuth({ credentialId: undefined }); + deps.providersConfig = { openai: { codexOauth: auth } }; + deps.policyDenied = true; + expect((await service.startDesktopFlow({ accountId: "default" })).success).toBe(false); + expect((await service.getValidAuth("default", { credentialId: undefined })).success).toBe( + true + ); + expect(deps.providersConfig.openai?.codexOauth).toEqual(auth); + deps.policyDenied = false; + const retry = await service.startDesktopFlow({ accountId: "default" }); + if (!retry.success) throw new Error(retry.error); + await service.cancelDesktopFlow(retry.data.flowId); + }); + + it("rejects a revoked account as a new global default", async () => { + deps.providersConfig = { + openai: { + codexOauth: validAuth(), + codexOauthDefaultAccountId: "default", + codexOauthAccounts: { + work: { label: "Work", credentials: validAuth({ invalidReason: "invalid_grant" }) }, + }, + }, + }; + expect((await service.setDefaultAccount("work")).success).toBe(false); + expect(deps.providersConfig.openai?.codexOauthDefaultAccountId).toBe("default"); + }); + + it.each(["default", "work"])( + "rejects an older %s startup that completes after its retry", + async (accountId) => { + const auth = validAuth(); + deps.providersConfig = { + openai: + accountId === "default" + ? { codexOauth: auth } + : { codexOauthAccounts: { work: { label: "Work", credentials: auth } } }, + }; + deviceFetch(); + const fetch = globalThis.fetch; + const firstStarted = createDeferred(); + const firstResponse = createDeferred(); + let starts = 0; + mockFetch((input, init) => { + if (input === CODEX_OAUTH_DEVICE_USERCODE_URL && starts++ === 0) { + firstStarted.resolve(undefined); + return firstResponse.promise; + } + return fetch(input, init); + }); + const older = service.startDeviceFlow({ accountId }); + await firstStarted.promise; + const newer = await service.startDeviceFlow({ accountId }); + if (!newer.success) throw new Error(newer.error); + firstResponse.resolve( + mockRefreshResponse({ device_auth_id: "older", user_code: "OLD", interval: 1 }) + ); + expect((await older).success).toBe(false); + expect(await service.waitForDeviceFlow(newer.data.flowId)).toEqual(Ok(undefined)); + } + ); + + it("keeps the active device login when a newer device startup fails", async () => { + deps.providersConfig = { openai: { codexOauth: validAuth() } }; + deviceFetch(); + const first = await service.startDeviceFlow({ accountId: "default" }); + if (!first.success) throw new Error(first.error); + const workingFetch = globalThis.fetch; + mockFetch((input, init) => + input === CODEX_OAUTH_DEVICE_USERCODE_URL + ? Promise.reject(new Error("Device startup failed")) + : workingFetch(input, init) + ); + expect((await service.startDeviceFlow({ accountId: "default" })).success).toBe(false); + expect(await service.waitForDeviceFlow(first.data.flowId)).toEqual(Ok(undefined)); + }); + + it("keeps the active desktop login when the new listener cannot start", async () => { + deps.providersConfig = { openai: { codexOauth: validAuth() } }; + deviceFetch(); + const first = await service.startDesktopFlow({ accountId: "default" }); + if (!first.success) throw new Error(first.error); + const createServer = spyOn(http, "createServer").mockImplementationOnce(() => { + throw new Error("Callback listener unavailable"); + }); + try { + expect((await service.startDesktopFlow({ accountId: "default" })).success).toBe(false); + } finally { + createServer.mockRestore(); + } + const response = await originalFetch( + "http://localhost:1455/auth/callback?code=code&state=" + first.data.flowId + ); + await response.text(); + expect(response.ok).toBe(true); + expect(await service.waitForDesktopFlow(first.data.flowId)).toEqual(Ok(undefined)); + }); + + it("completes a reconnect after a concurrent refresh persists first", async () => { + deps.providersConfig = { + openai: { + codexOauthAccounts: { + work: { label: "Work", credentials: expiredAuth({ refresh: "old" }) }, + }, + }, + }; + const refreshStarted = createDeferred(); + const refreshResponse = createDeferred(); + const exchangeStarted = createDeferred(); + const exchangeResponse = createDeferred(); + deviceFetch((init) => { + if (new URLSearchParams(requestBody(init)).get("grant_type") === "refresh_token") { + refreshStarted.resolve(undefined); + return refreshResponse.promise; + } + exchangeStarted.resolve(undefined); + return exchangeResponse.promise; + }); + const refresh = service.getValidAuth("work"); + await refreshStarted.promise; + const flow = await service.startDeviceFlow({ accountId: "work" }); + if (!flow.success) throw new Error(flow.error); + const reconnect = service.waitForDeviceFlow(flow.data.flowId); + await exchangeStarted.promise; + refreshResponse.resolve( + mockRefreshResponse({ + access_token: "refreshed", + refresh_token: "rotated", + expires_in: 3600, + }) + ); + expect((await refresh).success).toBe(true); + exchangeResponse.resolve( + mockRefreshResponse({ + access_token: "reconnected", + refresh_token: "new-login", + expires_in: 3600, + }) + ); + expect(await reconnect).toEqual(Ok(undefined)); + expect(getCodexOauthAuth(deps.providersConfig.openai, "work")?.refresh).toBe("new-login"); + }); + + it.each([{ accounts: [] }, { accounts: null }, { accounts: 42 }])( + "repairs a malformed account map during named login: %j", + async ({ accounts }) => { + const provider = new ProviderService(new Config(deps.rootDir)); + const store = provider.providersConfigStore; + store.saveProvidersConfig({ openai: { codexOauthAccounts: accounts, apiKey: "keep-key" } }); + const realService = new CodexOauthService(store, provider); + deviceFetch(); + try { + const flow = await realService.startDeviceFlow({ label: "Work" }); + if (!flow.success) throw new Error(flow.error); + expect(await realService.waitForDeviceFlow(flow.data.flowId)).toEqual(Ok(undefined)); + const openai = store.loadProvidersConfig()?.openai; + const connected = getCodexOauthAccounts(openai); + expect(connected).toHaveLength(1); + expect(getCodexOauthAccountId(openai)).toBe(connected[0].id); + expect(openai?.apiKey).toBe("keep-key"); + } finally { + await realService.dispose(); + } + } + ); + + it("does not leave a partial first login when saving the selected account fails", async () => { + const provider = new ProviderService(new Config(deps.rootDir)); + const store = provider.providersConfigStore; + store.saveProvidersConfig({ openai: { apiKey: "keep-key" } }); + const save = store.saveProvidersConfig.bind(store); + const saveSpy = spyOn(store, "saveProvidersConfig").mockImplementation((config) => { + // Fail the write that makes the new account the selected account. + const openai = config.openai; + if ( + getCodexOauthAccounts(openai).some( + (account) => account.id === getCodexOauthAccountId(openai) + ) + ) { + throw new Error("Disk unavailable"); + } + save(config); + }); + const realService = new CodexOauthService(store, provider); + deviceFetch(); + try { + const flow = await realService.startDeviceFlow({ label: "Work" }); + if (!flow.success) throw new Error(flow.error); + expect((await realService.waitForDeviceFlow(flow.data.flowId)).success).toBe(false); + expect(store.loadProvidersConfig()?.openai).toEqual({ apiKey: "keep-key" }); + } finally { + saveSpy.mockRestore(); + await realService.dispose(); + } + }); + + it("keeps a successful reconnect when an older refresh completes", async () => { + deps.providersConfig = { + openai: { + codexOauthAccounts: { + work: { label: "Work", credentials: expiredAuth({ refresh: "old" }) }, + }, + }, + }; + const refreshStarted = createDeferred(); + const refreshResponse = createDeferred(); + deviceFetch((init) => { + const body = new URLSearchParams(requestBody(init)); + if (body.get("grant_type") === "refresh_token") { + refreshStarted.resolve(undefined); + return refreshResponse.promise; + } + return Promise.resolve( + mockRefreshResponse({ + access_token: "reconnected", + refresh_token: "new", + expires_in: 3600, + }) + ); + }); + const refresh = service.getValidAuth("work"); + await refreshStarted.promise; + const flow = await service.startDeviceFlow({ accountId: "work" }); + if (!flow.success) throw new Error(flow.error); + expect(await service.waitForDeviceFlow(flow.data.flowId)).toEqual(Ok(undefined)); + refreshResponse.resolve(mockRefreshResponse({ access_token: "stale", expires_in: 3600 })); + expect((await refresh).success).toBe(false); + expect(getCodexOauthAuth(deps.providersConfig.openai, "work")?.access).toBe("reconnected"); + }); + + it("lets a reconnect replace credentials without changing its label or another slot", async () => { + const legacy = validAuth({ access: "legacy" }); + deps.providersConfig = { + openai: { + codexOauth: legacy, + codexOauthAccounts: { work: { label: "Work", credentials: validAuth() } }, + }, + }; + deviceFetch(); + const flow = await service.startDeviceFlow({ accountId: "work" }); + if (!flow.success) throw new Error(flow.error); + await service.renameAccount("work", "Team"); + expect(await service.waitForDeviceFlow(flow.data.flowId)).toEqual(Ok(undefined)); + expect(getCodexOauthAccounts(deps.providersConfig.openai)[1].label).toBe("Team"); + expect(getCodexOauthAuth(deps.providersConfig.openai, "work")?.access).toBe( + "access-device-1" + ); + expect(await service.getValidAuth()).toEqual(Ok(legacy)); + }); + }); }); diff --git a/src/node/services/codexOauthService.ts b/src/node/services/codexOauthService.ts index 4b03ca842c1..ff7e1668f1c 100644 --- a/src/node/services/codexOauthService.ts +++ b/src/node/services/codexOauthService.ts @@ -25,7 +25,13 @@ import { CODEX_OAUTH_DEVICE_VERIFY_URL, CODEX_OAUTH_TOKEN_URL, } from "@/common/constants/codexOAuth"; -import type { ProvidersConfigStore } from "@/node/config"; +import { + CODEX_OAUTH_DEFAULT_ACCOUNT_ID, + CODEX_OAUTH_ACCOUNT_LABEL_MAX_LENGTH, + CODEX_OAUTH_REFRESH_TIMEOUT_MS, + CODEX_OAUTH_START_TIMEOUT_MS, +} from "@/common/constants/codexOauthAccounts"; +import { FileLeaseManager, type ProvidersConfigStore } from "@/node/config"; import type { ProviderService } from "@/node/services/providerService"; import type { WindowService } from "@/node/services/windowService"; import { log } from "@/node/services/log"; @@ -33,7 +39,11 @@ import { sleepWithAbort } from "@/node/utils/abort"; import { AsyncMutex } from "@/node/utils/concurrency/asyncMutex"; import { extractAccountIdFromTokens, + getCodexOauthAccounts, + getCodexOauthAccountId, + getCodexOauthAuth, isCodexOauthAuthExpired, + isValidCodexOauthAccountId, parseCodexOauthAuth, type CodexOauthAuth, } from "@/node/utils/codexOauthAuth"; @@ -46,7 +56,27 @@ const DEFAULT_DESKTOP_TIMEOUT_MS = 5 * 60 * 1000; const DEFAULT_DEVICE_TIMEOUT_MS = 15 * 60 * 1000; const COMPLETED_FLOW_TTL_MS = 60 * 1000; +export interface CodexOauthLoginOptions { + label?: string; + accountId?: string; +} + +interface CodexOauthCredentialSnapshot { + // An undefined ID pins a legacy credential. An omitted snapshot does not constrain the ID. + credentialId: string | undefined; +} + +interface AccountSelection { + accountId: string; + revision: number; + credentialId?: string; + auth: CodexOauthAuth | null; + label?: string; + selectAsDefault: boolean; +} + interface DeviceFlow { + destination: AccountSelection; flowId: string; deviceAuthId: string; userCode: string; @@ -118,50 +148,117 @@ export class CodexOauthError extends Schema.TaggedError()("Code reason: Schema.String, }) {} +function createNamedAccount(accountId: string, label: unknown, credentials: CodexOauthAuth) { + const trimmedLabel = typeof label === "string" ? label.trim() : ""; + // Write only known fields. Disk entries can contain unsafe token copies outside credentials. + return { label: trimmedLabel || accountId, credentials }; +} + +function isValidLabel(label: string): boolean { + return label.trim().length > 0 && label.trim().length <= CODEX_OAUTH_ACCOUNT_LABEL_MAX_LENGTH; +} + +function matchesAuth(actual: CodexOauthAuth | null, expected: CodexOauthAuth | null): boolean { + if (!actual || !expected) return actual === expected; + return ( + actual.access === expected.access && + actual.refresh === expected.refresh && + actual.expires === expected.expires && + actual.accountId === expected.accountId && + actual.credentialId === expected.credentialId && + actual.legacyCredentialId === expected.legacyCredentialId && + actual.invalidReason === expected.invalidReason + ); +} + export class CodexOauthService { private readonly desktopFlows = new OAuthFlowManager(); private readonly deviceFlows = new Map(); - private readonly refreshMutex = new AsyncMutex(); - - // In-memory cache so getValidAuth() skips disk reads when tokens are valid. - // Invalidated on every write (exchange, refresh, disconnect). - private cachedAuth: CodexOauthAuth | null = null; + private readonly refreshMutexes = new Map(); + private readonly accountRevisions = new Map(); + private readonly loginSelections = new Map(); + private readonly loginStartupGenerations = new Map(); + private nextLoginStartupGeneration = 0; + private readonly authMutationMutexes = new Map(); constructor( private readonly providersConfigStore: ProvidersConfigStore, private readonly providerService: ProviderService, - private readonly windowService?: WindowService + private readonly windowService?: WindowService, + private readonly fileLeaseManager = new FileLeaseManager(providersConfigStore.rootDir) ) {} - async disconnect(): Promise> { - return Effect.runPromise(this.disconnectEffect()); + async disconnect(accountId?: string): Promise> { + return Effect.runPromise(this.disconnectEffect(accountId)); } - /** - * Wire-shaped Effect surface for handlerGen router handlers. Uninterruptible - * (mirrors asAtomicMutation in providerService.ts): a client abort must not - * skip the persisted-credential clear after the in-memory cache was already - * invalidated. - */ - disconnectEffect(): Effect.Effect> { - // eslint-disable-next-line @typescript-eslint/no-this-alias -- Effect.gen generator bodies do not inherit `this` - const self = this; - return Effect.uninterruptible( - Effect.gen(function* () { - // Clear stored ChatGPT OAuth tokens so Codex-only models are hidden again. - self.cachedAuth = null; - // setConfigValue resolves with a wire Result; a rejection stays a - // defect, matching the previously un-caught await. - return yield* Effect.promise(() => - self.providerService.setConfigValue("openai", ["codexOauth"], undefined) + // No argument retains legacy behavior. Account-aware callers pass the selected slot ID. + disconnectEffect( + accountId = CODEX_OAUTH_DEFAULT_ACCOUNT_ID + ): Effect.Effect> { + return Effect.suspend(() => { + if (!isValidCodexOauthAccountId(accountId)) + return Effect.succeed(Err("Invalid Codex OAuth account ID")); + return this.withAccountMutationEffect( + accountId, + this.updateConfigValueEffect(this.accountPath(accountId), () => ({ + value: undefined, + })).pipe( + Effect.map((result) => { + // Failed deletion must preserve active requests and logins. Fence stale writes before releasing the mutex. + if (result.success) { + this.accountRevisions.set(accountId, this.getAccountRevision(accountId) + 1); + } + return result; + }) + ) + ); + }); + } + + async setDefaultAccount(accountId: string): Promise> { + return Effect.runPromise(this.setDefaultAccountEffect(accountId)); + } + + setDefaultAccountEffect(accountId: string): Effect.Effect> { + return Effect.suspend(() => { + if (!isValidCodexOauthAccountId(accountId)) + return Effect.succeed(Err("Invalid Codex OAuth account ID")); + return this.updateConfigValueEffect(["codexOauthDefaultAccountId"], () => { + const auth = this.readStoredAuth(accountId); + return auth && !auth.invalidReason ? { value: accountId } : null; + }); + }); + } + + async renameAccount(accountId: string, label: string): Promise> { + return Effect.runPromise(this.renameAccountEffect(accountId, label)); + } + + renameAccountEffect(accountId: string, label: string): Effect.Effect> { + return Effect.suspend(() => { + if (!isValidCodexOauthAccountId(accountId)) + return Effect.succeed(Err("Invalid Codex OAuth account ID")); + if (!isValidLabel(label)) return Effect.succeed(Err("Invalid Codex OAuth account label")); + if (accountId === CODEX_OAUTH_DEFAULT_ACCOUNT_ID) { + return this.updateConfigValueEffect(["codexOauthLabel"], () => + this.readStoredAuth(accountId) ? { value: label.trim() } : null ); - }) - ); + } + return this.updateConfigValueEffect(this.accountPath(accountId), (current) => { + const credentials = isPlainObject(current) + ? parseCodexOauthAuth(current.credentials) + : null; + return credentials ? { value: createNamedAccount(accountId, label, credentials) } : null; + }); + }); } - async startDesktopFlow(): Promise> { - return Effect.runPromise(this.startDesktopFlowEffect()); + async startDesktopFlow( + options?: CodexOauthLoginOptions + ): Promise> { + return Effect.runPromise(this.startDesktopFlowEffect(options)); } /** @@ -172,16 +269,15 @@ export class CodexOauthService { * and local, so running it to completion on abort is cheap; an abandoned * flow still self-cleans via the registered timeout. */ - startDesktopFlowEffect(): Effect.Effect< - Result<{ flowId: string; authorizeUrl: string }, string> - > { - return Effect.uninterruptible(toWireResult(this.launchDesktopFlowEffect())); + startDesktopFlowEffect( + options?: CodexOauthLoginOptions + ): Effect.Effect> { + return Effect.uninterruptible(toWireResult(this.launchDesktopFlowEffect(options))); } - private launchDesktopFlowEffect(): Effect.Effect< - { flowId: string; authorizeUrl: string }, - CodexOauthError - > { + private launchDesktopFlowEffect( + options?: CodexOauthLoginOptions + ): Effect.Effect<{ flowId: string; authorizeUrl: string }, CodexOauthError> { // eslint-disable-next-line @typescript-eslint/no-this-alias -- Effect.gen generator bodies do not inherit `this` const self = this; return Effect.gen(function* () { @@ -191,27 +287,37 @@ export class CodexOauthService { const codeChallenge = sha256Base64Url(codeVerifier); const redirectUri = CODEX_OAUTH_BROWSER_REDIRECT_URI; - const loopback = yield* Effect.tryPromise({ - try: () => - startLoopbackServer({ - port: 1455, - host: "localhost", - callbackPath: "/auth/callback", - validateLoopback: true, - expectedState: flowId, - deferSuccessResponse: true, - }), - catch: (error) => - new CodexOauthError({ - reason: `Failed to start OAuth callback listener: ${getErrorMessage(error)}`, - }), - }); + const { destination, resource: loopback } = yield* self.startLogin( + options, + Effect.tryPromise({ + try: () => + startLoopbackServer({ + port: 1455, + host: "localhost", + callbackPath: "/auth/callback", + validateLoopback: true, + expectedState: flowId, + deferSuccessResponse: true, + }), + catch: (error) => + new CodexOauthError({ + reason: `Failed to start OAuth callback listener: ${getErrorMessage(error)}`, + }), + }), + (loopback) => loopback.cancel() + ); const resultDeferred = createDeferred>(); self.desktopFlows.register(flowId, { server: loopback.server, - resultDeferred, + resultDeferred: { + ...resultDeferred, + resolve: (result) => { + self.clearLoginSelection(destination); + resultDeferred.resolve(result); + }, + }, // Keep server-side timeout tied to flow lifetime so abandoned flows // (e.g. callers that never invoke waitForDesktopFlow) still self-clean. timeoutHandle: setTimeout(() => { @@ -233,6 +339,7 @@ export class CodexOauthService { // cancelled. Effect.runFork( self.desktopCallbackPipeline({ + destination, flowId, redirectUri, codeVerifier, @@ -254,6 +361,7 @@ export class CodexOauthService { * dangling on loopback.result. */ private desktopCallbackPipeline(args: { + destination: AccountSelection; flowId: string; redirectUri: string; codeVerifier: string; @@ -277,6 +385,8 @@ export class CodexOauthService { const exchangeResult: Result = yield* toWireResult( self.handleDesktopCallbackAndExchange({ + destination: args.destination, + isActive: () => self.desktopFlows.has(args.flowId), flowId: args.flowId, redirectUri: args.redirectUri, codeVerifier: args.codeVerifier, @@ -338,7 +448,7 @@ export class CodexOauthService { ); } - async startDeviceFlow(): Promise< + async startDeviceFlow(options?: CodexOauthLoginOptions): Promise< Result< { flowId: string; @@ -349,7 +459,7 @@ export class CodexOauthService { string > > { - return Effect.runPromise(this.startDeviceFlowEffect()); + return Effect.runPromise(this.startDeviceFlowEffect(options)); } /** @@ -359,7 +469,9 @@ export class CodexOauthService { * flow record (and its expiry timeout) that lets callers re-attach or the * flow self-clean. */ - startDeviceFlowEffect(): Effect.Effect< + startDeviceFlowEffect( + options?: CodexOauthLoginOptions + ): Effect.Effect< Result<{ flowId: string; userCode: string; verifyUrl: string; intervalSeconds: number }, string> > { // eslint-disable-next-line @typescript-eslint/no-this-alias -- Effect.gen generator bodies do not inherit `this` @@ -369,8 +481,12 @@ export class CodexOauthService { Effect.gen(function* () { const flowId = randomBase64Url(); - const { deviceAuthId, userCode, intervalSeconds, expiresAtMs } = - yield* self.requestDeviceUserCode(); + const { + destination, + resource: { deviceAuthId, userCode, intervalSeconds, expiresAtMs }, + } = yield* self.startLogin(options, self.requestDeviceUserCode(), () => + Promise.resolve() + ); const verifyUrl = CODEX_OAUTH_DEVICE_VERIFY_URL; const { promise: resultPromise, resolve: resolveResult } = @@ -387,6 +503,7 @@ export class CodexOauthService { }, timeoutMs); self.deviceFlows.set(flowId, { + destination, flowId, deviceAuthId, userCode, @@ -503,47 +620,93 @@ export class CodexOauthService { ); } - async getValidAuth(): Promise> { - return Effect.runPromise(this.getValidAuthEffect()); + async getValidAuth( + accountId?: string, + expectedCredential?: CodexOauthCredentialSnapshot + ): Promise> { + return Effect.runPromise(this.getValidAuthEffect(accountId, expectedCredential)); } - getValidAuthEffect(): Effect.Effect> { - // eslint-disable-next-line @typescript-eslint/no-this-alias -- Effect.gen generator bodies do not inherit `this` + getValidAuthEffect( + accountId?: string, + expectedCredential?: CodexOauthCredentialSnapshot + ): Effect.Effect> { + // eslint-disable-next-line @typescript-eslint/no-this-alias -- Effect generators do not inherit this. const self = this; return Effect.gen(function* () { - const stored = self.readStoredAuth(); - if (!stored) { - return Err("Codex OAuth is not configured"); - } + // Pin both the slot and credential. A replacement must not change an active request's account. + const selectedId = getCodexOauthAccountId(self.readOpenaiConfig(), accountId); + if (!isValidCodexOauthAccountId(selectedId)) return Err("Invalid Codex OAuth account ID"); + const stored = self.readStoredAuth(selectedId); + const selection = { + accountId: selectedId, + revision: self.getAccountRevision(selectedId), + credentialId: expectedCredential ? expectedCredential.credentialId : stored?.credentialId, + }; + const initial = self.validateRequestAuth(stored, selection); + if (!initial.success || !isCodexOauthAuthExpired(initial.data)) return initial; - if (!isCodexOauthAuthExpired(stored)) { - return Ok(stored); + let mutex = self.refreshMutexes.get(selectedId); + if (!mutex) { + mutex = new AsyncMutex(); + self.refreshMutexes.set(selectedId, mutex); } - - // acquireUseRelease guarantees the mutex is released on every exit path - // (refresh success/failure, defects, interruption) — the Effect - // equivalent of the pre-Effect `await using` lock. + const refreshMutex = mutex; return yield* Effect.acquireUseRelease( - Effect.promise(() => self.refreshMutex.acquire()), + Effect.promise(() => refreshMutex.acquire()), () => Effect.gen(function* () { - // Re-read after acquiring lock in case another caller refreshed first. - const latest = self.readStoredAuth(); - if (!latest) { - return Err("Codex OAuth is not configured"); - } - - if (!isCodexOauthAuthExpired(latest)) { - return Ok(latest); - } - - return yield* toWireResult(self.refreshTokens(latest)); + const afterMutex = self.validateRequestAuth(self.readStoredAuth(selectedId), selection); + if (!afterMutex.success || !isCodexOauthAuthExpired(afterMutex.data)) return afterMutex; + // Hold the file lease through persistence. Other processes adopt only the same credential's rotation. + return yield* Effect.tryPromise({ + try: () => + self.fileLeaseManager.withCodexOauthRefreshLock(selectedId, async () => { + const afterLease = self.validateRequestAuth( + self.readStoredAuth(selectedId), + selection + ); + if (!afterLease.success || !isCodexOauthAuthExpired(afterLease.data)) + return afterLease; + return await Effect.runPromise( + toWireResult( + self.refreshTokens( + { ...selection, auth: afterLease.data, selectAsDefault: false }, + afterLease.data + ) + ) + ); + }), + catch: (error) => getErrorMessage(error), + }).pipe( + Effect.catch((error) => Effect.succeed(Err(`Codex OAuth refresh failed: ${error}`))) + ); }), (lock) => Effect.promise(() => lock[Symbol.asyncDispose]()) ); }); } + private validateRequestAuth( + auth: CodexOauthAuth | null, + selection: Pick + ): Result { + if (!auth) return Err(`Codex OAuth account "${selection.accountId}" is not configured`); + const retainsLegacySnapshot = + selection.credentialId === undefined && + auth.legacyCredentialId !== undefined && + auth.legacyCredentialId === auth.credentialId; + if ( + (auth.credentialId !== selection.credentialId && !retainsLegacySnapshot) || + this.getAccountRevision(selection.accountId) !== selection.revision + ) { + return Err("Codex OAuth account changed during the request"); + } + if (auth.invalidReason) + return Err(`Codex OAuth account "${selection.accountId}" needs reconnect`); + return Ok(auth); + } + async dispose(): Promise { await this.desktopFlows.shutdownAll(); @@ -562,35 +725,344 @@ export class CodexOauthService { this.deviceFlows.clear(); } - private readStoredAuth(): CodexOauthAuth | null { - if (this.cachedAuth) { - return this.cachedAuth; + private readOpenaiConfig(): unknown { + return this.providersConfigStore.loadProvidersConfig()?.openai; + } + + private readStoredAuth(accountId: string): CodexOauthAuth | null { + // Read storage so another process cannot leave this service with stale credentials. + return getCodexOauthAuth(this.readOpenaiConfig(), accountId); + } + + private getAccountRevision(accountId: string): number { + return this.accountRevisions.get(accountId) ?? 0; + } + + private accountPath(accountId: string): string[] { + // Do not mirror named credentials into the legacy slot. Older versions refresh that slot independently. + // Keep existing legacy credentials until explicit disconnect; downgrade support does not include named accounts. + return accountId === CODEX_OAUTH_DEFAULT_ACCOUNT_ID + ? ["codexOauth"] + : ["codexOauthAccounts", accountId]; + } + + private updateConfigValueEffect( + keyPath: string[], + update: (current: unknown) => { value: unknown } | null + ): Effect.Effect> { + return this.configMutationEffect(() => + this.providerService.updateConfigValue("openai", keyPath, update, { enforcePolicy: true }) + ); + } + + private configMutationEffect( + mutation: () => Promise> + ): Effect.Effect> { + return Effect.uninterruptible( + Effect.tryPromise({ + try: mutation, + catch: (error) => getErrorMessage(error), + }).pipe( + Effect.map((result): Result => { + if (!result.success) return result; + return result.data.applied + ? Ok(undefined) + : Err("Codex OAuth account changed or is not configured"); + }), + Effect.catch((error) => Effect.succeed(Err(error))) + ) + ); + } + + private withAccountMutationEffect( + accountId: string, + mutation: Effect.Effect + ): Effect.Effect { + let mutex = this.authMutationMutexes.get(accountId); + if (!mutex) { + mutex = new AsyncMutex(); + this.authMutationMutexes.set(accountId, mutex); } - const providersConfig = this.providersConfigStore.loadProvidersConfig() ?? {}; - const openaiConfig = providersConfig.openai as Record | undefined; - const auth = parseCodexOauthAuth(openaiConfig?.codexOauth); - this.cachedAuth = auth; - return auth; + const mutationMutex = mutex; + // Keep persistence and local snapshot updates together. Network requests stay outside this lock. + return Effect.uninterruptible( + Effect.acquireUseRelease( + Effect.promise(() => mutationMutex.acquire()), + () => mutation, + (lock) => Effect.promise(() => lock[Symbol.asyncDispose]()) + ) + ); } - private persistAuth(auth: CodexOauthAuth): Effect.Effect> { - // eslint-disable-next-line @typescript-eslint/no-this-alias -- Effect.gen generator bodies do not inherit `this` + private startLogin( + options: CodexOauthLoginOptions | undefined, + start: Effect.Effect, + cleanup: (resource: T) => Promise + ): Effect.Effect<{ destination: AccountSelection; resource: T }, CodexOauthError> { + return Effect.tryPromise({ + try: async () => { + const initial = await Effect.runPromise(this.selectLoginDestination(options)); + const generation = ++this.nextLoginStartupGeneration; + this.loginStartupGenerations.set(initial.accountId, generation); + const assertCurrentStartup = () => { + if (this.loginStartupGenerations.get(initial.accountId) !== generation) { + throw new Error("Codex OAuth login startup was superseded"); + } + }; + const prepare = async (destination: AccountSelection) => { + const resource = await Effect.runPromise(start); + try { + assertCurrentStartup(); + const current = this.readStoredAuth(destination.accountId); + if ( + this.getAccountRevision(destination.accountId) !== destination.revision || + current?.credentialId !== destination.credentialId || + (!destination.credentialId && !matchesAuth(current, destination.auth)) + ) { + throw new Error("Codex OAuth account changed during login startup"); + } + if (destination.auth && !destination.credentialId) { + const auth = await this.initializeCredentialId(destination); + destination = { ...destination, auth, credentialId: auth.credentialId }; + } + // Claim the selection before yielding. An older startup must not replace a newer invocation. + assertCurrentStartup(); + this.loginSelections.set(destination.accountId, destination); + return { destination, resource }; + } catch (error) { + await cleanup(resource); + throw error; + } + }; + try { + if (!initial.auth || initial.credentialId) return await prepare(initial); + // Backfill an ID for cross-process reconnect checks without invalidating active legacy snapshots. + // Hold rotations until startup and stamping finish. + return await this.fileLeaseManager.withCodexOauthRefreshLock( + initial.accountId, + async () => { + assertCurrentStartup(); + const destination = await Effect.runPromise(this.selectLoginDestination(options)); + if (destination.revision !== initial.revision) { + throw new Error("Codex OAuth account changed during login startup"); + } + return prepare(destination); + } + ); + } finally { + if (this.loginStartupGenerations.get(initial.accountId) === generation) { + this.loginStartupGenerations.delete(initial.accountId); + } + } + }, + catch: (error) => new CodexOauthError({ reason: getErrorMessage(error) }), + }); + } + + /** The caller holds the refresh lease through startup and this conditional write. */ + private async initializeCredentialId(destination: AccountSelection): Promise { + let selected: CodexOauthAuth | null = null; + const { accountId } = destination; + const result = await this.providerService.updateProviderSection( + "openai", + (section) => { + const current = getCodexOauthAuth(section, accountId); + if ( + !current || + !matchesAuth(current, destination.auth) || + this.getAccountRevision(accountId) !== destination.revision + ) { + throw new Error("Codex OAuth account changed during login startup"); + } + const credentialId = crypto.randomUUID(); + selected = { ...current, credentialId, legacyCredentialId: credentialId }; + const next = { ...section }; + if (accountId === CODEX_OAUTH_DEFAULT_ACCOUNT_ID) { + next.codexOauth = selected; + } else { + const accounts = isPlainObject(section?.codexOauthAccounts) + ? section.codexOauthAccounts + : {}; + const entry = accounts[accountId]; + next.codexOauthAccounts = { + ...accounts, + [accountId]: createNamedAccount( + accountId, + isPlainObject(entry) ? entry.label : undefined, + selected + ), + }; + } + return { value: next }; + }, + { enforcePolicy: true } + ); + if (!result.success) throw new Error(result.error); + if (!selected) throw new Error("Codex OAuth account is not configured"); + return selected; + } + + private selectLoginDestination( + options?: CodexOauthLoginOptions + ): Effect.Effect { + // eslint-disable-next-line @typescript-eslint/no-this-alias -- Effect generators do not inherit this. const self = this; return Effect.gen(function* () { - // setConfigValue resolves with a wire Result; a rejection stays a - // defect, matching the previously un-caught await. - const result = yield* Effect.promise(() => - self.providerService.setConfigValue("openai", ["codexOauth"], auth) - ); - // Invalidate cache so the next readStoredAuth() picks up the persisted value from disk. - // We clear rather than set because setConfigValue may have side-effects (e.g. file-write - // failures) and we want the next read to be authoritative. - self.cachedAuth = null; - return result; + if (options?.accountId !== undefined && options.label !== undefined) { + return yield* Effect.fail( + new CodexOauthError({ reason: "Specify an account ID or a label, not both" }) + ); + } + if (options?.label !== undefined && !isValidLabel(options.label)) { + return yield* Effect.fail( + new CodexOauthError({ reason: "Invalid Codex OAuth account label" }) + ); + } + const accountId = + options?.accountId ?? + (options?.label !== undefined ? crypto.randomUUID() : CODEX_OAUTH_DEFAULT_ACCOUNT_ID); + if (!isValidCodexOauthAccountId(accountId)) { + return yield* Effect.fail( + new CodexOauthError({ reason: "Invalid Codex OAuth account ID" }) + ); + } + const auth = self.readStoredAuth(accountId); + if (options?.accountId !== undefined && !auth) { + return yield* Effect.fail( + new CodexOauthError({ reason: "Codex OAuth account is not configured" }) + ); + } + const revision = self.getAccountRevision(accountId); + return { + accountId, + revision, + credentialId: auth?.credentialId, + auth, + label: options?.label?.trim(), + selectAsDefault: + options?.label !== undefined && + getCodexOauthAccounts(self.readOpenaiConfig()).length === 0, + }; }); } + private persistAuth( + selection: AccountSelection, + auth: CodexOauthAuth + ): Effect.Effect> { + return this.withAccountMutationEffect( + selection.accountId, + this.updateConfigValueEffect(this.accountPath(selection.accountId), (current) => { + const legacy = selection.accountId === CODEX_OAUTH_DEFAULT_ACCOUNT_ID; + const stored = parseCodexOauthAuth( + legacy ? current : isPlainObject(current) ? current.credentials : undefined + ); + // Compare under the file lock. Old refreshes must not restore deleted or reconnected slots. + if ( + this.getAccountRevision(selection.accountId) !== selection.revision || + !matchesAuth(stored, selection.auth) + ) + return null; + if (legacy) return { value: auth }; + return { + value: createNamedAccount( + selection.accountId, + isPlainObject(current) ? current.label : undefined, + auth + ), + }; + }) + ); + } + + private persistLoginAuth( + selection: AccountSelection, + auth: CodexOauthAuth, + isActive: () => boolean + ): Effect.Effect> { + // Only successful authorization replaces identity. Cancelled logins retain the legacy alias. + const nextAuth = { + ...auth, + credentialId: crypto.randomUUID(), + legacyCredentialId: undefined, + invalidReason: undefined, + }; + return this.withAccountMutationEffect( + selection.accountId, + this.configMutationEffect(() => + this.providerService.updateProviderSection( + "openai", + (section) => { + if ( + !isActive() || + this.loginSelections.get(selection.accountId) !== selection || + this.getAccountRevision(selection.accountId) !== selection.revision + ) + return null; + const current = isPlainObject(section) ? section : {}; + const legacy = selection.accountId === CODEX_OAUTH_DEFAULT_ACCOUNT_ID; + // Hand-edited arrays cannot retain named properties when JSON serializes the config. + const accounts = isPlainObject(current.codexOauthAccounts) + ? current.codexOauthAccounts + : {}; + const entry = accounts[selection.accountId]; + const stored = parseCodexOauthAuth( + legacy ? current.codexOauth : isPlainObject(entry) ? entry.credentials : undefined + ); + // Token rotation preserves the login ID. Deletion, replacement, or an older writer cannot match it. + if ( + selection.credentialId === undefined + ? stored !== null + : stored?.credentialId !== selection.credentialId + ) + return null; + const next = { ...current }; + if (legacy) { + next.codexOauth = nextAuth; + } else { + // Do not mirror named credentials into the legacy slot. + // Older versions refresh that slot independently and cannot honor project selection. + // Named login preserves existing legacy credentials until explicit disconnect. + next.codexOauthAccounts = { + ...accounts, + // Legacy config readers redact credentials, including nested account identity fields. + [selection.accountId]: createNamedAccount( + selection.accountId, + isPlainObject(entry) ? entry.label : selection.label, + nextAuth + ), + }; + // Commit the first slot and its selection together. A failed write must leave neither field. + if ( + selection.selectAsDefault && + current.codexOauthDefaultAccountId === undefined && + !parseCodexOauthAuth(current.codexOauth) + ) { + next.codexOauthDefaultAccountId = selection.accountId; + } + } + return { value: next }; + }, + { enforcePolicy: true } + ) + ).pipe( + Effect.map((result) => { + if (result.success) { + this.accountRevisions.set( + selection.accountId, + this.getAccountRevision(selection.accountId) + 1 + ); + this.clearLoginSelection(selection); + } + return result; + }) + ) + ); + } + private handleDesktopCallbackAndExchange(input: { + destination: AccountSelection; + isActive: () => boolean; flowId: string; redirectUri: string; codeVerifier: string; @@ -618,7 +1090,7 @@ export class CodexOauthService { codeVerifier: input.codeVerifier, }); - const persistResult = yield* self.persistAuth(auth); + const persistResult = yield* self.persistLoginAuth(input.destination, auth, input.isActive); if (!persistResult.success) { return yield* Effect.fail(new CodexOauthError({ reason: persistResult.error })); } @@ -712,7 +1184,10 @@ export class CodexOauthService { }); } - private refreshTokens(current: CodexOauthAuth): Effect.Effect { + private refreshTokens( + selection: AccountSelection, + current: CodexOauthAuth + ): Effect.Effect { // eslint-disable-next-line @typescript-eslint/no-this-alias -- Effect.gen generator bodies do not inherit `this` const self = this; return Effect.gen(function* () { @@ -724,6 +1199,7 @@ export class CodexOauthService { method: "POST", headers: { "Content-Type": "application/x-www-form-urlencoded" }, body: buildCodexRefreshBody({ refreshToken: current.refresh }), + signal: AbortSignal.timeout(CODEX_OAUTH_REFRESH_TIMEOUT_MS), }), catch: (error) => new CodexOauthError({ @@ -734,15 +1210,14 @@ export class CodexOauthService { if (!response.ok) { const errorText = yield* Effect.promise(() => response.text().catch(() => "")); - // When the refresh token is invalid/revoked, clear persisted auth so subsequent - // requests fall back to the existing "not connected" behavior. + // Keep the credential identity so an in-progress reconnect can replace rejected tokens. if (isInvalidGrantError(errorText)) { - log.debug("[Codex OAuth] Refresh token rejected; clearing stored auth"); - const disconnectResult = yield* self.disconnectEffect(); - if (!disconnectResult.success) { - log.warn( - `[Codex OAuth] Failed to clear stored auth after refresh failure: ${disconnectResult.error}` - ); + const invalidationResult = yield* self.persistAuth(selection, { + ...current, + invalidReason: "invalid_grant", + }); + if (!invalidationResult.success) { + log.warn(`[Codex OAuth] Failed to mark rejected auth: ${invalidationResult.error}`); } } @@ -782,22 +1257,29 @@ export class CodexOauthService { ); } - const accountId = extractAccountIdFromTokens({ accessToken, idToken }) ?? current.accountId; + // Refresh cannot change the ChatGPT identity for this local slot. + const accountId = + current.accountId ?? extractAccountIdFromTokens({ accessToken, idToken }) ?? undefined; const next: CodexOauthAuth = { type: "oauth", + credentialId: current.credentialId, + legacyCredentialId: current.legacyCredentialId, access: accessToken, refresh: refreshToken ?? current.refresh, expires: Date.now() + Math.max(0, Math.floor(expiresIn * 1000)), accountId, }; - const persistResult = yield* self.persistAuth(next); + const persistResult = yield* self.persistAuth(selection, next); if (!persistResult.success) { return yield* Effect.fail(new CodexOauthError({ reason: persistResult.error })); } - return next; + const validated = self.validateRequestAuth(next, selection); + if (!validated.success) + return yield* Effect.fail(new CodexOauthError({ reason: validated.error })); + return validated.data; }).pipe( // Mirror the pre-Effect whole-body try/catch: an unexpected throw — // e.g. a rejected persistAuth/disconnect config write, which @@ -828,6 +1310,7 @@ export class CodexOauthService { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ client_id: CODEX_OAUTH_CLIENT_ID }), + signal: AbortSignal.timeout(CODEX_OAUTH_START_TIMEOUT_MS), }), catch: (error) => new CodexOauthError({ @@ -906,7 +1389,11 @@ export class CodexOauthService { const attempt = yield* self.pollDeviceTokenOnce(flow); if (attempt.kind === "success") { - const persistResult = yield* self.persistAuth(attempt.auth); + const persistResult = yield* self.persistLoginAuth( + flow.destination, + attempt.auth, + () => !flow.settled + ); if (!persistResult.success) { yield* self.finishDeviceFlowEffect(flowId, Err(persistResult.error)); return; @@ -1019,6 +1506,13 @@ export class CodexOauthService { ); } + private clearLoginSelection(destination: AccountSelection): void { + // An older flow must not release a newer login for the same slot. + if (this.loginSelections.get(destination.accountId) === destination) { + this.loginSelections.delete(destination.accountId); + } + } + /** Idempotent device-flow finish: all-sync bookkeeping + deferred resolve. */ private finishDeviceFlowEffect( flowId: string, @@ -1031,6 +1525,7 @@ export class CodexOauthService { } flow.settled = true; + this.clearLoginSelection(flow.destination); clearTimeout(flow.timeout); flow.abortController.abort(); diff --git a/src/node/services/compactionMonitor.test.ts b/src/node/services/compactionMonitor.test.ts index 10ed4f1694d..400ec2f9174 100644 --- a/src/node/services/compactionMonitor.test.ts +++ b/src/node/services/compactionMonitor.test.ts @@ -147,6 +147,34 @@ describe("CompactionMonitor", () => { ).toBe(false); }); + test("project account selection caps both pre-send and mid-stream pressure", () => { + const providersConfig: ProvidersConfigMap = { + openai: { apiKeySet: true, isEnabled: true, isConfigured: true }, + }; + const model = "openai:gpt-5.5"; + const params = { model, providersConfig, use1MContext: false }; + const beforeSend = { + ...params, + usage: { lastContextUsage: createUsageDisplay(260_000, model) }, + }; + const midStream = { ...params, usage: createMidStreamUsage(260_000) }; + const { monitor } = createMonitor(); + expect(monitor.checkBeforeSend(beforeSend).shouldForceCompact).toBe(false); + expect(monitor.checkMidStream(midStream)).toBe(false); + // A missing project selection must retain the OAuth cap instead of falling back to API billing. + expect( + monitor.checkBeforeSend({ ...beforeSend, codexOauthAccountId: "missing" }).shouldForceCompact + ).toBe(true); + expect(monitor.checkMidStream({ ...midStream, codexOauthAccountId: "missing" })).toBe(true); + + providersConfig.openai.codexOauthDefaultAuth = "apiKey"; + monitor.resetForNewStream(); + expect( + monitor.checkBeforeSend({ ...beforeSend, codexOauthAccountId: "missing" }).shouldForceCompact + ).toBe(false); + expect(monitor.checkMidStream({ ...midStream, codexOauthAccountId: "missing" })).toBe(false); + }); + test("checkMidStream stays disabled when threshold is set to 1.0", () => { const { monitor, statusEvents } = createMonitor(); monitor.setThreshold(1); diff --git a/src/node/services/compactionMonitor.ts b/src/node/services/compactionMonitor.ts index 10092b410f5..8efb9b8af6e 100644 --- a/src/node/services/compactionMonitor.ts +++ b/src/node/services/compactionMonitor.ts @@ -11,7 +11,7 @@ import { type AutoCompactionUsageState, } from "@/common/utils/compaction/autoCompactionCheck"; import { getEffectiveContextLimit } from "@/common/utils/compaction/contextLimit"; -import type { OpenAIWireFormat } from "@/common/types/providerOptions"; +import type { CodexOauthRoutingOptions } from "@/common/utils/providers/codexOauthRouting"; export type CompactionStatusEvent = | { @@ -24,22 +24,18 @@ export type CompactionStatusEvent = newUsagePercent: number; }; -interface CheckBeforeSendParams { +interface CheckBeforeSendParams extends CodexOauthRoutingOptions { model: string | null; usage: AutoCompactionUsageState | undefined; use1MContext: boolean; providersConfig: ProvidersConfigMap | null; - /** Request-level OpenAI wire format; decides whether the Codex OAuth cap applies. */ - openaiWireFormat?: OpenAIWireFormat | null; } -interface CheckMidStreamParams { +interface CheckMidStreamParams extends CodexOauthRoutingOptions { model: string; usage: LanguageModelV2Usage; use1MContext: boolean; providersConfig: ProvidersConfigMap | null; - /** Request-level OpenAI wire format; decides whether the Codex OAuth cap applies. */ - openaiWireFormat?: OpenAIWireFormat | null; } /** @@ -77,7 +73,7 @@ export class CompactionMonitor { this.threshold, undefined, params.providersConfig, - { openaiWireFormat: params.openaiWireFormat } + { openaiWireFormat: params.openaiWireFormat, codexOauthAccountId: params.codexOauthAccountId } ); } @@ -108,7 +104,7 @@ export class CompactionMonitor { params.model, params.use1MContext, params.providersConfig, - { openaiWireFormat: params.openaiWireFormat } + { openaiWireFormat: params.openaiWireFormat, codexOauthAccountId: params.codexOauthAccountId } ); // Defensive: malformed provider overrides can yield invalid/non-positive limits. // Treat those as "no compaction signal" instead of throwing inside usage-delta handlers. diff --git a/src/node/services/continuousCompactionSummary.ts b/src/node/services/continuousCompactionSummary.ts index d18a08e47b7..86767183405 100644 --- a/src/node/services/continuousCompactionSummary.ts +++ b/src/node/services/continuousCompactionSummary.ts @@ -20,6 +20,7 @@ import { import { enforceThinkingPolicy } from "@/common/utils/thinking/policy"; import { buildCompactionMessageText } from "@/common/utils/compaction/compactionPrompt"; import { getEffectiveContextLimit } from "@/common/utils/compaction/contextLimit"; +import type { ModelRoutingSnapshot } from "./modelRoutingSnapshot"; import { estimateMuxMessageTokens } from "@/common/utils/messages/keepRecentTail"; import { SUMMARIZER_INPUT_FRACTION } from "@/constants/continuousCompaction"; import type { Config } from "@/node/config"; @@ -50,21 +51,40 @@ export async function summarizeContinuousCompaction(args: { context: ContinuousCompactionContext; baseOptions: SendMessageOptions; compactOptions: SendMessageOptions; + modelRoutingSnapshot?: ModelRoutingSnapshot; }): Promise<{ text: string; model: string } | null> { args.signal.throwIfAborted(); - const providersConfig = args.aiService.getProvidersConfig(); + const modelRoutingSnapshot = + args.modelRoutingSnapshot ?? args.aiService.captureModelRoutingSnapshot(args.workspaceId); + const providersConfig = modelRoutingSnapshot.metadata; let options = args.compactOptions; + const codexOauthAccountId = modelRoutingSnapshot.codexOauthSelection.explicit + ? modelRoutingSnapshot.codexOauthSelection.accountId + : undefined; const compactLimit = getEffectiveContextLimit( options.model, isAnthropic1MEffectivelyEnabled(options.model, options.providerOptions, providersConfig), - providersConfig + providersConfig, + // Headless model construction uses stored wire format, not request-level options. + { codexOauthAccountId } ); const headTokens = args.head.reduce((total, row) => total + estimateMuxMessageTokens(row), 0); if (!compactLimit || headTokens > compactLimit * SUMMARIZER_INPUT_FRACTION) { // Do not truncate the head to fit a cheaper compact model: use the active // model's configured route, or leave the old compaction safety net in charge. options = args.baseOptions; - if (headTokens > args.context.contextWindowTokens * SUMMARIZER_INPUT_FRACTION) return null; + const baseLimit = getEffectiveContextLimit( + options.model, + isAnthropic1MEffectivelyEnabled(options.model, options.providerOptions, providersConfig), + providersConfig, + { codexOauthAccountId } + ); + // The active turn may use Chat Completions, while headless construction uses stored wire format. + if ( + !baseLimit || + headTokens > Math.min(baseLimit, args.context.contextWindowTokens) * SUMMARIZER_INPUT_FRACTION + ) + return null; } const modelString = options.model; const thinkingLevel = enforceThinkingPolicy( @@ -118,6 +138,7 @@ export async function summarizeContinuousCompaction(args: { const created = await args.aiService.createModelWithPinnedMetadata(modelString, { workspaceId: args.workspaceId, agentInitiated: true, + modelRoutingSnapshot, }); if (!created.success) throw new Error(`Cannot create compact model: ${created.error.type}`); try { diff --git a/src/node/services/modelRoutingSnapshot.ts b/src/node/services/modelRoutingSnapshot.ts new file mode 100644 index 00000000000..5b6438896c5 --- /dev/null +++ b/src/node/services/modelRoutingSnapshot.ts @@ -0,0 +1,15 @@ +import type { ProvidersConfig } from "@/common/config/schemas/providersConfig"; +import type { ProvidersConfigMap } from "@/common/orpc/types"; +import type { ProjectsConfig } from "@/common/types/project"; + +export type RouteConfigSnapshot = Required< + Pick +>; + +/** Internal turn state. Never persist this snapshot or send its credentials to the renderer. */ +export interface ModelRoutingSnapshot { + readonly providersConfig: ProvidersConfig; + readonly routeConfig: RouteConfigSnapshot; + readonly metadata: ProvidersConfigMap | null; + readonly codexOauthSelection: { accountId: string; explicit: boolean }; +} diff --git a/src/node/services/projectService.test.ts b/src/node/services/projectService.test.ts index 167fe14dc6a..31c05fd4b6a 100644 --- a/src/node/services/projectService.test.ts +++ b/src/node/services/projectService.test.ts @@ -2280,6 +2280,64 @@ exit 1 }); describe("project settings mutations", () => { + it("persists independent Codex accounts and clears only the selected override", async () => { + const firstPath = path.join(tempDir, "first"); + const secondPath = path.join(tempDir, "second"); + await config.editConfig((current) => { + current.projects.set(firstPath, { workspaces: [] }); + current.projects.set(secondPath, { workspaces: [] }); + return current; + }); + + const results = await Promise.all([ + service.setCodexOauthAccount(firstPath + "/", "personal"), + service.setCodexOauthAccount(secondPath, "work"), + service.setDisplayName(firstPath, "First project"), + ]); + expect(results[0]).toEqual(Ok(undefined)); + expect(results[1]).toEqual(Ok(undefined)); + const reloaded = new Config(tempDir).loadConfigOrDefault(); + expect(reloaded.projects.get(firstPath)).toMatchObject({ + codexOauthAccountId: "personal", + displayName: "First project", + }); + expect(reloaded.projects.get(secondPath)?.codexOauthAccountId).toBe("work"); + + expect(await service.setCodexOauthAccount(firstPath, null)).toEqual(Ok(undefined)); + const afterClear = config.loadConfigOrDefault(); + expect(afterClear.projects.get(firstPath)?.codexOauthAccountId).toBeUndefined(); + expect(afterClear.projects.get(secondPath)?.codexOauthAccountId).toBe("work"); + }); + + it.each(["work", null])( + "rejects an unpersisted Codex account selection: %s", + async (accountId) => { + const projectPath = path.join(tempDir, "write-failure"); + await config.editConfig((current) => { + current.projects.set(projectPath, { workspaces: [], codexOauthAccountId: "personal" }); + return current; + }); + // Config can complete the transform without persisting its result. + spyOn(config, "editConfig").mockImplementationOnce((transform) => { + transform(config.loadConfigOrDefault()); + return Promise.resolve(); + }); + + const result = await service.setCodexOauthAccount(projectPath, accountId); + expect(result.success).toBe(false); + expect(config.loadConfigOrDefault().projects.get(projectPath)?.codexOauthAccountId).toBe( + "personal" + ); + } + ); + + it("rejects a Codex account override for an unknown project", async () => { + const missing = path.join(tempDir, "missing"); + const result = await service.setCodexOauthAccount(missing, "work"); + expect(result.success).toBe(false); + expect(config.loadConfigOrDefault().projects.has(missing)).toBe(false); + }); + it("propagates parent trust to MCP state", async () => { const parentPath = path.join(tempDir, "project"); const childPath = path.join(parentPath, "packages", "api"); diff --git a/src/node/services/projectService.ts b/src/node/services/projectService.ts index 28a81f401ea..d6a5a297750 100644 --- a/src/node/services/projectService.ts +++ b/src/node/services/projectService.ts @@ -1995,6 +1995,32 @@ export class ProjectService { }); } + /** Select a Codex account, or inherit the global default with null. */ + async setCodexOauthAccount(projectPath: string, accountId: string | null): Promise> { + const normalizedPath = stripTrailingSlashes(projectPath); + try { + // Edit fresh state so concurrent project settings remain intact. + await this.config.editConfig((config) => { + const project = config.projects.get(normalizedPath); + if (!project) throw new Error(`Project not found: ${normalizedPath}`); + if (accountId === null) { + delete project.codexOauthAccountId; + } else { + project.codexOauthAccountId = accountId; + } + return config; + }); + // Config swallows write failures. Verify the billing selection before reporting success. + const persisted = this.config.loadConfigOrDefault().projects.get(normalizedPath); + if (!persisted || persisted.codexOauthAccountId !== (accountId ?? undefined)) { + return Err(`Failed to persist Codex account selection for ${normalizedPath}`); + } + return Ok(undefined); + } catch (error) { + return Err(`Failed to set Codex account for ${normalizedPath}: ${getErrorMessage(error)}`); + } + } + async setCustomInstructions( projectPath: string, customInstructions: string | null | undefined diff --git a/src/node/services/providerModelFactory.test.ts b/src/node/services/providerModelFactory.test.ts index 9eb498cb9c6..aaa55561c19 100644 --- a/src/node/services/providerModelFactory.test.ts +++ b/src/node/services/providerModelFactory.test.ts @@ -1062,6 +1062,434 @@ describe("ProviderModelFactory GitHub Copilot", () => { }); }); + it.each(["preference", "wire-format"])( + "uses pre-send routing config after a %s change", + async (change) => { + await withTempConfig(async (config, factory, oauth, store) => { + const auth = { + type: "oauth" as const, + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + }; + const requests: Array<{ url: string; authorization: string | null }> = []; + const providersConfig = { + openai: { + apiKey: "test-api-key", + codexOauthAccounts: { work: { label: "Work", credentials: auth } }, + }, + }; + const fetchStub = Object.assign( + (input: RequestInfo | URL, init?: RequestInit) => { + requests.push({ + url: input instanceof Request ? input.url : String(input), + authorization: new Headers(init?.headers).get("authorization"), + }); + return Promise.reject(new Error("Request captured")); + }, + { preconnect: () => undefined } + ); + const fetchSpy = spyOn(globalThis, "fetch").mockImplementation(fetchStub); + const codexOauthService = new CodexOauthService( + store, + new ProviderService(config, undefined, store) + ); + oauth.codexOauthService = codexOauthService; + try { + store.saveProvidersConfig(providersConfig); + await config.editConfig((cfg) => { + cfg.projects.set("/project", { + codexOauthAccountId: "work", + workspaces: [{ id: "snapshot-ws", name: "snapshot-ws", path: "/project/ws" }], + }); + return cfg; + }); + const options = { + workspaceId: "snapshot-ws", + providersConfig, + codexOauthSelection: { accountId: "work", explicit: true }, + }; + store.saveProvidersConfig({ + openai: { + ...providersConfig.openai, + ...(change === "preference" + ? { codexOauthDefaultAuth: "apiKey" } + : { wireFormat: "chatCompletions" }), + }, + }); + await config.editConfig((cfg) => { + delete cfg.projects.get("/project")!.codexOauthAccountId; + return cfg; + }); + const pinned = await factory.resolveAndCreateModel( + "openai:gpt-5.5", + "off", + undefined, + options + ); + expect(pinned.success).toBe(true); + if (!pinned.success) return; + expect(pinned.data.codexOauthAccountId).toBe("work"); + expect(pinned.data.model).toMatchObject({ provider: "openai.responses" }); + const next = await factory.resolveAndCreateModel("openai:gpt-5.5", "off", undefined, { + workspaceId: "snapshot-ws", + }); + expect(next.success).toBe(true); + if (!next.success) return; + // Verify authentication at the request boundary, not through cost markers. + const cases = [ + { + model: pinned.data.model, + expected: { + url: "https://chatgpt.com/backend-api/codex/responses", + authorization: "Bearer access", + }, + }, + { + model: next.data.model, + expected: { + url: `https://api.openai.com/v1/${change === "wire-format" ? "chat/completions" : "responses"}`, + authorization: "Bearer test-api-key", + }, + }, + ]; + for (const { model, expected } of cases) { + requests.length = 0; + // eslint-disable-next-line @typescript-eslint/await-thenable -- bun-types mistype .rejects.toThrow as void. + await expect(generateText({ model, prompt: "Hello", maxRetries: 0 })).rejects.toThrow( + "Request captured" + ); + expect(requests.length).toBeGreaterThan(0); + for (const request of requests) expect(request).toEqual(expected); + } + } finally { + fetchSpy.mockRestore(); + await codexOauthService.dispose(); + } + }); + } + ); + + it("pins project Codex accounts across requests and rejects deleted slots", async () => { + await withTempConfig(async (config, factory, oauth, providersConfigStore) => { + const personal = { + type: "oauth" as const, + access: "personal-access", + refresh: "personal-refresh", + expires: Date.now() + 3_600_000, + accountId: "chatgpt-personal", + }; + const work = { + ...personal, + credentialId: "1c9c50b0-d777-4dd2-998c-09c156ba9754", + access: "work-access", + accountId: "chatgpt-work", + }; + providersConfigStore.saveProvidersConfig({ + openai: { + apiKey: "api-key-must-not-win", + codexOauth: personal, + codexOauthAccounts: { work: { label: "Work", credentials: work } }, + }, + }); + await config.editConfig((current) => { + current.projects.set("/personal", { + codexOauthAccountId: "default", + workspaces: [{ id: "personal-ws", name: "personal", path: "/personal/ws" }], + }); + current.projects.set("/work", { + codexOauthAccountId: "work", + workspaces: [{ id: "work-ws", name: "work", path: "/work/ws" }], + }); + current.projects.set("/personal/sub", { + parentProjectPath: "/personal", + codexOauthAccountId: "work", + workspaces: [], + }); + current.projects.get("/personal")!.workspaces.push({ + id: "sub-ws", + name: "sub", + path: "/personal/sub-ws", + subProjectPath: "/personal/sub", + }); + current.projects.set("_multi", { + workspaces: [ + { + id: "multi-ws", + name: "multi", + path: "/multi/ws", + projects: [ + { projectPath: "/work", projectName: "work" }, + { projectPath: "/personal", projectName: "personal" }, + ], + }, + ], + }); + return current; + }); + oauth.codexOauthService = new CodexOauthService( + providersConfigStore, + new ProviderService(config, undefined, providersConfigStore) + ); + const originalRegistry = PROVIDER_REGISTRY.openai; + const requests: Headers[] = []; + let capturedFetch: typeof fetch | undefined; + PROVIDER_REGISTRY.openai = async () => { + const module = await originalRegistry(); + return { + ...module, + createOpenAI: (options) => { + capturedFetch = options?.fetch; + return module.createOpenAI(options); + }, + }; + }; + const createFetch = async (context: { workspaceId?: string; projectPath?: string } = {}) => { + const stored = providersConfigStore.loadProvidersConfig() ?? {}; + const result = await factory.createModel("openai:gpt-5.3-codex", undefined, { + ...context, + providersConfig: { + ...stored, + openai: { + ...stored.openai, + fetch: (_input: RequestInfo | URL, init?: RequestInit) => { + requests.push(new Headers(init?.headers)); + return Promise.resolve( + new Response("{}", { headers: { "content-type": "application/json" } }) + ); + }, + }, + }, + }); + expect(result.success).toBe(true); + if (!result.success || !capturedFetch) throw new Error("Expected an OAuth model"); + expect(result.data).toMatchObject({ provider: "openai.responses" }); + return capturedFetch; + }; + const send = (providerFetch: typeof fetch) => + providerFetch("https://api.openai.com/v1/responses", { + method: "POST", + headers: { "ChatGPT-Account-Id": "stale-configured-account" }, + body: JSON.stringify({ model: "gpt-5.3-codex", input: [] }), + }); + try { + const personalFetch = await createFetch({ workspaceId: "personal-ws" }); + const workFetch = await createFetch({ workspaceId: "work-ws" }); + await send(personalFetch); + await send(workFetch); + expect(requests.map((headers) => headers.get("authorization"))).toEqual([ + "Bearer personal-access", + "Bearer work-access", + ]); + expect(requests.map((headers) => headers.get("chatgpt-account-id"))).toEqual([ + "chatgpt-personal", + "chatgpt-work", + ]); + for (const context of [ + { workspaceId: "sub-ws" }, + { workspaceId: "multi-ws" }, + { projectPath: "/work" }, + ]) { + await send(await createFetch(context)); + expect(requests.at(-1)?.get("authorization")).toBe("Bearer work-access"); + } + // Legacy requests use the default slot when no selection exists. + await send(await createFetch()); + expect(requests.at(-1)?.get("authorization")).toBe("Bearer personal-access"); + + const updatedWork = { ...work, access: "work-refreshed", accountId: undefined }; + providersConfigStore.saveProvidersConfig({ + openai: { + codexOauth: personal, + codexOauthDefaultAccountId: "work", + codexOauthAccounts: { work: { label: "Work", credentials: updatedWork } }, + }, + }); + await config.editConfig((current) => { + current.projects.get("/work")!.codexOauthAccountId = "default"; + return current; + }); + await send(workFetch); + expect(requests.at(-1)?.get("authorization")).toBe("Bearer work-refreshed"); + expect(requests.at(-1)?.has("chatgpt-account-id")).toBe(false); + await send(personalFetch); + expect(requests.at(-1)?.get("authorization")).toBe("Bearer personal-access"); + await send(await createFetch()); + expect(requests.at(-1)?.get("authorization")).toBe("Bearer work-refreshed"); + await config.editConfig((current) => { + delete current.projects.get("/personal/sub")!.codexOauthAccountId; + return current; + }); + await send(await createFetch({ workspaceId: "sub-ws" })); + expect(requests.at(-1)?.get("authorization")).toBe("Bearer work-refreshed"); + const resolved = await factory.resolveAndCreateModel( + "openai:gpt-5.3-codex", + "off", + undefined, + { + workspaceId: "personal-ws", + } + ); + expect(resolved.success).toBe(true); + if (resolved.success) expect(resolved.data.codexOauthAccountId).toBe("default"); + + providersConfigStore.saveProvidersConfig({ + openai: { + apiKey: "api-key-must-not-win", + codexOauth: personal, + codexOauthDefaultAccountId: "work", + }, + }); + const sentCount = requests.length; + const rejected = await send(workFetch).then( + () => false, + () => true + ); + expect(rejected).toBe(true); + expect(requests).toHaveLength(sentCount); + expect(await factory.createModel("openai:gpt-5.3-codex")).toMatchObject({ + success: false, + error: { type: "oauth_not_connected" }, + }); + } finally { + PROVIDER_REGISTRY.openai = originalRegistry; + } + }); + }); + + for (const credentialId of [undefined, "1c9c50b0-d777-4dd2-998c-09c156ba9754"]) { + for (const beforeFirstFetch of [true, false]) { + it( + "rejects replaced " + + (credentialId ? "identified" : "legacy") + + " Codex credentials " + + (beforeFirstFetch ? "before first fetch" : "between fetches"), + async () => { + await withTempConfig(async (config, factory, oauth, store) => { + const auth = { + type: "oauth" as const, + credentialId, + access: "original-access", + refresh: "original-refresh", + expires: Date.now() + 3_600_000, + }; + const saveAuth = (next: typeof auth) => + store.saveProvidersConfig({ + openai: { codexOauth: next }, + }); + saveAuth(auth); + oauth.codexOauthService = new CodexOauthService( + store, + new ProviderService(config, undefined, store) + ); + const originalRegistry = PROVIDER_REGISTRY.openai; + let providerFetch: typeof fetch | undefined; + const sentTokens: Array = []; + PROVIDER_REGISTRY.openai = async () => { + const module = await originalRegistry(); + return { + ...module, + createOpenAI: (options) => { + providerFetch = options?.fetch; + return module.createOpenAI(options); + }, + }; + }; + try { + const model = await factory.createModel("openai:gpt-5.3-codex", undefined, { + providersConfig: { + openai: { + codexOauth: auth, + fetch: (_input: RequestInfo | URL, init?: RequestInit) => { + sentTokens.push(new Headers(init?.headers).get("authorization")); + return Promise.resolve(new Response("{}")); + }, + }, + }, + }); + expect(model.success).toBe(true); + if (!providerFetch) throw new Error("Expected an OAuth fetch wrapper"); + const send = () => + providerFetch!("https://api.openai.com/v1/responses", { + method: "POST", + body: JSON.stringify({ input: [] }), + }); + if (!beforeFirstFetch) { + await send(); + expect(sentTokens).toEqual(["Bearer original-access"]); + } + saveAuth({ + ...auth, + credentialId: "6fb7157c-c5a4-4ea7-852f-c46d0b090ff5", + access: "replacement-access", + }); + const count = sentTokens.length; + // eslint-disable-next-line @typescript-eslint/await-thenable -- bun-types mistype .rejects.toThrow as void. + await expect(send()).rejects.toThrow("account changed"); + expect(sentTokens).toHaveLength(count); + } finally { + PROVIDER_REGISTRY.openai = originalRegistry; + await oauth.codexOauthService.dispose(); + } + }); + } + ); + } + } + + it("rejects explicit missing Codex selections without changing API-key precedence", async () => { + await withTempConfig(async (config, factory, _oauth, providersConfigStore) => { + await config.editConfig((current) => { + current.projects.set("/missing", { codexOauthAccountId: "deleted", workspaces: [] }); + return current; + }); + for (const selection of [undefined, { projectPath: "/missing" }]) { + providersConfigStore.saveProvidersConfig({ + openai: { + apiKey: "test-key", + codexOauthDefaultAuth: "oauth", + ...(selection ? {} : { codexOauthDefaultAccountId: "deleted" }), + }, + }); + expect( + await factory.createModel("openai:gpt-5.3-codex", undefined, selection) + ).toMatchObject({ + success: false, + error: { type: "oauth_not_connected" }, + }); + const chat = await factory.createModel( + "openai:gpt-5.3-codex", + { + openai: { wireFormat: "chatCompletions" }, + }, + selection + ); + expect(chat.success).toBe(true); + if (chat.success) expect(chat.data).toMatchObject({ provider: "openai.chat" }); + } + providersConfigStore.saveProvidersConfig({ + openai: { apiKey: "test-key", codexOauthDefaultAuth: "apiKey" }, + }); + const api = await factory.createModel("openai:gpt-5.5", undefined, { + projectPath: "/missing", + }); + expect(api.success).toBe(true); + if (api.success) expect(api.data).toMatchObject({ provider: "openai.responses" }); + providersConfigStore.saveProvidersConfig({ + openai: { codexOauthDefaultAccountId: "deleted" }, + openrouter: { apiKey: "gateway-key" }, + }); + await saveRoutePriority(config, ["openrouter", "direct"]); + const gateway = await factory.resolveAndCreateModel("openai:gpt-5.5", "off", undefined, { + projectPath: "/missing", + }); + expectSuccessfulRouteResult(gateway, { + effectiveModelString: "openrouter:openai/gpt-5.5", + routeProvider: "openrouter", + }); + }); + }); + it("normalizes Request bodies for the Codex OAuth responses endpoint", async () => { await withTempConfig(async (_config, factory, oauth, providersConfigStore) => { const originalOpenAIRegistry = PROVIDER_REGISTRY.openai; @@ -1288,6 +1716,88 @@ describe("ProviderModelFactory GitHub Copilot", () => { }); }); +describe("ProviderModelFactory route config snapshots", () => { + it.each(["priority", "override"] as const)( + "keeps all resolution stages on captured %s settings", + async (change) => { + await withTempConfig(async (config, factory, _oauth, store) => { + const providersConfig = { + openai: { apiKey: "openai-key" }, + anthropic: { apiKey: "anthropic-key" }, + openrouter: { apiKey: "openrouter-key" }, + }; + store.saveProvidersConfig(providersConfig); + const routeConfig = { + routePriority: ["direct"], + routeOverrides: { "anthropic:claude-sonnet-4-5": "openrouter" }, + }; + await saveRoutePriority( + config, + change === "priority" ? ["openrouter", "direct"] : ["direct"], + { + routeOverrides: change === "override" ? { "openai:gpt-5.5": "openrouter" } : {}, + } + ); + for (const modelString of ["openai:gpt-5.5", "coder:openai/gpt-5.5"]) { + const resolved = await factory.resolveAndCreateModel(modelString, "off", undefined, { + providersConfig, + routeConfig, + }); + expectSuccessfulRouteResult(resolved, { + effectiveModelString: "openai:gpt-5.5", + routeProvider: "openai", + }); + if (!resolved.success) throw new Error("Expected a model"); + expect((resolved.data.model as { provider?: unknown }).provider).toBe("openai.responses"); + const direct = await factory.createModel(modelString, undefined, { + providersConfig, + routeConfig, + }); + expect(direct.success).toBe(true); + if (direct.success) + expect((direct.data as { provider?: unknown }).provider).toBe("openai.responses"); + expect( + factory.resolveEffectiveModelString( + modelString, + undefined, + providersConfig, + routeConfig + ) + ).toBe("openai:gpt-5.5"); + expect( + factory.resolveGatewayModelString( + modelString, + undefined, + undefined, + providersConfig, + routeConfig + ) + ).toBe("openai:gpt-5.5"); + const current = await factory.resolveAndCreateModel(modelString, "off"); + expectSuccessfulRouteResult(current, { + effectiveModelString: "openrouter:openai/gpt-5.5", + routeProvider: "openrouter", + }); + } + // Nested models resolve their own overrides, not a single route chosen for the parent. + const nested = await factory.resolveAndCreateModel( + "anthropic:claude-sonnet-4-5", + "off", + undefined, + { + providersConfig, + routeConfig, + } + ); + expectSuccessfulRouteResult(nested, { + effectiveModelString: "openrouter:anthropic/claude-sonnet-4-5", + routeProvider: "openrouter", + }); + }); + } + ); +}); + describe("ProviderModelFactory OpenAI WebSocket transport", () => { it("attaches cleanup when enabled for Responses models", async () => { await withOpenAIBaseUrlEnvUnset(async () => @@ -1900,45 +2410,172 @@ describe("ProviderModelFactory routing", () => { }); }); - it("treats OpenAI as available for routing when only Codex OAuth is configured", async () => { - // Temporarily remove OPENAI_API_KEY so the test only succeeds via Codex OAuth, - // not by falling through to an env-var credential path. - const savedKey = process.env.OPENAI_API_KEY; - delete process.env.OPENAI_API_KEY; - try { - await withTempConfig(async (config, factory) => { - new ProvidersConfigStore(config.rootDir).saveProvidersConfig({ - openai: { - // No apiKey — only Codex OAuth credentials. - codexOauth: { - type: "oauth", - access: "test-access-token", - refresh: "test-refresh-token", - expires: Date.now() + 60_000, + it.each(["default", "work"])( + "treats OpenAI as available with only the %s OAuth slot", + async (accountId) => { + // Temporarily remove OPENAI_API_KEY so the test only succeeds via Codex OAuth, + // not by falling through to an env-var credential path. + const savedKey = process.env.OPENAI_API_KEY; + delete process.env.OPENAI_API_KEY; + try { + await withTempConfig(async (config, factory) => { + new ProvidersConfigStore(config.rootDir).saveProvidersConfig({ + openai: { + // No apiKey — only Codex OAuth credentials. + ...(accountId === "default" + ? { + codexOauth: { + type: "oauth" as const, + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + }, + } + : { + codexOauthDefaultAccountId: accountId, + codexOauthAccounts: { + [accountId]: { + label: "Work", + credentials: { + type: "oauth" as const, + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + }, + }, + }, + }), }, - }, - openrouter: { - apiKey: "or-test", - }, + openrouter: { + apiKey: "or-test", + }, + }); + + await saveRoutePriority(config, ["direct", "openrouter"]); + + // Direct OpenAI should win because Codex OAuth makes it available for routing. + // Use a model from CODEX_OAUTH_ALLOWED_MODELS so createModel can route through OAuth. + const result = await factory.resolveAndCreateModel("openai:gpt-5.2", "off"); + expectSuccessfulRouteResult(result, { + effectiveModelString: "openai:gpt-5.2", + routeProvider: "openai", + routedThroughGateway: false, + }); }); + } finally { + if (savedKey !== undefined) { + process.env.OPENAI_API_KEY = savedKey; + } + } + } + ); - await saveRoutePriority(config, ["direct", "openrouter"]); + it.each(["default", "work"])( + "skips a revoked %s OAuth slot when a gateway is configured", + async (accountId) => { + const savedKey = process.env.OPENAI_API_KEY; + delete process.env.OPENAI_API_KEY; + try { + await withTempConfig(async (config, factory) => { + const auth = { + type: "oauth" as const, + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + invalidReason: "invalid_grant" as const, + }; + new ProvidersConfigStore(config.rootDir).saveProvidersConfig({ + openai: { + ...(accountId === "default" + ? { codexOauth: auth } + : { codexOauthAccounts: { work: { label: "Work", credentials: auth } } }), + codexOauthDefaultAccountId: accountId, + }, + openrouter: { apiKey: "or-test" }, + }); - // Direct OpenAI should win because Codex OAuth makes it available for routing. - // Use a model from CODEX_OAUTH_ALLOWED_MODELS so createModel can route through OAuth. - const result = await factory.resolveAndCreateModel("openai:gpt-5.2", "off"); - expectSuccessfulRouteResult(result, { - effectiveModelString: "openai:gpt-5.2", - routeProvider: "openai", - routedThroughGateway: false, + for (const routeOverrides of [{}, { "openai:gpt-5.2": "direct" }]) { + await saveRoutePriority(config, ["direct", "openrouter"], { routeOverrides }); + const result = await factory.resolveAndCreateModel("openai:gpt-5.2", "off"); + expectSuccessfulRouteResult(result, { + effectiveModelString: "openrouter:openai/gpt-5.2", + routeProvider: "openrouter", + }); + } }); - }); - } finally { - if (savedKey !== undefined) { - process.env.OPENAI_API_KEY = savedKey; + } finally { + if (savedKey !== undefined) process.env.OPENAI_API_KEY = savedKey; } } - }); + ); + + it.each(["default", "work"])( + "uses a gateway for models outside the %s OAuth slot's model support", + async (accountId) => { + const savedKey = process.env.OPENAI_API_KEY; + delete process.env.OPENAI_API_KEY; + try { + await withTempConfig(async (config, factory) => { + const auth = { + type: "oauth" as const, + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + }; + const openai = { + ...(accountId === "default" + ? { codexOauth: auth } + : { codexOauthAccounts: { work: { label: "Work", credentials: auth } } }), + codexOauthDefaultAccountId: accountId, + }; + const store = new ProvidersConfigStore(config.rootDir); + store.saveProvidersConfig({ openai, openrouter: { apiKey: "or-test" } }); + + for (const routeOverrides of [{}, { "openai:gpt-4.1": "direct" }]) { + await saveRoutePriority(config, ["direct", "openrouter"], { routeOverrides }); + expect(factory.resolveGatewayModelString("openai:gpt-4.1")).toBe( + "openrouter:openai/gpt-4.1" + ); + expect((await factory.createModel("openai:gpt-4.1")).success).toBe(true); + const result = await factory.resolveAndCreateModel("openai:gpt-4.1", "off"); + expectSuccessfulRouteResult(result, { + effectiveModelString: "openrouter:openai/gpt-4.1", + routeProvider: "openrouter", + }); + } + + store.saveProvidersConfig({ + openai: { + ...openai, + models: [{ id: "team-codex", mappedToModel: KNOWN_MODELS.GPT_53_CODEX.id }], + }, + openrouter: { apiKey: "or-test" }, + }); + expectSuccessfulRouteResult( + await factory.resolveAndCreateModel("openai:team-codex", "off"), + { + effectiveModelString: "openai:team-codex", + routeProvider: "openai", + } + ); + + store.saveProvidersConfig({ + openai: { ...openai, apiKey: "sk-test" }, + openrouter: { apiKey: "or-test" }, + }); + expectSuccessfulRouteResult( + await factory.resolveAndCreateModel("openai:gpt-4.1", "off"), + { + effectiveModelString: "openai:gpt-4.1", + routeProvider: "openai", + } + ); + }); + } finally { + if (savedKey !== undefined) process.env.OPENAI_API_KEY = savedKey; + } + } + ); it("leaves direct-provider model strings unchanged when direct routing wins", async () => { await withTempConfig(async (config, factory) => { diff --git a/src/node/services/providerModelFactory.ts b/src/node/services/providerModelFactory.ts index b9e5fb643b9..c6766fe51f1 100644 --- a/src/node/services/providerModelFactory.ts +++ b/src/node/services/providerModelFactory.ts @@ -19,12 +19,17 @@ import { isCodexOauthAllowedModel, isCodexOauthRequiredModel, } from "@/common/constants/codexOAuth"; -import { parseCodexOauthAuth } from "@/node/utils/codexOauthAuth"; +import { + getCodexOauthAccountId, + getCodexOauthAccounts, + getCodexOauthAuth, +} from "@/node/utils/codexOauthAuth"; import type { Config, ProviderConfig, ProvidersConfig } from "@/node/config"; import { ProvidersConfigStore } from "@/node/config"; import type { MuxProviderOptions } from "@/common/types/providerOptions"; import type { ServiceTier, XAIServiceTier } from "@/common/config/schemas/providersConfig"; import { resolveConfigBaseUrl } from "@/common/utils/providers/baseUrl"; +import { getCodexOauthProjectPath } from "@/common/utils/providers/codexOauthRouting"; import { isProviderDisabledInConfig } from "@/common/utils/providers/isProviderDisabled"; import { customProviderWireOrigin, @@ -45,6 +50,7 @@ import { CopilotResponsesLanguageModel } from "@/node/services/copilot/copilotRe import type { PolicyService } from "@/node/services/policyService"; import type { ProviderService } from "@/node/services/providerService"; import type { CodexOauthService } from "@/node/services/codexOauthService"; +import type { RouteConfigSnapshot } from "./modelRoutingSnapshot"; import type { CoderOauthService } from "@/node/services/coderOauthService"; import { coderAibridgeBaseUrl, @@ -1099,6 +1105,8 @@ export interface ResolveAndCreateModelResult { * retagged instance type diverging from the created fallback model. */ coderSelectedInstance?: { name: string; type: string }; + /** Local account slot pinned for request routing and context limits. */ + codexOauthAccountId?: string; /** Whether the request is being routed through the Xum gateway. */ routedThroughGateway: boolean; /** Route provider chosen by backend routing (direct provider or gateway). */ @@ -1108,7 +1116,13 @@ export interface ResolveAndCreateModelResult { interface CreateModelOptions { agentInitiated?: boolean; workspaceId?: string; + /** Project context for requests before workspace creation. */ + projectPath?: string; + /** Account selection snapshot from resolveAndCreateModel. */ + codexOauthSelection?: { accountId: string; explicit: boolean }; routeContext?: RouteContext; + /** Captured routing rules keep accepted work independent from settings changes. */ + routeConfig?: RouteConfigSnapshot; /** * Providers-config snapshot to create the model from. Passed by * resolveAndCreateModel so routing, the returned coderWire snapshot, @@ -1192,10 +1206,30 @@ export class ProviderModelFactory { return forced ? { ...providerConfig, deploymentUrl: forced } : providerConfig; } + private resolveCodexOauthSelection( + providerConfig: ProviderConfigRaw, + opts?: CreateModelOptions + ): { accountId: string; explicit: boolean } { + if (opts?.codexOauthSelection) { + return opts.codexOauthSelection; + } + const workspace = opts?.workspaceId ? this.config.findWorkspace(opts.workspaceId) : null; + const projectPath = getCodexOauthProjectPath(workspace) ?? opts?.projectPath; + const projectAccountId = projectPath + ? this.config.loadConfigOrDefault().projects.get(projectPath)?.codexOauthAccountId + : undefined; + return { + accountId: getCodexOauthAccountId(providerConfig, projectAccountId), + explicit: + projectAccountId !== undefined || providerConfig.codexOauthDefaultAccountId !== undefined, + }; + } + private isProviderAvailableForRouting( provider: ProviderName, providersConfig: ProvidersConfig, - config: ReturnType + config: ReturnType, + canonicalModel: string ): boolean { const rawProviderConfig = providersConfig[provider] ?? {}; const providerConfig = @@ -1204,11 +1238,12 @@ export class ProviderModelFactory { : rawProviderConfig; const credentials = resolveProviderCredentials(provider, providerConfig); - // OpenAI Codex OAuth is a valid credential path even without an API key; - // routing should treat it as available so direct OpenAI routes are honored. + // OAuth cannot serve every OpenAI model. Unsupported models must retain configured gateway routes. + // Rejected slots remain for reconnect, but must not hide configured gateway routes. const hasCodexOauth = provider === "openai" && - parseCodexOauthAuth((providerConfig as { codexOauth?: unknown }).codexOauth) !== null; + isCodexOauthAllowedModel(canonicalModel, providersConfig) && + getCodexOauthAccounts(providerConfig).some(({ auth }) => auth.invalidReason === undefined); if (!credentials.isConfigured && !hasCodexOauth) { return false; @@ -1284,11 +1319,7 @@ export class ProviderModelFactory { private createModelCoreEffect( modelString: string, muxProviderOptions?: MuxProviderOptions, - opts?: { - agentInitiated?: boolean; - routeContext?: RouteContext; - providersConfig?: ProvidersConfig; - } + opts?: CreateModelOptions ): Effect.Effect> { // eslint-disable-next-line @typescript-eslint/no-this-alias -- Effect.gen generator bodies do not inherit `this` const self = this; @@ -1302,7 +1333,8 @@ export class ProviderModelFactory { modelString = self.resolveEffectiveModelString( modelString, opts?.routeContext, - opts?.providersConfig + opts?.providersConfig, + opts?.routeConfig ); // Parse model string (format: "provider:model-id") @@ -1568,17 +1600,18 @@ export class ProviderModelFactory { const codexOauthAllowed = isCodexOauthAllowedModel(fullModelId, providersConfig); const codexOauthRequired = isCodexOauthRequiredModel(fullModelId, providersConfig); - const storedCodexOauth = parseCodexOauthAuth( - (providerConfig as { codexOauth?: unknown }).codexOauth - ); + // Pin the local slot, not the ChatGPT header ID, across refreshes and retries. + const selection = self.resolveCodexOauthSelection(providerConfig, opts); + const codexOauthAccountId = selection.accountId; + const storedCodexOauth = getCodexOauthAuth(providerConfig, codexOauthAccountId); + // Pin the credential before the first fetch. Retries and tool steps must reject replacement logins. + const codexOauthCredential = { credentialId: storedCodexOauth?.credentialId }; // Resolve credentials from config + env so we can decide whether to // route through Codex OAuth or fall back to API key auth. const creds = resolveProviderCredentials("openai", providerConfig); - // When a model requires Codex OAuth but the user hasn't connected it, - // fall back to their API key instead of blocking entirely. If the model - // truly only works through OAuth, OpenAI's API will return a clear error. + // Preserve the legacy API-key path when no OAuth account has been selected. if (codexOauthRequired && !storedCodexOauth && !creds.isConfigured) { return Err({ type: "oauth_not_connected", provider: providerName }); } @@ -1594,32 +1627,20 @@ export class ProviderModelFactory { (providerConfig.wireFormat as string | undefined) ?? muxProviderOptions?.openai?.wireFormat; - // Codex OAuth routing: - // - Chat Completions never routes through OAuth when an API key exists. - // - Required models route through ChatGPT OAuth when connected. - // - If OAuth is not connected, fall back to API key (if available). - // - Allowed models route through OAuth only when: - // - no API key is configured, OR - // - the user prefers OAuth when both are set. - const shouldRouteThroughCodexOauth = (() => { - if (!codexOauthAllowed || !storedCodexOauth) { - return false; - } - - if (earlyWireFormat === "chatCompletions" && creds.isConfigured) { - return false; - } - - if (codexOauthRequired) { - return true; - } - - if (!creds.isConfigured) { - return true; - } + // A missing selected slot must not switch billing to an API key or another account. + const prefersCodexOauth = + codexOauthAllowed && + !(earlyWireFormat === "chatCompletions" && creds.isConfigured) && + (codexOauthRequired || !creds.isConfigured || codexOauthDefaultAuth === "oauth"); + if ( + prefersCodexOauth && + !storedCodexOauth && + (selection.explicit || getCodexOauthAccounts(providerConfig).length > 0) + ) { + return Err({ type: "oauth_not_connected", provider: providerName }); + } - return codexOauthDefaultAuth === "oauth"; - })(); + const shouldRouteThroughCodexOauth = prefersCodexOauth && storedCodexOauth !== null; // OAuth requests use a placeholder key and override auth headers in fetch(). const resolvedApiKey = shouldRouteThroughCodexOauth ? undefined : creds.apiKey; @@ -1734,7 +1755,10 @@ export class ProviderModelFactory { throw new Error("Codex OAuth service not initialized"); } - const authResult = await codexOauthService.getValidAuth(); + const authResult = await codexOauthService.getValidAuth( + codexOauthAccountId, + codexOauthCredential + ); if (!authResult.success) { throw new Error(authResult.error); } @@ -1743,6 +1767,9 @@ export class ProviderModelFactory { headers.set("Authorization", `Bearer ${authResult.data.access}`); if (authResult.data.accountId) { headers.set("ChatGPT-Account-Id", authResult.data.accountId); + } else { + // A configured header must not select another ChatGPT account. + headers.delete("ChatGPT-Account-Id"); } nextInput = CODEX_ENDPOINT; @@ -2544,7 +2571,15 @@ export class ProviderModelFactory { modelString: string, thinkingLevel: ThinkingLevel, muxProviderOptions?: MuxProviderOptions, - opts?: { agentInitiated?: boolean; workspaceId?: string } + opts?: Pick< + CreateModelOptions, + | "agentInitiated" + | "workspaceId" + | "projectPath" + | "providersConfig" + | "codexOauthSelection" + | "routeConfig" + > ): Promise> { return Effect.runPromise( this.resolveAndCreateModelEffect(modelString, thinkingLevel, muxProviderOptions, opts) @@ -2555,7 +2590,15 @@ export class ProviderModelFactory { modelString: string, thinkingLevel: ThinkingLevel, muxProviderOptions?: MuxProviderOptions, - opts?: { agentInitiated?: boolean; workspaceId?: string } + opts?: Pick< + CreateModelOptions, + | "agentInitiated" + | "workspaceId" + | "projectPath" + | "providersConfig" + | "codexOauthSelection" + | "routeConfig" + > ): Effect.Effect> { // eslint-disable-next-line @typescript-eslint/no-this-alias -- Effect.gen generator bodies do not inherit `this` const self = this; @@ -2567,7 +2610,9 @@ export class ProviderModelFactory { // through the built-in machinery instead of the user's custom endpoint. // The equivalent guard in resolveGatewayModelString only protects callers // that pass raw strings. - const providersConfigForShadowCheck = self.providersConfigStore.loadProvidersConfig() ?? {}; + // Pre-send compaction and model construction must use the same routing snapshot. + const providersConfigForShadowCheck = + opts?.providersConfig ?? self.providersConfigStore.loadProvidersConfig() ?? {}; const [rawProviderName] = parseModelString(modelString); const rawPrefixShadowedByCustomProvider = rawProviderName.length > 0 && @@ -2644,7 +2689,8 @@ export class ProviderModelFactory { const routeContext = self.resolveModelRoute( routeSeedModelString, - providersConfigForShadowCheck + providersConfigForShadowCheck, + opts?.routeConfig ); if (rawCoderGatewayModelId != null) { const appConfig = self.config.loadConfigOrDefault(); @@ -2655,7 +2701,8 @@ export class ProviderModelFactory { const coderProviderRoutable = self.isProviderAvailableForRouting( "coder", providersConfigForShadowCheck, - appConfig + appConfig, + routeSeedModelString ); const coderModelAccessible = isGatewayModelAccessible("coder", rawCoderGatewayModelId); if (coderProviderRoutable && coderModelAccessible) { @@ -2695,7 +2742,8 @@ export class ProviderModelFactory { routeSeedModelString, routeContext, undefined, - providersConfigForShadowCheck + providersConfigForShadowCheck, + opts?.routeConfig ); } } else { @@ -2703,7 +2751,8 @@ export class ProviderModelFactory { effectiveModelString, routeContext, explicitGateway, - providersConfigForShadowCheck + providersConfigForShadowCheck, + opts?.routeConfig ); } @@ -2790,8 +2839,13 @@ export class ProviderModelFactory { } } + const codexOauthSelection = self.resolveCodexOauthSelection( + providersConfigForShadowCheck.openai ?? {}, + opts + ); const modelResult = yield* self.createModelEffect(effectiveModelString, muxProviderOptions, { ...opts, + codexOauthSelection, routeContext, // ONE config snapshot for the whole resolve+create: the wire snapshot // above and the SDK model must come from the same providers.jsonc read, @@ -2835,6 +2889,7 @@ export class ProviderModelFactory { wireProviderName, coderWire, coderSelectedInstance, + codexOauthAccountId: codexOauthSelection.accountId, routedThroughGateway, // Custom adapters are direct routes: the raw custom id is not a // ProviderName, and leaking it as routeProvider makes downstream @@ -2847,7 +2902,8 @@ export class ProviderModelFactory { private resolveModelRoute( canonicalModel: string, - providersConfigSnapshot?: ProvidersConfig + providersConfigSnapshot?: ProvidersConfig, + routeConfigSnapshot?: RouteConfigSnapshot ): RouteContext { const config = this.config.loadConfigOrDefault(); // resolveAndCreateModel passes its snapshot so route availability, @@ -2861,8 +2917,8 @@ export class ProviderModelFactory { ); return resolveRoute( canonicalModel, - config.routePriority ?? ["direct"], - config.routeOverrides ?? {}, + routeConfigSnapshot?.routePriority ?? config.routePriority ?? ["direct"], + routeConfigSnapshot?.routeOverrides ?? config.routeOverrides ?? {}, (provider) => { if (!Object.hasOwn(PROVIDER_REGISTRY, provider)) { return false; @@ -2871,7 +2927,8 @@ export class ProviderModelFactory { return this.isProviderAvailableForRouting( provider as ProviderName, providersConfig, - config + config, + canonicalModel ); }, isGatewayModelAccessible @@ -2889,14 +2946,16 @@ export class ProviderModelFactory { resolveEffectiveModelString( modelString: string, routeContext?: RouteContext, - providersConfig?: ProvidersConfig + providersConfig?: ProvidersConfig, + routeConfig?: RouteConfigSnapshot ): string { const explicitGateway = getExplicitGatewayProvider(modelString); return this.resolveGatewayModelString( modelString, routeContext, explicitGateway, - providersConfig + providersConfig, + routeConfig ); } @@ -2904,7 +2963,8 @@ export class ProviderModelFactory { modelString: string, modelKeyOrRouteContext?: string | RouteContext, explicitGatewayOrLegacyFlag?: ProviderName | boolean, - providersConfigSnapshot?: ProvidersConfig + providersConfigSnapshot?: ProvidersConfig, + routeConfigSnapshot?: RouteConfigSnapshot ): string { // Legacy callers may still pass boolean true to mean an explicit mux-gateway request. const explicitGateway: ProviderName | undefined = @@ -2958,15 +3018,17 @@ export class ProviderModelFactory { providersConfig, this.policyService ); + const routingModel = + typeof modelKeyOrRouteContext === "string" + ? normalizeToCanonical(modelKeyOrRouteContext) + : canonicalModelString; const routeContext = typeof modelKeyOrRouteContext === "object" && modelKeyOrRouteContext != null ? modelKeyOrRouteContext : resolveRoute( - typeof modelKeyOrRouteContext === "string" - ? normalizeToCanonical(modelKeyOrRouteContext) - : canonicalModelString, - config.routePriority ?? ["direct"], - config.routeOverrides ?? {}, + routingModel, + routeConfigSnapshot?.routePriority ?? config.routePriority ?? ["direct"], + routeConfigSnapshot?.routeOverrides ?? config.routeOverrides ?? {}, (provider) => { if (!Object.hasOwn(PROVIDER_REGISTRY, provider)) { return false; @@ -2975,7 +3037,8 @@ export class ProviderModelFactory { return this.isProviderAvailableForRouting( provider as ProviderName, providersConfig, - config + config, + routingModel ); }, isGatewayModelAccessible @@ -2988,7 +3051,12 @@ export class ProviderModelFactory { // gateway selections from being silently rewritten after canonicalization. if ( explicitGateway != null && - this.isProviderAvailableForRouting(explicitGateway, providersConfig, config) + this.isProviderAvailableForRouting( + explicitGateway, + providersConfig, + config, + canonicalModelString + ) ) { const explicitGatewayDefinition = PROVIDER_DEFINITIONS[explicitGateway]; if (explicitGatewayDefinition.kind === "gateway") { diff --git a/src/node/services/providerService.test.ts b/src/node/services/providerService.test.ts index 60a89ace63e..ad30c5c1f67 100644 --- a/src/node/services/providerService.test.ts +++ b/src/node/services/providerService.test.ts @@ -7,6 +7,13 @@ import { writeFile } from "node:fs/promises"; import * as os from "os"; import * as path from "path"; import { CUSTOM_PROVIDER_TYPES } from "@/common/utils/providers/customProviders"; +import { KNOWN_MODELS } from "@/common/constants/knownModels"; +import { + hasCodexOauthTokens, + resolveCodexOauthRouting, +} from "@/common/utils/providers/codexOauthRouting"; +import { getEffectiveContextLimit } from "@/common/utils/compaction/contextLimit"; +import { openaiProModeAvailable } from "@/common/utils/ai/proMode"; import type { ProviderModelEntry } from "@/common/orpc/types"; import { WORKSPACE_DEFAULTS } from "@/constants/workspaceDefaults"; import { Config } from "@/node/config"; @@ -314,6 +321,152 @@ describe("ProviderService.getConfig", () => { }); }); + it("exposes account labels without exposing credentials", () => { + withTempConfig((config, service) => { + const auth = { + type: "oauth", + access: "secret-access", + refresh: "secret-refresh", + expires: 12345, + }; + new ProvidersConfigStore(config.rootDir).saveProvidersConfig({ + openai: { + codexOauth: auth, + codexOauthLabel: "Personal", + codexOauthAccounts: { + work: { label: "Work", credentials: auth }, + broken: { label: "Broken", credentials: {} }, + }, + codexOauthDefaultAccountId: "work", + }, + }); + const result = service.getConfig().openai; + expect(result.codexOauthSet).toBe(true); + expect(result.codexOauthAccounts).toEqual([ + { id: "default", label: "Personal" }, + { id: "work", label: "Work" }, + ]); + expect(result.codexOauthDefaultAccountId).toBe("work"); + expect(JSON.stringify(result)).not.toContain("secret-access"); + expect(JSON.stringify(result)).not.toContain("secret-refresh"); + }); + }); + + it("preserves implicit and explicit account selections for browser routing", () => { + withTempConfig((config, service) => { + saveOpenAIConfig(config); + expect(resolveCodexOauthRouting(KNOWN_MODELS.GPT.id, service.getConfig())).toBe("other"); + + saveOpenAIConfig(config, { codexOauthDefaultAccountId: "default" }); + expect(resolveCodexOauthRouting(KNOWN_MODELS.GPT.id, service.getConfig())).toBe( + "missing-account" + ); + }); + }); + + it.each(["default", "work"])( + "retains invalid %s credentials without reporting usable OAuth", + (accountId) => { + withProviderEnv({}, () => + withTempConfig((config, service) => { + const store = new ProvidersConfigStore(config.rootDir); + const auth = { + type: "oauth", + access: "private-access", + refresh: "private-refresh", + expires: 12345, + invalidReason: "invalid_grant", + }; + const accountConfig = + accountId === "default" + ? { codexOauth: auth } + : { codexOauthAccounts: { work: { label: "Work", credentials: auth } } }; + store.saveProvidersConfig({ + openai: { ...accountConfig, codexOauthDefaultAccountId: accountId }, + }); + const view = service.getConfig(); + expect(view.openai.codexOauthSet).toBe(false); + expect(view.openai.isConfigured).toBe(false); + expect(view.openai.codexOauthAccounts).toEqual([ + { + id: accountId, + label: accountId === "default" ? "Default" : "Work", + reconnectRequired: true, + }, + ]); + expect(view.openai.codexOauthDefaultAccountId).toBe(accountId); + expect(JSON.stringify(view)).not.toContain("private-access"); + expect(JSON.stringify(view)).not.toContain("private-refresh"); + expect(JSON.stringify(view)).not.toContain("invalid_grant"); + expect(hasCodexOauthTokens(view.openai)).toBe(false); + expect(hasCodexOauthTokens(accountConfig, accountId)).toBe(false); + expect(resolveCodexOauthRouting("openai:gpt-5.6-sol", view)).toBe("missing-account"); + expect(getEffectiveContextLimit("openai:gpt-5.6-sol", false, view)).toBe(372_000); + expect(openaiProModeAvailable("openai:gpt-5.6-sol", { providersConfig: view })).toBe( + false + ); + + store.saveProvidersConfig({ + openai: { + ...accountConfig, + codexOauthDefaultAccountId: accountId, + apiKey: "api-key", + codexOauthDefaultAuth: "apiKey", + }, + }); + const apiView = service.getConfig(); + expect(apiView.openai.codexOauthSet).toBe(false); + expect(apiView.openai.isConfigured).toBe(true); + expect(resolveCodexOauthRouting("openai:gpt-5.6-sol", apiView)).toBe("other"); + }) + ); + } + ); + + it("keeps other accounts available without substituting them for an invalid selection", () => { + withProviderEnv({}, () => + withTempConfig((config, service) => { + const auth = { type: "oauth", access: "access", refresh: "refresh", expires: 12345 }; + new ProvidersConfigStore(config.rootDir).saveProvidersConfig({ + openai: { + codexOauth: { ...auth, invalidReason: "invalid_grant" }, + codexOauthAccounts: { work: { label: "Work", credentials: auth } }, + }, + }); + const view = service.getConfig(); + expect(view.openai.codexOauthSet).toBe(true); + expect(view.openai.isConfigured).toBe(true); + expect(hasCodexOauthTokens(view.openai)).toBe(false); + expect(hasCodexOauthTokens(view.openai, "work")).toBe(true); + expect(resolveCodexOauthRouting("openai:gpt-5.6-sol", view)).toBe("missing-account"); + expect( + resolveCodexOauthRouting("openai:gpt-5.6-sol", view, { codexOauthAccountId: "work" }) + ).toBe("oauth"); + }) + ); + }); + + it("reports named accounts as connected without legacy credentials", () => { + withTempConfig((config, service) => { + new ProvidersConfigStore(config.rootDir).saveProvidersConfig({ + openai: { + codexOauthAccounts: { + work: { + label: "Work", + credentials: { type: "oauth", access: "access", refresh: "refresh", expires: 12345 }, + }, + }, + codexOauthDefaultAccountId: "work", + }, + }); + expect(service.getConfig().openai).toMatchObject({ + codexOauthSet: true, + isConfigured: true, + codexOauthAccounts: [{ id: "work", label: "Work" }], + }); + }); + }); + it("treats disabled OpenAI as unconfigured even when Codex OAuth tokens are stored", () => { withTempConfig((config, service) => { new ProvidersConfigStore(config.rootDir).saveProvidersConfig({ @@ -621,6 +774,31 @@ describe("ProviderService.getConfig", () => { ); }); + it("checks policy for account edits but retains internal credential cleanup", async () => { + await withTempPolicyProviderService( + { policy_format_version: "0.1", provider_access: [{ id: "anthropic" }] }, + async (config, service) => { + const store = new ProvidersConfigStore(config.rootDir); + store.saveProvidersConfig({ openai: { codexOauthDefaultAccountId: "personal" } }); + const denied = await service.updateConfigValue( + "openai", + ["codexOauthDefaultAccountId"], + () => ({ value: "work" }), + { enforcePolicy: true } + ); + expect(denied.success).toBe(false); + expect(store.loadProvidersConfig()?.openai?.codexOauthDefaultAccountId).toBe("personal"); + const cleaned = await service.updateConfigValue( + "openai", + ["codexOauthDefaultAccountId"], + () => ({ value: undefined }) + ); + expect(cleaned).toEqual({ success: true, data: { applied: true } }); + expect(store.loadProvidersConfig()?.openai?.codexOauthDefaultAccountId).toBeUndefined(); + } + ); + }); + it("revalidates policy inside the providers file lock in setModels", async () => { // Regression: policy can refresh while another process holds the // cross-process providers lock. A check done only before the lock wait @@ -2069,6 +2247,59 @@ describe("ProviderService.updateConfigValue", () => { }); }); +describe("ProviderService.updateProviderSection policy", () => { + it("denies user section writes but preserves internal mutation behavior", async () => { + await withTempPolicyProviderService( + { policy_format_version: "0.1", provider_access: [{ id: "anthropic" }] }, + async (config, service) => { + const denied = await service.updateProviderSection( + "openai", + () => ({ value: { codexOauthDefaultAccountId: "work" } }), + { enforcePolicy: true } + ); + expect(denied.success).toBe(false); + expect( + new ProvidersConfigStore(config.rootDir).loadProvidersConfig()?.openai + ).toBeUndefined(); + const internal = await service.updateProviderSection("openai", () => ({ + value: { codexOauthDefaultAccountId: "work" }, + })); + expect(internal).toEqual({ success: true, data: { applied: true } }); + } + ); + }); + + it("allows unchanged locked URLs but rejects section writes that change them", async () => { + await withTempPolicyProviderService( + { + policy_format_version: "0.1", + provider_access: [{ id: "openai", base_url: "https://locked.example.com" }], + }, + async (config, service) => { + const store = new ProvidersConfigStore(config.rootDir); + store.saveProvidersConfig({ + openai: { baseUrl: "https://old.example.com", apiKey: "keep-key" }, + }); + expect( + await service.updateProviderSection( + "openai", + (section) => ({ value: { ...section, codexOauthDefaultAccountId: "work" } }), + { enforcePolicy: true } + ) + ).toEqual({ success: true, data: { applied: true } }); + const before = store.loadProvidersConfig()?.openai; + const denied = await service.updateProviderSection( + "openai", + (section) => ({ value: { ...section, baseUrl: "https://new.example.com" } }), + { enforcePolicy: true } + ); + expect(denied.success).toBe(false); + expect(store.loadProvidersConfig()?.openai).toEqual(before); + } + ); + }); +}); + describe("ProviderService gateway lifecycle", () => { it("auto-inserts gateway into routePriority when configured", async () => { await withTempConfigAsync(async (config, service) => { diff --git a/src/node/services/providerService.ts b/src/node/services/providerService.ts index 6c8088d8e97..ab236e7edc5 100644 --- a/src/node/services/providerService.ts +++ b/src/node/services/providerService.ts @@ -29,7 +29,7 @@ import { SUPPORTED_PROVIDERS, type ProviderName, } from "@/common/constants/providers"; -import type { BaseProviderConfig } from "@/common/config/schemas/providersConfig"; +import type { BaseProviderConfig, ProvidersConfig } from "@/common/config/schemas/providersConfig"; import type { Result } from "@/common/types/result"; import type { AddCustomProviderInput, @@ -63,7 +63,7 @@ import { isProviderAutoRouteEligible, resolveProviderCredentials, } from "@/node/utils/providerRequirements"; -import { parseCodexOauthAuth } from "@/node/utils/codexOauthAuth"; +import { getCodexOauthAccounts } from "@/node/utils/codexOauthAuth"; import { normalizeCoderDeploymentUrl, parseCoderGatewayProviders, @@ -363,8 +363,9 @@ export class ProviderService { /** * Get the full providers config with safe info (no actual API keys) */ - public getConfig(): ProvidersConfigMap { - const providersConfig = this.providersConfigStore.loadProvidersConfig() ?? {}; + public getConfig( + providersConfig: ProvidersConfig = this.providersConfigStore.loadProvidersConfig() ?? {} + ): ProvidersConfigMap { const mainConfig = this.config.loadConfigOrDefault(); const result: ProvidersConfigMap = {}; const shadowedCustomProviderIds = this.detectAndLogShadowedProviders(providersConfig); @@ -425,8 +426,8 @@ export class ProviderService { config.models === undefined ? undefined : normalizeProviderModelEntries(config.models); const filteredModels = filterProviderModelsByPolicy(normalizedModels, allowedModels); - const codexOauthSet = - provider === "openai" && parseCodexOauthAuth(config.codexOauth) !== null; + const codexOauthAccounts = provider === "openai" ? getCodexOauthAccounts(config) : []; + const codexOauthSet = codexOauthAccounts.some(({ auth }) => auth.invalidReason === undefined); let isEnabled = !isProviderDisabledInConfig(config); if (provider === "mux-gateway" && mainConfig.muxGatewayEnabled === false) { isEnabled = false; @@ -506,6 +507,19 @@ export class ProviderService { if (provider === "openai") { providerInfo.codexOauthSet = codexOauthSet; + providerInfo.codexOauthAccounts = codexOauthAccounts.map(({ id, label, auth }) => ({ + id, + label, + // Keep invalid slot identities without exposing credentials or provider error details. + ...(auth.invalidReason !== undefined ? { reconnectRequired: true } : {}), + })); + // Preserve an unset selection. A synthetic default would block API-key-only routing in the renderer. + if ( + "codexOauthDefaultAccountId" in config && + typeof config.codexOauthDefaultAccountId === "string" + ) { + providerInfo.codexOauthDefaultAccountId = config.codexOauthDefaultAccountId; + } const codexOauthDefaultAuth = config.codexOauthDefaultAuth; if (codexOauthDefaultAuth === "oauth" || codexOauthDefaultAuth === "apiKey") { @@ -1447,22 +1461,23 @@ export class ProviderService { * logins/refreshes (e.g. Coder OAuth token rotation across the desktop app * and `mux run`/`mux workflow`). * - * Unlike setConfigValue, this path skips policy gating: it is an internal - * credential-management primitive (clearing dead tokens, persisting - * rotations), not a user-driven config edit. + * Internal credential updates skip policy gating by default. + * User-driven edits must set enforcePolicy to check policy under the file lock. */ public updateConfigValue( provider: string, keyPath: string[], - update: (current: unknown) => { value: unknown } | null + update: (current: unknown) => { value: unknown } | null, + options?: { enforcePolicy?: boolean } ): Promise> { - return Effect.runPromise(this.updateConfigValueEffect(provider, keyPath, update)); + return Effect.runPromise(this.updateConfigValueEffect(provider, keyPath, update, options)); } private updateConfigValueEffect( provider: string, keyPath: string[], - update: (current: unknown) => { value: unknown } | null + update: (current: unknown) => { value: unknown } | null, + options?: { enforcePolicy?: boolean } ): Effect.Effect> { // eslint-disable-next-line @typescript-eslint/no-this-alias -- Effect.gen generator bodies do not inherit `this` const self = this; @@ -1479,6 +1494,10 @@ export class ProviderService { } const applied = yield* self.providersFileLockEffect(() => { + if (options?.enforcePolicy) { + const denial = self.validateProviderEditPolicy(provider, keyPath); + if (denial != null) return denial; + } // Load, decide, and write under the lock — no awaits in between, so // the predicate result cannot be invalidated by any cooperating writer. const providersConfig = self.providersConfigStore.loadProvidersConfig() ?? {}; @@ -1517,6 +1536,7 @@ export class ProviderService { return true; }); + if (typeof applied === "string") return { success: false as const, error: applied }; yield* self.afterAppliedMutationEffect(provider, applied); return { success: true as const, data: { applied } }; }).pipe( @@ -1566,28 +1586,33 @@ export class ProviderService { * that fetched it is still the stored credential, and disconnect clears * tokens + models in one write. * - * Internal credential-management primitive: skips policy gating like - * updateConfigValue. + * Internal callers can omit policy checks. User-driven mutations must set enforcePolicy. */ public updateProviderSection( provider: string, update: ( section: Record | undefined - ) => { value: Record } | null + ) => { value: Record } | null, + options?: { enforcePolicy?: boolean } ): Promise> { - return Effect.runPromise(this.updateProviderSectionEffect(provider, update)); + return Effect.runPromise(this.updateProviderSectionEffect(provider, update, options)); } private updateProviderSectionEffect( provider: string, update: ( section: Record | undefined - ) => { value: Record } | null + ) => { value: Record } | null, + options?: { enforcePolicy?: boolean } ): Effect.Effect> { // eslint-disable-next-line @typescript-eslint/no-this-alias -- Effect.gen generator bodies do not inherit `this` const self = this; return Effect.gen(function* () { const applied = yield* self.providersFileLockEffect(() => { + if (options?.enforcePolicy) { + const denial = self.validateProviderEditPolicy(provider, []); + if (denial != null) return denial; + } const providersConfig = self.providersConfigStore.loadProvidersConfig() ?? {}; const section = providersConfig[provider] as Record | undefined; @@ -1596,6 +1621,14 @@ export class ProviderService { return false; } + if (options?.enforcePolicy) { + for (const key of ["baseUrl", "baseURL"]) { + if (decision.value[key] === section?.[key]) continue; + const denial = self.validateProviderEditPolicy(provider, [key]); + if (denial != null) return denial; + } + } + const deniedKey = Object.keys(decision.value).find((key) => DENIED_KEY_PATH_SEGMENTS.has(key) ); @@ -1608,6 +1641,7 @@ export class ProviderService { return true; }); + if (typeof applied === "string") return { success: false as const, error: applied }; // Best-effort: a landed write must not be reported as failed (see // afterAppliedMutationEffect). yield* self.afterAppliedMutationEffect(provider, applied); diff --git a/src/node/services/streamManager.continuousCompaction.test.ts b/src/node/services/streamManager.continuousCompaction.test.ts index 691c246e086..0c1a44613a9 100644 --- a/src/node/services/streamManager.continuousCompaction.test.ts +++ b/src/node/services/streamManager.continuousCompaction.test.ts @@ -772,8 +772,8 @@ describe("continuous prefix prepareStep and journal", () => { sliced || mode === "journal-failure" || mode === "ambiguous-anchor") - ? ["prefix-swap-invalidated"] - : [] + ? ["stream-model-update", "prefix-swap-invalidated"] + : ["stream-model-update"] ); if ( consumed && diff --git a/src/node/services/streamManager.test.ts b/src/node/services/streamManager.test.ts index 74e7e75dd30..2f851d562bd 100644 --- a/src/node/services/streamManager.test.ts +++ b/src/node/services/streamManager.test.ts @@ -5,7 +5,13 @@ import * as path from "node:path"; import { KNOWN_MODELS } from "@/common/constants/knownModels"; import type { ProvidersConfigMap } from "@/common/orpc/types"; -import { StreamEndEventSchema, ToolCallStartEventSchema } from "@/common/orpc/schemas/stream"; +import { + StreamEndEventSchema, + StreamStartEventSchema, + StreamModelUpdateEventSchema, + ToolCallStartEventSchema, + UsageDeltaEventSchema, +} from "@/common/orpc/schemas/stream"; import type { CompletedMessagePart, ToolCallEndEvent, @@ -3473,6 +3479,40 @@ describe("StreamManager - Concurrent Stream Prevention", () => { }); describe("StreamManager - exact step indices", () => { + test.each([272_000, null])( + "publishes the accepted limit when the stream starts: %s", + async (limit) => { + const streamManager = new StreamManager(historyService); + const workspaceId = "start-context-workspace"; + const messageId = "start-context-message"; + await appendPartialAssistantForTests(workspaceId, messageId, 1); + Reflect.set(streamManager, "tokenTracker", { + setModel: () => Promise.resolve(undefined), + countTokens: () => Promise.resolve(0), + }); + const starts: Array> = []; + onTurnEngineEvent(streamManager, "stream-start", (event) => starts.push(event)); + const streamInfo = createStreamInfoForTests({ + messageId, + effectiveContextLimit: limit, + streamResult: createStreamResultForTests( + (async function* () { + await Promise.resolve(); + yield { type: "finish", finishReason: "stop" }; + })() + ), + }); + await getProcessStreamWithCleanupForTests(streamManager).call( + streamManager, + workspaceId, + streamInfo, + 1 + ); + expect(starts).toHaveLength(1); + expect(StreamStartEventSchema.parse(starts[0]).effectiveContextLimit).toBe(limit); + } + ); + test("persists exact tool-only step boundaries through successful completion", async () => { const streamManager = new StreamManager(historyService); const workspaceId = "step-indices-workspace"; @@ -4149,9 +4189,126 @@ describe("StreamManager - empty stream completions", () => { expect(committed?.metadata?.usage).toBeUndefined(); }); + test.each([false, true])( + "publishes every fallback before usage while preserving parts: %s", + async (preserveParts) => { + const streamManager = new StreamManager(historyService); + const workspaceId = "fallback-metadata-workspace"; + const messageId = "fallback-metadata-message"; + const models = ["openai:gpt-5.5", "openai:gpt-5.6-sol"]; + const limits = [272_000, null]; + const entered = models.map(() => Promise.withResolvers()); + const release = models.map(() => Promise.withResolvers()); + const events: TurnEngineEvent[] = []; + streamManager.setEventSink((event) => { + events.push(event); + }); + await appendPartialAssistantForTests(workspaceId, messageId, 1); + Reflect.set(streamManager, "tokenTracker", { + setModel: () => Promise.resolve(undefined), + countTokens: () => Promise.resolve(0), + }); + let attempt = 0; + Reflect.set(streamManager, "createStreamResult", () => { + const index = attempt++; + return createStreamResultForTests( + (async function* () { + entered[index].resolve(); + await release[index].promise; + if (index === 0) { + yield { type: "finish", finishReason: "content-filter" }; + } else { + yield { type: "text-delta", text: "fallback answer" }; + yield { type: "finish", finishReason: "stop" }; + } + })() + ); + }); + const streamInfo = createStreamInfoForTests({ + messageId, + model: KNOWN_MODELS.SONNET.id, + effectiveContextLimit: 200_000, + initialMetadata: { routedThroughGateway: true, routeProvider: "mux-gateway" }, + streamResult: createStreamResultForTests( + (async function* () { + await Promise.resolve(); + if (preserveParts) yield { type: "text-delta", text: "partial answer" }; + yield { + type: "finish-step", + usage: { inputTokens: 1000, outputTokens: 0, totalTokens: 1000 }, + }; + yield { type: "finish", finishReason: "content-filter" }; + })() + ), + modelFallback: { + options: { + chain: models, + prepare: (modelString: string) => + Promise.resolve( + Ok({ + model: createTestLanguageModel(modelString), + modelString, + effectiveContextLimit: limits[models.indexOf(modelString)], + messages: [], + system: "fallback", + tools: undefined, + }) + ), + }, + requestedModel: KNOWN_MODELS.SONNET.id, + refusedModels: [], + original: { maxOutputTokens: undefined }, + }, + }); + getWorkspaceStreamsForTests(streamManager).set(workspaceId, streamInfo); + const processing = getProcessStreamWithCleanupForTests(streamManager).call( + streamManager, + workspaceId, + streamInfo, + 1 + ); + try { + for (let index = 0; index < models.length; index++) { + await entered[index].promise; + const updates = events.filter((event) => event.type === "stream-model-update"); + expect(updates).toHaveLength(index + 1); + const update = StreamModelUpdateEventSchema.parse(updates[index]); + expect(update.model).toBe(models[index]); + expect(update.metadataModel).toBe(models[index]); + expect(update.effectiveContextLimit).toBe(limits[index]); + expect(update.routedThroughGateway).toBe(false); + expect(update.routeProvider).toBeUndefined(); + expect(update.modelFallback.refusedModels).toHaveLength(index + 1); + expect(events.filter((event) => event.type === "usage-delta")).toHaveLength(1); + expect( + events.filter((event) => event.type === "stream-start" && !event.replay) + ).toHaveLength(1); + expect((streamInfo.parts as Array<{ text: string }>).map((part) => part.text)).toEqual( + preserveParts ? ["partial answer"] : [] + ); + await streamManager.replayStream(workspaceId); + const replayStart = events.findLast((event) => event.type === "stream-start"); + expect(StreamStartEventSchema.parse(replayStart)).toMatchObject({ + model: models[index], + effectiveContextLimit: limits[index], + modelFallback: update.modelFallback, + }); + release[index].resolve(); + } + } finally { + for (const gate of release) gate.resolve(); + await processing; + } + } + ); + test("zero-output refusal with a configured fallback chain swaps models without any error event", async () => { const streamManager = new StreamManager(historyService); const errorEvents: unknown[] = []; + const limits: Array = []; + onTurnEngineEvent(streamManager, "usage-delta", (event) => + limits.push(event.effectiveContextLimit) + ); const streamEndEvents: Array<{ metadata?: { model?: string; @@ -4188,6 +4345,7 @@ describe("StreamManager - empty stream completions", () => { (async function* () { await Promise.resolve(); yield { type: "text-delta", text: "fallback answer" }; + yield { type: "finish-step", usage: { inputTokens: 5, outputTokens: 3, totalTokens: 8 } }; yield { type: "finish", finishReason: "stop" }; })(), { inputTokens: 5, outputTokens: 3, totalTokens: 8 } @@ -4211,6 +4369,7 @@ describe("StreamManager - empty stream completions", () => { Ok({ model: fallbackLanguageModel, modelString: nextModelString, + effectiveContextLimit: 272_000, messages: [], system: "fallback system", tools: fallbackTools, @@ -4241,6 +4400,7 @@ describe("StreamManager - empty stream completions", () => { startTime, lastPartTimestamp: startTime, model: KNOWN_MODELS.SONNET.id, + effectiveContextLimit: 200_000, metadataModel: KNOWN_MODELS.SONNET.id, historySequence, initialMetadata: { agentId: "plan" }, @@ -4258,6 +4418,7 @@ describe("StreamManager - empty stream completions", () => { // No terminal failure: TaskService and waiters never observe the refusal. expect(errorEvents).toHaveLength(0); + expect(limits).toEqual([200_000, 272_000]); expect(prepare).toHaveBeenCalledTimes(1); expect(prepare.mock.calls[0]?.[0]).toBe(fallbackModel); expect(prepare.mock.calls[0]?.[1]).toBeUndefined(); @@ -6213,6 +6374,39 @@ describe("StreamManager - replayStream", () => { ]); }); + test.each([272_000, null])( + "replays the accepted limit before usage exists: %s", + async (limit) => { + const streamManager = createReplayStreamManager(); + const workspaceId = "replay-context-before-usage"; + const starts: Array> = []; + const usageEvents: unknown[] = []; + onTurnEngineEvent(streamManager, "stream-start", (event) => starts.push(event)); + onTurnEngineEvent(streamManager, "usage-delta", (event) => usageEvents.push(event)); + setReplayStreamInfo(streamManager, workspaceId, { + state: "streaming", + messageId: "replay-context-message", + model: "openai:gpt-5.5", + effectiveContextLimit: limit, + historySequence: 1, + startTime: 123, + initialMetadata: {}, + toolCompletionTimestamps: new Map(), + parts: [], + }); + stubReplayTokenTracker(streamManager); + await streamManager.replayStream(workspaceId); + await streamManager.replayStream(workspaceId, { afterTimestamp: 123 }); + expect(starts).toHaveLength(2); + for (const start of starts) { + const parsed = StreamStartEventSchema.parse(start); + expect(parsed.effectiveContextLimit).toBe(limit); + expect(parsed.replay).toBe(true); + } + expect(usageEvents).toHaveLength(0); + } + ); + test("replayStream emits replay usage-delta from tracked step/cumulative usage", async () => { const streamManager = createReplayStreamManager(); @@ -6226,6 +6420,7 @@ describe("StreamManager - replayStream", () => { setReplayStreamInfo(streamManager, workspaceId, { state: "streaming", messageId: "msg-usage", + effectiveContextLimit: 100_000, model: "claude-sonnet-4", metadataModel: "claude-sonnet-4", historySequence: 1, @@ -6245,6 +6440,7 @@ describe("StreamManager - replayStream", () => { expect(usageEvents).toHaveLength(1); expect(usageEvents[0]?.replay).toBe(true); + expect(UsageDeltaEventSchema.parse(usageEvents[0]).effectiveContextLimit).toBe(100_000); expect(usageEvents[0]?.usage).toEqual({ inputTokens: 21, outputTokens: 3, totalTokens: 24 }); expect(usageEvents[0]?.providerMetadata).toEqual({ anthropic: { cacheReadInputTokens: 2 }, diff --git a/src/node/services/streamManager.ts b/src/node/services/streamManager.ts index 6929ff22e12..55972d699ec 100644 --- a/src/node/services/streamManager.ts +++ b/src/node/services/streamManager.ts @@ -29,6 +29,7 @@ import { Ok, Err } from "@/common/types/result"; import { log, type Logger } from "./log"; import type { StreamStartEvent, + StreamModelUpdateEvent, StreamDeltaEvent, StreamEndEvent, StreamAbortEvent, @@ -192,6 +193,7 @@ type StreamToken = string & { __brand: "StreamToken" }; export type TurnEngineEvent = | StreamStartEvent + | StreamModelUpdateEvent | StreamDeltaEvent | StreamEndEvent | StreamAbortEvent @@ -267,6 +269,7 @@ interface StreamRequestOptions { rebuildProviderOptionsForThinkingLevel?: RebuildProviderOptionsForThinkingLevel; forcedFirstStepToolNames?: string[]; providersConfigSnapshot?: ProvidersConfigMap; + effectiveContextLimit?: number | null; rebuildFirstStepForThinkingLevel?: RebuildFirstStepForThinkingLevel; } @@ -397,6 +400,7 @@ interface PreparedModelFallback { * usage identity. */ providersConfig?: ProvidersConfigMap; + effectiveContextLimit?: number | null; } export interface ModelFallbackPrepareOptions { @@ -571,6 +575,7 @@ function buildUsageDeltaEvent(opts: { cumulativeUsage: LanguageModelV2Usage; cumulativeProviderMetadata: Record | undefined; costsIncluded: boolean | undefined; + effectiveContextLimit?: number | null; replay?: true; }): UsageDeltaEvent { return { @@ -579,6 +584,7 @@ function buildUsageDeltaEvent(opts: { messageId: opts.messageId, ...(opts.replay ? { replay: true } : {}), usage: opts.usage, + effectiveContextLimit: opts.effectiveContextLimit, providerMetadata: opts.providerMetadata, cumulativeUsage: opts.cumulativeUsage, cumulativeProviderMetadata: markProviderMetadataCostsIncluded( @@ -661,6 +667,7 @@ interface WorkspaceStreamInfo { model: string; /** Metadata model resolved from provider mapping for cost/token metadata lookups. */ metadataModel: string; + effectiveContextLimit?: number | null; /** Effective thinking level after model policy clamping */ thinkingLevel?: string; initialMetadata?: Partial; @@ -2650,6 +2657,7 @@ export class StreamManager { pendingToolExecutionStarts: new Map(), model: modelString, metadataModel, + effectiveContextLimit: options.effectiveContextLimit, thinkingLevel, initialMetadata, toolModelUsages: [], @@ -2997,6 +3005,8 @@ export class StreamManager { // diverge from the backend ledger when a Coder catalog refresh // removes/retags the instance mid-stream. metadataModel: streamInfo.metadataModel, + effectiveContextLimit: streamInfo.effectiveContextLimit, + modelFallback: streamInfo.initialMetadata?.modelFallback, routedThroughGateway, ...(routeProvider != null && { routeProvider }), historySequence, @@ -3525,6 +3535,7 @@ export class StreamManager { streamInfo.reasoningBackfillStartIndex = preserveParts ? streamInfo.parts.length : undefined; streamInfo.model = prepared.data.modelString; + streamInfo.effectiveContextLimit = prepared.data.effectiveContextLimit ?? null; streamInfo.metadataModel = this.resolveMetadataModel( prepared.data.modelString, prepared.data.providersConfig @@ -3532,15 +3543,21 @@ export class StreamManager { if (prepared.data.thinkingLevel !== undefined) { streamInfo.thinkingLevel = prepared.data.thinkingLevel; } + const modelFallback = { + requestedModel: fallbackState.requestedModel, + refusedModels: [...fallbackState.refusedModels], + }; // Final stream-end metadata spreads initialMetadata, so route attribution // corrections and the fallback record propagate automatically. streamInfo.initialMetadata = { ...streamInfo.initialMetadata, ...prepared.data.initialMetadataPatch, - modelFallback: { - requestedModel: fallbackState.requestedModel, - refusedModels: [...fallbackState.refusedModels], - }, + // Missing fallback route fields must not retain the refused gateway's attribution. + routedThroughGateway: + prepared.data.initialMetadataPatch?.routedThroughGateway ?? + prepared.data.modelString.startsWith("mux-gateway:"), + routeProvider: prepared.data.initialMetadataPatch?.routeProvider, + modelFallback, }; // Release the refused model's transport resources now: the stream-exit // finally only cleans the final request's model, so without this the @@ -3548,6 +3565,19 @@ export class StreamManager { runLanguageModelCleanup(streamInfo.request.model); streamInfo.request = nextRequest; streamInfo.streamResult = nextStreamResult; + // Publish each accepted attempt before its first step, without restarting stream lifecycle consumers. + this.emitTurnEvent({ + type: "stream-model-update", + workspaceId, + messageId: streamInfo.messageId, + model: metadataModelIdentity(streamInfo.model), + metadataModel: streamInfo.metadataModel, + effectiveContextLimit: streamInfo.effectiveContextLimit, + modelFallback, + routedThroughGateway: streamInfo.initialMetadata.routedThroughGateway ?? false, + routeProvider: streamInfo.initialMetadata.routeProvider, + thinkingLevel: streamInfo.thinkingLevel as ThinkingLevel | undefined, + }); await this.tokenTracker.setModel(streamInfo.model, streamInfo.metadataModel); if ( consumedSwap && @@ -4134,6 +4164,7 @@ export class StreamManager { cumulativeProviderMetadata: streamInfo.cumulativeProviderMetadata, // Preserve gateway-billed zero-cost behavior throughout the stream. costsIncluded: streamInfo.initialMetadata?.costsIncluded, + effectiveContextLimit: streamInfo.effectiveContextLimit, }); streamInfo.currentStepStartIndex = streamInfo.parts.length; this.emitTurnEvent(usageEvent); @@ -5754,6 +5785,7 @@ export class StreamManager { providerMetadata: streamInfo.lastStepProviderMetadata, cumulativeUsage: streamInfo.cumulativeUsage, cumulativeProviderMetadata: streamInfo.cumulativeProviderMetadata, + effectiveContextLimit: streamInfo.effectiveContextLimit, // Replays must preserve gateway-billed zero-cost behavior from the original stream. costsIncluded: streamInfo.initialMetadata?.costsIncluded, }); diff --git a/src/node/services/taskWorkspaceSeam.ts b/src/node/services/taskWorkspaceSeam.ts index 2a4d358f75f..5b779f6fb4b 100644 --- a/src/node/services/taskWorkspaceSeam.ts +++ b/src/node/services/taskWorkspaceSeam.ts @@ -12,6 +12,7 @@ import type { } from "@/common/types/message"; import type { Result } from "@/common/types/result"; import type { StreamErrorRecoveryOutcome } from "@/node/services/agentSession"; +import type { ModelRoutingSnapshot } from "@/node/services/modelRoutingSnapshot"; import type { RuntimeConfig } from "@/common/types/runtime"; import type { FrontendWorkspaceMetadata, WorkspaceMetadata } from "@/common/types/workspace"; import type { AgentAiSettingsLayerValues } from "@/common/types/agentAiSettings"; @@ -310,6 +311,8 @@ export interface WorkspaceLiveActivity { } export interface SendMessageInternalOptions { + /** In-memory routing for same-session, idle-only continuations. Never persist these credentials. */ + modelRoutingSnapshot?: ModelRoutingSnapshot; allowQueuedAgentTask?: boolean; skipAutoResumeReset?: boolean; synthetic?: boolean; diff --git a/src/node/services/tools/shared/configRedaction.test.ts b/src/node/services/tools/shared/configRedaction.test.ts new file mode 100644 index 00000000000..c107a66c437 --- /dev/null +++ b/src/node/services/tools/shared/configRedaction.test.ts @@ -0,0 +1,67 @@ +import { describe, expect, it } from "bun:test"; +import { redactConfigDocument, REDACTED_SECRET_VALUE } from "./configRedaction"; + +describe("Codex account redaction", () => { + it("protects named credentials through the legacy generic redaction rule", () => { + const credentials = { + type: "oauth", + access: "private-access", + refresh: "private-refresh", + accountId: "private-chatgpt-account", + credentialId: "eb5beccb-f8a9-4dc0-b8ce-2bd2954f4e42", + expires: 12345, + }; + const accounts = { work: { label: "Work account", credentials } }; + // Commit 5605e7852 has the same generic rules, without the codexOauthAccounts explicit key. + // An unknown container exercises those rules without the new account-name protection. + const document = { openai: { futureAccountSlots: accounts } }; + const redacted = redactConfigDocument("providers", document); + expect(redacted).toEqual({ + openai: { + futureAccountSlots: { + work: { label: "Work account", credentials: REDACTED_SECRET_VALUE }, + }, + }, + }); + const serialized = JSON.stringify(redacted); + for (const secret of [ + credentials.access, + credentials.refresh, + credentials.accountId, + credentials.credentialId, + ]) { + expect(serialized).not.toContain(secret); + } + expect(serialized).not.toContain("accountId"); + expect(serialized).not.toContain("credentialId"); + expect(accounts.work.credentials).toEqual(credentials); + }); + + it("removes all account credentials without changing the source document", () => { + const auth = { + type: "oauth", + access: "private-access", + refresh: "private-refresh", + expires: 12345, + }; + const document = { + openai: { + codexOauth: auth, + codexOauthAccounts: { work: { label: "Work", credentials: auth } }, + codexOauthDefaultAccountId: "work", + }, + }; + + const redacted = redactConfigDocument("providers", document); + expect(redacted).toEqual({ + openai: { + codexOauth: REDACTED_SECRET_VALUE, + codexOauthAccounts: REDACTED_SECRET_VALUE, + codexOauthDefaultAccountId: "work", + }, + }); + expect(JSON.stringify(redacted)).not.toContain(auth.access); + expect(JSON.stringify(redacted)).not.toContain(auth.refresh); + expect(document.openai.codexOauthAccounts.work.credentials).toEqual(auth); + }); +}); diff --git a/src/node/services/tools/shared/configRedaction.ts b/src/node/services/tools/shared/configRedaction.ts index 9367bb16708..adc2b170579 100644 --- a/src/node/services/tools/shared/configRedaction.ts +++ b/src/node/services/tools/shared/configRedaction.ts @@ -13,6 +13,7 @@ const PROVIDER_SECRET_KEYS = new Set([ "couponCode", "voucher", "codexOauth", + "codexOauthAccounts", ]); const APP_SECRET_KEYS = new Set(["muxGovernorToken"]); diff --git a/src/node/services/turnRequestBuilder.test.ts b/src/node/services/turnRequestBuilder.test.ts index 191fafccb8d..81c9938d733 100644 --- a/src/node/services/turnRequestBuilder.test.ts +++ b/src/node/services/turnRequestBuilder.test.ts @@ -251,6 +251,196 @@ describe("TurnRequestBuilder tool scope", () => { }); describe("TurnRequestBuilder model attempt preparation", () => { + it("computes limits from accepted auth settings instead of current provider settings", async () => { + const harness = await createPreparationHarness(); + try { + const accepted: ProvidersConfigMap = { + openai: { + apiKeySet: true, + isEnabled: true, + isConfigured: true, + codexOauthSet: true, + codexOauthDefaultAccountId: "work", + codexOauthAccounts: [{ id: "work", label: "Work" }], + }, + }; + harness.providersConfigStore.saveProvidersConfig({ + openai: { codexOauthDefaultAuth: "apiKey" }, + }); + const model = "openai:gpt-5.5"; + const options = preparationOptions(accepted, { + rawModelString: model, + canonicalModelString: model, + canonicalProviderName: "openai", + effectiveModelString: model, + optionsModelString: model, + wireProviderName: "openai", + }); + const live = harness.builder.prepareModelAttempt(options); + const next = harness.builder.prepareModelAttempt({ + ...options, + providersConfigSnapshot: { + openai: { ...accepted.openai, codexOauthDefaultAuth: "apiKey" }, + }, + }); + const chatCompletions = harness.builder.prepareModelAttempt({ + ...options, + muxProviderOptions: { openai: { wireFormat: "chatCompletions" } }, + }); + expect(live.effectiveContextLimit).toBe(272_000); + expect(next.effectiveContextLimit).toBeGreaterThan(272_000); + expect(chatCompletions.effectiveContextLimit).toBe(next.effectiveContextLimit); + } finally { + await harness.cleanup(); + } + }); + + it("does not apply direct OpenAI OAuth caps to an automatic gateway route", async () => { + const harness = await createPreparationHarness(); + try { + const model = "openai:gpt-5.5"; + const options = preparationOptions( + { + openai: { + apiKeySet: true, + isEnabled: true, + isConfigured: true, + codexOauthSet: true, + codexOauthDefaultAccountId: "work", + codexOauthAccounts: [{ id: "work", label: "Work" }], + models: [{ id: "gpt-5.5", contextWindowTokens: 500_000 }], + }, + }, + { + rawModelString: model, + canonicalModelString: model, + canonicalProviderName: "openai", + effectiveModelString: model, + optionsModelString: model, + wireProviderName: "openai", + } + ); + const direct = harness.builder.prepareModelAttempt(options); + const gateway = harness.builder.prepareModelAttempt({ + ...options, + effectiveModelString: "openrouter:openai/gpt-5.5", + routeProvider: "openrouter", + }); + // A rejected gateway selection can resolve to the direct provider on a later attempt. + const fallback = harness.builder.prepareModelAttempt({ + ...options, + rawModelString: "openrouter:openai/gpt-5.5", + routeProvider: "openai", + }); + expect(direct.effectiveContextLimit).toBe(272_000); + expect(gateway.effectiveContextLimit).toBe(500_000); + expect(fallback.effectiveContextLimit).toBe(272_000); + } finally { + await harness.cleanup(); + } + }); + + it.each([ + { + routeProvider: "anthropic" as const, + effectiveModelString: "anthropic:claude-sonnet-4-5", + expectedLimit: 1_000_000, + }, + { + routeProvider: "openrouter" as const, + effectiveModelString: "openrouter:anthropic/claude-sonnet-4-5", + expectedLimit: 100_000, + }, + { + routeProvider: "mux-gateway" as const, + effectiveModelString: "mux-gateway:anthropic/claude-sonnet-4-5", + expectedLimit: 1_000_000, + }, + ])("uses the emitted 1M beta capability for the $routeProvider route", async (testCase) => { + const harness = await createPreparationHarness(); + try { + const options = preparationOptions( + { + anthropic: { + apiKeySet: true, + isEnabled: true, + isConfigured: true, + models: [{ id: "claude-sonnet-4-5", contextWindowTokens: 100_000 }], + }, + }, + { + routeProvider: testCase.routeProvider, + effectiveModelString: testCase.effectiveModelString, + muxProviderOptions: { + anthropic: { use1MContextModels: ["anthropic:claude-sonnet-4-5"] }, + }, + } + ); + const prepared = harness.builder.prepareModelAttempt(options); + expect(prepared.effectiveContextLimit).toBe(testCase.expectedLimit); + expect(prepared.requestHeaders?.["anthropic-beta"] !== undefined).toBe( + testCase.expectedLimit === 1_000_000 + ); + } finally { + await harness.cleanup(); + } + }); + + it("uses the direct context limit when a Coder selection falls back to Anthropic", async () => { + const harness = await createPreparationHarness(); + try { + const prepared = harness.builder.prepareModelAttempt( + preparationOptions( + { + anthropic: { + apiKeySet: true, + isEnabled: true, + isConfigured: true, + models: [{ id: "claude-sonnet-4-5", contextWindowTokens: 100_000 }], + }, + coder: { + apiKeySet: false, + isEnabled: true, + isConfigured: false, + additionalProviders: [{ name: "production", type: "anthropic" }], + models: [{ id: "production/claude-sonnet-4-5", contextWindowTokens: 80_000 }], + }, + }, + { + rawModelString: "coder:production/claude-sonnet-4-5", + routeProvider: "anthropic", + } + ) + ); + expect(prepared.effectiveContextLimit).toBe(100_000); + } finally { + await harness.cleanup(); + } + }); + + it("honors accepted custom limits and the request's 1M context option", async () => { + const harness = await createPreparationHarness(); + try { + const options = preparationOptions({ + anthropic: { + apiKeySet: true, + isEnabled: true, + isConfigured: true, + models: [{ id: "claude-sonnet-4-5", contextWindowTokens: 100_000 }], + }, + }); + expect(harness.builder.prepareModelAttempt(options).effectiveContextLimit).toBe(100_000); + expect( + harness.builder.prepareModelAttempt({ + ...options, + muxProviderOptions: { anthropic: { use1MContextModels: [options.rawModelString] } }, + }).effectiveContextLimit + ).toBe(1_000_000); + } finally { + await harness.cleanup(); + } + }); + it("merges call settings and provider extras at the resolved namespace", async () => { const harness = await createPreparationHarness(); try { diff --git a/src/node/services/turnRequestBuilder.ts b/src/node/services/turnRequestBuilder.ts index 51c292eb2e3..2f37af1c818 100644 --- a/src/node/services/turnRequestBuilder.ts +++ b/src/node/services/turnRequestBuilder.ts @@ -53,6 +53,7 @@ import type { InitStateManager } from "./initStateManager"; import { runLanguageModelCleanup } from "./languageModelCleanup"; import { log } from "./log"; import type { StreamManager } from "./streamManager"; +import { getEffectiveContextLimit } from "@/common/utils/compaction/contextLimit"; import { type ModelFallbackOptions, type StreamTextOnChunk, @@ -92,6 +93,7 @@ import type { WorkspaceMCPOverrides } from "@/common/types/mcp"; import { isExecLikeEditingCapableInResolvedChain } from "@/common/utils/agentTools"; import { resolveModelParameterOverrides } from "@/common/utils/ai/modelParameterOverrides"; import { + ANTHROPIC_1M_CONTEXT_HEADER, buildProviderOptions, buildRequestHeaders, resolveProviderOptionsNamespaceKey, @@ -229,6 +231,7 @@ import type { ErrorEvent } from "@/common/types/stream"; import type { ToolPolicy } from "@/common/utils/tools/toolPolicy"; import type { FileState } from "@/node/services/agentSession"; import type { ActiveTurnThinkingOverride } from "@/node/services/thinkingOverride"; +import type { ModelRoutingSnapshot } from "./modelRoutingSnapshot"; import type { WorkspaceGoalService } from "@/node/services/workspaceGoalService"; /** Options used to prepare and execute a turn. */ @@ -246,6 +249,8 @@ export interface StreamMessageOptions { additionalSystemInstructions?: string; maxOutputTokens?: number; muxProviderOptions?: MuxProviderOptions; + /** Internal routing state shared with the pre-send compaction check. */ + modelRoutingSnapshot?: ModelRoutingSnapshot; /** Internal-only flag for Copilot billing attribution; never sourced from IPC schemas. */ agentInitiated?: boolean; agentId?: string; @@ -424,6 +429,7 @@ interface WorkflowResultContinuationSender { options: SendMessageOptions, internal?: { skipAutoResumeReset?: boolean; + modelRoutingSnapshot?: ModelRoutingSnapshot; synthetic?: boolean; agentInitiated?: boolean; /** When true, reject instead of queueing if the workspace is busy. */ @@ -517,7 +523,12 @@ interface TurnRequestBuilderDependencies { createModel: ( modelString: string, muxProviderOptions?: MuxProviderOptions, - opts?: { agentInitiated?: boolean; workspaceId?: string; providersConfig?: ProvidersConfig } + opts?: { + agentInitiated?: boolean; + workspaceId?: string; + providersConfig?: ProvidersConfig; + modelRoutingSnapshot?: ModelRoutingSnapshot; + } ) => Promise>; isStreaming: (workspaceId: string) => boolean; trackPendingDevToolsRunMetadata: ( @@ -549,6 +560,7 @@ export interface PrepareModelAttemptOptions { } interface PreparedModelAttempt { + effectiveContextLimit: number | null; providerOptions: Record; requestHeaders: Record | undefined; resolvedOverrides: ReturnType; @@ -706,6 +718,13 @@ export class TurnRequestBuilder { }; options.recordStartupPhaseTiming?.("buildRequestConfigMs", buildRequestConfigStartedAt); return { + // Use the resolved route and emitted beta header, not the requested direct-provider identity. + effectiveContextLimit: getEffectiveContextLimit( + options.effectiveModelString, + requestHeaders?.["anthropic-beta"] === ANTHROPIC_1M_CONTEXT_HEADER, + options.providersConfigSnapshot, + { openaiWireFormat: options.muxProviderOptions.openai?.wireFormat } + ), providerOptions: mergeExtras(providerOptions), requestHeaders, resolvedOverrides, @@ -879,7 +898,8 @@ export class TurnRequestBuilder { } const requestedThinkingLevel = options.requestedThinkingLevel ?? THINKING_LEVEL_OFF; - const preliminaryProvidersConfig = this.dependencies.providerService.getConfig(); + const preliminaryProvidersConfig = + opts.modelRoutingSnapshot?.metadata ?? this.dependencies.providerService.getConfig(); const preliminaryMinThinkingLevel = resolveMinimumThinkingLevel( options.rawModelString, options.minimumThinkingLevelOverride, @@ -903,7 +923,15 @@ export class TurnRequestBuilder { options.rawModelString, preliminaryThinkingLevel, effectiveMuxProviderOptions, - { agentInitiated, workspaceId } + { + agentInitiated, + workspaceId, + ...(opts.modelRoutingSnapshot && { + providersConfig: opts.modelRoutingSnapshot.providersConfig, + codexOauthSelection: opts.modelRoutingSnapshot.codexOauthSelection, + routeConfig: opts.modelRoutingSnapshot.routeConfig, + }), + } ); if (options.recordTiming) { recordStartupPhaseTiming("resolveAndCreateModelMs", resolveAndCreateModelStartedAt); @@ -912,11 +940,25 @@ export class TurnRequestBuilder { return resolved; } - const providersConfig = pinCoderInstanceProvidersConfig( - this.dependencies.providerService.getConfig(), + let providersConfig = pinCoderInstanceProvidersConfig( + opts.modelRoutingSnapshot?.metadata ?? this.dependencies.providerService.getConfig(), options.rawModelString, resolved.data.coderSelectedInstance ); + // Context-limit mirrors must use the account that the model selected. + if ( + providersConfig.openai && + resolved.data.codexOauthAccountId != null && + (!opts.modelRoutingSnapshot || opts.modelRoutingSnapshot.codexOauthSelection.explicit) + ) { + providersConfig = { + ...providersConfig, + openai: { + ...providersConfig.openai, + codexOauthDefaultAccountId: resolved.data.codexOauthAccountId, + }, + }; + } const minThinkingLevel = resolveMinimumThinkingLevel( options.rawModelString, options.minimumThinkingLevelOverride, @@ -1774,6 +1816,8 @@ export class TurnRequestBuilder { return; } if (this.dependencies.bindings.taskService != null) { + // Durable attention can combine multiple origins and recover after restart. + // It starts a new turn with current routing; only the one-origin fallback keeps this snapshot. this.dependencies.bindings.taskService.noteWorkflowRunTerminalAttention({ ownerWorkspaceId: workspaceId, runId, @@ -1846,6 +1890,8 @@ export class TurnRequestBuilder { }, { skipAutoResumeReset: true, + // Only this live callback retains credentials. Restart recovery captures current routing. + modelRoutingSnapshot: opts.modelRoutingSnapshot, synthetic: true, agentInitiated: true, requireIdle: true, @@ -1878,7 +1924,7 @@ export class TurnRequestBuilder { const assistantMessageId = createAssistantMessageId(); const allowLegacyInvalidWorkflowAgentOutputSchema = await this.dependencies.shouldAllowLegacyInvalidWorkflowAgentOutputSchema(metadata); - // Share creation-time provider/pricing snapshots for both headless tools. + // Headless tools retain accepted turn routing, but use their own model options. const createToolModel = async (ms: string) => { const toolModelString = ms.trim(); assert( @@ -1890,11 +1936,13 @@ export class TurnRequestBuilder { // a catalog refresh land between them, running the request // on one wire while recording usage under another type. const toolProvidersConfig = - this.dependencies.providersConfigStore.loadProvidersConfig() ?? {}; - // View snapshot captured at creation time for option - // building (buildProviderOptions takes the oRPC view, not - // the raw config shape). - const toolOptionsProvidersConfig = this.dependencies.providerService.getConfig(); + opts.modelRoutingSnapshot?.providersConfig ?? + this.dependencies.providersConfigStore.loadProvidersConfig() ?? + {}; + // Option building uses the public provider view from the same snapshot. + const toolOptionsProvidersConfig = + opts.modelRoutingSnapshot?.metadata ?? + this.dependencies.providerService.getConfig(toolProvidersConfig); // Let the factory pin provider-level defaults (especially the OpenAI wire // format) without inheriting any options from the parent chat. const toolMuxProviderOptions: MuxProviderOptions = {}; @@ -1904,6 +1952,7 @@ export class TurnRequestBuilder { { workspaceId, providersConfig: toolProvidersConfig, + modelRoutingSnapshot: opts.modelRoutingSnapshot, agentInitiated: true, } ); @@ -1919,7 +1968,8 @@ export class TurnRequestBuilder { this.dependencies.providerModelFactory.resolveEffectiveModelString( toolModelString, undefined, - toolProvidersConfig + toolProvidersConfig, + opts.modelRoutingSnapshot?.routeConfig ); const toolOnCoderRoute = toolEffectiveModelString.startsWith("coder:"); // Creation-time identity from the SAME snapshot the model @@ -2516,6 +2566,7 @@ export class TurnRequestBuilder { engineTools: attemptPayload.tools ?? attemptTools, toolNamesForSentinel, forcedFirstStepToolNames, + effectiveContextLimit: preparedAttempt.effectiveContextLimit, providerOptions: preparedAttempt.providerOptions, headers: preparedAttempt.requestHeaders, resolvedOverrides: preparedAttempt.resolvedOverrides, @@ -2843,6 +2894,7 @@ export class TurnRequestBuilder { rebuildProviderOptionsForThinkingLevel: nextRequest.rebuildProviderOptionsForThinkingLevel, providersConfig: nextRequest.providersConfig, + effectiveContextLimit: nextRequest.effectiveContextLimit, initialMetadataPatch: { routedThroughGateway: nextRequest.routedThroughGateway, ...(nextRequest.routeProvider != null @@ -2944,6 +2996,7 @@ export class TurnRequestBuilder { rebuildProviderOptionsForThinkingLevel, forcedFirstStepToolNames, providersConfigSnapshot: requestProvidersConfig, + effectiveContextLimit: primaryRequest.effectiveContextLimit, onStreamConstructed: emitPrimaryEnvelope, rebuildFirstStepForThinkingLevel: primaryRequest.rebuildFirstStepForThinkingLevel, }; diff --git a/src/node/services/workspaceService.test.ts b/src/node/services/workspaceService.test.ts index 0a4defe9bc2..09ddd491981 100644 --- a/src/node/services/workspaceService.test.ts +++ b/src/node/services/workspaceService.test.ts @@ -7,6 +7,7 @@ import type { AgentSession } from "./agentSession"; import { CONTEXT_MUTATION_SEND_BLOCKED_MESSAGE } from "./agentSession"; import { createAgentSessionHarness, + createModelRoutingSnapshotMock, createStartedTurnHandle, createStreamLifecycleMocks, } from "./agentSession.testHarness"; @@ -6356,6 +6357,7 @@ describe("WorkspaceService truncateHistory goal acknowledgment", () => { aiServiceOverride ?? ({ ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on: mock(() => undefined), isStreaming: mock(() => false), } as unknown as AIService); @@ -7137,6 +7139,7 @@ describe("WorkspaceService truncateHistory goal acknowledgment", () => { let streaming = true; const aiService = { ...createStreamLifecycleMocks(), + captureModelRoutingSnapshot: createModelRoutingSnapshotMock(), on: mock(() => undefined), isStreaming: mock(() => streaming), } as unknown as AIService; @@ -7389,6 +7392,12 @@ describe("WorkspaceService truncateHistory goal acknowledgment", () => { }), ]; const summaryAiService: BranchSummaryAiService = { + captureModelRoutingSnapshot: () => ({ + providersConfig: {}, + metadata: null, + routeConfig: { routePriority: ["direct"], routeOverrides: {} }, + codexOauthSelection: { accountId: "default", explicit: false }, + }), createModelWithPinnedMetadata: (modelString: string) => Promise.resolve( Ok({ @@ -8827,6 +8836,41 @@ describe("WorkspaceService sendMessage status clearing", () => { await cleanupHistory(); }); + test("forwards workflow routing only through internal idle-send options", async () => { + fakeSession.isBusy.mockReturnValue(false); + const modelRoutingSnapshot: ReturnType = { + providersConfig: { openai: { apiKey: "snapshot-secret" } }, + metadata: null, + codexOauthSelection: { accountId: "work", explicit: true }, + routeConfig: { routePriority: ["direct"], routeOverrides: {} }, + }; + const options = { model: "openai:gpt-5.5", agentId: "exec", skipAiSettingsPersistence: true }; + expect( + ( + await workspaceService.sendMessage("test-workspace", "Workflow completed", options, { + synthetic: true, + agentInitiated: true, + requireIdle: true, + modelRoutingSnapshot, + }) + ).success + ).toBe(true); + expect(fakeSession.sendMessage).toHaveBeenCalledWith( + "Workflow completed", + options, + expect.objectContaining({ modelRoutingSnapshot }) + ); + expect(JSON.stringify(fakeSession.sendMessage.mock.calls[0]?.[1])).not.toContain( + "snapshot-secret" + ); + expect( + (await workspaceService.sendMessage("test-workspace", "New user turn", options)).success + ).toBe(true); + expect(fakeSession.sendMessage.mock.calls[1]?.[2]).toEqual( + expect.objectContaining({ modelRoutingSnapshot: undefined }) + ); + }); + test("delegates manual pricing rejections to AgentSession so user input is preserved", async () => { fakeSession.isBusy.mockReturnValue(false); const pricingError: SendMessageError = { type: "unknown", raw: "unpriced model" }; @@ -10138,6 +10182,7 @@ describe("WorkspaceService pending auto-title", () => { expect(metadata?.title).toBe("Harden auth flow"); expect(metadata?.pendingAutoTitle).toBeUndefined(); expect(generateIdentitySpy.mock.calls[0]?.[0]).toBe("Continue with auth hardening"); + expect(generateIdentitySpy.mock.calls[0]?.[5]).toEqual({ workspaceId }); } finally { generateIdentitySpy.mockRestore(); } @@ -17458,6 +17503,7 @@ describe("WorkspaceService regenerateTitle", () => { const call = generateIdentitySpy.mock.calls[0]; expect(call?.[3]).toBeUndefined(); expect(call?.[4]).toBe("Fix CI"); + expect(call?.[5]).toEqual({ workspaceId }); expect(updateTitleSpy).toHaveBeenCalledWith(workspaceId, "Fix CI"); } finally { updateTitleSpy.mockRestore(); diff --git a/src/node/services/workspaceService.ts b/src/node/services/workspaceService.ts index 819ee590112..5e9885770ea 100644 --- a/src/node/services/workspaceService.ts +++ b/src/node/services/workspaceService.ts @@ -7346,7 +7346,14 @@ export class WorkspaceService extends EventEmitter implements WorkspaceHost { try { const candidates = await this.getWorkspaceTitleModelCandidates(workspaceId); - const result = await generateWorkspaceIdentity(trimmedMessage, candidates, this.aiService); + const result = await generateWorkspaceIdentity( + trimmedMessage, + candidates, + this.aiService, + undefined, + undefined, + { workspaceId } + ); if (result.success) { const persistResult = await this.updateWorkspaceTitleState(workspaceId, { title: result.data.title, @@ -7711,7 +7718,8 @@ export class WorkspaceService extends EventEmitter implements WorkspaceHost { candidates, this.aiService, conversationContext, - latestUserText + latestUserText, + { workspaceId } ); if (!result.success) { return Err("Title generation failed"); @@ -11210,6 +11218,7 @@ export class WorkspaceService extends EventEmitter implements WorkspaceHost { // paths never fire the callback; the scoped disposal releases on return. const result = await session.sendMessage(message, continuationSendState.options, { onTurnAdmissionCommitted: () => sessionInvisiblePreflight.release(), + modelRoutingSnapshot: internal?.modelRoutingSnapshot, synthetic: internal?.synthetic, agentInitiated: internal?.agentInitiated, goalKind: internal?.goalKind, diff --git a/src/node/services/workspaceStatusGenerator.test.ts b/src/node/services/workspaceStatusGenerator.test.ts index 4ff982f307e..87aacc91e47 100644 --- a/src/node/services/workspaceStatusGenerator.test.ts +++ b/src/node/services/workspaceStatusGenerator.test.ts @@ -1,4 +1,9 @@ -import { describe, expect, test } from "bun:test"; +import * as aiSdk from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { afterEach, describe, expect, mock, spyOn, test } from "bun:test"; +import { Err, Ok } from "@/common/types/result"; +import type { AIService } from "./aiService"; +import type { ModelRoutingSnapshot } from "./modelRoutingSnapshot"; import { buildWorkspaceStatusPrompt, generateWorkspaceStatus } from "./workspaceStatusGenerator"; describe("buildWorkspaceStatusPrompt", () => { @@ -56,6 +61,99 @@ describe("buildWorkspaceStatusPrompt", () => { }); describe("generateWorkspaceStatus error paths", () => { + afterEach(() => mock.restore()); + + function createRoutingSnapshot(): ModelRoutingSnapshot { + return { + providersConfig: { openai: { apiKey: "test-key", codexOauthDefaultAuth: "oauth" } }, + routeConfig: { routePriority: ["direct"], routeOverrides: {} }, + metadata: null, + codexOauthSelection: { accountId: "work", explicit: true }, + }; + } + test("preserves workspace account context across candidate failures", async () => { + const createModelWithPinnedMetadata = mock(() => + Promise.resolve(Err({ type: "oauth_not_connected", provider: "openai" })) + ); + const modelRoutingSnapshot = createRoutingSnapshot(); + const captureModelRoutingSnapshot = mock(() => modelRoutingSnapshot); + const aiService = { + createModelWithPinnedMetadata, + captureModelRoutingSnapshot, + } as unknown as AIService; + const result = await generateWorkspaceStatus( + "Run tests", + ["openai:gpt-5.5", "openai:gpt-5.3-codex"], + aiService, + { workspaceId: "status-workspace" } + ); + expect(result.success).toBe(false); + expect(createModelWithPinnedMetadata).toHaveBeenCalledTimes(2); + for (const call of createModelWithPinnedMetadata.mock.calls) { + expect(call[1]).toEqual({ + workspaceId: "status-workspace", + agentInitiated: true, + modelRoutingSnapshot, + }); + } + }); + + test.each(["account", "preference"] as const)( + "pins status routing across a provider failure and %s change", + async (change) => { + const settings = createRoutingSnapshot(); + const captureModelRoutingSnapshot = mock(() => + structuredClone(settings) + ); + const createModelWithPinnedMetadata = mock( + (model) => Promise.resolve(Ok({ model: new MockLanguageModelV3(), metadataModel: model })) + ); + const aiService = { + captureModelRoutingSnapshot, + createModelWithPinnedMetadata, + } as unknown as AIService; + const stream = spyOn(aiSdk, "streamText").mockReturnValue({ + toolResults: Promise.resolve([ + { + dynamic: false, + toolName: "propose_status", + output: { emoji: "", message: "Run tests" }, + }, + ]), + } as unknown as ReturnType); + stream.mockImplementationOnce(() => { + if (change === "account") settings.codexOauthSelection.accountId = "personal"; + else settings.providersConfig.openai!.codexOauthDefaultAuth = "apiKey"; + throw new Error("First status provider fails"); + }); + const generate = () => + generateWorkspaceStatus( + "Run tests", + ["openai:gpt-5.5", "openai:gpt-5.3-codex"], + aiService, + { workspaceId: "status-workspace" } + ); + expect((await generate()).success).toBe(true); + expect(captureModelRoutingSnapshot).toHaveBeenCalledTimes(1); + expect(captureModelRoutingSnapshot).toHaveBeenCalledWith("status-workspace"); + expect(createModelWithPinnedMetadata).toHaveBeenCalledTimes(2); + const original = createModelWithPinnedMetadata.mock.calls[0]?.[1]?.modelRoutingSnapshot; + expect(original?.codexOauthSelection.accountId).toBe("work"); + expect(original?.providersConfig.openai?.codexOauthDefaultAuth).toBe("oauth"); + expect(createModelWithPinnedMetadata.mock.calls[1]?.[1]?.modelRoutingSnapshot).toBe(original); + + expect((await generate()).success).toBe(true); + expect(captureModelRoutingSnapshot).toHaveBeenCalledTimes(2); + expect(createModelWithPinnedMetadata).toHaveBeenCalledTimes(3); + const next = createModelWithPinnedMetadata.mock.calls[2]?.[1]?.modelRoutingSnapshot; + expect(next).not.toBe(original); + expect(next?.codexOauthSelection.accountId).toBe(change === "account" ? "personal" : "work"); + expect(next?.providersConfig.openai?.codexOauthDefaultAuth).toBe( + change === "preference" ? "apiKey" : "oauth" + ); + } + ); + test("returns a configuration error when no candidates are provided", async () => { const fakeAiService = { // Asserting this never gets called is the real point of this test — diff --git a/src/node/services/workspaceStatusGenerator.ts b/src/node/services/workspaceStatusGenerator.ts index 2d51b32bd07..21ec202fbe8 100644 --- a/src/node/services/workspaceStatusGenerator.ts +++ b/src/node/services/workspaceStatusGenerator.ts @@ -121,6 +121,7 @@ export function buildWorkspaceStatusPrompt( } export interface GenerateWorkspaceStatusOptions extends BuildWorkspaceStatusPromptOptions { + workspaceId?: string; /** * Best-effort cost telemetry: status generation bypasses StreamManager, * so the caller records the successful candidate's usage into @@ -188,6 +189,8 @@ function generateWorkspaceStatusEffect( }); } + // Retries belong to one operation. Keep its billing identity when settings change. + const modelRoutingSnapshot = aiService.captureModelRoutingSnapshot(options.workspaceId); const maxAttempts = Math.min(candidates.length, 3); let lastError: NameGenerationError | null = null; // Track whether any candidate's createModel call succeeded — i.e., whether @@ -208,6 +211,8 @@ function generateWorkspaceStatusEffect( const modelResult = yield* Effect.promise(async () => aiService.createModelWithPinnedMetadata(modelString, { agentInitiated: true, + workspaceId: options.workspaceId, + modelRoutingSnapshot, }) ); if (!modelResult.success) { diff --git a/src/node/services/workspaceTitleGenerator.test.ts b/src/node/services/workspaceTitleGenerator.test.ts index 089a6b08b2f..6908c444017 100644 --- a/src/node/services/workspaceTitleGenerator.test.ts +++ b/src/node/services/workspaceTitleGenerator.test.ts @@ -7,8 +7,9 @@ import { mapModelCreationError, mapNameGenerationError, } from "./workspaceTitleGenerator"; -import { Ok } from "@/common/types/result"; +import { Err, Ok } from "@/common/types/result"; import type { AIService } from "./aiService"; +import type { ModelRoutingSnapshot } from "./modelRoutingSnapshot"; import { attachLanguageModelCleanup } from "./languageModelCleanup"; afterEach(() => { @@ -67,6 +68,15 @@ const createApiCallError = ( responseBody: overrides?.responseBody, }); +function createRoutingSnapshot(): ModelRoutingSnapshot { + return { + providersConfig: { openai: { apiKey: "test-key", codexOauthDefaultAuth: "oauth" } }, + routeConfig: { routePriority: ["direct"], routeOverrides: {} }, + metadata: null, + codexOauthSelection: { accountId: "work", explicit: true }, + }; +} + describe("generateWorkspaceIdentity cleanup", () => { function createTitleModel(modelId = "title-model"): LanguageModel { return { @@ -80,9 +90,97 @@ describe("generateWorkspaceIdentity cleanup", () => { } function createTitleAIService(model: LanguageModel): AIService { - return { createModel: () => Promise.resolve(Ok(model)) } as unknown as AIService; + return { + captureModelRoutingSnapshot: createRoutingSnapshot, + createModel: () => Promise.resolve(Ok(model)), + } as unknown as AIService; } + test.each([{ workspaceId: "title-workspace" }, { projectPath: "/new-project" }])( + "preserves title account context across candidate failures: %j", + async (context) => { + const createModel = mock(() => + Promise.resolve(Err({ type: "oauth_not_connected", provider: "openai" })) + ); + const modelRoutingSnapshot = createRoutingSnapshot(); + const captureModelRoutingSnapshot = mock(() => modelRoutingSnapshot); + const aiService = { createModel, captureModelRoutingSnapshot } as unknown as AIService; + const result = await generateWorkspaceIdentity( + "Add setting", + ["openai:gpt-5.5", "openai:gpt-5.3-codex"], + aiService, + undefined, + undefined, + context + ); + expect(result.success).toBe(false); + expect(createModel).toHaveBeenCalledTimes(2); + for (const call of createModel.mock.calls) { + expect(call[2]).toEqual({ ...context, agentInitiated: true, modelRoutingSnapshot }); + } + } + ); + + test.each([ + { context: { workspaceId: "title-workspace" }, change: "account" }, + { context: { workspaceId: "title-workspace" }, change: "preference" }, + { context: { projectPath: "/new-project" }, change: "account" }, + { context: { projectPath: "/new-project" }, change: "preference" }, + ] as const)("pins title routing across provider failures: %j", async ({ context, change }) => { + const settings = createRoutingSnapshot(); + const captureModelRoutingSnapshot = mock(() => + structuredClone(settings) + ); + const createModel = mock(() => + Promise.resolve(Ok(createTitleModel())) + ); + const aiService = { captureModelRoutingSnapshot, createModel } as unknown as AIService; + const stream = spyOn(aiSdk, "streamText").mockReturnValue({ + toolResults: Promise.resolve([ + { + dynamic: false, + toolName: "propose_name", + output: { name: "settings", title: "Add setting" }, + }, + ]), + } as unknown as ReturnType); + stream.mockImplementationOnce(() => { + if (change === "account") settings.codexOauthSelection.accountId = "personal"; + else settings.providersConfig.openai!.codexOauthDefaultAuth = "apiKey"; + throw new Error("First title provider fails"); + }); + const generate = () => + generateWorkspaceIdentity( + "Add setting", + ["openai:gpt-5.5", "openai:gpt-5.3-codex"], + aiService, + undefined, + undefined, + context + ); + expect((await generate()).success).toBe(true); + expect(captureModelRoutingSnapshot).toHaveBeenCalledTimes(1); + expect(captureModelRoutingSnapshot).toHaveBeenCalledWith( + "workspaceId" in context ? context.workspaceId : undefined, + "projectPath" in context ? context.projectPath : undefined + ); + expect(createModel).toHaveBeenCalledTimes(2); + const original = createModel.mock.calls[0]?.[2]?.modelRoutingSnapshot; + expect(original?.codexOauthSelection.accountId).toBe("work"); + expect(original?.providersConfig.openai?.codexOauthDefaultAuth).toBe("oauth"); + expect(createModel.mock.calls[1]?.[2]?.modelRoutingSnapshot).toBe(original); + + expect((await generate()).success).toBe(true); + expect(captureModelRoutingSnapshot).toHaveBeenCalledTimes(2); + expect(createModel).toHaveBeenCalledTimes(3); + const next = createModel.mock.calls[2]?.[2]?.modelRoutingSnapshot; + expect(next).not.toBe(original); + expect(next?.codexOauthSelection.accountId).toBe(change === "account" ? "personal" : "work"); + expect(next?.providersConfig.openai?.codexOauthDefaultAuth).toBe( + change === "preference" ? "apiKey" : "oauth" + ); + }); + test("cleans up the model after a successful title stream", async () => { let cleanupCalls = 0; const model = createTitleModel(); @@ -145,6 +243,7 @@ describe("generateWorkspaceIdentity cleanup", () => { secondCleanupCalls += 1; }); const aiService = { + captureModelRoutingSnapshot: createRoutingSnapshot, createModel: mock((modelString: string) => Promise.resolve(Ok(modelString.includes("first") ? firstModel : secondModel)) ), diff --git a/src/node/services/workspaceTitleGenerator.ts b/src/node/services/workspaceTitleGenerator.ts index 0cea63887d3..f946f2a16a3 100644 --- a/src/node/services/workspaceTitleGenerator.ts +++ b/src/node/services/workspaceTitleGenerator.ts @@ -185,12 +185,19 @@ export async function generateWorkspaceIdentity( /** Optional conversation turns context used for regenerate-title prompts. */ conversationContext?: string, /** Optional most recent user message; included as additional context only — not given precedence over older turns. */ - latestUserMessage?: string + latestUserMessage?: string, + context?: { workspaceId?: string; projectPath?: string } ): Promise> { if (candidates.length === 0) { return Err({ type: "unknown", raw: "No model candidates provided for name generation" }); } + // Retries belong to one operation. Keep its billing identity when settings change. + const modelRoutingSnapshot = aiService.captureModelRoutingSnapshot( + context?.workspaceId, + context?.projectPath + ); + // Try up to 3 candidates const maxAttempts = Math.min(candidates.length, 3); @@ -201,7 +208,9 @@ export async function generateWorkspaceIdentity( const modelString = candidates[i]; const modelResult = await aiService.createModel(modelString, undefined, { + ...context, agentInitiated: true, + modelRoutingSnapshot, }); if (!modelResult.success) { lastError = mapModelCreationError(modelResult.error, modelString); diff --git a/src/node/utils/codexOauthAuth.test.ts b/src/node/utils/codexOauthAuth.test.ts index 614461e1b1f..1906f5972e2 100644 --- a/src/node/utils/codexOauthAuth.test.ts +++ b/src/node/utils/codexOauthAuth.test.ts @@ -2,6 +2,9 @@ import { describe, it, expect } from "bun:test"; import { parseCodexOauthAuth, + getCodexOauthAccounts, + getCodexOauthAccountId, + getCodexOauthAuth, isCodexOauthAuthExpired, parseJwtClaims, extractAccountIdFromClaims, @@ -48,6 +51,63 @@ describe("parseCodexOauthAuth", () => { expect(result).toEqual(input); }); + it("preserves credentials with missing or malformed login IDs", () => { + const auth = { type: "oauth" as const, access: "access", refresh: "refresh", expires: 1000 }; + expect(parseCodexOauthAuth(auth)).not.toBeNull(); + const credentialId = "1c9c50b0-d777-4dd2-998c-09c156ba9754"; + expect(parseCodexOauthAuth({ ...auth, credentialId })?.credentialId).toBe(credentialId); + for (const invalid of [null, "", "not-a-uuid", 42]) { + const stored = { ...auth, credentialId: invalid }; + expect(parseCodexOauthAuth(stored)).toEqual(auth); + expect( + getCodexOauthAccounts({ + codexOauthAccounts: { work: { label: "Work", credentials: stored } }, + }) + ).toHaveLength(1); + expect(stored.credentialId).toBe(invalid); + } + }); + + it("retains only a legacy alias that matches the current credential ID", () => { + const credentialId = "1c9c50b0-d777-4dd2-998c-09c156ba9754"; + const auth = { + type: "oauth", + access: "access", + refresh: "refresh", + expires: 1000, + credentialId, + }; + for (const legacyCredentialId of [ + credentialId, + undefined, + null, + "", + 42, + "not-a-uuid", + "50e00a32-b964-4ce2-b131-6b53356ce2db", + ]) { + const parsed = parseCodexOauthAuth({ ...auth, legacyCredentialId }); + expect(parsed?.access).toBe(auth.access); + expect(parsed?.legacyCredentialId).toBe( + legacyCredentialId === credentialId ? credentialId : undefined + ); + } + expect( + parseCodexOauthAuth({ ...auth, credentialId: undefined, legacyCredentialId: credentialId }) + ?.legacyCredentialId + ).toBeUndefined(); + }); + + it("preserves only the supported invalid credential marker", () => { + const auth = { type: "oauth" as const, access: "access", refresh: "refresh", expires: 1000 }; + const marked = parseCodexOauthAuth({ ...auth, invalidReason: "invalid_grant" }); + expect(marked).toEqual({ ...auth, invalidReason: "invalid_grant" }); + expect(getCodexOauthAccounts({ codexOauth: marked })).toHaveLength(1); + for (const invalidReason of [null, "other", 42]) { + expect(parseCodexOauthAuth({ ...auth, invalidReason })).toBeNull(); + } + }); + it("returns null for non-object values", () => { expect(parseCodexOauthAuth(null)).toBeNull(); expect(parseCodexOauthAuth(undefined)).toBeNull(); @@ -238,3 +298,116 @@ describe("extractAccountIdFromTokens", () => { expect(extractAccountIdFromTokens({ accessToken, idToken })).toBe("from_access_token"); }); }); + +describe("Codex OAuth account slots", () => { + const legacy = { type: "oauth", access: "legacy", refresh: "legacy-refresh", expires: 1000 }; + const work = { + type: "oauth", + access: "work", + refresh: "work-refresh", + expires: 2000, + accountId: "remote-chatgpt-id", + }; + + it("reads the legacy slot and named slots with separate local identities", () => { + const config = { + codexOauth: legacy, + codexOauthLabel: "Personal", + codexOauthAccounts: { work: { label: "Work", credentials: work } }, + }; + expect(getCodexOauthAccounts(config).map(({ id, label }) => ({ id, label }))).toEqual([ + { id: "default", label: "Personal" }, + { id: "work", label: "Work" }, + ]); + expect(getCodexOauthAuth(config, "work")?.accountId).toBe("remote-chatgpt-id"); + expect(getCodexOauthAuth(config, "remote-chatgpt-id")).toBeNull(); + expect(getCodexOauthAuth(config)?.access).toBe("legacy"); + }); + + it("uses explicit selection before the global default and never substitutes missing slots", () => { + const config = { + codexOauth: legacy, + codexOauthAccounts: { work: { label: "Work", credentials: work } }, + codexOauthDefaultAccountId: "work", + }; + expect(getCodexOauthAuth(config)?.access).toBe("work"); + expect(getCodexOauthAuth(config, "default")?.access).toBe("legacy"); + expect(getCodexOauthAuth(config, "missing")).toBeNull(); + expect(getCodexOauthAuth({ ...config, codexOauthDefaultAccountId: "missing" })).toBeNull(); + expect(getCodexOauthAccountId(undefined)).toBe("default"); + expect(getCodexOauthAccountId(config, "missing")).toBe("missing"); + }); + + it.each([ + { damage: "missing", label: undefined }, + { damage: "null", label: null }, + { damage: "number", label: 42 }, + { damage: "boolean", label: false }, + { damage: "object", label: {} }, + { damage: "array", label: [] }, + { damage: "empty", label: "" }, + { damage: "whitespace", label: " \t\n" }, + ])("preserves named credentials when the label is $damage", ({ label }) => { + const auth = { ...work, access: "damaged-label-access" }; + const config = { + codexOauth: legacy, + codexOauthDefaultAccountId: "damaged", + codexOauthAccounts: { + damaged: { label, credentials: auth }, + work: { label: " Work ", credentials: work }, + }, + }; + expect(getCodexOauthAccounts(config).map(({ id, label }) => ({ id, label }))).toEqual([ + { id: "default", label: "Default" }, + { id: "damaged", label: "damaged" }, + { id: "work", label: "Work" }, + ]); + expect(getCodexOauthAuth(config)).toMatchObject(auth); + expect(getCodexOauthAuth(config, "damaged")).toMatchObject(auth); + }); + + it("rejects malformed named credentials regardless of label validity", () => { + const config = { + codexOauthAccounts: { + missingAuth: {}, + invalidType: { label: "Valid label", credentials: { ...work, type: "apiKey" } }, + invalidAccess: { label: 42, credentials: { ...work, access: null } }, + invalidRefresh: { credentials: { ...work, refresh: "" } }, + invalidExpiry: { label: " ", credentials: { ...work, expires: Infinity } }, + work: { label: "Work", credentials: work }, + }, + }; + expect(getCodexOauthAccounts(config).map(({ id }) => id)).toEqual(["work"]); + for (const id of Object.keys(config.codexOauthAccounts)) { + if (id !== "work") expect(getCodexOauthAuth(config, id)).toBeNull(); + } + }); + + it("excludes stored IDs that cannot pass account mutation validation", () => { + const invalidIds = ["", "__proto__", "constructor", "prototype", "../bad", "x".repeat(201)]; + const accounts = Object.fromEntries( + invalidIds.map((id) => [id, { label: "Invalid", credentials: work }]) + ); + const config = { + codexOauthAccounts: { ...accounts, work: { label: "Work", credentials: work } }, + }; + expect(getCodexOauthAccounts(config).map((account) => account.id)).toEqual(["work"]); + for (const id of invalidIds) expect(getCodexOauthAuth(config, id)).toBeNull(); + expect(getCodexOauthAuth({ ...config, codexOauthDefaultAccountId: "../bad" })).toBeNull(); + }); + + it("filters malformed slots without accepting a duplicate legacy slot", () => { + const config = { + codexOauth: legacy, + codexOauthAccounts: { + default: { label: "Duplicate", credentials: work }, + broken: { label: "Broken", credentials: {} }, + blank: { label: " ", credentials: null }, + work: { label: "Work", credentials: work }, + }, + }; + expect(getCodexOauthAccounts(config).map((account) => account.id)).toEqual(["default", "work"]); + expect(getCodexOauthAccounts(null)).toEqual([]); + expect(getCodexOauthAccounts({ codexOauthAccounts: [] })).toEqual([]); + }); +}); diff --git a/src/node/utils/codexOauthAuth.ts b/src/node/utils/codexOauthAuth.ts index 5f188a1984e..a7289159309 100644 --- a/src/node/utils/codexOauthAuth.ts +++ b/src/node/utils/codexOauthAuth.ts @@ -5,8 +5,25 @@ * extract non-sensitive claims (e.g. ChatGPT-Account-Id) from OAuth responses. */ +import { z } from "zod"; + +import { + CODEX_OAUTH_DEFAULT_ACCOUNT_ID, + CODEX_OAUTH_ACCOUNT_ID_MAX_LENGTH, + CODEX_OAUTH_ACCOUNT_ID_PATTERN, + CODEX_OAUTH_RESERVED_ACCOUNT_IDS, +} from "@/common/constants/codexOauthAccounts"; + +const credentialIdSchema = z.string().uuid().optional(); + export interface CodexOauthAuth { type: "oauth"; + /** Identifies this login across token rotations and processes. */ + credentialId?: string; + /** Matches a backfilled ID while requests still pin the original undefined identity. */ + legacyCredentialId?: string; + /** Blocks requests while retaining the login identity for reconnect. */ + invalidReason?: "invalid_grant"; /** OAuth access token (JWT). */ access: string; /** OAuth refresh token. */ @@ -35,6 +52,9 @@ export function parseCodexOauthAuth(value: unknown): CodexOauthAuth | null { const refresh = value.refresh; const expires = value.expires; const accountId = value.accountId; + const credentialId = credentialIdSchema.safeParse(value.credentialId); + const legacyCredentialId = credentialIdSchema.safeParse(value.legacyCredentialId); + const invalidReason = value.invalidReason; if (type !== "oauth") return null; if (typeof access !== "string" || !access) return null; @@ -45,7 +65,87 @@ export function parseCodexOauthAuth(value: unknown): CodexOauthAuth | null { if (typeof accountId !== "string" || !accountId) return null; } - return { type: "oauth", access, refresh, expires, accountId }; + if (invalidReason !== undefined && invalidReason !== "invalid_grant") return null; + + return { + type: "oauth", + access, + refresh, + expires, + accountId, + // Treat a damaged optional ID as legacy state so reconnect can assign a valid ID. + credentialId: credentialId.success ? credentialId.data : undefined, + // A mismatched or malformed alias must not authorize an old request. + legacyCredentialId: + credentialId.success && + legacyCredentialId.success && + credentialId.data === legacyCredentialId.data + ? legacyCredentialId.data + : undefined, + invalidReason, + }; +} + +/** Validate local slot IDs at input and storage boundaries. */ +export function isValidCodexOauthAccountId(accountId: string): boolean { + return ( + accountId.length <= CODEX_OAUTH_ACCOUNT_ID_MAX_LENGTH && + CODEX_OAUTH_ACCOUNT_ID_PATTERN.test(accountId) && + !CODEX_OAUTH_RESERVED_ACCOUNT_IDS.has(accountId) + ); +} + +/** Read stored slots, including invalid credentials that need reconnect. */ +export function getCodexOauthAccounts(config: unknown): Array<{ + id: string; + label: string; + auth: CodexOauthAuth; +}> { + if (!isPlainObject(config)) return []; + const accounts: Array<{ id: string; label: string; auth: CodexOauthAuth }> = []; + const legacy = parseCodexOauthAuth(config.codexOauth); + if (legacy) { + accounts.push({ + id: CODEX_OAUTH_DEFAULT_ACCOUNT_ID, + label: + typeof config.codexOauthLabel === "string" && config.codexOauthLabel.trim() + ? config.codexOauthLabel.trim() + : "Default", + auth: legacy, + }); + } + if (isPlainObject(config.codexOauthAccounts)) { + for (const [id, entry] of Object.entries(config.codexOauthAccounts)) { + // The legacy slot owns this ID, even if malformed disk data repeats it. + if ( + id === CODEX_OAUTH_DEFAULT_ACCOUNT_ID || + !isValidCodexOauthAccountId(id) || + !isPlainObject(entry) + ) + continue; + const auth = parseCodexOauthAuth(entry.credentials); + if (!auth) continue; + // Damaged display labels must not hide credentials from reconnect, rename, or disconnect. + const label = typeof entry.label === "string" ? entry.label.trim() : ""; + accounts.push({ id, label: label || id, auth }); + } + } + return accounts; +} + +/** Resolve a local slot ID without substituting another connected account. */ +export function getCodexOauthAccountId(config: unknown, override?: string): string { + if (override !== undefined) return override; + if (isPlainObject(config) && typeof config.codexOauthDefaultAccountId === "string") { + return config.codexOauthDefaultAccountId; + } + return CODEX_OAUTH_DEFAULT_ACCOUNT_ID; +} + +/** Read the selected slot. Missing selections do not fall back to another slot. */ +export function getCodexOauthAuth(config: unknown, accountId?: string): CodexOauthAuth | null { + const selected = getCodexOauthAccountId(config, accountId); + return getCodexOauthAccounts(config).find((account) => account.id === selected)?.auth ?? null; } export function isCodexOauthAuthExpired( diff --git a/src/node/utils/providerRequirements.test.ts b/src/node/utils/providerRequirements.test.ts index 6fb8ce9c874..db4854c671f 100644 --- a/src/node/utils/providerRequirements.test.ts +++ b/src/node/utils/providerRequirements.test.ts @@ -76,6 +76,27 @@ describe("hasAnyConfiguredProvider", () => { expect(hasAnyConfiguredProvider(providers)).toBe(true); }); + it("accepts an additional OAuth account even when the global selection is missing", () => { + expect( + hasAnyConfiguredProvider({ + openai: { + codexOauthDefaultAccountId: "deleted", + codexOauthAccounts: { + work: { + label: "Work", + credentials: { + type: "oauth", + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + }, + }, + }, + }, + }) + ).toBe(true); + }); + it("returns true for OpenAI Codex OAuth-only configuration", () => { const providers: ProvidersConfig = { openai: { @@ -92,6 +113,49 @@ describe("hasAnyConfiguredProvider", () => { expect(hasAnyConfiguredProvider(providers)).toBe(true); }); + it.each(["default", "work"])("does not count a revoked %s account as configured", (accountId) => { + const savedKey = process.env.OPENAI_API_KEY; + delete process.env.OPENAI_API_KEY; + try { + const auth = { + type: "oauth" as const, + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + invalidReason: "invalid_grant" as const, + }; + const openai = + accountId === "default" + ? { codexOauth: auth } + : { codexOauthAccounts: { work: { label: "Work", credentials: auth } } }; + expect(hasAnyConfiguredProvider({ openai })).toBe(false); + expect(hasAnyConfiguredProvider({ openai, openrouter: { apiKey: "or-test" } })).toBe(true); + } finally { + if (savedKey !== undefined) process.env.OPENAI_API_KEY = savedKey; + } + }); + + it.each(["default", "work"])( + "does not count a disabled %s account as configured", + (accountId) => { + const auth = { + type: "oauth" as const, + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + }; + const openai = { + enabled: false, + ...(accountId === "default" + ? { codexOauth: auth } + : { codexOauthAccounts: { work: { label: "Work", credentials: auth } } }), + }; + expect(hasAnyConfiguredProvider({ openai })).toBe(false); + expect(hasAnyConfiguredProvider({ openai, openrouter: { apiKey: "or-test" } })).toBe(true); + expect(hasAnyConfiguredProvider({ openai: { ...openai, enabled: true } })).toBe(true); + } + ); + it("returns true for keyless providers with explicit config", () => { const providers: ProvidersConfig = { ollama: { diff --git a/src/node/utils/providerRequirements.ts b/src/node/utils/providerRequirements.ts index af26dccbdd0..e05c332d204 100644 --- a/src/node/utils/providerRequirements.ts +++ b/src/node/utils/providerRequirements.ts @@ -23,7 +23,7 @@ import type { OpenAIProviderConfig, } from "@/common/config/schemas/providersConfig"; import type { ProviderConfig, ProvidersConfig } from "@/node/config"; -import { parseCodexOauthAuth } from "@/node/utils/codexOauthAuth"; +import { getCodexOauthAccounts } from "@/node/utils/codexOauthAuth"; import { parseCoderOauthAuth } from "@/node/utils/coderOauthAuth"; import { normalizeCoderDeploymentUrl } from "@/common/constants/coderOAuth"; @@ -608,20 +608,21 @@ export function hasAnyConfiguredProvider(providers: ProvidersConfig | null | und continue; } + // Disabled providers cannot satisfy the CLI startup credential check. + if (isProviderDisabledInConfig(rawConfig)) { + continue; + } + // OpenAI Codex OAuth is a valid credential path even without apiKey. if ( providerKey === "openai" && - parseCodexOauthAuth((rawConfig as { codexOauth?: unknown }).codexOauth) !== null + getCodexOauthAccounts(rawConfig).some(({ auth }) => auth.invalidReason === undefined) ) { return true; } if (!(providerKey in PROVIDER_DEFINITIONS)) { - if ( - isCustomProviderConfig(rawConfig) && - !isProviderDisabledInConfig(rawConfig) && - resolveConfigBaseUrl(rawConfig) !== undefined - ) { + if (isCustomProviderConfig(rawConfig) && resolveConfigBaseUrl(rawConfig) !== undefined) { return true; } diff --git a/tests/ui/agents/thinkingSelector.test.ts b/tests/ui/agents/thinkingSelector.test.ts index ed9a6108285..d966dc9c8e7 100644 --- a/tests/ui/agents/thinkingSelector.test.ts +++ b/tests/ui/agents/thinkingSelector.test.ts @@ -95,7 +95,11 @@ describeIntegration("Thinking selector", () => { const harness = await createAppHarness({ branchPrefix: "thinking-selector", beforeRenderEnvironment: async (env) => { - await setupProviders(env, { xai: { apiKey: "dummy" } }); + // Pro and Fast controls require a usable direct API route. + await setupProviders(env, { + openai: { apiKey: "dummy" }, + xai: { apiKey: "dummy" }, + }); }, });