nanoPyCodeAgent 0.8.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

nanopycodeagent/agent.py

@@ -11,18 +11,18 @@ or overwrite them, an ``edit`` tool to replace part of one, and a ``bash``
 tool to run shell commands; every call and its output are echoed to the
 terminal as they happen.
 
-The interactive loop handles only the happy path: anything unexpected — a
-network error, a Ctrl-C mid-turn — crashes the session, and restarting it is
-the recovery. That trade keeps the core flow readable; the hardened variant
-it replaced is preserved at the ``hardened-agent-loop`` tag. A headless run
-has no one to restart it, so it catches API errors, reports them verbatim,
-and turns them into an exit code.
+Both modes retry interrupted response streams before changing conversation
+history or running tools. Other unexpected failures and Ctrl-C mid-turn still
+end an interactive session. Headless mode catches exhausted API/transport
+errors, reports them verbatim, and turns them into an exit code.
 """
 
+import json
 import os
 import sys
 import time
 import uuid
+from copy import copy
 from importlib.metadata import PackageNotFoundError, version
 from pathlib import Path
 
@@ -37,34 +37,47 @@ except ImportError:  # pragma: no cover - platform without readline
     pass
 
 import anthropic
-import httpx
 from anthropic.types import MessageParam, ToolResultBlockParam, ToolUseBlock
 
 from .atif import project_atif, write_atif
-from .bash_tool import BASH_TOOL, run_bash
+from .bash_tool import run_bash
 from .cost import (
+    estimated_cost,
     pending_cost,
     resolve_generation_cost,
     usage_cost,
 )
-from .edit_tool import EDIT_TOOL, edit_preview, run_edit
+from .edit_tool import edit_preview, run_edit
+from .deadline import DeadlineExceeded, check_deadline_support, wall_clock_limit
 from .event_journal import (
     EventEmitter,
     EventJournal,
     JsonObject,
     JsonValue,
     NativeEvent,
+    RunOutcome,
     utc_now,
 )
-from .read_tool import READ_TOOL, run_read
-from .settings import load_settings_env
+from .read_tool import run_read
+from .settings import DEFAULT_MAX_TOKENS, load_settings_env, resolve_max_tokens
 from .terminal import Spinner, print_tool_output, print_tool_use
-from .write_tool import WRITE_TOOL, content_preview, run_write
+from .tool_validation import TOOLS, tool_input_error
+from .transport import HTTP_ERRORS, RETRYABLE_STREAM_ERRORS
+from .write_tool import content_preview, run_write
 
 # The model used when ANTHROPIC_MODEL is set in neither the environment nor
 # the config file.
 DEFAULT_MODEL = "claude-sonnet-4-6"
-MAX_TOKENS = 8192
+
+# The SDK retries failures before streaming begins. Recover interrupted response
+# bodies here, before committing a reply to history or executing any of its tools.
+STREAM_RETRY_DELAYS = (1.0, 2.0)
+STREAM_RETRY_WINDOW_SECONDS = 300.0
+
+_TRUNCATION_NOTICE = (
+    "[response truncated: reached max_tokens; stopped without finishing the task. "
+    "Tool calls from this response were not executed.]"
+)
 
 # How many model replies one headless task may spend before the run stops on
 # its own. The interactive loop needs no such cap — a human watching the
@@ -73,6 +86,91 @@ MAX_TOKENS = 8192
 # refuses it.
 DEFAULT_MAX_TURNS = 50
 
+# A headless run may also be given a wall-clock budget. When it is, the loop
+# tells the model how much time is left and stops before the harness's own
+# timeout can kill the process with nothing written. The last stretch is
+# reserved so the model still has room to write the task's output file.
+_FINALIZATION_RESERVE_SECONDS = 180
+_FINALIZATION_RESERVE_FRACTION = 0.15
+_FINALIZATION_RESERVE_TURNS = 10
+_COST_RECONCILIATION_SECONDS = 30
+
+
+def _format_duration(total_seconds: float) -> str:
+    total = max(0, int(total_seconds))
+    return f"{total // 60}:{total % 60:02d}"
+
+
+def _time_budget_note(
+    *, turn: int, elapsed: float, budget: float | None, remaining: float | None,
+    max_turns: int | None = None,
+) -> str:
+    """Tell the model about both limits before either prevents finalization."""
+    replies_left = None if max_turns is None else max_turns - turn + 1
+    turn_note = f"Turn {turn}. " if max_turns is None else (
+        f"Turn {turn} of {max_turns}; {replies_left} replies remaining including "
+        "this one. The last reply must be a final summary: its tool calls "
+        "will not execute. "
+    )
+    time_low = remaining is not None and remaining <= max(
+        _FINALIZATION_RESERVE_FRACTION * budget, _FINALIZATION_RESERVE_SECONDS
+    )
+    note = "[runtime budget] " + turn_note
+    if remaining is not None:
+        note += (
+            f"Only {_format_duration(remaining)} of {_format_duration(budget)} left. "
+            if time_low else
+            f"Elapsed {_format_duration(elapsed)} of {_format_duration(budget)}; "
+            f"{_format_duration(remaining)} remaining. "
+        )
+    if time_low or (replies_left is not None and replies_left <= _FINALIZATION_RESERVE_TURNS):
+        note += (
+            "Finalize now. Stop investigating new approaches. Complete and save "
+            "the required deliverables, perform only the necessary checks, then "
+            "reply with a short summary and no further tool calls. If incomplete, "
+            "save useful progress and state the remaining limitation honestly."
+        )
+    else:
+        note += (
+            "Keep required deliverables up to date. Once the requirements are "
+            "satisfied and checked, finish immediately; unused budget is not "
+            "a reason to continue investigating."
+        )
+    return note
+
+
+def _append_budget_note(
+    messages: list[MessageParam],
+    note: str,
+    *,
+    emitter: EventEmitter,
+    model_call_id: str,
+    reason: str = "time_budget",
+) -> None:
+    """Append a wall-clock reminder to the tail of the conversation.
+
+    The reminder goes into the most recent user message — the initial task, or
+    the tool results — rather than the system prompt. Rewriting the system
+    prompt each turn changes the very first tokens of every request and defeats
+    the provider's prefix cache; appending to the tail keeps each request an
+    extension of the previous one, so the cached prefix survives.
+    """
+    last = messages[-1]
+    content = last["content"]
+    if isinstance(content, str):
+        last["content"] = f"{content}\n\n{note}"
+    else:
+        content.append({"type": "text", "text": note})
+    emitter.emit(
+        "input.injected",
+        {
+            "model_call_id": model_call_id,
+            "content": note,
+            "reason": reason,
+            "source_timestamp": utc_now(),
+        },
+    )
+
 # Shared by both system prompts: which tool to reach for is the same question
 # whoever is asking.
 _TOOL_GUIDANCE = (
@@ -98,13 +196,12 @@ HEADLESS_SYSTEM_PROMPT = (
     "out. Work the task through to the end, then check the result with the "
     "tools instead of assuming it worked. When it is done, answer with a "
     "short summary and no further tool calls: that reply is what ends the "
-    "run. "
+    "run. Runtime budget reminders report remaining time and model replies; "
+    "when either is running low, prioritize saving the required deliverables, "
+    "necessary verification, and a final summary. Do not start optional work "
+    "after the requirements are satisfied. "
 ) + _TOOL_GUIDANCE
 
-# Every tool offered to the model on each request.
-TOOLS = [READ_TOOL, WRITE_TOOL, EDIT_TOOL, BASH_TOOL]
-
-
 def _json_value(value: object) -> JsonValue:
     """Convert an SDK value into the provider-neutral event representation."""
     if value is None or isinstance(value, bool | int | float | str):
@@ -139,12 +236,19 @@ def _native_content_blocks(value: object) -> list[JsonValue]:
         if source_type == "text":
             content.append({"type": "text", "text": source_block.get("text", "")})
         elif source_type == "tool_use":
+            arguments = source_block.get("input")
             content.append(
                 {
                     "type": "tool_call",
                     "tool_call_id": source_block.get("id"),
                     "tool_name": source_block.get("name"),
-                    "input": source_block.get("input"),
+                    # Journal/ATIF require object arguments. Keep rejected
+                    # non-object input separately instead of discarding it.
+                    "input": arguments if isinstance(arguments, dict) else {},
+                    **(
+                        {"raw_input": arguments}
+                        if not isinstance(arguments, dict) else {}
+                    ),
                 }
             )
         else:
@@ -192,12 +296,25 @@ class _TextOutputProjector:
             model_call_id = str(event.payload["model_call_id"])
             if model_call_id in self._model_calls_with_text:
                 print()
+            if event.payload["stop_reason"] == "max_tokens":
+                print(_TRUNCATION_NOTICE, file=sys.stderr)
+        elif event.type == "model.failed" and event.payload["will_retry"]:
+            if str(event.payload["model_call_id"]) in self._model_calls_with_text:
+                print()
+            print(
+                f"[response interrupted: {event.payload['error_type']}; "
+                f"retrying in {event.payload['retry_delay_seconds']:g}s]",
+                file=sys.stderr,
+            )
         elif event.type == "tool.started":
             tool_name = str(event.payload["tool_name"])
             arguments = event.payload["input"]
-            if not isinstance(arguments, dict):
-                raise TypeError("tool input event payload must be an object")
-            if tool_name == "read":
+            error = event.payload.get("input_error") or tool_input_error(
+                tool_name, arguments
+            )
+            if error:
+                print_tool_use(f"[{tool_name}] (invalid arguments; not executed)")
+            elif tool_name == "read":
                 print_tool_use(f"[read] {arguments['path']}")
             elif tool_name == "write":
                 content = str(arguments["content"])
@@ -211,7 +328,7 @@ class _TextOutputProjector:
                     f"[edit] {arguments['path']}\n"
                     f"{edit_preview(old_text, new_text)}"
                 )
-            else:
+            elif tool_name == "bash":
                 print_tool_use(f"[bash]$ {arguments['command']}")
         elif event.type == "tool.completed":
             result = event.payload["result"]
@@ -237,24 +354,29 @@ def _run_one_tool(
     block: ToolUseBlock,
     emitter: EventEmitter,
     model_call_id: str,
+    *,
+    input_error: str | None = None,
+    remaining_seconds: float | None = None,
 ) -> ToolResultBlockParam:
     """Execute one ``tool_use`` block and emit its runtime facts."""
     tool_input = _json_value(block.input)
-    if not isinstance(tool_input, dict):
-        raise TypeError("tool input must be an object")
+    input_error = input_error or tool_input_error(block.name, tool_input)
     emitter.emit(
         "tool.started",
         {
             "model_call_id": model_call_id,
             "tool_call_id": block.id,
             "tool_name": block.name,
-            "input": tool_input,
+            "input": tool_input if isinstance(tool_input, dict) else {},
+            **({"input_error": input_error} if input_error else {}),
             "source_timestamp": utc_now(),
         },
     )
     tool_started_ns = time.perf_counter_ns()
     try:
-        if block.name == "read":
+        if input_error:
+            output, is_error = input_error, True
+        elif block.name == "read":
             path = block.input["path"]
             output, is_error = run_read(
                 path,
@@ -275,10 +397,13 @@ def _run_one_tool(
                 new_text,
                 replace_all=block.input.get("replace_all", False),
             )
-        else:  # bash — the only other tool offered
+        else:  # bash; unknown names have already been rejected
             command = block.input["command"]
             with Spinner("Running..."):
-                output, is_error = run_bash(command)
+                output, is_error = run_bash(command, **(
+                    {"timeout_seconds": remaining_seconds}
+                    if remaining_seconds is not None else {}
+                ))
     except BaseException as exc:
         emitter.emit(
             "tool.completed",
@@ -303,6 +428,10 @@ def _run_one_tool(
             "tool_name": block.name,
             "result": output,
             "is_error": is_error,
+            **(
+                {"error": {"type": "ToolInputError", "message": input_error}}
+                if input_error else {}
+            ),
             "duration_ms": (time.perf_counter_ns() - tool_started_ns) / 1_000_000,
             "source_timestamp": utc_now(),
         },
@@ -352,17 +481,21 @@ def _run_exchange(
     system: str,
     *,
     max_turns: int | None = None,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+    time_budget_seconds: int | None = None,
     reply_prefix: str = "\nAgent> ",
     trajectory_path: Path | None = None,
-) -> bool:
+) -> RunOutcome:
     """Reply to the conversation so far, running tools until the model stops.
 
