fix(assistant): send Gemini's thought signatures back with its tool calls

Gemini 3 signs a tool call on its OpenAI-compatible endpoint (extra_content.google.thought_signature) and refuses the next request with a 400 when the call comes back without it, so every turn in which the model called a tool failed. The Chat Completions client now keeps the signature per call as the step's provider state and sends it back with the call, to the model that signed it only. An own server entered with Gemini's endpoint now also gets Gemini's narrower tool schemas, which it refused before.
2026-11-05 04:04:07 +00:00 · 2026-10-27 21:06:29 +01:00
4 changed files with 217 additions and 14 deletions
@@ -375,6 +375,12 @@ const aiProviderStateSchema = z.discriminatedUnion('adapter', [
    adapter: z.literal('anthropic-messages'),
    blocks: z.array(z.record(z.string(), z.unknown())),
  }),
+  // Gemini's thought signature per tool call id, which its Chat Completions endpoint wants back with the call.
+  z.object({
+    adapter: z.literal('openai-chat'),
+    model: z.string(),
+    signatures: z.record(z.string(), z.string()),
+  }),
 ]);
 export type AiProviderState = z.infer<typeof aiProviderStateSchema>;

@@ -26,16 +26,29 @@ import type { AiConnectionKind } from '@trek/shared';
 /** The end of an OpenAI-compatible stream. */
 const DONE = '[DONE]';

+/** The host of Gemini's OpenAI-compatible endpoint, which an admin may also enter as their own server. */
+const GEMINI_HOST = 'generativelanguage.googleapis.com';
+
+function isGeminiHost(baseUrl: string): boolean {
+  try {
+    return new URL(baseUrl).hostname === GEMINI_HOST;
+  } catch {
+    return false;
+  }
+}
+
 /**
 * The provider-neutral history as Chat Completions messages, the instructions
 * and the trip card first. A tool message takes no image, so the images of a
- * step's tool results follow its last result as one user message.
+ * step's tool results follow its last result as one user message. `model` is
+ * the one the request goes to: a thought signature goes back only to the model
+ * that signed it, since another one cannot read it.
 */
