feat(agent): add remote session RPC

This commit is contained in:
Armin Ronacher 2026-08-18 10:54:51 +02:00
parent 4266bddc08
commit fb33e0db24
24 changed files with 1673 additions and 139 deletions

View file

@ -1,6 +1,6 @@
# Harness RPC Design Notes
> **Status:** Design input, not a normative contract or implementation handoff. Commit `3fbfcf48f` added context-last parameters and mechanically threads them through the current harness, session implementations, hooks, events, and backends. RPC integration, telemetry carriers, and cancellation semantics remain design work. Resolve the open decisions and fold the accepted behavior into `harness.md` before creating a work package.
> **Status:** Design input, not a normative contract or implementation handoff. Required trailing `Context` parameters and the untyped RPC transport have landed. `RemoteSession`/`RemoteSessionTree`, connection-owned remote mutation capabilities, and client-to-server request cancellation are implemented. The hosted Harness and worker protocols remain hand-written projections, worker-hop cancellation and trace carriers are still missing, and Pi request ingress still starts from `TODO_CONTEXT`. Resolve the remaining open decisions and fold the accepted behavior into `harness.md` before creating another work package.
## Goal
@ -168,7 +168,15 @@ receive request
A cancel frame aborts only the matching request controller. Disconnect aborts every active request and closes every subscription owned by that connection.
`packages/protocol/src/rpc.ts` already separates a server implementation context from serialized arguments through `RpcImplementation<TManifest, TContext>`. It can inform this design, but the current client API has no invocation-context/cancellation/trace integration and the manifest model does not yet cover remote object references or subscriptions.
`packages/protocol/src/rpc.ts` already separates a server implementation context from serialized arguments through `RpcImplementation<TManifest, TContext>`. Its implementation context is dispatcher-owned and context-first; an adapter then calls the real receiver with its required trailing `Context`. The manifest client API still has no invocation-context, cancellation, or trace integration, and the manifest model does not yet cover remote object references or subscriptions.
## Current hosted-harness adapter
The current Pi service protocol is an explicit adapter, not the generic proxy described below. `PiServerHost`, `HostedHarnessHandle`, `HostedHarnessAttachment`, and `HostedHarnessWatch` receive trailing contexts, and `HostedHarnessManager` forwards one context through session discovery, creation, attachment, prompting, watch lifecycle, and close.
At Pi request ingress, `PiServer` still starts from `TODO_CONTEXT`, but derives a request context with a per-request abort signal. `PiClient.invoke()` sends an untyped method string plus unknown arguments and maps its signal to a cancel frame. The server routes `session.*` methods through `RemoteSessionManager`, which owns exclusive Session handles and mutation capabilities. The experimental coding-agent client opens a `RemoteSession`; prompting raises until `AgentHarness` runtime2 prompting is implemented.
The coding-agent worker manager accepts contexts on its hosted interfaces, but worker operation frames carry neither cancellation nor trace metadata and the worker reconstructs calls with `TODO_CONTEXT`. Hosted Harness operations therefore still have interface propagation only across that hop.
## Remote object identity
@ -246,7 +254,7 @@ Members requiring adaptation:
- `view(lane)`: synchronous remote-object return; use a deterministic client proxy or adapter projection;
- `mutate(callback)`: not ordinary RPC because it contains executable callback code and holds the session mutation line.
If remote mutation is required, design a separate transaction protocol with explicit open/read/commit/close, timeout, cancellation, and disconnect behavior. Do not transparently serialize or reverse-call `Session.mutate()`.
`RemoteSession.mutate()` executes the callback locally against a connection-owned `RemoteSessionMutator`. Its untyped protocol explicitly begins a mutation, performs reads and at most one commit by mutation ID, then finishes it. The server implements begin by entering the real `Session.mutate()` callback and retains the lane until finish, preserving client-local commit-then-publish ordering. Disconnect releases every owned mutation; a server-owned lease expires a wedged mutation after admitted commits drain. Callbacks are never serialized or reverse-called, and uncertain commits are never retried.
### `AgentLane` and `AgentHarness`
@ -319,7 +327,7 @@ A future client facade may reconstruct the familiar local `Events`/`WatchHandle`
## Request cancellation
The adapter maps each invocation's `abortSignal`, when present, to one request ID:
Request cancellation is implemented between `PiClient` and `PiServer`. Worker transport cancellation remains missing. The adapter maps each invocation's `abortSignal`, when present, to one request ID:
```text
client context.abortSignal aborts
@ -387,37 +395,43 @@ Compilation proves signature coverage only. It does not prove correct trace pare
## Required tests for the later handoff
Current tests cover typed service-adapter validation, untyped transport calls, pre-abort, request cancellation, ordinary Pi service calls, attachment/watch lifecycle, snapshot-first event buffering, remote Session/tree behavior, callback-scoped mutation, lease expiry, disconnect cleanup, and use of `RemoteSession` by `AgentHarness.create()`. Remaining generic RPC and context coverage includes:
- static rejection of non-JSON ordinary interface members;
- runtime rejection of functions, cycles, non-finite numbers, sparse arrays, and class instances;
- context is removed from serialized arguments and reconstructed server-side;
- landed context-last calls, plus the selected JSON-safe optional-argument strategy;
- trailing context is removed from serialized arguments and reconstructed server-side;
- omitted optional arguments are normalized without serializing `undefined`;
- ordinary transparent method calls;
- nested remote references and stable proxy identity;
- adapter override invokes/decorates the real method;
- property and synchronous remote-object projections;
- unknown target/method and malformed envelope rejection;
- pre-aborted request starts no server work;
- one request cancellation does not abort another;
- disconnect aborts requests and closes subscriptions;
- client/server telemetry parentage;
- malformed/missing trace carrier degrades safely;
- event subscription establishment race;
- event source telemetry reconstruction;
- unsubscribe, disconnect, ordering, buffering, and backpressure;
- snapshot plus buffered events has no observation gap;
- drive joiner cancellation does not cancel shared execution;
- invocation cancellation writes no durable cancellation state.
## Resolved migration decisions
- Context-aware receiver methods use one required trailing `Context`.
- `Context`, `AbortSignal`, and `TelemetryContext` objects are never serialized.
- RPC transport calls use an untyped method string and unknown arguments/result; domain facades own typing.
- Remote mutation callbacks execute locally over an explicit connection-owned mutation capability.
- Remote mutation leases are server-owned policy, distinct from caller cancellation.
- The current Pi service remains an explicit adapter projection rather than a generic `AgentHarness` proxy.
## Open decisions before `harness.md`
- whether to retain context last or migrate context first;
- if context remains last, the JSON-safe trailing-`undefined`/arity protocol;
- whether all ordinary results must be JSON-compatible or some interfaces need RPC projections;
- static contract-checking API and test location;
- object-reference ownership, lifetime, and nested result encoding;
- exact adapter descriptor/override API;
- event flow-control and reconnection policy;
- whether remote transaction or hook protocols are required at all;
- whether a remote hook protocol is required at all;
- telemetry carrier adapter ownership;
- drive ownership policy under multiple callers;
- exact package boundaries for generic RPC, harness adapters, and protocol schemas.

View file

@ -7,7 +7,9 @@ export interface ContextKey<T> {
readonly valueType?: (value: T) => T;
}
const ABORT_SIGNAL_CONTEXT_KEY: ContextKey<AbortSignal> = Object.freeze({ token: Symbol("pi.abortSignal") });
const ABORT_SIGNAL_CONTEXT_KEY: ContextKey<AbortSignal | undefined> = Object.freeze({
token: Symbol("pi.abortSignal"),
});
const TELEMETRY_CONTEXT_KEY: ContextKey<TelemetryContext> = Object.freeze({ token: Symbol("pi.telemetryContext") });
/** Immutable invocation-scoped values passed explicitly through harness operations. */
@ -96,6 +98,11 @@ export function withAbortSignal(signal: AbortSignal, context: Context): Context
return withContextValue(ABORT_SIGNAL_CONTEXT_KEY, combined, context);
}
/** Derive a context retaining all values except caller cancellation. Intended for mandatory cleanup only. */
export function withoutAbortSignal(context: Context): Context {
return withContextValue(ABORT_SIGNAL_CONTEXT_KEY, undefined, context);
}
/** Derive an independently cancellable child context. */
export function withCancel(context: Context): {
readonly context: Context;

View file

@ -21,6 +21,14 @@ export {
} from "./jsonl/index.ts";
export type { MemorySessionRepoOptions } from "./memory.ts";
export { MemorySessionRepo } from "./memory.ts";
export { StorageBackedSession } from "./session.ts";
export { RemoteSession, type RemoteSessionRpc, RemoteSessionTree } from "./remote.ts";
export {
SessionInvalidLaneError,
SessionInvariantError,
SessionLaneExistsError,
SessionPendingAssistantMessageError,
SessionUnknownTargetError,
StorageBackedSession,
} from "./session.ts";
export * from "./types.ts";
export * from "./values.ts";

View file

@ -0,0 +1,488 @@
import { uuidv7 } from "@earendil-works/pi-ai";
import type { Context } from "../context.ts";
import { withoutAbortSignal } from "../context.ts";
import {
SessionInvalidLaneError,
SessionInvariantError,
SessionLaneExistsError,
SessionPendingAssistantMessageError,
SessionUnknownTargetError,
} from "./session.ts";
import type {
BranchScan,
CommitResult,
Entry,
EntryQuery,
IdGenerator,
JsonValue,
LaneConfiguration,
Session,
SessionMetadata,
SessionMutator,
SessionStats,
SessionTree,
Write,
} from "./types.ts";
import type { ListElement, ListReadOptions, StoredValue, Value, ValueList } from "./values.ts";
/** Untyped invocation boundary used by the remote Session facade. */
export interface RemoteSessionRpc {
invoke(method: string, args: unknown, context: Context): Promise<unknown>;
}
interface RemoteSessionOpenResult {
handle: string;
metadata: SessionMetadata;
}
/** SessionTree facade whose methods execute against one remote Session lane. */
export class RemoteSessionTree implements SessionTree {
protected readonly rpc: RemoteSessionRpc;
protected readonly handle: string;
protected readonly lane: string;
constructor(rpc: RemoteSessionRpc, handle: string, lane: string) {
this.rpc = rpc;
this.handle = handle;
this.lane = lane;
}
async getLeafId(context: Context): Promise<string | null> {
return expectNullableString(await this.invoke("session.tree.getLeafId", {}, context), "leaf ID");
}
async getEntry(id: string, context: Context): Promise<Entry | undefined> {
const result = await this.invoke("session.tree.getEntry", { id }, context);
return result === null ? undefined : (expectObject(result, "entry") as unknown as Entry);
}
async getStats(context: Context): Promise<SessionStats> {
return expectObject(
await this.invoke("session.tree.getStats", {}, context),
"Session stats",
) as unknown as SessionStats;
}
async getValue<T>(address: Value<T>, context: Context): Promise<StoredValue<T> | undefined> {
const result = await this.invoke("session.tree.getValue", { address }, context);
return result === null ? undefined : (expectObject(result, "stored value") as unknown as StoredValue<T>);
}
async scanValues<T>(prefix: Value<T>, context: Context): Promise<StoredValue<T>[]> {
return expectArray(
await this.invoke("session.tree.scanValues", { prefix }, context),
"stored values",
) as StoredValue<T>[];
}
async readList<T>(
address: ValueList<T>,
options: ListReadOptions | undefined,
context: Context,
): Promise<ListElement<T>[]> {
return expectArray(
await this.invoke("session.tree.readList", { address, options: options ?? null }, context),
"list elements",
) as ListElement<T>[];
}
async setValue<T>(address: Value<T>, next: NoInfer<T>, context: Context): Promise<void> {
expectNull(await this.invoke("session.tree.setValue", { address, next }, context));
}
async deleteValue<T>(address: Value<T>, context: Context): Promise<void> {
expectNull(await this.invoke("session.tree.deleteValue", { address }, context));
}
async appendList<T>(address: ValueList<T>, element: NoInfer<T>, context: Context): Promise<void> {
expectNull(await this.invoke("session.tree.appendList", { address, element }, context));
}
async deleteList<T>(address: ValueList<T>, context: Context): Promise<void> {
expectNull(await this.invoke("session.tree.deleteList", { address }, context));
}
async getName(context: Context): Promise<string | undefined> {
const value = expectNullableString(await this.invoke("session.tree.getName", {}, context), "Session name");
return value === null ? undefined : value;
}
async setName(name: string | undefined, context: Context): Promise<void> {
expectNull(await this.invoke("session.tree.setName", { name: name ?? null }, context));
}
async getLabel(targetId: string, context: Context): Promise<string | undefined> {
const value = expectNullableString(
await this.invoke("session.tree.getLabel", { targetId }, context),
"entry label",
);
return value === null ? undefined : value;
}
async setLabel(targetId: string, label: string | undefined, context: Context): Promise<void> {
expectNull(await this.invoke("session.tree.setLabel", { targetId, label: label ?? null }, context));
}
async findEntries(query: EntryQuery | undefined, context: Context): Promise<Entry[]> {
return expectArray(
await this.invoke("session.tree.findEntries", { query: query ?? null }, context),
"entries",
) as Entry[];
}
async findEntry(query: EntryQuery | undefined, context: Context): Promise<Entry | undefined> {
const result = await this.invoke("session.tree.findEntry", { query: query ?? null }, context);
return result === null ? undefined : (expectObject(result, "entry") as unknown as Entry);
}
async findEntriesOnBranch(query: BranchScan | undefined, context: Context): Promise<Entry[]> {
return expectArray(
await this.invoke("session.tree.findEntriesOnBranch", { query: query ?? null }, context),
"branch entries",
) as Entry[];
}
async findEntryOnBranch(query: BranchScan | undefined, context: Context): Promise<Entry | undefined> {
const result = await this.invoke("session.tree.findEntryOnBranch", { query: query ?? null }, context);
return result === null ? undefined : (expectObject(result, "entry") as unknown as Entry);
}
async appendMessage(message: Parameters<SessionTree["appendMessage"]>[0], context: Context): Promise<string> {
return expectString(await this.invoke("session.tree.appendMessage", { message }, context), "entry ID");
}
async appendCustomEntry(customType: string, data: JsonValue | undefined, context: Context): Promise<string> {
return expectString(
await this.invoke(
"session.tree.appendCustomEntry",
{ customType, ...(data === undefined ? {} : { data }) },
context,
),
"entry ID",
);
}
protected async invoke(method: string, args: Record<string, unknown>, context: Context): Promise<unknown> {
try {
return await this.rpc.invoke(method, { handle: this.handle, lane: this.lane, ...args }, context);
} catch (error) {
const code = remoteErrorCode(error);
if (code === "session_pending_message") throw new SessionPendingAssistantMessageError();
if (code === "session_invariant") throw new SessionInvariantError(remoteErrorMessage(error));
throw error;
}
}
}
/** Session facade backed entirely by untyped RPC calls. */
export class RemoteSession<TMetadata extends SessionMetadata = SessionMetadata>
extends RemoteSessionTree
implements Session<TMetadata>
{
readonly metadata: TMetadata;
readonly idGenerator: IdGenerator = { next: uuidv7 };
private closed = false;
constructor(rpc: RemoteSessionRpc, handle: string, metadata: TMetadata) {
super(rpc, handle, "main");
this.metadata = metadata;
}
static async open(
rpc: RemoteSessionRpc,
sessionId: string,
context: Context,
): Promise<RemoteSession<SessionMetadata>> {
const result = expectObject(await rpc.invoke("session.open", { sessionId }, context), "remote Session");
const opened: RemoteSessionOpenResult = {
handle: expectString(result.handle, "remote Session handle"),
metadata: decodeMetadata(result.metadata),
};
return new RemoteSession(rpc, opened.handle, opened.metadata);
}
view(lane: string): SessionTree {
this.assertOpen();
return new RemoteSessionTree(this.rpc, this.handle, lane);
}
async getEntries(ids: string[], context: Context): Promise<Map<string, Entry>> {
const pairs = expectArray(await this.invoke("session.getEntries", { ids }, context), "entry pairs");
const entries = new Map<string, Entry>();
for (const pair of pairs) {
if (!Array.isArray(pair) || pair.length !== 2)
throw new TypeError("Remote Session returned an invalid entry pair");
entries.set(expectString(pair[0], "entry ID"), expectObject(pair[1], "entry") as unknown as Entry);
}
return entries;
}
async getValue<T>(address: Value<T>, context: Context): Promise<StoredValue<T> | undefined> {
const result = await this.invoke("session.getValue", { address }, context);
return result === null ? undefined : (expectObject(result, "stored value") as unknown as StoredValue<T>);
}
async scanValues<T>(prefix: Value<T>, context: Context): Promise<StoredValue<T>[]> {
return expectArray(
await this.invoke("session.scanValues", { prefix }, context),
"stored values",
) as StoredValue<T>[];
}
async readList<T>(
address: ValueList<T>,
options: ListReadOptions | undefined,
context: Context,
): Promise<ListElement<T>[]> {
return expectArray(
await this.invoke("session.readList", { address, options: options ?? null }, context),
"list elements",
) as ListElement<T>[];
}
async mutate<T>(
lane: string,
mutation: (mutator: SessionMutator, context: Context) => T | Promise<T>,
context: Context,
): Promise<T> {
this.assertOpen();
const opened = expectObject(
await this.invoke("session.mutation.begin", { mutationLane: lane }, context),
"remote mutation",
);
const mutationId = expectString(opened.mutationId, "remote mutation ID");
const mutator = new RemoteSessionMutator(this.rpc, this.handle, mutationId, lane);
let outcome: { ok: true; value: T } | { ok: false; error: unknown };
try {
outcome = { ok: true, value: await mutation(mutator, context) };
} catch (error) {
outcome = { ok: false, error };
}
await mutator.settle();
mutator.invalidate();
let finishFailure: { error: unknown } | undefined;
try {
expectNull(
await this.rpc.invoke(
"session.mutation.finish",
{ handle: this.handle, mutationId },
withoutAbortSignal(context),
),
);
} catch (error) {
finishFailure = { error };
}
if (!outcome.ok) {
if (finishFailure !== undefined) {
throw new AggregateError(
[outcome.error, finishFailure.error],
"Remote Session mutation and cleanup failed",
);
}
throw outcome.error;
}
if (finishFailure !== undefined) throw finishFailure.error;
return outcome.value;
}
async createLane(
name: string,
at: string | null,
configuration: LaneConfiguration,
context: Context,
): Promise<SessionTree> {
this.assertOpen();
try {
expectNull(await this.invoke("session.createLane", { name, at, configuration }, context));
} catch (error) {
const code = remoteErrorCode(error);
if (code === "session_invalid_lane") throw new SessionInvalidLaneError(name, remoteErrorMessage(error));
if (code === "session_lane_exists") throw new SessionLaneExistsError(name);
if (code === "session_unknown_target" && at !== null) throw new SessionUnknownTargetError(at);
throw error;
}
return this.view(name);
}
async close(context: Context): Promise<void> {
if (this.closed) return;
this.closed = true;
expectNull(await this.rpc.invoke("session.close", { handle: this.handle }, context));
}
private assertOpen(): void {
if (this.closed) throw new Error("Session is closed");
}
}
class RemoteSessionMutator implements SessionMutator {
readonly lane: string;
private readonly rpc: RemoteSessionRpc;
private readonly handle: string;
private readonly mutationId: string;
private active = true;
private commitPromise: Promise<CommitResult> | undefined;
constructor(rpc: RemoteSessionRpc, handle: string, mutationId: string, lane: string) {
this.rpc = rpc;
this.handle = handle;
this.mutationId = mutationId;
this.lane = lane;
}
commit(writes: Write[], context: Context): Promise<CommitResult> {
this.assertActive();
if (this.commitPromise !== undefined) return Promise.reject(new Error("SessionMutator commit already attempted"));
this.commitPromise = this.rpc
.invoke("session.mutation.commit", { handle: this.handle, mutationId: this.mutationId, writes }, context)
.then(decodeCommitResult);
return this.commitPromise;
}
async getEntries(ids: string[], context: Context): Promise<Map<string, Entry>> {
this.assertActive();
const pairs = expectArray(
await this.rpc.invoke(
"session.mutation.getEntries",
{ handle: this.handle, mutationId: this.mutationId, ids },
context,
),
"entry pairs",
);
const entries = new Map<string, Entry>();
for (const pair of pairs) {
if (!Array.isArray(pair) || pair.length !== 2)
throw new TypeError("Remote Session returned an invalid entry pair");
entries.set(expectString(pair[0], "entry ID"), expectObject(pair[1], "entry") as unknown as Entry);
}
return entries;
}
async getValue<T>(address: Value<T>, context: Context): Promise<StoredValue<T> | undefined> {
this.assertActive();
const result = await this.rpc.invoke(
"session.mutation.getValue",
{ handle: this.handle, mutationId: this.mutationId, address },
context,
);
return result === null ? undefined : (expectObject(result, "stored value") as unknown as StoredValue<T>);
}
async scanValues<T>(prefix: Value<T>, context: Context): Promise<StoredValue<T>[]> {
this.assertActive();
return expectArray(
await this.rpc.invoke(
"session.mutation.scanValues",
{ handle: this.handle, mutationId: this.mutationId, prefix },
context,
),
"stored values",
) as StoredValue<T>[];
}
async readList<T>(
address: ValueList<T>,
options: ListReadOptions | undefined,
context: Context,
): Promise<ListElement<T>[]> {
this.assertActive();
return expectArray(
await this.rpc.invoke(
"session.mutation.readList",
{ handle: this.handle, mutationId: this.mutationId, address, options: options ?? null },
context,
),
"list elements",
) as ListElement<T>[];
}
settle(): Promise<void> {
return (
this.commitPromise?.then(
() => undefined,
() => undefined,
) ?? Promise.resolve()
);
}
invalidate(): void {
this.active = false;
}
private assertActive(): void {
if (!this.active) throw new Error("SessionMutator cannot be used outside its mutation callback");
}
}
function decodeMetadata(value: unknown): SessionMetadata {
const metadata = expectObject(value, "Session metadata");
return {
id: expectString(metadata.id, "Session ID"),
createdAt: expectNumber(metadata.createdAt, "Session creation time"),
storageVersion: expectNumber(metadata.storageVersion, "Session storage version"),
...(metadata.cwd === undefined ? {} : { cwd: expectString(metadata.cwd, "Session cwd") }),
...(metadata.parentSessionId === undefined
? {}
: { parentSessionId: expectString(metadata.parentSessionId, "parent Session ID") }),
...(metadata.legacyParentSessionPath === undefined
? {}
: { legacyParentSessionPath: expectString(metadata.legacyParentSessionPath, "legacy parent Session path") }),
};
}
function decodeCommitResult(value: unknown): CommitResult {
const result = expectObject(value, "commit result");
const seqs = expectArray(result.seqs, "commit sequences").map((seq) => expectNumber(seq, "commit sequence"));
return {
firstSeq: expectNumber(result.firstSeq, "first commit sequence"),
seqs,
timestamp: expectNumber(result.timestamp, "commit timestamp"),
};
}
function expectObject(value: unknown, description: string): Record<string, unknown> {
if (typeof value !== "object" || value === null || Array.isArray(value)) {
throw new TypeError(`Remote Session returned an invalid ${description}`);
}
const prototype = Object.getPrototypeOf(value);
if (prototype !== Object.prototype && prototype !== null) {
throw new TypeError(`Remote Session returned an invalid ${description}`);
}
return value as Record<string, unknown>;
}
function expectArray(value: unknown, description: string): unknown[] {
if (!Array.isArray(value)) throw new TypeError(`Remote Session returned invalid ${description}`);
return value;
}
function expectString(value: unknown, description: string): string {
if (typeof value !== "string") throw new TypeError(`Remote Session returned an invalid ${description}`);
return value;
}
function expectNullableString(value: unknown, description: string): string | null {
if (value === null) return null;
return expectString(value, description);
}
function expectNumber(value: unknown, description: string): number {
if (typeof value !== "number" || !Number.isFinite(value)) {
throw new TypeError(`Remote Session returned an invalid ${description}`);
}
return value;
}
function expectNull(value: unknown): void {
if (value !== null) throw new TypeError("Remote Session operation returned an invalid result");
}
function remoteErrorCode(error: unknown): string | undefined {
return typeof error === "object" && error !== null && "code" in error && typeof error.code === "string"
? error.code
: undefined;
}
function remoteErrorMessage(error: unknown): string {
return error instanceof Error ? error.message : String(error);
}

View file

@ -6,6 +6,7 @@ import {
withAbortSignal,
withCancel,
withContextValue,
withoutAbortSignal,
withTelemetryContext,
} from "../../src/harness/context.ts";
import { InMemoryTelemetryContext } from "../../src/index.ts";
@ -59,6 +60,18 @@ describe("Context", () => {
expect(sibling.context.abortSignal?.reason).toBe("parent");
});
it("masks caller cancellation for mandatory cleanup", () => {
const controller = new AbortController();
const key = createContextKey<string>("value");
const context = withContextValue(key, "preserved", withAbortSignal(controller.signal, BACKGROUND_CONTEXT));
const cleanup = withoutAbortSignal(context);
controller.abort();
expect(context.abortSignal?.aborted).toBe(true);
expect(cleanup.abortSignal).toBeUndefined();
expect(cleanup.value(key)).toBe("preserved");
});
it("carries telemetry as an ordinary context value", async () => {
const telemetry = new InMemoryTelemetryContext();
const context = withTelemetryContext(telemetry, BACKGROUND_CONTEXT);

View file

@ -7,6 +7,7 @@ import {
type PromptArguments,
type PromptImage,
type PromptMessage,
type ProtocolRpcCall,
ProtocolValidationError,
type ResponseEnvelope,
type ServerHello,
@ -24,6 +25,7 @@ import type { ConnectionState, ConnectionStateChange, PiClientOptions, PiLaneWat
interface PendingRequest {
resolve(result: unknown): void;
reject(error: Error): void;
cleanup(): void;
}
interface ActiveWatchListener {
@ -112,6 +114,12 @@ export class PiClient {
return () => this.#connectionStateListeners.delete(listener);
}
/** Invoke an untyped RPC method. Domain-specific facades own argument and result typing. */
invoke(method: string, args: unknown, signal?: AbortSignal): Promise<unknown> {
if (method.length === 0) return Promise.reject(new TypeError("RPC method must not be empty"));
return this.#request({ method, args }, signal);
}
listSessions(): Promise<readonly SessionMetadata[]> {
return this.#rpc.list();
}
@ -178,12 +186,44 @@ export class PiClient {
};
}
#request(call: ServiceRpcCall): Promise<unknown> {
#request(call: ProtocolRpcCall | ServiceRpcCall, signal?: AbortSignal): Promise<unknown> {
if (this.#disposed) return Promise.reject(new PiClientDisposedError());
if (!this.connected) return Promise.reject(new PiDisconnectedError());
if (signal?.aborted) return Promise.reject(abortError(signal));
const id = `request-${++this.#requestSequence}`;
const { promise, resolve, reject } = createPromiseResolvers<unknown>();
this.#pendingRequests.set(id, { resolve, reject });
let sent = false;
let aborted = false;
let onAbort: (() => void) | undefined;
const sendCancel = (): void => {
if (!sent || !this.connected) return;
try {
this.#connection.send(
encodeClientMessage(
{ type: "cancel", id, serverId: this.#options.serverId },
{ maxFrameLength: this.#connection.maxFrameLength },
),
);
} catch (error) {
this.#connection.fail(toError(error));
}
};
if (signal !== undefined) {
onAbort = () => {
if (aborted) return;
aborted = true;
reject(abortError(signal));
sendCancel();
};
signal.addEventListener("abort", onAbort, { once: true });
}
this.#pendingRequests.set(id, {
resolve,
reject,
cleanup: () => {
if (signal !== undefined && onAbort !== undefined) signal.removeEventListener("abort", onAbort);
},
});
let frame: Uint8Array;
try {
frame = encodeClientMessage(
@ -195,6 +235,8 @@ export class PiClient {
return promise;
}
this.#connection.send(frame);
sent = true;
if (aborted) sendCancel();
return promise;
}
@ -236,14 +278,20 @@ export class PiClient {
#takePendingRequest(id: string): PendingRequest | undefined {
const request = this.#pendingRequests.get(id);
if (request) this.#pendingRequests.delete(id);
if (request) {
this.#pendingRequests.delete(id);
request.cleanup();
}
return request;
}
#rejectPendingRequests(error: Error): void {
const requests = [...this.#pendingRequests.values()];
this.#pendingRequests.clear();
for (const request of requests) request.reject(error);
for (const request of requests) {
request.cleanup();
request.reject(error);
}
}
dispose(): Promise<void> {
@ -276,3 +324,8 @@ export class PiClient {
}
}
}
function abortError(signal: AbortSignal): Error {
const reason: unknown = signal.reason;
return reason instanceof Error ? reason : new DOMException("The operation was aborted", "AbortError");
}

