test(server): consolidate reusable test infrastructure

This commit is contained in:
Christian Klotz 2026-07-31 17:59:17 +03:00
parent e36b150f71
commit 546e002353
No known key found for this signature in database
12 changed files with 770 additions and 1113 deletions

View file

@ -42,11 +42,15 @@ await server.start();
`PiServer` composes transport listeners through the `PiServerListener` interface. The Unix submodule exports the `createUnixListener()` building block and `createUnixServer()` preset, keeping the common case concise without coupling the primary server to Unix sockets. The listener uses authenticated, length-prefixed CBOR messages from `@earendil-works/pi-protocol`. It does not yet replace the legacy JSONL IPC control plane, child-process supervisor, standalone `server` CLI, or Radius presence integration.
## Transport testing
Custom transports can use `@earendil-works/pi-server/testing` for deterministic protocol conformance tests. It exports `createTestServer()`, `TestSessionBackend`, `ProtocolTestClient`, and the transport-neutral `WireChannel` contract. `connectUnixTestClient()` is provided for Unix transport tests.
## `pi-ai` protocol bridge
`@earendil-works/pi-ai` domain objects and `@earendil-works/pi-protocol` wire DTOs remain independent. This package owns their boundary and exports `toProtocolModelMetadata()`, `toProtocolAssistantMessage()`, `toProtocolUserMessage()`, and `toProtocolToolResultMessage()`.
The adapters normalize values that cannot cross the wire and exhaustively handle closed `pi-ai` unions. The protocol mirrors `pi-ai` vocabulary such as `toolCall` and `toolUse` where the semantics are identical. Compile-time assertions cover shared thinking-level and model-input vocabularies. Tests encode adapter output through the protocol runtime schemas so incompatible changes fail in the bridging package.
The adapters reject invalid tool inputs, explicitly sanitize diagnostic details, and exhaustively handle closed `pi-ai` unions. The protocol mirrors `pi-ai` vocabulary such as `toolCall` and `toolUse` where the semantics are identical. Compile-time assertions cover shared thinking-level and model-input vocabularies. Tests encode adapter output through the protocol runtime schemas so incompatible changes fail in the bridging package.
## Legacy server migration

View file

