pythinker-code 2.0.0__py3-none-any.whl

pythinker_code/background/models.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+import time
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+type TaskKind = Literal["bash", "agent"]
+type TaskStatus = Literal[
+    "created",
+    "starting",
+    "running",
+    "awaiting_approval",
+    "completed",
+    "failed",
+    "killed",
+    "lost",
+]
+type TaskOwnerRole = Literal["root", "subagent"]
+
+TERMINAL_TASK_STATUSES: tuple[TaskStatus, ...] = ("completed", "failed", "killed", "lost")
+
+
+def is_terminal_status(status: TaskStatus) -> bool:
+    return status in TERMINAL_TASK_STATUSES
+
+
+class TaskSpec(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    version: int = 1
+    id: str
+    kind: TaskKind
+    session_id: str
+    description: str
+    tool_call_id: str
+    owner_role: TaskOwnerRole = "root"
+    created_at: float = Field(default_factory=time.time)
+
+    @field_validator("owner_role", mode="before")
+    @classmethod
+    def _normalize_owner_role(cls, v: str) -> str:
+        if v in ("fixed_subagent", "dynamic_subagent"):
+            return "subagent"
+        return v
+
+    # Bash-specific fields for V1. Future task types can use kind_payload.
+    command: str | None = None
+    shell_name: str | None = None
+    shell_path: str | None = None
+    cwd: str | None = None
+    timeout_s: int | None = None
+    kind_payload: dict[str, Any] | None = None
+
+
+class TaskRuntime(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    status: TaskStatus = "created"
+    worker_pid: int | None = None
+    child_pid: int | None = None
+    child_pgid: int | None = None
+    started_at: float | None = None
+    heartbeat_at: float | None = None
+    updated_at: float = Field(default_factory=time.time)
+    finished_at: float | None = None
+    exit_code: int | None = None
+    interrupted: bool = False
+    timed_out: bool = False
+    failure_reason: str | None = None
+
+
+class TaskControl(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    kill_requested_at: float | None = None
+    kill_reason: str | None = None
+    force: bool = False
+
+
+class TaskConsumerState(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    last_seen_output_size: int = 0
+    last_viewed_at: float | None = None
+
+
+class TaskView(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    spec: TaskSpec
+    runtime: TaskRuntime
+    control: TaskControl
+    consumer: TaskConsumerState
+
+
+class TaskOutputChunk(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    task_id: str
+    offset: int
+    next_offset: int
+    text: str
+    eof: bool
+    status: TaskStatus

pythinker_code/background/store.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+import os
+import re
+from pathlib import Path
+
+from pydantic import BaseModel, ValidationError
+
+from pythinker_code.utils.io import atomic_json_write
+from pythinker_code.utils.logging import logger
+
+from .models import (
+    TaskConsumerState,
+    TaskControl,
+    TaskOutputChunk,
+    TaskRuntime,
+    TaskSpec,
+    TaskStatus,
+    TaskView,
+)
+
+_VALID_TASK_ID = re.compile(r"^[a-z0-9][a-z0-9\-]{1,24}$")
+
+
+def _validate_task_id(task_id: str) -> None:
+    if not _VALID_TASK_ID.match(task_id):
+        raise ValueError(f"Invalid task_id: {task_id!r}")
+
+
+class BackgroundTaskStore:
+    SPEC_FILE = "spec.json"
+    RUNTIME_FILE = "runtime.json"
+    CONTROL_FILE = "control.json"
+    CONSUMER_FILE = "consumer.json"
+    OUTPUT_FILE = "output.log"
+
+    def __init__(self, root: Path):
+        self._root = root
+
+    @property
+    def root(self) -> Path:
+        return self._root
+
+    def _ensure_root(self) -> Path:
+        """Return the root directory, creating it if it does not exist."""
+        self._root.mkdir(parents=True, exist_ok=True)
+        return self._root
+
+    def task_dir(self, task_id: str) -> Path:
+        _validate_task_id(task_id)
+        path = self._ensure_root() / task_id
+        path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    def task_path(self, task_id: str) -> Path:
+        _validate_task_id(task_id)
+        return self.root / task_id
+
+    def spec_path(self, task_id: str) -> Path:
+        return self.task_path(task_id) / self.SPEC_FILE
+
+    def runtime_path(self, task_id: str) -> Path:
+        return self.task_path(task_id) / self.RUNTIME_FILE
+
+    def control_path(self, task_id: str) -> Path:
+        return self.task_path(task_id) / self.CONTROL_FILE
+
+    def consumer_path(self, task_id: str) -> Path:
+        return self.task_path(task_id) / self.CONSUMER_FILE
+
+    def output_path(self, task_id: str) -> Path:
+        return self.task_path(task_id) / self.OUTPUT_FILE
+
+    def create_task(self, spec: TaskSpec) -> None:
+        task_dir = self.task_dir(spec.id)
+        atomic_json_write(spec.model_dump(mode="json"), task_dir / self.SPEC_FILE)
+        atomic_json_write(TaskRuntime().model_dump(mode="json"), task_dir / self.RUNTIME_FILE)
+        atomic_json_write(TaskControl().model_dump(mode="json"), task_dir / self.CONTROL_FILE)
+        atomic_json_write(
+            TaskConsumerState().model_dump(mode="json"),
+            task_dir / self.CONSUMER_FILE,
+        )
+        self.output_path(spec.id).touch(exist_ok=True)
+
+    def list_task_ids(self) -> list[str]:
+        if not self.root.exists():
+            return []
+        task_ids: list[str] = []
+        for path in sorted(self.root.iterdir()):
+            if not path.is_dir():
+                continue
+            if not (path / self.SPEC_FILE).exists():
+                continue
+            task_ids.append(path.name)
+        return task_ids
+
+    def write_spec(self, spec: TaskSpec) -> None:
+        atomic_json_write(spec.model_dump(mode="json"), self.spec_path(spec.id))
+
+    def read_spec(self, task_id: str) -> TaskSpec:
+        return TaskSpec.model_validate_json(self.spec_path(task_id).read_text(encoding="utf-8"))
+
+    def write_runtime(self, task_id: str, runtime: TaskRuntime) -> None:
+        atomic_json_write(runtime.model_dump(mode="json"), self.runtime_path(task_id))
+
+    def read_runtime(self, task_id: str) -> TaskRuntime:
+        path = self.runtime_path(task_id)
+        if not path.exists():
+            return TaskRuntime()
+        return _read_json_model(
+            path,
+            TaskRuntime,
+            fallback=TaskRuntime(updated_at=0),
+            artifact="task runtime",
+        )
+
+    def write_control(self, task_id: str, control: TaskControl) -> None:
+        atomic_json_write(control.model_dump(mode="json"), self.control_path(task_id))
+
+    def read_control(self, task_id: str) -> TaskControl:
+        path = self.control_path(task_id)
+        if not path.exists():
+            return TaskControl()
+        return _read_json_model(
+            path,
+            TaskControl,
+            fallback=TaskControl(),
+            artifact="task control",
+        )
+
+    def write_consumer(self, task_id: str, consumer: TaskConsumerState) -> None:
+        atomic_json_write(consumer.model_dump(mode="json"), self.consumer_path(task_id))
+
+    def read_consumer(self, task_id: str) -> TaskConsumerState:
+        path = self.consumer_path(task_id)
+        if not path.exists():
+            return TaskConsumerState()
+        return _read_json_model(
+            path,
+            TaskConsumerState,
+            fallback=TaskConsumerState(),
+            artifact="task consumer state",
+        )
+
+    def merged_view(self, task_id: str) -> TaskView:
+        return TaskView(
+            spec=self.read_spec(task_id),
+            runtime=self.read_runtime(task_id),
+            control=self.read_control(task_id),
+            consumer=self.read_consumer(task_id),
+        )
+
+    def list_views(self) -> list[TaskView]:
+        views: list[TaskView] = []
+        for task_id in self.list_task_ids():
+            try:
+                views.append(self.merged_view(task_id))
+            except (OSError, ValidationError, ValueError, UnicodeDecodeError) as exc:
+                logger.warning(
+                    "Skipping invalid background task {task_id} from {path}: {error}",
+                    task_id=task_id,
+                    path=self.root / task_id / self.SPEC_FILE,
+                    error=exc,
+                )
+        views.sort(
+            key=lambda view: view.runtime.updated_at or view.spec.created_at,
+            reverse=True,
+        )
+        return views
+
+    def read_output(
+        self,
+        task_id: str,
+        offset: int,
+        max_bytes: int,
+        *,
+        status: TaskStatus,
+        path_override: Path | None = None,
+    ) -> TaskOutputChunk:
+        path = path_override if path_override is not None else self.output_path(task_id)
+        if not path.exists():
+            return TaskOutputChunk(
+                task_id=task_id,
+                offset=offset,
+                next_offset=offset,
+                text="",
+                eof=True,
+                status=status,
+            )
+
+        with path.open("rb") as f:
+            f.seek(0, os.SEEK_END)
+            total_size = f.tell()
+            bounded_offset = min(max(offset, 0), total_size)
+            f.seek(bounded_offset)
+            content = f.read(max_bytes)
+
+        next_offset = bounded_offset + len(content)
+        return TaskOutputChunk(
+            task_id=task_id,
+            offset=bounded_offset,
+            next_offset=next_offset,
+            text=content.decode("utf-8", errors="replace"),
+            eof=next_offset >= total_size,
+            status=status,
+        )
+
+    def tail_output(self, task_id: str, max_bytes: int, max_lines: int) -> str:
+        path = self.output_path(task_id)
+        if not path.exists():
+            return ""
+
+        with path.open("rb") as f:
+            f.seek(0, os.SEEK_END)
+            total_size = f.tell()
+            start = max(0, total_size - max_bytes)
+            f.seek(start)
+            content = f.read()
+
+        text = content.decode("utf-8", errors="replace")
+        lines = text.splitlines()
+        if len(lines) > max_lines:
+            lines = lines[-max_lines:]
+        return "\n".join(lines)
+
+
+def _read_json_model[T: BaseModel](path: Path, model: type[T], *, fallback: T, artifact: str) -> T:
+    try:
+        return model.model_validate_json(path.read_text(encoding="utf-8"))
+    except (OSError, ValidationError, ValueError, UnicodeDecodeError) as exc:
+        logger.warning(
+            "Failed to read {artifact} from {path}; using defaults: {error}",
+            artifact=artifact,
+            path=path,
+            error=exc,
+        )
+        return fallback

pythinker_code/background/summary.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from .manager import BackgroundTaskManager
+from .models import TaskView, is_terminal_status
+
+
+def list_task_views(
+    manager: BackgroundTaskManager,
+    *,
+    active_only: bool = True,
+    limit: int = 20,
+) -> list[TaskView]:
+    views = manager.list_tasks(limit=None)
+    if active_only:
+        views = [view for view in views if not is_terminal_status(view.runtime.status)]
+    return views[:limit]
+
+
+def format_task(view: TaskView, *, include_command: bool = False) -> str:
+    lines = [
+        f"task_id: {view.spec.id}",
+        f"kind: {view.spec.kind}",
+        f"status: {view.runtime.status}",
+        f"description: {view.spec.description}",
+    ]
+    if view.spec.kind == "agent" and view.spec.kind_payload:
+        if agent_id := view.spec.kind_payload.get("agent_id"):
+            lines.append(f"agent_id: {agent_id}")
+        if subagent_type := view.spec.kind_payload.get("subagent_type"):
+            lines.append(f"subagent_type: {subagent_type}")
+    if include_command and view.spec.command:
+        lines.append(f"command: {view.spec.command}")
+    if view.runtime.exit_code is not None:
+        lines.append(f"exit_code: {view.runtime.exit_code}")
+    if view.runtime.failure_reason:
+        lines.append(f"reason: {view.runtime.failure_reason}")
+    return "\n".join(lines)
+
+
+def format_task_list(
+    views: list[TaskView],
+    *,
+    active_only: bool = True,
+    include_command: bool = True,
+) -> str:
+    header = "active_background_tasks" if active_only else "background_tasks"
+    if not views:
+        return f"{header}: 0\n[no tasks]"
+
+    lines = [f"{header}: {len(views)}", ""]
+    for index, view in enumerate(views, start=1):
+        lines.extend([f"[{index}]", format_task(view, include_command=include_command), ""])
+    return "\n".join(lines).rstrip()
+
+
+def build_active_task_snapshot(manager: BackgroundTaskManager, *, limit: int = 20) -> str | None:
+    views = list_task_views(manager, active_only=True, limit=limit)
+    if not views:
+        return None
+    return "\n".join(
+        [
+            "<active-background-tasks>",
+            format_task_list(views, active_only=True, include_command=False),
+            "</active-background-tasks>",
+        ]
+    )

pythinker_code/background/worker.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import os
+import signal
+import subprocess
+import time
+from pathlib import Path
+from typing import Any
+
+from pythinker_code.utils.logging import logger
+from pythinker_code.utils.subprocess_env import get_clean_env
+
+from .models import TaskControl
+from .store import BackgroundTaskStore
+
+
+def terminate_process_tree_windows(pid: int, *, force: bool) -> None:
+    args = ["taskkill", "/PID", str(pid), "/T"]
+    if force:
+        args.append("/F")
+    subprocess.run(
+        args,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        check=False,
+    )
+
+
+async def run_background_task_worker(
+    task_dir: Path,
+    *,
+    heartbeat_interval_ms: int = 5000,
+    control_poll_interval_ms: int = 500,
+    kill_grace_period_ms: int = 2000,
+) -> None:
+    task_dir = task_dir.expanduser().resolve()
+    task_id = task_dir.name
+    store = BackgroundTaskStore(task_dir.parent)
+    spec = store.read_spec(task_id)
+    runtime = store.read_runtime(task_id)
+
+    runtime.status = "starting"
+    runtime.worker_pid = os.getpid()
+    runtime.started_at = time.time()
+    runtime.heartbeat_at = runtime.started_at
+    runtime.updated_at = runtime.started_at
+    store.write_runtime(task_id, runtime)
+
+    control = store.read_control(task_id)
+    if control.kill_requested_at is not None:
+        runtime.status = "killed"
+        runtime.interrupted = True
+        runtime.finished_at = time.time()
+        runtime.updated_at = runtime.finished_at
+        runtime.failure_reason = control.kill_reason or "Killed before command start"
+        store.write_runtime(task_id, runtime)
+        return
+
+    if spec.command is None or spec.shell_path is None or spec.cwd is None:
+        runtime.status = "failed"
+        runtime.finished_at = time.time()
+        runtime.updated_at = runtime.finished_at
+        runtime.failure_reason = "Task spec is incomplete for bash worker"
+        store.write_runtime(task_id, runtime)
+        return
+
+    process: asyncio.subprocess.Process | None = None
+    control_task: asyncio.Task[None] | None = None
+    heartbeat_task: asyncio.Task[None] | None = None
+    stop_event = asyncio.Event()
+    kill_sent_at: float | None = None
+    timed_out = False
+    timeout_reason: str | None = None
+
+    async def _heartbeat_loop() -> None:
+        while not stop_event.is_set():
+            await asyncio.sleep(heartbeat_interval_ms / 1000)
+            current = store.read_runtime(task_id)
+            if current.finished_at is not None:
+                return
+            current.heartbeat_at = time.time()
+            current.updated_at = current.heartbeat_at
+            store.write_runtime(task_id, current)
+
+    async def _terminate_process(force: bool = False) -> None:
+        nonlocal kill_sent_at
+        if process is None or process.returncode is not None:
+            return
+        kill_sent_at = kill_sent_at or time.time()
+
+        try:
+            if os.name == "nt":
+                terminate_process_tree_windows(process.pid, force=force)
+                return
+
+            target_pgid = process.pid
+            if force:
+                os.killpg(target_pgid, signal.SIGKILL)
+            else:
+                os.killpg(target_pgid, signal.SIGTERM)
+        except ProcessLookupError:
+            pass
+
+    async def _control_loop() -> None:
+        nonlocal kill_sent_at
+        while not stop_event.is_set():
+            await asyncio.sleep(control_poll_interval_ms / 1000)
+            current_control: TaskControl = store.read_control(task_id)
+            if current_control.kill_requested_at is not None:
+                await _terminate_process(force=current_control.force)
+                if (
+                    kill_sent_at is not None
+                    and process is not None
+                    and process.returncode is None
+                    and time.time() - kill_sent_at >= kill_grace_period_ms / 1000
+                ):
+                    await _terminate_process(force=True)
+
+    try:
+        output_path = store.output_path(task_id)
+        with output_path.open("ab") as output_file:
+            spawn_kwargs: dict[str, Any] = {
+                "stdin": subprocess.DEVNULL,
+                "stdout": output_file,
+                "stderr": output_file,
+                "cwd": spec.cwd,
+                "env": get_clean_env(),
+            }
+            if os.name == "nt":
+                spawn_kwargs["creationflags"] = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
+            else:
+                spawn_kwargs["start_new_session"] = True
+
+            args = (
+                (spec.shell_path, "-command", spec.command)
+                if spec.shell_name == "Windows PowerShell"
+                else (spec.shell_path, "-c", spec.command)
+            )
+            process = await asyncio.create_subprocess_exec(*args, **spawn_kwargs)
+
+            runtime = store.read_runtime(task_id)
+            runtime.status = "running"
+            runtime.child_pid = process.pid
+            runtime.child_pgid = process.pid if os.name != "nt" else None
+            runtime.updated_at = time.time()
+            runtime.heartbeat_at = runtime.updated_at
+            store.write_runtime(task_id, runtime)
+            last_known_runtime = runtime
+
+            heartbeat_task = asyncio.create_task(_heartbeat_loop())
+            control_task = asyncio.create_task(_control_loop())
+            if spec.timeout_s is None:
+                returncode = await process.wait()
+            else:
+                try:
+                    returncode = await asyncio.wait_for(process.wait(), timeout=spec.timeout_s)
+                except TimeoutError:
+                    timed_out = True
+                    timeout_reason = f"Command timed out after {spec.timeout_s}s"
+                    await _terminate_process(force=False)
+                    try:
+                        returncode = await asyncio.wait_for(
+                            process.wait(),
+                            timeout=kill_grace_period_ms / 1000,
+                        )
+                    except TimeoutError:
+                        await _terminate_process(force=True)
+                        returncode = await process.wait()
+    except Exception as exc:
+        logger.exception("Background task worker failed")
+        runtime = store.read_runtime(task_id)
+        runtime.status = "failed"
+        runtime.finished_at = time.time()
+        runtime.updated_at = runtime.finished_at
+        runtime.failure_reason = str(exc)
+        store.write_runtime(task_id, runtime)
+        return
+    finally:
+        stop_event.set()
+        for task in (heartbeat_task, control_task):
+            if task is not None:
+                task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await task
+
+    runtime = last_known_runtime.model_copy()
+    control = store.read_control(task_id)
+    runtime.finished_at = time.time()
+    runtime.updated_at = runtime.finished_at
+    runtime.exit_code = returncode
+    runtime.heartbeat_at = runtime.finished_at
+    if timed_out:
+        runtime.status = "failed"
+        runtime.interrupted = True
+        runtime.timed_out = True
+        runtime.failure_reason = timeout_reason
+    elif control.kill_requested_at is not None:
+        runtime.status = "killed"
+        runtime.interrupted = True
+        runtime.failure_reason = control.kill_reason or "Killed"
+    elif returncode == 0:
+        runtime.status = "completed"
+        runtime.failure_reason = None
+    else:
+        runtime.status = "failed"
+        runtime.failure_reason = f"Command failed with exit code {returncode}"
+    store.write_runtime(task_id, runtime)