mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
Merge pull request #7364 from earendil-works/fix/codex-account-websocket
fix(ai): scope Codex sockets to accounts
This commit is contained in:
Christian Klotz
GitHub
commit
784653468c
2 changed files with 125 additions and 19 deletions
|
|
@ -328,6 +328,7 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons
|
|||
httpTimeoutMs,
|
||||
websocketConnectTimeoutMs,
|
||||
cacheSessionId,
|
||||
accountId,
|
||||
grammarToolInputProperties,
|
||||
options,
|
||||
);
|
||||
|
|
@ -868,7 +869,7 @@ export interface OpenAICodexWebSocketDebugStats {
|
|||
lastWebSocketError?: string;
|
||||
}
|
||||
|
||||
const websocketSessionCache = new Map<string, CachedWebSocketConnection>();
|
||||
const websocketSessionCache = new Map<string, Map<string, CachedWebSocketConnection>>();
|
||||
const websocketDebugStats = new Map<string, OpenAICodexWebSocketDebugStats>();
|
||||
const websocketSseFallbackSessions = new Set<string>();
|
||||
|
||||
|
|
@ -913,13 +914,12 @@ export function closeOpenAICodexWebSocketSessions(sessionId?: string): void {
|
|||
closeWebSocketSilently(entry.socket, 1000, "debug_close");
|
||||
};
|
||||
if (sessionId) {
|
||||
const entry = websocketSessionCache.get(sessionId);
|
||||
if (entry) closeEntry(entry);
|
||||
for (const entry of websocketSessionCache.get(sessionId)?.values() ?? []) closeEntry(entry);
|
||||
websocketSessionCache.delete(sessionId);
|
||||
return;
|
||||
}
|
||||
for (const entry of websocketSessionCache.values()) {
|
||||
closeEntry(entry);
|
||||
for (const accountEntries of websocketSessionCache.values()) {
|
||||
for (const entry of accountEntries.values()) closeEntry(entry);
|
||||
}
|
||||
websocketSessionCache.clear();
|
||||
}
|
||||
|
|
@ -1021,14 +1021,16 @@ function closeWebSocketSilently(socket: WebSocketLike, code = 1000, reason = "do
|
|||
} catch {}
|
||||
}
|
||||
|
||||
function scheduleSessionWebSocketExpiry(sessionId: string, entry: CachedWebSocketConnection): void {
|
||||
function scheduleSessionWebSocketExpiry(sessionId: string, accountId: string, entry: CachedWebSocketConnection): void {
|
||||
if (entry.idleTimer) {
|
||||
clearTimeout(entry.idleTimer);
|
||||
}
|
||||
entry.idleTimer = setTimeout(() => {
|
||||
if (entry.busy) return;
|
||||
closeWebSocketSilently(entry.socket, 1000, "idle_timeout");
|
||||
websocketSessionCache.delete(sessionId);
|
||||
const accountEntries = websocketSessionCache.get(sessionId);
|
||||
if (accountEntries?.get(accountId) === entry) accountEntries.delete(accountId);
|
||||
if (accountEntries?.size === 0) websocketSessionCache.delete(sessionId);
|
||||
}, SESSION_WEBSOCKET_CACHE_TTL_MS);
|
||||
}
|
||||
|
||||
|
|
@ -1114,6 +1116,7 @@ async function acquireWebSocket(
|
|||
url: string,
|
||||
headers: Headers,
|
||||
sessionId: string | undefined,
|
||||
accountId: string,
|
||||
signal?: AbortSignal,
|
||||
connectTimeoutMs?: number,
|
||||
env?: ProviderEnv,
|
||||
|
|
@ -1132,7 +1135,8 @@ async function acquireWebSocket(
|
|||
};
|
||||
}
|
||||
|
||||
const cached = websocketSessionCache.get(sessionId);
|
||||
let accountEntries = websocketSessionCache.get(sessionId);
|
||||
const cached = accountEntries?.get(accountId);
|
||||
if (cached) {
|
||||
if (cached.idleTimer) {
|
||||
clearTimeout(cached.idleTimer);
|
||||
|
|
@ -1140,7 +1144,8 @@ async function acquireWebSocket(
|
|||
}
|
||||
if (!cached.busy && isWebSocketSessionExpired(cached)) {
|
||||
closeWebSocketSilently(cached.socket, 1000, "connection_age_limit");
|
||||
websocketSessionCache.delete(sessionId);
|
||||
accountEntries?.delete(accountId);
|
||||
if (accountEntries?.size === 0) websocketSessionCache.delete(sessionId);
|
||||
} else if (!cached.busy && isWebSocketReusable(cached.socket)) {
|
||||
cached.busy = true;
|
||||
return {
|
||||
|
|
@ -1150,11 +1155,13 @@ async function acquireWebSocket(
|
|||
release: ({ keep } = {}) => {
|
||||
if (!keep || !isWebSocketReusable(cached.socket)) {
|
||||
closeWebSocketSilently(cached.socket);
|
||||
websocketSessionCache.delete(sessionId);
|
||||
const currentEntries = websocketSessionCache.get(sessionId);
|
||||
if (currentEntries?.get(accountId) === cached) currentEntries.delete(accountId);
|
||||
if (currentEntries?.size === 0) websocketSessionCache.delete(sessionId);
|
||||
return;
|
||||
}
|
||||
cached.busy = false;
|
||||
scheduleSessionWebSocketExpiry(sessionId, cached);
|
||||
scheduleSessionWebSocketExpiry(sessionId, accountId, cached);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
|
@ -1170,13 +1177,19 @@ async function acquireWebSocket(
|
|||
}
|
||||
if (!isWebSocketReusable(cached.socket)) {
|
||||
closeWebSocketSilently(cached.socket);
|
||||
websocketSessionCache.delete(sessionId);
|
||||
accountEntries?.delete(accountId);
|
||||
if (accountEntries?.size === 0) websocketSessionCache.delete(sessionId);
|
||||
}
|
||||
}
|
||||
|
||||
const socket = await connectWebSocket(url, headers, signal, connectTimeoutMs, env);
|
||||
const entry: CachedWebSocketConnection = { socket, busy: true, createdAt: Date.now() };
|
||||
websocketSessionCache.set(sessionId, entry);
|
||||
accountEntries = websocketSessionCache.get(sessionId);
|
||||
if (!accountEntries) {
|
||||
accountEntries = new Map();
|
||||
websocketSessionCache.set(sessionId, accountEntries);
|
||||
}
|
||||
accountEntries.set(accountId, entry);
|
||||
return {
|
||||
socket,
|
||||
entry,
|
||||
|
|
@ -1185,13 +1198,13 @@ async function acquireWebSocket(
|
|||
if (!keep || !isWebSocketReusable(entry.socket)) {
|
||||
closeWebSocketSilently(entry.socket);
|
||||
if (entry.idleTimer) clearTimeout(entry.idleTimer);
|
||||
if (websocketSessionCache.get(sessionId) === entry) {
|
||||
websocketSessionCache.delete(sessionId);
|
||||
}
|
||||
const currentEntries = websocketSessionCache.get(sessionId);
|
||||
if (currentEntries?.get(accountId) === entry) currentEntries.delete(accountId);
|
||||
if (currentEntries?.size === 0) websocketSessionCache.delete(sessionId);
|
||||
return;
|
||||
}
|
||||
entry.busy = false;
|
||||
scheduleSessionWebSocketExpiry(sessionId, entry);
|
||||
scheduleSessionWebSocketExpiry(sessionId, accountId, entry);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
|
@ -1450,6 +1463,7 @@ async function processWebSocketStream(
|
|||
idleTimeoutMs: number | undefined,
|
||||
websocketConnectTimeoutMs: number | undefined,
|
||||
cacheSessionId: string | undefined,
|
||||
accountId: string,
|
||||
grammarToolInputProperties: ReadonlyMap<string, string>,
|
||||
options?: OpenAICodexResponsesOptions,
|
||||
): Promise<void> {
|
||||
|
|
@ -1457,6 +1471,7 @@ async function processWebSocketStream(
|
|||
url,
|
||||
headers,
|
||||
cacheSessionId,
|
||||
accountId,
|
||||
options?.signal,
|
||||
websocketConnectTimeoutMs,
|
||||
options?.env,
|
||||
|
|
|
|||
|
|
@ -28,9 +28,9 @@ afterEach(() => {
|
|||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
function mockToken(): string {
|
||||
function mockToken(accountId = "acc_test"): string {
|
||||
const payload = Buffer.from(
|
||||
JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "acc_test" } }),
|
||||
JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: accountId } }),
|
||||
"utf8",
|
||||
).toString("base64");
|
||||
return `aaa.${payload}.bbb`;
|
||||
|
|
@ -1334,6 +1334,97 @@ describe("openai-codex streaming", () => {
|
|||
});
|
||||
});
|
||||
|
||||
it("scopes cached websockets to the authenticated account", async () => {
|
||||
// Regression for #7284: rotating accounts must not reuse a socket authenticated by another account.
|
||||
const connectedHeaders: Record<string, string>[] = [];
|
||||
let responseId = 0;
|
||||
|
||||
class MockWebSocket {
|
||||
static OPEN = 1;
|
||||
readyState = MockWebSocket.OPEN;
|
||||
private listeners = new Map<string, Set<(event: unknown) => void>>();
|
||||
|
||||
constructor(_url: string, protocols?: string | string[] | { headers?: Record<string, string> }) {
|
||||
const headers =
|
||||
protocols && typeof protocols === "object" && !Array.isArray(protocols) ? protocols.headers : undefined;
|
||||
connectedHeaders.push(headers ?? {});
|
||||
queueMicrotask(() => this.dispatch("open", {}));
|
||||
}
|
||||
|
||||
addEventListener(type: string, listener: (event: unknown) => void): void {
|
||||
let listeners = this.listeners.get(type);
|
||||
if (!listeners) {
|
||||
listeners = new Set();
|
||||
this.listeners.set(type, listeners);
|
||||
}
|
||||
listeners.add(listener);
|
||||
}
|
||||
|
||||
removeEventListener(type: string, listener: (event: unknown) => void): void {
|
||||
this.listeners.get(type)?.delete(listener);
|
||||
}
|
||||
|
||||
send(): void {
|
||||
queueMicrotask(() => {
|
||||
this.dispatch("message", {
|
||||
data: JSON.stringify({
|
||||
type: "response.completed",
|
||||
response: {
|
||||
id: `resp_${++responseId}`,
|
||||
status: "completed",
|
||||
usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 },
|
||||
},
|
||||
}),
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
close(): void {
|
||||
this.readyState = 3;
|
||||
}
|
||||
|
||||
private dispatch(type: string, event: unknown): void {
|
||||
for (const listener of this.listeners.get(type) ?? []) listener(event);
|
||||
}
|
||||
}
|
||||
|
||||
vi.stubGlobal("WebSocket", MockWebSocket);
|
||||
vi.stubGlobal(
|
||||
"fetch",
|
||||
vi.fn(async () => new Response("unexpected fetch", { status: 500 })),
|
||||
);
|
||||
|
||||
const model: Model<"openai-codex-responses"> = {
|
||||
id: "gpt-5.1-codex",
|
||||
name: "GPT-5.1 Codex",
|
||||
api: "openai-codex-responses",
|
||||
provider: "openai-codex",
|
||||
baseUrl: "https://chatgpt.com/backend-api",
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 400000,
|
||||
maxTokens: 128000,
|
||||
};
|
||||
const context: Context = { systemPrompt: "", messages: [] };
|
||||
const options = { sessionId: "shared-session", transport: "websocket-cached" as const };
|
||||
|
||||
await streamOpenAICodexResponses(model, context, { ...options, apiKey: mockToken("account-a") }).result();
|
||||
await streamOpenAICodexResponses(model, context, { ...options, apiKey: mockToken("account-b") }).result();
|
||||
await streamOpenAICodexResponses(model, context, { ...options, apiKey: mockToken("account-a") }).result();
|
||||
|
||||
expect(connectedHeaders.map((headers) => headers["chatgpt-account-id"])).toEqual(["account-a", "account-b"]);
|
||||
expect(connectedHeaders.map((headers) => headers.authorization)).toEqual([
|
||||
`Bearer ${mockToken("account-a")}`,
|
||||
`Bearer ${mockToken("account-b")}`,
|
||||
]);
|
||||
expect(global.fetch).not.toHaveBeenCalled();
|
||||
expect(getOpenAICodexWebSocketDebugStats("shared-session")).toMatchObject({
|
||||
connectionsCreated: 2,
|
||||
connectionsReused: 1,
|
||||
});
|
||||
});
|
||||
|
||||
it("closes one-shot websockets when cacheRetention is none", async () => {
|
||||
const token = mockToken();
|
||||
const sentBodies: Array<{ prompt_cache_key?: string }> = [];
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue