fix: recover durable Gateway tasks by session

server.py

@@ -171,11 +171,44 @@ def _execute_durable_task(task: dict[str, Any]) -> dict[str, Any]:
     }
 
 
-_task_worker = TaskWorker(
-    TaskManager(_agent.memory_bank.store, workspace_id=_agent.memory_bank.workspace_id,
-                user_id=_SERVER_ACTOR_ID),
-    _execute_durable_task,
-)
+def _task_manager(session_id: str | None) -> TaskManager:
+    """Build a task manager bound to the authenticated deployment scope."""
+    return TaskManager(
+        _agent.memory_bank.store,
+        workspace_id=_agent.memory_bank.workspace_id,
+        session_id=session_id,
+        user_id=_SERVER_ACTOR_ID,
+    )
+
+
+def _task_worker(session_id: str | None) -> TaskWorker:
+    """Build a bounded worker that can only claim one exact session scope."""
+    return TaskWorker(_task_manager(session_id), _execute_durable_task, max_tasks=1)
+
+
+def _task_scopes(statuses: set[str]) -> set[str | None]:
+    rows = _agent.memory_bank.store.list_tasks(
+        workspace_id=_agent.memory_bank.workspace_id,
+        statuses=statuses,
+        limit=10000,
+        user_id=_SERVER_ACTOR_ID,
+    )
+    return {row.get("session_id") for row in rows}
+
+
+def _recover_task_scopes() -> list[dict[str, Any]]:
+    """Recover every unfinished task without collapsing session boundaries."""
+    scopes = _task_scopes({"pending", "running", "failed", "blocked"})
+    recovered: list[dict[str, Any]] = []
+    for session_id in scopes:
+        recovered.extend(_task_manager(session_id).recover())
+    return recovered
+
+
+def _run_task_worker_cycle() -> None:
+    """Run at most one safe action per pending session in this scheduler tick."""
+    for session_id in sorted(_task_scopes({"pending"}), key=lambda item: item or ""):
+        _task_worker(session_id).run_once()
 
 
 def _run_scheduled_job(job: dict[str, Any]) -> None:
@@ -184,7 +217,7 @@ def _run_scheduled_job(job: dict[str, Any]) -> None:
         _agent.dispatch_learning_cycle(surface="web", trigger="scheduled", max_features=4)
         return
     if payload.get("type") == "task_worker":
-        _task_worker.run_once()
+        _run_task_worker_cycle()
         return
     if payload.get("type") == "chat":
         session_id = _normalise_session_id(payload.get("session_id"))
@@ -196,6 +229,7 @@ def _run_scheduled_job(job: dict[str, Any]) -> None:
 
 _scheduler.register_callback("learning_cycle", _run_scheduled_job)
 _scheduler.register_callback("chat", _run_scheduled_job)
+_scheduler.register_callback("task_worker", _run_scheduled_job)
 
 
 def _normalise_session_id(raw_session_id: Any) -> str:
@@ -683,8 +717,7 @@ async def api_v2_memory_forget_claim(claim_id: str):
 @app.get("/api/v2/tasks", dependencies=[Depends(require_api_access)])
 async def api_v2_tasks(session_id: str | None = None):
     session = _normalise_session_id(session_id) if session_id else None
-    manager = TaskManager(_agent.memory_bank.store, workspace_id=_agent.memory_bank.workspace_id,
-                          session_id=session, user_id=_SERVER_ACTOR_ID)
+    manager = _task_manager(session)
     return {"tasks": manager.store.list_tasks(workspace_id=manager.workspace_id, session_id=session,
                                                 user_id=_SERVER_ACTOR_ID)}
 
@@ -695,8 +728,7 @@ async def api_v2_create_task(request: Request):
     if not isinstance(body, dict) or not str(body.get("description", "")).strip():
         raise HTTPException(400, "description is required")
     session = _normalise_session_id(body.get("session_id"))
-    manager = TaskManager(_agent.memory_bank.store, workspace_id=_agent.memory_bank.workspace_id,
-                          session_id=session, user_id=_SERVER_ACTOR_ID)
+    manager = _task_manager(session)
     checkpoint: dict[str, Any] = {}
     if body.get("action") is not None:
         action = _agent.TOOL_ALIASES.get(str(body.get("action")).strip(), str(body.get("action")).strip())
@@ -722,10 +754,23 @@ async def api_v2_create_task(request: Request):
 
 @app.post("/api/v2/tasks/{task_id}/resume", dependencies=[Depends(require_api_access)])
 async def api_v2_resume_task(task_id: str):
