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 e591b366c0..3f9fe77c43 100644 --- a/packages/server/src/server/agent/providers/omp/agent.test.ts +++ b/packages/server/src/server/agent/providers/omp/agent.test.ts @@ -7,6 +7,13 @@ import type { OmpUsagePollScheduler } from "./usage-poller.js"; import { resolveOmpProviderParams } from "./provider-config.js"; import { OmpHarness } from "./test-utils/omp-harness.js"; +// OMP's RPC events and Paseo's provider-state checks settle on microtasks. +async function settle(): Promise { + for (let index = 0; index < 6; index += 1) { + await waitForImmediate(); + } +} + test("OMP ready timeout defaults to 20 seconds and RPC timeout overrides both", () => { expect(resolveOmpProviderParams({}).runtimeProviderParams).toMatchObject({ readyTimeoutMs: 20_000, @@ -647,6 +654,131 @@ describe("OMP agent client and session", () => { ]); }); + test("a replacement prompt keeps its client correlation when OMP ends the aborted turn late", async () => { + const omp = new OmpHarness(); + await omp.start(); + + await omp.requireStartTurnFromClient("first prompt", "client-a"); + const runtime = omp.runtime(); + runtime.beginTurn(); + runtime.acceptPrompt("first prompt", "omp-user-a"); + runtime.streamAssistantText("working on it", "omp-assistant-a"); + await settle(); + + // OMP takes no steer, so Paseo interrupts and replaces the foreground turn. + await omp.interrupt(); + await omp.requireStartTurnFromClient("second prompt", "client-b"); + runtime.beginTurn(); + await settle(); + + // OMP flushes the aborted turn's terminal event after the replacement started. + runtime.emit({ + type: "agent_end", + messages: [{ role: "assistant", content: [], responseId: "omp-assistant-a" }], + }); + await settle(); + + runtime.acceptPrompt("second prompt", "omp-user-b"); + runtime.streamAssistantText("second answer", "omp-assistant-b"); + runtime.finishTurn(); + await settle(); + + expect(omp.timeline().filter((item) => item.type === "user_message")).toEqual([ + { + type: "user_message", + text: "first prompt", + messageId: "omp-user-a", + clientMessageId: "client-a", + }, + { + type: "user_message", + text: "second prompt", + messageId: "omp-user-b", + clientMessageId: "client-b", + }, + ]); + }); + + test("a prompt echoed after a hidden custom message completes the turn keeps its client correlation", async () => { + const omp = new OmpHarness(); + await omp.start(); + + await omp.requireStartTurnFromClient("do the thing", "client-a"); + const runtime = omp.runtime(); + runtime.beginTurn(); + runtime.emit({ + type: "message_end", + message: { role: "custom", content: "xd://inventory", display: false }, + }); + await settle(); + + runtime.acceptPrompt("do the thing", "omp-user-a"); + runtime.streamAssistantText("done", "omp-assistant-a"); + runtime.finishTurn(); + await settle(); + + expect(omp.timeline().filter((item) => item.type === "user_message")).toEqual([ + { + type: "user_message", + text: "do the thing", + messageId: "omp-user-a", + clientMessageId: "client-a", + }, + ]); + }); + + test("a re-emitted echo does not consume a later submission's client correlation", async () => { + const omp = new OmpHarness(); + await omp.start(); + + await omp.requireStartTurnFromClient("again", "client-a"); + const runtime = omp.runtime(); + runtime.beginTurn(); + runtime.acceptPrompt("again", "omp-user-a"); + runtime.streamAssistantText("first answer", "omp-assistant-a"); + runtime.finishTurn(); + await settle(); + + await omp.requireStartTurnFromClient("again", "client-b"); + runtime.beginTurn(); + // OMP re-sends the first turn's entry before echoing the new submission. + runtime.acceptPrompt("again", "omp-user-a"); + runtime.acceptPrompt("again", "omp-user-b"); + runtime.streamAssistantText("second answer", "omp-assistant-b"); + runtime.finishTurn(); + await settle(); + + expect(omp.timeline().filter((item) => item.type === "user_message")).toEqual([ + { type: "user_message", text: "again", messageId: "omp-user-a", clientMessageId: "client-a" }, + { type: "user_message", text: "again", messageId: "omp-user-b", clientMessageId: "client-b" }, + ]); + }); + + test("the same prompt text on two turns correlates to its own submission", async () => { + const omp = new OmpHarness(); + await omp.start(); + + await omp.requireStartTurnFromClient("again", "client-a"); + const runtime = omp.runtime(); + runtime.beginTurn(); + runtime.acceptPrompt("again", "omp-user-a"); + runtime.streamAssistantText("first answer", "omp-assistant-a"); + runtime.finishTurn(); + await settle(); + + await omp.requireStartTurnFromClient("again", "client-b"); + runtime.beginTurn(); + runtime.acceptPrompt("again", "omp-user-b"); + runtime.streamAssistantText("second answer", "omp-assistant-b"); + runtime.finishTurn(); + await settle(); + + expect(omp.timeline().filter((item) => item.type === "user_message")).toEqual([ + { type: "user_message", text: "again", messageId: "omp-user-a", clientMessageId: "client-a" }, + { type: "user_message", text: "again", messageId: "omp-user-b", clientMessageId: "client-b" }, + ]); + }); + test("re-emitted user message_end frames dedupe by native entry id", async () => { const omp = new OmpHarness(); await omp.start(); diff --git a/packages/server/src/server/agent/providers/omp/agent.ts b/packages/server/src/server/agent/providers/omp/agent.ts index 5f85b7c7d0..380a155be7 100644 --- a/packages/server/src/server/agent/providers/omp/agent.ts +++ b/packages/server/src/server/agent/providers/omp/agent.ts @@ -322,6 +322,9 @@ function renderTextOnlyImageHint(image: { data: string; mimeType: string }): str } } +// Bounds the dispatched prompts whose client correlation is still unclaimed. +const MAX_PENDING_CLIENT_CORRELATIONS = 16; + function convertPromptInput( prompt: AgentPromptInput, options: { model: OmpModel | null | undefined }, @@ -850,7 +853,12 @@ export class OmpAgentSession implements AgentSession { private activeAskUserDialog: ActiveAskUserDialog | null = null; private pendingCombinedAskUserResponse: PendingCombinedAskUserResponse | null = null; private activeTurnId: string | null = null; - private activeClientMessageId: string | null = null; + // A dispatched prompt owns its client correlation, not the turn that carried + // it: OMP echoes the prompt as a user message that can land after Paseo has + // already completed the turn (a replacement turn, or a turn completed by a + // custom message). Correlations are consumed by the echo they belong to, so + // the same text submitted on different turns stays separate. + private readonly pendingClientCorrelations: Array<{ text: string; clientMessageId: string }> = []; private activeAssistantMessageId: string | null = null; private activeTurnTerminalAssistantMessage: OmpAgentMessage | null = null; private activeTurnStarted = false; @@ -957,9 +965,11 @@ export class OmpAgentSession implements AgentSession { const payload = convertPromptInput(prompt, { model: this.state.model }); const turnId = randomUUID(); + if (options?.clientMessageId) { + this.rememberClientCorrelation(payload.text, options.clientMessageId); + } this.live = true; this.activeTurnId = turnId; - this.activeClientMessageId = options?.clientMessageId ?? null; this.activeAssistantMessageId = null; this.activeTurnTerminalAssistantMessage = null; this.activeTurnStarted = false; @@ -994,7 +1004,6 @@ export class OmpAgentSession implements AgentSession { } this.usagePoller.stopTurn(); this.activeTurnId = null; - this.activeClientMessageId = null; this.activeTurnStarted = false; this.activeTurnHasUserMessage = false; this.activeAssistantMessageId = null; @@ -1127,7 +1136,6 @@ export class OmpAgentSession implements AgentSession { this.terminalizeActiveWork(); this.usagePoller.stopTurn(); this.activeTurnId = null; - this.activeClientMessageId = null; this.activeTurnStarted = false; this.activeTurnHasUserMessage = false; this.activeAssistantMessageId = null; @@ -1288,6 +1296,24 @@ export class OmpAgentSession implements AgentSession { return this.activeTurnId ?? undefined; } + private rememberClientCorrelation(text: string, clientMessageId: string): void { + this.pendingClientCorrelations.push({ text, clientMessageId }); + // OMP never echoes some prompts (a local-only command, a session that dies + // mid-turn). Keep the queue bounded so those cannot accumulate. + if (this.pendingClientCorrelations.length > MAX_PENDING_CLIENT_CORRELATIONS) { + this.pendingClientCorrelations.shift(); + } + } + + private takeClientCorrelation(text: string): string | null { + const index = this.pendingClientCorrelations.findIndex((entry) => entry.text === text); + if (index === -1) { + return null; + } + const [entry] = this.pendingClientCorrelations.splice(index, 1); + return entry?.clientMessageId ?? null; + } + private scheduleNoTurnPromptCompletion(turnId: string): void { this.cancelNoTurnPromptCompletion(); const abort = new AbortController(); @@ -1341,6 +1367,7 @@ export class OmpAgentSession implements AgentSession { const outputs = this.pendingNoTurnOutputs.filter((output) => output.turnId === turnId); this.clearNoTurnBuffers(); if (promptText) { + const clientMessageId = this.takeClientCorrelation(promptText); this.emit({ type: "timeline", provider: this.provider, @@ -1348,7 +1375,7 @@ export class OmpAgentSession implements AgentSession { item: { type: "user_message", text: promptText, - ...(this.activeClientMessageId ? { clientMessageId: this.activeClientMessageId } : {}), + ...(clientMessageId ? { clientMessageId } : {}), }, }); } @@ -1786,7 +1813,6 @@ export class OmpAgentSession implements AgentSession { } const turnId = this.activeTurnId; this.activeTurnId = null; - this.activeClientMessageId = null; this.activeTurnStarted = false; this.activeTurnHasUserMessage = false; this.activeTurnTerminalAssistantMessage = null; @@ -2043,7 +2069,6 @@ export class OmpAgentSession implements AgentSession { } const nativeMessage = event.message as OmpAgentMessage & { id?: unknown; entryId?: unknown }; const messageId = readNativeMessageId(nativeMessage); - const clientMessageId = this.activeClientMessageId; const emitUserMessage = (resolvedMessageId?: string): void => { if (resolvedMessageId) { // OMP re-emits user message_end frames for entries it has already @@ -2054,6 +2079,9 @@ export class OmpAgentSession implements AgentSession { } this.emittedUserMessageIds.add(resolvedMessageId); } + // Claimed only once a row is really emitted, so a re-emitted frame cannot + // consume the correlation of a later submission with the same text. + const clientMessageId = this.takeClientCorrelation(text); this.emit({ type: "timeline", provider: this.provider, @@ -2127,7 +2155,6 @@ export class OmpAgentSession implements AgentSession { private completeTurn(turnId: string | undefined, messages: OmpAgentMessage[]): void { this.activeTurnId = null; - this.activeClientMessageId = null; this.activeAssistantMessageId = null; this.activeTurnTerminalAssistantMessage = null; this.activeTurnStarted = false; 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 30d75bb8f4..0700571467 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 @@ -466,6 +466,12 @@ export class OmpHarness { await promptStarted; } + async requireStartTurnFromClient(message: string, clientMessageId: string): Promise { + const promptStarted = this.omp.latestSession().nextPrompt(); + await this.requireSession().startTurn(message, { clientMessageId }); + await promptStarted; + } + async interrupt(): Promise { await this.requireSession().interrupt(); }