pidriver 0.0.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- pidriver-0.0.1/.gitignore +30 -0
- pidriver-0.0.1/LICENSE +21 -0
- pidriver-0.0.1/PKG-INFO +490 -0
- pidriver-0.0.1/README.md +466 -0
- pidriver-0.0.1/examples/basic_prompt.py +53 -0
- pidriver-0.0.1/pyproject.toml +55 -0
- pidriver-0.0.1/src/pidriver/__init__.py +132 -0
- pidriver-0.0.1/src/pidriver/_transport.py +272 -0
- pidriver-0.0.1/src/pidriver/client.py +92 -0
- pidriver-0.0.1/src/pidriver/config.py +259 -0
- pidriver-0.0.1/src/pidriver/errors.py +77 -0
- pidriver-0.0.1/src/pidriver/events.py +441 -0
- pidriver-0.0.1/src/pidriver/interaction.py +214 -0
- pidriver-0.0.1/src/pidriver/manager.py +192 -0
- pidriver-0.0.1/src/pidriver/py.typed +0 -0
- pidriver-0.0.1/src/pidriver/session.py +255 -0
- pidriver-0.0.1/src/pidriver/usage.py +98 -0
- pidriver-0.0.1/tests/conftest.py +80 -0
- pidriver-0.0.1/tests/fixtures/README.md +91 -0
- pidriver-0.0.1/tests/fixtures/_capture_driver.py +140 -0
- pidriver-0.0.1/tests/fixtures/frames/agent_end.json +115 -0
- pidriver-0.0.1/tests/fixtures/frames/agent_start.json +3 -0
- pidriver-0.0.1/tests/fixtures/frames/extension_ui_request.confirm.json +7 -0
- pidriver-0.0.1/tests/fixtures/frames/extension_ui_request.editor.json +7 -0
- pidriver-0.0.1/tests/fixtures/frames/extension_ui_request.input.json +7 -0
- pidriver-0.0.1/tests/fixtures/frames/extension_ui_request.select.approval.json +10 -0
- pidriver-0.0.1/tests/fixtures/frames/extension_ui_request.select.ask.json +12 -0
- pidriver-0.0.1/tests/fixtures/frames/extension_ui_request.setWidget.json +6 -0
- pidriver-0.0.1/tests/fixtures/frames/extension_ui_response.approve.json +5 -0
- pidriver-0.0.1/tests/fixtures/frames/extension_ui_response.confirm.json +5 -0
- pidriver-0.0.1/tests/fixtures/frames/extension_ui_response.editor.json +5 -0
- pidriver-0.0.1/tests/fixtures/frames/extension_ui_response.input.json +5 -0
- pidriver-0.0.1/tests/fixtures/frames/extension_ui_response.value.json +5 -0
- pidriver-0.0.1/tests/fixtures/frames/message_start_assistant.json +37 -0
- pidriver-0.0.1/tests/fixtures/frames/message_start_user.json +14 -0
- pidriver-0.0.1/tests/fixtures/frames/message_update.text_delta.json +82 -0
- pidriver-0.0.1/tests/fixtures/frames/message_update.thinking_delta.json +76 -0
- pidriver-0.0.1/tests/fixtures/frames/message_update.toolcall_end.json +103 -0
- pidriver-0.0.1/tests/fixtures/frames/message_update.toolcall_start.json +95 -0
- pidriver-0.0.1/tests/fixtures/frames/ready.json +3 -0
- pidriver-0.0.1/tests/fixtures/frames/response_prompt_ack.json +6 -0
- pidriver-0.0.1/tests/fixtures/frames/tool_execution_end.json +18 -0
- pidriver-0.0.1/tests/fixtures/frames/tool_execution_start.json +9 -0
- pidriver-0.0.1/tests/fixtures/frames/tool_execution_update.json +18 -0
- pidriver-0.0.1/tests/fixtures/frames/turn_end.json +65 -0
- pidriver-0.0.1/tests/fixtures/frames/turn_start.json +3 -0
- pidriver-0.0.1/tests/fixtures/payloads/get_session_stats.response.json +23 -0
- pidriver-0.0.1/tests/fixtures/payloads/get_state.response.json +3984 -0
- pidriver-0.0.1/tests/test_config.py +160 -0
- pidriver-0.0.1/tests/test_events.py +287 -0
- pidriver-0.0.1/tests/test_fixtures_replay.py +51 -0
- pidriver-0.0.1/tests/test_manager.py +95 -0
- pidriver-0.0.1/tests/test_replay_session.py +139 -0
- pidriver-0.0.1/tests/test_session.py +283 -0
- pidriver-0.0.1/tests/test_transport.py +119 -0
- pidriver-0.0.1/tests/test_ui_methods.py +74 -0
- pidriver-0.0.1/tests/test_usage.py +62 -0
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
# Python
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*.egg-info/
|
|
5
|
+
.eggs/
|
|
6
|
+
build/
|
|
7
|
+
dist/
|
|
8
|
+
.venv/
|
|
9
|
+
venv/
|
|
10
|
+
|
|
11
|
+
# Tooling caches
|
|
12
|
+
.pytest_cache/
|
|
13
|
+
.mypy_cache/
|
|
14
|
+
.ruff_cache/
|
|
15
|
+
.coverage
|
|
16
|
+
htmlcov/
|
|
17
|
+
|
|
18
|
+
# uv
|
|
19
|
+
uv.lock
|
|
20
|
+
|
|
21
|
+
# Editor / OS
|
|
22
|
+
.DS_Store
|
|
23
|
+
*.swp
|
|
24
|
+
.idea/
|
|
25
|
+
.vscode/
|
|
26
|
+
|
|
27
|
+
# pidriver runtime (isolated pi installs, sessions, workspaces)
|
|
28
|
+
.pi/
|
|
29
|
+
*.jsonl
|
|
30
|
+
scratch/
|
pidriver-0.0.1/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Grigory Bakunov
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
pidriver-0.0.1/PKG-INFO
ADDED
|
@@ -0,0 +1,490 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: pidriver
|
|
3
|
+
Version: 0.0.1
|
|
4
|
+
Summary: Async Python driver for the pi coding agent (pi --mode rpc), built for isolated, headless project work.
|
|
5
|
+
Project-URL: Homepage, https://github.com/bobuk/pidriver
|
|
6
|
+
Project-URL: Repository, https://github.com/bobuk/pidriver
|
|
7
|
+
Author-email: Grigory Bakunov <thebobuk@gmail.com>
|
|
8
|
+
License: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Keywords: automation,coding-agent,llm,pi,rpc
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
16
|
+
Classifier: Typing :: Typed
|
|
17
|
+
Requires-Python: >=3.12
|
|
18
|
+
Provides-Extra: dev
|
|
19
|
+
Requires-Dist: mypy>=1.11; extra == 'dev'
|
|
20
|
+
Requires-Dist: pytest-asyncio>=0.24; extra == 'dev'
|
|
21
|
+
Requires-Dist: pytest>=8; extra == 'dev'
|
|
22
|
+
Requires-Dist: ruff>=0.6; extra == 'dev'
|
|
23
|
+
Description-Content-Type: text/markdown
|
|
24
|
+
|
|
25
|
+
# pidriver
|
|
26
|
+
|
|
27
|
+
Async Python driver for the pi coding agent, built for **isolated, headless project work**. It
|
|
28
|
+
drives any `pi`-compatible CLI in `--mode rpc` over its JSONL stdin/stdout protocol and exposes a
|
|
29
|
+
small, typed Python API — with no third-party dependencies.
|
|
30
|
+
|
|
31
|
+
The recommended CLI is [**oh-my-pi**](https://bun.sh) (the `omp` binary); the original
|
|
32
|
+
[`pi`](https://www.npmjs.com/package/@earendil-works/pi-coding-agent) also works. Which one runs
|
|
33
|
+
is just a `PiConfig` field (`binary`), so the same code drives either.
|
|
34
|
+
|
|
35
|
+
Embed an agent inside another application (a chat bot, a scheduler, a service) and have it
|
|
36
|
+
develop a project autonomously, in an environment that is **isolated** from your personal pi
|
|
37
|
+
configuration and secrets.
|
|
38
|
+
|
|
39
|
+
Start with [`PiClient`](#high-level-api-piclient--pisession) for typed events and interaction
|
|
40
|
+
handling; drop to the [raw transport](#transport) only when you need it.
|
|
41
|
+
|
|
42
|
+
## Requirements
|
|
43
|
+
|
|
44
|
+
- Python **3.12+**
|
|
45
|
+
- A `pi`-compatible CLI on `PATH`. Recommended — **oh-my-pi** (the `omp` binary) via
|
|
46
|
+
[Bun](https://bun.sh), or the official installer:
|
|
47
|
+
|
|
48
|
+
```sh
|
|
49
|
+
bun install -g @oh-my-pi/pi-coding-agent # provides the `omp` binary
|
|
50
|
+
# or: curl -fsSL https://omp.sh/install | sh
|
|
51
|
+
omp --version
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
The original `pi` CLI works too: `npm i -g @earendil-works/pi-coding-agent`. Tell `PiConfig`
|
|
55
|
+
which one to launch with `binary` (`"omp"` or `"pi"`); they share the `--mode rpc` protocol.
|
|
56
|
+
- A provider API key — passed explicitly through `PiConfig`, see [Isolation](#isolation).
|
|
57
|
+
|
|
58
|
+
## Install
|
|
59
|
+
|
|
60
|
+
```sh
|
|
61
|
+
uv add pidriver # once published
|
|
62
|
+
# or, from a checkout:
|
|
63
|
+
uv pip install -e ".[dev]"
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
## Quick start
|
|
67
|
+
|
|
68
|
+
Configure a client once, then `start()` a session per task and iterate its typed events:
|
|
69
|
+
|
|
70
|
+
```python
|
|
71
|
+
import asyncio, os
|
|
72
|
+
from pidriver import PiClient, PiConfig, AutoApprove, MessageDelta, ToolStart, AgentEnd
|
|
73
|
+
|
|
74
|
+
async def main():
|
|
75
|
+
client = PiClient(PiConfig(
|
|
76
|
+
binary="omp", # which CLI to launch (oh-my-pi); "pi" also works
|
|
77
|
+
provider="openai",
|
|
78
|
+
model="gpt-4o",
|
|
79
|
+
api_key=os.environ["OPENAI_API_KEY"],
|
|
80
|
+
# Isolation defaults (scrubbed env, no host extensions/skills) are already on.
|
|
81
|
+
))
|
|
82
|
+
|
|
83
|
+
# AutoApprove answers permission prompts itself, so the run is unattended.
|
|
84
|
+
session = await client.start(
|
|
85
|
+
"List the files here and summarize the project.",
|
|
86
|
+
cwd="/srv/projects/acme", # the project the agent works in
|
|
87
|
+
interaction_handler=AutoApprove(),
|
|
88
|
+
)
|
|
89
|
+
async with session:
|
|
90
|
+
async for event in session:
|
|
91
|
+
match event:
|
|
92
|
+
case MessageDelta(text=text):
|
|
93
|
+
print(text, end="", flush=True)
|
|
94
|
+
case ToolStart(name=name, arguments=args):
|
|
95
|
+
print(f"\n[tool] {name} {args}")
|
|
96
|
+
case AgentEnd(reason=reason):
|
|
97
|
+
print(f"\n[done: {reason}]")
|
|
98
|
+
|
|
99
|
+
asyncio.run(main())
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
Need the raw protocol instead? Drive [`SubprocessTransport`](#transport) directly — see
|
|
103
|
+
[`examples/basic_prompt.py`](examples/basic_prompt.py).
|
|
104
|
+
|
|
105
|
+
## Architecture
|
|
106
|
+
|
|
107
|
+
`pi --mode rpc` speaks line-delimited JSON on stdin/stdout. pidriver is layered so each concern
|
|
108
|
+
is swappable and testable in isolation:
|
|
109
|
+
|
|
110
|
+
```
|
|
111
|
+
PiConfig ──► SubprocessTransport ──► pi --mode rpc
|
|
112
|
+
(argv+env) (JSONL framing only) (child process)
|
|
113
|
+
│ │
|
|
114
|
+
│ raw JSON dicts (events + responses, undifferentiated)
|
|
115
|
+
│ ▼
|
|
116
|
+
│ PiClient/PiSession ── typed events, id-correlated commands,
|
|
117
|
+
│ interaction handling ── respond()
|
|
118
|
+
▼
|
|
119
|
+
PiSessionManager ── registry · idle reaper · stop_all() over many sessions
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
| Module | Role |
|
|
123
|
+
| --- | --- |
|
|
124
|
+
| `pidriver.config` | `PiConfig` — isolation knobs, `to_argv()` / `to_env()` |
|
|
125
|
+
| `pidriver._transport` | `PiTransport` protocol + `SubprocessTransport` (the swap boundary) |
|
|
126
|
+
| `pidriver.events` | the typed `Event` hierarchy + `parse_event` |
|
|
127
|
+
| `pidriver.interaction` | interaction handlers/policies (`AskHost`, `AutoApprove`, …) |
|
|
128
|
+
| `pidriver.session` | `PiSession` (back-compat alias `AgentSession`) — the live RPC session |
|
|
129
|
+
| `pidriver.client` | `PiClient` — starts sessions |
|
|
130
|
+
| `pidriver.manager` | `PiSessionManager` — registry, idle reaper, `stop_all()` |
|
|
131
|
+
| `pidriver.usage` | `UsageTotals` — token/cost accounting |
|
|
132
|
+
| `pidriver.errors` | exception hierarchy |
|
|
133
|
+
|
|
134
|
+
Everything below is re-exported from the package root: `from pidriver import ...`.
|
|
135
|
+
|
|
136
|
+
## `PiConfig`
|
|
137
|
+
|
|
138
|
+
The single source of truth for **how** a pi process is spawned — binary, provider/model, tools,
|
|
139
|
+
session, and isolation. Immutable (frozen dataclass). The transport consumes two derived
|
|
140
|
+
outputs: `to_argv()` (the command line) and `to_env()` (the child environment).
|
|
141
|
+
|
|
142
|
+
```python
|
|
143
|
+
PiConfig(provider="anthropic", model="claude-sonnet-4-6",
|
|
144
|
+
api_key="sk-...", workspace="/path/to/project")
|
|
145
|
+
```
|
|
146
|
+
|
|
147
|
+
**Binary & workspace**
|
|
148
|
+
|
|
149
|
+
| Field | Default | Purpose |
|
|
150
|
+
| --- | --- | --- |
|
|
151
|
+
| `binary` | `"pi"` | Which `pi`-compatible CLI to launch — set to `"omp"` for oh-my-pi. |
|
|
152
|
+
| `workspace` | `None` | Project directory; becomes the subprocess `cwd`. |
|
|
153
|
+
|
|
154
|
+
**Model / provider**
|
|
155
|
+
|
|
156
|
+
| Field | Default | Purpose |
|
|
157
|
+
| --- | --- | --- |
|
|
158
|
+
| `provider` | `None` | `--provider` (e.g. `"anthropic"`, `"openai"`). |
|
|
159
|
+
| `model` | `None` | `--model`. |
|
|
160
|
+
| `thinking` | `None` | `--thinking` (`off`/`minimal`/`low`/`medium`/`high`/`xhigh`). |
|
|
161
|
+
| `api_key` | `None` | Secret, injected into the provider's env var by `to_env()`. |
|
|
162
|
+
| `api_key_env` | `None` | Override the env var name `api_key` is injected under. |
|
|
163
|
+
|
|
164
|
+
Prefer a **fully-qualified** `model` id (e.g. `openai/gpt-4o`, `anthropic/claude-haiku-4-5`) —
|
|
165
|
+
bare fuzzy names can mis-route to the wrong backend.
|
|
166
|
+
|
|
167
|
+
**Tools** — `tools` (`--tools` allowlist), `exclude_tools` (`--exclude-tools`), `no_tools`,
|
|
168
|
+
`no_builtin_tools`.
|
|
169
|
+
|
|
170
|
+
**Session** — `session` (False → `--no-session`), `session_id`, `session_path` (`--session`),
|
|
171
|
+
`session_name` (`--name`), `continue_session` (`--continue`), `system_prompt`,
|
|
172
|
+
`append_system_prompt`.
|
|
173
|
+
|
|
174
|
+
**Isolation** — see [Isolation](#isolation): `agent_dir`, `session_dir`, `no_extensions`,
|
|
175
|
+
`no_skills`, `no_context_files`, `no_prompt_templates`, `no_themes`, `offline`,
|
|
176
|
+
`skip_version_check`, `inherit_env`, `env_passthrough`.
|
|
177
|
+
|
|
178
|
+
**Escape hatches** — `extra_args` (appended to argv), `extra_env` (merged into the child env).
|
|
179
|
+
|
|
180
|
+
Methods:
|
|
181
|
+
|
|
182
|
+
- `to_argv() -> list[str]` — full command line, starting `[binary, "--mode", "rpc", ...]`.
|
|
183
|
+
- `to_env(base_environ=None) -> dict[str, str]` — the child environment (scrubbed by default).
|
|
184
|
+
- `api_key_env_name() -> str | None` — which env var `api_key` resolves to (provider-mapped via
|
|
185
|
+
`PROVIDER_API_KEY_ENV`, falling back to `<PROVIDER>_API_KEY`, then `ANTHROPIC_API_KEY`).
|
|
186
|
+
|
|
187
|
+
## Isolation
|
|
188
|
+
|
|
189
|
+
This is the reason pidriver exists. By default a `PiConfig` keeps an agent that's developing a
|
|
190
|
+
project from reading or mutating your global pi setup:
|
|
191
|
+
|
|
192
|
+
- **Scrubbed environment** (`inherit_env=False`, default). The child starts from an *empty*
|
|
193
|
+
environment plus only `env_passthrough` (`PATH`, `HOME`, `LANG`, `LC_ALL`, `TERM`, `TZ`) — so
|
|
194
|
+
a stray `*_API_KEY` in your shell never leaks into the agent. The configured `api_key` is then
|
|
195
|
+
injected under exactly one provider variable.
|
|
196
|
+
- **Private agent dir.** `agent_dir` sets `PI_CODING_AGENT_DIR`, pointing pi at a private config
|
|
197
|
+
directory instead of `~/.pi`, so it loads no personal extensions/skills/settings.
|
|
198
|
+
- **No host discovery** (all default **True**): `no_extensions`, `no_skills`,
|
|
199
|
+
`no_context_files`, `no_prompt_templates`, `no_themes` give a clean, reproducible agent.
|
|
200
|
+
- **Separate session storage.** `session_dir` (`--session-dir`) keeps session `.jsonl` files out
|
|
201
|
+
of the default location.
|
|
202
|
+
- `offline` (`--offline` + `PI_OFFLINE=1`) and `skip_version_check` (`PI_SKIP_VERSION_CHECK=1`,
|
|
203
|
+
default True) round out a hermetic run.
|
|
204
|
+
|
|
205
|
+
```python
|
|
206
|
+
cfg = PiConfig(
|
|
207
|
+
provider="anthropic", model="claude-sonnet-4-6", api_key=KEY,
|
|
208
|
+
workspace="/srv/projects/acme",
|
|
209
|
+
agent_dir="/var/lib/myapp/pi-home", # private ~/.pi replacement
|
|
210
|
+
session_dir="/var/lib/myapp/pi-sessions",
|
|
211
|
+
extra_env={"HOME": "/var/lib/myapp"}, # see note below — also override HOME
|
|
212
|
+
# inherit_env=False and no_* discovery flags are already the defaults.
|
|
213
|
+
)
|
|
214
|
+
```
|
|
215
|
+
|
|
216
|
+
> **Override `HOME` for full isolation.** `agent_dir`/`PI_CODING_AGENT_DIR` alone is **not**
|
|
217
|
+
> enough: omp derives its **log** root from `$HOME/.omp`, and the default `env_passthrough`
|
|
218
|
+
> copies `HOME` from the host — so an agent can still write logs into your `~/.omp`. To leave
|
|
219
|
+
> your personal setup byte-for-byte untouched, also point `HOME` at the private dir (via
|
|
220
|
+
> `extra_env`, as above) or drop it from `env_passthrough`. (Verified empirically against
|
|
221
|
+
> omp 15.7.3.)
|
|
222
|
+
|
|
223
|
+
## Transport
|
|
224
|
+
|
|
225
|
+
### `SubprocessTransport(config, *, on_stderr=None, spawn=..., terminate_timeout=5.0)`
|
|
226
|
+
|
|
227
|
+
A deliberately **dumb** transport: it spawns `pi --mode rpc`, writes JSON commands as JSONL
|
|
228
|
+
lines, and yields parsed JSON dicts from stdout. It does **not** correlate request/response ids
|
|
229
|
+
or interpret commands — that's the session layer's job. This thinness is what makes it the swap
|
|
230
|
+
boundary (a future omp-rpc transport can satisfy the same `PiTransport` protocol).
|
|
231
|
+
|
|
232
|
+
```python
|
|
233
|
+
await transport.start() # resolve pi binary + spawn
|
|
234
|
+
await transport.send({"type": "prompt", ...}) # write one JSONL command
|
|
235
|
+
obj = await transport.receive() # next JSON dict, or None at EOF
|
|
236
|
+
async for obj in transport: ... # iterate until EOF (single consumer)
|
|
237
|
+
await transport.aclose() # close stdin, SIGTERM→SIGKILL, reap
|
|
238
|
+
transport.pid, transport.returncode # process introspection
|
|
239
|
+
```
|
|
240
|
+
|
|
241
|
+
Contract: exactly **one** consumer iterates at a time; `receive()` returns `None` at EOF;
|
|
242
|
+
`send`/`receive` raise `TransportClosedError` before `start()` or after the process exits.
|
|
243
|
+
`on_stderr` receives pi's diagnostic lines (never interleaved into the dict stream). Framing
|
|
244
|
+
splits on `\n` only (a trailing `\r` is stripped), never on U+2028/U+2029, so JSON strings
|
|
245
|
+
containing those survive intact.
|
|
246
|
+
|
|
247
|
+
`PiTransport` is the `@runtime_checkable` Protocol that `SubprocessTransport` implements — type
|
|
248
|
+
against it and inject a fake transport in tests.
|
|
249
|
+
|
|
250
|
+
## Session manager
|
|
251
|
+
|
|
252
|
+
### `PiSessionManager(factory=None, *, idle_timeout=1800.0, reap_interval=60.0, max_sessions=None, clock=...)`
|
|
253
|
+
|
|
254
|
+
A registry + idle reaper + bulk shutdown for many concurrent sessions. Decoupled from the
|
|
255
|
+
concrete session class — it depends only on the structural `ManagedSession` protocol
|
|
256
|
+
(`session_id`, `last_activity`, `aclose()`).
|
|
257
|
+
|
|
258
|
+
```python
|
|
259
|
+
async with PiSessionManager(idle_timeout=1800) as mgr:
|
|
260
|
+
await mgr.register(session) # add an already-started session
|
|
261
|
+
# or: await mgr.create(...) # build via the injected factory
|
|
262
|
+
mgr.get(session_id) # -> session (raises SessionNotFoundError)
|
|
263
|
+
mgr.ids(); len(mgr); sid in mgr # introspection
|
|
264
|
+
await mgr.stop(session_id) # close + drop one
|
|
265
|
+
await mgr.stop_all() # close all concurrently
|
|
266
|
+
# context exit stops the reaper and tears everything down
|
|
267
|
+
```
|
|
268
|
+
|
|
269
|
+
The background reaper closes sessions idle past `idle_timeout` (set `None` to disable);
|
|
270
|
+
`max_sessions` caps the registry (`register`/`create` raise `RuntimeError` when exceeded).
|
|
271
|
+
|
|
272
|
+
## Usage accounting
|
|
273
|
+
|
|
274
|
+
### `UsageTotals`
|
|
275
|
+
|
|
276
|
+
An immutable token + cost tally that normalizes pi's two usage shapes into one addable value:
|
|
277
|
+
|
|
278
|
+
```python
|
|
279
|
+
from pidriver import UsageTotals
|
|
280
|
+
|
|
281
|
+
a = UsageTotals.from_session_stats(stats_data) # get_session_stats response
|
|
282
|
+
b = UsageTotals.from_assistant_usage(msg["usage"]) # per-message usage block
|
|
283
|
+
total = a + b # aggregate with +
|
|
284
|
+
total.total_tokens # input+output+cache_read+cache_write
|
|
285
|
+
total.with_cost(0.42) # copy with cost replaced (stats may lack cost)
|
|
286
|
+
total.as_dict() # plain dict for logging/serialization
|
|
287
|
+
```
|
|
288
|
+
|
|
289
|
+
Fields: `input_tokens`, `output_tokens`, `cache_read_tokens`, `cache_write_tokens`, `cost_usd`.
|
|
290
|
+
|
|
291
|
+
## Errors
|
|
292
|
+
|
|
293
|
+
All exceptions derive from `PiDriverError`, so one `except` catches the family:
|
|
294
|
+
|
|
295
|
+
| Exception | Raised when |
|
|
296
|
+
| --- | --- |
|
|
297
|
+
| `PiDriverError` | Base class for everything below. |
|
|
298
|
+
| `PiNotFoundError` | The `pi` executable can't be found or run. |
|
|
299
|
+
| `PiStartError` | The `pi --mode rpc` process failed to start. |
|
|
300
|
+
| `PiProtocolError` | A line from pi wasn't valid JSON / violated the contract (`.raw`). |
|
|
301
|
+
| `PiCommandError` | An RPC command returned `success: false` (`.command`, `.error`, `.data`). |
|
|
302
|
+
| `PiTimeoutError` | A response or awaited condition didn't arrive in time. |
|
|
303
|
+
| `TransportClosedError` | An op was attempted on a closed/exited transport. |
|
|
304
|
+
| `SessionNotFoundError` | A session id isn't registered with the manager (`.session_id`). |
|
|
305
|
+
|
|
306
|
+
## High-level API: `PiClient` / `PiSession`
|
|
307
|
+
|
|
308
|
+
The transport yields raw, undifferentiated JSON. The session layer adds the ergonomic surface
|
|
309
|
+
this library is ultimately for: typed events, an auto-answer policy for the agent's questions,
|
|
310
|
+
and a session object that plugs straight into `PiSessionManager`.
|
|
311
|
+
|
|
312
|
+
### `PiClient`
|
|
313
|
+
|
|
314
|
+
```python
|
|
315
|
+
client = PiClient(config: PiConfig, *,
|
|
316
|
+
transport_factory=None, # build a transport from a config (tests)
|
|
317
|
+
interaction_handler=None) # default handler for sessions it starts
|
|
318
|
+
|
|
319
|
+
session = await client.start(
|
|
320
|
+
task: str, *,
|
|
321
|
+
cwd: str | Path | None = None, # project dir; overrides config.workspace
|
|
322
|
+
config: PiConfig | None = None, # per-session config override
|
|
323
|
+
interaction_handler: InteractionHandler | None = None, # default: client's, else AskHost
|
|
324
|
+
resume: str | None = None, # prior pi session path/id → --continue
|
|
325
|
+
session_id: str | None = None, # stable registry key (auto uuid4 if omitted)
|
|
326
|
+
send_initial: bool = True, # False → open without sending `task` yet
|
|
327
|
+
) -> PiSession # AgentSession is a back-compat alias of the same class
|
|
328
|
+
```
|
|
329
|
+
|
|
330
|
+
`start()` spawns the subprocess and (by default) sends `task` as the first prompt. Its signature
|
|
331
|
+
is exactly what `PiSessionManager(factory=...)` expects, so you can wire them together:
|
|
332
|
+
|
|
333
|
+
```python
|
|
334
|
+
manager = PiSessionManager(factory=client.start)
|
|
335
|
+
session = await manager.create("Add a healthcheck endpoint", cwd="/srv/proj")
|
|
336
|
+
```
|
|
337
|
+
|
|
338
|
+
### `PiSession`
|
|
339
|
+
|
|
340
|
+
The live session (`AgentSession` is a back-compat alias for the same class) — async-iterate it
|
|
341
|
+
for events, command it with the methods below. It's an async context manager (`async with
|
|
342
|
+
session:` guarantees the subprocess is reaped) and satisfies `ManagedSession`, so it drops into
|
|
343
|
+
`PiSessionManager`.
|
|
344
|
+
|
|
345
|
+
| Member | Description |
|
|
346
|
+
| --- | --- |
|
|
347
|
+
| `async for event in session` / `session.events()` | Yields typed [`Event`](#events) objects. Consuming twice raises. |
|
|
348
|
+
| `await prompt(message)` | Send a follow-up user turn (completes at the next `AgentEnd`). |
|
|
349
|
+
| `await respond(request_id, value)` | Answer an interaction (see [below](#answering-the-agent)). |
|
|
350
|
+
| `await cancel()` | Interrupt the current turn (keeps the process alive). |
|
|
351
|
+
| `await aclose()` | Idempotently terminate the session and reap the subprocess. |
|
|
352
|
+
| `session.pending` | The current unanswered `InteractionRequest`, or `None`. |
|
|
353
|
+
| `session.usage` | Running `UsageTotals` accumulated from inline usage blocks. |
|
|
354
|
+
| `session.ended` | `True` once an `AgentEnd` has been seen. |
|
|
355
|
+
| `session.session_id` | Stable local id (registry key). |
|
|
356
|
+
| `session.pi_session_id` | The id pi reports in `agent_start` (pass to `resume`). |
|
|
357
|
+
|
|
358
|
+
### Events
|
|
359
|
+
|
|
360
|
+
Each RPC record maps to a frozen dataclass via `parse_event`. Unknown event types degrade to a
|
|
361
|
+
bare `Event` (original payload in `.raw`) instead of raising — a new pi/omp release never crashes
|
|
362
|
+
the driver. Every event carries `.type` (the raw tag) and `.raw` (the decoded dict).
|
|
363
|
+
|
|
364
|
+
| Event | Key fields | Meaning |
|
|
365
|
+
| --- | --- | --- |
|
|
366
|
+
| `AgentStart` | `session_id`, `model`, `cwd` | The agent run has begun. |
|
|
367
|
+
| `MessageDelta` | `text`, `channel` | A streaming fragment of assistant (or `thinking`) text. |
|
|
368
|
+
| `MessageComplete` | `text`, `channel` | A full message at a turn boundary. |
|
|
369
|
+
| `ToolStart` | `tool_call_id`, `name`, `arguments` | The agent invoked a tool. |
|
|
370
|
+
| `ToolUpdate` | `tool_call_id`, `name`, `output` | Partial output while a tool runs. |
|
|
371
|
+
| `ToolEnd` | `tool_call_id`, `name`, `result`, `is_error` | A tool finished. |
|
|
372
|
+
| `Usage` | `totals` (a `UsageTotals`) | Token/cost for a step; also folded into `session.usage`. |
|
|
373
|
+
| `AgentEnd` | `reason`, `final_text` | The run finished (`completed` / `cancelled` / `error` / `limit`). |
|
|
374
|
+
| `Error` | `message`, `code`, `fatal` | An error (or a failed command ack) surfaced by pi. |
|
|
375
|
+
| `InteractionRequest` | `request_id`, `kind`, `prompt`, `options`, `tool_call_id`, `default` | The agent is blocked waiting for the host. |
|
|
376
|
+
| `Event` | `type`, `raw` | Base / fallback for unrecognized records. |
|
|
377
|
+
|
|
378
|
+
`channel` is a `Channel` enum (`ASSISTANT` / `THINKING`).
|
|
379
|
+
|
|
380
|
+
### Answering the agent
|
|
381
|
+
|
|
382
|
+
When the agent needs the host it emits an `InteractionRequest` whose `kind` is an
|
|
383
|
+
`InteractionKind`:
|
|
384
|
+
|
|
385
|
+
| `InteractionKind` | What the agent wants | Answer `value` |
|
|
386
|
+
| --- | --- | --- |
|
|
387
|
+
| `PERMISSION` | Approve/deny a gated tool call (`tool_call_id` set) | a `Decision`, or `True`/`False` |
|
|
388
|
+
| `QUESTION` | A free-text answer | a `str` |
|
|
389
|
+
| `CHOICE` | Pick from `options` | the chosen `str`, or its `int` index |
|
|
390
|
+
|
|
391
|
+
`Decision` values: `ALLOW`, `ALLOW_ALWAYS`, `DENY`. `session.respond()` takes the request **id**
|
|
392
|
+
and coerces the value:
|
|
393
|
+
|
|
394
|
+
```python
|
|
395
|
+
await session.respond(req.request_id, Decision.ALLOW) # explicit decision
|
|
396
|
+
await session.respond(req.request_id, True) # bool → allow / deny
|
|
397
|
+
await session.respond(req.request_id, "use postgres") # free-text answer
|
|
398
|
+
await session.respond(req.request_id, 0) # choice by index
|
|
399
|
+
```
|
|
400
|
+
|
|
401
|
+
An **`InteractionHandler`** — any `async (request, session) -> InteractionResponse | None` — can
|
|
402
|
+
answer requests automatically. Returning `None` **defers** to the host (the request still surfaces
|
|
403
|
+
through the event stream). Built-ins:
|
|
404
|
+
|
|
405
|
+
| Handler | Behavior |
|
|
406
|
+
| --- | --- |
|
|
407
|
+
| `AskHost()` | **Default.** Never auto-answers — every request surfaces to your loop. |
|
|
408
|
+
| `AutoApprove(allow=None, *, always=False)` | Auto-approves `PERMISSION` (optionally filtered by an `allow(request)` predicate; `always=True` sends `ALLOW_ALWAYS`). **Questions/choices are deferred** — no safe default. |
|
|
409
|
+
| `DenyAll()` | Denies every `PERMISSION`; defers other kinds. A read-only sandbox. |
|
|
410
|
+
| `chain(h1, h2, ...)` | Composes handlers; first non-`None` response wins, else defers. |
|
|
411
|
+
|
|
412
|
+
Build responses directly with `InteractionResponse(req.request_id, value)`,
|
|
413
|
+
`InteractionResponse.allow(req.request_id, always=True)`, or `.deny(req.request_id)`.
|
|
414
|
+
|
|
415
|
+
> **Surfacing permission prompts.** omp defaults to `--approval-mode yolo` (auto-approve
|
|
416
|
+
> everything), so it won't emit `PERMISSION` requests at all. To make a human-in-the-loop
|
|
417
|
+
> (`AskHost`) flow meaningful, run the CLI in a stricter mode — pass `--approval-mode write`
|
|
418
|
+
> (or `always-ask`) via `PiConfig.extra_args`.
|
|
419
|
+
|
|
420
|
+
### Interaction examples
|
|
421
|
+
|
|
422
|
+
**AskHost — human in the loop.** The default handler defers everything, so each request surfaces
|
|
423
|
+
in your loop and you `respond()`:
|
|
424
|
+
|
|
425
|
+
```python
|
|
426
|
+
from pidriver import (
|
|
427
|
+
PiClient, PiConfig, AskHost,
|
|
428
|
+
MessageDelta, InteractionRequest, InteractionKind, Decision, AgentEnd,
|
|
429
|
+
)
|
|
430
|
+
|
|
431
|
+
client = PiClient(PiConfig(binary="omp", provider="openai", api_key=KEY))
|
|
432
|
+
session = await client.start("Refactor utils.py", cwd=project, interaction_handler=AskHost())
|
|
433
|
+
|
|
434
|
+
async with session:
|
|
435
|
+
async for event in session:
|
|
436
|
+
match event:
|
|
437
|
+
case MessageDelta(text=text):
|
|
438
|
+
print(text, end="", flush=True)
|
|
439
|
+
|
|
440
|
+
case InteractionRequest(kind=InteractionKind.PERMISSION) as req:
|
|
441
|
+
ok = input(f"\nAllow? {req.prompt} [y/N] ").lower() == "y"
|
|
442
|
+
await session.respond(req.request_id, Decision.ALLOW if ok else Decision.DENY)
|
|
443
|
+
|
|
444
|
+
case InteractionRequest(kind=InteractionKind.QUESTION) as req:
|
|
445
|
+
await session.respond(req.request_id, input(f"\n{req.prompt} "))
|
|
446
|
+
|
|
447
|
+
case InteractionRequest(kind=InteractionKind.CHOICE) as req:
|
|
448
|
+
for i, opt in enumerate(req.options):
|
|
449
|
+
print(f" {i}. {opt}")
|
|
450
|
+
await session.respond(req.request_id, int(input("> "))) # answer by index
|
|
451
|
+
|
|
452
|
+
case AgentEnd(reason=reason):
|
|
453
|
+
print(f"\n[{reason}]")
|
|
454
|
+
```
|
|
455
|
+
|
|
456
|
+
(`AskHost` is the default, so you can omit `interaction_handler=` for this behavior.)
|
|
457
|
+
|
|
458
|
+
**AutoApprove — autonomous.** The handler answers permission prompts itself, so an unattended run
|
|
459
|
+
just consumes output. Free-text questions and choices are still deferred — if your task might
|
|
460
|
+
trigger them, handle `InteractionRequest` in the loop too, or combine handlers with `chain`:
|
|
461
|
+
|
|
462
|
+
```python
|
|
463
|
+
from pidriver import PiClient, PiConfig, AutoApprove, DenyAll, chain, AgentEnd
|
|
464
|
+
|
|
465
|
+
session = await client.start("Run the test suite and fix failures", cwd=project,
|
|
466
|
+
interaction_handler=AutoApprove()) # or AutoApprove(always=True)
|
|
467
|
+
async with session:
|
|
468
|
+
async for event in session:
|
|
469
|
+
if isinstance(event, AgentEnd):
|
|
470
|
+
print(event.reason)
|
|
471
|
+
|
|
472
|
+
# Approve only safe (read-only) tools, deny the rest:
|
|
473
|
+
handler = chain(
|
|
474
|
+
AutoApprove(allow=lambda r: "read" in r.prompt.lower()),
|
|
475
|
+
DenyAll(),
|
|
476
|
+
)
|
|
477
|
+
```
|
|
478
|
+
|
|
479
|
+
## Development
|
|
480
|
+
|
|
481
|
+
```sh
|
|
482
|
+
uv pip install -e ".[dev]"
|
|
483
|
+
uv run pytest
|
|
484
|
+
uv run mypy src
|
|
485
|
+
uv run ruff check
|
|
486
|
+
```
|
|
487
|
+
|
|
488
|
+
## License
|
|
489
|
+
|
|
490
|
+
MIT © Grigory Bakunov
|