diff --git a/packages/server/src/server/agent/providers/omp/agent.test.ts b/packages/server/src/server/agent/providers/omp/agent.test.ts index c0e6193bb5..a1dc9130a8 100644 --- a/packages/server/src/server/agent/providers/omp/agent.test.ts +++ b/packages/server/src/server/agent/providers/omp/agent.test.ts @@ -215,6 +215,55 @@ function createToolCatalog(): PaseoToolCatalog { } describe("OMP agent client and session", () => { + test("failed manual compact closes its progress row and shows the error", async () => { + const omp = new OmpHarness(); + await omp.start(); + + const events = await omp.compact({}, new Error("compaction refused")); + + expect(events).toEqual([ + expect.objectContaining({ + item: { type: "compaction", status: "loading", trigger: "manual" }, + }), + expect.objectContaining({ + item: { type: "compaction", status: "completed", trigger: "manual" }, + }), + expect.objectContaining({ + item: { + type: "assistant_message", + text: "[Error] Failed to compact context: compaction refused", + }, + }), + ]); + }); + + test("manual compact shows its progress and publishes the compacted context usage", async () => { + const omp = new OmpHarness(); + await omp.start(); + + const events = await omp.compact({ + contextUsage: { tokens: 13_447, contextWindow: 272_000 }, + }); + + expect(events).toEqual([ + expect.objectContaining({ + type: "timeline", + item: { type: "compaction", status: "loading", trigger: "manual" }, + }), + expect.objectContaining({ + type: "timeline", + item: { type: "compaction", status: "completed", trigger: "manual" }, + }), + expect.objectContaining({ + type: "usage_updated", + usage: expect.objectContaining({ + contextWindowUsedTokens: 13_447, + contextWindowMaxTokens: 272_000, + }), + }), + ]); + }); + test("owns launch configuration and registers native host tools", async () => { const omp = new OmpHarness(); await omp.start({ modeId: "ask" }, createToolCatalog()); diff --git a/packages/server/src/server/agent/providers/omp/agent.ts b/packages/server/src/server/agent/providers/omp/agent.ts index 2dc6c58dea..5413364cdb 100644 --- a/packages/server/src/server/agent/providers/omp/agent.ts +++ b/packages/server/src/server/agent/providers/omp/agent.ts @@ -714,9 +714,7 @@ export class OmpAgentSession implements AgentSession { private readonly pendingPromptResults = new Map(); private pendingNoTurnCompletionAbort: AbortController | null = null; private lastKnownThinkingOptionId: string | null; - private outOfBandCompactionEmit: ((event: AgentStreamEvent) => void) | null = null; - private outOfBandCompactionStarted = false; - private outOfBandCompactionCompleted = false; + private manualCompactionRunning = false; private commandCache: AgentSlashCommand[] | null = null; private readonly subagentIndex = new OmpSubagentIndex(); private readonly subagentCardTracker: OmpSubagentCardTracker; @@ -1407,30 +1405,24 @@ export class OmpAgentSession implements AgentSession { customInstructions: string | undefined, emit: (event: AgentStreamEvent) => void, ): Promise { - if (this.outOfBandCompactionEmit) { + if (this.manualCompactionRunning) { throw new Error("An OMP compact command is already running"); } - this.outOfBandCompactionEmit = emit; - this.outOfBandCompactionStarted = false; - this.outOfBandCompactionCompleted = false; + // OMP answers a manual compact with its RPC response alone, without + // compaction_start or compaction_end, so the request owns both rows. + this.manualCompactionRunning = true; + const emitCompaction = (status: "loading" | "completed") => + emit({ + type: "timeline", + provider: this.provider, + item: { type: "compaction", status, trigger: "manual" }, + }); + emitCompaction("loading"); try { await this.runtimeSession.compact(customInstructions); } catch (error) { const message = error instanceof Error ? error.message : String(error); - if ( - this.outOfBandCompactionEmit === emit && - this.outOfBandCompactionStarted && - !this.outOfBandCompactionCompleted - ) { - this.emitCompactionTimeline({ - turnId: undefined, - item: { - type: "compaction", - status: "completed", - trigger: "manual", - }, - }); - } + emitCompaction("completed"); emit({ type: "timeline", provider: this.provider, @@ -1439,13 +1431,12 @@ export class OmpAgentSession implements AgentSession { text: `[Error] Failed to compact context: ${message}`, }, }); + return; } finally { - if (this.outOfBandCompactionEmit === emit && !this.outOfBandCompactionStarted) { - this.outOfBandCompactionEmit = null; - this.outOfBandCompactionStarted = false; - this.outOfBandCompactionCompleted = false; - } + this.manualCompactionRunning = false; } + emitCompaction("completed"); + await this.usagePoller.refresh(); } private async executeAutoCompactCommand( @@ -1853,24 +1844,8 @@ export class OmpAgentSession implements AgentSession { return; } case "compaction_start": - this.emitCompactionTimeline({ - turnId, - item: { - type: "compaction", - status: "loading", - trigger: event.reason === "manual" ? "manual" : "auto", - }, - }); - return; case "compaction_end": - this.emitCompactionTimeline({ - turnId, - item: { - type: "compaction", - status: "completed", - trigger: event.reason === "manual" ? "manual" : "auto", - }, - }); + this.handleCompactionEvent(event, turnId); return; case "agent_end": { const messages = event.messages ?? []; @@ -1924,35 +1899,24 @@ export class OmpAgentSession implements AgentSession { } } - private emitCompactionTimeline(input: { - turnId: string | undefined; - item: Extract["item"]; - }): void { - const emitOutOfBand = this.outOfBandCompactionEmit; - if (emitOutOfBand && input.item.type === "compaction") { - if (input.item.status === "loading") { - this.outOfBandCompactionStarted = true; - } - if (input.item.status === "completed") { - this.outOfBandCompactionCompleted = true; - } - } - const event: AgentStreamEvent = { - type: "timeline", - provider: this.provider, - ...(emitOutOfBand ? {} : { turnId: input.turnId }), - item: input.item, - }; - if (emitOutOfBand) { - emitOutOfBand(event); - if (input.item.type === "compaction" && input.item.status === "completed") { - this.outOfBandCompactionEmit = null; - this.outOfBandCompactionStarted = false; - this.outOfBandCompactionCompleted = false; - } + private handleCompactionEvent( + event: Extract, + turnId: string | undefined, + ): void { + // A manual compact emits its own rows around the request. + if (this.manualCompactionRunning) { return; } - this.emit(event); + this.emit({ + type: "timeline", + provider: this.provider, + turnId, + item: { + type: "compaction", + status: event.type === "compaction_start" ? "loading" : "completed", + trigger: event.reason === "manual" ? "manual" : "auto", + }, + }); } private handleMessageUpdate( diff --git a/packages/server/src/server/agent/providers/omp/test-utils/fake-omp.ts b/packages/server/src/server/agent/providers/omp/test-utils/fake-omp.ts index 403a389647..58688c245d 100644 --- a/packages/server/src/server/agent/providers/omp/test-utils/fake-omp.ts +++ b/packages/server/src/server/agent/providers/omp/test-utils/fake-omp.ts @@ -168,7 +168,6 @@ export class FakeOmpSession implements OmpRuntimeSession { subagents: FakeOmpSubagentSnapshot[] = []; readonly subagentSubscriptionErrors = new Map(); compactError: Error | null = null; - emitCompactEnd = true; getStateError: Error | null = null; promptAck: OmpPromptAck = {}; branchResponse: { text?: string; cancelled?: boolean } = { text: "" }; @@ -266,11 +265,8 @@ export class FakeOmpSession implements OmpRuntimeSession { } async compact(customInstructions?: string): Promise { + // Like OMP, answer a manual compact with the response alone, no compaction events. this.compactRequests.push(customInstructions === undefined ? {} : { customInstructions }); - this.emit({ type: "compaction_start", reason: "manual" }); - if (this.emitCompactEnd) { - this.emit({ type: "compaction_end", reason: "manual" }); - } if (this.compactError) { throw this.compactError; } diff --git a/packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts b/packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts index 1820c187c9..8be53f1742 100644 --- a/packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts +++ b/packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts @@ -19,7 +19,12 @@ import { type OmpProviderIdleScheduler, } from "../agent.js"; import type { OmpUsagePollScheduler } from "../usage-poller.js"; -import type { OmpAgentMessage, OmpRpcSlashCommand, OmpRuntimeEvent } from "../rpc-types.js"; +import type { + OmpAgentMessage, + OmpRpcSlashCommand, + OmpRuntimeEvent, + OmpSessionStats, +} from "../rpc-types.js"; import { FakeOmp } from "./fake-omp.js"; const CWD = "/tmp/paseo-omp-agent-test"; @@ -169,6 +174,18 @@ export class OmpHarness { return events; } + async compact(statsAfter: OmpSessionStats, error?: Error): Promise { + const handler = this.requireSession().tryHandleOutOfBand("/compact"); + if (!handler) throw new Error("OMP session did not handle /compact out-of-band"); + const runtime = this.omp.latestSession(); + const sessionEventCount = this.events.length; + const emitted: AgentStreamEvent[] = []; + runtime.stats = statsAfter; + runtime.compactError = error ?? null; + await handler.run({ emit: (event) => emitted.push(event) }); + return [...emitted, ...this.events.slice(sessionEventCount)]; + } + async runPrompt( input: string, output: string, diff --git a/packages/server/src/server/agent/providers/omp/usage-poller.test.ts b/packages/server/src/server/agent/providers/omp/usage-poller.test.ts index 9750ec59ef..56b9dc667a 100644 --- a/packages/server/src/server/agent/providers/omp/usage-poller.test.ts +++ b/packages/server/src/server/agent/providers/omp/usage-poller.test.ts @@ -36,6 +36,27 @@ function deferred() { } describe("OMP usage poller", () => { + test("refresh publishes usage without ending an active turn's polling", async () => { + const scheduler = new ManualPollScheduler(); + const used: Array = []; + let stats: OmpSessionStats = { contextUsage: { contextWindow: 200_000, tokens: 9_000 } }; + const poller = new OmpUsagePoller({ + scheduler, + readStats: async () => stats, + onUsage: (update) => used.push(update.contextWindowUsedTokens), + onPollError: (error) => { + throw error; + }, + }); + + poller.startTurn(); + stats = { contextUsage: { contextWindow: 200_000, tokens: 2_000 } }; + await poller.refresh(); + + expect(used).toEqual([2_000]); + expect(scheduler.activePollCount()).toBe(1); + }); + test("emits only changed usage while active", async () => { const scheduler = new ManualPollScheduler(); const updates: AgentUsage[] = []; diff --git a/packages/server/src/server/agent/providers/omp/usage-poller.ts b/packages/server/src/server/agent/providers/omp/usage-poller.ts index 4e41708b5a..bc3bc1c968 100644 --- a/packages/server/src/server/agent/providers/omp/usage-poller.ts +++ b/packages/server/src/server/agent/providers/omp/usage-poller.ts @@ -105,6 +105,17 @@ export class OmpUsagePoller { this.publishUsage(usage, turnId); } + async refresh(): Promise { + let usage: AgentUsage | undefined; + try { + usage = toAgentUsage(await this.options.readStats()); + } catch (error) { + if (!this.closed) this.options.onPollError(error); + return; + } + if (!this.closed) this.publishUsage(usage); + } + close(): void { this.closed = true; this.active = false; diff --git a/packages/server/src/server/daemon-e2e/omp.real.e2e.test.ts b/packages/server/src/server/daemon-e2e/omp.real.e2e.test.ts index 7149b16864..ab407e31bc 100644 --- a/packages/server/src/server/daemon-e2e/omp.real.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/omp.real.e2e.test.ts @@ -131,10 +131,10 @@ async function timeline(client: DaemonClient, agentId: string): Promise entry.item); } -function compactionFiller(): string { +function compactionFiller(records = 1200): string { const teas = ["assam", "darjeeling", "sencha", "oolong"]; return Array.from( - { length: 1200 }, + { length: records }, (_, index) => `Record ${String(index).padStart(5, "0")}: tea=${teas[index % teas.length]} origin=region-${index % 37} notes=lorem ipsum dolor sit amet`, ).join("\n"); @@ -234,6 +234,27 @@ afterEach(() => { roots.clear(); }); +async function waitForTimelineItem( + client: DaemonClient, + agentId: string, + predicate: (item: AgentTimelineItem) => boolean, +): Promise { + const deadline = Date.now() + 60_000; + let items: AgentTimelineItem[] = []; + while (Date.now() < deadline) { + items = await timeline(client, agentId); + const item = items.find(predicate); + if (item) return item; + await new Promise((resolve) => setTimeout(resolve, 250)); + } + throw new Error(`Timed out waiting for OMP timeline item: ${JSON.stringify(items)}`); +} + +async function contextWindowUsedTokens(client: DaemonClient, agentId: string) { + const result = await client.fetchAgent(agentId); + return result?.agent.lastUsage?.contextWindowUsedTokens; +} + describe("daemon E2E (real OMP)", () => { test( "a prompt OMP rejects before its agent runs fails the turn with OMP's error", @@ -258,6 +279,48 @@ describe("daemon E2E (real OMP)", () => { TIMEOUT_MS, ); + test( + "manual compact shows its progress and refreshes context usage", + async () => { + const harness = await createHarness(); + try { + const agent = await createAgent(harness, "manual-compact"); + await promptAndFinish(harness, agent.id, "Reply exactly: ready"); + await promptAndFinish( + harness, + agent.id, + `${compactionFiller(2500)}\n\nWhich tea is in record 00007? Reply with one word.`, + ); + await promptAndFinish(harness, agent.id, "Reply exactly: after-filler"); + const usedBefore = await contextWindowUsedTokens(harness.client, agent.id); + expect(usedBefore).toBeGreaterThan(20_000); + + await harness.client.sendMessage(agent.id, "/compact"); + await waitForTimelineItem( + harness.client, + agent.id, + (item) => item.type === "compaction" && item.status === "completed", + ); + + expect(await timeline(harness.client, agent.id)).toEqual( + expect.arrayContaining([ + expect.objectContaining({ type: "compaction", status: "loading", trigger: "manual" }), + expect.objectContaining({ type: "compaction", status: "completed", trigger: "manual" }), + ]), + ); + await harness.client.waitForAgentUpsert( + agent.id, + (snapshot) => + (snapshot.lastUsage?.contextWindowUsedTokens ?? Infinity) < (usedBefore ?? 0), + 60_000, + ); + } finally { + await closeHarness(harness); + } + }, + TIMEOUT_MS, + ); + test( "prompt, native tool, and resumed follow-up", async () => {