-export function toChatMessages(request: AiChatRequest): Record<string, unknown>[] {
+export function toChatMessages(request: AiChatRequest, model?: string): Record<string, unknown>[] {
  const system = [request.system, request.context].filter(Boolean).join('\n\n');
  const out: Record<string, unknown>[] = [{ role: 'system', content: system }];
  request.messages.forEach((message, index) => {
-    out.push(toChatMessage(message));
+    out.push(toChatMessage(message, model));
    const withImages = toolImagesEndingAt(request.messages, index);
    if (withImages.length === 0) return;
    const content = withImages.flatMap((tool) => [
@@ -51,12 +64,14 @@ function imagePart(image: AiImage): Record<string, unknown> {
  return { type: 'image_url', image_url: { url: imageDataUrl(image), detail: image.detail ?? 'high' } };
 }

-function toChatMessage(message: AiChatMessage): Record<string, unknown> {
+function toChatMessage(message: AiChatMessage, model: string | undefined): Record<string, unknown> {
  if (message.role === 'user') {
    if (!message.images?.length) return { role: 'user', content: message.text };
    return { role: 'user', content: [{ type: 'text', text: message.text }, ...message.images.map(imagePart)] };
  }
  if (message.role === 'tool') return { role: 'tool', tool_call_id: message.toolCallId, content: message.output };
+  const state = message.providerState;
+  const signatures = state?.adapter === 'openai-chat' && state.model === model ? state.signatures : {};
  return {
    role: 'assistant',
    content: message.text || null,
@@ -66,19 +81,32 @@ function toChatMessage(message: AiChatMessage): Record<string, unknown> {
            id: call.id,
            type: 'function',
            function: { name: call.name, arguments: JSON.stringify(call.args) },
+            ...(signatures[call.id] ? { extra_content: { google: { thought_signature: signatures[call.id] } } } : {}),
          })),
        }
      : {}),
  };
 }

+/** Gemini signs a tool call on its first piece; the call has to go back with it or the next request is refused. */
+function thoughtSignatureOf(piece: Record<string, unknown>): string {
+  return asText(asRecord(asRecord(piece.extra_content).google).thought_signature);
+}
+
+/** The id a call goes on with; a provider that names none gets one from the call's position. */
+function callId(index: number, id: string): string {
+  return id || `call_${index}`;
+}
+
 const FINISH_REASONS: Record<string, AiFinishReason> = { tool_calls: 'tool-calls', length: 'length' };

 /**
 * The OpenAI-compatible Chat Completions API, streamed: OpenAI itself, Gemini
 * (its `/v1beta/openai` endpoint), OpenRouter and an admin's own server. The
 * kind decides only what a request may carry: Gemini gets its narrower tool
- * schemas, OpenAI the prompt cache key, which the others do not know.
+ * schemas, also when an admin enters its endpoint as their own server, and
+ * OpenAI the prompt cache key, which the others do not know. Gemini's thought
+ * signatures come back as the turn's provider state.
 */
 export class OpenAiChatClient implements AiChatClient {
  /** OpenRouter searches the web for any model through its web plugin; the others have no search here. */
@@ -86,20 +114,25 @@ export class OpenAiChatClient implements AiChatClient {
    return this.kind === 'openrouter';
  }

+  /** The kind whose tool schemas the endpoint takes. */
+  private readonly schemaKind: AiConnectionKind;
+
  constructor(
    readonly kind: AiConnectionKind,
    readonly model: string,
    private readonly baseUrl: string,
    private readonly apiKey: string | null,
    private readonly fetcher: AiFetch,
-  ) {}
+  ) {
+    this.schemaKind = kind === 'custom' && isGeminiHost(baseUrl) ? 'gemini' : kind;
+  }

  async *stream(request: AiChatRequest, signal: AbortSignal): AsyncGenerator<AiChatEvent> {
    const body = {
      model: this.model,
      stream: true,
      stream_options: { include_usage: true },
-      messages: toChatMessages(request),
+      messages: toChatMessages(request, this.model),
      ...(request.tools.length > 0
        ? {
            tools: request.tools.map((tool) => ({
@@ -107,7 +140,7 @@ export class OpenAiChatClient implements AiChatClient {
              function: {
                name: tool.name,
                description: tool.description,
-                parameters: toProviderSchema(tool.parameters, this.kind),
+                parameters: toProviderSchema(tool.parameters, this.schemaKind),
              },
            })),
            tool_choice: 'auto',
@@ -153,10 +186,22 @@ export class OpenAiChatClient implements AiChatClient {
      cachedInputTokens: turn.cachedInputTokens,
      outputTokens: turn.outputTokens,
    };
+    const signatures = turn.signatures();
+    if (Object.keys(signatures).length > 0) {
+      yield { type: 'provider-state', state: { adapter: 'openai-chat', model: this.model, signatures } };
+    }
    yield { type: 'finish', reason: turn.finishReason() };
  }
 }

+/** A tool call while it streams in: its id, name and arguments so far, and Gemini's signature. */
+interface PendingCall {
+  id: string;
+  name: string;
+  args: string;
+  signature: string;
+}
+
 /**
 * What one streamed Chat Completions turn has said so far. A tool call
 * arrives in pieces keyed by its index (its id and name in the first, the
@@ -168,7 +213,7 @@ class ChatTurn {
  outputTokens = 0;
  reason: string | null = null;
  done = false;
-  private readonly pending = new Map<number, { id: string; name: string; args: string }>();
+  private readonly pending = new Map<number, PendingCall>();
  private readonly cited = new Set<string>();
  private fresh: AiSource[] = [];

@@ -199,10 +244,11 @@ class ChatTurn {
      const piece = asRecord(raw);
      const index = typeof piece.index === 'number' ? piece.index : position;
      const fn = asRecord(piece.function);
-      const call = this.pending.get(index) ?? { id: '', name: '', args: '' };
+      const call = this.pending.get(index) ?? { id: '', name: '', args: '', signature: '' };
      call.id ||= asText(piece.id);
      call.name ||= asText(fn.name);
      call.args += asText(fn.arguments);
+      call.signature ||= thoughtSignatureOf(piece);
      this.pending.set(index, call);
    });
    return asText(delta.content);
@@ -216,10 +262,24 @@ class ChatTurn {
  }

  calls(): AiToolCall[] {
-    return [...this.pending.entries()]
-      .sort(([a], [b]) => a - b)
-      .filter(([, call]) => call.name)
-      .map(([index, call]) => ({ id: call.id || `call_${index}`, name: call.name, args: parseToolArgs(call.args) }));
+    return this.named().map(([index, call]) => ({
+      id: callId(index, call.id),
+      name: call.name,
+      args: parseToolArgs(call.args),
+    }));
+  }
+
+  /** The thought signature of each signed call, by the id the call goes back with. */
+  signatures(): Record<string, string> {
+    return Object.fromEntries(
+      this.named()
+        .filter(([, call]) => call.signature)
+        .map(([index, call]) => [callId(index, call.id), call.signature]),
+    );
+  }
+
+  private named(): [number, PendingCall][] {
+    return [...this.pending.entries()].sort(([a], [b]) => a - b).filter(([, call]) => call.name);
  }

  finishReason(): AiFinishReason {
@@ -529,6 +529,125 @@ describe('OpenAiChatClient', () => {
      await failure(client(streamResponse([{ choices: [{ delta: { content: 'half' } }] }])).stream(REQUEST, signal())),
    ).toMatchObject({ providerCode: 'stream_cut' });
  });
+
+  // Gemini 3 signs the first call of a step on its first piece, the way its Chat Completions endpoint streams it.
+  const SIGNED = [
+    {
+      choices: [
+        {
+          delta: {
+            tool_calls: [
+              {
+                index: 0,
+                id: 'function-call-1',
+                type: 'function',
+                function: { name: 'create_place', arguments: '{"name":"Nijo"}' },
+                extra_content: { google: { thought_signature: 'sig-A' } },
+              },
+            ],
+          },
+        },
+      ],
+    },
+    {
+      choices: [
+        {
+          delta: {
+            tool_calls: [{ index: 1, function: { name: 'list_places', arguments: '{}' } }],
+          },
+          finish_reason: 'tool_calls',
+        },
+      ],
+    },
+    '[DONE]',
+  ];
+
+  it('OACH-006: a signed call comes back as the provider state, keyed by the id it goes on with', async () => {
+    const events = await run(
+      new OpenAiChatClient(
+        'gemini',
+        'gemini-3-flash',
+        'https://g',
+        'k',
+        fetchAnswering(streamResponse(SIGNED)).fetcher,
+      ).stream(REQUEST, signal()),
+    );
+    expect(events.filter((event) => event.type === 'tool-call')).toEqual([
+      { type: 'tool-call', call: { id: 'function-call-1', name: 'create_place', args: { name: 'Nijo' } } },
+      { type: 'tool-call', call: { id: 'call_1', name: 'list_places', args: {} } },
+    ]);
+    expect(events.slice(-2)).toEqual([
+      {
+        type: 'provider-state',
+        state: { adapter: 'openai-chat', model: 'gemini-3-flash', signatures: { 'function-call-1': 'sig-A' } },
+      },
+      { type: 'finish', reason: 'tool-calls' },
+    ]);
+
+    // An unsigned turn keeps no state: nothing is stored for a provider that signs nothing.
+    const plain = await run(
+      new OpenAiChatClient('openai', 'm', 'https://x', 'k', fetchAnswering(streamResponse(STREAM)).fetcher).stream(
+        REQUEST,
+        signal(),
+      ),
+    );
+    expect(plain.some((event) => event.type === 'provider-state')).toBe(false);
+  });
+
+  it('OACH-007: the signature goes back with its call to the model that signed it, and to no other', async () => {
+    const signed: AiChatRequest = {
+      ...REQUEST,
+      messages: [
+        { role: 'user', text: 'Add Nijo' },
+        {
+          role: 'assistant',
+          text: '',
+          toolCalls: [
+            { id: 'function-call-1', name: 'create_place', args: { name: 'Nijo' } },
+            { id: 'call_1', name: 'list_places', args: {} },
+          ],
+          providerState: {
+            adapter: 'openai-chat',
+            model: 'gemini-3-flash',
+            signatures: { 'function-call-1': 'sig-A' },
+          },
+        },
+        { role: 'tool', toolCallId: 'function-call-1', name: 'create_place', output: '{"id":7}' },
+        { role: 'tool', toolCallId: 'call_1', name: 'list_places', output: '[]' },
+      ],
+    };
+    const http = fetchAnswering(streamResponse(['[DONE]']));
+    await run(
+      new OpenAiChatClient('gemini', 'gemini-3-flash', 'https://g', 'k', http.fetcher).stream(signed, signal()),
+    );
+    expect(dig(http.body(), 'messages', 2, 'tool_calls')).toEqual([
+      {
+        id: 'function-call-1',
+        type: 'function',
+        function: { name: 'create_place', arguments: '{"name":"Nijo"}' },
+        extra_content: { google: { thought_signature: 'sig-A' } },
+      },
+      { id: 'call_1', type: 'function', function: { name: 'list_places', arguments: '{}' } },
+    ]);
+
+    // A thread that moved on to another model sends the calls plain: no other model reads the signature.
+    for (const messages of [toChatMessages(signed, 'gpt-5-mini'), toChatMessages(signed)]) {
+      expect(JSON.stringify(messages)).not.toContain('extra_content');
+    }
+  });
+
+  it('OACH-008: an own server on Gemini’s endpoint gets Gemini’s narrower schemas, any other the full ones', async () => {
+    const parameters = async (baseUrl: string) => {
+      const http = fetchAnswering(streamResponse(['[DONE]']));
+      await run(new OpenAiChatClient('custom', 'm', baseUrl, 'k', http.fetcher).stream(REQUEST, signal()));
+      return dig(http.body(), 'tools', 0, 'function', 'parameters');
+    };
+    expect(await parameters('https://generativelanguage.googleapis.com/v1beta/openai')).toEqual(
+      toProviderSchema(TOOL.parameters, 'gemini'),
+    );
+    expect(await parameters('http://ollama:11434/v1')).toEqual(toProviderSchema(TOOL.parameters, 'custom'));
+    expect(await parameters('not a url')).toEqual(toProviderSchema(TOOL.parameters, 'custom'));
+  });
 });

 describe('OpenAiResponsesClient', () => {
@@ -8,6 +8,7 @@ import type { AssistantMessageRow } from '../../../../src/db/repositories/Assist
 import type { AssistantToolCallRow } from '../../../../src/db/repositories/AssistantToolCalls.repository';
 import type { AiChatMessage } from '../../../../src/nest/ai-connections/ai-chat-client.interface';
 import { toAnthropicMessages } from '../../../../src/nest/ai-connections/providers/anthropic-messages.client';
+import { toChatMessages } from '../../../../src/nest/ai-connections/providers/openai-chat.client';
 import { toResponsesInput } from '../../../../src/nest/ai-connections/providers/openai-responses.client';
 import {
  messageIdOf,
@@ -573,6 +574,23 @@ describe('toChatHistory: calls the model made of the server tools (#2845)', () =
    expect(messages.flatMap((message) => message.content).some((block) => block.text === '')).toBe(false);
  });

+  it('ASTHIS-018: Gemini gets every stored call back with its thought signature, an applied change included', () => {
+    const state = (id: string, signature: string) => ({
+      adapter: 'openai-chat',
+      model: 'gemini-3-flash',
+      signatures: { [id]: signature },
+    });
+    type SentCall = { id: string; extra_content?: { google: { thought_signature: string } } };
+    const history = appliedAfterANote(state('call_note', 'sig-note'), state('call_move', 'sig-move'));
+    const sent = toChatMessages({ system: '', messages: history, tools: [] }, 'gemini-3-flash')
+      .flatMap((message) => (message.tool_calls ?? []) as SentCall[])
+      .map((call) => [call.id, call.extra_content?.google.thought_signature]);
+    expect(sent).toEqual([
+      ['call_note', 'sig-note'],
+      ['call_move', 'sig-move'],
+    ]);
+  });
+
  it('ASTHIS-017: a step that only kept a note is the model call it was, answered with what was kept', () => {
    const history = toChatHistory(
      [