-    Appends every assistant reply and tool result to ``messages`` in place.
-    Returns True when the model ended a reply without asking for tools, and
-    False when ``max_turns`` replies were spent while it was still calling
-    them — the caller decides what an exhausted budget means.
+    Appends assistant replies and tool results to ``messages`` in place.
+    A truncated reply retains only its text and a notice in request history;
+    the original response is kept in the journal. Returns the stopping outcome,
+    distinguishing completion, turn-budget exhaustion, and response truncation.
     """
     run_id = f"run-{uuid.uuid4()}"
+    if time_budget_seconds is not None:
+        check_deadline_support()
     run_started_ns = time.perf_counter_ns()
     projector = _TextOutputProjector(reply_prefix)
     with EventJournal.create(run_id) as journal:
@@ -370,9 +503,15 @@ def _run_exchange(
         emitter.emit(
             "run.started",
             {
-                "mode": "headless" if max_turns is not None else "interactive",
+                "mode": (
+                    "headless"
+                    if max_turns is not None or time_budget_seconds is not None
+                    else "interactive"
+                ),
                 "model": model,
                 "max_turns": max_turns,
+                "max_tokens": max_tokens,
+                "time_budget_seconds": time_budget_seconds,
                 "producer": {
                     "name": "nanoPyCodeAgent",
                     "version": _package_version(),
@@ -390,16 +529,21 @@ def _run_exchange(
             },
         )
         try:
-            finished = _run_model_loop(
+            outcome = _run_model_loop(
                 client,
                 model,
                 messages,
                 system,
                 emitter=emitter,
                 max_turns=max_turns,
+                max_tokens=max_tokens,
+                time_budget_seconds=time_budget_seconds,
             )
         except BaseException as exc:
-            cost_reconciliation = _reconcile_costs(client, journal, emitter)
+            cost_reconciliation = _reconcile_costs(
+                client, journal, emitter,
+                max_seconds=_COST_RECONCILIATION_SECONDS if time_budget_seconds else None,
+            )
             emitter.emit(
                 "run.failed",
                 {
@@ -417,11 +561,14 @@ def _run_exchange(
             )
             raise
         else:
-            cost_reconciliation = _reconcile_costs(client, journal, emitter)
+            cost_reconciliation = _reconcile_costs(
+                client, journal, emitter,
+                max_seconds=_COST_RECONCILIATION_SECONDS if time_budget_seconds else None,
+            )
             emitter.emit(
                 "run.completed",
                 {
-                    "outcome": "completed" if finished else "max_turns_exhausted",
+                    "outcome": outcome,
                     "duration_ms": (time.perf_counter_ns() - run_started_ns)
                     / 1_000_000,
                     **(
@@ -438,7 +585,7 @@ def _run_exchange(
                     project_atif(EventJournal.replay(journal.path)),
                     trajectory_path,
                 )
-        return finished
+        return outcome
 
 
 def _run_model_loop(
@@ -449,14 +596,42 @@ def _run_model_loop(
     *,
     emitter: EventEmitter,
     max_turns: int | None,
-) -> bool:
+    max_tokens: int,
+    time_budget_seconds: int | None = None,
+) -> RunOutcome:
     """Run model replies and tool calls for an already-started Agent Run."""
     turns = 0
+    retries = 0
+    retry_deadline = None
+    deadline = (
+        time.monotonic() + time_budget_seconds
+        if time_budget_seconds is not None and time_budget_seconds > 0
+        else None
+    )
     while True:
+        model_call_id = f"model-{uuid.uuid4()}"
+        remaining = None
+        if deadline is not None:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                return "time_budget_exhausted"
+        if deadline is not None or (max_turns is not None and retries == 0):
+            _append_budget_note(
+                messages,
+                _time_budget_note(
+                    turn=turns + 1,
+                    elapsed=time_budget_seconds - remaining if deadline is not None else 0,
+                    budget=time_budget_seconds,
+                    remaining=remaining,
+                    max_turns=max_turns,
+                ),
+                emitter=emitter,
+                model_call_id=model_call_id,
+                reason="time_budget" if deadline is not None else "turn_budget",
+            )
         # A spinner marks the wait for the reply; the first streamed
         # token replaces it with the reply prefix. A tool-only reply
         # streams no text, so the prefix is skipped for it entirely.
-        model_call_id = f"model-{uuid.uuid4()}"
         emitter.emit(
             "model.started",
             {
@@ -468,28 +643,111 @@ def _run_model_loop(
         model_started_ns = time.perf_counter_ns()
         # Stream the reply so text shows up as it is generated, then grab
         # the accumulated message for the conversation history.
-        with Spinner() as spinner, client.messages.stream(
-            model=model,
-            max_tokens=MAX_TOKENS,
-            system=system,
-            tools=TOOLS,
-            messages=messages,
-        ) as stream:
-            for text in stream.text_stream:
-                spinner.stop()
-                emitter.emit(
-                    "model.output_delta",
-                    {
-                        "model_call_id": model_call_id,
-                        "delta": text,
-                        "source_timestamp": utc_now(),
-                    },
-                )
-            message = stream.get_final_message()
-            generation_id = _response_header(stream, "x-generation-id")
-            model_completed_ns = time.perf_counter_ns()
+        generation_id = None
+        stream_entered = False
+        try:
+            with wall_clock_limit(remaining), Spinner() as spinner, client.messages.stream(
+                model=model,
+                max_tokens=max_tokens,
+                system=system,
+                tools=TOOLS,
+                messages=messages,
+                **({"timeout": remaining} if remaining is not None else {}),
+            ) as stream:
+                stream_entered = True
+                generation_id = _response_header(stream, "x-generation-id")
+                input_json: dict[int, list[str]] = {}
+                for event in stream:
+                    if event.type == "text":
+                        spinner.stop()
+                        emitter.emit(
+                            "model.output_delta",
+                            {
+                                "model_call_id": model_call_id,
+                                "delta": event.text,
+                                "source_timestamp": utc_now(),
+                            },
+                        )
+                    elif (
+                        event.type == "content_block_delta"
+                        and event.delta.type == "input_json_delta"
+                    ):
+                        input_json.setdefault(event.index, []).append(
+                            event.delta.partial_json
+                        )
+                message = stream.get_final_message()
+                model_completed_ns = time.perf_counter_ns()
+        except DeadlineExceeded as exc:
+            emitter.emit("model.failed", {
+                "model_call_id": model_call_id,
+                "error_type": type(exc).__name__, "message": str(exc),
+                "generation_id": generation_id,
+                "duration_ms": (time.perf_counter_ns() - model_started_ns) / 1_000_000,
+                "will_retry": False, "retry_delay_seconds": 0,
+                "source_timestamp": utc_now(),
+            })
+            return "time_budget_exhausted"
+        except (anthropic.APIError, *HTTP_ERRORS) as exc:
+            now = time.monotonic()
+            if retry_deadline is None:
+                retry_deadline = now + STREAM_RETRY_WINDOW_SECONDS
+            delay = STREAM_RETRY_DELAYS[retries] if retries < len(STREAM_RETRY_DELAYS) else 0
+            retryable = (
+                stream_entered
+                and isinstance(exc, RETRYABLE_STREAM_ERRORS)
+                and retries < len(STREAM_RETRY_DELAYS)
+                and now + delay <= retry_deadline
+            )
+            will_retry = retryable and (deadline is None or now + delay < deadline)
+            emitter.emit(
+                "model.failed",
+                {
+                    "model_call_id": model_call_id,
+                    "error_type": type(exc).__name__,
+                    "message": str(exc),
+                    "generation_id": generation_id,
+                    "duration_ms": (time.perf_counter_ns() - model_started_ns) / 1_000_000,
+                    "will_retry": will_retry,
+                    "retry_delay_seconds": delay if will_retry else 0,
+                    "source_timestamp": utc_now(),
+                },
+            )
+            if not will_retry:
+                if deadline is not None and (now >= deadline or (retryable and now + delay >= deadline)):
+                    return "time_budget_exhausted"
+                raise
+            # The retry delay fits inside both the recovery window and budget.
+            time.sleep(delay)
+            retries += 1
+            continue
+
+        retries = 0
+        retry_deadline = None
 
         content = _native_content_blocks(message.content)
+        input_errors: dict[str, str] = {}
+        invalid_json_ids: set[str] = set()
+        for index, block in enumerate(message.content):
+            if block.type != "tool_use":
+                continue
+            error = tool_input_error(block.name, block.input)
+            if index in input_json:
+                raw_json = "".join(input_json[index])
+                try:
+                    json.loads(raw_json)
+                except json.JSONDecodeError:
+                    # The SDK parses partial JSON while streaming. Even a
+                    # complete-looking dict is not permission to execute an
+                    # unfinished call after a provider reports tool_use.
+                    error = (
+                        "Invalid tool argument JSON: incomplete or malformed. "
+                        "Resend a complete JSON object."
+                    )
+                    invalid_json_ids.add(block.id)
+                    content[index]["input_json"] = raw_json
+            if error:
+                input_errors[block.id] = error
+                content[index]["input_error"] = error
         tool_calls = [
             item
             for item in content
@@ -510,6 +768,10 @@ def _run_model_loop(
             ),
             "generation_id": generation_id,
             "cost": usage_cost(usage if isinstance(usage, dict) else None)
+            or estimated_cost(
+                str(getattr(message, "model", None) or model),
+                usage if isinstance(usage, dict) else None,
+            )
             or pending_cost(generation_id),
             "duration_ms": (model_completed_ns - model_started_ns) / 1_000_000,
             "source_timestamp": utc_now(),
@@ -517,20 +779,58 @@ def _run_model_loop(
         emitter.emit("model.completed", payload)
 
         turns += 1
-        messages.append({"role": "assistant", "content": message.content})
+        if deadline is not None and time.monotonic() >= deadline:
+            return "time_budget_exhausted"
+        if message.stop_reason == "max_tokens":
+            # Do not execute partial tool calls or replay them without results.
+            # Thinking may also be cut off before its signature arrives. Keep
+            # visible text and an explicit notice for the next interactive turn.
+            text = "".join(
+                block.text for block in message.content if block.type == "text"
+            )
+            messages.append(
+                {
+                    "role": "assistant",
+                    "content": (f"{text}\n\n" if text else "") + _TRUNCATION_NOTICE,
+                }
+            )
+            return "response_truncated"
+        request_content = []
+        for block in message.content:
+            if block.type == "tool_use" and (
+                block.id in invalid_json_ids or not isinstance(block.input, dict)
+            ):
+                block = copy(block)
+                block.input = {}
+            request_content.append(block)
+        messages.append({"role": "assistant", "content": request_content})
         if message.stop_reason != "tool_use":
-            return True
+            return "completed"
         if max_turns is not None and turns >= max_turns:
             # Stop before running the tools: their results would only be
             # useful to a reply this budget can no longer pay for.
-            return False
+            return "max_turns_exhausted"
         # Every tool_use block needs a matching tool_result in the next
         # user message, or the API rejects the request.
-        results = [
-            _run_one_tool(block, emitter, model_call_id)
-            for block in message.content
-            if block.type == "tool_use"
-        ]
+        results = []
+        for block in message.content:
+            if block.type != "tool_use":
+                continue
+            if deadline is not None and time.monotonic() >= deadline:
+                # The reply or a preceding tool consumed the remaining time.
+                # Stop the run without starting another tool or model call.
+                return "time_budget_exhausted"
+            remaining = None if deadline is None else deadline - time.monotonic()
+            try:
+                with wall_clock_limit(remaining):
+                    result = _run_one_tool(
+                        block, emitter, model_call_id,
+                        input_error=input_errors.get(block.id),
+                        remaining_seconds=remaining,
+                    )
+            except DeadlineExceeded:
+                return "time_budget_exhausted"
+            results.append(result)
         messages.append({"role": "user", "content": results})
 
 
@@ -538,13 +838,16 @@ def _reconcile_costs(
     client: anthropic.Anthropic,
     journal: EventJournal,
     emitter: EventEmitter,
+    *,
+    max_seconds: float | None = None,
 ) -> list[JsonObject]:
-    """Append resolved OpenRouter costs without affecting the run outcome."""
+    """Reconcile pending or estimated costs without affecting the run outcome."""
     base_url = getattr(client, "base_url", "")
     credential = client.api_key or client.auth_token
     if not isinstance(credential, str) or not credential:
         return []
     outcomes: list[JsonObject] = []
+    deadline = None if max_seconds is None else time.monotonic() + max_seconds
     entries = EventJournal.replay(journal.path)
     already_resolved = {
         str(entry.payload["generation_id"])
@@ -552,24 +855,33 @@ def _reconcile_costs(
         if entry.type == "model.cost_resolved"
     }
     for entry in entries:
-        if entry.type != "model.completed":
+        if entry.type not in {"model.completed", "model.failed"}:
             continue
         generation_id = entry.payload.get("generation_id")
-        cost = entry.payload.get("cost")
+        cost = (
+            pending_cost(generation_id)
+            if entry.type == "model.failed"
+            else entry.payload.get("cost")
+        )
         if (
             not isinstance(generation_id, str)
             or generation_id in already_resolved
             or not isinstance(cost, dict)
-            or cost.get("status") != "pending"
+            or (cost.get("status") != "pending" and cost.get("kind") != "estimated")
         ):
             continue
         diagnostics: list[JsonObject] = []
-        resolved = resolve_generation_cost(
-            base_url,
-            generation_id,
-            credential,
-            diagnostics=diagnostics,
-        )
+        if deadline is not None and time.monotonic() >= deadline:
+            break
+        try:
+            with wall_clock_limit(None if deadline is None else deadline - time.monotonic()):
+                resolved = resolve_generation_cost(
+                    base_url, generation_id, credential, diagnostics=diagnostics,
+                )
+        except DeadlineExceeded:
+            outcomes.append({"generation_id": generation_id, "status": "unresolved",
+                             "attempts": diagnostics, "reason": "finalization_deadline"})
+            break
         if resolved is not None:
             resolved["source_timestamp"] = utc_now()
             emitter.emit("model.cost_resolved", resolved)
@@ -583,20 +895,22 @@ def _reconcile_costs(
     return outcomes
 
 
-def run() -> int:
+def run(*, max_tokens: int | None = None) -> int:
     """Start the read → ask → answer loop until the user types ``/exit``.
 
     A reply may include tool calls; they are executed and their results
     fed back to the model until it finishes the turn without tool use.
     Returns the process exit code.
     """
+    max_tokens = resolve_max_tokens(max_tokens)
     client = _create_client()
     if client is None:
         return 1
 
     model = _resolve_model()
     print(
-        f"nanoPyCodeAgent v{_package_version()} — model {model} "
+        f"nanoPyCodeAgent v{_package_version()} — model {model}, "
+        f"max tokens {max_tokens} "
         "(set ANTHROPIC_MODEL to override)."
     )
     print("Type a message to chat, or /exit to quit.")
@@ -618,7 +932,7 @@ def run() -> int:
             break
 
         messages.append({"role": "user", "content": user_input})
-        _run_exchange(client, model, messages, SYSTEM_PROMPT)
+        _run_exchange(client, model, messages, SYSTEM_PROMPT, max_tokens=max_tokens)
 
     print("Bye!")
     return 0
@@ -628,6 +942,8 @@ def run_headless(
     task: str,
     *,
     max_turns: int = DEFAULT_MAX_TURNS,
+    max_tokens: int | None = None,
+    time_budget_seconds: int | None = None,
     trajectory_path: Path | None = None,
 ) -> int:
     """Work ``task`` to completion without a user, and return the exit code.
