agent: add runtime state retrieval from context (#768)

Author: Rememorio

agent/invocationcontext.go

@@ -55,6 +55,29 @@ func GetStateValueFromContext[T any](ctx context.Context, key string) (T, bool)
 	return GetStateValue[T](inv, key)
 }
 
+// GetRuntimeStateValueFromContext retrieves a typed value from the runtime state
+// stored in the invocation's RunOptions within the context.
+//
+// Returns the typed value and true if the invocation exists, the key exists in
+// RuntimeState, and the type matches, or the zero value and false otherwise.
+//
+// Example:
+//
+//	if userID, ok := GetRuntimeStateValueFromContext[string](ctx, "user_id"); ok {
+//	    log.Printf("User ID: %s", userID)
+//	}
+//	if roomID, ok := GetRuntimeStateValueFromContext[int](ctx, "room_id"); ok {
+//	    log.Printf("Room ID: %d", roomID)
+//	}
+func GetRuntimeStateValueFromContext[T any](ctx context.Context, key string) (T, bool) {
+	var zero T
+	inv, ok := InvocationFromContext(ctx)
+	if !ok || inv == nil {
+		return zero, false
+	}
+	return GetRuntimeStateValue[T](&inv.RunOptions, key)
+}
+
 // CheckContextCancelled check context cancelled
 func CheckContextCancelled(ctx context.Context) error {
 	select {

agent/invocationcontext_test.go

@@ -196,3 +196,149 @@ func TestGetStateValueFromContext(t *testing.T) {
 		assert.Equal(t, data.Metadata, retrieved.Metadata)
 	})
 }
+
+func TestGetRuntimeStateValueFromContext(t *testing.T) {
+	t.Run("context without invocation", func(t *testing.T) {
+		ctx := context.Background()
+		val, ok := GetRuntimeStateValueFromContext[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 := GetRuntimeStateValueFromContext[string](ctx, "nonexistent")
+		assert.False(t, ok)
+		assert.Equal(t, "", val)
+	})
+
+	t.Run("context with invocation but nil RuntimeState", func(t *testing.T) {
+		inv := NewInvocation()
+		ctx := NewInvocationContext(context.Background(), inv)
+		val, ok := GetRuntimeStateValueFromContext[string](ctx, "key")
+		assert.False(t, ok)
+		assert.Equal(t, "", val)
+	})
+
+	t.Run("context with invocation and matching type", func(t *testing.T) {
+		inv := NewInvocation(
+			WithInvocationRunOptions(RunOptions{
+				RuntimeState: map[string]any{
+					"user_id": "12345",
+					"room_id": 678,
+					"config":  true,
+					"score":   3.14,
+				},
+			}),
+		)
+		ctx := NewInvocationContext(context.Background(), inv)
+
+		// Test string value.
+		userID, ok := GetRuntimeStateValueFromContext[string](ctx, "user_id")
+		assert.True(t, ok)
+		assert.Equal(t, "12345", userID)
+
+		// Test int value.
+		roomID, ok := GetRuntimeStateValueFromContext[int](ctx, "room_id")
+		assert.True(t, ok)
+		assert.Equal(t, 678, roomID)
+
+		// Test bool value.
+		config, ok := GetRuntimeStateValueFromContext[bool](ctx, "config")
+		assert.True(t, ok)
+		assert.Equal(t, true, config)
+
+		// Test float64 value.
+		score, ok := GetRuntimeStateValueFromContext[float64](ctx, "score")
+		assert.True(t, ok)
+		assert.Equal(t, 3.14, score)
+	})
+
+	t.Run("context with invocation but type mismatch", func(t *testing.T) {
+		inv := NewInvocation(
+			WithInvocationRunOptions(RunOptions{
+				RuntimeState: map[string]any{
+					"value": "hello",
+				},
+			}),
+		)
+		ctx := NewInvocationContext(context.Background(), inv)
+
+		// Try to get as int when it's actually string.
+		intVal, ok := GetRuntimeStateValueFromContext[int](ctx, "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 := GetRuntimeStateValueFromContext[string](ctx, "key")
+		assert.False(t, ok)
+		assert.Equal(t, "", val)
+	})
+
+	t.Run("slice type", func(t *testing.T) {
+		inv := NewInvocation(
+			WithInvocationRunOptions(RunOptions{
+				RuntimeState: map[string]any{
+					"tags": []string{"tag1", "tag2", "tag3"},
+				},
+			}),
+		)
+		ctx := NewInvocationContext(context.Background(), inv)
+
+		tags, ok := GetRuntimeStateValueFromContext[[]string](ctx, "tags")
+		assert.True(t, ok)
+		assert.Equal(t, []string{"tag1", "tag2", "tag3"}, tags)
+	})
+
+	t.Run("map type", func(t *testing.T) {
+		inv := NewInvocation(
+			WithInvocationRunOptions(RunOptions{
+				RuntimeState: map[string]any{
+					"metadata": map[string]string{
+						"key1": "value1",
+						"key2": "value2",
+					},
+				},
+			}),
+		)
+		ctx := NewInvocationContext(context.Background(), inv)
+
+		metadata, ok := GetRuntimeStateValueFromContext[map[string]string](ctx, "metadata")
+		assert.True(t, ok)
+		assert.Equal(t, "value1", metadata["key1"])
+		assert.Equal(t, "value2", metadata["key2"])
+	})
+
+	t.Run("complex struct type", func(t *testing.T) {
+		type UserContext struct {
+			UserID   string
+			RoomID   int
+			Metadata map[string]string
+		}
+
+		userCtx := UserContext{
+			UserID: "user-123",
+			RoomID: 456,
+			Metadata: map[string]string{
+				"key1": "value1",
+			},
+		}
+		inv := NewInvocation(
+			WithInvocationRunOptions(RunOptions{
+				RuntimeState: map[string]any{
+					"user_context": userCtx,
+				},
+			}),
+		)
+		ctx := NewInvocationContext(context.Background(), inv)
+
+		retrieved, ok := GetRuntimeStateValueFromContext[UserContext](ctx, "user_context")
+		require.True(t, ok)
+		assert.Equal(t, userCtx.UserID, retrieved.UserID)
+		assert.Equal(t, userCtx.RoomID, retrieved.RoomID)
+		assert.Equal(t, userCtx.Metadata, retrieved.Metadata)
+	})
+}