feat(dictate): use local socket-activated Parakeet
This commit is contained in:
@@ -4,6 +4,7 @@ Personal extensions and skills for [Pi](https://github.com/earendil-works/pi-mon
|
||||
|
||||
## Included extensions
|
||||
|
||||
- **dictate** — Linux-only local voice dictation through a socket-activated NVIDIA Parakeet service.
|
||||
- **effort-controls** — `/effort` reasoning-level selection and `/fast` priority mode for eligible OpenAI models.
|
||||
- **update** — `/update` and `--update` support for updating Pi and configured Pi packages.
|
||||
|
||||
@@ -13,7 +14,9 @@ Personal extensions and skills for [Pi](https://github.com/earendil-works/pi-mon
|
||||
docs/ Usage guides
|
||||
extensions/ Pi extensions
|
||||
profiles/nono/ Repo-managed nono profiles, symlinked into ~/.config/nono/profiles
|
||||
services/ Host-side services used by extensions
|
||||
skills/ Pi skills (one directory with a SKILL.md per skill)
|
||||
systemd/user/ Repository-managed systemd user units
|
||||
```
|
||||
|
||||
The root `package.json` exposes the extensions and skills as a single Pi package.
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
# Parakeet and NeMo attribution
|
||||
|
||||
Pi Dictate uses NVIDIA's Parakeet-TDT 0.6B v3 model through NVIDIA NeMo. The repository does not redistribute the model weights. The installer downloads or copies the exact snapshot listed below.
|
||||
|
||||
## Model
|
||||
|
||||
- **Name:** `nvidia/parakeet-tdt-0.6b-v3`
|
||||
- **Source:** <https://huggingface.co/nvidia/parakeet-tdt-0.6b-v3>
|
||||
- **Revision:** `541d1f99c6b0c3cd0b11a95167540bb8edefd82b`
|
||||
- **Weight file:** `parakeet-tdt-0.6b-v3.nemo`
|
||||
- **SHA-256:** `3cbdc85877e668ca7b82d0d56770eb1fac76691f55d6b97545e8d61ca588d10d`
|
||||
- **License:** Creative Commons Attribution 4.0 International (CC BY 4.0), as stated by the model card
|
||||
|
||||
Users of the model must retain NVIDIA's attribution and comply with CC BY 4.0. Review the model card for intended use, limitations, and current license notices.
|
||||
|
||||
## Runtime
|
||||
|
||||
- **NVIDIA NeMo:** `nemo-toolkit[asr]` 3.0.0
|
||||
- **Source:** <https://github.com/NVIDIA/NeMo>
|
||||
- **License:** Apache License 2.0
|
||||
- **Locally cached source used during implementation:** OpenSrc package `pypi:nemo-toolkit@3.0.0`
|
||||
- **PyTorch:** 2.13.0 with CUDA 13.0 support
|
||||
- **Python:** CPython 3.12
|
||||
|
||||
The complete transitive dependency versions and artifact hashes are recorded in `services/dictate/uv.lock`. Those dependencies retain their own licenses. Installing Pi Dictate does not change or replace their terms.
|
||||
|
||||
## Primary documentation
|
||||
|
||||
- NeMo ASR documentation: <https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/asr/intro.html>
|
||||
- CC BY 4.0: <https://creativecommons.org/licenses/by/4.0/>
|
||||
- Apache License 2.0: <https://www.apache.org/licenses/LICENSE-2.0>
|
||||
@@ -0,0 +1,48 @@
|
||||
# Dictate Phase 0 results
|
||||
|
||||
Date: 26AUG2026
|
||||
|
||||
## Pinned runtime
|
||||
|
||||
- CPython 3.12.13
|
||||
- NVIDIA NeMo 3.0.0
|
||||
- PyTorch 2.13.0 with CUDA 13.0
|
||||
- NumPy 2.5.2
|
||||
- Model revision `541d1f99c6b0c3cd0b11a95167540bb8edefd82b`
|
||||
- Model SHA-256 `3cbdc85877e668ca7b82d0d56770eb1fac76691f55d6b97545e8d61ca588d10d`
|
||||
- NVIDIA GeForce RTX 4090
|
||||
|
||||
The complete dependency graph and artifact hashes are in `services/dictate/uv.lock`. Exact direct dependency sources were synchronized into `/srv/data/opensrc`.
|
||||
|
||||
## Offline and in-memory validation
|
||||
|
||||
The model loaded from the explicit local `.nemo` archive with Hugging Face, Transformers, and NeMo offline settings enabled. Audio was passed to NeMo as an in-memory 16 kHz mono float array. The benchmark created no audio files.
|
||||
|
||||
| Input | Inference | Real-time factor | Peak CUDA allocation |
|
||||
| --- | ---: | ---: | ---: |
|
||||
| Existing 2.78-second verification speech | 0.304 s | 0.109 | 5.07 GB |
|
||||
| Live 9.96-second microphone speech | 0.285 s | 0.0286 | 5.07 GB |
|
||||
| Repeated 60-second speech stress input | 0.431 s | 0.00718 | 5.18 GB |
|
||||
| Repeated five-minute speech stress input | 0.820 s | 0.00273 | 6.15 GB |
|
||||
|
||||
Cold base-model loading was consistently about 6.4 seconds. The production adapter with its native phrase-boost tree loaded in about 10 seconds and transcribed the known verification fixture exactly. The model archive occupies about 2.4 GB. The prototype environment occupies about 5.6 GB.
|
||||
|
||||
The 60-second and five-minute inputs repeat real speech in memory. They validate latency and bounded GPU behavior, not natural long-form accuracy. Five-minute signed-16 PCM is exactly 9,600,000 bytes; the float32 inference representation is 19,200,000 bytes.
|
||||
|
||||
## Accuracy observations
|
||||
|
||||
The verification sample transcribed exactly:
|
||||
|
||||
> The quick brown fox jumps over the lazy dog.
|
||||
|
||||
A normal live dictation sample was coherent and punctuated, and the user accepted its accuracy.
|
||||
|
||||
The initial code-terminology sample missed `JSON`, `npm`, and `git`. The user failed the original Phase 0 accuracy gate and approved a conservative host-owned coding dictionary using NeMo's native phrase boosting.
|
||||
|
||||
A 37-term dictionary and boost strengths from 0.25 through 2.0 were tested against one in-memory capture. Alpha `1.0` changed the end of the captured phrase from `run np` to `run npm` while preserving the earlier `TypeScript` and `JSON response` terms. Alpha `1.0` was selected as the conservative default. The capture ended before `git status`, so that phrase still requires a focused validation sample.
|
||||
|
||||
## Gate status
|
||||
|
||||
The packaging, latency, memory, offline, and normal-dictation gates pass. Alpha `1.0` produced the same expected clean transcript as baseline, the same degraded transcript on seeded noisy input, and an empty transcript for silence. No dictionary hallucination was observed.
|
||||
|
||||
The coding-term remediation passes the final focused live check through the installed socket-activated service. The spoken phrase “Run npm test, then git status.” produced exactly `Run npm test, then git status.` with native phrase boosting at alpha `1.0`.
|
||||
@@ -0,0 +1,149 @@
|
||||
# Pi Dictate local-service protocol
|
||||
|
||||
## Transport and version
|
||||
|
||||
Protocol version `1` uses UTF-8, newline-delimited JSON over one persistent Unix stream socket. Each line contains one JSON object. The systemd socket unit owns the socket path and permissions. The service accepts the inherited listening file descriptor and does not bind or unlink the path.
|
||||
|
||||
Clients must send `hello` before any other request. A connection negotiates one version only. Coding terms are host-owned service configuration and never cross this protocol.
|
||||
|
||||
## Limits
|
||||
|
||||
| Data | Limit |
|
||||
| --- | ---: |
|
||||
| JSON frame, excluding newline | 65,536 bytes |
|
||||
| `requestId` or `generationId` | 128 UTF-8 bytes |
|
||||
| Final transcript | 32,768 UTF-8 bytes |
|
||||
| Error detail | 2,048 UTF-8 bytes |
|
||||
| Simultaneous client connections | 16 |
|
||||
| Pending output per client | 65,536 bytes |
|
||||
| Recorder stderr tail | 65,536 bytes |
|
||||
| Mono 16 kHz signed-16 PCM | 9,600,000 bytes, five minutes |
|
||||
|
||||
IDs are non-empty, opaque strings generated by the client. The service compares them only for request correlation, ownership, and stale-generation rejection. It does not interpret IDs as paths. Request objects reject unknown or missing fields.
|
||||
|
||||
Malformed JSON, a frame over the limit, or an unsupported version produces one bounded `error` frame and then closes the connection. The malformed-frame error uses an empty `requestId` when no client ID can be recovered. Schema, state, busy, and ownership errors leave the connection open.
|
||||
|
||||
## Client requests
|
||||
|
||||
All examples are one line when transmitted.
|
||||
|
||||
### `hello`
|
||||
|
||||
```json
|
||||
{"version":1,"type":"hello","requestId":"hello-1"}
|
||||
```
|
||||
|
||||
Exact fields: `version`, `type`, `requestId`.
|
||||
|
||||
A successful response is `pong` with the same `requestId`. `ready` reports whether the model is warm.
|
||||
|
||||
### `ping`
|
||||
|
||||
```json
|
||||
{"version":1,"type":"ping","requestId":"ping-1"}
|
||||
```
|
||||
|
||||
Exact fields: `version`, `type`, `requestId`. `hello` is required first. The response is `pong` and may include `ready` and `busy`.
|
||||
|
||||
### `start`
|
||||
|
||||
```json
|
||||
{"version":1,"type":"start","requestId":"start-1","generationId":"generation-1"}
|
||||
```
|
||||
|
||||
Exact fields: `version`, `type`, `requestId`, `generationId`. The service accepts only one global generation. It first emits `state: loading` for a cold model or `state: listening` for a warm model. Recording starts only after model loading succeeds.
|
||||
|
||||
### `stop`
|
||||
|
||||
```json
|
||||
{"version":1,"type":"stop","requestId":"stop-1","generationId":"generation-1"}
|
||||
```
|
||||
|
||||
Exact fields: `version`, `type`, `requestId`, `generationId`. Only the owning connection may stop the active listening generation.
|
||||
|
||||
### `cancel`
|
||||
|
||||
```json
|
||||
{"version":1,"type":"cancel","requestId":"cancel-1","generationId":"generation-1"}
|
||||
```
|
||||
|
||||
Exact fields: `version`, `type`, `requestId`, `generationId`. A valid cancellation immediately emits `cancelled`. Cancellation tombstones the generation, so it can never emit a later `result` or begin recording after a cancelled load.
|
||||
|
||||
## Service events
|
||||
|
||||
All service events contain `version`, `type`, and a bounded client-generated `requestId`. Generation events also contain `generationId`.
|
||||
|
||||
### `pong`
|
||||
|
||||
```json
|
||||
{"version":1,"type":"pong","requestId":"hello-1","service":"pi-dictate","ready":false}
|
||||
```
|
||||
|
||||
### `state`
|
||||
|
||||
```json
|
||||
{"version":1,"type":"state","requestId":"start-1","generationId":"generation-1","state":"loading"}
|
||||
```
|
||||
|
||||
`state` is one of `loading`, `listening`, `stopping`, `transcribing`, or `idle`.
|
||||
|
||||
### `result`
|
||||
|
||||
```json
|
||||
{"version":1,"type":"result","requestId":"stop-1","generationId":"generation-1","transcript":"example text"}
|
||||
```
|
||||
|
||||
There is at most one result for a generation. An empty transcript is a successful result with `"transcript":""`.
|
||||
|
||||
### `cancelled`
|
||||
|
||||
```json
|
||||
{"version":1,"type":"cancelled","requestId":"cancel-1","generationId":"generation-1"}
|
||||
```
|
||||
|
||||
This is a terminal delivery guarantee: the generation cannot later produce a result.
|
||||
|
||||
### `error`
|
||||
|
||||
```json
|
||||
{"version":1,"type":"error","requestId":"stop-1","generationId":"generation-1","code":"recorder_timeout","detail":"pw-record did not stop after SIGTERM and was killed"}
|
||||
```
|
||||
|
||||
Stable codes include:
|
||||
|
||||
- framing and schema: `malformed_json`, `frame_too_large`, `unsupported_version`, `invalid_request`;
|
||||
- session: `hello_required`, `busy`, `not_owner`, `generation_mismatch`, `invalid_state`;
|
||||
- recorder and audio: `recorder_spawn`, `recorder_exit`, `recorder_timeout`, `duration_limit`, `invalid_audio`;
|
||||
- model and output: `model_failure`, `transcript_too_large`.
|
||||
|
||||
## Lifecycle and ownership
|
||||
|
||||
| Current state | Request or event | Outcome |
|
||||
| --- | --- | --- |
|
||||
| no generation | owner `start` | `loading`, then `listening`; warm models may emit `listening` directly |
|
||||
| any active generation | another `start` | `busy`; connection stays open |
|
||||
| `loading` | owner `cancel` | immediate `cancelled`; loading finishes in the worker but recording is suppressed |
|
||||
| `listening` | owner `stop` | `stopping`; signal recorder, drain stdout to EOF, await child close, convert PCM, then `transcribing` |
|
||||
| `listening` | owner `cancel` or disconnect | stop and drain recorder, discard audio, never infer |
|
||||
| `stopping` | owner `cancel` | immediate `cancelled`; completed audio is discarded |
|
||||
| `transcribing` | owner `cancel` or disconnect | tombstone immediately; inference finishes but its result is discarded |
|
||||
| `transcribing` | successful model result | one `result`, then `idle` |
|
||||
| any work state | terminal failure | one `error`, then `idle` |
|
||||
| active generation | non-owner `stop` or `cancel` | `not_owner`; connection stays open |
|
||||
| active generation | stale generation ID | `generation_mismatch`; connection stays open |
|
||||
|
||||
Cancelled model loading and inference are non-interruptible. The global service remains busy until the worker finishes even though `cancelled` was already emitted. Disconnecting a generation owner cancels capture or tombstones loading/inference under the same rules.
|
||||
|
||||
The recorder command is exactly:
|
||||
|
||||
```text
|
||||
pw-record --raw --format s16 --rate 16000 --channels 1 -
|
||||
```
|
||||
|
||||
The service launches `pw-record` in its own process group and marks the generation stopping before SIGTERM. It continuously drains bounded stdout and stderr, waits a bounded time for pipe EOF and child close, and only then converts or transcribes. A shutdown timeout escalates to a process-group kill and fails the generation. PipeWire's exit code `1` is accepted only after this service intentionally signalled a still-running recorder and stderr remained empty; an already-exited process or stderr-bearing exit remains a failure. Raw audio and recorder logs are held only in bounded memory and are not written to disk.
|
||||
|
||||
## Idle exit
|
||||
|
||||
A successfully parsed protocol request resets the 15-minute monotonic idle activity timer, including requests that later produce state or ownership errors. Accepting a socket, receiving malformed data, and disconnecting do not reset activity.
|
||||
|
||||
Idle exit is suspended while model loading, recording, recorder shutdown, PCM conversion, or inference is active. Once work becomes idle, the full timer starts again. Exiting closes accepted clients and the inherited listener but does not unlink the systemd-owned socket path.
|
||||
@@ -0,0 +1,444 @@
|
||||
# Dictate extension review and local Parakeet remediation plan
|
||||
|
||||
Status: implemented, with final host validation in progress. This document records the original review of `extensions/dictate/index.ts`, the agreed Linux implementation direction, and the acceptance criteria used for the replacement.
|
||||
|
||||
## Decisions
|
||||
|
||||
The replacement will target this machine and workflow:
|
||||
|
||||
- **Platform:** Linux only, specifically CachyOS on x86-64 with niri and Ghostty.
|
||||
- **Transcription:** local NVIDIA NeMo with `nvidia/parakeet-tdt-0.6b-v3` on the machine's RTX 4090. Deepgram and runtime cloud transcription are removed.
|
||||
- **Process boundary:** a Python transcription service runs outside Pi's nono sandbox. The Pi extension communicates with it only through a Unix-domain socket.
|
||||
- **Microphone:** the local service owns PipeWire capture. Pi receives no PipeWire socket, `/dev/snd`, or GPU access.
|
||||
- **Interaction:** push to start, push to stop, then transcribe the completed recording. Voice activity detection, partial hypotheses, and continuous dictation are out of scope.
|
||||
- **Activation:** shortcuts work while Pi's main editor has focus. A compositor-global dictation service and shortcut handling inside Pi modal dialogs are out of scope.
|
||||
- **Delivery:** completed text is inserted into Pi's main editor with the public `ctx.ui.pasteToEditor()` API. Do not inspect focused TUI components or synthesize input.
|
||||
- **Service lifecycle:** a systemd user socket activates the service on first use. The service keeps the model warm for 15 minutes after the last activity, then exits and releases GPU memory.
|
||||
- **Sandbox:** `profiles/nono/pi-local.json` gains access only to the exact transcription-service socket. It does not gain direct PipeWire, GPU, model-directory, or Deepgram access.
|
||||
- **Mediated capability:** access to the socket is access to request microphone recording and receive transcripts. The user explicitly accepts that any code running in a profile that inherits `pi-local` can invoke this protocol.
|
||||
- **Privacy:** the service intentionally creates no audio files and releases in-memory audio after transcription, cancellation, or failure. This is best-effort non-persistence, not a guarantee against swap, core dumps, driver copies, or forensic recovery.
|
||||
|
||||
These decisions isolate native audio and model execution from Pi, but the socket remains a mediated microphone capability. The agent sandbox cannot open PipeWire or CUDA devices directly; it can ask the host service to record through the documented protocol.
|
||||
|
||||
## Target environment
|
||||
|
||||
The following was observed on this machine during planning:
|
||||
|
||||
- CachyOS with `XDG_SESSION_TYPE=wayland` and `XDG_CURRENT_DESKTOP=niri`.
|
||||
- niri `26.04` and Ghostty with `TERM=xterm-256color`.
|
||||
- An NVIDIA GeForce RTX 4090 and 188 GiB of system memory.
|
||||
- PipeWire `pw-record` is installed and linked against PipeWire `1.6.8`.
|
||||
- `wl-copy` is installed, but clipboard delivery is excluded.
|
||||
- SoX, `wtype`, and `ydotool` are not installed.
|
||||
- The user is UID `1000`, with `XDG_RUNTIME_DIR=/run/user/1000`.
|
||||
- The current Pi sandbox cannot initialize NVML or connect to PipeWire. That is acceptable because the service will run outside it.
|
||||
- The installed niri default configuration uses `Mod+M` for maximizing a window. In a native niri session, `Mod` is Super, so this does not conflict with Pi's `Alt+M`. The user's live niri configuration must still be checked for custom `Alt+M` or `Alt+N` bindings.
|
||||
|
||||
## External runtime choice
|
||||
|
||||
NVIDIA NeMo is the canonical Python runtime for Parakeet. The selected v3 model is multilingual and published under CC-BY-4.0. NeMo is Apache-2.0. The implementation must retain model attribution, record the exact model revision, and comply with both licenses.
|
||||
|
||||
Parakeet v3 is not treated as a drop-in replacement for Deepgram Live. This design intentionally does not require stable partial transcripts or cache-aware streaming compatibility. The service records one bounded clip and invokes NeMo's offline transcription API after the user stops recording.
|
||||
|
||||
Primary references:
|
||||
|
||||
- NVIDIA Parakeet-TDT 0.6B v3 model card: <https://huggingface.co/nvidia/parakeet-tdt-0.6b-v3>
|
||||
- NVIDIA NeMo repository: <https://github.com/NVIDIA/NeMo>
|
||||
- NeMo ASR documentation: <https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/asr/intro.html>
|
||||
|
||||
The model and Python/CUDA dependency set are too large and volatile to approve from documentation alone. Phase 0 is therefore a throwaway prototype and benchmark gate. Do not alter the production extension until that gate passes.
|
||||
|
||||
## Review scope and verification
|
||||
|
||||
The review covered Pi and TUI integration, recorder and WebSocket lifecycle, transcript delivery, failure handling, Linux and Wayland portability, nono integration, test coverage, and repository typechecking.
|
||||
|
||||
Historical verification from the original review:
|
||||
|
||||
- `npm run typecheck` failed at `extensions/dictate/index.ts:386`.
|
||||
- The extension had no automated tests.
|
||||
|
||||
Static revalidation before implementation confirmed:
|
||||
|
||||
- The original implementation defects were present in `extensions/dictate/index.ts`.
|
||||
- Dictate had no automated test suite or README, and the root `npm test` command excluded it.
|
||||
- `extensions/dictate/LICENSE` was already deleted in the working tree; the user approved preserving that deletion.
|
||||
- `profiles/nono/pi-local.json` did not yet allow the planned local service socket.
|
||||
- This historical review section describes the starting point, not the current implementation state.
|
||||
|
||||
## Findings
|
||||
|
||||
### 1. The clipboard fallback can crash Pi and does not work on Linux
|
||||
|
||||
**Severity:** high
|
||||
|
||||
`extensions/dictate/index.ts:263-268` starts `pbcopy` without an `error` listener and immediately reports success. Spawn failures arrive asynchronously, outside the surrounding `try/catch`, and an unhandled child-process error can terminate Pi. `pbcopy` is also macOS-only.
|
||||
|
||||
**Fix:** remove clipboard delivery completely. Insert successful transcripts through `ctx.ui.pasteToEditor()` and notify the user. This public API preserves the editor's cursor and collapsed-paste bookkeeping, unlike reconstructing editor contents through `getEditorText()` and `setEditorText()`.
|
||||
|
||||
### 2. Recording starts before the transcription backend is ready
|
||||
|
||||
**Severity:** high
|
||||
|
||||
The current recorder and listening UI start before Deepgram connects (`index.ts:300-387`). A quick stop or slow connection can discard speech while the UI claimed to be listening.
|
||||
|
||||
**Fix:** replace the cloud connection with a versioned local-service handshake. On a cold start, show `loading model...` while systemd activates the service and NeMo loads. Start `pw-record` only after the service reports ready, then show `listening...`. Stopping or cancelling during loading must prevent recorder startup even if the model finishes loading later.
|
||||
|
||||
### 3. The extension fails repository typechecking
|
||||
|
||||
**Severity:** high
|
||||
|
||||
`index.ts:386` passes a Node `Buffer<ArrayBufferLike>` directly to the DOM-compatible `WebSocket.send()` API, which TypeScript rejects.
|
||||
|
||||
**Fix:** delete the Deepgram WebSocket and PCM-send path. The extension exchanges bounded newline-delimited JSON frames with the local service and never handles audio buffers. Keep `npm run typecheck` as a mandatory gate.
|
||||
|
||||
### 4. Transcript targeting can misroute or lose text
|
||||
|
||||
**Severity:** medium
|
||||
|
||||
`resolveTarget()` accepts every focused component with `handleInput()` as typable (`index.ts:209-226`). `flush()` sends the transcript as one synthetic keyboard event and assumes delivery succeeded (`index.ts:247-262`). Selectors can ignore the string or interpret it as navigation.
|
||||
|
||||
**Fix:** delete focused-target resolution and synthetic input. Successful results always go to Pi's main editor through `pasteToEditor()`.
|
||||
|
||||
### 5. Dictation depends on private TUI state and input hooks
|
||||
|
||||
**Severity:** medium
|
||||
|
||||
The extension captures the concrete TUI through an invisible widget, stores it as `any`, reads private `focusedComponent`, and installs a global input listener (`index.ts:90,209-215,488-522`).
|
||||
|
||||
**Fix:** remove the invisible widget, TUI handle, global input listener, target adapter, and related `any` values. Register `Alt+M` and `Alt+N` with `pi.registerShortcut()` and add `/dictate-cancel`.
|
||||
|
||||
Installed Pi routes extension shortcuts through the main editor, not modal dialogs. The supported contract is therefore explicit:
|
||||
|
||||
- `Alt+M`, `Alt+N`, and `/dictate-cancel` work while the main editor has focus.
|
||||
- They do not start or cancel dictation while a selector, confirmation, or custom modal owns input.
|
||||
- If a modal opens after recording starts, a later successful result still goes to the main editor.
|
||||
- The registered keys are reserved while the main editor is focused, including while dictate is idle.
|
||||
|
||||
### 6. Unexpected recorder exits can leave dictation active
|
||||
|
||||
**Severity:** medium
|
||||
|
||||
The current recorder exit handler cleans up only for non-null, nonzero codes (`index.ts:352-362`). Zero exits and signal exits can leave state active.
|
||||
|
||||
**Fix:** move recorder ownership into the local service. Every `pw-record` termination is classified as intentional or unexpected before handling its exit code and signal. The service captures a bounded stderr tail, emits one structured error, discards buffered audio, and returns to idle. A service crash or socket close makes the extension return to idle without crashing Pi.
|
||||
|
||||
### 7. Audio and transport buffering are unbounded
|
||||
|
||||
**Severity:** medium
|
||||
|
||||
The current extension does not inspect WebSocket backpressure (`index.ts:380-386`). A local service would introduce a different risk if it buffered microphone audio without a duration cap.
|
||||
|
||||
**Fix:** remove network audio transport and bound recording inside the service. Start with a five-minute maximum, equivalent to approximately 9.6 MB of mono 16 kHz signed 16-bit PCM. At the limit, stop capture, discard the clip, report a clear duration-limit error, and return to idle. Keep protocol frames and transcript size bounded independently.
|
||||
|
||||
### 8. Cleanup implicitly decides whether to deliver text
|
||||
|
||||
**Severity:** medium
|
||||
|
||||
Current `cleanup()` calls `flush()` before teardown, so recorder and socket failures can insert fragments without identifying them as incomplete (`index.ts:271-298,346-350,407-414`).
|
||||
|
||||
**Fix:** cleanup becomes idempotent and delivery-neutral. Because local Parakeet returns one offline result, terminal outcomes are simpler:
|
||||
|
||||
| Event | Outcome | Editor change |
|
||||
| --- | --- | --- |
|
||||
| Service reports a non-empty result for the active generation | finalized | Paste once into the main editor |
|
||||
| Service reports an empty result | empty | None; notify that no speech was recognized |
|
||||
| User cancels before inference | cancelled | None; service discards audio |
|
||||
| User cancels during inference | cancelled | None; tombstone the generation and suppress its eventual result while the non-interruptible worker finishes |
|
||||
| Recorder, model, protocol, service, or timeout failure | failed | None; service discards audio and reports the cause |
|
||||
| Pi session shuts down | cancelled | None; best-effort cancel or tombstone, then close the socket |
|
||||
| Result belongs to an old generation | stale | None; ignore it |
|
||||
|
||||
No path may deliver text from cleanup itself.
|
||||
|
||||
### 9. The recorder and setup instructions are macOS-specific
|
||||
|
||||
**Severity:** high for this machine
|
||||
|
||||
The extension spawns SoX's `rec` and recommends `brew install sox` (`index.ts:24,319-348`). SoX is absent on this CachyOS machine, while `pw-record` is installed.
|
||||
|
||||
**Fix:** the local service spawns exactly:
|
||||
|
||||
```text
|
||||
pw-record --raw --format s16 --rate 16000 --channels 1 -
|
||||
```
|
||||
|
||||
The service pipes stdout, bounds stderr, and lets PipeWire select the WirePlumber default input. It must set state to stopping before signaling the recorder, continue reading stdout until the stream ends, and wait for the child `close` event before decoding. A timeout escalates process termination and produces a failed outcome. This ordering prevents the final PCM chunks from being truncated.
|
||||
|
||||
The service validates even-byte PCM length, converts little-endian signed 16-bit samples to the float representation expected by NeMo, and never creates an audio file. A selectable microphone target is future work.
|
||||
|
||||
### 10. The current sandbox boundary is too broad for in-process transcription
|
||||
|
||||
**Severity:** high for this machine
|
||||
|
||||
Running capture and NeMo inside Pi would require PipeWire, model-directory, native-library, CUDA, and `/dev/nvidia*` access in the agent sandbox. Deepgram also requires cloud credentials and a network destination.
|
||||
|
||||
**Fix:** keep direct device and model capabilities in the host-side service. Add only `/run/user/1000/pi-dictate.sock` to `filesystem.unix_socket` in `profiles/nono/pi-local.json`. Do not add PipeWire configuration, `/run/user/1000/pipewire-0`, `/dev/snd`, `/dev/nvidia*`, model paths, Hugging Face cache paths, or `api.deepgram.com`.
|
||||
|
||||
The socket is deliberately a mediated microphone and transcription capability, not a harmless health endpoint. Any process running under `pi-local` or a descendant profile can request recording and receive the resulting text. The user accepts that broad profile scope; document it prominently and keep the service single-user with socket mode `0600`.
|
||||
|
||||
The resolved profile already contains inherited and local capabilities. Acceptance must verify that this feature adds only the one service socket relative to the previous resolved policy, not claim that the full profile grants only one capability.
|
||||
|
||||
### 11. Status rendering performs unnecessary continuous updates
|
||||
|
||||
**Severity:** low
|
||||
|
||||
The current meter updates every 60 ms even when its rendered value has not changed, and the finalizing spinner updates every 80 ms (`index.ts:97-105,172-207`).
|
||||
|
||||
**Fix:** the extension renders state changes reported by the service and does not receive raw audio levels. Use static status text for `loading model...`, `listening...`, `stopping...`, and `transcribing...`. If a meter is wanted later, add a rate-limited level event only after measuring the need.
|
||||
|
||||
### 12. Tests, installation, attribution, and user documentation are missing
|
||||
|
||||
**Severity:** medium
|
||||
|
||||
There is no dictate test suite or README, and the root test command excludes dictate. The repository also has no service package, model provisioning flow, systemd unit, protocol documentation, or Parakeet attribution.
|
||||
|
||||
**Fix:** add focused TypeScript and Python tests, a concise extension README, a service installation guide, and model/license attribution. Add dictate to the root README extension inventory because the package exposes the extensions directory.
|
||||
|
||||
## Local service architecture
|
||||
|
||||
### Repository layout
|
||||
|
||||
The implementation should introduce one explicit host-service boundary:
|
||||
|
||||
```text
|
||||
extensions/dictate/
|
||||
index.ts
|
||||
index.test.ts
|
||||
README.md
|
||||
services/dictate/
|
||||
pyproject.toml
|
||||
uv.lock
|
||||
src/pi_dictate_service/
|
||||
__init__.py
|
||||
audio.py
|
||||
model.py
|
||||
protocol.py
|
||||
server.py
|
||||
tests/
|
||||
systemd/user/
|
||||
pi-dictate.socket
|
||||
pi-dictate.service
|
||||
docs/dictate/
|
||||
protocol.md
|
||||
parakeet-attribution.md
|
||||
```
|
||||
|
||||
Use a dedicated uv environment for the service. Pin Python, NeMo, PyTorch, CUDA-compatible wheels, and transitive dependencies in `uv.lock`. Follow the repository package-source process when adding dependencies and synchronize relevant third-party source into the OpenSrc cache.
|
||||
|
||||
Install the service entry point and environment outside the repository under `~/.local/share/pi-dictate/`. Store the pinned model snapshot under `~/.local/share/pi-dictate/models/parakeet-tdt-0.6b-v3/`. The runtime must load that explicit local snapshot with offline mode enabled rather than resolving a floating model name.
|
||||
|
||||
### Socket activation and isolation
|
||||
|
||||
`systemd/user/pi-dictate.socket` should:
|
||||
|
||||
- listen at `%t/pi-dictate.sock`, which resolves to `/run/user/1000/pi-dictate.sock` on this machine;
|
||||
- create the socket with mode `0600`;
|
||||
- remove the socket only when the socket unit itself stops;
|
||||
- remain active while the service exits on idle;
|
||||
- activate one non-template service that accepts the inherited listening file descriptor.
|
||||
|
||||
The service must close its inherited descriptor on exit but never unlink `%t/pi-dictate.sock`; the socket unit owns that pathname and reactivation lifecycle.
|
||||
|
||||
`systemd/user/pi-dictate.service` should:
|
||||
|
||||
- run the installed Python entry point outside nono;
|
||||
- inherit the listening socket rather than opening TCP;
|
||||
- load the pinned model onto CUDA and fail clearly when CUDA is unavailable;
|
||||
- own PipeWire capture and the in-memory audio buffer;
|
||||
- serialize model loading and inference and allow only one active recording or transcription globally;
|
||||
- reject start requests as busy while another client owns capture or while a cancelled inference worker is still running;
|
||||
- cancel capture if the owning client disconnects;
|
||||
- tombstone cancellation during NeMo inference, suppress the eventual result, and keep the service alive until the worker finishes;
|
||||
- never unload the model or perform idle exit while a loading or inference worker is active;
|
||||
- exit after 15 minutes without recording, transcription, loading, or protocol activity;
|
||||
- release recorder, model, CUDA, accepted-client, and inherited-listener resources on shutdown without unlinking the socket-unit pathname;
|
||||
- disable runtime network access through supported systemd hardening after the local model-loading path is proven;
|
||||
- set `LimitCORE=0` and, where supported for the user service, `MemorySwapMax=0`;
|
||||
- expose no raw audio or transcript logs by default.
|
||||
|
||||
Model loading must not block protocol cancellation. Load NeMo in a worker while the socket loop remains responsive. Cancelling during loading must emit `cancelled` promptly and tombstone that generation. The worker may finish warming the model, but it must never start recording for the cancelled generation.
|
||||
|
||||
### Versioned protocol
|
||||
|
||||
Use bounded newline-delimited JSON over one persistent Unix stream connection. Document protocol version `1` before implementation.
|
||||
|
||||
Client requests:
|
||||
|
||||
- `hello`: negotiate protocol version and service capabilities;
|
||||
- `start`: create a generation and request model readiness plus recording;
|
||||
- `stop`: stop and transcribe the caller's active generation;
|
||||
- `cancel`: discard the caller's active generation;
|
||||
- `ping`: health and readiness check.
|
||||
|
||||
Service events:
|
||||
|
||||
- `state`: `loading | listening | stopping | transcribing | idle`;
|
||||
- `result`: one final transcript for a generation;
|
||||
- `error`: stable code plus bounded human-readable detail;
|
||||
- `cancelled`: confirmation that a generation cannot later produce a result;
|
||||
- `pong`: service and protocol version information.
|
||||
|
||||
Every frame includes protocol version, request or generation ID, and type. Set strict frame and transcript byte limits. Reject malformed JSON, unknown fields where unsafe, unsupported versions, invalid state transitions, and commands from a client that does not own the active generation. Never use filesystem paths supplied by the Pi client.
|
||||
|
||||
### Authoritative lifecycle
|
||||
|
||||
The extension state machine is:
|
||||
|
||||
```text
|
||||
idle -> connecting -> recording -> stopping -> transcribing -> idle
|
||||
```
|
||||
|
||||
Cancellation and failure can return any active extension state to idle. State transitions are authoritative only when received for the current connection and generation. A cancelled loading or inference worker may remain active inside the service after the extension returns to idle; it owns no future delivery right, blocks new starts as busy until it finishes, and prevents service idle exit.
|
||||
|
||||
The service stop sequence is mandatory:
|
||||
|
||||
1. Mark the generation stopping so recorder exit is intentional.
|
||||
2. Signal `pw-record`.
|
||||
3. Continue consuming stdout.
|
||||
4. Wait for stdout end and the child `close` event.
|
||||
5. Enforce a short shutdown timeout and kill on overrun.
|
||||
6. Validate and convert the complete PCM buffer.
|
||||
7. Mark the generation transcribing.
|
||||
8. Run one NeMo transcription job.
|
||||
9. If the generation was tombstoned during inference, discard the result and emit or preserve the earlier `cancelled` outcome; otherwise emit exactly one result or error.
|
||||
10. Release application-owned audio buffers and return to idle.
|
||||
|
||||
A clean recorder exit alone is not transcription success. Only a `result` event for the active, non-tombstoned generation authorizes editor insertion. Buffer release is best effort because NumPy, PyTorch, CUDA, swap, and crash handling can create copies outside the service's direct control.
|
||||
|
||||
## niri and Wayland integration
|
||||
|
||||
No niri configuration change is required. niri should pass `Alt+M` and `Alt+N` to Ghostty and Pi while the main editor is focused, provided the user's config does not bind those combinations.
|
||||
|
||||
During implementation, inspect `~/.config/niri/config.kdl` outside the current restrictive sandbox and verify:
|
||||
|
||||
1. `Alt+M` and `Alt+N` are not compositor bindings.
|
||||
2. `niri validate` passes before and after any user-performed keybinding adjustment.
|
||||
3. Both shortcuts reach Pi's main editor in Ghostty.
|
||||
4. The shortcuts do not claim to work while a Pi modal dialog owns input.
|
||||
|
||||
Do not add niri `spawn` bindings. Do not use `wl-copy`, `wtype`, `ydotool`, `xclip`, or `xdotool`.
|
||||
|
||||
## Recommended implementation plan
|
||||
|
||||
Each behavior change starts with a focused failing test in the same implementation slice. Do not postpone regression tests until after all fixes.
|
||||
|
||||
### Phase 0: Prove NeMo and Parakeet v3 outside Pi
|
||||
|
||||
Build a throwaway prototype outside the production extension and nono sandbox:
|
||||
|
||||
1. Create an isolated uv environment with a pinned Python version.
|
||||
2. Install a CUDA-compatible PyTorch and NeMo stack using the package-source workflow.
|
||||
3. Download and pin one exact `nvidia/parakeet-tdt-0.6b-v3` model revision outside runtime.
|
||||
4. Record short, realistic dictation samples through `pw-record` only after explicit user consent.
|
||||
5. Transcribe in-memory 16 kHz mono audio without creating an audio file.
|
||||
6. Measure model download size, cold load time, warm transcription latency, peak GPU memory, and real-time factor for 10-second, 60-second, and five-minute clips.
|
||||
7. Test punctuation, code terminology, silence, muted input, and at least one realistic noisy sample.
|
||||
8. Confirm that the model works with runtime network access disabled after provisioning.
|
||||
9. Record dependency versions, model revision, hashes where available, and license obligations.
|
||||
|
||||
**Gate:** proceed only if warm transcription is comfortably faster than real time, cold-start latency is acceptable for socket activation, five-minute memory use is bounded, accuracy is acceptable for normal coding dictation, and the pinned offline load is reproducible. If the gate fails, revise the runtime decision with the user rather than silently substituting sherpa-onnx, another model, or cloud transcription.
|
||||
|
||||
### Phase 1: Define and test the service core
|
||||
|
||||
1. Write `docs/dictate/protocol.md` with frame schemas, limits, ownership, errors, and the lifecycle table.
|
||||
2. Implement Python protocol parsing, state transitions, recorder control, PCM conversion, model abstraction, and idle shutdown behind injected fakes.
|
||||
3. Add Python tests for malformed and oversized frames, ownership, busy behavior, cancellation during model load and inference, tombstoned-result suppression, recorder spawn and exit variants, stdout drain ordering, timeout escalation, duration bounds, PCM conversion, empty result, model failure, stale generations, disconnect cleanup, worker-aware idle suppression, and 15-minute idle shutdown.
|
||||
4. Keep NeMo imports lazy so unit tests run without loading CUDA or the model.
|
||||
5. Implement the real NeMo adapter only after service-core tests pass.
|
||||
|
||||
**Completion criterion:** service unit tests run without GPU access, every event/state pair has one explicit outcome, recorder stdout always drains before transcription, and no raw audio path reaches disk or logs.
|
||||
|
||||
### Phase 2: Add socket activation and host installation
|
||||
|
||||
1. Add `pi-dictate.socket` and `pi-dictate.service` user units.
|
||||
2. Add a repeatable installer that creates the dedicated uv environment, installs the pinned package, provisions the pinned model, links the user units, reloads systemd, and enables the socket.
|
||||
3. Keep model download and dependency installation as explicit setup actions with network access; runtime remains offline.
|
||||
4. Add integration tests for socket activation, protocol negotiation, cold load, warm reuse, busy clients, client disconnect, cancellation during inference, service crash, idle shutdown, socket-path ownership, fresh-connection reactivation, and GPU release.
|
||||
5. Validate unit hardening without blocking PipeWire, CUDA, or the runtime socket.
|
||||
|
||||
**Completion criterion:** connecting to `/run/user/1000/pi-dictate.sock` activates the service; the first request loads the model; subsequent requests stay warm for 15 minutes; idle exit releases GPU memory; and a later connection reactivates cleanly.
|
||||
|
||||
### Phase 3: Replace the Pi extension internals
|
||||
|
||||
1. Remove SoX, Deepgram, WebSocket, clipboard, focus inspection, hidden widget, meter, spinner, and global input-listener code.
|
||||
2. Add a small versioned Unix-socket client with strict frame limits and generation checks. After service idle exit or disconnect, discard the accepted connection; the next start must create a fresh connection and repeat `hello` negotiation so the socket unit can reactivate the service.
|
||||
3. Register `Alt+M`, `Alt+N`, and `/dictate-cancel` through public Pi APIs.
|
||||
4. Render only service-reported state changes.
|
||||
5. Insert only a current-generation `result` through `ctx.ui.pasteToEditor()`.
|
||||
6. On session shutdown, send a best-effort cancel, close the socket, clear status, and invalidate the generation.
|
||||
7. Add TypeScript tests for registration, connection failure, version mismatch, every state event, stale events, malformed and oversized frames, empty result, `pasteToEditor()` delivery, modal-safe main-editor delivery, service disconnect, and shutdown cleanup.
|
||||
8. Run `npm run typecheck` and the focused tests after each slice.
|
||||
|
||||
**Completion criterion:** the extension has no microphone, GPU, model, cloud, clipboard, WebSocket, private TUI, or synthetic-input capability; it can only exchange bounded protocol messages and paste a successful current result into Pi's main editor.
|
||||
|
||||
### Phase 4: Add the minimum nono capability
|
||||
|
||||
1. Add `/run/user/1000/pi-dictate.sock` to `filesystem.unix_socket` in `profiles/nono/pi-local.json`.
|
||||
2. Do not add PipeWire, GPU, model, cache, Deepgram, or extra network permissions.
|
||||
3. Validate the draft and resolved policy according to the nono-profile workflow.
|
||||
4. Compare the resolved policy before and after and confirm this feature adds only the exact service socket.
|
||||
5. Restart Pi because an active nono sandbox cannot gain socket access in place.
|
||||
6. Verify the extension can activate and use the service but still cannot connect directly to PipeWire or initialize NVML.
|
||||
|
||||
**Completion criterion:** `pi-local` validates, the service socket works after restart, and Pi still lacks direct microphone and GPU access.
|
||||
|
||||
### Phase 5: Document and manually validate
|
||||
|
||||
Create `extensions/dictate/README.md` covering:
|
||||
|
||||
- Linux-only support and the tested CachyOS, niri, Ghostty, PipeWire, CUDA, and RTX 4090 environment;
|
||||
- service installation, model provisioning, systemd socket activation, status, restart, and removal;
|
||||
- model size, cold-start expectations, 15-minute warm period, and GPU-memory behavior measured in Phase 0;
|
||||
- Parakeet and NeMo licenses and attribution;
|
||||
- offline runtime and the one-time networked installation step;
|
||||
- best-effort raw-audio non-persistence, core/swap hardening, and the limits imposed by Python, NumPy, PyTorch, CUDA, and host memory management;
|
||||
- the accepted security consequence that any code inside `pi-local` can use the socket as a mediated microphone and transcription capability;
|
||||
- `Alt+M`, `Alt+N`, and `/dictate-cancel`, including main-editor focus requirements;
|
||||
- loading, listening, stopping, transcribing, failure, and cancellation states;
|
||||
- main-editor `pasteToEditor()` delivery;
|
||||
- niri shortcut conflict checks;
|
||||
- PipeWire default-source, CUDA, model, socket, systemd, and bounded logging troubleshooting.
|
||||
|
||||
Add dictate to the root README extension inventory. Resolve the pre-existing `extensions/dictate/LICENSE` deletion with the user and add model attribution without overwriting unrelated license work.
|
||||
|
||||
Manually validate from a restarted, profile-confined Pi session:
|
||||
|
||||
1. Cold start, observing socket activation, model loading, then recorder start.
|
||||
2. Warm start within 15 minutes.
|
||||
3. Start and stop from the normal main editor with existing text, a moved cursor, and collapsed pasted content.
|
||||
4. Confirm shortcuts do not claim to work while a Pi modal owns input.
|
||||
5. Open a modal after recording starts, then confirm the result still reaches only the main editor.
|
||||
6. Cancel during model loading, recording, stopping, and transcription.
|
||||
7. Empty speech, muted source, missing `pw-record`, PipeWire failure, CUDA failure, model failure, service crash, malformed protocol, and duration limit.
|
||||
8. Two Pi sessions competing for the one recorder.
|
||||
9. Pi reload and shutdown during every active state.
|
||||
10. Idle for 15 minutes, verify service exit and GPU-memory release, then confirm a fresh client connection reactivates the still-owned socket path.
|
||||
11. Narrow terminal and tmux behavior.
|
||||
12. Ghostty under niri, confirming no compositor-global shortcut or synthetic typing was introduced.
|
||||
13. Run the full repository typecheck and test suite.
|
||||
14. Have an agent that did not implement the feature perform the final code and security review.
|
||||
|
||||
**Completion criterion:** documentation matches measured behavior; every manual path returns both processes to a known state; idle shutdown releases the model; no recorder, audio buffer, socket client, timer, status, or stale generation remains; and the independent review has no unresolved blocker.
|
||||
|
||||
## Acceptance criteria
|
||||
|
||||
- Phase 0 proves the pinned Parakeet v3 and NeMo stack meets measured accuracy, latency, memory, offline, and packaging requirements before production changes begin.
|
||||
- Runtime transcription is local and makes no network request.
|
||||
- The service intentionally creates no audio files, bounds application-owned audio memory, applies core/swap hardening where supported, and releases its buffers after every terminal outcome; documentation does not promise forensic zeroization.
|
||||
- The host service, not Pi, exclusively owns PipeWire, CUDA, model files, and audio buffers.
|
||||
- The service is socket-activated, handles one recording or inference worker at a time, remains warm for 15 minutes, exits only when truly idle, leaves socket-path ownership to the socket unit, and releases GPU memory.
|
||||
- Recorder stdout drains and the child closes before transcription starts.
|
||||
- Five-minute audio, protocol frames, stderr, logs, and transcript output are bounded.
|
||||
- Only a non-empty `result` for the active generation changes the editor.
|
||||
- Text insertion uses `ctx.ui.pasteToEditor()` and preserves cursor and collapsed-paste behavior.
|
||||
- `Alt+M`, `Alt+N`, and `/dictate-cancel` are documented as main-editor-only controls.
|
||||
- Stop, cancellation, disconnect, failure, timeout, stale result, and session shutdown are deterministic in every state.
|
||||
- The extension contains no SoX, Deepgram, WebSocket, clipboard, X11 injection, Wayland synthetic typing, private TUI access, hidden TUI widget, global input listener, microphone capture, or GPU code.
|
||||
- `npm run typecheck` passes.
|
||||
- `npm test` includes and passes dictate's TypeScript tests, and the documented service test command passes its Python tests.
|
||||
- The systemd user units and installer work from a clean setup and can be removed cleanly.
|
||||
- The pinned model revision, dependency lock, source references, and license attribution are documented.
|
||||
- The `pi-local` resolved-policy diff adds only `/run/user/1000/pi-dictate.sock` for this feature, and documentation identifies that socket as a user-accepted mediated microphone and transcription capability available to every descendant profile.
|
||||
- A restarted Pi can use the service socket but still cannot connect directly to PipeWire or initialize NVML.
|
||||
- The user's niri configuration has no conflicting `Alt+M` or `Alt+N` binding, and no global dictate binding is added.
|
||||
- User documentation covers setup, privacy, model behavior, controls, lifecycle, failure recovery, troubleshooting, and uninstall.
|
||||
+4
-1
@@ -92,7 +92,8 @@ The skill stages candidates under `~/.config/nono/profile-drafts/` and validates
|
||||
|
||||
- read/write access to the launch working directory and temporary directory
|
||||
- read-only access to local Pi extension source and `$HOME/.bun`; the Orca runner invokes the selected Node executable explicitly rather than relying on `PATH`
|
||||
- pathname-mediated access to the system resolver Unix socket on Linux
|
||||
- pathname-mediated access to the system resolver Unix socket and the exact `/run/user/1000/pi-dictate.sock` mediated microphone socket on Linux
|
||||
- an explicit deny for `~/.local/share/pi-dictate`, which overrides `$WORKDIR` if Pi is launched from the service data directory and protects the model, dictionary, caches, and runtime
|
||||
- an Internet hostname allowlist for model providers, source hosts, package registries, and toolchain downloads
|
||||
|
||||
The parent profile contributes more policy, including Pi state access, security groups, credential route definitions, URL-opening origins, localhost URL opening, and service launching. Inspect the resolved profile before treating a child as narrower than its parent:
|
||||
@@ -212,6 +213,8 @@ Add exact sockets for required local services. Avoid granting a socket directory
|
||||
|
||||
`linux.af_unix_mediation: "pathname"` enables pathname-based Unix-socket mediation on Linux.
|
||||
|
||||
`pi-local` also grants `/run/user/1000/pi-dictate.sock`. This socket is a mediated microphone and transcription capability, not a health-only endpoint. Any code running under `pi-local` or an inheriting profile can request recording and receive transcripts, while direct PipeWire, audio-device, GPU, model, and dictionary access remains outside Pi.
|
||||
|
||||
### Network access
|
||||
|
||||
`network.block` and `network.allow_domain` define three useful tiers.
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
# Todos
|
||||
|
||||
## Investigate transient concurrent NeMo benchmark failures
|
||||
|
||||
During Dictate Phase 0 on 26AUG2026, repeated overlapping throwaway benchmark processes produced two one-off failures in the external `/srv/data/tools/parakeet-v3/.venv` environment:
|
||||
|
||||
1. `ASRModel.restore_from()` raised `MemoryError` while OmegaConf rebuilt `joint.vocabulary`. The exact clean/noisy/silence harness passed on its immediate rerun.
|
||||
2. A later production-adapter check raised `ValueError: bad marshal data (invalid reference)` while importing `torch._tensor`. An immediate standalone `import torch` succeeded, and the cached `_tensor.cpython-312.pyc` passed `marshal.loads()`.
|
||||
|
||||
The production service serializes model loading and inference and therefore does not intentionally launch concurrent NeMo processes. Before treating concurrent workers as supported, build a repeatable cold-load stress loop, monitor process and cgroup memory, disable bytecode writes, and determine whether the failures come from NeMo/OmegaConf resource pressure or concurrent access to the external environment's bytecode cache.
|
||||
@@ -1,21 +0,0 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2025 Amos Blomqvist
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,181 @@
|
||||
# Pi Dictate
|
||||
|
||||
Linux-only push-to-dictate for Pi. Audio capture and NVIDIA Parakeet run in a socket-activated host service. The Pi extension receives state and finalized text through a Unix socket, then inserts successful text with `pasteToEditor()`.
|
||||
|
||||
Tested on CachyOS with niri, Ghostty, PipeWire, an RTX 4090, CUDA 13.0, NeMo 3.0.0, and Parakeet-TDT 0.6B v3.
|
||||
|
||||
## Requirements
|
||||
|
||||
- Linux with PipeWire and `pw-record`
|
||||
- systemd user services
|
||||
- NVIDIA GPU with a working CUDA runtime
|
||||
- `uv`
|
||||
- Network access during initial dependency and model provisioning
|
||||
|
||||
Runtime is offline. The systemd service restricts networking and loads an explicit local model archive.
|
||||
|
||||
## Install
|
||||
|
||||
From the repository root:
|
||||
|
||||
```bash
|
||||
./services/dictate/install.sh
|
||||
```
|
||||
|
||||
The installer creates `~/.local/share/pi-dictate/`, installs the locked service as a non-editable wheel, downloads and verifies the pinned model snapshot, seeds the coding dictionary, installs host-owned copies of the systemd user units, and enables `pi-dictate.socket`. The installed runtime does not depend on the repository remaining at its installation path.
|
||||
|
||||
To reuse an already downloaded model archive:
|
||||
|
||||
```bash
|
||||
./services/dictate/install.sh \
|
||||
--model-source /path/to/parakeet-tdt-0.6b-v3.nemo
|
||||
```
|
||||
|
||||
Preview without changing the host:
|
||||
|
||||
```bash
|
||||
./services/dictate/install.sh --dry-run --model-source /path/to/model.nemo
|
||||
```
|
||||
|
||||
The pinned model revision and SHA-256 are recorded in [`../../docs/dictate/parakeet-attribution.md`](../../docs/dictate/parakeet-attribution.md).
|
||||
|
||||
## Sandbox access
|
||||
|
||||
`profiles/nono/pi-local.json` grants exactly:
|
||||
|
||||
```text
|
||||
/run/user/1000/pi-dictate.sock
|
||||
```
|
||||
|
||||
This is a mediated microphone and transcription capability. Any code running under `pi-local` or a descendant profile can ask the host service to record and can receive the transcript. Pi still receives no PipeWire, `/dev/snd`, CUDA device, model, cache, or dictionary access. An explicit filesystem deny protects `~/.local/share/pi-dictate` even if that directory is selected as `$WORKDIR`.
|
||||
|
||||
The active profile is repository-linked on the tested host. Start a new Pi sandbox after updating the profile because an existing nono process cannot gain the socket capability in place.
|
||||
|
||||
## Controls
|
||||
|
||||
| Control | Action |
|
||||
| --- | --- |
|
||||
| `Alt+M` | Start recording when idle; stop and transcribe while listening |
|
||||
| `Alt+N` | Cancel the active generation |
|
||||
| `/dictate-cancel` | Cancel the active generation |
|
||||
|
||||
Shortcuts are handled by Pi's main editor. They do not claim to work while a selector, confirmation, or custom modal owns input. If a modal opens after recording starts, a successful result still goes only to the main editor.
|
||||
|
||||
The footer shows static service states:
|
||||
|
||||
- `loading model...`
|
||||
- `listening...`
|
||||
- `stopping...`
|
||||
- `transcribing...`
|
||||
|
||||
Only a non-empty finalized result for the current generation changes the editor. Empty speech, cancellation, stale results, disconnects, and failures do not insert text.
|
||||
|
||||
## Coding dictionary
|
||||
|
||||
The installer seeds:
|
||||
|
||||
```text
|
||||
~/.local/share/pi-dictate/config/coding-terms.txt
|
||||
```
|
||||
|
||||
The host service reads this file at model load and applies NeMo's native phrase boosting. Pi cannot read or modify it. Add one term or phrase per line, use `#` for comments, and restart the service after editing:
|
||||
|
||||
```bash
|
||||
systemctl --user restart pi-dictate.service
|
||||
```
|
||||
|
||||
Reinstalling preserves an existing dictionary. The file is bounded to 64 KiB, 512 distinct terms, and 128 UTF-8 bytes per term. A conservative boost of `1.0` is configured in `pi-dictate.service`; aggressive values can create false positives.
|
||||
|
||||
## Model lifecycle and measured behavior
|
||||
|
||||
The socket remains active while the Python service exits after 15 minutes of true inactivity. A later connection activates a fresh service. Loading, recording, recorder shutdown, and inference suspend idle exit.
|
||||
|
||||
Measured on the tested RTX 4090 using the pinned offline model:
|
||||
|
||||
- Model archive: about 2.4 GB
|
||||
- Locked Python environment: about 5.6 GB in the Phase 0 prototype
|
||||
- Cold base-model load: about 6.4 seconds
|
||||
- Production-adapter load with phrase-boost setup: about 10 seconds, including the private Triton cache build
|
||||
- Ten-second microphone inference: about 0.29 seconds, 0.029 real-time factor
|
||||
- Sixty-second repeated-speech inference: about 0.43 seconds, 0.0072 real-time factor
|
||||
- Five-minute repeated-speech inference: about 0.82 seconds, 0.0027 real-time factor
|
||||
- Peak CUDA allocation: about 5.1 GB for short clips and 6.15 GB for five minutes
|
||||
- Five-minute raw PCM cap: 9.6 MB
|
||||
|
||||
The 60-second and five-minute stress inputs repeat a real speech clip in memory. They validate bounded duration, latency, and GPU allocation, not long-form accuracy.
|
||||
|
||||
## Privacy
|
||||
|
||||
The service passes this exact recorder command through pipes:
|
||||
|
||||
```text
|
||||
pw-record --raw --format s16 --rate 16000 --channels 1 -
|
||||
```
|
||||
|
||||
It intentionally creates no audio files and releases application-owned buffers after transcription, cancellation, or failure. `LimitCORE=0` and `MemorySwapMax=0` reduce additional persistence. This is best-effort non-persistence, not forensic zeroization. Python, NumPy, PyTorch, CUDA, drivers, host memory management, and crash handling can retain copies outside the service's direct control.
|
||||
|
||||
The service does not log raw audio or transcripts by default. Recorder stderr and protocol output are bounded. The server accepts at most 16 clients and bounds each client's pending output to 64 KiB.
|
||||
|
||||
## Service operations
|
||||
|
||||
```bash
|
||||
systemctl --user status pi-dictate.socket pi-dictate.service
|
||||
journalctl --user -u pi-dictate.service
|
||||
systemctl --user restart pi-dictate.service
|
||||
systemctl --user stop pi-dictate.service
|
||||
```
|
||||
|
||||
Stopping only the service leaves the socket active for later reactivation. Stopping `pi-dictate.socket` removes the socket path.
|
||||
|
||||
## niri and Ghostty
|
||||
|
||||
No compositor-global binding is required. Confirm the live niri config does not claim `Alt+M` or `Alt+N`, then validate it outside the Pi sandbox:
|
||||
|
||||
```bash
|
||||
niri validate
|
||||
```
|
||||
|
||||
The shortcuts should reach Pi only while its main editor has focus. No `wl-copy`, `wtype`, `ydotool`, X11 injection, or compositor spawn binding is used.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Socket or activation failure
|
||||
|
||||
```bash
|
||||
systemctl --user status pi-dictate.socket pi-dictate.service
|
||||
ls -l "$XDG_RUNTIME_DIR/pi-dictate.sock"
|
||||
journalctl --user -u pi-dictate.service -n 100
|
||||
```
|
||||
|
||||
The socket should be mode `0600` and owned by the current user. Restart Pi after changing nono policy.
|
||||
|
||||
### PipeWire or microphone failure
|
||||
|
||||
```bash
|
||||
wpctl status
|
||||
pw-record --raw --format s16 --rate 16000 --channels 1 - >/dev/null
|
||||
```
|
||||
|
||||
Press `Ctrl+C` after confirming capture. The service uses the WirePlumber default input. Selectable devices are not implemented.
|
||||
|
||||
### CUDA or model failure
|
||||
|
||||
```bash
|
||||
nvidia-smi
|
||||
systemctl --user restart pi-dictate.service
|
||||
journalctl --user -u pi-dictate.service -n 100
|
||||
```
|
||||
|
||||
Confirm the installed model checksum matches the attribution document and that runtime network access is not needed.
|
||||
|
||||
### Service reports busy
|
||||
|
||||
Only one recording, load worker, or inference worker is allowed globally. A cancelled load or inference remains busy until its non-interruptible worker finishes. Wait for it to settle or restart the service if the owning client crashed and the worker is stuck.
|
||||
|
||||
## Uninstall
|
||||
|
||||
```bash
|
||||
./services/dictate/install.sh uninstall
|
||||
```
|
||||
|
||||
The uninstaller disables the socket, removes only installer-managed unit copies, reloads systemd, and removes the fixed `~/.local/share/pi-dictate/` directory. It ignores install-root environment overrides and refuses to replace or remove unmanaged unit files.
|
||||
@@ -0,0 +1,432 @@
|
||||
import assert from "node:assert/strict";
|
||||
import test from "node:test";
|
||||
import type {
|
||||
ExtensionAPI,
|
||||
ExtensionCommandContext,
|
||||
ExtensionContext,
|
||||
} from "@earendil-works/pi-coding-agent";
|
||||
import {
|
||||
createDictateExtension,
|
||||
MAX_FRAME_BYTES,
|
||||
MAX_TRANSCRIPT_BYTES,
|
||||
PROTOCOL_VERSION,
|
||||
type DictateSocket,
|
||||
} from "./index.ts";
|
||||
|
||||
type ShortcutHandler = (ctx: ExtensionContext) => Promise<void> | void;
|
||||
type CommandHandler = (
|
||||
args: string,
|
||||
ctx: ExtensionCommandContext,
|
||||
) => Promise<void> | void;
|
||||
type EventHandler = (event: unknown, ctx: ExtensionContext) => Promise<void> | void;
|
||||
type Frame = Record<string, unknown>;
|
||||
|
||||
class FakeSocket implements DictateSocket {
|
||||
readonly writes: string[] = [];
|
||||
endCount = 0;
|
||||
destroyCount = 0;
|
||||
private connectListener: (() => void) | undefined;
|
||||
private dataListener: ((chunk: Uint8Array) => void) | undefined;
|
||||
private errorListener: ((error: Error) => void) | undefined;
|
||||
private closeListener: (() => void) | undefined;
|
||||
|
||||
onConnect(listener: () => void) {
|
||||
this.connectListener = listener;
|
||||
}
|
||||
onData(listener: (chunk: Uint8Array) => void) {
|
||||
this.dataListener = listener;
|
||||
}
|
||||
onError(listener: (error: Error) => void) {
|
||||
this.errorListener = listener;
|
||||
}
|
||||
onClose(listener: () => void) {
|
||||
this.closeListener = listener;
|
||||
}
|
||||
write(data: string) {
|
||||
this.writes.push(data);
|
||||
}
|
||||
end() {
|
||||
this.endCount++;
|
||||
}
|
||||
destroy() {
|
||||
this.destroyCount++;
|
||||
}
|
||||
connect() {
|
||||
this.connectListener?.();
|
||||
}
|
||||
data(data: string | Uint8Array) {
|
||||
this.dataListener?.(
|
||||
typeof data === "string" ? Buffer.from(data, "utf8") : data,
|
||||
);
|
||||
}
|
||||
frame(frame: Frame) {
|
||||
this.data(`${JSON.stringify(frame)}\n`);
|
||||
}
|
||||
error(message: string) {
|
||||
this.errorListener?.(new Error(message));
|
||||
}
|
||||
close() {
|
||||
this.closeListener?.();
|
||||
}
|
||||
}
|
||||
|
||||
function createHarness() {
|
||||
const shortcuts = new Map<string, ShortcutHandler>();
|
||||
const commands = new Map<string, CommandHandler>();
|
||||
const events = new Map<string, EventHandler>();
|
||||
const sockets: FakeSocket[] = [];
|
||||
const socketPaths: string[] = [];
|
||||
const statuses: Array<string | undefined> = [];
|
||||
const notifications: Array<{ message: string; type?: string }> = [];
|
||||
const pasted: string[] = [];
|
||||
const context = {
|
||||
ui: {
|
||||
setStatus: (_key: string, text: string | undefined) => statuses.push(text),
|
||||
notify: (message: string, type?: string) =>
|
||||
notifications.push({ message, type }),
|
||||
pasteToEditor: (text: string) => pasted.push(text),
|
||||
},
|
||||
} as unknown as ExtensionContext;
|
||||
const api = {
|
||||
registerShortcut(key: string, options: { handler: ShortcutHandler }) {
|
||||
shortcuts.set(key, options.handler);
|
||||
},
|
||||
registerCommand(name: string, options: { handler: CommandHandler }) {
|
||||
commands.set(name, options.handler);
|
||||
},
|
||||
on(name: string, handler: EventHandler) {
|
||||
events.set(name, handler);
|
||||
},
|
||||
} as unknown as ExtensionAPI;
|
||||
|
||||
createDictateExtension({
|
||||
socketPath: "/test/pi-dictate.sock",
|
||||
socketFactory: (path) => {
|
||||
socketPaths.push(path);
|
||||
const socket = new FakeSocket();
|
||||
sockets.push(socket);
|
||||
return socket;
|
||||
},
|
||||
})(api);
|
||||
|
||||
const shortcut = (key: string) => {
|
||||
const handler = shortcuts.get(key);
|
||||
assert.ok(handler, `shortcut ${key} was registered`);
|
||||
return handler;
|
||||
};
|
||||
const command = (name: string) => {
|
||||
const handler = commands.get(name);
|
||||
assert.ok(handler, `command ${name} was registered`);
|
||||
return handler;
|
||||
};
|
||||
const event = (name: string) => {
|
||||
const handler = events.get(name);
|
||||
assert.ok(handler, `event ${name} was registered`);
|
||||
return handler;
|
||||
};
|
||||
const sentFrames = (socket: FakeSocket) =>
|
||||
socket.writes.map((line) => JSON.parse(line) as Frame);
|
||||
|
||||
return {
|
||||
shortcuts,
|
||||
commands,
|
||||
events,
|
||||
sockets,
|
||||
socketPaths,
|
||||
statuses,
|
||||
notifications,
|
||||
pasted,
|
||||
context,
|
||||
shortcut,
|
||||
command,
|
||||
event,
|
||||
sentFrames,
|
||||
};
|
||||
}
|
||||
|
||||
type Harness = ReturnType<typeof createHarness>;
|
||||
|
||||
async function connectAndNegotiate(harness: Harness) {
|
||||
await harness.shortcut("alt+m")(harness.context);
|
||||
const socket = harness.sockets.at(-1);
|
||||
assert.ok(socket);
|
||||
socket.connect();
|
||||
const hello = harness.sentFrames(socket)[0];
|
||||
assert.equal(hello?.type, "hello");
|
||||
assert.equal(hello?.version, PROTOCOL_VERSION);
|
||||
assert.equal(typeof hello?.requestId, "string");
|
||||
assert.equal(harness.sentFrames(socket).length, 1, "start waits for hello");
|
||||
socket.frame({
|
||||
version: PROTOCOL_VERSION,
|
||||
type: "pong",
|
||||
requestId: hello?.requestId,
|
||||
});
|
||||
const start = harness.sentFrames(socket)[1];
|
||||
assert.equal(start?.type, "start");
|
||||
assert.equal(start?.version, PROTOCOL_VERSION);
|
||||
assert.equal(typeof start?.requestId, "string");
|
||||
assert.equal(typeof start?.generationId, "string");
|
||||
return { socket, generationId: start?.generationId as string };
|
||||
}
|
||||
|
||||
function serviceFrame(
|
||||
socket: FakeSocket,
|
||||
generationId: string,
|
||||
type: string,
|
||||
fields: Frame = {},
|
||||
) {
|
||||
socket.frame({
|
||||
version: PROTOCOL_VERSION,
|
||||
type,
|
||||
requestId: `event-${type}`,
|
||||
generationId,
|
||||
...fields,
|
||||
});
|
||||
}
|
||||
|
||||
test("registers main-editor shortcuts, cancel command, and shutdown cleanup", () => {
|
||||
const harness = createHarness();
|
||||
assert.deepEqual([...harness.shortcuts.keys()], ["alt+m", "alt+n"]);
|
||||
assert.deepEqual([...harness.commands.keys()], ["dictate-cancel"]);
|
||||
assert.ok(harness.events.has("session_shutdown"));
|
||||
});
|
||||
|
||||
test("connection errors return to idle and the next start uses a fresh socket", async () => {
|
||||
const harness = createHarness();
|
||||
await harness.shortcut("alt+m")(harness.context);
|
||||
harness.sockets[0]?.error("socket unavailable");
|
||||
|
||||
assert.match(harness.notifications[0]?.message ?? "", /connection failed.*socket unavailable/i);
|
||||
assert.equal(harness.notifications[0]?.type, "error");
|
||||
assert.equal(harness.statuses.at(-1), undefined);
|
||||
assert.equal(harness.sockets[0]?.destroyCount, 1);
|
||||
|
||||
await harness.shortcut("alt+m")(harness.context);
|
||||
assert.equal(harness.sockets.length, 2);
|
||||
assert.deepEqual(harness.socketPaths, ["/test/pi-dictate.sock", "/test/pi-dictate.sock"]);
|
||||
});
|
||||
|
||||
test("rejects a mismatched protocol version before sending start", async () => {
|
||||
const harness = createHarness();
|
||||
await harness.shortcut("alt+m")(harness.context);
|
||||
const socket = harness.sockets[0];
|
||||
assert.ok(socket);
|
||||
socket.connect();
|
||||
const hello = harness.sentFrames(socket)[0];
|
||||
socket.frame({ version: 2, type: "pong", requestId: hello?.requestId });
|
||||
|
||||
assert.equal(harness.sentFrames(socket).length, 1);
|
||||
assert.match(harness.notifications[0]?.message ?? "", /version mismatch/i);
|
||||
assert.equal(socket.destroyCount, 1);
|
||||
});
|
||||
|
||||
test("rejects a hello response with the wrong request ID", async () => {
|
||||
const harness = createHarness();
|
||||
await harness.shortcut("alt+m")(harness.context);
|
||||
const socket = harness.sockets[0];
|
||||
assert.ok(socket);
|
||||
socket.connect();
|
||||
socket.frame({
|
||||
version: PROTOCOL_VERSION,
|
||||
type: "pong",
|
||||
requestId: "not-the-hello-request",
|
||||
});
|
||||
|
||||
assert.equal(harness.sentFrames(socket).length, 1);
|
||||
assert.match(harness.notifications[0]?.message ?? "", /unexpected hello response/i);
|
||||
assert.equal(socket.destroyCount, 1);
|
||||
});
|
||||
|
||||
test("rejects generation events before hello negotiation", async () => {
|
||||
const harness = createHarness();
|
||||
await harness.shortcut("alt+m")(harness.context);
|
||||
const socket = harness.sockets[0];
|
||||
assert.ok(socket);
|
||||
socket.connect();
|
||||
serviceFrame(socket, "generation-1", "state", { state: "listening" });
|
||||
|
||||
assert.match(harness.notifications[0]?.message ?? "", /before hello negotiation/i);
|
||||
assert.equal(socket.destroyCount, 1);
|
||||
assert.deepEqual(harness.pasted, []);
|
||||
});
|
||||
|
||||
test("renders authoritative service states and sends a bounded stop command", async () => {
|
||||
const harness = createHarness();
|
||||
const { socket, generationId } = await connectAndNegotiate(harness);
|
||||
|
||||
serviceFrame(socket, generationId, "state", { state: "loading" });
|
||||
serviceFrame(socket, generationId, "state", { state: "listening" });
|
||||
await harness.shortcut("alt+m")(harness.context);
|
||||
const stop = harness.sentFrames(socket).at(-1);
|
||||
assert.equal(stop?.type, "stop");
|
||||
assert.equal(stop?.generationId, generationId);
|
||||
assert.ok(Buffer.byteLength(String(stop?.requestId)) <= 128);
|
||||
serviceFrame(socket, generationId, "state", { state: "stopping" });
|
||||
serviceFrame(socket, generationId, "state", { state: "transcribing" });
|
||||
|
||||
assert.deepEqual(harness.statuses, [
|
||||
"loading model...",
|
||||
"listening...",
|
||||
"stopping...",
|
||||
"transcribing...",
|
||||
]);
|
||||
});
|
||||
|
||||
test("ignores stale generations but rejects malformed and oversized frames", async () => {
|
||||
const staleHarness = createHarness();
|
||||
const { socket: staleSocket, generationId } = await connectAndNegotiate(staleHarness);
|
||||
serviceFrame(staleSocket, "old-generation", "state", { state: "listening" });
|
||||
assert.deepEqual(staleHarness.statuses, []);
|
||||
serviceFrame(staleSocket, generationId, "state", { state: "listening" });
|
||||
assert.deepEqual(staleHarness.statuses, ["listening..."]);
|
||||
|
||||
const malformedHarness = createHarness();
|
||||
const malformed = await connectAndNegotiate(malformedHarness);
|
||||
malformed.socket.data("{not json}\n");
|
||||
assert.match(malformedHarness.notifications[0]?.message ?? "", /malformed JSON/i);
|
||||
assert.equal(malformed.socket.destroyCount, 1);
|
||||
|
||||
const nonObjectHarness = createHarness();
|
||||
const nonObject = await connectAndNegotiate(nonObjectHarness);
|
||||
nonObject.socket.data("null\n");
|
||||
assert.match(nonObjectHarness.notifications[0]?.message ?? "", /frame must be an object/i);
|
||||
assert.equal(nonObject.socket.destroyCount, 1);
|
||||
|
||||
const oversizedHarness = createHarness();
|
||||
const oversized = await connectAndNegotiate(oversizedHarness);
|
||||
oversized.socket.data("x".repeat(MAX_FRAME_BYTES + 1));
|
||||
assert.match(oversizedHarness.notifications[0]?.message ?? "", /frame is too large/i);
|
||||
assert.equal(oversized.socket.destroyCount, 1);
|
||||
});
|
||||
|
||||
test("rejects invalid UTF-8 instead of altering transcript text", async () => {
|
||||
const harness = createHarness();
|
||||
const { socket, generationId } = await connectAndNegotiate(harness);
|
||||
serviceFrame(socket, generationId, "state", { state: "listening" });
|
||||
await harness.shortcut("alt+m")(harness.context);
|
||||
serviceFrame(socket, generationId, "state", { state: "transcribing" });
|
||||
const prefix = Buffer.from(
|
||||
JSON.stringify({
|
||||
version: PROTOCOL_VERSION,
|
||||
type: "result",
|
||||
requestId: "event-result",
|
||||
generationId,
|
||||
transcript: "",
|
||||
}).replace('"transcript":""}', '"transcript":"'),
|
||||
);
|
||||
const suffix = Buffer.from('"}\n');
|
||||
socket.data(Buffer.concat([prefix, Buffer.from([0xc3, 0x28]), suffix]));
|
||||
|
||||
assert.match(harness.notifications[0]?.message ?? "", /invalid UTF-8/i);
|
||||
assert.equal(socket.destroyCount, 1);
|
||||
assert.deepEqual(harness.pasted, []);
|
||||
});
|
||||
|
||||
test("service idle closes the accepted connection and forces a fresh handshake", async () => {
|
||||
const harness = createHarness();
|
||||
const { socket, generationId } = await connectAndNegotiate(harness);
|
||||
serviceFrame(socket, generationId, "state", { state: "loading" });
|
||||
serviceFrame(socket, generationId, "state", { state: "idle" });
|
||||
assert.equal(socket.endCount, 1);
|
||||
assert.equal(harness.statuses.at(-1), undefined);
|
||||
|
||||
await harness.shortcut("alt+m")(harness.context);
|
||||
assert.equal(harness.sockets.length, 2);
|
||||
harness.sockets[1]?.connect();
|
||||
assert.equal(harness.sentFrames(harness.sockets[1] as FakeSocket)[0]?.type, "hello");
|
||||
});
|
||||
|
||||
test("an empty result leaves the editor unchanged and reports no speech", async () => {
|
||||
const harness = createHarness();
|
||||
const { socket, generationId } = await connectAndNegotiate(harness);
|
||||
serviceFrame(socket, generationId, "state", { state: "listening" });
|
||||
await harness.shortcut("alt+m")(harness.context);
|
||||
serviceFrame(socket, generationId, "state", { state: "transcribing" });
|
||||
serviceFrame(socket, generationId, "result", { transcript: " " });
|
||||
|
||||
assert.deepEqual(harness.pasted, []);
|
||||
assert.match(harness.notifications[0]?.message ?? "", /no speech recognized/i);
|
||||
assert.equal(socket.endCount, 1);
|
||||
});
|
||||
|
||||
test("pastes exactly one non-empty current result through the public editor API", async () => {
|
||||
const harness = createHarness();
|
||||
const { socket, generationId } = await connectAndNegotiate(harness);
|
||||
serviceFrame(socket, generationId, "state", { state: "listening" });
|
||||
await harness.shortcut("alt+m")(harness.context);
|
||||
serviceFrame(socket, generationId, "state", { state: "stopping" });
|
||||
serviceFrame(socket, generationId, "state", { state: "transcribing" });
|
||||
serviceFrame(socket, generationId, "result", { transcript: "typed words" });
|
||||
serviceFrame(socket, generationId, "result", { transcript: "duplicate" });
|
||||
|
||||
assert.deepEqual(harness.pasted, ["typed words"]);
|
||||
assert.equal(harness.statuses.at(-1), undefined);
|
||||
|
||||
const oversizedTranscriptHarness = createHarness();
|
||||
const active = await connectAndNegotiate(oversizedTranscriptHarness);
|
||||
serviceFrame(active.socket, active.generationId, "state", { state: "listening" });
|
||||
await oversizedTranscriptHarness.shortcut("alt+m")(oversizedTranscriptHarness.context);
|
||||
serviceFrame(active.socket, active.generationId, "state", { state: "transcribing" });
|
||||
serviceFrame(active.socket, active.generationId, "result", {
|
||||
transcript: "t".repeat(MAX_TRANSCRIPT_BYTES + 1),
|
||||
});
|
||||
assert.deepEqual(oversizedTranscriptHarness.pasted, []);
|
||||
assert.match(oversizedTranscriptHarness.notifications[0]?.message ?? "", /transcript is too large/i);
|
||||
});
|
||||
|
||||
test("disconnects and service errors return safely to idle without delivery", async () => {
|
||||
const disconnected = createHarness();
|
||||
const first = await connectAndNegotiate(disconnected);
|
||||
first.socket.close();
|
||||
assert.match(disconnected.notifications[0]?.message ?? "", /disconnected/i);
|
||||
assert.deepEqual(disconnected.pasted, []);
|
||||
await disconnected.shortcut("alt+m")(disconnected.context);
|
||||
assert.equal(disconnected.sockets.length, 2);
|
||||
|
||||
const failed = createHarness();
|
||||
const second = await connectAndNegotiate(failed);
|
||||
serviceFrame(second.socket, second.generationId, "error", {
|
||||
code: "model_failed",
|
||||
detail: "model unavailable",
|
||||
});
|
||||
assert.match(failed.notifications[0]?.message ?? "", /model_failed.*model unavailable/i);
|
||||
assert.deepEqual(failed.pasted, []);
|
||||
});
|
||||
|
||||
test("cancellation during loading tombstones later events and the slash command also cancels", async () => {
|
||||
const harness = createHarness();
|
||||
const { socket, generationId } = await connectAndNegotiate(harness);
|
||||
serviceFrame(socket, generationId, "state", { state: "loading" });
|
||||
await harness.shortcut("alt+n")(harness.context);
|
||||
|
||||
assert.equal(harness.sentFrames(socket).at(-1)?.type, "cancel");
|
||||
assert.equal(socket.endCount, 1);
|
||||
serviceFrame(socket, generationId, "state", { state: "listening" });
|
||||
serviceFrame(socket, generationId, "result", { transcript: "must not appear" });
|
||||
assert.deepEqual(harness.pasted, []);
|
||||
assert.equal(harness.statuses.at(-1), undefined);
|
||||
|
||||
const next = await connectAndNegotiate(harness);
|
||||
serviceFrame(next.socket, next.generationId, "state", { state: "listening" });
|
||||
await harness.command("dictate-cancel")(
|
||||
"",
|
||||
harness.context as ExtensionCommandContext,
|
||||
);
|
||||
assert.equal(harness.sentFrames(next.socket).at(-1)?.type, "cancel");
|
||||
assert.equal(next.socket.endCount, 1);
|
||||
});
|
||||
|
||||
test("shutdown best-effort cancels, closes, clears status, and remains delivery-neutral", async () => {
|
||||
const harness = createHarness();
|
||||
const { socket, generationId } = await connectAndNegotiate(harness);
|
||||
serviceFrame(socket, generationId, "state", { state: "listening" });
|
||||
const shutdown = harness.event("session_shutdown");
|
||||
await shutdown({}, harness.context);
|
||||
await shutdown({}, harness.context);
|
||||
|
||||
assert.equal(harness.sentFrames(socket).at(-1)?.type, "cancel");
|
||||
assert.equal(socket.endCount, 1);
|
||||
assert.equal(harness.statuses.at(-1), undefined);
|
||||
serviceFrame(socket, generationId, "result", { transcript: "never paste cleanup" });
|
||||
assert.deepEqual(harness.pasted, []);
|
||||
});
|
||||
+374
-503
@@ -1,548 +1,419 @@
|
||||
/**
|
||||
* dictate — minimal voice dictation for pi.
|
||||
*
|
||||
* Press alt+m to start, press it again to stop.
|
||||
* Press alt+n to cancel and discard the in-flight transcript.
|
||||
*
|
||||
* Focus-aware: alt+m/alt+n are intercepted at the TUI input layer (before any
|
||||
* focused component), so dictation works inside ANY dialog — quiz popups,
|
||||
* ask_user_question, ctx.ui.editor()/input() — not just the main chat editor.
|
||||
*
|
||||
* Start rule: dictation only begins if some text-capable component is
|
||||
* focused; otherwise an ephemeral notification explains why nothing happened.
|
||||
* Opaque dialogs (quiz/ask selects) count as text-capable, but their internal
|
||||
* focus is invisible to us — Tab into the note/Other field first so the text
|
||||
* lands there.
|
||||
*
|
||||
* Stop rule: the delivery target is resolved fresh at stop time and the
|
||||
* transcript goes to whatever is focused THEN (editor-like components get a
|
||||
* direct setText append; opaque components get synthetic keystrokes). If
|
||||
* nothing text-capable is focused at stop, the transcript is copied to the
|
||||
* clipboard and a notification says so — a finished dictation is never lost.
|
||||
*
|
||||
* Requires:
|
||||
* - sox installed (`brew install sox` — provides the `rec` command)
|
||||
* - DEEPGRAM_API_KEY environment variable set
|
||||
*
|
||||
* Streaming model: audio is sent to Deepgram while you talk; the server
|
||||
* transcribes in real time and emits per-utterance "final" results. We
|
||||
* collect those finals and inject the concatenated text on stop. No
|
||||
* partials are shown in the editor (cosmetic-only), so quality is good
|
||||
* and the editor never shows revisable text.
|
||||
*/
|
||||
|
||||
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
|
||||
import { Key, matchesKey, isKeyRelease, isKeyRepeat } from "@earendil-works/pi-tui";
|
||||
import { spawn, type ChildProcessByStdio } from "node:child_process";
|
||||
import type { Readable } from "node:stream";
|
||||
import { appendFileSync } from "node:fs";
|
||||
import { createConnection } from "node:net";
|
||||
|
||||
// Optional forensic logging: run pi with DICTATE_DEBUG=1 to append timestamped
|
||||
// lifecycle events (listener hits, toggles, ws open/error/close with their
|
||||
// generation) to /tmp/dictate-debug.log.
|
||||
const DEBUG = !!process.env.DICTATE_DEBUG;
|
||||
const dbg = (msg: string) => {
|
||||
if (!DEBUG) return;
|
||||
try {
|
||||
appendFileSync("/tmp/dictate-debug.log", `${new Date().toISOString()} ${msg}\n`);
|
||||
} catch {}
|
||||
};
|
||||
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;
|
||||
|
||||
// Deepgram streaming endpoint. Tuning notes:
|
||||
// model=nova-3 — flagship, sub-300ms latency, best accuracy
|
||||
// encoding=linear16 — raw 16-bit PCM (what sox/rec gives us with -e signed-integer -b 16)
|
||||
// sample_rate=16000 — 16kHz mono is the standard low-bandwidth STT format
|
||||
// interim_results=false — we only want finals, never partials
|
||||
// smart_format=true — formats numbers, dates, currencies nicely
|
||||
// punctuate=true — adds commas/periods/question marks
|
||||
// endpointing=300 — 300ms of silence ends an utterance (faster finals)
|
||||
const DG_URL =
|
||||
"wss://api.deepgram.com/v1/listen" +
|
||||
"?model=nova-3" +
|
||||
"&encoding=linear16" +
|
||||
"&sample_rate=16000" +
|
||||
"&channels=1" +
|
||||
"&interim_results=false" +
|
||||
"&smart_format=true" +
|
||||
"&punctuate=true" +
|
||||
"&endpointing=300";
|
||||
const STATUS_KEY = "dictate";
|
||||
const MAX_ERROR_DETAIL_BYTES = 2 * 1024;
|
||||
const UTF8_DECODER = new TextDecoder("utf-8", { fatal: true });
|
||||
|
||||
type State = "idle" | "recording" | "stopping";
|
||||
type DictateState =
|
||||
| "idle"
|
||||
| "connecting"
|
||||
| "recording"
|
||||
| "stopping"
|
||||
| "transcribing";
|
||||
type ServiceState =
|
||||
| "loading"
|
||||
| "listening"
|
||||
| "stopping"
|
||||
| "transcribing"
|
||||
| "idle";
|
||||
|
||||
// ── Focus-aware delivery ──────────────────────────────────────────────────
|
||||
// The TUI handle is captured once via a zero-height widget factory (the only
|
||||
// extension-API surface that exposes it). With it we can:
|
||||
// 1. Listen to ALL terminal input via tui.addInputListener — listeners run
|
||||
// before the focused component, so alt+m works even while a custom
|
||||
// dialog has stolen focus from the main editor (extension shortcuts are
|
||||
// otherwise only matched by the main editor component).
|
||||
// 2. Inspect tui.focusedComponent to decide where the transcript goes.
|
||||
// `focusedComponent` is declared private in the typings but is a plain
|
||||
// runtime property — a benign peek, easily patched if pi internals change.
|
||||
interface EditorLike {
|
||||
getText(): string;
|
||||
setText(text: string): void;
|
||||
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;
|
||||
}
|
||||
type Target =
|
||||
| { kind: "editor"; editor: EditorLike }
|
||||
| { kind: "typable"; component: { handleInput(data: string): void } };
|
||||
|
||||
const asEditorLike = (value: any): EditorLike | null =>
|
||||
value && typeof value.getText === "function" && typeof value.setText === "function" ? value : null;
|
||||
export interface DictateExtensionOptions {
|
||||
socketFactory?: (path: string) => DictateSocket;
|
||||
socketPath?: string;
|
||||
}
|
||||
|
||||
// Same braille frames pi-tui's Loader uses.
|
||||
const SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
|
||||
const SPINNER_INTERVAL_MS = 80;
|
||||
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();
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
// Audio meter — a tiny rolling waveform rendered in the status row while recording.
|
||||
// Tweakable knobs:
|
||||
// METER_CELLS = how many bars wide
|
||||
// METER_TICK_MS = how often bars shift left (smaller = snappier, more renders)
|
||||
// METER_FLOOR_DB = level at which the bar is empty (more negative = more sensitive)
|
||||
// METER_CEILING_DB = level at which the bar is full (less negative = needs louder to peg)
|
||||
const METER_CELLS = 6;
|
||||
const METER_TICK_MS = 60;
|
||||
const PEAK_BLOCKS = ["▁", "▂", "▃", "▄", "▅", "▆", "▇", "█"];
|
||||
// const PEAK_BLOCKS = ["⠀", "⣀", "⣄", "⣤", "⣦", "⣶", "⣷", "⣿"];
|
||||
const METER_FLOOR_DB = -50;
|
||||
const METER_CEILING_DB = -10;
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === "object" && value !== null && !Array.isArray(value);
|
||||
}
|
||||
|
||||
/** Compute normalized RMS (0..1) over a buffer of signed 16-bit little-endian PCM samples. */
|
||||
function rmsFromPcm16(buf: Buffer): number {
|
||||
const sampleCount = Math.floor(buf.length / 2);
|
||||
if (sampleCount === 0) return 0;
|
||||
let sumSquares = 0;
|
||||
for (let i = 0; i < sampleCount * 2; i += 2) {
|
||||
const s = buf.readInt16LE(i);
|
||||
sumSquares += s * s;
|
||||
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 Math.sqrt(sumSquares / sampleCount) / 32768;
|
||||
return `${text}…`;
|
||||
}
|
||||
|
||||
/** Map a normalized RMS value to one of PEAK_BLOCKS by converting to dB and clamping into the visible range. */
|
||||
function rmsToBlock(rms: number): string {
|
||||
if (rms <= 0) return PEAK_BLOCKS[0]!;
|
||||
const db = 20 * Math.log10(rms);
|
||||
const t = Math.max(0, Math.min(1, (db - METER_FLOOR_DB) / (METER_CEILING_DB - METER_FLOOR_DB)));
|
||||
const idx = Math.floor(t * (PEAK_BLOCKS.length - 1));
|
||||
return PEAK_BLOCKS[idx]!;
|
||||
}
|
||||
export function createDictateExtension(options: DictateExtensionOptions = {}) {
|
||||
const socketFactory = options.socketFactory ?? nodeSocket;
|
||||
const socketPath = options.socketPath ?? DICTATE_SOCKET_PATH;
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
let state: State = "idle";
|
||||
let rec: ChildProcessByStdio<null, Readable, Readable> | null = null;
|
||||
let ws: WebSocket | null = null;
|
||||
let finals: string[] = [];
|
||||
let activeCtx: ExtensionContext | null = null;
|
||||
let flushed = false;
|
||||
let cancelled = false;
|
||||
let stopTimeout: NodeJS.Timeout | null = null;
|
||||
let spinnerTimer: NodeJS.Timeout | null = null;
|
||||
let spinnerFrame = 0;
|
||||
// Session generation: incremented on every start and every cleanup. All
|
||||
// rec/ws event handlers capture the generation they belong to and no-op
|
||||
// when it's stale — otherwise a PREVIOUS session's socket erroring/closing
|
||||
// late (e.g. one we aborted mid-handshake) would run cleanup() and tear
|
||||
// down the CURRENT live session.
|
||||
let generation = 0;
|
||||
// Audio meter state. `meter` is a ring of recent RMS values, newest at
|
||||
// index METER_CELLS-1. `currentLevel` is the most recent RMS observed from
|
||||
// any audio chunk — the meter tick just samples it. Crucially we never reset
|
||||
// it: empty ticks re-render the last observed value, so the bars never drop
|
||||
// to silence just because no chunk happened to arrive in that 60ms window.
|
||||
let meterTimer: NodeJS.Timeout | null = null;
|
||||
let meter: number[] = new Array(METER_CELLS).fill(0);
|
||||
let currentLevel = 0;
|
||||
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 setStatus = (msg: string | undefined) => {
|
||||
if (!activeCtx) return;
|
||||
activeCtx.ui.setStatus("dictate", msg);
|
||||
};
|
||||
const newId = (kind: string) => `${kind}-${++nextId}`;
|
||||
|
||||
const stopSpinner = () => {
|
||||
if (spinnerTimer) {
|
||||
clearInterval(spinnerTimer);
|
||||
spinnerTimer = null;
|
||||
}
|
||||
};
|
||||
|
||||
const stopMeter = () => {
|
||||
if (meterTimer) {
|
||||
clearInterval(meterTimer);
|
||||
meterTimer = null;
|
||||
}
|
||||
};
|
||||
|
||||
/** Start the meter ticking. Each tick shifts the ring and samples currentLevel. */
|
||||
const startMeter = () => {
|
||||
stopMeter();
|
||||
meter = new Array(METER_CELLS).fill(0);
|
||||
currentLevel = 0;
|
||||
// Recording dot: a text glyph colored via the theme, not an emoji — emoji
|
||||
// presentation renders double-width in its own baked-in color and visually
|
||||
// shouts in the footer. `●` is the same dot pi's own docs use for
|
||||
// indicators; theme "error" gives the red. (If you ever want strictly
|
||||
// ASCII, swap the glyph for "O".)
|
||||
const render = () => {
|
||||
const dot = activeCtx?.ui.theme.fg("error", "●") ?? "●";
|
||||
setStatus(`${dot} ${meter.map(rmsToBlock).join("")} listening…`);
|
||||
const setStatus = (status: string | undefined) => {
|
||||
activeContext?.ui.setStatus(STATUS_KEY, status);
|
||||
};
|
||||
render();
|
||||
meterTimer = setInterval(() => {
|
||||
meter.shift();
|
||||
meter.push(currentLevel);
|
||||
render();
|
||||
}, METER_TICK_MS);
|
||||
};
|
||||
|
||||
/** Animate the dictate status row with a braille spinner + suffix message. */
|
||||
const startSpinner = (suffix: string) => {
|
||||
stopSpinner();
|
||||
spinnerFrame = 0;
|
||||
setStatus(`${SPINNER_FRAMES[0]} ${suffix}`);
|
||||
spinnerTimer = setInterval(() => {
|
||||
spinnerFrame = (spinnerFrame + 1) % SPINNER_FRAMES.length;
|
||||
setStatus(`${SPINNER_FRAMES[spinnerFrame]} ${suffix}`);
|
||||
}, SPINNER_INTERVAL_MS);
|
||||
};
|
||||
|
||||
let tuiHandle: any = null;
|
||||
let removeInputListener: (() => void) | null = null;
|
||||
let lastCtx: ExtensionContext | null = null;
|
||||
|
||||
/** Resolve where dictated text would go RIGHT NOW, based on keyboard focus. */
|
||||
const resolveTarget = (): Target | null => {
|
||||
const focused = tuiHandle?.focusedComponent;
|
||||
if (!focused) return null;
|
||||
// Editor-like focus: the main chat editor, custom editors, and the
|
||||
// ctx.ui.editor()/input() popups (their inner pi-tui Editor hangs off
|
||||
// `.editor`). These accept a guaranteed direct setText append.
|
||||
const editor = asEditorLike(focused) ?? asEditorLike(focused.editor);
|
||||
if (editor) return { kind: "editor", editor };
|
||||
// Opaque component with input handling (quiz/ask selects, selectors):
|
||||
// we can type into it, but whether the text lands depends on its
|
||||
// internal focus (e.g. the quiz note field must be Tab-focused).
|
||||
if (typeof focused.handleInput === "function") return { kind: "typable", component: focused };
|
||||
return null;
|
||||
};
|
||||
|
||||
const flush = () => {
|
||||
if (flushed || !activeCtx) return;
|
||||
flushed = true;
|
||||
if (cancelled) return; // discard transcript on cancel
|
||||
const text = finals.join(" ").replace(/\s+/g, " ").trim();
|
||||
if (!text) return;
|
||||
|
||||
// Legacy fallback: no TUI handle captured (non-TUI mode / older pi) —
|
||||
// append to the main chat editor exactly as before.
|
||||
if (!tuiHandle) {
|
||||
const current = activeCtx.ui.getEditorText() ?? "";
|
||||
const sep = current && !/\s$/.test(current) ? " " : "";
|
||||
activeCtx.ui.setEditorText(current + sep + text);
|
||||
return;
|
||||
}
|
||||
|
||||
// Resolve the target NOW — focus may have changed while dictating.
|
||||
const target = resolveTarget();
|
||||
if (target?.kind === "editor") {
|
||||
const current = target.editor.getText() ?? "";
|
||||
const sep = current && !/\s$/.test(current) ? " " : "";
|
||||
target.editor.setText(current + sep + text);
|
||||
tuiHandle.requestRender?.();
|
||||
return;
|
||||
}
|
||||
if (target?.kind === "typable") {
|
||||
// Synthetic typing: the component routes the text wherever its
|
||||
// internal focus is. Text is plain printable words (whitespace
|
||||
// already normalized), so no keybindings/autocomplete can trigger.
|
||||
target.component.handleInput(text);
|
||||
tuiHandle.requestRender?.();
|
||||
return;
|
||||
}
|
||||
// Nothing to type into: don't throw the transcript away — stash it on
|
||||
// the clipboard and say so.
|
||||
try {
|
||||
const p = spawn("pbcopy", [], { stdio: ["pipe", "ignore", "ignore"] });
|
||||
p.stdin.end(text);
|
||||
} catch {}
|
||||
activeCtx.ui.notify("Dictation finished but no input field is focused — transcript copied to clipboard", "warning");
|
||||
};
|
||||
|
||||
const cleanup = () => {
|
||||
generation++; // invalidate the dying session's event handlers
|
||||
dbg(`cleanup → gen ${generation}`);
|
||||
flush();
|
||||
stopSpinner();
|
||||
stopMeter();
|
||||
if (stopTimeout) {
|
||||
clearTimeout(stopTimeout);
|
||||
stopTimeout = null;
|
||||
}
|
||||
if (rec) {
|
||||
const closeSocket = (current: DictateSocket, force: boolean) => {
|
||||
try {
|
||||
rec.kill("SIGTERM");
|
||||
} catch {}
|
||||
rec = null;
|
||||
}
|
||||
if (ws) {
|
||||
try {
|
||||
if (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING) {
|
||||
ws.close();
|
||||
}
|
||||
} catch {}
|
||||
ws = null;
|
||||
}
|
||||
finals = [];
|
||||
state = "idle";
|
||||
setStatus(undefined);
|
||||
activeCtx = null;
|
||||
flushed = false;
|
||||
cancelled = false;
|
||||
};
|
||||
if (force) current.destroy();
|
||||
else current.end();
|
||||
} catch {
|
||||
// Cleanup is best effort and must never deliver text.
|
||||
}
|
||||
};
|
||||
|
||||
const startDictation = (ctx: ExtensionContext) => {
|
||||
const apiKey = process.env.DEEPGRAM_API_KEY;
|
||||
if (!apiKey) {
|
||||
ctx.ui.notify("DEEPGRAM_API_KEY not set in environment", "error");
|
||||
return;
|
||||
}
|
||||
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);
|
||||
};
|
||||
|
||||
activeCtx = ctx;
|
||||
finals = [];
|
||||
flushed = false;
|
||||
cancelled = false;
|
||||
state = "recording";
|
||||
const myGeneration = ++generation;
|
||||
dbg(`start (gen ${myGeneration})`);
|
||||
startMeter();
|
||||
|
||||
// Spawn sox `rec` to capture 16kHz / 16-bit / mono PCM to stdout.
|
||||
let proc: ChildProcessByStdio<null, Readable, Readable>;
|
||||
try {
|
||||
proc = spawn(
|
||||
"rec",
|
||||
[
|
||||
"-q", // quiet
|
||||
// Shrink sox's IO buffer so stdout flushes ~every 16ms instead of
|
||||
// the default ~256ms. 512 bytes = 256 samples = 16ms at 16kHz/16-bit
|
||||
// mono. This is the dominant source of meter latency.
|
||||
"--buffer", "512",
|
||||
"-r", "16000",
|
||||
"-c", "1",
|
||||
"-b", "16",
|
||||
"-e", "signed-integer",
|
||||
"-t", "raw",
|
||||
"-", // stdout
|
||||
],
|
||||
{ stdio: ["ignore", "pipe", "pipe"] },
|
||||
const notifyError = (message: string) => {
|
||||
activeContext?.ui.notify(
|
||||
boundedText(message, MAX_ERROR_DETAIL_BYTES),
|
||||
"error",
|
||||
);
|
||||
} catch (e: any) {
|
||||
ctx.ui.notify(`Failed to spawn 'rec'. Install sox: brew install sox`, "error");
|
||||
cleanup();
|
||||
return;
|
||||
}
|
||||
rec = proc;
|
||||
};
|
||||
|
||||
proc.on("error", (err) => {
|
||||
if (myGeneration !== generation) return;
|
||||
ctx.ui.notify(`rec error: ${err.message} (install sox: brew install sox)`, "error");
|
||||
cleanup();
|
||||
});
|
||||
const fail = (message: string) => {
|
||||
notifyError(message);
|
||||
reset(true);
|
||||
};
|
||||
|
||||
proc.on("exit", (code) => {
|
||||
if (myGeneration !== generation) return; // stale recorder — a newer/ended session owns state
|
||||
// Natural exit on SIGTERM during stopDictation is fine. Anything else
|
||||
// mid-recording is a problem.
|
||||
if (state === "recording" && code !== null && code !== 0) {
|
||||
if (activeCtx) {
|
||||
activeCtx.ui.notify(`rec exited unexpectedly (code ${code})`, "warning");
|
||||
}
|
||||
cleanup();
|
||||
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;
|
||||
}
|
||||
});
|
||||
|
||||
// Open Deepgram WebSocket. Auth via subprotocol (portable across Node native
|
||||
// WebSocket and browsers): `new WebSocket(url, ["token", API_KEY])`.
|
||||
try {
|
||||
ws = new WebSocket(DG_URL, ["token", apiKey]);
|
||||
} catch (e: any) {
|
||||
ctx.ui.notify(`Deepgram WS failed: ${e.message}`, "error");
|
||||
cleanup();
|
||||
return;
|
||||
}
|
||||
|
||||
ws.addEventListener("open", () => {
|
||||
if (myGeneration !== generation) {
|
||||
dbg(`ws open (stale gen ${myGeneration}, current ${generation}) — ignored`);
|
||||
return;
|
||||
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;
|
||||
}
|
||||
dbg(`ws open (gen ${myGeneration})`);
|
||||
if (!rec || !ws) return;
|
||||
rec.stdout.on("data", (chunk: Buffer) => {
|
||||
// Track loudness for the meter (just the latest chunk's RMS — the meter
|
||||
// tick samples this), then forward to Deepgram.
|
||||
currentLevel = rmsFromPcm16(chunk);
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
ws.send(chunk);
|
||||
}
|
||||
};
|
||||
|
||||
const sendGenerationCommand = (
|
||||
type: "start" | "stop" | "cancel",
|
||||
): boolean => {
|
||||
if (!activeGenerationId || !negotiated) return false;
|
||||
return send({
|
||||
version: PROTOCOL_VERSION,
|
||||
type,
|
||||
requestId: newId(type),
|
||||
generationId: activeGenerationId,
|
||||
});
|
||||
});
|
||||
};
|
||||
|
||||
ws.addEventListener("message", (ev) => {
|
||||
if (myGeneration !== generation) return;
|
||||
try {
|
||||
const msg = JSON.parse(ev.data as string);
|
||||
if (msg.type === "Results" && msg.is_final) {
|
||||
const t = msg.channel?.alternatives?.[0]?.transcript;
|
||||
if (t) finals.push(t);
|
||||
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;
|
||||
}
|
||||
// We could also handle msg.type === "Metadata" (sent after CloseStream
|
||||
// finishes draining), but ws.close handles the same flush path.
|
||||
} catch {
|
||||
// ignore non-JSON frames
|
||||
}
|
||||
});
|
||||
|
||||
ws.addEventListener("error", () => {
|
||||
if (myGeneration !== generation) {
|
||||
dbg(`ws error (stale gen ${myGeneration}, current ${generation}) — ignored`);
|
||||
negotiated = true;
|
||||
sendGenerationCommand("start");
|
||||
return;
|
||||
}
|
||||
dbg(`ws error (gen ${myGeneration})`);
|
||||
if (activeCtx) activeCtx.ui.notify("Deepgram WebSocket error", "error");
|
||||
cleanup();
|
||||
});
|
||||
|
||||
ws.addEventListener("close", (ev) => {
|
||||
if (myGeneration !== generation) {
|
||||
dbg(`ws close (stale gen ${myGeneration}, current ${generation}, code ${ev.code}) — ignored`);
|
||||
if (!negotiated) {
|
||||
protocolError("generation event received before hello negotiation.");
|
||||
return;
|
||||
}
|
||||
dbg(`ws close (gen ${myGeneration}, code ${ev.code})`);
|
||||
// Server-initiated close (or our own close in cleanup): finalize.
|
||||
if (state === "recording" || state === "stopping") {
|
||||
cleanup();
|
||||
|
||||
const generationId = value.generationId;
|
||||
if (!isBoundedId(generationId)) {
|
||||
protocolError("generation ID is invalid.");
|
||||
return;
|
||||
}
|
||||
});
|
||||
};
|
||||
if (generationId !== activeGenerationId) return;
|
||||
|
||||
/** Stop dictation, finalize transcript, append to editor. */
|
||||
const stopDictation = () => {
|
||||
if (state !== "recording") return;
|
||||
state = "stopping";
|
||||
stopMeter();
|
||||
startSpinner("finalizing…");
|
||||
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;
|
||||
}
|
||||
|
||||
// Stop the mic first so no more audio enqueues.
|
||||
if (rec) {
|
||||
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 {
|
||||
rec.kill("SIGTERM");
|
||||
} catch {}
|
||||
}
|
||||
|
||||
// Tell Deepgram we're done; it will flush remaining finals then close.
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
try {
|
||||
ws.send(JSON.stringify({ type: "CloseStream" }));
|
||||
} catch {
|
||||
cleanup();
|
||||
current = socketFactory(socketPath);
|
||||
} catch (error) {
|
||||
const detail = error instanceof Error ? error.message : "connection failed";
|
||||
fail(`Unable to connect to dictation service: ${detail}`);
|
||||
return;
|
||||
}
|
||||
// Safety net: if Deepgram never closes the socket, force cleanup after 3s.
|
||||
stopTimeout = setTimeout(() => {
|
||||
if (state === "stopping") cleanup();
|
||||
}, 3000);
|
||||
} else {
|
||||
cleanup();
|
||||
}
|
||||
};
|
||||
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.");
|
||||
});
|
||||
};
|
||||
|
||||
/** Cancel dictation: discard any collected transcript and tear everything down immediately. */
|
||||
const cancelDictation = () => {
|
||||
if (state !== "recording" && state !== "stopping") return;
|
||||
cancelled = true;
|
||||
finals = [];
|
||||
// No need to wait for Deepgram to flush — we're throwing the result away.
|
||||
cleanup();
|
||||
};
|
||||
const stopDictation = () => {
|
||||
if (state !== "recording") return;
|
||||
state = "stopping";
|
||||
sendGenerationCommand("stop");
|
||||
};
|
||||
|
||||
/** Toggle dictation, gated on there being somewhere for the text to go. */
|
||||
const toggleDictation = (ctx: ExtensionContext) => {
|
||||
lastCtx = ctx;
|
||||
if (state === "idle") {
|
||||
if (tuiHandle && !resolveTarget()) {
|
||||
ctx.ui.notify("No input field is focused — dictation not started", "warning");
|
||||
return;
|
||||
}
|
||||
startDictation(ctx);
|
||||
} else if (state === "recording") {
|
||||
stopDictation();
|
||||
}
|
||||
// Ignore presses during the "stopping" state — Deepgram is finalizing.
|
||||
};
|
||||
const cancelDictation = () => {
|
||||
if (state === "idle") return;
|
||||
if (negotiated) sendGenerationCommand("cancel");
|
||||
reset();
|
||||
};
|
||||
|
||||
// Global input listener: catches alt+m/alt+n before ANY focused component,
|
||||
// which is what makes dictation work inside dialogs. Registered once the
|
||||
// TUI handle is captured (see session_start below).
|
||||
const onGlobalInput = (data: string) => {
|
||||
// Kitty flag-2 terminals send press + REPEAT + RELEASE events, and input
|
||||
// listeners run BEFORE the TUI's release filter (that filter only guards
|
||||
// dispatch to the focused component). matchesKey also ignores the Kitty
|
||||
// event type. Without this guard a single physical alt+m press toggles
|
||||
// TWICE: press starts dictation, release instantly stops it and closes
|
||||
// the WebSocket mid-handshake — which then surfaces as
|
||||
// "Deepgram WebSocket error" (and its stale error event can kill the NEXT
|
||||
// session). Filter to press events only.
|
||||
if (isKeyRelease(data) || isKeyRepeat(data)) return undefined;
|
||||
if (matchesKey(data, Key.alt("m"))) {
|
||||
dbg(`alt+m (data=${JSON.stringify(data)}) state=${state}`);
|
||||
if (lastCtx) toggleDictation(lastCtx);
|
||||
return { consume: true };
|
||||
}
|
||||
if (matchesKey(data, Key.alt("n"))) {
|
||||
dbg(`alt+n (data=${JSON.stringify(data)}) state=${state}`);
|
||||
cancelDictation();
|
||||
return { consume: true };
|
||||
}
|
||||
return undefined;
|
||||
};
|
||||
const toggleDictation = (ctx: ExtensionContext) => {
|
||||
if (state === "idle") startDictation(ctx);
|
||||
else if (state === "recording") stopDictation();
|
||||
};
|
||||
|
||||
pi.on("session_start", (_event, ctx) => {
|
||||
lastCtx = ctx;
|
||||
if (ctx.mode !== "tui" || tuiHandle) return;
|
||||
// Capture the TUI handle via an invisible zero-height widget. The
|
||||
// listener function reference is stable, so even if the factory re-runs
|
||||
// the TUI's listener Set de-dupes it.
|
||||
ctx.ui.setWidget("dictate-tui-handle", (tui: any) => {
|
||||
tuiHandle = tui;
|
||||
removeInputListener = tui.addInputListener(onGlobalInput);
|
||||
return { render: () => [], invalidate: () => {} };
|
||||
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();
|
||||
});
|
||||
};
|
||||
}
|
||||
|
||||
// Shortcut registrations kept as a fallback for contexts where the TUI
|
||||
// handle was never captured (non-TUI modes, older pi): they only fire when
|
||||
// the main editor is focused, but that's precisely the legacy path. When
|
||||
// the listener IS installed it consumes the key first, so no double-fire.
|
||||
pi.registerShortcut(Key.alt("m"), {
|
||||
description: "Toggle voice dictation (Deepgram)",
|
||||
handler: async (ctx) => {
|
||||
toggleDictation(ctx);
|
||||
},
|
||||
});
|
||||
|
||||
// Dedicated cancel binding. Dictation-only — a no-op when no dictation is
|
||||
// in flight, so it's safe to hammer without affecting anything else.
|
||||
pi.registerShortcut(Key.alt("n"), {
|
||||
description: "Cancel voice dictation (discard transcript)",
|
||||
handler: async () => {
|
||||
cancelDictation();
|
||||
},
|
||||
});
|
||||
|
||||
pi.on("session_shutdown", () => {
|
||||
if (state !== "idle") cleanup();
|
||||
removeInputListener?.();
|
||||
removeInputListener = null;
|
||||
});
|
||||
}
|
||||
export default createDictateExtension();
|
||||
|
||||
+2
-1
@@ -13,8 +13,9 @@
|
||||
"scripts": {
|
||||
"prepare": "effect-tsgo patch",
|
||||
"typecheck": "tsc --noEmit",
|
||||
"test": "npm run test:runner && npm run test:links && node --test --experimental-strip-types extensions/firecrawl-search/index.test.ts extensions/shared/*.test.ts extensions/subagents/by-the-way.test.ts extensions/subagents/context-usage.test.ts extensions/subagents/manager.test.ts extensions/subagents/result-delivery.test.ts extensions/subagents/takeover.test.ts extensions/_workflows/*.test.ts extensions/pi-workflow/*.test.ts && vitest run extensions/file-search/index.spec.ts",
|
||||
"test": "npm run test:runner && npm run test:dictate-service && npm run test:links && node --test --experimental-strip-types extensions/dictate/index.test.ts extensions/firecrawl-search/index.test.ts extensions/shared/*.test.ts extensions/subagents/by-the-way.test.ts extensions/subagents/context-usage.test.ts extensions/subagents/manager.test.ts extensions/subagents/result-delivery.test.ts extensions/subagents/takeover.test.ts extensions/_workflows/*.test.ts extensions/pi-workflow/*.test.ts && vitest run extensions/file-search/index.spec.ts",
|
||||
"test:runner": "python -m unittest discover -s scripts -p 'test_*.py'",
|
||||
"test:dictate-service": "PYTHONPATH=services/dictate/src python -m unittest discover -s services/dictate/tests -p 'test_*.py'",
|
||||
"test:links": "node --test scripts/pi-links.test.mjs",
|
||||
"test:live": "node --test --experimental-strip-types extensions/subagents/claude.test.ts extensions/subagents/codex.test.ts"
|
||||
},
|
||||
|
||||
@@ -20,6 +20,9 @@
|
||||
"allow_file": [
|
||||
"/dev/ptmx"
|
||||
],
|
||||
"deny": [
|
||||
"$HOME/.local/share/pi-dictate"
|
||||
],
|
||||
"read_file": [
|
||||
"$HOME/.local/bin/claude",
|
||||
"$HOME/.local/bin/codex"
|
||||
@@ -31,7 +34,8 @@
|
||||
],
|
||||
"unix_socket": [
|
||||
"/run/systemd/resolve/io.systemd.Resolve",
|
||||
"/home/jb/.config/herdr/herdr.sock"
|
||||
"/home/jb/.config/herdr/herdr.sock",
|
||||
"/run/user/1000/pi-dictate.sock"
|
||||
]
|
||||
},
|
||||
"linux": {
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
.venv/
|
||||
build/
|
||||
src/*.egg-info/
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
@@ -0,0 +1 @@
|
||||
3.12.13
|
||||
@@ -0,0 +1,37 @@
|
||||
API
|
||||
async
|
||||
await
|
||||
Bun
|
||||
CLI
|
||||
CUDA
|
||||
Git
|
||||
git commit
|
||||
git status
|
||||
GitHub
|
||||
GitLab
|
||||
Ghostty
|
||||
JSON
|
||||
JSON response
|
||||
NeMo
|
||||
niri
|
||||
Node.js
|
||||
nono
|
||||
npm
|
||||
npm install
|
||||
npm test
|
||||
Parakeet
|
||||
Pi
|
||||
PipeWire
|
||||
Python
|
||||
PyTorch
|
||||
stderr
|
||||
stdout
|
||||
Svelte
|
||||
SvelteKit
|
||||
systemd
|
||||
TUI
|
||||
TypeScript
|
||||
Unix socket
|
||||
Vite
|
||||
Wayland
|
||||
WebSocket
|
||||
Executable
+183
@@ -0,0 +1,183 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
MODEL_REVISION="541d1f99c6b0c3cd0b11a95167540bb8edefd82b"
|
||||
MODEL_SHA256="3cbdc85877e668ca7b82d0d56770eb1fac76691f55d6b97545e8d61ca588d10d"
|
||||
MODEL_FILENAME="parakeet-tdt-0.6b-v3.nemo"
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPOSITORY_ROOT="$(cd -- "$SCRIPT_DIR/../.." && pwd)"
|
||||
INSTALL_ROOT="$HOME/.local/share/pi-dictate"
|
||||
UNIT_DIR="${XDG_CONFIG_HOME:-$HOME/.config}/systemd/user"
|
||||
MODEL_DIR="$INSTALL_ROOT/models/parakeet-tdt-0.6b-v3"
|
||||
MODEL_PATH="$MODEL_DIR/$MODEL_FILENAME"
|
||||
DICTIONARY_SOURCE="$SCRIPT_DIR/coding-terms.txt"
|
||||
DICTIONARY_PATH="$INSTALL_ROOT/config/coding-terms.txt"
|
||||
DRY_RUN=false
|
||||
MODEL_SOURCE=""
|
||||
ACTION=install
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage: install.sh [--dry-run] [--model-source PATH]
|
||||
install.sh uninstall [--dry-run]
|
||||
|
||||
Installs the pinned service environment, model snapshot, and systemd user units.
|
||||
Without --model-source, installation downloads the exact pinned model revision.
|
||||
EOF
|
||||
}
|
||||
|
||||
run() {
|
||||
printf '+ '
|
||||
printf '%q ' "$@"
|
||||
printf '\n'
|
||||
if ! $DRY_RUN; then
|
||||
"$@"
|
||||
fi
|
||||
}
|
||||
|
||||
is_owned_unit() {
|
||||
local destination="$1"
|
||||
local source="$2"
|
||||
[[ -L "$destination" && "$(readlink -f -- "$destination")" == "$(readlink -f -- "$source")" ]] ||
|
||||
{ [[ -f "$destination" ]] && grep -qxF '# Managed by pi-dictate install.sh' "$destination"; }
|
||||
}
|
||||
|
||||
install_owned_unit() {
|
||||
local source="$1"
|
||||
local destination="$2"
|
||||
if [[ -L "$destination" && "$(readlink -f -- "$destination")" == "$(readlink -f -- "$source")" ]]; then
|
||||
run rm "$destination"
|
||||
elif [[ -e "$destination" || -L "$destination" ]]; then
|
||||
if [[ ! -f "$destination" ]] || ! grep -qxF '# Managed by pi-dictate install.sh' "$destination"; then
|
||||
printf 'Refusing to replace unmanaged systemd unit: %s\n' "$destination" >&2
|
||||
return 1
|
||||
fi
|
||||
fi
|
||||
run install -m 0644 "$source" "$destination"
|
||||
}
|
||||
|
||||
remove_owned_unit() {
|
||||
local destination="$1"
|
||||
local source="$2"
|
||||
if [[ -L "$destination" && "$(readlink -f -- "$destination")" == "$(readlink -f -- "$source")" ]]; then
|
||||
run rm "$destination"
|
||||
elif [[ -f "$destination" ]] && grep -qxF '# Managed by pi-dictate install.sh' "$destination"; then
|
||||
run rm "$destination"
|
||||
elif [[ -e "$destination" || -L "$destination" ]]; then
|
||||
printf 'Leaving unmanaged unit in place: %s\n' "$destination" >&2
|
||||
fi
|
||||
}
|
||||
|
||||
while (($#)); do
|
||||
case "$1" in
|
||||
uninstall) ACTION=uninstall ;;
|
||||
--dry-run) DRY_RUN=true ;;
|
||||
--model-source)
|
||||
shift
|
||||
(($#)) || { printf '%s\n' '--model-source requires a path' >&2; exit 2; }
|
||||
MODEL_SOURCE="$1"
|
||||
;;
|
||||
-h|--help) usage; exit 0 ;;
|
||||
*) printf 'Unknown argument: %s\n' "$1" >&2; usage >&2; exit 2 ;;
|
||||
esac
|
||||
shift
|
||||
done
|
||||
|
||||
socket_source="$REPOSITORY_ROOT/systemd/user/pi-dictate.socket"
|
||||
service_source="$REPOSITORY_ROOT/systemd/user/pi-dictate.service"
|
||||
socket_link="$UNIT_DIR/pi-dictate.socket"
|
||||
service_link="$UNIT_DIR/pi-dictate.service"
|
||||
|
||||
if [[ "$ACTION" == uninstall ]]; then
|
||||
for pair in "$socket_link:$socket_source" "$service_link:$service_source"; do
|
||||
destination="${pair%%:*}"
|
||||
source="${pair#*:}"
|
||||
if [[ ( -e "$destination" || -L "$destination" ) ]] && ! is_owned_unit "$destination" "$source"; then
|
||||
printf 'Refusing to disable or remove unmanaged systemd unit: %s\n' "$destination" >&2
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
run systemctl --user disable --now pi-dictate.socket
|
||||
remove_owned_unit "$socket_link" "$socket_source"
|
||||
remove_owned_unit "$service_link" "$service_source"
|
||||
run systemctl --user daemon-reload
|
||||
run rm -rf "$INSTALL_ROOT"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
command -v uv >/dev/null || { printf '%s\n' 'uv is required' >&2; exit 1; }
|
||||
command -v systemctl >/dev/null || { printf '%s\n' 'systemctl is required' >&2; exit 1; }
|
||||
|
||||
run mkdir -p "$INSTALL_ROOT" "$MODEL_DIR" "$(dirname "$DICTIONARY_PATH")" "$UNIT_DIR"
|
||||
run chmod 0700 "$INSTALL_ROOT" "$(dirname "$MODEL_DIR")" "$MODEL_DIR" "$(dirname "$DICTIONARY_PATH")"
|
||||
if $DRY_RUN; then
|
||||
printf '+ UV_PROJECT_ENVIRONMENT=%q uv sync --frozen --no-dev --no-editable --project %q\n' "$INSTALL_ROOT/venv" "$SCRIPT_DIR"
|
||||
else
|
||||
UV_PROJECT_ENVIRONMENT="$INSTALL_ROOT/venv" uv sync --frozen --no-dev --no-editable --project "$SCRIPT_DIR"
|
||||
fi
|
||||
|
||||
MODEL_TEMP="$MODEL_DIR/.${MODEL_FILENAME}.install.$$"
|
||||
cleanup_model_temp() {
|
||||
if ! $DRY_RUN; then
|
||||
rm -f -- "$MODEL_TEMP"
|
||||
fi
|
||||
}
|
||||
trap cleanup_model_temp EXIT
|
||||
|
||||
if [[ -n "$MODEL_SOURCE" ]]; then
|
||||
[[ -f "$MODEL_SOURCE" ]] || { printf 'Model source not found: %s\n' "$MODEL_SOURCE" >&2; exit 1; }
|
||||
run cp --reflink=auto "$MODEL_SOURCE" "$MODEL_TEMP"
|
||||
else
|
||||
if $DRY_RUN; then
|
||||
printf '+ download nvidia/parakeet-tdt-0.6b-v3@%s to %q\n' "$MODEL_REVISION" "$MODEL_TEMP"
|
||||
else
|
||||
MODEL_REVISION="$MODEL_REVISION" MODEL_PATH="$MODEL_TEMP" \
|
||||
"$INSTALL_ROOT/venv/bin/python" - <<'PY'
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
source = hf_hub_download(
|
||||
repo_id="nvidia/parakeet-tdt-0.6b-v3",
|
||||
filename="parakeet-tdt-0.6b-v3.nemo",
|
||||
revision=os.environ["MODEL_REVISION"],
|
||||
)
|
||||
destination = Path(os.environ["MODEL_PATH"])
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copyfile(source, destination)
|
||||
PY
|
||||
fi
|
||||
fi
|
||||
|
||||
if ! $DRY_RUN; then
|
||||
actual_sha256="$(sha256sum "$MODEL_TEMP" | cut -d' ' -f1)"
|
||||
if [[ "$actual_sha256" != "$MODEL_SHA256" ]]; then
|
||||
rm -f "$MODEL_TEMP"
|
||||
printf 'Model checksum mismatch: expected %s, got %s\n' "$MODEL_SHA256" "$actual_sha256" >&2
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
run chmod 0400 "$MODEL_TEMP"
|
||||
run mv -f "$MODEL_TEMP" "$MODEL_PATH"
|
||||
trap - EXIT
|
||||
|
||||
if [[ ! -e "$DICTIONARY_PATH" ]]; then
|
||||
run cp "$DICTIONARY_SOURCE" "$DICTIONARY_PATH"
|
||||
run chmod 0600 "$DICTIONARY_PATH"
|
||||
else
|
||||
printf 'Keeping existing coding dictionary: %s\n' "$DICTIONARY_PATH"
|
||||
fi
|
||||
|
||||
install_owned_unit "$socket_source" "$socket_link"
|
||||
install_owned_unit "$service_source" "$service_link"
|
||||
run systemctl --user daemon-reload
|
||||
run systemctl --user enable --now pi-dictate.socket
|
||||
|
||||
if $DRY_RUN; then
|
||||
printf 'Would install pi-dictate with model revision %s.\n' "$MODEL_REVISION"
|
||||
else
|
||||
printf 'Installed pi-dictate with model revision %s.\n' "$MODEL_REVISION"
|
||||
printf 'Check it with: systemctl --user status pi-dictate.socket pi-dictate.service\n'
|
||||
fi
|
||||
@@ -0,0 +1,23 @@
|
||||
[project]
|
||||
name = "pi-dictate-service"
|
||||
version = "0.1.0"
|
||||
description = "Socket-activated local transcription service for the Pi dictate extension"
|
||||
requires-python = ">=3.12,<3.13"
|
||||
dependencies = [
|
||||
"nemo-toolkit[asr]==3.0.0",
|
||||
"numpy==2.5.2",
|
||||
"torch==2.13.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
pi-dictate-service = "pi_dictate_service.server:main"
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools==84.0.0"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
|
||||
[tool.uv]
|
||||
package = true
|
||||
@@ -0,0 +1,5 @@
|
||||
"""GPU-isolated local dictate service."""
|
||||
|
||||
from .protocol import PROTOCOL_VERSION
|
||||
|
||||
__all__ = ["PROTOCOL_VERSION"]
|
||||
@@ -0,0 +1,22 @@
|
||||
"""In-memory PCM validation and conversion."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from array import array
|
||||
import struct
|
||||
|
||||
MAX_PCM_BYTES = 5 * 60 * 16_000 * 2
|
||||
|
||||
|
||||
class AudioError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
def pcm_s16le_to_floats(pcm: bytes, *, max_bytes: int = MAX_PCM_BYTES) -> array[float]:
|
||||
"""Convert mono little-endian signed-16 PCM to a bounded float32 array."""
|
||||
|
||||
if len(pcm) > max_bytes:
|
||||
raise AudioError("PCM exceeds the five-minute duration limit")
|
||||
if len(pcm) % 2:
|
||||
raise AudioError("PCM byte length must be even")
|
||||
return array("f", (sample[0] / 32768.0 for sample in struct.iter_unpack("<h", pcm)))
|
||||
@@ -0,0 +1,138 @@
|
||||
"""Lazy, offline NeMo adapter for an explicit local Parakeet archive."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from copy import deepcopy
|
||||
import hashlib
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Protocol, Sequence
|
||||
|
||||
MAX_TERMS_FILE_BYTES = 64 * 1024
|
||||
MAX_CODING_TERMS = 512
|
||||
MAX_TERM_BYTES = 128
|
||||
PINNED_MODEL_SHA256 = "3cbdc85877e668ca7b82d0d56770eb1fac76691f55d6b97545e8d61ca588d10d"
|
||||
|
||||
|
||||
class TranscriptionModel(Protocol):
|
||||
def load(self) -> None: ...
|
||||
def transcribe(self, samples: Sequence[float]) -> str: ...
|
||||
|
||||
|
||||
def load_coding_terms(path: str | os.PathLike[str]) -> list[str]:
|
||||
"""Load a bounded host-owned phrase list, ignoring comments and duplicates."""
|
||||
|
||||
terms_path = Path(path).expanduser()
|
||||
if not terms_path.is_file():
|
||||
raise ValueError("PI_DICTATE_TERMS_PATH must name an existing local dictionary")
|
||||
if terms_path.stat().st_size > MAX_TERMS_FILE_BYTES:
|
||||
raise ValueError(f"coding dictionary exceeds {MAX_TERMS_FILE_BYTES} bytes")
|
||||
|
||||
terms: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for raw_line in terms_path.read_text(encoding="utf-8").splitlines():
|
||||
term = raw_line.strip()
|
||||
if not term or term.startswith("#"):
|
||||
continue
|
||||
if len(term.encode("utf-8")) > MAX_TERM_BYTES:
|
||||
raise ValueError(f"coding dictionary terms may not exceed {MAX_TERM_BYTES} UTF-8 bytes")
|
||||
if term not in seen:
|
||||
terms.append(term)
|
||||
seen.add(term)
|
||||
if len(terms) > MAX_CODING_TERMS:
|
||||
raise ValueError(f"coding dictionary may not exceed {MAX_CODING_TERMS} terms")
|
||||
return terms
|
||||
|
||||
|
||||
class NemoParakeetModel:
|
||||
"""Load NeMo only when requested so core tests require neither CUDA nor NeMo."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str | os.PathLike[str],
|
||||
*,
|
||||
terms_path: str | os.PathLike[str] | None = None,
|
||||
boost_alpha: float = 1.0,
|
||||
expected_sha256: str = PINNED_MODEL_SHA256,
|
||||
) -> None:
|
||||
path = Path(model_path).expanduser()
|
||||
if path.suffix != ".nemo" or not path.is_file():
|
||||
raise ValueError("PI_DICTATE_MODEL_PATH must name an existing local .nemo file")
|
||||
if boost_alpha <= 0:
|
||||
raise ValueError("PI_DICTATE_BOOST_ALPHA must be positive")
|
||||
self._model_path = path.resolve()
|
||||
self._terms = load_coding_terms(terms_path) if terms_path is not None else []
|
||||
self._boost_alpha = boost_alpha
|
||||
self._expected_sha256 = expected_sha256
|
||||
self._model: object | None = None
|
||||
self._numpy: object | None = None
|
||||
|
||||
@property
|
||||
def loaded(self) -> bool:
|
||||
return self._model is not None
|
||||
|
||||
def load(self) -> None:
|
||||
if self._model is not None:
|
||||
return
|
||||
with self._model_path.open("rb") as model_file:
|
||||
actual_sha256 = hashlib.file_digest(model_file, "sha256").hexdigest()
|
||||
if actual_sha256 != self._expected_sha256:
|
||||
raise ValueError(
|
||||
"model checksum mismatch: expected the pinned Parakeet revision"
|
||||
)
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||||
os.environ["HF_DATASETS_OFFLINE"] = "1"
|
||||
os.environ["NEMO_OFFLINE"] = "1"
|
||||
|
||||
import numpy
|
||||
from nemo.collections.asr.models import ASRModel
|
||||
from omegaconf import open_dict
|
||||
|
||||
model = ASRModel.restore_from(
|
||||
restore_path=str(self._model_path),
|
||||
map_location="cuda",
|
||||
)
|
||||
model.to("cuda")
|
||||
model.eval()
|
||||
if self._terms:
|
||||
decoding = deepcopy(model.cfg.decoding)
|
||||
with open_dict(decoding):
|
||||
greedy = decoding["greedy"]
|
||||
greedy["boosting_tree"] = {
|
||||
"key_phrases_list": list(self._terms),
|
||||
"context_score": 1.0,
|
||||
"depth_scaling": 2.0,
|
||||
"unk_score": 0.0,
|
||||
"final_eos_score": 1.0,
|
||||
"score_per_phrase": 0.0,
|
||||
"source_lang": "en",
|
||||
"use_triton": True,
|
||||
"uniform_weights": False,
|
||||
"use_bpe_dropout": False,
|
||||
}
|
||||
greedy["boosting_tree_alpha"] = self._boost_alpha
|
||||
model.change_decoding_strategy(decoding, verbose=False)
|
||||
self._numpy = numpy
|
||||
self._model = model
|
||||
|
||||
def transcribe(self, samples: Sequence[float]) -> str:
|
||||
if self._model is None or self._numpy is None:
|
||||
raise RuntimeError("model is not loaded")
|
||||
numpy = self._numpy
|
||||
audio = numpy.asarray(samples, dtype=numpy.float32) # type: ignore[attr-defined]
|
||||
output: Sequence[object] = self._model.transcribe( # type: ignore[attr-defined]
|
||||
audio=[audio],
|
||||
batch_size=1,
|
||||
num_workers=0,
|
||||
verbose=False,
|
||||
)
|
||||
if not output:
|
||||
return ""
|
||||
hypothesis = output[0]
|
||||
if isinstance(hypothesis, str):
|
||||
return hypothesis
|
||||
text = getattr(hypothesis, "text", None)
|
||||
if isinstance(text, str):
|
||||
return text
|
||||
raise RuntimeError("NeMo returned an unsupported transcription result")
|
||||
@@ -0,0 +1,195 @@
|
||||
"""Bounded newline-delimited JSON protocol primitives for dictate protocol v1."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import json
|
||||
from typing import Any, Mapping
|
||||
|
||||
PROTOCOL_VERSION = 1
|
||||
MAX_FRAME_BYTES = 64 * 1024
|
||||
MAX_TRANSCRIPT_BYTES = 32 * 1024
|
||||
MAX_ERROR_DETAIL_BYTES = 2 * 1024
|
||||
MAX_ID_BYTES = 128
|
||||
|
||||
|
||||
class ProtocolError(ValueError):
|
||||
"""A stable protocol error suitable for returning to a client."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
code: str,
|
||||
detail: str,
|
||||
*,
|
||||
close_connection: bool = False,
|
||||
request_id: str = "",
|
||||
generation_id: str | None = None,
|
||||
) -> None:
|
||||
super().__init__(detail)
|
||||
self.code = code
|
||||
self.detail = detail
|
||||
self.close_connection = close_connection
|
||||
self.request_id = request_id
|
||||
self.generation_id = generation_id
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Request:
|
||||
version: int
|
||||
type: str
|
||||
request_id: str
|
||||
generation_id: str | None = None
|
||||
|
||||
|
||||
_REQUEST_FIELDS = {
|
||||
"hello": frozenset({"version", "type", "requestId"}),
|
||||
"ping": frozenset({"version", "type", "requestId"}),
|
||||
"start": frozenset({"version", "type", "requestId", "generationId"}),
|
||||
"stop": frozenset({"version", "type", "requestId", "generationId"}),
|
||||
"cancel": frozenset({"version", "type", "requestId", "generationId"}),
|
||||
}
|
||||
|
||||
|
||||
def _bounded_utf8(value: str, limit: int) -> str:
|
||||
raw = value.encode("utf-8", errors="replace")
|
||||
if len(raw) <= limit:
|
||||
return value
|
||||
return raw[:limit].decode("utf-8", errors="ignore")
|
||||
|
||||
|
||||
def _validate_id(value: Any, field: str) -> str:
|
||||
if not isinstance(value, str) or not value:
|
||||
raise ProtocolError("invalid_request", f"{field} must be a non-empty string")
|
||||
try:
|
||||
encoded = value.encode("utf-8")
|
||||
except UnicodeEncodeError as error:
|
||||
raise ProtocolError("invalid_request", f"{field} must contain valid Unicode") from error
|
||||
if len(encoded) > MAX_ID_BYTES:
|
||||
raise ProtocolError("invalid_request", f"{field} exceeds {MAX_ID_BYTES} UTF-8 bytes")
|
||||
return value
|
||||
|
||||
|
||||
def _safe_id(value: Any) -> str:
|
||||
try:
|
||||
return _validate_id(value, "identifier")
|
||||
except ProtocolError:
|
||||
return ""
|
||||
|
||||
|
||||
def parse_request(frame: Mapping[str, Any]) -> Request:
|
||||
"""Validate a decoded request and reject all fields outside its exact schema."""
|
||||
|
||||
if not isinstance(frame, Mapping):
|
||||
raise ProtocolError("invalid_request", "request must be a JSON object")
|
||||
|
||||
request_id_value = frame.get("requestId")
|
||||
request_id = _safe_id(request_id_value)
|
||||
version = frame.get("version")
|
||||
if not isinstance(version, int) or isinstance(version, bool):
|
||||
raise ProtocolError("invalid_request", "version must be an integer", request_id=request_id)
|
||||
if version != PROTOCOL_VERSION:
|
||||
raise ProtocolError(
|
||||
"unsupported_version",
|
||||
f"protocol version {version} is unsupported",
|
||||
close_connection=True,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
request_type = frame.get("type")
|
||||
if not isinstance(request_type, str) or request_type not in _REQUEST_FIELDS:
|
||||
raise ProtocolError("invalid_request", "unknown request type", request_id=request_id)
|
||||
|
||||
expected_fields = _REQUEST_FIELDS[request_type]
|
||||
if frozenset(frame) != expected_fields:
|
||||
raise ProtocolError(
|
||||
"invalid_request",
|
||||
f"{request_type} requires exactly: {', '.join(sorted(expected_fields))}",
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
request_id = _validate_id(request_id_value, "requestId")
|
||||
generation_id = None
|
||||
if "generationId" in expected_fields:
|
||||
generation_id = _validate_id(frame.get("generationId"), "generationId")
|
||||
|
||||
return Request(version, request_type, request_id, generation_id)
|
||||
|
||||
|
||||
class FrameDecoder:
|
||||
"""Incrementally decode bounded newline-delimited JSON objects."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._buffer = bytearray()
|
||||
self._failed = False
|
||||
|
||||
def feed(self, data: bytes) -> list[dict[str, Any]]:
|
||||
if self._failed:
|
||||
raise ProtocolError("connection_closed", "decoder is closed", close_connection=True)
|
||||
self._buffer.extend(data)
|
||||
frames: list[dict[str, Any]] = []
|
||||
|
||||
while True:
|
||||
newline = self._buffer.find(b"\n")
|
||||
if newline < 0:
|
||||
if len(self._buffer) > MAX_FRAME_BYTES:
|
||||
self._fail("frame_too_large", f"frame exceeds {MAX_FRAME_BYTES} bytes")
|
||||
return frames
|
||||
if newline > MAX_FRAME_BYTES:
|
||||
self._fail("frame_too_large", f"frame exceeds {MAX_FRAME_BYTES} bytes")
|
||||
|
||||
raw = bytes(self._buffer[:newline])
|
||||
del self._buffer[: newline + 1]
|
||||
try:
|
||||
value = json.loads(raw.decode("utf-8"))
|
||||
except (UnicodeDecodeError, ValueError):
|
||||
self._fail("malformed_json", "frame contains malformed JSON")
|
||||
if not isinstance(value, dict):
|
||||
self._fail("malformed_json", "frame must contain a JSON object")
|
||||
frames.append(value)
|
||||
|
||||
def _fail(self, code: str, detail: str) -> None:
|
||||
self._failed = True
|
||||
self._buffer.clear()
|
||||
raise ProtocolError(code, detail, close_connection=True)
|
||||
|
||||
|
||||
def encode_frame(frame: Mapping[str, Any]) -> bytes:
|
||||
"""Encode one bounded service frame, applying the error-detail output cap."""
|
||||
|
||||
payload = dict(frame)
|
||||
if payload.get("type") == "result":
|
||||
transcript = payload.get("transcript")
|
||||
if not isinstance(transcript, str):
|
||||
raise ProtocolError("invalid_event", "result transcript must be a string")
|
||||
if len(transcript.encode("utf-8")) > MAX_TRANSCRIPT_BYTES:
|
||||
raise ProtocolError("transcript_too_large", "transcript exceeds output limit")
|
||||
if payload.get("type") == "error":
|
||||
detail = payload.get("detail", "")
|
||||
if not isinstance(detail, str):
|
||||
detail = str(detail)
|
||||
payload["detail"] = _bounded_utf8(detail, MAX_ERROR_DETAIL_BYTES)
|
||||
payload["requestId"] = _safe_id(payload.get("requestId"))
|
||||
generation_id = _safe_id(payload.get("generationId"))
|
||||
if generation_id:
|
||||
payload["generationId"] = generation_id
|
||||
else:
|
||||
payload.pop("generationId", None)
|
||||
|
||||
encoded = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
||||
if len(encoded) > MAX_FRAME_BYTES:
|
||||
raise ProtocolError("frame_too_large", "encoded frame exceeds output limit")
|
||||
return encoded + b"\n"
|
||||
|
||||
|
||||
def error_event(error: ProtocolError) -> dict[str, Any]:
|
||||
event: dict[str, Any] = {
|
||||
"version": PROTOCOL_VERSION,
|
||||
"type": "error",
|
||||
"requestId": _safe_id(error.request_id),
|
||||
"code": error.code,
|
||||
"detail": _bounded_utf8(error.detail, MAX_ERROR_DETAIL_BYTES),
|
||||
}
|
||||
generation_id = _safe_id(error.generation_id)
|
||||
if generation_id:
|
||||
event["generationId"] = generation_id
|
||||
return event
|
||||
@@ -0,0 +1,297 @@
|
||||
"""PipeWire recorder adapter with bounded in-memory output."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from typing import BinaryIO, Callable, Protocol
|
||||
|
||||
from .audio import MAX_PCM_BYTES
|
||||
|
||||
PW_RECORD_COMMAND = [
|
||||
"pw-record",
|
||||
"--raw",
|
||||
"--format",
|
||||
"s16",
|
||||
"--rate",
|
||||
"16000",
|
||||
"--channels",
|
||||
"1",
|
||||
"-",
|
||||
]
|
||||
MAX_STDERR_BYTES = 64 * 1024
|
||||
|
||||
|
||||
class Process(Protocol):
|
||||
pid: int
|
||||
stdout: BinaryIO
|
||||
stderr: BinaryIO
|
||||
returncode: int | None
|
||||
|
||||
def poll(self) -> int | None: ...
|
||||
def terminate(self) -> None: ...
|
||||
def kill(self) -> None: ...
|
||||
def wait(self, timeout: float | None = None) -> int: ...
|
||||
|
||||
|
||||
ProcessFactory = Callable[..., Process]
|
||||
ProcessGroupSignal = Callable[[int, int], None]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Capture:
|
||||
pcm: bytes
|
||||
returncode: int
|
||||
signal: int | None
|
||||
stderr_tail: bytes
|
||||
|
||||
|
||||
class RecorderError(RuntimeError):
|
||||
def __init__(
|
||||
self,
|
||||
code: str,
|
||||
detail: str,
|
||||
*,
|
||||
returncode: int | None = None,
|
||||
signal_number: int | None = None,
|
||||
) -> None:
|
||||
super().__init__(detail)
|
||||
self.code = code
|
||||
self.detail = detail
|
||||
self.returncode = returncode
|
||||
self.signal_number = signal_number
|
||||
|
||||
|
||||
class PwRecorder:
|
||||
"""Capture raw PCM while continuously draining both child pipes."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
process_factory: ProcessFactory = subprocess.Popen,
|
||||
max_bytes: int = MAX_PCM_BYTES,
|
||||
stderr_limit: int = MAX_STDERR_BYTES,
|
||||
stop_timeout: float = 3.0,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
process_group_signal: ProcessGroupSignal = os.killpg,
|
||||
) -> None:
|
||||
self._process_factory = process_factory
|
||||
self._max_bytes = max_bytes
|
||||
self._stderr_limit = stderr_limit
|
||||
self._stop_timeout = stop_timeout
|
||||
self._clock = clock
|
||||
self._process_group_signal = process_group_signal
|
||||
self._process: Process | None = None
|
||||
self._stdout_thread: threading.Thread | None = None
|
||||
self._stderr_thread: threading.Thread | None = None
|
||||
self._pcm = bytearray()
|
||||
self._stderr = bytearray()
|
||||
self._lock = threading.Lock()
|
||||
self._intentional = False
|
||||
self._duration_limited = False
|
||||
self._duration_stop_at: float | None = None
|
||||
self._reported = False
|
||||
|
||||
@property
|
||||
def buffered_bytes(self) -> int:
|
||||
with self._lock:
|
||||
return len(self._pcm)
|
||||
|
||||
def start(self) -> None:
|
||||
if self._process is not None:
|
||||
raise RecorderError("invalid_state", "recorder is already started")
|
||||
try:
|
||||
process = self._process_factory(
|
||||
list(PW_RECORD_COMMAND),
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
start_new_session=True,
|
||||
)
|
||||
except OSError as error:
|
||||
raise RecorderError("recorder_spawn", f"could not start pw-record: {error}") from error
|
||||
if process.stdout is None or process.stderr is None:
|
||||
raise RecorderError("recorder_spawn", "pw-record pipes were not created")
|
||||
self._process = process
|
||||
self._stdout_thread = threading.Thread(target=self._drain_stdout, daemon=True)
|
||||
self._stderr_thread = threading.Thread(target=self._drain_stderr, daemon=True)
|
||||
self._stdout_thread.start()
|
||||
self._stderr_thread.start()
|
||||
|
||||
def stop(self) -> Capture:
|
||||
return self._finish(discard=False)
|
||||
|
||||
def cancel(self) -> None:
|
||||
try:
|
||||
self._finish(discard=True)
|
||||
finally:
|
||||
with self._lock:
|
||||
self._pcm.clear()
|
||||
|
||||
def poll_unexpected(self) -> RecorderError | None:
|
||||
process = self._require_process()
|
||||
returncode = process.poll()
|
||||
if returncode is None:
|
||||
with self._lock:
|
||||
duration_stop_at = self._duration_stop_at
|
||||
if (
|
||||
not self._intentional
|
||||
and not self._reported
|
||||
and duration_stop_at is not None
|
||||
and self._clock() - duration_stop_at >= self._stop_timeout
|
||||
):
|
||||
self._reported = True
|
||||
self._signal_group(process, signal.SIGKILL)
|
||||
process.wait(timeout=self._stop_timeout)
|
||||
self._join_drainers()
|
||||
with self._lock:
|
||||
self._pcm.clear()
|
||||
return RecorderError(
|
||||
"duration_limit",
|
||||
"recording exceeded the five-minute duration limit; pw-record was killed after ignoring SIGTERM",
|
||||
returncode=process.returncode,
|
||||
signal_number=signal.SIGKILL,
|
||||
)
|
||||
return None
|
||||
if self._intentional or self._reported:
|
||||
return None
|
||||
self._reported = True
|
||||
self._join_drainers()
|
||||
with self._lock:
|
||||
self._pcm.clear()
|
||||
if self._duration_limited:
|
||||
return RecorderError(
|
||||
"duration_limit",
|
||||
"recording exceeded the five-minute duration limit",
|
||||
returncode=returncode,
|
||||
signal_number=-returncode if returncode < 0 else None,
|
||||
)
|
||||
return self._exit_error(returncode, "pw-record exited unexpectedly")
|
||||
|
||||
def _finish(self, *, discard: bool) -> Capture:
|
||||
process = self._require_process()
|
||||
self._intentional = True
|
||||
terminated_running_process = process.poll() is None
|
||||
if terminated_running_process:
|
||||
self._signal_group(process, signal.SIGTERM)
|
||||
try:
|
||||
returncode = process.wait(timeout=self._stop_timeout)
|
||||
except subprocess.TimeoutExpired as error:
|
||||
self._signal_group(process, signal.SIGKILL)
|
||||
process.wait(timeout=self._stop_timeout)
|
||||
self._join_drainers()
|
||||
with self._lock:
|
||||
self._pcm.clear()
|
||||
raise RecorderError(
|
||||
"recorder_timeout",
|
||||
"pw-record did not stop after SIGTERM and was killed",
|
||||
returncode=process.returncode,
|
||||
signal_number=signal.SIGKILL,
|
||||
) from error
|
||||
|
||||
self._join_drainers()
|
||||
stderr_tail = self._stderr_bytes()
|
||||
if self._duration_limited:
|
||||
with self._lock:
|
||||
self._pcm.clear()
|
||||
raise RecorderError(
|
||||
"duration_limit",
|
||||
"recording exceeded the five-minute duration limit",
|
||||
returncode=returncode,
|
||||
signal_number=-returncode if returncode < 0 else None,
|
||||
)
|
||||
clean_pipewire_sigterm = (
|
||||
terminated_running_process and returncode == 1 and not stderr_tail
|
||||
)
|
||||
if returncode not in (0, -signal.SIGTERM) and not clean_pipewire_sigterm:
|
||||
with self._lock:
|
||||
self._pcm.clear()
|
||||
raise self._exit_error(returncode, "pw-record failed while stopping")
|
||||
with self._lock:
|
||||
pcm = b"" if discard else bytes(self._pcm)
|
||||
self._pcm.clear()
|
||||
return Capture(
|
||||
pcm=pcm,
|
||||
returncode=returncode,
|
||||
signal=-returncode if returncode < 0 else None,
|
||||
stderr_tail=stderr_tail,
|
||||
)
|
||||
|
||||
def _drain_stdout(self) -> None:
|
||||
process = self._require_process()
|
||||
while True:
|
||||
chunk = process.stdout.read(64 * 1024)
|
||||
if not chunk:
|
||||
return
|
||||
with self._lock:
|
||||
if self._duration_limited:
|
||||
continue
|
||||
if len(self._pcm) + len(chunk) > self._max_bytes:
|
||||
self._duration_limited = True
|
||||
self._duration_stop_at = self._clock()
|
||||
self._pcm.clear()
|
||||
else:
|
||||
self._pcm.extend(chunk)
|
||||
continue
|
||||
if process.poll() is None:
|
||||
self._signal_group(process, signal.SIGTERM)
|
||||
|
||||
def _drain_stderr(self) -> None:
|
||||
process = self._require_process()
|
||||
while True:
|
||||
chunk = process.stderr.read(16 * 1024)
|
||||
if not chunk:
|
||||
return
|
||||
with self._lock:
|
||||
self._stderr.extend(chunk)
|
||||
overflow = len(self._stderr) - self._stderr_limit
|
||||
if overflow > 0:
|
||||
del self._stderr[:overflow]
|
||||
|
||||
def _join_drainers(self) -> None:
|
||||
for thread in (self._stdout_thread, self._stderr_thread):
|
||||
if thread is None:
|
||||
continue
|
||||
thread.join(timeout=self._stop_timeout)
|
||||
if thread.is_alive():
|
||||
raise RecorderError(
|
||||
"recorder_timeout",
|
||||
"pw-record pipe did not close after the process group stopped",
|
||||
)
|
||||
|
||||
def _signal_group(self, process: Process, signal_number: int) -> None:
|
||||
try:
|
||||
self._process_group_signal(process.pid, signal_number)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
|
||||
def _stderr_bytes(self) -> bytes:
|
||||
with self._lock:
|
||||
return bytes(self._stderr)
|
||||
|
||||
def _exit_error(self, returncode: int, prefix: str) -> RecorderError:
|
||||
if returncode < 0:
|
||||
classification = f"signal {-returncode}"
|
||||
signal_number = -returncode
|
||||
else:
|
||||
classification = f"code {returncode}"
|
||||
signal_number = None
|
||||
stderr = self._stderr_bytes().decode("utf-8", errors="replace").strip()
|
||||
detail = f"{prefix} with {classification}"
|
||||
if stderr:
|
||||
detail = f"{detail}: {stderr}"
|
||||
return RecorderError(
|
||||
"recorder_exit",
|
||||
detail,
|
||||
returncode=returncode,
|
||||
signal_number=signal_number,
|
||||
)
|
||||
|
||||
def _require_process(self) -> Process:
|
||||
if self._process is None:
|
||||
raise RecorderError("invalid_state", "recorder has not started")
|
||||
return self._process
|
||||
@@ -0,0 +1,570 @@
|
||||
"""Dictate service state machine and Unix-domain socket loop."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from concurrent.futures import Executor, Future, ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
import os
|
||||
from queue import Empty, SimpleQueue
|
||||
import selectors
|
||||
import socket
|
||||
import time
|
||||
from typing import Callable, Protocol
|
||||
|
||||
from .audio import AudioError, pcm_s16le_to_floats
|
||||
from .model import NemoParakeetModel, TranscriptionModel
|
||||
from .protocol import (
|
||||
MAX_FRAME_BYTES,
|
||||
MAX_TRANSCRIPT_BYTES,
|
||||
PROTOCOL_VERSION,
|
||||
FrameDecoder,
|
||||
ProtocolError,
|
||||
Request,
|
||||
encode_frame,
|
||||
error_event,
|
||||
parse_request,
|
||||
)
|
||||
from .recorder import Capture, PwRecorder, RecorderError
|
||||
|
||||
DEFAULT_IDLE_SECONDS = 15 * 60
|
||||
MAX_CLIENT_CONNECTIONS = 16
|
||||
MAX_CONNECTION_OUTPUT_BYTES = MAX_FRAME_BYTES
|
||||
|
||||
|
||||
class Recorder(Protocol):
|
||||
def start(self) -> None: ...
|
||||
def stop(self) -> Capture: ...
|
||||
def cancel(self) -> None: ...
|
||||
def poll_unexpected(self) -> RecorderError | None: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Outbound:
|
||||
client_id: str
|
||||
frame: dict[str, object]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Generation:
|
||||
owner: str
|
||||
generation_id: str
|
||||
request_id: str
|
||||
state: str
|
||||
recorder: Recorder | None = None
|
||||
tombstoned: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Completion:
|
||||
kind: str
|
||||
generation_id: str
|
||||
future: Future[object]
|
||||
|
||||
|
||||
class DictateServiceCore:
|
||||
"""Single-generation protocol state machine with injected host boundaries."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
model: TranscriptionModel,
|
||||
recorder_factory: Callable[[], Recorder] = PwRecorder,
|
||||
executor: Executor | None = None,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
idle_seconds: float = DEFAULT_IDLE_SECONDS,
|
||||
) -> None:
|
||||
self._model = model
|
||||
self._recorder_factory = recorder_factory
|
||||
self._executor = executor or ThreadPoolExecutor(max_workers=1, thread_name_prefix="dictate")
|
||||
self._owns_executor = executor is None
|
||||
self._clock = clock
|
||||
self._idle_seconds = idle_seconds
|
||||
self._last_activity = clock()
|
||||
self._clients: dict[str, bool] = {}
|
||||
self._active: _Generation | None = None
|
||||
self._model_loaded = bool(getattr(model, "loaded", False))
|
||||
self._worker_kind: str | None = None
|
||||
self._completions: SimpleQueue[_Completion] = SimpleQueue()
|
||||
self._events: list[Outbound] = []
|
||||
|
||||
@property
|
||||
def active_generation_id(self) -> str | None:
|
||||
return self._active.generation_id if self._active else None
|
||||
|
||||
def accepted(self, client_id: str) -> None:
|
||||
"""Track a socket without treating acceptance as protocol activity."""
|
||||
self._clients.setdefault(client_id, False)
|
||||
|
||||
def handle_request(self, client_id: str, request: Request) -> None:
|
||||
"""Handle one already parsed request and queue its observable events."""
|
||||
self._last_activity = self._clock()
|
||||
self._clients.setdefault(client_id, False)
|
||||
|
||||
if request.type == "hello":
|
||||
self._clients[client_id] = True
|
||||
self._emit(
|
||||
client_id,
|
||||
{
|
||||
"version": PROTOCOL_VERSION,
|
||||
"type": "pong",
|
||||
"requestId": request.request_id,
|
||||
"service": "pi-dictate",
|
||||
"ready": self._model_loaded,
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
if not self._clients[client_id]:
|
||||
self._request_error(client_id, request, "hello_required", "hello is required first")
|
||||
return
|
||||
|
||||
if request.type == "ping":
|
||||
self._emit(
|
||||
client_id,
|
||||
{
|
||||
"version": PROTOCOL_VERSION,
|
||||
"type": "pong",
|
||||
"requestId": request.request_id,
|
||||
"service": "pi-dictate",
|
||||
"ready": self._model_loaded,
|
||||
"busy": self._active is not None,
|
||||
},
|
||||
)
|
||||
elif request.type == "start":
|
||||
self._start(client_id, request)
|
||||
elif request.type == "stop":
|
||||
self._stop(client_id, request)
|
||||
elif request.type == "cancel":
|
||||
self._cancel(client_id, request)
|
||||
|
||||
def poll(self) -> None:
|
||||
"""Apply completed workers and recorder exits without blocking the socket loop."""
|
||||
while True:
|
||||
try:
|
||||
completion = self._completions.get_nowait()
|
||||
except Empty:
|
||||
break
|
||||
self._apply_completion(completion)
|
||||
|
||||
active = self._active
|
||||
if active and active.state == "listening" and active.recorder and self._worker_kind is None:
|
||||
try:
|
||||
failure = active.recorder.poll_unexpected()
|
||||
except RecorderError as error:
|
||||
failure = error
|
||||
if failure is not None:
|
||||
self._async_error(active, failure.code, failure.detail)
|
||||
self._finish_generation(active, notify_idle=True)
|
||||
|
||||
def drain_events(self) -> list[Outbound]:
|
||||
events, self._events = self._events, []
|
||||
return events
|
||||
|
||||
def disconnect(self, client_id: str) -> None:
|
||||
self._clients.pop(client_id, None)
|
||||
active = self._active
|
||||
if active is None or active.owner != client_id:
|
||||
return
|
||||
active.tombstoned = True
|
||||
if active.state == "listening" and active.recorder is not None:
|
||||
active.state = "stopping"
|
||||
self._submit("cancel", active, active.recorder.cancel)
|
||||
|
||||
def should_exit(self) -> bool:
|
||||
if self._active is not None or self._worker_kind is not None:
|
||||
return False
|
||||
return self._clock() - self._last_activity >= self._idle_seconds
|
||||
|
||||
def close(self) -> None:
|
||||
active = self._active
|
||||
if active and active.recorder and active.state == "listening":
|
||||
try:
|
||||
active.recorder.cancel()
|
||||
except RecorderError:
|
||||
pass
|
||||
self._active = None
|
||||
if self._owns_executor:
|
||||
self._executor.shutdown(wait=True, cancel_futures=False)
|
||||
|
||||
def _start(self, client_id: str, request: Request) -> None:
|
||||
assert request.generation_id is not None
|
||||
if self._active is not None or self._worker_kind is not None:
|
||||
self._request_error(client_id, request, "busy", "another generation is active")
|
||||
return
|
||||
active = _Generation(client_id, request.generation_id, request.request_id, "loading")
|
||||
self._active = active
|
||||
if not self._model_loaded:
|
||||
self._state(active, "loading")
|
||||
self._submit("load", active, self._model.load)
|
||||
return
|
||||
self._begin_recording(active)
|
||||
|
||||
def _begin_recording(self, active: _Generation) -> None:
|
||||
if active.tombstoned:
|
||||
self._finish_generation(active, notify_idle=False)
|
||||
return
|
||||
try:
|
||||
recorder = self._recorder_factory()
|
||||
recorder.start()
|
||||
except RecorderError as error:
|
||||
self._async_error(active, error.code, error.detail)
|
||||
self._finish_generation(active, notify_idle=True)
|
||||
return
|
||||
except OSError as error:
|
||||
self._async_error(active, "recorder_spawn", str(error))
|
||||
self._finish_generation(active, notify_idle=True)
|
||||
return
|
||||
active.recorder = recorder
|
||||
self._state(active, "listening")
|
||||
|
||||
def _stop(self, client_id: str, request: Request) -> None:
|
||||
active = self._owned_generation(client_id, request)
|
||||
if active is None:
|
||||
return
|
||||
if active.state != "listening" or active.recorder is None:
|
||||
self._request_error(client_id, request, "invalid_state", f"cannot stop while {active.state}")
|
||||
return
|
||||
active.request_id = request.request_id
|
||||
try:
|
||||
failure = active.recorder.poll_unexpected()
|
||||
except RecorderError as error:
|
||||
failure = error
|
||||
if failure is not None:
|
||||
self._async_error(active, failure.code, failure.detail)
|
||||
self._finish_generation(active, notify_idle=True)
|
||||
return
|
||||
active.state = "stopping"
|
||||
self._state(active, "stopping", assign=False)
|
||||
self._submit("stop", active, self._stop_and_convert, active.recorder)
|
||||
|
||||
def _cancel(self, client_id: str, request: Request) -> None:
|
||||
active = self._owned_generation(client_id, request)
|
||||
if active is None:
|
||||
return
|
||||
if active.tombstoned:
|
||||
self._request_error(client_id, request, "invalid_state", "generation is already cancelled")
|
||||
return
|
||||
active.tombstoned = True
|
||||
active.request_id = request.request_id
|
||||
self._emit(
|
||||
client_id,
|
||||
{
|
||||
"version": PROTOCOL_VERSION,
|
||||
"type": "cancelled",
|
||||
"requestId": active.request_id,
|
||||
"generationId": active.generation_id,
|
||||
},
|
||||
)
|
||||
if active.state == "listening" and active.recorder is not None:
|
||||
active.state = "stopping"
|
||||
self._submit("cancel", active, active.recorder.cancel)
|
||||
|
||||
def _owned_generation(self, client_id: str, request: Request) -> _Generation | None:
|
||||
assert request.generation_id is not None
|
||||
active = self._active
|
||||
if active is None:
|
||||
self._request_error(client_id, request, "invalid_state", "no generation is active")
|
||||
return None
|
||||
if active.owner != client_id:
|
||||
self._request_error(client_id, request, "not_owner", "generation belongs to another client")
|
||||
return None
|
||||
if active.generation_id != request.generation_id:
|
||||
self._request_error(
|
||||
client_id,
|
||||
request,
|
||||
"generation_mismatch",
|
||||
"generation ID does not match the active generation",
|
||||
)
|
||||
return None
|
||||
return active
|
||||
|
||||
@staticmethod
|
||||
def _stop_and_convert(recorder: Recorder) -> Sequence[float]:
|
||||
capture = recorder.stop()
|
||||
return pcm_s16le_to_floats(capture.pcm)
|
||||
|
||||
def _submit(
|
||||
self,
|
||||
kind: str,
|
||||
active: _Generation,
|
||||
function: Callable[..., object],
|
||||
*args: object,
|
||||
) -> None:
|
||||
self._worker_kind = kind
|
||||
future = self._executor.submit(function, *args)
|
||||
future.add_done_callback(
|
||||
lambda completed, kind=kind, generation_id=active.generation_id: self._completions.put(
|
||||
_Completion(kind, generation_id, completed)
|
||||
)
|
||||
)
|
||||
|
||||
def _apply_completion(self, completion: _Completion) -> None:
|
||||
if self._worker_kind == completion.kind:
|
||||
self._worker_kind = None
|
||||
active = self._active
|
||||
if active is None or active.generation_id != completion.generation_id:
|
||||
return
|
||||
try:
|
||||
result = completion.future.result()
|
||||
except RecorderError as error:
|
||||
if not active.tombstoned:
|
||||
self._async_error(active, error.code, error.detail)
|
||||
self._finish_generation(active, notify_idle=not active.tombstoned)
|
||||
return
|
||||
except AudioError as error:
|
||||
if not active.tombstoned:
|
||||
self._async_error(active, "invalid_audio", str(error))
|
||||
self._finish_generation(active, notify_idle=not active.tombstoned)
|
||||
return
|
||||
except BaseException as error:
|
||||
if not active.tombstoned:
|
||||
code = "model_failure" if completion.kind in ("load", "infer") else "recorder_failure"
|
||||
self._async_error(active, code, str(error))
|
||||
self._finish_generation(active, notify_idle=not active.tombstoned)
|
||||
return
|
||||
|
||||
if completion.kind == "load":
|
||||
self._model_loaded = True
|
||||
if active.tombstoned:
|
||||
self._finish_generation(active, notify_idle=False)
|
||||
else:
|
||||
self._begin_recording(active)
|
||||
elif completion.kind == "stop":
|
||||
if active.tombstoned:
|
||||
self._finish_generation(active, notify_idle=False)
|
||||
return
|
||||
samples = result
|
||||
if not isinstance(samples, Sequence) or isinstance(samples, (str, bytes, bytearray)):
|
||||
self._async_error(active, "invalid_audio", "audio conversion returned invalid samples")
|
||||
self._finish_generation(active, notify_idle=True)
|
||||
return
|
||||
self._state(active, "transcribing")
|
||||
self._submit("infer", active, self._model.transcribe, samples)
|
||||
elif completion.kind == "infer":
|
||||
if active.tombstoned:
|
||||
self._finish_generation(active, notify_idle=False)
|
||||
return
|
||||
if not isinstance(result, str):
|
||||
self._async_error(active, "model_failure", "model returned a non-string transcript")
|
||||
elif len(result.encode("utf-8")) > MAX_TRANSCRIPT_BYTES:
|
||||
self._async_error(active, "transcript_too_large", "transcript exceeds output limit")
|
||||
else:
|
||||
self._emit(
|
||||
active.owner,
|
||||
{
|
||||
"version": PROTOCOL_VERSION,
|
||||
"type": "result",
|
||||
"requestId": active.request_id,
|
||||
"generationId": active.generation_id,
|
||||
"transcript": result,
|
||||
},
|
||||
)
|
||||
self._finish_generation(active, notify_idle=True)
|
||||
elif completion.kind == "cancel":
|
||||
self._finish_generation(active, notify_idle=False)
|
||||
|
||||
def _finish_generation(self, active: _Generation, *, notify_idle: bool) -> None:
|
||||
if self._active is not active:
|
||||
return
|
||||
self._active = None
|
||||
self._last_activity = self._clock()
|
||||
if notify_idle and active.owner in self._clients:
|
||||
self._state(active, "idle", assign=False)
|
||||
|
||||
def _state(self, active: _Generation, state: str, *, assign: bool = True) -> None:
|
||||
if assign:
|
||||
active.state = state
|
||||
self._emit(
|
||||
active.owner,
|
||||
{
|
||||
"version": PROTOCOL_VERSION,
|
||||
"type": "state",
|
||||
"requestId": active.request_id,
|
||||
"generationId": active.generation_id,
|
||||
"state": state,
|
||||
},
|
||||
)
|
||||
|
||||
def _request_error(self, client_id: str, request: Request, code: str, detail: str) -> None:
|
||||
error = ProtocolError(
|
||||
code,
|
||||
detail,
|
||||
request_id=request.request_id,
|
||||
generation_id=request.generation_id,
|
||||
)
|
||||
self._emit(client_id, error_event(error))
|
||||
|
||||
def _async_error(self, active: _Generation, code: str, detail: str) -> None:
|
||||
error = ProtocolError(
|
||||
code,
|
||||
detail,
|
||||
request_id=active.request_id,
|
||||
generation_id=active.generation_id,
|
||||
)
|
||||
self._emit(active.owner, error_event(error))
|
||||
|
||||
def _emit(self, client_id: str, frame: dict[str, object]) -> None:
|
||||
if client_id in self._clients:
|
||||
self._events.append(Outbound(client_id, frame))
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Connection:
|
||||
socket: socket.socket
|
||||
decoder: FrameDecoder
|
||||
output: bytearray
|
||||
close_after_write: bool = False
|
||||
|
||||
|
||||
def inherited_listener(environ: dict[str, str] | None = None) -> socket.socket:
|
||||
"""Return systemd's single inherited listener at fd 3."""
|
||||
environment = os.environ if environ is None else environ
|
||||
try:
|
||||
listen_pid = int(environment.get("LISTEN_PID", "0"))
|
||||
listen_fds = int(environment.get("LISTEN_FDS", "0"))
|
||||
except ValueError as error:
|
||||
raise RuntimeError("invalid systemd socket activation environment") from error
|
||||
if listen_pid != os.getpid() or listen_fds != 1:
|
||||
raise RuntimeError("expected exactly one systemd socket-activated listener")
|
||||
listener = socket.socket(fileno=3)
|
||||
listener.setblocking(False)
|
||||
return listener
|
||||
|
||||
|
||||
def serve(listener: socket.socket, core: DictateServiceCore) -> None:
|
||||
"""Serve protocol v1 on an explicit listener until the core becomes idle."""
|
||||
selector = selectors.DefaultSelector()
|
||||
connections: dict[str, _Connection] = {}
|
||||
sequence = 0
|
||||
listener.setblocking(False)
|
||||
selector.register(listener, selectors.EVENT_READ, None)
|
||||
try:
|
||||
while not core.should_exit():
|
||||
core.poll()
|
||||
for outbound in core.drain_events():
|
||||
connection = connections.get(outbound.client_id)
|
||||
if connection is None:
|
||||
continue
|
||||
try:
|
||||
encoded = encode_frame(outbound.frame)
|
||||
except ProtocolError:
|
||||
_close_connection(selector, connections, core, outbound.client_id)
|
||||
continue
|
||||
if len(connection.output) + len(encoded) > MAX_CONNECTION_OUTPUT_BYTES:
|
||||
_close_connection(selector, connections, core, outbound.client_id)
|
||||
continue
|
||||
connection.output.extend(encoded)
|
||||
selector.modify(
|
||||
connection.socket,
|
||||
selectors.EVENT_READ | selectors.EVENT_WRITE,
|
||||
outbound.client_id,
|
||||
)
|
||||
for key, mask in selector.select(timeout=0.1):
|
||||
if key.data is None:
|
||||
try:
|
||||
client, _address = listener.accept()
|
||||
except BlockingIOError:
|
||||
continue
|
||||
if len(connections) >= MAX_CLIENT_CONNECTIONS:
|
||||
client.close()
|
||||
continue
|
||||
client.setblocking(False)
|
||||
sequence += 1
|
||||
client_id = f"client-{sequence}"
|
||||
connections[client_id] = _Connection(client, FrameDecoder(), bytearray())
|
||||
core.accepted(client_id)
|
||||
selector.register(client, selectors.EVENT_READ, client_id)
|
||||
continue
|
||||
client_id = key.data
|
||||
connection = connections.get(client_id)
|
||||
if connection is None:
|
||||
continue
|
||||
if mask & selectors.EVENT_READ:
|
||||
try:
|
||||
data = connection.socket.recv(64 * 1024)
|
||||
except BlockingIOError:
|
||||
data = None
|
||||
except (BrokenPipeError, ConnectionResetError):
|
||||
_close_connection(selector, connections, core, client_id)
|
||||
continue
|
||||
if data == b"":
|
||||
_close_connection(selector, connections, core, client_id)
|
||||
continue
|
||||
if data:
|
||||
try:
|
||||
for frame in connection.decoder.feed(data):
|
||||
core.handle_request(client_id, parse_request(frame))
|
||||
except ProtocolError as error:
|
||||
connection.output.extend(encode_frame(error_event(error)))
|
||||
if error.close_connection:
|
||||
connection.close_after_write = True
|
||||
if mask & selectors.EVENT_WRITE and connection.output:
|
||||
try:
|
||||
sent = connection.socket.send(connection.output)
|
||||
except BlockingIOError:
|
||||
sent = 0
|
||||
except (BrokenPipeError, ConnectionResetError):
|
||||
_close_connection(selector, connections, core, client_id)
|
||||
continue
|
||||
del connection.output[:sent]
|
||||
if connection.close_after_write and not connection.output:
|
||||
_close_connection(selector, connections, core, client_id)
|
||||
elif client_id in connections:
|
||||
events = selectors.EVENT_READ
|
||||
if connection.output:
|
||||
events |= selectors.EVENT_WRITE
|
||||
selector.modify(connection.socket, events, client_id)
|
||||
finally:
|
||||
for client_id in list(connections):
|
||||
_close_connection(selector, connections, core, client_id)
|
||||
selector.unregister(listener)
|
||||
listener.close()
|
||||
selector.close()
|
||||
core.close()
|
||||
|
||||
|
||||
def _close_connection(
|
||||
selector: selectors.BaseSelector,
|
||||
connections: dict[str, _Connection],
|
||||
core: DictateServiceCore,
|
||||
client_id: str,
|
||||
) -> None:
|
||||
connection = connections.pop(client_id, None)
|
||||
if connection is None:
|
||||
return
|
||||
core.disconnect(client_id)
|
||||
try:
|
||||
selector.unregister(connection.socket)
|
||||
except KeyError:
|
||||
pass
|
||||
connection.socket.close()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
model_path = os.environ.get("PI_DICTATE_MODEL_PATH")
|
||||
if not model_path:
|
||||
raise SystemExit("PI_DICTATE_MODEL_PATH must name an explicit local .nemo file")
|
||||
terms_path = os.environ.get("PI_DICTATE_TERMS_PATH")
|
||||
if not terms_path:
|
||||
raise SystemExit("PI_DICTATE_TERMS_PATH must name the host coding dictionary")
|
||||
idle_value = os.environ.get("PI_DICTATE_IDLE_SECONDS", str(DEFAULT_IDLE_SECONDS))
|
||||
boost_value = os.environ.get("PI_DICTATE_BOOST_ALPHA", "1.0")
|
||||
try:
|
||||
idle_seconds = float(idle_value)
|
||||
boost_alpha = float(boost_value)
|
||||
except ValueError as error:
|
||||
raise SystemExit("PI_DICTATE_IDLE_SECONDS and PI_DICTATE_BOOST_ALPHA must be numeric") from error
|
||||
if idle_seconds <= 0:
|
||||
raise SystemExit("PI_DICTATE_IDLE_SECONDS must be positive")
|
||||
if boost_alpha <= 0:
|
||||
raise SystemExit("PI_DICTATE_BOOST_ALPHA must be positive")
|
||||
|
||||
model = NemoParakeetModel(model_path, terms_path=terms_path, boost_alpha=boost_alpha)
|
||||
core = DictateServiceCore(model=model, idle_seconds=idle_seconds)
|
||||
listener = inherited_listener()
|
||||
serve(listener, core)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,107 @@
|
||||
from pathlib import Path
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
|
||||
SCRIPT = Path(__file__).resolve().parents[1] / "install.sh"
|
||||
|
||||
|
||||
class InstallerTests(unittest.TestCase):
|
||||
def run_installer(
|
||||
self,
|
||||
home: Path,
|
||||
model: Path,
|
||||
*arguments: str,
|
||||
install_root_override: str | None = None,
|
||||
) -> subprocess.CompletedProcess[str]:
|
||||
environment = {
|
||||
**os.environ,
|
||||
"HOME": str(home),
|
||||
"XDG_CONFIG_HOME": str(home / ".config"),
|
||||
"PI_DICTATE_INSTALL_ROOT": install_root_override or str(
|
||||
home / ".local/share/pi-dictate"
|
||||
),
|
||||
}
|
||||
return subprocess.run(
|
||||
[str(SCRIPT), *arguments, "--dry-run", "--model-source", str(model)],
|
||||
env=environment,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
|
||||
def test_dry_run_seeds_dictionary_and_installs_owned_unit_copies(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
home = Path(directory)
|
||||
model = home / "model.nemo"
|
||||
model.write_bytes(b"model")
|
||||
|
||||
result = self.run_installer(home, model)
|
||||
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
self.assertIn("coding-terms.txt", result.stdout)
|
||||
self.assertIn("pi-dictate.socket", result.stdout)
|
||||
self.assertIn("install -m 0644", result.stdout)
|
||||
self.assertNotIn("ln -s", result.stdout)
|
||||
self.assertIn(".nemo.install.", result.stdout)
|
||||
self.assertIn("mv -f", result.stdout)
|
||||
self.assertIn("uv sync --frozen --no-dev --no-editable", result.stdout)
|
||||
self.assertIn("Would install pi-dictate", result.stdout)
|
||||
self.assertFalse((home / ".local/share/pi-dictate").exists())
|
||||
|
||||
def test_ignores_dangerous_install_root_environment_override(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
home = Path(directory)
|
||||
model = home / "model.nemo"
|
||||
model.write_bytes(b"model")
|
||||
|
||||
result = self.run_installer(
|
||||
home,
|
||||
model,
|
||||
"uninstall",
|
||||
install_root_override="/",
|
||||
)
|
||||
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
self.assertIn(str(home / ".local/share/pi-dictate"), result.stdout)
|
||||
self.assertNotIn("rm -rf / ", result.stdout)
|
||||
|
||||
def test_uninstall_refuses_to_disable_an_unmanaged_same_named_unit(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
home = Path(directory)
|
||||
model = home / "model.nemo"
|
||||
model.write_bytes(b"model")
|
||||
unit = home / ".config/systemd/user/pi-dictate.socket"
|
||||
unit.parent.mkdir(parents=True)
|
||||
unit.write_text("[Socket]\nListenStream=/tmp/unrelated.sock\n", encoding="utf-8")
|
||||
|
||||
result = self.run_installer(home, model, "uninstall")
|
||||
|
||||
self.assertEqual(result.returncode, 1)
|
||||
self.assertIn("unmanaged systemd unit", result.stderr)
|
||||
self.assertNotIn("systemctl --user disable", result.stdout)
|
||||
self.assertEqual(
|
||||
unit.read_text(encoding="utf-8"),
|
||||
"[Socket]\nListenStream=/tmp/unrelated.sock\n",
|
||||
)
|
||||
|
||||
def test_reinstall_preserves_an_existing_host_dictionary(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
home = Path(directory)
|
||||
model = home / "model.nemo"
|
||||
model.write_bytes(b"model")
|
||||
dictionary = home / ".local/share/pi-dictate/config/coding-terms.txt"
|
||||
dictionary.parent.mkdir(parents=True)
|
||||
dictionary.write_text("my private term\n", encoding="utf-8")
|
||||
|
||||
result = self.run_installer(home, model)
|
||||
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
self.assertIn("Keeping existing coding dictionary", result.stdout)
|
||||
self.assertEqual(dictionary.read_text(encoding="utf-8"), "my private term\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,143 @@
|
||||
import hashlib
|
||||
import os
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
from types import ModuleType, SimpleNamespace
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from pi_dictate_service.model import NemoParakeetModel, load_coding_terms
|
||||
|
||||
|
||||
class FakeNemoModel:
|
||||
def __init__(self) -> None:
|
||||
self.eval_called = False
|
||||
self.to_device: str | None = None
|
||||
self.audio = None
|
||||
self.transcribe_options: dict[str, object] = {}
|
||||
self.cfg = SimpleNamespace(decoding={"strategy": "greedy_batch", "greedy": {"max_symbols": 10}})
|
||||
self.decoding_config: dict[str, object] | None = None
|
||||
|
||||
def eval(self) -> None:
|
||||
self.eval_called = True
|
||||
|
||||
def to(self, device: str) -> None:
|
||||
self.to_device = device
|
||||
|
||||
def change_decoding_strategy(self, config: dict[str, object], verbose: bool) -> None:
|
||||
self.decoding_config = config
|
||||
|
||||
def transcribe(self, audio: list[object], **options: object) -> list[object]:
|
||||
self.audio = audio
|
||||
self.transcribe_options = options
|
||||
return [type("Hypothesis", (), {"text": "hello locally"})()]
|
||||
|
||||
|
||||
class FakeASRModel:
|
||||
restored_path: str | None = None
|
||||
restored: FakeNemoModel | None = None
|
||||
|
||||
@classmethod
|
||||
def restore_from(cls, restore_path: str, map_location: str) -> FakeNemoModel:
|
||||
cls.restored_path = restore_path
|
||||
self = FakeNemoModel()
|
||||
cls.restored = self
|
||||
return self
|
||||
|
||||
|
||||
class FakeNumpy(ModuleType):
|
||||
float32 = "float32"
|
||||
|
||||
def asarray(self, samples: list[float], dtype: object) -> tuple[list[float], object]:
|
||||
return samples, dtype
|
||||
|
||||
|
||||
class NemoParakeetModelTests(unittest.TestCase):
|
||||
def test_load_is_lazy_offline_applies_coding_terms_and_transcribes_in_memory_audio(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
model_path = Path(directory) / "parakeet.nemo"
|
||||
model_path.write_bytes(b"model")
|
||||
terms_path = Path(directory) / "coding-terms.txt"
|
||||
terms_path.write_text("JSON\nnpm test\n", encoding="utf-8")
|
||||
modules = self.fake_modules()
|
||||
with patch.dict(sys.modules, modules, clear=False), patch.dict(os.environ, {}, clear=True):
|
||||
adapter = NemoParakeetModel(
|
||||
model_path,
|
||||
terms_path=terms_path,
|
||||
boost_alpha=0.5,
|
||||
expected_sha256=hashlib.sha256(b"model").hexdigest(),
|
||||
)
|
||||
self.assertIsNone(FakeASRModel.restored_path)
|
||||
|
||||
adapter.load()
|
||||
transcript = adapter.transcribe([0.0, 0.5])
|
||||
|
||||
self.assertEqual(os.environ["HF_HUB_OFFLINE"], "1")
|
||||
self.assertEqual(os.environ["TRANSFORMERS_OFFLINE"], "1")
|
||||
self.assertEqual(FakeASRModel.restored_path, str(model_path))
|
||||
self.assertTrue(FakeASRModel.restored.eval_called)
|
||||
self.assertEqual(FakeASRModel.restored.to_device, "cuda")
|
||||
self.assertEqual(transcript, "hello locally")
|
||||
self.assertEqual(FakeASRModel.restored.audio, [([0.0, 0.5], "float32")])
|
||||
self.assertEqual(
|
||||
FakeASRModel.restored.decoding_config["greedy"]["boosting_tree"]["key_phrases_list"],
|
||||
["JSON", "npm test"],
|
||||
)
|
||||
self.assertEqual(
|
||||
FakeASRModel.restored.decoding_config["greedy"]["boosting_tree_alpha"],
|
||||
0.5,
|
||||
)
|
||||
self.assertEqual(
|
||||
FakeASRModel.restored.transcribe_options,
|
||||
{"batch_size": 1, "num_workers": 0, "verbose": False},
|
||||
)
|
||||
|
||||
def test_loads_a_bounded_deduplicated_host_dictionary(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
terms_path = Path(directory) / "coding-terms.txt"
|
||||
terms_path.write_text("# local terms\n JSON \n\nnpm\nJSON\n", encoding="utf-8")
|
||||
self.assertEqual(load_coding_terms(terms_path), ["JSON", "npm"])
|
||||
|
||||
terms_path.write_text("x" * 129, encoding="utf-8")
|
||||
with self.assertRaisesRegex(ValueError, "128 UTF-8 bytes"):
|
||||
load_coding_terms(terms_path)
|
||||
|
||||
def test_refuses_a_model_that_no_longer_matches_the_pinned_hash(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
model_path = Path(directory) / "parakeet.nemo"
|
||||
model_path.write_bytes(b"replaced")
|
||||
adapter = NemoParakeetModel(model_path)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "model checksum mismatch"):
|
||||
adapter.load()
|
||||
|
||||
def test_requires_an_existing_local_nemo_file(self) -> None:
|
||||
with self.assertRaisesRegex(ValueError, "local .nemo"):
|
||||
NemoParakeetModel("/does/not/exist")
|
||||
|
||||
@staticmethod
|
||||
def fake_modules() -> dict[str, ModuleType]:
|
||||
nemo = ModuleType("nemo")
|
||||
collections = ModuleType("nemo.collections")
|
||||
asr = ModuleType("nemo.collections.asr")
|
||||
models = ModuleType("nemo.collections.asr.models")
|
||||
models.ASRModel = FakeASRModel
|
||||
asr.models = models
|
||||
collections.asr = asr
|
||||
nemo.collections = collections
|
||||
omegaconf = ModuleType("omegaconf")
|
||||
omegaconf.open_dict = lambda _config: nullcontext()
|
||||
return {
|
||||
"nemo": nemo,
|
||||
"nemo.collections": collections,
|
||||
"nemo.collections.asr": asr,
|
||||
"nemo.collections.asr.models": models,
|
||||
"numpy": FakeNumpy("numpy"),
|
||||
"omegaconf": omegaconf,
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,162 @@
|
||||
import json
|
||||
import unittest
|
||||
|
||||
from pi_dictate_service.audio import AudioError, pcm_s16le_to_floats
|
||||
from pi_dictate_service.protocol import (
|
||||
MAX_ERROR_DETAIL_BYTES,
|
||||
MAX_FRAME_BYTES,
|
||||
MAX_ID_BYTES,
|
||||
MAX_TRANSCRIPT_BYTES,
|
||||
FrameDecoder,
|
||||
ProtocolError,
|
||||
encode_frame,
|
||||
parse_request,
|
||||
)
|
||||
|
||||
|
||||
class FrameDecoderTests(unittest.TestCase):
|
||||
def test_decodes_fragmented_and_multiple_frames(self) -> None:
|
||||
decoder = FrameDecoder()
|
||||
|
||||
self.assertEqual(decoder.feed(b'{"version":1'), [])
|
||||
frames = decoder.feed(
|
||||
b',"type":"ping","requestId":"one"}\n'
|
||||
b'{"version":1,"type":"ping","requestId":"two"}\n'
|
||||
)
|
||||
|
||||
self.assertEqual([frame["requestId"] for frame in frames], ["one", "two"])
|
||||
|
||||
def test_rejects_malformed_json_and_oversized_frames_as_fatal(self) -> None:
|
||||
with self.assertRaisesRegex(ProtocolError, "malformed JSON") as malformed:
|
||||
FrameDecoder().feed(b'{nope}\n')
|
||||
self.assertTrue(malformed.exception.close_connection)
|
||||
self.assertEqual(malformed.exception.code, "malformed_json")
|
||||
|
||||
with self.assertRaises(ProtocolError) as oversized:
|
||||
FrameDecoder().feed(b"x" * (MAX_FRAME_BYTES + 1))
|
||||
self.assertTrue(oversized.exception.close_connection)
|
||||
self.assertEqual(oversized.exception.code, "frame_too_large")
|
||||
|
||||
|
||||
class RequestValidationTests(unittest.TestCase):
|
||||
def test_requires_exact_fields_bounded_ids_and_supported_version(self) -> None:
|
||||
request = parse_request(
|
||||
{"version": 1, "type": "start", "requestId": "r1", "generationId": "g1"}
|
||||
)
|
||||
self.assertEqual(request.generation_id, "g1")
|
||||
|
||||
with self.assertRaises(ProtocolError) as unknown:
|
||||
parse_request(
|
||||
{
|
||||
"version": 1,
|
||||
"type": "start",
|
||||
"requestId": "r1",
|
||||
"generationId": "g1",
|
||||
"path": "/tmp/audio",
|
||||
}
|
||||
)
|
||||
self.assertEqual(unknown.exception.code, "invalid_request")
|
||||
self.assertFalse(unknown.exception.close_connection)
|
||||
|
||||
with self.assertRaises(ProtocolError) as long_id:
|
||||
parse_request(
|
||||
{
|
||||
"version": 1,
|
||||
"type": "ping",
|
||||
"requestId": "é" * ((MAX_ID_BYTES // 2) + 1),
|
||||
}
|
||||
)
|
||||
self.assertEqual(long_id.exception.code, "invalid_request")
|
||||
|
||||
with self.assertRaises(ProtocolError) as unsupported:
|
||||
parse_request({"version": 2, "type": "hello", "requestId": "r2"})
|
||||
self.assertEqual(unsupported.exception.code, "unsupported_version")
|
||||
self.assertTrue(unsupported.exception.close_connection)
|
||||
|
||||
def test_rejects_unhashable_types_invalid_unicode_and_huge_integers(self) -> None:
|
||||
for request_type in ([], {}):
|
||||
with self.subTest(request_type=request_type):
|
||||
with self.assertRaises(ProtocolError) as failed:
|
||||
parse_request(
|
||||
{"version": 1, "type": request_type, "requestId": "request"}
|
||||
)
|
||||
self.assertEqual(failed.exception.code, "invalid_request")
|
||||
|
||||
with self.assertRaises(ProtocolError) as surrogate:
|
||||
parse_request({"version": 1, "type": "ping", "requestId": "\ud800"})
|
||||
self.assertEqual(surrogate.exception.request_id, "")
|
||||
|
||||
huge_integer = b'{"version":' + (b"9" * 5000) + b',"type":"hello","requestId":"r"}\n'
|
||||
with self.assertRaises(ProtocolError) as huge:
|
||||
FrameDecoder().feed(huge_integer)
|
||||
self.assertEqual(huge.exception.code, "malformed_json")
|
||||
self.assertTrue(huge.exception.close_connection)
|
||||
|
||||
def test_error_echoes_only_valid_bounded_ids(self) -> None:
|
||||
oversized = "x" * (MAX_ID_BYTES + 1)
|
||||
with self.assertRaises(ProtocolError) as unsupported:
|
||||
parse_request(
|
||||
{"version": 2, "type": "hello", "requestId": oversized}
|
||||
)
|
||||
self.assertEqual(unsupported.exception.request_id, "")
|
||||
|
||||
encoded = encode_frame(
|
||||
{
|
||||
"version": 1,
|
||||
"type": "error",
|
||||
"requestId": oversized,
|
||||
"generationId": "\ud800",
|
||||
"code": "invalid_request",
|
||||
"detail": "bad request",
|
||||
}
|
||||
)
|
||||
payload = json.loads(encoded)
|
||||
self.assertEqual(payload["requestId"], "")
|
||||
self.assertNotIn("generationId", payload)
|
||||
|
||||
def test_encoded_transcripts_and_error_details_are_bounded(self) -> None:
|
||||
with self.assertRaises(ProtocolError) as transcript:
|
||||
encode_frame(
|
||||
{
|
||||
"version": 1,
|
||||
"type": "result",
|
||||
"generationId": "g1",
|
||||
"transcript": "é" * ((MAX_TRANSCRIPT_BYTES // 2) + 1),
|
||||
}
|
||||
)
|
||||
self.assertEqual(transcript.exception.code, "transcript_too_large")
|
||||
|
||||
encoded = encode_frame(
|
||||
{
|
||||
"version": 1,
|
||||
"type": "error",
|
||||
"requestId": "r1",
|
||||
"code": "failed",
|
||||
"detail": "é" * MAX_ERROR_DETAIL_BYTES,
|
||||
}
|
||||
)
|
||||
payload = json.loads(encoded)
|
||||
self.assertLessEqual(len(payload["detail"].encode("utf-8")), MAX_ERROR_DETAIL_BYTES)
|
||||
self.assertLessEqual(len(encoded) - 1, MAX_FRAME_BYTES)
|
||||
|
||||
|
||||
class AudioConversionTests(unittest.TestCase):
|
||||
def test_converts_little_endian_signed_16_bit_pcm(self) -> None:
|
||||
pcm = b"\x00\x80\x00\xc0\x00\x00\x00\x40\xff\x7f"
|
||||
|
||||
samples = pcm_s16le_to_floats(pcm)
|
||||
self.assertEqual(samples.itemsize, 4)
|
||||
self.assertEqual(
|
||||
list(samples),
|
||||
[-1.0, -0.5, 0.0, 0.5, 32767 / 32768.0],
|
||||
)
|
||||
|
||||
def test_rejects_odd_or_overlong_pcm(self) -> None:
|
||||
with self.assertRaisesRegex(AudioError, "even"):
|
||||
pcm_s16le_to_floats(b"\x00")
|
||||
with self.assertRaisesRegex(AudioError, "duration"):
|
||||
pcm_s16le_to_floats(b"\x00\x00" * 5, max_bytes=8)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,246 @@
|
||||
import io
|
||||
import subprocess
|
||||
import threading
|
||||
import unittest
|
||||
|
||||
from pi_dictate_service.recorder import PW_RECORD_COMMAND, PwRecorder, RecorderError
|
||||
|
||||
|
||||
class TrackingStream(io.BytesIO):
|
||||
def __init__(self, value: bytes, events: list[str], name: str) -> None:
|
||||
super().__init__(value)
|
||||
self._events = events
|
||||
self._name = name
|
||||
self.eof = threading.Event()
|
||||
|
||||
def read(self, size: int = -1) -> bytes:
|
||||
value = super().read(size)
|
||||
if not value and not self.eof.is_set():
|
||||
self._events.append(f"{self._name}-eof")
|
||||
self.eof.set()
|
||||
return value
|
||||
|
||||
|
||||
class FakeProcess:
|
||||
def __init__(
|
||||
self,
|
||||
stdout: TrackingStream,
|
||||
stderr: TrackingStream,
|
||||
*,
|
||||
returncode: int | None = None,
|
||||
timeout: bool = False,
|
||||
terminate_returncode: int = -15,
|
||||
) -> None:
|
||||
self.pid = 1234
|
||||
self.stdout = stdout
|
||||
self.stderr = stderr
|
||||
self.returncode = returncode
|
||||
self._timeout = timeout
|
||||
self._terminate_returncode = terminate_returncode
|
||||
self.events = stdout._events
|
||||
|
||||
def poll(self) -> int | None:
|
||||
return self.returncode
|
||||
|
||||
def terminate(self) -> None:
|
||||
self.events.append("terminate")
|
||||
if not self._timeout:
|
||||
self.returncode = self._terminate_returncode
|
||||
|
||||
def kill(self) -> None:
|
||||
self.events.append("kill")
|
||||
self.returncode = -9
|
||||
self._timeout = False
|
||||
|
||||
def wait(self, timeout: float | None = None) -> int:
|
||||
self.events.append("wait")
|
||||
if self._timeout:
|
||||
raise subprocess.TimeoutExpired("pw-record", timeout)
|
||||
if self.returncode is None:
|
||||
self.returncode = 0
|
||||
return self.returncode
|
||||
|
||||
|
||||
class ProcessFactory:
|
||||
def __init__(self, process: FakeProcess) -> None:
|
||||
self.process = process
|
||||
self.calls: list[tuple[list[str], dict[str, object]]] = []
|
||||
|
||||
def __call__(self, command: list[str], **kwargs: object) -> FakeProcess:
|
||||
self.calls.append((command, kwargs))
|
||||
return self.process
|
||||
|
||||
|
||||
class Clock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
|
||||
class PwRecorderTests(unittest.TestCase):
|
||||
def make_recorder(
|
||||
self,
|
||||
pcm: bytes = b"\x00\x00\x01\x00",
|
||||
stderr: bytes = b"",
|
||||
*,
|
||||
returncode: int | None = None,
|
||||
timeout: bool = False,
|
||||
max_bytes: int = 1024,
|
||||
stderr_limit: int = 64 * 1024,
|
||||
clock: Clock | None = None,
|
||||
terminate_returncode: int = -15,
|
||||
) -> tuple[PwRecorder, FakeProcess, ProcessFactory, list[str]]:
|
||||
events: list[str] = []
|
||||
process = FakeProcess(
|
||||
TrackingStream(pcm, events, "stdout"),
|
||||
TrackingStream(stderr, events, "stderr"),
|
||||
returncode=returncode,
|
||||
timeout=timeout,
|
||||
terminate_returncode=terminate_returncode,
|
||||
)
|
||||
factory = ProcessFactory(process)
|
||||
def signal_group(pid: int, signal_number: int) -> None:
|
||||
self.assertEqual(pid, process.pid)
|
||||
if signal_number == 15:
|
||||
process.terminate()
|
||||
elif signal_number == 9:
|
||||
process.kill()
|
||||
else:
|
||||
self.fail(f"unexpected signal: {signal_number}")
|
||||
|
||||
recorder = PwRecorder(
|
||||
process_factory=factory,
|
||||
max_bytes=max_bytes,
|
||||
stderr_limit=stderr_limit,
|
||||
stop_timeout=0.01,
|
||||
clock=clock or Clock(),
|
||||
process_group_signal=signal_group,
|
||||
)
|
||||
return recorder, process, factory, events
|
||||
|
||||
def test_spawns_exact_pw_record_command_and_drains_bounded_streams(self) -> None:
|
||||
recorder, process, factory, events = self.make_recorder(
|
||||
pcm=b"\x00\x00\x01\x00", stderr=b"0123456789", stderr_limit=8
|
||||
)
|
||||
|
||||
recorder.start()
|
||||
self.assertTrue(process.stdout.eof.wait(1))
|
||||
capture = recorder.stop()
|
||||
|
||||
self.assertEqual(factory.calls[0][0], PW_RECORD_COMMAND)
|
||||
self.assertIs(factory.calls[0][1]["stdout"], subprocess.PIPE)
|
||||
self.assertIs(factory.calls[0][1]["stderr"], subprocess.PIPE)
|
||||
self.assertIs(factory.calls[0][1]["start_new_session"], True)
|
||||
self.assertEqual(capture.pcm, b"\x00\x00\x01\x00")
|
||||
self.assertEqual(capture.stderr_tail, b"23456789")
|
||||
self.assertIn("stdout-eof", events)
|
||||
self.assertIn("wait", events)
|
||||
|
||||
def test_accepts_pw_record_code_one_after_intentional_sigterm(self) -> None:
|
||||
recorder, process, _factory, _events = self.make_recorder(terminate_returncode=1)
|
||||
recorder.start()
|
||||
|
||||
capture = recorder.stop()
|
||||
|
||||
self.assertEqual(capture.returncode, 1)
|
||||
self.assertIsNone(capture.signal)
|
||||
self.assertEqual(capture.pcm, b"\x00\x00\x01\x00")
|
||||
self.assertEqual(process.returncode, 1)
|
||||
|
||||
def test_rejects_code_one_if_process_already_exited_or_has_stderr(self) -> None:
|
||||
exited, exited_process, _factory, _events = self.make_recorder(returncode=1)
|
||||
exited.start()
|
||||
self.assertTrue(exited_process.stdout.eof.wait(1))
|
||||
with self.assertRaises(RecorderError):
|
||||
exited.stop()
|
||||
|
||||
noisy, _process, _factory, _events = self.make_recorder(
|
||||
stderr=b"pipewire failed",
|
||||
terminate_returncode=1,
|
||||
)
|
||||
noisy.start()
|
||||
with self.assertRaises(RecorderError):
|
||||
noisy.stop()
|
||||
|
||||
def test_shutdown_timeout_escalates_to_kill_and_fails(self) -> None:
|
||||
recorder, _process, _factory, events = self.make_recorder(timeout=True)
|
||||
recorder.start()
|
||||
|
||||
with self.assertRaises(RecorderError) as failed:
|
||||
recorder.stop()
|
||||
|
||||
self.assertEqual(failed.exception.code, "recorder_timeout")
|
||||
self.assertLess(events.index("terminate"), events.index("kill"))
|
||||
|
||||
def test_classifies_unexpected_exit_code_and_signal(self) -> None:
|
||||
recorder, process, _factory, _events = self.make_recorder(
|
||||
stderr=b"pipewire failed", returncode=7
|
||||
)
|
||||
recorder.start()
|
||||
self.assertTrue(process.stdout.eof.wait(1))
|
||||
|
||||
failure = recorder.poll_unexpected()
|
||||
|
||||
self.assertIsNotNone(failure)
|
||||
assert failure is not None
|
||||
self.assertEqual(failure.code, "recorder_exit")
|
||||
self.assertIn("code 7", failure.detail)
|
||||
self.assertIn("pipewire failed", failure.detail)
|
||||
|
||||
signalled, signal_process, _factory, _events = self.make_recorder(returncode=-11)
|
||||
signalled.start()
|
||||
self.assertTrue(signal_process.stdout.eof.wait(1))
|
||||
signal_failure = signalled.poll_unexpected()
|
||||
assert signal_failure is not None
|
||||
self.assertIn("signal 11", signal_failure.detail)
|
||||
|
||||
def test_duration_limit_discards_audio_and_reports_failure(self) -> None:
|
||||
recorder, process, _factory, _events = self.make_recorder(
|
||||
pcm=b"\x00\x00" * 10, max_bytes=8
|
||||
)
|
||||
recorder.start()
|
||||
self.assertTrue(process.stdout.eof.wait(1))
|
||||
|
||||
failure = recorder.poll_unexpected()
|
||||
|
||||
assert failure is not None
|
||||
self.assertEqual(failure.code, "duration_limit")
|
||||
self.assertEqual(recorder.buffered_bytes, 0)
|
||||
|
||||
def test_stop_reports_a_duration_limit_that_was_already_reached(self) -> None:
|
||||
recorder, process, _factory, _events = self.make_recorder(
|
||||
pcm=b"\x00\x00" * 10, max_bytes=8
|
||||
)
|
||||
recorder.start()
|
||||
self.assertTrue(process.stdout.eof.wait(1))
|
||||
|
||||
with self.assertRaises(RecorderError) as failed:
|
||||
recorder.stop()
|
||||
|
||||
self.assertEqual(failed.exception.code, "duration_limit")
|
||||
self.assertEqual(recorder.buffered_bytes, 0)
|
||||
|
||||
def test_duration_limit_kills_a_recorder_that_ignores_sigterm(self) -> None:
|
||||
clock = Clock()
|
||||
recorder, process, _factory, events = self.make_recorder(
|
||||
pcm=b"\x00\x00" * 10,
|
||||
max_bytes=8,
|
||||
timeout=True,
|
||||
clock=clock,
|
||||
)
|
||||
recorder.start()
|
||||
self.assertTrue(process.stdout.eof.wait(1))
|
||||
self.assertIsNone(recorder.poll_unexpected())
|
||||
|
||||
clock.now = 0.02
|
||||
failure = recorder.poll_unexpected()
|
||||
|
||||
assert failure is not None
|
||||
self.assertEqual(failure.code, "duration_limit")
|
||||
self.assertIn("kill", events)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,333 @@
|
||||
from concurrent.futures import Future
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
import unittest
|
||||
|
||||
from pi_dictate_service.protocol import MAX_TRANSCRIPT_BYTES, parse_request
|
||||
from pi_dictate_service.recorder import Capture, RecorderError
|
||||
from pi_dictate_service.server import DictateServiceCore, Outbound
|
||||
|
||||
|
||||
class Clock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
def advance(self, seconds: float) -> None:
|
||||
self.now += seconds
|
||||
|
||||
|
||||
@dataclass
|
||||
class PendingWork:
|
||||
future: Future[object]
|
||||
function: object
|
||||
args: tuple[object, ...]
|
||||
|
||||
|
||||
class ControlledExecutor:
|
||||
def __init__(self) -> None:
|
||||
self.pending: list[PendingWork] = []
|
||||
|
||||
def submit(self, function: object, *args: object) -> Future[object]:
|
||||
future: Future[object] = Future()
|
||||
self.pending.append(PendingWork(future, function, args))
|
||||
return future
|
||||
|
||||
def run_next(self) -> None:
|
||||
work = self.pending.pop(0)
|
||||
try:
|
||||
result = work.function(*work.args) # type: ignore[operator]
|
||||
except BaseException as error:
|
||||
work.future.set_exception(error)
|
||||
else:
|
||||
work.future.set_result(result)
|
||||
|
||||
def shutdown(self, wait: bool = True, cancel_futures: bool = False) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class FakeModel:
|
||||
def __init__(self) -> None:
|
||||
self.loaded = False
|
||||
self.transcript = "dictated text"
|
||||
self.load_error: BaseException | None = None
|
||||
self.transcribe_error: BaseException | None = None
|
||||
self.samples: Sequence[float] | None = None
|
||||
|
||||
def load(self) -> None:
|
||||
if self.load_error:
|
||||
raise self.load_error
|
||||
self.loaded = True
|
||||
|
||||
def transcribe(self, samples: Sequence[float]) -> str:
|
||||
self.samples = samples
|
||||
if self.transcribe_error:
|
||||
raise self.transcribe_error
|
||||
return self.transcript
|
||||
|
||||
|
||||
class FakeRecorder:
|
||||
def __init__(self) -> None:
|
||||
self.started = False
|
||||
self.cancelled = False
|
||||
self.capture = Capture(b"\x00\x00\x00\x40", -15, 15, b"")
|
||||
self.stop_error: RecorderError | None = None
|
||||
self.unexpected: RecorderError | None = None
|
||||
|
||||
def start(self) -> None:
|
||||
self.started = True
|
||||
|
||||
def stop(self) -> Capture:
|
||||
if self.stop_error:
|
||||
raise self.stop_error
|
||||
return self.capture
|
||||
|
||||
def cancel(self) -> None:
|
||||
self.cancelled = True
|
||||
|
||||
def poll_unexpected(self) -> RecorderError | None:
|
||||
failure, self.unexpected = self.unexpected, None
|
||||
return failure
|
||||
|
||||
|
||||
class RecorderFactory:
|
||||
def __init__(self) -> None:
|
||||
self.instances: list[FakeRecorder] = []
|
||||
|
||||
def __call__(self) -> FakeRecorder:
|
||||
recorder = FakeRecorder()
|
||||
self.instances.append(recorder)
|
||||
return recorder
|
||||
|
||||
|
||||
class CoreTestCase(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.clock = Clock()
|
||||
self.executor = ControlledExecutor()
|
||||
self.model = FakeModel()
|
||||
self.recorders = RecorderFactory()
|
||||
self.core = DictateServiceCore(
|
||||
model=self.model,
|
||||
recorder_factory=self.recorders,
|
||||
executor=self.executor,
|
||||
clock=self.clock,
|
||||
idle_seconds=900,
|
||||
)
|
||||
|
||||
def request(
|
||||
self,
|
||||
client: str,
|
||||
request_type: str,
|
||||
request_id: str,
|
||||
generation_id: str | None = None,
|
||||
) -> list[dict[str, object]]:
|
||||
frame: dict[str, object] = {
|
||||
"version": 1,
|
||||
"type": request_type,
|
||||
"requestId": request_id,
|
||||
}
|
||||
if generation_id is not None:
|
||||
frame["generationId"] = generation_id
|
||||
self.core.handle_request(client, parse_request(frame))
|
||||
return self.frames(self.core.drain_events(), client)
|
||||
|
||||
@staticmethod
|
||||
def frames(events: list[Outbound], client: str | None = None) -> list[dict[str, object]]:
|
||||
return [event.frame for event in events if client is None or event.client_id == client]
|
||||
|
||||
def hello(self, client: str = "a") -> None:
|
||||
events = self.request(client, "hello", f"hello-{client}")
|
||||
self.assertEqual(events[0]["type"], "pong")
|
||||
|
||||
def load_to_listening(self, client: str = "a", generation: str = "g1") -> FakeRecorder:
|
||||
self.hello(client)
|
||||
self.request(client, "start", "start", generation)
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
events = self.frames(self.core.drain_events(), client)
|
||||
self.assertEqual(events[-1]["state"], "listening")
|
||||
self.assertEqual(events[-1]["requestId"], "start")
|
||||
return self.recorders.instances[-1]
|
||||
|
||||
|
||||
class HandshakeAndOwnershipTests(CoreTestCase):
|
||||
def test_hello_is_required_and_valid_errors_keep_session_open(self) -> None:
|
||||
events = self.request("a", "start", "s1", "g1")
|
||||
self.assertEqual(events[0]["code"], "hello_required")
|
||||
|
||||
self.hello("a")
|
||||
self.request("a", "start", "s2", "g1")
|
||||
self.assertEqual(len(self.executor.pending), 1)
|
||||
|
||||
def test_one_generation_is_globally_owned(self) -> None:
|
||||
self.load_to_listening()
|
||||
self.hello("b")
|
||||
|
||||
busy = self.request("b", "start", "s2", "g2")
|
||||
self.assertEqual(busy[0]["code"], "busy")
|
||||
not_owner = self.request("b", "stop", "x", "g1")
|
||||
self.assertEqual(not_owner[0]["code"], "not_owner")
|
||||
stale = self.request("a", "stop", "x2", "old")
|
||||
self.assertEqual(stale[0]["code"], "generation_mismatch")
|
||||
|
||||
stopping = self.request("a", "stop", "x3", "g1")
|
||||
self.assertEqual(stopping[0]["state"], "stopping")
|
||||
|
||||
|
||||
class CancellationTests(CoreTestCase):
|
||||
def test_cancel_during_load_is_immediate_and_blocks_until_worker_finishes(self) -> None:
|
||||
self.hello("a")
|
||||
self.request("a", "start", "start", "g1")
|
||||
|
||||
cancelled = self.request("a", "cancel", "cancel", "g1")
|
||||
self.assertEqual(cancelled[0]["type"], "cancelled")
|
||||
self.hello("b")
|
||||
self.assertEqual(self.request("b", "start", "start-b", "g2")[0]["code"], "busy")
|
||||
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
self.assertEqual(self.recorders.instances, [])
|
||||
self.assertNotIn("result", [event["type"] for event in self.frames(self.core.drain_events())])
|
||||
|
||||
def test_cancel_during_inference_suppresses_later_result(self) -> None:
|
||||
self.load_to_listening()
|
||||
self.request("a", "stop", "stop", "g1")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
transcribing = self.frames(self.core.drain_events(), "a")
|
||||
self.assertEqual(transcribing[0]["state"], "transcribing")
|
||||
|
||||
cancelled = self.request("a", "cancel", "cancel", "g1")
|
||||
self.assertEqual(cancelled[0]["type"], "cancelled")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
self.assertNotIn("result", [event["type"] for event in self.frames(self.core.drain_events())])
|
||||
|
||||
def test_disconnect_cancels_capture_and_tombstones_workers(self) -> None:
|
||||
recorder = self.load_to_listening()
|
||||
|
||||
self.core.disconnect("a")
|
||||
self.assertEqual(len(self.executor.pending), 1)
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
|
||||
self.assertTrue(recorder.cancelled)
|
||||
self.assertEqual(self.core.active_generation_id, None)
|
||||
|
||||
|
||||
class StopAndInferenceTests(CoreTestCase):
|
||||
def test_stop_drains_then_converts_before_inference_and_emits_result(self) -> None:
|
||||
self.load_to_listening()
|
||||
|
||||
stopping = self.request("a", "stop", "stop", "g1")
|
||||
self.assertEqual(stopping[0]["state"], "stopping")
|
||||
self.assertEqual(stopping[0]["requestId"], "stop")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
transcribing = self.frames(self.core.drain_events(), "a")
|
||||
self.assertEqual(transcribing[0]["state"], "transcribing")
|
||||
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
completed = self.frames(self.core.drain_events(), "a")
|
||||
|
||||
self.assertEqual(list(self.model.samples or []), [0.0, 0.5])
|
||||
self.assertEqual(completed[0]["type"], "result")
|
||||
self.assertEqual(completed[0]["transcript"], "dictated text")
|
||||
self.assertEqual(completed[1]["state"], "idle")
|
||||
|
||||
def test_empty_transcript_is_a_result(self) -> None:
|
||||
self.model.transcript = ""
|
||||
self.load_to_listening()
|
||||
self.request("a", "stop", "stop", "g1")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
self.core.drain_events()
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
|
||||
completed = self.frames(self.core.drain_events(), "a")
|
||||
self.assertEqual(completed[0]["type"], "result")
|
||||
self.assertEqual(completed[0]["transcript"], "")
|
||||
|
||||
def test_model_and_transcript_failures_are_bounded_terminal_errors(self) -> None:
|
||||
self.model.load_error = RuntimeError("CUDA unavailable")
|
||||
self.hello()
|
||||
self.request("a", "start", "start", "g1")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
failed = self.frames(self.core.drain_events(), "a")
|
||||
self.assertEqual(failed[0]["code"], "model_failure")
|
||||
self.assertEqual(failed[-1]["state"], "idle")
|
||||
|
||||
model = FakeModel()
|
||||
model.loaded = True
|
||||
model.transcript = "x" * (MAX_TRANSCRIPT_BYTES + 1)
|
||||
core = DictateServiceCore(
|
||||
model=model,
|
||||
recorder_factory=self.recorders,
|
||||
executor=self.executor,
|
||||
clock=self.clock,
|
||||
)
|
||||
core.handle_request("a", parse_request({"version": 1, "type": "hello", "requestId": "h"}))
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({"version": 1, "type": "start", "requestId": "s", "generationId": "g"}))
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({"version": 1, "type": "stop", "requestId": "t", "generationId": "g"}))
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
core.drain_events()
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
failed = self.frames(core.drain_events(), "a")
|
||||
self.assertEqual(failed[0]["code"], "transcript_too_large")
|
||||
|
||||
def test_unexpected_recorder_exit_is_terminal(self) -> None:
|
||||
recorder = self.load_to_listening()
|
||||
recorder.unexpected = RecorderError("recorder_exit", "pw-record exited with code 3")
|
||||
|
||||
self.core.poll()
|
||||
failed = self.frames(self.core.drain_events(), "a")
|
||||
|
||||
self.assertEqual(failed[0]["code"], "recorder_exit")
|
||||
self.assertEqual(failed[-1]["state"], "idle")
|
||||
|
||||
def test_stop_does_not_reclassify_an_already_exited_recorder_as_intentional(self) -> None:
|
||||
recorder = self.load_to_listening()
|
||||
recorder.unexpected = RecorderError("recorder_exit", "pw-record exited with code 0")
|
||||
|
||||
failed = self.request("a", "stop", "stop", "g1")
|
||||
|
||||
self.assertEqual(failed[0]["code"], "recorder_exit")
|
||||
self.assertEqual(failed[-1]["state"], "idle")
|
||||
self.assertEqual(self.executor.pending, [])
|
||||
|
||||
|
||||
class IdleTests(CoreTestCase):
|
||||
def test_only_valid_requests_reset_idle_and_workers_suspend_then_restart_timer(self) -> None:
|
||||
self.core.accepted("a")
|
||||
self.clock.advance(899)
|
||||
self.assertFalse(self.core.should_exit())
|
||||
self.request("a", "hello", "h")
|
||||
self.clock.advance(899)
|
||||
self.assertFalse(self.core.should_exit())
|
||||
|
||||
self.request("a", "start", "s", "g")
|
||||
self.clock.advance(901)
|
||||
self.assertFalse(self.core.should_exit())
|
||||
self.request("a", "cancel", "c", "g")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
self.core.drain_events()
|
||||
self.assertFalse(self.core.should_exit())
|
||||
|
||||
self.clock.advance(899)
|
||||
self.assertFalse(self.core.should_exit())
|
||||
self.clock.advance(1)
|
||||
self.assertTrue(self.core.should_exit())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,38 @@
|
||||
from pathlib import Path
|
||||
import unittest
|
||||
|
||||
|
||||
REPOSITORY_ROOT = Path(__file__).resolve().parents[3]
|
||||
UNIT_DIR = REPOSITORY_ROOT / "systemd/user"
|
||||
|
||||
|
||||
class SystemdUnitTests(unittest.TestCase):
|
||||
def test_socket_unit_owns_the_exact_private_runtime_path(self) -> None:
|
||||
socket_unit = (UNIT_DIR / "pi-dictate.socket").read_text(encoding="utf-8")
|
||||
self.assertIn("ListenStream=%t/pi-dictate.sock", socket_unit)
|
||||
self.assertIn("SocketMode=0600", socket_unit)
|
||||
self.assertIn("Accept=no", socket_unit)
|
||||
self.assertIn("RemoveOnStop=yes", socket_unit)
|
||||
|
||||
def test_service_is_offline_hardened_and_uses_private_native_caches(self) -> None:
|
||||
service_unit = (UNIT_DIR / "pi-dictate.service").read_text(encoding="utf-8")
|
||||
for expected in (
|
||||
"PI_DICTATE_MODEL_PATH=%h/.local/share/pi-dictate/models/",
|
||||
"PI_DICTATE_TERMS_PATH=%h/.local/share/pi-dictate/config/coding-terms.txt",
|
||||
"PI_DICTATE_BOOST_ALPHA=1.0",
|
||||
"HF_HUB_OFFLINE=1",
|
||||
"TRANSFORMERS_OFFLINE=1",
|
||||
"TRITON_CACHE_DIR=/tmp/pi-dictate-triton",
|
||||
"NUMBA_CACHE_DIR=/tmp/pi-dictate-numba",
|
||||
"LimitCORE=0",
|
||||
"MemorySwapMax=0",
|
||||
"ProtectSystem=strict",
|
||||
"ProtectHome=read-only",
|
||||
"RestrictAddressFamilies=AF_UNIX",
|
||||
"IPAddressDeny=any",
|
||||
):
|
||||
self.assertIn(expected, service_unit)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,225 @@
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import socket
|
||||
import struct
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from pi_dictate_service.server import (
|
||||
MAX_CLIENT_CONNECTIONS,
|
||||
DictateServiceCore,
|
||||
inherited_listener,
|
||||
main,
|
||||
serve,
|
||||
)
|
||||
|
||||
|
||||
class LoadedModel:
|
||||
loaded = True
|
||||
|
||||
def load(self) -> None:
|
||||
return None
|
||||
|
||||
def transcribe(self, samples: list[float]) -> str:
|
||||
return ""
|
||||
|
||||
|
||||
class UnixLoopTests(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.temporary = tempfile.TemporaryDirectory()
|
||||
self.socket_path = Path(self.temporary.name) / "dictate.sock"
|
||||
self.listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
self.listener.bind(str(self.socket_path))
|
||||
self.listener.listen()
|
||||
self.core = DictateServiceCore(model=LoadedModel(), idle_seconds=0.4)
|
||||
self.failure: BaseException | None = None
|
||||
|
||||
def run() -> None:
|
||||
try:
|
||||
serve(self.listener, self.core)
|
||||
except BaseException as error:
|
||||
self.failure = error
|
||||
|
||||
self.thread = threading.Thread(target=run)
|
||||
self.thread.start()
|
||||
|
||||
def tearDown(self) -> None:
|
||||
self.thread.join(timeout=2)
|
||||
if self.thread.is_alive():
|
||||
self.fail("server loop did not exit after its idle timeout")
|
||||
if self.failure:
|
||||
raise self.failure
|
||||
self.temporary.cleanup()
|
||||
|
||||
def connect(self) -> socket.socket:
|
||||
client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
client.settimeout(1)
|
||||
client.connect(str(self.socket_path))
|
||||
return client
|
||||
|
||||
@staticmethod
|
||||
def receive_frame(client: socket.socket) -> dict[str, object]:
|
||||
data = bytearray()
|
||||
while not data.endswith(b"\n"):
|
||||
chunk = client.recv(4096)
|
||||
if not chunk:
|
||||
break
|
||||
data.extend(chunk)
|
||||
return json.loads(data)
|
||||
|
||||
def test_accepts_fragmented_requests_and_writes_persistent_responses(self) -> None:
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":"hello",')
|
||||
client.sendall(b'"requestId":"hello"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["type"], "pong")
|
||||
|
||||
client.sendall(b'{"version":1,"type":"ping","requestId":"ping"}\n')
|
||||
pong = self.receive_frame(client)
|
||||
self.assertEqual(pong["requestId"], "ping")
|
||||
self.assertEqual(pong["service"], "pi-dictate")
|
||||
|
||||
def test_valid_state_error_keeps_open_but_malformed_frame_errors_once_then_closes(self) -> None:
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":"hello","requestId":"hello"}\n')
|
||||
self.receive_frame(client)
|
||||
client.sendall(
|
||||
b'{"version":1,"type":"stop","requestId":"stop","generationId":"g"}\n'
|
||||
)
|
||||
self.assertEqual(self.receive_frame(client)["code"], "invalid_state")
|
||||
client.sendall(b'{"version":1,"type":"ping","requestId":"ping"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["type"], "pong")
|
||||
|
||||
client.sendall(b'{broken}\n')
|
||||
self.assertEqual(self.receive_frame(client)["code"], "malformed_json")
|
||||
self.assertEqual(client.recv(1), b"")
|
||||
|
||||
self.assertTrue(self.socket_path.exists(), "the service must not unlink a listener-owned path")
|
||||
|
||||
def test_malicious_json_values_do_not_crash_the_service(self) -> None:
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":[],"requestId":"bad-type"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["code"], "invalid_request")
|
||||
client.sendall(b'{"version":1,"type":"ping","requestId":"\\ud800"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["code"], "invalid_request")
|
||||
|
||||
with self.connect() as client:
|
||||
huge_integer = b"9" * 5000
|
||||
client.sendall(
|
||||
b'{"version":' + huge_integer + b',"type":"hello","requestId":"huge"}\n'
|
||||
)
|
||||
self.assertEqual(self.receive_frame(client)["code"], "malformed_json")
|
||||
self.assertEqual(client.recv(1), b"")
|
||||
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":"hello","requestId":"healthy"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["type"], "pong")
|
||||
|
||||
def test_closes_non_reading_client_before_output_queue_can_grow_unbounded(self) -> None:
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":"hello","requestId":"hello"}\n')
|
||||
self.receive_frame(client)
|
||||
ping = b'{"version":1,"type":"ping","requestId":"p"}\n'
|
||||
try:
|
||||
client.sendall(ping * 1200)
|
||||
except BrokenPipeError:
|
||||
pass
|
||||
received = bytearray()
|
||||
while True:
|
||||
try:
|
||||
chunk = client.recv(64 * 1024)
|
||||
except ConnectionResetError:
|
||||
break
|
||||
if not chunk:
|
||||
break
|
||||
received.extend(chunk)
|
||||
self.assertLessEqual(len(received), 64 * 1024)
|
||||
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":"hello","requestId":"healthy"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["type"], "pong")
|
||||
|
||||
def test_rejects_connections_above_the_fixed_limit(self) -> None:
|
||||
clients = [self.connect() for _ in range(MAX_CLIENT_CONNECTIONS)]
|
||||
try:
|
||||
time.sleep(0.1)
|
||||
overflow = self.connect()
|
||||
try:
|
||||
overflow.sendall(b'{"version":1,"type":"hello","requestId":"overflow"}\n')
|
||||
with self.assertRaises((ConnectionResetError, BrokenPipeError)):
|
||||
if overflow.recv(1) == b"":
|
||||
raise ConnectionResetError
|
||||
finally:
|
||||
overflow.close()
|
||||
finally:
|
||||
for client in clients:
|
||||
client.close()
|
||||
|
||||
def test_client_reset_does_not_crash_the_service(self) -> None:
|
||||
reset_client = self.connect()
|
||||
reset_client.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
|
||||
reset_client.sendall(b'{"version":1,"type":"hello","requestId":"abandoned"}\n')
|
||||
reset_client.close()
|
||||
time.sleep(0.15)
|
||||
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":"hello","requestId":"hello"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["type"], "pong")
|
||||
|
||||
|
||||
class ActivationTests(unittest.TestCase):
|
||||
def test_inherited_listener_requires_one_fd_for_this_process_and_wraps_fd_three(self) -> None:
|
||||
listener = Mock()
|
||||
with patch("pi_dictate_service.server.os.getpid", return_value=42), patch(
|
||||
"pi_dictate_service.server.socket.socket", return_value=listener
|
||||
) as socket_factory:
|
||||
self.assertIs(
|
||||
inherited_listener({"LISTEN_PID": "42", "LISTEN_FDS": "1"}),
|
||||
listener,
|
||||
)
|
||||
|
||||
socket_factory.assert_called_once_with(fileno=3)
|
||||
listener.setblocking.assert_called_once_with(False)
|
||||
|
||||
with patch("pi_dictate_service.server.os.getpid", return_value=42):
|
||||
with self.assertRaisesRegex(RuntimeError, "exactly one"):
|
||||
inherited_listener({"LISTEN_PID": "41", "LISTEN_FDS": "1"})
|
||||
|
||||
def test_main_requires_local_model_path_and_reads_idle_override(self) -> None:
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
with self.assertRaisesRegex(SystemExit, "PI_DICTATE_MODEL_PATH"):
|
||||
main()
|
||||
|
||||
model = object()
|
||||
listener = object()
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"PI_DICTATE_MODEL_PATH": "/models/parakeet.nemo",
|
||||
"PI_DICTATE_TERMS_PATH": "/config/coding-terms.txt",
|
||||
"PI_DICTATE_BOOST_ALPHA": "0.5",
|
||||
"PI_DICTATE_IDLE_SECONDS": "12.5",
|
||||
},
|
||||
clear=True,
|
||||
), patch("pi_dictate_service.server.NemoParakeetModel", return_value=model) as adapter, patch(
|
||||
"pi_dictate_service.server.DictateServiceCore"
|
||||
) as core_factory, patch(
|
||||
"pi_dictate_service.server.inherited_listener", return_value=listener
|
||||
), patch("pi_dictate_service.server.serve") as serve_loop:
|
||||
core = core_factory.return_value
|
||||
main()
|
||||
|
||||
adapter.assert_called_once_with(
|
||||
"/models/parakeet.nemo",
|
||||
terms_path="/config/coding-terms.txt",
|
||||
boost_alpha=0.5,
|
||||
)
|
||||
core_factory.assert_called_once_with(model=model, idle_seconds=12.5)
|
||||
serve_loop.assert_called_once_with(listener, core)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Generated
+2238
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,37 @@
|
||||
# Managed by pi-dictate install.sh
|
||||
[Unit]
|
||||
Description=Pi local Parakeet dictation service
|
||||
Requires=pi-dictate.socket
|
||||
After=pipewire.service
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
ExecStart=%h/.local/share/pi-dictate/venv/bin/pi-dictate-service
|
||||
Environment=PI_DICTATE_MODEL_PATH=%h/.local/share/pi-dictate/models/parakeet-tdt-0.6b-v3/parakeet-tdt-0.6b-v3.nemo
|
||||
Environment=PI_DICTATE_TERMS_PATH=%h/.local/share/pi-dictate/config/coding-terms.txt
|
||||
Environment=PI_DICTATE_BOOST_ALPHA=1.0
|
||||
Environment=PI_DICTATE_IDLE_SECONDS=900
|
||||
Environment=HF_HUB_OFFLINE=1
|
||||
Environment=TRANSFORMERS_OFFLINE=1
|
||||
Environment=NEMO_HOME=/tmp/pi-dictate-nemo
|
||||
Environment=HF_HOME=/tmp/pi-dictate-huggingface
|
||||
Environment=TRITON_CACHE_DIR=/tmp/pi-dictate-triton
|
||||
Environment=NUMBA_CACHE_DIR=/tmp/pi-dictate-numba
|
||||
Environment=PYTHONDONTWRITEBYTECODE=1
|
||||
UMask=0077
|
||||
LimitCORE=0
|
||||
MemorySwapMax=0
|
||||
NoNewPrivileges=yes
|
||||
PrivateTmp=yes
|
||||
ProtectClock=yes
|
||||
ProtectControlGroups=yes
|
||||
ProtectKernelLogs=yes
|
||||
ProtectKernelModules=yes
|
||||
ProtectKernelTunables=yes
|
||||
ProtectSystem=strict
|
||||
ProtectHome=read-only
|
||||
RestrictAddressFamilies=AF_UNIX
|
||||
IPAddressDeny=any
|
||||
LockPersonality=yes
|
||||
|
||||
# The socket unit owns %t/pi-dictate.sock. The service must not unlink it.
|
||||
@@ -0,0 +1,14 @@
|
||||
# Managed by pi-dictate install.sh
|
||||
[Unit]
|
||||
Description=Pi local dictation socket
|
||||
|
||||
[Socket]
|
||||
ListenStream=%t/pi-dictate.sock
|
||||
SocketMode=0600
|
||||
DirectoryMode=0700
|
||||
Accept=no
|
||||
Service=pi-dictate.service
|
||||
RemoveOnStop=yes
|
||||
|
||||
[Install]
|
||||
WantedBy=sockets.target
|
||||
Reference in New Issue
Block a user