mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3f970252d5 | ||
|
|
3930184e71 | ||
|
|
4fb0d53206 | ||
|
|
5f50df7583 | ||
|
|
30feaabb7c | ||
|
|
427f908e71 | ||
|
|
8bd30a970e | ||
|
|
5d1d504e11 | ||
|
|
9bb19a98f7 | ||
|
|
0b33ca9926 | ||
|
|
9ba10eac35 | ||
|
|
27e8d6f8dd | ||
|
|
f49f368355 |
@@ -0,0 +1,4 @@
|
||||
-- Add source_channel to conversations for cross-channel approval authorization.
|
||||
-- Tracks which channel originally created a conversation so that approval
|
||||
-- messages from other channels can be validated.
|
||||
ALTER TABLE conversations ADD COLUMN source_channel TEXT;
|
||||
+35
-2
@@ -843,7 +843,10 @@ impl Agent {
|
||||
{
|
||||
use crate::agent::session::Thread;
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(id, sess.id);
|
||||
// Bootstrap thread has no incoming message -- use the
|
||||
// "__bootstrap__" sentinel so approvals from any channel are
|
||||
// permitted. None means "deny by default" (fail-closed).
|
||||
let thread = Thread::with_id(id, sess.id, Some("__bootstrap__"));
|
||||
sess.active_thread = Some(id);
|
||||
sess.threads.entry(id).or_insert(thread);
|
||||
}
|
||||
@@ -1148,7 +1151,37 @@ impl Agent {
|
||||
.get_or_create_session(&message.user_id)
|
||||
.await;
|
||||
let mut sess = session.lock().await;
|
||||
if sess.threads.contains_key(&target_thread_id) {
|
||||
if let Some(thread) = sess.threads.get(&target_thread_id) {
|
||||
// Verify the thread actually has a pending approval before
|
||||
// allowing approval-shaped messages to target it. Without this
|
||||
// check, an attacker could use approval messages to hijack any
|
||||
// thread by UUID.
|
||||
if thread.pending_approval.is_none() {
|
||||
tracing::warn!(
|
||||
%target_thread_id,
|
||||
approval_channel = %message.channel,
|
||||
"Blocked approval for thread with no pending approval"
|
||||
);
|
||||
drop(sess);
|
||||
return Ok(Some("Error: no pending approval on this thread".into()));
|
||||
}
|
||||
|
||||
let authorized = crate::agent::session::is_approval_authorized(
|
||||
thread.source_channel.as_deref(),
|
||||
&message.channel,
|
||||
);
|
||||
if !authorized {
|
||||
tracing::warn!(
|
||||
%target_thread_id,
|
||||
source_channel = ?thread.source_channel,
|
||||
approval_channel = %message.channel,
|
||||
"Blocked cross-channel approval attempt"
|
||||
);
|
||||
drop(sess);
|
||||
return Ok(Some(
|
||||
"Error: approval not authorized for this channel".into(),
|
||||
));
|
||||
}
|
||||
sess.active_thread = Some(target_thread_id);
|
||||
sess.last_active_at = chrono::Utc::now();
|
||||
drop(sess);
|
||||
|
||||
@@ -319,7 +319,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_format_turns() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
thread.start_turn("Hello");
|
||||
thread.complete_turn("Hi there");
|
||||
thread.start_turn("How are you?");
|
||||
@@ -351,7 +351,7 @@ mod tests {
|
||||
/// Helper: build a thread with `n` completed turns.
|
||||
/// Turn `i` has user_input "msg-{i}" and response "resp-{i}".
|
||||
fn make_thread(n: usize) -> Thread {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
for i in 0..n {
|
||||
thread.start_turn(format!("msg-{}", i));
|
||||
thread.complete_turn(format!("resp-{}", i));
|
||||
@@ -457,7 +457,7 @@ mod tests {
|
||||
async fn test_compact_truncate_empty_turns() {
|
||||
let llm = Arc::new(StubLlm::new("unused"));
|
||||
let compactor = make_compactor(llm);
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
assert!(thread.turns.is_empty());
|
||||
|
||||
let result = compactor
|
||||
@@ -698,7 +698,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_format_turns_for_storage_with_tool_calls() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
thread.start_turn("Search for X");
|
||||
// Record a tool call on the current turn
|
||||
if let Some(turn) = thread.turns.last_mut() {
|
||||
@@ -719,7 +719,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_format_turns_for_storage_incomplete_turn() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
thread.start_turn("In progress message");
|
||||
// Don't complete the turn
|
||||
|
||||
|
||||
@@ -2299,7 +2299,7 @@ mod tests {
|
||||
// Initialize a thread in the session so the loop can record tool calls.
|
||||
let thread_id = {
|
||||
let mut sess = session.lock().await;
|
||||
sess.create_thread().id
|
||||
sess.create_thread(Some("test")).id
|
||||
};
|
||||
|
||||
let message = IncomingMessage::new("test", "test-user", "do something");
|
||||
@@ -2412,7 +2412,7 @@ mod tests {
|
||||
let session = Arc::new(Mutex::new(Session::new("test-user")));
|
||||
let thread_id = {
|
||||
let mut sess = session.lock().await;
|
||||
sess.create_thread().id
|
||||
sess.create_thread(Some("test")).id
|
||||
};
|
||||
|
||||
let message = IncomingMessage::new("test", "test-user", "keep calling tools");
|
||||
|
||||
+169
-54
@@ -68,8 +68,8 @@ impl Session {
|
||||
}
|
||||
|
||||
/// Create a new thread in this session.
|
||||
pub fn create_thread(&mut self) -> &mut Thread {
|
||||
let thread = Thread::new(self.id);
|
||||
pub fn create_thread(&mut self, channel: Option<&str>) -> &mut Thread {
|
||||
let thread = Thread::new(self.id, channel);
|
||||
let thread_id = thread.id;
|
||||
self.active_thread = Some(thread_id);
|
||||
self.last_active_at = Utc::now();
|
||||
@@ -87,9 +87,9 @@ impl Session {
|
||||
}
|
||||
|
||||
/// Get or create the active thread.
|
||||
pub fn get_or_create_thread(&mut self) -> &mut Thread {
|
||||
pub fn get_or_create_thread(&mut self, channel: Option<&str>) -> &mut Thread {
|
||||
match self.active_thread {
|
||||
None => self.create_thread(),
|
||||
None => self.create_thread(channel),
|
||||
Some(id) => {
|
||||
if self.threads.contains_key(&id) {
|
||||
// Entry existence confirmed by contains_key above.
|
||||
@@ -100,7 +100,7 @@ impl Session {
|
||||
} else {
|
||||
// Stale active_thread ID: create a new thread, which
|
||||
// updates self.active_thread to the new thread's ID.
|
||||
self.create_thread()
|
||||
self.create_thread(channel)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -225,6 +225,9 @@ pub struct Thread {
|
||||
/// Messages queued while the thread was processing a turn.
|
||||
#[serde(default, skip_serializing_if = "VecDeque::is_empty")]
|
||||
pub pending_messages: VecDeque<String>,
|
||||
/// Channel that created this thread (for approval authorization).
|
||||
#[serde(default)]
|
||||
pub source_channel: Option<String>,
|
||||
}
|
||||
|
||||
/// Maximum number of messages that can be queued while a thread is processing.
|
||||
@@ -233,9 +236,29 @@ pub struct Thread {
|
||||
/// rapid follow-ups. The drain loop processes them as one newline-delimited turn.
|
||||
pub const MAX_PENDING_MESSAGES: usize = 10;
|
||||
|
||||
/// Sentinel value for bootstrap threads that accept approvals from any channel.
|
||||
pub const BOOTSTRAP_SOURCE_CHANNEL: &str = "__bootstrap__";
|
||||
|
||||
/// Check whether an approval from `requesting_channel` is authorized for a
|
||||
/// thread whose `source_channel` is `source`.
|
||||
///
|
||||
/// Rules:
|
||||
/// - `None` (unknown origin) -> denied (fail-closed)
|
||||
/// - `Some("__bootstrap__")` -> authorized from any channel
|
||||
/// - `Some(src) == requesting` -> same channel, authorized
|
||||
/// - requesting is "web" or "gateway" -> always authorized (trusted UI)
|
||||
/// - Otherwise -> denied
|
||||
pub fn is_approval_authorized(source: Option<&str>, requesting: &str) -> bool {
|
||||
match source {
|
||||
None => false,
|
||||
Some(src) if src == BOOTSTRAP_SOURCE_CHANNEL => true,
|
||||
Some(src) => src == requesting || requesting == "web" || requesting == "gateway",
|
||||
}
|
||||
}
|
||||
|
||||
impl Thread {
|
||||
/// Create a new thread.
|
||||
pub fn new(session_id: Uuid) -> Self {
|
||||
pub fn new(session_id: Uuid, source_channel: Option<&str>) -> Self {
|
||||
let now = Utc::now();
|
||||
Self {
|
||||
id: Uuid::new_v4(),
|
||||
@@ -248,11 +271,12 @@ impl Thread {
|
||||
pending_approval: None,
|
||||
pending_auth: None,
|
||||
pending_messages: VecDeque::new(),
|
||||
source_channel: source_channel.map(String::from),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a thread with a specific ID (for DB hydration).
|
||||
pub fn with_id(id: Uuid, session_id: Uuid) -> Self {
|
||||
pub fn with_id(id: Uuid, session_id: Uuid, source_channel: Option<&str>) -> Self {
|
||||
let now = Utc::now();
|
||||
Self {
|
||||
id,
|
||||
@@ -265,6 +289,7 @@ impl Thread {
|
||||
pending_approval: None,
|
||||
pending_auth: None,
|
||||
pending_messages: VecDeque::new(),
|
||||
source_channel: source_channel.map(String::from),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -787,13 +812,13 @@ mod tests {
|
||||
let mut session = Session::new("user-123");
|
||||
assert!(session.active_thread.is_none());
|
||||
|
||||
session.create_thread();
|
||||
session.create_thread(None);
|
||||
assert!(session.active_thread.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thread_turns() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
thread.start_turn("Hello");
|
||||
assert_eq!(thread.state, ThreadState::Processing);
|
||||
@@ -806,7 +831,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_thread_messages() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
thread.start_turn("First message");
|
||||
thread.complete_turn("First response");
|
||||
@@ -829,7 +854,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_messages() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// First add some turns
|
||||
thread.start_turn("Original message");
|
||||
@@ -855,7 +880,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_messages_incomplete_turn() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Messages with incomplete last turn (no assistant response)
|
||||
let messages = vec![
|
||||
@@ -874,7 +899,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_enter_auth_mode() {
|
||||
let before = Utc::now();
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
assert!(thread.pending_auth.is_none());
|
||||
|
||||
thread.enter_auth_mode("telegram".to_string());
|
||||
@@ -887,7 +912,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_take_pending_auth() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
thread.enter_auth_mode("notion".to_string());
|
||||
|
||||
let pending = thread.take_pending_auth();
|
||||
@@ -902,7 +927,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_pending_auth_serialization() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
thread.enter_auth_mode("openai".to_string());
|
||||
|
||||
let json = serde_json::to_string(&thread).expect("should serialize");
|
||||
@@ -932,7 +957,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_pending_auth_default_none() {
|
||||
// Deserialization of old data without pending_auth should default to None
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
thread.pending_auth = None;
|
||||
let json = serde_json::to_string(&thread).expect("serialize");
|
||||
|
||||
@@ -946,7 +971,7 @@ mod tests {
|
||||
fn test_thread_with_id() {
|
||||
let specific_id = Uuid::new_v4();
|
||||
let session_id = Uuid::new_v4();
|
||||
let thread = Thread::with_id(specific_id, session_id);
|
||||
let thread = Thread::with_id(specific_id, session_id, None);
|
||||
|
||||
assert_eq!(thread.id, specific_id);
|
||||
assert_eq!(thread.session_id, session_id);
|
||||
@@ -958,7 +983,7 @@ mod tests {
|
||||
fn test_thread_with_id_restore_messages() {
|
||||
let thread_id = Uuid::new_v4();
|
||||
let session_id = Uuid::new_v4();
|
||||
let mut thread = Thread::with_id(thread_id, session_id);
|
||||
let mut thread = Thread::with_id(thread_id, session_id, None);
|
||||
|
||||
let messages = vec![
|
||||
ChatMessage::user("Hello from DB"),
|
||||
@@ -977,7 +1002,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_messages_empty() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Add a turn first, then restore with empty vec
|
||||
thread.start_turn("hello");
|
||||
@@ -993,7 +1018,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_messages_only_assistant_messages() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Only assistant messages (no user messages to anchor turns)
|
||||
let messages = vec![
|
||||
@@ -1010,7 +1035,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_messages_multiple_user_messages_in_a_row() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Two user messages with no assistant response between them
|
||||
let messages = vec![
|
||||
@@ -1037,8 +1062,8 @@ mod tests {
|
||||
fn test_thread_switch() {
|
||||
let mut session = Session::new("user-1");
|
||||
|
||||
let t1_id = session.create_thread().id;
|
||||
let t2_id = session.create_thread().id;
|
||||
let t1_id = session.create_thread(None).id;
|
||||
let t2_id = session.create_thread(None).id;
|
||||
|
||||
// After creating two threads, active should be the last one
|
||||
assert_eq!(session.active_thread, Some(t2_id));
|
||||
@@ -1058,8 +1083,8 @@ mod tests {
|
||||
fn test_get_or_create_thread_idempotent() {
|
||||
let mut session = Session::new("user-1");
|
||||
|
||||
let tid1 = session.get_or_create_thread().id;
|
||||
let tid2 = session.get_or_create_thread().id;
|
||||
let tid1 = session.get_or_create_thread(None).id;
|
||||
let tid2 = session.get_or_create_thread(None).id;
|
||||
|
||||
// Should return the same thread (not create a new one each time)
|
||||
assert_eq!(tid1, tid2);
|
||||
@@ -1068,7 +1093,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_truncate_turns() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
for i in 0..5 {
|
||||
thread.start_turn(format!("msg-{}", i));
|
||||
@@ -1092,7 +1117,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_truncate_turns_noop_when_fewer() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
thread.start_turn("only one");
|
||||
thread.complete_turn("response");
|
||||
@@ -1104,7 +1129,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_thread_interrupt_and_resume() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
thread.start_turn("do something");
|
||||
assert_eq!(thread.state, ThreadState::Processing);
|
||||
@@ -1122,7 +1147,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_resume_only_from_interrupted() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Idle thread: resume should be a no-op
|
||||
assert_eq!(thread.state, ThreadState::Idle);
|
||||
@@ -1138,7 +1163,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_turn_fail() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
thread.start_turn("risky operation");
|
||||
thread.fail_turn("connection timed out");
|
||||
@@ -1154,7 +1179,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_messages_with_incomplete_last_turn() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
thread.start_turn("first");
|
||||
thread.complete_turn("first reply");
|
||||
@@ -1170,7 +1195,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_thread_serialization_round_trip() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
thread.start_turn("hello");
|
||||
thread.complete_turn("world");
|
||||
@@ -1188,7 +1213,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_session_serialization_round_trip() {
|
||||
let mut session = Session::new("user-ser");
|
||||
session.create_thread();
|
||||
session.create_thread(None);
|
||||
session.auto_approve_tool("echo");
|
||||
|
||||
let json = serde_json::to_string(&session).unwrap();
|
||||
@@ -1226,7 +1251,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_turn_number_increments() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Before any turns, turn_number() is 1 (1-indexed for display)
|
||||
assert_eq!(thread.turn_number(), 1);
|
||||
@@ -1241,7 +1266,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_complete_turn_on_empty_thread() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Completing a turn when there are no turns should be a safe no-op
|
||||
thread.complete_turn("phantom response");
|
||||
@@ -1251,7 +1276,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_fail_turn_on_empty_thread() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Failing a turn when there are no turns should be a safe no-op
|
||||
thread.fail_turn("phantom error");
|
||||
@@ -1261,7 +1286,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_pending_approval_flow() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
let approval = PendingApproval {
|
||||
request_id: Uuid::new_v4(),
|
||||
@@ -1288,7 +1313,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_clear_pending_approval() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
let approval = PendingApproval {
|
||||
request_id: Uuid::new_v4(),
|
||||
@@ -1317,7 +1342,7 @@ mod tests {
|
||||
assert!(session.active_thread().is_none());
|
||||
assert!(session.active_thread_mut().is_none());
|
||||
|
||||
let tid = session.create_thread().id;
|
||||
let tid = session.create_thread(None).id;
|
||||
|
||||
assert!(session.active_thread().is_some());
|
||||
assert_eq!(session.active_thread().unwrap().id, tid);
|
||||
@@ -1334,7 +1359,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_messages_includes_tool_calls() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
thread.start_turn("Search for X");
|
||||
{
|
||||
@@ -1366,7 +1391,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_messages_multiple_tool_calls_per_turn() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
thread.start_turn("Do two things");
|
||||
{
|
||||
@@ -1393,7 +1418,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_messages_with_tool_calls() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Build a message sequence with tool calls
|
||||
let tc = ToolCall {
|
||||
@@ -1425,7 +1450,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_messages_with_tool_error() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
let tc = ToolCall {
|
||||
id: "call_0".to_string(),
|
||||
@@ -1456,7 +1481,7 @@ mod tests {
|
||||
fn test_messages_round_trip_with_tools() {
|
||||
// Build a thread with tool calls, get messages(), restore, get messages() again
|
||||
// The two message sequences should be equivalent.
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
thread.start_turn("Do search");
|
||||
{
|
||||
@@ -1469,7 +1494,7 @@ mod tests {
|
||||
let messages_original = thread.messages();
|
||||
|
||||
// Restore into a new thread
|
||||
let mut thread2 = Thread::new(Uuid::new_v4());
|
||||
let mut thread2 = Thread::new(Uuid::new_v4(), None);
|
||||
thread2.restore_from_messages(messages_original.clone());
|
||||
|
||||
let messages_restored = thread2.messages();
|
||||
@@ -1491,7 +1516,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_restore_multi_stage_tool_calls() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
let tc1 = ToolCall {
|
||||
id: "call_a".to_string(),
|
||||
@@ -1534,7 +1559,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_messages_truncates_large_tool_results() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
thread.start_turn("Read big file");
|
||||
{
|
||||
@@ -1557,7 +1582,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_thread_message_queue() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Queue is initially empty
|
||||
assert!(thread.pending_messages.is_empty());
|
||||
@@ -1593,7 +1618,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_thread_message_queue_serialization() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Empty queue should not appear in serialization (skip_serializing_if)
|
||||
let json = serde_json::to_string(&thread).unwrap();
|
||||
@@ -1613,7 +1638,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_thread_message_queue_default_on_old_data() {
|
||||
// Deserialization of old data without pending_messages should default to empty
|
||||
let thread = Thread::new(Uuid::new_v4());
|
||||
let thread = Thread::new(Uuid::new_v4(), None);
|
||||
let json = serde_json::to_string(&thread).unwrap();
|
||||
|
||||
// The field is absent (skip_serializing_if), simulating old data
|
||||
@@ -1624,7 +1649,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_interrupt_clears_pending_messages() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Start a turn so there's something to interrupt
|
||||
thread.start_turn("initial input");
|
||||
@@ -1643,7 +1668,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_thread_state_idle_after_full_drain() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Simulate a full drain cycle: start turn, queue messages, complete turn,
|
||||
// then drain all queued messages as a single merged turn (#259).
|
||||
@@ -1671,7 +1696,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_drain_pending_messages_merges_with_newlines() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Empty queue returns None
|
||||
assert!(thread.drain_pending_messages().is_none());
|
||||
@@ -1700,7 +1725,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_requeue_drained_preserves_content_at_front() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
|
||||
// Re-queue into empty queue
|
||||
thread.requeue_drained("failed batch".to_string());
|
||||
@@ -1811,4 +1836,94 @@ mod tests {
|
||||
&serde_json::json!("done")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thread_new_stores_source_channel() {
|
||||
let thread = Thread::new(Uuid::new_v4(), Some("telegram"));
|
||||
assert_eq!(thread.source_channel.as_deref(), Some("telegram"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thread_new_none_channel() {
|
||||
let thread = Thread::new(Uuid::new_v4(), None);
|
||||
assert!(thread.source_channel.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_source_channel_serde_backcompat() {
|
||||
// Simulate deserializing a Thread from older DB records that lack source_channel.
|
||||
let thread = Thread::new(Uuid::new_v4(), Some("cli"));
|
||||
let json = serde_json::to_string(&thread).unwrap();
|
||||
|
||||
// Remove the source_channel field to simulate an old record.
|
||||
let mut value: serde_json::Value = serde_json::from_str(&json).unwrap();
|
||||
value.as_object_mut().unwrap().remove("source_channel");
|
||||
let old_json = serde_json::to_string(&value).unwrap();
|
||||
|
||||
let deserialized: Thread = serde_json::from_str(&old_json).unwrap();
|
||||
assert!(
|
||||
deserialized.source_channel.is_none(),
|
||||
"missing source_channel should deserialize as None"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_approval_authorized_same_channel() {
|
||||
assert!(
|
||||
is_approval_authorized(Some("telegram"), "telegram"),
|
||||
"same channel should be authorized"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_approval_authorized_different_channel_blocked() {
|
||||
assert!(
|
||||
!is_approval_authorized(Some("telegram"), "http"),
|
||||
"different channel should be blocked"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_approval_authorized_web_always_allowed() {
|
||||
assert!(
|
||||
is_approval_authorized(Some("telegram"), "web"),
|
||||
"web channel should always be authorized"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_approval_authorized_gateway_always_allowed() {
|
||||
assert!(
|
||||
is_approval_authorized(Some("telegram"), "gateway"),
|
||||
"gateway channel should always be authorized"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_approval_authorized_none_denied() {
|
||||
assert!(
|
||||
!is_approval_authorized(None, "telegram"),
|
||||
"None source_channel should be denied (fail-closed)"
|
||||
);
|
||||
assert!(
|
||||
!is_approval_authorized(None, "web"),
|
||||
"None source_channel should be denied even for web"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_approval_authorized_bootstrap_any_channel() {
|
||||
assert!(
|
||||
is_approval_authorized(Some(BOOTSTRAP_SOURCE_CHANNEL), "telegram"),
|
||||
"__bootstrap__ should be authorized from any channel"
|
||||
);
|
||||
assert!(
|
||||
is_approval_authorized(Some(BOOTSTRAP_SOURCE_CHANNEL), "http"),
|
||||
"__bootstrap__ should be authorized from any channel"
|
||||
);
|
||||
assert!(
|
||||
is_approval_authorized(Some(BOOTSTRAP_SOURCE_CHANNEL), "cli"),
|
||||
"__bootstrap__ should be authorized from any channel"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -200,7 +200,7 @@ impl SessionManager {
|
||||
// Create new thread (always create a new one for a new key)
|
||||
let thread_id = {
|
||||
let mut sess = session.lock().await;
|
||||
let thread = sess.create_thread();
|
||||
let thread = sess.create_thread(Some(channel));
|
||||
thread.id
|
||||
};
|
||||
|
||||
@@ -476,7 +476,7 @@ mod tests {
|
||||
let session = Arc::new(Mutex::new(Session::new("user-hydrate")));
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(thread_id, sess.id);
|
||||
let thread = Thread::with_id(thread_id, sess.id, None);
|
||||
sess.threads.insert(thread_id, thread);
|
||||
sess.active_thread = Some(thread_id);
|
||||
}
|
||||
@@ -600,7 +600,7 @@ mod tests {
|
||||
// Simulate hydration: create thread with a known UUID
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(known_uuid, session_id);
|
||||
let thread = Thread::with_id(known_uuid, session_id, None);
|
||||
sess.threads.insert(known_uuid, thread);
|
||||
}
|
||||
|
||||
@@ -627,7 +627,7 @@ mod tests {
|
||||
let session = Arc::new(Mutex::new(Session::new("user-idem")));
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(tid, sess.id);
|
||||
let thread = Thread::with_id(tid, sess.id, None);
|
||||
sess.threads.insert(tid, thread);
|
||||
}
|
||||
|
||||
@@ -656,7 +656,7 @@ mod tests {
|
||||
let session = Arc::new(Mutex::new(Session::new("user-undo")));
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(tid, sess.id);
|
||||
let thread = Thread::with_id(tid, sess.id, None);
|
||||
sess.threads.insert(tid, thread);
|
||||
}
|
||||
|
||||
@@ -680,7 +680,7 @@ mod tests {
|
||||
let session = Arc::new(Mutex::new(Session::new("user-new")));
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(tid, sess.id);
|
||||
let thread = Thread::with_id(tid, sess.id, None);
|
||||
sess.threads.insert(tid, thread);
|
||||
}
|
||||
|
||||
@@ -788,7 +788,7 @@ mod tests {
|
||||
let session = Arc::new(Mutex::new(Session::new("user-cross")));
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(tid, sess.id);
|
||||
let thread = Thread::with_id(tid, sess.id, None);
|
||||
sess.threads.insert(tid, thread);
|
||||
}
|
||||
|
||||
@@ -815,7 +815,7 @@ mod tests {
|
||||
let session = Arc::new(Mutex::new(Session::new("user-cross")));
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(tid, sess.id);
|
||||
let thread = Thread::with_id(tid, sess.id, None);
|
||||
sess.threads.insert(tid, thread);
|
||||
}
|
||||
|
||||
@@ -966,7 +966,7 @@ mod tests {
|
||||
let adopted_id = Uuid::new_v4();
|
||||
{
|
||||
let mut sess = session1.lock().await;
|
||||
let thread = Thread::with_id(adopted_id, sess.id);
|
||||
let thread = Thread::with_id(adopted_id, sess.id, None);
|
||||
sess.threads.insert(adopted_id, thread);
|
||||
}
|
||||
// Resolve with the UUID as external_thread_id -- should adopt it
|
||||
@@ -992,7 +992,7 @@ mod tests {
|
||||
let session = Arc::new(Mutex::new(Session::new("user-direct")));
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(tid, sess.id);
|
||||
let thread = Thread::with_id(tid, sess.id, None);
|
||||
sess.threads.insert(tid, thread);
|
||||
}
|
||||
{
|
||||
@@ -1030,7 +1030,7 @@ mod tests {
|
||||
let known_id = Uuid::new_v4();
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(known_id, sess.id);
|
||||
let thread = Thread::with_id(known_id, sess.id, None);
|
||||
sess.threads.insert(known_id, thread);
|
||||
}
|
||||
|
||||
@@ -1057,7 +1057,7 @@ mod tests {
|
||||
let known_id = Uuid::new_v4();
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(known_id, sess.id);
|
||||
let thread = Thread::with_id(known_id, sess.id, None);
|
||||
sess.threads.insert(known_id, thread);
|
||||
}
|
||||
|
||||
@@ -1081,7 +1081,7 @@ mod tests {
|
||||
let known_id = Uuid::new_v4();
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(known_id, sess.id);
|
||||
let thread = Thread::with_id(known_id, sess.id, None);
|
||||
sess.threads.insert(known_id, thread);
|
||||
}
|
||||
|
||||
@@ -1102,4 +1102,19 @@ mod tests {
|
||||
"should NOT adopt UUID when external_thread_id is None"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_thread_stores_source_channel() {
|
||||
let manager = SessionManager::new();
|
||||
|
||||
let (session, thread_id) = manager.resolve_thread("user-1", "telegram", None).await;
|
||||
|
||||
let sess = session.lock().await;
|
||||
let thread = sess.threads.get(&thread_id).unwrap();
|
||||
assert_eq!(
|
||||
thread.source_channel.as_deref(),
|
||||
Some("telegram"),
|
||||
"resolve_thread should store source_channel from the channel parameter"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+25
-9
@@ -135,13 +135,29 @@ impl Agent {
|
||||
msg_count = 0;
|
||||
}
|
||||
|
||||
// Create thread with the historical ID and restore messages
|
||||
// Create thread with the historical ID and restore messages.
|
||||
// Read source_channel from DB so the authorization check uses the
|
||||
// original creator's channel, not the requesting message's channel.
|
||||
let db_source_channel = if let Some(store) = self.store() {
|
||||
store
|
||||
.get_conversation_source_channel(thread_uuid)
|
||||
.await
|
||||
.unwrap_or(None)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let effective_source_channel = db_source_channel.as_deref();
|
||||
|
||||
let session_id = {
|
||||
let sess = session.lock().await;
|
||||
sess.id
|
||||
};
|
||||
|
||||
let mut thread = crate::agent::session::Thread::with_id(thread_uuid, session_id);
|
||||
let mut thread = crate::agent::session::Thread::with_id(
|
||||
thread_uuid,
|
||||
session_id,
|
||||
effective_source_channel,
|
||||
);
|
||||
if !chat_messages.is_empty() {
|
||||
thread.restore_from_messages(chat_messages);
|
||||
}
|
||||
@@ -636,7 +652,7 @@ impl Agent {
|
||||
user_id: &str,
|
||||
) -> bool {
|
||||
match store
|
||||
.ensure_conversation(thread_id, channel, user_id, None)
|
||||
.ensure_conversation(thread_id, channel, user_id, None, Some(channel))
|
||||
.await
|
||||
{
|
||||
Ok(true) => true,
|
||||
@@ -1781,7 +1797,7 @@ impl Agent {
|
||||
.get_or_create_session(&message.user_id)
|
||||
.await;
|
||||
let mut sess = session.lock().await;
|
||||
let thread = sess.create_thread();
|
||||
let thread = sess.create_thread(Some(&message.channel));
|
||||
let thread_id = thread.id;
|
||||
Ok(SubmissionResult::ok_with_message(format!(
|
||||
"New thread: {}",
|
||||
@@ -2117,7 +2133,7 @@ mod tests {
|
||||
|
||||
let session_id = Uuid::new_v4();
|
||||
let thread_id = Uuid::new_v4();
|
||||
let mut thread = Thread::with_id(thread_id, session_id);
|
||||
let mut thread = Thread::with_id(thread_id, session_id, None);
|
||||
|
||||
// Set thread to AwaitingApproval with a pending tool approval
|
||||
let pending = PendingApproval {
|
||||
@@ -2185,7 +2201,7 @@ mod tests {
|
||||
use crate::agent::session::{MAX_PENDING_MESSAGES, Thread, ThreadState};
|
||||
use uuid::Uuid;
|
||||
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
thread.start_turn("processing something");
|
||||
assert_eq!(thread.state, ThreadState::Processing);
|
||||
|
||||
@@ -2211,7 +2227,7 @@ mod tests {
|
||||
use crate::agent::session::{Thread, ThreadState};
|
||||
use uuid::Uuid;
|
||||
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
let mut thread = Thread::new(Uuid::new_v4(), None);
|
||||
thread.start_turn("processing");
|
||||
|
||||
thread.queue_message("pending-1".to_string());
|
||||
@@ -2241,7 +2257,7 @@ mod tests {
|
||||
|
||||
let thread_id = Uuid::new_v4();
|
||||
let session_id = Uuid::new_v4();
|
||||
let mut thread = Thread::with_id(thread_id, session_id);
|
||||
let mut thread = Thread::with_id(thread_id, session_id, None);
|
||||
thread.start_turn("working");
|
||||
assert_eq!(thread.state, ThreadState::Processing);
|
||||
|
||||
@@ -2268,7 +2284,7 @@ mod tests {
|
||||
|
||||
let thread_id = Uuid::new_v4();
|
||||
let session_id = Uuid::new_v4();
|
||||
let mut thread = Thread::with_id(thread_id, session_id);
|
||||
let mut thread = Thread::with_id(thread_id, session_id, None);
|
||||
thread.start_turn("working");
|
||||
assert_eq!(thread.state, ThreadState::Processing);
|
||||
|
||||
|
||||
@@ -34,6 +34,7 @@ pub async fn setup_wasm_channels(
|
||||
secrets_store: &Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||
extension_manager: Option<&Arc<ExtensionManager>>,
|
||||
database: Option<&Arc<dyn Database>>,
|
||||
registered_channel_names: &[String],
|
||||
) -> Option<WasmChannelSetup> {
|
||||
let runtime = match WasmChannelRuntime::new(WasmChannelRuntimeConfig::default()) {
|
||||
Ok(r) => Arc::new(r),
|
||||
@@ -71,7 +72,46 @@ pub async fn setup_wasm_channels(
|
||||
let mut channels: Vec<(String, Box<dyn crate::channels::Channel>)> = Vec::new();
|
||||
let mut channel_names: Vec<String> = Vec::new();
|
||||
|
||||
// Reserved channel names that WASM modules must not claim.
|
||||
// A malicious module could otherwise register as a trusted built-in
|
||||
// channel and bypass cross-channel authorization checks.
|
||||
// This list must cover every built-in channel name to prevent a WASM
|
||||
// module from impersonating a built-in and satisfying same-channel
|
||||
// approval checks.
|
||||
const RESERVED_CHANNEL_NAMES: &[&str] = &[
|
||||
"web",
|
||||
"gateway",
|
||||
"cli",
|
||||
"repl",
|
||||
"http",
|
||||
"signal",
|
||||
"slack-relay",
|
||||
"secret_save",
|
||||
];
|
||||
|
||||
for loaded in results.loaded {
|
||||
let name_lower = loaded.name().to_ascii_lowercase();
|
||||
if RESERVED_CHANNEL_NAMES.contains(&name_lower.as_str()) {
|
||||
tracing::warn!(
|
||||
channel = %loaded.name(),
|
||||
"Rejected WASM channel with reserved name"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
// Also reject any name that collides with an already-registered
|
||||
// channel to prevent a WASM module from shadowing a channel that
|
||||
// was registered earlier in the startup sequence.
|
||||
if registered_channel_names
|
||||
.iter()
|
||||
.any(|n| n.to_ascii_lowercase() == name_lower)
|
||||
{
|
||||
tracing::warn!(
|
||||
channel = %loaded.name(),
|
||||
"Rejected WASM channel that collides with already-registered channel"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
let (name, channel) = register_channel(
|
||||
loaded,
|
||||
config,
|
||||
|
||||
@@ -574,7 +574,7 @@ pub async fn chat_new_thread_handler(
|
||||
.await;
|
||||
let (thread_id, info) = {
|
||||
let mut sess = session.lock().await;
|
||||
let thread = sess.create_thread();
|
||||
let thread = sess.create_thread(Some("web"));
|
||||
let id = thread.id;
|
||||
let info = ThreadInfo {
|
||||
id: thread.id,
|
||||
@@ -593,7 +593,13 @@ pub async fn chat_new_thread_handler(
|
||||
// so that the subsequent loadThreads() call from the frontend sees it.
|
||||
if let Some(ref store) = state.store {
|
||||
match store
|
||||
.ensure_conversation(thread_id, "gateway", &identity.user_id, None)
|
||||
.ensure_conversation(
|
||||
thread_id,
|
||||
"gateway",
|
||||
&identity.user_id,
|
||||
None,
|
||||
Some("gateway"),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(true) => {}
|
||||
|
||||
@@ -15,6 +15,14 @@ use crate::channels::web::auth::AuthenticatedUser;
|
||||
use crate::channels::web::server::GatewayState;
|
||||
use crate::channels::web::types::*;
|
||||
|
||||
fn db_error(context: &str, e: impl std::fmt::Display) -> (StatusCode, String) {
|
||||
tracing::error!(%e, context, "Database error in jobs handler");
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"Internal database error".to_string(),
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn jobs_list_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
AuthenticatedUser(user): AuthenticatedUser,
|
||||
@@ -213,10 +221,7 @@ pub async fn jobs_detail_handler(
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(e) => {
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Database error: {}", e),
|
||||
));
|
||||
return Err(db_error("jobs_handler", e));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -257,10 +262,7 @@ pub async fn jobs_detail_handler(
|
||||
}))
|
||||
}
|
||||
Ok(None) => Err((StatusCode::NOT_FOUND, "Job not found".to_string())),
|
||||
Err(e) => Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Database error: {}", e),
|
||||
)),
|
||||
Err(e) => Err(db_error("jobs_handler", e)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -304,10 +306,7 @@ pub async fn jobs_cancel_handler(
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(e) => {
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Database error: {}", e),
|
||||
));
|
||||
return Err(db_error("jobs_handler", e));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -350,10 +349,7 @@ pub async fn jobs_cancel_handler(
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(e) => {
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Database error: {}", e),
|
||||
));
|
||||
return Err(db_error("jobs_handler", e));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -471,10 +467,7 @@ pub async fn jobs_restart_handler(
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(e) => {
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Database error: {}", e),
|
||||
));
|
||||
return Err(db_error("jobs_handler", e));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -530,10 +523,7 @@ pub async fn jobs_restart_handler(
|
||||
})))
|
||||
}
|
||||
Ok(None) => Err((StatusCode::NOT_FOUND, "Job not found".to_string())),
|
||||
Err(e) => Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Database error: {}", e),
|
||||
)),
|
||||
Err(e) => Err(db_error("jobs_handler", e)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -609,10 +599,7 @@ pub async fn jobs_prompt_handler(
|
||||
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||
}
|
||||
Err(e) => {
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Database error: {}", e),
|
||||
));
|
||||
return Err(db_error("jobs_handler", e));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -667,10 +654,7 @@ pub async fn jobs_events_handler(
|
||||
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
|
||||
}
|
||||
Err(e) => {
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Database error: {}", e),
|
||||
));
|
||||
return Err(db_error("jobs_handler", e));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -823,3 +807,17 @@ pub async fn job_files_read_handler(
|
||||
content,
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_db_error_does_not_leak_details() {
|
||||
let (status, body) = db_error("test_context", "relation \"jobs\" does not exist");
|
||||
assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
|
||||
assert_eq!(body, "Internal database error");
|
||||
assert!(!body.contains("relation"));
|
||||
assert!(!body.contains("does not exist"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2014,7 +2014,7 @@ async fn chat_new_thread_handler(
|
||||
let session = session_manager.get_or_create_session(&user.user_id).await;
|
||||
let (thread_id, info) = {
|
||||
let mut sess = session.lock().await;
|
||||
let thread = sess.create_thread();
|
||||
let thread = sess.create_thread(Some("gateway"));
|
||||
let id = thread.id;
|
||||
let info = ThreadInfo {
|
||||
id: thread.id,
|
||||
@@ -2033,7 +2033,7 @@ async fn chat_new_thread_handler(
|
||||
// so that the subsequent loadThreads() call from the frontend sees it.
|
||||
if let Some(ref store) = state.store {
|
||||
match store
|
||||
.ensure_conversation(thread_id, "gateway", &user.user_id, None)
|
||||
.ensure_conversation(thread_id, "gateway", &user.user_id, None, Some("gateway"))
|
||||
.await
|
||||
{
|
||||
Ok(true) => {}
|
||||
|
||||
+122
-4
@@ -569,6 +569,42 @@ pub async fn sweep_expired_flows(registry: &PendingOAuthRegistry) {
|
||||
const HOSTED_STATE_PREFIX: &str = "ic2";
|
||||
const HOSTED_STATE_CHECKSUM_BYTES: usize = 12;
|
||||
|
||||
/// Maximum length for a legacy flow ID or instance name.
|
||||
const LEGACY_STATE_MAX_LEN: usize = 128;
|
||||
/// Minimum length for a legacy flow ID.
|
||||
const LEGACY_STATE_MIN_LEN: usize = 8;
|
||||
|
||||
/// Validate that a legacy state component (flow_id or instance_name) contains
|
||||
/// only safe characters: alphanumeric, dash, underscore.
|
||||
fn is_valid_legacy_state_component(s: &str) -> bool {
|
||||
!s.is_empty()
|
||||
&& s.len() <= LEGACY_STATE_MAX_LEN
|
||||
&& s.bytes()
|
||||
.all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_')
|
||||
}
|
||||
|
||||
fn validate_legacy_flow_id(flow_id: &str) -> Result<(), String> {
|
||||
if flow_id.len() < LEGACY_STATE_MIN_LEN {
|
||||
return Err(format!(
|
||||
"Legacy OAuth flow_id too short ({} chars, minimum {LEGACY_STATE_MIN_LEN})",
|
||||
flow_id.len()
|
||||
));
|
||||
}
|
||||
if flow_id.len() > LEGACY_STATE_MAX_LEN {
|
||||
return Err(format!(
|
||||
"Legacy OAuth flow_id too long ({} chars, maximum {LEGACY_STATE_MAX_LEN})",
|
||||
flow_id.len()
|
||||
));
|
||||
}
|
||||
if !flow_id
|
||||
.bytes()
|
||||
.all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_')
|
||||
{
|
||||
return Err("Legacy OAuth flow_id contains invalid characters".to_string());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct DecodedHostedOAuthState {
|
||||
pub flow_id: String,
|
||||
@@ -653,6 +689,17 @@ pub fn decode_hosted_oauth_state(state: &str) -> Result<DecodedHostedOAuthState,
|
||||
if flow_id.is_empty() {
|
||||
return Err("Hosted OAuth legacy state is missing flow_id".to_string());
|
||||
}
|
||||
validate_legacy_flow_id(flow_id)?;
|
||||
if !instance_name.is_empty() && !is_valid_legacy_state_component(instance_name) {
|
||||
return Err(format!(
|
||||
"Legacy OAuth instance name contains invalid characters or exceeds max length ({LEGACY_STATE_MAX_LEN})"
|
||||
));
|
||||
}
|
||||
tracing::debug!(
|
||||
flow_id,
|
||||
instance_name,
|
||||
"Decoded legacy prefixed OAuth state"
|
||||
);
|
||||
return Ok(DecodedHostedOAuthState {
|
||||
flow_id: flow_id.to_string(),
|
||||
instance_name: if instance_name.is_empty() {
|
||||
@@ -668,6 +715,9 @@ pub fn decode_hosted_oauth_state(state: &str) -> Result<DecodedHostedOAuthState,
|
||||
return Err("Hosted OAuth state is empty".to_string());
|
||||
}
|
||||
|
||||
validate_legacy_flow_id(state)?;
|
||||
tracing::debug!(flow_id = state, "Decoded legacy raw OAuth state");
|
||||
|
||||
Ok(DecodedHostedOAuthState {
|
||||
flow_id: state.to_string(),
|
||||
instance_name: None,
|
||||
@@ -1734,13 +1784,13 @@ mod tests {
|
||||
fn test_decode_hosted_oauth_state_accepts_legacy_formats() {
|
||||
use crate::cli::oauth_defaults::decode_hosted_oauth_state;
|
||||
|
||||
let decoded = decode_hosted_oauth_state("kind-deer:abc123").expect("legacy prefixed");
|
||||
assert_eq!(decoded.flow_id, "abc123");
|
||||
let decoded = decode_hosted_oauth_state("kind-deer:abc12345").expect("legacy prefixed");
|
||||
assert_eq!(decoded.flow_id, "abc12345");
|
||||
assert_eq!(decoded.instance_name.as_deref(), Some("kind-deer"));
|
||||
assert!(decoded.is_legacy);
|
||||
|
||||
let decoded = decode_hosted_oauth_state("abc123").expect("legacy raw");
|
||||
assert_eq!(decoded.flow_id, "abc123");
|
||||
let decoded = decode_hosted_oauth_state("abc12345").expect("legacy raw");
|
||||
assert_eq!(decoded.flow_id, "abc12345");
|
||||
assert_eq!(decoded.instance_name, None);
|
||||
assert!(decoded.is_legacy);
|
||||
}
|
||||
@@ -1864,4 +1914,72 @@ mod tests {
|
||||
assert_eq!(decoded_no_instance.instance_name, None);
|
||||
assert!(!decoded_no_instance.is_legacy);
|
||||
}
|
||||
|
||||
/// Legacy flow IDs that are too short must be rejected (#1443).
|
||||
#[test]
|
||||
fn test_legacy_state_rejects_short_flow_id() {
|
||||
use crate::cli::oauth_defaults::decode_hosted_oauth_state;
|
||||
|
||||
let err = decode_hosted_oauth_state("abc").expect_err("short raw flow_id");
|
||||
assert!(err.contains("too short"), "unexpected error: {err}");
|
||||
|
||||
let err = decode_hosted_oauth_state("inst:abc").expect_err("short prefixed flow_id");
|
||||
assert!(err.contains("too short"), "unexpected error: {err}");
|
||||
}
|
||||
|
||||
/// Legacy flow IDs with invalid characters must be rejected (#1443).
|
||||
#[test]
|
||||
fn test_legacy_state_rejects_invalid_characters() {
|
||||
use crate::cli::oauth_defaults::decode_hosted_oauth_state;
|
||||
|
||||
let err = decode_hosted_oauth_state("flow id with spaces!").expect_err("spaces in flow_id");
|
||||
assert!(
|
||||
err.contains("invalid characters"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
|
||||
let err = decode_hosted_oauth_state("inst:flow/id?bad=yes")
|
||||
.expect_err("special chars in prefixed flow_id");
|
||||
assert!(
|
||||
err.contains("invalid characters"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
/// Legacy instance names with invalid characters must be rejected (#1444).
|
||||
#[test]
|
||||
fn test_legacy_state_rejects_invalid_instance_name() {
|
||||
use crate::cli::oauth_defaults::decode_hosted_oauth_state;
|
||||
|
||||
let err = decode_hosted_oauth_state("bad instance!:valid-flow-id-12345")
|
||||
.expect_err("invalid instance name");
|
||||
assert!(err.contains("instance name"), "unexpected error: {err}");
|
||||
}
|
||||
|
||||
/// Excessively long legacy flow IDs must be rejected (#1443).
|
||||
#[test]
|
||||
fn test_legacy_state_rejects_oversized_flow_id() {
|
||||
use crate::cli::oauth_defaults::decode_hosted_oauth_state;
|
||||
|
||||
let long_id = "a".repeat(200);
|
||||
let err = decode_hosted_oauth_state(&long_id).expect_err("oversized flow_id");
|
||||
assert!(err.contains("too long"), "unexpected error: {err}");
|
||||
}
|
||||
|
||||
/// Valid legacy flow IDs at boundary lengths are accepted.
|
||||
#[test]
|
||||
fn test_legacy_state_accepts_boundary_lengths() {
|
||||
use crate::cli::oauth_defaults::decode_hosted_oauth_state;
|
||||
|
||||
// Exactly 8 chars (minimum)
|
||||
let decoded = decode_hosted_oauth_state("abcd1234").expect("8-char flow_id");
|
||||
assert_eq!(decoded.flow_id, "abcd1234");
|
||||
assert!(decoded.is_legacy);
|
||||
|
||||
// Exactly 128 chars (maximum)
|
||||
let max_id = "a".repeat(128);
|
||||
let decoded = decode_hosted_oauth_state(&max_id).expect("128-char flow_id");
|
||||
assert_eq!(decoded.flow_id, max_id);
|
||||
assert!(decoded.is_legacy);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -67,19 +67,20 @@ impl ConversationStore for LibSqlBackend {
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
source_channel: Option<&str>,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let now = fmt_ts(&Utc::now());
|
||||
let affected = conn
|
||||
.execute(
|
||||
r#"
|
||||
INSERT INTO conversations (id, channel, user_id, thread_id, started_at, last_activity)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?5)
|
||||
INSERT INTO conversations (id, channel, user_id, thread_id, source_channel, started_at, last_activity)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?6)
|
||||
ON CONFLICT (id) DO UPDATE SET last_activity = excluded.last_activity
|
||||
WHERE conversations.user_id = excluded.user_id
|
||||
AND conversations.channel = excluded.channel
|
||||
"#,
|
||||
params![id.to_string(), channel, user_id, opt_text(thread_id), now],
|
||||
params![id.to_string(), channel, user_id, opt_text(thread_id), opt_text(source_channel), now],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
@@ -565,6 +566,28 @@ impl ConversationStore for LibSqlBackend {
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
Ok(found.is_some())
|
||||
}
|
||||
|
||||
async fn get_conversation_source_channel(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
) -> Result<Option<String>, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
"SELECT source_channel FROM conversations WHERE id = ?1",
|
||||
params![conversation_id.to_string()],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
match rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
{
|
||||
Some(row) => Ok(get_opt_text(&row, 0)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
+17
-13
@@ -17,7 +17,6 @@ mod workspace;
|
||||
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, NaiveDateTime, Utc};
|
||||
@@ -34,8 +33,6 @@ use crate::workspace::MemoryDocument;
|
||||
|
||||
use crate::db::libsql_migrations;
|
||||
|
||||
static NAIVE_TIMESTAMP_LOGGED: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
/// Explicit column list for routines table (matches positional access in `row_to_routine_libsql`).
|
||||
pub(crate) const ROUTINE_COLUMNS: &str = "\
|
||||
id, name, description, user_id, enabled, \
|
||||
@@ -167,13 +164,11 @@ impl LibSqlBackend {
|
||||
///
|
||||
/// Returns an error if none of the formats match.
|
||||
pub(crate) fn parse_timestamp(s: &str) -> Result<DateTime<Utc>, String> {
|
||||
let log_naive_timestamp_once = || {
|
||||
if !NAIVE_TIMESTAMP_LOGGED.swap(true, Ordering::Relaxed) {
|
||||
tracing::debug!(
|
||||
timestamp = %s,
|
||||
"parsed naive timestamp without timezone; assuming UTC for backward compatibility"
|
||||
);
|
||||
}
|
||||
let log_naive_timestamp = || {
|
||||
tracing::warn!(
|
||||
timestamp = %s,
|
||||
"parsed naive timestamp, assuming UTC — consider migrating to RFC 3339"
|
||||
);
|
||||
};
|
||||
|
||||
// RFC 3339 (our canonical write format)
|
||||
@@ -182,12 +177,12 @@ pub(crate) fn parse_timestamp(s: &str) -> Result<DateTime<Utc>, String> {
|
||||
}
|
||||
// Naive with fractional seconds (legacy or SQLite datetime() output)
|
||||
if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f") {
|
||||
log_naive_timestamp_once();
|
||||
log_naive_timestamp();
|
||||
return Ok(ndt.and_utc());
|
||||
}
|
||||
// Naive without fractional seconds (legacy format)
|
||||
if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
|
||||
log_naive_timestamp_once();
|
||||
log_naive_timestamp();
|
||||
return Ok(ndt.and_utc());
|
||||
}
|
||||
Err(format!("unparseable timestamp: {:?}", s))
|
||||
@@ -439,7 +434,7 @@ mod tests {
|
||||
use chrono::{TimeZone, Utc};
|
||||
|
||||
use crate::db::Database;
|
||||
use crate::db::libsql::{LibSqlBackend, normalize_notify_user, parse_timestamp};
|
||||
use crate::db::libsql::{LibSqlBackend, fmt_ts, normalize_notify_user, parse_timestamp};
|
||||
|
||||
#[test]
|
||||
fn test_normalize_notify_user_treats_legacy_default_as_missing() {
|
||||
@@ -468,6 +463,15 @@ mod tests {
|
||||
assert_eq!(naive_without_millis, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fmt_ts_roundtrips_through_parse_timestamp() {
|
||||
let original = Utc.with_ymd_and_hms(2026, 6, 15, 8, 30, 45).unwrap()
|
||||
+ chrono::Duration::milliseconds(123);
|
||||
let formatted = fmt_ts(&original);
|
||||
let parsed = parse_timestamp(&formatted).unwrap();
|
||||
assert_eq!(parsed, original);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_libsql_now_format_is_rfc3339_and_parseable() {
|
||||
let backend = LibSqlBackend::new_memory().await.unwrap();
|
||||
|
||||
@@ -981,7 +981,7 @@ mod tests {
|
||||
assert!(alice_stats.last_active_at.is_some());
|
||||
|
||||
// Bob has no LLM calls so doesn't appear in summary stats
|
||||
assert!(stats.iter().find(|s| s.user_id == "bob").is_none());
|
||||
assert!(!stats.iter().any(|s| s.user_id == "bob"));
|
||||
|
||||
// Filter to single user
|
||||
let alice_only = db.user_summary_stats(Some("alice")).await.unwrap();
|
||||
|
||||
@@ -785,6 +785,14 @@ CREATE TABLE IF NOT EXISTS api_tokens (
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_api_tokens_user ON api_tokens(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_api_tokens_hash ON api_tokens(token_hash);
|
||||
"#,
|
||||
),
|
||||
(
|
||||
15,
|
||||
"conversation_source_channel",
|
||||
// Add source_channel to conversations for cross-channel approval authorization.
|
||||
r#"
|
||||
ALTER TABLE conversations ADD COLUMN source_channel TEXT;
|
||||
"#,
|
||||
),
|
||||
];
|
||||
|
||||
@@ -373,6 +373,7 @@ pub trait ConversationStore: Send + Sync {
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
source_channel: Option<&str>,
|
||||
) -> Result<bool, DatabaseError>;
|
||||
async fn list_conversations_with_preview(
|
||||
&self,
|
||||
@@ -431,6 +432,11 @@ pub trait ConversationStore: Send + Sync {
|
||||
conversation_id: Uuid,
|
||||
user_id: &str,
|
||||
) -> Result<bool, DatabaseError>;
|
||||
/// Get the source_channel for a conversation (the channel that created it).
|
||||
async fn get_conversation_source_channel(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
) -> Result<Option<String>, DatabaseError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
|
||||
+11
-1
@@ -99,9 +99,10 @@ impl ConversationStore for PgBackend {
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
source_channel: Option<&str>,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
self.store
|
||||
.ensure_conversation(id, channel, user_id, thread_id)
|
||||
.ensure_conversation(id, channel, user_id, thread_id, source_channel)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -212,6 +213,15 @@ impl ConversationStore for PgBackend {
|
||||
.conversation_belongs_to_user(conversation_id, user_id)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_conversation_source_channel(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
) -> Result<Option<String>, DatabaseError> {
|
||||
self.store
|
||||
.get_conversation_source_channel(conversation_id)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== JobStore ====================
|
||||
|
||||
+649
-1
@@ -53,6 +53,38 @@ struct HostedOAuthFlowStart {
|
||||
flow: crate::cli::oauth_defaults::PendingOAuthFlow,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct SecretCleanupPlan {
|
||||
base_secrets: HashSet<String>,
|
||||
companion_secrets: HashMap<String, HashSet<String>>,
|
||||
}
|
||||
|
||||
impl SecretCleanupPlan {
|
||||
fn add_base_secret(&mut self, secret_name: impl AsRef<str>) {
|
||||
self.base_secrets
|
||||
.insert(secret_name.as_ref().to_lowercase());
|
||||
}
|
||||
|
||||
fn add_companion_secret(
|
||||
&mut self,
|
||||
base_secret_name: impl AsRef<str>,
|
||||
companion_secret_name: impl AsRef<str>,
|
||||
) {
|
||||
self.companion_secrets
|
||||
.entry(base_secret_name.as_ref().to_lowercase())
|
||||
.or_default()
|
||||
.insert(companion_secret_name.as_ref().to_lowercase());
|
||||
}
|
||||
}
|
||||
|
||||
fn oauth_refresh_secret_name(secret_name: &str) -> String {
|
||||
format!("{}_refresh_token", secret_name.to_lowercase())
|
||||
}
|
||||
|
||||
fn oauth_scopes_secret_name(secret_name: &str) -> String {
|
||||
format!("{}_scopes", secret_name.to_lowercase())
|
||||
}
|
||||
|
||||
fn normalize_oauth_callback_path(path: &str) -> String {
|
||||
let trimmed_path = path.trim_end_matches('/');
|
||||
if trimmed_path.is_empty() {
|
||||
@@ -1602,6 +1634,10 @@ impl ExtensionManager {
|
||||
|
||||
match kind {
|
||||
ExtensionKind::McpServer => {
|
||||
let cleanup_plan = self
|
||||
.collect_secret_cleanup_plan(name, kind, user_id)
|
||||
.await?;
|
||||
|
||||
// Unregister tools with this server's prefix
|
||||
let tool_names: Vec<String> = self
|
||||
.tool_registry
|
||||
@@ -1623,6 +1659,9 @@ impl ExtensionManager {
|
||||
.await
|
||||
.map_err(|e| ExtensionError::Config(e.to_string()))?;
|
||||
|
||||
self.cleanup_uninstalled_extension_secrets(cleanup_plan, user_id)
|
||||
.await;
|
||||
|
||||
Ok(format!(
|
||||
"Removed MCP server '{}' and {} tool(s)",
|
||||
name,
|
||||
@@ -1630,6 +1669,10 @@ impl ExtensionManager {
|
||||
))
|
||||
}
|
||||
ExtensionKind::WasmTool => {
|
||||
let cleanup_plan = self
|
||||
.collect_secret_cleanup_plan(name, kind, user_id)
|
||||
.await?;
|
||||
|
||||
// Unregister from tool registry
|
||||
self.tool_registry.unregister(name).await;
|
||||
|
||||
@@ -1674,9 +1717,16 @@ impl ExtensionManager {
|
||||
let _ = tokio::fs::remove_file(&cap_path).await;
|
||||
}
|
||||
|
||||
self.cleanup_uninstalled_extension_secrets(cleanup_plan, user_id)
|
||||
.await;
|
||||
|
||||
Ok(format!("Removed WASM tool '{}'", name))
|
||||
}
|
||||
ExtensionKind::WasmChannel => {
|
||||
let cleanup_plan = self
|
||||
.collect_secret_cleanup_plan(name, kind, user_id)
|
||||
.await?;
|
||||
|
||||
// Remove from active set and persist
|
||||
self.active_channel_names.write().await.remove(name);
|
||||
self.persist_active_channels(user_id).await;
|
||||
@@ -1702,6 +1752,9 @@ impl ExtensionManager {
|
||||
let _ = tokio::fs::remove_file(&cap_path).await;
|
||||
}
|
||||
|
||||
self.cleanup_uninstalled_extension_secrets(cleanup_plan, user_id)
|
||||
.await;
|
||||
|
||||
Ok(format!(
|
||||
"Removed channel '{}'. Restart IronClaw for the change to take effect.",
|
||||
name
|
||||
@@ -2999,6 +3052,258 @@ impl ExtensionManager {
|
||||
crate::tools::wasm::CapabilitiesFile::from_bytes(&cap_bytes).ok()
|
||||
}
|
||||
|
||||
async fn load_channel_capabilities(
|
||||
&self,
|
||||
name: &str,
|
||||
) -> Option<crate::channels::wasm::ChannelCapabilitiesFile> {
|
||||
let cap_path = self
|
||||
.wasm_channels_dir
|
||||
.join(format!("{}.capabilities.json", name));
|
||||
let cap_bytes = tokio::fs::read(&cap_path).await.ok()?;
|
||||
crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes).ok()
|
||||
}
|
||||
|
||||
async fn collect_secret_cleanup_plan(
|
||||
&self,
|
||||
name: &str,
|
||||
kind: ExtensionKind,
|
||||
user_id: &str,
|
||||
) -> Result<SecretCleanupPlan, ExtensionError> {
|
||||
let mut plan = SecretCleanupPlan::default();
|
||||
|
||||
match kind {
|
||||
ExtensionKind::WasmTool => {
|
||||
if let Some(cap) = self.load_tool_capabilities(name).await {
|
||||
for secret_name in Self::tool_secret_names(&cap) {
|
||||
plan.add_base_secret(secret_name);
|
||||
}
|
||||
|
||||
if let Some(auth) = cap.auth {
|
||||
plan.add_base_secret(&auth.secret_name);
|
||||
plan.add_companion_secret(
|
||||
&auth.secret_name,
|
||||
oauth_refresh_secret_name(&auth.secret_name),
|
||||
);
|
||||
plan.add_companion_secret(
|
||||
&auth.secret_name,
|
||||
oauth_scopes_secret_name(&auth.secret_name),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
ExtensionKind::WasmChannel => {
|
||||
if let Some(cap) = self.load_channel_capabilities(name).await {
|
||||
for secret_name in Self::channel_secret_names(&cap) {
|
||||
plan.add_base_secret(secret_name);
|
||||
}
|
||||
}
|
||||
}
|
||||
ExtensionKind::McpServer => {
|
||||
let server = self
|
||||
.get_mcp_server(name, user_id)
|
||||
.await
|
||||
.map_err(|e| ExtensionError::Config(e.to_string()))?;
|
||||
let token_secret_name = server.token_secret_name();
|
||||
plan.add_base_secret(&token_secret_name);
|
||||
plan.add_base_secret(server.client_id_secret_name());
|
||||
// MCP OAuth can persist companion secrets through two paths:
|
||||
// the MCP auth helper uses `mcp_<name>_refresh_token`, while the
|
||||
// hosted gateway callback stores companions alongside the access
|
||||
// token secret (`<token_secret>_refresh_token` / `_scopes`).
|
||||
plan.add_companion_secret(&token_secret_name, server.refresh_token_secret_name());
|
||||
plan.add_companion_secret(
|
||||
&token_secret_name,
|
||||
oauth_refresh_secret_name(&token_secret_name),
|
||||
);
|
||||
plan.add_companion_secret(
|
||||
&token_secret_name,
|
||||
oauth_scopes_secret_name(&token_secret_name),
|
||||
);
|
||||
}
|
||||
ExtensionKind::ChannelRelay => {}
|
||||
}
|
||||
|
||||
Ok(plan)
|
||||
}
|
||||
|
||||
async fn cleanup_uninstalled_extension_secrets(&self, plan: SecretCleanupPlan, user_id: &str) {
|
||||
let referenced_secrets = match self.collect_referenced_secret_names(user_id).await {
|
||||
Ok(secret_names) => secret_names,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
user_id,
|
||||
error,
|
||||
"Failed to determine which secrets are still referenced; keeping secrets"
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
for base_secret in &plan.base_secrets {
|
||||
if referenced_secrets.contains(base_secret) {
|
||||
continue;
|
||||
}
|
||||
|
||||
self.delete_secret_best_effort(user_id, base_secret).await;
|
||||
|
||||
if let Some(companion_secrets) = plan.companion_secrets.get(base_secret) {
|
||||
for companion_secret in companion_secrets {
|
||||
if !referenced_secrets.contains(companion_secret) {
|
||||
self.delete_secret_best_effort(user_id, companion_secret)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_secret_best_effort(&self, user_id: &str, secret_name: &str) {
|
||||
if let Err(error) = self.secrets.delete(user_id, secret_name).await {
|
||||
tracing::warn!(
|
||||
user_id,
|
||||
secret_name,
|
||||
error = %error,
|
||||
"Failed to delete secret while uninstalling extension"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async fn collect_referenced_secret_names(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<HashSet<String>, String> {
|
||||
let mut referenced_secret_names = HashSet::new();
|
||||
|
||||
let tools = discover_tools(&self.wasm_tools_dir)
|
||||
.await
|
||||
.map_err(|e| format!("discover tools: {e}"))?;
|
||||
for (tool_name, discovered_tool) in &tools {
|
||||
let cap = self
|
||||
.load_tool_capabilities(tool_name)
|
||||
.await
|
||||
.ok_or_else(|| {
|
||||
let path = discovered_tool
|
||||
.capabilities_path
|
||||
.as_ref()
|
||||
.map(|path| path.display().to_string())
|
||||
.unwrap_or_else(|| format!("{} (missing)", tool_name));
|
||||
format!("load tool capabilities for {tool_name}: {path}")
|
||||
})?;
|
||||
referenced_secret_names.extend(Self::tool_secret_names(&cap));
|
||||
}
|
||||
|
||||
let channels = crate::channels::wasm::discover_channels(&self.wasm_channels_dir)
|
||||
.await
|
||||
.map_err(|e| format!("discover channels: {e}"))?;
|
||||
for (channel_name, discovered_channel) in &channels {
|
||||
let cap = self
|
||||
.load_channel_capabilities(channel_name)
|
||||
.await
|
||||
.ok_or_else(|| {
|
||||
let path = discovered_channel
|
||||
.capabilities_path
|
||||
.as_ref()
|
||||
.map(|path| path.display().to_string())
|
||||
.unwrap_or_else(|| format!("{} (missing)", channel_name));
|
||||
format!("load channel capabilities for {channel_name}: {path}")
|
||||
})?;
|
||||
referenced_secret_names.extend(Self::channel_secret_names(&cap));
|
||||
}
|
||||
|
||||
let mcp_servers = self
|
||||
.load_mcp_servers(user_id)
|
||||
.await
|
||||
.map_err(|e| format!("load MCP servers: {e}"))?;
|
||||
for server in &mcp_servers.servers {
|
||||
referenced_secret_names.extend(Self::mcp_server_secret_names(server));
|
||||
}
|
||||
|
||||
Ok(referenced_secret_names)
|
||||
}
|
||||
|
||||
fn tool_secret_names(cap: &crate::tools::wasm::CapabilitiesFile) -> HashSet<String> {
|
||||
let mut names = HashSet::new();
|
||||
|
||||
if let Some(auth) = &cap.auth {
|
||||
names.insert(auth.secret_name.to_lowercase());
|
||||
}
|
||||
if let Some(setup) = &cap.setup {
|
||||
names.extend(
|
||||
setup
|
||||
.required_secrets
|
||||
.iter()
|
||||
.map(|secret| secret.name.to_lowercase()),
|
||||
);
|
||||
}
|
||||
if let Some(http) = &cap.http {
|
||||
names.extend(
|
||||
http.credentials
|
||||
.values()
|
||||
.map(|credential| credential.secret_name.to_lowercase()),
|
||||
);
|
||||
}
|
||||
if let Some(webhook) = &cap.webhook {
|
||||
if let Some(secret_name) = &webhook.secret_name {
|
||||
names.insert(secret_name.to_lowercase());
|
||||
}
|
||||
if let Some(secret_name) = &webhook.signature_key_secret_name {
|
||||
names.insert(secret_name.to_lowercase());
|
||||
}
|
||||
if let Some(secret_name) = &webhook.hmac_secret_name {
|
||||
names.insert(secret_name.to_lowercase());
|
||||
}
|
||||
}
|
||||
|
||||
names
|
||||
}
|
||||
|
||||
fn channel_secret_names(
|
||||
cap: &crate::channels::wasm::ChannelCapabilitiesFile,
|
||||
) -> HashSet<String> {
|
||||
let mut names: HashSet<String> = cap
|
||||
.setup
|
||||
.required_secrets
|
||||
.iter()
|
||||
.map(|secret| secret.name.to_lowercase())
|
||||
.collect();
|
||||
|
||||
if let Some(http) = cap.capabilities.tool.http.as_ref() {
|
||||
names.extend(
|
||||
http.credentials
|
||||
.values()
|
||||
.map(|credential| credential.secret_name.to_lowercase()),
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(webhook) = cap
|
||||
.capabilities
|
||||
.channel
|
||||
.as_ref()
|
||||
.and_then(|channel| channel.webhook.as_ref())
|
||||
{
|
||||
if webhook.secret_header.is_some() || webhook.secret_name.is_some() {
|
||||
names.insert(cap.webhook_secret_name().to_lowercase());
|
||||
}
|
||||
if let Some(secret_name) = cap.signature_key_secret_name() {
|
||||
names.insert(secret_name.to_lowercase());
|
||||
}
|
||||
if let Some(secret_name) = cap.hmac_secret_name() {
|
||||
names.insert(secret_name.to_lowercase());
|
||||
}
|
||||
}
|
||||
|
||||
names
|
||||
}
|
||||
|
||||
fn mcp_server_secret_names(server: &McpServerConfig) -> HashSet<String> {
|
||||
[
|
||||
server.token_secret_name().to_lowercase(),
|
||||
server.client_id_secret_name().to_lowercase(),
|
||||
]
|
||||
.into_iter()
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Collect merged OAuth scopes from all installed tools sharing the same secret_name.
|
||||
///
|
||||
/// When multiple tools share an OAuth provider (e.g., google-calendar and google-drive
|
||||
@@ -6033,6 +6338,8 @@ mod tests {
|
||||
ExtensionError, ExtensionKind, ExtensionSource, InstallResult, VerificationChallenge,
|
||||
};
|
||||
use crate::pairing::PairingStore;
|
||||
use crate::secrets::CreateSecretParams;
|
||||
use crate::tools::mcp::McpServerConfig;
|
||||
|
||||
fn require(condition: bool, message: impl Into<String>) -> Result<(), String> {
|
||||
if condition {
|
||||
@@ -6353,6 +6660,38 @@ mod tests {
|
||||
tools_dir
|
||||
}
|
||||
|
||||
fn write_test_channel(
|
||||
dir: &std::path::Path,
|
||||
name: &str,
|
||||
capabilities_json: &str,
|
||||
) -> std::path::PathBuf {
|
||||
let channels_dir = dir.join("channels");
|
||||
std::fs::create_dir_all(&channels_dir).expect("channels dir");
|
||||
std::fs::write(
|
||||
channels_dir.join(format!("{name}.wasm")),
|
||||
b"not-a-real-wasm",
|
||||
)
|
||||
.expect("wasm");
|
||||
std::fs::write(
|
||||
channels_dir.join(format!("{name}.capabilities.json")),
|
||||
capabilities_json,
|
||||
)
|
||||
.expect("capabilities");
|
||||
channels_dir
|
||||
}
|
||||
|
||||
async fn store_test_secret(
|
||||
manager: &crate::extensions::manager::ExtensionManager,
|
||||
name: &str,
|
||||
value: &str,
|
||||
) {
|
||||
manager
|
||||
.secrets
|
||||
.create("test", CreateSecretParams::new(name, value))
|
||||
.await
|
||||
.expect("store secret");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_setting_value_is_present() {
|
||||
assert!(
|
||||
@@ -7423,7 +7762,13 @@ mod tests {
|
||||
// Regression: remove() only checked channel_runtime for shutdown, missing
|
||||
// relay-only mode where only relay_channel_manager is set.
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let mgr = make_test_manager(None, dir.path().to_path_buf());
|
||||
let (store, _db_dir) = make_test_store().await;
|
||||
let mgr = make_test_manager_with_dirs(
|
||||
None,
|
||||
dir.path().join("tools"),
|
||||
dir.path().join("channels"),
|
||||
Some(store),
|
||||
);
|
||||
|
||||
// Set up relay channel manager with a stub channel
|
||||
let cm = Arc::new(crate::channels::ChannelManager::new());
|
||||
@@ -7450,6 +7795,8 @@ mod tests {
|
||||
.await
|
||||
.expect("store team_id");
|
||||
}
|
||||
store_test_secret(&mgr, "relay:slack-relay:oauth_state", "nonce").await;
|
||||
store_test_secret(&mgr, "relay:slack-relay:stream_token", "legacy-token").await;
|
||||
|
||||
// Verify channel exists before removal
|
||||
assert!(cm.get_channel("slack-relay").await.is_some());
|
||||
@@ -7478,6 +7825,30 @@ mod tests {
|
||||
cm.get_channel("slack-relay").await.is_none(),
|
||||
"relay channel should be removed from the channel manager"
|
||||
);
|
||||
assert!(
|
||||
!mgr.secrets
|
||||
.exists("test", "relay:slack-relay:oauth_state")
|
||||
.await
|
||||
.expect("oauth state exists query"),
|
||||
"relay oauth_state secret should be removed"
|
||||
);
|
||||
assert!(
|
||||
!mgr.secrets
|
||||
.exists("test", "relay:slack-relay:stream_token")
|
||||
.await
|
||||
.expect("stream token exists query"),
|
||||
"relay legacy stream token should be removed"
|
||||
);
|
||||
assert_eq!(
|
||||
mgr.store
|
||||
.as_ref()
|
||||
.expect("store")
|
||||
.get_setting("test", "relay:slack-relay:team_id")
|
||||
.await
|
||||
.expect("team_id query"),
|
||||
None,
|
||||
"relay team_id setting should be removed"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -7585,6 +7956,185 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_remove_wasm_tool_deletes_unique_secrets() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let tools_dir = write_test_tool(
|
||||
dir.path(),
|
||||
"github",
|
||||
r#"{
|
||||
"name": "github",
|
||||
"auth": { "secret_name": "github_token" },
|
||||
"setup": {
|
||||
"required_secrets": [
|
||||
{ "name": "github_client_secret", "prompt": "GitHub client secret for testing cleanup behavior." }
|
||||
]
|
||||
},
|
||||
"http": {
|
||||
"credentials": {
|
||||
"service_token": {
|
||||
"secret_name": "github_service_token",
|
||||
"location": { "type": "bearer" }
|
||||
}
|
||||
}
|
||||
},
|
||||
"webhook": {
|
||||
"hmac_secret_name": "github_webhook_secret"
|
||||
}
|
||||
}"#,
|
||||
);
|
||||
let mgr = make_test_manager_with_dirs(None, tools_dir, dir.path().join("channels"), None);
|
||||
|
||||
store_test_secret(&mgr, "github_token", "access-token").await;
|
||||
store_test_secret(&mgr, "github_token_refresh_token", "refresh-token").await;
|
||||
store_test_secret(&mgr, "github_token_scopes", "repo workflow").await;
|
||||
store_test_secret(&mgr, "github_client_secret", "client-secret").await;
|
||||
store_test_secret(&mgr, "github_service_token", "service-token").await;
|
||||
store_test_secret(&mgr, "github_webhook_secret", "webhook-secret").await;
|
||||
|
||||
mgr.remove("github", "test")
|
||||
.await
|
||||
.expect("remove should succeed");
|
||||
|
||||
for secret_name in [
|
||||
"github_token",
|
||||
"github_token_refresh_token",
|
||||
"github_token_scopes",
|
||||
"github_client_secret",
|
||||
"github_service_token",
|
||||
"github_webhook_secret",
|
||||
] {
|
||||
assert!(
|
||||
!mgr.secrets
|
||||
.exists("test", secret_name)
|
||||
.await
|
||||
.expect("exists query"),
|
||||
"secret {secret_name} should be deleted"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_remove_wasm_tool_keeps_secrets_when_other_tool_capabilities_missing() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let tools_dir = write_test_tool(
|
||||
dir.path(),
|
||||
"github",
|
||||
r#"{
|
||||
"name": "github",
|
||||
"auth": { "secret_name": "shared_token" }
|
||||
}"#,
|
||||
);
|
||||
std::fs::write(tools_dir.join("broken.wasm"), b"fake-tool").expect("write tool");
|
||||
|
||||
let mgr = make_test_manager_with_dirs(None, tools_dir, dir.path().join("channels"), None);
|
||||
store_test_secret(&mgr, "shared_token", "access-token").await;
|
||||
store_test_secret(&mgr, "shared_token_refresh_token", "refresh-token").await;
|
||||
store_test_secret(&mgr, "shared_token_scopes", "repo").await;
|
||||
|
||||
mgr.remove("github", "test")
|
||||
.await
|
||||
.expect("remove should succeed");
|
||||
|
||||
for secret_name in [
|
||||
"shared_token",
|
||||
"shared_token_refresh_token",
|
||||
"shared_token_scopes",
|
||||
] {
|
||||
assert!(
|
||||
mgr.secrets
|
||||
.exists("test", secret_name)
|
||||
.await
|
||||
.expect("exists query"),
|
||||
"secret {secret_name} should be retained when reference detection is uncertain"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_remove_wasm_tool_keeps_shared_secrets_until_last_extension() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
write_test_tool(
|
||||
dir.path(),
|
||||
"google-calendar",
|
||||
r#"{
|
||||
"name": "google-calendar",
|
||||
"auth": { "secret_name": "google_oauth_token" },
|
||||
"setup": {
|
||||
"required_secrets": [
|
||||
{ "name": "google_oauth_client_id", "prompt": "Google OAuth client id for cleanup testing." },
|
||||
{ "name": "google_oauth_client_secret", "prompt": "Google OAuth client secret for cleanup testing." }
|
||||
]
|
||||
}
|
||||
}"#,
|
||||
);
|
||||
let tools_dir = write_test_tool(
|
||||
dir.path(),
|
||||
"google-drive",
|
||||
r#"{
|
||||
"name": "google-drive",
|
||||
"auth": { "secret_name": "google_oauth_token" },
|
||||
"setup": {
|
||||
"required_secrets": [
|
||||
{ "name": "google_oauth_client_id", "prompt": "Google OAuth client id for cleanup testing." },
|
||||
{ "name": "google_oauth_client_secret", "prompt": "Google OAuth client secret for cleanup testing." }
|
||||
]
|
||||
}
|
||||
}"#,
|
||||
);
|
||||
let mgr = make_test_manager_with_dirs(None, tools_dir, dir.path().join("channels"), None);
|
||||
|
||||
for (secret_name, value) in [
|
||||
("google_oauth_token", "access-token"),
|
||||
("google_oauth_token_refresh_token", "refresh-token"),
|
||||
("google_oauth_token_scopes", "calendar drive"),
|
||||
("google_oauth_client_id", "client-id"),
|
||||
("google_oauth_client_secret", "client-secret"),
|
||||
] {
|
||||
store_test_secret(&mgr, secret_name, value).await;
|
||||
}
|
||||
|
||||
mgr.remove("google-calendar", "test")
|
||||
.await
|
||||
.expect("first remove should succeed");
|
||||
|
||||
for secret_name in [
|
||||
"google_oauth_token",
|
||||
"google_oauth_token_refresh_token",
|
||||
"google_oauth_token_scopes",
|
||||
"google_oauth_client_id",
|
||||
"google_oauth_client_secret",
|
||||
] {
|
||||
assert!(
|
||||
mgr.secrets
|
||||
.exists("test", secret_name)
|
||||
.await
|
||||
.expect("exists query"),
|
||||
"shared secret {secret_name} should remain while google-drive is still installed"
|
||||
);
|
||||
}
|
||||
|
||||
mgr.remove("google-drive", "test")
|
||||
.await
|
||||
.expect("second remove should succeed");
|
||||
|
||||
for secret_name in [
|
||||
"google_oauth_token",
|
||||
"google_oauth_token_refresh_token",
|
||||
"google_oauth_token_scopes",
|
||||
"google_oauth_client_id",
|
||||
"google_oauth_client_secret",
|
||||
] {
|
||||
assert!(
|
||||
!mgr.secrets
|
||||
.exists("test", secret_name)
|
||||
.await
|
||||
.expect("exists query"),
|
||||
"shared secret {secret_name} should be deleted after the last tool is removed"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_remove_wasm_channel_clears_activation_error_and_deletes_files() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
@@ -7619,6 +8169,104 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_remove_wasm_channel_deletes_setup_secrets() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let channels_dir = write_test_channel(
|
||||
dir.path(),
|
||||
"telegram",
|
||||
r#"{
|
||||
"type": "channel",
|
||||
"name": "telegram",
|
||||
"setup": {
|
||||
"required_secrets": [
|
||||
{
|
||||
"name": "telegram_bot_token",
|
||||
"prompt": "Telegram bot token used to verify uninstall cleanup behavior."
|
||||
}
|
||||
]
|
||||
},
|
||||
"capabilities": {
|
||||
"http": {
|
||||
"credentials": {
|
||||
"tenant_token": {
|
||||
"secret_name": "telegram_service_token",
|
||||
"location": { "type": "bearer" }
|
||||
}
|
||||
}
|
||||
},
|
||||
"channel": {
|
||||
"webhook": {
|
||||
"secret_header": "X-Telegram-Bot-Api-Secret-Token",
|
||||
"secret_name": "telegram_webhook_secret"
|
||||
}
|
||||
}
|
||||
}
|
||||
}"#,
|
||||
);
|
||||
let mgr = make_test_manager_with_dirs(None, dir.path().join("tools"), channels_dir, None);
|
||||
|
||||
store_test_secret(&mgr, "telegram_bot_token", "123:telegram-token").await;
|
||||
store_test_secret(&mgr, "telegram_service_token", "tenant-service-token").await;
|
||||
store_test_secret(&mgr, "telegram_webhook_secret", "webhook-secret").await;
|
||||
|
||||
mgr.remove("telegram", "test")
|
||||
.await
|
||||
.expect("remove should succeed");
|
||||
|
||||
for secret_name in [
|
||||
"telegram_bot_token",
|
||||
"telegram_service_token",
|
||||
"telegram_webhook_secret",
|
||||
] {
|
||||
assert!(
|
||||
!mgr.secrets
|
||||
.exists("test", secret_name)
|
||||
.await
|
||||
.expect("exists query"),
|
||||
"channel secret {secret_name} should be deleted"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_remove_mcp_server_deletes_stored_secrets() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let (store, _db_dir) = make_test_store().await;
|
||||
let mgr = make_test_manager_with_dirs(
|
||||
None,
|
||||
dir.path().join("tools"),
|
||||
dir.path().join("channels"),
|
||||
Some(Arc::clone(&store)),
|
||||
);
|
||||
let server = McpServerConfig::new("notion", "https://example.com/mcp");
|
||||
mgr.add_mcp_server(server.clone(), "test")
|
||||
.await
|
||||
.expect("add mcp server");
|
||||
|
||||
store_test_secret(&mgr, &server.token_secret_name(), "access-token").await;
|
||||
store_test_secret(&mgr, &server.refresh_token_secret_name(), "refresh-token").await;
|
||||
store_test_secret(&mgr, &server.client_id_secret_name(), "client-id").await;
|
||||
|
||||
mgr.remove("notion", "test")
|
||||
.await
|
||||
.expect("remove should succeed");
|
||||
|
||||
for secret_name in [
|
||||
server.token_secret_name(),
|
||||
server.refresh_token_secret_name(),
|
||||
server.client_id_secret_name(),
|
||||
] {
|
||||
assert!(
|
||||
!mgr.secrets
|
||||
.exists("test", &secret_name)
|
||||
.await
|
||||
.expect("exists query"),
|
||||
"MCP secret {secret_name} should be deleted"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_url_with_query_params() {
|
||||
let url = "https://api.example.com/path?api_key=secret123&token=abc";
|
||||
|
||||
+19
-3
@@ -1581,19 +1581,20 @@ impl Store {
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
source_channel: Option<&str>,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let affected = conn
|
||||
.execute(
|
||||
r#"
|
||||
INSERT INTO conversations (id, channel, user_id, thread_id)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
INSERT INTO conversations (id, channel, user_id, thread_id, source_channel)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
ON CONFLICT (id) DO UPDATE
|
||||
SET last_activity = NOW()
|
||||
WHERE conversations.user_id = EXCLUDED.user_id
|
||||
AND conversations.channel = EXCLUDED.channel
|
||||
"#,
|
||||
&[&id, &channel, &user_id, &thread_id],
|
||||
&[&id, &channel, &user_id, &thread_id, &source_channel],
|
||||
)
|
||||
.await?;
|
||||
Ok(affected > 0)
|
||||
@@ -1892,6 +1893,21 @@ impl Store {
|
||||
Ok(row.is_some())
|
||||
}
|
||||
|
||||
/// Get the source_channel for a conversation.
|
||||
pub async fn get_conversation_source_channel(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
) -> Result<Option<String>, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let row = conn
|
||||
.query_opt(
|
||||
"SELECT source_channel FROM conversations WHERE id = $1",
|
||||
&[&conversation_id],
|
||||
)
|
||||
.await?;
|
||||
Ok(row.and_then(|r| r.get::<_, Option<String>>(0)))
|
||||
}
|
||||
|
||||
/// Load messages for a conversation with cursor-based pagination.
|
||||
///
|
||||
/// Returns `(messages_oldest_first, has_more)`.
|
||||
|
||||
@@ -449,6 +449,7 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
&components.secrets_store,
|
||||
components.extension_manager.as_ref(),
|
||||
components.db.as_ref(),
|
||||
&channel_names,
|
||||
)
|
||||
.await;
|
||||
|
||||
|
||||
+1
-1
@@ -278,7 +278,7 @@ impl TenantScope {
|
||||
thread_id: Option<&str>,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
self.inner
|
||||
.ensure_conversation(id, channel, &self.user_id, thread_id)
|
||||
.ensure_conversation(id, channel, &self.user_id, thread_id, None)
|
||||
.await
|
||||
}
|
||||
|
||||
|
||||
+3
-3
@@ -753,7 +753,7 @@ mod tests {
|
||||
|
||||
// ensure_conversation should create the row.
|
||||
assert!(
|
||||
db.ensure_conversation(conv_id, "web", "carol", None)
|
||||
db.ensure_conversation(conv_id, "web", "carol", None, Some("web"))
|
||||
.await
|
||||
.expect("ensure first"),
|
||||
"first ensure_conversation should create the row"
|
||||
@@ -761,7 +761,7 @@ mod tests {
|
||||
|
||||
// Calling again with the same ID should not error.
|
||||
assert!(
|
||||
db.ensure_conversation(conv_id, "web", "carol", None)
|
||||
db.ensure_conversation(conv_id, "web", "carol", None, Some("web"))
|
||||
.await
|
||||
.expect("ensure second (idempotent)"),
|
||||
"second ensure_conversation should touch owned row"
|
||||
@@ -806,7 +806,7 @@ mod tests {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
|
||||
|
||||
assert!(
|
||||
!db.ensure_conversation(conv_id, "web", "mallory", None)
|
||||
!db.ensure_conversation(conv_id, "web", "mallory", None, None)
|
||||
.await
|
||||
.expect("foreign ensure should not error"),
|
||||
"foreign ensure_conversation should report not ensured"
|
||||
|
||||
@@ -124,8 +124,10 @@ impl WasmToolLoader {
|
||||
let wasm_bytes = fs::read(wasm_path).await?;
|
||||
|
||||
// Read capabilities (optional) and extract OAuth refresh config
|
||||
// and tool description. Parameter schema is auto-derived from the
|
||||
// WASM module's schema() export (see WasmToolSchemas::compact_schema).
|
||||
// and tool description. Parameter schema is NOT read from the
|
||||
// capabilities file — it is auto-derived from the WASM module's
|
||||
// schema() export at prepare time (see WasmToolSchemas::compact_schema),
|
||||
// so no schema override is needed here.
|
||||
let (capabilities, oauth_refresh, description) = if let Some(cap_path) = capabilities_path {
|
||||
if cap_path.exists() {
|
||||
let cap_bytes = fs::read(cap_path).await?;
|
||||
|
||||
@@ -759,14 +759,31 @@ impl WasmToolSchemas {
|
||||
}
|
||||
|
||||
let kept: serde_json::Map<String, serde_json::Value> = all_properties
|
||||
.into_iter()
|
||||
.iter()
|
||||
.filter(|(name, prop)| {
|
||||
required.contains(name) || prop.get("enum").is_some() || prop.get("const").is_some()
|
||||
required.contains(name.as_str())
|
||||
|| prop.get("enum").is_some()
|
||||
|| prop.get("const").is_some()
|
||||
})
|
||||
.map(|(k, v)| (k.clone(), v.clone()))
|
||||
.collect();
|
||||
|
||||
if kept.is_empty() {
|
||||
return Self::permissive_schema();
|
||||
// When the schema has typed properties but none survived the
|
||||
// required/enum filter, include all typed properties so the LLM
|
||||
// sees meaningful parameter hints instead of permissive `{}`.
|
||||
let typed: serde_json::Map<String, serde_json::Value> = all_properties
|
||||
.into_iter()
|
||||
.filter(|(_, prop)| schema_is_typed_property(prop))
|
||||
.collect();
|
||||
if typed.is_empty() {
|
||||
return Self::permissive_schema();
|
||||
}
|
||||
return serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": typed,
|
||||
"additionalProperties": true,
|
||||
});
|
||||
}
|
||||
|
||||
let kept_required: Vec<serde_json::Value> = required
|
||||
@@ -1991,6 +2008,58 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_typed_schema_without_required_is_advertised() {
|
||||
// Regression test for #1303: when a WASM tool exports a typed schema
|
||||
// with no required/enum fields, the advertised schema should still
|
||||
// contain the typed properties instead of falling back to permissive {}.
|
||||
let discovery_schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": { "type": "string" },
|
||||
"limit": { "type": "integer" }
|
||||
}
|
||||
});
|
||||
|
||||
let runtime = Arc::new(WasmToolRuntime::new(WasmRuntimeConfig::for_testing()).unwrap());
|
||||
let prepared = runtime
|
||||
.prepare("typed_search", b"\0asm\x0d\0\x01\0", None)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut wrapper =
|
||||
super::WasmToolWrapper::new(Arc::clone(&runtime), prepared, Capabilities::default());
|
||||
wrapper.schemas = super::WasmToolSchemas::new(discovery_schema.clone());
|
||||
wrapper.description = "Typed search tool".to_string();
|
||||
|
||||
let advertised = wrapper.parameters_schema();
|
||||
let props = advertised["properties"].as_object().unwrap();
|
||||
|
||||
// Both typed properties should be preserved in the advertised schema
|
||||
assert!(
|
||||
props.contains_key("query"),
|
||||
"advertised schema should contain 'query' property"
|
||||
);
|
||||
assert!(
|
||||
props.contains_key("limit"),
|
||||
"advertised schema should contain 'limit' property"
|
||||
);
|
||||
assert_eq!(props.len(), 2);
|
||||
|
||||
// The schema should NOT be permissive
|
||||
assert!(
|
||||
!super::WasmToolSchemas::is_permissive_schema(&advertised),
|
||||
"advertised schema should not be permissive when typed properties exist"
|
||||
);
|
||||
|
||||
// No tool_info hint needed since typed properties are visible
|
||||
let schema = wrapper.schema();
|
||||
assert!(
|
||||
!schema.description.contains("tool_info"),
|
||||
"description should not contain tool_info hint: {}",
|
||||
schema.description
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compact_schema_keeps_required_and_enum_properties() {
|
||||
let schema = serde_json::json!({
|
||||
@@ -2028,8 +2097,8 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compact_schema_falls_back_to_permissive_when_empty() {
|
||||
// No required, no enum → permissive fallback
|
||||
fn test_compact_schema_preserves_typed_properties_when_no_required() {
|
||||
// No required, no enum, but typed properties → keep all typed props
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
@@ -2038,6 +2107,24 @@ mod tests {
|
||||
}
|
||||
});
|
||||
|
||||
let compacted = super::WasmToolSchemas::compact_schema(&schema);
|
||||
let props = compacted["properties"].as_object().unwrap();
|
||||
assert_eq!(props.len(), 2);
|
||||
assert!(props.contains_key("query"));
|
||||
assert!(props.contains_key("limit"));
|
||||
assert_eq!(compacted["additionalProperties"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compact_schema_falls_back_to_permissive_when_no_typed_properties() {
|
||||
// Properties with no type info → permissive fallback
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"data": {}
|
||||
}
|
||||
});
|
||||
|
||||
let compacted = super::WasmToolSchemas::compact_schema(&schema);
|
||||
assert!(compacted["properties"].as_object().unwrap().is_empty());
|
||||
}
|
||||
|
||||
@@ -53,6 +53,7 @@ HEADED=1 pytest scenarios/
|
||||
| `test_skills.py` | Skills tab UI visibility, ClawHub search (skipped if registry unreachable), install + remove lifecycle |
|
||||
| `test_sse_reconnect.py` | SSE reconnects after programmatic `eventSource.close()` + `connectSSE()`; history is reloaded after reconnect |
|
||||
| `test_tool_approval.py` | Approval card appears, buttons disable on approve/deny, parameters toggle via `page.evaluate("showApproval(...)")`; the waiting-approval regression uses a real HTTP tool call |
|
||||
| `test_extension_uninstall_cleanup.py` | Real install/setup/remove coverage for WASM tools, WASM channels, OAuth-backed shared Google tools, and MCP servers; verifies uninstall deletes stored secrets from the libSQL `secrets` table while preserving shared credentials until the last referencing extension is removed |
|
||||
| `test_oauth_refresh.py` | Hosted Gmail OAuth regression: complete setup via `/oauth/callback`, expire the stored access token in libSQL, trigger a real `gmail` tool call through `/api/chat/send`, and verify refresh goes through the mock `/oauth/refresh` proxy without forwarding `client_secret` |
|
||||
|
||||
## `helpers.py`
|
||||
@@ -77,6 +78,7 @@ All fixtures are defined in `tests/e2e/conftest.py`. Running `pytest scenarios/`
|
||||
| `mock_llm_server` | Starts `mock_llm.py --port 0`, reads the assigned port from stdout, waits for `/v1/models` to return 200. Yields the base URL. |
|
||||
| `ironclaw_server` | Starts the ironclaw binary with a minimal env (see below), waits for `/api/health` (timeout 60s). Yields the base URL. On teardown sends **SIGINT** (not SIGTERM) so the tokio ctrl_c handler triggers a graceful shutdown and LLVM coverage data is flushed. |
|
||||
| `hosted_oauth_refresh_server` | Starts a second ironclaw instance with a dedicated libSQL DB and `GOOGLE_OAUTH_CLIENT_ID=hosted-google-client-id`, while still pointing `IRONCLAW_OAUTH_EXCHANGE_URL` at `mock_llm.py`. Yields a dict with `base_url`, `db_path`, `gateway_user_id`, and `mock_llm_url` for the hosted refresh regression scenario. |
|
||||
| `extension_cleanup_server` | Starts an isolated ironclaw instance with its own temp DB/home/WASM dirs, `SECRETS_MASTER_KEY`, and hosted-style OAuth env so uninstall-cleanup scenarios can inspect the `secrets` table without interfering with the shared E2E server state. |
|
||||
| `browser` | Launches a single Chromium instance (headless by default; set `HEADED=1` for headed). Shared across all tests. |
|
||||
|
||||
### Function-scoped fixtures
|
||||
|
||||
@@ -443,6 +443,115 @@ async def hosted_oauth_refresh_server(
|
||||
home_tmpdir.cleanup()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
async def extension_cleanup_server(
|
||||
ironclaw_binary,
|
||||
mock_llm_server,
|
||||
):
|
||||
"""Start an isolated ironclaw instance for uninstall secret cleanup E2E tests."""
|
||||
reserved = _reserve_loopback_sockets(2)
|
||||
db_tmpdir = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-cleanup-db-")
|
||||
home_tmpdir = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-cleanup-home-")
|
||||
tools_tmpdir = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-cleanup-tools-")
|
||||
channels_tmpdir = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-cleanup-channels-")
|
||||
|
||||
try:
|
||||
gateway_port = reserved[0].getsockname()[1]
|
||||
http_port = reserved[1].getsockname()[1]
|
||||
for sock in reserved:
|
||||
if sock.fileno() != -1:
|
||||
sock.close()
|
||||
|
||||
db_path = os.path.join(db_tmpdir.name, "extension-cleanup.db")
|
||||
home_dir = home_tmpdir.name
|
||||
env = {
|
||||
"PATH": os.environ.get("PATH", "/usr/bin:/bin"),
|
||||
"HOME": home_dir,
|
||||
"IRONCLAW_BASE_DIR": os.path.join(home_dir, ".ironclaw"),
|
||||
"RUST_LOG": "ironclaw=info",
|
||||
"RUST_BACKTRACE": "1",
|
||||
"IRONCLAW_OWNER_ID": OWNER_SCOPE_ID,
|
||||
"GATEWAY_ENABLED": "true",
|
||||
"GATEWAY_HOST": "127.0.0.1",
|
||||
"GATEWAY_PORT": str(gateway_port),
|
||||
"GATEWAY_AUTH_TOKEN": AUTH_TOKEN,
|
||||
"GATEWAY_USER_ID": OWNER_SCOPE_ID,
|
||||
"HTTP_HOST": "127.0.0.1",
|
||||
"HTTP_PORT": str(http_port),
|
||||
"HTTP_WEBHOOK_SECRET": HTTP_WEBHOOK_SECRET,
|
||||
"CLI_ENABLED": "false",
|
||||
"LLM_BACKEND": "openai_compatible",
|
||||
"LLM_BASE_URL": mock_llm_server,
|
||||
"LLM_MODEL": "mock-model",
|
||||
"DATABASE_BACKEND": "libsql",
|
||||
"LIBSQL_PATH": db_path,
|
||||
"SECRETS_MASTER_KEY": "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
|
||||
"SANDBOX_ENABLED": "false",
|
||||
"SKILLS_ENABLED": "true",
|
||||
"ROUTINES_ENABLED": "true",
|
||||
"HEARTBEAT_ENABLED": "false",
|
||||
"EMBEDDING_ENABLED": "false",
|
||||
"WASM_ENABLED": "true",
|
||||
"WASM_TOOLS_DIR": tools_tmpdir.name,
|
||||
"WASM_CHANNELS_DIR": channels_tmpdir.name,
|
||||
"ONBOARD_COMPLETED": "true",
|
||||
"IRONCLAW_OAUTH_CALLBACK_URL": "https://oauth.test.example/oauth/callback",
|
||||
"IRONCLAW_OAUTH_EXCHANGE_URL": mock_llm_server,
|
||||
"GOOGLE_OAUTH_CLIENT_ID": "hosted-google-client-id",
|
||||
}
|
||||
_forward_coverage_env(env)
|
||||
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
ironclaw_binary, "--no-onboard",
|
||||
stdin=asyncio.subprocess.DEVNULL,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=env,
|
||||
)
|
||||
startup_kill_attempted = False
|
||||
base_url = f"http://127.0.0.1:{gateway_port}"
|
||||
try:
|
||||
await wait_for_ready(f"{base_url}/api/health", timeout=60)
|
||||
yield {
|
||||
"base_url": base_url,
|
||||
"db_path": db_path,
|
||||
"gateway_user_id": OWNER_SCOPE_ID,
|
||||
"mock_llm_url": mock_llm_server,
|
||||
}
|
||||
except TimeoutError:
|
||||
if proc.returncode is None:
|
||||
startup_kill_attempted = True
|
||||
await _stop_process(proc, timeout=2)
|
||||
returncode = proc.returncode
|
||||
stderr_bytes = b""
|
||||
if proc.stderr:
|
||||
try:
|
||||
stderr_bytes = await asyncio.wait_for(proc.stderr.read(8192), timeout=2)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
stderr_text = stderr_bytes.decode("utf-8", errors="replace")
|
||||
pytest.fail(
|
||||
f"extension cleanup server failed to start on port {gateway_port} "
|
||||
f"(returncode={returncode}).\nstderr:\n{stderr_text}"
|
||||
)
|
||||
finally:
|
||||
if proc.returncode is None:
|
||||
if startup_kill_attempted:
|
||||
await _stop_process(proc, timeout=2)
|
||||
else:
|
||||
await _stop_process(proc, sig=signal.SIGINT, timeout=10)
|
||||
if proc.returncode is None:
|
||||
await _stop_process(proc, timeout=2)
|
||||
finally:
|
||||
for sock in reserved:
|
||||
if sock.fileno() != -1:
|
||||
sock.close()
|
||||
db_tmpdir.cleanup()
|
||||
home_tmpdir.cleanup()
|
||||
tools_tmpdir.cleanup()
|
||||
channels_tmpdir.cleanup()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
async def http_channel_server(ironclaw_server, server_ports):
|
||||
"""HTTP webhook channel base URL."""
|
||||
|
||||
@@ -0,0 +1,266 @@
|
||||
"""Extension uninstall secret cleanup E2E tests.
|
||||
|
||||
Exercises real install/setup/auth/remove flows and verifies the backing
|
||||
secrets table is cleaned up when extensions are uninstalled.
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
from helpers import api_get, api_post
|
||||
|
||||
|
||||
def _extract_state(auth_url: str) -> str:
|
||||
parsed = urlparse(auth_url)
|
||||
state = parse_qs(parsed.query).get("state", [None])[0]
|
||||
assert state, f"auth_url should include state: {auth_url}"
|
||||
return state
|
||||
|
||||
|
||||
def _secret_exists(db_path: str, user_id: str, name: str) -> bool:
|
||||
with sqlite3.connect(db_path) as conn:
|
||||
row = conn.execute(
|
||||
"SELECT 1 FROM secrets WHERE user_id = ?1 AND name = ?2 LIMIT 1",
|
||||
(user_id, name),
|
||||
).fetchone()
|
||||
return row is not None
|
||||
|
||||
|
||||
def _secret_names(db_path: str, user_id: str) -> set[str]:
|
||||
with sqlite3.connect(db_path) as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT name FROM secrets WHERE user_id = ?1",
|
||||
(user_id,),
|
||||
).fetchall()
|
||||
return {row[0] for row in rows}
|
||||
|
||||
|
||||
async def _get_extension(base_url: str, name: str) -> dict | None:
|
||||
response = await api_get(base_url, "/api/extensions", timeout=15)
|
||||
response.raise_for_status()
|
||||
for extension in response.json().get("extensions", []):
|
||||
if extension["name"] == name:
|
||||
return extension
|
||||
return None
|
||||
|
||||
|
||||
async def _ensure_removed(base_url: str, name: str) -> None:
|
||||
extension = await _get_extension(base_url, name)
|
||||
if extension is not None:
|
||||
response = await api_post(base_url, f"/api/extensions/{name}/remove", timeout=30)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json().get("success") is True, response.text
|
||||
|
||||
|
||||
async def _install_extension(
|
||||
base_url: str,
|
||||
name: str,
|
||||
*,
|
||||
kind: str | None = None,
|
||||
url: str | None = None,
|
||||
) -> None:
|
||||
payload = {"name": name}
|
||||
if kind is not None:
|
||||
payload["kind"] = kind
|
||||
if url is not None:
|
||||
payload["url"] = url
|
||||
|
||||
response = await api_post(
|
||||
base_url,
|
||||
"/api/extensions/install",
|
||||
json=payload,
|
||||
timeout=180,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json().get("success") is True, response.text
|
||||
|
||||
|
||||
async def test_remove_wasm_tool_deletes_unique_secret(extension_cleanup_server):
|
||||
server = extension_cleanup_server["base_url"]
|
||||
db_path = extension_cleanup_server["db_path"]
|
||||
user_id = extension_cleanup_server["gateway_user_id"]
|
||||
|
||||
await _ensure_removed(server, "web-search")
|
||||
|
||||
await _install_extension(server, "web-search")
|
||||
|
||||
setup_response = await api_post(
|
||||
server,
|
||||
"/api/extensions/web-search/setup",
|
||||
json={"secrets": {"brave_api_key": "cleanup-test-key"}},
|
||||
timeout=30,
|
||||
)
|
||||
assert setup_response.status_code == 200, setup_response.text
|
||||
assert setup_response.json().get("success") is True, setup_response.text
|
||||
assert _secret_exists(db_path, user_id, "brave_api_key")
|
||||
|
||||
remove_response = await api_post(
|
||||
server,
|
||||
"/api/extensions/web-search/remove",
|
||||
timeout=30,
|
||||
)
|
||||
assert remove_response.status_code == 200, remove_response.text
|
||||
assert remove_response.json().get("success") is True, remove_response.text
|
||||
assert not _secret_exists(db_path, user_id, "brave_api_key")
|
||||
|
||||
|
||||
async def test_remove_wasm_channel_deletes_setup_secrets(extension_cleanup_server):
|
||||
server = extension_cleanup_server["base_url"]
|
||||
db_path = extension_cleanup_server["db_path"]
|
||||
user_id = extension_cleanup_server["gateway_user_id"]
|
||||
|
||||
await _ensure_removed(server, "discord")
|
||||
|
||||
await _install_extension(server, "discord", kind="wasm_channel")
|
||||
|
||||
setup_response = await api_post(
|
||||
server,
|
||||
"/api/extensions/discord/setup",
|
||||
json={
|
||||
"secrets": {
|
||||
"discord_bot_token": "cleanup-discord-bot-token",
|
||||
"discord_public_key": "cleanup-discord-public-key",
|
||||
}
|
||||
},
|
||||
timeout=30,
|
||||
)
|
||||
assert setup_response.status_code == 200, setup_response.text
|
||||
assert setup_response.json().get("success") is True, setup_response.text
|
||||
assert _secret_exists(db_path, user_id, "discord_bot_token")
|
||||
assert _secret_exists(db_path, user_id, "discord_public_key")
|
||||
|
||||
remove_response = await api_post(
|
||||
server,
|
||||
"/api/extensions/discord/remove",
|
||||
timeout=30,
|
||||
)
|
||||
assert remove_response.status_code == 200, remove_response.text
|
||||
assert remove_response.json().get("success") is True, remove_response.text
|
||||
assert not _secret_exists(db_path, user_id, "discord_bot_token")
|
||||
assert not _secret_exists(db_path, user_id, "discord_public_key")
|
||||
|
||||
|
||||
async def test_remove_shared_google_oauth_secrets_after_last_tool(extension_cleanup_server):
|
||||
server = extension_cleanup_server["base_url"]
|
||||
db_path = extension_cleanup_server["db_path"]
|
||||
user_id = extension_cleanup_server["gateway_user_id"]
|
||||
|
||||
await _ensure_removed(server, "gmail")
|
||||
await _ensure_removed(server, "google-drive")
|
||||
|
||||
await _install_extension(server, "gmail")
|
||||
await _install_extension(server, "google-drive")
|
||||
|
||||
setup_response = await api_post(
|
||||
server,
|
||||
"/api/extensions/gmail/setup",
|
||||
json={"secrets": {}},
|
||||
timeout=30,
|
||||
)
|
||||
assert setup_response.status_code == 200, setup_response.text
|
||||
auth_url = setup_response.json().get("auth_url")
|
||||
assert auth_url, setup_response.text
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
callback_response = await client.get(
|
||||
f"{server}/oauth/callback",
|
||||
params={"code": "mock_auth_code", "state": _extract_state(auth_url)},
|
||||
timeout=30,
|
||||
follow_redirects=True,
|
||||
)
|
||||
assert callback_response.status_code == 200, callback_response.text[:400]
|
||||
|
||||
shared_secrets = [
|
||||
"google_oauth_token",
|
||||
"google_oauth_token_refresh_token",
|
||||
"google_oauth_token_scopes",
|
||||
]
|
||||
for secret_name in shared_secrets:
|
||||
assert _secret_exists(db_path, user_id, secret_name), f"expected {secret_name} to exist"
|
||||
|
||||
gmail_remove_response = await api_post(
|
||||
server,
|
||||
"/api/extensions/gmail/remove",
|
||||
timeout=30,
|
||||
)
|
||||
assert gmail_remove_response.status_code == 200, gmail_remove_response.text
|
||||
assert gmail_remove_response.json().get("success") is True, gmail_remove_response.text
|
||||
for secret_name in shared_secrets:
|
||||
assert _secret_exists(db_path, user_id, secret_name), (
|
||||
f"{secret_name} should remain while google-drive is still installed"
|
||||
)
|
||||
|
||||
drive_remove_response = await api_post(
|
||||
server,
|
||||
"/api/extensions/google-drive/remove",
|
||||
timeout=30,
|
||||
)
|
||||
assert drive_remove_response.status_code == 200, drive_remove_response.text
|
||||
assert drive_remove_response.json().get("success") is True, drive_remove_response.text
|
||||
for secret_name in shared_secrets:
|
||||
assert not _secret_exists(db_path, user_id, secret_name), (
|
||||
f"{secret_name} should be deleted after the last Google tool is removed"
|
||||
)
|
||||
|
||||
|
||||
async def test_remove_mcp_server_deletes_stored_secrets(extension_cleanup_server):
|
||||
server = extension_cleanup_server["base_url"]
|
||||
db_path = extension_cleanup_server["db_path"]
|
||||
user_id = extension_cleanup_server["gateway_user_id"]
|
||||
mcp_url = f"{extension_cleanup_server['mock_llm_url']}/mcp"
|
||||
|
||||
await _ensure_removed(server, "mock-mcp")
|
||||
|
||||
await _install_extension(server, "mock-mcp", kind="mcp_server", url=mcp_url)
|
||||
|
||||
setup_response = await api_post(
|
||||
server,
|
||||
"/api/extensions/mock-mcp/setup",
|
||||
json={"secrets": {}},
|
||||
timeout=30,
|
||||
)
|
||||
assert setup_response.status_code == 200, setup_response.text
|
||||
auth_url = setup_response.json().get("auth_url")
|
||||
if auth_url is None:
|
||||
activate_response = await api_post(
|
||||
server,
|
||||
"/api/extensions/mock-mcp/activate",
|
||||
timeout=30,
|
||||
)
|
||||
assert activate_response.status_code == 200, activate_response.text
|
||||
auth_url = activate_response.json().get("auth_url")
|
||||
assert auth_url, "mock-mcp should require OAuth in E2E"
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
callback_response = await client.get(
|
||||
f"{server}/oauth/callback",
|
||||
params={"code": "mock_mcp_code", "state": _extract_state(auth_url)},
|
||||
timeout=30,
|
||||
follow_redirects=True,
|
||||
)
|
||||
assert callback_response.status_code == 200, callback_response.text[:400]
|
||||
|
||||
expected_mcp_secrets = [
|
||||
"mcp_mock-mcp_access_token",
|
||||
"mcp_mock-mcp_client_id",
|
||||
]
|
||||
stored_secret_names = _secret_names(db_path, user_id)
|
||||
for secret_name in expected_mcp_secrets:
|
||||
assert secret_name in stored_secret_names, (
|
||||
f"expected {secret_name} to exist; stored secrets were {sorted(stored_secret_names)}"
|
||||
)
|
||||
|
||||
remove_response = await api_post(
|
||||
server,
|
||||
"/api/extensions/mock-mcp/remove",
|
||||
timeout=30,
|
||||
)
|
||||
assert remove_response.status_code == 200, remove_response.text
|
||||
assert remove_response.json().get("success") is True, remove_response.text
|
||||
remaining_secret_names = _secret_names(db_path, user_id)
|
||||
assert not any(name.startswith("mcp_mock-mcp_") for name in remaining_secret_names), (
|
||||
f"mock-mcp secrets should be deleted on remove; remaining secrets were "
|
||||
f"{sorted(remaining_secret_names)}"
|
||||
)
|
||||
@@ -48,7 +48,13 @@ mod tests {
|
||||
let store = rig.database();
|
||||
assert!(
|
||||
store
|
||||
.ensure_conversation(foreign_thread_id, "gateway", "victim-user", None)
|
||||
.ensure_conversation(
|
||||
foreign_thread_id,
|
||||
"gateway",
|
||||
"victim-user",
|
||||
None,
|
||||
Some("gateway")
|
||||
)
|
||||
.await
|
||||
.expect("failed to create victim conversation"),
|
||||
"test setup failed: victim conversation was not created"
|
||||
|
||||
Reference in New Issue
Block a user