Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions src/adapters/google-truncation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,3 +11,14 @@ export function vertexTruncationErrorMessage(reason?: string): string {
const suffix = reason ? ` (${redactSecretString(reason).slice(0, 160)})` : "";
return `Vertex AI response truncated upstream before the turn completed${suffix}`;
}

/**
* Whether a finished turn must fail closed. A truncation reason arriving mid tool call always
* does. MALFORMED_FUNCTION_CALL fails closed even with zero started calls: the malformed call
* is dropped upstream and usually never materializes as a part, so the turn is incomplete
* despite looking empty. MAX_TOKENS with no started call stays a plain token-limit stop.
*/
export function isVertexTruncatedTurn(finishReason: string | undefined, toolCallsStarted: number): boolean {
if (!isVertexTruncationReason(finishReason)) return false;
return toolCallsStarted > 0 || finishReason === "MALFORMED_FUNCTION_CALL";
}
6 changes: 3 additions & 3 deletions src/adapters/google.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ import { contentPartsToText, parseDataUrl } from "./image";
import { getVertexAccessToken } from "../lib/gcp-adc";
import { fetchAntigravityWithRetry, fetchVertexWithRetry } from "./google-http";
import { safeAntigravityHttpErrorMessage, safeVertexHttpErrorMessage } from "./google-errors";
import { isVertexTruncationReason, vertexTruncationErrorMessage } from "./google-truncation";
import { isVertexTruncatedTurn, vertexTruncationErrorMessage } from "./google-truncation";
import { ANTIGRAVITY_REQUEST_UA, antigravitySessionId, isLikelyRealThoughtSignature, sanitizeAntigravityClaudeSignatures } from "./google-antigravity-wire";
import { compileGoogleWireBody } from "./google-wire-compiler";
import { identifyRoutedModel } from "./identity";
Expand Down Expand Up @@ -596,7 +596,7 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte
// Fail-closed: a turn cut off mid tool call (MAX_TOKENS / MALFORMED_FUNCTION_CALL) surfaces
// an error instead of a silently-incomplete done. Mirrors kiro-truncation.
if ((provider.googleMode === "vertex" || provider.googleMode === "cloud-code-assist")
&& toolCallsStarted > 0 && isVertexTruncationReason(lastFinishReason)) {
&& isVertexTruncatedTurn(lastFinishReason, toolCallsStarted)) {
yield { type: "error", message: vertexTruncationErrorMessage(lastFinishReason) };
return;
}
Expand Down Expand Up @@ -752,7 +752,7 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte
// Fail-closed truncation, same as the stream path: a non-stream turn cut off mid tool call
// (MAX_TOKENS / MALFORMED_FUNCTION_CALL) surfaces an error instead of a silent done.
if ((provider.googleMode === "vertex" || provider.googleMode === "cloud-code-assist")
&& toolCallsStarted > 0 && isVertexTruncationReason(candidates?.[0]?.finishReason)) {
&& isVertexTruncatedTurn(candidates?.[0]?.finishReason, toolCallsStarted)) {
return finish([{ type: "error", message: vertexTruncationErrorMessage(candidates?.[0]?.finishReason) }]);
}

Expand Down
29 changes: 28 additions & 1 deletion tests/google-vertex-stream.test.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { describe, expect, test } from "bun:test";
import { createGoogleAdapter as createGoogleAdapterProduction } from "../src/adapters/google";
import { isVertexTruncationReason, vertexTruncationErrorMessage } from "../src/adapters/google-truncation";
import { isVertexTruncatedTurn, isVertexTruncationReason, vertexTruncationErrorMessage } from "../src/adapters/google-truncation";
import { bridgeToResponsesSSE } from "../src/bridge";
import type { AdapterEvent, OcxProviderConfig } from "../src/types";
import { withTestTranslatorBudget } from "./helpers/translator-budget";
Expand Down Expand Up @@ -30,6 +30,14 @@ describe("vertex truncation helpers", () => {
expect(isVertexTruncationReason(undefined)).toBe(false);
expect(vertexTruncationErrorMessage("MAX_TOKENS")).toContain("truncated upstream");
});

test("MALFORMED_FUNCTION_CALL fails closed with zero started calls; MAX_TOKENS does not", () => {
expect(isVertexTruncatedTurn("MALFORMED_FUNCTION_CALL", 0)).toBe(true);
expect(isVertexTruncatedTurn("MAX_TOKENS", 0)).toBe(false);
expect(isVertexTruncatedTurn("MAX_TOKENS", 1)).toBe(true);
expect(isVertexTruncatedTurn("STOP", 5)).toBe(false);
expect(isVertexTruncatedTurn(undefined, 5)).toBe(false);
});
});

describe("vertex parseStream fail-closed truncation", () => {
Expand Down Expand Up @@ -77,6 +85,17 @@ describe("vertex parseStream fail-closed truncation", () => {
expect(text).toContain('"incomplete_details":{"reason":"max_output_tokens"}');
});

test("MALFORMED_FUNCTION_CALL with NO emitted call part yields a terminal error, not done", async () => {
// The malformed call is dropped upstream, so the final chunk usually carries only the
// finishReason. Without the guard this surfaced as a clean empty completion.
const events = await collect(vertexProvider, [
{ candidates: [{ finishReason: "MALFORMED_FUNCTION_CALL" }], usageMetadata: { promptTokenCount: 5, candidatesTokenCount: 0 } },
]);
const last = events[events.length - 1];
expect(last.type).toBe("error");
expect(events.some(e => e.type === "done")).toBe(false);
});

test("usage-only final chunk (no candidates) is not dropped", async () => {
const events = await collect(vertexProvider, [
{ candidates: [{ content: { parts: [{ text: "hi" }] } }] },
Expand All @@ -98,6 +117,14 @@ describe("vertex parseResponse fail-closed truncation (non-streaming)", () => {
expect(events.some(e => e.type === "done")).toBe(false);
});

test("MALFORMED_FUNCTION_CALL with no call part yields a terminal error, not done", async () => {
const adapter = createGoogleAdapter(vertexProvider);
const body = JSON.stringify({ candidates: [{ finishReason: "MALFORMED_FUNCTION_CALL" }], usageMetadata: { promptTokenCount: 5, candidatesTokenCount: 0 } });
const events = await adapter.parseResponse!(new Response(body, { status: 200 }));
expect(events[events.length - 1].type).toBe("error");
expect(events.some(e => e.type === "done")).toBe(false);
});

test("clean STOP non-stream response yields done", async () => {
const adapter = createGoogleAdapter(vertexProvider);
const body = JSON.stringify({ candidates: [{ content: { parts: [{ text: "ok" }] }, finishReason: "STOP" }], usageMetadata: { promptTokenCount: 2, candidatesTokenCount: 1 } });
Expand Down
Loading