diff --git a/server.py b/server.py index 301f541..4c4396a 100644 --- a/server.py +++ b/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()): diff --git a/tests/test_durable_tasks_gateway.py b/tests/test_durable_tasks_gateway.py new file mode 100644 index 0000000..24af247 --- /dev/null +++ b/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() diff --git a/tests/test_learning_dispatcher.py b/tests/test_learning_dispatcher.py index 9d3919c..f396a51 100644 --- a/tests/test_learning_dispatcher.py +++ b/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"}}, ]), \ diff --git a/tests/test_server.py b/tests/test_server.py index e3ada17..2ea3d42 100644 --- a/tests/test_server.py +++ b/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"}}, ]):