mirror of
https://github.com/block/buzz.git
synced 2026-11-04 00:12:02 +00:00
Merge remote-tracking branch 'origin/main' into pulse-reaction
This commit is contained in:
3
Cargo.lock
generated
3
Cargo.lock
generated
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
248
crates/sprout-acp/src/engram_fetch.rs
Normal file
248
crates/sprout-acp/src/engram_fetch.rs
Normal 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:?}");
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
364
crates/sprout-cli/src/commands/mem.rs
Normal file
364
crates/sprout-cli/src/commands/mem.rs
Normal 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,
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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!({
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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 }
|
||||
|
||||
835
crates/sprout-core/src/engram.rs
Normal file
835
crates/sprout-core/src/engram.rs
Normal 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(_)));
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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))) => {
|
||||
|
||||
@@ -11,6 +11,7 @@
|
||||
|
||||
pub mod handler;
|
||||
pub mod room;
|
||||
pub mod wire;
|
||||
|
||||
pub use handler::ws_audio_handler;
|
||||
pub use room::AudioRoomManager;
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
168
crates/sprout-relay/src/audio/wire.rs
Normal file
168
crates/sprout-relay/src/audio/wire.rs
Normal 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());
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
337
desktop/src-tauri/Cargo.lock
generated
337
desktop/src-tauri/Cargo.lock
generated
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
117
desktop/src-tauri/examples/pocket_bench.rs
Normal file
117
desktop/src-tauri/examples/pocket_bench.rs
Normal 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"
|
||||
}
|
||||
);
|
||||
}
|
||||
148
desktop/src-tauri/examples/pocket_onset_probe.rs
Normal file
148
desktop/src-tauri/examples/pocket_onset_probe.rs
Normal 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
|
||||
);
|
||||
}
|
||||
}
|
||||
218
desktop/src-tauri/examples/prod_probe.rs
Normal file
218
desktop/src-tauri/examples/prod_probe.rs
Normal 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()
|
||||
}
|
||||
267
desktop/src-tauri/src/huddle/jitter.rs
Normal file
267
desktop/src-tauri/src/huddle/jitter.rs
Normal 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)",
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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}");
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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()));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
};
|
||||
|
||||
362
desktop/src-tauri/src/huddle/playout.rs
Normal file
362
desktop/src-tauri/src/huddle/playout.rs
Normal 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,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
872
desktop/src-tauri/src/huddle/pocket.rs
Normal file
872
desktop/src-tauri/src/huddle/pocket.rs
Normal 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"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
|
||||
292
desktop/src-tauri/src/huddle/wire.rs
Normal file
292
desktop/src-tauri/src/huddle/wire.rs
Normal 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);
|
||||
}
|
||||
}
|
||||
@@ -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).
|
||||
|
||||
142
desktop/src-tauri/tests/rodio_mixer_diagnostic.rs
Normal file
142
desktop/src-tauri/tests/rodio_mixer_diagnostic.rs
Normal 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}",
|
||||
);
|
||||
}
|
||||
@@ -70,6 +70,7 @@ export function AppShellOverlays({
|
||||
<SearchDialog
|
||||
channels={channels}
|
||||
currentPubkey={currentPubkey}
|
||||
onOpenChannel={onSelectChannel}
|
||||
onOpenResult={onOpenSearchResult}
|
||||
onOpenChange={onSearchOpenChange}
|
||||
open={true}
|
||||
|
||||
72
desktop/src/features/huddle/components/HeadphonesGate.tsx
Normal file
72
desktop/src/features/huddle/components/HeadphonesGate.tsx
Normal 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>
|
||||
);
|
||||
}
|
||||
41
desktop/src/features/huddle/components/HeadphonesNotice.tsx
Normal file
41
desktop/src/features/huddle/components/HeadphonesNotice.tsx
Normal 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>
|
||||
);
|
||||
}
|
||||
@@ -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>
|
||||
);
|
||||
|
||||
@@ -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}
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
78
desktop/src/features/huddle/lib/useHeadphonesGate.ts
Normal file
78
desktop/src/features/huddle/lib/useHeadphonesGate.ts
Normal 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 };
|
||||
}
|
||||
@@ -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 },
|
||||
],
|
||||
);
|
||||
});
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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" />
|
||||
|
||||
@@ -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 “{payload.topic}”</>,
|
||||
};
|
||||
case "purpose_changed":
|
||||
return {
|
||||
title: actorLabel,
|
||||
title: actorName,
|
||||
action: <>changed the purpose to “{payload.purpose}”</>,
|
||||
};
|
||||
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}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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>
|
||||
)}
|
||||
|
||||
212
desktop/src/features/search/ui/SearchResultItem.tsx
Normal file
212
desktop/src/features/search/ui/SearchResultItem.tsx
Normal 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>
|
||||
);
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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 &&
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
}) => {
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user