feat: expose user-defined state in MultiAgent Graph

[aditya270520]

🎯 Feature: Expose User Defined State in MultiAgent Graph

Summary

This PR implements the requested feature from issue #665 to expose user-defined state in MultiAgent Graph, similar to how Agent has AgentState and Swarm has SharedContext. This enables users to store and access information across graph nodes without including it in the LLM context.

Changes Made

1. Add SharedContext class to multiagent.base for unified state management

• File: src/strands/multiagent/base.py
• Purpose: Centralize shared state management across all multi-agent implementations
• Features:
  • JSON serialization validation for all stored values
  • Node isolation to prevent interference between different nodes
  • Deep copying to prevent mutation of returned objects
  • Comprehensive error handling for invalid keys and non-serializable values

2. Add shared_context property to Graph class for easy access

• File: src/strands/multiagent/graph.py
• Purpose: Provide simple, intuitive API for accessing shared context
• Usage: graph.shared_context.add_context("node_id", "key", "value")
• Benefits: Consistent with existing patterns in the codebase

3. Update GraphState to include shared_context field

• File: src/strands/multiagent/graph.py
• Purpose: Enable persistent storage of user-defined state across graph execution
• Implementation: Added shared_context: SharedContext field to GraphState class
• Benefits: Maintains state between node executions and follows same pattern as other state classes

4. Refactor Swarm to use SharedContext from base module

• File: src/strands/multiagent/swarm.py
• Purpose: Eliminate code duplication and ensure consistent behavior
• Changes:
  • Updated imports to use SharedContext from base module
  • Fixed API usage to work with new implementation
  • Maintained backward compatibility for existing users

5. Add comprehensive tests for SharedContext functionality

• Files:
  • tests/strands/multiagent/test_base.py - Unit tests for SharedContext
  • tests/strands/multiagent/test_graph.py - Integration tests for Graph
• Coverage:
  • Basic operations (add/get context)
  • Input validation (keys, values)
  • Node isolation
  • Deep copying semantics
  • Error handling

6. Support JSON serialization validation and deep copying

• Features:
  • Automatic validation that all values are JSON serializable
  • Deep copying to prevent unintended mutations
  • Comprehensive error messages for validation failures
  • Support for all JSON-compatible types (str, int, float, bool, list, dict, None)

Related Issues

• Resolves: [FEATURE] Expose user defined state in MultiAgent Graph #665 - [FEATURE] Expose user defined state in MultiAgent Graph

Documentation PR

No documentation PR needed at this time. The feature is self-contained and follows existing patterns.

Type of Change

New feature - This adds new functionality to the Graph class without breaking existing functionality.

Testing

How have you tested the change?

1. Local Testing: Successfully tested implementation on macOS with Python 3.10
2. Unit Tests: All new tests pass successfully
3. Integration Tests: Graph integration tests pass
4. Code Quality: All code passes linting and formatting checks
5. Backward Compatibility: Existing Swarm functionality continues to work

Verify that the changes do not break functionality or introduce warnings in consuming repositories

• agents-docs: No impact - this is a new feature addition
• agents-tools: No impact - no changes to tool interfaces
• agents-cli: No impact - no changes to CLI functionality

Test Results

# All tests passing
tests/strands/multiagent/test_base.py::test_shared_context_initialization PASSED
tests/strands/multiagent/test_base.py::test_shared_context_add_context PASSED
tests/strands/multiagent/test_base.py::test_shared_context_get_context PASSED
tests/strands/multiagent/test_base.py::test_shared_context_validation PASSED
tests/strands/multiagent/test_base.py::test_shared_context_isolation PASSED
tests/strands/multiagent/test_base.py::test_shared_context_copy_semantics PASSED
tests/strands/multiagent/test_graph.py::test_graph_shared_context PASSED
tests/strands/multiagent/test_graph.py::test_graph_shared_context_validation PASSED

Code Quality Checks

• Linting: All code passes ruff check with no errors
• Formatting: All code is properly formatted with ruff format
• Type Safety: Full type hints and validation implemented
• Documentation: Comprehensive docstrings for all new methods

Usage Examples

Basic Usage

from strands import Agent
from strands.multiagent import Graph, GraphBuilder

