feat: improve audio streaming and add wait-for-agent cancellation support

packages/server/src/server/agent/mcp-server.ts

@@ -24,6 +24,10 @@ import { curateAgentActivity } from "./activity-curator.js";
 import { AGENT_PROVIDER_DEFINITIONS } from "./provider-manifest.js";
 import { AgentRegistry } from "./agent-registry.js";
 import { createWorktree } from "../../utils/worktree.js";
+import {
+  WaitForAgentTracker,
+  type WaitForAgentCanceler,
+} from "./wait-for-agent-tracker.js";
 
 export interface AgentMcpServerOptions {
   agentManager: AgentManager;
@@ -103,6 +107,7 @@ async function serializeSnapshotWithMetadata(
 async function waitForAgentEvent(
   agentManager: AgentManager,
   agentId: string,
+  tracker: WaitForAgentTracker,
   signal?: AbortSignal
 ): Promise<{
   status: AgentLifecycleStatus;
@@ -142,6 +147,19 @@ async function waitForAgentEvent(
       resolve({ status: currentStatus, permission });
     };
 
+    const unregister = tracker.register(agentId, (reason) => {
+      cleanup();
+      reject(
+        Object.assign(
+          new Error(reason ?? "wait_for_agent cancelled"),
+          {
+            name: "AbortError",
+          }
+        )
+      );
+    });
+    cleanupFns.push(unregister);
+
     const unsubscribe = agentManager.subscribe(
       (event) => {
         if (event.type === "agent_state") {
@@ -329,7 +347,12 @@ export async function createAgentMcpServer(
       },
     },
     async ({ agentId }, { signal }) => {
-      const result = await waitForAgentEvent(agentManager, agentId, signal);
+      const result = await waitForAgentEvent(
+        agentManager,
+        agentId,
+        waitTracker,
+        signal
+      );
       const activity = buildActivityPayload(agentManager, agentId);
 
       return {

packages/server/src/server/agent/wait-for-agent-tracker.test.ts

@@ -0,0 +1,40 @@
+import { describe, expect, it, vi } from "vitest";
+import { WaitForAgentTracker } from "./wait-for-agent-tracker";
+
+describe("WaitForAgentTracker", () => {
+  it("registers and cancels waiters per agent", () => {
+    const tracker = new WaitForAgentTracker();
+    const cancelA = vi.fn();
+    const cancelB = vi.fn();
+
+    const unregisterA = tracker.register("agent-a", cancelA);
+    tracker.register("agent-a", cancelB);
+
+    expect(tracker.cancel("agent-a", "test")).toBe(true);
+    expect(cancelA).toHaveBeenCalledWith("test");
+    expect(cancelB).toHaveBeenCalledWith("test");
+
+    // Subsequent cancels should no-op once cleared
+    expect(tracker.cancel("agent-a")).toBe(false);
+
+    // Unregister removes callback without triggering it
+    const cancelC = vi.fn();
+    const unregisterC = tracker.register("agent-b", cancelC);
+    unregisterC();
+    expect(tracker.cancel("agent-b")).toBe(false);
+    expect(cancelC).not.toHaveBeenCalled();
+  });
+
+  it("supports cancelling all waiters", () => {
+    const tracker = new WaitForAgentTracker();
+    const cancelA = vi.fn();
+    const cancelB = vi.fn();
+
+    tracker.register("agent-a", cancelA);
+    tracker.register("agent-b", cancelB);
+
+    expect(tracker.cancelAll()).toBe(2);
+    expect(cancelA).toHaveBeenCalled();
+    expect(cancelB).toHaveBeenCalled();
+  });
+});

packages/server/src/server/agent/wait-for-agent-tracker.ts

@@ -0,0 +1,58 @@
+export type WaitForAgentCanceler = (agentId: string, reason?: string) => boolean;
+
+/**
+ * Tracks long-running wait_for_agent tool calls so they can be cancelled
+ * explicitly (e.g., when a realtime barge-in aborts the current turn).
+ */
+export class WaitForAgentTracker {
+  private waiters = new Map<string, Set<(reason?: string) => void>>();
+
+  register(agentId: string, cancel: (reason?: string) => void): () => void {
+    if (!this.waiters.has(agentId)) {
+      this.waiters.set(agentId, new Set());
+    }
+    const waitersForAgent = this.waiters.get(agentId)!;
+    waitersForAgent.add(cancel);
+
+    return () => {
+      const current = this.waiters.get(agentId);
+      if (!current) {
+        return;
+      }
+      current.delete(cancel);
+      if (current.size === 0) {
+        this.waiters.delete(agentId);
+      }
+    };
+  }
+
+  cancel(agentId: string, reason?: string): boolean {
+    const waitersForAgent = this.waiters.get(agentId);
+    if (!waitersForAgent || waitersForAgent.size === 0) {
+      return false;
+    }
+
+    this.waiters.delete(agentId);
+    for (const cancel of waitersForAgent) {
+      try {
+        cancel(reason);
+      } catch (error) {
+        console.warn(
+          `[WaitForAgentTracker] Cancel callback failed for ${agentId}:`,
+          error
+        );
+      }
+    }
+    return true;
+  }
+
+  cancelAll(reason?: string): number {
+    let cancelled = 0;
+    for (const agentId of Array.from(this.waiters.keys())) {
+      if (this.cancel(agentId, reason)) {
+        cancelled += 1;
+      }
+    }
+    return cancelled;
+  }
+}

packages/server/src/server/session.ts

@@ -1587,7 +1587,23 @@ export class Session {
       `[Session ${this.clientId}] Buffered audio chunk (${chunkBuffer.length} bytes, chunks: ${this.audioBuffer.chunks.length}${this.audioBuffer.isPCM ? `, PCM bytes: ${this.audioBuffer.totalPCMBytes}` : ""})`
     );
 
+    // In realtime mode, only process audio when the user has finished speaking (isLast = true)
+    // This prevents partial transcriptions from being sent to the LLM
+    if (this.isRealtimeMode) {
+      if (!msg.isLast) {
+        console.log(
+          `[Session ${this.clientId}] Realtime mode: buffering audio, waiting for speech end`
+        );
+        return;
+      }
+      console.log(
+        `[Session ${this.clientId}] Realtime mode: speech ended, processing complete audio`
+      );
+    }
+
+    // In non-realtime mode, use streaming threshold to process chunks
     const reachedStreamingThreshold =
+      !this.isRealtimeMode &&
       this.audioBuffer.isPCM &&
       this.audioBuffer.totalPCMBytes >= MIN_STREAMING_SEGMENT_BYTES;