diff --git a/FEATURE_PARITY.md b/FEATURE_PARITY.md index 15ae45cbab..68d36f6a90 100644 --- a/FEATURE_PARITY.md +++ b/FEATURE_PARITY.md @@ -113,7 +113,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O |---------|----------|----------|-------| | Streaming draft replies | ✅ | ❌ | Partial replies via draft message updates | | Configurable stream modes | ✅ | ❌ | Per-channel stream behavior | -| Thread ownership | ✅ | 🚧 | Reply participation memory now persists with TTL-bounded tracking; full thread-level ownership tracking is still missing | +| Thread ownership | ✅ | 🚧 | Reply participation memory is restart-stable and TTL-bounded; once the bot joins a thread, follow-ups inherit channel visibility. Full thread-level ownership tracking is still missing | | Download-file action | ✅ | ❌ | On-demand attachment downloads via message actions | ### Mattermost-Specific Features (since Mar 2026) diff --git a/channels-src/slack/slack.capabilities.json b/channels-src/slack/slack.capabilities.json index 7035d9253b..5444fc066f 100644 --- a/channels-src/slack/slack.capabilities.json +++ b/channels-src/slack/slack.capabilities.json @@ -47,6 +47,9 @@ "messages_per_minute": 100, "messages_per_hour": 5000 }, + "durable_workspace_paths": [ + "state/active_threads" + ], "webhook": { "hmac_secret_name": "slack_signing_secret" } diff --git a/channels-src/slack/src/lib.rs b/channels-src/slack/src/lib.rs index 3b6b212c6b..62109ee04f 100644 --- a/channels-src/slack/src/lib.rs +++ b/channels-src/slack/src/lib.rs @@ -23,7 +23,7 @@ wit_bindgen::generate!({ }); use serde::{Deserialize, Serialize}; -use std::collections::BTreeMap; +use std::collections::HashMap; // Re-export generated types use exports::near::agent::channel::{ @@ -116,6 +116,21 @@ struct SlackMessageMetadata { team_id: Option, } +#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)] +struct ActiveSlackThreadKey { + team_id: Option, + channel: String, + thread_ts: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct ActiveSlackThread { + #[serde(flatten)] + key: ActiveSlackThreadKey, + #[serde(default)] + last_seen_ms: u64, +} + /// Slack API response for chat.postMessage. #[derive(Debug, Deserialize)] struct SlackPostMessageResponse { @@ -130,16 +145,54 @@ const OWNER_ID_PATH: &str = "state/owner_id"; const DM_POLICY_PATH: &str = "state/dm_policy"; /// Workspace path for persisting allow_from (JSON array) across WASM callbacks. const ALLOW_FROM_PATH: &str = "state/allow_from"; -/// Workspace path for tracking recently active Slack threads. -const ACTIVE_THREADS_PATH: &str = "state/active_threads.json"; -/// Recently active threads expire after 24 hours to avoid reviving stale threads forever. +/// Workspace path for thread timestamps the bot has already joined. +const ACTIVE_THREADS_PATH: &str = "state/active_threads"; +/// Threads expire after 24h of inactivity so the participation cache stays bounded. const ACTIVE_THREAD_TTL_MS: u64 = 24 * 60 * 60 * 1000; -/// Cap stored thread markers so the workspace state stays bounded. +/// Hard cap on remembered threads per workspace. const ACTIVE_THREAD_MAX_ENTRIES: usize = 256; /// Channel name for pairing store (used by pairing host APIs). const CHANNEL_NAME: &str = "slack"; -type ActiveThreads = BTreeMap; +#[cfg(not(test))] +fn host_workspace_read(path: &str) -> Option { + channel_host::workspace_read(path) +} + +#[cfg(test)] +fn host_workspace_read(path: &str) -> Option { + test_host::workspace_read(path) +} + +#[cfg(not(test))] +fn host_workspace_write(path: &str, content: &str) -> Result<(), String> { + channel_host::workspace_write(path, content) +} + +#[cfg(test)] +fn host_workspace_write(path: &str, content: &str) -> Result<(), String> { + test_host::workspace_write(path, content) +} + +#[cfg(not(test))] +fn host_emit_message(message: &EmittedMessage) { + channel_host::emit_message(message); +} + +#[cfg(test)] +fn host_emit_message(message: &EmittedMessage) { + test_host::emit_message(message); +} + +#[cfg(not(test))] +fn host_now_millis() -> u64 { + channel_host::now_millis() +} + +#[cfg(test)] +fn host_now_millis() -> u64 { + test_host::now_millis() +} /// Channel configuration from capabilities file. #[derive(Debug, Deserialize)] @@ -174,22 +227,22 @@ impl Guest for SlackChannel { // Persist owner_id so subsequent callbacks can read it if let Some(ref owner_id) = config.owner_id { - let _ = channel_host::workspace_write(OWNER_ID_PATH, owner_id); + let _ = host_workspace_write(OWNER_ID_PATH, owner_id); channel_host::log( channel_host::LogLevel::Info, &format!("Owner restriction enabled: user {}", owner_id), ); } else { - let _ = channel_host::workspace_write(OWNER_ID_PATH, ""); + let _ = host_workspace_write(OWNER_ID_PATH, ""); } // Persist dm_policy and allow_from for DM pairing let dm_policy = config.dm_policy.as_deref().unwrap_or("pairing"); - let _ = channel_host::workspace_write(DM_POLICY_PATH, dm_policy); + let _ = host_workspace_write(DM_POLICY_PATH, dm_policy); let allow_from_json = serde_json::to_string(&config.allow_from.unwrap_or_default()) .unwrap_or_else(|_| "[]".to_string()); - let _ = channel_host::workspace_write(ALLOW_FROM_PATH, &allow_from_json); + let _ = host_workspace_write(ALLOW_FROM_PATH, &allow_from_json); Ok(ChannelConfig { display_name: "Slack".to_string(), @@ -265,8 +318,7 @@ impl Guest for SlackChannel { let metadata: SlackMessageMetadata = serde_json::from_str(&response.metadata_json) .map_err(|e| format!("Failed to parse metadata: {}", e))?; - let thread_ts = response.thread_id.or(metadata.thread_ts); - + let thread_ts = response.thread_id.clone().or(metadata.thread_ts.clone()); let ts = post_slack_message( &metadata.channel, &response.content, @@ -274,7 +326,11 @@ impl Guest for SlackChannel { )?; if let Some(thread_ts) = thread_ts { - if let Err(e) = track_active_thread(&metadata.channel, &thread_ts) { + if let Err(e) = remember_active_slack_thread( + metadata.team_id.as_deref(), + &metadata.channel, + &thread_ts, + ) { channel_host::log( channel_host::LogLevel::Warn, &format!("Failed to track active thread: {}", e), @@ -511,14 +567,15 @@ fn handle_slack_event(event: SlackEvent, team_id: Option, _event_id: Opt event.ts.clone(), ) { let is_dm = channel.starts_with('D'); + let is_active_thread = event.thread_ts.as_deref().is_some_and(|thread_ts| { + is_active_slack_thread(team_id.as_deref(), &channel, thread_ts) + }); - // Check if this is a reply in a thread where we previously participated - let is_active_thread = !is_dm - && event - .thread_ts - .as_ref() - .is_some_and(|thread_ts| is_active_thread(&channel, thread_ts)); - + // DMs are always processed. For channel threads, once the bot + // has already replied in a thread we intentionally allow + // follow-ups from that thread without re-running DM pairing or + // allow_from checks. This matches Slack's app_mention behavior: + // the thread stays as visible as the surrounding channel. if is_dm || is_active_thread { if !check_sender_permission(&user, &channel, is_dm) { return; @@ -545,6 +602,209 @@ fn handle_slack_event(event: SlackEvent, team_id: Option, _event_id: Opt } } +/// Emit a message to the agent. +fn emit_message( + user_id: String, + text: String, + channel: String, + thread_ts: Option, + team_id: Option, + attachments: Vec, +) { + let message_ts = thread_ts.clone().unwrap_or_default(); + + let metadata = SlackMessageMetadata { + channel: channel.clone(), + thread_ts: thread_ts.clone(), + message_ts: message_ts.clone(), + team_id, + }; + + let metadata_json = serde_json::to_string(&metadata).unwrap_or_else(|e| { + channel_host::log( + channel_host::LogLevel::Error, + &format!("Failed to serialize Slack metadata: {}", e), + ); + "{}".to_string() + }); + + // Strip @ mentions of the bot from the text for cleaner messages + let cleaned_text = strip_bot_mention(&text); + + host_emit_message(&EmittedMessage { + user_id, + user_name: None, // Could fetch from Slack API if needed + content: cleaned_text, + thread_id: thread_ts, + metadata_json, + attachments, + }); +} + +fn active_slack_thread_key( + team_id: Option<&str>, + channel: &str, + thread_ts: &str, +) -> ActiveSlackThreadKey { + ActiveSlackThreadKey { + team_id: team_id.map(str::to_string), + channel: channel.to_string(), + thread_ts: thread_ts.to_string(), + } +} + +fn active_slack_thread_entry( + team_id: Option<&str>, + channel: &str, + thread_ts: &str, + last_seen_ms: u64, +) -> ActiveSlackThread { + ActiveSlackThread { + key: active_slack_thread_key(team_id, channel, thread_ts), + last_seen_ms, + } +} + +fn parse_active_slack_threads( + raw: Option<&str>, + now_ms: u64, +) -> HashMap { + raw.and_then(|value| serde_json::from_str::>(value).ok()) + .map(|threads| { + threads + .into_iter() + .map(|thread| { + ( + thread.key, + if thread.last_seen_ms == 0 { + now_ms + } else { + thread.last_seen_ms + }, + ) + }) + .collect() + }) + .or_else(|| { + raw.and_then(|value| serde_json::from_str::>(value).ok()) + .map(|legacy| { + legacy + .into_iter() + .map(|thread_ts| (active_slack_thread_key(None, "", &thread_ts), now_ms)) + .collect() + }) + }) + .unwrap_or_default() +} + +fn serialize_active_slack_threads(threads: &HashMap) -> String { + let mut sorted: Vec<_> = threads + .iter() + .map(|(key, last_seen_ms)| { + active_slack_thread_entry( + key.team_id.as_deref(), + &key.channel, + &key.thread_ts, + *last_seen_ms, + ) + }) + .collect(); + sorted.sort_unstable_by(|left, right| { + left.key + .team_id + .cmp(&right.key.team_id) + .then(left.key.channel.cmp(&right.key.channel)) + .then(left.key.thread_ts.cmp(&right.key.thread_ts)) + }); + serde_json::to_string(&sorted).unwrap_or_else(|_| "[]".to_string()) +} + +fn prune_active_slack_threads(threads: &mut HashMap, now_ms: u64) { + let cutoff = now_ms.saturating_sub(ACTIVE_THREAD_TTL_MS); + threads.retain(|_, last_seen_ms| *last_seen_ms >= cutoff); + + if threads.len() <= ACTIVE_THREAD_MAX_ENTRIES { + return; + } + + let mut entries: Vec<_> = threads + .iter() + .map(|(key, last_seen_ms)| (key.clone(), *last_seen_ms)) + .collect(); + entries.sort_unstable_by(|left, right| { + right + .1 + .cmp(&left.1) + .then(left.0.team_id.cmp(&right.0.team_id)) + .then(left.0.channel.cmp(&right.0.channel)) + .then(left.0.thread_ts.cmp(&right.0.thread_ts)) + }); + entries.truncate(ACTIVE_THREAD_MAX_ENTRIES); + *threads = entries.into_iter().collect(); +} + +fn load_active_slack_threads_from_workspace() -> HashMap { + let raw = host_workspace_read(ACTIVE_THREADS_PATH); + let now_ms = host_now_millis(); + let mut threads = parse_active_slack_threads(raw.as_deref(), now_ms); + prune_active_slack_threads(&mut threads, now_ms); + + let serialized = serialize_active_slack_threads(&threads); + let should_persist = raw.as_deref().is_some_and(|existing| existing != serialized) + || (raw.is_none() && !threads.is_empty()); + if should_persist { + let _ = host_workspace_write(ACTIVE_THREADS_PATH, &serialized); + } + + threads +} + +fn active_slack_thread_is_known( + raw: Option<&str>, + team_id: Option<&str>, + channel: &str, + thread_ts: &str, + now_ms: u64, +) -> bool { + let mut threads = parse_active_slack_threads(raw, now_ms); + prune_active_slack_threads(&mut threads, now_ms); + threads.contains_key(&active_slack_thread_key(team_id, channel, thread_ts)) + || threads.contains_key(&active_slack_thread_key(None, channel, thread_ts)) + || threads.contains_key(&active_slack_thread_key(None, "", thread_ts)) +} + +fn is_active_slack_thread(team_id: Option<&str>, channel: &str, thread_ts: &str) -> bool { + let threads = load_active_slack_threads_from_workspace(); + threads.contains_key(&active_slack_thread_key(team_id, channel, thread_ts)) + || threads.contains_key(&active_slack_thread_key(None, channel, thread_ts)) + || threads.contains_key(&active_slack_thread_key(None, "", thread_ts)) +} + +fn track_active_thread(channel: &str, thread_ts: &str) -> Result<(), String> { + remember_active_slack_thread(None, channel, thread_ts) +} + +fn remember_active_slack_thread( + team_id: Option<&str>, + channel: &str, + thread_ts: &str, +) -> Result<(), String> { + if channel.starts_with('D') { + return Ok(()); + } + + let now_ms = host_now_millis(); + let mut threads = load_active_slack_threads_from_workspace(); + let key = active_slack_thread_key(team_id, channel, thread_ts); + threads.insert(key, now_ms); + threads.remove(&active_slack_thread_key(None, "", thread_ts)); + prune_active_slack_threads(&mut threads, now_ms); + + host_workspace_write(ACTIVE_THREADS_PATH, &serialize_active_slack_threads(&threads)) +} + +type ActiveThreads = HashMap; + fn active_thread_key(channel: &str, thread_ts: &str) -> String { format!("{channel}/{thread_ts}") } @@ -582,95 +842,6 @@ fn prune_active_threads(active_threads: &mut ActiveThreads, now_millis: u64) -> changed } -fn load_active_threads() -> ActiveThreads { - let Some(raw) = channel_host::workspace_read(ACTIVE_THREADS_PATH) else { - return ActiveThreads::new(); - }; - - match serde_json::from_str(&raw) { - Ok(active_threads) => active_threads, - Err(e) => { - channel_host::log( - channel_host::LogLevel::Warn, - &format!("Failed to parse active thread state: {e}"), - ); - ActiveThreads::new() - } - } -} - -fn persist_active_threads(active_threads: &ActiveThreads) -> Result<(), String> { - let serialized = serde_json::to_string(active_threads) - .map_err(|e| format!("Failed to serialize active thread state: {e}"))?; - channel_host::workspace_write(ACTIVE_THREADS_PATH, &serialized) - .map_err(|e| format!("Failed to persist active thread state: {e}")) -} - -fn track_active_thread(channel: &str, thread_ts: &str) -> Result<(), String> { - let now_millis = channel_host::now_millis(); - let mut active_threads = load_active_threads(); - prune_active_threads(&mut active_threads, now_millis); - active_threads.insert(active_thread_key(channel, thread_ts), now_millis); - prune_active_threads(&mut active_threads, now_millis); - persist_active_threads(&active_threads) -} - -fn is_active_thread(channel: &str, thread_ts: &str) -> bool { - let now_millis = channel_host::now_millis(); - let mut active_threads = load_active_threads(); - let changed = prune_active_threads(&mut active_threads, now_millis); - - if changed { - if let Err(e) = persist_active_threads(&active_threads) { - channel_host::log( - channel_host::LogLevel::Warn, - &format!("Failed to prune active thread state: {e}"), - ); - } - } - - active_threads.contains_key(&active_thread_key(channel, thread_ts)) -} - -/// Emit a message to the agent. -fn emit_message( - user_id: String, - text: String, - channel: String, - thread_ts: Option, - team_id: Option, - attachments: Vec, -) { - let message_ts = thread_ts.clone().unwrap_or_default(); - - let metadata = SlackMessageMetadata { - channel: channel.clone(), - thread_ts: thread_ts.clone(), - message_ts: message_ts.clone(), - team_id, - }; - - let metadata_json = serde_json::to_string(&metadata).unwrap_or_else(|e| { - channel_host::log( - channel_host::LogLevel::Error, - &format!("Failed to serialize Slack metadata: {}", e), - ); - "{}".to_string() - }); - - // Strip @ mentions of the bot from the text for cleaner messages - let cleaned_text = strip_bot_mention(&text); - - channel_host::emit_message(&EmittedMessage { - user_id, - user_name: None, // Could fetch from Slack API if needed - content: cleaned_text, - thread_id: thread_ts, - metadata_json, - attachments, - }); -} - // ============================================================================ // Permission & Pairing // ============================================================================ @@ -679,7 +850,7 @@ fn emit_message( /// For pairing mode, sends a pairing code DM if denied. fn check_sender_permission(user_id: &str, channel_id: &str, is_dm: bool) -> bool { // 1. Owner check (highest priority, applies to all contexts) - let owner_id = channel_host::workspace_read(OWNER_ID_PATH).filter(|s| !s.is_empty()); + let owner_id = host_workspace_read(OWNER_ID_PATH).filter(|s| !s.is_empty()); if let Some(ref owner) = owner_id { if user_id != owner { channel_host::log( @@ -699,15 +870,14 @@ fn check_sender_permission(user_id: &str, channel_id: &str, is_dm: bool) -> bool return true; // Channel messages bypass DM policy } - let dm_policy = - channel_host::workspace_read(DM_POLICY_PATH).unwrap_or_else(|| "pairing".to_string()); + let dm_policy = host_workspace_read(DM_POLICY_PATH).unwrap_or_else(|| "pairing".to_string()); if dm_policy == "open" { return true; } // 3. Build merged allow list: config allow_from + pairing store - let mut allowed: Vec = channel_host::workspace_read(ALLOW_FROM_PATH) + let mut allowed: Vec = host_workspace_read(ALLOW_FROM_PATH) .and_then(|s| serde_json::from_str(&s).ok()) .unwrap_or_default(); @@ -907,10 +1077,95 @@ fn json_response(status: u16, value: serde_json::Value) -> OutgoingHttpResponse // Export the component export!(SlackChannel); +#[cfg(test)] +mod test_host { + use super::*; + use std::cell::RefCell; + use std::collections::HashMap; + + #[derive(Debug, Clone, PartialEq, Eq)] + pub struct RecordedMessage { + pub user_id: String, + pub content: String, + pub thread_id: Option, + pub metadata_json: String, + } + + #[derive(Default)] + struct TestHostState { + workspace: HashMap, + emitted_messages: Vec, + now_millis: u64, + } + + std::thread_local! { + static STATE: RefCell = RefCell::new(TestHostState::default()); + } + + pub fn reset() { + STATE.with(|state| *state.borrow_mut() = TestHostState::default()); + } + + pub fn set_now_millis(now_millis: u64) { + STATE.with(|state| state.borrow_mut().now_millis = now_millis); + } + + pub fn now_millis() -> u64 { + STATE.with(|state| state.borrow().now_millis) + } + + pub fn workspace_read(path: &str) -> Option { + STATE.with(|state| state.borrow().workspace.get(path).cloned()) + } + + pub fn workspace_write(path: &str, content: &str) -> Result<(), String> { + STATE.with(|state| { + state + .borrow_mut() + .workspace + .insert(path.to_string(), content.to_string()); + }); + Ok(()) + } + + pub fn set_workspace(path: &str, content: &str) { + let _ = workspace_write(path, content); + } + + pub fn emit_message(message: &EmittedMessage) { + STATE.with(|state| { + state.borrow_mut().emitted_messages.push(RecordedMessage { + user_id: message.user_id.clone(), + content: message.content.clone(), + thread_id: message.thread_id.clone(), + metadata_json: message.metadata_json.clone(), + }); + }); + } + + pub fn take_emitted_messages() -> Vec { + STATE.with(|state| std::mem::take(&mut state.borrow_mut().emitted_messages)) + } +} + #[cfg(test)] mod tests { use super::*; + fn sample_thread_message_event(thread_ts: &str) -> SlackEvent { + SlackEvent { + event_type: "message".to_string(), + user: Some("U123".to_string()), + channel: Some("C123".to_string()), + text: Some("follow up".to_string()), + thread_ts: Some(thread_ts.to_string()), + ts: Some("1710000000.000002".to_string()), + bot_id: None, + subtype: None, + files: None, + } + } + #[test] fn test_extract_slack_attachments_with_files() { let files = Some(vec![ @@ -1016,10 +1271,131 @@ mod tests { #[test] fn test_max_download_size_constant() { - // Verify the constant is 20 MB assert_eq!(MAX_DOWNLOAD_SIZE_BYTES, 20 * 1024 * 1024); } + #[test] + fn test_active_slack_threads_round_trip() { + let now_ms = 1_710_000_000_000_u64; + let raw = format!( + r#"[{{"team_id":null,"channel":"G2","thread_ts":"678.90","last_seen_ms":{now_ms}}},{{"team_id":"T1","channel":"C1","thread_ts":"123.45","last_seen_ms":{now_ms}}}]"# + ); + let threads = parse_active_slack_threads(Some(&raw), now_ms); + assert_eq!( + threads.get(&active_slack_thread_key(Some("T1"), "C1", "123.45")), + Some(&now_ms) + ); + assert_eq!( + threads.get(&active_slack_thread_key(None, "G2", "678.90")), + Some(&now_ms) + ); + assert!(active_slack_thread_is_known( + Some(&raw), + Some("T1"), + "C1", + "123.45", + now_ms, + )); + assert!(!active_slack_thread_is_known( + Some(&raw), + Some("T1"), + "C2", + "123.45", + now_ms, + )); + assert_eq!(serialize_active_slack_threads(&threads), raw); + } + + #[test] + fn test_active_slack_threads_accept_legacy_timestamps() { + let now_ms = 1_710_000_000_000_u64; + let raw = r#"["123.45","678.90"]"#; + let threads = parse_active_slack_threads(Some(raw), now_ms); + assert_eq!( + threads.get(&active_slack_thread_key(None, "", "123.45")), + Some(&now_ms) + ); + assert!(active_slack_thread_is_known( + Some(raw), + Some("T1"), + "C1", + "123.45", + now_ms, + )); + assert!(!active_slack_thread_is_known( + Some(raw), + Some("T1"), + "C1", + "999.99", + now_ms, + )); + } + + #[test] + fn test_active_slack_threads_prune_expired_and_cap_entries() { + let now_ms = ACTIVE_THREAD_TTL_MS + 10_000; + let mut threads = HashMap::new(); + threads.insert(active_slack_thread_key(Some("T1"), "C1", "expired"), 1); + for idx in 0..(ACTIVE_THREAD_MAX_ENTRIES + 10) { + threads.insert( + active_slack_thread_key(Some("T1"), "C1", &format!("live-{idx}")), + now_ms.saturating_add(idx as u64), + ); + } + + prune_active_slack_threads(&mut threads, now_ms); + assert_eq!(threads.len(), ACTIVE_THREAD_MAX_ENTRIES); + assert!(!threads.contains_key(&active_slack_thread_key(Some("T1"), "C1", "expired"))); + assert!(!threads.contains_key(&active_slack_thread_key(Some("T1"), "C1", "live-0"))); + } + + #[test] + fn test_active_slack_threads_ignore_invalid_json() { + assert!(parse_active_slack_threads(Some("not-json"), 123).is_empty()); + assert!(parse_active_slack_threads(None, 123).is_empty()); + } + + #[test] + fn test_handle_slack_event_emits_for_known_active_thread() { + test_host::reset(); + test_host::set_now_millis(1_710_000_000_000_u64); + + let threads = HashMap::from([( + active_slack_thread_key(Some("T1"), "C123", "1710000000.000001"), + 1_710_000_000_000_u64, + )]); + test_host::set_workspace(ACTIVE_THREADS_PATH, &serialize_active_slack_threads(&threads)); + + handle_slack_event( + sample_thread_message_event("1710000000.000001"), + Some("T1".to_string()), + None, + ); + + let emitted = test_host::take_emitted_messages(); + assert_eq!(emitted.len(), 1); + assert_eq!(emitted[0].user_id, "U123"); + assert_eq!(emitted[0].content, "follow up"); + assert_eq!( + emitted[0].thread_id.as_deref(), + Some("1710000000.000001") + ); + } + + #[test] + fn test_handle_slack_event_skips_unknown_active_thread() { + test_host::reset(); + test_host::set_now_millis(1_710_000_000_000_u64); + + handle_slack_event( + sample_thread_message_event("1710000000.000001"), + Some("T1".to_string()), + None, + ); + + assert!(test_host::take_emitted_messages().is_empty()); + } + #[test] fn test_active_thread_key_scopes_by_channel_and_thread() { assert_eq!( diff --git a/src/channels/wasm/capabilities.rs b/src/channels/wasm/capabilities.rs index 2e3f47b3a7..512360d95b 100644 --- a/src/channels/wasm/capabilities.rs +++ b/src/channels/wasm/capabilities.rs @@ -41,6 +41,13 @@ pub struct ChannelCapabilities { /// Example: "channels/slack/" means writes to "state.json" become "channels/slack/state.json". pub workspace_prefix: String, + /// Workspace paths that are safe to persist across restarts. + /// + /// Paths are stored with the channel prefix already applied. + /// This allowlist exists so we do not accidentally persist secrets or + /// short-lived tokens that some channels keep in their callback workspace. + pub durable_workspace_paths: Vec, + /// Rate limiting for emit_message calls. pub emit_rate_limit: EmitRateLimitConfig, @@ -59,6 +66,7 @@ impl Default for ChannelCapabilities { allow_polling: false, min_poll_interval_ms: MIN_POLL_INTERVAL_MS, workspace_prefix: String::new(), + durable_workspace_paths: Vec::new(), emit_rate_limit: EmitRateLimitConfig::default(), max_message_size: 64 * 1024, // 64 KB callback_timeout: Duration::from_secs(30), @@ -88,6 +96,15 @@ impl ChannelCapabilities { self } + /// Set workspace paths that are safe to persist across restarts. + pub fn with_durable_workspace_paths(mut self, paths: Vec) -> Self { + self.durable_workspace_paths = paths + .into_iter() + .map(|path| self.prefix_workspace_path(&path)) + .collect(); + self + } + /// Set the emit rate limit. pub fn with_emit_rate_limit(mut self, rate_limit: EmitRateLimitConfig) -> Self { self.emit_rate_limit = rate_limit; @@ -133,6 +150,13 @@ impl ChannelCapabilities { } } + /// Returns true if the fully-prefixed workspace path is restart-durable. + pub fn is_durable_workspace_path(&self, full_path: &str) -> bool { + self.durable_workspace_paths + .iter() + .any(|path| path == full_path) + } + /// Check if a workspace path is valid for this channel. /// /// Paths cannot escape the channel's namespace. @@ -299,6 +323,19 @@ mod tests { assert!(result.is_err()); } + #[test] + fn test_durable_workspace_paths_are_prefixed() { + let caps = ChannelCapabilities::for_channel("slack") + .with_durable_workspace_paths(vec!["state/active_threads".to_string()]); + + assert_eq!( + caps.durable_workspace_paths, + vec!["channels/slack/state/active_threads".to_string()] + ); + assert!(caps.is_durable_workspace_path("channels/slack/state/active_threads")); + assert!(!caps.is_durable_workspace_path("channels/slack/state/owner_id")); + } + #[test] fn test_http_endpoint_config() { let endpoint = HttpEndpointConfig::post_webhook("/webhook/slack"); diff --git a/src/channels/wasm/host.rs b/src/channels/wasm/host.rs index e8a59a88c1..466cdf78db 100644 --- a/src/channels/wasm/host.rs +++ b/src/channels/wasm/host.rs @@ -540,6 +540,26 @@ impl ChannelWorkspaceStore { } } + /// Restore a previously-persisted subset of workspace values. + pub fn restore_snapshot(&self, snapshot: &std::collections::HashMap) { + if snapshot.is_empty() { + return; + } + if let Ok(mut data) = self.data.write() { + for (path, content) in snapshot { + data.insert(path.clone(), content.clone()); + } + } + } + + /// Take a point-in-time snapshot of the current workspace store. + pub fn snapshot(&self) -> std::collections::HashMap { + self.data + .read() + .map(|data| data.clone()) + .unwrap_or_default() + } + /// Append a text frame to a JSON queue stored at `path`. /// /// The queue is stored as a JSON array of strings and bounded to the most diff --git a/src/channels/wasm/schema.rs b/src/channels/wasm/schema.rs index 8ced07defc..d0a321f485 100644 --- a/src/channels/wasm/schema.rs +++ b/src/channels/wasm/schema.rs @@ -234,6 +234,14 @@ impl ChannelCapabilitiesSchema { caps.workspace_prefix = prefix.clone(); } + if !channel.durable_workspace_paths.is_empty() { + caps.durable_workspace_paths = channel + .durable_workspace_paths + .iter() + .map(|path| caps.prefix_workspace_path(path)) + .collect(); + } + if let Some(rate) = &channel.emit_rate_limit { caps.emit_rate_limit = rate.to_emit_rate_limit(); } @@ -270,6 +278,10 @@ pub struct ChannelSpecificCapabilitiesSchema { #[serde(default)] pub workspace_prefix: Option, + /// Workspace paths that are safe to persist across restarts. + #[serde(default)] + pub durable_workspace_paths: Vec, + /// Rate limiting for emit_message. #[serde(default)] pub emit_rate_limit: Option, @@ -590,6 +602,27 @@ mod tests { assert_eq!(caps.workspace_prefix, "integrations/custom/"); } + #[test] + fn test_durable_workspace_paths_are_prefixed() { + let json = r#"{ + "name": "slack", + "capabilities": { + "channel": { + "workspace_prefix": "channels/slack/", + "durable_workspace_paths": ["state/active_threads"] + } + } + }"#; + + let file = ChannelCapabilitiesFile::from_json(json).unwrap(); + let caps = file.to_capabilities(); + + assert_eq!( + caps.durable_workspace_paths, + vec!["channels/slack/state/active_threads".to_string()] + ); + } + #[test] fn test_emit_rate_limit() { let json = r#"{ diff --git a/src/channels/wasm/wrapper.rs b/src/channels/wasm/wrapper.rs index fd0103f86c..df532ab246 100644 --- a/src/channels/wasm/wrapper.rs +++ b/src/channels/wasm/wrapper.rs @@ -791,6 +791,12 @@ pub struct WasmChannel { /// Ensures WASM channels can maintain state (e.g., polling offsets) between ticks. workspace_store: Arc, + /// Serializes callback execution for a single channel instance. + /// + /// Some channel state is read-modify-written through the shared workspace + /// store, so overlapping callbacks can otherwise lose updates. + callback_lock: Arc>, + /// Last-seen message metadata (contains chat_id for broadcast routing). /// Populated from incoming messages so `broadcast()` knows where to send. last_broadcast_metadata: Arc>>, @@ -839,6 +845,117 @@ async fn do_update_broadcast_metadata( } } +fn durable_workspace_settings_key(channel_name: &str) -> String { + format!("channels.wasm_workspace.{}", channel_name) +} + +async fn do_persist_durable_workspace( + channel_name: &str, + owner_scope_id: &str, + workspace_store: &ChannelWorkspaceStore, + durable_paths: &[String], + settings_store: Option<&Arc>, +) { + if durable_paths.is_empty() { + return; + } + + let Some(store) = settings_store else { + return; + }; + + let snapshot = workspace_store.snapshot(); + let durable_snapshot: HashMap = durable_paths + .iter() + .filter_map(|path| { + snapshot + .get(path) + .cloned() + .map(|value| (path.clone(), value)) + }) + .collect(); + let key = durable_workspace_settings_key(channel_name); + + let result = if durable_snapshot.is_empty() { + store.delete_setting(owner_scope_id, &key).await.map(|_| ()) + } else { + store + .set_setting(owner_scope_id, &key, &serde_json::json!(durable_snapshot)) + .await + }; + + if let Err(e) = result { + tracing::warn!( + channel = %channel_name, + "Failed to persist durable workspace state: {}", + e + ); + } +} + +async fn do_load_durable_workspace( + channel_name: &str, + owner_scope_id: &str, + workspace_store: &ChannelWorkspaceStore, + durable_paths: &[String], + settings_store: Option<&Arc>, +) { + if durable_paths.is_empty() { + return; + } + + let Some(store) = settings_store else { + return; + }; + + let key = durable_workspace_settings_key(channel_name); + let load_value = match store.get_setting(owner_scope_id, &key).await { + Ok(value) => value, + Err(e) => { + tracing::warn!( + channel = %channel_name, + "Failed to load durable workspace state: {}", + e + ); + None + } + }; + + let load_value = if load_value.is_none() && owner_scope_id != "default" { + match store.get_setting("default", &key).await { + Ok(value) => value, + Err(e) => { + tracing::warn!( + channel = %channel_name, + "Failed to load legacy durable workspace state: {}", + e + ); + None + } + } + } else { + load_value + }; + + let Some(value) = load_value else { + return; + }; + + let Ok(snapshot) = serde_json::from_value::>(value) else { + tracing::warn!( + channel = %channel_name, + "Ignoring invalid durable workspace snapshot" + ); + return; + }; + + let filtered: HashMap = snapshot + .into_iter() + .filter(|(path, _)| durable_paths.iter().any(|durable| durable == path)) + .collect(); + workspace_store.restore_snapshot(&filtered); +} + fn resolve_message_scope( owner_scope_id: &str, owner_actor_id: Option<&str>, @@ -970,6 +1087,7 @@ impl WasmChannel { typing_task: RwLock::new(None), pairing_store, workspace_store: Arc::new(ChannelWorkspaceStore::new()), + callback_lock: Arc::new(tokio::sync::Mutex::new(())), last_broadcast_metadata: Arc::new(tokio::sync::RwLock::new(None)), settings_store, owner_scope_id: owner_scope_id.into(), @@ -1139,6 +1257,43 @@ impl WasmChannel { format!("channel_broadcast_metadata_{}", self.name) } + fn durable_workspace_paths(&self) -> &[String] { + &self.capabilities.durable_workspace_paths + } + + async fn load_durable_workspace_snapshot(&self) { + do_load_durable_workspace( + &self.name, + &self.owner_scope_id, + &self.workspace_store, + self.durable_workspace_paths(), + self.settings_store.as_ref(), + ) + .await; + } + + async fn persist_durable_workspace_snapshot_if_needed(&self, committed_paths: &[String]) { + if committed_paths.is_empty() { + return; + } + + if !committed_paths + .iter() + .any(|path| self.capabilities.is_durable_workspace_path(path)) + { + return; + } + + do_persist_durable_workspace( + &self.name, + &self.owner_scope_id, + &self.workspace_store, + self.durable_workspace_paths(), + self.settings_store.as_ref(), + ) + .await; + } + /// Update broadcast metadata in memory and persist if changed (best-effort). /// /// Compares with the current value to avoid redundant DB writes on every @@ -1257,6 +1412,7 @@ impl WasmChannel { let pairing_store = self.pairing_store.clone(); let callback_timeout = self.runtime.config().callback_timeout; let workspace_store = self.workspace_store.clone(); + let callback_lock = self.callback_lock.clone(); let last_broadcast_metadata = self.last_broadcast_metadata.clone(); let settings_store = self.settings_store.clone(); let owner_scope_id = self.owner_scope_id.clone(); @@ -1399,6 +1555,7 @@ impl WasmChannel { credentials: Arc::clone(&credentials), pairing_store: pairing_store.clone(), workspace_store: workspace_store.clone(), + callback_lock: Arc::clone(&callback_lock), message_tx: message_tx.clone(), rate_limiter: Arc::clone(&rate_limiter), last_broadcast_metadata: Arc::clone(&last_broadcast_metadata), @@ -1580,6 +1737,19 @@ impl WasmChannel { ) } + fn commit_callback_workspace_writes( + host_state: &mut ChannelHostState, + workspace_store: &ChannelWorkspaceStore, + ) -> Vec { + let pending_writes = host_state.take_pending_writes(); + let committed_paths = pending_writes + .iter() + .map(|write| write.path.clone()) + .collect(); + workspace_store.commit_writes(&pending_writes); + committed_paths + } + fn log_on_start_host_state(&self, host_state: &mut ChannelHostState) { for entry in host_state.take_logs() { match entry.level { @@ -1599,6 +1769,9 @@ impl WasmChannel { async fn execute_on_start_with_state( &self, ) -> Result<(Result, ChannelHostState), WasmChannelError> { + let _callback_guard = self.callback_lock.lock().await; + self.load_durable_workspace_snapshot().await; + let runtime = Arc::clone(&self.runtime); let prepared = Arc::clone(&self.prepared); let capabilities = Self::inject_workspace_reader(&self.capabilities, &self.workspace_store); @@ -1615,48 +1788,54 @@ impl WasmChannel { let pairing_store = self.pairing_store.clone(); let workspace_store = self.workspace_store.clone(); - tokio::time::timeout(timeout, async move { - tokio::task::spawn_blocking(move || { - let mut store = Self::create_store( - &runtime, - &prepared, - &capabilities, - credentials, - host_credentials, - pairing_store, - )?; - let instance = Self::instantiate_component(&runtime, &prepared, &mut store)?; + let (config_result, host_state, committed_paths) = + tokio::time::timeout(timeout, async move { + tokio::task::spawn_blocking(move || { + let mut store = Self::create_store( + &runtime, + &prepared, + &capabilities, + credentials, + host_credentials, + pairing_store, + )?; + let instance = Self::instantiate_component(&runtime, &prepared, &mut store)?; - let channel_iface = instance.near_agent_channel(); - let config_result = channel_iface - .call_on_start(&mut store, &config_json) - .map_err(|e| Self::map_wasm_error(e, &prepared.name, prepared.limits.fuel)) - .and_then(|wasm_result| match wasm_result { - Ok(wit_config) => Ok(convert_channel_config(wit_config)), - Err(err_msg) => Err(WasmChannelError::CallbackFailed { - name: prepared.name.clone(), - reason: err_msg, - }), - }); + let channel_iface = instance.near_agent_channel(); + let config_result = channel_iface + .call_on_start(&mut store, &config_json) + .map_err(|e| Self::map_wasm_error(e, &prepared.name, prepared.limits.fuel)) + .and_then(|wasm_result| match wasm_result { + Ok(wit_config) => Ok(convert_channel_config(wit_config)), + Err(err_msg) => Err(WasmChannelError::CallbackFailed { + name: prepared.name.clone(), + reason: err_msg, + }), + }); - let mut host_state = - Self::extract_host_state(&mut store, &prepared.name, &capabilities); - let pending_writes = host_state.take_pending_writes(); - workspace_store.commit_writes(&pending_writes); + let mut host_state = + Self::extract_host_state(&mut store, &prepared.name, &capabilities); + let committed_paths = + Self::commit_callback_workspace_writes(&mut host_state, &workspace_store); - Ok::<_, WasmChannelError>((config_result, host_state)) + Ok::<_, WasmChannelError>((config_result, host_state, committed_paths)) + }) + .await + .map_err(|e| WasmChannelError::ExecutionPanicked { + name: channel_name.clone(), + reason: e.to_string(), + })? }) .await - .map_err(|e| WasmChannelError::ExecutionPanicked { - name: channel_name.clone(), - reason: e.to_string(), - })? - }) - .await - .map_err(|_| WasmChannelError::Timeout { - name: self.name.clone(), - callback: "on_start".to_string(), - })? + .map_err(|_| WasmChannelError::Timeout { + name: self.name.clone(), + callback: "on_start".to_string(), + })??; + + self.persist_durable_workspace_snapshot_if_needed(&committed_paths) + .await; + + Ok((config_result, host_state)) } /// Execute the on_start callback. @@ -1706,6 +1885,8 @@ impl WasmChannel { body: &[u8], secret_validated: bool, ) -> Result { + let _callback_guard = self.callback_lock.lock().await; + tracing::info!( channel = %self.name, method = method, @@ -1801,10 +1982,10 @@ impl WasmChannel { Self::extract_host_state(&mut store, &prepared.name, &capabilities); // Commit pending workspace writes to the persistent store - let pending_writes = host_state.take_pending_writes(); - workspace_store.commit_writes(&pending_writes); + let committed_paths = + Self::commit_callback_workspace_writes(&mut host_state, &workspace_store); - Ok((response, host_state)) + Ok((response, host_state, committed_paths)) }) .await .map_err(|e| WasmChannelError::ExecutionPanicked { @@ -1816,7 +1997,9 @@ impl WasmChannel { let channel_name = self.name.clone(); match result { - Ok(Ok((response, mut host_state))) => { + Ok(Ok((response, mut host_state, committed_paths))) => { + self.persist_durable_workspace_snapshot_if_needed(&committed_paths) + .await; // Process emitted messages let emitted = host_state.take_emitted_messages(); self.process_emitted_messages(emitted).await?; @@ -1840,6 +2023,8 @@ impl WasmChannel { /// /// Called periodically if polling is configured. pub async fn call_on_poll(&self) -> Result<(), WasmChannelError> { + let _callback_guard = self.callback_lock.lock().await; + // If no WASM bytes, do nothing (for testing) if self.prepared.component().is_none() { tracing::debug!( @@ -1887,10 +2072,10 @@ impl WasmChannel { Self::extract_host_state(&mut store, &prepared.name, &capabilities); // Commit pending workspace writes to the persistent store - let pending_writes = host_state.take_pending_writes(); - workspace_store.commit_writes(&pending_writes); + let committed_paths = + Self::commit_callback_workspace_writes(&mut host_state, &workspace_store); - Ok(((), host_state)) + Ok(((), host_state, committed_paths)) }) .await .map_err(|e| WasmChannelError::ExecutionPanicked { @@ -1902,7 +2087,9 @@ impl WasmChannel { let channel_name = self.name.clone(); match result { - Ok(Ok(((), mut host_state))) => { + Ok(Ok(((), mut host_state, committed_paths))) => { + self.persist_durable_workspace_snapshot_if_needed(&committed_paths) + .await; let _ = drain_guest_logs(&channel_name, "on_poll", &mut host_state); // Process emitted messages @@ -1934,6 +2121,8 @@ impl WasmChannel { metadata_json: &str, attachments: &[String], ) -> Result<(), WasmChannelError> { + let _callback_guard = self.callback_lock.lock().await; + tracing::info!( channel = %self.name, message_id = %message_id, @@ -2047,10 +2236,10 @@ impl WasmChannel { let mut host_state = Self::extract_host_state(&mut store, &prepared.name, &capabilities); - let pending_writes = host_state.take_pending_writes(); - workspace_store.commit_writes(&pending_writes); + let committed_paths = + Self::commit_callback_workspace_writes(&mut host_state, &workspace_store); tracing::info!("on_respond WASM execution completed successfully"); - Ok(((), host_state)) + Ok(((), host_state, committed_paths)) }) .await .map_err(|e| { @@ -2065,7 +2254,9 @@ impl WasmChannel { let channel_name = self.name.clone(); match result { - Ok(Ok(((), _host_state))) => { + Ok(Ok(((), _host_state, committed_paths))) => { + self.persist_durable_workspace_snapshot_if_needed(&committed_paths) + .await; tracing::debug!( channel = %channel_name, message_id = %message_id, @@ -2091,6 +2282,8 @@ impl WasmChannel { thread_id: Option<&str>, attachments: &[String], ) -> Result<(), WasmChannelError> { + let _callback_guard = self.callback_lock.lock().await; + tracing::info!( channel = %self.name, user_id = %user_id, @@ -2175,10 +2368,10 @@ impl WasmChannel { let mut host_state = Self::extract_host_state(&mut store, &prepared.name, &capabilities); - let pending_writes = host_state.take_pending_writes(); - workspace_store.commit_writes(&pending_writes); + let committed_paths = + Self::commit_callback_workspace_writes(&mut host_state, &workspace_store); tracing::info!("on_broadcast WASM execution completed successfully"); - Ok(((), host_state)) + Ok(((), host_state, committed_paths)) }) .await .map_err(|e| WasmChannelError::ExecutionPanicked { @@ -2190,7 +2383,9 @@ impl WasmChannel { let channel_name = self.name.clone(); match result { - Ok(Ok(((), _host_state))) => { + Ok(Ok(((), _host_state, committed_paths))) => { + self.persist_durable_workspace_snapshot_if_needed(&committed_paths) + .await; tracing::debug!( channel = %channel_name, "WASM channel on_broadcast completed" @@ -2213,6 +2408,8 @@ impl WasmChannel { status: &StatusUpdate, metadata: &serde_json::Value, ) -> Result<(), WasmChannelError> { + let _callback_guard = self.callback_lock.lock().await; + // If no WASM bytes, do nothing (for testing) if self.prepared.component().is_none() { return Ok(()); @@ -2256,10 +2453,10 @@ impl WasmChannel { let mut host_state = Self::extract_host_state(&mut store, &prepared.name, &capabilities); - let pending_writes = host_state.take_pending_writes(); - workspace_store.commit_writes(&pending_writes); + let committed_paths = + Self::commit_callback_workspace_writes(&mut host_state, &workspace_store); - Ok(()) + Ok(committed_paths) }) .await .map_err(|e| WasmChannelError::ExecutionPanicked { @@ -2270,7 +2467,9 @@ impl WasmChannel { .await; match result { - Ok(Ok(())) => { + Ok(Ok(committed_paths)) => { + self.persist_durable_workspace_snapshot_if_needed(&committed_paths) + .await; tracing::debug!( channel = %self.name, "WASM channel on_status completed" @@ -2292,26 +2491,34 @@ impl WasmChannel { #[allow(clippy::too_many_arguments)] async fn execute_status( channel_name: &str, + owner_scope_id: &str, runtime: &Arc, prepared: &Arc, capabilities: &ChannelCapabilities, credentials: &RwLock>, host_credentials: Vec, pairing_store: Arc, - timeout: Duration, workspace_store: &Arc, + callback_lock: &Arc>, + settings_store: Option<&Arc>, + timeout: Duration, wit_update: wit_channel::StatusUpdate, ) -> Result<(), WasmChannelError> { if prepared.component().is_none() { return Ok(()); } + let _callback_guard = callback_lock.lock().await; + let runtime = Arc::clone(runtime); let prepared = Arc::clone(prepared); let capabilities = Self::inject_workspace_reader(capabilities, workspace_store); + let durable_workspace_paths = capabilities.durable_workspace_paths.clone(); let credentials_snapshot = credentials.read().await.clone(); let channel_name_owned = channel_name.to_string(); + let owner_scope_id_owned = owner_scope_id.to_string(); let workspace_store = Arc::clone(workspace_store); + let workspace_store_for_callback = Arc::clone(&workspace_store); let result = tokio::time::timeout(timeout, async move { tokio::task::spawn_blocking(move || { @@ -2332,10 +2539,12 @@ impl WasmChannel { let mut host_state = Self::extract_host_state(&mut store, &prepared.name, &capabilities); - let pending_writes = host_state.take_pending_writes(); - workspace_store.commit_writes(&pending_writes); + let committed_paths = Self::commit_callback_workspace_writes( + &mut host_state, + &workspace_store_for_callback, + ); - Ok(()) + Ok(committed_paths) }) .await .map_err(|e| WasmChannelError::ExecutionPanicked { @@ -2346,7 +2555,23 @@ impl WasmChannel { .await; match result { - Ok(Ok(())) => Ok(()), + Ok(Ok(committed_paths)) => { + if committed_paths.iter().any(|path| { + durable_workspace_paths + .iter() + .any(|durable| durable == path) + }) { + do_persist_durable_workspace( + channel_name, + &owner_scope_id_owned, + &workspace_store, + &durable_workspace_paths, + settings_store, + ) + .await; + } + Ok(()) + } Ok(Err(e)) => Err(e), Err(_) => Err(WasmChannelError::Timeout { name: channel_name.to_string(), @@ -2403,10 +2628,10 @@ impl WasmChannel { // Spawn background repeater let channel_name = self.name.clone(); + let owner_scope_id = self.owner_scope_id.clone(); let runtime = Arc::clone(&self.runtime); let prepared = Arc::clone(&self.prepared); let capabilities = self.capabilities.clone(); - let workspace_store = self.workspace_store.clone(); let credentials = self.credentials.clone(); // Pre-resolve host credentials once for the lifetime of the repeater. // Channels tokens rarely change, so a snapshot per-repeater is correct. @@ -2417,6 +2642,9 @@ impl WasmChannel { ) .await; let pairing_store = self.pairing_store.clone(); + let workspace_store = self.workspace_store.clone(); + let callback_lock = self.callback_lock.clone(); + let settings_store = self.settings_store.clone(); let callback_timeout = self.runtime.config().callback_timeout; let Some(wit_update) = status_to_wit(&status, metadata) else { return Ok(()); @@ -2435,14 +2663,17 @@ impl WasmChannel { if let Err(e) = Self::execute_status( &channel_name, + &owner_scope_id, &runtime, &prepared, &capabilities, &credentials, hc, pairing_store.clone(), - callback_timeout, &workspace_store, + &callback_lock, + settings_store.as_ref(), + callback_timeout, wit_update_clone, ) .await @@ -2709,6 +2940,7 @@ impl WasmChannel { let pairing_store = self.pairing_store.clone(); let callback_timeout = self.runtime.config().callback_timeout; let workspace_store = self.workspace_store.clone(); + let callback_lock = self.callback_lock.clone(); let last_broadcast_metadata = self.last_broadcast_metadata.clone(); let settings_store = self.settings_store.clone(); let poll_secrets_store = self.secrets_store.clone(); @@ -2737,14 +2969,17 @@ impl WasmChannel { // Execute on_poll with fresh WASM instance let result = Self::execute_poll( &channel_name, + &owner_scope_id, &runtime, &prepared, &capabilities, &credentials, host_credentials, pairing_store.clone(), - callback_timeout, &workspace_store, + &callback_lock, + settings_store.as_ref(), + callback_timeout, ).await; match result { @@ -2817,14 +3052,17 @@ impl WasmChannel { #[allow(clippy::too_many_arguments)] async fn execute_poll( channel_name: &str, + owner_scope_id: &str, runtime: &Arc, prepared: &Arc, capabilities: &ChannelCapabilities, credentials: &RwLock>, host_credentials: Vec, pairing_store: Arc, - timeout: Duration, workspace_store: &Arc, + callback_lock: &Arc>, + settings_store: Option<&Arc>, + timeout: Duration, ) -> Result, WasmChannelError> { // Skip if no WASM bytes (testing mode) if prepared.component().is_none() { @@ -2835,12 +3073,17 @@ impl WasmChannel { return Ok(Vec::new()); } + let _callback_guard = callback_lock.lock().await; + let runtime = Arc::clone(runtime); let prepared = Arc::clone(prepared); let capabilities = Self::inject_workspace_reader(capabilities, workspace_store); + let durable_workspace_paths = capabilities.durable_workspace_paths.clone(); let credentials_snapshot = credentials.read().await.clone(); let channel_name_owned = channel_name.to_string(); + let owner_scope_id_owned = owner_scope_id.to_string(); let workspace_store = Arc::clone(workspace_store); + let workspace_store_for_callback = Arc::clone(&workspace_store); // Execute in blocking task with timeout let result = tokio::time::timeout(timeout, async move { @@ -2865,10 +3108,12 @@ impl WasmChannel { Self::extract_host_state(&mut store, &prepared.name, &capabilities); // Commit pending workspace writes to the persistent store - let pending_writes = host_state.take_pending_writes(); - workspace_store.commit_writes(&pending_writes); + let committed_paths = Self::commit_callback_workspace_writes( + &mut host_state, + &workspace_store_for_callback, + ); - Ok(host_state) + Ok((host_state, committed_paths)) }) .await .map_err(|e| WasmChannelError::ExecutionPanicked { @@ -2879,7 +3124,21 @@ impl WasmChannel { .await; match result { - Ok(Ok(mut host_state)) => { + Ok(Ok((mut host_state, committed_paths))) => { + if committed_paths.iter().any(|path| { + durable_workspace_paths + .iter() + .any(|durable| durable == path) + }) { + do_persist_durable_workspace( + channel_name, + &owner_scope_id_owned, + &workspace_store, + &durable_workspace_paths, + settings_store, + ) + .await; + } let _ = drain_guest_logs(channel_name, "on_poll", &mut host_state); let emitted = host_state.take_emitted_messages(); tracing::debug!( @@ -3540,6 +3799,7 @@ struct WebsocketPollContext { credentials: Arc>>, pairing_store: Arc, workspace_store: Arc, + callback_lock: Arc>, message_tx: Arc>>>, rate_limiter: Arc>, last_broadcast_metadata: Arc>>, @@ -3586,14 +3846,17 @@ fn spawn_websocket_poll(poll_guard: tokio::sync::OwnedMutexGuard<()>, ctx: Webso match WasmChannel::execute_poll( &ctx.channel_name, + &ctx.owner_scope_id, &ctx.runtime, &ctx.prepared, &ctx.capabilities, &ctx.credentials, host_credentials, ctx.pairing_store.clone(), - ctx.callback_timeout, &ctx.workspace_store, + &ctx.callback_lock, + ctx.settings_store.as_ref(), + ctx.callback_timeout, ) .await { @@ -4586,6 +4849,34 @@ mod tests { ) } + #[cfg(feature = "libsql")] + fn create_test_slack_channel_with_settings_store( + settings_store: Arc, + owner_scope_id: &str, + ) -> WasmChannel { + let config = WasmChannelRuntimeConfig::for_testing(); + let runtime = Arc::new(WasmChannelRuntime::new(config).unwrap()); + let prepared = Arc::new(PreparedChannelModule { + name: "slack".to_string(), + description: "Slack test channel".to_string(), + component: None, + limits: ResourceLimits::default(), + }); + let capabilities = ChannelCapabilities::for_channel("slack") + .with_path("/webhook/slack") + .with_durable_workspace_paths(vec!["state/active_threads".to_string()]); + + WasmChannel::new( + runtime, + prepared, + capabilities, + owner_scope_id, + "{}".to_string(), + Arc::new(PairingStore::new_noop()), + Some(settings_store), + ) + } + struct RecordingSettingsStore { values: tokio::sync::RwLock>>, lookups: tokio::sync::RwLock>, @@ -5070,17 +5361,21 @@ mod tests { let timeout = std::time::Duration::from_secs(5); let workspace_store = Arc::new(crate::channels::wasm::host::ChannelWorkspaceStore::new()); + let callback_lock = Arc::new(tokio::sync::Mutex::new(())); let result = WasmChannel::execute_poll( "poll-test", + "default", &runtime, &prepared, &capabilities, &credentials, Vec::new(), // no host credentials in test Arc::new(PairingStore::new_noop()), - timeout, &workspace_store, + &callback_lock, + None, + timeout, ) .await; @@ -5088,6 +5383,105 @@ mod tests { assert!(result.unwrap().is_empty()); } + #[test] + fn test_on_respond_workspace_write_commit_survives_later_callback() { + use crate::channels::wasm::host::{ChannelHostState, ChannelWorkspaceStore}; + use crate::tools::wasm::{WorkspaceCapability, WorkspaceReader}; + + let workspace_store = ChannelWorkspaceStore::new(); + let mut respond_state = + ChannelHostState::new("slack", ChannelCapabilities::for_channel("slack")); + respond_state + .workspace_write( + "state/active_threads", + r#"[{"team_id":"T1","channel":"C1","thread_ts":"1710000000.000001"}]"#.to_string(), + ) + .expect("on_respond write should be accepted"); + + WasmChannel::commit_callback_workspace_writes(&mut respond_state, &workspace_store); + assert_eq!(respond_state.pending_writes_count(), 0); + + let workspace_store = Arc::new(workspace_store); + let mut later_caps = ChannelCapabilities::for_channel("slack"); + later_caps.tool_capabilities.workspace_read = Some(WorkspaceCapability { + allowed_prefixes: vec![], + reader: Some(Arc::clone(&workspace_store) as Arc), + }); + let later_state = ChannelHostState::new("slack", later_caps); + + assert_eq!( + later_state + .workspace_read("state/active_threads") + .expect("later callback read should not fail"), + Some( + r#"[{"team_id":"T1","channel":"C1","thread_ts":"1710000000.000001"}]"#.to_string() + ) + ); + } + + #[test] + fn test_inject_workspace_reader_supports_broadcast_callback_state_reads() { + use crate::channels::wasm::host::{ + ChannelHostState, ChannelWorkspaceStore, PendingWorkspaceWrite, + }; + + let workspace_store = Arc::new(ChannelWorkspaceStore::new()); + workspace_store.commit_writes(&[PendingWorkspaceWrite { + path: "channels/feishu/state/api_base".to_string(), + content: "https://open.feishu.cn".to_string(), + }]); + + let callback_caps = WasmChannel::inject_workspace_reader( + &ChannelCapabilities::for_channel("feishu"), + &workspace_store, + ); + let callback_state = ChannelHostState::new("feishu", callback_caps); + + assert_eq!( + callback_state + .workspace_read("state/api_base") + .expect("callback read should not fail"), + Some("https://open.feishu.cn".to_string()) + ); + } + + #[cfg(feature = "libsql")] + #[tokio::test] + async fn test_durable_workspace_paths_restore_across_channel_restart() { + use crate::channels::wasm::host::PendingWorkspaceWrite; + use crate::testing::test_db; + + let (db, _temp_dir) = test_db().await; + let settings_store = Arc::clone(&db) as Arc; + let channel = create_test_slack_channel_with_settings_store( + Arc::clone(&settings_store), + "owner-scope", + ); + + channel.workspace_store.commit_writes(&[PendingWorkspaceWrite { + path: "channels/slack/state/active_threads".to_string(), + content: r#"[{"team_id":"T1","channel":"C1","thread_ts":"1710000000.000001","last_seen_ms":1710000000000}]"#.to_string(), + }]); + channel + .persist_durable_workspace_snapshot_if_needed(&[ + "channels/slack/state/active_threads".to_string() + ]) + .await; + + let restored = create_test_slack_channel_with_settings_store(settings_store, "owner-scope"); + restored.load_durable_workspace_snapshot().await; + + assert_eq!( + crate::tools::wasm::WorkspaceReader::read( + &*restored.workspace_store, + "channels/slack/state/active_threads", + ), + Some( + r#"[{"team_id":"T1","channel":"C1","thread_ts":"1710000000.000001","last_seen_ms":1710000000000}]"#.to_string() + ) + ); + } + #[tokio::test] async fn test_dispatch_emitted_messages_sends_to_channel() { use crate::channels::wasm::host::EmittedMessage;