diff --git a/internal/api/openai.go b/internal/api/openai.go index df8a551..a7b1112 100644 --- a/internal/api/openai.go +++ b/internal/api/openai.go @@ -245,11 +245,18 @@ func chatMessageContent(message chatMessage) (aistudio.Content, error) { if !json.Valid(arguments) { return aistudio.Content{}, fmt.Errorf("tool call %q arguments must be JSON", call.Function.Name) } + signature := call.ExtraContent.Google.ThoughtSignature + if signature == "" { + // A client that does not echo extra_content (every ordinary + // OpenAI client, litellm included) still gets its round trip: + // the signature was kept when this call was emitted. + signature = thoughtSignatureFor(call.ID) + } parts = append(parts, aistudio.Part{FunctionCall: &aistudio.FunctionCall{ ID: call.ID, Name: call.Function.Name, Arguments: arguments, - ThoughtSignature: call.ExtraContent.Google.ThoughtSignature, + ThoughtSignature: signature, }}) } return aistudio.Content{Role: role, Parts: parts}, nil @@ -684,6 +691,7 @@ func (s *server) streamChatCompletion(w http.ResponseWriter, r *http.Request, re return nil } call := event.ToolCall + rememberThoughtSignature(call.ID, call.ThoughtSignature) toolCall := map[string]any{ "index": toolIndex, "id": call.ID, @@ -804,6 +812,7 @@ func openAIFinishReason(reason string, hasTools bool) string { } func openAIToolCallOutput(calls []aistudio.FunctionCall) []map[string]any { + rememberThoughtSignatures(calls) output := make([]map[string]any, 0, len(calls)) for _, call := range calls { item := map[string]any{ diff --git a/internal/api/thoughtsig.go b/internal/api/thoughtsig.go new file mode 100644 index 0000000..3c9c32b --- /dev/null +++ b/internal/api/thoughtsig.go @@ -0,0 +1,75 @@ +package api + +// Gemini 3 refuses a functionCall that comes back in history without the +// thought signature it was issued with: +// +// AI Studio GenerateContent 返回 HTTP 400、协议错误码 3: +// [original: beyond::dependency::INVALID_ARGUMENT] +// Function call is missing a thought signature. (qos=CRITICAL_PLUS) +// +// The signature does reach the client - this proxy hands it over as +// extra_content.google.thought_signature on every tool call - but no ordinary +// OpenAI client (OpenAI SDKs, litellm, most agent frameworks) echoes unknown +// fields back. By the next turn of a tool round trip the signature is gone, +// so every agent loop dies on the second request while both tool-call checks +// pass. +// +// Keep it server-side instead, keyed by the tool call id the client saw, and +// put it back when the assistant message comes home without one. A client that +// does echo extra_content still wins: chatMessageContent only consults this +// store when the field is missing, and a call this process never emitted still +// falls through to the skip_thought_signature_validator placeholder in +// internal/aistudio/request.go. + +import ( + "sync" + + "github.com/Mag1cFall/AIStudio2API/internal/aistudio" +) + +// thoughtSignatureLimit bounds the store. Tool call ids are unique per call +// and a client only ever echoes back ids from the recent past, so a plain FIFO +// cap is enough - it keeps a long-lived process from growing one entry per +// tool call forever. +const thoughtSignatureLimit = 1024 + +var thoughtSignatures = struct { + mu sync.Mutex + entries map[string]string + order []string +}{entries: make(map[string]string)} + +func rememberThoughtSignature(id, signature string) { + if id == "" || signature == "" { + return + } + thoughtSignatures.mu.Lock() + defer thoughtSignatures.mu.Unlock() + if _, seen := thoughtSignatures.entries[id]; !seen { + thoughtSignatures.order = append(thoughtSignatures.order, id) + } + thoughtSignatures.entries[id] = signature + for len(thoughtSignatures.order) > thoughtSignatureLimit { + oldest := thoughtSignatures.order[0] + thoughtSignatures.order = thoughtSignatures.order[1:] + delete(thoughtSignatures.entries, oldest) + } +} + +func thoughtSignatureFor(id string) string { + if id == "" { + return "" + } + thoughtSignatures.mu.Lock() + defer thoughtSignatures.mu.Unlock() + return thoughtSignatures.entries[id] +} + +// rememberThoughtSignatures stores every signed call of one model response. +// Calls the model did not sign are skipped, so a parallel response can leave +// some ids unknown - that is the model's choice, not something to invent here. +func rememberThoughtSignatures(calls []aistudio.FunctionCall) { + for _, call := range calls { + rememberThoughtSignature(call.ID, call.ThoughtSignature) + } +} diff --git a/internal/api/thoughtsig_test.go b/internal/api/thoughtsig_test.go new file mode 100644 index 0000000..3dc330c --- /dev/null +++ b/internal/api/thoughtsig_test.go @@ -0,0 +1,43 @@ +package api + +import ( + "encoding/json" + "testing" +) + +func TestChatMessageContentRestoresStoredThoughtSignature(t *testing.T) { + rememberThoughtSignature("call_sig_1", "sig-abc") + + var message chatMessage + raw := `{"role":"assistant","content":null,"tool_calls":[{"id":"call_sig_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"Berlin\"}"}}]}` + if err := json.Unmarshal([]byte(raw), &message); err != nil { + t.Fatal(err) + } + content, err := chatMessageContent(message) + if err != nil { + t.Fatal(err) + } + if got := content.Parts[0].FunctionCall.ThoughtSignature; got != "sig-abc" { + t.Fatalf("stored signature not restored: got %q", got) + } + + message.ToolCalls[0].ExtraContent.Google.ThoughtSignature = "sig-from-client" + content, err = chatMessageContent(message) + if err != nil { + t.Fatal(err) + } + if got := content.Parts[0].FunctionCall.ThoughtSignature; got != "sig-from-client" { + t.Fatalf("client-supplied signature must win: got %q", got) + } +} + +func TestThoughtSignatureStoreIsBounded(t *testing.T) { + for i := 0; i < thoughtSignatureLimit+10; i++ { + rememberThoughtSignature(string(rune('a'+i%26))+string(rune(i)), "s") + } + thoughtSignatures.mu.Lock() + defer thoughtSignatures.mu.Unlock() + if len(thoughtSignatures.entries) > thoughtSignatureLimit { + t.Fatalf("store grew past its limit: %d", len(thoughtSignatures.entries)) + } +}