fix: address codex review findings

This commit is contained in:
Cita
2026-03-21 19:45:29 +08:00
parent 1d8889001f
commit d9e550c279
10 changed files with 638 additions and 167 deletions
+122
View File
@@ -19,6 +19,16 @@ function buildJwt(payload: Record<string, unknown>) {
return `${encode({ alg: 'none', typ: 'JWT' })}.${encode(payload)}.signature`;
}
function createDeferred<T>() {
let resolve!: (value: T | PromiseLike<T>) => void;
let reject!: (reason?: unknown) => void;
const promise = new Promise<T>((res, rej) => {
resolve = res;
reject = rej;
});
return { promise, resolve, reject };
}
describe('oauth routes', { timeout: 15_000 }, () => {
let app: FastifyInstance;
let db: DbModule['db'];
@@ -599,6 +609,118 @@ describe('oauth routes', { timeout: 15_000 }, () => {
idToken: originalJwt,
},
});
const accounts = await db.select().from(schema.accounts).all();
expect(accounts).toHaveLength(1);
expect(accounts[0]?.oauthAccountKey).toBe('chatgpt-account-existing');
});
it('keeps the existing codex account non-active while rebind model discovery is still pending', async () => {
const originalJwt = buildJwt({
email: 'codex-existing@example.com',
'https://api.openai.com/auth': {
chatgpt_account_id: 'chatgpt-account-existing',
chatgpt_plan_type: 'team',
},
});
const site = await db.insert(schema.sites).values({
name: 'ChatGPT Codex OAuth',
url: 'https://chatgpt.com/backend-api/codex',
platform: 'codex',
status: 'active',
}).returning().get();
const existing = await db.insert(schema.accounts).values({
siteId: site.id,
username: 'codex-existing@example.com',
accessToken: 'stable-access-token',
apiToken: null,
status: 'active',
oauthProvider: 'codex',
oauthAccountKey: 'chatgpt-account-existing',
extraConfig: JSON.stringify({
credentialMode: 'session',
oauth: {
provider: 'codex',
accountId: 'chatgpt-account-existing',
accountKey: 'chatgpt-account-existing',
email: 'codex-existing@example.com',
planType: 'team',
refreshToken: 'stable-refresh-token',
idToken: originalJwt,
},
}),
}).returning().get();
const reboundJwt = buildJwt({
email: 'codex-existing@example.com',
'https://api.openai.com/auth': {
chatgpt_account_id: 'chatgpt-account-rebound',
chatgpt_plan_type: 'plus',
},
});
const discoveryGate = createDeferred<ResponseLike>();
fetchMock
.mockResolvedValueOnce({
ok: true,
status: 200,
json: async () => ({
access_token: 'rebound-access-token',
refresh_token: 'rebound-refresh-token',
id_token: reboundJwt,
expires_in: 3600,
token_type: 'Bearer',
}),
text: async () => JSON.stringify({ ok: true }),
})
.mockImplementationOnce(() => discoveryGate.promise);
const startResponse = await app.inject({
method: 'POST',
url: '/api/oauth/providers/codex/start',
headers: {
host: 'metapi.example',
'x-forwarded-proto': 'https',
},
payload: {
accountId: existing.id,
},
});
const startBody = startResponse.json() as { state: string };
const callbackPromise = app.inject({
method: 'POST',
url: `/api/oauth/sessions/${encodeURIComponent(startBody.state)}/manual-callback`,
payload: {
callbackUrl: `http://localhost:1455/auth/callback?state=${encodeURIComponent(startBody.state)}&code=oauth-code-rebind-pending`,
},
});
while (fetchMock.mock.calls.length < 2) {
await new Promise((resolve) => setTimeout(resolve, 0));
}
const pendingRow = await db.select().from(schema.accounts).where(eq(schema.accounts.id, existing.id)).get();
expect(pendingRow?.status).toBe('disabled');
discoveryGate.resolve({
ok: true,
status: 200,
json: async () => ({
models: [{ id: 'gpt-5.4' }],
}),
text: async () => JSON.stringify({ ok: true }),
} as ResponseLike);
const callbackResponse = await callbackPromise;
expect(callbackResponse.statusCode).toBe(200);
const stored = await db.select().from(schema.accounts).where(eq(schema.accounts.id, existing.id)).get();
expect(stored).toMatchObject({
id: existing.id,
status: 'active',
accessToken: 'rebound-access-token',
oauthAccountKey: 'chatgpt-account-rebound',
});
});
it('fails codex oauth onboarding when token exchange does not expose chatgpt_account_id', async () => {
@@ -255,6 +255,33 @@ describe('responses proxy codex oauth refresh', () => {
expect(response.json()?.output_text).toContain('ok after codex forbidden refresh');
});
it('does not refresh codex oauth token on non-auth 403 responses', async () => {
fetchMock.mockResolvedValueOnce(new Response(JSON.stringify({
error: { message: 'quota exceeded for workspace', type: 'usage_limit_reached' },
}), {
status: 403,
headers: { 'content-type': 'application/json' },
}));
const response = await app.inject({
method: 'POST',
url: '/v1/responses',
payload: {
model: 'gpt-5.2-codex',
input: 'hello codex',
},
});
expect(response.statusCode).toBe(403);
expect(refreshOauthAccessTokenSingleflightMock).not.toHaveBeenCalled();
expect(fetchMock).toHaveBeenCalledTimes(1);
expect(response.json()).toMatchObject({
error: {
message: expect.stringContaining('quota exceeded for workspace'),
},
});
});
it('retries oauth responses requests with a normalized upstream URL after refresh', async () => {
selectChannelMock.mockReturnValue({
channel: { id: 11, routeId: 22 },
@@ -194,6 +194,16 @@ function waitForSocketMessages(socket: WebSocket, count: number, timeoutMs = 100
});
}
function createDeferred<T>() {
let resolve!: (value: T | PromiseLike<T>) => void;
let reject!: (reason?: unknown) => void;
const promise = new Promise<T>((res, rej) => {
resolve = res;
reject = rej;
});
return { promise, resolve, reject };
}
function createClientSocket(baseUrl: string, headers: Record<string, string> = {}) {
return new WebSocket(`${baseUrl}/v1/responses`, {
headers: {
@@ -203,6 +213,10 @@ function createClientSocket(baseUrl: string, headers: Record<string, string> = {
});
}
function createClientSocketForPath(path: string, headers: Record<string, string> = {}) {
return new WebSocket(path, { headers });
}
describe('responses websocket transport', () => {
const originalCodexResponsesWebsocketBeta = config.codexResponsesWebsocketBeta;
let app: FastifyInstance;
@@ -545,6 +559,54 @@ describe('responses websocket transport', () => {
expect(messages[0]?.response?.id).toBe('resp_http_fallback_401');
});
it('preserves query parameter auth when websocket transport falls back to the HTTP responses route', async () => {
rejectedUpgradeStatus = 401;
rejectedUpgradeStatusText = 'Unauthorized';
rejectedUpgradeBody = JSON.stringify({
error: {
message: 'expired token',
type: 'invalid_request_error',
},
});
const selectedChannel = createSelectedChannel({
siteUrl: rejectedUpgradeSiteUrl,
});
selectChannelMock.mockReturnValue(selectedChannel);
previewSelectedChannelMock.mockResolvedValue(selectedChannel);
authorizeDownstreamTokenMock.mockResolvedValueOnce({
ok: true,
source: 'global',
token: 'sk-query-auth',
key: null,
policy: {
supportedModels: [],
allowedRouteIds: [],
siteWeightMultipliers: {},
},
});
fetchMock.mockResolvedValueOnce(createSseResponse([
'event: response.completed\n',
'data: {"type":"response.completed","response":{"id":"resp_http_fallback_query","model":"gpt-5.4","status":"completed","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}\n\n',
'data: [DONE]\n\n',
]));
const socket = createClientSocketForPath(`${baseUrl}/v1/responses?key=sk-query-auth`);
await waitForSocketOpen(socket);
const messagesPromise = waitForSocketMessages(socket, 1);
socket.send(JSON.stringify({
type: 'response.create',
model: 'gpt-5.4',
input: [],
}));
const messages = await messagesPromise;
socket.close();
expect(fetchMock).toHaveBeenCalledTimes(1);
expect(messages[0]?.response?.id).toBe('resp_http_fallback_query');
});
it('rejects websocket turns whose model is blocked by the downstream key policy before channel selection', async () => {
authorizeDownstreamTokenMock.mockResolvedValueOnce({
ok: true,
@@ -678,6 +740,111 @@ describe('responses websocket transport', () => {
]);
});
it('serializes websocket messages per connection so follow-up turns wait for the previous HTTP fallback to finish', async () => {
const selectedChannel = createSelectedChannel({
sitePlatform: 'openai',
actualModel: 'gpt-4.1',
});
selectChannelMock.mockReturnValue(selectedChannel);
previewSelectedChannelMock.mockResolvedValue(selectedChannel);
const firstResponseGate = createDeferred<Response>();
fetchMock
.mockImplementationOnce(() => firstResponseGate.promise)
.mockResolvedValueOnce(createSseResponse([
'event: response.completed\n',
'data: {"type":"response.completed","response":{"id":"resp_ws_2","model":"gpt-4.1","status":"completed","output":[{"id":"msg_2","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"done"}]}],"usage":{"input_tokens":5,"output_tokens":1,"total_tokens":6}}}\n\n',
'data: [DONE]\n\n',
]));
const socket = createClientSocket(baseUrl);
await waitForSocketOpen(socket);
socket.send(JSON.stringify({
type: 'response.create',
model: 'gpt-4.1',
instructions: 'be helpful',
input: [
{
id: 'msg_user_1',
type: 'message',
role: 'user',
content: [{ type: 'input_text', text: 'call the tool' }],
},
],
}));
while (fetchMock.mock.calls.length < 1) {
await new Promise((resolve) => setTimeout(resolve, 0));
}
socket.send(JSON.stringify({
type: 'response.create',
previous_response_id: 'resp_ws_1',
input: [
{
id: 'tool_out_1',
type: 'function_call_output',
call_id: 'call_1',
output: 'tool result',
},
],
}));
await new Promise((resolve) => setTimeout(resolve, 0));
expect(fetchMock).toHaveBeenCalledTimes(1);
const messagesPromise = waitForSocketMessages(socket, 4);
firstResponseGate.resolve(createSseResponse([
'event: response.output_item.done\n',
'data: {"type":"response.output_item.done","output_index":0,"item":{"id":"fc_1","type":"function_call","call_id":"call_1"}}\n\n',
'event: response.output_item.done\n',
'data: {"type":"response.output_item.done","output_index":1,"item":{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"call tool"}]}}\n\n',
'event: response.completed\n',
'data: {"type":"response.completed","response":{"id":"resp_ws_1","model":"gpt-4.1","status":"completed","output":[{"id":"fc_1","type":"function_call","call_id":"call_1"},{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"call tool"}]}],"usage":{"input_tokens":3,"output_tokens":1,"total_tokens":4}}}\n\n',
'data: [DONE]\n\n',
]));
const messages = await messagesPromise;
socket.close();
expect(fetchMock).toHaveBeenCalledTimes(2);
expect(messages.map((message) => message.type)).toEqual([
'response.output_item.done',
'response.output_item.done',
'response.completed',
'response.completed',
]);
const [, secondOptions] = fetchMock.mock.calls[1] as [string, RequestInit];
const secondBody = JSON.parse(String(secondOptions.body));
expect(secondBody.input).toEqual([
{
id: 'msg_user_1',
type: 'message',
role: 'user',
content: [{ type: 'input_text', text: 'call the tool' }],
},
{
id: 'fc_1',
type: 'function_call',
call_id: 'call_1',
status: 'completed',
},
{
id: 'msg_1',
type: 'message',
role: 'assistant',
status: 'completed',
content: [{ type: 'output_text', text: 'call tool' }],
},
{
id: 'tool_out_1',
type: 'function_call_output',
call_id: 'call_1',
output: 'tool result',
},
]);
});
it('preserves incremental response.create payloads with previous_response_id for websocket-capable upstreams', async () => {
const selectedChannel = createSelectedChannel({
siteUrl: upstreamSiteUrl,
+153 -137
View File
@@ -339,15 +339,25 @@ async function forwardResponsesRequestViaHttp(input: {
request: IncomingMessage;
payload: Record<string, unknown>;
preserveIncrementalMode: boolean;
authToken: string;
}): Promise<unknown[] | null> {
const injectHeaders = {
...buildInjectHeaders(input.request),
[RESPONSES_WEBSOCKET_TRANSPORT_HEADER]: '1',
...(input.preserveIncrementalMode ? { [RESPONSES_WEBSOCKET_MODE_HEADER]: 'incremental' } : {}),
};
if (
!headerValueToTrimmedString(injectHeaders.authorization)
&& !headerValueToTrimmedString(injectHeaders['x-api-key'])
&& !headerValueToTrimmedString(injectHeaders['x-goog-api-key'])
) {
injectHeaders.authorization = `Bearer ${input.authToken}`;
}
const response = await input.app.inject({
method: 'POST',
url: '/v1/responses',
headers: {
...buildInjectHeaders(input.request),
[RESPONSES_WEBSOCKET_TRANSPORT_HEADER]: '1',
...(input.preserveIncrementalMode ? { [RESPONSES_WEBSOCKET_MODE_HEADER]: 'incremental' } : {}),
},
headers: injectHeaders,
payload: input.payload,
});
@@ -444,7 +454,7 @@ function writeUpgradeHttpError(socket: Duplex, status: number, message: string):
? 'Bad Request'
: 'Error';
const body = JSON.stringify({ error: message });
socket.write(
socket.end(
`HTTP/1.1 ${status} ${statusText}\r\n`
+ 'Content-Type: application/json\r\n'
+ `Content-Length: ${Buffer.byteLength(body)}\r\n`
@@ -452,7 +462,6 @@ function writeUpgradeHttpError(socket: Duplex, status: number, message: string):
+ '\r\n'
+ body,
);
socket.destroy();
}
async function supportsResponsesWebsocketIncrementalInput(
@@ -483,148 +492,155 @@ async function handleResponsesWebsocketConnection(
let lastRequest: Record<string, unknown> | null = null;
let lastResponseOutput: unknown[] = [];
let selectedChannel: SelectedChannel | null = null;
let messageQueue = Promise.resolve();
socket.once('close', () => {
void codexWebsocketRuntime.closeSession(websocketSessionId);
});
socket.on('message', async (raw) => {
try {
const parsed = parseJsonObject(raw);
if (!parsed) {
writeResponsesWebsocketError(socket, 400, 'Invalid websocket JSON payload');
return;
}
const requestModel = asTrimmedString(parsed.model) || asTrimmedString(lastRequest?.model);
if (requestModel && !await isModelAllowedByPolicyOrAllowedRoutes(requestModel, authContext.policy)) {
writeResponsesWebsocketError(socket, 403, 'model is not allowed for this downstream key');
return;
}
const supportsIncrementalInput = selectedChannelSupportsIncrementalInput(selectedChannel, requestModel)
|| await supportsResponsesWebsocketIncrementalInput(parsed, lastRequest, authContext);
const shouldHandleLocalPrewarm = shouldHandleResponsesWebsocketPrewarmLocally(
parsed,
lastRequest,
supportsIncrementalInput,
);
const normalized = normalizeResponsesWebsocketRequest(
parsed,
lastRequest,
lastResponseOutput,
supportsIncrementalInput,
);
if (!normalized.ok) {
writeResponsesWebsocketError(socket, normalized.status, normalized.message);
return;
}
if (authContext.source === 'managed' && authContext.key?.id) {
await consumeManagedKeyRequest(authContext.key.id);
}
lastRequest = normalized.nextRequestSnapshot;
if (shouldHandleLocalPrewarm) {
lastResponseOutput = [];
for (const payload of synthesizePrewarmResponsePayloads(normalized.request)) {
socket.send(JSON.stringify(payload));
}
return;
}
if (!shouldReuseSelectedChannel(selectedChannel, requestModel)) {
selectedChannel = requestModel
? await tokenRouter.selectChannel(requestModel, authContext.policy)
: null;
}
const codexWebsocketChannel = selectedChannelSupportsCodexWebsocketTransport(selectedChannel, requestModel)
? selectedChannel
: null;
if (codexWebsocketChannel) {
const downstreamHeaders: Record<string, unknown> = {
...(request.headers as Record<string, unknown>),
[RESPONSES_WEBSOCKET_TRANSPORT_HEADER]: '1',
...(supportsIncrementalInput ? { [RESPONSES_WEBSOCKET_MODE_HEADER]: 'incremental' } : {}),
};
const providerHeaders = buildOauthProviderHeaders({
extraConfig: typeof codexWebsocketChannel.account.extraConfig === 'string'
? codexWebsocketChannel.account.extraConfig
: null,
downstreamHeaders,
});
const prepared = buildUpstreamEndpointRequest({
endpoint: 'responses',
modelName: asTrimmedString(codexWebsocketChannel.actualModel) || requestModel,
stream: true,
tokenValue: codexWebsocketChannel.tokenValue,
sitePlatform: codexWebsocketChannel.site.platform,
siteUrl: codexWebsocketChannel.site.url,
openaiBody: normalized.request,
downstreamFormat: 'responses',
responsesOriginalBody: normalized.request,
downstreamHeaders,
providerHeaders,
codexExplicitSessionId: deriveCodexExplicitSessionId(normalized.request, websocketSessionId),
});
const requestUrl = `${codexWebsocketChannel.site.url.replace(/\/+$/, '')}${prepared.path}`;
socket.on('message', (raw) => {
messageQueue = messageQueue
.catch(() => undefined)
.then(async () => {
try {
const runtimeResult = await codexWebsocketRuntime.sendRequest({
sessionId: websocketSessionId,
requestUrl,
headers: prepared.headers,
body: prepared.body,
});
lastResponseOutput = collectResponsesOutput(runtimeResult.events);
for (const payload of runtimeResult.events) {
socket.send(JSON.stringify(payload));
const parsed = parseJsonObject(raw);
if (!parsed) {
writeResponsesWebsocketError(socket, 400, 'Invalid websocket JSON payload');
return;
}
} catch (error) {
const runtimeError = error instanceof CodexWebsocketRuntimeError
? error
: new CodexWebsocketRuntimeError('upstream websocket request failed');
if (runtimeError.status && runtimeError.events.length === 0) {
const fallbackOutput = await forwardResponsesRequestViaHttp({
app,
socket,
request,
payload: normalized.request,
preserveIncrementalMode: false,
});
if (fallbackOutput) {
lastResponseOutput = fallbackOutput;
const requestModel = asTrimmedString(parsed.model) || asTrimmedString(lastRequest?.model);
if (requestModel && !await isModelAllowedByPolicyOrAllowedRoutes(requestModel, authContext.policy)) {
writeResponsesWebsocketError(socket, 403, 'model is not allowed for this downstream key');
return;
}
const supportsIncrementalInput = selectedChannelSupportsIncrementalInput(selectedChannel, requestModel)
|| await supportsResponsesWebsocketIncrementalInput(parsed, lastRequest, authContext);
const shouldHandleLocalPrewarm = shouldHandleResponsesWebsocketPrewarmLocally(
parsed,
lastRequest,
supportsIncrementalInput,
);
const normalized = normalizeResponsesWebsocketRequest(
parsed,
lastRequest,
lastResponseOutput,
supportsIncrementalInput,
);
if (!normalized.ok) {
writeResponsesWebsocketError(socket, normalized.status, normalized.message);
return;
}
if (authContext.source === 'managed' && authContext.key?.id) {
await consumeManagedKeyRequest(authContext.key.id);
}
lastRequest = normalized.nextRequestSnapshot;
if (shouldHandleLocalPrewarm) {
lastResponseOutput = [];
for (const payload of synthesizePrewarmResponsePayloads(normalized.request)) {
socket.send(JSON.stringify(payload));
}
return;
}
lastResponseOutput = collectResponsesOutput(runtimeError.events);
for (const payload of runtimeError.events) {
socket.send(JSON.stringify(payload));
}
writeResponsesWebsocketError(
socket,
runtimeError.status || 408,
runtimeError.message,
runtimeError.payload,
);
}
return;
}
const forwardedOutput = await forwardResponsesRequestViaHttp({
app,
socket,
request,
payload: normalized.request,
preserveIncrementalMode: supportsIncrementalInput,
if (!shouldReuseSelectedChannel(selectedChannel, requestModel)) {
selectedChannel = requestModel
? await tokenRouter.selectChannel(requestModel, authContext.policy)
: null;
}
const codexWebsocketChannel = selectedChannelSupportsCodexWebsocketTransport(selectedChannel, requestModel)
? selectedChannel
: null;
if (codexWebsocketChannel) {
const downstreamHeaders: Record<string, unknown> = {
...(request.headers as Record<string, unknown>),
[RESPONSES_WEBSOCKET_TRANSPORT_HEADER]: '1',
...(supportsIncrementalInput ? { [RESPONSES_WEBSOCKET_MODE_HEADER]: 'incremental' } : {}),
};
const providerHeaders = buildOauthProviderHeaders({
extraConfig: typeof codexWebsocketChannel.account.extraConfig === 'string'
? codexWebsocketChannel.account.extraConfig
: null,
downstreamHeaders,
});
const prepared = buildUpstreamEndpointRequest({
endpoint: 'responses',
modelName: asTrimmedString(codexWebsocketChannel.actualModel) || requestModel,
stream: true,
tokenValue: codexWebsocketChannel.tokenValue,
sitePlatform: codexWebsocketChannel.site.platform,
siteUrl: codexWebsocketChannel.site.url,
openaiBody: normalized.request,
downstreamFormat: 'responses',
responsesOriginalBody: normalized.request,
downstreamHeaders,
providerHeaders,
codexExplicitSessionId: deriveCodexExplicitSessionId(normalized.request, websocketSessionId),
});
const requestUrl = `${codexWebsocketChannel.site.url.replace(/\/+$/, '')}${prepared.path}`;
try {
const runtimeResult = await codexWebsocketRuntime.sendRequest({
sessionId: websocketSessionId,
requestUrl,
headers: prepared.headers,
body: prepared.body,
});
lastResponseOutput = collectResponsesOutput(runtimeResult.events);
for (const payload of runtimeResult.events) {
socket.send(JSON.stringify(payload));
}
} catch (error) {
const runtimeError = error instanceof CodexWebsocketRuntimeError
? error
: new CodexWebsocketRuntimeError('upstream websocket request failed');
if (runtimeError.status && runtimeError.events.length === 0) {
const fallbackOutput = await forwardResponsesRequestViaHttp({
app,
socket,
request,
payload: normalized.request,
preserveIncrementalMode: false,
authToken: authContext.token,
});
if (fallbackOutput) {
lastResponseOutput = fallbackOutput;
}
return;
}
lastResponseOutput = collectResponsesOutput(runtimeError.events);
for (const payload of runtimeError.events) {
socket.send(JSON.stringify(payload));
}
writeResponsesWebsocketError(
socket,
runtimeError.status || 408,
runtimeError.message,
runtimeError.payload,
);
}
return;
}
const forwardedOutput = await forwardResponsesRequestViaHttp({
app,
socket,
request,
payload: normalized.request,
preserveIncrementalMode: supportsIncrementalInput,
authToken: authContext.token,
});
if (forwardedOutput) {
lastResponseOutput = forwardedOutput;
}
} catch {
writeResponsesWebsocketError(socket, 500, 'internal websocket proxy error');
}
});
if (forwardedOutput) {
lastResponseOutput = forwardedOutput;
}
} catch {
writeResponsesWebsocketError(socket, 500, 'internal websocket proxy error');
}
});
}