Merge branch 'main' into feature/model-gemini

Author: pherzheyu

agent/a2aagent/a2a_converter.go

@@ -274,9 +274,13 @@ func processDataPart(part protocol.Part) (content string, toolCall *model.ToolCa
 		return
 	}
 
+	// Try both standard "type" and ADK-compatible "adk_type" metadata keys
 	typeVal, hasType := d.Metadata[ia2a.DataPartMetadataTypeKey]
 	if !hasType {
-		return
+		typeVal, hasType = d.Metadata[ia2a.GetADKMetadataKey(ia2a.DataPartMetadataTypeKey)]
+		if !hasType {
+			return
+		}
 	}
 
 	// Convert typeVal to string for comparison

agent/a2aagent/a2a_converter_test.go

@@ -1250,6 +1250,177 @@ func TestConvertDataPartToToolResponse(t *testing.T) {
 	}
 }
 
+// TestProcessDataPart_ADKMetadataKey tests handling of ADK-compatible metadata keys
+func TestProcessDataPart_ADKMetadataKey(t *testing.T) {
+	type testCase struct {
+		name         string
+		dataPart     protocol.Part
+		validateFunc func(t *testing.T, content string, toolCall *model.ToolCall, toolResp *toolResponseInfo)
+	}
+
+	tests := []testCase{
+		{
+			name: "function call with adk_type metadata",
+			dataPart: &protocol.DataPart{
+				Data: map[string]any{
+					"id":   "call-adk-1",
+					"type": "function",
+					"name": "get_weather",
+					"args": `{"location": "Beijing"}`,
+				},
+				Metadata: map[string]any{
+					"adk_type": "function_call", // Using ADK-compatible key
+				},
+			},
+			validateFunc: func(t *testing.T, content string, toolCall *model.ToolCall, toolResp *toolResponseInfo) {
+				if toolCall == nil {
+					t.Fatal("expected tool call, got nil")
+				}
+				if toolCall.ID != "call-adk-1" {
+					t.Errorf("expected ID 'call-adk-1', got %s", toolCall.ID)
+				}
+				if toolCall.Function.Name != "get_weather" {
+					t.Errorf("expected name 'get_weather', got %s", toolCall.Function.Name)
+				}
+				if toolResp != nil {
+					t.Errorf("expected nil tool response, got %v", toolResp)
+				}
+			},
+		},
+		{
+			name: "function response with adk_type metadata",
+			dataPart: &protocol.DataPart{
+				Data: map[string]any{
+					"id":       "call-adk-2",
+					"name":     "get_weather",
+					"response": "Beijing: Sunny, 20°C",
+				},
+				Metadata: map[string]any{
+					"adk_type": "function_response", // Using ADK-compatible key
+				},
+			},
+			validateFunc: func(t *testing.T, content string, toolCall *model.ToolCall, toolResp *toolResponseInfo) {
+				if content != "Beijing: Sunny, 20°C" {
+					t.Errorf("expected content 'Beijing: Sunny, 20°C', got %s", content)
+				}
+				if toolResp == nil {
+					t.Fatal("expected tool response info, got nil")
+				}
+				if toolResp.id != "call-adk-2" {
+					t.Errorf("expected tool response ID 'call-adk-2', got %s", toolResp.id)
+				}
+				if toolResp.name != "get_weather" {
+					t.Errorf("expected tool response name 'get_weather', got %s", toolResp.name)
+				}
+				if toolCall != nil {
+					t.Errorf("expected nil tool call, got %v", toolCall)
+				}
+			},
+		},
+		{
+			name: "no type metadata",
+			dataPart: &protocol.DataPart{
+				Data: map[string]any{
+					"id":   "call-no-type",
+					"name": "some_function",
+				},
+				Metadata: map[string]any{
+					"other_key": "other_value",
+				},
+			},
+			validateFunc: func(t *testing.T, content string, toolCall *model.ToolCall, toolResp *toolResponseInfo) {
+				if content != "" {
+					t.Errorf("expected empty content, got %s", content)
+				}
+				if toolCall != nil {
+					t.Errorf("expected nil tool call, got %v", toolCall)
+				}
+				if toolResp != nil {
+					t.Errorf("expected nil tool response, got %v", toolResp)
+				}
+			},
+		},
+		{
+			name: "nil metadata",
+			dataPart: &protocol.DataPart{
+				Data: map[string]any{
+					"id":   "call-nil-meta",
+					"name": "some_function",
+				},
+				Metadata: nil,
+			},
+			validateFunc: func(t *testing.T, content string, toolCall *model.ToolCall, toolResp *toolResponseInfo) {
+				if content != "" {
+					t.Errorf("expected empty content, got %s", content)
+				}
+				if toolCall != nil {
+					t.Errorf("expected nil tool call, got %v", toolCall)
+				}
+				if toolResp != nil {
+					t.Errorf("expected nil tool response, got %v", toolResp)
+				}
+			},
+		},
+		{
+			name: "non-string type value",
+			dataPart: &protocol.DataPart{
+				Data: map[string]any{
+					"id":   "call-bad-type",
+					"name": "some_function",
+				},
+				Metadata: map[string]any{
+					"type": 12345, // Non-string type value
+				},
+			},
+			validateFunc: func(t *testing.T, content string, toolCall *model.ToolCall, toolResp *toolResponseInfo) {
+				if content != "" {
+					t.Errorf("expected empty content, got %s", content)
+				}
+				if toolCall != nil {
+					t.Errorf("expected nil tool call, got %v", toolCall)
+				}
+				if toolResp != nil {
+					t.Errorf("expected nil tool response, got %v", toolResp)
+				}
+			},
+		},
+		{
+			name: "standard type key takes precedence over adk_type",
+			dataPart: &protocol.DataPart{
+				Data: map[string]any{
+					"id":   "call-precedence",
+					"type": "function",
+					"name": "test_func",
+					"args": `{"x": 1}`,
+				},
+				Metadata: map[string]any{
+					"type":     "function_call",     // Standard key
+					"adk_type": "function_response", // ADK key (should be ignored)
+				},
+			},
+			validateFunc: func(t *testing.T, content string, toolCall *model.ToolCall, toolResp *toolResponseInfo) {
+				// Should process as function_call, not function_response
+				if toolCall == nil {
+					t.Fatal("expected tool call, got nil")
+				}
+				if toolCall.Function.Name != "test_func" {
+					t.Errorf("expected name 'test_func', got %s", toolCall.Function.Name)
+				}
+				if toolResp != nil {
+					t.Errorf("expected nil tool response (standard key should take precedence), got %v", toolResp)
+				}
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			content, toolCall, toolResp := processDataPart(tc.dataPart)
+			tc.validateFunc(t, content, toolCall, toolResp)
+		})
+	}
+}
+
 // Helper function to create string pointer
 func stringPtr(s string) *string {
 	return &s

agent/invocation.go

@@ -558,6 +558,35 @@ func (inv *Invocation) GetState(key string) (any, bool) {
 	return value, ok
 }
 
+// GetStateValue retrieves a typed value from the invocation state.
+//
+// Returns the typed value and true if the key exists and the type matches,
+// or the zero value and false otherwise.
+//
+// Example:
+//
+//	if startTime, ok := GetStateValue[time.Time](inv, "agent:start_time"); ok {
+//	    duration := time.Since(startTime)
+//	}
+//	if requestID, ok := GetStateValue[string](inv, "middleware:request_id"); ok {
+//	    log.Printf("Request ID: %s", requestID)
+//	}
+func GetStateValue[T any](inv *Invocation, key string) (T, bool) {
+	var zero T
+	if inv == nil {
+		return zero, false
+	}
+	val, ok := inv.GetState(key)
+	if !ok {
+		return zero, false
+	}
+	typedVal, ok := val.(T)
+	if !ok {
+		return zero, false
+	}
+	return typedVal, true
+}
+
 // GetOrCreateTimingInfo gets or creates timing info for this invocation.
 // Only the first LLM call will create and populate timing info; subsequent calls reuse it.
 // This ensures timing metrics only reflect the first LLM call in scenarios with multiple calls (e.g., tool calls).

agent/invocation_state_test.go

@@ -323,3 +323,124 @@ func TestInvocation_State_ComplexStruct(t *testing.T) {
 	assert.Equal(t, data.ID, retrieved.ID)
 	assert.Equal(t, data.Metadata, retrieved.Metadata)
 }
+
+func TestGetStateValue(t *testing.T) {
+	t.Run("key not found", func(t *testing.T) {
+		inv := NewInvocation()
+		val, ok := GetStateValue[string](inv, "nonexistent")
+		assert.False(t, ok)
+		assert.Equal(t, "", val)
+	})
+
+	t.Run("matching type", func(t *testing.T) {
+		inv := NewInvocation()
+		inv.SetState("agent:string", "hello")
+		inv.SetState("agent:int", 42)
+		inv.SetState("agent:float", 3.14)
+		inv.SetState("agent:bool", true)
+		inv.SetState("agent:time", time.Now())
+
+		// Test string.
+		strVal, ok := GetStateValue[string](inv, "agent:string")
+		assert.True(t, ok)
+		assert.Equal(t, "hello", strVal)
+
+		// Test int.
+		intVal, ok := GetStateValue[int](inv, "agent:int")
+		assert.True(t, ok)
+		assert.Equal(t, 42, intVal)
+
+		// Test float64.
+		floatVal, ok := GetStateValue[float64](inv, "agent:float")
+		assert.True(t, ok)
+		assert.Equal(t, 3.14, floatVal)
+
+		// Test bool.
+		boolVal, ok := GetStateValue[bool](inv, "agent:bool")
+		assert.True(t, ok)
+		assert.Equal(t, true, boolVal)
+
+		// Test time.Time.
+		timeVal, ok := GetStateValue[time.Time](inv, "agent:time")
+		assert.True(t, ok)
+		assert.IsType(t, time.Time{}, timeVal)
+	})
+
+	t.Run("type mismatch", func(t *testing.T) {
+		inv := NewInvocation()
+		inv.SetState("agent:value", "hello")
+
+		// Try to get as int when it's actually string.
+		intVal, ok := GetStateValue[int](inv, "agent:value")
+		assert.False(t, ok)
+		assert.Equal(t, 0, intVal)
+
+		// Try to get as string when it's actually int.
+		inv.SetState("agent:number", 42)
+		strVal, ok := GetStateValue[string](inv, "agent:number")
+		assert.False(t, ok)
+		assert.Equal(t, "", strVal)
+	})
+
+	t.Run("nil invocation", func(t *testing.T) {
+		var inv *Invocation
+		val, ok := GetStateValue[string](inv, "key")
+		assert.False(t, ok)
+		assert.Equal(t, "", val)
+	})
+
+	t.Run("complex struct type", func(t *testing.T) {
+		type CustomData struct {
+			ID        string
+			Timestamp time.Time
+			Metadata  map[string]string
+		}
+
+		inv := NewInvocation()
+		data := CustomData{
+			ID:        "test-123",
+			Timestamp: time.Now(),
+			Metadata: map[string]string{
+				"key1": "value1",
+				"key2": "value2",
+			},
+		}
+		inv.SetState("agent:custom_data", data)
+
+		retrieved, ok := GetStateValue[CustomData](inv, "agent:custom_data")
+		require.True(t, ok)
+		assert.Equal(t, data.ID, retrieved.ID)
+		assert.Equal(t, data.Metadata, retrieved.Metadata)
+	})
+
+	t.Run("pointer type", func(t *testing.T) {
+		inv := NewInvocation()
+		str := "hello"
+		inv.SetState("agent:ptr", &str)
+
+		ptrVal, ok := GetStateValue[*string](inv, "agent:ptr")
+		assert.True(t, ok)
+		require.NotNil(t, ptrVal)
+		assert.Equal(t, "hello", *ptrVal)
+	})
+
+	t.Run("slice type", func(t *testing.T) {
+		inv := NewInvocation()
+		slice := []int{1, 2, 3}
+		inv.SetState("agent:slice", slice)
+
+		sliceVal, ok := GetStateValue[[]int](inv, "agent:slice")
+		assert.True(t, ok)
+		assert.Equal(t, []int{1, 2, 3}, sliceVal)
+	})
+
+	t.Run("map type", func(t *testing.T) {
+		inv := NewInvocation()
+		m := map[string]int{"a": 1, "b": 2}
+		inv.SetState("agent:map", m)
+
+		mapVal, ok := GetStateValue[map[string]int](inv, "agent:map")
+		assert.True(t, ok)
+		assert.Equal(t, map[string]int{"a": 1, "b": 2}, mapVal)
+	})
+}

agent/invocationcontext.go

@@ -32,6 +32,29 @@ func InvocationFromContext(ctx context.Context) (*Invocation, bool) {
 	return invocation, ok
 }
 
+// GetStateValueFromContext retrieves a typed value from the invocation state
+// stored in the context.
+//
+// Returns the typed value and true if the invocation exists, the key exists,
+// and the type matches, or the zero value and false otherwise.
+//
+// Example:
+//
+//	if startTime, ok := GetStateValueFromContext[time.Time](ctx, "agent:start_time"); ok {
+//	    duration := time.Since(startTime)
+//	}
+//	if requestID, ok := GetStateValueFromContext[string](ctx, "middleware:request_id"); ok {
+//	    log.Printf("Request ID: %s", requestID)
+//	}
+func GetStateValueFromContext[T any](ctx context.Context, key string) (T, bool) {
+	var zero T
+	inv, ok := InvocationFromContext(ctx)
+	if !ok {
+		return zero, false
+	}
+	return GetStateValue[T](inv, key)
+}
+
 // CheckContextCancelled check context cancelled
 func CheckContextCancelled(ctx context.Context) error {
 	select {