mj_controller/
dictation.rs

1//! Shared subscription-backed web and native prompt dictation support.
2//!
3//! Shared credential discovery is kept in `mj_client` because both the terminal
4//! chat and authenticated web server need the same decision about which Codex
5//! profile may transcribe audio. HTTP handlers send typed requests here;
6//! filesystem access and the provider call stay off the controller event loop.
7
8use std::path::PathBuf;
9use std::time::Duration;
10
11use anvil_client::codex_client::CodexClient;
12use anvil_client::transcribe::TranscribeRequest;
13use axum::body::Bytes;
14use tokio::sync::oneshot;
15use tokio_util::sync::CancellationToken;
16
17pub use mj_client::auth::{auth_paths, available_auth};
18
19/// Maximum complete WAV upload accepted by the web endpoint.
20pub const MAX_AUDIO_BYTES: usize = 20 * 1024 * 1024;
21/// Maximum audio duration accepted by the web endpoint.
22pub const MAX_AUDIO_DURATION: Duration = Duration::from_secs(600);
23/// Maximum audio bytes in 16 kHz mono PCM16. This is stricter than the whole
24/// RIFF envelope limit and avoids accepting a file whose header claims a
25/// shorter duration than its sample payload actually contains.
26const MAX_PCM_DATA_BYTES: u64 = 16_000 * 2 * MAX_AUDIO_DURATION.as_secs();
27/// Whole-request deadline for credential inspection and provider work.
28pub const DICTATION_TIMEOUT: Duration = Duration::from_secs(120);
29
30/// An operation submitted by an authenticated HTTP surface.
31#[derive(Debug)]
32pub enum DictationOperation {
33    /// Probe the selected session's profile set without contacting the provider.
34    Availability,
35    /// Validate and transcribe one complete WAV upload.
36    Transcribe(Bytes),
37}
38
39/// A request crossing the HTTP/controller boundary.
40#[derive(Debug)]
41pub struct DictationRequest {
42    pub session_id: String,
43    pub operation: DictationOperation,
44    pub cancel: CancellationToken,
45    pub reply: oneshot::Sender<Result<DictationResponse, DictationError>>,
46}
47
48/// Result returned to an HTTP handler.
49#[derive(Debug, Clone, PartialEq, Eq)]
50pub enum DictationResponse {
51    Availability {
52        available: bool,
53        reason: Option<String>,
54    },
55    Transcript {
56        text: String,
57    },
58}
59
60/// Errors that cross the controller boundary. Provider details are retained
61/// for logs and diagnostics, while the HTTP layer maps them to safe messages.
62#[derive(Debug)]
63pub enum DictationError {
64    SessionNotFound,
65    CredentialsUnavailable,
66    InvalidAudio(&'static str),
67    Cancelled,
68    TimedOut,
69    CredentialProbe,
70    Provider(String),
71}
72
73impl std::fmt::Display for DictationError {
74    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        match self {
76            Self::SessionNotFound => formatter.write_str("unknown session"),
77            Self::CredentialsUnavailable => {
78                formatter.write_str("no Codex subscription credentials are available")
79            }
80            Self::InvalidAudio(message) => write!(formatter, "invalid WAV audio: {message}"),
81            Self::Cancelled => formatter.write_str("dictation was cancelled"),
82            Self::TimedOut => formatter.write_str("dictation timed out"),
83            Self::CredentialProbe => formatter.write_str("could not inspect dictation credentials"),
84            Self::Provider(message) => {
85                write!(formatter, "transcription provider failed: {message}")
86            }
87        }
88    }
89}
90
91impl std::error::Error for DictationError {}
92
93/// Validate a complete RIFF/WAVE file before any credentials or provider work
94/// is started. Unknown chunks are allowed, but every chunk length and padding
95/// byte must fit inside the declared RIFF envelope.
96pub fn validate_wav(audio: &Bytes) -> Result<(), DictationError> {
97    validate_wav_bytes(audio)
98}
99
100fn validate_wav_bytes(audio: &[u8]) -> Result<(), DictationError> {
101    if audio.is_empty() {
102        return Err(DictationError::InvalidAudio(
103            "audio upload must not be empty",
104        ));
105    }
106    if audio.len() > MAX_AUDIO_BYTES {
107        return Err(DictationError::InvalidAudio("audio upload is too large"));
108    }
109    if audio.len() < 12 || &audio[..4] != b"RIFF" || &audio[8..12] != b"WAVE" {
110        return Err(DictationError::InvalidAudio("expected a RIFF/WAVE file"));
111    }
112    let declared_size = u32::from_le_bytes(audio[4..8].try_into().unwrap()) as usize;
113    if declared_size != audio.len().saturating_sub(8) {
114        return Err(DictationError::InvalidAudio(
115            "RIFF length does not match the upload",
116        ));
117    }
118
119    let mut cursor = 12_usize;
120    let mut fmt_seen = false;
121    let mut data_bytes = 0_u64;
122    let mut data_seen = false;
123    while cursor < audio.len() {
124        if audio.len() - cursor < 8 {
125            return Err(DictationError::InvalidAudio("truncated WAV chunk header"));
126        }
127        let id = &audio[cursor..cursor + 4];
128        let chunk_size = u32::from_le_bytes(
129            audio[cursor + 4..cursor + 8]
130                .try_into()
131                .expect("WAV chunk size is four bytes"),
132        ) as usize;
133        let data_start = cursor + 8;
134        let data_end = data_start
135            .checked_add(chunk_size)
136            .ok_or(DictationError::InvalidAudio("WAV chunk length overflows"))?;
137        let padded_end = data_end
138            .checked_add(chunk_size & 1)
139            .ok_or(DictationError::InvalidAudio("WAV chunk padding overflows"))?;
140        if data_end > audio.len() || padded_end > audio.len() {
141            return Err(DictationError::InvalidAudio("truncated WAV chunk data"));
142        }
143        match id {
144            b"fmt " => {
145                if fmt_seen || chunk_size < 16 {
146                    return Err(DictationError::InvalidAudio("invalid WAV format chunk"));
147                }
148                fmt_seen = true;
149                let format =
150                    u16::from_le_bytes(audio[data_start..data_start + 2].try_into().unwrap());
151                let channels =
152                    u16::from_le_bytes(audio[data_start + 2..data_start + 4].try_into().unwrap());
153                let rate =
154                    u32::from_le_bytes(audio[data_start + 4..data_start + 8].try_into().unwrap());
155                let bytes_per_second =
156                    u32::from_le_bytes(audio[data_start + 8..data_start + 12].try_into().unwrap());
157                let alignment =
158                    u16::from_le_bytes(audio[data_start + 12..data_start + 14].try_into().unwrap());
159                let bits =
160                    u16::from_le_bytes(audio[data_start + 14..data_start + 16].try_into().unwrap());
161                if format != 1
162                    || channels != 1
163                    || rate != 16_000
164                    || alignment != 2
165                    || bits != 16
166                    || bytes_per_second != 32_000
167                {
168                    return Err(DictationError::InvalidAudio(
169                        "WAV must be mono 16-bit PCM at 16 kHz",
170                    ));
171                }
172            }
173            b"data" => {
174                if data_seen || !chunk_size.is_multiple_of(2) {
175                    return Err(DictationError::InvalidAudio(
176                        "WAV PCM data is not one even-sized sample chunk",
177                    ));
178                }
179                data_seen = true;
180                data_bytes = data_bytes
181                    .checked_add(chunk_size as u64)
182                    .ok_or(DictationError::InvalidAudio("WAV sample length overflows"))?;
183            }
184            _ => {}
185        }
186        cursor = padded_end;
187    }
188    if !fmt_seen {
189        return Err(DictationError::InvalidAudio("WAV format chunk is missing"));
190    }
191    if data_bytes == 0 {
192        return Err(DictationError::InvalidAudio("WAV sample data is missing"));
193    }
194    if data_bytes > MAX_PCM_DATA_BYTES {
195        return Err(DictationError::InvalidAudio(
196            "audio duration exceeds 600 seconds",
197        ));
198    }
199    Ok(())
200}
201
202/// Execute one request. This future is supervised by the controller's
203/// `JoinSet`; no provider task is detached when the HTTP client disconnects or
204/// the daemon shuts down.
205pub async fn execute(
206    request: DictationRequest,
207    auth_paths: Option<Vec<PathBuf>>,
208    shutdown: CancellationToken,
209) {
210    let DictationRequest {
211        session_id: _,
212        operation,
213        cancel,
214        mut reply,
215    } = request;
216    let result = tokio::select! {
217        biased;
218        _ = reply.closed() => return,
219        _ = shutdown.cancelled() => Err(DictationError::Cancelled),
220        _ = cancel.cancelled() => Err(DictationError::Cancelled),
221        result = tokio::time::timeout(DICTATION_TIMEOUT, execute_operation(operation, auth_paths, cancel.clone())) => {
222            result.unwrap_or(Err(DictationError::TimedOut))
223        },
224    };
225    let _ = reply.send(result);
226}
227
228async fn execute_operation(
229    operation: DictationOperation,
230    auth_paths: Option<Vec<PathBuf>>,
231    cancel: CancellationToken,
232) -> Result<DictationResponse, DictationError> {
233    let paths = auth_paths.ok_or(DictationError::SessionNotFound)?;
234    let audio = match &operation {
235        DictationOperation::Transcribe(audio) => Some(audio.clone()),
236        DictationOperation::Availability => None,
237    };
238    let probe = tokio::task::spawn_blocking(move || {
239        if let Some(audio) = audio {
240            validate_wav(&audio)?;
241        }
242        Ok::<_, DictationError>(available_auth(paths))
243    });
244    let auth_path = tokio::select! {
245        biased;
246        _ = cancel.cancelled() => return Err(DictationError::Cancelled),
247        result = probe => result.map_err(|error| {
248            tracing::warn!(%error, "Codex dictation credential probe task failed");
249            DictationError::CredentialProbe
250        })??,
251    };
252
253    match operation {
254        DictationOperation::Availability => Ok(DictationResponse::Availability {
255            available: auth_path.is_some(),
256            reason: auth_path
257                .is_none()
258                .then(|| "no Codex subscription credentials are available".to_owned()),
259        }),
260        DictationOperation::Transcribe(audio) => {
261            let auth_path = auth_path.ok_or(DictationError::CredentialsUnavailable)?;
262            let mut request = TranscribeRequest::new(audio, "audio.wav", "audio/wav");
263            request.cancel = cancel.clone();
264            request.timeout = DICTATION_TIMEOUT;
265            let transcription = CodexClient::with_auth_path(auth_path)
266                .transcribe(request)
267                .await
268                .map_err(|error| {
269                    let message = error.to_string();
270                    if message.to_ascii_lowercase().contains("timed out") {
271                        DictationError::TimedOut
272                    } else if cancel.is_cancelled() {
273                        DictationError::Cancelled
274                    } else {
275                        DictationError::Provider(message)
276                    }
277                })?;
278            Ok(DictationResponse::Transcript {
279                text: transcription.text,
280            })
281        }
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288
289    fn wav(sample_bytes: usize) -> Bytes {
290        let padded = sample_bytes + (sample_bytes & 1);
291        let riff_size = 36 + padded;
292        let mut bytes = Vec::with_capacity(8 + riff_size);
293        bytes.extend_from_slice(b"RIFF");
294        bytes.extend_from_slice(&(riff_size as u32).to_le_bytes());
295        bytes.extend_from_slice(b"WAVEfmt ");
296        bytes.extend_from_slice(&16_u32.to_le_bytes());
297        bytes.extend_from_slice(&1_u16.to_le_bytes());
298        bytes.extend_from_slice(&1_u16.to_le_bytes());
299        bytes.extend_from_slice(&16_000_u32.to_le_bytes());
300        bytes.extend_from_slice(&32_000_u32.to_le_bytes());
301        bytes.extend_from_slice(&2_u16.to_le_bytes());
302        bytes.extend_from_slice(&16_u16.to_le_bytes());
303        bytes.extend_from_slice(b"data");
304        bytes.extend_from_slice(&(sample_bytes as u32).to_le_bytes());
305        bytes.resize(bytes.len() + sample_bytes, 0);
306        if sample_bytes & 1 != 0 {
307            bytes.push(0);
308        }
309        Bytes::from(bytes)
310    }
311
312    #[test]
313    fn accepts_wav_with_unknown_padded_chunk() {
314        let mut audio = wav(64).to_vec();
315        audio.splice(12..12, b"JUNK\x01\x00\x00\x00x\x00".iter().copied());
316        let size = audio.len() - 8;
317        audio[4..8].copy_from_slice(&(size as u32).to_le_bytes());
318        assert!(validate_wav(&Bytes::from(audio)).is_ok());
319    }
320
321    #[test]
322    fn rejects_truncated_and_wrong_format_audio() {
323        assert!(matches!(
324            validate_wav(&Bytes::from_static(b"RIFF")),
325            Err(DictationError::InvalidAudio(_))
326        ));
327        let mut audio = wav(64).to_vec();
328        audio[22..24].copy_from_slice(&2_u16.to_le_bytes());
329        assert!(matches!(
330            validate_wav(&Bytes::from(audio)),
331            Err(DictationError::InvalidAudio(_))
332        ));
333    }
334
335    #[test]
336    fn rejects_data_longer_than_six_hundred_seconds() {
337        let audio = wav((MAX_PCM_DATA_BYTES + 2) as usize);
338        assert!(matches!(
339            validate_wav(&audio),
340            Err(DictationError::InvalidAudio(_))
341        ));
342    }
343
344    #[tokio::test]
345    async fn availability_and_transcription_without_credentials_do_not_call_provider() {
346        for operation in [
347            DictationOperation::Availability,
348            DictationOperation::Transcribe(wav(96_000)),
349        ] {
350            let availability = matches!(operation, DictationOperation::Availability);
351            let (reply, answer) = oneshot::channel();
352            execute(
353                DictationRequest {
354                    session_id: "session".into(),
355                    operation,
356                    cancel: CancellationToken::new(),
357                    reply,
358                },
359                Some(vec![]),
360                CancellationToken::new(),
361            )
362            .await;
363            let result = answer.await.unwrap();
364            if availability {
365                assert!(matches!(
366                    result,
367                    Ok(DictationResponse::Availability {
368                        available: false,
369                        reason: Some(_)
370                    })
371                ));
372            } else {
373                assert!(matches!(
374                    result,
375                    Err(DictationError::CredentialsUnavailable)
376                ));
377            }
378        }
379    }
380
381    #[tokio::test]
382    async fn shutdown_and_request_cancellation_preempt_credential_work() {
383        for shutdown_cancelled in [false, true] {
384            let cancel = CancellationToken::new();
385            let shutdown = CancellationToken::new();
386            if shutdown_cancelled {
387                shutdown.cancel();
388            } else {
389                cancel.cancel();
390            }
391            let (reply, answer) = oneshot::channel();
392            execute(
393                DictationRequest {
394                    session_id: "session".into(),
395                    operation: DictationOperation::Availability,
396                    cancel,
397                    reply,
398                },
399                None,
400                shutdown,
401            )
402            .await;
403            assert!(matches!(
404                answer.await.unwrap(),
405                Err(DictationError::Cancelled)
406            ));
407        }
408    }
409
410    #[test]
411    fn duplicate_or_odd_data_chunks_and_wrong_riff_lengths_are_rejected() {
412        let mut duplicate = wav(2).to_vec();
413        duplicate.extend_from_slice(b"data\x02\x00\x00\x00\x00\x00");
414        let length = duplicate.len() as u32 - 8;
415        duplicate[4..8].copy_from_slice(&length.to_le_bytes());
416        assert!(validate_wav(&Bytes::from(duplicate)).is_err());
417        assert!(validate_wav(&wav(3)).is_err());
418        let mut wrong_length = wav(2).to_vec();
419        wrong_length[4..8].copy_from_slice(&0_u32.to_le_bytes());
420        assert!(validate_wav(&Bytes::from(wrong_length)).is_err());
421        assert!(validate_wav(&wav(MAX_PCM_DATA_BYTES as usize)).is_ok());
422    }
423}
mj_controller/dictation.rs

mj_controller/
dictation.rs
