pi-durable-python 1.3.0__py3-none-any.whl

pi_durable/testing/types.py

@@ -0,0 +1,99 @@
+"""Public types and protocols of the runner-independent conformance suites."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Sequence
+from dataclasses import dataclass
+from typing import Any, NotRequired, Protocol, TypedDict
+
+from ..env import ExecutionEnv
+from ..types import Storage
+
+ExpectLike = Callable[[Any, str | None], Any]
+"""A Vitest/Jest-style ``expect`` callable: it returns a matcher for the actual value."""
+
+
+class StorageConformanceAssertions(Protocol):
+    """The assertion vocabulary the storage conformance cases use."""
+
+    def ok(self, value: object, message: str | None = None) -> None: ...
+
+    def strict_equal(self, actual: object, expected: object) -> None: ...
+
+    def deep_equal(self, actual: object, expected: object) -> None: ...
+
+    def partial_deep_equal(self, actual: object, expected: object) -> None: ...
+
+    def greater_than(self, actual: float, expected: float) -> None: ...
+
+    async def rejects(self, operation: Awaitable[object], message_includes: str) -> None: ...
+
+
+StorageConformanceProvider = Callable[[Callable[[Storage], Awaitable[None]]], Awaitable[None]]
+"""Runs one case against a fresh storage, calling and awaiting ``use`` exactly once."""
+
+
+class StorageConformanceOptions(TypedDict):
+    assertions: StorageConformanceAssertions
+    with_storage: StorageConformanceProvider
+
+
+@dataclass(frozen=True, slots=True)
+class StorageConformanceCase:
+    """One runner-independent conformance case."""
+
+    name: str
+    run: Callable[[], Awaitable[None]]
+
+
+class StorageConformanceRunner(Protocol):
+    """The subset of a Vitest/Jest-compatible runner the registration helper needs."""
+
+    expect: ExpectLike
+
+    def describe(self, name: str, suite: Callable[[], object]) -> object: ...
+
+    def it(
+        self,
+        name: str,
+        test: Callable[[], Awaitable[None]],
+        timeout_ms: int | None = None,
+    ) -> object: ...
+
+
+EnvConformanceAssertions = StorageConformanceAssertions
+"""The assertion vocabulary the environment conformance cases use."""
+
+EnvConformanceProvider = Callable[[Callable[[ExecutionEnv], Awaitable[None]]], Awaitable[None]]
+"""Calls ``use`` once with an environment whose ``cwd`` is a fresh, empty, writable directory."""
+
+
+class EnvConformanceOptions(TypedDict):
+    assertions: EnvConformanceAssertions
+    with_env: EnvConformanceProvider
+    shell: NotRequired[Sequence[str]]
+    symlinks: NotRequired[bool]
+    list_files: NotRequired[bool]
+
+
+@dataclass(frozen=True, slots=True)
+class EnvConformanceCase:
+    """One runner-independent environment conformance case."""
+
+    name: str
+    run: Callable[[], Awaitable[None]]
+    timeout_ms: int | None = None
+
+
+__all__ = [
+    "EnvConformanceAssertions",
+    "EnvConformanceCase",
+    "EnvConformanceOptions",
+    "EnvConformanceProvider",
+    "ExpectLike",
+    "StorageConformanceAssertions",
+    "StorageConformanceCase",
+    "StorageConformanceOptions",
+    "StorageConformanceProvider",
+    "StorageConformanceRunner",
+]

pi_durable/tool_search/__init__.py

@@ -0,0 +1,159 @@
+"""The ``tool_search`` tool, a copy of the coding agent's.
+
+The model searches the deferred tools it is not offered yet, and the matches are loaded with
+``control.addTools``, so the next request offers them. Loading writes the conversation's
+``modelTools`` (its ``add``), so loaded tools stay offered on that branch, also after a compaction.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping
+from typing import Any, TypedDict, cast
+
+from ..context import Context
+from ..harness.define import define_extension, define_tool
+from ..harness.types import Extension, ToolExecutionApi, ToolExecutionResult, ToolRegistration
+from ..json import JsonValue
+from ..search import SearchDocument, rank_documents
+
+TOOL_SEARCH_TOOL_NAME = "tool_search"
+DEFAULT_TOOL_SEARCH_LIMIT = 8
+
+#: Fixed, so it does not change while deferred tools are registered, for example as MCP servers
+#: connect.
+TOOL_SEARCH_DESCRIPTION = (
+    "# Tool discovery\n\n"
+    "Searches over deferred tool metadata with BM25 and exposes matching tools for the next model "
+    "call.\n\n"
+    "Some of the tools, such as tools of MCP servers, may not have been provided to you upfront, "
+    f"and you should use this tool (`{TOOL_SEARCH_TOOL_NAME}`) to search for the required tools. "
+    f"For MCP tool discovery, always use `{TOOL_SEARCH_TOOL_NAME}`."
+)
+
+TOOL_SEARCH_SCHEMA: dict[str, Any] = {
+    "type": "object",
+    "properties": {
+        "query": {
+            "type": "string",
+            "description": "Search query for deferred tools.",
+        },
+        "limit": {
+            "type": "number",
+            "description": (f"Maximum number of tools to return. Defaults to {DEFAULT_TOOL_SEARCH_LIMIT}."),
+        },
+    },
+    "required": ["query"],
+}
+
+
+class ToolSearchToolDetails(TypedDict):
+    """``tool_search`` result details: the tools it loaded."""
+
+    loaded: list[str]
+
+
+def _is_object(value: Any) -> bool:
+    return isinstance(value, Mapping)
+
+
+def _schema_text(schema: Any, parts: list[str]) -> None:
+    """Schema descriptions and property names, recursively."""
+
+    if not _is_object(schema):
+        return
+    description = schema.get("description")
+    if isinstance(description, str):
+        parts.append(description)
+    properties = schema.get("properties")
+    if _is_object(properties):
+        for name, property_schema in properties.items():
+            parts.append(name)
+            _schema_text(property_schema, parts)
+    _schema_text(schema.get("items"), parts)
+    for key in ("anyOf", "oneOf", "allOf"):
+        variants = schema.get(key)
+        if isinstance(variants, (list, tuple)):
+            for variant in variants:
+                _schema_text(variant, parts)
+
+
+def _search_text(tool: ToolRegistration) -> str:
+    """A tool's name, the name with ``_`` as spaces, its description, and its parameters' text."""
+
+    parts = [tool.name, tool.name.replace("_", " "), tool.description]
+    _schema_text(tool.parameters, parts)
+    return " ".join(part for part in parts if part.strip() != "")
+
+
+def _positive_integer(value: Any) -> int | None:
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        return None
+    if isinstance(value, float) and (not math.isfinite(value) or not value.is_integer()):
+        return None
+    return int(value) if value > 0 else None
+
+
+def _first_line(description: str) -> str:
+    return description.strip().split("\n", 1)[0].rstrip("\r")
+
+
+async def _execute(args: Mapping[str, Any], api: ToolExecutionApi, context: Context) -> ToolExecutionResult:
+    query = args.get("query")
+    if not isinstance(query, str) or query.strip() == "":
+        raise ValueError("query must not be empty")
+    raw_limit = args.get("limit")
+    maximum = DEFAULT_TOOL_SEARCH_LIMIT if raw_limit is None else _positive_integer(raw_limit)
+    if maximum is None:
+        raise ValueError("limit must be a positive integer")
+    candidates = list((await api.agent(context)).deferred)
+    documents: list[SearchDocument] = [{"name": tool.name, "text": _search_text(tool)} for tool in candidates]
+    loaded = rank_documents(query, documents, maximum)
+    if not loaded:
+        text = "No matching tools found."
+    else:
+        by_name = {tool.name: tool for tool in candidates}
+        lines = [f"- {name}: {_first_line(by_name[name].description)}" for name in loaded]
+        text = (
+            f"Loaded {len(loaded)} tool{'' if len(loaded) == 1 else 's'}. "
+            "They are available from your next call:\n" + "\n".join(lines)
+        )
+    details: ToolSearchToolDetails = {"loaded": loaded}
+    result: ToolExecutionResult = {
+        "output": [{"type": "text", "text": text}],
+        "details": cast(JsonValue, details),
+    }
+    if loaded:
+        result["control"] = {"addTools": loaded}
+    return result
+
+
+def create_tool_search_tool() -> ToolRegistration:
+    """The ``tool_search`` tool: rank ``Agent.deferred`` with BM25 and load the matches."""
+
+    return define_tool(
+        {
+            "name": TOOL_SEARCH_TOOL_NAME,
+            "description": TOOL_SEARCH_DESCRIPTION,
+            "parameters": TOOL_SEARCH_SCHEMA,
+            # Searching changes what the model is offered; scripts have their own `searchTools()`.
+            "callers": ["model"],
+            "replay": "safe",
+            "execute": _execute,
+        }
+    )
+
+
+def create_tool_search_extension() -> Extension:
+    """The ``tool_search`` extension: select it where the model should find deferred tools."""
+
+    return define_extension({"name": "tool-search", "tools": [create_tool_search_tool()]})
+
+
+__all__ = [
+    "DEFAULT_TOOL_SEARCH_LIMIT",
+    "TOOL_SEARCH_TOOL_NAME",
+    "ToolSearchToolDetails",
+    "create_tool_search_extension",
+    "create_tool_search_tool",
+]

