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); +});