fix: stateful-tools domain handling + preserve tool output in conversation history

sdk/python/e2e/test_suite14_stateful_domain.py

@@ -11,7 +11,7 @@
   - Stateful swarm handoff + check_transfer execute in domain
   - Pipeline sub-agent tools inherit parent's domain
   - Concurrent stateful executions are isolated (different domains)
-  - Agent without stateful flag works without domain
+  - Non-stateful agents work without domain (regression guard)
 
 Validation: all assertions inspect the workflow execution via server API.
 No mocks, no LLM output parsing, fully deterministic.
@@ -512,12 +512,13 @@ def _make_agent(suffix):
                 f"{[(t['taskDefName'], t.get('pollCount')) for t in scheduled]}"
             )
 
-    # ── Test 6: Agent without stateful flag has no domain ──────────
+    # ── Test 6: Non-stateful has no domain (regression) ────────────
 
     def test_non_stateful_no_domain(self, fresh_runtime, model):
-        """Agent without stateful=True works without domain assignment.
+        """Non-stateful agent works without domain assignment.
 
         Validates: taskToDomain is empty, tasks have no domain, execution completes.
+        This is a regression guard — the domain fix must not break non-stateful agents.
         """
         agent = Agent(
             name="e2e_s14_non_stateful",

sdk/python/src/agentspan/agents/runtime/runtime.py

@@ -417,6 +417,20 @@ def _clear_workflow_credentials(
         with _workflow_credentials_lock:
             _workflow_credentials.pop(execution_id, None)
 
+    def _resolve_worker_domain(self, execution_id: str, run_id: Optional[str]) -> Optional[str]:
+        """Return the domain workers should poll for this execution.
+
+        A fresh stateful start uses ``run_id`` as the task domain.  If the
+        server returns an existing execution for an idempotency key, that
+        execution already has its original ``taskToDomain`` mapping, so the
+        freshly generated ``run_id`` would be wrong.  Prefer the server's
+        recorded domain and fall back to the generated one for brand-new runs
+        or older servers.
+        """
+        if not run_id:
+            return None
+        return self._extract_domain(execution_id) or run_id
+
     def _pre_deploy_nested_skills(self, agent: Agent) -> list:
         """Pre-deploy any skill agents nested inside agent_tool wrappers.
 
@@ -716,10 +730,12 @@ def _prepare(self, agent: Agent) -> Any:
                     logger.debug("Starting workers for agent '%s'", agent.name)
                     self._worker_manager.start()
                     self._workers_started = True
-                elif new_workers:
-                    # Inject new workers into the running TaskHandler without
-                    # stopping existing ones.  This avoids the fork() deadlock
-                    # window caused by a full stop/restart cycle.
+                else:
+                    # New stateful runs can register the same task names under
+                    # a different domain. WorkerManager is domain-aware and
+                    # starts only missing (task_name, domain) pairs, so call it
+                    # even when the task-name set has not changed — this avoids
+                    # the fork() deadlock window of a full stop/restart cycle.
                     self._worker_manager.start()
 
         return wf
