crabcode 0.1.6__py3-none-win_amd64.whl

crabcode_gateway/routes/snapshot.py

@@ -0,0 +1,211 @@
+"""Snapshot and revert routes — /snapshot/*."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from fastapi import APIRouter, HTTPException, Request
+
+from crabcode_gateway.schemas import CheckpointRequest, RevertRequest
+from crabcode_gateway.task_registry import SessionOperationRejected, run_session_operation
+
+router = APIRouter(prefix="/snapshot", tags=["snapshot"])
+
+
+async def _checkpoint(session: Any, label: str) -> tuple[str | None, bool]:
+    checkpoint_id = await asyncio.to_thread(session.checkpoint, label=label)
+    # Read the status in the same coroutine immediately after the worker
+    # returns, so concurrent checkpoint requests cannot overwrite the status
+    # before this request packages its response.
+    return checkpoint_id, bool(
+        getattr(session, "last_checkpoint_snapshot_included", False)
+    )
+
+
+async def _list_checkpoints(session: Any) -> list[dict[str, Any]]:
+    return session.list_checkpoints()
+
+
+def _resolve_checkpoint_id(checkpoints: list[dict[str, Any]], selector: str) -> str | None:
+    """Resolve a full ID, unique prefix, or one-based checkpoint index."""
+    value = selector.strip()
+    if not value:
+        return None
+    try:
+        index = int(value) - 1
+    except ValueError:
+        index = -1
+    if 0 <= index < len(checkpoints):
+        return str(checkpoints[index].get("id") or "") or None
+    exact = [str(cp.get("id") or "") for cp in checkpoints if cp.get("id") == value]
+    if exact:
+        return exact[0]
+    matches = [str(cp.get("id") or "") for cp in checkpoints if str(cp.get("id") or "").startswith(value)]
+    return matches[0] if len(matches) == 1 else None
+
+
+async def _resolve_and_revert(
+    session: Any,
+    selector: str,
+) -> tuple[str, dict[str, Any]] | None:
+    checkpoint_id = _resolve_checkpoint_id(session.list_checkpoints(), selector)
+    if checkpoint_id is None:
+        return None
+    return checkpoint_id, session.revert(checkpoint_id)
+
+
+async def _resolve_and_rollback(
+    session: Any,
+    selector: str,
+) -> tuple[str, bool, int] | None:
+    checkpoint_id = _resolve_checkpoint_id(session.list_checkpoints(), selector)
+    if checkpoint_id is None:
+        return None
+    ok = bool(session.rollback(checkpoint_id))
+    return checkpoint_id, ok, len(getattr(session, "messages", ()))
+
+
+async def _undo_latest(session: Any) -> tuple[str, dict[str, Any]] | None:
+    checkpoints = session.list_checkpoints()
+    if not checkpoints:
+        return None
+    checkpoint_id = str(checkpoints[0].get("id") or "")
+    if not checkpoint_id:
+        return None
+    return checkpoint_id, session.revert(checkpoint_id)
+
+
+def _get_session(request: Request, session_id: str):
+    """Retrieve a CoreSession from app state."""
+    sessions: dict = request.app.state.sessions
+    if session_id is None or session_id not in sessions:
+        return None
+    if session_id in getattr(request.app.state, "closing_sessions", set()):
+        return None
+    return sessions[session_id]
+
+
+@router.post("/checkpoint")
+async def create_checkpoint(req: CheckpointRequest, request: Request) -> dict[str, Any]:
+    """Create a checkpoint with a file-system snapshot."""
+    session = _get_session(request, req.session_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+    try:
+        cp_id, snapshot_included = await run_session_operation(
+            request.app.state,
+            session,
+            lambda: _checkpoint(session, req.label),
+        )
+    except SessionOperationRejected as exc:
+        raise HTTPException(status_code=503, detail=str(exc)) from exc
+    if not cp_id:
+        raise HTTPException(status_code=400, detail="Failed to create checkpoint")
+    return {
+        "checkpoint_id": cp_id,
+        "snapshot_included": snapshot_included,
+    }
+
+
+@router.get("/list")
+async def list_checkpoints(session_id: str, request: Request) -> list[dict[str, Any]]:
+    """List checkpoints for a session, including file snapshot info."""
+    session = _get_session(request, session_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+    try:
+        return await run_session_operation(
+            request.app.state,
+            session,
+            lambda: _list_checkpoints(session),
+        )
+    except SessionOperationRejected as exc:
+        raise HTTPException(status_code=503, detail=str(exc)) from exc
+
+
+@router.post("/revert")
+async def revert_checkpoint(req: RevertRequest, request: Request) -> dict[str, Any]:
+    """Revert both files and conversation to a checkpoint."""
+    session = _get_session(request, req.session_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+    try:
+        resolved = await run_session_operation(
+            request.app.state,
+            session,
+            lambda: _resolve_and_revert(session, req.checkpoint_id),
+        )
+    except SessionOperationRejected as exc:
+        raise HTTPException(status_code=503, detail=str(exc)) from exc
+    if resolved is None:
+        raise HTTPException(status_code=400, detail="Checkpoint not found or prefix is ambiguous")
+    _checkpoint_id, result = resolved
+    if not result.get("success"):
+        raise HTTPException(status_code=400, detail="Revert failed or checkpoint not found")
+    return result
+
+
+@router.post("/rollback")
+async def rollback_checkpoint(req: RevertRequest, request: Request) -> dict[str, Any]:
+    """Rollback conversation only (no file restore) to a checkpoint."""
+    session = _get_session(request, req.session_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+    try:
+        resolved = await run_session_operation(
+            request.app.state,
+            session,
+            lambda: _resolve_and_rollback(session, req.checkpoint_id),
+        )
+    except SessionOperationRejected as exc:
+        raise HTTPException(status_code=503, detail=str(exc)) from exc
+    if resolved is None:
+        raise HTTPException(status_code=400, detail="Checkpoint not found or prefix is ambiguous")
+    _checkpoint_id, ok, message_count = resolved
+    if not ok:
+        raise HTTPException(status_code=400, detail="Rollback failed or checkpoint not found")
+    return {"success": True, "messages_count": message_count}
+
+
+@router.post("/undo")
+async def undo_last_checkpoint(request: Request) -> dict[str, Any]:
+    """Revert files and conversation to the most recent checkpoint."""
+    try:
+        body = await request.json()
+    except (ValueError, TypeError) as exc:
+        raise HTTPException(status_code=400, detail="Request body must be valid JSON") from exc
+    if body is None:
+        body = {}
+    if not isinstance(body, dict):
+        raise HTTPException(status_code=400, detail="Request body must be a JSON object")
+    session_id = body.get("session_id") if body else None
+
+    sessions: dict = request.app.state.sessions
+    sid = (
+        request.app.state.default_session_id
+        if session_id is None
+        else session_id
+    )
+    if (
+        not sid
+        or sid not in sessions
+        or sid in getattr(request.app.state, "closing_sessions", set())
+    ):
+        raise HTTPException(status_code=404, detail="Session not found")
+    session = sessions[sid]
+
+    try:
+        resolved = await run_session_operation(
+            request.app.state,
+            session,
+            lambda: _undo_latest(session),
+        )
+    except SessionOperationRejected as exc:
+        raise HTTPException(status_code=503, detail=str(exc)) from exc
+    if resolved is None:
+        raise HTTPException(status_code=400, detail="No checkpoints to undo")
+    checkpoint_id, result = resolved
+    if not result.get("success"):
+        raise HTTPException(status_code=400, detail="Undo failed")
+    return {"checkpoint_id": checkpoint_id, **result}

crabcode_gateway/routes/tasks.py

@@ -0,0 +1,318 @@
+"""Background monitor and managed-agent task routes."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from fastapi import APIRouter, HTTPException, Request
+
+from crabcode_gateway.schemas import BackgroundTaskInfo, TaskStopRequest
+from crabcode_gateway.safe_file import read_regular_file_tail
+from crabcode_gateway.session_registry import get_session_lock
+from crabcode_gateway.task_registry import SessionOperationRejected, run_session_operation
+
+
+router = APIRouter(tags=["tasks"])
+
+
+def _task_dicts(session: Any) -> list[dict[str, Any]]:
+    """Return one stable representation for monitors and managed agents."""
+    result: list[dict[str, Any]] = []
+    cwd = str(getattr(session, "cwd", "") or "")
+    for snapshot in session.list_monitor_tasks():
+        data = snapshot.to_dict() if hasattr(snapshot, "to_dict") else dict(snapshot)
+        data.setdefault("cwd", cwd)
+        result.append(data)
+    for snapshot in session.list_agents():
+        result.append(
+            {
+                "task_id": snapshot.agent_id,
+                "agent_id": snapshot.agent_id,
+                "session_id": snapshot.session_id,
+                "cwd": cwd,
+                "description": snapshot.title,
+                "task_type": "local_agent",
+                "source": "agent",
+                "status": snapshot.status,
+                "output_file": snapshot.transcript_path,
+                "created_at": snapshot.created_at,
+                "started_at": snapshot.started_at,
+                "finished_at": snapshot.finished_at,
+                "updated_at": snapshot.updated_at,
+                "error": snapshot.error,
+            }
+        )
+    return result
+
+
+def _task_info(data: dict[str, Any]) -> BackgroundTaskInfo:
+    return BackgroundTaskInfo(
+        task_id=str(data.get("task_id") or data.get("agent_id") or ""),
+        agent_id=data.get("agent_id"),
+        session_id=str(data.get("session_id") or ""),
+        cwd=str(data.get("cwd") or ""),
+        description=str(data.get("description") or ""),
+        task_type=str(data.get("task_type") or ""),
+        source=str(data.get("source") or ""),
+        status=str(data.get("status") or ""),
+        output_file=data.get("output_file"),
+        created_at=str(data.get("created_at") or ""),
+        started_at=data.get("started_at"),
+        finished_at=data.get("finished_at"),
+        updated_at=str(data.get("updated_at") or ""),
+        error=str(data.get("error") or ""),
+        exit_code=data.get("exit_code"),
+    )
+
+
+def _task_matches(session: Any, selector: str) -> bool:
+    value = str(selector or "").strip()
+    if not value:
+        return False
+    return any(
+        item.get("task_id", "") == value
+        or str(item.get("task_id", "")).startswith(value)
+        for item in _task_dicts(session)
+    )
+
+
+def _get_session(request: Request, session_id: str | None = None):
+    sessions: dict[str, Any] = request.app.state.sessions
+    closing = getattr(request.app.state, "closing_sessions", set())
+
+    def usable(session: Any) -> bool:
+        return session is not None and getattr(session, "session_id", None) not in closing
+
+    if session_id is not None:
+        session = sessions.get(session_id)
+        return session if usable(session) else None
+
+    default = sessions.get(getattr(request.app.state, "default_session_id", None))
+    return default if usable(default) else None
+
+
+async def _get_task_session(
+    request: Request,
+    session_id: str | None,
+    task_id: str,
+):
+    """Resolve cross-session task selectors without inspecting cold sessions."""
+    if session_id is not None:
+        return _get_session(request, session_id)
+
+    sessions: dict[str, Any] = request.app.state.sessions
+    closing = getattr(request.app.state, "closing_sessions", set())
+    default = _get_session(request)
+    candidates = ([default] if default is not None else []) + [
+        session
+        for session in sessions.values()
+        if session is not default
+        and getattr(session, "session_id", None) not in closing
+    ]
+
+    async def _matches(candidate: Any) -> bool:
+        initializer = getattr(candidate, "initialize", None)
+        if callable(initializer):
+            await initializer()
+        return _task_matches(candidate, task_id)
+
+    for candidate in candidates:
+        try:
+            if await run_session_operation(request.app.state, candidate, lambda: _matches(candidate)):
+                return candidate
+        except SessionOperationRejected:
+            continue
+    # Keep historical default-session behavior so callers get a task 404 rather
+    # than a misleading session 404 when the selector simply does not exist.
+    return default
+
+
+def _matching_tasks(session: Any, task_id: str) -> list[dict[str, Any]]:
+    tasks = _task_dicts(session)
+    matches = [item for item in tasks if item.get("task_id") == task_id]
+    if not matches:
+        matches = [
+            item for item in tasks
+            if str(item.get("task_id") or "").startswith(task_id)
+        ]
+    return matches
+
+
+@router.get("/tasks", response_model=list[BackgroundTaskInfo])
+@router.get("/tasks/list", response_model=list[BackgroundTaskInfo])
+@router.get("/task/list", response_model=list[BackgroundTaskInfo])
+async def list_tasks(
+    request: Request,
+    session_id: str | None = None,
+    scope: str | None = None,
+    status: str | None = None,
+) -> list[BackgroundTaskInfo]:
+    if scope == "global":
+        # The automation deck needs one stable view across every currently
+        # loaded project instead of following whichever chat is the default.
+        async with get_session_lock(request.app.state):
+            closing = getattr(request.app.state, "closing_sessions", set())
+            sessions = [
+                session
+                for session in request.app.state.sessions.values()
+                if getattr(session, "session_id", None) not in closing
+            ]
+
+        result: list[BackgroundTaskInfo] = []
+        for session in sessions:
+            async def _list_global() -> list[BackgroundTaskInfo]:
+                initializer = getattr(session, "initialize", None)
+                if callable(initializer):
+                    await initializer()
+                return [
+                    _task_info(item)
+                    for item in _task_dicts(session)
+                    if status is None or str(item.get("status") or "") == status
+                ]
+
+            try:
+                result.extend(
+                    await run_session_operation(
+                        request.app.state,
+                        session,
+                        _list_global,
+                    )
+                )
+            except SessionOperationRejected:
+                # A session can be closing while the deck refreshes; the next
+                # poll will reconcile it without failing the whole dashboard.
+                continue
+        return sorted(result, key=lambda item: item.updated_at, reverse=True)
+
+    session = _get_session(request, session_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+
+    async def _list() -> list[BackgroundTaskInfo]:
+        initializer = getattr(session, "initialize", None)
+        if callable(initializer):
+            await initializer()
+        return [
+            _task_info(item)
+            for item in _task_dicts(session)
+            if status is None or str(item.get("status") or "") == status
+        ]
+
+    try:
+        return await run_session_operation(request.app.state, session, _list)
+    except SessionOperationRejected as exc:
+        raise HTTPException(status_code=503, detail=str(exc)) from exc
+
+
+@router.get("/tasks/{task_id}", response_model=BackgroundTaskInfo)
+async def get_task(task_id: str, request: Request, session_id: str | None = None) -> BackgroundTaskInfo:
+    session = await _get_task_session(request, session_id, task_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Task not found")
+
+    async def _get() -> list[dict[str, Any]]:
+        initializer = getattr(session, "initialize", None)
+        if callable(initializer):
+            await initializer()
+        return _matching_tasks(session, task_id)
+
+    try:
+        matches = await run_session_operation(request.app.state, session, _get)
+    except SessionOperationRejected as exc:
+        raise HTTPException(status_code=503, detail=str(exc)) from exc
+    if len(matches) != 1:
+        raise HTTPException(status_code=404, detail="Task not found or selector is ambiguous")
+    return _task_info(matches[0])
+
+
+@router.post("/tasks/stop")
+@router.post("/task/stop")
+async def stop_task(req: TaskStopRequest, request: Request) -> dict[str, Any]:
+    session = await _get_task_session(request, req.session_id, req.task_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Task not found")
+    try:
+        stopped = await run_session_operation(
+            request.app.state,
+            session,
+            lambda: session.stop_background_task(req.task_id),
+        )
+    except SessionOperationRejected as exc:
+        raise HTTPException(status_code=503, detail=str(exc)) from exc
+    if not stopped:
+        raise HTTPException(status_code=404, detail=f"No running background task: {req.task_id}")
+    return {"status": "ok", "task_id": req.task_id, "stopped": True}
+
+
+@router.get("/tasks/{task_id}/output")
+async def task_output(task_id: str, request: Request, session_id: str | None = None) -> dict[str, Any]:
+    """Read a bounded tail of a monitor/agent output file."""
+    session = await _get_task_session(request, session_id, task_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Task not found")
+    try:
+        limit = max(1, min(10_000, int(request.query_params.get("lines", "200"))))
+    except (TypeError, ValueError):
+        raise HTTPException(status_code=400, detail="lines must be an integer")
+    async def _read_output() -> dict[str, Any] | None:
+        initializer = getattr(session, "initialize", None)
+        if callable(initializer):
+            await initializer()
+        matches = _matching_tasks(session, task_id)
+        if len(matches) != 1:
+            return None
+        path_value = matches[0].get("output_file")
+        if not path_value:
+            return {
+                "task_id": matches[0].get("task_id"),
+                "path": None,
+                "lines": [],
+                "truncated": False,
+            }
+        task = matches[0]
+        resolved_task_id = str(task.get("task_id") or "")
+        try:
+            from crabcode_core.session.storage import (
+                get_agent_transcript_path,
+                get_task_output_path,
+            )
+
+            if task.get("agent_id") or task.get("source") == "agent":
+                expected_path = get_agent_transcript_path(
+                    session.cwd,
+                    session.session_id,
+                    str(task.get("agent_id") or resolved_task_id),
+                )
+            else:
+                expected_path = get_task_output_path(
+                    session.cwd,
+                    session.session_id,
+                    resolved_task_id,
+                )
+            output_lines, truncated = read_regular_file_tail(
+                expected_path,
+                max_lines=limit,
+            )
+        except (OSError, TypeError, ValueError) as exc:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Task output unavailable: {exc}",
+            ) from exc
+        return {
+            "task_id": resolved_task_id,
+            "path": str(expected_path),
+            "lines": output_lines[-limit:],
+            "truncated": truncated,
+        }
+
+    try:
+        result = await run_session_operation(
+            request.app.state,
+            session,
+            _read_output,
+        )
+    except SessionOperationRejected as exc:
+        raise HTTPException(status_code=503, detail=str(exc)) from exc
+    if result is None:
+        raise HTTPException(status_code=404, detail="Task not found or selector is ambiguous")
+    return result