tenant.openagents/omega

1use anyhow::{Context as _, Result};
2use collections::HashMap;
3use cpal::{
4    DeviceDescription, DeviceId, default_host,
5    traits::{DeviceTrait, HostTrait},
6};
7use gpui::{App, AsyncApp, BorrowAppContext, Global};
8
9pub(super) use cpal::Sample;
10
11use rodio::{Decoder, DeviceSinkBuilder, MixerDeviceSink, Source, mixer::Mixer, source::Buffered};
12use settings::Settings;
13use std::io::Cursor;
14use util::ResultExt;
15
16mod echo_canceller;
17use echo_canceller::EchoCanceller;
18mod rodio_ext;
19pub use crate::audio_settings::AudioSettings;
20pub use rodio_ext::RodioExt;
21
22use crate::Sound;
23
24use super::{CHANNEL_COUNT, SAMPLE_RATE};
25pub const BUFFER_SIZE: usize = // echo canceller and livekit want 10ms of audio
26    (SAMPLE_RATE.get() as usize / 100) * CHANNEL_COUNT.get() as usize;
27
28pub fn init(_cx: &mut App) {}
29
30// TODO(jk): this is currently cached only once - we should observe and react instead
31pub fn ensure_devices_initialized(cx: &mut App) {
32    if cx.has_global::<AvailableAudioDevices>() {
33        return;
34    }
35    cx.default_global::<AvailableAudioDevices>();
36    let task = cx
37        .background_executor()
38        .spawn(async move { get_available_audio_devices() });
39    cx.spawn(async move |cx: &mut AsyncApp| {
40        let devices = task.await;
41        cx.update(|cx| cx.set_global(AvailableAudioDevices(devices)));
42        cx.refresh();
43    })
44    .detach();
45}
46
47#[derive(Default)]
48pub struct Audio {
49    output: Option<(MixerDeviceSink, Mixer)>,
50    pub echo_canceller: EchoCanceller,
51    source_cache: HashMap<Sound, Buffered<Decoder<Cursor<Vec<u8>>>>>,
52}
53
54impl Global for Audio {}
55
56impl Audio {
57    fn ensure_output_exists(&mut self, output_audio_device: Option<DeviceId>) -> Result<&Mixer> {
58        #[cfg(debug_assertions)]
59        log::warn!(
60            "Audio does not sound correct without optimizations. Use a release build to debug audio issues"
61        );
62
63        if self.output.is_none() {
64            let (output_handle, output_mixer) =
65                open_output_stream(output_audio_device, self.echo_canceller.clone())?;
66            self.output = Some((output_handle, output_mixer));
67        }
68
69        Ok(self
70            .output
71            .as_ref()
72            .map(|(_, mixer)| mixer)
73            .expect("we only get here if opening the outputstream succeeded"))
74    }
75
76    pub fn play_sound(sound: Sound, cx: &mut App) {
77        let output_audio_device = AudioSettings::get_global(cx).output_audio_device.clone();
78        cx.update_default_global(|this: &mut Self, cx| {
79            let source = this.sound_source(sound, cx).log_err()?;
80            let output_mixer = this
81                .ensure_output_exists(output_audio_device)
82                .context("Could not get output mixer")
83                .log_err()?;
84
85            output_mixer.add(source);
86            Some(())
87        });
88    }
89
90    pub fn end_call(cx: &mut App) {
91        cx.update_default_global(|this: &mut Self, _cx| {
92            this.output.take();
93        });
94    }
95
96    fn sound_source(&mut self, sound: Sound, cx: &App) -> Result<impl Source + use<>> {
97        if let Some(wav) = self.source_cache.get(&sound) {
98            return Ok(wav.clone());
99        }
100
101        let path = format!("sounds/{}.wav", sound.file());
102        let bytes = cx
103            .asset_source()
104            .load(&path)?
105            .map(anyhow::Ok)
106            .with_context(|| format!("No asset available for path {path}"))??
107            .into_owned();
108        let cursor = Cursor::new(bytes);
109        let source = Decoder::new(cursor)?.buffered();
110
111        self.source_cache.insert(sound, source.clone());
112
113        Ok(source)
114    }
115}
116
117pub fn open_input_stream(
118    device_id: Option<DeviceId>,
119) -> anyhow::Result<rodio::microphone::Microphone> {
120    let builder = rodio::microphone::MicrophoneBuilder::new();
121    let builder = if let Some(id) = device_id {
122        // TODO(jk): upstream patch
123        // if let Some(input_device) = default_host().device_by_id(id) {
124        //     builder.device(input_device);
125        // }
126        let mut found = None;
127        for input in rodio::microphone::available_inputs()? {
128            if input.clone().into_inner().id()? == id {
129                found = Some(builder.device(input));
130                break;
131            }
132        }
133        found.unwrap_or_else(|| builder.default_device())?
134    } else {
135        builder.default_device()?
136    };
137    let stream = builder
138        .default_config()?
139        .prefer_sample_rates([
140            SAMPLE_RATE,
141            SAMPLE_RATE.saturating_mul(rodio::nz!(2)),
142            SAMPLE_RATE.saturating_mul(rodio::nz!(3)),
143            SAMPLE_RATE.saturating_mul(rodio::nz!(4)),
144        ])
145        .prefer_channel_counts([rodio::nz!(1), rodio::nz!(2), rodio::nz!(3), rodio::nz!(4)])
146        .prefer_buffer_sizes(512..)
147        .open_stream()?;
148    log::info!("Opened microphone: {:?}", stream.config());
149    Ok(stream)
150}
151
152pub fn resolve_device(device_id: Option<&DeviceId>, input: bool) -> anyhow::Result<cpal::Device> {
153    if let Some(id) = device_id {
154        if let Some(device) = default_host().device_by_id(id) {
155            return Ok(device);
156        }
157        log::warn!("Selected audio device not found, falling back to default");
158    }
159    if input {
160        default_host()
161            .default_input_device()
162            .context("no audio input device available")
163    } else {
164        default_host()
165            .default_output_device()
166            .context("no audio output device available")
167    }
168}
169
170pub fn open_test_output(device_id: Option<DeviceId>) -> anyhow::Result<MixerDeviceSink> {
171    let device = resolve_device(device_id.as_ref(), false)?;
172    DeviceSinkBuilder::from_device(device)?
173        .open_stream()
174        .context("Could not open output stream")
175}
176
177pub fn open_output_stream(
178    device_id: Option<DeviceId>,
179    mut echo_canceller: EchoCanceller,
180) -> anyhow::Result<(MixerDeviceSink, Mixer)> {
181    let device = resolve_device(device_id.as_ref(), false)?;
182    let mut output_handle = DeviceSinkBuilder::from_device(device)?
183        .open_stream()
184        .context("Could not open output stream")?;
185    output_handle.log_on_drop(false);
186    log::info!("Output stream: {:?}", output_handle);
187
188    let (output_mixer, source) = rodio::mixer::mixer(CHANNEL_COUNT, SAMPLE_RATE);
189    // otherwise the mixer ends as it's empty
190    output_mixer.add(rodio::source::Zero::new(CHANNEL_COUNT, SAMPLE_RATE));
191    let echo_cancelling_source = source // apply echo cancellation just before output
192        .inspect_buffer::<BUFFER_SIZE, _>(move |buffer| {
193            let mut buf: [i16; _] = buffer.map(|s| s.to_sample());
194            echo_canceller.process_reverse_stream(&mut buf)
195        });
196    output_handle.mixer().add(echo_cancelling_source);
197
198    Ok((output_handle, output_mixer))
199}
200
201#[derive(Clone, Debug)]
202pub struct AudioDeviceInfo {
203    pub id: DeviceId,
204    pub desc: DeviceDescription,
205}
206
207impl AudioDeviceInfo {
208    pub fn matches_input(&self, is_input: bool) -> bool {
209        if is_input {
210            self.desc.supports_input()
211        } else {
212            self.desc.supports_output()
213        }
214    }
215
216    pub fn matches(&self, id: &DeviceId, is_input: bool) -> bool {
217        &self.id == id && self.matches_input(is_input)
218    }
219}
220
221impl std::fmt::Display for AudioDeviceInfo {
222    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
223        write!(f, "{} ({})", self.desc.name(), self.id)
224    }
225}
226
227fn get_available_audio_devices() -> Vec<AudioDeviceInfo> {
228    let Some(devices) = default_host().devices().ok() else {
229        return Vec::new();
230    };
231    devices
232        .filter_map(|device| {
233            let id = device.id().ok()?;
234            let desc = device.description().ok()?;
235            Some(AudioDeviceInfo { id, desc })
236        })
237        .collect()
238}
239
240#[derive(Default, Clone, Debug)]
241pub struct AvailableAudioDevices(pub Vec<AudioDeviceInfo>);
242
243impl Global for AvailableAudioDevices {}
244