From ae3781dc4234b86b23742147f9d9c65579d0eb8f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 28 Jul 2026 08:07:43 +0000 Subject: [PATCH 1/2] Return waiting instead of needs_replanning when a step times out with an active remote job --- .../agents/execution_agent/recovery.py | 80 ++++++++- .../execution_agent/step_executor_runner.py | 163 +++++++++++++++++- src/matcreator/agents/graph_logger.py | 2 +- 3 files changed, 233 insertions(+), 12 deletions(-) diff --git a/src/matcreator/agents/execution_agent/recovery.py b/src/matcreator/agents/execution_agent/recovery.py index 57a16be3..dcfe554e 100644 --- a/src/matcreator/agents/execution_agent/recovery.py +++ b/src/matcreator/agents/execution_agent/recovery.py @@ -231,6 +231,20 @@ def _mark_attempt_stale(latest_path: Path, attempt: dict[str, Any]) -> None: _write_attempt(attempt) +def _refresh_remote_job_reference(attempt: dict[str, Any]) -> None: + """Reload the persisted remote-job identity into an in-memory attempt. + + ``record_remote_job_reference`` writes through the durable record, so a live + attempt dict only learns about a submitted job on its next heartbeat. + """ + latest_path = attempt.get("_latest_path") + if not latest_path: + return + persisted = _read_json(Path(latest_path)) + if isinstance(persisted.get("remote_job"), dict): + attempt["remote_job"] = persisted["remote_job"] + + def _active_remote_job(attempt: dict[str, Any]) -> dict[str, Any] | None: reference = attempt.get("remote_job") if not isinstance(reference, dict) or not reference.get("job_id"): @@ -249,6 +263,28 @@ def _active_remote_job(attempt: dict[str, Any]) -> dict[str, Any] | None: return None +def remote_job_reference(job: dict[str, Any]) -> dict[str, Any]: + """Return the identity subset of a remote job stored on a graph node.""" + return { + "job_id": job["job_id"], + "provider": job["provider"], + "external_id": job["external_id"], + "status": job["status"], + } + + +def active_remote_job_for_attempt(attempt: dict[str, Any]) -> dict[str, Any] | None: + """Return the still-active remote job owned by a live attempt, if any. + + Used by the step executor when its wall-clock timeout fires: a timeout that + leaves a tracked remote job running is a handoff, not a failure. + """ + if not attempt: + return None + _refresh_remote_job_reference(attempt) + return _active_remote_job(attempt) + + def reconcile_recovery_state( state: Any, workspace_dir: str | Path, @@ -299,12 +335,7 @@ def reconcile_recovery_state( if remote_job is not None: node["status"] = "waiting" node["result"] = "Recovered active remote job; waiting for provider completion." - node["remote_job"] = { - "job_id": remote_job["job_id"], - "provider": remote_job["provider"], - "external_id": remote_job["external_id"], - "status": remote_job["status"], - } + node["remote_job"] = remote_job_reference(remote_job) node["recovery"] = { "attempt": attempt.get("attempt"), "status": "waiting_remote", @@ -323,6 +354,43 @@ def reconcile_recovery_state( actions.append({"node_id": node_id, "action": "reset_stale_running", "status": "pending"}) continue + # A step executor that timed out while its remote job was still running + # hands the node off as ``waiting`` rather than failing it. Keep it + # waiting until the provider settles, then let the node run again so a + # fresh executor can re-attach and collect the results. Skip nodes the + # planner has already resumed, otherwise the retry would be undone here. + resumed_by_planner = ( + node_status == "pending" + and isinstance(node.get("recovery"), dict) + and node["recovery"].get("status") == "waiting_remote" + ) + if ( + attempt_status == "waiting" + and node_status in ("pending", "running", "waiting") + and not resumed_by_planner + ): + remote_job = _active_remote_job(attempt) + if remote_job is not None: + node["status"] = "waiting" + node["result"] = _attempt_summary(attempt) + node["remote_job"] = remote_job_reference(remote_job) + node["recovery"] = { + "attempt": attempt.get("attempt"), + "status": "waiting_remote", + "recovered_at": _now(), + } + actions.append({"node_id": node_id, "action": "wait_for_remote_job", "status": "waiting"}) + else: + node["status"] = "pending" + node["result"] = "Remote job is no longer active; re-running node to collect its results." + node["recovery"] = { + "attempt": attempt.get("attempt"), + "status": "remote_job_settled", + "recovered_at": _now(), + } + actions.append({"node_id": node_id, "action": "resume_after_remote_job", "status": "pending"}) + continue + if node_status in ("pending", "running") and attempt_status in { "success", "failed", diff --git a/src/matcreator/agents/execution_agent/step_executor_runner.py b/src/matcreator/agents/execution_agent/step_executor_runner.py index ce74ea3e..6adc35b1 100644 --- a/src/matcreator/agents/execution_agent/step_executor_runner.py +++ b/src/matcreator/agents/execution_agent/step_executor_runner.py @@ -25,7 +25,13 @@ StepExecutorResult, build_step_executor_agent, ) -from .recovery import finish_node_attempt, heartbeat_node_attempt, start_node_attempt +from .recovery import ( + active_remote_job_for_attempt, + finish_node_attempt, + heartbeat_node_attempt, + remote_job_reference as _remote_job_reference, + start_node_attempt, +) from ..graph_logger import AgentGraphLogger from ..execution_graph_state import get_execution_graph, set_execution_graph from ..session_log import ( @@ -51,6 +57,11 @@ # Wall-clock timeout for a single step or sub-step execution. # A sub-step that exceeds this returns needs_replanning instead of hanging. _SUB_STEP_TIMEOUT = int(os.environ.get("SUB_STEP_TIMEOUT", "3600")) # seconds +# One-time bounded extension granted when the step timeout fires while a tracked +# remote job is still active. This is deliberately modest: the durable job store +# and the remote job monitor — not a long-lived LLM session — are what keep a +# remote job alive, so the extension only avoids churn for jobs about to settle. +_REMOTE_JOB_GRACE_TIMEOUT = int(os.environ.get("STEP_REMOTE_JOB_GRACE_TIMEOUT", "300")) # seconds _RECOVERY_HEARTBEAT_INTERVAL = int(os.environ.get("STEP_RECOVERY_HEARTBEAT_INTERVAL", "10")) # seconds _MAX_INPUT_IMAGE_ATTACHMENTS = int(os.environ.get("MATCREATOR_MAX_INPUT_IMAGE_ATTACHMENTS", "4")) _MAX_INPUT_IMAGE_BYTES = int(os.environ.get("MATCREATOR_MAX_INPUT_IMAGE_BYTES", str(5 * 1024 * 1024))) @@ -318,6 +329,98 @@ async def _heartbeat_recovery_attempt(attempt: dict) -> None: pass +def _mark_node_waiting_on_remote_job( + tool_context: ToolContext, + *, + node_id: Optional[str], + remote_job: dict, + summary: str, +) -> None: + """Record the remote-job handoff directly on the execution graph node. + + Writing the node state here — instead of relying on the orchestrator LLM to + call ``set_node_status`` — keeps the handoff deterministic, and preserves the + job identity the next executor needs in order to re-attach. + """ + if not node_id: + return + graph_state = get_execution_graph(tool_context.state) or {} + graph_nodes = graph_state.get("nodes") or {} + if node_id not in graph_nodes: + return + graph_nodes[node_id]["status"] = "waiting" + graph_nodes[node_id]["result"] = summary + graph_nodes[node_id]["remote_job"] = remote_job + graph_nodes[node_id]["recovery"] = { + "status": "waiting_remote", + "recorded_at": _now(), + } + set_execution_graph(tool_context.state, graph_state) + + +def _remote_job_prior_context(tool_context: ToolContext, node_id: Optional[str]) -> Optional[str]: + """Return re-attachment instructions when a node already owns a remote job. + + Injected into the executor's ``prior_context`` so re-attachment does not + depend on the planner remembering the job, nor on the submission idempotency + key happening to match. + """ + if not node_id: + return None + node = ((get_execution_graph(tool_context.state) or {}).get("nodes") or {}).get(node_id) + if not isinstance(node, dict): + return None + remote_job = node.get("remote_job") + if not isinstance(remote_job, dict) or not remote_job.get("job_id"): + return None + return ( + "REMOTE JOB ALREADY SUBMITTED for this step: " + f"job_id={remote_job['job_id']} provider={remote_job.get('provider')} " + f"sandbox_id={remote_job.get('external_id')} last_known_status={remote_job.get('status')}. " + "Call get_e2b_job_status with this job_id to re-attach. Do NOT call submit_e2b_sandbox " + "again for this step — that job is still tracked and must not be duplicated. If it is " + "still running, report needs_replanning explaining that the job has not finished yet." + ) + + +async def _await_step_completion( + inner_task: asyncio.Task, + recovery_attempt: dict, + *, + step_number: int, + session_id: str, +) -> tuple[bool, Optional[dict]]: + """Wait for the executor task, honouring one remote-job grace extension. + + Returns ``(timed_out, active_remote_job)``. ``active_remote_job`` is the + still-active tracked job that caused the final timeout, if any; the caller + treats that case as a handoff rather than a step failure. Unlike + ``asyncio.wait_for`` this never cancels ``inner_task`` on timeout, so the + grace window can reuse the same in-flight executor. + """ + timeout = float(_SUB_STEP_TIMEOUT) + grace_remaining = float(max(_REMOTE_JOB_GRACE_TIMEOUT, 0)) + while True: + done, _ = await asyncio.wait({inner_task}, timeout=timeout) + if done: + return False, None + + remote_job = await asyncio.to_thread(active_remote_job_for_attempt, recovery_attempt) + if remote_job is None: + inner_task.cancel() + return True, None + if grace_remaining <= 0: + inner_task.cancel() + return True, remote_job + + logger.info( + "[TIMEOUT] Step %d reached its timeout with remote job %s still %s; " + "granting a single %ds grace window (session=%s)", + step_number, remote_job["job_id"], remote_job["status"], int(grace_remaining), session_id, + ) + timeout, grace_remaining = grace_remaining, 0.0 + + async def _cleanup_step_runner( runner: Runner, tasks: tuple[asyncio.Task, ...], @@ -473,6 +576,12 @@ async def run_step_executor( # Use node_id for label when provided (DAG mode); fall back to step_number. effective_id = node_id if node_id else str(step_number) + # Deterministic re-attachment: a node that already owns a tracked remote job + # must poll it rather than submit a replacement sandbox. + reattach_context = _remote_job_prior_context(tool_context, node_id) + if reattach_context: + prior_context = f"{reattach_context}\n\n{prior_context}" if prior_context else reattach_context + # All steps CWD directly to the workspace root/session workdir so they can read shared inputs. # If configured, generated artifacts are constrained separately by output_dir. step_workspace = Path(tool_context.state.get("workspace_dir") or str(get_session_workdir(session_id))) @@ -645,17 +754,21 @@ async def run_step_executor( cancelled = False timed_out = False + waiting_remote_job: Optional[dict] = None runner_error: Optional[Exception] = None step_state_delta: dict = {} plot_paths: list[str] = [] artifact_paths: list[str] = [] event_log: dict = {"conversation": [], "tool_calls": []} try: - step_state_delta, plot_paths, artifact_paths, event_log = await asyncio.wait_for( - inner_task, timeout=_SUB_STEP_TIMEOUT + timed_out, waiting_remote_job = await _await_step_completion( + inner_task, + recovery_attempt, + step_number=step_number, + session_id=session_id, ) - except asyncio.TimeoutError: - timed_out = True + if not timed_out: + step_state_delta, plot_paths, artifact_paths, event_log = inner_task.result() except asyncio.CancelledError: cancelled = True except Exception as exc: @@ -671,6 +784,46 @@ async def run_step_executor( else: await _cleanup_step_runner(runner, cleanup_tasks) + if timed_out and waiting_remote_job is not None: + job_summary = ( + f"Step {step_number} reached its {_SUB_STEP_TIMEOUT}s executor timeout while tracked " + f"{waiting_remote_job['provider']} job {waiting_remote_job['job_id']} is still " + f"{waiting_remote_job['status']}. The executor was released; the job keeps running." + ) + logger.warning("[WAITING] %s (session=%s)", job_summary, session_id) + remote_job_reference = _remote_job_reference(waiting_remote_job) + _mark_node_waiting_on_remote_job( + tool_context, + node_id=node_id, + remote_job=remote_job_reference, + summary=job_summary, + ) + await asyncio.to_thread( + graph.log_node_complete, + step_id, "waiting", summary=job_summary, + ) + clear_step_cancellation(session_id, step_number) + append_session_log_entry(tool_context, { + "kind": "step_complete", + **step_input_log, + "status": "waiting", + "message": job_summary, + "remote_job": remote_job_reference, + "events": event_log, + }, artifacts=artifact_paths) + await asyncio.to_thread( + finish_node_attempt, + recovery_attempt, + status="waiting", + artifacts=artifact_paths, + message=job_summary, + ) + return { + "status": "waiting", + "remote_job": remote_job_reference, + "message": job_summary, + } + if timed_out: logger.warning( "[TIMEOUT] Step %d timed out after %ds (session=%s)", diff --git a/src/matcreator/agents/graph_logger.py b/src/matcreator/agents/graph_logger.py index bc196a5d..8e5c83a5 100644 --- a/src/matcreator/agents/graph_logger.py +++ b/src/matcreator/agents/graph_logger.py @@ -42,7 +42,7 @@ _session_locks: dict[str, threading.Lock] = {} _session_locks_mutex = threading.Lock() -NodeStatus = Literal["idle", "running", "success", "failed", "cancelled", "needs_replanning"] +NodeStatus = Literal["idle", "running", "waiting", "success", "failed", "cancelled", "needs_replanning"] NodeType = Literal["orchestrator", "planning", "execution", "tester", "step"] From b92bec5c6ecfd6a38440b863f7a7304118adda44 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 28 Jul 2026 08:14:28 +0000 Subject: [PATCH 2/2] Make remote-job re-attachment deterministic across planner, orchestrator and flash mode --- docs/remote_job_monitoring.md | 34 +++ .../agents/execution_agent/agent.py | 9 +- .../agents/execution_agent/recovery.py | 31 ++- .../agents/execution_agent/step_executor.py | 9 + src/matcreator/agents/thinking_agent/agent.py | 29 ++- .../agents/thinking_agent/planning.py | 22 +- tests/test_execution_recovery.py | 213 +++++++++++++++++- tests/test_step_executor_runner.py | 128 +++++++++++ 8 files changed, 464 insertions(+), 11 deletions(-) diff --git a/docs/remote_job_monitoring.md b/docs/remote_job_monitoring.md index 6ad190ce..f96fce1d 100644 --- a/docs/remote_job_monitoring.md +++ b/docs/remote_job_monitoring.md @@ -101,6 +101,40 @@ to the agent even when a monitor probe updates the job while that command runs. Strict revision checks remain in place for lifecycle transitions and provider reconciliation, where accepting stale state would be unsafe. +## Executor Timeout and Remote-Job Handoff + +A step executor is a bounded LLM session; a remote job is durable. The two have +independent lifetimes, so a step executor is never kept alive merely to babysit +a running job. + +When `SUB_STEP_TIMEOUT` (default 3600s) elapses, the runner checks the durable +job store for a job still owned by that node: + +- **No active job** — the step times out as before and returns + `needs_replanning`. +- **An active job** — the executor is granted a single bounded grace window + (`STEP_REMOTE_JOB_GRACE_TIMEOUT`, default 300s) to let a nearly finished step + complete. If it is still unfinished afterwards, the executor is released and + the step returns `waiting` rather than `needs_replanning`. This is a handoff, + not a failure: dependents are not blocked, and the job keeps running with no + executor attached. + +The runner writes `status: waiting` and the job identity onto the execution +graph node itself, so the handoff does not depend on the orchestrator LLM +calling `set_node_status`. `reconcile_recovery_state` then keeps the node +`waiting` while the job is still in progress and moves it back to `pending` +once the job settles, so a fresh executor can collect its results. The identical +path also covers a crashed or restarted executor, so there is one recovery +mechanism rather than two. + +Re-attachment is explicit rather than accidental. When a node that already owns +a job runs again, the runner injects the job's identity into the executor's +`prior_context` with instructions to call `get_e2b_job_status` and never call +`submit_e2b_sandbox` for that step. In Flash mode, which has no execution graph, +a step's node ID is derived from its label or a hash of its action, so a repeated +step keeps the same submission idempotency key and re-attaches instead of +creating a duplicate sandbox. + ## Controls and Ownership The middleware exposes owner-scoped controls: diff --git a/src/matcreator/agents/execution_agent/agent.py b/src/matcreator/agents/execution_agent/agent.py index 67c82415..ff5a4560 100644 --- a/src/matcreator/agents/execution_agent/agent.py +++ b/src/matcreator/agents/execution_agent/agent.py @@ -59,15 +59,22 @@ a. Call `set_node_status(node_id="...", status="failed", result=)`. b. Call `mark_dependents_blocked(failed_node_id="...")`. c. Call `to_planner(reason=)`. STOP — do not run any further nodes. + - `status == "waiting"`: the executor timed out but its remote job is STILL RUNNING. + This is NOT a failure. The node status and its `remote_job` are already recorded. + a. Do NOT call `set_node_status`, `mark_dependents_blocked`, or re-run the node. + b. Do NOT start a replacement job for it. + c. Continue processing the other results in this batch, then follow step 6. - `status == "cancelled"`: a. Call `to_planner("execution cancelled by user")`. STOP. 5. After handling all results from a batch, call `get_ready_nodes()` again. -6. If it returns count == 0 and any graph node has status `waiting`, call +6. If it returns count == 0 and `waiting_nodes` is non-empty, call `to_planner(reason="waiting for remote job completion")`. STOP. 7. Execution is complete only when every graph node has status `success`. ## Rules - NEVER execute code directly — all work goes through `run_node_executor`. +- A ready node that carries a `remote_job` was resumed from a remote-job handoff. Run it + normally: its executor re-attaches to that job automatically and must never resubmit it. - When a batch has both successes and one failure, process the success nodes first (validate_summarize + set_node_status), then handle the failure last (set_node_status, mark_dependents_blocked, to_planner). diff --git a/src/matcreator/agents/execution_agent/recovery.py b/src/matcreator/agents/execution_agent/recovery.py index dcfe554e..5329b83a 100644 --- a/src/matcreator/agents/execution_agent/recovery.py +++ b/src/matcreator/agents/execution_agent/recovery.py @@ -14,6 +14,23 @@ _RECOVERY_DIR = "recovery" _STALE_AFTER_SECONDS = int(os.environ.get("STEP_RECOVERY_STALE_AFTER", "60")) +# A job in one of these states is still doing work on the provider, so no +# executor can make progress on it yet. ``succeeded``/``collecting`` are +# deliberately excluded: those mean results are ready to be collected, which is +# exactly the point at which a node should run again. +_REMOTE_JOB_IN_PROGRESS_STATUSES = frozenset( + { + "created", + "submitting", + "queued", + "running", + "pause_requested", + "paused", + "resume_requested", + "resuming", + } +) + def _now() -> str: return datetime.now(timezone.utc).isoformat() @@ -245,7 +262,11 @@ def _refresh_remote_job_reference(attempt: dict[str, Any]) -> None: attempt["remote_job"] = persisted["remote_job"] -def _active_remote_job(attempt: dict[str, Any]) -> dict[str, Any] | None: +def _active_remote_job( + attempt: dict[str, Any], + *, + statuses: frozenset[str] = ACTIVE_REMOTE_JOB_STATUSES, +) -> dict[str, Any] | None: reference = attempt.get("remote_job") if not isinstance(reference, dict) or not reference.get("job_id"): return None @@ -257,7 +278,7 @@ def _active_remote_job(attempt: dict[str, Any]) -> dict[str, Any] | None: job and job.get("session_id") == attempt.get("session_id") and job.get("node_id") == attempt.get("node_id") - and job.get("status") in ACTIVE_REMOTE_JOB_STATUSES + and job.get("status") in statuses ): return job return None @@ -369,7 +390,9 @@ def reconcile_recovery_state( and node_status in ("pending", "running", "waiting") and not resumed_by_planner ): - remote_job = _active_remote_job(attempt) + remote_job = _active_remote_job( + attempt, statuses=_REMOTE_JOB_IN_PROGRESS_STATUSES + ) if remote_job is not None: node["status"] = "waiting" node["result"] = _attempt_summary(attempt) @@ -382,7 +405,7 @@ def reconcile_recovery_state( actions.append({"node_id": node_id, "action": "wait_for_remote_job", "status": "waiting"}) else: node["status"] = "pending" - node["result"] = "Remote job is no longer active; re-running node to collect its results." + node["result"] = "Remote job is no longer in progress; re-running node to collect its results." node["recovery"] = { "attempt": attempt.get("attempt"), "status": "remote_job_settled", diff --git a/src/matcreator/agents/execution_agent/step_executor.py b/src/matcreator/agents/execution_agent/step_executor.py index 07a53bdc..2e822f6c 100644 --- a/src/matcreator/agents/execution_agent/step_executor.py +++ b/src/matcreator/agents/execution_agent/step_executor.py @@ -135,6 +135,15 @@ def _fill_missing_fields(self) -> "StepExecutorResult": 3. **If submission.json exists but outputs are missing**, reuse the same submission file (dpdispatcher is idempotent — it skips completed tasks). Do NOT regenerate submission.json. 4. **Never resubmit a job that already completed** — this wastes GPU time and creates duplicate training runs. +## Re-attaching to an existing E2B job (CRITICAL) +If your `prior_context` contains "REMOTE JOB ALREADY SUBMITTED", a tracked sandbox job +for this exact step is already running: +1. Call `get_e2b_job_status` with the given `job_id` FIRST. +2. NEVER call `submit_e2b_sandbox` for that step — it would duplicate a running job. +3. If the job finished, collect its outputs with `download_e2b_output` and report success. +4. If it is still running, call `submit_step_result(status="needs_replanning", ...)` stating + that the job has not finished yet and quoting its job_id and status. + ## User controls for E2B sandboxes `get_e2b_job_status` may return `user_control` when the user paused or terminated the sandbox from the UI. This does not cancel your executor. Treat it as the diff --git a/src/matcreator/agents/thinking_agent/agent.py b/src/matcreator/agents/thinking_agent/agent.py index 351d19e2..ac5a5e53 100644 --- a/src/matcreator/agents/thinking_agent/agent.py +++ b/src/matcreator/agents/thinking_agent/agent.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import hashlib import logging import threading from typing import Optional @@ -214,13 +215,22 @@ async def run_flash_step( Args: action: What to do (same semantics as a DAG node action). suggested_skills: Skill names to preload in the executor. - label: Optional display name shown in the agent graph. + label: Optional display name shown in the agent graph. Reuse the SAME label + when re-running the same logical task so any remote job it submitted is + re-attached rather than duplicated. """ from ..execution_agent.step_executor_runner import run_step_executor counter = tool_context.state.get("_flash_step_counter", 0) + 1 tool_context.state["_flash_step_counter"] = counter - node_id = label.lower().replace(" ", "_")[:40] if label else f"flash_{counter}" + # The node ID is the identity a remote job is tracked under, so it must be + # stable across retries of the same logical task. Deriving the fallback from + # the action (rather than an incrementing counter) means a repeated step + # re-attaches to its existing job instead of submitting a duplicate one. + if label: + node_id = label.lower().replace(" ", "_")[:40] + else: + node_id = f"flash_{hashlib.sha256(action.encode()).hexdigest()[:12]}" return await run_step_executor( step_number=counter, @@ -259,6 +269,12 @@ async def run_flash_step( - Be concise and responsive. - Do NOT call `validate_graph` or `confirm_plan_and_start_execution`. - Quote exact error messages and propose concrete solutions when something fails. +- A `run_flash_step` result with `status == "waiting"` is NOT a failure: the executor hit + its timeout while the step's `remote_job` is still running on the provider. Report the + job identity to the user and stop. Do NOT re-run the step to "restart" the job. + To check on it later, call `run_flash_step` again with the SAME `label` and an action + that says to call `get_e2b_job_status` with that job_id and collect results if finished — + never to submit a new sandbox. """ _NORMAL_INSTRUCTION = """ @@ -336,6 +352,15 @@ async def run_flash_step( ## Reviewing execution history - Use the current `execution_graph` state for normal replanning; avoid re-running nodes that already succeeded. +- A node with status `waiting` is NOT a failure: its step executor was released at its + timeout while the node's `remote_job` is still running on the provider. The job is + durable and keeps running without any executor attached. + - Do NOT replan around it, do NOT add a replacement node, and do NOT plan a new job + submission for it — that would duplicate the running job. + - Tell the user the job is still running and report its `remote_job` identity. + - When the user asks to check on it or the job is expected to be done, call + `set_node_status(node_id="...", status="pending")` and hand control back so the + node runs again; its executor automatically re-attaches to the same job. - As a last-resort debug path, call `read_session_log(view="overview")` first to inspect the coarse executor graph. Only then call `read_session_log(view="detail", step_id="...")` or `node_id="..."` for one executor of interest; do not request bulk detail by default. diff --git a/src/matcreator/agents/thinking_agent/planning.py b/src/matcreator/agents/thinking_agent/planning.py index 7558870b..5c54e52c 100644 --- a/src/matcreator/agents/thinking_agent/planning.py +++ b/src/matcreator/agents/thinking_agent/planning.py @@ -326,14 +326,30 @@ def get_ready_nodes(tool_context: ToolContext) -> dict: nodes.get(pred, {}).get("status") == "success" for pred in predecessors.get(node_id, set()) ): - ready.append({ + ready_node = { "node_id": node_id, "label": node.get("label", node_id), "action": node.get("action", ""), "suggested_skills": node.get("suggested_skills", []), - }) + } + # A node resumed from a remote-job handoff carries the job identity so + # its executor re-attaches instead of submitting a duplicate job. + if isinstance(node.get("remote_job"), dict): + ready_node["remote_job"] = node["remote_job"] + ready.append(ready_node) + + waiting = [ + {"node_id": node_id, "remote_job": node.get("remote_job")} + for node_id, node in nodes.items() + if node.get("status") == "waiting" + ] - return {"status": "ok", "ready_nodes": ready, "count": len(ready)} + return { + "status": "ok", + "ready_nodes": ready, + "count": len(ready), + "waiting_nodes": waiting, + } # --------------------------------------------------------------------------- diff --git a/tests/test_execution_recovery.py b/tests/test_execution_recovery.py index 47e3ba77..a41072db 100644 --- a/tests/test_execution_recovery.py +++ b/tests/test_execution_recovery.py @@ -1,12 +1,15 @@ from __future__ import annotations +import asyncio from datetime import datetime, timedelta, timezone import json from pathlib import Path import sqlite3 from types import SimpleNamespace +from unittest import mock from matcreator.agents.execution_agent.recovery import ( + active_remote_job_for_attempt, finish_node_attempt, heartbeat_node_attempt, record_remote_job_reference, @@ -15,10 +18,11 @@ ) from matcreator.agents.execution_graph_state import get_execution_graph, set_execution_graph from matcreator.control_plane.remote_jobs import RemoteJobStore -from matcreator.agents.thinking_agent.planning import validate_graph +from matcreator.agents.thinking_agent.planning import get_ready_nodes, validate_graph from matcreator.agents.thinking_agent.agent import ( confirm_plan_and_start_execution, resume_execution, + run_flash_step, ) @@ -446,3 +450,210 @@ def test_reconcile_waits_for_active_remote_job_instead_of_resubmitting(tmp_path, "external_id": "sandbox-123", "status": "running", } + + +def _running_remote_job(store, *, session_id="session-1", node_id="node-a"): + job = store.create_job( + owner_id="alice", + session_id=session_id, + node_id=node_id, + provider="e2b", + idempotency_key=f"{session_id}:{node_id}:1", + ) + submitting = store.transition_job(job["job_id"], "submitting") + return store.transition_job( + job["job_id"], + "running", + external_id="sandbox-123", + expected_revision=submitting["state_revision"], + ) + + +def _waiting_attempt(tmp_path, recovery_dir, job): + attempt = start_node_attempt( + workspace_dir=tmp_path, + session_id="session-1", + node_id="node-a", + step_id="execution_0__node_node-a", + step_number=1, + action="do remote work", + suggested_skills=["e2b"], + prior_context=None, + recovery_base_dir=recovery_dir, + ) + record_remote_job_reference( + session_id="session-1", + node_id="node-a", + job_id=job["job_id"], + provider="e2b", + external_id="sandbox-123", + recovery_base_dir=recovery_dir, + ) + heartbeat_node_attempt(attempt) + finish_node_attempt( + attempt, + status="waiting", + message="Executor released; remote job still running.", + ) + return attempt + + +def test_active_remote_job_is_visible_to_a_live_attempt(tmp_path, monkeypatch): + recovery_dir = tmp_path / "adk-recovery" + adk_dir = tmp_path / "adk" + monkeypatch.setattr("matcreator.agents.execution_agent.recovery.ADK_DIR", adk_dir) + store = RemoteJobStore(adk_dir / "remote-jobs.db") + job = _running_remote_job(store) + attempt = start_node_attempt( + workspace_dir=tmp_path, + session_id="session-1", + node_id="node-a", + step_id="execution_0__node_node-a", + step_number=1, + action="do remote work", + suggested_skills=["e2b"], + prior_context=None, + recovery_base_dir=recovery_dir, + ) + record_remote_job_reference( + session_id="session-1", + node_id="node-a", + job_id=job["job_id"], + provider="e2b", + external_id="sandbox-123", + recovery_base_dir=recovery_dir, + ) + + # The in-memory attempt has not heartbeated since submission, so the lookup + # must reload the reference from the durable record. + assert "remote_job" not in attempt + active = active_remote_job_for_attempt(attempt) + + assert active is not None + assert active["job_id"] == job["job_id"] + + +def test_waiting_attempt_keeps_node_waiting_while_remote_job_is_active(tmp_path, monkeypatch): + recovery_dir = tmp_path / "adk-recovery" + adk_dir = tmp_path / "adk" + monkeypatch.setattr("matcreator.agents.execution_agent.recovery.ADK_DIR", adk_dir) + store = RemoteJobStore(adk_dir / "remote-jobs.db") + job = _running_remote_job(store) + _waiting_attempt(tmp_path, recovery_dir, job) + state = { + "session_id": "session-1", + "execution_graph": {"nodes": {"node-a": {"status": "running"}}, "edges": []}, + } + + recovered = reconcile_recovery_state( + state, tmp_path, stale_after_seconds=1, recovery_base_dir=recovery_dir + ) + + assert recovered == [{"node_id": "node-a", "action": "wait_for_remote_job", "status": "waiting"}] + node = get_execution_graph(state)["nodes"]["node-a"] + assert node["status"] == "waiting" + assert node["remote_job"]["job_id"] == job["job_id"] + + +def test_waiting_node_becomes_pending_once_its_remote_job_settles(tmp_path, monkeypatch): + recovery_dir = tmp_path / "adk-recovery" + adk_dir = tmp_path / "adk" + monkeypatch.setattr("matcreator.agents.execution_agent.recovery.ADK_DIR", adk_dir) + store = RemoteJobStore(adk_dir / "remote-jobs.db") + job = _running_remote_job(store) + _waiting_attempt(tmp_path, recovery_dir, job) + store.transition_job(job["job_id"], "succeeded") + state = { + "session_id": "session-1", + "execution_graph": {"nodes": {"node-a": {"status": "waiting"}}, "edges": []}, + } + + recovered = reconcile_recovery_state( + state, tmp_path, stale_after_seconds=1, recovery_base_dir=recovery_dir + ) + + assert recovered == [ + {"node_id": "node-a", "action": "resume_after_remote_job", "status": "pending"} + ] + assert get_execution_graph(state)["nodes"]["node-a"]["status"] == "pending" + + +def test_reconcile_does_not_undo_a_planner_resume_of_a_waiting_node(tmp_path, monkeypatch): + recovery_dir = tmp_path / "adk-recovery" + adk_dir = tmp_path / "adk" + monkeypatch.setattr("matcreator.agents.execution_agent.recovery.ADK_DIR", adk_dir) + store = RemoteJobStore(adk_dir / "remote-jobs.db") + job = _running_remote_job(store) + _waiting_attempt(tmp_path, recovery_dir, job) + # The planner explicitly moved the waiting node back to pending so a fresh + # executor can re-attach and poll the still-running job. + state = { + "session_id": "session-1", + "execution_graph": { + "nodes": { + "node-a": { + "status": "pending", + "recovery": {"status": "waiting_remote"}, + "remote_job": {"job_id": job["job_id"]}, + } + }, + "edges": [], + }, + } + + recovered = reconcile_recovery_state( + state, tmp_path, stale_after_seconds=1, recovery_base_dir=recovery_dir + ) + + assert recovered == [] + assert get_execution_graph(state)["nodes"]["node-a"]["status"] == "pending" + + +def test_ready_nodes_carry_the_remote_job_a_resumed_node_must_re_attach_to(): + tool_context = SimpleNamespace(state={}) + set_execution_graph(tool_context.state, { + "nodes": { + "step_train": { + "status": "pending", + "label": "Train", + "action": "train the model", + "suggested_skills": [], + "remote_job": { + "job_id": "job-1", + "provider": "e2b", + "external_id": "sandbox-123", + "status": "running", + }, + }, + "step_wait": {"status": "waiting", "remote_job": {"job_id": "job-2"}}, + }, + "edges": [], + }) + + result = get_ready_nodes(tool_context) + + assert result["count"] == 1 + assert result["ready_nodes"][0]["remote_job"]["job_id"] == "job-1" + assert result["waiting_nodes"] == [ + {"node_id": "step_wait", "remote_job": {"job_id": "job-2"}} + ] + + +def test_flash_step_node_id_is_stable_across_retries_of_the_same_action(): + captured: list[str] = [] + + async def fake_run_step_executor(**kwargs): + captured.append(kwargs["node_id"]) + return {"status": "success"} + + tool_context = SimpleNamespace(state={}) + with mock.patch( + "matcreator.agents.execution_agent.step_executor_runner.run_step_executor", + fake_run_step_executor, + ): + asyncio.run(run_flash_step("submit the training job", [], tool_context)) + asyncio.run(run_flash_step("submit the training job", [], tool_context)) + asyncio.run(run_flash_step("a different action", [], tool_context)) + + assert captured[0] == captured[1] + assert captured[2] != captured[0] diff --git a/tests/test_step_executor_runner.py b/tests/test_step_executor_runner.py index 98393346..b5663b0f 100644 --- a/tests/test_step_executor_runner.py +++ b/tests/test_step_executor_runner.py @@ -14,6 +14,8 @@ _schedule_step_runner_cleanup, _verify_step_result_artifacts, ) +from matcreator.agents.execution_agent import step_executor_runner +from matcreator.agents.execution_graph_state import get_execution_graph, set_execution_graph from matcreator.agents.session_log import SESSION_ARTIFACTS_KEY @@ -295,3 +297,129 @@ def test_step_executor_retries_only_malformed_streamed_tool_arguments(): assert agent.retry_config is not None assert agent.retry_config.max_attempts == 2 assert agent.retry_config.exceptions == ["JSONDecodeError"] + + +# --------------------------------------------------------------------------- +# Remote-job handoff on timeout +# --------------------------------------------------------------------------- + + +class _MutableState(dict): + """Minimal stand-in for ADK session state used by the graph state helpers.""" + + +class _FakeToolContext: + def __init__(self, state): + self.state = state + + +def _run_await_step_completion(monkeypatch, *, remote_job, grace, hold=0.2): + monkeypatch.setattr(step_executor_runner, "_SUB_STEP_TIMEOUT", 0.01) + monkeypatch.setattr(step_executor_runner, "_REMOTE_JOB_GRACE_TIMEOUT", grace) + monkeypatch.setattr( + step_executor_runner, + "active_remote_job_for_attempt", + lambda attempt: remote_job, + ) + + async def exercise(): + async def slow_task(): + await asyncio.sleep(hold) + return ({}, [], [], {}) + + task = asyncio.create_task(slow_task()) + try: + return await step_executor_runner._await_step_completion( + task, {}, step_number=1, session_id="session-1" + ) + finally: + task.cancel() + + return asyncio.run(exercise()) + + +def test_timeout_without_remote_job_is_a_plain_timeout(monkeypatch): + timed_out, remote_job = _run_await_step_completion( + monkeypatch, remote_job=None, grace=5 + ) + + assert timed_out is True + assert remote_job is None + + +def test_timeout_with_active_remote_job_reports_it_after_one_grace_window(monkeypatch): + job = {"job_id": "job-1", "provider": "e2b", "external_id": "sbx-1", "status": "running"} + + timed_out, remote_job = _run_await_step_completion( + monkeypatch, remote_job=job, grace=0.01 + ) + + assert timed_out is True + assert remote_job == job + + +def test_grace_window_lets_a_nearly_finished_step_complete(monkeypatch): + job = {"job_id": "job-1", "provider": "e2b", "external_id": "sbx-1", "status": "running"} + + timed_out, remote_job = _run_await_step_completion( + monkeypatch, remote_job=job, grace=5, hold=0.05 + ) + + assert timed_out is False + assert remote_job is None + + +def test_waiting_handoff_records_remote_job_on_the_graph_node(): + state = _MutableState() + set_execution_graph(state, {"nodes": {"step_train": {"status": "running"}}, "edges": []}) + tool_context = _FakeToolContext(state) + job = {"job_id": "job-1", "provider": "e2b", "external_id": "sbx-1", "status": "running"} + + step_executor_runner._mark_node_waiting_on_remote_job( + tool_context, + node_id="step_train", + remote_job=job, + summary="Executor released; job still running.", + ) + + node = get_execution_graph(state)["nodes"]["step_train"] + assert node["status"] == "waiting" + assert node["remote_job"] == job + assert node["recovery"]["status"] == "waiting_remote" + + +def test_resumed_node_receives_explicit_reattachment_instructions(): + state = _MutableState() + set_execution_graph(state, { + "nodes": { + "step_train": { + "status": "pending", + "remote_job": { + "job_id": "job-1", + "provider": "e2b", + "external_id": "sbx-1", + "status": "running", + }, + } + }, + "edges": [], + }) + + context = step_executor_runner._remote_job_prior_context( + _FakeToolContext(state), "step_train" + ) + + assert context is not None + assert "job-1" in context + assert "sbx-1" in context + assert "get_e2b_job_status" in context + assert "Do NOT call submit_e2b_sandbox" in context + + +def test_node_without_remote_job_gets_no_reattachment_instructions(): + state = _MutableState() + set_execution_graph(state, {"nodes": {"step_train": {"status": "pending"}}, "edges": []}) + + assert step_executor_runner._remote_job_prior_context( + _FakeToolContext(state), "step_train" + ) is None