fix(adk): reject malformed tool call arguments from OpenAI chat completions

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
 }

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)
+			})
+		}
+	}
+}