View file

@ -263,6 +263,49 @@ describe("PiClient service operations", () => {
await client.dispose();
});
test("does not send a pre-aborted untyped RPC request", async () => {
const server = new MemoryByteServer();
const client = await connectClient(server);
const controller = new AbortController();
const reason = new Error("already cancelled");
controller.abort(reason);
await expect(client.invoke("session.test", {}, controller.signal)).rejects.toBe(reason);
expect(server.messages).toHaveLength(1);
await client.dispose();
});
test("cancels one untyped RPC request without disconnecting", async () => {
const server = new MemoryByteServer();
const client = await connectClient(server);
const controller = new AbortController();
const reason = new Error("stop this request");
const pending = client.invoke("session.test", { value: 42 }, controller.signal);
await server.waitForMessages(2);
expect(server.messages[1]).toMatchObject({
type: "request",
id: "request-1",
call: { method: "session.test", args: { value: 42 } },
});
controller.abort(reason);
await expect(pending).rejects.toBe(reason);
await server.waitForMessages(3);
expect(server.messages[2]).toEqual({
type: "cancel",
id: "request-1",
serverId: "00000000-0000-4000-8000-000000000001",
});
server.send({
type: "response",
id: "request-1",
ok: false,
error: { code: "cancelled", message: "cancelled" },
});
expect(client.connected).toBe(true);
await client.dispose();
});
test("rejects pending requests after disconnect or disposal", async () => {
const server = new MemoryByteServer();
const client = await connectClient(server);

View file

@ -67,6 +67,7 @@ describe.runIf(process.platform !== "win32")("createUnixTransportFactory", () =>
for (const byte of frame) socket.write(Uint8Array.of(byte));
continue;
}
if (message.type === "cancel") continue;
receivedMethods.push(message.call.method);
const frame = encodeServerMessage({
type: "response",

View file

@ -1,4 +1,5 @@
import { basename } from "node:path";
import { BACKGROUND_CONTEXT, RemoteSession, type RemoteSessionRpc } from "@earendil-works/pi-agent-core";
import { PiClient } from "@earendil-works/pi-client";
import { createUnixTransportFactory, discoverUnixServers, type UnixServerRoute } from "@earendil-works/pi-client/unix";
import { isServerId, type LaneEvent } from "@earendil-works/pi-protocol";
@ -16,7 +17,7 @@ export type ClientResult =
export interface RunClientOptions {
/** Directory searched when --connect is omitted. Defaults to PI_SERVER_DIR or ~/.pi/server. */
readonly directory?: string;
/** Receives snapshot-ordered main-lane events while a prompt is running. */
/** Reserved for snapshot-ordered main-lane events once remote harness prompting is implemented. */
readonly onEvent?: (event: LaneEvent) => void | Promise<void>;
}
@ -109,34 +110,19 @@ export async function runClient(command: ClientCommand, options: RunClientOption
await match.client.createSession({ id: selectedSessionId, cwd: process.cwd() });
}
}
const attached = await match.client.attachSession(sessionId);
if (command.prompt === undefined) {
return { kind: "attached", serverId: match.route.serverId, sessionId: attached.sessionId };
}
const prompt = command.prompt;
const onEvent = options.onEvent;
const watch = onEvent === undefined ? undefined : await match.client.watchSession(attached.sessionId);
if (watch !== undefined && onEvent !== undefined) await watch.start(onEvent);
const result = await (async () => {
try {
return await match.client.promptSession(attached.sessionId, prompt);
} finally {
await watch?.dispose();
}
})();
if (!result.ok) throw new Error(result.error.message);
if (result.value.kind === "failed") throw new Error(result.value.error.message);
if (result.value.kind !== "completed") throw new Error(`Remote prompt ${result.value.kind}`);
if (!("finalMessage" in result.value)) throw new Error("Remote prompt completed without an assistant message");
const text = result.value.finalMessage.content
.flatMap((part) => (part.type === "text" ? [part.text] : []))
.join("");
return {
kind: "prompted",
serverId: match.route.serverId,
sessionId: attached.sessionId,
text,
const rpc: RemoteSessionRpc = {
invoke: (method, args, context) => match.client.invoke(method, args, context.abortSignal),
};
const session = await RemoteSession.open(rpc, sessionId, BACKGROUND_CONTEXT);
if (command.prompt === undefined) {
await session.close(BACKGROUND_CONTEXT);
return { kind: "attached", serverId: match.route.serverId, sessionId: session.metadata.id };
}
try {
throw new Error("Experimental client prompting through AgentHarness is not implemented");
} finally {
await session.close(BACKGROUND_CONTEXT);
}
} finally {
await Promise.all([...openedClients].map((client) => client.dispose()));
}

View file

@ -8,7 +8,7 @@ import { NodeExecutionEnv } from "@earendil-works/pi-agent-core/node";
import { PiClient, PiDisconnectedError, PiServerError } from "@earendil-works/pi-client";
import { createUnixTransportFactory, type UnixServerRoute } from "@earendil-works/pi-client/unix";
import { isServerId, type ServerId } from "@earendil-works/pi-protocol";
import type { PiServer, PiServerHost } from "@earendil-works/pi-server";
import { type PiServer, type PiServerHost, SessionInUseError } from "@earendil-works/pi-server";
import { createUnixServer, getUnixSocketPath } from "@earendil-works/pi-server/unix";
import lockfile from "proper-lockfile";
import { getAgentDir } from "../config.ts";
@ -346,6 +346,12 @@ async function startServerBackend(
await session.close(context);
}
},
open: (metadata, context) => {
if (workers.hasSession(metadata.path)) {
throw new SessionInUseError();
}
return repo.open(metadata, context);
},
},
createHarness: (metadata, context) => workers.createHarness(metadata, context),
};

