diff --git a/tests/gateway/test_gateway_command_dispatch_minimal.py b/tests/gateway/test_gateway_command_dispatch_minimal.py
index d0a9b96593..fa8ad0a13b 100644
--- a/tests/gateway/test_gateway_command_dispatch_minimal.py
+++ b/tests/gateway/test_gateway_command_dispatch_minimal.py
@@ -98,7 +98,7 @@ def _make_runner():
     runner._should_send_telegram_lobby_reminder = lambda _source: False
     runner._check_slash_access = lambda _source, _command: None
     runner._begin_session_run_generation = lambda _key: 1
-    runner._release_running_agent_state = lambda key: runner._running_agents.pop(key, None)
+    runner._release_running_agent_state = lambda key, run_generation=None: runner._running_agents.pop(key, None)
     return runner, adapter
 
 
diff --git a/tests/hermes_cli/test_pre_command_hook.py b/tests/hermes_cli/test_pre_command_hook.py
index c3a4079f5d..1504354ca8 100644
--- a/tests/hermes_cli/test_pre_command_hook.py
+++ b/tests/hermes_cli/test_pre_command_hook.py
@@ -284,7 +284,7 @@ def _make_runner():
     runner._check_slash_access = lambda _source, _command: None
     runner._begin_session_run_generation = lambda _key: 1
     runner._release_running_agent_state = (
-        lambda key: runner._running_agents.pop(key, None)
+        lambda key, run_generation=None: runner._running_agents.pop(key, None)
     )
     return runner, adapter