From 384325d9e3cbb187310bfddd49fda1716f22e625 Mon Sep 17 00:00:00 2001 From: thc1006 <84045975+thc1006@users.noreply.github.com> Date: Mon, 5 Oct 2026 23:38:40 +0800 Subject: [PATCH] fix(adk): reject malformed tool call arguments from OpenAI chat completions The OpenAI chat completions adapter discarded json.Unmarshal errors for tool call arguments, so a call whose arguments were cut off at the token limit still ran, with no arguments. Fail the turn with MALFORMED_FUNCTION_CALL instead, and run none of the tool calls in that response. Blank arguments still mean a call without arguments, unless the response was cut off at the token limit. Signed-off-by: thc1006 <84045975+thc1006@users.noreply.github.com> --- go/adk/pkg/models/openai_adk.go | 58 ++++++++++++++--- go/adk/pkg/models/openai_adk_test.go | 93 ++++++++++++++++++++++++++++ 2 files changed, 143 insertions(+), 8 deletions(-) diff --git a/go/adk/pkg/models/openai_adk.go b/go/adk/pkg/models/openai_adk.go index 327be072c2..00901990c2 100644 --- a/go/adk/pkg/models/openai_adk.go +++ b/go/adk/pkg/models/openai_adk.go @@ -108,6 +108,37 @@ func thoughtSignaturesByToolCallID(contents []*genai.Content) map[string][]byte return thoughtSignatures } +// functionCallArgs decodes tool call arguments. Blank arguments mean none, +// unless the response was cut off at the token limit. +func functionCallArgs(name, id, arguments, finishReason string) (map[string]any, error) { + var args map[string]any + if strings.TrimSpace(arguments) == "" { + if finishReason == openAIFinishLength { + return nil, fmt.Errorf("tool call %q (id %q) was cut off before its arguments", name, id) + } + return args, nil + } + if err := json.Unmarshal([]byte(arguments), &args); err != nil { + return nil, fmt.Errorf("failed to decode arguments of tool call %q (id %q): %w", name, id, err) + } + return args, nil +} + +// rejectMalformedFunctionCalls fails the turn and drops all of its tool calls, +// not only the malformed ones. +func rejectMalformedFunctionCalls(resp *model.LLMResponse, errs []error, finishReason string) { + resp.Content.Parts = slices.DeleteFunc(resp.Content.Parts, func(p *genai.Part) bool { return p.FunctionCall != nil }) + msgs := make([]string, len(errs)) + for i, err := range errs { + msgs[i] = err.Error() + } + resp.ErrorCode = string(genai.FinishReasonMalformedFunctionCall) + resp.ErrorMessage = strings.Join(msgs, "; ") + if finishReason != "" { + resp.ErrorMessage += fmt.Sprintf(" (finish reason %q)", finishReason) + } +} + func newFunctionCallPart(name string, args map[string]any, id string, thoughtSignature []byte) *genai.Part { part := genai.NewPartFromFunctionCall(name, args) if part.FunctionCall != nil { @@ -454,16 +485,18 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet if text != "" { finalParts = append(finalParts, &genai.Part{Text: text}) } + var callErrs []error for _, idx := range indices { tc := toolCallsAcc[idx] argsStr, _ := tc["arguments"].(string) - var args map[string]any - if argsStr != "" { - _ = json.Unmarshal([]byte(argsStr), &args) - } name, _ := tc["name"].(string) id, _ := tc["id"].(string) if name != "" || id != "" { + args, err := functionCallArgs(name, id, argsStr, finishReason) + if err != nil { + callErrs = append(callErrs, err) + continue + } thoughtSignature, _ := tc["thought_signature"].([]byte) p := newFunctionCallPart(name, args, id, thoughtSignature) finalParts = append(finalParts, p) @@ -485,6 +518,9 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet UsageMetadata: usage, Content: &genai.Content{Role: string(genai.RoleModel), Parts: finalParts}, } + if len(callErrs) > 0 { + rejectMalformedFunctionCalls(resp, callErrs, finishReason) + } _ = yield(resp, nil) } @@ -514,11 +550,13 @@ func chatCompletionToLLMResponse(completion *openai.ChatCompletion) *model.LLMRe if msg.Content != "" { parts = append(parts, &genai.Part{Text: msg.Content}) } + var callErrs []error for _, tc := range msg.ToolCalls { if tc.Type == openAIToolTypeFunction && tc.Function.Name != "" { - var args map[string]any - if tc.Function.Arguments != "" { - _ = json.Unmarshal([]byte(tc.Function.Arguments), &args) + args, err := functionCallArgs(tc.Function.Name, tc.ID, tc.Function.Arguments, choice.FinishReason) + if err != nil { + callErrs = append(callErrs, err) + continue } functionToolCall := tc.AsFunction() p := newFunctionCallPart( @@ -538,11 +576,15 @@ func chatCompletionToLLMResponse(completion *openai.ChatCompletion) *model.LLMRe TotalTokenCount: int32(completion.Usage.TotalTokens), } } - return &model.LLMResponse{ + resp := &model.LLMResponse{ Partial: false, TurnComplete: true, FinishReason: openAIFinishReasonToGenai(choice.FinishReason), UsageMetadata: usage, Content: &genai.Content{Role: string(genai.RoleModel), Parts: parts}, } + if len(callErrs) > 0 { + rejectMalformedFunctionCalls(resp, callErrs, choice.FinishReason) + } + return resp } diff --git a/go/adk/pkg/models/openai_adk_test.go b/go/adk/pkg/models/openai_adk_test.go index c5ba95d671..0fa1f3047c 100644 --- a/go/adk/pkg/models/openai_adk_test.go +++ b/go/adk/pkg/models/openai_adk_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/base64" "encoding/json" + "fmt" "io" "net/http" "net/http/httptest" @@ -13,6 +14,8 @@ import ( "github.com/openai/openai-go/v3" "github.com/openai/openai-go/v3/option" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "google.golang.org/adk/v2/model" "google.golang.org/genai" ) @@ -511,3 +514,93 @@ func TestExtractThoughtSignatureFromStreamingToolCallChunk(t *testing.T) { t.Fatalf("thoughtSignature = %q, want %q", string(thoughtSignature), "abc") } } + +func TestOpenAIModelGenerateContentRejectsMalformedToolArguments(t *testing.T) { + tests := []struct { + name string + content string + arguments []string + finishReason string + wantArgs []map[string]any + wantErr string + }{ + {name: "valid arguments", arguments: []string{`{"query":"cpu"}`}, finishReason: "tool_calls", wantArgs: []map[string]any{{"query": "cpu"}}}, + {name: "no arguments", arguments: []string{"", " ", "null"}, finishReason: "tool_calls", wantArgs: []map[string]any{nil, nil, nil}}, + {name: "arguments cut off", content: "Let me check.", arguments: []string{`{"query":"cp`}, finishReason: "length", wantErr: `"call_0"`}, + {name: "cut off before the arguments", arguments: []string{""}, finishReason: "length", wantErr: `"call_0"`}, + {name: "one of parallel calls cut off", arguments: []string{`{"query":"cpu"}`, `{"query":"me`}, finishReason: "length", wantErr: `"call_1"`}, + {name: "arguments not an object", arguments: []string{`["cpu"]`}, finishReason: "tool_calls", wantErr: `"call_0"`}, + } + for _, tt := range tests { + for _, stream := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/stream=%t", tt.name, stream), func(t *testing.T) { + var calls []any + for i, arguments := range tt.arguments { + calls = append(calls, map[string]any{ + "index": i, "id": fmt.Sprintf("call_%d", i), "type": "function", + "function": map[string]any{"name": "search", "arguments": arguments}, + }) + } + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + if !stream { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "id": "chatcmpl-1", "object": "chat.completion", "created": 1, "model": "gpt-4o", + "choices": []any{map[string]any{"index": 0, "finish_reason": tt.finishReason, + "message": map[string]any{"role": "assistant", "content": tt.content, "tool_calls": calls}}}, + }) + return + } + w.Header().Set("Content-Type", "text/event-stream") + for _, choice := range []map[string]any{ + {"index": 0, "delta": map[string]any{"role": "assistant", "content": tt.content}}, + {"index": 0, "delta": map[string]any{"tool_calls": calls}}, + {"index": 0, "delta": map[string]any{}, "finish_reason": tt.finishReason}, + } { + chunk, _ := json.Marshal(map[string]any{ + "id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o", + "choices": []any{choice}, + }) + _, _ = fmt.Fprintf(w, "data: %s\n\n", chunk) + } + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer server.Close() + llm := &OpenAIModel{ + Config: &OpenAIConfig{Model: "gpt-4o"}, + Client: openai.NewClient(option.WithAPIKey("test"), option.WithBaseURL(server.URL), option.WithHTTPClient(server.Client())), + Logger: slog.New(slog.DiscardHandler), + } + request := &model.LLMRequest{Contents: []*genai.Content{{Role: "user", Parts: []*genai.Part{{Text: "check cpu"}}}}} + + var final *model.LLMResponse + for resp, err := range llm.GenerateContent(t.Context(), request, stream) { + require.NoError(t, err) + if !resp.Partial { + final = resp + } + } + require.NotNil(t, final) + var text string + var gotArgs []map[string]any + for _, part := range final.Content.Parts { + text += part.Text + if part.FunctionCall != nil { + gotArgs = append(gotArgs, part.FunctionCall.Args) + } + } + assert.Equal(t, tt.content, text) + if tt.wantErr != "" { + // No tool may run when the model did not finish one of the calls. + assert.Empty(t, gotArgs) + assert.Equal(t, string(genai.FinishReasonMalformedFunctionCall), final.ErrorCode) + assert.Contains(t, final.ErrorMessage, tt.wantErr) + assert.Equal(t, openAIFinishReasonToGenai(tt.finishReason), final.FinishReason) + return + } + assert.Empty(t, final.ErrorCode) + assert.Equal(t, tt.wantArgs, gotArgs) + }) + } + } +}