Merge remote-tracking branch 'origin/main' into pulse-reaction

This commit is contained in:
Thomas Petersen
2026-05-19 20:21:59 -04:00
68 changed files with 7207 additions and 2018 deletions

3
Cargo.lock generated
View File

@@ -3920,6 +3920,7 @@ dependencies = [
"serde",
"serde_json",
"sha2 0.11.0",
"sprout-core",
"sprout-persona",
"sprout-sdk",
"thiserror 2.0.18",
@@ -3934,12 +3935,14 @@ version = "0.1.0"
dependencies = [
"chrono",
"hex",
"hmac 0.13.0",
"nostr",
"percent-encoding",
"rand 0.10.1",
"schemars",
"serde",
"serde_json",
"sha2 0.11.0",
"subtle",
"thiserror 2.0.18",
"url",

View File

@@ -225,6 +225,7 @@ sprout channels add-member --channel "$CHANNEL" --pubkey "$AGENT_PUBKEY" --role
export SPROUT_PRIVATE_KEY="$AGENT_SK"
export SPROUT_RELAY_URL=ws://localhost:3000 # match step 3 (e.g. ws://localhost:3030 if overridden)
export SPROUT_ACP_RESPOND_TO=anyone # default is owner-only; opens the gate for testing
export SPROUT_ACP_MEMORY=true # opt in to NIP-AE core-memory prompt injection
export SPROUT_ACP_MCP_COMMAND="$PWD/target/release/sprout-mcp-server" # explicit path beats $PATH
export GOOSE_MODE=auto # must be 'auto' or goose hangs on prompts

View File

@@ -328,6 +328,30 @@ pub struct CliArgs {
#[arg(long, env = "SPROUT_ACP_NO_TYPING")]
pub no_typing: bool,
/// Enable NIP-AE agent core memory injection.
///
/// Memory injection is off by default for now. When enabled, the harness
/// fetches the agent's per-session core engram and renders it as an
/// `[Agent Memory — core]` prompt section (or renders the onboarding nudge
/// when the relay confirms no core engram exists). The `sprout mem` CLI
/// and the relay's acceptance of kind:30174 engrams are unaffected — this
/// flag controls prompt-time injection in the ACP harness only.
#[arg(long, env = "SPROUT_ACP_MEMORY", conflicts_with = "no_memory")]
pub memory: bool,
/// Disable NIP-AE agent core memory injection.
///
/// Deprecated compatibility alias for the previous default-on behavior.
/// The flag/env var is still accepted, but memory injection is already off
/// unless `--memory` / `SPROUT_ACP_MEMORY=true` is provided.
#[arg(
long,
env = "SPROUT_ACP_NO_MEMORY",
conflicts_with = "memory",
hide = true
)]
pub no_memory: bool,
/// Desired LLM model ID. Applied to every new ACP session after creation.
/// Use `sprout-acp models` to discover available model IDs.
#[arg(long, env = "SPROUT_ACP_MODEL")]
@@ -416,6 +440,11 @@ pub struct Config {
pub max_turns_per_session: u32,
pub presence_enabled: bool,
pub typing_enabled: bool,
/// Whether NIP-AE agent core memory injection is enabled. When false,
/// the harness skips the per-session core engram fetch and renders no
/// `[Agent Memory — core]` section. Mirrors the `--memory` /
/// `SPROUT_ACP_MEMORY` opt-in.
pub memory_enabled: bool,
/// Desired LLM model ID. Applied after every `session_new_full()`.
pub model: Option<String>,
/// Permission mode to apply after session creation. `Default` = skip.
@@ -761,6 +790,7 @@ impl Config {
max_turns_per_session: args.max_turns_per_session,
presence_enabled: !args.no_presence,
typing_enabled: !args.no_typing,
memory_enabled: args.memory && !args.no_memory,
model,
permission_mode: args.permission_mode,
respond_to: args.respond_to,
@@ -782,7 +812,7 @@ impl Config {
other => format!("respond_to={other}"),
};
format!(
"relay={} pubkey={} agent_cmd={} {} mcp_cmd={} idle_timeout={}s max_turn={}s agents={} heartbeat={}s subscribe={:?} dedup={:?} meh={:?} ignore_self={} context_limit={} max_turns_per_session={} presence={} typing={} model={} permission_mode={} {}",
"relay={} pubkey={} agent_cmd={} {} mcp_cmd={} idle_timeout={}s max_turn={}s agents={} heartbeat={}s subscribe={:?} dedup={:?} meh={:?} ignore_self={} context_limit={} max_turns_per_session={} presence={} typing={} memory={} model={} permission_mode={} {}",
self.relay_url,
self.keys.public_key().to_hex(),
self.agent_command,
@@ -800,6 +830,7 @@ impl Config {
self.max_turns_per_session,
self.presence_enabled,
self.typing_enabled,
self.memory_enabled,
self.model.as_deref().unwrap_or("(agent default)"),
self.permission_mode,
respond_to_detail,
@@ -1122,6 +1153,7 @@ mod tests {
max_turns_per_session: 0,
presence_enabled: true,
typing_enabled: true,
memory_enabled: false,
model: None,
permission_mode: PermissionMode::BypassPermissions,
respond_to: RespondTo::Anyone,
@@ -1705,6 +1737,38 @@ channels = "ALL"
);
}
// ── memory toggle ───────────────────────────────────────────────────────
#[test]
fn test_memory_enabled_default_false() {
let config = test_config(SubscribeMode::Mentions);
assert!(
!config.memory_enabled,
"memory_enabled should default to false"
);
}
#[test]
fn test_summary_includes_memory_disabled() {
let config = test_config(SubscribeMode::Mentions);
let s = config.summary();
assert!(
s.contains("memory=false"),
"summary should include memory=false by default, got: {s}"
);
}
#[test]
fn test_summary_reflects_memory_enabled() {
let mut config = test_config(SubscribeMode::Mentions);
config.memory_enabled = true;
let s = config.summary();
assert!(
s.contains("memory=true"),
"summary should include memory=true when enabled, got: {s}"
);
}
// ── permission mode ─────────────────────────────────────────────────────
#[test]

View File

@@ -0,0 +1,248 @@
//! Fetch the agent's NIP-AE `core` engram at session creation and render it
//! into a prompt section.
//!
//! Scope per Tyler's spec:
//! - Fire one synchronous query for the core head when a *new* session is born.
//! - If a body is found, emit `[Agent Memory — core]\n<profile>`.
//! - If no body is found, emit an onboarding nudge so the agent learns how
//! to set its own core.
//! - On any *error* (transport, parse), log and emit nothing. We must not
//! mistake a relay outage for "no core" — that would invite the agent to
//! overwrite real, just-unreachable memory with a fresh profile.
//! - Either way, session creation is never blocked.
use nostr::{Event, Keys, PublicKey};
use sprout_core::engram::{conversation_key, d_tag, select_head, validate_and_decrypt, Body};
use sprout_core::kind::KIND_AGENT_ENGRAM;
use crate::relay::RestClient;
/// Section header rendered into the prompt.
const SECTION_LABEL: &str = "Agent Memory — core";
/// Onboarding nudge for new agents with no core yet.
///
/// Wording is from Tyler's brief: "No core memory found. Use `sprout mem`
/// to create a core memory. Ask your user about yourself."
pub const ONBOARDING_NUDGE: &str = "No core memory found. \
Use `sprout mem set core \"…\"` to create one (it will hold your identity, \
rules, and goals across sessions). Ask your user about yourself.";
/// Build the rendered prompt section for the agent's core.
///
/// Returns:
/// - `Some(profile_section)` when a valid core exists,
/// - `Some(nudge_section)` when the relay confirmed absence,
/// - `None` when the fetch failed (transport, parse, decrypt) — the caller
/// should inject no section in that case so the agent doesn't conclude
/// memory is empty.
pub async fn build_core_section(
rest: &RestClient,
agent_keys: &Keys,
owner: &PublicKey,
) -> Option<String> {
match fetch_core_body(rest, agent_keys, owner).await {
Ok(Some(profile)) => Some(format!("[{SECTION_LABEL}]\n{profile}")),
Ok(None) => Some(format!("[{SECTION_LABEL}]\n{ONBOARDING_NUDGE}")),
Err(reason) => {
tracing::warn!(
target: "engram::core",
"core fetch failed: {reason} — emitting no section to avoid \
confusing a relay outage with an absent core"
);
None
}
}
}
/// Query the relay for the core head and decode it. Returns:
/// - `Ok(Some(profile))` if a valid core body was found,
/// - `Ok(None)` only if the relay confirmed absence (empty result set),
/// - `Err(reason)` if the relay returned candidates we could not parse,
/// verify, or decrypt — those are NOT treated as absence (would let an
/// unreadable but real core be silently overwritten by the onboarding nudge),
/// - `Err` for transport / parse errors.
async fn fetch_core_body(
rest: &RestClient,
agent_keys: &Keys,
owner: &PublicKey,
) -> Result<Option<String>, String> {
let k_c = conversation_key(agent_keys.secret_key(), owner);
let d = d_tag(&k_c, sprout_core::engram::CORE_SLUG);
let filter = nostr::Filter::new()
.kind(nostr::Kind::Custom(KIND_AGENT_ENGRAM as u16))
.author(agent_keys.public_key())
.custom_tag(nostr::SingleLetterTag::lowercase(nostr::Alphabet::D), [d])
.custom_tag(
nostr::SingleLetterTag::lowercase(nostr::Alphabet::P),
[owner.to_hex()],
)
.limit(16);
let value = rest
.query(&[filter])
.await
.map_err(|e| format!("relay query failed: {e}"))?;
let arr = value
.as_array()
.ok_or_else(|| "relay query returned non-array".to_string())?;
decode_core_body(arr, agent_keys, owner)
}
/// Pure decoder: given the relay's JSON array, decide whether we have a
/// readable core, confirmed absence, or an ambiguous unreadable-state.
///
/// - Empty array → `Ok(None)` (confirmed absence; caller renders the nudge).
/// - At least one event decrypts → use the winning head's body.
/// * Body::Core → `Ok(Some(profile))`
/// * Body::Tombstone or unexpected shape → `Ok(None)` (treat as absent).
/// - Non-empty array but nothing decrypts → `Err` (fail closed; caller
/// emits no section, so the agent does not assume memory is empty and
/// try to overwrite a real-but-unreadable core).
fn decode_core_body(
arr: &[serde_json::Value],
agent_keys: &Keys,
owner: &PublicKey,
) -> Result<Option<String>, String> {
if arr.is_empty() {
return Ok(None);
}
let mut valid_with_body: Vec<(Event, Body)> = Vec::with_capacity(arr.len());
let mut candidates_seen = 0usize;
let mut last_decrypt_err: Option<String> = None;
for ev_json in arr {
let event: Event = match serde_json::from_value(ev_json.clone()) {
Ok(e) => e,
Err(_) => continue,
};
if event.verify().is_err() {
continue;
}
candidates_seen += 1;
match validate_and_decrypt(
&event,
&agent_keys.public_key(),
owner,
agent_keys.secret_key(),
owner,
) {
Ok(body) => valid_with_body.push((event, body)),
Err(e) => {
last_decrypt_err = Some(e.to_string());
continue;
}
}
}
if valid_with_body.is_empty() {
if candidates_seen > 0 {
return Err(format!(
"{candidates_seen} core candidate(s) returned but none decryptable (last error: {})",
last_decrypt_err.as_deref().unwrap_or("unknown")
));
}
return Err(
"relay returned core candidate(s) that could not be parsed or verified".to_string(),
);
}
let events: Vec<Event> = valid_with_body.iter().map(|(e, _)| e.clone()).collect();
// `select_head` returns `None` only on an empty iterator, which we
// ruled out above.
let Some(head) = select_head(events) else {
return Ok(None);
};
let head_id = head.id;
let body = valid_with_body
.into_iter()
.find(|(e, _)| e.id == head_id)
.map(|(_, b)| b);
match body {
Some(Body::Core { profile }) => Ok(Some(profile)),
// A tombstone or unexpectedly-shaped head means "no usable core."
_ => Ok(None),
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use sprout_core::engram::{build_event, Body};
/// Empty array → confirmed absence → Ok(None), so the caller emits the
/// onboarding nudge. This is the only path that maps to "no core."
#[test]
fn decode_empty_array_is_confirmed_absence() {
let agent = Keys::generate();
let owner = Keys::generate();
let out = decode_core_body(&[], &agent, &owner.public_key()).unwrap();
assert_eq!(out, None);
}
/// Happy path: a real, decryptable core event yields the profile.
#[test]
fn decode_valid_core_returns_profile() {
let agent = Keys::generate();
let owner = Keys::generate();
let body = Body::Core {
profile: "I am Sami.".to_string(),
};
let ev = build_event(&agent, &owner.public_key(), &body, 1_700_000_000).unwrap();
let arr = vec![serde_json::to_value(&ev).unwrap()];
let out = decode_core_body(&arr, &agent, &owner.public_key()).unwrap();
assert_eq!(out.as_deref(), Some("I am Sami."));
}
/// Regression: when the relay returns a kind:30174 event addressed to
/// this agent that we cannot decrypt (here: encrypted to a *different*
/// owner's key, so the MAC fails for this agent↔owner pair), we MUST
/// return Err and NOT Ok(None). Returning Ok(None) would cause the
/// harness to emit the onboarding nudge, inviting the agent to overwrite
/// a real-but-unreadable core.
#[test]
fn decode_undecryptable_candidate_is_err_not_absent() {
let agent = Keys::generate();
let owner = Keys::generate();
let wrong_owner = Keys::generate();
// Build an engram encrypted to wrong_owner (not owner). It will pass
// sig verification but fail MAC/decrypt for the agent↔owner pair.
let body = Body::Core {
profile: "secret".to_string(),
};
let ev = build_event(&agent, &wrong_owner.public_key(), &body, 1_700_000_000).unwrap();
let arr = vec![serde_json::to_value(&ev).unwrap()];
let result = decode_core_body(&arr, &agent, &owner.public_key());
assert!(result.is_err(), "expected Err, got: {result:?}");
let msg = result.unwrap_err();
assert!(msg.contains("decryptable"), "got: {msg}");
}
/// An unexpectedly-shaped head (here: a Memory body in what was supposed
/// to be the core slot) is a legitimate, decryptable "no usable core" —
/// Ok(None). Real `rm core` is refused at the CLI, so this is a defensive
/// branch for malformed data on the wire.
#[test]
fn decode_non_core_body_is_absent() {
let agent = Keys::generate();
let owner = Keys::generate();
let body = Body::Memory {
slug: "mem/x".to_string(),
value: None,
};
let ev = build_event(&agent, &owner.public_key(), &body, 1_700_000_000).unwrap();
let arr = vec![serde_json::to_value(&ev).unwrap()];
let out = decode_core_body(&arr, &agent, &owner.public_key()).unwrap();
assert_eq!(out, None);
}
/// Non-empty array with only garbage entries (not even parseable as
/// events) is also treated as a fetch error, not absence.
#[test]
fn decode_unparseable_candidates_is_err() {
let agent = Keys::generate();
let owner = Keys::generate();
let arr = vec![json!({"not": "an event"}), json!("garbage")];
let result = decode_core_body(&arr, &agent, &owner.public_key());
assert!(result.is_err(), "expected Err, got: {result:?}");
}
}

View File

@@ -2,6 +2,7 @@
mod acp;
mod config;
mod engram_fetch;
mod filter;
mod observer;
mod pool;
@@ -967,7 +968,7 @@ async fn tokio_main() -> Result<()> {
_ => {} // anyone/nobody don't depend on owner
}
}
let owner_cache = OwnerCache::new(startup_owner);
let owner_cache = OwnerCache::new(startup_owner.clone());
let mut relay_observer_control_rx = None;
let mut relay_observer_publisher_task = None;
@@ -1084,8 +1085,20 @@ async fn tokio_main() -> Result<()> {
context_message_limit: config.context_message_limit,
max_turns_per_session: config.max_turns_per_session,
permission_mode: config.permission_mode,
agent_keys: config.keys.clone(),
agent_owner_pubkey: startup_owner
.as_deref()
.and_then(|hex| nostr::PublicKey::from_hex(hex).ok()),
memory_enabled: config.memory_enabled,
});
if !config.memory_enabled {
tracing::info!(
target: "engram::core",
"NIP-AE core memory injection disabled by default (enable with --memory / SPROUT_ACP_MEMORY)"
);
}
// ── Step 6: Heartbeat timer ───────────────────────────────────────────────
let mut heartbeat = if config.heartbeat_interval_secs > 0 {
let interval = Duration::from_secs(config.heartbeat_interval_secs);
@@ -2748,6 +2761,7 @@ mod build_mcp_servers_tests {
max_turns_per_session: 0,
presence_enabled: true,
typing_enabled: true,
memory_enabled: false,
model: None,
permission_mode: config::PermissionMode::BypassPermissions,
respond_to: config::RespondTo::Anyone,

View File

@@ -82,6 +82,9 @@ pub struct SessionState {
pub turn_counts: HashMap<Uuid, u32>,
/// Turn counter for the heartbeat session.
pub heartbeat_turn_count: u32,
/// channel_id → rendered NIP-AE core prompt section, populated once at
/// session creation per Tyler's spec (no mid-session refresh).
pub core_sections: HashMap<Uuid, String>,
}
impl SessionState {
@@ -102,6 +105,7 @@ impl SessionState {
/// Returns `true` if the channel had an active session.
pub fn invalidate_channel(&mut self, channel_id: &Uuid) -> bool {
self.turn_counts.remove(channel_id);
self.core_sections.remove(channel_id);
self.sessions.remove(channel_id).is_some()
}
@@ -111,6 +115,7 @@ impl SessionState {
self.turn_counts.clear();
self.heartbeat_session = None;
self.heartbeat_turn_count = 0;
self.core_sections.clear();
}
}
@@ -198,6 +203,18 @@ pub struct PromptContext {
pub max_turns_per_session: u32,
/// Permission mode to apply after session creation. `Default` = skip.
pub permission_mode: PermissionMode,
/// Agent identity — used to derive the NIP-AE conversation key at
/// session creation for core injection.
pub agent_keys: nostr::Keys,
/// Owner pubkey (hex), if resolved at startup. When unset, NIP-AE core
/// injection is skipped entirely (no owner = no `(agent, owner)` pair).
pub agent_owner_pubkey: Option<nostr::PublicKey>,
/// Whether NIP-AE agent core memory injection is enabled. When false,
/// the per-session core engram fetch is skipped and `core_sections`
/// remains empty for every channel, so `format_prompt` renders no
/// `[Agent Memory — core]` section. Driven by `--memory` /
/// `SPROUT_ACP_MEMORY`.
pub memory_enabled: bool,
}
// ── AgentPool impl ────────────────────────────────────────────────────────────
@@ -737,6 +754,67 @@ pub async fn run_prompt_task(
}),
);
// ── NIP-AE: fetch core engram on new channel sessions ────────────────
//
// Fire one synchronous fetch + decrypt + render right after the session
// is born; cache the rendered section in `state.core_sections` keyed by
// channel id. Subsequent turns in the same session reuse the cached
// section. `format_prompt` reads it from the per-channel cache.
//
// Failure modes (all fail open — no crash, no block):
// * no owner configured → skip (no NIP-AE namespace exists)
// * confirmed absence → cache the onboarding nudge so the agent
// learns how to bootstrap itself.
// * transport / decrypt / parse error → inject nothing. We never
// mistake "relay slow or broken" for "no core" — that would invite
// the agent to overwrite real, just-unreachable memory.
// * fetch exceeds CORE_FETCH_TIMEOUT → inject nothing, same reason.
//
// Per Tyler's locked spec: NO mid-session refreshes. Re-fetch only
// happens when a session is invalidated and recreated (see
// `SessionState::invalidate_channel`).
//
// Operator opt-in: `--memory` / `SPROUT_ACP_MEMORY` enables the NIP-AE
// injection path. By default we skip the fetch outright and leave
// `state.core_sections` empty, so `format_prompt` renders no core
// section. The `sprout mem` CLI and the relay's acceptance of
// kind:30174 engrams are unaffected.
if is_new_session && ctx.memory_enabled {
if let (PromptSource::Channel(cid), Some(owner_pk)) =
(&source, ctx.agent_owner_pubkey.as_ref())
{
// Bounded — we'd rather start the session with no core hint
// than block prompt formatting on a stalled relay.
const CORE_FETCH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(3);
let fetch = crate::engram_fetch::build_core_section(
&ctx.rest_client,
&ctx.agent_keys,
owner_pk,
);
let section = match tokio::time::timeout(CORE_FETCH_TIMEOUT, fetch).await {
Ok(s) => s,
Err(_) => {
tracing::warn!(
target: "engram::core",
channel = %cid,
timeout_ms = CORE_FETCH_TIMEOUT.as_millis() as u64,
"core fetch timed out — emitting no section"
);
None
}
};
if let Some(rendered) = section {
tracing::info!(
target: "engram::core",
channel = %cid,
section_len = rendered.len(),
"injected NIP-AE core section"
);
agent.state.core_sections.insert(*cid, rendered);
}
}
}
// ── Send initial_message on new channel sessions ──────────────────────
if is_new_session {
@@ -871,9 +949,11 @@ pub async fn run_prompt_task(
let profile_lookup =
fetch_prompt_profile_lookup(b, conversation_context.as_ref(), &ctx.rest_client).await;
let agent_core_section = agent.state.core_sections.get(&b.channel_id).cloned();
crate::queue::format_prompt(
b,
ctx.system_prompt.as_deref(),
agent_core_section.as_deref(),
channel_info.as_ref(),
conversation_context.as_ref(),
profile_lookup.as_ref(),

View File

@@ -952,6 +952,7 @@ fn format_conversation_context(
pub fn format_prompt(
batch: &FlushBatch,
system_prompt: Option<&str>,
agent_core: Option<&str>,
channel_info: Option<&PromptChannelInfo>,
conversation_context: Option<&ConversationContext>,
profile_lookup: Option<&PromptProfileLookup>,
@@ -979,6 +980,11 @@ pub fn format_prompt(
sections.push(format!("[System]\n{sp}"));
}
// 1b. NIP-AE agent core memory (rendered by `engram_fetch::build_core_section`).
if let Some(core) = agent_core {
sections.push(core.to_string());
}
// 2. Context hints (with reply instruction for thread replies).
let triggering_event_id = if thread_tags.root_event_id.is_some() {
Some(last_event.event.id.to_hex())
@@ -1298,7 +1304,7 @@ mod tests {
cancelled_events: vec![],
};
let prompt = format_prompt(&batch, None, None, None, None);
let prompt = format_prompt(&batch, None, None, None, None, None);
// Should contain [Context] section before the event.
assert!(prompt.contains("[Context]"));
@@ -1394,7 +1400,7 @@ mod tests {
cancelled_events: vec![],
};
let prompt = format_prompt(&batch, None, None, None, None);
let prompt = format_prompt(&batch, None, None, None, None, None);
assert!(prompt.contains("[Context]"));
assert!(prompt.contains("[Sprout events — 3 events]"));
@@ -1423,10 +1429,58 @@ mod tests {
cancelled_events: vec![],
};
let prompt = format_prompt(&batch, Some("You are a triage bot."), None, None, None);
let prompt = format_prompt(
&batch,
Some("You are a triage bot."),
None,
None,
None,
None,
);
assert!(prompt.starts_with("[System]\nYou are a triage bot.\n\n[Context]"));
}
// ── Test 11b: agent_core section is injected after [System] ──────────────
#[test]
fn test_format_prompt_with_agent_core() {
let ch = Uuid::new_v4();
let event = make_event("hi");
let batch = FlushBatch {
channel_id: ch,
events: vec![BatchEvent {
event,
prompt_tag: "test".into(),
received_at: Instant::now(),
}],
cancelled_events: vec![],
};
let core = "[Agent Memory — core]\nbe helpful";
let prompt = format_prompt(&batch, Some("sys"), Some(core), None, None, None);
assert!(
prompt.contains("[System]\nsys\n\n[Agent Memory — core]\nbe helpful"),
"expected core block after [System]; got: {prompt}"
);
}
#[test]
fn test_format_prompt_without_system_prompts_core_first() {
let ch = Uuid::new_v4();
let event = make_event("hi");
let batch = FlushBatch {
channel_id: ch,
events: vec![BatchEvent {
event,
prompt_tag: "test".into(),
received_at: Instant::now(),
}],
cancelled_events: vec![],
};
let core = "[Agent Memory — core]\nbe helpful";
let prompt = format_prompt(&batch, None, Some(core), None, None, None);
assert!(prompt.starts_with("[Agent Memory — core]\nbe helpful\n\n[Context]"));
}
// ── Test 12: drop mode discards in-flight channel events ─────────────────
#[test]
@@ -1903,7 +1957,7 @@ mod tests {
channel_type: "stream".into(),
};
let prompt = format_prompt(&batch, None, Some(&ci), None, None);
let prompt = format_prompt(&batch, None, None, Some(&ci), None, None);
assert!(prompt.contains("engineering (#"));
assert!(prompt.contains("Scope: channel"));
}
@@ -1926,7 +1980,7 @@ mod tests {
channel_type: "dm".into(),
};
let prompt = format_prompt(&batch, None, Some(&ci), None, None);
let prompt = format_prompt(&batch, None, None, Some(&ci), None, None);
assert!(prompt.contains("Scope: dm"));
}
@@ -1952,7 +2006,7 @@ mod tests {
cancelled_events: vec![],
};
let prompt = format_prompt(&batch, None, None, None, None);
let prompt = format_prompt(&batch, None, None, None, None, None);
assert!(prompt.contains("Scope: thread"));
assert!(prompt.contains("Thread root: root123"));
}
@@ -1995,7 +2049,7 @@ mod tests {
truncated: true,
};
let prompt = format_prompt(&batch, None, None, Some(&ctx), None);
let prompt = format_prompt(&batch, None, None, None, Some(&ctx), None);
assert!(prompt.contains("[Thread Context (2 of 5 messages, truncated)]"));
assert!(prompt.contains("Let's refactor auth"));
assert!(prompt.contains("Thread context included below"));
@@ -2028,7 +2082,7 @@ mod tests {
truncated: false,
};
let prompt = format_prompt(&batch, None, Some(&ci), Some(&ctx), None);
let prompt = format_prompt(&batch, None, None, Some(&ci), Some(&ctx), None);
assert!(prompt.contains("Scope: dm"));
assert!(prompt.contains("[Conversation Context (1 of 1 messages)]"));
assert!(prompt.contains("Can you deploy?"));
@@ -2080,7 +2134,7 @@ mod tests {
),
]);
let prompt = format_prompt(&batch, None, None, Some(&ctx), Some(&profiles));
let prompt = format_prompt(&batch, None, None, None, Some(&ctx), Some(&profiles));
assert!(prompt.contains("From: Wes (npub:"));
assert!(prompt.contains(
@@ -2175,7 +2229,7 @@ mod tests {
truncated: false,
};
let prompt = format_prompt(&batch, None, Some(&ci), Some(&ctx), None);
let prompt = format_prompt(&batch, None, None, Some(&ci), Some(&ctx), None);
// Scope should be "dm", not "thread".
assert!(
prompt.contains("Scope: dm"),
@@ -2214,7 +2268,7 @@ mod tests {
};
// No context fetched — hints only.
let prompt = format_prompt(&batch, None, Some(&ci), None, None);
let prompt = format_prompt(&batch, None, None, Some(&ci), None, None);
assert!(prompt.contains("Scope: dm"));
assert!(
prompt.contains("get_channel_history()"),
@@ -2241,7 +2295,7 @@ mod tests {
cancelled_events: vec![],
};
let prompt = format_prompt(&batch, None, None, None, None);
let prompt = format_prompt(&batch, None, None, None, None, None);
assert!(
prompt.contains(&format!("Event ID: {event_id}")),
"prompt should contain the event ID"
@@ -2264,7 +2318,7 @@ mod tests {
cancelled_events: vec![],
};
let prompt = format_prompt(&batch, None, None, None, None);
let prompt = format_prompt(&batch, None, None, None, None, None);
assert!(
prompt.contains(&format!("From: {npub} (hex: {hex})")),
"prompt should contain both npub and hex"
@@ -2286,7 +2340,7 @@ mod tests {
cancelled_events: vec![],
};
let prompt = format_prompt(&batch, None, None, None, None);
let prompt = format_prompt(&batch, None, None, None, None, None);
assert!(
prompt.contains("Tags:"),
"tags should always be included, even for stream messages"
@@ -2610,7 +2664,7 @@ mod tests {
cancelled_events: vec![],
};
let prompt = format_prompt(&batch, None, None, None, None);
let prompt = format_prompt(&batch, None, None, None, None, None);
assert!(
prompt.contains(&format!("parent_event_id=\"{event_id}\"")),
"channel thread reply should include reply instruction with triggering event ID"
@@ -2644,7 +2698,7 @@ mod tests {
channel_type: "dm".into(),
};
let prompt = format_prompt(&batch, None, Some(&ci), None, None);
let prompt = format_prompt(&batch, None, None, Some(&ci), None, None);
assert!(
prompt.contains(&format!("parent_event_id=\"{event_id}\"")),
"DM thread reply should include reply instruction"
@@ -2665,7 +2719,7 @@ mod tests {
cancelled_events: vec![],
};
let prompt = format_prompt(&batch, None, None, None, None);
let prompt = format_prompt(&batch, None, None, None, None, None);
assert!(
!prompt.contains("parent_event_id"),
"top-level message should NOT include reply instruction"
@@ -2690,7 +2744,7 @@ mod tests {
channel_type: "dm".into(),
};
let prompt = format_prompt(&batch, None, Some(&ci), None, None);
let prompt = format_prompt(&batch, None, None, Some(&ci), None, None);
assert!(
!prompt.contains("parent_event_id"),
"DM non-reply should NOT include reply instruction"
@@ -2720,7 +2774,7 @@ mod tests {
cancelled_events: vec![],
};
let prompt = format_prompt(&batch, None, None, None, None);
let prompt = format_prompt(&batch, None, None, None, None, None);
// The instruction should use the triggering event's own ID — not root or parent.
assert!(
prompt.contains(&format!("parent_event_id=\"{event_id}\"")),
@@ -2763,7 +2817,7 @@ mod tests {
cancelled_events: vec![],
};
let prompt = format_prompt(&batch, None, None, None, None);
let prompt = format_prompt(&batch, None, None, None, None, None);
assert!(
prompt.contains(&format!("parent_event_id=\"{threaded_id}\"")),
"batched prompt should use last (threaded) event's ID"
@@ -2796,7 +2850,7 @@ mod tests {
cancelled_events: vec![],
};
let prompt = format_prompt(&batch, None, None, None, None);
let prompt = format_prompt(&batch, None, None, None, None, None);
assert!(
!prompt.contains("parent_event_id"),
"batched prompt where last event is top-level should NOT include reply instruction"

View File

@@ -40,6 +40,7 @@ uuid = { workspace = true }
# Typed event builders for all write operations
sprout-sdk = { workspace = true }
sprout-core = { workspace = true }
# Base64 encoding — NIP-98 event serialization for Authorization header
base64 = "0.22"

View File

@@ -161,6 +161,17 @@ impl SproutClient {
&self.relay_url
}
/// Return the owner pubkey carried by the NIP-OA auth tag, if any.
///
/// The auth tag is `["auth", owner_pubkey, conditions, sig]`; the
/// owner pubkey lives at index 1.
pub fn auth_tag_owner_hex(&self) -> Option<String> {
self.auth_tag
.as_ref()
.map(|t| t.as_slice())
.and_then(|slice| slice.get(1).cloned())
}
/// Sign an event builder, injecting the NIP-OA auth tag if configured.
///
/// All event creation should go through this method to ensure consistent

View File

@@ -0,0 +1,364 @@
//! `sprout mem` — agent-side engram management (NIP-AE).
//!
//! Subcommands:
//! - `sprout mem ls` — list non-tombstoned memories
//! - `sprout mem get <slug>` — print the value to stdout
//! - `sprout mem set <slug> <value|-> ` — write a value (use `-` for stdin)
//! - `sprout mem rm <slug>` — publish a tombstone
//!
//! The caller's `SPROUT_PRIVATE_KEY` is the agent's nsec. The agent's owner
//! pubkey is resolved from `SPROUT_AUTH_TAG` (NIP-OA attestation) or the
//! `--owner` flag. Every record is encrypted under the agent↔owner NIP-44
//! conversation key; both parties can decrypt.
use std::io::Read;
use std::time::SystemTime;
use nostr::PublicKey;
use sprout_core::engram::{
self, conversation_key, d_tag, normalize_slug, select_head, validate_and_decrypt, Body, Listing,
};
use sprout_core::kind::KIND_AGENT_ENGRAM;
use crate::client::SproutClient;
use crate::error::CliError;
/// Resolve the agent's owner pubkey: explicit `--owner` flag wins, otherwise
/// fall back to the NIP-OA `auth_tag` (which carries owner pubkey in slot 1).
fn resolve_owner(client: &SproutClient, owner_flag: Option<&str>) -> Result<PublicKey, CliError> {
if let Some(s) = owner_flag {
return PublicKey::from_hex(s)
.map_err(|e| CliError::Usage(format!("--owner must be a 64-hex pubkey: {e}")));
}
let tag = client.auth_tag_owner_hex().ok_or_else(|| {
CliError::Usage(
"owner pubkey required (set SPROUT_AUTH_TAG with a NIP-OA attestation or pass --owner)"
.into(),
)
})?;
PublicKey::from_hex(&tag)
.map_err(|e| CliError::Other(format!("auth_tag owner pubkey is not valid hex: {e}")))
}
fn now_secs() -> u64 {
SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
/// Submit a signed engram event and confirm the relay treated it as
/// authoritative. The relay returns `{accepted, message}` where the
/// `message` field starts with `"duplicate:"` when the write was rejected
/// as already-superseded by a later head (NIP-33 LWW). In that case we
/// surface a `Conflict` so callers don't lie about success.
async fn submit_engram(client: &SproutClient, event: nostr::Event) -> Result<(), CliError> {
let raw = client.submit_event(event).await?;
let parsed: serde_json::Value = serde_json::from_str(&raw)
.map_err(|e| CliError::Other(format!("relay response is not JSON: {e} ({raw})")))?;
let accepted = parsed
.get("accepted")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let message = parsed.get("message").and_then(|v| v.as_str()).unwrap_or("");
if !accepted {
return Err(CliError::Other(format!("relay rejected event: {message}")));
}
if message.starts_with("duplicate:") || message == "duplicate" {
return Err(CliError::Conflict(
"relay reported event as duplicate / dominated by a newer head".into(),
));
}
Ok(())
}
/// Parse a relay-response JSON array of events.
fn parse_events(json: &str) -> Result<Vec<nostr::Event>, CliError> {
let value: serde_json::Value = serde_json::from_str(json)
.map_err(|e| CliError::Other(format!("relay returned invalid JSON: {e}")))?;
let arr = value
.as_array()
.ok_or_else(|| CliError::Other("relay response is not an array".into()))?;
// Per NIP-AE head selection: discard events that fail any validation
// step and pick the head from the survivors. A single garbled response
// entry must not deny-of-service the whole listing.
let mut out = Vec::with_capacity(arr.len());
for ev in arr {
// Skip any event that fails to deserialize. Downstream validation
// (signature, decrypt, slug↔d) will discard further bad apples; we
// never want a single corrupt record to fail `mem ls`/`mem get`.
if let Ok(event) = serde_json::from_value::<nostr::Event>(ev.clone()) {
out.push(event);
}
}
Ok(out)
}
/// Fetch the head event for `slug`, returning `(Option<Event>, Option<Body>)`.
async fn fetch_head(
client: &SproutClient,
owner: &PublicKey,
slug: &str,
) -> Result<(Option<nostr::Event>, Option<Body>), CliError> {
let agent = client.keys();
let k_c = conversation_key(agent.secret_key(), owner);
let d = d_tag(&k_c, slug);
let filter = serde_json::json!({
"kinds": [KIND_AGENT_ENGRAM],
"authors": [agent.public_key().to_hex()],
"#d": [d],
"#p": [owner.to_hex()],
"limit": 16,
});
let raw = client.query(&filter).await?;
let events = parse_events(&raw)?;
let mut valid_with_body: Vec<(nostr::Event, Body)> = Vec::new();
for ev in events {
// Signature is validated by NIP-01 / NIP-44 — `nostr::Event::verify` is
// the conservative belt-and-suspenders check before decrypting.
if ev.verify().is_err() {
continue;
}
match validate_and_decrypt(&ev, &agent.public_key(), owner, agent.secret_key(), owner) {
Ok(body) => valid_with_body.push((ev, body)),
Err(_) => continue,
}
}
if valid_with_body.is_empty() {
return Ok((None, None));
}
let events: Vec<nostr::Event> = valid_with_body.iter().map(|(e, _)| e.clone()).collect();
// `select_head` returns `None` only on an empty iterator; we guarded
// that above, so the head is always present.
let Some(head) = select_head(events) else {
return Ok((None, None));
};
let body = valid_with_body
.into_iter()
.find(|(e, _)| e.id == head.id)
.map(|(_, b)| b);
Ok((Some(head), body))
}
/// `sprout mem ls` — list non-tombstoned memory entries.
pub async fn cmd_ls(
client: &SproutClient,
owner_flag: Option<&str>,
json: bool,
) -> Result<(), CliError> {
let owner = resolve_owner(client, owner_flag)?;
let agent = client.keys();
let filter = serde_json::json!({
"kinds": [KIND_AGENT_ENGRAM],
"authors": [agent.public_key().to_hex()],
"#p": [owner.to_hex()],
"limit": 5000,
});
let raw = client.query(&filter).await?;
let events = parse_events(&raw)?;
// Validate + decrypt + group by d tag.
use std::collections::HashMap;
let mut groups: HashMap<String, Vec<(nostr::Event, Body)>> = HashMap::new();
for ev in events {
if ev.verify().is_err() {
continue;
}
let Some(d_value) = ev
.tags
.iter()
.find(|t| t.kind().to_string() == "d")
.and_then(|t| t.content())
.map(|s| s.to_string())
else {
continue;
};
let body = match validate_and_decrypt(
&ev,
&agent.public_key(),
&owner,
agent.secret_key(),
&owner,
) {
Ok(b) => b,
Err(_) => continue,
};
groups.entry(d_value).or_default().push((ev, body));
}
let mut listings: Vec<Listing> = Vec::new();
for (_d, members) in groups {
let events: Vec<nostr::Event> = members.iter().map(|(e, _)| e.clone()).collect();
let Some(head) = select_head(events) else {
continue;
};
// `select_head` returns one of the events it received, so the head
// is always present in `members`. If something pathological breaks
// that invariant, skip the group rather than panic.
let Some((_, body)) = members.into_iter().find(|(e, _)| e.id == head.id) else {
continue;
};
// Drop tombstones and the core entry (per spec: listing excludes core).
match &body {
Body::Core { .. } => continue,
Body::Memory { value: None, .. } => continue,
Body::Memory { slug, .. } => {
listings.push(Listing {
slug: slug.clone(),
event_id: head.id.to_hex(),
created_at: head.created_at.as_u64(),
});
}
}
}
listings.sort_by(|a, b| a.slug.cmp(&b.slug));
if json {
println!("{}", serde_json::to_string(&listings).unwrap_or_default());
} else if listings.is_empty() {
eprintln!("(no memories)");
} else {
for l in &listings {
println!("{}\t{}\t{}", l.slug, l.created_at, l.event_id);
}
}
Ok(())
}
/// `sprout mem get <slug>` — print value (memory) or profile (core) to stdout.
///
/// Exit codes: 0 on found, 1 on absent or tombstoned.
pub async fn cmd_get(
client: &SproutClient,
raw_slug: &str,
owner_flag: Option<&str>,
) -> Result<(), CliError> {
let slug =
normalize_slug(raw_slug).map_err(|e| CliError::Usage(format!("invalid slug: {e}")))?;
let owner = resolve_owner(client, owner_flag)?;
let (_head, body) = fetch_head(client, &owner, &slug).await?;
use std::io::Write;
match body {
None => Err(CliError::NotFound(format!("not found: {slug}"))),
Some(Body::Memory { value: None, .. }) => {
Err(CliError::NotFound(format!("tombstoned: {slug}")))
}
Some(Body::Memory { value: Some(v), .. }) => {
// Raw stdout, no trailing newline — round-trips with `sprout mem set foo -`.
std::io::stdout()
.write_all(v.as_bytes())
.map_err(|e| CliError::Other(e.to_string()))
}
Some(Body::Core { profile }) => std::io::stdout()
.write_all(profile.as_bytes())
.map_err(|e| CliError::Other(e.to_string())),
}
}
/// `sprout mem set <slug> <value|->` — write a value or core profile.
///
/// Pass `-` to read the value from stdin.
pub async fn cmd_set(
client: &SproutClient,
raw_slug: &str,
raw_value: &str,
owner_flag: Option<&str>,
) -> Result<(), CliError> {
let slug =
normalize_slug(raw_slug).map_err(|e| CliError::Usage(format!("invalid slug: {e}")))?;
let value = if raw_value == "-" {
// Bound the stdin read so a runaway producer can't OOM us. We allow
// one extra byte over the NIP-44 plaintext cap so the build step can
// surface an exact `BodyTooLarge` if the cap is breached.
let limit = engram::NIP44_PLAINTEXT_MAX + 1;
let mut buf = String::new();
std::io::stdin()
.take(limit as u64)
.read_to_string(&mut buf)
.map_err(|e| CliError::Other(format!("stdin read failed: {e}")))?;
if buf.len() > engram::NIP44_PLAINTEXT_MAX {
return Err(CliError::Usage(format!(
"stdin value exceeds {}-byte NIP-44 plaintext limit",
engram::NIP44_PLAINTEXT_MAX
)));
}
buf
} else {
raw_value.to_string()
};
let owner = resolve_owner(client, owner_flag)?;
let body = if slug == engram::CORE_SLUG {
Body::Core { profile: value }
} else {
Body::Memory {
slug: slug.clone(),
value: Some(value),
}
};
let (head, _) = fetch_head(client, &owner, &slug).await?;
let prior_created_at = head.map(|e| e.created_at.as_u64());
let created_at = engram::monotonic_created_at(now_secs(), prior_created_at);
let agent = client.keys();
let event = engram::build_event(agent, &owner, &body, created_at)
.map_err(|e| CliError::Other(format!("build event failed: {e}")))?;
let id = event.id.to_hex();
submit_engram(client, event).await?;
eprintln!("wrote {slug} (event {id}, created_at {created_at})");
Ok(())
}
/// `sprout mem rm <slug>` — publish a tombstone (`value: null`).
///
/// `rm core` writes a tombstone-shaped body, but a core tombstone has no
/// well-defined semantics in NIP-AE (the spec only defines tombstones for
/// memory entries). We refuse it and tell the operator to overwrite `core`
/// with an empty profile instead.
pub async fn cmd_rm(
client: &SproutClient,
raw_slug: &str,
owner_flag: Option<&str>,
) -> Result<(), CliError> {
let slug =
normalize_slug(raw_slug).map_err(|e| CliError::Usage(format!("invalid slug: {e}")))?;
if slug == engram::CORE_SLUG {
return Err(CliError::Usage(
"core cannot be tombstoned; overwrite it with `sprout mem set core ''` instead".into(),
));
}
let owner = resolve_owner(client, owner_flag)?;
let body = Body::Memory {
slug: slug.clone(),
value: None,
};
let (head, _) = fetch_head(client, &owner, &slug).await?;
let prior_created_at = head.map(|e| e.created_at.as_u64());
let created_at = engram::monotonic_created_at(now_secs(), prior_created_at);
let agent = client.keys();
let event = engram::build_event(agent, &owner, &body, created_at)
.map_err(|e| CliError::Other(format!("build event failed: {e}")))?;
let id = event.id.to_hex();
submit_engram(client, event).await?;
eprintln!("tombstoned {slug} (event {id}, created_at {created_at})");
Ok(())
}
// ---------------------------------------------------------------------------
// Dispatch
// ---------------------------------------------------------------------------
pub async fn dispatch(cmd: crate::MemCmd, client: &SproutClient) -> Result<(), CliError> {
use crate::MemCmd;
match cmd {
MemCmd::Ls { owner, json } => cmd_ls(client, owner.as_deref(), json).await,
MemCmd::Get { slug, owner } => cmd_get(client, &slug, owner.as_deref()).await,
MemCmd::Set { slug, value, owner } => {
cmd_set(client, &slug, &value, owner.as_deref()).await
}
MemCmd::Rm { slug, owner } => cmd_rm(client, &slug, owner.as_deref()).await,
}
}

View File

@@ -1,6 +1,7 @@
pub mod channels;
pub mod dms;
pub mod feed;
pub mod mem;
pub mod messages;
pub mod pack;
pub mod reactions;

View File

@@ -22,13 +22,24 @@ pub enum CliError {
#[error("key error: {0}")]
Key(String),
/// Relay accepted the event but reported it as superseded by a newer
/// head — used by `sprout mem` set/rm to surface NIP-33 LWW conflicts.
#[error("conflict: {0}")]
Conflict(String),
/// Requested resource was absent or tombstoned (e.g. `sprout mem get`
/// for a slug with no head).
#[error("{0}")]
NotFound(String),
/// Catch-all for unexpected failures
#[error("{0}")]
Other(String),
}
/// Map CliError to process exit code.
/// 0=success (not an error), 1=user, 2=network/relay, 3=auth, 4=other
/// 0=success (not an error), 1=user/not-found, 2=network/relay, 3=auth,
/// 4=other, 5=write conflict (NIP-33 dominated head).
pub fn exit_code(e: &CliError) -> i32 {
match e {
CliError::Usage(_) => 1,
@@ -42,6 +53,8 @@ pub fn exit_code(e: &CliError) -> i32 {
CliError::Network(_) => 2,
CliError::Auth(_) => 3,
CliError::Key(_) => 3,
CliError::Conflict(_) => 5,
CliError::NotFound(_) => 1,
CliError::Other(_) => 4,
}
}
@@ -61,6 +74,8 @@ pub fn print_error(e: &CliError) {
CliError::Network(_) => "network_error",
CliError::Auth(_) => "auth_error",
CliError::Key(_) => "key_error",
CliError::Conflict(_) => "conflict",
CliError::NotFound(_) => "not_found",
CliError::Other(_) => "error",
};
let obj = serde_json::json!({

View File

@@ -63,7 +63,7 @@ Configuration (flags override env vars):
The 'pack' subcommand runs locally and does not require a relay connection.
Exit codes: 0=ok 1=bad input 2=relay/network error 3=auth error 4=other
Exit codes: 0=ok 1=bad input 2=relay/network error 3=auth error 4=other 5=write conflict
Errors are JSON on stderr: {\"error\": \"<category>\", \"message\": \"<detail>\"}"
)]
struct Cli {
@@ -184,6 +184,9 @@ enum Cmd {
/// Upload files to the relay's Blossom store
#[command(subcommand)]
Upload(UploadCmd),
/// Agent engram management — persistent memory per NIP-AE
#[command(subcommand)]
Mem(MemCmd),
/// Persona pack operations (local, no relay connection needed)
#[command(subcommand)]
Pack(PackCmd),
@@ -806,6 +809,43 @@ pub enum UploadCmd {
},
}
// ---------------------------------------------------------------------------
// Mem subcommands (NIP-AE)
// ---------------------------------------------------------------------------
/// Subcommands for `sprout mem`.
#[derive(Subcommand)]
pub enum MemCmd {
/// List non-tombstoned memory entries
Ls {
/// Owner pubkey (hex). Overrides SPROUT_AUTH_TAG.
#[arg(long)]
owner: Option<String>,
/// Emit JSON instead of tab-delimited lines.
#[arg(long, default_value_t = false)]
json: bool,
},
/// Print the value of a slug to stdout (no trailing newline)
Get {
slug: String,
#[arg(long)]
owner: Option<String>,
},
/// Set a slug's value. Pass `-` to read the value from stdin.
Set {
slug: String,
value: String,
#[arg(long)]
owner: Option<String>,
},
/// Publish a tombstone for a slug (cannot be used on `core`).
Rm {
slug: String,
#[arg(long)]
owner: Option<String>,
},
}
// ---------------------------------------------------------------------------
// Pack subcommands (local, no relay connection needed)
// ---------------------------------------------------------------------------
@@ -878,6 +918,7 @@ async fn run(cli: Cli) -> Result<(), CliError> {
Cmd::Social(sub) => commands::social::dispatch(sub, &client).await,
Cmd::Repos(sub) => commands::repos::dispatch(sub, &client).await,
Cmd::Upload(sub) => commands::upload::dispatch(sub, &client).await,
Cmd::Mem(sub) => commands::mem::dispatch(sub, &client).await,
Cmd::Pack(_) => unreachable!("handled above"),
}
}
@@ -904,6 +945,7 @@ mod tests {
"channels",
"dms",
"feed",
"mem",
"messages",
"pack",
"reactions",

View File

@@ -19,6 +19,8 @@ thiserror = { workspace = true }
uuid = { workspace = true }
chrono = { workspace = true }
hex = { workspace = true }
hmac = { workspace = true }
sha2 = { workspace = true }
schemars = { workspace = true, optional = true }
rand = { workspace = true }
subtle = { workspace = true }

View File

@@ -0,0 +1,835 @@
//! NIP-AE Agent Engrams — pure crypto + parsing primitives.
//!
//! See `docs/nips/NIP-AE.md` for the spec. This module is I/O-free: it does
//! not talk to relays or filesystems. Callers wire it to a transport and a
//! key source.
//!
//! Shared by `sprout-cli` (`sprout mem …`) and `sprout-acp` (core injection
//! at session creation).
use hmac::digest::KeyInit;
use hmac::{Hmac, Mac};
use nostr::nips::nip44::{self, v2::ConversationKey, Version};
use nostr::{Event, EventBuilder, Keys, Kind, PublicKey, SecretKey, Tag};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use crate::kind::KIND_AGENT_ENGRAM;
/// The reserved slug for the agent's core (identity) engram.
pub const CORE_SLUG: &str = "core";
/// Domain prefix for the `d`-tag HMAC. Followed by a `0x00` byte and the slug.
/// Versioned independently of the NIP number; future revisions MUST change it.
pub const D_TAG_DOMAIN: &[u8] = b"agent-memory/v1/d-tag";
/// NIP-44 plaintext limit (bytes). Bodies whose serialized JSON exceeds this
/// MUST NOT be encrypted (spec: *Encryption*).
pub const NIP44_PLAINTEXT_MAX: usize = 65_535;
/// Maximum slug length in bytes (spec: *Slugs*).
pub const SLUG_MAX_LEN: usize = 255;
/// Memory slug grammar prefix.
const MEM_PREFIX: &str = "mem/";
/// Errors from engram operations.
#[derive(Debug, thiserror::Error)]
pub enum EngramError {
/// Slug failed the *Slugs* grammar.
#[error("invalid slug: {0}")]
InvalidSlug(String),
/// Body parsing or shape check failed.
#[error("invalid body: {0}")]
InvalidBody(String),
/// Event failed *Head selection* rule (1) — tag shape or addressing.
#[error("invalid envelope: {0}")]
InvalidEnvelope(String),
/// NIP-44 decryption failed.
#[error("decrypt failed")]
Decrypt,
/// Encryption failed.
#[error("encrypt failed: {0}")]
Encrypt(String),
/// Body exceeds the NIP-44 plaintext cap.
#[error("body exceeds {NIP44_PLAINTEXT_MAX}-byte plaintext limit ({0} bytes)")]
BodyTooLarge(usize),
/// Signing error.
#[error("sign failed: {0}")]
Sign(String),
}
/// Validate a slug against the *Slugs* grammar.
///
/// Returns `Ok(())` if the slug is the reserved `core` or matches:
/// `^mem/[a-z0-9][a-z0-9_-]{0,63}(/[a-z0-9][a-z0-9_-]{0,63})*$` with total
/// length ≤ 255 bytes.
pub fn validate_slug(slug: &str) -> Result<(), EngramError> {
if slug == CORE_SLUG {
return Ok(());
}
if slug.len() > SLUG_MAX_LEN {
return Err(EngramError::InvalidSlug(format!(
"length {} exceeds {}",
slug.len(),
SLUG_MAX_LEN
)));
}
let Some(rest) = slug.strip_prefix(MEM_PREFIX) else {
return Err(EngramError::InvalidSlug(format!(
"expected `core` or `mem/…`, got {slug:?}"
)));
};
if rest.is_empty() {
return Err(EngramError::InvalidSlug("empty after `mem/`".into()));
}
for (i, segment) in rest.split('/').enumerate() {
validate_segment(segment).map_err(|why| {
EngramError::InvalidSlug(format!("segment {} ({:?}): {why}", i + 1, segment))
})?;
}
Ok(())
}
fn validate_segment(s: &str) -> Result<(), &'static str> {
if s.is_empty() {
return Err("empty");
}
if s.len() > 64 {
return Err("longer than 64 bytes");
}
let bytes = s.as_bytes();
let first = bytes[0];
if !is_lower_alnum(first) {
return Err("first byte must be [a-z0-9]");
}
for &b in &bytes[1..] {
if !(is_lower_alnum(b) || b == b'_' || b == b'-') {
return Err("only [a-z0-9_-] allowed after the first byte");
}
}
Ok(())
}
const fn is_lower_alnum(b: u8) -> bool {
matches!(b, b'a'..=b'z' | b'0'..=b'9')
}
/// Normalize a user-supplied shorthand to a full slug.
///
/// `core` is returned verbatim. Otherwise, if the input does not already
/// start with `mem/`, the prefix is added. The result is validated; the
/// caller still gets an `InvalidSlug` if the shorthand is not a valid slug.
pub fn normalize_slug(raw: &str) -> Result<String, EngramError> {
let slug = if raw == CORE_SLUG || raw.starts_with(MEM_PREFIX) {
raw.to_string()
} else {
format!("{MEM_PREFIX}{raw}")
};
validate_slug(&slug)?;
Ok(slug)
}
/// Derive the conversation key `K_c` for the agent ↔ owner pair (NIP-44 v2).
///
/// `K_c` is symmetric: `derive(seckey_a, pubkey_o) == derive(seckey_o, pubkey_a)`.
pub fn conversation_key(my_seckey: &SecretKey, their_pubkey: &PublicKey) -> ConversationKey {
ConversationKey::derive(my_seckey, their_pubkey)
}
/// Compute the `d` tag for a slug under a conversation key.
///
/// `d = lower_hex(HMAC-SHA256(K_c, "agent-memory/v1/d-tag" || 0x00 || slug))`,
/// 64 hex characters.
pub fn d_tag(k_c: &ConversationKey, slug: &str) -> String {
// HMAC-SHA256 accepts a key of any byte length; `new_from_slice` only
// returns `Err` for fixed-length MAC variants. This is infallible for
// SHA-256 and propagating it would just add noise at every call site.
let mut mac = <Hmac<Sha256> as KeyInit>::new_from_slice(k_c.as_bytes())
.expect("HMAC-SHA256 is keyed-prefix MAC; new_from_slice cannot fail");
mac.update(D_TAG_DOMAIN);
mac.update(&[0u8]);
mac.update(slug.as_bytes());
hex::encode(mac.finalize().into_bytes())
}
// ── Bodies ──────────────────────────────────────────────────────────────────
/// A decoded engram body. The slug discriminates the variant.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Body {
/// `slug == "core"` — agent identity surface.
Core {
/// Free-form UTF-8 maintained by the agent.
profile: String,
},
/// `slug == "mem/…"` — one logical entry. `value = None` is a tombstone.
Memory {
/// The slug this body addresses.
slug: String,
/// The entry's value, or `None` for a tombstone.
value: Option<String>,
},
}
impl Body {
/// Return the slug this body addresses.
pub fn slug(&self) -> &str {
match self {
Body::Core { .. } => CORE_SLUG,
Body::Memory { slug, .. } => slug,
}
}
/// `true` if this is a tombstone (memory with `value = null`).
pub fn is_tombstone(&self) -> bool {
matches!(self, Body::Memory { value: None, .. })
}
/// Serialize to the exact JSON encoding the spec specifies for the body
/// passed to NIP-44. Whitespace-free, slug first.
pub fn to_json_bytes(&self) -> Vec<u8> {
// We hand-roll the JSON so the *Reference test vectors* match
// byte-for-byte (`{"slug":"…","value":"…"}` / `{"slug":"core","profile":"…"}`).
let mut out = String::with_capacity(64);
out.push('{');
out.push_str("\"slug\":");
write_json_string(&mut out, self.slug());
match self {
Body::Memory { value: Some(v), .. } => {
out.push_str(",\"value\":");
write_json_string(&mut out, v);
}
Body::Memory { value: None, .. } => {
out.push_str(",\"value\":null");
}
Body::Core { profile } => {
out.push_str(",\"profile\":");
write_json_string(&mut out, profile);
}
}
out.push('}');
out.into_bytes()
}
/// Parse a body from its decrypted JSON bytes. Rejects duplicate object
/// member names (spec: *Head selection* rule (3)). Unknown fields are
/// ignored. The body's `slug` is validated.
pub fn from_json_bytes(bytes: &[u8]) -> Result<Self, EngramError> {
// serde_json by default *accepts* duplicate keys (last wins). The spec
// requires rejection, so we deserialize through a wrapper that walks
// the tree once and fails on the first repeated member name at any
// nesting depth.
let raw = parse_strict_json(bytes)?;
let obj = match raw {
serde_json::Value::Object(m) => m,
_ => return Err(EngramError::InvalidBody("top-level not an object".into())),
};
let slug = match obj.get("slug") {
Some(serde_json::Value::String(s)) => s.clone(),
Some(_) => return Err(EngramError::InvalidBody("`slug` is not a string".into())),
None => return Err(EngramError::InvalidBody("missing `slug`".into())),
};
validate_slug(&slug)?;
if slug == CORE_SLUG {
let profile = match obj.get("profile") {
Some(serde_json::Value::String(s)) => s.clone(),
Some(_) => {
return Err(EngramError::InvalidBody("`profile` is not a string".into()))
}
None => return Err(EngramError::InvalidBody("core missing `profile`".into())),
};
Ok(Body::Core { profile })
} else {
let value = match obj.get("value") {
Some(serde_json::Value::String(s)) => Some(s.clone()),
Some(serde_json::Value::Null) => None,
Some(_) => return Err(EngramError::InvalidBody("`value` is not a string".into())),
None => return Err(EngramError::InvalidBody("memory missing `value`".into())),
};
Ok(Body::Memory { slug, value })
}
}
}
/// Minimal JSON-string encoder per RFC 8259 §7. Escapes `"` `\` and the
/// required control chars (`\b \f \n \r \t` + `\u00XX` for the rest).
fn write_json_string(out: &mut String, s: &str) {
out.push('"');
for ch in s.chars() {
match ch {
'"' => out.push_str("\\\""),
'\\' => out.push_str("\\\\"),
'\u{08}' => out.push_str("\\b"),
'\u{0c}' => out.push_str("\\f"),
'\n' => out.push_str("\\n"),
'\r' => out.push_str("\\r"),
'\t' => out.push_str("\\t"),
c if (c as u32) < 0x20 => {
out.push_str(&format!("\\u{:04x}", c as u32));
}
c => out.push(c),
}
}
out.push('"');
}
/// Parse a JSON document, rejecting duplicate member names at any nesting
/// depth. Returns the parsed `serde_json::Value` on success.
///
/// `serde_json::from_slice` happily last-wins on duplicates, but *Head
/// selection* rule (3) requires strict rejection. We get correct behaviour
/// by feeding the deserializer a custom visitor for maps that tracks seen
/// keys; arrays / scalars fall through to the default `Value` visitor.
fn parse_strict_json(bytes: &[u8]) -> Result<serde_json::Value, EngramError> {
use serde::de::{DeserializeSeed, Deserializer, MapAccess, SeqAccess, Visitor};
use serde_json::Value;
use std::collections::HashSet;
use std::fmt;
struct StrictValue;
impl<'de> DeserializeSeed<'de> for StrictValue {
type Value = Value;
fn deserialize<D: Deserializer<'de>>(self, d: D) -> Result<Value, D::Error> {
d.deserialize_any(StrictValue)
}
}
impl<'de> Visitor<'de> for StrictValue {
type Value = Value;
fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str("any valid JSON value (objects must have unique keys)")
}
fn visit_bool<E>(self, v: bool) -> Result<Value, E> {
Ok(Value::Bool(v))
}
fn visit_i64<E>(self, v: i64) -> Result<Value, E> {
Ok(Value::Number(v.into()))
}
fn visit_u64<E>(self, v: u64) -> Result<Value, E> {
Ok(Value::Number(v.into()))
}
fn visit_f64<E: serde::de::Error>(self, v: f64) -> Result<Value, E> {
serde_json::Number::from_f64(v)
.map(Value::Number)
.ok_or_else(|| E::custom("non-finite float"))
}
fn visit_str<E>(self, v: &str) -> Result<Value, E> {
Ok(Value::String(v.to_owned()))
}
fn visit_string<E>(self, v: String) -> Result<Value, E> {
Ok(Value::String(v))
}
fn visit_unit<E>(self) -> Result<Value, E> {
Ok(Value::Null)
}
fn visit_none<E>(self) -> Result<Value, E> {
Ok(Value::Null)
}
fn visit_some<D: Deserializer<'de>>(self, d: D) -> Result<Value, D::Error> {
d.deserialize_any(StrictValue)
}
fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Value, A::Error> {
let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(0));
while let Some(v) = seq.next_element_seed(StrictValue)? {
out.push(v);
}
Ok(Value::Array(out))
}
fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Value, A::Error> {
use serde::de::Error;
let mut seen: HashSet<String> = HashSet::new();
let mut out = serde_json::Map::new();
while let Some(k) = map.next_key::<String>()? {
if !seen.insert(k.clone()) {
return Err(A::Error::custom(format!(
"duplicate object member name: {k}"
)));
}
let v = map.next_value_seed(StrictValue)?;
out.insert(k, v);
}
Ok(Value::Object(out))
}
}
let mut de = serde_json::Deserializer::from_slice(bytes);
let v = StrictValue
.deserialize(&mut de)
.map_err(|e| EngramError::InvalidBody(format!("invalid JSON body: {e}")))?;
de.end()
.map_err(|e| EngramError::InvalidBody(format!("trailing data after JSON body: {e}")))?;
Ok(v)
}
// ── Envelope build / parse ──────────────────────────────────────────────────
/// Build a signed `kind:30174` event for a given body.
///
/// * `created_at` is the timestamp to sign — callers MUST supply a value
/// respecting the *Writing* monotonic rule (`max(now, T_head + 1)`).
/// * Returns `BodyTooLarge` if the serialized body exceeds 65,535 bytes.
pub fn build_event(
agent_keys: &Keys,
owner_pubkey: &PublicKey,
body: &Body,
created_at: u64,
) -> Result<Event, EngramError> {
let plaintext = body.to_json_bytes();
if plaintext.len() > NIP44_PLAINTEXT_MAX {
return Err(EngramError::BodyTooLarge(plaintext.len()));
}
// `to_json_bytes` only emits ASCII control chars or `&str` bytes, so
// this is always Ok. We still verify rather than `.expect()` so a future
// change to the serializer can't silently introduce a panic on the hot
// path.
let plaintext_str = std::str::from_utf8(&plaintext)
.map_err(|e| EngramError::Encrypt(format!("body JSON not UTF-8: {e}")))?;
let k_c = conversation_key(agent_keys.secret_key(), owner_pubkey);
let ciphertext = nip44::encrypt(
agent_keys.secret_key(),
owner_pubkey,
plaintext_str,
Version::V2,
)
.map_err(|e| EngramError::Encrypt(e.to_string()))?;
let d = d_tag(&k_c, body.slug());
let tags = vec![
Tag::parse(&["d", &d]).map_err(|e| EngramError::Encrypt(e.to_string()))?,
Tag::parse(&["p", &owner_pubkey.to_hex()])
.map_err(|e| EngramError::Encrypt(e.to_string()))?,
];
EventBuilder::new(Kind::Custom(KIND_AGENT_ENGRAM as u16), ciphertext, tags)
.custom_created_at(nostr::Timestamp::from(created_at))
.sign_with_keys(agent_keys)
.map_err(|e| EngramError::Sign(e.to_string()))
}
/// Validate an event against *Head selection* rules (1) and (5) and return
/// the decoded body. Caller must verify the signature beforehand (NIP-44
/// requires outer-signature-before-decrypt; `nostr::Event::verify_signature`
/// or NIP-01 ingest path handles this).
///
/// * `event` — the candidate event.
/// * `expected_agent` — the agent pubkey the event's `pubkey` field must equal.
/// * `expected_owner` — the owner pubkey the event's `p` tag must contain.
/// * `my_seckey` / `their_pubkey` — the *NIP-44 ECDH pair* the caller holds.
/// For an agent reading its own engrams: `my_seckey = seckey_a`,
/// `their_pubkey = pubkey_o`. For an owner reading the agent's engrams:
/// `my_seckey = seckey_o`, `their_pubkey = pubkey_a`. Either yields the
/// same `K_c`.
pub fn validate_and_decrypt(
event: &Event,
expected_agent: &PublicKey,
expected_owner: &PublicKey,
my_seckey: &SecretKey,
their_pubkey: &PublicKey,
) -> Result<Body, EngramError> {
if event.kind.as_u16() as u32 != KIND_AGENT_ENGRAM {
return Err(EngramError::InvalidEnvelope(format!(
"wrong kind: {}",
event.kind.as_u16()
)));
}
if &event.pubkey != expected_agent {
return Err(EngramError::InvalidEnvelope(
"pubkey != expected_agent".into(),
));
}
let mut d_tags = event.tags.iter().filter(|t| t.kind().to_string() == "d");
let d_value = d_tags
.next()
.and_then(|t| t.content().map(|s| s.to_string()))
.ok_or_else(|| EngramError::InvalidEnvelope("missing d tag".into()))?;
if d_tags.next().is_some() {
return Err(EngramError::InvalidEnvelope("multiple d tags".into()));
}
// Spec: d = lower_hex(HMAC...). Anything else is non-canonical and we
// refuse to interoperate with it.
if d_value.len() != 64
|| !d_value
.bytes()
.all(|b| b.is_ascii_hexdigit() && !b.is_ascii_uppercase())
{
return Err(EngramError::InvalidEnvelope(
"d tag must be 64 lowercase hex chars".into(),
));
}
let mut p_tags = event.tags.iter().filter(|t| t.kind().to_string() == "p");
let p_value = p_tags
.next()
.and_then(|t| t.content().map(|s| s.to_string()))
.ok_or_else(|| EngramError::InvalidEnvelope("missing p tag".into()))?;
if p_tags.next().is_some() {
return Err(EngramError::InvalidEnvelope("multiple p tags".into()));
}
if p_value.eq_ignore_ascii_case(&expected_owner.to_hex()) {
// ok
} else {
return Err(EngramError::InvalidEnvelope(
"p tag != expected_owner".into(),
));
}
// Decrypt. `K_c` is symmetric per NIP-44, so the caller's `(my_seckey,
// their_pubkey)` pair yields the same conversation key regardless of
// whether the caller is the agent or the owner.
let plaintext = nip44::decrypt(my_seckey, their_pubkey, &event.content)
.map_err(|_| EngramError::Decrypt)?;
let body = Body::from_json_bytes(plaintext.as_bytes())?;
// Rule (4): body slug re-derives to the event's d tag.
let k_c = conversation_key(my_seckey, their_pubkey);
let derived = d_tag(&k_c, body.slug());
if derived != d_value {
return Err(EngramError::InvalidEnvelope(
"body slug does not re-derive to d tag".into(),
));
}
Ok(body)
}
/// Pick the head from a set of events targeting the same slug — greatest
/// `created_at`, ties broken by lowest event id per NIP-01.
///
/// Caller is responsible for validating each event (kind/tags/sig/decrypt)
/// before passing it in; this function only does the LWW tiebreak.
pub fn select_head<I>(events: I) -> Option<Event>
where
I: IntoIterator<Item = Event>,
{
events.into_iter().reduce(|a, b| {
let a_ts = a.created_at.as_u64();
let b_ts = b.created_at.as_u64();
if b_ts > a_ts {
return b;
}
if a_ts > b_ts {
return a;
}
// Tie — lower id wins (lexicographic on hex == byte-wise on bytes).
if b.id.to_hex() < a.id.to_hex() {
b
} else {
a
}
})
}
/// Compute `created_at` for a new write per the *Writing* monotonic rule:
/// `max(now, prior_head_created_at + 1)`.
pub fn monotonic_created_at(now: u64, prior_head: Option<u64>) -> u64 {
match prior_head {
Some(t) => now.max(t.saturating_add(1)),
None => now,
}
}
/// Wire representation for `sprout mem ls`: one entry per non-tombstone
/// memory slug (`core` is excluded by the listing procedure).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Listing {
/// The slug (e.g. `mem/notes/today`).
pub slug: String,
/// Event id of the current head.
pub event_id: String,
/// `created_at` of the current head.
pub created_at: u64,
}
#[cfg(test)]
mod tests {
use super::*;
fn keys_from_hex(s: &str) -> Keys {
Keys::parse(s).unwrap()
}
// ── Reference test vectors from docs/nips/NIP-AE.md §"Reference test
// vectors". Pinning these as CI invariants is the single best
// interop guarantee for this implementation. ──
const SECKEY_A: &str = "0000000000000000000000000000000000000000000000000000000000000001";
const SECKEY_O: &str = "0000000000000000000000000000000000000000000000000000000000000002";
const PUBKEY_A: &str = "79be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798";
const PUBKEY_O: &str = "c6047f9441ed7d6d3045406e95c07cd85c778e4b8cef3ca7abac09b95c709ee5";
const K_C_HEX: &str = "c41c775356fd92eadc63ff5a0dc1da211b268cbea22316767095b2871ea1412d";
const D_CORE: &str = "bdc233238ffe52e272b44cc233c8f33a2bc510b08be04495b225964283be4a90";
const D_EXAMPLE: &str = "72d4f9629106451505d7d341ea85bb3ebad4f654fcfd2aad100d5a35f8a85cba";
const D_NOTES: &str = "31651571a312780cfdc1f0b706b682ac9f3f51a053e8dca76fe57710bae5a4d4";
#[test]
fn pubkeys_match_spec() {
assert_eq!(keys_from_hex(SECKEY_A).public_key().to_hex(), PUBKEY_A);
assert_eq!(keys_from_hex(SECKEY_O).public_key().to_hex(), PUBKEY_O);
}
#[test]
fn conversation_key_matches_spec() {
let a = keys_from_hex(SECKEY_A);
let o = keys_from_hex(SECKEY_O);
let k_c_ao = conversation_key(a.secret_key(), &o.public_key());
let k_c_oa = conversation_key(o.secret_key(), &a.public_key());
assert_eq!(hex::encode(k_c_ao.as_bytes()), K_C_HEX, "agent-side K_c");
assert_eq!(hex::encode(k_c_oa.as_bytes()), K_C_HEX, "owner-side K_c");
}
#[test]
fn d_tags_match_spec() {
let a = keys_from_hex(SECKEY_A);
let o = keys_from_hex(SECKEY_O);
let k_c = conversation_key(a.secret_key(), &o.public_key());
assert_eq!(d_tag(&k_c, "core"), D_CORE);
assert_eq!(d_tag(&k_c, "mem/example"), D_EXAMPLE);
assert_eq!(d_tag(&k_c, "mem/notes/2026-05-12"), D_NOTES);
}
#[test]
fn body_round_trips_byte_exact() {
// Memory with value.
let b = Body::Memory {
slug: "mem/example".into(),
value: Some("hello, agent memory".into()),
};
assert_eq!(
b.to_json_bytes(),
br#"{"slug":"mem/example","value":"hello, agent memory"}"#.to_vec()
);
// Memory with reference.
let b = Body::Memory {
slug: "mem/notes/2026-05-12".into(),
value: Some("meeting note: [[mem/example]]".into()),
};
assert_eq!(
b.to_json_bytes(),
br#"{"slug":"mem/notes/2026-05-12","value":"meeting note: [[mem/example]]"}"#.to_vec()
);
// Tombstone.
let b = Body::Memory {
slug: "mem/example".into(),
value: None,
};
assert_eq!(
b.to_json_bytes(),
br#"{"slug":"mem/example","value":null}"#.to_vec()
);
// Core.
let b = Body::Core {
profile: "test agent. see [[mem/example]] and [[mem/notes/2026-05-12]].".into(),
};
assert_eq!(
b.to_json_bytes(),
br#"{"slug":"core","profile":"test agent. see [[mem/example]] and [[mem/notes/2026-05-12]]."}"#.to_vec()
);
}
#[test]
fn body_parse_rejects_duplicate_keys() {
let dup = br#"{"slug":"mem/x","slug":"mem/y","value":"v"}"#;
let err = Body::from_json_bytes(dup).unwrap_err();
let msg = format!("{err}");
assert!(msg.contains("duplicate"), "got: {msg}");
}
#[test]
fn body_parse_ignores_unknown_fields() {
let body = br#"{"slug":"mem/x","value":"v","unknown":42,"future":{"nested":"ok"}}"#;
let b = Body::from_json_bytes(body).unwrap();
assert!(matches!(b, Body::Memory { value: Some(_), .. }));
}
#[test]
fn body_parse_accepts_arrays_of_repeated_strings() {
// Repeats inside an *array* aren't object-member duplicates and must
// be accepted. The previous hand-rolled scanner incorrectly flagged
// this as a dup.
let body = br#"{"slug":"mem/x","value":"v","future":["slug","slug","value"]}"#;
let b = Body::from_json_bytes(body).unwrap();
assert!(matches!(b, Body::Memory { value: Some(_), .. }));
}
#[test]
fn body_parse_accepts_surrogate_pair_escapes() {
// `\uD83D\uDE00` (😀) is a valid UTF-16 surrogate pair that serde_json
// decodes correctly. The previous hand-rolled scanner rejected each
// half as an invalid codepoint.
let body = br#"{"slug":"mem/x","value":"hi \uD83D\uDE00"}"#;
let b = Body::from_json_bytes(body).unwrap();
match b {
Body::Memory {
value: Some(v),
slug,
} => {
assert_eq!(slug, "mem/x");
assert_eq!(v, "hi \u{1F600}");
}
_ => panic!("expected Memory"),
}
}
#[test]
fn body_parse_rejects_duplicates_at_any_depth() {
// Object inside an unknown field with duplicate keys: still rejected.
let dup = br#"{"slug":"mem/x","value":"v","nested":{"k":1,"k":2}}"#;
let err = Body::from_json_bytes(dup).unwrap_err();
assert!(format!("{err}").contains("duplicate"), "got: {err}");
}
#[test]
fn validate_slug_accepts_grammar() {
for ok in [
"core",
"mem/x",
"mem/x-y_z",
"mem/0",
"mem/notes/2026-05-12",
"mem/a/b/c",
] {
assert!(validate_slug(ok).is_ok(), "{ok} should be valid");
}
}
#[test]
fn validate_slug_rejects_garbage() {
for bad in [
"", "MEM/x", "mem/", "mem//x", "mem/-x", "mem/_x", "mem/x/-y", "mem/x/", "memx",
"mem/x/Y",
"mem/x.y",
// 64-byte boundary: first byte then 64 more = 65 -> too long
] {
assert!(validate_slug(bad).is_err(), "{bad:?} should be rejected");
}
}
#[test]
fn normalize_slug_adds_mem_prefix() {
assert_eq!(normalize_slug("foo").unwrap(), "mem/foo");
assert_eq!(normalize_slug("core").unwrap(), "core");
assert_eq!(normalize_slug("mem/bar").unwrap(), "mem/bar");
assert!(normalize_slug("Foo").is_err());
}
#[test]
fn monotonic_clock_rule() {
assert_eq!(monotonic_created_at(100, None), 100);
assert_eq!(monotonic_created_at(100, Some(50)), 100);
assert_eq!(monotonic_created_at(100, Some(100)), 101);
assert_eq!(monotonic_created_at(100, Some(200)), 201);
}
#[test]
fn select_head_lww_with_id_tiebreak() {
// We build three events with the same kind/pubkey/d tag and check
// (1) greater created_at wins, (2) ties broken by lower id.
let agent = keys_from_hex(SECKEY_A);
let owner = keys_from_hex(SECKEY_O);
let body = Body::Memory {
slug: "mem/example".into(),
value: Some("v".into()),
};
let e1 = build_event(&agent, &owner.public_key(), &body, 1_700_000_000).unwrap();
let e2 = build_event(&agent, &owner.public_key(), &body, 1_700_000_001).unwrap();
let pick = select_head([e1.clone(), e2.clone()]).unwrap();
assert_eq!(pick.id, e2.id);
}
#[test]
fn round_trip_event_validates_and_decrypts() {
let agent = keys_from_hex(SECKEY_A);
let owner = keys_from_hex(SECKEY_O);
let original = Body::Memory {
slug: "mem/example".into(),
value: Some("hello".into()),
};
let event = build_event(&agent, &owner.public_key(), &original, 1_700_000_000).unwrap();
let decoded = validate_and_decrypt(
&event,
&agent.public_key(),
&owner.public_key(),
owner.secret_key(),
&agent.public_key(),
)
.unwrap();
// And the agent-side decrypt path must also work.
let decoded_agent = validate_and_decrypt(
&event,
&agent.public_key(),
&owner.public_key(),
agent.secret_key(),
&owner.public_key(),
)
.unwrap();
assert_eq!(decoded_agent, decoded);
assert_eq!(decoded, original);
}
#[test]
fn validate_and_decrypt_rejects_non_lowercase_d_tag() {
// Build a perfectly valid event but tamper the d tag to uppercase
// before signing. `validate_and_decrypt` must refuse, because the
// spec requires lowercase hex and otherwise two stringly-different
// d-tags would map to the same head.
let agent = keys_from_hex(SECKEY_A);
let owner = keys_from_hex(SECKEY_O);
let body = Body::Memory {
slug: "mem/example".into(),
value: Some("hi".into()),
};
let ev = build_event(&agent, &owner.public_key(), &body, 1).unwrap();
// Re-sign with the d tag uppercased.
let d = ev
.tags
.iter()
.find(|t| t.kind().to_string() == "d")
.and_then(|t| t.content().map(|s| s.to_string()))
.unwrap();
let upper_d = d.to_uppercase();
let tags = vec![
Tag::parse(&["d", &upper_d]).unwrap(),
Tag::parse(&["p", &owner.public_key().to_hex()]).unwrap(),
];
let tampered = EventBuilder::new(
Kind::Custom(KIND_AGENT_ENGRAM as u16),
ev.content.clone(),
tags,
)
.custom_created_at(ev.created_at)
.sign_with_keys(&agent)
.unwrap();
let err = validate_and_decrypt(
&tampered,
&agent.public_key(),
&owner.public_key(),
agent.secret_key(),
&owner.public_key(),
)
.unwrap_err();
assert!(matches!(err, EngramError::InvalidEnvelope(_)));
}
#[test]
fn body_too_large_rejected_at_build_time() {
let agent = keys_from_hex(SECKEY_A);
let owner = keys_from_hex(SECKEY_O);
// Build a value whose JSON representation just barely exceeds the limit.
let huge = "a".repeat(NIP44_PLAINTEXT_MAX);
let body = Body::Memory {
slug: "mem/example".into(),
value: Some(huge),
};
let err = build_event(&agent, &owner.public_key(), &body, 1).unwrap_err();
assert!(matches!(err, EngramError::BodyTooLarge(_)));
}
}

View File

@@ -45,6 +45,13 @@ pub const KIND_HTTP_AUTH: u32 = 27235;
/// Agent metadata + owner reference (replaceable, agent-authored).
pub const KIND_AGENT_PROFILE: u32 = 10100;
/// NIP-AE: Agent Engram (parameterized replaceable, agent-authored).
///
/// Encrypted memory record for AI agents. Addressed by `(pubkey_a, kind, d_tag)`,
/// where `d_tag` is an HMAC over the agent↔owner conversation key. See
/// `docs/nips/NIP-AE.md` and [`crate::engram`].
pub const KIND_AGENT_ENGRAM: u32 = 30174;
// NIP-29 group admin events
/// NIP-29: Add a user to a group.
pub const KIND_NIP29_PUT_USER: u32 = 9000;
@@ -321,6 +328,7 @@ pub const ALL_KINDS: &[u32] = &[
KIND_GIFT_WRAP,
KIND_FILE_METADATA,
KIND_AGENT_PROFILE,
KIND_AGENT_ENGRAM,
KIND_NIP29_PUT_USER,
KIND_NIP29_REMOVE_USER,
KIND_NIP29_EDIT_METADATA,

View File

@@ -7,6 +7,9 @@
/// Channel and membership enums shared across crates.
pub mod channel;
/// NIP-AE Agent Engrams — slug grammar, conversation key, d-tag derivation,
/// body parse/serialize, envelope build/validate, head selection.
pub mod engram;
/// Relay-side error types.
pub mod error;
/// Relay-side event wrapper with verification tracking.

View File

@@ -200,6 +200,12 @@ pub async fn query_events(
"restricted: p-gated kinds require #p tag matching your pubkey",
));
}
if !crate::handlers::req::engram_filters_authorized(&filters, &authed_pubkey_hex) {
return Err(api_error(
StatusCode::FORBIDDEN,
"restricted: agent-engram reads require authors=[self] or #p=[self]",
));
}
// Get channels this user can access — same enforcement as WS REQ handler.
let accessible_channels = state
@@ -294,6 +300,12 @@ pub async fn count_events(
"restricted: p-gated kinds require #p tag matching your pubkey",
));
}
if !crate::handlers::req::engram_filters_authorized(&filters, &authed_pubkey_hex) {
return Err(api_error(
StatusCode::FORBIDDEN,
"restricted: agent-engram reads require authors=[self] or #p=[self]",
));
}
// Get channels this user can access.
let accessible_channels = state
@@ -380,6 +392,34 @@ pub async fn count_events(
// ── NIP-50 search via HTTP bridge ────────────────────────────────────────────
/// Decide whether a search hit should be returned to the caller.
///
/// Mirrors the WS NIP-50 path's post-filter step in `handlers/req.rs`:
/// Typesense receives only the kind/authors/time pushdown, so any other filter
/// constraint (`#p`, `#h`, `#e`, `#d`, `ids`, …) must be enforced here against
/// the full stored event. Without this, an authorized engram search such as
/// `{"kinds":[30174],"#p":[self]}` would leak text-matching envelopes whose
/// `#p` belongs to a different owner — the NIP-AE read gate at the filter
/// layer would be bypassed for `/query`.
///
/// `accessible_channels` is the caller's channel scope; channel-scoped hits
/// outside that set are rejected regardless of NIP-01 match.
fn search_hit_accepted(
filter: &nostr::Filter,
stored: &sprout_core::StoredEvent,
accessible_channels: &[uuid::Uuid],
) -> bool {
if !sprout_core::filter::filters_match(std::slice::from_ref(filter), stored) {
return false;
}
if let Some(ch_id) = stored.channel_id {
if !accessible_channels.contains(&ch_id) {
return false;
}
}
true
}
/// Handle search filters by routing to Typesense, then fetching full events from DB.
/// Returns first page of results (no pagination for bridge MVP).
async fn handle_bridge_search(
@@ -499,11 +539,8 @@ async fn handle_bridge_search(
Some(ev) => ev,
None => continue,
};
// Channel access post-filter.
if let Some(ch_id) = stored.channel_id {
if !accessible_channels.contains(&ch_id) {
continue;
}
if !search_hit_accepted(filter, stored, accessible_channels) {
continue;
}
// Dedup across filters.
if !seen_ids.insert(id_array) {
@@ -735,3 +772,107 @@ async fn synthesize_presence(state: &AppState, filters: &[nostr::Filter]) -> Opt
Some(events)
}
#[cfg(test)]
mod tests {
use super::*;
use nostr::{Alphabet, EventBuilder, Keys, Kind, SingleLetterTag, Tag};
/// Build a kind:30174 engram envelope authored by `agent`, tagged with `owner`.
fn engram_envelope(agent: &Keys, owner_hex: &str) -> sprout_core::StoredEvent {
let d_tag = Tag::custom(
nostr::TagKind::SingleLetter(SingleLetterTag::lowercase(Alphabet::D)),
["abcd1234"],
);
let p_tag = Tag::custom(
nostr::TagKind::SingleLetter(SingleLetterTag::lowercase(Alphabet::P)),
[owner_hex],
);
let ev = EventBuilder::new(Kind::Custom(30174), "engram body", [d_tag, p_tag])
.sign_with_keys(agent)
.expect("sign engram");
sprout_core::StoredEvent::new(ev, None)
}
/// Regression test for the NIP-AE `/query` search leak (PR #593 review).
///
/// Setup: two engram envelopes by different agents for different owners.
/// An authorized search for `{kinds:[30174], #p:[owner_a]}` would be
/// approved by the engram gate (owner_a is querying engrams addressed to
/// them). Typesense's pushdown only carries `kind:=[30174]`, so the
/// envelope for owner_b can come back as a text-match hit. The post-filter
/// in `search_hit_accepted` must reject it.
#[test]
fn search_hit_rejects_envelope_with_mismatched_p_tag() {
let agent_a = Keys::generate();
let agent_b = Keys::generate();
let owner_a = Keys::generate().public_key().to_hex();
let owner_b = Keys::generate().public_key().to_hex();
let env_for_a = engram_envelope(&agent_a, &owner_a);
let env_for_b = engram_envelope(&agent_b, &owner_b);
let p_tag = SingleLetterTag::lowercase(Alphabet::P);
let filter = nostr::Filter::new()
.kind(Kind::Custom(30174))
.custom_tag(p_tag, [&owner_a]);
assert!(
search_hit_accepted(&filter, &env_for_a, &[]),
"envelope addressed to owner_a must be returned"
);
assert!(
!search_hit_accepted(&filter, &env_for_b, &[]),
"envelope addressed to owner_b must NOT be returned for a #p=[owner_a] search"
);
}
/// `authors=[agent_a]` search must not return an envelope authored by agent_b,
/// even if Typesense's text match would otherwise surface it. (Typesense does
/// carry an `authors` pushdown today, so this is defence-in-depth; mirroring
/// the WS contract.)
#[test]
fn search_hit_rejects_event_with_mismatched_author() {
let agent_a = Keys::generate();
let agent_b = Keys::generate();
let owner = Keys::generate().public_key().to_hex();
let env_a = engram_envelope(&agent_a, &owner);
let env_b = engram_envelope(&agent_b, &owner);
let filter = nostr::Filter::new()
.kind(Kind::Custom(30174))
.author(agent_a.public_key());
assert!(search_hit_accepted(&filter, &env_a, &[]));
assert!(
!search_hit_accepted(&filter, &env_b, &[]),
"authors=[agent_a] search must not return events authored by agent_b"
);
}
/// Channel-scoped events outside the caller's accessible-channel set are
/// rejected by the post-filter regardless of NIP-01 match.
#[test]
fn search_hit_rejects_inaccessible_channel() {
let agent = Keys::generate();
let owner = Keys::generate().public_key().to_hex();
let mut stored = engram_envelope(&agent, &owner);
let scoped_channel = uuid::Uuid::new_v4();
stored.channel_id = Some(scoped_channel);
let p_tag = SingleLetterTag::lowercase(Alphabet::P);
let filter = nostr::Filter::new()
.kind(Kind::Custom(30174))
.custom_tag(p_tag, [&owner]);
assert!(
!search_hit_accepted(&filter, &stored, &[]),
"channel-scoped hit must be rejected when caller has no channel access"
);
assert!(
search_hit_accepted(&filter, &stored, &[scoped_channel]),
"channel-scoped hit must be accepted when caller has access to that channel"
);
}
}

View File

@@ -65,12 +65,27 @@ pub async fn ws_audio_handler(
// ── Auth message shape ────────────────────────────────────────────────────────
/// Highest huddle audio protocol version this relay understands. Clients are
/// allowed to negotiate any version in `1..=CURRENT_PROTOCOL_VERSION`; older
/// versions stay supported indefinitely for staged rollouts.
const CURRENT_PROTOCOL_VERSION: u8 = 2;
#[derive(Deserialize)]
struct AuthMsg {
#[serde(rename = "type")]
msg_type: String,
event: nostr::Event,
parent_channel_id: Option<Uuid>,
/// Huddle audio protocol version requested by the client. Defaults to 1
/// when missing so existing clients keep working without recompile. A
/// room is pinned to whichever version its first peer requested; later
/// peers must match or get `upgrade_required`.
#[serde(default = "default_protocol_version")]
protocol_version: u8,
}
fn default_protocol_version() -> u8 {
1
}
// ── Core connection lifecycle ─────────────────────────────────────────────────
@@ -208,16 +223,81 @@ async fn handle_audio_connection(socket: WebSocket, state: Arc<AppState>, channe
Ok(_) => {} // Channel exists and is not archived — proceed.
}
let (peer_id, peer_index, audio_rx, peer_ctrl_rx) = match room.add_peer(pubkey_hex.clone()) {
Some(v) => v,
None => {
// Reject unsupported future versions up-front so we don't accidentally
// pin a room to a version we can't speak. Versions 1..=CURRENT are OK.
let requested_version = auth_msg.protocol_version;
if requested_version == 0 || requested_version > CURRENT_PROTOCOL_VERSION {
warn!(
channel_id = %channel_id,
pubkey = %pubkey_hex,
requested_version,
current = CURRENT_PROTOCOL_VERSION,
"audio: client requested unsupported protocol version"
);
let _ = ws_send
.send(WsMessage::Text(
serde_json::json!({
"type": "error",
"code": "unsupported_version",
"message": format!(
"huddle audio protocol v{requested_version} not supported; relay max is v{CURRENT_PROTOCOL_VERSION}"
),
"current_version": CURRENT_PROTOCOL_VERSION,
})
.to_string()
.into(),
))
.await;
return;
}
let (peer_id, peer_index, audio_rx, peer_ctrl_rx) = match room
.add_peer(pubkey_hex.clone(), requested_version)
{
Ok(v) => v,
Err(crate::audio::room::AdmissionError::Full) => {
warn!(channel_id = %channel_id, "audio room full (255 peers exhausted)");
let _ = ws_send
.send(WsMessage::Text(
serde_json::json!({"type":"error","code":"room_full","message":"peer index space exhausted"})
.to_string().into(),
))
.await;
.send(WsMessage::Text(
serde_json::json!({"type":"error","code":"room_full","message":"peer index space exhausted"})
.to_string().into(),
))
.await;
return;
}
Err(crate::audio::room::AdmissionError::Ended) => {
debug!(channel_id = %channel_id, "room ended before admission");
let _ = ws_send
.send(WsMessage::Text(
serde_json::json!({"type":"error","code":"room_ended","message":"huddle has ended"})
.to_string().into(),
))
.await;
return;
}
Err(crate::audio::room::AdmissionError::VersionMismatch { pinned, requested }) => {
info!(
channel_id = %channel_id,
pubkey = %pubkey_hex,
pinned,
requested,
"audio: protocol version mismatch — upgrade required"
);
let _ = ws_send
.send(WsMessage::Text(
serde_json::json!({
"type": "error",
"code": "upgrade_required",
"message": format!(
"this huddle is using audio protocol v{pinned}; your client requested v{requested}"
),
"pinned_version": pinned,
"requested_version": requested,
})
.to_string()
.into(),
))
.await;
return;
}
};
@@ -288,6 +368,7 @@ async fn handle_audio_connection(socket: WebSocket, state: Arc<AppState>, channe
ws_recv,
Arc::clone(&room),
peer_id,
requested_version,
ctrl_tx,
Arc::clone(&missed_pongs),
cancel.clone(),
@@ -364,10 +445,13 @@ async fn recv_loop(
mut ws_recv: futures_util::stream::SplitStream<WebSocket>,
room: Arc<crate::audio::room::Room>,
peer_id: Uuid,
protocol_version: u8,
ctrl_tx: mpsc::Sender<WsMessage>,
missed_pongs: Arc<AtomicU8>,
cancel: CancellationToken,
) {
use crate::audio::wire::{FrameHeader, V2_HEADER_LEN};
loop {
tokio::select! {
biased;
@@ -379,6 +463,52 @@ async fn recv_loop(
warn!(peer_id = %peer_id, bytes = data.len(), "audio frame too large — dropping");
continue;
}
// Protocol v2 sanity-parse: validate the header is
// present and well-shaped, then forward opaquely.
// We never strip, rewrite, or re-encode bytes — the
// header is sender-authored telemetry only — but we
// do refuse to broadcast frames that are clearly
// malformed for the room's pinned protocol so we
// don't help v2 peers feed garbage to other v2 peers.
if protocol_version >= 2 {
// Frame must carry at least the 8-byte header
// plus a non-empty Opus payload.
if data.len() <= V2_HEADER_LEN {
warn!(
peer_id = %peer_id,
bytes = data.len(),
"v2 frame missing header or payload — dropping"
);
continue;
}
match FrameHeader::parse(&data) {
Some((header, payload)) if !payload.is_empty() => {
// Header is well-formed. `level_dbov` is
// already clamped by `parse` — bad values
// do not drop the frame, they just lose
// the metric (which the relay does not
// trust for anything anyway).
tracing::trace!(
peer_id = %peer_id,
seq = header.seq,
ts_48k = header.ts_48k,
level_dbov = header.level_dbov,
is_dtx = header.is_dtx(),
"v2 audio frame"
);
}
_ => {
warn!(
peer_id = %peer_id,
bytes = data.len(),
"v2 frame failed header parse — dropping"
);
continue;
}
}
}
room.broadcast_frame(peer_id, data);
}
Some(Ok(WsMessage::Text(text))) => {

View File

@@ -11,6 +11,7 @@
pub mod handler;
pub mod room;
pub mod wire;
pub use handler::ws_audio_handler;
pub use room::AudioRoomManager;

View File

@@ -47,6 +47,24 @@ const CTRL_CHANNEL_CAPACITY: usize = 32;
/// is reasonable. The 255 index space is the hard limit; this is the soft one.
const MAX_PEERS_PER_ROOM: usize = 25;
/// Reason a peer was refused entry to a room.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AdmissionError {
/// The room has been ended (or is shutting down) and no longer admits peers.
Ended,
/// The room has hit the soft peer cap or exhausted the 255-index space.
Full,
/// The room is pinned to a different protocol version than the requested one.
/// The caller should reply to the WS client with an `upgrade_required` error
/// and the room's actual `pinned` version, then close the socket.
VersionMismatch {
/// Version the room is currently pinned to.
pinned: u8,
/// Version the joining client requested.
requested: u8,
},
}
/// Peer index allocator + room lifecycle gate.
///
/// The `ended` flag and peer admission are synchronized under the same mutex.
@@ -58,6 +76,25 @@ struct AdmissionGuard {
next_fresh: u8,
free: Vec<u8>,
ended: bool,
/// Pinned huddle audio protocol version for this room.
///
/// `None` until the first peer admits. The first admission pins the room
/// to that version; subsequent peers MUST present the same version or
/// they're rejected with `AdmissionError::VersionMismatch`. The relay
/// forwards binary frames opaquely either way, so allowing a v1 client
/// into a v2 room would silently corrupt v2 peers' decode (they'd see
/// no header where one is expected, and vice versa).
///
/// Pin is per-`Room`-instance and clears when the manager evicts the
/// Room via [`AudioRoomManager::cleanup_if_empty`] — the next
/// `get_or_create` for the same channel id then constructs a fresh
/// `Room` with a fresh `AdmissionGuard` (and therefore `None` pin),
/// so a new generation of joiners can negotiate a new version.
/// A momentarily-empty-but-not-yet-cleaned-up Room keeps its pin so
/// reconnecting peers don't accidentally renegotiate mid-call. See
/// `version_pin_persists_across_peer_churn` for the test that pins
/// this behavior.
pinned_version: Option<u8>,
}
impl AdmissionGuard {
@@ -66,6 +103,7 @@ impl AdmissionGuard {
next_fresh: 0,
free: Vec::new(),
ended: false,
pinned_version: None,
}
}
@@ -126,26 +164,53 @@ impl Room {
}
}
/// Add a peer. Returns `(peer_id, peer_index, audio_rx, ctrl_rx)`, or
/// `None` if the room has ended, hit the peer cap, or exhausted the index space.
/// Add a peer. Returns `(peer_id, peer_index, audio_rx, ctrl_rx)` on
/// success, or an [`AdmissionError`] explaining why the peer was rejected.
///
/// The ended check, index allocation, and peer insert all happen under
/// the admission guard lock — mutually exclusive with `mark_ended`.
/// `requested_version` is the huddle audio protocol version the peer
/// negotiated in its WS auth message. The first successful admission
/// pins the room to that version; later admits must match the pin or
/// they receive [`AdmissionError::VersionMismatch`].
///
/// The cap check, ended check, version pin, index allocation, and peer
/// insert all happen under the admission guard lock — mutually exclusive
/// with `mark_ended` and with any concurrent `add_peer` that might race
/// the version pin.
///
/// Error precedence is deliberate: `Ended` > `Full` > `VersionMismatch`.
/// A "no seat available" error wins over version mismatch because a
/// client that couldn't join either way shouldn't learn the room's
/// pinned protocol version — that's a (mild) information leak. The cap
/// check lives inside the lock so two concurrent joiners can't both
/// pass it; the per-room index space (255) plus the soft cap
/// (`MAX_PEERS_PER_ROOM`) is then a single, race-free invariant.
pub fn add_peer(
&self,
pubkey: String,
) -> Option<(Uuid, u8, mpsc::Receiver<Bytes>, mpsc::Receiver<PeerCtrl>)> {
if self.peers.len() >= MAX_PEERS_PER_ROOM {
return None;
}
// Hold the guard across ended check + index alloc + peer insert.
// This makes add_peer mutually exclusive with mark_ended — the lock
// is the single synchronization point that closes the race.
let mut g = self.guard.lock().ok()?;
requested_version: u8,
) -> Result<(Uuid, u8, mpsc::Receiver<Bytes>, mpsc::Receiver<PeerCtrl>), AdmissionError> {
let mut g = self.guard.lock().map_err(
|_| AdmissionError::Ended, /* poisoned ≈ shutting down */
)?;
if g.ended {
return None;
return Err(AdmissionError::Ended);
}
let peer_index = g.alloc()?;
if self.peers.len() >= MAX_PEERS_PER_ROOM {
return Err(AdmissionError::Full);
}
if let Some(pinned) = g.pinned_version {
if pinned != requested_version {
return Err(AdmissionError::VersionMismatch {
pinned,
requested: requested_version,
});
}
}
let peer_index = g.alloc().ok_or(AdmissionError::Full)?;
// Pin the room version on the first successful index allocation. We
// pin *after* alloc so a Full error doesn't accidentally set the
// version for a peer that didn't actually join.
g.pinned_version.get_or_insert(requested_version);
let peer_id = Uuid::new_v4();
let (audio_tx, audio_rx) = mpsc::channel(AUDIO_CHANNEL_CAPACITY);
let (ctrl_tx, ctrl_rx) = mpsc::channel(CTRL_CHANNEL_CAPACITY);
@@ -159,7 +224,7 @@ impl Room {
},
);
drop(g); // Release lock after insert.
Some((peer_id, peer_index, audio_rx, ctrl_rx))
Ok((peer_id, peer_index, audio_rx, ctrl_rx))
}
/// Remove a peer and recycle its index.
@@ -288,3 +353,154 @@ impl Default for AudioRoomManager {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn fresh_room() -> Room {
Room::new(Uuid::new_v4())
}
/// First peer's `requested_version` becomes the room's pin; later peers
/// requesting the same version are admitted normally.
#[test]
fn first_admit_pins_version_and_matching_admits_succeed() {
let room = fresh_room();
let first = room
.add_peer("alice".to_string(), 2)
.expect("first peer admits");
let second = room
.add_peer("bob".to_string(), 2)
.expect("matching version admits");
assert_eq!(room.peers.len(), 2);
// peer_index allocation is monotonic from 0 inside a fresh room.
assert_eq!(first.1, 0);
assert_eq!(second.1, 1);
}
/// A peer requesting a different protocol version than the pinned one
/// is refused with `VersionMismatch` and never appears in the peer map.
#[test]
fn admit_rejects_mismatched_version() {
let room = fresh_room();
let _ = room
.add_peer("alice".to_string(), 2)
.expect("first peer admits");
let err = room
.add_peer("bob".to_string(), 1)
.expect_err("mismatched version must be rejected");
match err {
AdmissionError::VersionMismatch { pinned, requested } => {
assert_eq!(pinned, 2);
assert_eq!(requested, 1);
}
other => panic!("expected VersionMismatch, got {other:?}"),
}
// The rejected peer must not appear in the peer map, and the index
// space must not have been consumed by the failed admit.
assert_eq!(room.peers.len(), 1);
}
/// `add_peer` rejects requests after the room is marked ended even if the
/// version matches.
#[test]
fn admit_after_mark_ended_returns_ended() {
let room = fresh_room();
assert!(room.mark_ended());
let err = room
.add_peer("alice".to_string(), 1)
.expect_err("ended room must refuse");
assert!(matches!(err, AdmissionError::Ended));
}
/// Per Max's review checklist: an empty-and-cleaned-up room (via the
/// manager's `cleanup_if_empty`) becomes a fresh room on the next
/// `get_or_create`, with no pin carried over.
#[test]
fn manager_cleanup_resets_version_pin() {
let manager = AudioRoomManager::new();
let channel_id = Uuid::new_v4();
let room1 = manager.get_or_create(channel_id);
let (peer_id, _, _, _) = room1
.add_peer("alice".to_string(), 2)
.expect("first peer admits");
// Last peer leaves and ends the room atomically.
let (_, ended) = room1
.remove_peer_and_check_ended(peer_id)
.expect("peer existed");
assert!(ended, "single-peer room should end on its last departure");
assert!(manager.cleanup_if_empty(channel_id));
// Next joiner with a different version on the same channel id gets a
// brand-new room (no v=2 pin carried over from the prior generation).
let room2 = manager.get_or_create(channel_id);
let _ = room2
.add_peer("bob".to_string(), 1)
.expect("fresh room must accept any version");
}
/// Peer-index reuse: after a peer leaves, their index is released; a new
/// peer joining the same (still-pinned) room reuses the freed index.
/// Version pin must persist across this reuse — the room generation
/// hasn't ended.
#[test]
fn version_pin_persists_across_peer_churn() {
let room = fresh_room();
let (alice_id, alice_idx, _, _) =
room.add_peer("alice".to_string(), 2).expect("alice admits");
room.remove_peer(alice_id);
// Room is non-empty thanks to nothing yet — wait, alice left and
// nobody else is here. Add bob with the same version: should work.
// Then add carol with a different version: should fail with the
// *original* pin, even though alice already left.
let (_, bob_idx, _, _) = room
.add_peer("bob".to_string(), 2)
.expect("bob admits at v=2");
assert_eq!(
bob_idx, alice_idx,
"freed peer index should be recycled by the next admit",
);
let err = room
.add_peer("carol".to_string(), 1)
.expect_err("v=1 must still be refused — room is pinned v=2");
assert!(matches!(
err,
AdmissionError::VersionMismatch {
pinned: 2,
requested: 1
}
));
}
/// Per Sami/Perci's review: when a room is both at-capacity AND the
/// joiner's protocol version doesn't match the pin, the error must be
/// `Full` — not `VersionMismatch`. A client that couldn't get a seat
/// either way shouldn't learn the room's pinned protocol version.
/// This also pins the in-lock cap check (the old code's outside-lock
/// cap check meant the version error could win in a race).
#[test]
fn admit_full_wins_over_version_mismatch() {
let room = fresh_room();
// Fill the room with v=2 peers right up to the soft cap.
for i in 0..MAX_PEERS_PER_ROOM {
room.add_peer(format!("peer-{i}"), 2)
.expect("seed admit must succeed");
}
// Next joiner is BOTH over the cap AND requests the wrong version.
let err = room
.add_peer("over-cap-and-wrong-version".to_string(), 1)
.expect_err("over-cap + wrong-version joiner must be rejected");
assert!(
matches!(err, AdmissionError::Full),
"expected Full to win over VersionMismatch, got {err:?}",
);
// And the room state must be unchanged.
assert_eq!(room.peers.len(), MAX_PEERS_PER_ROOM);
}
}

View File

@@ -0,0 +1,168 @@
//! Huddle audio wire protocol — relay-side parse helpers.
//!
//! Mirrors the desktop client's `huddle::wire` for the half of the protocol
//! the relay cares about: validating that v2 frames carry an 8-byte header,
//! parsing it into structured metrics, and clamping
//! caller-authored telemetry to safe ranges. The relay never *generates*
//! v2 frames (clients author them) and never *re-encodes* them (broadcast
//! is opaque byte forwarding) — so only the parse half of the protocol
//! lives here, not the encode half.
//!
//! # Threat-model invariant
//!
//! `level_dbov` is client-authored telemetry. Anything we surface from it
//! (logs, active-speaker UI hints, dominant-talker decisions) must treat
//! it as untrusted. This module's [`FrameHeader::parse`] clamps out-of-range
//! values into the canonical `-127..=0` range but never drops the audio
//! frame on bad VU data — bad metadata must not cause audible loss.
//! Trust decisions (admission, moderation, kicks) MUST NOT consume
//! `level_dbov`.
/// Length of the v2 per-frame header in bytes. The wire layout is, in
/// network byte order:
///
/// ```text
/// byte 0..=1 : seq u16
/// byte 2..=5 : ts_48k u32
/// byte 6 : level_dbov i8 range [-127, 0]
/// byte 7 : flags u8 bit 0 = DTX; other bits reserved
/// ```
pub const V2_HEADER_LEN: usize = 8;
/// `flags & FLAG_DTX` indicates a DTX/comfort-noise frame.
pub const FLAG_DTX: u8 = 0x01;
/// Parsed v2 header view. Cheap (Copy).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FrameHeader {
/// Sender-authored sequence number; wraps every 2^16 frames.
pub seq: u16,
/// Sender-authored 48 kHz RTP-style media timestamp.
pub ts_48k: u32,
/// Audio level in dBov. Always parsed into the canonical `-127..=0`
/// range — out-of-range inputs are clamped to `-127` (silence floor).
/// Untrusted; for diagnostics only.
pub level_dbov: i8,
/// Raw flags byte. Use bit masks; reserved bits MAY be set.
pub flags: u8,
}
impl FrameHeader {
/// True if `FLAG_DTX` is set.
pub fn is_dtx(&self) -> bool {
(self.flags & FLAG_DTX) != 0
}
/// Parse a v2 header from the leading 8 bytes of `bytes`. Returns the
/// header and the remaining payload slice (the opaque Opus body the
/// relay forwards). On `None`, the caller should treat the frame as
/// malformed and drop it.
///
/// `level_dbov` is clamped into `-127..=0`; out-of-range inputs become
/// `-127`. The audio frame is never rejected for bad telemetry — only
/// the metric is suppressed.
pub fn parse(bytes: &[u8]) -> Option<(Self, &[u8])> {
if bytes.len() < V2_HEADER_LEN {
return None;
}
let seq = u16::from_be_bytes([bytes[0], bytes[1]]);
let ts_48k = u32::from_be_bytes([bytes[2], bytes[3], bytes[4], bytes[5]]);
let raw_level = bytes[6] as i8;
let level_dbov = if (-127..=0).contains(&raw_level) {
raw_level
} else {
-127
};
let flags = bytes[7];
Some((
Self {
seq,
ts_48k,
level_dbov,
flags,
},
&bytes[V2_HEADER_LEN..],
))
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Canonical layout: BE u16 seq, BE u32 ts_48k, i8 level, u8 flags.
/// This test pins the byte order so an accidental endianness flip on
/// either side of the protocol is caught immediately.
#[test]
fn parse_reads_network_byte_order() {
let bytes = [
0x01, 0x02, // seq = 0x0102
0x03, 0x04, 0x05, 0x06, // ts_48k = 0x0304_0506
0xFF, // level_dbov = -1
0x01, // flags = FLAG_DTX
0xAA, 0xBB, 0xCC, // opaque payload — what the relay forwards
];
let (h, payload) = FrameHeader::parse(&bytes).expect("parse");
assert_eq!(h.seq, 0x0102);
assert_eq!(h.ts_48k, 0x0304_0506);
assert_eq!(h.level_dbov, -1);
assert!(h.is_dtx());
assert_eq!(payload, &[0xAA, 0xBB, 0xCC]);
}
/// Anything shorter than the fixed 8-byte header is malformed.
#[test]
fn parse_rejects_short_input() {
for len in 0..V2_HEADER_LEN {
let buf = vec![0u8; len];
assert!(
FrameHeader::parse(&buf).is_none(),
"{len}-byte input must fail"
);
}
}
/// Out-of-range `level_dbov` must not drop the frame. The metric is
/// suppressed (clamped to -127, the silence floor) but the audio
/// payload is preserved for forwarding. This pins the "bad VU
/// metadata is not audible loss" invariant relay-side as well as
/// desktop-side.
#[test]
fn parse_clamps_out_of_range_level_keeps_frame() {
// i8 = +127 → outside [-127, 0]
let bytes = [
0x00, 0x07, // seq = 7
0x00, 0x00, 0x03, 0xC0, // ts_48k = 960
0x7F, // level_dbov raw = +127 (invalid)
0x00, // flags
b'o', b'p', b'u', b's',
];
let (h, payload) = FrameHeader::parse(&bytes).expect("parse must succeed");
assert_eq!(h.level_dbov, -127, "invalid level clamps to silence floor");
assert_eq!(h.seq, 7, "valid fields preserved alongside clamp");
assert_eq!(
payload, b"opus",
"audio payload still available for forwarding"
);
}
/// Reserved flag bits are passed through untouched. Receivers (the
/// other half of the protocol) are responsible for ignoring them;
/// the relay's only job is to forward the bytes faithfully.
#[test]
fn parse_preserves_reserved_flag_bits() {
let bytes = [
0x00,
0x00,
0x00,
0x00,
0x00,
0x00,
0x00,
0b1010_1010, // FLAG_DTX clear, reserved bits set
];
let (h, _) = FrameHeader::parse(&bytes).expect("parse");
assert_eq!(h.flags, 0b1010_1010);
assert!(!h.is_dtx());
}
}

View File

@@ -54,6 +54,13 @@ pub async fn handle_count(
));
return;
}
if !super::req::engram_filters_authorized(&filters, &authed_pubkey_hex) {
conn.send(RelayMessage::closed(
&sub_id,
"restricted: agent-engram reads require authors=[self] or #p=[self]",
));
return;
}
// Get channels this user can access — same enforcement as WS REQ handler.
let accessible_channels = match state.get_accessible_channel_ids_cached(&pubkey_bytes).await {

View File

@@ -12,23 +12,23 @@ use uuid::Uuid;
use nostr::Event;
use sprout_auth::Scope;
use sprout_core::kind::{
event_kind_u32, is_parameterized_replaceable, is_relay_admin_kind, KIND_APPROVAL_DENY,
KIND_APPROVAL_GRANT, KIND_AUTH, KIND_CANVAS, KIND_CONTACT_LIST, KIND_DELETION,
KIND_DM_ADD_MEMBER, KIND_DM_HIDE, KIND_DM_OPEN, KIND_FORUM_COMMENT, KIND_FORUM_POST,
KIND_FORUM_VOTE, KIND_GIFT_WRAP, KIND_GIT_ISSUE, KIND_GIT_PATCH, KIND_GIT_PR_UPDATE,
KIND_GIT_PULL_REQUEST, KIND_GIT_REPO_ANNOUNCEMENT, KIND_GIT_REPO_STATE, KIND_GIT_STATUS_CLOSED,
KIND_GIT_STATUS_DRAFT, KIND_GIT_STATUS_MERGED, KIND_GIT_STATUS_OPEN, KIND_HUDDLE_ENDED,
KIND_HUDDLE_GUIDELINES, KIND_HUDDLE_PARTICIPANT_JOINED, KIND_HUDDLE_PARTICIPANT_LEFT,
KIND_HUDDLE_RECORDING_AVAILABLE, KIND_HUDDLE_STARTED, KIND_HUDDLE_TRACK_PUBLISHED,
KIND_LONG_FORM, KIND_MEMBER_ADDED_NOTIFICATION, KIND_MEMBER_REMOVED_NOTIFICATION,
KIND_NIP29_CREATE_GROUP, KIND_NIP29_DELETE_EVENT, KIND_NIP29_DELETE_GROUP,
KIND_NIP29_EDIT_METADATA, KIND_NIP29_JOIN_REQUEST, KIND_NIP29_LEAVE_REQUEST,
KIND_NIP29_PUT_USER, KIND_NIP29_REMOVE_USER, KIND_NIP43_LEAVE_REQUEST, KIND_PRESENCE_UPDATE,
KIND_PROFILE, KIND_REACTION, KIND_READ_STATE, KIND_STREAM_MESSAGE,
KIND_STREAM_MESSAGE_BOOKMARKED, KIND_STREAM_MESSAGE_DIFF, KIND_STREAM_MESSAGE_EDIT,
KIND_STREAM_MESSAGE_PINNED, KIND_STREAM_MESSAGE_SCHEDULED, KIND_STREAM_MESSAGE_V2,
KIND_STREAM_REMINDER, KIND_TEXT_NOTE, KIND_USER_STATUS, KIND_WORKFLOW_DEF,
KIND_WORKFLOW_TRIGGER, RELAY_ADMIN_ADD_MEMBER, RELAY_ADMIN_CHANGE_ROLE,
event_kind_u32, is_parameterized_replaceable, is_relay_admin_kind, KIND_AGENT_ENGRAM,
KIND_APPROVAL_DENY, KIND_APPROVAL_GRANT, KIND_AUTH, KIND_CANVAS, KIND_CONTACT_LIST,
KIND_DELETION, KIND_DM_ADD_MEMBER, KIND_DM_HIDE, KIND_DM_OPEN, KIND_FORUM_COMMENT,
KIND_FORUM_POST, KIND_FORUM_VOTE, KIND_GIFT_WRAP, KIND_GIT_ISSUE, KIND_GIT_PATCH,
KIND_GIT_PR_UPDATE, KIND_GIT_PULL_REQUEST, KIND_GIT_REPO_ANNOUNCEMENT, KIND_GIT_REPO_STATE,
KIND_GIT_STATUS_CLOSED, KIND_GIT_STATUS_DRAFT, KIND_GIT_STATUS_MERGED, KIND_GIT_STATUS_OPEN,
KIND_HUDDLE_ENDED, KIND_HUDDLE_GUIDELINES, KIND_HUDDLE_PARTICIPANT_JOINED,
KIND_HUDDLE_PARTICIPANT_LEFT, KIND_HUDDLE_RECORDING_AVAILABLE, KIND_HUDDLE_STARTED,
KIND_HUDDLE_TRACK_PUBLISHED, KIND_LONG_FORM, KIND_MEMBER_ADDED_NOTIFICATION,
KIND_MEMBER_REMOVED_NOTIFICATION, KIND_NIP29_CREATE_GROUP, KIND_NIP29_DELETE_EVENT,
KIND_NIP29_DELETE_GROUP, KIND_NIP29_EDIT_METADATA, KIND_NIP29_JOIN_REQUEST,
KIND_NIP29_LEAVE_REQUEST, KIND_NIP29_PUT_USER, KIND_NIP29_REMOVE_USER,
KIND_NIP43_LEAVE_REQUEST, KIND_PRESENCE_UPDATE, KIND_PROFILE, KIND_REACTION, KIND_READ_STATE,
KIND_STREAM_MESSAGE, KIND_STREAM_MESSAGE_BOOKMARKED, KIND_STREAM_MESSAGE_DIFF,
KIND_STREAM_MESSAGE_EDIT, KIND_STREAM_MESSAGE_PINNED, KIND_STREAM_MESSAGE_SCHEDULED,
KIND_STREAM_MESSAGE_V2, KIND_STREAM_REMINDER, KIND_TEXT_NOTE, KIND_USER_STATUS,
KIND_WORKFLOW_DEF, KIND_WORKFLOW_TRIGGER, RELAY_ADMIN_ADD_MEMBER, RELAY_ADMIN_CHANGE_ROLE,
RELAY_ADMIN_REMOVE_MEMBER,
};
use sprout_core::verification::verify_event;
@@ -150,7 +150,9 @@ fn required_scope_for_kind(kind: u32, event: &Event) -> Result<Scope, &'static s
match kind {
KIND_PROFILE => Ok(Scope::UsersWrite),
KIND_TEXT_NOTE | KIND_LONG_FORM => Ok(Scope::MessagesWrite),
KIND_CONTACT_LIST | KIND_READ_STATE | KIND_USER_STATUS => Ok(Scope::UsersWrite),
KIND_CONTACT_LIST | KIND_READ_STATE | KIND_USER_STATUS | KIND_AGENT_ENGRAM => {
Ok(Scope::UsersWrite)
}
KIND_DELETION
| KIND_REACTION
| KIND_GIFT_WRAP
@@ -300,6 +302,8 @@ pub(crate) fn is_global_only_kind(kind: u32) -> bool {
| KIND_LONG_FORM
| KIND_USER_STATUS
| KIND_READ_STATE
// NIP-AE agent engrams are addressed by (pubkey_a, kind, d_tag); never channel-scoped.
| KIND_AGENT_ENGRAM
// NIP-34: git events use `a` tags (repo reference), not `h` tags (channel scope).
// Parameterized replaceable kinds are keyed by (pubkey, kind, d_tag).
| KIND_GIT_REPO_ANNOUNCEMENT
@@ -767,6 +771,143 @@ fn validate_diff_event(event: &Event) -> Result<(), String> {
Ok(())
}
/// Validate the public envelope of a NIP-AE `kind:30174` event before it
/// reaches NIP-33 parameterized replacement.
///
/// We deliberately do this here (not in the d-tag length check downstream)
/// because a malformed envelope can otherwise *replace* a valid head in
/// storage and then be invisible to readers querying `#p`. The relay sees
/// no plaintext, but it can — and must — enforce the public tag shape:
///
/// * exactly one `d` tag with a 64-hex value (`d_tag = lower_hex(HMAC...)`),
/// * exactly one `p` tag with a 64-hex pubkey (the owner counterparty).
///
/// Content is opaque NIP-44 ciphertext; we do not parse it.
fn validate_engram_envelope(event: &Event) -> Result<(), String> {
let mut d_tags: Vec<&str> = Vec::new();
let mut p_tags: Vec<&str> = Vec::new();
for tag in event.tags.iter() {
let parts = tag.as_slice();
if parts.len() < 2 {
continue;
}
match parts[0].as_str() {
"d" => d_tags.push(&parts[1]),
"p" => p_tags.push(&parts[1]),
_ => {}
}
}
if d_tags.len() != 1 {
return Err(format!(
"agent-engram event must have exactly one `d` tag (got {})",
d_tags.len()
));
}
if p_tags.len() != 1 {
return Err(format!(
"agent-engram event must have exactly one `p` tag (got {})",
p_tags.len()
));
}
let d = d_tags[0];
if d.len() != 64
|| !d
.bytes()
.all(|b| b.is_ascii_hexdigit() && !b.is_ascii_uppercase())
{
return Err("agent-engram `d` tag must be 64 lowercase hex chars".to_string());
}
let p = p_tags[0];
// Lowercase-only: readers query `#p` with `owner.to_hex()` (lowercase) and
// Nostr tag matching is byte-exact. Accepting uppercase here would let a
// submitter replace the lowercase head with an event that subsequent
// lowercase-`#p` queries cannot see — silently bricking the slug.
if p.len() != 64
|| !p
.bytes()
.all(|b| b.is_ascii_hexdigit() && !b.is_ascii_uppercase())
{
return Err("agent-engram `p` tag must be 64 lowercase hex chars (pubkey)".to_string());
}
// Content must be a syntactically plausible NIP-44 v2 payload. We do not
// (and cannot) verify the MAC at the relay, but we can reject obvious
// garbage so a malformed event cannot supersede a valid head via NIP-33
// replacement and then be silently discarded by readers.
validate_engram_nip44_content(&event.content)?;
Ok(())
}
/// Validate that `content` is a syntactically plausible NIP-44 v2 ciphertext.
///
/// Checks:
/// - Non-empty.
/// - Standard base64 alphabet only (A-Z, a-z, 0-9, +, /, =), with padding only
/// at the end and total length a multiple of 4.
/// - Decoded length >= 99 bytes (1 version + 32 nonce + 32 MAC + minimum 34
/// bytes of length-prefixed padded ciphertext required by NIP-44 v2).
/// - First decoded byte is `0x02` (NIP-44 version 2).
///
/// This is an envelope sanity check, not full validation: the MAC and actual
/// decryption happen at the reader. The intent is to refuse obvious junk so a
/// malformed event cannot win NIP-33 replacement against a valid head and then
/// be silently skipped by `validate_and_decrypt`. Mirrors the validator in
/// `sprout-pair-relay::validate_nip44_content`.
fn validate_engram_nip44_content(content: &str) -> Result<(), String> {
if content.is_empty() {
return Err("agent-engram content must not be empty (NIP-44 ciphertext)".to_string());
}
let bytes = content.as_bytes();
let len = bytes.len();
if !len.is_multiple_of(4) {
return Err("agent-engram content is not valid base64 (length)".to_string());
}
let mut pad_count = 0usize;
for (i, &b) in bytes.iter().enumerate() {
match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'+' | b'/' => {
if pad_count > 0 {
return Err("agent-engram content is not valid base64".to_string());
}
}
b'=' => {
if i < len - 2 {
return Err("agent-engram content is not valid base64".to_string());
}
pad_count += 1;
if pad_count > 2 {
return Err("agent-engram content is not valid base64".to_string());
}
}
_ => return Err("agent-engram content is not valid base64".to_string()),
}
}
let decoded_len = (len / 4) * 3 - pad_count;
if decoded_len < 99 {
return Err("agent-engram content too short for NIP-44 v2".to_string());
}
let b64_val = |c: u8| -> Option<u8> {
match c {
b'A'..=b'Z' => Some(c - b'A'),
b'a'..=b'z' => Some(c - b'a' + 26),
b'0'..=b'9' => Some(c - b'0' + 52),
b'+' => Some(62),
b'/' => Some(63),
_ => None,
}
};
let v0 =
b64_val(bytes[0]).ok_or_else(|| "agent-engram content is not valid base64".to_string())?;
let v1 =
b64_val(bytes[1]).ok_or_else(|| "agent-engram content is not valid base64".to_string())?;
let first_byte = (v0 << 2) | (v1 >> 4);
if first_byte != 0x02 {
return Err(
"agent-engram content is not NIP-44 v2 (expected 0x02 version prefix)".to_string(),
);
}
Ok(())
}
// ── The pipeline ─────────────────────────────────────────────────────────────
/// Ingest a signed Nostr event through the full validation pipeline.
@@ -1164,6 +1305,12 @@ pub async fn ingest_event(
validate_diff_event(&event).map_err(|e| IngestError::Rejected(format!("invalid: {e}")))?;
}
// ── 15a. Agent engram envelope (kind:30174) ──────────────────────────
if kind_u32 == KIND_AGENT_ENGRAM {
validate_engram_envelope(&event)
.map_err(|e| IngestError::Rejected(format!("invalid: {e}")))?;
}
// Track pre-created channel UUID for compensation on insert failure.
let mut pre_created_channel: Option<Uuid> = None;
@@ -1693,6 +1840,7 @@ mod tests {
KIND_FORUM_COMMENT,
KIND_LONG_FORM,
KIND_USER_STATUS,
KIND_AGENT_ENGRAM,
];
for kind in migrated {
assert!(
@@ -1842,4 +1990,147 @@ mod tests {
let event = make_event_with_tags(5, "", &[&["e", "a".repeat(64).as_str()]]);
assert_eq!(count_e_tags(&event), 1);
}
// ── NIP-AE envelope validation ───────────────────────────────────────
fn make_engram(tags: &[&[&str]], content: &str) -> Event {
make_event_with_tags(KIND_AGENT_ENGRAM, content, tags)
}
/// Minimal syntactically-plausible NIP-44 v2 payload (99 zero-filled bytes
/// with the 0x02 version prefix). Real ciphertexts are larger and have real
/// MACs; the relay only checks shape, not authenticity.
fn fake_nip44_v2() -> String {
// base64(b"\x02" + b"\x00" * 98) — 132 chars, decoded length 99,
// first byte 0x02.
let mut s = String::from("Ag");
s.push_str(&"A".repeat(130));
s
}
#[test]
fn engram_envelope_accepts_canonical() {
let d = "a".repeat(64);
let p = "b".repeat(64);
let ev = make_engram(&[&["d", &d], &["p", &p]], &fake_nip44_v2());
assert!(validate_engram_envelope(&ev).is_ok());
}
#[test]
fn engram_envelope_rejects_missing_p() {
let d = "a".repeat(64);
let ev = make_engram(&[&["d", &d]], &fake_nip44_v2());
let err = validate_engram_envelope(&ev).unwrap_err();
assert!(err.contains("`p` tag"), "got: {err}");
}
#[test]
fn engram_envelope_rejects_duplicate_p() {
let d = "a".repeat(64);
let p = "b".repeat(64);
let ev = make_engram(&[&["d", &d], &["p", &p], &["p", &p]], &fake_nip44_v2());
let err = validate_engram_envelope(&ev).unwrap_err();
assert!(err.contains("`p` tag"), "got: {err}");
}
#[test]
fn engram_envelope_rejects_short_d() {
let p = "b".repeat(64);
let ev = make_engram(&[&["d", "abcd"], &["p", &p]], &fake_nip44_v2());
let err = validate_engram_envelope(&ev).unwrap_err();
assert!(err.contains("`d` tag"), "got: {err}");
}
#[test]
fn engram_envelope_rejects_uppercase_d() {
let p = "b".repeat(64);
// 64 chars but uppercase — spec mandates lowercase hex.
let d = "A".repeat(64);
let ev = make_engram(&[&["d", &d], &["p", &p]], &fake_nip44_v2());
let err = validate_engram_envelope(&ev).unwrap_err();
assert!(err.contains("`d` tag"), "got: {err}");
}
/// Regression: uppercase `p` tag must be rejected at ingest. Readers query
/// `#p` lowercase; an uppercase-tagged event that wins NIP-33 replacement
/// becomes invisible to readers, silently bricking the slug.
#[test]
fn engram_envelope_rejects_uppercase_p() {
let d = "a".repeat(64);
let p = "B".repeat(64);
let ev = make_engram(&[&["d", &d], &["p", &p]], &fake_nip44_v2());
let err = validate_engram_envelope(&ev).unwrap_err();
assert!(err.contains("`p` tag"), "got: {err}");
}
#[test]
fn engram_envelope_rejects_short_p() {
let d = "a".repeat(64);
let ev = make_engram(&[&["d", &d], &["p", "abcd"]], &fake_nip44_v2());
let err = validate_engram_envelope(&ev).unwrap_err();
assert!(err.contains("`p` tag"), "got: {err}");
}
#[test]
fn engram_envelope_rejects_empty_content() {
let d = "a".repeat(64);
let p = "b".repeat(64);
let ev = make_engram(&[&["d", &d], &["p", &p]], "");
let err = validate_engram_envelope(&ev).unwrap_err();
assert!(err.contains("content"), "got: {err}");
}
/// Regression: non-base64 content must be rejected. Otherwise a signed
/// event with `content="x"` wins NIP-33 replacement against a valid head,
/// and the new head is then skipped by `validate_and_decrypt` — making the
/// slug appear absent to readers.
#[test]
fn engram_envelope_rejects_non_base64_content() {
let d = "a".repeat(64);
let p = "b".repeat(64);
let ev = make_engram(&[&["d", &d], &["p", &p]], "x");
let err = validate_engram_envelope(&ev).unwrap_err();
assert!(
err.contains("base64") || err.contains("too short"),
"got: {err}"
);
}
#[test]
fn engram_envelope_rejects_wrong_nip44_version() {
// 99 bytes of valid base64 alphabet, but first byte decodes to 0x00,
// not the NIP-44 v2 prefix 0x02. Length OK (132 chars / 99 decoded).
let d = "a".repeat(64);
let p = "b".repeat(64);
let bad = "A".repeat(132);
let ev = make_engram(&[&["d", &d], &["p", &p]], &bad);
let err = validate_engram_envelope(&ev).unwrap_err();
assert!(
err.contains("NIP-44 v2") || err.contains("0x02"),
"got: {err}"
);
}
#[test]
fn engram_envelope_rejects_short_content() {
// Base64 of "Ag==" decodes to 1 byte — version prefix correct but
// way under the 99-byte floor.
let d = "a".repeat(64);
let p = "b".repeat(64);
let ev = make_engram(&[&["d", &d], &["p", &p]], "Ag==");
let err = validate_engram_envelope(&ev).unwrap_err();
assert!(err.contains("too short"), "got: {err}");
}
#[test]
fn engram_envelope_rejects_bad_base64_alphabet() {
let d = "a".repeat(64);
let p = "b".repeat(64);
// Contains '!' which is not in the standard base64 alphabet. Length is
// a multiple of 4 to defeat the length check.
let bad = format!("Ag!!{}", "A".repeat(128));
let ev = make_engram(&[&["d", &d], &["p", &p]], &bad);
let err = validate_engram_envelope(&ev).unwrap_err();
assert!(err.contains("base64"), "got: {err}");
}
}

View File

@@ -9,7 +9,7 @@ use hex;
use nostr::Filter;
use sprout_core::filter::filters_match;
use sprout_core::kind::{
KIND_AGENT_OBSERVER_FRAME, KIND_GIFT_WRAP, KIND_MEMBER_ADDED_NOTIFICATION,
KIND_AGENT_ENGRAM, KIND_AGENT_OBSERVER_FRAME, KIND_GIFT_WRAP, KIND_MEMBER_ADDED_NOTIFICATION,
KIND_MEMBER_REMOVED_NOTIFICATION,
};
use sprout_db::EventQuery;
@@ -90,11 +90,38 @@ pub async fn handle_req(
let channel_id = extract_channel_id_from_filters(&filters);
// ── NIP-50 search: intercept BEFORE #p gating ────────────────────────────
// Search filters are one-shot (not registered as persistent subscriptions).
// They never deliver gift wraps (not indexed) or membership notifications
// (global, no channel_id), so the #p gate below is irrelevant for them.
// Intercepting here lets clients send `{"search":"foo"}` without `kinds`.
// ── #p / engram gating for globally-stored sensitive kinds ───────────────
// Applied BEFORE the NIP-50 search branch so that an authenticated member
// cannot use `{"search":"...","kinds":[30174]}` (or similar for p-gated
// kinds) to harvest indexed-but-globally-stored sensitive events. Search
// hits are looked up by event id and returned without the per-filter
// post-check the historical-delivery branch applies, so the gate must run
// here, up front. Only applies to GLOBAL subscriptions (channel_id = None):
// channel-scoped subs can never receive globally-stored events because of
// the fan_out() invariant in subscription.rs.
if channel_id.is_none() {
let authed_pubkey_hex = hex::encode(&pubkey_bytes);
if !p_gated_filters_authorized(&filters, &authed_pubkey_hex) {
conn.send(RelayMessage::closed(
&sub_id,
"restricted: p-gated events require #p matching your pubkey",
));
return;
}
if !engram_filters_authorized(&filters, &authed_pubkey_hex) {
conn.send(RelayMessage::closed(
&sub_id,
"restricted: agent-engram reads require authors=[self] or #p=[self]",
));
return;
}
}
// ── NIP-50 search: one-shot, no persistent subscription ──────────────────
// Search filters hit Typesense and return historical hits, then EOSE.
// They are not registered for fan-out. The sensitive-kind gates above
// already ran, so an authed member cannot use search to bypass author/#p
// rules for kind:30174 or other globally-stored gated kinds.
let has_search = filters.iter().any(|f| f.search.is_some());
if has_search {
if filters.iter().any(|f| f.search.is_none()) {
@@ -116,21 +143,6 @@ pub async fn handle_req(
return;
}
// ── #p gating for globally-stored sensitive kinds ─────────────────────────
// Only applies to GLOBAL subscriptions (channel_id = None). Channel-scoped
// subscriptions can never receive globally-stored events — the fan_out()
// invariant in subscription.rs prevents it.
if channel_id.is_none() {
let authed_pubkey_hex = hex::encode(&pubkey_bytes);
if !p_gated_filters_authorized(&filters, &authed_pubkey_hex) {
conn.send(RelayMessage::closed(
&sub_id,
"restricted: p-gated events require #p matching your pubkey",
));
return;
}
}
// Check channel access BEFORE registering the subscription.
if let Some(ch_id) = channel_id {
if !accessible_channels.contains(&ch_id) {
@@ -718,6 +730,59 @@ pub(crate) fn p_gated_filters_authorized(filters: &[Filter], authed_pubkey_hex:
})
}
/// Authorize read access for filters that can match KIND_AGENT_ENGRAM events.
///
/// NIP-AE engrams are global (no channel scope) and have encrypted content,
/// but their public `#p` (owner) and timestamps still leak who-pairs-with-whom
/// plus write-activity patterns. Only the agent (the event's author) or the
/// owner (the `#p` value) should be able to enumerate them.
///
/// A filter is authorized when at least one of:
/// - `authors` is non-empty and every entry equals the authed pubkey
/// (the agent reading its own engrams), OR
/// - `#p` is non-empty and every entry equals the authed pubkey
/// (the owner reading engrams addressed to them).
///
/// Filters with explicit `ids` are exempt — knowing the event id already
/// implies authorization (the engram event id is itself derived from the
/// signed envelope, which only the agent could have produced).
///
/// Mixed-kind filters (e.g. `{kinds:[30174, 9]}`) are evaluated under this
/// gate when KIND_AGENT_ENGRAM is present; matching events of other kinds in
/// the same filter is also restricted, but that is the conservative choice
/// — clients should query engrams in a dedicated filter.
pub(crate) fn engram_filters_authorized(filters: &[Filter], authed_pubkey_hex: &str) -> bool {
let p_tag = nostr::SingleLetterTag::lowercase(nostr::Alphabet::P);
filters.iter().all(|filter| {
// Specific-event lookups don't fish.
if filter.ids.as_ref().is_some_and(|ids| !ids.is_empty()) {
return true;
}
let can_match_engram = filter
.kinds
.as_ref()
.is_none_or(|ks| ks.iter().any(|k| k.as_u16() as u32 == KIND_AGENT_ENGRAM));
if !can_match_engram {
return true;
}
let authors_ok = filter.authors.as_ref().is_some_and(|authors| {
!authors.is_empty()
&& authors
.iter()
.all(|a| a.to_hex().eq_ignore_ascii_case(authed_pubkey_hex))
});
if authors_ok {
return true;
}
filter.generic_tags.get(&p_tag).is_some_and(|values| {
!values.is_empty() && values.iter().all(|v| v == authed_pubkey_hex)
})
})
}
#[cfg(test)]
mod tests {
use super::*;
@@ -881,4 +946,163 @@ mod tests {
"restricted tokens must not fall back to global search results"
);
}
// ── NIP-AE engram read gating ────────────────────────────────────────
/// Three real x-only pubkeys (valid for `PublicKey::from_hex`). Distinct,
/// so we can label them clearly in tests.
fn three_pubkeys() -> (String, String, String) {
let agent = nostr::Keys::generate().public_key().to_hex();
let owner = nostr::Keys::generate().public_key().to_hex();
let attacker = nostr::Keys::generate().public_key().to_hex();
(agent, owner, attacker)
}
#[test]
fn engram_gate_allows_agent_querying_own() {
let (agent, owner, _) = three_pubkeys();
let p_tag = SingleLetterTag::lowercase(Alphabet::P);
let f = Filter::new()
.kind(nostr::Kind::Custom(KIND_AGENT_ENGRAM as u16))
.author(nostr::PublicKey::from_hex(&agent).unwrap())
.custom_tag(p_tag, [&owner]);
assert!(engram_filters_authorized(&[f], &agent));
}
#[test]
fn engram_gate_allows_owner_querying() {
let (agent, owner, _) = three_pubkeys();
let p_tag = SingleLetterTag::lowercase(Alphabet::P);
// Owner-side read: knows the agent's pubkey, queries with #p=self.
let f = Filter::new()
.kind(nostr::Kind::Custom(KIND_AGENT_ENGRAM as u16))
.author(nostr::PublicKey::from_hex(&agent).unwrap())
.custom_tag(p_tag, [&owner]);
assert!(engram_filters_authorized(&[f], &owner));
}
#[test]
fn engram_gate_allows_owner_with_no_authors_filter() {
// Owner doesn't necessarily know the agent's pubkey ahead of time.
let (_, owner, _) = three_pubkeys();
let p_tag = SingleLetterTag::lowercase(Alphabet::P);
let f = Filter::new()
.kind(nostr::Kind::Custom(KIND_AGENT_ENGRAM as u16))
.custom_tag(p_tag, [&owner]);
assert!(engram_filters_authorized(&[f], &owner));
}
#[test]
fn engram_gate_rejects_unrelated_reader() {
let (agent, owner, attacker) = three_pubkeys();
let p_tag = SingleLetterTag::lowercase(Alphabet::P);
// Attacker tries to fish for engrams between agent and owner.
let f = Filter::new()
.kind(nostr::Kind::Custom(KIND_AGENT_ENGRAM as u16))
.author(nostr::PublicKey::from_hex(&agent).unwrap())
.custom_tag(p_tag, [&owner]);
assert!(!engram_filters_authorized(&[f], &attacker));
}
#[test]
fn engram_gate_rejects_bare_kind_filter() {
// {kinds:[30174]} with no authors and no #p — open fishing.
let (agent, _, _) = three_pubkeys();
let f = Filter::new().kind(nostr::Kind::Custom(KIND_AGENT_ENGRAM as u16));
assert!(!engram_filters_authorized(&[f], &agent));
}
#[test]
fn engram_gate_rejects_wildcard_kind_filter() {
// Filter with no kinds field at all — matches everything including
// engrams; must still be gated.
let (agent, _, _) = three_pubkeys();
let f = Filter::new();
assert!(!engram_filters_authorized(&[f], &agent));
}
#[test]
fn engram_gate_skips_non_engram_kinds() {
// Filter not targeting engrams — pass through; this gate is silent.
let (agent, _, _) = three_pubkeys();
let f = Filter::new().kind(nostr::Kind::Custom(9));
assert!(engram_filters_authorized(&[f], &agent));
}
#[test]
fn engram_gate_allows_ids_lookup() {
// Specific event ids — knowing the id implies prior authorization.
let (agent, _, _) = three_pubkeys();
let id = nostr::EventId::from_hex(
"abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789",
)
.unwrap();
let f = Filter::new()
.kind(nostr::Kind::Custom(KIND_AGENT_ENGRAM as u16))
.id(id);
assert!(engram_filters_authorized(&[f], &agent));
}
#[test]
fn engram_gate_rejects_mixed_authors_with_unauthed() {
// {authors:[self, attacker]} — must reject; an author-list with any
// non-self entry could let an attacker piggy-back on the agent's
// legitimate query path.
let (agent, other, _) = three_pubkeys();
let f = Filter::new()
.kind(nostr::Kind::Custom(KIND_AGENT_ENGRAM as u16))
.authors([
nostr::PublicKey::from_hex(&agent).unwrap(),
nostr::PublicKey::from_hex(&other).unwrap(),
]);
assert!(!engram_filters_authorized(&[f], &agent));
}
// ── NIP-50 search bypass regressions ─────────────────────────────────
// These filters are the shape an authenticated relay member would send
// to try to harvest indexed engram envelopes via the search path. The
// gate must reject them regardless of the presence of `search`.
#[test]
fn engram_gate_rejects_bare_kind_search_filter() {
// {"search":"*", "kinds":[30174]} — exactly the bypass codex found.
let (agent, _, _) = three_pubkeys();
let f = Filter::new()
.kind(nostr::Kind::Custom(KIND_AGENT_ENGRAM as u16))
.search("*");
assert!(!engram_filters_authorized(&[f], &agent));
}
#[test]
fn engram_gate_rejects_wildcard_kind_search_filter() {
// {"search":"foo"} — no `kinds` field at all matches engrams too.
let (agent, _, _) = three_pubkeys();
let f = Filter::new().search("foo");
assert!(!engram_filters_authorized(&[f], &agent));
}
#[test]
fn engram_gate_allows_authored_engram_search() {
// Agent searching their own engrams by content keyword is legitimate.
let (agent, _, _) = three_pubkeys();
let f = Filter::new()
.kind(nostr::Kind::Custom(KIND_AGENT_ENGRAM as u16))
.author(nostr::PublicKey::from_hex(&agent).unwrap())
.search("foo");
assert!(engram_filters_authorized(&[f], &agent));
}
#[test]
fn p_gate_rejects_bare_kind_search_filter_for_gift_wrap() {
// P-gated kinds (observer frames, member notifications) are indexed
// too. Same bypass shape: {"search":"x","kinds":[<p-gated kind>]}.
// Use KIND_AGENT_OBSERVER_FRAME — globally stored, p-gated, indexed.
let (agent, _, _) = three_pubkeys();
let f = Filter::new()
.kind(nostr::Kind::Custom(
sprout_core::kind::KIND_AGENT_OBSERVER_FRAME as u16,
))
.search("x");
assert!(!p_gated_filters_authorized(&[f], &agent));
}
}

View File

@@ -66,13 +66,13 @@ const overrides = new Map([
["src/features/settings/ui/ChannelTemplatesSettingsCard.tsx", 850], // template CRUD card + TemplateFormDialog (persona/team chip selectors + provider assignments + canvas template) + TemplateTeamSelector + ProviderAssignments + ProviderRow
["src/shared/api/types.ts", 620], // ... + RespondToMode + respondTo/respondToAllowlist on ManagedAgent/Create/Update inputs
["src-tauri/src/events.rs", 610], // event builders + build_huddle_guidelines (kind:48106) + post_event_raw transport helper + participant p-tag on join/leave + NIP-43 relay admin builders (add/remove/change-role) + check_relay_role + DM/presence/workflow command builders
["src-tauri/src/huddle/kokoro.rs", 980], // Kokoro ONNX TTS engine + three-tier G2P + ARPAbet→IPA + CoreML + synth_chunk() public API + style validation + hyphenated compound splitting + 23 unit tests
["src-tauri/src/huddle/mod.rs", 1020], // huddle state machine + Tauri commands + sync protocol doc; state/relay/pipeline extracted + emit_huddle_state_changed wiring
["src-tauri/src/huddle/models.rs", 900], // model download manager for Parakeet TDT-CTC STT + Kokoro TTS with streaming downloads + SHA-256 verification + Rust-native tar extraction + version manifest + atomic swap + hot-start signaling + CC-BY-4.0 attribution sidecar + idempotent legacy Moonshine dir cleanup
["src-tauri/src/huddle/models.rs", 950], // model download manager for Parakeet TDT-CTC STT + Pocket TTS with streaming downloads + SHA-256 verification + Rust-native tar extraction + version manifest + atomic swap + hot-start signaling + MODEL_LICENSE.txt sidecar (fail-closed readiness) + idempotent legacy Moonshine dir cleanup + tts_readiness_requires_license_sidecar test + Mary (VCTK p333) reference voice attribution block
["src-tauri/src/huddle/stt.rs", 580], // STT pipeline + PTT edge-detection flush + PTT gating (is_speech AND ptt_active) + barge-in for VAD mode + rubato resampler + earshot VAD + sherpa-onnx transcription
["src-tauri/src/huddle/pocket.rs", 900], // Pocket TTS engine wrapper + prepare_pocket_prompt (capitalize/punctuate/pad short inputs + insert ". . " sacrificial cold-start prefix, mirror upstream pocket-tts prepare_text_prompt) + build_generation_extra (only overrides max_frames, and only for ≤4-word inputs, to bound runaway "monster breathing" generation without clipping multi-clause sentences) + trim_leading_cold_start (post-synth strip of sacrificial audio on short prompts, workaround for kyutai-labs/pocket-tts #91 first-phoneme drop) + 18 unit tests including regressions for the static-burst-on-multi-clause-sentence bug and the first-phoneme-drop trim hazards
["src-tauri/src/huddle/preprocessing.rs", 670], // TTS text preprocessing pipeline + unified split_sentences + int_to_words 0-999999 + URL trailing punctuation preservation + 23 unit tests
["src-tauri/src/huddle/relay_api.rs", 520], // audio relay recv task + per-peer frame counting for remote human TTS interrupt + NIP-98 channel member query
["src-tauri/src/huddle/tts.rs", 1030], // TTS pipeline + session warmup + cancel/shutdown handling + apply_fades + 18 unit tests for remote interrupt mechanism
["src-tauri/src/huddle/tts.rs", 1380], // TTS pipeline + session warmup + cancel/shutdown handling + apply_fade_out (fade-out only — leading fade removed 2026-05-18 after onset-attenuation regression measured in examples/pocket_onset_probe.rs) + FIRST_APPEND_LEAD_IN_SAMPLES + build_sentence_append_plan (pure helper enforcing the lead-in fires exactly once per utterance, not per sentence — see lead_in_pad_fires_exactly_once_per_utterance regression test) + normalize_for_playback (per-sentence peak normalization to -3 dBFS ceiling with MAX_GAIN cap) + 30 unit tests (18 interrupt + 5 fade-out + 1 first-append-lead-in + 3 build-sentence-append-plan + 6 normalize)
["src-tauri/src/relay.rs", 510], // +4 lines for NIP-OA auth tag injection in profile sync (build_profile_event) + verification test
["src-tauri/src/commands/pairing.rs", 600], // NIP-AB pairing actor: 3 Tauri commands + background WS task + NIP-42 auth + NIP-43 probe + event parsing helpers
["src-tauri/src/lib.rs", 715], // +4 lines for PairingHandle managed state + 3 pairing command registrations

View File

@@ -736,9 +736,9 @@ dependencies = [
[[package]]
name = "cc"
version = "1.2.60"
version = "1.2.62"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "43c5703da9466b66a946814e1adf53ea2c90f10063b86290cc9eb67ce3478a20"
checksum = "a1dce859f0832a7d088c4f1119888ab94ef4b5d6795d1ce05afb7fe159d79f98"
dependencies = [
"find-msvc-tools",
"jobserver",
@@ -956,7 +956,7 @@ dependencies = [
"bitflags 2.11.1",
"core-foundation 0.10.1",
"core-graphics-types",
"foreign-types 0.5.0",
"foreign-types",
"libc",
]
@@ -1066,25 +1066,6 @@ dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-deque"
version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51"
dependencies = [
"crossbeam-epoch",
"crossbeam-utils",
]
[[package]]
name = "crossbeam-epoch"
version = "0.9.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-utils"
version = "0.8.21"
@@ -1236,16 +1217,6 @@ version = "0.1.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ac6b926516df9c60bfa16e107b21086399f8285a44ca9711344b9e553c5146e2"
[[package]]
name = "der"
version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "71fd89660b2dc699704064e59e9dba0147b903e85319429e131620d022be411b"
dependencies = [
"pem-rfc7468",
"zeroize",
]
[[package]]
name = "deranged"
version = "0.5.8"
@@ -1474,12 +1445,6 @@ dependencies = [
"libm",
]
[[package]]
name = "either"
version = "1.15.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
[[package]]
name = "embed-resource"
version = "3.0.8"
@@ -1660,15 +1625,6 @@ version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
[[package]]
name = "foreign-types"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1"
dependencies = [
"foreign-types-shared 0.1.1",
]
[[package]]
name = "foreign-types"
version = "0.5.0"
@@ -1676,7 +1632,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965"
dependencies = [
"foreign-types-macros",
"foreign-types-shared 0.3.1",
"foreign-types-shared",
]
[[package]]
@@ -1690,12 +1646,6 @@ dependencies = [
"syn 2.0.117",
]
[[package]]
name = "foreign-types-shared"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b"
[[package]]
name = "foreign-types-shared"
version = "0.3.1"
@@ -2199,9 +2149,9 @@ dependencies = [
[[package]]
name = "hashbrown"
version = "0.17.0"
version = "0.17.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51"
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
[[package]]
name = "heck"
@@ -2251,12 +2201,6 @@ dependencies = [
"digest 0.10.7",
]
[[package]]
name = "hmac-sha256"
version = "1.1.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ec9d92d097f4749b64e8cc33d924d9f40a2d4eb91402b458014b781f5733d60f"
[[package]]
name = "html5ever"
version = "0.29.1"
@@ -2563,7 +2507,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9"
dependencies = [
"equivalent",
"hashbrown 0.17.0",
"hashbrown 0.17.1",
"serde",
"serde_core",
]
@@ -2812,9 +2756,9 @@ dependencies = [
[[package]]
name = "libc"
version = "0.2.185"
version = "0.2.186"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "52ff2c0fe9bc6cb6b14a0592c2ff4fa9ceb83eea9db979b0487cd054946a2b8f"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "libdbus-sys"
@@ -2896,12 +2840,6 @@ dependencies = [
"crc",
]
[[package]]
name = "lzma-rust2"
version = "0.15.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1670343e58806300d87950e3401e820b519b9384281bbabfb15e3636689ffd69"
[[package]]
name = "lzma-sys"
version = "0.1.20"
@@ -2988,16 +2926,6 @@ version = "0.8.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
[[package]]
name = "matrixmultiply"
version = "0.3.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08"
dependencies = [
"autocfg",
"rawpointer",
]
[[package]]
name = "memchr"
version = "2.8.0"
@@ -3067,39 +2995,6 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "native-tls"
version = "0.2.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "465500e14ea162429d264d44189adc38b199b62b1c21eea9f69e4b73cb03bbf2"
dependencies = [
"libc",
"log",
"openssl",
"openssl-probe",
"openssl-sys",
"schannel",
"security-framework",
"security-framework-sys",
"tempfile",
]
[[package]]
name = "ndarray"
version = "0.17.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d"
dependencies = [
"matrixmultiply",
"num-complex",
"num-integer",
"num-traits",
"portable-atomic",
"portable-atomic-util",
"rawpointer",
"rayon",
]
[[package]]
name = "ndk"
version = "0.9.0"
@@ -3142,6 +3037,19 @@ version = "0.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "43a88da9dd148bbcdce323dd6ac47d369b4769d4a3b78c6c52389b9269f77932"
[[package]]
name = "neteq"
version = "0.8.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "57fbe9601485d95750187c92344fa056e75055027b00e375c53d9f4b8b883112"
dependencies = [
"log",
"ringbuf",
"serde",
"thiserror 1.0.69",
"web-time",
]
[[package]]
name = "new_debug_unreachable"
version = "1.0.6"
@@ -3608,50 +3516,12 @@ dependencies = [
"pathdiff",
]
[[package]]
name = "openssl"
version = "0.10.77"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bfe4646e360ec77dff7dde40ed3d6c5fee52d156ef4a62f53973d38294dad87f"
dependencies = [
"bitflags 2.11.1",
"cfg-if",
"foreign-types 0.3.2",
"libc",
"once_cell",
"openssl-macros",
"openssl-sys",
]
[[package]]
name = "openssl-macros"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "openssl-probe"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe"
[[package]]
name = "openssl-sys"
version = "0.9.113"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ad2f2c0eba47118757e4c6d2bff2838f3e0523380021356e7875e858372ce644"
dependencies = [
"cc",
"libc",
"pkg-config",
"vcpkg",
]
[[package]]
name = "option-ext"
version = "0.2.0"
@@ -3687,30 +3557,6 @@ dependencies = [
"pin-project-lite",
]
[[package]]
name = "ort"
version = "2.0.0-rc.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d7de3af33d24a745ffb8fab904b13478438d1cd52868e6f17735ef6e1f8bf133"
dependencies = [
"ndarray",
"ort-sys",
"smallvec",
"tracing",
"ureq 3.3.0",
]
[[package]]
name = "ort-sys"
version = "2.0.0-rc.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d7b497d21a8b6fbb4b5a544f8fadb77e801a09ae0add9e411d31c6f89e3c1e90"
dependencies = [
"hmac-sha256",
"lzma-rust2",
"ureq 3.3.0",
]
[[package]]
name = "osakit"
version = "0.3.1"
@@ -3806,15 +3652,6 @@ dependencies = [
"hmac",
]
[[package]]
name = "pem-rfc7468"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a6305423e0e7738146434843d1694d621cce767262b2a86910beab705e4493d9"
dependencies = [
"base64ct",
]
[[package]]
name = "percent-encoding"
version = "2.3.2"
@@ -4087,21 +3924,6 @@ dependencies = [
"universal-hash",
]
[[package]]
name = "portable-atomic"
version = "1.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49"
[[package]]
name = "portable-atomic-util"
version = "0.2.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "091397be61a01d4be58e7841595bd4bfedb15f1cd54977d79b8271e94ed799a3"
dependencies = [
"portable-atomic",
]
[[package]]
name = "potential_utf"
version = "0.1.5"
@@ -4457,32 +4279,6 @@ version = "0.6.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539"
[[package]]
name = "rawpointer"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
[[package]]
name = "rayon"
version = "1.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d"
dependencies = [
"either",
"rayon-core",
]
[[package]]
name = "rayon-core"
version = "1.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
dependencies = [
"crossbeam-deque",
"crossbeam-utils",
]
[[package]]
name = "realfft"
version = "3.5.0"
@@ -4690,6 +4486,15 @@ dependencies = [
"windows-sys 0.52.0",
]
[[package]]
name = "ringbuf"
version = "0.3.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "79abed428d1fd2a128201cec72c5f6938e2da607c6f3745f769fabea399d950a"
dependencies = [
"crossbeam-utils",
]
[[package]]
name = "rodio"
version = "0.22.2"
@@ -4810,9 +4615,9 @@ dependencies = [
[[package]]
name = "rustls-pki-types"
version = "1.14.0"
version = "1.14.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd"
checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9"
dependencies = [
"web-time",
"zeroize",
@@ -5312,7 +5117,7 @@ checksum = "384b71e9586b28493902080ed89e9b4457d2c684521d4c10df62c381b6fdcb51"
dependencies = [
"bzip2 0.4.4",
"tar",
"ureq 2.12.1",
"ureq",
]
[[package]]
@@ -5371,17 +5176,6 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "socks"
version = "0.3.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b"
dependencies = [
"byteorder",
"libc",
"winapi",
]
[[package]]
name = "softbuffer"
version = "0.4.8"
@@ -5446,11 +5240,10 @@ dependencies = [
"hex",
"infer",
"libc",
"ndarray",
"neteq",
"nostr 0.36.0",
"nostr 0.37.0",
"opus",
"ort",
"png 0.18.1",
"regex",
"reqwest 0.13.2",
@@ -6967,36 +6760,6 @@ dependencies = [
"webpki-roots 0.26.11",
]
[[package]]
name = "ureq"
version = "3.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0"
dependencies = [
"base64 0.22.1",
"der",
"log",
"native-tls",
"percent-encoding",
"rustls-pki-types",
"socks",
"ureq-proto",
"utf8-zero",
"webpki-root-certs",
]
[[package]]
name = "ureq-proto"
version = "0.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c"
dependencies = [
"base64 0.22.1",
"http",
"httparse",
"log",
]
[[package]]
name = "url"
version = "2.5.8"
@@ -7028,12 +6791,6 @@ version = "0.7.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9"
[[package]]
name = "utf8-zero"
version = "0.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b8c0a043c9540bae7c578c88f91dda8bd82e59ae27c21baca69c8b191aaf5a6e"
[[package]]
name = "utf8_iter"
version = "1.0.4"
@@ -7052,12 +6809,6 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "vcpkg"
version = "0.2.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426"
[[package]]
name = "version-compare"
version = "0.2.1"
@@ -7134,11 +6885,11 @@ checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
[[package]]
name = "wasip2"
version = "1.0.2+wasi-0.2.9"
version = "1.0.3+wasi-0.2.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9517f9239f02c069db75e65f174b3da828fe5f5b945c4dd26bd25d89c03ebcf5"
checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6"
dependencies = [
"wit-bindgen",
"wit-bindgen 0.57.1",
]
[[package]]
@@ -7147,7 +6898,7 @@ version = "0.4.0+wasi-0.3.0-rc-2026-01-06"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5"
dependencies = [
"wit-bindgen",
"wit-bindgen 0.51.0",
]
[[package]]
@@ -7330,9 +7081,9 @@ dependencies = [
[[package]]
name = "webpki-root-certs"
version = "1.0.6"
version = "1.0.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "804f18a4ac2676ffb4e8b5b5fa9ae38af06df08162314f96a68d2a363e21a8ca"
checksum = "f31141ce3fc3e300ae89b78c0dd67f9708061d1d2eda54b8209346fd6be9a92c"
dependencies = [
"rustls-pki-types",
]
@@ -7942,6 +7693,12 @@ dependencies = [
"wit-bindgen-rust-macro",
]
[[package]]
name = "wit-bindgen"
version = "0.57.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e"
[[package]]
name = "wit-bindgen-core"
version = "0.51.0"

View File

@@ -1,3 +1,5 @@
[workspace]
[package]
name = "sprout"
version = "0.1.0"
@@ -42,6 +44,7 @@ tokio-tungstenite = { version = "0.29", features = ["rustls-tls-webpki-roots"] }
tokio-util = { version = "0.7", features = ["rt"] }
futures-util = "0.3"
opus = "0.3"
neteq = { version = "0.8", default-features = false }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
nostr = { version = "0.37", features = ["nip44"] }
@@ -63,8 +66,6 @@ uuid = { version = "1", features = ["v4"] }
png = "0.18"
zip = "2"
sherpa-onnx = "1.12"
ort = { version = "=2.0.0-rc.12", default-features = false, features = ["std", "ndarray", "tracing", "download-binaries", "tls-native", "copy-dylibs", "api-23", "coreml"] }
ndarray = { version = "0.17", features = ["rayon"] }
regex = "1"
axum = "0.8"
rodio = "0.22"

View File

@@ -0,0 +1,117 @@
//! Cold-vs-warm latency bench for Pocket TTS.
//!
//! This duplicates the small config-building snippet from `huddle::pocket` so it
//! doesn't depend on changing module visibility for a one-off dev tool.
//! Keep in sync with `huddle::pocket::load_text_to_speech`.
//!
//! Run with the model files in a directory (defaults to /tmp/pocket-tts-bench):
//! cargo run --release --example pocket_bench
//! cargo run --release --example pocket_bench /path/to/pocket-tts
use std::path::PathBuf;
use std::time::Instant;
use sherpa_onnx::{
self, GenerationConfig, OfflineTts, OfflineTtsConfig, OfflineTtsModelConfig,
OfflineTtsPocketModelConfig, Wave,
};
const SAMPLE_RATE: u32 = 24_000;
const TEST_TEXT: &str =
"Hello, this is a test of the new Pocket TTS engine running on sherpa-onnx.";
fn main() {
let model_dir = std::env::args()
.nth(1)
.unwrap_or_else(|| "/tmp/pocket-tts-bench".to_string());
println!("Model dir: {model_dir}");
let dir = PathBuf::from(&model_dir);
let p = |name: &str| dir.join(name).to_string_lossy().into_owned();
let t0 = Instant::now();
let mut cfg = OfflineTtsConfig::default();
cfg.model = OfflineTtsModelConfig {
pocket: OfflineTtsPocketModelConfig {
lm_main: Some(p("lm_main.int8.onnx")),
lm_flow: Some(p("lm_flow.int8.onnx")),
encoder: Some(p("encoder.onnx")),
decoder: Some(p("decoder.int8.onnx")),
text_conditioner: Some(p("text_conditioner.onnx")),
vocab_json: Some(p("vocab.json")),
token_scores_json: Some(p("token_scores.json")),
voice_embedding_cache_capacity: 16,
},
num_threads: 1,
debug: false,
..Default::default()
};
let engine = OfflineTts::create(&cfg).expect("engine create");
let load_ms = t0.elapsed().as_secs_f32() * 1000.0;
println!("Engine load: {load_ms:.1} ms");
let t0 = Instant::now();
let voice_path = dir.join("reference_sample.wav");
let wave = Wave::read(voice_path.to_str().unwrap()).expect("voice WAV");
let samples = wave.samples().to_vec();
let sr = wave.sample_rate();
let voice_ms = t0.elapsed().as_secs_f32() * 1000.0;
println!("Voice load: {voice_ms:.1} ms");
let gen = || GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 0.0,
reference_audio: Some(samples.clone()),
reference_sample_rate: sr,
..Default::default()
};
// Cold synth — first call after engine creation.
let t0 = Instant::now();
let cold = engine
.generate_with_config(TEST_TEXT, &gen(), None::<fn(&[f32], f32) -> bool>)
.expect("cold synth");
let cold_ms = t0.elapsed().as_secs_f32() * 1000.0;
let cold_audio_ms = (cold.samples().len() as f32 / SAMPLE_RATE as f32) * 1000.0;
let cold_rtf_x = cold_audio_ms / cold_ms;
println!(
"Cold synth: {cold_ms:.1} ms → {cold_audio_ms:.1} ms audio → {cold_rtf_x:.2}× realtime"
);
// Warm synth — second call.
let t0 = Instant::now();
let warm = engine
.generate_with_config(TEST_TEXT, &gen(), None::<fn(&[f32], f32) -> bool>)
.expect("warm synth");
let warm_ms = t0.elapsed().as_secs_f32() * 1000.0;
let warm_audio_ms = (warm.samples().len() as f32 / SAMPLE_RATE as f32) * 1000.0;
let warm_rtf_x = warm_audio_ms / warm_ms;
println!(
"Warm synth: {warm_ms:.1} ms → {warm_audio_ms:.1} ms audio → {warm_rtf_x:.2}× realtime"
);
// Write warm output for listening.
let out_path = "/tmp/pocket_bench_out.wav";
let ok = sherpa_onnx::write(out_path, warm.samples(), SAMPLE_RATE as i32);
println!(
"Wrote {} ({} samples, ok={ok})",
out_path,
warm.samples().len()
);
let delta_ms = cold_ms - warm_ms;
let delta_pct = (delta_ms / warm_ms) * 100.0;
println!();
println!("Cold/warm delta: {delta_ms:+.1} ms ({delta_pct:+.1}%)");
println!(
"Decision: warmup {}.",
if delta_ms > 200.0 {
"RECOMMENDED — significant cold-call penalty"
} else if delta_ms > 50.0 {
"OPTIONAL — small cold-call penalty"
} else {
"UNNECESSARY — cold and warm essentially equal"
}
);
}

View File

@@ -0,0 +1,148 @@
//! Onset-attenuation probe for Pocket TTS.
//!
//! Synthesises a handful of short sentences and dumps per-sentence onset
//! statistics (samples[0], 1ms/5ms/20ms peak + RMS) so we can decide whether
//! the production `apply_fades` 8 ms fade-in is masking real audio.
//!
//! Also writes the raw (un-faded, un-normalised) audio of each sentence to
//! /tmp so they can be inspected in Audacity / aplay without rodio in the
//! loop.
//!
//! Run with model files in /tmp/pocket-tts-bench (override with arg 1):
//! cargo run --release --example pocket_onset_probe
//! cargo run --release --example pocket_onset_probe /path/to/pocket-tts
use std::path::PathBuf;
use sherpa_onnx::{
self, GenerationConfig, OfflineTts, OfflineTtsConfig, OfflineTtsModelConfig,
OfflineTtsPocketModelConfig, Wave,
};
const SAMPLE_RATE: u32 = 24_000;
/// Test prompts chosen to span different onsets:
/// - palatal glide 'Y' (soft onset)
/// - voiceless fricative 'H' (very soft onset)
/// - labio-velar glide 'W' (medium onset)
/// - voiceless stop 'T' (hard onset)
const PROMPTS: &[&str] = &[
"Yep, I can hear you.",
"Hello there friend.",
"What can I help with?",
"Try this experiment now.",
];
fn main() {
let model_dir = std::env::args()
.nth(1)
.unwrap_or_else(|| "/tmp/pocket-tts-bench".to_string());
eprintln!("Model dir: {model_dir}");
let dir = PathBuf::from(&model_dir);
let p = |name: &str| dir.join(name).to_string_lossy().into_owned();
let mut cfg = OfflineTtsConfig::default();
cfg.model = OfflineTtsModelConfig {
pocket: OfflineTtsPocketModelConfig {
lm_main: Some(p("lm_main.int8.onnx")),
lm_flow: Some(p("lm_flow.int8.onnx")),
encoder: Some(p("encoder.onnx")),
decoder: Some(p("decoder.int8.onnx")),
text_conditioner: Some(p("text_conditioner.onnx")),
vocab_json: Some(p("vocab.json")),
token_scores_json: Some(p("token_scores.json")),
voice_embedding_cache_capacity: 16,
},
num_threads: 1,
debug: false,
..Default::default()
};
let engine = OfflineTts::create(&cfg).expect("engine create");
let voice_path = dir.join("reference_sample.wav");
let wave = Wave::read(voice_path.to_str().unwrap()).expect("voice WAV");
let voice_samples = wave.samples().to_vec();
let voice_sr = wave.sample_rate();
// Warmup so we're not measuring cold-call jitter.
{
let cfg = GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 0.0,
reference_audio: Some(voice_samples.clone()),
reference_sample_rate: voice_sr,
..Default::default()
};
let _ = engine.generate_with_config("warmup.", &cfg, None::<fn(&[f32], f32) -> bool>);
}
println!(
"{:<28} | {:>10} | {:>10} {:>10} | {:>10} {:>10} | {:>10} {:>10}",
"prompt",
"samples[0]",
"peak@1ms",
"rms@1ms",
"peak@5ms",
"rms@5ms",
"peak@20ms",
"rms@20ms"
);
println!("{}", "-".repeat(120));
for prompt in PROMPTS {
// Mirror the production prompt-prep (capitalise + terminal punctuation).
// These prompts already have it, so this is just to match what
// sherpa-onnx sees in production.
let cfg = GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 0.0,
reference_audio: Some(voice_samples.clone()),
reference_sample_rate: voice_sr,
..Default::default()
};
let out = engine
.generate_with_config(prompt, &cfg, None::<fn(&[f32], f32) -> bool>)
.expect("synth");
let samples = out.samples();
let n_1ms = (SAMPLE_RATE as f32 * 0.001) as usize;
let n_5ms = (SAMPLE_RATE as f32 * 0.005) as usize;
let n_20ms = (SAMPLE_RATE as f32 * 0.020) as usize;
let stats = |range: &[f32]| -> (f32, f32) {
if range.is_empty() {
return (0.0, 0.0);
}
let peak = range.iter().fold(0.0_f32, |a, &x| a.max(x.abs()));
let sumsq: f32 = range.iter().map(|x| x * x).sum();
let rms = (sumsq / range.len() as f32).sqrt();
(peak, rms)
};
let first = samples.first().copied().unwrap_or(0.0);
let (p1, r1) = stats(&samples[..n_1ms.min(samples.len())]);
let (p5, r5) = stats(&samples[..n_5ms.min(samples.len())]);
let (p20, r20) = stats(&samples[..n_20ms.min(samples.len())]);
println!(
"{:<28} | {:>10.6} | {:>10.6} {:>10.6} | {:>10.6} {:>10.6} | {:>10.6} {:>10.6}",
prompt, first, p1, r1, p5, r5, p20, r20
);
// Dump raw WAV for inspection.
let safe: String = prompt
.chars()
.map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
.collect();
let out_path = format!("/tmp/pocket_onset_{}.wav", &safe[..safe.len().min(24)]);
let _ = sherpa_onnx::write(&out_path, samples, SAMPLE_RATE as i32);
eprintln!(
" → wrote {out_path} ({} samples = {:.3} s)",
samples.len(),
samples.len() as f32 / SAMPLE_RATE as f32
);
}
}

View File

@@ -0,0 +1,218 @@
//! Reproduction probe for the Pocket TTS first-phoneme-drop fix.
//!
//! Pocket TTS' FlowLM has an autoregressive cold-start that occasionally
//! smears or drops the first phoneme of short utterances — see
//! kyutai-labs/pocket-tts #91, #70 and sherpa-onnx #3180. The cure is to
//! prepend a sacrificial `". . "` cold-start absorber to short prompts
//! and trim the resulting leading audio. This example reproduces both
//! variants of generated audio so you can listen-test the fix at the
//! exact `GenerationConfig` we ship in production (`huddle::pocket`):
//!
//! - silence_scale: 0.0 (production)
//! - max_frames: 100 (short) / sherpa default 500 (long)
//! - num_steps: 1
//! - speed: 1.05
//!
//! Note: production does NOT override `frames_after_eos` — sherpa-onnx's
//! default of 3 is what we want. The previous attempt to override it for
//! long prompts caused the "first 'yep' is static" regression (commit
//! 1dbfa2c). This probe mirrors that decision.
//!
//! Run:
//! cargo run --release --example prod_probe
//! cargo run --release --example prod_probe /path/to/pocket-tts
//!
//! Output (per (label, seed) pair):
//! /tmp/prod_<label>_s<seed>_raw.wav — raw engine output
//! /tmp/prod_<label>_s<seed>_trimmed.wav — post-trim (what production
//! ships, for `_sac` labels)
//!
//! The "no sacrificial" variants have no trim applied (just _raw); they
//! show what production produces for long prompts. The "_sac" variants
//! show both raw and trimmed, which is what `huddle::pocket::synth_chunk`
//! returns for short prompts. Listen back with `afplay`.
use std::collections::HashMap;
use std::path::PathBuf;
use sherpa_onnx::{
self, GenerationConfig, OfflineTts, OfflineTtsConfig, OfflineTtsModelConfig,
OfflineTtsPocketModelConfig, Wave,
};
const SAMPLE_RATE: u32 = 24_000;
const SHORT_PROMPT_MAX_FRAMES: i32 = 100;
// Mirror of huddle::pocket trim constants so the probe stays in sync with
// production. If you change either side, change both.
const TRIM_SCAN_START_SAMPLES: usize = (SAMPLE_RATE as usize * 30) / 1000;
const TRIM_SILENCE_THRESHOLD: f32 = 0.02;
const TRIM_MIN_GAP_SAMPLES: usize = (SAMPLE_RATE as usize * 50) / 1000;
const TRIM_MAX_DROP_SAMPLES: usize = (SAMPLE_RATE as usize * 1200) / 1000;
/// Mirror of `huddle::pocket::trim_leading_cold_start` — keep in sync.
fn trim_leading_cold_start(samples: &mut Vec<f32>) {
if samples.len() <= TRIM_SCAN_START_SAMPLES {
return;
}
let mut silence_run_start: Option<usize> = None;
let mut gap_end: Option<usize> = None;
for (i, sample) in samples.iter().enumerate().skip(TRIM_SCAN_START_SAMPLES) {
if sample.abs() < TRIM_SILENCE_THRESHOLD {
silence_run_start.get_or_insert(i);
} else if let Some(start) = silence_run_start {
if i - start >= TRIM_MIN_GAP_SAMPLES {
gap_end = Some(i);
break;
}
silence_run_start = None;
}
}
let Some(end) = gap_end else {
return;
};
if end > TRIM_MAX_DROP_SAMPLES {
return;
}
samples.drain(..end);
}
// (label, raw_text_before_prep, sacrificial_prefix_to_add_after_pad)
const TESTS: &[(&str, &str, &str)] = &[
// Short, previously-failing — sacrificial applied
("imhappy_sac", "I'm happy.", ". . "),
("imsorry_sac", "I'm sorry.", ". . "),
("imready_sac", "I'm ready.", ". . "),
// Short, previously-OK — sacrificial applied
("yep_sac", "Yep.", ". . "),
("isee_sac", "I see you.", ". . "),
// Long — sacrificial NOT applied (per design)
("longer_nosac", "Hello, how can I help you today?", ""),
("multi_nosac", "Yes, that works. Let me try again.", ""),
];
fn main() {
let model_dir = std::env::args()
.nth(1)
.unwrap_or_else(|| "/tmp/pocket-tts-bench".to_string());
let dir = PathBuf::from(&model_dir);
let p = |name: &str| dir.join(name).to_string_lossy().into_owned();
let mut cfg = OfflineTtsConfig::default();
cfg.model = OfflineTtsModelConfig {
pocket: OfflineTtsPocketModelConfig {
lm_main: Some(p("lm_main.int8.onnx")),
lm_flow: Some(p("lm_flow.int8.onnx")),
encoder: Some(p("encoder.onnx")),
decoder: Some(p("decoder.int8.onnx")),
text_conditioner: Some(p("text_conditioner.onnx")),
vocab_json: Some(p("vocab.json")),
token_scores_json: Some(p("token_scores.json")),
voice_embedding_cache_capacity: 16,
},
num_threads: 1,
debug: false,
..Default::default()
};
let engine = OfflineTts::create(&cfg).expect("engine create");
let voice_path = dir.join("reference_sample.wav");
let wave = Wave::read(voice_path.to_str().unwrap()).expect("voice WAV");
let ref_samples = wave.samples().to_vec();
let ref_sr = wave.sample_rate();
let seeds: &[i32] = &[42, 1337, 99999, 7, 314159];
println!(
"{:18} | {:>6} | {:>5} | {:>7} | gap_search (50ms,0.02) | path",
"test", "seed", "kind", "len_ms"
);
println!("{}", "-".repeat(110));
for (label, raw_text, sacrificial) in TESTS {
// Mirror huddle::pocket::prepare_pocket_prompt:
// - cleaned text starts with capital, ends with punctuation
// - short (≤4 words) → pad + sacrificial; max_frames=100
// - long → unchanged; no max_frames override
let cleaned = raw_text.trim();
let word_count = cleaned.split_whitespace().count();
let is_short = word_count <= 4;
let pad = if is_short { " " } else { "" };
let prompt = format!("{pad}{sacrificial}{cleaned}");
// Only short prompts get the post-synth trim in production; long
// prompts pass through unmodified.
let trim_in_production = is_short && !sacrificial.is_empty();
for seed in seeds {
let mut extra: HashMap<String, serde_json::Value> = HashMap::new();
extra.insert("seed".to_string(), serde_json::Value::from(*seed));
if is_short {
extra.insert(
"max_frames".to_string(),
serde_json::Value::from(SHORT_PROMPT_MAX_FRAMES),
);
}
let gen = GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 0.0, // PRODUCTION SETTING
reference_audio: Some(ref_samples.clone()),
reference_sample_rate: ref_sr,
extra: Some(extra),
..Default::default()
};
let audio = engine
.generate_with_config(&prompt, &gen, None::<fn(&[f32], f32) -> bool>)
.expect("synth");
// Raw output (what the engine returned).
let raw_samples = audio.samples().to_vec();
let raw_ms = raw_samples.len() as f32 / SAMPLE_RATE as f32 * 1000.0;
let raw_gap = find_gap(&raw_samples, SAMPLE_RATE, 0.02, 50);
let raw_path = format!("/tmp/prod_{}_s{}_raw.wav", label, seed);
sherpa_onnx::write(&raw_path, &raw_samples, SAMPLE_RATE as i32);
println!(
"{:18} | {:>6} | {:>5} | {:>5.0}ms | {:>22} | {}",
label, seed, "raw", raw_ms, raw_gap, raw_path
);
// Trimmed output — what synth_chunk actually returns to tts.rs
// for short prompts. For long prompts the engine output is
// returned untrimmed, so we skip writing a separate file.
if trim_in_production {
let mut trimmed = raw_samples;
trim_leading_cold_start(&mut trimmed);
let trimmed_ms = trimmed.len() as f32 / SAMPLE_RATE as f32 * 1000.0;
let trimmed_path = format!("/tmp/prod_{}_s{}_trimmed.wav", label, seed);
sherpa_onnx::write(&trimmed_path, &trimmed, SAMPLE_RATE as i32);
println!(
"{:18} | {:>6} | {:>5} | {:>5.0}ms | {:>22} | {}",
label, seed, "trim", trimmed_ms, "(post-trim)", trimmed_path
);
}
}
println!();
}
}
fn find_gap(samples: &[f32], sr: u32, thresh: f32, min_ms: u32) -> String {
let scan_start = (sr as usize * 30) / 1000;
let min_samples = (sr as usize * min_ms as usize) / 1000;
let mut silence_from: Option<usize> = None;
for i in scan_start..samples.len() {
if samples[i].abs() < thresh {
silence_from.get_or_insert(i);
} else if let Some(start) = silence_from {
if i - start >= min_samples {
return format!(
"{:.0}..{:.0}ms ({:.0}ms)",
start as f32 / sr as f32 * 1000.0,
i as f32 / sr as f32 * 1000.0,
(i - start) as f32 / sr as f32 * 1000.0
);
}
silence_from = None;
}
}
"<no gap>".to_string()
}

View File

@@ -0,0 +1,267 @@
//! Per-peer NetEQ-style jitter buffer for huddle remote audio.
//!
//! Wraps the [`neteq`] crate's `NetEq` state machine and registers a thin
//! `AudioDecoder` impl over our existing `opus = "0.3"` dependency. Default
//! features on `neteq` are disabled so we don't pull in axum/cpal/clap; this
//! module owns the Opus side of the trait instead.
//!
//! ## Lifecycle
//!
//! One [`PeerJitterBuffer`] per active remote peer. It owns:
//! * a `NetEq` configured for 48 kHz mono with conservative jitter bounds,
//! * a stable synthetic RTP SSRC (the relay-assigned `peer_index`), and
//! * a 16-bit sequence number that wraps with the wire protocol's `seq`.
//!
//! On `joined` / first packet, construct one. On `left` / rejoin, drop and
//! recreate so NetEq's internal delay state doesn't carry over.
//!
//! ## Decode path
//!
//! Each inbound Opus frame becomes a `neteq::AudioPacket` with:
//! * sequence number from the protocol v2 header,
//! * 48 kHz RTP-style media timestamp (frame_index * 960 for 20 ms frames),
//! * a synthetic SSRC == peer_index,
//! * payload type [`OPUS_PAYLOAD_TYPE`] (any value, only NetEq uses it to
//! dispatch to its registered decoder),
//! * the raw Opus bytes as `payload`,
//! * `duration_ms = 20` (our encoder frame size).
//!
//! The playout side calls [`PeerJitterBuffer::get_audio`] on a 10 ms tick;
//! NetEq returns 10 ms of decoded PCM with adaptive-delay / expand / accelerate
//! decisions already applied. For 20 ms input frames at 48 kHz mono, that's
//! 480 f32 samples per call.
use neteq::{
codec::AudioDecoder, neteq::SpeechType, AudioPacket, NetEq, NetEqConfig, NetEqError, RtpHeader,
};
/// RTP payload type used internally. NetEq only needs it to look up the decoder
/// registered via `register_decoder`; the relay forwards Opus payloads opaquely
/// so the wire never carries a payload-type byte.
pub const OPUS_PAYLOAD_TYPE: u8 = 111;
/// Audio sample rate for huddle (Opus VoIP, mono).
pub const SAMPLE_RATE_HZ: u32 = 48_000;
/// One channel (mono).
pub const CHANNELS: u8 = 1;
/// Opus encoder frame size in milliseconds. The wire format is one Opus packet
/// per 20 ms frame; the receiver still drains NetEq at 10 ms granularity.
pub const FRAME_DURATION_MS: u32 = 20;
/// 48 kHz media-time increment per encoder frame.
pub const FRAME_TIMESTAMP_DELTA: u32 = SAMPLE_RATE_HZ / 1000 * FRAME_DURATION_MS; // 960
/// NetEq returns 10 ms frames; this is the sample count per `get_audio` call.
/// Kept as a documented constant for the protocol-v2 follow-up and consumers
/// that want to size their own buffers around the playout contract.
#[allow(dead_code)]
pub const PLAYOUT_SAMPLES: usize = (SAMPLE_RATE_HZ as usize / 1000) * 10; // 480
/// Minimum jitter delay NetEq is allowed to converge to (ms).
const MIN_DELAY_MS: u32 = 40;
/// Maximum jitter delay NetEq is allowed to converge to (ms).
const MAX_DELAY_MS: u32 = 200;
/// Buffer cap in packets — generous, but bounded so a runaway sender can't
/// drive memory growth.
const MAX_PACKETS_IN_BUFFER: usize = 50;
/// Thin `AudioDecoder` over `opus::Decoder` so `neteq` can drive Opus payloads.
///
/// NetEq's `decode` trait is sample-based, not FEC-aware (no `fec` bool). The
/// stock `neteq` `NativeOpusDecoder` (gated behind the `native` feature, which
/// pulls in cpal/axum) calls `decode_float(..., false)` unconditionally; this
/// matches that behavior so we don't need the feature.
///
/// Receive-side FEC ("decode this frame's redundant copy on a known-lost prior
/// frame") would require either a trait change upstream or NetEq calling decode
/// twice with different `fec` values — out of scope for the initial 10-person
/// fix. Tracked as a follow-up alongside encoder-side `set_inband_fec`.
struct OpusFrameDecoder {
inner: opus::Decoder,
sample_rate: u32,
channels: u8,
/// Scratch buffer reused across calls to avoid per-frame allocation.
scratch: Vec<f32>,
}
impl OpusFrameDecoder {
fn new(sample_rate: u32, channels: u8) -> Result<Self, String> {
let opus_channels = match channels {
1 => opus::Channels::Mono,
2 => opus::Channels::Stereo,
n => return Err(format!("unsupported channel count: {n}")),
};
let inner = opus::Decoder::new(sample_rate, opus_channels)
.map_err(|e| format!("opus decoder init: {e}"))?;
// Largest plausible frame: 60 ms at 48 kHz stereo. We always use
// 20 ms mono in practice; this is just a one-shot allocation cap.
let cap = (sample_rate as usize / 1000) * 60 * channels as usize;
Ok(Self {
inner,
sample_rate,
channels,
scratch: vec![0.0; cap],
})
}
}
impl AudioDecoder for OpusFrameDecoder {
fn sample_rate(&self) -> u32 {
self.sample_rate
}
fn channels(&self) -> u8 {
self.channels
}
fn decode(&mut self, encoded: &[u8]) -> Result<Vec<f32>, NetEqError> {
match self.inner.decode_float(encoded, &mut self.scratch, false) {
Ok(samples_per_channel) => {
let total = samples_per_channel * self.channels as usize;
Ok(self.scratch[..total].to_vec())
}
Err(e) => Err(NetEqError::DecoderError(format!("opus decode: {e}"))),
}
}
}
/// One peer's jitter buffer + decoder pair.
pub struct PeerJitterBuffer {
neteq: NetEq,
ssrc: u32,
}
impl PeerJitterBuffer {
/// Construct a new jitter buffer for the peer identified by `peer_index`.
///
/// `peer_index` is the relay-assigned 0..=254 stable index for the peer's
/// lifetime in the room; we promote it to a 32-bit synthetic SSRC so NetEq
/// has a stable stream identity for its delay-manager state.
pub fn new(peer_index: u8) -> Result<Self, NetEqError> {
let config = NetEqConfig {
sample_rate: SAMPLE_RATE_HZ,
channels: CHANNELS,
max_packets_in_buffer: MAX_PACKETS_IN_BUFFER,
min_delay_ms: MIN_DELAY_MS,
max_delay_ms: MAX_DELAY_MS,
..Default::default()
};
let mut neteq = NetEq::new(config)?;
let decoder =
OpusFrameDecoder::new(SAMPLE_RATE_HZ, CHANNELS).map_err(NetEqError::DecoderError)?;
neteq.register_decoder(OPUS_PAYLOAD_TYPE, Box::new(decoder));
Ok(Self {
neteq,
ssrc: peer_index as u32,
})
}
/// Insert a received Opus packet.
///
/// `seq` and `ts_48k` come from the protocol v2 header (sender-authored,
/// 16-bit wrapping sequence + 48 kHz RTP-style media time).
pub fn insert_packet(
&mut self,
seq: u16,
ts_48k: u32,
opus_payload: &[u8],
) -> Result<(), NetEqError> {
let header = RtpHeader::new(seq, ts_48k, self.ssrc, OPUS_PAYLOAD_TYPE, false);
let packet = AudioPacket::new(
header,
opus_payload.to_vec(),
SAMPLE_RATE_HZ,
CHANNELS,
FRAME_DURATION_MS,
);
self.neteq.insert_packet(packet)
}
/// Drain one 10 ms playout frame. Always succeeds: NetEq emits PLC /
/// comfort-noise / silence rather than an error when there's nothing to
/// play.
///
/// Returns `(samples, vad_active)`. `samples.len() == PLAYOUT_SAMPLES`.
pub fn get_audio(&mut self) -> Result<(Vec<f32>, bool), NetEqError> {
let frame = self.neteq.get_audio()?;
let vad = frame.vad_activity && !matches!(frame.speech_type, SpeechType::Expand);
Ok((frame.samples, vad))
}
/// True when NetEq's buffer has no packets queued and the next
/// `get_audio` call would emit expand/silence. Used by
/// `huddle::playout` as part of the idle-peer gate — if a peer has
/// gone quiet (`last_packet_at` is stale) *and* there's nothing
/// buffered to drain, we stop appending frames to that peer's rodio
/// Player. The check is conservative: a peer who just sent a burst
/// then disconnected may have buffered audio that should still play
/// out, so this method drives that decision.
pub fn is_empty(&self) -> bool {
self.neteq.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn opus_silence_frame() -> Vec<u8> {
// Encode 20 ms of silence at 48 kHz mono so the test doesn't rely on
// synthetic Opus byte sequences.
let mut encoder = opus::Encoder::new(
SAMPLE_RATE_HZ,
opus::Channels::Mono,
opus::Application::Voip,
)
.expect("opus encoder");
let pcm = vec![0.0_f32; FRAME_TIMESTAMP_DELTA as usize];
let mut out = vec![0u8; 4000];
let n = encoder
.encode_float(&pcm, &mut out)
.expect("opus encode silence");
out.truncate(n);
out
}
#[test]
fn jitter_buffer_construction_succeeds_for_each_peer_index() {
for idx in [0u8, 1, 7, 42, 254] {
let jb = PeerJitterBuffer::new(idx).expect("construct jitter buffer");
assert!(jb.is_empty(), "peer {idx} should start empty");
// SSRC is the peer_index as u32 — verifies the stable-identity contract.
assert_eq!(jb.ssrc, idx as u32);
}
}
#[test]
fn insert_packet_then_get_audio_returns_playout_frame() {
let mut jb = PeerJitterBuffer::new(3).expect("jitter buffer");
let payload = opus_silence_frame();
// Insert a few sequential packets so NetEq has enough buffer to start
// emitting Normal-decoded frames instead of pure Expand/silence.
for i in 0..6u16 {
let ts = (i as u32) * FRAME_TIMESTAMP_DELTA;
jb.insert_packet(i, ts, &payload).expect("insert");
}
// Pull a playout frame — should be 10 ms = 480 mono samples.
let (samples, _vad) = jb.get_audio().expect("get_audio");
assert_eq!(
samples.len(),
PLAYOUT_SAMPLES,
"NetEq must emit exactly one 10 ms frame per get_audio call",
);
}
#[test]
fn empty_buffer_get_audio_still_produces_a_frame() {
// Even with no packets inserted, NetEq's contract is to always emit
// 10 ms — as Expand/silence — so the playout clock keeps ticking.
let mut jb = PeerJitterBuffer::new(0).expect("jitter buffer");
let (samples, vad) = jb.get_audio().expect("get_audio");
assert_eq!(samples.len(), PLAYOUT_SAMPLES);
assert!(
!vad,
"empty buffer should not report voice activity (Expand frame)",
);
}
}

View File

@@ -1,973 +0,0 @@
//! Kokoro-82M ONNX TTS engine — single-session inference with IPA G2P.
//!
//! Mental model:
//!
//! load_text_to_speech(model_dir) → KokoroTTS
//! load_voice_style(path) → VoiceStyle
//! tts.synth_chunk(text, lang, &style, steps, speed) → Vec<f32> @ 24 kHz
//!
//! ┌──────────┐ G2P ┌──────────┐ tokenize ┌──────────┐
//! │ raw text │ ──────→ │ IPA str │ ─────────→ │ int64[] │
//! └──────────┘ lexicon └──────────┘ 115-char └────┬─────┘
//! │
//! ┌──────────┐ style ┌──────────┐ ONNX ┌────▼─────┐
//! │ .bin file│ ──────→ │ [1, 256] │ ─────────→ │ Vec<f32> │
//! └──────────┘ indexed └──────────┘ session └──────────┘
//! by token count 24 kHz PCM
//!
//! G2P strategy: dictionary lookup (us_gold.json, Apache-2.0 via misaki).
//! OOV words are spelled letter-by-letter using a static IPA table.
//! No espeak dependency — fully GPL-free.
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use ndarray::{Array1, Array2};
use ort::{session::Session, value::Value};
// ── Public constants ──────────────────────────────────────────────────────────
pub const SAMPLE_RATE: u32 = 24_000;
pub const DEFAULT_VOICE: &str = "af_heart";
// Maximum phoneme tokens before padding (model context = 512, minus 2 pad tokens).
const MAX_PHONEME_TOKENS: usize = 510;
// ── VoiceStyle ────────────────────────────────────────────────────────────────
/// Raw f32 voice embedding loaded from a `<name>.bin` file.
///
/// The binary is a flat array of shape `[-1, 256]` in row-major order.
/// Row `i` is the style vector for an utterance with `i` phoneme tokens.
/// This encodes both speaker identity and sequence-length-dependent prosody.
#[derive(Debug)]
pub struct VoiceStyle {
data: Vec<f32>, // flat: row i = data[i*256 .. (i+1)*256]
}
impl VoiceStyle {
/// Return the 256-dim style vector for a given phoneme token count.
/// Clamps to the last available row if `token_count` is out of range.
fn get(&self, token_count: usize) -> &[f32] {
let max_rows = self.data.len() / 256;
let idx = token_count.min(max_rows.saturating_sub(1));
&self.data[idx * 256..(idx + 1) * 256]
}
}
/// Load a voice style from a raw little-endian f32 binary file.
pub fn load_voice_style(path: &Path) -> Result<VoiceStyle, String> {
let bytes = fs::read(path).map_err(|e| format!("read voice {}: {e}", path.display()))?;
if bytes.len() % 4 != 0 {
return Err(format!(
"voice file {} has non-multiple-of-4 byte count ({})",
path.display(),
bytes.len()
));
}
let data: Vec<f32> = bytes
.chunks_exact(4)
.map(|b| f32::from_le_bytes(b.try_into().unwrap()))
.collect();
if data.len() < 256 {
return Err(format!(
"voice file {} too small ({} floats, need at least 256)",
path.display(),
data.len()
));
}
if data.len() % 256 != 0 {
return Err(format!(
"voice style has {} floats — expected a multiple of 256 (got {} remainder)",
data.len(),
data.len() % 256,
));
}
Ok(VoiceStyle { data })
}
// ── Tokenizer ─────────────────────────────────────────────────────────────────
/// Static 115-entry IPA char → int64 lookup table.
/// IDs are non-contiguous (0–177); unknown chars are silently dropped.
/// Pad token '$' = 0 is prepended and appended to every sequence.
fn build_vocab() -> HashMap<char, i64> {
// Source: onnx-community/Kokoro-82M-v1.0-ONNX tokenizer.json
#[rustfmt::skip]
let entries: &[(char, i64)] = &[
('$', 0),
(';', 1), (':', 2), (',', 3), ('.', 4), ('!', 5), ('?', 6),
('—', 9), ('…', 10), ('"', 11), ('(', 12), (')', 13), ('\u{201c}', 14), ('\u{201d}', 15),
(' ', 16), ('\u{0303}', 17),
('ʣ', 18), ('ʥ', 19), ('ʦ', 20), ('ʨ', 21), ('ᵝ', 22), ('ꭧ', 23),
('A', 24), ('I', 25), ('O', 31), ('Q', 33), ('S', 35), ('T', 36),
('W', 39), ('Y', 41), ('ᵊ', 42),
('a', 43), ('b', 44), ('c', 45), ('d', 46), ('e', 47), ('f', 48),
('h', 50), ('i', 51), ('j', 52), ('k', 53), ('l', 54), ('m', 55),
('n', 56), ('o', 57), ('p', 58), ('q', 59), ('r', 60), ('s', 61),
('t', 62), ('u', 63), ('v', 64), ('w', 65), ('x', 66), ('y', 67), ('z', 68),
('ɑ', 69), ('ɐ', 70), ('ɒ', 71), ('æ', 72), ('β', 75), ('ɔ', 76),
('ɕ', 77), ('ç', 78), ('ɖ', 80), ('ð', 81), ('ʤ', 82), ('ə', 83),
('ɚ', 85), ('ɛ', 86), ('ɜ', 87), ('ɟ', 90), ('ɡ', 92), ('ɥ', 99),
('ɨ', 101), ('ɪ', 102), ('ʝ', 103), ('ɯ', 110), ('ɰ', 111),
('ŋ', 112), ('ɳ', 113), ('ɲ', 114), ('ɴ', 115), ('ø', 116),
('ɸ', 118), ('θ', 119), ('œ', 120), ('ɹ', 123), ('ɾ', 125),
('ɻ', 126), ('ʁ', 128), ('ɽ', 129), ('ʂ', 130), ('ʃ', 131),
('ʈ', 132), ('ʧ', 133), ('ʊ', 135), ('ʋ', 136), ('ʌ', 138),
('ɣ', 139), ('ɤ', 140), ('χ', 142), ('ʎ', 143), ('ʒ', 147),
('ʔ', 148), ('ˈ', 156), ('ˌ', 157), ('ː', 158), ('ʰ', 162),
('ʲ', 164), ('↓', 169), ('→', 171), ('↗', 172), ('↘', 173), ('ᵻ', 177),
];
entries.iter().copied().collect()
}
/// Convert an IPA phoneme string to a padded int64 token sequence.
/// Returns `[0, id1, id2, ..., idN, 0]` clamped to MAX_PHONEME_TOKENS+2.
/// The pre-pad token count (ids.len() - 2) is used to index the style vector.
fn tokenize(phonemes: &str, vocab: &HashMap<char, i64>) -> Vec<i64> {
let mut ids: Vec<i64> = vec![0]; // BOS pad
for id in phonemes
.chars()
.filter_map(|c| vocab.get(&c).copied())
.take(MAX_PHONEME_TOKENS)
{
ids.push(id);
}
ids.push(0); // EOS pad
ids
}
// ── G2P Lexicon ───────────────────────────────────────────────────────────────
/// Grapheme-to-phoneme engine with a four-tier fallback chain:
///
/// 1. Misaki gold+silver dicts (183K words, Kokoro-native IPA)
/// 2. CMUdict (135K words, ARPAbet→Kokoro IPA) — covers inflected forms
/// 3. Morphological suffix stripping (-s/-ed/-ing) + retry tiers 1-2
/// 4. Letter-by-letter spelling
///
/// All dictionaries are Apache-2.0 or BSD licensed. No GPL.
struct Lexicon {
/// Misaki gold+silver merged dictionary (Kokoro-native IPA).
misaki: HashMap<String, String>,
/// CMU Pronouncing Dictionary (ARPAbet converted to Kokoro IPA at load time).
cmudict: HashMap<String, String>,
}
/// IPA pronunciations for individual letter names (used for OOV words).
fn letter_ipa(c: char) -> &'static str {
match c {
'a' => "ˈeɪ",
'b' => "bˈiː",
'c' => "sˈiː",
'd' => "dˈiː",
'e' => "ˈiː",
'f' => "ˈɛf",
'g' => "dʒˈiː",
'h' => "ˈeɪtʃ",
'i' => "ˈaɪ",
'j' => "dʒˈeɪ",
'k' => "kˈeɪ",
'l' => "ˈɛl",
'm' => "ˈɛm",
'n' => "ˈɛn",
'o' => "ˈoʊ",
'p' => "pˈiː",
'q' => "kjˈuː",
'r' => "ˈɑːɹ",
's' => "ˈɛs",
't' => "tˈiː",
'u' => "jˈuː",
'v' => "vˈiː",
'w' => "dˈʌbəljˌuː",
'x' => "ˈɛks",
'y' => "wˈaɪ",
'z' => "zˈiː",
_ => "",
}
}
/// Punctuation chars that are valid Kokoro vocab tokens and should pass through.
fn is_passthrough_punct(c: char) -> bool {
matches!(c, ';' | ':' | ',' | '.' | '!' | '?' | '—' | '…' | ' ')
}
/// Vowels that trigger US English /t/→/ɾ/ flapping (misaki's US_TAUS).
const US_TAUS: &str = "AIOWYiuæɑəɛɪɹʊʌ";
/// ARPAbet → Kokoro IPA conversion. Stress digit is stripped before lookup.
fn arpabet_to_ipa(phoneme: &str) -> &'static str {
match phoneme {
"AA" => "ɑ",
"AE" => "æ",
"AH" => "ʌ",
"AO" => "ɔ",
"AW" => "W",
"AY" => "I",
"EH" => "ɛ",
"ER" => "ɜɹ",
"EY" => "A",
"IH" => "ɪ",
"IY" => "i",
"OW" => "O",
"OY" => "Y",
"UH" => "ʊ",
"UW" => "u",
"B" => "b",
"CH" => "ʧ",
"D" => "d",
"DH" => "ð",
"F" => "f",
"G" => "ɡ",
"HH" => "h",
"JH" => "ʤ",
"K" => "k",
"L" => "l",
"M" => "m",
"N" => "n",
"NG" => "ŋ",
"P" => "p",
"R" => "ɹ",
"S" => "s",
"SH" => "ʃ",
"T" => "t",
"TH" => "θ",
"V" => "v",
"W" => "w",
"Y" => "j",
"Z" => "z",
"ZH" => "ʒ",
_ => "",
}
}
/// Convert a CMUdict ARPAbet pronunciation line to Kokoro IPA.
/// Input: "K R IY0 EY1 T AH0 D" → Output: "kɹiˈAtəd"
fn arpabet_line_to_ipa(arpabet: &str) -> String {
let mut out = String::new();
for token in arpabet.split_whitespace() {
// Split phoneme from stress digit (e.g., "EY1" → "EY", Some('1'))
let (base, stress) = if token.ends_with(|c: char| c.is_ascii_digit()) {
(&token[..token.len() - 1], token.as_bytes().last().copied())
} else {
(token, None)
};
// Stress marker goes BEFORE the vowel's IPA
match stress {
Some(b'1') => out.push('ˈ'), // primary
Some(b'2') => out.push('ˌ'), // secondary
_ => {}
}
// AH with stress=0 is schwa (ə), not ʌ
if base == "AH" && stress == Some(b'0') {
out.push('ə');
} else if base == "ER" && stress == Some(b'0') {
// Unstressed ER is just əɹ
out.push_str("əɹ");
} else {
out.push_str(arpabet_to_ipa(base));
}
}
out
}
impl Lexicon {
/// Load misaki gold+silver dicts and CMUdict.
fn load(gold_path: &Path, silver_path: &Path, cmudict_path: &Path) -> Result<Self, String> {
let mut misaki = Self::load_json(silver_path)?;
let gold = Self::load_json(gold_path)?;
misaki.extend(gold);
let cmudict = if cmudict_path.exists() {
Self::load_cmudict(cmudict_path)?
} else {
eprintln!(
"sprout-desktop: CMUdict not found at {} — inflected forms may be spelled out",
cmudict_path.display()
);
HashMap::new()
};
eprintln!(
"sprout-desktop: G2P loaded — misaki: {} words, cmudict: {} words",
misaki.len(),
cmudict.len()
);
Ok(Lexicon { misaki, cmudict })
}
fn load_json(path: &Path) -> Result<HashMap<String, String>, String> {
let content =
fs::read_to_string(path).map_err(|e| format!("read {}: {e}", path.display()))?;
let raw: serde_json::Value =
serde_json::from_str(&content).map_err(|e| format!("parse {}: {e}", path.display()))?;
let obj = raw
.as_object()
.ok_or_else(|| format!("{}: expected JSON object", path.display()))?;
let mut dict = HashMap::with_capacity(obj.len());
for (word, val) in obj {
let ipa = match val {
serde_json::Value::String(s) => s.clone(),
serde_json::Value::Object(m) => m
.get("DEFAULT")
.or_else(|| m.values().next())
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
_ => continue,
};
if !ipa.is_empty() {
dict.insert(word.to_lowercase(), ipa);
}
}
Ok(dict)
}
/// Load CMUdict and convert ARPAbet → Kokoro IPA at load time.
/// Format: "WORD PH1 PH2 PH3\n" (single space between word and phonemes).
/// Variant pronunciations like "WORD(2)" are skipped — we take the first.
fn load_cmudict(path: &Path) -> Result<HashMap<String, String>, String> {
let content = fs::read_to_string(path).map_err(|e| format!("read cmudict: {e}"))?;
let mut dict = HashMap::with_capacity(140_000);
for line in content.lines() {
// Skip comments and blank lines
if line.starts_with(";;;") || line.is_empty() {
continue;
}
// Split on first space
let (word, phonemes) = match line.find(' ') {
Some(i) => (&line[..i], line[i + 1..].trim()),
None => continue,
};
// Skip variant pronunciations like "WORD(2)"
if word.contains('(') {
continue;
}
let key = word.to_lowercase();
let ipa = arpabet_line_to_ipa(phonemes);
if !ipa.is_empty() {
dict.entry(key).or_insert(ipa);
}
}
Ok(dict)
}
/// Look up a word across all tiers. Returns None if not found anywhere.
fn lookup(&self, word: &str) -> Option<String> {
self.misaki
.get(word)
.cloned()
.or_else(|| self.cmudict.get(word).cloned())
}
/// Apply English -s/-es/-ies suffix phoneme rules (misaki's `_s`).
fn apply_s(stem_ipa: &str) -> String {
let last = stem_ipa.chars().last().unwrap_or(' ');
if "ptkfθ".contains(last) {
format!("{stem_ipa}s")
} else if "szʃʒʧʤ".contains(last) {
format!("{stem_ipa}ᵻz")
} else {
format!("{stem_ipa}z")
}
}
/// Apply English -ed suffix phoneme rules (misaki's `_ed`).
fn apply_ed(stem_ipa: &str) -> String {
let chars: Vec<char> = stem_ipa.chars().collect();
let last = *chars.last().unwrap_or(&' ');
if "pkfθʃsʧ".contains(last) {
format!("{stem_ipa}t")
} else if last == 'd' {
format!("{stem_ipa}ᵻd")
} else if last != 't' {
format!("{stem_ipa}d")
} else if chars.len() >= 2 && US_TAUS.contains(chars[chars.len() - 2]) {
// US flap: "created" → kɹiˈAɾᵻd
let mut out: String = chars[..chars.len() - 1].iter().collect();
out.push_str("ɾᵻd");
out
} else {
format!("{stem_ipa}ᵻd")
}
}
/// Apply English -ing suffix phoneme rules (misaki's `_ing`).
fn apply_ing(stem_ipa: &str) -> String {
let chars: Vec<char> = stem_ipa.chars().collect();
let last = *chars.last().unwrap_or(&' ');
if last == 't' && chars.len() >= 2 && US_TAUS.contains(chars[chars.len() - 2]) {
// US flap: "creating" → kɹiˈAɾɪŋ
let mut out: String = chars[..chars.len() - 1].iter().collect();
out.push_str("ɾɪŋ");
out
} else {
format!("{stem_ipa}ɪŋ")
}
}
/// Try stripping -s/-ed/-ing suffix, look up the base, and re-apply phonetically.
fn try_morphological(&self, word: &str) -> Option<String> {
// Try -s / -es / -ies
if word.len() >= 3 && word.ends_with('s') {
// -ies → base + y
if word.len() > 4 && word.ends_with("ies") {
if let Some(stem) = self.lookup(&format!("{}y", &word[..word.len() - 3])) {
return Some(Self::apply_s(&stem));
}
}
// -es → base
if word.len() > 4 && word.ends_with("es") && !word.ends_with("ies") {
if let Some(stem) = self.lookup(&word[..word.len() - 2]) {
return Some(Self::apply_s(&stem));
}
}
// -s → base
if !word.ends_with("ss") {
if let Some(stem) = self.lookup(&word[..word.len() - 1]) {
return Some(Self::apply_s(&stem));
}
}
}
// Try -ed / -d
if word.len() >= 4 && word.ends_with('d') {
// -ed → base (not -eed)
if word.len() > 4 && word.ends_with("ed") && !word.ends_with("eed") {
if let Some(stem) = self.lookup(&word[..word.len() - 2]) {
return Some(Self::apply_ed(&stem));
}
// -ed where base ends in e: "created" → "create"
if let Some(stem) = self.lookup(&format!("{}e", &word[..word.len() - 2])) {
return Some(Self::apply_ed(&stem));
}
}
// -d → base (e.g., "discovered" → strip "d" → "discovere" fails,
// but "configured" → strip "d" → "configure" works)
if !word.ends_with("dd") {
if let Some(stem) = self.lookup(&word[..word.len() - 1]) {
return Some(Self::apply_ed(&stem));
}
}
}
// Try -ing
if word.len() >= 5 && word.ends_with("ing") {
let base = &word[..word.len() - 3];
// -ing → base (e.g., "running" base = "runn" — won't match, need double-consonant)
if let Some(stem) = self.lookup(base) {
return Some(Self::apply_ing(&stem));
}
// -ing + e → base+e (e.g., "creating" → "creat" + "e" = "create")
if let Some(stem) = self.lookup(&format!("{base}e")) {
return Some(Self::apply_ing(&stem));
}
// Double consonant: "running" → "run"
if base.len() >= 2 {
let bytes = base.as_bytes();
if bytes[bytes.len() - 1] == bytes[bytes.len() - 2] {
if let Some(stem) = self.lookup(&base[..base.len() - 1]) {
return Some(Self::apply_ing(&stem));
}
}
}
}
None
}
/// Convert a single word to IPA using the full fallback chain.
fn word_to_ipa(&self, word: &str) -> String {
// Compound words: split on hyphens and underscores, process each part
// independently. "short-and-natural" → "short" + "and" + "natural",
// "parent_event_id" → "parent" + "event" + "id".
// Each part gets full dict lookup. Joined with a space (brief TTS pause).
if word.contains('-') || word.contains('_') {
let parts: Vec<String> = word
.split(|c: char| c == '-' || c == '_')
.filter(|p| !p.is_empty())
.map(|p| self.word_to_ipa(p))
.collect();
return parts.join(" ");
}
// Normalize curly quotes to straight apostrophes.
let normalized = word.replace('\u{2019}', "'").replace('\u{2018}', "'");
let stripped: String = normalized
.chars()
.filter(|c| c.is_alphabetic() || *c == '\'')
.collect::<String>()
.to_lowercase();
// Tier 1+2: misaki + CMUdict direct lookup
if let Some(ipa) = self.lookup(&stripped) {
return ipa;
}
// Contractions: "don't" → "don" + "'t"
if let Some(apos_idx) = stripped.find('\'') {
let base = &stripped[..apos_idx];
let suffix = &stripped[apos_idx..];
if let Some(base_ipa) = self.lookup(base) {
let suffix_ipa = self.lookup(suffix).unwrap_or_else(|| match suffix {
"'ve" => "v".to_string(),
"'re" => "ɹ".to_string(),
_ => String::new(),
});
if !suffix_ipa.is_empty() {
return format!("{base_ipa}{suffix_ipa}");
}
}
}
// Tier 3: morphological suffix stripping
if let Some(ipa) = self.try_morphological(&stripped) {
return ipa;
}
// Tier 4: letter-by-letter spelling
stripped
.chars()
.filter(|c| c.is_alphabetic())
.map(letter_ipa)
.collect()
}
/// Convert a full text chunk to an IPA phoneme string.
fn text_to_ipa(&self, text: &str) -> String {
let mut out = String::new();
for token in text.split_whitespace() {
if !out.is_empty() {
out.push(' ');
}
let leading: String = token
.chars()
.take_while(|c| is_passthrough_punct(*c))
.collect();
let trailing: String = token
.chars()
.rev()
.take_while(|c| is_passthrough_punct(*c))
.collect::<String>()
.chars()
.rev()
.collect();
let word = &token[leading.len()..token.len() - trailing.len()];
out.push_str(&leading);
if !word.is_empty() {
out.push_str(&self.word_to_ipa(word));
}
out.push_str(&trailing);
}
out
}
}
// ── KokoroTTS ─────────────────────────────────────────────────────────────────
pub struct KokoroTTS {
session: Session,
vocab: HashMap<char, i64>,
lexicon: Lexicon,
// Retained for potential future use (e.g., hot-reloading voices by path).
#[allow(dead_code)]
model_dir: PathBuf,
}
/// Load the Kokoro TTS engine from a model directory.
///
/// Expects:
/// `<model_dir>/model.onnx` (or model_quantized.onnx — tries both)
/// `<model_dir>/us_gold.json` (G2P dictionary)
///
/// CoreML execution provider is registered with auto-fallback to CPU.
/// The compiled CoreML model is cached in `<model_dir>/.coreml_cache/`.
pub fn load_text_to_speech(model_dir: &str) -> Result<KokoroTTS, String> {
let model_dir_path = PathBuf::from(model_dir);
// Try quantized model first for speed, fall back to full-precision.
let model_path = ["model_quantized.onnx", "model_q8f16.onnx", "model.onnx"]
.iter()
.map(|name| model_dir_path.join(name))
.find(|p| p.exists())
.ok_or_else(|| format!("no model.onnx found in {model_dir}"))?;
// ── Session threading options ────────────────────────────────────────
//
// parallel_execution — runs independent graph operators concurrently.
// Kokoro's graph is mostly sequential, so the benefit is modest, but
// it's safe and free to enable.
// intra_threads(0) — lets ONNX Runtime use all available CPU cores for
// parallelism within individual operators (e.g., large matmuls).
//
// NOTE: GraphOptimizationLevel::All is already the ort default — no
// need to set it explicitly. Accuracy-altering flags (approximate_gelu,
// flush_to_zero) are intentionally omitted — they need A/B audio
// validation before enabling for a TTS model. memory_pattern is omitted
// because Kokoro input lengths vary per sentence and the ORT docs warn
// against it for variable-size inputs.
// Try CoreML first (zero binary cost — macOS system framework).
// If the model has ops CoreML can't handle (common with quantized models),
// the EP registers fine but commit_from_file fails. Catch that and retry
// with CPU-only. This is the expected path for model_q8f16.onnx.
let session = {
let mut builder_with_coreml = Session::builder()
.map_err(|e| format!("session builder: {e}"))?
.with_parallel_execution(true)
.map_err(|e| format!("parallel execution: {e}"))?
.with_intra_threads(0)
.map_err(|e| format!("intra threads: {e}"))?
.with_execution_providers([ort::ep::CoreML::default()
.with_compute_units(ort::ep::coreml::ComputeUnits::All)
.with_model_format(ort::ep::coreml::ModelFormat::MLProgram)
.with_model_cache_dir(model_dir_path.join(".coreml_cache").to_string_lossy())
.build()])
.map_err(|e| format!("execution provider: {e}"))?;
match builder_with_coreml.commit_from_file(&model_path) {
Ok(s) => {
eprintln!("sprout-desktop: Kokoro loaded with CoreML acceleration");
s
}
Err(coreml_err) => {
eprintln!(
"sprout-desktop: CoreML failed for {}, falling back to CPU: {coreml_err}",
model_path.display()
);
// Retry without any execution providers — pure CPU.
Session::builder()
.map_err(|e| format!("session builder (CPU fallback): {e}"))?
.with_parallel_execution(true)
.map_err(|e| format!("parallel execution (CPU): {e}"))?
.with_intra_threads(0)
.map_err(|e| format!("intra threads (CPU): {e}"))?
.commit_from_file(&model_path)
.map_err(|e| format!("load model {} (CPU): {e}", model_path.display()))?
}
}
};
let gold_path = model_dir_path.join("us_gold.json");
let silver_path = model_dir_path.join("us_silver.json");
let cmudict_path = model_dir_path.join("cmudict.dict");
let lexicon = Lexicon::load(&gold_path, &silver_path, &cmudict_path)?;
Ok(KokoroTTS {
session,
vocab: build_vocab(),
lexicon,
model_dir: model_dir_path,
})
}
impl KokoroTTS {
/// Synthesize a single pre-split text chunk. Caller is responsible for sentence splitting.
/// This avoids double-splitting when the TTS pipeline has already split the text.
///
/// - `_lang` is accepted for API compatibility but currently unused
/// (Kokoro v1.0 language is selected by voice name prefix, e.g. `af_*`).
/// - `_steps` is accepted for API compatibility but currently unused
/// (Kokoro is not diffusion-based).
pub fn synth_chunk(
&mut self,
text: &str,
_lang: &str,
style: &VoiceStyle,
_steps: usize,
speed: f32,
) -> Result<Vec<f32>, String> {
// G2P: text → IPA phoneme string
let ipa = self.lexicon.text_to_ipa(text);
// Tokenize: IPA → int64 ids with BOS/EOS pad tokens
let token_ids = tokenize(&ipa, &self.vocab);
// Style vector is indexed by phoneme count (excluding the 2 pad tokens).
// Shape expected by model: [1, 256] (kokoro-js uses [1, 256], not [1, 1, 256]).
let phoneme_count = token_ids.len() - 2;
let style_slice = style.get(phoneme_count);
// Build ONNX input tensors.
let seq_len = token_ids.len();
let input_ids_arr = Array2::from_shape_vec((1, seq_len), token_ids)
.map_err(|e| format!("input_ids shape: {e}"))?;
let input_ids_val =
Value::from_array(input_ids_arr).map_err(|e| format!("input_ids Value: {e}"))?;
// Style: [1, 256]. Research notes [1, 1, 256] but kokoro-js uses [1, 256].
let style_arr = Array2::from_shape_vec((1, 256), style_slice.to_vec())
.map_err(|e| format!("style shape: {e}"))?;
let style_val = Value::from_array(style_arr).map_err(|e| format!("style Value: {e}"))?;
let speed_arr = Array1::from_vec(vec![speed]);
let speed_val = Value::from_array(speed_arr).map_err(|e| format!("speed Value: {e}"))?;
// Run inference. Output[0] = waveform float32[1, N_samples].
let outputs = self
.session
.run(ort::inputs! {
"input_ids" => &input_ids_val,
"style" => &style_val,
"speed" => &speed_val,
})
.map_err(|e| format!("onnx run: {e}"))?;
let (_, waveform) = outputs[0]
.try_extract_tensor::<f32>()
.map_err(|e| format!("extract waveform: {e}"))?;
Ok(waveform.to_vec())
}
}
#[cfg(test)]
mod tests {
use super::*;
// ── Tokenizer ─────────────────────────────────────────────────────────
#[test]
fn tokenize_empty_produces_bos_eos() {
let vocab = build_vocab();
let ids = tokenize("", &vocab);
assert_eq!(ids, vec![0, 0]); // BOS + EOS only
}
#[test]
fn tokenize_known_chars() {
let vocab = build_vocab();
let ids = tokenize("a", &vocab);
// 'a' maps to 43 in the vocab
assert_eq!(ids, vec![0, 43, 0]);
}
#[test]
fn tokenize_unknown_chars_dropped() {
let vocab = build_vocab();
let ids = tokenize("🎉", &vocab);
// Emoji not in vocab — should be dropped, leaving only BOS+EOS
assert_eq!(ids, vec![0, 0]);
}
#[test]
fn tokenize_respects_max_length() {
let vocab = build_vocab();
let long_input: String = "a".repeat(600); // exceeds MAX_PHONEME_TOKENS (510)
let ids = tokenize(&long_input, &vocab);
// Should be clamped: BOS + 510 tokens + EOS = 512
assert_eq!(ids.len(), 512);
assert_eq!(ids[0], 0); // BOS
assert_eq!(*ids.last().unwrap(), 0); // EOS
}
// ── ARPAbet conversion ────────────────────────────────────────────────
#[test]
fn arpabet_simple_word() {
// "HH AH0 L OW1" = hello
let ipa = arpabet_line_to_ipa("HH AH0 L OW1");
assert_eq!(ipa, "həlˈO");
}
#[test]
fn arpabet_stress_markers() {
// Primary stress before vowel, secondary stress before vowel
let ipa = arpabet_line_to_ipa("K R IY0 EY1 T");
// IY0 = unstressed 'i', EY1 = primary 'A'
assert!(ipa.contains('ˈ'), "should contain primary stress: {ipa}");
}
#[test]
fn arpabet_schwa() {
// AH0 should produce schwa (ə), not ʌ
let ipa = arpabet_line_to_ipa("AH0");
assert_eq!(ipa, "ə");
}
#[test]
fn arpabet_unstressed_er() {
// ER0 should produce əɹ
let ipa = arpabet_line_to_ipa("ER0");
assert_eq!(ipa, "əɹ");
}
// ── Letter IPA ────────────────────────────────────────────────────────
#[test]
fn letter_ipa_covers_alphabet() {
for c in 'a'..='z' {
let ipa = letter_ipa(c);
assert!(!ipa.is_empty(), "letter_ipa('{c}') returned empty");
}
}
#[test]
fn letter_ipa_non_alpha_empty() {
assert_eq!(letter_ipa('1'), "");
assert_eq!(letter_ipa('!'), "");
}
// ── Punctuation passthrough ───────────────────────────────────────────
#[test]
fn passthrough_punct_includes_expected() {
assert!(is_passthrough_punct('.'));
assert!(is_passthrough_punct('!'));
assert!(is_passthrough_punct('?'));
assert!(is_passthrough_punct(' '));
assert!(is_passthrough_punct(','));
}
#[test]
fn passthrough_punct_excludes_alpha() {
assert!(!is_passthrough_punct('a'));
assert!(!is_passthrough_punct('Z'));
}
// ── VoiceStyle ────────────────────────────────────────────────────────
#[test]
fn voice_style_get_clamps_to_last_row() {
// 2 rows of 256 floats
let data: Vec<f32> = (0..512).map(|i| i as f32).collect();
let style = VoiceStyle { data };
// Row 0
assert_eq!(style.get(0)[0], 0.0);
// Row 1
assert_eq!(style.get(1)[0], 256.0);
// Row 999 should clamp to row 1 (last available)
assert_eq!(style.get(999)[0], 256.0);
}
#[test]
fn load_voice_style_rejects_too_small() {
use std::io::Write;
let dir = std::env::temp_dir().join("kokoro_test_small");
let _ = std::fs::create_dir_all(&dir);
let path = dir.join("tiny.bin");
let mut f = std::fs::File::create(&path).unwrap();
// Write only 100 floats (need at least 256)
for i in 0..100u32 {
f.write_all(&(i as f32).to_le_bytes()).unwrap();
}
drop(f);
let result = load_voice_style(&path);
assert!(result.is_err(), "should reject file with < 256 floats");
assert!(result.unwrap_err().contains("too small"));
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn load_voice_style_rejects_non_multiple_of_256() {
use std::io::Write;
let dir = std::env::temp_dir().join("kokoro_test_nonaligned");
let _ = std::fs::create_dir_all(&dir);
let path = dir.join("nonaligned.bin");
let mut f = std::fs::File::create(&path).unwrap();
// Write 257 floats — not a multiple of 256 (remainder = 1)
for i in 0..257u32 {
f.write_all(&(i as f32).to_le_bytes()).unwrap();
}
drop(f);
let result = load_voice_style(&path);
assert!(
result.is_err(),
"should reject file with non-multiple-of-256 floats"
);
let err = result.unwrap_err();
assert!(
err.contains("257"),
"error should mention float count: {err}"
);
assert!(
err.contains("remainder"),
"error should mention remainder: {err}"
);
let _ = std::fs::remove_dir_all(&dir);
}
// ── Suffix rules ──────────────────────────────────────────────────────
#[test]
fn apply_s_voiceless() {
// After voiceless consonants: +s
assert!(Lexicon::apply_s("kæt").ends_with('s'));
}
#[test]
fn apply_s_sibilant() {
// After sibilants: +ᵻz
assert!(Lexicon::apply_s("bʌz").ends_with("ᵻz"));
}
#[test]
fn apply_s_voiced() {
// After voiced consonants: +z
assert!(Lexicon::apply_s("dɔɡ").ends_with('z'));
}
#[test]
fn apply_ed_voiceless() {
// After voiceless: +t
assert!(Lexicon::apply_ed("wɔk").ends_with('t'));
}
#[test]
fn apply_ed_d_ending() {
// After d: +ᵻd
assert!(Lexicon::apply_ed("æd").ends_with("ᵻd"));
}
#[test]
fn apply_ing_basic() {
assert!(Lexicon::apply_ing("rʌn").ends_with("ɪŋ"));
}
#[test]
fn hyphenated_word_splits_into_parts() {
let lex = Lexicon {
misaki: HashMap::new(),
cmudict: HashMap::new(),
};
let result = lex.word_to_ipa("short-and-natural");
let space_count = result.matches(' ').count();
assert_eq!(
space_count, 2,
"expected 2 spaces for 3 hyphenated parts, got {space_count}: {result}"
);
}
#[test]
fn underscored_word_splits_into_parts() {
let lex = Lexicon {
misaki: HashMap::new(),
cmudict: HashMap::new(),
};
let result = lex.word_to_ipa("parent_event_id");
let space_count = result.matches(' ').count();
assert_eq!(
space_count, 2,
"expected 2 spaces for 3 underscored parts, got {space_count}: {result}"
);
}
#[test]
fn compound_word_with_dict_lookup() {
let mut dict = HashMap::new();
dict.insert("parent".to_string(), "pɛɹənt".to_string());
dict.insert("event".to_string(), "ɪvɛnt".to_string());
dict.insert("id".to_string(), "aɪdiː".to_string());
let lex = Lexicon {
misaki: dict,
cmudict: HashMap::new(),
};
// Underscore compound
let result = lex.word_to_ipa("parent_event_id");
assert!(result.contains("pɛɹənt"), "parent not resolved: {result}");
assert!(result.contains("ɪvɛnt"), "event not resolved: {result}");
assert!(result.contains("aɪdiː"), "id not resolved: {result}");
// Hyphen compound
let result = lex.word_to_ipa("short-and-sweet");
assert_eq!(result.matches(' ').count(), 2, "hyphen split: {result}");
}
}

View File

@@ -25,14 +25,17 @@
pub mod agents;
pub mod audio_output;
pub mod kokoro;
pub mod jitter;
pub mod models;
pub mod pipeline;
pub mod playout;
pub mod pocket;
pub mod preprocessing;
pub mod relay_api;
pub mod state;
pub mod stt;
pub mod tts;
pub mod wire;
// ── Shared utilities ──────────────────────────────────────────────────────────
@@ -661,12 +664,12 @@ pub async fn check_pipeline_hotstart(state: State<'_, AppState>) -> Result<(), S
let stt_ready = models::global_model_manager()
.map(|m| m.take_stt_ready())
.unwrap_or(false);
let kokoro_ready = models::global_model_manager()
.map(|m| m.take_kokoro_ready())
let tts_ready = models::global_model_manager()
.map(|m| m.take_tts_ready())
.unwrap_or(false);
// Start TTS first (so STT can capture tts_cancel).
if !has_tts && (kokoro_ready || models::is_kokoro_ready()) {
if !has_tts && (tts_ready || models::is_tts_ready()) {
if let Err(e) = maybe_start_tts_pipeline(&state).await {
eprintln!("sprout-desktop: TTS hotstart failed: {e}");
}
@@ -751,7 +754,7 @@ pub async fn start_stt_pipeline(state: State<'_, AppState>) -> Result<(), String
}
}
/// Trigger a background download of voice models (Parakeet STT + Kokoro TTS).
/// Trigger a background download of voice models (Parakeet STT + Pocket TTS).
///
/// Returns immediately — downloads run in tokio background tasks.
/// Poll `get_model_status` to track progress.
@@ -761,7 +764,7 @@ pub async fn download_voice_models(state: State<'_, AppState>) -> Result<(), Str
let manager = models::global_model_manager()
.ok_or("model manager unavailable (home directory could not be resolved)")?;
manager.start_stt_download(state.http_client.clone());
manager.start_kokoro_download(state.http_client.clone());
manager.start_tts_download(state.http_client.clone());
Ok(())
}
@@ -772,14 +775,14 @@ pub fn get_model_status(_state: State<'_, AppState>) -> Result<models::VoiceMode
.ok_or("model manager unavailable (home directory could not be resolved)")?;
Ok(models::VoiceModelStatus {
stt: manager.stt_status(),
kokoro: manager.kokoro_status(),
tts: manager.tts_status(),
})
}
/// Enable or disable TTS output.
///
/// When disabled, the TTS pipeline is shut down and audio output stops.
/// When re-enabled, the pipeline is restarted if Kokoro models are available.
/// When re-enabled, the pipeline is restarted if TTS models are available.
///
/// Takes the pipeline handle out of the lock before calling shutdown() — the
/// thread join in Drop can block for ~200 ms (ONNX inference) and we don't

View File

@@ -1,10 +1,10 @@
//! Model download manager for STT (Parakeet TDT-CTC 110M) and TTS (Kokoro) models.
//! Model download manager for STT (Parakeet TDT-CTC 110M) and TTS (Pocket TTS) models.
//!
//! Mental model:
//! app launch → start_stt_download (background) → ~/.sprout/models/parakeet-tdt-ctc-110m-en/
//! app launch → start_kokoro_download (background) → ~/.sprout/models/kokoro/
//! app launch → start_tts_download (background) → ~/.sprout/models/pocket-tts/
//! STT pipeline → is_stt_ready() → stt_model_dir() → run inference
//! TTS pipeline → is_kokoro_ready() → kokoro_model_dir() → run synthesis
//! TTS pipeline → is_tts_ready() → tts_model_dir() → run synthesis
//!
//! Models are downloaded once and cached. A version manifest (`.sprout-model-manifest`)
//! is written alongside model files — if the on-disk version doesn't match the
@@ -38,18 +38,38 @@ use sha2::{Digest, Sha256};
/// Computed from a known-good download. Update when upgrading model versions.
const STT_ARCHIVE_SHA256: &str = "17f945007b52ccd8b7200ffc7c5652e9e8e961dfdf479cefcabd06cf5703630b";
/// SHA-256 hashes for individual Kokoro model files.
/// Computed from known-good downloads. Update when upgrading model versions.
/// HuggingFace base URL for the sherpa-onnx Pocket TTS int8 repackage.
///
/// model.onnx (model_q8f16.onnx, 86 MB):
/// curl -sL "https://huggingface.co/onnx-community/Kokoro-82M-v1.0-ONNX/resolve/main/onnx/model_q8f16.onnx" | shasum -a 256
/// Pinned to commit e715955cf50d18d919d37231513c0e914b83661a
/// (2026-02-10) for reproducible downloads.
const POCKET_HF_BASE: &str =
"https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-int8-2026-01-26/resolve/e715955cf50d18d919d37231513c0e914b83661a";
/// Reference voice WAV: "Mary (f, conversation)" from the Kyutai TTS demo
/// voice set — VCTK speaker p333, ai-coustics-enhanced. Pinned to
/// kyutai/tts-voices commit 323332d33f997de8394f24a193e1a76df720e01a.
///
/// Mapping comes from the speaker dropdown on <https://kyutai.org/tts>:
/// the Pocket TTS preset "Mary (f, conversation)" maps to
/// `vctk/p333_023_enhanced.wav`. We rename to `reference_sample.wav` on disk
/// so the rest of the engine code stays voice-agnostic; the friendly label
/// only matters for attribution and PR-body docs.
const POCKET_REFERENCE_WAV_URL: &str =
"https://huggingface.co/kyutai/tts-voices/resolve/323332d33f997de8394f24a193e1a76df720e01a/vctk/p333_023_enhanced.wav";
/// SHA-256 hashes for individual Pocket TTS model files.
/// Computed from known-good pinned downloads. Update when upgrading model versions.
#[rustfmt::skip]
const KOKORO_FILE_HASHES: &[(&str, &str)] = &[
("model.onnx", "04c658aec1b6008857c2ad10f8c589d4180d0ec427e7e6118ceb487e215c3cd0"),
("af_heart.bin", "d583ccff3cdca2f7fae535cb998ac07e9fcb90f09737b9a41fa2734ec44a8f0b"),
("us_gold.json", "dc414872a49a28ae6c141463d502fd945f3b2fde040484fdc47d00cc4612686f"),
("us_silver.json", "de8f67be911bb6c659187b4a65fd966b6a30e56350e0f790d763210b053ac475"),
("cmudict.dict", "81917843c7f44ce2b094ac63873c2c7a4cf802040792c455ba3ca406891c3d22"),
const TTS_FILE_HASHES: &[(&str, &str)] = &[
("decoder.int8.onnx", "12b0857402d31aead94df19d6783b4350d1f740e811f3a3202c70ad89ae11eea"),
("encoder.onnx", "e8f2f6d301ffb96e398b138a7dc6d3038622d236044636b73d920bab85890260"),
("lm_flow.int8.onnx", "8d627d235c44a597da908e1085ebe241cbbe358964c502c5a5063d18851a5529"),
("lm_main.int8.onnx", "bfc0c7e7e3d72864fa3bb2ee499f62f21ddc1474b885f5f3ca570f8be73e787e"),
("text_conditioner.onnx", "0b84e837d7bfaf2c896627b03e3f080320309f37f4fc7df7698c644f7ba5e6b1"),
("vocab.json", "6fb646346cf931016f70c4921aab0900ce7a304b893cb02135c74e294abfea01"),
("token_scores.json", "5be2f278caf9b9800741f0fd82bff677f4943ec764c356f907213434b622d958"),
("LICENSE", "fe7b4ce83b8381cc5b216bbb4af73c570688d1b819c73bbaed8ca401f4677cd6"),
("reference_sample.wav", "a35b0468382218e9f37a9a7494d1e4b74deaf18d7ced22265b4e325bb55c183f"),
];
// ── Model versioning ──────────────────────────────────────────────────────────
@@ -65,8 +85,13 @@ const KOKORO_FILE_HASHES: &[(&str, &str)] = &[
/// honest (each version tag identifies one specific set of model bytes).
const STT_MODEL_VERSION: &str = "2";
/// Model manifest version for Kokoro. Increment when upgrading model files.
const KOKORO_MODEL_VERSION: &str = "1";
/// Model manifest version for Pocket TTS. Increment when upgrading model files.
/// Bumped "1" → "2" when the bundled reference voice changed from KevinAHM's
/// anonymous 16 kHz sample to Mary (VCTK p333, 32 kHz, ai-coustics-enhanced)
/// from kyutai/tts-voices. The hash mismatch on `reference_sample.wav` would
/// fail readiness on its own, but the manifest bump makes the re-download
/// reason explicit and skips the failing-then-re-fetching transient state.
const TTS_MODEL_VERSION: &str = "2";
/// Filename for the version manifest written alongside model files.
const MANIFEST_FILENAME: &str = ".sprout-model-manifest";
@@ -76,8 +101,8 @@ const MANIFEST_FILENAME: &str = ".sprout-model-manifest";
/// Maximum expected STT archive size (200 MB — actual is ~100 MB).
const MAX_STT_DOWNLOAD_BYTES: u64 = 200 * 1024 * 1024;
/// Maximum expected Kokoro file size (200 MB per file — model is 86 MB).
const MAX_KOKORO_FILE_BYTES: u64 = 200 * 1024 * 1024;
/// Maximum expected Pocket TTS file size (200 MB per file — largest is ~73 MB).
const MAX_TTS_FILE_BYTES: u64 = 200 * 1024 * 1024;
/// NVIDIA Parakeet TDT-CTC 110M (English, int8) — packaged for sherpa-onnx by
/// k2-fsa. Single ONNX file (CTC head) + tokens.txt. Avg WER ~7.5% across
@@ -125,35 +150,59 @@ Provided \"AS IS\", without warranty of any kind, express or implied. See the
license text for full warranty disclaimer.
";
// ── Kokoro TTS model ─────────────────────────────────────────────────────────
/// HuggingFace base URL for Kokoro ONNX model files.
const KOKORO_HF_BASE: &str =
"https://huggingface.co/onnx-community/Kokoro-82M-v1.0-ONNX/resolve/main";
/// Misaki G2P lexicons — pinned to commit fba1236 for reproducibility.
/// Gold = curated pronunciations. Silver = broader coverage (93K words).
/// Both are needed: gold is checked first, silver catches common words gold misses.
const KOKORO_LEXICON_GOLD_URL: &str =
"https://raw.githubusercontent.com/hexgrad/misaki/fba1236/misaki/data/us_gold.json";
const KOKORO_LEXICON_SILVER_URL: &str =
"https://raw.githubusercontent.com/hexgrad/misaki/fba1236/misaki/data/us_silver.json";
/// CMU Pronouncing Dictionary — 135K entries including inflected forms.
/// BSD 2-Clause license (Carnegie Mellon University). Compatible with Apache-2.0.
const KOKORO_CMUDICT_URL: &str =
"https://raw.githubusercontent.com/cmusphinx/cmudict/master/cmudict.dict";
// ── Pocket TTS model ──────────────────────────────────────────────────────────
/// Final directory name under `~/.sprout/models/`.
const KOKORO_MODEL_DIR_NAME: &str = "kokoro";
const TTS_MODEL_DIR_NAME: &str = "pocket-tts";
/// All files that must be present for Kokoro to be considered ready.
const KOKORO_EXPECTED_FILES: &[&str] = &[
"model.onnx",
"af_heart.bin",
"us_gold.json",
"us_silver.json",
"cmudict.dict",
/// Attribution sidecar written next to the Pocket TTS model files.
const TTS_LICENSE_FILE_NAME: &str = "MODEL_LICENSE.txt";
/// CC-BY-4.0 §3(a)(1) attribution block for Pocket TTS, its ONNX packaging,
/// and the bundled reference voice WAV.
const TTS_LICENSE_TEXT: &str = "\
Pocket TTS
© Kyutai.
Licensed under the Creative Commons Attribution 4.0 International License
(CC-BY-4.0). License text: https://creativecommons.org/licenses/by/4.0/
Original model by Kyutai: https://huggingface.co/kyutai/pocket-tts
Paper: Charles, Roebel, et al., Pocket TTS (arXiv:2509.06926).
Mimi neural codec by Kyutai is bundled as part of the model.
ONNX export by KevinAHM: https://huggingface.co/KevinAHM/pocket-tts-onnx
Sherpa-onnx repackage by csukuangfj / k2-fsa:
https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-int8-2026-01-26
Bundled reference voice (reference_sample.wav):
\"Mary (f, conversation)\" preset from the Kyutai TTS demo voice catalogue
(https://kyutai.org/tts), distributed via
https://huggingface.co/kyutai/tts-voices as `vctk/p333_023_enhanced.wav`.
Original recording from the Voice Cloning Toolkit (VCTK) corpus, speaker p333:
https://datashare.ed.ac.uk/handle/10283/3443 (CC-BY-4.0).
Recording enhancement (denoise/dereverb) by ai-coustics:
https://ai-coustics.com/
Sprout ships all ONNX/model artifacts and the reference voice WAV unmodified,
renamed only by placement in the local model directory.
Provided \"AS IS\", without warranty of any kind, express or implied. See the
license text for full warranty disclaimer.
";
/// All files that must be present for Pocket TTS to be considered ready.
const TTS_EXPECTED_FILES: &[&str] = &[
"decoder.int8.onnx",
"encoder.onnx",
"lm_flow.int8.onnx",
"lm_main.int8.onnx",
"text_conditioner.onnx",
"vocab.json",
"token_scores.json",
"LICENSE",
"reference_sample.wav",
TTS_LICENSE_FILE_NAME,
];
// ── Status types ──────────────────────────────────────────────────────────────
@@ -176,7 +225,7 @@ pub enum ModelStatus {
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VoiceModelStatus {
pub stt: ModelStatus,
pub kokoro: ModelStatus,
pub tts: ModelStatus,
}
// ── Safe archive extraction ───────────────────────────────────────────────────
@@ -339,7 +388,7 @@ where
// ── ModelSlot ─────────────────────────────────────────────────────────────────
/// Per-model state + config. `ModelManager` owns two of these (stt, kokoro).
/// Per-model state + config. `ModelManager` owns two of these (stt, tts).
#[derive(Clone)]
struct ModelSlot {
dir_name: &'static str, // subdir under ~/.sprout/models/
@@ -493,7 +542,7 @@ pub struct ModelManager {
/// `~/.sprout/models/`
models_dir: PathBuf,
stt: ModelSlot,
kokoro: ModelSlot,
tts: ModelSlot,
}
impl ModelManager {
@@ -505,11 +554,7 @@ impl ModelManager {
Some(Self {
models_dir,
stt: ModelSlot::new(STT_MODEL_DIR_NAME, STT_EXPECTED_FILES, STT_MODEL_VERSION),
kokoro: ModelSlot::new(
KOKORO_MODEL_DIR_NAME,
KOKORO_EXPECTED_FILES,
KOKORO_MODEL_VERSION,
),
tts: ModelSlot::new(TTS_MODEL_DIR_NAME, TTS_EXPECTED_FILES, TTS_MODEL_VERSION),
})
}
@@ -532,23 +577,23 @@ impl ModelManager {
self.stt.take_ready()
}
// ── Kokoro accessors ──────────────────────────────────────────────────────
// ── TTS accessors ─────────────────────────────────────────────────────────
/// Path to the Kokoro model directory, or `None` if not ready.
pub fn kokoro_model_dir(&self) -> Option<PathBuf> {
self.kokoro.dir_if_ready(&self.models_dir)
/// Path to the TTS model directory, or `None` if not ready.
pub fn tts_model_dir(&self) -> Option<PathBuf> {
self.tts.dir_if_ready(&self.models_dir)
}
/// `true` if all Kokoro files are present and the manifest version matches.
pub fn is_kokoro_ready(&self) -> bool {
self.kokoro.is_ready(&self.models_dir)
/// `true` if all TTS files are present and the manifest version matches.
pub fn is_tts_ready(&self) -> bool {
self.tts.is_ready(&self.models_dir)
}
/// Current Kokoro download status.
pub fn kokoro_status(&self) -> ModelStatus {
self.kokoro.status()
/// Current TTS download status.
pub fn tts_status(&self) -> ModelStatus {
self.tts.status()
}
/// Returns `true` once when Kokoro just became ready. Resets the flag.
pub fn take_kokoro_ready(&self) -> bool {
self.kokoro.take_ready()
/// Returns `true` once when TTS just became ready. Resets the flag.
pub fn take_tts_ready(&self) -> bool {
self.tts.take_ready()
}
// ── Download triggers ─────────────────────────────────────────────────────
@@ -583,14 +628,14 @@ impl ModelManager {
}
}
/// Start a background Kokoro download (~87 MB). No-op if already ready or downloading.
pub fn start_kokoro_download(&self, http_client: reqwest::Client) {
/// Start a background Pocket TTS download (~189 MB). No-op if already ready or downloading.
pub fn start_tts_download(&self, http_client: reqwest::Client) {
let manager = self.clone();
self.kokoro.start_download(
self.tts.start_download(
&self.models_dir,
http_client,
"kokoro",
move |client| async move { manager.download_kokoro_model(client).await },
"tts",
move |client| async move { manager.download_tts_model(client).await },
);
}
@@ -697,40 +742,42 @@ impl ModelManager {
Ok(())
}
/// Download and verify the Kokoro TTS model files from HuggingFace and GitHub.
/// Download and verify the Pocket TTS model files from HuggingFace.
///
/// Downloads files into `~/.sprout/models/kokoro/`:
/// - `model.onnx` — Kokoro-82M mixed-precision ONNX (86 MB)
/// - `af_heart.bin` — best-quality American English voice embedding (510 KB)
/// - `us_gold.json` — Misaki G2P lexicon, pinned to commit fba1236 (3 MB)
/// Downloads files into `~/.sprout/models/pocket-tts/`:
/// - five ONNX sessions (Pocket TTS + Mimi codec)
/// - `vocab.json` / `token_scores.json` for sherpa-onnx text conditioning
/// - upstream `LICENSE` plus Sprout's `MODEL_LICENSE.txt` attribution sidecar
/// - `reference_sample.wav` as the bundled default voice
///
/// Files are written to a temp directory first, then moved atomically.
async fn download_kokoro_model(&self, http_client: reqwest::Client) -> Result<(), String> {
async fn download_tts_model(&self, http_client: reqwest::Client) -> Result<(), String> {
tokio::fs::create_dir_all(&self.models_dir)
.await
.map_err(|e| format!("create models dir: {e}"))?;
let temp_dir = self.models_dir.join("kokoro.tmp");
let temp_dir = self.models_dir.join("pocket-tts.tmp");
fresh_temp_dir(&temp_dir).await?;
// (url, local_filename)
let downloads: &[(&str, &str)] = &[
(
&format!("{KOKORO_HF_BASE}/onnx/model_q8f16.onnx"),
"model.onnx",
),
(
&format!("{KOKORO_HF_BASE}/voices/af_heart.bin"),
"af_heart.bin",
),
(KOKORO_LEXICON_GOLD_URL, "us_gold.json"),
(KOKORO_LEXICON_SILVER_URL, "us_silver.json"),
(KOKORO_CMUDICT_URL, "cmudict.dict"),
let model_files = [
"decoder.int8.onnx",
"encoder.onnx",
"lm_flow.int8.onnx",
"lm_main.int8.onnx",
"text_conditioner.onnx",
"vocab.json",
"token_scores.json",
"LICENSE",
];
let mut downloads: Vec<(String, &'static str)> = model_files
.iter()
.map(|filename| (format!("{POCKET_HF_BASE}/{filename}"), *filename))
.collect();
downloads.push((POCKET_REFERENCE_WAV_URL.to_string(), "reference_sample.wav"));
let total_files = downloads.len() as u32;
for (i, (url, filename)) in downloads.iter().enumerate() {
eprintln!("sprout-desktop: downloading Kokoro {filename} from {url}");
eprintln!("sprout-desktop: downloading Pocket TTS {filename} from {url}");
let response = fetch_url(&http_client, url, filename).await.map_err(|e| {
let _ = std::fs::remove_dir_all(&temp_dir);
@@ -738,12 +785,12 @@ impl ModelManager {
})?;
let dest = temp_dir.join(filename);
let slot = self.kokoro.clone();
let slot = self.tts.clone();
let file_index = i as u32;
let bytes = download_file(
response,
&dest,
MAX_KOKORO_FILE_BYTES,
MAX_TTS_FILE_BYTES,
filename,
|downloaded, content_length| {
if let Some(total) = content_length {
@@ -766,30 +813,36 @@ impl ModelManager {
})?;
eprintln!("sprout-desktop: downloaded {bytes} bytes ({filename}), wrote to disk");
// Verify file integrity against pinned hash.
if let Some(&(_, expected)) = KOKORO_FILE_HASHES.iter().find(|(n, _)| *n == *filename) {
let actual = sha256_file(&dest).await?;
if actual != expected {
let _ = tokio::fs::remove_dir_all(&temp_dir).await;
return Err(format!(
"Kokoro {filename} integrity check failed: expected {expected}, got {actual}"
));
}
let expected = TTS_FILE_HASHES
.iter()
.find(|(n, _)| *n == *filename)
.map(|(_, hash)| *hash)
.ok_or_else(|| format!("missing expected hash for Pocket TTS file: {filename}"))?;
let actual = sha256_file(&dest).await?;
if actual != expected {
let _ = tokio::fs::remove_dir_all(&temp_dir).await;
return Err(format!(
"Pocket TTS {filename} integrity check failed: expected {expected}, got {actual}"
));
}
// Ensure progress reflects file completion even without content-length.
let pct = (((i as u32 + 1) * 89) / total_files).min(89) as u8;
self.kokoro.set_status(ModelStatus::Downloading {
self.tts.set_status(ModelStatus::Downloading {
progress_percent: pct,
});
}
self.kokoro.set_status(ModelStatus::Downloading {
tokio::fs::write(temp_dir.join(TTS_LICENSE_FILE_NAME), TTS_LICENSE_TEXT)
.await
.map_err(|e| format!("write TTS model license sidecar: {e}"))?;
self.tts.set_status(ModelStatus::Downloading {
progress_percent: 90,
});
if let Err(e) = self
.kokoro
.tts
.verify_and_install(&self.models_dir, &temp_dir, None)
.await
{
@@ -798,8 +851,8 @@ impl ModelManager {
}
eprintln!(
"sprout-desktop: Kokoro model ready at {}",
self.kokoro.model_dir(&self.models_dir).display()
"sprout-desktop: Pocket TTS model ready at {}",
self.tts.model_dir(&self.models_dir).display()
);
Ok(())
}
@@ -856,14 +909,37 @@ async fn cleanup_legacy_moonshine_dir(models_dir: &Path) {
}
}
/// Path to the Kokoro model directory, or `None` if not ready.
pub fn kokoro_model_dir() -> Option<PathBuf> {
global_model_manager()?.kokoro_model_dir()
/// Path to the TTS model directory, or `None` if not ready.
pub fn tts_model_dir() -> Option<PathBuf> {
global_model_manager()?.tts_model_dir()
}
/// `true` if all expected Kokoro model files are present on disk.
pub fn is_kokoro_ready() -> bool {
/// `true` if all expected TTS model files are present on disk.
pub fn is_tts_ready() -> bool {
global_model_manager()
.map(|m| m.is_kokoro_ready())
.map(|m| m.is_tts_ready())
.unwrap_or(false)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn tts_readiness_requires_license_sidecar() {
let temp = tempfile::tempdir().expect("tempdir");
let slot = ModelSlot::new(TTS_MODEL_DIR_NAME, TTS_EXPECTED_FILES, TTS_MODEL_VERSION);
let model_dir = temp.path().join(TTS_MODEL_DIR_NAME);
std::fs::create_dir_all(&model_dir).expect("create model dir");
for file in TTS_EXPECTED_FILES {
std::fs::write(model_dir.join(file), b"test").expect("write expected file");
}
std::fs::write(model_dir.join(MANIFEST_FILENAME), TTS_MODEL_VERSION).expect("manifest");
assert!(slot.is_ready(temp.path()));
std::fs::remove_file(model_dir.join(TTS_LICENSE_FILE_NAME)).expect("remove sidecar");
assert!(!slot.is_ready(temp.path()));
}
}

View File

@@ -41,7 +41,7 @@ pub(crate) async fn post_connect_setup(
// Ensure voice models are downloading (idempotent).
if let Some(mgr) = models::global_model_manager() {
mgr.start_stt_download(state.http_client.clone());
mgr.start_kokoro_download(state.http_client.clone());
mgr.start_tts_download(state.http_client.clone());
}
// Connect audio relay WebSocket (Opus encode/decode pipeline).
@@ -153,7 +153,7 @@ pub(crate) async fn maybe_start_stt_pipeline(
Ok(true)
}
/// Attempt to start the TTS pipeline if Kokoro models are present and TTS is enabled.
/// Attempt to start the TTS pipeline if TTS models are present and TTS is enabled.
///
/// Returns `Ok(true)` if the pipeline was started, `Ok(false)` if preconditions
/// aren't met (model not ready, pipeline exists, TTS disabled), or `Err` on failure.
@@ -164,11 +164,11 @@ pub(crate) async fn maybe_start_stt_pipeline(
/// leaks ~200MB of ONNX sessions. The sentinel is set under the lock before
/// releasing it for the expensive construction step.
pub(crate) async fn maybe_start_tts_pipeline(state: &AppState) -> Result<bool, String> {
if !models::is_kokoro_ready() {
return Ok(false); // Kokoro not downloaded yet — TTS unavailable.
if !models::is_tts_ready() {
return Ok(false); // TTS model not downloaded yet — TTS unavailable.
}
let model_dir = match models::kokoro_model_dir() {
let model_dir = match models::tts_model_dir() {
Some(d) => d,
None => return Ok(false),
};

View File

@@ -0,0 +1,362 @@
//! Receive-side playout loop for the huddle audio relay.
//!
//! Owns the per-peer state map (one `NetEq` + one `rodio::Player` per remote
//! peer), the 10 ms playout clock, and the 500 ms active-speaker tick. Sibling
//! to [`relay_api`](super::relay_api), which keeps the encode/send half.
//!
//! ## Architecture
//!
//! ```text
//! WS binary frame ──► insert_packet ──► NetEq jitter buffer
//! │
//! playout_tick (10 ms) ──┘──► get_audio ─► per-peer
//! rodio::Player
//! │
//! ▼
//! device mixer (sums
//! concurrent peers)
//! ```
//!
//! The pre-fix shape used a single `rodio::Player` shared across every peer.
//! `Player` is a FIFO queue, so 3+ simultaneous speakers serialized into one
//! voice flipping speakers every 20 ms with unbounded queue growth. See
//! `desktop/src-tauri/tests/rodio_mixer_diagnostic.rs` for the deterministic
//! repro that pins this diagnosis in CI.
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message as WsMsg;
use tokio_util::sync::CancellationToken;
use super::jitter::{PeerJitterBuffer, SAMPLE_RATE_HZ};
use super::relay_api::{WsStream, REMOTE_SPEECH_THRESHOLD};
use super::wire::{FrameHeader, FLAG_DTX, V2_HEADER_LEN};
/// Speaker-tick window for emitting `huddle-active-speakers`. Active set is
/// cleared each tick — peers that didn't send a frame in the last window are
/// considered silent.
const SPEAKER_TICK_MS: u64 = 500;
/// Per-peer arrival window for the TTS interrupt frame counter.
const FRAME_WINDOW: std::time::Duration = std::time::Duration::from_millis(500);
/// Playout clock: NetEq emits 10 ms frames, so we tick at 10 ms.
const PLAYOUT_TICK_MS: u64 = 10;
/// How long after the last received packet we keep pulling frames out of a
/// peer's NetEq into its rodio Player. NetEq always emits a frame on every
/// `get_audio` call — silent/Expand when there are no packets — so without
/// this bound an idle peer (one that disconnected without sending `left`)
/// would have its Player queue 100 silence buffers/sec forever. We pull for
/// a short grace window past the last packet so brief DTX gaps still feed
/// NetEq's PLC/expand path normally.
const IDLE_PEER_GRACE: std::time::Duration = std::time::Duration::from_millis(500);
/// Drift bound on per-peer rodio `Player` queue depth.
///
/// The playout pipeline has two clocks: the producer is a `tokio` 10 ms
/// interval (this loop) that pulls from NetEq and appends to each peer's
/// `Player`; the consumer is the `cpal` audio callback that pulls samples
/// from the device `Mixer` at hardware sample rate. NetEq does rate-adapt
/// (accelerate / expand) but only to its own input pattern — it cannot
/// see the actual device-side consumption rate.
///
/// In steady state producer ≈ consumer, but scheduler jitter or small
/// clock skew can leave the producer slightly ahead, and rodio's
/// `Player` queue is an unbounded MPSC under the hood. Over a long call
/// that drift would accumulate as monotonic added latency (and eventually
/// memory).
///
/// We bound it explicitly: before each append, if the queue is already
/// at or above this threshold, drop the oldest queued frame with
/// `Player::skip_one()` so the new frame replaces it. 4 frames × 10 ms
/// = 40 ms, far below NetEq's `max_delay_ms = 200 ms`, so the audible
/// effect is negligible while the worst-case latency stays bounded.
const PLAYOUT_QUEUE_HIGH_WATER: usize = 4;
/// One remote peer's slot: jitter buffer + dedicated rodio Player.
///
/// Per-frame seq/timestamp come from the v2 wire header (sender-authored).
/// The relay forwards `peer_index | header | opus_bytes` opaquely; we parse
/// the header here and pass the sender's own monotonic seq + 48 kHz media
/// timestamp into NetEq.
struct PeerSlot {
jitter: PeerJitterBuffer,
player: rodio::Player,
/// Wall-clock time of the most recent inbound packet for this peer. Read
/// by the playout tick to decide whether to keep draining NetEq into the
/// Player. Updated on every successful `insert_packet`.
last_packet_at: tokio::time::Instant,
}
impl PeerSlot {
fn new(peer_idx: u8, sink_mixer: &rodio::mixer::Mixer) -> Option<Self> {
match PeerJitterBuffer::new(peer_idx) {
Ok(jitter) => Some(Self {
jitter,
player: rodio::Player::connect_new(sink_mixer),
last_packet_at: tokio::time::Instant::now(),
}),
Err(e) => {
eprintln!("sprout-desktop: jitter buffer init peer {peer_idx}: {e}");
None
}
}
}
/// Whether this peer is still actively sending — used by the playout tick
/// to gate the rodio append so disconnected peers don't pump silence
/// indefinitely.
///
/// The gate is `recent packet OR jitter buffer not empty`. The recent-
/// packet half covers the common case: brief speech gaps and DTX cadence
/// (≤400 ms) stay inside the [`IDLE_PEER_GRACE`] window so PLC/expand
/// frames keep flowing. The buffer-not-empty half is a safety net for
/// the edge case Mari called out: a peer who sends a burst then
/// disconnects has real audio queued in NetEq that should still play
/// out, even if `last_packet_at` ages past the grace before the buffer
/// finishes draining. The grace alone is enough today because NetEq's
/// `max_delay_ms` (200 ms) is well inside the grace (500 ms), but the
/// OR keeps the invariant robust against future config tuning.
fn is_active(&self) -> bool {
self.last_packet_at.elapsed() < IDLE_PEER_GRACE || !self.jitter.is_empty()
}
}
/// Drive the receive loop until cancelled or the WS closes.
///
/// `ws_tx_for_pongs` is shared with the encode-side task and only used here to
/// reply to Pings; it is locked briefly per Ping and never held across the
/// audio fast path.
#[allow(clippy::too_many_arguments)]
pub(crate) async fn run_playout_recv_loop(
mut ws_rx: futures_util::stream::SplitStream<WsStream>,
ws_tx_for_pongs: Arc<tokio::sync::Mutex<futures_util::stream::SplitSink<WsStream, WsMsg>>>,
sink_handle: rodio::MixerDeviceSink,
cancel: CancellationToken,
app_handle: Option<tauri::AppHandle>,
initial_peers: Vec<(u8, String)>,
tts_active: Arc<AtomicBool>,
tts_cancel: Arc<AtomicBool>,
) {
use rodio::buffer::SamplesBuffer;
use std::num::NonZero;
let mut peers: std::collections::HashMap<u8, PeerSlot> = std::collections::HashMap::new();
let channels = NonZero::new(1u16).expect("1 is non-zero");
let rate = NonZero::new(SAMPLE_RATE_HZ).expect("48k is non-zero");
let mut index_to_pubkey: std::collections::HashMap<u8, String> =
initial_peers.into_iter().collect();
let mut active_indices: std::collections::HashSet<u8> = std::collections::HashSet::new();
let mut frame_counts: std::collections::HashMap<u8, u16> = std::collections::HashMap::new();
let mut last_frame_reset = tokio::time::Instant::now();
let mut tts_was_active = false;
let mut speaker_tick = tokio::time::interval(std::time::Duration::from_millis(SPEAKER_TICK_MS));
speaker_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut playout_tick = tokio::time::interval(std::time::Duration::from_millis(PLAYOUT_TICK_MS));
// `Delay` (not `Skip`) so a brief stall in another select arm — e.g. the
// ws_tx_for_pongs mutex contending with the encode-side task on a Ping —
// doesn't drop a playout tick outright. Dropped ticks would leave the
// per-peer Player queues empty for 10 ms and the device mixer would
// produce audible silence. `Delay` catches up immediately when the loop
// returns to the select.
playout_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
loop {
tokio::select! {
biased;
_ = cancel.cancelled() => break,
_ = playout_tick.tick() => {
// Drain one 10 ms frame from each *active* peer's NetEq into
// its Player. NetEq always emits a frame (Expand/silence when
// empty), so for peers that recently sent we keep the device
// mixer fed without starving; for peers that have stopped
// sending — disconnected without a `left`, or simply quiet —
// we skip the append so we don't pump 100 silence buffers/sec
// per idle peer into rodio forever. `is_active` is a 500 ms
// grace past the last received packet, far longer than typical
// DTX comfort-noise cadence.
for (peer_idx, slot) in peers.iter_mut() {
if !slot.is_active() {
// Still drain the frame to keep NetEq's internal clock
// advancing; just don't enqueue it for playback.
let _ = slot.jitter.get_audio();
continue;
}
match slot.jitter.get_audio() {
Ok((samples, _vad)) => {
// Bound producer-vs-device-clock drift. If our
// tokio tick has gotten ahead of the audio
// callback's actual consumption rate, drop the
// oldest queued frame rather than letting the
// queue grow without bound.
if slot.player.len() >= PLAYOUT_QUEUE_HIGH_WATER {
eprintln!(
"sprout-desktop: playout queue high-water for peer {peer_idx} \
(depth={}) — dropping oldest frame",
slot.player.len(),
);
slot.player.skip_one();
}
slot.player.append(SamplesBuffer::new(channels, rate, samples));
}
Err(e) => {
eprintln!(
"sprout-desktop: jitter get_audio peer {peer_idx}: {e}"
);
}
}
}
}
_ = speaker_tick.tick() => {
if let Some(ref app) = app_handle {
use tauri::Emitter;
let pubkeys: Vec<String> = active_indices
.iter()
.filter_map(|idx| index_to_pubkey.get(idx).cloned())
.collect();
let _ = app.emit("huddle-active-speakers", &pubkeys);
}
active_indices.clear();
}
msg = ws_rx.next() => {
match msg {
Some(Ok(WsMsg::Binary(data))) => {
// Wire shape (v2): [peer_index: u8][header: 8 bytes][opus payload...]
// The minimum size is 1 (peer_index) + 8 (header) + ≥1 Opus byte.
if data.len() <= 1 + V2_HEADER_LEN {
continue;
}
let peer_idx = data[0];
let after_idx = &data[1..];
let Some((header, opus_bytes)) = FrameHeader::parse(after_idx)
else {
// Malformed v2 frame: header parse only fails when
// the slice is too short, which `if data.len() <= ...`
// already guards. Defensive log + drop.
eprintln!(
"sprout-desktop: dropping malformed audio frame from peer {peer_idx} ({} bytes)",
data.len(),
);
continue;
};
if opus_bytes.is_empty() {
continue;
}
let is_dtx = (header.flags & FLAG_DTX) != 0;
// Only count non-DTX arrivals toward the UI's
// active-speaker set. DTX/comfort packets are emitted
// by an idle peer to keep the codec alive — they
// don't mean the peer is speaking, and shouldn't
// make their tile flash for the 500 ms speaker tick.
if !is_dtx {
active_indices.insert(peer_idx);
}
// TTS interrupt frame counter — reset on TTS rising edge.
let tts_now = tts_active.load(Ordering::Acquire);
if tts_now && !tts_was_active {
frame_counts.clear();
last_frame_reset = tokio::time::Instant::now();
}
tts_was_active = tts_now;
let slot = match peers.entry(peer_idx) {
std::collections::hash_map::Entry::Occupied(e) => e.into_mut(),
std::collections::hash_map::Entry::Vacant(e) => {
let Some(slot) = PeerSlot::new(peer_idx, sink_handle.mixer())
else {
continue;
};
e.insert(slot)
}
};
// Sender-authored seq/ts: NetEq can detect real
// packet reordering & loss, not just arrival jitter.
if let Err(err) =
slot.jitter
.insert_packet(header.seq, header.ts_48k, opus_bytes)
{
eprintln!(
"sprout-desktop: jitter insert peer {peer_idx}: {err}"
);
} else {
// Heartbeat for the playout tick's idle-peer
// guard — only on successful insert so a stream
// of bad packets can't keep a dead peer "active".
slot.last_packet_at = tokio::time::Instant::now();
}
// Count remote-speech frame arrivals for the TTS
// interrupt. DTX/comfort frames don't count — they
// mean the peer is silent, just keeping the codec
// state alive.
if tts_now && !is_dtx {
if last_frame_reset.elapsed() >= FRAME_WINDOW {
frame_counts.clear();
last_frame_reset = tokio::time::Instant::now();
}
let count = frame_counts.entry(peer_idx).or_insert(0);
*count = count.saturating_add(1);
if *count >= REMOTE_SPEECH_THRESHOLD {
tts_cancel.store(true, Ordering::Release);
}
}
}
Some(Ok(WsMsg::Text(text))) => {
if let Ok(v) = serde_json::from_str::<serde_json::Value>(&text) {
match v["type"].as_str() {
Some("joined") => {
if let Some(peer_list) = v["peers"].as_array() {
for p in peer_list {
if let (Some(pk), Some(idx)) = (
p["pubkey"].as_str(),
p["peer_index"].as_u64(),
) {
let key = idx as u8;
// peer_index reuse with a new pubkey:
// flush the old peer's NetEq + Player so
// the next frame starts clean.
if index_to_pubkey
.get(&key)
.map(|s| s.as_str())
!= Some(pk)
{
peers.remove(&key);
frame_counts.remove(&key);
active_indices.remove(&key);
}
index_to_pubkey.insert(key, pk.to_string());
}
}
}
}
Some("left") => {
if let Some(idx) = v["peer_index"].as_u64() {
let key = idx as u8;
index_to_pubkey.remove(&key);
frame_counts.remove(&key);
// Dropping Player detaches its queue from the
// device mixer, freeing the per-peer slot.
peers.remove(&key);
}
}
_ => {}
}
}
}
Some(Ok(WsMsg::Ping(data))) => {
let mut tx = ws_tx_for_pongs.lock().await;
let _ = tx.send(WsMsg::Pong(data)).await;
}
Some(Ok(WsMsg::Close(_))) | None => break,
Some(Ok(_)) => {} // non-binary/text frame
Some(Err(_)) => break,
}
}
}
}
}

View File

@@ -0,0 +1,872 @@
//! Pocket TTS engine wrapper around sherpa-onnx's `OfflineTts`.
//!
//! Pocket TTS is a small (~189 MB int8 ONNX) zero-shot voice-cloning TTS
//! model from Kyutai. It runs quickly on CPU via sherpa-onnx, replacing the
//! previous Kokoro-82M engine that also required an espeak-free but
//! lexicon-heavy G2P pipeline (Misaki + CMUdict).
//!
//! ## Attribution
//!
//! - **Model**: Kyutai *Pocket TTS* — Charles, Roebel, et al., 2026.
//! arXiv:2509.06926. Original repository: <https://huggingface.co/kyutai/pocket-tts>.
//! Licensed CC-BY-4.0.
//! - **Mimi neural codec**: Kyutai, bundled in the same release. CC-BY-4.0.
//! - **ONNX export**: KevinAHM —
//! <https://huggingface.co/KevinAHM/pocket-tts-onnx>. CC-BY-4.0.
//! - **sherpa-onnx repackage**: csukuangfj / k2-fsa —
//! <https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-int8-2026-01-26>.
//! Repackages KevinAHM's export with the file layout sherpa-onnx's
//! `OfflineTtsPocketModelConfig` expects. CC-BY-4.0.
//! - **Reference voice WAV** (`reference_sample.wav`): the "Mary
//! (f, conversation)" preset from the Kyutai TTS demo
//! (<https://kyutai.org/tts>), which maps to `vctk/p333_023_enhanced.wav`
//! in <https://huggingface.co/kyutai/tts-voices>. CC-BY-4.0, base recording
//! from the VCTK corpus, enhanced by ai-coustics.
//!
//! Sprout ships these files unmodified; see the on-disk `MODEL_LICENSE.txt`
//! sidecar written by `huddle::models` during install for the canonical
//! CC-BY-4.0 §3(a)(1) attribution block.
//!
//! ## Engine-module contract (see `huddle::tts`)
//!
//! `pocket.rs` exposes a fixed surface used by `tts.rs`. Mirroring this
//! contract is what lets the TTS pipeline stay engine-agnostic:
//!
//! - `SAMPLE_RATE: u32` — engine output sample rate in Hz.
//! - `DEFAULT_VOICE: &str` — default voice name (without extension).
//! - `VOICE_FILE_EXT: &str` — extension for per-voice files on disk.
//! - `load_text_to_speech(model_dir)` → `Result<Engine, String>`
//! - `load_voice_style(path)` → `Result<VoiceStyle, String>`
//! - `Engine::synth_chunk(&self, text, lang, &VoiceStyle, steps, speed)`
//! → `Result<Vec<f32>, String>`
//!
//! `lang` and `steps` are accepted for API compatibility with the previous
//! Kokoro engine but are unused — Pocket TTS does its own language ID from
//! the input text and is not a diffusion model (consistency LM, one step).
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use sherpa_onnx::{GenerationConfig, OfflineTts, OfflineTtsConfig, Wave};
// ── Engine-module contract: public consts ─────────────────────────────────────
/// Pocket TTS emits 24 kHz mono PCM. Matches the previous Kokoro output rate,
/// so the rodio sink and inter-sentence silence buffer in `tts.rs` remain valid.
pub const SAMPLE_RATE: u32 = 24_000;
/// Name (without extension) of the bundled reference voice. The model directory
/// is expected to contain `<DEFAULT_VOICE>.<VOICE_FILE_EXT>` after install.
pub const DEFAULT_VOICE: &str = "reference_sample";
/// Voice files for Pocket TTS are reference audio (WAV). Distinct from the
/// Kokoro `.bin` style vectors — the model conditions on raw waveform samples,
/// not a precomputed embedding, so the extension change is honest.
pub const VOICE_FILE_EXT: &str = "wav";
// ── Tuning ────────────────────────────────────────────────────────────────────
/// Single-threaded ONNX execution for predictable CPU contention with the STT
/// pipeline. Matches `STT_NUM_THREADS` in `stt.rs`; raise only if a benchmark
/// argues for it.
const TTS_NUM_THREADS: i32 = 1;
/// LRU cache size for cloned voice embeddings inside the sherpa-onnx engine.
/// We bind to one voice per pipeline today, but the upstream example uses 16
/// and the cost is negligible — keep room for future multi-voice support.
const VOICE_EMBEDDING_CACHE_CAPACITY: i32 = 16;
/// Pocket TTS is a consistency-based LM. Generation quality saturates at one
/// denoising step — the upstream `GenerationConfig` default of 5 multiplies
/// synthesis time by ~5× with no audible benefit on this model.
const SYNTH_NUM_STEPS: i32 = 1;
/// Disable the upstream default 200 ms of pre/post silence padding. We splice
/// `INTER_SENTENCE_SILENCE` in `tts.rs` ourselves and don't want a double
/// helping of leading silence on every utterance.
const SYNTH_SILENCE_SCALE: f32 = 0.0;
/// sherpa-onnx upstream default for `max_frames` (LM steps), in
/// `offline-tts-pocket-impl.h:Generate`. 500 steps ≈ 40 s of audio at the
/// Mimi 12.5 Hz frame rate. Referenced only by the regression test below;
/// production code path never raises (or even reads) this value — we just
/// leave sherpa-onnx's own default in place by not setting the override.
#[cfg(test)]
const SHERPA_ONNX_MAX_FRAMES_DEFAULT: i32 = 500;
/// Tight `max_frames` we ask for on short, padded prompts to bound the
/// original "monster breathing" runaway. 100 LM steps ≈ 8 s of audio —
/// roomy for any one-to-four-word utterance the user is likely to elicit
/// while still well short of the 40 s upstream default. Chosen with slack so
/// we never *truncate* a legitimate short reply.
const SHORT_PROMPT_MAX_FRAMES: i32 = 100;
/// Word-count threshold (inclusive) below which we pad the prompt with
/// leading spaces and cap `max_frames` tighter than the upstream default.
/// Matches upstream `pocket_tts.models.tts_model.prepare_text_prompt`. Above
/// this threshold we leave sherpa-onnx's own defaults in place — overriding
/// them caused the "first 'yep' is just static" regression seen on
/// 2026-05-18, where dropping `frames_after_eos` below the upstream default
/// of 3 clipped the leading audio of multi-clause sentences.
const SHORT_PROMPT_WORD_THRESHOLD: usize = 4;
/// Number of leading spaces prepended to short prompts. The upstream Python
/// uses exactly 8 — keep parity rather than tuning blindly.
const SHORT_PROMPT_PAD_SPACES: usize = 8;
/// Sacrificial cold-start prefix appended *after* the leading space pad for
/// short prompts. Pocket TTS' FlowLM autoregressive generation has a 2–3
/// step "settle" period at the start where the first generated phoneme can
/// be smeared or dropped entirely (see kyutai-labs/pocket-tts #91, #70 and
/// sherpa-onnx #3180). For short utterances like "I'm happy." the first
/// phoneme is most of the first word — losing it produces "m happy".
///
/// Two periods separated by a space act as a "phantom utterance" that the
/// model commits to, absorbing the cold-start. The pair (rather than a
/// single period) was empirically the only variant in our probe — see
/// `examples/prod_probe.rs` — that produced a usable post-sacrificial
/// silence gap on every random seed. The resulting leading sacrificial
/// audio is then stripped from the output by [`trim_leading_cold_start`]
/// before the buffer is returned to the synth pipeline.
///
/// Long prompts (>4 words) don't need this — the first phoneme already has
/// enough downstream context to avoid the smear, and an early natural pause
/// (e.g. the comma in "Hello, how can I help you?") could be misdetected as
/// the trim boundary.
const SACRIFICIAL_PREFIX: &str = ". . ";
// ── Leading cold-start trim (post-synth) ──────────────────────────────────────
/// Skip this many samples at the start of the synth buffer before looking
/// for the sacrificial→main silence gap. The Mimi decoder cold-start
/// produces ~30 ms of low-amplitude noise that we *don't* want to treat as
/// the gap. 30 ms × 24 kHz = 720 samples.
const TRIM_SCAN_START_SAMPLES: usize = (SAMPLE_RATE as usize * 30) / 1000;
/// Amplitude threshold below which a sample is considered "silence" for the
/// purposes of finding the post-sacrificial gap. Tuned empirically against
/// production-config probe data — the engine's own `ScaleSilence` uses 0.01,
/// but our boundary detection wants a looser threshold so that the breath /
/// aspiration of the rendered periods (which sits around 0.005–0.015) is
/// treated as silence too.
const TRIM_SILENCE_THRESHOLD: f32 = 0.02;
/// A silence run must be at least this many samples long to be accepted as
/// the sacrificial→main word boundary. 50 ms is comfortably longer than
/// inter-syllable silence within a normal word at this speed (typically
/// 10–30 ms), so this guards against trimming into the middle of the real
/// utterance. 50 ms × 24 kHz = 1200 samples.
const TRIM_MIN_GAP_SAMPLES: usize = (SAMPLE_RATE as usize * 50) / 1000;
/// Hard cap on how much audio we'll trim from the start. Production probe
/// data (with `silence_scale = 0.0`) shows valid trim boundaries land
/// between 30 ms and ~450 ms; 1.2 s is a wide safety margin. If the
/// detector finds a "gap" past this point it's almost certainly an interior
/// pause inside an unusually long short-prompt utterance — bail out and
/// emit untrimmed audio rather than corrupt it. 1.2 s × 24 kHz = 28800
/// samples.
const TRIM_MAX_DROP_SAMPLES: usize = (SAMPLE_RATE as usize * 1200) / 1000;
/// sherpa-onnx's documented `frames_after_eos` default. We deliberately do
/// *not* override this knob — the previous attempt to bump it for short
/// inputs and lower it for long inputs lowered it below the upstream default
/// of 3, which clipped the leading audio of multi-clause sentences (the
/// "first 'yep' is static" regression). The constant exists only for the
/// regression test below. Source: `offline-tts-pocket-impl.h:Generate`.
#[cfg(test)]
const SHERPA_ONNX_FRAMES_AFTER_EOS_DEFAULT: i32 = 3;
// ── ONNX file names (five Pocket TTS sessions plus two JSON tables) ───────────
const FILE_LM_MAIN: &str = "lm_main.int8.onnx";
const FILE_LM_FLOW: &str = "lm_flow.int8.onnx";
const FILE_ENCODER: &str = "encoder.onnx";
const FILE_DECODER: &str = "decoder.int8.onnx";
const FILE_TEXT_COND: &str = "text_conditioner.onnx";
const FILE_VOCAB: &str = "vocab.json";
const FILE_TOKEN_SCORES: &str = "token_scores.json";
// ── Voice style ───────────────────────────────────────────────────────────────
/// Loaded reference voice — normalised f32 PCM samples plus their sample rate.
///
/// Pocket TTS takes a reference waveform per generation call (not a
/// precomputed style embedding), so we keep the samples in memory and clone
/// the small `Vec` into each `GenerationConfig` rather than re-reading the
/// WAV from disk on every sentence.
#[derive(Debug, Clone)]
pub struct VoiceStyle {
samples: Vec<f32>,
sample_rate: i32,
}
/// Load a reference voice WAV from disk.
///
/// Accepts any sample rate sherpa-onnx's `Wave::read` can decode — Pocket TTS
/// resamples internally using `reference_sample_rate`. The bundled
/// `reference_sample.wav` ("Mary" — VCTK p333, enhanced) is 32 kHz mono.
pub fn load_voice_style(path: &Path) -> Result<VoiceStyle, String> {
let path_str = path
.to_str()
.ok_or_else(|| format!("voice path is not valid UTF-8: {}", path.display()))?;
let wave = Wave::read(path_str)
.ok_or_else(|| format!("could not read voice WAV at {}", path.display()))?;
let samples = wave.samples().to_vec();
if samples.is_empty() {
return Err(format!("voice WAV is empty: {}", path.display()));
}
Ok(VoiceStyle {
samples,
sample_rate: wave.sample_rate(),
})
}
// ── Engine ────────────────────────────────────────────────────────────────────
/// Pocket TTS engine handle. Cheap to construct (one `OfflineTts::create`
/// call). Owned by the TTS worker thread for the lifetime of a huddle session.
///
/// `OfflineTts` does not implement `Debug`, so we don't derive it here — the
/// pipeline only needs to move the engine into the worker thread and call
/// `synth_chunk` on it, never to print it.
pub struct PocketTts {
inner: OfflineTts,
}
/// Build the Pocket TTS engine from the model directory installed by
/// `huddle::models`. Returns `Err` if any expected ONNX or JSON file is
/// missing — readiness is normally enforced by `is_tts_ready` upstream, but
/// the check is repeated here so a manually-modified model dir produces a
/// clear error string instead of an opaque sherpa-onnx `None`.
pub fn load_text_to_speech(model_dir: &str) -> Result<PocketTts, String> {
let dir = PathBuf::from(model_dir);
for name in [
FILE_LM_MAIN,
FILE_LM_FLOW,
FILE_ENCODER,
FILE_DECODER,
FILE_TEXT_COND,
FILE_VOCAB,
FILE_TOKEN_SCORES,
] {
let p = dir.join(name);
if !p.is_file() {
return Err(format!("missing Pocket TTS file: {}", p.display()));
}
}
let to_str = |name: &str| -> String { dir.join(name).to_string_lossy().into_owned() };
// Build the config by mutating defaults — mirrors `stt.rs` and stays
// resilient if sherpa-onnx adds unrelated model-family fields.
let mut cfg = OfflineTtsConfig::default();
cfg.model.pocket.lm_main = Some(to_str(FILE_LM_MAIN));
cfg.model.pocket.lm_flow = Some(to_str(FILE_LM_FLOW));
cfg.model.pocket.encoder = Some(to_str(FILE_ENCODER));
cfg.model.pocket.decoder = Some(to_str(FILE_DECODER));
cfg.model.pocket.text_conditioner = Some(to_str(FILE_TEXT_COND));
cfg.model.pocket.vocab_json = Some(to_str(FILE_VOCAB));
cfg.model.pocket.token_scores_json = Some(to_str(FILE_TOKEN_SCORES));
cfg.model.pocket.voice_embedding_cache_capacity = VOICE_EMBEDDING_CACHE_CAPACITY;
cfg.model.num_threads = TTS_NUM_THREADS;
// Explicit — defaults are not part of the API contract, and noisy debug
// logging in release builds would be expensive on every synthesized chunk.
cfg.model.debug = false;
let inner = OfflineTts::create(&cfg)
.ok_or_else(|| "OfflineTts::create returned None for Pocket TTS".to_string())?;
Ok(PocketTts { inner })
}
// ── Prompt preparation ────────────────────────────────────────────────────────
/// Result of [`prepare_pocket_prompt`]: a synthesizer-ready prompt plus the
/// per-call generation overrides derived from the original text.
///
/// `None` for either override means "leave sherpa-onnx's documented default
/// in place". The pipeline only sets `max_frames` (and only for short
/// padded inputs) so it can bound the original "monster breathing" runaway
/// without disturbing the rest of the LM sampling envelope.
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct PreparedPrompt {
/// Text to hand to `OfflineTts::generate_with_config`. Capitalized,
/// punctuation-terminated, and (for short inputs) left-padded with spaces
/// plus a sacrificial `". . "` cold-start prefix.
pub text: String,
/// Value to pass via `GenerationConfig.extra["max_frames"]`, or `None` to
/// keep the upstream default of 500 LM steps. We only override on short
/// padded prompts where we have a tight expectation on output length.
pub max_frames: Option<i32>,
/// `true` iff this prompt received the short-input treatment (leading
/// space pad + sacrificial `". . "` prefix). The synth pipeline uses
/// this to decide whether to apply [`trim_leading_cold_start`] to the
/// output: long prompts have no sacrificial audio to strip, and
/// trimming them risks deleting real speech at a natural early pause
/// (e.g. the comma in "Hello, how can I help you?").
pub is_short: bool,
}
/// Mirror of the *text-preparation* half of upstream
/// `pocket_tts.models.tts_model.prepare_text_prompt`. Sherpa-onnx's C++
/// Pocket TTS impl does not run these preparation steps, so short /
/// unpunctuated / lowercase inputs can trigger up to 40 s of runaway
/// generation when the EOS logit never crosses its threshold. We replicate
/// the upstream Python recipe here:
///
/// 1. Collapse interior whitespace (already done by `preprocess_for_tts`, but
/// cheap to re-check after sentence splitting).
/// 2. Capitalize the first letter.
/// 3. Append `.` if the text doesn't end in punctuation.
/// 4. If fewer than five words, prepend `SHORT_PROMPT_PAD_SPACES` spaces
/// followed by [`SACRIFICIAL_PREFIX`] (a `". . "` cold-start absorber —
/// see its docstring for the bug it works around), and return a tight
/// [`SHORT_PROMPT_MAX_FRAMES`] cap so the LM can't run away if EOS
/// still doesn't fire.
///
/// We do **not** override `frames_after_eos` — sherpa-onnx's default of 3
/// is what we want. An earlier version set it to 1 on long inputs, which
/// clipped the leading audio of multi-clause sentences ("first 'yep' is
/// just static" regression). Tests `prepare_prompt_never_lowers_frames_…`
/// lock this in.
///
/// Returns `None` only if the input is empty after trimming — caller should
/// skip synthesis in that case.
pub(crate) fn prepare_pocket_prompt(input: &str) -> Option<PreparedPrompt> {
let trimmed = input.trim();
if trimmed.is_empty() {
return None;
}
// Collapse stray double-spaces / embedded newlines that may slip past
// `preprocess_for_tts` when sentences are spliced back together.
let mut cleaned = String::with_capacity(trimmed.len());
let mut last_was_space = false;
for ch in trimmed.chars() {
let is_ws = ch.is_whitespace();
if is_ws {
if !last_was_space {
cleaned.push(' ');
}
last_was_space = true;
} else {
cleaned.push(ch);
last_was_space = false;
}
}
// Capitalize first character. Uses `to_uppercase` (multi-codepoint safe).
let first = cleaned.chars().next().expect("cleaned non-empty above");
if first.is_lowercase() {
let upper: String = first.to_uppercase().collect();
let mut iter = cleaned.chars();
iter.next();
cleaned = upper + iter.as_str();
}
// Ensure terminal punctuation. Anything not in `.!?;:,` gets a period.
// The upstream Python only checks `isalnum` → period, but for our agent
// text we already may end in `!` `?` `.` etc. — treat any of those as OK.
let last = cleaned
.chars()
.next_back()
.expect("cleaned non-empty above");
if !matches!(last, '.' | '!' | '?' | ';' | ':' | ',') {
cleaned.push('.');
}
// Word count of the *cleaned but not padded* text — padding is whitespace
// only and would just lie to the threshold check below.
let word_count = cleaned.split_whitespace().count();
let is_short = word_count <= SHORT_PROMPT_WORD_THRESHOLD;
let (final_text, max_frames) = if is_short {
let mut padded = String::with_capacity(
cleaned.len() + SHORT_PROMPT_PAD_SPACES + SACRIFICIAL_PREFIX.len(),
);
for _ in 0..SHORT_PROMPT_PAD_SPACES {
padded.push(' ');
}
padded.push_str(SACRIFICIAL_PREFIX);
padded.push_str(&cleaned);
(padded, Some(SHORT_PROMPT_MAX_FRAMES))
} else {
// For everything ≥5 words, fall back to upstream defaults. Overriding
// these is what caused the "first 'yep' is static" regression — the
// upstream LM has been tuned for `frames_after_eos = 3` and
// `max_frames = 500`, and there's no clear win in second-guessing.
(cleaned, None)
};
Some(PreparedPrompt {
text: final_text,
max_frames,
is_short,
})
}
/// Strip the leading "sacrificial" audio produced by the `". . "` cold-start
/// prefix from a short-prompt synthesis result. Only call this when
/// [`PreparedPrompt::is_short`] is `true` — the trim looks for a long
/// silence run at the head of the buffer, and an early natural pause inside
/// a long unsacrificed utterance (e.g. the comma in "Hello, how can I help
/// you?") would be misclassified as the sacrificial gap.
///
/// Algorithm:
/// 1. Skip the first [`TRIM_SCAN_START_SAMPLES`] of the buffer (Mimi
/// cold-start noise we shouldn't classify as silence).
/// 2. Scan forward for the first run of samples below
/// [`TRIM_SILENCE_THRESHOLD`] that lasts at least
/// [`TRIM_MIN_GAP_SAMPLES`] — that's the post-sacrificial boundary.
/// 3. If that boundary lies beyond [`TRIM_MAX_DROP_SAMPLES`], treat it as
/// "almost certainly an interior pause" and *do not trim* — the safe
/// fallback is to play the slightly-degraded raw audio rather than
/// delete real speech.
/// 4. Otherwise, drop the leading samples up to the end of the silence
/// run. We don't insert a zero lead-in here — `tts.rs` owns playback
/// cushioning by prepending `SENTENCE_LEAD_IN_SAMPLES` of zeros before
/// each appended sentence chunk.
///
/// If the scan never finds a long-enough gap (≈1% of generations in the
/// production-config probe), the function is a no-op — the model trajectory
/// missed the expected sacrificial→main structure and we'd rather play the
/// raw buffer than emit silence.
fn trim_leading_cold_start(samples: &mut Vec<f32>) {
if samples.len() <= TRIM_SCAN_START_SAMPLES {
return;
}
let mut silence_run_start: Option<usize> = None;
let mut gap_end: Option<usize> = None;
for (i, sample) in samples.iter().enumerate().skip(TRIM_SCAN_START_SAMPLES) {
if sample.abs() < TRIM_SILENCE_THRESHOLD {
silence_run_start.get_or_insert(i);
} else if let Some(start) = silence_run_start {
if i - start >= TRIM_MIN_GAP_SAMPLES {
gap_end = Some(i);
break;
}
silence_run_start = None;
}
}
let Some(end) = gap_end else {
// No gap found — model didn't produce the expected sacrificial→main
// structure. Bail out and let the caller play the raw buffer.
return;
};
if end > TRIM_MAX_DROP_SAMPLES {
// Boundary too far into the audio to plausibly be the sacrificial
// gap. Almost certainly an interior pause in the real utterance —
// leave the audio alone.
return;
}
samples.drain(..end);
}
/// Build the `GenerationConfig.extra` HashMap from a [`PreparedPrompt`].
///
/// Centralised so the regression test below can assert that we **never**
/// emit a `frames_after_eos` override — the previous attempt to override
/// that knob (setting it to 1 for ≥5-word inputs) clipped the leading
/// audio of multi-clause sentences (the "first 'yep' is static" bug on
/// 2026-05-18). The upstream sherpa-onnx default of 3 is what we want, and
/// the right way to keep it is to not set it at all.
fn build_generation_extra(prepared: &PreparedPrompt) -> Option<HashMap<String, serde_json::Value>> {
prepared.max_frames.map(|mf| {
let mut h: HashMap<String, serde_json::Value> = HashMap::with_capacity(1);
h.insert("max_frames".to_string(), serde_json::Value::from(mf));
h
})
}
impl PocketTts {
/// Synthesise `text` with the given reference voice.
///
/// `_lang` and `_steps` are accepted for API compatibility with the
/// previous Kokoro engine. Pocket TTS infers language from the input text
/// directly and is a one-step consistency model. Returns an empty buffer
/// for whitespace-only input.
pub fn synth_chunk(
&self,
text: &str,
_lang: &str,
style: &VoiceStyle,
_steps: usize,
speed: f32,
) -> Result<Vec<f32>, String> {
// Mirror upstream pocket-tts prompt prep — without this short or
// unpunctuated inputs can cause the LM's EOS logit to never trip,
// producing up to 40 s of "monster breathing" garbage on the first
// utterance. See `prepare_pocket_prompt` for the full recipe.
let prepared = match prepare_pocket_prompt(text) {
Some(p) => p,
None => return Ok(Vec::new()),
};
// Per-call generation hints sherpa-onnx forwards to
// `offline-tts-pocket-impl.h`. We only override `max_frames`, and
// only for short padded prompts where we have a tight expectation
// on output length — that bounds the original runaway without
// disturbing the rest of the LM sampling envelope. See
// `prepare_pocket_prompt` docs for the regression history.
let extra = build_generation_extra(&prepared);
let cfg = GenerationConfig {
speed,
num_steps: SYNTH_NUM_STEPS,
silence_scale: SYNTH_SILENCE_SCALE,
reference_audio: Some(style.samples.clone()),
reference_sample_rate: style.sample_rate,
extra,
..Default::default()
};
// No progress callback — synthesis is fast enough that returning the
// whole buffer at once keeps the lookahead pipelining in `tts.rs`
// simple. `None::<fn(...) -> bool>` pins the callback type for the
// `generate_with_config` generic parameter.
let audio = self
.inner
.generate_with_config(&prepared.text, &cfg, None::<fn(&[f32], f32) -> bool>)
.ok_or_else(|| {
format!(
"Pocket TTS synthesis failed for text ({} chars)",
prepared.text.len()
)
})?;
let sample_rate = audio.sample_rate();
if sample_rate != SAMPLE_RATE as i32 {
eprintln!(
"sprout-desktop: Pocket TTS returned unexpected sample rate {sample_rate}Hz \
(expected {SAMPLE_RATE}Hz); playback speed may be wrong"
);
}
let mut samples = audio.samples().to_vec();
// For short prompts the prepared text includes a sacrificial ". . "
// prefix to absorb FlowLM/Mimi cold-start (see `SACRIFICIAL_PREFIX`).
// Strip the leading sacrificial audio before returning. Long prompts
// are never trimmed — they have no sacrificial audio, and an early
// natural pause could be mis-detected as the trim boundary.
if prepared.is_short {
trim_leading_cold_start(&mut samples);
}
Ok(samples)
}
}
#[cfg(test)]
mod tests {
use super::*;
// ── prepare_pocket_prompt ────────────────────────────────────────────────
#[test]
fn prepare_prompt_returns_none_for_empty_input() {
assert!(prepare_pocket_prompt("").is_none());
assert!(prepare_pocket_prompt(" ").is_none());
assert!(prepare_pocket_prompt("\n\t ").is_none());
}
/// Helper: the exact leading sequence prepended to every short prompt —
/// 8 spaces of padding followed by the sacrificial `". . "` cold-start
/// absorber. Centralising this keeps the assertions readable.
fn short_prefix() -> String {
let mut s = " ".repeat(SHORT_PROMPT_PAD_SPACES);
s.push_str(SACRIFICIAL_PREFIX);
s
}
#[test]
fn prepare_prompt_pads_and_capitalizes_one_word() {
// The "yep" case Tyler hit in production — bare lowercase one-word
// utterance with no punctuation. Must be padded with the short-prompt
// prefix (8 spaces + ". . " sacrificial), capitalized, terminated,
// with a tight `max_frames` cap to bound runaway gen.
let out = prepare_pocket_prompt("yep").expect("non-empty");
assert_eq!(out.text, format!("{}Yep.", short_prefix()));
assert!(out.is_short, "1-word input is short");
assert_eq!(out.max_frames, Some(SHORT_PROMPT_MAX_FRAMES));
assert!(
SHORT_PROMPT_MAX_FRAMES < SHERPA_ONNX_MAX_FRAMES_DEFAULT,
"short cap must be tighter than the upstream default"
);
}
#[test]
fn prepare_prompt_preserves_existing_punctuation() {
let out = prepare_pocket_prompt("yes!").expect("non-empty");
assert_eq!(out.text, format!("{}Yes!", short_prefix())); // exclamation kept
let out = prepare_pocket_prompt("really?").expect("non-empty");
assert_eq!(out.text, format!("{}Really?", short_prefix()));
}
#[test]
fn prepare_prompt_threshold_is_inclusive_at_four_words() {
// 4 words = short (padded + sacrificial + tight max_frames); 5 words
// = long (no padding, no sacrificial, no overrides — upstream
// defaults stand).
let four = prepare_pocket_prompt("one two three four").expect("non-empty");
assert!(four.is_short, "four-word input should be short");
assert!(
four.text.starts_with(' '),
"four-word input should start with the space pad"
);
assert!(
four.text.contains(SACRIFICIAL_PREFIX),
"four-word input should contain the sacrificial prefix"
);
assert_eq!(four.max_frames, Some(SHORT_PROMPT_MAX_FRAMES));
let five = prepare_pocket_prompt("one two three four five").expect("non-empty");
assert!(!five.is_short, "five-word input should NOT be short");
assert!(
!five.text.starts_with(' '),
"five-word input should NOT be padded"
);
assert!(
!five.text.contains(SACRIFICIAL_PREFIX),
"five-word input must not receive the sacrificial prefix"
);
assert_eq!(
five.max_frames, None,
"long inputs must leave sherpa-onnx's max_frames default in place"
);
}
#[test]
fn prepare_prompt_does_not_pad_long_text() {
let long = "This is a longer sentence that the model should handle just fine.";
let out = prepare_pocket_prompt(long).expect("non-empty");
assert!(!out.is_short);
assert!(!out.text.starts_with(' '));
assert!(!out.text.contains(SACRIFICIAL_PREFIX));
assert_eq!(out.max_frames, None);
assert!(out.text.ends_with('.'));
}
#[test]
fn prepare_prompt_collapses_whitespace() {
let out = prepare_pocket_prompt("Hello world\n\nfriend").expect("non-empty");
// 3 words → short → padded + sacrificial. Interior whitespace
// collapsed.
assert_eq!(out.text, format!("{}Hello world friend.", short_prefix()));
}
#[test]
fn prepare_prompt_does_not_double_capitalize_already_uppercase() {
let out = prepare_pocket_prompt("HELLO there").expect("non-empty");
assert_eq!(out.text, format!("{}HELLO there.", short_prefix()));
}
#[test]
fn prepare_prompt_handles_non_ascii_first_letter() {
// Cyrillic lowercase 'д' → uppercase 'Д'. Must not panic / produce
// mojibake.
let out = prepare_pocket_prompt("дa").expect("non-empty");
assert!(out.text.contains("Дa."));
}
#[test]
fn prepare_prompt_inserts_sacrificial_prefix_only_for_short() {
// Pinning the exact ordering: pad, then ". . ", then cleaned text.
// If this ever flips, the trim algorithm's calibration breaks.
let out = prepare_pocket_prompt("I'm happy.").expect("non-empty");
assert!(out.is_short);
let pad = " ".repeat(SHORT_PROMPT_PAD_SPACES);
let expected = format!("{pad}. . I'm happy.");
assert_eq!(out.text, expected);
assert_eq!(SACRIFICIAL_PREFIX, ". . ");
}
// ── build_generation_extra ───────────────────────────────────────────────
//
// These tests pin down a behaviour we've now regressed twice on:
// 1) Not padding/punctuating short inputs → 40 s of "monster breathing"
// (pre-773a2a1).
// 2) Setting `frames_after_eos = 1` on long inputs → clipped leading
// audio of multi-clause sentences, e.g. "Yep, I can hear you. …"
// came out as a static burst (the 773a2a1 regression Tyler hit on
// 2026-05-18 ~14:30 UTC).
//
// The contract we enforce going forward: we **only** override
// `max_frames`, and only for ≤4-word inputs. Every other knob is left
// at sherpa-onnx's documented default (notably `frames_after_eos = 3`).
#[test]
fn build_extra_short_prompt_sets_only_max_frames() {
let prepared = prepare_pocket_prompt("yep").expect("non-empty");
let extra = build_generation_extra(&prepared).expect("short prompts get extra");
// Exactly one key — `max_frames` — and nothing else.
assert_eq!(extra.len(), 1, "extra has unexpected keys: {extra:?}");
assert_eq!(
extra.get("max_frames"),
Some(&serde_json::Value::from(SHORT_PROMPT_MAX_FRAMES))
);
assert!(
!extra.contains_key("frames_after_eos"),
"frames_after_eos must never be set — upstream default of {SHERPA_ONNX_FRAMES_AFTER_EOS_DEFAULT} is what we want"
);
}
#[test]
fn build_extra_long_prompt_is_none() {
// ≥5 words: no extras at all. This is the key fix for the "first
// 'yep' in 'Yep, I can hear you. …' is static" regression — we
// were previously forcing `frames_after_eos = 1` on this path.
let prepared = prepare_pocket_prompt("Yep, I can hear you.").expect("non-empty");
assert_eq!(
build_generation_extra(&prepared),
None,
"long prompts must not override any LM knob"
);
}
#[test]
fn build_extra_never_lowers_frames_after_eos_for_any_word_count() {
// Sweep a range of prompt lengths and assert the `extra` map (when
// present) never carries a `frames_after_eos` override that's lower
// than the upstream sherpa-onnx default. Implemented as a structural
// check — we just never set the key — but worth a property test in
// case someone reintroduces the override in the future.
let prompts: &[&str] = &[
"hi",
"hi there",
"yes please",
"one two three four",
"one two three four five",
"a slightly longer reply, hopefully fine",
"This is a multi-clause sentence. It has two parts.",
"really really really really really long prompt with lots of words just to be sure",
];
for &p in prompts {
let prepared = prepare_pocket_prompt(p).expect("non-empty");
if let Some(extra) = build_generation_extra(&prepared) {
if let Some(v) = extra.get("frames_after_eos") {
let n = v.as_i64().expect("frames_after_eos should be int");
assert!(
n >= SHERPA_ONNX_FRAMES_AFTER_EOS_DEFAULT as i64,
"prompt {p:?} set frames_after_eos={n}, below upstream default of {SHERPA_ONNX_FRAMES_AFTER_EOS_DEFAULT}"
);
}
}
}
}
#[test]
fn short_prompt_max_frames_is_below_upstream_default() {
// Sanity: the override only ever *lowers* the cap, never raises it.
assert!(SHORT_PROMPT_MAX_FRAMES < SHERPA_ONNX_MAX_FRAMES_DEFAULT);
// …and is still large enough for a one-to-four-word reply. At Mimi's
// 12.5 Hz frame rate, 100 frames = 8 s, which is roomy.
assert!(SHORT_PROMPT_MAX_FRAMES >= 50, "would risk truncation");
}
// ── trim_leading_cold_start ──────────────────────────────────────────────
/// Build a synthetic buffer shaped like a sacrificial-prefixed Pocket TTS
/// output: a bit of "sacrificial" energy at the head (modelled as
/// alternating ±0.05 — above [`TRIM_SILENCE_THRESHOLD`] so it isn't
/// classified as silence, matching what real probe WAVs look like in the
/// 0–50 ms window), then a flat silence of `gap_ms` ms, then `tail_ms`
/// of "real speech" at peak `tail_peak`.
fn synth_buffer(sacrificial_ms: u32, gap_ms: u32, tail_ms: u32, tail_peak: f32) -> Vec<f32> {
let sr = SAMPLE_RATE as usize;
let mut v = Vec::new();
for i in 0..(sr * sacrificial_ms as usize / 1000) {
v.push(if i % 2 == 0 { 0.05 } else { -0.05 });
}
// Silence gap (true zeros — below TRIM_SILENCE_THRESHOLD).
v.extend(std::iter::repeat_n(0.0_f32, sr * gap_ms as usize / 1000));
// Real speech.
for i in 0..(sr * tail_ms as usize / 1000) {
v.push(if i % 2 == 0 { tail_peak } else { -tail_peak });
}
v
}
#[test]
fn trim_strips_sacrificial_and_keeps_only_speech() {
// 60 ms sacrificial + 100 ms gap + 500 ms speech at peak 0.3.
// After trim, the output is just the speech tail.
let mut v = synth_buffer(60, 100, 500, 0.3);
trim_leading_cold_start(&mut v);
// First sample should be speech (|s| ≥ 0.2). No zero lead-in here
// because tts.rs's `first_append` lead-in handles the device cushion.
assert!(
v[0].abs() > 0.2,
"first sample after trim should be speech, got {}",
v[0]
);
let actual_ms = (v.len() as f32 / SAMPLE_RATE as f32) * 1000.0;
assert!(
(actual_ms - 500.0).abs() < 5.0,
"expected ~500 ms of trimmed audio, got {actual_ms} ms"
);
}
#[test]
fn trim_is_noop_when_no_long_silence_gap_exists() {
// Pure speech: every sample is real (no gap >= 50 ms). Trimmer must
// leave the buffer untouched so we don't truncate the utterance.
let mut v = synth_buffer(0, 0, 600, 0.3);
let before = v.clone();
trim_leading_cold_start(&mut v);
assert_eq!(v, before, "no gap → no trim");
}
#[test]
fn trim_is_noop_when_gap_is_shorter_than_threshold() {
// 40 ms gap is below TRIM_MIN_GAP_SAMPLES (50 ms). Must not trigger.
let mut v = synth_buffer(60, 40, 600, 0.3);
let before = v.clone();
trim_leading_cold_start(&mut v);
assert_eq!(v, before, "sub-threshold gap → no trim");
}
#[test]
fn trim_is_noop_when_gap_is_beyond_max_drop_bound() {
// Gap starts at 1500 ms (past TRIM_MAX_DROP_SAMPLES = 1200 ms).
// This represents an interior pause inside an unusually long
// utterance that slipped past the short-prompt predicate; we must
// not chop the first 1.5 s of real audio.
let mut v = synth_buffer(1500, 200, 400, 0.3);
let before = v.clone();
trim_leading_cold_start(&mut v);
assert_eq!(v, before, "gap past max-drop bound → no trim");
}
#[test]
fn trim_is_noop_on_buffer_smaller_than_scan_start() {
// 20 ms buffer is smaller than TRIM_SCAN_START_SAMPLES (30 ms).
// Trimmer must early-return without panicking.
let mut v = vec![0.5f32; (SAMPLE_RATE as usize * 20) / 1000];
let before = v.clone();
trim_leading_cold_start(&mut v);
assert_eq!(v, before);
}
#[test]
fn trim_constants_use_sane_units() {
// Pin the constants in milliseconds so anyone tuning later can see
// at a glance what they're changing.
assert_eq!(
TRIM_SCAN_START_SAMPLES,
(SAMPLE_RATE as usize * 30) / 1000,
"scan-start should be 30 ms"
);
assert_eq!(
TRIM_MIN_GAP_SAMPLES,
(SAMPLE_RATE as usize * 50) / 1000,
"min-gap should be 50 ms"
);
assert_eq!(
TRIM_MAX_DROP_SAMPLES,
(SAMPLE_RATE as usize * 1200) / 1000,
"max-drop should be 1.2 s"
);
}
}

View File

@@ -11,10 +11,7 @@
//! ```
use futures_util::{SinkExt, StreamExt};
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
};
use std::sync::{atomic::AtomicBool, Arc};
use tokio_tungstenite::{connect_async, tungstenite::Message as WsMsg};
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
@@ -112,6 +109,12 @@ pub(crate) async fn connect_audio_relay(
"type": "auth",
"event": event_json,
"parent_channel_id": parent_channel_id,
// Negotiate huddle audio protocol v2 (8-byte sender-authored header
// per Opus frame: seq | ts_48k | level_dbov | flags). See
// huddle::wire for the layout. The relay pins the first joiner's
// version per-room and rejects mismatched joiners with
// `upgrade_required`.
"protocol_version": super::wire::PROTOCOL_VERSION,
});
ws_tx
.send(WsMsg::Text(auth_msg.to_string().into()))
@@ -198,12 +201,12 @@ pub(crate) async fn connect_audio_relay(
}
/// Background Opus encode/decode pipeline spawned by `connect_audio_relay`.
type WsStream =
pub(crate) type WsStream =
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
async fn audio_relay_pipeline(
ws_tx: futures_util::stream::SplitSink<WsStream, WsMsg>,
mut ws_rx: futures_util::stream::SplitStream<WsStream>,
ws_rx: futures_util::stream::SplitStream<WsStream>,
mut pcm_rx: tokio::sync::mpsc::Receiver<Vec<u8>>,
cancel: CancellationToken,
app_handle: Option<tauri::AppHandle>,
@@ -222,11 +225,6 @@ async fn audio_relay_pipeline(
.map_err(|e| format!("opus dtx: {e}"))?;
let sink_handle = super::audio_output::open_output_sink_by_name(output_device_name.as_deref())?;
let player = rodio::Player::connect_new(&sink_handle.mixer());
let decoders: std::collections::HashMap<u8, opus::Decoder> = std::collections::HashMap::new();
const FRAME_SAMPLES: usize = 960; // 20 ms at 48 kHz
let decode_buf = vec![0f32; FRAME_SAMPLES];
use std::sync::Arc as StdArc;
let ws_tx = StdArc::new(tokio::sync::Mutex::new(ws_tx));
@@ -234,9 +232,15 @@ async fn audio_relay_pipeline(
let cancel_send = cancel.clone();
let send_task = tokio::spawn(async move {
use super::wire::{audio_level_dbov, FrameHeader, V2_HEADER_LEN};
let mut encoder = encoder; // Move encoder into task.
const FRAME_SAMPLES: usize = 960;
let mut out_buf = vec![0u8; 4000];
// Per-frame wire-protocol state. We send v2 frames now: each Opus
// payload is preceded by an 8-byte header carrying our own seq +
// 48 kHz timestamp + audio level + flags.
let mut seq: u16 = 0;
let mut ts_48k: u32 = 0;
loop {
let pcm_bytes = {
@@ -260,6 +264,10 @@ async fn audio_relay_pipeline(
let mut tx = ws_tx_send.lock().await;
for chunk in samples.chunks(FRAME_SAMPLES) {
// dBov is computed from the pre-encode PCM. Opus DTX may
// produce a 1-2 byte comfort packet; computing level from
// the encoded payload would be meaningless.
let level = audio_level_dbov(chunk);
let encode_result = if chunk.len() == FRAME_SAMPLES {
encoder.encode_float(chunk, &mut out_buf)
} else {
@@ -275,14 +283,28 @@ async fn audio_relay_pipeline(
}
};
if n > 0 {
// Send raw Opus bytes — the relay prepends the peer_index.
if tx
.send(WsMsg::Binary(out_buf[..n].to_vec().into()))
.await
.is_err()
{
// Opus DTX packets are very small (≤2 bytes). Flag them
// explicitly so the receiver can elide DTX from speaker
// detection without re-parsing the Opus payload.
let flags = if n <= 2 { super::wire::FLAG_DTX } else { 0 };
let header = FrameHeader {
seq,
ts_48k,
level_dbov: level,
flags,
}
.encode();
// Build the v2 wire frame: 8-byte header + Opus payload.
let mut frame = Vec::with_capacity(V2_HEADER_LEN + n);
frame.extend_from_slice(&header);
frame.extend_from_slice(&out_buf[..n]);
if tx.send(WsMsg::Binary(frame.into())).await.is_err() {
return; // WS closed.
}
seq = seq.wrapping_add(1);
ts_48k = ts_48k.wrapping_add(super::jitter::FRAME_TIMESTAMP_DELTA);
}
}
}
@@ -290,154 +312,16 @@ async fn audio_relay_pipeline(
let _ = tx.send(WsMsg::Close(None)).await;
});
let recv_task = tokio::spawn(async move {
let mut decoders = decoders;
let mut decode_buf = decode_buf;
let cancel_recv = cancel;
let ws_tx_recv = ws_tx;
// Active speaker tracking: peer_index → pubkey, emit every 500ms.
let mut index_to_pubkey: std::collections::HashMap<u8, String> =
initial_peers.into_iter().collect();
let mut active_indices: std::collections::HashSet<u8> = std::collections::HashSet::new();
// Per-peer frame counter with Instant-based window (starvation-proof).
let mut frame_counts: std::collections::HashMap<u8, u16> = std::collections::HashMap::new();
let mut last_frame_reset = tokio::time::Instant::now();
const FRAME_WINDOW: std::time::Duration = std::time::Duration::from_millis(500);
let mut tts_was_active = false;
let mut speaker_tick = tokio::time::interval(std::time::Duration::from_millis(500));
speaker_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
use futures_util::future::Either;
let cancelled = std::pin::pin!(cancel_recv.cancelled());
let next = std::pin::pin!(ws_rx.next());
let tick = std::pin::pin!(speaker_tick.tick());
// Three-way select: cancel, WS message, or speaker tick.
let ws_or_tick = std::pin::pin!(futures_util::future::select(next, tick));
match futures_util::future::select(cancelled, ws_or_tick).await {
Either::Left(_) => break, // Cancelled.
Either::Right((Either::Right((_, _)), _)) => {
// Speaker tick: emit active speakers and reset.
if let Some(ref app) = app_handle {
use tauri::Emitter;
let pubkeys: Vec<String> = active_indices
.iter()
.filter_map(|idx| index_to_pubkey.get(idx).cloned())
.collect();
let _ = app.emit("huddle-active-speakers", &pubkeys);
}
active_indices.clear();
continue;
}
Either::Right((Either::Left((Some(Ok(msg)), _)), _)) => {
// WS message — process below.
match msg {
WsMsg::Binary(data) => {
if data.len() < 2 {
continue;
}
let peer_idx = data[0];
let opus_bytes = &data[1..];
active_indices.insert(peer_idx);
// Track TTS transitions at frame level (not decode level).
let tts_now = tts_active.load(Ordering::Acquire);
if tts_now && !tts_was_active {
frame_counts.clear();
last_frame_reset = tokio::time::Instant::now();
}
tts_was_active = tts_now;
let decoder = match decoders.entry(peer_idx) {
std::collections::hash_map::Entry::Occupied(e) => e.into_mut(),
std::collections::hash_map::Entry::Vacant(e) => {
match opus::Decoder::new(48000, opus::Channels::Mono) {
Ok(d) => e.insert(d),
Err(err) => {
eprintln!("sprout-desktop: opus decoder init peer {peer_idx}: {err}");
continue;
}
}
}
};
match decoder.decode_float(opus_bytes, &mut decode_buf, false) {
Ok(n) if n > 0 => {
if tts_now {
if last_frame_reset.elapsed() >= FRAME_WINDOW {
frame_counts.clear();
last_frame_reset = tokio::time::Instant::now();
}
let count = frame_counts.entry(peer_idx).or_insert(0);
*count = count.saturating_add(1);
if *count >= REMOTE_SPEECH_THRESHOLD {
tts_cancel.store(true, Ordering::Release);
}
}
use rodio::buffer::SamplesBuffer;
use std::num::NonZero;
let channels = NonZero::new(1u16).unwrap();
let rate = NonZero::new(48000u32).unwrap();
player.append(SamplesBuffer::new(
channels,
rate,
decode_buf[..n].to_vec(),
));
}
Ok(_) => {} // DTX silence — not counted.
Err(e) => {
eprintln!("sprout-desktop: opus decode peer {peer_idx}: {e}");
}
}
}
WsMsg::Text(text) => {
if let Ok(v) = serde_json::from_str::<serde_json::Value>(&text) {
match v["type"].as_str() {
Some("joined") => {
if let Some(peers) = v["peers"].as_array() {
for p in peers {
if let (Some(pk), Some(idx)) =
(p["pubkey"].as_str(), p["peer_index"].as_u64())
{
let key = idx as u8;
if index_to_pubkey.get(&key).map(|s| s.as_str())
!= Some(pk)
{
decoders.remove(&key);
frame_counts.remove(&key);
active_indices.remove(&key);
}
index_to_pubkey.insert(key, pk.to_string());
}
}
}
}
Some("left") => {
if let Some(idx) = v["peer_index"].as_u64() {
let key = idx as u8;
index_to_pubkey.remove(&key);
frame_counts.remove(&key);
decoders.remove(&key);
}
}
_ => {}
}
}
}
WsMsg::Ping(data) => {
let mut tx = ws_tx_recv.lock().await;
let _ = tx.send(WsMsg::Pong(data)).await;
}
WsMsg::Close(_) => break,
_ => {}
}
}
// WS error or closed.
Either::Right((Either::Left(_), _)) => break,
}
}
});
let recv_task = tokio::spawn(super::playout::run_playout_recv_loop(
ws_rx,
ws_tx,
sink_handle,
cancel,
app_handle,
initial_peers,
tts_active,
tts_cancel,
));
// Wait for either task to finish, then abort the survivor.
use futures_util::future::Either;

View File

@@ -5,7 +5,7 @@
//! ```text
//! caller: pipeline.speak("Hello world. How are you?")
//! → bounded sync_channel (TEXT_QUEUE_DEPTH = 8)
//! → tts_worker thread (owns 1 Kokoro engine)
//! → tts_worker thread (owns 1 Pocket TTS engine)
//! 1. Preprocess text
//! 2. Split into sentences
//! 3. Synthesize each sentence individually → f32 PCM
@@ -35,7 +35,7 @@ use std::{
time::Duration,
};
use super::kokoro::{load_text_to_speech, load_voice_style, SAMPLE_RATE};
use super::pocket::{load_text_to_speech, load_voice_style, SAMPLE_RATE, VOICE_FILE_EXT};
use super::preprocessing::{preprocess_for_tts, split_sentences};
// ── Constants ─────────────────────────────────────────────────────────────────
@@ -48,21 +48,62 @@ const TEXT_QUEUE_DEPTH: usize = 8;
/// How long the worker waits on the text channel before checking the shutdown flag.
const RECV_TIMEOUT: Duration = Duration::from_millis(100);
/// Kokoro ignores denoising steps (not a diffusion model). Kept for API compat.
/// Pocket TTS is a one-step consistency model, not diffusion. Kept for API compat.
const SYNTH_STEPS: usize = 1;
/// Synthesis speed multiplier. Slightly faster than natural speech.
const SYNTH_SPEED: f32 = 1.05;
/// Volume boost applied after synthesis — Kokoro output is normalized.
/// Start at 1.5 and tune empirically.
const VOLUME_BOOST: f32 = 1.5;
/// Upper bound on per-sentence loudness normalization, in linear scale.
/// −3 dBFS = 10^(−3/20) ≈ 0.708. This is a *ceiling*, not a floor: any
/// sentence whose computed gain would push its peak above this is held at
/// the ceiling, and quieter utterances under the [`MAX_GAIN`] cap may land
/// below it. The ceiling leaves 3 dB of headroom above the loudest sample
/// so the subsequent fade-out and any system mixer gain don't soft-clip.
///
/// Tyler bumped this from the previous −6 dBFS (0.501) to give Pocket TTS
/// output some perceived loudness — the model's reference voice is quieter
/// than Kokoro's was, and the prior target landed normal utterances around
/// −9 dBFS once interior dynamics are accounted for.
const TARGET_PEAK: f32 = 0.707_945_8; // 10f32.powf(-3.0 / 20.0)
/// Fade in/out length in samples (8ms at 24kHz ≈ 192 samples).
/// Eliminates clicks/pops at sentence boundaries.
const FADE_SAMPLES: usize = (SAMPLE_RATE as f64 * 0.008) as usize;
/// Maximum gain applied by `normalize_for_playback`. Caps amplification on
/// near-silent buffers so a mid-utterance pause or a malformed synth doesn't
/// get amplified to full scale (which would surface any quantization noise).
///
/// Pocket TTS reference-voice output measured ~7.6% peak on a 75-character
/// utterance (`examples/pocket_bench`); a gain of 1/0.076 ≈ 6.6 lands that
/// sample at −6 dBFS and ≈9.3 at −3 dBFS. `8.0` covers normal utterances at
/// the new ceiling while still catching pathological near-silent buffers
/// (which will land at peak ≤ 0.076 × 8 = 0.61 — below the ceiling, which
/// is the intended behaviour).
const MAX_GAIN: f32 = 8.0;
/// Sentence-by-sentence synthesis for lower TTFA (≈200ms vs ≈600ms for 3-sentence batches).
/// Fade-out length in samples (8 ms at 24 kHz ≈ 192 samples).
///
/// Applied only at the *end* of each synthesised sentence to eliminate the
/// click that would otherwise occur when a non-zero waveform terminates
/// abruptly. **No fade-in is applied** — see `apply_fade_out` for the
/// rationale and `examples/pocket_onset_probe.rs` for the measurement that
/// motivated removing the leading fade.
const FADE_OUT_SAMPLES: usize = (SAMPLE_RATE as f64 * 0.008) as usize;
/// Length of the zero-sample cushion prepended before each synthesized
/// sentence chunk, so the OS audio device / rodio mixer has a fully-quiet
/// ramp-up window before the real onset hits.
///
/// This used to be applied only before the first sentence of a whole response.
/// That still left later sentence chunks vulnerable to first-syllable clipping
/// when their first phoneme was soft (notably `I'm` / `I've`) and rodio crossed
/// from an explicit silence buffer straight into non-zero speech. 20 ms ≈ 480
/// samples is enough to cover a CoreAudio buffer turnover without being audible
/// as latency. At sentence boundaries this lead-in is budgeted out of the
/// existing inter-sentence pause, so it does not lengthen multi-sentence gaps.
const SENTENCE_LEAD_IN_SAMPLES: usize = (SAMPLE_RATE as f64 * 0.020) as usize;
/// Sentence-by-sentence synthesis — keeps first-sentence latency low and lets
/// playback of sentence N overlap with synthesis of sentence N+1 (see the
/// lookahead pipelining note in the module doc-comment above).
const BATCH_SIZE: usize = 1;
/// Silence inserted between sentences by the TTS pipeline (seconds).
@@ -87,7 +128,7 @@ pub struct TtsPipeline {
/// Kept alive here so the Arc isn't dropped — the worker holds a clone.
#[allow(dead_code)]
cancel: Arc<AtomicBool>,
/// Voice name (e.g. "af_heart"). Stored for future voice-switching support.
/// Voice name (e.g. "reference_sample"). Stored for future voice-switching support.
#[allow(dead_code)]
voice: String,
/// Worker thread handle — taken on drop to join cleanly.
@@ -97,8 +138,8 @@ pub struct TtsPipeline {
impl TtsPipeline {
/// Spawn the TTS pipeline thread using the default voice.
///
/// `model_dir` must contain the Kokoro model files:
/// `model_quantized.onnx`, `tokenizer.json`, `voices/<name>.bin`
/// `model_dir` must contain the Pocket TTS files declared by `huddle::models`
/// (the five ONNX sessions, the two JSON tables, and `<voice>.wav`).
///
/// `tts_active` is set to `true` while audio is playing and `false` when idle.
/// Pass the same `Arc` to the STT pipeline to gate microphone input.
@@ -112,11 +153,13 @@ impl TtsPipeline {
cancel: Arc<AtomicBool>,
output_device: Option<String>,
) -> Result<Self, String> {
use super::kokoro::DEFAULT_VOICE;
use super::pocket::DEFAULT_VOICE;
Self::new_with_voice(model_dir, tts_active, cancel, DEFAULT_VOICE, output_device)
}
/// Spawn the TTS pipeline thread with a specific voice name (e.g. `"af_heart"`, `"am_michael"`).
/// Spawn the TTS pipeline thread with a specific voice name. Today only the
/// bundled default voice (see `pocket::DEFAULT_VOICE`) is shipped; other
/// names will surface a clear error from `load_voice_style`.
pub fn new_with_voice(
model_dir: PathBuf,
tts_active: Arc<AtomicBool>,
@@ -204,14 +247,14 @@ fn tts_worker(
cancel: Arc<AtomicBool>,
output_device: Option<String>,
) {
// ── 1. Initialise Kokoro engine ───────────────────────────────────────────
// ── 1. Initialise TTS engine ──────────────────────────────────────────────
let model_dir_str = model_dir.to_string_lossy().to_string();
let mut engine = match load_text_to_speech(&model_dir_str) {
let engine = match load_text_to_speech(&model_dir_str) {
Ok(e) => e,
Err(e) => {
eprintln!(
"sprout-desktop: TTS Kokoro init failed (model_dir={}): {e}. TTS disabled.",
"sprout-desktop: TTS engine init failed (model_dir={}): {e}. TTS disabled.",
model_dir.display()
);
drain_until_shutdown(text_rx, &shutdown);
@@ -220,7 +263,7 @@ fn tts_worker(
};
// ── 2. Load voice style ───────────────────────────────────────────────────
let voice_path = model_dir.join(format!("{voice_name}.bin"));
let voice_path = model_dir.join(format!("{voice_name}.{VOICE_FILE_EXT}"));
let style = match load_voice_style(&voice_path) {
Ok(s) => s,
Err(e) => {
@@ -234,9 +277,9 @@ fn tts_worker(
// ── 2b. Warmup inference ─────────────────────────────────────────────────
// The first ONNX inference on any session is significantly slower than
// subsequent ones — it triggers JIT compilation, memory pool allocation,
// and (on CoreML) lazy model compilation. Run a short dummy synthesis and
// discard the output so the first real utterance runs at warm-session speed.
// subsequent ones — it can trigger native session initialization, memory
// pool allocation, and graph-specific caches. Run a short dummy synthesis
// and discard the output so the first real utterance runs at warm-session speed.
{
let t = std::time::Instant::now();
match engine.synth_chunk("warmup", "en", &style, SYNTH_STEPS, SYNTH_SPEED) {
@@ -368,17 +411,23 @@ fn tts_worker(
match engine.synth_chunk(text, "en", &style, SYNTH_STEPS, SYNTH_SPEED) {
Ok(samples) if !samples.is_empty() => {
let mut boosted: Vec<f32> = samples
.iter()
.map(|&s| (s * VOLUME_BOOST).clamp(-1.0, 1.0))
.collect();
apply_fades(&mut boosted);
player.append(SamplesBuffer::new(channels, rate, boosted));
// Insert inter-sentence silence after each synthesized chunk.
player.append(SamplesBuffer::new(channels, rate, silence_buf.clone()));
if first_append {
let mut boosted = normalize_for_playback(samples);
// Fade-out only — fading-in would attenuate the consonant
// onset (see `apply_fade_out` docstring + the
// 2026-05-18 "first little sound is missing" regression).
apply_fade_out(&mut boosted);
// Build one contiguous buffer per synthesized sentence:
// lead-in cushion + audio + trailing gap. Keeping this as
// a single rodio source preserves the original queue/drain
// semantics (one append per sentence) while still giving
// every chunk a quiet device warm-up window.
let was_first = first_append;
let buf =
build_sentence_append_buffer(&mut first_append, boosted, silence_buf.len());
player.append(SamplesBuffer::new(channels, rate, buf));
if was_first {
tts_active.store(true, Ordering::Release);
first_append = false;
}
}
Ok(_) => {}
@@ -439,23 +488,96 @@ fn handle_cancel_or_shutdown(
false
}
/// Apply a short linear fade-in at the start and fade-out at the end of `samples`.
/// Per-sentence peak normalization. Scales the buffer so its loudest sample
/// lands at `TARGET_PEAK` (−6 dBFS), capped at `MAX_GAIN` to avoid amplifying
/// near-silent buffers into quantization noise. Returns the input unchanged on
/// empty input or pure-silence input.
///
/// Uses `FADE_SAMPLES` (8ms) or half the buffer length, whichever is smaller.
/// Eliminates clicks/pops at sentence boundaries.
fn apply_fades(samples: &mut Vec<f32>) {
let len = samples.len();
let fade = FADE_SAMPLES.min(len / 2);
// Fade in: ramp from 0 → 1 over `fade` samples.
for i in 0..fade {
samples[i] *= i as f32 / fade as f32;
/// Why normalize per-sentence rather than apply a fixed boost: Pocket TTS
/// reference-voice output measured ~7.6% peak on one utterance, but per-
/// sentence peak distribution isn't known in advance and may vary with text
/// and any future voice swap. A fixed multiplier either clips loud sentences
/// or under-amplifies quiet ones; normalization adapts.
fn normalize_for_playback(samples: Vec<f32>) -> Vec<f32> {
let peak = samples.iter().fold(0.0_f32, |acc, &s| acc.max(s.abs()));
if peak == 0.0 {
return samples;
}
// Fade out: ramp from 1 → 0 over the last `fade` samples.
let gain = (TARGET_PEAK / peak).min(MAX_GAIN);
samples
.into_iter()
.map(|s| (s * gain).clamp(-1.0, 1.0))
.collect()
}
/// Apply a short linear fade-out at the *end* of `samples`.
///
/// Uses `FADE_OUT_SAMPLES` (8 ms) or half the buffer length, whichever is
/// smaller. Eliminates the click that occurs when a non-zero waveform
/// terminates abruptly at a sentence boundary.
///
/// # Why no fade-in
///
/// An earlier revision (pre 2026-05) symmetrically faded *in* over the same
/// 8 ms window. That swallowed the leading consonant attack on every
/// sentence — Pocket TTS produces real audio energy inside the first
/// millisecond (RMS ≈ 0.02, peak ≈ 0.03 measured across four prompts in
/// `examples/pocket_onset_probe.rs`), and a linear 0→1 ramp over 192 samples
/// scales those onset samples by ≤50 % for the first ~4 ms. The result was
/// the "first little sound or two is missing" regression heard on
/// 2026-05-18.
///
/// The first sample of Pocket output measures ≈ 0.0018 (≈ −54 dBFS) — well
/// below the threshold at which a DC-jump would be audible as a click — so
/// no fade-in is needed. The OS audio device gets its quiet ramp-up window
/// from `SENTENCE_LEAD_IN_SAMPLES` instead, inserted as pure silence before
/// each sentence buffer.
fn apply_fade_out(samples: &mut [f32]) {
let len = samples.len();
let fade = FADE_OUT_SAMPLES.min(len / 2);
for i in 0..fade {
samples[len - 1 - i] *= i as f32 / fade as f32;
}
}
/// Build the single buffer appended to the rodio `Player` for one synthesised
/// sentence.
///
/// Every sentence chunk gets a short lead-in pad immediately before its audio.
/// This matters for chunks that start with soft first phonemes (`I'm`, `I've`):
/// the sacrificial-prefix trim intentionally returns audio whose first sample is
/// already speech, so the playback layer must provide the device/mixer cushion.
/// To keep the audible gap unchanged, the trailing silence after this chunk is
/// shortened by the same amount (`silence_buf_len - SENTENCE_LEAD_IN_SAMPLES`):
/// sentence N contributes 80 ms of post-speech silence and sentence N+1
/// contributes the remaining 20 ms of pre-speech cushion.
///
/// The lead-in, audio, and trailing silence are concatenated into one
/// `SamplesBuffer` before appending. This keeps rodio's queue shape at one
/// tracked source per synthesized sentence, avoiding source-boundary/drain
/// regressions from enqueueing the lead-in, audio, and tail as separate sounds.
///
/// `first_append` is still accepted and flipped on the first call because the
/// worker uses it to decide when actual playback has been queued and when it is
/// safe to set `tts_active` for echo gating.
fn build_sentence_append_buffer(
first_append: &mut bool,
boosted: Vec<f32>,
silence_buf_len: usize,
) -> Vec<f32> {
if *first_append {
*first_append = false;
}
let trailing_silence_len = silence_buf_len.saturating_sub(SENTENCE_LEAD_IN_SAMPLES);
let mut buf =
Vec::with_capacity(SENTENCE_LEAD_IN_SAMPLES + boosted.len() + trailing_silence_len);
buf.extend(std::iter::repeat_n(0.0_f32, SENTENCE_LEAD_IN_SAMPLES));
buf.extend(boosted);
buf.extend(std::iter::repeat_n(0.0_f32, trailing_silence_len));
buf
}
// drain_until_shutdown lives in super (huddle/mod.rs) — shared with stt.rs.
use super::drain_until_shutdown;
@@ -1001,28 +1123,241 @@ mod tests {
);
}
// ── apply_fades tests ─────────────────────────────────────────────────────
// ── apply_fade_out tests ──────────────────────────────────────────────────
/// The fade-out half of the old `apply_fades`: last sample is silenced
/// and the ramp is monotonic. Mid-buffer must be untouched.
#[test]
fn apply_fades_short_buffer() {
fn apply_fade_out_short_buffer() {
let mut samples = vec![1.0f32; 10];
apply_fades(&mut samples);
assert_eq!(samples[0], 0.0);
assert_eq!(samples[9], 0.0);
assert!(samples[5] > 0.5);
apply_fade_out(&mut samples);
assert_eq!(samples[9], 0.0, "last sample should be silenced");
assert!(samples[5] > 0.5, "mid-buffer should be near-untouched");
}
/// REGRESSION (2026-05-18): the *first* samples must NOT be attenuated.
/// An earlier `apply_fades` symmetrically faded in over 8 ms which
/// swallowed the consonant onset of every sentence.
/// Lock in: samples[0..FADE_OUT_SAMPLES] are byte-equal to input.
#[test]
fn apply_fade_out_does_not_touch_leading_samples() {
// Input long enough that fade window doesn't overlap (≫ 2× fade).
let n = FADE_OUT_SAMPLES * 4;
let input: Vec<f32> = (0..n).map(|i| 0.5 + (i as f32) * 1e-4).collect();
let mut samples = input.clone();
apply_fade_out(&mut samples);
for i in 0..FADE_OUT_SAMPLES {
assert_eq!(
samples[i], input[i],
"leading sample {i} must not be attenuated (was {} → {})",
input[i], samples[i]
);
}
// And the trailing fade still works.
assert_eq!(samples[n - 1], 0.0);
}
#[test]
fn apply_fades_empty_buffer() {
fn apply_fade_out_empty_buffer() {
let mut samples: Vec<f32> = vec![];
apply_fades(&mut samples);
apply_fade_out(&mut samples);
assert!(samples.is_empty());
}
#[test]
fn apply_fades_single_sample() {
fn apply_fade_out_single_sample() {
// fade = min(FADE_OUT_SAMPLES, len/2) = 0, so nothing changes.
let mut samples = vec![1.0f32];
apply_fades(&mut samples);
apply_fade_out(&mut samples);
assert_eq!(samples[0], 1.0);
}
/// Sanity-check the per-sentence cushion length: 20 ms at 24 kHz must
/// land at exactly 480 samples. This is a const computation, so the
/// real value of this test is documenting *why* 20 ms was chosen — it
/// covers a typical CoreAudio buffer turnover (256–1024 samples)
/// without being audible as user-facing latency.
#[test]
fn sentence_lead_in_is_sane() {
assert_eq!(SENTENCE_LEAD_IN_SAMPLES, 480, "20 ms × 24 kHz");
}
// ── build_sentence_append_buffer tests ───────────────────────────────────
/// REGRESSION: every chunk needs an onset cushion; trimmed short chunks
/// can start with speech at sample 0.
#[test]
fn lead_in_pad_is_present_for_every_sentence_chunk() {
const SENTENCE_AUDIO_LEN: usize = 1000;
const SILENCE_BUF_LEN: usize = 2400; // 100 ms at 24 kHz, like production
const N_SENTENCES: usize = 5;
let mut first = true;
for _ in 0..N_SENTENCES {
let buf = build_sentence_append_buffer(
&mut first,
vec![0.5_f32; SENTENCE_AUDIO_LEN],
SILENCE_BUF_LEN,
);
assert_eq!(buf.len(), SENTENCE_AUDIO_LEN + SILENCE_BUF_LEN);
assert!(
buf[..SENTENCE_LEAD_IN_SAMPLES].iter().all(|&s| s == 0.0),
"lead-in pad must be pure silence"
);
assert!(
buf[SENTENCE_LEAD_IN_SAMPLES..SENTENCE_LEAD_IN_SAMPLES + SENTENCE_AUDIO_LEN]
.iter()
.all(|&s| s == 0.5),
"sentence audio must immediately follow the lead-in"
);
assert!(
buf[SENTENCE_LEAD_IN_SAMPLES + SENTENCE_AUDIO_LEN..]
.iter()
.all(|&s| s == 0.0),
"trailing gap must be pure silence"
);
}
assert!(!first, "first_append flag must be cleared after first call");
}
/// `first_append` still flips on the first call for `tts_active` gating.
#[test]
fn build_sentence_append_buffer_flips_first_append() {
let mut first = true;
let _ = build_sentence_append_buffer(&mut first, vec![0.5; 100], 2400);
assert!(!first, "first call must flip the flag");
// Subsequent call: still has a per-sentence lead-in, flag stays false.
let buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], 2400);
assert!(buf[..SENTENCE_LEAD_IN_SAMPLES].iter().all(|&s| s == 0.0));
assert!(!first);
}
/// Leading silence is exactly the lead-in; no pre-audio gap is double-counted.
#[test]
fn first_sentence_leading_silence_is_exactly_lead_in() {
let mut first = true;
let buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], 2400);
assert!(buf[..SENTENCE_LEAD_IN_SAMPLES].iter().all(|&s| s == 0.0));
assert_eq!(buf[SENTENCE_LEAD_IN_SAMPLES], 0.5);
}
/// Tail silence plus the next lead-in preserves the 100 ms sentence gap.
#[test]
fn sentence_gap_budget_is_preserved() {
let mut first = true;
let silence_buf_len = 2400;
let first_buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], silence_buf_len);
let second_buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], silence_buf_len);
let first_tail = &first_buf[SENTENCE_LEAD_IN_SAMPLES + 100..];
let second_lead = &second_buf[..SENTENCE_LEAD_IN_SAMPLES];
assert_eq!(first_tail.len(), silence_buf_len - SENTENCE_LEAD_IN_SAMPLES);
assert_eq!(second_lead.len(), SENTENCE_LEAD_IN_SAMPLES);
assert_eq!(first_tail.len() + second_lead.len(), silence_buf_len);
}
/// Regression guard: one contiguous rodio source per synthesized sentence.
#[test]
fn sentence_append_buffer_is_one_contiguous_source() {
let mut first = true;
let buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], 2400);
assert_eq!(buf.len(), 2400 + 100);
assert!(buf[..SENTENCE_LEAD_IN_SAMPLES].iter().all(|&s| s == 0.0));
assert!(
buf[SENTENCE_LEAD_IN_SAMPLES..SENTENCE_LEAD_IN_SAMPLES + 100]
.iter()
.all(|&s| s == 0.5)
);
}
// ── normalize_for_playback tests ──────────────────────────────────────────
/// A moderately quiet buffer (peak comfortably under MAX_GAIN reach of
/// TARGET_PEAK) is scaled up to exactly TARGET_PEAK. With TARGET_PEAK at
/// −3 dBFS (0.708) and MAX_GAIN at 8.0, the gain saturation point is
/// peak = 0.708 / 8.0 ≈ 0.0885 — so a 0.1-peak buffer is in the linear
/// region and gets normalized cleanly.
#[test]
fn normalize_for_playback_hits_target_on_quiet_buffer() {
// peak 0.1 ⇒ ideal gain 7.08, under MAX_GAIN (8.0), so target hit.
let mut input: Vec<f32> = (0..100).map(|i| 0.1 * (i as f32 / 100.0)).collect();
input.push(0.1); // ensure exact peak
let out = normalize_for_playback(input);
let peak = out.iter().fold(0.0_f32, |a, &s| a.max(s.abs()));
assert!(
(peak - TARGET_PEAK).abs() < 1e-3,
"expected peak ~{TARGET_PEAK}, got {peak}"
);
}
/// A buffer whose ideal gain exceeds MAX_GAIN gets clamped — the peak
/// lands at input_peak × MAX_GAIN, below TARGET_PEAK. With the new
/// −3 dBFS target this captures the Pocket-bench-typical case
/// (input peak 0.076 → ideal gain ≈ 9.3 > MAX_GAIN, so we clamp).
#[test]
fn normalize_for_playback_clamps_at_max_gain_below_target() {
let mut input: Vec<f32> = (0..100).map(|i| 0.076 * (i as f32 / 100.0)).collect();
input.push(0.076);
let out = normalize_for_playback(input);
let peak = out.iter().fold(0.0_f32, |a, &s| a.max(s.abs()));
let expected = 0.076 * MAX_GAIN;
assert!(
(peak - expected).abs() < 1e-3,
"expected peak ~{expected}, got {peak}"
);
assert!(
peak < TARGET_PEAK,
"clamped peak {peak} should be below TARGET_PEAK {TARGET_PEAK}"
);
}
/// A near-silent buffer would need a huge gain to reach TARGET_PEAK;
/// `MAX_GAIN` caps the amplification so we don't bring quantization noise
/// up to full scale.
#[test]
fn normalize_for_playback_caps_gain_on_near_silent_buffer() {
// peak 0.001 ⇒ ideal gain 501; MAX_GAIN caps it at 8.0, so the
// resulting peak is 0.001 × 8.0 = 0.008, NOT TARGET_PEAK.
let input = vec![0.001_f32, -0.001, 0.001];
let out = normalize_for_playback(input);
let peak = out.iter().fold(0.0_f32, |a, &s| a.max(s.abs()));
assert!(
(peak - 0.008).abs() < 1e-6,
"expected peak 0.008 (gain-capped), got {peak}"
);
assert!(peak < TARGET_PEAK);
}
/// A pure-silence buffer is returned unchanged — no division by zero, no
/// amplification of nothing.
#[test]
fn normalize_for_playback_silence_is_unchanged() {
let input = vec![0.0_f32; 16];
let out = normalize_for_playback(input);
assert!(out.iter().all(|&s| s == 0.0));
assert_eq!(out.len(), 16);
}
/// Empty input round-trips to empty output. No panic on `peak == 0` path.
#[test]
fn normalize_for_playback_empty_buffer() {
let out = normalize_for_playback(Vec::new());
assert!(out.is_empty());
}
/// A buffer that would otherwise clip after gain is hard-clamped to ±1.0.
/// Belt-and-suspenders: `gain * peak` shouldn't exceed TARGET_PEAK by
/// construction, but the `.clamp(-1.0, 1.0)` is still asserted here in case
/// future changes loosen TARGET_PEAK or MAX_GAIN past full scale.
#[test]
fn normalize_for_playback_clamps_to_full_scale() {
let input = vec![10.0_f32, -10.0]; // already past full scale before gain
let out = normalize_for_playback(input);
assert!(out.iter().all(|&s| s.abs() <= 1.0));
}
}

View File

@@ -0,0 +1,292 @@
//! Huddle audio wire protocol — versioning + frame header.
//!
//! ## v1 (legacy)
//!
//! Client → relay: `<opus_bytes>`
//! Relay → client: `<peer_index: u8><opus_bytes>`
//!
//! No per-frame metadata; receiver synthesizes sequence/timestamp on arrival.
//! Kept for backward compatibility — relay still admits v1 clients into
//! v1-pinned rooms — but new clients always speak v2.
//!
//! ## v2 (this commit)
//!
//! Client → relay: `<header: [u8; 8]><opus_bytes>`
//! Relay → client: `<peer_index: u8><header: [u8; 8]><opus_bytes>`
//!
//! Header layout (8 bytes, network byte order, big-endian):
//!
//! ```text
//! byte 0..=1 : seq u16 wrapping sequence number, +1 per packet
//! byte 2..=5 : ts_48k u32 sender RTP-style 48 kHz media time,
//! +960 per 20 ms encoded frame
//! byte 6 : level_dbov i8 audio level in dBov, range [-127, 0]
//! byte 7 : flags u8 bit 0 = DTX/comfort frame; other bits
//! MUST be ignored on decode
//! ```
//!
//! Notes for reviewers (Max's spec):
//! * Header is fixed-size. No extension mechanism in v2 — future fields go
//! into a v3 with a new pinned version.
//! * `level_dbov` is client-authored telemetry. The relay parses it for
//! logging/active-speaker hints, clamps invalid values into range, and
//! **never** uses it for trust decisions (admission, moderation, etc.).
//! * Negotiation lives in the WS auth message (`protocol_version: 2`), not
//! in any bit of `flags`. Mixed-version rooms are rejected at the relay
//! with `upgrade_required`.
/// Wire protocol version this client speaks. Bumped only when the frame
/// layout itself changes; the relay tracks pinned per-room.
pub const PROTOCOL_VERSION: u8 = 2;
/// Length of the v2 per-frame header in bytes.
pub const V2_HEADER_LEN: usize = 8;
/// Flag bit indicating a DTX / comfort-noise frame. Receivers MAY skip
/// arrival accounting for these; encoders MUST set it for any Opus packet
/// they tag as DTX. Unset on normal speech frames.
pub const FLAG_DTX: u8 = 0x01;
/// Reserved bit-mask. Any bits outside this set MUST be ignored by the
/// receiver — they're available for the next minor protocol bump within v2.
#[allow(dead_code)] // Public constant for diagnostics / future flag additions.
pub const RESERVED_FLAG_MASK: u8 = !FLAG_DTX;
/// Parsed view of one v2 header. Cheap (Copy).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FrameHeader {
/// Sender-authored sequence number; wraps every 2^16 frames (~21.8 min
/// at 50 Hz).
pub seq: u16,
/// Sender-authored 48 kHz RTP-style media timestamp.
pub ts_48k: u32,
/// Audio level in dBov. Always parsed into the canonical `-127..=0`
/// range — out-of-range inputs are clamped to `-127` (silence).
pub level_dbov: i8,
/// Raw flags byte. Use [`Self::is_dtx`] / explicit masks; bits outside
/// `FLAG_DTX | reserved` are not inspected.
pub flags: u8,
}
impl FrameHeader {
/// Encode `self` into 8 bytes, network byte order. Suitable for direct
/// concatenation onto an Opus payload.
pub fn encode(self) -> [u8; V2_HEADER_LEN] {
let mut out = [0u8; V2_HEADER_LEN];
out[0..2].copy_from_slice(&self.seq.to_be_bytes());
out[2..6].copy_from_slice(&self.ts_48k.to_be_bytes());
out[6] = self.level_dbov as u8;
out[7] = self.flags;
out
}
/// Parse a v2 header from the leading 8 bytes of `bytes`. Returns
/// `(header, remainder)` so callers can pass the remainder straight to
/// the Opus decoder.
///
/// `level_dbov` is clamped into `-127..=0`; values outside that range
/// are coerced to `-127`. Per Max's spec, malformed VU metadata must
/// *not* cause the audio frame to be dropped — only the metric is
/// suppressed.
pub fn parse(bytes: &[u8]) -> Option<(Self, &[u8])> {
if bytes.len() < V2_HEADER_LEN {
return None;
}
let seq = u16::from_be_bytes([bytes[0], bytes[1]]);
let ts_48k = u32::from_be_bytes([bytes[2], bytes[3], bytes[4], bytes[5]]);
let raw_level = bytes[6] as i8;
let level_dbov = if (-127..=0).contains(&raw_level) {
raw_level
} else {
// Out-of-range telemetry is coerced to "silent floor"; the
// audio packet is still delivered. This matches Max's "don't
// drop audio on bad VU" guidance.
-127
};
let flags = bytes[7];
Some((
Self {
seq,
ts_48k,
level_dbov,
flags,
},
&bytes[V2_HEADER_LEN..],
))
}
/// True if [`FLAG_DTX`] is set.
#[allow(dead_code)]
pub fn is_dtx(&self) -> bool {
(self.flags & FLAG_DTX) != 0
}
}
/// Compute a dBov audio level for a normalized f32 PCM frame.
///
/// "dBov" is RMS expressed in dB relative to full scale (where full scale =
/// peak amplitude 1.0), clamped into the v2 header's canonical `-127..=0`
/// range. A silent frame returns `-127`; full-scale ±1.0 returns `0`.
///
/// This is cheap (one pass over the samples, one log) and runs once per
/// 20 ms frame, so it's nowhere near the audio fast-path budget.
pub fn audio_level_dbov(samples: &[f32]) -> i8 {
if samples.is_empty() {
return -127;
}
let mean_square: f64 = samples
.iter()
.map(|&s| {
let v = s as f64;
v * v
})
.sum::<f64>()
/ samples.len() as f64;
if mean_square <= 0.0 {
return -127;
}
let rms = mean_square.sqrt();
let db = 20.0 * rms.log10();
// Clamp into the canonical range. -127 = silence floor, 0 = full scale.
if !db.is_finite() || db <= -127.0 {
-127
} else if db >= 0.0 {
0
} else {
db.round() as i8
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn round_trip_preserves_fields() {
let h = FrameHeader {
seq: 0xABCD,
ts_48k: 0x12_34_56_78,
level_dbov: -40,
flags: FLAG_DTX,
};
let bytes = h.encode();
let (parsed, tail) = FrameHeader::parse(&bytes).expect("parse");
assert_eq!(parsed, h);
assert!(tail.is_empty());
assert!(parsed.is_dtx());
}
/// Header is fixed 8 bytes; anything shorter is rejected.
#[test]
fn parse_rejects_short_input() {
for len in 0..V2_HEADER_LEN {
let buf = vec![0u8; len];
assert!(
FrameHeader::parse(&buf).is_none(),
"{len} bytes must be rejected",
);
}
}
/// Trailing bytes pass through as the remainder, which is what the
/// receive path hands to the Opus decoder.
#[test]
fn parse_returns_payload_remainder() {
let mut buf = FrameHeader {
seq: 1,
ts_48k: 960,
level_dbov: -20,
flags: 0,
}
.encode()
.to_vec();
buf.extend_from_slice(b"opus-bytes");
let (_h, tail) = FrameHeader::parse(&buf).expect("parse");
assert_eq!(tail, b"opus-bytes");
}
/// Bytes in big-endian network order, matching Max's spec. This pins
/// the byte layout against accidental endianness changes.
#[test]
fn encoded_byte_order_is_network_byte_order() {
let h = FrameHeader {
seq: 0x0102,
ts_48k: 0x0304_0506,
level_dbov: -1,
flags: 0,
};
let bytes = h.encode();
assert_eq!(bytes[0..2], [0x01, 0x02]);
assert_eq!(bytes[2..6], [0x03, 0x04, 0x05, 0x06]);
assert_eq!(bytes[6], 0xFF); // -1 as i8 -> 0xFF
assert_eq!(bytes[7], 0x00);
}
/// Per Max: malformed `level_dbov` must not drop the frame; only the
/// metric is suppressed. We coerce to -127 (silence floor) and keep
/// the payload available.
#[test]
fn out_of_range_level_dbov_is_clamped_not_rejected() {
// i8(127) is +127 which is outside [-127, 0].
let mut bytes = FrameHeader {
seq: 7,
ts_48k: 7 * 960,
level_dbov: 0,
flags: 0,
}
.encode();
bytes[6] = 0x7F; // raw +127 — out of range
let (parsed, _) = FrameHeader::parse(&bytes).expect("parse must succeed");
assert_eq!(parsed.level_dbov, -127, "clamped to silence floor");
assert_eq!(parsed.seq, 7, "other fields preserved");
}
#[test]
fn audio_level_dbov_clamps_silence_to_minus_127() {
assert_eq!(audio_level_dbov(&[]), -127);
assert_eq!(audio_level_dbov(&[0.0_f32; 960]), -127);
}
#[test]
fn audio_level_dbov_full_scale_returns_zero() {
let frame: Vec<f32> = (0..960).map(|_| 1.0_f32).collect();
assert_eq!(audio_level_dbov(&frame), 0);
}
#[test]
fn audio_level_dbov_is_in_canonical_range_for_realistic_speech() {
// A 1 kHz sine at amplitude 0.3 — typical normalized speech peak.
// RMS ≈ 0.3 / sqrt(2) ≈ 0.212 → ≈ -13 dBov.
let frame: Vec<f32> = (0..960)
.map(|i| {
let t = i as f32 / 48_000.0;
0.3 * (2.0 * std::f32::consts::PI * 1_000.0 * t).sin()
})
.collect();
let db = audio_level_dbov(&frame);
assert!(
(-20..=-8).contains(&db),
"expected -20..=-8 dBov for 0.3-amp sine, got {db}",
);
}
/// Reserved flag bits MUST be ignored on parse (no error, just preserved
/// for forward-compat). The contract is that *receivers* don't inspect
/// them; the parse path keeps them so introspection still works.
#[test]
fn reserved_flag_bits_are_preserved_for_inspection() {
let mut bytes = FrameHeader {
seq: 0,
ts_48k: 0,
level_dbov: 0,
flags: 0,
}
.encode();
bytes[7] = 0b1010_1010; // FLAG_DTX bit clear, reserved bits set
let (parsed, _) = FrameHeader::parse(&bytes).expect("parse");
assert_eq!(parsed.flags, 0b1010_1010);
assert!(!parsed.is_dtx());
// Caller can still check which reserved bits were set if needed.
assert_eq!(parsed.flags & RESERVED_FLAG_MASK, 0b1010_1010);
}
}

View File

@@ -395,10 +395,10 @@ pub fn run() {
// Pre-download voice models in the background so they're ready
// when the user starts their first huddle. Idempotent — no-op if
// already downloaded. ~187 MB total (~100 MB Parakeet STT + ~87 MB Kokoro).
// already downloaded. ~289 MB total (~100 MB Parakeet STT + ~189 MB Pocket TTS).
if let Some(mgr) = huddle::models::global_model_manager() {
mgr.start_stt_download(state.http_client.clone());
mgr.start_kokoro_download(state.http_client.clone());
mgr.start_tts_download(state.http_client.clone());
}
// Register PTT global shortcut (Ctrl+Space).

View File

@@ -0,0 +1,142 @@
//! Diagnostic regression test: rodio's `Player` is a single-queue FIFO, not a summing mixer.
//!
//! This pins the failure mode we hit at 3+ simultaneous speakers in `huddle`:
//! `Player::connect_new(&device_mixer)` calls `device_mixer.add(queue_source)` exactly once,
//! and every `player.append(SamplesBuffer)` enqueues onto that single queue. Multiple peers'
//! 20 ms frames therefore serialize — peer B plays after peer A finishes, etc. — instead of
//! mixing.
//!
//! The fix is per-peer `Player` (each peer gets its own queue added to the device mixer), so
//! the device mixer actually sums across queues. This test asserts both halves of the
//! diagnosis with a deterministic, device-free harness:
//!
//! * A single queue receiving two sources → samples arrive **serially**.
//! * A summing mixer fed by two sources → samples arrive **concurrently and summed**.
//!
//! If rodio ever changes either invariant under us we want CI to scream, not the next
//! 3-person huddle.
use rodio::buffer::SamplesBuffer;
use rodio::mixer;
use rodio::queue;
use rodio::source::Source;
use std::num::NonZero;
const SR: u32 = 48_000;
const CH: u16 = 1;
/// 20 ms at 48 kHz mono — one Opus frame's worth of decoded samples.
const FRAME_SAMPLES: usize = 960;
fn channels() -> NonZero<u16> {
NonZero::new(CH).unwrap()
}
fn sample_rate() -> NonZero<u32> {
NonZero::new(SR).unwrap()
}
/// Build a `SamplesBuffer` of length `n` samples filled with `value`.
fn buf(value: f32, n: usize) -> SamplesBuffer {
SamplesBuffer::new(channels(), sample_rate(), vec![value; n])
}
/// Drain at most `limit` samples from any `Source` and return them.
fn drain<S: Source<Item = f32>>(src: S, limit: usize) -> Vec<f32> {
src.take(limit).collect()
}
/// A single `queue` plays sources **one after the other** (FIFO).
///
/// This is the shape `rodio::Player::connect_new` produces: one queue is added to the
/// device mixer at start, and every `Player::append` enqueues onto that one queue.
/// In the huddle path, every peer's decoded frame was appended here — serializing them.
#[test]
fn single_queue_serializes_sources() {
let (input, output) = queue::queue(false /* don't keep alive when empty */);
// Two "peers": A produces all 1.0s, B produces all -1.0s. If they were summed, the
// overlap would be ~0.0; if they're serialized we see a clean 1.0 → -1.0 transition.
input.append(buf(1.0, FRAME_SAMPLES));
input.append(buf(-1.0, FRAME_SAMPLES));
let samples = drain(output, FRAME_SAMPLES * 4);
// Total samples produced = sum of inputs (serial). If it were summed we'd only get
// FRAME_SAMPLES of output (and at amplitude 0.0).
assert_eq!(
samples.len(),
FRAME_SAMPLES * 2,
"single queue should drain both sources back-to-back (serial), not mix them",
);
// The first source's samples come first, then the second's. No interleaving, no summing.
assert!(
samples[..FRAME_SAMPLES].iter().all(|&s| s == 1.0),
"first half of queue output should be source A's samples (1.0), not a mix",
);
assert!(
samples[FRAME_SAMPLES..].iter().all(|&s| s == -1.0),
"second half of queue output should be source B's samples (-1.0)",
);
}
/// A `mixer` plays sources **concurrently** and **sums** overlapping samples.
///
/// This is the shape we want for huddle playout: each peer gets its own source added to
/// the device mixer, so simultaneous speakers actually mix instead of serializing.
#[test]
fn mixer_sums_overlapping_sources() {
let (controller, mixer_source) = mixer::mixer(channels(), sample_rate());
// Same two "peers" — 1.0 and -1.0. If the mixer sums correctly, overlap == 0.0.
controller.add(buf(1.0, FRAME_SAMPLES));
controller.add(buf(-1.0, FRAME_SAMPLES));
// Take exactly one frame's worth — both sources should be active across that window.
let samples = drain(mixer_source, FRAME_SAMPLES);
assert_eq!(
samples.len(),
FRAME_SAMPLES,
"mixer should produce one frame's worth of mixed output, not two frames serialized",
);
// Both sources are active over this window. Their sum is 0.0 at every sample.
// Allow for f32 rounding (rodio's UniformSourceIterator + sum can introduce a few ULPs).
let max_abs = samples.iter().map(|s| s.abs()).fold(0.0_f32, f32::max);
assert!(
max_abs < 1e-5,
"mixer should sum 1.0 + -1.0 ≈ 0.0 per sample; saw max |sample| = {max_abs} \
which means the sources weren't actually mixed concurrently",
);
}
/// One mixer fed by **two distinct queues** (= two distinct `Player`s on one device sink)
/// behaves the same way: the queues drain concurrently, the mixer sums them.
///
/// This is the exact shape the huddle fix moves to: `HashMap<peer_index, Player>`, each
/// added to `MixerDeviceSink::mixer()`. The test pins that this composes correctly.
#[test]
fn mixer_of_queues_mixes_per_peer_streams() {
let (mixer_in, mixer_source) = mixer::mixer(channels(), sample_rate());
// Two queues — one per "peer".
let (peer_a_in, peer_a_out) =
queue::queue(true /* keep alive — players outlive frames */);
let (peer_b_in, peer_b_out) = queue::queue(true);
mixer_in.add(peer_a_out);
mixer_in.add(peer_b_out);
// Each peer pushes one 20 ms frame at the same wall-clock moment.
peer_a_in.append(buf(1.0, FRAME_SAMPLES));
peer_b_in.append(buf(-1.0, FRAME_SAMPLES));
let samples = drain(mixer_source, FRAME_SAMPLES);
assert_eq!(samples.len(), FRAME_SAMPLES);
let max_abs = samples.iter().map(|s| s.abs()).fold(0.0_f32, f32::max);
assert!(
max_abs < 1e-5,
"per-peer queues into one mixer should sum concurrently; saw max |sample| = {max_abs}",
);
}

View File

@@ -70,6 +70,7 @@ export function AppShellOverlays({
<SearchDialog
channels={channels}
currentPubkey={currentPubkey}
onOpenChannel={onSelectChannel}
onOpenResult={onOpenSearchResult}
onOpenChange={onSearchOpenChange}
open={true}

View File

@@ -0,0 +1,72 @@
import { Headphones } from "lucide-react";
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/shared/ui/alert-dialog";
/**
* Pre-join "use headphones" confirmation dialog.
*
* Shown once per browser-tab session when the user clicks the start/join
* huddle button while echo cancellation is missing. The mic isn't opened
* and `start_huddle` / `join_huddle` is not called until the user clicks
* "Continue". After a confirmed continue (or an explicit cancel) the gate
* remembers its decision in `sessionStorage` so the same user isn't
* nagged on every channel hop.
*
* Pair with [`useHeadphonesGate`]. The hook owns the storage + decision
* state; this component is presentational.
*
* Removable in one diff alongside `HeadphonesNotice` when the WebAudio
* AEC follow-up flips `aecMissing` to false in `HuddleBar`.
*/
export function HeadphonesGate({
open,
onContinue,
onCancel,
}: {
open: boolean;
onContinue: () => void;
onCancel: () => void;
}) {
return (
<AlertDialog
open={open}
onOpenChange={(next) => {
// Radix calls onOpenChange(false) when the user clicks outside or
// presses Escape. Treat that as cancel — never as a silent join.
if (!next) onCancel();
}}
>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle className="flex items-center gap-2">
<Headphones className="h-4 w-4" />
Headphones recommended
</AlertDialogTitle>
<AlertDialogDescription>
Echo cancellation lands in the next release. Until then, two or more
people on speakers in the same room will hear themselves echo. Use
headphones for the best experience.
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel onClick={onCancel}>Cancel</AlertDialogCancel>
<AlertDialogAction
data-testid="huddle-headphones-gate-continue"
onClick={onContinue}
>
Continue
</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
);
}

View File

@@ -0,0 +1,41 @@
import { Headphones } from "lucide-react";
/**
* "Headphones recommended — echo cancellation lands in the next release."
*
* Session-dismissable banner shown on the active HuddleBar while the desktop
* client lacks an echo-cancellation reference. The current play path is
* native rodio (outside the WebView render graph), so the browser's WebRTC
* AEC cannot suppress local-speaker → mic feedback for users on speakers.
*
* The follow-up PR moves remote-peer playout into WebAudio inside the same
* `AudioContext` as `getUserMedia({ echoCancellation: true })`. When that
* lands, the parent component flips `aecMissing` to false and this banner
* (and this file) become removable in a single diff.
*
* State (`dismissed`, `aecMissing`) lives in the parent so the parent can
* decide whether to render at all. Keep this component dumb so the
* deletion later is mechanical.
*/
export function HeadphonesNotice({ onDismiss }: { onDismiss: () => void }) {
return (
<div
role="status"
data-testid="huddle-headphones-notice"
className="flex items-center gap-1.5 rounded bg-amber-500/10 px-2 py-1 text-xs text-amber-700 dark:text-amber-300"
>
<Headphones className="h-3 w-3" />
<span className="max-w-[260px] truncate">
Headphones recommended — echo cancellation lands in the next release.
</span>
<button
aria-label="Dismiss headphones notice for this session"
className="ml-1 opacity-60 hover:opacity-100"
onClick={onDismiss}
type="button"
>
✕
</button>
</div>
);
}

View File

@@ -8,6 +8,7 @@ import { Button } from "@/shared/ui/button";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/shared/ui/tooltip";
import { useHuddle } from "../HuddleContext";
import { AddAgentDialog, type AgentAddResult } from "./AddAgentDialog";
import { HeadphonesNotice } from "./HeadphonesNotice";
import { MicControls, SpeakerControls } from "./MicControls";
import { ParticipantList } from "./ParticipantList";
@@ -60,13 +61,23 @@ export function HuddleBar({ className }: HuddleBarProps) {
const isPttMode = voiceInputMode === "push_to_talk";
const [state, setState] = React.useState<HuddleState | null>(null);
const [isMuted, setIsMuted] = React.useState(false);
// Session-dismissable "use headphones" notice. The desktop client plays
// remote peer audio via native rodio, outside the WebView render graph,
// so the browser's echo canceller has no far-end reference to cancel
// against. Until the WebAudio AEC follow-up PR lands, two people on
// speakers in the same physical room will hear themselves echo. The
// banner stays visible across huddle start/stop within a single page
// load; it auto-removes when the AEC plumbing lands and the detection
// below flips to false.
const [headphonesNoticeDismissed, setHeadphonesNoticeDismissed] =
React.useState(false);
const ttsEnabled = state?.tts_enabled ?? true;
const [isLeaving, setIsLeaving] = React.useState(false);
const [showAddAgent, setShowAddAgent] = React.useState(false);
const [agentAddError, setAgentAddError] = React.useState<string | null>(null);
const [modelStatus, setModelStatus] = React.useState<{
stt: string;
kokoro: string;
tts: string;
} | null>(null);
// Huddle state: event-driven + 10s fallback poll.
React.useEffect(() => {
@@ -133,13 +144,13 @@ export function HuddleBar({ className }: HuddleBarProps) {
try {
const status = await invoke<{
stt: unknown;
kokoro: unknown;
tts: unknown;
}>("get_model_status");
if (cancelled) return;
setModelStatus({
stt: fmt(status.stt),
kokoro: fmt(status.kokoro),
tts: fmt(status.tts),
});
} catch {
// best-effort
@@ -165,6 +176,12 @@ export function HuddleBar({ className }: HuddleBarProps) {
if (!state || (state.phase !== "active" && state.phase !== "connected"))
return null;
// Self-removing detection: remote-peer audio plays through native rodio
// today (outside the WebView render graph), so the browser's AEC has no
// far-end reference. The AEC follow-up PR flips this constant in the
// same diff that moves playout into WebAudio.
const aecMissing = true;
async function handleLeave() {
if (isLeaving) return;
setIsLeaving(true);
@@ -226,16 +243,23 @@ export function HuddleBar({ className }: HuddleBarProps) {
</div>
)}
{/* Echo-cancellation pre-PR notice. Removed when AEC plumbing lands. */}
{aecMissing && !headphonesNoticeDismissed && (
<HeadphonesNotice
onDismiss={() => setHeadphonesNoticeDismissed(true)}
/>
)}
{/* Model download progress */}
{modelStatus &&
(modelStatus.stt !== "ready" || modelStatus.kokoro !== "ready") && (
(modelStatus.stt !== "ready" || modelStatus.tts !== "ready") && (
<output className="flex items-center gap-1 text-xs text-muted-foreground">
<span className="animate-pulse">
{modelStatus.stt !== "ready" && modelStatus.kokoro !== "ready"
? `Voice models: STT ${modelStatus.stt}, TTS ${modelStatus.kokoro}`
{modelStatus.stt !== "ready" && modelStatus.tts !== "ready"
? `Voice models: STT ${modelStatus.stt}, TTS ${modelStatus.tts}`
: modelStatus.stt !== "ready"
? `STT model: ${modelStatus.stt}`
: `TTS model: ${modelStatus.kokoro}`}
: `TTS model: ${modelStatus.tts}`}
</span>
</output>
)}
@@ -464,8 +488,8 @@ export function HuddleBar({ className }: HuddleBarProps) {
modelStatus.stt !== "ready" &&
`, STT model ${modelStatus.stt}`}
{modelStatus &&
modelStatus.kokoro !== "ready" &&
`, TTS model ${modelStatus.kokoro}`}
modelStatus.tts !== "ready" &&
`, TTS model ${modelStatus.tts}`}
</output>
</div>
);

View File

@@ -9,6 +9,8 @@ import { cn } from "@/shared/lib/cn";
import { Button } from "@/shared/ui/button";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/shared/ui/tooltip";
import { useHuddle } from "../HuddleContext";
import { useHeadphonesGate } from "../lib/useHeadphonesGate";
import { HeadphonesGate } from "./HeadphonesGate";
/** Huddle lifecycle event kinds */
const KIND_HUDDLE_STARTED = 48100;
@@ -48,6 +50,13 @@ export function HuddleIndicator({
);
const [isJoining, setIsJoining] = React.useState(false);
// Pre-join "use headphones" confirmation while echo cancellation is
// missing. Mirrors the `aecMissing` constant in HuddleBar — both flip
// to false in the AEC follow-up PR, after which this hook is a no-op
// and the gate UI/components are mechanically deletable.
const aecMissing = true;
const headphonesGate = useHeadphonesGate(aecMissing);
React.useEffect(() => {
if (!channelId) return;
@@ -203,22 +212,36 @@ export function HuddleIndicator({
};
}, []);
// Pre-join headphones confirmation. The same dialog instance serves
// both the start and the join paths — `gate` decides which deferred
// action to fire on Continue.
const gateDialog = (
<HeadphonesGate
open={headphonesGate.dialogOpen}
onContinue={headphonesGate.onContinue}
onCancel={headphonesGate.onCancel}
/>
);
// No active huddle — render the start button (if onStart provided).
if (!activeHuddle) {
if (!onStart) return null;
return (
<Button
aria-label="Start huddle"
className={cn("h-7 w-7 rounded-full", className)}
data-testid="channel-start-huddle-trigger"
disabled={startDisabled || isStarting}
onClick={() => onStart()}
size="icon"
type="button"
variant="outline"
>
<Headphones className="h-3 w-3" />
</Button>
<>
<Button
aria-label="Start huddle"
className={cn("h-7 w-7 rounded-full", className)}
data-testid="channel-start-huddle-trigger"
disabled={startDisabled || isStarting}
onClick={() => headphonesGate.gate(() => onStart())}
size="icon"
type="button"
variant="outline"
>
<Headphones className="h-3 w-3" />
</Button>
{gateDialog}
</>
);
}
@@ -227,7 +250,7 @@ export function HuddleIndicator({
// reconstructed set — floor at 1 to avoid showing "0 participants".
const participantCount = Math.max(1, activeHuddle.participants.size);
async function handleJoin() {
async function doJoin() {
if (!activeHuddle || isJoining) return;
setIsJoining(true);
try {
@@ -242,30 +265,33 @@ export function HuddleIndicator({
}
return (
<Tooltip>
<TooltipTrigger asChild>
<Button
aria-label={`Join active huddle (${participantCount} participant${participantCount !== 1 ? "s" : ""})`}
className={cn("relative h-7 w-7 rounded-full", className)}
disabled={isJoining || isStarting}
onClick={() => void handleJoin()}
size="icon"
type="button"
variant="outline"
>
<Headphones className="h-3 w-3 text-muted-foreground" />
<span className="absolute inset-0 animate-pulse rounded-full ring-2 ring-border/70" />
{/* Participant count badge */}
{participantCount > 0 && (
<span className="absolute -right-1 -top-1 flex h-3.5 min-w-3.5 items-center justify-center rounded-full border border-border bg-background px-0.5 text-[9px] font-bold text-muted-foreground">
{participantCount}
</span>
)}
</Button>
</TooltipTrigger>
<TooltipContent>
{`Huddle active — ${participantCount} participant${participantCount !== 1 ? "s" : ""}`}
</TooltipContent>
</Tooltip>
<>
<Tooltip>
<TooltipTrigger asChild>
<Button
aria-label={`Join active huddle (${participantCount} participant${participantCount !== 1 ? "s" : ""})`}
className={cn("relative h-7 w-7 rounded-full", className)}
disabled={isJoining || isStarting}
onClick={() => headphonesGate.gate(() => void doJoin())}
size="icon"
type="button"
variant="outline"
>
<Headphones className="h-3 w-3 text-muted-foreground" />
<span className="absolute inset-0 animate-pulse rounded-full ring-2 ring-border/70" />
{/* Participant count badge */}
{participantCount > 0 && (
<span className="absolute -right-1 -top-1 flex h-3.5 min-w-3.5 items-center justify-center rounded-full border border-border bg-background px-0.5 text-[9px] font-bold text-muted-foreground">
{participantCount}
</span>
)}
</Button>
</TooltipTrigger>
<TooltipContent>
{`Huddle active — ${participantCount} participant${participantCount !== 1 ? "s" : ""}`}
</TooltipContent>
</Tooltip>
{gateDialog}
</>
);
}

View File

@@ -0,0 +1,78 @@
import * as React from "react";
/**
* Owns the "should we show the pre-join headphones warning?" decision.
*
* Wraps an arbitrary click handler — `gate(action)` returns a function that
* either runs `action` directly (already confirmed this session) or shows
* the confirmation dialog and runs `action` only after the user clicks
* Continue. Cancel/Escape/click-outside all decline silently.
*
* The "confirmed this session" bit is stored in `sessionStorage`, so it
* resets when the desktop window is closed but persists across channel
* navigation in the same window. That matches Quinn's recipe in the
* design thread: "session-dismissable, self-removing once AEC lands."
*
* `aecMissing` is a prop, not a feature-detect inside the hook, because
* the AEC follow-up PR flips it from a parent constant; centralizing the
* predicate makes the future deletion mechanical.
*
* Returns an object with `dialogOpen`, `gate(action)`, `onContinue`, and
* `onCancel`. Render the `HeadphonesGate` component with the dialog props.
*/
const STORAGE_KEY = "huddle.headphones-gate-confirmed";
function readConfirmed(): boolean {
try {
return sessionStorage.getItem(STORAGE_KEY) === "1";
} catch {
// sessionStorage can throw in some embedded WebView contexts; treat
// failure as "not confirmed" so the gate still warns conservatively.
return false;
}
}
function writeConfirmed() {
try {
sessionStorage.setItem(STORAGE_KEY, "1");
} catch {
/* best-effort; if storage is denied the user just sees the gate next time */
}
}
export function useHeadphonesGate(aecMissing: boolean) {
const [dialogOpen, setDialogOpen] = React.useState(false);
const pendingActionRef = React.useRef<(() => void) | null>(null);
const gate = React.useCallback(
(action: () => void) => {
// Once AEC works (predicate flips to false), this hook is a no-op:
// run the action immediately. The component using this hook can stay
// in place; deleting the hook + gate + notice happens in one diff
// when the follow-up PR makes aecMissing always false.
if (!aecMissing || readConfirmed()) {
action();
return;
}
pendingActionRef.current = action;
setDialogOpen(true);
},
[aecMissing],
);
const onContinue = React.useCallback(() => {
writeConfirmed();
const action = pendingActionRef.current;
pendingActionRef.current = null;
setDialogOpen(false);
action?.();
}, []);
const onCancel = React.useCallback(() => {
pendingActionRef.current = null;
setDialogOpen(false);
}, []);
return { dialogOpen, gate, onContinue, onCancel };
}

View File

@@ -1,7 +1,10 @@
import assert from "node:assert/strict";
import test from "node:test";
import { buildMainTimelineEntries } from "./threadPanel.ts";
import {
buildMainTimelineEntries,
buildThreadPanelData,
} from "./threadPanel.ts";
function message(overrides) {
return {
@@ -56,3 +59,44 @@ test("buildMainTimelineEntries includes broadcast replies", () => {
["root", "broadcast-reply"],
);
});
test("buildThreadPanelData keeps direct comments unindented", () => {
const root = message({ id: "root", createdAt: 1 });
const directComment = message({
id: "direct-comment",
createdAt: 2,
parentId: "root",
rootId: "root",
depth: 1,
tags: [["e", "root", "", "reply"]],
});
const nestedReply = message({
id: "nested-reply",
createdAt: 3,
parentId: "direct-comment",
rootId: "root",
depth: 2,
tags: [
["e", "root", "", "root"],
["e", "direct-comment", "", "reply"],
],
});
const panelData = buildThreadPanelData(
[root, directComment, nestedReply],
"root",
"root",
new Set(["direct-comment"]),
);
assert.deepEqual(
panelData.visibleReplies.map((entry) => ({
id: entry.message.id,
depth: entry.message.depth,
})),
[
{ id: "direct-comment", depth: 0 },
{ id: "nested-reply", depth: 1 },
],
);
});

View File

@@ -16,6 +16,7 @@ export type TimelineThreadSummaryParticipant = {
export type TimelineThreadSummary = {
threadHeadId: string;
replyCount: number;
lastReplyAt: number | null;
participants: TimelineThreadSummaryParticipant[];
};
@@ -26,6 +27,7 @@ export type MainTimelineEntry = {
type ThreadDescendantStats = {
descendantCount: number;
lastReplyAt: number | null;
recentParticipantsNewestFirst: TimelineThreadSummaryParticipant[];
};
@@ -80,6 +82,7 @@ function buildDescendantStatsByMessageId(
message.id,
{
descendantCount: 0,
lastReplyAt: null,
recentParticipantsNewestFirst: [],
},
]),
@@ -115,6 +118,10 @@ function buildDescendantStatsByMessageId(
}
ancestorStats.descendantCount += 1;
ancestorStats.lastReplyAt = Math.max(
ancestorStats.lastReplyAt ?? 0,
message.createdAt,
);
if (
ancestorStats.recentParticipantsNewestFirst.length <
@@ -146,6 +153,7 @@ function buildSummaryForDirectReplies(
return {
threadHeadId: messageId,
replyCount: descendantStats.descendantCount,
lastReplyAt: descendantStats.lastReplyAt,
participants: [...descendantStats.recentParticipantsNewestFirst].reverse(),
};
}
@@ -207,7 +215,7 @@ function buildVisibleThreadReplies(params: {
appendExpandedReplies({
entries,
parentId: openThreadHeadId,
depth: 1,
depth: 0,
directChildrenByParentId,
descendantStatsByMessageId,
expandedReplyIds,

View File

@@ -10,7 +10,10 @@ import { cn } from "@/shared/lib/cn";
import { UserAvatar } from "@/shared/ui/UserAvatar";
import { useChannelNavigation } from "@/shared/context/ChannelNavigationContext";
import { parseImetaTags } from "@/features/messages/lib/parseImeta";
import { resolveMentionNames } from "@/shared/lib/resolveMentionNames";
import {
resolveMentionNames,
resolveMentionPubkeysByName,
} from "@/shared/lib/resolveMentionNames";
import { Markdown } from "@/shared/ui/markdown";
import { MessageActionBar } from "./MessageActionBar";
import { MessageTimestamp } from "./MessageTimestamp";
@@ -19,6 +22,9 @@ import { Tooltip, TooltipContent, TooltipTrigger } from "@/shared/ui/tooltip";
const DiffMessage = React.lazy(() => import("./DiffMessage"));
const DiffMessageExpanded = React.lazy(() => import("./DiffMessageExpanded"));
const MESSAGE_TEXT_OFFSET_PX = 54;
const NESTED_REPLY_OFFSET_PX = 28;
export const MessageRow = React.memo(
function MessageRow({
activeReplyTargetId = null,
@@ -64,6 +70,10 @@ export const MessageRow = React.memo(
() => resolveMentionNames(message.tags, profiles),
[profiles, message.tags],
);
const mentionPubkeysByName = React.useMemo(
() => resolveMentionPubkeysByName(message.tags, profiles),
[profiles, message.tags],
);
const imetaByUrl = React.useMemo(
() => (message.tags ? parseImetaTags(message.tags) : undefined),
@@ -77,11 +87,21 @@ export const MessageRow = React.memo(
);
const visibleDepth = Math.min(message.depth, 6);
const indentPx = visibleDepth * 28;
const indentPx =
visibleDepth > 0
? MESSAGE_TEXT_OFFSET_PX + (visibleDepth - 1) * NESTED_REPLY_OFFSET_PX
: 0;
const depthGuideOffsets = React.useMemo(() => {
return Array.from(
{ length: visibleDepth },
(_, index) => 14 + index * 28,
if (visibleDepth === 0) {
return [];
}
return Array.from({ length: visibleDepth }, (_, index) =>
index === 0
? MESSAGE_TEXT_OFFSET_PX / 2
: MESSAGE_TEXT_OFFSET_PX +
NESTED_REPLY_OFFSET_PX / 2 +
(index - 1) * NESTED_REPLY_OFFSET_PX,
);
}, [visibleDepth]);
const getTag = (name: string) =>
@@ -119,6 +139,7 @@ export const MessageRow = React.memo(
content={message.body}
imetaByUrl={imetaByUrl}
mentionNames={mentionNames}
mentionPubkeysByName={mentionPubkeysByName}
searchQuery={searchQuery}
tight
/>
@@ -128,12 +149,8 @@ export const MessageRow = React.memo(
const isThreadReplyLayout = layoutVariant === "thread-reply";
const guideBleedPx = isThreadReplyLayout ? 4 : 0;
const avatarSizeClass = isThreadReplyLayout
? "!h-5 !w-5 !rounded-md"
: "!h-9 !w-9";
const avatarButtonRadiusClass = isThreadReplyLayout
? "rounded-md"
: "rounded-xl";
const avatarSizeClass = "!h-9 !w-9";
const avatarButtonRadiusClass = "rounded-xl";
const respondToDotColor =
message.respondTo === "anyone"
@@ -283,7 +300,7 @@ export const MessageRow = React.memo(
<article
className={cn(
"group/message relative rounded-2xl px-2 py-1 transition-colors",
isThreadReplyLayout ? "space-y-1" : "flex items-start gap-2.5",
"flex items-start gap-2.5",
highlighted ? "bg-primary/10 ring-1 ring-primary/30" : "",
)}
data-message-id={message.id}
@@ -291,33 +308,43 @@ export const MessageRow = React.memo(
>
{isThreadReplyLayout ? (
<>
<div className="flex min-w-0 items-baseline gap-1.5">
{message.pubkey ? (
<UserProfilePopover
pubkey={message.pubkey}
role={message.role}
botIdenticonValue={message.author}
{message.pubkey ? (
<UserProfilePopover
pubkey={message.pubkey}
role={message.role}
botIdenticonValue={message.author}
>
<button
className={cn(
"flex shrink-0 items-start focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
avatarButtonRadiusClass,
)}
type="button"
>
<button
className={cn(
"flex shrink-0 items-start gap-1.5 focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
avatarButtonRadiusClass,
)}
type="button"
{avatarNode}
</button>
</UserProfilePopover>
) : (
<div className="flex shrink-0 items-start">{avatarNode}</div>
)}
<div className="-mt-1 min-w-0 flex-1 space-y-0">
<div className="flex min-w-0 flex-wrap items-baseline gap-x-2 gap-y-0">
{message.pubkey ? (
<UserProfilePopover
pubkey={message.pubkey}
role={message.role}
botIdenticonValue={message.author}
>
{avatarNode}
{authorNode}
</button>
</UserProfilePopover>
) : (
<>
<div className="flex shrink-0 items-start">
{avatarNode}
</div>
{authorNode}
</>
)}
<div className="flex min-w-0 flex-1 flex-wrap items-baseline gap-x-2 gap-y-0.5">
<button
className="truncate rounded focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
type="button"
>
{authorNode}
</button>
</UserProfilePopover>
) : (
authorNode
)}
{inlineMetadataNode}
{message.personaDisplayName &&
message.personaDisplayName !== message.author ? (
@@ -326,8 +353,8 @@ export const MessageRow = React.memo(
</span>
) : null}
</div>
<div className="-mt-0.5">{messageBodyNode}</div>
</div>
<div className="min-w-0 space-y-0.5">{messageBodyNode}</div>
</>
) : (
<>

View File

@@ -228,10 +228,13 @@ export function MessageThreadPanel({
data-testid="message-thread-replies"
>
{threadReplies.length > 0 ? (
<div className="space-y-2">
<div className="space-y-2.5">
{threadReplies.map((entry) => {
return (
<div key={entry.message.id}>
<div
className="flex flex-col gap-1"
key={entry.message.id}
>
<MessageRow
activeReplyTargetId={replyTargetId}
channelId={channelId}
@@ -257,7 +260,6 @@ export function MessageThreadPanel({
{entry.summary ? (
<MessageThreadSummaryRow
depth={entry.message.depth}
layoutVariant="thread-reply"
message={entry.message}
onOpenThread={onExpandReplies}
summary={entry.summary}

View File

@@ -5,6 +5,24 @@ import type {
import type { TimelineMessage } from "@/features/messages/types";
import { UserAvatar } from "@/shared/ui/UserAvatar";
const MESSAGE_TEXT_OFFSET_PX = 54;
const NESTED_REPLY_OFFSET_PX = 28;
function formatRelativeTime(unixSeconds: number): string {
const now = Date.now() / 1_000;
const diff = now - unixSeconds;
if (diff < 60) return "just now";
if (diff < 3_600) return `${Math.floor(diff / 60)}m`;
if (diff < 86_400) return `${Math.floor(diff / 3_600)}h`;
if (diff < 604_800) return `${Math.floor(diff / 86_400)}d`;
return new Date(unixSeconds * 1_000).toLocaleDateString(undefined, {
month: "short",
day: "numeric",
});
}
function ParticipantAvatar({
participant,
index,
@@ -30,24 +48,31 @@ function ParticipantAvatar({
export function MessageThreadSummaryRow({
depth = 0,
layoutVariant = "default",
message,
onOpenThread,
summary,
}: {
depth?: number;
layoutVariant?: "default" | "thread-reply";
message: TimelineMessage;
onOpenThread: (message: TimelineMessage) => void;
summary: TimelineThreadSummary;
}) {
const visibleDepth = Math.min(Math.max(depth, 0), 6);
const messageTextOffsetPx = layoutVariant === "thread-reply" ? 8 : 50;
const marginLeftPx = visibleDepth * 28 + messageTextOffsetPx;
const depthGuideOffsets = Array.from(
{ length: visibleDepth },
(_, index) => 14 + index * 28,
);
const indentPx =
visibleDepth > 0
? MESSAGE_TEXT_OFFSET_PX + (visibleDepth - 1) * NESTED_REPLY_OFFSET_PX
: 0;
const marginLeftPx = indentPx + MESSAGE_TEXT_OFFSET_PX;
const depthGuideOffsets =
visibleDepth === 0
? []
: Array.from({ length: visibleDepth }, (_, index) =>
index === 0
? MESSAGE_TEXT_OFFSET_PX / 2
: MESSAGE_TEXT_OFFSET_PX +
NESTED_REPLY_OFFSET_PX / 2 +
(index - 1) * NESTED_REPLY_OFFSET_PX,
);
return (
<div className="relative pb-1 pt-1">
@@ -71,7 +96,7 @@ export function MessageThreadSummaryRow({
) : null}
<button
className="inline-flex w-fit max-w-full items-center gap-1 rounded-full border border-border/70 bg-muted/70 py-0.5 pl-0.5 pr-2 text-left text-xs font-medium text-foreground/90 transition-colors hover:bg-accent hover:text-accent-foreground"
className="group inline-flex w-fit max-w-full items-center gap-1 text-left text-xs font-medium text-muted-foreground"
data-thread-head-id={message.id}
data-testid="message-thread-summary"
onClick={() => onOpenThread(message)}
@@ -89,10 +114,15 @@ export function MessageThreadSummaryRow({
</div>
<div className="min-w-0">
<div className="font-medium">
<span>
<span className="transition-colors group-hover:text-foreground">
{summary.replyCount}{" "}
{summary.replyCount === 1 ? "reply" : "replies"}
</span>
{summary.lastReplyAt ? (
<span className="ml-1 text-muted-foreground/70">
last {formatRelativeTime(summary.lastReplyAt)}
</span>
) : null}
</div>
</div>
</button>

View File

@@ -135,7 +135,7 @@ export const MessageTimeline = React.memo(function MessageTimeline({
ref={scrollContainerRef}
>
<div
className="flex w-full flex-col gap-2 pb-10 pt-16"
className="flex w-full flex-col gap-2 pb-10 pt-12"
ref={contentRef}
>
<div ref={topSentinelRef} aria-hidden className="h-px" />
@@ -152,7 +152,7 @@ export const MessageTimeline = React.memo(function MessageTimeline({
data-testid="message-timeline-beginning"
>
<Separator className="flex-1" />
<p className="text-xs font-semibold uppercase tracking-[0.22em] text-muted-foreground">
<p className="text-[11px] font-medium uppercase tracking-[0.08em] text-muted-foreground/75">
Beginning of conversation
</p>
<Separator className="flex-1" />

View File

@@ -8,6 +8,7 @@ import { MessageReactions } from "@/features/messages/ui/MessageReactions";
import { useReactionHandler } from "@/features/messages/ui/useReactionHandler";
import type { UserProfileLookup } from "@/features/profile/lib/identity";
import { resolveUserLabel } from "@/features/profile/lib/identity";
import { UserProfilePopover } from "@/features/profile/ui/UserProfilePopover";
import { cn } from "@/shared/lib/cn";
import { Button } from "@/shared/ui/button";
import { Popover, PopoverContent, PopoverTrigger } from "@/shared/ui/popover";
@@ -26,7 +27,7 @@ type SystemMessagePayload = {
type SystemMessageDescription = {
action: React.ReactNode;
title: string;
title: React.ReactNode;
};
function resolveLabel(
@@ -57,7 +58,7 @@ function resolveAvatarUrl(
return profiles[pubkey.toLowerCase()]?.avatarUrl ?? null;
}
function labelWithSuffix(
function resolveLabelWithSuffix(
pubkey: string | undefined,
currentPubkey: string | undefined,
profiles: UserProfileLookup | undefined,
@@ -66,6 +67,97 @@ function labelWithSuffix(
return `${resolveLabel(pubkey, currentPubkey, profiles)}${suffix}`;
}
function ProfileName({
children,
highlight = false,
pubkey,
}: {
children: React.ReactNode;
highlight?: boolean;
pubkey: string | undefined;
}) {
const node = (
<span
className={cn(
"rounded-sm transition-colors hover:text-foreground",
pubkey && "cursor-pointer",
highlight &&
"rounded-md bg-primary/15 px-1 py-0.5 font-medium text-primary hover:bg-primary/25 hover:text-primary/90",
)}
>
{highlight ? "@" : null}
{children}
</span>
);
return pubkey ? (
<UserProfilePopover pubkey={pubkey} triggerElement="span">
{node}
</UserProfilePopover>
) : (
node
);
}
function SystemMessageAvatar({
actorPubkey,
currentPubkey,
profiles,
targetPubkey,
}: {
actorPubkey: string | undefined;
currentPubkey: string | undefined;
profiles: UserProfileLookup | undefined;
targetPubkey: string | undefined;
}) {
const hasActorAndTarget =
actorPubkey && targetPubkey && actorPubkey !== targetPubkey;
const actorLabel = actorPubkey
? resolveUserLabel({
pubkey: actorPubkey,
currentPubkey,
profiles,
preferResolvedSelfLabel: true,
})
: "Someone";
if (!hasActorAndTarget) {
return (
<UserAvatar
avatarUrl={resolveAvatarUrl(actorPubkey ?? targetPubkey, profiles)}
className="!h-9 !w-9 shrink-0 text-[10px]"
displayName={actorLabel}
testId="system-message-avatar"
/>
);
}
const targetLabel = resolveUserLabel({
pubkey: targetPubkey,
currentPubkey,
profiles,
preferResolvedSelfLabel: true,
});
return (
<div
className="relative h-9 w-9 shrink-0"
data-testid="system-message-avatar"
>
<UserAvatar
avatarUrl={resolveAvatarUrl(actorPubkey, profiles)}
className="!h-7 !w-7 border-2 border-background text-[9px]"
displayName={actorLabel}
/>
<UserAvatar
avatarUrl={resolveAvatarUrl(targetPubkey, profiles)}
className="!absolute !bottom-0 !right-0 !h-7 !w-7 border-2 border-background text-[9px]"
displayName={targetLabel}
/>
</div>
);
}
function describeSystemEvent(
payload: SystemMessagePayload,
currentPubkey: string | undefined,
@@ -73,27 +165,39 @@ function describeSystemEvent(
personaLookup?: Map<string, string>,
): SystemMessageDescription | null {
const personaSuffix = resolvePersonaSuffix(payload.target, personaLookup);
const actorLabel = labelWithSuffix(payload.actor, currentPubkey, profiles);
const targetLabel = labelWithSuffix(
const actorLabel = resolveLabelWithSuffix(
payload.actor,
currentPubkey,
profiles,
);
const targetLabel = resolveLabelWithSuffix(
payload.target,
currentPubkey,
profiles,
personaSuffix,
);
const actorName = (
<ProfileName pubkey={payload.actor}>{actorLabel}</ProfileName>
);
const targetName = (
<ProfileName highlight pubkey={payload.target}>
{targetLabel}
</ProfileName>
);
switch (payload.type) {
case "member_joined": {
if (payload.actor === payload.target) {
return {
title: targetLabel,
title: targetName,
action: "joined the channel",
};
}
return {
title: actorLabel,
title: actorName,
action: (
<>
added <span className="font-medium">{targetLabel}</span> to the
added <span className="font-medium">{targetName}</span> to the
channel
</>
),
@@ -101,42 +205,42 @@ function describeSystemEvent(
}
case "member_left":
return {
title: actorLabel,
title: actorName,
action: "left the channel",
};
case "member_removed":
return {
title: actorLabel,
title: actorName,
action: (
<>
removed <span className="font-medium">{targetLabel}</span> from the
removed <span className="font-medium">{targetName}</span> from the
channel
</>
),
};
case "topic_changed":
return {
title: actorLabel,
title: actorName,
action: <>changed the topic to &ldquo;{payload.topic}&rdquo;</>,
};
case "purpose_changed":
return {
title: actorLabel,
title: actorName,
action: <>changed the purpose to &ldquo;{payload.purpose}&rdquo;</>,
};
case "channel_created":
return {
title: actorLabel,
title: actorName,
action: "created this channel",
};
case "channel_archived":
return {
title: actorLabel,
title: actorName,
action: "archived this channel",
};
case "channel_unarchived":
return {
title: actorLabel,
title: actorName,
action: "unarchived this channel",
};
default:
@@ -188,33 +292,23 @@ export const SystemMessageRow = React.memo(function SystemMessageRow({
return null;
}
const avatarPubkey = payload.actor ?? payload.target;
const avatarLabel = avatarPubkey
? resolveUserLabel({
pubkey: avatarPubkey,
currentPubkey,
profiles,
preferResolvedSelfLabel: true,
})
: "Someone";
return (
<div
className="group/message relative rounded-2xl px-2 py-1 transition-colors"
data-testid="system-message-row"
>
<div className="flex items-start gap-2.5">
<UserAvatar
avatarUrl={resolveAvatarUrl(avatarPubkey, profiles)}
className="!h-9 !w-9 shrink-0 text-[10px]"
displayName={avatarLabel}
testId="system-message-avatar"
<SystemMessageAvatar
actorPubkey={payload.actor}
currentPubkey={currentPubkey}
profiles={profiles}
targetPubkey={payload.target}
/>
<div className="min-w-0 flex-1">
<div className="flex min-w-0 flex-wrap items-baseline gap-x-2 gap-y-0.5">
<p className="truncate text-sm font-semibold leading-none tracking-tight text-foreground/90">
<div className="truncate text-sm font-semibold leading-none tracking-tight text-foreground/90">
{description.title}
</p>
</div>
<MessageTimestamp
createdAt={message.createdAt}
time={message.time}

View File

@@ -33,6 +33,7 @@ type OnboardingGateStage = "blocking" | "onboarding" | "ready";
type UseFirstRunOnboardingGateOptions = {
currentPubkey: string | null;
hasExistingProfile: boolean;
identityIsFetching: boolean;
identityStatus: QueryStatus;
profileStatus: QueryStatus;
@@ -125,6 +126,7 @@ function resolveOnboardingGateStage({
export function useFirstRunOnboardingGate({
currentPubkey,
hasExistingProfile,
identityIsFetching,
identityStatus,
profileStatus,
@@ -166,14 +168,37 @@ export function useFirstRunOnboardingGate({
return;
}
// If the relay already has a real profile for this pubkey, the user has
// been onboarded elsewhere — skip the gate and persist the completion so
// future launches in this data dir don't re-check.
if (hasExistingProfile) {
if (typeof window !== "undefined" && currentPubkey) {
window.localStorage.setItem(
onboardingCompletionStorageKey(currentPubkey),
"true",
);
}
}
setGateState((current) =>
updateActiveGateState(current, currentPubkey, (activeGateState) => ({
...activeGateState,
hasSettledCurrentPubkey: true,
isOpen: !activeGateState.hasCompletedCurrentPubkey,
})),
updateActiveGateState(current, currentPubkey, (activeGateState) => {
const alreadyOnboarded =
activeGateState.hasCompletedCurrentPubkey || hasExistingProfile;
return {
...activeGateState,
hasCompletedCurrentPubkey: alreadyOnboarded,
hasSettledCurrentPubkey: true,
isOpen: !alreadyOnboarded,
};
}),
);
}, [currentPubkey, hasSettledCurrentPubkey, identityStatus, profileStatus]);
}, [
currentPubkey,
hasExistingProfile,
hasSettledCurrentPubkey,
identityStatus,
profileStatus,
]);
const skipForNow = React.useCallback(() => {
setGateState((current) =>
@@ -213,6 +238,14 @@ export function useFirstRunOnboardingGate({
};
}
function hasRealDisplayName(displayName?: string | null): boolean {
if (!displayName) return false;
const trimmed = displayName.trim();
if (trimmed.length === 0) return false;
const lower = trimmed.toLowerCase();
return !lower.startsWith("npub1") && !lower.startsWith("nostr:npub1");
}
export function useAppOnboardingState() {
const queryClient = useQueryClient();
const identityQuery = useIdentityQuery();
@@ -221,6 +254,7 @@ export function useAppOnboardingState() {
const profileQuery = useProfileQuery();
const onboardingGate = useFirstRunOnboardingGate({
currentPubkey,
hasExistingProfile: hasRealDisplayName(profileQuery.data?.displayName),
identityIsFetching: identityQuery.fetchStatus === "fetching",
identityStatus: identityQuery.status,
profileStatus: profileQuery.status,

View File

@@ -19,6 +19,7 @@ import { BotIdenticon } from "@/features/messages/ui/BotIdenticon";
type UserProfilePopoverProps = {
children: React.ReactNode;
pubkey: string;
triggerElement?: "div" | "span";
/** When set to "bot", a BotIdenticon badge renders next to the display name. */
role?: string;
/** Value used to generate the BotIdenticon glyph (typically the author name). */
@@ -58,6 +59,7 @@ function truncatePubkey(pubkey: string) {
export function UserProfilePopover({
children,
pubkey,
triggerElement = "div",
role,
botIdenticonValue,
}: UserProfilePopoverProps) {
@@ -130,11 +132,12 @@ export function UserProfilePopover({
return () => clearHoverTimer();
}, [clearHoverTimer]);
const TriggerElement = triggerElement;
return (
<Popover onOpenChange={setOpen} open={open}>
<PopoverAnchor asChild>
{/* biome-ignore lint/a11y/useSemanticElements: wrapper div for hover/click behavior */}
<div
<TriggerElement
role="button"
tabIndex={0}
onClick={handleTriggerClick}
@@ -152,7 +155,7 @@ export function UserProfilePopover({
className="inline-flex"
>
{children}
</div>
</TriggerElement>
</PopoverAnchor>
<PopoverContent
align="start"

View File

@@ -1,21 +1,23 @@
import * as React from "react";
import {
ArrowRight,
FileText,
Hash,
LoaderCircle,
MessagesSquare,
Search,
type LucideIcon,
} from "lucide-react";
import {
resolveUserLabel,
resolveUserSecondaryLabel,
} from "@/features/profile/lib/identity";
import { useUsersBatchQuery } from "@/features/profile/hooks";
import { useSearchMessagesQuery } from "@/features/search/hooks";
import type { Channel, SearchHit } from "@/shared/api/types";
import {
ChannelResultBody,
MessageResultBody,
resultIcon,
resultKey,
resultTestId,
SearchResultShell,
type SearchResult,
} from "@/features/search/ui/SearchResultItem";
import {
Dialog,
DialogContent,
@@ -23,68 +25,11 @@ import {
DialogHeader,
DialogTitle,
} from "@/shared/ui/dialog";
import { Badge } from "@/shared/ui/badge";
import { Input } from "@/shared/ui/input";
import { Skeleton } from "@/shared/ui/skeleton";
import { UserAvatar } from "@/shared/ui/UserAvatar";
const MIN_QUERY_LENGTH = 2;
function describeSearchHit(hit: SearchHit) {
switch (hit.kind) {
case 1:
return "Note";
case 45001:
return "Forum post";
case 45003:
return "Forum reply";
case 43001:
return "Agent job";
case 43003:
return "Agent update";
case 46010:
return "Approval request";
default:
return "Message";
}
}
function truncateContent(content: string) {
const trimmed = content.trim();
if (trimmed.length === 0) {
return "No message body.";
}
if (trimmed.length <= 180) {
return trimmed;
}
return `${trimmed.slice(0, 177)}...`;
}
function formatRelativeTime(unixSeconds: number) {
const diff = Math.floor(Date.now() / 1_000) - unixSeconds;
if (diff < 60) {
return "just now";
}
if (diff < 60 * 60) {
return `${Math.floor(diff / 60)}m ago`;
}
if (diff < 60 * 60 * 24) {
return `${Math.floor(diff / (60 * 60))}h ago`;
}
return new Intl.DateTimeFormat("en-US", {
month: "short",
day: "numeric",
hour: "numeric",
minute: "2-digit",
}).format(new Date(unixSeconds * 1_000));
}
function SearchState({
icon: Icon,
title,
@@ -129,6 +74,7 @@ type SearchDialogProps = {
currentPubkey?: string;
open: boolean;
onOpenChange: (open: boolean) => void;
onOpenChannel: (channelId: string) => void;
onOpenResult: (hit: SearchHit) => void;
};
@@ -137,6 +83,7 @@ export function SearchDialog({
currentPubkey,
open,
onOpenChange,
onOpenChannel,
onOpenResult,
}: SearchDialogProps) {
const [query, setQuery] = React.useState("");
@@ -153,21 +100,69 @@ export function SearchDialog({
limit: 12,
});
const results = searchQuery.data?.hits ?? [];
const messageResults = searchQuery.data?.hits ?? [];
const channelResults = React.useMemo(() => {
if (debouncedQuery.length < MIN_QUERY_LENGTH) {
return [];
}
const normalizedQuery = debouncedQuery.toLowerCase();
return channels
.filter(
(channel) =>
channel.channelType !== "dm" &&
(channel.archivedAt
? channel.isMember
: channel.visibility === "open" || channel.isMember) &&
(channel.name.toLowerCase().includes(normalizedQuery) ||
channel.description.toLowerCase().includes(normalizedQuery)),
)
.sort((a, b) => {
const aNameMatches = a.name.toLowerCase().includes(normalizedQuery);
const bNameMatches = b.name.toLowerCase().includes(normalizedQuery);
if (aNameMatches !== bNameMatches) {
return aNameMatches ? -1 : 1;
}
return a.name.localeCompare(b.name);
})
.slice(0, 5);
}, [channels, debouncedQuery]);
const results = React.useMemo<SearchResult[]>(
() => [
...channelResults.map((channel) => ({
kind: "channel" as const,
channel,
})),
...messageResults.map((hit) => ({
kind: "message" as const,
hit,
})),
],
[channelResults, messageResults],
);
const resultProfilesQuery = useUsersBatchQuery(
results.map((hit) => hit.pubkey),
messageResults.map((hit) => hit.pubkey),
{
enabled: open && results.length > 0,
enabled: open && messageResults.length > 0,
},
);
const resultProfiles = resultProfilesQuery.data?.profiles;
const openResult = React.useCallback(
(hit: SearchHit) => {
(result: SearchResult) => {
onOpenChange(false);
onOpenResult(hit);
if (result.kind === "channel") {
onOpenChannel(result.channel.id);
return;
}
onOpenResult(result.hit);
},
[onOpenChange, onOpenResult],
[onOpenChange, onOpenChannel, onOpenResult],
);
React.useEffect(() => {
@@ -204,7 +199,7 @@ export function SearchDialog({
});
}, [results]);
const selectedHit = results[selectedIndex];
const selectedResult = results[selectedIndex];
return (
<Dialog onOpenChange={onOpenChange} open={open}>
@@ -254,10 +249,10 @@ export function SearchDialog({
if (
event.key === "Enter" &&
!event.nativeEvent.isComposing &&
selectedHit
selectedResult
) {
event.preventDefault();
openResult(selectedHit);
openResult(selectedResult);
}
}}
placeholder="Search messages, approvals, and forum posts"
@@ -277,9 +272,9 @@ export function SearchDialog({
icon={MessagesSquare}
title="Search message history"
/>
) : searchQuery.isLoading ? (
) : searchQuery.isLoading && results.length === 0 ? (
<SearchLoadingState />
) : searchQuery.error instanceof Error ? (
) : searchQuery.error instanceof Error && results.length === 0 ? (
<SearchState
description={searchQuery.error.message}
icon={LoaderCircle}
@@ -294,86 +289,35 @@ export function SearchDialog({
) : (
<div className="p-3" data-testid="search-results">
<div className="mb-3 flex items-center justify-between px-2 text-xs font-semibold uppercase tracking-[0.16em] text-muted-foreground">
<span>{searchQuery.data?.found ?? results.length} results</span>
<span>
{channelResults.length +
(searchQuery.data?.found ?? messageResults.length)}{" "}
results
</span>
<span>Enter to open</span>
</div>
<div className="space-y-2">
{results.map((hit, index) => {
const channel = hit.channelId
? channelLookup.get(hit.channelId)
: undefined;
const authorLabel = resolveUserLabel({
pubkey: hit.pubkey,
currentPubkey,
profiles: resultProfiles,
preferResolvedSelfLabel: true,
});
const authorSecondaryLabel = resolveUserSecondaryLabel({
pubkey: hit.pubkey,
profiles: resultProfiles,
});
return (
<button
className={
index === selectedIndex
? "w-full rounded-2xl border border-primary/30 bg-primary/10 px-4 py-4 text-left shadow-sm outline-none transition-colors"
: "w-full rounded-2xl border border-border/80 bg-card/60 px-4 py-4 text-left shadow-sm outline-none transition-colors hover:border-primary/20 hover:bg-accent"
}
data-testid={`search-result-${hit.eventId}`}
key={hit.eventId}
onClick={() => openResult(hit)}
onMouseEnter={() => setSelectedIndex(index)}
type="button"
>
<div className="flex items-start gap-3">
<div className="mt-0.5 flex h-9 w-9 shrink-0 items-center justify-center rounded-2xl bg-secondary text-secondary-foreground">
{channel?.channelType === "forum" ? (
<FileText className="h-4 w-4" />
) : (
<Hash className="h-4 w-4" />
)}
</div>
<div className="min-w-0 flex-1">
<div className="flex flex-wrap items-center gap-2">
<p className="text-sm font-semibold tracking-tight">
{hit.channelName}
</p>
<Badge variant="secondary">
{describeSearchHit(hit)}
</Badge>
<span className="flex items-center gap-1 text-xs text-muted-foreground">
<UserAvatar
avatarUrl={
resultProfiles?.[hit.pubkey.toLowerCase()]
?.avatarUrl ?? null
}
displayName={authorLabel}
size="xs"
/>
{authorLabel}
</span>
<p className="ml-auto whitespace-nowrap text-xs text-muted-foreground">
{formatRelativeTime(hit.createdAt)}
</p>
</div>
{authorSecondaryLabel ? (
<p className="mt-1 text-xs text-muted-foreground">
{authorSecondaryLabel}
</p>
) : null}
<p className="mt-2 text-sm leading-6 text-foreground">
{truncateContent(hit.content)}
</p>
</div>
<ArrowRight className="mt-1 h-4 w-4 shrink-0 text-muted-foreground" />
</div>
</button>
);
})}
{results.map((result, index) => (
<SearchResultShell
icon={resultIcon(result, channelLookup)}
isSelected={index === selectedIndex}
key={resultKey(result)}
onClick={() => openResult(result)}
onMouseEnter={() => setSelectedIndex(index)}
testId={resultTestId(result)}
>
{result.kind === "channel" ? (
<ChannelResultBody channel={result.channel} />
) : (
<MessageResultBody
currentPubkey={currentPubkey}
hit={result.hit}
resultProfiles={resultProfiles}
/>
)}
</SearchResultShell>
))}
</div>
</div>
)}

View File

@@ -0,0 +1,212 @@
import type * as React from "react";
import { ArrowRight, FileText, Hash, type LucideIcon } from "lucide-react";
import {
resolveUserLabel,
resolveUserSecondaryLabel,
type UserProfileLookup,
} from "@/features/profile/lib/identity";
import type { Channel, SearchHit } from "@/shared/api/types";
import { Badge } from "@/shared/ui/badge";
import { UserAvatar } from "@/shared/ui/UserAvatar";
export type SearchResult =
| { kind: "channel"; channel: Channel }
| { kind: "message"; hit: SearchHit };
export function resultKey(result: SearchResult) {
if (result.kind === "channel") {
return `channel-${result.channel.id}`;
}
return `message-${result.hit.eventId}`;
}
export function resultTestId(result: SearchResult) {
if (result.kind === "channel") {
return `search-result-channel-${result.channel.id}`;
}
return `search-result-${result.hit.eventId}`;
}
export function resultIcon(
result: SearchResult,
channelLookup: ReadonlyMap<string, Channel>,
) {
const channelType =
result.kind === "channel"
? result.channel.channelType
: result.hit.channelId
? channelLookup.get(result.hit.channelId)?.channelType
: undefined;
return channelType === "forum" ? FileText : Hash;
}
export function SearchResultShell({
children,
icon: Icon,
isSelected,
onClick,
onMouseEnter,
testId,
}: {
children: React.ReactNode;
icon: LucideIcon;
isSelected: boolean;
onClick: () => void;
onMouseEnter: () => void;
testId: string;
}) {
return (
<button
className={
isSelected
? "w-full rounded-2xl border border-primary/30 bg-primary/10 px-4 py-4 text-left shadow-sm outline-none transition-colors"
: "w-full rounded-2xl border border-border/80 bg-card/60 px-4 py-4 text-left shadow-sm outline-none transition-colors hover:border-primary/20 hover:bg-accent"
}
data-testid={testId}
onClick={onClick}
onMouseEnter={onMouseEnter}
type="button"
>
<div className="flex items-start gap-3">
<div className="mt-0.5 flex h-9 w-9 shrink-0 items-center justify-center rounded-2xl bg-secondary text-secondary-foreground">
<Icon className="h-4 w-4" />
</div>
{children}
<ArrowRight className="mt-1 h-4 w-4 shrink-0 text-muted-foreground" />
</div>
</button>
);
}
export function ChannelResultBody({ channel }: { channel: Channel }) {
return (
<div className="min-w-0 flex-1">
<div className="flex flex-wrap items-center gap-2">
<p className="text-sm font-semibold tracking-tight">{channel.name}</p>
<Badge variant="secondary">{channel.channelType}</Badge>
<p className="ml-auto whitespace-nowrap text-xs text-muted-foreground">
Channel
</p>
</div>
{channel.description ? (
<p className="mt-2 text-sm leading-6 text-foreground">
{channel.description}
</p>
) : null}
</div>
);
}
function describeSearchHit(hit: SearchHit) {
switch (hit.kind) {
case 1:
return "Note";
case 45001:
return "Forum post";
case 45003:
return "Forum reply";
case 43001:
return "Agent job";
case 43003:
return "Agent update";
case 46010:
return "Approval request";
default:
return "Message";
}
}
function truncateContent(content: string) {
const trimmed = content.trim();
if (trimmed.length === 0) {
return "No message body.";
}
if (trimmed.length <= 180) {
return trimmed;
}
return `${trimmed.slice(0, 177)}...`;
}
function formatRelativeTime(unixSeconds: number) {
const diff = Math.floor(Date.now() / 1_000) - unixSeconds;
if (diff < 60) {
return "just now";
}
if (diff < 60 * 60) {
return `${Math.floor(diff / 60)}m ago`;
}
if (diff < 60 * 60 * 24) {
return `${Math.floor(diff / (60 * 60))}h ago`;
}
return new Intl.DateTimeFormat("en-US", {
month: "short",
day: "numeric",
hour: "numeric",
minute: "2-digit",
}).format(new Date(unixSeconds * 1_000));
}
export function MessageResultBody({
currentPubkey,
hit,
resultProfiles,
}: {
currentPubkey?: string;
hit: SearchHit;
resultProfiles?: UserProfileLookup;
}) {
const authorLabel = resolveUserLabel({
pubkey: hit.pubkey,
currentPubkey,
profiles: resultProfiles,
preferResolvedSelfLabel: true,
});
const authorSecondaryLabel = resolveUserSecondaryLabel({
pubkey: hit.pubkey,
profiles: resultProfiles,
});
const avatarUrl =
resultProfiles?.[hit.pubkey.toLowerCase()]?.avatarUrl ?? null;
return (
<div className="min-w-0 flex-1">
<div className="flex flex-wrap items-center gap-2">
<p className="text-sm font-semibold tracking-tight">
{hit.channelName}
</p>
<Badge variant="secondary">{describeSearchHit(hit)}</Badge>
<span className="flex items-center gap-1 text-xs text-muted-foreground">
<UserAvatar
avatarUrl={avatarUrl}
displayName={authorLabel}
size="xs"
/>
{authorLabel}
</span>
<p className="ml-auto whitespace-nowrap text-xs text-muted-foreground">
{formatRelativeTime(hit.createdAt)}
</p>
</div>
{authorSecondaryLabel ? (
<p className="mt-1 text-xs text-muted-foreground">
{authorSecondaryLabel}
</p>
) : null}
<p className="mt-2 text-sm leading-6 text-foreground">
{truncateContent(hit.content)}
</p>
</div>
);
}

View File

@@ -33,3 +33,28 @@ export function resolveMentionNames(
return names.size > 0 ? [...names] : undefined;
}
export function resolveMentionPubkeysByName(
tags: string[][] | undefined,
profiles: Record<string, UserProfileSummary> | undefined,
): Record<string, string> | undefined {
if (!profiles || !tags) {
return undefined;
}
const pubkeysByName: Record<string, string> = {};
for (const tag of tags) {
if (tag[0] !== "p" || !tag[1]) {
continue;
}
const pubkey = tag[1].toLowerCase();
const displayName = profiles[pubkey]?.displayName?.trim();
if (displayName) {
pubkeysByName[displayName.toLowerCase()] = pubkey;
}
}
return Object.keys(pubkeysByName).length > 0 ? pubkeysByName : undefined;
}

View File

@@ -7,6 +7,7 @@ import remarkGfm from "remark-gfm";
import { toast } from "sonner";
import { useAppNavigation } from "@/app/navigation/useAppNavigation";
import { UserProfilePopover } from "@/features/profile/ui/UserProfilePopover";
import { invokeTauri } from "@/shared/api/tauri";
import type { Channel } from "@/shared/api/types";
import { useChannelNavigation } from "@/shared/context/ChannelNavigationContext";
@@ -34,6 +35,7 @@ type MarkdownProps = {
content: string;
imetaByUrl?: ImetaLookup;
mentionNames?: string[];
mentionPubkeysByName?: Record<string, string>;
searchQuery?: string;
tight?: boolean;
};
@@ -106,6 +108,7 @@ function createMarkdownComponents(
channels: Channel[],
onOpenChannel: (channelId: string) => void,
imetaByUrl?: ImetaLookup,
mentionPubkeysByName?: Record<string, string>,
): Components {
const paragraphClassName =
variant === "tight"
@@ -322,14 +325,27 @@ function createMarkdownComponents(
ul: ({ children }) => (
<ul className={cn("list-disc", listClassName)}>{children}</ul>
),
mention: ({ children }: { children?: React.ReactNode }) => (
<span
data-mention=""
className="rounded-md bg-primary/15 px-1 py-0.5 text-sm font-semibold text-primary"
>
{children}
</span>
),
mention: ({ children }: { children?: React.ReactNode }) => {
const mentionText = String(children ?? "");
const mentionName = mentionText.replace(/^@/, "").trim().toLowerCase();
const pubkey = mentionPubkeysByName?.[mentionName];
const mentionNode = (
<span
data-mention=""
className="cursor-pointer rounded-md bg-primary/15 px-1 py-0.5 text-sm font-semibold text-primary transition-colors hover:bg-primary/25 hover:text-primary/90"
>
{children}
</span>
);
return pubkey ? (
<UserProfilePopover pubkey={pubkey} triggerElement="span">
{mentionNode}
</UserProfilePopover>
) : (
mentionNode
);
},
"channel-link": ({ children }: { children?: React.ReactNode }) => {
const text = String(children ?? "");
const channelName = text.startsWith("#") ? text.slice(1) : text;
@@ -374,6 +390,7 @@ function MarkdownInner({
content,
imetaByUrl,
mentionNames,
mentionPubkeysByName,
searchQuery,
tight = false,
}: MarkdownProps) {
@@ -394,8 +411,9 @@ function MarkdownInner({
void goChannel(channelId);
},
imetaByUrl,
mentionPubkeysByName,
),
[goChannel, variant, channels, imetaByUrl],
[goChannel, variant, channels, imetaByUrl, mentionPubkeysByName],
);
// biome-ignore lint/suspicious/noExplicitAny: PluggableList type not directly importable
@@ -462,6 +480,7 @@ export const Markdown = React.memo(
prev.className === next.className &&
prev.compact === next.compact &&
prev.tight === next.tight &&
prev.mentionPubkeysByName === next.mentionPubkeysByName &&
shallowArrayEqual(prev.mentionNames, next.mentionNames) &&
shallowArrayEqual(prev.channelNames, next.channelNames) &&
prev.imetaByUrl === next.imetaByUrl &&

View File

@@ -883,7 +883,7 @@ const mockChannels: MockChannel[] = [
name: "design",
channel_type: "stream",
visibility: "open",
description: "Design system and UX discussions",
description: "Design system and UX discussions with engineering partners",
topic: null,
purpose: null,
last_message_at: isoMinutesAgo(120),

View File

@@ -20,9 +20,25 @@ test("keyboard shortcut opens the channel browser dialog", async ({ page }) => {
await page.goto("/");
await expect(page.getByTestId("app-sidebar")).toBeVisible();
await page.keyboard.press(
process.platform === "darwin" ? "Meta+Shift+O" : "Control+Shift+O",
const isMacBrowser = await page.evaluate(() =>
/mac|iphone|ipad|ipod/i.test(navigator.platform),
);
if (isMacBrowser) {
await page.evaluate(() => {
window.dispatchEvent(
new KeyboardEvent("keydown", {
bubbles: true,
cancelable: true,
key: "O",
metaKey: true,
shiftKey: true,
}),
);
});
} else {
await page.keyboard.press("Control+Shift+O");
}
await expect(page.getByTestId("channel-browser-dialog")).toBeVisible();
});

View File

@@ -9,6 +9,10 @@ const BLANK_TYLER_IDENTITY = {
...TEST_IDENTITIES.tyler,
username: "",
};
const FIRST_RUN_ALICE = {
...TEST_IDENTITIES.alice,
username: "",
};
type TestIdentity = {
privateKey: string;
@@ -139,17 +143,16 @@ test("page 1 accepts an avatar URL as the secondary avatar path", async ({
test("first-run onboarding keeps the shell hidden through both pages and only marks Home seen after finish", async ({
page,
}) => {
await seedActiveIdentity(page, TEST_IDENTITIES.alice);
await seedActiveIdentity(page, FIRST_RUN_ALICE);
await installMockBridge(page, undefined, { skipOnboardingSeed: true });
await page.goto("/");
await expect(page.getByTestId("onboarding-gate")).toBeVisible();
await expect(page.getByTestId("onboarding-page-1")).toBeVisible();
await expect(page.getByTestId("onboarding-display-name")).toHaveValue(
"alice",
);
await expect(page.getByTestId("onboarding-display-name")).toHaveValue("");
await expectNoHomeSeenEntries(page);
await page.getByTestId("onboarding-display-name").fill("Alice");
await continueToSetupPage(page);
await expectShellHidden(page);
await expect(page.getByTestId("onboarding-provider-goose")).toBeVisible();
@@ -161,6 +164,17 @@ test("first-run onboarding keeps the shell hidden through both pages and only ma
await expectHomeSeenCount(page, 2);
});
test("existing relay profile auto-skips onboarding without localStorage completion", async ({
page,
}) => {
await seedActiveIdentity(page, TEST_IDENTITIES.alice);
await installMockBridge(page, undefined, { skipOnboardingSeed: true });
await page.goto("/");
await expect(page.getByTestId("onboarding-gate")).toHaveCount(0);
await expect(page.getByTestId("chat-title")).toHaveText("Home");
});
test("finishing onboarding auto-joins the #general channel for a new member", async ({
page,
}) => {
@@ -179,7 +193,7 @@ test("finishing onboarding auto-joins the #general channel for a new member", as
test("page 2 falls back to Doctor guidance when ACP tools are not installed", async ({
page,
}) => {
await seedActiveIdentity(page, TEST_IDENTITIES.alice);
await seedActiveIdentity(page, FIRST_RUN_ALICE);
await installMockBridge(
page,
{
@@ -189,6 +203,7 @@ test("page 2 falls back to Doctor guidance when ACP tools are not installed", as
);
await page.goto("/");
await page.getByTestId("onboarding-display-name").fill("Alice");
await continueToSetupPage(page);
await expect(page.getByTestId("onboarding-acp-empty")).toBeVisible();
await expect(

View File

@@ -246,6 +246,36 @@ test("opens relay-backed search from the sidebar and loads the exact result", as
);
});
test("opens channel matches from search", async ({ page }) => {
await page.goto("/");
await openSearchDialogWithButton(page);
await page.getByTestId("search-input").fill("engineering");
const results = page.getByTestId("search-results");
await expect(results).toContainText("engineering");
await expect(results).toContainText("Engineering discussions");
await expect(results).toContainText(
"Design system and UX discussions with engineering partners",
);
await expect(
results.locator('[data-testid^="search-result-channel-"]').first(),
).toHaveAttribute(
"data-testid",
"search-result-channel-1c7e1c02-87bb-5e88-b2da-5a7a9432d0c9",
);
await results
.getByTestId("search-result-channel-1c7e1c02-87bb-5e88-b2da-5a7a9432d0c9")
.click();
await expect(page).toHaveURL(
/#\/channels\/1c7e1c02-87bb-5e88-b2da-5a7a9432d0c9$/,
);
await expect(page.getByTestId("chat-title")).toHaveText("engineering");
});
test("search results use your resolved profile label instead of You", async ({
page,
}) => {

View File

@@ -660,10 +660,10 @@ packages:
dependency: "direct main"
description:
name: lucide_icons_flutter
sha256: df29c1cf4f19f9309f5204d40e0ba95338d9946a8ffdb4d75cf9aaccf449ce41
sha256: f9fd5d49b93bf14b89e0ff4818658a74ab16899bdaf7aa745358ee1a34f54eed
url: "https://pub.dev"
source: hosted
version: "3.1.12"
version: "3.1.14+1"
matcher:
dependency: transitive
description:

View File

@@ -5,6 +5,7 @@
# SPROUT_RELAY_PORT, SPROUT_RELAY_URL
# SPROUT_INSTANCE_SLUG, SPROUT_WORKTREE_LABEL, VITE_DEV_BRANCH (worktrees only)
# SPROUT_TAURI_CONFIG
# SPROUT_PRIVATE_KEY (worktrees only, when SPROUT_SHARE_IDENTITY=1)
WORKTREE_ROOT=$(git rev-parse --show-toplevel 2>/dev/null || pwd)
@@ -24,13 +25,31 @@ unset VITE_DEV_BRANCH
# In worktrees, extract a label from the branch name and derive a unique app
# identity and icon so multiple local desktop instances can run side by side.
#
# Worktree detection: compare --git-dir to --git-common-dir. In the main
# working tree these are identical; in any worktree (whether under .worktrees/,
# .claude/worktrees/, or elsewhere on disk) they differ.
if git rev-parse --is-inside-work-tree &>/dev/null; then
GIT_DIR=$(git rev-parse --git-dir)
if [[ "$GIT_DIR" == *".git/worktrees/"* ]]; then
GIT_COMMON_DIR=$(git rev-parse --git-common-dir 2>/dev/null)
if [[ -n "$GIT_COMMON_DIR" && "$GIT_DIR" != "$GIT_COMMON_DIR" ]]; then
BRANCH_NAME=$(git rev-parse --abbrev-ref HEAD)
export SPROUT_WORKTREE_LABEL="${BRANCH_NAME##*/}"
export SPROUT_INSTANCE_SLUG=$(echo "$BRANCH_NAME" | tr '[:upper:]' '[:lower:]' | sed 's/[^a-z0-9]/-/g' | sed 's/--*/-/g' | sed 's/^-//' | sed 's/-$//')
# SPROUT_SHARE_IDENTITY=1: reuse the main dev checkout's Nostr key so
# worktrees skip onboarding and share the same identity. The per-worktree
# identifier is kept so concurrent instances don't collide on
# tauri-plugin-single-instance or the app data directory.
if [[ "${SPROUT_SHARE_IDENTITY:-0}" == "1" ]]; then
CANONICAL_KEY="$HOME/Library/Application Support/xyz.block.sprout.app.dev/identity.key"
if [[ -f "$CANONICAL_KEY" ]]; then
export SPROUT_PRIVATE_KEY="$(cat "$CANONICAL_KEY")"
else
echo "⚠ SPROUT_SHARE_IDENTITY=1 but no identity found at $CANONICAL_KEY — run Sprout from repo root first" >&2
fi
fi
ICON_DIR="$(pwd)/src-tauri/target/dev-icons"
mkdir -p "$ICON_DIR"
DEV_ICON="$ICON_DIR/icon.icns"