[FEATURE] Pass invocation_state to edge condition call

[MElkady]

Overview

Add support for passing invocation_state to edge condition functions in Graph execution, enabling conditional edge traversal based on runtime context provided during graph invocation.

Problem Statement

Currently, edge condition functions only have access to GraphState, limiting their ability to make decisions based on runtime context passed during graph invocation. This makes it difficult to conditionally enable nodes based on the invocation_state received in the graph invocation.

Implementation Requirements

Based on discussion and repository analysis, the implementation should use a Protocol-based Union type for backwards compatibility and future extensibility.

Technical Approach

1. Define Edge Condition Protocol

Create a new Protocol in src/strands/multiagent/graph.py that allows future expansion with kwargs:

from typing import Protocol, runtime_checkable

@runtime_checkable
class EdgeConditionWithContext(Protocol):
    """Protocol for edge conditions that receive invocation context.
    
    This protocol allows conditions to access both GraphState and invocation_state,
    and is designed to be expanded with additional kwargs in the future.
    """
    def __call__(
        self, 
        state: GraphState, 
        *, 
        invocation_state: dict[str, Any] | None = None,
        **kwargs: Any
    ) -> bool: ...

# Legacy condition type (backwards compatible)
LegacyEdgeCondition = Callable[[GraphState], bool]

# Union type supporting both old and new signatures
EdgeCondition = LegacyEdgeCondition | EdgeConditionWithContext

2. Update GraphEdge Dataclass

@dataclass
class GraphEdge:
        ...
        if new style:
            return self.condition(state, invocation_state=invocation_state)
        
        # Legacy: single positional parameter only
        return self.condition(state)

3. Update Call Sites

Update all locations that call should_traverse() to pass invocation_state:

• Graph._is_node_ready_with_conditions() (line ~837)
• Graph._compute_ready_nodes_for_resume() (line ~1190)
• Graph._build_node_input() (line ~1080)

4. Update GraphBuilder Type Hint

def add_edge(
    self,
    from_node: str | GraphNode,
    to_node: str | GraphNode,
    condition: EdgeCondition | None = None,
) -> GraphEdge:

Files to Modify

File	Changes
src/strands/multiagent/graph.py	Add Protocol, update types, update should_traverse(), update all call sites
tests/strands/multiagent/test_graph.py	Add tests for new-style conditions with invocation_state
tests_integ/test_multiagent_graph.py	Add integration tests for conditions with invocation_state

Backwards Compatibility

• Old-style conditions (def condition(state: GraphState) -> bool) continue to work unchanged
• New-style conditions (def condition(state: GraphState, *, invocation_state=None, **kwargs) -> bool) receive invocation context
• Runtime signature inspection ensures correct calling convention
• No deprecation warnings - both styles supported indefinitely

Example Usage

# Old-style (still works)
def check_completion(state: GraphState) -> bool:
    return "data_processor" in {n.node_id for n in state.completed_nodes}

# New-style (with invocation context)
def should_run_premium_node(
    state: GraphState, 
    *, 
    invocation_state: dict[str, Any] | None = None,
    **kwargs: Any
) -> bool:
    if invocation_state is None:
        return False
    return invocation_state.get("user_tier") == "premium"

# Graph setup
builder = GraphBuilder()
builder.add_node(start_agent, "start")
builder.add_node(premium_agent, "premium_processor")
builder.add_edge("start", "premium_processor", condition=should_run_premium_node)
graph = builder.build()

# Invocation with context
result = graph(task="process data", invocation_state={"user_tier": "premium"})

Acceptance Criteria

• EdgeConditionWithContext Protocol defined with kwargs support for future expansion
• EdgeCondition Union type supports both legacy and new condition signatures
• GraphEdge.should_traverse() accepts and passes invocation_state
• All call sites updated to pass invocation_state through the execution flow
• Existing tests pass without modification (backwards compatibility)
• New unit tests for conditions that use invocation_state
• New integration test demonstrating the use case
• Type hints pass mypy strict mode

Implementation Notes

• Use inspect.signature() for runtime detection of condition signature
• The Protocol approach allows adding more context parameters in the future without breaking changes
• Consider caching signature inspection results for performance if needed

[dbschmigelski]

/strands refine refine this proposal but consider the interface of condition which currently accepts 2 input params. We need to have a solution that is backwards compatible

[github-actions]

Repository Analysis & Clarifying Questions

