From 7558a13a762d04995cea573916120b59f82e881a Mon Sep 17 00:00:00 2001 From: breken-ai Date: Tue, 8 Sep 2026 23:22:55 -0700 Subject: [PATCH] fix: reset agent retry budget only on real progress The resume-delay budget (30/60/120s) resets whenever the dying agent printed ANY assistant line. A crash-looping agent (e.g. expired OAuth, one line then exit) makes identical progress every attempt, so the retry loop never terminates. Replace the made_progress boolean with an assistant-line count and reset the budget only when an attempt beats the best previous count. A crash loop drains the retries and returns the error; a genuinely progressing agent still earns resets. Fixes #4163 --- .../_internal/cli/services/presets/agent.py | 20 +++++-- .../cli/services/presets/test_agent.py | 59 +++++++++++++++++++ 2 files changed, 73 insertions(+), 6 deletions(-) diff --git a/src/dstack/_internal/cli/services/presets/agent.py b/src/dstack/_internal/cli/services/presets/agent.py index 70d22be98a..00e6134924 100644 --- a/src/dstack/_internal/cli/services/presets/agent.py +++ b/src/dstack/_internal/cli/services/presets/agent.py @@ -101,7 +101,10 @@ class PresetAgentProcessOutput: report_data: Optional[dict[str, Any]] = None error: Optional[str] = None session_id: Optional[str] = None - made_progress: bool = False + # Assistant lines seen this attempt. A count, not a flag: a crash-looping + # agent reaches the same point every attempt, and only an attempt that got + # FURTHER than any before it earns a fresh retry budget. + assistant_events: int = 0 def _get_claude_version(auth: "ClaudeAuth") -> Optional[str]: @@ -225,6 +228,7 @@ async def run_preset_agent( resume_session_id: Optional[str] = initial_resume_session_id attempt_prompt = prompt if resume_session_id is None else _RESUME_PROMPT retry_delays = list(_RESUME_DELAYS_SECONDS) + best_progress = 0 while True: command = _prepare_subprocess_command( _build_claude_command(auth=auth, resume_session_id=resume_session_id) @@ -245,10 +249,14 @@ async def run_preset_agent( # failure report from the agent returns immediately. if output.report_data is not None or error is None: return output - # Only reset the retry budget when the last attempt made progress; a - # run that keeps stalling exhausts its retries instead of retrying a - # stuck agent forever. - if output.made_progress: + # Only reset the retry budget when the last attempt made NEW + # progress; a run that keeps stalling exhausts its retries instead + # of retrying a stuck agent forever. An agent that dies at startup + # after printing one assistant line makes the same "progress" on + # every attempt, so a boolean flag would reset the budget forever; + # only an attempt that got further than every previous one resets it. + if output.assistant_events > best_progress: + best_progress = output.assistant_events retry_delays = list(_RESUME_DELAYS_SECONDS) # Another process marked this session interrupted; don't restart it. state = session.read_state() @@ -573,7 +581,7 @@ async def _read_process_stream( if event.model: session.record_agent_model(event.model) if event.type == "assistant": - output.made_progress = True + output.assistant_events += 1 if not isinstance(event, ClaudeResultEvent): continue if event.is_error: diff --git a/src/tests/_internal/cli/services/presets/test_agent.py b/src/tests/_internal/cli/services/presets/test_agent.py index f36940e879..671f9a4706 100644 --- a/src/tests/_internal/cli/services/presets/test_agent.py +++ b/src/tests/_internal/cli/services/presets/test_agent.py @@ -1352,3 +1352,62 @@ def advance(key: str) -> None: reloaded = OffsetStore(state) for key in ("agent_stdout", "agent_stderr", "runs", "trials"): assert reloaded.get(key) == 50 + + +class TestAgentRetryBudget: + @pytest.mark.asyncio + async def test_crash_loop_after_one_assistant_line_is_bounded(self, tmp_path, monkeypatch): + # Regression: an agent that dies without a report but prints one + # `assistant` line each attempt (e.g. claude starts then dies on expired + # OAuth) used to reset the retry budget on every attempt and loop + # forever. Only NEW progress (more assistant lines than any earlier + # attempt) may reset the budget. + counter = tmp_path / "attempts" + script = tmp_path / "fake_claude.py" + script.write_text( + """import json +import sys +from pathlib import Path + +counter = Path(sys.argv[1]) +attempts = int(counter.read_text()) + 1 if counter.exists() else 1 +counter.write_text(str(attempts)) +print(json.dumps({"type": "assistant"}), flush=True) +if attempts > 8: # unstick the run so a broken loop cannot hang the test + print(json.dumps({"type": "result", "structured_output": {"attempts": attempts}}), flush=True) +sys.exit(1) +""" + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent._build_claude_command", + lambda **_: [sys.executable, str(script), str(counter)], + ) + real_sleep = asyncio.sleep + + async def _no_sleep(*_): + await real_sleep(0) + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent.asyncio.sleep", + _no_sleep, + ) + + workspace = PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") + session_path = tmp_path / "session-running" + session_path.mkdir() + (session_path / "agent.log").touch() + (session_path / "trace.jsonl").touch() + session = PresetSession(path=session_path, preset_id="ab12cd34") + output = await run_preset_agent( + prompt="full preset prompt", + env=os.environ.copy(), + workspace=workspace, + auth=_claude_auth(), + redacted_values=(), + session=session, + ) + + attempts = int(counter.read_text()) + # Initial attempt plus the bounded resume delays (30s, 60s, 120s). + assert attempts == 4 + assert output.error is not None