mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-06 23:44:36 +00:00
1ceb9ea488
Make connection, protocol, storage, and lock failures fail fast. Align the remote TUI's spacing, activity indicator, and exit behavior with legacy mode.
327 lines
13 KiB
TypeScript
327 lines
13 KiB
TypeScript
import { lstat, mkdtemp, rm } from "node:fs/promises";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
import { encodeClientMessage, encodeFrame, PROTOCOL_VERSION } from "@earendil-works/pi-protocol";
|
|
import { afterEach, describe, expect, test } from "vitest";
|
|
import { PiServer, PiServerError, type PiServerOptions } from "../src/index.ts";
|
|
import { ConformanceBackend, connectUnix, Deferred, TEST_TOKEN, type WireClient } from "./support.ts";
|
|
|
|
const servers = new Set<PiServer>();
|
|
const clients = new Set<WireClient>();
|
|
const tempDirectories = new Set<string>();
|
|
|
|
async function startServer(
|
|
backend = new ConformanceBackend(),
|
|
overrides: Partial<PiServerOptions> = {},
|
|
): Promise<{ server: PiServer; backend: ConformanceBackend }> {
|
|
const directory = await mkdtemp(join(tmpdir(), "pi-server-conformance-"));
|
|
tempDirectories.add(directory);
|
|
const server = new PiServer(backend, {
|
|
token: TEST_TOKEN,
|
|
unix: { path: join(directory, "server.sock") },
|
|
...overrides,
|
|
});
|
|
servers.add(server);
|
|
await server.start();
|
|
return { server, backend };
|
|
}
|
|
|
|
async function connect(server: PiServer): Promise<WireClient> {
|
|
const client = await connectUnix(server.unixSocketPath!);
|
|
clients.add(client);
|
|
return client;
|
|
}
|
|
|
|
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("Unix transport conformance", () => {
|
|
test("accepts a transport-fragmented framed-CBOR hello", async () => {
|
|
const { server } = await startServer();
|
|
const client = await connect(server);
|
|
const response = client.next((message) => message.type === "hello");
|
|
await client.sendFragmentedMessage({ type: "hello", version: PROTOCOL_VERSION, token: TEST_TOKEN }, 2);
|
|
expect(await response).toMatchObject({ type: "hello", version: PROTOCOL_VERSION });
|
|
});
|
|
|
|
test("enforces authentication, version, and exactly one first-message hello", async () => {
|
|
const { server } = await startServer();
|
|
|
|
const badToken = await connect(server);
|
|
expect(await badToken.hello("wrong-token")).toMatchObject({
|
|
type: "hello_error",
|
|
error: { code: "auth" },
|
|
});
|
|
await badToken.waitForClose();
|
|
|
|
const badVersion = await connect(server);
|
|
expect(await badVersion.hello(TEST_TOKEN, PROTOCOL_VERSION + 1)).toMatchObject({
|
|
type: "hello_error",
|
|
error: { code: "version" },
|
|
});
|
|
await badVersion.waitForClose();
|
|
|
|
const requestFirst = await connect(server);
|
|
const firstError = requestFirst.next((message) => message.type === "hello_error");
|
|
await requestFirst.sendMessage({ type: "request", id: "too-early", request: { command: "list" } });
|
|
expect(await firstError).toMatchObject({ type: "hello_error", error: { code: "invalid_request" } });
|
|
await requestFirst.waitForClose();
|
|
|
|
const duplicate = await connect(server);
|
|
expect(await duplicate.hello()).toMatchObject({ type: "hello" });
|
|
const duplicateError = duplicate.next((message) => message.type === "hello_error");
|
|
await duplicate.sendMessage({ type: "hello", version: PROTOCOL_VERSION, token: TEST_TOKEN });
|
|
expect(await duplicateError).toMatchObject({ type: "hello_error", error: { code: "invalid_request" } });
|
|
await duplicate.waitForClose();
|
|
});
|
|
|
|
test("closes connections that do not complete hello before the timeout", async () => {
|
|
const { server } = await startServer(new ConformanceBackend(), { handshakeTimeoutMs: 20 });
|
|
const client = await connect(server);
|
|
await client.waitForClose();
|
|
expect(client.messages).toContainEqual(
|
|
expect.objectContaining({ type: "hello_error", error: expect.objectContaining({ code: "invalid_request" }) }),
|
|
);
|
|
});
|
|
|
|
test("keeps the handshake timeout active until the server hello is sent", async () => {
|
|
const backend = new ConformanceBackend();
|
|
const delay = backend.delayNextList();
|
|
const { server } = await startServer(backend, { handshakeTimeoutMs: 20 });
|
|
const client = await connect(server);
|
|
await client.sendMessage({ type: "hello", version: PROTOCOL_VERSION, token: TEST_TOKEN });
|
|
await delay.entered.promise;
|
|
await client.waitForClose();
|
|
delay.release.resolve(undefined);
|
|
expect(client.messages).toContainEqual(
|
|
expect.objectContaining({ type: "hello_error", error: expect.objectContaining({ code: "invalid_request" }) }),
|
|
);
|
|
});
|
|
|
|
test("bounds and closes malformed or oversized frames", async () => {
|
|
const malformedServer = await startServer();
|
|
const malformed = await connect(malformedServer.server);
|
|
const malformedError = malformed.next((message) => message.type === "hello_error");
|
|
await malformed.sendBytes(encodeFrame(Uint8Array.of(0xff)));
|
|
expect(await malformedError).toMatchObject({
|
|
type: "hello_error",
|
|
error: { code: "invalid_request" },
|
|
});
|
|
await malformed.waitForClose();
|
|
|
|
const boundedServer = await startServer(new ConformanceBackend(), { maxFrameLength: 128 });
|
|
const oversized = await connect(boundedServer.server);
|
|
const frame = new Uint8Array(4 + 129);
|
|
frame[3] = 129;
|
|
await oversized.sendBytes(frame);
|
|
await oversized.waitForClose();
|
|
expect(oversized.messages.some((message) => message.type === "hello")).toBe(false);
|
|
|
|
const outboundServer = await startServer(new ConformanceBackend(), { maxFrameLength: 128 });
|
|
const outbound = await connect(outboundServer.server);
|
|
await outbound.sendMessage({ type: "hello", version: PROTOCOL_VERSION, token: TEST_TOKEN });
|
|
await outbound.waitForClose();
|
|
expect(outbound.messages).toEqual([]);
|
|
});
|
|
|
|
test("catches up a handshaking client after a concurrent server change", async () => {
|
|
class RacingBackend extends ConformanceBackend {
|
|
readonly entered = new Deferred<void>();
|
|
readonly release = new Deferred<void>();
|
|
race = false;
|
|
|
|
override async listSessions() {
|
|
const sessions = await super.listSessions();
|
|
if (!this.race) return sessions;
|
|
this.entered.resolve(undefined);
|
|
await this.release.promise;
|
|
return sessions;
|
|
}
|
|
}
|
|
const backend = new RacingBackend();
|
|
backend.seed("shared");
|
|
const { server } = await startServer(backend);
|
|
const controller = await connect(server);
|
|
await controller.hello();
|
|
backend.race = true;
|
|
const joining = await connect(server);
|
|
const hello = joining.hello();
|
|
await backend.entered.promise;
|
|
await controller.request({ command: "attach", sessionId: "shared" });
|
|
backend.release.resolve(undefined);
|
|
const handshake = await hello;
|
|
if (handshake.type !== "hello") throw new Error("Expected server hello");
|
|
const catchup = await joining.next(
|
|
(message) =>
|
|
message.type === "event" &&
|
|
message.event.type === "server_snapshot" &&
|
|
message.event.snapshot.revision > handshake.snapshot.revision,
|
|
);
|
|
expect(catchup).toMatchObject({
|
|
type: "event",
|
|
event: { type: "server_snapshot", snapshot: { sessions: [{ id: "shared", locked: true }] } },
|
|
});
|
|
});
|
|
|
|
test("shares request, event, attachment, and disconnect behavior", async () => {
|
|
const backend = new ConformanceBackend();
|
|
backend.seed("first");
|
|
backend.seed("second");
|
|
const { server } = await startServer(backend);
|
|
const client = await connect(server);
|
|
expect(await client.hello()).toMatchObject({
|
|
type: "hello",
|
|
snapshot: { sessions: [{ id: "first" }, { id: "second" }] },
|
|
});
|
|
|
|
const listed = await client.request({ command: "list" });
|
|
expect(listed).toMatchObject({
|
|
ok: true,
|
|
result: { command: "list", sessions: [{ id: "first" }, { id: "second" }] },
|
|
});
|
|
expect(await client.request({ command: "attach", sessionId: "first" })).toMatchObject({
|
|
ok: true,
|
|
result: { command: "attach", session: { id: "first", attached: true } },
|
|
});
|
|
expect(await client.request({ command: "attach", sessionId: "second" })).toMatchObject({
|
|
ok: true,
|
|
result: { command: "attach", session: { id: "second", attached: true } },
|
|
});
|
|
|
|
const progress = {
|
|
type: "assistant_delta" as const,
|
|
messageId: "assistant-1",
|
|
contentIndex: 0,
|
|
kind: "text" as const,
|
|
delta: "hello",
|
|
};
|
|
const progressEvent = client.next(
|
|
(message) => message.type === "event" && message.event.type === "session_progress",
|
|
);
|
|
backend.latestRuntime("first").emitProgress(progress);
|
|
expect(await progressEvent).toEqual({
|
|
type: "event",
|
|
event: { type: "session_progress", sessionId: "first", progress },
|
|
});
|
|
|
|
expect(await client.request({ command: "detach", sessionId: "first" })).toMatchObject({
|
|
ok: true,
|
|
result: { command: "detach", sessionId: "first" },
|
|
});
|
|
expect(backend.latestRuntime("first").disposeCount).toBe(1);
|
|
expect(
|
|
await client.request({ command: "set_thinking", sessionId: "second", thinkingLevel: "high" }),
|
|
).toMatchObject({ ok: true, result: { session: { id: "second", thinkingLevel: "high" } } });
|
|
|
|
const secondRuntime = backend.latestRuntime("second");
|
|
await client.close();
|
|
await secondRuntime.disposed.promise;
|
|
expect(secondRuntime.disposeCount).toBe(1);
|
|
});
|
|
|
|
test("disconnects attached clients when a runtime reports a terminal error", async () => {
|
|
const backend = new ConformanceBackend();
|
|
backend.seed("terminal");
|
|
const errors: Error[] = [];
|
|
const { server } = await startServer(backend, { onError: (error) => errors.push(error) });
|
|
const client = await connect(server);
|
|
await client.hello();
|
|
await client.request({ command: "attach", sessionId: "terminal" });
|
|
const runtime = backend.latestRuntime("terminal");
|
|
|
|
runtime.emitError(new PiServerError("session_locked", "lock ownership lost"));
|
|
await client.waitForClose();
|
|
await runtime.disposed.promise;
|
|
expect(errors).toContainEqual(expect.objectContaining({ code: "session_locked" }));
|
|
});
|
|
|
|
test("does not expose unexpected backend errors to clients", async () => {
|
|
class FailingBackend extends ConformanceBackend {
|
|
private listCount = 0;
|
|
|
|
override async listSessions() {
|
|
this.listCount += 1;
|
|
if (this.listCount > 1) throw new Error("private backend detail");
|
|
return super.listSessions();
|
|
}
|
|
}
|
|
const errors: Error[] = [];
|
|
const backend = new FailingBackend();
|
|
const { server } = await startServer(backend, {
|
|
onError: (error) => {
|
|
errors.push(error);
|
|
throw new Error("observer failure");
|
|
},
|
|
});
|
|
const client = await connect(server);
|
|
await client.hello();
|
|
expect(await client.request({ command: "list" })).toMatchObject({
|
|
ok: false,
|
|
error: { code: "invalid_request", message: "Internal server error" },
|
|
});
|
|
expect(errors).toContainEqual(expect.objectContaining({ message: "private backend detail" }));
|
|
});
|
|
|
|
test("can respond out of request order after the handshake", async () => {
|
|
const backend = new ConformanceBackend();
|
|
backend.seed("first");
|
|
const { server } = await startServer(backend);
|
|
const client = await connect(server);
|
|
await client.hello();
|
|
|
|
const delay = backend.delayNextList();
|
|
const slow = client.request({ command: "list" }, "slow");
|
|
await delay.entered.promise;
|
|
const fast = client.request({ command: "attach", sessionId: "first" }, "fast");
|
|
expect(await fast).toMatchObject({ ok: true, id: "fast", result: { command: "attach" } });
|
|
expect(client.messages.some((message) => message.type === "response" && message.id === "slow")).toBe(false);
|
|
|
|
delay.release.resolve(undefined);
|
|
expect(await slow).toMatchObject({ ok: true, id: "slow", result: { command: "list" } });
|
|
const responseIds = client.messages
|
|
.filter((message) => message.type === "response" && (message.id === "slow" || message.id === "fast"))
|
|
.map((message) => (message.type === "response" ? message.id : ""));
|
|
expect(responseIds).toEqual(["fast", "slow"]);
|
|
});
|
|
|
|
test("gracefully closes connections, sessions, and listener resources", async () => {
|
|
const backend = new ConformanceBackend();
|
|
backend.seed("first");
|
|
const { server } = await startServer(backend);
|
|
const socketPath = server.unixSocketPath;
|
|
const client = await connect(server);
|
|
await client.hello();
|
|
await client.request({ command: "attach", sessionId: "first" });
|
|
const runtime = backend.latestRuntime("first");
|
|
const clientClosed = client.waitForClose();
|
|
|
|
await server.close();
|
|
await clientClosed;
|
|
expect(runtime.disposeCount).toBe(1);
|
|
expect(server.unixSocketPath).toBeUndefined();
|
|
if (socketPath) await expect(lstat(socketPath)).rejects.toMatchObject({ code: "ENOENT" });
|
|
await server.close();
|
|
});
|
|
});
|
|
|
|
test("Unix socket decodes multiple framed requests from one raw chunk", async () => {
|
|
const { server } = await startServer();
|
|
const client = await connect(server);
|
|
await client.hello();
|
|
const first = encodeClientMessage({ type: "request", id: "first", request: { command: "list" } });
|
|
const second = encodeClientMessage({ type: "request", id: "second", request: { command: "list" } });
|
|
const combined = new Uint8Array(first.byteLength + second.byteLength);
|
|
combined.set(first);
|
|
combined.set(second, first.byteLength);
|
|
const firstResponse = client.next((message) => message.type === "response" && message.id === "first");
|
|
const secondResponse = client.next((message) => message.type === "response" && message.id === "second");
|
|
await client.sendBytes(combined);
|
|
expect(await firstResponse).toMatchObject({ type: "response", id: "first", ok: true });
|
|
expect(await secondResponse).toMatchObject({ type: "response", id: "second", ok: true });
|
|
});
|