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..b1bfe9b9 --- /dev/null +++ b/src/ai/harness/__init__.py @@ -0,0 +1,41 @@ +"""A coding agent harness on top of ``ai.Agent``. See :mod:`.agent`.""" + +from . import events, prompts, tools +from .agent import ( + INTERRUPTED, + STEER, + Agent, + Budget, + Steer, + Subagent, + UnknownCostError, + compact, + subagent, + to_usd, +) +from .events import Compacted, HarnessEvent, LimitReached, Steered +from .prompts import PROMPT, SUMMARY_PREFIX, SUMMARY_PROMPT, system_message + +__all__ = [ + "INTERRUPTED", + "PROMPT", + "STEER", + "SUMMARY_PREFIX", + "SUMMARY_PROMPT", + "Agent", + "Budget", + "Compacted", + "HarnessEvent", + "LimitReached", + "Steer", + "Steered", + "Subagent", + "UnknownCostError", + "compact", + "events", + "prompts", + "subagent", + "system_message", + "to_usd", + "tools", +] diff --git a/src/ai/harness/__main__.py b/src/ai/harness/__main__.py new file mode 100644 index 00000000..17f42dd5 --- /dev/null +++ b/src/ai/harness/__main__.py @@ -0,0 +1,180 @@ +"""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 + 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 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 + +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_ +from . import events as harness_events +from . import prompts + + +async def _run( + agent: agent_.Agent, + 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 then adds the partial + turn, with a result for every tool call. + """ + async with 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 harness_events.Steered(): + print("\n[steered]", flush=True) + case harness_events.Compacted(): + print(f"\n[compacted {event.removed} messages]") + case harness_events.LimitReached(): + 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(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 + 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 = [prompts.system_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(agent, 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(): + # 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 and budget errors; the history + # has every finished turn. The loop's task groups may wrap them. + try: + run.result() + except* errors.AIError 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( + agent: agent_.Agent, model: models.Model, prompt: str +) -> None: + messages = [prompts.system_message(), builders.user_message(prompt)] + await _run(agent, model, messages, []) + + +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(agent, model, args.prompt)) + else: + asyncio.run(_repl(agent, 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..fe744916 --- /dev/null +++ b/src/ai/harness/agent.py @@ -0,0 +1,408 @@ +"""``ai.harness.Agent``: a coding agent loop for long-running processes. + +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. + +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`. + +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 ``stream.messages``. Each tool call that did not finish +gets an ``INTERRUPTED`` error result. +""" + +from __future__ import annotations + +import asyncio +import contextvars +import dataclasses +from collections.abc import AsyncGenerator, Sequence +from typing import TYPE_CHECKING, ClassVar, Literal + +import pydantic + +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 events as harness_events +from . import prompts, tools + +if TYPE_CHECKING: + import modelsdotdev + +STEER = "steer" + +# result of a tool call that was cut short by an interrupt +INTERRUPTED = "Interrupted by user" + +# 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( + "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. + + Empty when the interface has nothing to add. + """ + + messages: list[str] = [] + + +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 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 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 + 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 + + +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] + 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 + + # 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[harness_events.HarnessEvent]: + _model.set(context.model) + + 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 + ) + 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" + ) + + # 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 + # 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 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 = ( + "turns" + if self.MAX_TURNS is not None and turns >= self.MAX_TURNS + else budget.exceeded() + ) + if limit is not None: + yield harness_events.LimitReached( + limit=limit, usage=budget.usage, usd=budget.usd + ) + 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, summary, cost = await compact( + context, keep_recent_tokens=self.KEEP_RECENT_TOKENS + ) + if summary is not None: + context.messages[:] = compacted + budget.spend(cost, prices=prices, model_id=model_id) + last = None + yield harness_events.Compacted( + summary=summary, + 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, + ): + 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 + 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: + _add_partial_turn( + context, stream.message, tr.get_tool_message(), ended + ) + raise + context.add(stream.message) + context.add(tr.get_tool_message()) + + turns += 1 + last = stream.usage + + +async def compact( + context: agents.Context, *, keep_recent_tokens: int = 20_000 +) -> 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. + """ + 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 + # 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, 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 + ), + ) + response = await models.experimental_generate( + context.model, + [ + *system, + *history[:cut], + builders.user_message(prompts.SUMMARY_PROMPT), + ], + tools=context.tools, + params=params, + ) + + summary = builders.user_message( + f"{prompts.SUMMARY_PREFIX}{response.text}\n" + ) + return [*system, summary, *history[cut:]], summary, response.usage + + +class Subagent(Agent): + """The agent behind the ``subagent`` tool. + + 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. + """ + + TOOLS: ClassVar[list[agents.AgentTool]] = [tools.read, tools.bash] + STEERING = False + + +@agents.tool +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. + """ + 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, subagent] + +__all__ = [ + "INTERRUPTED", + "STEER", + "Agent", + "Budget", + "Steer", + "Subagent", + "UnknownCostError", + "compact", + "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/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..8121f888 --- /dev/null +++ b/tests/harness/test_agent.py @@ -0,0 +1,558 @@ +from __future__ import annotations + +import asyncio +import pathlib +import types +from collections.abc import AsyncGenerator, Callable, Sequence +from typing import Any, ClassVar + +import modelsdotdev +import pytest + +import ai +from ai import _modelsdev, harness +from ai.types import events as ai_events +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, + [harness.system_message(), 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_loop_adds_no_system_message() -> None: + model = ai.testing.FakeModel( + [ai.user_message("hi"), ai.assistant_message("hello")] + ) + 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: + 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_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( + [ + 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 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_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"), + 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 + + +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] == ["user"] + assert not model.calls + + +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) + ] + 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.to_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.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.to_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.subagent, 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.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")], + ) + + 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 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 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")