From dd59f9b2262b79f0e201352173b3b4e677c1035d Mon Sep 17 00:00:00 2001 From: dimastatz Date: Wed, 7 Oct 2026 07:27:57 +0300 Subject: [PATCH 1/2] Implement dictation end to end (spec 0001 on whisper-flow) Streams microphone audio from a sidecar recorder to a whisper-flow server and inserts partial and final transcripts into the editor. - whisper-flow client: /ready and version checks, ws streaming, stop/flush, API key from SecretStorage, handshake rejections classified as bad key or server full (uvicorn sends both as HTTP 403) - Recorder sidecar: sox, arecord or ffmpeg, or whisperCode.recorderCommand - Editor sink: dimmed pending range replaced on each partial, committed on final, stop, user edit or cursor move - Workspace vocabulary sent as Whisper's prompt and applied to final text - Context-aware formatting for comments and commit messages - Consent for non-local servers, TLS warning, four-state status bar, Retry - Fake whisper-flow server and recorder for tests; coverage 96%+ Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NcUyc2475ugBRGPm7yQMzy --- CHANGELOG.md | 23 ++ CLAUDE.md | 19 +- README.md | 40 +++- package-lock.json | 35 +++ package.json | 46 +++- src/audio/recorder.ts | 154 ++++++++++++++ src/dictation/editorSink.ts | 190 +++++++++++++++++ src/dictation/session.ts | 183 ++++++++++++++++ src/dictation/workspaceVocabulary.ts | 67 ++++++ src/extension.ts | 195 +++++++++++++++-- src/test/extension.test.ts | 293 +++++++++++++++++++++++-- src/test/fakeRecorder.ts | 60 ++++++ src/test/fakeServer.ts | 187 ++++++++++++++++ src/test/recorder.test.ts | 149 +++++++++++++ src/test/server.test.ts | 111 ++++++++++ src/test/session.test.ts | 305 +++++++++++++++++++++++++++ src/test/text.test.ts | 222 +++++++++++++++++++ src/text/context.ts | 125 +++++++++++ src/text/format.ts | 57 +++++ src/text/vocabulary.ts | 105 +++++++++ src/whisperflow/client.ts | 177 ++++++++++++++++ src/whisperflow/messages.ts | 76 +++++++ src/whisperflow/server.ts | 136 ++++++++++++ 23 files changed, 2920 insertions(+), 35 deletions(-) create mode 100644 src/audio/recorder.ts create mode 100644 src/dictation/editorSink.ts create mode 100644 src/dictation/session.ts create mode 100644 src/dictation/workspaceVocabulary.ts create mode 100644 src/test/fakeRecorder.ts create mode 100644 src/test/fakeServer.ts create mode 100644 src/test/recorder.test.ts create mode 100644 src/test/server.test.ts create mode 100644 src/test/session.test.ts create mode 100644 src/test/text.test.ts create mode 100644 src/text/context.ts create mode 100644 src/text/format.ts create mode 100644 src/text/vocabulary.ts create mode 100644 src/whisperflow/client.ts create mode 100644 src/whisperflow/messages.ts create mode 100644 src/whisperflow/server.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index f04cceb..c941a49 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,29 @@ All notable changes to this project are documented here. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project follows [Semantic Versioning](https://semver.org/). +## [Unreleased] + +Dictation works end to end ([spec 0001](docs/specs/0001-overview.md), on +[spec 0002](docs/specs/0002-whisper-flow-transcription.md)). + +### Added + +- Streaming transcription through a whisper-flow server (1.1.0+): `/ready` check, version + check, WebSocket streaming of 16 kHz PCM, and finalizing the last phrase on stop (1.2.0+). +- Microphone capture in a sidecar process: SoX, `arecord` or `ffmpeg`, or + `whisperCode.recorderCommand`. +- Live editor insertion: partial text is dimmed and replaced as you speak, each phrase is one + undo step, and the selection is replaced. +- Workspace vocabulary: identifiers from open editors are sent as Whisper's prompt (1.2.0+) + and respelled in the output. +- Context-aware formatting for code comments and commit messages. +- Settings `whisperCode.serverUrl`, `whisperCode.model`, `whisperCode.vocabulary.enabled` and + `whisperCode.recorderCommand`; commands **Set API Key** and **Clear API Key** (secret + storage). +- Confirmation before sending audio to a non-local server, and a warning without TLS. +- Status bar states: idle, connecting, recording, error. Errors for a bad API key, a full + server, and lost connections (with **Retry**). + ## [0.1.0] - 2026-10-05 First preview release. Sets up the extension and project tooling; audio capture and diff --git a/CLAUDE.md b/CLAUDE.md index ae17b13..9eec925 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -43,9 +43,18 @@ Press **F5** in VS Code ("Run Extension") to launch the Extension Development Ho ## Layout -- `src/extension.ts` — entry point (`activate`/`deactivate`); registers - `whisperCode.toggleDictation` and the status bar item. Dictation itself is still a stub. -- `src/test/*.test.ts` — integration tests (Mocha `suite`/`test`). +- `src/extension.ts` — entry point (`activate`/`deactivate`): commands, settings, consent + prompts for non-local servers, the four-state status bar item. +- `src/dictation/session.ts` — one dictation session (`/ready` → socket → recorder → sink). + No `vscode` import; `editorSink.ts` (pending range, decorations) and + `workspaceVocabulary.ts` (symbol index) are the `vscode` side. +- `src/whisperflow/` — whisper-flow protocol: URLs and `/ready` (`server.ts`), server + messages and close codes (`messages.ts`), the WebSocket client (`client.ts`). +- `src/audio/recorder.ts` — microphone sidecar process (sox / arecord / ffmpeg) and framing. +- `src/text/` — context detector, post-processor, vocabulary. Plain functions. +- `src/test/*.test.ts` — tests (Mocha `suite`/`test`). `fakeServer.ts` is a fake whisper-flow + (real HTTP + WebSocket) and `fakeRecorder.ts` a fake recorder process; no real server or + microphone is needed. - `docs/specs/` — design specs; start with `0001-overview.md`. - `docs/images/` — images referenced from docs. @@ -57,4 +66,8 @@ Press **F5** in VS Code ("Run Extension") to launch the Extension Development Ho - Keep coverage at or above 95%: add tests with new code rather than excluding files. - Logic that doesn't need the `vscode` API should live in plain modules so it can be tested directly. +- whisper-flow (uvicorn) rejects a handshake with HTTP 403 for both a bad API key and a full + server, never close codes 1008/1013; `classifyRejectedHandshake` tells them apart. +- The `undo` command is a no-op in the test window (it never has OS focus), so undo grouping + can't be asserted in tests. - Commit on a feature branch and open a PR against `main`; don't push directly to `main`. diff --git a/README.md b/README.md index 67f15a4..c2b90e9 100644 --- a/README.md +++ b/README.md @@ -2,9 +2,8 @@ Real-time voice dictation for VS Code, powered by self-hosted Whisper through [whisper-flow](https://github.com/dimastatz/whisper-flow). Biases transcription with your workspace symbols so it spells useMemo right, and adapts to whether you're dictating a comment, a commit message, or a prompt. -> **Status: early development.** The extension installs and the dictation command and -> status bar toggle work, but audio capture and transcription are not implemented yet. -> See [docs/specs/0001-overview.md](docs/specs/0001-overview.md) for the plan. +> **Status: preview.** Dictation works end to end against a whisper-flow server. Expect rough +> edges; see [docs/specs](docs/specs) for the design. ## Usage @@ -14,10 +13,39 @@ Real-time voice dictation for VS Code, powered by self-hosted Whisper through [w | | Or click the microphone **Dictate** item in the status bar | | | Or run **Whisper Code: Toggle Dictation** from the Command Palette | -While dictating, the status bar item changes to **Dictating**. Toggle again to stop. +The status bar shows **Connecting**, then **Dictating** while it records. Text appears dimmed as +you speak and settles when each phrase is final; toggle again to stop. Each phrase is one undo +step. -Transcription runs on a [whisper-flow](https://github.com/dimastatz/whisper-flow) server that you -run yourself, on `localhost` by default. No third-party cloud speech service is used. +Output adapts to where the cursor is: + +- **Code comment:** sentences start with a capital letter after the comment marker. +- **Commit message** (`COMMIT_EDITMSG`, git commit editor): the subject line has no trailing + period. +- **Anything else** (prose, prompts): inserted as spoken. + +Identifiers from your open editors (such as `useMemo` or `getUserById`) are sent to the server +as a vocabulary hint and respelled in the text, so "use memo" comes out as `useMemo`. + +## Requirements + +- A [whisper-flow](https://github.com/dimastatz/whisper-flow) server, 1.1.0 or newer (1.2.0+ + for the vocabulary hint and model choice). Run it yourself, for example with + `./run.sh -run-server` or Docker. No third-party cloud speech service is used. +- A command-line audio recorder: [SoX](https://sourceforge.net/projects/sox/) (`brew install sox`, + `apt install sox`), or `arecord` on Linux, or `ffmpeg` on macOS and Linux. + +## Settings + +| Setting | Default | Description | +| -------------------------------- | ----------------------- | ------------------------------------------------------------------------------------------------------- | +| `whisperCode.serverUrl` | `http://localhost:8181` | whisper-flow server. A non-local server asks for your confirmation once, and warns without `https`. | +| `whisperCode.model` | (server default) | Model from the server's `/ready` list, such as `base.en.pt`. | +| `whisperCode.vocabulary.enabled` | `true` | Bias transcription toward identifiers from open editors. | +| `whisperCode.recorderCommand` | (auto) | Command that writes raw 16 kHz mono 16-bit little-endian PCM to stdout, if the defaults don't suit you. | + +If the server sets `WF_API_KEY`, run **Whisper Code: Set API Key**. The key is kept in VS Code's +secret storage, not in `settings.json`. ## Running from source diff --git a/package-lock.json b/package-lock.json index e494420..41a0df9 100644 --- a/package-lock.json +++ b/package-lock.json @@ -8,10 +8,14 @@ "name": "whisper-code", "version": "0.1.0", "license": "MIT", + "dependencies": { + "ws": "^8.22.0" + }, "devDependencies": { "@types/mocha": "^10.0.10", "@types/node": "20.x", "@types/vscode": "^1.95.0", + "@types/ws": "^8.18.2", "@vscode/test-cli": "^0.0.15", "@vscode/test-electron": "^3.1.0", "eslint": "^9.25.1", @@ -403,6 +407,16 @@ "dev": true, "license": "MIT" }, + "node_modules/@types/ws": { + "version": "8.18.2", + "resolved": "https://registry.npmjs.org/@types/ws/-/ws-8.18.2.tgz", + "integrity": "sha512-67MQl+fpWKVTT1NYdnmo3U4sc/xPo/zQBncVnI74qmQa0z/b+1g6iYqNmGCPbxO+zz2aklb08a0oHfegiVd0/w==", + "dev": true, + "license": "MIT", + "dependencies": { + "@types/node": "*" + } + }, "node_modules/@typescript-eslint/eslint-plugin": { "version": "8.71.0", "resolved": "https://registry.npmjs.org/@typescript-eslint/eslint-plugin/-/eslint-plugin-8.71.0.tgz", @@ -3312,6 +3326,27 @@ "url": "https://github.com/chalk/ansi-styles?sponsor=1" } }, + "node_modules/ws": { + "version": "8.22.0", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.22.0.tgz", + "integrity": "sha512-Ydggc987+RO0AnWtZ/7Wq9FtNvcrL1b/RO0ud9mWjUPgDrsAAwQSF51sm2hm1XofbU/4jkpGEsLFsZZxU+1DOg==", + "license": "MIT", + "engines": { + "node": ">=10.0.0" + }, + "peerDependencies": { + "bufferutil": "^4.0.1", + "utf-8-validate": ">=5.0.2" + }, + "peerDependenciesMeta": { + "bufferutil": { + "optional": true + }, + "utf-8-validate": { + "optional": true + } + } + }, "node_modules/y18n": { "version": "5.0.8", "resolved": "https://registry.npmjs.org/y18n/-/y18n-5.0.8.tgz", diff --git a/package.json b/package.json index f5cacdd..5d88351 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { "name": "whisper-code", "displayName": "Whisper Code", - "description": "Real-time voice dictation for VS Code, powered by local Whisper.", + "description": "Real-time voice dictation for VS Code, powered by a self-hosted whisper-flow server.", "version": "0.1.0", "publisher": "dimastatz", "license": "MIT", @@ -25,6 +25,16 @@ "command": "whisperCode.toggleDictation", "title": "Toggle Dictation", "category": "Whisper Code" + }, + { + "command": "whisperCode.setApiKey", + "title": "Set API Key", + "category": "Whisper Code" + }, + { + "command": "whisperCode.clearApiKey", + "title": "Clear API Key", + "category": "Whisper Code" } ], "keybindings": [ @@ -33,7 +43,35 @@ "key": "ctrl+alt+d", "mac": "cmd+alt+d" } - ] + ], + "configuration": { + "title": "Whisper Code", + "properties": { + "whisperCode.serverUrl": { + "type": "string", + "default": "http://localhost:8181", + "markdownDescription": "URL of the [whisper-flow](https://github.com/dimastatz/whisper-flow) server. The WebSocket URL is derived from it (`http` \u2192 `ws`, `https` \u2192 `wss`, plus `/ws`). Audio is sent only to this server; a non-local server needs your confirmation once." + }, + "whisperCode.model": { + "type": "string", + "default": "", + "markdownDescription": "Model for each session, from the server's `GET /ready` `models` list (for example `base.en.pt`). Empty uses the server's default. Needs whisper-flow 1.2 or newer." + }, + "whisperCode.vocabulary.enabled": { + "type": "boolean", + "default": true, + "markdownDescription": "Bias transcription toward identifiers from open editors (such as `useMemo` or `getUserById`), and respell them in the inserted text. Sending the vocabulary to the server needs whisper-flow 1.2 or newer." + }, + "whisperCode.recorderCommand": { + "type": "array", + "items": { + "type": "string" + }, + "default": [], + "markdownDescription": "Command that records the microphone and writes raw 16 kHz mono signed 16-bit little-endian PCM to stdout, for example `[\"sox\", \"-q\", \"-d\", \"-t\", \"raw\", \"-r\", \"16000\", \"-e\", \"signed-integer\", \"-b\", \"16\", \"-c\", \"1\", \"-L\", \"-\"]`. Empty tries `sox`, then `arecord` (Linux) or `ffmpeg`." + } + } + } }, "scripts": { "vscode:prepublish": "npm run compile", @@ -53,11 +91,15 @@ "@types/mocha": "^10.0.10", "@types/node": "20.x", "@types/vscode": "^1.95.0", + "@types/ws": "^8.18.2", "@vscode/test-cli": "^0.0.15", "@vscode/test-electron": "^3.1.0", "eslint": "^9.25.1", "prettier": "^3.9.9", "typescript": "^5.8.3", "typescript-eslint": "^8.31.1" + }, + "dependencies": { + "ws": "^8.22.0" } } diff --git a/src/audio/recorder.ts b/src/audio/recorder.ts new file mode 100644 index 0000000..962f9a8 --- /dev/null +++ b/src/audio/recorder.ts @@ -0,0 +1,154 @@ +// Microphone capture in a sidecar process (spec 0002: AU-1, AU-2, AU-3, AU-4, ER-3). +import { spawn as nodeSpawn, type ChildProcess } from "node:child_process"; + +export const SAMPLE_RATE = 16000; +/** 1024 samples of 16-bit mono PCM, whisper-flow's default WF_CHUNK_SIZE (AU-2). */ +export const FRAME_BYTES = 2048; + +/** Cuts a byte stream into fixed-size frames, keeping the remainder for the next push. */ +export class FrameSplitter { + private rest = Buffer.alloc(0); + + push(chunk: Buffer): Buffer[] { + const data = this.rest.length ? Buffer.concat([this.rest, chunk]) : chunk; + const frames: Buffer[] = []; + let offset = 0; + for (; offset + FRAME_BYTES <= data.length; offset += FRAME_BYTES) { + frames.push(data.subarray(offset, offset + FRAME_BYTES)); + } + this.rest = Buffer.from(data.subarray(offset)); + return frames; + } +} + +/** + * Recorders tried in order when `whisperCode.recorderCommand` is empty. Each writes raw + * 16 kHz mono s16le PCM to stdout and needs no Python (AU-4). + */ +export function defaultRecorderCommands(platform: NodeJS.Platform): string[][] { + const rate = String(SAMPLE_RATE); + const sox = [ + "sox", + "-q", + "-d", + "-t", + "raw", + "-r", + rate, + "-e", + "signed-integer", + "-b", + "16", + "-c", + "1", + "-L", + "-", + ]; + const ffmpegOut = ["-ac", "1", "-ar", rate, "-f", "s16le", "-"]; + const ffmpegQuiet = ["ffmpeg", "-hide_banner", "-loglevel", "error", "-nostdin"]; + switch (platform) { + case "darwin": + return [sox, [...ffmpegQuiet, "-f", "avfoundation", "-i", ":default", ...ffmpegOut]]; + case "linux": + return [ + sox, + ["arecord", "-q", "-f", "S16_LE", "-r", rate, "-c", "1", "-t", "raw"], + [...ffmpegQuiet, "-f", "pulse", "-i", "default", ...ffmpegOut], + ]; + default: + return [sox]; + } +} + +export interface RecorderHandlers { + onFrame(frame: Buffer): void; + /** The recorder failed to start or exited on its own. Not called after `stop()`. */ + onError(message: string): void; +} + +export type Spawn = (command: string, args: string[]) => ChildProcess; + +export interface Recorder { + stop(): void; +} + +/** + * Starts the first recorder command that exists. A command that is not installed (ENOENT) + * falls through to the next one. + */ +export function startRecorder( + commands: string[][], + handlers: RecorderHandlers, + spawn: Spawn = (command, args) => nodeSpawn(command, args, { stdio: ["ignore", "pipe", "pipe"] }), +): Recorder { + let child: ChildProcess | undefined; + let stopped = false; + const missing: string[] = []; + + const tryStart = (index: number) => { + if (index >= commands.length || commands[index]?.length === 0) { + handlers.onError( + missing.length + ? `No audio recorder found (tried ${missing.join(", ")}). Install sox, or set whisperCode.recorderCommand.` + : "No audio recorder configured.", + ); + return; + } + const [program, ...args] = commands[index] as [string, ...string[]]; + const splitter = new FrameSplitter(); + let stderr = ""; + const current = spawn(program, args); + child = current; + + current.stdout?.on("data", (chunk: Buffer) => { + if (stopped) { + return; + } + for (const frame of splitter.push(chunk)) { + handlers.onFrame(frame); + } + }); + current.stderr?.on("data", (chunk: Buffer) => { + stderr = (stderr + chunk.toString("utf8")).slice(-500); + }); + current.once("error", (error: NodeJS.ErrnoException) => { + if (stopped) { + return; + } + if (error.code === "ENOENT") { + missing.push(program); + tryStart(index + 1); + } else { + handlers.onError(`Audio recorder ${program} failed: ${error.message}`); + } + }); + current.once("exit", (code, signal) => { + // ENOENT also emits exit on some platforms; the error handler moved on already. + if (stopped || child !== current || missing.includes(program)) { + return; + } + const detail = stderr.trim() || (signal ? `signal ${signal}` : `exit code ${String(code)}`); + handlers.onError(`Audio recorder ${program} stopped: ${detail}`); + }); + }; + + tryStart(0); + + return { + stop() { + if (stopped) { + return; + } + stopped = true; + const current = child; + if (current?.exitCode === null && current.signalCode === null) { + current.kill("SIGTERM"); + const force = setTimeout(() => current.kill("SIGKILL"), 1000); + force.unref(); + current.once("exit", () => { + clearTimeout(force); + }); + } + }, + }; +} diff --git a/src/dictation/editorSink.ts b/src/dictation/editorSink.ts new file mode 100644 index 0000000..0290e24 --- /dev/null +++ b/src/dictation/editorSink.ts @@ -0,0 +1,190 @@ +// Editor integration: pending range, decorations, undo grouping (spec 0002: RS-2..RS-6). +import * as vscode from "vscode"; +import { detectTarget } from "../text/context"; +import { formatFinal, formatPartial } from "../text/format"; +import type { Vocabulary } from "../text/vocabulary"; +import type { TranscriptSink } from "./session"; + +interface Segment { + editor: vscode.TextEditor; + /** Offset where the segment starts. */ + start: number; + /** Length of the text in the editor that the next result replaces. */ + length: number; + /** No edit has been made yet, so the next one starts a new undo step. */ + fresh: boolean; +} + +export class EditorSink implements TranscriptSink, vscode.Disposable { + private segment: Segment | undefined; + private applying = false; + /** Document version after our last edit, to tell our change events from the user's. */ + private ownVersion = -1; + private readonly decoration = vscode.window.createTextEditorDecorationType({ + opacity: "0.55", + fontStyle: "italic", + }); + private readonly subscriptions: vscode.Disposable[]; + + constructor(private readonly vocabulary: Vocabulary | undefined) { + this.subscriptions = [ + this.decoration, + vscode.workspace.onDidChangeTextDocument((event) => { + if ( + !this.applying && + event.document === this.segment?.editor.document && + event.document.version !== this.ownVersion && + event.contentChanges.length > 0 + ) { + // The user edited the document: keep what's there and start over at the cursor. + this.commit(); + } + }), + vscode.window.onDidChangeTextEditorSelection((event) => { + const segment = this.segment; + if (this.applying || event.textEditor !== segment?.editor) { + return; + } + const offset = segment.editor.document.offsetAt( + event.selections[0]?.active ?? new vscode.Position(0, 0), + ); + if (offset < segment.start || offset > segment.start + segment.length) { + this.commit(); + } + }), + vscode.window.onDidChangeActiveTextEditor(() => { + this.commit(); + }), + ]; + } + + async partial(text: string): Promise { + const segment = this.segment ?? this.begin(); + if (!segment) { + return; + } + const newText = formatPartial(text, this.linePrefix(segment)); + await this.replace(segment, newText, false); + this.decorate(segment); + } + + async final(text: string): Promise { + const segment = this.segment ?? this.begin(); + if (!segment) { + // No editor: type the text into whatever has focus (e.g. a chat or commit input box). + const typed = formatFinal(text, { + target: "prompt", + linePrefix: "", + firstLine: false, + applyVocabulary: this.applyVocabulary, + }); + if (typed) { + await vscode.commands.executeCommand("type", { text: typed }); + } + return; + } + const { document } = segment.editor; + const position = document.positionAt(segment.start); + const linePrefix = this.linePrefix(segment); + const newText = formatFinal(text, { + target: detectTarget({ + languageId: document.languageId, + uriScheme: document.uri.scheme, + fileName: document.fileName, + linePrefix, + }), + linePrefix, + firstLine: position.line === 0, + applyVocabulary: this.applyVocabulary, + }); + await this.replace(segment, newText, true); + // The next segment continues right after this one. + this.segment = { + editor: segment.editor, + start: segment.start + newText.length, + length: 0, + fresh: true, + }; + segment.editor.setDecorations(this.decoration, []); + } + + finish(): Promise { + this.commit(); + return Promise.resolve(); + } + + dispose(): void { + this.commit(); + for (const subscription of this.subscriptions) { + subscription.dispose(); + } + } + + /** Leaves pending text in place as committed text. */ + private commit(): void { + this.segment?.editor.setDecorations(this.decoration, []); + this.segment = undefined; + } + + /** Starts a segment at the cursor, replacing the selection if there is one. */ + private begin(): Segment | undefined { + const editor = vscode.window.activeTextEditor; + if (!editor) { + return undefined; + } + const { document, selection } = editor; + const start = document.offsetAt(selection.start); + this.segment = { editor, start, length: document.offsetAt(selection.end) - start, fresh: true }; + return this.segment; + } + + private readonly applyVocabulary = (text: string) => this.vocabulary?.apply(text) ?? text; + + private linePrefix(segment: Segment): string { + const position = segment.editor.document.positionAt(segment.start); + return segment.editor.document.lineAt(position.line).text.slice(0, position.character); + } + + /** Replaces the segment's text. Edits of one segment form a single undo step (RS-5). */ + private async replace(segment: Segment, text: string, last: boolean): Promise { + const { editor } = segment; + const { document } = editor; + const range = new vscode.Range( + document.positionAt(segment.start), + document.positionAt(segment.start + segment.length), + ); + if (document.getText(range) === text && !last) { + return; + } + this.applying = true; + try { + const applied = await editor.edit( + (builder) => { + builder.replace(range, text); + }, + { undoStopBefore: segment.fresh, undoStopAfter: last }, + ); + if (!applied) { + this.commit(); + return; + } + this.ownVersion = document.version; + segment.fresh = false; + segment.length = text.length; + const end = document.positionAt(segment.start + text.length); + editor.selection = new vscode.Selection(end, end); + } finally { + this.applying = false; + } + } + + private decorate(segment: Segment): void { + const { document } = segment.editor; + segment.editor.setDecorations(this.decoration, [ + new vscode.Range( + document.positionAt(segment.start), + document.positionAt(segment.start + segment.length), + ), + ]); + } +} diff --git a/src/dictation/session.ts b/src/dictation/session.ts new file mode 100644 index 0000000..2ed18a0 --- /dev/null +++ b/src/dictation/session.ts @@ -0,0 +1,183 @@ +// One dictation session: /ready check → socket → recorder → transcript sink. +// Spec 0002: CF-3, TR-4, RS-1..3, ER-1..3. Has no vscode dependency so it can be unit tested. +import { startRecorder, type Recorder, type Spawn } from "../audio/recorder"; +import { HandshakeError, WhisperFlowConnection } from "../whisperflow/client"; +import { describeClose, type ServerMessage } from "../whisperflow/messages"; +import { checkReady, errorMessage, supportsControlFrames, type Fetch } from "../whisperflow/server"; + +export type SessionState = "idle" | "connecting" | "recording" | "error"; + +export interface SessionError { + message: string; + /** Offer a Retry action (ER-2). */ + retry: boolean; +} + +/** Where transcripts go; the editor implementation lives in editorSink.ts. */ +export interface TranscriptSink { + /** Replace the pending text of the current segment (RS-2). */ + partial(text: string): Promise; + /** Replace the pending text with the final text and commit it (RS-3). */ + final(text: string): Promise; + /** Session over: commit whatever is pending as-is (RS-4). */ + finish(): Promise; +} + +export interface SessionOptions { + serverUrl: string; + apiKey: string | undefined; + model: string | undefined; + /** Vocabulary prompt for the `start` frame (whisper-flow 1.2+). */ + prompt: string | undefined; + recorderCommands: string[][]; + sink: TranscriptSink; + onState(state: SessionState, error?: SessionError): void; + fetch?: Fetch; + spawn?: Spawn; +} + +export class DictationSession { + private connection: WhisperFlowConnection | undefined; + private recorder: Recorder | undefined; + private queue: Promise = Promise.resolve(); + private ended = false; + private useStopFrame = false; + + constructor(private readonly options: SessionOptions) {} + + async start(): Promise { + const { options } = this; + options.onState("connecting"); + const fetchFn = options.fetch ?? fetch; + try { + const info = await checkReady(options.serverUrl, options.apiKey, fetchFn); + if (this.isEnded()) { + return; + } + this.useStopFrame = supportsControlFrames(info); + const model = + options.model && info.models.includes(options.model) ? options.model : undefined; + if (options.model && !model) { + throw new Error( + `Model ${options.model} isn't available on ${options.serverUrl} (available: ${info.models.join(", ") || "none"}).`, + ); + } + + const connection = new WhisperFlowConnection( + options.serverUrl, + options.apiKey, + { + onMessage: (message) => { + this.handleMessage(message); + }, + onClose: (code, reason) => { + this.handleClose(code, reason); + }, + }, + fetchFn, + ); + this.connection = connection; + await connection.open(); + if (this.isEnded()) { + connection.dispose(); + return; + } + if (this.useStopFrame && (options.prompt || model)) { + connection.sendStart({ prompt: options.prompt, model }); + } + } catch (error) { + const closed = + error instanceof HandshakeError ? describeClose(error.closeCode, "") : undefined; + const message = + error instanceof HandshakeError && error.closeCode !== 0 && closed + ? closed.message + : error instanceof HandshakeError + ? `Can't connect to the whisper-flow server at ${options.serverUrl} (${error.message}).` + : errorMessage(error); + await this.fail({ + message, + retry: !(error instanceof HandshakeError && error.closeCode !== 0), + }); + return; + } + + this.recorder = startRecorder( + options.recorderCommands, + { + onFrame: (frame) => this.connection?.sendAudio(frame), + onError: (message) => { + void this.fail({ message, retry: true }); + }, + }, + options.spawn, + ); + if (!this.isEnded()) { + options.onState("recording"); + } + } + + /** Stops recording, lets the server finalize the last segment, and commits it (ER-3). */ + async stop(): Promise { + if (this.ended) { + return; + } + this.ended = true; + this.recorder?.stop(); + await this.connection?.stop(this.useStopFrame); + this.connection?.dispose(); + await this.drain(); + this.options.onState("idle"); + } + + /** Ends the session without waiting for the server (extension shutdown). */ + dispose(): void { + this.ended = true; + this.recorder?.stop(); + this.connection?.dispose(); + } + + /** Read through a method: TypeScript keeps `this.ended` narrowed across awaits. */ + private isEnded(): boolean { + return this.ended; + } + + private handleMessage(message: ServerMessage): void { + if (message.kind === "error") { + // An unknown control frame; the session continues. + return; + } + const { sink } = this.options; + this.enqueue(() => (message.isPartial ? sink.partial(message.text) : sink.final(message.text))); + } + + private handleClose(code: number, reason: string): void { + if (this.ended) { + return; + } + const error = describeClose(code, reason) ?? { + message: "The whisper-flow server ended the session.", + retry: true, + }; + void this.fail(error); + } + + private async fail(error: SessionError): Promise { + if (this.ended) { + return; + } + this.ended = true; + this.recorder?.stop(); + this.connection?.dispose(); + await this.drain(); + this.options.onState("error", error); + } + + private enqueue(task: () => Promise): void { + this.queue = this.queue.then(task).catch(() => undefined); + } + + private async drain(): Promise { + this.enqueue(() => this.options.sink.finish()); + await this.queue; + } +} diff --git a/src/dictation/workspaceVocabulary.ts b/src/dictation/workspaceVocabulary.ts new file mode 100644 index 0000000..17dc61f --- /dev/null +++ b/src/dictation/workspaceVocabulary.ts @@ -0,0 +1,67 @@ +// Symbol index (spec 0001, goal 2): collects identifiers from the open workspace. +import * as vscode from "vscode"; +import { + countIdentifiers, + createVocabulary, + rankIdentifiers, + type Vocabulary, +} from "../text/vocabulary"; + +const MAX_DOCUMENT_CHARS = 1_000_000; +const SYMBOL_TIMEOUT_MS = 1500; + +/** + * Ranks identifiers from the active editor (highest weight), its document symbols, other + * visible editors and other open documents. + */ +export async function collectWorkspaceVocabulary(): Promise { + const counts = new Map(); + const active = vscode.window.activeTextEditor?.document; + const visible = new Set(vscode.window.visibleTextEditors.map((editor) => editor.document)); + + for (const document of vscode.workspace.textDocuments) { + if (document.uri.scheme === "output" || document.getText().length > MAX_DOCUMENT_CHARS) { + continue; + } + const weight = document === active ? 3 : visible.has(document) ? 2 : 1; + countIdentifiers(document.getText(), counts, weight); + } + + if (active) { + for (const name of await documentSymbolNames(active.uri)) { + countIdentifiers(name, counts, 5); + } + } + return createVocabulary(rankIdentifiers(counts)); +} + +async function documentSymbolNames(uri: vscode.Uri): Promise { + const symbols = await withTimeout( + vscode.commands.executeCommand< + (vscode.DocumentSymbol | vscode.SymbolInformation)[] | undefined + >("vscode.executeDocumentSymbolProvider", uri), + SYMBOL_TIMEOUT_MS, + ).catch(() => undefined); + const names: string[] = []; + const visit = (items: readonly (vscode.DocumentSymbol | vscode.SymbolInformation)[]) => { + for (const item of items) { + names.push(item.name); + if ("children" in item) { + visit(item.children); + } + } + }; + visit(symbols ?? []); + return names; +} + +function withTimeout(promise: Thenable, ms: number): Promise { + return Promise.race([ + Promise.resolve(promise), + new Promise((resolve) => + setTimeout(() => { + resolve(undefined); + }, ms), + ), + ]); +} diff --git a/src/extension.ts b/src/extension.ts index 1251b86..cafffb3 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -1,35 +1,206 @@ import * as vscode from "vscode"; +import { defaultRecorderCommands } from "./audio/recorder"; +import { EditorSink } from "./dictation/editorSink"; +import { DictationSession, type SessionError, type SessionState } from "./dictation/session"; +import { collectWorkspaceVocabulary } from "./dictation/workspaceVocabulary"; +import { + errorMessage, + isInsecureRemote, + isLocalServer, + parseServerUrl, +} from "./whisperflow/server"; -let recording = false; +const API_KEY_SECRET = "whisperCode.apiKey"; +const APPROVED_SERVERS = "whisperCode.approvedServers"; + +let context: vscode.ExtensionContext; let statusItem: vscode.StatusBarItem; +let state: SessionState = "idle"; +let session: DictationSession | undefined; +let sink: EditorSink | undefined; +/** Set while a toggle is being handled, so a double press doesn't start two sessions. */ +let busy = false; -export function activate(context: vscode.ExtensionContext) { +export function activate(extensionContext: vscode.ExtensionContext) { + context = extensionContext; statusItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Right, 100); statusItem.command = "whisperCode.toggleDictation"; - updateStatus(); + setState("idle"); statusItem.show(); context.subscriptions.push( statusItem, vscode.commands.registerCommand("whisperCode.toggleDictation", toggleDictation), + vscode.commands.registerCommand("whisperCode.setApiKey", setApiKey), + vscode.commands.registerCommand("whisperCode.clearApiKey", clearApiKey), + { dispose: endSession }, ); } export function deactivate() { - // Nothing to clean up yet; disposables are released via context.subscriptions. + // No sidecar may outlive the extension (ER-3). + endSession(); } +/** True while connecting or recording. */ export function isRecording(): boolean { - return recording; + return state === "connecting" || state === "recording"; +} + +export function getState(): SessionState { + return state; +} + +async function toggleDictation(): Promise { + if (busy) { + return; + } + busy = true; + try { + if (session && isRecording()) { + await session.stop(); + } else { + await startDictation(); + } + } finally { + busy = false; + } +} + +async function startDictation(): Promise { + const config = vscode.workspace.getConfiguration("whisperCode"); + const serverUrl = config.get("serverUrl", "http://localhost:8181"); + try { + parseServerUrl(serverUrl); + } catch (error) { + showError({ message: errorMessage(error), retry: false }); + return; + } + if (!(await confirmServer(serverUrl))) { + return; + } + + endSession(); + const useVocabulary = config.get("vocabulary.enabled", true); + const vocabulary = useVocabulary ? await collectWorkspaceVocabulary() : undefined; + const recorderCommand = config.get("recorderCommand", []); + const currentSink = new EditorSink(vocabulary); + sink = currentSink; + + const current: DictationSession = new DictationSession({ + serverUrl, + apiKey: await context.secrets.get(API_KEY_SECRET), + model: emptyToUndefined(config.get("model")), + prompt: emptyToUndefined(vocabulary?.prompt), + recorderCommands: recorderCommand.length + ? [recorderCommand] + : defaultRecorderCommands(process.platform), + sink: currentSink, + onState: (next, error) => { + if (session !== current) { + return; + } + setState(next, error); + if (next === "idle" || next === "error") { + currentSink.dispose(); + } + if (error) { + showError(error); + } + }, + }); + session = current; + // Not awaited, so pressing the toggle again while connecting cancels the session. + void current.start(); +} + +/** Ends any session immediately, without waiting for the server. */ +function endSession(): void { + session?.dispose(); + session = undefined; + sink?.dispose(); + sink = undefined; + setState("idle"); +} + +/** + * Asks once per URL before sending audio to a non-local host (CF-4), and warns about a + * non-local server without TLS (CF-5). + */ +async function confirmServer(serverUrl: string): Promise { + if (isLocalServer(serverUrl)) { + return true; + } + const host = parseServerUrl(serverUrl).host; + const approved = context.globalState.get(APPROVED_SERVERS, []); + if (!approved.includes(serverUrl)) { + const allow = "Send Audio"; + const answer = await vscode.window.showWarningMessage( + `Whisper Code will send your microphone audio to ${host} (${serverUrl}). Continue?`, + { modal: true }, + allow, + ); + if (answer !== allow) { + return false; + } + await context.globalState.update(APPROVED_SERVERS, [...approved, serverUrl]); + } + if (isInsecureRemote(serverUrl)) { + void vscode.window.showWarningMessage( + `The connection to ${host} isn't encrypted. Audio and your API key are sent in clear text; use https:// if the server supports it.`, + ); + } + return true; +} + +async function setApiKey(): Promise { + const key = await vscode.window.showInputBox({ + title: "whisper-flow API key", + prompt: "Stored in VS Code's secret storage and sent as the x-api-key header.", + password: true, + ignoreFocusOut: true, + }); + if (key !== undefined) { + await (key + ? context.secrets.store(API_KEY_SECRET, key) + : context.secrets.delete(API_KEY_SECRET)); + } } -function toggleDictation() { - recording = !recording; - updateStatus(); - // TODO: start/stop audio capture and the local Whisper runner (see docs/specs/0001-overview.md). +async function clearApiKey(): Promise { + await context.secrets.delete(API_KEY_SECRET); +} + +function showError(error: SessionError): void { + const retry = "Retry"; + const actions = error.retry ? [retry] : []; + void vscode.window + .showErrorMessage(`Whisper Code: ${error.message}`, ...actions) + .then((choice) => { + if (choice === retry) { + void toggleDictation(); + } + }); +} + +const STATUS: Record = { + idle: { text: "$(mic) Dictate", tooltip: "Start dictation" }, + connecting: { text: "$(sync~spin) Connecting", tooltip: "Connecting to whisper-flow…" }, + recording: { text: "$(record) Dictating", tooltip: "Stop dictation" }, + error: { + text: "$(error) Dictate", + tooltip: "Dictation stopped with an error. Click to start again.", + }, +}; + +function setState(next: SessionState, error?: SessionError): void { + state = next; + statusItem.text = STATUS[next].text; + statusItem.tooltip = error ? `${STATUS[next].tooltip}\n${error.message}` : STATUS[next].tooltip; + statusItem.backgroundColor = + next === "error" ? new vscode.ThemeColor("statusBarItem.errorBackground") : undefined; } -function updateStatus() { - statusItem.text = recording ? "$(record) Dictating" : "$(mic) Dictate"; - statusItem.tooltip = recording ? "Stop dictation" : "Start dictation"; +function emptyToUndefined(value: string | undefined): string | undefined { + return value === "" ? undefined : value; } diff --git a/src/test/extension.test.ts b/src/test/extension.test.ts index e7624c1..e896590 100644 --- a/src/test/extension.test.ts +++ b/src/test/extension.test.ts @@ -1,38 +1,307 @@ import * as assert from "assert"; import * as vscode from "vscode"; -import { deactivate, isRecording } from "../extension"; +import { deactivate, getState, isRecording } from "../extension"; +import { FakeWhisperFlow, final, partial, waitFor, type FakeServerOptions } from "./fakeServer"; const TOGGLE = "whisperCode.toggleDictation"; +/** A recorder that writes a frame of silence every 20 ms, run by the extension host's Node. */ +const FAKE_RECORDER = [ + process.execPath, + "-e", + "setInterval(() => process.stdout.write(Buffer.alloc(2048)), 20)", +]; + +type Window = typeof vscode.window; +const stubs: (() => void)[] = []; + +/** Replaces a `vscode.window` function until the end of the test. */ +function stub(name: K, replacement: Window[K]): void { + const original = vscode.window[name]; + (vscode.window as Record)[name] = replacement; + stubs.push(() => { + (vscode.window as Record)[name] = original; + }); +} + +function config() { + return vscode.workspace.getConfiguration("whisperCode"); +} + +let server: FakeWhisperFlow | undefined; + +async function useServer(options: FakeServerOptions): Promise { + server = await new FakeWhisperFlow(options).listen(); + await config().update("serverUrl", server.url, vscode.ConfigurationTarget.Global); + return server; +} + +async function openEditor(content: string, language = "typescript"): Promise { + const document = await vscode.workspace.openTextDocument({ language, content }); + const editor = await vscode.window.showTextDocument(document); + const end = document.positionAt(content.length); + editor.selection = new vscode.Selection(end, end); + return editor; +} + +async function startDictation(): Promise { + await vscode.commands.executeCommand(TOGGLE); + await waitFor(() => getState() === "recording" || getState() === "error"); +} + +async function stopDictation(): Promise { + await vscode.commands.executeCommand(TOGGLE); + await waitFor(() => !isRecording()); +} + suite("Extension", () => { suiteSetup(async () => { const extension = vscode.extensions.getExtension("dimastatz.whisper-code"); assert.ok(extension, "extension not found"); await extension.activate(); + await config().update("recorderCommand", FAKE_RECORDER, vscode.ConfigurationTarget.Global); }); teardown(async () => { if (isRecording()) { - await vscode.commands.executeCommand(TOGGLE); + await stopDictation(); + } + for (const restore of stubs.splice(0)) { + restore(); + } + await server?.close(); + server = undefined; + // Revert first so closing an edited untitled document doesn't ask to save it. + for (let i = 0; i < 10 && vscode.window.activeTextEditor; i++) { + await vscode.commands.executeCommand("workbench.action.revertAndCloseActiveEditor"); + } + }); + + suiteTeardown(async () => { + for (const key of ["serverUrl", "recorderCommand", "vocabulary.enabled", "model"]) { + await config().update(key, undefined, vscode.ConfigurationTarget.Global); } + await vscode.commands.executeCommand("whisperCode.clearApiKey"); }); - test("toggleDictation command is registered", async () => { + test("commands are registered", async () => { const commands = await vscode.commands.getCommands(true); - assert.ok(commands.includes(TOGGLE)); + for (const command of [TOGGLE, "whisperCode.setApiKey", "whisperCode.clearApiKey"]) { + assert.ok(commands.includes(command), command); + } }); - test("toggleDictation starts and stops recording", async () => { - assert.strictEqual(isRecording(), false); - await vscode.commands.executeCommand(TOGGLE); - assert.strictEqual(isRecording(), true); + test("dictates into a comment with workspace vocabulary", async () => { + const fake = await useServer({ + events: [ + partial(" fix the"), + partial(" fix the cash"), + final(" fix the cache in get user by id."), + ], + }); + const content = "const user = getUserById(1);\n// "; + const editor = await openEditor(content); + await startDictation(); + assert.strictEqual(getState(), "recording"); + const expected = "const user = getUserById(1);\n// Fix the cache in getUserById."; + await waitFor(() => editor.document.getText() === expected); + assert.match(String(fake.controls[0]?.prompt), /getUserById/); + + await stopDictation(); + assert.strictEqual(getState(), "idle"); + assert.deepStrictEqual(fake.controls.at(-1), { type: "stop" }); + }); + + test("formats a commit subject", async () => { + await useServer({ events: [final(" add retry to the client.")] }); + const editor = await openEditor("", "git-commit"); + await startDictation(); + await waitFor(() => editor.document.getText() === "Add retry to the client"); + }); + + test("replaces the selection and commits pending text on stop", async () => { + await useServer({ events: [partial(" half done")] }); + const editor = await openEditor("before REPLACE after"); + editor.selection = new vscode.Selection(0, 7, 0, 14); + await startDictation(); + await waitFor(() => editor.document.getText() === "before half done after"); + await stopDictation(); + assert.strictEqual(editor.document.getText(), "before half done after"); + }); + + test("user edits and cursor moves commit the pending text", async () => { + const fake = await useServer({ events: [partial(" one")] }); + const editor = await openEditor("Start."); + await startDictation(); + await waitFor(() => editor.document.getText() === "Start. one"); + + // An edit elsewhere keeps "one" and the next segment starts at the cursor. + await editor.edit((builder) => { + builder.insert(new vscode.Position(0, 0), "> "); + }); + fake.send(partial(" two")); + await waitFor(() => editor.document.getText() === "> Start. one two"); + + // Moving the cursor out of the pending range commits it too. + editor.selection = new vscode.Selection(0, 0, 0, 0); + await new Promise((resolve) => setTimeout(resolve, 50)); + fake.send(final(" three")); + await waitFor(() => editor.document.getText() === "Three> Start. one two"); + + // A partial identical to what's there is a no-op. + fake.send(partial("")); + fake.send(partial("")); + fake.send(final(" ")); + await new Promise((resolve) => setTimeout(resolve, 50)); + assert.strictEqual(editor.document.getText(), "Three> Start. one two"); + }); + + test("switching editors commits the pending text", async () => { + const fake = await useServer({ events: [partial(" first")] }); + const first = await openEditor(""); + await startDictation(); + await waitFor(() => first.document.getText() === "first"); + const second = await openEditor(""); + await new Promise((resolve) => setTimeout(resolve, 50)); + fake.send(final(" second")); + await waitFor(() => second.document.getText() === "Second"); + assert.strictEqual(first.document.getText(), "first"); + }); + + test("with no editor, final text is typed into the focused input", async () => { + const fake = await useServer({ events: [partial(" ignored")] }); + await startDictation(); + fake.send(final(" hello there")); + fake.send(final(" ")); + await new Promise((resolve) => setTimeout(resolve, 100)); + assert.strictEqual(getState(), "recording"); + }); + + test("vocabulary can be disabled", async () => { + const fake = await useServer({ events: [] }); + await config().update("vocabulary.enabled", false, vscode.ConfigurationTarget.Global); + try { + await openEditor("getUserById"); + await startDictation(); + await waitFor(() => fake.audioBytes > 0); + await stopDictation(); + assert.deepStrictEqual(fake.controls, [{ type: "stop" }]); + } finally { + await config().update("vocabulary.enabled", undefined, vscode.ConfigurationTarget.Global); + } + }); + + test("passes the model setting to the server", async () => { + const fake = await useServer({ events: [] }); + await config().update("model", "base.en.pt", vscode.ConfigurationTarget.Global); + try { + await startDictation(); + await waitFor(() => fake.controls.length > 0); + assert.strictEqual(fake.controls[0]?.model, "base.en.pt"); + } finally { + await config().update("model", undefined, vscode.ConfigurationTarget.Global); + } + }); + + test("API key from secret storage (CF-2)", async () => { + const fake = await useServer({ apiKey: "s3cret", events: [] }); + const answers = ["s3cret", "", undefined]; + stub("showInputBox", () => Promise.resolve(answers.shift())); + await vscode.commands.executeCommand("whisperCode.setApiKey"); + await startDictation(); + assert.strictEqual(getState(), "recording"); + await stopDictation(); + assert.strictEqual(fake.connections, 1); + + await vscode.commands.executeCommand("whisperCode.setApiKey"); + await vscode.commands.executeCommand("whisperCode.setApiKey"); + const errors: string[] = []; + stub("showErrorMessage", (message: string) => { + errors.push(message); + return Promise.resolve(undefined); + }); + await startDictation(); + assert.strictEqual(getState(), "error"); + assert.deepStrictEqual(errors, ["Whisper Code: Invalid or missing API key"]); + }); + + test("errors show in the status bar and offer Retry (ER-2, ER-4)", async () => { + const fake = await useServer({ events: [] }); + const url = fake.url; + await fake.close(); + server = undefined; + const errors: { message: string; actions: string[] }[] = []; + let retried = false; + stub("showErrorMessage", ((message: string, ...actions: string[]) => { + errors.push({ message, actions }); + const retry = !retried && actions.includes("Retry"); + retried ||= retry; + return Promise.resolve(retry ? "Retry" : undefined); + }) as unknown as Window["showErrorMessage"]); + await config().update("serverUrl", url, vscode.ConfigurationTarget.Global); + await startDictation(); + assert.strictEqual(getState(), "error"); + await waitFor(() => errors.length === 2); + assert.match(errors[0]?.message ?? "", /Can't reach the whisper-flow server/); + assert.deepStrictEqual(errors[0]?.actions, ["Retry"]); + }); + + test("rejects an invalid server URL", async () => { + const errors: string[] = []; + stub("showErrorMessage", (message: string) => { + errors.push(message); + return Promise.resolve(undefined); + }); + await config().update("serverUrl", "ftp://example.com", vscode.ConfigurationTarget.Global); await vscode.commands.executeCommand(TOGGLE); - assert.strictEqual(isRecording(), false); + assert.match(errors[0] ?? "", /must start with http/); + assert.ok(!isRecording()); }); - test("deactivate does not throw", () => { - assert.doesNotThrow(() => { - deactivate(); + test("asks before sending audio to a non-local server (CF-4, CF-5)", async () => { + const prompts: string[] = []; + const reply: { answer?: string } = {}; + stub("showWarningMessage", (message: string) => { + prompts.push(message); + return Promise.resolve(reply.answer); }); + stub("showErrorMessage", () => Promise.resolve(undefined)); + await config().update( + "serverUrl", + "http://whisper.invalid:8181", + vscode.ConfigurationTarget.Global, + ); + + await vscode.commands.executeCommand(TOGGLE); + assert.ok(!isRecording()); + assert.strictEqual(prompts.length, 1); + assert.match(prompts[0] ?? "", /send your microphone audio to whisper\.invalid:8181/); + + reply.answer = "Send Audio"; + await startDictation(); + assert.match(prompts[2] ?? "", /isn't encrypted/); + + // Approved once per URL: only the TLS warning this time. + await startDictation(); + assert.strictEqual(prompts.length, 4); + assert.match(prompts[3] ?? "", /isn't encrypted/); + }); + + test("a double press doesn't start two sessions", async () => { + const fake = await useServer({ events: [] }); + await Promise.all([ + vscode.commands.executeCommand(TOGGLE), + vscode.commands.executeCommand(TOGGLE), + ]); + await waitFor(() => getState() === "recording"); + await new Promise((resolve) => setTimeout(resolve, 50)); + assert.strictEqual(fake.connections, 1); + }); + + test("deactivate ends the session (ER-3)", async () => { + await useServer({ events: [] }); + await startDictation(); + deactivate(); + assert.strictEqual(getState(), "idle"); }); }); diff --git a/src/test/fakeRecorder.ts b/src/test/fakeRecorder.ts new file mode 100644 index 0000000..a7faa31 --- /dev/null +++ b/src/test/fakeRecorder.ts @@ -0,0 +1,60 @@ +// A fake recorder process for tests: an in-memory ChildProcess look-alike. +import { EventEmitter } from "node:events"; +import type { ChildProcess } from "node:child_process"; +import { PassThrough } from "node:stream"; +import type { Spawn } from "../audio/recorder"; + +export class FakeChild extends EventEmitter { + readonly stdout = new PassThrough(); + readonly stderr = new PassThrough(); + exitCode: number | null = null; + signalCode: NodeJS.Signals | null = null; + readonly signals: string[] = []; + /** Exit when sent SIGTERM; when false only SIGKILL ends it. */ + exitOnTerm = true; + + kill(signal: NodeJS.Signals = "SIGTERM"): boolean { + this.signals.push(signal); + if (signal === "SIGKILL" || this.exitOnTerm) { + this.exit(null, signal); + } + return true; + } + + exit(code: number | null, signal: NodeJS.Signals | null = null): void { + this.exitCode = code; + this.signalCode = signal; + this.emit("exit", code, signal); + } + + /** Writes `frames` frames of silence. */ + writeFrames(frames: number): void { + this.stdout.write(Buffer.alloc(2048 * frames)); + } +} + +/** A spawn that fails with ENOENT for commands in `missing` and returns FakeChilds otherwise. */ +export function fakeSpawn(missing: string[] = []): { + spawn: Spawn; + children: FakeChild[]; + commands: string[][]; +} { + const children: FakeChild[] = []; + const commands: string[][] = []; + const spawn: Spawn = (command, args) => { + commands.push([command, ...args]); + const child = new FakeChild(); + children.push(child); + if (missing.includes(command)) { + setImmediate(() => { + child.emit( + "error", + Object.assign(new Error(`spawn ${command} ENOENT`), { code: "ENOENT" }), + ); + child.exit(-2); + }); + } + return child as unknown as ChildProcess; + }; + return { spawn, children, commands }; +} diff --git a/src/test/fakeServer.ts b/src/test/fakeServer.ts new file mode 100644 index 0000000..5300ece --- /dev/null +++ b/src/test/fakeServer.ts @@ -0,0 +1,187 @@ +// A fake whisper-flow server for tests (spec 0002, TS-2). It mimics whisper-flow 1.2.0 over +// real HTTP and WebSocket, including uvicorn's HTTP 403 for rejected handshakes. +import * as http from "node:http"; +import type { AddressInfo } from "node:net"; +import { WebSocketServer, type WebSocket } from "ws"; + +export interface WhisperFlowEvent { + is_partial: boolean; + data: { text: string; segments?: unknown[]; language?: string }; + time: number; +} + +export const partial = (text: string): WhisperFlowEvent => ({ + is_partial: true, + data: { text, segments: [], language: "en" }, + time: 300, +}); +export const final = (text: string): WhisperFlowEvent => ({ + is_partial: false, + data: { text, segments: [], language: "en" }, + time: 300, +}); + +/** + * Events recorded from whisper-flow 1.2.0 (tiny.en) for "Call use memo here. Then fetch the + * user with get user by id." Partials are not monotonic ("your support" → "useMemo"). + */ +export const RECORDED_EVENTS: WhisperFlowEvent[] = [ + partial(" Call you"), + partial(" Call your support"), + partial(" Call useMemo"), + partial(" Call useMemo here."), + final(" Call useMemo here."), + partial(" with"), + partial(" with getUser"), + final(" with getUser by ID."), +]; + +export interface FakeServerOptions { + ready?: Record | "error"; + apiKey?: string; + full?: boolean; + /** Events sent after the first audio frame. */ + events?: (WhisperFlowEvent | { type: "error"; message: string })[]; + /** Sent in reply to a `stop` frame, before closing with 1000. */ + finalOnStop?: WhisperFlowEvent; + /** Ignore `stop` frames (to test the client's timeout). */ + ignoreStop?: boolean; + /** Close with this code after sending the events. */ + closeAfterEvents?: { code: number; reason?: string }; +} + +export const READY_1_2 = { + status: "ok", + version: "1.2.0", + protocol_version: 1, + models: ["tiny.en.pt", "base.en.pt"], + model_loaded: true, + active_sessions: 0, +}; + +export class FakeWhisperFlow { + readonly controls: Record[] = []; + audioBytes = 0; + connections = 0; + private readonly http: http.Server; + private readonly wss = new WebSocketServer({ noServer: true }); + private readonly sockets = new Set(); + + constructor(public options: FakeServerOptions = {}) { + this.http = http.createServer((request, response) => { + const key = request.headers["x-api-key"]; + if (request.url?.endsWith("/ready")) { + if (this.options.ready === "error") { + response.writeHead(500).end(); + return; + } + response.setHeader("content-type", "application/json"); + response.end(JSON.stringify(this.options.ready ?? READY_1_2)); + } else if (request.url === "/transcribe_pcm_chunk" && request.method === "POST") { + const authorized = !this.options.apiKey || key === this.options.apiKey; + response.writeHead(authorized ? 422 : 401).end(); + } else { + response.writeHead(404).end(); + } + }); + this.http.on("upgrade", (request: http.IncomingMessage, socket, head) => { + const key = request.headers["x-api-key"]; + if ((this.options.apiKey && key !== this.options.apiKey) || this.options.full) { + socket.end("HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\n\r\n"); + return; + } + if (request.url !== "/ws") { + socket.end("HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n"); + return; + } + this.wss.handleUpgrade(request, socket, head, (ws) => { + this.accept(ws); + }); + }); + } + + get url(): string { + return `http://127.0.0.1:${String((this.http.address() as AddressInfo).port)}`; + } + + listen(): Promise { + return new Promise((resolve) => + this.http.listen(0, "127.0.0.1", () => { + resolve(this); + }), + ); + } + + /** Sends an event to every open session. */ + send(event: WhisperFlowEvent): void { + for (const socket of this.sockets) { + socket.send(JSON.stringify(event)); + } + } + + /** Server-side close of every open session. */ + closeSessions(code: number, reason = ""): void { + for (const socket of this.sockets) { + socket.close(code, reason); + } + } + + close(): Promise { + for (const socket of this.sockets) { + socket.terminate(); + } + return new Promise((resolve) => + this.http.close(() => { + resolve(); + }), + ); + } + + private accept(ws: WebSocket): void { + this.connections++; + this.sockets.add(ws); + ws.on("close", () => this.sockets.delete(ws)); + let replayed = false; + ws.on("message", (data, isBinary) => { + if (isBinary) { + this.audioBytes += (data as Buffer).length; + if (!replayed) { + replayed = true; + void this.replay(ws); + } + return; + } + const control = JSON.parse((data as Buffer).toString("utf8")) as Record; + this.controls.push(control); + if (control.type === "stop" && !this.options.ignoreStop) { + if (this.options.finalOnStop) { + ws.send(JSON.stringify(this.options.finalOnStop)); + } + ws.close(1000); + } + }); + } + + private async replay(ws: WebSocket): Promise { + for (const event of this.options.events ?? []) { + await new Promise((resolve) => setTimeout(resolve, 5)); + ws.send(JSON.stringify(event)); + } + // Something whisper-flow never sends; clients must ignore it. + ws.send(Buffer.from([1, 2, 3])); + if (this.options.closeAfterEvents) { + ws.close(this.options.closeAfterEvents.code, this.options.closeAfterEvents.reason); + } + } +} + +/** Waits until `check` returns true, polling every 10 ms. */ +export async function waitFor(check: () => boolean, timeoutMs = 3000): Promise { + const deadline = Date.now() + timeoutMs; + while (!check()) { + if (Date.now() > deadline) { + throw new Error("timed out waiting for condition"); + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } +} diff --git a/src/test/recorder.test.ts b/src/test/recorder.test.ts new file mode 100644 index 0000000..2e0e805 --- /dev/null +++ b/src/test/recorder.test.ts @@ -0,0 +1,149 @@ +import * as assert from "assert"; +import { + defaultRecorderCommands, + FrameSplitter, + FRAME_BYTES, + startRecorder, +} from "../audio/recorder"; +import { fakeSpawn } from "./fakeRecorder"; + +function collect() { + const frames: Buffer[] = []; + const errors: string[] = []; + return { + frames, + errors, + handlers: { + onFrame: (frame: Buffer) => frames.push(frame), + onError: (message: string) => errors.push(message), + }, + }; +} + +const tick = () => new Promise((resolve) => setImmediate(resolve)); + +suite("Recorder", () => { + test("FrameSplitter emits 2048-byte frames and keeps the rest (AU-2)", () => { + const splitter = new FrameSplitter(); + assert.strictEqual(splitter.push(Buffer.alloc(1000)).length, 0); + const frames = splitter.push(Buffer.alloc(3200)); + assert.deepStrictEqual( + frames.map((frame) => frame.length), + [FRAME_BYTES, FRAME_BYTES], + ); + assert.strictEqual(splitter.push(Buffer.alloc(1944)).length, 1); + }); + + test("default commands record 16 kHz mono s16le without Python (AU-1, AU-4)", () => { + for (const platform of ["darwin", "linux", "win32"] as const) { + const commands = defaultRecorderCommands(platform); + assert.strictEqual(commands[0]?.[0], "sox"); + for (const command of commands) { + assert.ok(command.includes("16000"), command.join(" ")); + assert.ok(!command.some((part) => part.includes("python"))); + } + } + assert.deepStrictEqual( + defaultRecorderCommands("linux").map((command) => command[0]), + ["sox", "arecord", "ffmpeg"], + ); + }); + + test("streams frames from the sidecar's stdout (AU-3)", async () => { + const { spawn, children } = fakeSpawn(); + const { frames, errors, handlers } = collect(); + const recorder = startRecorder([["rec", "-x"]], handlers, spawn); + children[0]?.writeFrames(3); + await tick(); + assert.strictEqual(frames.length, 3); + recorder.stop(); + recorder.stop(); + children[0]?.writeFrames(1); + await tick(); + assert.strictEqual(frames.length, 3); + assert.deepStrictEqual(children[0]?.signals, ["SIGTERM"]); + assert.deepStrictEqual(errors, []); + }); + + test("falls through to the next command when one isn't installed", async () => { + const { spawn, commands } = fakeSpawn(["sox"]); + const { errors, handlers } = collect(); + startRecorder([["sox"], ["ffmpeg", "-i"]], handlers, spawn).stop(); + await tick(); + assert.deepStrictEqual(commands, [["sox"]]); + + const second = fakeSpawn(["sox"]); + const recorder = startRecorder([["sox"], ["ffmpeg", "-i"]], handlers, second.spawn); + await tick(); + assert.deepStrictEqual(second.commands, [["sox"], ["ffmpeg", "-i"]]); + recorder.stop(); + assert.deepStrictEqual(errors, []); + }); + + test("reports when no recorder is installed", async () => { + const { spawn } = fakeSpawn(["sox", "arecord"]); + const { errors, handlers } = collect(); + startRecorder([["sox"], ["arecord"]], handlers, spawn); + await tick(); + await tick(); + assert.deepStrictEqual(errors, [ + "No audio recorder found (tried sox, arecord). Install sox, or set whisperCode.recorderCommand.", + ]); + + const empty = collect(); + startRecorder([], empty.handlers, spawn); + startRecorder([[]], empty.handlers, spawn); + assert.deepStrictEqual(empty.errors, [ + "No audio recorder configured.", + "No audio recorder configured.", + ]); + }); + + test("uses the real spawn by default", async () => { + const { errors, handlers } = collect(); + startRecorder([["whisper-code-no-such-recorder"]], handlers); + await new Promise((resolve) => setTimeout(resolve, 200)); + assert.match(errors[0] ?? "", /tried whisper-code-no-such-recorder/); + }); + + test("reports spawn errors and unexpected exits", async () => { + const { spawn, children } = fakeSpawn(); + const { errors, handlers } = collect(); + startRecorder([["rec"]], handlers, spawn); + children[0]?.emit("error", Object.assign(new Error("EACCES"), { code: "EACCES" })); + startRecorder([["rec"]], handlers, spawn); + children[1]?.stderr.write("rec: no default audio device"); + await tick(); + children[1]?.exit(2); + startRecorder([["rec"]], handlers, spawn); + children[2]?.exit(null, "SIGSEGV"); + startRecorder([["rec"]], handlers, spawn); + children[3]?.exit(1); + assert.deepStrictEqual(errors, [ + "Audio recorder rec failed: EACCES", + "Audio recorder rec stopped: rec: no default audio device", + "Audio recorder rec stopped: signal SIGSEGV", + "Audio recorder rec stopped: exit code 1", + ]); + }); + + test("ignores errors after stop and force-kills a stuck recorder", async function () { + this.timeout(3000); + const { spawn, children } = fakeSpawn(); + const { errors, handlers } = collect(); + const recorder = startRecorder([["rec"]], handlers, spawn); + const child = children[0]; + assert.ok(child); + child.exitOnTerm = false; + recorder.stop(); + child.emit("error", new Error("late")); + await new Promise((resolve) => setTimeout(resolve, 1100)); + assert.deepStrictEqual(child.signals, ["SIGTERM", "SIGKILL"]); + assert.deepStrictEqual(errors, []); + + const exited = startRecorder([["rec"]], handlers, spawn); + children[1]?.exit(0); + exited.stop(); + assert.deepStrictEqual(children[1]?.signals, []); + }); +}); diff --git a/src/test/server.test.ts b/src/test/server.test.ts new file mode 100644 index 0000000..c82b700 --- /dev/null +++ b/src/test/server.test.ts @@ -0,0 +1,111 @@ +import * as assert from "assert"; +import { + checkReady, + compareVersions, + endpointUrl, + errorMessage, + isInsecureRemote, + isLocalServer, + parseServerUrl, + supportsControlFrames, + toWebSocketUrl, + type Fetch, +} from "../whisperflow/server"; + +const respond = + (body: unknown, status = 200): Fetch => + () => + Promise.resolve({ ok: status < 400, status, json: () => Promise.resolve(body) }); + +const READY = { version: "1.2.0", protocol_version: 1, models: ["tiny.en.pt"], model_loaded: true }; + +suite("whisper-flow server", () => { + test("parseServerUrl rejects non-http URLs", () => { + assert.throws(() => parseServerUrl("not a url"), /Invalid whisper-flow server URL/); + assert.throws(() => parseServerUrl("ws://localhost:8181"), /must start with http/); + assert.strictEqual(parseServerUrl(" http://localhost:8181 ").port, "8181"); + }); + + test("toWebSocketUrl derives the /ws URL (CF-1)", () => { + assert.strictEqual(toWebSocketUrl("http://localhost:8181"), "ws://localhost:8181/ws"); + assert.strictEqual(toWebSocketUrl("https://gpu.example.com/"), "wss://gpu.example.com/ws"); + assert.strictEqual( + toWebSocketUrl("https://example.com/whisper/?x=1#y"), + "wss://example.com/whisper/ws", + ); + }); + + test("endpointUrl appends to the server path", () => { + assert.strictEqual(endpointUrl("http://h:1/base/", "/ready"), "http://h:1/base/ready"); + assert.strictEqual(endpointUrl("http://h:1?q", "/ready"), "http://h:1/ready"); + }); + + test("local and insecure remote servers (CF-4, CF-5)", () => { + for (const url of ["http://localhost:8181", "http://127.0.0.1", "http://[::1]:8181"]) { + assert.ok(isLocalServer(url), url); + assert.ok(!isInsecureRemote(url), url); + } + assert.ok(!isLocalServer("https://gpu.example.com")); + assert.ok(!isInsecureRemote("https://gpu.example.com")); + assert.ok(isInsecureRemote("http://192.168.1.5:8181")); + }); + + test("compareVersions compares numerically", () => { + assert.strictEqual(compareVersions("1.10.0", "1.9.0"), 1); + assert.strictEqual(compareVersions("1.1", "1.1.0"), 0); + assert.strictEqual(compareVersions("1.0.9", "1.1.0"), -1); + assert.strictEqual(compareVersions("1.1.0-rc1", "1.1.0"), 0); + assert.strictEqual(compareVersions("x", "0.0.0"), 0); + }); + + test("checkReady returns server info and sends the API key", async () => { + let headers: Record = {}; + const info = await checkReady("http://localhost:8181", "k", (url, init) => { + assert.strictEqual(url, "http://localhost:8181/ready"); + headers = init.headers; + return respond(READY)(url, init); + }); + assert.deepStrictEqual(headers, { "x-api-key": "k" }); + assert.deepStrictEqual(info, { version: "1.2.0", protocolVersion: 1, models: ["tiny.en.pt"] }); + assert.ok(supportsControlFrames(info)); + }); + + test("checkReady accepts whisper-flow 1.1.0 without protocol_version", async () => { + const info = await checkReady( + "http://localhost:8181", + undefined, + respond({ version: "1.1.0", model_loaded: true }), + ); + assert.strictEqual(info.protocolVersion, undefined); + assert.deepStrictEqual(info.models, []); + assert.ok(!supportsControlFrames(info)); + }); + + test("checkReady errors name the server URL or versions (CF-3, TR-4)", async () => { + const url = "http://localhost:8181"; + const cases: [Fetch, RegExp][] = [ + [ + () => Promise.reject(new Error("ECONNREFUSED")), + /Can't reach .* at http:\/\/localhost:8181 \(ECONNREFUSED\)/, + ], + [respond({}, 503), /Can't reach .*\(HTTP 503\)/], + [respond("ok"), /Unexpected \/ready response .*localhost:8181/], + [respond(null), /Unexpected \/ready response/], + [ + respond({ ...READY, version: "1.0.2" }), + /version 1\.0\.2; whisper-code needs 1\.1\.0 or newer/, + ], + [respond({ ...READY, version: 3 }), /version unknown/], + [respond({ ...READY, protocol_version: 2 }), /protocol 2; this whisper-code supports 1/], + [respond({ ...READY, model_loaded: false }), /localhost:8181 has not loaded a model/], + ]; + for (const [fetchFn, expected] of cases) { + await assert.rejects(checkReady(url, undefined, fetchFn), expected); + } + }); + + test("errorMessage handles non-errors", () => { + assert.strictEqual(errorMessage(new Error("boom")), "boom"); + assert.strictEqual(errorMessage("text"), "text"); + }); +}); diff --git a/src/test/session.test.ts b/src/test/session.test.ts new file mode 100644 index 0000000..1f3921b --- /dev/null +++ b/src/test/session.test.ts @@ -0,0 +1,305 @@ +import * as assert from "assert"; +import { + DictationSession, + type SessionError, + type SessionState, + type TranscriptSink, +} from "../dictation/session"; +import { + classifyRejectedHandshake, + HandshakeError, + WhisperFlowConnection, +} from "../whisperflow/client"; +import { fakeSpawn } from "./fakeRecorder"; +import { + FakeWhisperFlow, + final, + partial, + RECORDED_EVENTS, + waitFor, + type FakeServerOptions, +} from "./fakeServer"; + +class RecordingSink implements TranscriptSink { + readonly calls: string[] = []; + partial(text: string) { + this.calls.push(`partial:${text}`); + return Promise.resolve(); + } + final(text: string) { + this.calls.push(`final:${text}`); + return Promise.resolve(); + } + finish() { + this.calls.push("finish"); + return Promise.resolve(); + } +} + +interface Harness { + server: FakeWhisperFlow; + session: DictationSession; + sink: RecordingSink; + states: SessionState[]; + errors: SessionError[]; + recorder: ReturnType; +} + +let servers: FakeWhisperFlow[] = []; + +async function harness( + serverOptions: FakeServerOptions, + sessionOptions: { apiKey?: string; model?: string; prompt?: string; serverUrl?: string } = {}, +): Promise { + const server = await new FakeWhisperFlow(serverOptions).listen(); + servers.push(server); + const sink = new RecordingSink(); + const states: SessionState[] = []; + const errors: SessionError[] = []; + const recorder = fakeSpawn(); + const session = new DictationSession({ + serverUrl: sessionOptions.serverUrl ?? server.url, + apiKey: sessionOptions.apiKey, + model: sessionOptions.model, + prompt: sessionOptions.prompt, + recorderCommands: [["rec"]], + sink, + spawn: recorder.spawn, + onState: (state, error) => { + states.push(state); + if (error) { + errors.push(error); + } + }, + }); + return { server, session, sink, states, errors, recorder }; +} + +suite("Dictation session", () => { + teardown(async () => { + await Promise.all(servers.map((server) => server.close())); + servers = []; + }); + + test("streams audio and forwards recorded partials and finals in order (TS-2)", async () => { + const h = await harness( + { + events: [...RECORDED_EVENTS, { type: "error", message: "invalid control" }], + finalOnStop: final(" Then stop."), + }, + { prompt: "useMemo, getUserById", model: "base.en.pt" }, + ); + await h.session.start(); + assert.deepStrictEqual(h.states, ["connecting", "recording"]); + h.recorder.children[0]?.writeFrames(2); + await waitFor(() => h.sink.calls.length === RECORDED_EVENTS.length); + assert.deepStrictEqual(h.server.controls[0], { + type: "start", + prompt: "useMemo, getUserById", + model: "base.en.pt", + }); + assert.strictEqual(h.server.audioBytes, 4096); + assert.deepStrictEqual(h.sink.calls.slice(0, 5), [ + "partial: Call you", + "partial: Call your support", + "partial: Call useMemo", + "partial: Call useMemo here.", + "final: Call useMemo here.", + ]); + + await h.session.stop(); + await h.session.stop(); + assert.deepStrictEqual(h.server.controls.at(-1), { type: "stop" }); + assert.deepStrictEqual(h.sink.calls.slice(-2), ["final: Then stop.", "finish"]); + assert.deepStrictEqual(h.states, ["connecting", "recording", "idle"]); + assert.deepStrictEqual(h.recorder.children[0]?.signals, ["SIGTERM"]); + }); + + test("whisper-flow 1.1.0: no control frames, stop closes the socket", async () => { + const h = await harness( + { ready: { version: "1.1.0", model_loaded: true }, events: [partial(" hi")] }, + { prompt: "useMemo" }, + ); + await h.session.start(); + h.recorder.children[0]?.writeFrames(1); + await waitFor(() => h.sink.calls.length === 1); + await h.session.stop(); + assert.deepStrictEqual(h.server.controls, []); + assert.deepStrictEqual(h.sink.calls, ["partial: hi", "finish"]); + }); + + test("sends no start frame when there is nothing to configure", async () => { + const h = await harness({ events: [] }); + await h.session.start(); + h.recorder.children[0]?.writeFrames(1); + await waitFor(() => h.server.audioBytes > 0); + await h.session.stop(); + assert.deepStrictEqual(h.server.controls, [{ type: "stop" }]); + }); + + test("refuses to start when /ready fails (CF-3)", async () => { + const h = await harness({ ready: "error" }); + await h.session.start(); + assert.deepStrictEqual(h.states, ["connecting", "error"]); + assert.match(h.errors[0]?.message ?? "", /Can't reach .*HTTP 500/); + assert.strictEqual(h.errors[0]?.retry, true); + assert.strictEqual(h.server.connections, 0); + assert.strictEqual(h.recorder.children.length, 0); + }); + + test("rejects a model the server doesn't have", async () => { + const h = await harness({}, { model: "large.pt" }); + await h.session.start(); + assert.match( + h.errors[0]?.message ?? "", + /Model large\.pt isn't available .*tiny\.en\.pt, base\.en\.pt/, + ); + const none = await harness({ ready: { version: "1.1.0", model_loaded: true } }, { model: "x" }); + await none.session.start(); + assert.match(none.errors[0]?.message ?? "", /available: none/); + }); + + test("invalid or missing API key (ER-1)", async () => { + const h = await harness({ apiKey: "secret" }, { apiKey: "wrong" }); + await h.session.start(); + assert.deepStrictEqual(h.errors, [{ message: "Invalid or missing API key", retry: false }]); + assert.deepStrictEqual(h.sink.calls, ["finish"]); + }); + + test("server at capacity (ER-1)", async () => { + const h = await harness({ apiKey: "secret", full: true }, { apiKey: "secret" }); + await h.session.start(); + assert.deepStrictEqual(h.errors, [{ message: "Server is at capacity", retry: false }]); + }); + + test("a failed handshake offers Retry", async () => { + const h = await harness({}); + const bad = await harness({}, { serverUrl: `${h.server.url}/nope` }); + await bad.session.start(); + assert.match(bad.errors[0]?.message ?? "", /Can't connect .*\(HTTP 404\)/); + assert.strictEqual(bad.errors[0]?.retry, true); + }); + + test("server close codes stop recording (ER-1, ER-2)", async () => { + const cases: [number, string, SessionError][] = [ + [1013, "", { message: "Server is at capacity", retry: false }], + [ + 4000, + "unknown language: xx", + { message: "Server rejected the session options: unknown language: xx", retry: false }, + ], + [ + 1011, + "", + { message: "Lost connection to the whisper-flow server (code 1011)", retry: true }, + ], + [1000, "", { message: "The whisper-flow server ended the session.", retry: true }], + ]; + for (const [code, reason, expected] of cases) { + const h = await harness({ events: [partial(" kept")], closeAfterEvents: { code, reason } }); + await h.session.start(); + h.recorder.children[0]?.writeFrames(1); + await waitFor(() => h.states.includes("error")); + assert.deepStrictEqual(h.errors, [expected]); + assert.deepStrictEqual(h.sink.calls, ["partial: kept", "finish"]); + assert.deepStrictEqual(h.recorder.children[0]?.signals, ["SIGTERM"]); + } + }); + + test("a recorder failure ends the session with Retry", async () => { + const h = await harness({}); + await h.session.start(); + h.recorder.children[0]?.exit(1); + await waitFor(() => h.states.includes("error")); + assert.deepStrictEqual(h.errors, [ + { message: "Audio recorder rec stopped: exit code 1", retry: true }, + ]); + await h.session.stop(); + assert.deepStrictEqual(h.states, ["connecting", "recording", "error"]); + }); + + test("stopping while connecting cancels the session", async () => { + const h = await harness({}); + const starting = h.session.start(); + await h.session.stop(); + await starting; + assert.deepStrictEqual(h.states, ["connecting", "idle"]); + assert.strictEqual(h.recorder.children.length, 0); + }); + + test("dispose ends a recording session immediately (ER-3)", async () => { + const h = await harness({}); + await h.session.start(); + h.session.dispose(); + assert.deepStrictEqual(h.recorder.children[0]?.signals, ["SIGTERM"]); + await h.session.stop(); + assert.deepStrictEqual(h.states, ["connecting", "recording"]); + }); +}); + +suite("whisper-flow connection", () => { + teardown(async () => { + await Promise.all(servers.map((server) => server.close())); + servers = []; + }); + + test("stop times out when the server ignores it", async () => { + const server = await new FakeWhisperFlow({ ignoreStop: true }).listen(); + servers.push(server); + const closes: number[] = []; + const connection = new WhisperFlowConnection(server.url, undefined, { + onMessage: () => undefined, + onClose: (code) => closes.push(code), + }); + await connection.open(); + const started = Date.now(); + await connection.stop(true, 100); + assert.ok(Date.now() - started < 1000); + await waitFor(() => closes.length === 1); + await connection.stop(true); + connection.dispose(); + connection.sendAudio(Buffer.alloc(2)); + }); + + test("stop without control frames closes normally", async () => { + const server = await new FakeWhisperFlow().listen(); + servers.push(server); + const closes: number[] = []; + const connection = new WhisperFlowConnection(server.url, "k", { + onMessage: () => undefined, + onClose: (code) => closes.push(code), + }); + await connection.open(); + connection.sendStart({}); + await connection.stop(false); + assert.deepStrictEqual(closes, [1000]); + assert.deepStrictEqual(server.controls, [{ type: "start" }]); + await new WhisperFlowConnection(server.url, undefined, { + onMessage: () => undefined, + onClose: () => undefined, + }).stop(true); + }); + + test("connection errors reject open()", async () => { + const server = await new FakeWhisperFlow().listen(); + const url = server.url; + await server.close(); + const connection = new WhisperFlowConnection(url, undefined, { + onMessage: () => undefined, + onClose: () => undefined, + }); + await assert.rejects( + connection.open(), + (error: unknown) => error instanceof HandshakeError && error.closeCode === 0, + ); + }); + + test("classifyRejectedHandshake treats probe failures as capacity", async () => { + assert.strictEqual( + await classifyRejectedHandshake("http://localhost:1", "k", () => + Promise.reject(new Error("down")), + ), + 1013, + ); + }); +}); diff --git a/src/test/text.test.ts b/src/test/text.test.ts new file mode 100644 index 0000000..eeeca24 --- /dev/null +++ b/src/test/text.test.ts @@ -0,0 +1,222 @@ +import * as assert from "assert"; +import { detectTarget, isInComment, type InsertionContext } from "../text/context"; +import { formatFinal, formatPartial, normalize, separatorBefore } from "../text/format"; +import { parseServerMessage, describeClose } from "../whisperflow/messages"; +import { + buildPrompt, + countIdentifiers, + createVocabulary, + isCompoundIdentifier, + rankIdentifiers, + splitIdentifier, +} from "../text/vocabulary"; + +const context = (overrides: Partial): InsertionContext => ({ + languageId: "typescript", + uriScheme: "file", + fileName: "/w/a.ts", + linePrefix: "", + ...overrides, +}); + +suite("Server messages", () => { + test("parses transcripts (RS-1) and server errors", () => { + assert.deepStrictEqual( + parseServerMessage('{"is_partial":true,"data":{"text":" hi","segments":[]},"time":3}'), + { kind: "transcript", isPartial: true, text: " hi" }, + ); + assert.deepStrictEqual(parseServerMessage('{"type":"error","message":"invalid control"}'), { + kind: "error", + message: "invalid control", + }); + assert.deepStrictEqual(parseServerMessage('{"type":"error"}'), { + kind: "error", + message: "unknown error", + }); + }); + + test("ignores anything else", () => { + for (const raw of [ + "nope", + "3", + "null", + '{"is_partial":"yes","data":{}}', + '{"is_partial":true}', + '{"is_partial":true,"data":{"text":1}}', + ]) { + assert.strictEqual(parseServerMessage(raw), undefined, raw); + } + }); + + test("describes close codes (ER-1, ER-2)", () => { + assert.strictEqual(describeClose(1000, ""), undefined); + assert.deepStrictEqual(describeClose(1008, ""), { + message: "Invalid or missing API key", + retry: false, + }); + assert.deepStrictEqual(describeClose(1013, ""), { + message: "Server is at capacity", + retry: false, + }); + assert.match( + describeClose(4000, "unknown language: xx")?.message ?? "", + /unknown language: xx/, + ); + assert.match(describeClose(4000, "")?.message ?? "", /invalid options/); + assert.deepStrictEqual(describeClose(1006, ""), { + message: "Lost connection to the whisper-flow server (code 1006)", + retry: true, + }); + assert.match(describeClose(1011, "oops")?.message ?? "", /code 1011: oops/); + }); +}); + +suite("Vocabulary", () => { + test("splits identifiers", () => { + assert.deepStrictEqual(splitIdentifier("getHTTPServer_v2"), ["get", "http", "server", "v2"]); + assert.deepStrictEqual(splitIdentifier("useMemo"), ["use", "memo"]); + assert.deepStrictEqual(splitIdentifier("MAX_PROMPT_CHARS"), ["max", "prompt", "chars"]); + assert.ok(isCompoundIdentifier("getUserById")); + assert.ok(!isCompoundIdentifier("Promise")); + assert.ok(!isCompoundIdentifier("x2")); + }); + + test("counts and ranks compound identifiers", () => { + const counts = new Map(); + countIdentifiers("const value = useMemo(() => getUserById(id)); useMemo", counts); + countIdentifiers("getUserById", counts, 3); + assert.deepStrictEqual(rankIdentifiers(counts), ["getUserById", "useMemo"]); + assert.deepStrictEqual(rankIdentifiers(counts, 1), ["getUserById"]); + }); + + test("builds a prompt with the best identifiers last, within the limit", () => { + assert.strictEqual(buildPrompt(["aaBb", "ccDd", "eeFf"]), "eeFf, ccDd, aaBb"); + assert.strictEqual(buildPrompt(["aaBb", "ccDd", "eeFf"], 10), "ccDd, aaBb"); + assert.strictEqual(buildPrompt([]), ""); + }); + + test("respells spoken identifiers in text", () => { + const vocabulary = createVocabulary(["getUserById", "useMemo", "get_user_by_id", "userBy"]); + assert.strictEqual(vocabulary.prompt, "userBy, get_user_by_id, useMemo, getUserById"); + assert.strictEqual( + vocabulary.apply("Wrap it in use memo and call Get user by ID, then use-memo again."), + "Wrap it in useMemo and call getUserById, then useMemo again.", + ); + assert.strictEqual(vocabulary.apply("reuse memory"), "reuse memory"); + }); +}); + +suite("Context detector", () => { + test("detects commit messages", () => { + assert.strictEqual(detectTarget(context({ languageId: "git-commit" })), "commit"); + assert.strictEqual(detectTarget(context({ uriScheme: "vscode-scm" })), "commit"); + assert.strictEqual( + detectTarget(context({ languageId: "plaintext", fileName: "/r/.git/COMMIT_EDITMSG" })), + "commit", + ); + }); + + test("detects chat input as a prompt", () => { + assert.strictEqual( + detectTarget(context({ uriScheme: "vscode-chat-input", linePrefix: "// " })), + "prompt", + ); + }); + + test("detects comments", () => { + assert.strictEqual(detectTarget(context({ linePrefix: " // " })), "comment"); + assert.strictEqual(detectTarget(context({ linePrefix: "const a = 1; // " })), "comment"); + assert.strictEqual(detectTarget(context({ linePrefix: " * " })), "comment"); + assert.strictEqual(detectTarget(context({ linePrefix: "/** " })), "comment"); + assert.strictEqual(detectTarget(context({ linePrefix: "x /* " })), "comment"); + assert.strictEqual( + detectTarget(context({ languageId: "python", linePrefix: "x = 1 # " })), + "comment", + ); + assert.strictEqual(detectTarget(context({ languageId: "sql", linePrefix: "-- " })), "comment"); + assert.strictEqual( + detectTarget(context({ languageId: "html", linePrefix: "