@ -10,6 +10,10 @@
"types": "./dist/index.d.ts",
"import": "./dist/index.js"
},
"./testing": {
"types": "./dist/testing/index.d.ts",
"import": "./dist/testing/index.js"
},
"./unix": {
"types": "./dist/transports/unix/index.d.ts",
"import": "./dist/transports/unix/index.js"

View file

@ -1,29 +1,20 @@
import { once } from "node:events";
import { createConnection, type Socket } from "node:net";
import {
type ClientMessage,
type Command,
encodeClientMessage,
type ModelMetadata,
type ModelRef,
PROTOCOL_VERSION,
type ResponseEnvelope,
type ServerMessage,
ServerMessageDecoder,
type SessionPhase,
type SessionSnapshot,
type SessionSummary,
type ThinkingLevel,
type TranscriptProgress,
import type {
ModelMetadata,
ModelRef,
SessionPhase,
SessionSnapshot,
SessionSummary,
ThinkingLevel,
TranscriptProgress,
} from "@earendil-works/pi-protocol";
import {
type CreateSessionOptions,
PiServerError,
type PiSessionBackend,
type PiSessionRuntime,
type PiSessionRuntimeEvent,
type PromptInput,
} from "../src/index.ts";
import { PiServerError } from "../errors.ts";
import type {
CreateSessionOptions,
PiSessionBackend,
PiSessionRuntime,
PiSessionRuntimeEvent,
PromptInput,
} from "../types.ts";
export const TEST_TOKEN = "server-conformance-token";
export const TEST_MODEL: ModelMetadata = {
@ -59,12 +50,14 @@ interface StoredSession {
snapshot: SessionSnapshot;
}
export class ConformanceRuntime implements PiSessionRuntime {
export class TestSessionRuntime implements PiSessionRuntime {
readonly disposed = new Deferred<void>();
disposeCount = 0;
readonly steers: PromptInput[] = [];
private readonly stored: StoredSession;
private readonly onDispose: () => void;
private readonly listeners = new Set<(event: PiSessionRuntimeEvent) => void>();
private pendingPrompt?: { input: PromptInput; done: Deferred<"complete" | "aborted"> };
constructor(stored: StoredSession, onDispose: () => void) {
this.stored = stored;
@ -80,8 +73,11 @@ export class ConformanceRuntime implements PiSessionRuntime {
}
async prompt(input: PromptInput): Promise<void> {
if (this.getPhase() !== "idle") throw new PiServerError("busy", "Session is busy");
if (this.getPhase() !== "idle") throw new PiServerError("busy", "A prompt is already running");
const done = new Deferred<"complete" | "aborted">();
this.pendingPrompt = { input, done };
this.update({
phase: "turn",
transcript: [
...this.stored.snapshot.transcript,
{
@ -92,21 +88,54 @@ export class ConformanceRuntime implements PiSessionRuntime {
},
],
});
const outcome = await done.promise;
this.update({
phase: "idle",
transcript: [
...this.stored.snapshot.transcript,
{
id: `assistant-${this.stored.snapshot.revision + 1}`,
role: "assistant",
content: [{ type: "text", text: outcome === "complete" ? `reply:${input.text}` : "" }],
status: outcome === "complete" ? "complete" : "aborted",
model: this.stored.snapshot.model,
stopReason: outcome === "complete" ? "stop" : "aborted",
timestamp: this.stored.snapshot.revision + 1,
},
],
});
this.pendingPrompt = undefined;
}
async steer(): Promise<void> {
throw new PiServerError("busy", "There is no active prompt to steer");
async steer(input: PromptInput): Promise<void> {
if (this.getPhase() === "idle") throw new PiServerError("busy", "There is no active prompt to steer");
this.steers.push(input);
this.update({
queuedSteerCount: this.stored.snapshot.queuedSteerCount + 1,
queuedSteer: [
...this.stored.snapshot.queuedSteer,
{
id: `steer-${this.stored.snapshot.revision + 1}`,
role: "user",
content: [{ type: "text", text: input.text }],
timestamp: this.stored.snapshot.revision + 1,
},
],
});
}
async abort(): Promise<void> {
throw new PiServerError("busy", "There is no active prompt to abort");
if (!this.pendingPrompt) throw new PiServerError("busy", "There is no active prompt to abort");
this.pendingPrompt.done.resolve("aborted");
}
async setModel(model: ModelRef): Promise<void> {
if (this.getPhase() !== "idle") throw new PiServerError("busy", "Session is busy");
this.update({ model });
}
async setThinking(thinkingLevel: ThinkingLevel): Promise<void> {
if (this.getPhase() !== "idle") throw new PiServerError("busy", "Session is busy");
this.update({ thinkingLevel });
}
@ -125,6 +154,11 @@ export class ConformanceRuntime implements PiSessionRuntime {
this.stored.snapshot = { ...this.stored.snapshot, phase };
}
finishPrompt(): void {
if (!this.pendingPrompt) throw new Error("No prompt is pending");
this.pendingPrompt.done.resolve("complete");
}
emitProgress(progress: TranscriptProgress): void {
for (const listener of this.listeners) listener({ type: "progress", progress });
}
@ -133,6 +167,10 @@ export class ConformanceRuntime implements PiSessionRuntime {
for (const listener of this.listeners) listener({ type: "error", error });
}
emitSnapshot(): void {
for (const listener of this.listeners) listener({ type: "snapshot" });
}
private update(updates: Partial<SessionSnapshot>): void {
this.stored.snapshot = {
...this.stored.snapshot,
@ -140,7 +178,7 @@ export class ConformanceRuntime implements PiSessionRuntime {
revision: this.stored.snapshot.revision + 1,
updatedAt: this.stored.snapshot.updatedAt + 1,
};
for (const listener of this.listeners) listener({ type: "snapshot" });
this.emitSnapshot();
}
}
@ -149,10 +187,11 @@ interface ListDelay {
release: Deferred<void>;
}
export class ConformanceBackend implements PiSessionBackend {
export class TestSessionBackend implements PiSessionBackend {
readonly sessions = new Map<string, StoredSession>();
readonly runtimes = new Map<string, ConformanceRuntime[]>();
readonly runtimes = new Map<string, TestSessionRuntime[]>();
readonly locked = new Set<string>();
lastCreatedId?: string;
private nextListDelay?: ListDelay;
async listSessions(): Promise<SessionSummary[]> {
@ -181,6 +220,7 @@ export class ConformanceBackend implements PiSessionBackend {
}
async createSession(options: CreateSessionOptions): Promise<PiSessionRuntime> {
this.lastCreatedId = options.id;
if (this.sessions.has(options.id)) throw new PiServerError("session_locked", "Session already exists");
this.seed(options.id, options.name, options.cwd, options.model, options.thinkingLevel);
return this.acquire(options.id);
@ -193,7 +233,7 @@ export class ConformanceBackend implements PiSessionBackend {
}
seed(
id: string,
id = "session-1",
name = `Session ${id}`,
cwd = "/tmp/pi-server-conformance",
model: ModelRef = { provider: TEST_MODEL.provider, id: TEST_MODEL.id },
@ -225,151 +265,20 @@ export class ConformanceBackend implements PiSessionBackend {
return delay;
}
latestRuntime(id: string): ConformanceRuntime {
latestRuntime(id: string): TestSessionRuntime {
const runtimes = this.runtimes.get(id);
if (!runtimes?.length) throw new Error(`No runtime for ${id}`);
return runtimes.at(-1)!;
}
private acquire(id: string): ConformanceRuntime {
private acquire(id: string): TestSessionRuntime {
const stored = this.sessions.get(id);
if (!stored) throw new Error(`Unknown session: ${id}`);
this.locked.add(id);
const runtime = new ConformanceRuntime(stored, () => this.locked.delete(id));
const runtime = new TestSessionRuntime(stored, () => this.locked.delete(id));
const runtimes = this.runtimes.get(id) ?? [];
runtimes.push(runtime);
this.runtimes.set(id, runtimes);
return runtime;
}
}
interface MessageWaiter {
predicate: (message: ServerMessage) => boolean;
resolve: (message: ServerMessage) => void;
reject: (error: Error) => void;
}
interface WireChannel {
send(chunk: Uint8Array): Promise<void>;
sendFragmented(chunk: Uint8Array, splitAt: number): Promise<void>;
close(): Promise<void>;
}
export class WireClient {
readonly messages: ServerMessage[] = [];
private readonly channel: WireChannel;
private readonly decoder = new ServerMessageDecoder();
private readonly waiters = new Set<MessageWaiter>();
private readonly closedDeferred = new Deferred<void>();
private requestSequence = 0;
private closedValue = false;
constructor(channel: WireChannel) {
this.channel = channel;
}
get closed(): boolean {
return this.closedValue;
}
hello(token = TEST_TOKEN, version: number = PROTOCOL_VERSION): Promise<ServerMessage> {
const response = this.next((message) => message.type === "hello" || message.type === "hello_error");
void this.sendMessage({ type: "hello", token, version });
return response;
}
async request(command: Command, id = `request-${++this.requestSequence}`): Promise<ResponseEnvelope> {
const response = this.next(
(message): message is ResponseEnvelope => message.type === "response" && message.id === id,
);
await this.sendMessage({ type: "request", id, request: command });
return (await response) as ResponseEnvelope;
}
sendMessage(message: ClientMessage): Promise<void> {
return this.channel.send(encodeClientMessage(message));
}
sendBytes(chunk: Uint8Array): Promise<void> {
return this.channel.send(chunk);
}
sendFragmentedMessage(message: ClientMessage, splitAt: number): Promise<void> {
return this.channel.sendFragmented(encodeClientMessage(message), splitAt);
}
next(predicate: (message: ServerMessage) => boolean): Promise<ServerMessage> {
const existing = this.messages.find(predicate);
if (existing) return Promise.resolve(existing);
if (this.closedValue) return Promise.reject(new Error("Wire client is closed"));
return new Promise((resolve, reject) => this.waiters.add({ predicate, resolve, reject }));
}
waitForClose(): Promise<void> {
return this.closedValue ? Promise.resolve() : this.closedDeferred.promise;
}
close(): Promise<void> {
return this.channel.close();
}
receive(chunk: Uint8Array): void {
try {
for (const message of this.decoder.push(chunk)) {
this.messages.push(message);
for (const waiter of this.waiters) {
if (!waiter.predicate(message)) continue;
this.waiters.delete(waiter);
waiter.resolve(message);
}
}
} catch (error) {
this.fail(error instanceof Error ? error : new Error(String(error)));
}
}
markClosed(): void {
if (this.closedValue) return;
this.closedValue = true;
this.closedDeferred.resolve(undefined);
this.fail(new Error("Wire connection closed"));
}
fail(error: Error): void {
for (const waiter of this.waiters) waiter.reject(error);
this.waiters.clear();
}
}
export async function connectUnix(path: string): Promise<WireClient> {
const socket = createConnection(path);
await once(socket, "connect");
const client = new WireClient({
send: (chunk) => writeSocket(socket, chunk),
async sendFragmented(chunk, splitAt) {
await writeSocket(socket, chunk.subarray(0, splitAt));
await writeSocket(socket, chunk.subarray(splitAt));
},
async close() {
if (socket.destroyed) return;
const closed = once(socket, "close");
socket.destroy();
await closed;
},
});
socket.on("data", (chunk) => {
client.receive(new Uint8Array(chunk.buffer, chunk.byteOffset, chunk.byteLength));
});
socket.on("error", (error) => client.fail(error));
socket.once("close", () => client.markClosed());
return client;
}
function writeSocket(socket: Socket, chunk: Uint8Array): Promise<void> {
return new Promise<void>((resolve, reject) => {
socket.write(chunk, (error) => {
if (error) reject(error);
else resolve();
});
});
}

View file

@ -0,0 +1,147 @@
import { once } from "node:events";
import { createConnection, type Socket } from "node:net";
import {
type ClientMessage,
type Command,
encodeClientMessage,
PROTOCOL_VERSION,
type ResponseEnvelope,
type ServerMessage,
ServerMessageDecoder,
} from "@earendil-works/pi-protocol";
import { Deferred, TEST_TOKEN } from "./backend.ts";
interface MessageWaiter {
predicate: (message: ServerMessage) => boolean;
resolve: (message: ServerMessage) => void;
reject: (error: Error) => void;
}
export interface WireChannel {
send(chunk: Uint8Array): Promise<void>;
sendFragmented(chunk: Uint8Array, splitAt: number): Promise<void>;
close(): Promise<void>;
}
export class ProtocolTestClient {
readonly messages: ServerMessage[] = [];
private readonly channel: WireChannel;
private readonly decoder = new ServerMessageDecoder();
private readonly waiters = new Set<MessageWaiter>();
private readonly closedDeferred = new Deferred<void>();
private requestSequence = 0;
private closedValue = false;
constructor(channel: WireChannel) {
this.channel = channel;
}
get closed(): boolean {
return this.closedValue;
}
hello(token = TEST_TOKEN, version: number = PROTOCOL_VERSION): Promise<ServerMessage> {
const response = this.next((message) => message.type === "hello" || message.type === "hello_error");
void this.sendMessage({ type: "hello", token, version });
return response;
}
async request(command: Command, id = `request-${++this.requestSequence}`): Promise<ResponseEnvelope> {
const response = this.next(
(message): message is ResponseEnvelope => message.type === "response" && message.id === id,
);
await this.sendMessage({ type: "request", id, request: command });
return (await response) as ResponseEnvelope;
}
sendMessage(message: ClientMessage): Promise<void> {
return this.channel.send(encodeClientMessage(message));
}
sendBytes(chunk: Uint8Array): Promise<void> {
return this.channel.send(chunk);
}
sendFragmentedMessage(message: ClientMessage, splitAt: number): Promise<void> {
return this.channel.sendFragmented(encodeClientMessage(message), splitAt);
}
next(predicate: (message: ServerMessage) => boolean): Promise<ServerMessage> {
return this.nextFrom(0, predicate);
}
nextFrom(index: number, predicate: (message: ServerMessage) => boolean): Promise<ServerMessage> {
const existing = this.messages.slice(index).find(predicate);
if (existing) return Promise.resolve(existing);
if (this.closedValue) return Promise.reject(new Error("Wire client is closed"));
return new Promise((resolve, reject) => this.waiters.add({ predicate, resolve, reject }));
}
waitForClose(): Promise<void> {
return this.closedValue ? Promise.resolve() : this.closedDeferred.promise;
}
close(): Promise<void> {
return this.channel.close();
}
receive(chunk: Uint8Array): void {
try {
for (const message of this.decoder.push(chunk)) {
this.messages.push(message);
for (const waiter of this.waiters) {
if (!waiter.predicate(message)) continue;
this.waiters.delete(waiter);
waiter.resolve(message);
}
}
} catch (error) {
this.fail(error instanceof Error ? error : new Error(String(error)));
}
}
markClosed(): void {
if (this.closedValue) return;
this.closedValue = true;
this.closedDeferred.resolve(undefined);
this.fail(new Error("Wire connection closed"));
}
fail(error: Error): void {
for (const waiter of this.waiters) waiter.reject(error);
this.waiters.clear();
}
}
export async function connectUnixTestClient(path: string): Promise<ProtocolTestClient> {
const socket = createConnection(path);
await once(socket, "connect");
const client = new ProtocolTestClient({
send: (chunk) => writeSocket(socket, chunk),
async sendFragmented(chunk, splitAt) {
await writeSocket(socket, chunk.subarray(0, splitAt));
await writeSocket(socket, chunk.subarray(splitAt));
},
async close() {
if (socket.destroyed) return;
const closed = once(socket, "close");
socket.destroy();
await closed;
},
});
socket.on("data", (chunk) => {
client.receive(new Uint8Array(chunk.buffer, chunk.byteOffset, chunk.byteLength));
});
socket.on("error", (error) => client.fail(error));
socket.once("close", () => client.markClosed());
return client;
}
function writeSocket(socket: Socket, chunk: Uint8Array): Promise<void> {
return new Promise<void>((resolve, reject) => {
socket.write(chunk, (error) => {
if (error) reject(error);
else resolve();
});
});
}

View file

@ -0,0 +1,5 @@
export { Deferred, TEST_MODEL, TEST_TOKEN, TestSessionBackend, TestSessionRuntime } from "./backend.ts";
export type { WireChannel } from "./client.ts";
export { connectUnixTestClient, ProtocolTestClient } from "./client.ts";
export type { TestServer, TestServerOptions } from "./server.ts";
export { createTestServer } from "./server.ts";

View file

@ -0,0 +1,29 @@
import { PiServer } from "../server.ts";
import type { PiServerOptions, PiSessionBackend } from "../types.ts";
import { TEST_TOKEN, TestSessionBackend } from "./backend.ts";
export interface TestServerOptions extends Omit<PiServerOptions, "token"> {
token?: string;
backend?: PiSessionBackend;
}
export interface TestServer {
server: PiServer;
backend: PiSessionBackend;
}
/** Create an unstarted PiServer with deterministic defaults for transport conformance tests. */
export function createTestServer(options: TestServerOptions): TestServer {
const backend = options.backend ?? new TestSessionBackend();
return {
server: new PiServer(backend, {
token: options.token ?? TEST_TOKEN,
listeners: options.listeners,
maxFrameLength: options.maxFrameLength,
handshakeTimeoutMs: options.handshakeTimeoutMs,
serverId: options.serverId,
onError: options.onError,
}),
backend,
};
}

View file

@ -4,17 +4,23 @@ import { join } from "node:path";
import { encodeClientMessage, encodeFrame, PROTOCOL_VERSION } from "@earendil-works/pi-protocol";
import { afterEach, describe, expect, test } from "vitest";
import { type PiServer, PiServerError } from "../src/index.ts";
import {
connectUnixTestClient,
Deferred,
type ProtocolTestClient,
TEST_TOKEN,
TestSessionBackend,
} from "../src/testing/index.ts";
import { createUnixServer, type UnixServerOptions } from "../src/transports/unix/index.ts";
import { ConformanceBackend, connectUnix, Deferred, TEST_TOKEN, type WireClient } from "./support.ts";
const servers = new Set<PiServer>();
const clients = new Set<WireClient>();
const clients = new Set<ProtocolTestClient>();
const tempDirectories = new Set<string>();
async function startServer(
backend = new ConformanceBackend(),
backend = new TestSessionBackend(),
overrides: Partial<UnixServerOptions> = {},
): Promise<{ server: PiServer; backend: ConformanceBackend }> {
): Promise<{ server: PiServer; backend: TestSessionBackend }> {
const directory = await mkdtemp(join(tmpdir(), "pis-"));
tempDirectories.add(directory);
const server = createUnixServer(backend, {
@ -27,8 +33,8 @@ async function startServer(
return { server, backend };
}
async function connect(server: PiServer): Promise<WireClient> {
const client = await connectUnix(server.addresses[0]!);
async function connect(server: PiServer): Promise<ProtocolTestClient> {
const client = await connectUnixTestClient(server.addresses[0]!);
clients.add(client);
return client;
}
@ -83,7 +89,7 @@ describe("Unix transport conformance", () => {
});
test("closes connections that do not complete hello before the timeout", async () => {
const { server } = await startServer(new ConformanceBackend(), { handshakeTimeoutMs: 20 });
const { server } = await startServer(new TestSessionBackend(), { handshakeTimeoutMs: 20 });
const client = await connect(server);
await client.waitForClose();
expect(client.messages).toContainEqual(
@ -92,7 +98,7 @@ describe("Unix transport conformance", () => {
});
test("keeps the handshake timeout active until the server hello is sent", async () => {
const backend = new ConformanceBackend();
const backend = new TestSessionBackend();
const delay = backend.delayNextList();
const { server } = await startServer(backend, { handshakeTimeoutMs: 20 });
const client = await connect(server);
@ -116,7 +122,7 @@ describe("Unix transport conformance", () => {
});
await malformed.waitForClose();
const boundedServer = await startServer(new ConformanceBackend(), { maxFrameLength: 128 });
const boundedServer = await startServer(new TestSessionBackend(), { maxFrameLength: 128 });
const oversized = await connect(boundedServer.server);
const frame = new Uint8Array(4 + 129);
frame[3] = 129;
@ -124,7 +130,7 @@ describe("Unix transport conformance", () => {
await oversized.waitForClose();
expect(oversized.messages.some((message) => message.type === "hello")).toBe(false);
const outboundServer = await startServer(new ConformanceBackend(), { maxFrameLength: 128 });
const outboundServer = await startServer(new TestSessionBackend(), { maxFrameLength: 128 });
const outbound = await connect(outboundServer.server);
await outbound.sendMessage({ type: "hello", version: PROTOCOL_VERSION, token: TEST_TOKEN });
await outbound.waitForClose();
@ -132,7 +138,7 @@ describe("Unix transport conformance", () => {
});
test("catches up a handshaking client after a concurrent server change", async () => {
class RacingBackend extends ConformanceBackend {
class RacingBackend extends TestSessionBackend {
readonly entered = new Deferred<void>();
readonly release = new Deferred<void>();
race = false;
@ -171,7 +177,7 @@ describe("Unix transport conformance", () => {
});
test("shares request, event, attachment, and disconnect behavior", async () => {
const backend = new ConformanceBackend();
const backend = new TestSessionBackend();
backend.seed("first");
backend.seed("second");
const { server } = await startServer(backend);
@ -227,7 +233,7 @@ describe("Unix transport conformance", () => {
});
test("disconnects attached clients when a runtime reports a terminal error", async () => {
const backend = new ConformanceBackend();
const backend = new TestSessionBackend();
backend.seed("terminal");
const errors: Error[] = [];
const { server } = await startServer(backend, { onError: (error) => errors.push(error) });
@ -254,7 +260,7 @@ describe("Unix transport conformance", () => {
});
test("does not expose unexpected backend errors to clients", async () => {
class FailingBackend extends ConformanceBackend {
class FailingBackend extends TestSessionBackend {
private listCount = 0;
override async listSessions() {
@ -281,7 +287,7 @@ describe("Unix transport conformance", () => {
});
test("can respond out of request order after the handshake", async () => {
const backend = new ConformanceBackend();
const backend = new TestSessionBackend();
backend.seed("first");
const { server } = await startServer(backend);
const client = await connect(server);
@ -303,7 +309,7 @@ describe("Unix transport conformance", () => {
});
test("gracefully closes connections, sessions, and listener resources", async () => {
const backend = new ConformanceBackend();
const backend = new TestSessionBackend();
backend.seed("first");
const { server } = await startServer(backend);
const socketPath = server.addresses[0];

View file

@ -1,23 +1,7 @@
import { describe, expect, test } from "vitest";
import type { ByteConnectionAcceptor } from "../src/connection.ts";
import type { PiServerListener } from "../src/listener.ts";
import { PiServer } from "../src/server.ts";
import type { PiSessionBackend } from "../src/types.ts";
const backend: PiSessionBackend = {
async listSessions() {
return [];
},
async listModels() {
return [];
},
async createSession() {
throw new Error("not used");
},
async openSession() {
throw new Error("not used");
},
};
import { createTestServer } from "../src/testing/index.ts";
class TestListener implements PiServerListener {
address: string | undefined;
@ -47,7 +31,7 @@ describe("PiServer listener composition", () => {
test("starts and closes every configured listener", async () => {
const first = new TestListener("first");
const second = new TestListener("second");
const server = new PiServer(backend, { token: "secret", listeners: [first, second] });
const { server } = createTestServer({ token: "secret", listeners: [first, second] });
await server.start();
expect(server.addresses).toEqual(["first", "second"]);
@ -64,7 +48,7 @@ describe("PiServer listener composition", () => {
const first = new TestListener("first");
const failure = new Error("listener failed");
const second = new TestListener("second", failure);
const server = new PiServer(backend, { token: "secret", listeners: [first, second] });
const { server } = createTestServer({ token: "secret", listeners: [first, second] });
await expect(server.start()).rejects.toBe(failure);
expect(first.closeCount).toBe(1);

View file

@ -1,768 +1,103 @@
import { once } from "node:events";
import { mkdtemp, rm } from "node:fs/promises";
import { createConnection, type Socket } from "node:net";
import { lstat, mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import {
type ClientMessage,
encodeClientMessage,
type ModelMetadata,
type ModelRef,
PROTOCOL_VERSION,
type ServerMessage,
ServerMessageDecoder,
type SessionPhase,
type SessionSnapshot,
type SessionSummary,
type ThinkingLevel,
type TranscriptProgress,
} from "@earendil-works/pi-protocol";
import { afterEach, describe, expect, test } from "vitest";
import {
type CreateSessionOptions,
type PiServer,
PiServerError,
type PiSessionBackend,
type PiSessionRuntime,
type PiSessionRuntimeEvent,
type PromptInput,
} from "../src/index.ts";
import { createUnixServer, type UnixServerOptions } from "../src/transports/unix/index.ts";
import { ServerMessageDecoder } from "@earendil-works/pi-protocol";
import { afterEach, expect, test, vi } from "vitest";
import type { ByteConnection } from "../src/connection.ts";
import { PiServer } from "../src/index.ts";
import { TestSessionBackend } from "../src/testing/index.ts";
import { createUnixServer } from "../src/transports/unix/index.ts";
const TOKEN = "correct horse battery staple";
const MODEL: ModelMetadata = {
provider: "test",
id: "small",
name: "Test Small",
api: "test-api",
reasoning: true,
input: ["text", "image"],
contextWindow: 16_000,
maxTokens: 2_000,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
supportedThinkingLevels: ["off", "medium", "high"],
authenticated: true,
};
const TOKEN = "transport-smoke-token";
class Deferred<T> {
readonly promise: Promise<T>;
private resolvePromise!: (value: T) => void;
const backend = new TestSessionBackend();
constructor() {
this.promise = new Promise<T>((resolve) => {
this.resolvePromise = resolve;
});
}
let server: PiServer | undefined;
let tempDirectory: string | undefined;
resolve(value: T): void {
this.resolvePromise(value);
}
}
interface StoredSession {
snapshot: SessionSnapshot;
}
class MemoryRuntime implements PiSessionRuntime {
readonly disposed = new Deferred<void>();
disposeCount = 0;
steers: PromptInput[] = [];
private readonly listeners = new Set<(event: PiSessionRuntimeEvent) => void>();
private readonly stored: StoredSession;
private readonly onDispose: () => void;
private pendingPrompt?: { input: PromptInput; done: Deferred<"complete" | "aborted"> };
constructor(stored: StoredSession, onDispose: () => void) {
this.stored = stored;
this.onDispose = onDispose;
}
snapshot(): SessionSnapshot {
return structuredClone(this.stored.snapshot);
}
getPhase(): SessionPhase {
return this.stored.snapshot.phase;
}
async prompt(input: PromptInput): Promise<void> {
if (this.getPhase() !== "idle") {
throw new PiServerError("busy", "A prompt is already running");
}
const done = new Deferred<"complete" | "aborted">();
this.pendingPrompt = { input, done };
this.update({
phase: "turn",
transcript: [
...this.stored.snapshot.transcript,
{
id: `user-${this.stored.snapshot.revision + 1}`,
role: "user",
content: [{ type: "text", text: input.text }],
timestamp: this.stored.snapshot.revision + 1,
},
],
});
const outcome = await done.promise;
const assistantId = `assistant-${this.stored.snapshot.revision + 1}`;
this.update({
phase: "idle",
transcript: [
...this.stored.snapshot.transcript,
{
id: assistantId,
role: "assistant",
content: [{ type: "text", text: outcome === "complete" ? `reply:${input.text}` : "" }],
status: outcome === "complete" ? "complete" : "aborted",
model: this.stored.snapshot.model,
stopReason: outcome === "complete" ? "stop" : "aborted",
timestamp: this.stored.snapshot.revision + 1,
},
],
});
this.pendingPrompt = undefined;
}
async steer(input: PromptInput): Promise<void> {
if (this.getPhase() === "idle") {
throw new PiServerError("busy", "There is no active prompt to steer");
}
this.steers.push(input);
this.update({
queuedSteerCount: this.stored.snapshot.queuedSteerCount + 1,
queuedSteer: [
...this.stored.snapshot.queuedSteer,
{
id: `steer-${this.stored.snapshot.revision + 1}`,
role: "user",
content: [{ type: "text", text: input.text }],
timestamp: this.stored.snapshot.revision + 1,
},
],
});
}
async abort(): Promise<void> {
if (!this.pendingPrompt) {
throw new PiServerError("busy", "There is no active prompt to abort");
}
this.pendingPrompt.done.resolve("aborted");
}
async setModel(model: ModelRef): Promise<void> {
if (this.getPhase() !== "idle") throw new PiServerError("busy", "Session is busy");
this.update({ model });
}
async setThinking(thinkingLevel: ThinkingLevel): Promise<void> {
if (this.getPhase() !== "idle") throw new PiServerError("busy", "Session is busy");
this.update({ thinkingLevel });
}
subscribe(listener: (event: PiSessionRuntimeEvent) => void): () => void {
this.listeners.add(listener);
return () => this.listeners.delete(listener);
}
async dispose(): Promise<void> {
this.disposeCount += 1;
this.onDispose();
this.disposed.resolve();
}
finishPrompt(): void {
if (!this.pendingPrompt) throw new Error("No prompt is pending");
this.pendingPrompt.done.resolve("complete");
}
emitProgress(progress: TranscriptProgress): void {
for (const listener of this.listeners) listener({ type: "progress", progress });
}
emitSnapshot(): void {
for (const listener of this.listeners) listener({ type: "snapshot" });
}
private update(updates: Partial<SessionSnapshot>): void {
this.stored.snapshot = {
...this.stored.snapshot,
...updates,
revision: this.stored.snapshot.revision + 1,
updatedAt: this.stored.snapshot.updatedAt + 1,
};
this.emitSnapshot();
}
}
class MemoryBackend implements PiSessionBackend {
readonly sessions = new Map<string, StoredSession>();
readonly runtimes = new Map<string, MemoryRuntime[]>();
readonly locked = new Set<string>();
lastCreatedId?: string;
async listSessions(): Promise<SessionSummary[]> {
return [...this.sessions.values()].map(({ snapshot }) => ({
id: snapshot.id,
name: snapshot.name,
cwd: snapshot.cwd,
createdAt: snapshot.createdAt,
updatedAt: snapshot.updatedAt,
phase: snapshot.phase,
model: snapshot.model,
thinkingLevel: snapshot.thinkingLevel,
attached: false,
locked: this.locked.has(snapshot.id),
}));
}
async listModels(): Promise<ModelMetadata[]> {
return [MODEL];
}
async createSession(options: CreateSessionOptions): Promise<PiSessionRuntime> {
this.lastCreatedId = options.id;
if (this.sessions.has(options.id)) throw new PiServerError("session_locked", "Session id already exists");
const stored: StoredSession = {
snapshot: {
id: options.id,
name: options.name,
queuedSteer: [],
cwd: options.cwd ?? "/tmp/pi-server-test",
createdAt: 1,
updatedAt: 1,
phase: "idle",
model: options.model ?? { provider: MODEL.provider, id: MODEL.id },
thinkingLevel: options.thinkingLevel ?? "off",
attached: false,
locked: true,
revision: 0,
transcript: [],
queuedSteerCount: 0,
},
};
this.sessions.set(options.id, stored);
return this.acquire(options.id, stored);
}
async openSession(sessionId: string): Promise<PiSessionRuntime> {
const stored = this.sessions.get(sessionId);
if (!stored) throw new PiServerError("not_found", `Unknown session: ${sessionId}`);
if (this.locked.has(sessionId)) throw new PiServerError("session_locked", `Session is locked: ${sessionId}`);
return this.acquire(sessionId, stored);
}
seed(id = "session-1"): StoredSession {
const stored: StoredSession = {
snapshot: {
id,
name: "Seed session",
queuedSteer: [],
cwd: "/tmp/pi-server-test",
createdAt: 1,
updatedAt: 1,
phase: "idle",
model: { provider: MODEL.provider, id: MODEL.id },
thinkingLevel: "off",
attached: false,
locked: false,
revision: 0,
transcript: [],
queuedSteerCount: 0,
},
};
this.sessions.set(id, stored);
return stored;
}
latestRuntime(id: string): MemoryRuntime {
const runtimes = this.runtimes.get(id);
if (!runtimes?.length) throw new Error(`No runtime for ${id}`);
return runtimes.at(-1)!;
}
private acquire(id: string, stored: StoredSession): MemoryRuntime {
this.locked.add(id);
const runtime = new MemoryRuntime(stored, () => this.locked.delete(id));
const runtimes = this.runtimes.get(id) ?? [];
runtimes.push(runtime);
this.runtimes.set(id, runtimes);
return runtime;
}
}
class OrderedSnapshotBackend extends MemoryBackend {
readonly firstStarted = new Deferred<void>();
readonly secondStarted = new Deferred<void>();
readonly firstRelease = new Deferred<void>();
readonly secondRelease = new Deferred<void>();
controlled = false;
startedCount = 0;
override async listModels(): Promise<ModelMetadata[]> {
if (!this.controlled) return super.listModels();
this.startedCount += 1;
if (this.startedCount === 1) {
this.firstStarted.resolve(undefined);
await this.firstRelease.promise;
} else if (this.startedCount === 2) {
this.secondStarted.resolve(undefined);
await this.secondRelease.promise;
}
return [MODEL];
}
}
class Client {
readonly messages: ServerMessage[] = [];
private readonly waiters = new Set<{
predicate: (message: ServerMessage) => boolean;
resolve: (message: ServerMessage) => void;
}>();
private requestNumber = 0;
private readonly decoder = new ServerMessageDecoder();
readonly socket: Socket;
private constructor(socket: Socket) {
this.socket = socket;
socket.on("data", (chunk) => {
for (const message of this.decoder.push(new Uint8Array(chunk.buffer, chunk.byteOffset, chunk.byteLength))) {
this.messages.push(message);
for (const waiter of this.waiters) {
if (!waiter.predicate(message)) continue;
this.waiters.delete(waiter);
waiter.resolve(message);
}
}
});
}
static async connect(path: string): Promise<Client> {
const socket = createConnection(path);
const client = new Client(socket);
await once(socket, "connect");
return client;
}
async hello(token = TOKEN, version: number = PROTOCOL_VERSION): Promise<ServerMessage> {
this.send({ type: "hello", token, version });
return this.next((message) => message.type === "hello" || message.type === "hello_error");
}
async request(request: Record<string, unknown>): Promise<Extract<ServerMessage, { type: "response" }>> {
const id = `request-${++this.requestNumber}`;
this.send({ type: "request", id, request } as ClientMessage);
return (await this.next((message) => message.type === "response" && message.id === id)) as Extract<
ServerMessage,
{ type: "response" }
>;
}
next(predicate: (message: ServerMessage) => boolean): Promise<ServerMessage> {
return this.nextFrom(0, predicate);
}
nextFrom(index: number, predicate: (message: ServerMessage) => boolean): Promise<ServerMessage> {
const existing = this.messages.slice(index).find(predicate);
if (existing) return Promise.resolve(existing);
return new Promise((resolve) => this.waiters.add({ predicate, resolve }));
}
async close(): Promise<void> {
if (this.socket.destroyed) return;
const closed = once(this.socket, "close");
this.socket.destroy();
await closed;
}
private send(message: ClientMessage): void {
this.socket.write(encodeClientMessage(message));
}
}
const servers = new Set<PiServer>();
const clients = new Set<Client>();
const tempDirectories = new Set<string>();
async function startServer(backend = new MemoryBackend(), options: Partial<UnixServerOptions> = {}) {
const directory = await mkdtemp(join(tmpdir(), "pi-server-test-"));
tempDirectories.add(directory);
const server = createUnixServer(backend, {
token: TOKEN,
path: join(directory, "server.sock"),
...options,
});
servers.add(server);
await server.start();
return { server, backend };
}
async function connect(server: PiServer): Promise<Client> {
const client = await Client.connect(server.addresses[0]!);
clients.add(client);
return client;
}
async function attach(client: Client, sessionId: string): Promise<SessionSnapshot> {
const response = await client.request({ command: "attach", sessionId });
if (!response.ok || response.result.command !== "attach") throw new Error("Attach failed");
return response.result.session;
async function makeSocketPath(): Promise<string> {
tempDirectory = await mkdtemp(join(tmpdir(), "pss-"));
return join(tempDirectory, "server.sock");
}
afterEach(async () => {
await Promise.all([...clients].map((client) => client.close()));
clients.clear();
await Promise.all([...servers].map((server) => server.close()));
servers.clear();
await Promise.all([...tempDirectories].map((directory) => rm(directory, { recursive: true, force: true })));
tempDirectories.clear();
await server?.close();
if (tempDirectory) await rm(tempDirectory, { recursive: true, force: true });
server = undefined;
tempDirectory = undefined;
});
describe("PiServer Unix integration", () => {
test("requires a valid first-message auth/version handshake", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server } = await startServer(backend);
test("requires explicit listeners", () => {
expect(() => Reflect.construct(PiServer, [backend, { token: TOKEN }])).toThrow(/listeners/);
});
const badToken = await connect(server);
const authError = await badToken.hello("wrong");
expect(authError).toMatchObject({ type: "hello_error", error: { code: "auth" } });
await once(badToken.socket, "close");
test("rejects Unix socket paths that cannot fit in sockaddr_un", () => {
expect(() => createUnixServer(backend, { token: TOKEN, path: `/tmp/${"x".repeat(512)}` })).toThrow(/too long/);
});
const badVersion = await connect(server);
const versionError = await badVersion.hello(TOKEN, PROTOCOL_VERSION + 1);
expect(versionError).toMatchObject({ type: "hello_error", error: { code: "version" } });
test("rejects an overlong derived private Unix bind path", async () => {
const maxLength = process.platform === "linux" ? 107 : 103;
const suffixLength = Buffer.byteLength("/tmp//s");
const path = `/tmp/${"x".repeat(maxLength - suffixLength)}/s`;
server = createUnixServer(backend, { token: TOKEN, path });
const client = await connect(server);
const hello = await client.hello();
expect(hello).toMatchObject({
type: "hello",
version: PROTOCOL_VERSION,
snapshot: { sessions: [{ id: "session-1" }], models: [MODEL] },
});
});
await expect(server.start()).rejects.toThrow(/private Unix bind path.*too long/);
});
test("times out clients that do not send a hello", async () => {
const { server } = await startServer(new MemoryBackend(), { handshakeTimeoutMs: 10 });
const client = await connect(server);
await once(client.socket, "close");
expect(client.socket.destroyed).toBe(true);
});
test("rejects concurrent start calls without leaking the Unix listener", async () => {
const path = await makeSocketPath();
server = createUnixServer(backend, { token: TOKEN, path });
const starting = server.start();
await expect(server.start()).rejects.toThrow(/starting/);
await starting;
await server.close();
expect(server.addresses[0]).toBeUndefined();
await expect(lstat(path)).rejects.toMatchObject({ code: "ENOENT" });
});
test("serializes server snapshot revisions", async () => {
const backend = new OrderedSnapshotBackend();
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
backend.controlled = true;
const messageIndex = client.messages.length;
test("handshake timeout cleanup does not wait for a blocked output queue", async () => {
class BlockedConnection implements ByteConnection {
closed = false;
finalChunk?: Uint8Array;
const firstCreate = client.request({ command: "create", name: "first" });
await backend.firstStarted.promise;
const secondCreate = client.request({ command: "create", name: "second" });
await Promise.resolve();
expect(backend.startedCount).toBe(1);
backend.firstRelease.resolve(undefined);
await backend.secondStarted.promise;
backend.secondRelease.resolve(undefined);
await Promise.all([firstCreate, secondCreate]);
await client.nextFrom(
messageIndex,
(message) =>
message.type === "event" &&
message.event.type === "server_snapshot" &&
message.event.snapshot.revision === 2,
);
const revisions = client.messages
.slice(messageIndex)
.flatMap((message) =>
message.type === "event" && message.event.type === "server_snapshot"
? [message.event.snapshot.revision]
: [],
);
expect(revisions).toEqual([1, 2]);
});
test("creates server-assigned durable IDs and supports list, attach, and detach", async () => {
const { server, backend } = await startServer();
const client = await connect(server);
await client.hello();
const created = await client.request({ command: "create", cwd: "/work", name: "Created" });
expect(created.ok).toBe(true);
if (!created.ok || created.result.command !== "create") throw new Error("Create failed");
expect(created.result.session).toMatchObject({
id: backend.lastCreatedId,
cwd: "/work",
name: "Created",
attached: true,
locked: true,
});
const listed = await client.request({ command: "list" });
expect(listed).toMatchObject({
ok: true,
result: { command: "list", sessions: [{ id: backend.lastCreatedId, attached: true, locked: true }] },
});
const detached = await client.request({ command: "detach", sessionId: backend.lastCreatedId });
expect(detached).toMatchObject({ ok: true, result: { command: "detach", sessionId: backend.lastCreatedId } });
expect(backend.latestRuntime(backend.lastCreatedId!).disposeCount).toBe(1);
const attached = await attach(client, backend.lastCreatedId!);
expect(attached.id).toBe(backend.lastCreatedId);
expect(backend.runtimes.get(backend.lastCreatedId!)?.length).toBe(2);
});
test("keeps multiple attachments on one connection independent", async () => {
const backend = new MemoryBackend();
backend.seed("first");
backend.seed("second");
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
await attach(client, "first");
await attach(client, "second");
await client.request({ command: "detach", sessionId: "first" });
expect(backend.latestRuntime("first").disposeCount).toBe(1);
expect(backend.latestRuntime("second").disposeCount).toBe(0);
const response = await client.request({
command: "set_thinking",
sessionId: "second",
thinkingLevel: "medium",
});
expect(response).toMatchObject({ ok: true, result: { session: { id: "second", thinkingLevel: "medium" } } });
});
test("broadcasts full snapshots and progress only to clients attached to that session", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server } = await startServer(backend);
const attachedClient = await connect(server);
const unattachedClient = await connect(server);
await attachedClient.hello();
await unattachedClient.hello();
await attach(attachedClient, "session-1");
const runtime = backend.latestRuntime("session-1");
const progress: TranscriptProgress = {
type: "assistant_delta",
messageId: "assistant-1",
contentIndex: 0,
kind: "text",
delta: "hello",
};
runtime.emitProgress(progress);
const progressMessage = await attachedClient.next(
(message) => message.type === "event" && message.event.type === "session_progress",
);
expect(progressMessage).toEqual({
type: "event",
event: { type: "session_progress", sessionId: "session-1", progress },
});
expect(
unattachedClient.messages.some(
(message) => message.type === "event" && message.event.type === "session_progress",
),
).toBe(false);
const messageCount = attachedClient.messages.length;
runtime.emitSnapshot();
const snapshotMessage = await attachedClient.nextFrom(
messageCount,
(message) =>
message.type === "event" &&
message.event.type === "session_snapshot" &&
message.event.snapshot.revision === runtime.snapshot().revision,
);
expect(snapshotMessage).toMatchObject({
type: "event",
event: { type: "session_snapshot", snapshot: { id: "session-1", attached: true, locked: true } },
});
expect(
unattachedClient.messages.some(
(message) => message.type === "event" && message.event.type === "session_snapshot",
),
).toBe(false);
});
test("allows every attached client to control a singleton live runtime", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server } = await startServer(backend);
const first = await connect(server);
const second = await connect(server);
await first.hello();
await second.hello();
await attach(first, "session-1");
const secondList = await second.request({ command: "list" });
expect(secondList).toMatchObject({
ok: true,
result: { sessions: [{ id: "session-1", attached: false, locked: true }] },
});
await attach(second, "session-1");
expect(backend.runtimes.get("session-1")?.length).toBe(1);
const modelResponse = await second.request({
command: "set_model",
sessionId: "session-1",
model: { provider: "test", id: "large" },
});
expect(modelResponse).toMatchObject({ ok: true, result: { session: { model: { id: "large" } } } });
await first.next(
(message) =>
message.type === "event" &&
message.event.type === "session_snapshot" &&
message.event.snapshot.model.id === "large",
);
const thinkingResponse = await first.request({
command: "set_thinking",
sessionId: "session-1",
thinkingLevel: "high",
});
expect(thinkingResponse).toMatchObject({ ok: true, result: { session: { thinkingLevel: "high" } } });
});
test("does not queue prompts and processes steer and abort while a prompt response is pending", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
await attach(client, "session-1");
const prompt = client.request({ command: "prompt", sessionId: "session-1", text: "first" });
await client.next(
(message) =>
message.type === "event" &&
message.event.type === "session_snapshot" &&
message.event.snapshot.phase === "turn",
);
const busy = await client.request({ command: "prompt", sessionId: "session-1", text: "second" });
expect(busy).toMatchObject({ ok: false, error: { code: "busy" } });
const steer = await client.request({ command: "steer", sessionId: "session-1", text: "adjust" });
expect(steer).toMatchObject({ ok: true, result: { command: "steer" } });
expect(backend.latestRuntime("session-1").steers).toEqual([{ text: "adjust" }]);
const abort = await client.request({ command: "abort", sessionId: "session-1" });
expect(abort).toMatchObject({ ok: true, result: { command: "abort" } });
expect(await prompt).toMatchObject({ ok: true, result: { command: "prompt", session: { phase: "idle" } } });
});
test("returns operation attachment state relative to the requesting connection", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server } = await startServer(backend);
const first = await connect(server);
const second = await connect(server);
await first.hello();
await second.hello();
await attach(first, "session-1");
await attach(second, "session-1");
const prompt = first.request({ command: "prompt", sessionId: "session-1", text: "hello" });
await first.next(
(message) =>
message.type === "event" &&
message.event.type === "session_snapshot" &&
message.event.snapshot.phase === "turn",
);
await first.request({ command: "detach", sessionId: "session-1" });
backend.latestRuntime("session-1").finishPrompt();
expect(await prompt).toMatchObject({
ok: true,
result: { command: "prompt", session: { id: "session-1", attached: false } },
});
});
test("keeps busy work alive after disconnect and disposes when it next becomes idle", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
await attach(client, "session-1");
void client.request({ command: "prompt", sessionId: "session-1", text: "survive" });
await client.next(
(message) =>
message.type === "event" &&
message.event.type === "session_snapshot" &&
message.event.snapshot.phase === "turn",
);
const runtime = backend.latestRuntime("session-1");
await client.close();
expect(runtime.disposeCount).toBe(0);
runtime.finishPrompt();
await runtime.disposed.promise;
expect(runtime.disposeCount).toBe(1);
const reconnect = await connect(server);
await reconnect.hello();
const snapshot = await attach(reconnect, "session-1");
expect(snapshot.transcript).toHaveLength(2);
expect(snapshot.transcript[1]).toMatchObject({ role: "assistant", content: [{ text: "reply:survive" }] });
});
test("restores persisted sessions lazily after a server restart", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server: firstServer } = await startServer(backend);
const firstClient = await connect(firstServer);
await firstClient.hello();
await attach(firstClient, "session-1");
await firstClient.request({ command: "set_thinking", sessionId: "session-1", thinkingLevel: "high" });
await firstClient.close();
await firstServer.close();
const { server: secondServer } = await startServer(backend);
expect(backend.runtimes.get("session-1")).toHaveLength(1);
const secondClient = await connect(secondServer);
await secondClient.hello();
const restored = await attach(secondClient, "session-1");
expect(restored.thinkingLevel).toBe("high");
expect(backend.runtimes.get("session-1")).toHaveLength(2);
});
test("rejects and disposes a backend runtime with the wrong server-assigned ID", async () => {
class WrongIdBackend extends MemoryBackend {
override async createSession(options: CreateSessionOptions): Promise<PiSessionRuntime> {
return super.createSession({ ...options, id: "wrong-id" });
}
send(): Promise<void> {
return new Promise(() => {});
}
const backend = new WrongIdBackend();
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
const response = await client.request({ command: "create" });
expect(response).toMatchObject({ ok: false, error: { code: "invalid_request" } });
expect(backend.latestRuntime("wrong-id").disposeCount).toBe(1);
});
test("maps backend lock errors and rejects control from unattached clients", async () => {
const backend = new MemoryBackend();
backend.seed("locked");
backend.locked.add("locked");
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
const locked = await client.request({ command: "attach", sessionId: "locked" });
expect(locked).toMatchObject({ ok: false, error: { code: "session_locked" } });
const unattached = await client.request({ command: "abort", sessionId: "locked" });
expect(unattached).toMatchObject({ ok: false, error: { code: "invalid_request" } });
close(finalChunk?: Uint8Array): void {
this.finalChunk = finalChunk;
this.closed = true;
}
}
const core = new PiServer(backend, {
listeners: [],
token: TOKEN,
maxFrameLength: 1024,
handshakeTimeoutMs: 10,
});
const connection = new BlockedConnection();
core.accept(connection);
test("gracefully closes clients and releases all live runtimes", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
await attach(client, "session-1");
const runtime = backend.latestRuntime("session-1");
const clientClosed = once(client.socket, "close");
await server.close();
await clientClosed;
expect(runtime.disposeCount).toBe(1);
expect(server.addresses).toEqual([]);
await server.close();
});
await vi.waitFor(() => expect(connection.closed).toBe(true));
expect(connection.finalChunk).toBeInstanceOf(Uint8Array);
const messages = new ServerMessageDecoder().push(connection.finalChunk!);
expect(messages).toMatchObject([{ type: "hello_error", error: { code: "invalid_request" } }]);
await core.close();
});
test("rejects timeout values above Node's maximum timer delay", () => {
const unix = { path: "/tmp/pi-server-timeout-test.sock" };
expect(() =>
createUnixServer(backend, { token: TOKEN, path: unix.path, handshakeTimeoutMs: 2_147_483_648 }),
).toThrow(/handshakeTimeoutMs/);
expect(() =>
createUnixServer(backend, { token: TOKEN, path: unix.path, gracefulCloseTimeoutMs: 2_147_483_648 }),
).toThrow(/gracefulCloseTimeoutMs/);
});
test("rejects pending-byte limits smaller than one maximum frame", async () => {
const path = await makeSocketPath();
expect(() => createUnixServer(backend, { token: TOKEN, path, maxFrameLength: 128, maxPendingBytes: 131 })).toThrow(
/maxPendingBytes/,
);
});

View file

@ -0,0 +1,387 @@
import { mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { ModelMetadata, SessionSnapshot, TranscriptProgress } from "@earendil-works/pi-protocol";
import { afterEach, describe, expect, test } from "vitest";
import type { CreateSessionOptions, PiServer, PiSessionRuntime } from "../src/index.ts";
import {
connectUnixTestClient,
Deferred,
type ProtocolTestClient,
TEST_MODEL,
TEST_TOKEN,
TestSessionBackend,
} from "../src/testing/index.ts";
import { createUnixServer, type UnixServerOptions } from "../src/transports/unix/index.ts";
const TOKEN = TEST_TOKEN;
const MODEL = TEST_MODEL;
type Client = ProtocolTestClient;
class MemoryBackend extends TestSessionBackend {}
class OrderedSnapshotBackend extends MemoryBackend {
readonly firstStarted = new Deferred<void>();
readonly secondStarted = new Deferred<void>();
readonly firstRelease = new Deferred<void>();
readonly secondRelease = new Deferred<void>();
controlled = false;
startedCount = 0;
override async listModels(): Promise<ModelMetadata[]> {
if (!this.controlled) return super.listModels();
this.startedCount += 1;
if (this.startedCount === 1) {
this.firstStarted.resolve(undefined);
await this.firstRelease.promise;
} else if (this.startedCount === 2) {
this.secondStarted.resolve(undefined);
await this.secondRelease.promise;
}
return [MODEL];
}
}
const servers = new Set<PiServer>();
const clients = new Set<Client>();
const tempDirectories = new Set<string>();
async function startServer(backend = new MemoryBackend(), options: Partial<UnixServerOptions> = {}) {
const directory = await mkdtemp(join(tmpdir(), "pst-"));
tempDirectories.add(directory);
const server = createUnixServer(backend, {
token: TOKEN,
path: join(directory, "server.sock"),
...options,
});
servers.add(server);
await server.start();
return { server, backend };
}
async function connect(server: PiServer): Promise<Client> {
const client = await connectUnixTestClient(server.addresses[0]!);
clients.add(client);
return client;
}
async function attach(client: Client, sessionId: string): Promise<SessionSnapshot> {
const response = await client.request({ command: "attach", sessionId });
if (!response.ok || response.result.command !== "attach") throw new Error("Attach failed");
return response.result.session;
}
afterEach(async () => {
await Promise.all([...clients].map((client) => client.close()));
clients.clear();
await Promise.all([...servers].map((server) => server.close()));
servers.clear();
await Promise.all([...tempDirectories].map((directory) => rm(directory, { recursive: true, force: true })));
tempDirectories.clear();
});
describe("PiServer Unix integration", () => {
test("serializes server snapshot revisions", async () => {
const backend = new OrderedSnapshotBackend();
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
backend.controlled = true;
const messageIndex = client.messages.length;
const firstCreate = client.request({ command: "create", name: "first" });
await backend.firstStarted.promise;
const secondCreate = client.request({ command: "create", name: "second" });
await Promise.resolve();
expect(backend.startedCount).toBe(1);
backend.firstRelease.resolve(undefined);
await backend.secondStarted.promise;
backend.secondRelease.resolve(undefined);
await Promise.all([firstCreate, secondCreate]);
await client.nextFrom(
messageIndex,
(message) =>
message.type === "event" &&
message.event.type === "server_snapshot" &&
message.event.snapshot.revision === 2,
);
const revisions = client.messages
.slice(messageIndex)
.flatMap((message) =>
message.type === "event" && message.event.type === "server_snapshot"
? [message.event.snapshot.revision]
: [],
);
expect(revisions).toEqual([1, 2]);
});
test("creates server-assigned durable IDs and supports list, attach, and detach", async () => {
const { server, backend } = await startServer();
const client = await connect(server);
await client.hello();
const created = await client.request({ command: "create", cwd: "/work", name: "Created" });
expect(created.ok).toBe(true);
if (!created.ok || created.result.command !== "create") throw new Error("Create failed");
const createdId = created.result.session.id;
expect(created.result.session).toMatchObject({
id: backend.lastCreatedId,
cwd: "/work",
name: "Created",
attached: true,
locked: true,
});
const listed = await client.request({ command: "list" });
expect(listed).toMatchObject({
ok: true,
result: { command: "list", sessions: [{ id: backend.lastCreatedId, attached: true, locked: true }] },
});
const detached = await client.request({ command: "detach", sessionId: createdId });
expect(detached).toMatchObject({ ok: true, result: { command: "detach", sessionId: createdId } });
expect(backend.latestRuntime(createdId).disposeCount).toBe(1);
const attached = await attach(client, createdId);
expect(attached.id).toBe(backend.lastCreatedId);
expect(backend.runtimes.get(createdId)?.length).toBe(2);
});
test("keeps multiple attachments on one connection independent", async () => {
const backend = new MemoryBackend();
backend.seed("first");
backend.seed("second");
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
await attach(client, "first");
await attach(client, "second");
await client.request({ command: "detach", sessionId: "first" });
expect(backend.latestRuntime("first").disposeCount).toBe(1);
expect(backend.latestRuntime("second").disposeCount).toBe(0);
const response = await client.request({
command: "set_thinking",
sessionId: "second",
thinkingLevel: "medium",
});
expect(response).toMatchObject({ ok: true, result: { session: { id: "second", thinkingLevel: "medium" } } });
});
test("broadcasts full snapshots and progress only to clients attached to that session", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server } = await startServer(backend);
const attachedClient = await connect(server);
const unattachedClient = await connect(server);
await attachedClient.hello();
await unattachedClient.hello();
await attach(attachedClient, "session-1");
const runtime = backend.latestRuntime("session-1");
const progress: TranscriptProgress = {
type: "assistant_delta",
messageId: "assistant-1",
contentIndex: 0,
kind: "text",
delta: "hello",
};
runtime.emitProgress(progress);
const progressMessage = await attachedClient.next(
(message) => message.type === "event" && message.event.type === "session_progress",
);
expect(progressMessage).toEqual({
type: "event",
event: { type: "session_progress", sessionId: "session-1", progress },
});
expect(
unattachedClient.messages.some(
(message) => message.type === "event" && message.event.type === "session_progress",
),
).toBe(false);
const messageCount = attachedClient.messages.length;
runtime.emitSnapshot();
const snapshotMessage = await attachedClient.nextFrom(
messageCount,
(message) =>
message.type === "event" &&
message.event.type === "session_snapshot" &&
message.event.snapshot.revision === runtime.snapshot().revision,
);
expect(snapshotMessage).toMatchObject({
type: "event",
event: { type: "session_snapshot", snapshot: { id: "session-1", attached: true, locked: true } },
});
expect(
unattachedClient.messages.some(
(message) => message.type === "event" && message.event.type === "session_snapshot",
),
).toBe(false);
});
test("allows every attached client to control a singleton live runtime", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server } = await startServer(backend);
const first = await connect(server);
const second = await connect(server);
await first.hello();
await second.hello();
await attach(first, "session-1");
const secondList = await second.request({ command: "list" });
expect(secondList).toMatchObject({
ok: true,
result: { sessions: [{ id: "session-1", attached: false, locked: true }] },
});
await attach(second, "session-1");
expect(backend.runtimes.get("session-1")?.length).toBe(1);
const modelResponse = await second.request({
command: "set_model",
sessionId: "session-1",
model: { provider: "test", id: "large" },
});
expect(modelResponse).toMatchObject({ ok: true, result: { session: { model: { id: "large" } } } });
await first.next(
(message) =>
message.type === "event" &&
message.event.type === "session_snapshot" &&
message.event.snapshot.model.id === "large",
);
const thinkingResponse = await first.request({
command: "set_thinking",
sessionId: "session-1",
thinkingLevel: "high",
});
expect(thinkingResponse).toMatchObject({ ok: true, result: { session: { thinkingLevel: "high" } } });
});
test("does not queue prompts and processes steer and abort while a prompt response is pending", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
await attach(client, "session-1");
const prompt = client.request({ command: "prompt", sessionId: "session-1", text: "first" });
await client.next(
(message) =>
message.type === "event" &&
message.event.type === "session_snapshot" &&
message.event.snapshot.phase === "turn",
);
const busy = await client.request({ command: "prompt", sessionId: "session-1", text: "second" });
expect(busy).toMatchObject({ ok: false, error: { code: "busy" } });
const steer = await client.request({ command: "steer", sessionId: "session-1", text: "adjust" });
expect(steer).toMatchObject({ ok: true, result: { command: "steer" } });
expect(backend.latestRuntime("session-1").steers).toEqual([{ text: "adjust" }]);
const abort = await client.request({ command: "abort", sessionId: "session-1" });
expect(abort).toMatchObject({ ok: true, result: { command: "abort" } });
expect(await prompt).toMatchObject({ ok: true, result: { command: "prompt", session: { phase: "idle" } } });
});
test("returns operation attachment state relative to the requesting connection", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server } = await startServer(backend);
const first = await connect(server);
const second = await connect(server);
await first.hello();
await second.hello();
await attach(first, "session-1");
await attach(second, "session-1");
const prompt = first.request({ command: "prompt", sessionId: "session-1", text: "hello" });
await first.next(
(message) =>
message.type === "event" &&
message.event.type === "session_snapshot" &&
message.event.snapshot.phase === "turn",
);
await first.request({ command: "detach", sessionId: "session-1" });
backend.latestRuntime("session-1").finishPrompt();
expect(await prompt).toMatchObject({
ok: true,
result: { command: "prompt", session: { id: "session-1", attached: false } },
});
});
test("keeps busy work alive after disconnect and disposes when it next becomes idle", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
await attach(client, "session-1");
const prompt = client.request({ command: "prompt", sessionId: "session-1", text: "survive" });
await client.next(
(message) =>
message.type === "event" &&
message.event.type === "session_snapshot" &&
message.event.snapshot.phase === "turn",
);
const runtime = backend.latestRuntime("session-1");
await client.close();
await expect(prompt).rejects.toThrow(Error);
expect(runtime.disposeCount).toBe(0);
runtime.finishPrompt();
await runtime.disposed.promise;
expect(runtime.disposeCount).toBe(1);
const reconnect = await connect(server);
await reconnect.hello();
const snapshot = await attach(reconnect, "session-1");
expect(snapshot.transcript).toHaveLength(2);
expect(snapshot.transcript[1]).toMatchObject({ role: "assistant", content: [{ text: "reply:survive" }] });
});
test("restores persisted sessions lazily after a server restart", async () => {
const backend = new MemoryBackend();
backend.seed();
const { server: firstServer } = await startServer(backend);
const firstClient = await connect(firstServer);
await firstClient.hello();
await attach(firstClient, "session-1");
await firstClient.request({ command: "set_thinking", sessionId: "session-1", thinkingLevel: "high" });
await firstClient.close();
await firstServer.close();
const { server: secondServer } = await startServer(backend);
expect(backend.runtimes.get("session-1")).toHaveLength(1);
const secondClient = await connect(secondServer);
await secondClient.hello();
const restored = await attach(secondClient, "session-1");
expect(restored.thinkingLevel).toBe("high");
expect(backend.runtimes.get("session-1")).toHaveLength(2);
});
test("rejects and disposes a backend runtime with the wrong server-assigned ID", async () => {
class WrongIdBackend extends MemoryBackend {
override async createSession(options: CreateSessionOptions): Promise<PiSessionRuntime> {
return super.createSession({ ...options, id: "wrong-id" });
}
}
const backend = new WrongIdBackend();
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
const response = await client.request({ command: "create" });
expect(response).toMatchObject({ ok: false, error: { code: "invalid_request" } });
expect(backend.latestRuntime("wrong-id").disposeCount).toBe(1);
});
test("maps backend lock errors and rejects control from unattached clients", async () => {
const backend = new MemoryBackend();
backend.seed("locked");
backend.locked.add("locked");
const { server } = await startServer(backend);
const client = await connect(server);
await client.hello();
const locked = await client.request({ command: "attach", sessionId: "locked" });
expect(locked).toMatchObject({ ok: false, error: { code: "session_locked" } });
const unattached = await client.request({ command: "abort", sessionId: "locked" });
expect(unattached).toMatchObject({ ok: false, error: { code: "invalid_request" } });
});
});

View file

@ -1,158 +0,0 @@
import { once } from "node:events";
import { lstat, mkdtemp, rm } from "node:fs/promises";
import { createConnection, type Socket } from "node:net";
import { tmpdir } from "node:os";
import { join } from "node:path";
import {
encodeClientMessage,
PROTOCOL_VERSION,
type ServerMessage,
ServerMessageDecoder,
} from "@earendil-works/pi-protocol";
import { afterEach, expect, test, vi } from "vitest";
import type { ByteConnection } from "../src/connection.ts";
import { PiServer, type PiSessionBackend } from "../src/index.ts";
import { createUnixServer } from "../src/transports/unix/index.ts";
const TOKEN = "transport-smoke-token";
const backend: PiSessionBackend = {
async listSessions() {
return [];
},
async listModels() {
return [];
},
async createSession() {
throw new Error("not used");
},
async openSession() {
throw new Error("not used");
},
};
let server: PiServer | undefined;
let socket: Socket | undefined;
let tempDirectory: string | undefined;
async function makeSocketPath(): Promise<string> {
tempDirectory = await mkdtemp(join(tmpdir(), "pi-server-smoke-"));
return join(tempDirectory, "server.sock");
}
afterEach(async () => {
socket?.destroy();
await server?.close();
if (tempDirectory) await rm(tempDirectory, { recursive: true, force: true });
server = undefined;
socket = undefined;
tempDirectory = undefined;
});
test("requires explicit listeners", () => {
expect(() => Reflect.construct(PiServer, [backend, { token: TOKEN }])).toThrow(/listeners/);
});
test("rejects Unix socket paths that cannot fit in sockaddr_un", () => {
expect(() => createUnixServer(backend, { token: TOKEN, path: `/tmp/${"x".repeat(512)}` })).toThrow(/too long/);
});
test("rejects an overlong derived private Unix bind path", async () => {
const maxLength = process.platform === "linux" ? 107 : 103;
const suffixLength = Buffer.byteLength("/tmp//s");
const path = `/tmp/${"x".repeat(maxLength - suffixLength)}/s`;
server = createUnixServer(backend, { token: TOKEN, path });
await expect(server.start()).rejects.toThrow(/private Unix bind path.*too long/);
});
test("Unix socket accepts a fragmented framed-CBOR hello", async () => {
const path = await makeSocketPath();
server = createUnixServer(backend, { token: TOKEN, path });
await server.start();
expect(server.addresses[0]).toBe(path);
socket = createConnection(path);
await once(socket, "connect");
const response = nextServerMessage(socket);
const hello = encodeClientMessage({ type: "hello", version: PROTOCOL_VERSION, token: TOKEN });
socket.write(hello.subarray(0, 2));
socket.write(hello.subarray(2));
expect(await response).toMatchObject({ type: "hello", version: PROTOCOL_VERSION });
});
test("rejects concurrent start calls without leaking the Unix listener", async () => {
const path = await makeSocketPath();
server = createUnixServer(backend, { token: TOKEN, path });
const starting = server.start();
await expect(server.start()).rejects.toThrow(/starting/);
await starting;
await server.close();
expect(server.addresses[0]).toBeUndefined();
await expect(lstat(path)).rejects.toMatchObject({ code: "ENOENT" });
});
test("handshake timeout cleanup does not wait for a blocked output queue", async () => {
class BlockedConnection implements ByteConnection {
closed = false;
finalChunk?: Uint8Array;
send(): Promise<void> {
return new Promise(() => {});
}
close(finalChunk?: Uint8Array): void {
this.finalChunk = finalChunk;
this.closed = true;
}
}
const core = new PiServer(backend, {
listeners: [],
token: TOKEN,
maxFrameLength: 1024,
handshakeTimeoutMs: 10,
});
const connection = new BlockedConnection();
core.accept(connection);
await vi.waitFor(() => expect(connection.closed).toBe(true));
expect(connection.finalChunk).toBeInstanceOf(Uint8Array);
const messages = new ServerMessageDecoder().push(connection.finalChunk!);
expect(messages).toMatchObject([{ type: "hello_error", error: { code: "invalid_request" } }]);
await core.close();
});
test("rejects timeout values above Node's maximum timer delay", () => {
const unix = { path: "/tmp/pi-server-timeout-test.sock" };
expect(() =>
createUnixServer(backend, { token: TOKEN, path: unix.path, handshakeTimeoutMs: 2_147_483_648 }),
).toThrow(/handshakeTimeoutMs/);
expect(() =>
createUnixServer(backend, { token: TOKEN, path: unix.path, gracefulCloseTimeoutMs: 2_147_483_648 }),
).toThrow(/gracefulCloseTimeoutMs/);
});
test("rejects pending-byte limits smaller than one maximum frame", async () => {
const path = await makeSocketPath();
expect(() => createUnixServer(backend, { token: TOKEN, path, maxFrameLength: 128, maxPendingBytes: 131 })).toThrow(
/maxPendingBytes/,
);
});
function nextServerMessage(target: Socket): Promise<ServerMessage> {
const decoder = new ServerMessageDecoder();
return new Promise((resolve, reject) => {
const onData = (chunk: Buffer): void => {
try {
const messages = decoder.push(new Uint8Array(chunk.buffer, chunk.byteOffset, chunk.byteLength));
if (messages.length === 0) return;
target.off("data", onData);
resolve(messages[0]!);
} catch (error) {
reject(error);
}
};
target.on("data", onData);
target.once("error", reject);
});
}

View file

@ -5,11 +5,16 @@ import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, describe, expect, test } from "vitest";
import type { PiServer } from "../src/index.ts";
import {
connectUnixTestClient,
type ProtocolTestClient,
TEST_TOKEN,
TestSessionBackend,
} from "../src/testing/index.ts";
import { createUnixServer } from "../src/transports/unix/index.ts";
import { ConformanceBackend, connectUnix, TEST_TOKEN, type WireClient } from "./support.ts";
const servers = new Set<PiServer>();
const clients = new Set<WireClient>();
const clients = new Set<ProtocolTestClient>();
const children = new Set<ChildProcess>();
const tempDirectories = new Set<string>();
@ -20,7 +25,7 @@ async function makeSocketPath(nested = false): Promise<string> {
}
function makeServer(path: string): PiServer {
const server = createUnixServer(new ConformanceBackend(), { token: TEST_TOKEN, path });
const server = createUnixServer(new TestSessionBackend(), { token: TEST_TOKEN, path });
servers.add(server);
return server;
}
@ -55,7 +60,7 @@ describe("Unix listener filesystem lifecycle", () => {
ino: firstIdentity.ino,
});
const client = await connectUnix(path);
const client = await connectUnixTestClient(path);
clients.add(client);
expect(await client.hello()).toMatchObject({ type: "hello" });
});
@ -110,7 +115,7 @@ describe("Unix listener filesystem lifecycle", () => {
await server.start();
const liveIdentity = await lstat(path);
expect(liveIdentity.isSocket()).toBe(true);
const client = await connectUnix(path);
const client = await connectUnixTestClient(path);
clients.add(client);
expect(await client.hello()).toMatchObject({ type: "hello" });
});