Skip to content

Commit 7e63893

Browse files
ndycodendycode
committed
Fix Codex TUI goal proxy handling
1 parent 3f335e4 commit 7e63893

6 files changed

Lines changed: 329 additions & 14 deletions

File tree

‎lib/runtime-rotation-proxy.ts‎

Lines changed: 84 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,12 @@ const ALLOWED_MODELS_PATHS = new Set([
141141
URL_PATHS.MODELS,
142142
`/v1${URL_PATHS.MODELS}`,
143143
]);
144+
const ALLOWED_THREAD_GOAL_PATHS = new Set([
145+
"/thread/goal/get",
146+
"/thread/goal/set",
147+
"/codex/thread/goal/get",
148+
"/codex/thread/goal/set",
149+
]);
144150

145151
function isResponsesPath(pathname: string): boolean {
146152
return ALLOWED_RESPONSES_PATHS.has(pathname);
@@ -150,6 +156,14 @@ function isModelsPath(pathname: string): boolean {
150156
return ALLOWED_MODELS_PATHS.has(pathname);
151157
}
152158

159+
function isThreadGoalPath(pathname: string): boolean {
160+
return ALLOWED_THREAD_GOAL_PATHS.has(pathname);
161+
}
162+
163+
function normalizeThreadGoalUpstreamPath(pathname: string): string {
164+
return pathname.startsWith("/codex/") ? pathname : `/codex${pathname}`;
165+
}
166+
153167
function headersFromIncoming(req: IncomingMessage): Headers {
154168
const headers = new Headers();
155169
for (const [key, value] of Object.entries(req.headers)) {
@@ -383,6 +397,30 @@ function buildModelsRequestContext(req: IncomingMessage): RequestContext {
383397
};
384398
}
385399

400+
function buildThreadGoalRequestContext(
401+
req: IncomingMessage,
402+
body: Buffer,
403+
pathname: string,
404+
): RequestContext {
405+
const headers = headersFromIncoming(req);
406+
const parsedBody = parseRequestBody(body);
407+
const sessionKey = parsedBody
408+
? (readStringRecordValue(parsedBody, "thread_id") ??
409+
readStringRecordValue(parsedBody, "threadId") ??
410+
resolveSessionKey(headers, parsedBody))
411+
: resolveSessionKey(headers, null);
412+
return {
413+
body,
414+
headers,
415+
method: req.method === "GET" ? "GET" : "POST",
416+
upstreamPath: normalizeThreadGoalUpstreamPath(pathname),
417+
model: null,
418+
family: "codex",
419+
stream: false,
420+
sessionKey,
421+
};
422+
}
423+
386424
function buildUpstreamUrl(
387425
req: IncomingMessage,
388426
upstreamBaseUrl: string,
@@ -668,7 +706,7 @@ function writeMethodOrPathError(res: ServerResponse): void {
668706
writeJson(res, 404, {
669707
error: {
670708
message:
671-
"Runtime rotation proxy only accepts Responses API and model discovery requests.",
709+
"Runtime rotation proxy only accepts Responses API, model discovery, and Codex thread goal requests.",
672710
code: "runtime_rotation_proxy_not_found",
673711
},
674712
});
@@ -832,6 +870,7 @@ export async function startRuntimeRotationProxy(
832870
lastAccountId: null,
833871
lastAccountUpdatedAt: null,
834872
};
873+
const threadGoalFallbacks = new Map<string, string | null>();
835874

836875
const handleRequest = async (
837876
req: IncomingMessage,
@@ -844,7 +883,10 @@ export async function startRuntimeRotationProxy(
844883
req.method === "POST" && isResponsesPath(incomingUrl.pathname);
845884
const isModelsRequest =
846885
req.method === "GET" && isModelsPath(incomingUrl.pathname);
847-
if (!isResponsesRequest && !isModelsRequest) {
886+
const isThreadGoalRequest =
887+
(req.method === "GET" || req.method === "POST") &&
888+
isThreadGoalPath(incomingUrl.pathname);
889+
if (!isResponsesRequest && !isModelsRequest && !isThreadGoalRequest) {
848890
writeMethodOrPathError(res);
849891
return;
850892
}
@@ -856,12 +898,15 @@ export async function startRuntimeRotationProxy(
856898
}
857899

858900
status.totalRequests += 1;
901+
const requestBody =
902+
isResponsesRequest || (isThreadGoalRequest && req.method === "POST")
903+
? await readRequestBody(req, maxRequestBodyBytes)
904+
: Buffer.alloc(0);
859905
const context = isModelsRequest
860906
? buildModelsRequestContext(req)
861-
: buildResponsesRequestContext(
862-
req,
863-
await readRequestBody(req, maxRequestBodyBytes),
864-
);
907+
: isThreadGoalRequest
908+
? buildThreadGoalRequestContext(req, requestBody, incomingUrl.pathname)
909+
: buildResponsesRequestContext(req, requestBody);
865910
const requestStartedAt = now();
866911
let policyDecision: RuntimePolicyDecision | null = null;
867912
let projectKey: string | null = null;
@@ -881,7 +926,11 @@ export async function startRuntimeRotationProxy(
881926
}
882927
usageRecorder = createRuntimeUsageRecorder({
883928
source: "runtime-proxy",
884-
operation: isModelsRequest ? "models" : "responses",
929+
operation: isModelsRequest
930+
? "models"
931+
: isThreadGoalRequest
932+
? "diagnostic"
933+
: "responses",
885934
model: context.model,
886935
projectKey,
887936
requestId: context.sessionKey,
@@ -1078,6 +1127,34 @@ export async function startRuntimeRotationProxy(
10781127
continue;
10791128
}
10801129

1130+
if (isThreadGoalRequest && upstream.status >= 400) {
1131+
await readErrorBody(upstream);
1132+
const parsedGoalBody = parseRequestBody(context.body);
1133+
const fallbackKey = context.sessionKey ?? "default";
1134+
const goal =
1135+
typeof parsedGoalBody?.goal === "string" ? parsedGoalBody.goal : null;
1136+
accountManager.recordSuccess(refreshed.account, context.family, context.model);
1137+
await persistRuntimeActiveAccount(
1138+
accountManager,
1139+
refreshed.account,
1140+
context.family,
1141+
);
1142+
await usageRecorder.record({
1143+
outcome: "success",
1144+
statusCode: HTTP_STATUS.OK,
1145+
account: refreshed.account,
1146+
});
1147+
if (context.upstreamPath.endsWith("/set")) {
1148+
threadGoalFallbacks.set(fallbackKey, goal);
1149+
writeJson(res, HTTP_STATUS.OK, { ok: true, goal });
1150+
return;
1151+
}
1152+
writeJson(res, HTTP_STATUS.OK, {
1153+
goal: threadGoalFallbacks.get(fallbackKey) ?? null,
1154+
});
1155+
return;
1156+
}
1157+
10811158
accountManager.recordSuccess(refreshed.account, context.family, context.model);
10821159
const nearExhaustionWaitMs = getQuotaNearExhaustionWaitMs(
10831160
upstream.headers,

‎lib/runtime/config-toml.ts‎

Lines changed: 78 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -75,14 +75,44 @@ export function rewriteTopLevelModelProvider(rawConfig: string): string {
7575
return output.join(lineEnding);
7676
}
7777

78-
function extractTopLevelModelProviderLine(rawConfig: string): string | null {
78+
export function enableTopLevelResponseStorage(rawConfig: string): string {
79+
const lineEnding = rawConfig.includes("\r\n") ? "\r\n" : "\n";
80+
const lines = rawConfig.length > 0 ? rawConfig.split(/\r?\n/) : [];
81+
const output: string[] = [];
82+
let inTopLevel = true;
83+
84+
for (const line of lines) {
85+
if (readTomlTableName(line) !== null) {
86+
inTopLevel = false;
87+
output.push(line);
88+
continue;
89+
}
90+
if (
91+
inTopLevel &&
92+
/^\s*disable_response_storage\s*=\s*true\s*(?:#.*)?$/i.test(line)
93+
) {
94+
output.push("disable_response_storage = false");
95+
continue;
96+
}
97+
output.push(line);
98+
}
99+
100+
return output.join(lineEnding);
101+
}
102+
103+
function extractTopLevelLine(rawConfig: string, key: string): string | null {
104+
const pattern = new RegExp(`^\\s*${key}\\s*=`);
79105
for (const line of rawConfig.split(/\r?\n/)) {
80106
if (readTomlTableName(line) !== null) return null;
81-
if (/^\s*model_provider\s*=/.test(line)) return line;
107+
if (pattern.test(line)) return line;
82108
}
83109
return null;
84110
}
85111

112+
function extractTopLevelModelProviderLine(rawConfig: string): string | null {
113+
return extractTopLevelLine(rawConfig, "model_provider");
114+
}
115+
86116
export function restoreTopLevelModelProvider(
87117
currentConfig: string,
88118
originalConfig: string,
@@ -108,6 +138,43 @@ export function restoreTopLevelModelProvider(
108138
return output.join(lineEnding);
109139
}
110140

141+
export function restoreTopLevelResponseStorage(
142+
currentConfig: string,
143+
originalConfig: string,
144+
): string {
145+
const lineEnding = currentConfig.includes("\r\n") ? "\r\n" : "\n";
146+
const originalLine = extractTopLevelLine(
147+
originalConfig,
148+
"disable_response_storage",
149+
);
150+
if (!originalLine) return currentConfig;
151+
const lines = currentConfig.length > 0 ? currentConfig.split(/\r?\n/) : [];
152+
const output: string[] = [];
153+
let handled = false;
154+
let inTopLevel = true;
155+
156+
for (const line of lines) {
157+
if (readTomlTableName(line) !== null) {
158+
inTopLevel = false;
159+
output.push(line);
160+
continue;
161+
}
162+
if (
163+
!handled &&
164+
inTopLevel &&
165+
/^\s*disable_response_storage\s*=/.test(line) &&
166+
readTomlTableName(line) === null
167+
) {
168+
output.push(originalLine);
169+
handled = true;
170+
continue;
171+
}
172+
output.push(line);
173+
}
174+
175+
return output.join(lineEnding);
176+
}
177+
111178
export function ensureTomlTrailingNewline(value: string): string {
112179
return value.replace(/[\r\n]*$/, "\n");
113180
}
@@ -147,20 +214,27 @@ export function rewriteConfigTomlForRuntimeRotationProvider(
147214
/[\r\n]*$/,
148215
"",
149216
);
217+
const withResponseStorage = enableTopLevelResponseStorage(
218+
withModelProvider,
219+
).replace(/[\r\n]*$/, "");
150220
const providerBlock = createRuntimeRotationProviderBlock(
151221
baseUrl,
152222
clientApiKey,
153223
).join(lineEnding);
154-
return `${withModelProvider}${lineEnding}${lineEnding}${providerBlock}${lineEnding}`;
224+
return `${withResponseStorage}${lineEnding}${lineEnding}${providerBlock}${lineEnding}`;
155225
}
156226

157227
export function restoreConfigTomlFromRuntimeRotationProvider(
158228
currentConfig: string,
159229
originalConfig: string,
160230
): string {
161231
const withoutProvider = removeRuntimeRotationProviderBlock(currentConfig);
232+
const withResponseStorage = restoreTopLevelResponseStorage(
233+
withoutProvider,
234+
originalConfig,
235+
);
162236
return ensureTomlTrailingNewline(
163-
restoreTopLevelModelProvider(withoutProvider, originalConfig).replace(
237+
restoreTopLevelModelProvider(withResponseStorage, originalConfig).replace(
164238
/[\r\n]*$/,
165239
"",
166240
),

‎scripts/codex.js‎

Lines changed: 47 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1902,6 +1902,21 @@ function mirrorFileIntoShadowHome(sourcePath, destinationPath, tightenFile) {
19021902
tightenFile(destinationPath);
19031903
}
19041904

1905+
function shouldMaterializeFileIntoShadowHome(name) {
1906+
return /\.sqlite(?:-(?:shm|wal))?$/i.test(name);
1907+
}
1908+
1909+
function materializeFileIntoShadowHome(sourcePath, destinationPath, tightenFile) {
1910+
try {
1911+
linkSync(sourcePath, destinationPath);
1912+
return;
1913+
} catch {
1914+
// Hard links make SQLite roots look local without copying an active DB.
1915+
}
1916+
copyFileSync(sourcePath, destinationPath);
1917+
tightenFile(destinationPath);
1918+
}
1919+
19051920
function collectShadowHomeSyncFileNames(shadowCodexHome, syncFileNames) {
19061921
try {
19071922
for (const entry of readdirSync(shadowCodexHome, { withFileTypes: true })) {
@@ -2075,6 +2090,7 @@ function createShadowHomeMirror(
20752090
continue;
20762091
}
20772092
const isKnownStateFile = SHADOW_HOME_STATE_FILE_SET.has(name);
2093+
const shouldMaterializeFile = shouldMaterializeFileIntoShadowHome(name);
20782094
const sourcePath = join(originalCodexHome, name);
20792095
const destinationPath = join(shadowCodexHome, name);
20802096
if (existsSync(destinationPath)) {
@@ -2103,6 +2119,8 @@ function createShadowHomeMirror(
21032119
if (isKnownStateFile) {
21042120
copyFileSync(sourcePath, destinationPath);
21052121
tightenFile(destinationPath);
2122+
} else if (shouldMaterializeFile) {
2123+
materializeFileIntoShadowHome(sourcePath, destinationPath, tightenFile);
21062124
} else {
21072125
mirrorFileIntoShadowHome(sourcePath, destinationPath, tightenFile);
21082126
}
@@ -2266,14 +2284,24 @@ function resolveRuntimeRotationProxyOriginalCodexHome(baseEnv) {
22662284
return override || resolveCodexHomeDir(baseEnv);
22672285
}
22682286

2287+
function createRuntimeRotationShadowHome(originalCodexHome) {
2288+
const shadowRoot = join(
2289+
originalCodexHome,
2290+
"multi-auth",
2291+
"runtime-shadow-homes",
2292+
);
2293+
mkdirSync(shadowRoot, { recursive: true });
2294+
return mkdtempSync(join(shadowRoot, "codex-multi-auth-runtime-home-"));
2295+
}
2296+
22692297
function createRuntimeRotationProxyCodexHome(
22702298
baseEnv,
22712299
proxyBaseUrl,
22722300
clientApiKey,
22732301
configTomlModule,
22742302
) {
22752303
const originalCodexHome = resolveRuntimeRotationProxyOriginalCodexHome(baseEnv);
2276-
const shadowCodexHome = mkdtempSync(join(tmpdir(), "codex-multi-auth-runtime-home-"));
2304+
const shadowCodexHome = createRuntimeRotationShadowHome(originalCodexHome);
22772305
let syncShadowHomeStateBack = () => {};
22782306
const cleanup = () => {
22792307
try {
@@ -2841,15 +2869,22 @@ function startRuntimeRotationAppHelper(baseContext) {
28412869
});
28422870
}
28432871

2844-
async function createRuntimeRotationAppHelperContext(baseContext, configTomlModule) {
2872+
async function createRuntimeRotationAppHelperContext(
2873+
baseContext,
2874+
configTomlModule,
2875+
options = {},
2876+
) {
28452877
const startedAt = Date.now();
28462878
const { helper, message } = await startRuntimeRotationAppHelper(baseContext);
28472879
const helperEnv = message.env ?? {};
28482880
const detachGraceMs = resolveRuntimeRotationAppHelperDetachGraceMs(baseContext.env);
28492881

28502882
const cleanup = async ({ exitCode } = {}) => {
28512883
const livedMs = Date.now() - startedAt;
2852-
if (exitCode === 0 && livedMs < detachGraceMs) {
2884+
if (
2885+
options.detachOnExit === true ||
2886+
(exitCode === 0 && livedMs < detachGraceMs)
2887+
) {
28532888
helper.stdout?.destroy();
28542889
helper.stderr?.destroy();
28552890
helper.unref();
@@ -2898,6 +2933,11 @@ async function createRuntimeRotationProxyContextIfEnabled(
28982933
if (isCodexAppCommand(rawArgs)) {
28992934
return createRuntimeRotationAppHelperContext(baseContext, configTomlModule);
29002935
}
2936+
if (isCodexInteractiveTuiCommand(rawArgs)) {
2937+
return createRuntimeRotationAppHelperContext(baseContext, configTomlModule, {
2938+
detachOnExit: true,
2939+
});
2940+
}
29012941

29022942
const proxyModule = await loadRuntimeRotationProxyModule();
29032943
if (!proxyModule) {
@@ -3092,6 +3132,10 @@ function isCodexAppServerCommand(rawArgs) {
30923132
return findForwardedCommand(rawArgs)?.command === "app-server";
30933133
}
30943134

3135+
function isCodexInteractiveTuiCommand(rawArgs) {
3136+
return findForwardedCommand(rawArgs) === null;
3137+
}
3138+
30953139
function shouldUseRuntimeRoutingForForwardedArgs(rawArgs) {
30963140
if (!Array.isArray(rawArgs) || rawArgs.length === 0) {
30973141
return true;

0 commit comments

Comments
 (0)