feat: resume / re-run for aborted and failed sessions

backend/routers/sessions.py

@@ -17,8 +17,19 @@
 from errors import capture_exception
 from models import Artifact, Experiment, LogEvent, Message, Metric, Task
 from models import Session as SessionModel
-from schemas import ClarificationReply, MessageCreate, TaskCreate, TaskUpdate
+from schemas import (
+    ClarificationReply,
+    MessageCreate,
+    SessionResume,
+    TaskCreate,
+    TaskUpdate,
+)
 from services.agent import abort_agent, run_agent
+from services.agent.resume import (
+    build_resume_context,
+    build_resume_prompt,
+    is_resumable_state,
+)
 from services.agent.tasks import is_agent_running, register_task
 from services.broadcaster import broadcaster
 from services.clarifications import list_pending, resolve as resolve_clarification
@@ -223,6 +234,109 @@ async def abort_session(session_id: str, db: AsyncSession = Depends(get_db)):
     return {"status": "not_running"}
 
 
+@router.post("/sessions/{session_id}/resume")
+async def resume_session(
+    session_id: str,
+    body: SessionResume | None = None,
+    db: AsyncSession = Depends(get_db),
+):
+    """Resume or retry an interrupted session without re-driving the chat.
+
+    Relaunches the agent with the prior conversation, persisted tool history,
+    task state, and the workspace file listing injected (see
+    services/agent/resume.py) so steps whose artifacts already exist on the
+    volume are skipped rather than redone.
+    """
+    body = body or SessionResume()
+    result = await db.execute(
+        select(SessionModel)
+        .where(SessionModel.id == session_id)
+        .options(selectinload(SessionModel.experiment).selectinload(Experiment.project))
+    )
+    session = result.scalar_one_or_none()
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+
+    if is_agent_running(session_id):
+        raise HTTPException(
+            status_code=409,
+            detail="An agent is already running in this session — abort it first "
+            "or wait for it to finish.",
+        )
+
+    prior_state = session.state or "created"
+    if not is_resumable_state(prior_state):
+        raise HTTPException(
+            status_code=400,
+            detail=f"Session state '{prior_state}' has no prior run to resume.",
+        )
+
+    resume_context = await build_resume_context(session_id, prior_state)
+    resume_prompt = build_resume_prompt(prior_state, body.mode)
+
+    # Capture experiment context before the DB session closes (mirrors
+    # send_message's follow-up launch).
+    stage = "chat"
+    experiment_id = session.experiment_id
+    dataset_ref = session.experiment.dataset_ref or "" if session.experiment else ""
+    instructions = session.experiment.instructions or "" if session.experiment else ""
+    selected_model = body.model or session.model
+    agent_models = body.agent_models or {}
+    agent_thinking = body.agent_thinking or {}
+    sandbox_config = {}
+    if session.experiment and session.experiment.project:
+        sandbox_config = session.experiment.project.sandbox_config or {}
+
+    session.state = f"{stage}_running"
+    await db.commit()
+
+    await broadcaster.publish(
+        session_id,
+        {
+            "type": "session_resumed",
+            "data": {"mode": body.mode, "prior_state": prior_state},
+        },
+    )
+
+    async def _run_resume():
+        try:
+            await run_agent(
+                session_id=session_id,
+                experiment_id=experiment_id,
+                stage=stage,
+                instructions=instructions,
+                dataset_ref=dataset_ref,
+                user_prompt=resume_prompt,
+                sandbox_config=sandbox_config,
+                model=selected_model,
+                agent_models=agent_models,
+                agent_thinking=agent_thinking,
+                resume_context=resume_context,
+            )
+            async with async_session() as fresh_db:
+                s = await fresh_db.get(SessionModel, session_id)
+                if s:
+                    s.state = "done"
+                    await fresh_db.commit()
+        except asyncio.CancelledError:
+            async with async_session() as fresh_db:
+                s = await fresh_db.get(SessionModel, session_id)
+                if s and s.state != "cancelled":
+                    s.state = "cancelled"
+                    await fresh_db.commit()
+        except Exception:
+            async with async_session() as fresh_db:
+                s = await fresh_db.get(SessionModel, session_id)
+                if s:
+                    s.state = "failed"
+                    await fresh_db.commit()
+
+    task = asyncio.create_task(_run_resume())
+    await register_task(session_id, task)
+
+    return {"status": "resumed", "mode": body.mode, "prior_state": prior_state}
+
+
 @router.get("/sessions/{session_id}/clarifications")
 async def get_pending_clarifications(session_id: str):
     """Return any clarifications currently waiting for a user reply."""

backend/schemas.py

@@ -134,6 +134,18 @@ class ClarificationReply(BaseModel):
     answer: str = Field(..., max_length=_CLARIFICATION_MAX)
 
 