I've analyzed the repository structure and the current implementation. Here's what I found and some questions to ensure proper implementation:

Repository Context

Current Implementation (in src/strands/multiagent/graph.py)

• Condition type: Callable[[GraphState], bool] | None (line 147)
• should_traverse() method: Only passes GraphState to condition (line 153-157)
• invocation_state: Already flows through graph execution methods but is NOT passed to edge conditions

Files Requiring Modification

• src/strands/multiagent/graph.py:
  • GraphEdge.condition type annotation
  • GraphEdge.should_traverse() method signature and call sites
  • All places that call should_traverse() or edge.condition()

Testing Files

• tests/strands/multiagent/test_graph.py - numerous condition tests exist
• tests_integ/test_multiagent_graph.py - integration tests with conditions

---

Clarifying Questions

1. Backwards Compatibility Approach

The current condition interface accepts 1 parameter (GraphState), not 2 as mentioned in your comment. To maintain backwards compatibility when adding invocation_state as a second parameter, which approach do you prefer?

Option A: Runtime signature inspection (using inspect.signature)

# Detect param count at runtime and call appropriately
import inspect

def should_traverse(self, state: GraphState, invocation_state: dict[str, Any] | None = None) -> bool:
    if self.condition is None:
        return True
    sig = inspect.signature(self.condition)
    if len(sig.parameters) == 1:
        return self.condition(state)  # Old-style: 1 param
    return self.condition(state, invocation_state)  # New-style: 2 params

Option B: Union type with explicit handling

EdgeCondition = Callable[[GraphState], bool] | Callable[[GraphState, dict[str, Any] | None], bool]

Option C: Require new conditions to use optional second param

# All conditions should be written as:
def my_condition(state: GraphState, invocation_state: dict[str, Any] | None = None) -> bool:
    ...

Which approach aligns best with the SDK's design principles?

2. Type Annotation Strategy

For the type hint, should we use:

• A Union type that explicitly shows both signatures?
• A Protocol to define the condition interface?
• Keep it simple with just the new 2-param signature (relying on runtime detection)?

3. Deprecation Strategy

Should we:

• Silently support both signatures indefinitely?
• Log a deprecation warning when 1-param conditions are detected?
• Document a migration path in release notes?

4. Use Case Confirmation

Please confirm the use case: You want to conditionally enable/disable edges based on values in invocation_state passed during graph invocation, like:

def should_run_premium_node(state: GraphState, invocation_state: dict[str, Any] | None = None) -> bool:
    if invocation_state is None:
        return False
    return invocation_state.get("user_tier") == "premium"

builder.add_edge("start", "premium_processor", condition=should_run_premium_node)

# At invocation time:
graph(task="process data", invocation_state={"user_tier": "premium"})

Is this the intended usage pattern?

---

Scope Assessment

This feature appears to be a reasonable scope for a single PR:

• Primary changes in one file (graph.py)
• Test updates required but straightforward
• No GitHub workflow changes needed
• Estimated implementation: ~3-5 days

Please respond when you have a chance so I can finalize the implementation requirements.

[dbschmigelski]

/strands lets do EdgeCondition with a Union type, but isntead of Callable[[GraphState, dict[str, Any] | None] lets define a protocol so we can expand this freely in the future with kwargs

[github-actions]

Task Refinement Complete

Summary

I've updated the issue description with comprehensive implementation requirements based on our discussion.

Key Decisions Made

1. Protocol-based Union Type: Using EdgeConditionWithContext Protocol combined with LegacyEdgeCondition Callable in a Union type
2. Future Extensibility: Protocol uses **kwargs to allow adding more context parameters without breaking changes
3. Backwards Compatibility: Runtime signature inspection to detect and handle both old-style (1-param) and new-style (kwargs) conditions
4. No Deprecation: Both condition styles will be supported indefinitely

Implementation Scope

• Primary file: src/strands/multiagent/graph.py
• Test files: tests/strands/multiagent/test_graph.py, tests_integ/test_multiagent_graph.py
• Estimated effort: 3-5 days
• Single PR: Yes, reasonable scope

Ready for Implementation

The issue now contains:

• ✅ Clear technical approach with code examples
• ✅ Files to modify identified
• ✅ Backwards compatibility strategy defined
• ✅ Example usage patterns documented
• ✅ Acceptance criteria specified

This issue is ready for implementation.