diff --git a/docs/rpc.md b/docs/rpc.md index 006c30395..53a1d37de 100644 --- a/docs/rpc.md +++ b/docs/rpc.md @@ -42,13 +42,15 @@ There is no envelope beyond the object shape itself. 3. `AgentSessionEvent` objects (`agent_start`, `message_update`, etc.) 4. `RpcExtensionUIRequest` (`{ type: "extension_ui_request", ... }`) 5. Host tool requests/cancellations (`host_tool_call`, `host_tool_cancel`) -6. Extension errors (`{ type: "extension_error", extensionPath, event, error }`) +6. Host URI requests/cancellations (`host_uri_request`, `host_uri_cancel`) +7. Extension errors (`{ type: "extension_error", extensionPath, event, error }`) ### Inbound frame categories (stdin) 1. `RpcCommand` 2. `RpcExtensionUIResponse` (`{ type: "extension_ui_response", ... }`) 3. Host tool updates/results (`host_tool_update`, `host_tool_result`) +4. Host URI results (`host_uri_result`) ## Request/Response Correlation @@ -81,6 +83,7 @@ Important edge behavior from runtime: - `{ id?, type: "get_state" }` - `{ id?, type: "set_todos", phases: TodoPhase[] }` - `{ id?, type: "set_host_tools", tools: RpcHostToolDefinition[] }` +- `{ id?, type: "set_host_uri_schemes", schemes: RpcHostUriSchemeDefinition[] }` ### Model @@ -247,6 +250,38 @@ The response payload is: These tools are added to the active session tool registry before the next model call. Re-sending `set_host_tools` replaces the previous host-owned set. +### `set_host_uri_schemes` payload + +Replaces the current set of host-owned URL schemes the RPC server should +dispatch reads/writes through: + +```json +{ + "id": "req_4", + "type": "set_host_uri_schemes", + "schemes": [ + { + "scheme": "db", + "description": "Virtual db row files", + "writable": true, + "immutable": false + } + ] +} +``` + +The response payload is: + +```json +{ + "schemes": ["db"] +} +``` + +Schemes are case-insensitive on the wire and normalized to lowercase before +the response is sent. Re-sending `set_host_uri_schemes` replaces the entire +previous set — schemes missing from the new list are unregistered. + ## Event Stream Schema RPC mode forwards `AgentSessionEvent` objects from `AgentSession.subscribe(...)`. @@ -417,6 +452,82 @@ Completion uses: Set top-level `isError: true` on `host_tool_result` to reject the pending host tool call and surface the returned text content as a tool error. +## Host URI Sub-Protocol + +RPC hosts can also own custom URL schemes (virtual files). After +`set_host_uri_schemes`, every read of `://…` and write of +`://…` (when registered as `writable`) is bounced back to the host +over the same transport. + +### Outbound request + +When a session tool resolves a host-owned URL, RPC mode emits: + +```json +{ + "type": "host_uri_request", + "id": "uri_1", + "operation": "read", + "url": "db://users/42" +} +``` + +Writes look the same with `"operation": "write"` and an additional +`"content": "..."` field carrying the full replacement bytes. + +If the request is later aborted (caller cancels, session ends), RPC mode +emits: + +```json +{ + "type": "host_uri_cancel", + "id": "uri_cancel_1", + "targetId": "uri_1" +} +``` + +### Inbound result + +For successful reads: + +```json +{ + "type": "host_uri_result", + "id": "uri_1", + "content": "id=42\nname=Alice\n", + "contentType": "text/plain", + "notes": ["fresh from cache"], + "immutable": false +} +``` + +For successful writes, omit content: + +```json +{ "type": "host_uri_result", "id": "uri_1" } +``` + +To reject the request, set `isError: true` and either populate `error` with +a message or fall back to `content` for textual error surfacing: + +```json +{ + "type": "host_uri_result", + "id": "uri_1", + "isError": true, + "error": "row 42 not found" +} +``` + +### Constraints + +- The agent's `edit` tool does not target host URIs. Hosts that want to + mutate virtual files expose `write` and let the model use the `write` tool + with replacement content. +- Schemes are global to the process; `set_host_uri_schemes` replaces the + previous set, unregistering anything not in the new list. +- Schemes are normalized to lowercase before registration. + ## Error Model and Recoverability ### Command-level failures diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index ed21b5edd..3aa4adaa6 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,9 +1,12 @@ # Changelog ## [Unreleased] - ### Added +- Added the `set_host_uri_schemes` RPC command so hosts can register and replace writable/read-only internal URI schemes with scheme metadata (`writable`, `immutable`) at runtime +- Enabled the `write` tool to dispatch `write(url, content)` to registered internal URL handlers, allowing edits to non-filesystem resources via host-managed URI schemes +- Added host-owned internal URI read/write over RPC, including abort support, so URI operations are resolved by the host transport for `read` and `write` requests +- Added handling of host URI request results in RPC mode so host services can stream completion frames for internal URI operations - Added scratch-directory awareness to the status-line `path` segment. When the project directory is inside an OS-level scratch root (the platform `os.tmpdir()`, `/tmp` and `/var/tmp` plus their macOS `/private/...` aliases, `~/tmp`, or — on Windows — `%TEMP%` / `%TMP%` / `%SystemRoot%\Temp`), the segment now (1) renders the new `icon.scratchFolder` symbol instead of `icon.folder`, and (2) strips the scratch root from the displayed path so only the trailing folder (and any subpath beneath it) is shown — mirroring how `/work` and `~/Projects` are already abbreviated. Both behaviors honor the existing `stripWorkPrefix` option. Icon defaults: 🗑 (emoji), `` (nf-fa-trash) for Nerd Font, `[T]` for ASCII, `◌` in the poimandres themes; themes can override `icon.scratchFolder` independently of `icon.folder`. ### Changed diff --git a/packages/coding-agent/src/internal-urls/router.ts b/packages/coding-agent/src/internal-urls/router.ts index 676252a27..6529d8063 100644 --- a/packages/coding-agent/src/internal-urls/router.ts +++ b/packages/coding-agent/src/internal-urls/router.ts @@ -50,6 +50,14 @@ export class InternalUrlRouter { this.#handlers.set(handler.scheme.toLowerCase(), handler); } + unregister(scheme: string): boolean { + return this.#handlers.delete(scheme.toLowerCase()); + } + + getHandler(scheme: string): ProtocolHandler | undefined { + return this.#handlers.get(scheme.toLowerCase()); + } + canHandle(input: string): boolean { const match = input.match(/^([a-z][a-z0-9+.-]*):\/\//i); if (!match) return false; diff --git a/packages/coding-agent/src/internal-urls/types.ts b/packages/coding-agent/src/internal-urls/types.ts index d0d29516b..8c0b71020 100644 --- a/packages/coding-agent/src/internal-urls/types.ts +++ b/packages/coding-agent/src/internal-urls/types.ts @@ -63,6 +63,18 @@ export interface ResolveContext { signal?: AbortSignal; } +/** + * Caller context for write operations dispatched to host-owned URI handlers. + * Mirrors {@link ResolveContext} so handlers that share read/write state can + * accept the same shape. + */ +export interface WriteContext { + /** Working directory of the calling session. */ + cwd?: string; + /** Caller's abort signal. */ + signal?: AbortSignal; +} + /** * Handler for a specific internal URL scheme (e.g., agent://, memory://, skill://, mcp://). */ @@ -86,4 +98,13 @@ export interface ProtocolHandler { * @throws Error with user-friendly message if resolution fails */ resolve(url: InternalUrl, context?: ResolveContext): Promise; + /** + * Optional write hook. When present, the write tool dispatches + * `write(url, content)` to this handler instead of writing to a filesystem + * path. The handler is responsible for any persistence and validation. + * + * Handlers that omit this method are treated as read-only; the write tool + * surfaces a clear "not writable" error when invoked against them. + */ + write?(url: InternalUrl, content: string, context?: WriteContext): Promise; } diff --git a/packages/coding-agent/src/modes/rpc/host-uris.ts b/packages/coding-agent/src/modes/rpc/host-uris.ts new file mode 100644 index 000000000..253da55e7 --- /dev/null +++ b/packages/coding-agent/src/modes/rpc/host-uris.ts @@ -0,0 +1,235 @@ +import { Snowflake } from "@oh-my-pi/pi-utils"; +import { InternalUrlRouter } from "../../internal-urls"; +import type { + InternalResource, + InternalUrl, + ProtocolHandler, + ResolveContext, + WriteContext, +} from "../../internal-urls/types"; +import type { + RpcHostUriCancelRequest, + RpcHostUriRequest, + RpcHostUriResult, + RpcHostUriSchemeDefinition, +} from "./rpc-types"; + +type RpcHostUriOutput = (frame: RpcHostUriRequest | RpcHostUriCancelRequest) => void; + +type PendingUriRequest = { + operation: "read" | "write"; + url: string; + resolve: (frame: RpcHostUriResult) => void; + reject: (error: Error) => void; +}; + +/** Type guard for inbound `host_uri_result` frames coming from the host. */ +export function isRpcHostUriResult(value: unknown): value is RpcHostUriResult { + if (!value || typeof value !== "object") return false; + const frame = value as { type?: unknown; id?: unknown }; + return frame.type === "host_uri_result" && typeof frame.id === "string"; +} + +/** + * One handler instance per host-registered scheme. Delegates reads and (when + * the scheme was registered as writable) writes to the bridge, which serializes + * them over the RPC transport. + */ +class RpcHostUriProtocolHandler implements ProtocolHandler { + readonly scheme: string; + readonly immutable: boolean; + readonly write?: (url: InternalUrl, content: string, context?: WriteContext) => Promise; + readonly #bridge: RpcHostUriBridge; + + constructor(definition: RpcHostUriSchemeDefinition, bridge: RpcHostUriBridge) { + this.scheme = definition.scheme; + this.immutable = definition.immutable === true; + this.#bridge = bridge; + if (definition.writable === true) { + this.write = (url, content, context) => this.#bridge.requestWrite(this.scheme, url, content, context); + } + } + + resolve(url: InternalUrl, context?: ResolveContext): Promise { + return this.#bridge.requestRead(this.scheme, url, context); + } +} + +/** + * Bidirectional bridge that lets the RPC host own a set of URI schemes. + * + * The host registers schemes via `set_host_uri_schemes`; the bridge installs + * a `RpcHostUriProtocolHandler` per scheme into the process-global + * {@link InternalUrlRouter}. Reads land on the read tool through the existing + * router; writes are intercepted by the write tool and dispatched through + * `requestWrite`. + */ +export class RpcHostUriBridge { + #output: RpcHostUriOutput; + #router: InternalUrlRouter; + #definitions = new Map(); + #pending = new Map(); + + constructor(output: RpcHostUriOutput, router: InternalUrlRouter = InternalUrlRouter.instance()) { + this.#output = output; + this.#router = router; + } + + getSchemes(): string[] { + return Array.from(this.#definitions.keys()); + } + + /** + * Replace the registered set of host URI schemes. Previously registered + * schemes that no longer appear in the new set are unregistered from the + * router; surviving and new schemes get fresh handler instances. + */ + setSchemes(schemes: RpcHostUriSchemeDefinition[]): string[] { + const normalized = new Map(); + for (const raw of schemes) { + const scheme = typeof raw?.scheme === "string" ? raw.scheme.trim().toLowerCase() : ""; + if (!scheme) { + throw new Error("Host URI scheme must be a non-empty string"); + } + if (!/^[a-z][a-z0-9+.-]*$/.test(scheme)) { + throw new Error(`Host URI scheme contains invalid characters: ${raw.scheme}`); + } + normalized.set(scheme, { + scheme, + description: typeof raw.description === "string" ? raw.description : undefined, + writable: raw.writable === true, + immutable: raw.immutable === true, + }); + } + + for (const previous of this.#definitions.keys()) { + if (!normalized.has(previous)) { + this.#router.unregister(previous); + } + } + for (const definition of normalized.values()) { + this.#router.register(new RpcHostUriProtocolHandler(definition, this)); + } + this.#definitions = normalized; + return Array.from(normalized.keys()); + } + + /** + * Unregister every host scheme from the router and reject any in-flight + * requests. Called on RPC shutdown to keep the global router clean for + * subsequent sessions in the same process (used by tests). + */ + clear(message: string = "Host URI bridge shut down"): void { + for (const scheme of this.#definitions.keys()) { + this.#router.unregister(scheme); + } + this.#definitions.clear(); + this.rejectAllPending(message); + } + + /** Resolve a pending request by id; called by `rpc-mode` on inbound results. */ + handleResult(frame: RpcHostUriResult): boolean { + const pending = this.#pending.get(frame.id); + if (!pending) return false; + this.#pending.delete(frame.id); + pending.resolve(frame); + return true; + } + + rejectAllPending(message: string): void { + const error = new Error(message); + const pending = Array.from(this.#pending.values()); + this.#pending.clear(); + for (const entry of pending) { + entry.reject(error); + } + } + + async requestRead(scheme: string, url: InternalUrl, context?: ResolveContext): Promise { + const result = await this.#dispatch("read", url.href, undefined, context?.signal); + if (result.isError) { + throw new Error(result.error || result.content || `Host URI read failed for ${url.href}`); + } + const content = result.content ?? ""; + const contentType = result.contentType ?? "text/plain"; + const definition = this.#definitions.get(scheme); + return { + url: url.href, + content, + contentType, + size: Buffer.byteLength(content, "utf-8"), + notes: result.notes && result.notes.length > 0 ? [...result.notes] : undefined, + immutable: result.immutable ?? definition?.immutable === true, + }; + } + + async requestWrite(_scheme: string, url: InternalUrl, content: string, context?: WriteContext): Promise { + const result = await this.#dispatch("write", url.href, content, context?.signal); + if (result.isError) { + throw new Error(result.error || result.content || `Host URI write failed for ${url.href}`); + } + } + + #dispatch( + operation: "read" | "write", + url: string, + content: string | undefined, + signal: AbortSignal | undefined, + ): Promise { + if (signal?.aborted) { + return Promise.reject(new Error(`Host URI ${operation} for ${url} was aborted`)); + } + + const id = Snowflake.next() as string; + const { promise, resolve, reject } = Promise.withResolvers(); + let settled = false; + + const cleanup = () => { + signal?.removeEventListener("abort", onAbort); + this.#pending.delete(id); + }; + + const onAbort = () => { + if (settled) return; + settled = true; + cleanup(); + this.#output({ + type: "host_uri_cancel", + id: Snowflake.next() as string, + targetId: id, + }); + reject(new Error(`Host URI ${operation} for ${url} was aborted`)); + }; + + signal?.addEventListener("abort", onAbort, { once: true }); + this.#pending.set(id, { + operation, + url, + resolve: frame => { + if (settled) return; + settled = true; + cleanup(); + resolve(frame); + }, + reject: err => { + if (settled) return; + settled = true; + cleanup(); + reject(err); + }, + }); + + const frame: RpcHostUriRequest = { + type: "host_uri_request", + id, + operation, + url, + }; + if (operation === "write") { + frame.content = content ?? ""; + } + this.#output(frame); + + return promise; + } +} diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index a02886066..0a2f69a01 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -21,6 +21,7 @@ import { type Theme, theme } from "../../modes/theme/theme"; import type { AgentSession } from "../../session/agent-session"; import { initializeExtensions } from "../runtime-init"; import { isRpcHostToolResult, isRpcHostToolUpdate, RpcHostToolBridge } from "./host-tools"; +import { isRpcHostUriResult, RpcHostUriBridge } from "./host-uris"; import type { RpcCommand, RpcExtensionUIRequest, @@ -28,6 +29,8 @@ import type { RpcHostToolCallRequest, RpcHostToolCancelRequest, RpcHostToolDefinition, + RpcHostUriCancelRequest, + RpcHostUriRequest, RpcResponse, RpcSessionState, } from "./rpc-types"; @@ -41,7 +44,14 @@ export type PendingExtensionRequest = { }; type RpcOutput = ( - obj: RpcResponse | RpcExtensionUIRequest | RpcHostToolCallRequest | RpcHostToolCancelRequest | object, + obj: + | RpcResponse + | RpcExtensionUIRequest + | RpcHostToolCallRequest + | RpcHostToolCancelRequest + | RpcHostUriRequest + | RpcHostUriCancelRequest + | object, ) => void; function normalizeHostToolDefinitions(tools: RpcHostToolDefinition[]): RpcHostToolDefinition[] { @@ -188,6 +198,7 @@ export async function runRpcMode( const pendingExtensionRequests = new Map(); const hostToolBridge = new RpcHostToolBridge(output); + const hostUriBridge = new RpcHostUriBridge(output); // Shutdown request flag (wrapped in object to allow mutation with const) const shutdownState = { requested: false }; @@ -533,6 +544,15 @@ export async function runRpcMode( return success(id, "set_host_tools", { toolNames: tools.map(tool => tool.name) }); } + case "set_host_uri_schemes": { + try { + const schemes = hostUriBridge.setSchemes(command.schemes); + return success(id, "set_host_uri_schemes", { schemes }); + } catch (err) { + return error(id, "set_host_uri_schemes", err instanceof Error ? err.message : String(err)); + } + } + // ================================================================= // Model // ================================================================= @@ -807,6 +827,11 @@ export async function runRpcMode( continue; } + if (isRpcHostUriResult(parsed)) { + hostUriBridge.handleResult(parsed); + continue; + } + // Handle regular commands const command = parsed as RpcCommand; const response = await handleCommand(command); @@ -821,5 +846,6 @@ export async function runRpcMode( // stdin closed — RPC client is gone, exit cleanly hostToolBridge.rejectAllPending("RPC client disconnected before host tool execution completed"); + hostUriBridge.clear("RPC client disconnected before host URI request completed"); process.exit(0); } diff --git a/packages/coding-agent/src/modes/rpc/rpc-types.ts b/packages/coding-agent/src/modes/rpc/rpc-types.ts index cb22b1cc8..c074796fd 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-types.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-types.ts @@ -29,6 +29,7 @@ export type RpcCommand = | { id?: string; type: "get_state" } | { id?: string; type: "set_todos"; phases: TodoPhase[] } | { id?: string; type: "set_host_tools"; tools: RpcHostToolDefinition[] } + | { id?: string; type: "set_host_uri_schemes"; schemes: RpcHostUriSchemeDefinition[] } // Model | { id?: string; type: "set_model"; provider: string; modelId: string } @@ -121,6 +122,7 @@ export type RpcResponse = | { id?: string; type: "response"; command: "get_state"; success: true; data: RpcSessionState } | { id?: string; type: "response"; command: "set_todos"; success: true; data: { todoPhases: TodoPhase[] } } | { id?: string; type: "response"; command: "set_host_tools"; success: true; data: { toolNames: string[] } } + | { id?: string; type: "response"; command: "set_host_uri_schemes"; success: true; data: { schemes: string[] } } // Model | { @@ -304,6 +306,61 @@ export interface RpcHostToolResult { isError?: boolean; } +// ============================================================================ +// Host URI Frames (bidirectional) +// ============================================================================ + +export interface RpcHostUriSchemeDefinition { + /** URL scheme without trailing `://` (e.g. `db`, `notion`). */ + scheme: string; + /** Optional human-readable description for logs/diagnostics. */ + description?: string; + /** When true, the write tool is allowed to dispatch writes to this scheme. */ + writable?: boolean; + /** When true, downstream callers suppress hashline anchors for resolved content. */ + immutable?: boolean; +} + +export type RpcHostUriOperation = "read" | "write"; + +/** Emitted by the RPC server when it needs the host to satisfy a URI operation. */ +export interface RpcHostUriRequest { + type: "host_uri_request"; + id: string; + operation: RpcHostUriOperation; + url: string; + /** Present for write operations. */ + content?: string; +} + +/** Emitted by the RPC server when a pending URI request should be aborted. */ +export interface RpcHostUriCancelRequest { + type: "host_uri_cancel"; + id: string; + targetId: string; +} + +/** Sent by the host to complete a pending URI request. */ +export interface RpcHostUriResult { + type: "host_uri_result"; + id: string; + /** + * Required for successful `read` results. Ignored for `write` success. + * Set on errors when a textual explanation accompanies `isError`. + */ + content?: string; + /** Defaults to `text/plain` when omitted. */ + contentType?: "text/markdown" | "application/json" | "text/plain"; + /** Optional resolution notes propagated to the read tool. */ + notes?: string[]; + /** Overrides the scheme-level `immutable` flag for this single resolution. */ + immutable?: boolean; + /** When true, surface the result content as an error to the caller. */ + isError?: boolean; + /** Optional error message; preferred over `content` for error surfacing. */ + error?: string; +} + // ============================================================================ // Extension UI Commands (stdin) // ============================================================================ diff --git a/packages/coding-agent/src/tools/write.ts b/packages/coding-agent/src/tools/write.ts index df274b1e4..d2b280b2b 100644 --- a/packages/coding-agent/src/tools/write.ts +++ b/packages/coding-agent/src/tools/write.ts @@ -8,6 +8,8 @@ import { isEnoent, isRecord, prompt, untilAborted } from "@oh-my-pi/pi-utils"; import { type Static, Type } from "@sinclair/typebox"; import { stripHashlinePrefixes } from "../edit"; import type { RenderResultOptions } from "../extensibility/custom-tools/types"; +import { InternalUrlRouter } from "../internal-urls"; +import { parseInternalUrl } from "../internal-urls/parse"; import { createLspWritethrough, type FileDiagnosticsResult, type WritethroughCallback, writethroughNoop } from "../lsp"; import { getLanguageFromPath, highlightCode, type Theme } from "../modes/theme/theme"; import writeDescription from "../prompts/tools/write.md" with { type: "text" }; @@ -658,6 +660,24 @@ export class WriteTool implements AgentTool { // Strip hashline display prefixes (LINE+ID|) if the model copied them from read output const { text: cleanContent, stripped } = stripWriteContent(this.session, content); + const internalRouter = InternalUrlRouter.instance(); + if (internalRouter.canHandle(path)) { + const parsed = parseInternalUrl(path); + const scheme = parsed.protocol.replace(/:$/, "").toLowerCase(); + const handler = internalRouter.getHandler(scheme); + if (handler?.write) { + await handler.write(parsed, cleanContent, { cwd: this.session.cwd, signal }); + let resultText = `Successfully wrote ${cleanContent.length} bytes to ${path}`; + if (stripped) { + resultText += `\nNote: auto-stripped hashline display prefixes from content before writing.`; + } + return { content: [{ type: "text", text: resultText }], details: {} }; + } + // Schemes without a `write` hook fall through to existing logic + // (local:// resolves to a backing file via plan-mode-guard) or are + // rejected downstream when no backing file exists. + } + const conflictUri = parseConflictUri(path); if (conflictUri) { if (conflictUri.scope) { diff --git a/packages/coding-agent/test/tools/bash-interceptor.test.ts b/packages/coding-agent/test/tools/bash-interceptor.test.ts index 3ccb578a3..27a949e2d 100644 --- a/packages/coding-agent/test/tools/bash-interceptor.test.ts +++ b/packages/coding-agent/test/tools/bash-interceptor.test.ts @@ -83,13 +83,9 @@ describe("BashTool head/tail stripping", () => { // `seq 1 100 | head -3` would emit "1\n2\n3"; stripped, it emits 1..100. // We assert on the tail of the output rather than head, so a successful // strip is observable: line "100" only appears when head is gone. - const result = await tool.execute( - "tool-call", - { command: "seq 1 100 | head -3" }, - undefined, - undefined, - { toolNames: ["bash"] } as AgentToolContext, - ); + const result = await tool.execute("tool-call", { command: "seq 1 100 | head -3" }, undefined, undefined, { + toolNames: ["bash"], + } as AgentToolContext); const text = result.content.find(b => b.type === "text")?.text ?? ""; expect(text).toContain("100"); expect(text).toContain("Stripped trailing `| head -3`"); @@ -97,13 +93,9 @@ describe("BashTool head/tail stripping", () => { it("does not strip when the setting is disabled", async () => { const tool = createBashToolWithStrip(false); - const result = await tool.execute( - "tool-call", - { command: "seq 1 100 | head -3" }, - undefined, - undefined, - { toolNames: ["bash"] } as AgentToolContext, - ); + const result = await tool.execute("tool-call", { command: "seq 1 100 | head -3" }, undefined, undefined, { + toolNames: ["bash"], + } as AgentToolContext); const text = result.content.find(b => b.type === "text")?.text ?? ""; expect(text).toContain("1\n2\n3"); expect(text).not.toContain("100"); diff --git a/python/omp-rpc/README.md b/python/omp-rpc/README.md index 18d8dc32d..51914f771 100644 --- a/python/omp-rpc/README.md +++ b/python/omp-rpc/README.md @@ -142,6 +142,47 @@ If you want runtime conversion into a richer Python type, pass `decode=` to `host_tool(...)`. That lets you keep the JSON Schema contract on the wire while parsing the incoming argument object into a dataclass or model in the handler. +## Host-Owned URI Schemes + +Hosts can also expose custom URL schemes that behave like virtual files. +Registered schemes are routed through the agent's `read` (and `write`) tools +over the same RPC transport — handlers do the actual I/O on the Python side: + +```python +from omp_rpc import RpcClient, host_uri + +rows: dict[str, str] = {"42": "id=42\nname=Alice\n"} + + +def read_row(url: str, _ctx) -> str: + row_id = url.removeprefix("db://users/") + return rows[row_id] + + +def write_row(url: str, content: str, _ctx) -> None: + row_id = url.removeprefix("db://users/") + rows[row_id] = content + + +with RpcClient( + no_session=True, + host_uris=( + host_uri( + scheme="db", + description="Virtual db row files", + read=read_row, + write=write_row, + ), + ), +) as client: + client.prompt_and_wait("Read db://users/42 and rewrite it with name=Bob") +``` + +Schemes registered as read-only (no `write=`) reject `write` calls with a +clear error. The agent's `edit` tool does not target host URIs — hosts that +want mutation expose `write` and the model uses the `write` tool with the +full replacement content. + ## Extension UI Requests Extensions in RPC mode can ask the host for input. Those requests are available as diff --git a/python/omp-rpc/src/omp_rpc/__init__.py b/python/omp-rpc/src/omp_rpc/__init__.py index 2e9459991..014da5814 100644 --- a/python/omp-rpc/src/omp_rpc/__init__.py +++ b/python/omp-rpc/src/omp_rpc/__init__.py @@ -17,6 +17,16 @@ from .client import ( UiRequestListener, ) from .host_tools import HostTool, HostToolContext, HostToolResultPayload, HostToolResultValue, host_tool +from .host_uris import ( + HostUri, + HostUriContentType, + HostUriContext, + HostUriReadHandler, + HostUriReadResult, + HostUriReadValue, + HostUriWriteHandler, + host_uri, +) from .protocol import ( AgentEndEvent, AgentMessage, @@ -106,6 +116,13 @@ __all__ = [ "HostToolContext", "HostToolResultPayload", "HostToolResultValue", + "HostUri", + "HostUriContentType", + "HostUriContext", + "HostUriReadHandler", + "HostUriReadResult", + "HostUriReadValue", + "HostUriWriteHandler", "HookMessage", "ImageContent", "ListenerErrorEvent", @@ -162,4 +179,5 @@ __all__ = [ "parse_session_state", "parse_todo_phases", "host_tool", + "host_uri", ] diff --git a/python/omp-rpc/src/omp_rpc/client.py b/python/omp-rpc/src/omp_rpc/client.py index 17b36f8dc..eaac8921f 100644 --- a/python/omp-rpc/src/omp_rpc/client.py +++ b/python/omp-rpc/src/omp_rpc/client.py @@ -11,6 +11,7 @@ from pathlib import Path from typing import Any, Callable, Generic, Mapping, Sequence, TypeVar, cast from .host_tools import HostTool, HostToolContext +from .host_uris import HostUri, HostUriContext, normalize_read_result from .protocol import ( AgentStartEvent, AgentEndEvent, @@ -215,6 +216,11 @@ class _PendingHostToolCall: cancel_event: threading.Event +@dataclass(slots=True) +class _PendingHostUriRequest: + cancel_event: threading.Event + + @dataclass(slots=True) class _BoundedHistory(Generic[THistoryItem]): limit: int | None @@ -277,6 +283,7 @@ class RpcClient: provider_session_id: str | None = None, tools: Sequence[str] | None = None, custom_tools: Sequence[HostTool[Any, Any]] | None = None, + host_uris: Sequence[HostUri[Any]] | None = None, no_session: bool = False, no_skills: bool = False, no_rules: bool = False, @@ -300,6 +307,7 @@ class RpcClient: self._provider_session_id = provider_session_id self._tools = tuple(tools) if tools is not None else None self._custom_tools = tuple(custom_tools) if custom_tools is not None else () + self._host_uris = tuple(host_uris) if host_uris is not None else () self._no_session = no_session self._no_skills = no_skills self._no_rules = no_rules @@ -320,6 +328,7 @@ class RpcClient: self._event_condition = threading.Condition() self._pending: dict[str, _PendingRequest] = {} self._pending_host_tool_calls: dict[str, _PendingHostToolCall] = {} + self._pending_host_uri_requests: dict[str, _PendingHostUriRequest] = {} self._request_id = 0 self._events = _BoundedHistory[JsonObject](self._max_event_history) self._async_errors = _BoundedHistory[BaseException](_DEFAULT_ERROR_HISTORY_LIMIT) @@ -426,6 +435,8 @@ class RpcClient: if self._custom_tools: self.set_custom_tools(self._custom_tools) + if self._host_uris: + self.set_host_uris(self._host_uris) return self def stop(self) -> None: @@ -436,6 +447,8 @@ class RpcClient: self._stopping = True for pending_call in self._pending_host_tool_calls.values(): pending_call.cancel_event.set() + for pending_uri in self._pending_host_uri_requests.values(): + pending_uri.cancel_event.set() try: if process.stdin is not None: @@ -464,6 +477,7 @@ class RpcClient: pass self._fail_pending(RpcProcessExitError("RPC process stopped")) self._pending_host_tool_calls.clear() + self._pending_host_uri_requests.clear() self._process = None self._ready.set() with self._event_condition: @@ -760,6 +774,27 @@ class RpcClient: raise RpcError("set_host_tools response did not include toolNames") return tuple(str(name) for name in tool_names) + def set_host_uris(self, host_uris: Sequence[HostUri[Any]]) -> tuple[str, ...]: + self._host_uris = tuple(host_uris) + if self._process is None: + return tuple(uri.scheme for uri in self._host_uris) + + schemes_payload: list[JsonObject] = [] + for uri in self._host_uris: + entry: JsonObject = {"scheme": uri.scheme, "writable": uri.writable, "immutable": uri.immutable} + if uri.description is not None: + entry["description"] = uri.description + schemes_payload.append(entry) + + payload = self._request( + "set_host_uri_schemes", + schemes=cast(JsonValue, schemes_payload), + ) + schemes = payload.get("schemes") or [] + if not isinstance(schemes, list): + raise RpcError("set_host_uri_schemes response did not include schemes") + return tuple(str(entry) for entry in schemes) + def prompt( self, message: str, @@ -1054,6 +1089,88 @@ class RpcClient: if pending_call is not None: pending_call.cancel_event.set() + def _send_host_uri_error(self, request_id: str, message: str) -> None: + self._send_notification( + { + "type": "host_uri_result", + "id": request_id, + "error": message, + "isError": True, + } + ) + + def _handle_host_uri_request(self, payload: JsonObject) -> None: + request_id = payload.get("id") + operation = payload.get("operation") + url = payload.get("url") + if not isinstance(request_id, str) or not isinstance(operation, str) or not isinstance(url, str): + return + if operation not in ("read", "write"): + self._send_host_uri_error(request_id, f"Unsupported host URI operation: {operation}") + return + + try: + from urllib.parse import urlparse + + parsed = urlparse(url) + except ValueError: + self._send_host_uri_error(request_id, f"Could not parse host URI: {url}") + return + scheme = (parsed.scheme or "").lower() + uri = next((candidate for candidate in self._host_uris if candidate.scheme == scheme), None) + if uri is None: + self._send_host_uri_error(request_id, f'Host URI scheme "{scheme}://" is not registered') + return + + if operation == "write" and uri.write is None: + self._send_host_uri_error( + request_id, f'Host URI scheme "{scheme}://" was not registered with a write handler' + ) + return + + pending = _PendingHostUriRequest(cancel_event=threading.Event()) + self._pending_host_uri_requests[request_id] = pending + + def run() -> None: + try: + context = HostUriContext(url=url, operation=cast(Any, operation), _cancel_event=pending.cancel_event) + if operation == "read": + value = uri.read(url, context) + if pending.cancel_event.is_set(): + return + result_fields = normalize_read_result(value) + self._send_notification( + { + "type": "host_uri_result", + "id": request_id, + **result_fields, + } + ) + else: + raw_content = payload.get("content") + content = str(raw_content) if raw_content is not None else "" + assert uri.write is not None + uri.write(url, content, context) + if pending.cancel_event.is_set(): + return + self._send_notification({"type": "host_uri_result", "id": request_id}) + except Exception as exc: + if pending.cancel_event.is_set(): + return + self._send_host_uri_error(request_id, str(exc)) + finally: + self._pending_host_uri_requests.pop(request_id, None) + + threading.Thread(target=run, name=f"omp-rpc-host-uri:{scheme}:{operation}", daemon=True).start() + + def _handle_host_uri_cancel(self, payload: JsonObject) -> None: + target_id = payload.get("targetId") + if not isinstance(target_id, str): + return + pending = self._pending_host_uri_requests.get(target_id) + if pending is not None: + pending.cancel_event.set() + def _add_typed_event_listener(self, event_type: str, listener: TEventListener) -> Callable[[], None]: listeners = self._typed_event_listeners.setdefault(event_type, []) typed_listener = cast(AgentEventListener, listener) @@ -1232,6 +1349,12 @@ class RpcClient: if payload.get("type") == "host_tool_cancel": self._handle_host_tool_cancel(payload) continue + if payload.get("type") == "host_uri_request": + self._handle_host_uri_request(payload) + continue + if payload.get("type") == "host_uri_cancel": + self._handle_host_uri_cancel(payload) + continue notification = parse_notification(payload) listener_notification = parse_notification(payload) diff --git a/python/omp-rpc/src/omp_rpc/host_uris.py b/python/omp-rpc/src/omp_rpc/host_uris.py new file mode 100644 index 000000000..87b470be5 --- /dev/null +++ b/python/omp-rpc/src/omp_rpc/host_uris.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import threading +from dataclasses import dataclass, field +from typing import Callable, Generic, Literal, TypeAlias, TypeVar, TypedDict + +from .protocol import JsonObject + +TPayload = TypeVar("TPayload") + + +HostUriContentType: TypeAlias = Literal["text/markdown", "application/json", "text/plain"] + + +class HostUriReadResult(TypedDict, total=False): + """Structured response a `read` handler may return. + + Plain strings are also accepted; they are normalized to `{"content": }`. + """ + + content: str + content_type: HostUriContentType + notes: list[str] + immutable: bool + + +HostUriReadValue: TypeAlias = HostUriReadResult | str + + +@dataclass(slots=True) +class HostUriContext: + """Per-request context passed to host URI handlers. + + Mirrors the cancellation hooks `HostToolContext` exposes for parity, so + handlers can poll for cancellation when serving long-running reads/writes. + """ + + url: str + operation: Literal["read", "write"] + _cancel_event: threading.Event = field(default_factory=threading.Event) + + @property + def cancelled(self) -> bool: + return self._cancel_event.is_set() + + +HostUriReadHandler: TypeAlias = Callable[[str, HostUriContext], HostUriReadValue] +HostUriWriteHandler: TypeAlias = Callable[[str, str, HostUriContext], None] + + +@dataclass(slots=True, frozen=True) +class HostUri(Generic[TPayload]): + """Definition of a custom URI scheme served by the Python host. + + Hosts register a `HostUri` per scheme. The bridge dispatches `://` + URLs the agent reads (and, when `write` is provided, writes) to the + matching callbacks. The agent's `edit` tool is not supported for virtual + URIs — hosts that want to mutate virtual files expose a `write` handler + and let the model use the `write` tool with the full replacement content. + """ + + scheme: str + read: HostUriReadHandler + write: HostUriWriteHandler | None = None + description: str | None = None + immutable: bool = False + + @property + def writable(self) -> bool: + return self.write is not None + + +def host_uri( + *, + scheme: str, + read: HostUriReadHandler, + write: HostUriWriteHandler | None = None, + description: str | None = None, + immutable: bool = False, +) -> HostUri[None]: + cleaned = (scheme or "").strip().lower() + if not cleaned: + raise ValueError("scheme must be a non-empty string") + return HostUri( + scheme=cleaned, + read=read, + write=write, + description=description, + immutable=immutable, + ) + + +def normalize_read_result(value: HostUriReadValue) -> JsonObject: + """Convert a handler's `read` return into the wire-frame fields. + + Returns a dict suitable for spreading into a `host_uri_result` payload. + """ + + if isinstance(value, str): + return {"content": value} + if not isinstance(value, dict): + raise TypeError("Host URI read handlers must return a string or a HostUriReadResult mapping") + + payload: JsonObject = {} + if "content" not in value: + raise ValueError("HostUriReadResult requires a 'content' field") + payload["content"] = str(value["content"]) + + content_type = value.get("content_type") + if content_type is not None: + if content_type not in ("text/markdown", "application/json", "text/plain"): + raise ValueError(f"Unsupported content_type: {content_type!r}") + payload["contentType"] = content_type + + notes = value.get("notes") + if notes is not None: + payload["notes"] = [str(item) for item in notes] + + if "immutable" in value: + payload["immutable"] = bool(value["immutable"]) + + return payload diff --git a/python/omp-rpc/tests/test_host_uris.py b/python/omp-rpc/tests/test_host_uris.py new file mode 100644 index 000000000..632d47c13 --- /dev/null +++ b/python/omp-rpc/tests/test_host_uris.py @@ -0,0 +1,235 @@ +from __future__ import annotations + +import sys +import textwrap +import threading +import time +import unittest + +from omp_rpc import RpcClient, host_uri +from omp_rpc.host_uris import HostUri, normalize_read_result + + +URI_SERVER = textwrap.dedent( + """ + import json + import sys + + print(json.dumps({"type": "ready"}), flush=True) + + pending_uri_id = 1 + + def respond(request_id, command, data=None, success=True, error=None): + frame = {"id": request_id, "type": "response", "command": command, "success": success} + if success: + if data is not None: + frame["data"] = data + else: + frame["error"] = error or "error" + print(json.dumps(frame), flush=True) + + for raw_line in sys.stdin: + raw_line = raw_line.strip() + if not raw_line: + continue + command = json.loads(raw_line) + command_type = command.get("type") + request_id = command.get("id") + + if command_type == "set_host_uri_schemes": + schemes = command.get("schemes", []) + respond( + request_id, + "set_host_uri_schemes", + {"schemes": [entry.get("scheme", "") for entry in schemes]}, + ) + elif command_type == "trigger_read": + print( + json.dumps( + { + "type": "host_uri_request", + "id": f"uri-req-{pending_uri_id}", + "operation": "read", + "url": command["url"], + } + ), + flush=True, + ) + pending_uri_id += 1 + respond(request_id, "trigger_read", {}) + elif command_type == "trigger_write": + print( + json.dumps( + { + "type": "host_uri_request", + "id": f"uri-req-{pending_uri_id}", + "operation": "write", + "url": command["url"], + "content": command["content"], + } + ), + flush=True, + ) + pending_uri_id += 1 + respond(request_id, "trigger_write", {}) + elif command_type == "host_uri_result": + # Echo back as response so the test can assert on the wire frame + print( + json.dumps( + { + "type": "response", + "command": "uri_echo", + "success": True, + "data": {"frame": command}, + } + ), + flush=True, + ) + else: + respond(request_id, command_type, success=False, error=f"unsupported: {command_type}") + """ +) + + +class HostUriHelperTests(unittest.TestCase): + def test_normalize_read_result_accepts_string(self) -> None: + self.assertEqual(normalize_read_result("hello"), {"content": "hello"}) + + def test_normalize_read_result_accepts_full_mapping(self) -> None: + result = normalize_read_result( + { + "content": "body", + "content_type": "application/json", + "notes": ["fresh"], + "immutable": True, + } + ) + self.assertEqual(result["content"], "body") + self.assertEqual(result["contentType"], "application/json") + self.assertEqual(result["notes"], ["fresh"]) + self.assertTrue(result["immutable"]) + + def test_normalize_read_result_requires_content(self) -> None: + with self.assertRaises(ValueError): + normalize_read_result({"content_type": "text/plain"}) # type: ignore[arg-type] + + def test_normalize_read_result_rejects_invalid_content_type(self) -> None: + with self.assertRaises(ValueError): + normalize_read_result({"content": "x", "content_type": "application/octet-stream"}) # type: ignore[arg-type] + + def test_host_uri_helper_normalizes_scheme(self) -> None: + uri = host_uri(scheme=" DB ", read=lambda url, ctx: "x") + self.assertEqual(uri.scheme, "db") + self.assertFalse(uri.writable) + + with self.assertRaises(ValueError): + host_uri(scheme="", read=lambda url, ctx: "x") + + def test_host_uri_writable_when_write_supplied(self) -> None: + uri = host_uri(scheme="db", read=lambda url, ctx: "x", write=lambda url, content, ctx: None) + self.assertTrue(uri.writable) + + +class RpcHostUriBridgeTests(unittest.TestCase): + def _make_client(self, **kwargs: object) -> RpcClient: + return RpcClient( + command=[sys.executable, "-u", "-c", URI_SERVER], + startup_timeout=2.0, + request_timeout=2.0, + **kwargs, + ) + + def test_set_host_uris_registers_schemes_on_start(self) -> None: + captured: list[tuple[str, str]] = [] + + def read_db(url: str, _ctx) -> str: + captured.append(("read", url)) + return "id=42" + + with self._make_client( + host_uris=(host_uri(scheme="db", read=read_db, description="test rows"),), + ) as client: + # No public list — we exercise the on-start side effect by hitting the wire. + payload = client._request("trigger_read", url="db://users/42") # type: ignore[attr-defined] + self.assertEqual(payload, {}) + + frame = self._await_echo(client) + self.assertEqual(frame["type"], "host_uri_result") + self.assertEqual(frame["content"], "id=42") + self.assertEqual(captured, [("read", "db://users/42")]) + + def test_read_handler_can_return_structured_result(self) -> None: + def read_db(_url: str, _ctx): + return { + "content": '{"name":"Alice"}', + "content_type": "application/json", + "notes": ["row fresh"], + "immutable": True, + } + + with self._make_client(host_uris=(host_uri(scheme="db", read=read_db),)) as client: + client._request("trigger_read", url="db://users/42") # type: ignore[attr-defined] + frame = self._await_echo(client) + self.assertEqual(frame["content"], '{"name":"Alice"}') + self.assertEqual(frame["contentType"], "application/json") + self.assertEqual(frame["notes"], ["row fresh"]) + self.assertTrue(frame["immutable"]) + + def test_write_handler_receives_content_and_succeeds(self) -> None: + seen: dict[str, str] = {} + + def write_db(url: str, content: str, _ctx) -> None: + seen[url] = content + + uri = host_uri(scheme="db", read=lambda url, ctx: "ignored", write=write_db) + with self._make_client(host_uris=(uri,)) as client: + client._request("trigger_write", url="db://users/42", content="name=Bob") # type: ignore[attr-defined] + frame = self._await_echo(client) + self.assertEqual(frame["type"], "host_uri_result") + self.assertNotIn("isError", frame) + self.assertEqual(seen, {"db://users/42": "name=Bob"}) + + def test_write_rejected_for_read_only_scheme(self) -> None: + with self._make_client( + host_uris=(host_uri(scheme="db", read=lambda url, ctx: "x"),), + ) as client: + client._request("trigger_write", url="db://users/42", content="ignored") # type: ignore[attr-defined] + frame = self._await_echo(client) + self.assertTrue(frame.get("isError")) + self.assertIn("write handler", frame["error"]) + + def test_unknown_scheme_is_rejected_with_error(self) -> None: + with self._make_client( + host_uris=(host_uri(scheme="db", read=lambda url, ctx: "x"),), + ) as client: + client._request("trigger_read", url="other://stuff") # type: ignore[attr-defined] + frame = self._await_echo(client) + self.assertTrue(frame.get("isError")) + self.assertIn("not registered", frame["error"]) + + def test_handler_exception_is_surfaced_as_error(self) -> None: + def read_db(_url: str, _ctx) -> str: + raise RuntimeError("boom") + + with self._make_client(host_uris=(host_uri(scheme="db", read=read_db),)) as client: + client._request("trigger_read", url="db://users/42") # type: ignore[attr-defined] + frame = self._await_echo(client) + self.assertTrue(frame.get("isError")) + self.assertEqual(frame["error"], "boom") + + def _await_echo(self, client: RpcClient) -> dict: + # The fake server echoes the host_uri_result frame back as an + # `uri_echo` response. We poll the events history to surface it. + deadline = time.time() + 2.0 + while time.time() < deadline: + with client._state_lock: # type: ignore[attr-defined] + events = client._events.snapshot() # type: ignore[attr-defined] + for event in events: + if event.get("command") == "uri_echo" and event.get("data"): + return event["data"]["frame"] + time.sleep(0.02) + self.fail("Timed out waiting for host_uri_result echo") + + +if __name__ == "__main__": + unittest.main()