@@ -310,7 +310,7 @@ async def test_graph_edge_cases(mock_strands_tracer, mock_use_span):
310
310
result = await graph .invoke_async ([{"text" : "Original task" }])
311
311
312
312
# Verify entry node was called with original task
313
- entry_agent .invoke_async .assert_called_once_with ([{"text" : "Original task" }])
313
+ entry_agent .invoke_async .assert_called_once_with ([{"text" : "Original task" }], invocation_state = {} )
314
314
assert result .status == Status .COMPLETED
315
315
mock_strands_tracer .start_multiagent_span .assert_called ()
316
316
mock_use_span .assert_called_once ()
@@ -906,7 +906,7 @@ def __init__(self, name):
906
906
self ._session_manager = None
907
907
self .hooks = HookRegistry ()
908
908
909
- async def invoke_async (self , input_data ):
909
+ async def invoke_async (self , input_data , invocation_state = None ):
910
910
# Increment execution count in state
911
911
count = self .state .get ("execution_count" ) or 0
912
912
self .state .set ("execution_count" , count + 1 )
@@ -1300,7 +1300,9 @@ async def test_graph_kwargs_passing_agent(mock_strands_tracer, mock_use_span):
1300
1300
test_invocation_state = {"custom_param" : "test_value" , "another_param" : 42 }
1301
1301
result = await graph .invoke_async ("Test kwargs passing" , test_invocation_state )
1302
1302
1303
- kwargs_agent .invoke_async .assert_called_once_with ([{"text" : "Test kwargs passing" }], ** test_invocation_state )
1303
+ kwargs_agent .invoke_async .assert_called_once_with (
1304
+ [{"text" : "Test kwargs passing" }], invocation_state = test_invocation_state
1305
+ )
1304
1306
assert result .status == Status .COMPLETED
1305
1307
1306
1308
@@ -1335,5 +1337,7 @@ def test_graph_kwargs_passing_sync(mock_strands_tracer, mock_use_span):
1335
1337
test_invocation_state = {"custom_param" : "test_value" , "another_param" : 42 }
1336
1338
result = graph ("Test kwargs passing sync" , test_invocation_state )
1337
1339
1338
- kwargs_agent .invoke_async .assert_called_once_with ([{"text" : "Test kwargs passing sync" }], ** test_invocation_state )
1340
+ kwargs_agent .invoke_async .assert_called_once_with (
1341
+ [{"text" : "Test kwargs passing sync" }], invocation_state = test_invocation_state
1342
+ )
1339
1343
assert result .status == Status .COMPLETED
0 commit comments