mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
feat(harness): add pi harness for litellm.agent() (#45467)
* feat(harness): add pi harness for litellm.agent()
Adds Harness.PI, which drives the pi coding agent CLI (`pi --mode json`) once per turn
against the per-session model endpoint. The endpoint is declared as a `litellm`
openai-completions provider in the models.json of a LiteLLM-owned PI_CODING_AGENT_DIR, so
the user's own pi settings, auth, extensions and MCP servers are never read, and the token
travels only through an env var referenced as `$LITELLM_HARNESS_TOKEN`
pi exits 0 after a failed provider call, so a turn fails when the last assistant message
ends with stopReason error/aborted. Permission modes map to a `--tools` allowlist,
disable_tools to `--exclude-tools`, and sessions resume with `--session` from a persisted
session dir. build_instructions and turn_prompt move from the OpenCode config into the
shared harness utils so both configs use them
Unit tests run on JSONL recorded from pi 1.1.0
* test(harness): align pi harness tests with the other CLI harness suites
Mirror the OpenCode suite's layout, fakes and helpers, and add the cases every CLI harness
covers: turn before start, early close kills the process, persist failure, endpoint model
fallback, long JSONL lines, skills without SKILL.md and a missing endpoint. Add an
endpoint_requests.jsonl fixture recording what pi sends to the session endpoint, and
re-record the turn fixtures in one run so their ids agree
* fix(harness): parametrize harness e2e only over harnesses with a model group
harness_params() iterated every Harness member, so Harness.TOOL_LOOP, which has no entry in
BINARIES or DEFAULT_MODEL_GROUPS, raised KeyError and the whole suite failed to collect
* feat(harness): accept pi config and MCP servers through PiOptions.config
PiOptions gains `config`, like the other harness options. `mcpServers` is written to the
private agent dir's mcp.json and every other key to its settings.json. Keys that would reroute
model calls, bypass permissions=, or load code as the host user (defaultProvider, defaultModel,
defaultTools, defaultProjectTrust, extensions, packages, ...) raise OptionsMismatch
A plain-name `--tools` allowlist drops MCP tools in pi 1.1.0, so when servers are configured
the allowlist also gets `mcp__*` plus the tool each exposure is reached through (`codemode` by
default, `tool_search` for `deferred`). disable_tools still excludes individual MCP tools
* refactor(harness): give harness options immutable mapping defaults
`field(default_factory=dict)` left every config/env default typed as
`dict[Unknown, Unknown]`, which basedpyright reports as
reportUnknownVariableType under strict mode: 9 errors across the five
options classes, two of them added by PiOptions. A shared module-level
`MappingProxyType({})` default is accepted by frozen dataclasses, matches
the no-mutation rule the rest of the harness follows, and drops the file
to zero errors so the rule shrinks against the merge base instead of
growing.
* fix(harness): reject managed env keys and a missing model in the pi harness
PiOptions.env could name PI_CODING_AGENT_DIR or LITELLM_HARNESS_TOKEN and
the value was dropped in silence, because transform_session_setup merges
the managed keys last. The ordering stays as defense in depth, but
validate_environment now rejects those keys up front the way the Claude
Code config does, so a bad option is an error rather than a setting that
appears to apply and does not.
transform_turn_request also stringified the model it resolved, so a call
with neither model= nor an endpoint model would have put the literal
"None" on pi's --model. Both call sites now go through session_model(),
which raises the same ValueError transform_session_setup already raised.
* test(harness): escape the regex dot in the pi skill manifest match
* fix(harness): keep harness options importable on Python 3.11
Python 3.11 rejects an unhashable dataclass default, and MappingProxyType
only became hashable in 3.12, so the shared MappingProxyType({}) defaults
made importing litellm.harness.options raise ValueError and broke every
harness on 3.11. The defaults now come from typed default factories that
still return read-only mappings
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* test(harness): address review feedback on the pi harness tests
Drop the endpoint-request fixture test, which only read a recorded file and
could not fail when PiHarnessConfig changed, along with its now unused
fixture. Remove the section-divider comments and a comment that restated
MANAGED_ENV_KEYS, and type the test helpers and fakes
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* test(harness): cover the last uncovered pi harness lines
Assert that a tool_execution_end event without a result yields an empty
ToolResult output, and that a PiOptions.config value JSON cannot encode
raises TypeError. The Python 3.11 default check now asserts on the list of
unhashable defaults instead of relying on hash() raising, which the
test-quality gate flagged as an assertion-free test
* test(harness): annotate harness option tests and stop rebinding helper params
Mark the option test locals Final and check read-only defaults through the
MutableMapping interface instead of writing into them. In the pi suite,
type the parametrized test arguments, and have started() and parse_all()
bind new names instead of reassigning their sandbox and state parameters
* test(harness): test harness option defaults by constructing the options
Replace the field-introspection check with one that builds each options
class twice and asserts every mapping default is empty, read-only and not
shared between instances. Sharing one MappingProxyType across instances is
what broke class creation on Python 3.11, and that sharing is observable on
every Python version, so the test fails on the old code without inspecting
dataclass declarations
* test(harness): stop mutating test helper inputs and assert whole values
make_ctx took model and endpoint out of **kwargs with pop, which mutated
its own parameter; they are explicit parameters now, with a frozen default
endpoint. assistant_end builds its message in one expression instead of
adding errorMessage afterwards. Two docstrings that explained an assertion
became that assertion's message, and the options test compares whole
values so a failure names the offending fields
* ci(unit): run the pi harness unit tests in the llms-h-to-z shard
main now lists every tests/unit/llms provider directory per shard, so the new
tests/unit/llms/pi directory was claimed by no shard and assert_ci_coverage
failed on both the shard and the test-file checks
* test(harness): mark pi test locals Final and freeze the test output model
Every single-assignment local and module constant in the pi suite is now
Final. The name-mapping loop became one comprehension compared as a whole
list, test_start_persists_sessions_dir takes the sandbox by index instead of
binding _ twice, and the Answer output model is frozen like every other
pydantic model
* fix(harness): send PiOptions.thinking to the gateway as reasoning_effort
pi clamps --thinking to "off" for a model that does not declare reasoning,
and the generated models.json only declared {"id": model}, so every
thinking level was dropped and the gateway never saw reasoning_effort.
xhigh and max are also only offered when thinkingLevelMap names them. The
model now declares reasoning plus that map whenever a level other than
"off" is requested, and nothing when thinking is unset, so pi does not send
its default effort on every call
The pi tests also drop **kwargs from make_ctx and started in favour of
explicit parameters, and annotate their helpers and the fake sandbox with
read-only types; the fake records calls by rebinding tuples instead of
appending to lists
* docs(harness): list pi among the harnesses in the module and agent() docstrings
* fix(harness): keep pi calls provider-agnostic and harness option defaults copyable
Declare supportsStore false so pi stops sending store, take contextWindow and maxTokens from the model map, treat a length stop or a silent exit as an error, and give every options class plain dict defaults again so deepcopy, pickle and asdict work
* fix(harness): look up pi gateway groups as litellm_proxy and ask pi for cache_control marks
A gateway session's model is the bare group name, so the model map lookup
now passes custom_llm_provider=litellm_proxy, the way the SDK prices it.
That keeps an unmapped group from printing get_llm_provider's provider
list banner into the caller's stdout.
pi only adds Anthropic cache_control marks for OpenRouter Anthropic models,
so every Claude call through the harness missed the prompt cache. Models
the map says support prompt caching now declare cacheControlFormat
anthropic, which LiteLLM translates for each provider.
* fix(harness): cap pi's maxTokens at its own default and fork sessions saved in another directory
pi sends maxTokens as max_completion_tokens on every call, and the gateway reserves that many
tokens against the key's TPM limit, so passing the model's full output cap (128000 for
claude-haiku-5-5) made keys with lower limits answer 429. maxTokens is now the lower of the
model's cap and pi's own 16384 default.
agent_resume() into another sandbox directory stopped at pi's interactive fork prompt. The
native session id now carries the directory pi saved the session in, and a resume from a
different directory passes --fork, so the conversation continues in the new sandbox.
* fix(harness): leave pi room for the thinking budget and price cached prompt tokens
pi's maxTokens now adds LiteLLM's thinking budget for the requested level to pi's 16384 default, within the model's output cap, so Anthropic budget thinking at max no longer answers 400. Harness turns now pass cache read and write tokens to cost_per_token, so cached prompts are priced at their cache rates on every harness.
* test(harness): cover cached prompt cost on the gateway and non-stream routes
* test(harness): mark the cost test constants Final
* test(harness): type the cached cost test helpers without Any
* fix(harness): price OpenAI cache write tokens in harness cost
---------
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
parent
66cd1a6516
commit
0ca29c7d84
31 changed files with 2056 additions and 74 deletions
2
.github/ci-coverage-allowlist.yml
vendored
2
.github/ci-coverage-allowlist.yml
vendored
|
|
@ -5,7 +5,7 @@ description: >-
|
|||
|
||||
test_paths:
|
||||
- reason: >-
|
||||
litellm.agent() end-to-end suite. It drives the real claude, codex and opencode CLIs and
|
||||
litellm.agent() end-to-end suite. It drives the real claude, codex, opencode and pi CLIs and
|
||||
deepagents against a live LiteLLM AI Gateway, so it needs those binaries on PATH plus
|
||||
LITELLM_PROXY_API_BASE / LITELLM_PROXY_API_KEY, and skips without them. Run manually
|
||||
before changing litellm/harness; the mocked coverage runs in tests/unit/harness and
|
||||
|
|
|
|||
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -322,6 +322,7 @@ jobs:
|
|||
tests/unit/llms/pass_through
|
||||
tests/unit/llms/perplexity
|
||||
tests/unit/llms/pg_vector
|
||||
tests/unit/llms/pi
|
||||
tests/unit/llms/publicai
|
||||
tests/unit/llms/ragflow
|
||||
tests/unit/llms/recraft
|
||||
|
|
|
|||
|
|
@ -287,7 +287,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
|
|||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode or Deep Agents on any model (Python SDK)</summary>
|
||||
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode, pi or Deep Agents on any model (Python SDK)</summary>
|
||||
|
||||
### Python SDK - Agents
|
||||
|
||||
|
|
@ -296,7 +296,7 @@ import litellm
|
|||
from litellm import Harness, sandbox
|
||||
|
||||
result = litellm.agent(
|
||||
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.DEEPAGENTS
|
||||
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.PI, Harness.DEEPAGENTS
|
||||
"Find why tests/test_router.py is flaky and fix it.",
|
||||
sandbox=sandbox.local("./repo"),
|
||||
model="litellm_proxy/claude-sonnet-4-5", # a model group on your AI Gateway
|
||||
|
|
@ -305,7 +305,7 @@ result = litellm.agent(
|
|||
print(result.text, result.cost, [f.path for f in result.files])
|
||||
```
|
||||
|
||||
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex` or `opencode`), or `deepagents langchain-litellm` for Deep Agents.
|
||||
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex`, `opencode` or `pi`), or `deepagents langchain-litellm` for Deep Agents.
|
||||
|
||||
[**Docs: Agent Harnesses**](https://docs.litellm.ai/docs/harness)
|
||||
|
||||
|
|
|
|||
|
|
@ -2385,6 +2385,7 @@ _AGENT_EXPORTS: Final = frozenset(
|
|||
"ClaudeCodeOptions",
|
||||
"CodexOptions",
|
||||
"OpenCodeOptions",
|
||||
"PiOptions",
|
||||
"DeepAgentsOptions",
|
||||
"ToolLoopOptions",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
"""Agent harnesses: run Claude Code, Codex, OpenCode, Deep Agents or Tool Loop on any LiteLLM model.
|
||||
"""Agent harnesses: run Claude Code, Codex, OpenCode, pi, Deep Agents or Tool Loop on any LiteLLM model.
|
||||
|
||||
The entrypoints live on the top-level package:
|
||||
|
||||
|
|
@ -30,6 +30,7 @@ from litellm.harness.options import (
|
|||
CodexOptions,
|
||||
DeepAgentsOptions,
|
||||
OpenCodeOptions,
|
||||
PiOptions,
|
||||
ToolLoopOptions,
|
||||
)
|
||||
from litellm.harness.runtime import (
|
||||
|
|
@ -78,6 +79,7 @@ __all__ = (
|
|||
"OpenCodeOptions",
|
||||
"OptionsMismatch",
|
||||
"OutputInvalid",
|
||||
"PiOptions",
|
||||
"Reasoning",
|
||||
"Result",
|
||||
"SandboxError",
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol
|
|||
|
||||
import httpx
|
||||
import openai
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
|
|
@ -172,16 +172,57 @@ def usage_from_mapping(usage: object) -> tuple[int, int]:
|
|||
return _as_int(input_tokens), _as_int(output_tokens)
|
||||
|
||||
|
||||
def usage_from_body(body: object) -> tuple[int, int]:
|
||||
"""Usage from a non-streaming JSON response body."""
|
||||
if not isinstance(body, Mapping):
|
||||
return 0, 0
|
||||
if isinstance(body.get("usage"), Mapping):
|
||||
return usage_from_mapping(body["usage"])
|
||||
response = body.get("response")
|
||||
if isinstance(response, Mapping):
|
||||
return usage_from_mapping(response.get("usage"))
|
||||
return 0, 0
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PromptCache:
|
||||
"""Prompt tokens read from or written to the provider's cache. Anthropic's usage counts them outside input_tokens,
|
||||
while OpenAI's prompt_tokens and the Responses API's input_tokens already include them."""
|
||||
|
||||
read_tokens: int = 0
|
||||
creation_tokens: int = 0
|
||||
outside_input: bool = False
|
||||
|
||||
def prompt_tokens(self, input_tokens: int) -> int:
|
||||
if self.outside_input:
|
||||
return input_tokens + self.read_tokens + self.creation_tokens
|
||||
return input_tokens
|
||||
|
||||
|
||||
NO_PROMPT_CACHE: Final = PromptCache()
|
||||
|
||||
|
||||
def json_object(value: object) -> Mapping[str, object]:
|
||||
try:
|
||||
return _JSON_OBJECT_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
def prompt_cache_from_mapping(usage: object) -> PromptCache:
|
||||
"""Cache token counts from a usage dict using OpenAI, Anthropic, or Responses field names."""
|
||||
fields: Final = json_object(usage)
|
||||
details: Final = json_object(fields.get("prompt_tokens_details", fields.get("input_tokens_details")))
|
||||
return PromptCache(
|
||||
read_tokens=_as_int(fields.get("cache_read_input_tokens") or details.get("cached_tokens")),
|
||||
creation_tokens=_as_int(
|
||||
fields.get("cache_creation_input_tokens")
|
||||
or details.get("cache_write_tokens")
|
||||
or details.get("cache_creation_tokens")
|
||||
),
|
||||
outside_input="prompt_tokens" not in fields and "input_tokens_details" not in fields,
|
||||
)
|
||||
|
||||
|
||||
def usage_mapping_of_body(body: object) -> Mapping[str, object]:
|
||||
fields: Final = json_object(body)
|
||||
usage: Final = json_object(fields.get("usage"))
|
||||
return usage if usage else json_object(json_object(fields.get("response")).get("usage"))
|
||||
|
||||
|
||||
def usage_mapping_of_json(raw: bytes) -> Mapping[str, object]:
|
||||
try:
|
||||
return usage_mapping_of_body(_JSON_VALUE_ADAPTER.validate_python(json.loads(raw)))
|
||||
except ValueError:
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
class SSEUsageParser:
|
||||
|
|
@ -190,6 +231,7 @@ class SSEUsageParser:
|
|||
def __init__(self) -> None:
|
||||
self.input_tokens = 0
|
||||
self.output_tokens = 0
|
||||
self.prompt_cache = NO_PROMPT_CACHE
|
||||
self._buffer = b""
|
||||
|
||||
def feed(self, chunk: bytes) -> None:
|
||||
|
|
@ -226,21 +268,23 @@ class SSEUsageParser:
|
|||
elif event_type == "response.completed":
|
||||
self._absorb_response_completed(event)
|
||||
elif isinstance(event.get("usage"), Mapping):
|
||||
self._set(*usage_from_mapping(event["usage"]))
|
||||
self._absorb_usage(event["usage"])
|
||||
|
||||
def _absorb_message_start(self, event: Mapping[str, object]) -> None:
|
||||
message = event.get("message")
|
||||
if isinstance(message, Mapping):
|
||||
self._set(*usage_from_mapping(message.get("usage")))
|
||||
self._absorb_usage(json_object(event.get("message")).get("usage"))
|
||||
|
||||
def _absorb_message_delta(self, event: Mapping[str, object]) -> None:
|
||||
# message_delta output_tokens is cumulative for the whole message.
|
||||
self._set(*usage_from_mapping(event.get("usage")))
|
||||
self._absorb_usage(event.get("usage"))
|
||||
|
||||
def _absorb_response_completed(self, event: Mapping[str, object]) -> None:
|
||||
response = event.get("response")
|
||||
if isinstance(response, Mapping):
|
||||
self._set(*usage_from_mapping(response.get("usage")))
|
||||
self._absorb_usage(json_object(event.get("response")).get("usage"))
|
||||
|
||||
def _absorb_usage(self, usage: object) -> None:
|
||||
self._set(*usage_from_mapping(usage))
|
||||
prompt_cache: Final = prompt_cache_from_mapping(usage)
|
||||
if prompt_cache.read_tokens or prompt_cache.creation_tokens:
|
||||
self.prompt_cache = prompt_cache
|
||||
|
||||
def _set(self, input_tokens: int, output_tokens: int) -> None:
|
||||
if input_tokens:
|
||||
|
|
@ -249,13 +293,19 @@ class SSEUsageParser:
|
|||
self.output_tokens = output_tokens
|
||||
|
||||
|
||||
def compute_cost(model: str | None, input_tokens: int, output_tokens: int) -> float:
|
||||
"""Cost from LiteLLM's price map. Never raises; unknown models cost 0.0."""
|
||||
def compute_cost(
|
||||
model: str | None, input_tokens: int, output_tokens: int, prompt_cache: PromptCache = NO_PROMPT_CACHE
|
||||
) -> float:
|
||||
"""Cost from LiteLLM's price map, cached prompt tokens at their cache prices. Never raises; unknown models cost 0.0."""
|
||||
if not model or not (input_tokens or output_tokens):
|
||||
return 0.0
|
||||
try:
|
||||
prompt_cost, completion_cost = litellm.cost_per_token(
|
||||
model=model, prompt_tokens=input_tokens, completion_tokens=output_tokens
|
||||
model=model,
|
||||
prompt_tokens=prompt_cache.prompt_tokens(input_tokens),
|
||||
completion_tokens=output_tokens,
|
||||
cache_read_input_tokens=prompt_cache.read_tokens,
|
||||
cache_creation_input_tokens=prompt_cache.creation_tokens,
|
||||
)
|
||||
return float(prompt_cost) + float(completion_cost)
|
||||
except Exception: # accounting must never break a call; the price-map lookup raises bare Exception
|
||||
|
|
@ -565,9 +615,10 @@ class ModelEndpoint:
|
|||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cost: float | None,
|
||||
prompt_cache: PromptCache,
|
||||
) -> None:
|
||||
if cost is None:
|
||||
cost = compute_cost(model, input_tokens, output_tokens)
|
||||
cost = compute_cost(model, input_tokens, output_tokens, prompt_cache)
|
||||
self.usage.add(input_tokens, output_tokens, cost)
|
||||
|
||||
async def _forward(self, request: Request, route: str, body: Mapping[str, object]) -> Response:
|
||||
|
|
@ -617,14 +668,12 @@ class ModelEndpoint:
|
|||
) -> None:
|
||||
if is_sse:
|
||||
parser.close()
|
||||
tokens = (parser.input_tokens, parser.output_tokens)
|
||||
else:
|
||||
try:
|
||||
body: Final[object] = _JSON_VALUE_ADAPTER.validate_python(json.loads(collected))
|
||||
tokens = usage_from_body(body)
|
||||
except ValueError:
|
||||
tokens = (0, 0)
|
||||
self._record(model, tokens[0], tokens[1], header_cost(upstream.headers))
|
||||
self._record(
|
||||
model, parser.input_tokens, parser.output_tokens, header_cost(upstream.headers), parser.prompt_cache
|
||||
)
|
||||
return
|
||||
usage: Final = usage_mapping_of_json(collected)
|
||||
self._record(model, *usage_from_mapping(usage), header_cost(upstream.headers), prompt_cache_from_mapping(usage))
|
||||
|
||||
def _sdk_kwargs(
|
||||
self, body: Mapping[str, object]
|
||||
|
|
@ -667,9 +716,9 @@ class ModelEndpoint:
|
|||
return self._responses.StreamingResponse(
|
||||
self._sdk_stream(route, response, model), media_type=SSE_MEDIA_TYPE
|
||||
)
|
||||
data = to_jsonable(response)
|
||||
input_tokens, output_tokens = usage_from_body(data)
|
||||
self._record(model, input_tokens, output_tokens, hidden_cost(response))
|
||||
data: Final = to_jsonable(response)
|
||||
usage: Final = usage_mapping_of_body(data)
|
||||
self._record(model, *usage_from_mapping(usage), hidden_cost(response), prompt_cache_from_mapping(usage))
|
||||
return self._json(data)
|
||||
|
||||
async def _sdk_stream(self, route: str, iterator: AsyncIterable[object], model: str | None) -> AsyncIterator[bytes]:
|
||||
|
|
@ -688,4 +737,4 @@ class ModelEndpoint:
|
|||
yield f"event: error\ndata: {json.dumps(error_body(e, message))}\n\n".encode()
|
||||
finally:
|
||||
parser.close()
|
||||
self._record(model, parser.input_tokens, parser.output_tokens, None)
|
||||
self._record(model, parser.input_tokens, parser.output_tokens, None, parser.prompt_cache)
|
||||
|
|
|
|||
|
|
@ -4,28 +4,37 @@ from __future__ import annotations
|
|||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Literal
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
PiThinkingLevel: TypeAlias = Literal["off", "minimal", "low", "medium", "high", "xhigh", "max"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClaudeCodeOptions:
|
||||
config: Mapping[str, object] = field(default_factory=dict)
|
||||
env: Mapping[str, str] = field(default_factory=dict)
|
||||
config: Mapping[str, object] = field(default_factory=dict[str, object])
|
||||
env: Mapping[str, str] = field(default_factory=dict[str, str])
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CodexOptions:
|
||||
reasoning_effort: Literal["low", "medium", "high", "xhigh"] | None = None
|
||||
web_search: bool = False
|
||||
config: Mapping[str, object] = field(default_factory=dict)
|
||||
env: Mapping[str, str] = field(default_factory=dict)
|
||||
config: Mapping[str, object] = field(default_factory=dict[str, object])
|
||||
env: Mapping[str, str] = field(default_factory=dict[str, str])
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OpenCodeOptions:
|
||||
agent: str = "build"
|
||||
config: Mapping[str, object] = field(default_factory=dict)
|
||||
env: Mapping[str, str] = field(default_factory=dict)
|
||||
config: Mapping[str, object] = field(default_factory=dict[str, object])
|
||||
env: Mapping[str, str] = field(default_factory=dict[str, str])
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PiOptions:
|
||||
thinking: PiThinkingLevel | None = None
|
||||
config: Mapping[str, object] = field(default_factory=dict[str, object])
|
||||
env: Mapping[str, str] = field(default_factory=dict[str, str])
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
|
@ -36,7 +45,9 @@ class DeepAgentsOptions:
|
|||
|
||||
@dataclass(frozen=True)
|
||||
class ToolLoopOptions:
|
||||
completion_kwargs: Mapping[str, object] = field(default_factory=dict)
|
||||
completion_kwargs: Mapping[str, object] = field(default_factory=dict[str, object])
|
||||
|
||||
|
||||
HarnessOptions = ClaudeCodeOptions | CodexOptions | OpenCodeOptions | DeepAgentsOptions | ToolLoopOptions
|
||||
HarnessOptions: TypeAlias = (
|
||||
ClaudeCodeOptions | CodexOptions | OpenCodeOptions | PiOptions | DeepAgentsOptions | ToolLoopOptions
|
||||
)
|
||||
|
|
|
|||
|
|
@ -411,7 +411,7 @@ def agent(
|
|||
options: HarnessOptions | None = None,
|
||||
install: bool = False,
|
||||
) -> Result | EventStream:
|
||||
"""Run an agent harness (Claude Code, Codex, OpenCode, Deep Agents, Tool Loop) on one prompt.
|
||||
"""Run an agent harness (Claude Code, Codex, OpenCode, pi, Deep Agents, Tool Loop) on one prompt.
|
||||
|
||||
Returns a Result. With stream=True it returns an iterator of events instead.
|
||||
Prefix the model with `litellm_proxy/` to route every model call through your
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ class Harness(Enum):
|
|||
CLAUDE_CODE = "claude_code"
|
||||
CODEX = "codex"
|
||||
OPENCODE = "opencode"
|
||||
PI = "pi"
|
||||
DEEPAGENTS = "deepagents"
|
||||
TOOL_LOOP = "tool_loop"
|
||||
|
||||
|
|
|
|||
|
|
@ -70,6 +70,22 @@ def structured_output_instruction(schema: Mapping[str, object]) -> str:
|
|||
)
|
||||
|
||||
|
||||
def build_instructions(ctx: SessionContext) -> str | None:
|
||||
"""System instructions plus the structured-output instruction, for runtimes that take one prompt file."""
|
||||
schema_part: Final = (
|
||||
structured_output_instruction(ctx.output.model_json_schema()) if ctx.output is not None else None
|
||||
)
|
||||
sections: Final = tuple(section for section in (ctx.instructions, schema_part) if section)
|
||||
return "\n\n".join(sections) if sections else None
|
||||
|
||||
|
||||
def turn_prompt(ctx: SessionContext, prompt: str) -> str:
|
||||
"""Repeat the schema instruction in the user turn; system instructions alone are too weak."""
|
||||
if ctx.output is None:
|
||||
return prompt
|
||||
return f"{prompt}\n\n{structured_output_instruction(ctx.output.model_json_schema())}"
|
||||
|
||||
|
||||
def strict_json_schema(schema: JSONValue, depth: int = 0) -> JSONValue:
|
||||
"""Make a JSON schema acceptable to OpenAI strict structured outputs.
|
||||
|
||||
|
|
|
|||
|
|
@ -40,11 +40,12 @@ from litellm.llms.base_llm.harness.transformation import (
|
|||
event_list,
|
||||
)
|
||||
from litellm.llms.base_llm.harness.utils import (
|
||||
build_instructions,
|
||||
last_json_object,
|
||||
native_tool_names,
|
||||
normalize_tool_name,
|
||||
stderr_tail_text,
|
||||
structured_output_instruction,
|
||||
turn_prompt,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -274,21 +275,6 @@ def build_opencode_config(
|
|||
return {**user, **managed, **dict(present)} # mutable-ok: opencode config JSON
|
||||
|
||||
|
||||
def build_instructions(ctx: SessionContext) -> str | None:
|
||||
schema_part: Final = (
|
||||
structured_output_instruction(ctx.output.model_json_schema()) if ctx.output is not None else None
|
||||
)
|
||||
sections: Final = tuple(section for section in (ctx.instructions, schema_part) if section)
|
||||
return "\n\n".join(sections) if sections else None
|
||||
|
||||
|
||||
def turn_prompt(ctx: SessionContext, prompt: str) -> str:
|
||||
"""Repeat the schema instruction in the user turn; system instructions alone are too weak."""
|
||||
if ctx.output is None:
|
||||
return prompt
|
||||
return f"{prompt}\n\n{structured_output_instruction(ctx.output.model_json_schema())}"
|
||||
|
||||
|
||||
class OpenCodeHarnessConfig(BaseCLIHarnessConfig):
|
||||
harness = Harness.OPENCODE
|
||||
options_type = OpenCodeOptions
|
||||
|
|
|
|||
0
litellm/llms/pi/__init__.py
Normal file
0
litellm/llms/pi/__init__.py
Normal file
0
litellm/llms/pi/harness/__init__.py
Normal file
0
litellm/llms/pi/harness/__init__.py
Normal file
569
litellm/llms/pi/harness/transformation.py
Normal file
569
litellm/llms/pi/harness/transformation.py
Normal file
|
|
@ -0,0 +1,569 @@
|
|||
"""
|
||||
pi harness config: `pi --mode json` (JSONL session events), once per turn.
|
||||
|
||||
Every model call goes to one custom provider (`litellm`, api=openai-completions) declared in
|
||||
the models.json of a LiteLLM-owned PI_CODING_AGENT_DIR, so the user's own pi settings, auth,
|
||||
extensions, MCP servers and skills are never read. PiOptions.config is pi's own config:
|
||||
`mcpServers` goes to that dir's mcp.json and every other key to its settings.json, minus the
|
||||
keys LiteLLM manages. The provider's apiKey is the `$`-reference
|
||||
pi resolves from LITELLM_HARNESS_TOKEN, so the token is never in argv or on disk. pi exits 0
|
||||
after a failed provider call, so failure is read from the last assistant message instead.
|
||||
Verified against @earendil-works/pi-coding-agent 1.1.0 and 1.3.0.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.constants import (
|
||||
ANTHROPIC_MIN_THINKING_BUDGET_TOKENS,
|
||||
DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
|
||||
DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET,
|
||||
DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET,
|
||||
DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
|
||||
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET,
|
||||
DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
|
||||
)
|
||||
from litellm.exceptions import ModelNotMappedError
|
||||
from litellm.harness.errors import CapabilityUnsupported, HarnessError, OptionsMismatch
|
||||
from litellm.harness.options import PiOptions, PiThinkingLevel
|
||||
from litellm.harness.types import (
|
||||
Capabilities,
|
||||
Event,
|
||||
Harness,
|
||||
PermissionMode,
|
||||
Reasoning,
|
||||
Text,
|
||||
ToolCall,
|
||||
ToolResult,
|
||||
)
|
||||
from litellm.llms.base_llm.harness.transformation import (
|
||||
BaseCLIHarnessConfig,
|
||||
HarnessSessionSetup,
|
||||
HarnessTurnError,
|
||||
HarnessTurnRequest,
|
||||
HarnessTurnResponse,
|
||||
event_list,
|
||||
)
|
||||
from litellm.llms.base_llm.harness.utils import (
|
||||
build_instructions,
|
||||
last_json_object,
|
||||
native_tool_names,
|
||||
normalize_tool_name,
|
||||
stderr_tail_text,
|
||||
turn_prompt,
|
||||
)
|
||||
from litellm.types.utils import ModelInfo
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.harness.context import SessionContext
|
||||
|
||||
PI_BINARY: Final = "pi"
|
||||
PI_PROVIDER_ID: Final = "litellm"
|
||||
PI_TOKEN_ENV: Final = "LITELLM_HARNESS_TOKEN"
|
||||
PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR"
|
||||
AGENT_DIRNAME: Final = "agent"
|
||||
SESSIONS_DIRNAME: Final = "sessions"
|
||||
SKILLS_DIRNAME: Final = "skills"
|
||||
INSTRUCTIONS_FILENAME: Final = "instructions.md"
|
||||
MODELS_FILENAME: Final = f"{AGENT_DIRNAME}/models.json"
|
||||
SETTINGS_FILENAME: Final = f"{AGENT_DIRNAME}/settings.json"
|
||||
MCP_FILENAME: Final = f"{AGENT_DIRNAME}/mcp.json"
|
||||
MCP_CONFIG_KEY: Final = "mcpServers"
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object])
|
||||
MCP_TOOL_PATTERN: Final = "mcp__*"
|
||||
EXPOSURE_TOOLS: Final = (("codemode", "codemode"), ("codemode-deferred", "codemode"), ("deferred", "tool_search"))
|
||||
|
||||
# settings.json keys that would reroute model calls, bypass permissions=, or load code as the host user.
|
||||
MANAGED_CONFIG_KEYS: Final = frozenset(
|
||||
{
|
||||
"defaultProvider",
|
||||
"defaultModel",
|
||||
"enabledModels",
|
||||
"defaultThinkingLevel",
|
||||
"defaultTools",
|
||||
"defaultProjectTrust",
|
||||
"sessionDir",
|
||||
"httpProxy",
|
||||
"packages",
|
||||
"extensions",
|
||||
"skills",
|
||||
}
|
||||
)
|
||||
|
||||
MANAGED_ENV_KEYS: Final = frozenset({PI_CONFIG_DIR_ENV, PI_TOKEN_ENV})
|
||||
|
||||
PI_ISOLATION_ENV: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"PI_OFFLINE": "1",
|
||||
"PI_SKIP_VERSION_CHECK": "1",
|
||||
"PI_TELEMETRY": "0",
|
||||
}
|
||||
)
|
||||
|
||||
PERMISSION_TOOLS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType(
|
||||
{
|
||||
"read-only": ("read", "grep", "find", "ls"),
|
||||
"edit": ("read", "edit", "write", "grep", "find", "ls"),
|
||||
"full": ("read", "bash", "edit", "write", "grep", "find", "ls"),
|
||||
}
|
||||
)
|
||||
|
||||
NORMALIZED_TO_NATIVE: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType(
|
||||
{
|
||||
"read": ("read",),
|
||||
"write": ("write",),
|
||||
"edit": ("edit",),
|
||||
"bash": ("bash", "powershell"),
|
||||
"glob": ("find",),
|
||||
"grep": ("grep",),
|
||||
"ls": ("ls",),
|
||||
}
|
||||
)
|
||||
|
||||
NATIVE_TO_NORMALIZED: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"read": "read",
|
||||
"write": "write",
|
||||
"edit": "edit",
|
||||
"bash": "bash",
|
||||
"powershell": "bash",
|
||||
"find": "glob",
|
||||
"grep": "grep",
|
||||
"ls": "ls",
|
||||
}
|
||||
)
|
||||
|
||||
PI_THINKING_LEVEL_MAP: Final[Mapping[str, str]] = MappingProxyType({"xhigh": "xhigh", "max": "max"})
|
||||
|
||||
# pi adds `store: false` to every call otherwise, which LiteLLM rejects for providers that have no `store`.
|
||||
PI_PROVIDER_COMPAT: Final[Mapping[str, object]] = MappingProxyType({"supportsStore": False})
|
||||
|
||||
PI_CACHE_CONTROL_COMPAT: Final[Mapping[str, str]] = MappingProxyType({"cacheControlFormat": "anthropic"})
|
||||
|
||||
PI_DEFAULT_MAX_TOKENS: Final = 16384
|
||||
|
||||
LITELLM_THINKING_BUDGETS: Final[Mapping[str, int]] = MappingProxyType(
|
||||
{
|
||||
"minimal": max(DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET, ANTHROPIC_MIN_THINKING_BUDGET_TOKENS),
|
||||
"low": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET,
|
||||
"medium": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
|
||||
"high": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
|
||||
"xhigh": DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
|
||||
"max": DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET,
|
||||
}
|
||||
)
|
||||
|
||||
SESSION_CWD_SEPARATOR: Final = "@"
|
||||
|
||||
PI_BUILTIN_TOOLS: Final = frozenset({*NATIVE_TO_NORMALIZED, "codemode", "tool_search"})
|
||||
FAILED_STOP_REASONS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"error": "the model call failed",
|
||||
"aborted": "aborted",
|
||||
"length": "the reply hit the model's output token limit",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _Frozen(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
|
||||
class _McpServer(_Frozen):
|
||||
exposure: str = "codemode"
|
||||
tool_exposure: Mapping[str, str] = Field(default_factory=lambda: MappingProxyType({}), alias="toolExposure")
|
||||
|
||||
def exposures(self) -> frozenset[str]:
|
||||
return frozenset({self.exposure, *self.tool_exposure.values()})
|
||||
|
||||
|
||||
_MCP_SERVERS: Final = TypeAdapter(Mapping[str, _McpServer])
|
||||
|
||||
|
||||
class _ContentBlock(_Frozen):
|
||||
type: str
|
||||
text: str = ""
|
||||
|
||||
|
||||
class _Role(_Frozen):
|
||||
role: str = ""
|
||||
|
||||
|
||||
class _AssistantMessage(_Frozen):
|
||||
content: tuple[_ContentBlock, ...] = ()
|
||||
stop_reason: str | None = Field(default=None, alias="stopReason")
|
||||
error_message: str | None = Field(default=None, alias="errorMessage")
|
||||
|
||||
|
||||
class _ToolOutput(_Frozen):
|
||||
content: tuple[_ContentBlock, ...] = ()
|
||||
|
||||
|
||||
class _ToolStart(_Frozen):
|
||||
tool_call_id: str = Field(alias="toolCallId")
|
||||
tool_name: str = Field(alias="toolName")
|
||||
args: Mapping[str, object] = Field(default_factory=lambda: MappingProxyType({}))
|
||||
|
||||
|
||||
class _ToolEnd(_Frozen):
|
||||
tool_call_id: str = Field(alias="toolCallId")
|
||||
result: _ToolOutput | str | None = None
|
||||
is_error: bool = Field(default=False, alias="isError")
|
||||
|
||||
|
||||
class _MessageUpdate(_Frozen):
|
||||
type: str
|
||||
delta: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class PiStreamState:
|
||||
"""What the parser has learned from one `pi --mode json` run."""
|
||||
|
||||
session_id: str | None = None
|
||||
session_cwd: str | None = None
|
||||
final_text: str = ""
|
||||
stop_reason: str | None = None
|
||||
error: str | None = None
|
||||
|
||||
def record_session(self, session_id: object, cwd: object) -> None:
|
||||
if isinstance(session_id, str) and self.session_id is None:
|
||||
self.session_id = session_id
|
||||
self.session_cwd = cwd if isinstance(cwd, str) else None
|
||||
|
||||
def record_message_end(self, message: object) -> None:
|
||||
if _Role.model_validate(message).role != "assistant":
|
||||
return
|
||||
assistant: Final = _AssistantMessage.model_validate(message)
|
||||
self.final_text = _text_of(assistant.content)
|
||||
self.stop_reason = assistant.stop_reason
|
||||
self.error = assistant.error_message
|
||||
|
||||
def mark_aborted(self) -> None:
|
||||
self.stop_reason = "aborted"
|
||||
|
||||
|
||||
def _text_of(blocks: Sequence[_ContentBlock]) -> str:
|
||||
return "".join(block.text for block in blocks if block.type == "text")
|
||||
|
||||
|
||||
def _tool_output_text(result: _ToolOutput | str | None) -> str:
|
||||
if result is None:
|
||||
return ""
|
||||
if isinstance(result, str):
|
||||
return result
|
||||
return _text_of(result.content)
|
||||
|
||||
|
||||
def _message_update_events(update: object) -> Sequence[Event]:
|
||||
try:
|
||||
event: Final = _MessageUpdate.model_validate(update)
|
||||
except ValidationError:
|
||||
return event_list()
|
||||
if not event.delta:
|
||||
return event_list()
|
||||
if event.type == "text_delta":
|
||||
return event_list(Text(delta=event.delta))
|
||||
if event.type == "thinking_delta":
|
||||
return event_list(Reasoning(delta=event.delta))
|
||||
return event_list()
|
||||
|
||||
|
||||
def _tool_start_events(line: Mapping[str, object]) -> Sequence[Event]:
|
||||
start: Final = _ToolStart.model_validate(line)
|
||||
return event_list(
|
||||
ToolCall(
|
||||
id=start.tool_call_id,
|
||||
name=normalize_tool_name(start.tool_name, NATIVE_TO_NORMALIZED),
|
||||
native_name=start.tool_name,
|
||||
input=start.args,
|
||||
builtin=start.tool_name in PI_BUILTIN_TOOLS,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _tool_end_events(line: Mapping[str, object]) -> Sequence[Event]:
|
||||
end: Final = _ToolEnd.model_validate(line)
|
||||
return event_list(ToolResult(id=end.tool_call_id, output=_tool_output_text(end.result), is_error=end.is_error))
|
||||
|
||||
|
||||
def tool_args(
|
||||
permissions: PermissionMode, disable_tools: Sequence[str], mcp_tools: Sequence[str] = ()
|
||||
) -> tuple[str, ...]:
|
||||
"""`--tools` allowlist for a permission mode and MCP servers, plus `--exclude-tools` for disable_tools."""
|
||||
if permissions not in PERMISSION_TOOLS:
|
||||
raise CapabilityUnsupported(
|
||||
f"Harness.PI does not support permissions={permissions!r} (supported: {sorted(PERMISSION_TOOLS)})"
|
||||
)
|
||||
denied: Final = native_tool_names(disable_tools, NORMALIZED_TO_NATIVE)
|
||||
exclude: Final = ("--exclude-tools", ",".join(denied)) if denied else ()
|
||||
return ("--tools", ",".join((*PERMISSION_TOOLS[permissions], *mcp_tools)), *exclude)
|
||||
|
||||
|
||||
def session_args(native_session_id: str | None, workdir: str) -> tuple[str, ...]:
|
||||
"""`--session` for a session pi saved in this workdir, `--fork` for one saved in another, since pi otherwise
|
||||
stops at an interactive "Fork this session into current directory?" prompt the harness cannot answer."""
|
||||
if native_session_id is None:
|
||||
return ()
|
||||
session_id, _, saved_cwd = native_session_id.partition(SESSION_CWD_SEPARATOR)
|
||||
return ("--fork" if saved_cwd and saved_cwd != workdir else "--session", session_id)
|
||||
|
||||
|
||||
def mcp_tool_entries(config: Mapping[str, object]) -> tuple[str, ...]:
|
||||
"""`--tools` entries that keep configured MCP servers reachable.
|
||||
|
||||
A plain-name `--tools` list drops MCP tools in pi 1.1.0, so `mcp__*` is listed whenever servers
|
||||
are configured, plus the tool each exposure is reached through: `codemode` (pi's default) or
|
||||
`tool_search` for `deferred`.
|
||||
"""
|
||||
servers: Final = _MCP_SERVERS.validate_python(config.get(MCP_CONFIG_KEY) or {})
|
||||
if not servers:
|
||||
return ()
|
||||
exposures: Final = frozenset(itertools.chain.from_iterable(server.exposures() for server in servers.values()))
|
||||
reach: Final = tuple(tool for exposure, tool in EXPOSURE_TOOLS if exposure in exposures)
|
||||
return (MCP_TOOL_PATTERN, *reach)
|
||||
|
||||
|
||||
def session_model(ctx: SessionContext) -> str:
|
||||
"""The model pi is pointed at, from agent(model=) or the session endpoint."""
|
||||
model: Final = ctx.model or (ctx.endpoint.model if ctx.endpoint else None)
|
||||
if not model:
|
||||
raise ValueError("Harness.PI needs model= (a gateway model group or litellm model)")
|
||||
return model
|
||||
|
||||
|
||||
def mapped_model_info(model: str, through_gateway: bool) -> ModelInfo | None:
|
||||
"""LiteLLM's model map entry. A gateway model group is looked up as litellm_proxy, the way the SDK prices it,
|
||||
which also keeps an unmapped group from printing get_llm_provider's provider-list banner on stdout."""
|
||||
try:
|
||||
return get_model_info(model, custom_llm_provider="litellm_proxy" if through_gateway else None)
|
||||
except ModelNotMappedError:
|
||||
return None
|
||||
|
||||
|
||||
def max_output_tokens(info: ModelInfo | None, thinking: PiThinkingLevel | None) -> int | None:
|
||||
"""pi's own 16384 default plus the thinking budget LiteLLM sends for the level, within the model's output cap.
|
||||
pi sends maxTokens as max_completion_tokens on every call and the gateway reserves that many tokens against the
|
||||
key's TPM limit, so the model's full output cap would starve keys with lower limits, while Anthropic budget
|
||||
thinking needs max_tokens above the budget, which is how pi sizes its own Anthropic calls."""
|
||||
budget: Final = LITELLM_THINKING_BUDGETS.get(thinking, 0) if thinking else 0
|
||||
max_output: Final = info.get("max_output_tokens") if info is not None else None
|
||||
if isinstance(max_output, int):
|
||||
return min(max_output, PI_DEFAULT_MAX_TOKENS + budget)
|
||||
return PI_DEFAULT_MAX_TOKENS + budget if budget else None
|
||||
|
||||
|
||||
def model_limits(info: ModelInfo | None, thinking: PiThinkingLevel | None) -> Mapping[str, int]:
|
||||
"""contextWindow from LiteLLM's model map, since pi assumes 128000 otherwise, and maxTokens per max_output_tokens."""
|
||||
limits: Final = (
|
||||
("contextWindow", info.get("max_input_tokens") if info is not None else None),
|
||||
("maxTokens", max_output_tokens(info, thinking)),
|
||||
)
|
||||
return MappingProxyType({key: value for key, value in limits if isinstance(value, int)})
|
||||
|
||||
|
||||
def model_compat(info: ModelInfo | None) -> Mapping[str, object]:
|
||||
"""Anthropic cache_control marks, which LiteLLM translates per provider, for models that support prompt caching."""
|
||||
caches: Final = info is not None and info.get("supports_prompt_caching") is True
|
||||
return MappingProxyType({"compat": PI_CACHE_CONTROL_COMPAT} if caches else {})
|
||||
|
||||
|
||||
def build_models_json(model: str, base_url: str, thinking: PiThinkingLevel | None, info: ModelInfo | None) -> str:
|
||||
# pi clamps --thinking to "off" unless the model declares reasoning, and offers xhigh/max only when
|
||||
# thinkingLevelMap names them. Declaring reasoning without a requested level would make pi send its
|
||||
# default reasoning_effort on every call, so it is declared only when a level is asked for
|
||||
reasoning: Final = (
|
||||
MappingProxyType({"reasoning": True, "thinkingLevelMap": PI_THINKING_LEVEL_MAP})
|
||||
if thinking is not None and thinking != "off"
|
||||
else MappingProxyType({})
|
||||
)
|
||||
provider: Final = MappingProxyType(
|
||||
{
|
||||
"baseUrl": base_url,
|
||||
"api": "openai-completions",
|
||||
"apiKey": f"${PI_TOKEN_ENV}",
|
||||
"compat": PI_PROVIDER_COMPAT,
|
||||
"models": (
|
||||
MappingProxyType({"id": model, **model_limits(info, thinking), **model_compat(info), **reasoning}),
|
||||
),
|
||||
}
|
||||
)
|
||||
return json.dumps({"providers": {PI_PROVIDER_ID: provider}}, default=_plain_json)
|
||||
|
||||
|
||||
def validate_user_config(config: Mapping[str, object]) -> None:
|
||||
"""Reject PiOptions.config keys LiteLLM manages, and a malformed mcpServers."""
|
||||
managed: Final = sorted(MANAGED_CONFIG_KEYS.intersection(config))
|
||||
if managed:
|
||||
raise OptionsMismatch(
|
||||
f"PiOptions.config may not set {', '.join(managed)}; LiteLLM manages it. Use the matching "
|
||||
"agent() argument (model=, permissions=, disable_tools=, skills=) or PiOptions.thinking instead"
|
||||
)
|
||||
try:
|
||||
_MCP_SERVERS.validate_python(config.get(MCP_CONFIG_KEY) or {})
|
||||
except ValidationError as e:
|
||||
raise OptionsMismatch(
|
||||
f"PiOptions.config[{MCP_CONFIG_KEY!r}] must map each server name to its config: {e}"
|
||||
) from e
|
||||
|
||||
|
||||
def _plain_json(value: object) -> object:
|
||||
if isinstance(value, Mapping):
|
||||
return _JSON_OBJECT.validate_python(value)
|
||||
raise TypeError(f"{type(value).__name__} in PiOptions.config is not JSON serializable")
|
||||
|
||||
|
||||
def config_files(config: Mapping[str, object]) -> tuple[tuple[str, bytes], ...]:
|
||||
"""settings.json for every key but mcpServers, and mcp.json for mcpServers; each only when non-empty."""
|
||||
settings: Final = MappingProxyType({key: value for key, value in config.items() if key != MCP_CONFIG_KEY})
|
||||
servers: Final = config.get(MCP_CONFIG_KEY)
|
||||
settings_file: Final = ((SETTINGS_FILENAME, json.dumps(settings, default=_plain_json)),) if settings else ()
|
||||
mcp_file: Final = ((MCP_FILENAME, json.dumps({MCP_CONFIG_KEY: servers}, default=_plain_json)),) if servers else ()
|
||||
return tuple((path, text.encode("utf-8")) for path, text in (*settings_file, *mcp_file))
|
||||
|
||||
|
||||
class PiHarnessConfig(BaseCLIHarnessConfig[PiOptions, PiStreamState]):
|
||||
harness = Harness.PI
|
||||
options_type = PiOptions
|
||||
capabilities = Capabilities(
|
||||
structured_output=True,
|
||||
tool_approval=False,
|
||||
tool_filtering=True,
|
||||
history=False,
|
||||
custom_tools=False,
|
||||
skills=True,
|
||||
resume=True,
|
||||
permission_modes=frozenset(PERMISSION_TOOLS),
|
||||
)
|
||||
|
||||
def get_binary(self) -> str:
|
||||
return PI_BINARY
|
||||
|
||||
def get_install_hint(self) -> str:
|
||||
return "npm install -g @earendil-works/pi-coding-agent"
|
||||
|
||||
def validate_environment(self, ctx: SessionContext) -> None:
|
||||
options: Final = self.get_options(ctx)
|
||||
clashing: Final = sorted(MANAGED_ENV_KEYS.intersection(options.env))
|
||||
if clashing:
|
||||
raise OptionsMismatch(f"PiOptions.env may not set {', '.join(clashing)}; LiteLLM manages it")
|
||||
validate_user_config(options.config)
|
||||
tool_args(ctx.permissions, ctx.disable_tools, mcp_tool_entries(options.config))
|
||||
|
||||
def transform_session_setup(self, ctx: SessionContext, private_dir: str) -> HarnessSessionSetup:
|
||||
if ctx.endpoint is None:
|
||||
raise HarnessError("pi needs the session model endpoint")
|
||||
model: Final = session_model(ctx)
|
||||
options: Final = self.get_options(ctx)
|
||||
base_url: Final = ctx.sandbox.host_url(ctx.endpoint.port).rstrip("/") + "/v1"
|
||||
info: Final = mapped_model_info(model, through_gateway=ctx.gateway is not None)
|
||||
models_file: Final = (
|
||||
MODELS_FILENAME,
|
||||
build_models_json(model, base_url, options.thinking, info).encode("utf-8"),
|
||||
)
|
||||
instructions: Final = build_instructions(ctx)
|
||||
instructions_file: Final = (
|
||||
((INSTRUCTIONS_FILENAME, instructions.encode("utf-8")),) if instructions is not None else ()
|
||||
)
|
||||
return HarnessSessionSetup(
|
||||
files=MappingProxyType(dict((models_file, *config_files(options.config), *instructions_file))),
|
||||
persisted_dirs=((SESSIONS_DIRNAME, "pi/sessions"),),
|
||||
skills_dir=SKILLS_DIRNAME,
|
||||
env=MappingProxyType(
|
||||
{
|
||||
**PI_ISOLATION_ENV,
|
||||
**options.env,
|
||||
PI_CONFIG_DIR_ENV: f"{private_dir}/{AGENT_DIRNAME}",
|
||||
PI_TOKEN_ENV: ctx.endpoint.token,
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
def transform_turn_request(
|
||||
self,
|
||||
ctx: SessionContext,
|
||||
setup: HarnessSessionSetup,
|
||||
private_dir: str,
|
||||
prompt: str,
|
||||
native_session_id: str | None,
|
||||
) -> HarnessTurnRequest:
|
||||
options: Final = self.get_options(ctx)
|
||||
# --no-approve: a repo's .pi/extensions would otherwise run as the host user at startup.
|
||||
# --no-skills: ~/.agents/skills is discovered outside PI_CODING_AGENT_DIR.
|
||||
argv: Final = (
|
||||
PI_BINARY,
|
||||
"--mode",
|
||||
"json",
|
||||
"--no-approve",
|
||||
"--no-skills",
|
||||
"--provider",
|
||||
PI_PROVIDER_ID,
|
||||
"--model",
|
||||
session_model(ctx),
|
||||
"--session-dir",
|
||||
f"{private_dir}/{SESSIONS_DIRNAME}",
|
||||
*session_args(native_session_id, ctx.sandbox.workdir),
|
||||
*tool_args(ctx.permissions, ctx.disable_tools, mcp_tool_entries(options.config)),
|
||||
*(("--thinking", options.thinking) if options.thinking else ()),
|
||||
*(("--skill", f"{private_dir}/{SKILLS_DIRNAME}") if ctx.skills else ()),
|
||||
*(
|
||||
("--append-system-prompt", f"{private_dir}/{INSTRUCTIONS_FILENAME}")
|
||||
if INSTRUCTIONS_FILENAME in setup.files
|
||||
else ()
|
||||
),
|
||||
)
|
||||
# The prompt goes on stdin; pi prepends piped stdin to the (empty) first message.
|
||||
return HarnessTurnRequest(argv=argv, env=setup.env, stdin=turn_prompt(ctx, prompt), cwd=ctx.sandbox.workdir)
|
||||
|
||||
def create_stream_state(self) -> PiStreamState:
|
||||
return PiStreamState()
|
||||
|
||||
def transform_stream_line(self, line: Mapping[str, object], state: PiStreamState) -> Sequence[Event]:
|
||||
"""Usage on message_update is ignored on purpose: the session endpoint accounts it."""
|
||||
event_type: Final = line.get("type")
|
||||
if event_type == "session":
|
||||
state.record_session(line.get("id"), line.get("cwd"))
|
||||
return event_list()
|
||||
if event_type == "message_update":
|
||||
return _message_update_events(line.get("assistantMessageEvent"))
|
||||
if event_type == "message_end":
|
||||
state.record_message_end(line.get("message"))
|
||||
return event_list()
|
||||
if event_type == "tool_execution_start":
|
||||
return _tool_start_events(line)
|
||||
if event_type == "tool_execution_end":
|
||||
return _tool_end_events(line)
|
||||
if event_type == "agent_settled" and line.get("aborted") is True:
|
||||
state.mark_aborted()
|
||||
return event_list()
|
||||
|
||||
def get_native_session_id(self, state: PiStreamState) -> str | None:
|
||||
if state.session_id is None or state.session_cwd is None:
|
||||
return state.session_id
|
||||
return f"{state.session_id}{SESSION_CWD_SEPARATOR}{state.session_cwd}"
|
||||
|
||||
def transform_turn_response(
|
||||
self,
|
||||
ctx: SessionContext,
|
||||
state: PiStreamState,
|
||||
exit_code: int,
|
||||
stderr_tail: Sequence[str],
|
||||
) -> HarnessTurnResponse:
|
||||
failure: Final = FAILED_STOP_REASONS.get(state.stop_reason or "")
|
||||
if failure is not None:
|
||||
raise HarnessTurnError(f"pi turn failed: {state.error or failure}")
|
||||
if exit_code != 0:
|
||||
raise HarnessTurnError(f"pi exited with code {exit_code}: {stderr_tail_text(stderr_tail) or 'no output'}")
|
||||
if state.stop_reason is None:
|
||||
raise HarnessTurnError(f"pi exited without a reply: {stderr_tail_text(stderr_tail) or 'no output'}")
|
||||
output_json: Final = last_json_object(state.final_text) if ctx.output is not None else None
|
||||
return HarnessTurnResponse(final_text=state.final_text, output_json=output_json)
|
||||
|
|
@ -10597,6 +10597,10 @@ class ProviderConfigManager:
|
|||
)
|
||||
|
||||
return OpenCodeHarnessConfig()
|
||||
if harness == _Harness.PI:
|
||||
from litellm.llms.pi.harness.transformation import PiHarnessConfig
|
||||
|
||||
return PiHarnessConfig()
|
||||
if harness == _Harness.DEEPAGENTS:
|
||||
from litellm.llms.deepagents.harness.transformation import (
|
||||
DeepAgentsHarnessConfig,
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ EXCLUDED_FOLDERS = {
|
|||
"claude_code",
|
||||
"codex",
|
||||
"opencode",
|
||||
"pi",
|
||||
"deepagents",
|
||||
"tool_loop",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
"""Fixtures for the litellm.agent() end-to-end tests.
|
||||
|
||||
These run the real harness runtimes (claude, codex, opencode, deepagents) against a real
|
||||
These run the real harness runtimes (claude, codex, opencode, pi, deepagents) against a real
|
||||
LiteLLM AI Gateway, routed with the `litellm_proxy/` model prefix. They skip unless
|
||||
LITELLM_PROXY_API_BASE and LITELLM_PROXY_API_KEY are set. Model groups can be overridden
|
||||
per harness with HARNESS_E2E_MODEL_<HARNESS>.
|
||||
|
|
@ -23,6 +23,7 @@ DEFAULT_MODEL_GROUPS = {
|
|||
Harness.CLAUDE_CODE: "claude-haiku-4-5-20251001",
|
||||
Harness.CODEX: "bedrock_mantle/openai.gpt-5.4",
|
||||
Harness.OPENCODE: "claude-haiku-4-5-20251001",
|
||||
Harness.PI: "claude-haiku-4-5-20251001",
|
||||
Harness.DEEPAGENTS: "claude-haiku-4-5-20251001",
|
||||
}
|
||||
|
||||
|
|
@ -30,6 +31,7 @@ BINARIES = {
|
|||
Harness.CLAUDE_CODE: "claude",
|
||||
Harness.CODEX: "codex",
|
||||
Harness.OPENCODE: "opencode",
|
||||
Harness.PI: "pi",
|
||||
}
|
||||
|
||||
requires_gateway = pytest.mark.skipif(
|
||||
|
|
@ -62,7 +64,7 @@ def harness_params() -> list:
|
|||
not harness_available(h), reason=f"{h.value} runtime not installed"
|
||||
),
|
||||
)
|
||||
for h in Harness
|
||||
for h in DEFAULT_MODEL_GROUPS
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import json
|
||||
import sys
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -13,13 +13,17 @@ from litellm.harness.endpoint import (
|
|||
SSEUsageParser,
|
||||
UsageTracker,
|
||||
compute_cost,
|
||||
usage_from_body,
|
||||
usage_from_mapping,
|
||||
usage_mapping_of_body,
|
||||
)
|
||||
from litellm.harness.errors import HarnessInstallFailed
|
||||
from litellm.harness.context import GatewayTarget
|
||||
from litellm.harness.types import Harness, Usage
|
||||
from litellm.types.utils import ModelResponse, ModelResponseStream
|
||||
|
||||
REAL_COST_PER_TOKEN: Final = litellm.cost_per_token
|
||||
CACHED_MODEL: Final = "claude-haiku-5-5"
|
||||
|
||||
GATEWAY = GatewayTarget(api_base="https://gw.example.com", api_key="sk-gateway-secret")
|
||||
|
||||
ANTHROPIC_SSE = (
|
||||
|
|
@ -82,7 +86,7 @@ def no_real_cost(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"cost_per_token",
|
||||
lambda model, prompt_tokens, completion_tokens: (
|
||||
lambda model, prompt_tokens, completion_tokens, **cache_tokens: (
|
||||
prompt_tokens * 0.001,
|
||||
completion_tokens * 0.002,
|
||||
),
|
||||
|
|
@ -294,6 +298,180 @@ async def test_sdk_messages_stream_handles_dicts_and_bytes(
|
|||
assert (ep.usage.input_tokens, ep.usage.output_tokens) == (4, 9)
|
||||
|
||||
|
||||
async def chunks_of(chunks: list[dict[str, object]]) -> AsyncIterator[dict[str, object]]:
|
||||
for chunk in chunks:
|
||||
yield chunk
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "sdk", "chunks"),
|
||||
[
|
||||
(
|
||||
"/v1/messages",
|
||||
(litellm.anthropic.messages, "acreate"),
|
||||
[
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"cache_read_input_tokens": 800,
|
||||
"cache_creation_input_tokens": 100,
|
||||
"output_tokens": 1,
|
||||
}
|
||||
},
|
||||
},
|
||||
{"type": "message_delta", "usage": {"output_tokens": 10}},
|
||||
],
|
||||
),
|
||||
(
|
||||
"/chat/completions",
|
||||
(litellm, "acompletion"),
|
||||
[
|
||||
{
|
||||
"choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": 1000,
|
||||
"completion_tokens": 10,
|
||||
"prompt_tokens_details": {"cached_tokens": 800, "cache_creation_tokens": 100},
|
||||
},
|
||||
}
|
||||
],
|
||||
),
|
||||
(
|
||||
"/v1/responses",
|
||||
(litellm, "aresponses"),
|
||||
[
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"usage": {
|
||||
"input_tokens": 1000,
|
||||
"output_tokens": 10,
|
||||
"input_tokens_details": {"cached_tokens": 800, "cache_creation_tokens": 100},
|
||||
}
|
||||
},
|
||||
}
|
||||
],
|
||||
),
|
||||
(
|
||||
"/chat/completions",
|
||||
(litellm, "acompletion"),
|
||||
[
|
||||
{
|
||||
"choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": 1000,
|
||||
"completion_tokens": 10,
|
||||
"prompt_tokens_details": {"cached_tokens": 800, "cache_write_tokens": 100},
|
||||
},
|
||||
}
|
||||
],
|
||||
),
|
||||
(
|
||||
"/v1/responses",
|
||||
(litellm, "aresponses"),
|
||||
[
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"usage": {
|
||||
"input_tokens": 1000,
|
||||
"output_tokens": 10,
|
||||
"input_tokens_details": {"cached_tokens": 800, "cache_write_tokens": 100},
|
||||
}
|
||||
},
|
||||
}
|
||||
],
|
||||
),
|
||||
],
|
||||
ids=["messages", "chat", "responses", "chat cache_write_tokens", "responses cache_write_tokens"],
|
||||
)
|
||||
async def test_sdk_stream_cost_prices_cached_prompt_tokens_like_litellm(
|
||||
monkeypatch: pytest.MonkeyPatch, path: str, sdk: tuple[object, str], chunks: list[dict[str, object]]
|
||||
) -> None:
|
||||
async def fake_sdk(**kwargs: object) -> AsyncIterator[dict[str, object]]:
|
||||
return chunks_of(chunks)
|
||||
|
||||
monkeypatch.setattr(litellm, "cost_per_token", REAL_COST_PER_TOKEN)
|
||||
monkeypatch.setattr(*sdk, fake_sdk)
|
||||
async with ModelEndpoint(Harness.PI, CACHED_MODEL, None) as ep:
|
||||
async with httpx.AsyncClient(base_url=ep.url) as client:
|
||||
await client.post(path, json={"stream": True}, headers=auth(ep))
|
||||
cached_price: Final = sum(
|
||||
REAL_COST_PER_TOKEN(
|
||||
model=CACHED_MODEL,
|
||||
prompt_tokens=1000,
|
||||
completion_tokens=10,
|
||||
usage_object=litellm.Usage(
|
||||
prompt_tokens=1000,
|
||||
completion_tokens=10,
|
||||
total_tokens=1010,
|
||||
prompt_tokens_details={"cached_tokens": 800, "cache_creation_tokens": 100},
|
||||
),
|
||||
)
|
||||
)
|
||||
assert cached_price < sum(REAL_COST_PER_TOKEN(model=CACHED_MODEL, prompt_tokens=1000, completion_tokens=10))
|
||||
assert ep.usage.cost == pytest.approx(cached_price), (
|
||||
"SDK mode priced cached prompt tokens at the full input price, overstating the cost of every cached turn"
|
||||
)
|
||||
|
||||
|
||||
ANTHROPIC_CACHED_USAGE: Final = {
|
||||
"input_tokens": 100,
|
||||
"cache_read_input_tokens": 800,
|
||||
"cache_creation_input_tokens": 100,
|
||||
"output_tokens": 10,
|
||||
}
|
||||
|
||||
|
||||
def anthropic_cached_sse() -> bytes:
|
||||
start: Final = {"type": "message_start", "message": {"usage": {**ANTHROPIC_CACHED_USAGE, "output_tokens": 1}}}
|
||||
delta: Final = {"type": "message_delta", "usage": {"output_tokens": 10}}
|
||||
return (
|
||||
f"event: message_start\ndata: {json.dumps(start)}\n\n"
|
||||
f"event: message_delta\ndata: {json.dumps(delta)}\n\n"
|
||||
).encode()
|
||||
|
||||
|
||||
async def anthropic_cached_message(**kwargs: object) -> dict[str, object]:
|
||||
return {"type": "message", "content": [], "usage": ANTHROPIC_CACHED_USAGE}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", ["gateway stream", "gateway json", "sdk json"])
|
||||
async def test_claude_code_cost_counts_anthropic_cached_prompt_tokens(
|
||||
monkeypatch: pytest.MonkeyPatch, route: str
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "cost_per_token", REAL_COST_PER_TOKEN)
|
||||
monkeypatch.setattr(litellm.anthropic.messages, "acreate", anthropic_cached_message)
|
||||
upstream: Final = (
|
||||
sse_response(anthropic_cached_sse())
|
||||
if route == "gateway stream"
|
||||
else httpx.Response(200, json={"usage": ANTHROPIC_CACHED_USAGE})
|
||||
)
|
||||
endpoint: Final = (
|
||||
ModelEndpoint(Harness.CLAUDE_CODE, CACHED_MODEL, None)
|
||||
if route == "sdk json"
|
||||
else gateway_endpoint(Recorder(upstream), model=CACHED_MODEL)
|
||||
)
|
||||
async with endpoint as ep:
|
||||
async with httpx.AsyncClient(base_url=ep.url) as client:
|
||||
await client.post("/v1/messages", json={"stream": route == "gateway stream"}, headers=auth(ep))
|
||||
cached_price: Final = sum(
|
||||
REAL_COST_PER_TOKEN(
|
||||
model=CACHED_MODEL,
|
||||
prompt_tokens=1000,
|
||||
completion_tokens=10,
|
||||
cache_read_input_tokens=800,
|
||||
cache_creation_input_tokens=100,
|
||||
)
|
||||
)
|
||||
assert cached_price > sum(REAL_COST_PER_TOKEN(model=CACHED_MODEL, prompt_tokens=100, completion_tokens=10))
|
||||
assert ep.usage.cost == pytest.approx(cached_price), (
|
||||
"Anthropic counts cached prompt tokens outside input_tokens, so leaving them out understates the turn's cost"
|
||||
)
|
||||
|
||||
|
||||
async def test_sdk_error_is_sanitized(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def failing(**kwargs: Any) -> Any:
|
||||
raise litellm.RateLimitError(
|
||||
|
|
@ -330,12 +508,12 @@ def test_load_server_deps_maps_import_error(monkeypatch: pytest.MonkeyPatch) ->
|
|||
|
||||
|
||||
def test_usage_helpers() -> None:
|
||||
assert usage_from_body({"usage": {"prompt_tokens": 1, "completion_tokens": 2}}) == (
|
||||
assert usage_from_mapping(usage_mapping_of_body({"usage": {"prompt_tokens": 1, "completion_tokens": 2}})) == (
|
||||
1,
|
||||
2,
|
||||
)
|
||||
assert usage_from_body({"response": {"usage": {"input_tokens": 3}}}) == (3, 0)
|
||||
assert usage_from_body("nope") == (0, 0)
|
||||
assert usage_from_mapping(usage_mapping_of_body({"response": {"usage": {"input_tokens": 3}}})) == (3, 0)
|
||||
assert usage_from_mapping(usage_mapping_of_body("nope")) == (0, 0)
|
||||
|
||||
parser = SSEUsageParser()
|
||||
for i in range(0, len(ANTHROPIC_SSE), 7): # split across arbitrary chunk borders
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ PUBLIC_NAMES = [
|
|||
"ClaudeCodeOptions",
|
||||
"CodexOptions",
|
||||
"OpenCodeOptions",
|
||||
"PiOptions",
|
||||
"DeepAgentsOptions",
|
||||
"ToolLoopOptions",
|
||||
"HarnessError",
|
||||
|
|
|
|||
24
tests/unit/harness/test_options.py
Normal file
24
tests/unit/harness/test_options.py
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
import copy
|
||||
import dataclasses
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, get_args
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.harness.options import HarnessOptions
|
||||
|
||||
|
||||
@pytest.mark.parametrize("options_type", get_args(HarnessOptions), ids=lambda t: t.__name__)
|
||||
def test_default_options_survive_deepcopy_and_asdict(options_type: type) -> None:
|
||||
options: Final = options_type()
|
||||
assert copy.deepcopy(options) == options
|
||||
assert dataclasses.asdict(options) == vars(options)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("options_type", get_args(HarnessOptions), ids=lambda t: t.__name__)
|
||||
def test_each_options_instance_gets_its_own_empty_mappings(options_type: type) -> None:
|
||||
first: Final = vars(options_type())
|
||||
second: Final = vars(options_type())
|
||||
mapping_names: Final = tuple(name for name, value in first.items() if isinstance(value, Mapping))
|
||||
assert {name: first[name] for name in mapping_names} == dict.fromkeys(mapping_names, {})
|
||||
assert [name for name in mapping_names if first[name] is second[name]] == [], "defaults shared between instances"
|
||||
|
|
@ -30,6 +30,7 @@ def test_harness_is_plain_enum():
|
|||
("codex", "Harness.CODEX"),
|
||||
("claude-code", "Harness.CLAUDE_CODE"),
|
||||
("OPENCODE", "Harness.OPENCODE"),
|
||||
("pi", "Harness.PI"),
|
||||
],
|
||||
)
|
||||
def test_require_harness_hint(given, hint):
|
||||
|
|
|
|||
0
tests/unit/llms/pi/__init__.py
Normal file
0
tests/unit/llms/pi/__init__.py
Normal file
0
tests/unit/llms/pi/harness/__init__.py
Normal file
0
tests/unit/llms/pi/harness/__init__.py
Normal file
0
tests/unit/llms/pi/harness/fixtures/__init__.py
Normal file
0
tests/unit/llms/pi/harness/fixtures/__init__.py
Normal file
12
tests/unit/llms/pi/harness/fixtures/api_error.jsonl
Normal file
12
tests/unit/llms/pi/harness/fixtures/api_error.jsonl
Normal file
File diff suppressed because one or more lines are too long
27
tests/unit/llms/pi/harness/fixtures/mcp_tool.jsonl
Normal file
27
tests/unit/llms/pi/harness/fixtures/mcp_tool.jsonl
Normal file
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
19
tests/unit/llms/pi/harness/fixtures/structured_output.jsonl
Normal file
19
tests/unit/llms/pi/harness/fixtures/structured_output.jsonl
Normal file
File diff suppressed because one or more lines are too long
42
tests/unit/llms/pi/harness/fixtures/turn1_write_read.jsonl
Normal file
42
tests/unit/llms/pi/harness/fixtures/turn1_write_read.jsonl
Normal file
File diff suppressed because one or more lines are too long
15
tests/unit/llms/pi/harness/fixtures/turn2_resume.jsonl
Normal file
15
tests/unit/llms/pi/harness/fixtures/turn2_resume.jsonl
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
{"type":"session","version":3,"id":"01a11cb0-87cf-752e-a3e3-365c7a169871","timestamp":"2026-10-08T18:04:53.583Z","cwd":"/work"}
|
||||
{"type":"agent_start"}
|
||||
{"type":"turn_start"}
|
||||
{"type":"message_start","message":{"role":"user","content":[{"type":"text","text":"do the task"}],"timestamp":1791482698176}}
|
||||
{"type":"message_end","message":{"role":"user","content":[{"type":"text","text":"do the task"}],"timestamp":1791482698176}}
|
||||
{"type":"message_start","message":{"role":"assistant","content":[],"api":"openai-completions","provider":"litellm","model":"claude-haiku-4-5-20251001","usage":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"totalTokens":0,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"stopReason":"pending","timestamp":1791482698355}}
|
||||
{"type":"message_update","usage":{"input":120,"output":9,"cacheRead":0,"cacheWrite":0,"reasoning":0,"totalTokens":129,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"assistantMessageEvent":{"type":"text_start","contentIndex":0}}
|
||||
{"type":"message_update","usage":{"input":120,"output":9,"cacheRead":0,"cacheWrite":0,"reasoning":0,"totalTokens":129,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":"You asked me"}}
|
||||
{"type":"message_update","usage":{"input":120,"output":9,"cacheRead":0,"cacheWrite":0,"reasoning":0,"totalTokens":129,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":" to create h"}}
|
||||
{"type":"message_update","usage":{"input":120,"output":9,"cacheRead":0,"cacheWrite":0,"reasoning":0,"totalTokens":129,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":"ello.txt."}}
|
||||
{"type":"message_update","usage":{"input":120,"output":9,"cacheRead":0,"cacheWrite":0,"reasoning":0,"totalTokens":129,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"assistantMessageEvent":{"type":"text_end","contentIndex":0,"content":"You asked me to create hello.txt."}}
|
||||
{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"You asked me to create hello.txt."}],"api":"openai-completions","provider":"litellm","model":"claude-haiku-4-5-20251001","usage":{"input":120,"output":9,"cacheRead":0,"cacheWrite":0,"reasoning":0,"totalTokens":129,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"stopReason":"stop","timestamp":1791482698355,"responseId":"chatcmpl-1","rawStopReason":"stop","durationMs":169,"thinkingLevel":"off"}}
|
||||
{"type":"turn_end","message":{"role":"assistant","content":[{"type":"text","text":"You asked me to create hello.txt."}],"api":"openai-completions","provider":"litellm","model":"claude-haiku-4-5-20251001","usage":{"input":120,"output":9,"cacheRead":0,"cacheWrite":0,"reasoning":0,"totalTokens":129,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"stopReason":"stop","timestamp":1791482698355,"responseId":"chatcmpl-1","rawStopReason":"stop","durationMs":169,"thinkingLevel":"off"},"toolResults":[]}
|
||||
{"type":"agent_end","messages":[{"role":"user","content":[{"type":"text","text":"do the task"}],"timestamp":1791482698176},{"role":"assistant","content":[{"type":"text","text":"You asked me to create hello.txt."}],"api":"openai-completions","provider":"litellm","model":"claude-haiku-4-5-20251001","usage":{"input":120,"output":9,"cacheRead":0,"cacheWrite":0,"reasoning":0,"totalTokens":129,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"stopReason":"stop","timestamp":1791482698355,"responseId":"chatcmpl-1","rawStopReason":"stop","durationMs":169,"thinkingLevel":"off"}],"willRetry":false}
|
||||
{"type":"agent_settled","aborted":false}
|
||||
991
tests/unit/llms/pi/harness/test_transformation.py
Normal file
991
tests/unit/llms/pi/harness/test_transformation.py
Normal file
|
|
@ -0,0 +1,991 @@
|
|||
"""Unit tests for the pi harness config. No network, no real CLI.
|
||||
|
||||
Fixtures under fixtures/ are sanitized `pi --mode json` output recorded from
|
||||
@earendil-works/pi-coding-agent 1.1.0 against a scripted OpenAI-compatible endpoint.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import itertools
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from litellm.harness.context import GatewayTarget, SessionContext
|
||||
from litellm.harness.errors import (
|
||||
CapabilityUnsupported,
|
||||
HarnessError,
|
||||
HarnessInstallFailed,
|
||||
OptionsMismatch,
|
||||
)
|
||||
from litellm.harness.handlers.cli_handler import PERSIST_DIR_SCRIPT, CLIHarnessHandler
|
||||
from litellm.harness.options import HarnessOptions, OpenCodeOptions, PiOptions, PiThinkingLevel
|
||||
from litellm.harness.sandbox.base import CompletedRun
|
||||
from litellm.harness.types import Event, Harness, PermissionMode, Reasoning, Text, ToolCall, ToolResult
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.base_llm.harness.transformation import (
|
||||
HarnessSessionSetup,
|
||||
HarnessTurnError,
|
||||
)
|
||||
from litellm.llms.pi.harness.transformation import (
|
||||
INSTRUCTIONS_FILENAME,
|
||||
MANAGED_CONFIG_KEYS,
|
||||
MANAGED_ENV_KEYS,
|
||||
MCP_FILENAME,
|
||||
MODELS_FILENAME,
|
||||
PI_CONFIG_DIR_ENV,
|
||||
PI_DEFAULT_MAX_TOKENS,
|
||||
PI_ISOLATION_ENV,
|
||||
PI_TOKEN_ENV,
|
||||
SETTINGS_FILENAME,
|
||||
PiHarnessConfig,
|
||||
PiStreamState,
|
||||
build_models_json,
|
||||
config_files,
|
||||
mcp_tool_entries,
|
||||
tool_args,
|
||||
validate_user_config,
|
||||
)
|
||||
from litellm.types.utils import ModelInfo
|
||||
from litellm.utils import ProviderConfigManager, get_model_info
|
||||
|
||||
FIXTURES: Final = Path(__file__).parent / "fixtures"
|
||||
TOKEN: Final = "tok-secret-123"
|
||||
SESSION: Final = "01a11cb0-87cf-752e-a3e3-365c7a169871"
|
||||
PRIVATE: Final = "/tmp/pi-1"
|
||||
MODEL: Final = "claude-haiku-4-5-20251001"
|
||||
FULL_TOOLS: Final = "read,bash,edit,write,grep,find,ls"
|
||||
MOYAI: Final = {"command": "/path/to/mcp-server", "args": [], "env": {}, "exposure": "direct"}
|
||||
CONFIG: Final = PiHarnessConfig()
|
||||
|
||||
|
||||
def load_fixture(name: str) -> Sequence[Mapping[str, object]]:
|
||||
return tuple(json.loads(line) for line in (FIXTURES / name).read_text().splitlines() if line)
|
||||
|
||||
|
||||
def parse(obj: Mapping[str, object], state: PiStreamState) -> Sequence[Event]:
|
||||
return CONFIG.transform_stream_line(obj, state)
|
||||
|
||||
|
||||
def parse_all(name: str, state: PiStreamState | None = None) -> tuple[Sequence[Event], PiStreamState]:
|
||||
run_state: Final = state or CONFIG.create_stream_state()
|
||||
events: Final = tuple(itertools.chain.from_iterable(parse(obj, run_state) for obj in load_fixture(name)))
|
||||
return events, run_state
|
||||
|
||||
|
||||
def assistant_end(text: str, stop_reason: str, error: str | None = None) -> Mapping[str, object]:
|
||||
error_field: Final = {"errorMessage": error} if error else {}
|
||||
return {
|
||||
"type": "message_end",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": text}],
|
||||
"stopReason": stop_reason,
|
||||
**error_field,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class FakeStdin:
|
||||
def __init__(self) -> None:
|
||||
self.data = b""
|
||||
self.closed = False
|
||||
|
||||
def write(self, data: bytes) -> None:
|
||||
self.data += data
|
||||
|
||||
async def drain(self) -> None:
|
||||
return None
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class FakeProcess:
|
||||
def __init__(self, stdout: bytes, stderr: bytes = b"", exit_code: int = 0) -> None:
|
||||
self.stdin = FakeStdin()
|
||||
self.stdout = asyncio.StreamReader()
|
||||
self.stdout.feed_data(stdout)
|
||||
self.stdout.feed_eof()
|
||||
self.stderr = asyncio.StreamReader()
|
||||
self.stderr.feed_data(stderr)
|
||||
self.stderr.feed_eof()
|
||||
self._exit_code = exit_code
|
||||
self.killed = False
|
||||
|
||||
async def wait(self) -> int:
|
||||
return self._exit_code
|
||||
|
||||
async def kill(self) -> None:
|
||||
self.killed = True
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Exec:
|
||||
cmd: Sequence[str]
|
||||
env: Mapping[str, str]
|
||||
cwd: str | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class FakeSandbox:
|
||||
workdir: str = "/work"
|
||||
has_binary: bool = True
|
||||
persist_ok: bool = True
|
||||
outputs: tuple[FakeProcess, ...] = ()
|
||||
files: Mapping[str, bytes] = field(default_factory=lambda: MappingProxyType({}))
|
||||
execs: tuple[Exec, ...] = ()
|
||||
runs: tuple[Sequence[str], ...] = ()
|
||||
tempdirs: int = 0
|
||||
|
||||
async def exec(
|
||||
self, cmd: Sequence[str], *, env: Mapping[str, str] | None = None, cwd: str | None = None
|
||||
) -> FakeProcess:
|
||||
self.execs = (*self.execs, Exec(cmd=cmd, env=MappingProxyType(dict(env or {})), cwd=cwd))
|
||||
return self.outputs[len(self.execs) - 1]
|
||||
|
||||
async def run(
|
||||
self,
|
||||
cmd: Sequence[str],
|
||||
*,
|
||||
env: Mapping[str, str] | None = None,
|
||||
cwd: str | None = None,
|
||||
timeout: float | None = None,
|
||||
) -> CompletedRun:
|
||||
self.runs = (*self.runs, cmd)
|
||||
if self.persist_ok:
|
||||
return CompletedRun("", "", 0)
|
||||
return CompletedRun("", "read-only fs", 1)
|
||||
|
||||
async def read(self, path: str) -> bytes:
|
||||
return self.files[path]
|
||||
|
||||
async def write(self, path: str, data: bytes) -> None:
|
||||
self.files = MappingProxyType({**self.files, path: data})
|
||||
|
||||
def host_url(self, port: int) -> str:
|
||||
return f"http://host.docker.internal:{port}"
|
||||
|
||||
async def which(self, binary: str) -> str | None:
|
||||
return f"/usr/bin/{binary}" if self.has_binary else None
|
||||
|
||||
async def tempdir(self) -> str:
|
||||
self.tempdirs += 1
|
||||
return f"/tmp/pi-{self.tempdirs}"
|
||||
|
||||
async def snapshot(self) -> Mapping[str, str]:
|
||||
return MappingProxyType({})
|
||||
|
||||
async def close(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FakeEndpoint:
|
||||
port: int = 4555
|
||||
token: str = TOKEN
|
||||
model: str | None = None
|
||||
|
||||
|
||||
DEFAULT_ENDPOINT: Final = FakeEndpoint()
|
||||
|
||||
|
||||
class Answer(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
city: str
|
||||
country: str
|
||||
|
||||
|
||||
def make_ctx(
|
||||
sandbox: FakeSandbox | None = None,
|
||||
*,
|
||||
model: str | None = MODEL,
|
||||
gateway: GatewayTarget | None = None,
|
||||
endpoint: FakeEndpoint | None = DEFAULT_ENDPOINT,
|
||||
options: HarnessOptions | None = None,
|
||||
permissions: PermissionMode = "full",
|
||||
instructions: str | None = None,
|
||||
output: type[BaseModel] | None = None,
|
||||
skills: Sequence[str] = (),
|
||||
disable_tools: Sequence[str] = (),
|
||||
) -> SessionContext:
|
||||
return SessionContext(
|
||||
harness=Harness.PI,
|
||||
sandbox=sandbox or FakeSandbox(),
|
||||
session_id="s1",
|
||||
model=model,
|
||||
gateway=gateway,
|
||||
endpoint=endpoint,
|
||||
options=options,
|
||||
permissions=permissions,
|
||||
instructions=instructions,
|
||||
output=output,
|
||||
skills=skills,
|
||||
disable_tools=disable_tools,
|
||||
)
|
||||
|
||||
|
||||
def setup_for(ctx: SessionContext) -> HarnessSessionSetup:
|
||||
return CONFIG.transform_session_setup(ctx, PRIVATE)
|
||||
|
||||
|
||||
def setup_models(setup: HarnessSessionSetup) -> Mapping[str, object]:
|
||||
return json.loads(setup.files[MODELS_FILENAME])
|
||||
|
||||
|
||||
def fixture_proc(name: str, exit_code: int = 0) -> FakeProcess:
|
||||
return FakeProcess((FIXTURES / name).read_bytes(), exit_code=exit_code)
|
||||
|
||||
|
||||
async def collect(handler: CLIHarnessHandler, ctx: SessionContext, prompt: str) -> Sequence[Event]:
|
||||
return [e async for e in handler.turn(ctx, prompt)]
|
||||
|
||||
|
||||
async def started(ctx: SessionContext) -> CLIHarnessHandler:
|
||||
handler: Final = CLIHarnessHandler(PiHarnessConfig())
|
||||
await handler.start(ctx)
|
||||
return handler
|
||||
|
||||
|
||||
def flag(argv: Sequence[str], name: str) -> str:
|
||||
args: Final = list(argv)
|
||||
return args[args.index(name) + 1]
|
||||
|
||||
|
||||
def written_models(sandbox: FakeSandbox, private: str = PRIVATE) -> Mapping[str, object]:
|
||||
return json.loads(sandbox.files[f"{private}/{MODELS_FILENAME}"])
|
||||
|
||||
|
||||
def test_parse_write_read_turn():
|
||||
events, state = parse_all("turn1_write_read.jsonl")
|
||||
assert state.session_id == SESSION
|
||||
assert [type(e) for e in events] == [
|
||||
Reasoning,
|
||||
ToolCall,
|
||||
ToolResult,
|
||||
ToolCall,
|
||||
ToolResult,
|
||||
Text,
|
||||
Text,
|
||||
Text,
|
||||
]
|
||||
assert events[0] == Reasoning(delta="I should write the file first.")
|
||||
write, write_result = events[1], events[2]
|
||||
assert write.name == "write" and write.native_name == "write"
|
||||
assert write.builtin is True
|
||||
assert write.input == {"path": "hello.txt", "content": "hi"}
|
||||
assert write_result.id == write.id == "call_w1"
|
||||
assert write_result.output == "Successfully wrote to hello.txt"
|
||||
assert write_result.is_error is False
|
||||
read, read_result = events[3], events[4]
|
||||
assert read.name == "read" and read_result.output == "hi"
|
||||
assert "".join(e.delta for e in events[5:]) == state.final_text
|
||||
assert state.final_text == "Done! hello.txt contains: hi"
|
||||
assert state.stop_reason == "stop" and state.error is None
|
||||
|
||||
|
||||
def test_parse_continued_session():
|
||||
events, state = parse_all("turn2_resume.jsonl")
|
||||
assert state.session_id == SESSION
|
||||
assert all(isinstance(e, Text) for e in events)
|
||||
assert state.final_text == "You asked me to create hello.txt."
|
||||
|
||||
|
||||
def test_parse_denied_tool_is_error_result():
|
||||
events, state = parse_all("readonly_denied_bash.jsonl")
|
||||
call: Final = next(e for e in events if isinstance(e, ToolCall))
|
||||
result: Final = next(e for e in events if isinstance(e, ToolResult))
|
||||
assert call.native_name == "bash" and call.input == {"command": "echo x > blocked.txt"}
|
||||
assert result == ToolResult(id="call_b1", output="Tool bash not found", is_error=True)
|
||||
assert state.final_text.startswith("FAILED")
|
||||
|
||||
|
||||
def test_parse_mcp_tool_turn():
|
||||
events, state = parse_all("mcp_tool.jsonl")
|
||||
call: Final = next(e for e in events if isinstance(e, ToolCall))
|
||||
result: Final = next(e for e in events if isinstance(e, ToolResult))
|
||||
assert call == ToolCall(
|
||||
id="call_m1",
|
||||
name="mcp__moyai__echo",
|
||||
native_name="mcp__moyai__echo",
|
||||
input={"text": "hello"},
|
||||
builtin=False,
|
||||
)
|
||||
assert result == ToolResult(id="call_m1", output="moyai says: hello", is_error=False)
|
||||
assert state.final_text == "The MCP tool replied."
|
||||
|
||||
|
||||
def test_parse_api_error_records_error():
|
||||
events, state = parse_all("api_error.jsonl")
|
||||
assert events == ()
|
||||
assert state.stop_reason == "error"
|
||||
assert "no healthy deployments" in state.error
|
||||
|
||||
|
||||
def test_parse_reasoning_and_tool_error_and_name_mapping():
|
||||
state: Final = PiStreamState()
|
||||
reasoning: Final = parse(
|
||||
{
|
||||
"type": "message_update",
|
||||
"assistantMessageEvent": {"type": "thinking_delta", "delta": "thinking hard"},
|
||||
},
|
||||
state,
|
||||
)
|
||||
assert reasoning == [Reasoning(delta="thinking hard")]
|
||||
failed: Final = parse(
|
||||
{
|
||||
"type": "tool_execution_end",
|
||||
"toolCallId": "c1",
|
||||
"toolName": "bash",
|
||||
"result": {"content": [{"type": "text", "text": "boom"}]},
|
||||
"isError": True,
|
||||
},
|
||||
state,
|
||||
)
|
||||
assert failed == [ToolResult(id="c1", output="boom", is_error=True)]
|
||||
calls: Final = [
|
||||
parse({"type": "tool_execution_start", "toolCallId": "x", "toolName": native, "args": {}}, state)[0]
|
||||
for native in ("find", "powershell", "grep", "ls", "edit")
|
||||
]
|
||||
assert [(call.name, call.builtin) for call in calls] == [
|
||||
("glob", True),
|
||||
("bash", True),
|
||||
("grep", True),
|
||||
("ls", True),
|
||||
("edit", True),
|
||||
]
|
||||
mcp: Final = parse(
|
||||
{"type": "tool_execution_start", "toolCallId": "m", "toolName": "mcp__gh__search", "args": {"q": 1}},
|
||||
state,
|
||||
)
|
||||
assert mcp[0].builtin is False and mcp[0].input == {"q": 1}
|
||||
plain: Final = parse(
|
||||
{"type": "tool_execution_end", "toolCallId": "m", "toolName": "mcp__gh__search", "result": "raw"},
|
||||
state,
|
||||
)
|
||||
assert plain == [ToolResult(id="m", output="raw", is_error=False)]
|
||||
no_result: Final = parse({"type": "tool_execution_end", "toolCallId": "e", "toolName": "bash"}, state)
|
||||
assert no_result == [ToolResult(id="e", output="", is_error=False)]
|
||||
|
||||
|
||||
def test_final_text_is_last_assistant_message():
|
||||
state: Final = PiStreamState()
|
||||
parse(assistant_end("working", "toolUse"), state)
|
||||
parse(assistant_end("answer", "stop"), state)
|
||||
parse(
|
||||
{"type": "message_end", "message": {"role": "user", "content": [{"type": "text", "text": "q"}]}},
|
||||
state,
|
||||
)
|
||||
parse({"type": "message_end", "message": {"role": "system", "content": ""}}, state)
|
||||
assert state.final_text == "answer"
|
||||
parse(
|
||||
{
|
||||
"type": "message_end",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"stopReason": "stop",
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": "hmm"},
|
||||
{"type": "text", "text": "a"},
|
||||
{"type": "toolCall", "id": "x", "name": "read", "arguments": {}},
|
||||
{"type": "text", "text": "b"},
|
||||
],
|
||||
},
|
||||
},
|
||||
state,
|
||||
)
|
||||
assert state.final_text == "ab"
|
||||
|
||||
|
||||
def test_message_update_ignores_non_delta_events():
|
||||
state: Final = PiStreamState()
|
||||
for update in [
|
||||
{"type": "text_start", "contentIndex": 0},
|
||||
{"type": "text_end", "content": "full"},
|
||||
{"type": "toolcall_delta", "delta": "{"},
|
||||
{"type": "text_delta", "delta": ""},
|
||||
]:
|
||||
assert parse({"type": "message_update", "assistantMessageEvent": update}, state) == []
|
||||
assert parse({"type": "message_update"}, state) == []
|
||||
|
||||
|
||||
def test_retried_error_then_success_is_not_an_error():
|
||||
state: Final = PiStreamState()
|
||||
parse(assistant_end("", "error", "529 overloaded"), state)
|
||||
parse(assistant_end("recovered", "stop"), state)
|
||||
assert state.stop_reason == "stop" and state.error is None
|
||||
assert CONFIG.transform_turn_response(make_ctx(), state, 0, []).final_text == "recovered"
|
||||
|
||||
|
||||
def test_aborted_run_is_an_error():
|
||||
state: Final = PiStreamState()
|
||||
parse(assistant_end("partial", "stop"), state)
|
||||
parse({"type": "agent_settled", "aborted": False}, state)
|
||||
assert state.stop_reason == "stop"
|
||||
parse({"type": "agent_settled", "aborted": True}, state)
|
||||
assert state.stop_reason == "aborted"
|
||||
|
||||
|
||||
def test_permission_mapping():
|
||||
assert tool_args("full", ()) == ("--tools", FULL_TOOLS)
|
||||
assert tool_args("read-only", ()) == ("--tools", "read,grep,find,ls")
|
||||
assert tool_args("edit", ()) == ("--tools", "read,edit,write,grep,find,ls")
|
||||
with pytest.raises(CapabilityUnsupported):
|
||||
tool_args("ask", ())
|
||||
|
||||
|
||||
def test_disable_tools_map_to_native_excludes():
|
||||
args: Final = tool_args("full", ["bash", "glob", "write", "mcp__x"])
|
||||
assert flag(args, "--tools") == FULL_TOOLS
|
||||
assert flag(args, "--exclude-tools") == "bash,powershell,find,write,mcp__x"
|
||||
|
||||
|
||||
def test_mcp_servers_keep_tools_reachable():
|
||||
assert mcp_tool_entries({}) == ()
|
||||
assert mcp_tool_entries({"mcpServers": {}}) == ()
|
||||
assert mcp_tool_entries({"mcpServers": {"moyai": MOYAI}}) == ("mcp__*",)
|
||||
assert mcp_tool_entries({"mcpServers": {"s": {"command": "x"}}}) == ("mcp__*", "codemode")
|
||||
assert mcp_tool_entries({"mcpServers": {"s": {"url": "u", "exposure": "deferred"}}}) == (
|
||||
"mcp__*",
|
||||
"tool_search",
|
||||
)
|
||||
assert mcp_tool_entries(
|
||||
{"mcpServers": {"s": {"command": "x", "exposure": "hidden", "toolExposure": {"a": "direct", "b": "deferred"}}}}
|
||||
) == ("mcp__*", "tool_search")
|
||||
args: Final = tool_args("read-only", ["mcp__moyai__drop"], mcp_tool_entries({"mcpServers": {"moyai": MOYAI}}))
|
||||
assert flag(args, "--tools") == "read,grep,find,ls,mcp__*"
|
||||
assert flag(args, "--exclude-tools") == "mcp__moyai__drop"
|
||||
|
||||
|
||||
def test_config_files_split_settings_and_mcp():
|
||||
files: Final = dict(config_files({"mcpServers": {"moyai": MOYAI}, "compaction": {"enabled": False}}))
|
||||
assert json.loads(files[MCP_FILENAME]) == {"mcpServers": {"moyai": MOYAI}}
|
||||
assert json.loads(files[SETTINGS_FILENAME]) == {"compaction": {"enabled": False}}
|
||||
assert dict(config_files({"mcpServers": {"moyai": MOYAI}})).keys() == {MCP_FILENAME}
|
||||
assert dict(config_files({"theme": "dark"})).keys() == {SETTINGS_FILENAME}
|
||||
assert config_files({}) == ()
|
||||
|
||||
|
||||
def test_config_files_reject_values_that_are_not_json():
|
||||
with pytest.raises(TypeError, match=r"set in PiOptions\.config is not JSON serializable"):
|
||||
config_files({"theme": {"dark", "light"}})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config",
|
||||
[
|
||||
*({key: "x"} for key in sorted(MANAGED_CONFIG_KEYS)),
|
||||
{"mcpServers": "not-a-mapping"},
|
||||
{"mcpServers": {"moyai": "not-a-mapping"}},
|
||||
],
|
||||
)
|
||||
def test_managed_keys_rejected(config: Mapping[str, object]) -> None:
|
||||
with pytest.raises(OptionsMismatch):
|
||||
validate_user_config(config)
|
||||
|
||||
|
||||
def test_options_config_cannot_load_code_or_reroute_model():
|
||||
for key in ("extensions", "packages", "defaultProjectTrust", "defaultProvider", "defaultModel", "httpProxy"):
|
||||
assert key in MANAGED_CONFIG_KEYS
|
||||
with pytest.raises(OptionsMismatch, match="extensions"):
|
||||
validate_user_config({"extensions": ["./evil.ts"]})
|
||||
|
||||
|
||||
def test_build_models_json():
|
||||
models: Final = json.loads(build_models_json("m1", "http://h:1/v1", None, None))
|
||||
assert models == {
|
||||
"providers": {
|
||||
"litellm": {
|
||||
"baseUrl": "http://h:1/v1",
|
||||
"api": "openai-completions",
|
||||
"apiKey": f"${PI_TOKEN_ENV}",
|
||||
"compat": {"supportsStore": False},
|
||||
"models": [{"id": "m1"}],
|
||||
}
|
||||
}
|
||||
}, (
|
||||
"an unmapped model keeps pi's own limits and no cache_control marks, "
|
||||
"and pi must not send `store`, which non-OpenAI providers reject"
|
||||
)
|
||||
|
||||
|
||||
def test_mapped_model_declares_litellm_context_and_output_limits() -> None:
|
||||
info: Final = get_model_info(MODEL)
|
||||
models: Final = setup_models(setup_for(make_ctx()))
|
||||
declared: Final = models["providers"]["litellm"]["models"][0]
|
||||
assert (declared["contextWindow"], declared["maxTokens"]) == (
|
||||
info["max_input_tokens"],
|
||||
min(info["max_output_tokens"], PI_DEFAULT_MAX_TOKENS),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("max_output_tokens", [PI_DEFAULT_MAX_TOKENS // 4, PI_DEFAULT_MAX_TOKENS * 8])
|
||||
def test_pi_asks_for_no_more_output_tokens_than_the_model_or_its_own_default(max_output_tokens: int) -> None:
|
||||
info: Final = ModelInfo(**{**get_model_info(MODEL), "max_output_tokens": max_output_tokens})
|
||||
declared: Final = json.loads(build_models_json(MODEL, "http://h:1/v1", None, info))["providers"]["litellm"]
|
||||
assert declared["models"][0]["maxTokens"] == min(max_output_tokens, PI_DEFAULT_MAX_TOKENS), (
|
||||
"pi sends maxTokens as max_completion_tokens on every call, and the gateway reserves it against the key's "
|
||||
"TPM limit, so a model's full output cap made keys under that limit answer 429"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("info", [get_model_info(MODEL), None], ids=["mapped", "unmapped"])
|
||||
@pytest.mark.parametrize("thinking", ["minimal", "low", "medium", "high", "xhigh", "max"])
|
||||
def test_thinking_budget_leaves_pi_its_default_room_for_the_answer(
|
||||
thinking: PiThinkingLevel, info: ModelInfo | None
|
||||
) -> None:
|
||||
declared: Final = json.loads(build_models_json(MODEL, "http://h:1/v1", thinking, info))["providers"]["litellm"]
|
||||
sent: Final = AnthropicConfig().map_openai_params(
|
||||
{"reasoning_effort": thinking, "max_completion_tokens": declared["models"][0]["maxTokens"]},
|
||||
{},
|
||||
MODEL,
|
||||
drop_params=False,
|
||||
)
|
||||
assert sent["max_tokens"] - sent["thinking"]["budget_tokens"] >= PI_DEFAULT_MAX_TOKENS, (
|
||||
"Anthropic answers 400 unless max_tokens is above the thinking budget LiteLLM sends for the level, "
|
||||
"so a budget as large as pi's default left thinking=max with no request it could make"
|
||||
)
|
||||
assert sent["max_tokens"] <= get_model_info(MODEL)["max_output_tokens"]
|
||||
|
||||
|
||||
def test_model_that_supports_prompt_caching_asks_pi_for_cache_control_marks() -> None:
|
||||
assert get_model_info(MODEL).get("supports_prompt_caching") is True
|
||||
declared: Final = setup_models(setup_for(make_ctx()))["providers"]["litellm"]["models"][0]
|
||||
assert declared["compat"] == {"cacheControlFormat": "anthropic"}, (
|
||||
"pi only marks cache_control for OpenRouter Anthropic models by itself, so every Claude call missed the cache"
|
||||
)
|
||||
|
||||
|
||||
GATEWAY: Final = GatewayTarget(api_base="http://gateway:4000", api_key="sk-gateway")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("group", "declared_limits"),
|
||||
[
|
||||
(MODEL, {"contextWindow": get_model_info(MODEL)["max_input_tokens"]}),
|
||||
("team-default-model", {}),
|
||||
],
|
||||
)
|
||||
def test_gateway_group_is_looked_up_without_printing_to_stdout(
|
||||
group: str, declared_limits: Mapping[str, int], capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
declared: Final = setup_models(setup_for(make_ctx(model=group, gateway=GATEWAY)))["providers"]["litellm"]["models"][
|
||||
0
|
||||
]
|
||||
assert {key: declared[key] for key in declared_limits} == declared_limits
|
||||
assert ("maxTokens" in declared) == bool(declared_limits)
|
||||
assert capsys.readouterr().out == "", "an unmapped gateway group printed LiteLLM's provider-list banner on stdout"
|
||||
|
||||
|
||||
MAPPED_MODEL: Final = {
|
||||
"id": MODEL,
|
||||
"contextWindow": get_model_info(MODEL)["max_input_tokens"],
|
||||
"maxTokens": min(get_model_info(MODEL)["max_output_tokens"], PI_DEFAULT_MAX_TOKENS),
|
||||
"compat": {"cacheControlFormat": "anthropic"},
|
||||
}
|
||||
REASONING_MODEL: Final = {**MAPPED_MODEL, "reasoning": True, "thinkingLevelMap": {"xhigh": "xhigh", "max": "max"}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("thinking", "declared_model"),
|
||||
[
|
||||
(None, MAPPED_MODEL),
|
||||
("off", MAPPED_MODEL),
|
||||
("minimal", REASONING_MODEL),
|
||||
("high", REASONING_MODEL),
|
||||
("xhigh", REASONING_MODEL),
|
||||
("max", REASONING_MODEL),
|
||||
],
|
||||
)
|
||||
def test_thinking_declares_a_reasoning_model_only_when_a_level_is_requested(
|
||||
thinking: PiThinkingLevel | None, declared_model: Mapping[str, object]
|
||||
) -> None:
|
||||
models: Final = setup_models(setup_for(make_ctx(options=PiOptions(thinking=thinking))))
|
||||
declared: Final = {key: value for key, value in models["providers"]["litellm"]["models"][0].items() if key != "maxTokens"}
|
||||
assert declared == {key: value for key, value in declared_model.items() if key != "maxTokens"}, (
|
||||
"pi drops --thinking to off unless the model declares reasoning, and drops xhigh/max to high "
|
||||
"unless thinkingLevelMap names them; declaring reasoning with no level makes pi send its default effort"
|
||||
)
|
||||
|
||||
|
||||
def test_config_metadata():
|
||||
assert CONFIG.get_binary() == "pi"
|
||||
assert "@earendil-works/pi-coding-agent" in CONFIG.get_install_hint()
|
||||
assert CONFIG.uses_model_endpoint is True
|
||||
assert CONFIG.capabilities.permission_modes == {"read-only", "edit", "full"}
|
||||
assert CONFIG.capabilities.resume is True
|
||||
assert CONFIG.capabilities.tool_filtering is True
|
||||
assert isinstance(ProviderConfigManager.get_provider_harness_config(Harness.PI), PiHarnessConfig)
|
||||
|
||||
|
||||
def test_validate_environment_rejects_wrong_options_managed_config_and_ask_mode():
|
||||
CONFIG.validate_environment(make_ctx())
|
||||
CONFIG.validate_environment(make_ctx(options=PiOptions(config={"mcpServers": {"moyai": MOYAI}})))
|
||||
with pytest.raises(OptionsMismatch):
|
||||
CONFIG.validate_environment(make_ctx(options=OpenCodeOptions()))
|
||||
with pytest.raises(OptionsMismatch, match="defaultModel"):
|
||||
CONFIG.validate_environment(make_ctx(options=PiOptions(config={"defaultModel": "openai/x"})))
|
||||
with pytest.raises(CapabilityUnsupported):
|
||||
CONFIG.validate_environment(make_ctx(permissions="ask"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key", sorted(MANAGED_ENV_KEYS))
|
||||
def test_managed_env_keys_rejected_instead_of_silently_ignored(key: str) -> None:
|
||||
assert key in {PI_CONFIG_DIR_ENV, PI_TOKEN_ENV}
|
||||
with pytest.raises(OptionsMismatch, match=key):
|
||||
CONFIG.validate_environment(make_ctx(options=PiOptions(env={key: "/tmp/evil"})))
|
||||
CONFIG.validate_environment(make_ctx(options=PiOptions(env={f"{key}_SUFFIX": "ok"})))
|
||||
|
||||
|
||||
def test_session_setup_token_only_in_env():
|
||||
setup: Final = setup_for(make_ctx())
|
||||
assert set(setup.files) == {MODELS_FILENAME}
|
||||
assert all(TOKEN.encode() not in data for data in setup.files.values())
|
||||
assert setup.env[PI_TOKEN_ENV] == TOKEN
|
||||
provider: Final = setup_models(setup)["providers"]["litellm"]
|
||||
assert provider["baseUrl"] == "http://host.docker.internal:4555/v1"
|
||||
assert provider["apiKey"] == f"${PI_TOKEN_ENV}"
|
||||
|
||||
|
||||
def test_session_setup_env_and_persisted_sessions():
|
||||
options: Final = PiOptions(env={"FOO": "1", PI_TOKEN_ENV: "evil", PI_CONFIG_DIR_ENV: "/home/u/.pi/agent"})
|
||||
setup: Final = setup_for(make_ctx(options=options))
|
||||
assert list(setup.persisted_dirs) == [("sessions", "pi/sessions")]
|
||||
assert setup.skills_dir == "skills"
|
||||
env: Final = setup.env
|
||||
assert env[PI_CONFIG_DIR_ENV] == f"{PRIVATE}/agent", "managed keys must win over PiOptions.env"
|
||||
assert env[PI_TOKEN_ENV] == TOKEN, "managed keys must win over PiOptions.env"
|
||||
for key, value in PI_ISOLATION_ENV.items():
|
||||
assert env[key] == value
|
||||
assert env["FOO"] == "1"
|
||||
|
||||
|
||||
def test_session_setup_writes_user_config_next_to_managed_models():
|
||||
options: Final = PiOptions(config={"mcpServers": {"moyai": MOYAI}, "compaction": {"enabled": False}})
|
||||
setup: Final = setup_for(make_ctx(options=options))
|
||||
assert set(setup.files) == {MODELS_FILENAME, MCP_FILENAME, SETTINGS_FILENAME}
|
||||
assert json.loads(setup.files[MCP_FILENAME]) == {"mcpServers": {"moyai": MOYAI}}
|
||||
assert setup_models(setup)["providers"]["litellm"]["apiKey"] == f"${PI_TOKEN_ENV}"
|
||||
|
||||
|
||||
def test_session_setup_errors():
|
||||
with pytest.raises(HarnessError):
|
||||
setup_for(make_ctx(endpoint=None))
|
||||
with pytest.raises(ValueError, match="needs model="):
|
||||
setup_for(make_ctx(model=None, endpoint=FakeEndpoint(model=None)))
|
||||
|
||||
|
||||
def test_turn_request_without_a_model_raises_instead_of_passing_none():
|
||||
ctx: Final = make_ctx(model=None, endpoint=FakeEndpoint(model=None))
|
||||
with pytest.raises(ValueError, match="needs model="):
|
||||
CONFIG.transform_turn_request(ctx, HarnessSessionSetup(), PRIVATE, "hi", None)
|
||||
|
||||
|
||||
def test_session_setup_instructions_and_skills():
|
||||
ctx: Final = make_ctx(instructions="Be terse.", output=Answer, skills=["/s/greeter"])
|
||||
setup: Final = setup_for(ctx)
|
||||
written: Final = setup.files[INSTRUCTIONS_FILENAME].decode()
|
||||
assert written.startswith("Be terse.")
|
||||
assert '"city"' in written and "single JSON object" in written
|
||||
argv: Final = CONFIG.transform_turn_request(ctx, setup, PRIVATE, "hi", None).argv
|
||||
assert flag(argv, "--append-system-prompt") == f"{PRIVATE}/instructions.md"
|
||||
assert flag(argv, "--skill") == f"{PRIVATE}/skills"
|
||||
plain: Final = make_ctx()
|
||||
plain_setup: Final = setup_for(plain)
|
||||
plain_argv: Final = CONFIG.transform_turn_request(plain, plain_setup, PRIVATE, "hi", None).argv
|
||||
assert INSTRUCTIONS_FILENAME not in plain_setup.files
|
||||
assert "--append-system-prompt" not in plain_argv and "--skill" not in plain_argv
|
||||
|
||||
|
||||
def test_turn_request_argv_and_session_continuation():
|
||||
ctx: Final = make_ctx(options=PiOptions(thinking="high"))
|
||||
setup: Final = setup_for(ctx)
|
||||
first: Final = CONFIG.transform_turn_request(ctx, setup, PRIVATE, "hello", None)
|
||||
assert list(first.argv) == [
|
||||
"pi",
|
||||
"--mode",
|
||||
"json",
|
||||
"--no-approve",
|
||||
"--no-skills",
|
||||
"--provider",
|
||||
"litellm",
|
||||
"--model",
|
||||
MODEL,
|
||||
"--session-dir",
|
||||
f"{PRIVATE}/sessions",
|
||||
"--tools",
|
||||
FULL_TOOLS,
|
||||
"--thinking",
|
||||
"high",
|
||||
]
|
||||
assert first.cwd == "/work"
|
||||
assert first.stdin == "hello"
|
||||
assert first.env == setup.env
|
||||
second: Final = CONFIG.transform_turn_request(ctx, setup, PRIVATE, "again", SESSION)
|
||||
assert flag(second.argv, "--session") == SESSION
|
||||
assert "again" not in " ".join(second.argv)
|
||||
|
||||
|
||||
def test_turn_prompt_repeats_schema_when_output_set():
|
||||
ctx: Final = make_ctx(output=Answer)
|
||||
request: Final = CONFIG.transform_turn_request(ctx, setup_for(ctx), PRIVATE, "hi", None)
|
||||
assert request.stdin.startswith("hi\n\n") and '"city"' in request.stdin
|
||||
|
||||
|
||||
def test_turn_response_paths():
|
||||
_, state = parse_all("structured_output.jsonl")
|
||||
ok: Final = CONFIG.transform_turn_response(make_ctx(output=Answer), state, 0, [])
|
||||
assert json.loads(ok.output_json) == {"city": "Paris", "country": "France"}
|
||||
plain: Final = CONFIG.transform_turn_response(make_ctx(), state, 0, [])
|
||||
assert plain.output_json is None and plain.final_text == state.final_text
|
||||
with pytest.raises(HarnessTurnError, match="boom"):
|
||||
CONFIG.transform_turn_response(make_ctx(), PiStreamState(stop_reason="error", error="boom"), 0, [])
|
||||
with pytest.raises(HarnessTurnError, match="aborted"):
|
||||
CONFIG.transform_turn_response(make_ctx(), PiStreamState(stop_reason="aborted"), 0, [])
|
||||
with pytest.raises(HarnessTurnError, match="code 3: no output"):
|
||||
CONFIG.transform_turn_response(make_ctx(), PiStreamState(), 3, [])
|
||||
|
||||
|
||||
def test_reply_cut_off_at_the_output_token_limit_is_an_error():
|
||||
state: Final = PiStreamState()
|
||||
parse(assistant_end("partial", "length"), state)
|
||||
with pytest.raises(HarnessTurnError, match="output token limit"):
|
||||
CONFIG.transform_turn_response(make_ctx(), state, 0, [])
|
||||
|
||||
|
||||
def test_exit_zero_without_a_reply_is_an_error():
|
||||
with pytest.raises(HarnessTurnError, match="without a reply: Session found in different project"):
|
||||
CONFIG.transform_turn_response(make_ctx(), PiStreamState(), 0, ["Session found in different project: /a"])
|
||||
|
||||
|
||||
async def test_start_writes_token_only_in_env():
|
||||
sandbox: Final = FakeSandbox(outputs=(fixture_proc("turn1_write_read.jsonl"),))
|
||||
ctx: Final = make_ctx(sandbox)
|
||||
handler: Final = await started(ctx)
|
||||
assert written_models(sandbox)["providers"]["litellm"]["apiKey"] == f"${PI_TOKEN_ENV}"
|
||||
assert all(TOKEN.encode() not in data for data in sandbox.files.values())
|
||||
await collect(handler, ctx, "create hello.txt containing hi then read it")
|
||||
assert TOKEN not in json.dumps(list(sandbox.execs[0].cmd))
|
||||
assert sandbox.execs[0].env[PI_TOKEN_ENV] == TOKEN
|
||||
|
||||
|
||||
async def test_start_persists_sessions_dir():
|
||||
sandbox: Final = FakeSandbox()
|
||||
await started(make_ctx(sandbox))
|
||||
assert sandbox.runs == (["sh", "-c", PERSIST_DIR_SCRIPT, "sh", "/tmp/pi-1/sessions", "pi/sessions"],)
|
||||
|
||||
|
||||
async def test_persist_failure_still_uses_private_sessions():
|
||||
sandbox: Final = FakeSandbox(persist_ok=False, outputs=(fixture_proc("turn1_write_read.jsonl"),))
|
||||
ctx: Final = make_ctx(sandbox)
|
||||
await collect(await started(ctx), ctx, "x")
|
||||
assert flag(sandbox.execs[0].cmd, "--session-dir") == "/tmp/pi-1/sessions"
|
||||
|
||||
|
||||
async def test_turn_argv_env_and_session_continuation():
|
||||
sandbox: Final = FakeSandbox(outputs=(fixture_proc("turn1_write_read.jsonl"), fixture_proc("turn2_resume.jsonl")))
|
||||
ctx: Final = make_ctx(sandbox, options=PiOptions(env={"FOO": "1"}))
|
||||
handler: Final = await started(ctx)
|
||||
events: Final = await collect(handler, ctx, "create hello.txt containing hi then read it")
|
||||
first: Final = sandbox.execs[0]
|
||||
assert first.cmd == [
|
||||
"pi",
|
||||
"--mode",
|
||||
"json",
|
||||
"--no-approve",
|
||||
"--no-skills",
|
||||
"--provider",
|
||||
"litellm",
|
||||
"--model",
|
||||
MODEL,
|
||||
"--session-dir",
|
||||
"/tmp/pi-1/sessions",
|
||||
"--tools",
|
||||
FULL_TOOLS,
|
||||
]
|
||||
assert first.cwd == "/work"
|
||||
assert first.env["PI_CODING_AGENT_DIR"] == "/tmp/pi-1/agent"
|
||||
assert first.env["PI_OFFLINE"] == "1" and first.env["PI_TELEMETRY"] == "0"
|
||||
assert first.env["FOO"] == "1"
|
||||
assert any(isinstance(e, ToolCall) for e in events)
|
||||
assert ctx.final_text == "Done! hello.txt contains: hi"
|
||||
|
||||
await collect(handler, ctx, "what file did you create?")
|
||||
second: Final = sandbox.execs[1].cmd
|
||||
assert flag(second, "--session") == SESSION
|
||||
assert "what file" not in " ".join(second)
|
||||
assert ctx.final_text == "You asked me to create hello.txt."
|
||||
|
||||
|
||||
async def test_prompt_is_sent_on_stdin_not_argv():
|
||||
proc: Final = fixture_proc("turn1_write_read.jsonl")
|
||||
sandbox: Final = FakeSandbox(outputs=(proc,))
|
||||
ctx: Final = make_ctx(sandbox)
|
||||
await collect(await started(ctx), ctx, "secret prompt text")
|
||||
assert proc.stdin.data == b"secret prompt text" and proc.stdin.closed
|
||||
assert "secret prompt text" not in sandbox.execs[0].cmd
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("workdir", "resume_flag"), [("/work", "--session"), ("/elsewhere", "--fork")])
|
||||
def test_session_saved_in_another_workdir_is_forked_into_this_one(workdir: str, resume_flag: str) -> None:
|
||||
_, state = parse_all("turn1_write_read.jsonl")
|
||||
ctx: Final = make_ctx(FakeSandbox(workdir=workdir))
|
||||
argv: Final = CONFIG.transform_turn_request(
|
||||
ctx, setup_for(ctx), PRIVATE, "again", CONFIG.get_native_session_id(state)
|
||||
).argv
|
||||
assert flag(argv, resume_flag) == SESSION
|
||||
assert ({"--session", "--fork"} - {resume_flag}).isdisjoint(argv), (
|
||||
"pi answers --session for a session saved in another directory with a fork prompt the harness never answers"
|
||||
)
|
||||
|
||||
|
||||
async def test_resume_sets_session():
|
||||
sandbox: Final = FakeSandbox(outputs=(fixture_proc("turn2_resume.jsonl"),))
|
||||
ctx: Final = make_ctx(sandbox)
|
||||
handler: Final = await started(ctx)
|
||||
await handler.resume(ctx, "prev-session")
|
||||
await collect(handler, ctx, "hi")
|
||||
assert flag(sandbox.execs[0].cmd, "--session") == "prev-session"
|
||||
|
||||
|
||||
async def test_read_only_and_disable_tools_argv():
|
||||
sandbox: Final = FakeSandbox(outputs=(fixture_proc("readonly_denied_bash.jsonl"),))
|
||||
ctx: Final = make_ctx(sandbox, permissions="read-only", disable_tools=("grep",))
|
||||
events: Final = await collect(await started(ctx), ctx, "x")
|
||||
cmd: Final = sandbox.execs[0].cmd
|
||||
assert flag(cmd, "--tools") == "read,grep,find,ls"
|
||||
assert flag(cmd, "--exclude-tools") == "grep"
|
||||
assert ToolResult(id="call_b1", output="Tool bash not found", is_error=True) in events
|
||||
|
||||
|
||||
async def test_instructions_and_structured_output():
|
||||
proc: Final = fixture_proc("structured_output.jsonl")
|
||||
sandbox: Final = FakeSandbox(outputs=(proc,))
|
||||
ctx: Final = make_ctx(sandbox, instructions="Be terse.", output=Answer)
|
||||
handler: Final = await started(ctx)
|
||||
written: Final = sandbox.files["/tmp/pi-1/instructions.md"].decode()
|
||||
assert written.startswith("Be terse.")
|
||||
assert '"city"' in written and "single JSON object" in written
|
||||
await collect(handler, ctx, "x")
|
||||
assert flag(sandbox.execs[0].cmd, "--append-system-prompt") == "/tmp/pi-1/instructions.md"
|
||||
assert '"city"' in proc.stdin.data.decode()
|
||||
assert json.loads(ctx.output_json) == {"city": "Paris", "country": "France"}
|
||||
|
||||
|
||||
async def test_skills_copied_to_private_skills_path(tmp_path: Path) -> None:
|
||||
skill: Final = tmp_path / "greeter"
|
||||
(skill / "ref").mkdir(parents=True)
|
||||
(skill / "SKILL.md").write_text("---\nname: greeter\ndescription: d\n---\nbody")
|
||||
(skill / "ref" / "notes.txt").write_text("n")
|
||||
sandbox: Final = FakeSandbox(outputs=(fixture_proc("turn2_resume.jsonl"),))
|
||||
ctx: Final = make_ctx(sandbox, skills=(str(skill),))
|
||||
handler: Final = await started(ctx)
|
||||
assert sandbox.files["/tmp/pi-1/skills/greeter/SKILL.md"].startswith(b"---")
|
||||
assert sandbox.files["/tmp/pi-1/skills/greeter/ref/notes.txt"] == b"n"
|
||||
await collect(handler, ctx, "x")
|
||||
cmd: Final = sandbox.execs[0].cmd
|
||||
assert flag(cmd, "--skill") == "/tmp/pi-1/skills"
|
||||
assert "--no-skills" in cmd
|
||||
|
||||
|
||||
async def test_skill_without_manifest_rejected(tmp_path: Path) -> None:
|
||||
with pytest.raises(ValueError, match=r"SKILL\.md"):
|
||||
await started(make_ctx(skills=(str(tmp_path),)))
|
||||
|
||||
|
||||
async def test_missing_binary():
|
||||
with pytest.raises(HarnessInstallFailed, match="@earendil-works/pi-coding-agent"):
|
||||
await started(make_ctx(FakeSandbox(has_binary=False)))
|
||||
|
||||
|
||||
async def test_start_missing_endpoint_raises():
|
||||
with pytest.raises(HarnessError):
|
||||
await started(make_ctx(endpoint=None))
|
||||
|
||||
|
||||
async def test_mcp_servers_written_and_reachable():
|
||||
sandbox: Final = FakeSandbox(outputs=(fixture_proc("mcp_tool.jsonl"),))
|
||||
ctx: Final = make_ctx(sandbox, options=PiOptions(config={"mcpServers": {"moyai": MOYAI}}))
|
||||
handler: Final = await started(ctx)
|
||||
assert json.loads(sandbox.files["/tmp/pi-1/agent/mcp.json"]) == {"mcpServers": {"moyai": MOYAI}}
|
||||
events: Final = await collect(handler, ctx, "call the moyai echo tool")
|
||||
assert flag(sandbox.execs[0].cmd, "--tools") == f"{FULL_TOOLS},mcp__*"
|
||||
assert ToolResult(id="call_m1", output="moyai says: hello", is_error=False) in events
|
||||
|
||||
|
||||
async def test_wrong_options_and_ask_mode_rejected():
|
||||
with pytest.raises(OptionsMismatch):
|
||||
await started(make_ctx(options=OpenCodeOptions()))
|
||||
with pytest.raises(OptionsMismatch):
|
||||
await started(make_ctx(options=PiOptions(config={"sessionDir": "/elsewhere"})))
|
||||
with pytest.raises(CapabilityUnsupported):
|
||||
await started(make_ctx(permissions="ask"))
|
||||
|
||||
|
||||
async def test_turn_before_start_raises():
|
||||
handler: Final = CLIHarnessHandler(PiHarnessConfig())
|
||||
with pytest.raises(RuntimeError, match="before start"):
|
||||
await collect(handler, make_ctx(), "x")
|
||||
|
||||
|
||||
async def test_api_error_raises_even_on_exit_zero():
|
||||
ctx: Final = make_ctx(FakeSandbox(outputs=(fixture_proc("api_error.jsonl", exit_code=0),)))
|
||||
handler: Final = await started(ctx)
|
||||
with pytest.raises(HarnessTurnError, match="no healthy deployments"):
|
||||
await collect(handler, ctx, "x")
|
||||
|
||||
|
||||
async def test_nonzero_exit_raises_with_stderr_tail():
|
||||
failing: Final = FakeProcess(b"", stderr=b"line1\nfatal: bad flag\n", exit_code=2)
|
||||
ctx: Final = make_ctx(FakeSandbox(outputs=(failing,)))
|
||||
handler: Final = await started(ctx)
|
||||
with pytest.raises(HarnessTurnError, match="code 2: line1\nfatal: bad flag"):
|
||||
await collect(handler, ctx, "x")
|
||||
|
||||
|
||||
async def test_early_close_kills_process_and_stop_is_idempotent():
|
||||
proc: Final = fixture_proc("turn1_write_read.jsonl")
|
||||
ctx: Final = make_ctx(FakeSandbox(outputs=(proc,)))
|
||||
handler: Final = await started(ctx)
|
||||
gen: Final = handler.turn(ctx, "x")
|
||||
await gen.__anext__()
|
||||
await gen.aclose()
|
||||
assert proc.killed
|
||||
await handler.stop(ctx)
|
||||
await handler.stop(ctx)
|
||||
|
||||
|
||||
async def test_long_jsonl_line_is_parsed():
|
||||
text: Final = "x" * 200_000
|
||||
lines: Final = (
|
||||
{"type": "session", "version": 3, "id": "s"},
|
||||
{"type": "message_update", "assistantMessageEvent": {"type": "text_delta", "delta": text}},
|
||||
assistant_end(text, "stop"),
|
||||
)
|
||||
stdout: Final = "\n".join(json.dumps(dict(line)) for line in lines).encode()
|
||||
ctx: Final = make_ctx(FakeSandbox(outputs=(FakeProcess(stdout),)))
|
||||
events: Final = await collect(await started(ctx), ctx, "x")
|
||||
assert events == [Text(delta=text)]
|
||||
assert ctx.final_text == text
|
||||
|
||||
|
||||
async def test_model_falls_back_to_endpoint_model():
|
||||
sandbox: Final = FakeSandbox(outputs=(fixture_proc("turn1_write_read.jsonl"),))
|
||||
ctx: Final = make_ctx(sandbox, model=None, endpoint=FakeEndpoint(model="gw-model"))
|
||||
await collect(await started(ctx), ctx, "x")
|
||||
assert flag(sandbox.execs[0].cmd, "--model") == "gw-model"
|
||||
assert written_models(sandbox)["providers"]["litellm"]["models"] == [{"id": "gw-model"}]
|
||||
|
||||
|
||||
def test_turn_request_never_trusts_project_files():
|
||||
ctx: Final = make_ctx()
|
||||
argv: Final = list(CONFIG.transform_turn_request(ctx, setup_for(ctx), PRIVATE, "hi", None).argv)
|
||||
assert argv[:4] == ["pi", "--mode", "json", "--no-approve"], (
|
||||
"without --no-approve a repo's .pi/extensions run as the host user at startup"
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue