332 lines
14 KiB
Python
332 lines
14 KiB
Python
"""Co4E run manager — tracks concurrent flow runs and their live status.
|
|
|
|
The app already runs several Cowork/Code turns at once (each on its own
|
|
``AgentWorker`` QThread). This brings the same to Co4E: any number of flows can
|
|
run in parallel, each on its own worker, with live per-run status
|
|
(running / done / error / stopped) and progress (done / total steps).
|
|
|
|
The manager is the single source of truth — the Co4E "Running flows" list reads
|
|
it and refreshes on every ``changed`` signal (and when the tab is re-shown), so
|
|
switching sub-tabs never loses, stales, or drops status. ``event`` re-emits each
|
|
run's node-level events tagged with the run id, so the canvas/chat can mirror
|
|
the run that is currently open.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional
|
|
|
|
from PySide6.QtCore import QObject, Signal
|
|
|
|
from .co4e import STEP_DONE, STEP_ERROR, STEP_PLANNED, Workflow
|
|
|
|
_TERMINAL_NODE = {STEP_DONE, STEP_ERROR, STEP_PLANNED}
|
|
_HISTORY_CAP = 500 # keep the most-recent N runs on disk
|
|
|
|
|
|
def _now_str() -> str:
|
|
from datetime import datetime
|
|
return datetime.now().strftime("%Y-%m-%d %H:%M")
|
|
|
|
|
|
def _current_user() -> str:
|
|
"""Best-effort creator name for a run (signed-in MS365 identity → OS user)."""
|
|
import os
|
|
return (os.environ.get("USERNAME") or os.environ.get("USER") or "you")
|
|
|
|
|
|
class RunHandle:
|
|
"""Live state for one flow run. Mutated by the manager as events arrive."""
|
|
|
|
def __init__(self, run_id: str, wf_id: str, name: str, total: int,
|
|
plan_mode: bool, manual: bool, created_by: str = "", created_at: str = "",
|
|
project_id: str = ""):
|
|
self.id = run_id
|
|
self.wf_id = wf_id
|
|
self.name = name
|
|
self.project_id = project_id # workspace this run belongs to (Flow Status is per-project)
|
|
self.total = max(0, total)
|
|
self.done = 0
|
|
self.status = "running" # running | done | error | stopped
|
|
self.plan_mode = plan_mode
|
|
self.manual = manual
|
|
self.created_by = created_by
|
|
self.created_at = created_at
|
|
self.error = ""
|
|
self.node_status: Dict[str, str] = {}
|
|
self.worker = None
|
|
self.wf = None # the Workflow this run ran — lets Runs reopen it
|
|
# even after its tab was closed / it was never saved
|
|
self.out_dir = "" # workspace folder this run wrote its files into
|
|
|
|
@property
|
|
def running(self) -> bool:
|
|
return self.status == "running"
|
|
|
|
def progress_text(self) -> str:
|
|
return f"{self.done}/{self.total}" if self.total else self.status
|
|
|
|
# ---- persistence ------------------------------------------------------
|
|
def to_record(self) -> dict:
|
|
"""Serialize for the on-disk run history. The workflow snapshot is kept
|
|
so a past run can be reopened even if its saved flow was later edited or
|
|
deleted."""
|
|
from .co4e import workflow_to_dict
|
|
return {
|
|
"id": self.id, "wf_id": self.wf_id, "name": self.name,
|
|
"total": self.total, "done": self.done, "status": self.status,
|
|
"plan_mode": self.plan_mode, "manual": self.manual,
|
|
"created_by": self.created_by, "created_at": self.created_at,
|
|
"error": self.error, "node_status": dict(self.node_status),
|
|
"wf": workflow_to_dict(self.wf) if self.wf is not None else None,
|
|
"out_dir": self.out_dir, "project_id": self.project_id,
|
|
}
|
|
|
|
@classmethod
|
|
def from_record(cls, rec: dict) -> "RunHandle":
|
|
from .co4e import workflow_from_dict
|
|
rec = dict(rec or {})
|
|
h = cls(str(rec.get("id", "")), str(rec.get("wf_id", "")),
|
|
rec.get("name", ""), int(rec.get("total", 0) or 0),
|
|
bool(rec.get("plan_mode")), bool(rec.get("manual")),
|
|
created_by=rec.get("created_by", ""), created_at=rec.get("created_at", ""))
|
|
h.done = int(rec.get("done", 0) or 0)
|
|
h.status = rec.get("status", "done")
|
|
# a run persisted as "running" means the app closed mid-run — its worker
|
|
# is gone, so it's no longer live: settle it as "stopped".
|
|
if h.status == "running":
|
|
h.status = "stopped"
|
|
h.error = rec.get("error", "")
|
|
h.node_status = dict(rec.get("node_status") or {})
|
|
h.out_dir = rec.get("out_dir", "")
|
|
h.project_id = rec.get("project_id", "")
|
|
wfd = rec.get("wf")
|
|
h.wf = workflow_from_dict(wfd) if wfd else None
|
|
return h
|
|
|
|
|
|
class Co4ERunManager(QObject):
|
|
changed = Signal() # any run's status/progress changed → refresh views
|
|
event = Signal(str, dict) # (run_id, ev) — node-level events, for mirroring
|
|
|
|
def __init__(self, ctx):
|
|
super().__init__()
|
|
self.ctx = ctx
|
|
self._runs: Dict[str, RunHandle] = {}
|
|
self._seq = 0
|
|
self._output_root: Optional[Path] = None # active workspace's co4e output base
|
|
self._project_id: str = "" # active workspace — Flow Status is filtered to it
|
|
self._load_history() # restore past runs so the Flow Status tab
|
|
# keeps its full history across restarts
|
|
# every status/progress change is persisted, so history is never lost
|
|
self.changed.connect(self._save_history)
|
|
|
|
# ---- persistence ------------------------------------------------------
|
|
def _history_path(self) -> Path:
|
|
from .co4e import CO4E_DIR
|
|
return CO4E_DIR / "run_history.json"
|
|
|
|
def _load_history(self) -> None:
|
|
path = self._history_path()
|
|
try:
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
except (OSError, ValueError):
|
|
return
|
|
max_seq = 0
|
|
for rec in data.get("runs", []):
|
|
try:
|
|
handle = RunHandle.from_record(rec)
|
|
except Exception:
|
|
continue
|
|
if not handle.id:
|
|
continue
|
|
self._runs[handle.id] = handle
|
|
if handle.id.startswith("run") and handle.id[3:].isdigit():
|
|
max_seq = max(max_seq, int(handle.id[3:]))
|
|
self._seq = max_seq # avoid minting ids that collide with history
|
|
|
|
def _save_history(self) -> None:
|
|
path = self._history_path()
|
|
runs = list(self._runs.values())[-_HISTORY_CAP:]
|
|
payload = {"runs": [h.to_record() for h in runs]}
|
|
try:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = path.with_suffix(".json.tmp")
|
|
tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2),
|
|
encoding="utf-8")
|
|
tmp.replace(path) # atomic — never leaves a half-written file
|
|
except OSError:
|
|
pass
|
|
|
|
# ---- lifecycle --------------------------------------------------------
|
|
def _next_id(self) -> str:
|
|
self._seq += 1
|
|
return f"run{self._seq}"
|
|
|
|
def start(self, wf: Workflow, *, skill_map: Optional[Dict[str, str]] = None,
|
|
plan_mode: bool = False, only_nodes: Optional[set] = None,
|
|
seed_outputs: Optional[Dict[str, str]] = None,
|
|
manual: bool = False, label: Optional[str] = None) -> str:
|
|
"""Launch a flow (or a subset via ``only_nodes``) on its own worker and
|
|
return the new run id. Runs concurrently with every other active run."""
|
|
from . import co4e_runner
|
|
from .worker import AgentWorker
|
|
|
|
import copy as _copy
|
|
|
|
run_id = self._next_id()
|
|
total = len(only_nodes) if only_nodes else len(wf.nodes)
|
|
handle = RunHandle(run_id, wf.id, label or wf.name, total, plan_mode, manual,
|
|
created_by=_current_user(), created_at=_now_str(),
|
|
project_id=self._project_id)
|
|
# Keep a DEEP COPY so Runs can reopen the flow exactly as it ran — even if
|
|
# the live canvas is later edited (nodes moved, edges removed) while the
|
|
# run is still tracked. A shared reference caused reopened runs to show a
|
|
# disconnected/blank graph.
|
|
handle.wf = _copy.deepcopy(wf)
|
|
nodes = list(wf.nodes)
|
|
edges = list(wf.edges)
|
|
out_dir = self._out_dir(wf)
|
|
handle.out_dir = str(out_dir) # where this run writes its files (workspace)
|
|
ctx = self.ctx
|
|
sk = dict(skill_map or {})
|
|
only = set(only_nodes) if only_nodes else None
|
|
seed = dict(seed_outputs or {})
|
|
|
|
run_label = handle.name
|
|
def job(worker: AgentWorker):
|
|
return co4e_runner.run_workflow(
|
|
ctx, nodes, edges, out_dir, worker.emit_event, worker.is_cancelled,
|
|
plan_mode=plan_mode, skill_map=sk, only_nodes=only, seed_outputs=seed,
|
|
usage_label=run_label)
|
|
|
|
worker = AgentWorker(job)
|
|
handle.worker = worker
|
|
worker.event.connect(lambda ev, rid=run_id: self._on_event(rid, ev))
|
|
worker.finished_ok.connect(lambda _r, rid=run_id: self._on_finished(rid))
|
|
worker.failed.connect(lambda e, rid=run_id: self._on_failed(rid, e))
|
|
self._runs[run_id] = handle
|
|
worker.start()
|
|
self.changed.emit()
|
|
return run_id
|
|
|
|
# ---- worker callbacks -------------------------------------------------
|
|
def _on_event(self, run_id: str, ev: dict) -> None:
|
|
handle = self._runs.get(run_id)
|
|
if handle is not None and isinstance(ev, dict):
|
|
t = ev.get("type")
|
|
if t == "node_status":
|
|
handle.node_status[ev.get("node_id")] = ev.get("status")
|
|
handle.done = sum(1 for s in handle.node_status.values() if s in _TERMINAL_NODE)
|
|
self.changed.emit()
|
|
elif t == "run_done":
|
|
if handle.status == "running":
|
|
handle.status = "done" if ev.get("ok", True) else "error"
|
|
self.changed.emit()
|
|
self.event.emit(run_id, ev)
|
|
|
|
def _on_finished(self, run_id: str) -> None:
|
|
handle = self._runs.get(run_id)
|
|
if handle is not None and handle.status == "running":
|
|
# job returned without a run_done event (shouldn't happen) — settle it
|
|
handle.status = "done"
|
|
self.changed.emit()
|
|
|
|
def _on_failed(self, run_id: str, err: str) -> None:
|
|
handle = self._runs.get(run_id)
|
|
if handle is not None:
|
|
handle.status = "error"
|
|
handle.error = str(err)
|
|
self.event.emit(run_id, {"type": "run_error", "error": str(err)})
|
|
self.changed.emit()
|
|
|
|
# ---- control ----------------------------------------------------------
|
|
def stop(self, run_id: str) -> None:
|
|
handle = self._runs.get(run_id)
|
|
if handle is not None and handle.worker is not None and handle.running:
|
|
handle.worker.request_stop()
|
|
handle.status = "stopped"
|
|
self.changed.emit()
|
|
|
|
def stop_all(self) -> None:
|
|
# Only the CURRENT workspace's runs (Flow Status is per-project).
|
|
for run_id in [r for r, h in self._runs.items() if self._belongs(h)]:
|
|
self.stop(run_id)
|
|
|
|
def rename(self, run_id: str, new_name: str) -> None:
|
|
"""Rename a run in the Flow Status history (and its kept workflow snapshot),
|
|
then persist + refresh views. No-op on a blank name / unknown run."""
|
|
handle = self._runs.get(run_id)
|
|
new_name = (new_name or "").strip()
|
|
if handle is None or not new_name or new_name == handle.name:
|
|
return
|
|
handle.name = new_name
|
|
if handle.wf is not None:
|
|
handle.wf.name = new_name
|
|
self.changed.emit()
|
|
|
|
def remove(self, run_id: str) -> None:
|
|
handle = self._runs.get(run_id)
|
|
if handle is not None and handle.running:
|
|
self.stop(run_id)
|
|
self._runs.pop(run_id, None)
|
|
self.changed.emit()
|
|
|
|
def clear_finished(self) -> None:
|
|
# Only clear finished runs of the CURRENT workspace.
|
|
for run_id in [r for r, h in self._runs.items() if not h.running and self._belongs(h)]:
|
|
self._runs.pop(run_id, None)
|
|
self.changed.emit()
|
|
|
|
# ---- queries ----------------------------------------------------------
|
|
def _belongs(self, h: RunHandle) -> bool:
|
|
"""Whether a run belongs to the currently-selected workspace."""
|
|
return getattr(h, "project_id", "") == self._project_id
|
|
|
|
def runs(self) -> List[RunHandle]:
|
|
"""Runs of the CURRENT workspace only — Flow Status is per-project."""
|
|
return [h for h in self._runs.values() if self._belongs(h)]
|
|
|
|
def all_runs(self) -> List[RunHandle]:
|
|
"""Every tracked run across all workspaces (background tracking)."""
|
|
return list(self._runs.values())
|
|
|
|
def get(self, run_id: str) -> Optional[RunHandle]:
|
|
return self._runs.get(run_id)
|
|
|
|
def active_count(self) -> int:
|
|
return sum(1 for h in self._runs.values() if h.running and self._belongs(h))
|
|
|
|
def set_current_project(self, project_id: str) -> None:
|
|
"""Filter Flow Status (and new runs) to this workspace. Runs started while
|
|
this is set are tagged with it; the Runs view shows only matching runs."""
|
|
pid = project_id or ""
|
|
if pid != self._project_id:
|
|
self._project_id = pid
|
|
self.changed.emit() # re-render Flow Status for the new workspace
|
|
|
|
def set_output_root(self, root: Optional[Path]) -> None:
|
|
"""Point flow outputs at the SELECTED workspace's co4e folder (set by the
|
|
Co4E tab when a project is chosen). ``None`` → fall back to the global
|
|
Cowork output dir."""
|
|
self._output_root = Path(root) if root else None
|
|
|
|
def _out_dir(self, wf: Workflow) -> Path:
|
|
# Flow deliverables are written into the SELECTED workspace (the active
|
|
# project's folder) so they land where the user works with files (Folder
|
|
# tab), not in the config/install folder. One subfolder per flow keeps
|
|
# runs tidy. Falls back to the global Cowork output dir when no workspace
|
|
# is selected.
|
|
from .co4e import slugify
|
|
base = self._output_root
|
|
if base is None:
|
|
try:
|
|
base = self.ctx.config.cowork_output_dir() / "co4e"
|
|
except Exception: # noqa: BLE001 - fall back to the config dir if unavailable
|
|
from .co4e import CO4E_DIR
|
|
base = CO4E_DIR / "runs" / "co4e"
|
|
d = Path(base) / slugify(wf.name or "flow")
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
return d
|