433 lines
16 KiB
TypeScript
433 lines
16 KiB
TypeScript
import assert from "node:assert/strict";
|
|
import test from "node:test";
|
|
import type {
|
|
ExtensionAPI,
|
|
ExtensionCommandContext,
|
|
ExtensionContext,
|
|
} from "@earendil-works/pi-coding-agent";
|
|
import {
|
|
createDictateExtension,
|
|
MAX_FRAME_BYTES,
|
|
MAX_TRANSCRIPT_BYTES,
|
|
PROTOCOL_VERSION,
|
|
type DictateSocket,
|
|
} from "./index.ts";
|
|
|
|
type ShortcutHandler = (ctx: ExtensionContext) => Promise<void> | void;
|
|
type CommandHandler = (
|
|
args: string,
|
|
ctx: ExtensionCommandContext,
|
|
) => Promise<void> | void;
|
|
type EventHandler = (event: unknown, ctx: ExtensionContext) => Promise<void> | void;
|
|
type Frame = Record<string, unknown>;
|
|
|
|
class FakeSocket implements DictateSocket {
|
|
readonly writes: string[] = [];
|
|
endCount = 0;
|
|
destroyCount = 0;
|
|
private connectListener: (() => void) | undefined;
|
|
private dataListener: ((chunk: Uint8Array) => void) | undefined;
|
|
private errorListener: ((error: Error) => void) | undefined;
|
|
private closeListener: (() => void) | undefined;
|
|
|
|
onConnect(listener: () => void) {
|
|
this.connectListener = listener;
|
|
}
|
|
onData(listener: (chunk: Uint8Array) => void) {
|
|
this.dataListener = listener;
|
|
}
|
|
onError(listener: (error: Error) => void) {
|
|
this.errorListener = listener;
|
|
}
|
|
onClose(listener: () => void) {
|
|
this.closeListener = listener;
|
|
}
|
|
write(data: string) {
|
|
this.writes.push(data);
|
|
}
|
|
end() {
|
|
this.endCount++;
|
|
}
|
|
destroy() {
|
|
this.destroyCount++;
|
|
}
|
|
connect() {
|
|
this.connectListener?.();
|
|
}
|
|
data(data: string | Uint8Array) {
|
|
this.dataListener?.(
|
|
typeof data === "string" ? Buffer.from(data, "utf8") : data,
|
|
);
|
|
}
|
|
frame(frame: Frame) {
|
|
this.data(`${JSON.stringify(frame)}\n`);
|
|
}
|
|
error(message: string) {
|
|
this.errorListener?.(new Error(message));
|
|
}
|
|
close() {
|
|
this.closeListener?.();
|
|
}
|
|
}
|
|
|
|
function createHarness() {
|
|
const shortcuts = new Map<string, ShortcutHandler>();
|
|
const commands = new Map<string, CommandHandler>();
|
|
const events = new Map<string, EventHandler>();
|
|
const sockets: FakeSocket[] = [];
|
|
const socketPaths: string[] = [];
|
|
const statuses: Array<string | undefined> = [];
|
|
const notifications: Array<{ message: string; type?: string }> = [];
|
|
const pasted: string[] = [];
|
|
const context = {
|
|
ui: {
|
|
setStatus: (_key: string, text: string | undefined) => statuses.push(text),
|
|
notify: (message: string, type?: string) =>
|
|
notifications.push({ message, type }),
|
|
pasteToEditor: (text: string) => pasted.push(text),
|
|
},
|
|
} as unknown as ExtensionContext;
|
|
const api = {
|
|
registerShortcut(key: string, options: { handler: ShortcutHandler }) {
|
|
shortcuts.set(key, options.handler);
|
|
},
|
|
registerCommand(name: string, options: { handler: CommandHandler }) {
|
|
commands.set(name, options.handler);
|
|
},
|
|
on(name: string, handler: EventHandler) {
|
|
events.set(name, handler);
|
|
},
|
|
} as unknown as ExtensionAPI;
|
|
|
|
createDictateExtension({
|
|
socketPath: "/test/pi-dictate.sock",
|
|
socketFactory: (path) => {
|
|
socketPaths.push(path);
|
|
const socket = new FakeSocket();
|
|
sockets.push(socket);
|
|
return socket;
|
|
},
|
|
})(api);
|
|
|
|
const shortcut = (key: string) => {
|
|
const handler = shortcuts.get(key);
|
|
assert.ok(handler, `shortcut ${key} was registered`);
|
|
return handler;
|
|
};
|
|
const command = (name: string) => {
|
|
const handler = commands.get(name);
|
|
assert.ok(handler, `command ${name} was registered`);
|
|
return handler;
|
|
};
|
|
const event = (name: string) => {
|
|
const handler = events.get(name);
|
|
assert.ok(handler, `event ${name} was registered`);
|
|
return handler;
|
|
};
|
|
const sentFrames = (socket: FakeSocket) =>
|
|
socket.writes.map((line) => JSON.parse(line) as Frame);
|
|
|
|
return {
|
|
shortcuts,
|
|
commands,
|
|
events,
|
|
sockets,
|
|
socketPaths,
|
|
statuses,
|
|
notifications,
|
|
pasted,
|
|
context,
|
|
shortcut,
|
|
command,
|
|
event,
|
|
sentFrames,
|
|
};
|
|
}
|
|
|
|
type Harness = ReturnType<typeof createHarness>;
|
|
|
|
async function connectAndNegotiate(harness: Harness) {
|
|
await harness.shortcut("alt+m")(harness.context);
|
|
const socket = harness.sockets.at(-1);
|
|
assert.ok(socket);
|
|
socket.connect();
|
|
const hello = harness.sentFrames(socket)[0];
|
|
assert.equal(hello?.type, "hello");
|
|
assert.equal(hello?.version, PROTOCOL_VERSION);
|
|
assert.equal(typeof hello?.requestId, "string");
|
|
assert.equal(harness.sentFrames(socket).length, 1, "start waits for hello");
|
|
socket.frame({
|
|
version: PROTOCOL_VERSION,
|
|
type: "pong",
|
|
requestId: hello?.requestId,
|
|
});
|
|
const start = harness.sentFrames(socket)[1];
|
|
assert.equal(start?.type, "start");
|
|
assert.equal(start?.version, PROTOCOL_VERSION);
|
|
assert.equal(typeof start?.requestId, "string");
|
|
assert.equal(typeof start?.generationId, "string");
|
|
return { socket, generationId: start?.generationId as string };
|
|
}
|
|
|
|
function serviceFrame(
|
|
socket: FakeSocket,
|
|
generationId: string,
|
|
type: string,
|
|
fields: Frame = {},
|
|
) {
|
|
socket.frame({
|
|
version: PROTOCOL_VERSION,
|
|
type,
|
|
requestId: `event-${type}`,
|
|
generationId,
|
|
...fields,
|
|
});
|
|
}
|
|
|
|
test("registers main-editor shortcuts, cancel command, and shutdown cleanup", () => {
|
|
const harness = createHarness();
|
|
assert.deepEqual([...harness.shortcuts.keys()], ["alt+m", "alt+n"]);
|
|
assert.deepEqual([...harness.commands.keys()], ["dictate-cancel"]);
|
|
assert.ok(harness.events.has("session_shutdown"));
|
|
});
|
|
|
|
test("connection errors return to idle and the next start uses a fresh socket", async () => {
|
|
const harness = createHarness();
|
|
await harness.shortcut("alt+m")(harness.context);
|
|
harness.sockets[0]?.error("socket unavailable");
|
|
|
|
assert.match(harness.notifications[0]?.message ?? "", /connection failed.*socket unavailable/i);
|
|
assert.equal(harness.notifications[0]?.type, "error");
|
|
assert.equal(harness.statuses.at(-1), undefined);
|
|
assert.equal(harness.sockets[0]?.destroyCount, 1);
|
|
|
|
await harness.shortcut("alt+m")(harness.context);
|
|
assert.equal(harness.sockets.length, 2);
|
|
assert.deepEqual(harness.socketPaths, ["/test/pi-dictate.sock", "/test/pi-dictate.sock"]);
|
|
});
|
|
|
|
test("rejects a mismatched protocol version before sending start", async () => {
|
|
const harness = createHarness();
|
|
await harness.shortcut("alt+m")(harness.context);
|
|
const socket = harness.sockets[0];
|
|
assert.ok(socket);
|
|
socket.connect();
|
|
const hello = harness.sentFrames(socket)[0];
|
|
socket.frame({ version: 2, type: "pong", requestId: hello?.requestId });
|
|
|
|
assert.equal(harness.sentFrames(socket).length, 1);
|
|
assert.match(harness.notifications[0]?.message ?? "", /version mismatch/i);
|
|
assert.equal(socket.destroyCount, 1);
|
|
});
|
|
|
|
test("rejects a hello response with the wrong request ID", async () => {
|
|
const harness = createHarness();
|
|
await harness.shortcut("alt+m")(harness.context);
|
|
const socket = harness.sockets[0];
|
|
assert.ok(socket);
|
|
socket.connect();
|
|
socket.frame({
|
|
version: PROTOCOL_VERSION,
|
|
type: "pong",
|
|
requestId: "not-the-hello-request",
|
|
});
|
|
|
|
assert.equal(harness.sentFrames(socket).length, 1);
|
|
assert.match(harness.notifications[0]?.message ?? "", /unexpected hello response/i);
|
|
assert.equal(socket.destroyCount, 1);
|
|
});
|
|
|
|
test("rejects generation events before hello negotiation", async () => {
|
|
const harness = createHarness();
|
|
await harness.shortcut("alt+m")(harness.context);
|
|
const socket = harness.sockets[0];
|
|
assert.ok(socket);
|
|
socket.connect();
|
|
serviceFrame(socket, "generation-1", "state", { state: "listening" });
|
|
|
|
assert.match(harness.notifications[0]?.message ?? "", /before hello negotiation/i);
|
|
assert.equal(socket.destroyCount, 1);
|
|
assert.deepEqual(harness.pasted, []);
|
|
});
|
|
|
|
test("renders authoritative service states and sends a bounded stop command", async () => {
|
|
const harness = createHarness();
|
|
const { socket, generationId } = await connectAndNegotiate(harness);
|
|
|
|
serviceFrame(socket, generationId, "state", { state: "loading" });
|
|
serviceFrame(socket, generationId, "state", { state: "listening" });
|
|
await harness.shortcut("alt+m")(harness.context);
|
|
const stop = harness.sentFrames(socket).at(-1);
|
|
assert.equal(stop?.type, "stop");
|
|
assert.equal(stop?.generationId, generationId);
|
|
assert.ok(Buffer.byteLength(String(stop?.requestId)) <= 128);
|
|
serviceFrame(socket, generationId, "state", { state: "stopping" });
|
|
serviceFrame(socket, generationId, "state", { state: "transcribing" });
|
|
|
|
assert.deepEqual(harness.statuses, [
|
|
"loading model...",
|
|
"listening...",
|
|
"stopping...",
|
|
"transcribing...",
|
|
]);
|
|
});
|
|
|
|
test("ignores stale generations but rejects malformed and oversized frames", async () => {
|
|
const staleHarness = createHarness();
|
|
const { socket: staleSocket, generationId } = await connectAndNegotiate(staleHarness);
|
|
serviceFrame(staleSocket, "old-generation", "state", { state: "listening" });
|
|
assert.deepEqual(staleHarness.statuses, []);
|
|
serviceFrame(staleSocket, generationId, "state", { state: "listening" });
|
|
assert.deepEqual(staleHarness.statuses, ["listening..."]);
|
|
|
|
const malformedHarness = createHarness();
|
|
const malformed = await connectAndNegotiate(malformedHarness);
|
|
malformed.socket.data("{not json}\n");
|
|
assert.match(malformedHarness.notifications[0]?.message ?? "", /malformed JSON/i);
|
|
assert.equal(malformed.socket.destroyCount, 1);
|
|
|
|
const nonObjectHarness = createHarness();
|
|
const nonObject = await connectAndNegotiate(nonObjectHarness);
|
|
nonObject.socket.data("null\n");
|
|
assert.match(nonObjectHarness.notifications[0]?.message ?? "", /frame must be an object/i);
|
|
assert.equal(nonObject.socket.destroyCount, 1);
|
|
|
|
const oversizedHarness = createHarness();
|
|
const oversized = await connectAndNegotiate(oversizedHarness);
|
|
oversized.socket.data("x".repeat(MAX_FRAME_BYTES + 1));
|
|
assert.match(oversizedHarness.notifications[0]?.message ?? "", /frame is too large/i);
|
|
assert.equal(oversized.socket.destroyCount, 1);
|
|
});
|
|
|
|
test("rejects invalid UTF-8 instead of altering transcript text", async () => {
|
|
const harness = createHarness();
|
|
const { socket, generationId } = await connectAndNegotiate(harness);
|
|
serviceFrame(socket, generationId, "state", { state: "listening" });
|
|
await harness.shortcut("alt+m")(harness.context);
|
|
serviceFrame(socket, generationId, "state", { state: "transcribing" });
|
|
const prefix = Buffer.from(
|
|
JSON.stringify({
|
|
version: PROTOCOL_VERSION,
|
|
type: "result",
|
|
requestId: "event-result",
|
|
generationId,
|
|
transcript: "",
|
|
}).replace('"transcript":""}', '"transcript":"'),
|
|
);
|
|
const suffix = Buffer.from('"}\n');
|
|
socket.data(Buffer.concat([prefix, Buffer.from([0xc3, 0x28]), suffix]));
|
|
|
|
assert.match(harness.notifications[0]?.message ?? "", /invalid UTF-8/i);
|
|
assert.equal(socket.destroyCount, 1);
|
|
assert.deepEqual(harness.pasted, []);
|
|
});
|
|
|
|
test("service idle closes the accepted connection and forces a fresh handshake", async () => {
|
|
const harness = createHarness();
|
|
const { socket, generationId } = await connectAndNegotiate(harness);
|
|
serviceFrame(socket, generationId, "state", { state: "loading" });
|
|
serviceFrame(socket, generationId, "state", { state: "idle" });
|
|
assert.equal(socket.endCount, 1);
|
|
assert.equal(harness.statuses.at(-1), undefined);
|
|
|
|
await harness.shortcut("alt+m")(harness.context);
|
|
assert.equal(harness.sockets.length, 2);
|
|
harness.sockets[1]?.connect();
|
|
assert.equal(harness.sentFrames(harness.sockets[1] as FakeSocket)[0]?.type, "hello");
|
|
});
|
|
|
|
test("an empty result leaves the editor unchanged and reports no speech", async () => {
|
|
const harness = createHarness();
|
|
const { socket, generationId } = await connectAndNegotiate(harness);
|
|
serviceFrame(socket, generationId, "state", { state: "listening" });
|
|
await harness.shortcut("alt+m")(harness.context);
|
|
serviceFrame(socket, generationId, "state", { state: "transcribing" });
|
|
serviceFrame(socket, generationId, "result", { transcript: " " });
|
|
|
|
assert.deepEqual(harness.pasted, []);
|
|
assert.match(harness.notifications[0]?.message ?? "", /no speech recognized/i);
|
|
assert.equal(socket.endCount, 1);
|
|
});
|
|
|
|
test("pastes exactly one non-empty current result through the public editor API", async () => {
|
|
const harness = createHarness();
|
|
const { socket, generationId } = await connectAndNegotiate(harness);
|
|
serviceFrame(socket, generationId, "state", { state: "listening" });
|
|
await harness.shortcut("alt+m")(harness.context);
|
|
serviceFrame(socket, generationId, "state", { state: "stopping" });
|
|
serviceFrame(socket, generationId, "state", { state: "transcribing" });
|
|
serviceFrame(socket, generationId, "result", { transcript: "typed words" });
|
|
serviceFrame(socket, generationId, "result", { transcript: "duplicate" });
|
|
|
|
assert.deepEqual(harness.pasted, ["typed words"]);
|
|
assert.equal(harness.statuses.at(-1), undefined);
|
|
|
|
const oversizedTranscriptHarness = createHarness();
|
|
const active = await connectAndNegotiate(oversizedTranscriptHarness);
|
|
serviceFrame(active.socket, active.generationId, "state", { state: "listening" });
|
|
await oversizedTranscriptHarness.shortcut("alt+m")(oversizedTranscriptHarness.context);
|
|
serviceFrame(active.socket, active.generationId, "state", { state: "transcribing" });
|
|
serviceFrame(active.socket, active.generationId, "result", {
|
|
transcript: "t".repeat(MAX_TRANSCRIPT_BYTES + 1),
|
|
});
|
|
assert.deepEqual(oversizedTranscriptHarness.pasted, []);
|
|
assert.match(oversizedTranscriptHarness.notifications[0]?.message ?? "", /transcript is too large/i);
|
|
});
|
|
|
|
test("disconnects and service errors return safely to idle without delivery", async () => {
|
|
const disconnected = createHarness();
|
|
const first = await connectAndNegotiate(disconnected);
|
|
first.socket.close();
|
|
assert.match(disconnected.notifications[0]?.message ?? "", /disconnected/i);
|
|
assert.deepEqual(disconnected.pasted, []);
|
|
await disconnected.shortcut("alt+m")(disconnected.context);
|
|
assert.equal(disconnected.sockets.length, 2);
|
|
|
|
const failed = createHarness();
|
|
const second = await connectAndNegotiate(failed);
|
|
serviceFrame(second.socket, second.generationId, "error", {
|
|
code: "model_failed",
|
|
detail: "model unavailable",
|
|
});
|
|
assert.match(failed.notifications[0]?.message ?? "", /model_failed.*model unavailable/i);
|
|
assert.deepEqual(failed.pasted, []);
|
|
});
|
|
|
|
test("cancellation during loading tombstones later events and the slash command also cancels", async () => {
|
|
const harness = createHarness();
|
|
const { socket, generationId } = await connectAndNegotiate(harness);
|
|
serviceFrame(socket, generationId, "state", { state: "loading" });
|
|
await harness.shortcut("alt+n")(harness.context);
|
|
|
|
assert.equal(harness.sentFrames(socket).at(-1)?.type, "cancel");
|
|
assert.equal(socket.endCount, 1);
|
|
serviceFrame(socket, generationId, "state", { state: "listening" });
|
|
serviceFrame(socket, generationId, "result", { transcript: "must not appear" });
|
|
assert.deepEqual(harness.pasted, []);
|
|
assert.equal(harness.statuses.at(-1), undefined);
|
|
|
|
const next = await connectAndNegotiate(harness);
|
|
serviceFrame(next.socket, next.generationId, "state", { state: "listening" });
|
|
await harness.command("dictate-cancel")(
|
|
"",
|
|
harness.context as ExtensionCommandContext,
|
|
);
|
|
assert.equal(harness.sentFrames(next.socket).at(-1)?.type, "cancel");
|
|
assert.equal(next.socket.endCount, 1);
|
|
});
|
|
|
|
test("shutdown best-effort cancels, closes, clears status, and remains delivery-neutral", async () => {
|
|
const harness = createHarness();
|
|
const { socket, generationId } = await connectAndNegotiate(harness);
|
|
serviceFrame(socket, generationId, "state", { state: "listening" });
|
|
const shutdown = harness.event("session_shutdown");
|
|
await shutdown({}, harness.context);
|
|
await shutdown({}, harness.context);
|
|
|
|
assert.equal(harness.sentFrames(socket).at(-1)?.type, "cancel");
|
|
assert.equal(socket.endCount, 1);
|
|
assert.equal(harness.statuses.at(-1), undefined);
|
|
serviceFrame(socket, generationId, "result", { transcript: "never paste cleanup" });
|
|
assert.deepEqual(harness.pasted, []);
|
|
});
|