pi_durable/tools/__init__.py

@@ -0,0 +1,81 @@
+"""Built-in coding tools.
+
+Port of the v1.3.0 ``src/tools/index.ts``. ``read``, ``write``, ``edit``, and ``bash`` use only
+``api.env``; nothing installs them automatically. PowerShell is excluded, and there is no image
+implementation: ``read`` optionally takes an :class:`~pi_durable.tools.ImageProcessor`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from ..harness.define import define_extension
+from ..harness.types import Extension
+from .bash import (
+    BashExecution,
+    BashPrepare,
+    BashToolInput,
+    BashToolOptions,
+    create_bash_tool,
+)
+from .edit import EditToolDetails, EditToolInput, create_edit_tool
+from .image_processor import (
+    DEFAULT_IMAGE_LIMITS,
+    ImageLimits,
+    ImageProcessor,
+    PreparedImage,
+)
+from .read import ReadToolDetails, ReadToolInput, ReadToolOptions, create_read_tool
+from .write import WriteToolInput, create_write_tool
+
+
+def create_coding_tools(
+    options: ReadToolOptions | Mapping[str, Any] | None = None,
+    *,
+    images: ImageProcessor | None = None,
+) -> Extension:
+    """The ``coding-tools`` extension: ``read``, ``write``, ``edit``, and ``bash``.
+
+    ``images`` prepares the images ``read`` returns.
+    """
+
+    return define_extension(
+        {
+            "name": "coding-tools",
+            "tools": [
+                create_read_tool(options, images=images),
+                create_write_tool(),
+                create_edit_tool(),
+                create_bash_tool(),
+            ],
+        }
+    )
+
+
+#: :func:`create_coding_tools` without an image processor.
+CodingTools = create_coding_tools()
+
+
+__all__ = [
+    "BashExecution",
+    "BashPrepare",
+    "BashToolInput",
+    "BashToolOptions",
+    "CodingTools",
+    "DEFAULT_IMAGE_LIMITS",
+    "EditToolDetails",
+    "EditToolInput",
+    "ImageLimits",
+    "ImageProcessor",
+    "PreparedImage",
+    "ReadToolDetails",
+    "ReadToolInput",
+    "ReadToolOptions",
+    "WriteToolInput",
+    "create_bash_tool",
+    "create_coding_tools",
+    "create_edit_tool",
+    "create_read_tool",
+    "create_write_tool",
+]

pi_durable/tools/_values.py

@@ -0,0 +1,43 @@
+"""Small readers shared by the coding tools."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping
+from typing import Any
+
+
+def field(value: Any, name: str, default: Any = None) -> Any:
+    """A field of a mapping or object, as the wire and Python shapes vary."""
+
+    if isinstance(value, Mapping):
+        return value.get(name, default)
+    return getattr(value, name, default)
+
+
+def field_either(value: Any, snake: str, camel: str, default: Any = None) -> Any:
+    """A snake_case field, falling back to the upstream camelCase spelling."""
+
+    found = field(value, snake)
+    return field(value, camel, default) if found is None else found
+
+
+def js_number_string(value: Any) -> str:
+    """``String(number)`` for the numbers tool arguments may hold."""
+
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    if isinstance(value, int):
+        return str(value)
+    if isinstance(value, float):
+        if math.isnan(value):
+            return "NaN"
+        if math.isinf(value):
+            return "Infinity" if value > 0 else "-Infinity"
+        if value.is_integer() and abs(value) < 1e21:
+            return str(int(value))
+        return repr(value)
+    return str(value)
+
+
+__all__ = ["field", "field_either", "js_number_string"]

