mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 23:16:26 +00:00
Compare commits
9
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2cce2e3155 | ||
|
|
7223ce8c9f | ||
|
|
82822d7b25 | ||
|
|
dcb2d89e3a | ||
|
|
f3da30a454 | ||
|
|
5901451603 | ||
|
|
d3d517fd67 | ||
|
|
01678be61d | ||
|
|
fb3548956b |
@@ -1,6 +1,11 @@
|
||||
name: E2E Tests
|
||||
on:
|
||||
workflow_call:
|
||||
inputs:
|
||||
ref:
|
||||
description: Commit SHA or ref to test
|
||||
required: false
|
||||
type: string
|
||||
schedule:
|
||||
- cron: "0 6 * * 1" # Weekly Monday 6 AM UTC
|
||||
workflow_dispatch:
|
||||
@@ -19,6 +24,8 @@ jobs:
|
||||
timeout-minutes: 30
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.sha }}
|
||||
|
||||
- uses: dtolnay/rust-toolchain@stable
|
||||
|
||||
@@ -59,6 +66,8 @@ jobs:
|
||||
files: "tests/e2e/scenarios/test_owner_scope.py tests/e2e/scenarios/test_routine_event_batch.py"
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.sha }}
|
||||
|
||||
- name: Download binary
|
||||
uses: actions/download-artifact@v4
|
||||
|
||||
@@ -62,7 +62,7 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: staging
|
||||
ref: ${{ github.sha }}
|
||||
fetch-depth: 0
|
||||
fetch-tags: true
|
||||
|
||||
@@ -117,6 +117,8 @@ jobs:
|
||||
needs: check-changes
|
||||
if: needs.check-changes.outputs.has_changes == 'true'
|
||||
uses: ./.github/workflows/test.yml
|
||||
with:
|
||||
ref: ${{ needs.check-changes.outputs.current_head }}
|
||||
|
||||
# ── Run E2E browser tests ────────────────────────────────────────
|
||||
e2e:
|
||||
@@ -124,6 +126,8 @@ jobs:
|
||||
needs: check-changes
|
||||
if: needs.check-changes.outputs.has_changes == 'true'
|
||||
uses: ./.github/workflows/e2e.yml
|
||||
with:
|
||||
ref: ${{ needs.check-changes.outputs.current_head }}
|
||||
|
||||
# ── Create promotion PR (triggers claude-review.yml on the PR) ──
|
||||
create-promotion-pr:
|
||||
@@ -137,7 +141,7 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: staging
|
||||
ref: ${{ needs.check-changes.outputs.current_head }}
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Generate GitHub App token
|
||||
@@ -163,7 +167,7 @@ jobs:
|
||||
PROMOTION_BASE: ${{ needs.resolve-promotion-base.outputs.promotion_base }}
|
||||
run: |
|
||||
git fetch origin "${PROMOTION_BASE}"
|
||||
AHEAD=$(git rev-list --count "origin/${PROMOTION_BASE}..origin/staging")
|
||||
AHEAD=$(git rev-list --count "origin/${PROMOTION_BASE}..HEAD")
|
||||
echo "commits_ahead=${AHEAD}" >> "$GITHUB_OUTPUT"
|
||||
if [ "$AHEAD" -eq 0 ]; then
|
||||
echo "Staging is not ahead of ${PROMOTION_BASE}. Nothing to promote."
|
||||
|
||||
@@ -1,6 +1,11 @@
|
||||
name: Run Tests
|
||||
on:
|
||||
workflow_call:
|
||||
inputs:
|
||||
ref:
|
||||
description: Commit SHA or ref to test
|
||||
required: false
|
||||
type: string
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
@@ -28,6 +33,8 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.sha }}
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
@@ -48,6 +55,8 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.sha }}
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
@@ -71,6 +80,8 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.sha }}
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
@@ -96,6 +107,8 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.sha }}
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
@@ -113,6 +126,8 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.sha }}
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
@@ -133,6 +148,8 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.sha }}
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
@@ -150,6 +167,8 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.sha }}
|
||||
- name: Build Docker image
|
||||
run: docker build -t ironclaw-test:ci .
|
||||
|
||||
@@ -161,6 +180,7 @@ jobs:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.sha }}
|
||||
fetch-depth: 0
|
||||
- name: Check version bumps for changed extensions
|
||||
env:
|
||||
|
||||
@@ -1055,10 +1055,11 @@ impl Agent {
|
||||
} else {
|
||||
drop(sess);
|
||||
self.session_manager
|
||||
.resolve_thread(
|
||||
.resolve_thread_with_parsed_uuid(
|
||||
&message.user_id,
|
||||
&message.channel,
|
||||
message.conversation_scope(),
|
||||
approval_thread_uuid,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -1139,9 +1140,9 @@ impl Agent {
|
||||
&& let Submission::UserInput { ref content } = submission
|
||||
&& let Some(engine) = self.routine_engine().await
|
||||
{
|
||||
let fired = engine
|
||||
.check_event_triggers(&message.user_id, &message.channel, content)
|
||||
.await;
|
||||
// Use post-hook content so that BeforeInbound hooks that rewrite
|
||||
// input are respected by event trigger matching.
|
||||
let fired = engine.check_event_triggers(message, content).await;
|
||||
if fired > 0 {
|
||||
tracing::debug!(
|
||||
channel = %message.channel,
|
||||
|
||||
+186
-19
@@ -24,7 +24,7 @@ use crate::agent::Scheduler;
|
||||
use crate::agent::routine::{
|
||||
NotifyConfig, Routine, RoutineAction, RoutineRun, RunStatus, Trigger, next_cron_fire,
|
||||
};
|
||||
use crate::channels::OutgoingResponse;
|
||||
use crate::channels::{IncomingMessage, OutgoingResponse};
|
||||
use crate::config::RoutineConfig;
|
||||
use crate::context::{JobContext, JobState};
|
||||
use crate::db::Database;
|
||||
@@ -56,6 +56,40 @@ pub enum SandboxReadiness {
|
||||
DockerUnavailable,
|
||||
}
|
||||
|
||||
/// Check whether an event-triggered routine's user/channel filters match an
|
||||
/// incoming message.
|
||||
///
|
||||
/// Returns `true` if:
|
||||
/// - The routine has an `Event` trigger (non-Event routines always return `false`)
|
||||
/// - The routine's `user_id` matches the message's user scope
|
||||
/// - The routine's channel filter (if any) matches the message channel
|
||||
/// case-insensitively
|
||||
///
|
||||
/// This is a pure function extracted from `check_event_triggers` so the
|
||||
/// filter logic can be unit-tested without async infrastructure.
|
||||
pub(crate) fn routine_matches_message(routine: &Routine, message: &IncomingMessage) -> bool {
|
||||
// Only Event-triggered routines can match incoming messages.
|
||||
if !matches!(routine.trigger, Trigger::Event { .. }) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// User ownership filter — only fire routines scoped to this user.
|
||||
if routine.user_id != message.user_id {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Channel filter (case-insensitive, matching emit_system_event behavior)
|
||||
if let Trigger::Event {
|
||||
channel: Some(ch), ..
|
||||
} = &routine.trigger
|
||||
&& !ch.eq_ignore_ascii_case(&message.channel)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// The routine execution engine.
|
||||
pub struct RoutineEngine {
|
||||
config: RoutineConfig,
|
||||
@@ -167,10 +201,7 @@ impl RoutineEngine {
|
||||
}
|
||||
|
||||
/// Check incoming message against event triggers. Returns number of routines fired.
|
||||
///
|
||||
/// Accepts only the three fields needed for matching (user scope, channel,
|
||||
/// message content) so callers never need to clone a full `IncomingMessage`.
|
||||
pub async fn check_event_triggers(&self, user_id: &str, channel: &str, content: &str) -> usize {
|
||||
pub async fn check_event_triggers(&self, message: &IncomingMessage, content: &str) -> usize {
|
||||
let cache = self.event_cache.read().await;
|
||||
|
||||
// Early return if there are no message matchers at all.
|
||||
@@ -208,16 +239,24 @@ impl RoutineEngine {
|
||||
EventMatcher::System { .. } => continue,
|
||||
};
|
||||
|
||||
if routine.user_id != user_id {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Channel filter
|
||||
if let Trigger::Event {
|
||||
channel: Some(ch), ..
|
||||
} = &routine.trigger
|
||||
&& ch != channel
|
||||
{
|
||||
// User ownership + channel filter (extracted for testability).
|
||||
if !routine_matches_message(routine, message) {
|
||||
// User mismatch is expected for multi-user setups — keep at
|
||||
// trace to avoid one log per routine per inbound message.
|
||||
if routine.user_id != message.user_id {
|
||||
tracing::trace!(
|
||||
routine = %routine.name,
|
||||
routine_user = %routine.user_id,
|
||||
message_user = %message.user_id,
|
||||
"Skipped: user scope mismatch"
|
||||
);
|
||||
} else {
|
||||
tracing::debug!(
|
||||
routine = %routine.name,
|
||||
channel = %message.channel,
|
||||
"Skipped: channel mismatch"
|
||||
);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -228,14 +267,14 @@ impl RoutineEngine {
|
||||
|
||||
// Cooldown check
|
||||
if !self.check_cooldown(routine) {
|
||||
tracing::trace!(routine = %routine.name, "Skipped: cooldown active");
|
||||
tracing::debug!(routine = %routine.name, "Skipped: cooldown active");
|
||||
continue;
|
||||
}
|
||||
|
||||
// Concurrent run check (using batch-loaded counts)
|
||||
let running_count = concurrent_counts.get(&routine.id).copied().unwrap_or(0);
|
||||
if running_count >= routine.guardrails.max_concurrent as i64 {
|
||||
tracing::trace!(routine = %routine.name, "Skipped: max concurrent reached");
|
||||
tracing::debug!(routine = %routine.name, "Skipped: max concurrent reached");
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -1781,6 +1820,13 @@ pub fn spawn_cron_ticker(
|
||||
engine.check_cron_triggers().await;
|
||||
|
||||
let mut ticker = tokio::time::interval(interval);
|
||||
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
// Periodic event cache refresh so web/CLI mutations are picked up
|
||||
// without requiring tool-path code to call refresh_event_cache().
|
||||
// Uses wall-clock elapsed time so the refresh cadence is stable
|
||||
// regardless of the cron tick interval configuration.
|
||||
let refresh_interval = Duration::from_secs(60);
|
||||
let mut last_refresh = tokio::time::Instant::now();
|
||||
|
||||
loop {
|
||||
ticker.tick().await;
|
||||
@@ -1788,7 +1834,11 @@ pub fn spawn_cron_ticker(
|
||||
// never races with FullJobWatcher instances from this process.
|
||||
engine.sync_dispatched_runs().await;
|
||||
engine.check_cron_triggers().await;
|
||||
engine.sync_dispatched_runs().await;
|
||||
|
||||
if last_refresh.elapsed() >= refresh_interval {
|
||||
engine.refresh_event_cache().await;
|
||||
last_refresh = tokio::time::Instant::now();
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -1854,7 +1904,13 @@ fn strip_html_tags(s: &str) -> String {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::agent::routine::{NotifyConfig, RunStatus};
|
||||
use chrono::Utc;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::agent::routine::{
|
||||
NotifyConfig, Routine, RoutineAction, RoutineGuardrails, RunStatus, Trigger,
|
||||
};
|
||||
use crate::channels::IncomingMessage;
|
||||
use crate::config::RoutineConfig;
|
||||
|
||||
#[test]
|
||||
@@ -2052,6 +2108,117 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper to build a test routine with the given user_id and trigger.
|
||||
fn make_routine(user_id: &str, trigger: Trigger) -> Routine {
|
||||
Routine {
|
||||
id: Uuid::new_v4(),
|
||||
name: "test".to_string(),
|
||||
description: String::new(),
|
||||
user_id: user_id.to_string(),
|
||||
enabled: true,
|
||||
trigger,
|
||||
action: RoutineAction::Lightweight {
|
||||
prompt: String::new(),
|
||||
context_paths: vec![],
|
||||
max_tokens: 1000,
|
||||
use_tools: false,
|
||||
max_tool_rounds: 0,
|
||||
},
|
||||
guardrails: RoutineGuardrails::default(),
|
||||
notify: Default::default(),
|
||||
last_run_at: None,
|
||||
next_fire_at: None,
|
||||
run_count: 0,
|
||||
consecutive_failures: 0,
|
||||
state: serde_json::Value::Null,
|
||||
created_at: Utc::now(),
|
||||
updated_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper to build a test IncomingMessage.
|
||||
fn make_message(user_id: &str, channel: &str, content: &str) -> IncomingMessage {
|
||||
IncomingMessage {
|
||||
id: Uuid::new_v4(),
|
||||
channel: channel.to_string(),
|
||||
user_id: user_id.to_string(),
|
||||
owner_id: user_id.to_string(),
|
||||
sender_id: user_id.to_string(),
|
||||
user_name: None,
|
||||
content: content.to_string(),
|
||||
thread_id: None,
|
||||
conversation_scope_id: None,
|
||||
received_at: Utc::now(),
|
||||
metadata: serde_json::Value::Null,
|
||||
timezone: None,
|
||||
attachments: vec![],
|
||||
is_internal: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Regression test for issue #1051: event triggers used case-sensitive
|
||||
/// channel comparison, so "Telegram" != "telegram" caused silent mismatch.
|
||||
/// Tests the actual `routine_matches_message` function used in `check_event_triggers`.
|
||||
#[test]
|
||||
fn test_channel_filter_is_case_insensitive() {
|
||||
let routine = make_routine(
|
||||
"user1",
|
||||
Trigger::Event {
|
||||
pattern: ".*".to_string(),
|
||||
channel: Some("Telegram".to_string()),
|
||||
},
|
||||
);
|
||||
let msg = make_message("user1", "telegram", "hello");
|
||||
|
||||
// Case-insensitive channel match must succeed
|
||||
assert!(super::routine_matches_message(&routine, &msg));
|
||||
|
||||
// Exact case must also work
|
||||
let msg_exact = make_message("user1", "Telegram", "hello");
|
||||
assert!(super::routine_matches_message(&routine, &msg_exact));
|
||||
|
||||
// Different channel must not match
|
||||
let msg_wrong = make_message("user1", "discord", "hello");
|
||||
assert!(!super::routine_matches_message(&routine, &msg_wrong));
|
||||
}
|
||||
|
||||
/// Regression test for issue #1051: event triggers did not filter by
|
||||
/// user_id, so routines from user A could fire on messages from user B.
|
||||
/// Tests the actual `routine_matches_message` function used in `check_event_triggers`.
|
||||
#[test]
|
||||
fn test_event_trigger_requires_user_match() {
|
||||
let routine = make_routine(
|
||||
"alice",
|
||||
Trigger::Event {
|
||||
pattern: ".*".to_string(),
|
||||
channel: None,
|
||||
},
|
||||
);
|
||||
|
||||
// Different user must not match
|
||||
let msg_bob = make_message("bob", "telegram", "hello");
|
||||
assert!(!super::routine_matches_message(&routine, &msg_bob));
|
||||
|
||||
// Same user must match
|
||||
let msg_alice = make_message("alice", "telegram", "hello");
|
||||
assert!(super::routine_matches_message(&routine, &msg_alice));
|
||||
}
|
||||
|
||||
/// When no channel filter is set, any channel should match (given user matches).
|
||||
#[test]
|
||||
fn test_no_channel_filter_matches_any_channel() {
|
||||
let routine = make_routine(
|
||||
"user1",
|
||||
Trigger::Event {
|
||||
pattern: ".*".to_string(),
|
||||
channel: None,
|
||||
},
|
||||
);
|
||||
|
||||
let msg = make_message("user1", "whatever_channel", "hello");
|
||||
assert!(super::routine_matches_message(&routine, &msg));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_routine_tool_denylist_blocks_self_management_tools() {
|
||||
let denylisted = vec
|
||||
/// with `parsed_uuid: None`.
|
||||
pub async fn resolve_thread(
|
||||
&self,
|
||||
user_id: &str,
|
||||
channel: &str,
|
||||
external_thread_id: Option<&str>,
|
||||
) -> (Arc<Mutex<Session>>, Uuid) {
|
||||
self.resolve_thread_with_parsed_uuid(user_id, channel, external_thread_id, None)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Like [`resolve_thread`](Self::resolve_thread), but accepts a pre-parsed
|
||||
/// UUID to skip redundant parsing when the caller has already validated
|
||||
/// the external thread ID as a UUID (e.g. the approval routing path).
|
||||
///
|
||||
/// Uses a single read-lock acquisition for both the key lookup and the UUID
|
||||
/// adoption check to reduce contention under concurrent approval load.
|
||||
pub async fn resolve_thread_with_parsed_uuid(
|
||||
&self,
|
||||
user_id: &str,
|
||||
channel: &str,
|
||||
external_thread_id: Option<&str>,
|
||||
parsed_uuid: Option<Uuid>,
|
||||
) -> (Arc<Mutex<Session>>, Uuid) {
|
||||
let session = self.get_or_create_session(user_id).await;
|
||||
|
||||
@@ -116,51 +135,65 @@ impl SessionManager {
|
||||
external_thread_id: external_thread_id.map(String::from),
|
||||
};
|
||||
|
||||
// Check if we have a mapping
|
||||
{
|
||||
// Use pre-parsed UUID if available, otherwise parse from string.
|
||||
let ext_uuid = parsed_uuid
|
||||
.or_else(|| external_thread_id.and_then(|ext_tid| Uuid::parse_str(ext_tid).ok()));
|
||||
|
||||
// Validate that parsed_uuid (if provided) is consistent with external_thread_id.
|
||||
#[cfg(debug_assertions)]
|
||||
if let (Some(parsed), Some(ext_tid)) = (&parsed_uuid, external_thread_id) {
|
||||
debug_assert_eq!(
|
||||
Uuid::parse_str(ext_tid).ok().as_ref(),
|
||||
Some(parsed),
|
||||
"parsed_uuid must be the parsed form of external_thread_id"
|
||||
);
|
||||
}
|
||||
|
||||
// Single read lock for both the key lookup and UUID adoption check
|
||||
let adoptable_uuid = {
|
||||
let thread_map = self.thread_map.read().await;
|
||||
|
||||
// Fast path: exact key match
|
||||
if let Some(&thread_id) = thread_map.get(&key) {
|
||||
// Verify thread still exists in session
|
||||
let sess = session.lock().await;
|
||||
if sess.threads.contains_key(&thread_id) {
|
||||
return (Arc::clone(&session), thread_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check if external_thread_id is itself a known thread UUID that
|
||||
// exists in the session but was never registered in the thread_map
|
||||
// (e.g. created by chat_new_thread_handler or hydrated from DB).
|
||||
// We only adopt it if no thread_map entry maps to this UUID —
|
||||
// otherwise it belongs to a different channel scope.
|
||||
if let Some(ext_tid) = external_thread_id
|
||||
&& let Ok(ext_uuid) = Uuid::parse_str(ext_tid)
|
||||
{
|
||||
let thread_map = self.thread_map.read().await;
|
||||
let mapped_elsewhere = thread_map.values().any(|&v| v == ext_uuid);
|
||||
drop(thread_map);
|
||||
// UUID adoption check (still under the same read lock).
|
||||
// If external_thread_id is a valid UUID not mapped elsewhere,
|
||||
// it may be a thread created by chat_new_thread_handler or
|
||||
// hydrated from DB that we can adopt.
|
||||
// Only attempt adoption when external_thread_id is Some, preserving
|
||||
// the invariant that None external_thread_id never triggers adoption.
|
||||
if external_thread_id.is_some() {
|
||||
ext_uuid.filter(|&uuid| !thread_map.values().any(|&v| v == uuid))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}; // Single read lock dropped here
|
||||
|
||||
if !mapped_elsewhere {
|
||||
let sess = session.lock().await;
|
||||
if sess.threads.contains_key(&ext_uuid) {
|
||||
drop(sess);
|
||||
// If we found an adoptable UUID, verify it exists in session and acquire write lock
|
||||
if let Some(ext_uuid) = adoptable_uuid {
|
||||
let sess = session.lock().await;
|
||||
if sess.threads.contains_key(&ext_uuid) {
|
||||
drop(sess);
|
||||
|
||||
let mut thread_map = self.thread_map.write().await;
|
||||
// Re-check after acquiring write lock to prevent race condition
|
||||
// where another task mapped this UUID between our read and write.
|
||||
if !thread_map.values().any(|&v| v == ext_uuid) {
|
||||
thread_map.insert(key, ext_uuid);
|
||||
drop(thread_map);
|
||||
// Ensure undo manager exists
|
||||
let mut undo_managers = self.undo_managers.write().await;
|
||||
undo_managers
|
||||
.entry(ext_uuid)
|
||||
.or_insert_with(|| Arc::new(Mutex::new(UndoManager::new())));
|
||||
return (session, ext_uuid);
|
||||
}
|
||||
// If it was mapped elsewhere while we were unlocked, fall through
|
||||
// to create a new thread, preserving channel isolation.
|
||||
let mut thread_map = self.thread_map.write().await;
|
||||
// Re-check after acquiring write lock to prevent race condition
|
||||
// where another task mapped this UUID between our read and write.
|
||||
if !thread_map.values().any(|&v| v == ext_uuid) {
|
||||
thread_map.insert(key, ext_uuid);
|
||||
drop(thread_map);
|
||||
// Ensure undo manager exists
|
||||
let mut undo_managers = self.undo_managers.write().await;
|
||||
undo_managers
|
||||
.entry(ext_uuid)
|
||||
.or_insert_with(|| Arc::new(Mutex::new(UndoManager::new())));
|
||||
return (session, ext_uuid);
|
||||
}
|
||||
// If mapped elsewhere while unlocked, fall through to create new thread
|
||||
}
|
||||
}
|
||||
|
||||
@@ -909,6 +942,44 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_thread_consolidates_read_path() {
|
||||
// Verify that resolve_thread still correctly handles:
|
||||
// 1. Fast path: key exists in thread_map
|
||||
// 2. UUID adoption: external_thread_id is a UUID in session but not in map
|
||||
// 3. New thread: neither path matches
|
||||
use crate::agent::session::Thread;
|
||||
|
||||
let manager = SessionManager::new();
|
||||
|
||||
// Case 1: Normal resolution creates thread and maps it
|
||||
let (session1, tid1) = manager
|
||||
.resolve_thread("user1", "chan1", Some("ext-1"))
|
||||
.await;
|
||||
// Resolving again with same key should return same thread (fast path)
|
||||
let (_, tid1_again) = manager
|
||||
.resolve_thread("user1", "chan1", Some("ext-1"))
|
||||
.await;
|
||||
assert_eq!(tid1, tid1_again);
|
||||
|
||||
// Case 2: UUID adoption - insert a thread directly into session
|
||||
let adopted_id = Uuid::new_v4();
|
||||
{
|
||||
let mut sess = session1.lock().await;
|
||||
let thread = Thread::with_id(adopted_id, sess.id);
|
||||
sess.threads.insert(adopted_id, thread);
|
||||
}
|
||||
// Resolve with the UUID as external_thread_id -- should adopt it
|
||||
let (_, resolved) = manager
|
||||
.resolve_thread("user1", "chan1", Some(&adopted_id.to_string()))
|
||||
.await;
|
||||
assert_eq!(resolved, adopted_id);
|
||||
|
||||
// Case 3: Different channel gets different thread
|
||||
let (_, tid2) = manager.resolve_thread("user1", "chan2", None).await;
|
||||
assert_ne!(tid1, tid2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_thread_finds_existing_session_thread_by_uuid() {
|
||||
use crate::agent::session::{Session, Thread};
|
||||
@@ -947,4 +1018,88 @@ mod tests {
|
||||
"should have exactly 1 thread, not a duplicate"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_thread_with_pre_parsed_uuid_adopts_thread() {
|
||||
use crate::agent::session::Thread;
|
||||
|
||||
let manager = SessionManager::new();
|
||||
let (session, _) = manager.resolve_thread("user1", "chan1", None).await;
|
||||
|
||||
// Manually insert a thread with a known UUID
|
||||
let known_id = Uuid::new_v4();
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(known_id, sess.id);
|
||||
sess.threads.insert(known_id, thread);
|
||||
}
|
||||
|
||||
// Resolve with pre-parsed UUID -- should adopt it without re-parsing
|
||||
let (_, resolved) = manager
|
||||
.resolve_thread_with_parsed_uuid(
|
||||
"user1",
|
||||
"chan1",
|
||||
Some(&known_id.to_string()),
|
||||
Some(known_id),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(resolved, known_id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_thread_with_parsed_uuid_none_delegates_to_parse() {
|
||||
use crate::agent::session::Thread;
|
||||
|
||||
let manager = SessionManager::new();
|
||||
let (session, _) = manager.resolve_thread("user2", "chan2", None).await;
|
||||
|
||||
// Insert a thread with a known UUID
|
||||
let known_id = Uuid::new_v4();
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(known_id, sess.id);
|
||||
sess.threads.insert(known_id, thread);
|
||||
}
|
||||
|
||||
// Resolve with parsed_uuid=None but a valid UUID string -- should
|
||||
// fall back to parsing the string and still adopt the thread
|
||||
let (_, resolved) = manager
|
||||
.resolve_thread_with_parsed_uuid("user2", "chan2", Some(&known_id.to_string()), None)
|
||||
.await;
|
||||
assert_eq!(resolved, known_id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_thread_with_none_external_thread_id_does_not_adopt() {
|
||||
use crate::agent::session::Thread;
|
||||
|
||||
let manager = SessionManager::new();
|
||||
let (session, default_tid) = manager.resolve_thread("user3", "chan3", None).await;
|
||||
|
||||
// Manually insert a thread with a known UUID (simulating a thread
|
||||
// created by chat_new_thread_handler)
|
||||
let known_id = Uuid::new_v4();
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(known_id, sess.id);
|
||||
sess.threads.insert(known_id, thread);
|
||||
}
|
||||
|
||||
// Resolve with external_thread_id=None but parsed_uuid=Some.
|
||||
// This should NOT adopt the UUID — the old code prevented adoption
|
||||
// when external_thread_id was None, and we preserve that invariant.
|
||||
let (_, resolved) = manager
|
||||
.resolve_thread_with_parsed_uuid("user3", "chan3", None, Some(known_id))
|
||||
.await;
|
||||
|
||||
// Should return the existing default thread, not the injected UUID
|
||||
assert_eq!(
|
||||
resolved, default_tid,
|
||||
"should return existing default thread when external_thread_id is None"
|
||||
);
|
||||
assert_ne!(
|
||||
resolved, known_id,
|
||||
"should NOT adopt UUID when external_thread_id is None"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+1
-7
@@ -312,13 +312,7 @@ impl AppBuilder {
|
||||
.create_provider(&self.config.llm.nearai.base_url, self.session.clone());
|
||||
|
||||
// Register memory tools if database is available
|
||||
let workspace_user_id = self
|
||||
.config
|
||||
.channels
|
||||
.gateway
|
||||
.as_ref()
|
||||
.map(|gw| gw.user_id.as_str())
|
||||
.unwrap_or("default");
|
||||
let workspace_user_id = self.config.owner_id.as_str();
|
||||
let workspace = if let Some(ref db) = self.db {
|
||||
let emb_cache_config = EmbeddingCacheConfig {
|
||||
max_entries: self.config.embeddings.cache_size,
|
||||
|
||||
@@ -114,7 +114,7 @@ pub async fn routines_detail_handler(
|
||||
trigger_type: run.trigger_type.clone(),
|
||||
started_at: run.started_at.to_rfc3339(),
|
||||
completed_at: run.completed_at.map(|dt| dt.to_rfc3339()),
|
||||
status: format!("{:?}", run.status),
|
||||
status: run.status.to_string(),
|
||||
result_summary: run.result_summary.clone(),
|
||||
tokens_used: run.tokens_used,
|
||||
job_id: run.job_id,
|
||||
@@ -324,7 +324,7 @@ pub async fn routines_runs_handler(
|
||||
trigger_type: run.trigger_type.clone(),
|
||||
started_at: run.started_at.to_rfc3339(),
|
||||
completed_at: run.completed_at.map(|dt| dt.to_rfc3339()),
|
||||
status: format!("{:?}", run.status),
|
||||
status: run.status.to_string(),
|
||||
result_summary: run.result_summary.clone(),
|
||||
tokens_used: run.tokens_used,
|
||||
job_id: run.job_id,
|
||||
|
||||
+22
-3
@@ -98,7 +98,8 @@ impl GatewayChannel {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: config.user_id.clone(),
|
||||
owner_id: config.user_id.clone(),
|
||||
default_sender_id: config.user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
@@ -121,6 +122,22 @@ impl GatewayChannel {
|
||||
}
|
||||
}
|
||||
|
||||
/// Rebind the single-user auth identity to the durable owner scope while
|
||||
/// preserving the configured gateway sender/routing identity.
|
||||
pub fn with_owner_scope(mut self, owner_id: impl Into<String>) -> Self {
|
||||
let owner_id = owner_id.into();
|
||||
let single_user_token = if self.config.user_tokens.is_none() {
|
||||
self.auth.first_token().map(ToOwned::to_owned)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if let Some(token) = single_user_token {
|
||||
self.auth = MultiAuthState::single(token, owner_id.clone());
|
||||
}
|
||||
self.rebuild_state(|s| s.owner_id = owner_id);
|
||||
self
|
||||
}
|
||||
|
||||
/// Create a gateway channel with a pre-built multi-user auth state.
|
||||
pub fn new_multi_auth(config: GatewayConfig, auth: MultiAuthState) -> Self {
|
||||
let state = Arc::new(GatewayState {
|
||||
@@ -137,7 +154,8 @@ impl GatewayChannel {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: config.user_id.clone(),
|
||||
owner_id: config.user_id.clone(),
|
||||
default_sender_id: config.user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
@@ -177,7 +195,8 @@ impl GatewayChannel {
|
||||
job_manager: self.state.job_manager.clone(),
|
||||
prompt_queue: self.state.prompt_queue.clone(),
|
||||
scheduler: self.state.scheduler.clone(),
|
||||
default_user_id: self.state.default_user_id.clone(),
|
||||
owner_id: self.state.owner_id.clone(),
|
||||
default_sender_id: self.state.default_sender_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: self.state.ws_tracker.clone(),
|
||||
llm_provider: self.state.llm_provider.clone(),
|
||||
|
||||
+18
-16
@@ -345,8 +345,10 @@ pub struct GatewayState {
|
||||
pub job_manager: Option<Arc<ContainerJobManager>>,
|
||||
/// Prompt queue for Claude Code follow-up prompts.
|
||||
pub prompt_queue: Option<PromptQueue>,
|
||||
/// Default user ID (fallback for non-request contexts like heartbeat/routines).
|
||||
pub default_user_id: String,
|
||||
/// Durable owner scope for persistence and unauthenticated callback flows.
|
||||
pub owner_id: String,
|
||||
/// Default sender/routing identity for gateway-originated messages.
|
||||
pub default_sender_id: String,
|
||||
/// Shutdown signal sender.
|
||||
pub shutdown_tx: tokio::sync::RwLock<Option<oneshot::Sender<()>>>,
|
||||
/// WebSocket connection tracker.
|
||||
@@ -775,7 +777,7 @@ async fn oauth_callback_handler(
|
||||
error = %error,
|
||||
"OAuth callback received with malformed state"
|
||||
);
|
||||
clear_auth_mode(&state, &state.default_user_id).await;
|
||||
clear_auth_mode(&state, &state.owner_id).await;
|
||||
return oauth_error_page("IronClaw");
|
||||
}
|
||||
};
|
||||
@@ -1136,7 +1138,7 @@ async fn slack_relay_oauth_callback_handler(
|
||||
let state_key = format!("relay:{}:oauth_state", DEFAULT_RELAY_NAME);
|
||||
let stored_state = match ext_mgr
|
||||
.secrets()
|
||||
.get_decrypted(&state.default_user_id, &state_key)
|
||||
.get_decrypted(&state.owner_id, &state_key)
|
||||
.await
|
||||
{
|
||||
Ok(secret) => secret.expose().to_string(),
|
||||
@@ -1160,10 +1162,7 @@ async fn slack_relay_oauth_callback_handler(
|
||||
}
|
||||
|
||||
// Delete the nonce (one-time use)
|
||||
let _ = ext_mgr
|
||||
.secrets()
|
||||
.delete(&state.default_user_id, &state_key)
|
||||
.await;
|
||||
let _ = ext_mgr.secrets().delete(&state.owner_id, &state_key).await;
|
||||
|
||||
let result: Result<(), String> = async {
|
||||
let store = state.store.as_ref().ok_or_else(|| {
|
||||
@@ -1174,16 +1173,12 @@ async fn slack_relay_oauth_callback_handler(
|
||||
// Store team_id in settings
|
||||
let team_id_key = format!("relay:{}:team_id", DEFAULT_RELAY_NAME);
|
||||
let _ = store
|
||||
.set_setting(
|
||||
&state.default_user_id,
|
||||
&team_id_key,
|
||||
&serde_json::json!(team_id),
|
||||
)
|
||||
.set_setting(&state.owner_id, &team_id_key, &serde_json::json!(team_id))
|
||||
.await;
|
||||
|
||||
// Activate the relay channel
|
||||
ext_mgr
|
||||
.activate_stored_relay(DEFAULT_RELAY_NAME, &state.default_user_id)
|
||||
.activate_stored_relay(DEFAULT_RELAY_NAME, &state.owner_id)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to activate relay channel: {}", e))?;
|
||||
|
||||
@@ -1303,6 +1298,9 @@ async fn chat_send_handler(
|
||||
}
|
||||
|
||||
let mut msg = IncomingMessage::new("gateway", &user.user_id, &req.content);
|
||||
if state.owner_id != state.default_sender_id && user.user_id == state.owner_id {
|
||||
msg = msg.with_sender_id(&state.default_sender_id);
|
||||
}
|
||||
// Prefer timezone from JSON body, fall back to X-Timezone header
|
||||
let tz = req
|
||||
.timezone
|
||||
@@ -1404,6 +1402,9 @@ async fn chat_approval_handler(
|
||||
})?;
|
||||
|
||||
let mut msg = IncomingMessage::new("gateway", &user.user_id, content);
|
||||
if state.owner_id != state.default_sender_id && user.user_id == state.owner_id {
|
||||
msg = msg.with_sender_id(&state.default_sender_id);
|
||||
}
|
||||
|
||||
if let Some(ref thread_id) = req.thread_id {
|
||||
msg = msg.with_thread(thread_id);
|
||||
@@ -2572,7 +2573,7 @@ async fn routines_runs_handler(
|
||||
trigger_type: run.trigger_type.clone(),
|
||||
started_at: run.started_at.to_rfc3339(),
|
||||
completed_at: run.completed_at.map(|dt| dt.to_rfc3339()),
|
||||
status: format!("{:?}", run.status),
|
||||
status: run.status.to_string(),
|
||||
result_summary: run.result_summary.clone(),
|
||||
tokens_used: run.tokens_used,
|
||||
job_id: run.job_id,
|
||||
@@ -2976,7 +2977,8 @@ mod tests {
|
||||
store: None,
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
default_user_id: "test".to_string(),
|
||||
owner_id: "test".to_string(),
|
||||
default_sender_id: "test".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: None,
|
||||
llm_provider: None,
|
||||
|
||||
@@ -4265,9 +4265,9 @@ function renderRoutineDetail(routine) {
|
||||
+ '<th>Trigger</th><th>Started</th><th>Completed</th><th>Status</th><th>Summary</th><th>Tokens</th>'
|
||||
+ '</tr></thead><tbody>';
|
||||
for (const run of routine.recent_runs) {
|
||||
const runStatusClass = run.status === 'Ok' ? 'completed'
|
||||
: run.status === 'Failed' ? 'failed'
|
||||
: run.status === 'Attention' ? 'stuck'
|
||||
const runStatusClass = run.status === 'ok' ? 'completed'
|
||||
: run.status === 'failed' ? 'failed'
|
||||
: run.status === 'attention' ? 'stuck'
|
||||
: 'in_progress';
|
||||
html += '<tr>'
|
||||
+ '<td>' + escapeHtml(run.trigger_type) + '</td>'
|
||||
|
||||
@@ -76,7 +76,8 @@ impl TestGatewayBuilder {
|
||||
store: None,
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
default_user_id: self.user_id,
|
||||
owner_id: self.user_id.clone(),
|
||||
default_sender_id: self.user_id,
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: self.llm_provider,
|
||||
|
||||
@@ -16,6 +16,7 @@ use axum::routing::{delete, get, post};
|
||||
use tower::ServiceExt;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::channels::web::GatewayChannel;
|
||||
use crate::channels::web::auth::{
|
||||
AuthenticatedUser, MultiAuthState, UserIdentity, auth_middleware,
|
||||
};
|
||||
@@ -23,6 +24,7 @@ use crate::channels::web::server::{
|
||||
ActiveConfigSnapshot, GatewayState, PerUserRateLimiter, PromptQueue, RateLimiter, WorkspacePool,
|
||||
};
|
||||
use crate::channels::web::sse::SseManager;
|
||||
use crate::config::GatewayConfig;
|
||||
|
||||
// ── Helpers ────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -64,7 +66,8 @@ fn build_state(
|
||||
store,
|
||||
job_manager: None,
|
||||
prompt_queue,
|
||||
default_user_id: "test".to_string(),
|
||||
owner_id: "test".to_string(),
|
||||
default_sender_id: "test".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: None,
|
||||
llm_provider: None,
|
||||
@@ -82,6 +85,40 @@ fn build_state(
|
||||
})
|
||||
}
|
||||
|
||||
fn gateway_config() -> GatewayConfig {
|
||||
GatewayConfig {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 3000,
|
||||
auth_token: Some("gateway-auth".to_string()),
|
||||
user_id: "gateway-sender".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
memory_layers: Vec::new(),
|
||||
user_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_owner_scope_updates_gateway_owner_scope_in_multi_user_mode() {
|
||||
let mut gateway = GatewayChannel::new(gateway_config());
|
||||
gateway.auth = two_user_auth();
|
||||
gateway.config.user_tokens = Some(HashMap::new());
|
||||
let gateway = gateway.with_owner_scope("owner-scope");
|
||||
|
||||
assert_eq!(gateway.state.owner_id, "owner-scope");
|
||||
assert_eq!(gateway.state.default_sender_id, "gateway-sender");
|
||||
|
||||
let alice = gateway
|
||||
.auth
|
||||
.authenticate("tok-alice")
|
||||
.expect("alice token should remain valid");
|
||||
let bob = gateway
|
||||
.auth
|
||||
.authenticate("tok-bob")
|
||||
.expect("bob token should remain valid");
|
||||
assert_eq!(alice.user_id, "alice");
|
||||
assert_eq!(bob.user_id, "bob");
|
||||
}
|
||||
|
||||
/// Create a libSQL-backed test database in a temporary directory.
|
||||
///
|
||||
/// Returns the database and a `TempDir` guard — the database file is
|
||||
|
||||
@@ -520,7 +520,8 @@ mod tests {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: "test".to_string(),
|
||||
owner_id: "test".to_string(),
|
||||
default_sender_id: "test".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
|
||||
+400
-24
@@ -62,6 +62,30 @@ pub fn builtin_client_id_override_env(secret_name: &str) -> Option<&'static str>
|
||||
}
|
||||
}
|
||||
|
||||
/// Suppress the baked-in desktop OAuth client secret when a hosted proxy is configured.
|
||||
///
|
||||
/// In hosted deployments, IronClaw may resolve the platform Google client ID from
|
||||
/// environment variables while still falling back to the baked-in desktop secret.
|
||||
/// That client_id/client_secret mismatch breaks Google token exchange and refresh.
|
||||
///
|
||||
/// When the proxy is configured, the platform will inject the correct server-side
|
||||
/// secret for matching platform credentials, so the baked-in secret must be omitted.
|
||||
pub fn hosted_proxy_client_secret(
|
||||
client_secret: &Option<String>,
|
||||
builtin: Option<&OAuthCredentials>,
|
||||
exchange_proxy_configured: bool,
|
||||
) -> Option<String> {
|
||||
if !exchange_proxy_configured {
|
||||
return client_secret.clone();
|
||||
}
|
||||
|
||||
let builtin_secret = builtin.map(|credentials| credentials.client_secret);
|
||||
match (client_secret, builtin_secret) {
|
||||
(Some(resolved), Some(baked_in)) if resolved == baked_in => None,
|
||||
_ => client_secret.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
// ── Shared callback server ──────────────────────────────────────────────
|
||||
|
||||
// Core OAuth callback infrastructure is defined in `crate::llm::oauth_helpers`
|
||||
@@ -661,6 +685,48 @@ pub struct ProxyTokenExchangeRequest<'a> {
|
||||
pub extra_token_params: &'a HashMap<String, String>,
|
||||
}
|
||||
|
||||
pub struct ProxyRefreshTokenRequest<'a> {
|
||||
pub proxy_url: &'a str,
|
||||
pub gateway_token: &'a str,
|
||||
pub token_url: &'a str,
|
||||
pub client_id: &'a str,
|
||||
pub client_secret: Option<&'a str>,
|
||||
pub refresh_token: &'a str,
|
||||
pub provider: Option<&'a str>,
|
||||
}
|
||||
|
||||
fn oauth_token_response_from_json(
|
||||
token_data: serde_json::Value,
|
||||
access_token_field: &str,
|
||||
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
|
||||
let access_token = token_data
|
||||
.get(access_token_field)
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
let fields: Vec<&str> = token_data
|
||||
.as_object()
|
||||
.map(|o| o.keys().map(|k| k.as_str()).collect())
|
||||
.unwrap_or_default();
|
||||
OAuthCallbackError::Io(format!(
|
||||
"No '{}' field in proxy response (fields present: {:?})",
|
||||
access_token_field, fields
|
||||
))
|
||||
})?
|
||||
.to_string();
|
||||
|
||||
let refresh_token = token_data
|
||||
.get("refresh_token")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from);
|
||||
let expires_in = token_data.get("expires_in").and_then(|v| v.as_u64());
|
||||
|
||||
Ok(OAuthTokenResponse {
|
||||
access_token,
|
||||
refresh_token,
|
||||
expires_in,
|
||||
})
|
||||
}
|
||||
|
||||
/// Exchange an OAuth authorization code via the platform's token exchange proxy.
|
||||
///
|
||||
/// Authenticated via the gateway auth token (Bearer header). The caller may
|
||||
@@ -682,6 +748,7 @@ pub async fn exchange_via_proxy(
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(60))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?;
|
||||
let mut params = vec![
|
||||
@@ -724,41 +791,350 @@ pub async fn exchange_via_proxy(
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| OAuthCallbackError::Io(format!("Failed to parse proxy response: {}", e)))?;
|
||||
oauth_token_response_from_json(token_data, request.access_token_field)
|
||||
}
|
||||
|
||||
let access_token = token_data
|
||||
.get(request.access_token_field)
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
let fields: Vec<&str> = token_data
|
||||
.as_object()
|
||||
.map(|o| o.keys().map(|k| k.as_str()).collect())
|
||||
.unwrap_or_default();
|
||||
OAuthCallbackError::Io(format!(
|
||||
"No '{}' field in proxy response (fields present: {:?})",
|
||||
request.access_token_field, fields
|
||||
))
|
||||
})?
|
||||
.to_string();
|
||||
/// Refresh an OAuth access token via the platform's token refresh proxy.
|
||||
///
|
||||
/// Authenticated via the gateway auth token (Bearer header). The caller may
|
||||
/// either rely on proxy-side secret lookup or forward a `client_secret` when
|
||||
/// the provider requires it.
|
||||
pub async fn refresh_token_via_proxy(
|
||||
request: ProxyRefreshTokenRequest<'_>,
|
||||
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
|
||||
if request.gateway_token.is_empty() {
|
||||
return Err(OAuthCallbackError::Io(
|
||||
"Gateway auth token is required for proxy token refresh".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let refresh_token = token_data
|
||||
.get("refresh_token")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from);
|
||||
let expires_in = token_data.get("expires_in").and_then(|v| v.as_u64());
|
||||
let refresh_url = format!("{}/oauth/refresh", request.proxy_url.trim_end_matches('/'));
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(15))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?;
|
||||
|
||||
Ok(OAuthTokenResponse {
|
||||
access_token,
|
||||
refresh_token,
|
||||
expires_in,
|
||||
})
|
||||
let mut params = vec![
|
||||
("refresh_token", request.refresh_token.to_string()),
|
||||
("token_url", request.token_url.to_string()),
|
||||
("client_id", request.client_id.to_string()),
|
||||
];
|
||||
if let Some(secret) = request.client_secret {
|
||||
params.push(("client_secret", secret.to_string()));
|
||||
}
|
||||
if let Some(provider) = request.provider {
|
||||
params.push(("provider", provider.to_string()));
|
||||
}
|
||||
|
||||
let response = client
|
||||
.post(&refresh_url)
|
||||
.bearer_auth(request.gateway_token)
|
||||
.form(¶ms)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
OAuthCallbackError::Io(format!("Token refresh proxy request failed: {}", e))
|
||||
})?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(OAuthCallbackError::Io(format!(
|
||||
"Token refresh proxy failed: {} - {}",
|
||||
status, body
|
||||
)));
|
||||
}
|
||||
|
||||
let token_data: serde_json::Value = response
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| OAuthCallbackError::Io(format!("Failed to parse proxy response: {}", e)))?;
|
||||
|
||||
oauth_token_response_from_json(token_data, "access_token")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::{Form, State};
|
||||
use axum::http::HeaderMap;
|
||||
use axum::response::Redirect;
|
||||
use axum::routing::post;
|
||||
use axum::{Json, Router};
|
||||
use serde_json::json;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::{Mutex, oneshot};
|
||||
|
||||
use crate::cli::oauth_defaults::{
|
||||
builtin_credentials, callback_host, callback_url, is_loopback_host, landing_html,
|
||||
};
|
||||
use crate::config::helpers::lock_env;
|
||||
use crate::testing::credentials::{TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
struct RecordedProxyRequest {
|
||||
authorization: Option<String>,
|
||||
form: HashMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct MockProxyState {
|
||||
requests: Arc<Mutex<Vec<RecordedProxyRequest>>>,
|
||||
exchange_redirect_target: String,
|
||||
refresh_redirect_target: String,
|
||||
}
|
||||
|
||||
struct MockProxyServer {
|
||||
addr: SocketAddr,
|
||||
requests: Arc<Mutex<Vec<RecordedProxyRequest>>>,
|
||||
shutdown_tx: Option<oneshot::Sender<()>>,
|
||||
server_task: Option<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl MockProxyServer {
|
||||
async fn start() -> Self {
|
||||
async fn exchange_handler(
|
||||
State(state): State<MockProxyState>,
|
||||
headers: HeaderMap,
|
||||
Form(form): Form<HashMap<String, String>>,
|
||||
) -> Json<serde_json::Value> {
|
||||
state.requests.lock().await.push(RecordedProxyRequest {
|
||||
authorization: headers
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string),
|
||||
form,
|
||||
});
|
||||
Json(json!({
|
||||
"access_token": "proxy-access-token",
|
||||
"refresh_token": "proxy-refresh-token",
|
||||
"expires_in": 7200
|
||||
}))
|
||||
}
|
||||
|
||||
async fn refresh_handler(
|
||||
State(state): State<MockProxyState>,
|
||||
headers: HeaderMap,
|
||||
Form(form): Form<HashMap<String, String>>,
|
||||
) -> Json<serde_json::Value> {
|
||||
state.requests.lock().await.push(RecordedProxyRequest {
|
||||
authorization: headers
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string),
|
||||
form,
|
||||
});
|
||||
Json(json!({
|
||||
"access_token": "proxy-access-token",
|
||||
"refresh_token": "proxy-refresh-token",
|
||||
"expires_in": 7200
|
||||
}))
|
||||
}
|
||||
|
||||
async fn exchange_redirect_handler(State(state): State<MockProxyState>) -> Redirect {
|
||||
Redirect::temporary(&state.exchange_redirect_target)
|
||||
}
|
||||
|
||||
async fn refresh_redirect_handler(State(state): State<MockProxyState>) -> Redirect {
|
||||
Redirect::temporary(&state.refresh_redirect_target)
|
||||
}
|
||||
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind mock proxy");
|
||||
let addr = listener.local_addr().expect("read mock proxy addr");
|
||||
let exchange_redirect_target = format!("http://{addr}/oauth/exchange");
|
||||
let refresh_redirect_target = format!("http://{addr}/oauth/refresh");
|
||||
let app = Router::new()
|
||||
.route("/oauth/exchange", post(exchange_handler))
|
||||
.route("/oauth/refresh", post(refresh_handler))
|
||||
.route("/redirect/oauth/exchange", post(exchange_redirect_handler))
|
||||
.route("/redirect/oauth/refresh", post(refresh_redirect_handler))
|
||||
.with_state(MockProxyState {
|
||||
requests: Arc::clone(&requests),
|
||||
exchange_redirect_target,
|
||||
refresh_redirect_target,
|
||||
});
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
|
||||
let server_task = tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async {
|
||||
let _ = shutdown_rx.await;
|
||||
})
|
||||
.await;
|
||||
});
|
||||
|
||||
Self {
|
||||
addr,
|
||||
requests,
|
||||
shutdown_tx: Some(shutdown_tx),
|
||||
server_task: Some(server_task),
|
||||
}
|
||||
}
|
||||
|
||||
fn base_url(&self) -> String {
|
||||
format!("http://{}", self.addr)
|
||||
}
|
||||
|
||||
fn redirecting_base_url(&self) -> String {
|
||||
format!("{}/redirect", self.base_url())
|
||||
}
|
||||
|
||||
async fn requests(&self) -> Vec<RecordedProxyRequest> {
|
||||
self.requests.lock().await.clone()
|
||||
}
|
||||
|
||||
async fn shutdown(mut self) {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
if let Some(task) = self.server_task.take() {
|
||||
let _ = task.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for MockProxyServer {
|
||||
fn drop(&mut self) {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
if let Some(task) = self.server_task.take() {
|
||||
task.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hosted_proxy_client_secret_suppresses_builtin_secret() {
|
||||
let builtin = builtin_credentials("google_oauth_token").expect("google builtin creds");
|
||||
let client_secret = Some(builtin.client_secret.to_string());
|
||||
|
||||
let result = super::hosted_proxy_client_secret(&client_secret, Some(&builtin), true);
|
||||
|
||||
assert_eq!(result, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hosted_proxy_client_secret_preserves_explicit_secret() {
|
||||
let builtin = builtin_credentials("google_oauth_token").expect("google builtin creds");
|
||||
let client_secret = Some("hosted-server-secret".to_string());
|
||||
|
||||
let result = super::hosted_proxy_client_secret(&client_secret, Some(&builtin), true);
|
||||
|
||||
assert_eq!(result, client_secret);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_token_via_proxy_sends_auth_and_form() {
|
||||
let server = MockProxyServer::start().await;
|
||||
|
||||
let response = super::refresh_token_via_proxy(super::ProxyRefreshTokenRequest {
|
||||
proxy_url: &server.base_url(),
|
||||
gateway_token: "gateway-test-token",
|
||||
token_url: "https://oauth2.googleapis.com/token",
|
||||
client_id: TEST_OAUTH_CLIENT_ID,
|
||||
client_secret: Some(TEST_OAUTH_CLIENT_SECRET),
|
||||
refresh_token: "refresh-token-123",
|
||||
provider: Some("google"),
|
||||
})
|
||||
.await
|
||||
.expect("proxy refresh succeeds");
|
||||
|
||||
assert_eq!(response.access_token, "proxy-access-token");
|
||||
assert_eq!(
|
||||
response.refresh_token.as_deref(),
|
||||
Some("proxy-refresh-token")
|
||||
);
|
||||
assert_eq!(response.expires_in, Some(7200));
|
||||
|
||||
let requests = server.requests().await;
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(
|
||||
requests[0].authorization.as_deref(),
|
||||
Some("Bearer gateway-test-token")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("token_url").map(String::as_str),
|
||||
Some("https://oauth2.googleapis.com/token")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("client_id").map(String::as_str),
|
||||
Some(TEST_OAUTH_CLIENT_ID)
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("client_secret").map(String::as_str),
|
||||
Some(TEST_OAUTH_CLIENT_SECRET)
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("refresh_token").map(String::as_str),
|
||||
Some("refresh-token-123")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("provider").map(String::as_str),
|
||||
Some("google")
|
||||
);
|
||||
|
||||
server.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_exchange_via_proxy_does_not_follow_redirects() {
|
||||
let server = MockProxyServer::start().await;
|
||||
|
||||
let error = match super::exchange_via_proxy(super::ProxyTokenExchangeRequest {
|
||||
proxy_url: &server.redirecting_base_url(),
|
||||
gateway_token: "gateway-test-token",
|
||||
code: "auth-code-123",
|
||||
redirect_uri: "http://localhost:3000/oauth/callback",
|
||||
token_url: "https://oauth2.googleapis.com/token",
|
||||
client_id: TEST_OAUTH_CLIENT_ID,
|
||||
client_secret: Some(TEST_OAUTH_CLIENT_SECRET),
|
||||
access_token_field: "access_token",
|
||||
code_verifier: Some("code-verifier-123"),
|
||||
extra_token_params: &HashMap::new(),
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(_) => panic!("redirected proxy exchange should fail"),
|
||||
Err(error) => error,
|
||||
};
|
||||
|
||||
assert!(error.to_string().contains("307"));
|
||||
assert!(server.requests().await.is_empty());
|
||||
|
||||
server.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_token_via_proxy_does_not_follow_redirects() {
|
||||
let server = MockProxyServer::start().await;
|
||||
|
||||
let error = match super::refresh_token_via_proxy(super::ProxyRefreshTokenRequest {
|
||||
proxy_url: &server.redirecting_base_url(),
|
||||
gateway_token: "gateway-test-token",
|
||||
token_url: "https://oauth2.googleapis.com/token",
|
||||
client_id: TEST_OAUTH_CLIENT_ID,
|
||||
client_secret: Some(TEST_OAUTH_CLIENT_SECRET),
|
||||
refresh_token: "refresh-token-123",
|
||||
provider: Some("google"),
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(_) => panic!("redirected proxy refresh should fail"),
|
||||
Err(error) => error,
|
||||
};
|
||||
|
||||
assert!(error.to_string().contains("307"));
|
||||
assert!(server.requests().await.is_empty());
|
||||
|
||||
server.shutdown().await;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_loopback_host() {
|
||||
|
||||
+19
-8
@@ -10,7 +10,7 @@ use clap::Subcommand;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::agent::routine::{
|
||||
NotifyConfig, Routine, RoutineAction, RoutineGuardrails, Trigger, next_cron_fire,
|
||||
NotifyConfig, Routine, RoutineAction, RoutineGuardrails, RunStatus, Trigger, next_cron_fire,
|
||||
};
|
||||
use crate::db::Database;
|
||||
|
||||
@@ -251,15 +251,26 @@ async fn list(
|
||||
);
|
||||
println!("{}", "-".repeat(130));
|
||||
|
||||
// Fetch last-run status for all routines in a single batch query
|
||||
let routine_ids: Vec<Uuid> = filtered.iter().map(|r| r.id).collect();
|
||||
let last_run_results = db
|
||||
.batch_get_last_run_status(&routine_ids)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
for r in &filtered {
|
||||
let status = if r.enabled {
|
||||
if r.consecutive_failures > 0 {
|
||||
format!("err({})", r.consecutive_failures)
|
||||
} else {
|
||||
"active".to_string()
|
||||
}
|
||||
} else {
|
||||
let last_run_status = last_run_results.get(&r.id).copied();
|
||||
|
||||
let status = if !r.enabled {
|
||||
"disabled".to_string()
|
||||
} else if last_run_status == Some(RunStatus::Running) {
|
||||
"running".to_string()
|
||||
} else if r.consecutive_failures > 0 {
|
||||
format!("err({})", r.consecutive_failures)
|
||||
} else if last_run_status == Some(RunStatus::Attention) {
|
||||
"attention".to_string()
|
||||
} else {
|
||||
"active".to_string()
|
||||
};
|
||||
|
||||
let next_fire = r
|
||||
|
||||
+5
-7
@@ -312,13 +312,11 @@ impl Config {
|
||||
let tunnel = TunnelConfig::resolve(settings)?;
|
||||
let channels = ChannelsConfig::resolve(settings, &owner_id)?;
|
||||
|
||||
// Resolve workspace config using the gateway user_id for default layers.
|
||||
let workspace_user_id = channels
|
||||
.gateway
|
||||
.as_ref()
|
||||
.map(|gw| gw.user_id.as_str())
|
||||
.unwrap_or("default");
|
||||
let workspace = WorkspaceConfig::resolve(workspace_user_id)?;
|
||||
// Resolve the startup workspace against the durable owner scope. The
|
||||
// gateway may expose a distinct sender identity, but the base runtime
|
||||
// workspace stays owner-scoped and per-user gateway workspaces are
|
||||
// handled separately by WorkspacePool.
|
||||
let workspace = WorkspaceConfig::resolve(&owner_id)?;
|
||||
|
||||
Ok(Self {
|
||||
owner_id: owner_id.clone(),
|
||||
|
||||
@@ -462,6 +462,56 @@ impl RoutineStore for LibSqlBackend {
|
||||
Ok(counts)
|
||||
}
|
||||
|
||||
async fn batch_get_last_run_status(
|
||||
&self,
|
||||
routine_ids: &[Uuid],
|
||||
) -> Result<HashMap<Uuid, RunStatus>, DatabaseError> {
|
||||
if routine_ids.is_empty() {
|
||||
return Ok(HashMap::new());
|
||||
}
|
||||
|
||||
let conn = self.connect().await?;
|
||||
|
||||
// SQLite doesn't support ANY($1), so we query all latest runs and filter in memory.
|
||||
// Uses a subquery to pick only the most recent run per routine.
|
||||
let mut rows = conn
|
||||
.query(
|
||||
"SELECT routine_id, status FROM routine_runs r1
|
||||
WHERE started_at = (
|
||||
SELECT MAX(started_at) FROM routine_runs r2
|
||||
WHERE r2.routine_id = r1.routine_id
|
||||
)
|
||||
GROUP BY routine_id",
|
||||
params![],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
DatabaseError::Query(format!("Failed to batch get last run status: {}", e))
|
||||
})?;
|
||||
|
||||
let routine_id_set: HashSet<Uuid> = routine_ids.iter().copied().collect();
|
||||
let mut statuses = HashMap::new();
|
||||
|
||||
while let Some(row) = rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
{
|
||||
let id_str: String = get_text(&row, 0);
|
||||
let id = Uuid::parse_str(&id_str)
|
||||
.map_err(|e| DatabaseError::Query(format!("Invalid routine UUID: {}", e)))?;
|
||||
|
||||
if routine_id_set.contains(&id) {
|
||||
let status_str: String = get_text(&row, 1);
|
||||
if let std::result::Result::Ok(status) = status_str.parse::<RunStatus>() {
|
||||
statuses.insert(id, status);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(statuses)
|
||||
}
|
||||
|
||||
async fn link_routine_run_to_job(
|
||||
&self,
|
||||
run_id: Uuid,
|
||||
|
||||
@@ -528,6 +528,15 @@ pub trait RoutineStore: Send + Sync {
|
||||
&self,
|
||||
routine_ids: &[Uuid],
|
||||
) -> Result<HashMap<Uuid, i64>, DatabaseError>;
|
||||
|
||||
/// Fetch the last run status for multiple routines in a single query.
|
||||
/// Returns a map from routine_id to its most recent RunStatus.
|
||||
/// Routines with no runs are omitted from the result.
|
||||
async fn batch_get_last_run_status(
|
||||
&self,
|
||||
routine_ids: &[Uuid],
|
||||
) -> Result<HashMap<Uuid, RunStatus>, DatabaseError>;
|
||||
|
||||
async fn link_routine_run_to_job(
|
||||
&self,
|
||||
run_id: Uuid,
|
||||
|
||||
@@ -510,6 +510,14 @@ impl RoutineStore for PgBackend {
|
||||
.await
|
||||
}
|
||||
|
||||
async fn batch_get_last_run_status(
|
||||
&self,
|
||||
routine_ids: &[Uuid],
|
||||
) -> Result<std::collections::HashMap<Uuid, crate::agent::routine::RunStatus>, DatabaseError>
|
||||
{
|
||||
self.store.batch_get_last_run_status(routine_ids).await
|
||||
}
|
||||
|
||||
async fn link_routine_run_to_job(
|
||||
&self,
|
||||
run_id: Uuid,
|
||||
|
||||
+61
-76
@@ -53,22 +53,6 @@ struct HostedOAuthFlowStart {
|
||||
flow: crate::cli::oauth_defaults::PendingOAuthFlow,
|
||||
}
|
||||
|
||||
fn hosted_proxy_client_secret(
|
||||
client_secret: &Option<String>,
|
||||
builtin: Option<&crate::cli::oauth_defaults::OAuthCredentials>,
|
||||
exchange_proxy_configured: bool,
|
||||
) -> Option<String> {
|
||||
if !exchange_proxy_configured {
|
||||
return client_secret.clone();
|
||||
}
|
||||
|
||||
let builtin_secret = builtin.map(|credentials| credentials.client_secret);
|
||||
match (client_secret, builtin_secret) {
|
||||
(Some(resolved), Some(baked_in)) if resolved == baked_in => None,
|
||||
_ => client_secret.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_oauth_callback_path(path: &str) -> String {
|
||||
let trimmed_path = path.trim_end_matches('/');
|
||||
if trimmed_path.is_empty() {
|
||||
@@ -891,24 +875,27 @@ impl ExtensionManager {
|
||||
*self.relay_channel_manager.write().await = Some(channel_manager);
|
||||
}
|
||||
|
||||
/// Check if a channel name corresponds to a relay extension (has stored stream token
|
||||
/// Check if a channel name corresponds to a relay extension (has stored team_id
|
||||
/// or is tracked in the installed relay extensions set).
|
||||
pub async fn is_relay_channel(&self, name: &str, user_id: &str) -> bool {
|
||||
// Check in-memory installed set first (supports no-store mode)
|
||||
if self.installed_relay_extensions.read().await.contains(name) {
|
||||
return true;
|
||||
}
|
||||
// Then check for stored stream token
|
||||
self.secrets
|
||||
.exists(user_id, &format!("relay:{}:stream_token", name))
|
||||
.await
|
||||
.unwrap_or(false)
|
||||
// Check for stored team_id (persisted across restarts by the OAuth callback)
|
||||
if let Some(ref store) = self.store {
|
||||
let key = format!("relay:{}:team_id", name);
|
||||
if let Ok(Some(v)) = store.get_setting(user_id, &key).await {
|
||||
return v.as_str().is_some_and(|s| !s.is_empty());
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// Restore persisted relay channels after startup.
|
||||
///
|
||||
/// Loads the persisted active channel list, filters to relay types (those with
|
||||
/// a stored stream token), and activates each via `activate_stored_relay()`.
|
||||
/// a stored team_id setting), and activates each via `activate_stored_relay()`.
|
||||
/// Skips channels that are already active.
|
||||
///
|
||||
/// Call this only after `set_relay_channel_manager()` or `set_channel_runtime()`.
|
||||
@@ -1428,9 +1415,11 @@ impl ExtensionManager {
|
||||
if kind_filter.is_none() || kind_filter == Some(ExtensionKind::ChannelRelay) {
|
||||
let installed = self.installed_relay_extensions.read().await;
|
||||
let active_names = self.active_channel_names.read().await;
|
||||
let errors = self.activation_errors.read().await;
|
||||
for name in installed.iter() {
|
||||
let active = active_names.contains(name);
|
||||
let has_token = self.is_relay_channel(name, user_id).await;
|
||||
let authenticated = self.is_relay_channel(name, user_id).await;
|
||||
let activation_error = errors.get(name).cloned();
|
||||
let registry_entry = self
|
||||
.registry
|
||||
.get_with_kind(name, Some(ExtensionKind::ChannelRelay))
|
||||
@@ -1443,13 +1432,13 @@ impl ExtensionManager {
|
||||
display_name,
|
||||
description,
|
||||
url: None,
|
||||
authenticated: has_token,
|
||||
authenticated,
|
||||
active,
|
||||
tools: Vec::new(),
|
||||
needs_setup: false,
|
||||
has_auth: true,
|
||||
installed: true,
|
||||
activation_error: None,
|
||||
activation_error,
|
||||
version: None,
|
||||
});
|
||||
}
|
||||
@@ -1626,7 +1615,22 @@ impl ExtensionManager {
|
||||
self.persist_active_channels(user_id).await;
|
||||
self.activation_errors.write().await.remove(name);
|
||||
|
||||
// Remove stored stream token
|
||||
// Remove stored team_id setting and clean up secrets
|
||||
if let Some(ref store) = self.store
|
||||
&& let Err(e) = store
|
||||
.delete_setting(user_id, &format!("relay:{}:team_id", name))
|
||||
.await
|
||||
{
|
||||
tracing::warn!(error = %e, name, "Failed to delete relay team_id setting on removal");
|
||||
}
|
||||
if let Err(e) = self
|
||||
.secrets
|
||||
.delete(user_id, &format!("relay:{}:oauth_state", name))
|
||||
.await
|
||||
{
|
||||
tracing::warn!(error = %e, name, "Failed to delete relay oauth_state secret on removal");
|
||||
}
|
||||
// Clean up legacy stream_token secret from pre-webhook installs
|
||||
let _ = self
|
||||
.secrets
|
||||
.delete(user_id, &format!("relay:{}:stream_token", name))
|
||||
@@ -3179,7 +3183,7 @@ impl ExtensionManager {
|
||||
// apps. Sending the desktop secret would cause a client_id/secret
|
||||
// mismatch because the container's GOOGLE_OAUTH_CLIENT_ID is the web
|
||||
// app, not the desktop app.
|
||||
let proxy_client_secret = hosted_proxy_client_secret(
|
||||
let proxy_client_secret = oauth_defaults::hosted_proxy_client_secret(
|
||||
&client_secret,
|
||||
builtin.as_ref(),
|
||||
oauth_defaults::exchange_proxy_url().is_some(),
|
||||
@@ -4181,13 +4185,13 @@ impl ExtensionManager {
|
||||
///
|
||||
/// For Slack: initiates OAuth flow (redirect-based).
|
||||
/// For Telegram: accepts a bot token, registers it with channel-relay,
|
||||
/// and stores the returned stream token.
|
||||
/// and stores the team_id setting.
|
||||
async fn auth_channel_relay(
|
||||
&self,
|
||||
name: &str,
|
||||
user_id: &str,
|
||||
) -> Result<AuthResult, ExtensionError> {
|
||||
// Check if already authenticated (stream token exists)
|
||||
// Check if already authenticated (team_id setting exists)
|
||||
if self.is_relay_channel(name, user_id).await {
|
||||
return Ok(AuthResult::authenticated(name, ExtensionKind::ChannelRelay));
|
||||
}
|
||||
@@ -4233,19 +4237,9 @@ impl ExtensionManager {
|
||||
name: &str,
|
||||
user_id: &str,
|
||||
) -> Result<ActivateResult, ExtensionError> {
|
||||
let token_key = format!("relay:{}:stream_token", name);
|
||||
let team_id_key = format!("relay:{}:team_id", name);
|
||||
|
||||
// Check if we have a stream token
|
||||
// Verify auth: stream token must exist (even though we don't use it in this constructor path)
|
||||
let _stream_token = match self.secrets.get_decrypted(user_id, &token_key).await {
|
||||
Ok(secret) => secret.expose().to_string(),
|
||||
Err(_) => {
|
||||
return Err(ExtensionError::AuthRequired);
|
||||
}
|
||||
};
|
||||
|
||||
// Get team_id from settings
|
||||
// Get team_id from settings (stored by the OAuth callback)
|
||||
let team_id = if let Some(ref store) = self.store {
|
||||
store
|
||||
.get_setting(user_id, &team_id_key)
|
||||
@@ -4258,6 +4252,10 @@ impl ExtensionManager {
|
||||
String::new()
|
||||
};
|
||||
|
||||
if team_id.is_empty() {
|
||||
return Err(ExtensionError::AuthRequired);
|
||||
}
|
||||
|
||||
// Use relay config captured at startup
|
||||
let relay_config = self.relay_config()?;
|
||||
|
||||
@@ -4367,11 +4365,11 @@ impl ExtensionManager {
|
||||
return Ok(ExtensionKind::WasmChannel);
|
||||
}
|
||||
|
||||
// Check channel-relay extensions (installed in memory or has stored token)
|
||||
// Check channel-relay extensions (installed in memory or has stored team_id)
|
||||
if self.installed_relay_extensions.read().await.contains(name) {
|
||||
return Ok(ExtensionKind::ChannelRelay);
|
||||
}
|
||||
// Also check if there's a stored stream token (persisted across restarts)
|
||||
// Also check if there's a stored team_id setting (persisted across restarts)
|
||||
if self.is_relay_channel(name, user_id).await {
|
||||
return Ok(ExtensionKind::ChannelRelay);
|
||||
}
|
||||
@@ -4999,11 +4997,7 @@ impl ExtensionManager {
|
||||
names.insert(server.token_secret_name());
|
||||
(names, Vec::new())
|
||||
}
|
||||
ExtensionKind::ChannelRelay => {
|
||||
let mut names = std::collections::HashSet::new();
|
||||
names.insert(format!("relay:{}:stream_token", name));
|
||||
(names, Vec::new())
|
||||
}
|
||||
ExtensionKind::ChannelRelay => (std::collections::HashSet::new(), Vec::new()),
|
||||
};
|
||||
|
||||
let allowed_fields: std::collections::HashSet<String> =
|
||||
@@ -5434,7 +5428,9 @@ impl ExtensionManager {
|
||||
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
|
||||
server.token_secret_name()
|
||||
}
|
||||
ExtensionKind::ChannelRelay => format!("relay:{}:stream_token", name),
|
||||
ExtensionKind::ChannelRelay => {
|
||||
return Err(ExtensionError::AuthRequired);
|
||||
}
|
||||
};
|
||||
|
||||
let mut secrets = std::collections::HashMap::new();
|
||||
@@ -5702,7 +5698,7 @@ mod tests {
|
||||
use crate::extensions::manager::{
|
||||
ChannelRuntimeState, FallbackDecision, TelegramBindingData, TelegramBindingResult,
|
||||
TelegramOwnerBindingState, build_wasm_channel_runtime_config_updates,
|
||||
combine_install_errors, fallback_decision, hosted_proxy_client_secret, infer_kind_from_url,
|
||||
combine_install_errors, fallback_decision, infer_kind_from_url,
|
||||
normalize_hosted_callback_url, send_telegram_text_message,
|
||||
telegram_message_matches_verification_code,
|
||||
};
|
||||
@@ -7043,7 +7039,7 @@ mod tests {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let mgr = make_test_manager(None, dir.path().to_path_buf());
|
||||
|
||||
// No token stored → not a relay channel
|
||||
// No store configured, no team_id → not a relay channel
|
||||
assert!(!mgr.is_relay_channel("slack-relay", "test").await);
|
||||
}
|
||||
|
||||
@@ -7862,19 +7858,13 @@ mod tests {
|
||||
.await
|
||||
.insert("test-relay".to_string());
|
||||
|
||||
// configure() should dispatch to activate_channel_relay(), not
|
||||
// activate_wasm_channel(). Both will fail (no runtime configured),
|
||||
// but the error should be about relay config, not WASM channels.
|
||||
let mut secrets = std::collections::HashMap::new();
|
||||
secrets.insert(
|
||||
"relay:test-relay:stream_token".to_string(),
|
||||
"tok".to_string(),
|
||||
);
|
||||
|
||||
// configure() with empty secrets should dispatch to
|
||||
// activate_channel_relay(), not activate_wasm_channel(). Relay auth
|
||||
// is OAuth-only so there are no manual secrets to pass.
|
||||
let result = mgr
|
||||
.configure(
|
||||
"test-relay",
|
||||
&secrets,
|
||||
&std::collections::HashMap::new(),
|
||||
&std::collections::HashMap::new(),
|
||||
"test",
|
||||
)
|
||||
@@ -7886,7 +7876,6 @@ mod tests {
|
||||
);
|
||||
|
||||
let result = result.unwrap();
|
||||
// Activation will fail (no relay config), but secrets should still be stored
|
||||
assert!(
|
||||
!result.activated,
|
||||
"activation should fail without relay config"
|
||||
@@ -7896,15 +7885,6 @@ mod tests {
|
||||
"error should not mention WASM — got: {}",
|
||||
result.message
|
||||
);
|
||||
|
||||
// Verify the secret was stored
|
||||
assert!(
|
||||
mgr.secrets
|
||||
.exists("test", "relay:test-relay:stream_token")
|
||||
.await
|
||||
.unwrap_or(false),
|
||||
"configure should have stored the relay stream token"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn test_validation_failed_is_distinct_error_variant() {
|
||||
@@ -7970,7 +7950,8 @@ mod tests {
|
||||
let builtin_ref = builtin.as_ref();
|
||||
let secret = Some(builtin_ref.unwrap().client_secret.to_string());
|
||||
|
||||
let result = hosted_proxy_client_secret(&secret, builtin_ref, true);
|
||||
let result =
|
||||
crate::cli::oauth_defaults::hosted_proxy_client_secret(&secret, builtin_ref, true);
|
||||
assert_eq!(
|
||||
result, None,
|
||||
"built-in desktop secret must be suppressed when the exchange proxy is configured"
|
||||
@@ -7982,7 +7963,8 @@ mod tests {
|
||||
let builtin = crate::cli::oauth_defaults::builtin_credentials("google_oauth_token");
|
||||
let secret = Some("user-entered-custom-secret".to_string());
|
||||
|
||||
let result = hosted_proxy_client_secret(&secret, builtin.as_ref(), true);
|
||||
let result =
|
||||
crate::cli::oauth_defaults::hosted_proxy_client_secret(&secret, builtin.as_ref(), true);
|
||||
assert_eq!(
|
||||
result,
|
||||
Some("user-entered-custom-secret".to_string()),
|
||||
@@ -7996,7 +7978,8 @@ mod tests {
|
||||
let builtin_ref = builtin.as_ref();
|
||||
let secret = Some(builtin_ref.unwrap().client_secret.to_string());
|
||||
|
||||
let result = hosted_proxy_client_secret(&secret, builtin_ref, false);
|
||||
let result =
|
||||
crate::cli::oauth_defaults::hosted_proxy_client_secret(&secret, builtin_ref, false);
|
||||
assert_eq!(
|
||||
result, secret,
|
||||
"built-in secret must be kept when the callback will exchange directly"
|
||||
@@ -8007,7 +7990,8 @@ mod tests {
|
||||
fn test_proxy_client_secret_none_stays_none() {
|
||||
let builtin = crate::cli::oauth_defaults::builtin_credentials("google_oauth_token");
|
||||
|
||||
let result = hosted_proxy_client_secret(&None, builtin.as_ref(), true);
|
||||
let result =
|
||||
crate::cli::oauth_defaults::hosted_proxy_client_secret(&None, builtin.as_ref(), true);
|
||||
assert_eq!(
|
||||
result, None,
|
||||
"None secret stays None even when the exchange proxy is configured"
|
||||
@@ -8021,7 +8005,8 @@ mod tests {
|
||||
assert!(builtin.is_none());
|
||||
|
||||
let secret = Some("dcr-secret".to_string());
|
||||
let result = hosted_proxy_client_secret(&secret, builtin.as_ref(), true);
|
||||
let result =
|
||||
crate::cli::oauth_defaults::hosted_proxy_client_secret(&secret, builtin.as_ref(), true);
|
||||
assert_eq!(
|
||||
result,
|
||||
Some("dcr-secret".to_string()),
|
||||
|
||||
@@ -1403,6 +1403,40 @@ impl Store {
|
||||
Ok(counts)
|
||||
}
|
||||
|
||||
/// Batch-load the most recent run status for multiple routines in a single query.
|
||||
/// Uses a window function to pick only the latest run per routine.
|
||||
#[cfg(feature = "postgres")]
|
||||
pub async fn batch_get_last_run_status(
|
||||
&self,
|
||||
routine_ids: &[Uuid],
|
||||
) -> Result<HashMap<Uuid, RunStatus>, DatabaseError> {
|
||||
if routine_ids.is_empty() {
|
||||
return Ok(HashMap::new());
|
||||
}
|
||||
|
||||
let conn = self.conn().await?;
|
||||
let rows = conn
|
||||
.query(
|
||||
"SELECT DISTINCT ON (routine_id) routine_id, status
|
||||
FROM routine_runs
|
||||
WHERE routine_id = ANY($1)
|
||||
ORDER BY routine_id, started_at DESC",
|
||||
&[&routine_ids],
|
||||
)
|
||||
.await?;
|
||||
|
||||
let mut statuses = HashMap::new();
|
||||
for row in rows {
|
||||
let id: Uuid = row.get("routine_id");
|
||||
let status_str: String = row.get("status");
|
||||
if let std::result::Result::Ok(status) = status_str.parse::<RunStatus>() {
|
||||
statuses.insert(id, status);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(statuses)
|
||||
}
|
||||
|
||||
/// Link a routine run to a dispatched job.
|
||||
pub async fn link_routine_run_to_job(
|
||||
&self,
|
||||
|
||||
@@ -611,6 +611,7 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
} else {
|
||||
GatewayChannel::new(gw_config.clone())
|
||||
};
|
||||
gw = gw.with_owner_scope(config.owner_id.clone());
|
||||
gw = gw.with_llm_provider(Arc::clone(&components.llm));
|
||||
if let Some(ref ws) = components.workspace {
|
||||
gw = gw.with_workspace(Arc::clone(ws));
|
||||
|
||||
@@ -418,6 +418,7 @@ fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option<OAuthRefr
|
||||
let oauth = auth.oauth.as_ref()?;
|
||||
|
||||
let builtin = crate::cli::oauth_defaults::builtin_credentials(&auth.secret_name);
|
||||
let exchange_proxy_url = crate::cli::oauth_defaults::exchange_proxy_url();
|
||||
|
||||
let client_id = oauth
|
||||
.client_id
|
||||
@@ -440,11 +441,21 @@ fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option<OAuthRefr
|
||||
.and_then(|env| std::env::var(env).ok())
|
||||
})
|
||||
.or_else(|| builtin.as_ref().map(|c| c.client_secret.to_string()));
|
||||
let client_secret = crate::cli::oauth_defaults::hosted_proxy_client_secret(
|
||||
&client_secret,
|
||||
builtin.as_ref(),
|
||||
exchange_proxy_url.is_some(),
|
||||
);
|
||||
let gateway_token = crate::config::helpers::env_or_override("GATEWAY_AUTH_TOKEN")
|
||||
.map(|token| token.trim().to_string())
|
||||
.filter(|token| !token.is_empty());
|
||||
|
||||
Some(OAuthRefreshConfig {
|
||||
token_url: oauth.token_url.clone(),
|
||||
client_id,
|
||||
client_secret,
|
||||
exchange_proxy_url,
|
||||
gateway_token,
|
||||
secret_name: auth.secret_name.clone(),
|
||||
provider: auth.provider.clone(),
|
||||
})
|
||||
@@ -711,9 +722,44 @@ mod tests {
|
||||
|
||||
use tempfile::TempDir;
|
||||
|
||||
use crate::config::helpers::lock_env;
|
||||
use crate::testing::credentials::{TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET};
|
||||
use crate::tools::wasm::loader::{WasmLoadError, check_wit_version_compat, discover_tools};
|
||||
|
||||
/// Restores a test-scoped env var override on drop.
|
||||
struct EnvVarGuard {
|
||||
key: String,
|
||||
previous: Option<String>,
|
||||
}
|
||||
|
||||
impl Drop for EnvVarGuard {
|
||||
fn drop(&mut self) {
|
||||
// SAFETY: Tests use lock_env() to serialize environment access.
|
||||
unsafe {
|
||||
if let Some(ref value) = self.previous {
|
||||
std::env::set_var(&self.key, value);
|
||||
} else {
|
||||
std::env::remove_var(&self.key);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn set_env_var(key: &str, value: Option<&str>) -> EnvVarGuard {
|
||||
let previous = std::env::var(key).ok();
|
||||
// SAFETY: Tests use lock_env() to serialize environment access.
|
||||
unsafe {
|
||||
match value {
|
||||
Some(value) => std::env::set_var(key, value),
|
||||
None => std::env::remove_var(key),
|
||||
}
|
||||
}
|
||||
EnvVarGuard {
|
||||
key: key.to_string(),
|
||||
previous,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wit_version_compat_none_is_ok() {
|
||||
// Pre-versioning extensions (no wit_version declared) should always pass
|
||||
@@ -871,6 +917,8 @@ mod tests {
|
||||
config.client_secret,
|
||||
Some(TEST_OAUTH_CLIENT_SECRET.to_string())
|
||||
);
|
||||
assert_eq!(config.exchange_proxy_url, None);
|
||||
assert_eq!(config.gateway_token, None);
|
||||
assert_eq!(config.secret_name, "google_oauth_token");
|
||||
assert_eq!(config.provider, Some("google".to_string()));
|
||||
}
|
||||
@@ -931,6 +979,10 @@ mod tests {
|
||||
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema,
|
||||
};
|
||||
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_EXCHANGE_URL", None);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", None);
|
||||
|
||||
// google_oauth_token should fall back to built-in credentials
|
||||
let caps = CapabilitiesFile {
|
||||
auth: Some(AuthCapabilitySchema {
|
||||
@@ -952,6 +1004,95 @@ mod tests {
|
||||
let config = config.unwrap();
|
||||
assert!(!config.client_id.is_empty());
|
||||
assert!(config.client_secret.is_some());
|
||||
assert_eq!(config.exchange_proxy_url, None);
|
||||
assert_eq!(config.gateway_token, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_oauth_refresh_config_hosted_proxy_populates_env_and_suppresses_builtin_secret()
|
||||
{
|
||||
use crate::tools::wasm::capabilities_schema::{
|
||||
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema,
|
||||
};
|
||||
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var(
|
||||
"IRONCLAW_OAUTH_EXCHANGE_URL",
|
||||
Some("https://compose-api.example.com"),
|
||||
);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token"));
|
||||
let _client_id_guard =
|
||||
set_env_var("GOOGLE_OAUTH_CLIENT_ID", Some("hosted-google-client-id"));
|
||||
|
||||
let caps = CapabilitiesFile {
|
||||
auth: Some(AuthCapabilitySchema {
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
provider: Some("google".to_string()),
|
||||
oauth: Some(OAuthConfigSchema {
|
||||
authorization_url: "https://accounts.google.com/o/oauth2/v2/auth".to_string(),
|
||||
token_url: "https://oauth2.googleapis.com/token".to_string(),
|
||||
client_id_env: Some("GOOGLE_OAUTH_CLIENT_ID".to_string()),
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = super::resolve_oauth_refresh_config(&caps).expect("hosted oauth config");
|
||||
assert_eq!(config.client_id, "hosted-google-client-id");
|
||||
assert_eq!(config.client_secret, None);
|
||||
assert_eq!(
|
||||
config.exchange_proxy_url.as_deref(),
|
||||
Some("https://compose-api.example.com")
|
||||
);
|
||||
assert_eq!(config.gateway_token.as_deref(), Some("gateway-test-token"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_oauth_refresh_config_hosted_proxy_preserves_explicit_secret() {
|
||||
use crate::tools::wasm::capabilities_schema::{
|
||||
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema,
|
||||
};
|
||||
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var(
|
||||
"IRONCLAW_OAUTH_EXCHANGE_URL",
|
||||
Some("https://compose-api.example.com"),
|
||||
);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token"));
|
||||
let _client_id_guard =
|
||||
set_env_var("GOOGLE_OAUTH_CLIENT_ID", Some("hosted-google-client-id"));
|
||||
let _client_secret_guard =
|
||||
set_env_var("GOOGLE_OAUTH_CLIENT_SECRET", Some("hosted-server-secret"));
|
||||
|
||||
let caps = CapabilitiesFile {
|
||||
auth: Some(AuthCapabilitySchema {
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
provider: Some("google".to_string()),
|
||||
oauth: Some(OAuthConfigSchema {
|
||||
authorization_url: "https://accounts.google.com/o/oauth2/v2/auth".to_string(),
|
||||
token_url: "https://oauth2.googleapis.com/token".to_string(),
|
||||
client_id_env: Some("GOOGLE_OAUTH_CLIENT_ID".to_string()),
|
||||
client_secret_env: Some("GOOGLE_OAUTH_CLIENT_SECRET".to_string()),
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = super::resolve_oauth_refresh_config(&caps).expect("hosted oauth config");
|
||||
assert_eq!(config.client_id, "hosted-google-client-id");
|
||||
assert_eq!(
|
||||
config.client_secret.as_deref(),
|
||||
Some("hosted-server-secret")
|
||||
);
|
||||
assert_eq!(
|
||||
config.exchange_proxy_url.as_deref(),
|
||||
Some("https://compose-api.example.com")
|
||||
);
|
||||
assert_eq!(config.gateway_token.as_deref(), Some("gateway-test-token"));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------
|
||||
|
||||
+443
-38
@@ -19,7 +19,7 @@ use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiView};
|
||||
use crate::context::JobContext;
|
||||
use crate::llm::recording::{HttpExchangeRequest, HttpExchangeResponse, HttpInterceptor};
|
||||
use crate::safety::LeakDetector;
|
||||
use crate::secrets::SecretsStore;
|
||||
use crate::secrets::{DecryptedSecret, SecretsStore};
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||
use crate::tools::wasm::capabilities::Capabilities;
|
||||
use crate::tools::wasm::credential_injector::{
|
||||
@@ -44,6 +44,7 @@ wasmtime::component::bindgen!({
|
||||
});
|
||||
|
||||
// Alias the export interface types for convenience.
|
||||
use crate::cli::oauth_defaults;
|
||||
use exports::near::agent::tool as wit_tool;
|
||||
|
||||
/// Configuration needed to refresh an expired OAuth access token.
|
||||
@@ -59,6 +60,10 @@ pub struct OAuthRefreshConfig {
|
||||
pub client_id: String,
|
||||
/// OAuth client_secret (optional, some providers use PKCE without a secret).
|
||||
pub client_secret: Option<String>,
|
||||
/// Hosted OAuth proxy base URL (e.g., "http://host.docker.internal:8080").
|
||||
pub exchange_proxy_url: Option<String>,
|
||||
/// Gateway auth token for authenticating with the hosted OAuth proxy.
|
||||
pub gateway_token: Option<String>,
|
||||
/// Secret name of the access token (e.g., "google_oauth_token").
|
||||
/// The refresh token lives at `{secret_name}_refresh_token`.
|
||||
pub secret_name: String,
|
||||
@@ -1210,6 +1215,53 @@ async fn refresh_oauth_token(
|
||||
user_id: &str,
|
||||
config: &OAuthRefreshConfig,
|
||||
) -> bool {
|
||||
let refresh_name = format!("{}_refresh_token", config.secret_name);
|
||||
|
||||
if let Some(proxy_url) = config.exchange_proxy_url.as_deref() {
|
||||
let Some(gateway_token) = config.gateway_token.as_deref() else {
|
||||
tracing::warn!(
|
||||
"OAuth refresh proxy is configured, but no gateway auth token is available"
|
||||
);
|
||||
return false;
|
||||
};
|
||||
|
||||
// In hosted mode, the configured exchange proxy owns the outbound token
|
||||
// refresh and validation policy for the provider token_url. Direct-mode
|
||||
// HTTPS/private-IP checks remain in place for self-hosted refreshes below.
|
||||
let refresh_secret = match load_oauth_refresh_secret(store, user_id, &refresh_name).await {
|
||||
Some(secret) => secret,
|
||||
None => return false,
|
||||
};
|
||||
let token_response = match oauth_defaults::refresh_token_via_proxy(
|
||||
oauth_defaults::ProxyRefreshTokenRequest {
|
||||
proxy_url,
|
||||
gateway_token,
|
||||
token_url: &config.token_url,
|
||||
client_id: &config.client_id,
|
||||
client_secret: config.client_secret.as_deref(),
|
||||
refresh_token: refresh_secret.expose(),
|
||||
provider: config.provider.as_deref(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
tracing::warn!(error = %error, "OAuth token refresh via proxy failed");
|
||||
return false;
|
||||
}
|
||||
};
|
||||
|
||||
return persist_refreshed_oauth_tokens(
|
||||
store,
|
||||
user_id,
|
||||
config,
|
||||
&refresh_name,
|
||||
token_response,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// SSRF defense: token_url comes from the tool's capabilities file.
|
||||
if !config.token_url.starts_with("https://") {
|
||||
tracing::warn!(
|
||||
@@ -1227,19 +1279,6 @@ async fn refresh_oauth_token(
|
||||
return false;
|
||||
}
|
||||
|
||||
let refresh_name = format!("{}_refresh_token", config.secret_name);
|
||||
let refresh_secret = match store.get_decrypted(user_id, &refresh_name).await {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
secret_name = %refresh_name,
|
||||
error = %e,
|
||||
"No refresh token available, skipping token refresh"
|
||||
);
|
||||
return false;
|
||||
}
|
||||
};
|
||||
|
||||
let client = match reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(15))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
@@ -1252,6 +1291,10 @@ async fn refresh_oauth_token(
|
||||
}
|
||||
};
|
||||
|
||||
let refresh_secret = match load_oauth_refresh_secret(store, user_id, &refresh_name).await {
|
||||
Some(secret) => secret,
|
||||
None => return false,
|
||||
};
|
||||
let mut params = vec![
|
||||
("grant_type", "refresh_token".to_string()),
|
||||
("refresh_token", refresh_secret.expose().to_string()),
|
||||
@@ -1287,22 +1330,55 @@ async fn refresh_oauth_token(
|
||||
return false;
|
||||
}
|
||||
};
|
||||
|
||||
let new_access_token = match token_data.get("access_token").and_then(|v| v.as_str()) {
|
||||
Some(t) => t,
|
||||
let token_response = match token_data.get("access_token").and_then(|v| v.as_str()) {
|
||||
Some(access_token) => oauth_defaults::OAuthTokenResponse {
|
||||
access_token: access_token.to_string(),
|
||||
refresh_token: token_data
|
||||
.get("refresh_token")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(str::to_string),
|
||||
expires_in: token_data.get("expires_in").and_then(|v| v.as_u64()),
|
||||
},
|
||||
None => {
|
||||
tracing::warn!("Token refresh response missing access_token field");
|
||||
return false;
|
||||
}
|
||||
};
|
||||
|
||||
// Store the new access token with expiry
|
||||
persist_refreshed_oauth_tokens(store, user_id, config, &refresh_name, token_response).await
|
||||
}
|
||||
|
||||
async fn load_oauth_refresh_secret(
|
||||
store: &(dyn SecretsStore + Send + Sync),
|
||||
user_id: &str,
|
||||
refresh_name: &str,
|
||||
) -> Option<DecryptedSecret> {
|
||||
match store.get_decrypted(user_id, refresh_name).await {
|
||||
Ok(secret) => Some(secret),
|
||||
Err(error) => {
|
||||
tracing::debug!(
|
||||
secret_name = %refresh_name,
|
||||
error = %error,
|
||||
"No refresh token available, skipping token refresh"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn persist_refreshed_oauth_tokens(
|
||||
store: &(dyn SecretsStore + Send + Sync),
|
||||
user_id: &str,
|
||||
config: &OAuthRefreshConfig,
|
||||
refresh_name: &str,
|
||||
token_response: oauth_defaults::OAuthTokenResponse,
|
||||
) -> bool {
|
||||
let mut access_params =
|
||||
crate::secrets::CreateSecretParams::new(&config.secret_name, new_access_token);
|
||||
crate::secrets::CreateSecretParams::new(&config.secret_name, &token_response.access_token);
|
||||
if let Some(ref provider) = config.provider {
|
||||
access_params = access_params.with_provider(provider);
|
||||
}
|
||||
if let Some(expires_in) = token_data.get("expires_in").and_then(|v| v.as_u64()) {
|
||||
if let Some(expires_in) = token_response.expires_in {
|
||||
let expires_at = chrono::Utc::now() + chrono::Duration::seconds(expires_in as i64);
|
||||
access_params = access_params.with_expiry(expires_at);
|
||||
}
|
||||
@@ -1312,10 +1388,8 @@ async fn refresh_oauth_token(
|
||||
return false;
|
||||
}
|
||||
|
||||
// Store rotated refresh token if the provider sent a new one
|
||||
if let Some(new_refresh) = token_data.get("refresh_token").and_then(|v| v.as_str()) {
|
||||
let mut refresh_params =
|
||||
crate::secrets::CreateSecretParams::new(&refresh_name, new_refresh);
|
||||
if let Some(new_refresh) = token_response.refresh_token.as_deref() {
|
||||
let mut refresh_params = crate::secrets::CreateSecretParams::new(refresh_name, new_refresh);
|
||||
if let Some(ref provider) = config.provider {
|
||||
refresh_params = refresh_params.with_provider(provider);
|
||||
}
|
||||
@@ -1664,9 +1738,18 @@ fn build_tool_usage_hint(tool_name: &str, schema: &serde_json::Value) -> String
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use axum::extract::{Form, State};
|
||||
use axum::http::HeaderMap;
|
||||
use axum::routing::post;
|
||||
use axum::{Json, Router};
|
||||
use serde_json::json;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::{Mutex as AsyncMutex, oneshot};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::context::JobContext;
|
||||
@@ -1756,6 +1839,95 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
struct RecordedProxyRequest {
|
||||
authorization: Option<String>,
|
||||
form: HashMap<String, String>,
|
||||
}
|
||||
|
||||
struct MockProxyServer {
|
||||
addr: SocketAddr,
|
||||
requests: Arc<AsyncMutex<Vec<RecordedProxyRequest>>>,
|
||||
shutdown_tx: Option<oneshot::Sender<()>>,
|
||||
server_task: Option<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl MockProxyServer {
|
||||
async fn start() -> Self {
|
||||
async fn refresh_handler(
|
||||
State(requests): State<Arc<AsyncMutex<Vec<RecordedProxyRequest>>>>,
|
||||
headers: HeaderMap,
|
||||
Form(form): Form<HashMap<String, String>>,
|
||||
) -> Json<serde_json::Value> {
|
||||
requests.lock().await.push(RecordedProxyRequest {
|
||||
authorization: headers
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string),
|
||||
form,
|
||||
});
|
||||
Json(json!({
|
||||
"access_token": "mock-refreshed-access-token",
|
||||
"refresh_token": "mock-rotated-refresh-token",
|
||||
"expires_in": 3600
|
||||
}))
|
||||
}
|
||||
|
||||
let requests = Arc::new(AsyncMutex::new(Vec::new()));
|
||||
let app = Router::new()
|
||||
.route("/oauth/refresh", post(refresh_handler))
|
||||
.with_state(Arc::clone(&requests));
|
||||
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind mock proxy");
|
||||
let addr = listener.local_addr().expect("read mock proxy addr");
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
|
||||
let server_task = tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async {
|
||||
let _ = shutdown_rx.await;
|
||||
})
|
||||
.await;
|
||||
});
|
||||
|
||||
Self {
|
||||
addr,
|
||||
requests,
|
||||
shutdown_tx: Some(shutdown_tx),
|
||||
server_task: Some(server_task),
|
||||
}
|
||||
}
|
||||
|
||||
fn base_url(&self) -> String {
|
||||
format!("http://{}", self.addr)
|
||||
}
|
||||
|
||||
async fn requests(&self) -> Vec<RecordedProxyRequest> {
|
||||
self.requests.lock().await.clone()
|
||||
}
|
||||
|
||||
async fn shutdown(mut self) {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
if let Some(task) = self.server_task.take() {
|
||||
let _ = task.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for MockProxyServer {
|
||||
fn drop(&mut self) {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
if let Some(task) = self.server_task.take() {
|
||||
task.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wrapper_creation() {
|
||||
// This test verifies the runtime can be created
|
||||
@@ -2094,8 +2266,6 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_bearer() {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::secrets::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||
};
|
||||
@@ -2141,8 +2311,6 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_owner_scope_bearer() {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::secrets::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||
};
|
||||
@@ -2188,8 +2356,6 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_execute_resolves_host_credentials_from_owner_scope_context() {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::secrets::{CredentialLocation, CredentialMapping};
|
||||
use crate::tools::wasm::capabilities::HttpCapability;
|
||||
|
||||
@@ -2239,8 +2405,6 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_missing_secret() {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::secrets::{CredentialLocation, CredentialMapping};
|
||||
use crate::tools::wasm::capabilities::HttpCapability;
|
||||
use crate::tools::wasm::wrapper::resolve_host_credentials;
|
||||
@@ -2272,8 +2436,6 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_skips_refresh_when_not_expired() {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::secrets::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||
};
|
||||
@@ -2315,6 +2477,8 @@ mod tests {
|
||||
token_url: "https://oauth2.googleapis.com/token".to_string(),
|
||||
client_id: TEST_OAUTH_CLIENT_ID.to_string(),
|
||||
client_secret: Some(TEST_OAUTH_CLIENT_SECRET.to_string()),
|
||||
exchange_proxy_url: None,
|
||||
gateway_token: None,
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
provider: Some("google".to_string()),
|
||||
};
|
||||
@@ -2331,8 +2495,6 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_skips_refresh_no_config() {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::secrets::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||
};
|
||||
@@ -2376,8 +2538,6 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_skips_refresh_no_expires_at() {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::secrets::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||
};
|
||||
@@ -2417,6 +2577,8 @@ mod tests {
|
||||
token_url: "https://oauth2.googleapis.com/token".to_string(),
|
||||
client_id: TEST_OAUTH_CLIENT_ID.to_string(),
|
||||
client_secret: Some(TEST_OAUTH_CLIENT_SECRET.to_string()),
|
||||
exchange_proxy_url: None,
|
||||
gateway_token: None,
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
provider: Some("google".to_string()),
|
||||
};
|
||||
@@ -2431,6 +2593,249 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_refreshes_via_proxy_without_direct_token_url_validation()
|
||||
{
|
||||
use crate::secrets::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||
};
|
||||
use crate::tools::wasm::capabilities::HttpCapability;
|
||||
use crate::tools::wasm::wrapper::{OAuthRefreshConfig, resolve_host_credentials};
|
||||
|
||||
let proxy = MockProxyServer::start().await;
|
||||
let store = test_secrets_store();
|
||||
|
||||
store
|
||||
.create(
|
||||
"user1",
|
||||
CreateSecretParams::new("google_oauth_token", "expired-access-token")
|
||||
.with_expiry(chrono::Utc::now() - chrono::Duration::hours(1)),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.create(
|
||||
"user1",
|
||||
CreateSecretParams::new("google_oauth_token_refresh_token", "stored-refresh-token"),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut credentials = HashMap::new();
|
||||
credentials.insert(
|
||||
"google_oauth_token".to_string(),
|
||||
CredentialMapping {
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
location: CredentialLocation::AuthorizationBearer,
|
||||
host_patterns: vec!["www.googleapis.com".to_string()],
|
||||
},
|
||||
);
|
||||
|
||||
let caps = Capabilities {
|
||||
http: Some(HttpCapability {
|
||||
credentials,
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let oauth_config = OAuthRefreshConfig {
|
||||
token_url: "http://127.0.0.1:9/provider-token-endpoint".to_string(),
|
||||
client_id: "hosted-google-client-id".to_string(),
|
||||
client_secret: None,
|
||||
exchange_proxy_url: Some(proxy.base_url()),
|
||||
gateway_token: Some("gateway-test-token".to_string()),
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
provider: Some("google".to_string()),
|
||||
};
|
||||
|
||||
let resolved =
|
||||
resolve_host_credentials(&caps, Some(&store), "user1", Some(&oauth_config)).await;
|
||||
assert_eq!(resolved.len(), 1);
|
||||
assert_eq!(
|
||||
resolved[0].headers.get("Authorization"),
|
||||
Some(&"Bearer mock-refreshed-access-token".to_string())
|
||||
);
|
||||
|
||||
let access_secret = store.get("user1", "google_oauth_token").await.unwrap();
|
||||
assert!(
|
||||
access_secret
|
||||
.expires_at
|
||||
.expect("refreshed access token expiry")
|
||||
> chrono::Utc::now()
|
||||
);
|
||||
let access_value = store
|
||||
.get_decrypted("user1", "google_oauth_token")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(access_value.expose(), "mock-refreshed-access-token");
|
||||
|
||||
let refresh_value = store
|
||||
.get_decrypted("user1", "google_oauth_token_refresh_token")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(refresh_value.expose(), "mock-rotated-refresh-token");
|
||||
|
||||
let requests = proxy.requests().await;
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(
|
||||
requests[0].authorization.as_deref(),
|
||||
Some("Bearer gateway-test-token")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("client_id").map(String::as_str),
|
||||
Some("hosted-google-client-id")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("token_url").map(String::as_str),
|
||||
Some("http://127.0.0.1:9/provider-token-endpoint")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("refresh_token").map(String::as_str),
|
||||
Some("stored-refresh-token")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("provider").map(String::as_str),
|
||||
Some("google")
|
||||
);
|
||||
assert!(!requests[0].form.contains_key("client_secret"));
|
||||
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_skips_refresh_token_lookup_without_gateway_token() {
|
||||
use crate::secrets::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||
};
|
||||
use crate::tools::wasm::capabilities::HttpCapability;
|
||||
use crate::tools::wasm::wrapper::{OAuthRefreshConfig, resolve_host_credentials};
|
||||
|
||||
let store = RecordingSecretsStore::new();
|
||||
|
||||
store
|
||||
.create(
|
||||
"user1",
|
||||
CreateSecretParams::new("google_oauth_token", "expired-access-token")
|
||||
.with_expiry(chrono::Utc::now() - chrono::Duration::hours(1)),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.create(
|
||||
"user1",
|
||||
CreateSecretParams::new("google_oauth_token_refresh_token", "stored-refresh-token"),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut credentials = HashMap::new();
|
||||
credentials.insert(
|
||||
"google_oauth_token".to_string(),
|
||||
CredentialMapping {
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
location: CredentialLocation::AuthorizationBearer,
|
||||
host_patterns: vec!["www.googleapis.com".to_string()],
|
||||
},
|
||||
);
|
||||
|
||||
let caps = Capabilities {
|
||||
http: Some(HttpCapability {
|
||||
credentials,
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let oauth_config = OAuthRefreshConfig {
|
||||
token_url: "https://oauth2.googleapis.com/token".to_string(),
|
||||
client_id: "hosted-google-client-id".to_string(),
|
||||
client_secret: None,
|
||||
exchange_proxy_url: Some("https://compose-api.example.com".to_string()),
|
||||
gateway_token: None,
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
provider: Some("google".to_string()),
|
||||
};
|
||||
|
||||
let resolved =
|
||||
resolve_host_credentials(&caps, Some(&store), "user1", Some(&oauth_config)).await;
|
||||
assert!(resolved.is_empty());
|
||||
|
||||
let lookups = store.decrypted_lookups();
|
||||
assert!(lookups.contains(&("user1".to_string(), "google_oauth_token".to_string())));
|
||||
assert!(!lookups.contains(&(
|
||||
"user1".to_string(),
|
||||
"google_oauth_token_refresh_token".to_string(),
|
||||
)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_skips_refresh_token_lookup_for_invalid_direct_token_url()
|
||||
{
|
||||
use crate::secrets::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||
};
|
||||
use crate::tools::wasm::capabilities::HttpCapability;
|
||||
use crate::tools::wasm::wrapper::{OAuthRefreshConfig, resolve_host_credentials};
|
||||
|
||||
let store = RecordingSecretsStore::new();
|
||||
|
||||
store
|
||||
.create(
|
||||
"user1",
|
||||
CreateSecretParams::new("google_oauth_token", "expired-access-token")
|
||||
.with_expiry(chrono::Utc::now() - chrono::Duration::hours(1)),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.create(
|
||||
"user1",
|
||||
CreateSecretParams::new("google_oauth_token_refresh_token", "stored-refresh-token"),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut credentials = HashMap::new();
|
||||
credentials.insert(
|
||||
"google_oauth_token".to_string(),
|
||||
CredentialMapping {
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
location: CredentialLocation::AuthorizationBearer,
|
||||
host_patterns: vec!["www.googleapis.com".to_string()],
|
||||
},
|
||||
);
|
||||
|
||||
let caps = Capabilities {
|
||||
http: Some(HttpCapability {
|
||||
credentials,
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let oauth_config = OAuthRefreshConfig {
|
||||
token_url: "http://127.0.0.1:9/provider-token-endpoint".to_string(),
|
||||
client_id: TEST_OAUTH_CLIENT_ID.to_string(),
|
||||
client_secret: Some(TEST_OAUTH_CLIENT_SECRET.to_string()),
|
||||
exchange_proxy_url: None,
|
||||
gateway_token: None,
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
provider: Some("google".to_string()),
|
||||
};
|
||||
|
||||
let resolved =
|
||||
resolve_host_credentials(&caps, Some(&store), "user1", Some(&oauth_config)).await;
|
||||
assert!(resolved.is_empty());
|
||||
|
||||
let lookups = store.decrypted_lookups();
|
||||
assert!(lookups.contains(&("user1".to_string(), "google_oauth_token".to_string())));
|
||||
assert!(!lookups.contains(&(
|
||||
"user1".to_string(),
|
||||
"google_oauth_token_refresh_token".to_string(),
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_private_ip_v4() {
|
||||
use std::net::IpAddr;
|
||||
|
||||
@@ -111,10 +111,23 @@ impl Tunnel for CloudflareTunnel {
|
||||
}
|
||||
}
|
||||
|
||||
// Drain stderr in the background to prevent SIGPIPE/buffer stalls.
|
||||
tokio::spawn(async move { while let Ok(Some(_)) = reader.next_line().await {} });
|
||||
if let Ok(mut guard) = self.url.write() {
|
||||
*guard = Some(public_url.clone());
|
||||
}
|
||||
|
||||
// Drain stdout silently.
|
||||
// We took ownership of cloudflared's stderr pipe above to parse the URL.
|
||||
// cloudflared continues writing logs for its entire lifetime. If we drop
|
||||
// the reader, the pipe closes and cloudflared gets SIGPIPE on its next
|
||||
// write. We can't just store the reader without reading — the OS pipe
|
||||
// buffer fills up and cloudflared blocks. So we drain it in a background
|
||||
// task. The task exits naturally when cloudflared is killed (EOF).
|
||||
let drain_handle = tokio::spawn(async move {
|
||||
while let Ok(Some(line)) = reader.next_line().await {
|
||||
tracing::trace!("cloudflared: {line}");
|
||||
}
|
||||
});
|
||||
|
||||
// Drain stdout silently to prevent SIGPIPE/buffer stalls.
|
||||
if let Some(stdout) = stdout {
|
||||
tokio::spawn(async move {
|
||||
let mut out_reader = tokio::io::BufReader::new(stdout).lines();
|
||||
@@ -122,12 +135,11 @@ impl Tunnel for CloudflareTunnel {
|
||||
});
|
||||
}
|
||||
|
||||
if let Ok(mut guard) = self.url.write() {
|
||||
*guard = Some(public_url.clone());
|
||||
}
|
||||
|
||||
let mut guard = self.proc.lock().await;
|
||||
*guard = Some(TunnelProcess { child });
|
||||
*guard = Some(TunnelProcess {
|
||||
child,
|
||||
_pipe_drain: Some(drain_handle),
|
||||
});
|
||||
|
||||
Ok(public_url)
|
||||
}
|
||||
|
||||
+18
-5
@@ -73,6 +73,7 @@ impl Tunnel for CustomTunnel {
|
||||
let stderr = child.stderr.take();
|
||||
|
||||
let mut public_url = format!("http://{local_host}:{local_port}");
|
||||
let mut drain_handle: Option<tokio::task::JoinHandle<()>> = None;
|
||||
|
||||
if self.url_pattern.is_some()
|
||||
&& let Some(stdout) = stdout
|
||||
@@ -103,17 +104,26 @@ impl Tunnel for CustomTunnel {
|
||||
Err(_) => {}
|
||||
}
|
||||
}
|
||||
// Drain remaining stdout to prevent SIGPIPE/buffer stalls.
|
||||
tokio::spawn(async move { while let Ok(Some(_)) = reader.next_line().await {} });
|
||||
// We took ownership of the process's stdout pipe above to parse the
|
||||
// URL. The process may continue writing to stdout for its lifetime.
|
||||
// If we drop the reader, the pipe closes and the process gets SIGPIPE.
|
||||
// We can't just store the reader without reading — the OS pipe buffer
|
||||
// fills up and the process blocks. So we drain it in a background task.
|
||||
// The task exits naturally when the process is killed (EOF).
|
||||
drain_handle = Some(tokio::spawn(async move {
|
||||
while let Ok(Some(line)) = reader.next_line().await {
|
||||
tracing::trace!("custom-tunnel: {line}");
|
||||
}
|
||||
}));
|
||||
} else if let Some(stdout) = stdout {
|
||||
// No url_pattern: still drain stdout to prevent pipe stalls.
|
||||
// No url_pattern: still drain stdout to prevent SIGPIPE/buffer stalls.
|
||||
tokio::spawn(async move {
|
||||
let mut reader = tokio::io::BufReader::new(stdout).lines();
|
||||
while let Ok(Some(_)) = reader.next_line().await {}
|
||||
});
|
||||
}
|
||||
|
||||
// Drain stderr silently.
|
||||
// Drain stderr to prevent SIGPIPE/buffer stalls.
|
||||
if let Some(stderr) = stderr {
|
||||
tokio::spawn(async move {
|
||||
let mut reader = tokio::io::BufReader::new(stderr).lines();
|
||||
@@ -126,7 +136,10 @@ impl Tunnel for CustomTunnel {
|
||||
}
|
||||
|
||||
let mut guard = self.proc.lock().await;
|
||||
*guard = Some(TunnelProcess { child });
|
||||
*guard = Some(TunnelProcess {
|
||||
child,
|
||||
_pipe_drain: drain_handle,
|
||||
});
|
||||
|
||||
Ok(public_url)
|
||||
}
|
||||
|
||||
+126
-16
@@ -66,6 +66,10 @@ pub trait Tunnel: Send + Sync {
|
||||
/// Wraps a spawned tunnel child process.
|
||||
pub(crate) struct TunnelProcess {
|
||||
pub child: tokio::process::Child,
|
||||
/// Background task that drains the process's output pipe (stdout or stderr).
|
||||
/// Must stay alive or the process dies (SIGPIPE from closed pipe) or hangs
|
||||
/// (OS pipe buffer fills up, blocking the process's writes).
|
||||
pub _pipe_drain: Option<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
pub(crate) type SharedProcess = Arc<Mutex<Option<TunnelProcess>>>;
|
||||
@@ -182,6 +186,22 @@ pub fn create_tunnel(config: &TunnelProviderConfig) -> Result<Option<Box<dyn Tun
|
||||
|
||||
// ── Managed tunnel startup ───────────────────────────────────────
|
||||
|
||||
/// Determine which local address the tunnel should forward traffic to.
|
||||
///
|
||||
/// Prefers the webhook server (`HTTP_PORT`) since that's where webhook routes
|
||||
/// (Telegram, etc.) are served. Falls back to the gateway port if configured,
|
||||
/// otherwise defaults to 0.0.0.0:8080 (the same fallback the webhook server
|
||||
/// uses in main.rs when no HTTP config is present).
|
||||
fn resolve_tunnel_target(channels: &crate::config::ChannelsConfig) -> (&str, u16) {
|
||||
if let Some(ref http) = channels.http {
|
||||
return (http.host.as_str(), http.port);
|
||||
}
|
||||
if let Some(ref gw) = channels.gateway {
|
||||
return (gw.host.as_str(), gw.port);
|
||||
}
|
||||
("0.0.0.0", 8080)
|
||||
}
|
||||
|
||||
/// Start a managed tunnel if configured and no static URL is already set.
|
||||
///
|
||||
/// Returns the (potentially mutated) config with `tunnel.public_url` set,
|
||||
@@ -201,28 +221,17 @@ pub async fn start_managed_tunnel(
|
||||
return (config, None);
|
||||
};
|
||||
|
||||
let gateway_port = config
|
||||
.channels
|
||||
.gateway
|
||||
.as_ref()
|
||||
.map(|g| g.port)
|
||||
.unwrap_or(3000);
|
||||
let gateway_host = config
|
||||
.channels
|
||||
.gateway
|
||||
.as_ref()
|
||||
.map(|g| g.host.as_str())
|
||||
.unwrap_or("127.0.0.1");
|
||||
let (tunnel_host, tunnel_port) = resolve_tunnel_target(&config.channels);
|
||||
|
||||
match create_tunnel(provider_config) {
|
||||
Ok(Some(tunnel)) => {
|
||||
tracing::debug!(
|
||||
"Starting {} tunnel on {}:{}...",
|
||||
tunnel.name(),
|
||||
gateway_host,
|
||||
gateway_port
|
||||
tunnel_host,
|
||||
tunnel_port
|
||||
);
|
||||
match tunnel.start(gateway_host, gateway_port).await {
|
||||
match tunnel.start(tunnel_host, tunnel_port).await {
|
||||
Ok(url) => {
|
||||
tracing::debug!("Tunnel started: {}", url);
|
||||
config.tunnel.public_url = Some(url);
|
||||
@@ -383,10 +392,111 @@ mod tests {
|
||||
|
||||
{
|
||||
let mut guard = proc.lock().await;
|
||||
*guard = Some(TunnelProcess { child });
|
||||
*guard = Some(TunnelProcess {
|
||||
child,
|
||||
_pipe_drain: None,
|
||||
});
|
||||
}
|
||||
|
||||
kill_shared(&proc).await.unwrap();
|
||||
assert!(proc.lock().await.is_none());
|
||||
}
|
||||
|
||||
// ── Port selection regression tests ──────────────────────────────
|
||||
|
||||
fn base_channels() -> crate::config::ChannelsConfig {
|
||||
crate::config::ChannelsConfig {
|
||||
cli: crate::config::CliConfig { enabled: false },
|
||||
http: None,
|
||||
gateway: None,
|
||||
signal: None,
|
||||
wasm_channels_dir: std::env::temp_dir().join("ironclaw-test-channels"),
|
||||
wasm_channels_enabled: false,
|
||||
wasm_channel_owner_ids: std::collections::HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn channels_with_http(host: &str, port: u16) -> crate::config::ChannelsConfig {
|
||||
let mut c = base_channels();
|
||||
c.http = Some(crate::config::HttpConfig {
|
||||
host: host.to_string(),
|
||||
port,
|
||||
webhook_secret: None,
|
||||
user_id: "test".to_string(),
|
||||
});
|
||||
c.gateway = Some(crate::config::GatewayConfig {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 3000,
|
||||
auth_token: None,
|
||||
user_id: "test".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
memory_layers: Vec::new(),
|
||||
user_tokens: None,
|
||||
});
|
||||
c
|
||||
}
|
||||
|
||||
fn channels_gateway_only(host: &str, port: u16) -> crate::config::ChannelsConfig {
|
||||
let mut c = base_channels();
|
||||
c.gateway = Some(crate::config::GatewayConfig {
|
||||
host: host.to_string(),
|
||||
port,
|
||||
auth_token: None,
|
||||
user_id: "test".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
memory_layers: Vec::new(),
|
||||
user_tokens: None,
|
||||
});
|
||||
c
|
||||
}
|
||||
|
||||
fn channels_neither() -> crate::config::ChannelsConfig {
|
||||
base_channels()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_target_prefers_http_port() {
|
||||
let channels = channels_with_http("0.0.0.0", 8080);
|
||||
let (host, port) = resolve_tunnel_target(&channels);
|
||||
assert_eq!(host, "0.0.0.0"); // safety: test-only
|
||||
assert_eq!(port, 8080); // safety: test-only
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_target_falls_back_to_gateway() {
|
||||
let channels = channels_gateway_only("10.0.0.1", 4000);
|
||||
let (host, port) = resolve_tunnel_target(&channels);
|
||||
assert_eq!(host, "10.0.0.1"); // safety: test-only
|
||||
assert_eq!(port, 4000); // safety: test-only
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_target_defaults_to_webhook_fallback() {
|
||||
let channels = channels_neither();
|
||||
let (host, port) = resolve_tunnel_target(&channels);
|
||||
// Matches the webhook server's hardcoded fallback in main.rs
|
||||
assert_eq!(host, "0.0.0.0"); // safety: test-only
|
||||
assert_eq!(port, 8080); // safety: test-only
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_target_http_takes_priority_over_gateway() {
|
||||
let channels = channels_with_http("192.168.1.1", 9090);
|
||||
let (host, port) = resolve_tunnel_target(&channels);
|
||||
// Should use HTTP config, not gateway's 127.0.0.1:3000
|
||||
assert_eq!(host, "192.168.1.1"); // safety: test-only
|
||||
assert_eq!(port, 9090); // safety: test-only
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_target_no_http_no_gateway_matches_webhook_fallback() {
|
||||
// When HTTP_PORT is not set and gateway is not configured (e.g. WASM
|
||||
// channels exist but no explicit HTTP config), the webhook server in
|
||||
// main.rs binds to 0.0.0.0:8080 as a hardcoded fallback. The tunnel
|
||||
// must target the same address so webhook traffic reaches the right
|
||||
// server.
|
||||
let channels = channels_neither();
|
||||
let (host, port) = resolve_tunnel_target(&channels);
|
||||
assert_eq!((host, port), ("0.0.0.0", 8080)); // safety: test-only
|
||||
}
|
||||
}
|
||||
|
||||
+21
-10
@@ -110,12 +110,24 @@ impl Tunnel for NgrokTunnel {
|
||||
}
|
||||
}
|
||||
|
||||
// Drain stdout silently — ngrok only emits low-level connection events
|
||||
// to stdout; the pipe must be consumed to prevent SIGPIPE/buffer stalls.
|
||||
tokio::spawn(async move { while let Ok(Some(_)) = reader.next_line().await {} });
|
||||
if let Ok(mut guard) = self.url.write() {
|
||||
*guard = Some(public_url.clone());
|
||||
}
|
||||
|
||||
// Drain stderr silently — with --log stdout all meaningful output goes
|
||||
// to stdout; stderr only needs to be consumed to prevent pipe stalls.
|
||||
// We took ownership of ngrok's stdout pipe above to parse the URL.
|
||||
// ngrok continues writing logs to stdout for its entire lifetime.
|
||||
// If we drop the reader, the pipe closes and ngrok gets SIGPIPE on
|
||||
// its next write → process dies. We can't just store the reader
|
||||
// without reading — the OS pipe buffer (~64KB) fills up and ngrok
|
||||
// blocks. So we drain it in a background task. The task exits
|
||||
// naturally when ngrok is killed (EOF on the pipe).
|
||||
let drain_handle = tokio::spawn(async move {
|
||||
while let Ok(Some(line)) = reader.next_line().await {
|
||||
tracing::trace!("ngrok: {line}");
|
||||
}
|
||||
});
|
||||
|
||||
// Drain stderr silently to prevent SIGPIPE/buffer stalls.
|
||||
if let Some(stderr) = stderr {
|
||||
tokio::spawn(async move {
|
||||
let mut err_reader = tokio::io::BufReader::new(stderr).lines();
|
||||
@@ -123,12 +135,11 @@ impl Tunnel for NgrokTunnel {
|
||||
});
|
||||
}
|
||||
|
||||
if let Ok(mut guard) = self.url.write() {
|
||||
*guard = Some(public_url.clone());
|
||||
}
|
||||
|
||||
let mut guard = self.proc.lock().await;
|
||||
*guard = Some(TunnelProcess { child });
|
||||
*guard = Some(TunnelProcess {
|
||||
child,
|
||||
_pipe_drain: Some(drain_handle),
|
||||
});
|
||||
|
||||
Ok(public_url)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,191 @@
|
||||
//! Tests for batch_get_last_run_status (#1469 N+1 fix).
|
||||
//!
|
||||
//! Verifies:
|
||||
//! 1. Empty input returns empty map
|
||||
//! 2. Returns the most recent run status per routine
|
||||
//! 3. Routines with no runs are omitted from result
|
||||
//! 4. Multiple routines with different statuses are correctly returned
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::{Duration, Utc};
|
||||
use uuid::Uuid;
|
||||
|
||||
use ironclaw::agent::routine::{
|
||||
Routine, RoutineAction, RoutineGuardrails, RoutineRun, RunStatus, Trigger,
|
||||
};
|
||||
use ironclaw::db::Database;
|
||||
|
||||
async fn create_test_db() -> (Arc<dyn Database>, tempfile::TempDir) {
|
||||
use ironclaw::db::libsql::LibSqlBackend;
|
||||
|
||||
let temp_dir = tempfile::tempdir().expect("tempdir");
|
||||
let db_path = temp_dir.path().join("test.db");
|
||||
let backend = LibSqlBackend::new_local(&db_path)
|
||||
.await
|
||||
.expect("LibSqlBackend");
|
||||
backend.run_migrations().await.expect("migrations");
|
||||
let db: Arc<dyn Database> = Arc::new(backend);
|
||||
(db, temp_dir)
|
||||
}
|
||||
|
||||
fn make_routine(id: Uuid) -> Routine {
|
||||
Routine {
|
||||
id,
|
||||
name: format!("test-routine-{}", id),
|
||||
description: "Test routine".to_string(),
|
||||
user_id: "default".to_string(),
|
||||
enabled: true,
|
||||
trigger: Trigger::Manual,
|
||||
action: RoutineAction::FullJob {
|
||||
title: "Test job".to_string(),
|
||||
description: "Test description".to_string(),
|
||||
max_iterations: 5,
|
||||
},
|
||||
guardrails: RoutineGuardrails {
|
||||
cooldown: std::time::Duration::from_secs(0),
|
||||
max_concurrent: 1,
|
||||
dedup_window: None,
|
||||
},
|
||||
notify: Default::default(),
|
||||
last_run_at: None,
|
||||
next_fire_at: None,
|
||||
run_count: 0,
|
||||
consecutive_failures: 0,
|
||||
state: serde_json::json!({}),
|
||||
created_at: Utc::now(),
|
||||
updated_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
fn make_run(
|
||||
routine_id: Uuid,
|
||||
status: RunStatus,
|
||||
started_at: chrono::DateTime<chrono::Utc>,
|
||||
) -> RoutineRun {
|
||||
RoutineRun {
|
||||
id: Uuid::new_v4(),
|
||||
routine_id,
|
||||
trigger_type: "manual".to_string(),
|
||||
trigger_detail: None,
|
||||
started_at,
|
||||
completed_at: if status == RunStatus::Running {
|
||||
None
|
||||
} else {
|
||||
Some(Utc::now())
|
||||
},
|
||||
status,
|
||||
result_summary: None,
|
||||
tokens_used: None,
|
||||
job_id: None,
|
||||
created_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_batch_get_last_run_status_empty_input() {
|
||||
let (db, _tmp) = create_test_db().await;
|
||||
let result = db
|
||||
.batch_get_last_run_status(&[])
|
||||
.await
|
||||
.expect("batch query");
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_batch_get_last_run_status_returns_latest() {
|
||||
let (db, _tmp) = create_test_db().await;
|
||||
|
||||
let routine_id = Uuid::new_v4();
|
||||
db.create_routine(&make_routine(routine_id))
|
||||
.await
|
||||
.expect("create routine");
|
||||
|
||||
// Create an older run with Ok status
|
||||
let older_run = make_run(routine_id, RunStatus::Ok, Utc::now() - Duration::hours(2));
|
||||
db.create_routine_run(&older_run)
|
||||
.await
|
||||
.expect("create older run");
|
||||
db.complete_routine_run(older_run.id, RunStatus::Ok, None, None)
|
||||
.await
|
||||
.expect("complete older run");
|
||||
|
||||
// Create a newer run with Attention status
|
||||
let newer_run = make_run(
|
||||
routine_id,
|
||||
RunStatus::Attention,
|
||||
Utc::now() - Duration::hours(1),
|
||||
);
|
||||
db.create_routine_run(&newer_run)
|
||||
.await
|
||||
.expect("create newer run");
|
||||
db.complete_routine_run(newer_run.id, RunStatus::Attention, None, None)
|
||||
.await
|
||||
.expect("complete newer run");
|
||||
|
||||
let result = db
|
||||
.batch_get_last_run_status(&[routine_id])
|
||||
.await
|
||||
.expect("batch query");
|
||||
assert_eq!(result.get(&routine_id), Some(&RunStatus::Attention));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_batch_get_last_run_status_omits_routines_without_runs() {
|
||||
let (db, _tmp) = create_test_db().await;
|
||||
|
||||
let with_runs = Uuid::new_v4();
|
||||
let without_runs = Uuid::new_v4();
|
||||
db.create_routine(&make_routine(with_runs))
|
||||
.await
|
||||
.expect("create routine");
|
||||
db.create_routine(&make_routine(without_runs))
|
||||
.await
|
||||
.expect("create routine");
|
||||
|
||||
let run = make_run(with_runs, RunStatus::Ok, Utc::now());
|
||||
db.create_routine_run(&run).await.expect("create run");
|
||||
db.complete_routine_run(run.id, RunStatus::Ok, None, None)
|
||||
.await
|
||||
.expect("complete run");
|
||||
|
||||
let result = db
|
||||
.batch_get_last_run_status(&[with_runs, without_runs])
|
||||
.await
|
||||
.expect("batch query");
|
||||
assert_eq!(result.get(&with_runs), Some(&RunStatus::Ok));
|
||||
assert_eq!(result.get(&without_runs), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_batch_get_last_run_status_multiple_routines() {
|
||||
let (db, _tmp) = create_test_db().await;
|
||||
|
||||
let r1 = Uuid::new_v4();
|
||||
let r2 = Uuid::new_v4();
|
||||
db.create_routine(&make_routine(r1))
|
||||
.await
|
||||
.expect("create r1");
|
||||
db.create_routine(&make_routine(r2))
|
||||
.await
|
||||
.expect("create r2");
|
||||
|
||||
let run1 = make_run(r1, RunStatus::Running, Utc::now());
|
||||
db.create_routine_run(&run1).await.expect("create run1");
|
||||
|
||||
let run2 = make_run(r2, RunStatus::Failed, Utc::now());
|
||||
db.create_routine_run(&run2).await.expect("create run2");
|
||||
db.complete_routine_run(run2.id, RunStatus::Failed, None, None)
|
||||
.await
|
||||
.expect("complete run2");
|
||||
|
||||
let result = db
|
||||
.batch_get_last_run_status(&[r1, r2])
|
||||
.await
|
||||
.expect("batch query");
|
||||
assert_eq!(result.get(&r1), Some(&RunStatus::Running));
|
||||
assert_eq!(result.get(&r2), Some(&RunStatus::Failed));
|
||||
}
|
||||
}
|
||||
@@ -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_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`
|
||||
|
||||
@@ -75,6 +76,7 @@ All fixtures are defined in `tests/e2e/conftest.py`. Running `pytest scenarios/`
|
||||
| `ironclaw_binary` | Checks `target/debug/ironclaw`; if absent, runs `cargo build --no-default-features --features libsql` (timeout 600s). |
|
||||
| `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. |
|
||||
| `browser` | Launches a single Chromium instance (headless by default; set `HEADED=1` for headed). Shared across all tests. |
|
||||
|
||||
### Function-scoped fixtures
|
||||
@@ -100,6 +102,8 @@ EMBEDDING_ENABLED=false, SKILLS_ENABLED=true
|
||||
ONBOARD_COMPLETED=true # prevents setup wizard
|
||||
```
|
||||
|
||||
The `hosted_oauth_refresh_server` fixture uses the same baseline, but with its own DB/home tempdirs and `GOOGLE_OAUTH_CLIENT_ID=hosted-google-client-id` so hosted OAuth flows exercise proxy credential injection instead of the baked-in desktop Google app.
|
||||
|
||||
The binary is also started with `--no-onboard`. Coverage env vars (`CARGO_LLVM_COV*`, `LLVM_*`, `CARGO_ENCODED_RUSTFLAGS`, `CARGO_INCREMENTAL`) are forwarded from the outer environment when present.
|
||||
|
||||
## Mock LLM (`mock_llm.py`)
|
||||
@@ -113,6 +117,11 @@ python mock_llm.py --port 0
|
||||
|
||||
It serves `POST /v1/chat/completions` (streaming + non-streaming) and `GET /v1/models`. Responses are pattern-matched from `CANNED_RESPONSES` against the last user message. Unmatched messages return `"I understand your request."`. The model name reported is always `"mock-model"`.
|
||||
|
||||
It also hosts OAuth test endpoints:
|
||||
- `POST /oauth/exchange` for hosted auth-code exchange
|
||||
- `POST /oauth/refresh` for hosted refresh-token exchange
|
||||
- `GET /__mock/oauth/state` and `POST /__mock/oauth/reset` so HTTP E2E scenarios can assert exact proxy payloads and reset counters between setup and refresh assertions
|
||||
|
||||
To add a new canned response:
|
||||
```python
|
||||
# In mock_llm.py
|
||||
|
||||
+169
-34
@@ -112,6 +112,39 @@ def _reserve_loopback_sockets(count: int) -> list[socket.socket]:
|
||||
sock.close()
|
||||
raise
|
||||
|
||||
async def _stop_process(
|
||||
proc: asyncio.subprocess.Process, *, sig: int | None = None, timeout: float
|
||||
) -> None:
|
||||
"""Signal a subprocess and wait briefly without masking exit races."""
|
||||
if proc.returncode is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
if sig is None:
|
||||
proc.kill()
|
||||
else:
|
||||
proc.send_signal(sig)
|
||||
except ProcessLookupError:
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=timeout)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
return
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=timeout)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
|
||||
def _forward_coverage_env(env: dict[str, str]) -> None:
|
||||
"""Forward cargo-llvm-cov env vars into child processes when present."""
|
||||
cov_env_prefixes = ("CARGO_LLVM_COV", "LLVM_")
|
||||
cov_env_extras = ("CARGO_ENCODED_RUSTFLAGS", "CARGO_INCREMENTAL")
|
||||
for key, val in os.environ.items():
|
||||
if key.startswith(cov_env_prefixes) or key in cov_env_extras:
|
||||
env[key] = val
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def ironclaw_binary():
|
||||
@@ -264,14 +297,7 @@ async def ironclaw_server(
|
||||
"IRONCLAW_OAUTH_CALLBACK_URL": "https://oauth.test.example/oauth/callback",
|
||||
"IRONCLAW_OAUTH_EXCHANGE_URL": mock_llm_server,
|
||||
}
|
||||
# Forward LLVM coverage instrumentation env vars when present
|
||||
# (allows cargo-llvm-cov to collect profraw data from E2E runs).
|
||||
# Use prefix matching to stay resilient to cargo-llvm-cov changes.
|
||||
COV_ENV_PREFIXES = ("CARGO_LLVM_COV", "LLVM_")
|
||||
COV_ENV_EXTRAS = ("CARGO_ENCODED_RUSTFLAGS", "CARGO_INCREMENTAL")
|
||||
for key, val in os.environ.items():
|
||||
if key.startswith(COV_ENV_PREFIXES) or key in COV_ENV_EXTRAS:
|
||||
env[key] = val
|
||||
_forward_coverage_env(env)
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
ironclaw_binary, "--no-onboard",
|
||||
stdin=asyncio.subprocess.DEVNULL,
|
||||
@@ -279,35 +305,145 @@ async def ironclaw_server(
|
||||
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
|
||||
except TimeoutError:
|
||||
# Dump stderr so CI logs show why the server failed to start
|
||||
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, Exception):
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
stderr_text = stderr_bytes.decode("utf-8", errors="replace")
|
||||
proc.kill()
|
||||
pytest.fail(
|
||||
f"ironclaw server failed to start on port {gateway_port} "
|
||||
f"(returncode={returncode}).\nstderr:\n{stderr_text}"
|
||||
)
|
||||
finally:
|
||||
if proc.returncode is None:
|
||||
# Use SIGINT (not SIGTERM) so tokio's ctrl_c handler triggers a
|
||||
# graceful shutdown. This lets the LLVM coverage runtime run its
|
||||
# atexit handler and flush .profraw files for cargo-llvm-cov.
|
||||
proc.send_signal(signal.SIGINT)
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=10)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
if startup_kill_attempted:
|
||||
await _stop_process(proc, timeout=2)
|
||||
else:
|
||||
# Use SIGINT (not SIGTERM) so tokio's ctrl_c handler triggers a
|
||||
# graceful shutdown. This lets the LLVM coverage runtime run its
|
||||
# atexit handler and flush .profraw files for cargo-llvm-cov.
|
||||
await _stop_process(proc, sig=signal.SIGINT, timeout=10)
|
||||
if proc.returncode is None:
|
||||
await _stop_process(proc, timeout=2)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
async def hosted_oauth_refresh_server(
|
||||
ironclaw_binary,
|
||||
mock_llm_server,
|
||||
wasm_tools_dir,
|
||||
):
|
||||
"""Start a hosted-mode ironclaw instance for OAuth refresh regression tests."""
|
||||
reserved = _reserve_loopback_sockets(2)
|
||||
db_tmpdir = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-hosted-oauth-db-")
|
||||
home_tmpdir = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-hosted-oauth-home-")
|
||||
|
||||
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, "hosted-oauth-refresh.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": wasm_tools_dir,
|
||||
"WASM_CHANNELS_DIR": _WASM_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"hosted oauth refresh 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()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
@@ -362,12 +498,7 @@ async def http_channel_server_without_secret(
|
||||
"IRONCLAW_OAUTH_CALLBACK_URL": "https://oauth.test.example/oauth/callback",
|
||||
"IRONCLAW_OAUTH_EXCHANGE_URL": mock_llm_server,
|
||||
}
|
||||
# Forward LLVM coverage instrumentation env vars when present
|
||||
COV_ENV_PREFIXES = ("CARGO_LLVM_COV", "LLVM_")
|
||||
COV_ENV_EXTRAS = ("CARGO_ENCODED_RUSTFLAGS", "CARGO_INCREMENTAL")
|
||||
for key, val in os.environ.items():
|
||||
if key.startswith(COV_ENV_PREFIXES) or key in COV_ENV_EXTRAS:
|
||||
env[key] = val
|
||||
_forward_coverage_env(env)
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
ironclaw_binary, "--no-onboard",
|
||||
stdin=asyncio.subprocess.DEVNULL,
|
||||
@@ -375,6 +506,7 @@ async def http_channel_server_without_secret(
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=env,
|
||||
)
|
||||
startup_kill_attempted = False
|
||||
gateway_url = f"http://127.0.0.1:{gateway_port}"
|
||||
http_base_url = f"http://127.0.0.1:{http_port}"
|
||||
try:
|
||||
@@ -383,15 +515,17 @@ async def http_channel_server_without_secret(
|
||||
yield http_base_url
|
||||
except TimeoutError:
|
||||
# Dump stderr so CI logs show why the server failed to start
|
||||
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, Exception):
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
stderr_text = stderr_bytes.decode("utf-8", errors="replace")
|
||||
proc.kill()
|
||||
pytest.fail(
|
||||
f"ironclaw server without webhook secret failed to start on ports "
|
||||
f"gateway={gateway_port}, http={http_port} "
|
||||
@@ -399,14 +533,15 @@ async def http_channel_server_without_secret(
|
||||
)
|
||||
finally:
|
||||
if proc.returncode is None:
|
||||
# Use SIGINT (not SIGTERM) so tokio's ctrl_c handler triggers a
|
||||
# graceful shutdown. This lets the LLVM coverage runtime run its
|
||||
# atexit handler and flush .profraw files for cargo-llvm-cov.
|
||||
proc.send_signal(signal.SIGINT)
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=10)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
if startup_kill_attempted:
|
||||
await _stop_process(proc, timeout=2)
|
||||
else:
|
||||
# Use SIGINT (not SIGTERM) so tokio's ctrl_c handler triggers a
|
||||
# graceful shutdown. This lets the LLVM coverage runtime run its
|
||||
# atexit handler and flush .profraw files for cargo-llvm-cov.
|
||||
await _stop_process(proc, sig=signal.SIGINT, timeout=10)
|
||||
if proc.returncode is None:
|
||||
await _stop_process(proc, timeout=2)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
|
||||
@@ -34,6 +34,15 @@ TOOL_CALL_PATTERNS = [
|
||||
"body": {"label": m.group("label")},
|
||||
},
|
||||
),
|
||||
(
|
||||
re.compile(r"check gmail unread|gmail unread", re.IGNORECASE),
|
||||
"gmail",
|
||||
lambda _: {
|
||||
"action": "list_messages",
|
||||
"query": "is:unread",
|
||||
"max_results": 1,
|
||||
},
|
||||
),
|
||||
(re.compile(r"what time|current time", re.IGNORECASE), "time", lambda _: {"operation": "now"}),
|
||||
(
|
||||
re.compile(
|
||||
@@ -91,6 +100,15 @@ TOOL_CALL_PATTERNS = [
|
||||
]
|
||||
|
||||
|
||||
def _new_oauth_state() -> dict:
|
||||
return {
|
||||
"exchange_count": 0,
|
||||
"refresh_count": 0,
|
||||
"last_exchange": None,
|
||||
"last_refresh": None,
|
||||
}
|
||||
|
||||
|
||||
def _last_user_content(messages: list[dict]) -> str:
|
||||
for msg in reversed(messages):
|
||||
if msg.get("role") == "user":
|
||||
@@ -272,6 +290,12 @@ async def oauth_exchange(request: web.Request) -> web.Response:
|
||||
specific token params such as RFC 8707 `resource` are forwarded here.
|
||||
"""
|
||||
data = await request.post()
|
||||
oauth_state = request.app["oauth_state"]
|
||||
oauth_state["exchange_count"] += 1
|
||||
oauth_state["last_exchange"] = {
|
||||
"authorization": request.headers.get("Authorization"),
|
||||
"form": dict(data),
|
||||
}
|
||||
code = data.get("code", "")
|
||||
access_token_field = data.get("access_token_field", "access_token")
|
||||
|
||||
@@ -290,6 +314,39 @@ async def oauth_exchange(request: web.Request) -> web.Response:
|
||||
})
|
||||
|
||||
|
||||
async def oauth_refresh(request: web.Request) -> web.Response:
|
||||
"""Mock OAuth token refresh proxy for hosted refresh E2E tests."""
|
||||
data = await request.post()
|
||||
oauth_state = request.app["oauth_state"]
|
||||
oauth_state["refresh_count"] += 1
|
||||
oauth_state["last_refresh"] = {
|
||||
"authorization": request.headers.get("Authorization"),
|
||||
"form": dict(data),
|
||||
}
|
||||
|
||||
if request.headers.get("Authorization") != "Bearer e2e-test-token":
|
||||
return web.json_response({"error": "invalid_gateway_auth"}, status=401)
|
||||
if data.get("client_id") != "hosted-google-client-id":
|
||||
return web.json_response({"error": "invalid_client_id"}, status=400)
|
||||
if "client_secret" in data:
|
||||
return web.json_response({"error": "unexpected_client_secret"}, status=400)
|
||||
|
||||
return web.json_response({
|
||||
"access_token": "mock-refreshed-access-token",
|
||||
"refresh_token": "mock-rotated-refresh-token",
|
||||
"expires_in": 3600,
|
||||
})
|
||||
|
||||
|
||||
async def oauth_state_handler(request: web.Request) -> web.Response:
|
||||
return web.json_response(request.app["oauth_state"])
|
||||
|
||||
|
||||
async def oauth_reset(request: web.Request) -> web.Response:
|
||||
request.app["oauth_state"] = _new_oauth_state()
|
||||
return web.json_response({"ok": True})
|
||||
|
||||
|
||||
async def models(_request: web.Request) -> web.Response:
|
||||
return web.json_response({
|
||||
"object": "list",
|
||||
@@ -424,12 +481,16 @@ def main():
|
||||
parser.add_argument("--port", type=int, default=0)
|
||||
args = parser.parse_args()
|
||||
app = web.Application()
|
||||
app["oauth_state"] = _new_oauth_state()
|
||||
# Register both /v1/ and non-/v1/ paths (rig-core omits the /v1/ prefix)
|
||||
app.router.add_post("/v1/chat/completions", chat_completions)
|
||||
app.router.add_post("/chat/completions", chat_completions)
|
||||
app.router.add_get("/v1/models", models)
|
||||
app.router.add_get("/models", models)
|
||||
app.router.add_post("/oauth/exchange", oauth_exchange)
|
||||
app.router.add_post("/oauth/refresh", oauth_refresh)
|
||||
app.router.add_get("/__mock/oauth/state", oauth_state_handler)
|
||||
app.router.add_post("/__mock/oauth/reset", oauth_reset)
|
||||
# Mock MCP server endpoints
|
||||
app.router.add_post("/mcp", mcp_endpoint)
|
||||
app.router.add_post("/mcp-400", mcp_endpoint_400)
|
||||
|
||||
@@ -0,0 +1,227 @@
|
||||
"""Hosted OAuth refresh HTTP regression test.
|
||||
|
||||
Runs a real ironclaw binary in hosted mode, expires a stored Gmail access
|
||||
token in the libSQL database, triggers a real gmail tool call through the
|
||||
chat API, and verifies that refresh uses the hosted proxy endpoint.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sqlite3
|
||||
from datetime import datetime, timezone
|
||||
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 _parse_timestamp(value: str | None) -> datetime | None:
|
||||
if value is None:
|
||||
return None
|
||||
return datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
|
||||
|
||||
def _expire_access_token(db_path: str, user_id: str, secret_name: str) -> None:
|
||||
with sqlite3.connect(db_path) as conn:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
UPDATE secrets
|
||||
SET expires_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '-1 hour')
|
||||
WHERE user_id = ?1 AND name = ?2
|
||||
""",
|
||||
(user_id, secret_name),
|
||||
)
|
||||
conn.commit()
|
||||
assert cursor.rowcount == 1, f"Expected one secret row for {user_id}/{secret_name}"
|
||||
|
||||
|
||||
def _find_secret_row(
|
||||
db_path: str,
|
||||
secret_name: str,
|
||||
) -> tuple[str, str | None, str | None]:
|
||||
with sqlite3.connect(db_path) as conn:
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT user_id, expires_at, updated_at
|
||||
FROM secrets
|
||||
WHERE name = ?1
|
||||
ORDER BY updated_at DESC
|
||||
LIMIT 1
|
||||
""",
|
||||
(secret_name,),
|
||||
).fetchone()
|
||||
assert row is not None, f"Missing secret row for {secret_name}"
|
||||
return row[0], row[1], row[2]
|
||||
|
||||
|
||||
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 _reset_mock_oauth_state(mock_base_url: str) -> None:
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(f"{mock_base_url}/__mock/oauth/reset", timeout=10)
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
async def _get_mock_oauth_state(mock_base_url: str) -> dict:
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(f"{mock_base_url}/__mock/oauth/state", timeout=10)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
async def _approve_pending_request(base_url: str, thread_id: str, request_id: str) -> None:
|
||||
response = await api_post(
|
||||
base_url,
|
||||
"/api/chat/approval",
|
||||
json={"request_id": request_id, "action": "approve", "thread_id": thread_id},
|
||||
timeout=15,
|
||||
)
|
||||
assert response.status_code == 202, (
|
||||
f"Approval submission failed: {response.status_code} {response.text[:400]}"
|
||||
)
|
||||
|
||||
|
||||
async def _wait_for_gmail_tool_call(base_url: str, thread_id: str, timeout: float = 30.0) -> dict:
|
||||
approved_request_ids = set()
|
||||
for _ in range(int(timeout * 2)):
|
||||
response = await api_get(
|
||||
base_url,
|
||||
f"/api/chat/history?thread_id={thread_id}",
|
||||
timeout=15,
|
||||
)
|
||||
response.raise_for_status()
|
||||
history = response.json()
|
||||
|
||||
pending = history.get("pending_approval")
|
||||
if pending and pending["request_id"] not in approved_request_ids:
|
||||
await _approve_pending_request(base_url, thread_id, pending["request_id"])
|
||||
approved_request_ids.add(pending["request_id"])
|
||||
|
||||
for turn in history.get("turns", []):
|
||||
for tool_call in turn.get("tool_calls", []):
|
||||
if tool_call.get("name") == "gmail":
|
||||
return history
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
raise AssertionError(f"Timed out waiting for gmail tool call in thread {thread_id}")
|
||||
|
||||
|
||||
async def _wait_for_refresh_request(mock_base_url: str, timeout: float = 20.0) -> dict:
|
||||
for _ in range(int(timeout * 2)):
|
||||
state = await _get_mock_oauth_state(mock_base_url)
|
||||
if state.get("refresh_count") == 1:
|
||||
return state
|
||||
await asyncio.sleep(0.5)
|
||||
raise AssertionError("Timed out waiting for exactly one OAuth refresh request")
|
||||
|
||||
|
||||
async def test_hosted_gmail_oauth_refresh_uses_proxy(hosted_oauth_refresh_server):
|
||||
server = hosted_oauth_refresh_server["base_url"]
|
||||
db_path = hosted_oauth_refresh_server["db_path"]
|
||||
mock_base_url = hosted_oauth_refresh_server["mock_llm_url"]
|
||||
|
||||
install_response = await api_post(
|
||||
server,
|
||||
"/api/extensions/install",
|
||||
json={"name": "gmail"},
|
||||
timeout=180,
|
||||
)
|
||||
assert install_response.status_code == 200, install_response.text
|
||||
assert install_response.json().get("success") is True
|
||||
|
||||
setup_response = await api_post(
|
||||
server,
|
||||
"/api/extensions/gmail/setup",
|
||||
json={"secrets": {}},
|
||||
timeout=30,
|
||||
)
|
||||
assert setup_response.status_code == 200, setup_response.text
|
||||
setup_data = setup_response.json()
|
||||
assert setup_data.get("success") is True, setup_data
|
||||
auth_url = setup_data.get("auth_url")
|
||||
assert auth_url, setup_data
|
||||
auth_params = parse_qs(urlparse(auth_url).query)
|
||||
assert auth_params.get("client_id") == ["hosted-google-client-id"]
|
||||
|
||||
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]
|
||||
callback_body = callback_response.text.lower()
|
||||
assert "connected" in callback_body or "success" in callback_body
|
||||
|
||||
gmail = await _get_extension(server, "gmail")
|
||||
assert gmail is not None, "gmail should be installed"
|
||||
assert gmail["authenticated"] is True, gmail
|
||||
assert "gmail" in gmail.get("tools", []), gmail
|
||||
|
||||
await _reset_mock_oauth_state(mock_base_url)
|
||||
|
||||
stored_user_id, expires_before, updated_before = _find_secret_row(
|
||||
db_path, "google_oauth_token"
|
||||
)
|
||||
assert _parse_timestamp(expires_before) is not None
|
||||
assert _parse_timestamp(updated_before) is not None
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
_expire_access_token(db_path, stored_user_id, "google_oauth_token")
|
||||
|
||||
thread_response = await api_post(server, "/api/chat/thread/new", timeout=15)
|
||||
assert thread_response.status_code == 200, thread_response.text
|
||||
thread_id = thread_response.json()["id"]
|
||||
|
||||
send_response = await api_post(
|
||||
server,
|
||||
"/api/chat/send",
|
||||
json={"content": "check gmail unread", "thread_id": thread_id},
|
||||
timeout=30,
|
||||
)
|
||||
assert send_response.status_code == 202, send_response.text
|
||||
|
||||
history = await _wait_for_gmail_tool_call(server, thread_id)
|
||||
assert any(
|
||||
tool_call.get("name") == "gmail"
|
||||
for turn in history.get("turns", [])
|
||||
for tool_call in turn.get("tool_calls", [])
|
||||
), history
|
||||
|
||||
oauth_state = await _wait_for_refresh_request(mock_base_url)
|
||||
assert oauth_state["refresh_count"] == 1, oauth_state
|
||||
last_refresh = oauth_state["last_refresh"]
|
||||
assert last_refresh is not None, oauth_state
|
||||
assert last_refresh["authorization"] == "Bearer e2e-test-token"
|
||||
assert last_refresh["form"]["client_id"] == "hosted-google-client-id"
|
||||
assert "client_secret" not in last_refresh["form"], last_refresh
|
||||
|
||||
refreshed_user_id, expires_after, updated_after = _find_secret_row(
|
||||
db_path, "google_oauth_token"
|
||||
)
|
||||
assert refreshed_user_id == stored_user_id
|
||||
expires_after_dt = _parse_timestamp(expires_after)
|
||||
updated_after_dt = _parse_timestamp(updated_after)
|
||||
updated_before_dt = _parse_timestamp(updated_before)
|
||||
assert expires_after_dt is not None
|
||||
assert updated_after_dt is not None
|
||||
assert updated_before_dt is not None
|
||||
assert expires_after_dt > datetime.now(timezone.utc)
|
||||
assert updated_after_dt > updated_before_dt
|
||||
@@ -561,11 +561,7 @@ mod tests {
|
||||
"deploy to production now",
|
||||
);
|
||||
let fired = engine
|
||||
.check_event_triggers(
|
||||
&matching_msg.user_id,
|
||||
&matching_msg.channel,
|
||||
&matching_msg.content,
|
||||
)
|
||||
.check_event_triggers(&matching_msg, &matching_msg.content)
|
||||
.await;
|
||||
assert!(
|
||||
fired >= 1,
|
||||
@@ -584,11 +580,7 @@ mod tests {
|
||||
"check the staging environment",
|
||||
);
|
||||
let fired_neg = engine
|
||||
.check_event_triggers(
|
||||
&non_matching_msg.user_id,
|
||||
&non_matching_msg.channel,
|
||||
&non_matching_msg.content,
|
||||
)
|
||||
.check_event_triggers(&non_matching_msg, &non_matching_msg.content)
|
||||
.await;
|
||||
assert_eq!(fired_neg, 0, "Expected 0 routines fired on non-match");
|
||||
}
|
||||
@@ -652,7 +644,7 @@ mod tests {
|
||||
"deploy to production now",
|
||||
);
|
||||
let guest_fired = engine
|
||||
.check_event_triggers(&guest_msg.user_id, &guest_msg.channel, &guest_msg.content)
|
||||
.check_event_triggers(&guest_msg, &guest_msg.content)
|
||||
.await;
|
||||
assert_eq!(
|
||||
guest_fired, 0,
|
||||
@@ -677,7 +669,7 @@ mod tests {
|
||||
"deploy to production now",
|
||||
);
|
||||
let owner_fired = engine
|
||||
.check_event_triggers(&owner_msg.user_id, &owner_msg.channel, &owner_msg.content)
|
||||
.check_event_triggers(&owner_msg, &owner_msg.content)
|
||||
.await;
|
||||
assert!(
|
||||
owner_fired >= 1,
|
||||
@@ -906,9 +898,7 @@ mod tests {
|
||||
"default",
|
||||
"test-cooldown trigger",
|
||||
);
|
||||
let fired1 = engine
|
||||
.check_event_triggers(&msg.user_id, &msg.channel, &msg.content)
|
||||
.await;
|
||||
let fired1 = engine.check_event_triggers(&msg, &msg.content).await;
|
||||
assert!(fired1 >= 1, "First fire should work");
|
||||
|
||||
// Give spawn time, then update last_run_at to simulate recent execution.
|
||||
@@ -923,9 +913,7 @@ mod tests {
|
||||
engine.refresh_event_cache().await;
|
||||
|
||||
// Second fire should be blocked by cooldown.
|
||||
let fired2 = engine
|
||||
.check_event_triggers(&msg.user_id, &msg.channel, &msg.content)
|
||||
.await;
|
||||
let fired2 = engine.check_event_triggers(&msg, &msg.content).await;
|
||||
assert_eq!(fired2, 0, "Second fire should be blocked by cooldown");
|
||||
}
|
||||
|
||||
@@ -1095,9 +1083,7 @@ mod tests {
|
||||
engine.refresh_event_cache().await;
|
||||
|
||||
let msg = IncomingMessage::new("test", "default", "DISABLE_ME");
|
||||
let fired_before = engine
|
||||
.check_event_triggers(&msg.user_id, &msg.channel, &msg.content)
|
||||
.await;
|
||||
let fired_before = engine.check_event_triggers(&msg, &msg.content).await;
|
||||
assert!(fired_before >= 1, "Expected routine to fire before disable");
|
||||
|
||||
// Simulate what routines_toggle_handler now does: update DB, then refresh.
|
||||
@@ -1106,9 +1092,7 @@ mod tests {
|
||||
db.update_routine(&routine).await.expect("update_routine");
|
||||
engine.refresh_event_cache().await;
|
||||
|
||||
let fired_after = engine
|
||||
.check_event_triggers(&msg.user_id, &msg.channel, &msg.content)
|
||||
.await;
|
||||
let fired_after = engine.check_event_triggers(&msg, &msg.content).await;
|
||||
assert_eq!(
|
||||
fired_after, 0,
|
||||
"Disabled routine must not fire after cache refresh"
|
||||
@@ -1134,10 +1118,7 @@ mod tests {
|
||||
|
||||
let msg = IncomingMessage::new("test", "default", "DELETE_ME");
|
||||
assert!(
|
||||
engine
|
||||
.check_event_triggers(&msg.user_id, &msg.channel, &msg.content)
|
||||
.await
|
||||
>= 1,
|
||||
engine.check_event_triggers(&msg, &msg.content).await >= 1,
|
||||
"Expected routine to fire before delete"
|
||||
);
|
||||
|
||||
@@ -1146,9 +1127,7 @@ mod tests {
|
||||
engine.refresh_event_cache().await;
|
||||
|
||||
assert_eq!(
|
||||
engine
|
||||
.check_event_triggers(&msg.user_id, &msg.channel, &msg.content)
|
||||
.await,
|
||||
engine.check_event_triggers(&msg, &msg.content).await,
|
||||
0,
|
||||
"Deleted routine must not fire after cache refresh"
|
||||
);
|
||||
@@ -1462,8 +1441,9 @@ mod tests {
|
||||
db.create_routine(&routine).await.expect("create_routine");
|
||||
engine.refresh_event_cache().await;
|
||||
|
||||
let trigger_msg = IncomingMessage::new("test", "default", "owner-gate");
|
||||
let fired = engine
|
||||
.check_event_triggers("default", "test", "owner-gate")
|
||||
.check_event_triggers(&trigger_msg, &trigger_msg.content)
|
||||
.await;
|
||||
assert_eq!(fired, 1, "expected one matching event routine");
|
||||
|
||||
|
||||
@@ -19,10 +19,13 @@ use axum::middleware;
|
||||
use axum::routing::{get, post};
|
||||
use tower::ServiceExt;
|
||||
|
||||
use ironclaw::channels::IncomingMessage;
|
||||
use ironclaw::channels::web::auth::{
|
||||
AuthenticatedUser, MultiAuthState, UserIdentity, auth_middleware,
|
||||
};
|
||||
use ironclaw::channels::web::server::{GatewayState, PerUserRateLimiter, RateLimiter};
|
||||
use ironclaw::channels::web::server::{
|
||||
GatewayState, PerUserRateLimiter, RateLimiter, start_server,
|
||||
};
|
||||
use ironclaw::channels::web::sse::SseManager;
|
||||
use ironclaw::channels::web::test_helpers::TestGatewayBuilder;
|
||||
use ironclaw::channels::web::ws::WsConnectionTracker;
|
||||
@@ -37,6 +40,9 @@ const ALICE_TOKEN: &str = "tok-alice-secret";
|
||||
const BOB_TOKEN: &str = "tok-bob-secret";
|
||||
const ALICE_USER_ID: &str = "alice";
|
||||
const BOB_USER_ID: &str = "bob";
|
||||
const OWNER_TOKEN: &str = "tok-owner-secret";
|
||||
const OWNER_SCOPE_ID: &str = "owner-scope";
|
||||
const GATEWAY_SENDER_ID: &str = "gateway-sender";
|
||||
|
||||
/// Build a MultiAuthState with two users.
|
||||
fn two_user_auth() -> MultiAuthState {
|
||||
@@ -537,7 +543,8 @@ fn gateway_state_has_multi_tenant_fields() {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: "fallback".to_string(), // Multi-tenant: renamed from user_id
|
||||
owner_id: "fallback".to_string(),
|
||||
default_sender_id: "fallback".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
@@ -553,7 +560,8 @@ fn gateway_state_has_multi_tenant_fields() {
|
||||
active_config: Default::default(),
|
||||
};
|
||||
|
||||
assert_eq!(state.default_user_id, "fallback");
|
||||
assert_eq!(state.owner_id, "fallback");
|
||||
assert_eq!(state.default_sender_id, "fallback");
|
||||
assert!(state.workspace_pool.is_none());
|
||||
}
|
||||
|
||||
@@ -572,6 +580,69 @@ async fn start_multi_user_server() -> (SocketAddr, Arc<GatewayState>) {
|
||||
.expect("Failed to start multi-user test server")
|
||||
}
|
||||
|
||||
async fn start_owner_scoped_sender_server() -> (
|
||||
SocketAddr,
|
||||
Arc<GatewayState>,
|
||||
tokio::sync::mpsc::Receiver<IncomingMessage>,
|
||||
) {
|
||||
let (agent_tx, agent_rx) = tokio::sync::mpsc::channel(64);
|
||||
|
||||
let mut tokens = HashMap::new();
|
||||
tokens.insert(
|
||||
OWNER_TOKEN.to_string(),
|
||||
UserIdentity {
|
||||
user_id: OWNER_SCOPE_ID.to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
},
|
||||
);
|
||||
tokens.insert(
|
||||
BOB_TOKEN.to_string(),
|
||||
UserIdentity {
|
||||
user_id: BOB_USER_ID.to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
},
|
||||
);
|
||||
|
||||
let state = Arc::new(GatewayState {
|
||||
msg_tx: tokio::sync::RwLock::new(Some(agent_tx)),
|
||||
sse: Arc::new(SseManager::new()),
|
||||
workspace: None,
|
||||
workspace_pool: None,
|
||||
session_manager: None,
|
||||
log_broadcaster: None,
|
||||
log_level_handle: None,
|
||||
extension_manager: None,
|
||||
tool_registry: None,
|
||||
store: None,
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
owner_id: OWNER_SCOPE_ID.to_string(),
|
||||
default_sender_id: GATEWAY_SENDER_ID.to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
skill_registry: None,
|
||||
skill_catalog: None,
|
||||
chat_rate_limiter: PerUserRateLimiter::new(30, 60),
|
||||
oauth_rate_limiter: RateLimiter::new(10, 60),
|
||||
webhook_rate_limiter: RateLimiter::new(10, 60),
|
||||
registry_entries: Vec::new(),
|
||||
cost_guard: None,
|
||||
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
|
||||
startup_time: std::time::Instant::now(),
|
||||
active_config: Default::default(),
|
||||
});
|
||||
|
||||
let auth = MultiAuthState::multi(tokens);
|
||||
let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
|
||||
let bound = start_server(addr, state.clone(), auth)
|
||||
.await
|
||||
.expect("Failed to start owner-scoped sender test server");
|
||||
|
||||
(bound, state, agent_rx)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_server_alice_can_access_protected_endpoint() {
|
||||
let (addr, _state) = start_multi_user_server().await;
|
||||
@@ -677,6 +748,49 @@ async fn full_server_chat_send_accepted_for_alice() {
|
||||
assert_eq!(msg.channel, "gateway");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_server_chat_send_rewrites_sender_only_for_owner_scope_rebind() {
|
||||
let (addr, _state, mut agent_rx) = start_owner_scoped_sender_server().await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let owner_resp = client
|
||||
.post(format!("http://{}/api/chat/send", addr))
|
||||
.header("Authorization", format!("Bearer {}", OWNER_TOKEN))
|
||||
.header("Content-Type", "application/json")
|
||||
.body(r#"{"content":"hello from owner"}"#)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(owner_resp.status(), 202);
|
||||
|
||||
let owner_msg = tokio::time::timeout(Duration::from_secs(2), agent_rx.recv())
|
||||
.await
|
||||
.expect("Timed out waiting for owner message")
|
||||
.expect("Agent channel closed");
|
||||
assert_eq!(owner_msg.user_id, OWNER_SCOPE_ID);
|
||||
assert_eq!(owner_msg.sender_id, GATEWAY_SENDER_ID);
|
||||
assert_eq!(owner_msg.content, "hello from owner");
|
||||
|
||||
let other_resp = client
|
||||
.post(format!("http://{}/api/chat/send", addr))
|
||||
.header("Authorization", format!("Bearer {}", BOB_TOKEN))
|
||||
.header("Content-Type", "application/json")
|
||||
.body(r#"{"content":"hello from bob"}"#)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(other_resp.status(), 202);
|
||||
|
||||
let other_msg = tokio::time::timeout(Duration::from_secs(2), agent_rx.recv())
|
||||
.await
|
||||
.expect("Timed out waiting for non-owner message")
|
||||
.expect("Agent channel closed");
|
||||
assert_eq!(other_msg.user_id, BOB_USER_ID);
|
||||
assert_eq!(other_msg.sender_id, BOB_USER_ID);
|
||||
assert_eq!(other_msg.content, "hello from bob");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_server_chat_send_rejected_without_auth() {
|
||||
let (addr, _state) = start_multi_user_server().await;
|
||||
@@ -888,7 +1002,8 @@ async fn start_multi_user_server_with_db() -> (
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: ALICE_USER_ID.to_string(),
|
||||
owner_id: ALICE_USER_ID.to_string(),
|
||||
default_sender_id: ALICE_USER_ID.to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
|
||||
@@ -203,7 +203,8 @@ async fn start_test_server_with_provider(
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: "test-user".to_string(),
|
||||
owner_id: "test-user".to_string(),
|
||||
default_sender_id: "test-user".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: Some(llm_provider),
|
||||
@@ -701,7 +702,8 @@ async fn test_no_llm_provider_returns_503() {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: "test-user".to_string(),
|
||||
owner_id: "test-user".to_string(),
|
||||
default_sender_id: "test-user".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None, // No LLM!
|
||||
|
||||
@@ -226,7 +226,8 @@ impl GatewayWorkflowHarness {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: Some(scheduler_slot.clone()),
|
||||
default_user_id: user_id.clone(),
|
||||
owner_id: user_id.clone(),
|
||||
default_sender_id: user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: Some(Arc::clone(&components.llm)),
|
||||
|
||||
@@ -51,7 +51,8 @@ async fn start_test_server() -> (
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: "test-user".to_string(),
|
||||
owner_id: "test-user".to_string(),
|
||||
default_sender_id: "test-user".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
|
||||
Reference in New Issue
Block a user