From 775258bcffdf221a37017c10e9daa623c97c5eaf Mon Sep 17 00:00:00 2001 From: 1998x-stack <2484230700@qq.com> Date: Thu, 12 Nov 2026 21:45:14 +0800 Subject: [PATCH] Save local project changes before workspace cleanup --- AGENTS.md | 29 ++-- CONTEXT.md | 10 +- README.md | 16 ++- docs/gotchas.md | 21 ++- pyproject.toml | 2 +- tests/test_agent.py | 169 +++++++++++++++++++++++ tests/test_cli.py | 46 ++++++- tests/test_config.py | 15 +++ tests/test_file_tools_edge.py | 17 +++ tests/test_hitl.py | 40 +++++- tests/test_integration.py | 7 +- tests/test_llm_provider.py | 245 ++++++++++++++++++++++++++++++++++ tests/test_loop.py | 28 ++++ tests/test_persist.py | 6 + tests/test_shell_tools.py | 1 + tests/test_skill_files.py | 36 +++++ tests/test_tools.py | 140 +++++++++++++++++++ tiny_harness/__init__.py | 6 +- tiny_harness/_config.py | 18 ++- tiny_harness/_core.py | 110 +++++++++------ tiny_harness/_events.py | 6 + tiny_harness/_hitl.py | 2 + tiny_harness/_llm.py | 219 +++++++++++++++--------------- tiny_harness/_llm_response.py | 21 +++ tiny_harness/_loop.py | 24 +++- tiny_harness/_messages.py | 3 + tiny_harness/_persist.py | 49 +++++-- tiny_harness/_sse.py | 114 ++++++++++++++++ tiny_harness/_tools.py | 74 ++++++---- tiny_harness/cli.py | 39 ++++-- tiny_harness/tools/files.py | 17 ++- tiny_harness/tools/search.py | 5 +- tiny_harness/tools/shell.py | 7 - tiny_harness/tui.py | 4 +- 34 files changed, 1278 insertions(+), 268 deletions(-) create mode 100644 tiny_harness/_llm_response.py create mode 100644 tiny_harness/_sse.py diff --git a/AGENTS.md b/AGENTS.md index c72bb50..4f9f9c8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -3,22 +3,22 @@ ## Setup ```bash -pip install -e ".[dev]" +python -m pip install -e ".[dev]" ``` ### Running tests ```bash # All unit tests (excludes integration tests needing API keys) -pytest tests/ --ignore=tests/test_integration.py +python -m pytest tests/ --ignore=tests/test_integration.py # Single test file -pytest tests/test_tools.py -v +python -m pytest tests/test_tools.py -v # Single test -pytest tests/test_tools.py::test_execute_tool_timeout -v +python -m pytest tests/test_tools.py::test_execute_tool_timeout -v # Integration tests (needs ANTHROPIC_API_KEY or DEEPSEEK_API_KEY) -pytest tests/test_integration.py -v -m integration +python -m pytest tests/test_integration.py -v -m integration # Run the agent CLI python -m tiny_harness "hello" @@ -35,18 +35,17 @@ ruff check --fix tiny_harness/ tests/ ## Architecture -**Public API** (`tiny_harness/__init__.py`): `Agent`, `Prompt`, `Config`, `ToolDef` +**Public API** (`tiny_harness/__init__.py`): `Agent`, `Prompt`, `Config`, `ToolDef`, `SessionStore` **Module dependency order** (no cycles): ``` -_config.py → _events.py → _guard.py → _llm.py → _messages.py → _tools.py → _loop.py → _core.py - ↑ │ - └─── cli.py, tui.py import only _core.py (never modules directly) ───────────────────┘ +_config.py → _events.py → _guard.py → _sse.py/_llm_response.py → _llm.py + → _messages.py → _hitl.py/_tools.py → _loop.py/_persist.py → _core.py ``` -- `_core.py` (117 lines) — Agent class, the only integrator. Creates all components at `__init__`. Never add new cross-module wiring elsewhere. -- `_loop.py` (109 lines) — AgentLoop state machine. Zero intelligence: calls LLM, executes tools, loops. Never add planning logic here. -- `_llm.py` (341 lines) — All provider code (AnthropicProvider + OpenAIProvider + types). If adding a new provider, match the `LLMProvider` ABC. +- `_core.py` (~300 lines) — Agent class, the only integrator. Creates all components at `__init__`. Never add new cross-module wiring elsewhere. +- `_loop.py` (~130 lines) — AgentLoop state machine. Zero intelligence: calls LLM, executes tools, loops. Never add planning logic here. +- `_llm.py` (~350 lines) — Provider adapters and types; SSE decoding lives in `_sse.py`. New providers must match `LLMProvider`. - `tools/` and `skills/` — each is standalone. `tools/` has handlers (`def my_tool(args: dict) -> str`). `skills/` has `register(agent)` that wires tools + prompt sections together. --- @@ -55,7 +54,7 @@ _config.py → _events.py → _guard.py → _llm.py → _messages.py → _tools. | Term | Meaning | |---|---| -| Agent | The whole runtime. Created, given tasks, discarded. | +| Agent | The whole runtime, scoped to a Session and reusable across multiple `run()` calls. | | Agent Core | Loop + Messages + Config + LLMProvider + empty ToolRegistry. Zero tools built in. | | Tool | External plugin registered into ToolRegistry. Handler signature: `def handler(args: dict) -> str`. | | Prompt | First-class artifact, separate from `Config`. `Prompt("base").append("section")`. | @@ -111,11 +110,15 @@ _config.py → _events.py → _guard.py → _llm.py → _messages.py → _tools. | `tiny_harness/_core.py` | Agent class — orchestrator | | `tiny_harness/_loop.py` | AgentLoop — while-loop state machine | | `tiny_harness/_llm.py` | All LLM providers + types | +| `tiny_harness/_sse.py` | Anthropic/OpenAI SSE decoding and tool-call buffering | +| `tiny_harness/_llm_response.py` | Anthropic response object adapters | | `tiny_harness/_tools.py` | ToolRegistry, ToolExecutor, schema validator, ToolResult | | `tiny_harness/_messages.py` | MessageManager, TokenBudget | | `tiny_harness/_config.py` | AgentConfig, Prompt | | `tiny_harness/_events.py` | StreamEvent, EventBus | | `tiny_harness/_guard.py` | FilesystemGuard | +| `tiny_harness/_hitl.py` | Human approval policy and session approvals | +| `tiny_harness/_persist.py` | Atomic JSONL Session persistence | | `tiny_harness/cli.py` | CLI entry point | | `tiny_harness/tui.py` | Rich TUI mode | | `tiny_harness/tools/files.py` | File tool handlers (7 tools) | diff --git a/CONTEXT.md b/CONTEXT.md index 1c8779b..4c037e7 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -5,7 +5,7 @@ ### Agent The complete runtime system that wraps an LLM, gives it tools, and executes the agent loop until a final answer is produced. The Agent owns: the Loop (while-loop orchestrator), the ToolRegistry (available tools), the LLMProvider (model interface), the MessageManager (context/memory), and the Config (identity + constraints). -Usage: `agent = Agent(config); result = await agent.run("Do X")` +Usage: `agent = Agent(prompt=Prompt("..."), config=config); result = await agent.run("Do X")` The Agent is the only user-facing entry point. Internally, the Loop is the mechanical orchestrator; the Agent is the whole thing. @@ -15,19 +15,19 @@ The Agent is the only user-facing entry point. Internally, the Loop is the mecha The components that ship with the Agent and are essential for it to function: the Loop, MessageManager, Config, LLMProvider, and an empty ToolRegistry. The Agent Core has zero tools by default — even file access must be registered. ### Tool -A capability registered into the Agent's ToolRegistry. Tools are external plugins — they do not ship with the Agent Core. A convenience helper (`register_file_tools(agent)`) provides standard filesystem tools, but the user always controls which tools are available. +A capability registered into the Agent's ToolRegistry. Tools are external plugins — they do not ship with the Agent Core. Built-in Skills such as `agent.load_skill("files")` provide standard tools, but the user controls which Skills are loaded. **Distinction**: The ToolRegistry is part of the Agent Core (it's the container). Individual Tools are not — they are registered into the registry by the user. ### Prompt The agent's personality and operational rules — the system-level instructions that define how the Agent behaves, when to use tools, and what constraints to follow. The Prompt is a first-class design artifact, separate from Config (which is parameters like model, iterations) and from MessageManager (which manages runtime conversation state). -Usage: A Prompt may be a string, a file reference, or a template. It is version-controlled, tested, and swapped independently of the Agent's runtime parameters. +Usage: A Prompt currently wraps text and appended sections. File references and templates are future extensions. **Not to be confused with**: the User Prompt (the task given to the Agent at runtime) or Config (numeric/boolean parameters like `max_iterations`). ### Workspace -The root directory that bounds the Agent's filesystem authority. All relative paths resolve against the Workspace. The FilesystemGuard enforces that no file operation escapes the Workspace. The Workspace is the Agent's "home" — it can read, write, and organize within it, but cannot access anything outside. +The root directory that bounds filesystem-aware Tools. All relative path fields resolve against the Workspace, and FilesystemGuard prevents those operations from escaping it. Workspace is not an OS sandbox: a deliberately enabled `dangerous` shell Tool still runs with the host process's authority and is controlled through HITL approval. Usage: Set once at Agent creation. The Agent treats the Workspace as `/` — it doesn't know (and shouldn't care) about the absolute path on the host filesystem. @@ -36,7 +36,7 @@ Usage: Set once at Agent creation. The Agent treats the Workspace as `/` — it ### Conversation The full ordered history of messages exchanged during an Agent run: the system-level Prompt, the user's task, the LLM's responses (text + tool call requests), and the tool results. The Conversation IS the Agent's memory — everything the LLM knows about the current interaction lives in this sequence. The MessageManager owns the Conversation; all other components read from it or append to it. -Usage: The Conversation starts with Prompt + user task and grows monotonically as the Loop executes. When the Agent completes, the Conversation is discarded (single-use Agent). Context management (truncation, compaction) operates on the Conversation. +Usage: The Conversation starts with Prompt + user task and grows across `run()` calls in the same Session. `clear()` resets it to the current system Prompt. ### Session A continuous interaction between the user and Agent that spans multiple user prompts. Within a Session, the Conversation persists — each new user prompt appends to the same Conversation, giving the Agent memory of what happened earlier in the Session. The Agent is Session-scoped: created when the Session starts, serves multiple user prompts, and discarded when the Session ends. diff --git a/README.md b/README.md index f6e76f6..51b45b5 100644 --- a/README.md +++ b/README.md @@ -3,11 +3,11 @@ **Minimal AI agent harness — wrap any LLM with tools and a streaming CLI.** [![Python](https://img.shields.io/badge/python-3.11%2B-blue)](https://www.python.org/) -[![Tests](https://img.shields.io/badge/tests-195%20passed-brightgreen)](https://github.com/1998x-stack/tiny-harness/actions) -[![Lines](https://img.shields.io/badge/code-1%2C350%20lines-lightgrey)](.) +[![Tests](https://img.shields.io/badge/tests-271%20passed-brightgreen)](.) +[![Lines](https://img.shields.io/badge/code-%7E2%2C000%20lines-lightgrey)](.) [![License](https://img.shields.io/badge/license-MIT-green)](LICENSE) -`tiny-harness` is a ~1,100-line Python package that turns an LLM into an AI agent — giving it tools, a conversation loop, streaming events, and a CLI session. One dependency. One import. +`tiny-harness` is a compact Python package that turns an LLM into an AI agent—giving it tools, a conversation loop, streaming events, and a CLI session. ```bash pip install tiny-harness @@ -46,8 +46,8 @@ tiny-harness "Create hello.py" - **CLI** — session REPL with streaming output, one-shot mode for scripts. Files skill loaded by default. - **Rich TUI** — optional `--tui` mode with panels, markdown rendering, color-coded messages (`pip install tiny-harness[tui]`) - **Persistence** — JSONL session history with `/save` and `/history` commands -- **Filesystem guard** — workspace boundary enforcement, path traversal protection -- **1,125 lines** — readable top-to-bottom, CleanRL-inspired code style +- **Filesystem guard** — canonical Workspace paths and traversal protection for filesystem-aware Tools +- **Compact modules** — source modules stay below 350 lines and are readable top-to-bottom ## Quick Start @@ -142,15 +142,17 @@ Agent.run(prompt) ``` tiny_harness/ -├── __init__.py # Public API: Agent, Prompt, Config, ToolDef +├── __init__.py # Public API: Agent, Prompt, Config, ToolDef, SessionStore ├── _core.py # Agent class — orchestrator, session management ├── _loop.py # AgentLoop — state machine, error budget, loop detection ├── _llm.py # LLMProvider ABC + AnthropicProvider + OpenAIProvider +├── _sse.py # Provider-native SSE delta decoding ├── _messages.py # MessageManager — conversation, token counting ├── _tools.py # ToolRegistry + ToolExecutor + schema validator ├── _config.py # AgentConfig, Prompt dataclasses ├── _events.py # StreamEvent types + EventBus ├── _guard.py # FilesystemGuard — workspace boundaries +├── _persist.py # Atomic JSONL Session snapshots ├── cli.py # CLI entry point ├── tools/files.py # File tool handlers (read, write, list, find, delete, mkdir, move) └── skills/files.py # register(agent) → file tools + prompt instructions @@ -201,7 +203,7 @@ agent.tools.register_from_def( description="Get current weather for a city.", parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, ), - handler=lambda city: f"Sunny, 22°C in {city}", + handler=lambda args: f"Sunny, 22°C in {args['city']}", ) # Load built-in skills diff --git a/docs/gotchas.md b/docs/gotchas.md index 41a2c21..200cc9f 100644 --- a/docs/gotchas.md +++ b/docs/gotchas.md @@ -47,11 +47,11 @@ This uses `difflib.get_close_matches` with a 0.6 cutoff. Names very different fr | `mutation` | Checks path is writable within workspace | | `destructive` | Strictest — checks path for delete operations | -The guard extracts the path from `args["path"]`, `args["source"]`, or `args["destination"]`. If your tool uses a different key, the guard won't apply. +The guard validates and canonicalizes every conventional path field present: `path`, `source`, `destination`, and `cwd`. If a custom Tool uses a different key, it must enforce its own path policy. ### Tool timeouts are per-call, not cumulative -Each tool call has its own timeout (default 30s). A task requiring 10 tool calls can run for 10 × 30 = 300 seconds total. The iteration cap (default 25) is the primary time-gate. +Each Tool call has its own timeout (default 30s). Synchronous handlers run in a worker thread so they do not block the Agent Loop. Timing out stops waiting; Python cannot forcibly terminate arbitrary work already running in that thread. --- @@ -69,7 +69,7 @@ Stream failed: 400 'The `reasoning_content` in the thinking mode must be passed ### DeepSeek streaming tool calls arrive as deltas -Unlike Anthropic where tool calls arrive complete in a single event, OpenAI-compatible APIs (DeepSeek, OpenAI) stream tool call arguments as JSON fragments. The `OpenAIProvider.generate_stream()` buffers these deltas and emits the complete `ToolCallRequest` at the end. +Anthropic and OpenAI-compatible APIs stream tool-call arguments as fragments. Each Provider Adapter buffers its native deltas and emits the same complete `ToolCallRequest` at `tool_call_end`. ### Anthropic uses separate `system` parameter @@ -185,7 +185,7 @@ Negative offsets are clamped to 1. Offsets exceeding file length return an error ### Shell cwd validated against workspace -The `run_command` tool validates its `cwd` parameter against the workspace boundary. This prevents LLM-injected `cwd="/etc"` from bypassing the FilesystemGuard. +The ToolExecutor validates and canonicalizes `run_command.cwd` against the configured Workspace before invoking the shell handler. ### find_files pattern validated against path escapes @@ -359,7 +359,7 @@ Shell builtins (`exit`, `cd`, `export`) work via shell=True fallback since they ### `/save` now dumps the full conversation to disk -The `/save` command calls `_dump_conversation()` which serializes all current messages (user prompts, assistant responses, tool calls, tool results) to a JSONL file. Previously, `/save` only initialized the session without writing existing messages — the file would be created empty. +The `/save` command calls the public `agent.save_session()` Interface. It atomically replaces the JSONL snapshot with the current Conversation, so repeated saves are idempotent. ### `start_session()` is idempotent @@ -371,24 +371,24 @@ assert sid1 == sid2 Calling `start_session()` multiple times returns the same session ID and does NOT reset `_chat_id`. -### `_save_turn()` is a no-op until `start_session()` is called +### Active Sessions are snapshotted after each run -Both `agent.run()` and `agent.run_stream()` call `_save_turn()` internally, but it returns early if `_session_id` is None. The CLI/TUI `/save` commands call `start_session()` first. In the Python API, call it explicitly: +Once a Session has started, `agent.run()` and `agent.run_stream()` snapshot the complete Conversation after each turn, including Tool calls and results. You can also save explicitly: ```python agent.start_session() await agent.run("hello") # turn saved await agent.run("world") # turn saved -agent._dump_conversation() # write all to file +session_id, turns = agent.save_session() ``` ### Tool call format differs between memory and storage -Messages store tool calls as `{"function": {"name": ..., "arguments": "..."}}` (LLM format). The SessionStore expects `{"name": ..., "arguments": {...}}` (plain dicts). `_dump_conversation()` handles this conversion internally. +Messages store tool calls in the common LLM format; Session persistence stores structured arguments and preserves Tool call ids. `save_session()` handles this conversion internally. ### The loop now adds assistant messages even without tool calls -Previously, when the LLM returned text without calling tools, the assistant response was returned as a string but never added to `_messages`. It's now added via `add_assistant()`, so `_dump_conversation()` sees the complete conversation. +Final and degraded answers are appended to the Conversation before it is persisted. --- @@ -437,4 +437,3 @@ config = Config( max_tokens=4096, # limit response size (default: 16384) ) ``` - diff --git a/pyproject.toml b/pyproject.toml index 80b5d62..71d6983 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,7 @@ requires-python = ">=3.11" dependencies = ["httpx>=0.27.0"] [project.optional-dependencies] -dev = ["pytest>=8.0.0", "pytest-asyncio>=0.23.0"] +dev = ["pytest>=8.0.0", "pytest-asyncio>=0.23.0", "rich>=13.0.0", "ruff>=0.8.0"] tui = ["rich>=13.0.0"] [project.scripts] diff --git a/tests/test_agent.py b/tests/test_agent.py index ca15539..0c47236 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -59,6 +59,15 @@ def test_agent_config_defaults(): assert config.max_errors == 10 +def test_agent_deepseek_default_endpoint(): + agent = Agent( + prompt=Prompt("test"), + config=Config(model="deepseek-chat", api_key="k", workspace="/tmp"), + ) + + assert agent._llm_provider._base_url == "https://api.deepseek.com/v1" + + def test_agent_config_unknown_provider(): prompt = Prompt("test") config = Config(model="m", api_key="k", workspace=".", provider="unknown") @@ -88,6 +97,52 @@ async def generate_stream(self, messages, tools=None): assert "Hello" in "".join(events) +@pytest.mark.asyncio +async def test_agent_run_stream_releases_internal_subscriber(): + class FakeLLM: + async def generate_stream(self, messages, tools=None): + from tiny_harness._llm import LLMStreamChunk + yield LLMStreamChunk(type="text_delta", content="ok") + + agent = Agent( + prompt=Prompt("test"), + config=Config(model="test", api_key="k", workspace="/tmp"), + ) + agent._llm_provider = FakeLLM() + + for prompt in ("one", "two"): + async for _ in agent.run_stream(prompt): + pass + + assert len(agent.events._handlers) == 0 + + +@pytest.mark.asyncio +async def test_closing_agent_stream_cancels_background_run(): + import asyncio + + class SlowLLM: + async def generate_stream(self, messages, tools=None): + from tiny_harness._llm import LLMStreamChunk + await asyncio.sleep(0.05) + yield LLMStreamChunk(type="text_delta", content="late") + + agent = Agent( + prompt=Prompt("test"), + config=Config(model="test", api_key="k", workspace="/tmp"), + ) + agent._llm_provider = SlowLLM() + stream = agent.run_stream("stop") + + first = await anext(stream) + assert first.type == "iteration" + await stream.aclose() + await asyncio.sleep(0.1) + + assert agent._running is False + assert [m for m in agent._messages.to_list() if m["role"] == "assistant"] == [] + + @pytest.mark.asyncio async def test_agent_concurrent_run_raises(): prompt = Prompt("You are helpful.") @@ -198,6 +253,40 @@ def test_agent_dump_conversation_writes_user_messages(): assert records[1]["content"] == "World" +def test_agent_dump_conversation_is_idempotent(tmp_path): + from tiny_harness._persist import SessionStore + + store = SessionStore(base_dir=str(tmp_path)) + agent = Agent( + prompt=Prompt("test"), + config=Config(model="t", api_key="k", workspace="/tmp"), + store=store, + ) + agent._messages.add_user("Hello") + agent.start_session() + + agent._dump_conversation() + agent._dump_conversation() + + assert len(store.load_session(agent.session_id)) == 1 + + +def test_agent_save_session_exposes_persistence_interface(tmp_path): + from tiny_harness._persist import SessionStore + + agent = Agent( + prompt=Prompt("test"), + config=Config(model="t", api_key="k", workspace="/tmp"), + store=SessionStore(base_dir=str(tmp_path)), + ) + agent._messages.add_user("Hello") + + session_id, turns = agent.save_session() + + assert turns == 1 + assert agent.store.load_session(session_id)[0]["content"] == "Hello" + + def test_agent_dump_conversation_with_tool_calls(): from tiny_harness._llm import ToolCallRequest agent = Agent(prompt=Prompt("test"), config=Config(model="t", api_key="k", workspace="/tmp")) @@ -306,3 +395,83 @@ async def generate_stream(self, messages, tools=None): assert "user" in roles user_msgs = [m for m in msgs if m["role"] == "user"] assert any("First question" in m.get("content", "") for m in user_msgs) + + +def test_agent_session_round_trips_tool_call_protocol(tmp_path): + import json + from tiny_harness._llm import ToolCallRequest + from tiny_harness._persist import SessionStore + + store = SessionStore(base_dir=str(tmp_path)) + config = Config(model="t", api_key="k", workspace="/tmp") + agent = Agent(prompt=Prompt("test"), config=config, store=store) + agent._messages.add_user("read") + agent._messages.add_assistant( + "checking", + [ToolCallRequest(id="call_original", name="read_file", arguments={"path": "x.txt"})], + ) + full_result = "x" * 500 + agent._messages.add_tool_result("call_original", full_result) + agent.start_session() + session_id = agent.session_id + agent._dump_conversation() + + restored = Agent(prompt=Prompt("test"), config=config, store=store) + restored.resume_session(session_id) + messages = restored._messages.to_list() + assistant = next(m for m in messages if m["role"] == "assistant") + tool_result = next(m for m in messages if m["role"] == "tool") + + assert assistant["tool_calls"][0]["id"] == "call_original" + assert json.loads(assistant["tool_calls"][0]["function"]["arguments"]) == {"path": "x.txt"} + assert tool_result == { + "role": "tool", + "tool_call_id": "call_original", + "content": full_result, + } + + +@pytest.mark.asyncio +async def test_active_session_persists_complete_tool_turn(tmp_path): + from tiny_harness._llm import LLMStreamChunk, ToolCallRequest + from tiny_harness._persist import SessionStore + + class ToolLLM: + def __init__(self): + self.calls = 0 + + async def generate_stream(self, messages, tools=None): + self.calls += 1 + if self.calls == 1: + yield LLMStreamChunk( + type="tool_call_end", + tool_call=ToolCallRequest(id="call_1", name="echo", arguments={"message": "hi"}), + ) + else: + yield LLMStreamChunk(type="text_delta", content="done") + + store = SessionStore(base_dir=str(tmp_path)) + agent = Agent( + prompt=Prompt("test"), + config=Config(model="t", api_key="k", workspace=str(tmp_path), no_hitl=True), + store=store, + ) + agent._llm_provider = ToolLLM() + agent.tools.register_tool( + name="echo", + description="Echo", + parameters={ + "type": "object", + "properties": {"message": {"type": "string"}}, + "required": ["message"], + }, + handler=lambda args: args["message"], + ) + agent.start_session() + + await agent.run("use echo") + + records = store.load_session(agent.session_id) + tool_turn = next(record for record in records if record.get("tool_calls")) + assert tool_turn["tool_calls"][0]["id"] == "call_1" + assert tool_turn["tool_results"][0] == {"id": "call_1", "content": "hi"} diff --git a/tests/test_cli.py b/tests/test_cli.py index 3153858..3690deb 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -6,6 +6,10 @@ from tiny_harness.cli import parse_args, _run_one_shot +def _close_coroutine(coroutine): + coroutine.close() + + def test_parse_args_one_shot(): with patch.object(sys, 'argv', ['tiny-harness', 'Hello']): args = parse_args() @@ -33,7 +37,7 @@ def test_parse_args_with_options(): assert args.max_iterations == 10 assert args.skills == 'files' assert args.provider == 'deepseek' - assert args.api_base_url == 'https://api.deepseek.com/v1' + assert args.api_base_url is None def test_parse_args_with_provider(): @@ -43,6 +47,15 @@ def test_parse_args_with_provider(): assert args.api_base_url == 'https://api.deepseek.com/v1' +def test_parse_args_derives_provider_specific_defaults(): + with patch.object(sys, "argv", ["tiny-harness", "hi", "--provider", "anthropic"]): + args = parse_args() + + assert args.model == "claude-sonnet-4-20250514" + assert args.api_key_env == "ANTHROPIC_API_KEY" + assert args.api_base_url is None + + def test_parse_args_tui_flag(): with patch.object(sys, 'argv', ['tiny-harness', '--tui']): args = parse_args() @@ -105,7 +118,7 @@ async def generate_stream(self, messages, tools=None): def test_main_routes_to_tui(): with patch.object(sys, 'argv', ['tiny-harness', '--tui', '--api-key-env', 'FAKE_KEY']): with patch.dict('os.environ', {'FAKE_KEY': 'test-key'}): - with patch('tiny_harness.cli.asyncio.run') as mock_run: + with patch('tiny_harness.cli.asyncio.run', side_effect=_close_coroutine) as mock_run: from tiny_harness.cli import main main() assert mock_run.called @@ -114,7 +127,34 @@ def test_main_routes_to_tui(): def test_main_routes_to_session(): with patch.object(sys, 'argv', ['tiny-harness', '--api-key-env', 'FAKE_KEY']): with patch.dict('os.environ', {'FAKE_KEY': 'test-key'}): - with patch('tiny_harness.cli.asyncio.run') as mock_run: + with patch('tiny_harness.cli.asyncio.run', side_effect=_close_coroutine) as mock_run: from tiny_harness.cli import main main() assert mock_run.called + + +def test_main_allows_explicitly_empty_skills(): + loaded = [] + + class FakeAgent: + def __init__(self, **kwargs): + pass + + def set_approval_handler(self, handler): + pass + + def load_skill(self, name): + loaded.append(name) + + def close_coroutine(coroutine): + coroutine.close() + + with patch.object(sys, "argv", [ + "tiny-harness", "hello", "--skills", "", "--api-key-env", "FAKE_KEY", + ]), patch.dict("os.environ", {"FAKE_KEY": "test-key"}), \ + patch("tiny_harness.Agent", FakeAgent), \ + patch("tiny_harness.cli.asyncio.run", side_effect=close_coroutine): + from tiny_harness.cli import main + main() + + assert loaded == [] diff --git a/tests/test_config.py b/tests/test_config.py index 31633a8..dd993e1 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,4 +1,5 @@ # tests/test_config.py +import pytest from tiny_harness._config import AgentConfig, Prompt @@ -46,3 +47,17 @@ def test_config_provider_default(): config = AgentConfig(model="m", api_key="k", workspace=".") assert config.provider == "deepseek" assert config.api_base_url is None + + +def test_agent_config_hides_api_key_from_repr(): + config = AgentConfig(model="m", api_key="super-secret", workspace=".") + + assert "super-secret" not in repr(config) + + +@pytest.mark.parametrize("field", ["max_iterations", "max_errors", "timeout_ms", "max_tokens"]) +def test_agent_config_rejects_non_positive_limits(field): + values = {field: 0} + + with pytest.raises(ValueError, match=field): + AgentConfig(model="m", api_key="k", workspace=".", **values) diff --git a/tests/test_file_tools_edge.py b/tests/test_file_tools_edge.py index 265fe35..912a822 100644 --- a/tests/test_file_tools_edge.py +++ b/tests/test_file_tools_edge.py @@ -58,6 +58,17 @@ def test_read_file_limit_zero(tmpdir): assert "limit must be positive" in result +def test_read_file_describes_binary_content(tmpdir): + path = os.path.join(tmpdir, "image.bin") + with open(path, "wb") as file: + file.write(b"\x00\x01\x02") + + result = read_file({"path": path}) + + assert "binary" in result.lower() + assert "3B" in result + + def test_list_directory_recursive_with_hidden(tmpdir): hidden = os.path.join(tmpdir, ".hidden_dir") os.makedirs(hidden) @@ -71,6 +82,12 @@ def test_find_files_recursive(): assert ".py" in result.lower() or "No files" in result +def test_find_files_rejects_absolute_glob(tmpdir): + result = find_files({"pattern": "/etc/*", "path": tmpdir}) + + assert result.startswith("Error:") + + def test_delete_readonly_file(tmpdir): path = os.path.join(tmpdir, "readonly.txt") write_file({"path": path, "content": "x"}) diff --git a/tests/test_hitl.py b/tests/test_hitl.py index 76c8000..0031eea 100644 --- a/tests/test_hitl.py +++ b/tests/test_hitl.py @@ -32,16 +32,16 @@ def test_session_approvals_cascade_upward(): def test_session_approvals_destructive_not_cascaded(): sa = SessionApprovals() sa.approve_for_session("destructive") - assert sa.is_approved("mutation") is True - assert sa.is_approved("destructive") is True + assert sa.is_approved("mutation") is False + assert sa.is_approved("destructive") is False assert sa.is_approved("dangerous") is False -def test_session_approvals_dangerous_covers_all(): +def test_session_approvals_dangerous_is_not_cached(): sa = SessionApprovals() sa.approve_for_session("dangerous") for level in ["safe", "read_only", "mutation", "destructive", "dangerous"]: - assert sa.is_approved(level) is True + assert sa.is_approved(level) is False def test_session_approvals_clear(): @@ -146,6 +146,24 @@ async def handler(req): assert session.is_approved("mutation") is True +@pytest.mark.asyncio +async def test_dangerous_tools_cannot_be_approved_for_session(): + calls = 0 + + async def handler(req): + nonlocal calls + calls += 1 + return ApprovalDecision(approved=True, session_approved=True) + + gate = ApprovalGate(handler=handler) + + await gate.check("run_command", {"command": "pwd"}, "dangerous") + await gate.check("run_command", {"command": "pwd"}, "dangerous") + + assert calls == 2 + assert gate.session.is_approved("dangerous") is False + + @pytest.mark.asyncio async def test_check_handler_receives_request_data(): async def handler(req): @@ -367,3 +385,17 @@ async def modify_handler(req): result = await executor.execute("write_file", {"path": "/tmp/original.txt"}, "tc8") assert result.success is False assert "outside" in result.content.lower() or "denied" in result.content.lower() + + +@pytest.mark.asyncio +async def test_executor_modified_args_are_revalidated(registry): + async def modify_handler(req): + return ApprovalDecision(approved=True, modified_args={"path": 123}) + + gate = ApprovalGate(handler=modify_handler) + executor = ToolExecutor(registry, FilesystemGuard("/tmp"), approval_gate=gate) + + result = await executor.execute("write_file", {"path": "/tmp/original.txt"}, "tc9") + + assert result.success is False + assert "invalid arguments" in result.content.lower() diff --git a/tests/test_integration.py b/tests/test_integration.py index a230a38..03fada6 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -12,7 +12,12 @@ async def test_agent_simple_prompt_no_tools(): pytest.skip("ANTHROPIC_API_KEY not set") prompt = Prompt("You are a helpful assistant. Be concise.") - config = Config(model="claude-sonnet-4-20250514", api_key=api_key, workspace=os.getcwd()) + config = Config( + model="claude-sonnet-4-20250514", + api_key=api_key, + workspace=os.getcwd(), + provider="anthropic", + ) agent = Agent(prompt=prompt, config=config) result = await agent.run("What is 2+2?") assert "4" in result diff --git a/tests/test_llm_provider.py b/tests/test_llm_provider.py index e1aca35..b126acf 100644 --- a/tests/test_llm_provider.py +++ b/tests/test_llm_provider.py @@ -1,5 +1,6 @@ # tests/test_llm_provider.py import json +import pytest from tiny_harness._llm import AnthropicProvider, LLMRetryConfig, OpenAIProvider @@ -16,6 +17,40 @@ def test_convert_messages_extracts_system(): assert result[1]["role"] == "assistant" +def test_anthropic_converts_tool_messages_to_content_blocks(): + provider = AnthropicProvider(api_key="test-key", model="claude-test") + messages = [ + {"role": "system", "content": "system"}, + {"role": "user", "content": "read"}, + { + "role": "assistant", + "content": "checking", + "tool_calls": [ + { + "id": "toolu_1", + "type": "function", + "function": {"name": "read_file", "arguments": '{"path":"x.txt"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "toolu_1", "content": "contents"}, + ] + + converted = provider._convert_messages(messages) + + assert converted[1] == { + "role": "assistant", + "content": [ + {"type": "text", "text": "checking"}, + {"type": "tool_use", "id": "toolu_1", "name": "read_file", "input": {"path": "x.txt"}}, + ], + } + assert converted[2] == { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "contents"}], + } + + def test_extract_system_returns_system_content(): provider = AnthropicProvider(api_key="test-key", model="claude-test") messages = [ @@ -119,6 +154,162 @@ def test_retry_config_defaults(): assert config.max_delay == 60.0 +@pytest.mark.asyncio +async def test_anthropic_stream_pairs_event_and_data_lines(monkeypatch): + import httpx + + lines = [ + "event: message_start", + 'data: {"type":"message_start","message":{"usage":{"input_tokens":3}}}', + "", + "event: content_block_delta", + 'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}', + "", + ] + + class FakeResponse: + status_code = 200 + + async def aiter_lines(self): + for line in lines: + yield line + + class AsyncContext: + def __init__(self, value): + self.value = value + + async def __aenter__(self): + return self.value + + async def __aexit__(self, *exc): + return False + + class FakeClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + def stream(self, *args, **kwargs): + return AsyncContext(FakeResponse()) + + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: FakeClient()) + provider = AnthropicProvider(api_key="test", model="claude-test") + + chunks = [chunk async for chunk in provider.generate_stream([{"role": "user", "content": "hi"}])] + + assert [(c.type, c.content) for c in chunks] == [ + ("usage", '{"input": 3}'), + ("text_delta", "hello"), + ] + + +@pytest.mark.asyncio +async def test_anthropic_stream_emits_complete_tool_call(monkeypatch): + import httpx + + lines = [ + "event: content_block_start", + 'data: {"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_1","name":"read_file"}}', + "event: content_block_delta", + 'data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\\"path\\":"}}', + "event: content_block_delta", + 'data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"\\"x.txt\\"}"}}', + "event: content_block_stop", + 'data: {"type":"content_block_stop","index":1}', + ] + + class FakeResponse: + status_code = 200 + + async def aiter_lines(self): + for line in lines: + yield line + + class AsyncContext: + async def __aenter__(self): + return FakeResponse() + + async def __aexit__(self, *exc): + return False + + class FakeClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + def stream(self, *args, **kwargs): + return AsyncContext() + + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: FakeClient()) + provider = AnthropicProvider(api_key="test", model="claude-test") + + chunks = [chunk async for chunk in provider.generate_stream([{"role": "user", "content": "read"}])] + completed = [chunk.tool_call for chunk in chunks if chunk.type == "tool_call_end"] + + assert len(completed) == 1 + assert completed[0].id == "toolu_1" + assert completed[0].name == "read_file" + assert completed[0].arguments == {"path": "x.txt"} + + +@pytest.mark.asyncio +async def test_anthropic_stream_retries_transient_status(monkeypatch): + import httpx + + statuses = [429, 200] + + class FakeResponse: + def __init__(self, status_code): + self.status_code = status_code + + async def aread(self): + return b"busy" + + async def aiter_lines(self): + yield "event: content_block_delta" + yield 'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}' + + class AsyncContext: + def __init__(self, response): + self.response = response + + async def __aenter__(self): + return self.response + + async def __aexit__(self, *exc): + return False + + class FakeClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + def stream(self, *args, **kwargs): + return AsyncContext(FakeResponse(statuses.pop(0))) + + async def no_sleep(delay): + pass + + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: FakeClient()) + monkeypatch.setattr("tiny_harness._llm.asyncio.sleep", no_sleep) + provider = AnthropicProvider( + api_key="test", + model="claude-test", + retry_config=LLMRetryConfig(max_retries=1, base_delay=0), + ) + + chunks = [chunk async for chunk in provider.generate_stream([{"role": "user", "content": "hi"}])] + + assert [chunk.content for chunk in chunks if chunk.type == "text_delta"] == ["ok"] + assert statuses == [] + + def test_openai_provider_convert_messages_passthrough(): provider = OpenAIProvider(api_key="test-key", model="gpt-4o") msgs = [{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}] @@ -163,3 +354,57 @@ def test_openai_provider_parse_response_with_tool_calls(): def test_openai_provider_convert_tools_none(): provider = OpenAIProvider(api_key="test-key", model="gpt-4o") assert provider._convert_tools(None) is None + + +@pytest.mark.asyncio +async def test_openai_stream_retries_transient_status(monkeypatch): + import httpx + + statuses = [503, 200] + + class FakeResponse: + def __init__(self, status_code): + self.status_code = status_code + + async def aread(self): + return b"busy" + + async def aiter_lines(self): + yield 'data: {"choices":[{"delta":{"content":"ok"},"finish_reason":"stop"}]}' + yield "data: [DONE]" + + class AsyncContext: + def __init__(self, response): + self.response = response + + async def __aenter__(self): + return self.response + + async def __aexit__(self, *exc): + return False + + class FakeClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + def stream(self, *args, **kwargs): + return AsyncContext(FakeResponse(statuses.pop(0))) + + async def no_sleep(delay): + pass + + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: FakeClient()) + monkeypatch.setattr("tiny_harness._llm.asyncio.sleep", no_sleep) + provider = OpenAIProvider( + api_key="test", + model="gpt-test", + retry_config=LLMRetryConfig(max_retries=1, base_delay=0), + ) + + chunks = [chunk async for chunk in provider.generate_stream([{"role": "user", "content": "hi"}])] + + assert [chunk.content for chunk in chunks if chunk.type == "text_delta"] == ["ok"] + assert statuses == [] diff --git a/tests/test_loop.py b/tests/test_loop.py index f7e1241..8b49eaa 100644 --- a/tests/test_loop.py +++ b/tests/test_loop.py @@ -149,6 +149,7 @@ async def test_loop_max_iterations_triggers_degraded_finish(): loop = AgentLoop(config, messages, provider, executor, events) result = await loop.run("do stuff") assert "final_after_degrade" in result + assert messages.to_list()[-1] == {"role": "assistant", "content": "final_after_degrade"} @pytest.mark.asyncio @@ -170,6 +171,31 @@ async def test_loop_error_budget_exhausted_degrade(): assert "final_after_degrade" in result +@pytest.mark.asyncio +async def test_error_budget_degrade_completes_all_declared_tool_calls(): + config = AgentConfig( + model="t", + api_key="k", + workspace="/tmp", + max_errors=1, + max_consecutive_errors=1, + ) + messages = MessageManager(Prompt("Be helpful.")) + first = ToolCallRequest(id="t1", name="failing", arguments={}) + second = ToolCallRequest(id="t2", name="never_run", arguments={}) + provider = FakeProvider(responses=[[ + LLMStreamChunk(type="tool_call_end", tool_call=first), + LLMStreamChunk(type="tool_call_end", tool_call=second), + ]]) + executor = FakeExecutor(results=[ToolResult.error("t1", "failed")]) + loop = AgentLoop(config, messages, provider, executor, EventBus()) + + await loop.run("do") + + tool_result_ids = [m["tool_call_id"] for m in messages.to_list() if m["role"] == "tool"] + assert tool_result_ids == ["t1", "t2"] + + @pytest.mark.asyncio async def test_loop_detector_intervenes(): config = AgentConfig(model="t", api_key="k", workspace="/tmp", max_iterations=10, max_errors=20, max_consecutive_errors=20) @@ -205,6 +231,8 @@ async def generate(self, messages, tools=None): loop = AgentLoop(config, messages, ErrorProvider(), FakeExecutor(), events) result = await loop.run("hi") assert "LLM error" in result or "stopped" in result.lower() + assert messages.to_list()[-1]["role"] == "assistant" + assert messages.to_list()[-1]["content"] == result @pytest.mark.asyncio diff --git a/tests/test_persist.py b/tests/test_persist.py index c78e6ab..07bed44 100644 --- a/tests/test_persist.py +++ b/tests/test_persist.py @@ -3,6 +3,12 @@ from tiny_harness._persist import SessionStore +def test_session_store_is_publicly_exported(): + from tiny_harness import SessionStore as PublicSessionStore + + assert PublicSessionStore is SessionStore + + def test_new_session_creates_id(): store = SessionStore() sid = store.new_session() diff --git a/tests/test_shell_tools.py b/tests/test_shell_tools.py index 01645f2..25415b3 100644 --- a/tests/test_shell_tools.py +++ b/tests/test_shell_tools.py @@ -14,6 +14,7 @@ def test_basic_echo(): def test_cwd(): result = run_command({"command": "pwd", "cwd": "/tmp"}) assert "/tmp" in result + assert not result.startswith("Error:") def test_exit_code_nonzero(): diff --git a/tests/test_skill_files.py b/tests/test_skill_files.py index b1c00ee..40d4771 100644 --- a/tests/test_skill_files.py +++ b/tests/test_skill_files.py @@ -25,6 +25,18 @@ def test_load_files_skill_appends_prompt(): assert len(updated) > len(original) +def test_load_files_skill_updates_conversation_system_prompt(): + agent = Agent( + prompt=Prompt("You are helpful."), + config=Config(model="test", api_key="sk-test", workspace="/tmp"), + ) + + agent.load_skill("files") + + system_message = agent._messages.to_list()[0] + assert "File Operations" in system_message["content"] + + def test_load_skill_not_found(): prompt = Prompt("test") config = Config(model="t", api_key="k", workspace="/tmp") @@ -33,6 +45,18 @@ def test_load_skill_not_found(): agent.load_skill("nonexistent_skill_xyz") +def test_load_skill_preserves_internal_import_error(tmp_path, monkeypatch): + (tmp_path / "broken_skill.py").write_text("import dependency_that_is_missing\n") + monkeypatch.syspath_prepend(str(tmp_path)) + agent = Agent( + prompt=Prompt("test"), + config=Config(model="t", api_key="k", workspace="/tmp"), + ) + + with pytest.raises(ModuleNotFoundError, match="dependency_that_is_missing"): + agent.load_skill("broken_skill") + + def test_load_skill_duplicate_does_not_reregister(): prompt = Prompt("test") config = Config(model="t", api_key="k", workspace="/tmp") @@ -43,6 +67,18 @@ def test_load_skill_duplicate_does_not_reregister(): assert len(agent.tools.names()) == count_before +def test_load_skill_aliases_are_idempotent(): + agent = Agent( + prompt=Prompt("test"), + config=Config(model="t", api_key="k", workspace="/tmp"), + ) + + agent.load_skill("files") + agent.load_skill("tiny_harness.skills.files") + + assert agent._messages.to_list()[0]["content"].count("## File Operations") == 1 + + def test_load_skill_prompt_accumulates(): prompt = Prompt("base.") config = Config(model="t", api_key="k", workspace="/tmp") diff --git a/tests/test_tools.py b/tests/test_tools.py index 09f9710..9d7fd28 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -41,6 +41,11 @@ def test_names(registry): assert "failing" in names +def test_tool_definition_rejects_unknown_risk_level(): + with pytest.raises(ValueError, match="risk level"): + ToolDef(name="unsafe", description="", parameters={}, risk_level="typo") + + @pytest.mark.asyncio async def test_execute_returns_success_result(registry): guard = FilesystemGuard("/tmp") @@ -89,6 +94,63 @@ async def slow_handler(args): assert "timed out" in result.content.lower() +@pytest.mark.asyncio +async def test_execute_sync_tool_timeout(registry): + import time + + def slow_handler(args): + time.sleep(0.2) + return "done" + + registry.register_from_def( + ToolDef(name="slow_sync", description="Slow sync tool", parameters={"type": "object", "properties": {}}), + slow_handler, + ) + executor = ToolExecutor(registry, FilesystemGuard("/tmp"), timeout_ms=10, max_output_chars=100) + + started = time.monotonic() + result = await executor.execute("slow_sync", {}, "tc1") + + assert result.success is False + assert "timed out" in result.content.lower() + assert time.monotonic() - started < 0.1 + + +@pytest.mark.asyncio +async def test_execute_builtin_error_string_is_failure(tmp_path): + from tiny_harness.tools.files import read_file + + registry = ToolRegistry() + registry.register_from_def( + ToolDef( + name="read_file", + description="Read", + parameters={ + "type": "object", + "properties": {"path": {"type": "string"}}, + "required": ["path"], + }, + ), + read_file, + ) + executor = ToolExecutor(registry, FilesystemGuard(str(tmp_path))) + + result = await executor.execute("read_file", {"path": "missing.txt"}, "tc1") + + assert result.success is False + assert "not found" in result.content.lower() + + +@pytest.mark.asyncio +async def test_execute_rejects_non_object_arguments(registry): + executor = ToolExecutor(registry, FilesystemGuard("/tmp")) + + result = await executor.execute("echo", ["not", "an", "object"], "tc1") + + assert result.success is False + assert "object" in result.content.lower() + + @pytest.mark.asyncio async def test_execute_result_truncation(registry): executor = ToolExecutor(registry, FilesystemGuard("/tmp"), timeout_ms=5000, max_output_chars=20) @@ -185,6 +247,75 @@ def sync_handler(args): assert "HELLO" in result.content +@pytest.mark.asyncio +async def test_execute_resolves_relative_path_against_workspace(tmp_path, monkeypatch): + from tiny_harness.tools.files import write_file + + workspace = tmp_path / "workspace" + process_cwd = tmp_path / "process-cwd" + workspace.mkdir() + process_cwd.mkdir() + monkeypatch.chdir(process_cwd) + registry = ToolRegistry() + registry.register_from_def( + ToolDef( + name="write_file", + description="Write", + parameters={ + "type": "object", + "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, + "required": ["path", "content"], + }, + risk_level="mutation", + ), + write_file, + ) + executor = ToolExecutor(registry, FilesystemGuard(str(workspace))) + + result = await executor.execute("write_file", {"path": "note.txt", "content": "safe"}, "tc1") + + assert result.success is True + assert (workspace / "note.txt").read_text() == "safe" + assert not (process_cwd / "note.txt").exists() + + +@pytest.mark.asyncio +async def test_execute_guards_every_path_for_move(tmp_path): + from tiny_harness.tools.files import move_file + + workspace = tmp_path / "workspace" + outside = tmp_path / "outside" + workspace.mkdir() + outside.mkdir() + source = workspace / "source.txt" + source.write_text("keep") + registry = ToolRegistry() + registry.register_from_def( + ToolDef( + name="move_file", + description="Move", + parameters={ + "type": "object", + "properties": {"source": {"type": "string"}, "destination": {"type": "string"}}, + "required": ["source", "destination"], + }, + risk_level="mutation", + ), + move_file, + ) + executor = ToolExecutor(registry, FilesystemGuard(str(workspace))) + + result = await executor.execute( + "move_file", + {"source": str(source), "destination": str(outside / "moved.txt")}, + "tc1", + ) + + assert result.success is False + assert source.exists() + assert not (outside / "moved.txt").exists() + + def test_validate_schema_required_missing(): schema = {"type": "object", "properties": {"x": {"type": "string"}}, "required": ["x"]} errors = validate_schema(schema, {}) @@ -199,6 +330,15 @@ def test_validate_schema_wrong_type(): assert "integer" in errors[0] +def test_validate_schema_rejects_boolean_as_integer(): + schema = {"type": "object", "properties": {"count": {"type": "integer"}}} + + errors = validate_schema(schema, {"count": True}) + + assert len(errors) == 1 + assert "integer" in errors[0] + + def test_validate_schema_enum(): schema = {"type": "object", "properties": {"x": {"type": "string", "enum": ["a", "b"]}}} errors = validate_schema(schema, {"x": "c"}) diff --git a/tiny_harness/__init__.py b/tiny_harness/__init__.py index 411c17f..bf85216 100644 --- a/tiny_harness/__init__.py +++ b/tiny_harness/__init__.py @@ -3,5 +3,9 @@ from tiny_harness._tools import ToolDef from tiny_harness._core import Agent from tiny_harness._hitl import ApprovalDecision, ApprovalHandler, SessionApprovals +from tiny_harness._persist import SessionStore -__all__ = ["Agent", "Prompt", "Config", "ToolDef", "ApprovalDecision", "ApprovalHandler", "SessionApprovals"] +__all__ = [ + "Agent", "Prompt", "Config", "ToolDef", "ApprovalDecision", + "ApprovalHandler", "SessionApprovals", "SessionStore", +] diff --git a/tiny_harness/_config.py b/tiny_harness/_config.py index f4626ca..0b28c06 100644 --- a/tiny_harness/_config.py +++ b/tiny_harness/_config.py @@ -1,11 +1,11 @@ # tiny_harness/_config.py -from dataclasses import dataclass +from dataclasses import dataclass, field @dataclass class AgentConfig: model: str - api_key: str + api_key: str = field(repr=False) workspace: str provider: str = "deepseek" api_base_url: str | None = None @@ -19,6 +19,20 @@ class AgentConfig: approval_timeout_ms: int = 120_000 no_hitl: bool = False + def __post_init__(self) -> None: + positive_limits = ( + "max_iterations", + "max_errors", + "max_consecutive_errors", + "timeout_ms", + "max_tool_result_chars", + "max_tokens", + "approval_timeout_ms", + ) + for name in positive_limits: + if getattr(self, name) <= 0: + raise ValueError(f"{name} must be positive") + class Prompt: def __init__(self, base: str): diff --git a/tiny_harness/_core.py b/tiny_harness/_core.py index 7adf6af..3574596 100644 --- a/tiny_harness/_core.py +++ b/tiny_harness/_core.py @@ -1,6 +1,7 @@ # tiny_harness/_core.py import importlib.util import asyncio +import json from pathlib import Path from types import ModuleType from collections.abc import AsyncIterator @@ -45,7 +46,12 @@ def _create_provider(self, config: AgentConfig): case "anthropic": return AnthropicProvider(config.api_key, config.model, max_tokens=config.max_tokens) case "openai" | "deepseek": - base_url = config.api_base_url or "https://api.openai.com/v1" + default_url = ( + "https://api.deepseek.com/v1" + if config.provider == "deepseek" + else "https://api.openai.com/v1" + ) + base_url = config.api_base_url or default_url return OpenAIProvider(config.api_key, config.model, base_url=base_url, max_tokens=config.max_tokens) case _: raise ValueError(f"Unknown provider: {config.provider}") @@ -75,23 +81,25 @@ async def filtered(event: StreamEvent): self._event_bus.subscribe(filtered) def load_skill(self, skill_ref: str) -> None: - if skill_ref in self._loaded_skills: - return module = self._resolve_skill(skill_ref) + spec = getattr(module, "__spec__", None) + identity = spec.origin if spec and spec.origin else module.__name__ + if identity in self._loaded_skills: + return if not hasattr(module, "register"): raise RuntimeError(f"Skill '{skill_ref}' has no register() function") module.register(self) - self._loaded_skills.append(skill_ref) + self._messages.update_prompt(self._prompt) + self._loaded_skills.append(identity) def _resolve_skill(self, ref: str) -> ModuleType: - try: - return importlib.import_module(f"tiny_harness.skills.{ref}") - except ImportError: - pass - try: - return importlib.import_module(ref) - except ImportError: - pass + for module_name in (f"tiny_harness.skills.{ref}", ref): + try: + return importlib.import_module(module_name) + except ModuleNotFoundError as e: + missing_candidate = e.name == module_name or module_name.startswith(f"{e.name}.") + if not missing_candidate: + raise path = Path(ref) if path.exists(): spec = importlib.util.spec_from_file_location(path.stem, path) @@ -104,27 +112,26 @@ async def run(self, prompt: str) -> str: if self._running: raise RuntimeError("Agent is already running a task") self._running = True - self._save_turn("user", content=prompt) try: loop = AgentLoop(self._config, self._messages, self._llm_provider, self._tool_executor, self._event_bus) result = await loop.run(prompt) - self._save_turn("assistant", content=result) return result finally: + if self._session_id is not None: + self._dump_conversation() self._running = False async def run_stream(self, prompt: str) -> AsyncIterator[StreamEvent]: if self._running: raise RuntimeError("Agent is already running a task") self._running = True - self._save_turn("user", content=prompt) queue: list[StreamEvent] = [] async def collector(event: StreamEvent): queue.append(event) self._event_bus.subscribe(collector) - result_text = "" + task: asyncio.Task | None = None try: loop = AgentLoop(self._config, self._messages, self._llm_provider, self._tool_executor, self._event_bus) task = asyncio.create_task(loop.run(prompt)) @@ -132,15 +139,21 @@ async def collector(event: StreamEvent): while not task.done() or last_yielded < len(queue): while last_yielded < len(queue): event = queue[last_yielded] - if event.type == "text_delta" and event.content: - result_text += event.content yield event last_yielded += 1 if not task.done(): await asyncio.sleep(0.01) await task - self._save_turn("assistant", content=result_text if result_text else None) finally: + if task is not None and not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + self._event_bus.unsubscribe(collector) + if self._session_id is not None: + self._dump_conversation() self._running = False def clear(self) -> None: @@ -154,6 +167,10 @@ def start_session(self, session_id: str | None = None) -> str: self._session_id = session_id or self._store.new_session() return self._session_id + def save_session(self) -> tuple[str, int]: + session_id = self.start_session() + return session_id, self._dump_conversation() + def resume_session(self, session_id: str) -> int: """Reload a conversation from a saved session. Returns number of turns restored.""" if self._store is None: @@ -211,7 +228,7 @@ def _dump_conversation(self) -> int: if self._store is None or self._session_id is None: return 0 messages = self._messages.to_list() - count = 0 + turns: list[dict] = [] i = 1 while i < len(messages): @@ -219,24 +236,31 @@ def _dump_conversation(self) -> int: role = msg.get("role", "") if role == "user": - count += 1 - self._store.save_turn( - session_id=self._session_id, - chat_id=count, - role="user", - content=msg.get("content", ""), - model=self._config.model, - ) + turns.append({ + "role": "user", + "content": msg.get("content", ""), + "model": self._config.model, + }) i += 1 elif role == "assistant": raw_tool_calls = msg.get("tool_calls") tool_calls = None if raw_tool_calls: - tool_calls = [ - {"name": tc.get("function", {}).get("name", "?"), "arguments": tc.get("function", {}).get("arguments", {})} - for tc in raw_tool_calls - ] + tool_calls = [] + for tc in raw_tool_calls: + function = tc.get("function", {}) + arguments = function.get("arguments", {}) + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except json.JSONDecodeError: + arguments = {} + tool_calls.append({ + "id": tc.get("id", ""), + "name": function.get("name", "?"), + "arguments": arguments, + }) tool_results = [] j = i + 1 while j < len(messages) and messages[j].get("role") == "tool": @@ -247,16 +271,13 @@ def _dump_conversation(self) -> int: }) j += 1 - count += 1 - self._store.save_turn( - session_id=self._session_id, - chat_id=count, - role="assistant", - content=msg.get("content"), - tool_calls=tool_calls, - tool_results=tool_results if tool_results else None, - model=self._config.model, - ) + turns.append({ + "role": "assistant", + "content": msg.get("content"), + "tool_calls": tool_calls, + "tool_results": tool_results if tool_results else None, + "model": self._config.model, + }) i = j elif role == "tool": @@ -265,8 +286,9 @@ def _dump_conversation(self) -> int: else: i += 1 - self._chat_id = count - return count + self._store.replace_session(self._session_id, turns) + self._chat_id = len(turns) + return len(turns) @property def max_iterations(self) -> int: diff --git a/tiny_harness/_events.py b/tiny_harness/_events.py index 5e58835..562941a 100644 --- a/tiny_harness/_events.py +++ b/tiny_harness/_events.py @@ -22,6 +22,12 @@ def __init__(self): def subscribe(self, handler: Callable[[StreamEvent], Awaitable[Any]]) -> None: self._handlers.append(handler) + def unsubscribe(self, handler: Callable[[StreamEvent], Awaitable[Any]]) -> None: + try: + self._handlers.remove(handler) + except ValueError: + pass + async def emit(self, event: StreamEvent) -> None: for handler in self._handlers: await handler(event) diff --git a/tiny_harness/_hitl.py b/tiny_harness/_hitl.py index 7b35817..62079e1 100644 --- a/tiny_harness/_hitl.py +++ b/tiny_harness/_hitl.py @@ -56,6 +56,8 @@ def __init__(self) -> None: def approve_for_session(self, risk_level: str) -> None: """Approve this risk level and all lower levels for the session.""" + if risk_level not in ("safe", "read_only", "mutation"): + return for level in self._RISK_ORDER: self._approved_levels.add(level) if level == risk_level: diff --git a/tiny_harness/_llm.py b/tiny_harness/_llm.py index deda00d..1702c1a 100644 --- a/tiny_harness/_llm.py +++ b/tiny_harness/_llm.py @@ -7,6 +7,9 @@ import asyncio import random +from tiny_harness._sse import anthropic_events, openai_events +from tiny_harness._llm_response import ResponseAdapter + @dataclass class ToolCallRequest: @@ -77,7 +80,43 @@ def __init__(self, api_key: str, model: str, retry_config: LLMRetryConfig | None self._max_tokens = max_tokens def _convert_messages(self, messages: list[dict]) -> list[dict]: - return [m for m in messages if m["role"] != "system"] + converted: list[dict] = [] + for message in messages: + role = message["role"] + if role == "system": + continue + if role == "assistant" and message.get("tool_calls"): + content = [] + if message.get("content"): + content.append({"type": "text", "text": message["content"]}) + for tool_call in message["tool_calls"]: + function = tool_call.get("function", {}) + arguments = function.get("arguments", {}) + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except json.JSONDecodeError: + arguments = {} + content.append({ + "type": "tool_use", + "id": tool_call.get("id", ""), + "name": function.get("name", ""), + "input": arguments, + }) + converted.append({"role": "assistant", "content": content}) + elif role == "tool": + block = { + "type": "tool_result", + "tool_use_id": message.get("tool_call_id", ""), + "content": message.get("content", ""), + } + if converted and converted[-1]["role"] == "user" and isinstance(converted[-1]["content"], list): + converted[-1]["content"].append(block) + else: + converted.append({"role": "user", "content": [block]}) + else: + converted.append({"role": role, "content": message.get("content", "")}) + return converted def _extract_system(self, messages: list[dict]) -> str: for m in messages: @@ -138,7 +177,7 @@ async def generate(self, messages, tools=None) -> LLMResponse: headers={"x-api-key": self._api_key, "anthropic-version": "2023-06-01", "content-type": "application/json"}, ) if resp.status_code == 200: - return self._parse_response(_ResponseAdapter(resp.json())) + return self._parse_response(ResponseAdapter(resp.json())) if resp.status_code in (429, 529) or resp.status_code >= 500: raise RetryableLLMError(f"Status {resp.status_code}") if resp.status_code in (401, 403): @@ -165,57 +204,41 @@ async def generate_stream(self, messages, tools=None): if converted_tools: body["tools"] = converted_tools - async with httpx.AsyncClient(timeout=120.0) as client: - async with client.stream( - "POST", "https://api.anthropic.com/v1/messages", - json=body, - headers={"x-api-key": self._api_key, "anthropic-version": "2023-06-01", "content-type": "application/json"}, - ) as response: - if response.status_code != 200: - body_bytes = await response.aread() - raise FatalLLMError(f"Stream failed: {response.status_code} {body_bytes}") - async for line in response.aiter_lines(): - parsed = self._parse_sse_line(line) - if parsed is None: - continue - event_type, value = parsed - if event_type == "message_start": - yield LLMStreamChunk(type="usage", content=json.dumps({"input": value.get("message", {}).get("usage", {}).get("input_tokens", 0)})) - elif event_type == "content_block_start": - block = value.get("content_block", {}) - if block.get("type") == "tool_use": - yield LLMStreamChunk(type="tool_call_start", content=json.dumps({"id": block["id"], "name": block["name"]})) - elif event_type == "content_block_delta": - delta = value.get("delta", {}) - if delta.get("type") == "text_delta": - yield LLMStreamChunk(type="text_delta", content=delta.get("text", "")) - elif delta.get("type") == "input_json_delta": - yield LLMStreamChunk(type="tool_call_delta", content=delta.get("partial_json", "")) - elif event_type == "content_block_stop": - yield LLMStreamChunk(type="tool_call_end") - - -class _ResponseAdapter: - def __init__(self, data: dict): - self.content = [_BlockAdapter(b) for b in data.get("content", [])] - self.usage = _UsageAdapter(data.get("usage", {})) - self.model = data.get("model", "") - self.stop_reason = data.get("stop_reason", "stop") - - -class _BlockAdapter: - def __init__(self, data: dict): - self.type = data.get("type", "") - self.text = data.get("text", "") - self.name = data.get("name", "") - self.input = data.get("input", {}) - self.id = data.get("id", "") - - -class _UsageAdapter: - def __init__(self, data: dict): - self.input_tokens = data.get("input_tokens", 0) - self.output_tokens = data.get("output_tokens", 0) + for attempt in range(self._retry_config.max_retries + 1): + emitted = False + try: + async with httpx.AsyncClient(timeout=120.0) as client: + async with client.stream( + "POST", "https://api.anthropic.com/v1/messages", + json=body, + headers={"x-api-key": self._api_key, "anthropic-version": "2023-06-01", "content-type": "application/json"}, + ) as response: + if response.status_code != 200: + body_bytes = await response.aread() + if response.status_code in (429, 529) or response.status_code >= 500: + raise RetryableLLMError(f"Status {response.status_code}") + raise FatalLLMError(f"Stream failed: {response.status_code} {body_bytes}") + async for chunk in self._stream_chunks(response): + emitted = True + yield chunk + return + except (RetryableLLMError, httpx.TimeoutException, httpx.ConnectError) as e: + if emitted or attempt == self._retry_config.max_retries: + raise FatalLLMError(f"Stream failed after retries: {e}") + delay = min( + self._retry_config.base_delay * (self._retry_config.backoff_factor ** attempt), + self._retry_config.max_delay, + ) + await asyncio.sleep(delay + random.uniform(0, delay * 0.5)) + + async def _stream_chunks(self, response): + async for event in anthropic_events(response.aiter_lines(), self._parse_sse_line): + tool_call = event.get("tool_call") + yield LLMStreamChunk( + type=event["type"], + content=event.get("content"), + tool_call=ToolCallRequest(**tool_call) if tool_call else None, + ) class OpenAIProvider(LLMProvider): @@ -286,58 +309,38 @@ async def generate_stream(self, messages, tools=None): if converted_tools: body["tools"] = converted_tools - async with httpx.AsyncClient(timeout=120.0) as client: - async with client.stream( - "POST", f"{self._base_url}/chat/completions", - json=body, - headers={"Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json"}, - ) as response: - if response.status_code != 200: - body_bytes = await response.aread() - raise FatalLLMError(f"Stream failed: {response.status_code} {body_bytes}") - - tool_buffers: dict[int, dict] = {} - - async for line in response.aiter_lines(): - line = line.strip() - if not line or not line.startswith("data: "): - continue - data_str = line[6:] - if data_str == "[DONE]": - break - try: - data = json.loads(data_str) - except json.JSONDecodeError: - continue - - delta = data.get("choices", [{}])[0].get("delta", {}) - if delta.get("content"): - yield LLMStreamChunk(type="text_delta", content=delta["content"]) - - if delta.get("tool_calls"): - for tc in delta["tool_calls"]: - idx = tc.get("index", 0) - if idx not in tool_buffers: - tool_buffers[idx] = {"id": "", "name": "", "arguments": ""} - - buf = tool_buffers[idx] - if "id" in tc and tc["id"]: - buf["id"] = tc["id"] - func = tc.get("function", {}) - if "name" in func and func["name"]: - buf["name"] = func["name"] - yield LLMStreamChunk(type="tool_call_start", content=json.dumps({"id": buf["id"], "name": buf["name"]})) - if "arguments" in func: - buf["arguments"] += func["arguments"] - - finish = data.get("choices", [{}])[0].get("finish_reason") - if finish in ("tool_calls", "stop"): - for idx in sorted(tool_buffers.keys()): - buf = tool_buffers[idx] - if buf["name"] and buf["arguments"]: - try: - args = json.loads(buf["arguments"]) - except json.JSONDecodeError: - args = {} - tc = ToolCallRequest(id=buf["id"] or f"call_{idx}", name=buf["name"], arguments=args) - yield LLMStreamChunk(type="tool_call_end", tool_call=tc) + for attempt in range(self._retry_config.max_retries + 1): + emitted = False + try: + async with httpx.AsyncClient(timeout=120.0) as client: + async with client.stream( + "POST", f"{self._base_url}/chat/completions", + json=body, + headers={"Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json"}, + ) as response: + if response.status_code != 200: + body_bytes = await response.aread() + if response.status_code in (429, 529) or response.status_code >= 500: + raise RetryableLLMError(f"Status {response.status_code}") + raise FatalLLMError(f"Stream failed: {response.status_code} {body_bytes}") + async for chunk in self._stream_chunks(response): + emitted = True + yield chunk + return + except (RetryableLLMError, httpx.TimeoutException, httpx.ConnectError) as e: + if emitted or attempt == self._retry_config.max_retries: + raise FatalLLMError(f"Stream failed after retries: {e}") + delay = min( + self._retry_config.base_delay * (self._retry_config.backoff_factor ** attempt), + self._retry_config.max_delay, + ) + await asyncio.sleep(delay + random.uniform(0, delay * 0.5)) + + async def _stream_chunks(self, response): + async for event in openai_events(response.aiter_lines()): + tool_call = event.get("tool_call") + yield LLMStreamChunk( + type=event["type"], + content=event.get("content"), + tool_call=ToolCallRequest(**tool_call) if tool_call else None, + ) diff --git a/tiny_harness/_llm_response.py b/tiny_harness/_llm_response.py new file mode 100644 index 0000000..1afcd83 --- /dev/null +++ b/tiny_harness/_llm_response.py @@ -0,0 +1,21 @@ +class ResponseAdapter: + def __init__(self, data: dict): + self.content = [BlockAdapter(block) for block in data.get("content", [])] + self.usage = UsageAdapter(data.get("usage", {})) + self.model = data.get("model", "") + self.stop_reason = data.get("stop_reason", "stop") + + +class BlockAdapter: + def __init__(self, data: dict): + self.type = data.get("type", "") + self.text = data.get("text", "") + self.name = data.get("name", "") + self.input = data.get("input", {}) + self.id = data.get("id", "") + + +class UsageAdapter: + def __init__(self, data: dict): + self.input_tokens = data.get("input_tokens", 0) + self.output_tokens = data.get("output_tokens", 0) diff --git a/tiny_harness/_loop.py b/tiny_harness/_loop.py index e59598e..534f5fc 100644 --- a/tiny_harness/_loop.py +++ b/tiny_harness/_loop.py @@ -72,7 +72,9 @@ async def run(self, user_prompt: str) -> str: tool_calls.append(chunk.tool_call) except Exception as e: await self._events.emit(StreamEvent(type="error", message=f"LLM error: {e}")) - return f"Agent stopped due to LLM error: {e}" + final_text = f"Agent stopped due to LLM error: {e}" + self._messages.add_assistant(text=final_text, tool_calls=None) + return final_text if not tool_calls: final_text = "".join(collected_text) @@ -82,7 +84,7 @@ async def run(self, user_prompt: str) -> str: self._messages.add_assistant(text="".join(collected_text) if collected_text else None, tool_calls=tool_calls) collected_text = [] - for tc in tool_calls: + for call_index, tc in enumerate(tool_calls): if not loop_detector.check(tc.name, tc.arguments): result = ToolResult.error(tc.id, f"You've called '{tc.name}' with the same arguments {loop_detector.max_repeats} times. Try a different approach.") else: @@ -90,14 +92,22 @@ async def run(self, user_prompt: str) -> str: result = await self._tools.execute(tc.name, tc.arguments, tc.id) await self._events.emit(StreamEvent(type="tool_end", tool_name=tc.name, content=result.content[:100])) + exhausted = False if result.success: error_budget.record_success() elif result.denied: pass else: if not error_budget.record_error(): - return await self._degraded_finish(collected_text) + exhausted = True self._messages.add_tool_result(result.tool_call_id, result.content) + if exhausted: + for pending in tool_calls[call_index + 1:]: + self._messages.add_tool_result( + pending.id, + "Tool call cancelled because the error budget was exhausted.", + ) + return await self._degraded_finish(collected_text) status = self._messages.check_context() if status == TokenStatus.NEAR_CAPACITY: @@ -109,6 +119,10 @@ async def _degraded_finish(self, collected_text: list[str]) -> str: self._messages.add_system_notice("You've reached a safety limit. Please provide your best final answer based on what you know, without using any tools.") try: result = await self._llm.generate(self._messages.to_list(), tools=[]) - return result.text or "".join(collected_text) + final_text = result.text or "".join(collected_text) except Exception: - return "".join(collected_text) or "Agent stopped." + final_text = "".join(collected_text) or "Agent stopped." + self._messages.add_assistant(text=final_text, tool_calls=None) + if final_text: + await self._events.emit(StreamEvent(type="text_delta", content=final_text)) + return final_text diff --git a/tiny_harness/_messages.py b/tiny_harness/_messages.py index 5ea988c..601b21e 100644 --- a/tiny_harness/_messages.py +++ b/tiny_harness/_messages.py @@ -64,3 +64,6 @@ def check_context(self) -> TokenStatus: def clear(self) -> None: system = self.messages[0] self.messages = [system] + + def update_prompt(self, prompt: Prompt) -> None: + self.messages[0] = {"role": "system", "content": prompt.to_string()} diff --git a/tiny_harness/_persist.py b/tiny_harness/_persist.py index f137e87..4211704 100644 --- a/tiny_harness/_persist.py +++ b/tiny_harness/_persist.py @@ -24,6 +24,42 @@ def save_turn(self, session_id: str, chat_id: int, role: str, model: str | None = None, iteration: int | None = None, metadata: dict | None = None) -> None: + record = self._build_record( + session_id, chat_id, role, content, tool_calls, tool_results, + token_usage, model, iteration, metadata, + ) + filepath = self._session_path(session_id) + with open(filepath, "a", encoding="utf-8") as f: + f.write(json.dumps(record, ensure_ascii=False) + "\n") + + def replace_session(self, session_id: str, turns: list[dict]) -> None: + filepath = self._session_path(session_id) + temporary = filepath.with_suffix(".jsonl.tmp") + with open(temporary, "w", encoding="utf-8") as f: + for chat_id, turn in enumerate(turns, 1): + record = self._build_record( + session_id=session_id, + chat_id=chat_id, + role=turn["role"], + content=turn.get("content"), + tool_calls=turn.get("tool_calls"), + tool_results=turn.get("tool_results"), + token_usage=turn.get("token_usage"), + model=turn.get("model"), + iteration=turn.get("iteration"), + metadata=turn.get("metadata"), + ) + f.write(json.dumps(record, ensure_ascii=False) + "\n") + os.replace(temporary, filepath) + + def _build_record(self, session_id: str, chat_id: int, role: str, + content: str | None = None, + tool_calls: list | None = None, + tool_results: list | None = None, + token_usage: dict | None = None, + model: str | None = None, + iteration: int | None = None, + metadata: dict | None = None) -> dict: record = { "session_id": session_id, "chat_id": chat_id, @@ -34,14 +70,14 @@ def save_turn(self, session_id: str, chat_id: int, role: str, record["content"] = content if tool_calls: record["tool_calls"] = [ - {"name": tc["name"], "arguments": tc.get("arguments", {})} - if isinstance(tc, dict) else {"name": tc.name, "arguments": tc.arguments} + {"id": tc.get("id", ""), "name": tc["name"], "arguments": tc.get("arguments", {})} + if isinstance(tc, dict) else {"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in tool_calls ] if tool_results: record["tool_results"] = [ - {"id": tr.get("tool_call_id", ""), "content": str(tr.get("content", ""))[:200]} - if isinstance(tr, dict) else {"id": tr.tool_call_id, "content": str(tr.content)[:200]} + {"id": tr.get("id", tr.get("tool_call_id", "")), "content": str(tr.get("content", ""))} + if isinstance(tr, dict) else {"id": tr.tool_call_id, "content": str(tr.content)} for tr in tool_results ] if token_usage: @@ -52,10 +88,7 @@ def save_turn(self, session_id: str, chat_id: int, role: str, record["iteration"] = iteration if metadata: record["metadata"] = metadata - - filepath = self._session_path(session_id) - with open(filepath, "a", encoding="utf-8") as f: - f.write(json.dumps(record, ensure_ascii=False) + "\n") + return record def load_session(self, session_id: str) -> list[dict]: filepath = self._session_path(session_id) diff --git a/tiny_harness/_sse.py b/tiny_harness/_sse.py new file mode 100644 index 0000000..a64aba7 --- /dev/null +++ b/tiny_harness/_sse.py @@ -0,0 +1,114 @@ +import json +from collections.abc import AsyncIterator, Callable +from typing import Any + + +async def anthropic_events( + lines: AsyncIterator[str], + parse_line: Callable[[str], tuple[str, Any] | None], +) -> AsyncIterator[dict]: + current_event: str | None = None + tool_buffers: dict[int, dict[str, str]] = {} + async for line in lines: + parsed = parse_line(line) + if parsed is None: + continue + event_type, value = parsed + if value is None: + current_event = event_type + continue + event_type = current_event or value.get("type", event_type) + current_event = None + if event_type == "message_start": + usage = value.get("message", {}).get("usage", {}).get("input_tokens", 0) + yield {"type": "usage", "content": json.dumps({"input": usage})} + elif event_type == "content_block_start": + index = value.get("index", 0) + block = value.get("content_block", {}) + if block.get("type") == "tool_use": + tool_buffers[index] = { + "id": block.get("id", ""), + "name": block.get("name", ""), + "arguments": "", + } + yield { + "type": "tool_call_start", + "content": json.dumps({"id": block["id"], "name": block["name"]}), + } + elif event_type == "content_block_delta": + index = value.get("index", 0) + delta = value.get("delta", {}) + if delta.get("type") == "text_delta": + yield {"type": "text_delta", "content": delta.get("text", "")} + elif delta.get("type") == "input_json_delta": + partial = delta.get("partial_json", "") + if index in tool_buffers: + tool_buffers[index]["arguments"] += partial + yield {"type": "tool_call_delta", "content": partial} + elif event_type == "content_block_stop": + buffer = tool_buffers.pop(value.get("index", 0), None) + if buffer is not None: + try: + arguments = json.loads(buffer["arguments"] or "{}") + except json.JSONDecodeError: + arguments = {} + yield { + "type": "tool_call_end", + "tool_call": { + "id": buffer["id"], + "name": buffer["name"], + "arguments": arguments, + }, + } + + +async def openai_events(lines: AsyncIterator[str]) -> AsyncIterator[dict]: + tool_buffers: dict[int, dict[str, str]] = {} + async for line in lines: + line = line.strip() + if not line or not line.startswith("data: "): + continue + data_str = line[6:] + if data_str == "[DONE]": + break + try: + data = json.loads(data_str) + except json.JSONDecodeError: + continue + + choice = data.get("choices", [{}])[0] + delta = choice.get("delta", {}) + if delta.get("content"): + yield {"type": "text_delta", "content": delta["content"]} + + for tool_call in delta.get("tool_calls", []): + index = tool_call.get("index", 0) + buffer = tool_buffers.setdefault(index, {"id": "", "name": "", "arguments": ""}) + if tool_call.get("id"): + buffer["id"] = tool_call["id"] + function = tool_call.get("function", {}) + if function.get("name"): + buffer["name"] = function["name"] + yield { + "type": "tool_call_start", + "content": json.dumps({"id": buffer["id"], "name": buffer["name"]}), + } + if "arguments" in function: + buffer["arguments"] += function["arguments"] + + if choice.get("finish_reason") in ("tool_calls", "stop"): + for index in sorted(tool_buffers): + buffer = tool_buffers[index] + if buffer["name"] and buffer["arguments"]: + try: + arguments = json.loads(buffer["arguments"]) + except json.JSONDecodeError: + arguments = {} + yield { + "type": "tool_call_end", + "tool_call": { + "id": buffer["id"] or f"call_{index}", + "name": buffer["name"], + "arguments": arguments, + }, + } diff --git a/tiny_harness/_tools.py b/tiny_harness/_tools.py index fd95691..a78907d 100644 --- a/tiny_harness/_tools.py +++ b/tiny_harness/_tools.py @@ -22,6 +22,11 @@ class ToolDef: parameters: dict risk_level: str = "read_only" + def __post_init__(self) -> None: + allowed = {"safe", "read_only", "mutation", "destructive", "dangerous"} + if self.risk_level not in allowed: + raise ValueError(f"Unknown tool risk level: {self.risk_level!r}") + @dataclass class Tool: @@ -98,16 +103,10 @@ async def execute(self, name: str, args: dict, call_id: str) -> ToolResult: if errors: return ToolResult.error(call_id, f"Invalid arguments for '{name}':\n" + "\n".join(f" - {e}" for e in errors)) - if self._guard and tool.definition.risk_level != "safe": - path = args.get("path") or args.get("source") or args.get("destination") - if not path: - path = args.get("cwd") - if path: - try: - op = "delete" if tool.definition.risk_level == "destructive" else "write" if tool.definition.risk_level == "mutation" else "read" - self._guard.guard(path, op) - except Exception as e: - return ToolResult.error(call_id, str(e)) + try: + args = self._guard_paths(args, tool.definition.risk_level) + except Exception as e: + return ToolResult.error(call_id, str(e)) if self._approval_gate is not None: decision = await self._approval_gate.check(name, args, tool.definition.risk_level) @@ -117,30 +116,50 @@ async def execute(self, name: str, args: dict, call_id: str) -> ToolResult: return ToolResult.denial(call_id, f"Tool '{name}' denied: {decision.reason}") if decision.modified_args is not None: args = decision.modified_args - if self._guard and tool.definition.risk_level != "safe": - path = args.get("path") or args.get("source") or args.get("destination") - if not path: - path = args.get("cwd") - if path: - try: - op = "delete" if tool.definition.risk_level == "destructive" else "write" if tool.definition.risk_level == "mutation" else "read" - self._guard.guard(path, op) - except Exception as e: - return ToolResult.error(call_id, str(e)) + errors = validate_schema(tool.definition.parameters, args) + if errors: + return ToolResult.error( + call_id, + f"Invalid arguments for '{name}':\n" + "\n".join(f" - {e}" for e in errors), + ) + try: + args = self._guard_paths(args, tool.definition.risk_level) + except Exception as e: + return ToolResult.error(call_id, str(e)) try: if asyncio.iscoroutinefunction(tool.handler): - raw = await asyncio.wait_for(tool.handler(args), timeout=self._timeout_ms / 1000) + invocation = tool.handler(args) else: - raw = tool.handler(args) + invocation = asyncio.to_thread(tool.handler, args) + raw = await asyncio.wait_for(invocation, timeout=self._timeout_ms / 1000) except asyncio.TimeoutError: return ToolResult.error(call_id, f"Tool '{name}' timed out after {self._timeout_ms/1000}s") except Exception as e: return ToolResult.error(call_id, f"Tool '{name}' failed: {e}") formatted = self._format(raw) + if _is_error_output(formatted): + return ToolResult.error(call_id, formatted) return ToolResult.ok(call_id, formatted) + def _guard_paths(self, args: dict, risk_level: str) -> dict: + if not self._guard or risk_level == "safe": + return args + guarded = dict(args) + default_op = "delete" if risk_level == "destructive" else "write" if risk_level == "mutation" else "read" + operations = { + "path": default_op, + "source": "read", + "destination": "write", + "cwd": "read", + } + for field, operation in operations.items(): + path = guarded.get(field) + if path: + guarded[field] = self._guard.guard(path, operation) + return guarded + def _format(self, raw) -> str: if raw is None: return "Success." @@ -160,6 +179,8 @@ def validate_schema(schema: dict, args: dict) -> list[str]: schema_type = schema.get("type") if schema_type != "object": return errors + if not isinstance(args, dict): + return [f"arguments should be an object, got {type(args).__name__}"] properties = schema.get("properties", {}) required = schema.get("required", []) for field in required: @@ -171,9 +192,9 @@ def validate_schema(schema: dict, args: dict) -> list[str]: expected_type = prop.get("type") if expected_type == "string" and not isinstance(value, str): errors.append(f"'{key}' should be a string, got {type(value).__name__}") - elif expected_type == "integer" and not isinstance(value, int): + elif expected_type == "integer" and (not isinstance(value, int) or isinstance(value, bool)): errors.append(f"'{key}' should be an integer, got {type(value).__name__}") - elif expected_type == "number" and not isinstance(value, (int, float)): + elif expected_type == "number" and (not isinstance(value, (int, float)) or isinstance(value, bool)): errors.append(f"'{key}' should be a number, got {type(value).__name__}") elif expected_type == "boolean" and not isinstance(value, bool): errors.append(f"'{key}' should be a boolean, got {type(value).__name__}") @@ -184,3 +205,8 @@ def validate_schema(schema: dict, args: dict) -> list[str]: if "enum" in prop and value not in prop["enum"]: errors.append(f"'{key}' must be one of {prop['enum']}, got {value!r}") return errors + + +def _is_error_output(output: str) -> bool: + stripped = output.lstrip() + return stripped.startswith("Error:") or stripped.startswith("[exit code:") diff --git a/tiny_harness/cli.py b/tiny_harness/cli.py index 3093a5d..3ddf025 100644 --- a/tiny_harness/cli.py +++ b/tiny_harness/cli.py @@ -10,16 +10,31 @@ def parse_args() -> Namespace: parser = ArgumentParser(prog="tiny-harness", description="AI agent harness with tools and streaming CLI") parser.add_argument("prompt", nargs="?", default=None, help="Prompt for one-shot mode") - parser.add_argument("--model", "-m", default="deepseek-chat", help="Model identifier") + parser.add_argument("--model", "-m", default=None, help="Model identifier") parser.add_argument("--workspace", "-w", default=os.getcwd(), help="Workspace directory") - parser.add_argument("--provider", default="deepseek", help="LLM provider: anthropic, openai, deepseek") - parser.add_argument("--api-base-url", default="https://api.deepseek.com/v1", help="Custom API base URL") + parser.add_argument( + "--provider", + choices=("anthropic", "openai", "deepseek"), + default="deepseek", + help="LLM provider", + ) + parser.add_argument("--api-base-url", default=None, help="Custom API base URL") parser.add_argument("--max-iterations", type=int, default=25, help="Max loop iterations") - parser.add_argument("--skills", default="", help="Comma-separated skill names (default: files)") + parser.add_argument("--skills", default=None, help="Comma-separated skill names (default: files; empty disables tools)") parser.add_argument("--no-hitl", action="store_true", help="Disable human-in-the-loop approval prompts") parser.add_argument("--tui", action="store_true", help="Use rich TUI mode (requires 'rich' package)") - parser.add_argument("--api-key-env", default="DEEPSEEK_API_KEY", help="Env var for API key") - return parser.parse_args() + parser.add_argument("--api-key-env", default=None, help="Env var for API key") + args = parser.parse_args() + defaults = { + "deepseek": ("deepseek-chat", "DEEPSEEK_API_KEY"), + "openai": ("gpt-4o-mini", "OPENAI_API_KEY"), + "anthropic": ("claude-sonnet-4-20250514", "ANTHROPIC_API_KEY"), + } + if args.provider in defaults: + default_model, default_key_env = defaults[args.provider] + args.model = args.model or default_model + args.api_key_env = args.api_key_env or default_key_env + return args def _get_api_key(args: Namespace) -> str: @@ -74,9 +89,7 @@ async def _run_session(args: Namespace, agent): print("Conversation cleared.") continue if user_input == "/save": - agent.start_session() - turns = agent._dump_conversation() - sid = agent.session_id + sid, turns = agent.save_session() print(f"Session saved: {sid} ({turns} turns) -> ~/.tiny-harness/sessions/{sid}.jsonl") continue if user_input == "/history": @@ -155,9 +168,11 @@ def main(): agent.set_approval_handler(_cli_approval_handler) # Load skills from --skills flag, or default to files if none specified - skill_names = [s.strip() for s in args.skills.split(",") if s.strip()] - if not skill_names: - skill_names = ["files"] + skill_names = ( + ["files"] + if args.skills is None + else [s.strip() for s in args.skills.split(",") if s.strip()] + ) for skill_name in skill_names: agent.load_skill(skill_name) diff --git a/tiny_harness/tools/files.py b/tiny_harness/tools/files.py index 5315af4..46f76f4 100644 --- a/tiny_harness/tools/files.py +++ b/tiny_harness/tools/files.py @@ -14,6 +14,10 @@ def read_file(args: dict) -> str: return f"Error: File '{path}' not found." if os.path.isdir(path): return f"Error: '{path}' is a directory." + with open(path, "rb") as f: + sample = f.read(8192) + if b"\x00" in sample: + return f"[{path}] Binary file ({_format_size(os.path.getsize(path))})." with open(path, "r", encoding="utf-8", errors="replace") as f: lines = f.readlines() total = len(lines) @@ -75,14 +79,15 @@ def find_files(args: dict) -> str: import os as _os base = _os.path.realpath(path) - search_path = _os.path.join(base, pattern) if _os.path.isabs(pattern) or ".." in pattern.split(_os.sep): - resolved = _os.path.realpath(search_path) if _os.path.exists(search_path) else _os.path.realpath(base) - if not resolved.startswith(base + _os.sep) and resolved != base: - return f"Error: pattern '{pattern}' escapes workspace." + return f"Error: pattern '{pattern}' escapes workspace." + search_path = _os.path.join(base, pattern) - for i, match in enumerate(glob.glob(search_path, recursive=True)): - if i >= max_results: + for match in glob.glob(search_path, recursive=True): + resolved = _os.path.realpath(match) + if resolved != base and not resolved.startswith(base + _os.sep): + continue + if len(matches) >= max_results: break matches.append(os.path.relpath(match, path)) if not matches: diff --git a/tiny_harness/tools/search.py b/tiny_harness/tools/search.py index d7a924c..98bbe11 100644 --- a/tiny_harness/tools/search.py +++ b/tiny_harness/tools/search.py @@ -16,10 +16,9 @@ def search_content(args: dict) -> str: "--glob", "!.git", "--glob", "!__pycache__", "--glob", "!*.pyc"] if file_pattern: cmd.extend(["--glob", file_pattern]) - cmd.extend([pattern, path if os.path.isdir(path) else os.path.dirname(path) or "."]) + cmd.extend([pattern, path]) result = subprocess.run( - cmd, - capture_output=True, text=True, timeout=10, cwd=path, + cmd, capture_output=True, text=True, timeout=10, ) if result.returncode not in (0, 1): raise RuntimeError(result.stderr) diff --git a/tiny_harness/tools/shell.py b/tiny_harness/tools/shell.py index 284a9b9..fedb2d6 100644 --- a/tiny_harness/tools/shell.py +++ b/tiny_harness/tools/shell.py @@ -44,13 +44,6 @@ def run_command(args: dict) -> str: cwd = args.get("cwd") timeout = args.get("timeout", 30) - if cwd: - import os as _os - resolved = _os.path.realpath(_os.path.expanduser(cwd)) - here = _os.path.realpath(".") - if _os.path.commonpath([resolved, here]) != here: - return f"Error: cwd '{cwd}' is outside workspace." - cmd, use_shell = _build_args(command) try: result = subprocess.run( diff --git a/tiny_harness/tui.py b/tiny_harness/tui.py index 2c709f5..78542ce 100644 --- a/tiny_harness/tui.py +++ b/tiny_harness/tui.py @@ -118,9 +118,7 @@ async def tui_approval_handler(request: ToolApprovalRequest) -> ApprovalDecision continue if prompt == "/save": - agent.start_session() - turns = agent._dump_conversation() - sid = agent.session_id + sid, turns = agent.save_session() console.print(f" Saved: {sid} ({turns} turns)", style=c["ok"]) console.print() continue