Skip to content

Commit 9595e6c

Browse files
github-actions[bot]ibetitsmikessncferreira
authored
fix: preserve gemini thought signatures (#25933) (#26169)
Backport of #25933 Original PR: #25933 — fix: preserve gemini thought signatures Merge commit: c349ea6 Requested by: @ssncferreira Co-authored-by: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Co-authored-by: Susana Cardoso Ferreira <susana@coder.com>
1 parent 6419f53 commit 9595e6c

8 files changed

Lines changed: 488 additions & 105 deletions

File tree

‎aibridge/intercept/chatcompletions/blocking.go‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -282,7 +282,15 @@ func (i *BlockingInterception) newChatCompletionWithKey(ctx context.Context, svc
282282
_, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...))
283283
defer tracing.EndSpanErr(span, &outErr)
284284

285-
return svc.New(ctx, i.req.ChatCompletionNewParams, opts...)
285+
requestOpts, overrideBody, err := i.chatCompletionRequestOptions(opts)
286+
if err != nil {
287+
return nil, xerrors.Errorf("prepare request body: %w", err)
288+
}
289+
params := i.req.ChatCompletionNewParams
290+
if overrideBody {
291+
params = openai.ChatCompletionNewParams{}
292+
}
293+
return svc.New(ctx, params, requestOpts...)
286294
}
287295

288296
// newChatCompletionWithKeyFailover walks the centralized key

‎aibridge/intercept/chatcompletions/google_openai_compat.go‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
package chatcompletions
2+
3+
import (
4+
"encoding/json"
5+
"slices"
6+
7+
"github.com/openai/openai-go/v3/option"
8+
9+
"github.com/coder/coder/v2/internal/googleopenai"
10+
)
11+
12+
func (i *interceptionBase) chatCompletionRequestBody() ([]byte, error) {
13+
body, err := json.Marshal(i.req.ChatCompletionNewParams)
14+
if err != nil {
15+
return nil, err
16+
}
17+
if !googleopenai.ShouldPatchGoogleUpstreamRequest(i.cfg.BaseURL) {
18+
return body, nil
19+
}
20+
patched, _, err := googleopenai.PatchThoughtSignatures(body)
21+
if err != nil {
22+
return nil, err
23+
}
24+
return patched, nil
25+
}
26+
27+
func (i *interceptionBase) chatCompletionRequestOptions(opts []option.RequestOption) ([]option.RequestOption, bool, error) {
28+
if !googleopenai.ShouldPatchGoogleUpstreamRequest(i.cfg.BaseURL) {
29+
return opts, false, nil
30+
}
31+
body, err := i.chatCompletionRequestBody()
32+
if err != nil {
33+
return nil, false, err
34+
}
35+
updated := slices.Clone(opts)
36+
return append(updated, option.WithRequestBody("application/json", body)), true, nil
37+
}

‎aibridge/intercept/chatcompletions/google_openai_compat_internal_test.go‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
package chatcompletions
2+
3+
import (
4+
"encoding/json"
5+
"testing"
6+
7+
"github.com/openai/openai-go/v3/option"
8+
"github.com/stretchr/testify/require"
9+
10+
"github.com/coder/coder/v2/aibridge/config"
11+
"github.com/coder/coder/v2/internal/googleopenai"
12+
)
13+
14+
func TestGoogleOpenAICompatThoughtSignaturePatchSurvivesParamRoundTrip(t *testing.T) {
15+
t.Parallel()
16+
17+
const originalSignature = "SIG123"
18+
raw := []byte(`{
19+
"model":"gemini-3.5-flash",
20+
"stream":true,
21+
"messages":[
22+
{"role":"user","content":"write a file"},
23+
{
24+
"role":"assistant",
25+
"content":"I'll search for available workspace templates.",
26+
"tool_calls":[
27+
{
28+
"id":"pbk491lp",
29+
"function":{"arguments":"{}","name":"list_templates"},
30+
"type":"function",
31+
"extra_content":{"google":{"thought_signature":"` + originalSignature + `"}}
32+
}
33+
]
34+
},
35+
{"role":"tool","tool_call_id":"pbk491lp","content":"{}"}
36+
]
37+
}`)
38+
39+
var req ChatCompletionNewParamsWrapper
40+
require.NoError(t, json.Unmarshal(raw, &req))
41+
42+
roundTripped, err := json.Marshal(req.ChatCompletionNewParams)
43+
require.NoError(t, err)
44+
require.Empty(t, googleThoughtSignatureFromBody(t, roundTripped, 1, 0),
45+
"openai-go drops extra_content during the typed param round-trip")
46+
47+
body, err := (&interceptionBase{
48+
req: &req,
49+
cfg: config.OpenAI{BaseURL: "https://generativelanguage.googleapis.com/v1beta/openai/"},
50+
}).chatCompletionRequestBody()
51+
require.NoError(t, err)
52+
require.Equal(t, googleopenai.DummyThoughtSignature, googleThoughtSignatureFromBody(t, body, 1, 0))
53+
}
54+
55+
func TestGoogleOpenAICompatChatCompletionRequestOptions(t *testing.T) {
56+
t.Parallel()
57+
58+
var req ChatCompletionNewParamsWrapper
59+
require.NoError(t, json.Unmarshal([]byte(`{
60+
"model":"gemini-3.5-flash",
61+
"messages":[
62+
{"role":"user","content":"current turn"},
63+
{
64+
"role":"assistant",
65+
"tool_calls":[{"id":"call-1","function":{"arguments":"{}","name":"list_templates"},"type":"function"}]
66+
}
67+
]
68+
}`), &req))
69+
70+
opts := make([]option.RequestOption, 1)
71+
updated, overrideBody, err := (&interceptionBase{
72+
req: &req,
73+
cfg: config.OpenAI{BaseURL: "https://generativelanguage.googleapis.com/v1beta/openai/"},
74+
}).chatCompletionRequestOptions(opts)
75+
require.NoError(t, err)
76+
require.True(t, overrideBody)
77+
require.Len(t, opts, 1)
78+
require.Len(t, updated, 2)
79+
}
80+
81+
func googleThoughtSignatureFromBody(t *testing.T, body []byte, messageIndex int, toolCallIndex int) string {
82+
t.Helper()
83+
84+
var payload map[string]any
85+
require.NoError(t, json.Unmarshal(body, &payload))
86+
messages, ok := payload["messages"].([]any)
87+
require.True(t, ok)
88+
require.Greater(t, len(messages), messageIndex)
89+
message, ok := messages[messageIndex].(map[string]any)
90+
require.True(t, ok)
91+
toolCalls, ok := message["tool_calls"].([]any)
92+
require.True(t, ok)
93+
require.Greater(t, len(toolCalls), toolCallIndex)
94+
toolCall, ok := toolCalls[toolCallIndex].(map[string]any)
95+
require.True(t, ok)
96+
extraContent, _ := toolCall["extra_content"].(map[string]any)
97+
google, _ := extraContent["google"].(map[string]any)
98+
signature, _ := google["thought_signature"].(string)
99+
return signature
100+
}

