mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-27 08:00:17 +00:00
1116 lines
37 KiB
Rust
1116 lines
37 KiB
Rust
//! Host state for WASM channel execution.
|
|
//!
|
|
//! Extends the base tool host state with channel-specific functionality:
|
|
//! - Message emission (queueing messages to send to the agent)
|
|
//! - Workspace write access (scoped to channel namespace)
|
|
//! - Rate limiting for message emission
|
|
|
|
use std::collections::HashMap;
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
|
|
|
use crate::channels::wasm::capabilities::{ChannelCapabilities, EmitRateLimitConfig};
|
|
use crate::channels::wasm::error::WasmChannelError;
|
|
use crate::tools::wasm::{HostState, LogLevel};
|
|
|
|
/// Maximum emitted messages per callback execution.
|
|
const MAX_EMITS_PER_EXECUTION: usize = 100;
|
|
|
|
/// Maximum message content size (64 KB).
|
|
const MAX_MESSAGE_CONTENT_SIZE: usize = 64 * 1024;
|
|
|
|
/// A file or media attachment on an incoming message.
|
|
#[derive(Debug, Clone)]
|
|
pub struct Attachment {
|
|
/// Unique identifier within the channel (e.g., Telegram file_id).
|
|
pub id: String,
|
|
/// MIME type (e.g., "image/jpeg", "audio/ogg", "application/pdf").
|
|
pub mime_type: String,
|
|
/// Original filename, if known.
|
|
pub filename: Option<String>,
|
|
/// File size in bytes, if known.
|
|
pub size_bytes: Option<u64>,
|
|
/// URL to download the file from the channel's API.
|
|
pub source_url: Option<String>,
|
|
/// Opaque key for host-side storage (e.g., after download/caching).
|
|
pub storage_key: Option<String>,
|
|
/// Extracted text content (e.g., OCR result, PDF text, audio transcript).
|
|
pub extracted_text: Option<String>,
|
|
/// Raw file bytes (for small files downloaded by the channel).
|
|
pub data: Vec<u8>,
|
|
/// Duration in seconds (for audio/video).
|
|
pub duration_secs: Option<u32>,
|
|
}
|
|
|
|
/// Maximum total attachment size per message (20 MB).
|
|
const MAX_ATTACHMENT_TOTAL_SIZE: u64 = 20 * 1024 * 1024;
|
|
|
|
/// Maximum number of attachments per message.
|
|
const MAX_ATTACHMENTS_PER_MESSAGE: usize = 10;
|
|
|
|
/// Allowed MIME type prefixes for attachments.
|
|
const ALLOWED_MIME_PREFIXES: &[&str] = &[
|
|
"image/",
|
|
"audio/",
|
|
"video/",
|
|
"application/pdf",
|
|
"application/vnd.",
|
|
"application/msword",
|
|
"application/rtf",
|
|
"text/",
|
|
"application/json",
|
|
"application/zip",
|
|
"application/gzip",
|
|
"application/x-tar",
|
|
"application/octet-stream",
|
|
];
|
|
/// Truncate a string to at most `max_bytes` without splitting UTF-8 code points.
|
|
fn truncate_utf8(s: &str, max_bytes: usize) -> &str {
|
|
let end = crate::util::floor_char_boundary(s, max_bytes);
|
|
&s[..end]
|
|
}
|
|
/// A message emitted by a WASM channel to be sent to the agent.
|
|
#[derive(Debug, Clone)]
|
|
pub struct EmittedMessage {
|
|
/// User identifier within the channel.
|
|
pub user_id: String,
|
|
|
|
/// Optional user display name.
|
|
pub user_name: Option<String>,
|
|
|
|
/// Message content.
|
|
pub content: String,
|
|
|
|
/// Optional thread ID for threaded conversations.
|
|
pub thread_id: Option<String>,
|
|
|
|
/// Channel-specific metadata as JSON string.
|
|
pub metadata_json: String,
|
|
|
|
/// File or media attachments on this message.
|
|
pub attachments: Vec<Attachment>,
|
|
|
|
/// Timestamp when the message was emitted.
|
|
pub emitted_at_millis: u64,
|
|
}
|
|
|
|
impl EmittedMessage {
|
|
/// Create a new emitted message.
|
|
pub fn new(user_id: impl Into<String>, content: impl Into<String>) -> Self {
|
|
Self {
|
|
user_id: user_id.into(),
|
|
user_name: None,
|
|
content: content.into(),
|
|
thread_id: None,
|
|
metadata_json: "{}".to_string(),
|
|
attachments: Vec::new(),
|
|
emitted_at_millis: SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.map(|d| d.as_millis() as u64)
|
|
.unwrap_or(0),
|
|
}
|
|
}
|
|
|
|
/// Set the user name.
|
|
pub fn with_user_name(mut self, name: impl Into<String>) -> Self {
|
|
self.user_name = Some(name.into());
|
|
self
|
|
}
|
|
|
|
/// Set the thread ID.
|
|
pub fn with_thread_id(mut self, thread_id: impl Into<String>) -> Self {
|
|
self.thread_id = Some(thread_id.into());
|
|
self
|
|
}
|
|
|
|
/// Set metadata JSON.
|
|
pub fn with_metadata(mut self, metadata_json: impl Into<String>) -> Self {
|
|
self.metadata_json = metadata_json.into();
|
|
self
|
|
}
|
|
|
|
/// Set attachments.
|
|
pub fn with_attachments(mut self, attachments: Vec<Attachment>) -> Self {
|
|
self.attachments = attachments;
|
|
self
|
|
}
|
|
}
|
|
|
|
/// A pending workspace write operation.
|
|
#[derive(Debug, Clone)]
|
|
pub struct PendingWorkspaceWrite {
|
|
/// Full path (already prefixed with channel namespace).
|
|
pub path: String,
|
|
|
|
/// Content to write.
|
|
pub content: String,
|
|
}
|
|
|
|
/// Host state for WASM channel callbacks.
|
|
///
|
|
/// Maintains all side effects during callback execution and enforces limits.
|
|
/// This is the channel-specific equivalent of HostState for tools.
|
|
pub struct ChannelHostState {
|
|
/// Base tool host state (logging, time, HTTP, etc.).
|
|
base: HostState,
|
|
|
|
/// Channel name (for error messages).
|
|
channel_name: String,
|
|
|
|
/// Channel capabilities.
|
|
capabilities: ChannelCapabilities,
|
|
|
|
/// Emitted messages (queued for delivery).
|
|
emitted_messages: Vec<EmittedMessage>,
|
|
|
|
/// Pending workspace writes.
|
|
pending_writes: Vec<PendingWorkspaceWrite>,
|
|
|
|
/// Emit count for rate limiting within this execution.
|
|
emit_count: u32,
|
|
|
|
/// Whether emit is still allowed (false after rate limit hit).
|
|
emit_enabled: bool,
|
|
|
|
/// Count of emits dropped due to rate limiting.
|
|
emits_dropped: usize,
|
|
|
|
/// Binary data stored for attachments via `store-attachment-data`.
|
|
/// Keyed by attachment ID, cleared after callback completes.
|
|
attachment_data: HashMap<String, Vec<u8>>,
|
|
|
|
/// Total bytes stored in attachment_data (for enforcing limits).
|
|
attachment_data_total: u64,
|
|
}
|
|
|
|
impl std::fmt::Debug for ChannelHostState {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("ChannelHostState")
|
|
.field("channel_name", &self.channel_name)
|
|
.field("emitted_messages_count", &self.emitted_messages.len())
|
|
.field("pending_writes_count", &self.pending_writes.len())
|
|
.field("emit_count", &self.emit_count)
|
|
.field("emit_enabled", &self.emit_enabled)
|
|
.field("emits_dropped", &self.emits_dropped)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
impl ChannelHostState {
|
|
/// Create a new channel host state.
|
|
pub fn new(channel_name: impl Into<String>, capabilities: ChannelCapabilities) -> Self {
|
|
let base = HostState::new(capabilities.tool_capabilities.clone());
|
|
|
|
Self {
|
|
base,
|
|
channel_name: channel_name.into(),
|
|
capabilities,
|
|
emitted_messages: Vec::new(),
|
|
pending_writes: Vec::new(),
|
|
emit_count: 0,
|
|
emit_enabled: true,
|
|
emits_dropped: 0,
|
|
attachment_data: HashMap::new(),
|
|
attachment_data_total: 0,
|
|
}
|
|
}
|
|
|
|
/// Get the channel name.
|
|
pub fn channel_name(&self) -> &str {
|
|
&self.channel_name
|
|
}
|
|
|
|
/// Get the capabilities.
|
|
pub fn capabilities(&self) -> &ChannelCapabilities {
|
|
&self.capabilities
|
|
}
|
|
|
|
/// Get the base host state for tool capabilities.
|
|
pub fn base(&self) -> &HostState {
|
|
&self.base
|
|
}
|
|
|
|
/// Get mutable access to the base host state.
|
|
pub fn base_mut(&mut self) -> &mut HostState {
|
|
&mut self.base
|
|
}
|
|
|
|
/// Emit a message from the channel.
|
|
///
|
|
/// Messages are queued and delivered after callback execution completes.
|
|
/// Rate limiting is enforced per-execution and globally.
|
|
/// Attachments are validated for count, total size, and MIME type.
|
|
pub fn emit_message(&mut self, msg: EmittedMessage) -> Result<(), WasmChannelError> {
|
|
// Check per-execution limit
|
|
if !self.emit_enabled {
|
|
self.emits_dropped += 1;
|
|
return Ok(()); // Silently drop, don't fail execution
|
|
}
|
|
|
|
if self.emitted_messages.len() >= MAX_EMITS_PER_EXECUTION {
|
|
self.emit_enabled = false;
|
|
self.emits_dropped += 1;
|
|
tracing::warn!(
|
|
channel = %self.channel_name,
|
|
limit = MAX_EMITS_PER_EXECUTION,
|
|
"Channel emit limit reached, further messages dropped"
|
|
);
|
|
return Ok(());
|
|
}
|
|
|
|
// Validate attachments
|
|
let msg = self.validate_attachments(msg);
|
|
|
|
// Validate message content size
|
|
if msg.content.len() > MAX_MESSAGE_CONTENT_SIZE {
|
|
tracing::warn!(
|
|
channel = %self.channel_name,
|
|
size = msg.content.len(),
|
|
max = MAX_MESSAGE_CONTENT_SIZE,
|
|
"Message content too large, truncating"
|
|
);
|
|
let mut truncated = truncate_utf8(&msg.content, MAX_MESSAGE_CONTENT_SIZE).to_string();
|
|
truncated.push_str("... (truncated)");
|
|
let msg = EmittedMessage {
|
|
content: truncated,
|
|
..msg
|
|
};
|
|
self.emitted_messages.push(msg);
|
|
} else {
|
|
self.emitted_messages.push(msg);
|
|
}
|
|
|
|
self.emit_count += 1;
|
|
Ok(())
|
|
}
|
|
|
|
/// Validate and sanitize attachments on an emitted message.
|
|
///
|
|
/// Enforces count limits, total size limits, and MIME type allowlist.
|
|
/// Invalid attachments are dropped with a warning.
|
|
fn validate_attachments(&self, mut msg: EmittedMessage) -> EmittedMessage {
|
|
if msg.attachments.is_empty() {
|
|
return msg;
|
|
}
|
|
|
|
// Enforce attachment count limit
|
|
if msg.attachments.len() > MAX_ATTACHMENTS_PER_MESSAGE {
|
|
tracing::warn!(
|
|
channel = %self.channel_name,
|
|
count = msg.attachments.len(),
|
|
max = MAX_ATTACHMENTS_PER_MESSAGE,
|
|
"Too many attachments, truncating"
|
|
);
|
|
msg.attachments.truncate(MAX_ATTACHMENTS_PER_MESSAGE);
|
|
}
|
|
|
|
// Filter by MIME type and enforce total size limit
|
|
let mut total_size: u64 = 0;
|
|
msg.attachments.retain(|att| {
|
|
let mime_ok = ALLOWED_MIME_PREFIXES
|
|
.iter()
|
|
.any(|prefix| att.mime_type.starts_with(prefix));
|
|
if !mime_ok {
|
|
tracing::warn!(
|
|
channel = %self.channel_name,
|
|
mime_type = %att.mime_type,
|
|
"Attachment MIME type not allowed, dropping"
|
|
);
|
|
return false;
|
|
}
|
|
|
|
// Use the larger of reported size_bytes and actual stored data size
|
|
// to prevent WASM channels from under-reporting to bypass limits.
|
|
let stored_size = self
|
|
.attachment_data
|
|
.get(&att.id)
|
|
.map(|d| d.len() as u64)
|
|
.unwrap_or(att.data.len() as u64);
|
|
let size = att
|
|
.size_bytes
|
|
.map(|reported| reported.max(stored_size))
|
|
.unwrap_or(stored_size);
|
|
if size > 0 {
|
|
total_size = total_size.saturating_add(size);
|
|
if total_size > MAX_ATTACHMENT_TOTAL_SIZE {
|
|
tracing::warn!(
|
|
channel = %self.channel_name,
|
|
total_size,
|
|
max = MAX_ATTACHMENT_TOTAL_SIZE,
|
|
"Attachment total size exceeded, dropping"
|
|
);
|
|
return false;
|
|
}
|
|
}
|
|
|
|
true
|
|
});
|
|
|
|
msg
|
|
}
|
|
|
|
/// Take all emitted messages (clears the queue).
|
|
pub fn take_emitted_messages(&mut self) -> Vec<EmittedMessage> {
|
|
std::mem::take(&mut self.emitted_messages)
|
|
}
|
|
|
|
/// Get the number of emitted messages.
|
|
pub fn emitted_count(&self) -> usize {
|
|
self.emitted_messages.len()
|
|
}
|
|
|
|
/// Get the number of emits dropped due to rate limiting.
|
|
pub fn emits_dropped(&self) -> usize {
|
|
self.emits_dropped
|
|
}
|
|
|
|
/// Store binary data for an attachment.
|
|
///
|
|
/// Called by WASM channels to associate downloaded bytes with an attachment ID.
|
|
/// The data is retrieved after callback completion and merged into `Attachment::data`.
|
|
pub fn store_attachment_data(
|
|
&mut self,
|
|
attachment_id: &str,
|
|
data: Vec<u8>,
|
|
) -> Result<(), WasmChannelError> {
|
|
const MAX_PER_ATTACHMENT: u64 = 20 * 1024 * 1024; // 20 MB
|
|
const MAX_TOTAL: u64 = 50 * 1024 * 1024; // 50 MB
|
|
|
|
let size = data.len() as u64;
|
|
if size > MAX_PER_ATTACHMENT {
|
|
return Err(WasmChannelError::CallbackFailed {
|
|
name: self.channel_name.clone(),
|
|
reason: format!(
|
|
"Attachment data too large: {} bytes (max {})",
|
|
size, MAX_PER_ATTACHMENT
|
|
),
|
|
});
|
|
}
|
|
|
|
// Subtract the old entry size (if overwriting) before adding new size
|
|
let old_size = self
|
|
.attachment_data
|
|
.get(attachment_id)
|
|
.map(|d| d.len() as u64)
|
|
.unwrap_or(0);
|
|
let adjusted_total = self.attachment_data_total.saturating_sub(old_size);
|
|
let new_total = adjusted_total.saturating_add(size);
|
|
if new_total > MAX_TOTAL {
|
|
return Err(WasmChannelError::CallbackFailed {
|
|
name: self.channel_name.clone(),
|
|
reason: format!(
|
|
"Total attachment data too large: {} bytes (max {})",
|
|
new_total, MAX_TOTAL
|
|
),
|
|
});
|
|
}
|
|
|
|
self.attachment_data_total = new_total;
|
|
self.attachment_data.insert(attachment_id.to_string(), data);
|
|
Ok(())
|
|
}
|
|
|
|
/// Remove stored binary data for a specific attachment ID.
|
|
pub fn remove_attachment_data(&mut self, id: &str) -> Option<Vec<u8>> {
|
|
if let Some(data) = self.attachment_data.remove(id) {
|
|
self.attachment_data_total =
|
|
self.attachment_data_total.saturating_sub(data.len() as u64);
|
|
Some(data)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Take all stored attachment data (clears the store).
|
|
pub fn take_attachment_data(&mut self) -> HashMap<String, Vec<u8>> {
|
|
self.attachment_data_total = 0;
|
|
std::mem::take(&mut self.attachment_data)
|
|
}
|
|
|
|
/// Write to workspace (scoped to channel namespace).
|
|
///
|
|
/// Writes are queued and committed after callback execution completes.
|
|
pub fn workspace_write(&mut self, path: &str, content: String) -> Result<(), WasmChannelError> {
|
|
// Validate and prefix path
|
|
let full_path = self
|
|
.capabilities
|
|
.validate_workspace_path(path)
|
|
.map_err(|reason| WasmChannelError::WorkspaceEscape {
|
|
name: self.channel_name.clone(),
|
|
path: reason,
|
|
})?;
|
|
|
|
self.pending_writes.push(PendingWorkspaceWrite {
|
|
path: full_path,
|
|
content,
|
|
});
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Take all pending workspace writes (clears the queue).
|
|
pub fn take_pending_writes(&mut self) -> Vec<PendingWorkspaceWrite> {
|
|
std::mem::take(&mut self.pending_writes)
|
|
}
|
|
|
|
/// Get the number of pending workspace writes.
|
|
pub fn pending_writes_count(&self) -> usize {
|
|
self.pending_writes.len()
|
|
}
|
|
|
|
/// Log a message (delegates to base).
|
|
pub fn log(
|
|
&mut self,
|
|
level: LogLevel,
|
|
message: String,
|
|
) -> Result<(), crate::tools::wasm::WasmError> {
|
|
self.base.log(level, message)
|
|
}
|
|
|
|
/// Get current timestamp in milliseconds (delegates to base).
|
|
pub fn now_millis(&self) -> u64 {
|
|
self.base.now_millis()
|
|
}
|
|
|
|
/// Read from workspace (delegates to base).
|
|
pub fn workspace_read(
|
|
&self,
|
|
path: &str,
|
|
) -> Result<Option<String>, crate::tools::wasm::WasmError> {
|
|
// Prefix the path with channel namespace before reading
|
|
let full_path = self.capabilities.prefix_workspace_path(path);
|
|
self.base.workspace_read(&full_path)
|
|
}
|
|
|
|
/// Check if a secret exists (delegates to base).
|
|
pub fn secret_exists(&self, name: &str) -> bool {
|
|
self.base.secret_exists(name)
|
|
}
|
|
|
|
/// Check if HTTP is allowed (delegates to base).
|
|
pub fn check_http_allowed(&self, url: &str, method: &str) -> Result<(), String> {
|
|
self.base.check_http_allowed(url, method)
|
|
}
|
|
|
|
/// Record an HTTP request (delegates to base).
|
|
pub fn record_http_request(&mut self) -> Result<(), String> {
|
|
self.base.record_http_request()
|
|
}
|
|
|
|
/// Take logs (delegates to base).
|
|
pub fn take_logs(&mut self) -> Vec<crate::tools::wasm::LogEntry> {
|
|
self.base.take_logs()
|
|
}
|
|
}
|
|
|
|
/// In-memory workspace store for WASM channels.
|
|
///
|
|
/// Persists workspace writes across callback invocations within a single
|
|
/// channel lifetime. This allows WASM channels to maintain state (e.g.,
|
|
/// Telegram polling offsets) between poll ticks without requiring a
|
|
/// full database-backed workspace.
|
|
///
|
|
/// Uses `std::sync::RwLock` (not tokio) because WASM execution runs
|
|
/// inside `spawn_blocking`.
|
|
pub struct ChannelWorkspaceStore {
|
|
data: std::sync::RwLock<std::collections::HashMap<String, String>>,
|
|
}
|
|
|
|
impl ChannelWorkspaceStore {
|
|
/// Create a new empty workspace store.
|
|
pub fn new() -> Self {
|
|
Self {
|
|
data: std::sync::RwLock::new(std::collections::HashMap::new()),
|
|
}
|
|
}
|
|
|
|
/// Commit pending writes from a callback execution into the store.
|
|
pub fn commit_writes(&self, writes: &[PendingWorkspaceWrite]) {
|
|
if writes.is_empty() {
|
|
return;
|
|
}
|
|
if let Ok(mut data) = self.data.write() {
|
|
for write in writes {
|
|
tracing::debug!(
|
|
path = %write.path,
|
|
content_len = write.content.len(),
|
|
"Committing workspace write to channel store"
|
|
);
|
|
data.insert(write.path.clone(), write.content.clone());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl crate::tools::wasm::WorkspaceReader for ChannelWorkspaceStore {
|
|
fn read(&self, path: &str) -> Option<String> {
|
|
self.data.read().ok()?.get(path).cloned()
|
|
}
|
|
}
|
|
|
|
/// Rate limiter for channel message emission.
|
|
///
|
|
/// Tracks emission rates across multiple executions.
|
|
pub struct ChannelEmitRateLimiter {
|
|
config: EmitRateLimitConfig,
|
|
minute_window: RateWindow,
|
|
hour_window: RateWindow,
|
|
}
|
|
|
|
struct RateWindow {
|
|
count: u32,
|
|
window_start: u64,
|
|
window_duration_ms: u64,
|
|
}
|
|
|
|
impl RateWindow {
|
|
fn new(duration_ms: u64) -> Self {
|
|
Self {
|
|
count: 0,
|
|
window_start: 0,
|
|
window_duration_ms: duration_ms,
|
|
}
|
|
}
|
|
|
|
fn check_and_record(&mut self, now_ms: u64, limit: u32) -> bool {
|
|
// Reset window if expired
|
|
if now_ms.saturating_sub(self.window_start) > self.window_duration_ms {
|
|
self.count = 0;
|
|
self.window_start = now_ms;
|
|
}
|
|
|
|
if self.count >= limit {
|
|
return false;
|
|
}
|
|
|
|
self.count += 1;
|
|
true
|
|
}
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
impl ChannelEmitRateLimiter {
|
|
/// Create a new rate limiter with the given config.
|
|
pub fn new(config: EmitRateLimitConfig) -> Self {
|
|
Self {
|
|
config,
|
|
minute_window: RateWindow::new(60_000), // 1 minute
|
|
hour_window: RateWindow::new(3_600_000), // 1 hour
|
|
}
|
|
}
|
|
|
|
/// Check if an emit is allowed and record it if so.
|
|
///
|
|
/// Returns true if the emit is allowed, false if rate limited.
|
|
pub fn check_and_record(&mut self) -> bool {
|
|
let now = SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.map(|d| d.as_millis() as u64)
|
|
.unwrap_or(0);
|
|
|
|
// Check both windows
|
|
let minute_ok = self
|
|
.minute_window
|
|
.check_and_record(now, self.config.messages_per_minute);
|
|
let hour_ok = self
|
|
.hour_window
|
|
.check_and_record(now, self.config.messages_per_hour);
|
|
|
|
minute_ok && hour_ok
|
|
}
|
|
|
|
/// Get the current emission count for the minute window.
|
|
pub fn minute_count(&self) -> u32 {
|
|
self.minute_window.count
|
|
}
|
|
|
|
/// Get the current emission count for the hour window.
|
|
pub fn hour_count(&self) -> u32 {
|
|
self.hour_window.count
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use crate::channels::wasm::capabilities::{ChannelCapabilities, EmitRateLimitConfig};
|
|
use crate::channels::wasm::host::{
|
|
Attachment, ChannelEmitRateLimiter, ChannelHostState, EmittedMessage,
|
|
MAX_ATTACHMENT_TOTAL_SIZE, MAX_ATTACHMENTS_PER_MESSAGE, MAX_EMITS_PER_EXECUTION,
|
|
MAX_MESSAGE_CONTENT_SIZE,
|
|
};
|
|
|
|
#[test]
|
|
fn test_emit_message_basic() {
|
|
let caps = ChannelCapabilities::for_channel("test");
|
|
let mut state = ChannelHostState::new("test", caps);
|
|
|
|
let msg = EmittedMessage::new("user123", "Hello, world!");
|
|
state.emit_message(msg).unwrap();
|
|
|
|
assert_eq!(state.emitted_count(), 1);
|
|
|
|
let messages = state.take_emitted_messages();
|
|
assert_eq!(messages.len(), 1);
|
|
assert_eq!(messages[0].user_id, "user123");
|
|
assert_eq!(messages[0].content, "Hello, world!");
|
|
|
|
// Queue should be cleared
|
|
assert_eq!(state.emitted_count(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_emit_message_with_metadata() {
|
|
let caps = ChannelCapabilities::for_channel("test");
|
|
let mut state = ChannelHostState::new("test", caps);
|
|
|
|
let msg = EmittedMessage::new("user123", "Hello")
|
|
.with_user_name("John Doe")
|
|
.with_thread_id("thread-1")
|
|
.with_metadata(r#"{"key": "value"}"#);
|
|
|
|
state.emit_message(msg).unwrap();
|
|
|
|
let messages = state.take_emitted_messages();
|
|
assert_eq!(messages[0].user_name, Some("John Doe".to_string()));
|
|
assert_eq!(messages[0].thread_id, Some("thread-1".to_string()));
|
|
assert_eq!(messages[0].metadata_json, r#"{"key": "value"}"#);
|
|
}
|
|
|
|
#[test]
|
|
fn test_emit_per_execution_limit() {
|
|
let caps = ChannelCapabilities::for_channel("test");
|
|
let mut state = ChannelHostState::new("test", caps);
|
|
|
|
// Fill up to limit
|
|
for i in 0..MAX_EMITS_PER_EXECUTION {
|
|
let msg = EmittedMessage::new("user", format!("Message {}", i));
|
|
state.emit_message(msg).unwrap();
|
|
}
|
|
|
|
// This should be dropped silently
|
|
let msg = EmittedMessage::new("user", "Should be dropped");
|
|
state.emit_message(msg).unwrap();
|
|
|
|
assert_eq!(state.emitted_count(), MAX_EMITS_PER_EXECUTION);
|
|
assert_eq!(state.emits_dropped(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_emit_message_truncates_utf8_safely() {
|
|
let caps = ChannelCapabilities::for_channel("test");
|
|
let mut state = ChannelHostState::new("test", caps);
|
|
|
|
let prefix = "a".repeat(MAX_MESSAGE_CONTENT_SIZE - 1);
|
|
let content = format!("{}🙂suffix", prefix);
|
|
let msg = EmittedMessage::new("user123", content);
|
|
|
|
state.emit_message(msg).unwrap();
|
|
let messages = state.take_emitted_messages();
|
|
assert_eq!(messages.len(), 1);
|
|
|
|
let emitted = &messages[0].content;
|
|
assert!(emitted.starts_with(&prefix));
|
|
assert!(emitted.ends_with("... (truncated)"));
|
|
assert!(!emitted.contains("🙂"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_workspace_write_prefixing() {
|
|
let caps = ChannelCapabilities::for_channel("slack");
|
|
let mut state = ChannelHostState::new("slack", caps);
|
|
|
|
state
|
|
.workspace_write("state.json", "{}".to_string())
|
|
.unwrap();
|
|
|
|
let writes = state.take_pending_writes();
|
|
assert_eq!(writes.len(), 1);
|
|
assert_eq!(writes[0].path, "channels/slack/state.json");
|
|
}
|
|
|
|
#[test]
|
|
fn test_workspace_write_path_traversal_blocked() {
|
|
let caps = ChannelCapabilities::for_channel("slack");
|
|
let mut state = ChannelHostState::new("slack", caps);
|
|
|
|
// Try to escape namespace
|
|
let result = state.workspace_write("../secrets.json", "{}".to_string());
|
|
assert!(result.is_err());
|
|
|
|
// Absolute path
|
|
let result = state.workspace_write("/etc/passwd", "{}".to_string());
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_rate_limiter_basic() {
|
|
let config = EmitRateLimitConfig {
|
|
messages_per_minute: 10,
|
|
messages_per_hour: 100,
|
|
};
|
|
let mut limiter = ChannelEmitRateLimiter::new(config);
|
|
|
|
// Should allow 10 messages
|
|
for _ in 0..10 {
|
|
assert!(limiter.check_and_record());
|
|
}
|
|
|
|
// 11th should be blocked
|
|
assert!(!limiter.check_and_record());
|
|
}
|
|
|
|
#[test]
|
|
fn test_channel_name() {
|
|
let caps = ChannelCapabilities::for_channel("telegram");
|
|
let state = ChannelHostState::new("telegram", caps);
|
|
|
|
assert_eq!(state.channel_name(), "telegram");
|
|
}
|
|
|
|
#[test]
|
|
fn test_channel_workspace_store_commit_and_read() {
|
|
use crate::channels::wasm::host::{ChannelWorkspaceStore, PendingWorkspaceWrite};
|
|
use crate::tools::wasm::WorkspaceReader;
|
|
|
|
let store = ChannelWorkspaceStore::new();
|
|
|
|
// Initially empty
|
|
assert!(store.read("channels/telegram/offset").is_none());
|
|
|
|
// Commit some writes
|
|
let writes = vec![
|
|
PendingWorkspaceWrite {
|
|
path: "channels/telegram/offset".to_string(),
|
|
content: "103".to_string(),
|
|
},
|
|
PendingWorkspaceWrite {
|
|
path: "channels/telegram/state.json".to_string(),
|
|
content: r#"{"ok":true}"#.to_string(),
|
|
},
|
|
];
|
|
store.commit_writes(&writes);
|
|
|
|
// Should be readable
|
|
assert_eq!(
|
|
store.read("channels/telegram/offset"),
|
|
Some("103".to_string())
|
|
);
|
|
assert_eq!(
|
|
store.read("channels/telegram/state.json"),
|
|
Some(r#"{"ok":true}"#.to_string())
|
|
);
|
|
|
|
// Overwrite a value
|
|
let writes2 = vec![PendingWorkspaceWrite {
|
|
path: "channels/telegram/offset".to_string(),
|
|
content: "200".to_string(),
|
|
}];
|
|
store.commit_writes(&writes2);
|
|
assert_eq!(
|
|
store.read("channels/telegram/offset"),
|
|
Some("200".to_string())
|
|
);
|
|
|
|
// Empty writes are a no-op
|
|
store.commit_writes(&[]);
|
|
assert_eq!(
|
|
store.read("channels/telegram/offset"),
|
|
Some("200".to_string())
|
|
);
|
|
}
|
|
|
|
// === QA Plan P2 - 2.3: WASM channel lifecycle tests ===
|
|
|
|
#[test]
|
|
fn test_workspace_write_then_read_round_trip() {
|
|
// Full lifecycle: write in one "callback", commit, then read in a
|
|
// subsequent "callback" using the same store as the workspace reader.
|
|
use crate::channels::wasm::host::ChannelWorkspaceStore;
|
|
use crate::tools::wasm::{WorkspaceCapability, WorkspaceReader};
|
|
use std::sync::Arc;
|
|
|
|
let store = Arc::new(ChannelWorkspaceStore::new());
|
|
|
|
// --- Callback 1: write workspace data ---
|
|
let caps = ChannelCapabilities::for_channel("telegram");
|
|
let mut state = ChannelHostState::new("telegram", caps);
|
|
|
|
state
|
|
.workspace_write("offset", "12345".to_string())
|
|
.unwrap();
|
|
state
|
|
.workspace_write("state.json", r#"{"ok":true}"#.to_string())
|
|
.unwrap();
|
|
|
|
let writes = state.take_pending_writes();
|
|
assert_eq!(writes.len(), 2);
|
|
store.commit_writes(&writes);
|
|
|
|
// --- Callback 2: read back the data written in callback 1 ---
|
|
// Build capabilities with the store as the workspace reader.
|
|
let mut caps2 = ChannelCapabilities::for_channel("telegram");
|
|
caps2.tool_capabilities.workspace_read = Some(WorkspaceCapability {
|
|
allowed_prefixes: vec![], // empty = all paths allowed
|
|
reader: Some(Arc::clone(&store) as Arc<dyn WorkspaceReader>),
|
|
});
|
|
let state2 = ChannelHostState::new("telegram", caps2);
|
|
|
|
// workspace_read prefixes path with "channels/telegram/" before delegating.
|
|
let offset = state2.workspace_read("offset").unwrap();
|
|
assert_eq!(offset, Some("12345".to_string()));
|
|
|
|
let json = state2.workspace_read("state.json").unwrap();
|
|
assert_eq!(json, Some(r#"{"ok":true}"#.to_string()));
|
|
|
|
// Non-existent key returns None.
|
|
let missing = state2.workspace_read("no_such_key").unwrap();
|
|
assert!(missing.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_workspace_overwrite_across_callbacks() {
|
|
// Verify that a second write to the same key overwrites the first.
|
|
use crate::channels::wasm::host::ChannelWorkspaceStore;
|
|
use crate::tools::wasm::{WorkspaceCapability, WorkspaceReader};
|
|
use std::sync::Arc;
|
|
|
|
let store = Arc::new(ChannelWorkspaceStore::new());
|
|
|
|
// Callback 1: write initial value.
|
|
let caps = ChannelCapabilities::for_channel("slack");
|
|
let mut state = ChannelHostState::new("slack", caps);
|
|
state.workspace_write("cursor", "100".to_string()).unwrap();
|
|
let writes = state.take_pending_writes();
|
|
store.commit_writes(&writes);
|
|
|
|
// Callback 2: overwrite the same key.
|
|
let caps2 = ChannelCapabilities::for_channel("slack");
|
|
let mut state2 = ChannelHostState::new("slack", caps2);
|
|
state2.workspace_write("cursor", "200".to_string()).unwrap();
|
|
let writes2 = state2.take_pending_writes();
|
|
store.commit_writes(&writes2);
|
|
|
|
// Callback 3: read back -- should see the overwritten value.
|
|
let mut caps3 = ChannelCapabilities::for_channel("slack");
|
|
caps3.tool_capabilities.workspace_read = Some(WorkspaceCapability {
|
|
allowed_prefixes: vec![],
|
|
reader: Some(Arc::clone(&store) as Arc<dyn WorkspaceReader>),
|
|
});
|
|
let state3 = ChannelHostState::new("slack", caps3);
|
|
|
|
let value = state3.workspace_read("cursor").unwrap();
|
|
assert_eq!(value, Some("200".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_emit_and_take_preserves_order_and_content() {
|
|
// Emit multiple messages, take them, verify order and content.
|
|
let caps = ChannelCapabilities::for_channel("discord");
|
|
let mut state = ChannelHostState::new("discord", caps);
|
|
|
|
let messages_data = vec![
|
|
("user-a", "Hello from A"),
|
|
("user-b", "Hello from B"),
|
|
("user-a", "Follow-up from A"),
|
|
];
|
|
for (uid, content) in &messages_data {
|
|
state
|
|
.emit_message(EmittedMessage::new(*uid, *content))
|
|
.unwrap();
|
|
}
|
|
|
|
assert_eq!(state.emitted_count(), 3);
|
|
|
|
let taken = state.take_emitted_messages();
|
|
assert_eq!(taken.len(), 3);
|
|
|
|
// Order preserved.
|
|
for (i, (uid, content)) in messages_data.iter().enumerate() {
|
|
assert_eq!(taken[i].user_id, *uid);
|
|
assert_eq!(taken[i].content, *content);
|
|
}
|
|
|
|
// Take empties the queue.
|
|
assert_eq!(state.emitted_count(), 0);
|
|
let taken2 = state.take_emitted_messages();
|
|
assert!(taken2.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_channels_have_isolated_namespaces() {
|
|
// Two channels writing to the same relative path should not collide.
|
|
use crate::channels::wasm::host::ChannelWorkspaceStore;
|
|
use crate::tools::wasm::{WorkspaceCapability, WorkspaceReader};
|
|
use std::sync::Arc;
|
|
|
|
let store = Arc::new(ChannelWorkspaceStore::new());
|
|
|
|
// Telegram writes "offset" = "100".
|
|
let caps_tg = ChannelCapabilities::for_channel("telegram");
|
|
let mut state_tg = ChannelHostState::new("telegram", caps_tg);
|
|
state_tg
|
|
.workspace_write("offset", "100".to_string())
|
|
.unwrap();
|
|
store.commit_writes(&state_tg.take_pending_writes());
|
|
|
|
// Slack writes "offset" = "200".
|
|
let caps_sl = ChannelCapabilities::for_channel("slack");
|
|
let mut state_sl = ChannelHostState::new("slack", caps_sl);
|
|
state_sl
|
|
.workspace_write("offset", "200".to_string())
|
|
.unwrap();
|
|
store.commit_writes(&state_sl.take_pending_writes());
|
|
|
|
// Reading back: each channel sees its own value.
|
|
let mut caps_tg_read = ChannelCapabilities::for_channel("telegram");
|
|
caps_tg_read.tool_capabilities.workspace_read = Some(WorkspaceCapability {
|
|
allowed_prefixes: vec![],
|
|
reader: Some(Arc::clone(&store) as Arc<dyn WorkspaceReader>),
|
|
});
|
|
let tg_reader = ChannelHostState::new("telegram", caps_tg_read);
|
|
assert_eq!(
|
|
tg_reader.workspace_read("offset").unwrap(),
|
|
Some("100".to_string())
|
|
);
|
|
|
|
let mut caps_sl_read = ChannelCapabilities::for_channel("slack");
|
|
caps_sl_read.tool_capabilities.workspace_read = Some(WorkspaceCapability {
|
|
allowed_prefixes: vec![],
|
|
reader: Some(Arc::clone(&store) as Arc<dyn WorkspaceReader>),
|
|
});
|
|
let sl_reader = ChannelHostState::new("slack", caps_sl_read);
|
|
assert_eq!(
|
|
sl_reader.workspace_read("offset").unwrap(),
|
|
Some("200".to_string())
|
|
);
|
|
}
|
|
|
|
// === Attachment validation tests ===
|
|
|
|
fn make_attachment(id: &str, mime: &str, size: Option<u64>) -> Attachment {
|
|
Attachment {
|
|
id: id.to_string(),
|
|
mime_type: mime.to_string(),
|
|
filename: None,
|
|
size_bytes: size,
|
|
source_url: None,
|
|
storage_key: None,
|
|
extracted_text: None,
|
|
data: Vec::new(),
|
|
duration_secs: None,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_emit_message_with_attachments() {
|
|
let caps = ChannelCapabilities::for_channel("test");
|
|
let mut state = ChannelHostState::new("test", caps);
|
|
|
|
let msg = EmittedMessage::new("user1", "Check this image")
|
|
.with_attachments(vec![make_attachment("file1", "image/jpeg", Some(1024))]);
|
|
|
|
state.emit_message(msg).unwrap();
|
|
|
|
let messages = state.take_emitted_messages();
|
|
assert_eq!(messages.len(), 1);
|
|
assert_eq!(messages[0].attachments.len(), 1);
|
|
assert_eq!(messages[0].attachments[0].id, "file1");
|
|
assert_eq!(messages[0].attachments[0].mime_type, "image/jpeg");
|
|
assert_eq!(messages[0].attachments[0].size_bytes, Some(1024));
|
|
}
|
|
|
|
#[test]
|
|
fn test_emit_message_no_attachments_backward_compat() {
|
|
let caps = ChannelCapabilities::for_channel("test");
|
|
let mut state = ChannelHostState::new("test", caps);
|
|
|
|
let msg = EmittedMessage::new("user1", "Just text");
|
|
state.emit_message(msg).unwrap();
|
|
|
|
let messages = state.take_emitted_messages();
|
|
assert_eq!(messages.len(), 1);
|
|
assert!(messages[0].attachments.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_attachment_count_limit() {
|
|
let caps = ChannelCapabilities::for_channel("test");
|
|
let mut state = ChannelHostState::new("test", caps);
|
|
|
|
let attachments: Vec<Attachment> = (0..MAX_ATTACHMENTS_PER_MESSAGE + 5)
|
|
.map(|i| make_attachment(&format!("file{}", i), "image/png", Some(100)))
|
|
.collect();
|
|
|
|
let msg = EmittedMessage::new("user1", "Many files").with_attachments(attachments);
|
|
state.emit_message(msg).unwrap();
|
|
|
|
let messages = state.take_emitted_messages();
|
|
assert_eq!(messages[0].attachments.len(), MAX_ATTACHMENTS_PER_MESSAGE);
|
|
}
|
|
|
|
#[test]
|
|
fn test_attachment_total_size_limit() {
|
|
let caps = ChannelCapabilities::for_channel("test");
|
|
let mut state = ChannelHostState::new("test", caps);
|
|
|
|
// Each file is 1/3 of the limit, so 3 fit but 4th does not
|
|
let chunk_size = MAX_ATTACHMENT_TOTAL_SIZE / 3;
|
|
let attachments = vec![
|
|
make_attachment("file1", "image/png", Some(chunk_size)),
|
|
make_attachment("file2", "image/png", Some(chunk_size)),
|
|
make_attachment("file3", "image/png", Some(chunk_size)),
|
|
make_attachment("file4", "image/png", Some(chunk_size)),
|
|
];
|
|
|
|
let msg = EmittedMessage::new("user1", "Big files").with_attachments(attachments);
|
|
state.emit_message(msg).unwrap();
|
|
|
|
let messages = state.take_emitted_messages();
|
|
// Only first 3 fit within the total size limit
|
|
assert_eq!(messages[0].attachments.len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_attachment_mime_type_filtering() {
|
|
let caps = ChannelCapabilities::for_channel("test");
|
|
let mut state = ChannelHostState::new("test", caps);
|
|
|
|
let attachments = vec![
|
|
make_attachment("ok1", "image/jpeg", Some(100)),
|
|
make_attachment("bad1", "application/x-executable", Some(100)),
|
|
make_attachment("ok2", "application/pdf", Some(100)),
|
|
make_attachment("bad2", "application/x-msdos-program", Some(100)),
|
|
make_attachment("ok3", "text/plain", Some(100)),
|
|
make_attachment("ok4", "audio/mpeg", Some(100)),
|
|
make_attachment("ok5", "video/mp4", Some(100)),
|
|
];
|
|
|
|
let msg = EmittedMessage::new("user1", "Mixed files").with_attachments(attachments);
|
|
state.emit_message(msg).unwrap();
|
|
|
|
let messages = state.take_emitted_messages();
|
|
let ids: Vec<&str> = messages[0]
|
|
.attachments
|
|
.iter()
|
|
.map(|a| a.id.as_str())
|
|
.collect();
|
|
assert_eq!(ids, vec!["ok1", "ok2", "ok3", "ok4", "ok5"]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_attachment_unknown_size_allowed() {
|
|
let caps = ChannelCapabilities::for_channel("test");
|
|
let mut state = ChannelHostState::new("test", caps);
|
|
|
|
let attachments = vec![
|
|
make_attachment("file1", "image/jpeg", None),
|
|
make_attachment("file2", "image/png", None),
|
|
];
|
|
|
|
let msg = EmittedMessage::new("user1", "No sizes").with_attachments(attachments);
|
|
state.emit_message(msg).unwrap();
|
|
|
|
let messages = state.take_emitted_messages();
|
|
assert_eq!(messages[0].attachments.len(), 2);
|
|
}
|
|
}
|