{agent, graph}: implement generic getter for typed state retrieval (#752)

- Added GetStateValue function to retrieve typed values from the
invocation state.
- Introduced GetStateValueFromContext to fetch typed values from the
invocation stored in the context.
- Introduced state_test.go to validate the behavior of GetStateValue for
various scenarios, including key not found, nil state, type matching,
type mismatch, and complex types.

Author: Rememorio

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 {

agent/invocationcontext_test.go

@@ -114,3 +114,85 @@ func TestInvocationFromContext(t *testing.T) {
 		})
 	}
 }
+
+func TestGetStateValueFromContext(t *testing.T) {
+	t.Run("context without invocation", func(t *testing.T) {
+		ctx := context.Background()
+		val, ok := GetStateValueFromContext[string](ctx, "key")
+		assert.False(t, ok)
+		assert.Equal(t, "", val)
+	})
+
+	t.Run("context with invocation but key not found", func(t *testing.T) {
+		inv := NewInvocation()
+		ctx := NewInvocationContext(context.Background(), inv)
+		val, ok := GetStateValueFromContext[string](ctx, "nonexistent")
+		assert.False(t, ok)
+		assert.Equal(t, "", val)
+	})
+
+	t.Run("context with invocation and matching type", func(t *testing.T) {
+		inv := NewInvocation()
+		inv.SetState("agent:string", "hello")
+		inv.SetState("agent:int", 42)
+		inv.SetState("agent:time", time.Now())
+		ctx := NewInvocationContext(context.Background(), inv)
+
+		// Test string value.
+		strVal, ok := GetStateValueFromContext[string](ctx, "agent:string")
+		assert.True(t, ok)
+		assert.Equal(t, "hello", strVal)
+
+		// Test int value.
+		intVal, ok := GetStateValueFromContext[int](ctx, "agent:int")
+		assert.True(t, ok)
+		assert.Equal(t, 42, intVal)
+
+		// Test time.Time value.
+		timeVal, ok := GetStateValueFromContext[time.Time](ctx, "agent:time")
+		assert.True(t, ok)
+		assert.IsType(t, time.Time{}, timeVal)
+	})
+
+	t.Run("context with invocation but type mismatch", func(t *testing.T) {
+		inv := NewInvocation()
+		inv.SetState("agent:value", "hello")
+		ctx := NewInvocationContext(context.Background(), inv)
+
+		// Try to get as int when it's actually string.
+		intVal, ok := GetStateValueFromContext[int](ctx, "agent:value")
+		assert.False(t, ok)
+		assert.Equal(t, 0, intVal)
+	})
+
+	t.Run("context with nil invocation", func(t *testing.T) {
+		ctx := NewInvocationContext(context.Background(), nil)
+		val, ok := GetStateValueFromContext[string](ctx, "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",
+			},
+		}
+		inv.SetState("agent:custom_data", data)
+		ctx := NewInvocationContext(context.Background(), inv)
+
+		retrieved, ok := GetStateValueFromContext[CustomData](ctx, "agent:custom_data")
+		require.True(t, ok)
+		assert.Equal(t, data.ID, retrieved.ID)
+		assert.Equal(t, data.Metadata, retrieved.Metadata)
+	})
+}

graph/state.go

@@ -53,6 +53,35 @@ const (
 // This is the shared data structure that flows between nodes.
 type State map[string]any
 
+// GetStateValue retrieves a typed value from the state.
+//
+// Returns the typed value and true if the key exists and the type matches,
+// or the zero value and false otherwise.
+//
+// Example:
+//
+//	if messages, ok := GetStateValue[[]model.Message](state, StateKeyMessages); ok {
+//	    // use messages
+//	}
+//	if userInput, ok := GetStateValue[string](state, StateKeyUserInput); ok {
+//	    // use userInput
+//	}
+func GetStateValue[T any](s State, key string) (T, bool) {
+	var zero T
+	if s == nil {
+		return zero, false
+	}
+	val, ok := s[key]
+	if !ok {
+		return zero, false
+	}
+	typedVal, ok := val.(T)
+	if !ok {
+		return zero, false
+	}
+	return typedVal, true
+}
+
 // Clone creates a deep copy of the state.
 func (s State) Clone() State {
 	clone := make(State, len(s))