Files
agent-harness/agent-sessions/agent.py

1352 lines
58 KiB
Python

import argparse
import hashlib
import json
import os
import atexit
from contextlib import ExitStack
from enum import Enum
from pathlib import Path
from openai import OpenAI
from tools import (
get_tool_registry,
get_tool_schemas,
DockerSandbox,
DockerSandboxError,
ACTION_TOOLS,
AuditLog,
ToolValidator,
bounded_schemas,
)
from tools.scratchpad import scratchpad as scratchpad_state
from dotenv import load_dotenv
import prompt_safety
import tool_policy
import resource_limits
import secret_management
import session_control
from sessions import SessionStore as MemoryStore
from tools.todo import todo_store
from tools.memory import configure_memory
from resource_limits import (
IterationCaps,
ContextBudget,
CostTracker,
cap_tool_result,
)
from session_control import AbortController, FileRollback, kill_in_flight
load_dotenv()
# Raw, in-process tool registry. Action tools are replaced below with
# wrappers that dispatch into the Docker sandbox container.
RAW_TOOL_REGISTRY = get_tool_registry()
# Limit output scope: the schemas exposed to the LLM carry bounds
# (min/max on offset/limit, maxLength on command/content/url, relative-
# path format on glob patterns). These bounds are enforced host-side by
# TOOL_VALIDATOR before any tool is allowed to run.
TOOL_SCHEMAS = bounded_schemas(get_tool_schemas())
TOOL_VALIDATOR = ToolValidator(TOOL_SCHEMAS)
# ---------------------------------------------------------------------------
# Permission modes
# ---------------------------------------------------------------------------
class PermissionMode(Enum):
DEFAULT = "default"
ACCEPT_EDITS = "acceptEdits"
DANGEROUSLY_SKIP_PERMISSIONS = "dangerouslySkipPermissions"
# Always allowed: read-only filesystem tools
READ_TOOLS = {"read_file", "glob_files", "grep"}
# Always allowed: internal planning/bookkeeping and user-interaction tools (no external side effects)
PLANNING_TOOLS = {"todo_append", "todo_list", "todo_update",
"read_scratchpad", "write_scratchpad", "ask_question",
"read_memory", "append_memory", "list_sessions",
"search_sessions", "read_session", "write_session_summary"}
# Conditionally allowed in acceptEdits mode when target is within working dir
WRITE_TOOLS = {"write_file", "edit_file"}
def _resolve_tool_path(tool_name: str, args: dict) -> str | None:
"""Return the file-path argument for write tools, or None if not applicable."""
if tool_name in WRITE_TOOLS:
return args.get("path")
return None
def _is_within_working_dir(path: str, working_dir: Path) -> bool:
"""Return True if *path* resolves to somewhere inside *working_dir*."""
try:
target = Path(path)
if not target.is_absolute():
target = working_dir / target
target.resolve().relative_to(working_dir.resolve())
return True
except ValueError:
return False
def _ask_permission(tool_name: str, args: dict) -> bool:
"""Interactively ask the user whether to allow a tool call.
Returns True if the user grants permission, False otherwise.
"""
print(f"\n [permission required] {tool_name}")
print(f" Arguments: {json.dumps(args, ensure_ascii=False)}")
while True:
try:
answer = input(" Allow this action? [y/n]: ").strip().lower()
except EOFError:
print(" (EOF, denying permission)")
return False
if answer in ("y", "yes"):
return True
if answer in ("n", "no"):
return False
print(" Please enter 'y' or 'n'.")
def check_permission(
tool_name: str,
args: dict,
mode: PermissionMode,
working_dir: Path,
) -> tuple[bool, str | None]:
"""Decide whether a tool call is permitted under the current mode.
Returns ``(allowed, reason)``. When ``allowed`` is False, ``reason``
explains why and is surfaced to the LLM. May interactively prompt
the user when a decision cannot be made automatically.
Three layers, evaluated in order:
1. Tool policy: a hard gate that no mode can
override: path scoping, shell denylist, SSRF guard.
2. Always-confirm patterns: destructive calls require
explicit confirmation even in acceptEdits, and selected
irreversible patterns refuse even in dangerouslySkipPermissions.
3. Mode-based decision (default / acceptEdits / skip).
Permission rules (layer 3)
--------------------------
default
Read tools and planning tools run freely. Every other tool
requires explicit user approval.
acceptEdits
Read tools and planning tools run freely. Write tools
(write_file, edit_file) run freely only when the target path is
inside the working directory; otherwise the user is prompted.
All other tools require explicit user approval.
dangerouslySkipPermissions
All tools run without any prompt, EXCEPT calls blocked by layer
1 (hard policy) or layer 2 (always-confirm / irreversible).
"""
# Layer 1: hard policy gate (path scope, shell, SSRF).
ok, reason = tool_policy.check_tool_policy(tool_name, args, working_dir)
if not ok:
return False, reason
# Layer 2: always-confirm & irreversible patterns.
# These override even dangerouslySkipPermissions for the worst cases.
confirm_required, confirm_reason = tool_policy.always_confirm_required(
tool_name, args)
if confirm_required:
if mode == PermissionMode.DANGEROUSLY_SKIP_PERMISSIONS:
# Irreversible calls are refused outright in skip mode; the
# user explicitly accepted risk for normal destructive ops,
# but not for e.g. `git push --force`.
return False, (
f"Blocked even in dangerouslySkipPermissions: {
confirm_reason or 'irreversible action'}. "
"Run this command manually outside the agent if it is truly intended."
)
# In default / acceptEdits, force an explicit prompt.
return _ask_permission(f"{tool_name} [DESTRUCTIVE]", args), None
# Planning and read tools are always free regardless of mode
if tool_name in READ_TOOLS or tool_name in PLANNING_TOOLS:
return True, None
if mode == PermissionMode.DANGEROUSLY_SKIP_PERMISSIONS:
return True, None
if mode == PermissionMode.ACCEPT_EDITS and tool_name in WRITE_TOOLS:
path = _resolve_tool_path(tool_name, args)
if path and _is_within_working_dir(path, working_dir):
# an edit that empties an existing file is a delete,
# require confirmation rather than auto-approving.
if _is_delete_via_write(tool_name, args, path):
return _ask_permission(f"{tool_name} [DELETE-via-empty]", args), None
return True, None # auto-approved, within the working directory
# Path is outside the working directory → fall through to ask
# Default mode, or acceptEdits for non-write / out-of-tree tools
return _ask_permission(tool_name, args), None
def _is_delete_via_write(tool_name: str, args: dict, path: str) -> bool:
"""Return True if a write_file/edit_file call effectively deletes a file.
- ``write_file`` overwriting an existing file with empty/whitespace
content is treated as a delete.
- ``edit_file`` replacing content with empty content on a file that
becomes empty is also flagged.
"""
if tool_name == "write_file":
content = args.get("content", "")
if not content.strip():
try:
if Path(path).exists() and Path(path).stat().st_size > 0:
return True
except OSError:
pass
return False
def build_tool_registry(
sandbox: DockerSandbox,
allowed_tools: set[str] | None = None,
) -> dict:
"""Return a tool registry routing action tools into the container.
Planning/stateful tools (todo, scratchpad, ask_question) run
in-process on the host so their state survives across calls.
Action tools (filesystem, shell, web) are dispatched into the
long-lived Docker sandbox container.
If *allowed_tools* is given, only tools in that set are exposed,
this is the principle of least privilege. ``None``
means "expose everything".
"""
registry = {}
for name, fn in RAW_TOOL_REGISTRY.items():
if allowed_tools is not None and name not in allowed_tools:
continue
if name in ACTION_TOOLS:
def _dispatch(_n=name, _s=sandbox, **kwargs):
return _s.run_tool(_n, kwargs)
registry[name] = _dispatch
else:
registry[name] = fn
return registry
def filter_tool_schemas(schemas: list[dict], allowed: set[str] | None) -> list[dict]:
"""Return only the schemas whose tool name is in *allowed*.
``None`` means "no filtering".
"""
if allowed is None:
return schemas
return [s for s in schemas if s["function"]["name"] in allowed]
# ---------------------------------------------------------------------------
# Agent loop
# ---------------------------------------------------------------------------
def get_llm_client():
return OpenAI(
base_url="http://localhost:11434/v1",
api_key="."
)
def handle_tool_calls(
tool_calls,
messages,
mode: PermissionMode,
working_dir: Path,
tool_registry: dict,
audit: AuditLog,
user_goal: str = "",
iteration_caps: IterationCaps | None = None,
tool_timeout: float = 120.0,
abort_controller: AbortController | None = None,
file_rollback: FileRollback | None = None,
container_name: str | None = None,
plan_state: dict | None = None,
):
"""Execute each tool the LLM requested and append the results to messages.
Every call passes through these layers, in order, and each decision
is recorded in the audit log:
0. Abort check + iteration cap.
1. Schema validation.
2. Permission gate, may prompt the user.
3. Execution: action tools run inside the Docker sandbox container;
planning tools run in-process on the host. A container failure
or timeout is reported back to the LLM as an error.
4. File rollback snapshot before write/edit.
"""
caps = iteration_caps or IterationCaps()
abort = abort_controller or AbortController()
rollback = file_rollback or FileRollback()
for tool_call in tool_calls:
# session-level abort: stop dispatching immediately.
if abort.triggered:
audit.log_session_abort(
f"tool dispatch halted: {abort.reason}"
)
messages.append({
"role": "user",
"content": (
"Session abort requested. Stop calling tools and "
"give the user a concise summary of what was done."
),
})
return
# session-level tool-call cap.
cap_reason = caps.bump_tool_call()
if cap_reason is not None:
audit.log("iteration_cap_hit", scope="session",
tool_calls=caps.tool_calls)
# Inject a single stop message and abort this batch.
messages.append({
"role": "user",
"content": cap_reason,
})
return
name = tool_call.function.name
# --approve-plan: if the user asked for plan approval and
# the model is now trying to run an action tool for the first
# time, pause and show the scratchpad + todo list. Planning
# tools (scratchpad/todo/ask_question) run freely so the model
# can build its plan before the gate fires.
if plan_state and not plan_state.get("approved"):
if name in ACTION_TOOLS and not plan_state.get("prompted"):
plan_state["prompted"] = True
scratch = scratchpad_state.read()
print("\n [plan approval required]")
print(f" The agent wants to run '{name}'.")
if scratch and scratch != "(empty)":
print(f" --- scratchpad ---\n{scratch[:1000]}")
print(" --- end plan ---")
try:
ans = input(" Approve this plan? [y/n]: ").strip().lower()
except EOFError:
ans = "n"
if ans in ("y", "yes"):
plan_state["approved"] = True
audit.log("plan_approved", tool=name)
else:
audit.log("plan_rejected", tool=name)
result = (
"The user rejected the plan. Do not proceed with "
"action tools. Revise the plan in the scratchpad "
"and ask the user for confirmation again."
)
messages.append({
"role": "tool",
"tool_call_id": tool_call.id,
"content": prompt_safety.wrap_tool_result(name, result),
})
continue
# Schema-validate tool inputs: parse and validate before
# any policy/permission check. A malformed call is reported
# back to the LLM with the specific errors so it can correct
# and retry, it never reaches the sandbox or the permission
# gate.
try:
args = json.loads(tool_call.function.arguments)
except json.JSONDecodeError as e:
print(f" [tool] {name}(<unparseable args>)")
audit.log(
"validation_error",
tool=name,
raw_arguments=tool_call.function.arguments,
errors=[f"invalid JSON: {e}"],
)
messages.append({
"role": "tool",
"tool_call_id": tool_call.id,
"content": prompt_safety.wrap_tool_result(
name,
f"Error: tool arguments were not valid JSON: {e}. "
"Resend the call with correctly formatted JSON arguments.",
),
})
continue
print(f" [tool] {name}({args})")
permission_allowed = False
permission_reason: str | None = None
container_error = False
container_reason = None
intent_drift = False
intent_reason: str | None = None
validation_errors: list[str] | None = None
result: str
# schema + bounds validation.
v_ok, v_errs = TOOL_VALIDATOR.validate(name, args)
if not v_ok:
validation_errors = v_errs
result = (
"Error: tool arguments failed schema validation:\n- "
+ "\n- ".join(v_errs)
+ "\nCheck the tool schema and retry with valid arguments."
)
print(f" [validation] {v_errs}")
audit.log(
"validation_error",
tool=name,
args=args,
errors=v_errs,
)
audit.log_tool_result(
tool_call_id=tool_call.id,
name=name,
args=args,
permission_allowed=False,
permission_reason="schema validation failed",
container_error=False,
container_reason=None,
intent_drift=False,
intent_reason=None,
result=result,
)
messages.append({
"role": "tool",
"tool_call_id": tool_call.id,
"content": prompt_safety.wrap_tool_result(name, result),
})
continue
# Re-validate intent after tool use: flag high-risk calls
# whose arguments reference resources not present in the user's
# original goal or the scratchpad. This catches the model
# acting on instructions injected via prior tool output.
ok, reason = prompt_safety.intent_check(
user_goal=user_goal,
scratchpad=scratchpad_state.read(),
tool_name=name,
tool_args=args,
)
if not ok:
intent_drift = True
intent_reason = reason
print(f" [intent drift] {reason}")
audit.log(
"intent_drift_suspected",
tool=name,
args=args,
reason=reason,
)
if name not in tool_registry:
result = (
f"Error: unknown tool '{name}'. "
f"Available tools: {list(tool_registry.keys())}"
)
else:
# When intent drift is detected, force an explicit user
# confirmation even in auto-approve modes (but honor
# dangerouslySkipPermissions, where the user has accepted
# all risk). The reason is shown in the prompt.
if intent_drift and mode != PermissionMode.DANGEROUSLY_SKIP_PERMISSIONS:
permission_allowed = _ask_permission(
f"{name} [INTENT DRIFT: {reason}]",
args,
)
permission_reason = None
else:
permission_allowed, permission_reason = check_permission(
name, args, mode, working_dir,
)
# log the permission decision as a standalone event.
audit.log_permission_decision(
tool=name,
args=args,
mode=mode.value,
allowed=permission_allowed,
reason=permission_reason,
)
if not permission_allowed:
result = (
f"Permission denied: {
permission_reason or 'the user did not allow ' + name + ' to run.'} "
"Do not retry this tool call without asking the user first."
)
else:
# File rollback: snapshot the original file before
# any write/edit so the session can offer to revert on
# abort. Only files inside the working dir are
# snapshotted (writes outside are already blocked by).
if name in ("write_file", "edit_file"):
target_path = args.get("path", "")
if target_path:
rollback.snapshot(target_path)
try:
result = tool_registry[name](**args)
except DockerSandboxError as e:
container_error = True
container_reason = str(e)
result = (
f"Container error: {e}. The sandbox container could "
"not execute this tool. Do not retry without changing "
"the approach or asking the user."
)
# distinguish timeouts from other container errors
# for the audit log.
if "timed out" in str(e).lower():
audit.log("tool_timeout", tool=name,
timeout_s=tool_timeout)
except TypeError as e:
result = (
f"Error: invalid arguments for tool '{name}': {e}. "
"Check the tool schema and retry with the correct arguments."
)
print(f" [tool result] {result[:200]}{
'...' if len(result) > 200 else ''}")
# Treat external data as data: mark web pages and files read
# from outside the working directory as <external_document> so
# the model does not treat their bytes as instructions.
result = prompt_safety.mark_external_content(
tool_name=name,
tool_args=args,
result=result,
working_dir=working_dir,
)
audit.log_tool_result(
tool_call_id=tool_call.id,
name=name,
args=args,
permission_allowed=permission_allowed,
permission_reason=permission_reason,
container_error=container_error,
container_reason=container_reason,
intent_drift=intent_drift,
intent_reason=intent_reason,
result=result,
)
# Delimit context clearly: wrap the tool result so the model
# can distinguish data coming back from a tool from the user's
# instructions. This is the core prompt-injection defense.
#
# cap the result length before insertion to bound context
# growth (the audit log above already has the full content).
capped = cap_tool_result(result)
messages.append({
"role": "tool",
"tool_call_id": tool_call.id,
"content": prompt_safety.wrap_tool_result(name, capped),
})
def agent_loop(
client,
mode: PermissionMode,
working_dir: Path,
tool_registry: dict,
audit: AuditLog,
tool_schemas: list[dict] | None = None,
iteration_caps: IterationCaps | None = None,
context_budget: ContextBudget | None = None,
cost_tracker: CostTracker | None = None,
llm_timeout: float | None = 120.0,
tool_timeout: float = 120.0,
abort_controller: AbortController | None = None,
file_rollback: FileRollback | None = None,
container_name: str | None = None,
approve_plan: bool = False,
memory_store: MemoryStore | None = None,
session_id: str | None = None,
resumed_messages: list[dict] | None = None,
):
# resource controls: the harness, never the model, owns these.
iteration_caps = iteration_caps or IterationCaps()
context_budget = context_budget or ContextBudget()
cost_tracker = cost_tracker or CostTracker()
# session abort & rollback
abort_controller = abort_controller or AbortController()
file_rollback = file_rollback or FileRollback()
sandbox_exec_timeout = tool_timeout
# --approve-plan: when True, the first time the model tries to
# run an action tool after writing to the scratchpad, we pause and
# ask the user to approve the plan before any action runs.
plan_approved = not approve_plan # if flag off, always approved
plan_approval_pending = False
# Mutable state shared across handle_tool_calls invocations within
# a single user turn.
plan_state: dict | None = {"approved": plan_approved,
"prompted": False} if approve_plan else None
# Statically audit the system prompt for secret interpolation
# before it ever reaches the model. This is a defense-in-depth
# check: the prompt below is a static string with no env-var
# interpolation, but if someone later edits it to inject a key,
# this raises immediately at startup.
_system_prompt_body = (
"You are a capable coding and research assistant.\n\n"
+ prompt_safety.TRUST_BOUNDARIES
+ "\n\n"
"## Available tools\n\n"
"Action tools: read_file, write_file, edit_file, glob_files, grep, run_bash, webfetch\n\n"
"Planning tools:\n"
"- Scratchpad (read_scratchpad / write_scratchpad): your private working memory. "
"Use it to think through an approach, store intermediate findings, or draft content "
"before committing. Each write fully replaces the previous content.\n"
"- To-do list (todo_append / todo_list / todo_update): a persistent task tracker. "
"Items carry a status: pending, in_progress, done, cancelled, or failed.\n"
"- Clarification (ask_question): ask the user a single focused question when you "
"are genuinely blocked and cannot reasonably infer the missing information from "
"context. Do not use it for progress updates or to confirm actions you can already "
"take, only ask when it is strictly necessary to proceed.\n\n"
"- Long-term memory (read_memory / append_memory): at the start of a session, read "
"both project and user memory. Project memory is for durable lessons, environment "
"facts, and conventions; user memory is for stable preferences and expectations. "
"Append only concise, useful facts that are likely to hold in future sessions. Never "
"store secrets, credentials, or transient task details.\n"
"- Past sessions (list_sessions / search_sessions / read_session): use these to find "
"and inspect prior work. Before your final response on a substantial task, call "
"write_session_summary with the goal, completed work, decisions, and remaining work.\n\n"
"## Execution environment (Docker sandbox)\n\n"
"Action tools: read_file, write_file, edit_file, glob_files, grep, run_bash, "
"webfetch, execute inside a Docker container. The user's project directory is "
"bind-mounted into that container at the same path as on the host, and the "
"container's working directory is the project root. You can only see and modify "
"files inside the project mount; the rest of the container's filesystem is a "
"minimal, read-only-by-convention Linux image. If the session was started with "
"network disabled, run_bash and webfetch that need network will fail, treat that "
"as expected, not as a bug to work around. If a tool returns 'Container error: "
"...', do not retry the same call, adjust the approach or ask the user to run it "
"manually. Planning tools run on the host, not in the container, so their state "
"persists across calls.\n\n"
"## Working directory\n\n"
"The current working directory is always the user's project root. "
"When asked to work on a project or codebase without a specified path, "
"start by exploring '.' with glob_files or run_bash. "
"Never ask the user to supply a path.\n\n"
"## How to plan\n\n"
"For complex or multi-step tasks (roughly 3 or more distinct steps, or when the "
"path forward is unclear):\n"
"1. Write your initial thinking and approach to the scratchpad before acting.\n"
"2. Break the work into concrete steps and add each one to the to-do list with "
"todo_append (status: pending).\n"
"3. Before starting a step, mark it in_progress with todo_update. "
"Keep only one item in_progress at a time.\n"
"4. Mark items done immediately after completing them, do not batch completions.\n"
"5. Call todo_list to review remaining work before moving to the next step.\n"
"6. Mark tasks cancelled if they become unnecessary.\n\n"
"For simple, single-step tasks: act directly without creating todos.\n\n"
"Planning tool calls (write_scratchpad, todo_append, todo_update, todo_list) "
"are internal bookkeeping, not responses to the user. After any planning tool "
"call, always continue working immediately, make your next tool call or, once "
"the task is fully complete, give a substantive final answer. "
"Never emit an empty or whitespace-only message.\n\n"
"## Replanning\n\n"
"After every tool result, check whether the outcome matched your expectation. "
"If a tool returns an error, unexpected output, or reveals information that "
"changes your understanding of the task, do not move to the next planned step, "
"replan first.\n\n"
"When a step fails:\n"
"1. Diagnose in the scratchpad: is this a recoverable input error (wrong path, "
"typo, wrong argument) or a deeper problem (wrong approach, wrong assumption)?\n"
"2. Mark the task failed: todo_update(id, status='failed').\n"
"3. Choose a recovery action:\n"
" - Retry: the failure is correctable. Fix the input and set the task back to "
"in_progress. The tool will report which retry attempt this is.\n"
" - Replace: the approach is wrong. Cancel the task and add a revised one.\n"
" - Reorder: new information makes a different task more urgent. Update the "
"pending items before continuing.\n"
"4. If todo_update reports that the retry limit has been reached, stop retrying. "
"Write a clear diagnosis in the scratchpad: what you tried, what failed each "
"time, and what you need, then give the user a concise escalation message "
"and wait for their input.\n\n"
"When a tool succeeds but returns information that changes the picture, pause "
"before acting. Call todo_list, reassess all pending items in the scratchpad, "
"and cancel or replace any tasks that no longer make sense.\n\n"
"## How to use the scratchpad\n\n"
"Before each tool call during a complex task, update the scratchpad with your "
"current thinking. Structure each entry around these five steps:\n\n"
"1. Restate the goal: write what you understand the task to be, in your own words. "
"This catches misreads before they compound into wasted work.\n"
"2. Survey what you know: note which files you have seen, what the code structure "
"looks like, and what constraints or requirements apply.\n"
"3. Evaluate options: reason through at least two approaches and explain why you "
"are choosing one over the other (e.g. 'I could rewrite the middleware, or wrap it. "
"Wrapping is safer because it leaves the existing call sites untouched.').\n"
"4. Anticipate failure modes: write down what could go wrong with the chosen "
"approach and how you would diagnose it (e.g. 'If the tests fail after this, the "
"most likely cause is that the session cookie name changed.').\n"
"5. Decide the next single action: commit to exactly one tool call. "
"Do not plan several calls at once; decide the next step only.\n\n"
"Re-read the scratchpad whenever you resume after a tool result to keep your "
"reasoning grounded in what you have already learned.\n\n"
"## Done detection\n\n"
"Do not give a final answer based on the task list being empty alone. "
"Before declaring the task complete, verify all three of the following:\n\n"
"1. Structural completion: call todo_list and confirm there are no pending, "
"in_progress, or failed items.\n"
"2. Verification: check the output against the original goal. For code tasks: "
"run the tests or build with run_bash and confirm they pass. For research tasks: "
"re-read the scratchpad and confirm the assembled answer addresses what was "
"actually asked.\n"
"3. Uncertainty check: read the scratchpad and ask: are there unresolved "
"questions, assumptions that were never validated, or tasks that were cancelled "
"rather than properly completed?\n\n"
"If all three are satisfied, give your final answer. If any are not, re-enter "
"the planning loop, add the outstanding items to the todo list and continue."
)
_prompt_warnings = secret_management.audit_system_prompt(
_system_prompt_body)
if _prompt_warnings:
for w in _prompt_warnings:
print(f" [secret-safety] {w}")
raise RuntimeError(
"System prompt failed secret-safety audit. Refusing to start. "
"See secret_management.audit_system_prompt for details."
)
initial_messages = [
{
"role": "system",
"content": _system_prompt_body,
}
]
messages = resumed_messages or initial_messages
if memory_store is not None and session_id is not None:
if resumed_messages is None:
memory_store.create_session(session_id, messages)
record = memory_store.load_session(session_id)
state = record.get("state", {})
scratchpad_state.write(state.get("scratchpad", ""))
todo_store._items = state.get("todos", [])
for field in ("total_tokens_in", "total_tokens_out", "total_cost_usd", "calls"):
if field in state:
setattr(cost_tracker, field, state[field])
iteration_caps.tool_calls = state.get("tool_calls", 0)
memory_store.update(session_id, status="active", project_root=str(working_dir))
configure_memory(memory_store, session_id)
def checkpoint():
if memory_store is not None and session_id is not None:
memory_store.save_messages(session_id, messages)
memory_store.update(session_id, state={
"scratchpad": scratchpad_state._content,
"todos": todo_store.read(include_completed=True),
"tool_calls": iteration_caps.tool_calls,
**{field: getattr(cost_tracker, field) for field in
("total_tokens_in", "total_tokens_out", "total_cost_usd", "calls")},
})
while True:
if abort_controller.triggered:
checkpoint()
break
try:
user_input = input("You: ")
except EOFError:
checkpoint()
break
if user_input.lower() == "\\exit":
break
if user_input.lower() == "\\sessions" and memory_store is not None:
sessions = memory_store.list_sessions()
if sessions:
for session in sessions:
print(f"{session['id']} | {session['updated_at']} | "
f"{session['summary'][:160] or '(no summary)'}")
else:
print("No saved sessions.")
continue
if user_input.startswith("\\") and memory_store is not None:
command, _, argument = user_input.partition(" ")
try:
if command == "\\compact":
replacement = memory_store.compact(
session_id, client, context_budget, cost_tracker, llm_timeout, force=True)
if replacement is not None:
messages = replacement
checkpoint()
print("Context compacted." if replacement else "No older turns to compact.")
elif command == "\\title" and argument.strip():
memory_store.update(session_id, title=argument.strip())
elif command == "\\fork":
checkpoint()
print(f"Fork created: {memory_store.fork(session_id)} (resume it separately)")
elif command == "\\info":
record = memory_store.load_session(session_id)
print(f"{session_id} | {record['title']} | "
f"{len(record['transcript'])} transcript messages | "
f"{len(record['compactions'])} compactions | {cost_tracker.summary()}")
elif command == "\\help":
print("\\exit, \\sessions, \\compact, \\title NAME, \\fork, \\info, \\help")
else:
print("Unknown command. Type \\help.")
except Exception as error:
print(f"Session operation failed: {error}")
continue
audit.log_user_message(user_input)
# Delimit context clearly: wrap user input so the model can
# tell it apart from tool output. Keep the raw goal for the
# intent check after tool calls.
current_user_goal = user_input
messages.append({
"role": "user",
"content": prompt_safety.wrap_user_input(user_input),
})
if memory_store is not None and session_id is not None:
memory_store.save_messages(session_id, messages)
# reset the per-turn counter for each new user message.
iteration_caps.reset_turn()
# Keep looping until the LLM stops calling tools and gives a final reply
while True:
# session-level abort: check before every LLM turn.
if abort_controller.triggered:
audit.log_session_abort(
f"loop halted between turns: {abort_controller.reason}"
)
print(f"\n [abort] {abort_controller.reason}")
print("Assistant: Session aborted. Summarizing what was done.")
break
# hard iteration cap: stop the loop if the model has
# taken too many turns for a single user message.
cap_reason = iteration_caps.bump_turn()
if cap_reason is not None:
audit.log("iteration_cap_hit", scope="per_turn",
turns=iteration_caps.turns,
tool_calls=iteration_caps.tool_calls)
messages.append({
"role": "user",
"content": cap_reason,
})
# cost circuit breaker: abort the session if spend
# exceeds the configured cap.
cost_reason = cost_tracker.check()
if cost_reason is not None:
audit.log("cost_limit_hit", **{
"cost_usd": cost_tracker.total_cost_usd,
"cap_usd": cost_tracker.max_cost_usd,
"tokens_in": cost_tracker.total_tokens_in,
"tokens_out": cost_tracker.total_tokens_out,
"calls": cost_tracker.calls,
})
print(f"\n [cost] {cost_reason}")
print(f"Assistant: I've hit the session cost limit "
f"(${cost_tracker.total_cost_usd:.4f}). "
"Stopping here; increase --max-cost-usd if you need more.")
break
# Never silently discard context. Persist compaction independently
# from the immutable conversation transcript.
if memory_store is not None and session_id is not None:
try:
checkpoint()
replacement = memory_store.compact(
session_id, client, context_budget, cost_tracker, llm_timeout)
if replacement is not None:
messages = replacement
checkpoint()
audit.log("session_compacted", message_count=len(messages))
except Exception as error:
audit.log("compaction_error", error=str(error))
print(f" [compaction] {error}")
if context_budget.estimate_total(messages) > context_budget.max_tokens:
print("Context exceeds the budget. Start a fork/new session or increase --max-context-tokens.")
break
if cost_tracker.check():
checkpoint()
break
# log the LLM request before sending.
active_schemas = tool_schemas if tool_schemas is not None else TOOL_SCHEMAS
audit.log_llm_request(
model="gemma4",
message_count=len(messages),
token_estimate=context_budget.estimate_total(messages),
has_tools=bool(active_schemas),
)
try:
response = client.chat.completions.create(
model="gemma4",
messages=messages,
tools=active_schemas,
temperature=0.7,
timeout=llm_timeout,
)
except Exception as e:
# surface an LLM-call failure (including timeout)
# to the user and the audit log rather than crashing.
audit.log("llm_call_error", error=str(e))
print(f"\n [llm error] {e}")
print("Assistant: I could not reach the model. "
"Please check the endpoint and retry.")
break
# record usage for cost tracking. Local backends
# (Ollama) may not populate usage; that's fine, the
# tracker simply records zero tokens for the call.
cost_tracker.record_usage(getattr(response, "usage", None))
message = response.choices[0].message
# log the LLM response metadata + content hash for
# forensic replay (without storing every token in the log).
_content_str = message.content or ""
_content_hash = hashlib.sha256(
_content_str.encode("utf-8", errors="replace")
).hexdigest() if _content_str else None
_finish = None
try:
_finish = response.choices[0].finish_reason
except (IndexError, AttributeError):
pass
_usage_dict = None
_raw_usage = getattr(response, "usage", None)
if _raw_usage:
_usage_dict = {
"prompt_tokens": getattr(_raw_usage, "prompt_tokens", None),
"completion_tokens": getattr(_raw_usage, "completion_tokens", None),
}
audit.log_llm_response(
model="gemma4",
finish_reason=_finish,
usage=_usage_dict,
message_hash=_content_hash,
tool_call_count=len(
message.tool_calls) if message.tool_calls else 0,
)
# Always append the assistant turn so the conversation stays intact
messages.append(message)
audit.log_assistant_message(message.content, message.tool_calls)
if memory_store is not None and session_id is not None:
memory_store.save_messages(session_id, messages)
if message.tool_calls:
# The LLM wants to use one or more tools, run them, then loop
handle_tool_calls(
message.tool_calls,
messages,
mode,
working_dir,
tool_registry,
audit,
user_goal=current_user_goal,
iteration_caps=iteration_caps,
tool_timeout=sandbox_exec_timeout,
abort_controller=abort_controller,
file_rollback=file_rollback,
container_name=container_name,
plan_state=plan_state,
)
if memory_store is not None and session_id is not None:
checkpoint()
messages = memory_store.recover(session_id)
elif not message.content or not message.content.strip():
# The model ended its turn with an empty message, most commonly
# happens after a planning-only tool call (scratchpad / todo).
# Nudge it to continue rather than silently stalling.
messages.append({
"role": "user",
"content": "Continue.",
})
else:
# No tool calls: we have the final answer.
# Sanitize model output before rendering: the CLI
# surface is plain text, so no HTML/JS escaping is
# needed here. If a web UI is added in the future, the
# assistant content MUST be passed through html.escape
# (or a template engine's auto-escaping) before being
# inserted into the DOM, never trust model output.
print(f"Assistant: {message.content}")
checkpoint()
break
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Coding agent with Docker-sandboxed tool execution and audit logging."
)
parser.add_argument(
"--mode",
choices=["default", "acceptEdits", "dangerouslySkipPermissions"],
default="default",
help=(
"Permission mode for tool execution. "
"'default': read tools are free, everything else requires approval. "
"'acceptEdits': read + write tools are free when inside the working directory, "
"everything else requires approval. "
"'dangerouslySkipPermissions': all tools run without any prompt."
),
)
parser.add_argument(
"--memory-dir",
default="~/.agent-harness/memory",
help="Directory for durable project/user memory and saved sessions.",
)
parser.add_argument(
"--list-sessions",
action="store_true",
help="List saved sessions in --memory-dir and exit.",
)
parser.add_argument(
"--resume",
metavar="SESSION_ID",
help="Resume a saved session by ID.",
)
management = parser.add_mutually_exclusive_group()
management.add_argument("--fork", metavar="SESSION_ID", help="Fork a session and print its new ID.")
management.add_argument("--rename", nargs=2, metavar=("SESSION_ID", "TITLE"))
management.add_argument("--archive", metavar="SESSION_ID")
management.add_argument("--unarchive", metavar="SESSION_ID")
management.add_argument("--delete", metavar="SESSION_ID")
management.add_argument("--export", metavar="SESSION_ID", help="Print a full JSON session export.")
parser.add_argument("--yes", action="store_true", help="Confirm session deletion.")
parser.add_argument("--include-archived", action="store_true")
parser.add_argument(
"--sandbox-dir",
default=None,
help=(
"Directory bind-mounted into the Docker sandbox container as the "
"project root. Defaults to the current working directory. Action "
"tools can only read and write files inside this mount."
),
)
parser.add_argument(
"--network",
choices=["bridge", "none"],
default="bridge",
help=(
"Network mode for the sandbox container. 'bridge' (default) gives "
"the container normal network access (needed for webfetch and for "
"run_bash commands that download things). 'none' fully isolates the "
"container from the network."
),
)
parser.add_argument(
"--log-dir",
default="./logs",
help=(
"Directory where the JSONL audit log for this session is written. "
"Defaults to './logs'. One file per session."
),
)
parser.add_argument(
"--tools",
default=None,
help=(
"Comma-separated allowlist of tool names to expose to the agent "
"(principle of least privilege). "
"Example: --tools read_file,glob_files,grep,run_bash. "
"Default: all tools are exposed."
),
)
parser.add_argument(
"--tool-timeout",
type=float,
default=120.0,
help=(
"Per-tool-call timeout in seconds. A tool that does not "
"finish in this time is killed and reported to the LLM as a "
"timeout. Default: 120."
),
)
parser.add_argument(
"--llm-timeout",
type=float,
default=120.0,
help=(
"Timeout for each LLM completion call in seconds. "
"Default: 120."
),
)
parser.add_argument(
"--max-turns",
type=int,
default=None,
help=(
"Max LLM turns per user message before the loop stops "
f"and asks for a summary. Default: {
resource_limits.DEFAULT_MAX_TURNS_PER_USER_MSG}."
),
)
parser.add_argument(
"--max-tool-calls",
type=int,
default=None,
help=(
"Max tool calls for the whole session. Default: "
f"{resource_limits.DEFAULT_MAX_TOOL_CALLS_PER_SESSION}."
),
)
parser.add_argument(
"--max-context-tokens",
type=int,
default=None,
help=(
"Token budget for the message history. When exceeded, "
"older messages are trimmed and replaced with a summary. "
f"Default: {resource_limits.DEFAULT_MAX_CONTEXT_TOKENS}."
),
)
parser.add_argument(
"--max-cost-usd",
type=float,
default=None,
help=(
"Cumulative API spend cap in USD. When exceeded, the "
"session aborts gracefully. Default: "
f"{resource_limits.DEFAULT_MAX_COST_USD:.2f} (no effect for "
"local backends that don't report usage)."
),
)
parser.add_argument(
"--approve-plan",
action="store_true",
default=False,
help=(
"Require human approval of the agent's plan before any "
"action tool runs. The agent builds its plan in the "
"scratchpad + todo list; the first time it tries to run an "
"action tool, the user is shown the plan and asked to approve."
),
)
cli_args = parser.parse_args()
if cli_args.resume and any((cli_args.list_sessions, cli_args.fork, cli_args.rename,
cli_args.archive, cli_args.unarchive, cli_args.delete, cli_args.export)):
parser.error("--resume cannot be combined with offline session operations.")
if cli_args.max_context_tokens is not None and cli_args.max_context_tokens <= 0:
parser.error("--max-context-tokens must be positive.")
memory_store = MemoryStore(Path(cli_args.memory_dir))
try:
if cli_args.fork:
print(memory_store.fork(cli_args.fork))
elif cli_args.rename:
memory_store.update(cli_args.rename[0], title=cli_args.rename[1])
elif cli_args.archive or cli_args.unarchive:
memory_store.update(cli_args.archive or cli_args.unarchive, archived=bool(cli_args.archive))
elif cli_args.delete:
if not cli_args.yes:
parser.error("--delete requires --yes; export the session first if needed.")
memory_store.delete(cli_args.delete)
elif cli_args.export:
print(memory_store.export(cli_args.export))
if any((cli_args.fork, cli_args.rename, cli_args.archive, cli_args.unarchive,
cli_args.delete, cli_args.export)):
raise SystemExit(0)
except (ValueError, OSError) as error:
parser.error(str(error))
if cli_args.list_sessions:
sessions = memory_store.list_sessions(include_archived=cli_args.include_archived)
if sessions:
for session in sessions:
print(f"{session['id']} | {session['title']} | {session['status']} | {session['updated_at']} | "
f"{session['summary'][:160] or '(no summary)'}")
else:
print("No saved sessions.")
raise SystemExit(0)
resumed_messages = None
leases = ExitStack()
atexit.register(leases.close)
if cli_args.resume:
try:
leases.enter_context(memory_store.lease(cli_args.resume))
record = memory_store.load_session(cli_args.resume)
if record["archived"]:
raise ValueError("Session is archived; --unarchive it before resuming.")
expected_root = str(Path(cli_args.sandbox_dir).resolve() if cli_args.sandbox_dir else Path.cwd())
if record.get("project_root") and record["project_root"] != expected_root:
raise ValueError("Session belongs to another project; use its original --sandbox-dir.")
resumed_messages = memory_store.recover(cli_args.resume)
except (ValueError, KeyError) as e:
print(f"Could not resume session: {e}")
raise SystemExit(2)
# Principle of least privilege: parse the optional tool allowlist.
allowed_tools: set[str] | None = None
if cli_args.tools:
requested = {t.strip() for t in cli_args.tools.split(",") if t.strip()}
unknown = requested - set(RAW_TOOL_REGISTRY.keys())
if unknown:
print(
f"Error: unknown tool(s) in --tools: {sorted(unknown)}. "
f"Available: {sorted(RAW_TOOL_REGISTRY.keys())}"
)
raise SystemExit(2)
allowed_tools = requested
mode = PermissionMode(cli_args.mode)
working_dir = Path.cwd()
sandbox_root = Path(cli_args.sandbox_dir).resolve(
) if cli_args.sandbox_dir else working_dir
tools_dir = Path(__file__).resolve().parent / "tools"
# Secret & credential management
# Warn about secret-looking env vars on the host (informational).
secret_env = secret_management.scan_environment_for_secrets()
if secret_env:
print(f" [secret-safety] {len(secret_env)
} host env var(s) look like secrets:")
for name, src in secret_env:
print(f" - {name} (from {src})")
print(" These will NOT be passed to the sandbox container "
"(only an allowlist is inherited).")
# Verify no credential mount paths would leak into the container.
cred_mounts = secret_management.check_credential_mounts()
if cred_mounts:
print(
f" [secret-safety] Host credential paths present (NOT mounted into container):")
for p in cred_mounts:
print(f" - {p}")
# Generate per-session credentials. The token is injected into
# the container env; it never appears in the system prompt or tool
# schemas.
session_creds = secret_management.SessionCredentials()
active_session_id = cli_args.resume or session_creds.session_id
leases.enter_context(memory_store.lease(active_session_id))
container_env = session_creds.container_env()
try:
sandbox = DockerSandbox(
project_root=sandbox_root,
tools_dir=tools_dir,
network=cli_args.network,
exec_timeout=cli_args.tool_timeout,
container_env=container_env,
)
except DockerSandboxError as e:
print(f"Could not start sandbox: {e}")
raise SystemExit(1)
tool_registry = build_tool_registry(sandbox, allowed_tools=allowed_tools)
active_schemas = filter_tool_schemas(TOOL_SCHEMAS, allowed_tools)
# resource controls
iteration_caps = IterationCaps(
max_turns_per_user_msg=cli_args.max_turns or resource_limits.DEFAULT_MAX_TURNS_PER_USER_MSG,
max_tool_calls_per_session=cli_args.max_tool_calls or resource_limits.DEFAULT_MAX_TOOL_CALLS_PER_SESSION,
)
context_budget = ContextBudget(
max_tokens=cli_args.max_context_tokens or resource_limits.DEFAULT_MAX_CONTEXT_TOKENS,
)
cost_tracker = CostTracker(
max_cost_usd=cli_args.max_cost_usd if cli_args.max_cost_usd is not None else resource_limits.DEFAULT_MAX_COST_USD,
)
audit = AuditLog(Path(cli_args.log_dir))
audit.log_config(
mode=mode.value,
working_dir=str(working_dir),
sandbox_root=str(sandbox_root),
container=sandbox.container,
network=cli_args.network,
tools_allowed=sorted(allowed_tools) if allowed_tools else "ALL",
tools_host_side=[n for n in tool_registry if n not in ACTION_TOOLS],
tools_container_side=sorted(
n for n in tool_registry if n in ACTION_TOOLS),
tool_timeout_s=cli_args.tool_timeout,
llm_timeout_s=cli_args.llm_timeout,
max_turns_per_user_msg=iteration_caps.max_turns_per_user_msg,
max_tool_calls_per_session=iteration_caps.max_tool_calls_per_session,
max_context_tokens=context_budget.max_tokens,
max_cost_usd=cost_tracker.max_cost_usd,
session_id=active_session_id,
secret_env_count=len(secret_env),
credential_mounts_found=[str(p) for p in cred_mounts],
container_env_allowlist=sorted(container_env.keys()),
)
print(f"Agent started in '{
mode.value}' mode (working dir: {working_dir})")
print(f"Sandbox container: {sandbox.container} (mount: {
sandbox_root}, network: {cli_args.network})")
print(f"Session: {active_session_id}")
print(f"Memory directory: {memory_store.root}")
print(f"Resource limits: turns={iteration_caps.max_turns_per_user_msg} "
f"tool_calls={iteration_caps.max_tool_calls_per_session} "
f"ctx_tokens={context_budget.max_tokens} "
f"cost=${cost_tracker.max_cost_usd:.2f} "
f"tool_timeout={cli_args.tool_timeout}s llm_timeout={cli_args.llm_timeout}s")
if cli_args.approve_plan:
print(f"Plan approval: ENABLED (action tools wait for human approval)")
print(f"Audit log: {audit.path}")
print("Type \\exit to quit, or \\sessions to list saved sessions. (Ctrl-C to abort a session)\n")
# session abort controller: installs SIGINT/SIGTERM handlers.
abort_controller = AbortController()
abort_controller.install_signal_handlers()
# file rollback: snapshots before write/edit for revert on abort.
file_rollback = FileRollback()
client = get_llm_client()
try:
agent_loop(
client, mode, sandbox_root, tool_registry, audit,
tool_schemas=active_schemas,
iteration_caps=iteration_caps,
context_budget=context_budget,
cost_tracker=cost_tracker,
llm_timeout=cli_args.llm_timeout,
tool_timeout=cli_args.tool_timeout,
abort_controller=abort_controller,
file_rollback=file_rollback,
container_name=sandbox.container,
approve_plan=cli_args.approve_plan,
memory_store=memory_store,
session_id=active_session_id,
resumed_messages=resumed_messages,
)
finally:
try:
if memory_store._session_path(active_session_id).exists():
memory_store.ensure_summary(active_session_id)
memory_store.update(active_session_id, status="aborted" if abort_controller.triggered else "saved")
except (ValueError, OSError) as error:
print(f"Could not finalize saved session: {error}")
finally:
leases.close()
abort_controller.remove_signal_handlers()
# offer to revert file changes on exit / abort.
if file_rollback.snapshot_count > 0:
file_rollback.offer_rollback()
file_rollback.cleanup()
session_creds.revoke()
sandbox.close()
audit.close()
print(f"\nSandbox container removed. Audit log written to {
audit.path}")
print(f"Session cost: {cost_tracker.summary()}")