-    manager = _task_worker.manager
+    candidates = _agent.memory_bank.store.list_tasks(
+        workspace_id=_agent.memory_bank.workspace_id,
+        statuses={"pending", "running", "failed", "blocked"},
+        limit=10000,
+        user_id=_SERVER_ACTOR_ID,
+    )
+    stored = next((item for item in candidates if item["id"] == str(task_id)), None)
+    if stored is None:
+        raise HTTPException(404, "Failed or blocked task not found in this scope")
+    manager = _task_manager(stored.get("session_id"))
+    manager.recover()
     task = manager.resume(task_id)
     if task is None:
         raise HTTPException(404, "Failed or blocked task not found in this scope")
+    if not any(job.get("payload", {}).get("type") == "task_worker" for job in _scheduler.list_jobs()):
+        _scheduler.schedule_every("durable-task-worker", 1.0, payload={"type": "task_worker"},
+                                  job_id="job_durable_task_worker", delay_seconds=0)
     return {"task": task, "queued": True}
 
 
@@ -1431,7 +1476,7 @@ async def startup():
         "on_startup", agent=_agent, surface="web", user_id=_SERVER_ACTOR_ID,
         workspace_id=_agent.memory_bank.workspace_id,
     )
-    recovered = _task_worker.recover()
+    recovered = _recover_task_scopes()
     if recovered:
         _audit("TASK_RECOVERY", f"recovered={len(recovered)} pending_or_resumable tasks")
     if not any(job.get("payload", {}).get("type") == "task_worker" for job in _scheduler.list_jobs()):

tests/test_durable_tasks_gateway.py