@@ -640,6 +956,7 @@ def run_headless(
     scores the result. Only a run that could not happen at all — no
     credentials, an API that keeps refusing — exits non-zero.
     """
+    max_tokens = resolve_max_tokens(max_tokens)
     client = _create_client()
     if client is None:
         return 1
@@ -649,32 +966,41 @@ def run_headless(
     # model's prose and the echoed tool calls, nothing else.
     print(
         f"nanoPyCodeAgent v{_package_version()} — model {model}, "
-        f"max turns {max_turns}",
+        f"max turns {max_turns}, max tokens {max_tokens}, "
+        f"time budget {time_budget_seconds if time_budget_seconds else 'none'}",
         file=sys.stderr,
     )
 
     messages: list[MessageParam] = [{"role": "user", "content": task}]
     try:
-        finished = _run_exchange(
+        outcome = _run_exchange(
             client,
             model,
             messages,
             HEADLESS_SYSTEM_PROMPT,
             max_turns=max_turns,
+            max_tokens=max_tokens,
+            time_budget_seconds=time_budget_seconds,
             reply_prefix="",
             trajectory_path=trajectory_path,
         )
-    except (anthropic.APIError, httpx.HTTPError) as exc:
+    except (anthropic.APIError, *HTTP_ERRORS) as exc:
         # Printed verbatim on purpose: a harness classifies a failed run by
         # pattern-matching this text (rate limit, overloaded, context length,
         # …) to decide whether retrying is worth anything. Rewording it, or
         # swallowing it, throws that away.
         print(f"API error: {exc}", file=sys.stderr)
         return 1
-    if not finished:
+    if outcome == "max_turns_exhausted":
         turns = "turn" if max_turns == 1 else "turns"
         print(
             f"[stopped after {max_turns} {turns} without finishing the task]",
             file=sys.stderr,
         )
+    elif outcome == "time_budget_exhausted":
+        print(
+            f"[stopped after the {time_budget_seconds}s time budget without "
+            "finishing the task]",
+            file=sys.stderr,
+        )
     return 0