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.
This commit is contained in:
paseo-bot[bot]
2026-09-30 03:05:07 +02:00
parent 5599f9e567
commit 4943b09764
6 changed files with 62 additions and 4 deletions
+1 -1
View File
@@ -381,7 +381,7 @@ const AgentCapabilityFlagsSchema: z.ZodType<AgentCapabilityFlags> = z
})
.catchall(z.boolean());
const AgentUsageSchema: z.ZodType<AgentUsage> = z.object({
export const AgentUsageSchema: z.ZodType<AgentUsage> = z.object({
inputTokens: z.number().optional(),
cachedInputTokens: z.number().optional(),
outputTokens: z.number().optional(),
@@ -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");
@@ -1310,6 +1310,7 @@ export class AgentManager {
workspaceId?: string;
owner?: AgentOwner;
attention?: AttentionState;
lastUsage?: AgentUsage;
},
resumeOptions?: AgentResumeSessionOptions,
): Promise<ManagedAgent> {
@@ -1342,6 +1343,7 @@ export class AgentManager {
workspaceId?: string;
owner?: AgentOwner;
attention?: AttentionState;
lastUsage?: AgentUsage;
},
resumeOptions?: AgentResumeSessionOptions,
): Promise<ManagedAgent> {
@@ -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,
@@ -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,
@@ -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(),
@@ -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);
});