View file

@ -120,6 +120,10 @@ export class SessionWorkerManager {
return [...this.#workersBySession.values()].map((worker) => worker.metadata);
}
hasSession(path: string): boolean {
return this.#workersBySession.has(path) || this.#pending.has(path);
}
async discover(peerIds: ReadonlySet<string>): Promise<void> {
if (this.#detached) return;
const undiscovered = new Set(

View file

@ -666,6 +666,7 @@ export async function main(args: string[], options?: MainOptions) {
}
if (await runExperimentalCommand(args)) {
if (args[0] === "client") process.exit(process.exitCode ?? 0);
return;
}

View file

@ -291,52 +291,16 @@ describe("experimental durable server composition", () => {
await expect(competing.attachSession("demo-1")).rejects.toMatchObject({ code: "session_in_use" });
});
// WP00 intentionally selects execution-incomplete runtime2. Re-enable with the no-tool run package.
test.skip("creates generated and requested Sessions for one-shot prompts", async ({ onTestFinished }) => {
const spawn = vi
.spyOn(processRuntime, "spawnInternalProcess")
.mockImplementation((role, args, options) =>
realSpawnInternalProcess(
role,
args,
role === "session-worker" ? { ...options, entryUrl: fauxWorkerEntryUrl } : options,
),
);
onTestFinished(() => spawn.mockRestore());
const { directory, runtime } = await makeServer();
const prompt = (sessionId?: string) =>
runClient(
{ command: "client", prompt: "question", ...(sessionId === undefined ? {} : { sessionId }) },
{ directory },
);
test("rejects prompting until AgentHarness prompting is implemented", async () => {
const { directory } = await makeServer();
const generated = await prompt();
expect(generated).toMatchObject({
kind: "prompted",
serverId: runtime.serverId,
sessionId: expect.any(String),
text: "deterministic remote answer",
await expect(
runClient({ command: "client", sessionId: "demo-1", prompt: "question" }, { directory }),
).rejects.toThrow("Experimental client prompting through AgentHarness is not implemented");
await expect(runClient({ command: "client", sessionId: "demo-1" }, { directory })).resolves.toMatchObject({
kind: "attached",
sessionId: "demo-1",
});
const requested = await prompt("created-by-prompt");
expect(requested).toEqual({
kind: "prompted",
serverId: runtime.serverId,
sessionId: "created-by-prompt",
text: "deterministic remote answer",
});
for (const result of [generated, requested]) {
if (result.kind !== "prompted") throw new Error("Expected a prompted client result");
await expect.poll(() => runtime.workerPids.has(result.sessionId)).toBe(false);
const { branch } = await readExperimentalSessionState(runtime.sessionDir, result.sessionId);
expect(branch).toHaveLength(2);
expect(branch[0]).toMatchObject({
message: { role: "user", content: [{ type: "text", text: "question" }] },
});
expect(branch[1]).toMatchObject({
message: { role: "assistant", content: [{ type: "text", text: "deterministic remote answer" }] },
});
}
});
// WP00 intentionally selects execution-incomplete runtime2. Re-enable with the no-tool run package.

View file

@ -87,10 +87,18 @@ export type ServiceRpcResultUnion = RpcResultUnion<ServiceRpcManifest>;
export const ServiceRpcCallSchema = Type.Unsafe<ServiceRpcCall>(createRpcCallSchema(ServiceRpc));
export const ServiceRpcResultSchema = Type.Unsafe<ServiceRpcResultUnion>(createRpcResultSchema(ServiceRpc));
export type ProtocolRpcCall = ServiceRpcCall;
export type ProtocolRpcResult = ServiceRpcResultUnion;
const ProtocolRpcCallSchema = ServiceRpcCallSchema;
const ProtocolRpcResultSchema = ServiceRpcResultSchema;
/** Untyped transport call. Domain facades own method typing and result decoding. */
export const ProtocolRpcCallSchema = StrictObject({
method: Type.String({ minLength: 1 }),
args: Type.Unknown(),
});
export type ProtocolRpcCall = Static<typeof ProtocolRpcCallSchema>;
export type ProtocolRpcResult = unknown;
const ProtocolRpcResultSchema = Type.Unknown();
export function isServiceRpcCall(call: ProtocolRpcCall): call is ServiceRpcCall {
return Check(ServiceRpcCallSchema, call);
}
export const ProtocolErrorCodeSchema = Type.Union([
Type.Literal("version"),
@ -104,6 +112,14 @@ export const ProtocolErrorCodeSchema = Type.Union([
Type.Literal("not_supported"),
Type.Literal("server_draining"),
Type.Literal("invalid_request"),
Type.Literal("cancelled"),
Type.Literal("session_invalid_lane"),
Type.Literal("session_lane_exists"),
Type.Literal("session_unknown_target"),
Type.Literal("session_pending_message"),
Type.Literal("session_invariant"),
Type.Literal("mutation_not_found"),
Type.Literal("mutation_expired"),
Type.Literal("internal_error"),
]);
export const ProtocolErrorSchema = StrictObject({
@ -126,8 +142,14 @@ export const RequestEnvelopeSchema = StrictObject({
serverId: ServerIdSchema,
call: ProtocolRpcCallSchema,
});
export const CancelEnvelopeSchema = StrictObject({
type: Type.Literal("cancel"),
id: IdSchema,
serverId: ServerIdSchema,
});
export type RequestEnvelope = Static<typeof RequestEnvelopeSchema>;
export const ClientMessageSchema = Type.Union([ClientHelloSchema, RequestEnvelopeSchema]);
export type CancelEnvelope = Static<typeof CancelEnvelopeSchema>;
export const ClientMessageSchema = Type.Union([ClientHelloSchema, RequestEnvelopeSchema, CancelEnvelopeSchema]);
export type ClientMessage = Static<typeof ClientMessageSchema>;
export const ServerHelloSchema = StrictObject({

View file

@ -29,6 +29,7 @@ import {
ServerMessageDecoder,
ServiceRpc,
type SessionMetadata,
SessionMetadataSchema,
} from "../src/index.ts";
const clientHello: ClientHello = { type: "hello", version: PROTOCOL_VERSION };
@ -191,7 +192,7 @@ describe("protocol validation", () => {
).toThrow(ProtocolValidationError);
});
test("validates service RPC calls with logical targets", () => {
test("keeps transport RPC calls untyped while validating their envelope", () => {
const list: ClientMessage = {
type: "request",
id: "request-1",
@ -220,18 +221,29 @@ describe("protocol validation", () => {
expect(parseClientMessage(create)).toEqual(create);
expect(parseClientMessage(attach)).toEqual(attach);
expect(parseClientMessage(prompt)).toEqual(prompt);
expect(() => parseClientMessage({ ...create, call: { method: "create", args: [{ cwd: "" }] } })).toThrow(
ProtocolValidationError,
);
expect(() => parseClientMessage({ ...attach, call: { method: "attach", args: [] } })).toThrow(
ProtocolValidationError,
);
expect(() => parseClientMessage({ ...attach, call: { method: "unknown", args: [] } })).toThrow(
ProtocolValidationError,
);
expect(() => parseClientMessage({ ...prompt, call: { method: "prompt", args: ["session-1", []] } })).toThrow(
ProtocolValidationError,
);
expect(parseClientMessage({ ...create, call: { method: "create", args: [{ cwd: "" }] } })).toMatchObject({
call: { method: "create" },
});
expect(parseClientMessage({ ...attach, call: { method: "attach", args: [] } })).toMatchObject({
call: { method: "attach", args: [] },
});
expect(parseClientMessage({ ...attach, call: { method: "unknown", args: { arbitrary: true } } })).toMatchObject({
call: { method: "unknown", args: { arbitrary: true } },
});
expect(parseClientMessage({ ...prompt, call: { method: "prompt", args: ["session-1", []] } })).toMatchObject({
call: { method: "prompt" },
});
});
test("validates request cancellation envelopes", () => {
const cancel: ClientMessage = {
type: "cancel",
id: "request-1",
serverId: "00000000-0000-4000-8000-000000000001",
};
expect(parseClientMessage(cancel)).toEqual(cancel);
expect(() => parseClientMessage({ ...cancel, id: "" })).toThrow(ProtocolValidationError);
expect(() => parseClientMessage({ ...cancel, extra: true })).toThrow(ProtocolValidationError);
});
test("validates recursively nested JSON values", () => {
@ -308,7 +320,9 @@ describe("protocol validation", () => {
expect(Check(PromptArgumentsSchema, promptArguments)).toBe(false);
});
test("validates protocol Session metadata", () => {
test("keeps Session metadata validation at the typed service boundary", () => {
expect(Check(SessionMetadataSchema, metadata)).toBe(true);
expect(Check(SessionMetadataSchema, { id: "session-1", createdAt: 1 })).toBe(false);
const message: ServerMessage = {
type: "response",
id: "request-1",
@ -316,9 +330,6 @@ describe("protocol validation", () => {
result: [metadata],
};
expect(parseServerMessage(message)).toEqual(message);
expect(() => parseServerMessage({ ...message, result: [{ id: "session-1", createdAt: 1 }] })).toThrow(
ProtocolValidationError,
);
});
test("validates lane watch snapshots and events", () => {
@ -449,15 +460,6 @@ describe("protocol validation", () => {
call: { method: "list", args: [] },
},
],
[
"empty session id",
{
type: "request",
id: "request-1",
serverId: "00000000-0000-4000-8000-000000000001",
call: { method: "attach", args: [""] },
},
],
[
"extra call field",
{
@ -488,6 +490,14 @@ describe("protocol validation", () => {
"watch_in_use",
"not_supported",
"server_draining",
"cancelled",
"session_invalid_lane",
"session_lane_exists",
"session_unknown_target",
"session_pending_message",
"session_invariant",
"mutation_not_found",
"mutation_expired",
"internal_error",
] as const)("accepts the %s error code", (code) => {
const message: ServerMessage = {

View file

@ -26,6 +26,7 @@ export interface ConnectionState {
disconnected: boolean;
handshake?: Promise<void>;
handshakeTimeout: NodeJS.Timeout;
activeRequests: Map<string, AbortController>;
}
export function isTerminalConnection(state: ConnectionState): boolean {

View file

@ -11,6 +11,8 @@ type PiServerOperationErrorCode = Extract<
| "watch_in_use"
| "not_supported"
| "server_draining"
| "mutation_not_found"
| "mutation_expired"
>;
export const INTERNAL_SERVER_ERROR_MESSAGE = "Internal server error";

View file

@ -69,6 +69,7 @@ interface HostedSession {
interface HostedHarnessManagerOptions<TMetadata extends SessionMetadata> {
host: PiServerHost<TMetadata>;
isClosing: () => boolean;
isSessionUnavailable?: (sessionId: string) => boolean;
reportError: (error: unknown) => void;
}
@ -115,6 +116,10 @@ export class HostedHarnessManager<TMetadata extends SessionMetadata = SessionMet
);
}
hasSession(sessionId: string): boolean {
return this.hostedSessions.has(sessionId) || this.openingSessions.has(sessionId);
}
executeCall(
call: ServiceRpcCall,
client: object,
@ -198,6 +203,7 @@ export class HostedHarnessManager<TMetadata extends SessionMetadata = SessionMet
const current = this.attachmentsByClient.get(client);
if (current?.session.id === sessionId) return { sessionId };
if (current) await this.releaseAttachment(current, context);
if (this.options.isSessionUnavailable?.(sessionId)) throw new SessionInUseError();
const hosted = await this.acquire(sessionId, context);
if (this.options.isClosing() || this.disconnectedClients.has(client)) throw new ServerDrainingError();

View file

@ -0,0 +1,515 @@
import { randomUUID } from "node:crypto";
import type {
AgentMessage,
BranchScan,
Context,
EntryQuery,
JsonValue,
LaneConfiguration,
ListReadOptions,
Session,
SessionMetadata,
SessionMutator,
Value,
ValueList,
Write,
} from "@earendil-works/pi-agent-core";
import { isJsonValue, ProtocolValidationError } from "@earendil-works/pi-protocol";
import {
NotSupportedError,
PiServerError,
ServerDrainingError,
SessionAmbiguousError,
SessionInUseError,
SessionNotFoundError,
} from "./errors.ts";
import type { PiServerHost } from "./types.ts";
export const DEFAULT_REMOTE_MUTATION_LEASE_MS = 30_000;
export class RemoteMutationNotFoundError extends PiServerError {
constructor() {
super("mutation_not_found", "Remote Session mutation was not found");
this.name = "RemoteMutationNotFoundError";
}
}
export class RemoteMutationExpiredError extends PiServerError {
constructor() {
super("mutation_expired", "Remote Session mutation lease expired");
this.name = "RemoteMutationExpiredError";
}
}
interface OpenSession<TMetadata extends SessionMetadata> {
readonly client: object;
readonly handle: string;
readonly session: Session<TMetadata>;
readonly mutations: Map<string, RemoteMutation>;
readonly expiredMutations: Set<string>;
closing?: Promise<void>;
}
interface RemoteMutation {
readonly id: string;
completion: Promise<void>;
readonly release: () => void;
mutator?: SessionMutator;
lease?: NodeJS.Timeout;
state: "opening" | "active" | "finishing" | "expired";
}
interface RemoteSessionManagerOptions<TMetadata extends SessionMetadata> {
readonly host: PiServerHost<TMetadata>;
readonly mutationLeaseMs: number;
readonly isSessionUnavailable?: (sessionId: string) => boolean;
}
/** Connection-scoped remote Session capabilities dispatched through untyped RPC method names. */
export class RemoteSessionManager<TMetadata extends SessionMetadata = SessionMetadata> {
readonly #options: RemoteSessionManagerOptions<TMetadata>;
readonly #sessionsByHandle = new Map<string, OpenSession<TMetadata>>();
readonly #sessionsById = new Map<string, OpenSession<TMetadata>>();
readonly #sessionsByClient = new Map<object, Set<OpenSession<TMetadata>>>();
readonly #closingSessions = new Set<Promise<void>>();
#openTail: Promise<void> = Promise.resolve();
#closing = false;
constructor(options: RemoteSessionManagerOptions<TMetadata>) {
this.#options = options;
}
hasSession(sessionId: string): boolean {
return this.#sessionsById.has(sessionId);
}
invoke(method: string, args: unknown, client: object, context: Context): Promise<unknown> {
if (method === "session.open")
return this.#open(expectString(expectArgs(args).sessionId, "sessionId"), client, context);
const params = expectArgs(args);
const owner = this.#requireOwner(expectString(params.handle, "handle"), client);
switch (method) {
case "session.close":
return this.#closeOwner(owner, context).then(() => null);
case "session.getEntries":
return owner.session
.getEntries(expectStringArray(params.ids, "ids"), context)
.then((entries) => [...entries]);
case "session.getValue":
return owner.session.getValue(expectValue(params.address), context).then((value) => value ?? null);
case "session.scanValues":
return owner.session.scanValues(expectValue(params.prefix), context);
case "session.readList":
return owner.session.readList(expectList(params.address), optionalOptions(params.options), context);
case "session.createLane":
return owner.session
.createLane(
expectString(params.name, "name"),
expectNullableString(params.at, "at"),
expectObject(params.configuration, "configuration") as unknown as LaneConfiguration,
context,
)
.then(() => null);
case "session.mutation.begin":
return this.#beginMutation(owner, expectString(params.mutationLane, "mutationLane"), context);
case "session.mutation.finish":
return this.#finishMutation(owner, expectString(params.mutationId, "mutationId")).then(() => null);
case "session.mutation.getEntries": {
const mutation = this.#requireMutation(owner, params);
return mutation
.mutator!.getEntries(expectStringArray(params.ids, "ids"), context)
.then((entries) => [...entries]);
}
case "session.mutation.getValue": {
const mutation = this.#requireMutation(owner, params);
return mutation.mutator!.getValue(expectValue(params.address), context).then((value) => value ?? null);
}
case "session.mutation.scanValues": {
const mutation = this.#requireMutation(owner, params);
return mutation.mutator!.scanValues(expectValue(params.prefix), context);
}
case "session.mutation.readList": {
const mutation = this.#requireMutation(owner, params);
return mutation.mutator!.readList(expectList(params.address), optionalOptions(params.options), context);
}
case "session.mutation.commit": {
const mutation = this.#requireMutation(owner, params);
if (!Array.isArray(params.writes)) throw new ProtocolValidationError("writes must be an array");
return mutation.mutator!.commit(params.writes as Write[], context);
}
}
const lane = expectString(params.lane, "lane");
const tree = owner.session.view(lane);
switch (method) {
case "session.tree.getLeafId":
return tree.getLeafId(context);
case "session.tree.getEntry":
return tree.getEntry(expectString(params.id, "id"), context).then((entry) => entry ?? null);
case "session.tree.getStats":
return tree.getStats(context);
case "session.tree.getValue":
return tree.getValue(expectValue(params.address), context).then((value) => value ?? null);
case "session.tree.scanValues":
return tree.scanValues(expectValue(params.prefix), context);
case "session.tree.readList":
return tree.readList(expectList(params.address), optionalOptions(params.options), context);
case "session.tree.setValue":
return tree.setValue(expectValue(params.address), required(params, "next"), context).then(() => null);
case "session.tree.deleteValue":
return tree.deleteValue(expectValue(params.address), context).then(() => null);
case "session.tree.appendList":
return tree.appendList(expectList(params.address), required(params, "element"), context).then(() => null);
case "session.tree.deleteList":
return tree.deleteList(expectList(params.address), context).then(() => null);
case "session.tree.getName":
return tree.getName(context).then((name) => name ?? null);
case "session.tree.setName":
return tree.setName(expectOptionalString(params.name, "name"), context).then(() => null);
case "session.tree.getLabel":
return tree.getLabel(expectString(params.targetId, "targetId"), context).then((label) => label ?? null);
case "session.tree.setLabel":
return tree
.setLabel(
expectString(params.targetId, "targetId"),
expectOptionalString(params.label, "label"),
context,
)
.then(() => null);
case "session.tree.findEntries":
return tree.findEntries(optionalEntryQuery(params.query), context);
case "session.tree.findEntry":
return tree.findEntry(optionalEntryQuery(params.query), context).then((entry) => entry ?? null);
case "session.tree.findEntriesOnBranch":
return tree.findEntriesOnBranch(optionalBranchScan(params.query), context);
case "session.tree.findEntryOnBranch":
return tree.findEntryOnBranch(optionalBranchScan(params.query), context).then((entry) => entry ?? null);
case "session.tree.appendMessage":
return tree.appendMessage(expectObject(params.message, "message") as unknown as AgentMessage, context);
case "session.tree.appendCustomEntry":
return tree.appendCustomEntry(
expectString(params.customType, "customType"),
optionalJsonValue(params, "data"),
context,
);
default:
return Promise.reject(new ProtocolValidationError(`Unknown RPC method ${method}`));
}
}
async disconnect(client: object, context: Context): Promise<void> {
const sessions = [...(this.#sessionsByClient.get(client) ?? [])];
const results = await Promise.allSettled(sessions.map((session) => this.#closeOwner(session, context)));
const errors = results.flatMap((result) => (result.status === "rejected" ? [result.reason] : []));
if (errors.length === 1) throw errors[0];
if (errors.length > 1) throw new AggregateError(errors, "Failed to close remote Sessions after disconnect");
}
async close(context: Context): Promise<void> {
this.#closing = true;
await this.#openTail.catch(() => {});
const sessions = [...this.#sessionsByHandle.values()];
const closing = sessions.map((session) => this.#closeOwner(session, context));
const results = await Promise.allSettled([...this.#closingSessions, ...closing]);
const errors = results.flatMap((result) => (result.status === "rejected" ? [result.reason] : []));
if (errors.length === 1) throw errors[0];
if (errors.length > 1) throw new AggregateError(errors, "Failed to close remote Sessions");
}
#open(sessionId: string, client: object, context: Context): Promise<unknown> {
return this.#serializeOpen(async () => {
if (this.#closing) throw new ServerDrainingError();
const existing = [...(this.#sessionsByClient.get(client) ?? [])].find(
(candidate) => candidate.session.metadata.id === sessionId,
);
if (existing) return { handle: existing.handle, metadata: toProtocolMetadata(existing.session.metadata) };
if (this.#sessionsById.has(sessionId) || this.#options.isSessionUnavailable?.(sessionId)) {
throw new SessionInUseError();
}
const open = this.#options.host.sessions.open;
if (open === undefined) throw new NotSupportedError("Host does not support remote Sessions");
const matches = (await this.#options.host.sessions.list(context)).filter(
(metadata) => metadata.id === sessionId,
);
if (matches.length === 0) throw new SessionNotFoundError(`Unknown session: ${sessionId}`);
if (matches.length > 1) throw new SessionAmbiguousError();
const session = await open(matches[0]!, context);
if (this.#closing) {
await session.close(context);
throw new ServerDrainingError();
}
const owner: OpenSession<TMetadata> = {
client,
handle: randomUUID(),
session,
mutations: new Map(),
expiredMutations: new Set(),
};
this.#sessionsByHandle.set(owner.handle, owner);
this.#sessionsById.set(sessionId, owner);
let clientSessions = this.#sessionsByClient.get(client);
if (clientSessions === undefined) {
clientSessions = new Set();
this.#sessionsByClient.set(client, clientSessions);
}
clientSessions.add(owner);
return { handle: owner.handle, metadata: toProtocolMetadata(session.metadata) };
});
}
#serializeOpen<T>(operation: () => Promise<T>): Promise<T> {
const result = this.#openTail.catch(() => {}).then(operation);
this.#openTail = result.then(
() => undefined,
() => undefined,
);
return result;
}
#requireOwner(handle: string, client: object): OpenSession<TMetadata> {
const owner = this.#sessionsByHandle.get(handle);
if (owner === undefined || owner.client !== client || owner.closing !== undefined) {
throw new SessionNotFoundError("Remote Session handle was not found");
}
return owner;
}
async #closeOwner(owner: OpenSession<TMetadata>, context: Context): Promise<void> {
if (owner.closing !== undefined) return owner.closing;
const closing = (async () => {
this.#sessionsByHandle.delete(owner.handle);
if (this.#sessionsById.get(owner.session.metadata.id) === owner) {
this.#sessionsById.delete(owner.session.metadata.id);
}
const clientSessions = this.#sessionsByClient.get(owner.client);
clientSessions?.delete(owner);
if (clientSessions?.size === 0) this.#sessionsByClient.delete(owner.client);
const mutations = [...owner.mutations.values()];
for (const mutation of mutations) mutation.release();
const errors = (await Promise.allSettled(mutations.map((mutation) => mutation.completion))).flatMap(
(result) => (result.status === "rejected" ? [result.reason] : []),
);
owner.expiredMutations.clear();
try {
await owner.session.close(context);
} catch (error) {
errors.push(error);
}
if (errors.length === 1) throw errors[0];
if (errors.length > 1) throw new AggregateError(errors, "Failed to close remote Session");
})();
owner.closing = closing;
this.#closingSessions.add(closing);
void closing.finally(() => this.#closingSessions.delete(closing)).catch(() => {});
return closing;
}
async #beginMutation(
owner: OpenSession<TMetadata>,
lane: string,
context: Context,
): Promise<{ mutationId: string }> {
const id = randomUUID();
let resolveEntered!: () => void;
let rejectEntered!: (error: unknown) => void;
const entered = new Promise<void>((resolve, reject) => {
resolveEntered = resolve;
rejectEntered = reject;
});
let resolveRelease!: () => void;
const released = new Promise<void>((resolve) => {
resolveRelease = resolve;
});
let releasedOnce = false;
const release = (): void => {
if (releasedOnce) return;
releasedOnce = true;
resolveRelease();
};
const resource: RemoteMutation = {
id,
completion: Promise.resolve(),
release,
state: "opening",
};
owner.mutations.set(id, resource);
const completion = owner.session
.mutate(
lane,
async (mutator) => {
if (resource.state !== "opening") return;
resource.mutator = mutator;
resource.state = "active";
resource.lease = setTimeout(() => {
if (resource.state !== "active") return;
resource.state = "expired";
owner.expiredMutations.add(resource.id);
resource.release();
}, this.#options.mutationLeaseMs);
resource.lease.unref();
resolveEntered();
await released;
},
context,
)
.catch((error: unknown) => {
rejectEntered(error);
throw error;
})
.finally(() => {
if (resource.lease) clearTimeout(resource.lease);
owner.mutations.delete(id);
});
resource.completion = completion;
void completion.catch(() => {});
const signal = context.abortSignal;
let onAbort: (() => void) | undefined;
const aborted =
signal === undefined
? undefined
: new Promise<never>((_resolve, reject) => {
onAbort = () => {
resource.state = "finishing";
resource.release();
reject(abortError(signal));
};
signal.addEventListener("abort", onAbort, { once: true });
});
try {
if (signal?.aborted) throw abortError(signal);
await (aborted === undefined ? entered : Promise.race([entered, aborted]));
if (signal?.aborted) throw abortError(signal);
return { mutationId: id };
} catch (error) {
resource.state = "finishing";
resource.release();
throw error;
} finally {
if (signal !== undefined && onAbort !== undefined) signal.removeEventListener("abort", onAbort);
}
}
#requireMutation(owner: OpenSession<TMetadata>, params: Record<string, unknown>): RemoteMutation {
const id = expectString(params.mutationId, "mutationId");
const mutation = owner.mutations.get(id);
if (mutation === undefined) {
if (owner.expiredMutations.has(id)) throw new RemoteMutationExpiredError();
throw new RemoteMutationNotFoundError();
}
if (mutation.state === "expired") throw new RemoteMutationExpiredError();
if (mutation.state !== "active" || mutation.mutator === undefined) throw new RemoteMutationNotFoundError();
return mutation;
}
async #finishMutation(owner: OpenSession<TMetadata>, id: string): Promise<void> {
const mutation = owner.mutations.get(id);
if (mutation === undefined) {
if (owner.expiredMutations.delete(id)) throw new RemoteMutationExpiredError();
throw new RemoteMutationNotFoundError();
}
if (mutation.state === "expired") {
await mutation.completion;
owner.expiredMutations.delete(id);
throw new RemoteMutationExpiredError();
}
mutation.state = "finishing";
mutation.release();
await mutation.completion;
}
}
function toProtocolMetadata(metadata: SessionMetadata): SessionMetadata {
return {
id: metadata.id,
createdAt: metadata.createdAt,
storageVersion: metadata.storageVersion,
...(metadata.cwd === undefined ? {} : { cwd: metadata.cwd }),
...(metadata.parentSessionId === undefined ? {} : { parentSessionId: metadata.parentSessionId }),
...(metadata.legacyParentSessionPath === undefined
? {}
: { legacyParentSessionPath: metadata.legacyParentSessionPath }),
};
}
function expectArgs(value: unknown): Record<string, unknown> {
return expectObject(value, "RPC arguments");
}
function expectObject(value: unknown, name: string): Record<string, unknown> {
if (typeof value !== "object" || value === null || Array.isArray(value)) {
throw new ProtocolValidationError(`${name} must be an object`);
}
const prototype = Object.getPrototypeOf(value);
if (prototype !== Object.prototype && prototype !== null) {
throw new ProtocolValidationError(`${name} must be a plain object`);
}
return value as Record<string, unknown>;
}
function expectString(value: unknown, name: string): string {
if (typeof value !== "string" || value.length === 0) {
throw new ProtocolValidationError(`${name} must be a non-empty string`);
}
return value;
}
function expectStringArray(value: unknown, name: string): string[] {
if (!Array.isArray(value) || value.some((item) => typeof item !== "string")) {
throw new ProtocolValidationError(`${name} must be an array of strings`);
}
return value;
}
function expectNullableString(value: unknown, name: string): string | null {
return value === null ? null : expectString(value, name);
}
function expectOptionalString(value: unknown, name: string): string | undefined {
return value === null ? undefined : expectString(value, name);
}
function expectValue(value: unknown): Value<unknown> {
const address = expectObject(value, "value address");
if (address.kind !== "value" || typeof address.namespace !== "string" || typeof address.key !== "string") {
throw new ProtocolValidationError("Invalid value address");
}
return address as unknown as Value<unknown>;
}
function expectList(value: unknown): ValueList<unknown> {
const address = expectObject(value, "list address");
if (address.kind !== "list" || typeof address.namespace !== "string" || typeof address.key !== "string") {
throw new ProtocolValidationError("Invalid list address");
}
return address as unknown as ValueList<unknown>;
}
function required(params: Record<string, unknown>, key: string): unknown {
if (!Object.hasOwn(params, key)) throw new ProtocolValidationError(`${key} is required`);
return params[key];
}
function optionalJsonValue(params: Record<string, unknown>, key: string): JsonValue | undefined {
if (!(key in params)) return undefined;
const value = params[key];
if (!isJsonValue(value)) throw new ProtocolValidationError(`${key} must be a JSON value`);
return value;
}
function optionalOptions(value: unknown): ListReadOptions | undefined {
return value === null ? undefined : (expectObject(value, "list options") as unknown as ListReadOptions);
}
function optionalEntryQuery(value: unknown): EntryQuery | undefined {
return value === null ? undefined : (expectObject(value, "entry query") as unknown as EntryQuery);
}
function optionalBranchScan(value: unknown): BranchScan | undefined {
return value === null ? undefined : (expectObject(value, "branch query") as unknown as BranchScan);
}
function abortError(signal: AbortSignal): Error {
const reason: unknown = signal.reason;
return reason instanceof Error ? reason : new DOMException("The operation was aborted", "AbortError");
}

View file

@ -1,11 +1,23 @@
import { BACKGROUND_CONTEXT, type SessionMetadata, TODO_CONTEXT } from "@earendil-works/pi-agent-core";
import {
BACKGROUND_CONTEXT,
SessionInvalidLaneError,
SessionInvariantError,
SessionLaneExistsError,
type SessionMetadata,
SessionPendingAssistantMessageError,
SessionUnknownTargetError,
TODO_CONTEXT,
withAbortSignal,
} from "@earendil-works/pi-agent-core";
import {
type CancelEnvelope,
type ClientHello,
type ClientMessage,
ClientMessageDecoder,
DEFAULT_MAX_FRAME_LENGTH,
encodeServerMessage,
isServerId,
isServiceRpcCall,
isSupportedProtocolVersion,
PROTOCOL_VERSION,
type ProtocolError,
@ -26,6 +38,7 @@ import {
import { INTERNAL_SERVER_ERROR_MESSAGE, PiServerError, WrongServerError } from "./errors.ts";
import { HostedHarnessManager } from "./hosted-harness-manager.ts";
import type { PiServerListener } from "./listener.ts";
import { DEFAULT_REMOTE_MUTATION_LEASE_MS, RemoteSessionManager } from "./remote-session-manager.ts";
import type { PiServerHost, PiServerOptions } from "./types.ts";
const DEFAULT_HANDSHAKE_TIMEOUT_MS = 5_000;
@ -44,6 +57,7 @@ export class PiServer<TMetadata extends SessionMetadata = SessionMetadata> {
private readonly onError: ((error: Error) => void) | undefined;
private readonly connections = new Set<ConnectionState>();
private readonly sessions: HostedHarnessManager<TMetadata>;
private readonly remoteSessions: RemoteSessionManager<TMetadata>;
private closing = false;
private closePromise?: Promise<void>;
private closedSettled = false;
@ -60,9 +74,15 @@ export class PiServer<TMetadata extends SessionMetadata = SessionMetadata> {
this.handshakeTimeoutMs = resolved.handshakeTimeoutMs;
this.onConnectionCountChanged = options.onConnectionCountChanged;
this.onError = options.onError;
this.remoteSessions = new RemoteSessionManager({
host,
mutationLeaseMs: resolved.remoteMutationLeaseMs,
isSessionUnavailable: (sessionId) => this.sessions.hasSession(sessionId),
});
this.sessions = new HostedHarnessManager({
host,
isClosing: () => this.closing,
isSessionUnavailable: (sessionId) => this.remoteSessions.hasSession(sessionId),
reportError: (error) => this.reportError(error),
});
this.closed = new Promise((resolve, reject) => {
@ -137,6 +157,7 @@ export class PiServer<TMetadata extends SessionMetadata = SessionMetadata> {
stage: "awaitingHello",
disconnected: false,
handshakeTimeout,
activeRequests: new Map(),
};
this.connections.add(state);
this.notifyConnectionCountChanged();
@ -223,14 +244,17 @@ export class PiServer<TMetadata extends SessionMetadata = SessionMetadata> {
}
if (state.stage === "ready") {
void this.handleRequest(state, message);
if (message.type === "cancel") this.handleCancel(state, message);
else void this.handleRequest(state, message);
return;
}
if (state.stage !== "handshaking") return;
const handshake = state.handshake;
if (!handshake) return;
void handshake.then(() => {
if (state.stage === "ready" && !state.disconnected) void this.handleRequest(state, message);
if (state.stage !== "ready" || state.disconnected) return;
if (message.type === "cancel") this.handleCancel(state, message);
else void this.handleRequest(state, message);
});
}
@ -255,17 +279,41 @@ export class PiServer<TMetadata extends SessionMetadata = SessionMetadata> {
}
}
private handleCancel(state: ConnectionState, envelope: CancelEnvelope): void {
if (envelope.serverId !== this.serverId) return;
state.activeRequests.get(envelope.id)?.abort(new DOMException("RPC request cancelled", "AbortError"));
}
private async handleRequest(state: ConnectionState, envelope: RequestEnvelope): Promise<void> {
if (state.activeRequests.has(envelope.id)) {
await this.sendMessage(state, {
type: "response",
id: envelope.id,
ok: false,
error: { code: "invalid_request", message: "Request ID is already active" },
} satisfies ResponseEnvelope);
return;
}
const controller = new AbortController();
state.activeRequests.set(envelope.id, controller);
const context = withAbortSignal(controller.signal, TODO_CONTEXT);
try {
if (envelope.serverId !== this.serverId) throw new WrongServerError();
const result: ProtocolRpcResult = await this.sessions.executeCall(
envelope.call,
state,
async (message, _context) => {
await this.sendMessage(state, message);
},
TODO_CONTEXT,
);
let result: ProtocolRpcResult;
if (isServiceRpcCall(envelope.call)) {
result = await this.sessions.executeCall(
envelope.call,
state,
async (message, _context) => {
await this.sendMessage(state, message);
},
context,
);
} else if (envelope.call.method.startsWith("session.")) {
result = await this.remoteSessions.invoke(envelope.call.method, envelope.call.args, state, context);
} else {
throw new ProtocolValidationError(`Unknown RPC method ${envelope.call.method}`);
}
await this.sendMessage(state, {
type: "response",
id: envelope.id,
@ -277,8 +325,12 @@ export class PiServer<TMetadata extends SessionMetadata = SessionMetadata> {
type: "response",
id: envelope.id,
ok: false,
error: this.toProtocolError(error),
error: controller.signal.aborted
? { code: "cancelled", message: "RPC request cancelled" }
: this.toProtocolError(error),
} satisfies ResponseEnvelope);
} finally {
if (state.activeRequests.get(envelope.id) === controller) state.activeRequests.delete(envelope.id);
}
}
@ -298,8 +350,17 @@ export class PiServer<TMetadata extends SessionMetadata = SessionMetadata> {
connection.disconnected = true;
connection.stage = "closed";
clearTimeout(connection.handshakeTimeout);
for (const controller of connection.activeRequests.values()) {
controller.abort(new Error("Client disconnected"));
}
connection.activeRequests.clear();
if (this.connections.delete(connection)) this.notifyConnectionCountChanged();
void this.sessions.disconnect(connection, TODO_CONTEXT).catch((error: unknown) => this.reportError(error));
void Promise.allSettled([
this.sessions.disconnect(connection, TODO_CONTEXT),
this.remoteSessions.disconnect(connection, TODO_CONTEXT),
]).then((results) => {
for (const result of results) if (result.status === "rejected") this.reportError(result.reason);
});
}
private async sendMessage(connection: ConnectionState, message: ServerMessage): Promise<boolean> {
@ -347,8 +408,14 @@ export class PiServer<TMetadata extends SessionMetadata = SessionMetadata> {
}
await Promise.all(connections.map((connection) => this.closeConnection(connection.connection)));
for (const connection of connections) this.disconnect(connection);
await this.sessions.close(BACKGROUND_CONTEXT);
const cleanup = await Promise.allSettled([
this.sessions.close(BACKGROUND_CONTEXT),
this.remoteSessions.close(BACKGROUND_CONTEXT),
]);
this.connections.clear();
const errors = cleanup.flatMap((result) => (result.status === "rejected" ? [result.reason] : []));
if (errors.length === 1) throw errors[0];
if (errors.length > 1) throw new AggregateError(errors, "Failed to close server Sessions");
}
private async closeConnection(connection: ByteConnection, finalChunk?: Uint8Array): Promise<void> {
@ -366,6 +433,13 @@ export class PiServer<TMetadata extends SessionMetadata = SessionMetadata> {
if (error instanceof ProtocolValidationError) {
return { code: "invalid_request", message: error.message };
}
if (error instanceof SessionInvalidLaneError) return { code: "session_invalid_lane", message: error.message };
if (error instanceof SessionLaneExistsError) return { code: "session_lane_exists", message: error.message };
if (error instanceof SessionUnknownTargetError) return { code: "session_unknown_target", message: error.message };
if (error instanceof SessionPendingAssistantMessageError) {
return { code: "session_pending_message", message: error.message };
}
if (error instanceof SessionInvariantError) return { code: "session_invariant", message: error.message };
this.reportError(error);
return { code: "internal_error", message: INTERNAL_SERVER_ERROR_MESSAGE };
}
@ -394,7 +468,11 @@ export class PiServer<TMetadata extends SessionMetadata = SessionMetadata> {
}
}
function resolveOptions(options: PiServerOptions): { maxFrameLength: number; handshakeTimeoutMs: number } {
function resolveOptions(options: PiServerOptions): {
maxFrameLength: number;
handshakeTimeoutMs: number;
remoteMutationLeaseMs: number;
} {
if (!Array.isArray(options.listeners)) throw new TypeError("PiServer listeners must be an array");
if (!isServerId(options.serverId)) {
throw new TypeError("PiServer serverId must be a canonical lowercase UUIDv4");
@ -411,5 +489,13 @@ function resolveOptions(options: PiServerOptions): { maxFrameLength: number; han
) {
throw new TypeError(`PiServer handshakeTimeoutMs must be an integer between 1 and ${MAX_TIMER_DELAY_MS}`);
}
return { maxFrameLength, handshakeTimeoutMs };
const remoteMutationLeaseMs = options.remoteMutationLeaseMs ?? DEFAULT_REMOTE_MUTATION_LEASE_MS;
if (
!Number.isSafeInteger(remoteMutationLeaseMs) ||
remoteMutationLeaseMs <= 0 ||
remoteMutationLeaseMs > MAX_TIMER_DELAY_MS
) {
throw new TypeError(`PiServer remoteMutationLeaseMs must be an integer between 1 and ${MAX_TIMER_DELAY_MS}`);
}
return { maxFrameLength, handshakeTimeoutMs, remoteMutationLeaseMs };
}

View file

@ -214,6 +214,7 @@ export class TestServerHost implements PiServerHost {
await session.close(context);
}
},
open: (metadata, context) => this.repo.open(metadata, context),
};
private nextListDelay?: ListDelay;
private nextCreateHarnessGate?: OpenGate;

View file

@ -1,4 +1,4 @@
import type { Context, SessionMetadata } from "@earendil-works/pi-agent-core";
import type { Context, Session, SessionMetadata } from "@earendil-works/pi-agent-core";
import type {
LaneEvent,
LaneSnapshot,
@ -16,6 +16,8 @@ export interface PiServerOptions {
handshakeTimeoutMs?: number;
onConnectionCountChanged?: (count: number) => void;
onError?: (error: Error) => void;
/** Maximum time an acquired remote Session mutation may hold its lane. */
remoteMutationLeaseMs?: number;
}
export type MaybePromise<T> = T | Promise<T>;
@ -50,6 +52,8 @@ export interface PiServerHost<TMetadata extends SessionMetadata = SessionMetadat
readonly sessions: {
list(context: Context): Promise<TMetadata[]>;
create(options: SessionCreateOptions, context: Context): Promise<TMetadata>;
/** Open one exclusive Session capability for a remote harness client. */
open?(metadata: TMetadata, context: Context): Promise<Session<TMetadata>>;
};
createHarness(metadata: TMetadata, context: Context): Promise<HostedHarnessHandle>;
}

View file

@ -102,6 +102,93 @@ describe("Session protocol", () => {
expect(host.harnesses.size).toBe(0);
});
test("routes untyped remote Session calls through a connection-owned handle", async () => {
const host = new TestServerHost();
await host.seed("session-1");
const client = connect(createServer(host));
await client.hello();
const opened = await client.request("00000000-0000-4000-8000-000000000001", {
method: "session.open",
args: { sessionId: "session-1" },
});
expect(opened).toMatchObject({ ok: true, result: { metadata: { id: "session-1" } } });
if (!opened.ok || typeof opened.result !== "object" || opened.result === null || !("handle" in opened.result)) {
throw new Error("Missing remote Session handle");
}
const handle = opened.result.handle;
expect(typeof handle).toBe("string");
await expect(
client.request("00000000-0000-4000-8000-000000000001", {
method: "session.tree.getLeafId",
args: { handle, lane: "main" },
}),
).resolves.toMatchObject({ ok: true, result: null });
await expect(
client.request("00000000-0000-4000-8000-000000000001", {
method: "session.close",
args: { handle },
}),
).resolves.toMatchObject({ ok: true, result: null });
});
test("cancels a queued remote mutation without cancelling its lane owner", async () => {
const host = new TestServerHost();
await host.seed("session-1");
const client = connect(createServer(host));
const serverId = "00000000-0000-4000-8000-000000000001";
await client.hello();
const opened = await client.request(serverId, {
method: "session.open",
args: { sessionId: "session-1" },
});
if (!opened.ok || typeof opened.result !== "object" || opened.result === null || !("handle" in opened.result)) {
throw new Error("Missing remote Session handle");
}
const handle = opened.result.handle;
if (typeof handle !== "string") throw new Error("Invalid remote Session handle");
const first = await client.request(serverId, {
method: "session.mutation.begin",
args: { handle, lane: "main", mutationLane: "main" },
});
if (!first.ok || typeof first.result !== "object" || first.result === null || !("mutationId" in first.result)) {
throw new Error("Missing remote mutation ID");
}
const firstMutationId = first.result.mutationId;
if (typeof firstMutationId !== "string") throw new Error("Invalid remote mutation ID");
const queued = client.request(
serverId,
{
method: "session.mutation.begin",
args: { handle, lane: "main", mutationLane: "main" },
},
"queued-mutation",
);
await Promise.resolve();
await client.sendMessage({ type: "cancel", id: "queued-mutation", serverId });
await expect(queued).resolves.toMatchObject({ ok: false, error: { code: "cancelled" } });
await expect(
client.request(serverId, {
method: "session.mutation.finish",
args: { handle, mutationId: firstMutationId },
}),
).resolves.toMatchObject({ ok: true, result: null });
const next = await client.request(serverId, {
method: "session.mutation.begin",
args: { handle, lane: "main", mutationLane: "main" },
});
expect(next).toMatchObject({ ok: true, result: { mutationId: expect.any(String) } });
if (!next.ok || typeof next.result !== "object" || next.result === null || !("mutationId" in next.result)) {
throw new Error("Missing replacement mutation ID");
}
await client.request(serverId, {
method: "session.mutation.finish",
args: { handle, mutationId: next.result.mutationId },
});
});
test("attach passes concrete repository metadata to the Harness host", async () => {
type BackendMetadata = SessionMetadata & { path: string; modifiedAt: number };
const metadata: BackendMetadata = {
@ -234,6 +321,38 @@ describe("Session protocol", () => {
expect(client.messages.slice(stoppedAt).some((message) => message.type === "event")).toBe(false);
});
test("keeps hosted Harness and remote Session ownership mutually exclusive", async () => {
const host = new TestServerHost();
await host.seed("session-1");
const server = createServer(host);
const remoteClient = connect(server);
const hostedClient = connect(server);
const serverId = "00000000-0000-4000-8000-000000000001";
await Promise.all([remoteClient.hello(), hostedClient.hello()]);
const opened = await remoteClient.request(serverId, {
method: "session.open",
args: { sessionId: "session-1" },
});
await expect(hostedClient.request(serverId, { method: "attach", args: ["session-1"] })).resolves.toMatchObject({
ok: false,
error: { code: "session_in_use" },
});
if (!opened.ok || typeof opened.result !== "object" || opened.result === null || !("handle" in opened.result)) {
throw new Error("Missing remote Session handle");
}
await remoteClient.request(serverId, {
method: "session.close",
args: { handle: opened.result.handle },
});
await expect(hostedClient.request(serverId, { method: "attach", args: ["session-1"] })).resolves.toMatchObject({
ok: true,
});
await expect(
remoteClient.request(serverId, { method: "session.open", args: { sessionId: "session-1" } }),
).resolves.toMatchObject({ ok: false, error: { code: "session_in_use" } });
});
test("permits only one client attachment per Session", async () => {
const host = new TestServerHost();
await host.seed("session-1");

View file

@ -0,0 +1,180 @@
import {
AgentHarness,
BACKGROUND_CONTEXT,
MemorySessionRepo,
RemoteSession,
setValue,
value,
} from "@earendil-works/pi-agent-core";
import { createModels, type Model } from "@earendil-works/pi-ai";
import { afterEach, describe, expect, test } from "vitest";
import { RemoteSessionManager } from "../src/remote-session-manager.ts";
import type { PiServerHost } from "../src/types.ts";
const repos = new Set<MemorySessionRepo>();
async function createRemoteSession(
id: string,
mutationLeaseMs = 1_000,
): Promise<{
remote: RemoteSession;
manager: RemoteSessionManager;
repo: MemorySessionRepo;
client: object;
}> {
const repo = new MemorySessionRepo();
repos.add(repo);
const created = await repo.create({ id }, BACKGROUND_CONTEXT);
await created.close(BACKGROUND_CONTEXT);
const host: PiServerHost = {
sessions: {
list: (context) => repo.list(undefined, context),
create: async (options, context) => {
const session = await repo.create(options, context);
try {
return session.metadata;
} finally {
await session.close(context);
}
},
open: (metadata, context) => repo.open(metadata, context),
},
createHarness: async () => {
throw new Error("Harness hosting is not used by remote Session tests");
},
};
const manager = new RemoteSessionManager({ host, mutationLeaseMs });
const client = {};
const remote = await RemoteSession.open(
{
invoke: (method, args, context) => manager.invoke(method, args, client, context),
},
id,
BACKGROUND_CONTEXT,
);
return { remote, manager, repo, client };
}
afterEach(async () => {
await Promise.all([...repos].map((repo) => repo.close(BACKGROUND_CONTEXT)));
repos.clear();
});
describe("RemoteSession", () => {
test("preserves callback-scoped mutation and tree behavior over untyped RPC", async () => {
const { remote } = await createRemoteSession("remote-session");
const applicationValue = value<string>("app.test", "value");
let capturedMutator: Parameters<Parameters<typeof remote.mutate>[1]>[0] | undefined;
await expect(
remote.mutate(
"main",
async (mutator, context) => {
capturedMutator = mutator;
await expect(mutator.getValue(applicationValue, context)).resolves.toBeUndefined();
const committed = mutator.commit([setValue(applicationValue, "committed")], context);
await expect(mutator.commit([], context)).rejects.toThrow("already attempted");
await committed;
return "local-result";
},
BACKGROUND_CONTEXT,
),
).resolves.toBe("local-result");
await expect(remote.getValue(applicationValue, BACKGROUND_CONTEXT)).resolves.toMatchObject({
value: "committed",
});
await expect(capturedMutator?.getValue(applicationValue, BACKGROUND_CONTEXT)).rejects.toThrow(
"outside its mutation callback",
);
await remote.setValue(applicationValue, "tree-write", BACKGROUND_CONTEXT);
await expect(remote.getValue(applicationValue, BACKGROUND_CONTEXT)).resolves.toMatchObject({
value: "tree-write",
});
const review = await remote.createLane(
"review",
null,
{
model: { provider: "test", modelId: "model" },
thinkingLevel: "off",
activeToolNames: [],
},
BACKGROUND_CONTEXT,
);
await expect(review.getLeafId(BACKGROUND_CONTEXT)).resolves.toBeNull();
await remote.close(BACKGROUND_CONTEXT);
});
test("expires a wedged mutation and releases its lane", async () => {
const { remote } = await createRemoteSession("remote-expiry", 10);
const applicationValue = value<string>("app.test", "expiry");
await expect(
remote.mutate(
"main",
async (mutator, context) => {
await new Promise<void>((resolve) => setTimeout(resolve, 25));
await mutator.getValue(applicationValue, context);
},
BACKGROUND_CONTEXT,
),
).rejects.toThrow(/mutation/i);
await remote.setValue(applicationValue, "after-expiry", BACKGROUND_CONTEXT);
await expect(remote.getValue(applicationValue, BACKGROUND_CONTEXT)).resolves.toMatchObject({
value: "after-expiry",
});
await remote.close(BACKGROUND_CONTEXT);
});
test("releases active mutations and the Session on disconnect", async () => {
const { remote, manager, repo, client } = await createRemoteSession("remote-disconnect");
let markEntered!: () => void;
const entered = new Promise<void>((resolve) => {
markEntered = resolve;
});
let releaseCallback!: () => void;
const callbackGate = new Promise<void>((resolve) => {
releaseCallback = resolve;
});
const mutation = remote.mutate(
"main",
async () => {
markEntered();
await callbackGate;
},
BACKGROUND_CONTEXT,
);
const rejectedMutation = expect(mutation).rejects.toThrow(/Session|mutation/i);
await entered;
await manager.disconnect(client, BACKGROUND_CONTEXT);
releaseCallback();
await rejectedMutation;
const metadata = (await repo.list(undefined, BACKGROUND_CONTEXT))[0]!;
const reopened = await repo.open(metadata, BACKGROUND_CONTEXT);
await reopened.close(BACKGROUND_CONTEXT);
});
test("can be passed to AgentHarness.create without local/remote branching", async () => {
const { remote } = await createRemoteSession("remote-harness");
const model = {
id: "model",
name: "Model",
api: "test",
provider: "test",
baseUrl: "https://example.test",
reasoning: false,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 10_000,
maxTokens: 1_000,
} satisfies Model<"test">;
const { harness } = await AgentHarness.create(
{ session: remote, models: createModels(), model, resources: {} },
BACKGROUND_CONTEXT,
);
await expect(harness.getModel(BACKGROUND_CONTEXT)).resolves.toBeUndefined();
await harness.close(BACKGROUND_CONTEXT);
});
});