aicortexpy 0.1.0__py3-none-any.whl

aicortexpy/daemon.py

@@ -0,0 +1,677 @@
+"""Core daemon loops: register, heartbeat actions, wakeup, claim/run, GC."""
+
+from __future__ import annotations
+
+import logging
+import signal
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Any
+
+from .artifacts import report_issue_artifacts
+from .auto_commit import auto_commit_and_push
+from .chat_acp_pool import ChatACPPool
+from .client import APIError, DaemonClient, Task, is_not_found
+from .config import Config, as_register_runtimes, detect_local_timezone
+from .gc import ActiveEnvRoots, gc_meta_for_task, run_gc_loop, write_gc_meta
+from .health import HealthServer
+from .local_skills import list_runtime_local_skills, load_runtime_local_skill_bundle
+from .models import list_models, model_selection_supported
+from .platform_info import platform_label
+from .poisoned import classify_poisoned_error, classify_poisoned_output
+from .repo import RepoAllowlist
+from .runner import run_agent
+from .wakeup import WakeupHub, run_wakeup_loop
+
+log = logging.getLogger("aicortexpy")
+
+
+def should_interrupt_agent(status: str, err: Exception | None) -> bool:
+    if err is not None:
+        return is_not_found(err, "task not found")
+    return status == "cancelled"
+
+
+class Daemon:
+    def __init__(self, cfg: Config):
+        self.cfg = cfg
+        self.client = DaemonClient(
+            cfg.server_base_url,
+            cfg.token,
+            version=cfg.cli_version,
+        )
+        self._stop = threading.Event()
+        self._wake = threading.Event()
+        self._lock = threading.Lock()
+        self.runtime_index: dict[str, dict[str, Any]] = {}
+        self.workspace_runtimes: dict[str, list[str]] = {}
+        self.workspace_names: dict[str, str] = {}
+        self._reregister_backoff: dict[str, float] = {}
+        self.active_tasks = 0
+        self._health: HealthServer | None = None
+        self._pool = ThreadPoolExecutor(max_workers=max(1, cfg.max_concurrent_tasks))
+        self.chat_acp_pool = ChatACPPool.from_config(cfg)
+        self.active_env_roots = ActiveEnvRoots()
+        self.repo_allowlist = RepoAllowlist()
+        self.wakeup_hub = WakeupHub()
+
+    def stop(self) -> None:
+        self._stop.set()
+        self._wake.set()
+        try:
+            self.chat_acp_pool.close_all()
+        except Exception:
+            pass
+
+    def _snapshot(self) -> dict[str, Any]:
+        with self._lock:
+            workspaces = [
+                {"id": wid, "runtimes": list(rids)}
+                for wid, rids in self.workspace_runtimes.items()
+            ]
+            agents = list(self.cfg.agents.keys())
+            active = self.active_tasks
+        return {
+            "daemon_id": self.cfg.daemon_id,
+            "device_name": self.cfg.device_name,
+            "server_url": self.cfg.server_base_url,
+            "cli_version": self.cfg.cli_version,
+            "implementation": "python",
+            "platform": platform_label(),
+            "active_task_count": active,
+            "agents": agents,
+            "workspaces": workspaces,
+        }
+
+    def _runtime_ids(self) -> list[str]:
+        with self._lock:
+            return list(self.runtime_index.keys())
+
+    def run(self) -> int:
+        log.info(
+            "starting python daemon version=%s server=%s agents=%s",
+            self.cfg.cli_version,
+            self.cfg.server_base_url,
+            list(self.cfg.agents.keys()),
+        )
+        self._health = HealthServer(
+            self.cfg.health_port,
+            self._snapshot,
+            workspaces_root=self.cfg.workspaces_root,
+            on_shutdown=self.stop,
+            repo_allowlist=self.repo_allowlist,
+            client=self.client,
+        )
+        try:
+            self._health.start()
+        except RuntimeError as e:
+            log.error("%s", e)
+            return 1
+
+        def _sig(_signum: int, _frame: Any) -> None:
+            log.info("signal received, shutting down")
+            self.stop()
+
+        signal.signal(signal.SIGINT, _sig)
+        signal.signal(signal.SIGTERM, _sig)
+        if hasattr(signal, "SIGHUP"):
+            try:
+                signal.signal(signal.SIGHUP, signal.SIG_IGN)
+            except (ValueError, OSError):
+                pass
+
+        try:
+            self._sync_workspaces()
+        except Exception as e:
+            log.error("initial workspace sync failed: %s", e)
+            self._health.stop()
+            return 1
+
+        threads = [
+            threading.Thread(target=self._heartbeat_loop, name="heartbeat", daemon=True),
+            threading.Thread(target=self._sync_loop, name="ws-sync", daemon=True),
+            threading.Thread(target=self._wakeup_loop, name="wakeup", daemon=True),
+            threading.Thread(
+                target=run_gc_loop,
+                kwargs={
+                    "stop": self._stop,
+                    "client": self.client,
+                    "workspaces_root": self.cfg.workspaces_root,
+                    "interval": self.cfg.gc_interval,
+                    "ttl": self.cfg.gc_ttl,
+                    "orphan_ttl": self.cfg.gc_orphan_ttl,
+                    "artifact_ttl": self.cfg.gc_artifact_ttl,
+                    "enabled": self.cfg.gc_enabled,
+                    "active_roots": self.active_env_roots,
+                },
+                name="gc",
+                daemon=True,
+            ),
+        ]
+        for t in threads:
+            t.start()
+
+        try:
+            self._poll_loop()
+        finally:
+            log.info("stopping…")
+            self.stop()
+            ids = list(self.runtime_index.keys())
+            if ids:
+                try:
+                    self.client.deregister(ids)
+                    log.info("deregistered %s runtimes", len(ids))
+                except Exception as e:
+                    log.warning("deregister failed: %s", e)
+            self._pool.shutdown(wait=False, cancel_futures=True)
+            if self._health:
+                self._health.stop()
+        return 0
+
+    def _wakeup_loop(self) -> None:
+        run_wakeup_loop(
+            stop=self._stop,
+            wake=self._wake,
+            get_runtime_ids=self._runtime_ids,
+            base_url=self.cfg.server_base_url,
+            token=self.cfg.token,
+            version=self.cfg.cli_version,
+            hub=self.wakeup_hub,
+        )
+
+    def _sync_loop(self) -> None:
+        while not self._stop.wait(30.0):
+            try:
+                self._sync_workspaces()
+            except Exception as e:
+                log.warning("workspace sync failed: %s", e)
+
+    def _sync_workspaces(self) -> None:
+        workspaces = self.client.list_workspaces()
+        api_ids = {str(w.get("id")) for w in workspaces if w.get("id")}
+        with self._lock:
+            known = set(self.workspace_runtimes.keys())
+
+        for w in workspaces:
+            wid = str(w.get("id") or "")
+            name = str(w.get("name") or wid)
+            if not wid:
+                continue
+            self.workspace_names[wid] = name
+            with self._lock:
+                has = wid in self.workspace_runtimes and self.workspace_runtimes[wid]
+            if has:
+                self._refresh_workspace_repos(wid)
+                continue
+            self._register_workspace(wid, name)
+
+        for wid in list(known - api_ids):
+            with self._lock:
+                rids = self.workspace_runtimes.pop(wid, [])
+                for rid in rids:
+                    self.runtime_index.pop(rid, None)
+            self.repo_allowlist.drop_workspace(wid)
+            log.info("dropped workspace %s", wid)
+
+    def _refresh_workspace_repos(self, workspace_id: str) -> None:
+        try:
+            urls = self.repo_allowlist.refresh_from_api(self.client, workspace_id)
+            log.debug("workspace %s repos allowlist=%s", workspace_id[:8], len(urls))
+        except Exception as e:
+            log.warning("workspace repos refresh %s: %s", workspace_id[:8], e)
+
+    def _register_workspace(self, workspace_id: str, name: str) -> None:
+        now = time.time()
+        with self._lock:
+            until = self._reregister_backoff.get(workspace_id, 0)
+            if now < until:
+                return
+        runtimes = as_register_runtimes(self.cfg)
+        if not runtimes:
+            raise RuntimeError("no agent runtimes to register")
+        try:
+            resp = self.client.register(
+                {
+                    "workspace_id": workspace_id,
+                    "daemon_id": self.cfg.daemon_id,
+                    "legacy_daemon_ids": [],
+                    "device_name": self.cfg.device_name,
+                    "cli_version": self.cfg.cli_version,
+                    "launched_by": "python-daemon",
+                    "timezone": detect_local_timezone(),
+                    "runtimes": runtimes,
+                }
+            )
+        except Exception as e:
+            with self._lock:
+                self._reregister_backoff[workspace_id] = time.time() + 30
+            raise e
+
+        rts = resp.get("runtimes") or []
+        ids: list[str] = []
+        with self._lock:
+            # replace runtime set for workspace
+            old = self.workspace_runtimes.get(workspace_id, [])
+            for oid in old:
+                self.runtime_index.pop(oid, None)
+            for rt in rts:
+                if not isinstance(rt, dict):
+                    continue
+                rid = str(rt.get("id") or "")
+                if not rid:
+                    continue
+                ids.append(rid)
+                self.runtime_index[rid] = rt
+            self.workspace_runtimes[workspace_id] = ids
+            self._reregister_backoff.pop(workspace_id, None)
+
+        for rid in ids:
+            try:
+                self.client.recover_orphans(rid)
+            except Exception as e:
+                log.warning("recover-orphans %s failed: %s", rid, e)
+
+        self._refresh_workspace_repos(workspace_id)
+        log.info(
+            "watching workspace id=%s name=%s runtimes=%s",
+            workspace_id,
+            name,
+            len(ids),
+        )
+        self._wake.set()
+
+    def _predict_env_roots(self, task: Task) -> list[Path]:
+        """Candidate env roots that GC must not reclaim during this run."""
+        roots: list[Path] = []
+        ws = task.workspace_id or "default"
+        predicted = self.cfg.workspaces_root / ws / task.id
+        roots.append(predicted)
+        prior = (task.prior_work_dir or "").strip()
+        if prior:
+            p = Path(prior)
+            roots.append(p.parent if p.name == "workdir" else p)
+        if task.project_id.strip():
+            from .pinned_workdir import extract_local_path, pinned_workdir_path
+
+            local = extract_local_path(task.project_resources)
+            if local:
+                roots.append(Path(local))
+            else:
+                roots.append(
+                    pinned_workdir_path(self.cfg.workspaces_root, ws, task.project_id)
+                )
+        # de-dupe while preserving order
+        seen: set[str] = set()
+        out: list[Path] = []
+        for r in roots:
+            key = str(r)
+            if key in seen:
+                continue
+            seen.add(key)
+            out.append(r)
+        return out
+
+    def _workspace_for_runtime(self, runtime_id: str) -> str | None:
+        with self._lock:
+            for wid, rids in self.workspace_runtimes.items():
+                if runtime_id in rids:
+                    return wid
+        return None
+
+    def _handle_runtime_gone(self, runtime_id: str) -> None:
+        wid = self._workspace_for_runtime(runtime_id)
+        with self._lock:
+            self.runtime_index.pop(runtime_id, None)
+            if wid and wid in self.workspace_runtimes:
+                self.workspace_runtimes[wid] = [
+                    x for x in self.workspace_runtimes[wid] if x != runtime_id
+                ]
+        if not wid:
+            return
+        name = self.workspace_names.get(wid, wid)
+        log.warning("runtime gone %s — re-register workspace %s", runtime_id[:8], wid)
+        try:
+            self._register_workspace(wid, name)
+        except Exception as e:
+            log.warning("re-register failed: %s", e)
+
+    def _heartbeat_loop(self) -> None:
+        while not self._stop.wait(self.cfg.heartbeat_interval):
+            with self._lock:
+                items = list(self.runtime_index.items())
+            for rid, rt in items:
+                if self._stop.is_set():
+                    return
+                try:
+                    resp = self.client.heartbeat(rid)
+                    self._handle_heartbeat_actions(rid, rt, resp)
+                except Exception as e:
+                    if is_not_found(e, "runtime not found"):
+                        self._handle_runtime_gone(rid)
+                    else:
+                        log.warning("heartbeat %s failed: %s", rid[:8], e)
+
+    def _handle_heartbeat_actions(
+        self, runtime_id: str, rt: dict[str, Any], resp: dict[str, Any]
+    ) -> None:
+        # skip pending_update (auto_update n/a)
+        provider = str(rt.get("provider") or rt.get("type") or "")
+        pml = resp.get("pending_model_list") or resp.get("pendingModelList")
+        if isinstance(pml, dict) and pml.get("id"):
+            threading.Thread(
+                target=self._handle_model_list,
+                args=(runtime_id, provider, str(pml["id"])),
+                daemon=True,
+            ).start()
+        pls = resp.get("pending_local_skills") or resp.get("pendingLocalSkills")
+        if isinstance(pls, dict) and pls.get("id"):
+            threading.Thread(
+                target=self._handle_local_skill_list,
+                args=(runtime_id, provider, str(pls["id"])),
+                daemon=True,
+            ).start()
+        pli = resp.get("pending_local_skill_import") or resp.get(
+            "pendingLocalSkillImport"
+        )
+        if isinstance(pli, dict) and pli.get("id"):
+            threading.Thread(
+                target=self._handle_local_skill_import,
+                args=(runtime_id, provider, pli),
+                daemon=True,
+            ).start()
+
+    def _handle_model_list(self, runtime_id: str, provider: str, request_id: str) -> None:
+        entry = self.cfg.agents.get(provider)
+        if not entry:
+            self.client.report_model_list_result(
+                runtime_id,
+                request_id,
+                {"status": "failed", "error": f"no agent for {provider}"},
+            )
+            return
+        models = list_models(provider, entry.path)
+        self.client.report_model_list_result(
+            runtime_id,
+            request_id,
+            {
+                "status": "completed",
+                "models": models,
+                "supported": model_selection_supported(provider),
+            },
+        )
+
+    def _handle_local_skill_list(
+        self, runtime_id: str, provider: str, request_id: str
+    ) -> None:
+        skills, supported = list_runtime_local_skills(provider)
+        self.client.report_local_skill_list_result(
+            runtime_id,
+            request_id,
+            {"status": "completed", "skills": skills, "supported": supported},
+        )
+
+    def _handle_local_skill_import(
+        self, runtime_id: str, provider: str, pending: dict[str, Any]
+    ) -> None:
+        request_id = str(pending.get("id") or "")
+        skill_key = str(pending.get("skill_key") or pending.get("skillKey") or "")
+        bundle, supported, err = load_runtime_local_skill_bundle(provider, skill_key)
+        if err or not bundle:
+            self.client.report_local_skill_import_result(
+                runtime_id,
+                request_id,
+                {
+                    "status": "failed",
+                    "error": err or "import failed",
+                    "supported": supported,
+                },
+            )
+            return
+        self.client.report_local_skill_import_result(
+            runtime_id,
+            request_id,
+            {"status": "completed", "skill": bundle, "supported": supported},
+        )
+
+    def _poll_loop(self) -> None:
+        while not self._stop.is_set():
+            self._claim_round()
+            self._wake.clear()
+            self._wake.wait(timeout=self.cfg.poll_interval)
+
+    def _claim_round(self) -> None:
+        with self._lock:
+            ids = list(self.runtime_index.keys())
+            active = self.active_tasks
+        capacity = self.cfg.max_concurrent_tasks - active
+        if capacity <= 0:
+            return
+        claimed = 0
+        for rid in ids:
+            if claimed >= capacity or self._stop.is_set():
+                break
+            try:
+                task = self.client.claim_task(rid)
+            except Exception as e:
+                log.warning("claim %s failed: %s", rid[:8], e)
+                continue
+            if not task:
+                continue
+            claimed += 1
+            with self._lock:
+                self.active_tasks += 1
+            self._pool.submit(self._handle_task, task)
+
+    def _handle_task(self, task: Task) -> None:
+        try:
+            self._run_one(task)
+        finally:
+            with self._lock:
+                self.active_tasks = max(0, self.active_tasks - 1)
+
+    def _watch_cancel(self, task_id: str, cancel_event: threading.Event) -> None:
+        while not self._stop.is_set() and not cancel_event.is_set():
+            try:
+                status = self.client.get_task_status(task_id)
+                err: Exception | None = None
+            except Exception as e:
+                status = ""
+                err = e
+            if should_interrupt_agent(status, err):
+                cancel_event.set()
+                return
+            time.sleep(5.0)
+
+    def _run_one(self, task: Task) -> None:
+        with self._lock:
+            rt = self.runtime_index.get(task.runtime_id) or {}
+        provider = str(rt.get("provider") or rt.get("type") or "")
+        if not provider:
+            provider = next(iter(self.cfg.agents), "")
+        entry = self.cfg.agents.get(provider)
+        agent_name = (task.agent or {}).get("name") or "agent"
+        log.info(
+            "picked task=%s issue=%s agent=%s provider=%s",
+            task.id[:8],
+            (task.issue_id or "-")[:8],
+            agent_name,
+            provider,
+        )
+
+        try:
+            self.client.start_task(task.id)
+        except Exception as e:
+            log.error("start task failed: %s", e)
+            try:
+                self.client.fail_task(task.id, f"start task failed: {e}")
+            except Exception:
+                pass
+            return
+
+        if not entry:
+            self.client.fail_task(task.id, f"no local agent for provider {provider!r}")
+            return
+
+        try:
+            self.client.report_progress(task.id, f"Launching {provider}", 1, 2)
+        except Exception:
+            pass
+
+        cancel_event = threading.Event()
+        threading.Thread(
+            target=self._watch_cancel,
+            args=(task.id, cancel_event),
+            name=f"cancel-{task.id[:8]}",
+            daemon=True,
+        ).start()
+
+        def on_messages(batch: list[dict]) -> None:
+            try:
+                self.client.report_task_messages(task.id, batch)
+            except Exception as e:
+                log.warning("report messages failed: %s", e)
+
+        def on_session(session_id: str, work_dir: str) -> None:
+            try:
+                self.client.pin_task_session(
+                    task.id, session_id=session_id, work_dir=work_dir
+                )
+            except Exception as e:
+                log.warning("pin session failed: %s", e)
+
+        # Task-scoped repos (e.g. project github_repo) join the checkout allowlist.
+        task_urls = [str(r.get("url")) for r in task.repos if r.get("url")]
+        if task_urls:
+            self.repo_allowlist.register_task_repos(task.workspace_id, task_urls)
+
+        marked_roots = self._predict_env_roots(task)
+        for root in marked_roots:
+            self.active_env_roots.mark(root)
+        result = None
+        try:
+            with self._lock:
+                task_slot = self.active_tasks
+            result = run_agent(
+                task,
+                provider,
+                entry,
+                self.cfg.workspaces_root,
+                self.cfg.agent_timeout,
+                client=self.client,
+                cfg=self.cfg,
+                task_slot=task_slot,
+                cancel_event=cancel_event,
+                on_messages=on_messages,
+                on_session=on_session,
+                acp_pool=self.chat_acp_pool,
+                idle_watchdog=self.cfg.agent_idle_watchdog,
+            )
+        except Exception as e:
+            log.exception("run_agent crashed: %s", e)
+            try:
+                self.client.fail_task(task.id, f"daemon run crashed: {e}")
+            except Exception:
+                pass
+            return
+        finally:
+            for root in marked_roots:
+                self.active_env_roots.unmark(root)
+            cancel_event.set()
+
+        if result is None:
+            return
+
+        if result.usage:
+            try:
+                self.client.report_task_usage(task.id, result.usage)
+            except Exception as e:
+                log.warning("report usage failed: %s", e)
+
+        if result.cancelled:
+            log.info("task %s cancelled, discarding result", task.id[:8])
+            return
+
+        try:
+            status = self.client.get_task_status(task.id)
+            if should_interrupt_agent(status, None):
+                log.info("task %s cancelled, discarding result", task.id[:8])
+                return
+        except APIError as e:
+            if should_interrupt_agent("", e):
+                log.info("task %s deleted, discarding result", task.id[:8])
+                return
+
+        try:
+            self.client.report_progress(task.id, "Finishing task", 2, 2)
+        except Exception:
+            pass
+
+        # poisoned classification
+        fail_reason = "agent_error"
+        if result.ok:
+            poisoned_reason, is_poisoned = classify_poisoned_output(result.output)
+            if is_poisoned:
+                result.ok = False
+                result.error = result.output or poisoned_reason
+                fail_reason = poisoned_reason
+        else:
+            poisoned_reason, is_poisoned = classify_poisoned_error(result.error)
+            if is_poisoned:
+                fail_reason = poisoned_reason
+
+        if result.ok:
+            try:
+                auto_commit_and_push(
+                    result.work_dir, task.id, str(agent_name)
+                )
+            except Exception as e:
+                log.warning("auto-commit: %s", e)
+            try:
+                report_issue_artifacts(
+                    result.work_dir,
+                    lambda arts: self.client.report_task_artifacts(task.id, arts),
+                )
+            except Exception as e:
+                log.warning("artifacts: %s", e)
+            try:
+                self.client.complete_task(
+                    task.id,
+                    result.output[:100_000],
+                    work_dir=result.work_dir,
+                    session_id=result.session_id,
+                )
+                log.info("task %s completed", task.id[:8])
+            except Exception as e:
+                log.error("complete task failed: %s", e)
+        else:
+            try:
+                self.client.fail_task(
+                    task.id,
+                    result.error or "agent failed",
+                    work_dir=result.work_dir,
+                    session_id=result.session_id,
+                    failure_reason=fail_reason,
+                )
+                log.info("task %s failed: %s", task.id[:8], result.error)
+            except Exception as e:
+                log.error("fail task callback failed: %s", e)
+
+        # GC meta for non-pinned task roots
+        if result.work_dir:
+            env_root = Path(result.work_dir)
+            if env_root.name == "workdir":
+                env_root = env_root.parent
+            if "projects" not in env_root.parts:
+                write_gc_meta(
+                    env_root,
+                    gc_meta_for_task(
+                        task.id,
+                        task.issue_id,
+                        task.chat_session_id,
+                        task.autopilot_run_id,
+                        workspace_id=task.workspace_id,
+                        quick_create=bool(task.quick_create_prompt),
+                    ),
+                )