@@ -0,0 +1,163 @@
+import json
+import os
+import socket
+import subprocess
+import sys
+import tempfile
+import time
+import unittest
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+from event_store import EventStore
+from task_engine import TaskManager
+
+
+class DurableTaskGatewayTests(unittest.TestCase):
+    @staticmethod
+    def _free_port() -> int:
+        with socket.socket() as sock:
+            sock.bind(("127.0.0.1", 0))
+            return int(sock.getsockname()[1])
+
+    @staticmethod
+    def _request(base_url: str, path: str, *, method: str = "GET", body: dict | None = None) -> dict:
+        data = json.dumps(body).encode("utf-8") if body is not None else None
+        request = urllib.request.Request(
+            base_url + path,
+            data=data,
+            method=method,
+            headers={"Content-Type": "application/json", "Authorization": "Bearer durable-test-token"},
+        )
+        with urllib.request.urlopen(request, timeout=5) as response:
+            return json.loads(response.read().decode("utf-8"))
+
+    def _start_server(self, workspace: Path, db_path: Path, skills_path: Path) -> tuple[subprocess.Popen, str]:
+        repository = Path(__file__).parents[1]
+        port = self._free_port()
+        env = os.environ.copy()
+        env.update({
+            "HOME": str(workspace / "home"),
+            "KYROZEN_DB_PATH": str(db_path),
+            "KYROZEN_DISABLE_VECTOR_INDEX": "1",
+            "KYROZEN_PROVIDER": "ollama",
+            "KYROZEN_BASE_URL": "http://127.0.0.1:11434/v1",
+            "KYROZEN_SERVER_TOKEN": "durable-test-token",
+            "KYROZEN_SKILLS_DIR": str(skills_path),
+        })
+        for variable in ("KYROZEN_API_KEY", "DEEPSEEK_API_KEY", "OPENAI_API_KEY", "ANTHROPIC_API_KEY"):
+            env.pop(variable, None)
+        process = subprocess.Popen(
+            [sys.executable, str(repository / "server.py"), "--host", "127.0.0.1", "--port", str(port)],
+            cwd=workspace,
+            env=env,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+        )
+        base_url = f"http://127.0.0.1:{port}"
+        deadline = time.monotonic() + 20
+        while time.monotonic() < deadline:
+            if process.poll() is not None:
+                output = process.stdout.read() if process.stdout else ""
+                self.fail(f"Gateway exited before becoming ready: {output}")
+            try:
+                health = self._request(base_url, "/api/health")
+                if health.get("status") in {"ok", "degraded"}:
+                    return process, base_url
+            except (OSError, urllib.error.URLError):
+                time.sleep(0.1)
+        process.terminate()
+        process.wait(timeout=5)
+        self.fail("Gateway did not become ready")
+
+    @staticmethod
+    def _stop_server(process: subprocess.Popen) -> None:
+        if process.poll() is None:
+            process.terminate()
+            process.wait(timeout=5)
+        if process.stdout:
+            process.stdout.close()
+
+    def test_gateway_recovers_scoped_running_task_and_executes_api_task(self):
+        with tempfile.TemporaryDirectory(prefix="openkyrozen-task-gateway-") as directory:
+            root = Path(directory)
+            workspace = root / "workspace"
+            workspace.mkdir()
+            (workspace / "home").mkdir()
+            db_path = root / "state.sqlite3"
+            skills_path = root / "skills"
+            store = EventStore(db_path)
+            recovery_manager = TaskManager(
+                store, user_id="local", workspace_id="default", session_id="restart-session"
+            )
+            recovery_index = recovery_manager.add_task(
+                "recover a running file task",
+                checkpoint={"action": "write_file", "args": "recovered.txt|after restart"},
+            )
+            recovery_manager.set_status(recovery_index, "running")
+            failed_manager = TaskManager(
+                store, user_id="local", workspace_id="default", session_id="resume-session"
+            )
+            failed_index = failed_manager.add_task("resume only when requested")
+            failed_manager.set_status(failed_index, "failed")
+            failed_id = failed_manager.tasks[failed_index]["id"]
+
+            process, base_url = self._start_server(workspace, db_path, skills_path)
+            try:
+                deadline = time.monotonic() + 10
+                recovered = self._request(base_url, "/api/v2/tasks?session_id=restart-session")
+                while time.monotonic() < deadline:
+                    recovered = self._request(base_url, "/api/v2/tasks?session_id=restart-session")
+                    if recovered["tasks"][0]["status"] == "succeeded":
+                        break
+                    time.sleep(0.1)
+                self.assertEqual(recovered["tasks"][0]["status"], "succeeded")
+                self.assertEqual((workspace / "recovered.txt").read_text(encoding="utf-8"), "after restart")
+
+                created = self._request(
+                    base_url,
+                    "/api/v2/tasks",
+                    method="POST",
+                    body={
+                        "description": "write a live task marker",
+                        "session_id": "task-live",
+                        "action": "write_file",
+                        "args": "task-runtime.txt|created by durable worker",
+                        "acceptance": ["marker exists"],
+                    },
+                )
+                live_id = created["task"]["id"]
+                deadline = time.monotonic() + 10
+                live = created["task"]
+                while time.monotonic() < deadline:
+                    live = self._request(base_url, "/api/v2/tasks?session_id=task-live")["tasks"][0]
+                    if live["status"] == "succeeded":
+                        break
+                    time.sleep(0.1)
+                self.assertEqual(live["id"], live_id)
+                self.assertEqual(live["status"], "succeeded")
+                self.assertEqual(
+                    (workspace / "task-runtime.txt").read_text(encoding="utf-8"),
+                    "created by durable worker",
+                )
+
+                before_resume = self._request(base_url, "/api/v2/tasks?session_id=resume-session")
+                self.assertEqual(before_resume["tasks"][0]["status"], "failed")
+                resumed = self._request(
+                    base_url, f"/api/v2/tasks/{failed_id}/resume", method="POST", body={}
+                )
+                self.assertEqual(resumed["task"]["id"], failed_id)
+                self.assertEqual(resumed["task"]["status"], "pending")
+
+                events = self._request(base_url, "/api/v2/events?session_id=restart-session")["events"]
+                event_types = {event["event_type"] for event in events}
+                self.assertIn("task.recovered", event_types)
+                self.assertIn("task.execution_completed", event_types)
+            finally:
+                self._stop_server(process)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_learning_dispatcher.py

@@ -188,7 +188,7 @@ def test_web_startup_persists_a_learning_cycle_job(self):
              patch.object(server._agent, "_load_project_files_into_memory"), \
              patch.object(server, "_load_plugins"), \
              patch.object(server, "_trigger_hook"), \
-             patch.object(server._task_worker, "recover", return_value=[]), \
+             patch.object(server, "_recover_task_scopes", return_value=[]), \
              patch.object(server._scheduler, "list_jobs", return_value=[
                  {"payload": {"type": "task_worker"}},
              ]), \

tests/test_server.py

@@ -192,7 +192,7 @@ def test_server_startup_initializes_headlessly(self):
                 with patch.object(server._agent, "_load_project_files_into_memory"):
                     with patch.object(server, "_load_plugins"):
                         with patch.object(server, "_trigger_hook"):
-                            with patch.object(server._task_worker, "recover", return_value=[]):
+                            with patch.object(server, "_recover_task_scopes", return_value=[]):
                                 with patch.object(server._scheduler, "list_jobs", return_value=[
                                     {"payload": {"type": "task_worker"}},
                                 ]):