fix(google): surface blocked Gemini prompts (#116756)

Co-authored-by: Peter Steinberger <steipete@macos.shared>
This commit is contained in:
Peter Steinberger
2026-07-31 04:29:06 -07:00
committed by GitHub
parent 95df2dfd6a
commit 0d40cfbb53
4 changed files with 161 additions and 22 deletions

View File

@@ -541,6 +541,61 @@ describe("google transport stream", () => {
expect(result.content[2]).toHaveProperty("thoughtSignature", "Y2FsbF9zaWdfMQ==");
});
it.each([
{
provider: "google",
feedback: { blockReason: "SAFETY" },
expectedCode: "SAFETY",
expectedMessage: "Google prompt blocked (SAFETY)",
},
{
provider: "google",
feedback: {},
expectedCode: "PROMPT_BLOCKED",
expectedMessage: "Google prompt blocked (PROMPT_BLOCKED)",
},
{
provider: "google-vertex",
feedback: {
blockReason: "PROHIBITED_CONTENT",
blockReasonMessage: "Prompt violates provider safety policy",
},
expectedCode: "PROHIBITED_CONTENT",
expectedMessage:
"Google prompt blocked (PROHIBITED_CONTENT): Prompt violates provider safety policy",
},
{
provider: "google-vertex",
feedback: { blockReasonMessage: "Prompt violates provider safety policy" },
expectedCode: "PROMPT_BLOCKED",
expectedMessage:
"Google prompt blocked (PROMPT_BLOCKED): Prompt violates provider safety policy",
},
])(
"surfaces blocked $provider prompts as typed stream errors",
async ({ provider, feedback, expectedCode, expectedMessage }) => {
guardedFetchMock.mockResolvedValueOnce(buildSseResponse([{ promptFeedback: feedback }]));
if (provider === "google-vertex") {
vi.stubEnv("GOOGLE_CLOUD_PROJECT", "vertex-project");
vi.stubEnv("GOOGLE_CLOUD_LOCATION", "global");
googleAuthGetAccessTokenMock.mockResolvedValueOnce("ya29.vertex-token");
}
const result =
provider === "google-vertex"
? await runGoogleVertexStreamResult({ fetch: guardedFetchMock })
: await runGeminiStreamResult({ options: { apiKey: "gemini-api-key" } });
expect(result).toMatchObject({
stopReason: "error",
errorCode: expectedCode,
errorType: "google_prompt_blocked",
errorMessage: expectedMessage,
content: [],
});
},
);
it("rotates Gemini LLM API keys when a pre-stream request is rate limited", async () => {
vi.stubEnv("OPENCLAW_LIVE_GEMINI_KEY", "");
vi.stubEnv("GEMINI_API_KEYS", "gemini-key-2");

View File

@@ -130,6 +130,10 @@ const GOOGLE_VERTEX_DEFAULT_API_VERSION = "v1";
type GoogleSseChunk = {
responseId?: string;
promptFeedback?: {
blockReason?: string;
blockReasonMessage?: string;
};
candidates?: Array<{
content?: {
parts?: Array<{
@@ -1343,6 +1347,17 @@ function createGoogleTransportStreamFn(kind: CanonicalGoogleTransportApi): Strea
output.responseId ||= chunk.responseId;
updateUsage(output, model, chunk);
const candidate = chunk.candidates?.[0];
const promptFeedback = chunk.promptFeedback;
if (!candidate && promptFeedback) {
const blockReason =
normalizeOptionalString(promptFeedback.blockReason) ?? "PROMPT_BLOCKED";
const blockMessage = normalizeOptionalString(promptFeedback.blockReasonMessage);
const message = `Google prompt blocked (${blockReason})${blockMessage ? `: ${blockMessage}` : ""}`;
throw Object.assign(new Error(message), {
code: blockReason,
type: "google_prompt_blocked",
});
}
if (candidate?.content?.parts) {
for (const part of candidate.content.parts) {
const hasThoughtSignature =

View File

@@ -264,6 +264,66 @@ describe("consumeGoogleGenerateContentStream", () => {
});
describe("runGoogleGenerateContentLifecycle", () => {
it.each([
{ api: "google-generative-ai", blockReason: "SAFETY" },
{ api: "google-generative-ai", blockReason: undefined },
{ api: "google-vertex", blockReason: "SAFETY" },
{ api: "google-vertex", blockReason: undefined },
] as const)(
"surfaces blocked $api prompts as typed stream errors when blockReason is $blockReason",
async ({ api, blockReason }) => {
const targetModel = {
...model,
api,
provider: api === "google-vertex" ? "google-vertex" : "google",
} satisfies Model<"google-generative-ai" | "google-vertex">;
const output: AssistantMessage = {
...createOutput(),
api: targetModel.api,
provider: targetModel.provider,
};
const stream = new AssistantMessageEventStream();
await runGoogleGenerateContentLifecycle({
stream,
model: targetModel,
output,
createClient: () => ({
models: {
generateContentStream: async () =>
chunks([
{
promptFeedback: {
...(blockReason ? { blockReason } : {}),
blockReasonMessage: "Prompt violates provider safety policy",
},
usageMetadata: {
promptTokenCount: 12,
cachedContentTokenCount: 2,
totalTokenCount: 12,
},
} as GenerateContentResponse,
]),
},
}),
buildParams: () => ({ model: targetModel.id, contents: [] }),
nextToolCallId: () => "call_1",
});
const result = await stream.result();
const expectedBlockReason = blockReason ?? "PROMPT_BLOCKED";
expect(result).toMatchObject({
stopReason: "error",
errorCode: expectedBlockReason,
errorType: "google_prompt_blocked",
errorMessage: `Google prompt blocked (${expectedBlockReason}): Prompt violates provider safety policy`,
content: [],
usage: { input: 10, cacheRead: 2, totalTokens: 12 },
});
expect(result.usage.cost.total).toBeGreaterThan(0);
},
);
it("surfaces HTTP response body text from Google-compatible errors", async () => {
const output = createOutput();
const stream = new AssistantMessageEventStream();

View File

@@ -779,7 +779,38 @@ export async function consumeGoogleGenerateContentStream<T extends GoogleApiType
for await (const chunk of params.chunks) {
params.output.responseId ||= chunk.responseId;
if (chunk.usageMetadata) {
params.output.usage = {
input:
(chunk.usageMetadata.promptTokenCount || 0) -
(chunk.usageMetadata.cachedContentTokenCount || 0),
output:
(chunk.usageMetadata.candidatesTokenCount || 0) +
(chunk.usageMetadata.thoughtsTokenCount || 0),
cacheRead: chunk.usageMetadata.cachedContentTokenCount || 0,
cacheWrite: 0,
totalTokens: chunk.usageMetadata.totalTokenCount || 0,
cost: {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
total: 0,
},
};
calculateCost(params.model, params.output.usage);
}
const candidate = chunk.candidates?.[0];
const promptFeedback = chunk.promptFeedback;
if (!candidate && promptFeedback) {
const blockReason = promptFeedback.blockReason ?? "PROMPT_BLOCKED";
const blockMessage = promptFeedback.blockReasonMessage?.trim();
params.output.errorCode = blockReason;
params.output.errorType = "google_prompt_blocked";
throw new Error(
`Google prompt blocked (${blockReason})${blockMessage ? `: ${blockMessage}` : ""}`,
);
}
if (candidate?.content?.parts) {
for (const part of candidate.content.parts) {
if (part.text !== undefined) {
@@ -881,28 +912,6 @@ export async function consumeGoogleGenerateContentStream<T extends GoogleApiType
params.output.stopReason = "toolUse";
}
}
if (chunk.usageMetadata) {
params.output.usage = {
input:
(chunk.usageMetadata.promptTokenCount || 0) -
(chunk.usageMetadata.cachedContentTokenCount || 0),
output:
(chunk.usageMetadata.candidatesTokenCount || 0) +
(chunk.usageMetadata.thoughtsTokenCount || 0),
cacheRead: chunk.usageMetadata.cachedContentTokenCount || 0,
cacheWrite: 0,
totalTokens: chunk.usageMetadata.totalTokenCount || 0,
cost: {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
total: 0,
},
};
calculateCost(params.model, params.output.usage);
}
}
endCurrentBlock();