+class SessionResume(BaseModel):
+    """Body for POST /sessions/{id}/resume. Everything is optional — the
+    default is 'relaunch with the same knobs the session already has'."""
+
+    # "resume" continues interrupted work; "retry" re-drives a failed stage.
+    # Both relaunch the same way — the mode only shades the agent prompt.
+    mode: Literal["resume", "retry"] = "resume"
+    model: Optional[str] = Field(default=None, max_length=_MODEL_ID_MAX)
+    agent_models: Optional[dict[str, str]] = Field(default=None)
+    agent_thinking: Optional[dict[str, str]] = Field(default=None)
+
+
 class TaskCreate(BaseModel):
     subject: str = Field(..., min_length=1, max_length=_NAME_MAX)
     short_description: str = Field(default="", max_length=_DESC_MAX)

backend/services/agent/resume.py

@@ -0,0 +1,187 @@
+"""Resume-context assembly for interrupted sessions.
+
+When a session is aborted, fails, or times out mid-run, the in-flight agent
+loop is gone but three durable traces of its progress remain:
+
+1. persisted tool history (``agent_thought`` Message rows — tool_use /
+   tool_result blocks, already truncated at write time),
+2. the session's task list (the ``tasks`` skill state), and
+3. artifacts already written to the session workspace on the volume.
+
+``build_resume_context`` folds those into a single prompt block the relaunched
+agent reads to skip completed steps instead of re-driving the whole
+conversation.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from sqlalchemy import select
+
+from db import async_session
+from models import Message, Task
+from services.volume import listdir_async, reload_volume_async
+
+logger = logging.getLogger(__name__)
+
+# Caps so a long session can't blow up the relaunched agent's context window.
+_MAX_TOOL_BLOCKS = 40
+_MAX_BLOCK_CHARS = 500
+_MAX_FILES = 100
+
+# States a session can be resumed from. `*_running` is included because a
+# backend restart can strand the DB in a running state with no live task.
+RESUMABLE_TERMINAL_STATES = {"failed", "cancelled", "timed_out", "done"}
+
+
+def is_resumable_state(state: str | None) -> bool:
+    state = state or ""
+    return (
+        state in RESUMABLE_TERMINAL_STATES
+        or state.endswith("_running")
+        or state.endswith("_done")
+    )
+
+
+def _prior_run_outcome(prior_state: str) -> str:
+    """How the previous run ended, for prompt wording. `done` / `*_done`
+    sessions are resumed as re-runs of finished work, not interruptions."""
+    if prior_state == "done" or prior_state.endswith("_done"):
+        return "completed"
+    return "stopped early"
+
+
+def _clip(text: str | None, limit: int = _MAX_BLOCK_CHARS) -> str:
+    if text is None:
+        return ""
+    text = str(text)
+    if len(text) <= limit:
+        return text
+    return text[:limit] + f"…[+{len(text) - limit} chars]"
+
+
+async def _load_tool_history(session_id: str) -> list[str]:
+    """Render the persisted agent_thought tool blocks as one line each."""
+    lines: list[str] = []
+    try:
+        async with async_session() as db:
+            result = await db.execute(
+                select(Message)
+                .where(Message.session_id == session_id)
+                .order_by(Message.id)
+            )
+            for msg in result.scalars().all():
+                meta = msg.metadata_ or {}
+                if meta.get("event_type") != "agent_thought":
+                    continue
+                block_type = meta.get("block_type")
+                if block_type == "tool_use":
+                    tool = meta.get("tool_name") or "tool"
+                    lines.append(f"- called `{tool}` with {_clip(msg.content)}")
+                elif block_type == "tool_result":
+                    marker = "ERROR result" if meta.get("is_error") else "result"
+                    lines.append(f"  - {marker}: {_clip(msg.content)}")
+    except Exception as e:
+        logger.warning("Resume: failed to load tool history for %s: %s", session_id, e)
+    # Keep the most recent blocks — the tail is what the agent needs to
+    # figure out where the previous run stopped.
+    return lines[-_MAX_TOOL_BLOCKS:]
+
+
+async def _load_task_state(session_id: str) -> list[str]:
+    lines: list[str] = []
+    try:
+        async with async_session() as db:
+            result = await db.execute(
+                select(Task).where(Task.session_id == session_id).order_by(Task.id)
+            )
+            for t in result.scalars().all():
+                lines.append(f"- [{t.status}] {t.subject}")
+    except Exception as e:
+        logger.warning("Resume: failed to load tasks for %s: %s", session_id, e)
+    return lines
+
+
+async def _load_workspace_files(session_id: str) -> list[str]:
+    paths: list[str] = []
+    workspace = f"/sessions/{session_id}"
+    try:
+        await reload_volume_async()
+        for entry in await listdir_async(workspace, recursive=True):
+            if entry.type.name != "FILE":
+                continue
+            paths.append(entry.path)
+    except FileNotFoundError:
+        pass
+    except Exception as e:
+        logger.warning(
+            "Resume: failed to list workspace files for %s: %s", session_id, e
+        )
+    paths.sort()
+    if len(paths) > _MAX_FILES:
+        extra = len(paths) - _MAX_FILES
+        paths = paths[:_MAX_FILES] + [f"…({extra} more)"]
+    return paths
+
+
+async def build_resume_context(session_id: str, prior_state: str) -> str:
+    """Assemble the resume-context prompt block for a relaunched agent."""
+    tool_lines = await _load_tool_history(session_id)
+    task_lines = await _load_task_state(session_id)
+    file_lines = await _load_workspace_files(session_id)
+
+    sections = [
+        "## Resumed session — recovered progress",
+        "",
+        f"This session's previous run {_prior_run_outcome(prior_state)} "
+        f"(last recorded state: `{prior_state}`). The evidence of its progress "
+        "is below. Use it to skip steps whose outputs already exist instead "
+        "of redoing them.",
+    ]
+
+    sections.append("\n### Task list at interruption")
+    if task_lines:
+        sections.append("\n".join(task_lines))
+    else:
+        sections.append("(no tasks were recorded)")
+
+    sections.append(
+        f"\n### Recent tool activity (last {_MAX_TOOL_BLOCKS} blocks, oldest first)"
+    )
+    if tool_lines:
+        sections.append("\n".join(tool_lines))
+    else:
+        sections.append("(no tool activity was recorded)")
+
+    sections.append("\n### Files already present in the session workspace")
+    if file_lines:
+        sections.append("\n".join(f"- {p}" for p in file_lines))
+    else:
+        sections.append("(the workspace is empty)")
+
+    sections.append(
+        "\n### Resume rules\n"
+        "- Treat files listed above as completed work: verify cheaply (read / "
+        "list) instead of regenerating them.\n"
+        "- Continue from the first incomplete step; mark tasks completed as "
+        "you confirm or finish them.\n"
+        "- If a step's artifact exists but looks partial or corrupt, redo "
+        "only that step."
+    )
+
+    return "\n".join(sections)
+
+
+def build_resume_prompt(prior_state: str, mode: str) -> str:
+    """The synthetic user prompt the relaunched agent is driven with."""
+    verb = "Retry the failed work" if mode == "retry" else "Resume the work"
+    return (
+        f"{verb} in this session. The previous run "
+        f"{_prior_run_outcome(prior_state)} (last recorded state: "
+        f"`{prior_state}`). Review the 'Resumed session — "
+        "recovered progress' section of your system prompt: skip steps whose "
+        "artifacts already exist in the workspace, then continue from the "
+        "first incomplete step through to completion. Do not re-ask questions "
+        "already answered in the prior conversation."
+    )

backend/services/agent/runner.py

@@ -953,9 +953,16 @@ async def run_agent(
     agent_id: str = "root",
     parent_agent_id: str | None = None,
     mentions: list[dict] | None = None,
+    resume_context: str | None = None,
 ):
     """Run an agent. agent_type maps to a YAML in agents/. Falls back to stage name.
 