# Create agents
data_processor = Agent(name="data_processor")
file_analyzer = Agent(name="file_analyzer")

# Build graph
builder = GraphBuilder()
builder.add_node(data_processor, "data_processor")
builder.add_node(file_analyzer, "file_analyzer")
builder.add_edge("data_processor", "file_analyzer")
builder.set_entry_point("data_processor")

graph = builder.build()

# Store information in shared context
graph.shared_context.add_context("data_processor", "file_path", "/path/to/large/data.csv")
graph.shared_context.add_context("data_processor", "config", {"batch_size": 1000, "timeout": 30})

# Retrieve context
file_path = graph.shared_context.get_context("data_processor", "file_path")
config = graph.shared_context.get_context("data_processor", "config")

Advanced Usage

# Store complex data structures
graph.shared_context.add_context("node1", "processed_files", [
    "/tmp/file1.csv",
    "/tmp/file2.csv"
])

# Store configuration
graph.shared_context.add_context("node1", "processing_config", {
    "batch_size": 1000,
    "timeout": 30,
    "retry_attempts": 3,
    "quality_threshold": 0.95
})

# Retrieve all context for a node
all_context = graph.shared_context.get_context("node1")

Benefits

1. 🎯 Solves Original Issue: Graph now has user-defined state like Agent and Swarm
2. 🔄 Consistency: Unified approach across all multi-agent implementations
3. 🔒 Backward Compatibility: Existing Swarm code continues to work
4. 📝 Clean API: Simple, intuitive interface for storing and retrieving shared data
5. ✅ Validation: Automatic JSON serialization validation prevents errors
6. 🚀 Performance: Avoids including large data in LLM context
7. �� Maintainability: Single implementation reduces code duplication

Migration Path

For Existing Swarm Users

No changes required - existing code continues to work. However, users are encouraged to import from the base module:

# Old (still works)
from strands.multiagent.swarm import SharedContext

# New (recommended)
from strands.multiagent.base import SharedContext

For New Graph Users

Simply use the new shared_context property:

graph = Graph(...)
graph.shared_context.add_context("node_id", "key", "value")
value = graph.shared_context.get_context("node_id", "key")

Checklist

• I have read the CONTRIBUTING document
• I have added necessary tests that prove my fix is effective or my feature works
• I have updated the documentation accordingly (no new docs needed for this feature)
• I have added appropriate examples to demonstrate the feature
• My changes generate no new warnings
• Any dependent changes have been merged and published (no dependencies)

Files Changed

• src/strands/multiagent/base.py - Added SharedContext class
• src/strands/multiagent/graph.py - Added shared_context to GraphState and Graph
• src/strands/multiagent/swarm.py - Updated to use base SharedContext
• tests/strands/multiagent/test_base.py - Tests for SharedContext
• tests/strands/multiagent/test_graph.py - Tests for Graph integration

Impact Assessment

• Breaking Changes: None
• New Dependencies: None
• Performance Impact: Minimal - only adds lightweight state management
• Memory Impact: Minimal - only stores explicitly added context
• API Changes: None - only adds new functionality

---

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[aditya270520]

Could you please take a moment to review the updated implementation? I've made every effort to address your feedback while maintaining the requested functionality for issue #665. The code is now ready for CI validation and should pass all checks.
If there are any remaining concerns or if you'd like me to make additional adjustments, I'm happy to iterate further. Your guidance has been invaluable in improving this implementation.
Thank you for your time and patience!

[dbschmigelski]

nit we can remove the hash from SwarmNode since we can now inherit

[dbschmigelski]

Sorry for the delay I need to take a look at this again more deeply.

What is strange to be is that we have different hierarchies for Swarm and Graph.

For Swarm we have SharedContext at the same level as SwarmState. For Graph we have SharedContext nested within GraphState.

Even if we kept it this way I'm not sure we need this convenience method to surface shared context at the GraphNode level. It feels strange to place the context on the state then surface it higher.

[dbschmigelski]

I'll come back tomorrow with a final decision on the path for this.

[dbschmigelski]

I'm not sure I'm set on this approach because of backwards compatibility.

If context were private _context then the deepcopy could make sense to me. Since its not, having a user access different objects though the two approaches feels wrong.

get_context(node, key) and context[node.node_id][key]