fix(omp): claim the newest pending correlation for a prompt echo

Paseo runs one foreground turn at a time, so two unclaimed correlations
with the same text mean the older prompt's turn ended without an echo.
The newest one belongs to the submission OMP is echoing now.

This replaces dropping a correlation on dispatch failure, which was
wrong for a request that timed out after OMP had already accepted the
prompt: OMP still echoes it, and the echo needs its id.

Tests wait on observable session state instead of a fixed number of
event-loop cycles.

packages/server/src/server/agent/providers/omp/agent.test.ts

@@ -7,13 +7,6 @@ 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<void> {
-  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,
@@ -663,25 +656,24 @@ describe("OMP agent client and session", () => {
     runtime.beginTurn();
     runtime.acceptPrompt("first prompt", "omp-user-a");
     runtime.streamAssistantText("working on it", "omp-assistant-a");
-    await settle();
+    await omp.waitForUserMessages(1);
 
     // 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();
+    await omp.waitForEvent("turn_completed");
 
     runtime.acceptPrompt("second prompt", "omp-user-b");
     runtime.streamAssistantText("second answer", "omp-assistant-b");
     runtime.finishTurn();
-    await settle();
+    await omp.waitForUserMessages(2);
 
     expect(omp.timeline().filter((item) => item.type === "user_message")).toEqual([
       {
@@ -710,12 +702,12 @@ describe("OMP agent client and session", () => {
       type: "message_end",
       message: { role: "custom", content: "xd://inventory", display: false },
     });
-    await settle();
+    await omp.waitForEvent("turn_completed");
 
     runtime.acceptPrompt("do the thing", "omp-user-a");
     runtime.streamAssistantText("done", "omp-assistant-a");
     runtime.finishTurn();
-    await settle();
+    await omp.waitForUserMessages(1);
 
     expect(omp.timeline().filter((item) => item.type === "user_message")).toEqual([
       {
@@ -737,7 +729,7 @@ describe("OMP agent client and session", () => {
     runtime.acceptPrompt("again", "omp-user-a");
     runtime.streamAssistantText("first answer", "omp-assistant-a");
     runtime.finishTurn();
-    await settle();
+    await omp.waitForUserMessages(1);
 
     await omp.requireStartTurnFromClient("again", "client-b");
     runtime.beginTurn();
@@ -746,7 +738,7 @@ describe("OMP agent client and session", () => {
     runtime.acceptPrompt("again", "omp-user-b");
     runtime.streamAssistantText("second answer", "omp-assistant-b");
     runtime.finishTurn();
-    await settle();
+    await omp.waitForUserMessages(2);
 
     expect(omp.timeline().filter((item) => item.type === "user_message")).toEqual([
       { type: "user_message", text: "again", messageId: "omp-user-a", clientMessageId: "client-a" },
@@ -754,6 +746,34 @@ describe("OMP agent client and session", () => {
     ]);
   });
 
+  test("a prompt OMP echoes after its dispatch failed keeps its client correlation", async () => {
+    const omp = new OmpHarness();
+    await omp.start();
+
+    const runtime = omp.runtime();
+    // The prompt reached OMP, but Paseo's request failed before it was answered.
+    runtime.promptError = new Error("OMP prompt request timed out");
+    await omp.requireStartTurnFromClient("slow prompt", "client-a");
+    await omp.waitForEvent("turn_failed");
+    runtime.promptError = null;
+
+    // OMP accepted it anyway and echoes it once it gets to the turn.
+    runtime.beginTurn();
+    runtime.acceptPrompt("slow prompt", "omp-user-a");
+    runtime.streamAssistantText("late answer", "omp-assistant-a");
+    runtime.finishTurn();
+    await omp.waitForUserMessages(1);
+
+    expect(omp.timeline().filter((item) => item.type === "user_message")).toEqual([
+      {
+        type: "user_message",
+        text: "slow prompt",
+        messageId: "omp-user-a",
+        clientMessageId: "client-a",
+      },
+    ]);
+  });
+
   test("a prompt that never reaches OMP leaves no correlation for a later submission", async () => {
     const omp = new OmpHarness();
     await omp.start();
@@ -761,8 +781,7 @@ describe("OMP agent client and session", () => {
     const runtime = omp.runtime();
     runtime.promptError = new Error("OMP prompt dispatch failed");
     await omp.requireStartTurnFromClient("retry me", "client-a");
-    await settle();
-    expect(omp.eventTypes()).toContain("turn_failed");
+    await omp.waitForEvent("turn_failed");
     runtime.promptError = null;
 
     await omp.requireStartTurnFromClient("retry me", "client-b");
@@ -770,7 +789,7 @@ describe("OMP agent client and session", () => {
     runtime.acceptPrompt("retry me", "omp-user-b");
     runtime.streamAssistantText("done", "omp-assistant-b");
     runtime.finishTurn();
-    await settle();
+    await omp.waitForUserMessages(1);
 
     expect(omp.timeline().filter((item) => item.type === "user_message")).toEqual([
       {
@@ -792,14 +811,14 @@ describe("OMP agent client and session", () => {
     runtime.acceptPrompt("again", "omp-user-a");
     runtime.streamAssistantText("first answer", "omp-assistant-a");
     runtime.finishTurn();
-    await settle();
+    await omp.waitForUserMessages(1);
 
     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();
+    await omp.waitForUserMessages(2);
 
     expect(omp.timeline().filter((item) => item.type === "user_message")).toEqual([
       { type: "user_message", text: "again", messageId: "omp-user-a", clientMessageId: "client-a" },

packages/server/src/server/agent/providers/omp/agent.ts

@@ -999,8 +999,6 @@ export class OmpAgentSession implements AgentSession {
           return;
         }
       } catch (error) {
-        // The prompt never reached OMP, so no echo will ever claim it.
-        this.forgetClientCorrelation(options?.clientMessageId);
         if (this.activeTurnId !== turnId) {
           return;
         }
@@ -1307,20 +1305,13 @@ export class OmpAgentSession implements AgentSession {
     }
   }
 
-  private forgetClientCorrelation(clientMessageId: string | undefined): void {
-    if (!clientMessageId) {
-      return;
-    }
-    const index = this.pendingClientCorrelations.findIndex(
-      (entry) => entry.clientMessageId === clientMessageId,
-    );
-    if (index !== -1) {
-      this.pendingClientCorrelations.splice(index, 1);
-    }
-  }
-
+  // Paseo runs one foreground turn at a time, so two unclaimed correlations with
+  // the same text mean the older prompt's turn ended without an echo (a failed
+  // or timed-out dispatch). The newest one belongs to the submission OMP is
+  // echoing now; the older one stays queued in case OMP echoes it late, and
+  // falls off the bounded queue otherwise.
   private takeClientCorrelation(text: string): string | null {
-    const index = this.pendingClientCorrelations.findIndex((entry) => entry.text === text);
+    const index = this.pendingClientCorrelations.findLastIndex((entry) => entry.text === text);
     if (index === -1) {
       return null;
     }

packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts

@@ -466,6 +466,27 @@ export class OmpHarness {
     await promptStarted;
   }
 
+  // Drains the microtask queue until the session has reached the state the test
+  // is waiting on, so no test depends on how many async steps a path takes.
+  async waitUntil(condition: () => boolean, description: string): Promise<void> {
+    for (let index = 0; index < 100; index += 1) {
+      if (condition()) return;
+      await waitForImmediate();
+    }
+    throw new Error(`OMP harness timed out waiting for ${description}`);
+  }
+
+  async waitForUserMessages(count: number): Promise<void> {
+    await this.waitUntil(
+      () => this.timeline().filter((item) => item.type === "user_message").length >= count,
+      `${count} user message(s)`,
+    );
+  }
+
+  async waitForEvent(type: AgentStreamEvent["type"]): Promise<void> {
+    await this.waitUntil(() => this.eventTypes().includes(type), `a ${type} event`);
+  }
+
   async requireStartTurnFromClient(message: string, clientMessageId: string): Promise<void> {
     const promptStarted = this.omp.latestSession().nextPrompt();
     await this.requireSession().startTurn(message, { clientMessageId });