‎aibridge/intercept/chatcompletions/streaming.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,9 @@ func (i *StreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Re
186186
// We take control of request body here and pass it to the SDK as a raw byte slice.
187187
// This is because the SDK's serialization applies hidden request options that result in
188188
// unexpected, breaking behavior. See https://github.com/coder/aibridge/pull/164
189-
body, err := json.Marshal(i.req.ChatCompletionNewParams)
189+
// chatCompletionRequestBody also applies provider-specific
190+
// compatibility patches to the exact body sent upstream.
191+
body, err := i.chatCompletionRequestBody()
190192
if err != nil {
191193
return xerrors.Errorf("marshal request body: %w", err)
192194
}

‎coderd/x/chatd/chatprovider/openai_compat_patches.go‎

Lines changed: 4 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -6,16 +6,13 @@ import (
66
"io"
77
"net/http"
88
"strings"
9+
10+
"github.com/coder/coder/v2/internal/googleopenai"
911
)
1012

1113
// OpenAI-compatible providers share an API shape but differ in the exact JSON
1214
// they accept. These patches adjust Fantasy's serialized request body at the
1315
// transport boundary so higher-level generation code can stay provider agnostic.
14-
//
15-
// googleOpenAICompatDummyThoughtSignature is Google's documented last-resort
16-
// bypass for callers that cannot preserve a real Gemini thought signature.
17-
// See https://ai.google.dev/gemini-api/docs/thought-signatures.
18-
const googleOpenAICompatDummyThoughtSignature = "skip_thought_signature_validator"
1916

2017
func withOpenAICompatRequestPatches(
2118
client *http.Client,
@@ -91,8 +88,8 @@ func patchOpenAICompatChatCompletionsBody(body []byte, baseURL string, modelID s
9188
}
9289

9390
changed := rewriteOpenAICompatSingleToolChoice(payload)
94-
if shouldAddGoogleOpenAICompatThoughtSignatures(baseURL, modelID) {
95-
changed = addGoogleOpenAICompatThoughtSignatures(payload) || changed
91+
if googleopenai.ShouldPatchOpenAICompatRequest(baseURL, modelID) {
92+
changed = googleopenai.AddThoughtSignaturesToLatestTurn(payload) || changed
9693
}
9794
if !changed {
9895
return body
@@ -144,93 +141,3 @@ func rewriteOpenAICompatSingleToolChoice(payload map[string]any) bool {
144141
payload["tool_choice"] = "required"
145142
return true
146143
}
147-
148-
// shouldAddGoogleOpenAICompatThoughtSignatures detects direct Gemini OpenAI
149-
// endpoints and Coder AI Bridge Gemini routes. Other gateways, such as Vercel,
150-
// keep their own provider-specific compatibility behavior.
151-
func shouldAddGoogleOpenAICompatThoughtSignatures(baseURL string, modelID string) bool {
152-
parsed, ok := parseProviderBaseURL(baseURL)
153-
if !ok {
154-
return false
155-
}
156-
host := strings.ToLower(parsed.Hostname())
157-
path := strings.ToLower(parsed.EscapedPath())
158-
if host == "generativelanguage.googleapis.com" && strings.Contains(path, "/openai") {
159-
return true
160-
}
161-
return host == "coder-aibridge" && isGeminiModelID(modelID)
162-
}
163-
164-
func isGeminiModelID(modelID string) bool {
165-
modelID = strings.ToLower(strings.TrimSpace(modelID))
166-
return strings.HasPrefix(modelID, "gemini-") || strings.Contains(modelID, "/gemini-")
167-
}
168-
169-
// addGoogleOpenAICompatThoughtSignatures adds a dummy thought signature to the
170-
// first tool call on each assistant tool-call message in the latest user turn.
171-
// Gemini validates tool-call history with thought signatures, but
172-
// OpenAI-compatible serialization can drop the original provider metadata.
173-
func addGoogleOpenAICompatThoughtSignatures(payload map[string]any) bool {
174-
messages, ok := payload["messages"].([]any)
175-
if !ok {
176-
return false
177-
}
178-
179-
currentTurnStart := -1
180-
for i, raw := range messages {
181-
message, ok := raw.(map[string]any)
182-
if !ok {
183-
continue
184-
}
185-
if role, _ := message["role"].(string); role == "user" {
186-
currentTurnStart = i
187-
}
188-
}
189-
190-
if currentTurnStart == -1 {
191-
return false
192-
}
193-
194-
changed := false
195-
for _, raw := range messages[currentTurnStart+1:] {
196-
message, ok := raw.(map[string]any)
197-
if !ok || !isOpenAICompatAssistantRole(message["role"]) {
198-
continue
199-
}
200-
toolCalls, ok := message["tool_calls"].([]any)
201-
if !ok || len(toolCalls) == 0 {
202-
continue
203-
}
204-
firstToolCall, ok := toolCalls[0].(map[string]any)
205-
if !ok {
206-
continue
207-
}
208-
if ensureGoogleOpenAICompatThoughtSignature(firstToolCall) {
209-
changed = true
210-
}
211-
}
212-
return changed
213-
}
214-
215-
func isOpenAICompatAssistantRole(role any) bool {
216-
roleValue, _ := role.(string)
217-
return roleValue == "assistant" || roleValue == "model"
218-
}
219-
220-
func ensureGoogleOpenAICompatThoughtSignature(toolCall map[string]any) bool {
221-
extraContent, _ := toolCall["extra_content"].(map[string]any)
222-
google, _ := extraContent["google"].(map[string]any)
223-
if signature, _ := google["thought_signature"].(string); signature != "" {
224-
return false
225-
}
226-
if extraContent == nil {
227-
extraContent = map[string]any{}
228-
toolCall["extra_content"] = extraContent
229-
}
230-
if google == nil {
231-
google = map[string]any{}
232-
extraContent["google"] = google
233-
}
234-
google["thought_signature"] = googleOpenAICompatDummyThoughtSignature
235-
return true
236-
}

‎coderd/x/chatd/chatprovider/openai_compat_patches_test.go‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,9 @@ import (
1212
"github.com/stretchr/testify/require"
1313

1414
"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
15+
"github.com/coder/coder/v2/internal/googleopenai"
1516
)
1617

17-
const dummyThoughtSignature = "skip_thought_signature_validator"
18-
1918
func TestModelFromConfig_GeminiOpenAICompatThoughtSignatures(t *testing.T) {
2019
t.Parallel()
2120

@@ -26,9 +25,9 @@ func TestModelFromConfig_GeminiOpenAICompatThoughtSignatures(t *testing.T) {
2625
messages := body["messages"].([]any)
2726

2827
require.Empty(t, thoughtSignature(t, messages[1], 0))
29-
require.Equal(t, dummyThoughtSignature, thoughtSignature(t, messages[4], 0))
30-
require.Empty(t, thoughtSignature(t, messages[4], 1))
31-
require.Equal(t, dummyThoughtSignature, thoughtSignature(t, messages[6], 0))
28+
require.Equal(t, googleopenai.DummyThoughtSignature, thoughtSignature(t, messages[4], 0))
29+
require.Equal(t, googleopenai.DummyThoughtSignature, thoughtSignature(t, messages[4], 1))
30+
require.Equal(t, googleopenai.DummyThoughtSignature, thoughtSignature(t, messages[6], 0))
3231
})
3332

3433
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) {
3736
body := generateOpenAICompatRequest(t, "http://coder-aibridge/v1", "gemini-3.5-flash")
3837
messages := body["messages"].([]any)
3938

40-
require.Equal(t, dummyThoughtSignature, thoughtSignature(t, messages[4], 0))
39+
require.Equal(t, googleopenai.DummyThoughtSignature, thoughtSignature(t, messages[4], 0))
4140
})
4241

4342
t.Run("Vercel OpenAI-compatible Gemini route is unchanged", func(t *testing.T) {

0 commit comments

Comments
 (0)