talon-agent 3.34.0 → 3.35.0

package/src/backend/codex/handler/rollout-accounting.ts

@@ -0,0 +1,137 @@
+/**
+ * Rollout-JSONL usage accounting for one Codex turn.
+ *
+ * Codex's `turn.completed.usage` is CUMULATIVE across every API call in
+ * the turn — never in the per-turn units the shared stream state (and
+ * everything downstream: /status, the companion's per-message counts)
+ * speaks. The rollout JSONL's totals diffed against the pre-turn baseline
+ * are this turn's real usage; that is the ONLY authoritative source. The
+ * SDK figure is a last-resort fallback when the rollout can't be read,
+ * and it overstates multi-call turns.
+ *
+ * The same baseline serves the throttled mid-turn poll (live /status) and
+ * the final settlement, which the success path and the terminal-failure
+ * path both run so failed turns account for the tokens they burned.
+ */
+
+import type { Usage } from "@openai/codex-sdk";
+import {
+  pushLiveUsage,
+  recordTokens,
+  type StreamState,
+} from "../../shared/index.js";
+import { CODEX_LIVE_POLL_INTERVAL_MS } from "../constants.js";
+import {
+  readLastRolloutSnapshot,
+  type CodexRolloutSnapshot,
+} from "../token-usage.js";
+
+type RolloutTotals = NonNullable<CodexRolloutSnapshot["totals"]>;
+
+export type RolloutAccounting = {
+  /** Thread id once `thread.started` lands; polls are no-ops before it. */
+  threadId: string | undefined;
+  /** Throttled mid-turn rollout tail — fire-and-forget, never throws. */
+  pollLive(): void;
+  /** Final authoritative settlement; `usage` is the SDK's fallback figure. */
+  settle(usage: Usage | null): Promise<void>;
+};
+
+/**
+ * Capture the cumulative token totals BEFORE the turn runs.
+ * `total_token_usage` accumulates across the whole session file, so this
+ * turn's usage = post-turn totals minus this baseline. Fresh threads have
+ * no rollout yet → zero baseline. `null` = resumed thread whose baseline
+ * couldn't be read.
+ */
+export async function createRolloutAccounting(inputs: {
+  state: StreamState;
+  sessionId: string | undefined;
+}): Promise<RolloutAccounting> {
+  const { state } = inputs;
+  const baseline: RolloutTotals | null = inputs.sessionId
+    ? ((await readLastRolloutSnapshot(inputs.sessionId).catch(() => null))
+        ?.totals ?? null)
+    : { inputTokens: 0, cachedInputTokens: 0, outputTokens: 0 };
+
+  let pollInFlight = false;
+  let lastPollAt = 0;
+
+  const accounting: RolloutAccounting = {
+    threadId: undefined,
+
+    // The Codex CLI appends a `token_count` event to the rollout JSONL
+    // after every API call, so tailing it during the turn gives live
+    // context-fill / token / API-call stats long before `turn.completed`.
+    // In-flight guard — never blocks the event loop.
+    pollLive() {
+      if (!accounting.threadId || pollInFlight) return;
+      const now = Date.now();
+      if (now - lastPollAt < CODEX_LIVE_POLL_INTERVAL_MS) return;
+      pollInFlight = true;
+      lastPollAt = now;
+      readLastRolloutSnapshot(accounting.threadId)
+        .then((snap) => {
+          if (!snap) return;
+          applyContext(state, snap);
+          // Same delta-vs-baseline math as the final settlement, which
+          // recomputes and overwrites — a torn mid-turn read can't
+          // corrupt the committed numbers.
+          if (snap.totals && baseline) {
+            state.sdkInputTokens = Math.max(
+              0,
+              snap.totals.inputTokens - baseline.inputTokens,
+            );
+            state.sdkOutputTokens = Math.max(
+              0,
+              snap.totals.outputTokens - baseline.outputTokens,
+            );
+            state.sdkCacheRead = Math.max(
+              0,
+              snap.totals.cachedInputTokens - baseline.cachedInputTokens,
+            );
+          }
+          pushLiveUsage(state);
+        })
+        .catch(() => {})
+        .finally(() => {
+          pollInFlight = false;
+        });
+    },
+
+    async settle(usage) {
+      const last = accounting.threadId
+        ? await readLastRolloutSnapshot(accounting.threadId).catch(() => null)
+        : null;
+      if (last) applyContext(state, last);
+      if (last?.totals && baseline) {
+        recordTokens(state, {
+          inputTokens: last.totals.inputTokens - baseline.inputTokens,
+          outputTokens: last.totals.outputTokens - baseline.outputTokens,
+          cacheRead: last.totals.cachedInputTokens - baseline.cachedInputTokens,
+          cacheWrite: 0, // Codex doesn't report cache writes
+        });
+      } else if (usage) {
+        recordTokens(state, {
+          inputTokens: usage.input_tokens,
+          outputTokens: usage.output_tokens,
+          cacheRead: usage.cached_input_tokens,
+          cacheWrite: 0, // Codex doesn't report cache writes
+        });
+      }
+    },
+  };
+  return accounting;
+}
+
+function applyContext(state: StreamState, snap: CodexRolloutSnapshot): void {
+  if (snap.usage) {
+    state.contextTokens = snap.usage.contextTokens;
+    if (snap.usage.contextWindow) {
+      state.contextWindow = snap.usage.contextWindow;
+    }
+  }
+  if (typeof snap.numApiCalls === "number") {
+    state.numApiCalls = snap.numApiCalls;
+  }
+}

