From dac929b9760150b0743c0dcd0aa01f89e1da37be Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Wed, 4 Nov 2026 15:51:54 -0800 Subject: [PATCH 1/4] Add ai.harness A coding agent on top of ai.Agent, for long-running processes. The loop is the standard loop plus: - a built system prompt: env facts, AGENTS.md files, a skill list - steering: before each model request, and before stopping, the loop awaits the `steer` hook; the interface answers with its queued messages (possibly none) - turn and token limits - compaction by a client-side summary Tools: read, bash, and task (runs a subagent that cannot nest or steer). `python -m ai.harness` is a stdlib REPL to try it. --- pyproject.toml | 3 + src/ai/__init__.py | 2 + src/ai/harness/__init__.py | 35 ++++ src/ai/harness/__main__.py | 155 +++++++++++++++ src/ai/harness/agent.py | 367 ++++++++++++++++++++++++++++++++++++ src/ai/harness/tools.py | 90 +++++++++ tests/harness/__init__.py | 0 tests/harness/test_agent.py | 248 ++++++++++++++++++++++++ tests/harness/test_tools.py | 84 +++++++++ 9 files changed, 984 insertions(+) create mode 100644 src/ai/harness/__init__.py create mode 100644 src/ai/harness/__main__.py create mode 100644 src/ai/harness/agent.py create mode 100644 src/ai/harness/tools.py create mode 100644 tests/harness/__init__.py create mode 100644 tests/harness/test_agent.py create mode 100644 tests/harness/test_tools.py diff --git a/pyproject.toml b/pyproject.toml index 6bfd5968..59451ad5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -183,6 +183,9 @@ extend-ignore = [ "RUF029", # example async tools often demonstrate async entry points "T201", # examples are command-line scripts and demos ] +"src/ai/harness/__main__.py" = [ + "T201", # the harness REPL prints to the terminal +] "tests/**/*.py" = [ "D", # docstrings not required in tests "FBT", # tests intentionally cover boolean-typed tool parameters diff --git a/src/ai/__init__.py b/src/ai/__init__.py index b10faf3e..7e432207 100644 --- a/src/ai/__init__.py +++ b/src/ai/__init__.py @@ -2,6 +2,7 @@ durable, errors, experimental_telemetry, + harness, models, ops, providers, @@ -205,6 +206,7 @@ "get_hook_registry", "get_model", "get_provider", + "harness", "hook", "mcp", "message", diff --git a/src/ai/harness/__init__.py b/src/ai/harness/__init__.py new file mode 100644 index 00000000..7075ff67 --- /dev/null +++ b/src/ai/harness/__init__.py @@ -0,0 +1,35 @@ +"""A coding agent harness on top of ``ai.Agent``. See :mod:`.agent`.""" + +from . import tools +from .agent import ( + PROMPT, + STEER, + SUMMARY_PREFIX, + SUMMARY_PROMPT, + Agent, + Compacted, + HarnessEvent, + LimitReached, + Steer, + Steered, + Subagent, + compact, + task, +) + +__all__ = [ + "PROMPT", + "STEER", + "SUMMARY_PREFIX", + "SUMMARY_PROMPT", + "Agent", + "Compacted", + "HarnessEvent", + "LimitReached", + "Steer", + "Steered", + "Subagent", + "compact", + "task", + "tools", +] diff --git a/src/ai/harness/__main__.py b/src/ai/harness/__main__.py new file mode 100644 index 00000000..756cb3ba --- /dev/null +++ b/src/ai/harness/__main__.py @@ -0,0 +1,155 @@ +"""Test interface for the harness, standard library only. + + python -m ai.harness [-m MODEL] # REPL + python -m ai.harness [-m MODEL] -p PROMPT # headless, one run + +In the REPL, a line typed during a run is queued as steering: the run +takes it before its next model request. Ctrl+C stops the run; Ctrl+C or +Ctrl+D at the prompt exits. +""" + +from __future__ import annotations + +import argparse +import asyncio +import contextlib +import signal +import sys +import threading + +from .. import agents, errors, models +from ..types import builders, events +from ..types import messages as messages_ +from . import agent as agent_ + + +async def _run( + model: models.Model, + history: list[messages_.Message], + steering: list[str], +) -> None: + """Run the agent on *history* and print its events. + + Every steering ask is answered with the lines in *steering*, which + the caller may append to while the run goes. *history* is updated in + place, also when the run is stopped: the loop adds only whole turns. + """ + async with agent_.Agent().run(model, history) as stream: + try: + async for event in stream: + match event: + case events.PendingHookEvent() if ( + event.hook.hook_id == agent_.STEER + ): + agents.resolve_hook( + event.hook, agent_.Steer(messages=steering.copy()) + ) + steering.clear() + case events.TextDelta(): + print(event.chunk, end="", flush=True) + case events.ToolEnd(): + tc = event.tool_call + print(f"\n> {tc.tool_name} {tc.tool_args}", flush=True) + case events.ToolCallResult(): + for r in event.results: + text = str(r.result).strip().splitlines() + head = text[0][:100] if text else "" + more = f" (+{len(text) - 1} lines)" if text else "" + print(f"< {head}{more}", flush=True) + case agent_.Steered(): + print("\n[steered]", flush=True) + case agent_.Compacted(): + print(f"\n[compacted {event.removed} messages]") + case agent_.LimitReached(): + print(f"\n[limit reached: {event.limit}]") + print() + finally: + history[:] = stream.messages + + +async def _repl(model: models.Model) -> None: + loop = asyncio.get_running_loop() + + # stdin is read on a thread, so lines arrive while a run is going + lines: asyncio.Queue[str | None] = asyncio.Queue() + + def read_stdin() -> None: + for line in sys.stdin: + loop.call_soon_threadsafe(lines.put_nowait, line.rstrip("\n")) + loop.call_soon_threadsafe(lines.put_nowait, None) # EOF + + threading.Thread(target=read_stdin, daemon=True).start() + + history: list[messages_.Message] = [] + steering: list[str] = [] + eof = False + while True: + # lines typed after the run's last ask start the next run as is + if not steering: + if eof: + return + print("> ", end="", flush=True) + line = await lines.get() + if line is None: + return + if not line.strip(): + continue + history.append(builders.user_message(line)) + + # ── one run: Ctrl+C cancels it, typed lines steer it ── + run = asyncio.create_task(_run(model, history, steering)) + loop.add_signal_handler(signal.SIGINT, run.cancel) + try: + while not run.done(): + if eof: # stdin is closed; let the run finish + await asyncio.wait({run}) + break + get = asyncio.create_task(lines.get()) + await asyncio.wait({run, get}, return_when="FIRST_COMPLETED") + if not get.done(): + get.cancel() + elif (line := get.result()) is None: + eof = True + elif line.strip(): + steering.append(line) + finally: + loop.remove_signal_handler(signal.SIGINT) + + if run.cancelled(): + print("\n[stopped]") + steering.clear() + continue + # keep the session on provider errors; the history has every + # finished turn. The loop's task groups wrap them in a group. + try: + run.result() + except* errors.ProviderError as group: + todo: list[BaseException] = [group] + while todo: # the groups nest; print the leaves + exc = todo.pop() + if isinstance(exc, BaseExceptionGroup): + todo.extend(exc.exceptions) + else: + print(f"\n[error] {type(exc).__name__}: {exc}") + + +async def _headless(model: models.Model, prompt: str) -> None: + await _run(model, [builders.user_message(prompt)], []) + + +def main() -> None: + parser = argparse.ArgumentParser(prog="python -m ai.harness") + parser.add_argument("-m", "--model", help="model id, e.g. openai:gpt-5") + parser.add_argument("-p", "--prompt", help="run once, headless") + args = parser.parse_args() + model = models.get_model(args.model) + + with contextlib.suppress(KeyboardInterrupt): + if args.prompt is not None: + asyncio.run(_headless(model, args.prompt)) + else: + asyncio.run(_repl(model)) + + +if __name__ == "__main__": + main() diff --git a/src/ai/harness/agent.py b/src/ai/harness/agent.py new file mode 100644 index 00000000..f49c1c08 --- /dev/null +++ b/src/ai/harness/agent.py @@ -0,0 +1,367 @@ +"""``ai.harness.Agent``: a coding agent loop for long-running processes. + +The loop is the standard ``ai.Agent`` loop plus, in order: a built system +prompt (env facts, ``AGENTS.md`` files, a skill list), steering through +the ``steer`` hook, turn/token limits, and compaction by a client-side +summary. The ``task`` tool runs a subagent. Settings are class variables. + +Steering: before each model request the loop asks the interface for new +messages with ``ai.hook("steer", payload=Steer)``. The interface keeps +its own buffer and answers every ask, with an empty list if it has +nothing:: + + if isinstance(event, ai.events.PendingHookEvent) and ( + event.hook.hook_id == ai.harness.STEER + ): + ai.resolve_hook(event.hook, ai.harness.Steer(messages=queued)) + queued.clear() + +The run waits until the ask is answered. +""" + +from __future__ import annotations + +import contextvars +import dataclasses +import datetime +import pathlib +import platform +from collections.abc import AsyncGenerator +from typing import ClassVar, Literal + +import pydantic + +from .. import _modelsdev, agents, models, util +from .. import experimental_telemetry as telemetry +from ..types import builders, events +from ..types import messages as messages_ +from ..types import usage as usage_ +from . import tools + +STEER = "steer" + +PROMPT = """\ +You are a coding agent that works in the user's local environment. +Use your tools to inspect files and run commands; do not guess. +- read: read files and images. +- bash: run shell commands (search, edit, write, git, curl, ...). +- task: hand a self-contained job to a subagent with a fresh context. +Be brief. Show file paths clearly.""" + +SUMMARY_PROMPT = """\ +The messages above are a conversation to summarize. Create a structured \ +context checkpoint summary that another LLM will use to continue the work. + +Use this EXACT format: + +## Goal +[What is the user trying to accomplish?] + +## Constraints & Preferences +- [Constraints, preferences, or requirements from the user, or "(none)"] + +## Progress +### Done +- [x] [Completed tasks/changes] + +### In Progress +- [ ] [Current work] + +### Blocked +- [Issues preventing progress, if any] + +## Key Decisions +- **[Decision]**: [Brief rationale] + +## Next Steps +1. [Ordered list of what should happen next] + +## Critical Context +- [Data, examples, or references needed to continue, or "(none)"] + +Keep each section concise. Preserve exact file paths, function names, \ +and error messages.""" + +# pi's prefix: the summary goes into history as a `user` message +SUMMARY_PREFIX = ( + "The conversation history before this point was compacted into the " + "following summary:\n\n\n" +) + +# The model of the current run, for the `task` tool. `loop` sets it. +# `decouple` runs the loop in its own task, so the value does not leak +# out of the run, and tool tasks created by the loop inherit it. +_model: contextvars.ContextVar[models.Model] = contextvars.ContextVar( + "harness_model" +) + + +class Steer(pydantic.BaseModel): + """Answer to the ``steer`` hook: user messages for the next request. + + Empty when the interface has nothing to add. + """ + + messages: list[str] = [] + + +class Steered(events.BaseEvent): + """Steering messages were added to the history.""" + + messages: tuple[messages_.Message, ...] + kind: Literal["harness_steered"] = "harness_steered" + + +class Compacted(events.BaseEvent): + """Old history was replaced by a summary.""" + + summary: messages_.Message + removed: int + usage: usage_.Usage | None + kind: Literal["harness_compacted"] = "harness_compacted" + + +class LimitReached(events.BaseEvent): + """The run stopped at ``MAX_TURNS`` or ``MAX_TOKENS``.""" + + limit: Literal["turns", "tokens"] + usage: usage_.Usage + kind: Literal["harness_limit_reached"] = "harness_limit_reached" + + +type HarnessEvent = events.AgentEvent | Steered | Compacted | LimitReached + + +class Agent(agents.Agent): + TOOLS: ClassVar[list[agents.AgentTool]] = [tools.read, tools.bash] # + task + STEERING: ClassVar[bool] = True + MAX_TURNS: ClassVar[int | None] = None + MAX_TOKENS: ClassVar[int | None] = None + RESERVE_TOKENS: ClassVar[int] = 16_384 + KEEP_RECENT_TOKENS: ClassVar[int] = 20_000 + SKILL_DIRS: ClassVar[tuple[str, ...]] = ( + ".agents/skills", + "~/.agents/skills", + ) + + # The loop yields harness events too; they join `AgentEvent` later. + async def loop( # type: ignore[override] # ty: ignore[invalid-method-override] + self, context: agents.Context + ) -> AsyncGenerator[HarnessEvent]: + _model.set(context.model) + + # ── system prompt: built only if the caller gave none ── + if not context.messages or context.messages[0].role != "system": + cwd = pathlib.Path.cwd() + sections = [ + PROMPT, + f"cwd: {cwd}\nos: {platform.system()} {platform.release()}\n" + f"date: {datetime.date.today().isoformat()}", + ] + # AGENTS.md from the filesystem root down to cwd, so the + # most specific file comes last and wins + for d in reversed([cwd, *cwd.parents]): + if (f := d / "AGENTS.md").is_file(): + sections.append( + f"Instructions from {f}:\n\n{f.read_text()}" + ) + # skills: name + description from SKILL.md frontmatter; the + # model reads the file itself with `read` + skills = [] + for root in self.SKILL_DIRS: + for f in sorted( + pathlib.Path(root).expanduser().glob("*/SKILL.md") + ): + meta = {"name": f.parent.name, "description": ""} + head = f.read_text().split("---") + if len(head) >= 3: # "", frontmatter, body... + for line in head[1].splitlines(): + key, sep, value = line.partition(":") + if sep and key.strip() in meta: + meta[key.strip()] = value.strip() + skills.append( + f" \n" + f" {meta['name']}\n" + f" {meta['description']}" + "\n" + f" {f.resolve()}\n" + f" " + ) + if skills: + sections.append( + "Skills give task-specific instructions. When a task " + "matches a skill, read its SKILL.md first.\n" + "\n" + + "\n".join(skills) + + "\n" + ) + context.messages.insert( + 0, builders.system_message("\n\n".join(sections)) + ) + + info = _modelsdev.get_model_by_id( + f"{context.model.provider.name}:{context.model.id}" + ) + window = ( + info.limits.context if info and info.limits.context else 128_000 + ) + # size of the last request, for the compaction trigger + last = next( + (m.usage for m in reversed(context.messages) if m.usage), None + ) + usage, turns = usage_.Usage(), 0 + + while True: + # ── steering: ask the interface before each request, and + # once more before stopping, so a late message continues + if self.STEERING: + steer = await agents.hook(STEER, payload=Steer) + if steer.messages: + new = [builders.user_message(t) for t in steer.messages] + context.add(new) + yield Steered(messages=tuple(new)) + if not context.keep_running(): + return + + # ── limits ── + if self.MAX_TURNS is not None and turns >= self.MAX_TURNS: + yield LimitReached(limit="turns", usage=usage) + return + if ( + self.MAX_TOKENS is not None + and usage.total_tokens >= self.MAX_TOKENS + ): + yield LimitReached(limit="tokens", usage=usage) + return + + # ── compaction, when the last request was near the window ── + if ( + last is not None + and last.total_tokens > window - self.RESERVE_TOKENS + ): + before = len(context.messages) + compacted, cost = await compact( + context, keep_recent_tokens=self.KEEP_RECENT_TOKENS + ) + if len(compacted) < before: + context.messages[:] = compacted + if cost is not None: + usage += cost + last = None + yield Compacted( + summary=compacted[1], + removed=before - len(compacted) + 1, + usage=cost, + ) + + # ── one turn: the standard loop body ── + async with ( + telemetry.span(telemetry.LoopTurnSpanData()), + models.stream(context=context) as stream, + agents.ToolRunner() as tr, + util.merge(stream, tr.events()) as merged, + ): + async for event in merged: + yield event + if isinstance(event, events.ToolEnd): + tr.schedule(self.resolve(event.tool_call)) + context.add(stream.message) + context.add(tr.get_tool_message()) + + turns += 1 + last = stream.usage + if last is not None: + usage += last + + +async def compact( + context: agents.Context, *, keep_recent_tokens: int = 20_000 +) -> tuple[list[messages_.Message], usage_.Usage | None]: + """Replace old history with a summary. Keeps ``messages[0]`` (system). + + Returns the new messages and the usage of the summary request, or + the old messages and None if there is nothing to cut. + """ + system, history = context.messages[:1], context.messages[1:] + + # ── cut point: walk back from the end, about `keep_recent_tokens` + # deep. Cut only before a user/assistant message, so a tool call + # and its result stay on the same side. history[0] always goes. + cut, kept = len(history), 0 + for i in range(len(history) - 1, 0, -1): + kept += len(history[i].model_dump_json()) // 4 # rough token count + if history[i].role in ("user", "assistant"): + cut = i + if kept >= keep_recent_tokens: + break + if cut == len(history): + return context.messages, None + + # ── summary request: same prefix and tools as the last request, so + # the prompt cache hits; tool choice "none" forces a text answer + params = context.params or models.InferenceRequestParams() + if context.tools: + params = dataclasses.replace( + params, + tool_calling=models.ToolCallingParams( + tool_choice=models.ToolChoiceMode.NONE + ), + ) + async with models.stream( + context.model, + [*system, *history[:cut], builders.user_message(SUMMARY_PROMPT)], + tools=context.tools, + params=params, + ) as s: + async for _ in s: + pass + + summary = builders.user_message( + f"{SUMMARY_PREFIX}{s.message.text}\n" + ) + return [*system, summary, *history[cut:]], s.usage + + +class Subagent(Agent): + """The agent behind ``task``. + + No ``task`` tool, so no nesting. No steering: the subagent shares + the parent's hook registry, so its asks would take the parent's + messages, and parallel subagents would collide on the label. + """ + + TOOLS: ClassVar[list[agents.AgentTool]] = [tools.read, tools.bash] + STEERING = False + + +@agents.tool +async def task(prompt: str) -> agents.SubAgentTool: + """Give a self-contained task to a subagent with a fresh context. + + It returns its final answer. Use it for broad searches or + independent work. The prompt must contain everything it needs. + """ + async with Subagent().run( + _model.get(), [builders.user_message(prompt)] + ) as stream: + async for event in stream: + yield event + + +Agent.TOOLS = [*Agent.TOOLS, task] + +__all__ = [ + "PROMPT", + "STEER", + "SUMMARY_PREFIX", + "SUMMARY_PROMPT", + "Agent", + "Compacted", + "HarnessEvent", + "LimitReached", + "Steer", + "Steered", + "Subagent", + "compact", + "task", +] diff --git a/src/ai/harness/tools.py b/src/ai/harness/tools.py new file mode 100644 index 00000000..87f3e0de --- /dev/null +++ b/src/ai/harness/tools.py @@ -0,0 +1,90 @@ +"""The harness tools: ``read`` and ``bash``. Both work on the local machine.""" + +from __future__ import annotations + +import asyncio +import contextlib +import os +import pathlib +import signal + +from .. import agents +from ..types import builders, media +from ..types import messages as messages_ + + +@agents.tool +async def read( + path: str, offset: int | None = None, limit: int | None = None +) -> str | messages_.ContentOutput: + """Read a file. + + offset is the 1-based first line; limit is the number of lines. + Images (png, jpg, gif, webp) are returned as images. + """ + p = pathlib.Path(path).expanduser() + data = p.read_bytes() # missing file / directory errors go to the model + + # images: return the bytes so the provider sends a real image block + if mime := media.detect_image_media_type(data[:32]): + return builders.content_output( + f"Read image file [{mime}, {len(data)} bytes]", + builders.file_part(data, media_type=mime), + ) + + lines = data.decode("utf-8", errors="replace").split("\n") + start = max((offset or 1) - 1, 0) + if start >= len(lines): + raise ValueError( + f"offset {offset} is past the end of the file " + f"({len(lines)} lines)" + ) + end = len(lines) if limit is None else start + limit + return "\n".join(lines[start:end]) + + +@agents.tool +async def bash(command: str, timeout: float = 120) -> str: + """Run a bash command in the current directory. Returns stdout+stderr. + + timeout is in seconds. + """ + # own process group, so we can kill the command and all its children + proc = await asyncio.create_subprocess_exec( + "bash", + "-c", + command, + stdin=asyncio.subprocess.DEVNULL, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + start_new_session=True, + ) + stdout = proc.stdout + assert stdout is not None + chunks: list[bytes] = [] + + async def pump() -> None: + # collect into `chunks` so a timeout keeps the partial output + while chunk := await stdout.read(65536): + chunks.append(chunk) + await proc.wait() + + timed_out = False + try: + await asyncio.wait_for(pump(), timeout) + except TimeoutError: + timed_out = True + finally: + # timeout, or the run was stopped: kill the whole group. A no-op + # if it already exited. + with contextlib.suppress(ProcessLookupError, PermissionError): + os.killpg(proc.pid, signal.SIGKILL) + + output = b"".join(chunks).decode("utf-8", errors="replace") + if timed_out: + note = f"[timed out after {timeout}s]" + elif proc.returncode: + note = f"[exit code {proc.returncode}]" + else: + note = "" + return "\n\n".join(s for s in (output, note) if s) or "[no output]" diff --git a/tests/harness/__init__.py b/tests/harness/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/harness/test_agent.py b/tests/harness/test_agent.py new file mode 100644 index 00000000..678ceb16 --- /dev/null +++ b/tests/harness/test_agent.py @@ -0,0 +1,248 @@ +from __future__ import annotations + +import asyncio +import pathlib +from collections.abc import Sequence +from typing import Any, ClassVar + +import pytest + +import ai +from ai import harness +from ai.types import usage + + +async def _run( + agent: ai.Agent, + model: ai.Model, + messages: list[ai.messages.Message], + *, + steer: Sequence[Sequence[str]] = (), +) -> tuple[list[Any], list[ai.messages.Message]]: + """Run to the end. Answer the n-th steering ask with ``steer[n]``.""" + answers = list(steer) + got = [] + async with agent.run(model, messages) as stream: + async for event in stream: + got.append(event) + if ( + isinstance(event, ai.events.PendingHookEvent) + and event.hook.hook_id == harness.STEER + ): + texts = answers.pop(0) if answers else [] + ai.resolve_hook(event.hook, harness.Steer(messages=list(texts))) + assert not answers, "the run asked for steering fewer times than expected" + return got, stream.messages + + +def _asks(got: list[Any]) -> int: + return sum( + isinstance(e, ai.events.PendingHookEvent) + and e.hook.hook_id == harness.STEER + for e in got + ) + + +@pytest.fixture(autouse=True) +def _cwd(tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch) -> None: + # keep the real AGENTS.md / skills out of the system prompt + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + + +@ai.tool +async def fast() -> str: + """Return at once.""" + return "fast done" + + +async def test_system_prompt_has_env_agents_md_and_skills( + tmp_path: pathlib.Path, +) -> None: + (tmp_path / "AGENTS.md").write_text("use uv") + skill = tmp_path / ".agents/skills/pdf/SKILL.md" + skill.parent.mkdir(parents=True) + skill.write_text("---\nname: pdf\ndescription: Read PDFs.\n---\nbody") + model = ai.testing.FakeModel( + [ai.user_message("hi"), ai.assistant_message("hello")] + ) + + _, messages = await _run(harness.Agent(), model, [ai.user_message("hi")]) + + system = model.calls[0][0] + assert system.role == "system" + assert f"cwd: {tmp_path}" in system.text + assert ( + f"Instructions from {tmp_path / 'AGENTS.md'}:\n\nuse uv" in system.text + ) + assert "pdf" in system.text + assert "Read PDFs." in system.text + assert f"{skill.resolve()}" in system.text + assert messages[-1].text == "hello" + + +async def test_caller_system_message_is_kept() -> None: + model = ai.testing.FakeModel( + [ai.user_message("hi"), ai.assistant_message("hello")] + ) + await _run( + harness.Agent(), + model, + [ai.system_message("custom"), ai.user_message("hi")], + ) + assert [m.text for m in model.calls[0]] == ["custom", "hi"] + + +async def test_asks_before_each_request_and_before_stopping() -> None: + call = ai.testing.tool_call(fast) + model = ai.testing.FakeModel( + [ + ai.user_message("go"), + ai.assistant_message(call), + ai.assistant_message("done"), + ] + ) + + class A(harness.Agent): + TOOLS: ClassVar[list[ai.AgentTool]] = [fast] + + got, _ = await _run(A(), model, [ai.user_message("go")]) + + assert _asks(got) == 3 + assert not any(isinstance(e, harness.Steered) for e in got) + assert not model.unused + + +async def test_steering_before_the_first_request_keeps_order() -> None: + model = ai.testing.FakeModel( + [ + ai.user_message("hi"), + ai.user_message("a"), + ai.user_message("b"), + ai.assistant_message("ok"), + ] + ) + + got, _ = await _run( + harness.Agent(), model, [ai.user_message("hi")], steer=[["a", "b"]] + ) + + (steered,) = [e for e in got if isinstance(e, harness.Steered)] + assert [m.text for m in steered.messages] == ["a", "b"] + assert not model.unused + + +async def test_steering_lands_after_the_turns_tools() -> None: + call = ai.testing.tool_call(fast) + model = ai.testing.FakeModel( + [ + ai.user_message("go"), + ai.assistant_message(call), + ai.user_message("also this"), + ai.assistant_message("ok"), + ] + ) + + class A(harness.Agent): + TOOLS: ClassVar[list[ai.AgentTool]] = [fast] + + got, messages = await _run( + A(), model, [ai.user_message("go")], steer=[[], ["also this"]] + ) + + result = next(e for e in got if isinstance(e, ai.events.ToolCallResult)) + steered = next(e for e in got if isinstance(e, harness.Steered)) + assert got.index(result) < got.index(steered) + (part,) = [r for m in messages for r in m.tool_results] + assert part.result == "fast done" + assert not model.unused + + +async def test_steering_after_the_final_answer_continues_the_run() -> None: + model = ai.testing.FakeModel( + [ + ai.user_message("hi"), + ai.assistant_message("hello"), + ai.user_message("again"), + ai.assistant_message("ok"), + ] + ) + + _, messages = await _run( + harness.Agent(), model, [ai.user_message("hi")], steer=[[], ["again"]] + ) + + assert messages[-1].text == "ok" + assert not model.unused + + +async def test_max_turns() -> None: + call = ai.testing.tool_call(fast) + model = ai.testing.FakeModel( + [ + ai.user_message("go"), + ai.assistant_message(call), + ai.assistant_message("x"), + ] + ) + + class A(harness.Agent): + TOOLS: ClassVar[list[ai.AgentTool]] = [fast] + MAX_TURNS = 1 + + got, _ = await _run(A(), model, [ai.user_message("go")]) + + assert isinstance(got[-1], harness.LimitReached) + assert got[-1].limit == "turns" + + +async def test_compaction_replaces_old_history_with_summary() -> None: + first = ai.assistant_message("first").model_copy( + update={"usage": usage.Usage(input_tokens=1000)} + ) + summary = ai.user_message(f"{harness.SUMMARY_PREFIX}SUM\n") + model = ai.testing.FakeModel( + [ + ai.user_message("one"), + first, + ai.user_message(harness.SUMMARY_PROMPT), + ai.assistant_message("SUM"), + ], + [summary, ai.user_message("two"), ai.assistant_message("done")], + ) + + class A(harness.Agent): + RESERVE_TOKENS = 10**9 # always over the window + KEEP_RECENT_TOKENS = 0 + + got, messages = await _run( + A(), model, [ai.user_message("one"), first, ai.user_message("two")] + ) + + compacted = next(e for e in got if isinstance(e, harness.Compacted)) + assert compacted.removed == 2 + assert [m.text for m in messages[1:]] == [summary.text, "two", "done"] + assert len(model.calls) == 2 # the summary request, then the turn + + +async def test_task_runs_a_subagent_that_does_not_steer() -> None: + call = ai.testing.tool_call(harness.task, prompt="look around") + model = ai.testing.FakeModel( + [ + ai.user_message("go"), + ai.assistant_message(call), + ai.assistant_message("done"), + ], + [ai.user_message("look around"), ai.assistant_message("found it")], + ) + + # a steering ask from the subagent would never be answered here + got, messages = await asyncio.wait_for( + _run(harness.Agent(), model, [ai.user_message("go")]), 5 + ) + + (result,) = [r for m in messages for r in m.tool_results] + assert result.model_input == "found it" + assert messages[-1].text == "done" + assert _asks(got) == 3 + assert not model.unused diff --git a/tests/harness/test_tools.py b/tests/harness/test_tools.py new file mode 100644 index 00000000..0d636591 --- /dev/null +++ b/tests/harness/test_tools.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import asyncio +import os +import pathlib +import time + +import pytest + +from ai.harness import tools +from ai.types import messages as messages_ + +PNG = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01" + b"\x08\x06\x00\x00\x00\x1f\x15\xc4\x89" +) + + +async def test_read_window(tmp_path: pathlib.Path) -> None: + f = tmp_path / "a.txt" + f.write_text("\n".join(f"line {i}" for i in range(1, 11))) + + assert await tools.read.fn(str(f), offset=3, limit=2) == "line 3\nline 4" + assert await tools.read.fn(str(f), offset=10) == "line 10" + with pytest.raises(ValueError, match="past the end"): + await tools.read.fn(str(f), offset=11) + + +async def test_read_returns_the_whole_file(tmp_path: pathlib.Path) -> None: + f = tmp_path / "big.txt" + text = "\n".join(str(i) for i in range(5000)) + f.write_text(text) + + assert await tools.read.fn(str(f)) == text + + +async def test_read_image(tmp_path: pathlib.Path) -> None: + f = tmp_path / "a.png" + f.write_bytes(PNG) + + out = await tools.read.fn(str(f)) + + assert isinstance(out, messages_.ContentOutput) + assert isinstance(out.value[1], messages_.FilePart) + assert out.value[1].media_type == "image/png" + + +async def test_bash_output_and_exit_code() -> None: + assert await tools.bash.fn("echo hi; echo err >&2") == "hi\nerr\n" + assert await tools.bash.fn("echo no; exit 3") == "no\n\n\n[exit code 3]" + assert await tools.bash.fn("exit 3") == "[exit code 3]" + assert await tools.bash.fn("true") == "[no output]" + + +async def test_bash_returns_the_whole_output() -> None: + out = await tools.bash.fn("seq 1 5000") + assert out == "".join(f"{i}\n" for i in range(1, 5001)) + + +async def test_bash_timeout_keeps_partial_output() -> None: + out = await tools.bash.fn("echo started; sleep 10", timeout=0.3) + assert out == "started\n\n\n[timed out after 0.3s]" + + +async def test_bash_cancel_kills_process_group(tmp_path: pathlib.Path) -> None: + pidfile = tmp_path / "pid" + task = asyncio.create_task( + tools.bash.fn(f"sleep 30 & echo $! > {pidfile}; wait") + ) + while not pidfile.exists() or not pidfile.read_text().strip(): + await asyncio.sleep(0.01) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + child = int(pidfile.read_text()) + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + try: + os.kill(child, 0) + except ProcessLookupError: + return + await asyncio.sleep(0.01) + pytest.fail("the background child survived cancellation") From d13b9759ee5c4b6b1d07668590c77531c28668fd Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Thu, 5 Nov 2026 10:42:21 -0800 Subject: [PATCH 2/4] Add interrupts and shared budgets to ai.harness Interrupts: the interface cancels the run's task. The loop adds the partial turn before the cancel goes on: the text so far, the finished tool results, and an "Interrupted by user" error result for every other tool call whose ToolEnd it saw. Every tool call needs a result, or the next run would replay (re-run) it. Interrupt and redirect is a new run on that history; the REPL starts it with the lines typed during the stopped run. Budgets: MAX_TOKENS and the new MAX_USD cover a run and its subagents. The top-level run puts a Budget in a contextvar; subagent runs add to it and stop at the same limits. A request counts when its stream ends, so subagents started by the turn see it. MAX_USD is best-effort, from models.dev prices, and raises UnknownCostError when the model has no prices or a request reports no usage. LimitReached gains "usd" and reports the shared totals. The REPL gets --max-usd and keeps the session on any AIError. --- src/ai/harness/__init__.py | 8 ++ src/ai/harness/__main__.py | 54 +++++--- src/ai/harness/agent.py | 208 ++++++++++++++++++++++++++--- tests/harness/test_agent.py | 257 +++++++++++++++++++++++++++++++++++- 4 files changed, 487 insertions(+), 40 deletions(-) diff --git a/src/ai/harness/__init__.py b/src/ai/harness/__init__.py index 7075ff67..db2de6d2 100644 --- a/src/ai/harness/__init__.py +++ b/src/ai/harness/__init__.py @@ -2,34 +2,42 @@ from . import tools from .agent import ( + INTERRUPTED, PROMPT, STEER, SUMMARY_PREFIX, SUMMARY_PROMPT, Agent, + Budget, Compacted, HarnessEvent, LimitReached, Steer, Steered, Subagent, + UnknownCostError, compact, task, + usd, ) __all__ = [ + "INTERRUPTED", "PROMPT", "STEER", "SUMMARY_PREFIX", "SUMMARY_PROMPT", "Agent", + "Budget", "Compacted", "HarnessEvent", "LimitReached", "Steer", "Steered", "Subagent", + "UnknownCostError", "compact", "task", "tools", + "usd", ] diff --git a/src/ai/harness/__main__.py b/src/ai/harness/__main__.py index 756cb3ba..cf532eec 100644 --- a/src/ai/harness/__main__.py +++ b/src/ai/harness/__main__.py @@ -2,10 +2,13 @@ python -m ai.harness [-m MODEL] # REPL python -m ai.harness [-m MODEL] -p PROMPT # headless, one run + python -m ai.harness --max-usd 0.50 # budget per run In the REPL, a line typed during a run is queued as steering: the run -takes it before its next model request. Ctrl+C stops the run; Ctrl+C or -Ctrl+D at the prompt exits. +takes it before its next model request. Ctrl+C stops the run at once and +keeps the partial turn; if lines were typed during the run, a new run +starts with them (interrupt and redirect). Ctrl+C or Ctrl+D at the +prompt exits. """ from __future__ import annotations @@ -24,6 +27,7 @@ async def _run( + agent: agent_.Agent, model: models.Model, history: list[messages_.Message], steering: list[str], @@ -32,9 +36,10 @@ async def _run( Every steering ask is answered with the lines in *steering*, which the caller may append to while the run goes. *history* is updated in - place, also when the run is stopped: the loop adds only whole turns. + place, also when the run is stopped: the loop then adds the partial + turn, with a result for every tool call. """ - async with agent_.Agent().run(model, history) as stream: + async with agent.run(model, history) as stream: try: async for event in stream: match event: @@ -61,13 +66,17 @@ async def _run( case agent_.Compacted(): print(f"\n[compacted {event.removed} messages]") case agent_.LimitReached(): - print(f"\n[limit reached: {event.limit}]") + print( + f"\n[limit reached: {event.limit}, " + f"{event.usage.total_tokens} tokens, " + f"${event.usd:.4f}]" + ) print() finally: history[:] = stream.messages -async def _repl(model: models.Model) -> None: +async def _repl(agent: agent_.Agent, model: models.Model) -> None: loop = asyncio.get_running_loop() # stdin is read on a thread, so lines arrive while a run is going @@ -97,7 +106,7 @@ def read_stdin() -> None: history.append(builders.user_message(line)) # ── one run: Ctrl+C cancels it, typed lines steer it ── - run = asyncio.create_task(_run(model, history, steering)) + run = asyncio.create_task(_run(agent, model, history, steering)) loop.add_signal_handler(signal.SIGINT, run.cancel) try: while not run.done(): @@ -116,14 +125,15 @@ def read_stdin() -> None: loop.remove_signal_handler(signal.SIGINT) if run.cancelled(): - print("\n[stopped]") - steering.clear() + # the history has the partial turn. Lines typed during the + # run are kept: the next run starts with them (redirect). + print("\n[stopped]" + (", redirecting" if steering else "")) continue - # keep the session on provider errors; the history has every - # finished turn. The loop's task groups wrap them in a group. + # keep the session on provider and budget errors; the history + # has every finished turn. The loop's task groups may wrap them. try: run.result() - except* errors.ProviderError as group: + except* errors.AIError as group: todo: list[BaseException] = [group] while todo: # the groups nest; print the leaves exc = todo.pop() @@ -133,22 +143,34 @@ def read_stdin() -> None: print(f"\n[error] {type(exc).__name__}: {exc}") -async def _headless(model: models.Model, prompt: str) -> None: - await _run(model, [builders.user_message(prompt)], []) +async def _headless( + agent: agent_.Agent, model: models.Model, prompt: str +) -> None: + await _run(agent, model, [builders.user_message(prompt)], []) def main() -> None: parser = argparse.ArgumentParser(prog="python -m ai.harness") parser.add_argument("-m", "--model", help="model id, e.g. openai:gpt-5") parser.add_argument("-p", "--prompt", help="run once, headless") + parser.add_argument( + "--max-usd", + type=float, + help="budget per run in USD, subagents included (best-effort)", + ) args = parser.parse_args() model = models.get_model(args.model) + class Agent(agent_.Agent): + MAX_USD = args.max_usd + + agent = Agent() + with contextlib.suppress(KeyboardInterrupt): if args.prompt is not None: - asyncio.run(_headless(model, args.prompt)) + asyncio.run(_headless(agent, model, args.prompt)) else: - asyncio.run(_repl(model)) + asyncio.run(_repl(agent, model)) if __name__ == "__main__": diff --git a/src/ai/harness/agent.py b/src/ai/harness/agent.py index f49c1c08..e93187a2 100644 --- a/src/ai/harness/agent.py +++ b/src/ai/harness/agent.py @@ -5,6 +5,11 @@ the ``steer`` hook, turn/token limits, and compaction by a client-side summary. The ``task`` tool runs a subagent. Settings are class variables. +Budgets: ``MAX_TOKENS`` and ``MAX_USD`` cover the run and its subagents +(see :class:`Budget`). ``MAX_USD`` is best-effort, from models.dev +prices; without prices for the model the run raises +:class:`UnknownCostError` before its first request. + Steering: before each model request the loop asks the interface for new messages with ``ai.hook("steer", payload=Steer)``. The interface keeps its own buffer and answers every ask, with an empty list if it has @@ -17,29 +22,42 @@ queued.clear() The run waits until the ask is answered. + +Interrupting: cancel the task that runs the agent. The loop adds the +partial turn to the history (``stream.messages``) before the cancel goes +on: the text so far, the finished tool results, and an ``INTERRUPTED`` +error result for every tool call that did not finish. To interrupt and +redirect, start a new run on that history with the new user message. """ from __future__ import annotations +import asyncio import contextvars import dataclasses import datetime import pathlib import platform -from collections.abc import AsyncGenerator -from typing import ClassVar, Literal +from collections.abc import AsyncGenerator, Sequence +from typing import TYPE_CHECKING, ClassVar, Literal import pydantic -from .. import _modelsdev, agents, models, util +from .. import _modelsdev, agents, errors, models, util from .. import experimental_telemetry as telemetry from ..types import builders, events from ..types import messages as messages_ from ..types import usage as usage_ from . import tools +if TYPE_CHECKING: + import modelsdotdev + STEER = "steer" +# result of a tool call that was cut short by an interrupt +INTERRUPTED = "Interrupted by user" + PROMPT = """\ You are a coding agent that works in the user's local environment. Use your tools to inspect files and run commands; do not guess. @@ -95,6 +113,12 @@ "harness_model" ) +# The budget of the top-level run. `loop` sets it when it is not set, so +# subagent runs, which run in tool tasks of the parent, share it. +_budget: contextvars.ContextVar[Budget] = contextvars.ContextVar( + "harness_budget" +) + class Steer(pydantic.BaseModel): """Answer to the ``steer`` hook: user messages for the next request. @@ -122,13 +146,79 @@ class Compacted(events.BaseEvent): class LimitReached(events.BaseEvent): - """The run stopped at ``MAX_TURNS`` or ``MAX_TOKENS``.""" + """The run stopped at ``MAX_TURNS``, ``MAX_TOKENS`` or ``MAX_USD``. + + ``usage`` and ``usd`` are the budget's totals: this run and the runs + that share its budget (subagents). + """ - limit: Literal["turns", "tokens"] + limit: Literal["turns", "tokens", "usd"] usage: usage_.Usage + usd: float kind: Literal["harness_limit_reached"] = "harness_limit_reached" +class UnknownCostError(errors.AIError): + """``MAX_USD`` is set, but the cost of a request cannot be known. + + models.dev has no token prices for the model, or the provider did not + report the usage of a request. + """ + + +@dataclasses.dataclass +class Budget: + """What a run and its subagents spent, and the limits they share. + + The top-level run makes it from its ``MAX_TOKENS`` and ``MAX_USD``. + Subagent runs (``task``) find it in a contextvar, add their spend to + it, and stop at the same limits; their own ``MAX_TOKENS`` and + ``MAX_USD`` are not used. A request counts when its model stream + ends, before the turn's tools run. Limits are checked before each + request, so a run can go over by its last request, and parallel + subagents by one request each. Tools start when their call ends, so a + subagent can send its first request before the parent's request is + counted. A request cut short by an interrupt is not counted: its usage is + never reported. + """ + + max_tokens: int | None = None + max_usd: float | None = None + usage: usage_.Usage = dataclasses.field(default_factory=usage_.Usage) + usd: float = 0.0 + + +def usd(prices: Sequence[modelsdotdev.Cost], usage: usage_.Usage) -> float: + """Best-effort cost of one request in USD, from models.dev prices. + + Prices are per million tokens. The tier is picked by the request's + input size. It errs high: cache writes are charged on top of + ``input_tokens`` (Anthropic does not count them there, other + providers may), and reasoning is priced as output unless the model + has a reasoning price. + """ + tier = next( + (p for p in reversed(prices) if p.min_context <= usage.input_tokens), + prices[0], + ) + cache_read = usage.cache_read_tokens or 0 + cache_write = usage.cache_write_tokens or 0 + reasoning = usage.reasoning_tokens or 0 + return ( + max(usage.input_tokens - cache_read, 0) * tier.input + + cache_read + * (tier.input if tier.cache_read is None else tier.cache_read) + + cache_write + * (tier.input if tier.cache_write is None else tier.cache_write) + + ( + usage.output_tokens * tier.output + if tier.reasoning is None + else max(usage.output_tokens - reasoning, 0) * tier.output + + reasoning * tier.reasoning + ) + ) / 1_000_000 + + type HarnessEvent = events.AgentEvent | Steered | Compacted | LimitReached @@ -136,7 +226,8 @@ class Agent(agents.Agent): TOOLS: ClassVar[list[agents.AgentTool]] = [tools.read, tools.bash] # + task STEERING: ClassVar[bool] = True MAX_TURNS: ClassVar[int | None] = None - MAX_TOKENS: ClassVar[int | None] = None + MAX_TOKENS: ClassVar[int | None] = None # shared with subagents + MAX_USD: ClassVar[float | None] = None # shared with subagents RESERVE_TOKENS: ClassVar[int] = 16_384 KEEP_RECENT_TOKENS: ClassVar[int] = 20_000 SKILL_DIRS: ClassVar[tuple[str, ...]] = ( @@ -209,7 +300,32 @@ async def loop( # type: ignore[override] # ty: ignore[invalid-method-override] last = next( (m.usage for m in reversed(context.messages) if m.usage), None ) - usage, turns = usage_.Usage(), 0 + turns = 0 + + # ── budget: the top-level run makes it, subagents share it ── + budget = _budget.get(None) + if budget is None: + budget = Budget(max_tokens=self.MAX_TOKENS, max_usd=self.MAX_USD) + _budget.set(budget) + model_id = f"{context.model.provider.name}:{context.model.id}" + prices = info.cost if info else None + if budget.max_usd is not None and not prices: + raise UnknownCostError( + f"models.dev has no token prices for {model_id}; " + "MAX_USD cannot be enforced" + ) + + def spend(request: usage_.Usage | None) -> None: + if request is None: + if budget.max_usd is not None: + raise UnknownCostError( + f"{model_id} did not report usage; " + "MAX_USD cannot be enforced" + ) + return + budget.usage += request + if prices: + budget.usd += usd(prices, request) while True: # ── steering: ask the interface before each request, and @@ -223,15 +339,21 @@ async def loop( # type: ignore[override] # ty: ignore[invalid-method-override] if not context.keep_running(): return - # ── limits ── + # ── limits: turns are per run, the budget is shared ── + limit: Literal["turns", "tokens", "usd"] | None = None if self.MAX_TURNS is not None and turns >= self.MAX_TURNS: - yield LimitReached(limit="turns", usage=usage) - return - if ( - self.MAX_TOKENS is not None - and usage.total_tokens >= self.MAX_TOKENS + limit = "turns" + elif ( + budget.max_tokens is not None + and budget.usage.total_tokens >= budget.max_tokens ): - yield LimitReached(limit="tokens", usage=usage) + limit = "tokens" + elif budget.max_usd is not None and budget.usd >= budget.max_usd: + limit = "usd" + if limit is not None: + yield LimitReached( + limit=limit, usage=budget.usage, usd=budget.usd + ) return # ── compaction, when the last request was near the window ── @@ -245,8 +367,7 @@ async def loop( # type: ignore[override] # ty: ignore[invalid-method-override] ) if len(compacted) < before: context.messages[:] = compacted - if cost is not None: - usage += cost + spend(cost) last = None yield Compacted( summary=compacted[1], @@ -261,17 +382,56 @@ async def loop( # type: ignore[override] # ty: ignore[invalid-method-override] agents.ToolRunner() as tr, util.merge(stream, tr.events()) as merged, ): - async for event in merged: - yield event - if isinstance(event, events.ToolEnd): - tr.schedule(self.resolve(event.tool_call)) + ended: set[str] = set() # tool calls with complete args + try: + async for event in merged: + if ( + isinstance(event, events.StreamEnd) + and not event.replay + ): + # count the request now, not after the turn: + # its tools (subagents) check the same budget + spend(event.usage) + yield event + if isinstance(event, events.ToolEnd): + ended.add(event.tool_call.tool_call_id) + tr.schedule(self.resolve(event.tool_call)) + except asyncio.CancelledError: + # interrupted: keep the partial turn, but every tool + # call must have a result. A call without one would + # be replayed (run again) by the next run. + parts = [ + p + for p in stream.message.parts + if not isinstance(p, messages_.ToolCallPart) + or p.tool_call_id in ended + ] + if parts: + context.add( + stream.message.model_copy(update={"parts": parts}) + ) + done = tr.get_tool_message() + results = list(done.tool_results) if done else [] + answered = {r.tool_call_id for r in results} + results += [ + builders.tool_result_part( + p.tool_call_id, + tool_name=p.tool_name, + result=INTERRUPTED, + is_error=True, + ) + for p in parts + if isinstance(p, messages_.ToolCallPart) + and p.tool_call_id not in answered + ] + if results: + context.add(builders.tool_message(*results)) + raise context.add(stream.message) context.add(tr.get_tool_message()) turns += 1 last = stream.usage - if last is not None: - usage += last async def compact( @@ -351,17 +511,21 @@ async def task(prompt: str) -> agents.SubAgentTool: Agent.TOOLS = [*Agent.TOOLS, task] __all__ = [ + "INTERRUPTED", "PROMPT", "STEER", "SUMMARY_PREFIX", "SUMMARY_PROMPT", "Agent", + "Budget", "Compacted", "HarnessEvent", "LimitReached", "Steer", "Steered", "Subagent", + "UnknownCostError", "compact", "task", + "usd", ] diff --git a/tests/harness/test_agent.py b/tests/harness/test_agent.py index 678ceb16..7b5ab391 100644 --- a/tests/harness/test_agent.py +++ b/tests/harness/test_agent.py @@ -2,13 +2,16 @@ import asyncio import pathlib -from collections.abc import Sequence +import types +from collections.abc import AsyncGenerator, Callable, Sequence from typing import Any, ClassVar +import modelsdotdev import pytest import ai -from ai import harness +from ai import _modelsdev, harness +from ai.types import events as ai_events from ai.types import usage @@ -246,3 +249,253 @@ async def test_task_runs_a_subagent_that_does_not_steer() -> None: assert messages[-1].text == "done" assert _asks(got) == 3 assert not model.unused + + +async def _interrupt( + agent: ai.Agent, + model: ai.Model, + messages: list[ai.messages.Message], + when: Callable[[Any], bool], +) -> list[ai.messages.Message]: + """Run until ``when(event)``, then cancel, as the REPL does on Ctrl+C. + + Returns the history the interface keeps (``stream.messages``). + """ + history = list(messages) + + async def run() -> None: + async with agent.run(model, history) as stream: + try: + async for event in stream: + if when(event): + task = asyncio.current_task() + assert task is not None + task.cancel() + finally: + history[:] = stream.messages + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.create_task(run()), 5) + return history + + +slow_runs: list[str] = [] + + +@ai.tool +async def slow() -> str: + """Never finish.""" + slow_runs.append("slow") + await asyncio.sleep(3600) + return "slow done" + + +async def test_interrupt_keeps_partial_turn_and_closes_tool_calls() -> None: + fast_call, slow_call = ( + ai.testing.tool_call(fast), + ai.testing.tool_call(slow), + ) + model = ai.testing.FakeModel( + [ + ai.user_message("go"), + ai.assistant_message("on it", fast_call, slow_call), + ai.assistant_message("ok"), + ] + ) + + class A(harness.Agent): + TOOLS: ClassVar[list[ai.AgentTool]] = [fast, slow] + + seen: set[str] = set() + + def when_fast_is_done_and_slow_runs(event: Any) -> bool: + if ( + isinstance(event, ai.events.PendingHookEvent) + and event.hook.hook_id == harness.STEER + ): + ai.resolve_hook(event.hook, harness.Steer()) + elif isinstance(event, ai.events.ToolEnd): + seen.add(f"end {event.tool_call.tool_name}") + elif isinstance(event, ai.events.ToolCallResult): + seen.add("fast result") + return seen >= {"end slow", "fast result"} + + history = await _interrupt( + A(), model, [ai.user_message("go")], when_fast_is_done_and_slow_runs + ) + + assistant, tool = history[-2:] + assert assistant.text == "on it" + assert [p.tool_name for p in assistant.tool_calls] == ["fast", "slow"] + results = {r.tool_name: r for r in tool.tool_results} + assert results["fast"].result == "fast done" + assert results["slow"].result == harness.INTERRUPTED + assert results["slow"].is_error + + # the next run goes on from there; it does not run the tools again + slow_runs.clear() + _, messages = await _run(A(), model, history) + assert messages[-1].text == "ok" + assert not slow_runs + assert not model.unused + + +async def test_interrupt_while_waiting_for_steering_keeps_history() -> None: + model = ai.testing.FakeModel( + [ai.user_message("go"), ai.assistant_message("x")] + ) + + history = await _interrupt( + harness.Agent(), + model, + [ai.user_message("go")], + lambda e: isinstance(e, ai.events.PendingHookEvent), + ) + + assert [m.role for m in history] == ["system", "user"] + assert not model.calls + + +def test_usd_prices_cache_tiers_and_reasoning() -> None: + sonnet = [ + modelsdotdev.Cost(input=3, output=15, cache_read=0.3, cache_write=3.75) + ] + request = usage.Usage( + input_tokens=1000, # 400 of them cache reads + output_tokens=200, + cache_read_tokens=400, + cache_write_tokens=100, # charged on top + ) + assert harness.usd(sonnet, request) == pytest.approx( + (600 * 3 + 400 * 0.3 + 100 * 3.75 + 200 * 15) / 1e6 + ) + + tiers = [ + modelsdotdev.Cost(input=1, output=0), + modelsdotdev.Cost(input=2, output=0, min_context=200_000), + ] + assert harness.usd(tiers, usage.Usage(input_tokens=1_000_000)) == 2 + + priced = [modelsdotdev.Cost(input=0, output=10, reasoning=1)] + request = usage.Usage(output_tokens=1_000_000, reasoning_tokens=400_000) + assert harness.usd(priced, request) == pytest.approx(6 + 0.4) + + +def _priced(monkeypatch: pytest.MonkeyPatch) -> None: + """$1 per million input tokens for every model; FakeModel reports the + ``usage`` of each scripted message.""" + info = types.SimpleNamespace( + cost=[modelsdotdev.Cost(input=1, output=0)], + limits=types.SimpleNamespace(context=10**9), # no compaction + ) + monkeypatch.setattr(_modelsdev, "get_model_by_id", lambda _id: info) + + def stream( + self: Any, model: Any, messages: Any, **kwargs: Any + ) -> AsyncGenerator[ai.events.Event]: + response = self._next_response(messages) + + async def events() -> AsyncGenerator[ai.events.Event]: + async for event in ai_events._replay_message_events(response): + if isinstance(event, ai.events.StreamEnd): + event = event.model_copy(update={"usage": response.usage}) + yield event + + return events() + + monkeypatch.setattr(ai.testing._FakeProvider, "stream", stream) + + +def _costs(dollars: float, *parts: Any) -> ai.messages.Message: + """An assistant message whose request cost *dollars* (see `_priced`).""" + return ai.assistant_message(*parts).model_copy( + update={"usage": usage.Usage(input_tokens=int(dollars * 1e6))} + ) + + +def _leaves(exc: BaseException) -> list[BaseException]: + if isinstance(exc, BaseExceptionGroup): + return [leaf for e in exc.exceptions for leaf in _leaves(e)] + return [exc] + + +async def test_max_usd_without_prices_errors_before_the_first_request() -> None: + model = ai.testing.FakeModel( + [ai.user_message("go"), ai.assistant_message("x")] + ) + + class A(harness.Agent): + MAX_USD = 1.0 + + with pytest.raises(BaseException) as info: + await _run(A(), model, [ai.user_message("go")]) + + (error,) = _leaves(info.value) + assert isinstance(error, harness.UnknownCostError) + assert "fake-model" in str(error) + assert not model.calls + + +async def test_max_usd_errors_when_usage_is_not_reported( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _priced(monkeypatch) + model = ai.testing.FakeModel( + [ai.user_message("go"), ai.assistant_message("no usage")] + ) + + class A(harness.Agent): + MAX_USD = 1.0 + + with pytest.raises(BaseException) as info: + await _run(A(), model, [ai.user_message("go")]) + + (error,) = _leaves(info.value) + assert isinstance(error, harness.UnknownCostError) + + +async def test_subagent_spend_counts_against_the_parent_budget( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _priced(monkeypatch) + call = ai.testing.tool_call(harness.task, prompt="look around") + model = ai.testing.FakeModel( + [ + ai.user_message("go"), + _costs(0.5, call), + ai.assistant_message("done"), + ], + [ai.user_message("look around"), _costs(1.0, "found it")], + ) + + class A(harness.Agent): + MAX_USD = 1.5 # the parent alone spends only 0.5 + + got, _ = await _run(A(), model, [ai.user_message("go")]) + + assert isinstance(got[-1], harness.LimitReached) + assert got[-1].limit == "usd" + assert got[-1].usd == pytest.approx(1.5) + assert got[-1].usage.input_tokens == 1_500_000 + assert [m.text for m in model.unused] == ["done"] + + +async def test_subagent_stops_at_the_shared_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _priced(monkeypatch) + call = ai.testing.tool_call(harness.task, prompt="look around") + model = ai.testing.FakeModel( + [ai.user_message("go"), _costs(1.0, call), ai.assistant_message("x")], + [ai.user_message("look around"), _costs(1.0, "found it")], + ) + + class A(harness.Agent): + MAX_USD = 1.0 # spent by the parent's first turn + + got, _ = await asyncio.wait_for( + _run(A(), model, [ai.user_message("go")]), 5 + ) + + assert isinstance(got[-1], harness.LimitReached) + assert len(model.calls) == 1 # neither the subagent nor the parent asked From ce6f9f0ff608b79923602a22fe80d33d7b37806c Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Thu, 5 Nov 2026 11:04:53 -0800 Subject: [PATCH 3/4] Split ai.harness prompts and events into modules Move PROMPT, SUMMARY_PROMPT, SUMMARY_PREFIX and the system message builder to ai.harness.prompts, and Steered, Compacted, LimitReached to ai.harness.events. The system message is now built by callers with prompts.system_message(). Rename task -> subagent and usd -> to_usd. --- src/ai/harness/__init__.py | 22 +- src/ai/harness/__main__.py | 13 +- src/ai/harness/agent.py | 429 +++++++++++++----------------------- src/ai/harness/events.py | 41 ++++ src/ai/harness/prompts.py | 110 +++++++++ tests/harness/test_agent.py | 37 ++-- 6 files changed, 341 insertions(+), 311 deletions(-) create mode 100644 src/ai/harness/events.py create mode 100644 src/ai/harness/prompts.py diff --git a/src/ai/harness/__init__.py b/src/ai/harness/__init__.py index db2de6d2..b1bfe9b9 100644 --- a/src/ai/harness/__init__.py +++ b/src/ai/harness/__init__.py @@ -1,25 +1,20 @@ """A coding agent harness on top of ``ai.Agent``. See :mod:`.agent`.""" -from . import tools +from . import events, prompts, tools from .agent import ( INTERRUPTED, - PROMPT, STEER, - SUMMARY_PREFIX, - SUMMARY_PROMPT, Agent, Budget, - Compacted, - HarnessEvent, - LimitReached, Steer, - Steered, Subagent, UnknownCostError, compact, - task, - usd, + subagent, + to_usd, ) +from .events import Compacted, HarnessEvent, LimitReached, Steered +from .prompts import PROMPT, SUMMARY_PREFIX, SUMMARY_PROMPT, system_message __all__ = [ "INTERRUPTED", @@ -37,7 +32,10 @@ "Subagent", "UnknownCostError", "compact", - "task", + "events", + "prompts", + "subagent", + "system_message", + "to_usd", "tools", - "usd", ] diff --git a/src/ai/harness/__main__.py b/src/ai/harness/__main__.py index cf532eec..17f42dd5 100644 --- a/src/ai/harness/__main__.py +++ b/src/ai/harness/__main__.py @@ -24,6 +24,8 @@ from ..types import builders, events from ..types import messages as messages_ from . import agent as agent_ +from . import events as harness_events +from . import prompts async def _run( @@ -61,11 +63,11 @@ async def _run( head = text[0][:100] if text else "" more = f" (+{len(text) - 1} lines)" if text else "" print(f"< {head}{more}", flush=True) - case agent_.Steered(): + case harness_events.Steered(): print("\n[steered]", flush=True) - case agent_.Compacted(): + case harness_events.Compacted(): print(f"\n[compacted {event.removed} messages]") - case agent_.LimitReached(): + case harness_events.LimitReached(): print( f"\n[limit reached: {event.limit}, " f"{event.usage.total_tokens} tokens, " @@ -89,7 +91,7 @@ def read_stdin() -> None: threading.Thread(target=read_stdin, daemon=True).start() - history: list[messages_.Message] = [] + history = [prompts.system_message()] steering: list[str] = [] eof = False while True: @@ -146,7 +148,8 @@ def read_stdin() -> None: async def _headless( agent: agent_.Agent, model: models.Model, prompt: str ) -> None: - await _run(agent, model, [builders.user_message(prompt)], []) + messages = [prompts.system_message(), builders.user_message(prompt)] + await _run(agent, model, messages, []) def main() -> None: diff --git a/src/ai/harness/agent.py b/src/ai/harness/agent.py index e93187a2..fe744916 100644 --- a/src/ai/harness/agent.py +++ b/src/ai/harness/agent.py @@ -1,33 +1,19 @@ """``ai.harness.Agent``: a coding agent loop for long-running processes. -The loop is the standard ``ai.Agent`` loop plus, in order: a built system -prompt (env facts, ``AGENTS.md`` files, a skill list), steering through -the ``steer`` hook, turn/token limits, and compaction by a client-side -summary. The ``task`` tool runs a subagent. Settings are class variables. +The loop is the ``ai.Agent`` loop with steering, limits and compaction. +Settings are class variables. Make the system message with +:func:`.prompts.system_message` and put it first in the messages. -Budgets: ``MAX_TOKENS`` and ``MAX_USD`` cover the run and its subagents -(see :class:`Budget`). ``MAX_USD`` is best-effort, from models.dev -prices; without prices for the model the run raises -:class:`UnknownCostError` before its first request. +Steering: before each request, the loop asks the interface for user +messages with the ``steer`` hook. The interface must answer each ask. +If it has no messages, it answers with an empty :class:`Steer`. -Steering: before each model request the loop asks the interface for new -messages with ``ai.hook("steer", payload=Steer)``. The interface keeps -its own buffer and answers every ask, with an empty list if it has -nothing:: - - if isinstance(event, ai.events.PendingHookEvent) and ( - event.hook.hook_id == ai.harness.STEER - ): - ai.resolve_hook(event.hook, ai.harness.Steer(messages=queued)) - queued.clear() - -The run waits until the ask is answered. +Limits: ``MAX_TURNS`` is per run. ``MAX_TOKENS`` and ``MAX_USD`` are +shared with subagents (see :class:`Budget`). Interrupting: cancel the task that runs the agent. The loop adds the -partial turn to the history (``stream.messages``) before the cancel goes -on: the text so far, the finished tool results, and an ``INTERRUPTED`` -error result for every tool call that did not finish. To interrupt and -redirect, start a new run on that history with the new user message. +partial turn to ``stream.messages``. Each tool call that did not finish +gets an ``INTERRUPTED`` error result. """ from __future__ import annotations @@ -35,9 +21,6 @@ import asyncio import contextvars import dataclasses -import datetime -import pathlib -import platform from collections.abc import AsyncGenerator, Sequence from typing import TYPE_CHECKING, ClassVar, Literal @@ -48,7 +31,8 @@ from ..types import builders, events from ..types import messages as messages_ from ..types import usage as usage_ -from . import tools +from . import events as harness_events +from . import prompts, tools if TYPE_CHECKING: import modelsdotdev @@ -58,55 +42,7 @@ # result of a tool call that was cut short by an interrupt INTERRUPTED = "Interrupted by user" -PROMPT = """\ -You are a coding agent that works in the user's local environment. -Use your tools to inspect files and run commands; do not guess. -- read: read files and images. -- bash: run shell commands (search, edit, write, git, curl, ...). -- task: hand a self-contained job to a subagent with a fresh context. -Be brief. Show file paths clearly.""" - -SUMMARY_PROMPT = """\ -The messages above are a conversation to summarize. Create a structured \ -context checkpoint summary that another LLM will use to continue the work. - -Use this EXACT format: - -## Goal -[What is the user trying to accomplish?] - -## Constraints & Preferences -- [Constraints, preferences, or requirements from the user, or "(none)"] - -## Progress -### Done -- [x] [Completed tasks/changes] - -### In Progress -- [ ] [Current work] - -### Blocked -- [Issues preventing progress, if any] - -## Key Decisions -- **[Decision]**: [Brief rationale] - -## Next Steps -1. [Ordered list of what should happen next] - -## Critical Context -- [Data, examples, or references needed to continue, or "(none)"] - -Keep each section concise. Preserve exact file paths, function names, \ -and error messages.""" - -# pi's prefix: the summary goes into history as a `user` message -SUMMARY_PREFIX = ( - "The conversation history before this point was compacted into the " - "following summary:\n\n\n" -) - -# The model of the current run, for the `task` tool. `loop` sets it. +# The model of the current run, for the `subagent` tool. `loop` sets it. # `decouple` runs the loop in its own task, so the value does not leak # out of the run, and tool tasks created by the loop inherit it. _model: contextvars.ContextVar[models.Model] = contextvars.ContextVar( @@ -129,35 +65,6 @@ class Steer(pydantic.BaseModel): messages: list[str] = [] -class Steered(events.BaseEvent): - """Steering messages were added to the history.""" - - messages: tuple[messages_.Message, ...] - kind: Literal["harness_steered"] = "harness_steered" - - -class Compacted(events.BaseEvent): - """Old history was replaced by a summary.""" - - summary: messages_.Message - removed: int - usage: usage_.Usage | None - kind: Literal["harness_compacted"] = "harness_compacted" - - -class LimitReached(events.BaseEvent): - """The run stopped at ``MAX_TURNS``, ``MAX_TOKENS`` or ``MAX_USD``. - - ``usage`` and ``usd`` are the budget's totals: this run and the runs - that share its budget (subagents). - """ - - limit: Literal["turns", "tokens", "usd"] - usage: usage_.Usage - usd: float - kind: Literal["harness_limit_reached"] = "harness_limit_reached" - - class UnknownCostError(errors.AIError): """``MAX_USD`` is set, but the cost of a request cannot be known. @@ -171,15 +78,15 @@ class Budget: """What a run and its subagents spent, and the limits they share. The top-level run makes it from its ``MAX_TOKENS`` and ``MAX_USD``. - Subagent runs (``task``) find it in a contextvar, add their spend to - it, and stop at the same limits; their own ``MAX_TOKENS`` and - ``MAX_USD`` are not used. A request counts when its model stream - ends, before the turn's tools run. Limits are checked before each - request, so a run can go over by its last request, and parallel - subagents by one request each. Tools start when their call ends, so a - subagent can send its first request before the parent's request is - counted. A request cut short by an interrupt is not counted: its usage is - never reported. + Subagent runs find it in a contextvar, add their spend to it, and + stop at the same limits; their own ``MAX_TOKENS`` and ``MAX_USD`` are + not used. A request counts when its model stream ends, before the + turn's tools run. Limits are checked before each request, so a run + can go over by its last request, and parallel subagents by one + request each. Tools start when their call ends, so a subagent can + send its first request before the parent's request is counted. A + request cut short by an interrupt is not counted: its usage is never + reported. """ max_tokens: int | None = None @@ -187,8 +94,38 @@ class Budget: usage: usage_.Usage = dataclasses.field(default_factory=usage_.Usage) usd: float = 0.0 - -def usd(prices: Sequence[modelsdotdev.Cost], usage: usage_.Usage) -> float: + def spend( + self, + request: usage_.Usage | None, + *, + prices: Sequence[modelsdotdev.Cost] | None, + model_id: str, + ) -> None: + """Add the usage of one request. ``None``: it was not reported.""" + if request is None: + if self.max_usd is not None: + raise UnknownCostError( + f"{model_id} did not report usage; " + "MAX_USD cannot be enforced" + ) + return + self.usage += request + if prices: + self.usd += to_usd(prices, request) + + def exceeded(self) -> Literal["tokens", "usd"] | None: + """Return the limit that was reached, or None.""" + if ( + self.max_tokens is not None + and self.usage.total_tokens >= self.max_tokens + ): + return "tokens" + if self.max_usd is not None and self.usd >= self.max_usd: + return "usd" + return None + + +def to_usd(prices: Sequence[modelsdotdev.Cost], usage: usage_.Usage) -> float: """Best-effort cost of one request in USD, from models.dev prices. Prices are per million tokens. The tier is picked by the request's @@ -219,113 +156,84 @@ def usd(prices: Sequence[modelsdotdev.Cost], usage: usage_.Usage) -> float: ) / 1_000_000 -type HarnessEvent = events.AgentEvent | Steered | Compacted | LimitReached +def _shared_budget(max_tokens: int | None, max_usd: float | None) -> Budget: + """Get the shared budget. The top-level run makes and sets it.""" + budget = _budget.get(None) + if budget is None: + budget = Budget(max_tokens=max_tokens, max_usd=max_usd) + _budget.set(budget) + return budget + + +def _add_partial_turn( + context: agents.Context, + message: messages_.Message, + tool_message: messages_.Message | None, + ended: set[str], +) -> None: + """Add an interrupted turn to the history. + + Tool calls with incomplete args are dropped. Each kept call gets a + result: a call without one would be replayed (run again) by the + next run. + """ + parts = [ + p + for p in message.parts + if not isinstance(p, messages_.ToolCallPart) or p.tool_call_id in ended + ] + if parts: + context.add(message.model_copy(update={"parts": parts})) + results = list(tool_message.tool_results) if tool_message else [] + answered = {r.tool_call_id for r in results} + results += [ + builders.tool_result_part( + p.tool_call_id, + tool_name=p.tool_name, + result=INTERRUPTED, + is_error=True, + ) + for p in parts + if isinstance(p, messages_.ToolCallPart) + and p.tool_call_id not in answered + ] + if results: + context.add(builders.tool_message(*results)) class Agent(agents.Agent): - TOOLS: ClassVar[list[agents.AgentTool]] = [tools.read, tools.bash] # + task + TOOLS: ClassVar[list[agents.AgentTool]] = [tools.read, tools.bash] STEERING: ClassVar[bool] = True MAX_TURNS: ClassVar[int | None] = None MAX_TOKENS: ClassVar[int | None] = None # shared with subagents MAX_USD: ClassVar[float | None] = None # shared with subagents RESERVE_TOKENS: ClassVar[int] = 16_384 KEEP_RECENT_TOKENS: ClassVar[int] = 20_000 - SKILL_DIRS: ClassVar[tuple[str, ...]] = ( - ".agents/skills", - "~/.agents/skills", - ) # The loop yields harness events too; they join `AgentEvent` later. async def loop( # type: ignore[override] # ty: ignore[invalid-method-override] self, context: agents.Context - ) -> AsyncGenerator[HarnessEvent]: + ) -> AsyncGenerator[harness_events.HarnessEvent]: _model.set(context.model) - # ── system prompt: built only if the caller gave none ── - if not context.messages or context.messages[0].role != "system": - cwd = pathlib.Path.cwd() - sections = [ - PROMPT, - f"cwd: {cwd}\nos: {platform.system()} {platform.release()}\n" - f"date: {datetime.date.today().isoformat()}", - ] - # AGENTS.md from the filesystem root down to cwd, so the - # most specific file comes last and wins - for d in reversed([cwd, *cwd.parents]): - if (f := d / "AGENTS.md").is_file(): - sections.append( - f"Instructions from {f}:\n\n{f.read_text()}" - ) - # skills: name + description from SKILL.md frontmatter; the - # model reads the file itself with `read` - skills = [] - for root in self.SKILL_DIRS: - for f in sorted( - pathlib.Path(root).expanduser().glob("*/SKILL.md") - ): - meta = {"name": f.parent.name, "description": ""} - head = f.read_text().split("---") - if len(head) >= 3: # "", frontmatter, body... - for line in head[1].splitlines(): - key, sep, value = line.partition(":") - if sep and key.strip() in meta: - meta[key.strip()] = value.strip() - skills.append( - f" \n" - f" {meta['name']}\n" - f" {meta['description']}" - "\n" - f" {f.resolve()}\n" - f" " - ) - if skills: - sections.append( - "Skills give task-specific instructions. When a task " - "matches a skill, read its SKILL.md first.\n" - "\n" - + "\n".join(skills) - + "\n" - ) - context.messages.insert( - 0, builders.system_message("\n\n".join(sections)) - ) - - info = _modelsdev.get_model_by_id( - f"{context.model.provider.name}:{context.model.id}" - ) + model_id = f"{context.model.provider.name}:{context.model.id}" + info = _modelsdev.get_model_by_id(model_id) window = ( info.limits.context if info and info.limits.context else 128_000 ) - # size of the last request, for the compaction trigger - last = next( - (m.usage for m in reversed(context.messages) if m.usage), None - ) - turns = 0 - - # ── budget: the top-level run makes it, subagents share it ── - budget = _budget.get(None) - if budget is None: - budget = Budget(max_tokens=self.MAX_TOKENS, max_usd=self.MAX_USD) - _budget.set(budget) - model_id = f"{context.model.provider.name}:{context.model.id}" prices = info.cost if info else None + budget = _shared_budget(self.MAX_TOKENS, self.MAX_USD) if budget.max_usd is not None and not prices: raise UnknownCostError( f"models.dev has no token prices for {model_id}; " "MAX_USD cannot be enforced" ) - def spend(request: usage_.Usage | None) -> None: - if request is None: - if budget.max_usd is not None: - raise UnknownCostError( - f"{model_id} did not report usage; " - "MAX_USD cannot be enforced" - ) - return - budget.usage += request - if prices: - budget.usd += usd(prices, request) + # size of the last request, for the compaction trigger + last = next( + (m.usage for m in reversed(context.messages) if m.usage), None + ) + turns = 0 while True: # ── steering: ask the interface before each request, and @@ -335,23 +243,18 @@ def spend(request: usage_.Usage | None) -> None: if steer.messages: new = [builders.user_message(t) for t in steer.messages] context.add(new) - yield Steered(messages=tuple(new)) + yield harness_events.Steered(messages=tuple(new)) if not context.keep_running(): return # ── limits: turns are per run, the budget is shared ── - limit: Literal["turns", "tokens", "usd"] | None = None - if self.MAX_TURNS is not None and turns >= self.MAX_TURNS: - limit = "turns" - elif ( - budget.max_tokens is not None - and budget.usage.total_tokens >= budget.max_tokens - ): - limit = "tokens" - elif budget.max_usd is not None and budget.usd >= budget.max_usd: - limit = "usd" + limit: Literal["turns", "tokens", "usd"] | None = ( + "turns" + if self.MAX_TURNS is not None and turns >= self.MAX_TURNS + else budget.exceeded() + ) if limit is not None: - yield LimitReached( + yield harness_events.LimitReached( limit=limit, usage=budget.usage, usd=budget.usd ) return @@ -362,15 +265,15 @@ def spend(request: usage_.Usage | None) -> None: and last.total_tokens > window - self.RESERVE_TOKENS ): before = len(context.messages) - compacted, cost = await compact( + compacted, summary, cost = await compact( context, keep_recent_tokens=self.KEEP_RECENT_TOKENS ) - if len(compacted) < before: + if summary is not None: context.messages[:] = compacted - spend(cost) + budget.spend(cost, prices=prices, model_id=model_id) last = None - yield Compacted( - summary=compacted[1], + yield harness_events.Compacted( + summary=summary, removed=before - len(compacted) + 1, usage=cost, ) @@ -391,41 +294,17 @@ def spend(request: usage_.Usage | None) -> None: ): # count the request now, not after the turn: # its tools (subagents) check the same budget - spend(event.usage) + budget.spend( + event.usage, prices=prices, model_id=model_id + ) yield event if isinstance(event, events.ToolEnd): ended.add(event.tool_call.tool_call_id) tr.schedule(self.resolve(event.tool_call)) except asyncio.CancelledError: - # interrupted: keep the partial turn, but every tool - # call must have a result. A call without one would - # be replayed (run again) by the next run. - parts = [ - p - for p in stream.message.parts - if not isinstance(p, messages_.ToolCallPart) - or p.tool_call_id in ended - ] - if parts: - context.add( - stream.message.model_copy(update={"parts": parts}) - ) - done = tr.get_tool_message() - results = list(done.tool_results) if done else [] - answered = {r.tool_call_id for r in results} - results += [ - builders.tool_result_part( - p.tool_call_id, - tool_name=p.tool_name, - result=INTERRUPTED, - is_error=True, - ) - for p in parts - if isinstance(p, messages_.ToolCallPart) - and p.tool_call_id not in answered - ] - if results: - context.add(builders.tool_message(*results)) + _add_partial_turn( + context, stream.message, tr.get_tool_message(), ended + ) raise context.add(stream.message) context.add(tr.get_tool_message()) @@ -436,13 +315,17 @@ def spend(request: usage_.Usage | None) -> None: async def compact( context: agents.Context, *, keep_recent_tokens: int = 20_000 -) -> tuple[list[messages_.Message], usage_.Usage | None]: - """Replace old history with a summary. Keeps ``messages[0]`` (system). - - Returns the new messages and the usage of the summary request, or - the old messages and None if there is nothing to cut. +) -> tuple[ + list[messages_.Message], messages_.Message | None, usage_.Usage | None +]: + """Replace old history with a summary. Keeps a leading system message. + + Returns the new messages, the summary message, and the usage of the + summary request. If there is nothing to cut, returns the old + messages, None and None. """ - system, history = context.messages[:1], context.messages[1:] + n = 1 if context.messages and context.messages[0].role == "system" else 0 + system, history = context.messages[:n], context.messages[n:] # ── cut point: walk back from the end, about `keep_recent_tokens` # deep. Cut only before a user/assistant message, so a tool call @@ -455,7 +338,7 @@ async def compact( if kept >= keep_recent_tokens: break if cut == len(history): - return context.messages, None + return context.messages, None, None # ── summary request: same prefix and tools as the last request, so # the prompt cache hits; tool choice "none" forces a text answer @@ -467,25 +350,27 @@ async def compact( tool_choice=models.ToolChoiceMode.NONE ), ) - async with models.stream( + response = await models.experimental_generate( context.model, - [*system, *history[:cut], builders.user_message(SUMMARY_PROMPT)], + [ + *system, + *history[:cut], + builders.user_message(prompts.SUMMARY_PROMPT), + ], tools=context.tools, params=params, - ) as s: - async for _ in s: - pass + ) summary = builders.user_message( - f"{SUMMARY_PREFIX}{s.message.text}\n" + f"{prompts.SUMMARY_PREFIX}{response.text}\n" ) - return [*system, summary, *history[cut:]], s.usage + return [*system, summary, *history[cut:]], summary, response.usage class Subagent(Agent): - """The agent behind ``task``. + """The agent behind the ``subagent`` tool. - No ``task`` tool, so no nesting. No steering: the subagent shares + No ``subagent`` tool, so no nesting. No steering: the subagent shares the parent's hook registry, so its asks would take the parent's messages, and parallel subagents would collide on the label. """ @@ -495,37 +380,29 @@ class Subagent(Agent): @agents.tool -async def task(prompt: str) -> agents.SubAgentTool: +async def subagent(prompt: str) -> agents.SubAgentTool: """Give a self-contained task to a subagent with a fresh context. It returns its final answer. Use it for broad searches or independent work. The prompt must contain everything it needs. """ - async with Subagent().run( - _model.get(), [builders.user_message(prompt)] - ) as stream: + messages = [prompts.system_message(), builders.user_message(prompt)] + async with Subagent().run(_model.get(), messages) as stream: async for event in stream: yield event -Agent.TOOLS = [*Agent.TOOLS, task] +Agent.TOOLS = [*Agent.TOOLS, subagent] __all__ = [ "INTERRUPTED", - "PROMPT", "STEER", - "SUMMARY_PREFIX", - "SUMMARY_PROMPT", "Agent", "Budget", - "Compacted", - "HarnessEvent", - "LimitReached", "Steer", - "Steered", "Subagent", "UnknownCostError", "compact", - "task", - "usd", + "subagent", + "to_usd", ] diff --git a/src/ai/harness/events.py b/src/ai/harness/events.py new file mode 100644 index 00000000..01604c11 --- /dev/null +++ b/src/ai/harness/events.py @@ -0,0 +1,41 @@ +"""The events the harness loop yields on top of ``ai.events.AgentEvent``.""" + +from __future__ import annotations + +from typing import Literal + +from ..types import events +from ..types import messages as messages_ +from ..types import usage as usage_ + + +class Steered(events.BaseEvent): + """Steering messages were added to the history.""" + + messages: tuple[messages_.Message, ...] + kind: Literal["harness_steered"] = "harness_steered" + + +class Compacted(events.BaseEvent): + """Old history was replaced by a summary.""" + + summary: messages_.Message + removed: int + usage: usage_.Usage | None + kind: Literal["harness_compacted"] = "harness_compacted" + + +class LimitReached(events.BaseEvent): + """The run stopped at ``MAX_TURNS``, ``MAX_TOKENS`` or ``MAX_USD``. + + ``usage`` and ``usd`` are the budget's totals: this run and the runs + that share its budget (subagents). + """ + + limit: Literal["turns", "tokens", "usd"] + usage: usage_.Usage + usd: float + kind: Literal["harness_limit_reached"] = "harness_limit_reached" + + +type HarnessEvent = events.AgentEvent | Steered | Compacted | LimitReached diff --git a/src/ai/harness/prompts.py b/src/ai/harness/prompts.py new file mode 100644 index 00000000..0e8958c9 --- /dev/null +++ b/src/ai/harness/prompts.py @@ -0,0 +1,110 @@ +"""The harness prompts, and the system message built from them.""" + +from __future__ import annotations + +import datetime +import pathlib +import platform +from collections.abc import Sequence + +from ..types import builders +from ..types import messages as messages_ + +PROMPT = """\ +You are a coding agent that works in the user's local environment. +Use your tools to inspect files and run commands; do not guess. +- read: read files and images. +- bash: run shell commands (search, edit, write, git, curl, ...). +- subagent: hand a self-contained job to a subagent with a fresh context. +Be brief. Show file paths clearly.""" + +SUMMARY_PROMPT = """\ +The messages above are a conversation to summarize. Create a structured \ +context checkpoint summary that another LLM will use to continue the work. + +Use this EXACT format: + +## Goal +[What is the user trying to accomplish?] + +## Constraints & Preferences +- [Constraints, preferences, or requirements from the user, or "(none)"] + +## Progress +### Done +- [x] [Completed tasks/changes] + +### In Progress +- [ ] [Current work] + +### Blocked +- [Issues preventing progress, if any] + +## Key Decisions +- **[Decision]**: [Brief rationale] + +## Next Steps +1. [Ordered list of what should happen next] + +## Critical Context +- [Data, examples, or references needed to continue, or "(none)"] + +Keep each section concise. Preserve exact file paths, function names, \ +and error messages.""" + +# pi's prefix: the summary goes into history as a `user` message +SUMMARY_PREFIX = ( + "The conversation history before this point was compacted into the " + "following summary:\n\n\n" +) + +SKILL_DIRS = (".agents/skills", "~/.agents/skills") + + +def system_message( + prompt: str = PROMPT, *, skill_dirs: Sequence[str] = SKILL_DIRS +) -> messages_.Message: + """Build the system message: *prompt*, env facts, AGENTS.md, skills. + + Put it first in the messages you give to ``Agent.run``. + """ + cwd = pathlib.Path.cwd() + sections = [ + prompt, + f"cwd: {cwd}\nos: {platform.system()} {platform.release()}\n" + f"date: {datetime.date.today().isoformat()}", + ] + + # AGENTS.md from the filesystem root down to cwd, so the most + # specific file comes last and wins + for d in reversed([cwd, *cwd.parents]): + if (f := d / "AGENTS.md").is_file(): + sections.append(f"Instructions from {f}:\n\n{f.read_text()}") + + # skills: name + description from SKILL.md frontmatter; the model + # reads the file itself with `read` + skills = [] + for root in skill_dirs: + for f in sorted(pathlib.Path(root).expanduser().glob("*/SKILL.md")): + meta = {"name": f.parent.name, "description": ""} + head = f.read_text().split("---") + if len(head) >= 3: # "", frontmatter, body... + for line in head[1].splitlines(): + key, sep, value = line.partition(":") + if sep and key.strip() in meta: + meta[key.strip()] = value.strip() + skills.append( + f" \n" + f" {meta['name']}\n" + f" {meta['description']}\n" + f" {f.resolve()}\n" + f" " + ) + if skills: + sections.append( + "Skills give task-specific instructions. When a task matches " + "a skill, read its SKILL.md first.\n" + "\n" + "\n".join(skills) + "\n" + ) + + return builders.system_message("\n\n".join(sections)) diff --git a/tests/harness/test_agent.py b/tests/harness/test_agent.py index 7b5ab391..35767748 100644 --- a/tests/harness/test_agent.py +++ b/tests/harness/test_agent.py @@ -70,7 +70,11 @@ async def test_system_prompt_has_env_agents_md_and_skills( [ai.user_message("hi"), ai.assistant_message("hello")] ) - _, messages = await _run(harness.Agent(), model, [ai.user_message("hi")]) + _, messages = await _run( + harness.Agent(), + model, + [harness.system_message(), ai.user_message("hi")], + ) system = model.calls[0][0] assert system.role == "system" @@ -84,16 +88,12 @@ async def test_system_prompt_has_env_agents_md_and_skills( assert messages[-1].text == "hello" -async def test_caller_system_message_is_kept() -> None: +async def test_loop_adds_no_system_message() -> None: model = ai.testing.FakeModel( [ai.user_message("hi"), ai.assistant_message("hello")] ) - await _run( - harness.Agent(), - model, - [ai.system_message("custom"), ai.user_message("hi")], - ) - assert [m.text for m in model.calls[0]] == ["custom", "hi"] + await _run(harness.Agent(), model, [ai.user_message("hi")]) + assert [m.text for m in model.calls[0]] == ["hi"] async def test_asks_before_each_request_and_before_stopping() -> None: @@ -224,12 +224,13 @@ class A(harness.Agent): compacted = next(e for e in got if isinstance(e, harness.Compacted)) assert compacted.removed == 2 - assert [m.text for m in messages[1:]] == [summary.text, "two", "done"] + assert compacted.summary.text == summary.text + assert [m.text for m in messages] == [summary.text, "two", "done"] assert len(model.calls) == 2 # the summary request, then the turn -async def test_task_runs_a_subagent_that_does_not_steer() -> None: - call = ai.testing.tool_call(harness.task, prompt="look around") +async def test_subagent_tool_runs_a_subagent_that_does_not_steer() -> None: + call = ai.testing.tool_call(harness.subagent, prompt="look around") model = ai.testing.FakeModel( [ ai.user_message("go"), @@ -352,11 +353,11 @@ async def test_interrupt_while_waiting_for_steering_keeps_history() -> None: lambda e: isinstance(e, ai.events.PendingHookEvent), ) - assert [m.role for m in history] == ["system", "user"] + assert [m.role for m in history] == ["user"] assert not model.calls -def test_usd_prices_cache_tiers_and_reasoning() -> None: +def test_to_usd_prices_cache_tiers_and_reasoning() -> None: sonnet = [ modelsdotdev.Cost(input=3, output=15, cache_read=0.3, cache_write=3.75) ] @@ -366,7 +367,7 @@ def test_usd_prices_cache_tiers_and_reasoning() -> None: cache_read_tokens=400, cache_write_tokens=100, # charged on top ) - assert harness.usd(sonnet, request) == pytest.approx( + assert harness.to_usd(sonnet, request) == pytest.approx( (600 * 3 + 400 * 0.3 + 100 * 3.75 + 200 * 15) / 1e6 ) @@ -374,11 +375,11 @@ def test_usd_prices_cache_tiers_and_reasoning() -> None: modelsdotdev.Cost(input=1, output=0), modelsdotdev.Cost(input=2, output=0, min_context=200_000), ] - assert harness.usd(tiers, usage.Usage(input_tokens=1_000_000)) == 2 + assert harness.to_usd(tiers, usage.Usage(input_tokens=1_000_000)) == 2 priced = [modelsdotdev.Cost(input=0, output=10, reasoning=1)] request = usage.Usage(output_tokens=1_000_000, reasoning_tokens=400_000) - assert harness.usd(priced, request) == pytest.approx(6 + 0.4) + assert harness.to_usd(priced, request) == pytest.approx(6 + 0.4) def _priced(monkeypatch: pytest.MonkeyPatch) -> None: @@ -458,7 +459,7 @@ async def test_subagent_spend_counts_against_the_parent_budget( monkeypatch: pytest.MonkeyPatch, ) -> None: _priced(monkeypatch) - call = ai.testing.tool_call(harness.task, prompt="look around") + call = ai.testing.tool_call(harness.subagent, prompt="look around") model = ai.testing.FakeModel( [ ai.user_message("go"), @@ -484,7 +485,7 @@ async def test_subagent_stops_at_the_shared_limit( monkeypatch: pytest.MonkeyPatch, ) -> None: _priced(monkeypatch) - call = ai.testing.tool_call(harness.task, prompt="look around") + call = ai.testing.tool_call(harness.subagent, prompt="look around") model = ai.testing.FakeModel( [ai.user_message("go"), _costs(1.0, call), ai.assistant_message("x")], [ai.user_message("look around"), _costs(1.0, "found it")], From 9c631e13019ebc782113c1d6b47bdcbf97708ddf Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Thu, 5 Nov 2026 12:41:04 -0800 Subject: [PATCH 4/4] Test steering twice in ai.harness and its REPL --- tests/harness/test_agent.py | 56 +++++++++++++++++ tests/harness/test_main.py | 116 ++++++++++++++++++++++++++++++++++++ 2 files changed, 172 insertions(+) create mode 100644 tests/harness/test_main.py diff --git a/tests/harness/test_agent.py b/tests/harness/test_agent.py index 35767748..8121f888 100644 --- a/tests/harness/test_agent.py +++ b/tests/harness/test_agent.py @@ -179,6 +179,62 @@ async def test_steering_after_the_final_answer_continues_the_run() -> None: assert not model.unused +async def test_steering_twice_in_one_run() -> None: + model = ai.testing.FakeModel( + [ + ai.user_message("go"), + ai.assistant_message(ai.testing.tool_call(fast)), + ai.user_message("s1"), + ai.assistant_message(ai.testing.tool_call(fast)), + ai.user_message("s2"), + ai.assistant_message("ok"), + ] + ) + + class A(harness.Agent): + TOOLS: ClassVar[list[ai.AgentTool]] = [fast] + + got, messages = await _run( + A(), model, [ai.user_message("go")], steer=[[], ["s1"], ["s2"]] + ) + + steered = [e for e in got if isinstance(e, harness.Steered)] + assert [[m.text for m in e.messages] for e in steered] == [["s1"], ["s2"]] + assert messages[-1].text == "ok" + assert not model.unused + + +async def test_steering_in_each_of_two_runs() -> None: + model = ai.testing.FakeModel( + [ + ai.user_message("go"), + ai.assistant_message(ai.testing.tool_call(fast)), + ai.user_message("s1"), + ai.assistant_message("ok1"), + ai.user_message("go2"), + ai.assistant_message(ai.testing.tool_call(fast)), + ai.user_message("s2"), + ai.assistant_message("ok2"), + ] + ) + + class A(harness.Agent): + TOOLS: ClassVar[list[ai.AgentTool]] = [fast] + + got1, history = await _run( + A(), model, [ai.user_message("go")], steer=[[], ["s1"]] + ) + got2, history = await _run( + A(), model, [*history, ai.user_message("go2")], steer=[[], ["s2"]] + ) + + for got, text in [(got1, "s1"), (got2, "s2")]: + (steered,) = [e for e in got if isinstance(e, harness.Steered)] + assert [m.text for m in steered.messages] == [text] + assert history[-1].text == "ok2" + assert not model.unused + + async def test_max_turns() -> None: call = ai.testing.tool_call(fast) model = ai.testing.FakeModel( diff --git a/tests/harness/test_main.py b/tests/harness/test_main.py new file mode 100644 index 00000000..aceac1e3 --- /dev/null +++ b/tests/harness/test_main.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import asyncio +import os +import pathlib +import sys +from typing import Any, ClassVar + +import pytest + +import ai +from ai import harness +from ai.harness import __main__ as main + + +@pytest.fixture(autouse=True) +def _cwd(tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch) -> None: + # keep the real AGENTS.md / skills out of the system prompt + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path / "home")) + + +@pytest.fixture +def typist(monkeypatch: pytest.MonkeyPatch) -> ai.AgentTool: + """A tool that types lines into the REPL's stdin, like a user mid-run. + + stdin starts with the prompt ``go``. The tool returns once the REPL + has queued the lines as steering; with ``last`` it closes stdin, so + the REPL exits after the run. + """ + r, w = os.pipe() + monkeypatch.setattr(sys, "stdin", os.fdopen(r)) + os.write(w, b"go\n") + + queued: list[list[str]] = [[]] + run = main._run + + async def spy(*args: Any) -> None: + queued[0] = args[3] # the REPL's steering list + await run(*args) + + monkeypatch.setattr(main, "_run", spy) + + @ai.tool + async def type_(lines: list[str], last: bool = False) -> str: + """Type lines into the REPL.""" + os.write(w, "".join(f"{line}\n" for line in lines).encode()) + if last: + os.close(w) + async with asyncio.timeout(5): + while queued[0] != lines: + await asyncio.sleep(0.01) + return "typed" + + return type_ + + +async def _repl(model: ai.Model, tool: ai.AgentTool) -> None: + class A(harness.Agent): + TOOLS: ClassVar[list[ai.AgentTool]] = [tool] + + await asyncio.wait_for(main._repl(A(), model), 10) + + +async def test_steering_in_two_turns_prints_two_markers( + typist: ai.AgentTool, capsys: pytest.CaptureFixture[str] +) -> None: + model = ai.testing.FakeModel( + [ + ai.user_message("go"), + ai.assistant_message(ai.testing.tool_call(typist, lines=["s1"])), + ai.user_message("s1"), + ai.assistant_message( + ai.testing.tool_call(typist, lines=["s2"], last=True) + ), + ai.user_message("s2"), + ai.assistant_message("ok"), + ] + ) + + await _repl(model, typist) + + assert capsys.readouterr().out.count("[steered]") == 2 + assert [m.text for m in model.calls[-1] if m.role == "user"] == [ + "go", + "s1", + "s2", + ] + assert not model.unused + + +async def test_steering_twice_in_one_turn_is_one_batch( + typist: ai.AgentTool, capsys: pytest.CaptureFixture[str] +) -> None: + # both lines land, but in one Steered event: one marker for two lines + model = ai.testing.FakeModel( + [ + ai.user_message("go"), + ai.assistant_message( + ai.testing.tool_call(typist, lines=["s1", "s2"], last=True) + ), + ai.user_message("s1"), + ai.user_message("s2"), + ai.assistant_message("ok"), + ] + ) + + await _repl(model, typist) + + assert capsys.readouterr().out.count("[steered]") == 1 + assert [m.text for m in model.calls[-1] if m.role == "user"] == [ + "go", + "s1", + "s2", + ] + assert not model.unused