Patch thought signatures for thinking

internal/llminternal/functions.go

@@ -43,20 +43,33 @@ func generateRequestConfirmationEvent(
 
 	parts := []*genai.Part{}
 	longRunningToolIDs := []string{}
-	functionCalls := make(map[string]*genai.FunctionCall, len(functionCallEvent.Content.Parts))
-	for _, call := range utils.FunctionCalls(functionCallEvent.Content) {
-		functionCalls[call.ID] = call
+
+	// Index the original Parts by function call ID so we can carry over
+	// ThoughtSignature. Gemini thinking models require every model-role
+	// function call Part to include its thought signature; without it the
+	// API returns INVALID_ARGUMENT.
+	type originalCall struct {
+		call             *genai.FunctionCall
+		thoughtSignature []byte
+	}
+	originals := make(map[string]originalCall, len(functionCallEvent.Content.Parts))
+	for _, p := range functionCallEvent.Content.Parts {
+		if p.FunctionCall != nil {
+			originals[p.FunctionCall.ID] = originalCall{
+				call:             p.FunctionCall,
+				thoughtSignature: p.ThoughtSignature,
+			}
+		}
 	}
 
 	for funcID, confirmation := range functionResponseEvent.Actions.RequestedToolConfirmations {
-		originalFunctionCall, ok := functionCalls[funcID]
-		if !ok || originalFunctionCall == nil {
+		orig, ok := originals[funcID]
+		if !ok || orig.call == nil {
 			continue
 		}
 
-		// Prepare arguments for the adk_request_confirmation call
 		args := map[string]any{
-			"originalFunctionCall": originalFunctionCall,
+			"originalFunctionCall": orig.call,
 			"toolConfirmation":     confirmation,
 		}
 
@@ -67,7 +80,8 @@ func generateRequestConfirmationEvent(
 		}
 
 		parts = append(parts, &genai.Part{
-			FunctionCall: requestConfirmationFC,
+			FunctionCall:     requestConfirmationFC,
+			ThoughtSignature: orig.thoughtSignature,
 		})
 		longRunningToolIDs = append(longRunningToolIDs, requestConfirmationFC.ID)
 	}

internal/llminternal/functions_test.go

@@ -194,6 +194,104 @@ func TestGenerateRequestConfirmationEvent(t *testing.T) {
 	}
 }
 
+// TestGenerateRequestConfirmationEventPreservesThoughtSignature verifies that
+// ThoughtSignature from the original function call Part is carried over to the
+// adk_request_confirmation Part. Gemini thinking models require every model-role
+// function call Part to include a thought_signature — without it the API returns
+// INVALID_ARGUMENT.
+func TestGenerateRequestConfirmationEventPreservesThoughtSignature(t *testing.T) {
+	sig := []byte("opaque-thought-signature")
+
+	ctx := &mockInvocationContext{
+		invocationID: "inv_1",
+		agentName:    "agent_1",
+		branch:       "main",
+	}
+
+	functionCallEvent := &session.Event{
+		LLMResponse: model.LLMResponse{
+			Content: &genai.Content{
+				Parts: []*genai.Part{
+					{
+						FunctionCall: &genai.FunctionCall{
+							ID:   "call_1",
+							Name: "test_tool",
+							Args: map[string]any{"arg": "val"},
+						},
+						ThoughtSignature: sig,
+					},
+				},
+			},
+		},
+	}
+
+	functionResponseEvent := &session.Event{
+		Actions: session.EventActions{
+			RequestedToolConfirmations: map[string]toolconfirmation.ToolConfirmation{
+				"call_1": {Hint: "confirm?"},
+			},
+		},
+	}
+
+	got := generateRequestConfirmationEvent(ctx, functionCallEvent, functionResponseEvent)
+	if got == nil {
+		t.Fatal("expected non-nil event")
+	}
+	if len(got.Content.Parts) != 1 {
+		t.Fatalf("expected 1 part, got %d", len(got.Content.Parts))
+	}
+
+	part := got.Content.Parts[0]
+	if part.FunctionCall == nil || part.FunctionCall.Name != toolconfirmation.FunctionCallName {
+		t.Fatal("expected adk_request_confirmation function call")
+	}
+	if string(part.ThoughtSignature) != string(sig) {
+		t.Errorf("ThoughtSignature not preserved: got %q, want %q", part.ThoughtSignature, sig)
+	}
+}
+
+// TestGenerateRequestConfirmationEventNoThoughtSignature verifies the function
+// works normally when the original Part has no ThoughtSignature.
+func TestGenerateRequestConfirmationEventNoThoughtSignature(t *testing.T) {
+	ctx := &mockInvocationContext{
+		invocationID: "inv_1",
+		agentName:    "agent_1",
+	}
+
+	functionCallEvent := &session.Event{
+		LLMResponse: model.LLMResponse{
+			Content: &genai.Content{
+				Parts: []*genai.Part{
+					{
+						FunctionCall: &genai.FunctionCall{
+							ID:   "call_1",
+							Name: "test_tool",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	functionResponseEvent := &session.Event{
+		Actions: session.EventActions{
+			RequestedToolConfirmations: map[string]toolconfirmation.ToolConfirmation{
+				"call_1": {Hint: "confirm?"},
+			},
+		},
+	}
+
+	got := generateRequestConfirmationEvent(ctx, functionCallEvent, functionResponseEvent)
+	if got == nil {
+		t.Fatal("expected non-nil event")
+	}
+
+	part := got.Content.Parts[0]
+	if len(part.ThoughtSignature) != 0 {
+		t.Errorf("expected no ThoughtSignature, got %q", part.ThoughtSignature)
+	}
+}
+
 // TestGenerateRequestConfirmationEventHasID verifies that the event returned
 // by generateRequestConfirmationEvent always has a non-empty ID.
 //

session/vertexai/vertexai_client.go

@@ -441,14 +441,16 @@ func aiplatformToGenaiContent(rpcResp *aiplatformpb.SessionEvent) *genai.Content
 					Data:     v.InlineData.Data,
 				}
 			case *aiplatformpb.Part_FunctionCall:
-				argsMap := v.FunctionCall.Args.AsMap() // Converts *structpb.Struct -> map[string]any
+				argsMap := v.FunctionCall.Args.AsMap()
 				part.FunctionCall = &genai.FunctionCall{
+					ID:   v.FunctionCall.Id,
 					Name: v.FunctionCall.Name,
 					Args: argsMap,
 				}
 			case *aiplatformpb.Part_FunctionResponse:
-				responseMap := v.FunctionResponse.Response.AsMap() // Converts *structpb.Struct -> map[string]any
+				responseMap := v.FunctionResponse.Response.AsMap()
 				part.FunctionResponse = &genai.FunctionResponse{
+					ID:       v.FunctionResponse.Id,
 					Name:     v.FunctionResponse.Name,
 					Response: responseMap,
 				}

session/vertexai/vertexai_test.go

@@ -16,8 +16,106 @@ package vertexai
 
 import (
 	"testing"
+
+	aiplatformpb "cloud.google.com/go/aiplatform/apiv1beta1/aiplatformpb"
+	"google.golang.org/genai"
+	"google.golang.org/protobuf/types/known/structpb"
+
+	"google.golang.org/adk/model"
+	"google.golang.org/adk/session"
 )
 
+func TestContentRoundTrip(t *testing.T) {
+	sig := []byte("opaque-thought-signature-bytes")
+
+	event := &session.Event{
+		LLMResponse: model.LLMResponse{
+			Content: &genai.Content{
+				Role: genai.RoleModel,
+				Parts: []*genai.Part{
+					{
+						Text:             "thinking...",
+						Thought:          true,
+						ThoughtSignature: sig,
+					},
+					{
+						FunctionCall: &genai.FunctionCall{
+							ID:   "call-123",
+							Name: "my_tool",
+							Args: map[string]any{"key": "val"},
+						},
+						ThoughtSignature: sig,
+					},
+				},
+			},
+		},
+	}
+
+	pb, err := createAiplatformpbContent(event)
+	if err != nil {
+		t.Fatalf("createAiplatformpbContent: %v", err)
+	}
+
+	rpcResp := &aiplatformpb.SessionEvent{Content: pb}
+	got := aiplatformToGenaiContent(rpcResp)
+
+	if len(got.Parts) != 2 {
+		t.Fatalf("expected 2 parts, got %d", len(got.Parts))
+	}
+
+	// Thought text part.
+	p0 := got.Parts[0]
+	if p0.Text != "thinking..." {
+		t.Errorf("part[0].Text = %q, want %q", p0.Text, "thinking...")
+	}
+	if !p0.Thought {
+		t.Error("part[0].Thought should be true")
+	}
+	if string(p0.ThoughtSignature) != string(sig) {
+		t.Errorf("part[0].ThoughtSignature lost in round-trip")
+	}
+
+	// FunctionCall part.
+	p1 := got.Parts[1]
+	if p1.FunctionCall == nil {
+		t.Fatal("part[1].FunctionCall is nil")
+	}
+	if p1.FunctionCall.ID != "call-123" {
+		t.Errorf("part[1].FunctionCall.ID = %q, want %q", p1.FunctionCall.ID, "call-123")
+	}
+	if p1.FunctionCall.Name != "my_tool" {
+		t.Errorf("part[1].FunctionCall.Name = %q, want %q", p1.FunctionCall.Name, "my_tool")
+	}
+	if string(p1.ThoughtSignature) != string(sig) {
+		t.Errorf("part[1].ThoughtSignature lost in round-trip")
+	}
+}
+
+func TestContentRoundTripFunctionResponseID(t *testing.T) {
+	args, _ := structpb.NewStruct(map[string]any{"result": "ok"})
+	rpcResp := &aiplatformpb.SessionEvent{
+		Content: &aiplatformpb.Content{
+			Role: "user",
+			Parts: []*aiplatformpb.Part{
+				{
+					Data: &aiplatformpb.Part_FunctionResponse{
+						FunctionResponse: &aiplatformpb.FunctionResponse{
+							Id:       "resp-456",
+							Name:     "my_tool",
+							Response: args,
+						},
+					},
+				},
+			},
+		},
+	}
+
+	got := aiplatformToGenaiContent(rpcResp)
+	if got.Parts[0].FunctionResponse.ID != "resp-456" {
+		t.Errorf("FunctionResponse.ID = %q, want %q", got.Parts[0].FunctionResponse.ID, "resp-456")
+	}
+}
+
 func TestGetReasoningEngineID(t *testing.T) {
 	tests := []struct {
 		name             string