@@ -832,10 +848,12 @@ def _prepare_workers(
                     logger.debug("Starting workers for agent '%s'", agent.name)
                     self._worker_manager.start()
                     self._workers_started = True
-                elif new_workers:
-                    # Inject new workers into the running TaskHandler without
-                    # stopping existing ones.  This avoids the fork() deadlock
-                    # window caused by a full stop/restart cycle.
+                else:
+                    # New stateful runs can register the same task names under
+                    # a different domain. WorkerManager is domain-aware and
+                    # starts only missing (task_name, domain) pairs, so call it
+                    # even when the task-name set has not changed — this avoids
+                    # the fork() deadlock window of a full stop/restart cycle.
                     self._worker_manager.start()
 
     def _collect_worker_names(
@@ -942,8 +960,9 @@ def _collect_worker_names(
         ):
             names.add(f"{agent.name}_router_fn")
 
-        # Handoff check (swarm with handoff conditions)
-        if agent.handoffs:
+        # Handoff check — needed for any SWARM parent (server always generates
+        # the task) or any agent with explicit handoff conditions.
+        if agent.handoffs or (agent.strategy == "swarm" and agent.agents):
             names.add(f"{agent.name}_handoff_check")
 
         # Swarm transfer workers — prefixed with SOURCE agent name
@@ -1127,8 +1146,9 @@ def _server_needs(task_name: str) -> bool:
             if _server_needs(task_name):
                 self._register_router_worker(agent, domain=domain)
 
-        # 7. Handoff check (swarm with handoff conditions)
-        if agent.handoffs:
+        # 7. Handoff check — needed for any SWARM parent (server always
+        #    generates the task) or any agent with explicit handoff conditions.
+        if agent.handoffs or (agent.strategy == "swarm" and agent.agents):
             task_name = f"{agent.name}_handoff_check"
             if _server_needs(task_name):
                 self._register_handoff_worker(agent, domain=domain)
@@ -2545,8 +2565,10 @@ def run(
             run_id=run_id,
         )
 
-        self._prepare_workers(agent, required_workers=required_workers, domain=run_id)
-        self._register_and_start_skill_workers(pre_deployed_skills, domain=run_id)
+        worker_domain = self._resolve_worker_domain(execution_id, run_id)
+
+        self._prepare_workers(agent, required_workers=required_workers, domain=worker_domain)
+        self._register_and_start_skill_workers(pre_deployed_skills, domain=worker_domain)
 
         self._register_workflow_credentials(execution_id, credentials)
 
@@ -3675,8 +3697,10 @@ def start(
             run_id=run_id,
         )
 
-        self._prepare_workers(agent, required_workers=required_workers, domain=run_id)
-        self._register_and_start_skill_workers(pre_deployed_skills, domain=run_id)
+        worker_domain = self._resolve_worker_domain(execution_id, run_id)
+
+        self._prepare_workers(agent, required_workers=required_workers, domain=worker_domain)
+        self._register_and_start_skill_workers(pre_deployed_skills, domain=worker_domain)
 
         return AgentHandle(
             execution_id=execution_id, runtime=self, correlation_id=correlation_id, run_id=run_id
@@ -4070,8 +4094,10 @@ async def run_async(
             run_id=run_id,
         )
 
-        self._prepare_workers(agent, required_workers=required_workers, domain=run_id)
-        self._register_and_start_skill_workers(pre_deployed_skills, domain=run_id)
+        worker_domain = self._resolve_worker_domain(execution_id, run_id)
+
+        self._prepare_workers(agent, required_workers=required_workers, domain=worker_domain)
+        self._register_and_start_skill_workers(pre_deployed_skills, domain=worker_domain)
         self._register_workflow_credentials(execution_id, credentials)
 
         effective_timeout = timeout or (
@@ -4206,8 +4232,10 @@ async def start_async(
             run_id=run_id,
         )
 
-        self._prepare_workers(agent, required_workers=required_workers, domain=run_id)
-        self._register_and_start_skill_workers(pre_deployed_skills, domain=run_id)
+        worker_domain = self._resolve_worker_domain(execution_id, run_id)
+
+        self._prepare_workers(agent, required_workers=required_workers, domain=worker_domain)
+        self._register_and_start_skill_workers(pre_deployed_skills, domain=worker_domain)
 
         return AgentHandle(
             execution_id=execution_id, runtime=self, correlation_id=correlation_id, run_id=run_id

sdk/python/tests/unit/test_runtime.py

@@ -8,6 +8,7 @@
 """
 
 import logging
+import threading
 import uuid
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -713,6 +714,74 @@ def test_sub_agent_with_string_tools_does_not_raise(self):
         assert _has_stateful_tools(parent) is False
 
 
+class TestStatefulWorkerDomains:
+    """Stateful workers must use the execution's real task domain."""
+
+    def test_resolve_worker_domain_prefers_server_domain(self):
+        from agentspan.agents.runtime.runtime import AgentRuntime
+
+        rt = AgentRuntime.__new__(AgentRuntime)
+        rt._extract_domain = lambda execution_id: "original-domain"
+
+        assert rt._resolve_worker_domain("wf-1", "fresh-domain") == "original-domain"
+
+    def test_resolve_worker_domain_falls_back_to_generated_run_id(self):
+        from agentspan.agents.runtime.runtime import AgentRuntime
+
+        rt = AgentRuntime.__new__(AgentRuntime)
+        rt._extract_domain = lambda execution_id: None
+
+        assert rt._resolve_worker_domain("wf-1", "fresh-domain") == "fresh-domain"
+
+    def test_resolve_worker_domain_returns_none_for_stateless_execution(self):
+        from agentspan.agents.runtime.runtime import AgentRuntime
+
+        rt = AgentRuntime.__new__(AgentRuntime)
+        rt._extract_domain = lambda execution_id: "should-not-be-used"
+
+        assert rt._resolve_worker_domain("wf-1", None) is None
+
+    def test_prepare_workers_starts_worker_manager_when_only_domain_changes(self):
+        """Same task name under a new domain still needs a new polling process."""
+        from types import SimpleNamespace
+
+        from agentspan.agents.runtime.runtime import AgentRuntime
+
+        class FakeWorkerManager:
+            def __init__(self):
+                self.starts = 0
+
+            def start(self):
+                self.starts += 1
+
+        rt = AgentRuntime.__new__(AgentRuntime)
+        rt._config = SimpleNamespace(
+            auto_register_integrations=False,
+            auto_start_workers=True,
+        )
+        rt._worker_start_lock = threading.Lock()
+        rt._registered_tool_names = {"write_architecture"}
+        rt._workers_started = True
+        rt._worker_manager = FakeWorkerManager()
+        rt._associate_templates_with_models = lambda agent: None
+        registered_domains = []
+        rt._register_workers = lambda agent, required_workers=None, domain=None: (
+            registered_domains.append(domain)
+        )
+        rt._has_worker_tools = lambda agent: True
+        rt._collect_worker_names = lambda agent, required_workers=None: {"write_architecture"}
+
+        agent = Agent(name="pipeline", model="openai/gpt-4o")
+        rt._prepare_workers(
+            agent,
+            required_workers={"write_architecture"},
+            domain="original-domain",
+        )
+
+        assert registered_domains == ["original-domain"]
+        assert rt._worker_manager.starts == 1
+
+
 # ── _extract_token_usage ────────────────────────────────────────────────