pi_durable/tools/bash.py

@@ -0,0 +1,250 @@
+"""The ``bash`` coding tool.
+
+Port of the v1.3.0 ``src/tools/bash.ts`` without the excluded PowerShell tool. The command runs
+through the call's execution environment, streams into the Harness's retained output window, spills
+oversized output, and reports a nonzero exit as an error result.
+"""
+
+from __future__ import annotations
+
+import inspect
+import math
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from dataclasses import dataclass
+from dataclasses import field as dataclass_field
+from typing import Any, TypedDict, cast
+
+from ..context import Context
+from ..env import ExecutionError, Result, ShellExecOptions, ShellExecResult
+from ..harness.define import define_tool
+from ..harness.types import ToolDiagnostic, ToolExecutionApi, ToolExecutionResult, ToolRegistration
+from ..truncate import DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES
+from ._values import field, js_number_string
+from .env import require_env
+
+MAX_TIMEOUT_SECONDS = 2_147_483_647 / 1000
+
+
+class BashToolInput(TypedDict):
+    """Arguments of one ``bash`` call."""
+
+    command: str
+    timeout: float
+
+
+@dataclass(slots=True)
+class BashExecution:
+    """A command about to run, which ``prepare`` may change."""
+
+    command: str = ""
+    cwd: str = ""
+    env: dict[str, str] = dataclass_field(default_factory=dict)
+    inherit_env: bool = True
+
+
+BashPrepare = Callable[[BashExecution, ToolExecutionApi, Context], Awaitable[None] | None]
+
+
+class BashToolOptions(TypedDict, total=False):
+    """Options of :func:`create_bash_tool`."""
+
+    command_prefix: str
+    prepare: BashPrepare
+
+
+BASH_SCHEMA: dict[str, Any] = {
+    "type": "object",
+    "properties": {
+        "command": {"type": "string", "description": "Bash command to execute"},
+        "timeout": {
+            "type": "number",
+            "description": "Timeout in seconds (optional, no default timeout)",
+        },
+    },
+    "required": ["command"],
+}
+
+SHELL_STRUCTURED_OUTPUT_SCHEMA: dict[str, Any] = {
+    "type": "object",
+    "properties": {
+        "output": {
+            "type": "string",
+            "description": "Combined stdout and stderr: the retained tail, as the model sees it",
+        },
+        "truncated": {
+            "type": "boolean",
+            "description": "Whether earlier output was dropped",
+        },
+        "fullOutputPath": {
+            "type": "string",
+            "description": "File with the complete output, when it was spilled",
+        },
+        "exitCode": {"type": "number"},
+    },
+    "required": ["output", "truncated", "exitCode"],
+}
+
+
+def _validate_timeout(timeout: float | None) -> None:
+    if timeout is None:
+        return
+    if (
+        isinstance(timeout, bool)
+        or not isinstance(timeout, (int, float))
+        or not math.isfinite(timeout)
+        or timeout <= 0
+    ):
+        raise ValueError("Invalid timeout: must be a finite number of seconds")
+    if timeout > MAX_TIMEOUT_SECONDS:
+        raise ValueError(f"Invalid timeout: maximum is {js_number_string(MAX_TIMEOUT_SECONDS)} seconds")
+
+
+async def _prepare_execution(
+    command: str,
+    options: Mapping[str, Any] | None,
+    api: ToolExecutionApi,
+    context: Context,
+) -> BashExecution:
+    values = dict(options or {})
+    prefix = values.get("command_prefix", values.get("commandPrefix"))
+    execution = BashExecution(
+        command=f"{prefix}\n{command}" if prefix else command,
+        cwd=require_env(api).cwd,
+        env={},
+        inherit_env=True,
+    )
+    prepare = values.get("prepare")
+    if prepare is not None:
+        prepared = prepare(execution, api, context)
+        if inspect.isawaitable(prepared):
+            await prepared
+    return execution
+
+
+async def _run_command(
+    commands: Sequence[str | Sequence[str]],
+    execution: BashExecution,
+    timeout: float | None,
+    api: ToolExecutionApi,
+    context: Context,
+) -> ToolExecutionResult:
+    env = require_env(api)
+    result: Result[ShellExecResult, ExecutionError] | None = None
+    for command in commands:
+        options: dict[str, Any] = {
+            "cwd": execution.cwd,
+            "env": dict(execution.env),
+            "inherit_env": execution.inherit_env,
+            "on_output": lambda text, _context, info: api.output(text, field(info, "skipped")),
+            "spill": {"after_bytes": DEFAULT_MAX_BYTES, "after_lines": DEFAULT_MAX_LINES},
+        }
+        if timeout is not None:
+            options["timeout"] = timeout
+        # An environment may then omit output outside the retained tail and report the omission.
+        output_window = getattr(api, "output_window", None)
+        if output_window is not None:
+            options["window"] = output_window
+        result = await env.exec(command, cast(ShellExecOptions, options), context)
+        # A program that could not start produced no output; the next form may.
+        if result.ok or result.error.code != "spawn_error":
+            break
+    if result is None:
+        raise RuntimeError("No command to run")
+    spill_path = (
+        field(result.value, "spillPath", field(result.value, "spill_path"))
+        if result.ok
+        else field(result.error, "spill_path", field(result.error, "spillPath"))
+    )
+    if spill_path is not None:
+        api.diagnostic(
+            ToolDiagnostic(severity="info", code="full_output", message=f"Full output: {spill_path}")
+        )
+    if not result.ok:
+        error = result.error
+        if error.code == "aborted" and context.abort_signal is not None and context.abort_signal.aborted:
+            raise error
+        if error.code == "timeout":
+            raise RuntimeError(f"Command timed out after {js_number_string(timeout)} seconds")
+        if error.code == "aborted":
+            raise RuntimeError("Command aborted")
+        raise error
+    exit_code = result.value.get("exitCode")
+    retained = api.retained_output()
+    structured: dict[str, Any] = {
+        "output": field(retained, "text", ""),
+        "truncated": bool(field(retained, "truncated", False)),
+        "exitCode": exit_code,
+    }
+    if spill_path is not None:
+        structured["fullOutputPath"] = spill_path
+    if exit_code == 0:
+        return {"structuredOutput": structured}
+    return {
+        "structuredOutput": structured,
+        "isError": True,
+        "diagnostics": [
+            ToolDiagnostic(
+                severity="error",
+                code="exit_code",
+                message=f"Command exited with code {js_number_string(exit_code)}",
+            )
+        ],
+    }
+
+
+def create_bash_tool(
+    options: BashToolOptions | Mapping[str, Any] | None = None,
+    *,
+    command_prefix: str | None = None,
+    prepare: BashPrepare | None = None,
+) -> ToolRegistration:
+    """Run a command through the environment's shell.
+
+    Its output streams to ``api.output()``, where the Harness keeps the tail within the default
+    limits; the result content is that retained output. Output beyond the limits is spilled to a
+    file whose path is reported as a diagnostic. A nonzero exit is an error result with the exit
+    code; a timeout throws.
+    """
+
+    values: dict[str, Any] = dict(options or {})
+    if command_prefix is not None:
+        values["command_prefix"] = command_prefix
+    if prepare is not None:
+        values["prepare"] = prepare
+
+    async def execute(
+        args: Mapping[str, Any], api: ToolExecutionApi, context: Context
+    ) -> ToolExecutionResult:
+        timeout = args.get("timeout")
+        _validate_timeout(timeout)
+        execution = await _prepare_execution(str(args["command"]), values, api, context)
+        return await _run_command([execution.command], execution, timeout, api, context)
+
+    return define_tool(
+        {
+            "name": "bash",
+            "description": (
+                "Execute a bash command in the current working directory. Returns combined stdout "
+                "and stderr. Output is truncated to last "
+                f"{DEFAULT_MAX_LINES} lines or {DEFAULT_MAX_BYTES // 1024}KB (whichever is hit "
+                "first). If truncated, full output is saved to a temp file. Optionally provide a "
+                "timeout in seconds."
+            ),
+            "parameters": BASH_SCHEMA,
+            "structured_output_schema": SHELL_STRUCTURED_OUTPUT_SCHEMA,
+            "output_limits": {"retain": "tail"},
+            "execute": execute,
+        }
+    )
+
+
+__all__ = [
+    "BASH_SCHEMA",
+    "MAX_TIMEOUT_SECONDS",
+    "SHELL_STRUCTURED_OUTPUT_SCHEMA",
+    "BashExecution",
+    "BashPrepare",
+    "BashToolInput",
+    "BashToolOptions",
+    "create_bash_tool",
+]