import asyncio
import importlib.util
import inspect
import itertools
import os
import platform
import re
import secrets
import time
from datetime import datetime, timedelta
from io import StringIO
from llm import Chat, DEFAULT_MODEL, available_models, model_alias, resolve_model
from utils import execute_python, format_stats, installed_packages, parse_duration
CLOCK_RESOLUTION_MARGIN = 0.02
ATTR_RE = re.compile(r'(\w+)\s*=\s*(?:"([^"]*)"|([^\s>]+))')
def parse_attrs(text):
return {name: quoted or bare for name, quoted, bare in ATTR_RE.findall(text)}
def find_tags(message, tag):
"""Yield (attrs, body) for every <tag ...>...</tag> that starts a line in message. The close only counts
when it too starts a line, so a </tag> sitting mid-line inside the body (in a string, a comment) is inert
and the block still parses. A boundary="x" on the tag additionally protects a </tag> that lands at the
start of a line inside the body. The `(?=[ >])` guards run_code against also matching run_code_result."""
for match in re.finditer(rf"^[ \t]*<{tag}(?=[ >])(?P<attrs>[^>]*)>", message, re.MULTILINE):
attrs = parse_attrs(match.group("attrs"))
boundary = attrs.get("boundary")
close = f'</{tag} boundary="{boundary}">' if boundary else f"</{tag}>"
found = re.search(rf"^[ \t]*{re.escape(close)}", message[match.end():], re.MULTILINE)
if found:
yield attrs, message[match.end():match.end() + found.start()]
def wrap(tag, attrs, body):
"""Emit <tag attrs>body</tag>. If body contains the closing tag at the start of a line (the only thing a
reader could mistake for the real close), add a boundary="x" (a short random token absent from body) to
both delimiters so the true end stays findable — like a multipart boundary. Otherwise plain tags."""
if re.search(rf"^[ \t]*</{tag}>", body, re.MULTILINE):
boundary = secrets.token_hex(3)
while boundary in body:
boundary = secrets.token_hex(3)
return f'<{tag}{attrs} boundary="{boundary}">\n{body}\n</{tag} boundary="{boundary}">\n'
return f"<{tag}{attrs}>\n{body}\n</{tag}>\n"
SYSTEM = """You are an AI agent that solves tasks by writing and running Python code and starting sub-agents.
Each reply is one or more <run_code>/<run_agent> blocks, or plain text to wait while running blocks and sub-agents keep working. Run code like this:
<run_code id="x" timeout="30s" heartbeat="5s">
print("result")
</run_code>
Only what you print comes back. Attributes: id (echoed on the result), timeout (default 10s), heartbeat (how often partial output streams while it runs).
A finished block returns <run_code_result id=... duration=...>output</run_code_result>. With a heartbeat it also streams <running id=... current=...>partial output</running> as it runs. That is only a status. Reply in plain text to let it keep going. Left alone it runs to its deadline, then is killed.
Boundaries: a closing tag ends a block only at the start of a line, so </run_code> mid-line (in a string or comment) is inert.
Rules:
- Work in small steps. Each block makes a little progress.
- Blocks share one persistent namespace and run concurrently. Variables, imports and functions persist.
- Your code runs in an asyncio loop. Use top-level await directly.
- Prefer your provided tools over reimplementing them.
- self.compact_knowledge("...") replaces your whole history with that text. Keep every fact, path, value and next step you still need. Drop the rest.
- When the task is solved and nothing of yours is running, reply in plain text. That is your final answer.
Sub-agents run in a fresh context. Start one with <run_agent name="sub" namespace="a,b" model="flash" heartbeat="30s">task</run_agent>. It returns instantly. Its answer arrives as <agent_result name=...>...</agent_result>, and heartbeat streams <agent_status name=...> while it works. namespace lists objects to pass from your namespace. model picks its model; send bulk or simple work to a fast one. Delegate any subtask whose main cost is reading a lot (pages, big files, logs): the sub-agent returns only the result and keeps your context small.
"""
ENVIRONMENT = """You are running on the '{model_alias}' model ({model}). When you spawn a sub-agent you can give it a different model with the model="alias" attribute on <run_agent> (or the model= argument of self.run_agent): pick a cheap fast model for simple or bulk work and a stronger model for hard reasoning. Omit it to reuse your own model. Available models:
{available_models}
Environment:
Runtime: Python {python} on {os}
CPU cores: {cpu}{ram}
Top-level installed packages: {packages}{namespace}"""
def build_system(namespace, model):
"""The agent's system prompt: the static format guide (SYSTEM) followed by the ENVIRONMENT template filled with
the model note and the live environment (runtime, packages, and every tool in the namespace).
AGENT_MINIMAL_SYSTEM=1 drops the SYSTEM guide, keeping only the environment — for models fine-tuned to know
the format intrinsically (see --drop-system in train_qlora.py)."""
ram = ""
if importlib.util.find_spec("psutil"):
import psutil
memory = psutil.virtual_memory()
ram = f"\nRAM: {memory.available // 1024**3} GB available of {memory.total // 1024**3} GB"
tools = "\n".join(f"- {describe_object(name, value)}" for name, value in namespace.items())
namespace_block = f"\nThese objects are already available by name in your code (no import needed):\n{tools}" if namespace else ""
intro = "" if os.environ.get("AGENT_MINIMAL_SYSTEM") == "1" else SYSTEM + "\n"
return intro + ENVIRONMENT.format(
model_alias=model_alias(model),
model=model,
available_models=available_models(),
python=platform.python_version(),
os=f"{platform.system()} {platform.release()}",
cpu=os.cpu_count(),
ram=ram,
packages=", ".join(installed_packages()),
namespace=namespace_block,
)
def describe_object(name, value):
"""One entry in the namespace listing: a callable as its signature + docstring, an object as its public
methods, or a plain value as a truncated repr."""
if callable(value):
prefix = "async " if inspect.iscoroutinefunction(value) else ""
return f"{prefix}{name}{inspect.signature(value)}: {inspect.getdoc(value) or ''}"
if type(value).__module__ != "builtins":
methods = [(attr, getattr(value, attr)) for attr in dir(value) if not attr.startswith("_") and callable(getattr(value, attr, None))]
if methods:
doc = inspect.getdoc(value)
lines = [f"{name}: {type(value).__name__}" + (f" — {doc.splitlines()[0]}" if doc else "")]
for attr, method in methods:
try:
sig = str(inspect.signature(method))
except (TypeError, ValueError):
sig = "(...)"
prefix = "async " if inspect.iscoroutinefunction(method) else ""
lines.append(f" {prefix}{name}.{attr}{sig}: {inspect.getdoc(method) or ''}")
return "\n".join(lines)
preview = repr(value)
if len(preview) > 80:
preview = preview[:80] + "…"
return f"{name}: {type(value).__name__} = {preview}"
class CodeBlock:
counter = itertools.count(1)
@classmethod
def find_all(cls, message):
blocks = []
for attrs, code in find_tags(message, "run_code"):
heartbeat = parse_duration(attrs["heartbeat"]) if "heartbeat" in attrs else None
blocks.append(cls(code, attrs.get("id"), parse_duration(attrs.get("timeout")), heartbeat))
return blocks
def __init__(self, code, id, timeout, heartbeat):
self.code = code
self.id = id or f"b{next(CodeBlock.counter)}"
self.timeout = timeout
self.heartbeat = heartbeat
self.task = None
self.buffer = None
self.start = None
self.deadline = None
self.started_at = None
self.next_tick = None
self.result_suffix = ""
def launch(self, namespace, display, name, on_heartbeat):
self.display = display
self.name = name
self.buffer = StringIO()
self.start = time.time()
self.deadline = time.time() + self.timeout
self.started_at = datetime.now()
self.task = asyncio.create_task(self.run(namespace, on_heartbeat))
def cancel(self):
if self.task and not self.task.done():
self.result_suffix = "[cancelled]"
self.task.cancel()
def is_running(self):
return not self.task.done()
async def run(self, namespace, on_heartbeat):
python_task = asyncio.create_task(execute_python(self.code, namespace, self.buffer))
try:
while not python_task.done():
timeout = self.deadline - time.time()
if self.heartbeat:
timeout = min(timeout, self.heartbeat)
await asyncio.wait([python_task], timeout=max(0, timeout) + CLOCK_RESOLUTION_MARGIN)
if python_task.done():
break
if time.time() >= self.deadline:
self.result_suffix = "[timed out]"
python_task.cancel()
await asyncio.wait([python_task])
else:
self.display.event("progress", name=self.name, id=self.id, current=time.time() - self.start, output=self.output)
on_heartbeat(self.format())
self.display.event("finished", name=self.name, id=self.id, duration=time.time() - self.start, cancelled=bool(self.result_suffix), output=self.output)
elapsed = time.time() - self.start
return wrap("run_code_result", f" id={self.id} duration={elapsed:.2f}s", self.output)
except asyncio.CancelledError:
if not python_task.done():
python_task.cancel()
await asyncio.wait([python_task])
raise
@property
def output(self):
return self.buffer.getvalue() + self.result_suffix
def format(self):
elapsed = time.time() - self.start
deadline = (self.started_at + timedelta(seconds=self.timeout)).strftime("%H:%M:%S")
clock = datetime.now().strftime("%H:%M:%S")
return wrap("running", f" id={self.id} current={elapsed:.0f}s now={clock} deadline={deadline}", self.buffer.getvalue())
class AgentDisplay:
"""Everything an Agent reports flows through one method: event(type, **fields). The base is a no-op.
Event types and their fields:
system(name, prompt) · sent(name, message, count) · replied(name, message, stats, latency)
launched(name, id, timeout, heartbeat, code) · progress(name, id, current, output)
finished(name, id, duration, cancelled, output) · cancelled(name, label) · failed(name, message)"""
def event(self, type, **fields):
pass
def accepted_kwargs(handler, fields):
"""The subset of fields that handler declares by name (or all of them when it takes **kwargs) — lets an event
carry an optional field like model that only some renderers render, without the others raising on it."""
params = inspect.signature(handler).parameters.values()
if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params):
return fields
names = {param.name for param in params}
return {key: value for key, value in fields.items() if key in names}
class DispatchDisplay(AgentDisplay):
"""For renderers: routes each event to a same-named method (def launched(self, name, id, ...)) if defined,
ignoring event types the subclass doesn't care about and event fields the handler doesn't declare."""
def event(self, type, **fields):
handler = getattr(self, type, None)
if callable(handler):
handler(**accepted_kwargs(handler, fields))
class FanoutDisplay(AgentDisplay):
"""Sends every event to several displays at once, e.g. a console renderer plus a JSONL recorder."""
def __init__(self, displays):
self.displays = displays
def event(self, type, **fields):
for display in self.displays:
display.event(type, **fields)
class Agent:
"""Conversation loop, code-block runtime, and API exposed to code the LLM writes."""
@classmethod
def find_all(cls, message):
agents = []
for attrs, task in find_tags(message, "run_agent"):
agent = cls(name=attrs.get("name", "sub"), context=attrs.get("context", ""))
agent.assignment = task.strip()
agent.namespace_names = [name.strip() for name in attrs.get("namespace", "").split(",") if name.strip()]
agent.heartbeat = attrs.get("heartbeat", "")
agent.model = attrs.get("model", "")
agents.append(agent)
return agents
def __init__(self, llm=None, namespace=None, display=None, depth=0, max_depth=3, name="main", context=""):
self.llm = llm
self.display = display or AgentDisplay()
self.depth = depth
self.max_depth = max_depth
self.name = name
self.context = context
self.intro = f"Context handed to you by your caller — use it, do not re-derive it:\n{context}\n\n" if context else ""
self.namespace = {**(namespace or {}), "asyncio": asyncio, "self": self}
self.base_names = set(self.namespace)
self.namespace_names = []
self.namespace_values = {}
self.assignment = ""
self.heartbeat = ""
self.model = ""
self.blocks = []
self.sub_agents = []
self.inbox = []
self.wake = asyncio.Event()
self.task = None
self.start = None
self.heartbeat_interval = None
self._parent_on_heartbeat = None
self.supervisor = None
self.result_value = None
def on_heartbeat(self, report):
self.inbox.append(report)
self.wake.set()
def process_llm_response(self, reply):
"""Launches everything the LLM's reply asks for: starts each <run_code> block and <run_agent> sub-agent,
and queues an error into the inbox for malformed attributes or a truncated (unclosed) tag. Does not wait."""
try:
code_blocks = CodeBlock.find_all(reply)
agents = Agent.find_all(reply)
except Exception as error:
code_blocks, agents = [], []
self.inbox.append(wrap("error", "", f"Could not parse the attributes of your <run_code>/<run_agent> tags: {error}. "
f"Fix the malformed attribute and resend; nothing was run."))
for block in code_blocks:
self.start_block(block)
for agent in agents:
try:
self.start_agent(agent)
except Exception as error:
self.inbox.append(wrap("error", f" name={agent.name}", f"Could not start this sub-agent: {error}"))
attempted = bool(re.search(r"^[ \t]*<run_(?:code|agent)(?=[ >])", reply, re.MULTILINE))
if not code_blocks and not agents and attempted:
self.inbox.append(wrap("error", "", "Your reply was cut off: it opens a <run_code>/<run_agent> tag but never closes it, "
"so nothing ran. Resend a complete, closed block."))
def next_message(self):
"""Drains everything to send back to the model this turn: queued inbox items (heartbeats, errors, steer
notes) plus the results of any blocks and sub-agents that just finished, which are then dropped."""
done = [w for w in self.blocks + self.sub_agents if w.task.done()]
self.blocks = [b for b in self.blocks if not b.task.done()]
self.sub_agents = [a for a in self.sub_agents if not a.task.done()]
events, self.inbox = self.inbox, []
return "".join(events + [w.task.result() for w in done])
def is_completed(self):
"""True when nothing of this agent's is still running: no blocks, no sub-agents, no inbox messages."""
return not self.blocks and not self.sub_agents and not self.inbox
async def wait_first_event(self):
"""Blocks until the first event: a block or sub-agent task finishing, or a running one streaming a
heartbeat (on_heartbeat sets self.wake). Returns at once if output is already pending."""
tasks = [w.task for w in self.blocks + self.sub_agents]
if not self.inbox and tasks:
self.wake.clear()
waker = asyncio.ensure_future(self.wake.wait())
await asyncio.wait([*tasks, waker], return_when=asyncio.FIRST_COMPLETED)
waker.cancel()
def steer(self, text):
self.inbox.append(wrap("steer", "", text))
self.wake.set()
def start_block(self, block):
if any(b.id == block.id and b.is_running() for b in self.blocks):
self.display.event("failed", name=self.name, message=f"id={block.id} already running, not restarted")
self.inbox.append(wrap("error", f" id={block.id}", "already running, not restarted; reply in plain text to wait for it"))
return
block.launch(self.namespace, self.display, self.name, self.on_heartbeat)
self.display.event("launched", name=self.name, id=block.id, timeout=block.timeout, heartbeat=block.heartbeat, code=block.code)
self.blocks.append(block)
def start_agent(self, agent):
if self.depth >= self.max_depth:
raise RuntimeError(f"max delegation depth {self.max_depth} reached — do this part yourself instead of delegating")
taken = {child.name for child in self.sub_agents}
base, suffix = agent.name, 2
while agent.name in taken:
agent.name = f"{base}{suffix}"
suffix += 1
namespace = agent.namespace_values or {name: self.namespace[name] for name in agent.namespace_names}
model = resolve_model(agent.model) if agent.model else getattr(self.llm, "model", None)
child = create_agent(namespace=namespace, display=self.display, depth=self.depth + 1, max_depth=self.max_depth, name=agent.name, context=agent.context, model=model)
child.start_task(agent.assignment, parse_duration(agent.heartbeat) if agent.heartbeat else None, self.on_heartbeat, supervisor=self)
self.sub_agents.append(child)
return child
def shutdown(self):
for block in self.blocks:
block.cancel()
for agent in self.sub_agents:
if not agent.task.done():
agent.task.cancel()
def start_task(self, task, heartbeat, on_heartbeat, supervisor=None):
"""Starts this agent as a background sub-agent and returns self as the control handle."""
self.supervisor = supervisor or self.supervisor
self.start = time.time()
self.heartbeat_interval = heartbeat
self._parent_on_heartbeat = on_heartbeat
self.result_value = None
self.task = asyncio.create_task(self.run_as_subagent(task, heartbeat, on_heartbeat))
return self
async def run(self, task, max_messages=200):
"""The conversation loop: send a message, run the reply's blocks/agents, feed the results back, until the
agent answers in plain text (process_llm_response returns None) or runs out of messages."""
try:
message = self.intro + task
for _ in range(max_messages):
reply = await self.llm.send(message)
self.process_llm_response(reply)
if self.is_completed():
return reply
await self.wait_first_event()
await asyncio.sleep(0)
message = self.next_message() + self.context_note()
raise RuntimeError(f"Task unfinished after {max_messages} messages")
finally:
self.shutdown()
async def run_as_subagent(self, task, heartbeat, on_heartbeat, max_messages=200):
"""Runs run() in the background as a sub-agent: streams <agent_status> heartbeats to the parent while it
works, then wraps the final answer (or cancellation/failure) in <agent_result>."""
agent_task = asyncio.create_task(self.run(task, max_messages))
try:
while not agent_task.done():
await asyncio.wait([agent_task], timeout=heartbeat)
if not agent_task.done() and heartbeat:
on_heartbeat(wrap("agent_status", f" name={self.name} elapsed={time.time() - self.start:.0f}s", self.status()))
except asyncio.CancelledError:
if not agent_task.done():
agent_task.cancel()
await asyncio.wait([agent_task])
raise
if agent_task.cancelled():
body = "[cancelled]"
elif agent_task.exception():
body = f"[failed] {agent_task.exception()!r}"
else:
self.result_value = agent_task.result()
body = self.result_value
return wrap("agent_result", f" name={self.name} duration={time.time() - self.start:.0f}s", body)
def done(self):
return self.task is None or self.task.done()
def result(self):
if self.result_value is not None:
return self.result_value
if self.task is None:
return None
return self.task.result()
def cancel(self):
"""Cancels this sub-agent."""
self.display.event("cancelled", name=self.name, label="cancelled by caller")
if self.task is not None:
self.task.cancel()
def new_task(self, task):
"""Asks a follow-up that continues this sub-agent's existing context and namespace."""
if self.task is not None and not self.task.done():
raise RuntimeError(f"sub-agent '{self.name}' is still working; wait for its result before new_task")
self.start_task(task, self.heartbeat_interval, self._parent_on_heartbeat)
self.supervisor.sub_agents.append(self)
return self
def namespace_note(self):
items = []
for name, value in self.namespace.items():
if name in self.base_names or name.startswith("_") or inspect.ismodule(value):
continue
descr = type(value).__name__
if isinstance(value, (str, bytes, list, dict, tuple, set)):
descr += f"({len(value)})"
items.append(f"{name}:{descr}")
return f" vars=[{', '.join(items[:25])}]" if items else ""
def context_note(self):
note = f"<status context_tokens={self.llm.context} spent=${self.llm.cost:.4f} last_reply={self.llm.last_latency:.1f}s{self.namespace_note()}/>\n"
if self.llm.context > 16000:
note += ('<reminder>Your context is growing and every turn now costs more. Call '
'self.compact_knowledge("...") NOW, before your next action — pass a SELF-CONTAINED summary that '
'keeps every exact URL, file path you created, value/credential/deduction found, the item you are '
'on and the next step. Everything you leave out is GONE. Do it this turn, not in a few turns.</reminder>\n')
return note
def compact_knowledge(self, summary):
"""Replaces your whole working context with `summary`, keeping what you learned but dropping the raw
history to free context. Put every fact, decision and next step you still need into it — the rest is lost."""
self.llm.compact(summary)
def cancel_subagents(self):
"""Cancels every sub-agent you spawned that is still running."""
alive = [agent for agent in self.sub_agents if not agent.task.done()]
self.display.event("cancelled", name=self.name, label=f"cancel_subagents ({len(alive)} running)")
for agent in alive:
agent.task.cancel()
def run_agent(self, task, name="sub", context="", namespace=None, heartbeat="", model=""):
"""Starts task in a fresh background sub-agent (its own context) and returns its handle instantly — never
await it. Its final answer arrives by itself as <agent_result name=...>; heartbeat="30s" streams
<agent_status> heartbeats while it works. namespace={...} gives it the tools/objects its code needs;
context="..." seeds it with what an earlier agent learned so it doesn't start from scratch;
model="flash"/"gpt"/... runs it on a different model (default: your own)."""
agent = Agent(name=name, context=context)
agent.assignment = task
agent.namespace_names = list((namespace or {}).keys())
agent.namespace_values = namespace or {}
agent.heartbeat = heartbeat
agent.model = model
return self.start_agent(agent)
def cost(self):
"""Resource use so far: {context, generated, cost, seconds}. Works on yourself (self.cost()) or a handle."""
return self.llm.stats
def status(self):
return f"[{self.name} - {format_stats(self.llm.stats)}]"
def create_agent(functions=None, namespace=None, display=None, depth=0, max_depth=3, name="main", context="", model=None):
"""Wires a production agent together: system prompt, LLM conversation, runtime namespace, and display."""
os.makedirs("files", exist_ok=True)
tools = {**(namespace or {}), **(functions or {})}
display = display or AgentDisplay()
model = model or DEFAULT_MODEL
llm = Chat(build_system(tools, model), display=display, name=name, model=model)
return Agent(llm=llm, namespace=tools, display=display, depth=depth, max_depth=max_depth, name=name, context=context)
if __name__ == "__main__":
from display import RichDisplay
repo_files = ["agent.py", "utils.py", "test_agent.py", "display.py", "llm.py", "CLAUDE.md", "AGENTS.md"]
async def read_repo_file(path):
if path not in repo_files:
raise ValueError(f"allowed files: {repo_files}")
with open(path, encoding="utf-8") as f:
return f.read()
example_tasks = {
"contract": """
Use <run_agent> to split this review into two sub-agents. Give both namespace="read_repo_file,repo_files".
1. One sub-agent should inspect agent.py and report the exact runtime contract for <run_code> and <run_agent>: parsing, launch path, result format, heartbeat behavior, and cancellation behavior.
2. One sub-agent should inspect test_agent.py, utils.py, and CLAUDE.md and report where tests or guidelines do not cover that contract.
Then combine their <agent_result> outputs into a concrete review: what is coherent, what is under-tested, and the next two tests/refactors you would do.
""",
"concurrency": """
Use <run_agent> to split this investigation into two sub-agents. Give both namespace="read_repo_file,repo_files".
1. One sub-agent should inspect CodeBlock in agent.py and explain the lifecycle of a code block from parsing to final <run_code_result>, including timeout and heartbeat updates.
2. One sub-agent should inspect utils.py and test_agent.py and decide whether concurrent code blocks can safely capture stdout/stderr without mixing outputs.
Then combine their <agent_result> outputs into a short technical conclusion with evidence and any remaining race-condition risk.
""",
"tag": """
Use <run_agent> to split this feature audit into two sub-agents. Give both namespace="read_repo_file,repo_files".
1. One sub-agent should inspect all code paths related to <run_agent>, from Agent.find_all to Agent.start_agent and Agent.run.
2. One sub-agent should inspect SYSTEM and tests to verify what the LLM is told versus what the implementation actually supports.
Then combine their <agent_result> outputs into a precise answer: supported attributes, unsupported-but-documented behavior, documented-but-untested behavior, and one minimal patch recommendation.
""",
}
agent = create_agent(
namespace={
"read_repo_file": read_repo_file,
"repo_files": repo_files,
},
display=RichDisplay(),
)
asyncio.run(
agent.run(
example_tasks[os.environ.get("AGENT_EXAMPLE", "contract")]
)
)