diff --git a/src/cli.mjs b/src/cli.mjs index 89e35e2..baf7ecc 100755 --- a/src/cli.mjs +++ b/src/cli.mjs @@ -42,6 +42,7 @@ if (options.command === "snapshot") { cdpPort: options.cdpPort || 9334, signal: controller.signal, }); + if (controller.signal.aborted) process.exitCode = 130; } else if (options.command === "remove") { const modulePath = new URL("./codex/injector.mjs", import.meta.url); const { removeCodexMeter } = await import(modulePath); diff --git a/src/codex/cdp-client.mjs b/src/codex/cdp-client.mjs index 09f3cd4..bb5a348 100644 --- a/src/codex/cdp-client.mjs +++ b/src/codex/cdp-client.mjs @@ -2,7 +2,6 @@ function withTimeout(promise, timeoutMs, message) { let timer; const timeout = new Promise((_, reject) => { timer = setTimeout(() => reject(new Error(message)), timeoutMs); - timer.unref?.(); }); return Promise.race([promise, timeout]).finally(() => clearTimeout(timer)); } @@ -55,6 +54,11 @@ export class CdpClient { } async call(method, params = {}, { timeoutMs = 5_000 } = {}) { + if (this.socket.readyState !== 1) { + throw new Error( + `CDP socket is not open (readyState ${this.socket.readyState})`, + ); + } const id = this.nextId++; const response = new Promise((resolve, reject) => { this.pending.set(id, { resolve, reject }); diff --git a/src/codex/injector.mjs b/src/codex/injector.mjs index 2290e7e..9ef07e3 100644 --- a/src/codex/injector.mjs +++ b/src/codex/injector.mjs @@ -323,56 +323,89 @@ export async function runCodexInjector({ try { await verifyMacListenerOwner(cdpPort, { appPath }); + let cdpVerified = true; + let lastCdpErrorLoggedMs = 0; while (!stopping) { - const nowMs = Date.now(); - if (nowMs - lastDiscoveryMs >= targetDiscoveryIntervalMs) { - const targets = await listTargets(cdpPort); - for (const target of targets) { - if (attached.has(target.id)) continue; - const connection = await attachCodexTarget(target, payload).catch(() => null); - if (connection) attached.set(target.id, connection); + try { + if (!cdpVerified) { + await verifyMacListenerOwner(cdpPort, { appPath }); + cdpVerified = true; + console.error( + `[token-meter] CDP listener on port ${cdpPort} is available again; resuming.`, + ); } - lastDiscoveryMs = nowMs; - } - - const probes = []; - for (const connection of attached.values()) { - try { - connection.probe = await connection.client.evaluate( - buildSessionProbeExpression(), + const nowMs = Date.now(); + if (nowMs - lastDiscoveryMs >= targetDiscoveryIntervalMs) { + const targets = await listTargets(cdpPort); + for (const target of targets) { + if (attached.has(target.id)) continue; + const connection = await attachCodexTarget(target, payload).catch(() => null); + if (connection) attached.set(target.id, connection); + } + lastDiscoveryMs = nowMs; + } + } catch (error) { + cdpVerified = false; + for (const connection of attached.values()) connection.failed = true; + const nowMs = Date.now(); + if (nowMs - lastCdpErrorLoggedMs >= 60_000) { + lastCdpErrorLoggedMs = nowMs; + console.error( + `[token-meter] CDP temporarily unavailable (${error?.message ?? error}); retrying.`, ); - probes.push(connection.probe); - } catch { - connection.failed = true; } } - for (const [id, connection] of attached) { - if (connection.failed || !connection.probe?.eligible) { - connection.client.close(); - attached.delete(id); + + if (cdpVerified && !stopping) { + const probes = []; + for (const connection of attached.values()) { + try { + connection.probe = await connection.client.evaluate( + buildSessionProbeExpression(), + ); + probes.push(connection.probe); + } catch { + connection.failed = true; + } + } + for (const [id, connection] of attached) { + if (connection.failed || !connection.probe?.eligible) { + connection.client.close(); + attached.delete(id); + } } - } - const activeThreadIds = probes.map((probe) => probe.threadId).filter(Boolean); - store.historyFileLimit = warmedHistoryFileLimit; - const files = await store.refresh({ activeThreadIds }); - for (const connection of attached.values()) { - const snapshot = engine.snapshot(files, { - threadId: connection.probe.threadId, - nowMs, - }); - snapshot.binding = { - source: connection.probe.bindingSource, - exact: Boolean(connection.probe.threadId), - }; - await connection.client.evaluate(updateExpression(snapshot)).catch(() => { - connection.failed = true; - }); + const activeThreadIds = probes.map((probe) => probe.threadId).filter(Boolean); + store.historyFileLimit = warmedHistoryFileLimit; + try { + const files = await store.refresh({ activeThreadIds }); + for (const connection of attached.values()) { + const snapshot = engine.snapshot(files, { + threadId: connection.probe.threadId, + nowMs: Date.now(), + }); + snapshot.binding = { + source: connection.probe.bindingSource, + exact: Boolean(connection.probe.threadId), + }; + await connection.client.evaluate(updateExpression(snapshot)).catch(() => { + connection.failed = true; + }); + } + } catch (error) { + const nowMs = Date.now(); + if (nowMs - lastCdpErrorLoggedMs >= 60_000) { + lastCdpErrorLoggedMs = nowMs; + console.error( + `[token-meter] poll error (${error?.message ?? error}); continuing.`, + ); + } + } + warmedHistoryFileLimit = Math.min( + historyFileLimit, + warmedHistoryFileLimit + historyFilesPerPoll, + ); } - warmedHistoryFileLimit = Math.min( - historyFileLimit, - warmedHistoryFileLimit + historyFilesPerPoll, - ); await waitForNextPoll(pollIntervalMs, signal); } diff --git a/src/codex/session-probe.mjs b/src/codex/session-probe.mjs index 3de21f0..31e56cb 100644 --- a/src/codex/session-probe.mjs +++ b/src/codex/session-probe.mjs @@ -35,11 +35,26 @@ export function buildSessionProbeExpression() { const routeId = normalizeThreadId( routeMatch == null ? null : decodeURIComponent(routeMatch[1]) ); + const optimisticMatch = String( + activeRow?.getAttribute('data-app-action-sidebar-thread-id') ?? '' + ) + .replace(/^local:/, '') + .match(/^client-new-thread:([0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12})$/i); + const contentConversationId = normalizeThreadId( + document.querySelector('[data-response-annotation-conversation]') + ?.getAttribute('data-response-annotation-conversation') ?? + document.querySelector('[data-above-composer-conversation-id]') + ?.getAttribute('data-above-composer-conversation-id') ?? + null + ); let threadId = null; let bindingSource = null; if (activeId != null) { threadId = activeId; bindingSource = 'active-sidebar-row'; + } else if (optimisticMatch != null && contentConversationId != null) { + threadId = contentConversationId; + bindingSource = 'active-composer-conversation'; } else if (routeId != null) { threadId = routeId; bindingSource = 'thread-route'; @@ -54,6 +69,7 @@ export function buildSessionProbeExpression() { 'textarea, [contenteditable="true"], [data-app-action-composer]' )), activeThread: Boolean(activeRow), + conversationId: Boolean(contentConversationId), }; const markerCount = Object.values(markers).filter(Boolean).length; const eligible = diff --git a/src/core/rollout-store.mjs b/src/core/rollout-store.mjs index 30839d7..cfd265d 100644 --- a/src/core/rollout-store.mjs +++ b/src/core/rollout-store.mjs @@ -76,7 +76,13 @@ export function parseRolloutLine(line) { } async function defaultReadRange(filePath, { length, position }) { - const handle = await open(filePath, "r"); + let handle; + try { + handle = await open(filePath, "r"); + } catch (error) { + if (error?.code === "ENOENT") return Buffer.alloc(0); + throw error; + } try { const buffer = Buffer.allocUnsafe(length); const { bytesRead } = await handle.read(buffer, 0, length, position); @@ -86,6 +92,15 @@ async function defaultReadRange(filePath, { length, position }) { } } +async function statIfExists(targetPath) { + try { + return await stat(targetPath); + } catch (error) { + if (error?.code === "ENOENT") return null; + throw error; + } +} + function createFileState(filePath, discoveredId, modifiedMs) { return { path: filePath, @@ -121,7 +136,8 @@ async function walk(directory, result) { } const match = entry.name.match(ROLLOUT_FILE); if (!entry.isFile() || match == null) return; - const fileStat = await stat(fullPath); + const fileStat = await statIfExists(fullPath); + if (fileStat == null) return; result.push({ path: fullPath, discoveredId: match[1], @@ -264,7 +280,8 @@ export class RolloutStore { async #readMetadata(file) { if (file.meta != null) return; - const fileStat = await stat(file.path); + const fileStat = await statIfExists(file.path); + if (fileStat == null) return; const decoder = new StringDecoder("utf8"); let position = 0; let source = ""; @@ -284,7 +301,8 @@ export class RolloutStore { } async #readAppended(file) { - const fileStat = await stat(file.path); + const fileStat = await statIfExists(file.path); + if (fileStat == null) return; if (fileStat.size < file.offset) { file.offset = 0; file.remainder = ""; diff --git a/test/cdp-client.test.mjs b/test/cdp-client.test.mjs index 3dea7a2..1572274 100644 --- a/test/cdp-client.test.mjs +++ b/test/cdp-client.test.mjs @@ -1,6 +1,30 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { isLoopbackWebSocketUrl } from "../src/codex/cdp-client.mjs"; +import { + CdpClient, + isLoopbackWebSocketUrl, +} from "../src/codex/cdp-client.mjs"; + +class FakeSocket extends EventTarget { + constructor() { + super(); + this.readyState = 1; // WebSocket.OPEN + this.sent = []; + } + + send(data) { + this.sent.push(String(data)); + } + + close() { + this.readyState = 3; // WebSocket.CLOSED + this.dispatchEvent(new Event("close")); + } + + receive(raw) { + this.dispatchEvent(new MessageEvent("message", { data: raw })); + } +} test("CDP client accepts loopback targets only", () => { assert.equal(isLoopbackWebSocketUrl("ws://127.0.0.1:9334/devtools/page/1"), true); @@ -8,3 +32,42 @@ test("CDP client accepts loopback targets only", () => { assert.equal(isLoopbackWebSocketUrl("ws://192.168.1.4:9334/devtools/page/1"), false); assert.equal(isLoopbackWebSocketUrl("wss://example.com/devtools/page/1"), false); }); + +test("CDP client rejects a pending call when the socket closes", async () => { + const socket = new FakeSocket(); + const client = new CdpClient(socket); + const call = client.call("Runtime.evaluate", { expression: "1" }); + socket.close(); + await assert.rejects(call, /CDP target closed/); +}); + +test("CDP client refuses calls on a closed socket", async () => { + const socket = new FakeSocket(); + const client = new CdpClient(socket); + socket.close(); + await assert.rejects( + client.call("Runtime.evaluate", { expression: "1" }), + /not open/, + ); +}); + +test("CDP client resolves calls from matching responses", async () => { + const socket = new FakeSocket(); + const client = new CdpClient(socket); + const call = client.call("Runtime.evaluate", { expression: "1" }); + const sent = JSON.parse(socket.sent[0]); + assert.equal(sent.method, "Runtime.evaluate"); + socket.receive( + JSON.stringify({ id: sent.id, result: { result: { value: 42 } } }), + ); + assert.equal((await call).result.value, 42); +}); + +test("CDP client times out hung calls instead of hanging forever", async () => { + const socket = new FakeSocket(); + const client = new CdpClient(socket); + await assert.rejects( + client.call("Runtime.evaluate", { expression: "1" }, { timeoutMs: 50 }), + /timed out/, + ); +}); diff --git a/test/rollout-store.test.mjs b/test/rollout-store.test.mjs index bb29da6..4c35647 100644 --- a/test/rollout-store.test.mjs +++ b/test/rollout-store.test.mjs @@ -276,3 +276,57 @@ test("a filesystem notification makes a new child Agent discoverable immediately [rootId, childId], ); }); + +test("refresh survives a rollout file removed between discovery and read", async (context) => { + const directory = await mkdtemp(path.join(os.tmpdir(), "token-meter-rollout-")); + context.after(() => rm(directory, { recursive: true, force: true })); + const threadId = "019fc0bf-d10c-7472-bb0e-fd6f0df8ab3e"; + const filePath = path.join( + directory, + `rollout-2026-08-01T21-33-44-${threadId}.jsonl`, + ); + await writeFile( + filePath, + `${JSON.stringify({ + timestamp: "2026-08-01T21:33:44.000Z", + type: "session_meta", + payload: { + id: threadId, + session_id: threadId, + source: "vscode", + thread_source: "user", + }, + })}\n`, + ); + + const store = new RolloutStore({ + sessionsDirectory: directory, + discoveryIntervalMs: 60_000, + }); + + // Discover first, then remove the file before any metadata/appended read. + await store.discover({ force: true }); + await rm(filePath); + const files = await store.refresh({ activeThreadIds: [threadId] }); + assert.deepEqual(files, []); + + // Also exercise the appended-read path: cached state, then file disappears. + await writeFile( + filePath, + `${JSON.stringify({ + timestamp: "2026-08-01T21:33:44.000Z", + type: "session_meta", + payload: { + id: threadId, + session_id: threadId, + source: "vscode", + thread_source: "user", + }, + })}\n`, + ); + await store.markDiscoveryDirty(); + await store.refresh({ activeThreadIds: [threadId] }); + await rm(filePath); + const cached = await store.refresh({ activeThreadIds: [threadId] }); + assert.equal(cached.length, 1); +}); diff --git a/test/session-probe.test.mjs b/test/session-probe.test.mjs index d29db0c..d8134b4 100644 --- a/test/session-probe.test.mjs +++ b/test/session-probe.test.mjs @@ -7,13 +7,24 @@ import { const threadId = "11111111-2222-4333-8444-555555555555"; -function runProbe({ activeId = threadId, pathname = `/thread/${threadId}` } = {}) { +function runProbe({ + activeId = threadId, + pathname = `/thread/${threadId}`, + contentConversationId = null, +} = {}) { const mainSurface = {}; const activeRow = { getAttribute(name) { return name === "data-app-action-sidebar-thread-id" ? activeId : null; }, }; + const contentElement = contentConversationId == null ? null : { + getAttribute(name) { + return name === "data-response-annotation-conversation" + ? contentConversationId + : null; + }, + }; const document = { querySelector(selector) { if (selector.includes("avatar-overlay")) return null; @@ -21,6 +32,8 @@ function runProbe({ activeId = threadId, pathname = `/thread/${threadId}` } = {} if (selector === "aside.app-shell-left-panel") return {}; if (selector.includes("main[data-app-shell-main-surface]")) return mainSurface; if (selector.startsWith("textarea")) return {}; + if (selector === "[data-response-annotation-conversation]") return contentElement; + if (selector === "[data-above-composer-conversation-id]") return contentElement; return null; }, }; @@ -55,6 +68,27 @@ test("session probe normalizes Codex local thread identifiers", () => { assert.equal(result.bindingSource, "active-sidebar-row"); }); +test("session probe resolves the real conversation id for a new optimistic thread", () => { + const realId = "22222222-3333-4444-8555-666666666666"; + const result = runProbe({ + activeId: `local:client-new-thread:${threadId}`, + pathname: "/index.html", + contentConversationId: realId, + }); + assert.equal(result.eligible, true); + assert.equal(result.threadId, realId); + assert.equal(result.bindingSource, "active-composer-conversation"); +}); + +test("session probe refuses an optimistic thread without a real conversation id", () => { + const result = runProbe({ + activeId: `local:client-new-thread:${threadId}`, + pathname: "/index.html", + }); + assert.equal(result.threadId, null); + assert.equal(result.eligible, true); +}); + test("session probe falls back to the exact route and rejects ambiguous labels", () => { const result = runProbe({ activeId: "not-a-thread-id" }); assert.equal(result.threadId, threadId);