fix(omp): complete turns when agent_end omits messages (#2261)

* fix(omp): complete turns after empty agent_end

* test(omp): remove timing flush from empty end regression

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

@@ -132,6 +132,20 @@ describe("OMP agent client and session", () => {
     ]);
   });
 
+  test("completes a streamed assistant turn when agent_end omits messages", async () => {
+    const omp = new OmpHarness();
+    await omp.start();
+
+    const { completion } = await omp.startPromptWithEmptyAgentEnd(
+      "hello OMP",
+      "empty terminal payload recovered",
+    );
+    await expect(completion).resolves.toMatchObject({
+      finalText: "empty terminal payload recovered",
+    });
+    expect(omp.completedTurnCount()).toBe(1);
+  });
+
   test("does not accept a follow-up until OMP reports stable idle", async () => {
     const omp = new OmpHarness();
     await omp.start();

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

@@ -899,6 +899,7 @@ export class OmpAgentSession implements AgentSession {
   private activeTurnId: string | null = null;
   private activeClientMessageId: string | null = null;
   private activeAssistantMessageId: string | null = null;
+  private activeTurnTerminalAssistantMessage: OmpAgentMessage | null = null;
   private activeTurnStarted = false;
   private activeTurnHasUserMessage = false;
   private activeNoTurnPromptText: string | null = null;
@@ -987,6 +988,7 @@ export class OmpAgentSession implements AgentSession {
     this.activeTurnId = turnId;
     this.activeClientMessageId = options?.clientMessageId ?? null;
     this.activeAssistantMessageId = null;
+    this.activeTurnTerminalAssistantMessage = null;
     this.activeTurnStarted = false;
     this.activeTurnHasUserMessage = false;
     this.activePromptRequestId = null;
@@ -1017,6 +1019,7 @@ export class OmpAgentSession implements AgentSession {
         this.activeTurnStarted = false;
         this.activeTurnHasUserMessage = false;
         this.activeAssistantMessageId = null;
+        this.activeTurnTerminalAssistantMessage = null;
         this.clearNoTurnBuffers();
         if (isOmpRequestAbortError(error)) {
           this.emit({
@@ -1151,6 +1154,7 @@ export class OmpAgentSession implements AgentSession {
       this.activeTurnStarted = false;
       this.activeTurnHasUserMessage = false;
       this.activeAssistantMessageId = null;
+      this.activeTurnTerminalAssistantMessage = null;
       this.clearNoTurnBuffers();
       this.emit({
         type: "turn_canceled",
@@ -1773,6 +1777,7 @@ export class OmpAgentSession implements AgentSession {
     this.activeClientMessageId = null;
     this.activeTurnStarted = false;
     this.activeTurnHasUserMessage = false;
+    this.activeTurnTerminalAssistantMessage = null;
     this.clearNoTurnBuffers();
     this.emit({
       type: "turn_failed",
@@ -1857,17 +1862,25 @@ export class OmpAgentSession implements AgentSession {
           },
         });
         return;
-      case "agent_end":
+      case "agent_end": {
+        const messages = event.messages ?? [];
+        let terminalMessages: OmpAgentMessage[] | null = null;
+        if (messages.some((message) => message.role === "assistant")) {
+          terminalMessages = messages;
+        } else if (this.activeTurnTerminalAssistantMessage) {
+          terminalMessages = [this.activeTurnTerminalAssistantMessage];
+        }
         // OMP can end an internal extension-notice cycle before it starts the
-        // model turn for the same prompt. That cycle has no assistant message
-        // and is not the foreground turn's terminal event.
-        if (!(event.messages ?? []).some((message) => message.role === "assistant")) {
+        // model turn for the same prompt. Ignore only cycles where neither the
+        // terminal payload nor the live stream contained an assistant message.
+        if (!terminalMessages) {
           return;
         }
         // A state request is processed after OMP's RPC loop becomes promptable,
         // so do not advertise Paseo idle until it reports that transition.
-        void this.completeTurnAfterProviderIdle(turnId, event.messages ?? []);
+        void this.completeTurnAfterProviderIdle(turnId, terminalMessages);
         return;
+      }
       default:
         return;
     }
@@ -1983,6 +1996,9 @@ export class OmpAgentSession implements AgentSession {
   ): void {
     if (event.message.role === "assistant") {
       this.activeAssistantMessageId = null;
+      if (turnId) {
+        this.activeTurnTerminalAssistantMessage = event.message;
+      }
       return;
     }
     if (event.message.role === "custom") {
@@ -2097,6 +2113,7 @@ export class OmpAgentSession implements AgentSession {
     this.activeTurnId = null;
     this.activeClientMessageId = null;
     this.activeAssistantMessageId = null;
+    this.activeTurnTerminalAssistantMessage = null;
     this.activeTurnStarted = false;
     this.activeTurnHasUserMessage = false;
     this.clearNoTurnBuffers();

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

@@ -413,6 +413,12 @@ export class FakeOmpSession implements OmpRuntimeSession {
     this.emit({ type: "agent_end", messages: this.messages });
   }
 
+  finishTurnWithEmptyAgentEnd(message: OmpAgentMessage = { role: "assistant", content: [] }): void {
+    this.messages = [...this.messages, message];
+    this.emit({ type: "message_end", message });
+    this.emit({ type: "agent_end", messages: [] });
+  }
+
   beginTurn(): void {
     this.emit({ type: "turn_start" });
   }

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

@@ -173,6 +173,22 @@ export class OmpHarness {
     return await run;
   }
 
+  async startPromptWithEmptyAgentEnd(
+    input: string,
+    output: string,
+  ): Promise<{ completion: Promise<unknown> }> {
+    const session = this.requireSession();
+    const promptStarted = this.omp.latestSession().nextPrompt();
+    const completion = session.run(input);
+    await promptStarted;
+    const runtime = this.omp.latestSession();
+    runtime.beginTurn();
+    runtime.acceptPrompt(input, "user-1");
+    runtime.streamAssistantText(output);
+    runtime.finishTurnWithEmptyAgentEnd();
+    return { completion };
+  }
+
   async runPromptAfterExtensionNotice(input: string, output: string): Promise<unknown> {
     const session = this.requireSession();
     const promptStarted = this.omp.latestSession().nextPrompt();