+    resume_context, when set, marks this run as a resume/retry of an
+    interrupted session: the block (built by services.agent.resume) is
+    appended to the system prompt so the agent can skip steps whose
+    artifacts already exist, and the full prior conversation is injected
+    (a resume has no freshly-persisted user message to exclude).
+
     agent_models is a per-agent model override map: {"eda": "claude-haiku-4-5", ...}
     agent_thinking is the parallel reasoning-level map: {"eda": "high", ...}.
     Levels are abstract — services/llm/thinking.py translates them per provider.
@@ -1058,6 +1065,11 @@ async def _publish(
         if "execute-code" in get_agent_skills(agent_type):
             system_prompt += "\n\n" + _format_compute_env(sandbox_config)
 
+        # Resume/retry runs carry the recovered-progress block so the agent
+        # skips steps whose artifacts already exist on the volume.
+        if resume_context:
+            system_prompt += "\n\n" + resume_context
+
         # Pre-flight training controls (issue #104) — only meaningful for
         # agents that can open a training window. Empty config renders to ""
         # so unconfigured projects get a byte-identical prompt.
@@ -1109,12 +1121,16 @@ async def _publish(
             )
             thinking_level = normalize_level(chosen) if chosen else None
 
-        # Load conversation history for follow-up messages
+        # Load conversation history for follow-up messages. A normal
+        # follow-up excludes the last message (it's the just-persisted user
+        # prompt this run is answering); a resume run has no fresh user
+        # message in the DB, so the full history is injected.
         if user_prompt:
             history = await _load_conversation_history(session_id)
+            history_for_context = history if resume_context else history[:-1]
             if history:
                 context_parts = []
-                for msg in history[:-1]:
+                for msg in history_for_context:
                     prefix = "User" if msg["role"] == "user" else "Assistant"
                     context_parts.append(f"{prefix}: {msg['content']}")
                 if context_parts: