fix(google): bound native realtime tool ownership

2026-08-12 21:53:00 -06:00 · 2026-08-02 11:02:04 +08:00
parent 86a95b459c
commit 5be95ad4ba
2 changed files with 426 additions and 14 deletions
@@ -859,6 +859,172 @@ describe("buildGoogleRealtimeVoiceProvider", () => {
    ]);
  });

+  it("preserves tool ownership while reusing a resumption handle", async () => {
+    vi.useFakeTimers();
+    const provider = buildGoogleRealtimeVoiceProvider();
+    const onToolCall = vi.fn();
+    const bridge = provider.createBridge({
+      providerConfig: { apiKey: "gemini-key" },
+      onAudio: vi.fn(),
+      onClearAudio: vi.fn(),
+      onToolCall,
+    });
+
+    await bridge.connect();
+    const firstSession = lastConnectParams().callbacks;
+    firstSession.onmessage({
+      sessionResumptionUpdate: { resumable: true, newHandle: "resume-1" },
+      toolCall: {
+        functionCalls: [{ id: "call-1", name: "lookup", args: { query: "before" } }],
+      },
+    });
+    firstSession.onclose({ code: 1011, reason: "temporary" });
+    void bridge.submitToolResult("call-1", { result: "ok" });
+    expect(session.sendToolResponse).not.toHaveBeenCalled();
+    await vi.advanceTimersByTimeAsync(250);
+
+    const resumedSession = lastConnectParams().callbacks;
+    resumedSession.onmessage({
+      toolCall: {
+        functionCalls: [{ id: "call-1", name: "different", args: { query: "replay" } }],
+      },
+    });
+    resumedSession.onopen();
+    resumedSession.onmessage({ setupComplete: {} });
+
+    expect(lastConnectParams().config.sessionResumption).toEqual({ handle: "resume-1" });
+    expect(onToolCall).toHaveBeenCalledOnce();
+    expect(session.sendToolResponse).toHaveBeenCalledWith({
+      functionResponses: [
+        {
+          id: "call-1",
+          name: "lookup",
+          response: { result: "ok" },
+        },
+      ],
+    });
+  });
+
+  it("fails closed when resumable tool responses exceed the reconnect buffer", async () => {
+    vi.useFakeTimers();
+    const provider = buildGoogleRealtimeVoiceProvider();
+    const onError = vi.fn();
+    const onClose = vi.fn();
+    const bridge = provider.createBridge({
+      providerConfig: { apiKey: "gemini-key" },
+      onAudio: vi.fn(),
+      onClearAudio: vi.fn(),
+      onToolCall: vi.fn(),
+      onError,
+      onClose,
+    });
+
+    await bridge.connect();
+    const firstSession = lastConnectParams().callbacks;
+    firstSession.onmessage({
+      sessionResumptionUpdate: { resumable: true, newHandle: "resume-1" },
+      toolCall: {
+        functionCalls: [{ id: "call-1", name: "lookup", args: {} }],
+      },
+    });
+    firstSession.onclose({ code: 1011, reason: "temporary" });
+    onError.mockClear();
+
+    void bridge.submitToolResult("call-1", { result: "x".repeat(1024 * 1024) });
+
+    expect(requireFirstError(onError).message).toBe(
+      "Google Live reconnect tool-response buffer limit exceeded",
+    );
+    expect(onClose).toHaveBeenCalledWith("error");
+    expect(session.sendToolResponse).not.toHaveBeenCalled();
+  });
+
+  it("drops queued reconnect responses when the resumed session cancels their call", async () => {
+    vi.useFakeTimers();
+    const provider = buildGoogleRealtimeVoiceProvider();
+    const onEvent = vi.fn();
+    const bridge = provider.createBridge({
+      providerConfig: { apiKey: "gemini-key" },
+      onAudio: vi.fn(),
+      onClearAudio: vi.fn(),
+      onToolCall: vi.fn(),
+      onEvent,
+    });
+
+    await bridge.connect();
+    const firstSession = lastConnectParams().callbacks;
+    firstSession.onmessage({
+      sessionResumptionUpdate: { resumable: true, newHandle: "resume-1" },
+      toolCall: {
+        functionCalls: [{ id: "call-1", name: "lookup", args: {} }],
+      },
+    });
+    firstSession.onclose({ code: 1011, reason: "temporary" });
+    await vi.advanceTimersByTimeAsync(250);
+
+    const resumedSession = lastConnectParams().callbacks;
+    void bridge.submitToolResult("call-1", { result: "stale" });
+    expect(session.sendToolResponse).not.toHaveBeenCalled();
+    resumedSession.onopen();
+    resumedSession.onmessage({
+      setupComplete: {},
+      toolCallCancellation: { ids: ["call-1"] },
+    });
+
+    expect(session.sendToolResponse).not.toHaveBeenCalled();
+    expect(onEvent).toHaveBeenCalledWith({
+      direction: "server",
+      type: "tool.call.cancelled",
+      itemId: "call-1",
+    });
+  });
+
+  it("resets tool ownership before a fresh automatic reconnect", async () => {
+    vi.useFakeTimers();
+    const provider = buildGoogleRealtimeVoiceProvider();
+    const onToolCall = vi.fn();
+    const onEvent = vi.fn();
+    const bridge = provider.createBridge({
+      providerConfig: { apiKey: "gemini-key" },
+      onAudio: vi.fn(),
+      onClearAudio: vi.fn(),
+      onToolCall,
+      onEvent,
+    });
+
+    await bridge.connect();
+    const firstSession = lastConnectParams().callbacks;
+    firstSession.onmessage({
+      toolCall: {
+        functionCalls: [{ id: "call-1", name: "old_lookup", args: {} }],
+      },
+    });
+    firstSession.onclose({ code: 1011, reason: "temporary" });
+    await vi.advanceTimersByTimeAsync(250);
+
+    lastConnectParams().callbacks.onmessage({
+      toolCall: {
+        functionCalls: [{ id: "call-1", name: "new_lookup", args: {} }],
+      },
+    });
+    void bridge.submitToolResult("call-1", { result: "ok" });
+
+    expect(onEvent).toHaveBeenCalledWith({
+      direction: "client",
+      type: "session.continuity.reset",
+    });
+    expect(onToolCall).toHaveBeenCalledTimes(2);
+    expect(session.sendToolResponse).toHaveBeenCalledWith({
+      functionResponses: [
+        {
+          id: "call-1",
+          name: "new_lookup",
+          response: { result: "ok" },
+        },
+      ],
+    });
+  });
+
  it("preserves continuity when resumability recovers before reconnect", async () => {
    vi.useFakeTimers();
    const provider = buildGoogleRealtimeVoiceProvider();
@@ -1952,6 +2118,114 @@ describe("buildGoogleRealtimeVoiceProvider", () => {
    });
  });

+  it("deduplicates replayed Google Live tool calls by call id", async () => {
+    const provider = buildGoogleRealtimeVoiceProvider();
+    const onToolCall = vi.fn();
+    const bridge = provider.createBridge({
+      providerConfig: { apiKey: "gemini-key" },
+      onAudio: vi.fn(),
+      onClearAudio: vi.fn(),
+      onToolCall,
+    });
+
+    await bridge.connect();
+    const callbacks = lastConnectParams().callbacks;
+    callbacks.onmessage({
+      toolCall: {
+        functionCalls: [{ id: "call-1", name: "lookup", args: { query: "first" } }],
+      },
+    });
+    callbacks.onmessage({
+      toolCall: {
+        functionCalls: [{ id: "call-1", name: "different", args: { query: "replay" } }],
+      },
+    });
+
+    expect(onToolCall).toHaveBeenCalledOnce();
+    void bridge.submitToolResult("call-1", { result: "ok" });
+    expect(session.sendToolResponse).toHaveBeenCalledWith({
+      functionResponses: [
+        {
+          id: "call-1",
+          name: "lookup",
+          response: { result: "ok" },
+        },
+      ],
+    });
+    callbacks.onmessage({
+      toolCall: {
+        functionCalls: [{ id: "call-1", name: "lookup", args: { query: "late replay" } }],
+      },
+    });
+    expect(onToolCall).toHaveBeenCalledOnce();
+  });
+
+  it("ignores late results after Google cancels a tool call", async () => {
+    const provider = buildGoogleRealtimeVoiceProvider();
+    const onEvent = vi.fn();
+    const onError = vi.fn();
+    const bridge = provider.createBridge({
+      providerConfig: { apiKey: "gemini-key" },
+      onAudio: vi.fn(),
+      onClearAudio: vi.fn(),
+      onToolCall: vi.fn(),
+      onEvent,
+      onError,
+    });
+
+    await bridge.connect();
+    const callbacks = lastConnectParams().callbacks;
+    callbacks.onmessage({
+      toolCall: {
+        functionCalls: [{ id: "call-1", name: "lookup", args: { query: "hi" } }],
+      },
+    });
+    callbacks.onmessage({
+      toolCallCancellation: { ids: ["call-1"] },
+    });
+
+    void bridge.submitToolResult("call-1", { result: "late" });
+
+    expect(session.sendToolResponse).not.toHaveBeenCalled();
+    expect(onError).not.toHaveBeenCalled();
+    expect(onEvent).toHaveBeenCalledWith({
+      direction: "server",
+      type: "tool.call.cancelled",
+      itemId: "call-1",
+    });
+  });
+
+  it("fails closed when Google exceeds the tool-call session limit", async () => {
+    const provider = buildGoogleRealtimeVoiceProvider();
+    const onToolCall = vi.fn();
+    const onError = vi.fn();
+    const onClose = vi.fn();
+    const bridge = provider.createBridge({
+      providerConfig: { apiKey: "gemini-key" },
+      onAudio: vi.fn(),
+      onClearAudio: vi.fn(),
+      onToolCall,
+      onError,
+      onClose,
+    });
+
+    await bridge.connect();
+    lastConnectParams().callbacks.onmessage({
+      toolCall: {
+        functionCalls: Array.from({ length: 1_025 }, (_, index) => ({
+          id: `call-${index}`,
+          name: "lookup",
+          args: {},
+        })),
+      },
+    });
+
+    expect(onToolCall).toHaveBeenCalledTimes(1_024);
+    expect(requireFirstError(onError).message).toBe("Google Live tool-call session limit exceeded");
+    expect(session.close).toHaveBeenCalledOnce();
+    expect(onClose).toHaveBeenCalledWith("error");
+  });
+
  it("keeps Google Live consult calls open after continuing tool responses", async () => {
    const provider = buildGoogleRealtimeVoiceProvider();
    const bridge = provider.createBridge({
@@ -68,6 +68,9 @@ const GOOGLE_REALTIME_BROWSER_NEW_SESSION_TTL_MS = 60 * 1000;
 const GOOGLE_REALTIME_RECONNECT_MAX_ATTEMPTS = 3;
 const GOOGLE_REALTIME_RECONNECT_BASE_DELAY_MS = 250;
 const GOOGLE_REALTIME_RECONNECT_MAX_DELAY_MS = 2_000;
+const GOOGLE_REALTIME_MAX_TOOL_CALL_IDS = 1_024;
+const GOOGLE_REALTIME_MAX_PENDING_TOOL_RESPONSES = 1_024;
+const GOOGLE_REALTIME_MAX_PENDING_TOOL_RESPONSE_BYTES = 1024 * 1024;
 const GOOGLE_REALTIME_MAX_PENDING_TRANSCRIPT_BYTES = 256 * 1024;
 const GOOGLE_REALTIME_TRANSCRIPT_OVERFLOW_MESSAGE =
  "Google Live transcript exceeded the 256 KiB UTF-8 pending buffer limit";
@@ -462,6 +465,12 @@ type GoogleLiveConnectionAttempt = {
  cancel: () => void;
 };

+type GooglePendingToolResponse = {
+  callId: string;
+  payload: string;
+  byteLength: number;
+};
+
 class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
  readonly supportsToolResultContinuation: boolean;
  readonly supportsToolResultSuppression = false;
@@ -477,9 +486,13 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
  private consecutiveSilenceMs = 0;
  private audioStreamEnded = false;
  private pendingFunctionNames = new Map<string, string>();
+  private seenFunctionCallIds = new Set<string>();
+  private pendingToolResponses: GooglePendingToolResponse[] = [];
+  private pendingToolResponseBytes = 0;
  private readonly audioFormat: RealtimeVoiceAudioFormat;
  private readonly model: string;
  private resumptionHandle: string | undefined;
+  private resumingSession = false;
  private reconnectAttempts = 0;
  private reconnectTimer: ReturnType<typeof setTimeout> | undefined;
  private hasConnectedSession = false;
@@ -539,7 +552,12 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
    this.sessionReadyFired = false;
    this.consecutiveSilenceMs = 0;
    this.audioStreamEnded = false;
-    this.pendingFunctionNames.clear();
+    const resumesExistingSession =
+      this.config.sessionResumption !== false && Boolean(this.resumptionHandle);
+    this.resumingSession = resumesExistingSession;
+    if (!resumesExistingSession) {
+      this.resetToolCallOwnership();
+    }
    const ai = createGoogleGenAI({
      apiKey: this.config.apiKey,
      httpOptions: {
@@ -596,7 +614,6 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
            this.connected = false;
            this.setupCompleteReceived = false;
            this.sessionConfigured = false;
-            this.pendingFunctionNames.clear();
            this.session = null;
            if (this.terminalError) {
              this.notifyClose("error");
@@ -610,6 +627,7 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
            if (this.scheduleReconnect(closeDetails)) {
              return;
            }
+            this.resetToolCallOwnership();
            // Transport failure is not an utterance boundary. Preserve transcript
            // fragments across reconnects and finalize only when recovery is exhausted.
            this.flushPendingTranscripts();
@@ -712,11 +730,11 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
    result: unknown,
    options?: RealtimeVoiceToolResultOptions,
  ): void {
-    if (!this.session) {
-      return;
-    }
    const name = this.pendingFunctionNames.get(callId);
    if (!name) {
+      if (this.seenFunctionCallIds.has(callId)) {
+        return;
+      }
      this.config.onError?.(
        new Error(
          `Google Live function response is missing a matching function call for ${callId}`,
@@ -753,16 +771,28 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
        );
        return;
      }
-      this.session.sendToolResponse({
-        functionResponses: [functionResponse],
-      });
+      const session = this.session;
+      const canSendImmediately = Boolean(
+        session && (!this.resumingSession || this.sessionConfigured),
+      );
+      if (session && canSendImmediately) {
+        session.sendToolResponse({
+          functionResponses: [functionResponse],
+        });
+      } else {
+        this.queueToolResponseForReconnect(callId, functionResponse);
+      }
      if (options?.willContinue !== true) {
        this.pendingFunctionNames.delete(callId);
      }
    } catch (error) {
-      this.config.onError?.(
-        error instanceof Error ? error : new Error("Failed to send Google Live function response"),
-      );
+      const sendError =
+        error instanceof Error ? error : new Error("Failed to send Google Live function response");
+      if (this.session && (!this.resumingSession || this.sessionConfigured)) {
+        this.config.onError?.(sendError);
+      } else {
+        this.failConnection(sendError);
+      }
    }
  }

@@ -783,7 +813,7 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
    this.clearPendingAudio();
    this.consecutiveSilenceMs = 0;
    this.audioStreamEnded = false;
-    this.pendingFunctionNames.clear();
+    this.resetToolCallOwnership();
    this.flushPendingTranscripts();
    const owner = this.connectionOwner;
    this.connectionOwner = undefined;
@@ -840,6 +870,14 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
    if (message.toolCall) {
      this.handleToolCall(message.toolCall);
    }
+    if (message.toolCallCancellation) {
+      this.handleToolCallCancellation(message.toolCallCancellation.ids);
+    }
+    if (message.setupComplete) {
+      // Apply cancellation and tool facts from the same server message before
+      // setup activation flushes responses retained across a resumable reconnect.
+      this.maybeActivateSession();
+    }
  }

  private captureSessionLifecycle(message: LiveServerMessage): void {
@@ -868,7 +906,6 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
      this.continuityResetEmitted = false;
    }
    this.setupCompleteReceived = true;
-    this.maybeActivateSession();
  }

  private maybeActivateSession(): void {
@@ -879,6 +916,10 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
    }
    this.sessionConfigured = true;
    this.reconnectAttempts = 0;
+    if (!this.flushPendingToolResponses()) {
+      return;
+    }
+    this.resumingSession = false;
    for (const chunk of this.pendingAudio.drain()) {
      this.sendAudio(chunk);
    }
@@ -995,7 +1036,7 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = undefined;
    }
-    this.pendingFunctionNames.clear();
+    this.resetToolCallOwnership();
    this.flushPendingTranscripts();
    const owner = this.connectionOwner;
    this.connectionOwner = undefined;
@@ -1043,6 +1084,16 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
        continue;
      }
      const callId = call.id?.trim() || `google-live-${randomUUID()}`;
+      if (this.seenFunctionCallIds.has(callId)) {
+        continue;
+      }
+      // The Live protocol defines no replay window, so dropping old IDs could execute
+      // a very late duplicate. End an extreme session instead of weakening dedupe.
+      if (this.seenFunctionCallIds.size >= GOOGLE_REALTIME_MAX_TOOL_CALL_IDS) {
+        this.failConnection(new Error("Google Live tool-call session limit exceeded"));
+        return;
+      }
+      this.seenFunctionCallIds.add(callId);
      this.pendingFunctionNames.set(callId, name);
      this.config.onToolCall?.({
        itemId: callId,
@@ -1053,6 +1104,91 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
    }
  }

+  private handleToolCallCancellation(ids: string[] | undefined): void {
+    for (const rawId of ids ?? []) {
+      const callId = rawId.trim();
+      if (!callId) {
+        continue;
+      }
+      const removedPendingCall = this.pendingFunctionNames.delete(callId);
+      const removedQueuedResponse = this.removePendingToolResponses(callId);
+      if (!removedPendingCall && !removedQueuedResponse) {
+        continue;
+      }
+      // Provider cancellation invalidates any late consumer result for this call ID.
+      this.config.onEvent?.({
+        direction: "server",
+        type: "tool.call.cancelled",
+        itemId: callId,
+      });
+    }
+  }
+
+  private resetToolCallOwnership(): void {
+    this.pendingFunctionNames.clear();
+    this.seenFunctionCallIds.clear();
+    this.pendingToolResponses = [];
+    this.pendingToolResponseBytes = 0;
+  }
+
+  private queueToolResponseForReconnect(callId: string, functionResponse: FunctionResponse): void {
+    const payload = JSON.stringify(functionResponse);
+    const payloadBytes = Buffer.byteLength(payload, "utf8");
+    if (
+      this.pendingToolResponses.length >= GOOGLE_REALTIME_MAX_PENDING_TOOL_RESPONSES ||
+      this.pendingToolResponseBytes + payloadBytes > GOOGLE_REALTIME_MAX_PENDING_TOOL_RESPONSE_BYTES
+    ) {
+      throw new Error("Google Live reconnect tool-response buffer limit exceeded");
+    }
+    // Store the serialized wire shape so a stalled reconnect cannot retain an
+    // arbitrarily large caller-owned object graph through the tool result.
+    this.pendingToolResponses.push({ callId, payload, byteLength: payloadBytes });
+    this.pendingToolResponseBytes += payloadBytes;
+  }
+
+  private removePendingToolResponses(callId: string): boolean {
+    const retained: GooglePendingToolResponse[] = [];
+    let removed = false;
+    for (const response of this.pendingToolResponses) {
+      if (response.callId === callId) {
+        this.pendingToolResponseBytes -= response.byteLength;
+        removed = true;
+      } else {
+        retained.push(response);
+      }
+    }
+    this.pendingToolResponses = retained;
+    return removed;
+  }
+
+  private flushPendingToolResponses(): boolean {
+    const session = this.session;
+    if (!session) {
+      return false;
+    }
+    try {
+      while (this.pendingToolResponses.length > 0) {
+        const response = this.pendingToolResponses[0];
+        if (!response) {
+          break;
+        }
+        session.sendToolResponse({
+          functionResponses: [JSON.parse(response.payload) as FunctionResponse],
+        });
+        this.pendingToolResponses.shift();
+        this.pendingToolResponseBytes -= response.byteLength;
+      }
+      return true;
+    } catch (error) {
+      this.failConnection(
+        error instanceof Error
+          ? error
+          : new Error("Failed to flush Google Live function responses"),
+      );
+      return false;
+    }
+  }
+
  private scheduleReconnect(closeDetails: string): boolean {
    if (this.reconnectAttempts >= GOOGLE_REALTIME_RECONNECT_MAX_ATTEMPTS) {
      return false;
@@ -1064,6 +1200,7 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
      // consumers before backoff so stale work cannot finish into the replacement.
      this.continuityResetEmitted = true;
      this.resetPendingTranscripts();
+      this.resetToolCallOwnership();
      this.config.onEvent?.({
        direction: "client",
        type: "session.continuity.reset",
@@ -1088,6 +1225,7 @@ class GoogleRealtimeVoiceBridge implements RealtimeVoiceBridge {
        const message = error instanceof Error ? error.message : String(error);
        this.config.onError?.(error instanceof Error ? error : new Error(message));
        if (!this.scheduleReconnect(`connect failed: ${message}`)) {
+          this.resetToolCallOwnership();
          this.flushPendingTranscripts();
          this.notifyClose("error");
        }