From c349ea6b78b85f43ce476b785fb1fe949ca66003 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 9 Jun 2026 13:11:43 +0200 Subject: [PATCH] fix: preserve gemini thought signatures (#25933) AI Bridge reserializes OpenAI chat-completions requests before sending them upstream. For Gemini OpenAI-compatible routes, that OpenAI typed-parameter round trip drops `tool_calls[].extra_content.google.thought_signature`, so Google rejects tool-result continuations with `Function call is missing a thought_signature`. This PR: - patches the AI Bridge upstream serialization boundary for Gemini OpenAI-compatible chat completions - shares the Gemini thought-signature patching helpers with chatd's OpenAI-compatible transport patch to keep behavior consistent - treats direct Google OpenAI-compatible upstream endpoints as Gemini-scoped even when the request model is an alias - adds the Google fallback thought signature to every assistant tool call in the active turn, including parallel tool calls - covers the regression that `extra_content` is dropped before the upstream body is patched > Mux updated this PR description on behalf of Mike. --------- Co-authored-by: Susana Cardoso Ferreira --- .../intercept/chatcompletions/blocking.go | 10 +- .../chatcompletions/google_openai_compat.go | 37 ++++ .../google_openai_compat_internal_test.go | 100 +++++++++++ .../intercept/chatcompletions/streaming.go | 4 +- .../chatprovider/openai_compat_patches.go | 101 +---------- .../openai_compat_patches_test.go | 11 +- internal/googleopenai/thought_signature.go | 163 +++++++++++++++++ .../googleopenai/thought_signature_test.go | 167 ++++++++++++++++++ 8 files changed, 488 insertions(+), 105 deletions(-) create mode 100644 aibridge/intercept/chatcompletions/google_openai_compat.go create mode 100644 aibridge/intercept/chatcompletions/google_openai_compat_internal_test.go create mode 100644 internal/googleopenai/thought_signature.go create mode 100644 internal/googleopenai/thought_signature_test.go diff --git a/aibridge/intercept/chatcompletions/blocking.go b/aibridge/intercept/chatcompletions/blocking.go index c35cfcc86a..29cc8bc704 100644 --- a/aibridge/intercept/chatcompletions/blocking.go +++ b/aibridge/intercept/chatcompletions/blocking.go @@ -291,7 +291,15 @@ func (i *BlockingInterception) newChatCompletionWithKey(ctx context.Context, svc _, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) defer tracing.EndSpanErr(span, &outErr) - return svc.New(ctx, i.req.ChatCompletionNewParams, opts...) + requestOpts, overrideBody, err := i.chatCompletionRequestOptions(opts) + if err != nil { + return nil, xerrors.Errorf("prepare request body: %w", err) + } + params := i.req.ChatCompletionNewParams + if overrideBody { + params = openai.ChatCompletionNewParams{} + } + return svc.New(ctx, params, requestOpts...) } // newChatCompletionWithKeyFailover walks the centralized key pool, trying each diff --git a/aibridge/intercept/chatcompletions/google_openai_compat.go b/aibridge/intercept/chatcompletions/google_openai_compat.go new file mode 100644 index 0000000000..251cbc71a0 --- /dev/null +++ b/aibridge/intercept/chatcompletions/google_openai_compat.go @@ -0,0 +1,37 @@ +package chatcompletions + +import ( + "encoding/json" + "slices" + + "github.com/openai/openai-go/v3/option" + + "github.com/coder/coder/v2/internal/googleopenai" +) + +func (i *interceptionBase) chatCompletionRequestBody() ([]byte, error) { + body, err := json.Marshal(i.req.ChatCompletionNewParams) + if err != nil { + return nil, err + } + if !googleopenai.ShouldPatchGoogleUpstreamRequest(i.cfg.BaseURL) { + return body, nil + } + patched, _, err := googleopenai.PatchThoughtSignatures(body) + if err != nil { + return nil, err + } + return patched, nil +} + +func (i *interceptionBase) chatCompletionRequestOptions(opts []option.RequestOption) ([]option.RequestOption, bool, error) { + if !googleopenai.ShouldPatchGoogleUpstreamRequest(i.cfg.BaseURL) { + return opts, false, nil + } + body, err := i.chatCompletionRequestBody() + if err != nil { + return nil, false, err + } + updated := slices.Clone(opts) + return append(updated, option.WithRequestBody("application/json", body)), true, nil +} diff --git a/aibridge/intercept/chatcompletions/google_openai_compat_internal_test.go b/aibridge/intercept/chatcompletions/google_openai_compat_internal_test.go new file mode 100644 index 0000000000..826dba07b6 --- /dev/null +++ b/aibridge/intercept/chatcompletions/google_openai_compat_internal_test.go @@ -0,0 +1,100 @@ +package chatcompletions + +import ( + "encoding/json" + "testing" + + "github.com/openai/openai-go/v3/option" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/aibridge/config" + "github.com/coder/coder/v2/internal/googleopenai" +) + +func TestGoogleOpenAICompatThoughtSignaturePatchSurvivesParamRoundTrip(t *testing.T) { + t.Parallel() + + const originalSignature = "SIG123" + raw := []byte(`{ + "model":"gemini-3.5-flash", + "stream":true, + "messages":[ + {"role":"user","content":"write a file"}, + { + "role":"assistant", + "content":"I'll search for available workspace templates.", + "tool_calls":[ + { + "id":"pbk491lp", + "function":{"arguments":"{}","name":"list_templates"}, + "type":"function", + "extra_content":{"google":{"thought_signature":"` + originalSignature + `"}} + } + ] + }, + {"role":"tool","tool_call_id":"pbk491lp","content":"{}"} + ] + }`) + + var req ChatCompletionNewParamsWrapper + require.NoError(t, json.Unmarshal(raw, &req)) + + roundTripped, err := json.Marshal(req.ChatCompletionNewParams) + require.NoError(t, err) + require.Empty(t, googleThoughtSignatureFromBody(t, roundTripped, 1, 0), + "openai-go drops extra_content during the typed param round-trip") + + body, err := (&interceptionBase{ + req: &req, + cfg: config.OpenAI{BaseURL: "https://generativelanguage.googleapis.com/v1beta/openai/"}, + }).chatCompletionRequestBody() + require.NoError(t, err) + require.Equal(t, googleopenai.DummyThoughtSignature, googleThoughtSignatureFromBody(t, body, 1, 0)) +} + +func TestGoogleOpenAICompatChatCompletionRequestOptions(t *testing.T) { + t.Parallel() + + var req ChatCompletionNewParamsWrapper + require.NoError(t, json.Unmarshal([]byte(`{ + "model":"gemini-3.5-flash", + "messages":[ + {"role":"user","content":"current turn"}, + { + "role":"assistant", + "tool_calls":[{"id":"call-1","function":{"arguments":"{}","name":"list_templates"},"type":"function"}] + } + ] + }`), &req)) + + opts := make([]option.RequestOption, 1) + updated, overrideBody, err := (&interceptionBase{ + req: &req, + cfg: config.OpenAI{BaseURL: "https://generativelanguage.googleapis.com/v1beta/openai/"}, + }).chatCompletionRequestOptions(opts) + require.NoError(t, err) + require.True(t, overrideBody) + require.Len(t, opts, 1) + require.Len(t, updated, 2) +} + +func googleThoughtSignatureFromBody(t *testing.T, body []byte, messageIndex int, toolCallIndex int) string { + t.Helper() + + var payload map[string]any + require.NoError(t, json.Unmarshal(body, &payload)) + messages, ok := payload["messages"].([]any) + require.True(t, ok) + require.Greater(t, len(messages), messageIndex) + message, ok := messages[messageIndex].(map[string]any) + require.True(t, ok) + toolCalls, ok := message["tool_calls"].([]any) + require.True(t, ok) + require.Greater(t, len(toolCalls), toolCallIndex) + toolCall, ok := toolCalls[toolCallIndex].(map[string]any) + require.True(t, ok) + extraContent, _ := toolCall["extra_content"].(map[string]any) + google, _ := extraContent["google"].(map[string]any) + signature, _ := google["thought_signature"].(string) + return signature +} diff --git a/aibridge/intercept/chatcompletions/streaming.go b/aibridge/intercept/chatcompletions/streaming.go index a03ee09768..694b09893f 100644 --- a/aibridge/intercept/chatcompletions/streaming.go +++ b/aibridge/intercept/chatcompletions/streaming.go @@ -193,7 +193,9 @@ func (i *StreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Re // We take control of request body here and pass it to the SDK as a raw byte slice. // This is because the SDK's serialization applies hidden request options that result in // unexpected, breaking behavior. See https://github.com/coder/aibridge/pull/164 - body, err := json.Marshal(i.req.ChatCompletionNewParams) + // chatCompletionRequestBody also applies provider-specific + // compatibility patches to the exact body sent upstream. + body, err := i.chatCompletionRequestBody() if err != nil { return xerrors.Errorf("marshal request body: %w", err) } diff --git a/coderd/x/chatd/chatprovider/openai_compat_patches.go b/coderd/x/chatd/chatprovider/openai_compat_patches.go index 26a1f80631..8c60f2a16b 100644 --- a/coderd/x/chatd/chatprovider/openai_compat_patches.go +++ b/coderd/x/chatd/chatprovider/openai_compat_patches.go @@ -6,16 +6,13 @@ import ( "io" "net/http" "strings" + + "github.com/coder/coder/v2/internal/googleopenai" ) // OpenAI-compatible providers share an API shape but differ in the exact JSON // they accept. These patches adjust Fantasy's serialized request body at the // transport boundary so higher-level generation code can stay provider agnostic. -// -// googleOpenAICompatDummyThoughtSignature is Google's documented last-resort -// bypass for callers that cannot preserve a real Gemini thought signature. -// See https://ai.google.dev/gemini-api/docs/thought-signatures. -const googleOpenAICompatDummyThoughtSignature = "skip_thought_signature_validator" func withOpenAICompatRequestPatches( client *http.Client, @@ -91,8 +88,8 @@ func patchOpenAICompatChatCompletionsBody(body []byte, baseURL string, modelID s } changed := rewriteOpenAICompatSingleToolChoice(payload) - if shouldAddGoogleOpenAICompatThoughtSignatures(baseURL, modelID) { - changed = addGoogleOpenAICompatThoughtSignatures(payload) || changed + if googleopenai.ShouldPatchOpenAICompatRequest(baseURL, modelID) { + changed = googleopenai.AddThoughtSignaturesToLatestTurn(payload) || changed } if !changed { return body @@ -144,93 +141,3 @@ func rewriteOpenAICompatSingleToolChoice(payload map[string]any) bool { payload["tool_choice"] = "required" return true } - -// shouldAddGoogleOpenAICompatThoughtSignatures detects direct Gemini OpenAI -// endpoints and Coder AI Bridge Gemini routes. Other gateways, such as Vercel, -// keep their own provider-specific compatibility behavior. -func shouldAddGoogleOpenAICompatThoughtSignatures(baseURL string, modelID string) bool { - parsed, ok := parseProviderBaseURL(baseURL) - if !ok { - return false - } - host := strings.ToLower(parsed.Hostname()) - path := strings.ToLower(parsed.EscapedPath()) - if host == "generativelanguage.googleapis.com" && strings.Contains(path, "/openai") { - return true - } - return host == "coder-aibridge" && isGeminiModelID(modelID) -} - -func isGeminiModelID(modelID string) bool { - modelID = strings.ToLower(strings.TrimSpace(modelID)) - return strings.HasPrefix(modelID, "gemini-") || strings.Contains(modelID, "/gemini-") -} - -// addGoogleOpenAICompatThoughtSignatures adds a dummy thought signature to the -// first tool call on each assistant tool-call message in the latest user turn. -// Gemini validates tool-call history with thought signatures, but -// OpenAI-compatible serialization can drop the original provider metadata. -func addGoogleOpenAICompatThoughtSignatures(payload map[string]any) bool { - messages, ok := payload["messages"].([]any) - if !ok { - return false - } - - currentTurnStart := -1 - for i, raw := range messages { - message, ok := raw.(map[string]any) - if !ok { - continue - } - if role, _ := message["role"].(string); role == "user" { - currentTurnStart = i - } - } - - if currentTurnStart == -1 { - return false - } - - changed := false - for _, raw := range messages[currentTurnStart+1:] { - message, ok := raw.(map[string]any) - if !ok || !isOpenAICompatAssistantRole(message["role"]) { - continue - } - toolCalls, ok := message["tool_calls"].([]any) - if !ok || len(toolCalls) == 0 { - continue - } - firstToolCall, ok := toolCalls[0].(map[string]any) - if !ok { - continue - } - if ensureGoogleOpenAICompatThoughtSignature(firstToolCall) { - changed = true - } - } - return changed -} - -func isOpenAICompatAssistantRole(role any) bool { - roleValue, _ := role.(string) - return roleValue == "assistant" || roleValue == "model" -} - -func ensureGoogleOpenAICompatThoughtSignature(toolCall map[string]any) bool { - extraContent, _ := toolCall["extra_content"].(map[string]any) - google, _ := extraContent["google"].(map[string]any) - if signature, _ := google["thought_signature"].(string); signature != "" { - return false - } - if extraContent == nil { - extraContent = map[string]any{} - toolCall["extra_content"] = extraContent - } - if google == nil { - google = map[string]any{} - extraContent["google"] = google - } - google["thought_signature"] = googleOpenAICompatDummyThoughtSignature - return true -} diff --git a/coderd/x/chatd/chatprovider/openai_compat_patches_test.go b/coderd/x/chatd/chatprovider/openai_compat_patches_test.go index c6042c0c63..abb97079c6 100644 --- a/coderd/x/chatd/chatprovider/openai_compat_patches_test.go +++ b/coderd/x/chatd/chatprovider/openai_compat_patches_test.go @@ -12,10 +12,9 @@ import ( "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" + "github.com/coder/coder/v2/internal/googleopenai" ) -const dummyThoughtSignature = "skip_thought_signature_validator" - func TestModelFromConfig_GeminiOpenAICompatThoughtSignatures(t *testing.T) { t.Parallel() @@ -26,9 +25,9 @@ func TestModelFromConfig_GeminiOpenAICompatThoughtSignatures(t *testing.T) { messages := body["messages"].([]any) require.Empty(t, thoughtSignature(t, messages[1], 0)) - require.Equal(t, dummyThoughtSignature, thoughtSignature(t, messages[4], 0)) - require.Empty(t, thoughtSignature(t, messages[4], 1)) - require.Equal(t, dummyThoughtSignature, thoughtSignature(t, messages[6], 0)) + require.Equal(t, googleopenai.DummyThoughtSignature, thoughtSignature(t, messages[4], 0)) + require.Equal(t, googleopenai.DummyThoughtSignature, thoughtSignature(t, messages[4], 1)) + require.Equal(t, googleopenai.DummyThoughtSignature, thoughtSignature(t, messages[6], 0)) }) t.Run("Coder AI Bridge Gemini route receives current turn thought signature", func(t *testing.T) { @@ -37,7 +36,7 @@ func TestModelFromConfig_GeminiOpenAICompatThoughtSignatures(t *testing.T) { body := generateOpenAICompatRequest(t, "http://coder-aibridge/v1", "gemini-3.5-flash") messages := body["messages"].([]any) - require.Equal(t, dummyThoughtSignature, thoughtSignature(t, messages[4], 0)) + require.Equal(t, googleopenai.DummyThoughtSignature, thoughtSignature(t, messages[4], 0)) }) t.Run("Vercel OpenAI-compatible Gemini route is unchanged", func(t *testing.T) { diff --git a/internal/googleopenai/thought_signature.go b/internal/googleopenai/thought_signature.go new file mode 100644 index 0000000000..84467bec18 --- /dev/null +++ b/internal/googleopenai/thought_signature.go @@ -0,0 +1,163 @@ +// Package googleopenai contains compatibility helpers for Google's +// OpenAI-compatible Gemini APIs. +package googleopenai + +import ( + "encoding/json" + "net/url" + "strings" +) + +// DummyThoughtSignature is Google's documented last-resort bypass for callers +// that cannot preserve a real Gemini thought signature through OpenAI-compatible +// serialization. See https://ai.google.dev/gemini-api/docs/thought-signatures. +const DummyThoughtSignature = "skip_thought_signature_validator" + +// ShouldPatchOpenAICompatRequest reports whether a client-side +// OpenAI-compatible request should carry Gemini thought signatures. +func ShouldPatchOpenAICompatRequest(baseURL string, modelID string) bool { + // Direct Google endpoints are already provider-scoped. Patch them even when + // the configured model ID is an alias without a Gemini prefix. + if isDirectGeminiOpenAIEndpoint(baseURL) { + return true + } + return isCoderAIBridgeEndpoint(baseURL) && isGeminiModelID(modelID) +} + +// ShouldPatchGoogleUpstreamRequest reports whether an AI Bridge upstream +// OpenAI-compatible request should carry Gemini thought signatures. +func ShouldPatchGoogleUpstreamRequest(baseURL string) bool { + return isDirectGeminiOpenAIEndpoint(baseURL) +} + +// Vertex AI has different hosts and paths. Add it here only with a fixture that +// confirms it accepts the same thought-signature fallback shape. +func isDirectGeminiOpenAIEndpoint(baseURL string) bool { + parsed, ok := parseBaseURL(baseURL) + if !ok { + return false + } + host := strings.ToLower(parsed.Hostname()) + path := strings.ToLower(parsed.EscapedPath()) + return host == "generativelanguage.googleapis.com" && strings.Contains(path, "/openai") +} + +func isCoderAIBridgeEndpoint(baseURL string) bool { + parsed, ok := parseBaseURL(baseURL) + if !ok { + return false + } + return strings.ToLower(parsed.Hostname()) == "coder-aibridge" +} + +// parseBaseURL parses a provider base URL, handling bare hostnames without +// a scheme by prepending "https://". +func parseBaseURL(baseURL string) (*url.URL, bool) { + baseURL = strings.TrimSpace(baseURL) + if baseURL == "" { + return nil, false + } + parsed, err := url.Parse(baseURL) + if err == nil && parsed.Hostname() == "" && !strings.Contains(baseURL, "://") { + parsed, err = url.Parse("https://" + baseURL) + } + if err != nil { + return nil, false + } + return parsed, true +} + +func isGeminiModelID(modelID string) bool { + modelID = strings.ToLower(strings.TrimSpace(modelID)) + return strings.HasPrefix(modelID, "gemini-") || strings.Contains(modelID, "/gemini-") +} + +// PatchThoughtSignatures adds fallback thought signatures to Gemini tool-call +// history in body. It returns changed=false when no patch is needed. +func PatchThoughtSignatures(body []byte) ([]byte, bool, error) { + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + return nil, false, err + } + if !AddThoughtSignaturesToLatestTurn(payload) { + return body, false, nil + } + patched, err := json.Marshal(payload) + if err != nil { + return nil, false, err + } + return patched, true, nil +} + +// AddThoughtSignaturesToLatestTurn patches only the current turn because +// completed tool-call/result pairs from earlier turns are not validated by +// Google as active function calls. +func AddThoughtSignaturesToLatestTurn(payload map[string]any) bool { + messages, ok := payload["messages"].([]any) + if !ok { + return false + } + + currentTurnStart := -1 + for i, raw := range messages { + message, ok := raw.(map[string]any) + if !ok { + continue + } + if role, _ := message["role"].(string); role == "user" { + currentTurnStart = i + } + } + if currentTurnStart == -1 { + return false + } + + changed := false + for _, raw := range messages[currentTurnStart+1:] { + message, ok := raw.(map[string]any) + if !ok || !isAssistantRole(message["role"]) { + continue + } + toolCalls, ok := message["tool_calls"].([]any) + if !ok || len(toolCalls) == 0 { + continue + } + // Every tool call in parallel batches needs a signature, + // not just the first one. + for _, rawToolCall := range toolCalls { + toolCall, ok := rawToolCall.(map[string]any) + if !ok { + continue + } + if ensureThoughtSignature(toolCall) { + changed = true + } + } + } + return changed +} + +// Gemini can serialize assistant messages with its native "model" role. +func isAssistantRole(role any) bool { + roleValue, _ := role.(string) + return roleValue == "assistant" || roleValue == "model" +} + +// Real provider signatures are preserved when present. +func ensureThoughtSignature(toolCall map[string]any) bool { + extraContent, _ := toolCall["extra_content"].(map[string]any) + google, _ := extraContent["google"].(map[string]any) + if signature, _ := google["thought_signature"].(string); signature != "" { + return false + } + if extraContent == nil { + extraContent = map[string]any{} + toolCall["extra_content"] = extraContent + } + if google == nil { + google = map[string]any{} + extraContent["google"] = google + } + google["thought_signature"] = DummyThoughtSignature + return true +} diff --git a/internal/googleopenai/thought_signature_test.go b/internal/googleopenai/thought_signature_test.go new file mode 100644 index 0000000000..c73bf25ee7 --- /dev/null +++ b/internal/googleopenai/thought_signature_test.go @@ -0,0 +1,167 @@ +package googleopenai_test + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/internal/googleopenai" +) + +func TestShouldPatchOpenAICompatRequest(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + baseURL string + modelID string + want bool + }{ + { + name: "direct endpoint with gemini model", + baseURL: "https://generativelanguage.googleapis.com/v1beta/openai/", + modelID: "gemini-3.5-flash", + want: true, + }, + { + name: "direct endpoint does not require gemini model name", + baseURL: "https://generativelanguage.googleapis.com/v1beta/openai/", + modelID: "gpt-4o", + want: true, + }, + { + name: "coder aibridge gemini route", + baseURL: "http://coder-aibridge/v1", + modelID: "gemini-3.5-flash", + want: true, + }, + { + name: "aibridge endpoint requires gemini model", + baseURL: "http://coder-aibridge/v1", + modelID: "gpt-4o", + }, + { + name: "other gateway unchanged", + baseURL: "https://gateway.vercel.ai/v1", + modelID: "google/gemini-3.5-flash", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + require.Equal(t, tt.want, googleopenai.ShouldPatchOpenAICompatRequest(tt.baseURL, tt.modelID)) + }) + } +} + +func TestShouldPatchGoogleUpstreamRequest(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + baseURL string + want bool + }{ + { + name: "gemini api openai endpoint", + baseURL: "https://generativelanguage.googleapis.com/v1beta/openai/", + want: true, + }, + { + name: "openai endpoint", + baseURL: "https://api.openai.com/v1/", + }, + { + name: "vertex endpoint not enabled without fixture", + baseURL: "https://us-central1-aiplatform.googleapis.com/v1/", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + require.Equal(t, tt.want, googleopenai.ShouldPatchGoogleUpstreamRequest(tt.baseURL)) + }) + } +} + +func TestAddThoughtSignaturesToLatestTurn(t *testing.T) { + t.Parallel() + + payload := decodePayload(t, []byte(`{ + "messages":[ + {"role":"user","content":"previous turn"}, + { + "role":"assistant", + "tool_calls":[{"id":"old-call","type":"function","function":{"name":"old","arguments":"{}"}}] + }, + {"role":"tool","tool_call_id":"old-call","content":"{}"}, + {"role":"user","content":"current turn"}, + { + "role":"model", + "tool_calls":[ + {"id":"call-1","type":"function","function":{"name":"list_templates","arguments":"{}"}}, + {"id":"call-2","type":"function","function":{"name":"read_template","arguments":"{}"}} + ] + } + ] + }`)) + + require.True(t, googleopenai.AddThoughtSignaturesToLatestTurn(payload)) + require.Empty(t, thoughtSignature(t, payload, 1, 0), "previous turns should stay unchanged") + require.Equal(t, googleopenai.DummyThoughtSignature, thoughtSignature(t, payload, 4, 0)) + require.Equal(t, googleopenai.DummyThoughtSignature, thoughtSignature(t, payload, 4, 1)) +} + +func TestAddThoughtSignaturesToLatestTurnPreservesRealSignature(t *testing.T) { + t.Parallel() + + payload := decodePayload(t, []byte(`{ + "messages":[ + {"role":"user","content":"current turn"}, + { + "role":"assistant", + "tool_calls":[{ + "id":"call-1", + "type":"function", + "function":{"name":"list_templates","arguments":"{}"}, + "extra_content":{"google":{"thought_signature":"real-signature"}} + }] + } + ] + }`)) + + require.False(t, googleopenai.AddThoughtSignaturesToLatestTurn(payload)) + require.Equal(t, "real-signature", thoughtSignature(t, payload, 1, 0)) +} + +func decodePayload(t *testing.T, body []byte) map[string]any { + t.Helper() + + var payload map[string]any + require.NoError(t, json.Unmarshal(body, &payload)) + return payload +} + +func thoughtSignature(t *testing.T, payload map[string]any, messageIndex int, toolCallIndex int) string { + t.Helper() + + messages, ok := payload["messages"].([]any) + require.True(t, ok) + require.Greater(t, len(messages), messageIndex) + message, ok := messages[messageIndex].(map[string]any) + require.True(t, ok) + toolCalls, ok := message["tool_calls"].([]any) + require.True(t, ok) + require.Greater(t, len(toolCalls), toolCallIndex) + toolCall, ok := toolCalls[toolCallIndex].(map[string]any) + require.True(t, ok) + extraContent, _ := toolCall["extra_content"].(map[string]any) + google, _ := extraContent["google"].(map[string]any) + signature, _ := google["thought_signature"].(string) + return signature +}