package/src/backend/codex/plan-usage.ts

@@ -11,7 +11,7 @@
  * missing or expired token, or any transport failure.
  */
 
-import { readFile } from "node:fs/promises";
+import { readFile, stat } from "node:fs/promises";
 import { homedir } from "node:os";
 import { join } from "node:path";
 import { logWarn } from "../../util/log.js";
@@ -26,6 +26,12 @@ const CACHE_TTL_MS = 60_000;
 
 let cache: { value: PlanUsage; fetchedAt: number } | undefined;
 let inFlight: Promise<PlanUsage | undefined> | undefined;
+/**
+ * `auth.json` mtime of a token the endpoint rejected with 401. An
+ * invalidated login stays invalid until `codex login` rewrites the file,
+ * so the poller warns once and stops calling until the mtime moves.
+ */
+let rejectedAuthMtimeMs: number | undefined;
 
 function authPath(): string {
   const home = process.env.CODEX_HOME?.trim();
@@ -37,11 +43,18 @@ function authPath(): string {
 interface CodexAuth {
   accessToken: string;
   accountId?: string;
+  /** `auth.json` mtime — identifies the login the token came from. */
+  mtimeMs: number;
 }
 
 async function readAuth(): Promise<CodexAuth | undefined> {
   try {
-    const parsed = JSON.parse(await readFile(authPath(), "utf8")) as {
+    const path = authPath();
+    const [raw, stats] = await Promise.all([
+      readFile(path, "utf8"),
+      stat(path),
+    ]);
+    const parsed = JSON.parse(raw) as {
       auth_mode?: string;
       tokens?: { access_token?: string; account_id?: string };
     };
@@ -52,6 +65,7 @@ async function readAuth(): Promise<CodexAuth | undefined> {
       ...(parsed.tokens?.account_id
         ? { accountId: parsed.tokens.account_id }
         : {}),
+      mtimeMs: stats.mtimeMs,
     };
   } catch {
     return undefined;
@@ -123,6 +137,8 @@ export function parseCodexUsage(body: unknown): PlanUsage | undefined {
 async function load(): Promise<PlanUsage | undefined> {
   const auth = await readAuth();
   if (!auth) return undefined;
+  if (auth.mtimeMs === rejectedAuthMtimeMs) return undefined;
+  rejectedAuthMtimeMs = undefined;
 
   try {
     const res = await fetch(USAGE_ENDPOINT, {
@@ -133,6 +149,15 @@ async function load(): Promise<PlanUsage | undefined> {
       },
       signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
     });
+    if (res.status === 401) {
+      rejectedAuthMtimeMs = auth.mtimeMs;
+      logWarn(
+        "agent",
+        "codex usage: endpoint returned 401 — Codex login expired; " +
+          "run `codex login` (not retried until auth.json changes)",
+      );
+      return undefined;
+    }
     if (!res.ok) {
       logWarn("agent", `codex usage: endpoint returned ${res.status}`);
       return undefined;

package/src/backend/kilo/handler/message.ts

@@ -7,6 +7,7 @@ import type { QueryParams, QueryResult } from "../../shared/handler-types.js";
 import { runRemoteChatTurn } from "../../remote-server/chat-turn.js";
 import {
   ensureServer,
+  trackActiveTurn,
   ensureSession,
   ensureChatMcpServer,
   ensurePluginMcpServers,
@@ -26,6 +27,7 @@ export function handleMessage(params: QueryParams): Promise<QueryResult> {
       label: "Kilo",
       getConfig,
       ensureServer,
+      trackActiveTurn,
       parseModelSelection: parseStoredKiloModelSelection,
       resolveProviderID,
       ensureSession,

package/src/backend/kilo/server.ts

@@ -82,6 +82,7 @@ export const stopKiloServer = kilo.stop;
 export const {
   onServerStop,
   ensureServer,
+  trackActiveTurn,
   ensureChatMcpServer,
   ensurePluginMcpServers,
   buildToolOverrides,