mirror of
https://github.com/getpaseo/paseo.git
synced 2026-10-14 13:22:37 +00:00
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:
@@ -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);
|
||||
});
|
||||
Reference in New Issue
Block a user