Preserve Gemini thought_signature across OpenAI-compatible tool-call round trips (#1631)

## Summary

This fixes a loss of Gemini thinking-model state in the
OpenAI-compatible adapter.

When kagent is configured to use Gemini through an OpenAI-compatible
endpoint, the adapter was flattening tool calls and tool results into
standard OpenAI chat messages without preserving Gemini-specific
`extra_content.google.thought_signature`. That caused thinking-capable
Gemini models to fail on the post-tool turn, even though the tool call
itself succeeded.

This change preserves the thought signature in both directions:
- OpenAI-compatible response -> ADK Part
- ADK Content -> OpenAI-compatible follow-up tool/result messages

It also applies the same preservation in the streaming tool-call
accumulation path.

## Why

We traced a production stream where:

1. The first model turn succeeded,
2. The model issued a tool call,
3. The tool executed successfully,
4. The follow-up request failed because the original function call’s
`thought_signature` was no longer present.

The lossy step was in kagent’s Python OpenAI-compatible adapter.

The adapter was reconstructing only standard OpenAI fields like text,
tool calls, and tool results. Gemini thinking models require the opaque
thought signature to survive the round-trip, so dropping it breaks the
second turn.

## What Changed

- Added logic to extract extra_content.google.thought_signature from
OpenAI-compatible tool calls.
- Mapped that value onto ADK Part.thought_signature.
- Preserved Part.thought_signature when converting function-call parts
back into OpenAI-compatible tool_calls.
- Preserved the same signature on the paired tool-result message sent
after tool execution.
- Applied the same preservation in the streaming tool-call aggregation
path.
- Added regression tests for:
    - response conversion preserving thought signatures,
    - outbound message conversion preserving thought signatures,
    - full round-trip preservation across tool call + tool result.

## Scope

This PR does not change public config or CRD APIs. It only fixes
internal adapter behavior for OpenAI-compatible model flows.

## Validation

Ran:
`uv run pytest
python/packages/kagent-adk/tests/unittests/models/test_openai.py`

Result:

`26 passed`

---------

Signed-off-by: adminturneddevops <mlevan1992@gmail.com>
Co-authored-by: Eitan Yarmush <eitan.yarmush@solo.io>

Author: AdminTurnedDevOps

go/adk/pkg/models/openai_adk.go

@@ -14,6 +14,7 @@ import (
 	"github.com/kagent-dev/kagent/go/adk/pkg/telemetry"
 	"github.com/openai/openai-go/v3"
 	"github.com/openai/openai-go/v3/packages/param"
+	"github.com/openai/openai-go/v3/packages/respjson"
 	"github.com/openai/openai-go/v3/shared"
 	"github.com/openai/openai-go/v3/shared/constant"
 	"google.golang.org/adk/model"
@@ -28,6 +29,7 @@ const (
 	openAIFinishLength        = "length"
 	openAIFinishContentFilter = "content_filter"
 	openAIToolTypeFunction    = "function"
+	openAIExtraContentKey     = "extra_content"
 )
 
 // openAIFinishReasonToGenai maps OpenAI finish_reason to genai.FinishReason.
@@ -42,6 +44,82 @@ func openAIFinishReasonToGenai(reason string) genai.FinishReason {
 	}
 }
 
+type openAIThoughtSignatureExtra struct {
+	Google struct {
+		ThoughtSignature string `json:"thought_signature"`
+	} `json:"google"`
+}
+
+func extractThoughtSignatureFromRaw(raw string) []byte {
+	if raw == "" {
+		return nil
+	}
+
+	var extra openAIThoughtSignatureExtra
+	if err := json.Unmarshal([]byte(raw), &extra); err != nil {
+		return nil
+	}
+	if extra.Google.ThoughtSignature == "" {
+		return nil
+	}
+
+	decoded, err := base64.StdEncoding.DecodeString(extra.Google.ThoughtSignature)
+	if err != nil {
+		return nil
+	}
+	return decoded
+}
+
+func extractThoughtSignatureFromExtraFields(extraFields map[string]respjson.Field) []byte {
+	if len(extraFields) == 0 {
+		return nil
+	}
+	field, ok := extraFields[openAIExtraContentKey]
+	if !ok {
+		return nil
+	}
+	return extractThoughtSignatureFromRaw(field.Raw())
+}
+
+func openAIExtraContentForThoughtSignature(thoughtSignature []byte) map[string]any {
+	if len(thoughtSignature) == 0 {
+		return nil
+	}
+
+	return map[string]any{
+		"google": map[string]any{
+			"thought_signature": base64.StdEncoding.EncodeToString(thoughtSignature),
+		},
+	}
+}
+
+func thoughtSignaturesByToolCallID(contents []*genai.Content) map[string][]byte {
+	thoughtSignatures := make(map[string][]byte)
+	for _, content := range contents {
+		if content == nil || content.Parts == nil {
+			continue
+		}
+		for _, part := range content.Parts {
+			if part == nil || part.FunctionCall == nil || len(part.ThoughtSignature) == 0 {
+				continue
+			}
+			thoughtSignatures[part.FunctionCall.ID] = part.ThoughtSignature
+		}
+	}
+	return thoughtSignatures
+}
+
+func newFunctionCallPart(name string, args map[string]any, id string, thoughtSignature []byte) *genai.Part {
+	part := genai.NewPartFromFunctionCall(name, args)
+	if part.FunctionCall != nil {
+		part.FunctionCall.ID = id
+	}
+	if len(thoughtSignature) > 0 {
+		part.ThoughtSignature = thoughtSignature
+	}
+	return part
+}
+
 // Name implements model.LLM.
 func (m *OpenAIModel) Name() string {
 	return m.Config.Model
@@ -126,6 +204,7 @@ func genaiContentsToOpenAIMessages(contents []*genai.Content, config *genai.Gene
 	systemInstruction := strings.TrimSpace(systemBuilder.String())
 
 	functionResponses := make(map[string]*genai.FunctionResponse)
+	thoughtSignatures := thoughtSignaturesByToolCallID(contents)
 	for _, c := range contents {
 		if c == nil || c.Parts == nil {
 			continue
@@ -167,21 +246,35 @@ func genaiContentsToOpenAIMessages(contents []*genai.Content, config *genai.Gene
 			var toolResponseMessages []openai.ChatCompletionMessageParamUnion
 			for _, fc := range functionCalls {
 				argsJSON, _ := json.Marshal(fc.Args)
-				toolCalls = append(toolCalls, openai.ChatCompletionMessageToolCallUnionParam{
-					OfFunction: &openai.ChatCompletionMessageFunctionToolCallParam{
-						ID:   fc.ID,
-						Type: constant.Function(openAIToolTypeFunction),
-						Function: openai.ChatCompletionMessageFunctionToolCallFunctionParam{
-							Name:      fc.Name,
-							Arguments: string(argsJSON),
-						},
+				toolCall := openai.ChatCompletionMessageFunctionToolCallParam{
+					ID:   fc.ID,
+					Type: constant.Function(openAIToolTypeFunction),
+					Function: openai.ChatCompletionMessageFunctionToolCallFunctionParam{
+						Name:      fc.Name,
+						Arguments: string(argsJSON),
 					},
+				}
+				if extraContent := openAIExtraContentForThoughtSignature(thoughtSignatures[fc.ID]); extraContent != nil {
+					toolCall.SetExtraFields(map[string]any{openAIExtraContentKey: extraContent})
+				}
+				toolCalls = append(toolCalls, openai.ChatCompletionMessageToolCallUnionParam{
+					OfFunction: &toolCall,
 				})
 				contentStr := "No response available for this function call."
 				if fr := functionResponses[fc.ID]; fr != nil {
 					contentStr = functionResponseContentString(fr.Response)
 				}
-				toolResponseMessages = append(toolResponseMessages, openai.ToolMessage(contentStr, fc.ID))
+				toolMessage := openai.ChatCompletionToolMessageParam{
+					Content: openai.ChatCompletionToolMessageParamContentUnion{
+						OfString: param.NewOpt(contentStr),
+					},
+					ToolCallID: fc.ID,
+					Role:       constant.Tool("tool"),
+				}
+				if extraContent := openAIExtraContentForThoughtSignature(thoughtSignatures[fc.ID]); extraContent != nil {
+					toolMessage.SetExtraFields(map[string]any{openAIExtraContentKey: extraContent})
+				}
+				toolResponseMessages = append(toolResponseMessages, openai.ChatCompletionMessageParamUnion{OfTool: &toolMessage})
 			}
 			textContent := strings.Join(textParts, "\n")
 			asst := openai.ChatCompletionAssistantMessageParam{
@@ -357,7 +450,7 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet
 		for _, tc := range delta.ToolCalls {
 			idx := tc.Index
 			if toolCallsAcc[idx] == nil {
-				toolCallsAcc[idx] = map[string]any{"id": "", "name": "", "arguments": ""}
+				toolCallsAcc[idx] = map[string]any{"id": "", "name": "", "arguments": "", "thought_signature": []byte(nil)}
 			}
 			if tc.ID != "" {
 				toolCallsAcc[idx]["id"] = tc.ID
@@ -369,6 +462,9 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet
 				prev, _ := toolCallsAcc[idx]["arguments"].(string)
 				toolCallsAcc[idx]["arguments"] = prev + tc.Function.Arguments
 			}
+			if thoughtSignature := extractThoughtSignatureFromExtraFields(tc.JSON.ExtraFields); len(thoughtSignature) > 0 {
+				toolCallsAcc[idx]["thought_signature"] = thoughtSignature
+			}
 		}
 		if choice.FinishReason != "" {
 			finishReason = choice.FinishReason
@@ -404,8 +500,8 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet
 		name, _ := tc["name"].(string)
 		id, _ := tc["id"].(string)
 		if name != "" || id != "" {
-			p := genai.NewPartFromFunctionCall(name, args)
-			p.FunctionCall.ID = id
+			thoughtSignature, _ := tc["thought_signature"].([]byte)
+			p := newFunctionCallPart(name, args, id, thoughtSignature)
 			finalParts = append(finalParts, p)
 		}
 	}
@@ -438,6 +534,12 @@ func runNonStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComp
 		yield(&model.LLMResponse{ErrorCode: "API_ERROR", ErrorMessage: "No choices in response"}, nil)
 		return
 	}
+	resp := chatCompletionToLLMResponse(completion)
+	telemetry.SetLLMResponseAttributes(ctx, resp)
+	yield(resp, nil)
+}
+
+func chatCompletionToLLMResponse(completion *openai.ChatCompletion) *model.LLMResponse {
 	choice := completion.Choices[0]
 	msg := choice.Message
 	nParts := 0
@@ -455,8 +557,13 @@ func runNonStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComp
 			if tc.Function.Arguments != "" {
 				_ = json.Unmarshal([]byte(tc.Function.Arguments), &args)
 			}
-			p := genai.NewPartFromFunctionCall(tc.Function.Name, args)
-			p.FunctionCall.ID = tc.ID
+			functionToolCall := tc.AsFunction()
+			p := newFunctionCallPart(
+				tc.Function.Name,
+				args,
+				tc.ID,
+				extractThoughtSignatureFromExtraFields(functionToolCall.JSON.ExtraFields),
+			)
 			parts = append(parts, p)
 		}
 	}
@@ -467,13 +574,11 @@ func runNonStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComp
 			CandidatesTokenCount: int32(completion.Usage.CompletionTokens),
 		}
 	}
-	resp := &model.LLMResponse{
+	return &model.LLMResponse{
 		Partial:       false,
 		TurnComplete:  true,
 		FinishReason:  openAIFinishReasonToGenai(choice.FinishReason),
 		UsageMetadata: usage,
 		Content:       &genai.Content{Role: string(genai.RoleModel), Parts: parts},
 	}
-	telemetry.SetLLMResponseAttributes(ctx, resp)
-	yield(resp, nil)
 }

go/adk/pkg/models/openai_adk_test.go

@@ -1,6 +1,8 @@
 package models
 
 import (
+	"encoding/base64"
+	"encoding/json"
 	"testing"
 
 	"github.com/openai/openai-go/v3"
@@ -246,3 +248,163 @@ func TestApplyOpenAIConfig(t *testing.T) {
 		}
 	})
 }
+
+func TestGenaiContentsToOpenAIMessages_PreservesThoughtSignatureOnToolCallAndToolResult(t *testing.T) {
+	thoughtSignature := []byte("abc")
+
+	functionCall := genai.NewPartFromFunctionCall("add", map[string]any{"a": 2, "b": 2})
+	functionCall.FunctionCall.ID = "call_1"
+	functionCall.ThoughtSignature = thoughtSignature
+
+	functionResponse := genai.NewPartFromFunctionResponse("add", map[string]any{"result": "4"})
+	functionResponse.FunctionResponse.ID = "call_1"
+
+	messages, _ := genaiContentsToOpenAIMessages([]*genai.Content{
+		{
+			Role:  string(genai.RoleModel),
+			Parts: []*genai.Part{functionCall},
+		},
+		{
+			Role:  string(genai.RoleUser),
+			Parts: []*genai.Part{functionResponse},
+		},
+	}, nil)
+
+	if len(messages) != 2 {
+		t.Fatalf("len(messages) = %d, want 2", len(messages))
+	}
+
+	assistantJSON, err := json.Marshal(messages[0].OfAssistant)
+	if err != nil {
+		t.Fatalf("json.Marshal(assistant) error = %v", err)
+	}
+	toolJSON, err := json.Marshal(messages[1].OfTool)
+	if err != nil {
+		t.Fatalf("json.Marshal(tool) error = %v", err)
+	}
+
+	want := base64.StdEncoding.EncodeToString(thoughtSignature)
+	assertThoughtSignature := func(name string, payload []byte) {
+		t.Helper()
+		var obj map[string]any
+		if err := json.Unmarshal(payload, &obj); err != nil {
+			t.Fatalf("%s json.Unmarshal error = %v", name, err)
+		}
+		extra, ok := obj["extra_content"].(map[string]any)
+		if !ok {
+			t.Fatalf("%s missing extra_content: %s", name, string(payload))
+		}
+		googleExtra, ok := extra["google"].(map[string]any)
+		if !ok {
+			t.Fatalf("%s missing google extra content: %s", name, string(payload))
+		}
+		if got, _ := googleExtra["thought_signature"].(string); got != want {
+			t.Fatalf("%s thought_signature = %q, want %q", name, got, want)
+		}
+	}
+
+	var assistantObj map[string]any
+	if err := json.Unmarshal(assistantJSON, &assistantObj); err != nil {
+		t.Fatalf("assistant json.Unmarshal error = %v", err)
+	}
+	toolCalls, ok := assistantObj["tool_calls"].([]any)
+	if !ok || len(toolCalls) != 1 {
+		t.Fatalf("assistant tool_calls = %#v, want 1 tool call", assistantObj["tool_calls"])
+	}
+	firstToolCall, ok := toolCalls[0].(map[string]any)
+	if !ok {
+		t.Fatalf("assistant tool call = %#v, want object", toolCalls[0])
+	}
+	firstToolCallJSON, err := json.Marshal(firstToolCall)
+	if err != nil {
+		t.Fatalf("json.Marshal(firstToolCall) error = %v", err)
+	}
+
+	assertThoughtSignature("assistant tool_call", firstToolCallJSON)
+	assertThoughtSignature("tool message", toolJSON)
+}
+
+func TestChatCompletionToLLMResponse_PreservesThoughtSignature(t *testing.T) {
+	raw := []byte(`{
+		"id":"chatcmpl-1",
+		"object":"chat.completion",
+		"created":123,
+		"model":"gemini-2.5-flash",
+		"choices":[{
+			"index":0,
+			"finish_reason":"tool_calls",
+			"message":{
+				"role":"assistant",
+				"tool_calls":[{
+					"id":"call_1",
+					"type":"function",
+					"function":{
+						"name":"add",
+						"arguments":"{\"a\":2,\"b\":2}"
+					},
+					"extra_content":{
+						"google":{
+							"thought_signature":"YWJj"
+						}
+					}
+				}]
+			}
+		}],
+		"usage":{
+			"prompt_tokens":3,
+			"completion_tokens":4,
+			"total_tokens":7
+		}
+	}`)
+
+	var completion openai.ChatCompletion
+	if err := json.Unmarshal(raw, &completion); err != nil {
+		t.Fatalf("json.Unmarshal(ChatCompletion) error = %v", err)
+	}
+
+	resp := chatCompletionToLLMResponse(&completion)
+	if resp.Content == nil || len(resp.Content.Parts) != 1 {
+		t.Fatalf("response parts = %#v, want 1 function-call part", resp.Content)
+	}
+
+	part := resp.Content.Parts[0]
+	if part.FunctionCall == nil {
+		t.Fatalf("part.FunctionCall = nil, want function call")
+	}
+	if part.FunctionCall.Name != "add" {
+		t.Fatalf("part.FunctionCall.Name = %q, want %q", part.FunctionCall.Name, "add")
+	}
+	if string(part.ThoughtSignature) != "abc" {
+		t.Fatalf("part.ThoughtSignature = %q, want %q", string(part.ThoughtSignature), "abc")
+	}
+	if resp.UsageMetadata == nil || resp.UsageMetadata.PromptTokenCount != 3 || resp.UsageMetadata.CandidatesTokenCount != 4 {
+		t.Fatalf("usage metadata = %#v, want prompt=3 completion=4", resp.UsageMetadata)
+	}
+}
+
+func TestExtractThoughtSignatureFromStreamingToolCallChunk(t *testing.T) {
+	raw := []byte(`{
+		"index":0,
+		"id":"call_1",
+		"type":"function",
+		"function":{
+			"name":"add",
+			"arguments":"{\"a\":2"
+		},
+		"extra_content":{
+			"google":{
+				"thought_signature":"YWJj"
+			}
+		}
+	}`)
+
+	var toolCall openai.ChatCompletionChunkChoiceDeltaToolCall
+	if err := json.Unmarshal(raw, &toolCall); err != nil {
+		t.Fatalf("json.Unmarshal(ChatCompletionChunkChoiceDeltaToolCall) error = %v", err)
+	}
+
+	thoughtSignature := extractThoughtSignatureFromExtraFields(toolCall.JSON.ExtraFields)
+	if string(thoughtSignature) != "abc" {
+		t.Fatalf("thoughtSignature = %q, want %q", string(thoughtSignature), "abc")
+	}
+}