From 4943b097649f226281cc273435aa9d1d60f94cd0 Mon Sep 17 00:00:00 2001 From: "paseo-bot[bot]" <266920839+paseo-bot[bot]@users.noreply.github.com> Date: Wed, 30 Sep 2026 03:05:07 +0200 Subject: [PATCH] fix(server): keep an agent's context usage across daemon restarts The last usage report lived only on the in-memory agent. The stored agent record never saved it, so after a restart both the stored snapshot and the resumed agent reported no usage and the composer hid the context meter until the next turn. Save it on the record and carry it into the stored snapshot and the resumed session. --- packages/protocol/src/messages.ts | 2 +- .../server/src/server/agent/agent-loading.ts | 6 ++- .../server/src/server/agent/agent-manager.ts | 4 +- .../src/server/agent/agent-projections.ts | 2 + .../server/src/server/agent/agent-storage.ts | 3 +- .../usage-survives-restart.e2e.test.ts | 49 +++++++++++++++++++ 6 files changed, 62 insertions(+), 4 deletions(-) create mode 100644 packages/server/src/server/daemon-e2e/usage-survives-restart.e2e.test.ts diff --git a/packages/protocol/src/messages.ts b/packages/protocol/src/messages.ts index 589137d720..e68cd32443 100644 --- a/packages/protocol/src/messages.ts +++ b/packages/protocol/src/messages.ts @@ -381,7 +381,7 @@ const AgentCapabilityFlagsSchema: z.ZodType = z }) .catchall(z.boolean()); -const AgentUsageSchema: z.ZodType = z.object({ +export const AgentUsageSchema: z.ZodType = z.object({ inputTokens: z.number().optional(), cachedInputTokens: z.number().optional(), outputTokens: z.number().optional(), diff --git a/packages/server/src/server/agent/agent-loading.ts b/packages/server/src/server/agent/agent-loading.ts index 248e8d2b1b..4593f81e22 100644 --- a/packages/server/src/server/agent/agent-loading.ts +++ b/packages/server/src/server/agent/agent-loading.ts @@ -110,7 +110,11 @@ export async function ensureAgentLoaded( handle, buildConfigOverrides(record), agentId, - { ...extractTimestamps(record), attention: extractAttention(record) }, + { + ...extractTimestamps(record), + attention: extractAttention(record), + lastUsage: record.lastUsage, + }, record.archivedAt ? { purpose: "history" } : undefined, ); deps.logger.info({ agentId, provider: record.provider }, "Agent resumed from persistence"); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 87d3b290aa..ee809b648a 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -1310,6 +1310,7 @@ export class AgentManager { workspaceId?: string; owner?: AgentOwner; attention?: AttentionState; + lastUsage?: AgentUsage; }, resumeOptions?: AgentResumeSessionOptions, ): Promise { @@ -1342,6 +1343,7 @@ export class AgentManager { workspaceId?: string; owner?: AgentOwner; attention?: AttentionState; + lastUsage?: AgentUsage; }, resumeOptions?: AgentResumeSessionOptions, ): Promise { @@ -1904,7 +1906,7 @@ export class AgentManager { persistence: record.persistence ?? null, historyPrimed: true, lastUserMessageAt: record.lastUserMessageAt ? new Date(record.lastUserMessageAt) : null, - lastUsage: undefined, + lastUsage: record.lastUsage, lastError: record.lastError ?? undefined, attention, internal: record.internal, diff --git a/packages/server/src/server/agent/agent-projections.ts b/packages/server/src/server/agent/agent-projections.ts index 05f932aea9..49c06124d7 100644 --- a/packages/server/src/server/agent/agent-projections.ts +++ b/packages/server/src/server/agent/agent-projections.ts @@ -90,6 +90,7 @@ export function toStoredAgentRecord( runtimeInfo, features: normalizeFeatures(agent.features), persistence, + lastUsage: sanitizeUsage(agent.lastUsage), lastError: agent.lastError ?? undefined, requiresAttention: agent.attention.requiresAttention, attentionReason: agent.attention.requiresAttention ? agent.attention.attentionReason : null, @@ -243,6 +244,7 @@ export function buildStoredAgentPayload( pendingPermissions: [], persistence, title: record.title ?? null, + ...(record.lastUsage ? { lastUsage: record.lastUsage } : {}), requiresAttention: record.requiresAttention ?? false, attentionReason: record.attentionReason ?? null, attentionTimestamp: record.attentionTimestamp ?? null, diff --git a/packages/server/src/server/agent/agent-storage.ts b/packages/server/src/server/agent/agent-storage.ts index b41bc75cfc..1696ace296 100644 --- a/packages/server/src/server/agent/agent-storage.ts +++ b/packages/server/src/server/agent/agent-storage.ts @@ -4,7 +4,7 @@ import { z } from "zod"; import type { Logger } from "pino"; import { writeJsonFileAtomic } from "../atomic-file.js"; -import { AgentFeatureSchema, AgentStatusSchema } from "../messages.js"; +import { AgentFeatureSchema, AgentStatusSchema, AgentUsageSchema } from "../messages.js"; import { toStoredAgentRecord } from "./agent-projections.js"; import type { ManagedAgent } from "./agent-manager.js"; import type { AgentSessionConfig } from "./agent-sdk-types.js"; @@ -68,6 +68,7 @@ const STORED_AGENT_SCHEMA = z.object({ .optional(), features: z.array(AgentFeatureSchema).optional(), persistence: PERSISTENCE_HANDLE_SCHEMA, + lastUsage: AgentUsageSchema.optional(), lastError: z.string().nullable().optional(), requiresAttention: z.boolean().optional(), attentionReason: z.enum(["finished", "error", "permission"]).nullable().optional(), diff --git a/packages/server/src/server/daemon-e2e/usage-survives-restart.e2e.test.ts b/packages/server/src/server/daemon-e2e/usage-survives-restart.e2e.test.ts new file mode 100644 index 0000000000..2e90af3ded --- /dev/null +++ b/packages/server/src/server/daemon-e2e/usage-survives-restart.e2e.test.ts @@ -0,0 +1,49 @@ +import { afterEach, describe, expect, test } from "vitest"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import path from "node:path"; + +import { createDaemonTestContext, type DaemonTestContext } from "../test-utils/index.js"; + +describe("daemon E2E - usage after restart", () => { + let ctx: DaemonTestContext | null = null; + const dirs: string[] = []; + + afterEach(async () => { + await ctx?.cleanup(); + ctx = null; + for (const dir of dirs.splice(0)) { + rmSync(dir, { recursive: true, force: true }); + } + }, 60_000); + + test("a finished agent reports its last usage after the daemon restarts", async () => { + const cwd = mkdtempSync(path.join(tmpdir(), "usage-restart-cwd-")); + const paseoHomeRoot = mkdtempSync(path.join(tmpdir(), "usage-restart-home-")); + dirs.push(cwd, paseoHomeRoot); + + ctx = await createDaemonTestContext({ paseoHomeRoot, cleanup: false }); + const agent = await ctx.client.createAgent({ + provider: "codex", + cwd, + title: "Usage restart agent", + modeId: "full-access", + }); + await ctx.client.sendMessage(agent.id, "Say 'done' and nothing else"); + const finished = await ctx.client.waitForFinish(agent.id, 5_000); + expect(finished.status).toBe("idle"); + const usageBeforeRestart = finished.final?.lastUsage; + expect(usageBeforeRestart).toBeDefined(); + + await ctx.cleanup(); + ctx = await createDaemonTestContext({ paseoHomeRoot, cleanup: false }); + + const stored = await ctx.client.fetchAgent(agent.id); + expect(stored?.agent.lastUsage).toEqual(usageBeforeRestart); + + await ctx.client.fetchAgentTimeline(agent.id, { direction: "tail", limit: 0 }); + const reopened = await ctx.client.fetchAgent(agent.id); + expect(reopened?.agent.status).toBe("idle"); + expect(reopened?.agent.lastUsage).toEqual(usageBeforeRestart); + }, 30_000); +});