420 lines
12 KiB
TypeScript
420 lines
12 KiB
TypeScript
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
|
|
import { createConnection } from "node:net";
|
|
|
|
export const DICTATE_SOCKET_PATH = "/run/user/1000/pi-dictate.sock";
|
|
export const PROTOCOL_VERSION = 1;
|
|
export const MAX_FRAME_BYTES = 64 * 1024;
|
|
export const MAX_TRANSCRIPT_BYTES = 32 * 1024;
|
|
export const MAX_ID_BYTES = 128;
|
|
|
|
const STATUS_KEY = "dictate";
|
|
const MAX_ERROR_DETAIL_BYTES = 2 * 1024;
|
|
const UTF8_DECODER = new TextDecoder("utf-8", { fatal: true });
|
|
|
|
type DictateState =
|
|
| "idle"
|
|
| "connecting"
|
|
| "recording"
|
|
| "stopping"
|
|
| "transcribing";
|
|
type ServiceState =
|
|
| "loading"
|
|
| "listening"
|
|
| "stopping"
|
|
| "transcribing"
|
|
| "idle";
|
|
|
|
export interface DictateSocket {
|
|
onConnect(listener: () => void): void;
|
|
onData(listener: (chunk: Uint8Array) => void): void;
|
|
onError(listener: (error: Error) => void): void;
|
|
onClose(listener: () => void): void;
|
|
write(data: string): void;
|
|
end(): void;
|
|
destroy(): void;
|
|
}
|
|
|
|
export interface DictateExtensionOptions {
|
|
socketFactory?: (path: string) => DictateSocket;
|
|
socketPath?: string;
|
|
}
|
|
|
|
function nodeSocket(path: string): DictateSocket {
|
|
const socket = createConnection(path);
|
|
return {
|
|
onConnect: (listener) => socket.on("connect", listener),
|
|
onData: (listener) => socket.on("data", listener),
|
|
onError: (listener) => socket.on("error", listener),
|
|
onClose: (listener) => socket.on("close", listener),
|
|
write: (data) => {
|
|
socket.write(data);
|
|
},
|
|
end: () => {
|
|
socket.end();
|
|
},
|
|
destroy: () => {
|
|
socket.destroy();
|
|
},
|
|
};
|
|
}
|
|
|
|
function isRecord(value: unknown): value is Record<string, unknown> {
|
|
return typeof value === "object" && value !== null && !Array.isArray(value);
|
|
}
|
|
|
|
function utf8Bytes(value: string): number {
|
|
return Buffer.byteLength(value, "utf8");
|
|
}
|
|
|
|
function isBoundedId(value: unknown): value is string {
|
|
return (
|
|
typeof value === "string" &&
|
|
value.length > 0 &&
|
|
utf8Bytes(value) <= MAX_ID_BYTES
|
|
);
|
|
}
|
|
|
|
function boundedText(value: string, maxBytes: number): string {
|
|
if (utf8Bytes(value) <= maxBytes) return value;
|
|
let text = value;
|
|
while (text.length > 0 && utf8Bytes(`${text}…`) > maxBytes) {
|
|
text = text.slice(0, -1);
|
|
}
|
|
return `${text}…`;
|
|
}
|
|
|
|
export function createDictateExtension(options: DictateExtensionOptions = {}) {
|
|
const socketFactory = options.socketFactory ?? nodeSocket;
|
|
const socketPath = options.socketPath ?? DICTATE_SOCKET_PATH;
|
|
|
|
return function registerDictate(pi: ExtensionAPI) {
|
|
let state: DictateState = "idle";
|
|
let socket: DictateSocket | null = null;
|
|
let activeContext: ExtensionContext | null = null;
|
|
let activeGenerationId: string | null = null;
|
|
let helloRequestId: string | null = null;
|
|
let negotiated = false;
|
|
let receiveBuffer = Buffer.alloc(0);
|
|
let nextId = 0;
|
|
let resultDelivered = false;
|
|
|
|
const newId = (kind: string) => `${kind}-${++nextId}`;
|
|
|
|
const setStatus = (status: string | undefined) => {
|
|
activeContext?.ui.setStatus(STATUS_KEY, status);
|
|
};
|
|
|
|
const closeSocket = (current: DictateSocket, force: boolean) => {
|
|
try {
|
|
if (force) current.destroy();
|
|
else current.end();
|
|
} catch {
|
|
// Cleanup is best effort and must never deliver text.
|
|
}
|
|
};
|
|
|
|
const reset = (forceClose = false) => {
|
|
const current = socket;
|
|
socket = null;
|
|
activeGenerationId = null;
|
|
helloRequestId = null;
|
|
negotiated = false;
|
|
receiveBuffer = Buffer.alloc(0);
|
|
resultDelivered = false;
|
|
state = "idle";
|
|
setStatus(undefined);
|
|
activeContext = null;
|
|
if (current) closeSocket(current, forceClose);
|
|
};
|
|
|
|
const notifyError = (message: string) => {
|
|
activeContext?.ui.notify(
|
|
boundedText(message, MAX_ERROR_DETAIL_BYTES),
|
|
"error",
|
|
);
|
|
};
|
|
|
|
const fail = (message: string) => {
|
|
notifyError(message);
|
|
reset(true);
|
|
};
|
|
|
|
const send = (frame: Record<string, string | number>): boolean => {
|
|
if (!socket) return false;
|
|
const encoded = `${JSON.stringify(frame)}\n`;
|
|
if (utf8Bytes(encoded) - 1 > MAX_FRAME_BYTES) {
|
|
fail("Dictation protocol error: outgoing frame is too large.");
|
|
return false;
|
|
}
|
|
try {
|
|
socket.write(encoded);
|
|
return true;
|
|
} catch (error) {
|
|
const detail = error instanceof Error ? error.message : "write failed";
|
|
fail(`Dictation service write failed: ${detail}`);
|
|
return false;
|
|
}
|
|
};
|
|
|
|
const sendGenerationCommand = (
|
|
type: "start" | "stop" | "cancel",
|
|
): boolean => {
|
|
if (!activeGenerationId || !negotiated) return false;
|
|
return send({
|
|
version: PROTOCOL_VERSION,
|
|
type,
|
|
requestId: newId(type),
|
|
generationId: activeGenerationId,
|
|
});
|
|
};
|
|
|
|
const protocolError = (detail: string) => {
|
|
fail(`Dictation protocol error: ${detail}`);
|
|
};
|
|
|
|
const acceptServiceState = (serviceState: ServiceState): boolean => {
|
|
if (serviceState === "idle") {
|
|
reset();
|
|
return true;
|
|
}
|
|
if (serviceState === "loading" && state === "connecting") {
|
|
setStatus("loading model...");
|
|
return true;
|
|
}
|
|
if (serviceState === "listening" && state === "connecting") {
|
|
state = "recording";
|
|
setStatus("listening...");
|
|
return true;
|
|
}
|
|
if (serviceState === "stopping" && state === "stopping") {
|
|
setStatus("stopping...");
|
|
return true;
|
|
}
|
|
if (serviceState === "transcribing" && state === "stopping") {
|
|
state = "transcribing";
|
|
setStatus("transcribing...");
|
|
return true;
|
|
}
|
|
if (
|
|
(serviceState === "loading" && state === "connecting") ||
|
|
(serviceState === "listening" && state === "recording") ||
|
|
(serviceState === "stopping" && state === "stopping") ||
|
|
(serviceState === "transcribing" && state === "transcribing")
|
|
) {
|
|
return true;
|
|
}
|
|
return false;
|
|
};
|
|
|
|
const handleFrame = (value: unknown, current: DictateSocket) => {
|
|
if (socket !== current) return;
|
|
if (!isRecord(value)) {
|
|
protocolError("frame must be an object.");
|
|
return;
|
|
}
|
|
if (value.version !== PROTOCOL_VERSION) {
|
|
fail("Dictation service protocol version mismatch.");
|
|
return;
|
|
}
|
|
if (typeof value.type !== "string" || !isBoundedId(value.requestId)) {
|
|
protocolError("frame type or request ID is invalid.");
|
|
return;
|
|
}
|
|
|
|
if (value.type === "pong") {
|
|
if (
|
|
negotiated ||
|
|
state !== "connecting" ||
|
|
value.requestId !== helloRequestId
|
|
) {
|
|
protocolError("unexpected hello response.");
|
|
return;
|
|
}
|
|
negotiated = true;
|
|
sendGenerationCommand("start");
|
|
return;
|
|
}
|
|
if (!negotiated) {
|
|
protocolError("generation event received before hello negotiation.");
|
|
return;
|
|
}
|
|
|
|
const generationId = value.generationId;
|
|
if (!isBoundedId(generationId)) {
|
|
protocolError("generation ID is invalid.");
|
|
return;
|
|
}
|
|
if (generationId !== activeGenerationId) return;
|
|
|
|
if (value.type === "state") {
|
|
const serviceState = value.state;
|
|
if (
|
|
typeof serviceState !== "string" ||
|
|
![
|
|
"loading",
|
|
"listening",
|
|
"stopping",
|
|
"transcribing",
|
|
"idle",
|
|
].includes(serviceState)
|
|
) {
|
|
protocolError("service state is invalid.");
|
|
return;
|
|
}
|
|
if (!acceptServiceState(serviceState as ServiceState)) {
|
|
protocolError("service state transition is invalid.");
|
|
}
|
|
return;
|
|
}
|
|
|
|
if (value.type === "result") {
|
|
if (state !== "transcribing" || typeof value.transcript !== "string") {
|
|
protocolError("result is invalid for the current state.");
|
|
return;
|
|
}
|
|
if (utf8Bytes(value.transcript) > MAX_TRANSCRIPT_BYTES) {
|
|
protocolError("transcript is too large.");
|
|
return;
|
|
}
|
|
if (resultDelivered) return;
|
|
resultDelivered = true;
|
|
const context = activeContext;
|
|
const transcript = value.transcript;
|
|
reset();
|
|
if (transcript.trim().length === 0) {
|
|
context?.ui.notify("No speech recognized.", "info");
|
|
} else {
|
|
context?.ui.pasteToEditor(transcript);
|
|
}
|
|
return;
|
|
}
|
|
|
|
if (value.type === "cancelled") {
|
|
reset();
|
|
return;
|
|
}
|
|
|
|
if (value.type === "error") {
|
|
const code = typeof value.code === "string" ? value.code : "service_error";
|
|
const detail = typeof value.detail === "string" ? value.detail : "Dictation failed";
|
|
fail(`Dictation service error (${boundedText(code, 64)}): ${boundedText(detail, 384)}`);
|
|
return;
|
|
}
|
|
|
|
protocolError("unknown frame type.");
|
|
};
|
|
|
|
const consumeData = (chunk: Uint8Array, current: DictateSocket) => {
|
|
if (socket !== current) return;
|
|
receiveBuffer = Buffer.concat([receiveBuffer, Buffer.from(chunk)]);
|
|
while (socket === current) {
|
|
const newline = receiveBuffer.indexOf(0x0a);
|
|
if (newline < 0) {
|
|
if (receiveBuffer.length > MAX_FRAME_BYTES) {
|
|
protocolError("frame is too large.");
|
|
}
|
|
return;
|
|
}
|
|
if (newline > MAX_FRAME_BYTES) {
|
|
protocolError("frame is too large.");
|
|
return;
|
|
}
|
|
const line = receiveBuffer.subarray(0, newline);
|
|
receiveBuffer = receiveBuffer.subarray(newline + 1);
|
|
if (line.length === 0) {
|
|
protocolError("empty frame.");
|
|
return;
|
|
}
|
|
let decoded: string;
|
|
try {
|
|
decoded = UTF8_DECODER.decode(line);
|
|
} catch {
|
|
protocolError("frame contains invalid UTF-8.");
|
|
return;
|
|
}
|
|
try {
|
|
handleFrame(JSON.parse(decoded), current);
|
|
} catch {
|
|
protocolError("malformed JSON frame.");
|
|
return;
|
|
}
|
|
}
|
|
};
|
|
|
|
const startDictation = (ctx: ExtensionContext) => {
|
|
if (state !== "idle") return;
|
|
state = "connecting";
|
|
activeContext = ctx;
|
|
activeGenerationId = newId("generation");
|
|
helloRequestId = newId("hello");
|
|
negotiated = false;
|
|
resultDelivered = false;
|
|
receiveBuffer = Buffer.alloc(0);
|
|
|
|
let current: DictateSocket;
|
|
try {
|
|
current = socketFactory(socketPath);
|
|
} catch (error) {
|
|
const detail = error instanceof Error ? error.message : "connection failed";
|
|
fail(`Unable to connect to dictation service: ${detail}`);
|
|
return;
|
|
}
|
|
socket = current;
|
|
current.onConnect(() => {
|
|
if (socket !== current || !helloRequestId) return;
|
|
send({
|
|
version: PROTOCOL_VERSION,
|
|
type: "hello",
|
|
requestId: helloRequestId,
|
|
});
|
|
});
|
|
current.onData((chunk) => consumeData(chunk, current));
|
|
current.onError((error) => {
|
|
if (socket !== current) return;
|
|
fail(`Dictation service connection failed: ${error.message}`);
|
|
});
|
|
current.onClose(() => {
|
|
if (socket !== current) return;
|
|
fail("Dictation service disconnected.");
|
|
});
|
|
};
|
|
|
|
const stopDictation = () => {
|
|
if (state !== "recording") return;
|
|
state = "stopping";
|
|
sendGenerationCommand("stop");
|
|
};
|
|
|
|
const cancelDictation = () => {
|
|
if (state === "idle") return;
|
|
if (negotiated) sendGenerationCommand("cancel");
|
|
reset();
|
|
};
|
|
|
|
const toggleDictation = (ctx: ExtensionContext) => {
|
|
if (state === "idle") startDictation(ctx);
|
|
else if (state === "recording") stopDictation();
|
|
};
|
|
|
|
pi.registerShortcut("alt+m", {
|
|
description: "Start or stop local voice dictation",
|
|
handler: async (ctx) => toggleDictation(ctx),
|
|
});
|
|
pi.registerShortcut("alt+n", {
|
|
description: "Cancel local voice dictation",
|
|
handler: async () => cancelDictation(),
|
|
});
|
|
pi.registerCommand("dictate-cancel", {
|
|
description: "Cancel local voice dictation",
|
|
handler: async () => cancelDictation(),
|
|
});
|
|
pi.on("session_shutdown", async () => {
|
|
if (state === "idle") return;
|
|
if (negotiated) sendGenerationCommand("cancel");
|
|
reset();
|
|
});
|
|
};
|
|
}
|
|
|
|
export default createDictateExtension();
|