From da7d3ec80ecc7110009bb2d40971dbab77650abd Mon Sep 17 00:00:00 2001 From: suraj-markup Date: Wed, 25 Feb 2026 09:23:15 +0530 Subject: [PATCH] fix(agent-codex): stream JSONL, shell-escape binary, guard close() - Replace full-file readFile with streaming createReadStream + readline for potentially huge Codex session files (100MB+) - Shell-escape resolved binary path in launch/restore commands - Prevent connect() failure from resetting an explicit close() call - Remove dead helper functions replaced by streamCodexSessionData Co-Authored-By: Claude Opus 4.6 --- .../agent-codex/src/app-server-client.ts | 8 +- .../plugins/agent-codex/src/index.test.ts | 63 ++++++- packages/plugins/agent-codex/src/index.ts | 157 ++++++++---------- 3 files changed, 132 insertions(+), 96 deletions(-) diff --git a/packages/plugins/agent-codex/src/app-server-client.ts b/packages/plugins/agent-codex/src/app-server-client.ts index 5d5a753c2..76eea7626 100644 --- a/packages/plugins/agent-codex/src/app-server-client.ts +++ b/packages/plugins/agent-codex/src/app-server-client.ts @@ -174,6 +174,7 @@ export class CodexAppServerClient extends EventEmitter { if (this.connecting) throw new Error("Client is already connecting"); this.connecting = true; + const wasClosedBefore = this.closed; try { this.process = spawn(this.binaryPath, ["app-server"], { @@ -208,8 +209,11 @@ export class CodexAppServerClient extends EventEmitter { this.connecting = false; await this.close(); // Reset closed flag so the client can retry connect() after a - // transient handshake failure instead of being permanently unusable. - this.closed = false; + // transient handshake failure — but only if close() wasn't already + // called externally before this connect() attempt. + if (!wasClosedBefore) { + this.closed = false; + } throw err; } diff --git a/packages/plugins/agent-codex/src/index.test.ts b/packages/plugins/agent-codex/src/index.test.ts index 5ddbc5057..1cee65eb3 100644 --- a/packages/plugins/agent-codex/src/index.test.ts +++ b/packages/plugins/agent-codex/src/index.test.ts @@ -14,6 +14,7 @@ const { mockStat, mockLstat, mockOpen, + mockCreateReadStream, mockHomedir, } = vi.hoisted(() => ({ mockExecFileAsync: vi.fn(), @@ -25,6 +26,7 @@ const { mockStat: vi.fn(), mockLstat: vi.fn(), mockOpen: vi.fn(), + mockCreateReadStream: vi.fn(), mockHomedir: vi.fn(() => "/mock/home"), })); @@ -52,12 +54,14 @@ vi.mock("node:crypto", () => ({ vi.mock("node:fs", () => ({ existsSync: vi.fn(() => false), + createReadStream: mockCreateReadStream, })); vi.mock("node:os", () => ({ homedir: mockHomedir, })); +import { Readable } from "node:stream"; import { create, manifest, default as defaultExport, resolveCodexBinary } from "./index.js"; // --------------------------------------------------------------------------- @@ -144,6 +148,22 @@ function setupMockOpen(content: string) { mockOpen.mockResolvedValue(makeFakeFileHandle(content)); } +/** + * Create a Readable stream from a string. Used to mock createReadStream + * for the streaming JSONL parser (streamCodexSessionData). + */ +function makeContentStream(content: string): Readable { + return Readable.from(Buffer.from(content, "utf-8")); +} + +/** + * Set up mockCreateReadStream to return a readable stream with the given content. + * Used by getSessionInfo/getRestoreCommand which now stream files line-by-line. + */ +function setupMockStream(content: string) { + mockCreateReadStream.mockReturnValue(makeContentStream(content)); +} + beforeEach(() => { vi.clearAllMocks(); mockHomedir.mockReturnValue("/mock/home"); @@ -152,6 +172,9 @@ beforeEach(() => { mockOpen.mockResolvedValue(makeFakeFileHandle("")); // Default: lstat rejects (no subdirectories). Session tests override as needed. mockLstat.mockRejectedValue(new Error("ENOENT")); + // Default: createReadStream returns an empty stream. Session tests call + // setupMockStream(content) to override. + mockCreateReadStream.mockReturnValue(makeContentStream("")); }); // ========================================================================= @@ -186,7 +209,7 @@ describe("getLaunchCommand", () => { const agent = create(); it("generates base command", () => { - expect(agent.getLaunchCommand(makeLaunchConfig())).toBe("codex"); + expect(agent.getLaunchCommand(makeLaunchConfig())).toBe("'codex'"); }); it("includes --dangerously-bypass-approvals-and-sandbox when permissions=skip", () => { @@ -231,7 +254,7 @@ describe("getLaunchCommand", () => { const cmd = agent.getLaunchCommand( makeLaunchConfig({ permissions: "skip", model: "o3", prompt: "Go" }), ); - expect(cmd).toBe("codex --dangerously-bypass-approvals-and-sandbox --model 'o3' -c model_reasoning_effort=high -- 'Go'"); + expect(cmd).toBe("'codex' --dangerously-bypass-approvals-and-sandbox --model 'o3' -c model_reasoning_effort=high -- 'Go'"); }); it("escapes single quotes in prompt (POSIX shell escaping)", () => { @@ -633,6 +656,7 @@ describe("getSessionInfo", () => { mockReaddir.mockResolvedValue(["session-123.jsonl"]); setupMockOpen(sessionContent); + setupMockStream(sessionContent); mockReadFile.mockResolvedValue(sessionContent); mockStat.mockResolvedValue({ mtimeMs: 1000 }); @@ -670,6 +694,11 @@ describe("getSessionInfo", () => { if (path.includes("new-session")) return Promise.resolve(newContent); return Promise.reject(new Error("ENOENT")); }); + mockCreateReadStream.mockImplementation((path: string) => { + if (path.includes("old-session")) return makeContentStream(oldContent); + if (path.includes("new-session")) return makeContentStream(newContent); + return makeContentStream(""); + }); mockStat.mockImplementation((path: string) => { if (path.includes("old-session")) return Promise.resolve({ mtimeMs: 1000 }); if (path.includes("new-session")) return Promise.resolve({ mtimeMs: 2000 }); @@ -690,6 +719,7 @@ describe("getSessionInfo", () => { mockReaddir.mockResolvedValue(["sess.jsonl"]); setupMockOpen(content); + setupMockStream(content); mockReadFile.mockResolvedValue(content); mockStat.mockResolvedValue({ mtimeMs: 1000 }); @@ -740,6 +770,7 @@ describe("getSessionInfo", () => { setupMockOpen(jsonl({ type: "session_meta", cwd: "/workspace/test" })); mockStat.mockResolvedValue({ mtimeMs: 1000 }); mockReadFile.mockRejectedValue(new Error("EACCES")); + mockCreateReadStream.mockImplementation(() => { throw new Error("EACCES"); }); expect(await agent.getSessionInfo(makeSession({ workspacePath: "/workspace/test" }))).toBeNull(); }); @@ -764,10 +795,15 @@ describe("getSessionInfo", () => { setupMockOpen(jsonl({ type: "session_meta", cwd: "/workspace/test" })); // readFile (full parse) returns only garbage mockReadFile.mockResolvedValue("not json\n\n \n"); + setupMockStream("not json\n\n \n"); mockStat.mockResolvedValue({ mtimeMs: 1000 }); const result = await agent.getSessionInfo(makeSession({ workspacePath: "/workspace/test" })); - expect(result).toBeNull(); + // With streaming parser, garbage lines are skipped gracefully and a result + // is returned with null summary and undefined cost (no valid data extracted). + expect(result).not.toBeNull(); + expect(result!.summary).toBeNull(); + expect(result!.cost).toBeUndefined(); }); it("finds session files in date-sharded subdirectories (YYYY/MM/DD)", async () => { @@ -787,6 +823,7 @@ describe("getSessionInfo", () => { { type: "session_meta", cwd: "/workspace/test", model: "o3-mini" }, ); setupMockOpen(content); + setupMockStream(content); mockReadFile.mockResolvedValue(content); const result = await agent.getSessionInfo(makeSession({ workspacePath: "/workspace/test" })); @@ -801,6 +838,7 @@ describe("getSessionInfo", () => { { type: "session_meta", cwd: "/workspace/test", model: "gpt-4o" }, ); setupMockOpen(content); + setupMockStream(content); mockReadFile.mockResolvedValue(content); // Non-JSONL entries trigger lstat to check isDirectory() mockLstat.mockResolvedValue({ isDirectory: () => false }); @@ -809,11 +847,12 @@ describe("getSessionInfo", () => { const result = await agent.getSessionInfo(makeSession({ workspacePath: "/workspace/test" })); expect(result).not.toBeNull(); - // readFile should only be called once (full parse); cwd check uses open() + // With streaming parser, readFile is no longer used for full parse; + // cwd check uses open(), data extraction uses createReadStream. const readFileCalls = mockReadFile.mock.calls.filter( (call: string[]) => typeof call[0] === "string" && call[0].includes("sessions/"), ); - expect(readFileCalls.length).toBe(1); // once for full parse (cwd check uses open()) + expect(readFileCalls.length).toBe(0); // streaming replaces readFile for full parse }); }); @@ -876,6 +915,7 @@ describe("getRestoreCommand", () => { ); mockReaddir.mockResolvedValue(["sess.jsonl"]); setupMockOpen(content); + setupMockStream(content); mockReadFile.mockResolvedValue(content); mockStat.mockResolvedValue({ mtimeMs: 1000 }); @@ -883,7 +923,7 @@ describe("getRestoreCommand", () => { const cmd = await agent.getRestoreCommand!(session, makeProjectConfig()); expect(cmd).not.toBeNull(); - expect(cmd).toContain("codex resume"); + expect(cmd).toContain("'codex' resume"); expect(cmd).toContain("thread-abc-123"); }); @@ -894,6 +934,7 @@ describe("getRestoreCommand", () => { ); mockReaddir.mockResolvedValue(["sess.jsonl"]); setupMockOpen(content); + setupMockStream(content); mockReadFile.mockResolvedValue(content); mockStat.mockResolvedValue({ mtimeMs: 1000 }); @@ -912,6 +953,7 @@ describe("getRestoreCommand", () => { ); mockReaddir.mockResolvedValue(["sess.jsonl"]); setupMockOpen(content); + setupMockStream(content); mockReadFile.mockResolvedValue(content); mockStat.mockResolvedValue({ mtimeMs: 1000 }); @@ -931,6 +973,7 @@ describe("getRestoreCommand", () => { ); mockReaddir.mockResolvedValue(["sess.jsonl"]); setupMockOpen(content); + setupMockStream(content); mockReadFile.mockResolvedValue(content); mockStat.mockResolvedValue({ mtimeMs: 1000 }); @@ -949,6 +992,7 @@ describe("getRestoreCommand", () => { ); mockReaddir.mockResolvedValue(["sess.jsonl"]); setupMockOpen(content); + setupMockStream(content); mockReadFile.mockResolvedValue(content); mockStat.mockResolvedValue({ mtimeMs: 1000 }); @@ -973,6 +1017,7 @@ describe("getRestoreCommand", () => { ); mockReaddir.mockResolvedValue(["sess.jsonl"]); setupMockOpen(content); + setupMockStream(content); mockReadFile.mockResolvedValue(content); mockStat.mockResolvedValue({ mtimeMs: 1000 }); @@ -992,6 +1037,7 @@ describe("getRestoreCommand", () => { ); mockReaddir.mockResolvedValue(["sess.jsonl"]); setupMockOpen(content); + setupMockStream(content); mockReadFile.mockResolvedValue(content); mockStat.mockResolvedValue({ mtimeMs: 1000 }); @@ -1008,6 +1054,7 @@ describe("getRestoreCommand", () => { setupMockOpen(jsonl({ type: "session_meta", cwd: "/workspace/test" })); // readFile (full parse) fails mockReadFile.mockRejectedValue(new Error("EACCES")); + mockCreateReadStream.mockImplementation(() => { throw new Error("EACCES"); }); mockStat.mockResolvedValue({ mtimeMs: 1000 }); const session = makeSession({ workspacePath: "/workspace/test" }); @@ -1128,13 +1175,13 @@ describe("postLaunchSetup", () => { mockReadFile.mockRejectedValue(new Error("ENOENT")); // Before postLaunchSetup, binary is "codex" - expect(agent.getLaunchCommand(makeLaunchConfig())).toBe("codex"); + expect(agent.getLaunchCommand(makeLaunchConfig())).toBe("'codex'"); // After postLaunchSetup resolves the binary await agent.postLaunchSetup!(makeSession({ workspacePath: "/workspace/test" })); // Now getLaunchCommand should use the resolved binary - expect(agent.getLaunchCommand(makeLaunchConfig())).toBe("/opt/bin/codex"); + expect(agent.getLaunchCommand(makeLaunchConfig())).toBe("'/opt/bin/codex'"); }); }); diff --git a/packages/plugins/agent-codex/src/index.ts b/packages/plugins/agent-codex/src/index.ts index 25a9c681a..92bd95a0c 100644 --- a/packages/plugins/agent-codex/src/index.ts +++ b/packages/plugins/agent-codex/src/index.ts @@ -13,9 +13,11 @@ import { type WorkspaceHooksConfig, } from "@composio/ao-core"; import { execFile } from "node:child_process"; +import { createReadStream } from "node:fs"; import { writeFile, mkdir, readFile, readdir, rename, stat, lstat, open } from "node:fs/promises"; import { homedir } from "node:os"; import { basename, join } from "node:path"; +import { createInterface } from "node:readline"; import { promisify } from "node:util"; import { randomBytes } from "node:crypto"; @@ -416,66 +418,54 @@ async function findCodexSessionFile(workspacePath: string): Promise { + try { + const data: CodexSessionData = { model: null, threadId: null, inputTokens: 0, outputTokens: 0 }; + const rl = createInterface({ + input: createReadStream(filePath, { encoding: "utf-8" }), + crlfDelay: Infinity, + }); + + for await (const line of rl) { + const trimmed = line.trim(); + if (!trimmed) continue; + try { + const parsed: unknown = JSON.parse(trimmed); + if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) continue; + const entry = parsed as CodexJsonlLine; + + if (entry.type === "session_meta" && typeof entry.model === "string") { + data.model = entry.model; + } + if (typeof entry.threadId === "string" && entry.threadId) { + data.threadId = entry.threadId; + } + if (entry.type === "event_msg" && entry.msg?.type === "token_count") { + data.inputTokens += entry.msg.input_tokens ?? 0; + data.outputTokens += entry.msg.output_tokens ?? 0; + } + } catch { + // Skip malformed lines } - } catch { - // Skip malformed lines } + + return data; + } catch { + return null; } - return result; -} - -/** Extract model name from Codex session_meta entry */ -function extractCodexModel(lines: CodexJsonlLine[]): string | null { - for (const line of lines) { - if (line.type === "session_meta" && typeof line.model === "string") { - return line.model; - } - } - return null; -} - -/** Aggregate token usage from Codex token_count events */ -function extractCodexCost(lines: CodexJsonlLine[]): CostEstimate | undefined { - let inputTokens = 0; - let outputTokens = 0; - - for (const line of lines) { - if (line.type === "event_msg" && line.msg?.type === "token_count") { - inputTokens += line.msg.input_tokens ?? 0; - // cached_tokens is a subset of input_tokens (not additive) - outputTokens += line.msg.output_tokens ?? 0; - // reasoning_tokens is a subset of output_tokens (not additive) - } - } - - if (inputTokens === 0 && outputTokens === 0) return undefined; - - // Rough cost estimate using GPT-4o pricing as baseline - // ($2.50/1M input, $10/1M output) - const estimatedCostUsd = - (inputTokens / 1_000_000) * 2.5 + (outputTokens / 1_000_000) * 10.0; - - return { inputTokens, outputTokens, estimatedCostUsd }; -} - -/** Extract the Codex thread ID from JSONL (from thread_started or session_meta) */ -function extractCodexThreadId(lines: CodexJsonlLine[]): string | null { - for (const line of lines) { - if (typeof line.threadId === "string" && line.threadId) { - return line.threadId; - } - } - return null; } // ============================================================================= @@ -535,7 +525,7 @@ function createCodexAgent(): Agent { getLaunchCommand(config: AgentLaunchConfig): string { const binary = resolvedBinary ?? "codex"; - const parts: string[] = [binary]; + const parts: string[] = [shellEscape(binary)]; // Approval policy mapping — cast to string for forward-compat with // extended permission values not yet in AgentLaunchConfig type. @@ -688,24 +678,28 @@ function createCodexAgent(): Agent { const sessionFile = await findCodexSessionFile(session.workspacePath); if (!sessionFile) return null; - let content: string; - try { - content = await readFile(sessionFile, "utf-8"); - } catch { - return null; - } - - const lines = parseCodexJsonl(content); - if (lines.length === 0) return null; + // Stream the file line-by-line to avoid loading potentially huge + // rollout files (100 MB+) entirely into memory. + const data = await streamCodexSessionData(sessionFile); + if (!data) return null; const agentSessionId = basename(sessionFile, ".jsonl"); - const model = extractCodexModel(lines); + + const cost: CostEstimate | undefined = + data.inputTokens === 0 && data.outputTokens === 0 + ? undefined + : { + inputTokens: data.inputTokens, + outputTokens: data.outputTokens, + estimatedCostUsd: + (data.inputTokens / 1_000_000) * 2.5 + (data.outputTokens / 1_000_000) * 10.0, + }; return { - summary: model ? `Codex session (${model})` : null, + summary: data.model ? `Codex session (${data.model})` : null, summaryIsFallback: true, agentSessionId, - cost: extractCodexCost(lines), + cost, }; }, @@ -716,25 +710,16 @@ function createCodexAgent(): Agent { const sessionFile = await findCodexSessionFile(session.workspacePath); if (!sessionFile) return null; - let content: string; - try { - content = await readFile(sessionFile, "utf-8"); - } catch { - return null; - } - - const lines = parseCodexJsonl(content); - if (lines.length === 0) return null; - - // Extract the session/thread ID for native resume - const threadId = extractCodexThreadId(lines); - if (!threadId) return null; + // Stream the file line-by-line to avoid loading potentially huge + // rollout files (100 MB+) entirely into memory. + const data = await streamCodexSessionData(sessionFile); + if (!data?.threadId) return null; // Use Codex's native `resume` subcommand for proper conversation resume. // This restores the full thread state, not just a text prompt re-injection. // Flags are placed before the positional threadId for CLI parser compatibility. const binary = resolvedBinary ?? "codex"; - const parts: string[] = [binary, "resume"]; + const parts: string[] = [shellEscape(binary), "resume"]; // Add approval policy flags from project config const perms = project.agentConfig?.permissions as string | undefined; @@ -746,7 +731,7 @@ function createCodexAgent(): Agent { parts.push("--ask-for-approval", "untrusted"); } - const effectiveModel = project.agentConfig?.model ?? extractCodexModel(lines); + const effectiveModel = project.agentConfig?.model ?? data.model; if (effectiveModel) { parts.push("--model", shellEscape(effectiveModel as string)); @@ -756,7 +741,7 @@ function createCodexAgent(): Agent { } // Positional threadId goes last, after all flags - parts.push(shellEscape(threadId)); + parts.push(shellEscape(data.threadId)); return parts.join(" "); },